1 //===- LowerGpuOpsToNVVMOps.cpp - MLIR GPU to NVVM lowering passes --------===// 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 generate NVVMIR operations for higher-level 10 // GPU operations. 11 // 12 //===----------------------------------------------------------------------===// 13 14 #include "mlir/Conversion/GPUToNVVM/GPUToNVVMPass.h" 15 16 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h" 17 #include "mlir/Dialect/GPU/GPUDialect.h" 18 #include "mlir/Dialect/GPU/Passes.h" 19 #include "mlir/Dialect/LLVMIR/NVVMDialect.h" 20 #include "mlir/IR/BlockAndValueMapping.h" 21 #include "mlir/Pass/Pass.h" 22 #include "mlir/Transforms/DialectConversion.h" 23 #include "llvm/Support/FormatVariadic.h" 24 25 #include "../GPUCommon/IndexIntrinsicsOpLowering.h" 26 #include "../GPUCommon/OpToFuncCallLowering.h" 27 28 using namespace mlir; 29 30 namespace { 31 32 33 struct GPUShuffleOpLowering : public ConvertToLLVMPattern { 34 explicit GPUShuffleOpLowering(LLVMTypeConverter &lowering_) 35 : ConvertToLLVMPattern(gpu::ShuffleOp::getOperationName(), 36 lowering_.getDialect()->getContext(), lowering_) {} 37 38 /// Lowers a shuffle to the corresponding NVVM op. 39 /// 40 /// Convert the `width` argument into an activeMask (a bitmask which specifies 41 /// which threads participate in the shuffle) and a maskAndClamp (specifying 42 /// the highest lane which participates in the shuffle). 43 /// 44 /// %one = llvm.constant(1 : i32) : !llvm.i32 45 /// %shl = llvm.shl %one, %width : !llvm.i32 46 /// %active_mask = llvm.sub %shl, %one : !llvm.i32 47 /// %mask_and_clamp = llvm.sub %width, %one : !llvm.i32 48 /// %shfl = nvvm.shfl.sync.bfly %active_mask, %value, %offset, 49 /// %mask_and_clamp : !llvm<"{ float, i1 }"> 50 /// %shfl_value = llvm.extractvalue %shfl[0 : index] : 51 /// !llvm<"{ float, i1 }"> 52 /// %shfl_pred = llvm.extractvalue %shfl[1 : index] : 53 /// !llvm<"{ float, i1 }"> 54 LogicalResult 55 matchAndRewrite(Operation *op, ArrayRef<Value> operands, 56 ConversionPatternRewriter &rewriter) const override { 57 Location loc = op->getLoc(); 58 gpu::ShuffleOpOperandAdaptor adaptor(operands); 59 60 auto dialect = typeConverter.getDialect(); 61 auto valueTy = adaptor.value().getType().cast<LLVM::LLVMType>(); 62 auto int32Type = LLVM::LLVMType::getInt32Ty(dialect); 63 auto predTy = LLVM::LLVMType::getInt1Ty(dialect); 64 auto resultTy = LLVM::LLVMType::getStructTy(dialect, {valueTy, predTy}); 65 66 Value one = rewriter.create<LLVM::ConstantOp>( 67 loc, int32Type, rewriter.getI32IntegerAttr(1)); 68 // Bit mask of active lanes: `(1 << activeWidth) - 1`. 69 Value activeMask = rewriter.create<LLVM::SubOp>( 70 loc, int32Type, 71 rewriter.create<LLVM::ShlOp>(loc, int32Type, one, adaptor.width()), 72 one); 73 // Clamp lane: `activeWidth - 1` 74 Value maskAndClamp = 75 rewriter.create<LLVM::SubOp>(loc, int32Type, adaptor.width(), one); 76 77 auto returnValueAndIsValidAttr = rewriter.getUnitAttr(); 78 Value shfl = rewriter.create<NVVM::ShflBflyOp>( 79 loc, resultTy, activeMask, adaptor.value(), adaptor.offset(), 80 maskAndClamp, returnValueAndIsValidAttr); 81 Value shflValue = rewriter.create<LLVM::ExtractValueOp>( 82 loc, valueTy, shfl, rewriter.getIndexArrayAttr(0)); 83 Value isActiveSrcLane = rewriter.create<LLVM::ExtractValueOp>( 84 loc, predTy, shfl, rewriter.getIndexArrayAttr(1)); 85 86 rewriter.replaceOp(op, {shflValue, isActiveSrcLane}); 87 return success(); 88 } 89 }; 90 91 struct GPUFuncOpLowering : ConvertToLLVMPattern { 92 explicit GPUFuncOpLowering(LLVMTypeConverter &typeConverter) 93 : ConvertToLLVMPattern(gpu::GPUFuncOp::getOperationName(), 94 typeConverter.getDialect()->getContext(), 95 typeConverter) {} 96 97 LogicalResult 98 matchAndRewrite(Operation *op, ArrayRef<Value> operands, 99 ConversionPatternRewriter &rewriter) const override { 100 assert(operands.empty() && "func op is not expected to have operands"); 101 auto gpuFuncOp = cast<gpu::GPUFuncOp>(op); 102 Location loc = gpuFuncOp.getLoc(); 103 104 SmallVector<LLVM::GlobalOp, 3> workgroupBuffers; 105 workgroupBuffers.reserve(gpuFuncOp.getNumWorkgroupAttributions()); 106 for (auto en : llvm::enumerate(gpuFuncOp.getWorkgroupAttributions())) { 107 Value attribution = en.value(); 108 109 auto type = attribution.getType().dyn_cast<MemRefType>(); 110 assert(type && type.hasStaticShape() && "unexpected type in attribution"); 111 112 uint64_t numElements = type.getNumElements(); 113 114 auto elementType = typeConverter.convertType(type.getElementType()) 115 .cast<LLVM::LLVMType>(); 116 auto arrayType = LLVM::LLVMType::getArrayTy(elementType, numElements); 117 std::string name = std::string( 118 llvm::formatv("__wg_{0}_{1}", gpuFuncOp.getName(), en.index())); 119 auto globalOp = rewriter.create<LLVM::GlobalOp>( 120 gpuFuncOp.getLoc(), arrayType, /*isConstant=*/false, 121 LLVM::Linkage::Internal, name, /*value=*/Attribute(), 122 gpu::GPUDialect::getWorkgroupAddressSpace()); 123 workgroupBuffers.push_back(globalOp); 124 } 125 126 // Rewrite the original GPU function to an LLVM function. 127 auto funcType = typeConverter.convertType(gpuFuncOp.getType()) 128 .cast<LLVM::LLVMType>() 129 .getPointerElementTy(); 130 131 // Remap proper input types. 132 TypeConverter::SignatureConversion signatureConversion( 133 gpuFuncOp.front().getNumArguments()); 134 typeConverter.convertFunctionSignature( 135 gpuFuncOp.getType(), /*isVariadic=*/false, signatureConversion); 136 137 // Create the new function operation. Only copy those attributes that are 138 // not specific to function modeling. 139 SmallVector<NamedAttribute, 4> attributes; 140 for (const auto &attr : gpuFuncOp.getAttrs()) { 141 if (attr.first.is(SymbolTable::getSymbolAttrName()) || 142 attr.first.is(impl::getTypeAttrName()) || 143 attr.first.is(gpu::GPUFuncOp::getNumWorkgroupAttributionsAttrName())) 144 continue; 145 attributes.push_back(attr); 146 } 147 auto llvmFuncOp = rewriter.create<LLVM::LLVMFuncOp>( 148 gpuFuncOp.getLoc(), gpuFuncOp.getName(), funcType, 149 LLVM::Linkage::External, attributes); 150 151 { 152 // Insert operations that correspond to converted workgroup and private 153 // memory attributions to the body of the function. This must operate on 154 // the original function, before the body region is inlined in the new 155 // function to maintain the relation between block arguments and the 156 // parent operation that assigns their semantics. 157 OpBuilder::InsertionGuard guard(rewriter); 158 159 // Rewrite workgroup memory attributions to addresses of global buffers. 160 rewriter.setInsertionPointToStart(&gpuFuncOp.front()); 161 unsigned numProperArguments = gpuFuncOp.getNumArguments(); 162 auto i32Type = LLVM::LLVMType::getInt32Ty(typeConverter.getDialect()); 163 164 Value zero = nullptr; 165 if (!workgroupBuffers.empty()) 166 zero = rewriter.create<LLVM::ConstantOp>(loc, i32Type, 167 rewriter.getI32IntegerAttr(0)); 168 for (auto en : llvm::enumerate(workgroupBuffers)) { 169 LLVM::GlobalOp global = en.value(); 170 Value address = rewriter.create<LLVM::AddressOfOp>(loc, global); 171 auto elementType = global.getType().getArrayElementType(); 172 Value memory = rewriter.create<LLVM::GEPOp>( 173 loc, elementType.getPointerTo(global.addr_space().getZExtValue()), 174 address, ArrayRef<Value>{zero, zero}); 175 176 // Build a memref descriptor pointing to the buffer to plug with the 177 // existing memref infrastructure. This may use more registers than 178 // otherwise necessary given that memref sizes are fixed, but we can try 179 // and canonicalize that away later. 180 Value attribution = gpuFuncOp.getWorkgroupAttributions()[en.index()]; 181 auto type = attribution.getType().cast<MemRefType>(); 182 auto descr = MemRefDescriptor::fromStaticShape( 183 rewriter, loc, typeConverter, type, memory); 184 signatureConversion.remapInput(numProperArguments + en.index(), descr); 185 } 186 187 // Rewrite private memory attributions to alloca'ed buffers. 188 unsigned numWorkgroupAttributions = 189 gpuFuncOp.getNumWorkgroupAttributions(); 190 auto int64Ty = LLVM::LLVMType::getInt64Ty(typeConverter.getDialect()); 191 for (auto en : llvm::enumerate(gpuFuncOp.getPrivateAttributions())) { 192 Value attribution = en.value(); 193 auto type = attribution.getType().cast<MemRefType>(); 194 assert(type && type.hasStaticShape() && 195 "unexpected type in attribution"); 196 197 // Explicitly drop memory space when lowering private memory 198 // attributions since NVVM models it as `alloca`s in the default 199 // memory space and does not support `alloca`s with addrspace(5). 200 auto ptrType = typeConverter.convertType(type.getElementType()) 201 .cast<LLVM::LLVMType>() 202 .getPointerTo(); 203 Value numElements = rewriter.create<LLVM::ConstantOp>( 204 gpuFuncOp.getLoc(), int64Ty, 205 rewriter.getI64IntegerAttr(type.getNumElements())); 206 Value allocated = rewriter.create<LLVM::AllocaOp>( 207 gpuFuncOp.getLoc(), ptrType, numElements, /*alignment=*/0); 208 auto descr = MemRefDescriptor::fromStaticShape( 209 rewriter, loc, typeConverter, type, allocated); 210 signatureConversion.remapInput( 211 numProperArguments + numWorkgroupAttributions + en.index(), descr); 212 } 213 } 214 215 // Move the region to the new function, update the entry block signature. 216 rewriter.inlineRegionBefore(gpuFuncOp.getBody(), llvmFuncOp.getBody(), 217 llvmFuncOp.end()); 218 rewriter.applySignatureConversion(&llvmFuncOp.getBody(), 219 signatureConversion); 220 221 rewriter.eraseOp(gpuFuncOp); 222 return success(); 223 } 224 }; 225 226 struct GPUReturnOpLowering : public ConvertToLLVMPattern { 227 GPUReturnOpLowering(LLVMTypeConverter &typeConverter) 228 : ConvertToLLVMPattern(gpu::ReturnOp::getOperationName(), 229 typeConverter.getDialect()->getContext(), 230 typeConverter) {} 231 232 LogicalResult 233 matchAndRewrite(Operation *op, ArrayRef<Value> operands, 234 ConversionPatternRewriter &rewriter) const override { 235 rewriter.replaceOpWithNewOp<LLVM::ReturnOp>(op, operands); 236 return success(); 237 } 238 }; 239 240 /// Import the GPU Ops to NVVM Patterns. 241 #include "GPUToNVVM.cpp.inc" 242 243 /// A pass that replaces all occurrences of GPU device operations with their 244 /// corresponding NVVM equivalent. 245 /// 246 /// This pass only handles device code and is not meant to be run on GPU host 247 /// code. 248 class LowerGpuOpsToNVVMOpsPass 249 : public OperationPass<LowerGpuOpsToNVVMOpsPass, gpu::GPUModuleOp> { 250 public: 251 /// Include the generated pass utilities. 252 #define GEN_PASS_ConvertGpuOpsToNVVMOps 253 #include "mlir/Conversion/Passes.h.inc" 254 255 void runOnOperation() override { 256 gpu::GPUModuleOp m = getOperation(); 257 258 /// MemRef conversion for GPU to NVVM lowering. The GPU dialect uses memory 259 /// space 5 for private memory attributions, but NVVM represents private 260 /// memory allocations as local `alloca`s in the default address space. This 261 /// converter drops the private memory space to support the use case above. 262 LLVMTypeConverter converter(m.getContext()); 263 converter.addConversion([&](MemRefType type) -> Optional<Type> { 264 if (type.getMemorySpace() != gpu::GPUDialect::getPrivateAddressSpace()) 265 return llvm::None; 266 return converter.convertType(MemRefType::Builder(type).setMemorySpace(0)); 267 }); 268 269 OwningRewritePatternList patterns; 270 271 // Apply in-dialect lowering first. In-dialect lowering will replace ops 272 // which need to be lowered further, which is not supported by a single 273 // conversion pass. 274 populateGpuRewritePatterns(m.getContext(), patterns); 275 applyPatternsGreedily(m, patterns); 276 patterns.clear(); 277 278 populateStdToLLVMConversionPatterns(converter, patterns); 279 populateGpuToNVVMConversionPatterns(converter, patterns); 280 LLVMConversionTarget target(getContext()); 281 target.addIllegalDialect<gpu::GPUDialect>(); 282 target.addIllegalOp<LLVM::CosOp, LLVM::ExpOp, LLVM::FAbsOp, LLVM::FCeilOp, 283 LLVM::LogOp, LLVM::Log10Op, LLVM::Log2Op>(); 284 target.addIllegalOp<FuncOp>(); 285 target.addLegalDialect<NVVM::NVVMDialect>(); 286 // TODO(csigg): Remove once we support replacing non-root ops. 287 target.addLegalOp<gpu::YieldOp, gpu::GPUModuleOp, gpu::ModuleEndOp>(); 288 if (failed(applyPartialConversion(m, target, patterns, &converter))) 289 signalPassFailure(); 290 } 291 }; 292 293 } // anonymous namespace 294 295 void mlir::populateGpuToNVVMConversionPatterns( 296 LLVMTypeConverter &converter, OwningRewritePatternList &patterns) { 297 populateWithGenerated(converter.getDialect()->getContext(), &patterns); 298 patterns 299 .insert<GPUIndexIntrinsicOpLowering<gpu::ThreadIdOp, NVVM::ThreadIdXOp, 300 NVVM::ThreadIdYOp, NVVM::ThreadIdZOp>, 301 GPUIndexIntrinsicOpLowering<gpu::BlockDimOp, NVVM::BlockDimXOp, 302 NVVM::BlockDimYOp, NVVM::BlockDimZOp>, 303 GPUIndexIntrinsicOpLowering<gpu::BlockIdOp, NVVM::BlockIdXOp, 304 NVVM::BlockIdYOp, NVVM::BlockIdZOp>, 305 GPUIndexIntrinsicOpLowering<gpu::GridDimOp, NVVM::GridDimXOp, 306 NVVM::GridDimYOp, NVVM::GridDimZOp>, 307 GPUShuffleOpLowering, GPUFuncOpLowering, GPUReturnOpLowering>( 308 converter); 309 patterns.insert<OpToFuncCallLowering<AbsFOp>>(converter, "__nv_fabsf", 310 "__nv_fabs"); 311 patterns.insert<OpToFuncCallLowering<CeilFOp>>(converter, "__nv_ceilf", 312 "__nv_ceil"); 313 patterns.insert<OpToFuncCallLowering<CosOp>>(converter, "__nv_cosf", 314 "__nv_cos"); 315 patterns.insert<OpToFuncCallLowering<ExpOp>>(converter, "__nv_expf", 316 "__nv_exp"); 317 patterns.insert<OpToFuncCallLowering<LogOp>>(converter, "__nv_logf", 318 "__nv_log"); 319 patterns.insert<OpToFuncCallLowering<Log10Op>>(converter, "__nv_log10f", 320 "__nv_log10"); 321 patterns.insert<OpToFuncCallLowering<Log2Op>>(converter, "__nv_log2f", 322 "__nv_log2"); 323 patterns.insert<OpToFuncCallLowering<TanhOp>>(converter, "__nv_tanhf", 324 "__nv_tanh"); 325 } 326 327 std::unique_ptr<OpPassBase<gpu::GPUModuleOp>> 328 mlir::createLowerGpuOpsToNVVMOpsPass() { 329 return std::make_unique<LowerGpuOpsToNVVMOpsPass>(); 330 } 331