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 = {/*useBarePtrCallConv =*/false, 112 /*emitCWrappers =*/true, 113 /*indexBitwidth =*/indexBitwidth, 114 /*useAlignedAlloc =*/false}; 115 116 /// MemRef conversion for GPU to NVVM lowering. The GPU dialect uses memory 117 /// space 5 for private memory attributions, but NVVM represents private 118 /// memory allocations as local `alloca`s in the default address space. This 119 /// converter drops the private memory space to support the use case above. 120 LLVMTypeConverter converter(m.getContext(), options); 121 converter.addConversion([&](MemRefType type) -> Optional<Type> { 122 if (type.getMemorySpaceAsInt() != 123 gpu::GPUDialect::getPrivateAddressSpace()) 124 return llvm::None; 125 return converter.convertType(MemRefType::Builder(type).setMemorySpace(0)); 126 }); 127 128 RewritePatternSet patterns(m.getContext()); 129 RewritePatternSet llvmPatterns(m.getContext()); 130 131 // Apply in-dialect lowering first. In-dialect lowering will replace ops 132 // which need to be lowered further, which is not supported by a single 133 // conversion pass. 134 populateGpuRewritePatterns(patterns); 135 (void)applyPatternsAndFoldGreedily(m, std::move(patterns)); 136 137 populateStdToLLVMConversionPatterns(converter, llvmPatterns); 138 populateGpuToNVVMConversionPatterns(converter, llvmPatterns); 139 LLVMConversionTarget target(getContext()); 140 configureGpuToNVVMConversionLegality(target); 141 if (failed(applyPartialConversion(m, target, std::move(llvmPatterns)))) 142 signalPassFailure(); 143 } 144 }; 145 146 } // anonymous namespace 147 148 void mlir::configureGpuToNVVMConversionLegality(ConversionTarget &target) { 149 target.addIllegalOp<FuncOp>(); 150 target.addLegalDialect<::mlir::LLVM::LLVMDialect>(); 151 target.addLegalDialect<::mlir::NVVM::NVVMDialect>(); 152 target.addIllegalDialect<gpu::GPUDialect>(); 153 target.addIllegalOp<LLVM::CosOp, LLVM::ExpOp, LLVM::FAbsOp, LLVM::FCeilOp, 154 LLVM::FFloorOp, LLVM::LogOp, LLVM::Log10Op, LLVM::Log2Op, 155 LLVM::PowOp, LLVM::SinOp, LLVM::SqrtOp>(); 156 157 // TODO: Remove once we support replacing non-root ops. 158 target.addLegalOp<gpu::YieldOp, gpu::GPUModuleOp, gpu::ModuleEndOp>(); 159 } 160 161 void mlir::populateGpuToNVVMConversionPatterns(LLVMTypeConverter &converter, 162 RewritePatternSet &patterns) { 163 populateWithGenerated(patterns); 164 patterns 165 .add<GPUIndexIntrinsicOpLowering<gpu::ThreadIdOp, NVVM::ThreadIdXOp, 166 NVVM::ThreadIdYOp, NVVM::ThreadIdZOp>, 167 GPUIndexIntrinsicOpLowering<gpu::BlockDimOp, NVVM::BlockDimXOp, 168 NVVM::BlockDimYOp, NVVM::BlockDimZOp>, 169 GPUIndexIntrinsicOpLowering<gpu::BlockIdOp, NVVM::BlockIdXOp, 170 NVVM::BlockIdYOp, NVVM::BlockIdZOp>, 171 GPUIndexIntrinsicOpLowering<gpu::GridDimOp, NVVM::GridDimXOp, 172 NVVM::GridDimYOp, NVVM::GridDimZOp>, 173 GPUShuffleOpLowering, GPUReturnOpLowering>(converter); 174 175 // Explicitly drop memory space when lowering private memory 176 // attributions since NVVM models it as `alloca`s in the default 177 // memory space and does not support `alloca`s with addrspace(5). 178 patterns.add<GPUFuncOpLowering>( 179 converter, /*allocaAddrSpace=*/0, 180 Identifier::get(NVVM::NVVMDialect::getKernelFuncAttrName(), 181 &converter.getContext())); 182 183 patterns.add<OpToFuncCallLowering<AbsFOp>>(converter, "__nv_fabsf", 184 "__nv_fabs"); 185 patterns.add<OpToFuncCallLowering<math::AtanOp>>(converter, "__nv_atanf", 186 "__nv_atan"); 187 patterns.add<OpToFuncCallLowering<math::Atan2Op>>(converter, "__nv_atan2f", 188 "__nv_atan2"); 189 patterns.add<OpToFuncCallLowering<CeilFOp>>(converter, "__nv_ceilf", 190 "__nv_ceil"); 191 patterns.add<OpToFuncCallLowering<math::CosOp>>(converter, "__nv_cosf", 192 "__nv_cos"); 193 patterns.add<OpToFuncCallLowering<math::ExpOp>>(converter, "__nv_expf", 194 "__nv_exp"); 195 patterns.add<OpToFuncCallLowering<math::ExpM1Op>>(converter, "__nv_expm1f", 196 "__nv_expm1"); 197 patterns.add<OpToFuncCallLowering<FloorFOp>>(converter, "__nv_floorf", 198 "__nv_floor"); 199 patterns.add<OpToFuncCallLowering<math::LogOp>>(converter, "__nv_logf", 200 "__nv_log"); 201 patterns.add<OpToFuncCallLowering<math::Log1pOp>>(converter, "__nv_log1pf", 202 "__nv_log1p"); 203 patterns.add<OpToFuncCallLowering<math::Log10Op>>(converter, "__nv_log10f", 204 "__nv_log10"); 205 patterns.add<OpToFuncCallLowering<math::Log2Op>>(converter, "__nv_log2f", 206 "__nv_log2"); 207 patterns.add<OpToFuncCallLowering<math::PowFOp>>(converter, "__nv_powf", 208 "__nv_pow"); 209 patterns.add<OpToFuncCallLowering<math::RsqrtOp>>(converter, "__nv_rsqrtf", 210 "__nv_rsqrt"); 211 patterns.add<OpToFuncCallLowering<math::SinOp>>(converter, "__nv_sinf", 212 "__nv_sin"); 213 patterns.add<OpToFuncCallLowering<math::SqrtOp>>(converter, "__nv_sqrtf", 214 "__nv_sqrt"); 215 patterns.add<OpToFuncCallLowering<math::TanhOp>>(converter, "__nv_tanhf", 216 "__nv_tanh"); 217 } 218 219 std::unique_ptr<OperationPass<gpu::GPUModuleOp>> 220 mlir::createLowerGpuOpsToNVVMOpsPass(unsigned indexBitwidth) { 221 return std::make_unique<LowerGpuOpsToNVVMOpsPass>(indexBitwidth); 222 } 223