diff --git a/mlir/lib/Conversion/GPUToVulkan/ConvertGPULaunchFuncToVulkanLaunchFunc.cpp b/mlir/lib/Conversion/GPUToVulkan/ConvertGPULaunchFuncToVulkanLaunchFunc.cpp --- a/mlir/lib/Conversion/GPUToVulkan/ConvertGPULaunchFuncToVulkanLaunchFunc.cpp +++ b/mlir/lib/Conversion/GPUToVulkan/ConvertGPULaunchFuncToVulkanLaunchFunc.cpp @@ -65,6 +65,11 @@ /// 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 @@ -93,14 +98,23 @@ 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 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 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 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"; } @@ -147,10 +161,18 @@ if (failed(declareVulkanLaunchFunc(loc, launchOp))) return signalPassFailure(); + SmallVector gpuLaunchOperands(launchOp.getOperands()); + SmallVector vulkanLaunchOperands( + gpuLaunchOperands.begin(), + gpuLaunchOperands.begin() + kVulkanLaunchNumConfigOperands); + vulkanLaunchOperands.append(gpuLaunchOperands.begin() + + gpu::LaunchOp::kNumConfigOperands, + gpuLaunchOperands.end()); + // Create vulkan launch call op. auto vulkanLaunchCallOp = builder.create( loc, ArrayRef{}, builder.getSymbolRefAttr(kVulkanLaunch), - launchOp.getOperands()); + vulkanLaunchOperands); // Set SPIR-V binary shader data as an attribute. vulkanLaunchCallOp.setAttr( diff --git a/mlir/lib/Conversion/GPUToVulkan/ConvertLaunchFuncToVulkanCalls.cpp b/mlir/lib/Conversion/GPUToVulkan/ConvertLaunchFuncToVulkanCalls.cpp --- a/mlir/lib/Conversion/GPUToVulkan/ConvertLaunchFuncToVulkanCalls.cpp +++ b/mlir/lib/Conversion/GPUToVulkan/ConvertLaunchFuncToVulkanCalls.cpp @@ -16,7 +16,6 @@ #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" @@ -121,7 +120,7 @@ /// 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 @@ -129,7 +128,7 @@ 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 @@ -162,6 +161,9 @@ // TODO: Use an associative array to support multiple vulkan launch calls. std::pair spirvAttributes; + /// The number of vulkan launch configuration operands, placed at the leading + /// positions of the operand list. + static constexpr unsigned kVulkanLaunchNumConfigOperands = 3; }; } // anonymous namespace @@ -209,7 +211,7 @@ void VulkanLaunchFuncToVulkanCallsPass::createBindMemRefCalls( LLVM::CallOp cInterfaceVulkanLaunchCallOp, Value vulkanRuntime) { if (cInterfaceVulkanLaunchCallOp.getNumOperands() == - gpu::LaunchOp::kNumConfigOperands) + kVulkanLaunchNumConfigOperands) return; OpBuilder builder(cInterfaceVulkanLaunchCallOp); Location loc = cInterfaceVulkanLaunchCallOp.getLoc(); @@ -222,7 +224,7 @@ 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( loc, getInt32Type(), builder.getI32IntegerAttr(en.index())); diff --git a/mlir/test/Conversion/GPUToVulkan/invoke-vulkan.mlir b/mlir/test/Conversion/GPUToVulkan/invoke-vulkan.mlir --- a/mlir/test/Conversion/GPUToVulkan/invoke-vulkan.mlir +++ b/mlir/test/Conversion/GPUToVulkan/invoke-vulkan.mlir @@ -40,11 +40,11 @@ %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] }"> @@ -54,8 +54,8 @@ %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] }*">) } diff --git a/mlir/test/Conversion/GPUToVulkan/lower-gpu-launch-vulkan-launch.mlir b/mlir/test/Conversion/GPUToVulkan/lower-gpu-launch-vulkan-launch.mlir --- a/mlir/test/Conversion/GPUToVulkan/lower-gpu-launch-vulkan-launch.mlir +++ b/mlir/test/Conversion/GPUToVulkan/lower-gpu-launch-vulkan-launch.mlir @@ -2,7 +2,7 @@ // 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 {