[mlir][spirv] Add pass ConvertGpuLaunchFuncToVulkanCallsPass.

[denis13]

Summary

Implement a pass to convert gpu.launch_func op into a sequence of
Vulkan runtime calls. The Vulkan runtime API surface is huge so currently we
don't expose separate external functions in IR for each of them, instead we
expose a few external functions to wrapper libraries which manages
Vulkan runtime.

[denis13]

Fix typo in comment.

[denis13]

Rebased on master.

[antiagainst]

Nice, thanks Denis! Reviews already happened in D72696 so I just have two more nits here.

mlir/lib/Conversion/GPUToVulkan/ConvertLaunchFuncToVulkanCalls.cpp

104

Move this into initializeCachedTypes or put all of them into the constructor of GpuLaunchFuncToVulkanCalssPass?

158

Could you document why we are putting a mysterious \0 in the middle here?

[denis13]

Addressed review comments.

mlir/lib/Conversion/GPUToVulkan/ConvertLaunchFuncToVulkanCalls.cpp

158

I was looking how llvm.mlir globals lowering to llvm ir, and as far as I undertood in the current approach the AddNull is set to false
https://llvm.org/doxygen/classllvm_1_1ConstantDataArray.html#a3edef3fa47c611d3d10606591213e57b
, so we need to explicitly append \0 to align with C style string. Please fix me I'm wrong.

if (op.getValueOrNull()) {
     // String attributes are treated separately because they cannot appear as
     // in-function constants and are thus not supported by getLLVMConstant.
     if (auto strAttr = op.getValueOrNull().dyn_cast_or_null<StringAttr>()) {
       cst = llvm::ConstantDataArray::getString(
           llvmModule->getContext(), strAttr.getValue(), /*AddNull=*/false);
       type = cst->getType();
     } else {
       cst = getLLVMConstant(type, op.getValueOrNull(), op.getLoc());
     }
   } else if (Block *initializer = op.getInitialize

[antiagainst]

mlir/lib/Conversion/GPUToVulkan/ConvertLaunchFuncToVulkanCalls.cpp

158

Yeah I was just saying that we should document along the code. :)

[rriddle]

mlir/lib/Conversion/GPUToVulkan/ConvertLaunchFuncToVulkanCalls.cpp

84

SmallVectorImpl

157

SmallString ? That would remove the explicit StringRef conversion below.

163

Seems easier to just write (name + "_spv_entry_point_name").str()

176

Return the error directly it auto converts to failure.

180

Drop all of the trivial braces.

198

Same here.

211

///

[denis13]

@rriddle thanks, I'll address your comments D72696

mlir/include/mlir/Conversion/GPUToVulkan/ConvertGPUToVulkanPass.h

+//===- ConvertGPUToVulkanPass.h - GPU to Vulkan conversion pass -*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// The file declares a pass to convert GPU dialect ops to to Vulkan runtime
+// calls.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_CONVERSION_GPUTOVULKAN_CONVERTGPUTOVULKANPASS_H
+#define MLIR_CONVERSION_GPUTOVULKAN_CONVERTGPUTOVULKANPASS_H
+
+#include "mlir/Support/LLVM.h"
+
+#include <memory>
+
+namespace mlir {
+
+class ModuleOp;
+template <typename T> class OpPassBase;
+
+std::unique_ptr<OpPassBase<ModuleOp>>
+createConvertGpuLaunchFuncToVulkanCallsPass();
+
+} // namespace mlir
+#endif // MLIR_CONVERSION_GPUTOVULKAN_CONVERTGPUTOVULKANPASS_H

mlir/lib/Conversion/CMakeLists.txt

 add_subdirectory(AffineToStandard)
 add_subdirectory(GPUToCUDA)
 add_subdirectory(GPUToNVVM)
 add_subdirectory(GPUToROCDL)
 add_subdirectory(GPUToSPIRV)
+add_subdirectory(GPUToVulkan)
 add_subdirectory(LinalgToLLVM)
 add_subdirectory(LinalgToSPIRV)
 add_subdirectory(LoopsToGPU)
 add_subdirectory(LoopToStandard)
 add_subdirectory(StandardToLLVM)
 add_subdirectory(StandardToSPIRV)
 add_subdirectory(VectorToLLVM)
 add_subdirectory(VectorToLoops)

