SME: Run small and thin products on the NEON kernel from the GEMM driver libeigen/eigen!3370 Closes #3194 Co-authored-by: Rasmus Munk Larsen <rmlarsen@gmail.com> Co-authored-by: michael tesch <tesch1@gmail.com>
diff --git a/Eigen/src/Core/arch/SME/GeneralBlockPanelKernel.h b/Eigen/src/Core/arch/SME/GeneralBlockPanelKernel.h index 32fdd21..d56c6f8 100644 --- a/Eigen/src/Core/arch/SME/GeneralBlockPanelKernel.h +++ b/Eigen/src/Core/arch/SME/GeneralBlockPanelKernel.h
@@ -2520,30 +2520,31 @@ }; // One (NPack * PacketSize) x NCol micro-tile of a real block; NPack == 0 is a -// single scalar row. Pointers arrive offset to the tile's first row and column. +// single scalar row. Pointers arrive offset to the tile's first row and column; +// a depth step of A is `as` scalars apart, of B `bs`. template <typename Scalar, typename Index, int NPack, int NCol> struct sme_neon_tile { using Packet = typename packet_traits<Scalar>::type; using Cols = sme_neon_cols<Scalar, NCol>; static constexpr int PS = packet_traits<Scalar>::size; - // Narrow tiles split the depth over KU accumulator chains to hide FMA latency. - static constexpr int KU = (NPack * NCol >= 4) ? 1 : 4 / (NPack * NCol); + // Narrow tiles split the depth over KU accumulator sets, for at least 8 FMA chains. + static constexpr int KU = (NPack * NCol >= 8) ? 1 : 8 / (NPack * NCol); static EIGEN_ALWAYS_INLINE void step(Packet (&acc)[NPack][NCol], const Scalar* a, const Scalar* b) { Packet av[NPack]; for (int p = 0; p < NPack; ++p) av[p] = ploadu<Packet>(a + p * PS); sme_neon_col_loop<0, NCol>::template real<Cols, Packet, NPack>(acc, av, Cols::load(b)); } static EIGEN_ALWAYS_INLINE void run(Scalar* C, Index rs, Index cs, const Scalar* blA, const Scalar* blB, Index depth, - Scalar alpha, Index pw, Index cw) { + Scalar alpha, Index as, Index bs) { Packet acc[KU][NPack][NCol]; for (int u = 0; u < KU; ++u) for (int p = 0; p < NPack; ++p) for (int c = 0; c < NCol; ++c) acc[u][p][c] = pset1<Packet>(Scalar(0)); Index k = 0; for (; k + KU <= depth; k += KU) { - for (int u = 0; u < KU; ++u) step(acc[u], blA + (k + u) * pw, blB + (k + u) * cw); + for (int u = 0; u < KU; ++u) step(acc[u], blA + (k + u) * as, blB + (k + u) * bs); } - for (; k < depth; ++k) step(acc[0], blA + k * pw, blB + k * cw); + for (; k < depth; ++k) step(acc[0], blA + k * as, blB + k * bs); for (int u = 1; u < KU; ++u) for (int p = 0; p < NPack; ++p) for (int c = 0; c < NCol; ++c) acc[0][p][c] = padd(acc[0][p][c], acc[u][p][c]); @@ -2564,7 +2565,7 @@ template <typename Scalar, typename Index, int NCol> struct sme_neon_tile<Scalar, Index, 0, NCol> { static EIGEN_ALWAYS_INLINE void run(Scalar* C, Index rs, Index cs, const Scalar* blA, const Scalar* blB, Index depth, - Scalar alpha, Index pw, Index cw) { + Scalar alpha, Index as, Index bs) { EIGEN_UNUSED_VARIABLE(rs); constexpr int KU = (NCol >= 4) ? 1 : 4 / NCol; Scalar acc[KU][NCol]; @@ -2573,14 +2574,14 @@ Index k = 0; for (; k + KU <= depth; k += KU) { for (int u = 0; u < KU; ++u) { - const Scalar a = blA[(k + u) * pw]; - const Scalar* b = blB + (k + u) * cw; + const Scalar a = blA[(k + u) * as]; + const Scalar* b = blB + (k + u) * bs; for (int c = 0; c < NCol; ++c) acc[u][c] += a * b[c]; } } for (; k < depth; ++k) { - const Scalar a = blA[k * pw]; - const Scalar* b = blB + k * cw; + const Scalar a = blA[k * as]; + const Scalar* b = blB + k * bs; for (int c = 0; c < NCol; ++c) acc[0][c] += a * b[c]; } for (int u = 1; u < KU; ++u) @@ -2636,32 +2637,48 @@ } } } + // Narrow tiles split the depth over KU accumulator sets, for at least 16 FMA chains (two per cell). + static constexpr int KU = (NPack * NCol >= 8) ? 1 : 8 / (NPack * NCol); + static EIGEN_ALWAYS_INLINE void step(Packet (&acc_re)[NPack][NCol], Packet (&acc_im)[NPack][NCol], + const RealScalar* a, const RealScalar* b, Index pw, Index cw) { + Packet are[NPack], aim[NPack]; + for (int p = 0; p < NPack; ++p) { + are[p] = ploadu<Packet>(a + p * PS); + aim[p] = ploadu<Packet>(a + pw + p * PS); + } + sme_neon_col_loop<0, NCol>::template cplx<ConjLhs, ConjRhs, Cols, Packet, NPack>(acc_re, acc_im, are, aim, + Cols::load(b), Cols::load(b + cw)); + } static EIGEN_ALWAYS_INLINE void run(Scalar* C, Index rs, Index cs, const RealScalar* rA, const RealScalar* rB, Index depth, Scalar alpha, Index pw, Index cw) { - Packet acc_re[NPack][NCol], acc_im[NPack][NCol]; - for (int p = 0; p < NPack; ++p) { - for (int c = 0; c < NCol; ++c) { - acc_re[p][c] = pset1<Packet>(RealScalar(0)); - acc_im[p][c] = pset1<Packet>(RealScalar(0)); + Packet acc_re[KU][NPack][NCol], acc_im[KU][NPack][NCol]; + for (int u = 0; u < KU; ++u) { + for (int p = 0; p < NPack; ++p) { + for (int c = 0; c < NCol; ++c) { + acc_re[u][p][c] = pset1<Packet>(RealScalar(0)); + acc_im[u][p][c] = pset1<Packet>(RealScalar(0)); + } } } - for (Index k = 0; k < depth; ++k) { - const RealScalar* a = rA + k * 2 * pw; - const RealScalar* b = rB + k * 2 * cw; - Packet are[NPack], aim[NPack]; + Index k = 0; + for (; k + KU <= depth; k += KU) { + for (int u = 0; u < KU; ++u) step(acc_re[u], acc_im[u], rA + (k + u) * 2 * pw, rB + (k + u) * 2 * cw, pw, cw); + } + for (; k < depth; ++k) step(acc_re[0], acc_im[0], rA + k * 2 * pw, rB + k * 2 * cw, pw, cw); + for (int u = 1; u < KU; ++u) { for (int p = 0; p < NPack; ++p) { - are[p] = ploadu<Packet>(a + p * PS); - aim[p] = ploadu<Packet>(a + pw + p * PS); + for (int c = 0; c < NCol; ++c) { + acc_re[0][p][c] = padd(acc_re[0][p][c], acc_re[u][p][c]); + acc_im[0][p][c] = padd(acc_im[0][p][c], acc_im[u][p][c]); + } } - sme_neon_col_loop<0, NCol>::template cplx<ConjLhs, ConjRhs, Cols, Packet, NPack>( - acc_re, acc_im, are, aim, Cols::load(b), Cols::load(b + cw)); } const RealScalar ar = numext::real(alpha), ai = numext::imag(alpha); if (rs == 1) { if (ai == RealScalar(0)) { - store_unit_stride<true>(C, cs, acc_re, acc_im, ar, ai); + store_unit_stride<true>(C, cs, acc_re[0], acc_im[0], ar, ai); } else { - store_unit_stride<false>(C, cs, acc_re, acc_im, ar, ai); + store_unit_stride<false>(C, cs, acc_re[0], acc_im[0], ar, ai); } return; } @@ -2669,8 +2686,8 @@ for (int p = 0; p < NPack; ++p) { RealScalar* pc = reinterpret_cast<RealScalar*>(C + Index(p * PS) * rs + Index(c) * cs); RealScalar re[PS], im[PS]; - pstoreu(re, acc_re[p][c]); - pstoreu(im, acc_im[p][c]); + pstoreu(re, acc_re[0][p][c]); + pstoreu(im, acc_im[0][p][c]); for (int l = 0; l < PS; ++l) sme_neon_cstore_scalar(pc + Index(l) * 2 * rs, re[l], im[l], ar, ai); } } @@ -2709,20 +2726,58 @@ } }; -// Row loop of one column group: 3, 2 and 1 packets of rows, then scalar rows. -template <bool ConjLhs, bool ConjRhs, int NCol, typename Scalar, typename Index> -EIGEN_ALWAYS_INLINE void sme_neon_column_group(Scalar* C, Index rs, Index cs, const Scalar* blA, const Scalar* blB, - Index depth, Scalar alpha, Index pw, Index cw) { +// Rows r .. rows - 1, fewer than a packet. When the group's columns fill whole packets, C^T = B^T A^T runs on the +// tile with the operands and C's strides swapped: vectorized over the columns, two rows (Cols lanes) at a time. +template <int NCol, typename Scalar, typename Index> +EIGEN_ALWAYS_INLINE void sme_neon_row_tail(std::true_type, Scalar* C, Index rs, Index cs, const Scalar* A, Index as, + const Scalar* B, Index bs, Index depth, Scalar alpha, Index r, Index rows) { + constexpr int NV = NCol / packet_traits<Scalar>::size; + for (; r + 2 <= rows; r += 2) + sme_neon_tile<Scalar, Index, NV, 2>::run(C + r * rs, cs, rs, B, A + r, depth, alpha, bs, as); + if (r < rows) sme_neon_tile<Scalar, Index, NV, 1>::run(C + r * rs, cs, rs, B, A + r, depth, alpha, bs, as); +} +template <int NCol, typename Scalar, typename Index> +EIGEN_ALWAYS_INLINE void sme_neon_row_tail(std::false_type, Scalar* C, Index rs, Index cs, const Scalar* A, Index as, + const Scalar* B, Index bs, Index depth, Scalar alpha, Index r, Index rows) { + for (; r < rows; ++r) sme_neon_tile<Scalar, Index, 0, NCol>::run(C + r * rs, rs, cs, A + r, B, depth, alpha, as, bs); +} + +// Row loop of one column group: 3, 2 and 1 packets of rows, then the row tail. A depth step of A is `as` scalars +// apart, of B `bs`: the panel widths for packed panels, or the leading dimensions for operands read in place. +template <int NCol, typename Scalar, typename Index> +EIGEN_ALWAYS_INLINE void sme_neon_column_group(Scalar* C, Index rs, Index cs, const Scalar* A, Index as, + const Scalar* B, Index bs, Index depth, Scalar alpha, Index rows) { constexpr Index PS = Index(packet_traits<Scalar>::size); Index r = 0; - for (; r + 3 * PS <= pw; r += 3 * PS) - sme_neon_tile<Scalar, Index, 3, NCol>::run(C + r * rs, rs, cs, blA + r, blB, depth, alpha, pw, cw); - for (; r + 2 * PS <= pw; r += 2 * PS) - sme_neon_tile<Scalar, Index, 2, NCol>::run(C + r * rs, rs, cs, blA + r, blB, depth, alpha, pw, cw); - for (; r + PS <= pw; r += PS) - sme_neon_tile<Scalar, Index, 1, NCol>::run(C + r * rs, rs, cs, blA + r, blB, depth, alpha, pw, cw); + for (; r + 3 * PS <= rows; r += 3 * PS) + sme_neon_tile<Scalar, Index, 3, NCol>::run(C + r * rs, rs, cs, A + r, B, depth, alpha, as, bs); + for (; r + 2 * PS <= rows; r += 2 * PS) + sme_neon_tile<Scalar, Index, 2, NCol>::run(C + r * rs, rs, cs, A + r, B, depth, alpha, as, bs); + for (; r + PS <= rows; r += PS) + sme_neon_tile<Scalar, Index, 1, NCol>::run(C + r * rs, rs, cs, A + r, B, depth, alpha, as, bs); + sme_neon_row_tail<NCol>(bool_constant<NCol % PS == 0>(), C, rs, cs, A, as, B, bs, depth, alpha, r, rows); +} + +template <bool ConjLhs, bool ConjRhs, int NCol, typename RealScalar, typename Index> +EIGEN_ALWAYS_INLINE void sme_neon_crow_tail(std::true_type, std::complex<RealScalar>* C, Index rs, Index cs, + const RealScalar* rA, const RealScalar* rB, Index depth, + std::complex<RealScalar> alpha, Index r, Index pw, Index cw) { + constexpr int NV = NCol / packet_traits<RealScalar>::size; + // conj(a) b = b conj(a): the swapped tile swaps the conjugations too. + for (; r + 2 <= pw; r += 2) + sme_neon_ctile<RealScalar, Index, ConjRhs, ConjLhs, NV, 2>::run(C + r * rs, cs, rs, rB, rA + r, depth, alpha, cw, + pw); + if (r < pw) + sme_neon_ctile<RealScalar, Index, ConjRhs, ConjLhs, NV, 1>::run(C + r * rs, cs, rs, rB, rA + r, depth, alpha, cw, + pw); +} +template <bool ConjLhs, bool ConjRhs, int NCol, typename RealScalar, typename Index> +EIGEN_ALWAYS_INLINE void sme_neon_crow_tail(std::false_type, std::complex<RealScalar>* C, Index rs, Index cs, + const RealScalar* rA, const RealScalar* rB, Index depth, + std::complex<RealScalar> alpha, Index r, Index pw, Index cw) { for (; r < pw; ++r) - sme_neon_tile<Scalar, Index, 0, NCol>::run(C + r * rs, rs, cs, blA + r, blB, depth, alpha, pw, cw); + sme_neon_ctile<RealScalar, Index, ConjLhs, ConjRhs, 0, NCol>::run(C + r * rs, rs, cs, rA + r, rB, depth, alpha, pw, + cw); } template <bool ConjLhs, bool ConjRhs, int NCol, typename RealScalar, typename Index> @@ -2738,24 +2793,27 @@ for (; r + PS <= pw; r += PS) sme_neon_ctile<RealScalar, Index, ConjLhs, ConjRhs, 1, NCol>::run(C + r * rs, rs, cs, rA + r, rB, depth, alpha, pw, cw); - for (; r < pw; ++r) - sme_neon_ctile<RealScalar, Index, ConjLhs, ConjRhs, 0, NCol>::run(C + r * rs, rs, cs, rA + r, rB, depth, alpha, pw, - cw); + sme_neon_crow_tail<ConjLhs, ConjRhs, NCol>(bool_constant<NCol % PS == 0>(), C, rs, cs, rA, rB, depth, alpha, r, pw, + cw); } -// One pw x cw block: columns in groups of eight, four, two, then single columns. +// One rows x cols real block: columns in groups of eight, four, two, then single columns. +template <typename Scalar, typename Index> +EIGEN_ALWAYS_INLINE void sme_neon_block_strided(Scalar* C, Index rs, Index cs, const Scalar* A, Index as, + const Scalar* B, Index bs, Index depth, Scalar alpha, Index rows, + Index cols) { + Index c = 0; + for (; c + 8 <= cols; c += 8) sme_neon_column_group<8>(C + c * cs, rs, cs, A, as, B + c, bs, depth, alpha, rows); + for (; c + 4 <= cols; c += 4) sme_neon_column_group<4>(C + c * cs, rs, cs, A, as, B + c, bs, depth, alpha, rows); + for (; c + 2 <= cols; c += 2) sme_neon_column_group<2>(C + c * cs, rs, cs, A, as, B + c, bs, depth, alpha, rows); + for (; c < cols; ++c) sme_neon_column_group<1>(C + c * cs, rs, cs, A, as, B + c, bs, depth, alpha, rows); +} + +// One pw x cw block of packed panels. template <bool ConjLhs, bool ConjRhs, typename Scalar, typename Index> EIGEN_ALWAYS_INLINE void sme_neon_block(Scalar* C, Index rs, Index cs, const Scalar* blA, const Scalar* blB, Index depth, Scalar alpha, Index pw, Index cw) { - Index c = 0; - for (; c + 8 <= cw; c += 8) - sme_neon_column_group<ConjLhs, ConjRhs, 8>(C + c * cs, rs, cs, blA, blB + c, depth, alpha, pw, cw); - for (; c + 4 <= cw; c += 4) - sme_neon_column_group<ConjLhs, ConjRhs, 4>(C + c * cs, rs, cs, blA, blB + c, depth, alpha, pw, cw); - for (; c + 2 <= cw; c += 2) - sme_neon_column_group<ConjLhs, ConjRhs, 2>(C + c * cs, rs, cs, blA, blB + c, depth, alpha, pw, cw); - for (; c < cw; ++c) - sme_neon_column_group<ConjLhs, ConjRhs, 1>(C + c * cs, rs, cs, blA, blB + c, depth, alpha, pw, cw); + sme_neon_block_strided(C, rs, cs, blA, pw, blB, cw, depth, alpha, pw, cw); } template <bool ConjLhs, bool ConjRhs, typename RealScalar, typename Index> @@ -3297,6 +3355,103 @@ return true; } +// Largest rows * cols * depth of a whole product the NEON kernel beats ZA on (fitted on Apple M4 Pro). +#ifndef EIGEN_SME_NEON_SMALL_PRODUCT +template <typename Scalar> +struct sme_neon_small_product : std::integral_constant<int, 10240> {}; +template <> +struct sme_neon_small_product<float> : std::integral_constant<int, 28672> {}; +template <> +struct sme_neon_small_product<std::complex<double>> : std::integral_constant<int, 5120> {}; +#else +template <typename Scalar> +struct sme_neon_small_product : std::integral_constant<int, EIGEN_SME_NEON_SMALL_PRODUCT> {}; +#endif + +template <typename Scalar, typename Index> +EIGEN_ALWAYS_INLINE bool sme_small_gemm_wins(Index rows, Index cols, Index depth) { +#if defined(EIGEN_SME_NO_NEON_SMALL_BLOCKS) || defined(EIGEN_SME_FORCE_NEON_SMALL_BLOCKS) || defined(EIGEN_USE_BLAS) + EIGEN_UNUSED_VARIABLE(rows); + EIGEN_UNUSED_VARIABLE(cols); + EIGEN_UNUSED_VARIABLE(depth); + return false; +#else + constexpr bool kReal = !NumTraits<Scalar>::IsComplex; + constexpr Index kRhs = Index(65536 / sizeof(Scalar)); + constexpr int kSize = sme_neon_small_product<Scalar>::value; + if (rows <= 0 || cols <= 0 || depth <= 0) return false; + // The narrow ZA kernel reads a tall real LHS in place. + if (kReal && cols <= 3 && rows >= 3 * Index(sme_block<Scalar>::mr)) return false; + // Rows of at most two NEON vectors leave most of a ZA tile idle: NEON up to a 64 KB RHS, real ones 128 deep. The + // factors are bounded before each product, which then fits a 32-bit Index. + if (rows <= Index(32 / sizeof(Scalar)) && cols <= kRhs && depth <= (kReal ? Index(128) : kRhs) && + cols * depth <= kRhs) + return true; + // Each product is bounded by kSize^2. Widen before multiplying: kSize is an int, including overrides. + return rows <= kSize && cols <= kSize && int64_t(rows) * int64_t(cols) <= kSize && depth <= kSize && + int64_t(rows) * int64_t(cols) * int64_t(depth) <= kSize; +#endif +} + +// Complex operands: both sides packed, as the kernel reads real and imaginary parts from separate halves. +template <bool ConjLhs, bool ConjRhs, bool LhsInPlace, bool RhsInPlace, typename Scalar, typename Index> +EIGEN_ALWAYS_INLINE void sme_small_gemm_kernel(std::false_type, Scalar* C, Index incr, Index ldc, const Scalar* A, + Index, const Scalar* B, Index, Index rows, Index depth, Index cols, + Scalar alpha) { + sme_gebp_neon<Scalar, ConjLhs, ConjRhs>(C, incr, ldc, A, B, rows, depth, cols, alpha, depth, depth, Index(0), + Index(0)); +} +// Real operands: a side read in place is one panel with its leading dimension as the depth step. +template <bool ConjLhs, bool ConjRhs, bool LhsInPlace, bool RhsInPlace, typename Scalar, typename Index> +EIGEN_ALWAYS_INLINE void sme_small_gemm_kernel(std::true_type, Scalar* C, Index incr, Index ldc, const Scalar* A, + Index lda, const Scalar* B, Index ldb, Index rows, Index depth, + Index cols, Scalar alpha) { + static_assert(!ConjLhs && !ConjRhs, "the SME kernel does not support conjugation of real scalars"); + const Index mr = LhsInPlace ? rows : Index(sme_block<Scalar>::mr), + nr = RhsInPlace ? cols : Index(sme_block<Scalar>::nr); + for (Index j = 0; j < cols; j += nr) { + const Index cw = numext::mini(nr, cols - j); + for (Index i = 0; i < rows; i += mr) { + const Index pw = numext::mini(mr, rows - i); + sme_neon_block_strided(C + i * incr + j * ldc, incr, ldc, LhsInPlace ? A + i : A + i * depth, + LhsInPlace ? lda : pw, RhsInPlace ? B + j : B + j * depth, RhsInPlace ? ldb : cw, depth, + alpha, pw, cw); + } + } +} + +// The whole product on the NEON kernel. A real side already in the packed layout, a ColMajor LHS (a depth step's rows +// contiguous) or a RowMajor RHS (its columns), is read in place; the others are packed with NEON, so the kernel never +// reads a streaming-packed panel. +template <typename Scalar, int LhsOrder, int RhsOrder, bool ConjLhs, bool ConjRhs, int ResInnerStride, typename Index> +EIGEN_DONT_INLINE void sme_small_gemm(Index rows, Index cols, Index depth, const Scalar* lhs, Index lhsStride, + const Scalar* rhs, Index rhsStride, Scalar* res, Index resIncr, Index resStride, + Scalar alpha, Scalar* blockA_, Scalar* blockB_) { + constexpr int MR = sme_block<Scalar>::mr; + constexpr int NR = sme_block<Scalar>::nr; + constexpr bool kReal = !NumTraits<Scalar>::IsComplex; + constexpr bool kLhsInPlace = kReal && LhsOrder == ColMajor, kRhsInPlace = kReal && RhsOrder == RowMajor; + using LhsMapper = const_blas_data_mapper<Scalar, Index, LhsOrder>; + using RhsMapper = const_blas_data_mapper<Scalar, Index, RhsOrder>; + using PackLhs = + std::conditional_t<LhsOrder == ColMajor, sme_pack_lhs_colmajor<Scalar, MR, Index, LhsMapper, false, false>, + sme_pack_lhs_rowmajor<Scalar, MR, Index, LhsMapper, false, false>>; + using PackRhs = + std::conditional_t<RhsOrder == ColMajor, sme_pack_rhs_colmajor<Scalar, NR, Index, RhsMapper, false, false>, + sme_pack_rhs_rowmajor<Scalar, NR, Index, RhsMapper, false, false>>; + ei_declare_aligned_stack_constructed_variable(Scalar, blockA, + kLhsInPlace ? 0 : std::size_t(rows) * std::size_t(depth), blockA_); + ei_declare_aligned_stack_constructed_variable(Scalar, blockB, + kRhsInPlace ? 0 : std::size_t(depth) * std::size_t(cols), blockB_); + if (!kLhsInPlace) PackLhs::pack_neon(blockA, lhs, lhsStride, depth, rows, 0, 0); + if (!kRhsInPlace) PackRhs::pack_neon(blockB, rhs, rhsStride, depth, cols, 0, 0); + // A compile-time unit row stride, as the result mapper gives the kernel, keeps its contiguous column stores. + const Index incr = ResInnerStride == Dynamic ? resIncr : Index(ResInnerStride); + sme_small_gemm_kernel<ConjLhs, ConjRhs, kLhsInPlace, kRhsInPlace>( + bool_constant<kReal>(), res, incr, resStride, kLhsInPlace ? lhs : blockA, lhsStride, kRhsInPlace ? rhs : blockB, + rhsStride, rows, depth, cols, alpha); +} + } // namespace internal } // namespace Eigen
diff --git a/Eigen/src/Core/products/GeneralMatrixMatrix.h b/Eigen/src/Core/products/GeneralMatrixMatrix.h index 8b93e6b..62c1e5f 100644 --- a/Eigen/src/Core/products/GeneralMatrixMatrix.h +++ b/Eigen/src/Core/products/GeneralMatrixMatrix.h
@@ -128,6 +128,31 @@ const Scalar*, Index, Index, Index, ResScalar) { return false; } +template <typename Scalar, typename Index> +EIGEN_ALWAYS_INLINE bool sme_small_gemm_wins(Index rows, Index cols, Index depth); +template <typename Scalar, int LhsOrder, int RhsOrder, bool ConjLhs, bool ConjRhs, int ResInnerStride, typename Index> +EIGEN_DONT_INLINE void sme_small_gemm(Index rows, Index cols, Index depth, const Scalar* lhs, Index lhsStride, + const Scalar* rhs, Index rhsStride, Scalar* res, Index resIncr, Index resStride, + Scalar alpha, Scalar* blockA, Scalar* blockB); +// The test stays inline: a call into the NEON path's frame dirties stack lines the ZA path then packs into. +template <int LhsOrder, int RhsOrder, bool ConjLhs, bool ConjRhs, int ResInnerStride, typename Scalar, typename Index, + typename Blocking> +EIGEN_ALWAYS_INLINE bool sme_run_small_gemm(std::true_type, Index rows, Index cols, Index depth, const Scalar* lhs, + Index lhsStride, const Scalar* rhs, Index rhsStride, Scalar* res, + Index resIncr, Index resStride, Scalar alpha, Blocking& blocking) { + if (!sme_small_gemm_wins<Scalar>(rows, cols, depth)) return false; + const bool fits = depth <= blocking.kc() && rows <= blocking.mc() && cols <= blocking.nc(); + sme_small_gemm<Scalar, LhsOrder, RhsOrder, ConjLhs, ConjRhs, ResInnerStride>( + rows, cols, depth, lhs, lhsStride, rhs, rhsStride, res, resIncr, resStride, alpha, + fits ? blocking.blockA() : nullptr, fits ? blocking.blockB() : nullptr); + return true; +} +template <int LhsOrder, int RhsOrder, bool ConjLhs, bool ConjRhs, int ResInnerStride, typename LhsScalar, + typename RhsScalar, typename ResScalar, typename Index, typename Blocking> +EIGEN_ALWAYS_INLINE bool sme_run_small_gemm(std::false_type, Index, Index, Index, const LhsScalar*, Index, + const RhsScalar*, Index, ResScalar*, Index, Index, ResScalar, Blocking&) { + return false; +} #endif // RHS-first loop order: nc -> kc -> mc. Used by SME to stream ColMajor result @@ -208,9 +233,13 @@ if (numext::is_exactly_zero(alpha)) return; #ifdef EIGEN_VECTORIZE_SME // A parallel session needs every thread to take the packed path; scaleAndAddTo runs tiny results on one thread. - if (info == nullptr && sme_run_tiny_gemm<LhsStorageOrder, RhsStorageOrder>( - bool_constant<sme_tiny_gemm_pair<LhsScalar, RhsScalar>::value>(), rows, cols, depth, - lhs_, lhsStride, rhs_, rhsStride, res_, resIncr, resStride, alpha)) + if (info == nullptr && + (sme_run_tiny_gemm<LhsStorageOrder, RhsStorageOrder>( + bool_constant<sme_tiny_gemm_pair<LhsScalar, RhsScalar>::value>(), rows, cols, depth, lhs_, lhsStride, rhs_, + rhsStride, res_, resIncr, resStride, alpha) || + sme_run_small_gemm<LhsStorageOrder, RhsStorageOrder, ConjugateLhs, ConjugateRhs, ResInnerStride>( + bool_constant<sme_has_gebp_kernel<LhsScalar, RhsScalar>::value>(), rows, cols, depth, lhs_, lhsStride, + rhs_, rhsStride, res_, resIncr, resStride, alpha, blocking))) return; #endif
diff --git a/benchmarks/Core/bench_gemm.cpp b/benchmarks/Core/bench_gemm.cpp index 50c479a..35ec76c 100644 --- a/benchmarks/Core/bench_gemm.cpp +++ b/benchmarks/Core/bench_gemm.cpp
@@ -89,7 +89,11 @@ /* threshold (k*n*sizeof(Lhs) near L1). mc blocking only takes effect when */ \ /* m is large; hitting the threshold requires sweeping small n=k with large m. */ \ ->Args({4096, 96, 96})->Args({4096, 128, 128})->Args({4096, 144, 144}) \ - ->Args({4096, 160, 160})->Args({4096, 176, 176})->Args({8192, 128, 128}) + ->Args({4096, 160, 160})->Args({4096, 176, 176})->Args({8192, 128, 128}) \ + /* Small and thin products on both sides of the SME backend's NEON small-product bound. */ \ + ->Args({17, 17, 17})->Args({20, 20, 20})->Args({25, 25, 25})->Args({28, 28, 28})->Args({31, 31, 31}) \ + ->Args({4, 32, 32})->Args({8, 32, 32})->Args({32, 4, 32})->Args({32, 8, 32})->Args({16, 16, 32}) \ + ->Args({24, 24, 48})->Args({12, 12, 64})->Args({2, 64, 64})->Args({8, 64, 128})->Args({64, 8, 128}) // Depth sweeps through the packers' transposer, which packs b and a row-major a; their leading dimension is k, so // these also sweep the source stride. At SVL 512, k = 41 (complex<double>) and 81 (double) end in a paired step and a
diff --git a/doc/PreprocessorDirectives.dox b/doc/PreprocessorDirectives.dox index cc7f710..86c6e95 100644 --- a/doc/PreprocessorDirectives.dox +++ b/doc/PreprocessorDirectives.dox
@@ -151,7 +151,14 @@ the panel bound; a block of such panels no wider than the width bound, or the shallow width up to the shallow depth, or narrower than the thin bound on one side). One value overrides the per-scalar defaults, which are fitted for Apple M4. \c EIGEN_SME_NO_NEON_SMALL_BLOCKS and \c EIGEN_SME_FORCE_NEON_SMALL_BLOCKS pin the path for A/B runs; - both also turn off the NEON kernel for results of at most 8 x 8 (float) or 4 x 8 (double). + both also turn off the NEON kernel for results of at most 8 x 8 (float) or 4 x 8 (double), and the small-product + path below. + - \b \c EIGEN_SME_NEON_SMALL_PRODUCT - Largest rows x cols x depth at which the ARM SME backend runs a whole + general matrix product on the NEON kernel, without packing it for the ZA tiles. A result of at most two NEON + vectors of rows also takes it while the right-hand side is at most 64 KB (and, for real scalars, the depth at most + 128), and a real result of at most three columns and at least three LHS panels of rows never does. One value + overrides the per-scalar defaults (28672 for float, 10240 for double and complex<float>, 5120 for + complex<double>), which are fitted for Apple M4 Pro and checked on Apple M4. - \b \c EIGEN_SME_UNITS - Number of ARM SME units a product on the SME GEMM kernel spreads over; more threads than units only contend for them. Detected from the core topology by default (see \c Eigen::nbSmeUnits() and \c Eigen::setNbSmeUnits()): on macOS from the performance clusters, and 0 (no cap) elsewhere. Define it on other
diff --git a/test/product_sme.cpp b/test/product_sme.cpp index 4cc29d6..7a43822 100644 --- a/test/product_sme.cpp +++ b/test/product_sme.cpp
@@ -16,8 +16,28 @@ // and they keep the generic kernel, so their subtests pack their cases through // that path instead. +#ifdef EIGEN_TEST_PART_6 +#undef EIGEN_SME_NEON_SMALL_PRODUCT +#define EIGEN_SME_NEON_SMALL_PRODUCT 65536 +#undef EIGEN_DEFAULT_DENSE_INDEX_TYPE +#define EIGEN_DEFAULT_DENSE_INDEX_TYPE int +#endif + +#ifdef EIGEN_TEST_PART_7 +#ifndef EIGEN_SME_NO_NEON_SMALL_BLOCKS +#define EIGEN_SME_NO_NEON_SMALL_BLOCKS +#endif +#endif + +#ifdef EIGEN_TEST_PART_8 +#ifndef EIGEN_SME_FORCE_NEON_SMALL_BLOCKS +#define EIGEN_SME_FORCE_NEON_SMALL_BLOCKS +#endif +#endif + #include "product.h" +#include <array> #include <cfenv> // Without the right -march flags, __ARM_FEATURE_SME is undefined and @@ -559,6 +579,102 @@ } } +// Whole products on the NEON kernel (sme_small_gemm) and just past its bound, every storage order of A, B and C. +template <typename Scalar, int AO, int BO, int CO> +static void verify_small_products(Index m, Index n, Index k) { + using MA = Matrix<Scalar, Dynamic, Dynamic, AO>; + using MB = Matrix<Scalar, Dynamic, Dynamic, BO>; + using MC = Matrix<Scalar, Dynamic, Dynamic, CO>; + const MA A = MA::Random(m, k); + const MB B = MB::Random(k, n); + const MC ref = A.lazyProduct(B); + const Scalar alpha = nontrivial_alpha_impl<Scalar>::run(); + MC C = MC::Random(m, n); + const MC C0 = C; + C.noalias() = A * B; + VERIFY_IS_APPROX(C, ref); + C = C0; + C.noalias() += alpha * A * B; + VERIFY_IS_APPROX(C, C0 + alpha * ref); + const SmeColMajorMat<Scalar> Ac = A.conjugate(), Bc = B.conjugate(); + C.noalias() = A.conjugate() * B; + VERIFY_IS_APPROX(SmeColMajorMat<Scalar>(C), SmeColMajorMat<Scalar>(Ac.lazyProduct(B))); + C.noalias() = A * B.conjugate(); + VERIFY_IS_APPROX(SmeColMajorMat<Scalar>(C), SmeColMajorMat<Scalar>(A.lazyProduct(Bc))); + const MA At = A.adjoint(); + C.noalias() = At.adjoint() * B; + VERIFY_IS_APPROX(C, ref); +} + +template <typename Scalar> +static void test_small_products() { +#if defined(EIGEN_SME_NO_NEON_SMALL_BLOCKS) || defined(EIGEN_SME_FORCE_NEON_SMALL_BLOCKS) || defined(EIGEN_USE_BLAS) + constexpr bool kSmallProductEnabled = false; +#else + constexpr bool kSmallProductEnabled = true; +#endif + const Index size = internal::sme_neon_small_product<Scalar>::value; + const Index thin = Index(32 / sizeof(Scalar)); + // Both sides of the RHS bound for a thin LHS and of the size bound past it, a tall real LHS against three columns, + // then squares and thin shapes on either side. + std::vector<std::array<Index, 3>> shapes; + const Index rhs_depth = + numext::mini(Index(65536 / sizeof(Scalar)) / 37, Index(NumTraits<Scalar>::IsComplex ? 4096 : 128)); + shapes.push_back({thin, 37, rhs_depth}); + shapes.push_back({thin, 37, rhs_depth + 1}); + // Deeper than the blocking's kc, which packs into a buffer of its own. + shapes.push_back({thin, 4, 2048}); + for (Index n : {Index(16), Index(37)}) { + const Index k = size / ((thin + 1) * n); + shapes.push_back({thin + 1, n, k}); + shapes.push_back({thin + 1, n, k + 1}); + } + shapes.push_back({3 * sme_mr<Scalar>(), 3, 16}); + for (Index s : {12, 17, 20, 25, 29, 31}) shapes.push_back({s, s, s}); + for (auto s : std::vector<std::array<Index, 3>>{ + {2, 64, 48}, {4, 32, 64}, {8, 64, 64}, {32, 4, 32}, {64, 2, 64}, {16, 16, 32}, {12, 12, 64}, {3, 33, 97}}) + shapes.push_back(s); + int on_neon = 0, past = 0; + for (const auto& s : shapes) { + const bool neon = internal::sme_small_gemm_wins<Scalar>(s[0], s[1], s[2]); + on_neon += neon; + past += !neon; + verify_small_products<Scalar, ColMajor, ColMajor, ColMajor>(s[0], s[1], s[2]); + verify_small_products<Scalar, RowMajor, ColMajor, ColMajor>(s[0], s[1], s[2]); + verify_small_products<Scalar, ColMajor, RowMajor, ColMajor>(s[0], s[1], s[2]); + verify_small_products<Scalar, RowMajor, RowMajor, ColMajor>(s[0], s[1], s[2]); + verify_small_products<Scalar, ColMajor, ColMajor, RowMajor>(s[0], s[1], s[2]); + verify_small_products<Scalar, RowMajor, RowMajor, RowMajor>(s[0], s[1], s[2]); + } + VERIFY_IS_EQUAL(on_neon > 0, kSmallProductEnabled); + VERIFY(past > 0); + // Sub-blocks of larger operands and a result with an inner stride. + const Index m = thin + 3, n = 9, k = 33; + VERIFY_IS_EQUAL(internal::sme_small_gemm_wins<Scalar>(m, n, k), kSmallProductEnabled && m * n * k <= size); + const SmeColMajorMat<Scalar> Abig = SmeColMajorMat<Scalar>::Random(m + 5, k + 2); + const SmeRowMajorMat<Scalar> Bbig = SmeRowMajorMat<Scalar>::Random(k + 3, n + 4); + const auto A = Abig.block(2, 1, m, k); + const auto B = Bbig.block(1, 2, k, n); + SmeColMajorMat<Scalar> Cbig = SmeColMajorMat<Scalar>::Random(2 * m, n); + SmeColMajorMat<Scalar> ref = Cbig; + Map<SmeColMajorMat<Scalar>, 0, Stride<Dynamic, 2>> C(Cbig.data(), m, n, Stride<Dynamic, 2>(2 * m, 2)); + C.noalias() -= Scalar(3) * A * B; + Map<SmeColMajorMat<Scalar>, 0, Stride<Dynamic, 2>>(ref.data(), m, n, Stride<Dynamic, 2>(2 * m, 2)) -= + Scalar(3) * A.lazyProduct(B); + VERIFY_IS_APPROX(Cbig, ref); +} + +template <typename TestIndex> +static void test_small_product_overflow() { + // With a 65536 threshold, both the two-factor and three-factor products can wrap a 32-bit index to zero. + VERIFY(!internal::sme_small_gemm_wins<float>(TestIndex(65536), TestIndex(65536), TestIndex(1))); + VERIFY(!internal::sme_small_gemm_wins<float>(TestIndex(256), TestIndex(256), TestIndex(65536))); + VERIFY(!internal::sme_small_gemm_wins<float>(TestIndex(16), TestIndex(16), TestIndex(257))); +#if !defined(EIGEN_SME_NO_NEON_SMALL_BLOCKS) && !defined(EIGEN_SME_FORCE_NEON_SMALL_BLOCKS) && !defined(EIGEN_USE_BLAS) + VERIFY(internal::sme_small_gemm_wins<float>(TestIndex(16), TestIndex(16), TestIndex(256))); +#endif +} + // Every tail width of a deep block, which the ZA transposer packs: a ColMajor RHS tail and, through a RowMajor // product, a transposed LHS tail. template <typename Scalar> @@ -1389,6 +1505,7 @@ CALL_SUBTEST_1(test_pack_direct<float>()); CALL_SUBTEST_1(test_mapper_fallback<float>()); CALL_SUBTEST_1(test_neon_small_blocks<float>()); + CALL_SUBTEST_1(test_small_products<float>()); CALL_SUBTEST_1(test_direct_lhs<float>()); CALL_SUBTEST_1(test_small_k_and_single_panel_rhs<float>()); CALL_SUBTEST_1(test_deep_tail_panels<float>()); @@ -1408,6 +1525,7 @@ CALL_SUBTEST_2(test_pack_direct<double>()); CALL_SUBTEST_2(test_mapper_fallback<double>()); CALL_SUBTEST_2(test_neon_small_blocks<double>()); + CALL_SUBTEST_2(test_small_products<double>()); CALL_SUBTEST_2(test_direct_lhs<double>()); CALL_SUBTEST_2(test_small_k_and_single_panel_rhs<double>()); CALL_SUBTEST_2(test_deep_tail_panels<double>()); @@ -1422,6 +1540,7 @@ CALL_SUBTEST_3(test_pack_direct<std::complex<float>>()); CALL_SUBTEST_3(test_mapper_fallback<std::complex<float>>()); CALL_SUBTEST_3(test_neon_small_blocks<std::complex<float>>()); + CALL_SUBTEST_3(test_small_products<std::complex<float>>()); CALL_SUBTEST_3(test_fp_flags<std::complex<float>>()); CALL_SUBTEST_3(test_disjoint_parts<std::complex<float>>()); CALL_SUBTEST_3((test_complex_alpha_range<float, ColMajor>())); @@ -1437,6 +1556,7 @@ CALL_SUBTEST_4(test_pack_direct<std::complex<double>>()); CALL_SUBTEST_4(test_mapper_fallback<std::complex<double>>()); CALL_SUBTEST_4(test_neon_small_blocks<std::complex<double>>()); + CALL_SUBTEST_4(test_small_products<std::complex<double>>()); CALL_SUBTEST_4(test_fp_flags<std::complex<double>>()); CALL_SUBTEST_4(test_disjoint_parts<std::complex<double>>()); CALL_SUBTEST_4((test_complex_alpha_range<double, ColMajor>())); @@ -1448,4 +1568,17 @@ // 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))); + + CALL_SUBTEST_6(test_small_product_overflow<int>()); + CALL_SUBTEST_6(test_small_product_overflow<int64_t>()); + CALL_SUBTEST_6(test_small_products<float>()); + +#if defined(EIGEN_TEST_PART_7) || defined(EIGEN_TEST_PART_8) + CALL_SUBTEST(test_small_products<float>()); + CALL_SUBTEST(test_small_products<std::complex<float>>()); +#ifdef EIGEN_VECTORIZE_SME_F64F64 + CALL_SUBTEST(test_small_products<double>()); + CALL_SUBTEST(test_small_products<std::complex<double>>()); +#endif +#endif }