blob: cb8d41cee8764c33954f9ea8a80e77994980a788 [file]
// 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 a compiler targeting SME2, e.g. -march=armv9.2-a+sme2, which selects
// the backend on its own; -DEIGEN_ARM64_USE_SME then makes a toolchain that
// cannot provide it fail the build rather than fall back to NEON.
// double and complex<double> additionally need FEAT_SME_F64F64 (+sme-f64f64, or
// a -mcpu that implies it); without it EIGEN_VECTORIZE_SME_F64F64 is undefined
// and they keep the generic kernel, so their subtests pack their cases through
// that path instead.
#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 " \
"(see -DEIGEN_TEST_SME=ON in test/CMakeLists.txt for the typical CMake " \
"invocation)."
#endif
// The backend is chosen from the compiler's target features; EIGEN_ARM64_USE_SME,
// which this test's CMake option passes, only turns a toolchain that cannot
// provide it into a build error. Assert that here, so a regression making
// selection depend on the opt-in again is caught rather than masked by it.
#if !defined(EIGEN_ARM64_SME_SELECTED)
#error "the SME backend was forced by EIGEN_ARM64_USE_SME rather than selected from the compiler's target features."
#endif
template <typename Scalar>
using SmeColMajorMat = Matrix<Scalar, Dynamic, Dynamic, ColMajor>;
template <typename Scalar>
using SmeRowMajorMat = Matrix<Scalar, Dynamic, Dynamic, RowMajor>;
template <typename Scalar>
using SmeVector = Matrix<Scalar, Dynamic, 1>;
template <typename Scalar>
using SmeColMajorStridedMat = Map<SmeColMajorMat<Scalar>, 0, Stride<Dynamic, Dynamic>>;
template <typename Scalar>
using SmeRowMajorStridedMat = Map<SmeRowMajorMat<Scalar>, 0, Stride<Dynamic, Dynamic>>;
// The logical micro-kernel block for Scalar (kSmeMr/kSmeNr for float and so
// on). Sizes below are expressed in terms of it so each scalar sweeps its own
// block boundaries; sme_tile() is the ZA tile side at the SVL=512 design point,
// which is where the intra-block splits fall. Complex scalars pair two tiles
// per accumulator, so their block is not square.
template <typename Scalar>
static constexpr int sme_mr() {
return internal::sme_block<Scalar>::mr;
}
template <typename Scalar>
static constexpr int sme_nr() {
return internal::sme_block<Scalar>::nr;
}
template <typename Scalar>
static constexpr int sme_tile() {
return internal::sme_block<Scalar>::mr / internal::sme_block<Scalar>::kGridRows;
}
// svwhilelt_b* is overloaded only on the four fixed-width types, so forwarding
// the caller's own type would leave Index ambiguous wherever it is `long` while
// int64_t is `long long`, and deducing one type from both arguments would reject
// the literal-plus-Index bounds the packers pass. Name every spelling Index
// takes, mix the two argument types, and check the predicate so the traits'
// widening cannot change which lanes are active.
template <typename Scalar, typename Begin, typename End = Begin>
__arm_locally_streaming static bool sme_whilelt_covers_first_lane_only() {
using Traits = internal::sme_traits<Scalar>;
const svbool_t pg = Traits::whilelt(Begin(0), End(1));
return svptest_first(Traits::ptrue(), pg) && svcntp_b8(Traits::ptrue(), pg) == 1;
}
template <typename Scalar>
static void test_whilelt_operand_types() {
VERIFY((sme_whilelt_covers_first_lane_only<Scalar, int>()));
VERIFY((sme_whilelt_covers_first_lane_only<Scalar, long>()));
VERIFY((sme_whilelt_covers_first_lane_only<Scalar, long long>()));
VERIFY((sme_whilelt_covers_first_lane_only<Scalar, Index>()));
VERIFY((sme_whilelt_covers_first_lane_only<Scalar, int, Index>()));
}
// Write one element into a packed panel of width w, in the layout the SME
// kernel reads: a real scalar lands at dst[k*w + r], a complex one splits into
// the depth step's real and imaginary halves.
template <typename Scalar>
static void set_packed(Scalar* panel, Index w, Index k, Index r, const Scalar& v) {
panel[k * w + r] = v;
}
template <typename RealScalar>
static void set_packed(std::complex<RealScalar>* panel, Index w, Index k, Index r, const std::complex<RealScalar>& v) {
RealScalar* p = reinterpret_cast<RealScalar*>(panel + k * w);
p[r] = numext::real(v);
p[w + r] = numext::imag(v);
}
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 Scalar, 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.
SmeColMajorMat<Scalar> storage = SmeColMajorMat<Scalar>::Zero(2 * n, n);
verify_strided_result<InputMat, SmeColMajorMat<Scalar>, SmeColMajorStridedMat<Scalar>>(
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) == Scalar(0));
}
}
}
template <typename Scalar>
static void test_rowmajor_strided_result(int n) {
// RowMajor C path: inner stride is one, with padded columns after the Map.
SmeRowMajorMat<Scalar> storage = SmeRowMajorMat<Scalar>::Zero(n, 2 * n);
verify_strided_result<SmeRowMajorMat<Scalar>, SmeRowMajorMat<Scalar>, SmeRowMajorStridedMat<Scalar>>(
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) == Scalar(0));
}
}
}
// A non-trivial alpha: complex where the scalar is, so the store's rotation of
// the accumulator is exercised rather than a plain rescale.
template <typename Scalar>
struct nontrivial_alpha_impl {
static Scalar run() { return Scalar(1.375); }
};
template <typename RealScalar>
struct nontrivial_alpha_impl<std::complex<RealScalar>> {
static std::complex<RealScalar> run() { return std::complex<RealScalar>(RealScalar(1.375), RealScalar(-0.625)); }
};
// Conjugated operands and a non-real alpha.
//
// blas_traits folds .adjoint()/.conjugate() and a scalar factor into the
// gebp_kernel's ConjugateLhs/ConjugateRhs and alpha rather than materializing a
// temporary, so these expressions are the only thing that reaches the complex
// kernel's FMOPA/FMOPS sign choices and its alpha-scaled store. The references
// materialize the conjugation first and multiply coefficient-wise, so they do
// not share a code path with what they check.
template <typename Scalar, typename LhsMat, typename RhsMat, typename ResMat>
static void verify_conjugated_products(int n) {
const LhsMat A = LhsMat::Random(n, n);
const RhsMat B = RhsMat::Random(n, n);
const SmeColMajorMat<Scalar> Ap = A, Bp = B;
const SmeColMajorMat<Scalar> Aa = Ap.adjoint().eval(), Ac = Ap.conjugate().eval();
const SmeColMajorMat<Scalar> Ba = Bp.adjoint().eval(), Bc = Bp.conjugate().eval();
const Scalar alpha = nontrivial_alpha_impl<Scalar>::run();
ResMat C(n, n);
C.setZero();
C.noalias() += A.adjoint() * B;
VERIFY_IS_APPROX(SmeColMajorMat<Scalar>(C), SmeColMajorMat<Scalar>(Aa.lazyProduct(Bp)));
C.setZero();
C.noalias() += A * B.adjoint();
VERIFY_IS_APPROX(SmeColMajorMat<Scalar>(C), SmeColMajorMat<Scalar>(Ap.lazyProduct(Ba)));
C.setZero();
C.noalias() += A.adjoint() * B.adjoint();
VERIFY_IS_APPROX(SmeColMajorMat<Scalar>(C), SmeColMajorMat<Scalar>(Aa.lazyProduct(Ba)));
C.setZero();
C.noalias() += A.conjugate() * B;
VERIFY_IS_APPROX(SmeColMajorMat<Scalar>(C), SmeColMajorMat<Scalar>(Ac.lazyProduct(Bp)));
C.setZero();
C.noalias() += A * B.conjugate();
VERIFY_IS_APPROX(SmeColMajorMat<Scalar>(C), SmeColMajorMat<Scalar>(Ap.lazyProduct(Bc)));
C.setZero();
C.noalias() += A.conjugate() * B.conjugate();
VERIFY_IS_APPROX(SmeColMajorMat<Scalar>(C), SmeColMajorMat<Scalar>(Ac.lazyProduct(Bc)));
C.setZero();
C.noalias() += alpha * (A.adjoint() * B);
VERIFY_IS_APPROX(SmeColMajorMat<Scalar>(C), SmeColMajorMat<Scalar>(alpha * Aa.lazyProduct(Bp)));
}
// Exercise the kc split path just above the SME blocking heuristic's depth cap
// (sme_max_kc in GeneralBlockPanelKernel.h, scaled by the scalar width).
template <typename Scalar>
static void test_deep_k_split() {
constexpr int rows = 64;
const int depth = int(2 * (Index(EIGEN_SME_MAX_KC) * Index(sizeof(float)) / Index(sizeof(Scalar)))) + 2;
constexpr int cols = 64;
SmeColMajorMat<Scalar> A = SmeColMajorMat<Scalar>::Random(rows, depth);
SmeColMajorMat<Scalar> B = SmeColMajorMat<Scalar>::Random(depth, cols);
SmeColMajorMat<Scalar> C = SmeColMajorMat<Scalar>::Random(rows, cols);
SmeColMajorMat<Scalar> c_before = C;
C.noalias() += A * B;
VERIFY_IS_APPROX(C, c_before + (A.lazyProduct(B)).eval());
}
// ---------------------------------------------------------------------------
// Raw packed-buffer tests.
//
// The product tests above validate the packers only transitively through a full
// product, where the gebp_kernel can mask a mispack, or where a bug only shows
// at a specific SVL/region boundary. The tests below call the SME packers
// directly and compare the packed buffer exactly against a scalar reference
// (equality is exact). This pins every region -- in particular a dropped
// row-group in the two-pass trailing transpose, and a SYMM packer that copies
// the unused triangle instead of mirroring -- at whatever SVL the run uses.
// ---------------------------------------------------------------------------
// A distinctive marker for buffer cells the packer must leave untouched, and
// for the unused triangle of a lower-triangular operand. Random values live in
// [-1, 1], so it never collides with a real packed value. The complex marker's
// imaginary part differs from its real one, so a packer that fills only the
// real half of a split depth step is caught.
template <typename Scalar>
struct pack_sentinel_impl {
static Scalar run() { return Scalar(98765); }
};
template <typename RealScalar>
struct pack_sentinel_impl<std::complex<RealScalar>> {
static std::complex<RealScalar> run() { return std::complex<RealScalar>(RealScalar(98765), RealScalar(-54321)); }
};
template <typename Scalar>
static Scalar pack_sentinel() {
return pack_sentinel_impl<Scalar>::run();
}
// Lower-triangular n x n operand plus the dense selfadjoint reference the packer
// must emit. The unused triangle is filled with the sentinel so a packer that
// copies the dense matrix and never mirrors fails VERIFY_IS_EQUAL.
// product_selfadjoint_matrix stores the valid triangle where row >= col
// (after the Upper/RowMajor xor), so the packer must read stored(row,col)
// below the diagonal and conj(stored(col,row)) above it. Averaging with the
// adjoint cancels the diagonal's imaginary part exactly, which is what a
// selfadjoint view defines it to be.
template <typename Scalar, int StorageOrder>
static void make_lower_stored_selfadjoint(Index n, Matrix<Scalar, Dynamic, Dynamic, StorageOrder>& stored,
Matrix<Scalar, Dynamic, Dynamic, StorageOrder>& full) {
full = Matrix<Scalar, Dynamic, Dynamic, StorageOrder>::Random(n, n);
full = ((full + full.adjoint()) * Scalar(0.5)).eval();
stored = Matrix<Scalar, Dynamic, Dynamic, StorageOrder>::Constant(n, n, pack_sentinel<Scalar>());
for (Index i = 0; i < n; ++i)
for (Index j = 0; j <= i; ++j) stored(i, j) = full(i, j);
}
// LHS SYMM packer: a square selfadjoint diagonal block of size n, packed into
// uniform mr-wide depth-major panels. Reference: full(i+r, k).
template <typename Scalar, int StorageOrder>
static void verify_symm_pack_lhs(Index n) {
const Index MR = sme_mr<Scalar>();
Matrix<Scalar, Dynamic, Dynamic, StorageOrder> stored, full;
make_lower_stored_selfadjoint<Scalar, StorageOrder>(n, stored, full);
SmeVector<Scalar> packed = SmeVector<Scalar>::Constant(n * n, pack_sentinel<Scalar>());
SmeVector<Scalar> ref = SmeVector<Scalar>::Constant(n * n, pack_sentinel<Scalar>());
for (Index i = 0; i < n; i += MR) {
const Index w = numext::mini(MR, n - i);
for (Index k = 0; k < n; ++k)
for (Index r = 0; r < w; ++r) set_packed(ref.data() + i * n, w, k, r, Scalar(full(i + r, k)));
}
internal::symm_pack_lhs<Scalar, Index, sme_mr<Scalar>(), 1, StorageOrder> pack;
pack(packed.data(), stored.data(), stored.outerStride(), /*cols(depth)=*/n, /*rows=*/n);
VERIFY_IS_EQUAL(packed, ref);
}
// RHS SYMM packer: a depth block [k2, k2 + depth) x cols columns of an N x N
// selfadjoint matrix, packed into nr-wide depth-major panels. Reference:
// full(k2 + k, j + c). A k2 > 0 offset makes the transposed region non-empty,
// so partial-width panels reach the two-pass transpose.
template <typename Scalar, int StorageOrder>
static void verify_symm_pack_rhs(Index N, Index depth, Index cols, Index k2) {
eigen_assert(k2 + depth <= N && cols <= N);
const Index NR = sme_nr<Scalar>();
Matrix<Scalar, Dynamic, Dynamic, StorageOrder> stored, full;
make_lower_stored_selfadjoint<Scalar, StorageOrder>(N, stored, full);
SmeVector<Scalar> packed = SmeVector<Scalar>::Constant(cols * depth, pack_sentinel<Scalar>());
SmeVector<Scalar> ref = SmeVector<Scalar>::Constant(cols * depth, pack_sentinel<Scalar>());
for (Index j = 0; j < cols; j += NR) {
const Index w = numext::mini(NR, cols - j);
for (Index k = 0; k < depth; ++k)
for (Index c = 0; c < w; ++c) set_packed(ref.data() + j * depth, w, k, c, Scalar(full(k2 + k, j + c)));
}
internal::symm_pack_rhs<Scalar, Index, sme_nr<Scalar>(), StorageOrder> pack;
pack(packed.data(), stored.data(), stored.outerStride(), /*rows(depth)=*/depth, /*cols=*/cols, k2);
VERIFY_IS_EQUAL(packed, ref);
}
template <typename Scalar>
static void test_symm_pack() {
// The last panel width sweeps a range of partial widths; at each SVL the
// two-pass trailing transpose (the if->loop fix) fires when a partial width
// leaves a trailing row-group remainder in (svl, 2*svl). The spread below
// hits that for svl in {2, 4, 8, 16, 32, 64} -- fp32 SVL 128..2048 and the
// fp64 lane counts, which are half of those.
const int sizes[] = {1, 5, 7, 17, 31, 32, 33, 37, 39, 45, 48, 49, 55, 57, 63, 64, 65, 79, 96, 97};
for (int n : sizes) {
verify_symm_pack_lhs<Scalar, ColMajor>(n);
verify_symm_pack_lhs<Scalar, RowMajor>(n);
// RHS, single depth block anchored at the diagonal (k2 == 0).
verify_symm_pack_rhs<Scalar, ColMajor>(n, n, n, 0);
verify_symm_pack_rhs<Scalar, RowMajor>(n, n, n, 0);
}
// RHS depth blocks offset from the diagonal (k2 > 0): the transposed region is
// non-empty, so the RowMajor operand drives partial-width panels through the
// two-pass transpose and the ColMajor operand through the partial copy.
struct RhsCase {
int N, depth, cols, k2;
};
const RhsCase rhs_cases[] = {
{100, 32, 39, 16}, {100, 24, 39, 32}, {100, 40, 64, 8}, {100, 39, 39, 33}, {128, 57, 57, 40}, {128, 33, 45, 60},
};
for (const RhsCase& c : rhs_cases) {
verify_symm_pack_rhs<Scalar, ColMajor>(c.N, c.depth, c.cols, c.k2);
verify_symm_pack_rhs<Scalar, RowMajor>(c.N, c.depth, c.cols, c.k2);
}
}
// ---------------------------------------------------------------------------
// Mapper-based packing fallback (sme_pack_{lhs,rhs}_fallback).
//
// Taken by by-value tensor sub-mappers and inner-strided blas mappers, which
// the raw pointer + stride packers cannot walk. The product suite reaches these
// only through tensor contractions (nightly, SVL=512) and TriangularSolver
// (random sizes), so here we drive the fallback directly with both mapper
// families and compare the packed buffer against a scalar reference.
// ---------------------------------------------------------------------------
// Minimal stand-ins for by-value sub-mappers. ColMajor packets advance the
// first index, while RowMajor packets follow the normal storage-inner second
// index. operator() returns by value so both take the no-direct-access dispatch.
template <typename Scalar>
struct ByValueColMajorLhsMapper {
const Scalar* data;
Index stride; // element(i, k) = data[i + k * stride], contiguous in i
Scalar operator()(Index i, Index k) const { return data[i + k * stride]; }
template <typename Packet>
EIGEN_ALWAYS_INLINE Packet loadPacket(Index i, Index k) const {
return internal::ploadu<Packet>(data + i + k * stride);
}
};
template <typename Scalar>
struct ByValueRowMajorLhsMapper {
const Scalar* data;
Index stride; // element(i, k) = data[i * stride + k], contiguous in k
Scalar operator()(Index i, Index k) const { return data[i * stride + k]; }
template <typename Packet>
EIGEN_ALWAYS_INLINE Packet loadPacket(Index i, Index k) const {
return internal::ploadu<Packet>(data + i * stride + k);
}
};
template <typename Scalar>
struct ByValueColMajorRhsMapper {
struct LinearMapper {
const Scalar* p; // &element(0, col); contiguous in depth
Scalar operator()(Index k) const { return p[k]; }
template <typename Packet>
EIGEN_ALWAYS_INLINE Packet loadPacket(Index k) const {
return internal::ploadu<Packet>(p + k);
}
};
const Scalar* data;
Index stride; // element(k, col) = data[k + col * stride], contiguous in k
Scalar operator()(Index k, Index col) const { return data[k + col * stride]; }
LinearMapper getLinearMapper(Index k, Index col) const { return LinearMapper{data + k + col * stride}; }
};
template <typename Scalar>
struct ByValueRowMajorRhsMapper {
struct LinearMapper {
const Scalar* p; // &element(row, col); packet offsets advance columns
Scalar operator()(Index offset) const { return p[offset]; }
template <typename Packet>
EIGEN_ALWAYS_INLINE Packet loadPacket(Index offset) const {
return internal::ploadu<Packet>(p + offset);
}
};
const Scalar* data;
Index stride; // element(k, col) = data[k * stride + col], contiguous in col
Scalar operator()(Index k, Index col) const { return data[k * stride + col]; }
LinearMapper getLinearMapper(Index k, Index col) const { return LinearMapper{data + k * stride + col}; }
};
// Length of the packed LHS buffer for the given panel-mode layout, matching the
// dst_panel formula in sme_pack_lhs_fallback.
template <bool PanelMode>
static Index packed_len(Index outer, Index depth, Index unit, Index dst_stride, Index dst_offset) {
// `outer` is rows (LHS) or cols (RHS); `unit` is the panel width mr or nr.
if (!PanelMode) return outer * depth;
Index end = 0;
for (Index i = 0; i < outer; i += unit) {
const Index w = numext::mini(unit, outer - i);
end = numext::maxi(end, i * dst_stride + dst_offset * w + depth * w);
}
return end;
}
template <typename Scalar, bool PanelMode, bool Conjugate, typename MatrixType>
static void fill_lhs_ref(SmeVector<Scalar>& ref, const MatrixType& V, Index rows, Index depth, Index dst_stride,
Index dst_offset) {
const Index MR = sme_mr<Scalar>();
ref.setConstant(pack_sentinel<Scalar>());
for (Index i = 0; i < rows; i += MR) {
const Index w = numext::mini(MR, rows - i);
const Index base = PanelMode ? i * dst_stride + dst_offset * w : i * depth;
for (Index k = 0; k < depth; ++k)
for (Index r = 0; r < w; ++r) {
const Scalar v = V(i + r, k);
set_packed(ref.data() + base, w, k, r, Conjugate ? numext::conj(v) : v);
}
}
}
template <typename Scalar, bool PanelMode, bool Conjugate, typename MatrixType>
static void fill_rhs_ref(SmeVector<Scalar>& ref, const MatrixType& V, Index cols, Index depth, Index dst_stride,
Index dst_offset) {
const Index NR = sme_nr<Scalar>();
ref.setConstant(pack_sentinel<Scalar>());
for (Index j = 0; j < cols; j += NR) {
const Index w = numext::mini(NR, cols - j);
const Index base = PanelMode ? j * dst_stride + dst_offset * w : j * depth;
for (Index k = 0; k < depth; ++k)
for (Index c = 0; c < w; ++c) {
const Scalar v = V(k, j + c);
set_packed(ref.data() + base, w, k, c, Conjugate ? numext::conj(v) : v);
}
}
}
// Every packer check below is the same nine lines -- fill a sentinel buffer,
// build the scalar reference, run the packer, compare exactly -- over a
// different mapper. The mapper is what each check is really about, so it is
// built by the caller and the rest lives here once.
template <typename Scalar, int StorageOrder, bool PanelMode, bool Conjugate, typename Mapper, typename MatrixType>
static void check_packed_lhs(const Mapper& mapper, const MatrixType& V, Index rows, Index depth) {
const Index dst_stride = PanelMode ? depth + 5 : 0;
const Index dst_offset = PanelMode ? 3 : 0;
const Index len = packed_len<PanelMode>(rows, depth, sme_mr<Scalar>(), dst_stride, dst_offset);
SmeVector<Scalar> packed = SmeVector<Scalar>::Constant(len, pack_sentinel<Scalar>());
SmeVector<Scalar> ref(len);
fill_lhs_ref<Scalar, PanelMode, Conjugate>(ref, V, rows, depth, dst_stride, dst_offset);
internal::gemm_pack_lhs<Scalar, Index, Mapper, sme_mr<Scalar>(), 1, typename internal::packet_traits<Scalar>::type,
StorageOrder, Conjugate, PanelMode>
pack;
pack(packed.data(), mapper, depth, rows, dst_stride, dst_offset);
VERIFY_IS_EQUAL(packed, ref);
}
template <typename Scalar, int StorageOrder, bool PanelMode, bool Conjugate, typename Mapper, typename MatrixType>
static void check_packed_rhs(const Mapper& mapper, const MatrixType& V, Index depth, Index cols) {
const Index dst_stride = PanelMode ? depth + 5 : 0;
const Index dst_offset = PanelMode ? 3 : 0;
const Index len = packed_len<PanelMode>(cols, depth, sme_nr<Scalar>(), dst_stride, dst_offset);
SmeVector<Scalar> packed = SmeVector<Scalar>::Constant(len, pack_sentinel<Scalar>());
SmeVector<Scalar> ref(len);
fill_rhs_ref<Scalar, PanelMode, Conjugate>(ref, V, cols, depth, dst_stride, dst_offset);
internal::gemm_pack_rhs<Scalar, Index, Mapper, sme_nr<Scalar>(), StorageOrder, Conjugate, PanelMode> pack;
pack(packed.data(), mapper, depth, cols, dst_stride, dst_offset);
VERIFY_IS_EQUAL(packed, ref);
}
// The raw pointer + stride packers, taken whenever the mapper grants direct
// unit-inner-stride access.
template <typename Scalar, int StorageOrder, bool PanelMode, bool Conjugate>
static void verify_pack_lhs_direct(Index rows, Index depth) {
using MatrixType = Matrix<Scalar, Dynamic, Dynamic, StorageOrder>;
using Mapper = internal::const_blas_data_mapper<Scalar, Index, StorageOrder>;
MatrixType V = MatrixType::Random(rows, depth);
check_packed_lhs<Scalar, StorageOrder, PanelMode, Conjugate>(Mapper(V.data(), V.outerStride()), V, rows, depth);
}
template <typename Scalar, int StorageOrder, bool PanelMode, bool Conjugate>
static void verify_pack_rhs_direct(Index depth, Index cols) {
using MatrixType = Matrix<Scalar, Dynamic, Dynamic, StorageOrder>;
using Mapper = internal::const_blas_data_mapper<Scalar, Index, StorageOrder>;
MatrixType V = MatrixType::Random(depth, cols);
check_packed_rhs<Scalar, StorageOrder, PanelMode, Conjugate>(Mapper(V.data(), V.outerStride()), V, depth, cols);
}
// Inner-strided blas mapper LHS: element(i, k) laid out with inner stride
// `incr`. ColMajor takes the vectorised gather path; RowMajor takes the scalar
// path (its packets would run along depth, not rows).
template <typename Scalar, int StorageOrder, bool PanelMode, bool Conjugate>
static void verify_fallback_lhs_strided(Index rows, Index depth, Index incr) {
using Mapper = internal::blas_data_mapper<Scalar, Index, StorageOrder, Unaligned, Dynamic>;
Matrix<Scalar, Dynamic, Dynamic> V = Matrix<Scalar, Dynamic, Dynamic>::Random(rows, depth);
const Index mstride = (StorageOrder == ColMajor ? rows : depth) * incr;
SmeVector<Scalar> buf = SmeVector<Scalar>::Zero((StorageOrder == ColMajor ? depth : rows) * mstride + incr);
for (Index k = 0; k < depth; ++k)
for (Index i = 0; i < rows; ++i)
buf[StorageOrder == ColMajor ? i * incr + k * mstride : k * incr + i * mstride] = V(i, k);
check_packed_lhs<Scalar, StorageOrder, PanelMode, Conjugate>(Mapper(buf.data(), mstride, incr), V, rows, depth);
}
// By-value LHS mappers exercise both packet directions. RowMajor must stay
// scalar because its packets advance depth rather than rows.
template <typename Scalar, int StorageOrder, bool PanelMode, bool Conjugate>
static void verify_fallback_lhs_byvalue(Index rows, Index depth) {
using MatrixType = Matrix<Scalar, Dynamic, Dynamic, StorageOrder>;
using Mapper = typename std::conditional<StorageOrder == ColMajor, ByValueColMajorLhsMapper<Scalar>,
ByValueRowMajorLhsMapper<Scalar>>::type;
MatrixType V = MatrixType::Random(rows, depth);
check_packed_lhs<Scalar, StorageOrder, PanelMode, Conjugate>(Mapper{V.data(), V.outerStride()}, V, rows, depth);
}
// Inner-strided blas mapper RHS: element(k, col) with inner stride `incr`.
// ColMajor takes the vectorised transpose path; RowMajor takes the scalar path.
template <typename Scalar, int StorageOrder, bool PanelMode, bool Conjugate>
static void verify_fallback_rhs_strided(Index depth, Index cols, Index incr) {
using Mapper = internal::blas_data_mapper<Scalar, Index, StorageOrder, Unaligned, Dynamic>;
Matrix<Scalar, Dynamic, Dynamic> V = Matrix<Scalar, Dynamic, Dynamic>::Random(depth, cols);
const Index mstride = (StorageOrder == ColMajor ? depth : cols) * incr;
SmeVector<Scalar> buf = SmeVector<Scalar>::Zero((StorageOrder == ColMajor ? cols : depth) * mstride + incr);
for (Index col = 0; col < cols; ++col)
for (Index k = 0; k < depth; ++k)
buf[StorageOrder == ColMajor ? k * incr + col * mstride : col * incr + k * mstride] = V(k, col);
check_packed_rhs<Scalar, StorageOrder, PanelMode, Conjugate>(Mapper(buf.data(), mstride, incr), V, depth, cols);
}
// By-value RHS mappers likewise cover both packet directions. RowMajor packets
// advance columns, so the depth-oriented transpose fallback must stay scalar.
template <typename Scalar, int StorageOrder, bool PanelMode, bool Conjugate>
static void verify_fallback_rhs_byvalue(Index depth, Index cols) {
using MatrixType = Matrix<Scalar, Dynamic, Dynamic, StorageOrder>;
using Mapper = typename std::conditional<StorageOrder == ColMajor, ByValueColMajorRhsMapper<Scalar>,
ByValueRowMajorRhsMapper<Scalar>>::type;
MatrixType V = MatrixType::Random(depth, cols);
check_packed_rhs<Scalar, StorageOrder, PanelMode, Conjugate>(Mapper{V.data(), V.outerStride()}, V, depth, cols);
}
// Each packer is swept over all four (PanelMode, Conjugate) combinations.
// Conjugate=true is the identity on real scalars; for complex ones it negates
// the packed imaginary half, and reaches the direct packers from the SYMM
// above-diagonal transposed pack and the fallbacks from an inner-strided
// selfadjoint operand.
template <typename Scalar, int SO>
static void sweep_pack_direct(Index n, Index depth) {
verify_pack_lhs_direct<Scalar, SO, false, false>(n, depth);
verify_pack_lhs_direct<Scalar, SO, false, true>(n, depth);
verify_pack_lhs_direct<Scalar, SO, true, false>(n, depth);
verify_pack_lhs_direct<Scalar, SO, true, true>(n, depth);
verify_pack_rhs_direct<Scalar, SO, false, false>(depth, n);
verify_pack_rhs_direct<Scalar, SO, false, true>(depth, n);
verify_pack_rhs_direct<Scalar, SO, true, false>(depth, n);
verify_pack_rhs_direct<Scalar, SO, true, true>(depth, n);
}
template <typename Scalar, int SO>
static void sweep_fallback_strided(Index n, Index depth, Index incr) {
verify_fallback_lhs_strided<Scalar, SO, false, false>(n, depth, incr);
verify_fallback_lhs_strided<Scalar, SO, false, true>(n, depth, incr);
verify_fallback_lhs_strided<Scalar, SO, true, false>(n, depth, incr);
verify_fallback_lhs_strided<Scalar, SO, true, true>(n, depth, incr);
verify_fallback_rhs_strided<Scalar, SO, false, false>(depth, n, incr);
verify_fallback_rhs_strided<Scalar, SO, false, true>(depth, n, incr);
verify_fallback_rhs_strided<Scalar, SO, true, false>(depth, n, incr);
verify_fallback_rhs_strided<Scalar, SO, true, true>(depth, n, incr);
}
template <typename Scalar, int SO>
static void sweep_fallback_byvalue(Index n, Index depth) {
verify_fallback_lhs_byvalue<Scalar, SO, false, false>(n, depth);
verify_fallback_lhs_byvalue<Scalar, SO, false, true>(n, depth);
verify_fallback_lhs_byvalue<Scalar, SO, true, false>(n, depth);
verify_fallback_lhs_byvalue<Scalar, SO, true, true>(n, depth);
verify_fallback_rhs_byvalue<Scalar, SO, false, false>(depth, n);
verify_fallback_rhs_byvalue<Scalar, SO, false, true>(depth, n);
verify_fallback_rhs_byvalue<Scalar, SO, true, false>(depth, n);
verify_fallback_rhs_byvalue<Scalar, SO, true, true>(depth, n);
}
template <typename Scalar>
static void test_pack_direct() {
const int TILE = sme_tile<Scalar>();
const int MR = sme_mr<Scalar>();
const int NR = sme_nr<Scalar>();
// Widths around the tile side and both panel widths, which differ for
// complex scalars.
const int widths[] = {1, TILE - 1, TILE, TILE + 1, MR, MR + 1, NR, NR + 1, 2 * NR + 1};
const int depths[] = {1, 3, 8, 35};
for (int d : depths) {
for (int n : widths) {
sweep_pack_direct<Scalar, ColMajor>(n, d);
sweep_pack_direct<Scalar, RowMajor>(n, d);
}
}
}
template <typename Scalar>
static void test_mapper_fallback() {
const int MR = sme_mr<Scalar>();
const int widths[] = {4, 5, MR, MR + 1, 2 * MR + 1}; // rows/cols around 4 and the panel width
const int depths[] = {1, 3, 8, 35}; // depth remainders 1..3 and larger
for (int n : widths) {
for (int d : depths) {
// RowMajor mappers take the scalar path: their packets advance the index
// the fallback does not transpose.
for (int incr : {2, 3}) {
sweep_fallback_strided<Scalar, ColMajor>(n, d, incr);
sweep_fallback_strided<Scalar, RowMajor>(n, d, incr);
}
sweep_fallback_byvalue<Scalar, ColMajor>(n, d);
sweep_fallback_byvalue<Scalar, RowMajor>(n, d);
}
}
}
// ---------------------------------------------------------------------------
// Product-level coverage, swept relative to the scalar's own block width.
// ---------------------------------------------------------------------------
// Sizes that land just on and off the block tails and the intra-block ZA-tile
// splits, for both panel widths (they differ for complex scalars) and the tile
// side at the SVL=512 design point.
template <typename Scalar>
static std::vector<int> sme_edge_sizes() {
const int T = sme_tile<Scalar>();
const int MR = sme_mr<Scalar>();
const int NR = sme_nr<Scalar>();
std::vector<int> sizes = {1, T - 1, T, T + 1, MR - 1, MR, MR + 1, NR - 1,
NR, NR + 1, 2 * MR - 1, 2 * MR, 2 * MR + 1, 2 * NR - 1, 2 * NR, 2 * NR + 1};
std::sort(sizes.begin(), sizes.end());
sizes.erase(std::unique(sizes.begin(), sizes.end()), sizes.end());
return sizes;
}
// The dimension-sum and output-area bounds in GeneralProduct.h route a small
// square product to the coeff-based evaluator, so the n x n sweeps above stop
// reaching the SME kernel once n gets small -- which is exactly where the
// intra-block ZA tile-split edges (T-1, T, T+1) live for double and
// complex<double>. Repeat the same widths over a depth that clears both
// bounds, so every m and n edge still reaches the kernel, and check the
// conjugated forms too: at those sizes test_conjugated_products compares a
// lazyProduct against a lazyProduct and cannot fail.
template <typename Scalar>
static void test_edge_sizes_deep_k() {
const int depth = 4 * sme_mr<Scalar>();
for (int n : sme_edge_sizes<Scalar>()) {
if (n < 2) continue;
const SmeColMajorMat<Scalar> A = SmeColMajorMat<Scalar>::Random(n, depth);
const SmeColMajorMat<Scalar> B = SmeColMajorMat<Scalar>::Random(depth, n);
const SmeColMajorMat<Scalar> Aa = A.adjoint().eval();
const SmeColMajorMat<Scalar> Ba = B.adjoint().eval();
const SmeColMajorMat<Scalar> c_before = SmeColMajorMat<Scalar>::Random(n, n);
const SmeColMajorMat<Scalar> expected = c_before + A.lazyProduct(B);
SmeColMajorMat<Scalar> C = c_before;
C.noalias() += A * B;
VERIFY_IS_APPROX(C, expected);
C = c_before;
C.noalias() += Aa.adjoint() * B;
VERIFY_IS_APPROX(C, expected);
C = c_before;
C.noalias() += A * Ba.adjoint();
VERIFY_IS_APPROX(C, expected);
}
}
template <typename Scalar>
static void test_conjugated_products() {
for (int n : sme_edge_sizes<Scalar>()) {
verify_conjugated_products<Scalar, SmeColMajorMat<Scalar>, SmeColMajorMat<Scalar>, SmeColMajorMat<Scalar>>(n);
verify_conjugated_products<Scalar, SmeRowMajorMat<Scalar>, SmeColMajorMat<Scalar>, SmeColMajorMat<Scalar>>(n);
verify_conjugated_products<Scalar, SmeColMajorMat<Scalar>, SmeRowMajorMat<Scalar>, SmeRowMajorMat<Scalar>>(n);
verify_conjugated_products<Scalar, SmeRowMajorMat<Scalar>, SmeRowMajorMat<Scalar>, SmeColMajorMat<Scalar>>(n);
}
}
template <typename Scalar>
static void test_products() {
const int MR = sme_mr<Scalar>();
// Square edge cases around the block and tile boundaries.
for (int n : sme_edge_sizes<Scalar>()) product(SmeColMajorMat<Scalar>(n, n));
// Thin / wide rectangular cases (M x 1, 1 x N) and non-square cases that
// exercise tail paths for both M and N.
product(SmeColMajorMat<Scalar>(MR, 1));
product(SmeColMajorMat<Scalar>(1, MR));
product(SmeColMajorMat<Scalar>(1, 2 * MR));
product(SmeColMajorMat<Scalar>(2 * MR, 1));
product(SmeColMajorMat<Scalar>(MR + 1, 2 * MR + 1));
product(SmeColMajorMat<Scalar>(2 * MR + 1, MR + 1));
product(SmeColMajorMat<Scalar>(MR - 1, 2 * MR - 1));
product(SmeColMajorMat<Scalar>(MR + 1, 7));
product(SmeColMajorMat<Scalar>(7, MR + 1));
product(SmeColMajorMat<Scalar>(4 * MR, 3));
product(SmeColMajorMat<Scalar>(3, 4 * MR));
test_deep_k_split<Scalar>();
test_edge_sizes_deep_k<Scalar>();
// Random sizes
for (int i = 0; i < g_repeat; i++) {
product(SmeColMajorMat<Scalar>(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).
for (int n : sme_edge_sizes<Scalar>()) {
if (n > 1) product(SmeRowMajorMat<Scalar>(n, n));
}
product(SmeRowMajorMat<Scalar>(MR + 1, 2 * MR + 1));
product(SmeRowMajorMat<Scalar>(2 * MR + 1, MR + 1));
for (int i = 0; i < g_repeat; i++) {
product(SmeRowMajorMat<Scalar>(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 : sme_edge_sizes<Scalar>()) {
if (n < 2) continue;
test_general_strided_result<Scalar, SmeColMajorMat<Scalar>>(n);
test_general_strided_result<Scalar, SmeRowMajorMat<Scalar>>(n);
test_rowmajor_strided_result<Scalar>(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 : sme_edge_sizes<Scalar>()) {
if (n < 2) continue;
SmeRowMajorMat<Scalar> A = SmeRowMajorMat<Scalar>::Random(n, n);
SmeRowMajorMat<Scalar> B = SmeRowMajorMat<Scalar>::Random(n, n);
SmeColMajorMat<Scalar> C = SmeColMajorMat<Scalar>::Zero(n, n);
C.noalias() += A * B;
VERIFY_IS_APPROX(C, (A.lazyProduct(B)).eval());
}
}
EIGEN_DECLARE_TEST(product_sme) {
CALL_SUBTEST_1(test_whilelt_operand_types<float>());
CALL_SUBTEST_1(test_products<float>());
CALL_SUBTEST_1(test_conjugated_products<float>());
CALL_SUBTEST_1(test_symm_pack<float>());
CALL_SUBTEST_1(test_pack_direct<float>());
CALL_SUBTEST_1(test_mapper_fallback<float>());
// double reaches the SME kernel and packers only with FEAT_SME_F64F64; the
// product sweep is meaningful either way, but the packed-layout tests name
// specializations that only exist when it is available.
CALL_SUBTEST_2(test_products<double>());
CALL_SUBTEST_2(test_conjugated_products<double>());
#ifdef EIGEN_VECTORIZE_SME_F64F64
CALL_SUBTEST_2(test_whilelt_operand_types<double>());
CALL_SUBTEST_2(test_symm_pack<double>());
CALL_SUBTEST_2(test_pack_direct<double>());
CALL_SUBTEST_2(test_mapper_fallback<double>());
#endif
CALL_SUBTEST_3(test_products<std::complex<float>>());
CALL_SUBTEST_3(test_conjugated_products<std::complex<float>>());
CALL_SUBTEST_3(test_symm_pack<std::complex<float>>());
CALL_SUBTEST_3(test_pack_direct<std::complex<float>>());
CALL_SUBTEST_3(test_mapper_fallback<std::complex<float>>());
// complex<double> accumulates into ZA.D tiles, so it needs FEAT_SME_F64F64
// exactly as double does.
CALL_SUBTEST_4(test_products<std::complex<double>>());
CALL_SUBTEST_4(test_conjugated_products<std::complex<double>>());
#ifdef EIGEN_VECTORIZE_SME_F64F64
CALL_SUBTEST_4(test_symm_pack<std::complex<double>>());
CALL_SUBTEST_4(test_pack_direct<std::complex<double>>());
CALL_SUBTEST_4(test_mapper_fallback<std::complex<double>>());
#endif
// A scalar type SME does not specialize, proving it still routes through the
// generic product path inside an SME build -- where packet traits, alignment
// and cache blocking all differ from a plain NEON build.
CALL_SUBTEST_5(product(Matrix<long double, Dynamic, Dynamic>(33, 17)));
}