[Mlir-commits] [mlir] [MLIR][Vector] Refactor tests for contract -> OP transforms (2/N) (PR #73445)

[Andrzej Warzyński]

https://github.com/banach-space created https://github.com/llvm/llvm-project/pull/73445

This is a direct follow-up of #73348. The matvec trait that's used for
`@matvec_m_mk_k` was incorrectly updated from:

```
  affine_map<(m, k) -> (m)>,
  affine_map<(m, k) -> (m, k)>,
  affine_map<(m, k) -> (k)>
]
  indexing_maps = #redpar_vecmattrans_accesses,
  iterator_types = ["reduction", "parallel"]
}
```

to:

```
  affine_map<(m, k) -> (k)>,
  affine_map<(m, k) -> (k, m)>,
  affine_map<(m, k) -> (m)>
]
  indexing_maps = #matvec_accesses_4,
  iterator_types = ["parallel", "reduction"]
}
```

Note that these traits describe identical matvec operation, hence the
`CHECK` lines are identical for both.

Also, `#redpar_vecmattrans_trait` is identical to `#matvec_trait_8`
that's already present in:

  vector-contract-to-outerproduct-matvec-transforms.mlir

For this reason:
  * `@matvec_m_mk_k` is moved near other tests that already use `#matvec_trait_8`,
  * `#redpar_vecmattrans_trait` is replaced `#matvec_trait_8`.

This is a part of a larger effort to add cases with scalable vectors to
tests for the Vector dialect. I am refactoring these test as a
preparation for follow-up patches.

Implements #72834.


>From 51009fc1e199a5428e29af3ef04243757f954765 Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Sun, 26 Nov 2023 11:53:19 +0000
Subject: [PATCH] [MLIR][Vector] Refactor tests for contract -> OP transforms
 (2/N)

This is a direct follow-up of #73348. The matvec trait that's used for
`@matvec_m_mk_k` was incorrectly updated from:

```
  affine_map<(m, k) -> (m)>,
  affine_map<(m, k) -> (m, k)>,
  affine_map<(m, k) -> (k)>
]
  indexing_maps = #redpar_vecmattrans_accesses,
  iterator_types = ["reduction", "parallel"]
}
```

to:

```
  affine_map<(m, k) -> (k)>,
  affine_map<(m, k) -> (k, m)>,
  affine_map<(m, k) -> (m)>
]
  indexing_maps = #matvec_accesses_4,
  iterator_types = ["parallel", "reduction"]
}
```

Note that these traits describe identical matvec operation, hence the
`CHECK` lines are identical for both.

Also, `#redpar_vecmattrans_trait` is identical to `#matvec_trait_8`
that's already present in:

  vector-contract-to-outerproduct-matvec-transforms.mlir

For this reason:
  * `@matvec_m_mk_k` is moved near other tests that already use `#matvec_trait_8`,
  * `#redpar_vecmattrans_trait` is replaced `#matvec_trait_8`.

This is a part of a larger effort to add cases with scalable vectors to
tests for the Vector dialect. I am refactoring these test as a
preparation for follow-up patches.

Implements #72834.
---
 ...act-to-outerproduct-matvec-transforms.mlir | 54 ++++++++-----------
 1 file changed, 22 insertions(+), 32 deletions(-)

diff --git a/mlir/test/Dialect/Vector/vector-contract-to-outerproduct-matvec-transforms.mlir b/mlir/test/Dialect/Vector/vector-contract-to-outerproduct-matvec-transforms.mlir
index 3ca3d344c1abe04..7456e122e946f7d 100644
--- a/mlir/test/Dialect/Vector/vector-contract-to-outerproduct-matvec-transforms.mlir
+++ b/mlir/test/Dialect/Vector/vector-contract-to-outerproduct-matvec-transforms.mlir
@@ -46,13 +46,13 @@
   iterator_types = ["parallel", "reduction"]
 }
 
-#redpar_vecmattrans_accesses = [
-  affine_map<(m, k) -> (m)>,
-  affine_map<(m, k) -> (m, k)>,
-  affine_map<(m, k) -> (k)>
+#matvec_accesses_8 = [
+  affine_map<(k, m) -> (k)>,
+  affine_map<(k, m) -> (k, m)>,
+  affine_map<(k, m) -> (m)>
 ]
