Use llvm.func to define functions with wrapped LLVM IR function type

This function-like operation allows one to define functions that have wrapped
LLVM IR function type, in particular variadic functions. The operation was
added in parallel to the existing lowering flow, this commit only switches the
flow to use it.

Using a custom function type makes the LLVM IR dialect type system more
consistent and avoids complex conversion rules for functions that previously
had to use the built-in function type instead of a wrapped LLVM IR dialect type
and perform conversions during the analysis.

PiperOrigin-RevId: 273910855

2 files changed, 8 insertions, 7 deletions

diff --git a/mlir/test/Conversion/GPUToCUDA/lower-launch-func-to-cuda.mlir b/mlir/test/Conversion/GPUToCUDA/lower-launch-func-to-cuda.mlir
index fcc8438bd93..7647d29fd45 100644
--- a/mlir/test/Conversion/GPUToCUDA/lower-launch-func-to-cuda.mlir
+++ b/mlir/test/Conversion/GPUToCUDA/lower-launch-func-to-cuda.mlir
@@ -10,7 +10,7 @@ module attributes {gpu.container_module} {
         attributes { gpu.kernel }
   }
 
-  func @foo() {
+  llvm.func @foo() {
     %0 = "op"() : () -> (!llvm.float)
     %1 = "op"() : () -> (!llvm<"float*">)
     %cst = constant 8 : index
@@ -29,7 +29,7 @@ module attributes {gpu.container_module} {
     "gpu.launch_func"(%cst, %cst, %cst, %cst, %cst, %cst, %0, %1) { kernel = "kernel", kernel_module = @kernel_module }
         : (index, index, index, index, index, index, !llvm.float, !llvm<"float*">) -> ()
 
-    return
+    llvm.return
   }
 
 }
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 90b33913058..62fe2b99338 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,9 @@
 // RUN: mlir-opt %s --test-kernel-to-cubin -split-input-file | FileCheck %s
 
 // CHECK: attributes {gpu.kernel_module, nvvm.cubin = "CUBIN"}
-module @kernels attributes {gpu.kernel_module} {
-  func @kernel(%arg0 : !llvm.float, %arg1 : !llvm<"float*">)
+module @foo attributes {gpu.kernel_module} {
+  llvm.func @kernel(%arg0 : !llvm.float, %arg1 : !llvm<"float*">)
+    // CHECK: attributes  {gpu.kernel}
     attributes  { gpu.kernel } {
     llvm.return
   }
@@ -10,15 +11,15 @@ module @kernels attributes {gpu.kernel_module} {
 
 // -----
 
-module attributes {gpu.kernel_module} {
+module @bar attributes {gpu.kernel_module} {
   // CHECK: func @kernel_a
-  func @kernel_a()
+  llvm.func @kernel_a()
     attributes  { gpu.kernel } {
     llvm.return
   }
 
   // CHECK: func @kernel_b
-  func @kernel_b()
+  llvm.func @kernel_b()
     attributes  { gpu.kernel } {
     llvm.return
   }