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 /// NVVM memory space identifiers.
46 enum NVVMMemorySpace {
47   /// Global memory space identifier.
48   kGlobalMemorySpace = 1,
49   /// Shared memory space identifier.
50   kSharedMemorySpace = 3
51 };
52 
53 /// Convert gpu dialect shfl mode enum to the equivalent nvvm one.
54 static NVVM::ShflKind convertShflKind(gpu::ShuffleMode mode) {
55   switch (mode) {
56   case gpu::ShuffleMode::XOR:
57     return NVVM::ShflKind::bfly;
58   case gpu::ShuffleMode::UP:
59     return NVVM::ShflKind::up;
60   case gpu::ShuffleMode::DOWN:
61     return NVVM::ShflKind::down;
62   case gpu::ShuffleMode::IDX:
63     return NVVM::ShflKind::idx;
64   }
65   llvm_unreachable("unknown shuffle mode");
66 }
67 
68 struct GPUShuffleOpLowering : public ConvertOpToLLVMPattern<gpu::ShuffleOp> {
69   using ConvertOpToLLVMPattern<gpu::ShuffleOp>::ConvertOpToLLVMPattern;
70 
71   /// Lowers a shuffle to the corresponding NVVM op.
72   ///
73   /// Convert the `width` argument into an activeMask (a bitmask which specifies
74   /// which threads participate in the shuffle) and a maskAndClamp (specifying
75   /// the highest lane which participates in the shuffle).
76   ///
77   ///     %one = llvm.constant(1 : i32) : i32
78   ///     %minus_one = llvm.constant(-1 : i32) : i32
79   ///     %thirty_two = llvm.constant(32 : i32) : i32
80   ///     %num_lanes = llvm.sub %thirty_two, %width : i32
81   ///     %active_mask = llvm.lshr %minus_one, %num_lanes : i32
82   ///     %mask_and_clamp = llvm.sub %width, %one : i32
83   ///     %shfl = nvvm.shfl.sync.bfly %active_mask, %value, %offset,
84   ///         %mask_and_clamp : !llvm<"{ float, i1 }">
85   ///     %shfl_value = llvm.extractvalue %shfl[0 : index] :
86   ///         !llvm<"{ float, i1 }">
87   ///     %shfl_pred = llvm.extractvalue %shfl[1 : index] :
88   ///         !llvm<"{ float, i1 }">
89   LogicalResult
90   matchAndRewrite(gpu::ShuffleOp op, OpAdaptor adaptor,
91                   ConversionPatternRewriter &rewriter) const override {
92     Location loc = op->getLoc();
93 
94     auto valueTy = adaptor.value().getType();
95     auto int32Type = IntegerType::get(rewriter.getContext(), 32);
96     auto predTy = IntegerType::get(rewriter.getContext(), 1);
97     auto resultTy = LLVM::LLVMStructType::getLiteral(rewriter.getContext(),
98                                                      {valueTy, predTy});
99 
100     Value one = rewriter.create<LLVM::ConstantOp>(
101         loc, int32Type, rewriter.getI32IntegerAttr(1));
102     Value minusOne = rewriter.create<LLVM::ConstantOp>(
103         loc, int32Type, rewriter.getI32IntegerAttr(-1));
104     Value thirtyTwo = rewriter.create<LLVM::ConstantOp>(
105         loc, int32Type, rewriter.getI32IntegerAttr(32));
106     Value numLeadInactiveLane = rewriter.create<LLVM::SubOp>(
107         loc, int32Type, thirtyTwo, adaptor.width());
108     // Bit mask of active lanes: `(-1) >> (32 - activeWidth)`.
109     Value activeMask = rewriter.create<LLVM::LShrOp>(loc, int32Type, minusOne,
110                                                      numLeadInactiveLane);
111     Value maskAndClamp;
112     if (op.mode() == gpu::ShuffleMode::UP) {
113       // Clamp lane: `32 - activeWidth`
114       maskAndClamp = numLeadInactiveLane;
115     } else {
116       // Clamp lane: `activeWidth - 1`
117       maskAndClamp =
118           rewriter.create<LLVM::SubOp>(loc, int32Type, adaptor.width(), one);
119     }
120 
121     auto returnValueAndIsValidAttr = rewriter.getUnitAttr();
122     Value shfl = rewriter.create<NVVM::ShflOp>(
123         loc, resultTy, activeMask, adaptor.value(), adaptor.offset(),
124         maskAndClamp, convertShflKind(op.mode()), returnValueAndIsValidAttr);
125     Value shflValue = rewriter.create<LLVM::ExtractValueOp>(
126         loc, valueTy, shfl, rewriter.getIndexArrayAttr(0));
127     Value isActiveSrcLane = rewriter.create<LLVM::ExtractValueOp>(
128         loc, predTy, shfl, rewriter.getIndexArrayAttr(1));
129 
130     rewriter.replaceOp(op, {shflValue, isActiveSrcLane});
131     return success();
132   }
133 };
134 
135 struct GPUAsyncCopyLowering
136     : public ConvertOpToLLVMPattern<gpu::DeviceAsyncCopyOp> {
137   using ConvertOpToLLVMPattern<gpu::DeviceAsyncCopyOp>::ConvertOpToLLVMPattern;
138 
139   LogicalResult
140   matchAndRewrite(gpu::DeviceAsyncCopyOp op, OpAdaptor adaptor,
141                   ConversionPatternRewriter &rewriter) const override {
142     Location loc = op->getLoc();
143     auto dstMemrefType = op.dst().getType().cast<MemRefType>();
144     Value dstPtr = getStridedElementPtr(loc, dstMemrefType, adaptor.dst(),
145                                         adaptor.dstIndices(), rewriter);
146     auto i8Ty = IntegerType::get(op.getContext(), 8);
147     auto dstPointerType =
148         LLVM::LLVMPointerType::get(i8Ty, dstMemrefType.getMemorySpaceAsInt());
149     dstPtr = rewriter.create<LLVM::BitcastOp>(loc, dstPointerType, dstPtr);
150 
151     auto srcMemrefType = op.src().getType().cast<MemRefType>();
152 
153     Value scrPtr = getStridedElementPtr(loc, srcMemrefType, adaptor.src(),
154                                         adaptor.srcIndices(), rewriter);
155     auto srcPointerType =
156         LLVM::LLVMPointerType::get(i8Ty, srcMemrefType.getMemorySpaceAsInt());
157     scrPtr = rewriter.create<LLVM::BitcastOp>(loc, srcPointerType, scrPtr);
158     // Intrinsics takes a global pointer so we need an address space cast.
159     auto srcPointerGlobalType =
160         LLVM::LLVMPointerType::get(i8Ty, NVVMMemorySpace::kGlobalMemorySpace);
161     scrPtr = rewriter.create<LLVM::AddrSpaceCastOp>(loc, srcPointerGlobalType,
162                                                     scrPtr);
163     int64_t numElements = adaptor.numElements().getZExtValue();
164     int64_t sizeInBytes =
165         (dstMemrefType.getElementTypeBitWidth() / 8) * numElements;
166     rewriter.create<NVVM::CpAsyncOp>(loc, dstPtr, scrPtr,
167                                      rewriter.getI32IntegerAttr(sizeInBytes));
168 
169     // Drop the result token.
170     Value zero = rewriter.create<LLVM::ConstantOp>(
171         op->getLoc(), IntegerType::get(op.getContext(), 32),
172         rewriter.getI32IntegerAttr(0));
173     rewriter.replaceOp(op, zero);
174     return success();
175   }
176 };
177 
178 struct GPUAsyncCreateGroupLowering
179     : public ConvertOpToLLVMPattern<gpu::DeviceAsyncCreateGroupOp> {
180   using ConvertOpToLLVMPattern<
181       gpu::DeviceAsyncCreateGroupOp>::ConvertOpToLLVMPattern;
182 
183   LogicalResult
184   matchAndRewrite(gpu::DeviceAsyncCreateGroupOp op, OpAdaptor adaptor,
185                   ConversionPatternRewriter &rewriter) const override {
186     rewriter.create<NVVM::CpAsyncCommitGroupOp>(op.getLoc());
187     // Drop the result token.
188     Value zero = rewriter.create<LLVM::ConstantOp>(
189         op->getLoc(), IntegerType::get(op.getContext(), 32),
190         rewriter.getI32IntegerAttr(0));
191     rewriter.replaceOp(op, zero);
192     return success();
193   }
194 };
195 
196 struct GPUAsyncWaitLowering
197     : public ConvertOpToLLVMPattern<gpu::DeviceAsyncWaitOp> {
198   using ConvertOpToLLVMPattern<gpu::DeviceAsyncWaitOp>::ConvertOpToLLVMPattern;
199 
200   LogicalResult
201   matchAndRewrite(gpu::DeviceAsyncWaitOp op, OpAdaptor adaptor,
202                   ConversionPatternRewriter &rewriter) const override {
203     // If numGroup is not present pick 0 as a conservative correct value.
204     int32_t numGroups = adaptor.numGroups() ? *adaptor.numGroups() : 0;
205     rewriter.create<NVVM::CpAsyncWaitGroupOp>(op.getLoc(), numGroups);
206     rewriter.eraseOp(op);
207     return success();
208   }
209 };
210 
211 /// Import the GPU Ops to NVVM Patterns.
212 #include "GPUToNVVM.cpp.inc"
213 
214 /// A pass that replaces all occurrences of GPU device operations with their
215 /// corresponding NVVM equivalent.
216 ///
217 /// This pass only handles device code and is not meant to be run on GPU host
218 /// code.
219 struct LowerGpuOpsToNVVMOpsPass
220     : public ConvertGpuOpsToNVVMOpsBase<LowerGpuOpsToNVVMOpsPass> {
221   LowerGpuOpsToNVVMOpsPass() = default;
222   LowerGpuOpsToNVVMOpsPass(unsigned indexBitwidth) {
223     this->indexBitwidth = indexBitwidth;
224   }
225 
226   void runOnOperation() override {
227     gpu::GPUModuleOp m = getOperation();
228 
229     /// Customize the bitwidth used for the device side index computations.
230     LowerToLLVMOptions options(
231         m.getContext(),
232         DataLayout(cast<DataLayoutOpInterface>(m.getOperation())));
233     options.emitCWrappers = true;
234     if (indexBitwidth != kDeriveIndexBitwidthFromDataLayout)
235       options.overrideIndexBitwidth(indexBitwidth);
236 
237     /// MemRef conversion for GPU to NVVM lowering. The GPU dialect uses memory
238     /// space 5 for private memory attributions, but NVVM represents private
239     /// memory allocations as local `alloca`s in the default address space. This
240     /// converter drops the private memory space to support the use case above.
241     LLVMTypeConverter converter(m.getContext(), options);
242     converter.addConversion([&](MemRefType type) -> Optional<Type> {
243       if (type.getMemorySpaceAsInt() !=
244           gpu::GPUDialect::getPrivateAddressSpace())
245         return llvm::None;
246       return converter.convertType(MemRefType::Builder(type).setMemorySpace(0));
247     });
248     /// device-side async tokens cannot be materialized in nvvm. We just convert
249     /// them to a dummy i32 type in order to easily drop them during conversion.
250     converter.addConversion([&](gpu::DeviceAsyncTokenType type) -> Type {
251       return converter.convertType(IntegerType::get(type.getContext(), 32));
252     });
253     // Lowering for MMAMatrixType.
254     converter.addConversion([&](gpu::MMAMatrixType type) -> Type {
255       return convertMMAToLLVMType(type);
256     });
257     RewritePatternSet patterns(m.getContext());
258     RewritePatternSet llvmPatterns(m.getContext());
259 
260     // Apply in-dialect lowering first. In-dialect lowering will replace ops
261     // which need to be lowered further, which is not supported by a single
262     // conversion pass.
263     populateGpuRewritePatterns(patterns);
264     (void)applyPatternsAndFoldGreedily(m, std::move(patterns));
265 
266     arith::populateArithmeticToLLVMConversionPatterns(converter, llvmPatterns);
267     cf::populateControlFlowToLLVMConversionPatterns(converter, llvmPatterns);
268     populateFuncToLLVMConversionPatterns(converter, llvmPatterns);
269     populateMemRefToLLVMConversionPatterns(converter, llvmPatterns);
270     populateGpuToNVVMConversionPatterns(converter, llvmPatterns);
271     populateGpuWMMAToNVVMConversionPatterns(converter, llvmPatterns);
272     LLVMConversionTarget target(getContext());
273     configureGpuToNVVMConversionLegality(target);
274     if (failed(applyPartialConversion(m, target, std::move(llvmPatterns))))
275       signalPassFailure();
276   }
277 };
278 
279 } // namespace
280 
281 void mlir::configureGpuToNVVMConversionLegality(ConversionTarget &target) {
282   target.addIllegalOp<FuncOp>();
283   target.addLegalDialect<::mlir::LLVM::LLVMDialect>();
284   target.addLegalDialect<::mlir::NVVM::NVVMDialect>();
285   target.addIllegalDialect<gpu::GPUDialect>();
286   target.addIllegalOp<LLVM::CosOp, LLVM::ExpOp, LLVM::Exp2Op, LLVM::FAbsOp,
287                       LLVM::FCeilOp, LLVM::FFloorOp, LLVM::LogOp, LLVM::Log10Op,
288                       LLVM::Log2Op, LLVM::PowOp, LLVM::SinOp, LLVM::SqrtOp>();
289 
290   // TODO: Remove once we support replacing non-root ops.
291   target.addLegalOp<gpu::YieldOp, gpu::GPUModuleOp, gpu::ModuleEndOp>();
292 }
293 
294 void mlir::populateGpuToNVVMConversionPatterns(LLVMTypeConverter &converter,
295                                                RewritePatternSet &patterns) {
296   populateWithGenerated(patterns);
297   patterns
298       .add<GPUIndexIntrinsicOpLowering<gpu::ThreadIdOp, NVVM::ThreadIdXOp,
299                                        NVVM::ThreadIdYOp, NVVM::ThreadIdZOp>,
300            GPUIndexIntrinsicOpLowering<gpu::BlockDimOp, NVVM::BlockDimXOp,
301                                        NVVM::BlockDimYOp, NVVM::BlockDimZOp>,
302            GPUIndexIntrinsicOpLowering<gpu::BlockIdOp, NVVM::BlockIdXOp,
303                                        NVVM::BlockIdYOp, NVVM::BlockIdZOp>,
304            GPUIndexIntrinsicOpLowering<gpu::GridDimOp, NVVM::GridDimXOp,
305                                        NVVM::GridDimYOp, NVVM::GridDimZOp>,
306            GPUShuffleOpLowering, GPUReturnOpLowering>(converter);
307 
308   // Explicitly drop memory space when lowering private memory
309   // attributions since NVVM models it as `alloca`s in the default
310   // memory space and does not support `alloca`s with addrspace(5).
311   patterns.add<GPUFuncOpLowering>(
312       converter, /*allocaAddrSpace=*/0,
313       StringAttr::get(&converter.getContext(),
314                       NVVM::NVVMDialect::getKernelFuncAttrName()));
315 
316   patterns.add<OpToFuncCallLowering<math::AbsOp>>(converter, "__nv_fabsf",
317                                                   "__nv_fabs");
318   patterns.add<OpToFuncCallLowering<math::AtanOp>>(converter, "__nv_atanf",
319                                                    "__nv_atan");
320   patterns.add<OpToFuncCallLowering<math::Atan2Op>>(converter, "__nv_atan2f",
321                                                     "__nv_atan2");
322   patterns.add<OpToFuncCallLowering<math::CeilOp>>(converter, "__nv_ceilf",
323                                                    "__nv_ceil");
324   patterns.add<OpToFuncCallLowering<math::CosOp>>(converter, "__nv_cosf",
325                                                   "__nv_cos");
326   patterns.add<OpToFuncCallLowering<math::ExpOp>>(converter, "__nv_expf",
327                                                   "__nv_exp");
328   patterns.add<OpToFuncCallLowering<math::Exp2Op>>(converter, "__nv_exp2f",
329                                                    "__nv_exp2");
330   patterns.add<OpToFuncCallLowering<math::ExpM1Op>>(converter, "__nv_expm1f",
331                                                     "__nv_expm1");
332   patterns.add<OpToFuncCallLowering<math::FloorOp>>(converter, "__nv_floorf",
333                                                     "__nv_floor");
334   patterns.add<OpToFuncCallLowering<math::LogOp>>(converter, "__nv_logf",
335                                                   "__nv_log");
336   patterns.add<OpToFuncCallLowering<math::Log1pOp>>(converter, "__nv_log1pf",
337                                                     "__nv_log1p");
338   patterns.add<OpToFuncCallLowering<math::Log10Op>>(converter, "__nv_log10f",
339                                                     "__nv_log10");
340   patterns.add<OpToFuncCallLowering<math::Log2Op>>(converter, "__nv_log2f",
341                                                    "__nv_log2");
342   patterns.add<OpToFuncCallLowering<math::PowFOp>>(converter, "__nv_powf",
343                                                    "__nv_pow");
344   patterns.add<OpToFuncCallLowering<math::RsqrtOp>>(converter, "__nv_rsqrtf",
345                                                     "__nv_rsqrt");
346   patterns.add<OpToFuncCallLowering<math::SinOp>>(converter, "__nv_sinf",
347                                                   "__nv_sin");
348   patterns.add<OpToFuncCallLowering<math::SqrtOp>>(converter, "__nv_sqrtf",
349                                                    "__nv_sqrt");
350   patterns.add<OpToFuncCallLowering<math::TanhOp>>(converter, "__nv_tanhf",
351                                                    "__nv_tanh");
352   patterns.add<GPUAsyncCopyLowering, GPUAsyncCreateGroupLowering,
353                GPUAsyncWaitLowering>(converter);
354 }
355 
356 std::unique_ptr<OperationPass<gpu::GPUModuleOp>>
357 mlir::createLowerGpuOpsToNVVMOpsPass(unsigned indexBitwidth) {
358   return std::make_unique<LowerGpuOpsToNVVMOpsPass>(indexBitwidth);
359 }
360