[Mlir-commits] [mlir] [mlir][vector] Constrain broadcast->shape_cast folding (PR #190230)

[Andrzej Warzyński]

================
@@ -6725,6 +6725,22 @@ class ShapeCastBroadcastFolder final : public OpRewritePattern<ShapeCastOp> {
     // to
     // %1 = vector.broadcast %in : vector<3xf32> to vector<8x3xf32>
     VectorType dstVectorType = shapeCastOp.getResultVectorType();
+    VectorType intermediateType = broadcastOp.getResultVectorType();
+    // 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.
+    if (!srcIsScalar && srcVectorType.getRank() >= 2) {
+      auto hasUnitDim = [](ArrayRef<int64_t> shape) {
+        return llvm::any_of(shape, [](int64_t d) { return d == 1; });
+      };
+      if (hasUnitDim(srcVectorType.getShape()) &&
+          intermediateType.getRank() > dstVectorType.getRank()) {
+        ArrayRef<int64_t> droppedTrailingDims =
+            intermediateType.getShape().drop_front(dstVectorType.getRank());
+        if (hasUnitDim(droppedTrailingDims))
+          return failure();
+      }
+    }
----------------
banach-space wrote:

I think that this would be easier to follow if you introduces a helper hook, e.g.: `isDuplicatingBcast`, and then rewrote this logic as follows:

```cpp
if (isDuplicatingBcast(oldBcast) != isDuplicatingBcast(newBcast))
  return rewriter.notifyMatchFailure("Folding would result in broadcast mode change");
```

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