[STF] use driver's memory pool for locality domains - #11202
Conversation
The localized per-(device, locality-domain) memory pools are created with the default release threshold of 0, and the device default pools get their threshold set only inside the peer-access loop — so on single-device machines neither pool retains freed memory: every synchronization returns freed pages to the OS, and every subsequent large stream-ordered allocation pays page re-backing (~2 ms/GB per allocation cycle). - exec/locality_domain.cuh: set CU_MEMPOOL_ATTR_RELEASE_THRESHOLD = UINT64_MAX on the localized pools at creation, mirroring the device default pools. - machine.cuh: set the device default-pool threshold unconditionally, before the peer loop, so single-device machines get it too. Measured on a 268M-key fp32 sort over 2 locality domains (an n-sized per-call scratch consumer): 10.4 ms -> 5.1 ms end to end, restoring parity with a whole-device sort. Any algorithm allocating large stream-ordered temporaries per call pays the same mechanism proportionally. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
/ok to test f654606 |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 SummarySummary by CodeRabbit
WalkthroughThe locality-domain implementation shares allocation-location construction and retrieves the appropriate driver default memory pool. The private pool cache is removed. A regression test verifies that locality-domain allocation does not change the device default pool release threshold. ChangesMemory pool selection
Merge Risk: 🟡 Moderate · up to The new default-pool path can select the wrong locality domain for invalid large IDs and adds driver topology queries to every allocation; a separate failure path may leave the caller on the wrong CUDA device. These issues should be addressed or explicitly accepted before merge. ✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (1)
cudax/include/cuda/experimental/__places/exec/locality_domain.cuh (1)
240-240: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winsuggestion: Add a direct integer header in both changed files.
The new threshold code uses
UINT64_MAXin both files, andmachine.cuhalso usesuint64_t. Add#include <cstdint>directly in each file instead of relying on transitive CUDA includes.
cudax/include/cuda/experimental/__places/exec/locality_domain.cuh#L240-L240: add the direct integer header forUINT64_MAX.cudax/include/cuda/experimental/__places/machine.cuh#L84-L85: add the direct integer header foruint64_tandUINT64_MAX.As per coding guidelines, include all headers needed by the symbols being used and do not rely on transitive includes.
Source: Coding guidelines
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 9f5c938a-5236-4cb4-b8e4-7f7b328bef6f
📒 Files selected for processing (2)
cudax/include/cuda/experimental/__places/exec/locality_domain.cuhcudax/include/cuda/experimental/__places/machine.cuh
Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.
| uint64_t threshold = UINT64_MAX; | ||
| cuda_try(cudaMemPoolSetAttribute(mempool, cudaMemPoolAttrReleaseThreshold, &threshold)); |
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟠 Major | ⚡ Quick win
🔎 Supported by static analysis
🏁 Script executed:
#!/bin/bash
set -eu
printf '%s\n' '--- target file map ---'
ast-grep outline cudax/include/cuda/experimental/__places/machine.cuh
printf '%s\n' '--- target source ---'
sed -n '1,150p' cudax/include/cuda/experimental/__places/machine.cuh
printf '%s\n' '--- referenced helper locations ---'
rg -n -C 12 'cudaMemPool(Set|Get)Attribute|ReleaseThreshold|release.?threshold|UINT64_MAX' python/cuda_cccl/cuda/compute/_utils/temp_storage_buffer.py cudax/include/cuda/experimental/__places/machine.cuhRepository: NVIDIA/cccl
Length of output: 11739
🤖 get_repo_knowledge executed:
get_repo_knowledge NVIDIA/cccl /tmp/coderabbit-repo-knowledge/nvidia-cccl-2f773195/architecture /tmp/coderabbit-repo-knowledge/nvidia-cccl-2f773195/learnings /tmp/coderabbit-repo-knowledge/nvidia-cccl-2f773195/conventions
Length of output: 47511
🏁 Script executed:
#!/bin/bash
set -eu
printf '%s\n' '--- complete threshold helper ---'
sed -n '20,45p' python/cuda_cccl/cuda/compute/_utils/temp_storage_buffer.py
printf '%s\n' '--- machine references and nearby API context ---'
rg -n -C 8 'places::reserved::machine|reserved::machine|enable_peer_accesses\(|get_ith_closest_node\(' cudax test python docsRepository: NVIDIA/cccl
Length of output: 10968
🌐 Web query:
CUDA Programming Guide stream-ordered memory allocation cudaMemPoolAttrReleaseThreshold UINT64_MAX release threshold synchronization pool trimming
💡 Result:
In CUDA's stream-ordered memory allocator, the cudaMemPoolAttrReleaseThreshold attribute determines the amount of reserved memory in bytes that a memory pool will retain before attempting to release excess memory back to the operating system [1][2]. By default, this threshold is set to 0, meaning the allocator attempts to minimize the physical memory footprint of the pool by releasing memory after every synchronization event (such as stream, event, or device synchronization) if the held amount exceeds the threshold [3][1]. Setting the cudaMemPoolAttrReleaseThreshold attribute to UINT64_MAX effectively disables this automatic shrinking behavior, preventing the driver from attempting to release memory back to the OS after synchronizations [3][4]. This is useful for applications that wish to maintain a steady memory footprint and avoid the overhead of repeated OS-level memory management calls [3]. Because setting this threshold to UINT64_MAX prevents the pool from shrinking automatically, you can use the cudaMemPoolTrimTo API to manually manage the memory footprint [3][4]. The cudaMemPoolTrimTo function allows you to explicitly instruct the pool to release memory back to the OS until it reaches a specified size (the minBytesToKeep parameter), enabling precise control over memory usage during different phases of an application's execution [5][4]. Example usage to disable automatic shrinking: cuuint64_t setVal = UINT64_MAX; cudaMemPoolSetAttribute(memPool, cudaMemPoolAttrReleaseThreshold, &setVal); [3][4]
Citations:
- 1: https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__MEMORY__POOLS.html
- 2: https://docs.nvidia.com/cuda/developer-preview/13.4/cuda-runtime-api/cuda_runtime_api/group__CUDART__MEMORY__POOLS.html
- 3: https://docs.nvidia.com/cuda/cuda-programming-guide/04-special-topics/stream-ordered-memory-allocation.html
- 4: https://docs.nvidia.com/cuda/archive/13.3.0/cuda-programming-guide/04-special-topics/stream-ordered-memory-allocation.html
- 5: https://docs.nvidia.com/cuda/archive/12.9.1/cuda-runtime-api/group__CUDART__MEMORY__POOLS.html
important: Preserve caller-configured release thresholds. machine::machine() invokes enable_peer_accesses(), which unconditionally sets each default pool's cudaMemPoolAttrReleaseThreshold to UINT64_MAX. This can override a caller's nonzero threshold and prevent automatic pool shrinking, increasing memory pressure. Read the existing attribute and set UINT64_MAX only when it is zero, or make the override an explicit contract.
Sources: Path instructions, MCP tools
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🔎 Supported by static analysis
🏁 Script executed:
#!/bin/bash
set -e
printf '%s\n' '--- machine.cuh relevant lines ---'
cat -n cudax/include/cuda/experimental/__places/machine.cuh | sed -n '60,135p'
printf '%s\n' '--- direct uses of current_dev and cudaSetDevice in file ---'
rg -n -C 3 'current_dev|cudaSetDevice|cudaMemPoolSetAttribute|cuda_try|scope_guard' cudax/include/cuda/experimental/__places/machine.cuh
printf '%s\n' '--- nearby helper definitions ---'
rg -n -C 4 'cuda_try|scope_guard|make_scope|current device|cudaGetDevice' cudax/include/cuda/experimental cudax/include | head -240Repository: NVIDIA/cccl
Length of output: 29231
🤖 get_repo_knowledge executed:
get_repo_knowledge NVIDIA/cccl /tmp/coderabbit-repo-knowledge/nvidia-cccl-2f773195/architecture /tmp/coderabbit-repo-knowledge/nvidia-cccl-2f773195/learnings /tmp/coderabbit-repo-knowledge/nvidia-cccl-2f773195/conventions
Length of output: 47883
🏁 Script executed:
#!/bin/bash
set -e
printf '%s\n' '--- cuda_try binding and implementation ---'
rg -n -C 8 'cuda_try' cudax/include/cuda/experimental/__stf cudax/include/cuda/experimental/__places | head -260
printf '%s\n' '--- machine header prologue and callers ---'
cat -n cudax/include/cuda/experimental/__places/machine.cuh | sed -n '1,65p'
rg -n -C 5 'enable_peer_accesses|machine\(' cudax/include/cuda/experimental/__places cudax/test | head -220Repository: NVIDIA/cccl
Length of output: 32667
🏁 Script executed:
#!/bin/bash
set -e
printf '%s\n' '--- cuda_safe_call definitions ---'
cat -n cudax/include/cuda/experimental/__stf/utility/cuda_safe_call.cuh | sed -n '1,220p'
printf '%s\n' '--- scope guard definitions used by cudax ---'
rg -n -C 6 '`#define` SCOPE|class.*scope|struct.*scope|SCOPE\(exit\)' cudax/include/cuda/experimental/__stf cudax/include/cuda/experimental/__places | head -220Repository: NVIDIA/cccl
Length of output: 28938
🏁 Script executed:
#!/bin/bash
set -e
printf '%s\n' '--- cuda_try and cuda_safe_call implementation ---'
rg -n 'cuda_try|cuda_safe_call|template.*auto|template.*typename' cudax/include/cuda/experimental/__stf/utility/cuda_safe_call.cuh
cat -n cudax/include/cuda/experimental/__stf/utility/cuda_safe_call.cuh | sed -n '220,430p'
printf '%s\n' '--- established device-restoration pattern ---'
cat -n cudax/include/cuda/experimental/__places/exec/green_context.cuh | sed -n '95,150p'
cat -n cudax/include/cuda/experimental/__places/data_place_impl.cuh | sed -n '245,300p'Repository: NVIDIA/cccl
Length of output: 15655
important: Restore the caller's device on exception.
cuda_try throws on CUDA errors. A failure after cudaSetDevice(d) skips the restore at line 118 and leaves the caller on d. Use a SCOPE(exit) guard to restore current_dev on every exit path. Add a multi-device failure-injection test that checks cudaGetDevice() after the exception.
Sources: Path instructions, MCP tools
This comment has been minimized.
This comment has been minimized.
Drop the unconditional device DEFAULT-pool threshold: machine's initialization runs from every STF context creation (backend_ctx calls machine::instance()), so an unbounded retention threshold on the process-global default pool changes allocator behavior for every consumer of the process — and on shared single-GPU CI runners, several test workers retaining freed memory indefinitely can starve one another (observed as an allocation-path failure in the Python STF suite on this PR's CI). The default pool keeps main's existing behavior. The localized per-(device, locality-domain) pools keep the threshold: the places layer creates and owns those pools, so retention there has no footprint beyond its own allocations — and it alone restores the measured regression on real domain topologies (268M-key sort over 2 locality domains: 10.4 ms -> 5.3 ms, parity with a whole-device sort; the default-pool half only served the no-VMM whole-device fallback path). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…e ones Replace the pool cache's cuMemPoolCreate + manual release-threshold configuration with cuda::__get_default_memory_pool on the same location (localized when available, whole-device degrade otherwise, same PINNED allocation type). The default-pool policy — including repairing a zero release threshold to unbounded so freed memory is retained across calls — then lives in its one library-wide site instead of being duplicated here, the pool is shared with every other consumer of that location in the process, and this layer no longer owns any pool lifetime at all. Measured: identical to the private-pool configuration on the localized path (268M-key fp32 sort over 2 locality domains: parity with a whole-device sort, 0.99x), and the whole-device fallback path is now covered by the same central policy as well (1.07x across sizes, from ~1.8x before). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
With the pools coming from `cuda::__get_default_memory_pool`, the layer's own pool infrastructure has nothing left to do: the driver keeps one default pool per (location, allocation type) and returns the same handle on every call, so there is no pool to create, configure, own, cache or destroy here. - Remove `locality_domain_mem_pool_cache` entirely (singleton, map, mutex, copy/assign deletions) and call the accessor directly from `allocate`. Measured at 144 ns per call — below the cost of the stream-ordered allocation it precedes, and comparable to the map-plus-mutex lookup it replaces. - Factor the location computation the allocation and VMM paths share into `__pool_location()` (same locality-domain / whole-device degrade rule as before, now stated once). - Update the file's documentation to describe what the layer actually does: it obtains the driver's default pool for the domain's location and owns no pool of its own; the release-threshold policy lives in libcu++'s accessor. (`<map>` and `<mutex>` stay: the green-context cache uses them.) Certified: places tests (include_only, data_place_alloc, data_place_vmm, exec_place_scope) under the strict warning set in C++17 and C++20, on native multi-domain and single-domain topologies. Substrate parity unchanged (268M-key fp32 sort: 1.01x whole-device at 2 domains, 1.07x at one). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
There was a problem hiding this comment.
Actionable comments posted: 2
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: d1fe16e3-42a1-4707-b6ea-ce3401a0ea22
📒 Files selected for processing (1)
cudax/include/cuda/experimental/__places/exec/locality_domain.cuh
Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.
| * answer the locality-domain query (whole-device degrade — the localized | ||
| * location type would be rejected). | ||
| */ | ||
| CUmemLocation __pool_location() const |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
important: Add the required declarations to the new functions.
Line 608 and Line 625 omit _CCCL_HOST_API. Add it to both functions. Mark __pool_location [[nodiscard]] and noexcept. Mark mem_create noexcept.
As per path instructions, follow the common CCCL style guidance. As per coding guidelines, functions require a CCCL API annotation and non-throwing functions require noexcept.
Also applies to: 625-625
Sources: Coding guidelines, Path instructions
… path A locality-domain data place falls back to whole-device memory when localization is disabled or the driver reports no domains. On that path the location is `CU_MEM_LOCATION_TYPE_DEVICE`, whose default pool is the process-global one every `cudaMallocAsync` user shares — so routing it through the policy-applying accessor made this place set an unbounded release threshold on a pool it does not own, on exactly the machines that have no locality domains to begin with. Split the accessor by whose pool it is: - locality-domain location: `cuda::__get_default_memory_pool`, which settles the release-threshold policy for that pool. This is the workload's own location, and retention there is what keeps algorithm-scale scratch from being re-backed on every call. - whole-device degrade location: the raw driver getter, inheriting whatever policy the process already has. Retention on the shared pool is not this place's decision, and the behavior on that path is now identical to before this branch. Adds a regression test asserting the device default pool's release threshold is unchanged across a locality-domain allocation (it reproduces the issue against the previous revision: 0 vs SIZE_MAX). Certified: places tests under the strict warning set in C++17 and C++20, on native multi-domain, single-domain and degraded (localization disabled) configurations. Substrate parity unchanged (268M-key fp32 sort: 1.00x whole-device at 2 domains). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
There was a problem hiding this comment.
Actionable comments posted: 2
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: e3795f7e-7a98-4c95-8876-281c7cd7144b
📒 Files selected for processing (2)
cudax/include/cuda/experimental/__places/exec/locality_domain.cuhcudax/test/places/data_place_alloc.cu
Included review availability: Your plan provides up to 12 included reviews per hour; 10 remain after this review.
| // otherwise be configuring a pool it does not own. (The domain's own default | ||
| // pool is a different location and is configured by the library-wide | ||
| // accessor.) | ||
| void test_device_default_pool_policy_untouched() |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
important: Mark test_device_default_pool_policy_untouched with _CCCL_HOST_API inline. This new non-template, non-constexpr function requires the CCCL API annotation and inline.
As per coding guidelines, “Functions must be marked with _CCCL_HOST_API, _CCCL_DEVICE_API, _CCCL_HOST_DEVICE_API, _CCCL_TILE_API, or _CCCL_API” and “Non-template, non-constexpr functions must be declared inline.”
Source: Coding guidelines
| cudaStream_t stream; | ||
| cuda_try(cudaStreamCreate(&stream)); | ||
|
|
||
| auto place = data_place::locality_domain(0, 0); |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
important: Declare place and ptr const. Neither variable is reassigned after initialization. Use const auto place and void* const ptr.
As per coding guidelines, “All variables that are not modified must be declared const.”
Also applies to: 208-208
Source: Coding guidelines
- `__pool_location` gains `[[nodiscard]] _CCCL_HOST_API ... noexcept`: it only fills a struct from two non-throwing queries (the environment check reads a cached getenv result; the domain-count query checks every CUresult and reports 0 on failure), and its result is meaningless to discard. - Move the "create physical memory localized to this domain" comment block back onto `mem_create`, where it belongs; introducing the helper had left it stranded one declaration above. `mem_create` itself is intentionally left as it was. `noexcept` would not hold across both backends: the whole-device fallback implementation of the same member calls through `data_place::device(...)`, which can throw, so the specification would differ by build configuration — and a throw there should stay a catchable error rather than a terminate. `_CCCL_HOST_API` is likewise skipped on the override: the base declaration and every sibling `mem_create` in `data_place_impl.cuh` are written without it, and changing that is a convention decision for the whole data-place family rather than for one override. Certified: places tests under the strict warning set in C++17 and C++20, on native multi-domain, single-domain and degraded configurations; substrate parity unchanged (268M-key fp32 sort: 0.99x whole-device at 2 domains). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
cudax/include/cuda/experimental/__places/exec/locality_domain.cuh (1)
664-668: 🚀 Performance & Scalability | 🟠 Major | 🏗️ Heavy liftimportant: Cache the raw locality-domain count per device.
__pool_location()repeatscuInit(0),cuDeviceGet(), andcuDeviceGetAttribute()for everymem_create()andallocate(). Reuse a shared per-device result, asexec_place_resources::raw_counts_already does.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: f2ab161c-a0f2-40b6-9007-c3da14c57449
📒 Files selected for processing (1)
cudax/include/cuda/experimental/__places/exec/locality_domain.cuh
Included review availability: Your plan provides up to 12 included reviews per hour; 9 remain after this review.
| * which memory this place refers to. | ||
| */ | ||
| CUresult mem_create(CUmemGenericAllocationHandle* handle, size_t size) const override | ||
| [[nodiscard]] _CCCL_HOST_API CUmemLocation __pool_location() const noexcept |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🔎 Supported by static analysis
🏁 Script executed:
#!/bin/bash
set -eu
file='cudax/include/cuda/experimental/__places/exec/locality_domain.cuh'
printf '%s\n' '--- changed hunk ---'
git diff --unified=30 -- "$file"
printf '%s\n' '--- relevant definitions and callers ---'
rg -n -C 8 '__pool_location|domain_id|mem_create|allocate|CUmemLocation' "$file"Repository: NVIDIA/cccl
Length of output: 15780
🏁 Script executed:
#!/bin/bash
set -eu
file='cudax/include/cuda/experimental/__places/exec/locality_domain.cuh'
printf '%s\n' '--- view, count, and native backend contracts ---'
sed -n '180,340p' "$file"
sed -n '450,540p' "$file"
printf '%s\n' '--- data-place construction and allocation path ---'
sed -n '550,680p' "$file"
printf '%s\n' '--- locality-domain view declarations and public API ---'
rg -n -C 12 'struct locality_domain_view|class locality_domain_view|locality_domain_view\(|locality_domain_count|locality_domain_native_raw_count|locality_domain_memory_disabled' "$file"Repository: NVIDIA/cccl
Length of output: 29842
important: Validate view_.domain_id before narrowing it to unsigned char. data_place::locality_domain(view) accepts the token without validation, and the native locality_domain_data_place_impl uses the localized path when domains are available. For domain_id = 256 + k, the cast wraps to k, so mem_create() or allocate() can target the wrong domain. Preserve the separate CUresult and exception contracts.
caugonnet
changed the title
| // accessor.) | ||
| void test_device_default_pool_policy_untouched() | ||
| { | ||
| printf("Testing that a locality-domain place leaves the device default pool alone...\n"); |
There was a problem hiding this comment.
remove this print, tests are silent
|
pre-commit.ci autofix |
|
/ok to test 99f145d |
🥳 CI Workflow Results🟩 Finished in 1h 17m: Pass: 100%/63 | Total: 1d 17h | Max: 1h 11m | Hits: 9%/242119See results here. |
1 participant
Stop creating our own singleton with per domain memory pools, which have bad properties such as a 0 threshold, but use CCCL constructs. We don't want to build over the locality domain abstractions because we have our own contexts, and that we merely describe the data place as a pair of ints, which is a name, not the pool of some actual context.
Description
closes
Checklist