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> 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 : vector<4xi64> | all leading positions 43 /// %3 = vector.insert %2, %result [0, 0, 0] | 44 /// : vector<8xi64> into vector<1x2x3x8xi64> ┘ 45 /// ``` 46 /// 47 /// Note: If any leading dimension before the `targetRank` is scalable the 48 /// unrolling will stop before the scalable dimension. 49 class UnrollInterleaveOp : public OpRewritePattern<vector::InterleaveOp> { 50 public: 51 UnrollInterleaveOp(int64_t targetRank, MLIRContext *context, 52 PatternBenefit benefit = 1) 53 : OpRewritePattern(context, benefit), targetRank(targetRank){}; 54 55 LogicalResult matchAndRewrite(vector::InterleaveOp op, 56 PatternRewriter &rewriter) const override { 57 VectorType resultType = op.getResultVectorType(); 58 auto unrollIterator = vector::createUnrollIterator(resultType, targetRank); 59 if (!unrollIterator) 60 return failure(); 61 62 auto loc = op.getLoc(); 63 Value result = rewriter.create<arith::ConstantOp>( 64 loc, resultType, rewriter.getZeroAttr(resultType)); 65 for (auto position : *unrollIterator) { 66 Value extractLhs = rewriter.create<ExtractOp>(loc, op.getLhs(), position); 67 Value extractRhs = rewriter.create<ExtractOp>(loc, op.getRhs(), position); 68 Value interleave = 69 rewriter.create<InterleaveOp>(loc, extractLhs, extractRhs); 70 result = rewriter.create<InsertOp>(loc, interleave, result, position); 71 } 72 73 rewriter.replaceOp(op, result); 74 return success(); 75 } 76 77 private: 78 int64_t targetRank = 1; 79 }; 80 81 /// Rewrite vector.interleave op into an equivalent vector.shuffle op, when 82 /// applicable: `sourceType` must be 1D and non-scalable. 83 /// 84 /// Example: 85 /// 86 /// ```mlir 87 /// vector.interleave %a, %b : vector<7xi16> 88 /// ``` 89 /// 90 /// Is rewritten into: 91 /// 92 /// ```mlir 93 /// vector.shuffle %arg0, %arg1 [0, 7, 1, 8, 2, 9, 3, 10, 4, 11, 5, 12, 6, 13] 94 /// : vector<7xi16>, vector<7xi16> 95 /// ``` 96 class InterleaveToShuffle : public OpRewritePattern<vector::InterleaveOp> { 97 public: 98 InterleaveToShuffle(MLIRContext *context, PatternBenefit benefit = 1) 99 : OpRewritePattern(context, benefit) {}; 100 101 LogicalResult matchAndRewrite(vector::InterleaveOp op, 102 PatternRewriter &rewriter) const override { 103 VectorType sourceType = op.getSourceVectorType(); 104 if (sourceType.getRank() != 1 || sourceType.isScalable()) { 105 return failure(); 106 } 107 int64_t n = sourceType.getNumElements(); 108 auto seq = llvm::seq<int64_t>(2 * n); 109 auto zip = llvm::to_vector(llvm::map_range( 110 seq, [n](int64_t i) { return (i % 2 ? n : 0) + i / 2; })); 111 rewriter.replaceOpWithNewOp<ShuffleOp>(op, op.getLhs(), op.getRhs(), zip); 112 return success(); 113 } 114 }; 115 116 } // namespace 117 118 void mlir::vector::populateVectorInterleaveLoweringPatterns( 119 RewritePatternSet &patterns, int64_t targetRank, PatternBenefit benefit) { 120 patterns.add<UnrollInterleaveOp>(targetRank, patterns.getContext(), benefit); 121 } 122 123 void mlir::vector::populateVectorInterleaveToShufflePatterns( 124 RewritePatternSet &patterns, PatternBenefit benefit) { 125 patterns.add<InterleaveToShuffle>(patterns.getContext(), benefit); 126 } 127