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