[Mlir-commits] [mlir] [MLIR][Vector] Fix ShapeCastBroadcastFolder when shape_cast changes non-unit dim strides (PR #227574)

[Federico Bruzzone]

================
@@ -7194,16 +7178,17 @@ class ShapeCastBroadcastFolder final : public OpRewritePattern<ShapeCastOp> {
           BroadcastableToResult::Success) {
         return failure();
       }
-      // Avoid folding if this would result in switching between the two
-      // distinct semantic modes of vector.broadcast (duplication vs
-      // stretching). See https://github.com/llvm/llvm-project/issues/190614.
-      // This is detected by a change in the stretching factor. However if the
-      // source has a single element, there is no ambiguity.
-      if (srcVectorType.getNumElements() != 1) {
-        if (getBroadcastStretchingFactor(srcShape, dstShape) !=
-            getBroadcastStretchingFactor(srcShape, broadcastShape)) {
+      // `shape_cast` preserves the flattened 1-D (row-major) element order.
+      // Thus, `shape_cast(broadcast(src, broadcastShape), dstShape)` produces
+      // the exact same elements as `broadcast(src, dstShape)` if and only if
+      // every non-unit dimension of `src` has the same linear stride in both
+      // `broadcastShape` and `dstShape`.
+      int64_t bcastStride = 1, dstStride = 1;
+      for (size_t i = 1, e = srcShape.size(); i <= e; ++i) {
+        if (srcShape[e - i] != 1 && bcastStride != dstStride)
           return failure();
-        }
+        bcastStride *= broadcastShape.take_back(i).front();
+        dstStride *= dstShape.take_back(i).front();
       }
     }
----------------
FedericoBruzzone wrote:

Are scalable dims handled correctly here? The strides are compared as static integers, so a scalable dim in `broadcastShape` and a fixed dim in `dstShape` with the same numeric size would compare equal (?)

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