[mlir][vulkan-runner] Simplify vulkan launch call op.

[denis13]

Summary

Workgroup size is written into the kernel. So to properly modelling
vulkan launch, we have to skip local workgroup size for vulkan launch
call op.

[mravishankar]

mlir/lib/Conversion/GPUToVulkan/ConvertGPULaunchFuncToVulkanLaunchFunc.cpp

107

Super Nit : I think this usage is different from the comment on what kVulkanLaunchNumConfigOperands says. The comment implies, off the argument list to the vulkan call how many of the arguments relate to configuration. Here it is used in the context that the first 3 arguments of the arguments to the gpu.launch_func are needed for the vulkan launch. Had to read this a couple of times to understand. Maybe just add a comment that says that the first kVulkanLaunchNumConfigOperands of the gpu.launch_func op are the same as the config operands for the vulkan call.

[denis13]

mlir/lib/Conversion/GPUToVulkan/ConvertGPULaunchFuncToVulkanLaunchFunc.cpp

107

Thanks for review! Sure, will fix it!

mlir/lib/Conversion/GPUToVulkan/ConvertGPULaunchFuncToVulkanLaunchFunc.cpp

if (auto memRefType = type.dyn_cast_or_null<MemRefType>())
              (memRefType.getRank() >= 1 && memRefType.getRank() <= 3);
     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);
+
+private:
+  /// The number of vulkan launch configuration operands, placed at the leading
+  /// positions of the operand list.
+  static constexpr unsigned kVulkanLaunchNumConfigOperands = 3;
 };
 
 } // anonymous namespace
 
 void ConvertGpuLaunchFuncToVulkanLaunchFunc::runOnOperation() {
   bool done = false;
   getOperation().walk([this, &done](gpu::LaunchFuncOp op) {
     if (done) {
void ConvertGpuLaunchFuncToVulkanLaunchFunc::runOnOperation() {
   for (auto spirvModule :
        llvm::make_early_inc_range(getOperation().getOps<spirv::ModuleOp>()))
     spirvModule.erase();
 }
 
 LogicalResult ConvertGpuLaunchFuncToVulkanLaunchFunc::declareVulkanLaunchFunc(
     Location loc, gpu::LaunchFuncOp launchOp) {
   OpBuilder builder(getOperation().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.
+  // Workgroup size is written into the kernel. So to properly modelling
+  // vulkan launch, we have to skip local workgroup size configuration here.
+  SmallVector<Type, 8> gpuLaunchTypes(launchOp.getOperandTypes());
+  // The first kVulkanLaunchNumConfigOperands of the gpu.launch_func op are the
+  // same as the config operands for the vulkan launch call op.
+  SmallVector<Type, 8> vulkanLaunchTypes(gpuLaunchTypes.begin(),
+                                         gpuLaunchTypes.begin() +
+                                             kVulkanLaunchNumConfigOperands);
+  vulkanLaunchTypes.append(gpuLaunchTypes.begin() +
+                               gpu::LaunchOp::kNumConfigOperands,
+                           gpuLaunchTypes.end());
+
+  // Check that all operands have supported types except those for the
+  // launch configuration.
   for (auto type :
-       llvm::drop_begin(vulkanLaunchTypes, gpu::LaunchOp::kNumConfigOperands)) {
+       llvm::drop_begin(vulkanLaunchTypes, kVulkanLaunchNumConfigOperands)) {
     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()),
void ConvertGpuLaunchFuncToVulkanLaunchFunc::convertGpuLaunchFunc(
   std::vector<char> binary;
   if (failed(createBinaryShader(module, binary)))
     return signalPassFailure();
 
   // Declare vulkan launch function.
   if (failed(declareVulkanLaunchFunc(loc, launchOp)))
     return signalPassFailure();
 
+  SmallVector<Value, 8> gpuLaunchOperands(launchOp.getOperands());
+  SmallVector<Value, 8> vulkanLaunchOperands(
+      gpuLaunchOperands.begin(),
+      gpuLaunchOperands.begin() + kVulkanLaunchNumConfigOperands);
+  vulkanLaunchOperands.append(gpuLaunchOperands.begin() +
+                                  gpu::LaunchOp::kNumConfigOperands,
+                              gpuLaunchOperands.end());
+
   // Create vulkan launch call op.
   auto vulkanLaunchCallOp = builder.create<CallOp>(
       loc, ArrayRef<Type>{}, builder.getSymbolRefAttr(kVulkanLaunch),
-      launchOp.getOperands());
+      vulkanLaunchOperands);
 
   // Set SPIR-V binary shader data as an attribute.
   vulkanLaunchCallOp.setAttr(
       kSPIRVBlobAttrName,
       StringAttr::get({binary.data(), binary.size()}, loc->getContext()));
 
   // Set entry point name as an attribute.
   vulkanLaunchCallOp.setAttr(

mlir/lib/Conversion/GPUToVulkan/ConvertLaunchFuncToVulkanCalls.cpp

 // 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.
 //
 //===----------------------------------------------------------------------===//
 
 #include "../PassDetail.h"
 #include "mlir/Conversion/GPUToVulkan/ConvertGPUToVulkanPass.h"
-#include "mlir/Dialect/GPU/GPUDialect.h"
 #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 "llvm/ADT/SmallString.h"
 #include "llvm/Support/FormatVariadic.h"
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() >= gpu::LaunchOp::kNumConfigOperands);
+            callOp.getNumOperands() >= kVulkanLaunchNumConfigOperands);
   }
 
   /// Checks whether the given LLVM::CallOp is a "ci_face" vulkan launch call
   /// op.
   bool isCInterfaceVulkanLaunchCallOp(LLVM::CallOp callOp) {
     return (callOp.callee() &&
             callOp.callee().getValue() == kCInterfaceVulkanLaunch &&
-            callOp.getNumOperands() >= gpu::LaunchOp::kNumConfigOperands);
+            callOp.getNumOperands() >= kVulkanLaunchNumConfigOperands);
   }
 
   /// Translates the given `vulkanLaunchCallOp` to the sequence of Vulkan
   /// runtime calls.
   void translateVulkanLaunchCall(LLVM::CallOp vulkanLaunchCallOp);
 
   /// Creates call to `bindMemRef` for each memref operand.
   void createBindMemRefCalls(LLVM::CallOp vulkanLaunchCallOp,
private:
   LLVM::LLVMType llvmInt32Type;
   LLVM::LLVMType llvmInt64Type;
   LLVM::LLVMType llvmMemRef1DFloat;
   LLVM::LLVMType llvmMemRef2DFloat;
   LLVM::LLVMType llvmMemRef3DFloat;
 
   // TODO: Use an associative array to support multiple vulkan launch calls.
   std::pair<StringAttr, StringAttr> spirvAttributes;
+  /// The number of vulkan launch configuration operands, placed at the leading
+  /// positions of the operand list.
+  static constexpr unsigned kVulkanLaunchNumConfigOperands = 3;
 };
 
 } // anonymous namespace
 
 void VulkanLaunchFuncToVulkanCallsPass::runOnOperation() {
   initializeCachedTypes();
 
   // Collect SPIR-V attributes such as `spirv_blob` and
void VulkanLaunchFuncToVulkanCallsPass::collectSPIRVAttributes(
   }
 
   spirvAttributes = std::make_pair(spirvBlobAttr, spirvEntryPointNameAttr);
 }
 
 void VulkanLaunchFuncToVulkanCallsPass::createBindMemRefCalls(
     LLVM::CallOp cInterfaceVulkanLaunchCallOp, Value vulkanRuntime) {
   if (cInterfaceVulkanLaunchCallOp.getNumOperands() ==
-      gpu::LaunchOp::kNumConfigOperands)
+      kVulkanLaunchNumConfigOperands)
     return;
   OpBuilder builder(cInterfaceVulkanLaunchCallOp);
   Location loc = cInterfaceVulkanLaunchCallOp.getLoc();
 
   // Create LLVM constant for the descriptor set index.
   // 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));
 
   for (auto en :
        llvm::enumerate(cInterfaceVulkanLaunchCallOp.getOperands().drop_front(
-           gpu::LaunchOp::kNumConfigOperands))) {
+           kVulkanLaunchNumConfigOperands))) {
     // 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()

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

llvm.func @foo() {
     %14 = llvm.insertvalue %0, %12[3, 0] : !llvm<"{ float*, float*, i64, [1 x i64], [1 x i64] }">
     %15 = llvm.insertvalue %13, %14[4, 0] : !llvm<"{ float*, float*, i64, [1 x i64], [1 x i64] }">
     %16 = llvm.mlir.constant(1 : index) : !llvm.i64
     %17 = llvm.extractvalue %15[0] : !llvm<"{ float*, float*, i64, [1 x i64], [1 x i64] }">
     %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.call @vulkanLaunch(%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<"float*">, !llvm<"float*">, !llvm.i64, !llvm.i64, !llvm.i64) -> ()
     llvm.return
   }
-  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) {
+  llvm.func @vulkanLaunch(%arg0: !llvm.i64, %arg1: !llvm.i64, %arg2: !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.call @_mlir_ciface_vulkanLaunch(%arg0, %arg1, %arg2, %7) : (!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] }*">)
+  llvm.func @_mlir_ciface_vulkanLaunch(!llvm.i64, !llvm.i64, !llvm.i64, !llvm<"{ float*, float*, i64, [1 x i64], [1 x i64] }*">)
 }

