[Mlir-commits] [mlir] [mlir][llvm] Add c interface ext wrapper for host bare ptr as well. (PR #72629)

[Jungwook Park]

https://github.com/jungpark-mlir created https://github.com/llvm/llvm-project/pull/72629

Quick addition to the funcs need to emit wrapper for the c interface. Doesn't convert anything but providing wrappers with "_mlir_ciface_" prefix.

>From 85cf652defb1bd5a16046df7f1585075e30dec14 Mon Sep 17 00:00:00 2001
From: jungpark-mlir <jungwook at jungwook-22.04>
Date: Fri, 17 Nov 2023 09:42:49 +0000
Subject: [PATCH] Add c interface ext wrapper for host bare ptr as well.

Quick addition to the funcs need to emit wrapper for the c interface.
Doesn't convert anything but providing wrappers with "_mlir_ciface_" prefix.
---
 mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp | 41 +++++++++++++++++++
 .../GPUCommon/lower-launch-func-bare-ptr.mlir |  6 ++-
 2 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp b/mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp
index bd50c67fb87958a..b41e707ef130515 100644
--- a/mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp
+++ b/mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp
@@ -269,6 +269,37 @@ static void wrapExternalFunction(OpBuilder &builder, Location loc,
   }
 }
 
+// Creating external wrappers with UseBarePtrCallConv=true.
+static void wrapForExternalCallers(OpBuilder &rewriter, Location loc,
+                                   LLVM::LLVMFuncOp newFuncOp) {
+  // Create the auxiliary function.
+  auto wrapperFunc = rewriter.cloneWithoutRegions(newFuncOp);
+  wrapperFunc.setSymName(
+      llvm::formatv("_mlir_ciface_{0}", newFuncOp.getName()).str());
+
+  OpBuilder::InsertionGuard guard(rewriter);
+  rewriter.setInsertionPointToStart(wrapperFunc.addEntryBlock());
+  auto call =
+      rewriter.create<LLVM::CallOp>(loc, newFuncOp, wrapperFunc.getArguments());
+  rewriter.create<LLVM::ReturnOp>(loc, call.getResults());
+}
+
+static void wrapExternalFunction(OpBuilder &builder, Location loc,
+                                 LLVM::LLVMFuncOp newFuncOp) {
+  OpBuilder::InsertionGuard guard(builder);
+  // Create the auxiliary function.
+  auto wrapperFunc = builder.cloneWithoutRegions(newFuncOp);
+  wrapperFunc.setSymName(
+      llvm::formatv("_mlir_ciface_{0}", newFuncOp.getName()).str());
+
+  // This wrapper should only be visible in this module.
+  newFuncOp.setLinkage(LLVM::Linkage::Private);
+  builder.setInsertionPointToStart(newFuncOp.addEntryBlock());
+  auto call =
+      builder.create<LLVM::CallOp>(loc, wrapperFunc, newFuncOp.getArguments());
+  builder.create<LLVM::ReturnOp>(loc, call.getResults());
+}
+
 /// Modifies the body of the function to construct the `MemRefDescriptor` from
 /// the bare pointer calling convention lowering of `memref` types.
 static void modifyFuncOpToUseBarePtrCallingConv(
@@ -502,6 +533,16 @@ struct FuncOpConversion : public FuncOpConversionBase {
       modifyFuncOpToUseBarePtrCallingConv(rewriter, funcOp->getLoc(),
                                           *getTypeConverter(), *newFuncOp,
                                           funcOp.getFunctionType().getInputs());
+      if (funcOp->getAttrOfType<UnitAttr>(
+              LLVM::LLVMDialect::getEmitCWrapperAttrName())) {
+        if (newFuncOp->isVarArg())
+          return funcOp->emitError("C interface for variadic functions is not "
+                                   "supported yet.");
+        if (newFuncOp->isExternal())
+          wrapExternalFunction(rewriter, newFuncOp->getLoc(), *newFuncOp);
+        else
+          wrapForExternalCallers(rewriter, funcOp->getLoc(), *newFuncOp);
+      }
     }
 
     rewriter.eraseOp(funcOp);
diff --git a/mlir/test/Conversion/GPUCommon/lower-launch-func-bare-ptr.mlir b/mlir/test/Conversion/GPUCommon/lower-launch-func-bare-ptr.mlir
index 5e1c3b797235f24..8a0eb38a450764b 100644
--- a/mlir/test/Conversion/GPUCommon/lower-launch-func-bare-ptr.mlir
+++ b/mlir/test/Conversion/GPUCommon/lower-launch-func-bare-ptr.mlir
@@ -1,4 +1,4 @@
-// RUN: mlir-opt %s --gpu-to-llvm="use-bare-pointers-for-kernels=1" -split-input-file | FileCheck %s
+// RUN: mlir-opt %s --gpu-to-llvm="use-bare-pointers-for-kernels=1 use-bare-pointers-for-host=1" -split-input-file | FileCheck %s
 
 module attributes {gpu.container_module} {
   gpu.module @kernels [#nvvm.target]  {
@@ -15,7 +15,8 @@ module attributes {gpu.container_module} {
       llvm.return
     }
   }
-  func.func @foo() {
+  // CHECK: @foo
+  func.func @foo() attributes {llvm.emit_c_interface} {
     // CHECK: [[MEMREF:%.*]] = gpu.alloc () : memref<10xf32, 1>
     // CHECK: [[DESCRIPTOR:%.*]] = builtin.unrealized_conversion_cast [[MEMREF]] : memref<10xf32, 1> to !llvm.struct<(ptr<1>, ptr<1>, i64, array<1 x i64>, array<1 x i64>)>
     // CHECK: [[PTR:%.*]] = llvm.extractvalue [[DESCRIPTOR]][1] : !llvm.struct<(ptr<1>, ptr<1>, i64, array<1 x i64>, array<1 x i64>)>
@@ -28,3 +29,4 @@ module attributes {gpu.container_module} {
     return
   }
 }
+// CHECK: @_mlir_ciface_foo