178 lines
8.0 KiB
C++
178 lines
8.0 KiB
C++
// This file is part of Eigen, a lightweight C++ template library
|
|
// for linear algebra.
|
|
//
|
|
// This Source Code Form is subject to the terms of the Mozilla
|
|
// Public License v. 2.0. If a copy of the MPL was not distributed
|
|
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
|
|
// SPDX-FileCopyrightText: The Eigen Authors
|
|
// SPDX-License-Identifier: MPL-2.0
|
|
|
|
// SME GEMM kernel tests.
|
|
// Requires compiler flags: -march=armv9.2-a+sme2 and -DEIGEN_ARM64_USE_SME.
|
|
|
|
#include "product.h"
|
|
|
|
// Without the right -march flags, __ARM_FEATURE_SME is undefined and
|
|
// EIGEN_VECTORIZE_SME never fires - the test would silently compile
|
|
// against the NEON GEBP kernel and pass, making this a useless no-op.
|
|
// Fail the build instead.
|
|
#if !defined(EIGEN_VECTORIZE_SME)
|
|
#error \
|
|
"product_sme requires the SME backend. Build with -march=armv9.2-a+sme2 " \
|
|
"-DEIGEN_ARM64_USE_SME (see -DEIGEN_TEST_SME=ON in test/CMakeLists.txt for " \
|
|
"the typical CMake invocation)."
|
|
#endif
|
|
|
|
using SmeColMajorMatF = Matrix<float, Dynamic, Dynamic, ColMajor>;
|
|
using SmeRowMajorMatF = Matrix<float, Dynamic, Dynamic, RowMajor>;
|
|
using SmeColMajorStridedMatF = Map<SmeColMajorMatF, 0, Stride<Dynamic, Dynamic>>;
|
|
using SmeRowMajorStridedMatF = Map<SmeRowMajorMatF, 0, Stride<Dynamic, Dynamic>>;
|
|
|
|
template <typename InputMat, typename ResultMat, typename ResultMap>
|
|
static void verify_strided_result(int n, ResultMat& storage, const Stride<Dynamic, Dynamic>& stride) {
|
|
InputMat A = InputMat::Random(n, n);
|
|
InputMat B = InputMat::Random(n, n);
|
|
ResultMap C(storage.data(), n, n, stride);
|
|
C = ResultMat::Random(n, n);
|
|
ResultMat c_before = C.eval();
|
|
|
|
C.noalias() += A * B;
|
|
|
|
ResultMat ref = c_before + (A.lazyProduct(B)).eval();
|
|
ResultMat got = C;
|
|
VERIFY_IS_APPROX(got, ref);
|
|
}
|
|
|
|
template <typename InputMat>
|
|
static void test_general_strided_result(int n) {
|
|
// General-stride C path: InputMat selects the source packers, while both C
|
|
// strides are non-unit so sme_store_za_tile uses scalar scatter.
|
|
SmeColMajorMatF storage = SmeColMajorMatF::Zero(2 * n, n);
|
|
verify_strided_result<InputMat, SmeColMajorMatF, SmeColMajorStridedMatF>(
|
|
n, storage, Stride<Dynamic, Dynamic>(/*outer=*/2 * n, /*inner=*/2));
|
|
|
|
// Padding rows skipped by the strided Map should not be touched.
|
|
for (int i = 0; i < n; ++i) {
|
|
for (int j = 0; j < n; ++j) {
|
|
VERIFY(storage(2 * i + 1, j) == float(0));
|
|
}
|
|
}
|
|
}
|
|
|
|
static void test_rowmajor_strided_result(int n) {
|
|
// RowMajor C path: inner stride is one, with padded columns after the Map.
|
|
SmeRowMajorMatF storage = SmeRowMajorMatF::Zero(n, 2 * n);
|
|
verify_strided_result<SmeRowMajorMatF, SmeRowMajorMatF, SmeRowMajorStridedMatF>(
|
|
n, storage, Stride<Dynamic, Dynamic>(/*outer=*/2 * n, /*inner=*/1));
|
|
|
|
// Padding columns skipped by the strided Map should not be touched.
|
|
for (int i = 0; i < n; ++i) {
|
|
for (int j = n; j < 2 * n; ++j) {
|
|
VERIFY(storage(i, j) == float(0));
|
|
}
|
|
}
|
|
}
|
|
|
|
static void test_deep_k_split() {
|
|
constexpr int rows = 64;
|
|
constexpr int depth = 2050;
|
|
constexpr int cols = 64;
|
|
SmeColMajorMatF A = SmeColMajorMatF::Random(rows, depth);
|
|
SmeColMajorMatF B = SmeColMajorMatF::Random(depth, cols);
|
|
SmeColMajorMatF C = SmeColMajorMatF::Random(rows, cols);
|
|
SmeColMajorMatF c_before = C;
|
|
|
|
C.noalias() += A * B;
|
|
|
|
VERIFY_IS_APPROX(C, c_before + (A.lazyProduct(B)).eval());
|
|
}
|
|
|
|
EIGEN_DECLARE_TEST(product_sme) {
|
|
// Square edge cases around the block and tile boundaries (the block is
|
|
// kSmeMr x kSmeNr and a ZA tile is svlw x svlw, so the sizes below land
|
|
// just on/off the intra-tile splits and the block tails at SVL=512).
|
|
CALL_SUBTEST_1(product(Matrix<float, Dynamic, Dynamic>(1, 1)));
|
|
CALL_SUBTEST_1(product(Matrix<float, Dynamic, Dynamic>(15, 15)));
|
|
CALL_SUBTEST_1(product(Matrix<float, Dynamic, Dynamic>(16, 16)));
|
|
CALL_SUBTEST_1(product(Matrix<float, Dynamic, Dynamic>(17, 17)));
|
|
CALL_SUBTEST_1(product(Matrix<float, Dynamic, Dynamic>(31, 31)));
|
|
CALL_SUBTEST_1(product(Matrix<float, Dynamic, Dynamic>(33, 33)));
|
|
CALL_SUBTEST_1(product(Matrix<float, Dynamic, Dynamic>(63, 63)));
|
|
CALL_SUBTEST_1(product(Matrix<float, Dynamic, Dynamic>(64, 64)));
|
|
CALL_SUBTEST_1(product(Matrix<float, Dynamic, Dynamic>(65, 65)));
|
|
|
|
// Thin / wide rectangular cases (M x 1, 1 x N)
|
|
CALL_SUBTEST_1(product(Matrix<float, Dynamic, Dynamic>(32, 1)));
|
|
CALL_SUBTEST_1(product(Matrix<float, Dynamic, Dynamic>(1, 32)));
|
|
CALL_SUBTEST_1(product(Matrix<float, Dynamic, Dynamic>(1, 64)));
|
|
CALL_SUBTEST_1(product(Matrix<float, Dynamic, Dynamic>(64, 1)));
|
|
|
|
// Non-float scalar smoke tests: SME only specializes fp32, so these prove
|
|
// unsupported scalar types still route through the generic product path.
|
|
CALL_SUBTEST_2(product(Matrix<double, Dynamic, Dynamic>(33, 17)));
|
|
CALL_SUBTEST_3(product(Matrix<std::complex<float>, Dynamic, Dynamic>(33, 17)));
|
|
|
|
// Non-square cases that exercise tail paths for both M and N
|
|
CALL_SUBTEST_1(product(Matrix<float, Dynamic, Dynamic>(17, 65)));
|
|
CALL_SUBTEST_1(product(Matrix<float, Dynamic, Dynamic>(65, 17)));
|
|
CALL_SUBTEST_1(product(Matrix<float, Dynamic, Dynamic>(15, 63)));
|
|
CALL_SUBTEST_1(product(Matrix<float, Dynamic, Dynamic>(33, 7)));
|
|
CALL_SUBTEST_1(product(Matrix<float, Dynamic, Dynamic>(7, 33)));
|
|
CALL_SUBTEST_1(product(Matrix<float, Dynamic, Dynamic>(128, 3)));
|
|
CALL_SUBTEST_1(product(Matrix<float, Dynamic, Dynamic>(3, 128)));
|
|
|
|
// Exercise the kc split path just above the SME blocking heuristic's depth
|
|
// cap (sme_max_kc in GeneralBlockPanelKernel.h).
|
|
test_deep_k_split();
|
|
|
|
// Random sizes
|
|
for (int i = 0; i < g_repeat; i++) {
|
|
CALL_SUBTEST_1(product(Matrix<float, Dynamic, Dynamic>(internal::random<int>(1, EIGEN_TEST_MAX_SIZE),
|
|
internal::random<int>(1, EIGEN_TEST_MAX_SIZE))));
|
|
}
|
|
|
|
// Exercise the RowMajor packers and RowMajor result path. When the input
|
|
// MatrixType is RowMajor, product() instantiates m1/m2/m3/res in RowMajor,
|
|
// so every matrix-matrix product in the suite flows through:
|
|
// - the RowMajor LHS packer (gemm_pack_lhs<..., RowMajor>)
|
|
// - the RowMajor RHS packer (gemm_pack_rhs<..., RowMajor>)
|
|
// - the RowMajor-C dispatch in GeneralMatrixMatrix.h (which transposes
|
|
// the computation: C^T = B^T * A^T).
|
|
CALL_SUBTEST_1(product(SmeRowMajorMatF(15, 15)));
|
|
CALL_SUBTEST_1(product(SmeRowMajorMatF(16, 16)));
|
|
CALL_SUBTEST_1(product(SmeRowMajorMatF(17, 17)));
|
|
CALL_SUBTEST_1(product(SmeRowMajorMatF(31, 31)));
|
|
CALL_SUBTEST_1(product(SmeRowMajorMatF(32, 32)));
|
|
CALL_SUBTEST_1(product(SmeRowMajorMatF(33, 33)));
|
|
CALL_SUBTEST_1(product(SmeRowMajorMatF(64, 64)));
|
|
CALL_SUBTEST_1(product(SmeRowMajorMatF(65, 65)));
|
|
CALL_SUBTEST_1(product(SmeRowMajorMatF(17, 65)));
|
|
CALL_SUBTEST_1(product(SmeRowMajorMatF(65, 17)));
|
|
for (int i = 0; i < g_repeat; i++) {
|
|
CALL_SUBTEST_1(product(
|
|
SmeRowMajorMatF(internal::random<int>(1, EIGEN_TEST_MAX_SIZE), internal::random<int>(1, EIGEN_TEST_MAX_SIZE))));
|
|
}
|
|
|
|
// Exercise the general-stride branch of sme_store_za_tile: fires when both
|
|
// C_stride_row != 1 and C_stride_col != 1, e.g. a Map<Matrix> with an
|
|
// explicit non-unit inner stride. product.h never builds such a result, so
|
|
// without this subtest the scalar-scatter path is effectively untested.
|
|
for (int n : {15, 16, 17, 31, 32, 33, 63, 64, 65}) {
|
|
test_general_strided_result<SmeColMajorMatF>(n);
|
|
test_general_strided_result<SmeRowMajorMatF>(n);
|
|
test_rowmajor_strided_result(n);
|
|
}
|
|
|
|
// Row-LHS x Row-RHS -> Col-C: the one LHS/RHS/C storage combination that
|
|
// product.h's transpose-style expressions never build directly (it always
|
|
// flips one side of the multiplication). The code paths are the same as
|
|
// other combinations via Eigen's dispatch, but exercise them explicitly.
|
|
for (int n : {15, 16, 17, 31, 32, 33, 63, 64, 65}) {
|
|
Matrix<float, Dynamic, Dynamic, RowMajor> A = Matrix<float, Dynamic, Dynamic, RowMajor>::Random(n, n);
|
|
Matrix<float, Dynamic, Dynamic, RowMajor> B = Matrix<float, Dynamic, Dynamic, RowMajor>::Random(n, n);
|
|
SmeColMajorMatF C = SmeColMajorMatF::Zero(n, n);
|
|
C.noalias() += A * B;
|
|
VERIFY_IS_APPROX(C, (A.lazyProduct(B)).eval());
|
|
}
|
|
}
|