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

[Andrzej Warzyński]

================
@@ -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>
+}
----------------
banach-space wrote:

Is this example correct? In this cast, the dynamic dims could be anything, right? (*)
```mlir
    %cast = tensor.cast %src : tensor<1x3x8x1xi32> to tensor<?x3x?x1xi32>
```

Subsequently, `%tile_size` could be anything, yet the folder happily assumes it is `8` from `tensor<1x3x8x1xi32> `. 

Perhaps I am overthinking this, so let me ask a more specific question:
* Can we safely conclude that `?` and `?` in `tensor<?x3x?x1xi32>` are `1` and `8` from `tensor<1x3x8x1xi32>`?
* Are the semantics of `tensor.cast` clear from the docs for you?

Thanks!

(*) Looking at https://mlir.llvm.org/docs/Dialects/TensorOps/#tensorcast-tensorcastop, that's not super clear.

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