[mlir][Vector] Add pattern to reorder elementwise and broadcast ops

[awarzynski]

Summary

The new pattern will replace elementwise(broadcast) with
broadcast(elementwise) when safe.

[kuhar]

mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp

906

Member casts are deprecated, use free cast functions: https://mlir.llvm.org/deprecation/

927–933

This is the third loop that inspects operands and casts them to broadcasts. Why not do the checks and collect all the broadcast elements in just one loop?

935

nit: Type vectorTy = *op->result_type_begin(); or op->getResultTypes()[0];?

1366–1367

Why did you decide to add the new pattern here? Is it because the other patterns tend to create code that needs this cleanup around elementwise operations on broadcasts?

[awarzynski]

Thanks for taking a look, I will send an update shortly.

mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp

906

Thanks!

927–933

Good point, though I would still seperate "checking" from "processing".

935

Thanks, that's neat!

Btw, for consistency with the rest of this file, I will stick with vectorType (instead of using vectorTy).

1366–1367

Why did you decide to add the new pattern here?

ReorderElementwiseOpsOnTranspose is the closest to what I am trying to do. But perhaps a new category of patterns would be better 🤔 . Any suggestions?

Is it because the other patterns tend to create code that needs this cleanup around elementwise operations on broadcasts?

Not really. This is fairly generic.

[awarzynski]

Address comments from Jakub

[dcaballe]

Thanks a lot for working on this! I think this will have an impact beyond the nD tensor extract! A few comments.

mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp

907

We should also check for some properties of the op. There is a trait called ElementWiseMappable something or similar that we can use here.
In this regard, I would add a test with a vector.contract where all the inputs are broadcasted from a scalar... We do not want this transformation to trigger on this one.

mlir/test/Dialect/Vector/vector-reduce-to-contract.mlir

403 ↗

(On Diff #530907)

• Move this to a different file? Maybe to a dedicated one?
• Add test with vector1D to vector 2D broadcast?
• Add test with vector.contract of broadcasted scalar operands (it shouldn't trigger)?
• Any other negative tests?

[kuhar]

mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp

912–915

nit: Create a local variable for the broadcast so that we don't have to cast twice.

1366–1367

I'm not very familiar with patterns in this file, but this strikes me like a very specific choice, while this patterns seems like a more general cleanup that someone may want to run without all the other reduction to contract patterns. If this is intentional, could you explain this in the commit description, and/or mention that we may want a better 'home' populate function for this pattern?

[dcaballe]

mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp

1366–1367

Sorry, I had missed this. Yes, I think it makes sense to create a different populateVectorBroadcast... to get more finer-grain control on this pattern. I can't find any other populate where this would make sense.

[awarzynski]

Addressing comments from Diego an Jakub:

• added more tests and moved to a dedicated file
• created a dedicated populateVector(...)Patterns
• added support for broadcasts of shaped types (also added a test to demonstrate)

[kuhar]

mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp

922

You don't need the isa now, just return bcast && bcast.get..;

mlir/test/Dialect/Vector/vector-remove-broadcast.mlir

29–30 ↗

(On Diff #531318)

Can you also add some test(s) that demonstrate that your precondition checks work? For instance the one that checks that *all* arguments should be broadcasts

[awarzynski]

mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp

922

Argh, I missed that, Thank you :)

mlir/test/Dialect/Vector/vector-remove-broadcast.mlir

29–30 ↗

(On Diff #531318)

one that checks that *all* arguments should be broadcasts

Thanks for the suggestion! I'm actually struggling to find other good examples. I could add "vector.broadcast + vector.extractelement" (it's not ElementwiseMappable), but it doesn't feel right. And everything in Arith is ElementwiseMappable (most of the example that I can think of are a mix of Arith + Vector). Happy to add more if you have more ideas.

[awarzynski]

Add a test, remove redundant llvm::isa

[kuhar]

LGTM, thanks for the changes. You may want to wait for a thumbs up from @dcaballe before submitting.

[dcaballe]

Thanks for addressing the feedback. LGTM! Just some minor comments that you can address before submitting.

mlir/test/Dialect/Linalg/vectorize-tensor-extract.mlir

321–322

Instead of removing the DAG here, could we reduce this test to match only what it's important? That would help maintainability

mlir/test/Dialect/Vector/vector-remove-broadcast.mlir

1 ↗

(On Diff #531345)

personal opinion, feel free to ignore: maybe "sink-vector-broadcast" is more accurate?

62–64 ↗

(On Diff #531345)

good to add a -DAG here

mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp

377

personal opinion, feel free to ignore: ->SinkVectorBroadcast?

[awarzynski]

Thank you both for reviewing and for helping me with this!

@dcaballe , I will address your comments in the final version that I am about to merge.

mlir/test/Dialect/Linalg/vectorize-tensor-extract.mlir

321–322

Agreed. I've tried to trim it a couple of times, but always get stuck. I will just remove pretty much everything apart from the scalar and the contiguous loads, which are the key bits in this test.

mlir/test/Dialect/Vector/vector-remove-broadcast.mlir

1 ↗

(On Diff #531345)

Naming is hard, but gets better with experience. That's something that you have tons of, so will happily incorporate your suggestion :)

mlir/include/mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h

 
 /// Collect a set of patterns to reduce the rank of the operands of vector
 /// transfer ops to operate on the largest contigious vector.
 /// These patterns are useful when lowering to dialects with 1d vector type
 /// such as llvm and it will result fewer memory reads.
 void populateVectorTransferCollapseInnerMostContiguousDimsPatterns(
     RewritePatternSet &patterns, PatternBenefit benefit = 1);
 
+/// Patterns that remove redundant vector broadcasts.
+void populateSinkVectorBroadcastPatterns(RewritePatternSet &patterns,
+                                          PatternBenefit benefit = 1);
+
 /// Populate `patterns` with the following patterns.
 ///
 /// [DecomposeDifferentRankInsertStridedSlice]
 /// ==========================================
 /// RewritePattern for InsertStridedSliceOp where source and destination vectors
 /// have different ranks.
 ///
 /// When ranks are different, InsertStridedSlice needs to extract a properly

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

transform::VectorizeOp::applyToOne(Operation *target,
   patterns.add<VectorizationPattern>(ctx, getVectorizeNdExtract());
 
   if (!getDisableTransferPermutationMapLoweringPatterns())
     vector::populateVectorTransferPermutationMapLoweringPatterns(patterns);
 
   if (!getDisableMultiReductionToContractPatterns())
     vector::populateVectorReductionToContractPatterns(patterns);
 
+  vector::populateSinkVectorBroadcastPatterns(patterns);
+
   patterns.add<linalg::LinalgCopyVTRForwardingPattern,
                linalg::LinalgCopyVTWForwardingPattern>(ctx,
                                                        /*benefit=*/2);
   vector::TransferReadOp::getCanonicalizationPatterns(patterns, ctx);
   vector::TransferWriteOp::getCanonicalizationPatterns(patterns, ctx);
   vector::populateVectorTransferTensorSliceTransforms(patterns);
 
   patterns.add<CopyVectorizationPattern>(ctx);

mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp

LogicalResult matchAndRewrite(vector::BitCastOp bitcastOp,
     rewriter.replaceOp(bitcastOp, res);
     return success();
   }
 
 private:
   std::function<bool(BitCastOp)> controlFn;
 };
 
+/// Reorders elementwise(broadcast) to broadcast(elementwise). Ex:
+/// ```
+/// %a = vector.broadcast %arg1 : index to vector<1x4xindex>
+/// %b = vector.broadcast %arg2 : index to vector<1x4xindex>
+/// %r = arith.addi %a, %b : vector<1x4xindex>
+/// ```
+/// Gets converted to:
+/// ```
+/// %r = arith.addi %arg0, %arg1 : index
+/// %b = vector.broadcast %r : index to vector<1x4xindex>
+/// ```
+struct ReorderElementwiseOpsOnBroadcast final
+    : public OpTraitRewritePattern<OpTrait::Elementwise> {
+  using OpTraitRewritePattern::OpTraitRewritePattern;
+  LogicalResult matchAndRewrite(Operation *op,
+                                PatternRewriter &rewriter) const override {
+    if (op->getNumResults() != 1)
+      return failure();
+    if (!llvm::isa<ShapedType>(op->getResults()[0].getType()))
+      return failure();
+    if (!OpTrait::hasElementwiseMappableTraits(op))
+      return failure();
+
+    // Get the type of the first operand
+    auto firstBcast = op->getOperand(0).getDefiningOp<vector::BroadcastOp>();
+    if (!firstBcast)
+      return failure();
+    auto firstOpType = firstBcast.getOperand().getType();
+
+    // Make sure that operands are "broadcast"ed from identical (scalar or
+    // vector) types. That indicates that it's safe to skip the broadcasting of
+    // operands.
+    if (!llvm::all_of(op->getOperands(), [&firstOpType](Value val) {
+          auto bcast = val.getDefiningOp<vector::BroadcastOp>();
+          return (bcast && (bcast.getOperand().getType() == firstOpType));
+        })) {
+      return failure();
+    }
+
+    // Collect the source values
+    SmallVector<Value> srcValues;
+    srcValues.reserve(op->getNumOperands());
+
+    for (Value operand : op->getOperands()) {
+      srcValues.push_back(
+          operand.getDefiningOp<vector::BroadcastOp>().getOperand());
+    }
+
+    Operation *elementwiseOp =
+        rewriter.create(op->getLoc(), op->getName().getIdentifier(), srcValues,
+                        firstOpType, op->getAttrs());
+
+    auto vectorType = op->getResultTypes()[0];
+    rewriter.replaceOpWithNewOp<vector::BroadcastOp>(
+        op, vectorType, elementwiseOp->getResults());
+
+    return success();
+  }
+};
+
 // Helper that returns a vector comparison that constructs a mask:
 //     mask = [0,1,..,n-1] + [o,o,..,o] < [b,b,..,b]
 //
 // If `dim == 0` then the result will be a 0-D vector.
 //
 // NOTE: The LLVM::GetActiveLaneMaskOp intrinsic would provide an alternative,
 //       much more compact, IR for this operation, but LLVM eventually
 //       generates more elaborate instructions for this intrinsic since it

 }
 
 void mlir::vector::populateVectorReductionToContractPatterns(
     RewritePatternSet &patterns, PatternBenefit benefit) {
   patterns.add<MultiReduceToContract, CombineContractBroadcast,
                CombineContractABTranspose, CombineContractResultTranspose,
                ReorderCastOpsOnBroadcast, ReorderElementwiseOpsOnTranspose>(
       patterns.getContext(), benefit);
 }
 
 void mlir::vector::
     populateVectorTransferCollapseInnerMostContiguousDimsPatterns(
         RewritePatternSet &patterns, PatternBenefit benefit) {
   patterns.add<DropInnerMostUnitDims>(patterns.getContext(), benefit);
 }
 
+void mlir::vector::populateSinkVectorBroadcastPatterns(
+    RewritePatternSet &patterns, PatternBenefit benefit) {
+  patterns.add<ReorderElementwiseOpsOnBroadcast>(patterns.getContext(),
+                                                 benefit);
+}
+
 //===----------------------------------------------------------------------===//
 // TableGen'd enum attribute definitions
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Vector/Transforms/VectorTransformsEnums.cpp.inc"

mlir/test/Dialect/Linalg/vectorize-tensor-extract.mlir

^bb0(%out: f32):
     %30 = arith.addi %arg1, %29 : index
     %31 = arith.addi %30, %arg4 : index
     %extracted = tensor.extract %6[%28, %c79, %31] : tensor<45x80x16xf32>
     linalg.yield %extracted : f32
   } -> tensor<1x4xf32>
   return %25 : tensor<1x4xf32>
 }
 
-// CHECK-LABEL:   func.func @vectorize_nd_tensor_extract_transfer_read_complex
+
+// CHECK-LABEL:   func.func @vectorize_nd_tensor_extract_transfer_read_complex(
 // CHECK-SAME:      %[[VAL_0:.*]]: tensor<45x80x16xf32>,
-// CHECK-SAME:      {{.*}}: index,
+// CHECK-SAME:      %[[VAL_1:.*]]: index, %[[VAL_2:.*]]: index, %[[VAL_3:.*]]: index, %[[VAL_4:.*]]: index,
 // CHECK-SAME:      %[[VAL_5:.*]]: tensor<1x4xf32>) -> tensor<1x4xf32> {
 // CHECK:           %[[VAL_6:.*]] = arith.constant dense<[0, 1, 2, 3]> : vector<4xindex>
 // CHECK:           %[[VAL_7:.*]] = arith.constant 0 : i32
 // CHECK:           %[[VAL_8:.*]] = arith.constant 0.000000e+00 : f32
 // CHECK:           %[[VAL_9:.*]] = arith.constant 0 : index
 // CHECK:           %[[VAL_10:.*]] = arith.constant 79 : index
-// CHECK:           %[[VAL_11:.*]] = vector.broadcast %{{.*}} : index to vector<1x4xindex>
-// CHECK:           %[[VAL_12:.*]] = vector.broadcast %{{.*}} : index to vector<1x4xindex>
-// CHECK:           %[[VAL_13:.*]] = arith.addi %[[VAL_11]], %[[VAL_12]] : vector<1x4xindex>
-// CHECK:           %[[VAL_14:.*]] = vector.broadcast %{{.*}} : index to vector<4xindex>
-// CHECK:           %[[VAL_15:.*]] = arith.addi %[[VAL_14]], %[[VAL_6]] : vector<4xindex>
-// CHECK:           %[[VAL_16:.*]] = vector.broadcast %{{.*}} : index to vector<4xindex>
-// CHECK:           %[[VAL_17:.*]] = arith.addi %[[VAL_15]], %[[VAL_16]] : vector<4xindex>
-// CHECK:           %[[VAL_18:.*]] = vector.shape_cast %[[VAL_13]] : vector<1x4xindex> to vector<4xindex>
-// CHECK:           %[[VAL_19:.*]] = vector.extractelement %[[VAL_18]]{{\[}}%[[VAL_7]] : i32] : vector<4xindex>
-// CHECK:           %[[VAL_20:.*]] = vector.extractelement %[[VAL_17]]{{\[}}%[[VAL_7]] : i32] : vector<4xindex>
-// CHECK:           %[[VAL_21:.*]] = vector.transfer_read %[[VAL_0]]{{\[}}%[[VAL_19]], %[[VAL_10]], %[[VAL_20]]], %[[VAL_8]] {in_bounds = [true, true]} : tensor<45x80x16xf32>, vector<1x4xf32>
-// CHECK:           %[[VAL_22:.*]] = vector.transfer_write %[[VAL_21]], %[[VAL_5]]{{\[}}%[[VAL_9]], %[[VAL_9]]] {in_bounds = [true, true]} : vector<1x4xf32>, tensor<1x4xf32>
+// CHECK:           %[[VAL_11:.*]] = arith.addi %[[VAL_1]], %[[VAL_2]] : index
+// CHECK:           %[[VAL_12:.*]] = vector.broadcast %[[VAL_11]] : index to vector<1x4xindex>
+// CHECK:           %[[VAL_13:.*]] = vector.broadcast %[[VAL_3]] : index to vector<4xindex>
+// CHECK:           %[[VAL_14:.*]] = arith.addi %[[VAL_13]], %[[VAL_6]] : vector<4xindex>
+// CHECK:           %[[VAL_15:.*]] = vector.broadcast %[[VAL_4]] : index to vector<4xindex>
+// CHECK:           %[[VAL_16:.*]] = arith.addi %[[VAL_14]], %[[VAL_15]] : vector<4xindex>
+// CHECK:           %[[VAL_17:.*]] = vector.shape_cast %[[VAL_12]] : vector<1x4xindex> to vector<4xindex>
+// CHECK:           %[[VAL_18:.*]] = vector.extractelement %[[VAL_17]]{{\[}}%[[VAL_7]] : i32] : vector<4xindex>
+// CHECK:           %[[VAL_19:.*]] = vector.extractelement %[[VAL_16]]{{\[}}%[[VAL_7]] : i32] : vector<4xindex>
+// CHECK:           %[[VAL_20:.*]] = vector.transfer_read %[[VAL_0]]{{\[}}%[[VAL_18]], %[[VAL_10]], %[[VAL_19]]], %[[VAL_8]] {in_bounds = [true, true]} : tensor<45x80x16xf32>, vector<1x4xf32>
+// CHECK:           %[[VAL_21:.*]] = vector.transfer_write %[[VAL_20]], %[[VAL_5]]{{\[}}%[[VAL_9]], %[[VAL_9]]] {in_bounds = [true, true]} : vector<1x4xf32>, tensor<1x4xf32>
+// CHECK:           return %[[VAL_21]] : tensor<1x4xf32>
+// CHECK:         }
 
 transform.sequence failures(propagate) {
  ^bb1(%arg1: !transform.any_op):
    %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op
    %1 = get_closest_isolated_parent %0 : (!transform.any_op) -> !transform.any_op
    %2 = transform.structured.vectorize %1 { vectorize_nd_extract } : (!transform.any_op) -> !transform.any_op
  }
 
^bb0(%out: f32):
     %19 = arith.maxsi %18, %c0 : index
     %20 = arith.minsi %19, %c256 : index
     %extracted_1 = tensor.extract %input_2[%20, %17] : tensor<257x24xf32>
     linalg.yield %extracted_1 : f32
   } -> tensor<1x1x4xf32>
   return %1 : tensor<1x1x4xf32>
 }
 
 // CHECK-LABEL:   func.func @vectorize_nd_tensor_extract_with_tensor_extract(
