Allocator handle for custom allocators used in shared library (#84)

Dian-Lun (Aaron) Lin · web-flow · commit a88c8e621995 · 2025-02-14T11:51:48.000-08:00
This PR is specifically designed to support a custom allocator that
works with a shared library. It offers an interface that enables the
implementation of a custom allocator not included in the shared library.
Which custom allocator to use is determined at runtime.

This PR provides three utility functions:
1. `make_allocator_handle(Alloc alloc)`: Accepts an allocator and
returns an `AllocatorHandle` to wrap the allocator.
2. `make_blocked_allocator_handle(Alloc alloc)`: Accepts an allocator,
wraps it in an `AllocatorHandle`, and returns a `Blocked` object that
manages the allocator.
3. `make_blocked_allocator_handle(const BlockingParameters&amp; parameters,
Alloc alloc)`: Accepts an allocator, wraps it in an `AllocatorHandle`,
and returns a `Blocked` object configured with the specified parameters.

The `AllocatorHandle` provides `allocate` and `deallocate` functions to
use the wrapped custom allocator for memory allocation.
diff --git a/include/svs/core/allocator.h b/include/svs/core/allocator.h
@@ -507,4 +507,123 @@ class MemoryMapper {
         return MMapPtr<void>(base, bytes.value());
     }
 };
+
+namespace detail {
+
+template <typename Alloc>
+concept HasValueType = requires { typename Alloc::value_type; };
+
+// clang-format off
+template <typename Alloc>
+concept Allocator = HasValueType<Alloc> && std::is_copy_constructible_v<Alloc>
+&& requires(Alloc& alloc, typename Alloc::value_type* ptr, size_t n) {
+    { alloc.allocate(n) } -> std::same_as<typename Alloc::value_type*>;
+
+    { alloc.deallocate(ptr, n) };
+};
+// clang-format on
+
+} // namespace detail
+
+/////
+///// Type erasure allocator handle implementation
+/////
+
+class AllocatorInterface {
+  public:
+    virtual ~AllocatorInterface() = default;
+    virtual void* allocate(size_t n) = 0;
+    virtual void deallocate(void* ptr, size_t n) = 0;
+
+    // covariant return type
+    virtual AllocatorInterface* clone() const = 0;
+};
+
+template <detail::Allocator Impl> class AllocatorImpl : public AllocatorInterface {
+  public:
+    // pass by value due to clone()
+    explicit AllocatorImpl(Impl impl)
+        : AllocatorInterface{}
+        , impl_{std::move(impl)} {}
+
+    void* allocate(size_t n) override { return static_cast<void*>(impl_.allocate(n)); }
+
+    void deallocate(void* ptr, size_t n) override {
+        impl_.deallocate(static_cast<typename Impl::value_type*>(ptr), n);
+    }
+
+    AllocatorImpl<Impl>* clone() const override { return new AllocatorImpl(impl_); }
+
+  private:
+    Impl impl_;
+};
+
+template <typename T> class AllocatorHandle {
+  public:
+    using value_type = T;
+
+    template <detail::Allocator Impl>
+    explicit AllocatorHandle(Impl&& impl)
+        requires(!std::is_same_v<Impl, AllocatorHandle>) &&
+                std::is_rvalue_reference_v<Impl&&> &&
+                std::is_same_v<typename Impl::value_type, T>
+        : impl_{new AllocatorImpl(std::move(impl))} {}
+
+    AllocatorHandle() {}
+    AllocatorHandle(const AllocatorHandle& other)
+        : impl_{other.impl_->clone()} {}
+    AllocatorHandle(AllocatorHandle&&) = default;
+    AllocatorHandle& operator=(const AllocatorHandle& other) {
+        impl_.reset(other.impl_->clone());
+        return *this;
+    }
+    AllocatorHandle& operator=(AllocatorHandle&&) = default;
+    ~AllocatorHandle() = default;
+
+    T* allocate(size_t n) {
+        if (impl_.get() == nullptr) {
+            throw ANNEXCEPTION("Empty allocator handle!");
+        }
+        return static_cast<T*>(impl_->allocate(n));
+    }
+
+    void deallocate(T* ptr, size_t n) {
+        if (impl_.get() == nullptr) {
+            throw ANNEXCEPTION("Empty allocator handle!");
+        }
+        impl_->deallocate(static_cast<void*>(ptr), n);
+    }
+
+  private:
+    std::unique_ptr<AllocatorInterface> impl_;
+};
+
+///
+/// @class allocator_handle
+///
+/// AllocatorHandle
+/// ================
+/// `AllocatorHandle` is designed to support a custom allocator that works with the
+/// shared library. It offers an interface that enables the
+/// implementation of a custom allocator not included in the shared library. Which custom
+/// allocator to use is determined at runtime.
+///
+
+///
+/// @brief Creates an `AllocatorHandle` from a given allocator.
+///
+/// @tparam Alloc The type of the allocator, which must satisfy the `Allocator` concept.
+//
+/// @param alloc The allocator to be wrapped in an `AllocatorHandle`.
+/// @return An `AllocatorHandle` that manages the given allocator.
+///
+/// @copydoc allocator_handle
+///
+/// @sa make_blocked_allocator_handle
+///
+template <detail::Allocator Alloc>
+AllocatorHandle<typename Alloc::value_type> make_allocator_handle(Alloc alloc) {
+    return AllocatorHandle<typename Alloc::value_type>{std::move(alloc)};
+}
+
 } // namespace svs
