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