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