1 //===- ConvertGPULaunchFuncToVulkanLaunchFunc.cpp - MLIR conversion pass --===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // This file implements a pass to convert gpu launch function into a vulkan
10 // launch function. Creates a SPIR-V binary shader from the `spirv::ModuleOp`
11 // using `spirv::serialize` function, attaches binary data and entry point name
12 // as an attributes to vulkan launch call op.
13 //
14 //===----------------------------------------------------------------------===//
15
16 #include "../PassDetail.h"
17 #include "mlir/Conversion/GPUToVulkan/ConvertGPUToVulkanPass.h"
18 #include "mlir/Dialect/GPU/GPUDialect.h"
19 #include "mlir/Dialect/SPIRV/SPIRVDialect.h"
20 #include "mlir/Dialect/SPIRV/SPIRVOps.h"
21 #include "mlir/Dialect/SPIRV/Serialization.h"
22 #include "mlir/Dialect/StandardOps/IR/Ops.h"
23 #include "mlir/IR/Attributes.h"
24 #include "mlir/IR/Builders.h"
25 #include "mlir/IR/BuiltinOps.h"
26 #include "mlir/IR/BuiltinTypes.h"
27
28 using namespace mlir;
29
30 static constexpr const char *kSPIRVBlobAttrName = "spirv_blob";
31 static constexpr const char *kSPIRVEntryPointAttrName = "spirv_entry_point";
32 static constexpr const char *kVulkanLaunch = "vulkanLaunch";
33
34 namespace {
35
36 /// A pass to convert gpu launch op to vulkan launch call op, by creating a
37 /// SPIR-V binary shader from `spirv::ModuleOp` using `spirv::serialize`
38 /// function and attaching binary data and entry point name as an attributes to
39 /// created vulkan launch call op.
40 class ConvertGpuLaunchFuncToVulkanLaunchFunc
41 : public ConvertGpuLaunchFuncToVulkanLaunchFuncBase<
42 ConvertGpuLaunchFuncToVulkanLaunchFunc> {
43 public:
44 void runOnOperation() override;
45
46 private:
47 /// Creates a SPIR-V binary shader from the given `module` using
48 /// `spirv::serialize` function.
49 LogicalResult createBinaryShader(ModuleOp module,
50 std::vector<char> &binaryShader);
51
52 /// Converts the given `launchOp` to vulkan launch call.
53 void convertGpuLaunchFunc(gpu::LaunchFuncOp launchOp);
54
55 /// Checks where the given type is supported by Vulkan runtime.
isSupportedType(Type type)56 bool isSupportedType(Type type) {
57 if (auto memRefType = type.dyn_cast_or_null<MemRefType>()) {
58 auto elementType = memRefType.getElementType();
59 return memRefType.hasRank() &&
60 (memRefType.getRank() >= 1 && memRefType.getRank() <= 3) &&
61 (elementType.isIntOrFloat());
62 }
63 return false;
64 }
65
66 /// Declares the vulkan launch function. Returns an error if the any type of
67 /// operand is unsupported by Vulkan runtime.
68 LogicalResult declareVulkanLaunchFunc(Location loc,
69 gpu::LaunchFuncOp launchOp);
70
71 private:
72 /// The number of vulkan launch configuration operands, placed at the leading
73 /// positions of the operand list.
74 static constexpr unsigned kVulkanLaunchNumConfigOperands = 3;
75 };
76
77 } // anonymous namespace
78
runOnOperation()79 void ConvertGpuLaunchFuncToVulkanLaunchFunc::runOnOperation() {
80 bool done = false;
81 getOperation().walk([this, &done](gpu::LaunchFuncOp op) {
82 if (done) {
83 op.emitError("should only contain one 'gpu::LaunchFuncOp' op");
84 return signalPassFailure();
85 }
86 done = true;
87 convertGpuLaunchFunc(op);
88 });
89
90 // Erase `gpu::GPUModuleOp` and `spirv::Module` operations.
91 for (auto gpuModule :
92 llvm::make_early_inc_range(getOperation().getOps<gpu::GPUModuleOp>()))
93 gpuModule.erase();
94
95 for (auto spirvModule :
96 llvm::make_early_inc_range(getOperation().getOps<spirv::ModuleOp>()))
97 spirvModule.erase();
98 }
99
declareVulkanLaunchFunc(Location loc,gpu::LaunchFuncOp launchOp)100 LogicalResult ConvertGpuLaunchFuncToVulkanLaunchFunc::declareVulkanLaunchFunc(
101 Location loc, gpu::LaunchFuncOp launchOp) {
102 OpBuilder builder(getOperation().getBody()->getTerminator());
103
104 // Workgroup size is written into the kernel. So to properly modelling
105 // vulkan launch, we have to skip local workgroup size configuration here.
106 SmallVector<Type, 8> gpuLaunchTypes(launchOp.getOperandTypes());
107 // The first kVulkanLaunchNumConfigOperands of the gpu.launch_func op are the
108 // same as the config operands for the vulkan launch call op.
109 SmallVector<Type, 8> vulkanLaunchTypes(gpuLaunchTypes.begin(),
110 gpuLaunchTypes.begin() +
111 kVulkanLaunchNumConfigOperands);
112 vulkanLaunchTypes.append(gpuLaunchTypes.begin() +
113 gpu::LaunchOp::kNumConfigOperands,
114 gpuLaunchTypes.end());
115
116 // Check that all operands have supported types except those for the
117 // launch configuration.
118 for (auto type :
119 llvm::drop_begin(vulkanLaunchTypes, kVulkanLaunchNumConfigOperands)) {
120 if (!isSupportedType(type))
121 return launchOp.emitError() << type << " is unsupported to run on Vulkan";
122 }
123
124 // Declare vulkan launch function.
125 auto funcType = FunctionType::get(vulkanLaunchTypes, {}, loc->getContext());
126 builder.create<FuncOp>(loc, kVulkanLaunch, funcType).setPrivate();
127
128 return success();
129 }
130
createBinaryShader(ModuleOp module,std::vector<char> & binaryShader)131 LogicalResult ConvertGpuLaunchFuncToVulkanLaunchFunc::createBinaryShader(
132 ModuleOp module, std::vector<char> &binaryShader) {
133 bool done = false;
134 SmallVector<uint32_t, 0> binary;
135 for (auto spirvModule : module.getOps<spirv::ModuleOp>()) {
136 if (done)
137 return spirvModule.emitError("should only contain one 'spv.module' op");
138 done = true;
139
140 if (failed(spirv::serialize(spirvModule, binary)))
141 return failure();
142 }
143 binaryShader.resize(binary.size() * sizeof(uint32_t));
144 std::memcpy(binaryShader.data(), reinterpret_cast<char *>(binary.data()),
145 binaryShader.size());
146 return success();
147 }
148
convertGpuLaunchFunc(gpu::LaunchFuncOp launchOp)149 void ConvertGpuLaunchFuncToVulkanLaunchFunc::convertGpuLaunchFunc(
150 gpu::LaunchFuncOp launchOp) {
151 ModuleOp module = getOperation();
152 OpBuilder builder(launchOp);
153 Location loc = launchOp.getLoc();
154
155 // Serialize `spirv::Module` into binary form.
156 std::vector<char> binary;
157 if (failed(createBinaryShader(module, binary)))
158 return signalPassFailure();
159
160 // Declare vulkan launch function.
161 if (failed(declareVulkanLaunchFunc(loc, launchOp)))
162 return signalPassFailure();
163
164 SmallVector<Value, 8> gpuLaunchOperands(launchOp.getOperands());
165 SmallVector<Value, 8> vulkanLaunchOperands(
166 gpuLaunchOperands.begin(),
167 gpuLaunchOperands.begin() + kVulkanLaunchNumConfigOperands);
168 vulkanLaunchOperands.append(gpuLaunchOperands.begin() +
169 gpu::LaunchOp::kNumConfigOperands,
170 gpuLaunchOperands.end());
171
172 // Create vulkan launch call op.
173 auto vulkanLaunchCallOp = builder.create<CallOp>(
174 loc, TypeRange{}, builder.getSymbolRefAttr(kVulkanLaunch),
175 vulkanLaunchOperands);
176
177 // Set SPIR-V binary shader data as an attribute.
178 vulkanLaunchCallOp.setAttr(
179 kSPIRVBlobAttrName,
180 StringAttr::get({binary.data(), binary.size()}, loc->getContext()));
181
182 // Set entry point name as an attribute.
183 vulkanLaunchCallOp.setAttr(
184 kSPIRVEntryPointAttrName,
185 StringAttr::get(launchOp.getKernelName(), loc->getContext()));
186
187 launchOp.erase();
188 }
189
190 std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
createConvertGpuLaunchFuncToVulkanLaunchFuncPass()191 mlir::createConvertGpuLaunchFuncToVulkanLaunchFuncPass() {
192 return std::make_unique<ConvertGpuLaunchFuncToVulkanLaunchFunc>();
193 }
194