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( 111 m.getContext(), 112 DataLayout(cast<DataLayoutOpInterface>(m.getOperation()))); 113 options.emitCWrappers = true; 114 if (indexBitwidth != kDeriveIndexBitwidthFromDataLayout) 115 options.overrideIndexBitwidth(indexBitwidth); 116 117 /// MemRef conversion for GPU to NVVM lowering. The GPU dialect uses memory 118 /// space 5 for private memory attributions, but NVVM represents private 119 /// memory allocations as local `alloca`s in the default address space. This 120 /// converter drops the private memory space to support the use case above. 121 LLVMTypeConverter converter(m.getContext(), options); 122 converter.addConversion([&](MemRefType type) -> Optional<Type> { 123 if (type.getMemorySpaceAsInt() != 124 gpu::GPUDialect::getPrivateAddressSpace()) 125 return llvm::None; 126 return converter.convertType(MemRefType::Builder(type).setMemorySpace(0)); 127 }); 128 129 // Lowering for MMAMatrixType. 130 converter.addConversion([&](gpu::MMAMatrixType type) -> Type { 131 // The number of items in structToReturn are dependent on the the dataType 132 // and the MMA operand that this operation is associated with. 133 llvm::DenseMap<StringRef, int64_t> numElemsPerThreadF16, 134 numElemsPerThreadF32; 135 numElemsPerThreadF16["AOp"] = 8; 136 numElemsPerThreadF16["BOp"] = 8; 137 numElemsPerThreadF16["COp"] = 4; 138 numElemsPerThreadF32["AOp"] = 8; 139 numElemsPerThreadF32["BOp"] = 8; 140 numElemsPerThreadF32["COp"] = 8; 141 Type structToReturn; 142 if (type.getElementType().isF16()) { 143 // Number of f16's in 32-bit. 144 unsigned vecSize = 2; 145 Type vec = VectorType::get(vecSize, FloatType::getF16(&getContext())); 146 unsigned size = numElemsPerThreadF16[type.getOperand()]; 147 SmallVector<Type> elements(size, vec); 148 structToReturn = 149 LLVM::LLVMStructType::getLiteral(&getContext(), elements); 150 } else if (type.getElementType().isF32()) { 151 unsigned size = numElemsPerThreadF32[type.getOperand()]; 152 SmallVector<Type> elements(size, FloatType::getF32(&getContext())); 153 structToReturn = 154 LLVM::LLVMStructType::getLiteral(&getContext(), elements); 155 } 156 return structToReturn; 157 }); 158 159 RewritePatternSet patterns(m.getContext()); 160 RewritePatternSet llvmPatterns(m.getContext()); 161 162 // Apply in-dialect lowering first. In-dialect lowering will replace ops 163 // which need to be lowered further, which is not supported by a single 164 // conversion pass. 165 populateGpuRewritePatterns(patterns); 166 (void)applyPatternsAndFoldGreedily(m, std::move(patterns)); 167 168 populateStdToLLVMConversionPatterns(converter, llvmPatterns); 169 populateGpuToNVVMConversionPatterns(converter, llvmPatterns); 170 populateGpuWMMAToNVVMConversionPatterns(converter, llvmPatterns); 171 LLVMConversionTarget target(getContext()); 172 configureGpuToNVVMConversionLegality(target); 173 if (failed(applyPartialConversion(m, target, std::move(llvmPatterns)))) 174 signalPassFailure(); 175 } 176 }; 177 178 } // anonymous namespace 179 180 void mlir::configureGpuToNVVMConversionLegality(ConversionTarget &target) { 181 target.addIllegalOp<FuncOp>(); 182 target.addLegalDialect<::mlir::LLVM::LLVMDialect>(); 183 target.addLegalDialect<::mlir::NVVM::NVVMDialect>(); 184 target.addIllegalDialect<gpu::GPUDialect>(); 185 target.addIllegalOp<LLVM::CosOp, LLVM::ExpOp, LLVM::FAbsOp, LLVM::FCeilOp, 186 LLVM::FFloorOp, LLVM::LogOp, LLVM::Log10Op, LLVM::Log2Op, 187 LLVM::PowOp, LLVM::SinOp, LLVM::SqrtOp>(); 188 189 // TODO: Remove once we support replacing non-root ops. 190 target.addLegalOp<gpu::YieldOp, gpu::GPUModuleOp, gpu::ModuleEndOp>(); 191 } 192 193 void mlir::populateGpuToNVVMConversionPatterns(LLVMTypeConverter &converter, 194 RewritePatternSet &patterns) { 195 populateWithGenerated(patterns); 196 patterns 197 .add<GPUIndexIntrinsicOpLowering<gpu::ThreadIdOp, NVVM::ThreadIdXOp, 198 NVVM::ThreadIdYOp, NVVM::ThreadIdZOp>, 199 GPUIndexIntrinsicOpLowering<gpu::BlockDimOp, NVVM::BlockDimXOp, 200 NVVM::BlockDimYOp, NVVM::BlockDimZOp>, 201 GPUIndexIntrinsicOpLowering<gpu::BlockIdOp, NVVM::BlockIdXOp, 202 NVVM::BlockIdYOp, NVVM::BlockIdZOp>, 203 GPUIndexIntrinsicOpLowering<gpu::GridDimOp, NVVM::GridDimXOp, 204 NVVM::GridDimYOp, NVVM::GridDimZOp>, 205 GPUShuffleOpLowering, GPUReturnOpLowering>(converter); 206 207 // Explicitly drop memory space when lowering private memory 208 // attributions since NVVM models it as `alloca`s in the default 209 // memory space and does not support `alloca`s with addrspace(5). 210 patterns.add<GPUFuncOpLowering>( 211 converter, /*allocaAddrSpace=*/0, 212 Identifier::get(NVVM::NVVMDialect::getKernelFuncAttrName(), 213 &converter.getContext())); 214 215 patterns.add<OpToFuncCallLowering<AbsFOp>>(converter, "__nv_fabsf", 216 "__nv_fabs"); 217 patterns.add<OpToFuncCallLowering<math::AtanOp>>(converter, "__nv_atanf", 218 "__nv_atan"); 219 patterns.add<OpToFuncCallLowering<math::Atan2Op>>(converter, "__nv_atan2f", 220 "__nv_atan2"); 221 patterns.add<OpToFuncCallLowering<CeilFOp>>(converter, "__nv_ceilf", 222 "__nv_ceil"); 223 patterns.add<OpToFuncCallLowering<math::CosOp>>(converter, "__nv_cosf", 224 "__nv_cos"); 225 patterns.add<OpToFuncCallLowering<math::ExpOp>>(converter, "__nv_expf", 226 "__nv_exp"); 227 patterns.add<OpToFuncCallLowering<math::ExpM1Op>>(converter, "__nv_expm1f", 228 "__nv_expm1"); 229 patterns.add<OpToFuncCallLowering<FloorFOp>>(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