Report actual requested memory on allocation failure; add memory_pool::free()

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'

src/cosma/context.cpp

@@ -66,6 +66,11 @@ memory_pool<Scalar> &cosma_context<Scalar>::get_memory_pool() {
     return memory_pool_;
 }
 
+template <typename Scalar>
+void cosma_context<Scalar>::free_memory_pool() {
+    memory_pool_.free();
+}
+
 template <typename Scalar>
 long long cosma_context<Scalar>::get_cpu_memory_limit() {
     return cpu_memory_limit;
@@ -157,6 +162,11 @@ global_context<Scalar> get_context_instance() {
     return ctxt.get();
 }
 
+template <typename Scalar>
+void free_memory_pool() {
+    get_context_instance<Scalar>()->free_memory_pool();
+}
+
 using zfloat = std::complex<float>;
 using zdouble = std::complex<double>;
 
@@ -198,4 +208,10 @@ template global_context<float> get_context_instance();
 template global_context<double> get_context_instance();
 template global_context<zfloat> get_context_instance();
 template global_context<zdouble> get_context_instance();
+
+// template instantiation for free_memory_pool
+template void free_memory_pool<float>();
+template void free_memory_pool<double>();
+template void free_memory_pool<zfloat>();
+template void free_memory_pool<zdouble>();
 } // namespace cosma

src/cosma/context.hpp

@@ -30,6 +30,13 @@ class cosma_context {
     void register_state(MPI_Comm comm, const Strategy strategy);
 
     memory_pool<Scalar> &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<Scalar> *get_gpu_context();
 #endif
@@ -89,4 +96,13 @@ context<Scalar> make_context(size_t cpu_mem_limit,
 //     for completion of the initialization
 template <typename Scalar>
 global_context<Scalar> 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 <typename Scalar>
+void free_memory_pool();
 } // namespace cosma

src/cosma/memory_pool.cpp

@@ -1,8 +1,29 @@
 #include <cassert>
 #include <complex>
 #include <cosma/memory_pool.hpp>
+#include <iomanip>
 #include <iostream>
 #include <mpi.h>
+#include <sstream>
+
+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 <typename T>
+std::string human_readable_size(size_t n_elements) {
+    double bytes = static_cast<double>(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 <typename T>
 cosma::memory_pool<T>::memory_pool() {}
@@ -75,13 +96,29 @@ void cosma::memory_pool<T>::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<T>(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<T>(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<T>(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<T>::reset() {
     already_pinned = false;
 }
 
+template <typename T>
+void cosma::memory_pool<T>::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 <typename T>
 T* cosma::memory_pool<T>::get_pool_pointer() {
     return pool_.data();
@@ -112,6 +164,11 @@ size_t cosma::memory_pool<T>::size() {
     return pool_size_;
 }
 
+template <typename T>
+size_t cosma::memory_pool<T>::capacity() {
+    return pool_capacity_;
+}
+
 template <typename T>
 void cosma::memory_pool<T>::reserve(std::vector<size_t>& buffer_sizes) {
     auto alignment = aligned_allocator<T>::get_alignment();
@@ -137,13 +194,29 @@ void cosma::memory_pool<T>::reserve(std::vector<size_t>& 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<T>(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<T>(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<T>(pool_capacity_)
+                       << " per rank, potentially a bug. Please inform us of "
+                          "the test-case."
+                       << std::endl;
             throw;
         }
     }

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<size_t>& buffer_sizes);
 
     void pin(T* ptr, std::size_t size);

tests/CMakeLists.txt

@@ -29,10 +29,17 @@ 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()
 
 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)

tests/memory_pool.cpp

@@ -0,0 +1,72 @@
+#include <cosma/memory_pool.hpp>
+
+#include <gtest/gtest.h>
+
+#include <limits>
+#include <vector>
+
+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<double> pool;
+    pool.amortization = 1.0;
+
+    std::vector<size_t> 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<double> pool;
+    pool.amortization = 1.0;
+
+    std::vector<size_t> 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<size_t> 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<double> pool;
+    pool.amortization = 1.0;
+
+    std::vector<size_t> 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<double> 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<size_t> buffer_sizes = {
+        std::numeric_limits<size_t>::max() / 2};
+    EXPECT_ANY_THROW(pool.reserve(buffer_sizes));
+}

tests/memory_pool_mpi.cpp

@@ -0,0 +1,25 @@
+#include "../utils/cosma_utils.hpp"
+
+#include <gtest/gtest.h>
+#include <gtest_mpi/gtest_mpi.hpp>
+
+// 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<double>();
+
+    bool first_ok = test_cosma<double>(strategy, ctx, MPI_COMM_WORLD, 1e-8, 0);
+    ASSERT_TRUE(first_ok);
+
+    ctx->free_memory_pool();
+
+    bool second_ok = test_cosma<double>(strategy, ctx, MPI_COMM_WORLD, 1e-8, 1);
+    ASSERT_TRUE(second_ok);
+}