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" 105c0c51a9SNicolas Vasilache #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h" 115c0c51a9SNicolas Vasilache #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h" 125c0c51a9SNicolas Vasilache #include "mlir/Dialect/LLVMIR/LLVMDialect.h" 1365678d93SNicolas Vasilache #include "mlir/Dialect/StandardOps/Ops.h" 145c0c51a9SNicolas Vasilache #include "mlir/Dialect/VectorOps/VectorOps.h" 155c0c51a9SNicolas Vasilache #include "mlir/IR/Attributes.h" 165c0c51a9SNicolas Vasilache #include "mlir/IR/Builders.h" 175c0c51a9SNicolas Vasilache #include "mlir/IR/MLIRContext.h" 185c0c51a9SNicolas Vasilache #include "mlir/IR/Module.h" 195c0c51a9SNicolas Vasilache #include "mlir/IR/Operation.h" 205c0c51a9SNicolas Vasilache #include "mlir/IR/PatternMatch.h" 215c0c51a9SNicolas Vasilache #include "mlir/IR/StandardTypes.h" 225c0c51a9SNicolas Vasilache #include "mlir/IR/Types.h" 235c0c51a9SNicolas Vasilache #include "mlir/Pass/Pass.h" 245c0c51a9SNicolas Vasilache #include "mlir/Pass/PassManager.h" 255c0c51a9SNicolas Vasilache #include "mlir/Transforms/DialectConversion.h" 265c0c51a9SNicolas Vasilache #include "mlir/Transforms/Passes.h" 275c0c51a9SNicolas Vasilache 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, 395c0c51a9SNicolas Vasilache LLVMTypeConverter &lowering) { 405c0c51a9SNicolas Vasilache return lowering.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, 59e62a6956SRiver Riddle LLVMTypeConverter &lowering, Location loc, Value val1, 60e62a6956SRiver Riddle Value val2, Type llvmType, int64_t rank, int64_t pos) { 611c81adf3SAart Bik if (rank == 1) { 621c81adf3SAart Bik auto idxType = rewriter.getIndexType(); 631c81adf3SAart Bik auto constant = rewriter.create<LLVM::ConstantOp>( 641c81adf3SAart Bik loc, lowering.convertType(idxType), 651c81adf3SAart Bik rewriter.getIntegerAttr(idxType, pos)); 661c81adf3SAart Bik return rewriter.create<LLVM::InsertElementOp>(loc, llvmType, val1, val2, 671c81adf3SAart Bik constant); 681c81adf3SAart Bik } 691c81adf3SAart Bik return rewriter.create<LLVM::InsertValueOp>(loc, llvmType, val1, val2, 701c81adf3SAart Bik rewriter.getI64ArrayAttr(pos)); 711c81adf3SAart Bik } 721c81adf3SAart Bik 732d515e49SNicolas Vasilache // Helper that picks the proper sequence for inserting. 742d515e49SNicolas Vasilache static Value insertOne(PatternRewriter &rewriter, Location loc, Value from, 752d515e49SNicolas Vasilache Value into, int64_t offset) { 762d515e49SNicolas Vasilache auto vectorType = into.getType().cast<VectorType>(); 772d515e49SNicolas Vasilache if (vectorType.getRank() > 1) 782d515e49SNicolas Vasilache return rewriter.create<InsertOp>(loc, from, into, offset); 792d515e49SNicolas Vasilache return rewriter.create<vector::InsertElementOp>( 802d515e49SNicolas Vasilache loc, vectorType, from, into, 812d515e49SNicolas Vasilache rewriter.create<ConstantIndexOp>(loc, offset)); 822d515e49SNicolas Vasilache } 832d515e49SNicolas Vasilache 841c81adf3SAart Bik // Helper that picks the proper sequence for extracting. 85e62a6956SRiver Riddle static Value extractOne(ConversionPatternRewriter &rewriter, 86e62a6956SRiver Riddle LLVMTypeConverter &lowering, Location loc, Value val, 87e62a6956SRiver Riddle Type llvmType, int64_t rank, int64_t pos) { 881c81adf3SAart Bik if (rank == 1) { 891c81adf3SAart Bik auto idxType = rewriter.getIndexType(); 901c81adf3SAart Bik auto constant = rewriter.create<LLVM::ConstantOp>( 911c81adf3SAart Bik loc, lowering.convertType(idxType), 921c81adf3SAart Bik rewriter.getIntegerAttr(idxType, pos)); 931c81adf3SAart Bik return rewriter.create<LLVM::ExtractElementOp>(loc, llvmType, val, 941c81adf3SAart Bik constant); 951c81adf3SAart Bik } 961c81adf3SAart Bik return rewriter.create<LLVM::ExtractValueOp>(loc, llvmType, val, 971c81adf3SAart Bik rewriter.getI64ArrayAttr(pos)); 981c81adf3SAart Bik } 991c81adf3SAart Bik 1002d515e49SNicolas Vasilache // Helper that picks the proper sequence for extracting. 1012d515e49SNicolas Vasilache static Value extractOne(PatternRewriter &rewriter, Location loc, Value vector, 1022d515e49SNicolas Vasilache int64_t offset) { 1032d515e49SNicolas Vasilache auto vectorType = vector.getType().cast<VectorType>(); 1042d515e49SNicolas Vasilache if (vectorType.getRank() > 1) 1052d515e49SNicolas Vasilache return rewriter.create<ExtractOp>(loc, vector, offset); 1062d515e49SNicolas Vasilache return rewriter.create<vector::ExtractElementOp>( 1072d515e49SNicolas Vasilache loc, vectorType.getElementType(), vector, 1082d515e49SNicolas Vasilache rewriter.create<ConstantIndexOp>(loc, offset)); 1092d515e49SNicolas Vasilache } 1102d515e49SNicolas Vasilache 1112d515e49SNicolas Vasilache // Helper that returns a subset of `arrayAttr` as a vector of int64_t. 1122d515e49SNicolas Vasilache // TODO(rriddle): Better support for attribute subtype forwarding + slicing. 1132d515e49SNicolas Vasilache static SmallVector<int64_t, 4> getI64SubArray(ArrayAttr arrayAttr, 1142d515e49SNicolas Vasilache unsigned dropFront = 0, 1152d515e49SNicolas Vasilache unsigned dropBack = 0) { 1162d515e49SNicolas Vasilache assert(arrayAttr.size() > dropFront + dropBack && "Out of bounds"); 1172d515e49SNicolas Vasilache auto range = arrayAttr.getAsRange<IntegerAttr>(); 1182d515e49SNicolas Vasilache SmallVector<int64_t, 4> res; 1192d515e49SNicolas Vasilache res.reserve(arrayAttr.size() - dropFront - dropBack); 1202d515e49SNicolas Vasilache for (auto it = range.begin() + dropFront, eit = range.end() - dropBack; 1212d515e49SNicolas Vasilache it != eit; ++it) 1222d515e49SNicolas Vasilache res.push_back((*it).getValue().getSExtValue()); 1232d515e49SNicolas Vasilache return res; 1242d515e49SNicolas Vasilache } 1252d515e49SNicolas Vasilache 12690c01357SBenjamin Kramer namespace { 127*e83b7b99Saartbik 128b36aaeafSAart Bik class VectorBroadcastOpConversion : public LLVMOpLowering { 129b36aaeafSAart Bik public: 130b36aaeafSAart Bik explicit VectorBroadcastOpConversion(MLIRContext *context, 131b36aaeafSAart Bik LLVMTypeConverter &typeConverter) 132b36aaeafSAart Bik : LLVMOpLowering(vector::BroadcastOp::getOperationName(), context, 133b36aaeafSAart Bik typeConverter) {} 134b36aaeafSAart Bik 135b36aaeafSAart Bik PatternMatchResult 136e62a6956SRiver Riddle matchAndRewrite(Operation *op, ArrayRef<Value> operands, 137b36aaeafSAart Bik ConversionPatternRewriter &rewriter) const override { 138b36aaeafSAart Bik auto broadcastOp = cast<vector::BroadcastOp>(op); 139b36aaeafSAart Bik VectorType dstVectorType = broadcastOp.getVectorType(); 140b36aaeafSAart Bik if (lowering.convertType(dstVectorType) == nullptr) 141b36aaeafSAart Bik return matchFailure(); 142b36aaeafSAart Bik // Rewrite when the full vector type can be lowered (which 143b36aaeafSAart Bik // implies all 'reduced' types can be lowered too). 1441c81adf3SAart Bik auto adaptor = vector::BroadcastOpOperandAdaptor(operands); 145b36aaeafSAart Bik VectorType srcVectorType = 146b36aaeafSAart Bik broadcastOp.getSourceType().dyn_cast<VectorType>(); 147b36aaeafSAart Bik rewriter.replaceOp( 1481c81adf3SAart Bik op, expandRanks(adaptor.source(), // source value to be expanded 149b36aaeafSAart Bik op->getLoc(), // location of original broadcast 150b36aaeafSAart Bik srcVectorType, dstVectorType, rewriter)); 151b36aaeafSAart Bik return matchSuccess(); 152b36aaeafSAart Bik } 153b36aaeafSAart Bik 154b36aaeafSAart Bik private: 155b36aaeafSAart Bik // Expands the given source value over all the ranks, as defined 156b36aaeafSAart Bik // by the source and destination type (a null source type denotes 157b36aaeafSAart Bik // expansion from a scalar value into a vector). 158b36aaeafSAart Bik // 159b36aaeafSAart Bik // TODO(ajcbik): consider replacing this one-pattern lowering 160b36aaeafSAart Bik // with a two-pattern lowering using other vector 161b36aaeafSAart Bik // ops once all insert/extract/shuffle operations 162fc817b09SKazuaki Ishizaki // are available with lowering implementation. 163b36aaeafSAart Bik // 164e62a6956SRiver Riddle Value expandRanks(Value value, Location loc, VectorType srcVectorType, 165b36aaeafSAart Bik VectorType dstVectorType, 166b36aaeafSAart Bik ConversionPatternRewriter &rewriter) const { 167b36aaeafSAart Bik assert((dstVectorType != nullptr) && "invalid result type in broadcast"); 168b36aaeafSAart Bik // Determine rank of source and destination. 169b36aaeafSAart Bik int64_t srcRank = srcVectorType ? srcVectorType.getRank() : 0; 170b36aaeafSAart Bik int64_t dstRank = dstVectorType.getRank(); 171b36aaeafSAart Bik int64_t curDim = dstVectorType.getDimSize(0); 172b36aaeafSAart Bik if (srcRank < dstRank) 173b36aaeafSAart Bik // Duplicate this rank. 174b36aaeafSAart Bik return duplicateOneRank(value, loc, srcVectorType, dstVectorType, dstRank, 175b36aaeafSAart Bik curDim, rewriter); 176b36aaeafSAart Bik // If all trailing dimensions are the same, the broadcast consists of 177b36aaeafSAart Bik // simply passing through the source value and we are done. Otherwise, 178b36aaeafSAart Bik // any non-matching dimension forces a stretch along this rank. 179b36aaeafSAart Bik assert((srcVectorType != nullptr) && (srcRank > 0) && 180b36aaeafSAart Bik (srcRank == dstRank) && "invalid rank in broadcast"); 181b36aaeafSAart Bik for (int64_t r = 0; r < dstRank; r++) { 182b36aaeafSAart Bik if (srcVectorType.getDimSize(r) != dstVectorType.getDimSize(r)) { 183b36aaeafSAart Bik return stretchOneRank(value, loc, srcVectorType, dstVectorType, dstRank, 184b36aaeafSAart Bik curDim, rewriter); 185b36aaeafSAart Bik } 186b36aaeafSAart Bik } 187b36aaeafSAart Bik return value; 188b36aaeafSAart Bik } 189b36aaeafSAart Bik 190b36aaeafSAart Bik // Picks the best way to duplicate a single rank. For the 1-D case, a 191b36aaeafSAart Bik // single insert-elt/shuffle is the most efficient expansion. For higher 192b36aaeafSAart Bik // dimensions, however, we need dim x insert-values on a new broadcast 193b36aaeafSAart Bik // with one less leading dimension, which will be lowered "recursively" 194b36aaeafSAart Bik // to matching LLVM IR. 195b36aaeafSAart Bik // For example: 196b36aaeafSAart Bik // v = broadcast s : f32 to vector<4x2xf32> 197b36aaeafSAart Bik // becomes: 198b36aaeafSAart Bik // x = broadcast s : f32 to vector<2xf32> 199b36aaeafSAart Bik // v = [x,x,x,x] 200b36aaeafSAart Bik // becomes: 201b36aaeafSAart Bik // x = [s,s] 202b36aaeafSAart Bik // v = [x,x,x,x] 203e62a6956SRiver Riddle Value duplicateOneRank(Value value, Location loc, VectorType srcVectorType, 204e62a6956SRiver Riddle VectorType dstVectorType, int64_t rank, int64_t dim, 205b36aaeafSAart Bik ConversionPatternRewriter &rewriter) const { 206b36aaeafSAart Bik Type llvmType = lowering.convertType(dstVectorType); 207b36aaeafSAart Bik assert((llvmType != nullptr) && "unlowerable vector type"); 208b36aaeafSAart Bik if (rank == 1) { 209e62a6956SRiver Riddle Value undef = rewriter.create<LLVM::UndefOp>(loc, llvmType); 210e62a6956SRiver Riddle Value expand = 2111c81adf3SAart Bik insertOne(rewriter, lowering, loc, undef, value, llvmType, rank, 0); 212b36aaeafSAart Bik SmallVector<int32_t, 4> zeroValues(dim, 0); 213b36aaeafSAart Bik return rewriter.create<LLVM::ShuffleVectorOp>( 214b36aaeafSAart Bik loc, expand, undef, rewriter.getI32ArrayAttr(zeroValues)); 215b36aaeafSAart Bik } 216e62a6956SRiver Riddle Value expand = expandRanks(value, loc, srcVectorType, 2179826fe5cSAart Bik reducedVectorTypeFront(dstVectorType), rewriter); 218e62a6956SRiver Riddle Value result = rewriter.create<LLVM::UndefOp>(loc, llvmType); 219b36aaeafSAart Bik for (int64_t d = 0; d < dim; ++d) { 2201c81adf3SAart Bik result = 2211c81adf3SAart Bik insertOne(rewriter, lowering, loc, result, expand, llvmType, rank, d); 222b36aaeafSAart Bik } 223b36aaeafSAart Bik return result; 224b36aaeafSAart Bik } 225b36aaeafSAart Bik 226b36aaeafSAart Bik // Picks the best way to stretch a single rank. For the 1-D case, a 227b36aaeafSAart Bik // single insert-elt/shuffle is the most efficient expansion when at 228b36aaeafSAart Bik // a stretch. Otherwise, every dimension needs to be expanded 229b36aaeafSAart Bik // individually and individually inserted in the resulting vector. 230b36aaeafSAart Bik // For example: 231b36aaeafSAart Bik // v = broadcast w : vector<4x1x2xf32> to vector<4x2x2xf32> 232b36aaeafSAart Bik // becomes: 233b36aaeafSAart Bik // a = broadcast w[0] : vector<1x2xf32> to vector<2x2xf32> 234b36aaeafSAart Bik // b = broadcast w[1] : vector<1x2xf32> to vector<2x2xf32> 235b36aaeafSAart Bik // c = broadcast w[2] : vector<1x2xf32> to vector<2x2xf32> 236b36aaeafSAart Bik // d = broadcast w[3] : vector<1x2xf32> to vector<2x2xf32> 237b36aaeafSAart Bik // v = [a,b,c,d] 238b36aaeafSAart Bik // becomes: 239b36aaeafSAart Bik // x = broadcast w[0][0] : vector<2xf32> to vector <2x2xf32> 240b36aaeafSAart Bik // y = broadcast w[1][0] : vector<2xf32> to vector <2x2xf32> 241b36aaeafSAart Bik // a = [x, y] 242b36aaeafSAart Bik // etc. 243e62a6956SRiver Riddle Value stretchOneRank(Value value, Location loc, VectorType srcVectorType, 244e62a6956SRiver Riddle VectorType dstVectorType, int64_t rank, int64_t dim, 245b36aaeafSAart Bik ConversionPatternRewriter &rewriter) const { 246b36aaeafSAart Bik Type llvmType = lowering.convertType(dstVectorType); 247b36aaeafSAart Bik assert((llvmType != nullptr) && "unlowerable vector type"); 248e62a6956SRiver Riddle Value result = rewriter.create<LLVM::UndefOp>(loc, llvmType); 249b36aaeafSAart Bik bool atStretch = dim != srcVectorType.getDimSize(0); 250b36aaeafSAart Bik if (rank == 1) { 2511c81adf3SAart Bik assert(atStretch); 252b36aaeafSAart Bik Type redLlvmType = lowering.convertType(dstVectorType.getElementType()); 253e62a6956SRiver Riddle Value one = 2541c81adf3SAart Bik extractOne(rewriter, lowering, loc, value, redLlvmType, rank, 0); 255e62a6956SRiver Riddle Value expand = 2561c81adf3SAart Bik insertOne(rewriter, lowering, loc, result, one, llvmType, rank, 0); 257b36aaeafSAart Bik SmallVector<int32_t, 4> zeroValues(dim, 0); 258b36aaeafSAart Bik return rewriter.create<LLVM::ShuffleVectorOp>( 259b36aaeafSAart Bik loc, expand, result, rewriter.getI32ArrayAttr(zeroValues)); 260b36aaeafSAart Bik } 2619826fe5cSAart Bik VectorType redSrcType = reducedVectorTypeFront(srcVectorType); 2629826fe5cSAart Bik VectorType redDstType = reducedVectorTypeFront(dstVectorType); 263b36aaeafSAart Bik Type redLlvmType = lowering.convertType(redSrcType); 264b36aaeafSAart Bik for (int64_t d = 0; d < dim; ++d) { 265b36aaeafSAart Bik int64_t pos = atStretch ? 0 : d; 266e62a6956SRiver Riddle Value one = 2671c81adf3SAart Bik extractOne(rewriter, lowering, loc, value, redLlvmType, rank, pos); 268e62a6956SRiver Riddle Value expand = expandRanks(one, loc, redSrcType, redDstType, rewriter); 2691c81adf3SAart Bik result = 2701c81adf3SAart Bik insertOne(rewriter, lowering, loc, result, expand, llvmType, rank, d); 271b36aaeafSAart Bik } 272b36aaeafSAart Bik return result; 273b36aaeafSAart Bik } 2741c81adf3SAart Bik }; 275b36aaeafSAart Bik 276*e83b7b99Saartbik class VectorReductionOpConversion : public LLVMOpLowering { 277*e83b7b99Saartbik public: 278*e83b7b99Saartbik explicit VectorReductionOpConversion(MLIRContext *context, 279*e83b7b99Saartbik LLVMTypeConverter &typeConverter) 280*e83b7b99Saartbik : LLVMOpLowering(vector::ReductionOp::getOperationName(), context, 281*e83b7b99Saartbik typeConverter) {} 282*e83b7b99Saartbik 283*e83b7b99Saartbik PatternMatchResult 284*e83b7b99Saartbik matchAndRewrite(Operation *op, ArrayRef<Value> operands, 285*e83b7b99Saartbik ConversionPatternRewriter &rewriter) const override { 286*e83b7b99Saartbik auto reductionOp = cast<vector::ReductionOp>(op); 287*e83b7b99Saartbik auto kind = reductionOp.kind(); 288*e83b7b99Saartbik Type eltType = reductionOp.dest().getType(); 289*e83b7b99Saartbik Type llvmType = lowering.convertType(eltType); 290*e83b7b99Saartbik if (eltType.isInteger(32) || eltType.isInteger(64)) { 291*e83b7b99Saartbik // Integer reductions: add/mul/min/max/and/or/xor. 292*e83b7b99Saartbik if (kind == "add") 293*e83b7b99Saartbik rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_add>( 294*e83b7b99Saartbik op, llvmType, operands[0]); 295*e83b7b99Saartbik else if (kind == "mul") 296*e83b7b99Saartbik rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_mul>( 297*e83b7b99Saartbik op, llvmType, operands[0]); 298*e83b7b99Saartbik else if (kind == "min") 299*e83b7b99Saartbik rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_smin>( 300*e83b7b99Saartbik op, llvmType, operands[0]); 301*e83b7b99Saartbik else if (kind == "max") 302*e83b7b99Saartbik rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_smax>( 303*e83b7b99Saartbik op, llvmType, operands[0]); 304*e83b7b99Saartbik else if (kind == "and") 305*e83b7b99Saartbik rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_and>( 306*e83b7b99Saartbik op, llvmType, operands[0]); 307*e83b7b99Saartbik else if (kind == "or") 308*e83b7b99Saartbik rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_or>( 309*e83b7b99Saartbik op, llvmType, operands[0]); 310*e83b7b99Saartbik else if (kind == "xor") 311*e83b7b99Saartbik rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_xor>( 312*e83b7b99Saartbik op, llvmType, operands[0]); 313*e83b7b99Saartbik else 314*e83b7b99Saartbik return matchFailure(); 315*e83b7b99Saartbik return matchSuccess(); 316*e83b7b99Saartbik 317*e83b7b99Saartbik } else if (eltType.isF32() || eltType.isF64()) { 318*e83b7b99Saartbik // Floating-point reductions: add/mul/min/max 319*e83b7b99Saartbik if (kind == "add") { 320*e83b7b99Saartbik Value zero = rewriter.create<LLVM::ConstantOp>( 321*e83b7b99Saartbik op->getLoc(), llvmType, rewriter.getZeroAttr(eltType)); 322*e83b7b99Saartbik rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_v2_fadd>( 323*e83b7b99Saartbik op, llvmType, zero, operands[0]); 324*e83b7b99Saartbik } else if (kind == "mul") { 325*e83b7b99Saartbik Value one = rewriter.create<LLVM::ConstantOp>( 326*e83b7b99Saartbik op->getLoc(), llvmType, rewriter.getFloatAttr(eltType, 1.0)); 327*e83b7b99Saartbik rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_v2_fmul>( 328*e83b7b99Saartbik op, llvmType, one, operands[0]); 329*e83b7b99Saartbik } else if (kind == "min") 330*e83b7b99Saartbik rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_fmin>( 331*e83b7b99Saartbik op, llvmType, operands[0]); 332*e83b7b99Saartbik else if (kind == "max") 333*e83b7b99Saartbik rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_fmax>( 334*e83b7b99Saartbik op, llvmType, operands[0]); 335*e83b7b99Saartbik else 336*e83b7b99Saartbik return matchFailure(); 337*e83b7b99Saartbik return matchSuccess(); 338*e83b7b99Saartbik } 339*e83b7b99Saartbik return matchFailure(); 340*e83b7b99Saartbik } 341*e83b7b99Saartbik }; 342*e83b7b99Saartbik 3431c81adf3SAart Bik class VectorShuffleOpConversion : public LLVMOpLowering { 3441c81adf3SAart Bik public: 3451c81adf3SAart Bik explicit VectorShuffleOpConversion(MLIRContext *context, 3461c81adf3SAart Bik LLVMTypeConverter &typeConverter) 3471c81adf3SAart Bik : LLVMOpLowering(vector::ShuffleOp::getOperationName(), context, 3481c81adf3SAart Bik typeConverter) {} 3491c81adf3SAart Bik 3501c81adf3SAart Bik PatternMatchResult 351e62a6956SRiver Riddle matchAndRewrite(Operation *op, ArrayRef<Value> operands, 3521c81adf3SAart Bik ConversionPatternRewriter &rewriter) const override { 3531c81adf3SAart Bik auto loc = op->getLoc(); 3541c81adf3SAart Bik auto adaptor = vector::ShuffleOpOperandAdaptor(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(); 3591c81adf3SAart Bik Type llvmType = lowering.convertType(vectorType); 3601c81adf3SAart Bik auto maskArrayAttr = shuffleOp.mask(); 3611c81adf3SAart Bik 3621c81adf3SAart Bik // Bail if result type cannot be lowered. 3631c81adf3SAart Bik if (!llvmType) 3641c81adf3SAart Bik return matchFailure(); 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); 3781c81adf3SAart Bik return matchSuccess(); 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 } 391e62a6956SRiver Riddle Value extract = 3921c81adf3SAart Bik extractOne(rewriter, lowering, loc, value, llvmType, rank, extPos); 3931c81adf3SAart Bik insert = insertOne(rewriter, lowering, loc, insert, extract, llvmType, 3941c81adf3SAart Bik rank, insPos++); 3951c81adf3SAart Bik } 3961c81adf3SAart Bik rewriter.replaceOp(op, insert); 3971c81adf3SAart Bik return matchSuccess(); 398b36aaeafSAart Bik } 399b36aaeafSAart Bik }; 400b36aaeafSAart Bik 401cd5dab8aSAart Bik class VectorExtractElementOpConversion : public LLVMOpLowering { 402cd5dab8aSAart Bik public: 403cd5dab8aSAart Bik explicit VectorExtractElementOpConversion(MLIRContext *context, 404cd5dab8aSAart Bik LLVMTypeConverter &typeConverter) 405cd5dab8aSAart Bik : LLVMOpLowering(vector::ExtractElementOp::getOperationName(), context, 406cd5dab8aSAart Bik typeConverter) {} 407cd5dab8aSAart Bik 408cd5dab8aSAart Bik PatternMatchResult 409e62a6956SRiver Riddle matchAndRewrite(Operation *op, ArrayRef<Value> operands, 410cd5dab8aSAart Bik ConversionPatternRewriter &rewriter) const override { 411cd5dab8aSAart Bik auto adaptor = vector::ExtractElementOpOperandAdaptor(operands); 412cd5dab8aSAart Bik auto extractEltOp = cast<vector::ExtractElementOp>(op); 413cd5dab8aSAart Bik auto vectorType = extractEltOp.getVectorType(); 414cd5dab8aSAart Bik auto llvmType = lowering.convertType(vectorType.getElementType()); 415cd5dab8aSAart Bik 416cd5dab8aSAart Bik // Bail if result type cannot be lowered. 417cd5dab8aSAart Bik if (!llvmType) 418cd5dab8aSAart Bik return matchFailure(); 419cd5dab8aSAart Bik 420cd5dab8aSAart Bik rewriter.replaceOpWithNewOp<LLVM::ExtractElementOp>( 421cd5dab8aSAart Bik op, llvmType, adaptor.vector(), adaptor.position()); 422cd5dab8aSAart Bik return matchSuccess(); 423cd5dab8aSAart Bik } 424cd5dab8aSAart Bik }; 425cd5dab8aSAart Bik 4269826fe5cSAart Bik class VectorExtractOpConversion : public LLVMOpLowering { 4275c0c51a9SNicolas Vasilache public: 4289826fe5cSAart Bik explicit VectorExtractOpConversion(MLIRContext *context, 4295c0c51a9SNicolas Vasilache LLVMTypeConverter &typeConverter) 430d37f2725SAart Bik : LLVMOpLowering(vector::ExtractOp::getOperationName(), context, 4315c0c51a9SNicolas Vasilache typeConverter) {} 4325c0c51a9SNicolas Vasilache 4335c0c51a9SNicolas Vasilache PatternMatchResult 434e62a6956SRiver Riddle matchAndRewrite(Operation *op, ArrayRef<Value> operands, 4355c0c51a9SNicolas Vasilache ConversionPatternRewriter &rewriter) const override { 4365c0c51a9SNicolas Vasilache auto loc = op->getLoc(); 437d37f2725SAart Bik auto adaptor = vector::ExtractOpOperandAdaptor(operands); 438d37f2725SAart Bik auto extractOp = cast<vector::ExtractOp>(op); 4399826fe5cSAart Bik auto vectorType = extractOp.getVectorType(); 4402bdf33ccSRiver Riddle auto resultType = extractOp.getResult().getType(); 4415c0c51a9SNicolas Vasilache auto llvmResultType = lowering.convertType(resultType); 4425c0c51a9SNicolas Vasilache auto positionArrayAttr = extractOp.position(); 4439826fe5cSAart Bik 4449826fe5cSAart Bik // Bail if result type cannot be lowered. 4459826fe5cSAart Bik if (!llvmResultType) 4469826fe5cSAart Bik return matchFailure(); 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); 4535c0c51a9SNicolas Vasilache return matchSuccess(); 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>( 4655c0c51a9SNicolas Vasilache loc, lowering.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>(); 4711d47564aSAart Bik auto i64Type = LLVM::LLVMType::getInt64Ty(lowering.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 4775c0c51a9SNicolas Vasilache return matchSuccess(); 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 /// ``` 491681f929fSNicolas Vasilache /// llvm.intr.fma %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 /// ``` 495681f929fSNicolas Vasilache class VectorFMAOp1DConversion : public LLVMOpLowering { 496681f929fSNicolas Vasilache public: 497681f929fSNicolas Vasilache explicit VectorFMAOp1DConversion(MLIRContext *context, 498681f929fSNicolas Vasilache LLVMTypeConverter &typeConverter) 499681f929fSNicolas Vasilache : LLVMOpLowering(vector::FMAOp::getOperationName(), context, 500681f929fSNicolas Vasilache typeConverter) {} 501681f929fSNicolas Vasilache 502681f929fSNicolas Vasilache PatternMatchResult 503681f929fSNicolas Vasilache matchAndRewrite(Operation *op, ArrayRef<Value> operands, 504681f929fSNicolas Vasilache ConversionPatternRewriter &rewriter) const override { 505681f929fSNicolas Vasilache auto adaptor = vector::FMAOpOperandAdaptor(operands); 506681f929fSNicolas Vasilache vector::FMAOp fmaOp = cast<vector::FMAOp>(op); 507681f929fSNicolas Vasilache VectorType vType = fmaOp.getVectorType(); 508681f929fSNicolas Vasilache if (vType.getRank() != 1) 509681f929fSNicolas Vasilache return matchFailure(); 510681f929fSNicolas Vasilache rewriter.replaceOpWithNewOp<LLVM::FMAOp>(op, adaptor.lhs(), adaptor.rhs(), 511681f929fSNicolas Vasilache adaptor.acc()); 512681f929fSNicolas Vasilache return matchSuccess(); 513681f929fSNicolas Vasilache } 514681f929fSNicolas Vasilache }; 515681f929fSNicolas Vasilache 516cd5dab8aSAart Bik class VectorInsertElementOpConversion : public LLVMOpLowering { 517cd5dab8aSAart Bik public: 518cd5dab8aSAart Bik explicit VectorInsertElementOpConversion(MLIRContext *context, 519cd5dab8aSAart Bik LLVMTypeConverter &typeConverter) 520cd5dab8aSAart Bik : LLVMOpLowering(vector::InsertElementOp::getOperationName(), context, 521cd5dab8aSAart Bik typeConverter) {} 522cd5dab8aSAart Bik 523cd5dab8aSAart Bik PatternMatchResult 524e62a6956SRiver Riddle matchAndRewrite(Operation *op, ArrayRef<Value> operands, 525cd5dab8aSAart Bik ConversionPatternRewriter &rewriter) const override { 526cd5dab8aSAart Bik auto adaptor = vector::InsertElementOpOperandAdaptor(operands); 527cd5dab8aSAart Bik auto insertEltOp = cast<vector::InsertElementOp>(op); 528cd5dab8aSAart Bik auto vectorType = insertEltOp.getDestVectorType(); 529cd5dab8aSAart Bik auto llvmType = lowering.convertType(vectorType); 530cd5dab8aSAart Bik 531cd5dab8aSAart Bik // Bail if result type cannot be lowered. 532cd5dab8aSAart Bik if (!llvmType) 533cd5dab8aSAart Bik return matchFailure(); 534cd5dab8aSAart Bik 535cd5dab8aSAart Bik rewriter.replaceOpWithNewOp<LLVM::InsertElementOp>( 536cd5dab8aSAart Bik op, llvmType, adaptor.dest(), adaptor.source(), adaptor.position()); 537cd5dab8aSAart Bik return matchSuccess(); 538cd5dab8aSAart Bik } 539cd5dab8aSAart Bik }; 540cd5dab8aSAart Bik 5419826fe5cSAart Bik class VectorInsertOpConversion : public LLVMOpLowering { 5429826fe5cSAart Bik public: 5439826fe5cSAart Bik explicit VectorInsertOpConversion(MLIRContext *context, 5449826fe5cSAart Bik LLVMTypeConverter &typeConverter) 5459826fe5cSAart Bik : LLVMOpLowering(vector::InsertOp::getOperationName(), context, 5469826fe5cSAart Bik typeConverter) {} 5479826fe5cSAart Bik 5489826fe5cSAart Bik PatternMatchResult 549e62a6956SRiver Riddle matchAndRewrite(Operation *op, ArrayRef<Value> operands, 5509826fe5cSAart Bik ConversionPatternRewriter &rewriter) const override { 5519826fe5cSAart Bik auto loc = op->getLoc(); 5529826fe5cSAart Bik auto adaptor = vector::InsertOpOperandAdaptor(operands); 5539826fe5cSAart Bik auto insertOp = cast<vector::InsertOp>(op); 5549826fe5cSAart Bik auto sourceType = insertOp.getSourceType(); 5559826fe5cSAart Bik auto destVectorType = insertOp.getDestVectorType(); 5569826fe5cSAart Bik auto llvmResultType = lowering.convertType(destVectorType); 5579826fe5cSAart Bik auto positionArrayAttr = insertOp.position(); 5589826fe5cSAart Bik 5599826fe5cSAart Bik // Bail if result type cannot be lowered. 5609826fe5cSAart Bik if (!llvmResultType) 5619826fe5cSAart Bik return matchFailure(); 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); 5699826fe5cSAart Bik return matchSuccess(); 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>( 5839826fe5cSAart Bik loc, lowering.convertType(oneDVectorType), extracted, 5849826fe5cSAart Bik nMinusOnePositionAttrs); 5859826fe5cSAart Bik } 5869826fe5cSAart Bik 5879826fe5cSAart Bik // Insertion of an element into a 1-D LLVM vector. 5881d47564aSAart Bik auto i64Type = LLVM::LLVMType::getInt64Ty(lowering.getDialect()); 5891d47564aSAart Bik auto constant = rewriter.create<LLVM::ConstantOp>(loc, i64Type, position); 590e62a6956SRiver Riddle Value inserted = rewriter.create<LLVM::InsertElementOp>( 5919826fe5cSAart Bik loc, lowering.convertType(oneDVectorType), extracted, adaptor.source(), 5929826fe5cSAart Bik 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); 6049826fe5cSAart Bik return matchSuccess(); 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 633681f929fSNicolas Vasilache PatternMatchResult matchAndRewrite(FMAOp op, 634681f929fSNicolas Vasilache PatternRewriter &rewriter) const override { 635681f929fSNicolas Vasilache auto vType = op.getVectorType(); 636681f929fSNicolas Vasilache if (vType.getRank() < 2) 637681f929fSNicolas Vasilache return matchFailure(); 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); 652681f929fSNicolas Vasilache return matchSuccess(); 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 6762d515e49SNicolas Vasilache PatternMatchResult 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()) 6822d515e49SNicolas Vasilache return matchFailure(); 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) 6882d515e49SNicolas Vasilache return matchFailure(); 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)); 7072d515e49SNicolas Vasilache return matchSuccess(); 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 7252d515e49SNicolas Vasilache PatternMatchResult 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()) 7312d515e49SNicolas Vasilache return matchFailure(); 7322d515e49SNicolas Vasilache 7332d515e49SNicolas Vasilache int64_t rankDiff = dstType.getRank() - srcType.getRank(); 7342d515e49SNicolas Vasilache assert(rankDiff >= 0); 7352d515e49SNicolas Vasilache if (rankDiff != 0) 7362d515e49SNicolas Vasilache return matchFailure(); 7372d515e49SNicolas Vasilache 7382d515e49SNicolas Vasilache if (srcType == dstType) { 7392d515e49SNicolas Vasilache rewriter.replaceOp(op, op.source()); 7402d515e49SNicolas Vasilache return matchSuccess(); 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. 7622d515e49SNicolas Vasilache InsertStridedSliceOp insertStridedSliceOp = 7632d515e49SNicolas Vasilache rewriter.create<InsertStridedSliceOp>( 7642d515e49SNicolas Vasilache loc, extractedSource, extractedDest, 7652d515e49SNicolas Vasilache getI64SubArray(op.offsets(), /* dropFront=*/1), 7662d515e49SNicolas Vasilache getI64SubArray(op.strides(), /* dropFront=*/1)); 7672d515e49SNicolas Vasilache // Call matchAndRewrite recursively from within the pattern. This 7682d515e49SNicolas Vasilache // circumvents the current limitation that a given pattern cannot 7692d515e49SNicolas Vasilache // be called multiple times by the PatternRewrite infrastructure (to 7702d515e49SNicolas Vasilache // avoid infinite recursion, but in this case, infinite recursion 7712d515e49SNicolas Vasilache // cannot happen because the rank is strictly decreasing). 7722d515e49SNicolas Vasilache // TODO(rriddle, nicolasvasilache) Implement something like a hook for 7732d515e49SNicolas Vasilache // a potential function that must decrease and allow the same pattern 7742d515e49SNicolas Vasilache // multiple times. 7752d515e49SNicolas Vasilache auto success = matchAndRewrite(insertStridedSliceOp, rewriter); 7762d515e49SNicolas Vasilache (void)success; 7772d515e49SNicolas Vasilache assert(success && "Unexpected failure"); 7782d515e49SNicolas Vasilache extractedSource = insertStridedSliceOp; 7792d515e49SNicolas Vasilache } 7802d515e49SNicolas Vasilache // 4. Insert the extractedSource into the res vector. 7812d515e49SNicolas Vasilache res = insertOne(rewriter, loc, extractedSource, res, off); 7822d515e49SNicolas Vasilache } 7832d515e49SNicolas Vasilache 7842d515e49SNicolas Vasilache rewriter.replaceOp(op, res); 7852d515e49SNicolas Vasilache return matchSuccess(); 7862d515e49SNicolas Vasilache } 7872d515e49SNicolas Vasilache }; 7882d515e49SNicolas Vasilache 7895c0c51a9SNicolas Vasilache class VectorOuterProductOpConversion : public LLVMOpLowering { 7905c0c51a9SNicolas Vasilache public: 7915c0c51a9SNicolas Vasilache explicit VectorOuterProductOpConversion(MLIRContext *context, 7925c0c51a9SNicolas Vasilache LLVMTypeConverter &typeConverter) 7935c0c51a9SNicolas Vasilache : LLVMOpLowering(vector::OuterProductOp::getOperationName(), context, 7945c0c51a9SNicolas Vasilache typeConverter) {} 7955c0c51a9SNicolas Vasilache 7965c0c51a9SNicolas Vasilache PatternMatchResult 797e62a6956SRiver Riddle matchAndRewrite(Operation *op, ArrayRef<Value> operands, 7985c0c51a9SNicolas Vasilache ConversionPatternRewriter &rewriter) const override { 7995c0c51a9SNicolas Vasilache auto loc = op->getLoc(); 8005c0c51a9SNicolas Vasilache auto adaptor = vector::OuterProductOpOperandAdaptor(operands); 8015c0c51a9SNicolas Vasilache auto *ctx = op->getContext(); 8022bdf33ccSRiver Riddle auto vLHS = adaptor.lhs().getType().cast<LLVM::LLVMType>(); 8032bdf33ccSRiver Riddle auto vRHS = adaptor.rhs().getType().cast<LLVM::LLVMType>(); 8045c0c51a9SNicolas Vasilache auto rankLHS = vLHS.getUnderlyingType()->getVectorNumElements(); 8055c0c51a9SNicolas Vasilache auto rankRHS = vRHS.getUnderlyingType()->getVectorNumElements(); 8065c0c51a9SNicolas Vasilache auto llvmArrayOfVectType = lowering.convertType( 8072bdf33ccSRiver Riddle cast<vector::OuterProductOp>(op).getResult().getType()); 808e62a6956SRiver Riddle Value desc = rewriter.create<LLVM::UndefOp>(loc, llvmArrayOfVectType); 809e62a6956SRiver Riddle Value a = adaptor.lhs(), b = adaptor.rhs(); 810e62a6956SRiver Riddle Value acc = adaptor.acc().empty() ? nullptr : adaptor.acc().front(); 811e62a6956SRiver Riddle SmallVector<Value, 8> lhs, accs; 8125c0c51a9SNicolas Vasilache lhs.reserve(rankLHS); 8135c0c51a9SNicolas Vasilache accs.reserve(rankLHS); 8145c0c51a9SNicolas Vasilache for (unsigned d = 0, e = rankLHS; d < e; ++d) { 8155c0c51a9SNicolas Vasilache // shufflevector explicitly requires i32. 8165c0c51a9SNicolas Vasilache auto attr = rewriter.getI32IntegerAttr(d); 8175c0c51a9SNicolas Vasilache SmallVector<Attribute, 4> bcastAttr(rankRHS, attr); 8185c0c51a9SNicolas Vasilache auto bcastArrayAttr = ArrayAttr::get(bcastAttr, ctx); 819e62a6956SRiver Riddle Value aD = nullptr, accD = nullptr; 8205c0c51a9SNicolas Vasilache // 1. Broadcast the element a[d] into vector aD. 8215c0c51a9SNicolas Vasilache aD = rewriter.create<LLVM::ShuffleVectorOp>(loc, a, a, bcastArrayAttr); 8225c0c51a9SNicolas Vasilache // 2. If acc is present, extract 1-d vector acc[d] into accD. 8235c0c51a9SNicolas Vasilache if (acc) 8245c0c51a9SNicolas Vasilache accD = rewriter.create<LLVM::ExtractValueOp>( 8255c0c51a9SNicolas Vasilache loc, vRHS, acc, rewriter.getI64ArrayAttr(d)); 8265c0c51a9SNicolas Vasilache // 3. Compute aD outer b (plus accD, if relevant). 827e62a6956SRiver Riddle Value aOuterbD = 828499ad458SNicolas Vasilache accD 829499ad458SNicolas Vasilache ? rewriter.create<LLVM::FMAOp>(loc, vRHS, aD, b, accD).getResult() 8305c0c51a9SNicolas Vasilache : rewriter.create<LLVM::FMulOp>(loc, aD, b).getResult(); 8315c0c51a9SNicolas Vasilache // 4. Insert as value `d` in the descriptor. 8325c0c51a9SNicolas Vasilache desc = rewriter.create<LLVM::InsertValueOp>(loc, llvmArrayOfVectType, 8335c0c51a9SNicolas Vasilache desc, aOuterbD, 8345c0c51a9SNicolas Vasilache rewriter.getI64ArrayAttr(d)); 8355c0c51a9SNicolas Vasilache } 8365c0c51a9SNicolas Vasilache rewriter.replaceOp(op, desc); 8375c0c51a9SNicolas Vasilache return matchSuccess(); 8385c0c51a9SNicolas Vasilache } 8395c0c51a9SNicolas Vasilache }; 8405c0c51a9SNicolas Vasilache 8415c0c51a9SNicolas Vasilache class VectorTypeCastOpConversion : public LLVMOpLowering { 8425c0c51a9SNicolas Vasilache public: 8435c0c51a9SNicolas Vasilache explicit VectorTypeCastOpConversion(MLIRContext *context, 8445c0c51a9SNicolas Vasilache LLVMTypeConverter &typeConverter) 8455c0c51a9SNicolas Vasilache : LLVMOpLowering(vector::TypeCastOp::getOperationName(), context, 8465c0c51a9SNicolas Vasilache typeConverter) {} 8475c0c51a9SNicolas Vasilache 8485c0c51a9SNicolas Vasilache PatternMatchResult 849e62a6956SRiver Riddle matchAndRewrite(Operation *op, ArrayRef<Value> operands, 8505c0c51a9SNicolas Vasilache ConversionPatternRewriter &rewriter) const override { 8515c0c51a9SNicolas Vasilache auto loc = op->getLoc(); 8525c0c51a9SNicolas Vasilache vector::TypeCastOp castOp = cast<vector::TypeCastOp>(op); 8535c0c51a9SNicolas Vasilache MemRefType sourceMemRefType = 8542bdf33ccSRiver Riddle castOp.getOperand().getType().cast<MemRefType>(); 8555c0c51a9SNicolas Vasilache MemRefType targetMemRefType = 8562bdf33ccSRiver Riddle castOp.getResult().getType().cast<MemRefType>(); 8575c0c51a9SNicolas Vasilache 8585c0c51a9SNicolas Vasilache // Only static shape casts supported atm. 8595c0c51a9SNicolas Vasilache if (!sourceMemRefType.hasStaticShape() || 8605c0c51a9SNicolas Vasilache !targetMemRefType.hasStaticShape()) 8615c0c51a9SNicolas Vasilache return matchFailure(); 8625c0c51a9SNicolas Vasilache 8635c0c51a9SNicolas Vasilache auto llvmSourceDescriptorTy = 8642bdf33ccSRiver Riddle operands[0].getType().dyn_cast<LLVM::LLVMType>(); 8655c0c51a9SNicolas Vasilache if (!llvmSourceDescriptorTy || !llvmSourceDescriptorTy.isStructTy()) 8665c0c51a9SNicolas Vasilache return matchFailure(); 8675c0c51a9SNicolas Vasilache MemRefDescriptor sourceMemRef(operands[0]); 8685c0c51a9SNicolas Vasilache 8695c0c51a9SNicolas Vasilache auto llvmTargetDescriptorTy = lowering.convertType(targetMemRefType) 8705c0c51a9SNicolas Vasilache .dyn_cast_or_null<LLVM::LLVMType>(); 8715c0c51a9SNicolas Vasilache if (!llvmTargetDescriptorTy || !llvmTargetDescriptorTy.isStructTy()) 8725c0c51a9SNicolas Vasilache return matchFailure(); 8735c0c51a9SNicolas Vasilache 8745c0c51a9SNicolas Vasilache int64_t offset; 8755c0c51a9SNicolas Vasilache SmallVector<int64_t, 4> strides; 8765c0c51a9SNicolas Vasilache auto successStrides = 8775c0c51a9SNicolas Vasilache getStridesAndOffset(sourceMemRefType, strides, offset); 8785c0c51a9SNicolas Vasilache bool isContiguous = (strides.back() == 1); 8795c0c51a9SNicolas Vasilache if (isContiguous) { 8805c0c51a9SNicolas Vasilache auto sizes = sourceMemRefType.getShape(); 8815c0c51a9SNicolas Vasilache for (int index = 0, e = strides.size() - 2; index < e; ++index) { 8825c0c51a9SNicolas Vasilache if (strides[index] != strides[index + 1] * sizes[index + 1]) { 8835c0c51a9SNicolas Vasilache isContiguous = false; 8845c0c51a9SNicolas Vasilache break; 8855c0c51a9SNicolas Vasilache } 8865c0c51a9SNicolas Vasilache } 8875c0c51a9SNicolas Vasilache } 8885c0c51a9SNicolas Vasilache // Only contiguous source tensors supported atm. 8895c0c51a9SNicolas Vasilache if (failed(successStrides) || !isContiguous) 8905c0c51a9SNicolas Vasilache return matchFailure(); 8915c0c51a9SNicolas Vasilache 8925c0c51a9SNicolas Vasilache auto int64Ty = LLVM::LLVMType::getInt64Ty(lowering.getDialect()); 8935c0c51a9SNicolas Vasilache 8945c0c51a9SNicolas Vasilache // Create descriptor. 8955c0c51a9SNicolas Vasilache auto desc = MemRefDescriptor::undef(rewriter, loc, llvmTargetDescriptorTy); 8965c0c51a9SNicolas Vasilache Type llvmTargetElementTy = desc.getElementType(); 8975c0c51a9SNicolas Vasilache // Set allocated ptr. 898e62a6956SRiver Riddle Value allocated = sourceMemRef.allocatedPtr(rewriter, loc); 8995c0c51a9SNicolas Vasilache allocated = 9005c0c51a9SNicolas Vasilache rewriter.create<LLVM::BitcastOp>(loc, llvmTargetElementTy, allocated); 9015c0c51a9SNicolas Vasilache desc.setAllocatedPtr(rewriter, loc, allocated); 9025c0c51a9SNicolas Vasilache // Set aligned ptr. 903e62a6956SRiver Riddle Value ptr = sourceMemRef.alignedPtr(rewriter, loc); 9045c0c51a9SNicolas Vasilache ptr = rewriter.create<LLVM::BitcastOp>(loc, llvmTargetElementTy, ptr); 9055c0c51a9SNicolas Vasilache desc.setAlignedPtr(rewriter, loc, ptr); 9065c0c51a9SNicolas Vasilache // Fill offset 0. 9075c0c51a9SNicolas Vasilache auto attr = rewriter.getIntegerAttr(rewriter.getIndexType(), 0); 9085c0c51a9SNicolas Vasilache auto zero = rewriter.create<LLVM::ConstantOp>(loc, int64Ty, attr); 9095c0c51a9SNicolas Vasilache desc.setOffset(rewriter, loc, zero); 9105c0c51a9SNicolas Vasilache 9115c0c51a9SNicolas Vasilache // Fill size and stride descriptors in memref. 9125c0c51a9SNicolas Vasilache for (auto indexedSize : llvm::enumerate(targetMemRefType.getShape())) { 9135c0c51a9SNicolas Vasilache int64_t index = indexedSize.index(); 9145c0c51a9SNicolas Vasilache auto sizeAttr = 9155c0c51a9SNicolas Vasilache rewriter.getIntegerAttr(rewriter.getIndexType(), indexedSize.value()); 9165c0c51a9SNicolas Vasilache auto size = rewriter.create<LLVM::ConstantOp>(loc, int64Ty, sizeAttr); 9175c0c51a9SNicolas Vasilache desc.setSize(rewriter, loc, index, size); 9185c0c51a9SNicolas Vasilache auto strideAttr = 9195c0c51a9SNicolas Vasilache rewriter.getIntegerAttr(rewriter.getIndexType(), strides[index]); 9205c0c51a9SNicolas Vasilache auto stride = rewriter.create<LLVM::ConstantOp>(loc, int64Ty, strideAttr); 9215c0c51a9SNicolas Vasilache desc.setStride(rewriter, loc, index, stride); 9225c0c51a9SNicolas Vasilache } 9235c0c51a9SNicolas Vasilache 9245c0c51a9SNicolas Vasilache rewriter.replaceOp(op, {desc}); 9255c0c51a9SNicolas Vasilache return matchSuccess(); 9265c0c51a9SNicolas Vasilache } 9275c0c51a9SNicolas Vasilache }; 9285c0c51a9SNicolas Vasilache 929d9b500d3SAart Bik class VectorPrintOpConversion : public LLVMOpLowering { 930d9b500d3SAart Bik public: 931d9b500d3SAart Bik explicit VectorPrintOpConversion(MLIRContext *context, 932d9b500d3SAart Bik LLVMTypeConverter &typeConverter) 933d9b500d3SAart Bik : LLVMOpLowering(vector::PrintOp::getOperationName(), context, 934d9b500d3SAart Bik typeConverter) {} 935d9b500d3SAart Bik 936d9b500d3SAart Bik // Proof-of-concept lowering implementation that relies on a small 937d9b500d3SAart Bik // runtime support library, which only needs to provide a few 938d9b500d3SAart Bik // printing methods (single value for all data types, opening/closing 939d9b500d3SAart Bik // bracket, comma, newline). The lowering fully unrolls a vector 940d9b500d3SAart Bik // in terms of these elementary printing operations. The advantage 941d9b500d3SAart Bik // of this approach is that the library can remain unaware of all 942d9b500d3SAart Bik // low-level implementation details of vectors while still supporting 943d9b500d3SAart Bik // output of any shaped and dimensioned vector. Due to full unrolling, 944d9b500d3SAart Bik // this approach is less suited for very large vectors though. 945d9b500d3SAart Bik // 946d9b500d3SAart Bik // TODO(ajcbik): rely solely on libc in future? something else? 947d9b500d3SAart Bik // 948d9b500d3SAart Bik PatternMatchResult 949e62a6956SRiver Riddle matchAndRewrite(Operation *op, ArrayRef<Value> operands, 950d9b500d3SAart Bik ConversionPatternRewriter &rewriter) const override { 951d9b500d3SAart Bik auto printOp = cast<vector::PrintOp>(op); 952d9b500d3SAart Bik auto adaptor = vector::PrintOpOperandAdaptor(operands); 953d9b500d3SAart Bik Type printType = printOp.getPrintType(); 954d9b500d3SAart Bik 955d9b500d3SAart Bik if (lowering.convertType(printType) == nullptr) 956d9b500d3SAart Bik return matchFailure(); 957d9b500d3SAart Bik 958d9b500d3SAart Bik // Make sure element type has runtime support (currently just Float/Double). 959d9b500d3SAart Bik VectorType vectorType = printType.dyn_cast<VectorType>(); 960d9b500d3SAart Bik Type eltType = vectorType ? vectorType.getElementType() : printType; 961d9b500d3SAart Bik int64_t rank = vectorType ? vectorType.getRank() : 0; 962d9b500d3SAart Bik Operation *printer; 963e52414b1Saartbik if (eltType.isInteger(32)) 964e52414b1Saartbik printer = getPrintI32(op); 965e52414b1Saartbik else if (eltType.isInteger(64)) 966e52414b1Saartbik printer = getPrintI64(op); 967e52414b1Saartbik else if (eltType.isF32()) 968d9b500d3SAart Bik printer = getPrintFloat(op); 969d9b500d3SAart Bik else if (eltType.isF64()) 970d9b500d3SAart Bik printer = getPrintDouble(op); 971d9b500d3SAart Bik else 972d9b500d3SAart Bik return matchFailure(); 973d9b500d3SAart Bik 974d9b500d3SAart Bik // Unroll vector into elementary print calls. 975d9b500d3SAart Bik emitRanks(rewriter, op, adaptor.source(), vectorType, printer, rank); 976d9b500d3SAart Bik emitCall(rewriter, op->getLoc(), getPrintNewline(op)); 977d9b500d3SAart Bik rewriter.eraseOp(op); 978d9b500d3SAart Bik return matchSuccess(); 979d9b500d3SAart Bik } 980d9b500d3SAart Bik 981d9b500d3SAart Bik private: 982d9b500d3SAart Bik void emitRanks(ConversionPatternRewriter &rewriter, Operation *op, 983e62a6956SRiver Riddle Value value, VectorType vectorType, Operation *printer, 984d9b500d3SAart Bik int64_t rank) const { 985d9b500d3SAart Bik Location loc = op->getLoc(); 986d9b500d3SAart Bik if (rank == 0) { 987d9b500d3SAart Bik emitCall(rewriter, loc, printer, value); 988d9b500d3SAart Bik return; 989d9b500d3SAart Bik } 990d9b500d3SAart Bik 991d9b500d3SAart Bik emitCall(rewriter, loc, getPrintOpen(op)); 992d9b500d3SAart Bik Operation *printComma = getPrintComma(op); 993d9b500d3SAart Bik int64_t dim = vectorType.getDimSize(0); 994d9b500d3SAart Bik for (int64_t d = 0; d < dim; ++d) { 995d9b500d3SAart Bik auto reducedType = 996d9b500d3SAart Bik rank > 1 ? reducedVectorTypeFront(vectorType) : nullptr; 997d9b500d3SAart Bik auto llvmType = lowering.convertType( 998d9b500d3SAart Bik rank > 1 ? reducedType : vectorType.getElementType()); 999e62a6956SRiver Riddle Value nestedVal = 1000d9b500d3SAart Bik extractOne(rewriter, lowering, loc, value, llvmType, rank, d); 1001d9b500d3SAart Bik emitRanks(rewriter, op, nestedVal, reducedType, printer, rank - 1); 1002d9b500d3SAart Bik if (d != dim - 1) 1003d9b500d3SAart Bik emitCall(rewriter, loc, printComma); 1004d9b500d3SAart Bik } 1005d9b500d3SAart Bik emitCall(rewriter, loc, getPrintClose(op)); 1006d9b500d3SAart Bik } 1007d9b500d3SAart Bik 1008d9b500d3SAart Bik // Helper to emit a call. 1009d9b500d3SAart Bik static void emitCall(ConversionPatternRewriter &rewriter, Location loc, 1010d9b500d3SAart Bik Operation *ref, ValueRange params = ValueRange()) { 1011d9b500d3SAart Bik rewriter.create<LLVM::CallOp>(loc, ArrayRef<Type>{}, 1012d9b500d3SAart Bik rewriter.getSymbolRefAttr(ref), params); 1013d9b500d3SAart Bik } 1014d9b500d3SAart Bik 1015d9b500d3SAart Bik // Helper for printer method declaration (first hit) and lookup. 1016d9b500d3SAart Bik static Operation *getPrint(Operation *op, LLVM::LLVMDialect *dialect, 1017d9b500d3SAart Bik StringRef name, ArrayRef<LLVM::LLVMType> params) { 1018d9b500d3SAart Bik auto module = op->getParentOfType<ModuleOp>(); 1019d9b500d3SAart Bik auto func = module.lookupSymbol<LLVM::LLVMFuncOp>(name); 1020d9b500d3SAart Bik if (func) 1021d9b500d3SAart Bik return func; 1022d9b500d3SAart Bik OpBuilder moduleBuilder(module.getBodyRegion()); 1023d9b500d3SAart Bik return moduleBuilder.create<LLVM::LLVMFuncOp>( 1024d9b500d3SAart Bik op->getLoc(), name, 1025d9b500d3SAart Bik LLVM::LLVMType::getFunctionTy(LLVM::LLVMType::getVoidTy(dialect), 1026d9b500d3SAart Bik params, /*isVarArg=*/false)); 1027d9b500d3SAart Bik } 1028d9b500d3SAart Bik 1029d9b500d3SAart Bik // Helpers for method names. 1030e52414b1Saartbik Operation *getPrintI32(Operation *op) const { 1031e52414b1Saartbik LLVM::LLVMDialect *dialect = lowering.getDialect(); 1032e52414b1Saartbik return getPrint(op, dialect, "print_i32", 1033e52414b1Saartbik LLVM::LLVMType::getInt32Ty(dialect)); 1034e52414b1Saartbik } 1035e52414b1Saartbik Operation *getPrintI64(Operation *op) const { 1036e52414b1Saartbik LLVM::LLVMDialect *dialect = lowering.getDialect(); 1037e52414b1Saartbik return getPrint(op, dialect, "print_i64", 1038e52414b1Saartbik LLVM::LLVMType::getInt64Ty(dialect)); 1039e52414b1Saartbik } 1040d9b500d3SAart Bik Operation *getPrintFloat(Operation *op) const { 1041d9b500d3SAart Bik LLVM::LLVMDialect *dialect = lowering.getDialect(); 1042d9b500d3SAart Bik return getPrint(op, dialect, "print_f32", 1043d9b500d3SAart Bik LLVM::LLVMType::getFloatTy(dialect)); 1044d9b500d3SAart Bik } 1045d9b500d3SAart Bik Operation *getPrintDouble(Operation *op) const { 1046d9b500d3SAart Bik LLVM::LLVMDialect *dialect = lowering.getDialect(); 1047d9b500d3SAart Bik return getPrint(op, dialect, "print_f64", 1048d9b500d3SAart Bik LLVM::LLVMType::getDoubleTy(dialect)); 1049d9b500d3SAart Bik } 1050d9b500d3SAart Bik Operation *getPrintOpen(Operation *op) const { 1051d9b500d3SAart Bik return getPrint(op, lowering.getDialect(), "print_open", {}); 1052d9b500d3SAart Bik } 1053d9b500d3SAart Bik Operation *getPrintClose(Operation *op) const { 1054d9b500d3SAart Bik return getPrint(op, lowering.getDialect(), "print_close", {}); 1055d9b500d3SAart Bik } 1056d9b500d3SAart Bik Operation *getPrintComma(Operation *op) const { 1057d9b500d3SAart Bik return getPrint(op, lowering.getDialect(), "print_comma", {}); 1058d9b500d3SAart Bik } 1059d9b500d3SAart Bik Operation *getPrintNewline(Operation *op) const { 1060d9b500d3SAart Bik return getPrint(op, lowering.getDialect(), "print_newline", {}); 1061d9b500d3SAart Bik } 1062d9b500d3SAart Bik }; 1063d9b500d3SAart Bik 106465678d93SNicolas Vasilache /// Progressive lowering of StridedSliceOp to either: 106565678d93SNicolas Vasilache /// 1. extractelement + insertelement for the 1-D case 106665678d93SNicolas Vasilache /// 2. extract + optional strided_slice + insert for the n-D case. 10672d515e49SNicolas Vasilache class VectorStridedSliceOpConversion : public OpRewritePattern<StridedSliceOp> { 106865678d93SNicolas Vasilache public: 106965678d93SNicolas Vasilache using OpRewritePattern<StridedSliceOp>::OpRewritePattern; 107065678d93SNicolas Vasilache 107165678d93SNicolas Vasilache PatternMatchResult matchAndRewrite(StridedSliceOp op, 107265678d93SNicolas Vasilache PatternRewriter &rewriter) const override { 107365678d93SNicolas Vasilache auto dstType = op.getResult().getType().cast<VectorType>(); 107465678d93SNicolas Vasilache 107565678d93SNicolas Vasilache assert(!op.offsets().getValue().empty() && "Unexpected empty offsets"); 107665678d93SNicolas Vasilache 107765678d93SNicolas Vasilache int64_t offset = 107865678d93SNicolas Vasilache op.offsets().getValue().front().cast<IntegerAttr>().getInt(); 107965678d93SNicolas Vasilache int64_t size = op.sizes().getValue().front().cast<IntegerAttr>().getInt(); 108065678d93SNicolas Vasilache int64_t stride = 108165678d93SNicolas Vasilache op.strides().getValue().front().cast<IntegerAttr>().getInt(); 108265678d93SNicolas Vasilache 108365678d93SNicolas Vasilache auto loc = op.getLoc(); 108465678d93SNicolas Vasilache auto elemType = dstType.getElementType(); 108565678d93SNicolas Vasilache assert(elemType.isIntOrIndexOrFloat()); 108665678d93SNicolas Vasilache Value zero = rewriter.create<ConstantOp>(loc, elemType, 108765678d93SNicolas Vasilache rewriter.getZeroAttr(elemType)); 108865678d93SNicolas Vasilache Value res = rewriter.create<SplatOp>(loc, dstType, zero); 108965678d93SNicolas Vasilache for (int64_t off = offset, e = offset + size * stride, idx = 0; off < e; 109065678d93SNicolas Vasilache off += stride, ++idx) { 109165678d93SNicolas Vasilache Value extracted = extractOne(rewriter, loc, op.vector(), off); 109265678d93SNicolas Vasilache if (op.offsets().getValue().size() > 1) { 109365678d93SNicolas Vasilache StridedSliceOp stridedSliceOp = rewriter.create<StridedSliceOp>( 109465678d93SNicolas Vasilache loc, extracted, getI64SubArray(op.offsets(), /* dropFront=*/1), 109565678d93SNicolas Vasilache getI64SubArray(op.sizes(), /* dropFront=*/1), 109665678d93SNicolas Vasilache getI64SubArray(op.strides(), /* dropFront=*/1)); 109765678d93SNicolas Vasilache // Call matchAndRewrite recursively from within the pattern. This 109865678d93SNicolas Vasilache // circumvents the current limitation that a given pattern cannot 109965678d93SNicolas Vasilache // be called multiple times by the PatternRewrite infrastructure (to 110065678d93SNicolas Vasilache // avoid infinite recursion, but in this case, infinite recursion 110165678d93SNicolas Vasilache // cannot happen because the rank is strictly decreasing). 110265678d93SNicolas Vasilache // TODO(rriddle, nicolasvasilache) Implement something like a hook for 110365678d93SNicolas Vasilache // a potential function that must decrease and allow the same pattern 110465678d93SNicolas Vasilache // multiple times. 110565678d93SNicolas Vasilache auto success = matchAndRewrite(stridedSliceOp, rewriter); 110665678d93SNicolas Vasilache (void)success; 110765678d93SNicolas Vasilache assert(success && "Unexpected failure"); 110865678d93SNicolas Vasilache extracted = stridedSliceOp; 110965678d93SNicolas Vasilache } 111065678d93SNicolas Vasilache res = insertOne(rewriter, loc, extracted, res, idx); 111165678d93SNicolas Vasilache } 111265678d93SNicolas Vasilache rewriter.replaceOp(op, {res}); 111365678d93SNicolas Vasilache return matchSuccess(); 111465678d93SNicolas Vasilache } 111565678d93SNicolas Vasilache }; 111665678d93SNicolas Vasilache 1117df186507SBenjamin Kramer } // namespace 1118df186507SBenjamin Kramer 11195c0c51a9SNicolas Vasilache /// Populate the given list with patterns that convert from Vector to LLVM. 11205c0c51a9SNicolas Vasilache void mlir::populateVectorToLLVMConversionPatterns( 11215c0c51a9SNicolas Vasilache LLVMTypeConverter &converter, OwningRewritePatternList &patterns) { 112265678d93SNicolas Vasilache MLIRContext *ctx = converter.getDialect()->getContext(); 1123681f929fSNicolas Vasilache patterns.insert<VectorFMAOpNDRewritePattern, 1124681f929fSNicolas Vasilache VectorInsertStridedSliceOpDifferentRankRewritePattern, 11252d515e49SNicolas Vasilache VectorInsertStridedSliceOpSameRankRewritePattern, 11262d515e49SNicolas Vasilache VectorStridedSliceOpConversion>(ctx); 1127*e83b7b99Saartbik patterns.insert<VectorBroadcastOpConversion, VectorReductionOpConversion, 1128*e83b7b99Saartbik VectorShuffleOpConversion, VectorExtractElementOpConversion, 1129*e83b7b99Saartbik VectorExtractOpConversion, VectorFMAOp1DConversion, 1130*e83b7b99Saartbik VectorInsertElementOpConversion, VectorInsertOpConversion, 1131*e83b7b99Saartbik VectorOuterProductOpConversion, VectorTypeCastOpConversion, 1132*e83b7b99Saartbik VectorPrintOpConversion>(ctx, converter); 11335c0c51a9SNicolas Vasilache } 11345c0c51a9SNicolas Vasilache 11355c0c51a9SNicolas Vasilache namespace { 11365c0c51a9SNicolas Vasilache struct LowerVectorToLLVMPass : public ModulePass<LowerVectorToLLVMPass> { 11375c0c51a9SNicolas Vasilache void runOnModule() override; 11385c0c51a9SNicolas Vasilache }; 11395c0c51a9SNicolas Vasilache } // namespace 11405c0c51a9SNicolas Vasilache 11415c0c51a9SNicolas Vasilache void LowerVectorToLLVMPass::runOnModule() { 1142459cf6e5Saartbik // Perform progressive lowering of operations on "slices". 1143459cf6e5Saartbik // Folding and DCE get rid of all non-leaking tuple ops. 1144459cf6e5Saartbik { 11455c0c51a9SNicolas Vasilache OwningRewritePatternList patterns; 1146459cf6e5Saartbik populateVectorSlicesLoweringPatterns(patterns, &getContext()); 1147459cf6e5Saartbik applyPatternsGreedily(getModule(), patterns); 1148459cf6e5Saartbik } 1149459cf6e5Saartbik 1150459cf6e5Saartbik // Convert to the LLVM IR dialect. 11515c0c51a9SNicolas Vasilache LLVMTypeConverter converter(&getContext()); 1152459cf6e5Saartbik OwningRewritePatternList patterns; 11535c0c51a9SNicolas Vasilache populateVectorToLLVMConversionPatterns(converter, patterns); 11545c0c51a9SNicolas Vasilache populateStdToLLVMConversionPatterns(converter, patterns); 11555c0c51a9SNicolas Vasilache 11565c0c51a9SNicolas Vasilache ConversionTarget target(getContext()); 11575c0c51a9SNicolas Vasilache target.addLegalDialect<LLVM::LLVMDialect>(); 11585c0c51a9SNicolas Vasilache target.addDynamicallyLegalOp<FuncOp>( 11595c0c51a9SNicolas Vasilache [&](FuncOp op) { return converter.isSignatureLegal(op.getType()); }); 11605c0c51a9SNicolas Vasilache if (failed( 11615c0c51a9SNicolas Vasilache applyPartialConversion(getModule(), target, patterns, &converter))) { 11625c0c51a9SNicolas Vasilache signalPassFailure(); 11635c0c51a9SNicolas Vasilache } 11645c0c51a9SNicolas Vasilache } 11655c0c51a9SNicolas Vasilache 11665c0c51a9SNicolas Vasilache OpPassBase<ModuleOp> *mlir::createLowerVectorToLLVMPass() { 11675c0c51a9SNicolas Vasilache return new LowerVectorToLLVMPass(); 11685c0c51a9SNicolas Vasilache } 11695c0c51a9SNicolas Vasilache 11705c0c51a9SNicolas Vasilache static PassRegistration<LowerVectorToLLVMPass> 11715c0c51a9SNicolas Vasilache pass("convert-vector-to-llvm", 11725c0c51a9SNicolas Vasilache "Lower the operations from the vector dialect into the LLVM dialect"); 1173