1 //===- ConvertLaunchFuncToLLVMCalls.cpp - MLIR GPU launch to LLVM 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 passes to convert `gpu.launch_func` op into a sequence
10 // of LLVM calls that emulate the host and device sides.
11 //
12 //===----------------------------------------------------------------------===//
13 
14 #include "../PassDetail.h"
15 #include "mlir/Conversion/SPIRVToLLVM/SPIRVToLLVM.h"
16 #include "mlir/Conversion/SPIRVToLLVM/SPIRVToLLVMPass.h"
17 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h"
18 #include "mlir/Dialect/GPU/GPUDialect.h"
19 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
20 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
21 #include "mlir/Dialect/StandardOps/IR/Ops.h"
22 #include "mlir/IR/BuiltinOps.h"
23 #include "mlir/IR/SymbolTable.h"
24 #include "mlir/Transforms/DialectConversion.h"
25 
26 #include "llvm/ADT/DenseMap.h"
27 #include "llvm/ADT/StringExtras.h"
28 #include "llvm/Support/FormatVariadic.h"
29 
30 using namespace mlir;
31 
32 static constexpr const char kSPIRVModule[] = "__spv__";
33 
34 //===----------------------------------------------------------------------===//
35 // Utility functions
36 //===----------------------------------------------------------------------===//
37 
38 /// Returns the string name of the `DescriptorSet` decoration.
39 static std::string descriptorSetName() {
40   return llvm::convertToSnakeFromCamelCase(
41       stringifyDecoration(spirv::Decoration::DescriptorSet));
42 }
43 
44 /// Returns the string name of the `Binding` decoration.
45 static std::string bindingName() {
46   return llvm::convertToSnakeFromCamelCase(
47       stringifyDecoration(spirv::Decoration::Binding));
48 }
49 
50 /// Calculates the index of the kernel's operand that is represented by the
51 /// given global variable with the `bind` attribute. We assume that the index of
52 /// each kernel's operand is mapped to (descriptorSet, binding) by the map:
53 ///   i -> (0, i)
54 /// which is implemented under `LowerABIAttributesPass`.
55 static unsigned calculateGlobalIndex(spirv::GlobalVariableOp op) {
56   IntegerAttr binding = op->getAttrOfType<IntegerAttr>(bindingName());
57   return binding.getInt();
58 }
59 
60 /// Copies the given number of bytes from src to dst pointers.
61 static void copy(Location loc, Value dst, Value src, Value size,
62                  OpBuilder &builder) {
63   MLIRContext *context = builder.getContext();
64   auto llvmI1Type = IntegerType::get(context, 1);
65   Value isVolatile = builder.create<LLVM::ConstantOp>(
66       loc, llvmI1Type, builder.getBoolAttr(false));
67   builder.create<LLVM::MemcpyOp>(loc, dst, src, size, isVolatile);
68 }
69 
70 /// Encodes the binding and descriptor set numbers into a new symbolic name.
71 /// The name is specified by
72 ///   {kernel_module_name}_{variable_name}_descriptor_set{ds}_binding{b}
73 /// to avoid symbolic conflicts, where 'ds' and 'b' are descriptor set and
74 /// binding numbers.
75 static std::string
76 createGlobalVariableWithBindName(spirv::GlobalVariableOp op,
77                                  StringRef kernelModuleName) {
78   IntegerAttr descriptorSet =
79       op->getAttrOfType<IntegerAttr>(descriptorSetName());
80   IntegerAttr binding = op->getAttrOfType<IntegerAttr>(bindingName());
81   return llvm::formatv("{0}_{1}_descriptor_set{2}_binding{3}",
82                        kernelModuleName.str(), op.sym_name().str(),
83                        std::to_string(descriptorSet.getInt()),
84                        std::to_string(binding.getInt()));
85 }
86 
87 /// Returns true if the given global variable has both a descriptor set number
88 /// and a binding number.
89 static bool hasDescriptorSetAndBinding(spirv::GlobalVariableOp op) {
90   IntegerAttr descriptorSet =
91       op->getAttrOfType<IntegerAttr>(descriptorSetName());
92   IntegerAttr binding = op->getAttrOfType<IntegerAttr>(bindingName());
93   return descriptorSet && binding;
94 }
95 
96 /// Fills `globalVariableMap` with SPIR-V global variables that represent kernel
97 /// arguments from the given SPIR-V module. We assume that the module contains a
98 /// single entry point function. Hence, all `spv.GlobalVariable`s with a bind
99 /// attribute are kernel arguments.
100 static LogicalResult getKernelGlobalVariables(
101     spirv::ModuleOp module,
102     DenseMap<uint32_t, spirv::GlobalVariableOp> &globalVariableMap) {
103   auto entryPoints = module.getOps<spirv::EntryPointOp>();
104   if (!llvm::hasSingleElement(entryPoints)) {
105     return module.emitError(
106         "The module must contain exactly one entry point function");
107   }
108   auto globalVariables = module.getOps<spirv::GlobalVariableOp>();
109   for (auto globalOp : globalVariables) {
110     if (hasDescriptorSetAndBinding(globalOp))
111       globalVariableMap[calculateGlobalIndex(globalOp)] = globalOp;
112   }
113   return success();
114 }
115 
116 /// Encodes the SPIR-V module's symbolic name into the name of the entry point
117 /// function.
118 static LogicalResult encodeKernelName(spirv::ModuleOp module) {
119   StringRef spvModuleName = module.sym_name().getValue();
120   // We already know that the module contains exactly one entry point function
121   // based on `getKernelGlobalVariables()` call. Update this function's name
122   // to:
123   //   {spv_module_name}_{function_name}
124   auto entryPoint = *module.getOps<spirv::EntryPointOp>().begin();
125   StringRef funcName = entryPoint.fn();
126   auto funcOp = module.lookupSymbol<spirv::FuncOp>(funcName);
127   std::string newFuncName = spvModuleName.str() + "_" + funcName.str();
128   if (failed(SymbolTable::replaceAllSymbolUses(funcOp, newFuncName, module)))
129     return failure();
130   SymbolTable::setSymbolName(funcOp, newFuncName);
131   return success();
132 }
133 
134 //===----------------------------------------------------------------------===//
135 // Conversion patterns
136 //===----------------------------------------------------------------------===//
137 
138 namespace {
139 
140 /// Structure to group information about the variables being copied.
141 struct CopyInfo {
142   Value dst;
143   Value src;
144   Value size;
145 };
146 
147 /// This pattern emulates a call to the kernel in LLVM dialect. For that, we
148 /// copy the data to the global variable (emulating device side), call the
149 /// kernel as a normal void LLVM function, and copy the data back (emulating the
150 /// host side).
151 class GPULaunchLowering : public ConvertOpToLLVMPattern<gpu::LaunchFuncOp> {
152   using ConvertOpToLLVMPattern<gpu::LaunchFuncOp>::ConvertOpToLLVMPattern;
153 
154   LogicalResult
155   matchAndRewrite(gpu::LaunchFuncOp launchOp, ArrayRef<Value> operands,
156                   ConversionPatternRewriter &rewriter) const override {
157     auto *op = launchOp.getOperation();
158     MLIRContext *context = rewriter.getContext();
159     auto module = launchOp->getParentOfType<ModuleOp>();
160 
161     // Get the SPIR-V module that represents the gpu kernel module. The module
162     // is named:
163     //   __spv__{kernel_module_name}
164     // based on GPU to SPIR-V conversion.
165     StringRef kernelModuleName = launchOp.getKernelModuleName();
166     std::string spvModuleName = kSPIRVModule + kernelModuleName.str();
167     auto spvModule = module.lookupSymbol<spirv::ModuleOp>(spvModuleName);
168     if (!spvModule) {
169       return launchOp.emitOpError("SPIR-V kernel module '")
170              << spvModuleName << "' is not found";
171     }
172 
173     // Declare kernel function in the main module so that it later can be linked
174     // with its definition from the kernel module. We know that the kernel
175     // function would have no arguments and the data is passed via global
176     // variables. The name of the kernel will be
177     //   {spv_module_name}_{kernel_function_name}
178     // to avoid symbolic name conflicts.
179     StringRef kernelFuncName = launchOp.getKernelName();
180     std::string newKernelFuncName = spvModuleName + "_" + kernelFuncName.str();
181     auto kernelFunc = module.lookupSymbol<LLVM::LLVMFuncOp>(newKernelFuncName);
182     if (!kernelFunc) {
183       OpBuilder::InsertionGuard guard(rewriter);
184       rewriter.setInsertionPointToStart(module.getBody());
185       kernelFunc = rewriter.create<LLVM::LLVMFuncOp>(
186           rewriter.getUnknownLoc(), newKernelFuncName,
187           LLVM::LLVMFunctionType::get(LLVM::LLVMVoidType::get(context),
188                                       ArrayRef<Type>()));
189       rewriter.setInsertionPoint(launchOp);
190     }
191 
192     // Get all global variables associated with the kernel operands.
193     DenseMap<uint32_t, spirv::GlobalVariableOp> globalVariableMap;
194     if (failed(getKernelGlobalVariables(spvModule, globalVariableMap)))
195       return failure();
196 
197     // Traverse kernel operands that were converted to MemRefDescriptors. For
198     // each operand, create a global variable and copy data from operand to it.
199     Location loc = launchOp.getLoc();
200     SmallVector<CopyInfo, 4> copyInfo;
201     auto numKernelOperands = launchOp.getNumKernelOperands();
202     auto kernelOperands = operands.take_back(numKernelOperands);
203     for (auto operand : llvm::enumerate(kernelOperands)) {
204       // Check if the kernel's operand is a ranked memref.
205       auto memRefType = launchOp.getKernelOperand(operand.index())
206                             .getType()
207                             .dyn_cast<MemRefType>();
208       if (!memRefType)
209         return failure();
210 
211       // Calculate the size of the memref and get the pointer to the allocated
212       // buffer.
213       SmallVector<Value, 4> sizes;
214       SmallVector<Value, 4> strides;
215       Value sizeBytes;
216       getMemRefDescriptorSizes(loc, memRefType, {}, rewriter, sizes, strides,
217                                sizeBytes);
218       MemRefDescriptor descriptor(operand.value());
219       Value src = descriptor.allocatedPtr(rewriter, loc);
220 
221       // Get the global variable in the SPIR-V module that is associated with
222       // the kernel operand. Construct its new name and create a corresponding
223       // LLVM dialect global variable.
224       spirv::GlobalVariableOp spirvGlobal = globalVariableMap[operand.index()];
225       auto pointeeType =
226           spirvGlobal.type().cast<spirv::PointerType>().getPointeeType();
227       auto dstGlobalType = typeConverter->convertType(pointeeType);
228       if (!dstGlobalType)
229         return failure();
230       std::string name =
231           createGlobalVariableWithBindName(spirvGlobal, spvModuleName);
232       // Check if this variable has already been created.
233       auto dstGlobal = module.lookupSymbol<LLVM::GlobalOp>(name);
234       if (!dstGlobal) {
235         OpBuilder::InsertionGuard guard(rewriter);
236         rewriter.setInsertionPointToStart(module.getBody());
237         dstGlobal = rewriter.create<LLVM::GlobalOp>(
238             loc, dstGlobalType,
239             /*isConstant=*/false, LLVM::Linkage::Linkonce, name, Attribute(),
240             /*alignment=*/0);
241         rewriter.setInsertionPoint(launchOp);
242       }
243 
244       // Copy the data from src operand pointer to dst global variable. Save
245       // src, dst and size so that we can copy data back after emulating the
246       // kernel call.
247       Value dst = rewriter.create<LLVM::AddressOfOp>(loc, dstGlobal);
248       copy(loc, dst, src, sizeBytes, rewriter);
249 
250       CopyInfo info;
251       info.dst = dst;
252       info.src = src;
253       info.size = sizeBytes;
254       copyInfo.push_back(info);
255     }
256     // Create a call to the kernel and copy the data back.
257     rewriter.replaceOpWithNewOp<LLVM::CallOp>(op, kernelFunc,
258                                               ArrayRef<Value>());
259     for (CopyInfo info : copyInfo)
260       copy(loc, info.src, info.dst, info.size, rewriter);
261     return success();
262   }
263 };
264 
265 class LowerHostCodeToLLVM
266     : public LowerHostCodeToLLVMBase<LowerHostCodeToLLVM> {
267 public:
268   void runOnOperation() override {
269     ModuleOp module = getOperation();
270 
271     // Erase the GPU module.
272     for (auto gpuModule :
273          llvm::make_early_inc_range(module.getOps<gpu::GPUModuleOp>()))
274       gpuModule.erase();
275 
276     // Specify options to lower Standard to LLVM and pull in the conversion
277     // patterns.
278     LowerToLLVMOptions options(module.getContext());
279     options.emitCWrappers = true;
280     auto *context = module.getContext();
281     RewritePatternSet patterns(context);
282     LLVMTypeConverter typeConverter(context, options);
283     populateStdToLLVMConversionPatterns(typeConverter, patterns);
284     patterns.add<GPULaunchLowering>(typeConverter);
285 
286     // Pull in SPIR-V type conversion patterns to convert SPIR-V global
287     // variable's type to LLVM dialect type.
288     populateSPIRVToLLVMTypeConversion(typeConverter);
289 
290     ConversionTarget target(*context);
291     target.addLegalDialect<LLVM::LLVMDialect>();
292     if (failed(applyPartialConversion(module, target, std::move(patterns))))
293       signalPassFailure();
294 
295     // Finally, modify the kernel function in SPIR-V modules to avoid symbolic
296     // conflicts.
297     for (auto spvModule : module.getOps<spirv::ModuleOp>())
298       (void)encodeKernelName(spvModule);
299   }
300 };
301 } // namespace
302 
303 std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
304 mlir::createLowerHostCodeToLLVMPass() {
305   return std::make_unique<LowerHostCodeToLLVM>();
306 }
307