From 7a0d0c6ffc2a7915ce882adc6b39e458f65bc1c2 Mon Sep 17 00:00:00 2001 From: Andrzej Warzynski 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 9bc6eab39d69fd..5810e1136f0d9c 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 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 5810e1136f0d9c..2f14bf147de91a 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 e80df2067197a6c5759f26ab6dfb99db0c172f6d Mon Sep 17 00:00:00 2001 From: Andrzej Warzynski 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 2f14bf147de91a..d86ff56d79e33c 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> + %mv1_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"