[mlir][vulkan-runner] Use C-compatible wrapper emission.

[denis13]

Summary

A memref argument is converted into a pointer-to-struct argument
of type {T*, T*, i64, i64[N], i64[N]}* in the wrapper function,
where T is the converted element type and N is the memref rank.

[denis13]

This patch eliminates VulkanLaunchOpOperandAdaptor and enables emitCWrappers while lowering to LLVM dialect. By doing this we follow the official guideline for calling convention https://mlir.llvm.org/docs/ConversionToLLVMDialect/#calling-convention.
emitCWrappers enables the following steps:

1. Declares a new function _mlir_ciface_vulkanLaunch where memref arguments are converted to pointer-to-struct and the remaining arguments are converted as usual.
2. Adds a body to the original function (making it non-external) that
   1. allocates a memref descriptor,
   2. populates it, and
   3. passes the pointer to it into the newly declared interface function
   4. collects the result of the call and returns it to the caller.

[rriddle]

mlir/lib/Conversion/GPUToVulkan/ConvertLaunchFuncToVulkanCalls.cpp

205

nit: You shouldn't need to construct a vector here, just use the operand_range directly. You can then change the llvm::drop_begin to operands.drop_front(...).

209

Why not use llvm::enumerate here? That would remove the need for the extra index variable.

[denis13]

Fix LINT issues.

[denis13]

mlir/lib/Conversion/GPUToVulkan/ConvertLaunchFuncToVulkanCalls.cpp

205

Thanks for point on this!

209

Thanks!

[denis13]

Address review comments.

[antiagainst]

Awesome! Thanks for the fix, Denis! Sorry for the delay. Overall LGTM; just a few nits. :)

mlir/lib/Conversion/GPUToVulkan/ConvertLaunchFuncToVulkanCalls.cpp

118

Nit: isCInterfaceVulkanLaunchCallOp? Just spell it out to make it clear. Applies to other parameter names, etc.

149

I think we can use gpu::LaunchOp::kNumConfigOperands instead of having our own here.

157

Nit: SPIR-V

[denis13]

Address review comments.

[denis13]

mlir/lib/Conversion/GPUToVulkan/ConvertLaunchFuncToVulkanCalls.cpp

118

Thanks, fixed!

149

Thanks for pointing on this! Fixed.

157

Thanks!

[mehdi_amini]

@denis13 FYI if you haven't already, you really can ask for commit access: https://llvm.org/docs/DeveloperPolicy.html#obtaining-commit-access

[denis13]

@mehdi_amini thanks, recently obtained it.

