1 //===- LowerVectorBroadcast.cpp - Lower 'vector.broadcast' 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.broadcast' operation. 11 // 12 //===----------------------------------------------------------------------===// 13 14 #include "mlir/Dialect/Affine/IR/AffineOps.h" 15 #include "mlir/Dialect/Arith/IR/Arith.h" 16 #include "mlir/Dialect/Arith/Utils/Utils.h" 17 #include "mlir/Dialect/Linalg/IR/Linalg.h" 18 #include "mlir/Dialect/MemRef/IR/MemRef.h" 19 #include "mlir/Dialect/SCF/IR/SCF.h" 20 #include "mlir/Dialect/Tensor/IR/Tensor.h" 21 #include "mlir/Dialect/Utils/IndexingUtils.h" 22 #include "mlir/Dialect/Utils/StructuredOpsUtils.h" 23 #include "mlir/Dialect/Vector/IR/VectorOps.h" 24 #include "mlir/Dialect/Vector/Transforms/LoweringPatterns.h" 25 #include "mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h" 26 #include "mlir/Dialect/Vector/Utils/VectorUtils.h" 27 #include "mlir/IR/BuiltinAttributeInterfaces.h" 28 #include "mlir/IR/BuiltinTypes.h" 29 #include "mlir/IR/ImplicitLocOpBuilder.h" 30 #include "mlir/IR/Location.h" 31 #include "mlir/IR/Matchers.h" 32 #include "mlir/IR/PatternMatch.h" 33 #include "mlir/IR/TypeUtilities.h" 34 #include "mlir/Interfaces/VectorInterfaces.h" 35 #include "mlir/Support/LogicalResult.h" 36 37 #define DEBUG_TYPE "vector-broadcast-lowering" 38 39 using namespace mlir; 40 using namespace mlir::vector; 41 42 namespace { 43 /// Progressive lowering of BroadcastOp. 44 class BroadcastOpLowering : public OpRewritePattern<vector::BroadcastOp> { 45 public: 46 using OpRewritePattern::OpRewritePattern; 47 48 LogicalResult matchAndRewrite(vector::BroadcastOp op, 49 PatternRewriter &rewriter) const override { 50 auto loc = op.getLoc(); 51 VectorType dstType = op.getResultVectorType(); 52 VectorType srcType = op.getSourceType().dyn_cast<VectorType>(); 53 Type eltType = dstType.getElementType(); 54 55 // Scalar to any vector can use splat. 56 if (!srcType) { 57 rewriter.replaceOpWithNewOp<vector::SplatOp>(op, dstType, op.getSource()); 58 return success(); 59 } 60 61 // Determine rank of source and destination. 62 int64_t srcRank = srcType.getRank(); 63 int64_t dstRank = dstType.getRank(); 64 65 // Stretching scalar inside vector (e.g. vector<1xf32>) can use splat. 66 if (srcRank <= 1 && dstRank == 1) { 67 Value ext; 68 if (srcRank == 0) 69 ext = rewriter.create<vector::ExtractElementOp>(loc, op.getSource()); 70 else 71 ext = rewriter.create<vector::ExtractOp>(loc, op.getSource(), 0); 72 rewriter.replaceOpWithNewOp<vector::SplatOp>(op, dstType, ext); 73 return success(); 74 } 75 76 // Duplicate this rank. 77 // For example: 78 // %x = broadcast %y : k-D to n-D, k < n 79 // becomes: 80 // %b = broadcast %y : k-D to (n-1)-D 81 // %x = [%b,%b,%b,%b] : n-D 82 // becomes: 83 // %b = [%y,%y] : (n-1)-D 84 // %x = [%b,%b,%b,%b] : n-D 85 if (srcRank < dstRank) { 86 // Duplication. 87 VectorType resType = 88 VectorType::get(dstType.getShape().drop_front(), eltType); 89 Value bcst = 90 rewriter.create<vector::BroadcastOp>(loc, resType, op.getSource()); 91 Value result = rewriter.create<arith::ConstantOp>( 92 loc, dstType, rewriter.getZeroAttr(dstType)); 93 for (int64_t d = 0, dim = dstType.getDimSize(0); d < dim; ++d) 94 result = rewriter.create<vector::InsertOp>(loc, bcst, result, d); 95 rewriter.replaceOp(op, result); 96 return success(); 97 } 98 99 // Find non-matching dimension, if any. 100 assert(srcRank == dstRank); 101 int64_t m = -1; 102 for (int64_t r = 0; r < dstRank; r++) 103 if (srcType.getDimSize(r) != dstType.getDimSize(r)) { 104 m = r; 105 break; 106 } 107 108 // All trailing dimensions are the same. Simply pass through. 109 if (m == -1) { 110 rewriter.replaceOp(op, op.getSource()); 111 return success(); 112 } 113 114 // Any non-matching dimension forces a stretch along this rank. 115 // For example: 116 // %x = broadcast %y : vector<4x1x2xf32> to vector<4x2x2xf32> 117 // becomes: 118 // %a = broadcast %y[0] : vector<1x2xf32> to vector<2x2xf32> 119 // %b = broadcast %y[1] : vector<1x2xf32> to vector<2x2xf32> 120 // %c = broadcast %y[2] : vector<1x2xf32> to vector<2x2xf32> 121 // %d = broadcast %y[3] : vector<1x2xf32> to vector<2x2xf32> 122 // %x = [%a,%b,%c,%d] 123 // becomes: 124 // %u = broadcast %y[0][0] : vector<2xf32> to vector <2x2xf32> 125 // %v = broadcast %y[1][0] : vector<2xf32> to vector <2x2xf32> 126 // %a = [%u, %v] 127 // .. 128 // %x = [%a,%b,%c,%d] 129 VectorType resType = 130 VectorType::get(dstType.getShape().drop_front(), eltType); 131 Value result = rewriter.create<arith::ConstantOp>( 132 loc, dstType, rewriter.getZeroAttr(dstType)); 133 if (m == 0) { 134 // Stetch at start. 135 Value ext = rewriter.create<vector::ExtractOp>(loc, op.getSource(), 0); 136 Value bcst = rewriter.create<vector::BroadcastOp>(loc, resType, ext); 137 for (int64_t d = 0, dim = dstType.getDimSize(0); d < dim; ++d) 138 result = rewriter.create<vector::InsertOp>(loc, bcst, result, d); 139 } else { 140 // Stetch not at start. 141 for (int64_t d = 0, dim = dstType.getDimSize(0); d < dim; ++d) { 142 Value ext = rewriter.create<vector::ExtractOp>(loc, op.getSource(), d); 143 Value bcst = rewriter.create<vector::BroadcastOp>(loc, resType, ext); 144 result = rewriter.create<vector::InsertOp>(loc, bcst, result, d); 145 } 146 } 147 rewriter.replaceOp(op, result); 148 return success(); 149 } 150 }; 151 } // namespace 152 153 void mlir::vector::populateVectorBroadcastLoweringPatterns( 154 RewritePatternSet &patterns, PatternBenefit benefit) { 155 patterns.add<BroadcastOpLowering>(patterns.getContext(), benefit); 156 } 157