Automated rollback of commit 5684a12434f923d03b6870f2aa16226bfb0b38b6

PiperOrigin-RevId: 270126672

4 files changed, 44 insertions, 93 deletions

diff --git a/mlir/test/Conversion/GPUToCUDA/insert-cubin-getter.mlir b/mlir/test/Conversion/GPUToCUDA/insert-cubin-getter.mlir
index 9e0907f7477..d2e291f57e7 100644
--- a/mlir/test/Conversion/GPUToCUDA/insert-cubin-getter.mlir
+++ b/mlir/test/Conversion/GPUToCUDA/insert-cubin-getter.mlir
@@ -2,14 +2,9 @@
 
 // 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.cubingetter = @[[getter:.*]]}
-  attributes  {gpu.kernel}
+// CHECK: attributes  {gpu.kernel, nvvm.cubin = "CUBIN", nvvm.cubingetter = @[[getter:.*]]}
+  attributes  {gpu.kernel, nvvm.cubin = "CUBIN"}
 
 // 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 b6e19989203..8ddfc1996ef 100644
--- a/mlir/test/Conversion/GPUToCUDA/lower-nvvm-kernel-to-cubin.mlir
+++ b/mlir/test/Conversion/GPUToCUDA/lower-nvvm-kernel-to-cubin.mlir
@@ -1,26 +1,8 @@
-// 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
-  }
-}
+// 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
diff --git a/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir b/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir
index 02637376622..cf8e7ed1113 100644
--- a/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir
+++ b/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir
@@ -1,37 +1,35 @@
 // RUN: mlir-opt %s -lower-gpu-ops-to-nvvm-ops | FileCheck %s
 
