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                   OpAdaptor adaptor,
73                   ConversionPatternRewriter &rewriter) const override {
74     Operation *op = subgroupMmaLoadMatrixOp.getOperation();
75     if (failed(areAllLLVMTypes(op, adaptor.getOperands(), 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     // MemRefDescriptor to extract alignedPtr and offset.
92     MemRefDescriptor promotedSrcOp(adaptor.srcMemref());
93 
94     // Emit ops which compute the load offset using `srcOffsetI`,
95     // `srcOffsetJ`. The actualOffset is (memrefOffset + (alignedPtr +
96     // ((leadDimension * srcOffsetI) + srcOffsetJ)). The memrefs here are
97     // assumed to be normalized and hence the simple conversion works.
98     SmallVector<Value> indices(adaptor.indices());
99     Value srcOffsetIVal = indices[0];
100     Value srcOffsetJVal = indices[1];
101     Type i32Ty = rewriter.getI32Type();
102     Value leadingDim32 =
103         rewriter.create<LLVM::ConstantOp>(loc, i32Ty, leadDimension);
104     Value numElemsLeadDim =
105         rewriter.create<LLVM::MulOp>(loc, i32Ty, leadingDim32, srcOffsetIVal);
106     Value loadOffset = rewriter.create<LLVM::AddOp>(loc, i32Ty, numElemsLeadDim,
107                                                     srcOffsetJVal);
108 
109     Value promotedSrcOpToUse;
110     promotedSrcOpToUse = promotedSrcOp.offset(rewriter, loc);
111     Value actualOffset = rewriter.create<LLVM::AddOp>(loc, i32Ty, loadOffset,
112                                                       promotedSrcOpToUse);
113     Value loadAddress = rewriter.create<LLVM::GEPOp>(
114         loc, promotedSrcOp.getElementPtrType(),
115         promotedSrcOp.alignedPtr(rewriter, loc), ArrayRef<Value>{actualOffset});
116 
117     // Bitcast the base address pointer of the destination memref, So that
118     // values can be stored in chunks of 32-bits and semantics match with the
119     // intrinsic exposed by NVPTX backend.
120     Value loadAddressCasted = rewriter.create<LLVM::BitcastOp>(
121         loc,
122         LLVM::LLVMPointerType::get(
123             i32Ty, promotedSrcOp.getElementPtrType().getAddressSpace()),
124         loadAddress);
125 
126     // Get the shape of the MMAMatrix type being returned. The shape will
127     // choose which intrinsic this op will be lowered to.
128     gpu::MMAMatrixType retType =
129         subgroupMmaLoadMatrixOp.res().getType().cast<gpu::MMAMatrixType>();
130     ArrayRef<int64_t> retTypeShape = retType.getShape();
131 
132     Type resType = convertMMAToLLVMType(retType);
133     StringRef operandStr = retType.getOperand();
134 
135     // Create nvvm.mma_load op according to the operand types.
136     SmallVector<Value, 2> loadOpOperands({loadAddressCasted, leadingDim32});
137     if (operandStr.equals("AOp")) {
138       if (retTypeShape[0] == 16 && retTypeShape[1] == 16) {
139         rewriter.replaceOpWithNewOp<NVVM::WMMALoadAM16N16K16Op>(op, resType,
140                                                                 loadOpOperands);
141       } else {
142         return rewriter.notifyMatchFailure(op, kInvalidCaseStr);
143       }
144     } else if (operandStr.equals("BOp")) {
145       if (retTypeShape[0] == 16 && retTypeShape[1] == 16) {
146         rewriter.replaceOpWithNewOp<NVVM::WMMALoadBM16N16K16Op>(op, resType,
147                                                                 loadOpOperands);
148       } else {
149         return rewriter.notifyMatchFailure(op, kInvalidCaseStr);
150       }
151     } else {
152       if (retTypeShape[0] == 16 && retTypeShape[1] == 16) {
153         if (retType.getElementType().isF16()) {
154           rewriter.replaceOpWithNewOp<NVVM::WMMALoadCF16M16N16K16Op>(
155               op, resType, loadOpOperands);
156         } else if (retType.getElementType().isF32()) {
157           rewriter.replaceOpWithNewOp<NVVM::WMMALoadCF32M16N16K16Op>(
158               op, resType, loadOpOperands);
159         }
160       } else {
161         return rewriter.notifyMatchFailure(op, kInvalidCaseStr);
162       }
163     }
164     return success();
165   }
166 };
167 
168 /// This class implements the conversion of GPU MMA storeOp to wmma.store op
169 /// in the NVVM dialect. The conversion not only emits the NVVM op but also
170 /// emits code that is necessary to unpack the data in the source and
171 /// convert the data in the format that is needed by the NVVM op.
172 struct WmmaStoreOpToNVVMLowering
173     : public ConvertOpToLLVMPattern<gpu::SubgroupMmaStoreMatrixOp> {
174   using ConvertOpToLLVMPattern<
175       gpu::SubgroupMmaStoreMatrixOp>::ConvertOpToLLVMPattern;
176 
177   LogicalResult
178   matchAndRewrite(gpu::SubgroupMmaStoreMatrixOp subgroupMmaStoreMatrixOp,
179                   OpAdaptor adaptor,
180                   ConversionPatternRewriter &rewriter) const override {
181     Operation *op = subgroupMmaStoreMatrixOp.getOperation();
182     if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)))
183       return failure();
184 
185     unsigned indexTypeBitwidth =
186         this->getTypeConverter()->getIndexTypeBitwidth();
187     // The corresponding intrinsics expects leadDimension to be a 32-bit
188     // integer, so all the calculations of linearizing the store address
189     // must also follow this restriction.
190     if (indexTypeBitwidth != 32)
191       return rewriter.notifyMatchFailure(
192           op, "expected indices to the memref to be 32-bit wide.");
193 
194     Location loc = op->getLoc();
195 
196     // MemRefDescriptor to extract alignedPtr and offset.
197     MemRefDescriptor promotedDstOp(adaptor.dstMemref());
198 
199     auto leadDimension = subgroupMmaStoreMatrixOp.leadDimensionAttr();
200 
201     // Emit ops which compute the store offset using `dstOffsetI`,
202     // `dstOffsetJ`. The actualOffset is (memrefOffset + (alignedPtr +
203     // ((leadDimension * dstOffsetI) + dstOffsetJ)).
204     SmallVector<Value> indices(adaptor.indices());
205     Value dstOffsetIVal = indices[0];
206     Value dstOffsetJVal = indices[1];
207     Type i32Ty = rewriter.getI32Type();
208     Value leadingDim32 =
209         rewriter.create<LLVM::ConstantOp>(loc, i32Ty, leadDimension);
210     Value numElemsLeadDim =
211         rewriter.create<LLVM::MulOp>(loc, i32Ty, leadingDim32, dstOffsetIVal);
212     Value loadOffset = rewriter.create<LLVM::AddOp>(loc, i32Ty, numElemsLeadDim,
213                                                     dstOffsetJVal);
214 
215     Value promotedDstOpToUse;
216     promotedDstOpToUse = promotedDstOp.offset(rewriter, loc);
217     Value actualOffset = rewriter.create<LLVM::AddOp>(loc, i32Ty, loadOffset,
218                                                       promotedDstOpToUse);
219     Value storeAddress = rewriter.create<LLVM::GEPOp>(
220         loc, promotedDstOp.getElementPtrType(),
221         promotedDstOp.alignedPtr(rewriter, loc), ArrayRef<Value>{actualOffset});
222 
223     // Bitcast the base address pointer of the destination memref, So that
224     // values can be stored in chunks of 32-bits and semantics match with the
225     // intrinsic exposed by NVPTX backend.
226     Value storeAddressCasted = rewriter.create<LLVM::BitcastOp>(
227         loc,
228         LLVM::LLVMPointerType::get(
229             i32Ty, promotedDstOp.getElementPtrType().getAddressSpace()),
230         storeAddress);
231 
232     SmallVector<Value, 4> storeOpOperands;
233     storeOpOperands.push_back(storeAddressCasted);
234 
235     // Get the shape of the MMAMatrix type being stored. The shape will
236     // choose which intrinsic this op will be lowered to.
237     gpu::MMAMatrixType srcType =
238         subgroupMmaStoreMatrixOp.src().getType().cast<gpu::MMAMatrixType>();
239     ArrayRef<int64_t> srcTypeShape = srcType.getShape();
240 
241     auto matrixType = adaptor.src().getType().cast<LLVM::LLVMStructType>();
242     for (unsigned i = 0, e = matrixType.getBody().size(); i < e; ++i) {
243       Value toUse = rewriter.create<LLVM::ExtractValueOp>(
244           loc, matrixType.getBody()[i], adaptor.src(),
245           rewriter.getI32ArrayAttr(i));
246       storeOpOperands.push_back(toUse);
247     }
248     storeOpOperands.push_back(leadingDim32);
249     // Unpack the results from the source.
250     if (srcType.getElementType().isF16()) {
251       // Create nvvm.mma_store op.
252       if (srcTypeShape[0] == 16 && srcTypeShape[1] == 16) {
253         rewriter.create<NVVM::WMMAStoreF16M16N16K16Op>(loc, storeOpOperands);
254       } else {
255         return rewriter.notifyMatchFailure(op, kInvalidCaseStr);
256       }
257       rewriter.eraseOp(op);
258       return success();
259     }
260     if (srcType.getElementType().isF32()) {
261       // Create nvvm.mma_store op.
262       if (srcTypeShape[0] == 16 && srcTypeShape[1] == 16)
263         rewriter.create<NVVM::WMMAStoreF32M16N16K16Op>(loc, storeOpOperands);
264       else {
265         return rewriter.notifyMatchFailure(op, kInvalidCaseStr);
266       }
267       rewriter.eraseOp(op);
268       return success();
269     }
270     return failure();
271   }
272 };
273 
274 /// This class implements the conversion of GPU MMA computeOp to wmma.mma op
275 /// in the NVVM dialect.
276 struct WmmaMmaOpToNVVMLowering
277     : public ConvertOpToLLVMPattern<gpu::SubgroupMmaComputeOp> {
278   using ConvertOpToLLVMPattern<
279       gpu::SubgroupMmaComputeOp>::ConvertOpToLLVMPattern;
280 
281   LogicalResult
282   matchAndRewrite(gpu::SubgroupMmaComputeOp subgroupMmaComputeOp,
283                   OpAdaptor adaptor,
284                   ConversionPatternRewriter &rewriter) const override {
285     Operation *op = subgroupMmaComputeOp.getOperation();
286     if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)))
287       return failure();
288 
289     Location loc = op->getLoc();
290 
291     // The wmma.mma intrinsic in llvm requires the operands as individual
292     // values. So individual elements from the memrefs need to be extracted and
293     // then passed on to the intrinsic call. Emit llvm ops to extract individual
294     // values form lowered memrefs.
295     SmallVector<Value> unpackedOps;
296 
297     auto unpackOp = [&](Value operand) {
298       auto structType = operand.getType().cast<LLVM::LLVMStructType>();
299       for (size_t i = 0, e = structType.getBody().size(); i < e; ++i) {
300         Value toUse = rewriter.create<LLVM::ExtractValueOp>(
301             loc, structType.getBody()[i], operand, rewriter.getI32ArrayAttr(i));
302         unpackedOps.push_back(toUse);
303       }
304     };
305 
306     // Get the shapes of the MMAMatrix type being used. The shapes will
307     // choose which intrinsic this op will be lowered to.
308     gpu::MMAMatrixType aType =
309         subgroupMmaComputeOp.opA().getType().cast<gpu::MMAMatrixType>();
310     ArrayRef<int64_t> aTypeShape = aType.getShape();
311     gpu::MMAMatrixType bType =
312         subgroupMmaComputeOp.opB().getType().cast<gpu::MMAMatrixType>();
313     ArrayRef<int64_t> bTypeShape = bType.getShape();
314     gpu::MMAMatrixType cType =
315         subgroupMmaComputeOp.opC().getType().cast<gpu::MMAMatrixType>();
316     ArrayRef<int64_t> cTypeShape = cType.getShape();
317 
318     unpackOp(adaptor.opA());
319     unpackOp(adaptor.opB());
320     unpackOp(adaptor.opC());
321 
322     if (cType.getElementType().isF16()) {
323       if (aTypeShape[0] == 16 && aTypeShape[1] == 16 && bTypeShape[0] == 16 &&
324           bTypeShape[1] == 16 && cTypeShape[0] == 16 && cTypeShape[1] == 16) {
325         // Create nvvm.wmma.mma op.
326         rewriter.replaceOpWithNewOp<NVVM::WMMAMmaF16F16M16N16K16Op>(
327             op, adaptor.opC().getType(), unpackedOps);
328 
329         return success();
330       }
331       return rewriter.notifyMatchFailure(op, kInvalidCaseStr);
332     }
333     if (cType.getElementType().isF32()) {
334       if (aTypeShape[0] == 16 && aTypeShape[1] == 16 && bTypeShape[0] == 16 &&
335           bTypeShape[1] == 16 && cTypeShape[0] == 16 && cTypeShape[1] == 16) {
336         // Create nvvm.wmma.mma op.
337         rewriter.replaceOpWithNewOp<NVVM::WMMAMmaF32F32M16N16K16Op>(
338             op, adaptor.opC().getType(), unpackedOps);
339 
340         return success();
341       }
342       return rewriter.notifyMatchFailure(op, kInvalidCaseStr);
343     }
344     return failure();
345   }
346 };
347 
348 /// Convert GPU MMA ConstantMatrixOp to a chain of InsertValueOp.
349 struct WmmaConstantOpToNVVMLowering
350     : public ConvertOpToLLVMPattern<gpu::SubgroupMmaConstantMatrixOp> {
351   using ConvertOpToLLVMPattern<
352       gpu::SubgroupMmaConstantMatrixOp>::ConvertOpToLLVMPattern;
353 
354   LogicalResult
355   matchAndRewrite(gpu::SubgroupMmaConstantMatrixOp subgroupMmaConstantOp,
356                   OpAdaptor adaptor,
357                   ConversionPatternRewriter &rewriter) const override {
358     if (failed(areAllLLVMTypes(subgroupMmaConstantOp.getOperation(),
359                                adaptor.getOperands(), rewriter)))
360       return failure();
361     Location loc = subgroupMmaConstantOp.getLoc();
362     Value cst = adaptor.getOperands()[0];
363     LLVM::LLVMStructType type = convertMMAToLLVMType(
364         subgroupMmaConstantOp.getType().cast<gpu::MMAMatrixType>());
365     // If the element type is a vector create a vector from the operand.
366     if (auto vecType = type.getBody()[0].dyn_cast<VectorType>()) {
367       Value vecCst = rewriter.create<LLVM::UndefOp>(loc, vecType);
368       for (int64_t vecEl = 0; vecEl < vecType.getNumElements(); vecEl++) {
369         Value idx = rewriter.create<LLVM::ConstantOp>(
370             loc, typeConverter->convertType(rewriter.getIntegerType(32)),
371             rewriter.getI32IntegerAttr(vecEl));
372         vecCst = rewriter.create<LLVM::InsertElementOp>(loc, vecType, vecCst,
373                                                         cst, idx);
374       }
375       cst = vecCst;
376     }
377     Value matrixStruct = rewriter.create<LLVM::UndefOp>(loc, type);
378     for (size_t i : llvm::seq(size_t(0), type.getBody().size())) {
379       matrixStruct = rewriter.create<LLVM::InsertValueOp>(
380           loc, matrixStruct, cst, rewriter.getI32ArrayAttr(i));
381     }
382     rewriter.replaceOp(subgroupMmaConstantOp, matrixStruct);
383     return success();
384   }
385 };
386 
387 } // anonymous namespace
388 
389 namespace mlir {
390 void populateGpuWMMAToNVVMConversionPatterns(LLVMTypeConverter &converter,
391                                              RewritePatternSet &patterns) {
392   patterns.insert<WmmaLoadOpToNVVMLowering, WmmaMmaOpToNVVMLowering,
393                   WmmaStoreOpToNVVMLowering, WmmaConstantOpToNVVMLowering>(
394       converter);
395 }
396 } // namespace mlir
397