[MLIR][ROCDL] Add conversion for gpu.lane_id to ROCDL

[sjw36]

Summary

Creates rocdl.lane_id op with llvm conversion to:

__device__ static unsigned int __lane_id() {
    return  __builtin_amdgcn_mbcnt_hi(
               -1, __builtin_amdgcn_mbcnt_lo(-1, 0));
}

[krzysz00]

As a high-level note, wouldn't it make more sense to define a rocdl wrapper around the mbcnt intrinsic and then rewrite to that so than we're not hiding a substantial bit of translation in the LLVM IR builder? I've generally seen the rocdl dialect as the place for 1:1 wrappers around LLVM functionality.

Second, as a minor note, when you have the final lane_id number, would it be possible to put range metadata on it - probably a conservative [0, 63] value, but, still, that'll allow for optimizations.

[arsenm]

In D154666#4480824, @krzysz00 wrote:

I've generally seen the rocdl dialect as the place for 1:1 wrappers around LLVM functionality.

Let's not keep doing that. Just call intrinsics directly. All of those straight intrinsic wrappers introduce more trouble

[arsenm]

In D154666#4480836, @arsenm wrote:

In D154666#4480824, @krzysz00 wrote:

I've generally seen the rocdl dialect as the place for 1:1 wrappers around LLVM functionality.

Let's not keep doing that. Just call intrinsics directly. All of those straight intrinsic wrappers introduce more trouble

I missed the "dialect" part here. Just don't call the ockl wrapper for this

[krzysz00]

The ROCDL dialect is meant to directly represent the AMDGPU-specific intrinsics in LLVM IR within MLIR, and to work with the LLVM dialect (which is LLVM-IR-in-MLIR that you can run through a simple translation layer)

[sjw36]

• updated for review feedback

• in the future, the backend should provide a lane_id intrinsic if the HW ever adds it

[krzysz00]

This overall design works, but I've got minor nitpicks

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

89

Would it be possible to instead explicitly spell out the arguments and the result type, given that they're known from the LLVM?

mlir/test/Target/LLVMIR/rocdl.mlir

60

Could we get a variable capture in this? %[[loCount:.+]] = call i32 ...

[sjw36]

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

89

Hoping this is short lived and very limited usage. So I am inclined to leave it.

[krzysz00]

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

89

I'm thinking this might have some other use in the future we don't know about yet so let's do it right

[sjw36]

• updated per review

[krzysz00]

Overall, looks good

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

91

Nit: We don't need to list the argument types here since they're statically known, but this is a weird rare intrinsic so it's fine.

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

 //===----------------------------------------------------------------------===//
 // ROCDL op definitions
 //===----------------------------------------------------------------------===//
 
 class ROCDL_Op<string mnemonic, list<Trait> traits = []> :
   LLVM_OpBase<ROCDL_Dialect, mnemonic, traits> {
 }
 
+class ROCDL_IntrPure1Op<string mnemonic> :
+  LLVM_IntrOpBase<ROCDL_Dialect, mnemonic,
+  "amdgcn_" # !subst(".", "_", mnemonic), [], [], [Pure], 1>;
+
 //===----------------------------------------------------------------------===//
 // ROCDL special register op definitions
 //===----------------------------------------------------------------------===//
 
 class ROCDL_SpecialRegisterOp<string mnemonic,
     list<Trait> traits = []> :
   ROCDL_Op<mnemonic, !listconcat(traits, [Pure])>,
   Results<(outs LLVM_Type:$res)>, Arguments<(ins)> {
   string llvmBuilder = "$res = createIntrinsicCallWithRange(builder,"
     # "llvm::Intrinsic::amdgcn_" # !subst(".","_", mnemonic)
     # ", op->getAttrOfType<::mlir::DenseI32ArrayAttr>(\"range\"));";
   let assemblyFormat = "attr-dict `:` type($res)";
 }
 
 class ROCDL_DeviceFunctionOp<string mnemonic, string device_function,
                              int parameter, list<Trait> traits = []> :
   ROCDL_Op<mnemonic, !listconcat(traits, [Pure])>,
   Results<(outs LLVM_Type:$res)>, Arguments<(ins)> {
   string llvmBuilder = "$res = createDeviceFunctionCall(builder, \""
   # device_function # "\", " # parameter # ");";
   let assemblyFormat = "attr-dict `:` type($res)";
 }
 
 //===----------------------------------------------------------------------===//
