Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged

Conversation

banach-space
Copy link
Contributor

@banach-space banach-space commented Nov 26, 2023

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

#redpar_vecmattrans_accesses = [
  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:

#matvec_accesses_4 = [
  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 tests so that it's
easier to identify what cases are tested and where to add tests for
scalable vectors.

Implements #72834.

@llvmbot
Copy link
Member

llvmbot commented Nov 26, 2023

@llvm/pr-subscribers-mlir

Author: Andrzej Warzyński (banach-space)

Changes

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&lt;(m, k) -&gt; (m)&gt;,
  affine_map&lt;(m, k) -&gt; (m, k)&gt;,
  affine_map&lt;(m, k) -&gt; (k)&gt;
]
  indexing_maps = #redpar_vecmattrans_accesses,
  iterator_types = ["reduction", "parallel"]
}

to:

  affine_map&lt;(m, k) -&gt; (k)&gt;,
  affine_map&lt;(m, k) -&gt; (k, m)&gt;,
  affine_map&lt;(m, k) -&gt; (m)&gt;
]
  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.


Full diff: https://github.com/llvm/llvm-project/pull/73445.diff

1 Files Affected:

  • (modified) mlir/test/Dialect/Vector/vector-contract-to-outerproduct-matvec-transforms.mlir (+22-32)
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

@llvmbot
Copy link
Member

llvmbot commented Nov 26, 2023

@llvm/pr-subscribers-mlir-vector

Author: Andrzej Warzyński (banach-space)

Changes

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&lt;(m, k) -&gt; (m)&gt;,
  affine_map&lt;(m, k) -&gt; (m, k)&gt;,
  affine_map&lt;(m, k) -&gt; (k)&gt;
]
  indexing_maps = #redpar_vecmattrans_accesses,
  iterator_types = ["reduction", "parallel"]
}

to:

  affine_map&lt;(m, k) -&gt; (k)&gt;,
  affine_map&lt;(m, k) -&gt; (k, m)&gt;,
  affine_map&lt;(m, k) -&gt; (m)&gt;
]
  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.


Full diff: https://github.com/llvm/llvm-project/pull/73445.diff

1 Files Affected:

  • (modified) mlir/test/Dialect/Vector/vector-contract-to-outerproduct-matvec-transforms.mlir (+22-32)
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

This is a direct follow-up of llvm#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 tests so that it's
easier to identify what cases are tested and where to add tests for
scalable vectors.

Implements llvm#72834.
Copy link
Member

@MacDue MacDue left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, thanks 👍

@banach-space banach-space merged commit eb1d506 into llvm:main Nov 28, 2023
@banach-space banach-space deleted the andrzej/update_contract_test_v2_p2 branch March 8, 2024 14:42
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants