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