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

[Federico Bruzzone]

================
@@ -7173,37 +7157,65 @@ class ShapeCastBroadcastFolder final : public OpRewritePattern<ShapeCastOp> {
     if (!broadcastOp)
       return failure();
 
-    auto srcVectorType = dyn_cast<VectorType>(broadcastOp.getSourceType());
-    bool srcIsScalar = !srcVectorType;
-
     // Replace Y = ShapeCast(Broadcast(X)) with Y = Broadcast(X)
     // Example
     // %0 = vector.broadcast %in : vector<3xf32> to vector<2x4x3xf32>
     // %1 = vector.shape_cast %0 : vector<2x4x3xf32> to vector<8x3xf32>
     // to
     // %1 = vector.broadcast %in : vector<3xf32> to vector<8x3xf32>
+    auto srcVectorType = dyn_cast<VectorType>(broadcastOp.getSourceType());
+    bool srcIsScalar = !srcVectorType;
     VectorType dstVectorType = shapeCastOp.getResultVectorType();
-    ArrayRef<int64_t> dstShape = dstVectorType.getShape();
-    ArrayRef<int64_t> srcShape =
-        srcIsScalar ? ArrayRef<int64_t>{} : srcVectorType.getShape();
-    ArrayRef<int64_t> broadcastShape =
-        broadcastOp.getResultVectorType().getShape();
 
     if (!srcIsScalar) {
       if (isBroadcastableTo(srcVectorType, dstVectorType) !=
           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`.
+      //
+      // A scalable dimension of static size D has runtime size D * vscale.
+      // A linear stride through trailing dimensions is thus:
+      //   stride = staticStride * (vscale ^ numScalableDims)
+      // Since vscale is an unknown runtime quantity (vscale >= 1), both
+      // staticStride and numScalableDims must match to guarantee stride
+      // equivalence across all targets.
+      VectorType broadcastVectorType = broadcastOp.getResultVectorType();
+      SmallVector<int64_t> bcastStrides =
+          computeStrides(broadcastVectorType.getShape());
+      SmallVector<int64_t> dstStrides =
+          computeStrides(dstVectorType.getShape());
+      ArrayRef<bool> bcastScalable = broadcastVectorType.getScalableDims();
+      ArrayRef<bool> dstScalable = dstVectorType.getScalableDims();
+      ArrayRef<bool> srcScalable = srcVectorType.getScalableDims();
+
+      // Iterate from the innermost (trailing) dimension outward. `llvm::zip`
+      // stops once all `src` dimensions are visited, aligning each `src`
+      // dimension with the corresponding trailing dimension of `broadcast` and
+      // `dst`.
+      for (auto [idx, dims] : llvm::enumerate(llvm::zip(
+               llvm::reverse(srcVectorType.getShape()),
+               llvm::reverse(srcScalable), llvm::reverse(bcastStrides),
+               llvm::reverse(dstStrides)))) {
----------------
FedericoBruzzone wrote:

This is  hard to follow now. Explicit running products would make the invariant clear I believe...

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