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

[Renato Golin]

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

>From the docs:

> Convert a tensor from one type to an equivalent type without changing any data elements. 

This is not explicit that the shape must be the same either, but `cast` is _pure_ so does not involve memory effects, so the elements must not change neither value nor location.

> The source and destination types must both be tensor types with the same element type.

Not relevant here.

> If both are ranked, then the rank should be the same and static dimensions should match. 

This is still true for dynamic dims and the patch is not changing that, so _check_.

> The operation is invalid if converting to a mismatching constant dimension.

This is what the assert is about. If the dim is constant and they're not the same, crash. I'm not sure this is being retained after this patch.

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.

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