1 //===------ WmmaOpsToNVVM.cpp - WMMA LD/ST/Compute to NVVM lowering -------===//
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 contains definitions of patterns to lower GPU Subgroup MMA ops to
10 // NVVM Dialect.
11 //
12 //===----------------------------------------------------------------------===//
13 
14 #include "mlir/Conversion/LLVMCommon/Pattern.h"
15 #include "mlir/Dialect/GPU/GPUDialect.h"
16 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
17 #include "mlir/Dialect/LLVMIR/NVVMDialect.h"
18 
19 using namespace mlir;
20 
21 namespace {
22 
23 /// Checks if all the operands of the op being lowered are of LLVM Types. The
24 /// types are expected to be converted by the `LLVMTypeConverter` before the op
25 /// is actually lowered. If the type of an operands is not already converted it
26 /// hints a missing typeConversion and failure is returned in that case.
27 static LogicalResult areAllLLVMTypes(Operation *op, ValueRange operands,
28                                      ConversionPatternRewriter &rewriter) {
29   if (!llvm::all_of(operands, [](Value value) {
30         return LLVM::isCompatibleType(value.getType());
31       })) {
32     return rewriter.notifyMatchFailure(
33         op, "cannot convert if operands aren't of LLVM type.");
34   }
35 
36   return success();
37 }
38 
39 /// Error string to emit when unimplemented WMMA variant is encountered.
40 static constexpr StringRef kInvalidCaseStr =
41     "Unimplemented WMMA variant, Only M16N16K16 version implemented.";
42 
43 /// Return the LLVMStructureType corresponding to the MMAMatrixType `type`.
44 static LLVM::LLVMStructType convertMMAToLLVMType(gpu::MMAMatrixType type) {
45   StringRef operandStr = type.getOperand();
46   assert(type.getElementType().isa<FloatType>());
47   Type baseType = type.getElementType().isF16()
48                       ? VectorType::get(2, type.getElementType())
49                       : type.getElementType();
50   auto getLLVMType = [&](int64_t numElements) {
51     return LLVM::LLVMStructType::getLiteral(
52         type.getContext(), SmallVector<Type, 8>(numElements, baseType));
53   };
54   if (operandStr.equals("AOp") || operandStr.equals("BOp"))
55     return getLLVMType(8);
56   if (type.getElementType().isF16())
57     return getLLVMType(4);
58   return getLLVMType(8);
59 }
60 
61 /// This class implements the conversion of GPU MMA loadOp to wmma.load op
62 /// in the NVVM dialect. The conversion not only emits the NVVM op but also
63 /// emits code that is necessary to store the data in the destination memref
64 /// after it has been loaded.
65 struct WmmaLoadOpToNVVMLowering
66     : public ConvertOpToLLVMPattern<gpu::SubgroupMmaLoadMatrixOp> {
67   using ConvertOpToLLVMPattern<
68       gpu::SubgroupMmaLoadMatrixOp>::ConvertOpToLLVMPattern;
69 
70   LogicalResult
71   matchAndRewrite(gpu::SubgroupMmaLoadMatrixOp subgroupMmaLoadMatrixOp,
72                   ArrayRef<Value> operands,
73                   ConversionPatternRewriter &rewriter) const override {
74     Operation *op = subgroupMmaLoadMatrixOp.getOperation();
75     if (failed(areAllLLVMTypes(op, operands, rewriter)))
76       return failure();
77 
78     unsigned indexTypeBitwidth =
79         this->getTypeConverter()->getIndexTypeBitwidth();
80 
81     // The corresponding intrinsics expects leadDimension to be a 32-bit
82     // integer, so all the calculations of linearizing the load address
83     // must also follow this restriction.
84     if (indexTypeBitwidth != 32)
85       return rewriter.notifyMatchFailure(
86           op, "Expected indices to the memref to be 32-bit wide.");
87     Location loc = op->getLoc();
88 
89     auto leadDimension = subgroupMmaLoadMatrixOp.leadDimensionAttr();
90 
91     gpu::SubgroupMmaLoadMatrixOpAdaptor adaptor(operands);
92     // MemRefDescriptor to extract alignedPtr and offset.
93     MemRefDescriptor promotedSrcOp(adaptor.srcMemref());
94 
95     // Emit ops which compute the load offset using `srcOffsetI`,
96     // `srcOffsetJ`. The actualOffset is (memrefOffset + (alignedPtr +
97     // ((leadDimension * srcOffsetI) + srcOffsetJ)). The memrefs here are
98     // assumed to be normalized and hence the simple conversion works.
99     SmallVector<Value> indices(adaptor.indices());
100     Value srcOffsetIVal = indices[0];
101     Value srcOffsetJVal = indices[1];
102     Type i32Ty = rewriter.getI32Type();
103     Value leadingDim32 =
104         rewriter.create<LLVM::ConstantOp>(loc, i32Ty, leadDimension);
105     Value numElemsLeadDim =
106         rewriter.create<LLVM::MulOp>(loc, i32Ty, leadingDim32, srcOffsetIVal);
107     Value loadOffset = rewriter.create<LLVM::AddOp>(loc, i32Ty, numElemsLeadDim,
108                                                     srcOffsetJVal);
109 
110     Value promotedSrcOpToUse;
111     promotedSrcOpToUse = promotedSrcOp.offset(rewriter, loc);
112     Value actualOffset = rewriter.create<LLVM::AddOp>(loc, i32Ty, loadOffset,
113                                                       promotedSrcOpToUse);
114     Value loadAddress = rewriter.create<LLVM::GEPOp>(
115         loc, promotedSrcOp.getElementPtrType(),
116         promotedSrcOp.alignedPtr(rewriter, loc), ArrayRef<Value>{actualOffset});
117 
118     // Bitcast the base address pointer of the destination memref, So that
119     // values can be stored in chunks of 32-bits and semantics match with the
120     // intrinsic exposed by NVPTX backend.
121     Value loadAddressCasted = rewriter.create<LLVM::BitcastOp>(
122         loc,
123         LLVM::LLVMPointerType::get(
124             i32Ty, promotedSrcOp.getElementPtrType().getAddressSpace()),
125         loadAddress);
126 
127     // Get the shape of the MMAMatrix type being returned. The shape will
128     // choose which intrinsic this op will be lowered to.
129     gpu::MMAMatrixType retType =
130         subgroupMmaLoadMatrixOp.res().getType().cast<gpu::MMAMatrixType>();
131     ArrayRef<int64_t> retTypeShape = retType.getShape();
132 
133     Type resType = convertMMAToLLVMType(retType);
134     StringRef operandStr = retType.getOperand();
135 
136     // Create nvvm.mma_load op according to the operand types.
137     SmallVector<Value, 2> loadOpOperands({loadAddressCasted, leadingDim32});
138     if (operandStr.equals("AOp")) {
139       if (retTypeShape[0] == 16 && retTypeShape[1] == 16) {
140         rewriter.replaceOpWithNewOp<NVVM::WMMALoadAM16N16K16Op>(op, resType,
141                                                                 loadOpOperands);
142       } else {
143         return rewriter.notifyMatchFailure(op, kInvalidCaseStr);
144       }
145     } else if (operandStr.equals("BOp")) {
146       if (retTypeShape[0] == 16 && retTypeShape[1] == 16) {
147         rewriter.replaceOpWithNewOp<NVVM::WMMALoadBM16N16K16Op>(op, resType,
148                                                                 loadOpOperands);
149       } else {
150         return rewriter.notifyMatchFailure(op, kInvalidCaseStr);
151       }
152     } else {
153       if (retTypeShape[0] == 16 && retTypeShape[1] == 16) {
154         if (retType.getElementType().isF16()) {
155           rewriter.replaceOpWithNewOp<NVVM::WMMALoadCF16M16N16K16Op>(
156               op, resType, loadOpOperands);
157         } else if (retType.getElementType().isF32()) {
158           rewriter.replaceOpWithNewOp<NVVM::WMMALoadCF32M16N16K16Op>(
159               op, resType, loadOpOperands);
160         }
161       } else {
162         return rewriter.notifyMatchFailure(op, kInvalidCaseStr);
163       }
164     }
165     return success();
166   }
167 };
168 
169 /// This class implements the conversion of GPU MMA storeOp to wmma.store op
170 /// in the NVVM dialect. The conversion not only emits the NVVM op but also
171 /// emits code that is necessary to unpack the data in the source and
172 /// convert the data in the format that is needed by the NVVM op.
173 struct WmmaStoreOpToNVVMLowering
174     : public ConvertOpToLLVMPattern<gpu::SubgroupMmaStoreMatrixOp> {
175   using ConvertOpToLLVMPattern<
176       gpu::SubgroupMmaStoreMatrixOp>::ConvertOpToLLVMPattern;
177 
178   LogicalResult
179   matchAndRewrite(gpu::SubgroupMmaStoreMatrixOp subgroupMmaStoreMatrixOp,
180                   ArrayRef<Value> operands,
181                   ConversionPatternRewriter &rewriter) const override {
182     Operation *op = subgroupMmaStoreMatrixOp.getOperation();
183     if (failed(areAllLLVMTypes(op, operands, rewriter)))
184       return failure();
185 
186     unsigned indexTypeBitwidth =
187         this->getTypeConverter()->getIndexTypeBitwidth();
188     // The corresponding intrinsics expects leadDimension to be a 32-bit
189     // integer, so all the calculations of linearizing the store address
190     // must also follow this restriction.
191     if (indexTypeBitwidth != 32)
192       return rewriter.notifyMatchFailure(
193           op, "expected indices to the memref to be 32-bit wide.");
194 
195     Location loc = op->getLoc();
196 
197     gpu::SubgroupMmaStoreMatrixOpAdaptor adaptor(operands);
198     // MemRefDescriptor to extract alignedPtr and offset.
199     MemRefDescriptor promotedDstOp(adaptor.dstMemref());
200 
201     auto leadDimension = subgroupMmaStoreMatrixOp.leadDimensionAttr();
202 
203     // Emit ops which compute the store offset using `dstOffsetI`,
204     // `dstOffsetJ`. The actualOffset is (memrefOffset + (alignedPtr +
205     // ((leadDimension * dstOffsetI) + dstOffsetJ)).
206     SmallVector<Value> indices(adaptor.indices());
207     Value dstOffsetIVal = indices[0];
208     Value dstOffsetJVal = indices[1];
209     Type i32Ty = rewriter.getI32Type();
210     Value leadingDim32 =
211         rewriter.create<LLVM::ConstantOp>(loc, i32Ty, leadDimension);
212     Value numElemsLeadDim =
213         rewriter.create<LLVM::MulOp>(loc, i32Ty, leadingDim32, dstOffsetIVal);
214     Value loadOffset = rewriter.create<LLVM::AddOp>(loc, i32Ty, numElemsLeadDim,
215                                                     dstOffsetJVal);
216 
217     Value promotedDstOpToUse;
218     promotedDstOpToUse = promotedDstOp.offset(rewriter, loc);
219     Value actualOffset = rewriter.create<LLVM::AddOp>(loc, i32Ty, loadOffset,
220                                                       promotedDstOpToUse);
221     Value storeAddress = rewriter.create<LLVM::GEPOp>(
222         loc, promotedDstOp.getElementPtrType(),
223         promotedDstOp.alignedPtr(rewriter, loc), ArrayRef<Value>{actualOffset});
224 
225     // Bitcast the base address pointer of the destination memref, So that
226     // values can be stored in chunks of 32-bits and semantics match with the
227     // intrinsic exposed by NVPTX backend.
228     Value storeAddressCasted = rewriter.create<LLVM::BitcastOp>(
229         loc,
230         LLVM::LLVMPointerType::get(
231             i32Ty, promotedDstOp.getElementPtrType().getAddressSpace()),
232         storeAddress);
233 
234     SmallVector<Value, 4> storeOpOperands;
235     storeOpOperands.push_back(storeAddressCasted);
236 
237     // Get the shape of the MMAMatrix type being stored. The shape will
238     // choose which intrinsic this op will be lowered to.
239     gpu::MMAMatrixType srcType =
240         subgroupMmaStoreMatrixOp.src().getType().cast<gpu::MMAMatrixType>();
241     ArrayRef<int64_t> srcTypeShape = srcType.getShape();
242 
243     auto matrixType = adaptor.src().getType().cast<LLVM::LLVMStructType>();
244     for (unsigned i = 0, e = matrixType.getBody().size(); i < e; ++i) {
245       Value toUse = rewriter.create<LLVM::ExtractValueOp>(
246           loc, matrixType.getBody()[i], adaptor.src(),
247           rewriter.getI32ArrayAttr(i));
248       storeOpOperands.push_back(toUse);
249     }
250     storeOpOperands.push_back(leadingDim32);
251     // Unpack the results from the source.
252     if (srcType.getElementType().isF16()) {
253       // Create nvvm.mma_store op.
254       if (srcTypeShape[0] == 16 && srcTypeShape[1] == 16) {
255         rewriter.create<NVVM::WMMAStoreF16M16N16K16Op>(loc, storeOpOperands);
256       } else {
257         return rewriter.notifyMatchFailure(op, kInvalidCaseStr);
258       }
259       rewriter.eraseOp(op);
260       return success();
261     }
262     if (srcType.getElementType().isF32()) {
263       // Create nvvm.mma_store op.
264       if (srcTypeShape[0] == 16 && srcTypeShape[1] == 16)
265         rewriter.create<NVVM::WMMAStoreF32M16N16K16Op>(loc, storeOpOperands);
266       else {
267         return rewriter.notifyMatchFailure(op, kInvalidCaseStr);
268       }
269       rewriter.eraseOp(op);
270       return success();
271     }
272     return failure();
273   }
274 };
275 
276 /// This class implements the conversion of GPU MMA computeOp to wmma.mma op
277 /// in the NVVM dialect.
278 struct WmmaMmaOpToNVVMLowering
279     : public ConvertOpToLLVMPattern<gpu::SubgroupMmaComputeOp> {
280   using ConvertOpToLLVMPattern<
281       gpu::SubgroupMmaComputeOp>::ConvertOpToLLVMPattern;
282 
283   LogicalResult
284   matchAndRewrite(gpu::SubgroupMmaComputeOp subgroupMmaComputeOp,
285                   ArrayRef<Value> operands,
286                   ConversionPatternRewriter &rewriter) const override {
287     Operation *op = subgroupMmaComputeOp.getOperation();
288     if (failed(areAllLLVMTypes(op, operands, rewriter)))
289       return failure();
290 
291     Location loc = op->getLoc();
292 
293     // The wmma.mma intrinsic in llvm requires the operands as individual
294     // values. So individual elements from the memrefs need to be extracted and
295     // then passed on to the intrinsic call. Emit llvm ops to extract individual
296     // values form lowered memrefs.
297     SmallVector<Value> unpackedOps;
298 
299     auto unpackOp = [&](Value operand) {
300       auto structType = operand.getType().cast<LLVM::LLVMStructType>();
301       for (size_t i = 0, e = structType.getBody().size(); i < e; ++i) {
302         Value toUse = rewriter.create<LLVM::ExtractValueOp>(
303             loc, structType.getBody()[i], operand, rewriter.getI32ArrayAttr(i));
304         unpackedOps.push_back(toUse);
305       }
306     };
307 
308     // Get the shapes of the MMAMatrix type being used. The shapes will
309     // choose which intrinsic this op will be lowered to.
310     gpu::MMAMatrixType aType =
311         subgroupMmaComputeOp.opA().getType().cast<gpu::MMAMatrixType>();
312     ArrayRef<int64_t> aTypeShape = aType.getShape();
313     gpu::MMAMatrixType bType =
314         subgroupMmaComputeOp.opB().getType().cast<gpu::MMAMatrixType>();
315     ArrayRef<int64_t> bTypeShape = bType.getShape();
316     gpu::MMAMatrixType cType =
317         subgroupMmaComputeOp.opC().getType().cast<gpu::MMAMatrixType>();
318     ArrayRef<int64_t> cTypeShape = cType.getShape();
319 
320     gpu::SubgroupMmaComputeOpAdaptor transformedOperands(operands);
321     unpackOp(transformedOperands.opA());
322     unpackOp(transformedOperands.opB());
323     unpackOp(transformedOperands.opC());
324 
325     if (cType.getElementType().isF16()) {
326       if (aTypeShape[0] == 16 && aTypeShape[1] == 16 && bTypeShape[0] == 16 &&
327           bTypeShape[1] == 16 && cTypeShape[0] == 16 && cTypeShape[1] == 16) {
328         // Create nvvm.wmma.mma op.
329         rewriter.replaceOpWithNewOp<NVVM::WMMAMmaF16F16M16N16K16Op>(
330             op, transformedOperands.opC().getType(), unpackedOps);
331 
332         return success();
333       }
334       return rewriter.notifyMatchFailure(op, kInvalidCaseStr);
335     }
336     if (cType.getElementType().isF32()) {
337       if (aTypeShape[0] == 16 && aTypeShape[1] == 16 && bTypeShape[0] == 16 &&
338           bTypeShape[1] == 16 && cTypeShape[0] == 16 && cTypeShape[1] == 16) {
339         // Create nvvm.wmma.mma op.
340         rewriter.replaceOpWithNewOp<NVVM::WMMAMmaF32F32M16N16K16Op>(
341             op, transformedOperands.opC().getType(), unpackedOps);
342 
343         return success();
344       }
345       return rewriter.notifyMatchFailure(op, kInvalidCaseStr);
346     }
347     return failure();
348   }
349 };
350 
351 /// Convert GPU MMA ConstantMatrixOp to a chain of InsertValueOp.
352 struct WmmaConstantOpToNVVMLowering
353     : public ConvertOpToLLVMPattern<gpu::SubgroupMmaConstantMatrixOp> {
354   using ConvertOpToLLVMPattern<
355       gpu::SubgroupMmaConstantMatrixOp>::ConvertOpToLLVMPattern;
356 
357   LogicalResult
358   matchAndRewrite(gpu::SubgroupMmaConstantMatrixOp subgroupMmaConstantOp,
359                   ArrayRef<Value> operands,
360                   ConversionPatternRewriter &rewriter) const override {
361     if (failed(areAllLLVMTypes(subgroupMmaConstantOp.getOperation(), operands,
362                                rewriter)))
363       return failure();
364     Location loc = subgroupMmaConstantOp.getLoc();
365     Value cst = operands[0];
366     LLVM::LLVMStructType type = convertMMAToLLVMType(
367         subgroupMmaConstantOp.getType().cast<gpu::MMAMatrixType>());
368     // If the element type is a vector create a vector from the operand.
369     if (auto vecType = type.getBody()[0].dyn_cast<VectorType>()) {
370       Value vecCst = rewriter.create<LLVM::UndefOp>(loc, vecType);
371       for (int64_t vecEl = 0; vecEl < vecType.getNumElements(); vecEl++) {
372         Value idx = rewriter.create<LLVM::ConstantOp>(
373             loc, typeConverter->convertType(rewriter.getIntegerType(32)),
374             rewriter.getI32IntegerAttr(vecEl));
375         vecCst = rewriter.create<LLVM::InsertElementOp>(loc, vecType, vecCst,
376                                                         cst, idx);
377       }
378       cst = vecCst;
379     }
380     Value matrixStruct = rewriter.create<LLVM::UndefOp>(loc, type);
381     for (size_t i : llvm::seq(size_t(0), type.getBody().size())) {
382       matrixStruct = rewriter.create<LLVM::InsertValueOp>(
383           loc, matrixStruct, cst, rewriter.getI32ArrayAttr(i));
384     }
385     rewriter.replaceOp(subgroupMmaConstantOp, matrixStruct);
386     return success();
387   }
388 };
389 
390 } // anonymous namespace
391 
392 namespace mlir {
393 void populateGpuWMMAToNVVMConversionPatterns(LLVMTypeConverter &converter,
394                                              RewritePatternSet &patterns) {
395   patterns.insert<WmmaLoadOpToNVVMLowering, WmmaMmaOpToNVVMLowering,
396                   WmmaStoreOpToNVVMLowering, WmmaConstantOpToNVVMLowering>(
397       converter);
398 }
399 } // namespace mlir
400