[Mlir-commits] [mlir] [mlir][SVE] Add more e2e test for vector.contract (PR #70367)
Andrzej Warzyński
llvmlistbot at llvm.org
Fri Oct 27 04:50:04 PDT 2023
https://github.com/banach-space updated https://github.com/llvm/llvm-project/pull/70367
>From 7a0d0c6ffc2a7915ce882adc6b39e458f65bc1c2 Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Thu, 26 Oct 2023 12:49:32 +0000
Subject: [PATCH 1/3] [mlir][SVE] Add more e2e test for vector.contract
Adds basic integration tests for `vector.contract` for the dot product
and matvec operations. These tests excercise scalable vectors.
---
.../Vector/CPU/ArmSVE/test-contraction.mlir | 82 ++++++++++++++++++-
1 file changed, 81 insertions(+), 1 deletion(-)
diff --git a/mlir/test/Integration/Dialect/Vector/CPU/ArmSVE/test-contraction.mlir b/mlir/test/Integration/Dialect/Vector/CPU/ArmSVE/test-contraction.mlir
index 9bc6eab39d69fd9..5810e1136f0d9c7 100644
--- a/mlir/test/Integration/Dialect/Vector/CPU/ArmSVE/test-contraction.mlir
+++ b/mlir/test/Integration/Dialect/Vector/CPU/ArmSVE/test-contraction.mlir
@@ -13,12 +13,38 @@
// REDEFINE: %{entry} = matmul_f32
// RUN: %{run} | FileCheck %s --check-prefix=F32
+// REDEFINE: %{entry} = dot_product_i32
+// RUN: %{run} | FileCheck %s --check-prefix=DP
+
+// REDEFINE: %{entry} = matvec_i32
+// RUN: %{run} | FileCheck %s --check-prefix=MV
+
// NOTE: These tests are meant to complement the integration tests from:
// * ../test-contraction.mlir
// (tests with fixed width vectors). Rather than duplicating those tests, this
// file focuses on excercissing scalable vectors in a few most common cases.
-// TODO: Masks + matvec + dot product
+// TODO: Masks
+
+#dotp_accesses = [
+ affine_map<(i) -> (i)>,
+ affine_map<(i) -> (i)>,
+ affine_map<(i) -> ()>
+]
+#dotp_trait = {
+ indexing_maps = #dotp_accesses,
+ iterator_types = ["reduction"]
+}
+
+#matvec_accesses = [
+ affine_map<(i, j) -> (i, j)>,
+ affine_map<(i, j) -> (j)>,
+ affine_map<(i, j) -> (i)>
+]
+#matvec_trait = {
+ indexing_maps = #matvec_accesses,
+ iterator_types = ["parallel", "reduction"]
+}
#matmat_accesses = [
affine_map<(i, j, k) -> (i, k)>,
@@ -30,6 +56,60 @@
iterator_types = ["parallel", "parallel", "reduction"]
}
+// Contraction: dot-product a x b.
+func.func @dot_product_i32() {
+ %acc = arith.constant 0: i32
+
+ %vector_a = arith.constant dense<123> : vector<[4]xi32>
+ %vector_b = arith.constant dense<314> : vector<[4]xi32>
+ %vector_c = arith.constant dense<0> : vector<[4]xi32>
+
+ // The result of this dot-product will depend
+ // on the vector length, so we are unable to verify it.
+ %dp1 = vector.contract #dotp_trait %vector_a, %vector_b, %acc
+ : vector<[4]xi32>, vector<[4]xi32> into i32
+ // DP: {{[0-9]*}}
+ vector.print %dp1 : i32
+
+ // The result of this dot-product should be 0.
+ %dp2 = vector.contract #dotp_trait %vector_a, %vector_c, %acc
+ : vector<[4]xi32>, vector<[4]xi32> into i32
+ // DP: 0
+ vector.print %dp2 : i32
+
+ // DP: SVE: END OF TEST OUTPUT
+ vector.print str "SVE: END OF TEST OUTPUT"
+
+ return
+}
+
+// Contraction: matrix-vector A x c
+func.func @matvec_i32() {
+ %acc = arith.constant dense<0>: vector<3xi32>
+
+ %vector_a = arith.constant dense<123> : vector<3x[4]xi32>
+ %vector_b = arith.constant dense<314> : vector<[4]xi32>
+ %vector_c = arith.constant dense<0> : vector<[4]xi32>
+
+ // The result of this matvec will depend on the vector length, so we are
+ // unable to verify it.
+ %dp1 = vector.contract #matvec_trait %vector_a, %vector_b, %acc
+ : vector<3x[4]xi32>, vector<[4]xi32> into vector<3xi32>
+ // MV: {{[0-9]*}}, {{[0-9]*}}, {{[0-9]*}}
+ vector.print %dp1 : vector<3xi32>
+
+ // The result of this matvc should be a vector of 0s.
+ %dp2 = vector.contract #matvec_trait %vector_a, %vector_c, %acc
+ : vector<3x[4]xi32>, vector<[4]xi32> into vector<3xi32>
+ // MV: 0, 0, 0
+ vector.print %dp2 : vector<3xi32>
+
+ // MV: SVE: END OF TEST OUTPUT
+ vector.print str "SVE: END OF TEST OUTPUT"
+
+ return
+}
+
func.func @matmul_i32() {
// Setup vector A:
%vector_a = arith.constant dense<123> : vector<3x5xi32>
>From 2558a030521f7bc07f445d84a5fe4083de04f4b4 Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Fri, 27 Oct 2023 10:04:14 +0000
Subject: [PATCH 2/3] fixup! [mlir][SVE] Add more e2e test for vector.contract
Update the dot product teset to check the actual result (thanks Ben)
---
.../Dialect/Vector/CPU/ArmSVE/test-contraction.mlir | 9 +++++++--
1 file changed, 7 insertions(+), 2 deletions(-)
diff --git a/mlir/test/Integration/Dialect/Vector/CPU/ArmSVE/test-contraction.mlir b/mlir/test/Integration/Dialect/Vector/CPU/ArmSVE/test-contraction.mlir
index 5810e1136f0d9c7..2f14bf147de91a1 100644
--- a/mlir/test/Integration/Dialect/Vector/CPU/ArmSVE/test-contraction.mlir
+++ b/mlir/test/Integration/Dialect/Vector/CPU/ArmSVE/test-contraction.mlir
@@ -68,8 +68,13 @@ func.func @dot_product_i32() {
// on the vector length, so we are unable to verify it.
%dp1 = vector.contract #dotp_trait %vector_a, %vector_b, %acc
: vector<[4]xi32>, vector<[4]xi32> into i32
- // DP: {{[0-9]*}}
- vector.print %dp1 : i32
+ // Dot product should be (123 * 314) * 4 * vscale, so ...
+ %vscale = vector.vscale
+ %vscale_i32 = arith.index_cast %vscale : index to i32
+ %dp1_divvl = arith.divui %dp1, %vscale_i32 : i32
+ // ... %dp/%vscale = 123 * 314 * 4 = 154488
+ // DP: 154488
+ vector.print %dp1_divvl : i32
// The result of this dot-product should be 0.
%dp2 = vector.contract #dotp_trait %vector_a, %vector_c, %acc
>From 0ac670719bdffb605a29dbb3a0e04c0d989fa297 Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Fri, 27 Oct 2023 11:49:26 +0000
Subject: [PATCH 3/3] fixup! [mlir][SVE] Add more e2e test for vector.contract
Further generalization
---
.../Vector/CPU/ArmSVE/test-contraction.mlir | 41 ++++++++++++-------
1 file changed, 26 insertions(+), 15 deletions(-)
diff --git a/mlir/test/Integration/Dialect/Vector/CPU/ArmSVE/test-contraction.mlir b/mlir/test/Integration/Dialect/Vector/CPU/ArmSVE/test-contraction.mlir
index 2f14bf147de91a1..f4f3bad504915a4 100644
--- a/mlir/test/Integration/Dialect/Vector/CPU/ArmSVE/test-contraction.mlir
+++ b/mlir/test/Integration/Dialect/Vector/CPU/ArmSVE/test-contraction.mlir
@@ -64,18 +64,20 @@ func.func @dot_product_i32() {
%vector_b = arith.constant dense<314> : vector<[4]xi32>
%vector_c = arith.constant dense<0> : vector<[4]xi32>
- // The result of this dot-product will depend
- // on the vector length, so we are unable to verify it.
+ // DOT PRODUCT 1
%dp1 = vector.contract #dotp_trait %vector_a, %vector_b, %acc
: vector<[4]xi32>, vector<[4]xi32> into i32
- // Dot product should be (123 * 314) * 4 * vscale, so ...
+ // Dot product should be:
+ // * val = (123 * 314) * 4 * vscale,
+ // so ...
%vscale = vector.vscale
%vscale_i32 = arith.index_cast %vscale : index to i32
- %dp1_divvl = arith.divui %dp1, %vscale_i32 : i32
- // ... %dp/%vscale = 123 * 314 * 4 = 154488
+ %dp1_div = arith.divui %dp1, %vscale_i32 : i32
+ // ... val / %vscale = 123 * 314 * 4 = 154488
// DP: 154488
- vector.print %dp1_divvl : i32
+ vector.print %dp1_div : i32
+ // DOT PRODUCT 2
// The result of this dot-product should be 0.
%dp2 = vector.contract #dotp_trait %vector_a, %vector_c, %acc
: vector<[4]xi32>, vector<[4]xi32> into i32
@@ -96,18 +98,27 @@ func.func @matvec_i32() {
%vector_b = arith.constant dense<314> : vector<[4]xi32>
%vector_c = arith.constant dense<0> : vector<[4]xi32>
- // The result of this matvec will depend on the vector length, so we are
- // unable to verify it.
- %dp1 = vector.contract #matvec_trait %vector_a, %vector_b, %acc
+ // MATVEC 1
+ %mv1 = vector.contract #matvec_trait %vector_a, %vector_b, %acc
: vector<3x[4]xi32>, vector<[4]xi32> into vector<3xi32>
- // MV: {{[0-9]*}}, {{[0-9]*}}, {{[0-9]*}}
- vector.print %dp1 : vector<3xi32>
-
- // The result of this matvc should be a vector of 0s.
- %dp2 = vector.contract #matvec_trait %vector_a, %vector_c, %acc
+ // Every element in the output vector is a result of a dot product, for
+ // which:
+ // val = (123 * 314) * 4 * vscale
+ // so ...
+ %vscale = vector.vscale
+ %vscale_v = vector.splat %vscale : vector<3xindex>
+ %vscale_i32 = arith.index_cast %vscale_v : vector<3xindex> to vector<3xi32>
+ %dp1_div = arith.divui %mv1, %vscale_i32 : vector<3xi32>
+ // ... val/%vscale = 123 * 314 * 4 = 154488
+ // MV: 154488, 154488, 154488
+ vector.print %mv1_div : vector<3xi32>
+
+ // MATVEC 2
+ // The result of this matvec should be a vector of 0s.
+ %mv2 = vector.contract #matvec_trait %vector_a, %vector_c, %acc
: vector<3x[4]xi32>, vector<[4]xi32> into vector<3xi32>
// MV: 0, 0, 0
- vector.print %dp2 : vector<3xi32>
+ vector.print %mv2 : vector<3xi32>
// MV: SVE: END OF TEST OUTPUT
vector.print str "SVE: END OF TEST OUTPUT"
More information about the Mlir-commits
mailing list