[Mlir-commits] [mlir] [mlir][vector] Fold transpose(broadcast(shape_cast)) to broadcast (PR #215940)

[Federico Bruzzone]

================
@@ -7777,13 +7788,118 @@ class FoldTransposeBroadcast : public OpRewritePattern<vector::TransposeOp> {
   }
 };
 
+/// Folds transpose(broadcast(shape_cast(x))) to broadcast(x) when:
+///  - x is broadcastable to the transpose's result type.
+///  - The shape_cast only inserts, removes, or moves fixed size-1 dimensions.
+///    All other dimensions retain their sizes, scalability, and relative order.
+///  - The transpose preserves the relative order of all dimensions inherited
+///    from the shape_cast result, including size-1 dimensions. Dimensions
+///    prepended by the broadcast may be freely permuted around them.
+///  - Each non-unit dimension of x ends up at the position it would occupy in
+///    a direct broadcast of x. Scalable dimensions count as non-unit.
+///
+/// These conditions ensure that the direct broadcast replicates the same
+/// source elements along the same result dimensions as the original chain.
+///
+/// For example, the broadcast prepends a dimension that the transpose moves
+/// to the back, where x already has a size-1 dimension:
+/// ```
+///  %0 = vector.shape_cast %x : vector<1x32x1xf32> to vector<1x32xf32>
+///  %1 = vector.broadcast %0 : vector<1x32xf32> to vector<64x1x32xf32>
+///  %2 = vector.transpose %1, [1, 2, 0] : vector<64x1x32xf32>
+///                                                    to vector<1x32x64xf32>
+/// ```
+/// This rewrites to:
+/// ```
+///  %2 = vector.broadcast %x : vector<1x32x1xf32> to vector<1x32x64xf32>
+/// ```
+///
+/// The order-preservation check is conservative: a size-1 dimension expanded
+/// by the broadcast is still an inherited dimension, not a prepended one.
+/// The following chain therefore does not fold, even though broadcasting x
+/// directly to vector<3x4xf32> would produce the same result:
+/// ```
+///  %0 = vector.shape_cast %x : vector<1x4xf32> to vector<4x1xf32>
+///  %1 = vector.broadcast %0 : vector<4x1xf32> to vector<4x3xf32>
+///  %2 = vector.transpose %1, [1, 0] : vector<4x3xf32> to vector<3x4xf32>
+/// ```
+class FoldTransposeShapeCastBroadcast
+    : public OpRewritePattern<vector::TransposeOp> {
+public:
+  using Base::Base;
+
+  LogicalResult matchAndRewrite(vector::TransposeOp transpose,
+                                PatternRewriter &rewriter) const override {
+    auto broadcast = transpose.getVector().getDefiningOp<vector::BroadcastOp>();
+    if (!broadcast)
+      return rewriter.notifyMatchFailure(transpose, "not a broadcast source");
+    auto shapeCast = broadcast.getSource().getDefiningOp<vector::ShapeCastOp>();
+    if (!shapeCast)
+      return rewriter.notifyMatchFailure(transpose, "not a shape_cast source");
+
+    VectorType srcType = shapeCast.getSourceVectorType(); // x
+    VectorType midType =
+        shapeCast.getResultVectorType(); // y, the broadcast input
+    VectorType outType = transpose.getResultVectorType();
+
+    if (vector::isBroadcastableTo(srcType, outType) !=
+        vector::BroadcastableToResult::Success)
+      return rewriter.notifyMatchFailure(transpose, "not broadcastable");
+
+    // Size-1-dim-only shape_cast: x and y must have identical non-broadcast
+    // dims in the same order. E.g. rejects 5x4 -> 4x5x1, which reorders 5
+    // and 4.
+    SmallVector<int64_t> srcAxes = nonUnitAxes(srcType);
+    SmallVector<int64_t> midAxes = nonUnitAxes(midType);
+    if (srcAxes.size() != midAxes.size())
+      return rewriter.notifyMatchFailure(transpose,
+                                         "reshapes a non-broadcast dim");
----------------
FedericoBruzzone wrote:

```suggestion
                                         "reshapes a non-unit dim");
```

https://github.com/llvm/llvm-project/pull/215940