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/Dialect/Math/IR/Math.h" 21 #include "mlir/Dialect/MemRef/IR/MemRef.h" 22 #include "mlir/IR/BlockAndValueMapping.h" 23 #include "mlir/Transforms/DialectConversion.h" 24 #include "mlir/Transforms/GreedyPatternRewriteDriver.h" 25 #include "llvm/Support/FormatVariadic.h" 26 27 #include "../GPUCommon/GPUOpsLowering.h" 28 #include "../GPUCommon/IndexIntrinsicsOpLowering.h" 29 #include "../GPUCommon/OpToFuncCallLowering.h" 30 #include "../PassDetail.h" 31 32 using namespace mlir; 33 34 namespace { 35 36 struct GPUShuffleOpLowering : public ConvertOpToLLVMPattern<gpu::ShuffleOp> { 37 using ConvertOpToLLVMPattern<gpu::ShuffleOp>::ConvertOpToLLVMPattern; 38 39 /// Lowers a shuffle to the corresponding NVVM op. 40 /// 41 /// Convert the `width` argument into an activeMask (a bitmask which specifies 42 /// which threads participate in the shuffle) and a maskAndClamp (specifying 43 /// the highest lane which participates in the shuffle). 44 /// 45 /// %one = llvm.constant(1 : i32) : i32 46 /// %shl = llvm.shl %one, %width : i32 47 /// %active_mask = llvm.sub %shl, %one : i32 48 /// %mask_and_clamp = llvm.sub %width, %one : i32 49 /// %shfl = nvvm.shfl.sync.bfly %active_mask, %value, %offset, 50 /// %mask_and_clamp : !llvm<"{ float, i1 }"> 51 /// %shfl_value = llvm.extractvalue %shfl[0 : index] : 52 /// !llvm<"{ float, i1 }"> 53 /// %shfl_pred = llvm.extractvalue %shfl[1 : index] : 54 /// !llvm<"{ float, i1 }"> 55 LogicalResult 56 matchAndRewrite(gpu::ShuffleOp op, ArrayRef<Value> operands, 57 ConversionPatternRewriter &rewriter) const override { 58 Location loc = op->getLoc(); 59 gpu::ShuffleOpAdaptor adaptor(operands); 60 61 auto valueTy = adaptor.value().getType(); 62 auto int32Type = IntegerType::get(rewriter.getContext(), 32); 63 auto predTy = IntegerType::get(rewriter.getContext(), 1); 64 auto resultTy = LLVM::LLVMStructType::getLiteral(rewriter.getContext(), 65 {valueTy, predTy}); 66 67 Value one = rewriter.create<LLVM::ConstantOp>( 68 loc, int32Type, rewriter.getI32IntegerAttr(1)); 69 // Bit mask of active lanes: `(1 << activeWidth) - 1`. 70 Value activeMask = rewriter.create<LLVM::SubOp>( 71 loc, int32Type, 72 rewriter.create<LLVM::ShlOp>(loc, int32Type, one, adaptor.width()), 73 one); 74 // Clamp lane: `activeWidth - 1` 75 Value maskAndClamp = 76 rewriter.create<LLVM::SubOp>(loc, int32Type, adaptor.width(), one); 77 78 auto returnValueAndIsValidAttr = rewriter.getUnitAttr(); 79 Value shfl = rewriter.create<NVVM::ShflBflyOp>( 80 loc, resultTy, activeMask, adaptor.value(), adaptor.offset(), 81 maskAndClamp, returnValueAndIsValidAttr); 82 Value shflValue = rewriter.create<LLVM::ExtractValueOp>( 83 loc, valueTy, shfl, rewriter.getIndexArrayAttr(0)); 84 Value isActiveSrcLane = rewriter.create<LLVM::ExtractValueOp>( 85 loc, predTy, shfl, rewriter.getIndexArrayAttr(1)); 86 87 rewriter.replaceOp(op, {shflValue, isActiveSrcLane}); 88 return success(); 89 } 90 }; 91 92 /// Import the GPU Ops to NVVM Patterns. 93 #include "GPUToNVVM.cpp.inc" 94 95 /// A pass that replaces all occurrences of GPU device operations with their 96 /// corresponding NVVM equivalent. 97 /// 98 /// This pass only handles device code and is not meant to be run on GPU host 99 /// code. 100 struct LowerGpuOpsToNVVMOpsPass 101 : public ConvertGpuOpsToNVVMOpsBase<LowerGpuOpsToNVVMOpsPass> { 102 LowerGpuOpsToNVVMOpsPass() = default; 103 LowerGpuOpsToNVVMOpsPass(unsigned indexBitwidth) { 104 this->indexBitwidth = indexBitwidth; 105 } 106 107 void runOnOperation() override { 108 gpu::GPUModuleOp m = getOperation(); 109 110 /// Customize the bitwidth used for the device side index computations. 111 LowerToLLVMOptions options( 112 m.getContext(), 113 DataLayout(cast<DataLayoutOpInterface>(m.getOperation()))); 114 options.emitCWrappers = true; 115 if (indexBitwidth != kDeriveIndexBitwidthFromDataLayout) 116 options.overrideIndexBitwidth(indexBitwidth); 117 118 /// MemRef conversion for GPU to NVVM lowering. The GPU dialect uses memory 119 /// space 5 for private memory attributions, but NVVM represents private 120 /// memory allocations as local `alloca`s in the default address space. This 121 /// converter drops the private memory space to support the use case above. 122 LLVMTypeConverter converter(m.getContext(), options); 123 converter.addConversion([&](MemRefType type) -> Optional<Type> { 124 if (type.getMemorySpaceAsInt() != 125 gpu::GPUDialect::getPrivateAddressSpace()) 126 return llvm::None; 127 return converter.convertType(MemRefType::Builder(type).setMemorySpace(0)); 128 }); 129 130 RewritePatternSet patterns(m.getContext()); 131 RewritePatternSet llvmPatterns(m.getContext()); 132 133 // Apply in-dialect lowering first. In-dialect lowering will replace ops 134 // which need to be lowered further, which is not supported by a single 135 // conversion pass. 136 populateGpuRewritePatterns(patterns); 137 (void)applyPatternsAndFoldGreedily(m, std::move(patterns)); 138 139 populateStdToLLVMConversionPatterns(converter, llvmPatterns); 140 populateGpuToNVVMConversionPatterns(converter, llvmPatterns); 141 LLVMConversionTarget target(getContext()); 142 configureGpuToNVVMConversionLegality(target); 143 if (failed(applyPartialConversion(m, target, std::move(llvmPatterns)))) 144 signalPassFailure(); 145 } 146 }; 147 148 } // anonymous namespace 149 150 void mlir::configureGpuToNVVMConversionLegality(ConversionTarget &target) { 151 target.addIllegalOp<FuncOp>(); 152 target.addLegalDialect<::mlir::LLVM::LLVMDialect>(); 153 target.addLegalDialect<::mlir::NVVM::NVVMDialect>(); 154 target.addIllegalDialect<gpu::GPUDialect>(); 155 target.addIllegalOp<LLVM::CosOp, LLVM::ExpOp, LLVM::FAbsOp, LLVM::FCeilOp, 156 LLVM::FFloorOp, LLVM::LogOp, LLVM::Log10Op, LLVM::Log2Op, 157 LLVM::PowOp, LLVM::SinOp, LLVM::SqrtOp>(); 158 159 // TODO: Remove once we support replacing non-root ops. 160 target.addLegalOp<gpu::YieldOp, gpu::GPUModuleOp, gpu::ModuleEndOp>(); 161 } 162 163 void mlir::populateGpuToNVVMConversionPatterns(LLVMTypeConverter &converter, 164 RewritePatternSet &patterns) { 165 populateWithGenerated(patterns); 166 patterns 167 .add<GPUIndexIntrinsicOpLowering<gpu::ThreadIdOp, NVVM::ThreadIdXOp, 168 NVVM::ThreadIdYOp, NVVM::ThreadIdZOp>, 169 GPUIndexIntrinsicOpLowering<gpu::BlockDimOp, NVVM::BlockDimXOp, 170 NVVM::BlockDimYOp, NVVM::BlockDimZOp>, 171 GPUIndexIntrinsicOpLowering<gpu::BlockIdOp, NVVM::BlockIdXOp, 172 NVVM::BlockIdYOp, NVVM::BlockIdZOp>, 173 GPUIndexIntrinsicOpLowering<gpu::GridDimOp, NVVM::GridDimXOp, 174 NVVM::GridDimYOp, NVVM::GridDimZOp>, 175 GPUShuffleOpLowering, GPUReturnOpLowering>(converter); 176 177 // Explicitly drop memory space when lowering private memory 178 // attributions since NVVM models it as `alloca`s in the default 179 // memory space and does not support `alloca`s with addrspace(5). 180 patterns.add<GPUFuncOpLowering>( 181 converter, /*allocaAddrSpace=*/0, 182 Identifier::get(NVVM::NVVMDialect::getKernelFuncAttrName(), 183 &converter.getContext())); 184 185 patterns.add<OpToFuncCallLowering<AbsFOp>>(converter, "__nv_fabsf", 186 "__nv_fabs"); 187 patterns.add<OpToFuncCallLowering<math::AtanOp>>(converter, "__nv_atanf", 188 "__nv_atan"); 189 patterns.add<OpToFuncCallLowering<math::Atan2Op>>(converter, "__nv_atan2f", 190 "__nv_atan2"); 191 patterns.add<OpToFuncCallLowering<CeilFOp>>(converter, "__nv_ceilf", 192 "__nv_ceil"); 193 patterns.add<OpToFuncCallLowering<math::CosOp>>(converter, "__nv_cosf", 194 "__nv_cos"); 195 patterns.add<OpToFuncCallLowering<math::ExpOp>>(converter, "__nv_expf", 196 "__nv_exp"); 197 patterns.add<OpToFuncCallLowering<math::ExpM1Op>>(converter, "__nv_expm1f", 198 "__nv_expm1"); 199 patterns.add<OpToFuncCallLowering<FloorFOp>>(converter, "__nv_floorf", 200 "__nv_floor"); 201 patterns.add<OpToFuncCallLowering<math::LogOp>>(converter, "__nv_logf", 202 "__nv_log"); 203 patterns.add<OpToFuncCallLowering<math::Log1pOp>>(converter, "__nv_log1pf", 204 "__nv_log1p"); 205 patterns.add<OpToFuncCallLowering<math::Log10Op>>(converter, "__nv_log10f", 206 "__nv_log10"); 207 patterns.add<OpToFuncCallLowering<math::Log2Op>>(converter, "__nv_log2f", 208 "__nv_log2"); 209 patterns.add<OpToFuncCallLowering<math::PowFOp>>(converter, "__nv_powf", 210 "__nv_pow"); 211 patterns.add<OpToFuncCallLowering<math::RsqrtOp>>(converter, "__nv_rsqrtf", 212 "__nv_rsqrt"); 213 patterns.add<OpToFuncCallLowering<math::SinOp>>(converter, "__nv_sinf", 214 "__nv_sin"); 215 patterns.add<OpToFuncCallLowering<math::SqrtOp>>(converter, "__nv_sqrtf", 216 "__nv_sqrt"); 217 patterns.add<OpToFuncCallLowering<math::TanhOp>>(converter, "__nv_tanhf", 218 "__nv_tanh"); 219 } 220 221 std::unique_ptr<OperationPass<gpu::GPUModuleOp>> 222 mlir::createLowerGpuOpsToNVVMOpsPass(unsigned indexBitwidth) { 223 return std::make_unique<LowerGpuOpsToNVVMOpsPass>(indexBitwidth); 224 } 225