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

Ensure that `gpu.func` is only used within the dedicated `gpu.module`.
Implement the constraint to the GPU dialect and adopt test cases.

Differential Revision: https://reviews.llvm.org/D78541

diff --git a/mlir/include/mlir/Dialect/GPU/GPUBase.td b/mlir/include/mlir/Dialect/GPU/GPUBase.td
index 39e2f1a..16ce93f 100644 (file)
--- a/mlir/include/mlir/Dialect/GPU/GPUBase.td
+++ b/mlir/include/mlir/Dialect/GPU/GPUBase.td
@@ -33,9 +33,6 @@ def GPU_Dialect : Dialect {
     /// 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);

diff --git a/mlir/include/mlir/Dialect/GPU/GPUOps.td b/mlir/include/mlir/Dialect/GPU/GPUOps.td
index 342b36b..7254f0f 100644 (file)
--- a/mlir/include/mlir/Dialect/GPU/GPUOps.td
+++ b/mlir/include/mlir/Dialect/GPU/GPUOps.td
@@ -85,7 +85,8 @@ def GPU_ThreadIdOp : GPU_IndexOp<"thread_id"> {
   }];
 }
 
-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";
 

diff --git a/mlir/test/Dialect/GPU/all-reduce-max.mlir b/mlir/test/Dialect/GPU/all-reduce-max.mlir
index 9c227a8..5c94bd4 100644 (file)
--- a/mlir/test/Dialect/GPU/all-reduce-max.mlir
+++ b/mlir/test/Dialect/GPU/all-reduce-max.mlir
@@ -1,8 +1,8 @@
 // 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 {

diff --git a/mlir/test/Dialect/GPU/all-reduce.mlir b/mlir/test/Dialect/GPU/all-reduce.mlir
index 94ddf8c..ff79863 100644 (file)
--- a/mlir/test/Dialect/GPU/all-reduce.mlir
+++ b/mlir/test/Dialect/GPU/all-reduce.mlir
@@ -1,8 +1,8 @@
 // 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 {

diff --git a/mlir/test/Dialect/GPU/invalid.mlir b/mlir/test/Dialect/GPU/invalid.mlir
index be02dec..b0cc4dd 100644 (file)
--- a/mlir/test/Dialect/GPU/invalid.mlir
+++ b/mlir/test/Dialect/GPU/invalid.mlir
@@ -88,6 +88,17 @@ module attributes {gpu.container_module} {
 
 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) {
@@ -336,7 +347,7 @@ module {
 // -----
 
 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
@@ -347,7 +358,7 @@ module {
 // -----
 
 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
@@ -358,7 +369,7 @@ module {
 // -----
 
 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
@@ -369,7 +380,7 @@ module {
 // -----
 
 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
@@ -380,7 +391,7 @@ module {
 // -----
 
 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

diff --git a/mlir/test/Dialect/GPU/ops.mlir b/mlir/test/Dialect/GPU/ops.mlir
index f500d71..0a996a5 100644 (file)
--- a/mlir/test/Dialect/GPU/ops.mlir
+++ b/mlir/test/Dialect/GPU/ops.mlir
@@ -83,7 +83,7 @@ module attributes {gpu.container_module} {
     return
   }
 
-  module @gpu_funcs attributes {gpu.kernel_module} {
+  gpu.module @gpu_funcs {
     // CHECK-LABEL: gpu.func @kernel_1({{.*}}: f32)
     // CHECK:       workgroup
     // CHECK:       private

diff --git a/mlir/test/Dialect/GPU/promotion.mlir b/mlir/test/Dialect/GPU/promotion.mlir
index bb5a93c..7a51048 100644 (file)
--- a/mlir/test/Dialect/GPU/promotion.mlir
+++ b/mlir/test/Dialect/GPU/promotion.mlir
@@ -1,6 +1,7 @@
-// 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>
@@ -49,7 +50,8 @@ module @foo attributes {gpu.kernel_module} {
 
 // -----
 
-module @foo attributes {gpu.kernel_module} {
+gpu.module @foo {
+
   // Verify that the attribution was indeed introduced
   // CHECK-LABEL: @memref5d
   // CHECK-SAME: (%[[arg:.*]]: memref<8x7x6x5x4xf32>
@@ -101,7 +103,8 @@ module @foo attributes {gpu.kernel_module} {
 
 // -----
 
-module @foo attributes {gpu.kernel_module} {
+gpu.module @foo {
+
   // Check that attribution insertion works fine.
   // CHECK-LABEL: @insert
   // CHECK-SAME: (%{{.*}}: memref<4xf32>