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/Func/IR/FuncOps.h"
19 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
20 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
21 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
22 #include "mlir/IR/Attributes.h"
23 #include "mlir/IR/Builders.h"
24 #include "mlir/IR/BuiltinOps.h"
25 #include "mlir/IR/BuiltinTypes.h"
26 #include "mlir/Target/SPIRV/Serialization.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.
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 } // namespace
78 
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 
100 LogicalResult ConvertGpuLaunchFuncToVulkanLaunchFunc::declareVulkanLaunchFunc(
101     Location loc, gpu::LaunchFuncOp launchOp) {
102   auto builder = OpBuilder::atBlockEnd(getOperation().getBody());
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 = builder.getFunctionType(vulkanLaunchTypes, {});
126   builder.create<func::FuncOp>(loc, kVulkanLaunch, funcType).setPrivate();
127 
128   return success();
129 }
130 
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 
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<func::CallOp>(
174       loc, TypeRange{}, SymbolRefAttr::get(builder.getContext(), kVulkanLaunch),
175       vulkanLaunchOperands);
176 
177   // Set SPIR-V binary shader data as an attribute.
178   vulkanLaunchCallOp->setAttr(
179       kSPIRVBlobAttrName,
180       builder.getStringAttr(StringRef(binary.data(), binary.size())));
181 
182   // Set entry point name as an attribute.
183   vulkanLaunchCallOp->setAttr(kSPIRVEntryPointAttrName,
184                               launchOp.getKernelName());
185 
186   launchOp.erase();
187 }
188 
189 std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
190 mlir::createConvertGpuLaunchFuncToVulkanLaunchFuncPass() {
191   return std::make_unique<ConvertGpuLaunchFuncToVulkanLaunchFunc>();
192 }
193