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. 1132d515e49SNicolas Vasilache // TODO(rriddle): 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 12790c01357SBenjamin Kramer namespace { 128e83b7b99Saartbik 129870c1fd4SAlex Zinenko class VectorBroadcastOpConversion : public ConvertToLLVMPattern { 130b36aaeafSAart Bik public: 131b36aaeafSAart Bik explicit VectorBroadcastOpConversion(MLIRContext *context, 132b36aaeafSAart Bik LLVMTypeConverter &typeConverter) 133870c1fd4SAlex Zinenko : ConvertToLLVMPattern(vector::BroadcastOp::getOperationName(), context, 134b36aaeafSAart Bik typeConverter) {} 135b36aaeafSAart Bik 1363145427dSRiver Riddle LogicalResult 137e62a6956SRiver Riddle matchAndRewrite(Operation *op, ArrayRef<Value> operands, 138b36aaeafSAart Bik ConversionPatternRewriter &rewriter) const override { 139b36aaeafSAart Bik auto broadcastOp = cast<vector::BroadcastOp>(op); 140b36aaeafSAart Bik VectorType dstVectorType = broadcastOp.getVectorType(); 1410f04384dSAlex Zinenko if (typeConverter.convertType(dstVectorType) == nullptr) 1423145427dSRiver Riddle return failure(); 143b36aaeafSAart Bik // Rewrite when the full vector type can be lowered (which 144b36aaeafSAart Bik // implies all 'reduced' types can be lowered too). 1451c81adf3SAart Bik auto adaptor = vector::BroadcastOpOperandAdaptor(operands); 146b36aaeafSAart Bik VectorType srcVectorType = 147b36aaeafSAart Bik broadcastOp.getSourceType().dyn_cast<VectorType>(); 148b36aaeafSAart Bik rewriter.replaceOp( 1491c81adf3SAart Bik op, expandRanks(adaptor.source(), // source value to be expanded 150b36aaeafSAart Bik op->getLoc(), // location of original broadcast 151b36aaeafSAart Bik srcVectorType, dstVectorType, rewriter)); 1523145427dSRiver Riddle return success(); 153b36aaeafSAart Bik } 154b36aaeafSAart Bik 155b36aaeafSAart Bik private: 156b36aaeafSAart Bik // Expands the given source value over all the ranks, as defined 157b36aaeafSAart Bik // by the source and destination type (a null source type denotes 158b36aaeafSAart Bik // expansion from a scalar value into a vector). 159b36aaeafSAart Bik // 160b36aaeafSAart Bik // TODO(ajcbik): consider replacing this one-pattern lowering 161b36aaeafSAart Bik // with a two-pattern lowering using other vector 162b36aaeafSAart Bik // ops once all insert/extract/shuffle operations 163fc817b09SKazuaki Ishizaki // are available with lowering implementation. 164b36aaeafSAart Bik // 165e62a6956SRiver Riddle Value expandRanks(Value value, Location loc, VectorType srcVectorType, 166b36aaeafSAart Bik VectorType dstVectorType, 167b36aaeafSAart Bik ConversionPatternRewriter &rewriter) const { 168b36aaeafSAart Bik assert((dstVectorType != nullptr) && "invalid result type in broadcast"); 169b36aaeafSAart Bik // Determine rank of source and destination. 170b36aaeafSAart Bik int64_t srcRank = srcVectorType ? srcVectorType.getRank() : 0; 171b36aaeafSAart Bik int64_t dstRank = dstVectorType.getRank(); 172b36aaeafSAart Bik int64_t curDim = dstVectorType.getDimSize(0); 173b36aaeafSAart Bik if (srcRank < dstRank) 174b36aaeafSAart Bik // Duplicate this rank. 175b36aaeafSAart Bik return duplicateOneRank(value, loc, srcVectorType, dstVectorType, dstRank, 176b36aaeafSAart Bik curDim, rewriter); 177b36aaeafSAart Bik // If all trailing dimensions are the same, the broadcast consists of 178b36aaeafSAart Bik // simply passing through the source value and we are done. Otherwise, 179b36aaeafSAart Bik // any non-matching dimension forces a stretch along this rank. 180b36aaeafSAart Bik assert((srcVectorType != nullptr) && (srcRank > 0) && 181b36aaeafSAart Bik (srcRank == dstRank) && "invalid rank in broadcast"); 182b36aaeafSAart Bik for (int64_t r = 0; r < dstRank; r++) { 183b36aaeafSAart Bik if (srcVectorType.getDimSize(r) != dstVectorType.getDimSize(r)) { 184b36aaeafSAart Bik return stretchOneRank(value, loc, srcVectorType, dstVectorType, dstRank, 185b36aaeafSAart Bik curDim, rewriter); 186b36aaeafSAart Bik } 187b36aaeafSAart Bik } 188b36aaeafSAart Bik return value; 189b36aaeafSAart Bik } 190b36aaeafSAart Bik 191b36aaeafSAart Bik // Picks the best way to duplicate a single rank. For the 1-D case, a 192b36aaeafSAart Bik // single insert-elt/shuffle is the most efficient expansion. For higher 193b36aaeafSAart Bik // dimensions, however, we need dim x insert-values on a new broadcast 194b36aaeafSAart Bik // with one less leading dimension, which will be lowered "recursively" 195b36aaeafSAart Bik // to matching LLVM IR. 196b36aaeafSAart Bik // For example: 197b36aaeafSAart Bik // v = broadcast s : f32 to vector<4x2xf32> 198b36aaeafSAart Bik // becomes: 199b36aaeafSAart Bik // x = broadcast s : f32 to vector<2xf32> 200b36aaeafSAart Bik // v = [x,x,x,x] 201b36aaeafSAart Bik // becomes: 202b36aaeafSAart Bik // x = [s,s] 203b36aaeafSAart Bik // v = [x,x,x,x] 204e62a6956SRiver Riddle Value duplicateOneRank(Value value, Location loc, VectorType srcVectorType, 205e62a6956SRiver Riddle VectorType dstVectorType, int64_t rank, int64_t dim, 206b36aaeafSAart Bik ConversionPatternRewriter &rewriter) const { 2070f04384dSAlex Zinenko Type llvmType = typeConverter.convertType(dstVectorType); 208b36aaeafSAart Bik assert((llvmType != nullptr) && "unlowerable vector type"); 209b36aaeafSAart Bik if (rank == 1) { 210e62a6956SRiver Riddle Value undef = rewriter.create<LLVM::UndefOp>(loc, llvmType); 2110f04384dSAlex Zinenko Value expand = insertOne(rewriter, typeConverter, loc, undef, value, 2120f04384dSAlex Zinenko llvmType, rank, 0); 213b36aaeafSAart Bik SmallVector<int32_t, 4> zeroValues(dim, 0); 214b36aaeafSAart Bik return rewriter.create<LLVM::ShuffleVectorOp>( 215b36aaeafSAart Bik loc, expand, undef, rewriter.getI32ArrayAttr(zeroValues)); 216b36aaeafSAart Bik } 217e62a6956SRiver Riddle Value expand = expandRanks(value, loc, srcVectorType, 2189826fe5cSAart Bik reducedVectorTypeFront(dstVectorType), rewriter); 219e62a6956SRiver Riddle Value result = rewriter.create<LLVM::UndefOp>(loc, llvmType); 220b36aaeafSAart Bik for (int64_t d = 0; d < dim; ++d) { 2210f04384dSAlex Zinenko result = insertOne(rewriter, typeConverter, loc, result, expand, llvmType, 2220f04384dSAlex Zinenko rank, d); 223b36aaeafSAart Bik } 224b36aaeafSAart Bik return result; 225b36aaeafSAart Bik } 226b36aaeafSAart Bik 227b36aaeafSAart Bik // Picks the best way to stretch a single rank. For the 1-D case, a 228b36aaeafSAart Bik // single insert-elt/shuffle is the most efficient expansion when at 229b36aaeafSAart Bik // a stretch. Otherwise, every dimension needs to be expanded 230b36aaeafSAart Bik // individually and individually inserted in the resulting vector. 231b36aaeafSAart Bik // For example: 232b36aaeafSAart Bik // v = broadcast w : vector<4x1x2xf32> to vector<4x2x2xf32> 233b36aaeafSAart Bik // becomes: 234b36aaeafSAart Bik // a = broadcast w[0] : vector<1x2xf32> to vector<2x2xf32> 235b36aaeafSAart Bik // b = broadcast w[1] : vector<1x2xf32> to vector<2x2xf32> 236b36aaeafSAart Bik // c = broadcast w[2] : vector<1x2xf32> to vector<2x2xf32> 237b36aaeafSAart Bik // d = broadcast w[3] : vector<1x2xf32> to vector<2x2xf32> 238b36aaeafSAart Bik // v = [a,b,c,d] 239b36aaeafSAart Bik // becomes: 240b36aaeafSAart Bik // x = broadcast w[0][0] : vector<2xf32> to vector <2x2xf32> 241b36aaeafSAart Bik // y = broadcast w[1][0] : vector<2xf32> to vector <2x2xf32> 242b36aaeafSAart Bik // a = [x, y] 243b36aaeafSAart Bik // etc. 244e62a6956SRiver Riddle Value stretchOneRank(Value value, Location loc, VectorType srcVectorType, 245e62a6956SRiver Riddle VectorType dstVectorType, int64_t rank, int64_t dim, 246b36aaeafSAart Bik ConversionPatternRewriter &rewriter) const { 2470f04384dSAlex Zinenko Type llvmType = typeConverter.convertType(dstVectorType); 248b36aaeafSAart Bik assert((llvmType != nullptr) && "unlowerable vector type"); 249e62a6956SRiver Riddle Value result = rewriter.create<LLVM::UndefOp>(loc, llvmType); 250b36aaeafSAart Bik bool atStretch = dim != srcVectorType.getDimSize(0); 251b36aaeafSAart Bik if (rank == 1) { 2521c81adf3SAart Bik assert(atStretch); 2530f04384dSAlex Zinenko Type redLlvmType = 2540f04384dSAlex Zinenko typeConverter.convertType(dstVectorType.getElementType()); 255e62a6956SRiver Riddle Value one = 2560f04384dSAlex Zinenko extractOne(rewriter, typeConverter, loc, value, redLlvmType, rank, 0); 2570f04384dSAlex Zinenko Value expand = insertOne(rewriter, typeConverter, loc, result, one, 2580f04384dSAlex Zinenko llvmType, rank, 0); 259b36aaeafSAart Bik SmallVector<int32_t, 4> zeroValues(dim, 0); 260b36aaeafSAart Bik return rewriter.create<LLVM::ShuffleVectorOp>( 261b36aaeafSAart Bik loc, expand, result, rewriter.getI32ArrayAttr(zeroValues)); 262b36aaeafSAart Bik } 2639826fe5cSAart Bik VectorType redSrcType = reducedVectorTypeFront(srcVectorType); 2649826fe5cSAart Bik VectorType redDstType = reducedVectorTypeFront(dstVectorType); 2650f04384dSAlex Zinenko Type redLlvmType = typeConverter.convertType(redSrcType); 266b36aaeafSAart Bik for (int64_t d = 0; d < dim; ++d) { 267b36aaeafSAart Bik int64_t pos = atStretch ? 0 : d; 2680f04384dSAlex Zinenko Value one = extractOne(rewriter, typeConverter, loc, value, redLlvmType, 2690f04384dSAlex Zinenko rank, pos); 270e62a6956SRiver Riddle Value expand = expandRanks(one, loc, redSrcType, redDstType, rewriter); 2710f04384dSAlex Zinenko result = insertOne(rewriter, typeConverter, loc, result, expand, llvmType, 2720f04384dSAlex Zinenko rank, d); 273b36aaeafSAart Bik } 274b36aaeafSAart Bik return result; 275b36aaeafSAart Bik } 2761c81adf3SAart Bik }; 277b36aaeafSAart Bik 27863b683a8SNicolas Vasilache /// Conversion pattern for a vector.matrix_multiply. 27963b683a8SNicolas Vasilache /// This is lowered directly to the proper llvm.intr.matrix.multiply. 28063b683a8SNicolas Vasilache class VectorMatmulOpConversion : public ConvertToLLVMPattern { 28163b683a8SNicolas Vasilache public: 28263b683a8SNicolas Vasilache explicit VectorMatmulOpConversion(MLIRContext *context, 28363b683a8SNicolas Vasilache LLVMTypeConverter &typeConverter) 28463b683a8SNicolas Vasilache : ConvertToLLVMPattern(vector::MatmulOp::getOperationName(), context, 28563b683a8SNicolas Vasilache typeConverter) {} 28663b683a8SNicolas Vasilache 2873145427dSRiver Riddle LogicalResult 28863b683a8SNicolas Vasilache matchAndRewrite(Operation *op, ArrayRef<Value> operands, 28963b683a8SNicolas Vasilache ConversionPatternRewriter &rewriter) const override { 29063b683a8SNicolas Vasilache auto matmulOp = cast<vector::MatmulOp>(op); 29163b683a8SNicolas Vasilache auto adaptor = vector::MatmulOpOperandAdaptor(operands); 29263b683a8SNicolas Vasilache rewriter.replaceOpWithNewOp<LLVM::MatrixMultiplyOp>( 29363b683a8SNicolas Vasilache op, typeConverter.convertType(matmulOp.res().getType()), adaptor.lhs(), 29463b683a8SNicolas Vasilache adaptor.rhs(), matmulOp.lhs_rows(), matmulOp.lhs_columns(), 29563b683a8SNicolas Vasilache matmulOp.rhs_columns()); 2963145427dSRiver Riddle return success(); 29763b683a8SNicolas Vasilache } 29863b683a8SNicolas Vasilache }; 29963b683a8SNicolas Vasilache 300870c1fd4SAlex Zinenko class VectorReductionOpConversion : public ConvertToLLVMPattern { 301e83b7b99Saartbik public: 302e83b7b99Saartbik explicit VectorReductionOpConversion(MLIRContext *context, 303e83b7b99Saartbik LLVMTypeConverter &typeConverter) 304870c1fd4SAlex Zinenko : ConvertToLLVMPattern(vector::ReductionOp::getOperationName(), context, 305e83b7b99Saartbik typeConverter) {} 306e83b7b99Saartbik 3073145427dSRiver Riddle LogicalResult 308e83b7b99Saartbik matchAndRewrite(Operation *op, ArrayRef<Value> operands, 309e83b7b99Saartbik ConversionPatternRewriter &rewriter) const override { 310e83b7b99Saartbik auto reductionOp = cast<vector::ReductionOp>(op); 311e83b7b99Saartbik auto kind = reductionOp.kind(); 312e83b7b99Saartbik Type eltType = reductionOp.dest().getType(); 3130f04384dSAlex Zinenko Type llvmType = typeConverter.convertType(eltType); 31435b68527SLei Zhang if (eltType.isSignlessInteger(32) || eltType.isSignlessInteger(64)) { 315e83b7b99Saartbik // Integer reductions: add/mul/min/max/and/or/xor. 316e83b7b99Saartbik if (kind == "add") 317e83b7b99Saartbik rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_add>( 318e83b7b99Saartbik op, llvmType, operands[0]); 319e83b7b99Saartbik else if (kind == "mul") 320e83b7b99Saartbik rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_mul>( 321e83b7b99Saartbik op, llvmType, operands[0]); 322e83b7b99Saartbik else if (kind == "min") 323e83b7b99Saartbik rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_smin>( 324e83b7b99Saartbik op, llvmType, operands[0]); 325e83b7b99Saartbik else if (kind == "max") 326e83b7b99Saartbik rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_smax>( 327e83b7b99Saartbik op, llvmType, operands[0]); 328e83b7b99Saartbik else if (kind == "and") 329e83b7b99Saartbik rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_and>( 330e83b7b99Saartbik op, llvmType, operands[0]); 331e83b7b99Saartbik else if (kind == "or") 332e83b7b99Saartbik rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_or>( 333e83b7b99Saartbik op, llvmType, operands[0]); 334e83b7b99Saartbik else if (kind == "xor") 335e83b7b99Saartbik rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_xor>( 336e83b7b99Saartbik op, llvmType, operands[0]); 337e83b7b99Saartbik else 3383145427dSRiver Riddle return failure(); 3393145427dSRiver Riddle return success(); 340e83b7b99Saartbik 341e83b7b99Saartbik } else if (eltType.isF32() || eltType.isF64()) { 342e83b7b99Saartbik // Floating-point reductions: add/mul/min/max 343e83b7b99Saartbik if (kind == "add") { 3440d924700Saartbik // Optional accumulator (or zero). 3450d924700Saartbik Value acc = operands.size() > 1 ? operands[1] 3460d924700Saartbik : rewriter.create<LLVM::ConstantOp>( 3470d924700Saartbik op->getLoc(), llvmType, 3480d924700Saartbik rewriter.getZeroAttr(eltType)); 349e83b7b99Saartbik rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_v2_fadd>( 3500d924700Saartbik op, llvmType, acc, operands[0]); 351e83b7b99Saartbik } else if (kind == "mul") { 3520d924700Saartbik // Optional accumulator (or one). 3530d924700Saartbik Value acc = operands.size() > 1 3540d924700Saartbik ? operands[1] 3550d924700Saartbik : rewriter.create<LLVM::ConstantOp>( 3560d924700Saartbik op->getLoc(), llvmType, 3570d924700Saartbik rewriter.getFloatAttr(eltType, 1.0)); 358e83b7b99Saartbik rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_v2_fmul>( 3590d924700Saartbik op, llvmType, acc, operands[0]); 360e83b7b99Saartbik } else if (kind == "min") 361e83b7b99Saartbik rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_fmin>( 362e83b7b99Saartbik op, llvmType, operands[0]); 363e83b7b99Saartbik else if (kind == "max") 364e83b7b99Saartbik rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_fmax>( 365e83b7b99Saartbik op, llvmType, operands[0]); 366e83b7b99Saartbik else 3673145427dSRiver Riddle return failure(); 3683145427dSRiver Riddle return success(); 369e83b7b99Saartbik } 3703145427dSRiver Riddle return failure(); 371e83b7b99Saartbik } 372e83b7b99Saartbik }; 373e83b7b99Saartbik 374870c1fd4SAlex Zinenko class VectorShuffleOpConversion : public ConvertToLLVMPattern { 3751c81adf3SAart Bik public: 3761c81adf3SAart Bik explicit VectorShuffleOpConversion(MLIRContext *context, 3771c81adf3SAart Bik LLVMTypeConverter &typeConverter) 378870c1fd4SAlex Zinenko : ConvertToLLVMPattern(vector::ShuffleOp::getOperationName(), context, 3791c81adf3SAart Bik typeConverter) {} 3801c81adf3SAart Bik 3813145427dSRiver Riddle LogicalResult 382e62a6956SRiver Riddle matchAndRewrite(Operation *op, ArrayRef<Value> operands, 3831c81adf3SAart Bik ConversionPatternRewriter &rewriter) const override { 3841c81adf3SAart Bik auto loc = op->getLoc(); 3851c81adf3SAart Bik auto adaptor = vector::ShuffleOpOperandAdaptor(operands); 3861c81adf3SAart Bik auto shuffleOp = cast<vector::ShuffleOp>(op); 3871c81adf3SAart Bik auto v1Type = shuffleOp.getV1VectorType(); 3881c81adf3SAart Bik auto v2Type = shuffleOp.getV2VectorType(); 3891c81adf3SAart Bik auto vectorType = shuffleOp.getVectorType(); 3900f04384dSAlex Zinenko Type llvmType = typeConverter.convertType(vectorType); 3911c81adf3SAart Bik auto maskArrayAttr = shuffleOp.mask(); 3921c81adf3SAart Bik 3931c81adf3SAart Bik // Bail if result type cannot be lowered. 3941c81adf3SAart Bik if (!llvmType) 3953145427dSRiver Riddle return failure(); 3961c81adf3SAart Bik 3971c81adf3SAart Bik // Get rank and dimension sizes. 3981c81adf3SAart Bik int64_t rank = vectorType.getRank(); 3991c81adf3SAart Bik assert(v1Type.getRank() == rank); 4001c81adf3SAart Bik assert(v2Type.getRank() == rank); 4011c81adf3SAart Bik int64_t v1Dim = v1Type.getDimSize(0); 4021c81adf3SAart Bik 4031c81adf3SAart Bik // For rank 1, where both operands have *exactly* the same vector type, 4041c81adf3SAart Bik // there is direct shuffle support in LLVM. Use it! 4051c81adf3SAart Bik if (rank == 1 && v1Type == v2Type) { 406e62a6956SRiver Riddle Value shuffle = rewriter.create<LLVM::ShuffleVectorOp>( 4071c81adf3SAart Bik loc, adaptor.v1(), adaptor.v2(), maskArrayAttr); 4081c81adf3SAart Bik rewriter.replaceOp(op, shuffle); 4093145427dSRiver Riddle return success(); 410b36aaeafSAart Bik } 411b36aaeafSAart Bik 4121c81adf3SAart Bik // For all other cases, insert the individual values individually. 413e62a6956SRiver Riddle Value insert = rewriter.create<LLVM::UndefOp>(loc, llvmType); 4141c81adf3SAart Bik int64_t insPos = 0; 4151c81adf3SAart Bik for (auto en : llvm::enumerate(maskArrayAttr)) { 4161c81adf3SAart Bik int64_t extPos = en.value().cast<IntegerAttr>().getInt(); 417e62a6956SRiver Riddle Value value = adaptor.v1(); 4181c81adf3SAart Bik if (extPos >= v1Dim) { 4191c81adf3SAart Bik extPos -= v1Dim; 4201c81adf3SAart Bik value = adaptor.v2(); 421b36aaeafSAart Bik } 4220f04384dSAlex Zinenko Value extract = extractOne(rewriter, typeConverter, loc, value, llvmType, 4230f04384dSAlex Zinenko rank, extPos); 4240f04384dSAlex Zinenko insert = insertOne(rewriter, typeConverter, loc, insert, extract, 4250f04384dSAlex Zinenko llvmType, rank, insPos++); 4261c81adf3SAart Bik } 4271c81adf3SAart Bik rewriter.replaceOp(op, insert); 4283145427dSRiver Riddle return success(); 429b36aaeafSAart Bik } 430b36aaeafSAart Bik }; 431b36aaeafSAart Bik 432870c1fd4SAlex Zinenko class VectorExtractElementOpConversion : public ConvertToLLVMPattern { 433cd5dab8aSAart Bik public: 434cd5dab8aSAart Bik explicit VectorExtractElementOpConversion(MLIRContext *context, 435cd5dab8aSAart Bik LLVMTypeConverter &typeConverter) 436870c1fd4SAlex Zinenko : ConvertToLLVMPattern(vector::ExtractElementOp::getOperationName(), 437870c1fd4SAlex Zinenko context, typeConverter) {} 438cd5dab8aSAart Bik 4393145427dSRiver Riddle LogicalResult 440e62a6956SRiver Riddle matchAndRewrite(Operation *op, ArrayRef<Value> operands, 441cd5dab8aSAart Bik ConversionPatternRewriter &rewriter) const override { 442cd5dab8aSAart Bik auto adaptor = vector::ExtractElementOpOperandAdaptor(operands); 443cd5dab8aSAart Bik auto extractEltOp = cast<vector::ExtractElementOp>(op); 444cd5dab8aSAart Bik auto vectorType = extractEltOp.getVectorType(); 4450f04384dSAlex Zinenko auto llvmType = typeConverter.convertType(vectorType.getElementType()); 446cd5dab8aSAart Bik 447cd5dab8aSAart Bik // Bail if result type cannot be lowered. 448cd5dab8aSAart Bik if (!llvmType) 4493145427dSRiver Riddle return failure(); 450cd5dab8aSAart Bik 451cd5dab8aSAart Bik rewriter.replaceOpWithNewOp<LLVM::ExtractElementOp>( 452cd5dab8aSAart Bik op, llvmType, adaptor.vector(), adaptor.position()); 4533145427dSRiver Riddle return success(); 454cd5dab8aSAart Bik } 455cd5dab8aSAart Bik }; 456cd5dab8aSAart Bik 457870c1fd4SAlex Zinenko class VectorExtractOpConversion : public ConvertToLLVMPattern { 4585c0c51a9SNicolas Vasilache public: 4599826fe5cSAart Bik explicit VectorExtractOpConversion(MLIRContext *context, 4605c0c51a9SNicolas Vasilache LLVMTypeConverter &typeConverter) 461870c1fd4SAlex Zinenko : ConvertToLLVMPattern(vector::ExtractOp::getOperationName(), context, 4625c0c51a9SNicolas Vasilache typeConverter) {} 4635c0c51a9SNicolas Vasilache 4643145427dSRiver Riddle LogicalResult 465e62a6956SRiver Riddle matchAndRewrite(Operation *op, ArrayRef<Value> operands, 4665c0c51a9SNicolas Vasilache ConversionPatternRewriter &rewriter) const override { 4675c0c51a9SNicolas Vasilache auto loc = op->getLoc(); 468d37f2725SAart Bik auto adaptor = vector::ExtractOpOperandAdaptor(operands); 469d37f2725SAart Bik auto extractOp = cast<vector::ExtractOp>(op); 4709826fe5cSAart Bik auto vectorType = extractOp.getVectorType(); 4712bdf33ccSRiver Riddle auto resultType = extractOp.getResult().getType(); 4720f04384dSAlex Zinenko auto llvmResultType = typeConverter.convertType(resultType); 4735c0c51a9SNicolas Vasilache auto positionArrayAttr = extractOp.position(); 4749826fe5cSAart Bik 4759826fe5cSAart Bik // Bail if result type cannot be lowered. 4769826fe5cSAart Bik if (!llvmResultType) 4773145427dSRiver Riddle return failure(); 4789826fe5cSAart Bik 4795c0c51a9SNicolas Vasilache // One-shot extraction of vector from array (only requires extractvalue). 4805c0c51a9SNicolas Vasilache if (resultType.isa<VectorType>()) { 481e62a6956SRiver Riddle Value extracted = rewriter.create<LLVM::ExtractValueOp>( 4825c0c51a9SNicolas Vasilache loc, llvmResultType, adaptor.vector(), positionArrayAttr); 4835c0c51a9SNicolas Vasilache rewriter.replaceOp(op, extracted); 4843145427dSRiver Riddle return success(); 4855c0c51a9SNicolas Vasilache } 4865c0c51a9SNicolas Vasilache 4879826fe5cSAart Bik // Potential extraction of 1-D vector from array. 4885c0c51a9SNicolas Vasilache auto *context = op->getContext(); 489e62a6956SRiver Riddle Value extracted = adaptor.vector(); 4905c0c51a9SNicolas Vasilache auto positionAttrs = positionArrayAttr.getValue(); 4915c0c51a9SNicolas Vasilache if (positionAttrs.size() > 1) { 4929826fe5cSAart Bik auto oneDVectorType = reducedVectorTypeBack(vectorType); 4935c0c51a9SNicolas Vasilache auto nMinusOnePositionAttrs = 4945c0c51a9SNicolas Vasilache ArrayAttr::get(positionAttrs.drop_back(), context); 4955c0c51a9SNicolas Vasilache extracted = rewriter.create<LLVM::ExtractValueOp>( 4960f04384dSAlex Zinenko loc, typeConverter.convertType(oneDVectorType), extracted, 4975c0c51a9SNicolas Vasilache nMinusOnePositionAttrs); 4985c0c51a9SNicolas Vasilache } 4995c0c51a9SNicolas Vasilache 5005c0c51a9SNicolas Vasilache // Remaining extraction of element from 1-D LLVM vector 5015c0c51a9SNicolas Vasilache auto position = positionAttrs.back().cast<IntegerAttr>(); 5020f04384dSAlex Zinenko auto i64Type = LLVM::LLVMType::getInt64Ty(typeConverter.getDialect()); 5031d47564aSAart Bik auto constant = rewriter.create<LLVM::ConstantOp>(loc, i64Type, position); 5045c0c51a9SNicolas Vasilache extracted = 5055c0c51a9SNicolas Vasilache rewriter.create<LLVM::ExtractElementOp>(loc, extracted, constant); 5065c0c51a9SNicolas Vasilache rewriter.replaceOp(op, extracted); 5075c0c51a9SNicolas Vasilache 5083145427dSRiver Riddle return success(); 5095c0c51a9SNicolas Vasilache } 5105c0c51a9SNicolas Vasilache }; 5115c0c51a9SNicolas Vasilache 512681f929fSNicolas Vasilache /// Conversion pattern that turns a vector.fma on a 1-D vector 513681f929fSNicolas Vasilache /// into an llvm.intr.fmuladd. This is a trivial 1-1 conversion. 514681f929fSNicolas Vasilache /// This does not match vectors of n >= 2 rank. 515681f929fSNicolas Vasilache /// 516681f929fSNicolas Vasilache /// Example: 517681f929fSNicolas Vasilache /// ``` 518681f929fSNicolas Vasilache /// vector.fma %a, %a, %a : vector<8xf32> 519681f929fSNicolas Vasilache /// ``` 520681f929fSNicolas Vasilache /// is converted to: 521681f929fSNicolas Vasilache /// ``` 522681f929fSNicolas Vasilache /// llvm.intr.fma %va, %va, %va: 523681f929fSNicolas Vasilache /// (!llvm<"<8 x float>">, !llvm<"<8 x float>">, !llvm<"<8 x float>">) 524681f929fSNicolas Vasilache /// -> !llvm<"<8 x float>"> 525681f929fSNicolas Vasilache /// ``` 526870c1fd4SAlex Zinenko class VectorFMAOp1DConversion : public ConvertToLLVMPattern { 527681f929fSNicolas Vasilache public: 528681f929fSNicolas Vasilache explicit VectorFMAOp1DConversion(MLIRContext *context, 529681f929fSNicolas Vasilache LLVMTypeConverter &typeConverter) 530870c1fd4SAlex Zinenko : ConvertToLLVMPattern(vector::FMAOp::getOperationName(), context, 531681f929fSNicolas Vasilache typeConverter) {} 532681f929fSNicolas Vasilache 5333145427dSRiver Riddle LogicalResult 534681f929fSNicolas Vasilache matchAndRewrite(Operation *op, ArrayRef<Value> operands, 535681f929fSNicolas Vasilache ConversionPatternRewriter &rewriter) const override { 536681f929fSNicolas Vasilache auto adaptor = vector::FMAOpOperandAdaptor(operands); 537681f929fSNicolas Vasilache vector::FMAOp fmaOp = cast<vector::FMAOp>(op); 538681f929fSNicolas Vasilache VectorType vType = fmaOp.getVectorType(); 539681f929fSNicolas Vasilache if (vType.getRank() != 1) 5403145427dSRiver Riddle return failure(); 541681f929fSNicolas Vasilache rewriter.replaceOpWithNewOp<LLVM::FMAOp>(op, adaptor.lhs(), adaptor.rhs(), 542681f929fSNicolas Vasilache adaptor.acc()); 5433145427dSRiver Riddle return success(); 544681f929fSNicolas Vasilache } 545681f929fSNicolas Vasilache }; 546681f929fSNicolas Vasilache 547870c1fd4SAlex Zinenko class VectorInsertElementOpConversion : public ConvertToLLVMPattern { 548cd5dab8aSAart Bik public: 549cd5dab8aSAart Bik explicit VectorInsertElementOpConversion(MLIRContext *context, 550cd5dab8aSAart Bik LLVMTypeConverter &typeConverter) 551870c1fd4SAlex Zinenko : ConvertToLLVMPattern(vector::InsertElementOp::getOperationName(), 552870c1fd4SAlex Zinenko context, typeConverter) {} 553cd5dab8aSAart Bik 5543145427dSRiver Riddle LogicalResult 555e62a6956SRiver Riddle matchAndRewrite(Operation *op, ArrayRef<Value> operands, 556cd5dab8aSAart Bik ConversionPatternRewriter &rewriter) const override { 557cd5dab8aSAart Bik auto adaptor = vector::InsertElementOpOperandAdaptor(operands); 558cd5dab8aSAart Bik auto insertEltOp = cast<vector::InsertElementOp>(op); 559cd5dab8aSAart Bik auto vectorType = insertEltOp.getDestVectorType(); 5600f04384dSAlex Zinenko auto llvmType = typeConverter.convertType(vectorType); 561cd5dab8aSAart Bik 562cd5dab8aSAart Bik // Bail if result type cannot be lowered. 563cd5dab8aSAart Bik if (!llvmType) 5643145427dSRiver Riddle return failure(); 565cd5dab8aSAart Bik 566cd5dab8aSAart Bik rewriter.replaceOpWithNewOp<LLVM::InsertElementOp>( 567cd5dab8aSAart Bik op, llvmType, adaptor.dest(), adaptor.source(), adaptor.position()); 5683145427dSRiver Riddle return success(); 569cd5dab8aSAart Bik } 570cd5dab8aSAart Bik }; 571cd5dab8aSAart Bik 572870c1fd4SAlex Zinenko class VectorInsertOpConversion : public ConvertToLLVMPattern { 5739826fe5cSAart Bik public: 5749826fe5cSAart Bik explicit VectorInsertOpConversion(MLIRContext *context, 5759826fe5cSAart Bik LLVMTypeConverter &typeConverter) 576870c1fd4SAlex Zinenko : ConvertToLLVMPattern(vector::InsertOp::getOperationName(), context, 5779826fe5cSAart Bik typeConverter) {} 5789826fe5cSAart Bik 5793145427dSRiver Riddle LogicalResult 580e62a6956SRiver Riddle matchAndRewrite(Operation *op, ArrayRef<Value> operands, 5819826fe5cSAart Bik ConversionPatternRewriter &rewriter) const override { 5829826fe5cSAart Bik auto loc = op->getLoc(); 5839826fe5cSAart Bik auto adaptor = vector::InsertOpOperandAdaptor(operands); 5849826fe5cSAart Bik auto insertOp = cast<vector::InsertOp>(op); 5859826fe5cSAart Bik auto sourceType = insertOp.getSourceType(); 5869826fe5cSAart Bik auto destVectorType = insertOp.getDestVectorType(); 5870f04384dSAlex Zinenko auto llvmResultType = typeConverter.convertType(destVectorType); 5889826fe5cSAart Bik auto positionArrayAttr = insertOp.position(); 5899826fe5cSAart Bik 5909826fe5cSAart Bik // Bail if result type cannot be lowered. 5919826fe5cSAart Bik if (!llvmResultType) 5923145427dSRiver Riddle return failure(); 5939826fe5cSAart Bik 5949826fe5cSAart Bik // One-shot insertion of a vector into an array (only requires insertvalue). 5959826fe5cSAart Bik if (sourceType.isa<VectorType>()) { 596e62a6956SRiver Riddle Value inserted = rewriter.create<LLVM::InsertValueOp>( 5979826fe5cSAart Bik loc, llvmResultType, adaptor.dest(), adaptor.source(), 5989826fe5cSAart Bik positionArrayAttr); 5999826fe5cSAart Bik rewriter.replaceOp(op, inserted); 6003145427dSRiver Riddle return success(); 6019826fe5cSAart Bik } 6029826fe5cSAart Bik 6039826fe5cSAart Bik // Potential extraction of 1-D vector from array. 6049826fe5cSAart Bik auto *context = op->getContext(); 605e62a6956SRiver Riddle Value extracted = adaptor.dest(); 6069826fe5cSAart Bik auto positionAttrs = positionArrayAttr.getValue(); 6079826fe5cSAart Bik auto position = positionAttrs.back().cast<IntegerAttr>(); 6089826fe5cSAart Bik auto oneDVectorType = destVectorType; 6099826fe5cSAart Bik if (positionAttrs.size() > 1) { 6109826fe5cSAart Bik oneDVectorType = reducedVectorTypeBack(destVectorType); 6119826fe5cSAart Bik auto nMinusOnePositionAttrs = 6129826fe5cSAart Bik ArrayAttr::get(positionAttrs.drop_back(), context); 6139826fe5cSAart Bik extracted = rewriter.create<LLVM::ExtractValueOp>( 6140f04384dSAlex Zinenko loc, typeConverter.convertType(oneDVectorType), extracted, 6159826fe5cSAart Bik nMinusOnePositionAttrs); 6169826fe5cSAart Bik } 6179826fe5cSAart Bik 6189826fe5cSAart Bik // Insertion of an element into a 1-D LLVM vector. 6190f04384dSAlex Zinenko auto i64Type = LLVM::LLVMType::getInt64Ty(typeConverter.getDialect()); 6201d47564aSAart Bik auto constant = rewriter.create<LLVM::ConstantOp>(loc, i64Type, position); 621e62a6956SRiver Riddle Value inserted = rewriter.create<LLVM::InsertElementOp>( 6220f04384dSAlex Zinenko loc, typeConverter.convertType(oneDVectorType), extracted, 6230f04384dSAlex Zinenko adaptor.source(), constant); 6249826fe5cSAart Bik 6259826fe5cSAart Bik // Potential insertion of resulting 1-D vector into array. 6269826fe5cSAart Bik if (positionAttrs.size() > 1) { 6279826fe5cSAart Bik auto nMinusOnePositionAttrs = 6289826fe5cSAart Bik ArrayAttr::get(positionAttrs.drop_back(), context); 6299826fe5cSAart Bik inserted = rewriter.create<LLVM::InsertValueOp>(loc, llvmResultType, 6309826fe5cSAart Bik adaptor.dest(), inserted, 6319826fe5cSAart Bik nMinusOnePositionAttrs); 6329826fe5cSAart Bik } 6339826fe5cSAart Bik 6349826fe5cSAart Bik rewriter.replaceOp(op, inserted); 6353145427dSRiver Riddle return success(); 6369826fe5cSAart Bik } 6379826fe5cSAart Bik }; 6389826fe5cSAart Bik 639681f929fSNicolas Vasilache /// Rank reducing rewrite for n-D FMA into (n-1)-D FMA where n > 1. 640681f929fSNicolas Vasilache /// 641681f929fSNicolas Vasilache /// Example: 642681f929fSNicolas Vasilache /// ``` 643681f929fSNicolas Vasilache /// %d = vector.fma %a, %b, %c : vector<2x4xf32> 644681f929fSNicolas Vasilache /// ``` 645681f929fSNicolas Vasilache /// is rewritten into: 646681f929fSNicolas Vasilache /// ``` 647681f929fSNicolas Vasilache /// %r = splat %f0: vector<2x4xf32> 648681f929fSNicolas Vasilache /// %va = vector.extractvalue %a[0] : vector<2x4xf32> 649681f929fSNicolas Vasilache /// %vb = vector.extractvalue %b[0] : vector<2x4xf32> 650681f929fSNicolas Vasilache /// %vc = vector.extractvalue %c[0] : vector<2x4xf32> 651681f929fSNicolas Vasilache /// %vd = vector.fma %va, %vb, %vc : vector<4xf32> 652681f929fSNicolas Vasilache /// %r2 = vector.insertvalue %vd, %r[0] : vector<4xf32> into vector<2x4xf32> 653681f929fSNicolas Vasilache /// %va2 = vector.extractvalue %a2[1] : vector<2x4xf32> 654681f929fSNicolas Vasilache /// %vb2 = vector.extractvalue %b2[1] : vector<2x4xf32> 655681f929fSNicolas Vasilache /// %vc2 = vector.extractvalue %c2[1] : vector<2x4xf32> 656681f929fSNicolas Vasilache /// %vd2 = vector.fma %va2, %vb2, %vc2 : vector<4xf32> 657681f929fSNicolas Vasilache /// %r3 = vector.insertvalue %vd2, %r2[1] : vector<4xf32> into vector<2x4xf32> 658681f929fSNicolas Vasilache /// // %r3 holds the final value. 659681f929fSNicolas Vasilache /// ``` 660681f929fSNicolas Vasilache class VectorFMAOpNDRewritePattern : public OpRewritePattern<FMAOp> { 661681f929fSNicolas Vasilache public: 662681f929fSNicolas Vasilache using OpRewritePattern<FMAOp>::OpRewritePattern; 663681f929fSNicolas Vasilache 6643145427dSRiver Riddle LogicalResult matchAndRewrite(FMAOp op, 665681f929fSNicolas Vasilache PatternRewriter &rewriter) const override { 666681f929fSNicolas Vasilache auto vType = op.getVectorType(); 667681f929fSNicolas Vasilache if (vType.getRank() < 2) 6683145427dSRiver Riddle return failure(); 669681f929fSNicolas Vasilache 670681f929fSNicolas Vasilache auto loc = op.getLoc(); 671681f929fSNicolas Vasilache auto elemType = vType.getElementType(); 672681f929fSNicolas Vasilache Value zero = rewriter.create<ConstantOp>(loc, elemType, 673681f929fSNicolas Vasilache rewriter.getZeroAttr(elemType)); 674681f929fSNicolas Vasilache Value desc = rewriter.create<SplatOp>(loc, vType, zero); 675681f929fSNicolas Vasilache for (int64_t i = 0, e = vType.getShape().front(); i != e; ++i) { 676681f929fSNicolas Vasilache Value extrLHS = rewriter.create<ExtractOp>(loc, op.lhs(), i); 677681f929fSNicolas Vasilache Value extrRHS = rewriter.create<ExtractOp>(loc, op.rhs(), i); 678681f929fSNicolas Vasilache Value extrACC = rewriter.create<ExtractOp>(loc, op.acc(), i); 679681f929fSNicolas Vasilache Value fma = rewriter.create<FMAOp>(loc, extrLHS, extrRHS, extrACC); 680681f929fSNicolas Vasilache desc = rewriter.create<InsertOp>(loc, fma, desc, i); 681681f929fSNicolas Vasilache } 682681f929fSNicolas Vasilache rewriter.replaceOp(op, desc); 6833145427dSRiver Riddle return success(); 684681f929fSNicolas Vasilache } 685681f929fSNicolas Vasilache }; 686681f929fSNicolas Vasilache 6872d515e49SNicolas Vasilache // When ranks are different, InsertStridedSlice needs to extract a properly 6882d515e49SNicolas Vasilache // ranked vector from the destination vector into which to insert. This pattern 6892d515e49SNicolas Vasilache // only takes care of this part and forwards the rest of the conversion to 6902d515e49SNicolas Vasilache // another pattern that converts InsertStridedSlice for operands of the same 6912d515e49SNicolas Vasilache // rank. 6922d515e49SNicolas Vasilache // 6932d515e49SNicolas Vasilache // RewritePattern for InsertStridedSliceOp where source and destination vectors 6942d515e49SNicolas Vasilache // have different ranks. In this case: 6952d515e49SNicolas Vasilache // 1. the proper subvector is extracted from the destination vector 6962d515e49SNicolas Vasilache // 2. a new InsertStridedSlice op is created to insert the source in the 6972d515e49SNicolas Vasilache // destination subvector 6982d515e49SNicolas Vasilache // 3. the destination subvector is inserted back in the proper place 6992d515e49SNicolas Vasilache // 4. the op is replaced by the result of step 3. 7002d515e49SNicolas Vasilache // The new InsertStridedSlice from step 2. will be picked up by a 7012d515e49SNicolas Vasilache // `VectorInsertStridedSliceOpSameRankRewritePattern`. 7022d515e49SNicolas Vasilache class VectorInsertStridedSliceOpDifferentRankRewritePattern 7032d515e49SNicolas Vasilache : public OpRewritePattern<InsertStridedSliceOp> { 7042d515e49SNicolas Vasilache public: 7052d515e49SNicolas Vasilache using OpRewritePattern<InsertStridedSliceOp>::OpRewritePattern; 7062d515e49SNicolas Vasilache 7073145427dSRiver Riddle LogicalResult matchAndRewrite(InsertStridedSliceOp op, 7082d515e49SNicolas Vasilache PatternRewriter &rewriter) const override { 7092d515e49SNicolas Vasilache auto srcType = op.getSourceVectorType(); 7102d515e49SNicolas Vasilache auto dstType = op.getDestVectorType(); 7112d515e49SNicolas Vasilache 7122d515e49SNicolas Vasilache if (op.offsets().getValue().empty()) 7133145427dSRiver Riddle return failure(); 7142d515e49SNicolas Vasilache 7152d515e49SNicolas Vasilache auto loc = op.getLoc(); 7162d515e49SNicolas Vasilache int64_t rankDiff = dstType.getRank() - srcType.getRank(); 7172d515e49SNicolas Vasilache assert(rankDiff >= 0); 7182d515e49SNicolas Vasilache if (rankDiff == 0) 7193145427dSRiver Riddle return failure(); 7202d515e49SNicolas Vasilache 7212d515e49SNicolas Vasilache int64_t rankRest = dstType.getRank() - rankDiff; 7222d515e49SNicolas Vasilache // Extract / insert the subvector of matching rank and InsertStridedSlice 7232d515e49SNicolas Vasilache // on it. 7242d515e49SNicolas Vasilache Value extracted = 7252d515e49SNicolas Vasilache rewriter.create<ExtractOp>(loc, op.dest(), 7262d515e49SNicolas Vasilache getI64SubArray(op.offsets(), /*dropFront=*/0, 7272d515e49SNicolas Vasilache /*dropFront=*/rankRest)); 7282d515e49SNicolas Vasilache // A different pattern will kick in for InsertStridedSlice with matching 7292d515e49SNicolas Vasilache // ranks. 7302d515e49SNicolas Vasilache auto stridedSliceInnerOp = rewriter.create<InsertStridedSliceOp>( 7312d515e49SNicolas Vasilache loc, op.source(), extracted, 7322d515e49SNicolas Vasilache getI64SubArray(op.offsets(), /*dropFront=*/rankDiff), 733c8fc76a9Saartbik getI64SubArray(op.strides(), /*dropFront=*/0)); 7342d515e49SNicolas Vasilache rewriter.replaceOpWithNewOp<InsertOp>( 7352d515e49SNicolas Vasilache op, stridedSliceInnerOp.getResult(), op.dest(), 7362d515e49SNicolas Vasilache getI64SubArray(op.offsets(), /*dropFront=*/0, 7372d515e49SNicolas Vasilache /*dropFront=*/rankRest)); 7383145427dSRiver Riddle return success(); 7392d515e49SNicolas Vasilache } 7402d515e49SNicolas Vasilache }; 7412d515e49SNicolas Vasilache 7422d515e49SNicolas Vasilache // RewritePattern for InsertStridedSliceOp where source and destination vectors 7432d515e49SNicolas Vasilache // have the same rank. In this case, we reduce 7442d515e49SNicolas Vasilache // 1. the proper subvector is extracted from the destination vector 7452d515e49SNicolas Vasilache // 2. a new InsertStridedSlice op is created to insert the source in the 7462d515e49SNicolas Vasilache // destination subvector 7472d515e49SNicolas Vasilache // 3. the destination subvector is inserted back in the proper place 7482d515e49SNicolas Vasilache // 4. the op is replaced by the result of step 3. 7492d515e49SNicolas Vasilache // The new InsertStridedSlice from step 2. will be picked up by a 7502d515e49SNicolas Vasilache // `VectorInsertStridedSliceOpSameRankRewritePattern`. 7512d515e49SNicolas Vasilache class VectorInsertStridedSliceOpSameRankRewritePattern 7522d515e49SNicolas Vasilache : public OpRewritePattern<InsertStridedSliceOp> { 7532d515e49SNicolas Vasilache public: 7542d515e49SNicolas Vasilache using OpRewritePattern<InsertStridedSliceOp>::OpRewritePattern; 7552d515e49SNicolas Vasilache 7563145427dSRiver Riddle LogicalResult matchAndRewrite(InsertStridedSliceOp op, 7572d515e49SNicolas Vasilache PatternRewriter &rewriter) const override { 7582d515e49SNicolas Vasilache auto srcType = op.getSourceVectorType(); 7592d515e49SNicolas Vasilache auto dstType = op.getDestVectorType(); 7602d515e49SNicolas Vasilache 7612d515e49SNicolas Vasilache if (op.offsets().getValue().empty()) 7623145427dSRiver Riddle return failure(); 7632d515e49SNicolas Vasilache 7642d515e49SNicolas Vasilache int64_t rankDiff = dstType.getRank() - srcType.getRank(); 7652d515e49SNicolas Vasilache assert(rankDiff >= 0); 7662d515e49SNicolas Vasilache if (rankDiff != 0) 7673145427dSRiver Riddle return failure(); 7682d515e49SNicolas Vasilache 7692d515e49SNicolas Vasilache if (srcType == dstType) { 7702d515e49SNicolas Vasilache rewriter.replaceOp(op, op.source()); 7713145427dSRiver Riddle return success(); 7722d515e49SNicolas Vasilache } 7732d515e49SNicolas Vasilache 7742d515e49SNicolas Vasilache int64_t offset = 7752d515e49SNicolas Vasilache op.offsets().getValue().front().cast<IntegerAttr>().getInt(); 7762d515e49SNicolas Vasilache int64_t size = srcType.getShape().front(); 7772d515e49SNicolas Vasilache int64_t stride = 7782d515e49SNicolas Vasilache op.strides().getValue().front().cast<IntegerAttr>().getInt(); 7792d515e49SNicolas Vasilache 7802d515e49SNicolas Vasilache auto loc = op.getLoc(); 7812d515e49SNicolas Vasilache Value res = op.dest(); 7822d515e49SNicolas Vasilache // For each slice of the source vector along the most major dimension. 7832d515e49SNicolas Vasilache for (int64_t off = offset, e = offset + size * stride, idx = 0; off < e; 7842d515e49SNicolas Vasilache off += stride, ++idx) { 7852d515e49SNicolas Vasilache // 1. extract the proper subvector (or element) from source 7862d515e49SNicolas Vasilache Value extractedSource = extractOne(rewriter, loc, op.source(), idx); 7872d515e49SNicolas Vasilache if (extractedSource.getType().isa<VectorType>()) { 7882d515e49SNicolas Vasilache // 2. If we have a vector, extract the proper subvector from destination 7892d515e49SNicolas Vasilache // Otherwise we are at the element level and no need to recurse. 7902d515e49SNicolas Vasilache Value extractedDest = extractOne(rewriter, loc, op.dest(), off); 7912d515e49SNicolas Vasilache // 3. Reduce the problem to lowering a new InsertStridedSlice op with 7922d515e49SNicolas Vasilache // smaller rank. 793bd1ccfe6SRiver Riddle extractedSource = rewriter.create<InsertStridedSliceOp>( 7942d515e49SNicolas Vasilache loc, extractedSource, extractedDest, 7952d515e49SNicolas Vasilache getI64SubArray(op.offsets(), /* dropFront=*/1), 7962d515e49SNicolas Vasilache getI64SubArray(op.strides(), /* dropFront=*/1)); 7972d515e49SNicolas Vasilache } 7982d515e49SNicolas Vasilache // 4. Insert the extractedSource into the res vector. 7992d515e49SNicolas Vasilache res = insertOne(rewriter, loc, extractedSource, res, off); 8002d515e49SNicolas Vasilache } 8012d515e49SNicolas Vasilache 8022d515e49SNicolas Vasilache rewriter.replaceOp(op, res); 8033145427dSRiver Riddle return success(); 8042d515e49SNicolas Vasilache } 805bd1ccfe6SRiver Riddle /// This pattern creates recursive InsertStridedSliceOp, but the recursion is 806bd1ccfe6SRiver Riddle /// bounded as the rank is strictly decreasing. 807bd1ccfe6SRiver Riddle bool hasBoundedRewriteRecursion() const final { return true; } 8082d515e49SNicolas Vasilache }; 8092d515e49SNicolas Vasilache 810870c1fd4SAlex Zinenko class VectorTypeCastOpConversion : public ConvertToLLVMPattern { 8115c0c51a9SNicolas Vasilache public: 8125c0c51a9SNicolas Vasilache explicit VectorTypeCastOpConversion(MLIRContext *context, 8135c0c51a9SNicolas Vasilache LLVMTypeConverter &typeConverter) 814870c1fd4SAlex Zinenko : ConvertToLLVMPattern(vector::TypeCastOp::getOperationName(), context, 8155c0c51a9SNicolas Vasilache typeConverter) {} 8165c0c51a9SNicolas Vasilache 8173145427dSRiver Riddle LogicalResult 818e62a6956SRiver Riddle matchAndRewrite(Operation *op, ArrayRef<Value> operands, 8195c0c51a9SNicolas Vasilache ConversionPatternRewriter &rewriter) const override { 8205c0c51a9SNicolas Vasilache auto loc = op->getLoc(); 8215c0c51a9SNicolas Vasilache vector::TypeCastOp castOp = cast<vector::TypeCastOp>(op); 8225c0c51a9SNicolas Vasilache MemRefType sourceMemRefType = 8232bdf33ccSRiver Riddle castOp.getOperand().getType().cast<MemRefType>(); 8245c0c51a9SNicolas Vasilache MemRefType targetMemRefType = 8252bdf33ccSRiver Riddle castOp.getResult().getType().cast<MemRefType>(); 8265c0c51a9SNicolas Vasilache 8275c0c51a9SNicolas Vasilache // Only static shape casts supported atm. 8285c0c51a9SNicolas Vasilache if (!sourceMemRefType.hasStaticShape() || 8295c0c51a9SNicolas Vasilache !targetMemRefType.hasStaticShape()) 8303145427dSRiver Riddle return failure(); 8315c0c51a9SNicolas Vasilache 8325c0c51a9SNicolas Vasilache auto llvmSourceDescriptorTy = 8332bdf33ccSRiver Riddle operands[0].getType().dyn_cast<LLVM::LLVMType>(); 8345c0c51a9SNicolas Vasilache if (!llvmSourceDescriptorTy || !llvmSourceDescriptorTy.isStructTy()) 8353145427dSRiver Riddle return failure(); 8365c0c51a9SNicolas Vasilache MemRefDescriptor sourceMemRef(operands[0]); 8375c0c51a9SNicolas Vasilache 8380f04384dSAlex Zinenko auto llvmTargetDescriptorTy = typeConverter.convertType(targetMemRefType) 8395c0c51a9SNicolas Vasilache .dyn_cast_or_null<LLVM::LLVMType>(); 8405c0c51a9SNicolas Vasilache if (!llvmTargetDescriptorTy || !llvmTargetDescriptorTy.isStructTy()) 8413145427dSRiver Riddle return failure(); 8425c0c51a9SNicolas Vasilache 8435c0c51a9SNicolas Vasilache int64_t offset; 8445c0c51a9SNicolas Vasilache SmallVector<int64_t, 4> strides; 8455c0c51a9SNicolas Vasilache auto successStrides = 8465c0c51a9SNicolas Vasilache getStridesAndOffset(sourceMemRefType, strides, offset); 8475c0c51a9SNicolas Vasilache bool isContiguous = (strides.back() == 1); 8485c0c51a9SNicolas Vasilache if (isContiguous) { 8495c0c51a9SNicolas Vasilache auto sizes = sourceMemRefType.getShape(); 8505c0c51a9SNicolas Vasilache for (int index = 0, e = strides.size() - 2; index < e; ++index) { 8515c0c51a9SNicolas Vasilache if (strides[index] != strides[index + 1] * sizes[index + 1]) { 8525c0c51a9SNicolas Vasilache isContiguous = false; 8535c0c51a9SNicolas Vasilache break; 8545c0c51a9SNicolas Vasilache } 8555c0c51a9SNicolas Vasilache } 8565c0c51a9SNicolas Vasilache } 8575c0c51a9SNicolas Vasilache // Only contiguous source tensors supported atm. 8585c0c51a9SNicolas Vasilache if (failed(successStrides) || !isContiguous) 8593145427dSRiver Riddle return failure(); 8605c0c51a9SNicolas Vasilache 8610f04384dSAlex Zinenko auto int64Ty = LLVM::LLVMType::getInt64Ty(typeConverter.getDialect()); 8625c0c51a9SNicolas Vasilache 8635c0c51a9SNicolas Vasilache // Create descriptor. 8645c0c51a9SNicolas Vasilache auto desc = MemRefDescriptor::undef(rewriter, loc, llvmTargetDescriptorTy); 8655c0c51a9SNicolas Vasilache Type llvmTargetElementTy = desc.getElementType(); 8665c0c51a9SNicolas Vasilache // Set allocated ptr. 867e62a6956SRiver Riddle Value allocated = sourceMemRef.allocatedPtr(rewriter, loc); 8685c0c51a9SNicolas Vasilache allocated = 8695c0c51a9SNicolas Vasilache rewriter.create<LLVM::BitcastOp>(loc, llvmTargetElementTy, allocated); 8705c0c51a9SNicolas Vasilache desc.setAllocatedPtr(rewriter, loc, allocated); 8715c0c51a9SNicolas Vasilache // Set aligned ptr. 872e62a6956SRiver Riddle Value ptr = sourceMemRef.alignedPtr(rewriter, loc); 8735c0c51a9SNicolas Vasilache ptr = rewriter.create<LLVM::BitcastOp>(loc, llvmTargetElementTy, ptr); 8745c0c51a9SNicolas Vasilache desc.setAlignedPtr(rewriter, loc, ptr); 8755c0c51a9SNicolas Vasilache // Fill offset 0. 8765c0c51a9SNicolas Vasilache auto attr = rewriter.getIntegerAttr(rewriter.getIndexType(), 0); 8775c0c51a9SNicolas Vasilache auto zero = rewriter.create<LLVM::ConstantOp>(loc, int64Ty, attr); 8785c0c51a9SNicolas Vasilache desc.setOffset(rewriter, loc, zero); 8795c0c51a9SNicolas Vasilache 8805c0c51a9SNicolas Vasilache // Fill size and stride descriptors in memref. 8815c0c51a9SNicolas Vasilache for (auto indexedSize : llvm::enumerate(targetMemRefType.getShape())) { 8825c0c51a9SNicolas Vasilache int64_t index = indexedSize.index(); 8835c0c51a9SNicolas Vasilache auto sizeAttr = 8845c0c51a9SNicolas Vasilache rewriter.getIntegerAttr(rewriter.getIndexType(), indexedSize.value()); 8855c0c51a9SNicolas Vasilache auto size = rewriter.create<LLVM::ConstantOp>(loc, int64Ty, sizeAttr); 8865c0c51a9SNicolas Vasilache desc.setSize(rewriter, loc, index, size); 8875c0c51a9SNicolas Vasilache auto strideAttr = 8885c0c51a9SNicolas Vasilache rewriter.getIntegerAttr(rewriter.getIndexType(), strides[index]); 8895c0c51a9SNicolas Vasilache auto stride = rewriter.create<LLVM::ConstantOp>(loc, int64Ty, strideAttr); 8905c0c51a9SNicolas Vasilache desc.setStride(rewriter, loc, index, stride); 8915c0c51a9SNicolas Vasilache } 8925c0c51a9SNicolas Vasilache 8935c0c51a9SNicolas Vasilache rewriter.replaceOp(op, {desc}); 8943145427dSRiver Riddle return success(); 8955c0c51a9SNicolas Vasilache } 8965c0c51a9SNicolas Vasilache }; 8975c0c51a9SNicolas Vasilache 8988345b86dSNicolas Vasilache template <typename ConcreteOp> 8998345b86dSNicolas Vasilache void replaceTransferOp(ConversionPatternRewriter &rewriter, 9008345b86dSNicolas Vasilache LLVMTypeConverter &typeConverter, Location loc, 9018345b86dSNicolas Vasilache Operation *op, ArrayRef<Value> operands, Value dataPtr, 9028345b86dSNicolas Vasilache Value mask); 9038345b86dSNicolas Vasilache 9048345b86dSNicolas Vasilache template <> 9058345b86dSNicolas Vasilache void replaceTransferOp<TransferReadOp>(ConversionPatternRewriter &rewriter, 9068345b86dSNicolas Vasilache LLVMTypeConverter &typeConverter, 9078345b86dSNicolas Vasilache Location loc, Operation *op, 9088345b86dSNicolas Vasilache ArrayRef<Value> operands, Value dataPtr, 9098345b86dSNicolas Vasilache Value mask) { 9108345b86dSNicolas Vasilache auto xferOp = cast<TransferReadOp>(op); 9118345b86dSNicolas Vasilache auto toLLVMTy = [&](Type t) { return typeConverter.convertType(t); }; 9128345b86dSNicolas Vasilache VectorType fillType = xferOp.getVectorType(); 9138345b86dSNicolas Vasilache Value fill = rewriter.create<SplatOp>(loc, fillType, xferOp.padding()); 9148345b86dSNicolas Vasilache fill = rewriter.create<LLVM::DialectCastOp>(loc, toLLVMTy(fillType), fill); 9158345b86dSNicolas Vasilache 9168345b86dSNicolas Vasilache auto vecTy = toLLVMTy(xferOp.getVectorType()).template cast<LLVM::LLVMType>(); 9178345b86dSNicolas Vasilache rewriter.replaceOpWithNewOp<LLVM::MaskedLoadOp>( 9188345b86dSNicolas Vasilache op, vecTy, dataPtr, mask, ValueRange{fill}, 9198345b86dSNicolas Vasilache rewriter.getI32IntegerAttr(1)); 9208345b86dSNicolas Vasilache } 9218345b86dSNicolas Vasilache 9228345b86dSNicolas Vasilache template <> 9238345b86dSNicolas Vasilache void replaceTransferOp<TransferWriteOp>(ConversionPatternRewriter &rewriter, 9248345b86dSNicolas Vasilache LLVMTypeConverter &typeConverter, 9258345b86dSNicolas Vasilache Location loc, Operation *op, 9268345b86dSNicolas Vasilache ArrayRef<Value> operands, Value dataPtr, 9278345b86dSNicolas Vasilache Value mask) { 9288345b86dSNicolas Vasilache auto adaptor = TransferWriteOpOperandAdaptor(operands); 9298345b86dSNicolas Vasilache rewriter.replaceOpWithNewOp<LLVM::MaskedStoreOp>( 9308345b86dSNicolas Vasilache op, adaptor.vector(), dataPtr, mask, rewriter.getI32IntegerAttr(1)); 9318345b86dSNicolas Vasilache } 9328345b86dSNicolas Vasilache 9338345b86dSNicolas Vasilache static TransferReadOpOperandAdaptor 9348345b86dSNicolas Vasilache getTransferOpAdapter(TransferReadOp xferOp, ArrayRef<Value> operands) { 9358345b86dSNicolas Vasilache return TransferReadOpOperandAdaptor(operands); 9368345b86dSNicolas Vasilache } 9378345b86dSNicolas Vasilache 9388345b86dSNicolas Vasilache static TransferWriteOpOperandAdaptor 9398345b86dSNicolas Vasilache getTransferOpAdapter(TransferWriteOp xferOp, ArrayRef<Value> operands) { 9408345b86dSNicolas Vasilache return TransferWriteOpOperandAdaptor(operands); 9418345b86dSNicolas Vasilache } 9428345b86dSNicolas Vasilache 9438345b86dSNicolas Vasilache /// Conversion pattern that converts a 1-D vector transfer read/write op in a 9448345b86dSNicolas Vasilache /// sequence of: 9458345b86dSNicolas Vasilache /// 1. Bitcast to vector form. 9468345b86dSNicolas Vasilache /// 2. Create an offsetVector = [ offset + 0 .. offset + vector_length - 1 ]. 9478345b86dSNicolas Vasilache /// 3. Create a mask where offsetVector is compared against memref upper bound. 9488345b86dSNicolas Vasilache /// 4. Rewrite op as a masked read or write. 9498345b86dSNicolas Vasilache template <typename ConcreteOp> 9508345b86dSNicolas Vasilache class VectorTransferConversion : public ConvertToLLVMPattern { 9518345b86dSNicolas Vasilache public: 9528345b86dSNicolas Vasilache explicit VectorTransferConversion(MLIRContext *context, 9538345b86dSNicolas Vasilache LLVMTypeConverter &typeConv) 9548345b86dSNicolas Vasilache : ConvertToLLVMPattern(ConcreteOp::getOperationName(), context, 9558345b86dSNicolas Vasilache typeConv) {} 9568345b86dSNicolas Vasilache 9578345b86dSNicolas Vasilache LogicalResult 9588345b86dSNicolas Vasilache matchAndRewrite(Operation *op, ArrayRef<Value> operands, 9598345b86dSNicolas Vasilache ConversionPatternRewriter &rewriter) const override { 9608345b86dSNicolas Vasilache auto xferOp = cast<ConcreteOp>(op); 9618345b86dSNicolas Vasilache auto adaptor = getTransferOpAdapter(xferOp, operands); 9628345b86dSNicolas Vasilache if (xferOp.getMemRefType().getRank() != 1) 9638345b86dSNicolas Vasilache return failure(); 9648345b86dSNicolas Vasilache if (!xferOp.permutation_map().isIdentity()) 9658345b86dSNicolas Vasilache return failure(); 9668345b86dSNicolas Vasilache 9678345b86dSNicolas Vasilache auto toLLVMTy = [&](Type t) { return typeConverter.convertType(t); }; 9688345b86dSNicolas Vasilache 9698345b86dSNicolas Vasilache Location loc = op->getLoc(); 9708345b86dSNicolas Vasilache Type i64Type = rewriter.getIntegerType(64); 9718345b86dSNicolas Vasilache MemRefType memRefType = xferOp.getMemRefType(); 9728345b86dSNicolas Vasilache 9738345b86dSNicolas Vasilache // 1. Get the source/dst address as an LLVM vector pointer. 9748345b86dSNicolas Vasilache // TODO: support alignment when possible. 9758345b86dSNicolas Vasilache Value dataPtr = getDataPtr(loc, memRefType, adaptor.memref(), 9768345b86dSNicolas Vasilache adaptor.indices(), rewriter, getModule()); 9778345b86dSNicolas Vasilache auto vecTy = 9788345b86dSNicolas Vasilache toLLVMTy(xferOp.getVectorType()).template cast<LLVM::LLVMType>(); 9798345b86dSNicolas Vasilache auto vectorDataPtr = 9808345b86dSNicolas Vasilache rewriter.create<LLVM::BitcastOp>(loc, vecTy.getPointerTo(), dataPtr); 9818345b86dSNicolas Vasilache 9828345b86dSNicolas Vasilache // 2. Create a vector with linear indices [ 0 .. vector_length - 1 ]. 9838345b86dSNicolas Vasilache unsigned vecWidth = vecTy.getVectorNumElements(); 9848345b86dSNicolas Vasilache VectorType vectorCmpType = VectorType::get(vecWidth, i64Type); 9858345b86dSNicolas Vasilache SmallVector<int64_t, 8> indices; 9868345b86dSNicolas Vasilache indices.reserve(vecWidth); 9878345b86dSNicolas Vasilache for (unsigned i = 0; i < vecWidth; ++i) 9888345b86dSNicolas Vasilache indices.push_back(i); 9898345b86dSNicolas Vasilache Value linearIndices = rewriter.create<ConstantOp>( 9908345b86dSNicolas Vasilache loc, vectorCmpType, 9918345b86dSNicolas Vasilache DenseElementsAttr::get(vectorCmpType, ArrayRef<int64_t>(indices))); 9928345b86dSNicolas Vasilache linearIndices = rewriter.create<LLVM::DialectCastOp>( 9938345b86dSNicolas Vasilache loc, toLLVMTy(vectorCmpType), linearIndices); 9948345b86dSNicolas Vasilache 9958345b86dSNicolas Vasilache // 3. Create offsetVector = [ offset + 0 .. offset + vector_length - 1 ]. 9968345b86dSNicolas Vasilache Value offsetIndex = *(xferOp.indices().begin()); 9978345b86dSNicolas Vasilache offsetIndex = rewriter.create<IndexCastOp>( 9988345b86dSNicolas Vasilache loc, vectorCmpType.getElementType(), offsetIndex); 9998345b86dSNicolas Vasilache Value base = rewriter.create<SplatOp>(loc, vectorCmpType, offsetIndex); 10008345b86dSNicolas Vasilache Value offsetVector = rewriter.create<AddIOp>(loc, base, linearIndices); 10018345b86dSNicolas Vasilache 10028345b86dSNicolas Vasilache // 4. Let dim the memref dimension, compute the vector comparison mask: 10038345b86dSNicolas Vasilache // [ offset + 0 .. offset + vector_length - 1 ] < [ dim .. dim ] 10048345b86dSNicolas Vasilache Value dim = rewriter.create<DimOp>(loc, xferOp.memref(), 0); 10058345b86dSNicolas Vasilache dim = 10068345b86dSNicolas Vasilache rewriter.create<IndexCastOp>(loc, vectorCmpType.getElementType(), dim); 10078345b86dSNicolas Vasilache dim = rewriter.create<SplatOp>(loc, vectorCmpType, dim); 10088345b86dSNicolas Vasilache Value mask = 10098345b86dSNicolas Vasilache rewriter.create<CmpIOp>(loc, CmpIPredicate::slt, offsetVector, dim); 10108345b86dSNicolas Vasilache mask = rewriter.create<LLVM::DialectCastOp>(loc, toLLVMTy(mask.getType()), 10118345b86dSNicolas Vasilache mask); 10128345b86dSNicolas Vasilache 10138345b86dSNicolas Vasilache // 5. Rewrite as a masked read / write. 10148345b86dSNicolas Vasilache replaceTransferOp<ConcreteOp>(rewriter, typeConverter, loc, op, operands, 10158345b86dSNicolas Vasilache vectorDataPtr, mask); 10168345b86dSNicolas Vasilache 10178345b86dSNicolas Vasilache return success(); 10188345b86dSNicolas Vasilache } 10198345b86dSNicolas Vasilache }; 10208345b86dSNicolas Vasilache 1021870c1fd4SAlex Zinenko class VectorPrintOpConversion : public ConvertToLLVMPattern { 1022d9b500d3SAart Bik public: 1023d9b500d3SAart Bik explicit VectorPrintOpConversion(MLIRContext *context, 1024d9b500d3SAart Bik LLVMTypeConverter &typeConverter) 1025870c1fd4SAlex Zinenko : ConvertToLLVMPattern(vector::PrintOp::getOperationName(), context, 1026d9b500d3SAart Bik typeConverter) {} 1027d9b500d3SAart Bik 1028d9b500d3SAart Bik // Proof-of-concept lowering implementation that relies on a small 1029d9b500d3SAart Bik // runtime support library, which only needs to provide a few 1030d9b500d3SAart Bik // printing methods (single value for all data types, opening/closing 1031d9b500d3SAart Bik // bracket, comma, newline). The lowering fully unrolls a vector 1032d9b500d3SAart Bik // in terms of these elementary printing operations. The advantage 1033d9b500d3SAart Bik // of this approach is that the library can remain unaware of all 1034d9b500d3SAart Bik // low-level implementation details of vectors while still supporting 1035d9b500d3SAart Bik // output of any shaped and dimensioned vector. Due to full unrolling, 1036d9b500d3SAart Bik // this approach is less suited for very large vectors though. 1037d9b500d3SAart Bik // 1038d9b500d3SAart Bik // TODO(ajcbik): rely solely on libc in future? something else? 1039d9b500d3SAart Bik // 10403145427dSRiver Riddle LogicalResult 1041e62a6956SRiver Riddle matchAndRewrite(Operation *op, ArrayRef<Value> operands, 1042d9b500d3SAart Bik ConversionPatternRewriter &rewriter) const override { 1043d9b500d3SAart Bik auto printOp = cast<vector::PrintOp>(op); 1044d9b500d3SAart Bik auto adaptor = vector::PrintOpOperandAdaptor(operands); 1045d9b500d3SAart Bik Type printType = printOp.getPrintType(); 1046d9b500d3SAart Bik 10470f04384dSAlex Zinenko if (typeConverter.convertType(printType) == nullptr) 10483145427dSRiver Riddle return failure(); 1049d9b500d3SAart Bik 1050d9b500d3SAart Bik // Make sure element type has runtime support (currently just Float/Double). 1051d9b500d3SAart Bik VectorType vectorType = printType.dyn_cast<VectorType>(); 1052d9b500d3SAart Bik Type eltType = vectorType ? vectorType.getElementType() : printType; 1053d9b500d3SAart Bik int64_t rank = vectorType ? vectorType.getRank() : 0; 1054d9b500d3SAart Bik Operation *printer; 105535b68527SLei Zhang if (eltType.isSignlessInteger(32)) 1056e52414b1Saartbik printer = getPrintI32(op); 105735b68527SLei Zhang else if (eltType.isSignlessInteger(64)) 1058e52414b1Saartbik printer = getPrintI64(op); 1059e52414b1Saartbik else if (eltType.isF32()) 1060d9b500d3SAart Bik printer = getPrintFloat(op); 1061d9b500d3SAart Bik else if (eltType.isF64()) 1062d9b500d3SAart Bik printer = getPrintDouble(op); 1063d9b500d3SAart Bik else 10643145427dSRiver Riddle return failure(); 1065d9b500d3SAart Bik 1066d9b500d3SAart Bik // Unroll vector into elementary print calls. 1067d9b500d3SAart Bik emitRanks(rewriter, op, adaptor.source(), vectorType, printer, rank); 1068d9b500d3SAart Bik emitCall(rewriter, op->getLoc(), getPrintNewline(op)); 1069d9b500d3SAart Bik rewriter.eraseOp(op); 10703145427dSRiver Riddle return success(); 1071d9b500d3SAart Bik } 1072d9b500d3SAart Bik 1073d9b500d3SAart Bik private: 1074d9b500d3SAart Bik void emitRanks(ConversionPatternRewriter &rewriter, Operation *op, 1075e62a6956SRiver Riddle Value value, VectorType vectorType, Operation *printer, 1076d9b500d3SAart Bik int64_t rank) const { 1077d9b500d3SAart Bik Location loc = op->getLoc(); 1078d9b500d3SAart Bik if (rank == 0) { 1079d9b500d3SAart Bik emitCall(rewriter, loc, printer, value); 1080d9b500d3SAart Bik return; 1081d9b500d3SAart Bik } 1082d9b500d3SAart Bik 1083d9b500d3SAart Bik emitCall(rewriter, loc, getPrintOpen(op)); 1084d9b500d3SAart Bik Operation *printComma = getPrintComma(op); 1085d9b500d3SAart Bik int64_t dim = vectorType.getDimSize(0); 1086d9b500d3SAart Bik for (int64_t d = 0; d < dim; ++d) { 1087d9b500d3SAart Bik auto reducedType = 1088d9b500d3SAart Bik rank > 1 ? reducedVectorTypeFront(vectorType) : nullptr; 10890f04384dSAlex Zinenko auto llvmType = typeConverter.convertType( 1090d9b500d3SAart Bik rank > 1 ? reducedType : vectorType.getElementType()); 1091e62a6956SRiver Riddle Value nestedVal = 10920f04384dSAlex Zinenko extractOne(rewriter, typeConverter, loc, value, llvmType, rank, d); 1093d9b500d3SAart Bik emitRanks(rewriter, op, nestedVal, reducedType, printer, rank - 1); 1094d9b500d3SAart Bik if (d != dim - 1) 1095d9b500d3SAart Bik emitCall(rewriter, loc, printComma); 1096d9b500d3SAart Bik } 1097d9b500d3SAart Bik emitCall(rewriter, loc, getPrintClose(op)); 1098d9b500d3SAart Bik } 1099d9b500d3SAart Bik 1100d9b500d3SAart Bik // Helper to emit a call. 1101d9b500d3SAart Bik static void emitCall(ConversionPatternRewriter &rewriter, Location loc, 1102d9b500d3SAart Bik Operation *ref, ValueRange params = ValueRange()) { 1103d9b500d3SAart Bik rewriter.create<LLVM::CallOp>(loc, ArrayRef<Type>{}, 1104d9b500d3SAart Bik rewriter.getSymbolRefAttr(ref), params); 1105d9b500d3SAart Bik } 1106d9b500d3SAart Bik 1107d9b500d3SAart Bik // Helper for printer method declaration (first hit) and lookup. 1108d9b500d3SAart Bik static Operation *getPrint(Operation *op, LLVM::LLVMDialect *dialect, 1109d9b500d3SAart Bik StringRef name, ArrayRef<LLVM::LLVMType> params) { 1110d9b500d3SAart Bik auto module = op->getParentOfType<ModuleOp>(); 1111d9b500d3SAart Bik auto func = module.lookupSymbol<LLVM::LLVMFuncOp>(name); 1112d9b500d3SAart Bik if (func) 1113d9b500d3SAart Bik return func; 1114d9b500d3SAart Bik OpBuilder moduleBuilder(module.getBodyRegion()); 1115d9b500d3SAart Bik return moduleBuilder.create<LLVM::LLVMFuncOp>( 1116d9b500d3SAart Bik op->getLoc(), name, 1117d9b500d3SAart Bik LLVM::LLVMType::getFunctionTy(LLVM::LLVMType::getVoidTy(dialect), 1118d9b500d3SAart Bik params, /*isVarArg=*/false)); 1119d9b500d3SAart Bik } 1120d9b500d3SAart Bik 1121d9b500d3SAart Bik // Helpers for method names. 1122e52414b1Saartbik Operation *getPrintI32(Operation *op) const { 11230f04384dSAlex Zinenko LLVM::LLVMDialect *dialect = typeConverter.getDialect(); 1124e52414b1Saartbik return getPrint(op, dialect, "print_i32", 1125e52414b1Saartbik LLVM::LLVMType::getInt32Ty(dialect)); 1126e52414b1Saartbik } 1127e52414b1Saartbik Operation *getPrintI64(Operation *op) const { 11280f04384dSAlex Zinenko LLVM::LLVMDialect *dialect = typeConverter.getDialect(); 1129e52414b1Saartbik return getPrint(op, dialect, "print_i64", 1130e52414b1Saartbik LLVM::LLVMType::getInt64Ty(dialect)); 1131e52414b1Saartbik } 1132d9b500d3SAart Bik Operation *getPrintFloat(Operation *op) const { 11330f04384dSAlex Zinenko LLVM::LLVMDialect *dialect = typeConverter.getDialect(); 1134d9b500d3SAart Bik return getPrint(op, dialect, "print_f32", 1135d9b500d3SAart Bik LLVM::LLVMType::getFloatTy(dialect)); 1136d9b500d3SAart Bik } 1137d9b500d3SAart Bik Operation *getPrintDouble(Operation *op) const { 11380f04384dSAlex Zinenko LLVM::LLVMDialect *dialect = typeConverter.getDialect(); 1139d9b500d3SAart Bik return getPrint(op, dialect, "print_f64", 1140d9b500d3SAart Bik LLVM::LLVMType::getDoubleTy(dialect)); 1141d9b500d3SAart Bik } 1142d9b500d3SAart Bik Operation *getPrintOpen(Operation *op) const { 11430f04384dSAlex Zinenko return getPrint(op, typeConverter.getDialect(), "print_open", {}); 1144d9b500d3SAart Bik } 1145d9b500d3SAart Bik Operation *getPrintClose(Operation *op) const { 11460f04384dSAlex Zinenko return getPrint(op, typeConverter.getDialect(), "print_close", {}); 1147d9b500d3SAart Bik } 1148d9b500d3SAart Bik Operation *getPrintComma(Operation *op) const { 11490f04384dSAlex Zinenko return getPrint(op, typeConverter.getDialect(), "print_comma", {}); 1150d9b500d3SAart Bik } 1151d9b500d3SAart Bik Operation *getPrintNewline(Operation *op) const { 11520f04384dSAlex Zinenko return getPrint(op, typeConverter.getDialect(), "print_newline", {}); 1153d9b500d3SAart Bik } 1154d9b500d3SAart Bik }; 1155d9b500d3SAart Bik 115665678d93SNicolas Vasilache /// Progressive lowering of StridedSliceOp to either: 115765678d93SNicolas Vasilache /// 1. extractelement + insertelement for the 1-D case 115865678d93SNicolas Vasilache /// 2. extract + optional strided_slice + insert for the n-D case. 11592d515e49SNicolas Vasilache class VectorStridedSliceOpConversion : public OpRewritePattern<StridedSliceOp> { 116065678d93SNicolas Vasilache public: 116165678d93SNicolas Vasilache using OpRewritePattern<StridedSliceOp>::OpRewritePattern; 116265678d93SNicolas Vasilache 11633145427dSRiver Riddle LogicalResult matchAndRewrite(StridedSliceOp op, 116465678d93SNicolas Vasilache PatternRewriter &rewriter) const override { 116565678d93SNicolas Vasilache auto dstType = op.getResult().getType().cast<VectorType>(); 116665678d93SNicolas Vasilache 116765678d93SNicolas Vasilache assert(!op.offsets().getValue().empty() && "Unexpected empty offsets"); 116865678d93SNicolas Vasilache 116965678d93SNicolas Vasilache int64_t offset = 117065678d93SNicolas Vasilache op.offsets().getValue().front().cast<IntegerAttr>().getInt(); 117165678d93SNicolas Vasilache int64_t size = op.sizes().getValue().front().cast<IntegerAttr>().getInt(); 117265678d93SNicolas Vasilache int64_t stride = 117365678d93SNicolas Vasilache op.strides().getValue().front().cast<IntegerAttr>().getInt(); 117465678d93SNicolas Vasilache 117565678d93SNicolas Vasilache auto loc = op.getLoc(); 117665678d93SNicolas Vasilache auto elemType = dstType.getElementType(); 117735b68527SLei Zhang assert(elemType.isSignlessIntOrIndexOrFloat()); 117865678d93SNicolas Vasilache Value zero = rewriter.create<ConstantOp>(loc, elemType, 117965678d93SNicolas Vasilache rewriter.getZeroAttr(elemType)); 118065678d93SNicolas Vasilache Value res = rewriter.create<SplatOp>(loc, dstType, zero); 118165678d93SNicolas Vasilache for (int64_t off = offset, e = offset + size * stride, idx = 0; off < e; 118265678d93SNicolas Vasilache off += stride, ++idx) { 118365678d93SNicolas Vasilache Value extracted = extractOne(rewriter, loc, op.vector(), off); 118465678d93SNicolas Vasilache if (op.offsets().getValue().size() > 1) { 1185bd1ccfe6SRiver Riddle extracted = rewriter.create<StridedSliceOp>( 118665678d93SNicolas Vasilache loc, extracted, getI64SubArray(op.offsets(), /* dropFront=*/1), 118765678d93SNicolas Vasilache getI64SubArray(op.sizes(), /* dropFront=*/1), 118865678d93SNicolas Vasilache getI64SubArray(op.strides(), /* dropFront=*/1)); 118965678d93SNicolas Vasilache } 119065678d93SNicolas Vasilache res = insertOne(rewriter, loc, extracted, res, idx); 119165678d93SNicolas Vasilache } 119265678d93SNicolas Vasilache rewriter.replaceOp(op, {res}); 11933145427dSRiver Riddle return success(); 119465678d93SNicolas Vasilache } 1195bd1ccfe6SRiver Riddle /// This pattern creates recursive StridedSliceOp, but the recursion is 1196bd1ccfe6SRiver Riddle /// bounded as the rank is strictly decreasing. 1197bd1ccfe6SRiver Riddle bool hasBoundedRewriteRecursion() const final { return true; } 119865678d93SNicolas Vasilache }; 119965678d93SNicolas Vasilache 1200df186507SBenjamin Kramer } // namespace 1201df186507SBenjamin Kramer 12025c0c51a9SNicolas Vasilache /// Populate the given list with patterns that convert from Vector to LLVM. 12035c0c51a9SNicolas Vasilache void mlir::populateVectorToLLVMConversionPatterns( 12045c0c51a9SNicolas Vasilache LLVMTypeConverter &converter, OwningRewritePatternList &patterns) { 120565678d93SNicolas Vasilache MLIRContext *ctx = converter.getDialect()->getContext(); 12068345b86dSNicolas Vasilache // clang-format off 1207681f929fSNicolas Vasilache patterns.insert<VectorFMAOpNDRewritePattern, 1208681f929fSNicolas Vasilache VectorInsertStridedSliceOpDifferentRankRewritePattern, 12092d515e49SNicolas Vasilache VectorInsertStridedSliceOpSameRankRewritePattern, 12102d515e49SNicolas Vasilache VectorStridedSliceOpConversion>(ctx); 12118345b86dSNicolas Vasilache patterns 12128345b86dSNicolas Vasilache .insert<VectorBroadcastOpConversion, 12138345b86dSNicolas Vasilache VectorReductionOpConversion, 12148345b86dSNicolas Vasilache VectorShuffleOpConversion, 12158345b86dSNicolas Vasilache VectorExtractElementOpConversion, 12168345b86dSNicolas Vasilache VectorExtractOpConversion, 12178345b86dSNicolas Vasilache VectorFMAOp1DConversion, 12188345b86dSNicolas Vasilache VectorInsertElementOpConversion, 12198345b86dSNicolas Vasilache VectorInsertOpConversion, 12208345b86dSNicolas Vasilache VectorPrintOpConversion, 12218345b86dSNicolas Vasilache VectorTransferConversion<TransferReadOp>, 12228345b86dSNicolas Vasilache VectorTransferConversion<TransferWriteOp>, 12238345b86dSNicolas Vasilache VectorTypeCastOpConversion>(ctx, converter); 12248345b86dSNicolas Vasilache // clang-format on 12255c0c51a9SNicolas Vasilache } 12265c0c51a9SNicolas Vasilache 122763b683a8SNicolas Vasilache void mlir::populateVectorToLLVMMatrixConversionPatterns( 122863b683a8SNicolas Vasilache LLVMTypeConverter &converter, OwningRewritePatternList &patterns) { 122963b683a8SNicolas Vasilache MLIRContext *ctx = converter.getDialect()->getContext(); 123063b683a8SNicolas Vasilache patterns.insert<VectorMatmulOpConversion>(ctx, converter); 123163b683a8SNicolas Vasilache } 123263b683a8SNicolas Vasilache 12335c0c51a9SNicolas Vasilache namespace { 1234722f909fSRiver Riddle struct LowerVectorToLLVMPass 12351834ad4aSRiver Riddle : public ConvertVectorToLLVMBase<LowerVectorToLLVMPass> { 1236722f909fSRiver Riddle void runOnOperation() override; 12375c0c51a9SNicolas Vasilache }; 12385c0c51a9SNicolas Vasilache } // namespace 12395c0c51a9SNicolas Vasilache 1240722f909fSRiver Riddle void LowerVectorToLLVMPass::runOnOperation() { 1241078776a6Saartbik // Perform progressive lowering of operations on slices and 1242b21c7999Saartbik // all contraction operations. Also applies folding and DCE. 1243459cf6e5Saartbik { 12445c0c51a9SNicolas Vasilache OwningRewritePatternList patterns; 1245459cf6e5Saartbik populateVectorSlicesLoweringPatterns(patterns, &getContext()); 1246b21c7999Saartbik populateVectorContractLoweringPatterns(patterns, &getContext()); 1247*a5b9316bSUday Bondhugula applyPatternsAndFoldGreedily(getOperation(), patterns); 1248459cf6e5Saartbik } 1249459cf6e5Saartbik 1250459cf6e5Saartbik // Convert to the LLVM IR dialect. 12515c0c51a9SNicolas Vasilache LLVMTypeConverter converter(&getContext()); 1252459cf6e5Saartbik OwningRewritePatternList patterns; 125363b683a8SNicolas Vasilache populateVectorToLLVMMatrixConversionPatterns(converter, patterns); 12545c0c51a9SNicolas Vasilache populateVectorToLLVMConversionPatterns(converter, patterns); 1255bbf3ef85SNicolas Vasilache populateVectorToLLVMMatrixConversionPatterns(converter, patterns); 12565c0c51a9SNicolas Vasilache populateStdToLLVMConversionPatterns(converter, patterns); 12575c0c51a9SNicolas Vasilache 12582a00ae39STim Shen LLVMConversionTarget target(getContext()); 12595c0c51a9SNicolas Vasilache target.addDynamicallyLegalOp<FuncOp>( 12605c0c51a9SNicolas Vasilache [&](FuncOp op) { return converter.isSignatureLegal(op.getType()); }); 1261722f909fSRiver Riddle if (failed(applyPartialConversion(getOperation(), target, patterns, 1262722f909fSRiver Riddle &converter))) { 12635c0c51a9SNicolas Vasilache signalPassFailure(); 12645c0c51a9SNicolas Vasilache } 12655c0c51a9SNicolas Vasilache } 12665c0c51a9SNicolas Vasilache 126780aca1eaSRiver Riddle std::unique_ptr<OperationPass<ModuleOp>> mlir::createConvertVectorToLLVMPass() { 12682fae7878SNicolas Vasilache return std::make_unique<LowerVectorToLLVMPass>(); 12695c0c51a9SNicolas Vasilache } 1270