1edd9515bSthomasraoux //===- VectorToGPU.cpp - Convert vector to GPU dialect ----------*- C++ -*-===//
2edd9515bSthomasraoux //
3edd9515bSthomasraoux // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4edd9515bSthomasraoux // See https://llvm.org/LICENSE.txt for license information.
5edd9515bSthomasraoux // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6edd9515bSthomasraoux //
7edd9515bSthomasraoux //===----------------------------------------------------------------------===//
8edd9515bSthomasraoux //
9edd9515bSthomasraoux // This file implements lowering of vector operations to GPU dialect ops.
10edd9515bSthomasraoux //
11edd9515bSthomasraoux //===----------------------------------------------------------------------===//
12edd9515bSthomasraoux 
13edd9515bSthomasraoux #include <type_traits>
14edd9515bSthomasraoux 
15edd9515bSthomasraoux #include "mlir/Conversion/VectorToGPU/VectorToGPU.h"
16edd9515bSthomasraoux 
17edd9515bSthomasraoux #include "../PassDetail.h"
18edd9515bSthomasraoux #include "mlir/Analysis/SliceAnalysis.h"
19a54f4eaeSMogball #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
20edd9515bSthomasraoux #include "mlir/Dialect/GPU/GPUDialect.h"
2166f878ceSMatthias Springer #include "mlir/Dialect/MemRef/IR/MemRef.h"
221a865592Sthomasraoux #include "mlir/Dialect/SCF/SCF.h"
23edd9515bSthomasraoux #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
2499ef9eebSMatthias Springer #include "mlir/Dialect/Vector/IR/VectorOps.h"
2599ef9eebSMatthias Springer #include "mlir/Dialect/Vector/Utils/VectorUtils.h"
26edd9515bSthomasraoux #include "mlir/IR/Builders.h"
27edd9515bSthomasraoux #include "mlir/Pass/Pass.h"
28edd9515bSthomasraoux #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
29edd9515bSthomasraoux #include "mlir/Transforms/Passes.h"
30edd9515bSthomasraoux 
31edd9515bSthomasraoux using namespace mlir;
32edd9515bSthomasraoux 
33edd9515bSthomasraoux // Return true if the contract op can be convert to MMA matmul.
34edd9515bSthomasraoux static bool contractSupportsMMAMatrixType(vector::ContractionOp contract) {
35edd9515bSthomasraoux   if (llvm::size(contract.masks()) != 0)
36edd9515bSthomasraoux     return false;
37edd9515bSthomasraoux 
38edd9515bSthomasraoux   using MapList = ArrayRef<ArrayRef<AffineExpr>>;
39edd9515bSthomasraoux   auto infer = [](MapList m) { return AffineMap::inferFromExprList(m); };
40edd9515bSthomasraoux   AffineExpr m, n, k;
41edd9515bSthomasraoux   bindDims(contract.getContext(), m, n, k);
42edd9515bSthomasraoux   auto iteratorTypes = contract.iterator_types().getValue();
43edd9515bSthomasraoux   if (!(isParallelIterator(iteratorTypes[0]) &&
44edd9515bSthomasraoux         isParallelIterator(iteratorTypes[1]) &&
45edd9515bSthomasraoux         isReductionIterator(iteratorTypes[2])))
46edd9515bSthomasraoux     return false;
47edd9515bSthomasraoux 
48edd9515bSthomasraoux   // The contract needs to represent a matmul to be able to convert to
49edd9515bSthomasraoux   // MMAMatrix matmul.
50edd9515bSthomasraoux   if (contract.getIndexingMaps() != infer({{m, k}, {k, n}, {m, n}}))
51edd9515bSthomasraoux     return false;
52edd9515bSthomasraoux 
53edd9515bSthomasraoux   return true;
54edd9515bSthomasraoux }
55edd9515bSthomasraoux 
56edd9515bSthomasraoux // Return the stide for the dimension 0 of |type| if it is a memref and has a
57edd9515bSthomasraoux // constant stride.
58edd9515bSthomasraoux static llvm::Optional<int64_t>
59edd9515bSthomasraoux getMemrefConstantHorizontalStride(ShapedType type) {
60edd9515bSthomasraoux   auto memrefType = type.dyn_cast<MemRefType>();
61edd9515bSthomasraoux   if (!memrefType)
62edd9515bSthomasraoux     return false;
63a57ccad5SThomas Raoux   // If the memref is 0 or 1D the horizontal stride is 0.
64a57ccad5SThomas Raoux   if(memrefType.getRank() < 2)
65a57ccad5SThomas Raoux     return 0;
66edd9515bSthomasraoux   int64_t offset = 0;
67edd9515bSthomasraoux   SmallVector<int64_t, 2> strides;
68*d77f4836SThomas Raoux   if (failed(getStridesAndOffset(memrefType, strides, offset)) ||
69*d77f4836SThomas Raoux       strides.back() != 1)
70edd9515bSthomasraoux     return llvm::None;
71a57ccad5SThomas Raoux   int64_t stride = strides[strides.size() - 2];
72a57ccad5SThomas Raoux   if (stride == ShapedType::kDynamicStrideOrOffset)
73edd9515bSthomasraoux     return llvm::None;
74a57ccad5SThomas Raoux   return stride;
75edd9515bSthomasraoux }
76edd9515bSthomasraoux 
77edd9515bSthomasraoux // Return true if the transfer op can be converted to a MMA matrix load.
78edd9515bSthomasraoux static bool transferReadSupportsMMAMatrixType(vector::TransferReadOp readOp) {
79edd9515bSthomasraoux   if (readOp.mask() || readOp.hasOutOfBoundsDim() ||
80edd9515bSthomasraoux       readOp.getVectorType().getRank() != 2)
81edd9515bSthomasraoux     return false;
82edd9515bSthomasraoux   if (!getMemrefConstantHorizontalStride(readOp.getShapedType()))
83edd9515bSthomasraoux     return false;
84e7969240SThomas Raoux   AffineMap map = readOp.permutation_map();
85e7969240SThomas Raoux   OpBuilder b(readOp.getContext());
86e7969240SThomas Raoux   AffineExpr innerDim = b.getAffineDimExpr(map.getNumDims() - 1);
87e7969240SThomas Raoux   AffineExpr zero = b.getAffineConstantExpr(0);
88e7969240SThomas Raoux   auto broadcastInnerDim = AffineMap::get(map.getNumDims(), 0, {zero, innerDim},
89e7969240SThomas Raoux                                           readOp.getContext());
90edd9515bSthomasraoux   // TODO: Support transpose once it is added to GPU dialect ops.
91e7969240SThomas Raoux   // For now we only support (d0, d1) -> (d0, d1) and (d0, d1) -> (0, d1).
926786d7e4SMehdi Amini   return !(!map.isMinorIdentity() && map != broadcastInnerDim);
93edd9515bSthomasraoux }
94edd9515bSthomasraoux 
95edd9515bSthomasraoux // Return true if the transfer op can be converted to a MMA matrix store.
96edd9515bSthomasraoux static bool
97edd9515bSthomasraoux transferWriteSupportsMMAMatrixType(vector::TransferWriteOp writeOp) {
98c537a943SNicolas Vasilache   // TODO: support 0-d corner case.
99c537a943SNicolas Vasilache   if (writeOp.getTransferRank() == 0)
100c537a943SNicolas Vasilache     return false;
101c537a943SNicolas Vasilache 
102edd9515bSthomasraoux   if (writeOp.mask() || writeOp.hasOutOfBoundsDim() ||
103edd9515bSthomasraoux       writeOp.getVectorType().getRank() != 2)
104edd9515bSthomasraoux     return false;
105edd9515bSthomasraoux   if (!getMemrefConstantHorizontalStride(writeOp.getShapedType()))
106edd9515bSthomasraoux     return false;
107edd9515bSthomasraoux   // TODO: Support transpose once it is added to GPU dialect ops.
108edd9515bSthomasraoux   if (!writeOp.permutation_map().isMinorIdentity())
109edd9515bSthomasraoux     return false;
110edd9515bSthomasraoux   return true;
111edd9515bSthomasraoux }
112edd9515bSthomasraoux 
1136413226dSthomasraoux /// Return true if the constant is a splat to a 2D vector so that it can be
1146413226dSthomasraoux /// converted to a MMA constant matrix op.
115a54f4eaeSMogball static bool constantSupportsMMAMatrixType(arith::ConstantOp constantOp) {
1166413226dSthomasraoux   auto vecType = constantOp.getType().dyn_cast<VectorType>();
1176413226dSthomasraoux   if (!vecType || vecType.getRank() != 2)
1186413226dSthomasraoux     return false;
119cfb72fd3SJacques Pienaar   return constantOp.getValue().isa<SplatElementsAttr>();
1206413226dSthomasraoux }
1216413226dSthomasraoux 
12243928419Sthomasraoux /// Return true if this is a broadcast from scalar to a 2D vector.
12343928419Sthomasraoux static bool broadcastSupportsMMAMatrixType(vector::BroadcastOp broadcastOp) {
12443928419Sthomasraoux   return broadcastOp.getVectorType().getRank() == 2 &&
12543928419Sthomasraoux          broadcastOp.source().getType().isa<FloatType>();
12643928419Sthomasraoux }
12743928419Sthomasraoux 
1287fbb0678Sthomasraoux /// Return the MMA elementwise enum associated with `op` if it is supported.
1297fbb0678Sthomasraoux /// Return `llvm::None` otherwise.
1307fbb0678Sthomasraoux static llvm::Optional<gpu::MMAElementwiseOp>
1317fbb0678Sthomasraoux convertElementwiseOpToMMA(Operation *op) {
1327fbb0678Sthomasraoux   if (isa<arith::AddFOp>(op))
1337fbb0678Sthomasraoux     return gpu::MMAElementwiseOp::ADDF;
1347fbb0678Sthomasraoux   if (isa<arith::MulFOp>(op))
1357fbb0678Sthomasraoux     return gpu::MMAElementwiseOp::MULF;
1369b1d90e8SAlexander Belyaev   if (isa<arith::MaxFOp>(op))
1377fbb0678Sthomasraoux     return gpu::MMAElementwiseOp::MAXF;
1389b1d90e8SAlexander Belyaev   if (isa<arith::MinFOp>(op))
1397fbb0678Sthomasraoux     return gpu::MMAElementwiseOp::MINF;
140e7969240SThomas Raoux   if (isa<arith::DivFOp>(op))
141e7969240SThomas Raoux     return gpu::MMAElementwiseOp::DIVF;
1427fbb0678Sthomasraoux   return llvm::None;
1437fbb0678Sthomasraoux }
1447fbb0678Sthomasraoux 
1457fbb0678Sthomasraoux /// Return true if the op is supported as elementwise op on MMAMatrix type.
1467fbb0678Sthomasraoux static bool elementwiseSupportsMMAMatrixType(Operation *op) {
1477fbb0678Sthomasraoux   return convertElementwiseOpToMMA(op).hasValue();
1487fbb0678Sthomasraoux }
1497fbb0678Sthomasraoux 
150edd9515bSthomasraoux static bool supportsMMaMatrixType(Operation *op) {
1511a865592Sthomasraoux   if (isa<scf::ForOp, scf::YieldOp>(op))
1521a865592Sthomasraoux     return true;
153edd9515bSthomasraoux   if (auto transferRead = dyn_cast<vector::TransferReadOp>(op))
154edd9515bSthomasraoux     return transferReadSupportsMMAMatrixType(transferRead);
155edd9515bSthomasraoux   if (auto transferWrite = dyn_cast<vector::TransferWriteOp>(op))
156edd9515bSthomasraoux     return transferWriteSupportsMMAMatrixType(transferWrite);
157edd9515bSthomasraoux   if (auto contract = dyn_cast<vector::ContractionOp>(op))
158edd9515bSthomasraoux     return contractSupportsMMAMatrixType(contract);
159a54f4eaeSMogball   if (auto constant = dyn_cast<arith::ConstantOp>(op))
1606413226dSthomasraoux     return constantSupportsMMAMatrixType(constant);
16143928419Sthomasraoux   if (auto broadcast = dyn_cast<vector::BroadcastOp>(op))
16243928419Sthomasraoux     return broadcastSupportsMMAMatrixType(broadcast);
1637fbb0678Sthomasraoux   return elementwiseSupportsMMAMatrixType(op);
164edd9515bSthomasraoux }
165edd9515bSthomasraoux 
166e7969240SThomas Raoux /// Return an unsorted slice handling scf.for region differently than
167e7969240SThomas Raoux /// `getSlice`. In scf.for we only want to include as part of the slice elements
168e7969240SThomas Raoux /// that are part of the use/def chain.
169e7969240SThomas Raoux static SetVector<Operation *> getSliceContract(Operation *op,
170e7969240SThomas Raoux                                                TransitiveFilter backwardFilter,
171e7969240SThomas Raoux                                                TransitiveFilter forwardFilter) {
172e7969240SThomas Raoux   SetVector<Operation *> slice;
173e7969240SThomas Raoux   slice.insert(op);
174e7969240SThomas Raoux   unsigned currentIndex = 0;
175e7969240SThomas Raoux   SetVector<Operation *> backwardSlice;
176e7969240SThomas Raoux   SetVector<Operation *> forwardSlice;
177e7969240SThomas Raoux   while (currentIndex != slice.size()) {
178e7969240SThomas Raoux     auto *currentOp = (slice)[currentIndex];
179e7969240SThomas Raoux     // Compute and insert the backwardSlice starting from currentOp.
180e7969240SThomas Raoux     backwardSlice.clear();
181e7969240SThomas Raoux     getBackwardSlice(currentOp, &backwardSlice, backwardFilter);
182e7969240SThomas Raoux     slice.insert(backwardSlice.begin(), backwardSlice.end());
183e7969240SThomas Raoux 
184e7969240SThomas Raoux     // Compute and insert the forwardSlice starting from currentOp.
185e7969240SThomas Raoux     forwardSlice.clear();
186e7969240SThomas Raoux     // Special case for ForOp, we don't want to include the whole region but
187e7969240SThomas Raoux     // only the value using the region arguments.
188e7969240SThomas Raoux     // TODO: We should refine this to only care about the region arguments being
189e7969240SThomas Raoux     // converted to matrix type.
190e7969240SThomas Raoux     if (auto forOp = dyn_cast<scf::ForOp>(currentOp)) {
191e7969240SThomas Raoux       for (Value forOpResult : forOp.getResults())
192e7969240SThomas Raoux         getForwardSlice(forOpResult, &forwardSlice, forwardFilter);
193e7969240SThomas Raoux       for (BlockArgument &arg : forOp.getRegionIterArgs())
194e7969240SThomas Raoux         getForwardSlice(arg, &forwardSlice, forwardFilter);
195e7969240SThomas Raoux     } else {
196e7969240SThomas Raoux       getForwardSlice(currentOp, &forwardSlice, forwardFilter);
197e7969240SThomas Raoux     }
198e7969240SThomas Raoux     slice.insert(forwardSlice.begin(), forwardSlice.end());
199e7969240SThomas Raoux     ++currentIndex;
200e7969240SThomas Raoux   }
201e7969240SThomas Raoux   return slice;
202e7969240SThomas Raoux }
203e7969240SThomas Raoux 
204edd9515bSthomasraoux // Analyze slice of operations based on convert op to figure out if the whole
205edd9515bSthomasraoux // slice can be converted to MMA operations.
206edd9515bSthomasraoux static SetVector<Operation *> getOpToConvert(mlir::Operation *op) {
207edd9515bSthomasraoux   auto hasVectorDest = [](Operation *op) {
20843928419Sthomasraoux     return llvm::any_of(op->getResultTypes(),
20943928419Sthomasraoux                         [](Type t) { return t.isa<VectorType>(); });
21043928419Sthomasraoux   };
21143928419Sthomasraoux   auto hasVectorSrc = [](Operation *op) {
21243928419Sthomasraoux     return llvm::any_of(op->getOperandTypes(),
213edd9515bSthomasraoux                         [](Type t) { return t.isa<VectorType>(); });
214edd9515bSthomasraoux   };
215edd9515bSthomasraoux   SetVector<Operation *> opToConvert;
216edd9515bSthomasraoux   op->walk([&](vector::ContractionOp contract) {
217edd9515bSthomasraoux     if (opToConvert.contains(contract.getOperation()))
218edd9515bSthomasraoux       return;
219edd9515bSthomasraoux     SetVector<Operation *> dependentOps =
220e7969240SThomas Raoux         getSliceContract(contract, hasVectorDest, hasVectorSrc);
221edd9515bSthomasraoux     // If any instruction cannot use MMA matrix type drop the whole
222e7969240SThomas Raoux     // chain. MMA matrix are stored in an opaque type so they cannot be used
223edd9515bSthomasraoux     // by all operations.
224edd9515bSthomasraoux     if (llvm::any_of(dependentOps,
225edd9515bSthomasraoux                      [](Operation *op) { return !supportsMMaMatrixType(op); }))
226edd9515bSthomasraoux       return;
227edd9515bSthomasraoux     opToConvert.insert(dependentOps.begin(), dependentOps.end());
228edd9515bSthomasraoux   });
229e7969240SThomas Raoux   // Sort the operations so that we can convert them in topological order.
230e7969240SThomas Raoux   return topologicalSort(opToConvert);
231edd9515bSthomasraoux }
232edd9515bSthomasraoux 
233edd9515bSthomasraoux namespace {
234edd9515bSthomasraoux // Transform contract into (m, k)x(k, n)x(m, n) form so that it can be converted
235edd9515bSthomasraoux // to MMA matmul.
236edd9515bSthomasraoux struct PrepareContractToGPUMMA
237edd9515bSthomasraoux     : public OpRewritePattern<vector::ContractionOp> {
238edd9515bSthomasraoux   using OpRewritePattern<vector::ContractionOp>::OpRewritePattern;
239edd9515bSthomasraoux 
240edd9515bSthomasraoux   LogicalResult matchAndRewrite(vector::ContractionOp op,
241edd9515bSthomasraoux                                 PatternRewriter &rewriter) const override {
242edd9515bSthomasraoux     Location loc = op.getLoc();
243edd9515bSthomasraoux     Value lhs = op.lhs(), rhs = op.rhs(), res = op.acc();
244edd9515bSthomasraoux 
245edd9515bSthomasraoux     // Set up the parallel/reduction structure in right form.
246edd9515bSthomasraoux     using MapList = ArrayRef<ArrayRef<AffineExpr>>;
247edd9515bSthomasraoux     auto infer = [](MapList m) { return AffineMap::inferFromExprList(m); };
248edd9515bSthomasraoux     AffineExpr m, n, k;
249edd9515bSthomasraoux     bindDims(rewriter.getContext(), m, n, k);
250edd9515bSthomasraoux     static constexpr std::array<int64_t, 2> perm = {1, 0};
251edd9515bSthomasraoux     auto iteratorTypes = op.iterator_types().getValue();
252edd9515bSthomasraoux     SmallVector<AffineMap, 4> maps = op.getIndexingMaps();
253edd9515bSthomasraoux     if (!(isParallelIterator(iteratorTypes[0]) &&
254edd9515bSthomasraoux           isParallelIterator(iteratorTypes[1]) &&
255edd9515bSthomasraoux           isReductionIterator(iteratorTypes[2])))
256edd9515bSthomasraoux       return failure();
257edd9515bSthomasraoux     //
258edd9515bSthomasraoux     // Two outer parallel, one inner reduction (matmat flavor).
259edd9515bSthomasraoux     //
260edd9515bSthomasraoux     if (maps == infer({{m, k}, {k, n}, {m, n}})) {
261edd9515bSthomasraoux       // This is the classical row-major matmul, nothing to do.
262edd9515bSthomasraoux       return failure();
263edd9515bSthomasraoux     }
264edd9515bSthomasraoux     if (maps == infer({{m, k}, {n, k}, {m, n}})) {
265edd9515bSthomasraoux       rhs = rewriter.create<vector::TransposeOp>(loc, rhs, perm);
266edd9515bSthomasraoux     } else if (maps == infer({{k, m}, {k, n}, {m, n}})) {
267edd9515bSthomasraoux       lhs = rewriter.create<vector::TransposeOp>(loc, lhs, perm);
268edd9515bSthomasraoux     } else if (maps == infer({{k, m}, {n, k}, {m, n}})) {
269edd9515bSthomasraoux       rhs = rewriter.create<vector::TransposeOp>(loc, rhs, perm);
270edd9515bSthomasraoux       lhs = rewriter.create<vector::TransposeOp>(loc, lhs, perm);
271edd9515bSthomasraoux     } else if (maps == infer({{m, k}, {k, n}, {n, m}})) {
272edd9515bSthomasraoux       std::swap(rhs, lhs);
273edd9515bSthomasraoux       rhs = rewriter.create<vector::TransposeOp>(loc, rhs, perm);
274edd9515bSthomasraoux       lhs = rewriter.create<vector::TransposeOp>(loc, lhs, perm);
275edd9515bSthomasraoux     } else if (maps == infer({{m, k}, {n, k}, {n, m}})) {
276edd9515bSthomasraoux       std::swap(rhs, lhs);
277edd9515bSthomasraoux       rhs = rewriter.create<vector::TransposeOp>(loc, rhs, perm);
278edd9515bSthomasraoux     } else if (maps == infer({{k, m}, {k, n}, {n, m}})) {
279edd9515bSthomasraoux       std::swap(lhs, rhs);
280edd9515bSthomasraoux       lhs = rewriter.create<vector::TransposeOp>(loc, lhs, perm);
281edd9515bSthomasraoux     } else if (maps == infer({{k, m}, {n, k}, {n, m}})) {
282edd9515bSthomasraoux       std::swap(lhs, rhs);
283edd9515bSthomasraoux     } else {
284edd9515bSthomasraoux       return failure();
285edd9515bSthomasraoux     }
286edd9515bSthomasraoux     rewriter.replaceOpWithNewOp<vector::ContractionOp>(
287edd9515bSthomasraoux         op, lhs, rhs, res,
288edd9515bSthomasraoux         rewriter.getAffineMapArrayAttr(infer({{m, k}, {k, n}, {m, n}})),
289edd9515bSthomasraoux         op.iterator_types());
290edd9515bSthomasraoux     return success();
291edd9515bSthomasraoux   }
292edd9515bSthomasraoux };
293edd9515bSthomasraoux 
294edd9515bSthomasraoux // Merge transpose op into the transfer read op. Transpose are not supported on
295edd9515bSthomasraoux // MMA types but MMA load can transpose the matrix when loading.
296edd9515bSthomasraoux struct CombineTransferReadOpTranspose final
297edd9515bSthomasraoux     : public OpRewritePattern<vector::TransposeOp> {
298edd9515bSthomasraoux   using OpRewritePattern<vector::TransposeOp>::OpRewritePattern;
299edd9515bSthomasraoux 
300edd9515bSthomasraoux   LogicalResult matchAndRewrite(vector::TransposeOp op,
301edd9515bSthomasraoux                                 PatternRewriter &rewriter) const override {
302edd9515bSthomasraoux     auto transferReadOp = op.vector().getDefiningOp<vector::TransferReadOp>();
303edd9515bSthomasraoux     if (!transferReadOp)
304edd9515bSthomasraoux       return failure();
305c537a943SNicolas Vasilache 
306c537a943SNicolas Vasilache     // TODO: support 0-d corner case.
307c537a943SNicolas Vasilache     if (transferReadOp.getTransferRank() == 0)
308c537a943SNicolas Vasilache       return failure();
309c537a943SNicolas Vasilache 
310edd9515bSthomasraoux     if (transferReadOp.mask() || transferReadOp.hasOutOfBoundsDim())
311edd9515bSthomasraoux       return failure();
312edd9515bSthomasraoux     SmallVector<int64_t, 2> perm;
313edd9515bSthomasraoux     op.getTransp(perm);
314edd9515bSthomasraoux     SmallVector<unsigned, 2> permU;
315edd9515bSthomasraoux     for (int64_t o : perm)
316edd9515bSthomasraoux       permU.push_back(unsigned(o));
317edd9515bSthomasraoux     AffineMap permutationMap =
318edd9515bSthomasraoux         AffineMap::getPermutationMap(permU, op.getContext());
319edd9515bSthomasraoux     AffineMap newMap = permutationMap.compose(transferReadOp.permutation_map());
320edd9515bSthomasraoux     rewriter.replaceOpWithNewOp<vector::TransferReadOp>(
321edd9515bSthomasraoux         op, op.getType(), transferReadOp.source(), transferReadOp.indices(),
322c537a943SNicolas Vasilache         AffineMapAttr::get(newMap), transferReadOp.padding(),
323c537a943SNicolas Vasilache         transferReadOp.mask(), transferReadOp.in_boundsAttr());
324edd9515bSthomasraoux     return success();
325edd9515bSthomasraoux   }
326edd9515bSthomasraoux };
327edd9515bSthomasraoux 
328edd9515bSthomasraoux } // namespace
329edd9515bSthomasraoux 
330edd9515bSthomasraoux // MMA types have different layout based on how they are used in matmul ops.
3316413226dSthomasraoux // Figure the right layout to use by looking at op uses.
332edd9515bSthomasraoux // TODO: Change the GPU dialect to abstract the layout at the this level and
333edd9515bSthomasraoux // only care about it during lowering to NVVM.
3346413226dSthomasraoux template <typename OpTy>
3356413226dSthomasraoux static const char *inferFragType(OpTy op) {
336edd9515bSthomasraoux   for (Operation *users : op->getUsers()) {
337edd9515bSthomasraoux     auto contract = dyn_cast<vector::ContractionOp>(users);
338edd9515bSthomasraoux     if (!contract)
339edd9515bSthomasraoux       continue;
340edd9515bSthomasraoux     if (contract.lhs() == op.getResult())
341edd9515bSthomasraoux       return "AOp";
342edd9515bSthomasraoux     if (contract.rhs() == op.getResult())
343edd9515bSthomasraoux       return "BOp";
344edd9515bSthomasraoux   }
345edd9515bSthomasraoux   return "COp";
346edd9515bSthomasraoux }
347edd9515bSthomasraoux 
348edd9515bSthomasraoux static void convertTransferReadOp(vector::TransferReadOp op,
349edd9515bSthomasraoux                                   llvm::DenseMap<Value, Value> &valueMapping) {
350c537a943SNicolas Vasilache   assert(op.getTransferRank() > 0 && "unexpected 0-d transfer");
351edd9515bSthomasraoux   assert(transferReadSupportsMMAMatrixType(op));
352edd9515bSthomasraoux   Optional<int64_t> stride =
353edd9515bSthomasraoux       getMemrefConstantHorizontalStride(op.getShapedType());
354e7969240SThomas Raoux   AffineMap map = op.permutation_map();
355e7969240SThomas Raoux   // Handle broadcast by setting the stride to 0.
356e7969240SThomas Raoux   if (map.getResult(0).isa<AffineConstantExpr>()) {
357e7969240SThomas Raoux     assert(map.getResult(0).cast<AffineConstantExpr>().getValue() == 0);
358e7969240SThomas Raoux     stride = 0;
359e7969240SThomas Raoux   }
360edd9515bSthomasraoux   assert(stride);
361edd9515bSthomasraoux   const char *fragType = inferFragType(op);
362edd9515bSthomasraoux   gpu::MMAMatrixType type =
363edd9515bSthomasraoux       gpu::MMAMatrixType::get(op.getVectorType().getShape(),
364edd9515bSthomasraoux                               op.getVectorType().getElementType(), fragType);
365edd9515bSthomasraoux   OpBuilder b(op);
366edd9515bSthomasraoux   Value load = b.create<gpu::SubgroupMmaLoadMatrixOp>(
367edd9515bSthomasraoux       op.getLoc(), type, op.source(), op.indices(), b.getIndexAttr(*stride));
368edd9515bSthomasraoux   valueMapping[op.getResult()] = load;
369edd9515bSthomasraoux }
370edd9515bSthomasraoux 
371edd9515bSthomasraoux static void convertTransferWriteOp(vector::TransferWriteOp op,
372edd9515bSthomasraoux                                    llvm::DenseMap<Value, Value> &valueMapping) {
373edd9515bSthomasraoux   assert(transferWriteSupportsMMAMatrixType(op));
374edd9515bSthomasraoux   Optional<int64_t> stride =
375edd9515bSthomasraoux       getMemrefConstantHorizontalStride(op.getShapedType());
376edd9515bSthomasraoux   assert(stride);
377edd9515bSthomasraoux   OpBuilder b(op);
378edd9515bSthomasraoux   Value matrix = valueMapping.find(op.vector())->second;
379edd9515bSthomasraoux   b.create<gpu::SubgroupMmaStoreMatrixOp>(
380edd9515bSthomasraoux       op.getLoc(), matrix, op.source(), op.indices(), b.getIndexAttr(*stride));
381edd9515bSthomasraoux   op.erase();
382edd9515bSthomasraoux }
383edd9515bSthomasraoux 
384edd9515bSthomasraoux static void convertContractOp(vector::ContractionOp op,
385edd9515bSthomasraoux                               llvm::DenseMap<Value, Value> &valueMapping) {
386edd9515bSthomasraoux   OpBuilder b(op);
387edd9515bSthomasraoux   Value opA = valueMapping.find(op.lhs())->second;
388edd9515bSthomasraoux   Value opB = valueMapping.find(op.rhs())->second;
389edd9515bSthomasraoux   Value opC = valueMapping.find(op.acc())->second;
390edd9515bSthomasraoux   Value matmul = b.create<gpu::SubgroupMmaComputeOp>(op.getLoc(), opC.getType(),
391edd9515bSthomasraoux                                                      opA, opB, opC);
392edd9515bSthomasraoux   valueMapping[op.getResult()] = matmul;
393edd9515bSthomasraoux }
394edd9515bSthomasraoux 
3956413226dSthomasraoux /// Convert a 2D splat ConstantOp to a SubgroupMmaConstantMatrix op.
396a54f4eaeSMogball static void convertConstantOp(arith::ConstantOp op,
3976413226dSthomasraoux                               llvm::DenseMap<Value, Value> &valueMapping) {
3986413226dSthomasraoux   assert(constantSupportsMMAMatrixType(op));
3996413226dSthomasraoux   OpBuilder b(op);
400937e40a8SRiver Riddle   Attribute splat =
401937e40a8SRiver Riddle       op.getValue().cast<SplatElementsAttr>().getSplatValue<Attribute>();
4026413226dSthomasraoux   auto scalarConstant =
403a54f4eaeSMogball       b.create<arith::ConstantOp>(op.getLoc(), splat.getType(), splat);
4046413226dSthomasraoux   const char *fragType = inferFragType(op);
4056413226dSthomasraoux   auto vecType = op.getType().cast<VectorType>();
4066413226dSthomasraoux   gpu::MMAMatrixType type = gpu::MMAMatrixType::get(
4076413226dSthomasraoux       vecType.getShape(), vecType.getElementType(), llvm::StringRef(fragType));
4086413226dSthomasraoux   auto matrix = b.create<gpu::SubgroupMmaConstantMatrixOp>(op.getLoc(), type,
4096413226dSthomasraoux                                                            scalarConstant);
4106413226dSthomasraoux   valueMapping[op.getResult()] = matrix;
4116413226dSthomasraoux }
4126413226dSthomasraoux 
41343928419Sthomasraoux /// Convert a vector.broadcast from scalar to a SubgroupMmaConstantMatrix op.
41443928419Sthomasraoux static void convertBroadcastOp(vector::BroadcastOp op,
41543928419Sthomasraoux                                llvm::DenseMap<Value, Value> &valueMapping) {
41643928419Sthomasraoux   assert(broadcastSupportsMMAMatrixType(op));
41743928419Sthomasraoux   OpBuilder b(op);
41843928419Sthomasraoux   const char *fragType = inferFragType(op);
41943928419Sthomasraoux   auto vecType = op.getVectorType();
42043928419Sthomasraoux   gpu::MMAMatrixType type = gpu::MMAMatrixType::get(
42143928419Sthomasraoux       vecType.getShape(), vecType.getElementType(), llvm::StringRef(fragType));
42243928419Sthomasraoux   auto matrix = b.create<gpu::SubgroupMmaConstantMatrixOp>(op.getLoc(), type,
42343928419Sthomasraoux                                                            op.source());
42443928419Sthomasraoux   valueMapping[op.getResult()] = matrix;
42543928419Sthomasraoux }
42643928419Sthomasraoux 
4271a865592Sthomasraoux // Replace ForOp with a new ForOp with extra operands. The YieldOp is not
4281a865592Sthomasraoux // updated and needs to be updated separatly for the loop to be correct.
4291a865592Sthomasraoux static scf::ForOp replaceForOpWithNewSignature(OpBuilder &b, scf::ForOp loop,
4301a865592Sthomasraoux                                                ValueRange newIterOperands) {
4311a865592Sthomasraoux   // Create a new loop before the existing one, with the extra operands.
4321a865592Sthomasraoux   OpBuilder::InsertionGuard g(b);
4331a865592Sthomasraoux   b.setInsertionPoint(loop);
4341a865592Sthomasraoux   auto operands = llvm::to_vector<4>(loop.getIterOperands());
4351a865592Sthomasraoux   operands.append(newIterOperands.begin(), newIterOperands.end());
4361a865592Sthomasraoux   scf::ForOp newLoop =
437c0342a2dSJacques Pienaar       b.create<scf::ForOp>(loop.getLoc(), loop.getLowerBound(),
438c0342a2dSJacques Pienaar                            loop.getUpperBound(), loop.getStep(), operands);
4391a865592Sthomasraoux   newLoop.getBody()->erase();
4401a865592Sthomasraoux   newLoop.getLoopBody().getBlocks().splice(
4411a865592Sthomasraoux       newLoop.getLoopBody().getBlocks().begin(),
4421a865592Sthomasraoux       loop.getLoopBody().getBlocks());
443e084679fSRiver Riddle   for (Value operand : newIterOperands)
444e084679fSRiver Riddle     newLoop.getBody()->addArgument(operand.getType(), operand.getLoc());
4451a865592Sthomasraoux 
4461a865592Sthomasraoux   for (auto it : llvm::zip(loop.getResults(), newLoop.getResults().take_front(
4471a865592Sthomasraoux                                                   loop.getNumResults())))
4481a865592Sthomasraoux     std::get<0>(it).replaceAllUsesWith(std::get<1>(it));
4491a865592Sthomasraoux   loop.erase();
4501a865592Sthomasraoux   return newLoop;
4511a865592Sthomasraoux }
4521a865592Sthomasraoux 
4531a865592Sthomasraoux static void convertForOp(scf::ForOp op,
4541a865592Sthomasraoux                          llvm::DenseMap<Value, Value> &valueMapping) {
4551a865592Sthomasraoux   SmallVector<Value> newOperands;
4561a865592Sthomasraoux   SmallVector<std::pair<size_t, size_t>> argMapping;
457e4853be2SMehdi Amini   for (const auto &operand : llvm::enumerate(op.getIterOperands())) {
4581a865592Sthomasraoux     auto it = valueMapping.find(operand.value());
4591a865592Sthomasraoux     if (it == valueMapping.end())
4601a865592Sthomasraoux       continue;
4611a865592Sthomasraoux     argMapping.push_back(std::make_pair(
4621a865592Sthomasraoux         operand.index(), op.getNumIterOperands() + newOperands.size()));
4631a865592Sthomasraoux     newOperands.push_back(it->second);
4641a865592Sthomasraoux   }
4651a865592Sthomasraoux   OpBuilder b(op);
4661a865592Sthomasraoux   scf::ForOp newForOp = replaceForOpWithNewSignature(b, op, newOperands);
4671a865592Sthomasraoux   Block &loopBody = *newForOp.getBody();
4681a865592Sthomasraoux   for (auto mapping : argMapping) {
4691a865592Sthomasraoux     valueMapping[newForOp.getResult(mapping.first)] =
4701a865592Sthomasraoux         newForOp.getResult(mapping.second);
4711a865592Sthomasraoux     valueMapping[loopBody.getArgument(mapping.first +
4721a865592Sthomasraoux                                       newForOp.getNumInductionVars())] =
4731a865592Sthomasraoux         loopBody.getArgument(mapping.second + newForOp.getNumInductionVars());
4741a865592Sthomasraoux   }
4751a865592Sthomasraoux }
4761a865592Sthomasraoux 
4771a865592Sthomasraoux static void convertYieldOp(scf::YieldOp op,
4781a865592Sthomasraoux                            llvm::DenseMap<Value, Value> &valueMapping) {
4791a865592Sthomasraoux   OpBuilder b(op);
4801a865592Sthomasraoux   auto loop = cast<scf::ForOp>(op->getParentOp());
4811a865592Sthomasraoux   auto yieldOperands = llvm::to_vector<4>(op.getOperands());
482e4853be2SMehdi Amini   for (const auto &operand : llvm::enumerate(op.getOperands())) {
4831a865592Sthomasraoux     auto it = valueMapping.find(operand.value());
4841a865592Sthomasraoux     if (it == valueMapping.end())
4851a865592Sthomasraoux       continue;
4861a865592Sthomasraoux     // Replace the yield of old value with the for op argument to make it easier
4871a865592Sthomasraoux     // to remove the dead code.
4881a865592Sthomasraoux     yieldOperands[operand.index()] = loop.getIterOperands()[operand.index()];
4891a865592Sthomasraoux     yieldOperands.push_back(it->second);
4901a865592Sthomasraoux   }
4911a865592Sthomasraoux   b.create<scf::YieldOp>(op.getLoc(), yieldOperands);
4921a865592Sthomasraoux   op.erase();
4931a865592Sthomasraoux }
4941a865592Sthomasraoux 
4957fbb0678Sthomasraoux /// Convert an elementwise op to the equivalent elementwise op on MMA matrix.
4967fbb0678Sthomasraoux static void convertElementwiseOp(Operation *op, gpu::MMAElementwiseOp opType,
4977fbb0678Sthomasraoux                                  llvm::DenseMap<Value, Value> &valueMapping) {
4987fbb0678Sthomasraoux   OpBuilder b(op);
4997fbb0678Sthomasraoux   SmallVector<Value> matrixOperands;
5007fbb0678Sthomasraoux   for (Value operand : op->getOperands())
5017fbb0678Sthomasraoux     matrixOperands.push_back(valueMapping.find(operand)->second);
5027fbb0678Sthomasraoux   Value newOp = b.create<gpu::SubgroupMmaElementwiseOp>(
5037fbb0678Sthomasraoux       op->getLoc(), matrixOperands[0].getType(), matrixOperands, opType);
5047fbb0678Sthomasraoux   valueMapping[op->getResult(0)] = newOp;
5057fbb0678Sthomasraoux }
5067fbb0678Sthomasraoux 
50747f175b0SRiver Riddle void mlir::populatePrepareVectorToMMAPatterns(RewritePatternSet &patterns) {
508edd9515bSthomasraoux   patterns.add<PrepareContractToGPUMMA, CombineTransferReadOpTranspose>(
509edd9515bSthomasraoux       patterns.getContext());
510edd9515bSthomasraoux }
511edd9515bSthomasraoux 
51247f175b0SRiver Riddle void mlir::convertVectorToMMAOps(Operation *rootOp) {
51347f175b0SRiver Riddle   SetVector<Operation *> ops = getOpToConvert(rootOp);
514edd9515bSthomasraoux   llvm::DenseMap<Value, Value> valueMapping;
515edd9515bSthomasraoux   for (Operation *op : ops) {
516edd9515bSthomasraoux     if (auto transferRead = dyn_cast<vector::TransferReadOp>(op)) {
517edd9515bSthomasraoux       convertTransferReadOp(transferRead, valueMapping);
518edd9515bSthomasraoux     } else if (auto transferWrite = dyn_cast<vector::TransferWriteOp>(op)) {
519edd9515bSthomasraoux       convertTransferWriteOp(transferWrite, valueMapping);
520edd9515bSthomasraoux     } else if (auto contractOp = dyn_cast<vector::ContractionOp>(op)) {
521edd9515bSthomasraoux       convertContractOp(contractOp, valueMapping);
522a54f4eaeSMogball     } else if (auto constantOp = dyn_cast<arith::ConstantOp>(op)) {
5236413226dSthomasraoux       convertConstantOp(constantOp, valueMapping);
52443928419Sthomasraoux     } else if (auto broadcastOp = dyn_cast<vector::BroadcastOp>(op)) {
52543928419Sthomasraoux       convertBroadcastOp(broadcastOp, valueMapping);
5261a865592Sthomasraoux     } else if (auto forOp = dyn_cast<scf::ForOp>(op)) {
5271a865592Sthomasraoux       convertForOp(forOp, valueMapping);
5281a865592Sthomasraoux     } else if (auto yiledOp = dyn_cast<scf::YieldOp>(op)) {
5291a865592Sthomasraoux       convertYieldOp(yiledOp, valueMapping);
5307fbb0678Sthomasraoux     } else if (auto elementwiseType = convertElementwiseOpToMMA(op)) {
5317fbb0678Sthomasraoux       convertElementwiseOp(op, *elementwiseType, valueMapping);
532edd9515bSthomasraoux     }
533edd9515bSthomasraoux   }
534edd9515bSthomasraoux }
535edd9515bSthomasraoux 
536edd9515bSthomasraoux namespace {
537edd9515bSthomasraoux 
538edd9515bSthomasraoux struct ConvertVectorToGPUPass
539edd9515bSthomasraoux     : public ConvertVectorToGPUBase<ConvertVectorToGPUPass> {
54041574554SRiver Riddle   void runOnOperation() override {
54147f175b0SRiver Riddle     RewritePatternSet patterns(&getContext());
542edd9515bSthomasraoux     populatePrepareVectorToMMAPatterns(patterns);
54341574554SRiver Riddle     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
544edd9515bSthomasraoux 
54541574554SRiver Riddle     convertVectorToMMAOps(getOperation());
546edd9515bSthomasraoux   }
547edd9515bSthomasraoux };
548edd9515bSthomasraoux 
549edd9515bSthomasraoux } // namespace
550edd9515bSthomasraoux 
551edd9515bSthomasraoux std::unique_ptr<Pass> mlir::createConvertVectorToGPUPass() {
552edd9515bSthomasraoux   return std::make_unique<ConvertVectorToGPUPass>();
553edd9515bSthomasraoux }
554