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 "mlir/Conversion/GPUToVulkan/ConvertGPUToVulkanPass.h"
17 #include "mlir/Dialect/GPU/GPUDialect.h"
18 #include "mlir/Dialect/SPIRV/SPIRVOps.h"
19 #include "mlir/Dialect/SPIRV/Serialization.h"
20 #include "mlir/Dialect/StandardOps/IR/Ops.h"
21 #include "mlir/IR/Attributes.h"
22 #include "mlir/IR/Builders.h"
23 #include "mlir/IR/Function.h"
24 #include "mlir/IR/Module.h"
25 #include "mlir/IR/StandardTypes.h"
26 #include "mlir/Pass/Pass.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 ModulePass<ConvertGpuLaunchFuncToVulkanLaunchFunc> {
42 public:
43 /// Include the generated pass utilities.
44 #define GEN_PASS_ConvertGpuLaunchFuncToVulkanLaunchFunc
45 #include "mlir/Conversion/Passes.h.inc"
46 
47   void runOnModule() override;
48 
49 private:
50   /// Creates a SPIR-V binary shader from the given `module` using
51   /// `spirv::serialize` function.
52   LogicalResult createBinaryShader(ModuleOp module,
53                                    std::vector<char> &binaryShader);
54 
55   /// Converts the given `launchOp` to vulkan launch call.
56   void convertGpuLaunchFunc(gpu::LaunchFuncOp launchOp);
57 
58   /// Checks where the given type is supported by Vulkan runtime.
59   bool isSupportedType(Type type) {
60     // TODO(denis0x0D): Handle other types.
61     if (auto memRefType = type.dyn_cast_or_null<MemRefType>())
62       return memRefType.hasRank() &&
63              (memRefType.getRank() >= 1 && memRefType.getRank() <= 3);
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 };
73 
74 } // anonymous namespace
75 
76 void ConvertGpuLaunchFuncToVulkanLaunchFunc::runOnModule() {
77   bool done = false;
78   getModule().walk([this, &done](gpu::LaunchFuncOp op) {
79     if (done) {
80       op.emitError("should only contain one 'gpu::LaunchFuncOp' op");
81       return signalPassFailure();
82     }
83     done = true;
84     convertGpuLaunchFunc(op);
85   });
86 
87   // Erase `gpu::GPUModuleOp` and `spirv::Module` operations.
88   for (auto gpuModule :
89        llvm::make_early_inc_range(getModule().getOps<gpu::GPUModuleOp>()))
90     gpuModule.erase();
91 
92   for (auto spirvModule :
93        llvm::make_early_inc_range(getModule().getOps<spirv::ModuleOp>()))
94     spirvModule.erase();
95 }
96 
97 LogicalResult ConvertGpuLaunchFuncToVulkanLaunchFunc::declareVulkanLaunchFunc(
98     Location loc, gpu::LaunchFuncOp launchOp) {
99   OpBuilder builder(getModule().getBody()->getTerminator());
100   // TODO: Workgroup size is written into the kernel. So to properly modelling
101   // vulkan launch, we cannot have the local workgroup size configuration here.
102   SmallVector<Type, 8> vulkanLaunchTypes{launchOp.getOperandTypes()};
103 
104   // Check that all operands have supported types except those for the launch
105   // configuration.
106   for (auto type :
107        llvm::drop_begin(vulkanLaunchTypes, gpu::LaunchOp::kNumConfigOperands)) {
108     if (!isSupportedType(type))
109       return launchOp.emitError() << type << " is unsupported to run on Vulkan";
110   }
111 
112   // Declare vulkan launch function.
113   builder.create<FuncOp>(
114       loc, kVulkanLaunch,
115       FunctionType::get(vulkanLaunchTypes, ArrayRef<Type>{}, loc->getContext()),
116       ArrayRef<NamedAttribute>{});
117 
118   return success();
119 }
120 
121 LogicalResult ConvertGpuLaunchFuncToVulkanLaunchFunc::createBinaryShader(
122     ModuleOp module, std::vector<char> &binaryShader) {
123   bool done = false;
124   SmallVector<uint32_t, 0> binary;
125   for (auto spirvModule : module.getOps<spirv::ModuleOp>()) {
126     if (done)
127       return spirvModule.emitError("should only contain one 'spv.module' op");
128     done = true;
129 
130     if (failed(spirv::serialize(spirvModule, binary)))
131       return failure();
132   }
133   binaryShader.resize(binary.size() * sizeof(uint32_t));
134   std::memcpy(binaryShader.data(), reinterpret_cast<char *>(binary.data()),
135               binaryShader.size());
136   return success();
137 }
138 
139 void ConvertGpuLaunchFuncToVulkanLaunchFunc::convertGpuLaunchFunc(
140     gpu::LaunchFuncOp launchOp) {
141   ModuleOp module = getModule();
142   OpBuilder builder(launchOp);
143   Location loc = launchOp.getLoc();
144 
145   // Serialize `spirv::Module` into binary form.
146   std::vector<char> binary;
147   if (failed(createBinaryShader(module, binary)))
148     return signalPassFailure();
149 
150   // Declare vulkan launch function.
151   if (failed(declareVulkanLaunchFunc(loc, launchOp)))
152     return signalPassFailure();
153 
154   // Create vulkan launch call op.
155   auto vulkanLaunchCallOp = builder.create<CallOp>(
156       loc, ArrayRef<Type>{}, builder.getSymbolRefAttr(kVulkanLaunch),
157       launchOp.getOperands());
158 
159   // Set SPIR-V binary shader data as an attribute.
160   vulkanLaunchCallOp.setAttr(
161       kSPIRVBlobAttrName,
162       StringAttr::get({binary.data(), binary.size()}, loc->getContext()));
163 
164   // Set entry point name as an attribute.
165   vulkanLaunchCallOp.setAttr(
166       kSPIRVEntryPointAttrName,
167       StringAttr::get(launchOp.kernel(), loc->getContext()));
168 
169   launchOp.erase();
170 }
171 
172 std::unique_ptr<mlir::OpPassBase<mlir::ModuleOp>>
173 mlir::createConvertGpuLaunchFuncToVulkanLaunchFuncPass() {
174   return std::make_unique<ConvertGpuLaunchFuncToVulkanLaunchFunc>();
175 }
176