-module attributes {gpu.kernel_module} {
-  // CHECK-LABEL: func @gpu_index_ops()
-  func @gpu_index_ops()
-      attributes { gpu.kernel } {
-    // CHECK: = nvvm.read.ptx.sreg.tid.x : !llvm.i32
-    %tIdX = "gpu.thread_id"() {dimension = "x"} : () -> (index)
-    // CHECK: = nvvm.read.ptx.sreg.tid.y : !llvm.i32
-    %tIdY = "gpu.thread_id"() {dimension = "y"} : () -> (index)
-    // CHECK: = nvvm.read.ptx.sreg.tid.z : !llvm.i32
-    %tIdZ = "gpu.thread_id"() {dimension = "z"} : () -> (index)
+// CHECK-LABEL: func @gpu_index_ops()
+func @gpu_index_ops()
+    attributes { gpu.kernel } {
+  // CHECK: = nvvm.read.ptx.sreg.tid.x : !llvm.i32
+  %tIdX = "gpu.thread_id"() {dimension = "x"} : () -> (index)
+  // CHECK: = nvvm.read.ptx.sreg.tid.y : !llvm.i32
+  %tIdY = "gpu.thread_id"() {dimension = "y"} : () -> (index)
+  // CHECK: = nvvm.read.ptx.sreg.tid.z : !llvm.i32
+  %tIdZ = "gpu.thread_id"() {dimension = "z"} : () -> (index)
 
-    // CHECK: = nvvm.read.ptx.sreg.ntid.x : !llvm.i32
-    %bDimX = "gpu.block_dim"() {dimension = "x"} : () -> (index)
-    // CHECK: = nvvm.read.ptx.sreg.ntid.y : !llvm.i32
-    %bDimY = "gpu.block_dim"() {dimension = "y"} : () -> (index)
-    // CHECK: = nvvm.read.ptx.sreg.ntid.z : !llvm.i32
-    %bDimZ = "gpu.block_dim"() {dimension = "z"} : () -> (index)
+  // CHECK: = nvvm.read.ptx.sreg.ntid.x : !llvm.i32
+  %bDimX = "gpu.block_dim"() {dimension = "x"} : () -> (index)
+  // CHECK: = nvvm.read.ptx.sreg.ntid.y : !llvm.i32
+  %bDimY = "gpu.block_dim"() {dimension = "y"} : () -> (index)
+  // CHECK: = nvvm.read.ptx.sreg.ntid.z : !llvm.i32
+  %bDimZ = "gpu.block_dim"() {dimension = "z"} : () -> (index)
 
-    // CHECK: = nvvm.read.ptx.sreg.ctaid.x : !llvm.i32
-    %bIdX = "gpu.block_id"() {dimension = "x"} : () -> (index)
-    // CHECK: = nvvm.read.ptx.sreg.ctaid.y : !llvm.i32
-    %bIdY = "gpu.block_id"() {dimension = "y"} : () -> (index)
-    // CHECK: = nvvm.read.ptx.sreg.ctaid.z : !llvm.i32
-    %bIdZ = "gpu.block_id"() {dimension = "z"} : () -> (index)
+  // CHECK: = nvvm.read.ptx.sreg.ctaid.x : !llvm.i32
+  %bIdX = "gpu.block_id"() {dimension = "x"} : () -> (index)
+  // CHECK: = nvvm.read.ptx.sreg.ctaid.y : !llvm.i32
+  %bIdY = "gpu.block_id"() {dimension = "y"} : () -> (index)
+  // CHECK: = nvvm.read.ptx.sreg.ctaid.z : !llvm.i32
+  %bIdZ = "gpu.block_id"() {dimension = "z"} : () -> (index)
 
-    // CHECK: = nvvm.read.ptx.sreg.nctaid.x : !llvm.i32
-    %gDimX = "gpu.grid_dim"() {dimension = "x"} : () -> (index)
-    // CHECK: = nvvm.read.ptx.sreg.nctaid.y : !llvm.i32
-    %gDimY = "gpu.grid_dim"() {dimension = "y"} : () -> (index)
-    // CHECK: = nvvm.read.ptx.sreg.nctaid.z : !llvm.i32
-    %gDimZ = "gpu.grid_dim"() {dimension = "z"} : () -> (index)
+  // CHECK: = nvvm.read.ptx.sreg.nctaid.x : !llvm.i32
+  %gDimX = "gpu.grid_dim"() {dimension = "x"} : () -> (index)
+  // CHECK: = nvvm.read.ptx.sreg.nctaid.y : !llvm.i32
+  %gDimY = "gpu.grid_dim"() {dimension = "y"} : () -> (index)
+  // CHECK: = nvvm.read.ptx.sreg.nctaid.z : !llvm.i32
+  %gDimZ = "gpu.grid_dim"() {dimension = "z"} : () -> (index)
 
-    std.return
-  }
+  std.return
 }
diff --git a/mlir/test/Dialect/GPU/outlining.mlir b/mlir/test/Dialect/GPU/outlining.mlir
index fdfe8d08115..07499a305ee 100644
--- a/mlir/test/Dialect/GPU/outlining.mlir
+++ b/mlir/test/Dialect/GPU/outlining.mlir
@@ -1,4 +1,4 @@
-// RUN: mlir-opt -gpu-kernel-outlining -split-input-file -verify-diagnostics %s | FileCheck %s
+// RUN: mlir-opt -gpu-kernel-outlining -split-input-file %s | FileCheck %s
 
 // CHECK-LABEL: func @launch()
 func @launch() {
@@ -35,11 +35,7 @@ func @launch() {
 }
 
 // CHECK-LABEL: func @launch_kernel
-// CHECK-SAME: (f32, memref<?xf32, 1>)
-// CHECK-NEXT: attributes {gpu.kernel}
-
-// CHECK-LABEL: func @launch_kernel
-// CHECK-SAME: (%[[KERNEL_ARG0:.*]]: f32, %[[KERNEL_ARG1:.*]]: memref<?xf32, 1>)
+// CHECK-SAME: (%[[ARG0:.*]]: f32, %[[ARG1:.*]]: memref<?xf32, 1>)
 // CHECK-NEXT: attributes {gpu.kernel}
 // CHECK-NEXT: %[[BID:.*]] = "gpu.block_id"() {dimension = "x"} : () -> index
 // CHECK-NEXT: = "gpu.block_id"() {dimension = "y"} : () -> index
@@ -53,9 +49,9 @@ func @launch() {
 // CHECK-NEXT: %[[BDIM:.*]] = "gpu.block_dim"() {dimension = "x"} : () -> index
 // CHECK-NEXT: = "gpu.block_dim"() {dimension = "y"} : () -> index
 // CHECK-NEXT: = "gpu.block_dim"() {dimension = "z"} : () -> index
-// CHECK-NEXT: "use"(%[[KERNEL_ARG0]]) : (f32) -> ()
+// CHECK-NEXT: "use"(%[[ARG0]]) : (f32) -> ()
 // CHECK-NEXT: "some_op"(%[[BID]], %[[BDIM]]) : (index, index) -> ()
-// CHECK-NEXT: = load %[[KERNEL_ARG1]][%[[TID]]] : memref<?xf32, 1>
+// CHECK-NEXT: = load %[[ARG1]][%[[TID]]] : memref<?xf32, 1>
 
 // -----
 
@@ -79,8 +75,8 @@ func @multiple_launches() {
   return
 }
 
-// CHECK: func @multiple_launches_kernel()
-// CHECK: func @multiple_launches_kernel_0()
+// CHECK-LABEL: func @multiple_launches_kernel()
+// CHECK-LABEL: func @multiple_launches_kernel_0()
 
 // -----
 
@@ -104,23 +100,3 @@ func @extra_constants(%arg0 : memref<?xf32>) {
 // CHECK-LABEL: func @extra_constants_kernel(%{{.*}}: memref<?xf32>)
 // CHECK: constant
 // CHECK: constant
-
-// -----
-
-func @function_call(%arg0 : memref<?xf32>) {
-  %cst = constant 8 : index
-  gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %cst, %grid_y = %cst,
-                                       %grid_z = %cst)
-             threads(%tx, %ty, %tz) in (%block_x = %cst, %block_y = %cst,
-                                        %block_z = %cst) {
-    // TODO(b/141098412): Support function calls.
-    // expected-error @+1 {{'device_function' does not reference a valid function}}
-    call @device_function() : () -> ()
-    gpu.return
-  }
-  return
-}
-
-func @device_function() {
-  gpu.return
-}