[Mlir-commits] [mlir] [mlir][tensor] Improve `FoldTensorCastProducerOp` (dynamic shapes) (PR #114559)

[Andrzej Warzyński]

================
@@ -4756,8 +4756,50 @@ struct FoldTensorCastProducerOp
         newResultTypes[dpsInitIdx++] = newOperands.back().getType();
     }
 
-    // Clone op.
-    Operation *newOp = clone(rewriter, op, newResultTypes, newOperands);
+    // For ops that have sizes-like attribute, update these accordingly.
+    // For now, only `tensor.pack` is supported.
+    // TODO: Generalize to make it work with other ops as well (e.g.
+    // `tensor.unpack`)
+    SmallVector<OpFoldResult> newMixedTileSizes;
----------------
banach-space wrote:

> This seems like it is mixing something related to PackOp into a method that is for general destination passing style operations. 

No, this logic applies to all Tensor ops that take size-like attribute. Here's an example with `tensor.extract_slice`:
```mlir
%1 = tensor.extract_slice %t[%c0, %c0, %c0][%c2, %c2, %c2][%c1, %c1, %c1] : tensor<8x16x4xf32> to tensor<?x?x?xf32>
%res = tensor.cast  tensor<?x?x?xf32> to tensor<2x2x2xf32>
```

The folder should return this (i.e update the result type **and** the sizes attribute):
```mlir
%res = tensor.extract_slice %t[%c0, %c0, %c0][2, 2, 2][%c1, %c1, %c1] : tensor<8x16x4xf32> to tensor<2x2x2xf32>
```

instead of (i.e. update the result type):
```mlir
%res = tensor.extract_slice %t[%c0, %c0, %c0][%c2, %c2, %c2][%c1, %c1, %c1] : tensor<8x16x4xf32> to tensor<2x2x2xf32>
```

The latter fails verification (that's the correct and intended behaviour, as discussed in this [thread](https://discourse.llvm.org/t/tensor-ops-with-dynamic-sizes-which-behaviour-is-more-correct) ). 

I only happen to implement the logic for `tensor.pack` as that's what's blocking me right now. I am happy to extend this to other Ops, but want to avoid sending a large patch - these tend to get longer to get reviewed ;-)

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