Outline GPU kernel function into a nested module.

Roll forward of commit 5684a12.

When outlining GPU kernels, put the kernel function inside a nested module. Then use a nested pipeline to generate the cubins, independently per kernel. In a final pass, move the cubins back to the parent module.

PiperOrigin-RevId: 270639748

2 files changed, 33 insertions, 10 deletions

diff --git a/mlir/test/Conversion/GPUToCUDA/insert-cubin-getter.mlir b/mlir/test/Conversion/GPUToCUDA/insert-cubin-getter.mlir
index d2e291f57e7..9e0907f7477 100644
--- a/mlir/test/Conversion/GPUToCUDA/insert-cubin-getter.mlir
+++ b/mlir/test/Conversion/GPUToCUDA/insert-cubin-getter.mlir
@@ -2,9 +2,14 @@
 
 // CHECK: llvm.mlir.global constant @[[global:.*]]("CUBIN")
 
+module attributes {gpu.kernel_module} {
+  func @kernel(!llvm.float, !llvm<"float*">)
+  attributes  {nvvm.cubin = "CUBIN"}
+}
+
 func @kernel(!llvm.float, !llvm<"float*">)
-// CHECK: attributes  {gpu.kernel, nvvm.cubin = "CUBIN", nvvm.cubingetter = @[[getter:.*]]}
-  attributes  {gpu.kernel, nvvm.cubin = "CUBIN"}
+// CHECK: attributes  {gpu.kernel, nvvm.cubingetter = @[[getter:.*]]}
+  attributes  {gpu.kernel}
 
 // CHECK: func @[[getter]]() -> !llvm<"i8*">
 // CHECK: %[[addressof:.*]] = llvm.mlir.addressof @[[global]]
diff --git a/mlir/test/Conversion/GPUToCUDA/lower-nvvm-kernel-to-cubin.mlir b/mlir/test/Conversion/GPUToCUDA/lower-nvvm-kernel-to-cubin.mlir
index 8ddfc1996ef..b6e19989203 100644
--- a/mlir/test/Conversion/GPUToCUDA/lower-nvvm-kernel-to-cubin.mlir
+++ b/mlir/test/Conversion/GPUToCUDA/lower-nvvm-kernel-to-cubin.mlir
@@ -1,8 +1,26 @@
-// RUN: mlir-opt %s --test-kernel-to-cubin | FileCheck %s
-
-func @kernel(%arg0 : !llvm.float, %arg1 : !llvm<"float*">)
-// CHECK: attributes  {gpu.kernel, nvvm.cubin = "CUBIN"}
-  attributes  { gpu.kernel } {
-// CHECK-NOT: llvm.return
-  llvm.return
-}
\ No newline at end of file
+// RUN: mlir-opt %s --test-kernel-to-cubin -split-input-file | FileCheck %s
+
+module attributes {gpu.kernel_module} {
+  func @kernel(%arg0 : !llvm.float, %arg1 : !llvm<"float*">)
+    // CHECK: attributes  {gpu.kernel, nvvm.cubin = "CUBIN"}
+    attributes  { gpu.kernel } {
+    // CHECK-NOT: llvm.return
+    llvm.return
+  }
+}
+
+// -----
+
+module attributes {gpu.kernel_module} {
+  // CHECK: func @kernel_a
+  func @kernel_a()
+    attributes  { gpu.kernel } {
+    llvm.return
+  }
+
+  // CHECK: func @kernel_b
+  func @kernel_b()
+    attributes  { gpu.kernel } {
+    llvm.return
+  }
+}