+// Wave-level primitives
+
+class ROCDL_MbcntOp<string mnemonic> :
+    ROCDL_IntrPure1Op<"mbcnt." # mnemonic>,
+  Arguments<(ins I32:$in0, I32:$in1)> {
+  let assemblyFormat = [{
+    $in0 `,` $in1  attr-dict `:` `(` type($in0) `,` type($in1) `)` `->` type($res)
+   }];
+}
+
+def ROCDL_MbcntLoOp : ROCDL_MbcntOp<"lo">;
+def ROCDL_MbcntHiOp : ROCDL_MbcntOp<"hi">;
+
+//===----------------------------------------------------------------------===//
 // Thread index and Block index
 
 def ROCDL_ThreadIdXOp : ROCDL_SpecialRegisterOp<"workitem.id.x">;
 def ROCDL_ThreadIdYOp : ROCDL_SpecialRegisterOp<"workitem.id.y">;
 def ROCDL_ThreadIdZOp : ROCDL_SpecialRegisterOp<"workitem.id.z">;
 
 def ROCDL_BlockIdXOp : ROCDL_SpecialRegisterOp<"workgroup.id.x">;
 def ROCDL_BlockIdYOp : ROCDL_SpecialRegisterOp<"workgroup.id.y">;

mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp

static bool canBeCalledWithBarePointers(gpu::GPUFuncOp func) {
   bool canBeBare = true;
   for (Type type : func.getArgumentTypes())
     if (auto memrefTy = dyn_cast<BaseMemRefType>(type))
       canBeBare &= LLVMTypeConverter::canConvertToBarePtr(memrefTy);
   return canBeBare;
 }
 
 namespace {
+struct GPULaneIdOpToROCDL : ConvertOpToLLVMPattern<gpu::LaneIdOp> {
+  using ConvertOpToLLVMPattern<gpu::LaneIdOp>::ConvertOpToLLVMPattern;
+
+  LogicalResult
+  matchAndRewrite(gpu::LaneIdOp op, gpu::LaneIdOp::Adaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto loc = op->getLoc();
+    MLIRContext *context = rewriter.getContext();
+    // convert to:  %mlo = call @llvm.amdgcn.mbcnt.lo(-1, 0)
+    // followed by: %lid = call @llvm.amdgcn.mbcnt.hi(-1, %mlo)
+
+    Type intTy = IntegerType::get(context, 32);
+    Value zero = rewriter.createOrFold<arith::ConstantIntOp>(loc, 0, 32);
+    Value minus1 = rewriter.createOrFold<arith::ConstantIntOp>(loc, -1, 32);
+    Value mbcntLo =
+        rewriter.create<ROCDL::MbcntLoOp>(loc, intTy, ValueRange{minus1, zero});
+    Value laneId = rewriter.create<ROCDL::MbcntHiOp>(
+        loc, intTy, ValueRange{minus1, mbcntLo});
+    // Truncate or extend the result depending on the index bitwidth specified
+    // by the LLVMTypeConverter options.
+    const unsigned indexBitwidth = getTypeConverter()->getIndexTypeBitwidth();
+    if (indexBitwidth > 32) {
+      laneId = rewriter.create<LLVM::SExtOp>(
+          loc, IntegerType::get(context, indexBitwidth), laneId);
+    } else if (indexBitwidth < 32) {
+      laneId = rewriter.create<LLVM::TruncOp>(
+          loc, IntegerType::get(context, indexBitwidth), laneId);
+    }
+    rewriter.replaceOp(op, {laneId});
+    return success();
+  }
+};
 
 /// Import the GPU Ops to ROCDL Patterns.
 #include "GPUToROCDL.cpp.inc"
 
 // A pass that replaces all occurrences of GPU device operations with their
 // corresponding ROCDL equivalent.
 //
 // This pass only handles device code and is not meant to be run on GPU host
patterns.add<GPUFuncOpLowering>(
                       ROCDL::ROCDLDialect::getKernelFuncAttrName()));
   if (Runtime::HIP == runtime) {
     patterns.add<GPUPrintfOpToHIPLowering>(converter);
   } else if (Runtime::OpenCL == runtime) {
     // Use address space = 4 to match the OpenCL definition of printf()
     patterns.add<GPUPrintfOpToLLVMCallLowering>(converter, /*addressSpace=*/4);
   }
 
