[mlir][vulkan-runner] Add support for 2D memref.

[denis13]

Summary

This patch adds support for 2D memref in mlir-vulkan-runner.

[antiagainst]

Nice, thanks!

mlir/lib/Conversion/GPUToVulkan/ConvertLaunchFuncToVulkanCalls.cpp

102

This is similar enough with the above; but we only have two instance here so I'm fine with it atm. Might want to dedup by having a function later if there are more cases.

[denis13]

mlir/lib/Conversion/GPUToVulkan/ConvertLaunchFuncToVulkanCalls.cpp

102

oops, I see, thanks!

mlir/lib/Conversion/GPUToVulkan/ConvertGPULaunchFuncToVulkanLaunchFunc.cpp

private:
 
   /// Converts the given `luanchOp` to vulkan launch call.
   void convertGpuLaunchFunc(gpu::LaunchFuncOp launchOp);
 
   /// Checks where the given type is supported by Vulkan runtime.
   bool isSupportedType(Type type) {
     // TODO(denis0x0D): Handle other types.
     if (auto memRefType = type.dyn_cast_or_null<MemRefType>())
-      return memRefType.hasRank() && memRefType.getRank() == 1;
+      return memRefType.hasRank() &&
+             (memRefType.getRank() == 1 || memRefType.getRank() == 2);
     return false;
   }
 
   /// Declares the vulkan launch function. Returns an error if the any type of
   /// operand is unsupported by Vulkan runtime.
   LogicalResult declareVulkanLaunchFunc(Location loc,
                                         gpu::LaunchFuncOp launchOp);
 
