[Mlir-commits] [mlir] [mlir][linalg] Fix crash in FoldTensorCastUnPackOp with dynamic non-constant tile size (PR #189071)

[Mehdi Amini]

================
@@ -1952,6 +1952,29 @@ func.func @fold_cast_unpack_dynamic_tile_size(
 
 // -----
 
+// Regression test: FoldTensorCastUnPackOp must not crash when a tile size is a
+// dynamic (non-constant) SSA value and the cast makes the packed dim static.
+// The static dim value from the cast should be used as the new tile size.
+// CHECK-LABEL: func.func @fold_cast_unpack_nonconstant_dynamic_tile(
+// CHECK-SAME:     %[[SRC:.*]]: tensor<1x3x8x1xi32>,
+// CHECK-SAME:     %[[TILE:.*]]: index,
+// CHECK-SAME:     %[[DEST:.*]]: tensor<7x3xi32>) -> tensor<7x3xi32> {
+// CHECK:          %[[RES:.*]] = linalg.unpack %[[SRC]] inner_dims_pos = [0, 1] inner_tiles = [8, 1] into %[[DEST]] : tensor<1x3x8x1xi32> -> tensor<7x3xi32>
+// CHECK:          return %[[RES]] : tensor<7x3xi32>
+func.func @fold_cast_unpack_nonconstant_dynamic_tile(
+  %src: tensor<1x3x8x1xi32>,
+  %tile_size: index,
+  %dest: tensor<7x3xi32>) -> tensor<7x3xi32> {
+    %cast = tensor.cast %src : tensor<1x3x8x1xi32> to tensor<?x3x?x1xi32>
+    %unpack = linalg.unpack %cast
+      inner_dims_pos = [0, 1]
+      inner_tiles = [%tile_size, 1]
+      into %dest : tensor<?x3x?x1xi32> -> tensor<7x3xi32>
+    return %unpack : tensor<7x3xi32>
+}
----------------
joker-eph wrote:

> Can we safely conclude that ? and ? in tensor<?x3x?x1xi32> are 1 and 8 from tensor<1x3x8x1xi32>?

Yes, such cast is just here to please the type-system by "erasing" the static dim. 

> cast is pure so does not involve memory effects, so the elements must not change neither value nor location.

Since we're operating on Tensors, you could move elements (or change their value) in a Pure fashion actually.
(I'm not saying `cast` should be allowed to do so, just that "Pure" doesn't prevent it here).

> I'd add tests where they're constant and don't match, where the source is dynamic and the dest is not, as well to make sure the patch is not ignoring something else.

I added @fold_cast_unpack_constant_tile_mismatch and @fold_cast_pack_nonconstant_dynamic_tile to provide more coverage.

The op we're folding to for each of these tests will look like:
```
    %unpack = linalg.unpack %arg0 inner_dims_pos = [0, 1] inner_tiles = [8, 1] into %arg1 : tensor<1x3x8x1xi32> -> tensor<7x3xi32>
```

and:

```
    %pack = linalg.pack %arg0 inner_dims_pos = [0, 1] inner_tiles = [8, 1] into %arg2 : tensor<8x3xi32> -> tensor<1x3x8x1xi32>
```

Did you think of another kind of test?

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