diff --git a/include/svs/core/data/simple.h b/include/svs/core/data/simple.h
@@ -826,4 +826,44 @@ template <typename T, size_t Extent = Dynamic, typename Alloc = lib::Allocator<T
 using BlockedData = SimpleData<T, Extent, Blocked<Alloc>>;
 
 } // namespace data
+
+///
+/// @brief Creates a `Blocked` object with an `AllocatorHandle` from a given allocator.
+///
+/// @tparam Alloc The type of the allocator, which must satisfy the `Allocator` concept.
+/// @param alloc The allocator to be wrapped in an `AllocatorHandle` and managed by a
+/// `Blocked` object.
+/// @return A `Blocked` object that manages an `AllocatorHandle` wrapping the given
+/// allocator.
+///
+/// @copydoc allocator_handle
+///
+/// @sa make_allocator_handle
+///
+template <detail::Allocator Alloc>
+data::Blocked<AllocatorHandle<typename Alloc::value_type>>
+make_blocked_allocator_handle(Alloc alloc) {
+    return data::Blocked{make_allocator_handle(std::move(alloc))};
+}
+
+///
+/// @brief Creates a `Blocked` object with an `AllocatorHandle` from a given allocator.
+///
+/// @tparam Alloc The type of the allocator, which must satisfy the `Allocator` concept.
+//  @param parameters The blocking parameters used to configure the `Blocked` object.
+/// @param alloc The allocator to be wrapped in an `AllocatorHandle and managed by a
+/// `Blocked` object.
+/// @return A `Blocked` object that manages an `AllocatorHandle` wrapping the given
+/// allocator.
+///
+/// @copydoc allocator_handle
+///
+/// @sa make_allocator_handle
+///
+template <detail::Allocator Alloc>
+data::Blocked<AllocatorHandle<typename Alloc::value_type>>
+make_blocked_allocator_handle(const data::BlockingParameters& parameters, Alloc alloc) {
+    return data::Blocked{parameters, make_allocator_handle(std::move(alloc))};
+}
+
 } // namespace svs
diff --git a/tests/svs/core/allocator.cpp b/tests/svs/core/allocator.cpp
@@ -137,4 +137,64 @@ CATCH_TEST_CASE("Testing Allocator", "[allocators]") {
             }
         }
     }