-#redpar_vecmattrans_trait = {
-  indexing_maps = #redpar_vecmattrans_accesses,
+#matvec_trait_8 = {
+  indexing_maps = #matvec_accesses_8,
   iterator_types = ["reduction", "parallel"]
 }
 
@@ -321,23 +321,6 @@ func.func @matvec_k_km_m(%A: vector<2x2xf32>,
   return %0 : vector<2xf32>
 }
 
-// CHECK-LABEL: func @matvec_m_mk_k
-// CHECK-SAME: %[[A:.*0]]: vector<2x2xf32>
-// CHECK-SAME: %[[B:.*1]]: vector<2xf32>
-// CHECK-SAME: %[[C:.*2]]: vector<2xf32>
-// CHECK: %[[T3:.*]] = vector.extract %[[A]][0] : vector<2xf32> from vector<2x2xf32>
-// CHECK: %[[T4:.*]] = vector.extract %[[B]][0] : f32 from vector<2xf32>
-// CHECK: %[[T5:.*]] = vector.outerproduct %[[T3]], %[[T4]], %[[C]] {kind = #vector.kind<add>} : vector<2xf32>, f32
-// CHECK: %[[T6:.*]] = vector.extract %[[A]][1] : vector<2xf32> from vector<2x2xf32>
-// CHECK: %[[T7:.*]] = vector.extract %[[B]][1] : f32 from vector<2xf32>
-// CHECK: %[[T8:.*]] = vector.outerproduct %[[T6]], %[[T7]], %[[T5]] {kind = #vector.kind<add>} : vector<2xf32>, f32
-func.func @matvec_m_mk_k(%A: vector<2x2xf32>,
-                         %x: vector<2xf32>,
-                         %b: vector<2xf32>) -> vector<2xf32> {
-  %0 = vector.contract #matvec_trait_4 %x, %A, %b : vector<2xf32>, vector<2x2xf32> into vector<2xf32>
-  return %0 : vector<2xf32>
-}
-
 // ============================================================================
 //  Matvec 5 (masked + scalable)
 // ============================================================================
@@ -474,16 +457,23 @@ func.func @masked_tmatvec_k_mk_m_scalable_parallel_dim(%arg0: vector<[4]x2xf32>,
 }
 
 // ============================================================================
-//  Matvec 8 (masked + scalable)
+//  Matvec 8 (plain + masked + scalable)
 // ============================================================================
-#matvec_accesses_8 = [
-  affine_map<(k, m) -> (k)>,
-  affine_map<(k, m) -> (k, m)>,
-  affine_map<(k, m) -> (m)>
-]
-#matvec_trait_8 = {
-  indexing_maps = #matvec_accesses_8,
-  iterator_types = ["reduction", "parallel"]
+// CHECK-LABEL: func @matvec_m_mk_k
+// CHECK-SAME: %[[A:.*0]]: vector<2x2xf32>
+// CHECK-SAME: %[[B:.*1]]: vector<2xf32>
+// CHECK-SAME: %[[C:.*2]]: vector<2xf32>
+// CHECK: %[[T3:.*]] = vector.extract %[[A]][0] : vector<2xf32> from vector<2x2xf32>
+// CHECK: %[[T4:.*]] = vector.extract %[[B]][0] : f32 from vector<2xf32>
+// CHECK: %[[T5:.*]] = vector.outerproduct %[[T3]], %[[T4]], %[[C]] {kind = #vector.kind<add>} : vector<2xf32>, f32
+// CHECK: %[[T6:.*]] = vector.extract %[[A]][1] : vector<2xf32> from vector<2x2xf32>
+// CHECK: %[[T7:.*]] = vector.extract %[[B]][1] : f32 from vector<2xf32>
+// CHECK: %[[T8:.*]] = vector.outerproduct %[[T6]], %[[T7]], %[[T5]] {kind = #vector.kind<add>} : vector<2xf32>, f32
+func.func @matvec_m_mk_k(%A: vector<2x2xf32>,
+                         %x: vector<2xf32>,
+                         %b: vector<2xf32>) -> vector<2xf32> {
+  %0 = vector.contract #matvec_trait_8 %x, %A, %b : vector<2xf32>, vector<2x2xf32> into vector<2xf32>
+  return %0 : vector<2xf32>
 }
 
 // CHECK-LABEL: @masked_tmatvec_k_km_m