Skip to content

Commit

Permalink
Provide explicit pool sizes and avoid RMM detail APIs (#4086)
Browse files Browse the repository at this point in the history
This PR fixes up cuGraph to avoid usage that will soon be deprecated in RMM.

Depends on rapidsai/rmm#1417

Fixes #4066

Authors:
  - Mark Harris (https://github.com/harrism)

Approvers:
  - Chuck Hastings (https://github.com/ChuckHastings)

URL: #4086
  • Loading branch information
harrism authored Jan 17, 2024
1 parent 8672534 commit eacdf58
Show file tree
Hide file tree
Showing 2 changed files with 6 additions and 7 deletions.
6 changes: 3 additions & 3 deletions cpp/include/cugraph/mtmg/resource_manager.hpp
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* Copyright (c) 2023, NVIDIA CORPORATION.
* Copyright (c) 2023-2024, NVIDIA CORPORATION.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
Expand Down Expand Up @@ -106,9 +106,9 @@ class resource_manager_t {
auto per_device_it = per_device_rmm_resources_.insert(
std::pair{global_rank, std::make_shared<rmm::mr::cuda_memory_resource>()});
#else
auto const [free, total] = rmm::detail::available_device_memory();
auto const [free, total] = rmm::available_device_memory();
auto const min_alloc =
rmm::detail::align_down(std::min(free, total / 6), rmm::detail::CUDA_ALLOCATION_ALIGNMENT);
rmm::align_down(std::min(free, total / 6), rmm::CUDA_ALLOCATION_ALIGNMENT);

auto per_device_it = per_device_rmm_resources_.insert(
std::pair{global_rank,
Expand Down
7 changes: 3 additions & 4 deletions cpp/tests/utilities/base_fixture.hpp
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* Copyright (c) 2020-2023, NVIDIA CORPORATION.
* Copyright (c) 2020-2024, NVIDIA CORPORATION.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
Expand Down Expand Up @@ -73,9 +73,8 @@ inline auto make_pool()
// run more than 2 tests in parallel at the same time. Changes to this value could
// effect the maximum amount of parallel tests, and therefore `tests/CMakeLists.txt`
// `_CUGRAPH_TEST_PERCENT` default value will need to be audited.
auto const [free, total] = rmm::detail::available_device_memory();
auto const min_alloc =
rmm::detail::align_down(std::min(free, total / 6), rmm::detail::CUDA_ALLOCATION_ALIGNMENT);
auto const [free, total] = rmm::available_device_memory();
auto const min_alloc = rmm::align_down(std::min(free, total / 6), rmm::CUDA_ALLOCATION_ALIGNMENT);
return rmm::mr::make_owning_wrapper<rmm::mr::pool_memory_resource>(make_cuda(), min_alloc);
}

Expand Down

0 comments on commit eacdf58

Please sign in to comment.