Core: Let GEMM overwrite the destination for C = A*B instead of zeroing it first libeigen/eigen!3371
diff --git a/Eigen/src/Core/arch/SME/GeneralBlockPanelKernel.h b/Eigen/src/Core/arch/SME/GeneralBlockPanelKernel.h index d81759e..59f345b 100644 --- a/Eigen/src/Core/arch/SME/GeneralBlockPanelKernel.h +++ b/Eigen/src/Core/arch/SME/GeneralBlockPanelKernel.h
@@ -1476,6 +1476,45 @@ #undef EIGEN_SME_DECLARE_GEMM_PACKERS +// Overwriting stores (gebp_overwrite_t) still load C, then clear the loaded bits and accumulate onto that +0: the +// result is what accumulating into a zeroed C gives, whatever C held, and the memory accesses stay those of +// accumulation. On Apple M4, stores into lines the SME unit has not just loaded measure slower than the load, FMA +// and store they replace, so dropping the loads loses at many sizes. The clearing predicate is opaque to the +// compiler, which would otherwise fold the predicated zero and the loads away. +template <bool Overwrite> +EIGEN_ALWAYS_INLINE svbool_t sme_c_mask() __arm_streaming EIGEN_SME_ZA_AGNOSTIC { + // b32 sets the predicate bit of every 32-bit and every 64-bit lane. + svbool_t m = Overwrite ? svptrue_b32() : svpfalse_b(); + EIGEN_IF_CONSTEXPR (Overwrite) __asm__("" : "+Upa"(m)); + return m; +} +template <bool Overwrite> +EIGEN_ALWAYS_INLINE svfloat32_t sme_c_value(svbool_t m, svfloat32_t v) __arm_streaming EIGEN_SME_ZA_AGNOSTIC { + return Overwrite ? svsel_f32(m, svdup_f32(0.0f), v) : v; +} +template <bool Overwrite> +EIGEN_ALWAYS_INLINE svfloat64_t sme_c_value(svbool_t m, svfloat64_t v) __arm_streaming EIGEN_SME_ZA_AGNOSTIC { + return Overwrite ? svsel_f64(m, svdup_f64(0.0), v) : v; +} + +// Zeroes a rows x cols block of C with streaming stores, so that the SME unit, which accumulates into it next, is the +// first to write its lines. +template <typename Scalar, typename Index> +EIGEN_DONT_INLINE __arm_locally_streaming void sme_zero_c(Scalar* C, Index rs, Index cs, Index rows, Index cols) { + if (rs == 1 || cs == 1) { + const Index inner = rs == 1 ? rows : cols, outer = rs == 1 ? cols : rows, stride = rs == 1 ? cs : rs; + const int64_t bytes = int64_t(inner) * int64_t(sizeof(Scalar)), vl = int64_t(svcntb()); + const svuint8_t z = svdup_u8(0); + for (Index o = 0; o < outer; ++o) { + uint8_t* p = reinterpret_cast<uint8_t*>(C + o * stride); + for (int64_t b = 0; b < bytes; b += vl) svst1_u8(svwhilelt_b8_s64(b, bytes), p + b, z); + } + } else { + for (Index c = 0; c < cols; ++c) + for (Index r = 0; r < rows; ++r) C[r * rs + c * cs] = Scalar(0); + } +} + /***************************************************************************** * sme_store_za_tile -- Store one ZA tile back to C with alpha scaling. * @@ -1483,12 +1522,13 @@ * (both <= the runtime svl). *****************************************************************************/ -template <typename Scalar, int TileId, typename Index> +template <typename Scalar, int TileId, bool Overwrite = false, typename Index> EIGEN_ALWAYS_INLINE void sme_store_za_tile(Scalar* EIGEN_RESTRICT C, Index C_stride_row, Index C_stride_col, Scalar alpha, Index row_start, int pw, Index col_start, int cw) __arm_streaming __arm_inout("za") { using Traits = sme_packet_traits<Scalar>; using Vec = typename Traits::type; + const svbool_t cmask = sme_c_mask<Overwrite>(); const svbool_t pg_m = Traits::whilelt(0, pw); const svbool_t pg_n = Traits::whilelt(0, cw); // FMLA and FADD have equal latency/throughput on ARMv9 cores, and @@ -1509,14 +1549,14 @@ for (; ci + 2 <= cw; ci += 2) { Scalar* p0 = C + row_start + (col_start + ci) * C_stride_col; Scalar* p1 = p0 + C_stride_col; - Vec c0 = ploadu(pg_m, p0); - Vec c1 = ploadu(pg_m, p1); + Vec c0 = sme_c_value<Overwrite>(cmask, ploadu(pg_m, p0)); + Vec c1 = sme_c_value<Overwrite>(cmask, ploadu(pg_m, p1)); pstoreu(pg_m, p0, pmadd(pg_m, sme_read_ver_za<TileId>(vzero, pg_m, (uint32_t)ci), valpha, c0)); pstoreu(pg_m, p1, pmadd(pg_m, sme_read_ver_za<TileId>(vzero, pg_m, (uint32_t)(ci + 1)), valpha, c1)); } if (ci < cw) { Scalar* pC = C + row_start + (col_start + ci) * C_stride_col; - Vec vc = ploadu(pg_m, pC); + Vec vc = sme_c_value<Overwrite>(cmask, ploadu(pg_m, pC)); pstoreu(pg_m, pC, pmadd(pg_m, sme_read_ver_za<TileId>(vzero, pg_m, (uint32_t)ci), valpha, vc)); } } else if (C_stride_col == 1) { @@ -1525,14 +1565,14 @@ for (; ri + 2 <= pw; ri += 2) { Scalar* p0 = C + (row_start + ri) * C_stride_row + col_start; Scalar* p1 = p0 + C_stride_row; - Vec c0 = ploadu(pg_n, p0); - Vec c1 = ploadu(pg_n, p1); + Vec c0 = sme_c_value<Overwrite>(cmask, ploadu(pg_n, p0)); + Vec c1 = sme_c_value<Overwrite>(cmask, ploadu(pg_n, p1)); pstoreu(pg_n, p0, pmadd(pg_n, sme_read_hor_za<TileId>(vzero, pg_n, (uint32_t)ri), valpha, c0)); pstoreu(pg_n, p1, pmadd(pg_n, sme_read_hor_za<TileId>(vzero, pg_n, (uint32_t)(ri + 1)), valpha, c1)); } if (ri < pw) { Scalar* pC = C + (row_start + ri) * C_stride_row + col_start; - Vec vc = ploadu(pg_n, pC); + Vec vc = sme_c_value<Overwrite>(cmask, ploadu(pg_n, pC)); pstoreu(pg_n, pC, pmadd(pg_n, sme_read_hor_za<TileId>(vzero, pg_n, (uint32_t)ri), valpha, vc)); } } else { @@ -1546,7 +1586,8 @@ vres = pmul(pg_n, vres, valpha); pstoreu(pg_n, scratch, vres); for (int ci = 0; ci < cw; ++ci) { - C[(row_start + ri) * C_stride_row + (col_start + ci) * C_stride_col] += scratch[ci]; + Scalar& c = C[(row_start + ri) * C_stride_row + (col_start + ci) * C_stride_col]; + c = (Overwrite ? Scalar(0) : c) + scratch[ci]; } } } @@ -1566,13 +1607,14 @@ // Full 2*svl x 2*svl block into column-major C: four columns per vertical four-slice read of the two tiles that // share them, each column loaded and stored whole with two-vector accesses, all four loaded before any is stored. -template <int TileLo, int TileHi, typename Scalar, typename Index> +template <int TileLo, int TileHi, bool Overwrite = false, typename Scalar, typename Index> EIGEN_ALWAYS_INLINE void sme_store_tile_pair_colmajor(Scalar* EIGEN_RESTRICT C, Index ldc, Scalar alpha, int svl) __arm_streaming __arm_inout("za") { using Traits = sme_packet_traits<Scalar>; const svcount_t pn = Traits::ptrue_c(); const svbool_t pg = Traits::ptrue(); const typename Traits::type valpha = pset1<typename Traits::type>(alpha); + const svbool_t m = sme_c_mask<Overwrite>(); for (int c = 0; c < svl; c += 4) { // GCC 14 can overwrite live ZA tuple components with C loads. #if !EIGEN_COMP_GNUC_STRICT @@ -1587,35 +1629,40 @@ const auto hi = sme_read_ver_za_vg4<TileHi, Scalar>(uint32_t(c)); #endif pstoreu_x2(pn, p, - pcreate(pmadd(pg, pget<0>(lo), valpha, pget<0>(c0)), pmadd(pg, pget<0>(hi), valpha, pget<1>(c0)))); + pcreate(pmadd(pg, pget<0>(lo), valpha, sme_c_value<Overwrite>(m, pget<0>(c0))), + pmadd(pg, pget<0>(hi), valpha, sme_c_value<Overwrite>(m, pget<1>(c0))))); pstoreu_x2(pn, p + ldc, - pcreate(pmadd(pg, pget<1>(lo), valpha, pget<0>(c1)), pmadd(pg, pget<1>(hi), valpha, pget<1>(c1)))); + pcreate(pmadd(pg, pget<1>(lo), valpha, sme_c_value<Overwrite>(m, pget<0>(c1))), + pmadd(pg, pget<1>(hi), valpha, sme_c_value<Overwrite>(m, pget<1>(c1))))); pstoreu_x2(pn, p + 2 * ldc, - pcreate(pmadd(pg, pget<2>(lo), valpha, pget<0>(c2)), pmadd(pg, pget<2>(hi), valpha, pget<1>(c2)))); + pcreate(pmadd(pg, pget<2>(lo), valpha, sme_c_value<Overwrite>(m, pget<0>(c2))), + pmadd(pg, pget<2>(hi), valpha, sme_c_value<Overwrite>(m, pget<1>(c2))))); pstoreu_x2(pn, p + 3 * ldc, - pcreate(pmadd(pg, pget<3>(lo), valpha, pget<0>(c3)), pmadd(pg, pget<3>(hi), valpha, pget<1>(c3)))); + pcreate(pmadd(pg, pget<3>(lo), valpha, sme_c_value<Overwrite>(m, pget<0>(c3))), + pmadd(pg, pget<3>(hi), valpha, sme_c_value<Overwrite>(m, pget<1>(c3))))); } } -template <typename Scalar, typename Index> +template <bool Overwrite = false, typename Scalar, typename Index> EIGEN_ALWAYS_INLINE void sme_store_2x2_grid(Scalar* EIGEN_RESTRICT C, Index C_stride_row, Index C_stride_col, Scalar alpha, Index row_start, int rlo, int rhi, Index col_start, int clo, int chi) __arm_streaming __arm_inout("za") { const int svl = sme_packet_traits<Scalar>::size(); if (svl >= 4 && C_stride_row == 1 && rlo == svl && rhi == svl && clo == svl && chi == svl) { Scalar* c0 = C + row_start + col_start * C_stride_col; - sme_store_tile_pair_colmajor<0, 2>(c0, C_stride_col, alpha, svl); - sme_store_tile_pair_colmajor<1, 3>(c0 + Index(svl) * C_stride_col, C_stride_col, alpha, svl); + sme_store_tile_pair_colmajor<0, 2, Overwrite>(c0, C_stride_col, alpha, svl); + sme_store_tile_pair_colmajor<1, 3, Overwrite>(c0 + Index(svl) * C_stride_col, C_stride_col, alpha, svl); return; } - sme_store_za_tile<Scalar, 0>(C, C_stride_row, C_stride_col, alpha, row_start, rlo, col_start, clo); + sme_store_za_tile<Scalar, 0, Overwrite>(C, C_stride_row, C_stride_col, alpha, row_start, rlo, col_start, clo); if (chi > 0) { - sme_store_za_tile<Scalar, 1>(C, C_stride_row, C_stride_col, alpha, row_start, rlo, col_start + svl, chi); + sme_store_za_tile<Scalar, 1, Overwrite>(C, C_stride_row, C_stride_col, alpha, row_start, rlo, col_start + svl, chi); } if (rhi > 0) { - sme_store_za_tile<Scalar, 2>(C, C_stride_row, C_stride_col, alpha, row_start + svl, rhi, col_start, clo); + sme_store_za_tile<Scalar, 2, Overwrite>(C, C_stride_row, C_stride_col, alpha, row_start + svl, rhi, col_start, clo); if (chi > 0) { - sme_store_za_tile<Scalar, 3>(C, C_stride_row, C_stride_col, alpha, row_start + svl, rhi, col_start + svl, chi); + sme_store_za_tile<Scalar, 3, Overwrite>(C, C_stride_row, C_stride_col, alpha, row_start + svl, rhi, + col_start + svl, chi); } } } @@ -1909,8 +1956,8 @@ template <bool ConjLhs, bool ConjRhs, typename Scalar, typename Index> EIGEN_ALWAYS_INLINE void sme_process(Scalar* EIGEN_RESTRICT C, Index C_stride_row, Index C_stride_col, const Scalar* EIGEN_RESTRICT blA, const Scalar* EIGEN_RESTRICT blB, Index depth, - Scalar alpha, Index row_start, int pw, Index col_start, int cw, - Index a_step) __arm_streaming __arm_inout("za") { + Scalar alpha, Index row_start, int pw, Index col_start, int cw, Index a_step, + bool overwrite) __arm_streaming __arm_inout("za") { // Conjugation is the identity on real scalars, so this overload ignores it. // a_step is the distance between A's depth steps: pw for a packed panel, the // column stride for a ColMajor source read in place. @@ -1998,7 +2045,12 @@ } // Store the (up to) 2x2 grid of tiles for this sub-block pass. - sme_store_2x2_grid(C, C_stride_row, C_stride_col, alpha, row_start + rt, rlo, rhi, col_start + ct, clo, chi); + if (overwrite) + sme_store_2x2_grid<true>(C, C_stride_row, C_stride_col, alpha, row_start + rt, rlo, rhi, col_start + ct, clo, + chi); + else + sme_store_2x2_grid<false>(C, C_stride_row, C_stride_col, alpha, row_start + rt, rlo, rhi, col_start + ct, clo, + chi); } } } @@ -2120,7 +2172,8 @@ const Scalar* EIGEN_RESTRICT blA0, Index step0, const Scalar* EIGEN_RESTRICT blA1, Index step1, int pw1, const Scalar* EIGEN_RESTRICT blB, int cw, Index depth, Scalar alpha, - Index row_start, Index col_start) __arm_streaming __arm_inout("za") { + Index row_start, Index col_start, + bool overwrite) __arm_streaming __arm_inout("za") { using Traits = sme_packet_traits<Scalar>; using Vec = typename Traits::type; const int svl = Traits::size(); @@ -2140,12 +2193,28 @@ sme_mopa<2>(p1lo, pn, a1lo, b); sme_mopa<3>(p1hi, pn, a1hi, b); } - sme_store_za_tile<Scalar, 0>(C, C_stride_row, C_stride_col, alpha, row_start, svl, col_start, cw); - sme_store_za_tile<Scalar, 1>(C, C_stride_row, C_stride_col, alpha, row_start + svl, svl, col_start, cw); - sme_store_za_tile<Scalar, 2>(C, C_stride_row, C_stride_col, alpha, row_start + 2 * svl, sme_min(pw1, svl), col_start, - cw); - if (pw1 > svl) - sme_store_za_tile<Scalar, 3>(C, C_stride_row, C_stride_col, alpha, row_start + 3 * svl, pw1 - svl, col_start, cw); + if (overwrite) + sme_store_za_tile<Scalar, 0, true>(C, C_stride_row, C_stride_col, alpha, row_start, svl, col_start, cw); + else + sme_store_za_tile<Scalar, 0, false>(C, C_stride_row, C_stride_col, alpha, row_start, svl, col_start, cw); + if (overwrite) + sme_store_za_tile<Scalar, 1, true>(C, C_stride_row, C_stride_col, alpha, row_start + svl, svl, col_start, cw); + else + sme_store_za_tile<Scalar, 1, false>(C, C_stride_row, C_stride_col, alpha, row_start + svl, svl, col_start, cw); + if (overwrite) + sme_store_za_tile<Scalar, 2, true>(C, C_stride_row, C_stride_col, alpha, row_start + 2 * svl, sme_min(pw1, svl), + col_start, cw); + else + sme_store_za_tile<Scalar, 2, false>(C, C_stride_row, C_stride_col, alpha, row_start + 2 * svl, sme_min(pw1, svl), + col_start, cw); + if (pw1 > svl) { + if (overwrite) + sme_store_za_tile<Scalar, 3, true>(C, C_stride_row, C_stride_col, alpha, row_start + 3 * svl, pw1 - svl, + col_start, cw); + else + sme_store_za_tile<Scalar, 3, false>(C, C_stride_row, C_stride_col, alpha, row_start + 3 * svl, pw1 - svl, + col_start, cw); + } } // Tile Dst += tile Src, four columns per step. @@ -2166,7 +2235,7 @@ EIGEN_ALWAYS_INLINE void sme_process_split(Scalar* EIGEN_RESTRICT C, Index C_stride_row, Index C_stride_col, const Scalar* EIGEN_RESTRICT blA, const Scalar* EIGEN_RESTRICT blB, Index depth, Scalar alpha, Index row_start, int pw, Index col_start, int cw, - Index a_step) __arm_streaming __arm_inout("za") { + Index a_step, bool overwrite) __arm_streaming __arm_inout("za") { using Traits = sme_packet_traits<Scalar>; using Vec = typename Traits::type; const int svl = Traits::size(); @@ -2189,7 +2258,10 @@ sme_fold_tile<0, 1, Scalar>(); sme_fold_tile<2, 3, Scalar>(); sme_fold_tile<0, 2, Scalar>(); - sme_store_za_tile<Scalar, 0>(C, C_stride_row, C_stride_col, alpha, row_start, pw, col_start, cw); + if (overwrite) + sme_store_za_tile<Scalar, 0, true>(C, C_stride_row, C_stride_col, alpha, row_start, pw, col_start, cw); + else + sme_store_za_tile<Scalar, 0, false>(C, C_stride_row, C_stride_col, alpha, row_start, pw, col_start, cw); } else if (pw <= svl) { // Tiles 0 and 1 take the even depth steps, 2 and 3 the odd ones. for (; k + 2 <= depth; k += 2) { @@ -2208,7 +2280,10 @@ } sme_fold_tile<0, 2, Scalar>(); sme_fold_tile<1, 3, Scalar>(); - sme_store_2x2_grid(C, C_stride_row, C_stride_col, alpha, row_start, pw, 0, col_start, svl, cw - svl); + if (overwrite) + sme_store_2x2_grid<true>(C, C_stride_row, C_stride_col, alpha, row_start, pw, 0, col_start, svl, cw - svl); + else + sme_store_2x2_grid<false>(C, C_stride_row, C_stride_col, alpha, row_start, pw, 0, col_start, svl, cw - svl); } else { // cw <= svl < pw: tiles 0 and 2 take the even depth steps, 1 and 3 the odd ones. for (; k + 2 <= depth; k += 2) { @@ -2227,29 +2302,35 @@ } sme_fold_tile<0, 1, Scalar>(); sme_fold_tile<2, 3, Scalar>(); - sme_store_2x2_grid(C, C_stride_row, C_stride_col, alpha, row_start, svl, pw - svl, col_start, cw, 0); + if (overwrite) + sme_store_2x2_grid<true>(C, C_stride_row, C_stride_col, alpha, row_start, svl, pw - svl, col_start, cw, 0); + else + sme_store_2x2_grid<false>(C, C_stride_row, C_stride_col, alpha, row_start, svl, pw - svl, col_start, cw, 0); } } // One pw x cw block: sme_process_split when it fills at most two tiles of the 2 x 2 grid; complex blocks keep -// sme_process. `split_ok` (sme_split_ok) holds once per call, outside the block loops. +// sme_process. `split_ok` (sme_split_ok) holds once per call, outside the block loops. `overwrite` assigns the block +// instead of accumulating into it; a complex overwrite zeroes C up front (sme_zero_c) and then accumulates. The flag is +// a runtime argument and each store picks its compile-time form, so one kernel serves both modes, accumulation runs +// the plain accumulating stores, and only the store epilogues are compiled twice. template <bool ConjLhs, bool ConjRhs, bool FoldRealAlpha, typename Scalar, typename Index> EIGEN_ALWAYS_INLINE void sme_process_block(Scalar* C, Index rs, Index cs, const Scalar* blA, const Scalar* blB, Index depth, Scalar alpha, Index row_start, int pw, Index col_start, int cw, - Index a_step, bool split_ok, Index, - Index) __arm_streaming __arm_inout("za") { + Index a_step, bool split_ok, Index, Index, + bool overwrite) __arm_streaming __arm_inout("za") { const int svl = sme_packet_traits<Scalar>::size(); if (split_ok && (pw <= svl || cw <= svl)) - sme_process_split(C, rs, cs, blA, blB, depth, alpha, row_start, pw, col_start, cw, a_step); + sme_process_split(C, rs, cs, blA, blB, depth, alpha, row_start, pw, col_start, cw, a_step, overwrite); else - sme_process<ConjLhs, ConjRhs>(C, rs, cs, blA, blB, depth, alpha, row_start, pw, col_start, cw, a_step); + sme_process<ConjLhs, ConjRhs>(C, rs, cs, blA, blB, depth, alpha, row_start, pw, col_start, cw, a_step, overwrite); } template <bool ConjLhs, bool ConjRhs, bool FoldRealAlpha, typename RealScalar, typename Index> EIGEN_ALWAYS_INLINE void sme_process_block(std::complex<RealScalar>* C, Index rs, Index cs, const std::complex<RealScalar>* blA, const std::complex<RealScalar>* blB, Index depth, std::complex<RealScalar> alpha, Index row_start, int pw, - Index col_start, int cw, Index a_step, bool, Index rows, - Index cols) __arm_streaming __arm_inout("za") { + Index col_start, int cw, Index a_step, bool, Index rows, Index cols, + bool) __arm_streaming __arm_inout("za") { sme_process<ConjLhs, ConjRhs, FoldRealAlpha>(C, rs, cs, blA, blB, depth, alpha, row_start, pw, col_start, cw, a_step, rows, cols); } @@ -2271,16 +2352,17 @@ EIGEN_ALWAYS_INLINE void sme_process_narrow_dispatch(Scalar* C, Index rs, Index cs, const Scalar* blA0, Index step0, const Scalar* blA1, Index step1, int pw1, const Scalar* blB, int cw, Index depth, Scalar alpha, Index row_start, - Index col_start) __arm_streaming __arm_inout("za") { - sme_process_narrow(C, rs, cs, blA0, step0, blA1, step1, pw1, blB, cw, depth, alpha, row_start, col_start); + Index col_start, + bool overwrite) __arm_streaming __arm_inout("za") { + sme_process_narrow(C, rs, cs, blA0, step0, blA1, step1, pw1, blB, cw, depth, alpha, row_start, col_start, overwrite); } template <typename RealScalar, typename Index> EIGEN_ALWAYS_INLINE void sme_process_narrow_dispatch(std::complex<RealScalar>*, Index, Index, const std::complex<RealScalar>*, Index, const std::complex<RealScalar>*, Index, int, const std::complex<RealScalar>*, int, Index, - std::complex<RealScalar>, Index, - Index) __arm_streaming __arm_inout("za") {} + std::complex<RealScalar>, Index, Index, + bool) __arm_streaming __arm_inout("za") {} // Core-side prefetch into L2 of the column-major C block the next sme_process call reads and writes. template <typename Scalar, typename Index> @@ -2296,9 +2378,9 @@ } template <typename Scalar, bool ConjLhs, bool ConjRhs, bool FoldRealAlpha = false, typename Index> -EIGEN_DONT_INLINE __arm_locally_streaming __arm_new("za") void sme_gebp_impl( +EIGEN_DONT_INLINE EIGEN_GEMM_NO_CLONE __arm_locally_streaming __arm_new("za") void sme_gebp_impl( Scalar* C, Index C_stride_row, Index C_stride_col, const Scalar* blockA, const Scalar* blockB, Index rows, - Index depth, Index cols, Scalar alpha, Index strideA, Index strideB, Index offsetA, Index offsetB) { + Index depth, Index cols, Scalar alpha, Index strideA, Index strideB, Index offsetA, Index offsetB, bool overwrite) { constexpr int MR = sme_block<Scalar>::mr; constexpr int NR = sme_block<Scalar>::nr; @@ -2323,7 +2405,7 @@ const int pw1 = static_cast<int>(sme_min(rows - i - MR, Index(MR))); sme_process_narrow_dispatch(C, C_stride_row, C_stride_col, blockA + i * strideA + offsetA * MR, Index(MR), blockA + (i + MR) * strideA + offsetA * pw1, Index(pw1), pw1, blB, cw, depth, - alpha, i, j); + alpha, i, j, overwrite); } } for (; i < rows; i += MR) { @@ -2333,7 +2415,7 @@ sme_prefetch_next_c(C, C_stride_row, C_stride_col, i + MR < rows ? i + MR : Index(0), i + MR < rows ? j : j + NR, rows, cols, MR, NR); sme_process_block<ConjLhs, ConjRhs, FoldRealAlpha>(C, C_stride_row, C_stride_col, blA, blB, depth, alpha, i, pw, - j, cw, Index(pw), split_ok, rows, cols); + j, cw, Index(pw), split_ok, rows, cols, overwrite); } } } @@ -2347,9 +2429,9 @@ template <typename Index> static void run(Scalar* C, Index C_stride_row, Index C_stride_col, const Scalar* blockA, const Scalar* blockB, Index rows, Index depth, Index cols, Scalar alpha, Index strideA, Index strideB, Index offsetA, - Index offsetB) { + Index offsetB, bool overwrite = false) { sme_gebp_impl<Scalar, ConjLhs, ConjRhs, false>(C, C_stride_row, C_stride_col, blockA, blockB, rows, depth, cols, - alpha, strideA, strideB, offsetA, offsetB); + alpha, strideA, strideB, offsetA, offsetB, overwrite); } }; @@ -2358,14 +2440,14 @@ template <typename Index> static void run(Scalar* C, Index C_stride_row, Index C_stride_col, const Scalar* blockA, const Scalar* blockB, Index rows, Index depth, Index cols, Scalar alpha, Index strideA, Index strideB, Index offsetA, - Index offsetB) { + Index offsetB, bool overwrite = false) { using RealScalar = typename NumTraits<Scalar>::Real; if (numext::imag(alpha) == RealScalar(0) && numext::real(alpha) != RealScalar(1)) { sme_gebp_impl<Scalar, ConjLhs, ConjRhs, true>(C, C_stride_row, C_stride_col, blockA, blockB, rows, depth, cols, - alpha, strideA, strideB, offsetA, offsetB); + alpha, strideA, strideB, offsetA, offsetB, overwrite); } else { sme_gebp_impl<Scalar, ConjLhs, ConjRhs, false>(C, C_stride_row, C_stride_col, blockA, blockB, rows, depth, cols, - alpha, strideA, strideB, offsetA, offsetB); + alpha, strideA, strideB, offsetA, offsetB, overwrite); } } }; @@ -2373,9 +2455,9 @@ // gebp with the LHS read from a ColMajor source: rows i..i+pw of column k are // at lhs + i + k * lda, the packed layout with a_step = lda. Real scalars only. template <typename Scalar, typename Index> -EIGEN_DONT_INLINE __arm_locally_streaming __arm_new("za") void sme_gebp_impl_direct_lhs( +EIGEN_DONT_INLINE EIGEN_GEMM_NO_CLONE __arm_locally_streaming __arm_new("za") void sme_gebp_impl_direct_lhs( Scalar* C, Index C_stride_row, Index C_stride_col, const Scalar* lhs, Index lda, const Scalar* blockB, Index rows, - Index depth, Index cols, Scalar alpha, Index strideB, Index offsetB) { + Index depth, Index cols, Scalar alpha, Index strideB, Index offsetB, bool overwrite) { constexpr int MR = sme_block<Scalar>::mr; constexpr int NR = sme_block<Scalar>::nr; const bool split_ok = sme_split_ok<Scalar>(), narrow_ok = sme_narrow_ok<Scalar>(); @@ -2384,13 +2466,15 @@ const Scalar* blB = blockB + j * strideB + offsetB * cw; Index i = 0; if (narrow_ok && cw <= sme_packet_traits<Scalar>::size()) - for (; i + MR < rows; i += 2 * MR) - sme_process_narrow(C, C_stride_row, C_stride_col, lhs + i, lda, lhs + i + MR, lda, - static_cast<int>(sme_min(rows - i - MR, Index(MR))), blB, cw, depth, alpha, i, j); + for (; i + MR < rows; i += 2 * MR) { + const int pw1 = static_cast<int>(sme_min(rows - i - MR, Index(MR))); + sme_process_narrow(C, C_stride_row, C_stride_col, lhs + i, lda, lhs + i + MR, lda, pw1, blB, cw, depth, alpha, + i, j, overwrite); + } for (; i < rows; i += MR) { const int pw = static_cast<int>(sme_min(rows - i, Index(MR))); sme_process_block<false, false, false>(C, C_stride_row, C_stride_col, lhs + i, blB, depth, alpha, i, pw, j, cw, - lda, split_ok, rows, cols); + lda, split_ok, rows, cols, overwrite); } } } @@ -2917,10 +3001,50 @@ template <typename Scalar, typename Index, typename DataMapper, int mr, int nr, bool ConjugateLhs, bool ConjugateRhs> struct sme_gebp_kernel { using ResScalar = Scalar; + // Real kernels overwrite in their stores (sme_c_mask). Complex ones zero their C block with streaming stores + // (sme_zero_c), then accumulate: on M4 that beats the driver's non-streaming zeroing for a small result. + static constexpr bool CanOverwrite = true; + // Whether overwriting a rows x cols result pays over zeroing it first (gebp_overwrite_pays). The NEON small-block + // path has nothing to gain: it zeroes and accumulates either way. Otherwise, on Apple M4, overwriting pays for + // 32-bit scalars at every size, and for 64-bit and complex ones up to a 128 KB result, past which zeroing up front + // measures faster. + static bool overwrite_pays(Index rows, Index cols, Index depth) { + if (rows <= 0 || cols <= 0 || sme_kernel_with_neon<Scalar>(rows, cols, depth, depth, depth)) return false; + // rows * cols <= 128 KB / sizeof(Scalar), without forming a product a narrow Index cannot hold. + return sizeof(Scalar) == 4 || rows <= Index((128 * 1024) / sizeof(Scalar)) / cols; + } EIGEN_DONT_INLINE void operator()(const DataMapper& res, const Scalar* blockA, const Scalar* blockB, Index rows, Index depth, Index cols, ResScalar alpha, Index strideA = -1, Index strideB = -1, Index offsetA = 0, Index offsetB = 0) { + run(false, res, blockA, blockB, rows, depth, cols, alpha, strideA, strideB, offsetA, offsetB); + } + // C = alpha * A' * B' over the block, whatever C held (see gebp_overwrite_t). + EIGEN_DONT_INLINE void operator()(gebp_overwrite_t, const DataMapper& res, const Scalar* blockA, const Scalar* blockB, + Index rows, Index depth, Index cols, ResScalar alpha, Index strideA = -1, + Index strideB = -1, Index offsetA = 0, Index offsetB = 0) { + run(true, res, blockA, blockB, rows, depth, cols, alpha, strideA, strideB, offsetA, offsetB); + } + + // The LHS block read from its ColMajor source (see sme_direct_lhs_ok); only + // instantiated for real scalars. + EIGEN_DONT_INLINE void run_direct_lhs(const DataMapper& res, const Scalar* lhs, Index lhsStride, const Scalar* blockB, + Index rows, Index depth, Index cols, ResScalar alpha, Index strideB = -1, + Index offsetB = 0) { + run_direct_lhs_impl(false, res, lhs, lhsStride, blockB, rows, depth, cols, alpha, strideB, offsetB); + } + EIGEN_DONT_INLINE void run_direct_lhs(gebp_overwrite_t, const DataMapper& res, const Scalar* lhs, Index lhsStride, + const Scalar* blockB, Index rows, Index depth, Index cols, ResScalar alpha, + Index strideB = -1, Index offsetB = 0) { + run_direct_lhs_impl(true, res, lhs, lhsStride, blockB, rows, depth, cols, alpha, strideB, offsetB); + } + + private: + // One body for both modes, so that the kernels are compiled once (see sme_process_block). + EIGEN_DONT_INLINE EIGEN_GEMM_NO_CLONE void run(bool overwrite, const DataMapper& res, const Scalar* blockA, + const Scalar* blockB, Index rows, Index depth, Index cols, + ResScalar alpha, Index strideA, Index strideB, Index offsetA, + Index offsetB) { // Real scalars never reach the kernel conjugated (conj_helper folds it into // the identity long before), so the real path stays free of the flags. static_assert(NumTraits<Scalar>::IsComplex || (!ConjugateLhs && !ConjugateRhs), @@ -2931,36 +3055,50 @@ if (strideA == -1) strideA = depth; if (strideB == -1) strideB = depth; - if (rows <= 0 || cols <= 0 || depth <= 0) return; + if (rows <= 0 || cols <= 0) return; Scalar* C_base = const_cast<Scalar*>(&res(0, 0)); const Index C_stride_row = &res(1, 0) - &res(0, 0); const Index C_stride_col = &res(0, 1) - &res(0, 0); + if (depth <= 0) { + if (overwrite) gemm_zero_result(C_base, rows, cols, C_stride_row, C_stride_col); + return; + } if (sme_kernel_with_neon<Scalar>(rows, cols, depth, strideA, strideB)) { + if (overwrite) gemm_zero_result(C_base, rows, cols, C_stride_row, C_stride_col); sme_gebp_neon<Scalar, ConjugateLhs, ConjugateRhs>(C_base, C_stride_row, C_stride_col, blockA, blockB, rows, depth, cols, alpha, strideA, strideB, offsetA, offsetB); return; } sme_fpsr_guard fpsr; + const bool zero_first = overwrite && NumTraits<Scalar>::IsComplex; + if (zero_first) sme_zero_c(C_base, C_stride_row, C_stride_col, rows, cols); sme_gebp_dispatch<Scalar, ConjugateLhs, ConjugateRhs>::run(C_base, C_stride_row, C_stride_col, blockA, blockB, rows, - depth, cols, alpha, strideA, strideB, offsetA, offsetB); + depth, cols, alpha, strideA, strideB, offsetA, offsetB, + overwrite && !zero_first); } - // The LHS block read from its ColMajor source (see sme_direct_lhs_ok); only - // instantiated for real scalars. - EIGEN_DONT_INLINE void run_direct_lhs(const DataMapper& res, const Scalar* lhs, Index lhsStride, const Scalar* blockB, - Index rows, Index depth, Index cols, ResScalar alpha, Index strideB = -1, - Index offsetB = 0) { + EIGEN_DONT_INLINE EIGEN_GEMM_NO_CLONE void run_direct_lhs_impl(bool overwrite, const DataMapper& res, + const Scalar* lhs, Index lhsStride, + const Scalar* blockB, Index rows, Index depth, + Index cols, ResScalar alpha, Index strideB, + Index offsetB) { if (strideB == -1) strideB = depth; - if (rows <= 0 || cols <= 0 || depth <= 0) return; + if (rows <= 0 || cols <= 0) return; Scalar* C_base = const_cast<Scalar*>(&res(0, 0)); const Index C_stride_row = &res(1, 0) - &res(0, 0); const Index C_stride_col = &res(0, 1) - &res(0, 0); + if (depth <= 0) { + if (overwrite) gemm_zero_result(C_base, rows, cols, C_stride_row, C_stride_col); + return; + } sme_fpsr_guard fpsr; + const bool zero_first = overwrite && NumTraits<Scalar>::IsComplex; + if (zero_first) sme_zero_c(C_base, C_stride_row, C_stride_col, rows, cols); sme_gebp_impl_direct_lhs<Scalar, Index>(C_base, C_stride_row, C_stride_col, lhs, lhsStride, blockB, rows, depth, - cols, alpha, strideB, offsetB); + cols, alpha, strideB, offsetB, overwrite && !zero_first); } };
diff --git a/Eigen/src/Core/products/GeneralBlockPanelKernel.h b/Eigen/src/Core/products/GeneralBlockPanelKernel.h index 3bebe2a..4dc123c 100644 --- a/Eigen/src/Core/products/GeneralBlockPanelKernel.h +++ b/Eigen/src/Core/products/GeneralBlockPanelKernel.h
@@ -132,6 +132,41 @@ * * \sa setCpuCacheSizes */ +// GCC clones a function for constant arguments (IPA-CP, IPA-SRA), which would compile a runtime overwrite flag's +// body twice again; clang 18 and Apple clang 21 did not in our builds. +#if EIGEN_COMP_GNUC_STRICT +#define EIGEN_GEMM_NO_CLONE __attribute__((noclone)) +#else +#define EIGEN_GEMM_NO_CLONE +#endif + +// Tag for the gebp_kernel call that assigns C = alpha * A' * B' over the block instead of accumulating into it, +// independently of C's prior contents. A kernel offers it by declaring `static constexpr bool CanOverwrite = true` +// next to that overload; every other kernel, including user specializations, keeps the accumulate-only contract and +// the GEMM driver zeroes the result before calling it. A specialization inheriting a true CanOverwrite while hiding +// the tagged overload must declare it false. +struct gebp_overwrite_t {}; +template <typename Gebp, typename = void> +struct gebp_kernel_overwrites : std::false_type {}; +template <typename Gebp> +struct gebp_kernel_overwrites<Gebp, std::enable_if_t<Gebp::CanOverwrite>> : std::true_type {}; + +// Whether an overwriting kernel should overwrite a rows x cols x depth product rather than have the driver zero the +// result first: Gebp::overwrite_pays where the kernel declares it, true otherwise. +template <typename Gebp, typename Index> +EIGEN_ALWAYS_INLINE bool gebp_overwrite_pays_impl(int, Index rows, Index cols, Index depth, + decltype(Gebp::overwrite_pays(rows, cols, depth))* = nullptr) { + return Gebp::overwrite_pays(rows, cols, depth); +} +template <typename Gebp, typename Index> +EIGEN_ALWAYS_INLINE bool gebp_overwrite_pays_impl(long, Index, Index, Index) { + return true; +} +template <typename Gebp, typename Index> +EIGEN_ALWAYS_INLINE bool gebp_overwrite_pays(Index rows, Index cols, Index depth) { + return gebp_overwrite_pays_impl<Gebp>(0, rows, cols, depth); +} + #ifdef EIGEN_VECTORIZE_SME // True for the scalar pairs the SME gebp_kernel specializes (see // arch/SME/GeneralBlockPanelKernel.h, which static_asserts that it agrees with @@ -1107,6 +1142,13 @@ * |real |cplx | no vectorization yet, would require to pack A with duplication * |cplx |real | easy vectorization */ +// The result scalars whose kernel stores may overwrite: assigning one is known to give what adding it onto a zero +// does. Other scalars, including custom ones over a built-in real type, keep accumulating. +template <typename Scalar> +struct gebp_overwritable_scalar : bool_constant<std::is_arithmetic<Scalar>::value> {}; +template <typename Real> +struct gebp_overwritable_scalar<std::complex<Real>> : bool_constant<std::is_arithmetic<Real>::value> {}; + template <typename LhsScalar, typename RhsScalar, typename Index, typename DataMapper, int mr, int nr, bool ConjugateLhs, bool ConjugateRhs> struct gebp_kernel { @@ -1153,11 +1195,52 @@ ResPacketSize = Traits::ResPacketSize }; - EIGEN_DONT_INLINE void operator()(const DataMapper& res, const LhsScalar* blockA, const RhsScalar* blockB, Index rows, - Index depth, Index cols, ResScalar alpha, Index strideA = -1, Index strideB = -1, - Index offsetA = 0, Index offsetB = 0) const; + // Overwriting starts each result packet from zero instead of loading it, so it computes exactly what accumulating + // into a zeroed result does (see gebp_overwritable_scalar for the scalars that holds for). + static constexpr bool CanOverwrite = gebp_overwritable_scalar<ResScalar>::value; + + EIGEN_STRONG_INLINE void operator()(const DataMapper& res, const LhsScalar* blockA, const RhsScalar* blockB, + Index rows, Index depth, Index cols, ResScalar alpha, Index strideA = -1, + Index strideB = -1, Index offsetA = 0, Index offsetB = 0) const { + run(false, res, blockA, blockB, rows, depth, cols, alpha, strideA, strideB, offsetA, offsetB); + } + // C = alpha * A' * B' over the block, whatever C held (see gebp_overwrite_t). + EIGEN_STRONG_INLINE void operator()(gebp_overwrite_t, const DataMapper& res, const LhsScalar* blockA, + const RhsScalar* blockB, Index rows, Index depth, Index cols, ResScalar alpha, + Index strideA = -1, Index strideB = -1, Index offsetA = 0, + Index offsetB = 0) const { + run(true, res, blockA, blockB, rows, depth, cols, alpha, strideA, strideB, offsetA, offsetB); + } + + // One body for both: `overwrite` is a runtime flag, so that assignment and accumulation share one + // instantiation. + EIGEN_DONT_INLINE EIGEN_GEMM_NO_CLONE void run(bool overwrite, const DataMapper& res, const LhsScalar* blockA, + const RhsScalar* blockB, Index rows, Index depth, Index cols, + ResScalar alpha, Index strideA, Index strideB, Index offsetA, + Index offsetB) const; }; +// One scalar result, alpha * c: added to the destination, or, overwriting, assigned as the sum onto zero that +// accumulating into a zeroed destination computes. A scalar that cannot overwrite (CanOverwrite) takes an overload +// without the overwriting form, so that accumulation asks of it only the `dst += alpha * c` it always did. +template <typename ResScalar, typename Alpha, typename Acc> +EIGEN_ALWAYS_INLINE void gebp_store_scalar_impl(std::false_type, bool, ResScalar& dst, const Alpha& alpha, + const Acc& c) { + dst += alpha * c; +} +template <typename ResScalar, typename Alpha, typename Acc> +EIGEN_ALWAYS_INLINE void gebp_store_scalar_impl(std::true_type, bool overwrite, ResScalar& dst, const Alpha& alpha, + const Acc& c) { + if (overwrite) + dst = ResScalar(0) + alpha * c; + else + dst += alpha * c; +} +template <bool CanOverwrite, typename ResScalar, typename Alpha, typename Acc> +EIGEN_ALWAYS_INLINE void gebp_store_scalar(bool overwrite, ResScalar& dst, const Alpha& alpha, const Acc& c) { + gebp_store_scalar_impl(bool_constant<CanOverwrite>(), overwrite, dst, alpha, c); +} + template <typename LhsScalar, typename RhsScalar, typename Index, typename DataMapper, int mr, int nr, bool ConjugateLhs, bool ConjugateRhs, int SwappedLhsProgress = @@ -1174,7 +1257,8 @@ EIGEN_STRONG_INLINE void operator()(const DataMapper& res, SwappedTraits& straits, const LhsScalar* blA, const RhsScalar* blB, Index depth, const Index endk, Index i, Index j2, - ResScalar alpha, SAccPacket& C0) const { + ResScalar alpha, SAccPacket& C0, bool overwrite = false) const { + EIGEN_UNUSED_VARIABLE(overwrite); EIGEN_UNUSED_VARIABLE(res); EIGEN_UNUSED_VARIABLE(straits); EIGEN_UNUSED_VARIABLE(blA); @@ -1202,13 +1286,14 @@ EIGEN_STRONG_INLINE void operator()(const DataMapper& res, SwappedTraits& straits, const LhsScalar* blA, const RhsScalar* blB, Index depth, const Index endk, Index i, Index j2, - ResScalar alpha, SAccPacket& C0) const { + ResScalar alpha, SAccPacket& C0, bool overwrite = false) const { using SResPacketQuarter = typename unpacket_traits<typename unpacket_traits<SResPacket>::half>::half; using SLhsPacketQuarter = typename unpacket_traits<typename unpacket_traits<SLhsPacket>::half>::half; using SRhsPacketQuarter = typename unpacket_traits<typename unpacket_traits<SRhsPacket>::half>::half; using SAccPacketQuarter = typename unpacket_traits<typename unpacket_traits<SAccPacket>::half>::half; - SResPacketQuarter R = res.template gatherPacket<SResPacketQuarter>(i, j2); + SResPacketQuarter R = + overwrite ? pset1<SResPacketQuarter>(ResScalar(0)) : res.template gatherPacket<SResPacketQuarter>(i, j2); SResPacketQuarter alphav = pset1<SResPacketQuarter>(alpha); if (depth - endk > 0) { @@ -1444,7 +1529,8 @@ EIGEN_ALWAYS_INLINE void gebp_micro_panel_impl(GEBPTraits& traits, const DataMapper_& res, const LhsScalar_* blockA, const RhsScalar_* blockB, ResScalar_ alpha, Index_ i, Index_ j2, Index_ depth, Index_ strideA, Index_ strideB, Index_ offsetA, - Index_ offsetB, int prefetch_res_offset, Index_ peeled_kc, int pk) { + Index_ offsetB, int prefetch_res_offset, Index_ peeled_kc, int pk, + bool overwrite) { using LhsPacketLocal = typename GEBPTraits::LhsPacket; using RhsPacketLocal = typename GEBPTraits::RhsPacket; using ResPacketLocal = typename GEBPTraits::ResPacket; @@ -1544,7 +1630,8 @@ LinearMapper_ r = res.getLinearMapper(i, j2 + j); EIGEN_GEBP_UNROLL_ACCUMULATORS for (int p = 0; p < MrPackets; ++p) { - alignas(ResPacketLocal) ResPacketLocal R = r.template loadPacket<ResPacketLocal>(p * ResPacketSz); + alignas(ResPacketLocal) ResPacketLocal R = + overwrite ? pset1<ResPacketLocal>(ResScalar_(0)) : r.template loadPacket<ResPacketLocal>(p * ResPacketSz); traits.acc(C[j + p * NrCols], alphav, R); r.storePacket(p * ResPacketSz, R); } @@ -1568,11 +1655,12 @@ #endif template <typename LhsScalar, typename RhsScalar, typename Index, typename DataMapper, int mr, int nr, bool ConjugateLhs, bool ConjugateRhs> -EIGEN_DONT_INLINE void gebp_kernel<LhsScalar, RhsScalar, Index, DataMapper, mr, nr, ConjugateLhs, - ConjugateRhs>::operator()(const DataMapper& res, const LhsScalar* blockA, - const RhsScalar* blockB, Index rows, Index depth, - Index cols, ResScalar alpha, Index strideA, Index strideB, - Index offsetA, Index offsetB) const { +EIGEN_DONT_INLINE EIGEN_GEMM_NO_CLONE void +gebp_kernel<LhsScalar, RhsScalar, Index, DataMapper, mr, nr, ConjugateLhs, ConjugateRhs>::run( + bool overwrite_, const DataMapper& res, const LhsScalar* blockA, const RhsScalar* blockB, Index rows, Index depth, + Index cols, ResScalar alpha, Index strideA, Index strideB, Index offsetA, Index offsetB) const { + // Constant false where the scalar cannot overwrite, so that the overwriting stores fold away. + const bool overwrite = CanOverwrite && overwrite_; Traits traits; SwappedTraits straits; @@ -1607,7 +1695,7 @@ using LTraits = std::remove_reference_t<decltype(local_traits)>; gebp_micro_panel_impl<MrP, NrC, LTraits, LhsScalar, RhsScalar, ResScalar, Index, DataMapper, LinearMapper, LhsPacket>(local_traits, res, blockA, blockB, alpha, i, j2, depth, strideA, strideB, offsetA, - offsetB, prefetch_res_offset, peeled_kc, pk); + offsetB, prefetch_res_offset, peeled_kc, pk, overwrite); }; // Budget (in bytes) for co-residency of the RHS block and a strip of the @@ -1746,18 +1834,18 @@ for (Index j2 = 0; j2 < packet_cols8; j2 += 8) { gebp_micro_panel_impl<1, 8, HalfTraits, LhsScalar, RhsScalar, ResScalar, Index, DataMapper, LinearMapper, LhsPacket>(half_traits, res, blockA, blockB, alpha, i, j2, depth, strideA, strideB, - offsetA, offsetB, prefetch_res_offset, peeled_kc, pk); + offsetA, offsetB, prefetch_res_offset, peeled_kc, pk, overwrite); } } for (Index j2 = packet_cols8; j2 < packet_cols4; j2 += 4) { gebp_micro_panel_impl<1, 4, HalfTraits, LhsScalar, RhsScalar, ResScalar, Index, DataMapper, LinearMapper, LhsPacket>(half_traits, res, blockA, blockB, alpha, i, j2, depth, strideA, strideB, - offsetA, offsetB, prefetch_res_offset, peeled_kc, pk); + offsetA, offsetB, prefetch_res_offset, peeled_kc, pk, overwrite); } for (Index j2 = packet_cols4; j2 < cols; j2++) { gebp_micro_panel_impl<1, 1, HalfTraits, LhsScalar, RhsScalar, ResScalar, Index, DataMapper, LinearMapper, LhsPacket>(half_traits, res, blockA, blockB, alpha, i, j2, depth, strideA, strideB, - offsetA, offsetB, prefetch_res_offset, peeled_kc, pk); + offsetA, offsetB, prefetch_res_offset, peeled_kc, pk, overwrite); } } } @@ -1770,18 +1858,18 @@ for (Index j2 = 0; j2 < packet_cols8; j2 += 8) { gebp_micro_panel_impl<1, 8, QuarterTraits, LhsScalar, RhsScalar, ResScalar, Index, DataMapper, LinearMapper, LhsPacket>(quarter_traits, res, blockA, blockB, alpha, i, j2, depth, strideA, strideB, - offsetA, offsetB, prefetch_res_offset, peeled_kc, pk); + offsetA, offsetB, prefetch_res_offset, peeled_kc, pk, overwrite); } } for (Index j2 = packet_cols8; j2 < packet_cols4; j2 += 4) { gebp_micro_panel_impl<1, 4, QuarterTraits, LhsScalar, RhsScalar, ResScalar, Index, DataMapper, LinearMapper, LhsPacket>(quarter_traits, res, blockA, blockB, alpha, i, j2, depth, strideA, strideB, - offsetA, offsetB, prefetch_res_offset, peeled_kc, pk); + offsetA, offsetB, prefetch_res_offset, peeled_kc, pk, overwrite); } for (Index j2 = packet_cols4; j2 < cols; j2++) { gebp_micro_panel_impl<1, 1, QuarterTraits, LhsScalar, RhsScalar, ResScalar, Index, DataMapper, LinearMapper, LhsPacket>(quarter_traits, res, blockA, blockB, alpha, i, j2, depth, strideA, strideB, - offsetA, offsetB, prefetch_res_offset, peeled_kc, pk); + offsetA, offsetB, prefetch_res_offset, peeled_kc, pk, overwrite); } } } @@ -1828,14 +1916,14 @@ blB += 8; } - res(i, j2 + 0) += alpha * C0; - res(i, j2 + 1) += alpha * C1; - res(i, j2 + 2) += alpha * C2; - res(i, j2 + 3) += alpha * C3; - res(i, j2 + 4) += alpha * C4; - res(i, j2 + 5) += alpha * C5; - res(i, j2 + 6) += alpha * C6; - res(i, j2 + 7) += alpha * C7; + gebp_store_scalar<CanOverwrite>(overwrite, res(i, j2 + 0), alpha, C0); + gebp_store_scalar<CanOverwrite>(overwrite, res(i, j2 + 1), alpha, C1); + gebp_store_scalar<CanOverwrite>(overwrite, res(i, j2 + 2), alpha, C2); + gebp_store_scalar<CanOverwrite>(overwrite, res(i, j2 + 3), alpha, C3); + gebp_store_scalar<CanOverwrite>(overwrite, res(i, j2 + 4), alpha, C4); + gebp_store_scalar<CanOverwrite>(overwrite, res(i, j2 + 5), alpha, C5); + gebp_store_scalar<CanOverwrite>(overwrite, res(i, j2 + 6), alpha, C6); + gebp_store_scalar<CanOverwrite>(overwrite, res(i, j2 + 7), alpha, C7); } } } @@ -1924,7 +2012,8 @@ SAccPacket> SAccPacketHalf; - SResPacketHalf R = res.template gatherPacket<SResPacketHalf>(i, j2); + SResPacketHalf R = + overwrite ? pset1<SResPacketHalf>(ResScalar(0)) : res.template gatherPacket<SResPacketHalf>(i, j2); SResPacketHalf alphav = pset1<SResPacketHalf>(alpha); if (depth - endk > 0) { @@ -1946,9 +2035,9 @@ // template form, so that LhsProgress < 16 paths don't // fail to compile last_row_process_16_packets<LhsScalar, RhsScalar, Index, DataMapper, mr, nr, ConjugateLhs, ConjugateRhs> p; - p(res, straits, blA, blB, depth, endk, i, j2, alpha, C0); + p(res, straits, blA, blB, depth, endk, i, j2, alpha, C0, overwrite); } else { - SResPacket R = res.template gatherPacket<SResPacket>(i, j2); + SResPacket R = overwrite ? pset1<SResPacket>(ResScalar(0)) : res.template gatherPacket<SResPacket>(i, j2); SResPacket alphav = pset1<SResPacket>(alpha); straits.acc(C0, alphav, R); res.scatterPacket(i, j2, R); @@ -1976,10 +2065,10 @@ blB += 4; } - res(i, j2 + 0) += alpha * C0; - res(i, j2 + 1) += alpha * C1; - res(i, j2 + 2) += alpha * C2; - res(i, j2 + 3) += alpha * C3; + gebp_store_scalar<CanOverwrite>(overwrite, res(i, j2 + 0), alpha, C0); + gebp_store_scalar<CanOverwrite>(overwrite, res(i, j2 + 1), alpha, C1); + gebp_store_scalar<CanOverwrite>(overwrite, res(i, j2 + 2), alpha, C2); + gebp_store_scalar<CanOverwrite>(overwrite, res(i, j2 + 3), alpha, C3); } } } @@ -1997,7 +2086,7 @@ RhsScalar B_0 = blB[k]; C0 = cj.pmadd(A0, B_0, C0); } - res(i, j2) += alpha * C0; + gebp_store_scalar<CanOverwrite>(overwrite, res(i, j2), alpha, C0); } } }
diff --git a/Eigen/src/Core/products/GeneralMatrixMatrix.h b/Eigen/src/Core/products/GeneralMatrixMatrix.h index 541c9a8..c6507fc 100644 --- a/Eigen/src/Core/products/GeneralMatrixMatrix.h +++ b/Eigen/src/Core/products/GeneralMatrixMatrix.h
@@ -76,15 +76,43 @@ RhsScalar* blockB() const { return m_blockB; } }; +// Runs one gebp call, as C = alpha * A' * B' when `overwrite` holds and the kernel can (gebp_kernel_overwrites), and +// as C += alpha * A' * B' otherwise. It and the loop policies are forced inline into the driver: GCC 13 otherwise +// clones a policy (IPA-SRA) with the kernel call still in it, and emits kernels that only that call's dead branches +// reach (the AVX-512 kernel's non-unit-increment ones). +template <typename Gebp, typename... Args> +EIGEN_ALWAYS_INLINE void gebp_run(std::true_type, Gebp& gebp, bool overwrite, const Args&... args) { + if (overwrite) + gebp(gebp_overwrite_t(), args...); + else + gebp(args...); +} +template <typename Gebp, typename... Args> +EIGEN_ALWAYS_INLINE void gebp_run(std::false_type, Gebp& gebp, bool, const Args&... args) { + gebp(args...); +} + +// Zeroes the rows x cols column-major result block at `res`. +template <typename Scalar, typename Index> +void gemm_zero_result(Scalar* res, Index rows, Index cols, Index resIncr, Index resStride) { + using Block = Matrix<Scalar, Dynamic, Dynamic, ColMajor>; + if (resIncr == 1 && resStride == rows) + Map<Matrix<Scalar, Dynamic, 1>>(res, rows * cols).setZero(); + else if (resIncr == 1) + Map<Block, 0, OuterStride<>>(res, rows, cols, OuterStride<>(resStride)).setZero(); + else + Map<Block, 0, Stride<Dynamic, Dynamic>>(res, rows, cols, Stride<Dynamic, Dynamic>(resStride, resIncr)).setZero(); +} + // LHS-first loop order: mc -> kc -> nc. This is Eigen's usual sequential // schedule, with a fast path that can pack RHS once for tall skinny blocking. struct gemm_pack_lhs_first_loop_policy { template <typename Index, typename LhsScalar, typename RhsScalar, typename ResScalar, typename LhsMapper, typename RhsMapper, typename ResMapper, typename PackLhs, typename PackRhs, typename Gebp> - static EIGEN_STRONG_INLINE void run(Index rows, Index cols, Index depth, Index kc, Index mc, Index nc, + static EIGEN_ALWAYS_INLINE void run(Index rows, Index cols, Index depth, Index kc, Index mc, Index nc, const LhsMapper& lhs, const RhsMapper& rhs, ResMapper& res, PackLhs& pack_lhs, PackRhs& pack_rhs, Gebp& gebp, LhsScalar* blockA, RhsScalar* blockB, - ResScalar alpha) { + ResScalar alpha, bool overwrite) { const bool pack_rhs_once = mc != rows && kc == depth && nc == cols; // For each horizontal panel of the rhs, and corresponding panel of the lhs... @@ -110,7 +138,8 @@ if ((!pack_rhs_once) || i2 == 0) pack_rhs(blockB, rhs.getSubMapper(k2, j2), actual_kc, actual_nc); // Everything is packed, we can now call the panel * block kernel: - gebp(res.getSubMapper(i2, j2), blockA, blockB, actual_mc, actual_kc, actual_nc, alpha); + gebp_run(bool_constant<gebp_kernel_overwrites<Gebp>::value>(), gebp, overwrite && k2 == 0, + res.getSubMapper(i2, j2), blockA, blockB, actual_mc, actual_kc, actual_nc, alpha); } } } @@ -131,12 +160,25 @@ : bool_constant<!NumTraits<Scalar>::IsComplex> {}; // Runs the block on the in-place LHS when the kernel can take it; the false // overload keeps the call out of every other instantiation. +// The in-place call, tagged to overwrite only for a kernel that can (gebp_run). +template <typename Gebp, typename... Args> +EIGEN_ALWAYS_INLINE void sme_direct_lhs_call(std::true_type, Gebp& gebp, bool overwrite, const Args&... args) { + if (overwrite) + gebp.run_direct_lhs(gebp_overwrite_t(), args...); + else + gebp.run_direct_lhs(args...); +} +template <typename Gebp, typename... Args> +EIGEN_ALWAYS_INLINE void sme_direct_lhs_call(std::false_type, Gebp& gebp, bool, const Args&... args) { + gebp.run_direct_lhs(args...); +} template <typename Gebp, typename ResMapper, typename LhsMapper, typename Scalar, typename ResScalar, typename Index> EIGEN_ALWAYS_INLINE bool sme_run_direct_lhs(std::true_type, Gebp& gebp, const ResMapper& res, const LhsMapper& lhs, Index i2, Index k2, const Scalar* blockB, Index mc, Index kc, Index nc, - ResScalar alpha) { + ResScalar alpha, bool overwrite) { if (!sme_direct_lhs_ok<Scalar>(lhs.stride(), mc, kc, nc)) return false; - gebp.run_direct_lhs(res, &lhs(i2, k2), lhs.stride(), blockB, mc, kc, nc, alpha); + sme_direct_lhs_call(bool_constant<gebp_kernel_overwrites<Gebp>::value>(), gebp, overwrite, res, &lhs(i2, k2), + lhs.stride(), blockB, mc, kc, nc, alpha); return true; } template <typename Scalar, int LhsOrder, int RhsOrder, typename Index> @@ -180,7 +222,7 @@ } template <typename Gebp, typename ResMapper, typename LhsMapper, typename Scalar, typename ResScalar, typename Index> EIGEN_ALWAYS_INLINE bool sme_run_direct_lhs(std::false_type, Gebp&, const ResMapper&, const LhsMapper&, Index, Index, - const Scalar*, Index, Index, Index, ResScalar) { + const Scalar*, Index, Index, Index, ResScalar, bool) { return false; } template <typename Scalar, typename Index> @@ -192,10 +234,12 @@ // 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) { +EIGEN_ALWAYS_INLINE bool sme_run_small_gemm(std::true_type, bool overwrite, 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; + if (overwrite) gemm_zero_result(res, rows, cols, resIncr, resStride); 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, @@ -204,7 +248,7 @@ } 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, +EIGEN_ALWAYS_INLINE bool sme_run_small_gemm(std::false_type, bool, Index, Index, Index, const LhsScalar*, Index, const RhsScalar*, Index, ResScalar*, Index, Index, ResScalar, Blocking&) { return false; } @@ -216,10 +260,10 @@ struct gemm_pack_rhs_first_loop_policy { template <typename Index, typename LhsScalar, typename RhsScalar, typename ResScalar, typename LhsMapper, typename RhsMapper, typename ResMapper, typename PackLhs, typename PackRhs, typename Gebp> - static EIGEN_STRONG_INLINE void run(Index rows, Index cols, Index depth, Index kc, Index mc, Index nc, + static EIGEN_ALWAYS_INLINE void run(Index rows, Index cols, Index depth, Index kc, Index mc, Index nc, const LhsMapper& lhs, const RhsMapper& rhs, ResMapper& res, PackLhs& pack_lhs, PackRhs& pack_rhs, Gebp& gebp, LhsScalar* blockA, RhsScalar* blockB, - ResScalar alpha) { + ResScalar alpha, bool overwrite) { // Mirror of pack_rhs_once: reuse one full LHS panel across column blocks. The in-place read is decided per column // block, so a block that falls back to packing packs unless an earlier one already did. const bool pack_lhs_once = nc != cols && kc == depth && mc == rows; @@ -238,20 +282,63 @@ const Index actual_mc = (std::min)(i2 + mc, rows) - i2; #ifdef EIGEN_VECTORIZE_SME if (sme_run_direct_lhs(bool_constant<sme_direct_lhs_mapper<LhsMapper>::value>(), gebp, - res.getSubMapper(i2, j2), lhs, i2, k2, blockB, actual_mc, actual_kc, actual_nc, alpha)) + res.getSubMapper(i2, j2), lhs, i2, k2, blockB, actual_mc, actual_kc, actual_nc, alpha, + overwrite && k2 == 0)) continue; #endif if (!pack_lhs_once || !lhs_packed) { pack_lhs(blockA, lhs.getSubMapper(i2, k2), actual_kc, actual_mc); lhs_packed = true; } - gebp(res.getSubMapper(i2, j2), blockA, blockB, actual_mc, actual_kc, actual_nc, alpha); + gebp_run(bool_constant<gebp_kernel_overwrites<Gebp>::value>(), gebp, overwrite && k2 == 0, + res.getSubMapper(i2, j2), blockA, blockB, actual_mc, actual_kc, actual_nc, alpha); } } } } }; +// Whether Gemm::run_overwrite writes the result without zeroing it first (Gemm::CanOverwrite), so that assignment +// should take it instead of zeroing the destination and accumulating. +template <typename Gemm, typename = void> +struct gemm_can_overwrite : std::false_type {}; +template <typename Gemm> +struct gemm_can_overwrite<Gemm, std::enable_if_t<Gemm::CanOverwrite>> : std::true_type {}; + +// Overwrites the result through Gemm::run_overwrite, or, for a general_matrix_matrix_product specialization without +// one, by zeroing it and accumulating with Gemm::run. +template <typename Gemm, typename = void> +struct gemm_has_run_overwrite : std::false_type {}; +template <typename Gemm> +struct gemm_has_run_overwrite<Gemm, void_t<decltype(&Gemm::run_overwrite)>> : std::true_type {}; +template <typename Gemm, typename Index, typename LhsScalar, typename RhsScalar, typename ResScalar, typename Blocking> +EIGEN_STRONG_INLINE void gemm_run_overwrite_impl(std::true_type, Index rows, Index cols, Index depth, + const LhsScalar* lhs, Index lhsStride, const RhsScalar* rhs, + Index rhsStride, ResScalar* res, Index resIncr, Index resStride, + ResScalar alpha, Blocking& blocking, GemmParallelInfo<Index>* info) { + Gemm::run_overwrite(rows, cols, depth, lhs, lhsStride, rhs, rhsStride, res, resIncr, resStride, alpha, blocking, + info); +} +template <typename Gemm, typename Index, typename LhsScalar, typename RhsScalar, typename ResScalar, typename Blocking> +EIGEN_STRONG_INLINE void gemm_run_overwrite_impl(std::false_type, Index rows, Index cols, Index depth, + const LhsScalar* lhs, Index lhsStride, const RhsScalar* rhs, + Index rhsStride, ResScalar* res, Index resIncr, Index resStride, + ResScalar alpha, Blocking& blocking, GemmParallelInfo<Index>* info) { + // A cooperative session needs every thread in Gemm::run, so only a sequential call returns on an empty result. + if (info == nullptr && (rows <= 0 || cols <= 0)) return; + gemm_zero_result(res, rows, cols, resIncr, resStride); + if (numext::is_exactly_zero(alpha) || depth <= 0) return; + Gemm::run(rows, cols, depth, lhs, lhsStride, rhs, rhsStride, res, resIncr, resStride, alpha, blocking, info); +} +template <typename Gemm, typename Index, typename LhsScalar, typename RhsScalar, typename ResScalar, typename Blocking> +EIGEN_STRONG_INLINE void gemm_run_overwrite(Index rows, Index cols, Index depth, const LhsScalar* lhs, Index lhsStride, + const RhsScalar* rhs, Index rhsStride, ResScalar* res, Index resIncr, + Index resStride, ResScalar alpha, Blocking& blocking, + GemmParallelInfo<Index>* info) { + gemm_run_overwrite_impl<Gemm>(gemm_has_run_overwrite<Gemm>(), rows, cols, depth, lhs, lhsStride, rhs, rhsStride, res, + resIncr, resStride, alpha, blocking, info); +} + /* Specialization for a row-major destination matrix => simple transposition of the product */ template <typename Index, typename LhsScalar, int LhsStorageOrder, bool ConjugateLhs, typename RhsScalar, int RhsStorageOrder, bool ConjugateRhs, int ResInnerStride> @@ -260,6 +347,9 @@ using Traits = gebp_traits<RhsScalar, LhsScalar>; using ResScalar = typename ScalarBinaryOpTraits<LhsScalar, RhsScalar>::ReturnType; + static constexpr bool CanOverwrite = gemm_can_overwrite<general_matrix_matrix_product< + Index, RhsScalar, RhsStorageOrder == RowMajor ? ColMajor : RowMajor, ConjugateRhs, LhsScalar, + LhsStorageOrder == RowMajor ? ColMajor : RowMajor, ConjugateLhs, ColMajor, ResInnerStride>>::value; static EIGEN_STRONG_INLINE void run(Index rows, Index cols, Index depth, const LhsScalar* lhs, Index lhsStride, const RhsScalar* rhs, Index rhsStride, ResScalar* res, Index resIncr, Index resStride, ResScalar alpha, level3_blocking<RhsScalar, LhsScalar>& blocking, @@ -270,6 +360,17 @@ ResInnerStride>::run(cols, rows, depth, rhs, rhsStride, lhs, lhsStride, res, resIncr, resStride, alpha, blocking, info); } + // As run, but overwrites the result instead of accumulating into it. + static EIGEN_STRONG_INLINE void run_overwrite(Index rows, Index cols, Index depth, const LhsScalar* lhs, + Index lhsStride, const RhsScalar* rhs, Index rhsStride, ResScalar* res, + Index resIncr, Index resStride, ResScalar alpha, + level3_blocking<RhsScalar, LhsScalar>& blocking, + GemmParallelInfo<Index>* info = 0) { + gemm_run_overwrite<general_matrix_matrix_product< + Index, RhsScalar, RhsStorageOrder == RowMajor ? ColMajor : RowMajor, ConjugateRhs, LhsScalar, + LhsStorageOrder == RowMajor ? ColMajor : RowMajor, ConjugateLhs, ColMajor, ResInnerStride>>( + cols, rows, depth, rhs, rhsStride, lhs, lhsStride, res, resIncr, resStride, alpha, blocking, info); + } }; /* Specialization for a col-major destination matrix @@ -280,22 +381,64 @@ ConjugateRhs, ColMajor, ResInnerStride> { using Traits = gebp_traits<LhsScalar, RhsScalar>; + private: + using Gebp = gebp_kernel<LhsScalar, RhsScalar, Index, + blas_data_mapper<typename Traits::ResScalar, Index, ColMajor, Unaligned, ResInnerStride>, + Traits::mr, Traits::nr, ConjugateLhs, ConjugateRhs>; + + public: using ResScalar = typename ScalarBinaryOpTraits<LhsScalar, RhsScalar>::ReturnType; + static constexpr bool CanOverwrite = gebp_kernel_overwrites<Gebp>::value; static void run(Index rows, Index cols, Index depth, const LhsScalar* lhs_, Index lhsStride, const RhsScalar* rhs_, Index rhsStride, ResScalar* res_, Index resIncr, Index resStride, ResScalar alpha, level3_blocking<LhsScalar, RhsScalar>& blocking, GemmParallelInfo<Index>* info = 0) { - // BLAS contract: if alpha == 0, the result is unchanged (and lhs/rhs need not be read). - if (numext::is_exactly_zero(alpha)) return; + run_impl(false, rows, cols, depth, lhs_, lhsStride, rhs_, rhsStride, res_, resIncr, resStride, alpha, blocking, + info); + } + // As run, but assigns the result instead of accumulating into it, independently of what res held before. + static void run_overwrite(Index rows, Index cols, Index depth, const LhsScalar* lhs_, Index lhsStride, + const RhsScalar* rhs_, Index rhsStride, ResScalar* res_, Index resIncr, Index resStride, + ResScalar alpha, level3_blocking<LhsScalar, RhsScalar>& blocking, + GemmParallelInfo<Index>* info = 0) { + run_impl(true, rows, cols, depth, lhs_, lhsStride, rhs_, rhsStride, res_, resIncr, resStride, alpha, blocking, + info); + } + + private: + // One body for both: `overwrite` is a runtime flag, so that assignment and accumulation share one + // instantiation. + static EIGEN_DONT_INLINE EIGEN_GEMM_NO_CLONE void run_impl(bool overwrite, Index rows, Index cols, Index depth, + const LhsScalar* lhs_, Index lhsStride, + const RhsScalar* rhs_, Index rhsStride, ResScalar* res_, + Index resIncr, Index resStride, ResScalar alpha, + level3_blocking<LhsScalar, RhsScalar>& blocking, + GemmParallelInfo<Index>* info) { + // A cooperative session needs every thread to pack and publish its part of A', so only a sequential call returns. + if (info == nullptr && (rows <= 0 || cols <= 0)) return; + // BLAS contract: if alpha == 0, the result is unchanged (and lhs/rhs need not be read); overwriting, it is zero. + // A kernel that only accumulates or finds it does not pay, and an empty depth, overwrite by zeroing first. + if (overwrite && (numext::is_exactly_zero(alpha) || depth == 0 || !gebp_kernel_overwrites<Gebp>::value || + !gebp_overwrite_pays<Gebp>(rows, cols, depth))) { + gemm_zero_result(res_, rows, cols, resIncr, resStride); + overwrite = false; + } + if (numext::is_exactly_zero(alpha) || depth == 0) 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) || - 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; + // The tiny-result and NEON kernels accumulate, so overwriting zeroes their result first; the tiny one runs exactly + // when sme_tiny_gemm_wins holds. + if (info == nullptr) { + constexpr bool tiny_pair = sme_tiny_gemm_pair<LhsScalar, RhsScalar>::value; + if (overwrite && tiny_pair && sme_tiny_gemm_wins<ResScalar>(rows, cols, depth)) + gemm_zero_result(res_, rows, cols, resIncr, resStride); + if (sme_run_tiny_gemm<LhsStorageOrder, RhsStorageOrder>(bool_constant<tiny_pair>(), 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>(), overwrite, rows, cols, depth, lhs_, + lhsStride, rhs_, rhsStride, res_, resIncr, resStride, alpha, blocking)) + return; + } #endif using LhsMapper = const_blas_data_mapper<LhsScalar, Index, LhsStorageOrder>; @@ -313,9 +456,10 @@ LhsStorageOrder> pack_lhs; gemm_pack_rhs<RhsScalar, Index, RhsMapper, Traits::nr, RhsStorageOrder> pack_rhs; - gebp_kernel<LhsScalar, RhsScalar, Index, ResMapper, Traits::mr, Traits::nr, ConjugateLhs, ConjugateRhs> gebp; + Gebp gebp; #if !defined(EIGEN_USE_BLAS) && (defined(EIGEN_HAS_OPENMP) || defined(EIGEN_GEMM_THREADPOOL)) + constexpr bool kOverwrites = gebp_kernel_overwrites<Gebp>::value; if (info) { // this is the parallel version! int tid = info->logical_thread_id; @@ -366,8 +510,9 @@ } } - gebp(res.getSubMapper(info->task_info[i].lhs_start, 0), blockA + info->task_info[i].lhs_start * actual_kc, - blockB, info->task_info[i].lhs_length, actual_kc, nc, alpha); + gebp_run(bool_constant<kOverwrites>(), gebp, overwrite && k == 0, + res.getSubMapper(info->task_info[i].lhs_start, 0), blockA + info->task_info[i].lhs_start * actual_kc, + blockB, info->task_info[i].lhs_length, actual_kc, nc, alpha); } // Then keep going as usual with the remaining B' @@ -378,7 +523,8 @@ pack_rhs(blockB, rhs.getSubMapper(k, j), actual_kc, actual_nc); // C_j += A' * B' - gebp(res.getSubMapper(0, j), blockA, blockB, rows, actual_kc, actual_nc, alpha); + gebp_run(bool_constant<kOverwrites>(), gebp, overwrite && k == 0, res.getSubMapper(0, j), blockA, blockB, + rows, actual_kc, actual_nc, alpha); } // Release all the sub blocks A'_i of A' for the current thread, @@ -410,7 +556,7 @@ #endif SequentialGemmLoop::run(rows, cols, depth, kc, mc, nc, lhs, rhs, res, pack_lhs, pack_rhs, gebp, blockA, blockB, - alpha); + alpha, overwrite); } } }; @@ -420,11 +566,18 @@ * implementation of the high level wrapper to general_matrix_matrix_product **********************************************************************************/ +// overwrite: the product is assigned to dest, independently of its prior contents, instead of added to it. template <typename Scalar, typename Index, typename Gemm, typename Lhs, typename Rhs, typename Dest, typename BlockingType> struct gemm_functor { - gemm_functor(const Lhs& lhs, const Rhs& rhs, Dest& dest, const Scalar& actualAlpha, BlockingType& blocking) - : m_lhs(lhs), m_rhs(rhs), m_dest(dest), m_actualAlpha(actualAlpha), m_blocking(blocking) {} + gemm_functor(const Lhs& lhs, const Rhs& rhs, Dest& dest, const Scalar& actualAlpha, BlockingType& blocking, + bool overwrite = false) + : m_lhs(lhs), + m_rhs(rhs), + m_dest(dest), + m_actualAlpha(actualAlpha), + m_blocking(blocking), + m_overwrite(overwrite) {} void initParallelSession(Index num_threads) const { m_blocking.initParallel(m_lhs.rows(), m_rhs.cols(), m_lhs.cols(), num_threads); @@ -437,9 +590,15 @@ void operator()(Index row, Index rows, Index col = 0, Index cols = -1, GemmParallelInfo<Index>* info = 0) const { if (cols == -1) cols = m_rhs.cols(); - Gemm::run(rows, cols, m_lhs.cols(), &m_lhs.coeffRef(row, 0), m_lhs.outerStride(), &m_rhs.coeffRef(0, col), - m_rhs.outerStride(), (Scalar*)&(m_dest.coeffRef(row, col)), m_dest.innerStride(), m_dest.outerStride(), - m_actualAlpha, m_blocking, info); + if (m_overwrite) { + gemm_run_overwrite<Gemm>(rows, cols, m_lhs.cols(), &m_lhs.coeffRef(row, 0), m_lhs.outerStride(), + &m_rhs.coeffRef(0, col), m_rhs.outerStride(), (Scalar*)&(m_dest.coeffRef(row, col)), + m_dest.innerStride(), m_dest.outerStride(), m_actualAlpha, m_blocking, info); + } else { + Gemm::run(rows, cols, m_lhs.cols(), &m_lhs.coeffRef(row, 0), m_lhs.outerStride(), &m_rhs.coeffRef(0, col), + m_rhs.outerStride(), (Scalar*)&(m_dest.coeffRef(row, col)), m_dest.innerStride(), m_dest.outerStride(), + m_actualAlpha, m_blocking, info); + } } using Traits = typename Gemm::Traits; @@ -450,6 +609,7 @@ Dest& m_dest; Scalar m_actualAlpha; BlockingType& m_blocking; + bool m_overwrite; }; template <int StorageOrder, typename LhsScalar, typename RhsScalar, int MaxRows, int MaxCols, int MaxDepth, @@ -687,10 +847,8 @@ } if (useRuntimeCoeffBasedProduct(dst, rhs)) lazyproduct::eval_dynamic(dst, lhs, rhs, internal::assign_op<typename Dst::Scalar, Scalar>()); - else { - dst.setZero(); - scaleAndAddTo(dst, lhs, rhs, Scalar(1)); - } + else + scaleAndEvalTo(dst, lhs, rhs, Scalar(1)); } template <typename Dst> @@ -713,8 +871,34 @@ template <typename Dest> static void scaleAndAddTo(Dest& dst, const Lhs& a_lhs, const Rhs& a_rhs, const Scalar& alpha) { + scaleAndAddOrEvalTo(dst, a_lhs, a_rhs, alpha, false); + } + + // dst = alpha * a_lhs * a_rhs, independently of dst's prior contents and, where the GEMM kernel can overwrite, + // without zeroing it first. + template <typename Dest> + static void scaleAndEvalTo(Dest& dst, const Lhs& a_lhs, const Rhs& a_rhs, const Scalar& alpha) { + scaleAndAddOrEvalTo(dst, a_lhs, a_rhs, alpha, true); + } + + private: + // One body for both: `overwrite` is a runtime flag, so that assignment and accumulation share one + // instantiation. + template <typename Dest> + static void scaleAndAddOrEvalTo(Dest& dst, const Lhs& a_lhs, const Rhs& a_rhs, const Scalar& alpha, bool overwrite) { EIGEN_INTERNAL_ASSERT_SAME_ROWS(dst, a_lhs); EIGEN_INTERNAL_ASSERT_SAME_COLS(dst, a_rhs); + using Gemm = internal::general_matrix_matrix_product< + Index, LhsScalar, (ActualLhsTypeCleaned::Flags & RowMajorBit) ? RowMajor : ColMajor, + bool(LhsBlasTraits::NeedToConjugate), RhsScalar, + (ActualRhsTypeCleaned::Flags & RowMajorBit) ? RowMajor : ColMajor, bool(RhsBlasTraits::NeedToConjugate), + (Dest::Flags & RowMajorBit) ? RowMajor : ColMajor, Dest::InnerStrideAtCompileTime>; + // Empty operands, the GEMV fallbacks and a GEMM that cannot overwrite accumulate, so overwriting zeroes first. + if (overwrite && (!internal::gemm_can_overwrite<Gemm>::value || a_lhs.cols() == 0 || a_lhs.rows() == 0 || + a_rhs.cols() == 0 || dst.cols() == 1 || dst.rows() == 1)) { + dst.setZero(); + overwrite = false; + } if (a_lhs.cols() == 0 || a_lhs.rows() == 0 || a_rhs.cols() == 0) return; if (dst.cols() == 1) { @@ -738,24 +922,18 @@ internal::gemm_blocking_space<(Dest::Flags & RowMajorBit) ? RowMajor : ColMajor, LhsScalar, RhsScalar, Dest::MaxRowsAtCompileTime, Dest::MaxColsAtCompileTime, MaxDepthAtCompileTime>; - using GemmFunctor = internal::gemm_functor< - Scalar, Index, - internal::general_matrix_matrix_product< - Index, LhsScalar, (ActualLhsTypeCleaned::Flags & RowMajorBit) ? RowMajor : ColMajor, - bool(LhsBlasTraits::NeedToConjugate), RhsScalar, - (ActualRhsTypeCleaned::Flags & RowMajorBit) ? RowMajor : ColMajor, bool(RhsBlasTraits::NeedToConjugate), - (Dest::Flags & RowMajorBit) ? RowMajor : ColMajor, Dest::InnerStrideAtCompileTime>, - ActualLhsTypeCleaned, ActualRhsTypeCleaned, Dest, BlockingType>; + using GemmFunctor = + internal::gemm_functor<Scalar, Index, Gemm, ActualLhsTypeCleaned, ActualRhsTypeCleaned, Dest, BlockingType>; BlockingType blocking(dst.rows(), dst.cols(), lhs.cols(), 1, true); #ifdef EIGEN_VECTORIZE_SME // A result the tiny-result kernel takes is never worth a parallel session. if (tinyKernelWins(dst, a_rhs)) - return internal::parallelize_gemm<false>(GemmFunctor(lhs, rhs, dst, actualAlpha, blocking), a_lhs.rows(), - a_rhs.cols(), a_lhs.cols(), Dest::Flags & RowMajorBit); + return internal::parallelize_gemm<false>(GemmFunctor(lhs, rhs, dst, actualAlpha, blocking, overwrite), + a_lhs.rows(), a_rhs.cols(), a_lhs.cols(), Dest::Flags & RowMajorBit); #endif internal::parallelize_gemm<(Dest::MaxRowsAtCompileTime > 32 || Dest::MaxRowsAtCompileTime == Dynamic)>( - GemmFunctor(lhs, rhs, dst, actualAlpha, blocking), a_lhs.rows(), a_rhs.cols(), a_lhs.cols(), + GemmFunctor(lhs, rhs, dst, actualAlpha, blocking, overwrite), a_lhs.rows(), a_rhs.cols(), a_lhs.cols(), Dest::Flags & RowMajorBit); } };
diff --git a/benchmarks/Core/bench_gemm.cpp b/benchmarks/Core/bench_gemm.cpp index cb9329a..7da8d3f 100644 --- a/benchmarks/Core/bench_gemm.cpp +++ b/benchmarks/Core/bench_gemm.cpp
@@ -21,6 +21,11 @@ c.noalias() += a * b; } +template <typename A, typename B, typename C> +EIGEN_DONT_INLINE void gemm_assign(const A& a, const B& b, C& c) { + c.noalias() = a * b; +} + static void add_gflops_counter(benchmark::State& state) { const double m = state.range(0), n = state.range(1), p = state.range(2); state.counters["GFLOPS"] = benchmark::Counter( @@ -44,8 +49,12 @@ return true; } -// c += a * b with a stored in LhsOrder and b in RhsOrder; c is column-major. -template <int LhsOrder, int RhsOrder = ColMajor> +// How the timed loop forms c = a * b: zeroing c then accumulating (the cost assignment had before it could +// overwrite), assigning, or accumulating onto c as it is. +enum GemmMode { ZeroThenAdd, Assign, Accumulate }; + +// c (op)= a * b with a stored in LhsOrder and b in RhsOrder; c is column-major. +template <int LhsOrder, int RhsOrder = ColMajor, GemmMode Mode = ZeroThenAdd> static void eigen_gemm(benchmark::State& state) { int m = state.range(0); int n = state.range(1); @@ -54,15 +63,23 @@ a.setRandom(); Matrix<Scalar, Dynamic, Dynamic, RhsOrder> b(p, n); b.setRandom(); - Mat c = Mat::Zero(m, n); - gemm(a, b, c); + // Validated through the timed operation, assignment from a nonzero c so that a coefficient it misses shows. + Mat c = Mat::Constant(m, n, Scalar(Mode == Assign ? 7 : 0)); + if (Mode == Assign) + gemm_assign(a, b, c); + else + gemm(a, b, c); if (!gemm_result_ok(a, b, c)) { state.SkipWithError("GEMM result check failed"); return; } for (auto _ : state) { - c.setZero(); - gemm(a, b, c); + if (Mode == Assign) { + gemm_assign(a, b, c); + } else { + if (Mode == ZeroThenAdd) c.setZero(); + gemm(a, b, c); + } benchmark::DoNotOptimize(c.data()); benchmark::ClobberMemory(); } @@ -70,6 +87,8 @@ } static void BM_EigenGemm(benchmark::State& state) { eigen_gemm<ColMajor>(state); } +static void BM_EigenGemmAssign(benchmark::State& state) { eigen_gemm<ColMajor, ColMajor, Assign>(state); } +static void BM_EigenGemmAccumulate(benchmark::State& state) { eigen_gemm<ColMajor, ColMajor, Accumulate>(state); } static void BM_EigenGemmRowMajorLhs(benchmark::State& state) { eigen_gemm<RowMajor>(state); } static void BM_EigenGemmRowMajorRhs(benchmark::State& state) { eigen_gemm<ColMajor, RowMajor>(state); } @@ -122,6 +141,8 @@ // clang-format on BENCHMARK(BM_EigenGemm) GEMM_SIZES DEPTH_SWEEP_SIZES TAIL_PANEL_SIZES SHALLOW_DEPTH_SIZES; +BENCHMARK(BM_EigenGemmAssign) GEMM_SIZES DEPTH_SWEEP_SIZES TAIL_PANEL_SIZES SHALLOW_DEPTH_SIZES; +BENCHMARK(BM_EigenGemmAccumulate) GEMM_SIZES DEPTH_SWEEP_SIZES TAIL_PANEL_SIZES SHALLOW_DEPTH_SIZES; BENCHMARK(BM_EigenGemmRowMajorLhs)->Args({256, 256, 256})->Args({1024, 1024, 1024}) DEPTH_SWEEP_SIZES TAIL_PANEL_SIZES; BENCHMARK(BM_EigenGemmRowMajorRhs) ->Args({64, 64, 1024})
diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt index c859af8..39b12ac 100644 --- a/test/CMakeLists.txt +++ b/test/CMakeLists.txt
@@ -393,6 +393,7 @@ target_link_libraries(gemm_packing "-Wl,--wrap=malloc") endif() ei_add_test(blas_singleton_stride) +ei_add_test(product_assign) ei_add_test(product_threshold) if(EIGEN_TEST_SME) # EIGEN_TEST_SME defines EIGEN_ARM64_USE_SME (root CMakeLists.txt), which turns
diff --git a/test/product_assign.cpp b/test/product_assign.cpp new file mode 100644 index 0000000..d82400b --- /dev/null +++ b/test/product_assign.cpp
@@ -0,0 +1,409 @@ +// 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 + +// GEMM assignment, C = A * B, which kernels that can overwrite (gebp_kernel_overwrites) compute without zeroing C +// first. Every overwritten coefficient is alpha * acc + 0, so the result must be bitwise the one that zeroing C and +// accumulating gives; C starts out NaN, so a coefficient not written, or computed from C's old contents, fails the +// comparison. Small blocking budgets split the products over several depth, row and column blocks, of which only +// the first depth block may overwrite. + +#define EIGEN_DEBUG_SMALL_PRODUCT_BLOCKS +#include "main.h" +#include "product_test_helpers.h" + +#include <limits> + +template <typename Scalar> +Scalar nan_value() { + return Scalar(std::numeric_limits<typename NumTraits<Scalar>::Real>::quiet_NaN()); +} + +// Bitwise equality of every coefficient, so that +0 and -0 differ. +template <typename Real> +bool bits_equal(const Real& x, const Real& y) { + using Bits = typename numext::get_integer_by_size<sizeof(Real)>::unsigned_type; + return numext::bit_cast<Bits>(x) == numext::bit_cast<Bits>(y); +} +template <typename Real> +bool bits_equal(const std::complex<Real>& x, const std::complex<Real>& y) { + return bits_equal(x.real(), y.real()) && bits_equal(x.imag(), y.imag()); +} +template <typename A, typename B> +bool bitwise_equal(const DenseBase<A>& a, const DenseBase<B>& b) { + if (a.rows() != b.rows() || a.cols() != b.cols()) return false; + for (Index j = 0; j < a.cols(); ++j) + for (Index i = 0; i < a.rows(); ++i) + if (!bits_equal(a.derived().coeff(i, j), b.derived().coeff(i, j))) return false; + return true; +} + +// dst = scale * a * b through assignment, from a NaN-filled destination, against the zero-then-accumulate result and +// the coefficient-based product. The two must be bitwise equal where the product runs GEMM. The coefficient-based +// product and the GEMV fallbacks compile `=` and `+=` separately, and those may round differently. +template <typename Dst, typename Lhs, typename Rhs, typename Scalar> +void check_assign(Dst& dst, const Lhs& a, const Rhs& b, const Scalar& scale) { + using ScaledLhs = internal::remove_all_t<decltype(scale * a)>; + using Impl = internal::generic_product_impl<ScaledLhs, Rhs, DenseShape, DenseShape, GemmProduct>; + const bool gemm = dst.rows() > 1 && dst.cols() > 1 && !Impl::useRuntimeCoeffBasedProduct(dst, b); + // Accumulated into dst itself, so that it takes the same kernel (or external BLAS) as the assignment: a row-major or + // strided destination may take another, which rounds differently. + using Plain = typename Dst::PlainObject; + dst.setZero(); + dst.noalias() += scale * a * b; + const Plain expected = dst; + + dst.setConstant(nan_value<typename Dst::Scalar>()); + dst.noalias() = scale * a * b; + VERIFY(dst.allFinite()); + if (gemm) + VERIFY(bitwise_equal(dst, expected)); + else + VERIFY(verifyProduct(dst, expected, scale * a, b)); + + // Assignment without noalias goes through a temporary, whose storage order is the product's. + dst.setConstant(nan_value<typename Dst::Scalar>()); + dst = scale * a * b; + VERIFY(dst.allFinite()); + VERIFY(verifyProduct(dst, expected, scale * a, b)); + + if (a.cols() > 0) VERIFY(verifyProduct(expected, Plain((scale * a).lazyProduct(b)), scale * a, b)); +} + +// Sizes around the packet widths and the register and cache blocks of the kernels, at depths below, at and above +// the small blocking's depth block. +template <typename Scalar, int DstOrder> +void assign_shapes() { + using Mat = Matrix<Scalar, Dynamic, Dynamic, ColMajor>; + using DstMat = Matrix<Scalar, Dynamic, Dynamic, DstOrder>; + const Index sizes[] = {1, 3, 16, 17, 32, 33, 48, 65, 129}; + const Index depths[] = {1, 25, 64, 130, 300}; + for (Index m : sizes) + for (Index n : sizes) + for (Index k : depths) { + if (m * n * k > 1000000) continue; + const Mat a = Mat::Random(m, k), b = Mat::Random(k, n); + DstMat c(m, n); + check_assign(c, a, b, Scalar(1)); + if (m > 1 && n > 1) check_assign(c, a, b, Scalar(2)); + } + + // Transposed operands, and a column-major result read back through an outer stride with sentinel padding. + const Index m = 97, n = 70, k = 300; + const Mat at = Mat::Random(k, m), bt = Mat::Random(n, k); + DstMat c(m, n); + check_assign(c, at.transpose(), bt.transpose(), Scalar(1)); + + const Mat a = Mat::Random(m, k), b = Mat::Random(k, n); + Mat storage = Mat::Constant(m + 5, n, Scalar(7)); + Map<Mat, 0, OuterStride<>> mapped(storage.data(), m, n, OuterStride<>(m + 5)); + check_assign(mapped, a, b, Scalar(1)); + VERIFY((storage.bottomRows(5).array() == Scalar(7)).all()); +} + +// Outer stride 1 and inner stride n: a row-major layout seen through column-major strides, so the kernel's C has a +// unit column stride. +template <typename Scalar> +void assign_unit_col_stride() { + using Mat = Matrix<Scalar, Dynamic, Dynamic, ColMajor>; + for (Index m : {Index(17), Index(65), Index(129)}) + for (Index n : {Index(9), Index(40), Index(129)}) + for (Index k : {Index(25), Index(300)}) { + const Mat a = Mat::Random(m, k), b = Mat::Random(k, n); + Matrix<Scalar, Dynamic, 1> storage = Matrix<Scalar, Dynamic, 1>::Constant(m * n + 3, Scalar(7)); + Map<Mat, 0, Stride<Dynamic, Dynamic>> mapped(storage.data(), m, n, Stride<Dynamic, Dynamic>(1, n)); + check_assign(mapped, a, b, Scalar(1)); + for (Index i = m * n; i < storage.size(); ++i) VERIFY_IS_EQUAL(storage(i), Scalar(7)); + } +} + +// A depth past the blocking's depth block: only the first pass overwrites, the later ones accumulate. +template <typename Scalar> +void assign_depth_passes() { + using Mat = Matrix<Scalar, Dynamic, Dynamic, ColMajor>; + using Blocking = internal::gemm_blocking_space<ColMajor, Scalar, Scalar, Dynamic, Dynamic, Dynamic>; + for (Index m : {Index(33), Index(129)}) { + const Index n = 40; + Index k = 2 * Blocking(m, n, 1 << 20, 1, true).kc() + 7; + VERIFY(Blocking(m, n, k, 1, true).kc() < k); + const Mat a = Mat::Random(m, k), b = Mat::Random(k, n); + Mat c(m, n); + check_assign(c, a, b, Scalar(1)); + } +} + +// An inner stride of two: the result coefficients are every other scalar, and the ones in between stay untouched. +template <typename Scalar> +void assign_inner_stride() { + using Mat = Matrix<Scalar, Dynamic, Dynamic, ColMajor>; + const Index m = 65, n = 40, k = 150; + const Mat a = Mat::Random(m, k), b = Mat::Random(k, n); + Matrix<Scalar, Dynamic, 1> storage = Matrix<Scalar, Dynamic, 1>::Constant(2 * m * n, Scalar(7)); + Map<Mat, 0, Stride<Dynamic, 2>> mapped(storage.data(), m, n, Stride<Dynamic, 2>(2 * m, 2)); + check_assign(mapped, a, b, Scalar(1)); + for (Index i = 1; i < storage.size(); i += 2) VERIFY_IS_EQUAL(storage(i), Scalar(7)); +} + +// Conjugated operands and scalar factors on either side, and a real operand on either side of a complex product. +template <typename Scalar> +void assign_conjugate() { + using Mat = Matrix<Scalar, Dynamic, Dynamic, ColMajor>; + using RealScalar = typename NumTraits<Scalar>::Real; + using RealMat = Matrix<RealScalar, Dynamic, Dynamic, ColMajor>; + const Index m = 70, n = 45, k = 140; + const Mat a = Mat::Random(m, k), b = Mat::Random(k, n), at = Mat::Random(k, m); + Mat c(m, n); + const Scalar s = internal::random<Scalar>(); + check_assign(c, a.conjugate(), b, s); + check_assign(c, at.adjoint(), b, s); + check_assign(c, a, (s * b).conjugate(), Scalar(1)); + const RealMat ra = RealMat::Random(m, k), rb = RealMat::Random(k, n); + check_assign(c, ra, b, RealScalar(1)); + check_assign(c, a, rb, RealScalar(1)); +} + +// Accumulation is unchanged: += and -= onto a nonzero result. +template <typename Scalar> +void accumulate_unchanged() { + using Mat = Matrix<Scalar, Dynamic, Dynamic, ColMajor>; + const Index m = 66, n = 33, k = 200; + const Mat a = Mat::Random(m, k), b = Mat::Random(k, n), c0 = Mat::Random(m, n); + Mat c = c0; + c.noalias() += a * b; + VERIFY(verifyProduct(Mat(c - c0), Mat(a.lazyProduct(b)), a, b)); + c = c0; + c.noalias() -= a * b; + VERIFY(verifyProduct(Mat(c0 - c), Mat(a.lazyProduct(b)), a, b)); +} + +// alpha * acc is -0 for alpha = -1 and an all-zero product, and overwriting must keep the sign that accumulating into +// a zeroed result gives: +0 for real scalars (0 + (-0)); for complex ones it depends on the packet pmadd (the AVX and +// SSE fmaddsub forms give a -0 real part). Odd sizes reach the scalar row and column tails as well as the packet +// stores. +template <typename Scalar> +void assign_signed_zero() { + using Mat = Matrix<Scalar, Dynamic, Dynamic, ColMajor>; + for (Index size : {64, 67}) { + const Mat a = Mat::Ones(size, size), b = Mat::Zero(size, size); + Mat expected = Mat::Zero(size, size); + expected.noalias() += Scalar(-1) * a * b; + Mat c = Mat::Constant(size, size, nan_value<Scalar>()); + c.noalias() = Scalar(-1) * a * b; + VERIFY(bitwise_equal(c, expected)); + for (Index j = 0; j < size; ++j) + for (Index i = 0; i < size; ++i) + VERIFY(c(i, j) == Scalar(0) && (NumTraits<Scalar>::IsComplex || !(std::signbit)(numext::real(c(i, j))))); + } +} + +// The driver's own contract: overwriting assigns, so a zero factor or an empty depth zeroes the result. The result is +// a padded column-major block whatever the default storage order, and the padding must stay untouched. Through +// gemm_run_overwrite, which takes run_overwrite where the driver has one (not the external-BLAS one). +template <typename Scalar> +void driver_overwrite() { + using Mat = Matrix<Scalar, Dynamic, Dynamic, ColMajor>; + using Gemm = + internal::general_matrix_matrix_product<Index, Scalar, ColMajor, false, Scalar, ColMajor, false, ColMajor, 1>; + using Blocking = internal::gemm_blocking_space<ColMajor, Scalar, Scalar, Dynamic, Dynamic, Dynamic>; + const Index m = 40, n = 30, k = 20, ldc = m + 5; + const Mat a = Mat::Random(m, k), b = Mat::Random(k, n); + Mat storage(ldc, n); + internal::GemmParallelInfo<Index>* no_info = nullptr; + for (Index depth : {Index(0), k}) { + for (Scalar alpha : {Scalar(0), Scalar(1)}) { + Blocking blocking(m, n, depth, 1, true); + storage.setConstant(Scalar(7)); + storage.topRows(m).setConstant(nan_value<Scalar>()); + internal::gemm_run_overwrite<Gemm>(m, n, depth, a.data(), m, b.data(), k, storage.data(), Index(1), ldc, alpha, + blocking, no_info); + const Mat expected = (depth == 0 || alpha == Scalar(0)) ? Mat(Mat::Zero(m, n)) : Mat(a.lazyProduct(b)); + VERIFY(storage.topRows(m).allFinite()); + VERIFY(verifyProduct(Mat(storage.topRows(m)), expected, a, b)); + VERIFY((storage.bottomRows(ldc - m).array() == Scalar(7)).all()); + } + } + // An empty result is left alone. + Blocking blocking(m, 0, k, 1, true); + internal::gemm_run_overwrite<Gemm>(m, Index(0), k, a.data(), m, b.data(), k, storage.data(), Index(1), ldc, Scalar(1), + blocking, no_info); +} + +// A kernel from before the overwrite contract: only the untagged call, which accumulates. +struct accumulate_only_kernel { + int* calls; + template <typename... Args> + void operator()(const Args&...) const { + ++*calls; + } +}; + +// A general_matrix_matrix_product specialization from before the contract: no run_overwrite and no CanOverwrite. +struct accumulate_only_gemm { + static int calls; + template <typename Scalar, typename Blocking> + static void run(Index rows, Index cols, Index, const Scalar*, Index, const Scalar*, Index, Scalar* res, Index, + Index resStride, Scalar, Blocking&, internal::GemmParallelInfo<Index>*) { + ++calls; + for (Index j = 0; j < cols; ++j) + for (Index i = 0; i < rows; ++i) res[i + j * resStride] += Scalar(1); + } +}; +int accumulate_only_gemm::calls = 0; + +// A scalar from before the overwrite contract whose product is an expression: it supports `dst += a * c` but not +// adding that expression to a scalar, so accumulating stores must ask for nothing more. +struct et_scalar; +struct et_product { + double v; +}; +struct et_scalar { + double v; + et_scalar& operator+=(const et_product& p) { + v += p.v; + return *this; + } +}; +et_product operator*(const et_scalar& a, const et_scalar& b) { return et_product{a.v * b.v}; } + +// Kernels with and without the profitability hook. +struct pays_never { + static bool overwrite_pays(Index, Index, Index) { return false; } +}; +struct no_pays_hook {}; + +#ifdef EIGEN_VECTORIZE_SME +// A kernel from before the contract that reads its LHS in place: only the untagged run_direct_lhs. +struct accumulate_only_direct_kernel { + int* calls; + template <typename... Args> + void run_direct_lhs(const Args&...) const { + ++*calls; + } +}; +#endif + +// Old interfaces keep working: the kernel only ever gets the untagged call, and a driver without run_overwrite is +// overwritten by zeroing the result and accumulating into it. +void old_interfaces() { + et_scalar dst{1}, alpha{2}, acc{3}; + internal::gebp_store_scalar<false>(true, dst, alpha, acc); + VERIFY_IS_EQUAL(dst.v, 7.0); + + VERIFY((internal::gebp_overwrite_pays<no_pays_hook>(Index(64), Index(64), Index(64)))); + VERIFY(!(internal::gebp_overwrite_pays<pays_never>(Index(64), Index(64), Index(64)))); + +#ifdef EIGEN_VECTORIZE_SME + int direct_calls = 0; + accumulate_only_direct_kernel direct{&direct_calls}; + internal::sme_direct_lhs_call(std::false_type(), direct, true, 1, 2, 3); + VERIFY_IS_EQUAL(direct_calls, 1); +#endif + + int calls = 0; + accumulate_only_kernel kernel{&calls}; + internal::gebp_run(std::false_type(), kernel, true, 1, 2, 3); + VERIFY_IS_EQUAL(calls, 1); + STATIC_CHECK(!(internal::gebp_kernel_overwrites<accumulate_only_kernel>::value)); + + STATIC_CHECK(!(internal::gemm_can_overwrite<accumulate_only_gemm>::value)); + const Index m = 6, n = 4, k = 3, ldc = 8; + Matrix<float, Dynamic, 1> res = Matrix<float, Dynamic, 1>::Constant(ldc * n, nan_value<float>()); + internal::gemm_blocking_space<ColMajor, float, float, Dynamic, Dynamic, Dynamic> blocking(m, n, k, 1, true); + const float* none = nullptr; + internal::GemmParallelInfo<Index>* no_info = nullptr; + accumulate_only_gemm::calls = 0; + internal::gemm_run_overwrite<accumulate_only_gemm>(m, n, k, none, m, none, k, res.data(), Index(1), ldc, 1.f, + blocking, no_info); + VERIFY_IS_EQUAL(accumulate_only_gemm::calls, 1); + for (Index j = 0; j < n; ++j) + for (Index i = 0; i < ldc; ++i) VERIFY(i < m ? res(i + j * ldc) == 1.f : (numext::isnan)(res(i + j * ldc))); +} + +// A complex-like custom scalar over a built-in real type. +struct custom_complex { + double re, im; +}; + +// int and half have no architecture kernel, so they take the generic one: built-in arithmetic opts in, half and +// custom scalars do not (assigning them is not known to match accumulating onto zero). Architecture kernels opt in +// where they implement it. +template <typename Lhs, typename Rhs = Lhs> +using kernel_for = + internal::gebp_kernel<Lhs, Rhs, Index, internal::blas_data_mapper<Lhs, Index, ColMajor>, + internal::gebp_traits<Lhs, Rhs>::mr, internal::gebp_traits<Lhs, Rhs>::nr, false, false>; +void capability_trait() { + STATIC_CHECK((internal::gebp_kernel_overwrites<kernel_for<int>>::value)); + STATIC_CHECK(!(internal::gebp_kernel_overwrites<kernel_for<half>>::value)); + STATIC_CHECK((internal::gebp_overwritable_scalar<std::complex<double>>::value)); + STATIC_CHECK(!(internal::gebp_overwritable_scalar<custom_complex>::value)); +#ifdef EIGEN_USE_BLAS + // With beta = 0, ?gemm may form alpha * t and return -0 where zero-then-accumulate gives +0 (reference BLAS does + // in its transposed branches), so the external-BLAS driver keeps zeroing first. + STATIC_CHECK( + !(internal::gemm_can_overwrite<internal::general_matrix_matrix_product<Index, float, ColMajor, false, float, + ColMajor, false, ColMajor, 1>>::value)); + STATIC_CHECK( + !(internal::gemm_can_overwrite<internal::general_matrix_matrix_product<Index, double, RowMajor, false, double, + ColMajor, false, ColMajor, 1>>::value)); +#endif +#ifdef EIGEN_VECTORIZE_SME + STATIC_CHECK((internal::gebp_kernel_overwrites<kernel_for<float>>::value)); + STATIC_CHECK((internal::gebp_kernel_overwrites<kernel_for<std::complex<float>>>::value)); + // Complex results overwrite up to 128 KB. + VERIFY((kernel_for<std::complex<float>>::overwrite_pays(128, 128, 64))); + VERIFY(!(kernel_for<std::complex<float>>::overwrite_pays(128, 129, 64))); +#ifdef EIGEN_VECTORIZE_SME_F64F64 + // 64-bit results overwrite up to 128 KB; a 32-bit Index must not overflow forming the byte count. + using DoubleKernel = internal::gebp_kernel<double, double, int, internal::blas_data_mapper<double, int, ColMajor>, + internal::gebp_traits<double, double>::mr, + internal::gebp_traits<double, double>::nr, false, false>; + VERIFY(DoubleKernel::overwrite_pays(128, 128, 64)); + VERIFY(!DoubleKernel::overwrite_pays(128, 129, 64)); + VERIFY(!DoubleKernel::overwrite_pays(16384, 16384, 64)); + VERIFY((kernel_for<std::complex<double>>::overwrite_pays(128, 64, 64))); + VERIFY(!(kernel_for<std::complex<double>>::overwrite_pays(128, 65, 64))); +#endif +#endif +} + +EIGEN_DECLARE_TEST(product_assign) { + CALL_SUBTEST_1(capability_trait()); + CALL_SUBTEST_1(old_interfaces()); + CALL_SUBTEST_1((assign_shapes<float, ColMajor>())); + CALL_SUBTEST_1((assign_shapes<float, RowMajor>())); + CALL_SUBTEST_1(assign_inner_stride<float>()); + CALL_SUBTEST_1(assign_unit_col_stride<float>()); + CALL_SUBTEST_1(assign_depth_passes<float>()); + CALL_SUBTEST_1(assign_signed_zero<float>()); + CALL_SUBTEST_1(accumulate_unchanged<float>()); + CALL_SUBTEST_1(driver_overwrite<float>()); + CALL_SUBTEST_2((assign_shapes<double, ColMajor>())); + CALL_SUBTEST_2((assign_shapes<double, RowMajor>())); + CALL_SUBTEST_2(assign_inner_stride<double>()); + CALL_SUBTEST_2(assign_unit_col_stride<double>()); + CALL_SUBTEST_2(assign_depth_passes<double>()); + CALL_SUBTEST_2(assign_signed_zero<double>()); + CALL_SUBTEST_2(driver_overwrite<double>()); + CALL_SUBTEST_3((assign_shapes<std::complex<float>, ColMajor>())); + CALL_SUBTEST_3((assign_shapes<std::complex<float>, RowMajor>())); + CALL_SUBTEST_3(assign_conjugate<std::complex<float>>()); + CALL_SUBTEST_3(assign_inner_stride<std::complex<float>>()); + CALL_SUBTEST_3(assign_unit_col_stride<std::complex<float>>()); + CALL_SUBTEST_3(assign_depth_passes<std::complex<float>>()); + CALL_SUBTEST_3(assign_signed_zero<std::complex<float>>()); + CALL_SUBTEST_3(accumulate_unchanged<std::complex<float>>()); + CALL_SUBTEST_3(driver_overwrite<std::complex<float>>()); + CALL_SUBTEST_4((assign_shapes<std::complex<double>, ColMajor>())); + CALL_SUBTEST_4((assign_shapes<std::complex<double>, RowMajor>())); + CALL_SUBTEST_4(assign_conjugate<std::complex<double>>()); + CALL_SUBTEST_4(assign_inner_stride<std::complex<double>>()); + CALL_SUBTEST_4(assign_unit_col_stride<std::complex<double>>()); + CALL_SUBTEST_4(assign_depth_passes<std::complex<double>>()); + CALL_SUBTEST_4(assign_signed_zero<std::complex<double>>()); + CALL_SUBTEST_4(driver_overwrite<std::complex<double>>()); +}
diff --git a/test/product_threaded.cpp b/test/product_threaded.cpp index 6cb91d2..6ac2afe 100644 --- a/test/product_threaded.cpp +++ b/test/product_threaded.cpp
@@ -316,6 +316,49 @@ VERIFY_IS_APPROX(fc, fr); } +// Bitwise equality of every coefficient, so that +0 and -0 differ. +template <typename Real> +bool bits_equal(const Real& x, const Real& y) { + using Bits = typename numext::get_integer_by_size<sizeof(Real)>::unsigned_type; + return numext::bit_cast<Bits>(x) == numext::bit_cast<Bits>(y); +} +template <typename Real> +bool bits_equal(const std::complex<Real>& x, const std::complex<Real>& y) { + return bits_equal(x.real(), y.real()) && bits_equal(x.imag(), y.imag()); +} + +// Threaded assignment, c = a * b, which the kernels overwrite on the first depth pass instead of zeroing c, against +// zero-then-accumulate in the same pool: the same partition and the same summation order, so bitwise equal. c starts +// out NaN, and the depth spans several depth blocks, so the later passes accumulate. The complex<double> result is at +// most 128 KB, where the SME kernel overwrites rather than having the driver zero it. +template <typename Scalar> +void test_parallelize_gemm_assign() { + static ThreadPool pool(4); + Eigen::setGemmThreadPool(&pool); + using Mat = Matrix<Scalar, Dynamic, Dynamic>; + using Real = typename NumTraits<Scalar>::Real; + const Index m = NumTraits<Scalar>::IsComplex ? 64 : 517, n = NumTraits<Scalar>::IsComplex ? 128 : 283, k = 2049; + internal::gemm_blocking_space<ColMajor, Scalar, Scalar, Dynamic, Dynamic, Dynamic> blocking(m, n, k, 1, true); + VERIFY(blocking.kc() < k); + blocking.initParallel(m, n, k, 4); + VERIFY(blocking.kc() < k); + const Mat a = Mat::Random(m, k), b = Mat::Random(k, n); + Mat expected = Mat::Zero(m, n); + expected.noalias() += a * b; + Mat c = Mat::Constant(m, n, Scalar(std::numeric_limits<Real>::quiet_NaN())); + c.noalias() = a * b; + bool equal = true; + for (Index j = 0; j < n; ++j) + for (Index i = 0; i < m; ++i) equal = equal && bits_equal(c(i, j), expected(i, j)); + VERIFY(equal); + Mat serial(m, n); + { + ScopedSerialGemm guard; + serial.noalias() = a * b; + } + VERIFY_IS_APPROX(c, serial); +} + // Runs a test with the given SME unit count: a nonzero count splits products on the SME kernel into disjoint parts, // 0 keeps the shared parallel session, so both paths are covered on every host. template <typename Test> @@ -344,6 +387,8 @@ CALL_SUBTEST_8(with_sme_units(units, test_parallelize_gemm_parts_types<std::complex<float>>)); CALL_SUBTEST_8(with_sme_units(units, test_parallelize_gemm_parts_types<std::complex<double>>)); CALL_SUBTEST_9(with_sme_units(units, test_parallelize_gemm_parts_blocks<>)); + CALL_SUBTEST_10(with_sme_units(units, test_parallelize_gemm_assign<float>)); + CALL_SUBTEST_10(with_sme_units(units, test_parallelize_gemm_assign<std::complex<double>>)); } CALL_SUBTEST_3(test_balanced_gemm_range<>()); CALL_SUBTEST_5(test_parallelize_gemm_require_initialization<>());