Add configurable pool properties to PinnedMemoryResource by nirandaperera · Pull Request #851 · rapidsai/rapidsmpf

nirandaperera · 2026-02-09T15:48:31Z

This PR adds PinnedPoolProperties struct to allow configuration of pinned memory pool behavior via initial_pool_size and max_pool_size parameters.

Changes

New PinnedPoolProperties struct with two configurable fields:
- initial_pool_size: Pre-allocates pinned memory for improved performance
- max_pool_size: Limits maximum pool size (0 = unlimited)
Updated PinnedMemoryResource constructor to accept optional PinnedPoolProperties parameter with backward-compatible defaults

Initial allocation benchmark

Added a benchmark to test the impact of initial_pool_size. On my workstation (RTX A6000) it shows around 10x allocation performance when the pool is set an initial size.

---------------------------------------------------------------------------------------------------------------
Benchmark                                                     Time             CPU   Iterations UserCounters...
---------------------------------------------------------------------------------------------------------------
BM_PinnedFirstAlloc_InitialPoolSize/1/1/real_time          1301 us         1300 us          538 bytes_per_second=768.881M/s initial_pool_size=1048.58k
BM_PinnedFirstAlloc_InitialPoolSize/1/0/real_time          9518 us         9515 us           70 bytes_per_second=105.067M/s initial_pool_size=0
BM_PinnedFirstAlloc_InitialPoolSize/256/1/real_time       10532 us        10517 us           67 bytes_per_second=23.7374G/s initial_pool_size=268.435M
BM_PinnedFirstAlloc_InitialPoolSize/256/0/real_time       85740 us        85706 us            8 bytes_per_second=2.91578G/s initial_pool_size=0
BM_PinnedFirstAlloc_InitialPoolSize/1024/1/real_time      45428 us        45366 us           16 bytes_per_second=22.0129G/s initial_pool_size=1073.74M
BM_PinnedFirstAlloc_InitialPoolSize/1024/0/real_time     302111 us       301913 us            2 bytes_per_second=3.31004G/s initial_pool_size=0

…s_pinned_pool

pentschev · 2026-02-09T20:58:03Z

