Add a test for lowering GPU ops that cover cases where the symbol table isn't held by a ModuleOp (NFC)

PiperOrigin-RevId: 277752004

2 files changed, 45 insertions, 0 deletions

diff --git a/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir b/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir
index 205f1e664c6..6ad6becce39 100644
--- a/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir
+++ b/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir
@@ -100,3 +100,25 @@ module attributes {gpu.kernel_module} {
     std.return
   }
 }
+
+// -----
+
+// Test that we handled properly operation with SymbolTable other than module op
+module attributes {gpu.kernel_module} {
+  "test.symbol_scope"() ({
+  // CHECK: test.symbol_scope
+  // CHECK: llvm.func @__nv_expf(!llvm.float) -> !llvm.float
+  // CHECK: llvm.func @__nv_exp(!llvm.double) -> !llvm.double
+  // CHECK-LABEL: func @gpu_exp
+    func @gpu_exp(%arg_f32 : f32, %arg_f64 : f64) {
+      %exp_f32 = std.exp %arg_f32 : f32
+      // CHECK: llvm.call @__nv_expf(%{{.*}}) : (!llvm.float) -> !llvm.float
+      %result_f32 = std.exp %exp_f32 : f32
+      // CHECK: llvm.call @__nv_expf(%{{.*}}) : (!llvm.float) -> !llvm.float
+      %result64 = std.exp %arg_f64 : f64
+      // CHECK: llvm.call @__nv_exp(%{{.*}}) : (!llvm.double) -> !llvm.double
+      std.return
+    }
+    "test.finish" () : () -> ()
+  }) : () -> ()
+}
diff --git a/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl.mlir b/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl.mlir
index daa17ff734b..ccb931c7298 100644
--- a/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl.mlir
+++ b/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl.mlir
@@ -52,3 +52,26 @@ module attributes {gpu.kernel_module} {
     std.return
   }
 }
+
+
+// -----
+
+// Test that we handled properly operation with SymbolTable other than module op
+module attributes {gpu.kernel_module} {
+  "test.symbol_scope"() ({
+    // CHECK: test.symbol_scope
+    // CHECK: llvm.func @_ocml_exp_f32(!llvm.float) -> !llvm.float
+    // CHECK: llvm.func @_ocml_exp_f64(!llvm.double) -> !llvm.double
+    // CHECK-LABEL: func @gpu_exp
+    func @gpu_exp(%arg_f32 : f32, %arg_f64 : f64) {
+      %exp_f32 = std.exp %arg_f32 : f32
+      // CHECK: llvm.call @_ocml_exp_f32(%{{.*}}) : (!llvm.float) -> !llvm.float
+      %result_f32 = std.exp %exp_f32 : f32
+      // CHECK: llvm.call @_ocml_exp_f32(%{{.*}}) : (!llvm.float) -> !llvm.float
+      %result64 = std.exp %arg_f64 : f64
+      // CHECK: llvm.call @_ocml_exp_f64(%{{.*}}) : (!llvm.double) -> !llvm.double
+      std.return
+    }
+    "test.finish" () : () -> ()
+  }) : () -> ()
+}