Core: Coalesce heap packing buffers in GEMM and SYRK libeigen/eigen!3392 Co-authored-by: Rasmus Munk Larsen <rmlarsen@gmail.com>
diff --git a/Eigen/src/Core/products/GeneralMatrixMatrix.h b/Eigen/src/Core/products/GeneralMatrixMatrix.h index 62c1e5f..b80041e 100644 --- a/Eigen/src/Core/products/GeneralMatrixMatrix.h +++ b/Eigen/src/Core/products/GeneralMatrixMatrix.h
@@ -21,6 +21,60 @@ template <typename LhsScalar_, typename RhsScalar_> class level3_blocking; +// Coalesce two heap packing buffers so glibc's dynamic thresholds learn the whole workspace size (#3201). +// Leave stack allocations and caller-provided buffers on their existing paths. +template <typename LhsScalar, typename RhsScalar> +class gemm_packing_space { + LhsScalar* m_blockA; + RhsScalar* m_blockB; + aligned_stack_memory_handler<unsigned char> m_storage; + aligned_stack_memory_handler<LhsScalar> m_lhs_lifetime; + aligned_stack_memory_handler<RhsScalar> m_rhs_lifetime; + + static EIGEN_DONT_INLINE unsigned char* allocateCombined(std::size_t sizeA, std::size_t sizeB, LhsScalar*& blockA, + RhsScalar*& blockB) { + check_size_for_overflow<LhsScalar>(sizeA); + check_size_for_overflow<RhsScalar>(sizeB); + // Not numext::maxi: under SYCL_DEVICE_ONLY its cl_ulong (std::size_t) specialization is not constexpr. + constexpr std::size_t alignment = (std::max)(std::size_t{EIGEN_STACK_ALIGN_BYTES}, alignof(RhsScalar)); + const std::size_t max_bytes = (std::numeric_limits<std::ptrdiff_t>::max)(); + const std::size_t bytesA = sizeA * sizeof(LhsScalar); + const std::size_t bytesB = sizeB * sizeof(RhsScalar); + if (bytesA > max_bytes - (alignment - 1)) throw_std_bad_alloc(); + const std::size_t offsetB = first_multiple(bytesA, alignment); + if (bytesB > max_bytes - offsetB) throw_std_bad_alloc(); + unsigned char* storage = static_cast<unsigned char*>(scratch_malloc(offsetB + bytesB)); + blockA = reinterpret_cast<LhsScalar*>(storage); + blockB = reinterpret_cast<RhsScalar*>(storage + offsetB); + return storage; + } + + static unsigned char* allocate(std::size_t sizeA, std::size_t sizeB, LhsScalar*& blockA, RhsScalar*& blockB) { + if (blockA != nullptr || blockB != nullptr || sizeA == 0 || sizeB == 0) return nullptr; +#if defined(EIGEN_ALLOCA) && !defined(EIGEN_NO_ALLOCA) + if (sizeA <= EIGEN_STACK_ALLOCATION_LIMIT / sizeof(LhsScalar) || + sizeB <= EIGEN_STACK_ALLOCATION_LIMIT / sizeof(RhsScalar)) + return nullptr; +#endif + // scratch_malloc guarantees only EIGEN_STACK_ALIGN_BYTES, or malloc's alignment when that is 0. + constexpr std::size_t base_alignment = + EIGEN_STACK_ALIGN_BYTES > 0 ? EIGEN_STACK_ALIGN_BYTES : alignof(std::max_align_t); + if ((std::max)(alignof(LhsScalar), alignof(RhsScalar)) > base_alignment) return nullptr; + return allocateCombined(sizeA, sizeB, blockA, blockB); + } + + public: + gemm_packing_space(std::size_t sizeA, std::size_t sizeB, LhsScalar* blockA, RhsScalar* blockB) + : m_blockA(blockA), + m_blockB(blockB), + m_storage(allocate(sizeA, sizeB, m_blockA, m_blockB), 0, true), + m_lhs_lifetime(blockA == nullptr ? m_blockA : nullptr, sizeA, false), + m_rhs_lifetime(blockB == nullptr ? m_blockB : nullptr, sizeB, false) {} + + LhsScalar* blockA() const { return m_blockA; } + RhsScalar* blockB() const { return m_blockB; } +}; + // 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 { @@ -339,8 +393,9 @@ std::size_t sizeA = kc * mc; std::size_t sizeB = kc * nc; - ei_declare_aligned_stack_constructed_variable(LhsScalar, blockA, sizeA, blocking.blockA()); - ei_declare_aligned_stack_constructed_variable(RhsScalar, blockB, sizeB, blocking.blockB()); + gemm_packing_space<LhsScalar, RhsScalar> packing(sizeA, sizeB, blocking.blockA(), blocking.blockB()); + ei_declare_aligned_stack_constructed_variable(LhsScalar, blockA, sizeA, packing.blockA()); + ei_declare_aligned_stack_constructed_variable(RhsScalar, blockB, sizeB, packing.blockB()); // The SME kernel uses RHS-first order so consecutive gebp calls stream // through adjacent row panels of a ColMajor result. Other kernels --
diff --git a/Eigen/src/Core/products/GeneralMatrixMatrixTriangular.h b/Eigen/src/Core/products/GeneralMatrixMatrixTriangular.h index e19d4fc..019150f 100644 --- a/Eigen/src/Core/products/GeneralMatrixMatrixTriangular.h +++ b/Eigen/src/Core/products/GeneralMatrixMatrixTriangular.h
@@ -95,8 +95,9 @@ std::size_t sizeA = kc * mc; std::size_t sizeB = kc * size; - ei_declare_aligned_stack_constructed_variable(LhsScalar, blockA, sizeA, blocking.blockA()); - ei_declare_aligned_stack_constructed_variable(RhsScalar, blockB, sizeB, blocking.blockB()); + gemm_packing_space<LhsScalar, RhsScalar> packing(sizeA, sizeB, blocking.blockA(), blocking.blockB()); + ei_declare_aligned_stack_constructed_variable(LhsScalar, blockA, sizeA, packing.blockA()); + ei_declare_aligned_stack_constructed_variable(RhsScalar, blockB, sizeB, packing.blockB()); gemm_pack_lhs<LhsScalar, Index, LhsMapper, Traits::mr, Traits::LhsProgress, typename Traits::LhsPacket4Packing, LhsStorageOrder>
diff --git a/benchmarks/Core/CMakeLists.txt b/benchmarks/Core/CMakeLists.txt index b8d0edb..b02faf5 100644 --- a/benchmarks/Core/CMakeLists.txt +++ b/benchmarks/Core/CMakeLists.txt
@@ -95,6 +95,7 @@ eigen_add_benchmark(bench_diagonal_move bench_diagonal_move.cpp) eigen_add_benchmark(bench_triangular_product bench_triangular_product.cpp) eigen_add_benchmark(bench_selfadjoint_product bench_selfadjoint_product.cpp) +eigen_add_benchmark(bench_gemm_packing bench_gemm_packing.cpp) eigen_add_benchmark(bench_symv bench_symv.cpp) eigen_add_benchmark(bench_trmv bench_trmv.cpp) eigen_add_benchmark(bench_syr bench_syr.cpp)
diff --git a/benchmarks/Core/bench_gemm_packing.cpp b/benchmarks/Core/bench_gemm_packing.cpp new file mode 100644 index 0000000..1708459 --- /dev/null +++ b/benchmarks/Core/bench_gemm_packing.cpp
@@ -0,0 +1,72 @@ +// SPDX-FileCopyrightText: The Eigen Authors +// SPDX-License-Identifier: MPL-2.0 + +#include <benchmark/benchmark.h> +#include <Eigen/Core> + +#if defined(__linux__) +#include <sys/resource.h> +#endif + +template <typename Scalar, bool RankUpdate> +static void BM_Packing(benchmark::State& state) { + Eigen::setNbThreads(1); + using MatrixType = Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic>; + using Real = typename Eigen::NumTraits<Scalar>::Real; + const Eigen::Index n = state.range(0); + const MatrixType a = MatrixType::Random(n, n); + const MatrixType b = MatrixType::Random(n, n); + MatrixType c = MatrixType::Zero(n, n); + const auto product = [&]() { + if (RankUpdate) + c.template selfadjointView<Eigen::Lower>().rankUpdate(a); + else + c.noalias() += a * b; + }; + product(); + // Validate sampled coefficients without a reference allocation changing malloc's history. + for (Eigen::Index i : {Eigen::Index(0), n / 2, n - 1}) { + for (Eigen::Index j : {Eigen::Index(0), n / 2, n - 1}) { + long double expected = 0, scale = 0; + if (!RankUpdate || i >= j) { + for (Eigen::Index k = 0; k < n; ++k) { + const long double term = static_cast<long double>(a(i, k)) * (RankUpdate ? a(j, k) : b(k, j)); + expected += term; + scale += Eigen::numext::abs(term); + } + } + const long double bound = 8 * n * Eigen::NumTraits<Real>::epsilon() * scale; + if (!(Eigen::numext::abs(static_cast<long double>(c(i, j)) - expected) <= bound)) { + state.SkipWithError("matrix product validation failed"); + return; + } + } + } + c.setZero(); +#if defined(__linux__) + rusage before, after; + getrusage(RUSAGE_SELF, &before); +#endif + for (auto _ : state) { + product(); + benchmark::DoNotOptimize(c.data()); + benchmark::ClobberMemory(); + } +#if defined(__linux__) + getrusage(RUSAGE_SELF, &after); + state.counters["minor_faults_per_call"] = double(after.ru_minflt - before.ru_minflt) / state.iterations(); +#endif + const double flops = RankUpdate ? double(n) * (n + 1) * n : 2.0 * n * n * n; + state.counters["GFLOPS"] = + benchmark::Counter(flops, benchmark::Counter::kIsIterationInvariantRate, benchmark::Counter::kIs1000); +} + +// clang-format off +#define PACKING_SIZES ->Arg(128)->Arg(144)->Arg(160)->Arg(176)->Arg(184)->Arg(192)->Arg(200)->Arg(240) \ + ->Arg(256)->Arg(384)->Arg(448)->Arg(512)->Arg(768) +BENCHMARK(BM_Packing<float, false>)->Name("GEMM_float") PACKING_SIZES; +BENCHMARK(BM_Packing<double, false>)->Name("GEMM_double") PACKING_SIZES; +BENCHMARK(BM_Packing<float, true>)->Name("SYRK_float") PACKING_SIZES; +BENCHMARK(BM_Packing<double, true>)->Name("SYRK_double") PACKING_SIZES; +#undef PACKING_SIZES +// clang-format on
diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt index e4c5998..862a9f8 100644 --- a/test/CMakeLists.txt +++ b/test/CMakeLists.txt
@@ -397,6 +397,11 @@ ei_add_test(conservative_resize) ei_add_test(product_small) ei_add_test(product_large) +ei_add_test(gemm_packing) +if(CMAKE_SYSTEM_NAME STREQUAL "Linux" AND CMAKE_CXX_COMPILER_ID MATCHES "^(GNU|Clang)$") + target_compile_definitions(gemm_packing PRIVATE EIGEN_TEST_WRAP_MALLOC) + target_link_libraries(gemm_packing "-Wl,--wrap=malloc") +endif() 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/gemm_packing.cpp b/test/gemm_packing.cpp new file mode 100644 index 0000000..21f6a0c --- /dev/null +++ b/test/gemm_packing.cpp
@@ -0,0 +1,186 @@ +// SPDX-FileCopyrightText: The Eigen Authors +// SPDX-License-Identifier: MPL-2.0 + +// Real cache blocking lets these products exceed the stack allocation limit. +#define EIGEN_NO_DEBUG_SMALL_PRODUCT_BLOCKS +#include "main.h" +#include <Eigen/Core> + +#ifdef EIGEN_TEST_WRAP_MALLOC +static volatile bool count_packing_allocations = false; +static volatile Index packing_allocations = 0; +// Preserve the linker's --wrap symbols without reserved C++ identifiers. +extern "C" void* real_malloc(std::size_t) asm("__real_malloc"); +extern "C" void* wrap_malloc(std::size_t) asm("__wrap_malloc"); +extern "C" void* wrap_malloc(std::size_t bytes) { + if (count_packing_allocations) packing_allocations = packing_allocations + 1; + return real_malloc(bytes); +} +#endif + +template <typename LhsScalar, typename RhsScalar> +void packing_storage() { + const std::size_t sizeA = EIGEN_STACK_ALLOCATION_LIMIT / sizeof(LhsScalar) + 3; + const std::size_t sizeB = EIGEN_STACK_ALLOCATION_LIMIT / sizeof(RhsScalar) + 7; + internal::gemm_packing_space<LhsScalar, RhsScalar> packing(sizeA, sizeB, nullptr, nullptr); + VERIFY(packing.blockA() != nullptr); + VERIFY(packing.blockB() != nullptr); + const auto first = reinterpret_cast<std::uintptr_t>(packing.blockA()); + const auto second = reinterpret_cast<std::uintptr_t>(packing.blockB()); + const std::size_t alignment = (std::max)(std::size_t(EIGEN_STACK_ALIGN_BYTES), std::size_t(alignof(RhsScalar))); + VERIFY_IS_EQUAL(second - first, (sizeA * sizeof(LhsScalar) + alignment - 1) / alignment * alignment); + VERIFY_IS_EQUAL(first % alignof(LhsScalar), std::uintptr_t(0)); + VERIFY_IS_EQUAL(second % alignment, std::uintptr_t(0)); + packing.blockA()[0] = LhsScalar(1); + packing.blockA()[sizeA - 1] = LhsScalar(2); + packing.blockB()[0] = RhsScalar(3); + packing.blockB()[sizeB - 1] = RhsScalar(4); + VERIFY_IS_EQUAL(packing.blockA()[0], LhsScalar(1)); + VERIFY_IS_EQUAL(packing.blockA()[sizeA - 1], LhsScalar(2)); + + std::vector<LhsScalar> suppliedA(sizeA); + std::vector<RhsScalar> suppliedB(sizeB); + internal::gemm_packing_space<LhsScalar, RhsScalar> supplied(sizeA, sizeB, suppliedA.data(), suppliedB.data()); + VERIFY(supplied.blockA() == suppliedA.data()); + VERIFY(supplied.blockB() == suppliedB.data()); + internal::gemm_packing_space<LhsScalar, RhsScalar> partial(sizeA, sizeB, suppliedA.data(), nullptr); + VERIFY(partial.blockA() == suppliedA.data()); + VERIFY(partial.blockB() == nullptr); +#if defined(EIGEN_ALLOCA) && !defined(EIGEN_NO_ALLOCA) && EIGEN_STACK_ALLOCATION_LIMIT > 0 + internal::gemm_packing_space<LhsScalar, RhsScalar> small(1, 1, nullptr, nullptr); + VERIFY(small.blockA() == nullptr); + VERIFY(small.blockB() == nullptr); + internal::gemm_packing_space<LhsScalar, RhsScalar> mixed(1, sizeB, nullptr, nullptr); + VERIFY(mixed.blockA() == nullptr); + VERIFY(mixed.blockB() == nullptr); +#endif +} + +// Stricter than the alignment scratch_malloc guarantees, so either panel of this type keeps its own allocation. +constexpr std::size_t kScratchAlignment = EIGEN_STACK_ALIGN_BYTES > 0 ? EIGEN_STACK_ALIGN_BYTES + : alignof(std::max_align_t); +struct alignas(2 * kScratchAlignment) OverAlignedScalar { + double value = 0; +}; + +template <typename LhsScalar, typename RhsScalar> +void packing_alignment_fallback() { + const std::size_t sizeA = EIGEN_STACK_ALLOCATION_LIMIT / sizeof(LhsScalar) + 3; + const std::size_t sizeB = EIGEN_STACK_ALLOCATION_LIMIT / sizeof(RhsScalar) + 7; + internal::gemm_packing_space<LhsScalar, RhsScalar> packing(sizeA, sizeB, nullptr, nullptr); + VERIFY(packing.blockA() == nullptr); + VERIFY(packing.blockB() == nullptr); +} + +struct PackingScalar { + static Index live; + static Index remaining; + int value = 0; + PackingScalar() { +#ifdef EIGEN_EXCEPTIONS + if (remaining == 0) throw std::bad_alloc(); +#endif + --remaining; + ++live; + } + ~PackingScalar() { --live; } +}; +Index PackingScalar::live = 0; +Index PackingScalar::remaining = 0; + +namespace Eigen { +template <> +struct NumTraits<PackingScalar> : GenericNumTraits<PackingScalar> { + static constexpr bool RequireInitialization = true; +}; +} // namespace Eigen + +#ifdef EIGEN_EXCEPTIONS +template <typename LhsScalar, typename RhsScalar> +bool packing_throws_bad_alloc(std::size_t sizeA, std::size_t sizeB) { + try { + internal::gemm_packing_space<LhsScalar, RhsScalar> packing(sizeA, sizeB, nullptr, nullptr); + } catch (const std::bad_alloc&) { + return true; + } + return false; +} +#endif + +void packing_lifetimes() { + const std::size_t size = EIGEN_STACK_ALLOCATION_LIMIT / sizeof(PackingScalar) + 3; + PackingScalar::remaining = Index(2 * size); + { + internal::gemm_packing_space<PackingScalar, PackingScalar> packing(size, size, nullptr, nullptr); + VERIFY_IS_EQUAL(PackingScalar::live, Index(2 * size)); + } + VERIFY_IS_EQUAL(PackingScalar::live, Index(0)); +#ifdef EIGEN_EXCEPTIONS + // Throw while constructing A's first element, A's second, B's first, and B's second; size >= 3 even when + // EIGEN_STACK_ALLOCATION_LIMIT is 0. + for (Index count : {Index(0), Index(1), Index(size), Index(size + 1)}) { + PackingScalar::remaining = count; + VERIFY((packing_throws_bad_alloc<PackingScalar, PackingScalar>(size, size))); + VERIFY_IS_EQUAL(PackingScalar::live, Index(0)); + } + const std::size_t max_bytes = (std::numeric_limits<std::ptrdiff_t>::max)(); + for (std::size_t sizeA : {max_bytes, max_bytes / 2}) + VERIFY((packing_throws_bad_alloc<char, char>(sizeA, max_bytes / 2 + 2))); +#endif +} + +// Runs product(); where malloc is wrapped, returns whether it allocated at most once. +template <typename Product> +bool allocates_at_most_once(Product product) { +#ifdef EIGEN_TEST_WRAP_MALLOC + packing_allocations = 0; + count_packing_allocations = true; +#endif + product(); +#ifdef EIGEN_TEST_WRAP_MALLOC + count_packing_allocations = false; + return packing_allocations <= 1; +#else + return true; +#endif +} + +template <typename Scalar, int Order> +void packing_products() { + Eigen::setNbThreads(1); + using MatrixType = Matrix<Scalar, Dynamic, Dynamic, Order>; + using Real = typename NumTraits<Scalar>::Real; + for (Index n : {Index(145), Index(201)}) { + const MatrixType a = MatrixType::Random(n, n + 3); + const MatrixType b = MatrixType::Random(n + 3, n + 1); + MatrixType result = MatrixType::Zero(n, n + 1); + MatrixType reference = a.lazyProduct(b); + VERIFY(allocates_at_most_once([&] { result.noalias() += a * b; })); + const Real gemm_bound = Real(8) * Real(n + 3) * NumTraits<Real>::epsilon() * a.norm() * b.norm(); + VERIFY((result - reference).norm() <= gemm_bound); + + result.setZero(n, n); + reference = a.lazyProduct(a.adjoint()); + const Real syrk_bound = Real(8) * Real(n + 3) * NumTraits<Real>::epsilon() * a.squaredNorm(); + VERIFY(allocates_at_most_once([&] { result.template selfadjointView<Lower>().rankUpdate(a); })); + reference.template triangularView<StrictlyUpper>().setZero(); + VERIFY((result - reference).norm() <= syrk_bound); + result.setZero(); + result.template selfadjointView<Upper>().rankUpdate(a.conjugate()); + reference = a.conjugate().lazyProduct(a.transpose()); + reference.template triangularView<StrictlyLower>().setZero(); + VERIFY((result - reference).norm() <= syrk_bound); + } +} + +EIGEN_DECLARE_TEST(gemm_packing) { + CALL_SUBTEST((packing_storage<float, double>())); + CALL_SUBTEST((packing_storage<double, std::complex<double>>())); + CALL_SUBTEST((packing_alignment_fallback<OverAlignedScalar, float>())); + CALL_SUBTEST((packing_alignment_fallback<float, OverAlignedScalar>())); + CALL_SUBTEST(packing_lifetimes()); + CALL_SUBTEST((packing_products<float, ColMajor>())); + CALL_SUBTEST((packing_products<double, RowMajor>())); + CALL_SUBTEST((packing_products<std::complex<double>, ColMajor>())); + CALL_SUBTEST((packing_products<std::complex<double>, RowMajor>())); +}