Skip to content

Commit 24217de

Browse files
bdicegforsyth
bdice
and
gforsyth
authored
Migrate RMM usage to CCCL MR design (rapidsai#22008)
## Summary - Replace `device_memory_resource*` with `device_async_resource_ref` across all C++ headers, sources, benchmarks, examples, and tests. - In Cython `.pxd` declarations, change `device_memory_resource *mr` parameters to `device_async_resource_ref mr` (value type). In `.pyx` files, replace `mr.get_mr()` calls with `mr.c_ref.value()`. - Remove `cudf::set_current_device_resource` (pointer-based) wrapper, keeping only the ref-based `set_current_device_resource_ref`. Update return types of `set/reset_current_device_resource_ref` to `cuda::mr::any_resource<cuda::mr::device_accessible>`. - In `host_memory.cpp`, remove `device_memory_resource` inheritance from `pinned_pool_with_fallback_memory_resource`, remove the forward declaration workaround for `rmm::mr::pool_memory_resource` (no longer templated), and wrap non-copyable state in `shared_ptr` to satisfy the `any_resource` copyability requirement. Part of rapidsai/rmm#2011. Depends on rapidsai/rmm#2361. --------- Co-authored-by: Gil Forsyth <gforsyth@users.noreply.github.com>
1 parent 04b7b21 commit 24217de

147 files changed

Lines changed: 1374 additions & 1196 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎ci/run_cudf_benchmark_smoketests.sh‎

Lines changed: 17 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,25 @@
11
#!/bin/bash
2-
# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION.
2+
# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
33
# SPDX-License-Identifier: Apache-2.0
44

55
set -euo pipefail
66

7-
# Support customizing the ctests' install location
8-
cd "${INSTALL_PREFIX:-${CONDA_PREFIX:-/usr}}/bin/benchmarks/libcudf/";
7+
# Support customizing the benchmarks' install location
8+
# First, try the installed location (CI/conda environments)
9+
installed_benchmark_location="${INSTALL_PREFIX:-${CONDA_PREFIX:-/usr}}/bin/benchmarks/libcudf/"
10+
# Fall back to the build directory (devcontainer environments)
11+
devcontainers_benchmark_location="$(dirname "$(realpath "${BASH_SOURCE[0]}")")/../cpp/build/latest/benchmarks/"
12+
13+
if [[ -d "${installed_benchmark_location}" ]]; then
14+
cd "${installed_benchmark_location}"
15+
elif [[ -d "${devcontainers_benchmark_location}" ]]; then
16+
cd "${devcontainers_benchmark_location}"
17+
else
18+
echo "Error: Benchmark location not found. Searched:" >&2
19+
echo " - ${installed_benchmark_location}" >&2
20+
echo " - ${devcontainers_benchmark_location}" >&2
21+
exit 1
22+
fi
923

1024
# Ensure that benchmarks are runnable
1125
# Run a small nvbench benchmark

‎cpp/benchmarks/common/memory_stats.hpp‎

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* SPDX-FileCopyrightText: Copyright (c) 2019-2025, NVIDIA CORPORATION.
2+
* SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION.
33
* SPDX-License-Identifier: Apache-2.0
44
*/
55

@@ -15,14 +15,12 @@ namespace cudf {
1515
class memory_stats_logger {
1616
public:
1717
memory_stats_logger()
18-
: existing_mr(cudf::get_current_device_resource_ref()),
19-
statistics_mr(
20-
rmm::mr::statistics_resource_adaptor<rmm::device_async_resource_ref>(existing_mr))
18+
: existing_mr(cudf::get_current_device_resource_ref()), statistics_mr(existing_mr)
2119
{
22-
cudf::set_current_device_resource_ref(&statistics_mr);
20+
cudf::set_current_device_resource(statistics_mr);
2321
}
2422

25-
~memory_stats_logger() { cudf::set_current_device_resource_ref(existing_mr); }
23+
~memory_stats_logger() { cudf::set_current_device_resource(existing_mr); }
2624

2725
[[nodiscard]] size_t peak_memory_usage() const noexcept
2826
{
@@ -31,7 +29,7 @@ class memory_stats_logger {
3129

3230
private:
3331
rmm::device_async_resource_ref existing_mr;
34-
rmm::mr::statistics_resource_adaptor<rmm::device_async_resource_ref> statistics_mr;
32+
rmm::mr::statistics_resource_adaptor statistics_mr;
3533
};
3634

3735
} // namespace cudf

‎cpp/benchmarks/fixture/nvbench_fixture.hpp‎

Lines changed: 12 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* SPDX-FileCopyrightText: Copyright (c) 2021-2025, NVIDIA CORPORATION.
2+
* SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION.
33
* SPDX-License-Identifier: Apache-2.0
44
*/
55
#pragma once
@@ -14,10 +14,11 @@
1414
#include <rmm/mr/cuda_async_memory_resource.hpp>
1515
#include <rmm/mr/cuda_memory_resource.hpp>
1616
#include <rmm/mr/managed_memory_resource.hpp>
17-
#include <rmm/mr/owning_wrapper.hpp>
1817
#include <rmm/mr/pinned_host_memory_resource.hpp>
1918
#include <rmm/mr/pool_memory_resource.hpp>
2019

20+
#include <cuda/memory_resource>
21+
2122
#include <string>
2223

2324
namespace cudf {
@@ -34,30 +35,25 @@ static std::string cuio_host_mem_param{
3435
* Initializes the default memory resource to use the RMM pool device resource.
3536
*/
3637
struct nvbench_base_fixture {
37-
inline auto make_cuda() { return std::make_shared<rmm::mr::cuda_memory_resource>(); }
38+
inline auto make_cuda() { return rmm::mr::cuda_memory_resource{}; }
3839

3940
inline auto make_pool()
4041
{
41-
return rmm::mr::make_owning_wrapper<rmm::mr::pool_memory_resource>(
42-
make_cuda(), rmm::percent_of_free_device_memory(50));
42+
return rmm::mr::pool_memory_resource{make_cuda(), rmm::percent_of_free_device_memory(50)};
4343
}
4444

45-
inline auto make_async() { return std::make_shared<rmm::mr::cuda_async_memory_resource>(); }
45+
inline auto make_async() { return rmm::mr::cuda_async_memory_resource{}; }
4646

47-
inline auto make_managed() { return std::make_shared<rmm::mr::managed_memory_resource>(); }
47+
inline auto make_managed() { return rmm::mr::managed_memory_resource{}; }
4848

49-
inline auto make_arena()
50-
{
51-
return rmm::mr::make_owning_wrapper<rmm::mr::arena_memory_resource>(make_cuda());
52-
}
49+
inline auto make_arena() { return rmm::mr::arena_memory_resource{make_cuda()}; }
5350

5451
inline auto make_managed_pool()
5552
{
56-
return rmm::mr::make_owning_wrapper<rmm::mr::pool_memory_resource>(
57-
make_managed(), rmm::percent_of_free_device_memory(50));
53+
return rmm::mr::pool_memory_resource{make_managed(), rmm::percent_of_free_device_memory(50)};
5854
}
5955

60-
inline std::shared_ptr<rmm::mr::device_memory_resource> create_memory_resource(
56+
inline cuda::mr::any_resource<cuda::mr::device_accessible> create_memory_resource(
6157
std::string const& mode)
6258
{
6359
if (mode == "cuda") return make_cuda();
@@ -101,7 +97,7 @@ struct nvbench_base_fixture {
10197
}
10298

10399
mr = create_memory_resource(rmm_mode);
104-
cudf::set_current_device_resource(mr.get());
100+
cudf::set_current_device_resource(mr);
105101
std::cout << "RMM memory resource = " << rmm_mode << "\n";
106102

107103
cudf::set_pinned_memory_resource(create_cuio_host_memory_resource(cuio_host_mode));
@@ -114,7 +110,7 @@ struct nvbench_base_fixture {
114110
cudf::set_pinned_memory_resource(this->make_cuio_host_pinned());
115111
}
116112

117-
std::shared_ptr<rmm::mr::device_memory_resource> mr;
113+
cuda::mr::any_resource<cuda::mr::device_accessible> mr;
118114
std::string rmm_mode{"pool"};
119115

120116
std::string cuio_host_mode{"pinned_pool"};

‎cpp/benchmarks/io/parquet/experimental/hybrid_scan/hybrid_scan_composer.cpp‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -106,8 +106,7 @@ std::vector<cudf::size_type> apply_row_group_filters(
106106
bloom_filter_byte_ranges.size()) {
107107
// Fetch 32-byte aligned bloom filter data buffers from the input file buffer
108108
auto constexpr bloom_filter_alignment = rmm::CUDA_ALLOCATION_ALIGNMENT;
109-
auto aligned_mr = rmm::mr::aligned_resource_adaptor<rmm::mr::device_memory_resource>(
110-
mr, bloom_filter_alignment);
109+
auto aligned_mr = rmm::mr::aligned_resource_adaptor(mr, bloom_filter_alignment);
111110
auto [bloom_filter_buffers, bloom_filter_data, bloom_read_tasks] =
112111
cudf::io::parquet::fetch_byte_ranges_to_device_async(
113112
datasource, bloom_filter_byte_ranges, stream, aligned_mr);

‎cpp/benchmarks/ndsh/utilities.cpp‎

Lines changed: 4 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,6 @@
2222
#include <cudf/utilities/default_stream.hpp>
2323

2424
#include <rmm/mr/managed_memory_resource.hpp>
25-
#include <rmm/mr/owning_wrapper.hpp>
2625
#include <rmm/mr/pool_memory_resource.hpp>
2726

2827
#include <algorithm>
@@ -388,12 +387,6 @@ void write_to_parquet_device_buffer(std::unique_ptr<cudf::table> const& table,
388387
cudf::io::write_parquet(options, stream);
389388
}
390389

391-
inline auto make_managed_pool()
392-
{
393-
return rmm::mr::make_owning_wrapper<rmm::mr::pool_memory_resource>(
394-
std::make_shared<rmm::mr::managed_memory_resource>(), rmm::percent_of_free_device_memory(50));
395-
}
396-
397390
void generate_parquet_data_sources(double scale_factor,
398391
std::vector<std::string> const& table_names,
399392
std::unordered_map<std::string, cuio_source_sink_pair>& sources)
@@ -403,9 +396,9 @@ void generate_parquet_data_sources(double scale_factor,
403396
// Set the memory resource to the managed pool
404397
auto old_mr = cudf::get_current_device_resource_ref();
405398
// TODO: if old_mr is already managed pool or managed, don't create new one.
406-
using managed_pool_mr_t = decltype(make_managed_pool());
407-
managed_pool_mr_t managed_pool_mr = make_managed_pool();
408-
cudf::set_current_device_resource_ref(managed_pool_mr.get());
399+
rmm::mr::pool_memory_resource managed_pool_mr{rmm::mr::managed_memory_resource{},
400+
rmm::percent_of_free_device_memory(50)};
401+
cudf::set_current_device_resource(managed_pool_mr);
409402
// drawback: if already pool takes 50% of free memory, we are left with 50% of 50% of free memory
410403

411404
std::unordered_set<std::string> const requested_table_names = [&table_names]() {
@@ -469,5 +462,5 @@ void generate_parquet_data_sources(double scale_factor,
469462
}
470463

471464
// Restore the original memory resource
472-
cudf::set_current_device_resource_ref(old_mr);
465+
cudf::set_current_device_resource(old_mr);
473466
}

‎cpp/examples/basic/src/process_csv.cpp‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -67,18 +67,18 @@ int main(int argc, char** argv)
6767
{
6868
// Construct a CUDA memory resource using RAPIDS Memory Manager (RMM)
6969
// This is the default memory resource for libcudf for allocating device memory.
70-
rmm::mr::cuda_memory_resource cuda_mr{};
7170
// Construct a memory pool using the CUDA memory resource
7271
// Using a memory pool for device memory allocations is important for good performance in libcudf.
7372
// The pool defaults to allocating half of the available GPU memory.
74-
rmm::mr::pool_memory_resource mr{&cuda_mr, rmm::percent_of_free_device_memory(50)};
73+
rmm::mr::pool_memory_resource mr{rmm::mr::cuda_memory_resource{},
74+
rmm::percent_of_free_device_memory(50)};
7575

7676
// Set the pool resource to be used by default for all device memory allocations
7777
// Note: It is the user's responsibility to ensure the `mr` object stays alive for the duration of
7878
// it being set as the default
7979
// Also, call this before the first libcudf API call to ensure all data is allocated by the same
8080
// memory resource.
81-
cudf::set_current_device_resource(&mr);
81+
cudf::set_current_device_resource(mr);
8282

8383
// Read data
8484
auto stock_table_with_metadata = read_csv("4stock_5day.csv");

‎cpp/examples/billion_rows/brc.cpp‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION.
2+
* SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
33
* SPDX-License-Identifier: Apache-2.0
44
*/
55
#include "common.hpp"
@@ -34,9 +34,8 @@ int main(int argc, char const** argv)
3434

3535
auto const mr_name = std::string("pool");
3636
auto resource = create_memory_resource(mr_name);
37-
auto stats_mr =
38-
rmm::mr::statistics_resource_adaptor<rmm::mr::device_memory_resource>(resource.get());
39-
rmm::mr::set_current_device_resource(&stats_mr);
37+
auto stats_mr = rmm::mr::statistics_resource_adaptor{resource};
38+
rmm::mr::set_current_device_resource(stats_mr);
4039
auto stream = cudf::get_default_stream();
4140

4241
auto start = std::chrono::steady_clock::now();

‎cpp/examples/billion_rows/brc_chunks.cpp‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION.
2+
* SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
33
* SPDX-License-Identifier: Apache-2.0
44
*/
55
#include "common.hpp"
@@ -56,9 +56,8 @@ int main(int argc, char const** argv)
5656

5757
auto const mr_name = std::string("pool");
5858
auto resource = create_memory_resource(mr_name);
59-
auto stats_mr =
60-
rmm::mr::statistics_resource_adaptor<rmm::mr::device_memory_resource>(resource.get());
61-
rmm::mr::set_current_device_resource(&stats_mr);
59+
auto stats_mr = rmm::mr::statistics_resource_adaptor{resource};
60+
rmm::mr::set_current_device_resource(stats_mr);
6261
auto stream = cudf::get_default_stream();
6362

6463
std::filesystem::path p = input_file;

‎cpp/examples/billion_rows/brc_pipeline.cpp‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION.
2+
* SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
33
* SPDX-License-Identifier: Apache-2.0
44
*/
55
#include "common.hpp"
@@ -104,9 +104,8 @@ int main(int argc, char const** argv)
104104

105105
auto const mr_name = std::string("pool");
106106
auto resource = create_memory_resource(mr_name);
107-
auto stats_mr =
108-
rmm::mr::statistics_resource_adaptor<rmm::mr::device_memory_resource>(resource.get());
109-
rmm::mr::set_current_device_resource(&stats_mr);
107+
auto stats_mr = rmm::mr::statistics_resource_adaptor{resource};
108+
rmm::mr::set_current_device_resource(stats_mr);
110109
auto stream = cudf::get_default_stream();
111110

112111
std::filesystem::path p = input_file;

‎cpp/examples/billion_rows/common.hpp‎

Lines changed: 10 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -1,36 +1,26 @@
11
/*
2-
* SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION.
2+
* SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
33
* SPDX-License-Identifier: Apache-2.0
44
*/
55
#pragma once
66

77
#include <rmm/cuda_device.hpp>
88
#include <rmm/mr/cuda_memory_resource.hpp>
9-
#include <rmm/mr/device_memory_resource.hpp>
10-
#include <rmm/mr/owning_wrapper.hpp>
119
#include <rmm/mr/pool_memory_resource.hpp>
10+
#include <rmm/resource_ref.hpp>
1211

13-
#include <string>
12+
#include <cuda/memory_resource>
1413

15-
/**
16-
* @brief Create CUDA memory resource
17-
*/
18-
auto make_cuda_mr() { return std::make_shared<rmm::mr::cuda_memory_resource>(); }
19-
20-
/**
21-
* @brief Create a pool device memory resource
22-
*/
23-
auto make_pool_mr()
24-
{
25-
return rmm::mr::make_owning_wrapper<rmm::mr::pool_memory_resource>(
26-
make_cuda_mr(), rmm::percent_of_free_device_memory(50));
27-
}
14+
#include <string>
2815

2916
/**
3017
* @brief Create memory resource for libcudf functions
3118
*/
32-
std::shared_ptr<rmm::mr::device_memory_resource> create_memory_resource(std::string const& name)
19+
cuda::mr::any_resource<cuda::mr::device_accessible> create_memory_resource(std::string const& name)
3320
{
34-
if (name == "pool") { return make_pool_mr(); }
35-
return make_cuda_mr();
21+
rmm::mr::cuda_memory_resource cuda_mr{};
22+
if (name == "pool") {
23+
return rmm::mr::pool_memory_resource{cuda_mr, rmm::percent_of_free_device_memory(50)};
24+
}
25+
return cuda_mr;
3626
}

0 commit comments

Comments
 (0)