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