-// CHECK-SAME:    %[[VAL_0:.*]]: tensor<1x20xi32>,
-// CHECK-SAME:    %[[VAL_1:.*]]: tensor<257x24xf32>,
-// CHECK-SAME:       %[[VAL_2:.*]]: index, %[[VAL_3:.*]]: index, %[[VAL_4:.*]]: index, %[[VAL_5:.*]]: index) -> tensor<1x1x4xf32> {
-// CHECK-DAG:       %[[VAL_6:.*]] = arith.constant dense<0> : vector<1x1x4xindex>
-// CHECK-DAG:       %[[VAL_7:.*]] = arith.constant dense<[0, 1, 2, 3]> : vector<4xindex>
-// CHECK-DAG:       %[[VAL_8:.*]] = arith.constant 0 : i32
-// CHECK-DAG:       %[[VAL_9:.*]] = arith.constant dense<256> : vector<1x1x4xindex>
-// CHECK-DAG:       %[[VAL_10:.*]] = arith.constant 0.000000e+00 : f32
-// CHECK-DAG:       %[[VAL_11:.*]] = arith.constant 0 : index
-// CHECK:           %[[VAL_12:.*]] = tensor.empty() : tensor<1x1x4xf32>
-// CHECK:           %[[VAL_13:.*]] = vector.broadcast %[[VAL_2]] : index to vector<1x1x4xindex>
-// CHECK:           %[[VAL_14:.*]] = vector.broadcast %[[VAL_4]] : index to vector<1x1x4xindex>
-// CHECK:           %[[VAL_15:.*]] = arith.addi %[[VAL_13]], %[[VAL_14]] : vector<1x1x4xindex>
-// CHECK:           %[[VAL_16:.*]] = vector.broadcast %[[VAL_3]] : index to vector<1x1x4xindex>
-// CHECK:           %[[VAL_17:.*]] = vector.broadcast %[[VAL_7]] : vector<4xindex> to vector<1x1x4xindex>
-// CHECK:           %[[VAL_18:.*]] = arith.addi %[[VAL_16]], %[[VAL_17]] : vector<1x1x4xindex>
-// CHECK:           %[[VAL_19:.*]] = vector.broadcast %[[VAL_5]] : index to vector<1x1x4xindex>
-// CHECK:           %[[VAL_20:.*]] = arith.addi %[[VAL_18]], %[[VAL_19]] : vector<1x1x4xindex>
-// CHECK:           %[[VAL_21:.*]] = vector.shape_cast %[[VAL_15]] : vector<1x1x4xindex> to vector<4xindex>
-// CHECK:           %[[VAL_22:.*]] = vector.extractelement %[[VAL_21]][%[[VAL_8]] : i32] : vector<4xindex>
+// CHECK-SAME:      %[[INPUT_1:.*]]: tensor<1x20xi32>,
+// CHECK-SAME:      %[[INPUT_2:.*]]: tensor<257x24xf32>,
+// CHECK:           %[[EXTRACTED_0_IDX_0:.*]] = arith.constant 0 : index
+// CHECK:           %[[EXTRACTED_0_IDX_1:.*]] = vector.extractelement %{{.*}}[%{{.*}} : i32] : vector<4xindex>
 // First `tensor.extract` from the generic Op - loop invariant scalar load.
