Redo subgroup id

Author: lialan

mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td

@@ -204,10 +204,7 @@ def ROCDL_ReadlaneOp : ROCDL_IntrOp<"readlane", [], [0], [AllTypesMatch<["res",
    }];
 }
 
-// The LLVM intrinsic function name is rather mouthful,
-// so here we opt to use a shorter rocdl name.
-def ROCDL_WaveIdOp :  LLVM_IntrOpBase<ROCDL_Dialect, "wave_id",
-                        "amdgcn_s_get_waveid_in_workgroup", [], [], [], 1>,
+def ROCDL_WaveIdOp : ROCDL_IntrOp<"s.get.waveid.in.workgroup", [], [], [Pure], 1>,
   Arguments<(ins)> {
   let results = (outs LLVM_Type:$res);
   let assemblyFormat = "attr-dict `:` type($res)";

mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp

@@ -214,16 +214,32 @@ struct GPUSubgroupIdOpToROCDL final
   LogicalResult
   matchAndRewrite(gpu::SubgroupIdOp op, gpu::SubgroupIdOp::Adaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    if (chipset.majorVersion < 10) {
-      return rewriter.notifyMatchFailure(
-          op, "SubgroupIdOp is not yet supported on this architecture");
-    }
-
     auto int32Type = IntegerType::get(rewriter.getContext(), 32);
-    Value waveIdOp = rewriter.create<ROCDL::WaveIdOp>(op.getLoc(), int32Type);
-    waveIdOp = truncOrExtToLLVMType(rewriter, op.getLoc(), waveIdOp,
-                                    *getTypeConverter());
-    rewriter.replaceOp(op, {waveIdOp});
+    auto loc = op.getLoc();
+    LLVM::IntegerOverflowFlags flags =
+        LLVM::IntegerOverflowFlags::nsw | LLVM::IntegerOverflowFlags::nuw;
+    // w_id.x + w_dim.x * (w_id.y + w_dim.y * w_id.z)) / subgroup_size
+    Value workitemIdX = rewriter.create<ROCDL::ThreadIdXOp>(loc, int32Type);
+    Value workitemIdY = rewriter.create<ROCDL::ThreadIdYOp>(loc, int32Type);
+    Value workitemIdZ = rewriter.create<ROCDL::ThreadIdZOp>(loc, int32Type);
+    Value workitemDimX = rewriter.create<ROCDL::BlockDimXOp>(loc, int32Type);
+    Value workitemDimY = rewriter.create<ROCDL::BlockDimYOp>(loc, int32Type);
+    Value dimYxIdZ = rewriter.create<LLVM::MulOp>(loc, int32Type, workitemDimY,
+                                                  workitemIdZ, flags);
+    Value dimYxIdZPlusIdY = rewriter.create<LLVM::AddOp>(
+        loc, int32Type, dimYxIdZ, workitemIdY, flags);
+    Value dimYxIdZPlusIdYTimesDimX = rewriter.create<LLVM::MulOp>(
+        loc, int32Type, workitemDimX, dimYxIdZPlusIdY, flags);
+    Value workitemIdXPlusDimYxIdZPlusIdYTimesDimX =
+        rewriter.create<LLVM::AddOp>(loc, int32Type, workitemIdX,
+                                     dimYxIdZPlusIdYTimesDimX, flags);
+    Value subgroupSize = rewriter.create<LLVM::ConstantOp>(
+        loc, IntegerType::get(rewriter.getContext(), 32), 64);
+    Value waveIdOp = rewriter.create<LLVM::SDivOp>(
+        loc, workitemIdXPlusDimYxIdZPlusIdYTimesDimX, subgroupSize);
+
+    rewriter.replaceOp(op, {truncOrExtToLLVMType(rewriter, loc, waveIdOp,
+                                                 *getTypeConverter())});
     return success();
   }
 };

mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl-chipset.mlir

@@ -740,3 +740,25 @@ gpu.module @test_module {
 gpu.module @test_custom_data_layout attributes {llvm.data_layout = "e"} {
 
 }
+
+// -----
+
+gpu.module @test_module {
+  // CHECK-LABEL: func @gpu_subgroup_id()
+  func.func @gpu_subgroup_id() -> (index) {
+    // CHECK: %[[widx:.*]] = rocdl.workitem.id.x : i32
+    // CHECK: %[[widy:.*]] = rocdl.workitem.id.y : i32
+    // CHECK: %[[widz:.*]] = rocdl.workitem.id.z : i32
+    // CHECK: %[[dimx:.*]] = rocdl.workgroup.dim.x : i32
+    // CHECK: %[[dimy:.*]] = rocdl.workgroup.dim.y : i32
+    // CHECK: %[[int5:.*]] = llvm.mul %[[dimy]], %[[widz]] overflow<nsw, nuw> : i32
+    // CHECK: %[[int6:.*]] = llvm.add %[[int5]], %[[widy]] overflow<nsw, nuw> : i32
+    // CHECK: %[[int7:.*]] = llvm.mul %[[dimx]], %[[int6]] overflow<nsw, nuw> : i32
+    // CHECK: %[[int8:.*]] = llvm.add %[[widx]], %[[int7]] overflow<nsw, nuw> : i32
+    // CHECK: %[[ssize:.*]] = llvm.mlir.constant(64 : i32) : i32
+    // CHECK: = llvm.sdiv %[[int8]], %[[ssize]] : i32
+    // CHECK: = llvm.sext %10 : i32 to i64
+    %subgroupId = gpu.subgroup_id : index
+    func.return  %subgroupId :  index
+  }
+}