mlir/test/Conversion/GPUToVulkan/lower-gpu-launch-vulkan-launch.mlir

 // RUN: mlir-opt %s -convert-gpu-launch-to-vulkan-launch | FileCheck %s
 
 // CHECK: %[[resource:.*]] = alloc() : memref<12xf32>
 // CHECK: %[[index:.*]] = constant 1 : index
-// CHECK: call @vulkanLaunch(%[[index]], %[[index]], %[[index]], %[[index]], %[[index]], %[[index]], %[[resource]]) {spirv_blob = "{{.*}}", spirv_entry_point = "kernel"}
+// CHECK: call @vulkanLaunch(%[[index]], %[[index]], %[[index]], %[[resource]]) {spirv_blob = "{{.*}}", spirv_entry_point = "kernel"}
 
 module attributes {gpu.container_module} {
   spv.module Logical GLSL450 requires #spv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]> {
     spv.globalVariable @kernel_arg_0 bind(0, 0) : !spv.ptr<!spv.struct<!spv.array<12 x f32, stride=4> [0]>, StorageBuffer>
     spv.func @kernel() "None" attributes {workgroup_attributions = 0 : i64} {
       %0 = spv._address_of @kernel_arg_0 : !spv.ptr<!spv.struct<!spv.array<12 x f32, stride=4> [0]>, StorageBuffer>
       %2 = spv.constant 0 : i32
       %3 = spv._address_of @kernel_arg_0 : !spv.ptr<!spv.struct<!spv.array<12 x f32, stride=4> [0]>, StorageBuffer>