-// CHECK:           %[[VAL_23:.*]] = tensor.extract %[[VAL_0]][%[[VAL_11]], %[[VAL_22]]] : tensor<1x20xi32>
-// CHECK:           %[[VAL_24:.*]] = arith.index_cast %[[VAL_23]] : i32 to index
-// CHECK:           %[[VAL_25:.*]] = vector.broadcast %[[VAL_24]] : index to vector<1x1x4xindex>
-// CHECK:           %[[VAL_26:.*]] = arith.maxsi %[[VAL_25]], %[[VAL_6]] : vector<1x1x4xindex>
-// CHECK:           %[[VAL_27:.*]] = arith.minsi %[[VAL_26]], %[[VAL_9]] : vector<1x1x4xindex>
-// CHECK:           %[[VAL_28:.*]] = vector.shape_cast %[[VAL_27]] : vector<1x1x4xindex> to vector<4xindex>
-// CHECK:           %[[VAL_29:.*]] = vector.extractelement %[[VAL_28]][%[[VAL_8]] : i32] : vector<4xindex>
-// CHECK:           %[[VAL_30:.*]] = vector.shape_cast %[[VAL_20]] : vector<1x1x4xindex> to vector<4xindex>
-// CHECK:           %[[VAL_31:.*]] = vector.extractelement %[[VAL_30]][%[[VAL_8]] : i32] : vector<4xindex>
+// CHECK:           tensor.extract %[[INPUT_1]][%[[EXTRACTED_0_IDX_0]], %[[EXTRACTED_0_IDX_1]]] : tensor<1x20xi32>
 // The following `tensor.extract` from the generic Op s a contiguous load (all Ops used
 // for address calculation also satisfy the required conditions).
