[BUG] Missing memory-resource in device-allocating APIs and helpers
- Dominant language
- C++
- Stars
- 9.8k
- Forks
- 1.1k
- Avg merge
- 3d 6m
- Merged PRs (30d)
- 278
Description
## Is your feature request related to a problem? Please describe.
Several libcudf functions allocate device memory but do not accept a memory resource. Their callers cannot control those allocations, so the implementations fall back to the current/default resource.
This issue is limited to **adding an MR parameter where none exists and propagating it to allocations and callees**. It does not propose migrating existing MR parameters to `cudf::memory_resources` or classifying allocations as output, temporary, or persistent.
This problem surfaced while working on multi-output JIT transform examples (#23173), and an allocation was unaccounted for.
## Findings
### Public APIs missing an MR parameter
| API | Allocation evidence |
| --- | --- |
| `cudf::contains(column_view, scalar, stream)` | `cpp/src/search/contains_scalar.cu` creates a `column_device_view` and calls the allocating `cudf::detail::count_if`. |
| `cudf::is_sorted(table_view, column_order, null_precedence, stream)` | `cpp/src/sort/is_sorted.cu` builds a row comparator and, for nested input, an `rmm::device_uvector`. |
| `cudf::tables_equal(table_view, table_view, null_equality, stream)` | `cpp/src/table/table_equal.cu` builds a row comparator, allocates an `rmm::device_uvector`, and calls `cudf::detail::reduce`. |
| `cudf::has_nonempty_nulls(column_view, stream)` | `cpp/src/copying/purge_nonempty_nulls.cu` creates a `column_device_view` and calls `cudf::detail::count_if`. |
| `cudf::null_count` and `cudf::batch_null_count` | `cpp/src/bitmask/null_mask.cu::batch_count_set_bits` creates device vectors for bitmask pointers and counts. |
| `cudf::segmented_valid_count` and `cudf::segmented_null_count` | The segmented count implementation allocates device vectors and CUB storage. |
| `cudf::index_of_first_set_bit` | `cpp/src/bitmask/null_mask.cu` allocates a `cudf::detail::device_scalar`. |
| `column_view::null_count(begin, end, stream)` and `has_nulls(begin, end, stream)` | Delegate to the allocating bitmask count implementation. |
| `cudf::hash_join::inner_join_size(left, stream)` and `left_join_size(left, stream)` | `cpp/src/join/hash_join/size_impl.cuh` creates a `row::equality::preprocessed_table` using the current resource. |
`hash_join::full_join_size` and the conditional join-size APIs already accept an MR and are not part of this list.
### Reusable internal helpers missing an MR parameter
| Helper | File | Allocation/caller impact |
| --- | --- | --- |
| `cudf::detail::reduce` | `cpp/include/cudf/detail/algorithms/reduce.cuh` | Allocates a result device scalar and supplies an MR to CUB. Used by table equality and IO. |
| `cudf::detail::reduce_by_key` | Same | Allocates a run-count scalar and CUB storage. Used by groupby, IO, lists, and labeling. |
| `cudf::detail::reduce_by_key_async` | Same | Allocates CUB storage. Used by groupby, JSON, and lists. |
| `cudf::detail::transform_reduce` | Same | Allocates a result scalar and CUB storage. Used heavily by parquet. |
| `cudf::detail::{all_of, any_of, none_of, count_if}` | Same | Forward to `transform_reduce`; callers span copying, IO, joins, quantiles, rolling, search, table, and text. |
| Both `cudf::detail::copy_if` overloads | `cpp/include/cudf/detail/algorithms/copy_if.cuh` | Allocate a selected-count scalar and CUB storage. Callers span copying, IO, joins, lists, stream compaction, strings, and text. |
| Both `cudf::detail::copy_if_async` overloads | Same | Allocate CUB storage. |
| `cudf::detail::labels_to_offsets` | `cpp/include/cudf/detail/labeling/label_segments.cuh` | Allocates two device vectors and calls `reduce_by_key`. |
| `cudf::detail::scatter_to_gather` and `scatter_to_gather_complement` | `cpp/include/cudf/detail/scatter.cuh` | Return newly allocated device vectors and invoke Thrust without caller-provided MR control. |
### Regex helpers missing an MR parameter
Regex extraction is one visible consumer, but the missing resource control is shared by the regex infrastructure.
| Helper | File | Allocation |
| --- | --- | --- |
| `regex_device_builder::create_prog_device` / `detail::reprog_device::create` | `cpp/src/strings/regex/regex_program_impl.h`, `regexec.cpp` | Device regex program in an `rmm::device_uvector`. |
| `regex_device_builder::create_gkprog_device` / `detail::gkprog_device::create` | `cpp/src/strings/regex/regex_program_impl.h`, `gkexec.cpp` | Glushkov device program in an `rmm::device_uvector`. |
| `launch_for_each_kernel` | `cpp/src/strings/regex/utilities.cuh` | Per-launch working memory in an `rmm::device_buffer`. |
| `launch_transform_kernel` | Same | Per-launch working memory in an `rmm::device_buffer`. |
The builders are used by regex `contains`, `count_matches`, `extract`, `extract_single`, `extract_all_record`, regex/back-reference replacement, `findall`, and regex split. The MR should be added to the builders and launch helpers and propagated through every caller. This covers the uncontrollable allocations observed in regex extraction.
## Describe the solution you would like
1. Add a trailing `rmm::device_async_resource_ref mr` to each public API and reusable helper above. Default it only at public API boundaries.
2. Pass `mr` to every RMM allocation, CUB execution environment/storage allocation, Thrust policy, device-view factory, row-preprocessing helper, and nested allocating helper below that entry point.
3. Propagate it through dispatchers, recursive calls, and implementation wrappers so no hidden default/current-resource lookup remains.
4. Add focused tests that make the current resource unusable, invoke the API with an explicit working resource, and verify success and complete deallocation.
5. Update Doxygen for every public signature.
The exact MR type can follow the convention of the surrounding API. This issue does not require deciding whether a later migration should distinguish allocation lifetimes.
### Suggested implementation groups
1. Generic `reduce`/`copy_if` helpers and callers.
2. Null counting and `column_view` range helpers.
3. Regex helpers and callers.
4. `contains` and `has_nonempty_nulls`.
5. Table equality and sorting.
6. Hash-join size APIs.
7. Label/scatter helpers.
### Acceptance criteria
- Every API/helper listed above accepts and propagates an MR to every device allocation below it.
- No listed implementation relies on an RMM default constructor, implicit Thrust/CUB resource, or unparameterized allocating helper.
- Existing source compatibility is preserved with defaults only at public boundaries; detail helpers require an explicit MR.
- Tests validate explicit resource selection while the current resource is unusable.
## Describe alternatives you have considered
- **Continue relying on the current/default resource.** Callers cannot control or validate these allocations.
- **Only patch individual allocation expressions.** Adding MR parameters to shared helpers prevents every consumer from repeating the fallback.
- **No default memory resources and streams in cudf and RMM APIs** Aggressive and extensive breaking change. It will allow us to drop stream and memory resource tests and will eliminate the need for manually tracking which API uses the resources or not. Python bindings can be updated to use specific resources/streams if need be.
Contributor guide
Research direction
Start with cpp/include/cudf/detail/algorithms/reduce.cuh and copy_if.cuh, then trace one suggested group such as generic reductions through the callers listed in the issue. Inspect the allocation, CUB, Thrust, and nested-helper paths before adding focused resource-selection tests. Done means every listed API and helper propagates an explicit MR, preserves public defaults, and passes with the current resource unusable.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- backend-api-design, data-engineering, performance
- Issue type
- Bug
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 25/100