mlir/lib/Conversion/GPUToVulkan/ConvertGPULaunchFuncToVulkanLaunchFunc.cpp

 using namespace mlir;
 
 static constexpr const char *kSPIRVBlobAttrName = "spirv_blob";
 static constexpr const char *kSPIRVEntryPointAttrName = "spirv_entry_point";
 static constexpr const char *kVulkanLaunch = "vulkanLaunch";
 
 namespace {
 
-// A pass to convert gpu launch op to vulkan launch call op, by creating a
-// SPIR-V binary shader from `spirv::ModuleOp` using `spirv::serialize`
-// function and attaching binary data and entry point name as an attributes to
-// created vulkan launch call op.
+/// A pass to convert gpu launch op to vulkan launch call op, by creating a
+/// SPIR-V binary shader from `spirv::ModuleOp` using `spirv::serialize`
+/// function and attaching binary data and entry point name as an attributes to
+/// created vulkan launch call op.
 class ConvertGpuLaunchFuncToVulkanLaunchFunc
     : public ModulePass<ConvertGpuLaunchFuncToVulkanLaunchFunc> {
 public:
   void runOnModule() override;
 
 private:
   /// Creates a SPIR-V binary shader from the given `module` using
   /// `spirv::serialize` function.

mlir/lib/Conversion/GPUToVulkan/ConvertLaunchFuncToVulkanCalls.cpp

 #include "mlir/IR/Function.h"
 #include "mlir/IR/Module.h"
 #include "mlir/Pass/Pass.h"
 
 #include "llvm/ADT/SmallString.h"
 
 using namespace mlir;
 
-static constexpr const char *kBindResource = "bindResource";
+static constexpr const char *kBindMemRef1DFloat = "bindMemRef1DFloat";
+static constexpr const char *kCiFaceVulkanLaunch = "_mlir_ciface_vulkanLaunch";
 static constexpr const char *kDeinitVulkan = "deinitVulkan";
 static constexpr const char *kRunOnVulkan = "runOnVulkan";
 static constexpr const char *kInitVulkan = "initVulkan";
 static constexpr const char *kSetBinaryShader = "setBinaryShader";
 static constexpr const char *kSetEntryPoint = "setEntryPoint";
 static constexpr const char *kSetNumWorkGroups = "setNumWorkGroups";
 static constexpr const char *kSPIRVBinary = "SPIRV_BIN";
 static constexpr const char *kSPIRVBlobAttrName = "spirv_blob";
 static constexpr const char *kSPIRVEntryPointAttrName = "spirv_entry_point";
 static constexpr const char *kVulkanLaunch = "vulkanLaunch";
 
 namespace {
 
-/// A pass to convert vulkan launch func into a sequence of Vulkan
+/// A pass to convert vulkan launch call op into a sequence of Vulkan
 /// runtime calls in the following order:
 ///
 /// * initVulkan           -- initializes vulkan runtime
-/// * bindResource         -- binds resource
+/// * bindMemRef           -- binds memref
 /// * setBinaryShader      -- sets the binary shader data
 /// * setEntryPoint        -- sets the entry point name
 /// * setNumWorkGroups     -- sets the number of a local workgroups
 /// * runOnVulkan          -- runs vulkan runtime
 /// * deinitVulkan         -- deinitializes vulkan runtime
 ///
 class VulkanLaunchFuncToVulkanCallsPass
     : public ModulePass<VulkanLaunchFuncToVulkanCallsPass> {
 private:
   LLVM::LLVMDialect *getLLVMDialect() { return llvmDialect; }
 
   llvm::LLVMContext &getLLVMContext() {
     return getLLVMDialect()->getLLVMContext();
   }
 
   void initializeCachedTypes() {
     llvmDialect = getContext().getRegisteredDialect<LLVM::LLVMDialect>();
     llvmFloatType = LLVM::LLVMType::getFloatTy(llvmDialect);
     llvmVoidType = LLVM::LLVMType::getVoidTy(llvmDialect);
     llvmPointerType = LLVM::LLVMType::getInt8PtrTy(llvmDialect);
     llvmInt32Type = LLVM::LLVMType::getInt32Ty(llvmDialect);
     llvmInt64Type = LLVM::LLVMType::getInt64Ty(llvmDialect);
+    initializeMemRefTypes();
+  }
+
+  void initializeMemRefTypes() {
+    // According to the MLIR doc memref argument is converted into a
+    // pointer-to-struct argument of type:
+    // template <typename Elem, size_t Rank>
+    // struct {
+    //   Elem *allocated;
+    //   Elem *aligned;
+    //   int64_t offset;
+    //   int64_t sizes[Rank]; // omitted when rank == 0
+    //   int64_t strides[Rank]; // omitted when rank == 0
+    // };
+    auto llvmPtrToFloatType = getFloatType().getPointerTo();
+    auto llvmArrayOneElementSizeType =
+        LLVM::LLVMType::getArrayTy(getInt64Type(), 1);
+
+    // Create a type `!llvm<"{ float*, float*, i64, [1 x i64], [1 x i64]}">`.
+    llvmMemRef1DFloat = LLVM::LLVMType::getStructTy(
+        llvmDialect,
+        {llvmPtrToFloatType, llvmPtrToFloatType, getInt64Type(),
+         llvmArrayOneElementSizeType, llvmArrayOneElementSizeType});
   }
 
   LLVM::LLVMType getFloatType() { return llvmFloatType; }
   LLVM::LLVMType getVoidType() { return llvmVoidType; }
   LLVM::LLVMType getPointerType() { return llvmPointerType; }
   LLVM::LLVMType getInt32Type() { return llvmInt32Type; }
   LLVM::LLVMType getInt64Type() { return llvmInt64Type; }
+  LLVM::LLVMType getMemRef1DFloat() { return llvmMemRef1DFloat; }
 
   /// Creates a LLVM global for the given `name`.
   Value createEntryPointNameConstant(StringRef name, Location loc,
                                      OpBuilder &builder);
 
   /// Declares all needed runtime functions.
   void declareVulkanFunctions(Location loc);
 
   /// Checks whether the given LLVM::CallOp is a vulkan launch call op.
   bool isVulkanLaunchCallOp(LLVM::CallOp callOp) {
     return (callOp.callee() && callOp.callee().getValue() == kVulkanLaunch &&
-            callOp.getNumOperands() >= 6);
+            callOp.getNumOperands() >= kNumConfigOps);
+  }
+
+  /// Checks whether the given LLVM::CallOp is a "ci_face" vulkan launch call
+  /// op.
+  bool isCiFaceVulkanLaunchCallOp(LLVM::CallOp callOp) {
+    return (callOp.callee() &&
+            callOp.callee().getValue() == kCiFaceVulkanLaunch &&
+            callOp.getNumOperands() >= kNumConfigOps);
   }
 
   /// Translates the given `vulkanLaunchCallOp` to the sequence of Vulkan
   /// runtime calls.
   void translateVulkanLaunchCall(LLVM::CallOp vulkanLaunchCallOp);
 
-  /// Creates call to `bindResource` for each resource operand.
-  void createBindResourceCalls(LLVM::CallOp vulkanLaunchCallOp,
-                               Value vulkanRuntiem);
+  /// Creates call to `bindMemRef` for each memref operand.
+  void createBindMemRefCalls(LLVM::CallOp vulkanLaunchCallOp,
+                             Value vulkanRuntime);
+
+  /// Collects SPIRV attributes from the given `vulkanLaunchCallOp`.
+  void collectSPIRVAttributes(LLVM::CallOp vulkanLaunchCallOp);
 
 public:
   void runOnModule() override;
 
 private:
   LLVM::LLVMDialect *llvmDialect;
   LLVM::LLVMType llvmFloatType;
   LLVM::LLVMType llvmVoidType;
   LLVM::LLVMType llvmPointerType;
   LLVM::LLVMType llvmInt32Type;
   LLVM::LLVMType llvmInt64Type;
-};
-
-/// Represents operand adaptor for vulkan launch call operation, to simplify an
-/// access to the lowered memref.
-// TODO: We should use 'emit-c-wrappers' option to lower memref type:
-// https://mlir.llvm.org/docs/ConversionToLLVMDialect/#c-compatible-wrapper-emission.
-struct VulkanLaunchOpOperandAdaptor {
-  VulkanLaunchOpOperandAdaptor(ArrayRef<Value> values) { operands = values; }
-  VulkanLaunchOpOperandAdaptor(const VulkanLaunchOpOperandAdaptor &) = delete;
-  VulkanLaunchOpOperandAdaptor
-  operator=(const VulkanLaunchOpOperandAdaptor &) = delete;
-
-  /// Returns a tuple with a pointer to the memory and the size for the index-th
-  /// resource.
-  std::tuple<Value, Value> getResourceDescriptor1D(uint32_t index) {
-    assert(index < getResourceCount1D());
-    // 1D memref calling convention according to "ConversionToLLVMDialect.md":
-    // 0. Allocated pointer.
-    // 1. Aligned pointer.
-    // 2. Offset.
-    // 3. Size in dim 0.
-    // 4. Stride in dim 0.
-    auto offset = numConfigOps + index * loweredMemRefNumOps1D;
-    return std::make_tuple(operands[offset], operands[offset + 3]);
-  }
-
-  /// Returns the number of resources assuming all operands lowered from
-  /// 1D memref.
-  uint32_t getResourceCount1D() {
-    return (operands.size() - numConfigOps) / loweredMemRefNumOps1D;
-  }
+  LLVM::LLVMType llvmMemRef1DFloat;
 
-private:
-  /// The number of operands of lowered 1D memref.
-  static constexpr const uint32_t loweredMemRefNumOps1D = 5;
-  /// The number of the first config operands.
-  static constexpr const uint32_t numConfigOps = 6;
-  ArrayRef<Value> operands;
+  // TODO: Use an associative array to support multiple vulkan launch calls.
+  std::pair<StringAttr, StringAttr> spirvAttributes;
+  static constexpr const uint32_t kNumConfigOps = 6;
 };
 
 } // anonymous namespace
 
 void VulkanLaunchFuncToVulkanCallsPass::runOnModule() {
   initializeCachedTypes();
+
+  // Collect SPIRV attributes such as `spirv_blob` and `spirv_entry_point_name`.
   getModule().walk([this](LLVM::CallOp op) {
     if (isVulkanLaunchCallOp(op))
+      collectSPIRVAttributes(op);
+  });
+
+  // Convert vulkan launch call op into a sequence of Vulkan runtime calls.
+  getModule().walk([this](LLVM::CallOp op) {
+    if (isCiFaceVulkanLaunchCallOp(op))
       translateVulkanLaunchCall(op);
   });
 }
 
-void VulkanLaunchFuncToVulkanCallsPass::createBindResourceCalls(
-    LLVM::CallOp vulkanLaunchCallOp, Value vulkanRuntime) {
-  if (vulkanLaunchCallOp.getNumOperands() == 6)
+void VulkanLaunchFuncToVulkanCallsPass::collectSPIRVAttributes(
+    LLVM::CallOp vulkanLaunchCallOp) {
+  // Check that `kSPIRVBinary` and `kSPIRVEntryPoint` are present in attributes
+  // for the given vulkan launch call.
+  auto spirvBlobAttr =
+      vulkanLaunchCallOp.getAttrOfType<StringAttr>(kSPIRVBlobAttrName);
+  if (!spirvBlobAttr) {
+    vulkanLaunchCallOp.emitError()
+        << "missing " << kSPIRVBlobAttrName << " attribute";
+    return signalPassFailure();
+  }
+
+  auto spirvEntryPointNameAttr =
+      vulkanLaunchCallOp.getAttrOfType<StringAttr>(kSPIRVEntryPointAttrName);
+  if (!spirvEntryPointNameAttr) {
+    vulkanLaunchCallOp.emitError()
+        << "missing " << kSPIRVEntryPointAttrName << " attribute";
+    return signalPassFailure();
+  }
+
+  spirvAttributes = std::make_pair(spirvBlobAttr, spirvEntryPointNameAttr);
+}
+
+void VulkanLaunchFuncToVulkanCallsPass::createBindMemRefCalls(
+    LLVM::CallOp ciFaceVulkanLaunchCallOp, Value vulkanRuntime) {
+  if (ciFaceVulkanLaunchCallOp.getNumOperands() == kNumConfigOps)
     return;
-  OpBuilder builder(vulkanLaunchCallOp);
-  Location loc = vulkanLaunchCallOp.getLoc();
+  OpBuilder builder(ciFaceVulkanLaunchCallOp);
+  Location loc = ciFaceVulkanLaunchCallOp.getLoc();
 
   // Create LLVM constant for the descriptor set index.
-  // Bind all resources to the `0` descriptor set, the same way as `GPUToSPIRV`
+  // Bind all memrefs to the `0` descriptor set, the same way as `GPUToSPIRV`
   // pass does.
   Value descriptorSet = builder.create<LLVM::ConstantOp>(
       loc, getInt32Type(), builder.getI32IntegerAttr(0));
+  auto operands =
+      SmallVector<Value, 16>{ciFaceVulkanLaunchCallOp.getOperands()};
 
-  auto operands = SmallVector<Value, 32>{vulkanLaunchCallOp.getOperands()};
-  VulkanLaunchOpOperandAdaptor vkLaunchOperandAdaptor(operands);
-
-  for (auto resourceIdx :
-       llvm::seq<uint32_t>(0, vkLaunchOperandAdaptor.getResourceCount1D())) {
+  uint32_t operandIdx = 0;
+  for (const auto ptrToMemRefDescriptor :
+       llvm::drop_begin(operands, kNumConfigOps)) {
     // Create LLVM constant for the descriptor binding index.
     Value descriptorBinding = builder.create<LLVM::ConstantOp>(
-        loc, getInt32Type(), builder.getI32IntegerAttr(resourceIdx));
-    // Get a pointer to the memory and size of that memory.
-    auto resourceDescriptor =
-        vkLaunchOperandAdaptor.getResourceDescriptor1D(resourceIdx);
-    // Create call to `bindResource`.
+        loc, getInt32Type(), builder.getI32IntegerAttr(operandIdx));
+    // Create call to `bindMemRef`.
     builder.create<LLVM::CallOp>(
         loc, ArrayRef<Type>{getVoidType()},
-        builder.getSymbolRefAttr(kBindResource),
+        // TODO: Add support for memref with other ranks.
+        builder.getSymbolRefAttr(kBindMemRef1DFloat),
         ArrayRef<Value>{vulkanRuntime, descriptorSet, descriptorBinding,
-                        // Pointer to the memory.
-                        std::get<0>(resourceDescriptor),
-                        // Size of the memory.
-                        std::get<1>(resourceDescriptor)});
+                        ptrToMemRefDescriptor});
+    ++operandIdx;
   }
 }
 
 void VulkanLaunchFuncToVulkanCallsPass::declareVulkanFunctions(Location loc) {
   ModuleOp module = getModule();
   OpBuilder builder(module.getBody()->getTerminator());
 
   if (!module.lookupSymbol(kSetEntryPoint)) {
void VulkanLaunchFuncToVulkanCallsPass::declareVulkanFunctions(Location loc) {
 
   if (!module.lookupSymbol(kRunOnVulkan)) {
     builder.create<LLVM::LLVMFuncOp>(
         loc, kRunOnVulkan,
         LLVM::LLVMType::getFunctionTy(getVoidType(), {getPointerType()},
                                       /*isVarArg=*/false));
   }
 
-  if (!module.lookupSymbol(kBindResource)) {
+  if (!module.lookupSymbol(kBindMemRef1DFloat)) {
     builder.create<LLVM::LLVMFuncOp>(
-        loc, kBindResource,
-        LLVM::LLVMType::getFunctionTy(
-            getVoidType(),
-            {getPointerType(), getInt32Type(), getInt32Type(),
-             getFloatType().getPointerTo(), getInt64Type()},
+        loc, kBindMemRef1DFloat,
+        LLVM::LLVMType::getFunctionTy(getVoidType(),
+                                      {getPointerType(), getInt32Type(),
+                                       getInt32Type(),
+                                       getMemRef1DFloat().getPointerTo()},
                                       /*isVarArg=*/false));
   }
 
   if (!module.lookupSymbol(kInitVulkan)) {
     builder.create<LLVM::LLVMFuncOp>(
         loc, kInitVulkan,
         LLVM::LLVMType::getFunctionTy(getPointerType(), {},
                                       /*isVarArg=*/false));
   }
Value VulkanLaunchFuncToVulkanCallsPass::createEntryPointNameConstant(
 
   std::string entryPointGlobalName = (name + "_spv_entry_point_name").str();
   return LLVM::createGlobalString(loc, builder, entryPointGlobalName,
                                   shaderName, LLVM::Linkage::Internal,
                                   getLLVMDialect());
 }
 
 void VulkanLaunchFuncToVulkanCallsPass::translateVulkanLaunchCall(
-    LLVM::CallOp vulkanLaunchCallOp) {
-  OpBuilder builder(vulkanLaunchCallOp);
-  Location loc = vulkanLaunchCallOp.getLoc();
-
-  // Check that `kSPIRVBinary` and `kSPIRVEntryPoint` are present in attributes
-  // for the given vulkan launch call.
-  auto spirvBlobAttr =
-      vulkanLaunchCallOp.getAttrOfType<StringAttr>(kSPIRVBlobAttrName);
-  if (!spirvBlobAttr) {
-    vulkanLaunchCallOp.emitError()
-        << "missing " << kSPIRVBlobAttrName << " attribute";
-    return signalPassFailure();
-  }
-
-  auto entryPointNameAttr =
-      vulkanLaunchCallOp.getAttrOfType<StringAttr>(kSPIRVEntryPointAttrName);
-  if (!entryPointNameAttr) {
-    vulkanLaunchCallOp.emitError()
-        << "missing " << kSPIRVEntryPointAttrName << " attribute";
-    return signalPassFailure();
-  }
-
+    LLVM::CallOp ciFaceVulkanLaunchCallOp) {
+  OpBuilder builder(ciFaceVulkanLaunchCallOp);
+  Location loc = ciFaceVulkanLaunchCallOp.getLoc();
   // Create call to `initVulkan`.
   auto initVulkanCall = builder.create<LLVM::CallOp>(
       loc, ArrayRef<Type>{getPointerType()},
       builder.getSymbolRefAttr(kInitVulkan), ArrayRef<Value>{});
   // The result of `initVulkan` function is a pointer to Vulkan runtime, we
   // need to pass that pointer to each Vulkan runtime call.
   auto vulkanRuntime = initVulkanCall.getResult(0);
 
   // Create LLVM global with SPIR-V binary data, so we can pass a pointer with
   // that data to runtime call.
   Value ptrToSPIRVBinary = LLVM::createGlobalString(
-      loc, builder, kSPIRVBinary, spirvBlobAttr.getValue(),
+      loc, builder, kSPIRVBinary, spirvAttributes.first.getValue(),
       LLVM::Linkage::Internal, getLLVMDialect());
 
   // Create LLVM constant for the size of SPIR-V binary shader.
   Value binarySize = builder.create<LLVM::ConstantOp>(
       loc, getInt32Type(),
-      builder.getI32IntegerAttr(spirvBlobAttr.getValue().size()));
+      builder.getI32IntegerAttr(spirvAttributes.first.getValue().size()));
 
-  // Create call to `bindResource` for each resource operand.
-  createBindResourceCalls(vulkanLaunchCallOp, vulkanRuntime);
+  // Create call to `bindMemRef` for each memref operand.
+  createBindMemRefCalls(ciFaceVulkanLaunchCallOp, vulkanRuntime);
 
   // Create call to `setBinaryShader` runtime function with the given pointer to
   // SPIR-V binary and binary size.
   builder.create<LLVM::CallOp>(
       loc, ArrayRef<Type>{getVoidType()},
       builder.getSymbolRefAttr(kSetBinaryShader),
       ArrayRef<Value>{vulkanRuntime, ptrToSPIRVBinary, binarySize});
   // Create LLVM global with entry point name.
-  Value entryPointName =
-      createEntryPointNameConstant(entryPointNameAttr.getValue(), loc, builder);
+  Value entryPointName = createEntryPointNameConstant(
+      spirvAttributes.second.getValue(), loc, builder);
   // Create call to `setEntryPoint` runtime function with the given pointer to
   // entry point name.
   builder.create<LLVM::CallOp>(loc, ArrayRef<Type>{getVoidType()},
                                builder.getSymbolRefAttr(kSetEntryPoint),
                                ArrayRef<Value>{vulkanRuntime, entryPointName});
 
   // Create number of local workgroup for each dimension.
   builder.create<LLVM::CallOp>(
       loc, ArrayRef<Type>{getVoidType()},
       builder.getSymbolRefAttr(kSetNumWorkGroups),
-      ArrayRef<Value>{vulkanRuntime, vulkanLaunchCallOp.getOperand(0),
-                      vulkanLaunchCallOp.getOperand(1),
-                      vulkanLaunchCallOp.getOperand(2)});
+      ArrayRef<Value>{vulkanRuntime, ciFaceVulkanLaunchCallOp.getOperand(0),
+                      ciFaceVulkanLaunchCallOp.getOperand(1),
+                      ciFaceVulkanLaunchCallOp.getOperand(2)});
 
   // Create call to `runOnVulkan` runtime function.
   builder.create<LLVM::CallOp>(loc, ArrayRef<Type>{getVoidType()},
                                builder.getSymbolRefAttr(kRunOnVulkan),
                                ArrayRef<Value>{vulkanRuntime});
 
   // Create call to 'deinitVulkan' runtime function.
   builder.create<LLVM::CallOp>(loc, ArrayRef<Type>{getVoidType()},
                                builder.getSymbolRefAttr(kDeinitVulkan),
                                ArrayRef<Value>{vulkanRuntime});
 
   // Declare runtime functions.
   declareVulkanFunctions(loc);
 
-  vulkanLaunchCallOp.erase();
+  ciFaceVulkanLaunchCallOp.erase();
 }
 
 std::unique_ptr<mlir::OpPassBase<mlir::ModuleOp>>
 mlir::createConvertVulkanLaunchFuncToVulkanCallsPass() {
   return std::make_unique<VulkanLaunchFuncToVulkanCallsPass>();
 }
 
 static PassRegistration<VulkanLaunchFuncToVulkanCallsPass>
     pass("launch-func-to-vulkan",
          "Convert vulkanLaunch external call to Vulkan runtime external calls");

mlir/test/Conversion/GPUToVulkan/invoke-vulkan.mlir

 // RUN: mlir-opt %s -launch-func-to-vulkan | FileCheck %s
 
 // CHECK: llvm.mlir.global internal constant @kernel_spv_entry_point_name
 // CHECK: llvm.mlir.global internal constant @SPIRV_BIN
 // CHECK: %[[Vulkan_Runtime_ptr:.*]] = llvm.call @initVulkan() : () -> !llvm<"i8*">
 // CHECK: %[[addressof_SPIRV_BIN:.*]] = llvm.mlir.addressof @SPIRV_BIN
 // CHECK: %[[SPIRV_BIN_ptr:.*]] = llvm.getelementptr %[[addressof_SPIRV_BIN]]
 // CHECK: %[[SPIRV_BIN_size:.*]] = llvm.mlir.constant
-// CHECK: llvm.call @bindResource(%[[Vulkan_Runtime_ptr]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) : (!llvm<"i8*">, !llvm.i32, !llvm.i32, !llvm<"float*">, !llvm.i64) -> !llvm.void
+// CHECK: llvm.call @bindMemRef1DFloat(%[[Vulkan_Runtime_ptr]], %{{.*}}, %{{.*}}, %{{.*}}) : (!llvm<"i8*">, !llvm.i32, !llvm.i32, !llvm<"{ float*, float*, i64, [1 x i64], [1 x i64] }*">) -> !llvm.void
 // CHECK: llvm.call @setBinaryShader(%[[Vulkan_Runtime_ptr]], %[[SPIRV_BIN_ptr]], %[[SPIRV_BIN_size]]) : (!llvm<"i8*">, !llvm<"i8*">, !llvm.i32) -> !llvm.void
 // CHECK: %[[addressof_entry_point:.*]] = llvm.mlir.addressof @kernel_spv_entry_point_name
 // CHECK: %[[entry_point_ptr:.*]] = llvm.getelementptr %[[addressof_entry_point]]
 // CHECK: llvm.call @setEntryPoint(%[[Vulkan_Runtime_ptr]], %[[entry_point_ptr]]) : (!llvm<"i8*">, !llvm<"i8*">) -> !llvm.void
 // CHECK: llvm.call @setNumWorkGroups(%[[Vulkan_Runtime_ptr]], %{{.*}}, %{{.*}}, %{{.*}}) : (!llvm<"i8*">, !llvm.i64, !llvm.i64, !llvm.i64) -> !llvm.void
 // CHECK: llvm.call @runOnVulkan(%[[Vulkan_Runtime_ptr]]) : (!llvm<"i8*">) -> !llvm.void
 // CHECK: llvm.call @deinitVulkan(%[[Vulkan_Runtime_ptr]]) : (!llvm<"i8*">) -> !llvm.void
 
llvm.func @foo() {
     %18 = llvm.extractvalue %15[1] : !llvm<"{ float*, float*, i64, [1 x i64], [1 x i64] }">
     %19 = llvm.extractvalue %15[2] : !llvm<"{ float*, float*, i64, [1 x i64], [1 x i64] }">
     %20 = llvm.extractvalue %15[3, 0] : !llvm<"{ float*, float*, i64, [1 x i64], [1 x i64] }">
     %21 = llvm.extractvalue %15[4, 0] : !llvm<"{ float*, float*, i64, [1 x i64], [1 x i64] }">
     llvm.call @vulkanLaunch(%16, %16, %16, %16, %16, %16, %17, %18, %19, %20, %21) {spirv_blob = "\03\02#\07\00", spirv_entry_point = "kernel"}
     : (!llvm.i64, !llvm.i64, !llvm.i64, !llvm.i64, !llvm.i64, !llvm.i64, !llvm<"float*">, !llvm<"float*">, !llvm.i64, !llvm.i64, !llvm.i64) -> ()
     llvm.return
   }
-  llvm.func @vulkanLaunch(!llvm.i64, !llvm.i64, !llvm.i64, !llvm.i64, !llvm.i64, !llvm.i64, !llvm<"float*">, !llvm<"float*">, !llvm.i64, !llvm.i64, !llvm.i64)
+  llvm.func @vulkanLaunch(%arg0: !llvm.i64, %arg1: !llvm.i64, %arg2: !llvm.i64, %arg3: !llvm.i64, %arg4: !llvm.i64, %arg5: !llvm.i64, %arg6: !llvm<"float*">, %arg7: !llvm<"float*">, %arg8: !llvm.i64, %arg9: !llvm.i64, %arg10: !llvm.i64) {
+    %0 = llvm.mlir.undef : !llvm<"{ float*, float*, i64, [1 x i64], [1 x i64] }">
+    %1 = llvm.insertvalue %arg6, %0[0] : !llvm<"{ float*, float*, i64, [1 x i64], [1 x i64] }">
+    %2 = llvm.insertvalue %arg7, %1[1] : !llvm<"{ float*, float*, i64, [1 x i64], [1 x i64] }">
+    %3 = llvm.insertvalue %arg8, %2[2] : !llvm<"{ float*, float*, i64, [1 x i64], [1 x i64] }">
+    %4 = llvm.insertvalue %arg9, %3[3, 0] : !llvm<"{ float*, float*, i64, [1 x i64], [1 x i64] }">
+    %5 = llvm.insertvalue %arg10, %4[4, 0] : !llvm<"{ float*, float*, i64, [1 x i64], [1 x i64] }">
+    %6 = llvm.mlir.constant(1 : index) : !llvm.i64
+    %7 = llvm.alloca %6 x !llvm<"{ float*, float*, i64, [1 x i64], [1 x i64] }"> : (!llvm.i64) -> !llvm<"{ float*, float*, i64, [1 x i64], [1 x i64] }*">
+    llvm.store %5, %7 : !llvm<"{ float*, float*, i64, [1 x i64], [1 x i64] }*">
+    llvm.call @_mlir_ciface_vulkanLaunch(%arg0, %arg1, %arg2, %arg3, %arg4, %arg5, %7) : (!llvm.i64, !llvm.i64, !llvm.i64, !llvm.i64, !llvm.i64, !llvm.i64, !llvm<"{ float*, float*, i64, [1 x i64], [1 x i64] }*">) -> ()
+    llvm.return
+  }
+  llvm.func @_mlir_ciface_vulkanLaunch(!llvm.i64, !llvm.i64, !llvm.i64, !llvm.i64, !llvm.i64, !llvm.i64, !llvm<"{ float*, float*, i64, [1 x i64], [1 x i64] }*">)
 }

mlir/tools/mlir-vulkan-runner/VulkanRuntime.h

 #include "mlir/Dialect/SPIRV/SPIRVOps.h"
 #include "mlir/Dialect/SPIRV/Serialization.h"
 #include "mlir/IR/Module.h"
 #include "mlir/Support/LogicalResult.h"
 #include "mlir/Support/StringExtras.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/Support/ToolOutputFile.h"
 
-#include <vulkan/vulkan.h> // NOLINT
+#include <vulkan/vulkan.h>
 
 using namespace mlir;
 
 using DescriptorSetIndex = uint32_t;
 using BindingIndex = uint32_t;
 
 /// Struct containing information regarding to a device memory buffer.
 struct VulkanDeviceMemoryBuffer {

mlir/tools/mlir-vulkan-runner/mlir-vulkan-runner.cpp

static LogicalResult runMLIRPasses(ModuleOp module) {
   applyPassManagerCLOptions(passManager);
 
   passManager.addPass(createGpuKernelOutliningPass());
   passManager.addPass(createLegalizeStdOpsForSPIRVLoweringPass());
   passManager.addPass(createConvertGPUToSPIRVPass());
   OpPassManager &modulePM = passManager.nest<spirv::ModuleOp>();
   modulePM.addPass(spirv::createLowerABIAttributesPass());
   passManager.addPass(createConvertGpuLaunchFuncToVulkanLaunchFuncPass());
-  passManager.addPass(createLowerToLLVMPass());
+  passManager.addPass(createLowerToLLVMPass(/*useAlloca=*/false,
+                                            /*useBarePtrCallConv=*/false,
+                                            /*emitCWrappers=*/true));
   passManager.addPass(createConvertVulkanLaunchFuncToVulkanCallsPass());
   return passManager.run(module);
 }
 
 int main(int argc, char **argv) {
   llvm::llvm_shutdown_obj x;
   registerPassManagerCLOptions();
 
   mlir::registerAllDialects();
   llvm::InitLLVM y(argc, argv);
   llvm::InitializeNativeTarget();
   llvm::InitializeNativeTargetAsmPrinter();
   mlir::initializeLLVMPasses();
 
   return mlir::JitRunnerMain(argc, argv, &runMLIRPasses);
 }

mlir/tools/mlir-vulkan-runner/vulkan-runtime-wrappers.cpp

 
 private:
   VulkanRuntime vulkanRuntime;
   std::mutex mutex;
 };
 
 } // namespace
 
+template <typename T, int N>
+struct MemRefDescriptor {
+  T *allocated;
+  T *aligned;
+  int64_t offset;
+  int64_t sizes[N];
+  int64_t strides[N];
+};
+
 extern "C" {
-// Initializes `VulkanRuntimeManager` and returns a pointer to it.
+/// Initializes `VulkanRuntimeManager` and returns a pointer to it.
 void *initVulkan() { return new VulkanRuntimeManager(); }
 
-// Deinitializes `VulkanRuntimeManager` by the given pointer.
+/// Deinitializes `VulkanRuntimeManager` by the given pointer.
 void deinitVulkan(void *vkRuntimeManager) {
   delete reinterpret_cast<VulkanRuntimeManager *>(vkRuntimeManager);
 }
 
-/// Binds the given memref to the given descriptor set and descriptor index.
-void bindResource(void *vkRuntimeManager, DescriptorSetIndex setIndex,
-                  BindingIndex bindIndex, float *ptr, int64_t size) {
-  VulkanHostMemoryBuffer memBuffer{ptr,
-                                   static_cast<uint32_t>(size * sizeof(float))};
-  reinterpret_cast<VulkanRuntimeManager *>(vkRuntimeManager)
-      ->setResourceData(setIndex, bindIndex, memBuffer);
-}
-
 void runOnVulkan(void *vkRuntimeManager) {
   reinterpret_cast<VulkanRuntimeManager *>(vkRuntimeManager)->runOnVulkan();
 }
 
-/// Fills the given 1D float memref with the given float value.
-void fillResource1DFloat(float *allocated, float *aligned, int64_t offset,
-                         int64_t size, int64_t stride, float value) {
-  std::fill_n(allocated, size, value);
-}
-
 void setEntryPoint(void *vkRuntimeManager, const char *entryPoint) {
   reinterpret_cast<VulkanRuntimeManager *>(vkRuntimeManager)
       ->setEntryPoint(entryPoint);
 }
 
 void setNumWorkGroups(void *vkRuntimeManager, uint32_t x, uint32_t y,
                       uint32_t z) {
   reinterpret_cast<VulkanRuntimeManager *>(vkRuntimeManager)
       ->setNumWorkGroups({x, y, z});
 }
 
 void setBinaryShader(void *vkRuntimeManager, uint8_t *shader, uint32_t size) {
   reinterpret_cast<VulkanRuntimeManager *>(vkRuntimeManager)
       ->setShaderModule(shader, size);
 }
+
+/// Binds the given 1D float memref to the given descriptor set and descriptor
+/// index.
+void bindMemRef1DFloat(void *vkRuntimeManager, DescriptorSetIndex setIndex,
+                       BindingIndex bindIndex,
+                       MemRefDescriptor<float, 1> *ptr) {
+  VulkanHostMemoryBuffer memBuffer{
+      ptr->allocated, static_cast<uint32_t>(ptr->sizes[0] * sizeof(float))};
+  reinterpret_cast<VulkanRuntimeManager *>(vkRuntimeManager)
+      ->setResourceData(setIndex, bindIndex, memBuffer);
+}
+
+/// Fills the given 1D float memref with the given float value.
+void _mlir_ciface_fillResource1DFloat(MemRefDescriptor<float, 1> *ptr, // NOLINT
+                                      float value) {
+  std::fill_n(ptr->allocated, ptr->sizes[0], value);
+}
 }