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
39using namespace mlir;
40using namespace mlir::vector;
41
42namespace {
43/// Progressive lowering of BroadcastOp.
44class BroadcastOpLowering : public OpRewritePattern<vector::BroadcastOp> {
45public:
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 = dyn_cast<VectorType>(op.getSourceType());
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 = VectorType::Builder(dstType).dropDim(0);
88 Value bcst =
89 rewriter.create<vector::BroadcastOp>(loc, resType, op.getSource());
90 Value result = rewriter.create<arith::ConstantOp>(
91 loc, dstType, rewriter.getZeroAttr(dstType));
92 for (int64_t d = 0, dim = dstType.getDimSize(0); d < dim; ++d)
93 result = rewriter.create<vector::InsertOp>(loc, bcst, result, d);
94 rewriter.replaceOp(op, result);
95 return success();
96 }
97
98 // Find non-matching dimension, if any.
99 assert(srcRank == dstRank);
100 int64_t m = -1;
101 for (int64_t r = 0; r < dstRank; r++)
102 if (srcType.getDimSize(r) != dstType.getDimSize(r)) {
103 m = r;
104 break;
105 }
106
107 // All trailing dimensions are the same. Simply pass through.
108 if (m == -1) {
109 rewriter.replaceOp(op, op.getSource());
110 return success();
111 }
112
113 // Any non-matching dimension forces a stretch along this rank.
114 // For example:
115 // %x = broadcast %y : vector<4x1x2xf32> to vector<4x2x2xf32>
116 // becomes:
117 // %a = broadcast %y[0] : vector<1x2xf32> to vector<2x2xf32>
118 // %b = broadcast %y[1] : vector<1x2xf32> to vector<2x2xf32>
119 // %c = broadcast %y[2] : vector<1x2xf32> to vector<2x2xf32>
120 // %d = broadcast %y[3] : vector<1x2xf32> to vector<2x2xf32>
121 // %x = [%a,%b,%c,%d]
122 // becomes:
123 // %u = broadcast %y[0][0] : vector<2xf32> to vector <2x2xf32>
124 // %v = broadcast %y[1][0] : vector<2xf32> to vector <2x2xf32>
125 // %a = [%u, %v]
126 // ..
127 // %x = [%a,%b,%c,%d]
128 VectorType resType =
129 VectorType::get(dstType.getShape().drop_front(), eltType);
130 Value result = rewriter.create<arith::ConstantOp>(
131 loc, dstType, rewriter.getZeroAttr(dstType));
132 if (m == 0) {
133 // Stetch at start.
134 Value ext = rewriter.create<vector::ExtractOp>(loc, op.getSource(), 0);
135 Value bcst = rewriter.create<vector::BroadcastOp>(loc, resType, ext);
136 for (int64_t d = 0, dim = dstType.getDimSize(0); d < dim; ++d)
137 result = rewriter.create<vector::InsertOp>(loc, bcst, result, d);
138 } else {
139 // Stetch not at start.
140 for (int64_t d = 0, dim = dstType.getDimSize(0); d < dim; ++d) {
141 Value ext = rewriter.create<vector::ExtractOp>(loc, op.getSource(), d);
142 Value bcst = rewriter.create<vector::BroadcastOp>(loc, resType, ext);
143 result = rewriter.create<vector::InsertOp>(loc, bcst, result, d);
144 }
145 }
146 rewriter.replaceOp(op, result);
147 return success();
148 }
149};
150} // namespace
151
152void mlir::vector::populateVectorBroadcastLoweringPatterns(
153 RewritePatternSet &patterns, PatternBenefit benefit) {
154 patterns.add<BroadcastOpLowering>(arg: patterns.getContext(), args&: benefit);
155}
156

source code of mlir/lib/Dialect/Vector/Transforms/LowerVectorBroadcast.cpp