216 lines · c
1//===- GPUOpsLowering.h - GPU FuncOp / ReturnOp lowering -------*- C++ -*--===//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#ifndef MLIR_CONVERSION_GPUCOMMON_GPUOPSLOWERING_H_9#define MLIR_CONVERSION_GPUCOMMON_GPUOPSLOWERING_H_10 11#include "mlir/Conversion/LLVMCommon/Pattern.h"12#include "mlir/Dialect/GPU/IR/GPUDialect.h"13#include "mlir/Dialect/LLVMIR/LLVMAttrs.h"14#include "mlir/Dialect/LLVMIR/LLVMDialect.h"15 16namespace mlir {17 18//===----------------------------------------------------------------------===//19// Helper Functions20//===----------------------------------------------------------------------===//21 22/// Note that these functions don't take a `SymbolTable` because GPU module23/// lowerings can have name collisions as an intermediate state.24 25/// Find or create an external function declaration in the given module.26LLVM::LLVMFuncOp getOrDefineFunction(Operation *moduleOp, Location loc,27 OpBuilder &b, StringRef name,28 LLVM::LLVMFunctionType type);29 30/// Create a global that contains the given string. If a global with the same31/// string already exists in the module, return that global.32LLVM::GlobalOp getOrCreateStringConstant(OpBuilder &b, Location loc,33 Operation *moduleOp, Type llvmI8,34 StringRef namePrefix, StringRef str,35 uint64_t alignment = 0,36 unsigned addrSpace = 0);37 38//===----------------------------------------------------------------------===//39// Lowering Patterns40//===----------------------------------------------------------------------===//41 42/// Lowering for gpu.dynamic.shared.memory to LLVM dialect. The pattern first43/// create a 0-sized global array symbol similar as LLVM expects. It constructs44/// a memref descriptor with these values and return it.45struct GPUDynamicSharedMemoryOpLowering46 : public ConvertOpToLLVMPattern<gpu::DynamicSharedMemoryOp> {47 using ConvertOpToLLVMPattern<48 gpu::DynamicSharedMemoryOp>::ConvertOpToLLVMPattern;49 GPUDynamicSharedMemoryOpLowering(const LLVMTypeConverter &converter,50 unsigned alignmentBit = 0,51 PatternBenefit benefit = 1)52 : ConvertOpToLLVMPattern<gpu::DynamicSharedMemoryOp>(converter, benefit),53 alignmentBit(alignmentBit) {}54 55 LogicalResult56 matchAndRewrite(gpu::DynamicSharedMemoryOp op, OpAdaptor adaptor,57 ConversionPatternRewriter &rewriter) const override;58 59private:60 // Alignment bit61 unsigned alignmentBit;62};63 64struct GPUFuncOpLoweringOptions {65 /// The address space to use for `alloca`s in private memory.66 unsigned allocaAddrSpace;67 /// The address space to use declaring workgroup memory.68 unsigned workgroupAddrSpace;69 70 /// The attribute name to use instead of `gpu.kernel`. Null if no attribute71 /// should be used.72 StringAttr kernelAttributeName;73 /// The attribute name to to set block size. Null if no attribute should be74 /// used.75 StringAttr kernelBlockSizeAttributeName;76 77 /// The calling convention to use for kernel functions.78 LLVM::CConv kernelCallingConvention = LLVM::CConv::C;79 /// The calling convention to use for non-kernel functions.80 LLVM::CConv nonKernelCallingConvention = LLVM::CConv::C;81 82 /// Whether to encode workgroup attributions as additional arguments instead83 /// of a global variable.84 bool encodeWorkgroupAttributionsAsArguments = false;85};86 87struct GPUFuncOpLowering : ConvertOpToLLVMPattern<gpu::GPUFuncOp> {88 GPUFuncOpLowering(const LLVMTypeConverter &converter,89 const GPUFuncOpLoweringOptions &options,90 PatternBenefit benefit = 1)91 : ConvertOpToLLVMPattern<gpu::GPUFuncOp>(converter, benefit),92 allocaAddrSpace(options.allocaAddrSpace),93 workgroupAddrSpace(options.workgroupAddrSpace),94 kernelAttributeName(options.kernelAttributeName),95 kernelBlockSizeAttributeName(options.kernelBlockSizeAttributeName),96 kernelCallingConvention(options.kernelCallingConvention),97 nonKernelCallingConvention(options.nonKernelCallingConvention),98 encodeWorkgroupAttributionsAsArguments(99 options.encodeWorkgroupAttributionsAsArguments) {}100 101 LogicalResult102 matchAndRewrite(gpu::GPUFuncOp gpuFuncOp, OpAdaptor adaptor,103 ConversionPatternRewriter &rewriter) const override;104 105private:106 /// The address space to use for `alloca`s in private memory.107 unsigned allocaAddrSpace;108 /// The address space to use declaring workgroup memory.109 unsigned workgroupAddrSpace;110 111 /// The attribute name to use instead of `gpu.kernel`. Null if no attribute112 /// should be used.113 StringAttr kernelAttributeName;114 /// The attribute name to to set block size. Null if no attribute should be115 /// used.116 StringAttr kernelBlockSizeAttributeName;117 118 /// The calling convention to use for kernel functions119 LLVM::CConv kernelCallingConvention;120 /// The calling convention to use for non-kernel functions121 LLVM::CConv nonKernelCallingConvention;122 123 /// Whether to encode workgroup attributions as additional arguments instead124 /// of a global variable.125 bool encodeWorkgroupAttributionsAsArguments;126};127 128/// The lowering of gpu.printf to a call to HIP hostcalls129///130/// Simplifies llvm/lib/Transforms/Utils/AMDGPUEmitPrintf.cpp, as we don't have131/// to deal with %s (even if there were first-class strings in MLIR, they're not132/// legal input to gpu.printf) or non-constant format strings133struct GPUPrintfOpToHIPLowering : public ConvertOpToLLVMPattern<gpu::PrintfOp> {134 using ConvertOpToLLVMPattern<gpu::PrintfOp>::ConvertOpToLLVMPattern;135 136 LogicalResult137 matchAndRewrite(gpu::PrintfOp gpuPrintfOp, gpu::PrintfOpAdaptor adaptor,138 ConversionPatternRewriter &rewriter) const override;139};140 141/// The lowering of gpu.printf to a call to an external printf() function142///143/// This pass will add a declaration of printf() to the GPUModule if needed144/// and separate out the format strings into global constants. For some145/// runtimes, such as OpenCL on AMD, this is sufficient setup, as the compiler146/// will lower printf calls to appropriate device-side code.147/// However not all backends use the same calling convention and function148/// naming.149/// For example, the LLVM SPIRV backend requires calling convention150/// LLVM::cconv::CConv::SPIR_FUNC and function name needs to be151/// mangled as "_Z6printfPU3AS2Kcz".152/// Default callingConvention is LLVM::cconv::CConv::C and153/// funcName is "printf" but they can be customized as needed.154struct GPUPrintfOpToLLVMCallLowering155 : public ConvertOpToLLVMPattern<gpu::PrintfOp> {156 GPUPrintfOpToLLVMCallLowering(157 const LLVMTypeConverter &converter, int addressSpace = 0,158 LLVM::cconv::CConv callingConvention = LLVM::cconv::CConv::C,159 StringRef funcName = "printf")160 : ConvertOpToLLVMPattern<gpu::PrintfOp>(converter),161 addressSpace(addressSpace), callingConvention(callingConvention),162 funcName(funcName) {}163 164 LogicalResult165 matchAndRewrite(gpu::PrintfOp gpuPrintfOp, gpu::PrintfOpAdaptor adaptor,166 ConversionPatternRewriter &rewriter) const override;167 168private:169 int addressSpace;170 LLVM::cconv::CConv callingConvention;171 StringRef funcName;172};173 174/// Lowering of gpu.printf to a vprintf standard library.175struct GPUPrintfOpToVPrintfLowering176 : public ConvertOpToLLVMPattern<gpu::PrintfOp> {177 using ConvertOpToLLVMPattern<gpu::PrintfOp>::ConvertOpToLLVMPattern;178 179 LogicalResult180 matchAndRewrite(gpu::PrintfOp gpuPrintfOp, gpu::PrintfOpAdaptor adaptor,181 ConversionPatternRewriter &rewriter) const override;182};183 184struct GPUReturnOpLowering : public ConvertOpToLLVMPattern<gpu::ReturnOp> {185 using ConvertOpToLLVMPattern<gpu::ReturnOp>::ConvertOpToLLVMPattern;186 187 LogicalResult188 matchAndRewrite(gpu::ReturnOp op, OpAdaptor adaptor,189 ConversionPatternRewriter &rewriter) const override;190};191 192namespace impl {193/// Unrolls op to array/vector elements.194LogicalResult scalarizeVectorOp(Operation *op, ValueRange operands,195 ConversionPatternRewriter &rewriter,196 const LLVMTypeConverter &converter);197} // namespace impl198 199/// Unrolls SourceOp to array/vector elements.200template <typename SourceOp>201struct ScalarizeVectorOpLowering : public ConvertOpToLLVMPattern<SourceOp> {202public:203 using ConvertOpToLLVMPattern<SourceOp>::ConvertOpToLLVMPattern;204 205 LogicalResult206 matchAndRewrite(SourceOp op, typename SourceOp::Adaptor adaptor,207 ConversionPatternRewriter &rewriter) const override {208 return impl::scalarizeVectorOp(op, adaptor.getOperands(), rewriter,209 *this->getTypeConverter());210 }211};212 213} // namespace mlir214 215#endif // MLIR_CONVERSION_GPUCOMMON_GPUOPSLOWERING_H_216