[MLIR] Ensure `gpu.func` must be inside a `gpu.module`.

[frgossen]

Summary

Verify that gpu.func can only live inside a gpu.module.
Update tests accordingly.

[rriddle]

mlir/test/Dialect/GPU/promotion.mlir

100

You can remove all of the empty module wrappers, module is implicitly the top-level operation.

mlir/test/lib/Transforms/TestGpuMemoryPromotion.cpp

28 ↗

(On Diff #258911)

Can you just change the pipeline specification in the tests instead of changing this pass?

30 ↗

(On Diff #258911)

MLIR uses camelCase for variable names.

[herhut]

Nice!

mlir/test/Dialect/GPU/promotion.mlir

1–3

Maybe something like

// RUN: mlir-opt -allow-unregistered-dialect -pass-pipeline='module(gpu.module(gpu.func(test-gpu-memory-promotion)))' -split-input-file %s | FileCheck %s

is enough to direct the test pass to the functions through the extra wrapping modules.

55

Why this change?

mlir/test/lib/Transforms/TestGpuMemoryPromotion.cpp

28 ↗

(On Diff #258911)

See my above comment for how to do this. You can then revert these changes,

[frgossen]

Simplify test case.

[frgossen]

mlir/test/Dialect/GPU/promotion.mlir

55

Was not intended. Locally, I did not see this because all the indented lines appear as a change.

100

Didn't know that. Thanks!

mlir/test/lib/Transforms/TestGpuMemoryPromotion.cpp

28 ↗

(On Diff #258911)

That's much nicer. Thanks!

[herhut]

Thanks!

[mravishankar]

@ftynse , @herhut is this giving any added semantic advantage? I have been trying to see if there is a way to target gpu.func directly without going through gpu.launch + outlining. Having the requirement that a gpu.func always live in a gpu.module seems too restrictive. Just want to understand the advantage apart from the fact that all gpu.func need to live in a module "separate" from the host side.

[ftynse]

In D78541#1994763, @mravishankar wrote:

@ftynse , @herhut is this giving any added semantic advantage? I have been trying to see if there is a way to target gpu.func directly without going through gpu.launch + outlining. Having the requirement that a gpu.func always live in a gpu.module seems too restrictive. Just want to understand the advantage apart from the fact that all gpu.func need to live in a module "separate" from the host side.

We always had a requirement that gpu.funcs live in a somehow special module. Originally, it was a module with a kernel_module attribute. This is necessary to pick out the modules that should be translated to device-specific dialects. Later, gpu.module was introduced, but the old condition with the attribute was not removed. This change basically completes the transition.

I do see some value in having a dedicated gpu.module to clearly separate what is intended for device and what is not. Your issue with generating a gpu.func without gpu.launch_func seems mostly orthogonal: this is a matter of writing the conversion, not of the rules the resulting IR must follow. If you have another conversion pass (or just IR building code) that creates a gpu.func, it can equally create a gpu.module first. Not having a gpu.module would put the semantics in a weird state when, for example, the same function is called from both gpu.func and from std.func or if there is some global state mutation.

[frgossen]

@mravishankar, can I land this?

[mravishankar]

Sorry for the delay. Let's land this as is. If I have a more concrete issue/suggestion on this aspect, I can revisit this.

mlir/include/mlir/Dialect/GPU/GPUBase.td

let extraClassDeclaration = [{
     /// kernel modules.
     static StringRef getContainerModuleAttrName() {
       return "gpu.container_module";
     }
     /// Get the name of the attribute used to annotate external kernel
     /// functions.
     static StringRef getKernelFuncAttrName() { return "gpu.kernel"; }
 
-    /// Get the name of the attribute used to annotate kernel modules.
-    static StringRef getKernelModuleAttrName() { return "gpu.kernel_module"; }
-
     /// Returns whether the given function is a kernel function, i.e., has the
     /// 'gpu.kernel' attribute.
     static bool isKernel(Operation *op);
 
     /// Returns the number of workgroup (thread, block) dimensions supported in
     /// the GPU dialect.
     // TODO(zinenko,herhut): consider generalizing this.
     static unsigned getNumWorkgroupDimensions() { return 3; }

mlir/include/mlir/Dialect/GPU/GPUOps.td

let description = [{
     Example:
 
     ```mlir
     %tIdX = "gpu.thread_id"() {dimension = "x"} : () -> (index)
     ```
   }];
 }
 
-def GPU_GPUFuncOp : GPU_Op<"func", [AutomaticAllocationScope, FunctionLike,
+def GPU_GPUFuncOp : GPU_Op<"func", [HasParent<"GPUModuleOp">,
+                                    AutomaticAllocationScope, FunctionLike,
                                     IsolatedFromAbove, Symbol]> {
   let summary = "Function executable on a GPU";
 
   let description = [{
     Defines a function that can be executed on a GPU. This supports memory
     attribution and its body has a particular execution model.
 
     GPU functions are either kernels (as indicated by the `kernel` attribute) or

mlir/test/Dialect/GPU/all-reduce-max.mlir

 // RUN: mlir-opt -test-all-reduce-lowering %s | FileCheck %s
 
 // NOTE: Assertions have been autogenerated by utils/generate-test-checks.py
-// CHECK: module @kernels attributes {gpu.kernel_module} {
-module @kernels attributes {gpu.kernel_module} {
+// CHECK: gpu.module @kernels {
+gpu.module @kernels {
 
   // CHECK-LABEL: gpu.func @kernel(
   // CHECK-SAME: [[VAL_0:%.*]]: f32) workgroup([[VAL_1:%.*]] : memref<32xf32, 3>) kernel {
   gpu.func @kernel(%arg0 : f32) kernel {
     // CHECK:   [[VAL_2:%.*]] = constant 31 : i32
     // CHECK:   [[VAL_3:%.*]] = constant 0 : i32
     // CHECK:   [[VAL_4:%.*]] = constant 0 : index
     // CHECK:   [[VAL_5:%.*]] = constant 32 : i32

mlir/test/Dialect/GPU/all-reduce.mlir

 // RUN: mlir-opt -test-all-reduce-lowering %s | FileCheck %s
 
 // NOTE: Assertions have been autogenerated by utils/generate-test-checks.py
-// CHECK: module @kernels attributes {gpu.kernel_module} {
-module @kernels attributes {gpu.kernel_module} {
+// CHECK: gpu.module @kernels {
+gpu.module @kernels {
 
   // CHECK-LABEL: gpu.func @kernel(
   // CHECK-SAME: [[VAL_0:%.*]]: f32) workgroup([[VAL_1:%.*]] : memref<32xf32, 3>) kernel {
   gpu.func @kernel(%arg0 : f32) kernel {
     // CHECK:   [[VAL_2:%.*]] = constant 31 : i32
     // CHECK:   [[VAL_3:%.*]] = constant 0 : i32
     // CHECK:   [[VAL_4:%.*]] = constant 0 : index
     // CHECK:   [[VAL_5:%.*]] = constant 32 : i32

mlir/test/Dialect/GPU/invalid.mlir

func @launch_func_undefined_module(%sz : index) {
     return
   }
 }
 
 // -----
 
 module attributes {gpu.container_module} {
   module @kernels {
+    // expected-error@+1 {{'gpu.func' op expects parent op 'gpu.module'}}
+    gpu.func @kernel_1(%arg1 : !llvm<"float*">) {
+      gpu.return
+    }
+  }
+}
+
+// -----
+
+module attributes {gpu.container_module} {
+  module @kernels {
   }
 
   func @launch_func_missing_module_attribute(%sz : index) {
     // expected-error@+1 {{kernel module 'kernels' is undefined}}
     "gpu.launch_func"(%sz, %sz, %sz, %sz, %sz, %sz)
     { kernel = @kernels::@kernel_1 }
         : (index, index, index, index, index, index) -> ()
     return
^bb0(%arg0: f32):
       gpu.return
     }
   }
 }
 
 // -----
 
 module {
-  module @gpu_funcs attributes {gpu.kernel_module} {
+  gpu.module @gpu_funcs {
     // expected-error @+1 {{requires 'type' attribute of function type}}
     "gpu.func"() ({
       gpu.return
     }) {sym_name="kernel_1", type=f32} : () -> ()
   }
 }
 
 // -----
 
 module {
-  module @gpu_funcs attributes {gpu.kernel_module} {
+  gpu.module @gpu_funcs {
     // expected-error @+1 {{expected memref type in attribution}}
     gpu.func @kernel() workgroup(%0: i32) {
       gpu.return
     }
   }
 }
 
 // -----
 
 module {
-  module @gpu_funcs attributes {gpu.kernel_module} {
+  gpu.module @gpu_funcs {
     // expected-error @+1 {{expected memory space 3 in attribution}}
     gpu.func @kernel() workgroup(%0: memref<4xf32>) {
       gpu.return
     }
   }
 }
 
 // -----
 
 module {
-  module @gpu_funcs attributes {gpu.kernel_module} {
+  gpu.module @gpu_funcs {
     // expected-error @+1 {{expected memory space 5 in attribution}}
     gpu.func @kernel() private(%0: memref<4xf32>) {
       gpu.return
     }
   }
 }
 
 // -----
 
 module {
-  module @gpu_funcs attributes {gpu.kernel_module} {
+  gpu.module @gpu_funcs {
     // expected-error @+1 {{expected memory space 5 in attribution}}
     gpu.func @kernel() private(%0: memref<4xf32>) {
       gpu.return
     }
   }
 }
 
 // -----

mlir/test/Dialect/GPU/ops.mlir

func @foo() {
     // CHECK: "gpu.launch_func"(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) {kernel = @kernels::@kernel_2} : (index, index, index, index, index, index, f32, memref<?xf32, 1>) -> ()
     "gpu.launch_func"(%cst, %cst, %cst, %cst, %cst, %cst, %0, %1)
     { kernel = @kernels::@kernel_2}
         : (index, index, index, index, index, index, f32, memref<?xf32, 1>) -> ()
 
     return
   }
 
-  module @gpu_funcs attributes {gpu.kernel_module} {
+  gpu.module @gpu_funcs {
     // CHECK-LABEL: gpu.func @kernel_1({{.*}}: f32)
     // CHECK:       workgroup
     // CHECK:       private
     // CHECK:       attributes
     gpu.func @kernel_1(%arg0: f32)
         workgroup(%arg1: memref<42xf32, 3>)
         private(%arg2: memref<2xf32, 5>, %arg3: memref<1xf32, 5>)
         kernel

mlir/test/Dialect/GPU/promotion.mlir

-// RUN: mlir-opt -allow-unregistered-dialect -test-gpu-memory-promotion -split-input-file %s | FileCheck %s
+// RUN: mlir-opt -allow-unregistered-dialect -test-gpu-memory-promotion -pass-pipeline='gpu.module(gpu.func(test-gpu-memory-promotion))' -split-input-file %s | FileCheck %s
+
+gpu.module @foo {
 
-module @foo attributes {gpu.kernel_module} {
   // Verify that the attribution was indeed introduced
   // CHECK-LABEL: @memref3d
   // CHECK-SAME: (%[[arg:.*]]: memref<5x4xf32>
   // CHECK-SAME: workgroup(%[[promoted:.*]] : memref<5x4xf32, 3>)
   gpu.func @memref3d(%arg0: memref<5x4xf32> {gpu.test_promote_workgroup}) kernel {
     // Verify that loop bounds are emitted, the order does not matter.
     // CHECK-DAG: %[[c1:.*]] = constant 1
     // CHECK-DAG: %[[c4:.*]] = constant 4
gpu.func @memref3d(%arg0: memref<5x4xf32> {gpu.test_promote_workgroup}) kernel {
     // CHECK:       %[[v:.*]] = load %[[promoted]][%[[i1]], %[[i2]]]
     // CHECK:       store %[[v]], %[[arg]][%[[i1]], %[[i2]]]
     gpu.return
   }
 }
 
 // -----
 
-module @foo attributes {gpu.kernel_module} {
+gpu.module @foo {
+
   // Verify that the attribution was indeed introduced
   // CHECK-LABEL: @memref5d
   // CHECK-SAME: (%[[arg:.*]]: memref<8x7x6x5x4xf32>
   // CHECK-SAME: workgroup(%[[promoted:.*]] : memref<8x7x6x5x4xf32, 3>)
   gpu.func @memref5d(%arg0: memref<8x7x6x5x4xf32> {gpu.test_promote_workgroup}) kernel {
     // Verify that loop bounds are emitted, the order does not matter.
     // CHECK-DAG: %[[c0:.*]] = constant 0
     // CHECK-DAG: %[[c1:.*]] = constant 1
     // CHECK-DAG: %[[c4:.*]] = constant 4
gpu.func @memref5d(%arg0: memref<8x7x6x5x4xf32> {gpu.test_promote_workgroup}) kernel {
     // CHECK:   loop.for %[[i1:.*]] =
     // CHECK:     loop.for %[[i2:.*]] =
     // CHECK:       loop.for %[[i3:.*]] =
     // CHECK:         loop.for %[[i4:.*]] =
 
     // Verify that the copy is emitted.
     // CHECK:           %[[v:.*]] = load %[[promoted]][%[[i0]], %[[i1]], %[[i2]], %[[i3]], %[[i4]]]
     // CHECK:           store %[[v]], %[[arg]][%[[i0]], %[[i1]], %[[i2]], %[[i3]], %[[i4]]]
     gpu.return
   }
 }
 
 // -----
 
-module @foo attributes {gpu.kernel_module} {
+gpu.module @foo {
+
   // Check that attribution insertion works fine.
   // CHECK-LABEL: @insert
   // CHECK-SAME: (%{{.*}}: memref<4xf32>
   // CHECK-SAME: workgroup(%{{.*}}: memref<1x1xf64, 3>
   // CHECK-SAME: %[[wg2:.*]] : memref<4xf32, 3>)
   // CHECK-SAME: private(%{{.*}}: memref<1x1xi64, 5>)
   gpu.func @insert(%arg0: memref<4xf32> {gpu.test_promote_workgroup})
       workgroup(%arg1: memref<1x1xf64, 3>)
       private(%arg2: memref<1x1xi64, 5>)
       kernel {
     // CHECK: "use"(%[[wg2]])
     "use"(%arg0) : (memref<4xf32>) -> ()
     gpu.return
   }
 }