-// CHECK:           %[[VAL_32:.*]] = vector.transfer_read %[[VAL_1]][%[[VAL_29]], %[[VAL_31]]], %[[VAL_10]] {in_bounds = [true, true]} : tensor<257x24xf32>, vector<1x4xf32>
-// CHECK:           %[[VAL_33:.*]] = vector.broadcast %[[VAL_32]] : vector<1x4xf32> to vector<1x1x4xf32>
-// CHECK:           %[[VAL_34:.*]] = vector.transfer_write %[[VAL_33]], %[[VAL_12]][%[[VAL_11]], %[[VAL_11]], %[[VAL_11]]] {in_bounds = [true, true, true]} : vector<1x1x4xf32>, tensor<1x1x4xf32>
-// CHECK:           return %[[VAL_34]] : tensor<1x1x4xf32>
-// CHECK:         }
+// CHECK:           vector.transfer_read %[[INPUT_2]][%{{.*}}, %{{.*}}, %{{.*}} {in_bounds = [true, true]} : tensor<257x24xf32>, vector<1x4xf32>
+
 
 transform.sequence failures(propagate) {
  ^bb1(%arg1: !transform.any_op):
    %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op
    %1 = get_closest_isolated_parent %0 : (!transform.any_op) -> !transform.any_op
    %2 = transform.structured.vectorize %1 { vectorize_nd_extract } : (!transform.any_op) -> !transform.any_op
  }
 

