[Mlir-commits] [mlir] [mlir][scf] Fix ForOpTensorCastFolder on loops that change shape (PR #234422)

[Jared Hoberock]

================
@@ -1020,6 +1046,16 @@ struct ForOpTensorCastFolder : public OpRewritePattern<ForOp> {
         continue;
       if (!std::get<1>(it).hasOneUse())
         continue;
+      // The fold casts the yielded value to the static type, which is only
+      // valid if every iteration yields a value of that type.
+      int64_t resultNumber = std::get<1>(it).getResultNumber();
+      auto yieldedCast =
+          op.getYieldedValues()[resultNumber].getDefiningOp<tensor::CastOp>();
+      bool yieldsStaticType =
+          yieldedCast && yieldedCast.getSource().getType() ==
----------------
jaredhoberock wrote:

I think so. After the fold, the loop carries `tensor<?x16xf32>`, so the first dimension is free to change between iterations. The only size the fold asserts is the 16, and the yielded value already has that type.

For example, this is what the canonicalizer would produce on a case like this:

```
%0 = scf.for %i = %c0 to %c8 step %c1 iter_args(%arg2 = %arg0) -> (tensor<?x16xf32>) {
    %cast_0 = tensor.cast %arg2 : tensor<?x16xf32> to tensor<?x?xf32>
    %1 = func.call @do(%cast_0) : (tensor<?x?xf32>) -> tensor<?x16xf32>
    scf.yield %1 : tensor<?x16xf32>
}
%cast = tensor.cast %0 : tensor<?x16xf32> to tensor<?x?xf32>
```

I added this case as a test and also clarified the comment above the check.

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