MLIR  16.0.0git
ConvertLaunchFuncToLLVMCalls.cpp
Go to the documentation of this file.
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 
15 
27 #include "mlir/IR/BuiltinOps.h"
28 #include "mlir/IR/SymbolTable.h"
29 #include "mlir/Pass/Pass.h"
31 #include "llvm/ADT/DenseMap.h"
32 #include "llvm/ADT/StringExtras.h"
33 #include "llvm/Support/FormatVariadic.h"
34 
35 namespace mlir {
36 #define GEN_PASS_DEF_LOWERHOSTCODETOLLVM
37 #include "mlir/Conversion/Passes.h.inc"
38 } // namespace mlir
39 
40 using namespace mlir;
41 
42 static constexpr const char kSPIRVModule[] = "__spv__";
43 
44 //===----------------------------------------------------------------------===//
45 // Utility functions
46 //===----------------------------------------------------------------------===//
47 
48 /// Returns the string name of the `DescriptorSet` decoration.
49 static std::string descriptorSetName() {
50  return llvm::convertToSnakeFromCamelCase(
51  stringifyDecoration(spirv::Decoration::DescriptorSet));
52 }
53 
54 /// Returns the string name of the `Binding` decoration.
55 static std::string bindingName() {
56  return llvm::convertToSnakeFromCamelCase(
57  stringifyDecoration(spirv::Decoration::Binding));
58 }
59 
60 /// Calculates the index of the kernel's operand that is represented by the
61 /// given global variable with the `bind` attribute. We assume that the index of
62 /// each kernel's operand is mapped to (descriptorSet, binding) by the map:
63 /// i -> (0, i)
64 /// which is implemented under `LowerABIAttributesPass`.
65 static unsigned calculateGlobalIndex(spirv::GlobalVariableOp op) {
66  IntegerAttr binding = op->getAttrOfType<IntegerAttr>(bindingName());
67  return binding.getInt();
68 }
69 
70 /// Copies the given number of bytes from src to dst pointers.
71 static void copy(Location loc, Value dst, Value src, Value size,
72  OpBuilder &builder) {
73  MLIRContext *context = builder.getContext();
74  auto llvmI1Type = IntegerType::get(context, 1);
75  Value isVolatile = builder.create<LLVM::ConstantOp>(
76  loc, llvmI1Type, builder.getBoolAttr(false));
77  builder.create<LLVM::MemcpyOp>(loc, dst, src, size, isVolatile);
78 }
79 
80 /// Encodes the binding and descriptor set numbers into a new symbolic name.
81 /// The name is specified by
82 /// {kernel_module_name}_{variable_name}_descriptor_set{ds}_binding{b}
83 /// to avoid symbolic conflicts, where 'ds' and 'b' are descriptor set and
84 /// binding numbers.
85 static std::string
86 createGlobalVariableWithBindName(spirv::GlobalVariableOp op,
87  StringRef kernelModuleName) {
88  IntegerAttr descriptorSet =
89  op->getAttrOfType<IntegerAttr>(descriptorSetName());
90  IntegerAttr binding = op->getAttrOfType<IntegerAttr>(bindingName());
91  return llvm::formatv("{0}_{1}_descriptor_set{2}_binding{3}",
92  kernelModuleName.str(), op.getSymName().str(),
93  std::to_string(descriptorSet.getInt()),
94  std::to_string(binding.getInt()));
95 }
96 
97 /// Returns true if the given global variable has both a descriptor set number
98 /// and a binding number.
99 static bool hasDescriptorSetAndBinding(spirv::GlobalVariableOp op) {
100  IntegerAttr descriptorSet =
101  op->getAttrOfType<IntegerAttr>(descriptorSetName());
102  IntegerAttr binding = op->getAttrOfType<IntegerAttr>(bindingName());
103  return descriptorSet && binding;
104 }
105 
106 /// Fills `globalVariableMap` with SPIR-V global variables that represent kernel
107 /// arguments from the given SPIR-V module. We assume that the module contains a
108 /// single entry point function. Hence, all `spirv.GlobalVariable`s with a bind
109 /// attribute are kernel arguments.
111  spirv::ModuleOp module,
112  DenseMap<uint32_t, spirv::GlobalVariableOp> &globalVariableMap) {
113  auto entryPoints = module.getOps<spirv::EntryPointOp>();
114  if (!llvm::hasSingleElement(entryPoints)) {
115  return module.emitError(
116  "The module must contain exactly one entry point function");
117  }
118  auto globalVariables = module.getOps<spirv::GlobalVariableOp>();
119  for (auto globalOp : globalVariables) {
120  if (hasDescriptorSetAndBinding(globalOp))
121  globalVariableMap[calculateGlobalIndex(globalOp)] = globalOp;
122  }
123  return success();
124 }
125 
126 /// Encodes the SPIR-V module's symbolic name into the name of the entry point
127 /// function.
128 static LogicalResult encodeKernelName(spirv::ModuleOp module) {
129  StringRef spvModuleName = *module.getSymName();
130  // We already know that the module contains exactly one entry point function
131  // based on `getKernelGlobalVariables()` call. Update this function's name
132  // to:
133  // {spv_module_name}_{function_name}
134  auto entryPoint = *module.getOps<spirv::EntryPointOp>().begin();
135  StringRef funcName = entryPoint.getFn();
136  auto funcOp = module.lookupSymbol<spirv::FuncOp>(entryPoint.getFnAttr());
137  StringAttr newFuncName =
138  StringAttr::get(module->getContext(), spvModuleName + "_" + funcName);
139  if (failed(SymbolTable::replaceAllSymbolUses(funcOp, newFuncName, module)))
140  return failure();
141  SymbolTable::setSymbolName(funcOp, newFuncName);
142  return success();
143 }
144 
145 //===----------------------------------------------------------------------===//
146 // Conversion patterns
147 //===----------------------------------------------------------------------===//
148 
149 namespace {
150 
151 /// Structure to group information about the variables being copied.
152 struct CopyInfo {
153  Value dst;
154  Value src;
155  Value size;
156 };
157 
158 /// This pattern emulates a call to the kernel in LLVM dialect. For that, we
159 /// copy the data to the global variable (emulating device side), call the
160 /// kernel as a normal void LLVM function, and copy the data back (emulating the
161 /// host side).
162 class GPULaunchLowering : public ConvertOpToLLVMPattern<gpu::LaunchFuncOp> {
164 
166  matchAndRewrite(gpu::LaunchFuncOp launchOp, OpAdaptor adaptor,
167  ConversionPatternRewriter &rewriter) const override {
168  auto *op = launchOp.getOperation();
169  MLIRContext *context = rewriter.getContext();
170  auto module = launchOp->getParentOfType<ModuleOp>();
171 
172  // Get the SPIR-V module that represents the gpu kernel module. The module
173  // is named:
174  // __spv__{kernel_module_name}
175  // based on GPU to SPIR-V conversion.
176  StringRef kernelModuleName = launchOp.getKernelModuleName().getValue();
177  std::string spvModuleName = kSPIRVModule + kernelModuleName.str();
178  auto spvModule = module.lookupSymbol<spirv::ModuleOp>(
179  StringAttr::get(context, spvModuleName));
180  if (!spvModule) {
181  return launchOp.emitOpError("SPIR-V kernel module '")
182  << spvModuleName << "' is not found";
183  }
184 
185  // Declare kernel function in the main module so that it later can be linked
186  // with its definition from the kernel module. We know that the kernel
187  // function would have no arguments and the data is passed via global
188  // variables. The name of the kernel will be
189  // {spv_module_name}_{kernel_function_name}
190  // to avoid symbolic name conflicts.
191  StringRef kernelFuncName = launchOp.getKernelName().getValue();
192  std::string newKernelFuncName = spvModuleName + "_" + kernelFuncName.str();
193  auto kernelFunc = module.lookupSymbol<LLVM::LLVMFuncOp>(
194  StringAttr::get(context, newKernelFuncName));
195  if (!kernelFunc) {
196  OpBuilder::InsertionGuard guard(rewriter);
197  rewriter.setInsertionPointToStart(module.getBody());
198  kernelFunc = rewriter.create<LLVM::LLVMFuncOp>(
199  rewriter.getUnknownLoc(), newKernelFuncName,
200  LLVM::LLVMFunctionType::get(LLVM::LLVMVoidType::get(context),
201  ArrayRef<Type>()));
202  rewriter.setInsertionPoint(launchOp);
203  }
204 
205  // Get all global variables associated with the kernel operands.
207  if (failed(getKernelGlobalVariables(spvModule, globalVariableMap)))
208  return failure();
209 
210  // Traverse kernel operands that were converted to MemRefDescriptors. For
211  // each operand, create a global variable and copy data from operand to it.
212  Location loc = launchOp.getLoc();
213  SmallVector<CopyInfo, 4> copyInfo;
214  auto numKernelOperands = launchOp.getNumKernelOperands();
215  auto kernelOperands = adaptor.getOperands().take_back(numKernelOperands);
216  for (const auto &operand : llvm::enumerate(kernelOperands)) {
217  // Check if the kernel's operand is a ranked memref.
218  auto memRefType = launchOp.getKernelOperand(operand.index())
219  .getType()
220  .dyn_cast<MemRefType>();
221  if (!memRefType)
222  return failure();
223 
224  // Calculate the size of the memref and get the pointer to the allocated
225  // buffer.
226  SmallVector<Value, 4> sizes;
227  SmallVector<Value, 4> strides;
228  Value sizeBytes;
229  getMemRefDescriptorSizes(loc, memRefType, {}, rewriter, sizes, strides,
230  sizeBytes);
231  MemRefDescriptor descriptor(operand.value());
232  Value src = descriptor.allocatedPtr(rewriter, loc);
233 
234  // Get the global variable in the SPIR-V module that is associated with
235  // the kernel operand. Construct its new name and create a corresponding
236  // LLVM dialect global variable.
237  spirv::GlobalVariableOp spirvGlobal = globalVariableMap[operand.index()];
238  auto pointeeType =
239  spirvGlobal.getType().cast<spirv::PointerType>().getPointeeType();
240  auto dstGlobalType = typeConverter->convertType(pointeeType);
241  if (!dstGlobalType)
242  return failure();
243  std::string name =
244  createGlobalVariableWithBindName(spirvGlobal, spvModuleName);
245  // Check if this variable has already been created.
246  auto dstGlobal = module.lookupSymbol<LLVM::GlobalOp>(name);
247  if (!dstGlobal) {
248  OpBuilder::InsertionGuard guard(rewriter);
249  rewriter.setInsertionPointToStart(module.getBody());
250  dstGlobal = rewriter.create<LLVM::GlobalOp>(
251  loc, dstGlobalType,
252  /*isConstant=*/false, LLVM::Linkage::Linkonce, name, Attribute(),
253  /*alignment=*/0);
254  rewriter.setInsertionPoint(launchOp);
255  }
256 
257  // Copy the data from src operand pointer to dst global variable. Save
258  // src, dst and size so that we can copy data back after emulating the
259  // kernel call.
260  Value dst = rewriter.create<LLVM::AddressOfOp>(loc, dstGlobal);
261  copy(loc, dst, src, sizeBytes, rewriter);
262 
263  CopyInfo info;
264  info.dst = dst;
265  info.src = src;
266  info.size = sizeBytes;
267  copyInfo.push_back(info);
268  }
269  // Create a call to the kernel and copy the data back.
270  rewriter.replaceOpWithNewOp<LLVM::CallOp>(op, kernelFunc,
271  ArrayRef<Value>());
272  for (CopyInfo info : copyInfo)
273  copy(loc, info.src, info.dst, info.size, rewriter);
274  return success();
275  }
276 };
277 
278 class LowerHostCodeToLLVM
279  : public impl::LowerHostCodeToLLVMBase<LowerHostCodeToLLVM> {
280 public:
281  void runOnOperation() override {
282  ModuleOp module = getOperation();
283 
284  // Erase the GPU module.
285  for (auto gpuModule :
286  llvm::make_early_inc_range(module.getOps<gpu::GPUModuleOp>()))
287  gpuModule.erase();
288 
289  // Request C wrapper emission.
290  for (auto func : module.getOps<func::FuncOp>()) {
291  func->setAttr(LLVM::LLVMDialect::getEmitCWrapperAttrName(),
292  UnitAttr::get(&getContext()));
293  }
294 
295  // Specify options to lower to LLVM and pull in the conversion patterns.
296  LowerToLLVMOptions options(module.getContext());
297  auto *context = module.getContext();
298  RewritePatternSet patterns(context);
299  LLVMTypeConverter typeConverter(context, options);
301  populateMemRefToLLVMConversionPatterns(typeConverter, patterns);
302  populateFuncToLLVMConversionPatterns(typeConverter, patterns);
303  patterns.add<GPULaunchLowering>(typeConverter);
304 
305  // Pull in SPIR-V type conversion patterns to convert SPIR-V global
306  // variable's type to LLVM dialect type.
307  populateSPIRVToLLVMTypeConversion(typeConverter);
308 
309  ConversionTarget target(*context);
310  target.addLegalDialect<LLVM::LLVMDialect>();
311  if (failed(applyPartialConversion(module, target, std::move(patterns))))
312  signalPassFailure();
313 
314  // Finally, modify the kernel function in SPIR-V modules to avoid symbolic
315  // conflicts.
316  for (auto spvModule : module.getOps<spirv::ModuleOp>())
317  (void)encodeKernelName(spvModule);
318  }
319 };
320 } // namespace
321 
322 std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
324  return std::make_unique<LowerHostCodeToLLVM>();
325 }
static LogicalResult getKernelGlobalVariables(spirv::ModuleOp module, DenseMap< uint32_t, spirv::GlobalVariableOp > &globalVariableMap)
Fills globalVariableMap with SPIR-V global variables that represent kernel arguments from the given S...
Location getUnknownLoc()
Definition: Builders.cpp:26
Include the generated interface declarations.
Utility class for operation conversions targeting the LLVM dialect that match exactly one source oper...
Definition: Pattern.h:133
MLIRContext * getContext() const
Definition: Builders.h:54
RewritePatternSet & add(ConstructorArg &&arg, ConstructorArgs &&...args)
Add an instance of each of the pattern types &#39;Ts&#39; to the pattern list with the given arguments...
LogicalResult applyPartialConversion(ArrayRef< Operation *> ops, ConversionTarget &target, const FrozenRewritePatternSet &patterns, DenseSet< Operation *> *unconvertedOps=nullptr)
Below we define several entry points for operation conversion.
static std::string bindingName()
Returns the string name of the Binding decoration.
void setInsertionPoint(Block *block, Block::iterator insertPoint)
Set the insertion point to the specified location.
Definition: Builders.h:349
static std::string createGlobalVariableWithBindName(spirv::GlobalVariableOp op, StringRef kernelModuleName)
Encodes the binding and descriptor set numbers into a new symbolic name.
static void copy(Location loc, Value dst, Value src, Value size, OpBuilder &builder)
Copies the given number of bytes from src to dst pointers.
bool failed(LogicalResult result)
Utility function that returns true if the provided LogicalResult corresponds to a failure value...
Definition: LogicalResult.h:72
static std::string descriptorSetName()
Returns the string name of the DescriptorSet decoration.
static LLVMFunctionType get(Type result, ArrayRef< Type > arguments, bool isVarArg=false)
Gets or creates an instance of LLVM dialect function in the same context as the result type...
Definition: LLVMTypes.cpp:112
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Definition: Location.h:48
void populateMemRefToLLVMConversionPatterns(LLVMTypeConverter &converter, RewritePatternSet &patterns)
Collect a set of patterns to convert memory-related operations from the MemRef dialect to the LLVM di...
LogicalResult success(bool isSuccess=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:56
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
Definition: Builders.cpp:418
static void setSymbolName(Operation *symbol, StringAttr name)
Sets the name of the given symbol operation.
This class represents an efficient way to signal success or failure.
Definition: LogicalResult.h:26
LogicalResult failure(bool isFailure=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:62
Helper class to produce LLVM dialect operations extracting or inserting elements of a MemRef descript...
Definition: MemRefBuilder.h:33
Attributes are known-constant values of operations.
Definition: Attributes.h:25
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
Definition: Matchers.h:232
void populateArithToLLVMConversionPatterns(LLVMTypeConverter &converter, RewritePatternSet &patterns)
std::unique_ptr< OperationPass< ModuleOp > > createLowerHostCodeToLLVMPass()
Creates a pass to emulate gpu.launch_func call in LLVM dialect and lower the host module code to LLVM...
void addLegalDialect(StringRef name, Names... names)
Register the operations of the given dialects as legal.
static bool hasDescriptorSetAndBinding(spirv::GlobalVariableOp op)
Returns true if the given global variable has both a descriptor set number and a binding number...
static unsigned calculateGlobalIndex(spirv::GlobalVariableOp op)
Calculates the index of the kernel&#39;s operand that is represented by the given global variable with th...
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:85
static llvm::ManagedStatic< PassManagerOptions > options
void setInsertionPointToStart(Block *block)
Sets the insertion point to the start of the specified block.
Definition: Builders.h:382
RAII guard to reset the insertion point of the builder when destroyed.
Definition: Builders.h:299
OpTy replaceOpWithNewOp(Operation *op, Args &&...args)
Replaces the result op with a new op that is created without verification.
Definition: PatternMatch.h:451
Conversion from types to the LLVM IR dialect.
Definition: TypeConverter.h:30
BoolAttr getBoolAttr(bool value)
Definition: Builders.cpp:97
MLIRContext is the top-level object for a collection of MLIR operations.
Definition: MLIRContext.h:56
Options to control the LLVM lowering.
This class implements a pattern rewriter for use with ConversionPatterns.
This class describes a specific conversion target.
static LogicalResult replaceAllSymbolUses(StringAttr oldSymbol, StringAttr newSymbol, Operation *from)
Attempt to replace all uses of the given symbol &#39;oldSymbol&#39; with the provided symbol &#39;newSymbol&#39; that...
static LogicalResult encodeKernelName(spirv::ModuleOp module)
Encodes the SPIR-V module&#39;s symbolic name into the name of the entry point function.
void populateFuncToLLVMConversionPatterns(LLVMTypeConverter &converter, RewritePatternSet &patterns)
Collect the patterns to convert from the Func dialect to LLVM.
Definition: FuncToLLVM.cpp:692
void populateSPIRVToLLVMTypeConversion(LLVMTypeConverter &typeConverter)
Populates type conversions with additional SPIR-V types.
This class helps build Operations.
Definition: Builders.h:197
static constexpr const char kSPIRVModule[]