15c0c51a9SNicolas Vasilache //===- VectorToLLVM.cpp - Conversion from Vector to the LLVM dialect ------===// 25c0c51a9SNicolas Vasilache // 330857107SMehdi Amini // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. 456222a06SMehdi Amini // See https://llvm.org/LICENSE.txt for license information. 556222a06SMehdi Amini // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception 65c0c51a9SNicolas Vasilache // 756222a06SMehdi Amini //===----------------------------------------------------------------------===// 85c0c51a9SNicolas Vasilache 965678d93SNicolas Vasilache #include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h" 10870c1fd4SAlex Zinenko 111834ad4aSRiver Riddle #include "../PassDetail.h" 125c0c51a9SNicolas Vasilache #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h" 135c0c51a9SNicolas Vasilache #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h" 145c0c51a9SNicolas Vasilache #include "mlir/Dialect/LLVMIR/LLVMDialect.h" 1569d757c0SRob Suderman #include "mlir/Dialect/StandardOps/IR/Ops.h" 164d60f47bSRob Suderman #include "mlir/Dialect/Vector/VectorOps.h" 178345b86dSNicolas Vasilache #include "mlir/IR/AffineMap.h" 185c0c51a9SNicolas Vasilache #include "mlir/IR/Attributes.h" 195c0c51a9SNicolas Vasilache #include "mlir/IR/Builders.h" 205c0c51a9SNicolas Vasilache #include "mlir/IR/MLIRContext.h" 215c0c51a9SNicolas Vasilache #include "mlir/IR/Module.h" 225c0c51a9SNicolas Vasilache #include "mlir/IR/Operation.h" 235c0c51a9SNicolas Vasilache #include "mlir/IR/PatternMatch.h" 245c0c51a9SNicolas Vasilache #include "mlir/IR/StandardTypes.h" 255c0c51a9SNicolas Vasilache #include "mlir/IR/Types.h" 265c0c51a9SNicolas Vasilache #include "mlir/Transforms/DialectConversion.h" 275c0c51a9SNicolas Vasilache #include "mlir/Transforms/Passes.h" 285c0c51a9SNicolas Vasilache #include "llvm/IR/DerivedTypes.h" 295c0c51a9SNicolas Vasilache #include "llvm/IR/Module.h" 305c0c51a9SNicolas Vasilache #include "llvm/IR/Type.h" 315c0c51a9SNicolas Vasilache #include "llvm/Support/Allocator.h" 325c0c51a9SNicolas Vasilache #include "llvm/Support/ErrorHandling.h" 335c0c51a9SNicolas Vasilache 345c0c51a9SNicolas Vasilache using namespace mlir; 3565678d93SNicolas Vasilache using namespace mlir::vector; 365c0c51a9SNicolas Vasilache 375c0c51a9SNicolas Vasilache template <typename T> 385c0c51a9SNicolas Vasilache static LLVM::LLVMType getPtrToElementType(T containerType, 390f04384dSAlex Zinenko LLVMTypeConverter &typeConverter) { 400f04384dSAlex Zinenko return typeConverter.convertType(containerType.getElementType()) 415c0c51a9SNicolas Vasilache .template cast<LLVM::LLVMType>() 425c0c51a9SNicolas Vasilache .getPointerTo(); 435c0c51a9SNicolas Vasilache } 445c0c51a9SNicolas Vasilache 459826fe5cSAart Bik // Helper to reduce vector type by one rank at front. 469826fe5cSAart Bik static VectorType reducedVectorTypeFront(VectorType tp) { 479826fe5cSAart Bik assert((tp.getRank() > 1) && "unlowerable vector type"); 489826fe5cSAart Bik return VectorType::get(tp.getShape().drop_front(), tp.getElementType()); 499826fe5cSAart Bik } 509826fe5cSAart Bik 519826fe5cSAart Bik // Helper to reduce vector type by *all* but one rank at back. 529826fe5cSAart Bik static VectorType reducedVectorTypeBack(VectorType tp) { 539826fe5cSAart Bik assert((tp.getRank() > 1) && "unlowerable vector type"); 549826fe5cSAart Bik return VectorType::get(tp.getShape().take_back(), tp.getElementType()); 559826fe5cSAart Bik } 569826fe5cSAart Bik 571c81adf3SAart Bik // Helper that picks the proper sequence for inserting. 58e62a6956SRiver Riddle static Value insertOne(ConversionPatternRewriter &rewriter, 590f04384dSAlex Zinenko LLVMTypeConverter &typeConverter, Location loc, 600f04384dSAlex Zinenko Value val1, Value val2, Type llvmType, int64_t rank, 610f04384dSAlex Zinenko int64_t pos) { 621c81adf3SAart Bik if (rank == 1) { 631c81adf3SAart Bik auto idxType = rewriter.getIndexType(); 641c81adf3SAart Bik auto constant = rewriter.create<LLVM::ConstantOp>( 650f04384dSAlex Zinenko loc, typeConverter.convertType(idxType), 661c81adf3SAart Bik rewriter.getIntegerAttr(idxType, pos)); 671c81adf3SAart Bik return rewriter.create<LLVM::InsertElementOp>(loc, llvmType, val1, val2, 681c81adf3SAart Bik constant); 691c81adf3SAart Bik } 701c81adf3SAart Bik return rewriter.create<LLVM::InsertValueOp>(loc, llvmType, val1, val2, 711c81adf3SAart Bik rewriter.getI64ArrayAttr(pos)); 721c81adf3SAart Bik } 731c81adf3SAart Bik 742d515e49SNicolas Vasilache // Helper that picks the proper sequence for inserting. 752d515e49SNicolas Vasilache static Value insertOne(PatternRewriter &rewriter, Location loc, Value from, 762d515e49SNicolas Vasilache Value into, int64_t offset) { 772d515e49SNicolas Vasilache auto vectorType = into.getType().cast<VectorType>(); 782d515e49SNicolas Vasilache if (vectorType.getRank() > 1) 792d515e49SNicolas Vasilache return rewriter.create<InsertOp>(loc, from, into, offset); 802d515e49SNicolas Vasilache return rewriter.create<vector::InsertElementOp>( 812d515e49SNicolas Vasilache loc, vectorType, from, into, 822d515e49SNicolas Vasilache rewriter.create<ConstantIndexOp>(loc, offset)); 832d515e49SNicolas Vasilache } 842d515e49SNicolas Vasilache 851c81adf3SAart Bik // Helper that picks the proper sequence for extracting. 86e62a6956SRiver Riddle static Value extractOne(ConversionPatternRewriter &rewriter, 870f04384dSAlex Zinenko LLVMTypeConverter &typeConverter, Location loc, 880f04384dSAlex Zinenko Value val, Type llvmType, int64_t rank, int64_t pos) { 891c81adf3SAart Bik if (rank == 1) { 901c81adf3SAart Bik auto idxType = rewriter.getIndexType(); 911c81adf3SAart Bik auto constant = rewriter.create<LLVM::ConstantOp>( 920f04384dSAlex Zinenko loc, typeConverter.convertType(idxType), 931c81adf3SAart Bik rewriter.getIntegerAttr(idxType, pos)); 941c81adf3SAart Bik return rewriter.create<LLVM::ExtractElementOp>(loc, llvmType, val, 951c81adf3SAart Bik constant); 961c81adf3SAart Bik } 971c81adf3SAart Bik return rewriter.create<LLVM::ExtractValueOp>(loc, llvmType, val, 981c81adf3SAart Bik rewriter.getI64ArrayAttr(pos)); 991c81adf3SAart Bik } 1001c81adf3SAart Bik 1012d515e49SNicolas Vasilache // Helper that picks the proper sequence for extracting. 1022d515e49SNicolas Vasilache static Value extractOne(PatternRewriter &rewriter, Location loc, Value vector, 1032d515e49SNicolas Vasilache int64_t offset) { 1042d515e49SNicolas Vasilache auto vectorType = vector.getType().cast<VectorType>(); 1052d515e49SNicolas Vasilache if (vectorType.getRank() > 1) 1062d515e49SNicolas Vasilache return rewriter.create<ExtractOp>(loc, vector, offset); 1072d515e49SNicolas Vasilache return rewriter.create<vector::ExtractElementOp>( 1082d515e49SNicolas Vasilache loc, vectorType.getElementType(), vector, 1092d515e49SNicolas Vasilache rewriter.create<ConstantIndexOp>(loc, offset)); 1102d515e49SNicolas Vasilache } 1112d515e49SNicolas Vasilache 1122d515e49SNicolas Vasilache // Helper that returns a subset of `arrayAttr` as a vector of int64_t. 1139db53a18SRiver Riddle // TODO: Better support for attribute subtype forwarding + slicing. 1142d515e49SNicolas Vasilache static SmallVector<int64_t, 4> getI64SubArray(ArrayAttr arrayAttr, 1152d515e49SNicolas Vasilache unsigned dropFront = 0, 1162d515e49SNicolas Vasilache unsigned dropBack = 0) { 1172d515e49SNicolas Vasilache assert(arrayAttr.size() > dropFront + dropBack && "Out of bounds"); 1182d515e49SNicolas Vasilache auto range = arrayAttr.getAsRange<IntegerAttr>(); 1192d515e49SNicolas Vasilache SmallVector<int64_t, 4> res; 1202d515e49SNicolas Vasilache res.reserve(arrayAttr.size() - dropFront - dropBack); 1212d515e49SNicolas Vasilache for (auto it = range.begin() + dropFront, eit = range.end() - dropBack; 1222d515e49SNicolas Vasilache it != eit; ++it) 1232d515e49SNicolas Vasilache res.push_back((*it).getValue().getSExtValue()); 1242d515e49SNicolas Vasilache return res; 1252d515e49SNicolas Vasilache } 1262d515e49SNicolas Vasilache 1275f9e0466SNicolas Vasilache template <typename TransferOp> 1285f9e0466SNicolas Vasilache LogicalResult getVectorTransferAlignment(LLVMTypeConverter &typeConverter, 1295f9e0466SNicolas Vasilache TransferOp xferOp, unsigned &align) { 1305f9e0466SNicolas Vasilache Type elementTy = 1315f9e0466SNicolas Vasilache typeConverter.convertType(xferOp.getMemRefType().getElementType()); 1325f9e0466SNicolas Vasilache if (!elementTy) 1335f9e0466SNicolas Vasilache return failure(); 1345f9e0466SNicolas Vasilache 1355f9e0466SNicolas Vasilache auto dataLayout = typeConverter.getDialect()->getLLVMModule().getDataLayout(); 1365f9e0466SNicolas Vasilache align = dataLayout.getPrefTypeAlignment( 1375f9e0466SNicolas Vasilache elementTy.cast<LLVM::LLVMType>().getUnderlyingType()); 1385f9e0466SNicolas Vasilache return success(); 1395f9e0466SNicolas Vasilache } 1405f9e0466SNicolas Vasilache 1415f9e0466SNicolas Vasilache static LogicalResult 1425f9e0466SNicolas Vasilache replaceTransferOpWithLoadOrStore(ConversionPatternRewriter &rewriter, 1435f9e0466SNicolas Vasilache LLVMTypeConverter &typeConverter, Location loc, 1445f9e0466SNicolas Vasilache TransferReadOp xferOp, 1455f9e0466SNicolas Vasilache ArrayRef<Value> operands, Value dataPtr) { 146*affbc0cdSNicolas Vasilache unsigned align; 147*affbc0cdSNicolas Vasilache if (failed(getVectorTransferAlignment(typeConverter, xferOp, align))) 148*affbc0cdSNicolas Vasilache return failure(); 149*affbc0cdSNicolas Vasilache rewriter.replaceOpWithNewOp<LLVM::LoadOp>(xferOp, dataPtr, align); 1505f9e0466SNicolas Vasilache return success(); 1515f9e0466SNicolas Vasilache } 1525f9e0466SNicolas Vasilache 1535f9e0466SNicolas Vasilache static LogicalResult 1545f9e0466SNicolas Vasilache replaceTransferOpWithMasked(ConversionPatternRewriter &rewriter, 1555f9e0466SNicolas Vasilache LLVMTypeConverter &typeConverter, Location loc, 1565f9e0466SNicolas Vasilache TransferReadOp xferOp, ArrayRef<Value> operands, 1575f9e0466SNicolas Vasilache Value dataPtr, Value mask) { 1585f9e0466SNicolas Vasilache auto toLLVMTy = [&](Type t) { return typeConverter.convertType(t); }; 1595f9e0466SNicolas Vasilache VectorType fillType = xferOp.getVectorType(); 1605f9e0466SNicolas Vasilache Value fill = rewriter.create<SplatOp>(loc, fillType, xferOp.padding()); 1615f9e0466SNicolas Vasilache fill = rewriter.create<LLVM::DialectCastOp>(loc, toLLVMTy(fillType), fill); 1625f9e0466SNicolas Vasilache 1635f9e0466SNicolas Vasilache Type vecTy = typeConverter.convertType(xferOp.getVectorType()); 1645f9e0466SNicolas Vasilache if (!vecTy) 1655f9e0466SNicolas Vasilache return failure(); 1665f9e0466SNicolas Vasilache 1675f9e0466SNicolas Vasilache unsigned align; 1685f9e0466SNicolas Vasilache if (failed(getVectorTransferAlignment(typeConverter, xferOp, align))) 1695f9e0466SNicolas Vasilache return failure(); 1705f9e0466SNicolas Vasilache 1715f9e0466SNicolas Vasilache rewriter.replaceOpWithNewOp<LLVM::MaskedLoadOp>( 1725f9e0466SNicolas Vasilache xferOp, vecTy, dataPtr, mask, ValueRange{fill}, 1735f9e0466SNicolas Vasilache rewriter.getI32IntegerAttr(align)); 1745f9e0466SNicolas Vasilache return success(); 1755f9e0466SNicolas Vasilache } 1765f9e0466SNicolas Vasilache 1775f9e0466SNicolas Vasilache static LogicalResult 1785f9e0466SNicolas Vasilache replaceTransferOpWithLoadOrStore(ConversionPatternRewriter &rewriter, 1795f9e0466SNicolas Vasilache LLVMTypeConverter &typeConverter, Location loc, 1805f9e0466SNicolas Vasilache TransferWriteOp xferOp, 1815f9e0466SNicolas Vasilache ArrayRef<Value> operands, Value dataPtr) { 182*affbc0cdSNicolas Vasilache unsigned align; 183*affbc0cdSNicolas Vasilache if (failed(getVectorTransferAlignment(typeConverter, xferOp, align))) 184*affbc0cdSNicolas Vasilache return failure(); 1852d2c73c5SJacques Pienaar auto adaptor = TransferWriteOpAdaptor(operands); 186*affbc0cdSNicolas Vasilache rewriter.replaceOpWithNewOp<LLVM::StoreOp>(xferOp, adaptor.vector(), dataPtr, 187*affbc0cdSNicolas Vasilache align); 1885f9e0466SNicolas Vasilache return success(); 1895f9e0466SNicolas Vasilache } 1905f9e0466SNicolas Vasilache 1915f9e0466SNicolas Vasilache static LogicalResult 1925f9e0466SNicolas Vasilache replaceTransferOpWithMasked(ConversionPatternRewriter &rewriter, 1935f9e0466SNicolas Vasilache LLVMTypeConverter &typeConverter, Location loc, 1945f9e0466SNicolas Vasilache TransferWriteOp xferOp, ArrayRef<Value> operands, 1955f9e0466SNicolas Vasilache Value dataPtr, Value mask) { 1965f9e0466SNicolas Vasilache unsigned align; 1975f9e0466SNicolas Vasilache if (failed(getVectorTransferAlignment(typeConverter, xferOp, align))) 1985f9e0466SNicolas Vasilache return failure(); 1995f9e0466SNicolas Vasilache 2002d2c73c5SJacques Pienaar auto adaptor = TransferWriteOpAdaptor(operands); 2015f9e0466SNicolas Vasilache rewriter.replaceOpWithNewOp<LLVM::MaskedStoreOp>( 2025f9e0466SNicolas Vasilache xferOp, adaptor.vector(), dataPtr, mask, 2035f9e0466SNicolas Vasilache rewriter.getI32IntegerAttr(align)); 2045f9e0466SNicolas Vasilache return success(); 2055f9e0466SNicolas Vasilache } 2065f9e0466SNicolas Vasilache 2072d2c73c5SJacques Pienaar static TransferReadOpAdaptor getTransferOpAdapter(TransferReadOp xferOp, 2082d2c73c5SJacques Pienaar ArrayRef<Value> operands) { 2092d2c73c5SJacques Pienaar return TransferReadOpAdaptor(operands); 2105f9e0466SNicolas Vasilache } 2115f9e0466SNicolas Vasilache 2122d2c73c5SJacques Pienaar static TransferWriteOpAdaptor getTransferOpAdapter(TransferWriteOp xferOp, 2132d2c73c5SJacques Pienaar ArrayRef<Value> operands) { 2142d2c73c5SJacques Pienaar return TransferWriteOpAdaptor(operands); 2155f9e0466SNicolas Vasilache } 2165f9e0466SNicolas Vasilache 21790c01357SBenjamin Kramer namespace { 218e83b7b99Saartbik 21963b683a8SNicolas Vasilache /// Conversion pattern for a vector.matrix_multiply. 22063b683a8SNicolas Vasilache /// This is lowered directly to the proper llvm.intr.matrix.multiply. 22163b683a8SNicolas Vasilache class VectorMatmulOpConversion : public ConvertToLLVMPattern { 22263b683a8SNicolas Vasilache public: 22363b683a8SNicolas Vasilache explicit VectorMatmulOpConversion(MLIRContext *context, 22463b683a8SNicolas Vasilache LLVMTypeConverter &typeConverter) 22563b683a8SNicolas Vasilache : ConvertToLLVMPattern(vector::MatmulOp::getOperationName(), context, 22663b683a8SNicolas Vasilache typeConverter) {} 22763b683a8SNicolas Vasilache 2283145427dSRiver Riddle LogicalResult 22963b683a8SNicolas Vasilache matchAndRewrite(Operation *op, ArrayRef<Value> operands, 23063b683a8SNicolas Vasilache ConversionPatternRewriter &rewriter) const override { 23163b683a8SNicolas Vasilache auto matmulOp = cast<vector::MatmulOp>(op); 2322d2c73c5SJacques Pienaar auto adaptor = vector::MatmulOpAdaptor(operands); 23363b683a8SNicolas Vasilache rewriter.replaceOpWithNewOp<LLVM::MatrixMultiplyOp>( 23463b683a8SNicolas Vasilache op, typeConverter.convertType(matmulOp.res().getType()), adaptor.lhs(), 23563b683a8SNicolas Vasilache adaptor.rhs(), matmulOp.lhs_rows(), matmulOp.lhs_columns(), 23663b683a8SNicolas Vasilache matmulOp.rhs_columns()); 2373145427dSRiver Riddle return success(); 23863b683a8SNicolas Vasilache } 23963b683a8SNicolas Vasilache }; 24063b683a8SNicolas Vasilache 241c295a65dSaartbik /// Conversion pattern for a vector.flat_transpose. 242c295a65dSaartbik /// This is lowered directly to the proper llvm.intr.matrix.transpose. 243c295a65dSaartbik class VectorFlatTransposeOpConversion : public ConvertToLLVMPattern { 244c295a65dSaartbik public: 245c295a65dSaartbik explicit VectorFlatTransposeOpConversion(MLIRContext *context, 246c295a65dSaartbik LLVMTypeConverter &typeConverter) 247c295a65dSaartbik : ConvertToLLVMPattern(vector::FlatTransposeOp::getOperationName(), 248c295a65dSaartbik context, typeConverter) {} 249c295a65dSaartbik 250c295a65dSaartbik LogicalResult 251c295a65dSaartbik matchAndRewrite(Operation *op, ArrayRef<Value> operands, 252c295a65dSaartbik ConversionPatternRewriter &rewriter) const override { 253c295a65dSaartbik auto transOp = cast<vector::FlatTransposeOp>(op); 2542d2c73c5SJacques Pienaar auto adaptor = vector::FlatTransposeOpAdaptor(operands); 255c295a65dSaartbik rewriter.replaceOpWithNewOp<LLVM::MatrixTransposeOp>( 256c295a65dSaartbik transOp, typeConverter.convertType(transOp.res().getType()), 257c295a65dSaartbik adaptor.matrix(), transOp.rows(), transOp.columns()); 258c295a65dSaartbik return success(); 259c295a65dSaartbik } 260c295a65dSaartbik }; 261c295a65dSaartbik 262870c1fd4SAlex Zinenko class VectorReductionOpConversion : public ConvertToLLVMPattern { 263e83b7b99Saartbik public: 264e83b7b99Saartbik explicit VectorReductionOpConversion(MLIRContext *context, 265ceb1b327Saartbik LLVMTypeConverter &typeConverter, 266ceb1b327Saartbik bool reassociateFP) 267870c1fd4SAlex Zinenko : ConvertToLLVMPattern(vector::ReductionOp::getOperationName(), context, 268ceb1b327Saartbik typeConverter), 269ceb1b327Saartbik reassociateFPReductions(reassociateFP) {} 270e83b7b99Saartbik 2713145427dSRiver Riddle LogicalResult 272e83b7b99Saartbik matchAndRewrite(Operation *op, ArrayRef<Value> operands, 273e83b7b99Saartbik ConversionPatternRewriter &rewriter) const override { 274e83b7b99Saartbik auto reductionOp = cast<vector::ReductionOp>(op); 275e83b7b99Saartbik auto kind = reductionOp.kind(); 276e83b7b99Saartbik Type eltType = reductionOp.dest().getType(); 2770f04384dSAlex Zinenko Type llvmType = typeConverter.convertType(eltType); 27835b68527SLei Zhang if (eltType.isSignlessInteger(32) || eltType.isSignlessInteger(64)) { 279e83b7b99Saartbik // Integer reductions: add/mul/min/max/and/or/xor. 280e83b7b99Saartbik if (kind == "add") 281e83b7b99Saartbik rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_add>( 282e83b7b99Saartbik op, llvmType, operands[0]); 283e83b7b99Saartbik else if (kind == "mul") 284e83b7b99Saartbik rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_mul>( 285e83b7b99Saartbik op, llvmType, operands[0]); 286e83b7b99Saartbik else if (kind == "min") 287e83b7b99Saartbik rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_smin>( 288e83b7b99Saartbik op, llvmType, operands[0]); 289e83b7b99Saartbik else if (kind == "max") 290e83b7b99Saartbik rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_smax>( 291e83b7b99Saartbik op, llvmType, operands[0]); 292e83b7b99Saartbik else if (kind == "and") 293e83b7b99Saartbik rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_and>( 294e83b7b99Saartbik op, llvmType, operands[0]); 295e83b7b99Saartbik else if (kind == "or") 296e83b7b99Saartbik rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_or>( 297e83b7b99Saartbik op, llvmType, operands[0]); 298e83b7b99Saartbik else if (kind == "xor") 299e83b7b99Saartbik rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_xor>( 300e83b7b99Saartbik op, llvmType, operands[0]); 301e83b7b99Saartbik else 3023145427dSRiver Riddle return failure(); 3033145427dSRiver Riddle return success(); 304e83b7b99Saartbik 305e83b7b99Saartbik } else if (eltType.isF32() || eltType.isF64()) { 306e83b7b99Saartbik // Floating-point reductions: add/mul/min/max 307e83b7b99Saartbik if (kind == "add") { 3080d924700Saartbik // Optional accumulator (or zero). 3090d924700Saartbik Value acc = operands.size() > 1 ? operands[1] 3100d924700Saartbik : rewriter.create<LLVM::ConstantOp>( 3110d924700Saartbik op->getLoc(), llvmType, 3120d924700Saartbik rewriter.getZeroAttr(eltType)); 313e83b7b99Saartbik rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_v2_fadd>( 314ceb1b327Saartbik op, llvmType, acc, operands[0], 315ceb1b327Saartbik rewriter.getBoolAttr(reassociateFPReductions)); 316e83b7b99Saartbik } else if (kind == "mul") { 3170d924700Saartbik // Optional accumulator (or one). 3180d924700Saartbik Value acc = operands.size() > 1 3190d924700Saartbik ? operands[1] 3200d924700Saartbik : rewriter.create<LLVM::ConstantOp>( 3210d924700Saartbik op->getLoc(), llvmType, 3220d924700Saartbik rewriter.getFloatAttr(eltType, 1.0)); 323e83b7b99Saartbik rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_v2_fmul>( 324ceb1b327Saartbik op, llvmType, acc, operands[0], 325ceb1b327Saartbik rewriter.getBoolAttr(reassociateFPReductions)); 326e83b7b99Saartbik } else if (kind == "min") 327e83b7b99Saartbik rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_fmin>( 328e83b7b99Saartbik op, llvmType, operands[0]); 329e83b7b99Saartbik else if (kind == "max") 330e83b7b99Saartbik rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_fmax>( 331e83b7b99Saartbik op, llvmType, operands[0]); 332e83b7b99Saartbik else 3333145427dSRiver Riddle return failure(); 3343145427dSRiver Riddle return success(); 335e83b7b99Saartbik } 3363145427dSRiver Riddle return failure(); 337e83b7b99Saartbik } 338ceb1b327Saartbik 339ceb1b327Saartbik private: 340ceb1b327Saartbik const bool reassociateFPReductions; 341e83b7b99Saartbik }; 342e83b7b99Saartbik 343870c1fd4SAlex Zinenko class VectorShuffleOpConversion : public ConvertToLLVMPattern { 3441c81adf3SAart Bik public: 3451c81adf3SAart Bik explicit VectorShuffleOpConversion(MLIRContext *context, 3461c81adf3SAart Bik LLVMTypeConverter &typeConverter) 347870c1fd4SAlex Zinenko : ConvertToLLVMPattern(vector::ShuffleOp::getOperationName(), context, 3481c81adf3SAart Bik typeConverter) {} 3491c81adf3SAart Bik 3503145427dSRiver Riddle LogicalResult 351e62a6956SRiver Riddle matchAndRewrite(Operation *op, ArrayRef<Value> operands, 3521c81adf3SAart Bik ConversionPatternRewriter &rewriter) const override { 3531c81adf3SAart Bik auto loc = op->getLoc(); 3542d2c73c5SJacques Pienaar auto adaptor = vector::ShuffleOpAdaptor(operands); 3551c81adf3SAart Bik auto shuffleOp = cast<vector::ShuffleOp>(op); 3561c81adf3SAart Bik auto v1Type = shuffleOp.getV1VectorType(); 3571c81adf3SAart Bik auto v2Type = shuffleOp.getV2VectorType(); 3581c81adf3SAart Bik auto vectorType = shuffleOp.getVectorType(); 3590f04384dSAlex Zinenko Type llvmType = typeConverter.convertType(vectorType); 3601c81adf3SAart Bik auto maskArrayAttr = shuffleOp.mask(); 3611c81adf3SAart Bik 3621c81adf3SAart Bik // Bail if result type cannot be lowered. 3631c81adf3SAart Bik if (!llvmType) 3643145427dSRiver Riddle return failure(); 3651c81adf3SAart Bik 3661c81adf3SAart Bik // Get rank and dimension sizes. 3671c81adf3SAart Bik int64_t rank = vectorType.getRank(); 3681c81adf3SAart Bik assert(v1Type.getRank() == rank); 3691c81adf3SAart Bik assert(v2Type.getRank() == rank); 3701c81adf3SAart Bik int64_t v1Dim = v1Type.getDimSize(0); 3711c81adf3SAart Bik 3721c81adf3SAart Bik // For rank 1, where both operands have *exactly* the same vector type, 3731c81adf3SAart Bik // there is direct shuffle support in LLVM. Use it! 3741c81adf3SAart Bik if (rank == 1 && v1Type == v2Type) { 375e62a6956SRiver Riddle Value shuffle = rewriter.create<LLVM::ShuffleVectorOp>( 3761c81adf3SAart Bik loc, adaptor.v1(), adaptor.v2(), maskArrayAttr); 3771c81adf3SAart Bik rewriter.replaceOp(op, shuffle); 3783145427dSRiver Riddle return success(); 379b36aaeafSAart Bik } 380b36aaeafSAart Bik 3811c81adf3SAart Bik // For all other cases, insert the individual values individually. 382e62a6956SRiver Riddle Value insert = rewriter.create<LLVM::UndefOp>(loc, llvmType); 3831c81adf3SAart Bik int64_t insPos = 0; 3841c81adf3SAart Bik for (auto en : llvm::enumerate(maskArrayAttr)) { 3851c81adf3SAart Bik int64_t extPos = en.value().cast<IntegerAttr>().getInt(); 386e62a6956SRiver Riddle Value value = adaptor.v1(); 3871c81adf3SAart Bik if (extPos >= v1Dim) { 3881c81adf3SAart Bik extPos -= v1Dim; 3891c81adf3SAart Bik value = adaptor.v2(); 390b36aaeafSAart Bik } 3910f04384dSAlex Zinenko Value extract = extractOne(rewriter, typeConverter, loc, value, llvmType, 3920f04384dSAlex Zinenko rank, extPos); 3930f04384dSAlex Zinenko insert = insertOne(rewriter, typeConverter, loc, insert, extract, 3940f04384dSAlex Zinenko llvmType, rank, insPos++); 3951c81adf3SAart Bik } 3961c81adf3SAart Bik rewriter.replaceOp(op, insert); 3973145427dSRiver Riddle return success(); 398b36aaeafSAart Bik } 399b36aaeafSAart Bik }; 400b36aaeafSAart Bik 401870c1fd4SAlex Zinenko class VectorExtractElementOpConversion : public ConvertToLLVMPattern { 402cd5dab8aSAart Bik public: 403cd5dab8aSAart Bik explicit VectorExtractElementOpConversion(MLIRContext *context, 404cd5dab8aSAart Bik LLVMTypeConverter &typeConverter) 405870c1fd4SAlex Zinenko : ConvertToLLVMPattern(vector::ExtractElementOp::getOperationName(), 406870c1fd4SAlex Zinenko context, typeConverter) {} 407cd5dab8aSAart Bik 4083145427dSRiver Riddle LogicalResult 409e62a6956SRiver Riddle matchAndRewrite(Operation *op, ArrayRef<Value> operands, 410cd5dab8aSAart Bik ConversionPatternRewriter &rewriter) const override { 4112d2c73c5SJacques Pienaar auto adaptor = vector::ExtractElementOpAdaptor(operands); 412cd5dab8aSAart Bik auto extractEltOp = cast<vector::ExtractElementOp>(op); 413cd5dab8aSAart Bik auto vectorType = extractEltOp.getVectorType(); 4140f04384dSAlex Zinenko auto llvmType = typeConverter.convertType(vectorType.getElementType()); 415cd5dab8aSAart Bik 416cd5dab8aSAart Bik // Bail if result type cannot be lowered. 417cd5dab8aSAart Bik if (!llvmType) 4183145427dSRiver Riddle return failure(); 419cd5dab8aSAart Bik 420cd5dab8aSAart Bik rewriter.replaceOpWithNewOp<LLVM::ExtractElementOp>( 421cd5dab8aSAart Bik op, llvmType, adaptor.vector(), adaptor.position()); 4223145427dSRiver Riddle return success(); 423cd5dab8aSAart Bik } 424cd5dab8aSAart Bik }; 425cd5dab8aSAart Bik 426870c1fd4SAlex Zinenko class VectorExtractOpConversion : public ConvertToLLVMPattern { 4275c0c51a9SNicolas Vasilache public: 4289826fe5cSAart Bik explicit VectorExtractOpConversion(MLIRContext *context, 4295c0c51a9SNicolas Vasilache LLVMTypeConverter &typeConverter) 430870c1fd4SAlex Zinenko : ConvertToLLVMPattern(vector::ExtractOp::getOperationName(), context, 4315c0c51a9SNicolas Vasilache typeConverter) {} 4325c0c51a9SNicolas Vasilache 4333145427dSRiver Riddle LogicalResult 434e62a6956SRiver Riddle matchAndRewrite(Operation *op, ArrayRef<Value> operands, 4355c0c51a9SNicolas Vasilache ConversionPatternRewriter &rewriter) const override { 4365c0c51a9SNicolas Vasilache auto loc = op->getLoc(); 4372d2c73c5SJacques Pienaar auto adaptor = vector::ExtractOpAdaptor(operands); 438d37f2725SAart Bik auto extractOp = cast<vector::ExtractOp>(op); 4399826fe5cSAart Bik auto vectorType = extractOp.getVectorType(); 4402bdf33ccSRiver Riddle auto resultType = extractOp.getResult().getType(); 4410f04384dSAlex Zinenko auto llvmResultType = typeConverter.convertType(resultType); 4425c0c51a9SNicolas Vasilache auto positionArrayAttr = extractOp.position(); 4439826fe5cSAart Bik 4449826fe5cSAart Bik // Bail if result type cannot be lowered. 4459826fe5cSAart Bik if (!llvmResultType) 4463145427dSRiver Riddle return failure(); 4479826fe5cSAart Bik 4485c0c51a9SNicolas Vasilache // One-shot extraction of vector from array (only requires extractvalue). 4495c0c51a9SNicolas Vasilache if (resultType.isa<VectorType>()) { 450e62a6956SRiver Riddle Value extracted = rewriter.create<LLVM::ExtractValueOp>( 4515c0c51a9SNicolas Vasilache loc, llvmResultType, adaptor.vector(), positionArrayAttr); 4525c0c51a9SNicolas Vasilache rewriter.replaceOp(op, extracted); 4533145427dSRiver Riddle return success(); 4545c0c51a9SNicolas Vasilache } 4555c0c51a9SNicolas Vasilache 4569826fe5cSAart Bik // Potential extraction of 1-D vector from array. 4575c0c51a9SNicolas Vasilache auto *context = op->getContext(); 458e62a6956SRiver Riddle Value extracted = adaptor.vector(); 4595c0c51a9SNicolas Vasilache auto positionAttrs = positionArrayAttr.getValue(); 4605c0c51a9SNicolas Vasilache if (positionAttrs.size() > 1) { 4619826fe5cSAart Bik auto oneDVectorType = reducedVectorTypeBack(vectorType); 4625c0c51a9SNicolas Vasilache auto nMinusOnePositionAttrs = 4635c0c51a9SNicolas Vasilache ArrayAttr::get(positionAttrs.drop_back(), context); 4645c0c51a9SNicolas Vasilache extracted = rewriter.create<LLVM::ExtractValueOp>( 4650f04384dSAlex Zinenko loc, typeConverter.convertType(oneDVectorType), extracted, 4665c0c51a9SNicolas Vasilache nMinusOnePositionAttrs); 4675c0c51a9SNicolas Vasilache } 4685c0c51a9SNicolas Vasilache 4695c0c51a9SNicolas Vasilache // Remaining extraction of element from 1-D LLVM vector 4705c0c51a9SNicolas Vasilache auto position = positionAttrs.back().cast<IntegerAttr>(); 4710f04384dSAlex Zinenko auto i64Type = LLVM::LLVMType::getInt64Ty(typeConverter.getDialect()); 4721d47564aSAart Bik auto constant = rewriter.create<LLVM::ConstantOp>(loc, i64Type, position); 4735c0c51a9SNicolas Vasilache extracted = 4745c0c51a9SNicolas Vasilache rewriter.create<LLVM::ExtractElementOp>(loc, extracted, constant); 4755c0c51a9SNicolas Vasilache rewriter.replaceOp(op, extracted); 4765c0c51a9SNicolas Vasilache 4773145427dSRiver Riddle return success(); 4785c0c51a9SNicolas Vasilache } 4795c0c51a9SNicolas Vasilache }; 4805c0c51a9SNicolas Vasilache 481681f929fSNicolas Vasilache /// Conversion pattern that turns a vector.fma on a 1-D vector 482681f929fSNicolas Vasilache /// into an llvm.intr.fmuladd. This is a trivial 1-1 conversion. 483681f929fSNicolas Vasilache /// This does not match vectors of n >= 2 rank. 484681f929fSNicolas Vasilache /// 485681f929fSNicolas Vasilache /// Example: 486681f929fSNicolas Vasilache /// ``` 487681f929fSNicolas Vasilache /// vector.fma %a, %a, %a : vector<8xf32> 488681f929fSNicolas Vasilache /// ``` 489681f929fSNicolas Vasilache /// is converted to: 490681f929fSNicolas Vasilache /// ``` 4913bffe602SBenjamin Kramer /// llvm.intr.fmuladd %va, %va, %va: 492681f929fSNicolas Vasilache /// (!llvm<"<8 x float>">, !llvm<"<8 x float>">, !llvm<"<8 x float>">) 493681f929fSNicolas Vasilache /// -> !llvm<"<8 x float>"> 494681f929fSNicolas Vasilache /// ``` 495870c1fd4SAlex Zinenko class VectorFMAOp1DConversion : public ConvertToLLVMPattern { 496681f929fSNicolas Vasilache public: 497681f929fSNicolas Vasilache explicit VectorFMAOp1DConversion(MLIRContext *context, 498681f929fSNicolas Vasilache LLVMTypeConverter &typeConverter) 499870c1fd4SAlex Zinenko : ConvertToLLVMPattern(vector::FMAOp::getOperationName(), context, 500681f929fSNicolas Vasilache typeConverter) {} 501681f929fSNicolas Vasilache 5023145427dSRiver Riddle LogicalResult 503681f929fSNicolas Vasilache matchAndRewrite(Operation *op, ArrayRef<Value> operands, 504681f929fSNicolas Vasilache ConversionPatternRewriter &rewriter) const override { 5052d2c73c5SJacques Pienaar auto adaptor = vector::FMAOpAdaptor(operands); 506681f929fSNicolas Vasilache vector::FMAOp fmaOp = cast<vector::FMAOp>(op); 507681f929fSNicolas Vasilache VectorType vType = fmaOp.getVectorType(); 508681f929fSNicolas Vasilache if (vType.getRank() != 1) 5093145427dSRiver Riddle return failure(); 5103bffe602SBenjamin Kramer rewriter.replaceOpWithNewOp<LLVM::FMulAddOp>(op, adaptor.lhs(), 5113bffe602SBenjamin Kramer adaptor.rhs(), adaptor.acc()); 5123145427dSRiver Riddle return success(); 513681f929fSNicolas Vasilache } 514681f929fSNicolas Vasilache }; 515681f929fSNicolas Vasilache 516870c1fd4SAlex Zinenko class VectorInsertElementOpConversion : public ConvertToLLVMPattern { 517cd5dab8aSAart Bik public: 518cd5dab8aSAart Bik explicit VectorInsertElementOpConversion(MLIRContext *context, 519cd5dab8aSAart Bik LLVMTypeConverter &typeConverter) 520870c1fd4SAlex Zinenko : ConvertToLLVMPattern(vector::InsertElementOp::getOperationName(), 521870c1fd4SAlex Zinenko context, typeConverter) {} 522cd5dab8aSAart Bik 5233145427dSRiver Riddle LogicalResult 524e62a6956SRiver Riddle matchAndRewrite(Operation *op, ArrayRef<Value> operands, 525cd5dab8aSAart Bik ConversionPatternRewriter &rewriter) const override { 5262d2c73c5SJacques Pienaar auto adaptor = vector::InsertElementOpAdaptor(operands); 527cd5dab8aSAart Bik auto insertEltOp = cast<vector::InsertElementOp>(op); 528cd5dab8aSAart Bik auto vectorType = insertEltOp.getDestVectorType(); 5290f04384dSAlex Zinenko auto llvmType = typeConverter.convertType(vectorType); 530cd5dab8aSAart Bik 531cd5dab8aSAart Bik // Bail if result type cannot be lowered. 532cd5dab8aSAart Bik if (!llvmType) 5333145427dSRiver Riddle return failure(); 534cd5dab8aSAart Bik 535cd5dab8aSAart Bik rewriter.replaceOpWithNewOp<LLVM::InsertElementOp>( 536cd5dab8aSAart Bik op, llvmType, adaptor.dest(), adaptor.source(), adaptor.position()); 5373145427dSRiver Riddle return success(); 538cd5dab8aSAart Bik } 539cd5dab8aSAart Bik }; 540cd5dab8aSAart Bik 541870c1fd4SAlex Zinenko class VectorInsertOpConversion : public ConvertToLLVMPattern { 5429826fe5cSAart Bik public: 5439826fe5cSAart Bik explicit VectorInsertOpConversion(MLIRContext *context, 5449826fe5cSAart Bik LLVMTypeConverter &typeConverter) 545870c1fd4SAlex Zinenko : ConvertToLLVMPattern(vector::InsertOp::getOperationName(), context, 5469826fe5cSAart Bik typeConverter) {} 5479826fe5cSAart Bik 5483145427dSRiver Riddle LogicalResult 549e62a6956SRiver Riddle matchAndRewrite(Operation *op, ArrayRef<Value> operands, 5509826fe5cSAart Bik ConversionPatternRewriter &rewriter) const override { 5519826fe5cSAart Bik auto loc = op->getLoc(); 5522d2c73c5SJacques Pienaar auto adaptor = vector::InsertOpAdaptor(operands); 5539826fe5cSAart Bik auto insertOp = cast<vector::InsertOp>(op); 5549826fe5cSAart Bik auto sourceType = insertOp.getSourceType(); 5559826fe5cSAart Bik auto destVectorType = insertOp.getDestVectorType(); 5560f04384dSAlex Zinenko auto llvmResultType = typeConverter.convertType(destVectorType); 5579826fe5cSAart Bik auto positionArrayAttr = insertOp.position(); 5589826fe5cSAart Bik 5599826fe5cSAart Bik // Bail if result type cannot be lowered. 5609826fe5cSAart Bik if (!llvmResultType) 5613145427dSRiver Riddle return failure(); 5629826fe5cSAart Bik 5639826fe5cSAart Bik // One-shot insertion of a vector into an array (only requires insertvalue). 5649826fe5cSAart Bik if (sourceType.isa<VectorType>()) { 565e62a6956SRiver Riddle Value inserted = rewriter.create<LLVM::InsertValueOp>( 5669826fe5cSAart Bik loc, llvmResultType, adaptor.dest(), adaptor.source(), 5679826fe5cSAart Bik positionArrayAttr); 5689826fe5cSAart Bik rewriter.replaceOp(op, inserted); 5693145427dSRiver Riddle return success(); 5709826fe5cSAart Bik } 5719826fe5cSAart Bik 5729826fe5cSAart Bik // Potential extraction of 1-D vector from array. 5739826fe5cSAart Bik auto *context = op->getContext(); 574e62a6956SRiver Riddle Value extracted = adaptor.dest(); 5759826fe5cSAart Bik auto positionAttrs = positionArrayAttr.getValue(); 5769826fe5cSAart Bik auto position = positionAttrs.back().cast<IntegerAttr>(); 5779826fe5cSAart Bik auto oneDVectorType = destVectorType; 5789826fe5cSAart Bik if (positionAttrs.size() > 1) { 5799826fe5cSAart Bik oneDVectorType = reducedVectorTypeBack(destVectorType); 5809826fe5cSAart Bik auto nMinusOnePositionAttrs = 5819826fe5cSAart Bik ArrayAttr::get(positionAttrs.drop_back(), context); 5829826fe5cSAart Bik extracted = rewriter.create<LLVM::ExtractValueOp>( 5830f04384dSAlex Zinenko loc, typeConverter.convertType(oneDVectorType), extracted, 5849826fe5cSAart Bik nMinusOnePositionAttrs); 5859826fe5cSAart Bik } 5869826fe5cSAart Bik 5879826fe5cSAart Bik // Insertion of an element into a 1-D LLVM vector. 5880f04384dSAlex Zinenko auto i64Type = LLVM::LLVMType::getInt64Ty(typeConverter.getDialect()); 5891d47564aSAart Bik auto constant = rewriter.create<LLVM::ConstantOp>(loc, i64Type, position); 590e62a6956SRiver Riddle Value inserted = rewriter.create<LLVM::InsertElementOp>( 5910f04384dSAlex Zinenko loc, typeConverter.convertType(oneDVectorType), extracted, 5920f04384dSAlex Zinenko adaptor.source(), constant); 5939826fe5cSAart Bik 5949826fe5cSAart Bik // Potential insertion of resulting 1-D vector into array. 5959826fe5cSAart Bik if (positionAttrs.size() > 1) { 5969826fe5cSAart Bik auto nMinusOnePositionAttrs = 5979826fe5cSAart Bik ArrayAttr::get(positionAttrs.drop_back(), context); 5989826fe5cSAart Bik inserted = rewriter.create<LLVM::InsertValueOp>(loc, llvmResultType, 5999826fe5cSAart Bik adaptor.dest(), inserted, 6009826fe5cSAart Bik nMinusOnePositionAttrs); 6019826fe5cSAart Bik } 6029826fe5cSAart Bik 6039826fe5cSAart Bik rewriter.replaceOp(op, inserted); 6043145427dSRiver Riddle return success(); 6059826fe5cSAart Bik } 6069826fe5cSAart Bik }; 6079826fe5cSAart Bik 608681f929fSNicolas Vasilache /// Rank reducing rewrite for n-D FMA into (n-1)-D FMA where n > 1. 609681f929fSNicolas Vasilache /// 610681f929fSNicolas Vasilache /// Example: 611681f929fSNicolas Vasilache /// ``` 612681f929fSNicolas Vasilache /// %d = vector.fma %a, %b, %c : vector<2x4xf32> 613681f929fSNicolas Vasilache /// ``` 614681f929fSNicolas Vasilache /// is rewritten into: 615681f929fSNicolas Vasilache /// ``` 616681f929fSNicolas Vasilache /// %r = splat %f0: vector<2x4xf32> 617681f929fSNicolas Vasilache /// %va = vector.extractvalue %a[0] : vector<2x4xf32> 618681f929fSNicolas Vasilache /// %vb = vector.extractvalue %b[0] : vector<2x4xf32> 619681f929fSNicolas Vasilache /// %vc = vector.extractvalue %c[0] : vector<2x4xf32> 620681f929fSNicolas Vasilache /// %vd = vector.fma %va, %vb, %vc : vector<4xf32> 621681f929fSNicolas Vasilache /// %r2 = vector.insertvalue %vd, %r[0] : vector<4xf32> into vector<2x4xf32> 622681f929fSNicolas Vasilache /// %va2 = vector.extractvalue %a2[1] : vector<2x4xf32> 623681f929fSNicolas Vasilache /// %vb2 = vector.extractvalue %b2[1] : vector<2x4xf32> 624681f929fSNicolas Vasilache /// %vc2 = vector.extractvalue %c2[1] : vector<2x4xf32> 625681f929fSNicolas Vasilache /// %vd2 = vector.fma %va2, %vb2, %vc2 : vector<4xf32> 626681f929fSNicolas Vasilache /// %r3 = vector.insertvalue %vd2, %r2[1] : vector<4xf32> into vector<2x4xf32> 627681f929fSNicolas Vasilache /// // %r3 holds the final value. 628681f929fSNicolas Vasilache /// ``` 629681f929fSNicolas Vasilache class VectorFMAOpNDRewritePattern : public OpRewritePattern<FMAOp> { 630681f929fSNicolas Vasilache public: 631681f929fSNicolas Vasilache using OpRewritePattern<FMAOp>::OpRewritePattern; 632681f929fSNicolas Vasilache 6333145427dSRiver Riddle LogicalResult matchAndRewrite(FMAOp op, 634681f929fSNicolas Vasilache PatternRewriter &rewriter) const override { 635681f929fSNicolas Vasilache auto vType = op.getVectorType(); 636681f929fSNicolas Vasilache if (vType.getRank() < 2) 6373145427dSRiver Riddle return failure(); 638681f929fSNicolas Vasilache 639681f929fSNicolas Vasilache auto loc = op.getLoc(); 640681f929fSNicolas Vasilache auto elemType = vType.getElementType(); 641681f929fSNicolas Vasilache Value zero = rewriter.create<ConstantOp>(loc, elemType, 642681f929fSNicolas Vasilache rewriter.getZeroAttr(elemType)); 643681f929fSNicolas Vasilache Value desc = rewriter.create<SplatOp>(loc, vType, zero); 644681f929fSNicolas Vasilache for (int64_t i = 0, e = vType.getShape().front(); i != e; ++i) { 645681f929fSNicolas Vasilache Value extrLHS = rewriter.create<ExtractOp>(loc, op.lhs(), i); 646681f929fSNicolas Vasilache Value extrRHS = rewriter.create<ExtractOp>(loc, op.rhs(), i); 647681f929fSNicolas Vasilache Value extrACC = rewriter.create<ExtractOp>(loc, op.acc(), i); 648681f929fSNicolas Vasilache Value fma = rewriter.create<FMAOp>(loc, extrLHS, extrRHS, extrACC); 649681f929fSNicolas Vasilache desc = rewriter.create<InsertOp>(loc, fma, desc, i); 650681f929fSNicolas Vasilache } 651681f929fSNicolas Vasilache rewriter.replaceOp(op, desc); 6523145427dSRiver Riddle return success(); 653681f929fSNicolas Vasilache } 654681f929fSNicolas Vasilache }; 655681f929fSNicolas Vasilache 6562d515e49SNicolas Vasilache // When ranks are different, InsertStridedSlice needs to extract a properly 6572d515e49SNicolas Vasilache // ranked vector from the destination vector into which to insert. This pattern 6582d515e49SNicolas Vasilache // only takes care of this part and forwards the rest of the conversion to 6592d515e49SNicolas Vasilache // another pattern that converts InsertStridedSlice for operands of the same 6602d515e49SNicolas Vasilache // rank. 6612d515e49SNicolas Vasilache // 6622d515e49SNicolas Vasilache // RewritePattern for InsertStridedSliceOp where source and destination vectors 6632d515e49SNicolas Vasilache // have different ranks. In this case: 6642d515e49SNicolas Vasilache // 1. the proper subvector is extracted from the destination vector 6652d515e49SNicolas Vasilache // 2. a new InsertStridedSlice op is created to insert the source in the 6662d515e49SNicolas Vasilache // destination subvector 6672d515e49SNicolas Vasilache // 3. the destination subvector is inserted back in the proper place 6682d515e49SNicolas Vasilache // 4. the op is replaced by the result of step 3. 6692d515e49SNicolas Vasilache // The new InsertStridedSlice from step 2. will be picked up by a 6702d515e49SNicolas Vasilache // `VectorInsertStridedSliceOpSameRankRewritePattern`. 6712d515e49SNicolas Vasilache class VectorInsertStridedSliceOpDifferentRankRewritePattern 6722d515e49SNicolas Vasilache : public OpRewritePattern<InsertStridedSliceOp> { 6732d515e49SNicolas Vasilache public: 6742d515e49SNicolas Vasilache using OpRewritePattern<InsertStridedSliceOp>::OpRewritePattern; 6752d515e49SNicolas Vasilache 6763145427dSRiver Riddle LogicalResult matchAndRewrite(InsertStridedSliceOp op, 6772d515e49SNicolas Vasilache PatternRewriter &rewriter) const override { 6782d515e49SNicolas Vasilache auto srcType = op.getSourceVectorType(); 6792d515e49SNicolas Vasilache auto dstType = op.getDestVectorType(); 6802d515e49SNicolas Vasilache 6812d515e49SNicolas Vasilache if (op.offsets().getValue().empty()) 6823145427dSRiver Riddle return failure(); 6832d515e49SNicolas Vasilache 6842d515e49SNicolas Vasilache auto loc = op.getLoc(); 6852d515e49SNicolas Vasilache int64_t rankDiff = dstType.getRank() - srcType.getRank(); 6862d515e49SNicolas Vasilache assert(rankDiff >= 0); 6872d515e49SNicolas Vasilache if (rankDiff == 0) 6883145427dSRiver Riddle return failure(); 6892d515e49SNicolas Vasilache 6902d515e49SNicolas Vasilache int64_t rankRest = dstType.getRank() - rankDiff; 6912d515e49SNicolas Vasilache // Extract / insert the subvector of matching rank and InsertStridedSlice 6922d515e49SNicolas Vasilache // on it. 6932d515e49SNicolas Vasilache Value extracted = 6942d515e49SNicolas Vasilache rewriter.create<ExtractOp>(loc, op.dest(), 6952d515e49SNicolas Vasilache getI64SubArray(op.offsets(), /*dropFront=*/0, 6962d515e49SNicolas Vasilache /*dropFront=*/rankRest)); 6972d515e49SNicolas Vasilache // A different pattern will kick in for InsertStridedSlice with matching 6982d515e49SNicolas Vasilache // ranks. 6992d515e49SNicolas Vasilache auto stridedSliceInnerOp = rewriter.create<InsertStridedSliceOp>( 7002d515e49SNicolas Vasilache loc, op.source(), extracted, 7012d515e49SNicolas Vasilache getI64SubArray(op.offsets(), /*dropFront=*/rankDiff), 702c8fc76a9Saartbik getI64SubArray(op.strides(), /*dropFront=*/0)); 7032d515e49SNicolas Vasilache rewriter.replaceOpWithNewOp<InsertOp>( 7042d515e49SNicolas Vasilache op, stridedSliceInnerOp.getResult(), op.dest(), 7052d515e49SNicolas Vasilache getI64SubArray(op.offsets(), /*dropFront=*/0, 7062d515e49SNicolas Vasilache /*dropFront=*/rankRest)); 7073145427dSRiver Riddle return success(); 7082d515e49SNicolas Vasilache } 7092d515e49SNicolas Vasilache }; 7102d515e49SNicolas Vasilache 7112d515e49SNicolas Vasilache // RewritePattern for InsertStridedSliceOp where source and destination vectors 7122d515e49SNicolas Vasilache // have the same rank. In this case, we reduce 7132d515e49SNicolas Vasilache // 1. the proper subvector is extracted from the destination vector 7142d515e49SNicolas Vasilache // 2. a new InsertStridedSlice op is created to insert the source in the 7152d515e49SNicolas Vasilache // destination subvector 7162d515e49SNicolas Vasilache // 3. the destination subvector is inserted back in the proper place 7172d515e49SNicolas Vasilache // 4. the op is replaced by the result of step 3. 7182d515e49SNicolas Vasilache // The new InsertStridedSlice from step 2. will be picked up by a 7192d515e49SNicolas Vasilache // `VectorInsertStridedSliceOpSameRankRewritePattern`. 7202d515e49SNicolas Vasilache class VectorInsertStridedSliceOpSameRankRewritePattern 7212d515e49SNicolas Vasilache : public OpRewritePattern<InsertStridedSliceOp> { 7222d515e49SNicolas Vasilache public: 7232d515e49SNicolas Vasilache using OpRewritePattern<InsertStridedSliceOp>::OpRewritePattern; 7242d515e49SNicolas Vasilache 7253145427dSRiver Riddle LogicalResult matchAndRewrite(InsertStridedSliceOp op, 7262d515e49SNicolas Vasilache PatternRewriter &rewriter) const override { 7272d515e49SNicolas Vasilache auto srcType = op.getSourceVectorType(); 7282d515e49SNicolas Vasilache auto dstType = op.getDestVectorType(); 7292d515e49SNicolas Vasilache 7302d515e49SNicolas Vasilache if (op.offsets().getValue().empty()) 7313145427dSRiver Riddle return failure(); 7322d515e49SNicolas Vasilache 7332d515e49SNicolas Vasilache int64_t rankDiff = dstType.getRank() - srcType.getRank(); 7342d515e49SNicolas Vasilache assert(rankDiff >= 0); 7352d515e49SNicolas Vasilache if (rankDiff != 0) 7363145427dSRiver Riddle return failure(); 7372d515e49SNicolas Vasilache 7382d515e49SNicolas Vasilache if (srcType == dstType) { 7392d515e49SNicolas Vasilache rewriter.replaceOp(op, op.source()); 7403145427dSRiver Riddle return success(); 7412d515e49SNicolas Vasilache } 7422d515e49SNicolas Vasilache 7432d515e49SNicolas Vasilache int64_t offset = 7442d515e49SNicolas Vasilache op.offsets().getValue().front().cast<IntegerAttr>().getInt(); 7452d515e49SNicolas Vasilache int64_t size = srcType.getShape().front(); 7462d515e49SNicolas Vasilache int64_t stride = 7472d515e49SNicolas Vasilache op.strides().getValue().front().cast<IntegerAttr>().getInt(); 7482d515e49SNicolas Vasilache 7492d515e49SNicolas Vasilache auto loc = op.getLoc(); 7502d515e49SNicolas Vasilache Value res = op.dest(); 7512d515e49SNicolas Vasilache // For each slice of the source vector along the most major dimension. 7522d515e49SNicolas Vasilache for (int64_t off = offset, e = offset + size * stride, idx = 0; off < e; 7532d515e49SNicolas Vasilache off += stride, ++idx) { 7542d515e49SNicolas Vasilache // 1. extract the proper subvector (or element) from source 7552d515e49SNicolas Vasilache Value extractedSource = extractOne(rewriter, loc, op.source(), idx); 7562d515e49SNicolas Vasilache if (extractedSource.getType().isa<VectorType>()) { 7572d515e49SNicolas Vasilache // 2. If we have a vector, extract the proper subvector from destination 7582d515e49SNicolas Vasilache // Otherwise we are at the element level and no need to recurse. 7592d515e49SNicolas Vasilache Value extractedDest = extractOne(rewriter, loc, op.dest(), off); 7602d515e49SNicolas Vasilache // 3. Reduce the problem to lowering a new InsertStridedSlice op with 7612d515e49SNicolas Vasilache // smaller rank. 762bd1ccfe6SRiver Riddle extractedSource = rewriter.create<InsertStridedSliceOp>( 7632d515e49SNicolas Vasilache loc, extractedSource, extractedDest, 7642d515e49SNicolas Vasilache getI64SubArray(op.offsets(), /* dropFront=*/1), 7652d515e49SNicolas Vasilache getI64SubArray(op.strides(), /* dropFront=*/1)); 7662d515e49SNicolas Vasilache } 7672d515e49SNicolas Vasilache // 4. Insert the extractedSource into the res vector. 7682d515e49SNicolas Vasilache res = insertOne(rewriter, loc, extractedSource, res, off); 7692d515e49SNicolas Vasilache } 7702d515e49SNicolas Vasilache 7712d515e49SNicolas Vasilache rewriter.replaceOp(op, res); 7723145427dSRiver Riddle return success(); 7732d515e49SNicolas Vasilache } 774bd1ccfe6SRiver Riddle /// This pattern creates recursive InsertStridedSliceOp, but the recursion is 775bd1ccfe6SRiver Riddle /// bounded as the rank is strictly decreasing. 776bd1ccfe6SRiver Riddle bool hasBoundedRewriteRecursion() const final { return true; } 7772d515e49SNicolas Vasilache }; 7782d515e49SNicolas Vasilache 779870c1fd4SAlex Zinenko class VectorTypeCastOpConversion : public ConvertToLLVMPattern { 7805c0c51a9SNicolas Vasilache public: 7815c0c51a9SNicolas Vasilache explicit VectorTypeCastOpConversion(MLIRContext *context, 7825c0c51a9SNicolas Vasilache LLVMTypeConverter &typeConverter) 783870c1fd4SAlex Zinenko : ConvertToLLVMPattern(vector::TypeCastOp::getOperationName(), context, 7845c0c51a9SNicolas Vasilache typeConverter) {} 7855c0c51a9SNicolas Vasilache 7863145427dSRiver Riddle LogicalResult 787e62a6956SRiver Riddle matchAndRewrite(Operation *op, ArrayRef<Value> operands, 7885c0c51a9SNicolas Vasilache ConversionPatternRewriter &rewriter) const override { 7895c0c51a9SNicolas Vasilache auto loc = op->getLoc(); 7905c0c51a9SNicolas Vasilache vector::TypeCastOp castOp = cast<vector::TypeCastOp>(op); 7915c0c51a9SNicolas Vasilache MemRefType sourceMemRefType = 7922bdf33ccSRiver Riddle castOp.getOperand().getType().cast<MemRefType>(); 7935c0c51a9SNicolas Vasilache MemRefType targetMemRefType = 7942bdf33ccSRiver Riddle castOp.getResult().getType().cast<MemRefType>(); 7955c0c51a9SNicolas Vasilache 7965c0c51a9SNicolas Vasilache // Only static shape casts supported atm. 7975c0c51a9SNicolas Vasilache if (!sourceMemRefType.hasStaticShape() || 7985c0c51a9SNicolas Vasilache !targetMemRefType.hasStaticShape()) 7993145427dSRiver Riddle return failure(); 8005c0c51a9SNicolas Vasilache 8015c0c51a9SNicolas Vasilache auto llvmSourceDescriptorTy = 8022bdf33ccSRiver Riddle operands[0].getType().dyn_cast<LLVM::LLVMType>(); 8035c0c51a9SNicolas Vasilache if (!llvmSourceDescriptorTy || !llvmSourceDescriptorTy.isStructTy()) 8043145427dSRiver Riddle return failure(); 8055c0c51a9SNicolas Vasilache MemRefDescriptor sourceMemRef(operands[0]); 8065c0c51a9SNicolas Vasilache 8070f04384dSAlex Zinenko auto llvmTargetDescriptorTy = typeConverter.convertType(targetMemRefType) 8085c0c51a9SNicolas Vasilache .dyn_cast_or_null<LLVM::LLVMType>(); 8095c0c51a9SNicolas Vasilache if (!llvmTargetDescriptorTy || !llvmTargetDescriptorTy.isStructTy()) 8103145427dSRiver Riddle return failure(); 8115c0c51a9SNicolas Vasilache 8125c0c51a9SNicolas Vasilache int64_t offset; 8135c0c51a9SNicolas Vasilache SmallVector<int64_t, 4> strides; 8145c0c51a9SNicolas Vasilache auto successStrides = 8155c0c51a9SNicolas Vasilache getStridesAndOffset(sourceMemRefType, strides, offset); 8165c0c51a9SNicolas Vasilache bool isContiguous = (strides.back() == 1); 8175c0c51a9SNicolas Vasilache if (isContiguous) { 8185c0c51a9SNicolas Vasilache auto sizes = sourceMemRefType.getShape(); 8195c0c51a9SNicolas Vasilache for (int index = 0, e = strides.size() - 2; index < e; ++index) { 8205c0c51a9SNicolas Vasilache if (strides[index] != strides[index + 1] * sizes[index + 1]) { 8215c0c51a9SNicolas Vasilache isContiguous = false; 8225c0c51a9SNicolas Vasilache break; 8235c0c51a9SNicolas Vasilache } 8245c0c51a9SNicolas Vasilache } 8255c0c51a9SNicolas Vasilache } 8265c0c51a9SNicolas Vasilache // Only contiguous source tensors supported atm. 8275c0c51a9SNicolas Vasilache if (failed(successStrides) || !isContiguous) 8283145427dSRiver Riddle return failure(); 8295c0c51a9SNicolas Vasilache 8300f04384dSAlex Zinenko auto int64Ty = LLVM::LLVMType::getInt64Ty(typeConverter.getDialect()); 8315c0c51a9SNicolas Vasilache 8325c0c51a9SNicolas Vasilache // Create descriptor. 8335c0c51a9SNicolas Vasilache auto desc = MemRefDescriptor::undef(rewriter, loc, llvmTargetDescriptorTy); 8345c0c51a9SNicolas Vasilache Type llvmTargetElementTy = desc.getElementType(); 8355c0c51a9SNicolas Vasilache // Set allocated ptr. 836e62a6956SRiver Riddle Value allocated = sourceMemRef.allocatedPtr(rewriter, loc); 8375c0c51a9SNicolas Vasilache allocated = 8385c0c51a9SNicolas Vasilache rewriter.create<LLVM::BitcastOp>(loc, llvmTargetElementTy, allocated); 8395c0c51a9SNicolas Vasilache desc.setAllocatedPtr(rewriter, loc, allocated); 8405c0c51a9SNicolas Vasilache // Set aligned ptr. 841e62a6956SRiver Riddle Value ptr = sourceMemRef.alignedPtr(rewriter, loc); 8425c0c51a9SNicolas Vasilache ptr = rewriter.create<LLVM::BitcastOp>(loc, llvmTargetElementTy, ptr); 8435c0c51a9SNicolas Vasilache desc.setAlignedPtr(rewriter, loc, ptr); 8445c0c51a9SNicolas Vasilache // Fill offset 0. 8455c0c51a9SNicolas Vasilache auto attr = rewriter.getIntegerAttr(rewriter.getIndexType(), 0); 8465c0c51a9SNicolas Vasilache auto zero = rewriter.create<LLVM::ConstantOp>(loc, int64Ty, attr); 8475c0c51a9SNicolas Vasilache desc.setOffset(rewriter, loc, zero); 8485c0c51a9SNicolas Vasilache 8495c0c51a9SNicolas Vasilache // Fill size and stride descriptors in memref. 8505c0c51a9SNicolas Vasilache for (auto indexedSize : llvm::enumerate(targetMemRefType.getShape())) { 8515c0c51a9SNicolas Vasilache int64_t index = indexedSize.index(); 8525c0c51a9SNicolas Vasilache auto sizeAttr = 8535c0c51a9SNicolas Vasilache rewriter.getIntegerAttr(rewriter.getIndexType(), indexedSize.value()); 8545c0c51a9SNicolas Vasilache auto size = rewriter.create<LLVM::ConstantOp>(loc, int64Ty, sizeAttr); 8555c0c51a9SNicolas Vasilache desc.setSize(rewriter, loc, index, size); 8565c0c51a9SNicolas Vasilache auto strideAttr = 8575c0c51a9SNicolas Vasilache rewriter.getIntegerAttr(rewriter.getIndexType(), strides[index]); 8585c0c51a9SNicolas Vasilache auto stride = rewriter.create<LLVM::ConstantOp>(loc, int64Ty, strideAttr); 8595c0c51a9SNicolas Vasilache desc.setStride(rewriter, loc, index, stride); 8605c0c51a9SNicolas Vasilache } 8615c0c51a9SNicolas Vasilache 8625c0c51a9SNicolas Vasilache rewriter.replaceOp(op, {desc}); 8633145427dSRiver Riddle return success(); 8645c0c51a9SNicolas Vasilache } 8655c0c51a9SNicolas Vasilache }; 8665c0c51a9SNicolas Vasilache 8678345b86dSNicolas Vasilache /// Conversion pattern that converts a 1-D vector transfer read/write op in a 8688345b86dSNicolas Vasilache /// sequence of: 869be16075bSWen-Heng (Jack) Chung /// 1. Bitcast or addrspacecast to vector form. 8708345b86dSNicolas Vasilache /// 2. Create an offsetVector = [ offset + 0 .. offset + vector_length - 1 ]. 8718345b86dSNicolas Vasilache /// 3. Create a mask where offsetVector is compared against memref upper bound. 8728345b86dSNicolas Vasilache /// 4. Rewrite op as a masked read or write. 8738345b86dSNicolas Vasilache template <typename ConcreteOp> 8748345b86dSNicolas Vasilache class VectorTransferConversion : public ConvertToLLVMPattern { 8758345b86dSNicolas Vasilache public: 8768345b86dSNicolas Vasilache explicit VectorTransferConversion(MLIRContext *context, 8778345b86dSNicolas Vasilache LLVMTypeConverter &typeConv) 8788345b86dSNicolas Vasilache : ConvertToLLVMPattern(ConcreteOp::getOperationName(), context, 8798345b86dSNicolas Vasilache typeConv) {} 8808345b86dSNicolas Vasilache 8818345b86dSNicolas Vasilache LogicalResult 8828345b86dSNicolas Vasilache matchAndRewrite(Operation *op, ArrayRef<Value> operands, 8838345b86dSNicolas Vasilache ConversionPatternRewriter &rewriter) const override { 8848345b86dSNicolas Vasilache auto xferOp = cast<ConcreteOp>(op); 8858345b86dSNicolas Vasilache auto adaptor = getTransferOpAdapter(xferOp, operands); 886b2c79c50SNicolas Vasilache 887b2c79c50SNicolas Vasilache if (xferOp.getVectorType().getRank() > 1 || 888b2c79c50SNicolas Vasilache llvm::size(xferOp.indices()) == 0) 8898345b86dSNicolas Vasilache return failure(); 8905f9e0466SNicolas Vasilache if (xferOp.permutation_map() != 8915f9e0466SNicolas Vasilache AffineMap::getMinorIdentityMap(xferOp.permutation_map().getNumInputs(), 8925f9e0466SNicolas Vasilache xferOp.getVectorType().getRank(), 8935f9e0466SNicolas Vasilache op->getContext())) 8948345b86dSNicolas Vasilache return failure(); 8958345b86dSNicolas Vasilache 8968345b86dSNicolas Vasilache auto toLLVMTy = [&](Type t) { return typeConverter.convertType(t); }; 8978345b86dSNicolas Vasilache 8988345b86dSNicolas Vasilache Location loc = op->getLoc(); 8998345b86dSNicolas Vasilache Type i64Type = rewriter.getIntegerType(64); 9008345b86dSNicolas Vasilache MemRefType memRefType = xferOp.getMemRefType(); 9018345b86dSNicolas Vasilache 9028345b86dSNicolas Vasilache // 1. Get the source/dst address as an LLVM vector pointer. 903be16075bSWen-Heng (Jack) Chung // The vector pointer would always be on address space 0, therefore 904be16075bSWen-Heng (Jack) Chung // addrspacecast shall be used when source/dst memrefs are not on 905be16075bSWen-Heng (Jack) Chung // address space 0. 9068345b86dSNicolas Vasilache // TODO: support alignment when possible. 9078345b86dSNicolas Vasilache Value dataPtr = getDataPtr(loc, memRefType, adaptor.memref(), 9088345b86dSNicolas Vasilache adaptor.indices(), rewriter, getModule()); 9098345b86dSNicolas Vasilache auto vecTy = 9108345b86dSNicolas Vasilache toLLVMTy(xferOp.getVectorType()).template cast<LLVM::LLVMType>(); 911be16075bSWen-Heng (Jack) Chung Value vectorDataPtr; 912be16075bSWen-Heng (Jack) Chung if (memRefType.getMemorySpace() == 0) 913be16075bSWen-Heng (Jack) Chung vectorDataPtr = 9148345b86dSNicolas Vasilache rewriter.create<LLVM::BitcastOp>(loc, vecTy.getPointerTo(), dataPtr); 915be16075bSWen-Heng (Jack) Chung else 916be16075bSWen-Heng (Jack) Chung vectorDataPtr = rewriter.create<LLVM::AddrSpaceCastOp>( 917be16075bSWen-Heng (Jack) Chung loc, vecTy.getPointerTo(), dataPtr); 9188345b86dSNicolas Vasilache 9191870e787SNicolas Vasilache if (!xferOp.isMaskedDim(0)) 9201870e787SNicolas Vasilache return replaceTransferOpWithLoadOrStore(rewriter, typeConverter, loc, 9211870e787SNicolas Vasilache xferOp, operands, vectorDataPtr); 9221870e787SNicolas Vasilache 9238345b86dSNicolas Vasilache // 2. Create a vector with linear indices [ 0 .. vector_length - 1 ]. 9248345b86dSNicolas Vasilache unsigned vecWidth = vecTy.getVectorNumElements(); 9258345b86dSNicolas Vasilache VectorType vectorCmpType = VectorType::get(vecWidth, i64Type); 9268345b86dSNicolas Vasilache SmallVector<int64_t, 8> indices; 9278345b86dSNicolas Vasilache indices.reserve(vecWidth); 9288345b86dSNicolas Vasilache for (unsigned i = 0; i < vecWidth; ++i) 9298345b86dSNicolas Vasilache indices.push_back(i); 9308345b86dSNicolas Vasilache Value linearIndices = rewriter.create<ConstantOp>( 9318345b86dSNicolas Vasilache loc, vectorCmpType, 9328345b86dSNicolas Vasilache DenseElementsAttr::get(vectorCmpType, ArrayRef<int64_t>(indices))); 9338345b86dSNicolas Vasilache linearIndices = rewriter.create<LLVM::DialectCastOp>( 9348345b86dSNicolas Vasilache loc, toLLVMTy(vectorCmpType), linearIndices); 9358345b86dSNicolas Vasilache 9368345b86dSNicolas Vasilache // 3. Create offsetVector = [ offset + 0 .. offset + vector_length - 1 ]. 9379db53a18SRiver Riddle // TODO: when the leaf transfer rank is k > 1 we need the last 938b2c79c50SNicolas Vasilache // `k` dimensions here. 939b2c79c50SNicolas Vasilache unsigned lastIndex = llvm::size(xferOp.indices()) - 1; 940b2c79c50SNicolas Vasilache Value offsetIndex = *(xferOp.indices().begin() + lastIndex); 941b2c79c50SNicolas Vasilache offsetIndex = rewriter.create<IndexCastOp>(loc, i64Type, offsetIndex); 9428345b86dSNicolas Vasilache Value base = rewriter.create<SplatOp>(loc, vectorCmpType, offsetIndex); 9438345b86dSNicolas Vasilache Value offsetVector = rewriter.create<AddIOp>(loc, base, linearIndices); 9448345b86dSNicolas Vasilache 9458345b86dSNicolas Vasilache // 4. Let dim the memref dimension, compute the vector comparison mask: 9468345b86dSNicolas Vasilache // [ offset + 0 .. offset + vector_length - 1 ] < [ dim .. dim ] 947b2c79c50SNicolas Vasilache Value dim = rewriter.create<DimOp>(loc, xferOp.memref(), lastIndex); 948b2c79c50SNicolas Vasilache dim = rewriter.create<IndexCastOp>(loc, i64Type, dim); 9498345b86dSNicolas Vasilache dim = rewriter.create<SplatOp>(loc, vectorCmpType, dim); 9508345b86dSNicolas Vasilache Value mask = 9518345b86dSNicolas Vasilache rewriter.create<CmpIOp>(loc, CmpIPredicate::slt, offsetVector, dim); 9528345b86dSNicolas Vasilache mask = rewriter.create<LLVM::DialectCastOp>(loc, toLLVMTy(mask.getType()), 9538345b86dSNicolas Vasilache mask); 9548345b86dSNicolas Vasilache 9558345b86dSNicolas Vasilache // 5. Rewrite as a masked read / write. 9561870e787SNicolas Vasilache return replaceTransferOpWithMasked(rewriter, typeConverter, loc, xferOp, 957a99f62c4SAlex Zinenko operands, vectorDataPtr, mask); 9588345b86dSNicolas Vasilache } 9598345b86dSNicolas Vasilache }; 9608345b86dSNicolas Vasilache 961870c1fd4SAlex Zinenko class VectorPrintOpConversion : public ConvertToLLVMPattern { 962d9b500d3SAart Bik public: 963d9b500d3SAart Bik explicit VectorPrintOpConversion(MLIRContext *context, 964d9b500d3SAart Bik LLVMTypeConverter &typeConverter) 965870c1fd4SAlex Zinenko : ConvertToLLVMPattern(vector::PrintOp::getOperationName(), context, 966d9b500d3SAart Bik typeConverter) {} 967d9b500d3SAart Bik 968d9b500d3SAart Bik // Proof-of-concept lowering implementation that relies on a small 969d9b500d3SAart Bik // runtime support library, which only needs to provide a few 970d9b500d3SAart Bik // printing methods (single value for all data types, opening/closing 971d9b500d3SAart Bik // bracket, comma, newline). The lowering fully unrolls a vector 972d9b500d3SAart Bik // in terms of these elementary printing operations. The advantage 973d9b500d3SAart Bik // of this approach is that the library can remain unaware of all 974d9b500d3SAart Bik // low-level implementation details of vectors while still supporting 975d9b500d3SAart Bik // output of any shaped and dimensioned vector. Due to full unrolling, 976d9b500d3SAart Bik // this approach is less suited for very large vectors though. 977d9b500d3SAart Bik // 9789db53a18SRiver Riddle // TODO: rely solely on libc in future? something else? 979d9b500d3SAart Bik // 9803145427dSRiver Riddle LogicalResult 981e62a6956SRiver Riddle matchAndRewrite(Operation *op, ArrayRef<Value> operands, 982d9b500d3SAart Bik ConversionPatternRewriter &rewriter) const override { 983d9b500d3SAart Bik auto printOp = cast<vector::PrintOp>(op); 9842d2c73c5SJacques Pienaar auto adaptor = vector::PrintOpAdaptor(operands); 985d9b500d3SAart Bik Type printType = printOp.getPrintType(); 986d9b500d3SAart Bik 9870f04384dSAlex Zinenko if (typeConverter.convertType(printType) == nullptr) 9883145427dSRiver Riddle return failure(); 989d9b500d3SAart Bik 990d9b500d3SAart Bik // Make sure element type has runtime support (currently just Float/Double). 991d9b500d3SAart Bik VectorType vectorType = printType.dyn_cast<VectorType>(); 992d9b500d3SAart Bik Type eltType = vectorType ? vectorType.getElementType() : printType; 993d9b500d3SAart Bik int64_t rank = vectorType ? vectorType.getRank() : 0; 994d9b500d3SAart Bik Operation *printer; 995c9eeeb38Saartbik if (eltType.isSignlessInteger(1) || eltType.isSignlessInteger(32)) 996e52414b1Saartbik printer = getPrintI32(op); 99735b68527SLei Zhang else if (eltType.isSignlessInteger(64)) 998e52414b1Saartbik printer = getPrintI64(op); 999e52414b1Saartbik else if (eltType.isF32()) 1000d9b500d3SAart Bik printer = getPrintFloat(op); 1001d9b500d3SAart Bik else if (eltType.isF64()) 1002d9b500d3SAart Bik printer = getPrintDouble(op); 1003d9b500d3SAart Bik else 10043145427dSRiver Riddle return failure(); 1005d9b500d3SAart Bik 1006d9b500d3SAart Bik // Unroll vector into elementary print calls. 1007d9b500d3SAart Bik emitRanks(rewriter, op, adaptor.source(), vectorType, printer, rank); 1008d9b500d3SAart Bik emitCall(rewriter, op->getLoc(), getPrintNewline(op)); 1009d9b500d3SAart Bik rewriter.eraseOp(op); 10103145427dSRiver Riddle return success(); 1011d9b500d3SAart Bik } 1012d9b500d3SAart Bik 1013d9b500d3SAart Bik private: 1014d9b500d3SAart Bik void emitRanks(ConversionPatternRewriter &rewriter, Operation *op, 1015e62a6956SRiver Riddle Value value, VectorType vectorType, Operation *printer, 1016d9b500d3SAart Bik int64_t rank) const { 1017d9b500d3SAart Bik Location loc = op->getLoc(); 1018d9b500d3SAart Bik if (rank == 0) { 1019c9eeeb38Saartbik if (value.getType() == 1020c9eeeb38Saartbik LLVM::LLVMType::getInt1Ty(typeConverter.getDialect())) { 1021c9eeeb38Saartbik // Convert i1 (bool) to i32 so we can use the print_i32 method. 1022c9eeeb38Saartbik // This avoids the need for a print_i1 method with an unclear ABI. 1023c9eeeb38Saartbik auto i32Type = LLVM::LLVMType::getInt32Ty(typeConverter.getDialect()); 1024c9eeeb38Saartbik auto trueVal = rewriter.create<ConstantOp>( 1025c9eeeb38Saartbik loc, i32Type, rewriter.getI32IntegerAttr(1)); 1026c9eeeb38Saartbik auto falseVal = rewriter.create<ConstantOp>( 1027c9eeeb38Saartbik loc, i32Type, rewriter.getI32IntegerAttr(0)); 1028c9eeeb38Saartbik value = rewriter.create<SelectOp>(loc, value, trueVal, falseVal); 1029c9eeeb38Saartbik } 1030d9b500d3SAart Bik emitCall(rewriter, loc, printer, value); 1031d9b500d3SAart Bik return; 1032d9b500d3SAart Bik } 1033d9b500d3SAart Bik 1034d9b500d3SAart Bik emitCall(rewriter, loc, getPrintOpen(op)); 1035d9b500d3SAart Bik Operation *printComma = getPrintComma(op); 1036d9b500d3SAart Bik int64_t dim = vectorType.getDimSize(0); 1037d9b500d3SAart Bik for (int64_t d = 0; d < dim; ++d) { 1038d9b500d3SAart Bik auto reducedType = 1039d9b500d3SAart Bik rank > 1 ? reducedVectorTypeFront(vectorType) : nullptr; 10400f04384dSAlex Zinenko auto llvmType = typeConverter.convertType( 1041d9b500d3SAart Bik rank > 1 ? reducedType : vectorType.getElementType()); 1042e62a6956SRiver Riddle Value nestedVal = 10430f04384dSAlex Zinenko extractOne(rewriter, typeConverter, loc, value, llvmType, rank, d); 1044d9b500d3SAart Bik emitRanks(rewriter, op, nestedVal, reducedType, printer, rank - 1); 1045d9b500d3SAart Bik if (d != dim - 1) 1046d9b500d3SAart Bik emitCall(rewriter, loc, printComma); 1047d9b500d3SAart Bik } 1048d9b500d3SAart Bik emitCall(rewriter, loc, getPrintClose(op)); 1049d9b500d3SAart Bik } 1050d9b500d3SAart Bik 1051d9b500d3SAart Bik // Helper to emit a call. 1052d9b500d3SAart Bik static void emitCall(ConversionPatternRewriter &rewriter, Location loc, 1053d9b500d3SAart Bik Operation *ref, ValueRange params = ValueRange()) { 1054d9b500d3SAart Bik rewriter.create<LLVM::CallOp>(loc, ArrayRef<Type>{}, 1055d9b500d3SAart Bik rewriter.getSymbolRefAttr(ref), params); 1056d9b500d3SAart Bik } 1057d9b500d3SAart Bik 1058d9b500d3SAart Bik // Helper for printer method declaration (first hit) and lookup. 1059d9b500d3SAart Bik static Operation *getPrint(Operation *op, LLVM::LLVMDialect *dialect, 1060d9b500d3SAart Bik StringRef name, ArrayRef<LLVM::LLVMType> params) { 1061d9b500d3SAart Bik auto module = op->getParentOfType<ModuleOp>(); 1062d9b500d3SAart Bik auto func = module.lookupSymbol<LLVM::LLVMFuncOp>(name); 1063d9b500d3SAart Bik if (func) 1064d9b500d3SAart Bik return func; 1065d9b500d3SAart Bik OpBuilder moduleBuilder(module.getBodyRegion()); 1066d9b500d3SAart Bik return moduleBuilder.create<LLVM::LLVMFuncOp>( 1067d9b500d3SAart Bik op->getLoc(), name, 1068d9b500d3SAart Bik LLVM::LLVMType::getFunctionTy(LLVM::LLVMType::getVoidTy(dialect), 1069d9b500d3SAart Bik params, /*isVarArg=*/false)); 1070d9b500d3SAart Bik } 1071d9b500d3SAart Bik 1072d9b500d3SAart Bik // Helpers for method names. 1073e52414b1Saartbik Operation *getPrintI32(Operation *op) const { 10740f04384dSAlex Zinenko LLVM::LLVMDialect *dialect = typeConverter.getDialect(); 1075e52414b1Saartbik return getPrint(op, dialect, "print_i32", 1076e52414b1Saartbik LLVM::LLVMType::getInt32Ty(dialect)); 1077e52414b1Saartbik } 1078e52414b1Saartbik Operation *getPrintI64(Operation *op) const { 10790f04384dSAlex Zinenko LLVM::LLVMDialect *dialect = typeConverter.getDialect(); 1080e52414b1Saartbik return getPrint(op, dialect, "print_i64", 1081e52414b1Saartbik LLVM::LLVMType::getInt64Ty(dialect)); 1082e52414b1Saartbik } 1083d9b500d3SAart Bik Operation *getPrintFloat(Operation *op) const { 10840f04384dSAlex Zinenko LLVM::LLVMDialect *dialect = typeConverter.getDialect(); 1085d9b500d3SAart Bik return getPrint(op, dialect, "print_f32", 1086d9b500d3SAart Bik LLVM::LLVMType::getFloatTy(dialect)); 1087d9b500d3SAart Bik } 1088d9b500d3SAart Bik Operation *getPrintDouble(Operation *op) const { 10890f04384dSAlex Zinenko LLVM::LLVMDialect *dialect = typeConverter.getDialect(); 1090d9b500d3SAart Bik return getPrint(op, dialect, "print_f64", 1091d9b500d3SAart Bik LLVM::LLVMType::getDoubleTy(dialect)); 1092d9b500d3SAart Bik } 1093d9b500d3SAart Bik Operation *getPrintOpen(Operation *op) const { 10940f04384dSAlex Zinenko return getPrint(op, typeConverter.getDialect(), "print_open", {}); 1095d9b500d3SAart Bik } 1096d9b500d3SAart Bik Operation *getPrintClose(Operation *op) const { 10970f04384dSAlex Zinenko return getPrint(op, typeConverter.getDialect(), "print_close", {}); 1098d9b500d3SAart Bik } 1099d9b500d3SAart Bik Operation *getPrintComma(Operation *op) const { 11000f04384dSAlex Zinenko return getPrint(op, typeConverter.getDialect(), "print_comma", {}); 1101d9b500d3SAart Bik } 1102d9b500d3SAart Bik Operation *getPrintNewline(Operation *op) const { 11030f04384dSAlex Zinenko return getPrint(op, typeConverter.getDialect(), "print_newline", {}); 1104d9b500d3SAart Bik } 1105d9b500d3SAart Bik }; 1106d9b500d3SAart Bik 1107334a4159SReid Tatge /// Progressive lowering of ExtractStridedSliceOp to either: 110865678d93SNicolas Vasilache /// 1. extractelement + insertelement for the 1-D case 110965678d93SNicolas Vasilache /// 2. extract + optional strided_slice + insert for the n-D case. 1110334a4159SReid Tatge class VectorStridedSliceOpConversion 1111334a4159SReid Tatge : public OpRewritePattern<ExtractStridedSliceOp> { 111265678d93SNicolas Vasilache public: 1113334a4159SReid Tatge using OpRewritePattern<ExtractStridedSliceOp>::OpRewritePattern; 111465678d93SNicolas Vasilache 1115334a4159SReid Tatge LogicalResult matchAndRewrite(ExtractStridedSliceOp op, 111665678d93SNicolas Vasilache PatternRewriter &rewriter) const override { 111765678d93SNicolas Vasilache auto dstType = op.getResult().getType().cast<VectorType>(); 111865678d93SNicolas Vasilache 111965678d93SNicolas Vasilache assert(!op.offsets().getValue().empty() && "Unexpected empty offsets"); 112065678d93SNicolas Vasilache 112165678d93SNicolas Vasilache int64_t offset = 112265678d93SNicolas Vasilache op.offsets().getValue().front().cast<IntegerAttr>().getInt(); 112365678d93SNicolas Vasilache int64_t size = op.sizes().getValue().front().cast<IntegerAttr>().getInt(); 112465678d93SNicolas Vasilache int64_t stride = 112565678d93SNicolas Vasilache op.strides().getValue().front().cast<IntegerAttr>().getInt(); 112665678d93SNicolas Vasilache 112765678d93SNicolas Vasilache auto loc = op.getLoc(); 112865678d93SNicolas Vasilache auto elemType = dstType.getElementType(); 112935b68527SLei Zhang assert(elemType.isSignlessIntOrIndexOrFloat()); 113065678d93SNicolas Vasilache Value zero = rewriter.create<ConstantOp>(loc, elemType, 113165678d93SNicolas Vasilache rewriter.getZeroAttr(elemType)); 113265678d93SNicolas Vasilache Value res = rewriter.create<SplatOp>(loc, dstType, zero); 113365678d93SNicolas Vasilache for (int64_t off = offset, e = offset + size * stride, idx = 0; off < e; 113465678d93SNicolas Vasilache off += stride, ++idx) { 113565678d93SNicolas Vasilache Value extracted = extractOne(rewriter, loc, op.vector(), off); 113665678d93SNicolas Vasilache if (op.offsets().getValue().size() > 1) { 1137334a4159SReid Tatge extracted = rewriter.create<ExtractStridedSliceOp>( 113865678d93SNicolas Vasilache loc, extracted, getI64SubArray(op.offsets(), /* dropFront=*/1), 113965678d93SNicolas Vasilache getI64SubArray(op.sizes(), /* dropFront=*/1), 114065678d93SNicolas Vasilache getI64SubArray(op.strides(), /* dropFront=*/1)); 114165678d93SNicolas Vasilache } 114265678d93SNicolas Vasilache res = insertOne(rewriter, loc, extracted, res, idx); 114365678d93SNicolas Vasilache } 114465678d93SNicolas Vasilache rewriter.replaceOp(op, {res}); 11453145427dSRiver Riddle return success(); 114665678d93SNicolas Vasilache } 1147334a4159SReid Tatge /// This pattern creates recursive ExtractStridedSliceOp, but the recursion is 1148bd1ccfe6SRiver Riddle /// bounded as the rank is strictly decreasing. 1149bd1ccfe6SRiver Riddle bool hasBoundedRewriteRecursion() const final { return true; } 115065678d93SNicolas Vasilache }; 115165678d93SNicolas Vasilache 1152df186507SBenjamin Kramer } // namespace 1153df186507SBenjamin Kramer 11545c0c51a9SNicolas Vasilache /// Populate the given list with patterns that convert from Vector to LLVM. 11555c0c51a9SNicolas Vasilache void mlir::populateVectorToLLVMConversionPatterns( 1156ceb1b327Saartbik LLVMTypeConverter &converter, OwningRewritePatternList &patterns, 1157ceb1b327Saartbik bool reassociateFPReductions) { 115865678d93SNicolas Vasilache MLIRContext *ctx = converter.getDialect()->getContext(); 11598345b86dSNicolas Vasilache // clang-format off 1160681f929fSNicolas Vasilache patterns.insert<VectorFMAOpNDRewritePattern, 1161681f929fSNicolas Vasilache VectorInsertStridedSliceOpDifferentRankRewritePattern, 11622d515e49SNicolas Vasilache VectorInsertStridedSliceOpSameRankRewritePattern, 11632d515e49SNicolas Vasilache VectorStridedSliceOpConversion>(ctx); 1164ceb1b327Saartbik patterns.insert<VectorReductionOpConversion>( 1165ceb1b327Saartbik ctx, converter, reassociateFPReductions); 11668345b86dSNicolas Vasilache patterns 1167ceb1b327Saartbik .insert<VectorShuffleOpConversion, 11688345b86dSNicolas Vasilache VectorExtractElementOpConversion, 11698345b86dSNicolas Vasilache VectorExtractOpConversion, 11708345b86dSNicolas Vasilache VectorFMAOp1DConversion, 11718345b86dSNicolas Vasilache VectorInsertElementOpConversion, 11728345b86dSNicolas Vasilache VectorInsertOpConversion, 11738345b86dSNicolas Vasilache VectorPrintOpConversion, 11748345b86dSNicolas Vasilache VectorTransferConversion<TransferReadOp>, 11758345b86dSNicolas Vasilache VectorTransferConversion<TransferWriteOp>, 11768345b86dSNicolas Vasilache VectorTypeCastOpConversion>(ctx, converter); 11778345b86dSNicolas Vasilache // clang-format on 11785c0c51a9SNicolas Vasilache } 11795c0c51a9SNicolas Vasilache 118063b683a8SNicolas Vasilache void mlir::populateVectorToLLVMMatrixConversionPatterns( 118163b683a8SNicolas Vasilache LLVMTypeConverter &converter, OwningRewritePatternList &patterns) { 118263b683a8SNicolas Vasilache MLIRContext *ctx = converter.getDialect()->getContext(); 118363b683a8SNicolas Vasilache patterns.insert<VectorMatmulOpConversion>(ctx, converter); 1184c295a65dSaartbik patterns.insert<VectorFlatTransposeOpConversion>(ctx, converter); 118563b683a8SNicolas Vasilache } 118663b683a8SNicolas Vasilache 11875c0c51a9SNicolas Vasilache namespace { 1188722f909fSRiver Riddle struct LowerVectorToLLVMPass 11891834ad4aSRiver Riddle : public ConvertVectorToLLVMBase<LowerVectorToLLVMPass> { 11901bfdf7c7Saartbik LowerVectorToLLVMPass(const LowerVectorToLLVMOptions &options) { 11911bfdf7c7Saartbik this->reassociateFPReductions = options.reassociateFPReductions; 11921bfdf7c7Saartbik } 1193722f909fSRiver Riddle void runOnOperation() override; 11945c0c51a9SNicolas Vasilache }; 11955c0c51a9SNicolas Vasilache } // namespace 11965c0c51a9SNicolas Vasilache 1197722f909fSRiver Riddle void LowerVectorToLLVMPass::runOnOperation() { 1198078776a6Saartbik // Perform progressive lowering of operations on slices and 1199b21c7999Saartbik // all contraction operations. Also applies folding and DCE. 1200459cf6e5Saartbik { 12015c0c51a9SNicolas Vasilache OwningRewritePatternList patterns; 1202b1c688dbSaartbik populateVectorToVectorCanonicalizationPatterns(patterns, &getContext()); 1203459cf6e5Saartbik populateVectorSlicesLoweringPatterns(patterns, &getContext()); 1204b21c7999Saartbik populateVectorContractLoweringPatterns(patterns, &getContext()); 1205a5b9316bSUday Bondhugula applyPatternsAndFoldGreedily(getOperation(), patterns); 1206459cf6e5Saartbik } 1207459cf6e5Saartbik 1208459cf6e5Saartbik // Convert to the LLVM IR dialect. 12095c0c51a9SNicolas Vasilache LLVMTypeConverter converter(&getContext()); 1210459cf6e5Saartbik OwningRewritePatternList patterns; 121163b683a8SNicolas Vasilache populateVectorToLLVMMatrixConversionPatterns(converter, patterns); 1212ceb1b327Saartbik populateVectorToLLVMConversionPatterns(converter, patterns, 1213ceb1b327Saartbik reassociateFPReductions); 1214bbf3ef85SNicolas Vasilache populateVectorToLLVMMatrixConversionPatterns(converter, patterns); 12155c0c51a9SNicolas Vasilache populateStdToLLVMConversionPatterns(converter, patterns); 12165c0c51a9SNicolas Vasilache 12172a00ae39STim Shen LLVMConversionTarget target(getContext()); 12188d67d187SRiver Riddle if (failed(applyPartialConversion(getOperation(), target, patterns))) { 12195c0c51a9SNicolas Vasilache signalPassFailure(); 12205c0c51a9SNicolas Vasilache } 12215c0c51a9SNicolas Vasilache } 12225c0c51a9SNicolas Vasilache 12231bfdf7c7Saartbik std::unique_ptr<OperationPass<ModuleOp>> 12241bfdf7c7Saartbik mlir::createConvertVectorToLLVMPass(const LowerVectorToLLVMOptions &options) { 12251bfdf7c7Saartbik return std::make_unique<LowerVectorToLLVMPass>(options); 12265c0c51a9SNicolas Vasilache } 1227