+ */
+struct PinnedPoolProperties {
+    std::size_t initial_pool_size = 0;  ///< initial size of the pool. Initial size is
+                                        ///< important for pinned memory performance.


This comment seems contrary https://github.com/rapidsai/rapidsmpf/pull/851/changes#diff-eabfc19cbc8bce3fdbf90c6fa14736f3e52ae669bcd3822fa39b6c57d26a14d8R34-R35, which one is applicable now?

I think, priming pools (ie. make some allocations up front and deacllocating) has little effect on device pools. But for pinned memory pools, initial allocation and warming up is important. I extended the comment to include this.

What does "important for pinned memory performance" mean? How is it important?

pentschev · 2026-02-09T20:58:51Z

        // Before <https://github.com/NVIDIA/cccl/pull/6718>, the default
        // `release_threshold` was 0, which defeats the purpose of having a pool. We
        // now set it so the pool never releases unused pinned memory.


This comment seems outdated now, no?

pentschev · 2026-02-09T21:00:38Z

        // `release_threshold` was 0, which defeats the purpose of having a pool. We
        // now set it so the pool never releases unused pinned memory.
-        .release_threshold = std::numeric_limits<size_t>::max(),
+        .release_threshold = props.max_pool_size > 0 ? props.max_pool_size


max_pool_size seems to imply that the pool cannot go beyond that limit. However, AFAIU, release_threshold is something different, meaning it will start releasing unused memory back to the driver once that limit is reached. Can you clarify what max_pool_size really means, and if necessary rename it to release_threshold or something more accurate?

Thinking about this again, I feel like my change is wrong. I reverted this. Thank you for catching this.

pentschev · 2026-02-09T21:01:09Z

        // It was observed that priming async pools have little effect for performance.
        // See <https://github.com/rapidsai/rmm/issues/1931>.
-        .initial_pool_size = 0,
+        .initial_pool_size = props.initial_pool_size,
        // Before <https://github.com/NVIDIA/cccl/pull/6718>, the default
        // `release_threshold` was 0, which defeats the purpose of having a pool. We
        // now set it so the pool never releases unused pinned memory.
-        .release_threshold = std::numeric_limits<size_t>::max(),
+        .release_threshold = props.max_pool_size > 0 ? props.max_pool_size


All comments above are applicable here.

wence- · 2026-02-10T10:52:58Z

+ */
+struct PinnedPoolProperties {
+    std::size_t initial_pool_size = 0;  ///< initial size of the pool. Initial size is
+                                        ///< important for pinned memory performance.


What does "important for pinned memory performance" mean? How is it important?

wence- · 2026-02-10T11:01:04Z

+    // Create a PinnedMemoryResource with max pool size of 1MiB
+    auto pinned_mr = std::make_shared<rapidsmpf::PinnedMemoryResource>(
+        rapidsmpf::get_current_numa_node(),
+        rapidsmpf::PinnedPoolProperties{.initial_pool_size = 0, .max_pool_size = 1_MiB}
+    );
+    auto stream = cudf::get_default_stream();
+
+    void* ptr = pinned_mr->allocate(stream, 512_KiB);
+    EXPECT_NE(nullptr, ptr);
+    pinned_mr->deallocate(stream, ptr, 512_KiB);
+
+    // NOTE: currently cuda driver rounds up max size to 32MB. So we need to allocate 32MB
+    // + 1 byte.
+    EXPECT_THROW(
+        {
+            void* ptr2 = pinned_mr->allocate(stream, 32_MiB + 1);


Can we inspect the max size we get so that this test is robust?

wence- · 2026-02-11T15:01:33Z

 #if CCCL_MAJOR_VERSION > 3 || (CCCL_MAJOR_VERSION == 3 && CCCL_MINOR_VERSION >= 2)
-cuda::memory_pool_properties get_memory_pool_properties() {
+cuda::memory_pool_properties get_memory_pool_properties(


I note that we are now on CCCL 3.2.1, so we should be able to remove all of this macro-conditional stuff.

It's also no longer gated behind experimental, so I think we can avoid needing to turn on experimental mode in cccl (and can just get the version from RMM).

And, the headers now no longer require nvcc, so we can remove the pimpl idiom.

Before doing this refactoring, can we migrate to the CCCL >= 3.2 version of the memory pool resource. It's plausible that we can then delete huge swathes of this code anyway

Yes! Let me send a separate PR for it.

@wence- #856 I opened a new PR for this

…s_pinned_pool Signed-off-by: niranda perera <niranda.perera@gmail.com>

wence- · 2026-02-13T12:15:31Z

+/// Discover the actual pool size the driver creates when a small max is requested.
+/// Creates a pool with \p requested_max_pool_size (e.g. 1 MiB), then uses recursive
+/// doubling of allocation size until allocation fails; returns the last successful size.
+std::size_t discover_pinned_pool_actual_size(
+    rmm::cuda_stream_view stream, std::size_t requested_max_pool_size = 1_MiB
+) {
+    rapidsmpf::PinnedMemoryResource pinned_mr{
+        rapidsmpf::get_current_numa_node(),
+        rapidsmpf::PinnedPoolProperties{.max_pool_size = requested_max_pool_size}
+    };
+    std::size_t try_size = requested_max_pool_size;
+    while (true) {
+        try {
+            void* ptr = pinned_mr.allocate(stream, try_size);
+            pinned_mr.deallocate(stream, ptr, try_size);
+            try_size *= 2;
+        } catch (cuda::cuda_error const&) {
+            break;
+        }
+    }
+    return std::max(try_size / 2, requested_max_pool_size);
+}


This should probably do bisection search, in case the actual size is not a power of two.

…s_pinned_pool

Signed-off-by: niranda perera <niranda.perera@gmail.com>

madsbk

Update PinnedMemoryResource::from_options with a pinned_memory_initial_pool_size option (also update READSME.md).

madsbk · 2026-02-18T19:41:26Z

+    std::size_t initial_pool_size = 0;
+
+    /// @brief Maximum size of the pool. 0 means no limit.
+    std::size_t max_pool_size = 0;


Use std::optional instead of 0

…s_pinned_pool

Signed-off-by: niranda perera <niranda.perera@gmail.com>

madsbk · 2026-02-19T07:31:35Z

@nirandaperera, we should introduce an config option pinned_memory_initial_pool_size in PinnedMemoryResource::from_options.

…s_pinned_pool

nirandaperera · 2026-02-20T18:04:20Z

@nirandaperera, we should introduce an config option pinned_memory_initial_pool_size in PinnedMemoryResource::from_options.

@madsbk done

madsbk · 2026-02-23T14:53:25Z

+            .max_pool_size = options.get<std::optional<size_t>>(
+                "pinned_max_pool_size", [](auto const& s) {
+                    return s.empty() ? std::nullopt
+                                     : std::optional<size_t>(parse_string<size_t>(s));
+                }
+            )


Use parse_optional() to handle the optional case before parsing it to parse_string<size_t> like we do here:

rapidsmpf/cpp/src/memory/buffer_resource.cpp

Line 266 in c48f3a8

if (auto val = parse_optional(s); val.has_value()) {

madsbk · 2026-02-23T15:03:37Z

+                [](auto const& s) { return parse_string<size_t>(s.empty() ? "0" : s); }
+            ),
+            .max_pool_size = options.get<std::optional<size_t>>(
+                "pinned_max_pool_size", [](auto const& s) {


Docs the pinned_max_pool_size option:

rapidsmpf/docs/source/configuration.md

Line 33 in c48f3a8

### General

madsbk · 2026-02-23T15:03:55Z

+    if (pinned_memory) {
+        PinnedPoolProperties pool_properties{
+            .initial_pool_size = options.get<size_t>(
+                "pinned_initial_pool_size",


Docs the pinned_initial_pool_size option:

rapidsmpf/docs/source/configuration.md

Line 33 in c48f3a8

### General

Co-authored-by: Mads R. B. Kristensen <madsbk@gmail.com>

…apidsmpf into enable_props_pinned_pool

…s_pinned_pool

nirandaperera · 2026-03-02T20:07:48Z

@madsbk @wence- Let's get this in as well.

madsbk

LGTM

madsbk · 2026-03-03T14:08:03Z

+- **`pinned_initial_pool_size`**
+  - **Environment Variable**: `RAPIDSMPF_PINNED_INITIAL_POOL_SIZE`
+  - **Default**: `0`
+  - **Description**: Initial size (in bytes) of the pinned host memory pool when
+    `pinned_memory` is enabled. A value of `0` means the pool starts empty and grows
+    on demand. Accepts byte counts (e.g. `"1GiB"`, `"512MiB"`).
+
+- **`pinned_max_pool_size`**
+  - **Environment Variable**: `RAPIDSMPF_PINNED_MAX_POOL_SIZE`
+  - **Default**: `"disabled"`
+  - **Description**: Maximum size (in bytes) of the pinned host memory pool when
+    `pinned_memory` is enabled. When unset or empty, the pool is allowed to grow
+    without an upper bound. Accepts byte counts (e.g. `"4GiB"`, `"2048MiB"`).
+


move up to the pinned_memory section

Removed pinned memory pool size parameters from documentation.

nirandaperera · 2026-03-05T01:29:28Z

/merge

@wence-

@wence- I will do the requested changes in follow-up PR

nirandaperera added 2 commits February 9, 2026 07:42

enabling pinned pool properties

88754f7

Merge branch 'main' of github.com:rapidsai/rapidsmpf into enable_prop…

1924107

…s_pinned_pool

nirandaperera requested a review from a team as a code owner February 9, 2026 15:48

nirandaperera added improvement Improves an existing functionality non-breaking Introduces a non-breaking change labels Feb 9, 2026

fix test

a625a8d

pentschev reviewed Feb 9, 2026

View reviewed changes

addressing PR comments

8b745e1

nirandaperera requested a review from pentschev February 9, 2026 23:28

wence- requested changes Feb 11, 2026

View reviewed changes

wence- mentioned this pull request Feb 12, 2026

Removing cudax components #856

Merged

nirandaperera added 2 commits February 12, 2026 12:25

Merge branch 'main' of github.com:rapidsai/rapidsmpf into enable_prop…

2ab5c09

…s_pinned_pool Signed-off-by: niranda perera <niranda.perera@gmail.com>

add a bench for init alloc

33dd607

nirandaperera requested a review from wence- February 12, 2026 21:24

making tests robust

137f0f2

wence- reviewed Feb 13, 2026

View reviewed changes

pentschev reviewed Feb 16, 2026

View reviewed changes

Comment thread cpp/benchmarks/bench_memory_resources.cpp Outdated

nirandaperera added 2 commits February 17, 2026 13:00

addressing PR comments

628083b

Merge branch 'main' of github.com:rapidsai/rapidsmpf into enable_prop…

9d82b2b

…s_pinned_pool

nirandaperera requested review from pentschev and wence- February 17, 2026 21:00

using shared_resource

67fa83c

Signed-off-by: niranda perera <niranda.perera@gmail.com>

madsbk reviewed Feb 18, 2026

View reviewed changes

nirandaperera added 2 commits February 18, 2026 13:04

Merge branch 'main' of github.com:rapidsai/rapidsmpf into enable_prop…

16bf572

…s_pinned_pool

addressing PR comments

92d3dbe

Signed-off-by: niranda perera <niranda.perera@gmail.com>

nirandaperera requested a review from madsbk February 20, 2026 17:46

nirandaperera added 2 commits February 20, 2026 10:03

addressing PR comments

c9cda17

Merge branch 'main' of github.com:rapidsai/rapidsmpf into enable_prop…

ada1678

…s_pinned_pool

madsbk requested changes Feb 23, 2026

View reviewed changes

wence- previously requested changes Feb 24, 2026

View reviewed changes

Comment thread cpp/include/rapidsmpf/memory/pinned_memory_resource.hpp

Comment thread cpp/include/rapidsmpf/memory/pinned_memory_resource.hpp

Update cpp/tests/test_host_buffer.cpp

ba285dc

Co-authored-by: Mads R. B. Kristensen <madsbk@gmail.com>

nirandaperera requested review from madsbk and wence- March 2, 2026 20:04

nirandaperera added 3 commits March 2, 2026 12:05

addressing PR comments

4a213e9

Merge branch 'enable_props_pinned_pool' of github.com:nirandaperera/r…

be3961f

…apidsmpf into enable_props_pinned_pool

Merge branch 'main' of github.com:rapidsai/rapidsmpf into enable_prop…

8f383e6

…s_pinned_pool

madsbk approved these changes Mar 3, 2026

View reviewed changes

nirandaperera added 2 commits March 4, 2026 17:29

Remove pinned memory pool size configuration details

da436ae

Removed pinned memory pool size parameters from documentation.

Merge branch 'main' into enable_props_pinned_pool

6f98f3d

Merge branch 'main' into enable_props_pinned_pool

d4c4fe5

rapids-bot bot merged commit 42a0304 into rapidsai:main Mar 5, 2026
66 checks passed

wence- mentioned this pull request Mar 6, 2026

Adding PinnedPoolProperties cython bindings #904

Open

Conversation

nirandaperera commented Feb 9, 2026 • edited by wence- Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Changes

Initial allocation benchmark

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

madsbk left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

madsbk commented Feb 19, 2026

Uh oh!

nirandaperera commented Feb 20, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

nirandaperera commented Mar 2, 2026

Uh oh!

madsbk left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

nirandaperera commented Mar 5, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

nirandaperera commented Feb 9, 2026 •

edited by wence-

Loading