[Mlir-commits] [mlir] [mlir][vector] Update `castAwayContractionLeadingOneDim` to omit transposes solely on leading unit dims. (PR #85694)

[Andrzej Warzyński]

================
@@ -166,6 +166,29 @@ func.func @cast_away_contraction_leading_one_dims_nonleadingunitdim_rank4_acctra
 
 // -----
 
+// CHECK-DAG: #[[$MAP_0:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
+// CHECK-DAG: #[[$MAP_1:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+// CHECK-DAG: #[[$MAP_2:.+]] = affine_map<(d0, d1, d2) -> (d1)>
+
+// CHECK-LABEL:   func.func @cast_away_contraction_does_not_transpose_leading_unit_dims(
+// CHECK-SAME:                                 %[[LHS:.*]]: vector<1x1x8xi32>,
+// CHECK-SAME:                                 %[[RHS:.*]]: vector<1x8x8xi32>,
+// CHECK-SAME:                                 %[[ACC:.*]]: vector<1x8xi32>) -> vector<1x8xi32> {
+// CHECK:           %[[EXT_LHS:.*]] = vector.extract %[[LHS]][0] : vector<1x8xi32> from vector<1x1x8xi32>
+// CHECK:           %[[EXT_ACC:.*]] = vector.extract %[[ACC]][0] : vector<8xi32> from vector<1x8xi32>
+// CHECK:           %[[RES:.*]] = vector.contract {indexing_maps = [#[[$MAP_0]], #[[$MAP_1]], #[[$MAP_2]]], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %[[EXT_LHS]], %[[RHS]], %[[EXT_ACC]] : vector<1x8xi32>, vector<1x8x8xi32> into vector<8xi32>
+// CHECK:           %[[BROADCAST_RES:.*]] = vector.broadcast %[[RES]] : vector<8xi32> to vector<1x8xi32>
+// CHECK:           return %[[BROADCAST_RES]] : vector<1x8xi32>
+// CHECK:         }
+// CHECK-NOT        vector.transpose
----------------
banach-space wrote:

Currently this will only check that there's no `vector.transpose` after `vector.contract`. IIUC, that's not the intention here. See https://llvm.org/docs/CommandGuide/FileCheck.html#the-check-not-directive

Given that we only need to verify that there's no `vector.transpose` _before_ `vector.contract`, this should be sufficient:
```suggestion
// CHECK-LABEL:   func.func @cast_away_contraction_does_not_transpose_leading_unit_dims
// CHECK-NOT   vector.transpose
// CHECK:           vector.contract
```

Also, from https://mlir.llvm.org/getting_started/TestingGuide/#filecheck-best-practices:
> Tests should be minimal, and only check what is absolutely necessary.

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