+
+    CATCH_SECTION("Testing `AllocatorHandle`") {
+        size_t num_elements = 1024;
+        CATCH_SECTION("Allocator") {
+            auto alloc = svs::make_allocator_handle(svs::lib::Allocator<float>());
+            auto* ptr = alloc.allocate(num_elements);
+
+            alloc.deallocate(ptr, num_elements);
+
+            CATCH_STATIC_REQUIRE(std::is_same_v<decltype(ptr), float*>);
+        }
+        CATCH_SECTION("HugepageAllocator - std::byte") {
+            auto alloc = svs::make_allocator_handle(svs::HugepageAllocator<std::byte>());
+            auto* ptr = alloc.allocate(num_elements);
+
+            auto allocations = svs::detail::GenericHugepageAllocator::get_allocations();
+            CATCH_REQUIRE(allocations.size() == 1);
+            CATCH_REQUIRE(allocations.contains(ptr));
+            CATCH_REQUIRE(allocations.at(ptr) >= sizeof(std::byte) * num_elements);
+
+            alloc.deallocate(ptr, num_elements);
+            allocations = svs::detail::GenericHugepageAllocator::get_allocations();
+            CATCH_REQUIRE(allocations.size() == 0);
+            CATCH_REQUIRE(!allocations.contains(ptr));
+
+            CATCH_STATIC_REQUIRE(std::is_same_v<decltype(ptr), std::byte*>);
+        }
+        CATCH_SECTION("HugepageAllocator - int8_t") {
+            auto alloc = svs::make_allocator_handle(svs::HugepageAllocator<int8_t>());
+            auto* ptr = alloc.allocate(num_elements);
+
+            auto allocations = svs::detail::GenericHugepageAllocator::get_allocations();
+            CATCH_REQUIRE(allocations.size() == 1);
+            CATCH_REQUIRE(allocations.contains(ptr));
+            CATCH_REQUIRE(allocations.at(ptr) >= sizeof(int8_t) * num_elements);
+
+            alloc.deallocate(ptr, num_elements);
+            allocations = svs::detail::GenericHugepageAllocator::get_allocations();
+            CATCH_REQUIRE(allocations.size() == 0);
+            CATCH_REQUIRE(!allocations.contains(ptr));
+
+            CATCH_STATIC_REQUIRE(std::is_same_v<decltype(ptr), int8_t*>);
+        }
+        CATCH_SECTION("HugepageAllocator - svs::Float16") {
+            auto alloc = svs::make_allocator_handle(svs::HugepageAllocator<svs::Float16>());
+            auto* ptr = alloc.allocate(num_elements);
+
+            auto allocations = svs::detail::GenericHugepageAllocator::get_allocations();
+            CATCH_REQUIRE(allocations.size() == 1);
+            CATCH_REQUIRE(allocations.contains(ptr));
+            CATCH_REQUIRE(allocations.at(ptr) >= sizeof(svs::Float16) * num_elements);
+
+            alloc.deallocate(ptr, num_elements);
+            allocations = svs::detail::GenericHugepageAllocator::get_allocations();
+            CATCH_REQUIRE(allocations.size() == 0);
+            CATCH_REQUIRE(!allocations.contains(ptr));
+
+            CATCH_STATIC_REQUIRE(std::is_same_v<decltype(ptr), svs::Float16*>);
+        }
+    }
 }
diff --git a/tests/svs/core/compact.cpp b/tests/svs/core/compact.cpp
@@ -125,4 +125,39 @@ CATCH_TEST_CASE("Simple Data Compaction", "[core][compaction]") {
             CATCH_REQUIRE(check_line(data.get_datum(i), val));
         }
     }
+
+    CATCH_SECTION("Blocked Data with AllocatorHandle") {
+        auto data = svs::data::SimpleData<
+            float,
+            svs::Dynamic,
+            svs::data::Blocked<svs::AllocatorHandle<float>>>(
+            100, 20, svs::make_blocked_allocator_handle(svs::lib::Allocator<float>())
+        );
+        sequential_fill(data);
+        CATCH_REQUIRE(check_sequential(data));
+
+        auto new_to_old = std::vector<uint32_t>{};
+        for (size_t i = 0, imax = data.size() / 3; i < imax; ++i) {
+            new_to_old.push_back(3 * i);
+        }
+
+        // Single-threaded version
+        data.compact(svs::lib::as_const_span(new_to_old), 20);
+        for (size_t i = 0, imax = new_to_old.size(); i < imax; ++i) {
+            auto val = new_to_old.at(i);
+            CATCH_REQUIRE(check_line(data.get_datum(i), val));
+        }
+
+        // Multi-threaded version.
+        data.resize(0);
+        data.resize(100);
+        sequential_fill(data);
+        CATCH_REQUIRE(check_sequential(data));
+        auto tpool = svs::threads::DefaultThreadPool(2);
+        data.compact(svs::lib::as_const_span(new_to_old), tpool, 20);
+        for (size_t i = 0, imax = new_to_old.size(); i < imax; ++i) {
+            auto val = new_to_old.at(i);
+            CATCH_REQUIRE(check_line(data.get_datum(i), val));
+        }
+    }
 }
