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"
24*99ef9eebSMatthias Springer #include "mlir/Dialect/Vector/IR/VectorOps.h"
25*99ef9eebSMatthias 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;
68edd9515bSthomasraoux   if (failed(getStridesAndOffset(memrefType, strides, offset)))
69edd9515bSthomasraoux     return llvm::None;
70a57ccad5SThomas Raoux   int64_t stride = strides[strides.size() - 2];
71a57ccad5SThomas Raoux   if (stride == ShapedType::kDynamicStrideOrOffset)
72edd9515bSthomasraoux     return llvm::None;
73a57ccad5SThomas Raoux   return stride;
74edd9515bSthomasraoux }
75edd9515bSthomasraoux 
76edd9515bSthomasraoux // Return true if the transfer op can be converted to a MMA matrix load.
77edd9515bSthomasraoux static bool transferReadSupportsMMAMatrixType(vector::TransferReadOp readOp) {
78edd9515bSthomasraoux   if (readOp.mask() || readOp.hasOutOfBoundsDim() ||
79edd9515bSthomasraoux       readOp.getVectorType().getRank() != 2)
80edd9515bSthomasraoux     return false;
81edd9515bSthomasraoux   if (!getMemrefConstantHorizontalStride(readOp.getShapedType()))
82edd9515bSthomasraoux     return false;
83e7969240SThomas Raoux   AffineMap map = readOp.permutation_map();
84e7969240SThomas Raoux   OpBuilder b(readOp.getContext());
85e7969240SThomas Raoux   AffineExpr innerDim = b.getAffineDimExpr(map.getNumDims() - 1);
86e7969240SThomas Raoux   AffineExpr zero = b.getAffineConstantExpr(0);
87e7969240SThomas Raoux   auto broadcastInnerDim = AffineMap::get(map.getNumDims(), 0, {zero, innerDim},
88e7969240SThomas Raoux                                           readOp.getContext());
89edd9515bSthomasraoux   // TODO: Support transpose once it is added to GPU dialect ops.
90e7969240SThomas Raoux   // For now we only support (d0, d1) -> (d0, d1) and (d0, d1) -> (0, d1).
916786d7e4SMehdi Amini   return !(!map.isMinorIdentity() && map != broadcastInnerDim);
92edd9515bSthomasraoux }
93edd9515bSthomasraoux 
94edd9515bSthomasraoux // Return true if the transfer op can be converted to a MMA matrix store.
95edd9515bSthomasraoux static bool
96edd9515bSthomasraoux transferWriteSupportsMMAMatrixType(vector::TransferWriteOp writeOp) {
97c537a943SNicolas Vasilache   // TODO: support 0-d corner case.
98c537a943SNicolas Vasilache   if (writeOp.getTransferRank() == 0)
99c537a943SNicolas Vasilache     return false;
100c537a943SNicolas Vasilache 
101edd9515bSthomasraoux   if (writeOp.mask() || writeOp.hasOutOfBoundsDim() ||
102edd9515bSthomasraoux       writeOp.getVectorType().getRank() != 2)
103edd9515bSthomasraoux     return false;
104edd9515bSthomasraoux   if (!getMemrefConstantHorizontalStride(writeOp.getShapedType()))
105edd9515bSthomasraoux     return false;
106edd9515bSthomasraoux   // TODO: Support transpose once it is added to GPU dialect ops.
107edd9515bSthomasraoux   if (!writeOp.permutation_map().isMinorIdentity())
108edd9515bSthomasraoux     return false;
109edd9515bSthomasraoux   return true;
110edd9515bSthomasraoux }
111edd9515bSthomasraoux 
1126413226dSthomasraoux /// Return true if the constant is a splat to a 2D vector so that it can be
1136413226dSthomasraoux /// converted to a MMA constant matrix op.
114a54f4eaeSMogball static bool constantSupportsMMAMatrixType(arith::ConstantOp constantOp) {
1156413226dSthomasraoux   auto vecType = constantOp.getType().dyn_cast<VectorType>();
1166413226dSthomasraoux   if (!vecType || vecType.getRank() != 2)
1176413226dSthomasraoux     return false;
118cfb72fd3SJacques Pienaar   return constantOp.getValue().isa<SplatElementsAttr>();
1196413226dSthomasraoux }
1206413226dSthomasraoux 
12143928419Sthomasraoux /// Return true if this is a broadcast from scalar to a 2D vector.
12243928419Sthomasraoux static bool broadcastSupportsMMAMatrixType(vector::BroadcastOp broadcastOp) {
12343928419Sthomasraoux   return broadcastOp.getVectorType().getRank() == 2 &&
12443928419Sthomasraoux          broadcastOp.source().getType().isa<FloatType>();
12543928419Sthomasraoux }
12643928419Sthomasraoux 
1277fbb0678Sthomasraoux /// Return the MMA elementwise enum associated with `op` if it is supported.
1287fbb0678Sthomasraoux /// Return `llvm::None` otherwise.
1297fbb0678Sthomasraoux static llvm::Optional<gpu::MMAElementwiseOp>
1307fbb0678Sthomasraoux convertElementwiseOpToMMA(Operation *op) {
1317fbb0678Sthomasraoux   if (isa<arith::AddFOp>(op))
1327fbb0678Sthomasraoux     return gpu::MMAElementwiseOp::ADDF;
1337fbb0678Sthomasraoux   if (isa<arith::MulFOp>(op))
1347fbb0678Sthomasraoux     return gpu::MMAElementwiseOp::MULF;
1359b1d90e8SAlexander Belyaev   if (isa<arith::MaxFOp>(op))
1367fbb0678Sthomasraoux     return gpu::MMAElementwiseOp::MAXF;
1379b1d90e8SAlexander Belyaev   if (isa<arith::MinFOp>(op))
1387fbb0678Sthomasraoux     return gpu::MMAElementwiseOp::MINF;
139e7969240SThomas Raoux   if (isa<arith::DivFOp>(op))
140e7969240SThomas Raoux     return gpu::MMAElementwiseOp::DIVF;
1417fbb0678Sthomasraoux   return llvm::None;
1427fbb0678Sthomasraoux }
1437fbb0678Sthomasraoux 
1447fbb0678Sthomasraoux /// Return true if the op is supported as elementwise op on MMAMatrix type.
1457fbb0678Sthomasraoux static bool elementwiseSupportsMMAMatrixType(Operation *op) {
1467fbb0678Sthomasraoux   return convertElementwiseOpToMMA(op).hasValue();
1477fbb0678Sthomasraoux }
1487fbb0678Sthomasraoux 
149edd9515bSthomasraoux static bool supportsMMaMatrixType(Operation *op) {
1501a865592Sthomasraoux   if (isa<scf::ForOp, scf::YieldOp>(op))
1511a865592Sthomasraoux     return true;
152edd9515bSthomasraoux   if (auto transferRead = dyn_cast<vector::TransferReadOp>(op))
153edd9515bSthomasraoux     return transferReadSupportsMMAMatrixType(transferRead);
154edd9515bSthomasraoux   if (auto transferWrite = dyn_cast<vector::TransferWriteOp>(op))
155edd9515bSthomasraoux     return transferWriteSupportsMMAMatrixType(transferWrite);
156edd9515bSthomasraoux   if (auto contract = dyn_cast<vector::ContractionOp>(op))
157edd9515bSthomasraoux     return contractSupportsMMAMatrixType(contract);
158a54f4eaeSMogball   if (auto constant = dyn_cast<arith::ConstantOp>(op))
1596413226dSthomasraoux     return constantSupportsMMAMatrixType(constant);
16043928419Sthomasraoux   if (auto broadcast = dyn_cast<vector::BroadcastOp>(op))
16143928419Sthomasraoux     return broadcastSupportsMMAMatrixType(broadcast);
1627fbb0678Sthomasraoux   return elementwiseSupportsMMAMatrixType(op);
163edd9515bSthomasraoux }
164edd9515bSthomasraoux 
165e7969240SThomas Raoux /// Return an unsorted slice handling scf.for region differently than
166e7969240SThomas Raoux /// `getSlice`. In scf.for we only want to include as part of the slice elements
167e7969240SThomas Raoux /// that are part of the use/def chain.
168e7969240SThomas Raoux static SetVector<Operation *> getSliceContract(Operation *op,
169e7969240SThomas Raoux                                                TransitiveFilter backwardFilter,
170e7969240SThomas Raoux                                                TransitiveFilter forwardFilter) {
171e7969240SThomas Raoux   SetVector<Operation *> slice;
172e7969240SThomas Raoux   slice.insert(op);
173e7969240SThomas Raoux   unsigned currentIndex = 0;
174e7969240SThomas Raoux   SetVector<Operation *> backwardSlice;
175e7969240SThomas Raoux   SetVector<Operation *> forwardSlice;
176e7969240SThomas Raoux   while (currentIndex != slice.size()) {
177e7969240SThomas Raoux     auto *currentOp = (slice)[currentIndex];
178e7969240SThomas Raoux     // Compute and insert the backwardSlice starting from currentOp.
179e7969240SThomas Raoux     backwardSlice.clear();
180e7969240SThomas Raoux     getBackwardSlice(currentOp, &backwardSlice, backwardFilter);
181e7969240SThomas Raoux     slice.insert(backwardSlice.begin(), backwardSlice.end());
182e7969240SThomas Raoux 
183e7969240SThomas Raoux     // Compute and insert the forwardSlice starting from currentOp.
184e7969240SThomas Raoux     forwardSlice.clear();
185e7969240SThomas Raoux     // Special case for ForOp, we don't want to include the whole region but
186e7969240SThomas Raoux     // only the value using the region arguments.
187e7969240SThomas Raoux     // TODO: We should refine this to only care about the region arguments being
188e7969240SThomas Raoux     // converted to matrix type.
189e7969240SThomas Raoux     if (auto forOp = dyn_cast<scf::ForOp>(currentOp)) {
190e7969240SThomas Raoux       for (Value forOpResult : forOp.getResults())
191e7969240SThomas Raoux         getForwardSlice(forOpResult, &forwardSlice, forwardFilter);
192e7969240SThomas Raoux       for (BlockArgument &arg : forOp.getRegionIterArgs())
193e7969240SThomas Raoux         getForwardSlice(arg, &forwardSlice, forwardFilter);
194e7969240SThomas Raoux     } else {
195e7969240SThomas Raoux       getForwardSlice(currentOp, &forwardSlice, forwardFilter);
196e7969240SThomas Raoux     }
197e7969240SThomas Raoux     slice.insert(forwardSlice.begin(), forwardSlice.end());
198e7969240SThomas Raoux     ++currentIndex;
199e7969240SThomas Raoux   }
200e7969240SThomas Raoux   return slice;
201e7969240SThomas Raoux }
202e7969240SThomas Raoux 
203edd9515bSthomasraoux // Analyze slice of operations based on convert op to figure out if the whole
204edd9515bSthomasraoux // slice can be converted to MMA operations.
205edd9515bSthomasraoux static SetVector<Operation *> getOpToConvert(mlir::Operation *op) {
206edd9515bSthomasraoux   auto hasVectorDest = [](Operation *op) {
20743928419Sthomasraoux     return llvm::any_of(op->getResultTypes(),
20843928419Sthomasraoux                         [](Type t) { return t.isa<VectorType>(); });
20943928419Sthomasraoux   };
21043928419Sthomasraoux   auto hasVectorSrc = [](Operation *op) {
21143928419Sthomasraoux     return llvm::any_of(op->getOperandTypes(),
212edd9515bSthomasraoux                         [](Type t) { return t.isa<VectorType>(); });
213edd9515bSthomasraoux   };
214edd9515bSthomasraoux   SetVector<Operation *> opToConvert;
215edd9515bSthomasraoux   op->walk([&](vector::ContractionOp contract) {
216edd9515bSthomasraoux     if (opToConvert.contains(contract.getOperation()))
217edd9515bSthomasraoux       return;
218edd9515bSthomasraoux     SetVector<Operation *> dependentOps =
219e7969240SThomas Raoux         getSliceContract(contract, hasVectorDest, hasVectorSrc);
220edd9515bSthomasraoux     // If any instruction cannot use MMA matrix type drop the whole
221e7969240SThomas Raoux     // chain. MMA matrix are stored in an opaque type so they cannot be used
222edd9515bSthomasraoux     // by all operations.
223edd9515bSthomasraoux     if (llvm::any_of(dependentOps,
224edd9515bSthomasraoux                      [](Operation *op) { return !supportsMMaMatrixType(op); }))
225edd9515bSthomasraoux       return;
226edd9515bSthomasraoux     opToConvert.insert(dependentOps.begin(), dependentOps.end());
227edd9515bSthomasraoux   });
228e7969240SThomas Raoux   // Sort the operations so that we can convert them in topological order.
229e7969240SThomas Raoux   return topologicalSort(opToConvert);
230edd9515bSthomasraoux }
231edd9515bSthomasraoux 
232edd9515bSthomasraoux namespace {
233edd9515bSthomasraoux // Transform contract into (m, k)x(k, n)x(m, n) form so that it can be converted
234edd9515bSthomasraoux // to MMA matmul.
235edd9515bSthomasraoux struct PrepareContractToGPUMMA
236edd9515bSthomasraoux     : public OpRewritePattern<vector::ContractionOp> {
237edd9515bSthomasraoux   using OpRewritePattern<vector::ContractionOp>::OpRewritePattern;
238edd9515bSthomasraoux 
239edd9515bSthomasraoux   LogicalResult matchAndRewrite(vector::ContractionOp op,
240edd9515bSthomasraoux                                 PatternRewriter &rewriter) const override {
241edd9515bSthomasraoux     Location loc = op.getLoc();
242edd9515bSthomasraoux     Value lhs = op.lhs(), rhs = op.rhs(), res = op.acc();
243edd9515bSthomasraoux 
244edd9515bSthomasraoux     // Set up the parallel/reduction structure in right form.
245edd9515bSthomasraoux     using MapList = ArrayRef<ArrayRef<AffineExpr>>;
246edd9515bSthomasraoux     auto infer = [](MapList m) { return AffineMap::inferFromExprList(m); };
247edd9515bSthomasraoux     AffineExpr m, n, k;
248edd9515bSthomasraoux     bindDims(rewriter.getContext(), m, n, k);
249edd9515bSthomasraoux     static constexpr std::array<int64_t, 2> perm = {1, 0};
250edd9515bSthomasraoux     auto iteratorTypes = op.iterator_types().getValue();
251edd9515bSthomasraoux     SmallVector<AffineMap, 4> maps = op.getIndexingMaps();
252edd9515bSthomasraoux     if (!(isParallelIterator(iteratorTypes[0]) &&
253edd9515bSthomasraoux           isParallelIterator(iteratorTypes[1]) &&
254edd9515bSthomasraoux           isReductionIterator(iteratorTypes[2])))
255edd9515bSthomasraoux       return failure();
256edd9515bSthomasraoux     //
257edd9515bSthomasraoux     // Two outer parallel, one inner reduction (matmat flavor).
258edd9515bSthomasraoux     //
259edd9515bSthomasraoux     if (maps == infer({{m, k}, {k, n}, {m, n}})) {
260edd9515bSthomasraoux       // This is the classical row-major matmul, nothing to do.
261edd9515bSthomasraoux       return failure();
262edd9515bSthomasraoux     }
263edd9515bSthomasraoux     if (maps == infer({{m, k}, {n, k}, {m, n}})) {
264edd9515bSthomasraoux       rhs = rewriter.create<vector::TransposeOp>(loc, rhs, perm);
265edd9515bSthomasraoux     } else if (maps == infer({{k, m}, {k, n}, {m, n}})) {
266edd9515bSthomasraoux       lhs = rewriter.create<vector::TransposeOp>(loc, lhs, perm);
267edd9515bSthomasraoux     } else if (maps == infer({{k, m}, {n, k}, {m, n}})) {
268edd9515bSthomasraoux       rhs = rewriter.create<vector::TransposeOp>(loc, rhs, perm);
269edd9515bSthomasraoux       lhs = rewriter.create<vector::TransposeOp>(loc, lhs, perm);
270edd9515bSthomasraoux     } else if (maps == infer({{m, k}, {k, n}, {n, m}})) {
271edd9515bSthomasraoux       std::swap(rhs, lhs);
272edd9515bSthomasraoux       rhs = rewriter.create<vector::TransposeOp>(loc, rhs, perm);
273edd9515bSthomasraoux       lhs = rewriter.create<vector::TransposeOp>(loc, lhs, perm);
274edd9515bSthomasraoux     } else if (maps == infer({{m, k}, {n, k}, {n, m}})) {
275edd9515bSthomasraoux       std::swap(rhs, lhs);
276edd9515bSthomasraoux       rhs = rewriter.create<vector::TransposeOp>(loc, rhs, perm);
277edd9515bSthomasraoux     } else if (maps == infer({{k, m}, {k, n}, {n, m}})) {
278edd9515bSthomasraoux       std::swap(lhs, rhs);
279edd9515bSthomasraoux       lhs = rewriter.create<vector::TransposeOp>(loc, lhs, perm);
280edd9515bSthomasraoux     } else if (maps == infer({{k, m}, {n, k}, {n, m}})) {
281edd9515bSthomasraoux       std::swap(lhs, rhs);
282edd9515bSthomasraoux     } else {
283edd9515bSthomasraoux       return failure();
284edd9515bSthomasraoux     }
285edd9515bSthomasraoux     rewriter.replaceOpWithNewOp<vector::ContractionOp>(
286edd9515bSthomasraoux         op, lhs, rhs, res,
287edd9515bSthomasraoux         rewriter.getAffineMapArrayAttr(infer({{m, k}, {k, n}, {m, n}})),
288edd9515bSthomasraoux         op.iterator_types());
289edd9515bSthomasraoux     return success();
290edd9515bSthomasraoux   }
291edd9515bSthomasraoux };
292edd9515bSthomasraoux 
293edd9515bSthomasraoux // Merge transpose op into the transfer read op. Transpose are not supported on
294edd9515bSthomasraoux // MMA types but MMA load can transpose the matrix when loading.
295edd9515bSthomasraoux struct CombineTransferReadOpTranspose final
296edd9515bSthomasraoux     : public OpRewritePattern<vector::TransposeOp> {
297edd9515bSthomasraoux   using OpRewritePattern<vector::TransposeOp>::OpRewritePattern;
298edd9515bSthomasraoux 
299edd9515bSthomasraoux   LogicalResult matchAndRewrite(vector::TransposeOp op,
300edd9515bSthomasraoux                                 PatternRewriter &rewriter) const override {
301edd9515bSthomasraoux     auto transferReadOp = op.vector().getDefiningOp<vector::TransferReadOp>();
302edd9515bSthomasraoux     if (!transferReadOp)
303edd9515bSthomasraoux       return failure();
304c537a943SNicolas Vasilache 
305c537a943SNicolas Vasilache     // TODO: support 0-d corner case.
306c537a943SNicolas Vasilache     if (transferReadOp.getTransferRank() == 0)
307c537a943SNicolas Vasilache       return failure();
308c537a943SNicolas Vasilache 
309edd9515bSthomasraoux     if (transferReadOp.mask() || transferReadOp.hasOutOfBoundsDim())
310edd9515bSthomasraoux       return failure();
311edd9515bSthomasraoux     SmallVector<int64_t, 2> perm;
312edd9515bSthomasraoux     op.getTransp(perm);
313edd9515bSthomasraoux     SmallVector<unsigned, 2> permU;
314edd9515bSthomasraoux     for (int64_t o : perm)
315edd9515bSthomasraoux       permU.push_back(unsigned(o));
316edd9515bSthomasraoux     AffineMap permutationMap =
317edd9515bSthomasraoux         AffineMap::getPermutationMap(permU, op.getContext());
318edd9515bSthomasraoux     AffineMap newMap = permutationMap.compose(transferReadOp.permutation_map());
319edd9515bSthomasraoux     rewriter.replaceOpWithNewOp<vector::TransferReadOp>(
320edd9515bSthomasraoux         op, op.getType(), transferReadOp.source(), transferReadOp.indices(),
321c537a943SNicolas Vasilache         AffineMapAttr::get(newMap), transferReadOp.padding(),
322c537a943SNicolas Vasilache         transferReadOp.mask(), transferReadOp.in_boundsAttr());
323edd9515bSthomasraoux     return success();
324edd9515bSthomasraoux   }
325edd9515bSthomasraoux };
326edd9515bSthomasraoux 
327edd9515bSthomasraoux } // namespace
328edd9515bSthomasraoux 
329edd9515bSthomasraoux // MMA types have different layout based on how they are used in matmul ops.
3306413226dSthomasraoux // Figure the right layout to use by looking at op uses.
331edd9515bSthomasraoux // TODO: Change the GPU dialect to abstract the layout at the this level and
332edd9515bSthomasraoux // only care about it during lowering to NVVM.
3336413226dSthomasraoux template <typename OpTy>
3346413226dSthomasraoux static const char *inferFragType(OpTy op) {
335edd9515bSthomasraoux   for (Operation *users : op->getUsers()) {
336edd9515bSthomasraoux     auto contract = dyn_cast<vector::ContractionOp>(users);
337edd9515bSthomasraoux     if (!contract)
338edd9515bSthomasraoux       continue;
339edd9515bSthomasraoux     if (contract.lhs() == op.getResult())
340edd9515bSthomasraoux       return "AOp";
341edd9515bSthomasraoux     if (contract.rhs() == op.getResult())
342edd9515bSthomasraoux       return "BOp";
343edd9515bSthomasraoux   }
344edd9515bSthomasraoux   return "COp";
345edd9515bSthomasraoux }
346edd9515bSthomasraoux 
347edd9515bSthomasraoux static void convertTransferReadOp(vector::TransferReadOp op,
348edd9515bSthomasraoux                                   llvm::DenseMap<Value, Value> &valueMapping) {
349c537a943SNicolas Vasilache   assert(op.getTransferRank() > 0 && "unexpected 0-d transfer");
350edd9515bSthomasraoux   assert(transferReadSupportsMMAMatrixType(op));
351edd9515bSthomasraoux   Optional<int64_t> stride =
352edd9515bSthomasraoux       getMemrefConstantHorizontalStride(op.getShapedType());
353e7969240SThomas Raoux   AffineMap map = op.permutation_map();
354e7969240SThomas Raoux   // Handle broadcast by setting the stride to 0.
355e7969240SThomas Raoux   if (map.getResult(0).isa<AffineConstantExpr>()) {
356e7969240SThomas Raoux     assert(map.getResult(0).cast<AffineConstantExpr>().getValue() == 0);
357e7969240SThomas Raoux     stride = 0;
358e7969240SThomas Raoux   }
359edd9515bSthomasraoux   assert(stride);
360edd9515bSthomasraoux   const char *fragType = inferFragType(op);
361edd9515bSthomasraoux   gpu::MMAMatrixType type =
362edd9515bSthomasraoux       gpu::MMAMatrixType::get(op.getVectorType().getShape(),
363edd9515bSthomasraoux                               op.getVectorType().getElementType(), fragType);
364edd9515bSthomasraoux   OpBuilder b(op);
365edd9515bSthomasraoux   Value load = b.create<gpu::SubgroupMmaLoadMatrixOp>(
366edd9515bSthomasraoux       op.getLoc(), type, op.source(), op.indices(), b.getIndexAttr(*stride));
367edd9515bSthomasraoux   valueMapping[op.getResult()] = load;
368edd9515bSthomasraoux }
369edd9515bSthomasraoux 
370edd9515bSthomasraoux static void convertTransferWriteOp(vector::TransferWriteOp op,
371edd9515bSthomasraoux                                    llvm::DenseMap<Value, Value> &valueMapping) {
372edd9515bSthomasraoux   assert(transferWriteSupportsMMAMatrixType(op));
373edd9515bSthomasraoux   Optional<int64_t> stride =
374edd9515bSthomasraoux       getMemrefConstantHorizontalStride(op.getShapedType());
375edd9515bSthomasraoux   assert(stride);
376edd9515bSthomasraoux   OpBuilder b(op);
377edd9515bSthomasraoux   Value matrix = valueMapping.find(op.vector())->second;
378edd9515bSthomasraoux   b.create<gpu::SubgroupMmaStoreMatrixOp>(
379edd9515bSthomasraoux       op.getLoc(), matrix, op.source(), op.indices(), b.getIndexAttr(*stride));
380edd9515bSthomasraoux   op.erase();
381edd9515bSthomasraoux }
382edd9515bSthomasraoux 
383edd9515bSthomasraoux static void convertContractOp(vector::ContractionOp op,
384edd9515bSthomasraoux                               llvm::DenseMap<Value, Value> &valueMapping) {
385edd9515bSthomasraoux   OpBuilder b(op);
386edd9515bSthomasraoux   Value opA = valueMapping.find(op.lhs())->second;
387edd9515bSthomasraoux   Value opB = valueMapping.find(op.rhs())->second;
388edd9515bSthomasraoux   Value opC = valueMapping.find(op.acc())->second;
389edd9515bSthomasraoux   Value matmul = b.create<gpu::SubgroupMmaComputeOp>(op.getLoc(), opC.getType(),
390edd9515bSthomasraoux                                                      opA, opB, opC);
391edd9515bSthomasraoux   valueMapping[op.getResult()] = matmul;
392edd9515bSthomasraoux }
393edd9515bSthomasraoux 
3946413226dSthomasraoux /// Convert a 2D splat ConstantOp to a SubgroupMmaConstantMatrix op.
395a54f4eaeSMogball static void convertConstantOp(arith::ConstantOp op,
3966413226dSthomasraoux                               llvm::DenseMap<Value, Value> &valueMapping) {
3976413226dSthomasraoux   assert(constantSupportsMMAMatrixType(op));
3986413226dSthomasraoux   OpBuilder b(op);
399937e40a8SRiver Riddle   Attribute splat =
400937e40a8SRiver Riddle       op.getValue().cast<SplatElementsAttr>().getSplatValue<Attribute>();
4016413226dSthomasraoux   auto scalarConstant =
402a54f4eaeSMogball       b.create<arith::ConstantOp>(op.getLoc(), splat.getType(), splat);
4036413226dSthomasraoux   const char *fragType = inferFragType(op);
4046413226dSthomasraoux   auto vecType = op.getType().cast<VectorType>();
4056413226dSthomasraoux   gpu::MMAMatrixType type = gpu::MMAMatrixType::get(
4066413226dSthomasraoux       vecType.getShape(), vecType.getElementType(), llvm::StringRef(fragType));
4076413226dSthomasraoux   auto matrix = b.create<gpu::SubgroupMmaConstantMatrixOp>(op.getLoc(), type,
4086413226dSthomasraoux                                                            scalarConstant);
4096413226dSthomasraoux   valueMapping[op.getResult()] = matrix;
4106413226dSthomasraoux }
4116413226dSthomasraoux 
41243928419Sthomasraoux /// Convert a vector.broadcast from scalar to a SubgroupMmaConstantMatrix op.
41343928419Sthomasraoux static void convertBroadcastOp(vector::BroadcastOp op,
41443928419Sthomasraoux                                llvm::DenseMap<Value, Value> &valueMapping) {
41543928419Sthomasraoux   assert(broadcastSupportsMMAMatrixType(op));
41643928419Sthomasraoux   OpBuilder b(op);
41743928419Sthomasraoux   const char *fragType = inferFragType(op);
41843928419Sthomasraoux   auto vecType = op.getVectorType();
41943928419Sthomasraoux   gpu::MMAMatrixType type = gpu::MMAMatrixType::get(
42043928419Sthomasraoux       vecType.getShape(), vecType.getElementType(), llvm::StringRef(fragType));
42143928419Sthomasraoux   auto matrix = b.create<gpu::SubgroupMmaConstantMatrixOp>(op.getLoc(), type,
42243928419Sthomasraoux                                                            op.source());
42343928419Sthomasraoux   valueMapping[op.getResult()] = matrix;
42443928419Sthomasraoux }
42543928419Sthomasraoux 
4261a865592Sthomasraoux // Replace ForOp with a new ForOp with extra operands. The YieldOp is not
4271a865592Sthomasraoux // updated and needs to be updated separatly for the loop to be correct.
4281a865592Sthomasraoux static scf::ForOp replaceForOpWithNewSignature(OpBuilder &b, scf::ForOp loop,
4291a865592Sthomasraoux                                                ValueRange newIterOperands) {
4301a865592Sthomasraoux   // Create a new loop before the existing one, with the extra operands.
4311a865592Sthomasraoux   OpBuilder::InsertionGuard g(b);
4321a865592Sthomasraoux   b.setInsertionPoint(loop);
4331a865592Sthomasraoux   auto operands = llvm::to_vector<4>(loop.getIterOperands());
4341a865592Sthomasraoux   operands.append(newIterOperands.begin(), newIterOperands.end());
4351a865592Sthomasraoux   scf::ForOp newLoop =
436c0342a2dSJacques Pienaar       b.create<scf::ForOp>(loop.getLoc(), loop.getLowerBound(),
437c0342a2dSJacques Pienaar                            loop.getUpperBound(), loop.getStep(), operands);
4381a865592Sthomasraoux   newLoop.getBody()->erase();
4391a865592Sthomasraoux   newLoop.getLoopBody().getBlocks().splice(
4401a865592Sthomasraoux       newLoop.getLoopBody().getBlocks().begin(),
4411a865592Sthomasraoux       loop.getLoopBody().getBlocks());
442e084679fSRiver Riddle   for (Value operand : newIterOperands)
443e084679fSRiver Riddle     newLoop.getBody()->addArgument(operand.getType(), operand.getLoc());
4441a865592Sthomasraoux 
4451a865592Sthomasraoux   for (auto it : llvm::zip(loop.getResults(), newLoop.getResults().take_front(
4461a865592Sthomasraoux                                                   loop.getNumResults())))
4471a865592Sthomasraoux     std::get<0>(it).replaceAllUsesWith(std::get<1>(it));
4481a865592Sthomasraoux   loop.erase();
4491a865592Sthomasraoux   return newLoop;
4501a865592Sthomasraoux }
4511a865592Sthomasraoux 
4521a865592Sthomasraoux static void convertForOp(scf::ForOp op,
4531a865592Sthomasraoux                          llvm::DenseMap<Value, Value> &valueMapping) {
4541a865592Sthomasraoux   SmallVector<Value> newOperands;
4551a865592Sthomasraoux   SmallVector<std::pair<size_t, size_t>> argMapping;
456e4853be2SMehdi Amini   for (const auto &operand : llvm::enumerate(op.getIterOperands())) {
4571a865592Sthomasraoux     auto it = valueMapping.find(operand.value());
4581a865592Sthomasraoux     if (it == valueMapping.end())
4591a865592Sthomasraoux       continue;
4601a865592Sthomasraoux     argMapping.push_back(std::make_pair(
4611a865592Sthomasraoux         operand.index(), op.getNumIterOperands() + newOperands.size()));
4621a865592Sthomasraoux     newOperands.push_back(it->second);
4631a865592Sthomasraoux   }
4641a865592Sthomasraoux   OpBuilder b(op);
4651a865592Sthomasraoux   scf::ForOp newForOp = replaceForOpWithNewSignature(b, op, newOperands);
4661a865592Sthomasraoux   Block &loopBody = *newForOp.getBody();
4671a865592Sthomasraoux   for (auto mapping : argMapping) {
4681a865592Sthomasraoux     valueMapping[newForOp.getResult(mapping.first)] =
4691a865592Sthomasraoux         newForOp.getResult(mapping.second);
4701a865592Sthomasraoux     valueMapping[loopBody.getArgument(mapping.first +
4711a865592Sthomasraoux                                       newForOp.getNumInductionVars())] =
4721a865592Sthomasraoux         loopBody.getArgument(mapping.second + newForOp.getNumInductionVars());
4731a865592Sthomasraoux   }
4741a865592Sthomasraoux }
4751a865592Sthomasraoux 
4761a865592Sthomasraoux static void convertYieldOp(scf::YieldOp op,
4771a865592Sthomasraoux                            llvm::DenseMap<Value, Value> &valueMapping) {
4781a865592Sthomasraoux   OpBuilder b(op);
4791a865592Sthomasraoux   auto loop = cast<scf::ForOp>(op->getParentOp());
4801a865592Sthomasraoux   auto yieldOperands = llvm::to_vector<4>(op.getOperands());
481e4853be2SMehdi Amini   for (const auto &operand : llvm::enumerate(op.getOperands())) {
4821a865592Sthomasraoux     auto it = valueMapping.find(operand.value());
4831a865592Sthomasraoux     if (it == valueMapping.end())
4841a865592Sthomasraoux       continue;
4851a865592Sthomasraoux     // Replace the yield of old value with the for op argument to make it easier
4861a865592Sthomasraoux     // to remove the dead code.
4871a865592Sthomasraoux     yieldOperands[operand.index()] = loop.getIterOperands()[operand.index()];
4881a865592Sthomasraoux     yieldOperands.push_back(it->second);
4891a865592Sthomasraoux   }
4901a865592Sthomasraoux   b.create<scf::YieldOp>(op.getLoc(), yieldOperands);
4911a865592Sthomasraoux   op.erase();
4921a865592Sthomasraoux }
4931a865592Sthomasraoux 
4947fbb0678Sthomasraoux /// Convert an elementwise op to the equivalent elementwise op on MMA matrix.
4957fbb0678Sthomasraoux static void convertElementwiseOp(Operation *op, gpu::MMAElementwiseOp opType,
4967fbb0678Sthomasraoux                                  llvm::DenseMap<Value, Value> &valueMapping) {
4977fbb0678Sthomasraoux   OpBuilder b(op);
4987fbb0678Sthomasraoux   SmallVector<Value> matrixOperands;
4997fbb0678Sthomasraoux   for (Value operand : op->getOperands())
5007fbb0678Sthomasraoux     matrixOperands.push_back(valueMapping.find(operand)->second);
5017fbb0678Sthomasraoux   Value newOp = b.create<gpu::SubgroupMmaElementwiseOp>(
5027fbb0678Sthomasraoux       op->getLoc(), matrixOperands[0].getType(), matrixOperands, opType);
5037fbb0678Sthomasraoux   valueMapping[op->getResult(0)] = newOp;
5047fbb0678Sthomasraoux }
5057fbb0678Sthomasraoux 
506edd9515bSthomasraoux namespace mlir {
507edd9515bSthomasraoux 
508edd9515bSthomasraoux void populatePrepareVectorToMMAPatterns(RewritePatternSet &patterns) {
509edd9515bSthomasraoux   patterns.add<PrepareContractToGPUMMA, CombineTransferReadOpTranspose>(
510edd9515bSthomasraoux       patterns.getContext());
511edd9515bSthomasraoux }
512edd9515bSthomasraoux 
513edd9515bSthomasraoux void convertVectorToMMAOps(FuncOp funcOp) {
514edd9515bSthomasraoux   SetVector<Operation *> ops = getOpToConvert(funcOp);
515edd9515bSthomasraoux   llvm::DenseMap<Value, Value> valueMapping;
516edd9515bSthomasraoux   for (Operation *op : ops) {
517edd9515bSthomasraoux     if (auto transferRead = dyn_cast<vector::TransferReadOp>(op)) {
518edd9515bSthomasraoux       convertTransferReadOp(transferRead, valueMapping);
519edd9515bSthomasraoux     } else if (auto transferWrite = dyn_cast<vector::TransferWriteOp>(op)) {
520edd9515bSthomasraoux       convertTransferWriteOp(transferWrite, valueMapping);
521edd9515bSthomasraoux     } else if (auto contractOp = dyn_cast<vector::ContractionOp>(op)) {
522edd9515bSthomasraoux       convertContractOp(contractOp, valueMapping);
523a54f4eaeSMogball     } else if (auto constantOp = dyn_cast<arith::ConstantOp>(op)) {
5246413226dSthomasraoux       convertConstantOp(constantOp, valueMapping);
52543928419Sthomasraoux     } else if (auto broadcastOp = dyn_cast<vector::BroadcastOp>(op)) {
52643928419Sthomasraoux       convertBroadcastOp(broadcastOp, valueMapping);
5271a865592Sthomasraoux     } else if (auto forOp = dyn_cast<scf::ForOp>(op)) {
5281a865592Sthomasraoux       convertForOp(forOp, valueMapping);
5291a865592Sthomasraoux     } else if (auto yiledOp = dyn_cast<scf::YieldOp>(op)) {
5301a865592Sthomasraoux       convertYieldOp(yiledOp, valueMapping);
5317fbb0678Sthomasraoux     } else if (auto elementwiseType = convertElementwiseOpToMMA(op)) {
5327fbb0678Sthomasraoux       convertElementwiseOp(op, *elementwiseType, valueMapping);
533edd9515bSthomasraoux     }
534edd9515bSthomasraoux   }
535edd9515bSthomasraoux }
536edd9515bSthomasraoux 
537edd9515bSthomasraoux } // namespace mlir
538edd9515bSthomasraoux namespace {
539edd9515bSthomasraoux 
540edd9515bSthomasraoux struct ConvertVectorToGPUPass
541edd9515bSthomasraoux     : public ConvertVectorToGPUBase<ConvertVectorToGPUPass> {
54241574554SRiver Riddle   void runOnOperation() override {
54341574554SRiver Riddle     RewritePatternSet patterns(getOperation().getContext());
544edd9515bSthomasraoux     populatePrepareVectorToMMAPatterns(patterns);
54541574554SRiver Riddle     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
546edd9515bSthomasraoux 
54741574554SRiver Riddle     convertVectorToMMAOps(getOperation());
548edd9515bSthomasraoux   }
549edd9515bSthomasraoux };
550edd9515bSthomasraoux 
551edd9515bSthomasraoux } // namespace
552edd9515bSthomasraoux 
553edd9515bSthomasraoux std::unique_ptr<Pass> mlir::createConvertVectorToGPUPass() {
554edd9515bSthomasraoux   return std::make_unique<ConvertVectorToGPUPass>();
555edd9515bSthomasraoux }
556