+  patterns.add<GPULaneIdOpToROCDL>(converter);
+
   populateOpPatterns<math::AbsFOp>(converter, patterns, "__ocml_fabs_f32",
                                    "__ocml_fabs_f64");
   populateOpPatterns<math::AtanOp>(converter, patterns, "__ocml_atan_f32",
                                    "__ocml_atan_f64");
   populateOpPatterns<math::Atan2Op>(converter, patterns, "__ocml_atan2_f32",
                                     "__ocml_atan2_f64");
   populateOpPatterns<math::CbrtOp>(converter, patterns, "__ocml_cbrt_f32",
                                    "__ocml_cbrt_f64");

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

 // RUN: mlir-opt %s -convert-gpu-to-rocdl='use-opaque-pointers=1' -split-input-file | FileCheck %s
 // RUN: mlir-opt %s -convert-gpu-to-rocdl='index-bitwidth=32 use-opaque-pointers=1' -split-input-file | FileCheck --check-prefix=CHECK32 %s
 
 gpu.module @test_module {
   // CHECK-LABEL: func @gpu_index_ops()
   // CHECK32-LABEL: func @gpu_index_ops()
   func.func @gpu_index_ops()
       -> (index, index, index, index, index, index,
-          index, index, index, index, index, index) {
+          index, index, index, index, index, index,
+          index) {
     // CHECK32-NOT: = llvm.sext %{{.*}} : i32 to i64
 
     // CHECK: rocdl.workitem.id.x : i32
     // CHECK: = llvm.sext %{{.*}} : i32 to i64
     %tIdX = gpu.thread_id x
     // CHECK: rocdl.workitem.id.y : i32
     // CHECK: = llvm.sext %{{.*}} : i32 to i64
     %tIdY = gpu.thread_id y
func.func @gpu_index_ops()
     %gDimX = gpu.grid_dim x
     // CHECK: rocdl.grid.dim.y : i32
     // CHECK: = llvm.sext %{{.*}} : i32 to i64
     %gDimY = gpu.grid_dim y
     // CHECK: rocdl.grid.dim.z : i32
     // CHECK: = llvm.sext %{{.*}} : i32 to i64
     %gDimZ = gpu.grid_dim z
 
+    // CHECK: = rocdl.mbcnt.lo %{{.*}}, %{{.*}} : (i32, i32) -> i32
+    // CHECK: = rocdl.mbcnt.hi %{{.*}}, %{{.*}} : (i32, i32) -> i32
+    // CHECK: = llvm.sext %{{.*}} : i32 to i64
+    %laneId = gpu.lane_id
+
     func.return %tIdX, %tIdY, %tIdZ, %bDimX, %bDimY, %bDimZ,
-               %bIdX, %bIdY, %bIdZ, %gDimX, %gDimY, %gDimZ
+               %bIdX, %bIdY, %bIdZ, %gDimX, %gDimY, %gDimZ,
+               %laneId
         : index, index, index, index, index, index,
-          index, index, index, index, index, index
+          index, index, index, index, index, index,
+          index
   }
 }
 
 // -----
 
 gpu.module @test_module {
   // CHECK-LABEL: func @gpu_index_ops_range()
   // CHECK-SAME: rocdl.flat_work_group_size = "1536,1536"

mlir/test/Target/LLVMIR/rocdl.mlir

attributes {rocdl.kernel,
       rocdl.flat_work_group_size = "128,128",
       rocdl.reqd_work_group_size = array<i32: 16, 4, 2>} {
   // CHECK-LABEL: amdgpu_kernel void @known_block_sizes()
   // CHECK: #[[$KNOWN_BLOCK_SIZE_ATTRS:[0-9]+]]
   // CHECK: !reqd_work_group_size ![[$REQD_WORK_GROUP_SIZE:[0-9]+]]
   llvm.return
 }
 
+llvm.func @rocdl.lane_id() -> i32 {
+  // CHECK: [[mbcntlo:%.+]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+  // CHECK-NEXT: call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 [[mbcntlo]])
+  %0 = llvm.mlir.constant(-1 : i32) : i32
+  %1 = llvm.mlir.constant(0 : i32) : i32
+  %2 = rocdl.mbcnt.lo %0, %1 : (i32, i32) -> i32
+  %3 = rocdl.mbcnt.hi %0, %2 : (i32, i32) -> i32
+  llvm.return %3 : i32
+}
+
 llvm.func @rocdl.barrier() {
   // CHECK:      fence syncscope("workgroup") release
   // CHECK-NEXT: call void @llvm.amdgcn.s.barrier()
   // CHECK-NEXT: fence syncscope("workgroup") acquire
   rocdl.barrier
   llvm.return
 }