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

[Adam Siemieniuk]

================
@@ -7556,64 +7568,116 @@ class FoldTransposeBroadcast : public OpRewritePattern<vector::TransposeOp> {
 
   LogicalResult matchAndRewrite(vector::TransposeOp transpose,
                                 PatternRewriter &rewriter) const override {
-
     vector::BroadcastOp broadcast =
         transpose.getVector().getDefiningOp<vector::BroadcastOp>();
-    if (!broadcast) {
+    if (!broadcast)
       return rewriter.notifyMatchFailure(transpose,
                                          "not preceded by a broadcast");
-    }
-
-    auto inputType = dyn_cast<VectorType>(broadcast.getSourceType());
-    VectorType outputType = transpose.getResultVectorType();
 
-    // transpose(broadcast(scalar)) -> broadcast(scalar) is always valid
-    bool inputIsScalar = !inputType;
-    if (inputIsScalar) {
-      rewriter.replaceOpWithNewOp<vector::BroadcastOp>(transpose, outputType,
-                                                       broadcast.getSource());
-      return success();
-    }
-
-    ArrayRef<int64_t> permutation = transpose.getPermutation();
-    ArrayRef<int64_t> inputShape = inputType.getShape();
-    int64_t inputRank = inputType.getRank();
-    int64_t outputRank = transpose.getType().getRank();
-    int64_t deltaRank = outputRank - inputRank;
-
-    int low = 0;
-    for (int inputIndex = 0; inputIndex < inputRank; ++inputIndex) {
-      bool notOne = inputShape[inputIndex] != 1;
-      bool prevNotOne = (inputIndex != 0 && inputShape[inputIndex - 1] != 1);
-      bool groupEndFound = notOne || prevNotOne;
-      if (groupEndFound) {
-        int high = inputIndex + deltaRank;
-        // Return failure if not all permutation destinations for indices in
-        // [low, high) are in [low, high), i.e. the permutation is not local to
-        // the group.
-        for (int i = low; i < high; ++i) {
-          if (permutation[i] < low || permutation[i] >= high) {
-            return rewriter.notifyMatchFailure(
-                transpose, "permutation not local to group");
-          }
-        }
-        low = high;
-      }
+    VectorType outType = transpose.getResultVectorType();
+    auto srcType = dyn_cast<VectorType>(broadcast.getSourceType());
+
+    // transpose(broadcast(scalar)) always folds. Otherwise the source must be
+    // broadcastable to the result and the transpose must leave its non-broadcast
+    // dims in place, i.e. map each to itself (from == to).
+    if (srcType) {
+      if (vector::isBroadcastableTo(srcType, outType) !=
+          vector::BroadcastableToResult::Success)
+        return rewriter.notifyMatchFailure(transpose, "not broadcastable");
+
+      int64_t rankDelta = outType.getRank() - srcType.getRank();
+      SmallVector<int64_t> axes;
+      for (int64_t axis : nonBroadcastAxes(srcType))
+        axes.push_back(axis + rankDelta);
+      if (!transposeMapsAxes(axes, axes, transpose.getPermutation()))
+        return rewriter.notifyMatchFailure(transpose,
+                                           "not an order-preserving broadcast");
----------------
adam-smnk wrote:

nit: isn't it the "transpose" that has to be order-preserving?

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