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