[Mlir-commits] [mlir] [mlir][vector] Preserve dynamic mask bounds in shape cast folding (PR #229699)

[purnima shrivastava]

================
@@ -3590,6 +3607,40 @@ func.func @fold_shape_cast_with_mask_leading_unit(%arg0: tensor<1x?xf32>) -> vec
 
 // -----
 
+// CHECK-LABEL:   func.func @fold_shape_cast_with_mask_leading_unit_dynamic(
+// CHECK-SAME:     %[[N:.*]]: index) -> vector<8xi1> {
+func.func @fold_shape_cast_with_mask_leading_unit_dynamic(%n: index) -> vector<8xi1> {
+// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+// CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index
+// CHECK:     %[[MASK:.*]] = vector.create_mask %[[C1]], %[[N]], %[[C8]] : vector<1x1x8xi1>
+// CHECK:     %[[CAST:.*]] = vector.shape_cast %[[MASK]] : vector<1x1x8xi1> to vector<8xi1>
+// CHECK:     return %[[CAST]] : vector<8xi1>
+  %c1 = arith.constant 1 : index
+  %c8 = arith.constant 8 : index
+  %mask = vector.create_mask %c1, %n, %c8 : vector<1x1x8xi1>
+  %cast = vector.shape_cast %mask : vector<1x1x8xi1> to vector<8xi1>
+  return %cast : vector<8xi1>
+}
+
+// -----
+
+// CHECK-LABEL:   func.func @fold_shape_cast_with_mask_leading_unit_dynamic_mixed_bounds(
+// CHECK-SAME:     %[[N:.*]]: index, %[[M:.*]]: index) -> vector<8x4xi1> {
+func.func @fold_shape_cast_with_mask_leading_unit_dynamic_mixed_bounds(%n: index, %m: index) -> vector<8x4xi1> {
----------------
purnima-nlp wrote:

> Please check if it already exists :D

Thanks! I checked the existing tests, and this positive case is already covered by `fold_shape_cast_with_mask_trailing_unit` and `fold_shape_cast_with_mask_leading_unit`. Both have a dynamic `%dim` in a kept dimension and check that the `shape_cast` is folded while the dynamic bound is preserved in the resulting `vector.create_mask`.

So I think the existing positive coverage complements the new negative cases, where a dynamic bound is among the dimensions being dropped.

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