[MLIR][GPU] Make max flat work group size for ROCDL kernels configurable

[krzysz00]

Summary

While the default value for the amdgpu-flat-work-group-size attribute,
"1, 256", matches the defaults from Clang, some users of the ROCDL dialect,
namely Tensorflow, use larger workgroups, such as 1024. Therefore,
instead of hardcoding this value, we add a rocdl.max_flat_work_group_size
attribute that can be set on GPU kernels to override the default value.

[whchung]

How about add a unit test to check the attribute can be overriden?

mlir/lib/Target/LLVMIR/Dialect/ROCDL/ROCDLToLLVMIRTranslation.cpp

21

I assume we can live without this header file?

[krzysz00]

• Fix bug, add tests

[krzysz00]

I've added a unit test

mlir/lib/Target/LLVMIR/Dialect/ROCDL/ROCDLToLLVMIRTranslation.cpp

21

We do need one, since I'm using raw_svector_ostream below

mlir/lib/Target/LLVMIR/Dialect/ROCDL/ROCDLToLLVMIRTranslation.cpp

 
 #include "mlir/Target/LLVMIR/Dialect/ROCDL/ROCDLToLLVMIRTranslation.h"
 #include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Target/LLVMIR/ModuleTranslation.h"
 
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/IntrinsicsAMDGPU.h"
+#include "llvm/Support/raw_ostream.h"
 
 using namespace mlir;
 using namespace mlir::LLVM;
 using mlir::LLVM::detail::createIntrinsicCall;
 
 // Create a call to ROCm-Device-Library function
 // Currently this routine will work only for calling ROCDL functions that
 // take a single int32 argument. It is likely that the interface of this
amendOperation(Operation *op, NamedAttribute attribute,
                  LLVM::ModuleTranslation &moduleTranslation) const final {
     if (attribute.getName() == ROCDL::ROCDLDialect::getKernelFuncAttrName()) {
       auto func = dyn_cast<LLVM::LLVMFuncOp>(op);
       if (!func)
         return failure();
 
       // For GPU kernels,
       // 1. Insert AMDGPU_KERNEL calling convention.
-      // 2. Insert amdgpu-flat-workgroup-size(1, 256) attribute.
+      // 2. Insert amdgpu-flat-work-group-size(1, 256) attribute unless the user
+      // has overriden this value - 256 is the default in clang
       // 3. Insert amdgpu-implicitarg-num-bytes=56 (which must be set on OpenCL
       // and HIP kernels per Clang)
       llvm::Function *llvmFunc =
           moduleTranslation.lookupFunction(func.getName());
       llvmFunc->setCallingConv(llvm::CallingConv::AMDGPU_KERNEL);
+      if (!llvmFunc->hasFnAttribute("amdgpu-flat-work-group-size")) {
         llvmFunc->addFnAttr("amdgpu-flat-work-group-size", "1, 256");
+      }
       llvmFunc->addFnAttr("amdgpu-implicitarg-num-bytes", "56");
     }
+    // Override flat-work-group-size
+    if ("rocdl.max_flat_work_group_size" == attribute.getName()) {
+      auto func = dyn_cast<LLVM::LLVMFuncOp>(op);
+      if (!func)
+        return failure();
+      auto value = attribute.getValue().dyn_cast<IntegerAttr>();
+      if (!value)
+        return failure();
+
+      llvm::Function *llvmFunc =
+          moduleTranslation.lookupFunction(func.getName());
+      llvm::SmallString<8> llvmAttrValue;
+      llvm::raw_svector_ostream attrValueStream(llvmAttrValue);
+      attrValueStream << "1, " << value.getInt();
+      llvmFunc->addFnAttr("amdgpu-flat-work-group-size", llvmAttrValue);
+    }
     return success();
   }
 };
 } // namespace
 
 void mlir::registerROCDLDialectTranslation(DialectRegistry &registry) {
   registry.insert<ROCDL::ROCDLDialect>();
   registry.addDialectInterface<ROCDL::ROCDLDialect,
                                ROCDLDialectLLVMIRTranslationInterface>();
 }
 
 void mlir::registerROCDLDialectTranslation(MLIRContext &context) {
   DialectRegistry registry;
   registerROCDLDialectTranslation(registry);
   context.appendDialectRegistry(registry);
 }

mlir/test/Target/LLVMIR/rocdl.mlir

llvm.func @rocdl_special_regs() -> i32 {
   // CHECK: call i64 @__ockl_get_global_size(i32 1)
   %11 = rocdl.grid.dim.y : i64
   // CHECK: call i64 @__ockl_get_global_size(i32 2)
   %12 = rocdl.grid.dim.z : i64
   llvm.return %1 : i32
 }
 
 llvm.func @kernel_func() attributes {rocdl.kernel} {
-  // CHECK-LABEL: amdgpu_kernel void @kernel_func
+  // CHECK-LABEL: amdgpu_kernel void @kernel_func()
+  // CHECK: #[[$KERNEL_ATTRS:[0-9]+]]
+  llvm.return
+}
+
+llvm.func @kernel_func_workgroups()
+    attributes {rocdl.kernel, rocdl.max_flat_work_group_size = 1024 : index} {
+  // CHECK-LABEL: amdgpu_kernel void @kernel_func_workgroups()
+  // CHECK: #[[$KERNEL_WORKGROUP_ATTRS:[0-9]+]]
   llvm.return
 }
 
 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.func @rocdl.mubuf(%rsrc : vector<4xi32>, %vindex : i32,
   // CHECK: call void @llvm.amdgcn.buffer.store.v2f32(<2 x float> %{{.*}}, <4 x i32> %{{.*}}, i32 %{{.*}}, i32 %{{.*}}, i1 {{.*}}, i1 {{.*}})
   rocdl.buffer.store %vdata2, %rsrc, %vindex, %offset, %glc, %slc : vector<2xf32>
   // CHECK: call void @llvm.amdgcn.buffer.store.v4f32(<4 x float> %{{.*}}, <4 x i32> %{{.*}}, i32 %{{.*}}, i32 %{{.*}}, i1 {{.*}}, i1 {{.*}})
   rocdl.buffer.store %vdata4, %rsrc, %vindex, %offset, %glc, %slc : vector<4xf32>
 
   llvm.return
 }
 
+// CHECK-DAG: attributes #[[$KERNEL_ATTRS]] = { "amdgpu-flat-work-group-size"="1, 256" "amdgpu-implicitarg-num-bytes"="56" }
+// CHECK-DAG: attributes #[[$KERNEL_WORKGROUP_ATTRS]] = { "amdgpu-flat-work-group-size"="1, 1024"