mlir/lib/Conversion/GPUToVulkan/CMakeLists.txt

+add_llvm_library(MLIRGPUtoVulkanTransforms
+  ConvertLaunchFuncToVulkanCalls.cpp
+  )
+
+target_link_libraries(MLIRGPUtoVulkanTransforms
+  MLIRGPU
+  MLIRIR
+  MLIRLLVMIR
+  MLIRPass
+  MLIRSPIRV
+  MLIRSPIRVSerialization
+  MLIRStandardOps
+  MLIRSupport
+  MLIRTransforms
+  MLIRTranslation
+  )

mlir/lib/Conversion/GPUToVulkan/ConvertLaunchFuncToVulkanCalls.cpp

+//===- ConvertLaunchFuncToVulkanCalls.cpp - MLIR Vulkan conversion passes -===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements a pass to convert gpu.launch_func op into a sequence of
+// Vulkan runtime calls. The Vulkan runtime API surface is huge so currently we
+// don't expose separate external functions in IR for each of them, instead we
+// expose a few external functions to wrapper libraries which which manages
+// Vulkan runtime.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Conversion/GPUToVulkan/ConvertGPUToVulkanPass.h"
+#include "mlir/Dialect/GPU/GPUDialect.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/SPIRV/SPIRVOps.h"
+#include "mlir/Dialect/SPIRV/Serialization.h"
+#include "mlir/Dialect/StandardOps/Ops.h"
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/Function.h"
+#include "mlir/IR/Module.h"
+#include "mlir/IR/StandardTypes.h"
+#include "mlir/Pass/Pass.h"
+
+#include "llvm/Support/FormatVariadic.h"
+
+using namespace mlir;
+
+static constexpr const char *kSetBinaryShader = "setBinaryShader";
+static constexpr const char *kSetEntryPoint = "setEntryPoint";
+static constexpr const char *kSetNumWorkGroups = "setNumWorkGroups";
+static constexpr const char *kRunOnVulkan = "runOnVulkan";
+static constexpr const char *kSPIRVBinary = "SPIRV_BIN";
+
+namespace {
+
+/// A pass to convert gpu.launch_func operation into a sequence of Vulkan
+/// runtime calls.
+///
+/// * setBinaryShader      -- sets the binary shader data
+/// * setEntryPoint        -- sets the entry point name
+/// * setNumWorkGroups     -- sets the number of a local workgroups
+/// * runOnVulkan          -- runs vulkan runtime
+///
+class GpuLaunchFuncToVulkanCalssPass
+    : public ModulePass<GpuLaunchFuncToVulkanCalssPass> {
+private:
+  LLVM::LLVMDialect *getLLVMDialect() { return llvmDialect; }
+
+  llvm::LLVMContext &getLLVMContext() {
+    return getLLVMDialect()->getLLVMContext();
+  }
+
+  void initializeCachedTypes() {
+    llvmVoidType = LLVM::LLVMType::getVoidTy(llvmDialect);
+    llvmPointerType = LLVM::LLVMType::getInt8PtrTy(llvmDialect);
+    llvmInt32Type = LLVM::LLVMType::getInt32Ty(llvmDialect);
+  }
+
+  LLVM::LLVMType getVoidType() { return llvmVoidType; }
+  LLVM::LLVMType getPointerType() { return llvmPointerType; }
+  LLVM::LLVMType getInt32Type() { return llvmInt32Type; }
+
+  /// Creates a SPIR-V binary shader from the given `module` using
+  /// `spirv::serialize` function.
+  LogicalResult createBinaryShader(ModuleOp module,
+                                   std::vector<char> &binaryShader);
+
+  /// Creates a LLVM global for the given `name`.
+  Value createEntryPointNameConstant(StringRef name, Location loc,
+                                     OpBuilder &builder);
+
+  /// Creates a LLVM constant for each dimension of local workgroup and
+  /// populates the given `numWorkGroups`.
+  LogicalResult createNumWorkGroups(Location loc, OpBuilder &builder,
+                                    mlir::gpu::LaunchFuncOp launchOp,
+                                    SmallVector<Value, 3> &numWorkGroups);
+
+  /// Declares all needed runtime functions.
+  void declareVulkanFunctions(Location loc);
+
+  /// Translates the given `launcOp` op to the sequence of Vulkan runtime calls
+  void translateGpuLaunchCalls(mlir::gpu::LaunchFuncOp launchOp);
+
+public:
+  void runOnModule() override;
+
+private:
+  LLVM::LLVMDialect *llvmDialect;
+  LLVM::LLVMType llvmVoidType;
+  LLVM::LLVMType llvmPointerType;
+  LLVM::LLVMType llvmInt32Type;
+};
+
+} // anonymous namespace
+
+void GpuLaunchFuncToVulkanCalssPass::runOnModule() {
+  llvmDialect = getContext().getRegisteredDialect<LLVM::LLVMDialect>();
+  initializeCachedTypes();
+
+  getModule().walk(
+      [this](mlir::gpu::LaunchFuncOp op) { translateGpuLaunchCalls(op); });
+
+  // Erase `gpu::GPUModuleOp` and `spirv::Module` operations.
+  for (auto gpuModule :
+       llvm::make_early_inc_range(getModule().getOps<gpu::GPUModuleOp>()))
+    gpuModule.erase();
+
+  for (auto spirvModule :
+       llvm::make_early_inc_range(getModule().getOps<spirv::ModuleOp>()))
+    spirvModule.erase();
+}
+
+void GpuLaunchFuncToVulkanCalssPass::declareVulkanFunctions(Location loc) {
+  ModuleOp module = getModule();
+  OpBuilder builder(module.getBody()->getTerminator());
+
+  if (!module.lookupSymbol(kSetEntryPoint)) {
+    builder.create<LLVM::LLVMFuncOp>(
+        loc, kSetEntryPoint,
+        LLVM::LLVMType::getFunctionTy(getVoidType(), {getPointerType()},
+                                      /*isVarArg=*/false));
+  }
+
+  if (!module.lookupSymbol(kSetNumWorkGroups)) {
+    builder.create<LLVM::LLVMFuncOp>(
+        loc, kSetNumWorkGroups,
+        LLVM::LLVMType::getFunctionTy(
+            getVoidType(), {getInt32Type(), getInt32Type(), getInt32Type()},
+            /*isVarArg=*/false));
+  }
+
+  if (!module.lookupSymbol(kSetBinaryShader)) {
+    builder.create<LLVM::LLVMFuncOp>(
+        loc, kSetBinaryShader,
+        LLVM::LLVMType::getFunctionTy(getVoidType(),
+                                      {getPointerType(), getInt32Type()},
+                                      /*isVarArg=*/false));
+  }
+
+  if (!module.lookupSymbol(kRunOnVulkan)) {
+    builder.create<LLVM::LLVMFuncOp>(
+        loc, kRunOnVulkan,
+        LLVM::LLVMType::getFunctionTy(getVoidType(), {},
+                                      /*isVarArg=*/false));
+  }
+}
+
+Value GpuLaunchFuncToVulkanCalssPass::createEntryPointNameConstant(
+    StringRef name, Location loc, OpBuilder &builder) {
+  std::vector<char> shaderName(name.begin(), name.end());
+  shaderName.push_back('\0');
+
+  std::string entryPointGlobalName =
+      std::string(llvm::formatv("{0}_spv_entry_point_name", name));
+  return LLVM::createGlobalString(
+      loc, builder, entryPointGlobalName,
+      StringRef(shaderName.data(), shaderName.size()), LLVM::Linkage::Internal,
+      getLLVMDialect());
+}
+
+LogicalResult GpuLaunchFuncToVulkanCalssPass::createBinaryShader(
+    ModuleOp module, std::vector<char> &binaryShader) {
+  bool done = false;
+  SmallVector<uint32_t, 0> binary;
+  for (auto spirvModule : module.getOps<spirv::ModuleOp>()) {
+    if (done) {
+      spirvModule.emitError("should only contain one 'spv.module' op");
+      return failure();
+    }
+    done = true;
+    if (failed(spirv::serialize(spirvModule, binary))) {
+      return failure();
+    }
+  }
+
+  binaryShader.resize(binary.size() * sizeof(uint32_t));
+  std::memcpy(binaryShader.data(), reinterpret_cast<char *>(binary.data()),
+              binaryShader.size());
+  return success();
+}
+
+LogicalResult GpuLaunchFuncToVulkanCalssPass::createNumWorkGroups(
+    Location loc, OpBuilder &builder, mlir::gpu::LaunchFuncOp launchOp,
+    SmallVector<Value, 3> &numWorkGroups) {
+  for (auto index : llvm::seq(0, 3)) {
+    auto numWorkGroupDimConstant = dyn_cast_or_null<ConstantOp>(
+        launchOp.getOperand(index).getDefiningOp());
+
+    if (!numWorkGroupDimConstant) {
+      return failure();
+    }
+
+    auto numWorkGroupDimValue =
+        numWorkGroupDimConstant.getValue().cast<IntegerAttr>().getInt();
+    numWorkGroups.push_back(builder.create<LLVM::ConstantOp>(
+        loc, getInt32Type(), builder.getI32IntegerAttr(numWorkGroupDimValue)));
+  }
+
+  return success();
+}
+
+// Translates gpu launch op to the sequence of Vulkan runtime calls.
+void GpuLaunchFuncToVulkanCalssPass::translateGpuLaunchCalls(
+    mlir::gpu::LaunchFuncOp launchOp) {
+  ModuleOp module = getModule();
+  OpBuilder builder(launchOp);
+  Location loc = launchOp.getLoc();
+
+  // Serialize `spirv::Module` into binary form.
+  std::vector<char> binary;
+  if (failed(
+          GpuLaunchFuncToVulkanCalssPass::createBinaryShader(module, binary))) {
+    return signalPassFailure();
+  }
+
+  // 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, StringRef(binary.data(), binary.size()),
+      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(binary.size()));
+  // 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>{ptrToSPIRVBinary, binarySize});
+
+  // Create LLVM global with entry point name.
+  Value entryPointName =
+      createEntryPointNameConstant(launchOp.kernel(), 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>{entryPointName});
+
+  // Create number of local workgroup for each dimension.
+  SmallVector<Value, 3> numWorkGroups;
+  if (failed(createNumWorkGroups(loc, builder, launchOp, numWorkGroups))) {
+    return signalPassFailure();
+  }
+
+  // Create call `setNumWorkGroups` runtime function with the given numbers of
+  // local workgroup.
+  builder.create<LLVM::CallOp>(
+      loc, ArrayRef<Type>{getVoidType()},
+      builder.getSymbolRefAttr(kSetNumWorkGroups),
+      ArrayRef<Value>{numWorkGroups[0], numWorkGroups[1], numWorkGroups[2]});
+
+  // Create call to `runOnVulkan` runtime function.
+  builder.create<LLVM::CallOp>(loc, ArrayRef<Type>{getVoidType()},
+                               builder.getSymbolRefAttr(kRunOnVulkan),
+                               ArrayRef<Value>{});
+
+  // Declare runtime functions.
+  declareVulkanFunctions(loc);
+
+  launchOp.erase();
+}
+
+std::unique_ptr<mlir::OpPassBase<mlir::ModuleOp>>
+mlir::createConvertGpuLaunchFuncToVulkanCallsPass() {
+  return std::make_unique<GpuLaunchFuncToVulkanCalssPass>();
+}
+
+static PassRegistration<GpuLaunchFuncToVulkanCalssPass>
+    pass("launch-func-to-vulkan",
+         "Convert gpu.launch_func op to Vulkan runtime 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: %[[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 @setBinaryShader(%[[SPIRV_BIN_ptr]], %[[SPIRV_BIN_size]]) : (!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(%[[entry_point_ptr]]) : (!llvm<"i8*">) -> !llvm.void
+// CHECK: %[[Workgroup_X:.*]] = llvm.mlir.constant
+// CHECK: %[[Workgroup_Y:.*]] = llvm.mlir.constant
+// CHECK: %[[Workgroup_Z:.*]] = llvm.mlir.constant
+// CHECK: llvm.call @setNumWorkGroups(%[[Workgroup_X]], %[[Workgroup_Y]], %[[Workgroup_Z]]) : (!llvm.i32, !llvm.i32, !llvm.i32) -> !llvm.void
+// CHECK: llvm.call @runOnVulkan() : () -> !llvm.void
+
+module attributes {gpu.container_module} {
+  spv.module "Logical" "GLSL450" {
+    spv.globalVariable @kernel_arg_0 bind(0, 0) : !spv.ptr<!spv.struct<f32 [0]>, StorageBuffer>
+    spv.globalVariable @kernel_arg_1 bind(0, 1) : !spv.ptr<!spv.struct<!spv.array<12 x f32 [4]> [0]>, StorageBuffer>
+    func @kernel() attributes {workgroup_attributions = 0 : i64} {
+      %0 = spv._address_of @kernel_arg_1 : !spv.ptr<!spv.struct<!spv.array<12 x f32 [4]> [0]>, StorageBuffer>
+      %1 = spv._address_of @kernel_arg_0 : !spv.ptr<!spv.struct<f32 [0]>, StorageBuffer>
+      %2 = spv.constant 0 : i32
+      %3 = spv.AccessChain %1[%2] : !spv.ptr<!spv.struct<f32 [0]>, StorageBuffer>
+      %4 = spv.Load "StorageBuffer" %3 : f32
+      spv.Return
+    }
+    spv.EntryPoint "GLCompute" @kernel
+    spv.ExecutionMode @kernel "LocalSize", 1, 1, 1
+  } attributes {capabilities = ["Shader"], extensions = ["SPV_KHR_storage_buffer_storage_class"]}
+  gpu.module @kernels {
+    gpu.func @kernel(%arg0: f32, %arg1: memref<12xf32>) kernel {
+      gpu.return
+    }
+  }
+  func @foo() {
+    %0 = "op"() : () -> f32
+    %1 = "op"() : () -> memref<12xf32>
+    %c1 = constant 1 : index
+    "gpu.launch_func"(%c1, %c1, %c1, %c1, %c1, %c1, %0, %1) {kernel = "kernel", kernel_module = @kernels} : (index, index, index, index, index, index, f32, memref<12xf32>) -> ()
+    return
+  }
+}

mlir/tools/mlir-opt/CMakeLists.txt

set(FULL_LINK_LIBS
 
   MLIRLoopToStandard
   MLIREDSC
   MLIRFxpMathOps
   MLIRGPU
   MLIRGPUtoNVVMTransforms
   MLIRGPUtoROCDLTransforms
   MLIRGPUtoSPIRVTransforms
+  MLIRGPUtoVulkanTransforms
   MLIRLinalgOps
   MLIRLinalgAnalysis
   MLIRLinalgEDSC
   MLIRLinalgTransforms
   MLIRLinalgUtils
   MLIRLLVMIR
   MLIRLoopOps
   MLIRNVVMIR
   MLIROpenMP
   MLIROptMain
   MLIRParser
   MLIRPass
   MLIRQuantizerFxpMathConfig
   MLIRQuantizerSupport
   MLIRQuantizerTransforms
   MLIRQuantOps
   MLIRROCDLIR
   MLIRSPIRV
   MLIRLinalgToSPIRVTransforms
   MLIRStandardToSPIRVTransforms
   MLIRSPIRVTestPasses
   MLIRSPIRVTransforms
   MLIRStandardOps
   MLIRStandardToLLVM
+  MLIRSPIRVSerialization
   MLIRTransforms
   MLIRTransformUtils
   MLIRTestDialect
   MLIRTestIR
   MLIRTestPass
   MLIRTestTransforms
+  # MLIRTranslation is needed for MLIRSPIRVSerialization
+  MLIRTranslation
   MLIRSupport
   MLIRVectorOps
   MLIRVectorToLLVM
   MLIRVectorToLoops
 )
 set(LIBS
   MLIRIR
   MLIROptMain