From 6a53c5247499f245848e8ba38aa6291860ae8b8f Mon Sep 17 00:00:00 2001 From: Antoine Pitrou Date: Wed, 10 Apr 2024 17:27:37 +0200 Subject: [PATCH] EXPERIMENT: [C++] Access mimalloc through dynamically-resolved symbols --- cpp/src/arrow/CMakeLists.txt | 8 ++++++++ cpp/src/arrow/memory_pool.cc | 19 +++++++------------ cpp/src/arrow/memory_pool_benchmark.cc | 1 + cpp/src/arrow/memory_pool_internal.h | 21 +++++++++++++-------- cpp/src/arrow/memory_pool_test.cc | 5 ++++- cpp/src/arrow/memory_pool_test.h | 8 ++++++++ 6 files changed, 41 insertions(+), 21 deletions(-) diff --git a/cpp/src/arrow/CMakeLists.txt b/cpp/src/arrow/CMakeLists.txt index 026bb5c77e066..9e1c5ba3b7f2b 100644 --- a/cpp/src/arrow/CMakeLists.txt +++ b/cpp/src/arrow/CMakeLists.txt @@ -428,7 +428,15 @@ if(ARROW_JEMALLOC) PROPERTIES SKIP_PRECOMPILE_HEADERS ON SKIP_UNITY_BUILD_INCLUSION ON) endif() +if(ARROW_MIMALLOC) + list(APPEND ARROW_MEMORY_POOL_SRCS memory_pool_mimalloc.cc) + set_source_files_properties(memory_pool_mimalloc.cc + PROPERTIES SKIP_PRECOMPILE_HEADERS ON + SKIP_UNITY_BUILD_INCLUSION ON) +endif() + arrow_add_object_library(ARROW_MEMORY_POOL ${ARROW_MEMORY_POOL_SRCS}) + if(ARROW_JEMALLOC) foreach(ARROW_MEMORY_POOL_TARGET ${ARROW_MEMORY_POOL_TARGETS}) target_link_libraries(${ARROW_MEMORY_POOL_TARGET} PRIVATE jemalloc::jemalloc) diff --git a/cpp/src/arrow/memory_pool.cc b/cpp/src/arrow/memory_pool.cc index 2f8ce3a6fa8c7..a2af6bdf5ce77 100644 --- a/cpp/src/arrow/memory_pool.cc +++ b/cpp/src/arrow/memory_pool.cc @@ -54,16 +54,11 @@ #endif namespace arrow { - -namespace memory_pool { - -namespace internal { +namespace memory_pool::internal { alignas(kDefaultBufferAlignment) int64_t zero_size_area[1] = {kDebugXorSuffix}; -} // namespace internal - -} // namespace memory_pool +} // namespace memory_pool::internal namespace { @@ -394,15 +389,15 @@ class MimallocAllocator { *out = memory_pool::internal::kZeroSizeArea; return Status::OK(); } - *out = reinterpret_cast( - mi_malloc_aligned(static_cast(size), static_cast(alignment))); + *out = reinterpret_cast(arrow_mi_malloc_aligned( + static_cast(size), static_cast(alignment))); if (*out == NULL) { return Status::OutOfMemory("malloc of size ", size, " failed"); } return Status::OK(); } - static void ReleaseUnused() { mi_collect(true); } + static void ReleaseUnused() { arrow_mi_collect(true); } static Status ReallocateAligned(int64_t old_size, int64_t new_size, int64_t alignment, uint8_t** ptr) { @@ -417,7 +412,7 @@ class MimallocAllocator { return Status::OK(); } *ptr = reinterpret_cast( - mi_realloc_aligned(previous_ptr, static_cast(new_size), alignment)); + arrow_mi_realloc_aligned(previous_ptr, static_cast(new_size), alignment)); if (*ptr == NULL) { *ptr = previous_ptr; return Status::OutOfMemory("realloc of size ", new_size, " failed"); @@ -429,7 +424,7 @@ class MimallocAllocator { if (ptr == memory_pool::internal::kZeroSizeArea) { DCHECK_EQ(size, 0); } else { - mi_free(ptr); + arrow_mi_free(ptr); } } }; diff --git a/cpp/src/arrow/memory_pool_benchmark.cc b/cpp/src/arrow/memory_pool_benchmark.cc index c2e55314b56f9..6426c43100b55 100644 --- a/cpp/src/arrow/memory_pool_benchmark.cc +++ b/cpp/src/arrow/memory_pool_benchmark.cc @@ -17,6 +17,7 @@ #include "arrow/memory_pool.h" #include "arrow/result.h" +#include "arrow/util/config.h" #include "arrow/util/logging.h" #include "benchmark/benchmark.h" diff --git a/cpp/src/arrow/memory_pool_internal.h b/cpp/src/arrow/memory_pool_internal.h index 01500b3c1eae1..70ec29f4bab3e 100644 --- a/cpp/src/arrow/memory_pool_internal.h +++ b/cpp/src/arrow/memory_pool_internal.h @@ -19,12 +19,9 @@ #include "arrow/memory_pool.h" #include "arrow/util/config.h" +#include "arrow/util/macros.h" -namespace arrow { - -namespace memory_pool { - -namespace internal { +namespace arrow::memory_pool::internal { static constexpr int64_t kDebugXorSuffix = -0x181fe80e0b464188LL; @@ -48,8 +45,16 @@ class JemallocAllocator { #endif // defined(ARROW_JEMALLOC) -} // namespace internal +} // namespace arrow::memory_pool::internal + +#ifdef ARROW_MIMALLOC + +extern "C" { -} // namespace memory_pool +ARROW_NOINLINE void* arrow_mi_malloc_aligned(size_t size, size_t alignment); +ARROW_NOINLINE void* arrow_mi_realloc_aligned(void* p, size_t new_size, size_t alignment); +ARROW_NOINLINE void arrow_mi_free(void* p); +ARROW_NOINLINE void arrow_mi_collect(bool force); +} -} // namespace arrow +#endif // defined(ARROW_MIMALLOC) diff --git a/cpp/src/arrow/memory_pool_test.cc b/cpp/src/arrow/memory_pool_test.cc index 3f0a852876718..83b48b91042a9 100644 --- a/cpp/src/arrow/memory_pool_test.cc +++ b/cpp/src/arrow/memory_pool_test.cc @@ -77,7 +77,10 @@ TYPED_TEST_P(TestMemoryPool, Reallocate) { this->TestReallocate(); } TYPED_TEST_P(TestMemoryPool, Alignment) { this->TestAlignment(); } -REGISTER_TYPED_TEST_SUITE_P(TestMemoryPool, MemoryTracking, OOM, Reallocate, Alignment); +TYPED_TEST_P(TestMemoryPool, ReleaseUnused) { this->TestReleaseUnused(); } + +REGISTER_TYPED_TEST_SUITE_P(TestMemoryPool, MemoryTracking, OOM, Reallocate, Alignment, + ReleaseUnused); INSTANTIATE_TYPED_TEST_SUITE_P(Default, TestMemoryPool, DefaultMemoryPoolFactory); INSTANTIATE_TYPED_TEST_SUITE_P(System, TestMemoryPool, SystemMemoryPoolFactory); diff --git a/cpp/src/arrow/memory_pool_test.h b/cpp/src/arrow/memory_pool_test.h index 32f1cc5d1d310..87f4f2a152ec6 100644 --- a/cpp/src/arrow/memory_pool_test.h +++ b/cpp/src/arrow/memory_pool_test.h @@ -106,6 +106,14 @@ class TestMemoryPoolBase : public ::testing::Test { pool->Free(data512, 10, 512); } } + + void TestReleaseUnused() { + auto pool = memory_pool(); + const int64_t nbytes = pool->bytes_allocated(); + pool->ReleaseUnused(); + // Unfortunately there's not much that we can assert here + ASSERT_EQ(nbytes, pool->bytes_allocated()); + } }; } // namespace arrow