mlir/test/Dialect/Vector/sink-vector-broadcast.mlir

+// RUN: mlir-opt %s -test-sink-vector-broadcast -split-input-file | FileCheck %s
+
+// CHECK-LABEL:   func.func @broadcast_scalar(
+// CHECK-SAME:     %[[ARG_0:.*]]: index, %[[ARG_1:.*]]: index) -> vector<1x4xindex> {
+// CHECK:           %[[ADD:.*]] = arith.addi %[[ARG_0]], %[[ARG_1]] : index
+// CHECK:           %[[BCAST:.*]] = vector.broadcast %[[ADD]] : index to vector<1x4xindex>
+// CHECK:           return %[[BCAST]] : vector<1x4xindex>
+// CHECK:         }
+
+func.func @broadcast_scalar( %arg1: index, %arg2: index) -> vector<1x4xindex> {
+  %0 = vector.broadcast %arg1 : index to vector<1x4xindex>
+  %1 = vector.broadcast %arg2 : index to vector<1x4xindex>
+  %2 = arith.addi %0, %1 : vector<1x4xindex>
+  return %2 : vector<1x4xindex>
+}
+
+// -----
+
+// CHECK-LABEL:   func.func @broadcast_vector(
+// CHECK-SAME:      %[[ARG_0:.*]]: vector<4xf32>,
+// CHECK-SAME:      %[[ARG_1:.*]]: vector<4xf32>) -> vector<3x4xf32> {
+// CHECK:           %[[ADDF:.*]] = arith.addf %[[ARG_0]], %[[ARG_1]] : vector<4xf32>
+// CHECK:           %[[BCAST:.*]] = vector.broadcast %[[ADDF]] : vector<4xf32> to vector<3x4xf32>
+// CHECK:           return %[[BCAST]] : vector<3x4xf32>
+// CHECK:         }
+
+func.func @broadcast_vector( %arg1: vector<4xf32>, %arg2: vector<4xf32>) -> vector<3x4xf32> {
+  %arg1_bcast = vector.broadcast %arg1 : vector<4xf32> to vector<3x4xf32>
+  %arg2_bcast = vector.broadcast %arg2 : vector<4xf32> to vector<3x4xf32>
+  %2 = arith.addf %arg1_bcast, %arg2_bcast : vector<3x4xf32>
+  return %2 : vector<3x4xf32>
+}
+// -----
+
+// CHECK-LABEL:   func.func @broadcast_vector_and_scalar(
+// CHECK-SAME:      %[[ARG_0:.*]]: i32,
+// CHECK-SAME:      %[[ARG_1:.*]]: vector<4xi32>) -> vector<4xi32> {
+// CHECK:           %[[BCAST:.*]] = vector.broadcast %[[ARG_0]] : i32 to vector<4xi32>
+// CHECK:           %[[ADD:.*]] = arith.addi %[[BCAST]], %[[ARG_1]] : vector<4xi32>
+// CHECK:           return %[[ADD]] : vector<4xi32>
+// CHECK:         }
+
+func.func @broadcast_vector_and_scalar( %arg1: i32, %arg2: vector<4xi32>) -> vector<4xi32> {
+  %arg1_bcast = vector.broadcast %arg1 : i32 to vector<4xi32>
+  %2 = arith.addi %arg1_bcast, %arg2 : vector<4xi32>
+  return %2 : vector<4xi32>
+}
+
+// -----
+
+#matmat_accesses = [
+  affine_map<(i, j, k) -> (i, k)>,
+  affine_map<(i, j, k) -> (k, j)>,
+  affine_map<(i, j, k) -> (i, j)>
+]
+#matmat_trait = {
+  indexing_maps = #matmat_accesses,
+  iterator_types = ["parallel", "parallel", "reduction"]
+}
+
+// CHECK-LABEL:   func.func @broadcast_not_elementwise() -> vector<2x2xf32> {
+// CHECK-DAG:       %[[VAL_0:.*]] = arith.constant dense<1.000000e+00> : vector<2x2xf32>
+// CHECK-DAG:       %[[VAL_1:.*]] = arith.constant dense<2.000000e+00> : vector<2x2xf32>
+// CHECK-DAG:       %[[VAL_2:.*]] = arith.constant dense<3.000000e+00> : vector<2x2xf32>
+// CHECK:           %[[VAL_3:.*]] = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %[[VAL_0]], %[[VAL_1]], %[[VAL_2]] : vector<2x2xf32>, vector<2x2xf32> into vector<2x2xf32>
+func.func @broadcast_not_elementwise() -> vector<2x2xf32> {
+  %f1 = arith.constant 1.0: f32
+  %f2 = arith.constant 2.0: f32
+  %f3 = arith.constant 3.0: f32
+
+  %A = vector.broadcast %f1 : f32 to vector<2x2xf32>
+  %B = vector.broadcast %f2 : f32 to vector<2x2xf32>
+  %C = vector.broadcast %f3 : f32 to vector<2x2xf32>
+  %mm1 = vector.contract #matmat_trait %A, %B, %C
+    : vector<2x2xf32>, vector<2x2xf32> into vector<2x2xf32>
+
+  return %mm1 : vector<2x2xf32>
+}

mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp

struct TestVectorTransferCollapseInnerMostContiguousDims
 
   void runOnOperation() override {
     RewritePatternSet patterns(&getContext());
     populateVectorTransferCollapseInnerMostContiguousDimsPatterns(patterns);
     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
   }
 };
 
+struct TestSinkVectorBroadcast
+    : public PassWrapper<TestSinkVectorBroadcast, OperationPass<func::FuncOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestSinkVectorBroadcast)
+
+  TestSinkVectorBroadcast() = default;
+  TestSinkVectorBroadcast(const TestSinkVectorBroadcast &pass) = default;
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<memref::MemRefDialect, affine::AffineDialect>();
+  }
+
+  StringRef getArgument() const final { return "test-sink-vector-broadcast"; }
+
+  StringRef getDescription() const final {
+    return "Test lowering patterns that eliminate redundant brodacast "
+           "operations.";
+  }
+
+  void runOnOperation() override {
+    RewritePatternSet patterns(&getContext());
+    populateSinkVectorBroadcastPatterns(patterns);
+    (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
+  }
+};
+
 struct TestVectorReduceToContractPatternsPatterns
     : public PassWrapper<TestVectorReduceToContractPatternsPatterns,
                          OperationPass<func::FuncOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
       TestVectorReduceToContractPatternsPatterns)
 
   StringRef getArgument() const final {
     return "test-vector-reduction-to-contract-patterns";
void registerTestVectorLowerings() {
   PassRegistration<TestVectorTransferUnrollingPatterns>();
 
   PassRegistration<TestScalarVectorTransferLoweringPatterns>();
 
   PassRegistration<TestVectorTransferOpt>();
 
   PassRegistration<TestVectorTransferCollapseInnerMostContiguousDims>();
 
+  PassRegistration<TestSinkVectorBroadcast>();
+
   PassRegistration<TestVectorReduceToContractPatternsPatterns>();
 
   PassRegistration<TestFlattenVectorTransferPatterns>();
 
   PassRegistration<TestVectorScanLowering>();
 
   PassRegistration<TestVectorDistribution>();