25 #define DEBUG_TYPE "vector-broadcast-lowering"
36 LogicalResult matchAndRewrite(vector::BroadcastOp op,
38 auto loc = op.getLoc();
39 VectorType dstType = op.getResultVectorType();
40 VectorType srcType = dyn_cast<VectorType>(op.getSourceType());
41 Type eltType = dstType.getElementType();
50 int64_t srcRank = srcType.getRank();
51 int64_t dstRank = dstType.getRank();
54 if (srcRank <= 1 && dstRank == 1) {
55 Value ext = rewriter.
create<vector::ExtractOp>(loc, op.getSource());
69 if (srcRank < dstRank) {
73 rewriter.
create<vector::BroadcastOp>(loc, resType, op.getSource());
74 Value result = rewriter.
create<ub::PoisonOp>(loc, dstType);
75 for (int64_t d = 0, dim = dstType.getDimSize(0); d < dim; ++d)
76 result = rewriter.
create<vector::InsertOp>(loc, bcst, result, d);
82 assert(srcRank == dstRank);
84 for (int64_t r = 0; r < dstRank; r++)
85 if (srcType.getDimSize(r) != dstType.getDimSize(r)) {
113 dstType.getScalableDims().drop_front());
114 Value result = rewriter.
create<ub::PoisonOp>(loc, dstType);
117 Value ext = rewriter.
create<vector::ExtractOp>(loc, op.getSource(), 0);
118 Value bcst = rewriter.
create<vector::BroadcastOp>(loc, resType, ext);
119 for (int64_t d = 0, dim = dstType.getDimSize(0); d < dim; ++d)
120 result = rewriter.
create<vector::InsertOp>(loc, bcst, result, d);
123 if (dstType.getScalableDims()[0]) {
127 for (int64_t d = 0, dim = dstType.getDimSize(0); d < dim; ++d) {
128 Value ext = rewriter.
create<vector::ExtractOp>(loc, op.getSource(), d);
129 Value bcst = rewriter.
create<vector::BroadcastOp>(loc, resType, ext);
130 result = rewriter.
create<vector::InsertOp>(loc, bcst, result, d);
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
This class represents the benefit of a pattern match in a unitless scheme that ranges from 0 (very li...
A special type of RewriterBase that coordinates the application of a rewrite pattern on the current I...
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...
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
This is a builder type that keeps local references to arguments.
Builder & dropDim(unsigned pos)
Erase a dim from shape @pos.
void populateVectorBroadcastLoweringPatterns(RewritePatternSet &patterns, PatternBenefit benefit=1)
Populate the pattern set with the following patterns:
Include the generated interface declarations.
const FrozenRewritePatternSet & patterns
auto get(MLIRContext *context, Ts &&...params)
Helper method that injects context only if needed, this helps unify some of the attribute constructio...
OpRewritePattern is a wrapper around RewritePattern that allows for matching and rewriting against an...
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...