Introduce gpu.func

Introduce a new function-like operation to the GPU dialect to provide a
placeholder for the execution semantic description and to add support for GPU
memory hierarchy.  This aligns with the overall goal of the dialect to expose
the common abstraction layer for GPU devices, in particular by providing an
MLIR unit of semantics (i.e. an operation) for memory modeling.

This proposal has been discussed in the mailing list:
https://groups.google.com/a/tensorflow.org/d/msg/mlir/RfXNP7Hklsc/MBNN7KhjAgAJ
As decided, the "convergence" aspect of the execution model will be factored
out into a new discussion and therefore is not included in this commit. This
commit only introduces the operation but does not hook it up with the remaining
flow. The intention is to develop the new flow while keeping the old flow
operational and do the switch in a simple, separately reversible commit.

PiperOrigin-RevId: 282357599

2 files changed, 71 insertions, 0 deletions

diff --git a/mlir/test/Dialect/GPU/invalid.mlir b/mlir/test/Dialect/GPU/invalid.mlir
index 9dace254b42..6565c628377 100644
--- a/mlir/test/Dialect/GPU/invalid.mlir
+++ b/mlir/test/Dialect/GPU/invalid.mlir
@@ -360,3 +360,25 @@ func @reduce_incorrect_yield(%arg0 : f32) {
   }) : (f32) -> (f32)
 }
 
+// -----
+
+module {
+  module @gpu_funcs attributes {gpu.kernel_module} {
+    // expected-error @+1 {{custom op 'gpu.func' gpu.func requires named arguments}}
+    gpu.func @kernel_1(f32, f32) {
+    ^bb0(%arg0: f32):
+      gpu.return
+    }
+  }
+}
+
+// -----
+
+module {
+  module @gpu_funcs attributes {gpu.kernel_module} {
+    // expected-error @+1 {{requires 'type' attribute of function type}}
+    "gpu.func"() ({
+      gpu.return
+    }) {sym_name="kernel_1", type=f32} : () -> ()
+  }
+}
diff --git a/mlir/test/Dialect/GPU/ops.mlir b/mlir/test/Dialect/GPU/ops.mlir
index bfc0f154309..e2fd26f254b 100644
--- a/mlir/test/Dialect/GPU/ops.mlir
+++ b/mlir/test/Dialect/GPU/ops.mlir
@@ -112,4 +112,53 @@ module attributes {gpu.container_module} {
     return
   }
 
+  module @gpu_funcs attributes {gpu.kernel_module} {
+    // CHECK-LABEL: gpu.func @kernel_1({{.*}}: f32) -> f32
+    // CHECK:       workgroup
+    // CHECK:       private
+    // CHECK:       attributes
+    gpu.func @kernel_1(%arg0: f32) -> f32
+        workgroup(%arg1: memref<42xf32, 3>)
+        private(%arg2: memref<2xf32, 5>, %arg3: memref<1xf32, 5>)
+        kernel
+        attributes {foo="bar"} {
+      "use"(%arg1) : (memref<42xf32, 3>) -> ()
+      "use"(%arg2) : (memref<2xf32, 5>) -> ()
+      "use"(%arg3) : (memref<1xf32, 5>) -> ()
+      gpu.return
+    }
+
+    // CHECK-LABEL: gpu.func @no_attribution
+    // CHECK: {
+    gpu.func @no_attribution(%arg0: f32) {
+      gpu.return
+    }
+
+    // CHECK-LABEL: @no_attribution_attrs
+    // CHECK:       attributes
+    // CHECK:       {
+    gpu.func @no_attribution_attrs(%arg0: f32) attributes {foo="bar"} {
+      gpu.return
+    }
+
+    // CHECK-LABEL: @workgroup_only
+    // CHECK:       workgroup({{.*}}: {{.*}})
+    // CHECK:       {
+    gpu.func @workgroup_only() workgroup(%arg0: memref<42xf32, 3>) {
+      gpu.return
+    }
+    // CHECK-LABEL: @private_only
+    // CHECK:       private({{.*}}: {{.*}})
+    // CHECK:       {
+    gpu.func @private_only() private(%arg0: memref<2xf32, 5>) {
+      gpu.return
+    }
+
+    // CHECK-LABEL: @empty_attribution
+    // CHECK:       {
+    gpu.func @empty_attribution(%arg0: f32) workgroup() private() {
+      gpu.return
+    }
+  }
+
 }