From 12cfbcc731f64e1862d12f70662b8019b41294a3 Mon Sep 17 00:00:00 2001 From: mohitt31 Date: Tue, 11 Aug 2026 13:02:34 +0530 Subject: [PATCH 1/2] Report actual requested memory on allocation failure; add memory_pool::free() Addresses two of the three asks in #118: the bad_alloc/length_error handlers in memory_pool now report the exact size COSMA tried to allocate (in both human-readable units and element count) instead of a generic message, and memory_pool gains a free() method (plus thin wrappers on cosma_context and a top-level free_memory_pool()) that actually releases the pool's capacity back to the OS, unlike the existing reset() which only marks it logically empty. Adds tests/memory_pool.cpp (non-MPI) and tests/memory_pool_mpi.cpp covering both changes, including a regression test that frees the pool between two independent multiply() calls and checks the result. --- src/cosma/context.cpp | 16 ++++++++ src/cosma/context.hpp | 16 ++++++++ src/cosma/memory_pool.cpp | 85 ++++++++++++++++++++++++++++++++++++--- src/cosma/memory_pool.hpp | 16 ++++++++ tests/CMakeLists.txt | 7 ++++ tests/memory_pool.cpp | 72 +++++++++++++++++++++++++++++++++ tests/memory_pool_mpi.cpp | 25 ++++++++++++ 7 files changed, 231 insertions(+), 6 deletions(-) create mode 100644 tests/memory_pool.cpp create mode 100644 tests/memory_pool_mpi.cpp diff --git a/src/cosma/context.cpp b/src/cosma/context.cpp index 5e6462ca..e5500e4d 100644 --- a/src/cosma/context.cpp +++ b/src/cosma/context.cpp @@ -66,6 +66,11 @@ memory_pool &cosma_context::get_memory_pool() { return memory_pool_; } +template +void cosma_context::free_memory_pool() { + memory_pool_.free(); +} + template long long cosma_context::get_cpu_memory_limit() { return cpu_memory_limit; @@ -157,6 +162,11 @@ global_context get_context_instance() { return ctxt.get(); } +template +void free_memory_pool() { + get_context_instance()->free_memory_pool(); +} + using zfloat = std::complex; using zdouble = std::complex; @@ -198,4 +208,10 @@ template global_context get_context_instance(); template global_context get_context_instance(); template global_context get_context_instance(); template global_context get_context_instance(); + +// template instantiation for free_memory_pool +template void free_memory_pool(); +template void free_memory_pool(); +template void free_memory_pool(); +template void free_memory_pool(); } // namespace cosma diff --git a/src/cosma/context.hpp b/src/cosma/context.hpp index e15c7606..d950c5d9 100644 --- a/src/cosma/context.hpp +++ b/src/cosma/context.hpp @@ -30,6 +30,13 @@ class cosma_context { void register_state(MPI_Comm comm, const Strategy strategy); memory_pool &get_memory_pool(); + + // Releases all memory currently held by COSMA's internal memory + // pool back to the OS. Safe to call between multiply() invocations + // (the next multiplication that needs memory will reallocate it), + // but not while a CosmaMatrix/Buffer from a previous call is still + // alive - see memory_pool::free() for details. + void free_memory_pool(); #ifdef COSMA_HAVE_GPU gpu::mm_handle *get_gpu_context(); #endif @@ -89,4 +96,13 @@ context make_context(size_t cpu_mem_limit, // for completion of the initialization template global_context get_context_instance(); + +// Releases all memory currently held by the global COSMA context's +// memory pool back to the OS. Useful for freeing up memory between +// multiply() calls (e.g. in an application that interleaves COSMA +// with other memory-hungry libraries), without having to destroy +// and recreate the whole context. See cosma_context::free_memory_pool() +// for the lifetime constraint on when this is safe to call. +template +void free_memory_pool(); } // namespace cosma diff --git a/src/cosma/memory_pool.cpp b/src/cosma/memory_pool.cpp index 8fc097c5..fc6bb708 100644 --- a/src/cosma/memory_pool.cpp +++ b/src/cosma/memory_pool.cpp @@ -1,8 +1,29 @@ #include #include #include +#include #include #include +#include + +namespace { +// human-readable representation of a number of elements of type T, +// e.g. "3.42 GB", used to report how much memory COSMA actually +// tried (and failed) to allocate. +template +std::string human_readable_size(size_t n_elements) { + double bytes = static_cast(n_elements) * sizeof(T); + const char *units[] = {"B", "KB", "MB", "GB", "TB"}; + int unit = 0; + while (bytes >= 1024.0 && unit < 4) { + bytes /= 1024.0; + ++unit; + } + std::ostringstream oss; + oss << std::fixed << std::setprecision(2) << bytes << " " << units[unit]; + return oss.str(); +} +} template cosma::memory_pool::memory_pool() {} @@ -75,13 +96,29 @@ void cosma::memory_pool::resize(size_t capacity) { try { pool_.resize(capacity); } catch (const std::bad_alloc& e) { - std::cout << "COSMA (memory pool): not enough space. Try setting the CPU memory limit (see environment variable COSMA_CPU_MAX_MEMORY)." << std::endl; + std::cout << "COSMA (memory pool): failed to allocate " + << human_readable_size(capacity) + << " (" << capacity << " elements) per rank. " + << "Try lowering the CPU memory limit (see environment " + "variable COSMA_CPU_MAX_MEMORY) so that COSMA uses " + "more sequential steps and less memory per rank." + << std::endl; throw; } catch (const std::length_error& e) { - std::cout << "COSMA (memory pool): size >= max_size(). Try setting the CPU memory limit (see environment variable COSMA_CPU_MAX_MEMORY)." << std::endl; + std::cout << "COSMA (memory pool): requested size (" + << human_readable_size(capacity) + << ", " << capacity << " elements per rank) " + << "exceeds the container's max_size(). Try setting the " + "CPU memory limit (see environment variable " + "COSMA_CPU_MAX_MEMORY)." + << std::endl; throw; } catch (const std::exception& e) { - std::cout << "COSMA (memory pool): unknown exception, potentially a bug. Please inform us of the test-case." << std::endl; + std::cout << "COSMA (memory pool): unknown exception while " + "allocating " << human_readable_size(capacity) + << " per rank, potentially a bug. Please inform us of " + "the test-case." + << std::endl; throw; } pool_size_ = capacity; @@ -97,6 +134,21 @@ void cosma::memory_pool::reset() { already_pinned = false; } +template +void cosma::memory_pool::free() { + this->unpin_all(); + // swapping with a freshly-constructed, empty vector is the + // standard way to force the allocated capacity to be released; + // shrink_to_fit() is only a non-binding request and most + // implementations honor it, but it is not guaranteed to. + aligned_vector_t().swap(pool_); + pool_size_ = 0; + pool_capacity_ = 0; + n_buffers_ = 0; + resized = false; + already_pinned = false; +} + template T* cosma::memory_pool::get_pool_pointer() { return pool_.data(); @@ -112,6 +164,11 @@ size_t cosma::memory_pool::size() { return pool_size_; } +template +size_t cosma::memory_pool::capacity() { + return pool_capacity_; +} + template void cosma::memory_pool::reserve(std::vector& buffer_sizes) { auto alignment = aligned_allocator::get_alignment(); @@ -137,13 +194,29 @@ void cosma::memory_pool::reserve(std::vector& buffer_sizes) { try { pool_.reserve(pool_capacity_); } catch (const std::bad_alloc& e) { - std::cout << "COSMA (memory pool): not enough space. Try setting the CPU memory limit (see environment variable COSMA_CPU_MAX_MEMORY)." << std::endl; + std::cout << "COSMA (memory pool): failed to reserve " + << human_readable_size(pool_capacity_) + << " (" << pool_capacity_ << " elements) per rank. " + << "Try lowering the CPU memory limit (see environment " + "variable COSMA_CPU_MAX_MEMORY) so that COSMA uses " + "more sequential steps and less memory per rank." + << std::endl; throw; } catch (const std::length_error& e) { - std::cout << "COSMA (memory pool): size >= max_size(). Try setting the CPU memory limit (see environment variable COSMA_CPU_MAX_MEMORY)." << std::endl; + std::cout << "COSMA (memory pool): requested size (" + << human_readable_size(pool_capacity_) + << ", " << pool_capacity_ << " elements per rank) " + << "exceeds the container's max_size(). Try setting the " + "CPU memory limit (see environment variable " + "COSMA_CPU_MAX_MEMORY)." + << std::endl; throw; } catch (const std::exception& e) { - std::cout << "COSMA (memory pool): unknown exception, potentially a bug. Please inform us of the test-case." << std::endl; + std::cout << "COSMA (memory pool): unknown exception while " + "reserving " << human_readable_size(pool_capacity_) + << " per rank, potentially a bug. Please inform us of " + "the test-case." + << std::endl; throw; } } diff --git a/src/cosma/memory_pool.hpp b/src/cosma/memory_pool.hpp index d73e1581..fdb029e8 100644 --- a/src/cosma/memory_pool.hpp +++ b/src/cosma/memory_pool.hpp @@ -33,11 +33,27 @@ class memory_pool { void resize(size_t capacity); void reset(); + // Releases all memory currently held by the pool back to the OS + // (unlike reset(), which only marks the pool as empty but keeps + // its capacity allocated). Useful when COSMA's buffers are not + // needed for a while and other libraries/parts of the application + // need the memory in the meantime. The next multiplication that + // needs the pool will simply reallocate what it needs. + // + // WARNING: only call this when no CosmaMatrix/Buffer objects from + // a previous multiply() are still alive. Those objects cache raw + // pointers into the pool's backing storage, which free() releases; + // reusing them afterwards is a use-after-free. It is safe to call + // between independent multiply()/multiply_using_layout() calls, + // since those construct fresh matrix objects internally each time. + void free(); + T* get_pool_pointer(); void turn_on_output(); size_t size(); + size_t capacity(); void reserve(std::vector& buffer_sizes); void pin(T* ptr, std::size_t size); diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 6dc8fa6d..de351bcb 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -29,6 +29,12 @@ add_test(NAME test.mapper COMMAND test.mapper) add_dependencies(tests test.mapper) install(TARGETS test.mapper DESTINATION "${CMAKE_INSTALL_BINDIR}") +add_executable(test.memory_pool memory_pool.cpp) +target_link_libraries(test.memory_pool PRIVATE main_gtest gtest cosma) +add_test(NAME test.memory_pool COMMAND test.memory_pool) +add_dependencies(tests test.memory_pool) +install(TARGETS test.memory_pool DESTINATION "${CMAKE_INSTALL_BINDIR}") + if(NOT COSMA_SCALAPACK MATCHES "OFF") add_cosma_mpi_test(pdgemm 16 cosma_pxgemm_cpp) endif() @@ -36,3 +42,4 @@ endif() add_cosma_mpi_test(multiply_using_layout 4 cosma) add_cosma_mpi_test(multiply 16 cosma) add_cosma_mpi_test(scalar_matmul 8 cosma) +add_cosma_mpi_test(memory_pool_mpi 4 cosma) diff --git a/tests/memory_pool.cpp b/tests/memory_pool.cpp new file mode 100644 index 00000000..b133c4b3 --- /dev/null +++ b/tests/memory_pool.cpp @@ -0,0 +1,72 @@ +#include + +#include + +#include +#include + +using namespace cosma; + +// Requesting/releasing memory should not require any MPI ranks or +// communication, so this is tested as a plain (non-MPI) gtest, similar +// to test.mapper. + +TEST(memory_pool, free_releases_capacity) { + memory_pool pool; + pool.amortization = 1.0; + + std::vector buffer_sizes = {1000, 2000, 3000}; + pool.reserve(buffer_sizes); + + EXPECT_GT(pool.capacity(), 0u); + + pool.free(); + + EXPECT_EQ(pool.capacity(), 0u); + EXPECT_EQ(pool.size(), 0u); +} + +TEST(memory_pool, usable_after_free) { + memory_pool pool; + pool.amortization = 1.0; + + std::vector buffer_sizes = {500}; + pool.reserve(buffer_sizes); + pool.free(); + + // reserving/using the pool again after free() should work exactly + // as it would on a freshly-constructed pool. + std::vector buffer_sizes_2 = {500}; + EXPECT_NO_THROW(pool.reserve(buffer_sizes_2)); + EXPECT_NO_THROW(auto id = pool.get_buffer_id(500); pool.get_buffer_pointer(id)); +} + +TEST(memory_pool, reset_keeps_capacity_free_does_not) { + memory_pool pool; + pool.amortization = 1.0; + + std::vector buffer_sizes = {10000}; + pool.reserve(buffer_sizes); + size_t reserved_capacity = pool.capacity(); + ASSERT_GT(reserved_capacity, 0u); + + pool.reset(); + // reset() only marks the pool as logically empty, capacity is unchanged + EXPECT_EQ(pool.capacity(), reserved_capacity); + + pool.free(); + // free() actually gives the memory back + EXPECT_EQ(pool.capacity(), 0u); +} + +TEST(memory_pool, oversized_request_throws) { + memory_pool pool; + pool.amortization = 1.0; + + // a request this large cannot be satisfied on any real machine and + // should raise a length_error/bad_alloc (caught internally, message + // printed, and rethrown), rather than silently corrupting state. + std::vector buffer_sizes = { + std::numeric_limits::max() / 2}; + EXPECT_ANY_THROW(pool.reserve(buffer_sizes)); +} diff --git a/tests/memory_pool_mpi.cpp b/tests/memory_pool_mpi.cpp new file mode 100644 index 00000000..3f28349d --- /dev/null +++ b/tests/memory_pool_mpi.cpp @@ -0,0 +1,25 @@ +#include "../utils/cosma_utils.hpp" + +#include +#include + +// Regression test for cosma_context::free_memory_pool(): releasing the +// pool between two independent multiply() calls (each of which builds +// its own fresh CosmaMatrix/Buffer objects, as test_cosma() does) must +// not corrupt the pool nor the results of the following multiplication. +TEST(MemoryPool, FreeBetweenMultipliesIsSafe) { + int rank, P; + MPI_Comm_rank(MPI_COMM_WORLD, &rank); + MPI_Comm_size(MPI_COMM_WORLD, &P); + + Strategy strategy(200, 200, 200, P); + auto ctx = cosma::make_context(); + + bool first_ok = test_cosma(strategy, ctx, MPI_COMM_WORLD, 1e-8, 0); + ASSERT_TRUE(first_ok); + + ctx->free_memory_pool(); + + bool second_ok = test_cosma(strategy, ctx, MPI_COMM_WORLD, 1e-8, 1); + ASSERT_TRUE(second_ok); +} From 280c193678db3c8af61ed9d3a6e214ea7df2bc62 Mon Sep 17 00:00:00 2001 From: Simon Pintarelli Date: Thu, 20 Aug 2026 11:07:14 +0200 Subject: [PATCH 2/2] run new tests in CI --- ci/cscs.yml | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/ci/cscs.yml b/ci/cscs.yml index 39e3a1d2..0e263583 100644 --- a/ci/cscs.yml +++ b/ci/cscs.yml @@ -91,3 +91,20 @@ multiply_using_layout: variables: SLURM_JOB_NUM_NODES: 1 SLURM_NTASKS: 4 + +memory_pool: + extends: .run_tests + stage: test + script: /cosma-env-cuda/.spack-env/view/bin/test.memory_pool + variables: + SLURM_JOB_NUM_NODES: 1 + SLURM_NTASKS: 1 + +memory_pool_mpi: + extends: .run_tests + stage: test + script: /cosma-env-cuda/.spack-env/view/bin/test.memory_pool_mpi + variables: + SLURM_JOB_NUM_NODES: 1 + SLURM_NTASKS: 4 + USE_MPI: 'YES'