SME: Store tail-panel lines whole libeigen/eigen!3363
diff --git a/Eigen/src/Core/arch/SME/GeneralBlockPanelKernel.h b/Eigen/src/Core/arch/SME/GeneralBlockPanelKernel.h index 3a52e3a..df30f1c 100644 --- a/Eigen/src/Core/arch/SME/GeneralBlockPanelKernel.h +++ b/Eigen/src/Core/arch/SME/GeneralBlockPanelKernel.h
@@ -116,12 +116,13 @@ // packed panel of width `width`, for the depth sub-range [k0, k1). Both dst and // src are indexed by the absolute depth index k (dst[k*width+off], // src[k*src_stride+off]); the caller offsets `src` to the region's column base -// and `dst` to the panel base. Generalised over the runtime svl: the panel is -// covered in svl-wide column chunks, each streamed over the depth sub-range. -// The chunk loop is outermost so each chunk's predicate is computed once instead -// of per depth step (the runtime chunk count keeps the compiler from hoisting it -// on its own). The symm packers reuse this for the diagonal-split direct/ -// transposed regions (a contiguous depth sub-range at a depth offset). +// and `dst` to the panel base. A panel with svl < width <= 2*svl stores each packed +// line with one two-vector store; others store svl-wide predicated chunks, each +// streamed over the depth sub-range. The chunk loop is outermost so each chunk's +// predicate is computed once instead of per depth step (the runtime chunk count +// keeps the compiler from hoisting it on its own). The symm packers reuse this +// for the diagonal-split direct/transposed regions (a contiguous depth sub-range +// at a depth offset). template <bool Conjugate, typename Scalar, typename Index> static EIGEN_ALWAYS_INLINE void sve_copy_panel_range(Scalar* EIGEN_RESTRICT dst, const Scalar* EIGEN_RESTRICT src, Index src_stride, Index k0, Index k1, @@ -145,6 +146,23 @@ for (; k < k1; ++k) pstoreu_x2(pn, &dst[k * width], ploadu_x2(pn, &src[k * src_stride])); return; } + if (width > svl && width < 2 * svl) { + // Tail panel: each packed line in one counter-predicated two-vector store, four steps in flight. + const svcount_t pn = Traits::whilelt_c2(0, width); + Index k = k0; + for (; k + 4 <= k1; k += 4) { + const auto v0 = ploadu_x2(pn, &src[k * src_stride]); + const auto v1 = ploadu_x2(pn, &src[(k + 1) * src_stride]); + const auto v2 = ploadu_x2(pn, &src[(k + 2) * src_stride]); + const auto v3 = ploadu_x2(pn, &src[(k + 3) * src_stride]); + pstoreu_x2(pn, &dst[k * width], v0); + pstoreu_x2(pn, &dst[(k + 1) * width], v1); + pstoreu_x2(pn, &dst[(k + 2) * width], v2); + pstoreu_x2(pn, &dst[(k + 3) * width], v3); + } + for (; k < k1; ++k) pstoreu_x2(pn, &dst[k * width], ploadu_x2(pn, &src[k * src_stride])); + return; + } for (int off = 0; off < width; off += svl) { const svbool_t pred = sme_packet_traits<Scalar>::whilelt(off, width); for (Index k = k0; k < k1; ++k) { @@ -383,7 +401,7 @@ const svcount_t pn = Traits::ptrue_c(); const Vec zero = pset1<Vec>(RealScalar(0)); const int rows_lo = sme_min(width, svl), rows_hi = width - rows_lo; - const svbool_t pg_lo = Traits::whilelt(0, rows_lo), pg_hi = Traits::whilelt(0, rows_hi); + const svbool_t pg_lo = Traits::whilelt(0, rows_lo); const svbool_t pg_two = Traits::whilelt(0, 2 * width); Index k = k0; for (; k + w2 <= k1; k += w2) { @@ -419,6 +437,27 @@ } continue; } + if (rows_hi > 0) { + // Each line in one counter-predicated two-vector store. + const svcount_t pc = Traits::whilelt_c2(0, width); + for (int c = 0; c < svl; c += 4) { + const auto t0 = sme_read_ver_za_vg4<0, RealScalar>(uint32_t(c)), + t1 = sme_read_ver_za_vg4<1, RealScalar>(uint32_t(c)); + const auto t2 = sme_read_ver_za_vg4<2, RealScalar>(uint32_t(c)), + t3 = sme_read_ver_za_vg4<3, RealScalar>(uint32_t(c)); + RealScalar* d0 = dst + (k + c) * width; + RealScalar* d1 = d0 + Index(svl) * width; + pstoreu_x2(pc, d0, pcreate(pget<0>(t0), pget<0>(t2))); + pstoreu_x2(pc, d0 + width, pcreate(pget<1>(t0), pget<1>(t2))); + pstoreu_x2(pc, d0 + 2 * width, pcreate(pget<2>(t0), pget<2>(t2))); + pstoreu_x2(pc, d0 + 3 * width, pcreate(pget<3>(t0), pget<3>(t2))); + pstoreu_x2(pc, d1, pcreate(pget<0>(t1), pget<0>(t3))); + pstoreu_x2(pc, d1 + width, pcreate(pget<1>(t1), pget<1>(t3))); + pstoreu_x2(pc, d1 + 2 * width, pcreate(pget<2>(t1), pget<2>(t3))); + pstoreu_x2(pc, d1 + 3 * width, pcreate(pget<3>(t1), pget<3>(t3))); + } + continue; + } for (int c = 0; c < svl; c += 4) { const auto t0 = sme_read_ver_za_vg4<0, RealScalar>(uint32_t(c)), t1 = sme_read_ver_za_vg4<1, RealScalar>(uint32_t(c)); @@ -432,18 +471,6 @@ pstoreu(pg_lo, d1 + width, pget<1>(t1)); pstoreu(pg_lo, d1 + 2 * width, pget<2>(t1)); pstoreu(pg_lo, d1 + 3 * width, pget<3>(t1)); - if (rows_hi > 0) { - const auto t2 = sme_read_ver_za_vg4<2, RealScalar>(uint32_t(c)), - t3 = sme_read_ver_za_vg4<3, RealScalar>(uint32_t(c)); - pstoreu(pg_hi, d0 + svl, pget<0>(t2)); - pstoreu(pg_hi, d0 + width + svl, pget<1>(t2)); - pstoreu(pg_hi, d0 + 2 * width + svl, pget<2>(t2)); - pstoreu(pg_hi, d0 + 3 * width + svl, pget<3>(t2)); - pstoreu(pg_hi, d1 + svl, pget<0>(t3)); - pstoreu(pg_hi, d1 + width + svl, pget<1>(t3)); - pstoreu(pg_hi, d1 + 2 * width + svl, pget<2>(t3)); - pstoreu(pg_hi, d1 + 3 * width + svl, pget<3>(t3)); - } } } return k;
diff --git a/Eigen/src/Core/arch/SME/PacketMath.h b/Eigen/src/Core/arch/SME/PacketMath.h index 656cda4..a4a3cb5 100644 --- a/Eigen/src/Core/arch/SME/PacketMath.h +++ b/Eigen/src/Core/arch/SME/PacketMath.h
@@ -91,6 +91,14 @@ static EIGEN_ALWAYS_INLINE svbool_t whilelt(int64_t begin, int64_t end) __arm_streaming EIGEN_SME_ZA_AGNOSTIC { return svwhilelt_b32(begin, end); } + static EIGEN_ALWAYS_INLINE svcount_t whilelt_c2(int64_t begin, int64_t end) __arm_streaming EIGEN_SME_ZA_AGNOSTIC { +#if EIGEN_COMP_GNUC_STRICT + // GCC 14 and 15 crash folding a counter whilelt whose bounds are small constants, as a full panel's width + // inlined into a copy's tail branch is; keep `end` opaque to it. + EIGEN_OPTIMIZATION_BARRIER(end) +#endif + return svwhilelt_c32_s64(begin, end, 2); + } static EIGEN_ALWAYS_INLINE svcount_t whilelt_c4(int64_t begin, int64_t end) __arm_streaming EIGEN_SME_ZA_AGNOSTIC { return svwhilelt_c32_s64(begin, end, 4); } @@ -110,6 +118,13 @@ static EIGEN_ALWAYS_INLINE svbool_t whilelt(int64_t begin, int64_t end) __arm_streaming EIGEN_SME_ZA_AGNOSTIC { return svwhilelt_b64(begin, end); } + static EIGEN_ALWAYS_INLINE svcount_t whilelt_c2(int64_t begin, int64_t end) __arm_streaming EIGEN_SME_ZA_AGNOSTIC { +#if EIGEN_COMP_GNUC_STRICT + // See the float specialization. + EIGEN_OPTIMIZATION_BARRIER(end) +#endif + return svwhilelt_c64_s64(begin, end, 2); + } static EIGEN_ALWAYS_INLINE svcount_t whilelt_c4(int64_t begin, int64_t end) __arm_streaming EIGEN_SME_ZA_AGNOSTIC { return svwhilelt_c64_s64(begin, end, 4); }
diff --git a/benchmarks/Core/bench_gemm.cpp b/benchmarks/Core/bench_gemm.cpp index 7190356..50c479a 100644 --- a/benchmarks/Core/bench_gemm.cpp +++ b/benchmarks/Core/bench_gemm.cpp
@@ -28,8 +28,8 @@ } // Sampled coefficients of c = a * b: |c(i,j) - sum_q a(i,q) b(q,j)| <= 8 p eps sum_q |a(i,q)| |b(q,j)|, p the depth. -template <typename Lhs> -static bool gemm_result_ok(const Lhs& a, const Mat& b, const Mat& c) { +template <typename Lhs, typename Rhs> +static bool gemm_result_ok(const Lhs& a, const Rhs& b, const Mat& c) { for (Index i : {Index(0), a.rows() / 2, a.rows() - 1}) { for (Index j : {Index(0), b.cols() / 2, b.cols() - 1}) { Scalar expected(0); @@ -44,15 +44,15 @@ return true; } -// c += a * b with a stored in LhsOrder; b and c are column-major. -template <int LhsOrder> +// c += a * b with a stored in LhsOrder and b in RhsOrder; c is column-major. +template <int LhsOrder, int RhsOrder = ColMajor> static void eigen_gemm(benchmark::State& state) { int m = state.range(0); int n = state.range(1); int p = state.range(2); Matrix<Scalar, Dynamic, Dynamic, LhsOrder> a(m, p); a.setRandom(); - Mat b(p, n); + Matrix<Scalar, Dynamic, Dynamic, RhsOrder> b(p, n); b.setRandom(); Mat c = Mat::Zero(m, n); gemm(a, b, c); @@ -71,6 +71,7 @@ static void BM_EigenGemm(benchmark::State& state) { eigen_gemm<ColMajor>(state); } static void BM_EigenGemmRowMajorLhs(benchmark::State& state) { eigen_gemm<RowMajor>(state); } +static void BM_EigenGemmRowMajorRhs(benchmark::State& state) { eigen_gemm<ColMajor, RowMajor>(state); } // clang-format off #define GEMM_SIZES \ @@ -97,10 +98,20 @@ ->ArgsProduct({{64}, {64}, {41, 81, 120, 136, 256, 512, 2048, 4096}}) \ ->ArgsProduct({{2048}, {64}, {64, 128, 256, 512, 1024}}) \ ->Args({128, 128, 2048})->Args({256, 256, 2048})->Args({4096, 64, 512})->Args({2048, 128, 512}) + +// Tail panels: at SVL 512, m = 60, 120 and 252 leave a's last panel between one and two vectors wide (28, 24 and 28 +// floats; 12 doubles for 60 and 252), and n = 60 leaves b's. A column-major b and a row-major a pack through the ZA +// transposer, a row-major b through the copy; a column-major a is mostly read in place. +#define TAIL_PANEL_SIZES \ + ->Args({60, 60, 1024})->Args({60, 64, 1024})->Args({64, 60, 1024})->Args({120, 120, 1024})->Args({252, 252, 252}) // clang-format on -BENCHMARK(BM_EigenGemm) GEMM_SIZES DEPTH_SWEEP_SIZES; -BENCHMARK(BM_EigenGemmRowMajorLhs)->Args({256, 256, 256})->Args({1024, 1024, 1024}) DEPTH_SWEEP_SIZES; +BENCHMARK(BM_EigenGemm) GEMM_SIZES DEPTH_SWEEP_SIZES TAIL_PANEL_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}) + ->Args({256, 256, 256}) + ->Args({1024, 1024, 1024}) TAIL_PANEL_SIZES; // The reference below calls sgemm_/dgemm_ directly, so it is real-only. #if defined(HAVE_BLAS) && !defined(EIGEN_BENCH_COMPLEX_SCALAR)
diff --git a/test/product_sme.cpp b/test/product_sme.cpp index b4b8be7..1d50e1e 100644 --- a/test/product_sme.cpp +++ b/test/product_sme.cpp
@@ -1103,7 +1103,8 @@ // all but complex<float>. Below it, eligible real tails reach the paired transposer through tail_pack once deep // enough, and complex<float> has MR == 2*svl at SVL 256. Complex tails take the NEON tail_transpose_pack, so // complex<double> at SVL 256 reaches the four-chunk loop only through test_symm_pack's 2*svl-wide direct regions. - std::vector<int> deep_widths = {MR, MR + 1, NR, 2 * NR + 1, 2 * svl}; + // Real tails between svl and 2*svl wide store whole packed lines, through the copy or the ZA transposer. + std::vector<int> deep_widths = {MR, MR + 1, NR, 2 * NR + 1, 2 * svl, svl + 1, 2 * svl - 1}; std::sort(deep_widths.begin(), deep_widths.end()); deep_widths.erase(std::unique(deep_widths.begin(), deep_widths.end()), deep_widths.end()); for (int d : deep_depths) {