1 //===- LowerVectorInterleave.cpp - Lower 'vector.interleave' operation ----===// 2 // 3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. 4 // See https://llvm.org/LICENSE.txt for license information. 5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception 6 // 7 //===----------------------------------------------------------------------===// 8 // 9 // This file implements target-independent rewrites and utilities to lower the 10 // 'vector.interleave' operation. 11 // 12 //===----------------------------------------------------------------------===// 13 14 #include "mlir/Dialect/Vector/IR/VectorOps.h" 15 #include "mlir/Dialect/Vector/Transforms/LoweringPatterns.h" 16 #include "mlir/Dialect/Vector/Utils/VectorUtils.h" 17 #include "mlir/IR/BuiltinTypes.h" 18 #include "mlir/IR/PatternMatch.h" 19 #include "mlir/Support/LogicalResult.h" 20 21 #define DEBUG_TYPE "vector-interleave-lowering" 22 23 using namespace mlir; 24 using namespace mlir::vector; 25 26 namespace { 27 28 /// A one-shot unrolling of vector.interleave to the `targetRank`. 29 /// 30 /// Example: 31 /// 32 /// ```mlir 33 /// vector.interleave %a, %b : vector<1x2x3x4xi64> -> vector<1x2x3x8xi64> 34 /// ``` 35 /// Would be unrolled to: 36 /// ```mlir 37 /// %result = arith.constant dense<0> : vector<1x2x3x8xi64> 38 /// %0 = vector.extract %a[0, 0, 0] ─┐ 39 /// : vector<4xi64> from vector<1x2x3x4xi64> | 40 /// %1 = vector.extract %b[0, 0, 0] | 41 /// : vector<4xi64> from vector<1x2x3x4xi64> | - Repeated 6x for 42 /// %2 = vector.interleave %0, %1 : | all leading positions 43 /// : vector<4xi64> -> vector<8xi64> | 44 /// %3 = vector.insert %2, %result [0, 0, 0] | 45 /// : vector<8xi64> into vector<1x2x3x8xi64> ┘ 46 /// ``` 47 /// 48 /// Note: If any leading dimension before the `targetRank` is scalable the 49 /// unrolling will stop before the scalable dimension. 50 class UnrollInterleaveOp final : public OpRewritePattern<vector::InterleaveOp> { 51 public: 52 UnrollInterleaveOp(int64_t targetRank, MLIRContext *context, 53 PatternBenefit benefit = 1) 54 : OpRewritePattern(context, benefit), targetRank(targetRank){}; 55 56 LogicalResult matchAndRewrite(vector::InterleaveOp op, 57 PatternRewriter &rewriter) const override { 58 VectorType resultType = op.getResultVectorType(); 59 auto unrollIterator = vector::createUnrollIterator(resultType, targetRank); 60 if (!unrollIterator) 61 return failure(); 62 63 auto loc = op.getLoc(); 64 Value result = rewriter.create<arith::ConstantOp>( 65 loc, resultType, rewriter.getZeroAttr(resultType)); 66 for (auto position : *unrollIterator) { 67 Value extractLhs = rewriter.create<ExtractOp>(loc, op.getLhs(), position); 68 Value extractRhs = rewriter.create<ExtractOp>(loc, op.getRhs(), position); 69 Value interleave = 70 rewriter.create<InterleaveOp>(loc, extractLhs, extractRhs); 71 result = rewriter.create<InsertOp>(loc, interleave, result, position); 72 } 73 74 rewriter.replaceOp(op, result); 75 return success(); 76 } 77 78 private: 79 int64_t targetRank = 1; 80 }; 81 82 /// Rewrite vector.interleave op into an equivalent vector.shuffle op, when 83 /// applicable: `sourceType` must be 1D and non-scalable. 84 /// 85 /// Example: 86 /// 87 /// ```mlir 88 /// vector.interleave %a, %b : vector<7xi16> -> vector<14xi16> 89 /// ``` 90 /// 91 /// Is rewritten into: 92 /// 93 /// ```mlir 94 /// vector.shuffle %arg0, %arg1 [0, 7, 1, 8, 2, 9, 3, 10, 4, 11, 5, 12, 6, 13] 95 /// : vector<7xi16>, vector<7xi16> 96 /// ``` 97 struct InterleaveToShuffle final : OpRewritePattern<vector::InterleaveOp> { 98 using OpRewritePattern::OpRewritePattern; 99 100 LogicalResult matchAndRewrite(vector::InterleaveOp op, 101 PatternRewriter &rewriter) const override { 102 VectorType sourceType = op.getSourceVectorType(); 103 if (sourceType.getRank() != 1 || sourceType.isScalable()) { 104 return failure(); 105 } 106 int64_t n = sourceType.getNumElements(); 107 auto seq = llvm::seq<int64_t>(2 * n); 108 auto zip = llvm::to_vector(llvm::map_range( 109 seq, [n](int64_t i) { return (i % 2 ? n : 0) + i / 2; })); 110 rewriter.replaceOpWithNewOp<ShuffleOp>(op, op.getLhs(), op.getRhs(), zip); 111 return success(); 112 } 113 }; 114 115 } // namespace 116 117 void mlir::vector::populateVectorInterleaveLoweringPatterns( 118 RewritePatternSet &patterns, int64_t targetRank, PatternBenefit benefit) { 119 patterns.add<UnrollInterleaveOp>(targetRank, patterns.getContext(), benefit); 120 } 121 122 void mlir::vector::populateVectorInterleaveToShufflePatterns( 123 RewritePatternSet &patterns, PatternBenefit benefit) { 124 patterns.add<InterleaveToShuffle>(patterns.getContext(), benefit); 125 } 126