LogicalResult ConvertGpuLaunchFuncToVulkanLaunchFunc::declareVulkanLaunchFunc(
     Location loc, gpu::LaunchFuncOp launchOp) {
   OpBuilder builder(getModule().getBody()->getTerminator());
   // TODO: Workgroup size is written into the kernel. So to properly modelling
   // vulkan launch, we cannot have the local workgroup size configuration here.
   SmallVector<Type, 8> vulkanLaunchTypes{launchOp.getOperandTypes()};
 
   // Check that all operands have supported types except those for the launch
   // configuration.
-  for (auto type : llvm::drop_begin(vulkanLaunchTypes, 6)) {
+  for (auto type :
+       llvm::drop_begin(vulkanLaunchTypes, gpu::LaunchOp::kNumConfigOperands)) {
     if (!isSupportedType(type))
       return launchOp.emitError() << type << " is unsupported to run on Vulkan";
   }
 
   // Declare vulkan launch function.
   builder.create<FuncOp>(
       loc, kVulkanLaunch,
       FunctionType::get(vulkanLaunchTypes, ArrayRef<Type>{}, loc->getContext()),

mlir/lib/Conversion/GPUToVulkan/ConvertLaunchFuncToVulkanCalls.cpp

 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Function.h"
 #include "mlir/IR/Module.h"
 #include "mlir/Pass/Pass.h"
 
 #include "llvm/ADT/SmallString.h"
+#include "llvm/Support/FormatVariadic.h"
 
 using namespace mlir;
 
 static constexpr const char *kBindMemRef1DFloat = "bindMemRef1DFloat";
+static constexpr const char *kBindMemRef2DFloat = "bindMemRef2DFloat";
 static constexpr const char *kCInterfaceVulkanLaunch =
     "_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";
void initializeMemRefTypes() {
     //   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);
+    auto llvmArrayTwoElementSizeType =
+        LLVM::LLVMType::getArrayTy(getInt64Type(), 2);
 
     // Create a type `!llvm<"{ float*, float*, i64, [1 x i64], [1 x i64]}">`.
     llvmMemRef1DFloat = LLVM::LLVMType::getStructTy(
         llvmDialect,
         {llvmPtrToFloatType, llvmPtrToFloatType, getInt64Type(),
          llvmArrayOneElementSizeType, llvmArrayOneElementSizeType});
+
+    // Create a type `!llvm<"{ float*, float*, i64, [2 x i64], [2 x i64]}">`.
+    llvmMemRef2DFloat = LLVM::LLVMType::getStructTy(
+        llvmDialect,
+        {llvmPtrToFloatType, llvmPtrToFloatType, getInt64Type(),
+         llvmArrayTwoElementSizeType, llvmArrayTwoElementSizeType});
   }
 
   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; }
+  LLVM::LLVMType getMemRef2DFloat() { return llvmMemRef2DFloat; }
 
   /// 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);
 
private:
 
   /// 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);
 
+  /// Deduces a rank from the given 'ptrToMemRefDescriptor`.
+  LogicalResult deduceMemRefRank(Value ptrToMemRefDescriptor, uint32_t &rank);
+
 public:
   void runOnModule() override;
 
 private:
   LLVM::LLVMDialect *llvmDialect;
   LLVM::LLVMType llvmFloatType;
   LLVM::LLVMType llvmVoidType;
   LLVM::LLVMType llvmPointerType;
   LLVM::LLVMType llvmInt32Type;
   LLVM::LLVMType llvmInt64Type;
   LLVM::LLVMType llvmMemRef1DFloat;
+  LLVM::LLVMType llvmMemRef2DFloat;
 
   // TODO: Use an associative array to support multiple vulkan launch calls.
   std::pair<StringAttr, StringAttr> spirvAttributes;
 };
 
 } // anonymous namespace
 
 void VulkanLaunchFuncToVulkanCallsPass::runOnModule() {
Value descriptorSet = builder.create<LLVM::ConstantOp>(
       loc, getInt32Type(), builder.getI32IntegerAttr(0));
 
   for (auto en :
        llvm::enumerate(cInterfaceVulkanLaunchCallOp.getOperands().drop_front(
            gpu::LaunchOp::kNumConfigOperands))) {
     // Create LLVM constant for the descriptor binding index.
     Value descriptorBinding = builder.create<LLVM::ConstantOp>(
         loc, getInt32Type(), builder.getI32IntegerAttr(en.index()));
+
+    auto ptrToMemRefDescriptor = en.value();
+    uint32_t rank = 0;
+    if (failed(deduceMemRefRank(ptrToMemRefDescriptor, rank))) {
+      cInterfaceVulkanLaunchCallOp.emitError()
+          << "invalid memref descriptor " << ptrToMemRefDescriptor.getType();
+      return signalPassFailure();
+    }
+
+    auto symbolName = llvm::formatv("bindMemRef{0}DFloat", rank).str();
     // Create call to `bindMemRef`.
     builder.create<LLVM::CallOp>(
         loc, ArrayRef<Type>{getVoidType()},
-        // TODO: Add support for memref with other ranks.
-        builder.getSymbolRefAttr(kBindMemRef1DFloat),
+        builder.getSymbolRefAttr(
+            StringRef(symbolName.data(), symbolName.size())),
         ArrayRef<Value>{vulkanRuntime, descriptorSet, descriptorBinding,
-                        en.value()});
+                        ptrToMemRefDescriptor});
+  }
+}
+
+LogicalResult
+VulkanLaunchFuncToVulkanCallsPass::deduceMemRefRank(Value ptrToMemRefDescriptor,
+                                                    uint32_t &rank) {
+  auto llvmPtrDescriptorTy =
+      ptrToMemRefDescriptor.getType().dyn_cast<LLVM::LLVMType>();
+  if (!llvmPtrDescriptorTy)
+    return failure();
+
+  auto llvmDescriptorTy = llvmPtrDescriptorTy.getPointerElementTy();
+  // 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
+  // };
+  if (!llvmDescriptorTy || !llvmDescriptorTy.isStructTy())
+    return failure();
+  if (llvmDescriptorTy.getStructNumElements() == 3) {
+    rank = 0;
+    return success();
   }
+
+  rank = llvmDescriptorTy.getStructElementType(3).getArrayNumElements();
+  return success();
 }
 
 void VulkanLaunchFuncToVulkanCallsPass::declareVulkanFunctions(Location loc) {
   ModuleOp module = getModule();
   OpBuilder builder(module.getBody()->getTerminator());
 
   if (!module.lookupSymbol(kSetEntryPoint)) {
     builder.create<LLVM::LLVMFuncOp>(
builder.create<LLVM::LLVMFuncOp>(
         loc, kBindMemRef1DFloat,
         LLVM::LLVMType::getFunctionTy(getVoidType(),
                                       {getPointerType(), getInt32Type(),
                                        getInt32Type(),
                                        getMemRef1DFloat().getPointerTo()},
                                       /*isVarArg=*/false));
   }
 
+  if (!module.lookupSymbol(kBindMemRef2DFloat)) {
+    builder.create<LLVM::LLVMFuncOp>(
+        loc, kBindMemRef2DFloat,
+        LLVM::LLVMType::getFunctionTy(getVoidType(),
+                                      {getPointerType(), getInt32Type(),
+                                       getInt32Type(),
+                                       getMemRef2DFloat().getPointerTo()},
+                                      /*isVarArg=*/false));
+  }
+
   if (!module.lookupSymbol(kInitVulkan)) {
     builder.create<LLVM::LLVMFuncOp>(
         loc, kInitVulkan,
         LLVM::LLVMType::getFunctionTy(getPointerType(), {},
                                       /*isVarArg=*/false));
   }
 
   if (!module.lookupSymbol(kDeinitVulkan)) {

mlir/test/mlir-vulkan-runner/mulf.mlir

+// RUN: mlir-vulkan-runner %s --shared-libs=%vulkan_wrapper_library_dir/libvulkan-runtime-wrappers%shlibext,%linalg_test_lib_dir/libmlir_runner_utils%shlibext --entry-point-result=void | FileCheck %s
+
+// CHECK-COUNT-4: [6, 6, 6, 6]
+module attributes {
+  gpu.container_module,
+  spv.target_env = #spv.target_env<
+    #spv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>,
+    {max_compute_workgroup_invocations = 128 : i32,
+     max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>
+} {
+  gpu.module @kernels {
+    gpu.func @kernel_mul(%arg0 : memref<4x4xf32>, %arg1 : memref<4x4xf32>, %arg2 : memref<4x4xf32>)
+      attributes {gpu.kernel, spv.entry_point_abi = {local_size = dense<[1, 1, 1]>: vector<3xi32>}} {
+      %x = "gpu.block_id"() {dimension = "x"} : () -> index
+      %y = "gpu.block_id"() {dimension = "y"} : () -> index
+      %1 = load %arg0[%x, %y] : memref<4x4xf32>
+      %2 = load %arg1[%x, %y] : memref<4x4xf32>
+      %3 = mulf %1, %2 : f32
+      store %3, %arg2[%x, %y] : memref<4x4xf32>
+      gpu.return
+    }
+  }
+
+  func @main() {
+    %arg0 = alloc() : memref<4x4xf32>
+    %arg1 = alloc() : memref<4x4xf32>
+    %arg2 = alloc() : memref<4x4xf32>
+    %0 = constant 0 : i32
+    %1 = constant 1 : i32
+    %2 = constant 2 : i32
+    %value0 = constant 0.0 : f32
+    %value1 = constant 2.0 : f32
+    %value2 = constant 3.0 : f32
+    %arg3 = memref_cast %arg0 : memref<4x4xf32> to memref<?x?xf32>
+    %arg4 = memref_cast %arg1 : memref<4x4xf32> to memref<?x?xf32>
+    %arg5 = memref_cast %arg2 : memref<4x4xf32> to memref<?x?xf32>
+    call @fillResource2DFloat(%arg3, %value1) : (memref<?x?xf32>, f32) -> ()
+    call @fillResource2DFloat(%arg4, %value2) : (memref<?x?xf32>, f32) -> ()
+    call @fillResource2DFloat(%arg5, %value0) : (memref<?x?xf32>, f32) -> ()
+
+    %cst1 = constant 1 : index
+    %cst4 = constant 4 : index
+    "gpu.launch_func"(%cst4, %cst4, %cst1, %cst1, %cst1, %cst1, %arg0, %arg1, %arg2) { kernel = "kernel_mul", kernel_module = @kernels }
+        : (index, index, index, index, index, index, memref<4x4xf32>, memref<4x4xf32>, memref<4x4xf32>) -> ()
+    %arg6 = memref_cast %arg5 : memref<?x?xf32> to memref<*xf32>
+    call @print_memref_f32(%arg6) : (memref<*xf32>) -> ()
+    return
+  }
+  func @fillResource2DFloat(%0 : memref<?x?xf32>, %1 : f32)
+  func @print_memref_f32(%ptr : memref<*xf32>)
+}
+

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

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);
 }
 
+/// Binds the given 2D float memref to the given descriptor set and descriptor
+/// index.
+void bindMemRef2DFloat(void *vkRuntimeManager, DescriptorSetIndex setIndex,
+                       BindingIndex bindIndex,
+                       MemRefDescriptor<float, 2> *ptr) {
+  VulkanHostMemoryBuffer memBuffer{
+      ptr->allocated,
+      static_cast<uint32_t>(ptr->sizes[0] * ptr->sizes[1] * 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);
 }
+
+/// Fills the given 2D float memref with the given float value.
+void _mlir_ciface_fillResource2DFloat(MemRefDescriptor<float, 2> *ptr, // NOLINT
+                                      float value) {
+  std::fill_n(ptr->allocated, ptr->sizes[0] * ptr->sizes[1], value);
+}
 }