328 lines · cpp
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-exception6//7//===----------------------------------------------------------------------===//8//9// This file implements passes to convert `gpu.launch_func` op into a sequence10// of LLVM calls that emulate the host and device sides.11//12//===----------------------------------------------------------------------===//13 14#include "mlir/Conversion/ArithToLLVM/ArithToLLVM.h"15#include "mlir/Conversion/FuncToLLVM/ConvertFuncToLLVM.h"16#include "mlir/Conversion/LLVMCommon/LoweringOptions.h"17#include "mlir/Conversion/LLVMCommon/Pattern.h"18#include "mlir/Conversion/LLVMCommon/TypeConverter.h"19#include "mlir/Conversion/MemRefToLLVM/MemRefToLLVM.h"20#include "mlir/Conversion/SPIRVToLLVM/SPIRVToLLVM.h"21#include "mlir/Dialect/Func/IR/FuncOps.h"22#include "mlir/Dialect/GPU/IR/GPUDialect.h"23#include "mlir/Dialect/LLVMIR/LLVMDialect.h"24#include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"25#include "mlir/IR/BuiltinOps.h"26#include "mlir/IR/SymbolTable.h"27#include "mlir/Pass/Pass.h"28#include "mlir/Transforms/DialectConversion.h"29#include "llvm/ADT/DenseMap.h"30#include "llvm/ADT/StringExtras.h"31#include "llvm/Support/FormatVariadic.h"32 33namespace mlir {34#define GEN_PASS_DEF_LOWERHOSTCODETOLLVMPASS35#include "mlir/Conversion/Passes.h.inc"36} // namespace mlir37 38using namespace mlir;39 40static constexpr const char kSPIRVModule[] = "__spv__";41 42//===----------------------------------------------------------------------===//43// Utility functions44//===----------------------------------------------------------------------===//45 46/// Returns the string name of the `DescriptorSet` decoration.47static std::string descriptorSetName() {48 return llvm::convertToSnakeFromCamelCase(49 stringifyDecoration(spirv::Decoration::DescriptorSet));50}51 52/// Returns the string name of the `Binding` decoration.53static std::string bindingName() {54 return llvm::convertToSnakeFromCamelCase(55 stringifyDecoration(spirv::Decoration::Binding));56}57 58/// Calculates the index of the kernel's operand that is represented by the59/// given global variable with the `bind` attribute. We assume that the index of60/// each kernel's operand is mapped to (descriptorSet, binding) by the map:61/// i -> (0, i)62/// which is implemented under `LowerABIAttributesPass`.63static unsigned calculateGlobalIndex(spirv::GlobalVariableOp op) {64 IntegerAttr binding = op->getAttrOfType<IntegerAttr>(bindingName());65 return binding.getInt();66}67 68/// Copies the given number of bytes from src to dst pointers.69static void copy(Location loc, Value dst, Value src, Value size,70 OpBuilder &builder) {71 LLVM::MemcpyOp::create(builder, loc, dst, src, size, /*isVolatile=*/false);72}73 74/// Encodes the binding and descriptor set numbers into a new symbolic name.75/// The name is specified by76/// {kernel_module_name}_{variable_name}_descriptor_set{ds}_binding{b}77/// to avoid symbolic conflicts, where 'ds' and 'b' are descriptor set and78/// binding numbers.79static std::string80createGlobalVariableWithBindName(spirv::GlobalVariableOp op,81 StringRef kernelModuleName) {82 IntegerAttr descriptorSet =83 op->getAttrOfType<IntegerAttr>(descriptorSetName());84 IntegerAttr binding = op->getAttrOfType<IntegerAttr>(bindingName());85 return llvm::formatv("{0}_{1}_descriptor_set{2}_binding{3}",86 kernelModuleName.str(), op.getSymName().str(),87 std::to_string(descriptorSet.getInt()),88 std::to_string(binding.getInt()));89}90 91/// Returns true if the given global variable has both a descriptor set number92/// and a binding number.93static bool hasDescriptorSetAndBinding(spirv::GlobalVariableOp op) {94 IntegerAttr descriptorSet =95 op->getAttrOfType<IntegerAttr>(descriptorSetName());96 IntegerAttr binding = op->getAttrOfType<IntegerAttr>(bindingName());97 return descriptorSet && binding;98}99 100/// Fills `globalVariableMap` with SPIR-V global variables that represent kernel101/// arguments from the given SPIR-V module. We assume that the module contains a102/// single entry point function. Hence, all `spirv.GlobalVariable`s with a bind103/// attribute are kernel arguments.104static LogicalResult getKernelGlobalVariables(105 spirv::ModuleOp module,106 DenseMap<uint32_t, spirv::GlobalVariableOp> &globalVariableMap) {107 auto entryPoints = module.getOps<spirv::EntryPointOp>();108 if (!llvm::hasSingleElement(entryPoints)) {109 return module.emitError(110 "The module must contain exactly one entry point function");111 }112 auto globalVariables = module.getOps<spirv::GlobalVariableOp>();113 for (auto globalOp : globalVariables) {114 if (hasDescriptorSetAndBinding(globalOp))115 globalVariableMap[calculateGlobalIndex(globalOp)] = globalOp;116 }117 return success();118}119 120/// Encodes the SPIR-V module's symbolic name into the name of the entry point121/// function.122static LogicalResult encodeKernelName(spirv::ModuleOp module) {123 StringRef spvModuleName = module.getSymName().value_or(kSPIRVModule);124 // We already know that the module contains exactly one entry point function125 // based on `getKernelGlobalVariables()` call. Update this function's name126 // to:127 // {spv_module_name}_{function_name}128 auto entryPoints = module.getOps<spirv::EntryPointOp>();129 if (!llvm::hasSingleElement(entryPoints)) {130 return module.emitError(131 "The module must contain exactly one entry point function");132 }133 spirv::EntryPointOp entryPoint = *entryPoints.begin();134 StringRef funcName = entryPoint.getFn();135 auto funcOp = module.lookupSymbol<spirv::FuncOp>(entryPoint.getFnAttr());136 StringAttr newFuncName =137 StringAttr::get(module->getContext(), spvModuleName + "_" + funcName);138 if (failed(SymbolTable::replaceAllSymbolUses(funcOp, newFuncName, module)))139 return failure();140 SymbolTable::setSymbolName(funcOp, newFuncName);141 return success();142}143 144//===----------------------------------------------------------------------===//145// Conversion patterns146//===----------------------------------------------------------------------===//147 148namespace {149 150/// Structure to group information about the variables being copied.151struct CopyInfo {152 Value dst;153 Value src;154 Value size;155};156 157/// This pattern emulates a call to the kernel in LLVM dialect. For that, we158/// copy the data to the global variable (emulating device side), call the159/// kernel as a normal void LLVM function, and copy the data back (emulating the160/// host side).161class GPULaunchLowering : public ConvertOpToLLVMPattern<gpu::LaunchFuncOp> {162 using ConvertOpToLLVMPattern<gpu::LaunchFuncOp>::ConvertOpToLLVMPattern;163 164 LogicalResult165 matchAndRewrite(gpu::LaunchFuncOp launchOp, OpAdaptor adaptor,166 ConversionPatternRewriter &rewriter) const override {167 auto *op = launchOp.getOperation();168 MLIRContext *context = rewriter.getContext();169 auto module = launchOp->getParentOfType<ModuleOp>();170 171 // Get the SPIR-V module that represents the gpu kernel module. The module172 // is named:173 // __spv__{kernel_module_name}174 // based on GPU to SPIR-V conversion.175 StringRef kernelModuleName = launchOp.getKernelModuleName().getValue();176 std::string spvModuleName = kSPIRVModule + kernelModuleName.str();177 auto spvModule = module.lookupSymbol<spirv::ModuleOp>(178 StringAttr::get(context, spvModuleName));179 if (!spvModule) {180 return launchOp.emitOpError("SPIR-V kernel module '")181 << spvModuleName << "' is not found";182 }183 184 // Declare kernel function in the main module so that it later can be linked185 // with its definition from the kernel module. We know that the kernel186 // function would have no arguments and the data is passed via global187 // variables. The name of the kernel will be188 // {spv_module_name}_{kernel_function_name}189 // to avoid symbolic name conflicts.190 StringRef kernelFuncName = launchOp.getKernelName().getValue();191 std::string newKernelFuncName = spvModuleName + "_" + kernelFuncName.str();192 auto kernelFunc = module.lookupSymbol<LLVM::LLVMFuncOp>(193 StringAttr::get(context, newKernelFuncName));194 if (!kernelFunc) {195 OpBuilder::InsertionGuard guard(rewriter);196 rewriter.setInsertionPointToStart(module.getBody());197 kernelFunc = LLVM::LLVMFuncOp::create(198 rewriter, rewriter.getUnknownLoc(), newKernelFuncName,199 LLVM::LLVMFunctionType::get(LLVM::LLVMVoidType::get(context),200 ArrayRef<Type>()));201 rewriter.setInsertionPoint(launchOp);202 }203 204 // Get all global variables associated with the kernel operands.205 DenseMap<uint32_t, spirv::GlobalVariableOp> globalVariableMap;206 if (failed(getKernelGlobalVariables(spvModule, globalVariableMap)))207 return failure();208 209 // Traverse kernel operands that were converted to MemRefDescriptors. For210 // each operand, create a global variable and copy data from operand to it.211 Location loc = launchOp.getLoc();212 SmallVector<CopyInfo, 4> copyInfo;213 auto numKernelOperands = launchOp.getNumKernelOperands();214 auto kernelOperands = adaptor.getOperands().take_back(numKernelOperands);215 for (const auto &operand : llvm::enumerate(kernelOperands)) {216 // Check if the kernel's operand is a ranked memref.217 auto memRefType = dyn_cast<MemRefType>(218 launchOp.getKernelOperand(operand.index()).getType());219 if (!memRefType)220 return failure();221 222 // Calculate the size of the memref and get the pointer to the allocated223 // buffer.224 SmallVector<Value, 4> sizes;225 SmallVector<Value, 4> strides;226 Value sizeBytes;227 getMemRefDescriptorSizes(loc, memRefType, {}, rewriter, sizes, strides,228 sizeBytes);229 MemRefDescriptor descriptor(operand.value());230 Value src = descriptor.allocatedPtr(rewriter, loc);231 232 // Get the global variable in the SPIR-V module that is associated with233 // the kernel operand. Construct its new name and create a corresponding234 // LLVM dialect global variable.235 spirv::GlobalVariableOp spirvGlobal = globalVariableMap[operand.index()];236 auto pointeeType =237 cast<spirv::PointerType>(spirvGlobal.getType()).getPointeeType();238 auto dstGlobalType = typeConverter->convertType(pointeeType);239 if (!dstGlobalType)240 return failure();241 std::string name =242 createGlobalVariableWithBindName(spirvGlobal, spvModuleName);243 // Check if this variable has already been created.244 auto dstGlobal = module.lookupSymbol<LLVM::GlobalOp>(name);245 if (!dstGlobal) {246 OpBuilder::InsertionGuard guard(rewriter);247 rewriter.setInsertionPointToStart(module.getBody());248 dstGlobal = LLVM::GlobalOp::create(249 rewriter, loc, dstGlobalType,250 /*isConstant=*/false, LLVM::Linkage::Linkonce, name, Attribute(),251 /*alignment=*/0);252 rewriter.setInsertionPoint(launchOp);253 }254 255 // Copy the data from src operand pointer to dst global variable. Save256 // src, dst and size so that we can copy data back after emulating the257 // kernel call.258 Value dst = LLVM::AddressOfOp::create(259 rewriter, loc, typeConverter->convertType(spirvGlobal.getType()),260 dstGlobal.getSymName());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 278class LowerHostCodeToLLVM279 : public impl::LowerHostCodeToLLVMPassBase<LowerHostCodeToLLVM> {280public:281 using Base::Base;282 283 void runOnOperation() override {284 ModuleOp module = getOperation();285 286 // Erase the GPU module.287 for (auto gpuModule :288 llvm::make_early_inc_range(module.getOps<gpu::GPUModuleOp>()))289 gpuModule.erase();290 291 // Request C wrapper emission.292 for (auto func : module.getOps<func::FuncOp>()) {293 func->setAttr(LLVM::LLVMDialect::getEmitCWrapperAttrName(),294 UnitAttr::get(&getContext()));295 }296 297 // Specify options to lower to LLVM and pull in the conversion patterns.298 LowerToLLVMOptions options(module.getContext());299 300 auto *context = module.getContext();301 RewritePatternSet patterns(context);302 LLVMTypeConverter typeConverter(context, options);303 mlir::arith::populateArithToLLVMConversionPatterns(typeConverter, patterns);304 populateFinalizeMemRefToLLVMConversionPatterns(typeConverter, patterns);305 populateFuncToLLVMConversionPatterns(typeConverter, patterns);306 patterns.add<GPULaunchLowering>(typeConverter);307 308 // Pull in SPIR-V type conversion patterns to convert SPIR-V global309 // variable's type to LLVM dialect type.310 populateSPIRVToLLVMTypeConversion(typeConverter);311 312 ConversionTarget target(*context);313 target.addLegalDialect<LLVM::LLVMDialect>();314 if (failed(applyPartialConversion(module, target, std::move(patterns))))315 signalPassFailure();316 317 // Finally, modify the kernel function in SPIR-V modules to avoid symbolic318 // conflicts.319 for (auto spvModule : module.getOps<spirv::ModuleOp>()) {320 if (failed(encodeKernelName(spvModule))) {321 signalPassFailure();322 return;323 }324 }325 }326};327} // namespace328