MLIR  19.0.0git
LowerVectorBroadcast.cpp
Go to the documentation of this file.
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 
28 #include "mlir/IR/BuiltinTypes.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"
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:
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 
153  RewritePatternSet &patterns, PatternBenefit benefit) {
154  patterns.add<BroadcastOpLowering>(patterns.getContext(), benefit);
155 }
TypedAttr getZeroAttr(Type type)
Definition: Builders.cpp:331
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
Definition: Builders.cpp:464
Location getLoc()
The source location the operation was defined or derived from.
Definition: Operation.h:223
This class represents the benefit of a pattern match in a unitless scheme that ranges from 0 (very li...
Definition: PatternMatch.h:34
A special type of RewriterBase that coordinates the application of a rewrite pattern on the current I...
Definition: PatternMatch.h:785
MLIRContext * getContext() const
Definition: PatternMatch.h:823
RewritePatternSet & add(ConstructorArg &&arg, ConstructorArgs &&...args)
Add an instance of each of the pattern types 'Ts' to the pattern list with the given arguments.
Definition: PatternMatch.h:847
virtual void replaceOp(Operation *op, ValueRange newValues)
Replace the results of the given (original) operation with the specified list of values (replacements...
OpTy replaceOpWithNewOp(Operation *op, Args &&...args)
Replace the results of the given (original) op with a new op that is created without verification (re...
Definition: PatternMatch.h:536
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
Definition: Types.h:74
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:96
This is a builder type that keeps local references to arguments.
Definition: BuiltinTypes.h:305
Builder & dropDim(unsigned pos)
Erase a dim from shape @pos.
Definition: BuiltinTypes.h:330
void populateVectorBroadcastLoweringPatterns(RewritePatternSet &patterns, PatternBenefit benefit=1)
Populate the pattern set with the following patterns:
Include the generated interface declarations.
LogicalResult success(bool isSuccess=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:56
auto get(MLIRContext *context, Ts &&...params)
Helper method that injects context only if needed, this helps unify some of the attribute constructio...
This class represents an efficient way to signal success or failure.
Definition: LogicalResult.h:26
OpRewritePattern is a wrapper around RewritePattern that allows for matching and rewriting against an...
Definition: PatternMatch.h:358
OpRewritePattern(MLIRContext *context, PatternBenefit benefit=1, ArrayRef< StringRef > generatedNames={})
Patterns must specify the root operation name they match against, and can also specify the benefit of...
Definition: PatternMatch.h:362