Update RMM adaptors, containers and tests to use get/set_current_device_resource_ref() #1661

harrism · 2024-08-22T04:17:02Z

Description

Closes #1660.

This adds a constructor to each MR adaptor to take a resource_ref rather than an Upstream*. It also updates RMM to use get_current_device_resource_ref() everywhere: in containers, in tests, in adaptors, Thrust allocator, polymorphic allocator, execution_policy, etc.

Importantly, this PR also modifies set_current_device_resource() to basically call set_current_device_resource_ref(). This is necessary, because while RMM C++ uses get_current_device_resource_ref() everywhere, the Python API still uses the raw pointer API set_current_device_resource(). So we need the latter to update the state for the former. This is a temporary bootstrap to help with the refactoring.

Checklist

I am familiar with the Contributing Guidelines.
New or existing tests cover these changes.
The documentation is up to date with these changes.

Co-authored-by: Lawrence Mitchell <[email protected]>

Co-authored-by: Bradley Dice <[email protected]>

…ism/rmm into fea-get_current_device_resource_ref

…tic.

…source_ref state as well.

wence- · 2024-08-23T14:57:35Z

On the Python side, we need someone to own concrete objects (we can't just provide classes that wrap device_async_resource_ref since that is non-owning).

Adapters also need a way to keep their upstream alive.

We are removing the device_memory_resource inheritance on the C++ side, but I think until we have an owning any_resource available, we probably need to keep something on the Cython side.

Here is a very slide-ware sketch approach for discussion:

cdef extern from * nogil:
    """
    struct hack {
        std::function<rmm::device_async_resource_ref> _to_ref;
        hack() = default; // unsafe, but this is all hacky anyway
        rmm::device_async_resource_ref to_ref() { return _to_ref(); }
        template<typename T>
        hack(std::unique_ptr<T> obj) : _to_ref{[&obj]() -> rmm::device_async_resource_ref { return *obj; }} {};
    };
    """
       
cdef class DeviceMemoryResource:
   cdef hack obj;
   cdef device_async_resource_ref as_ref(self):
       return self.obj.to_ref()

cdef class CudaMemoryResource(DeviceMemoryResource):
   def __cinit__(self):
       self.obj = hack(move(make_unique[cuda_memory_resource]()))
   

cdef class PoolMemoryResource(DeviceMemoryResource):
   cdef DeviceMemoryResource upstream
   def __cinit__(self, DeviceMemoryResource upstream):
      self.upstream = upstream
      self.obj = hack(move(make_unique[pool_memory_resource](upstream.to_ref())))

This is effectively working around not (yet) having any_resource available by doing a minimal type-erased owning wrapper.

harrism · 2024-08-26T21:50:20Z

On the Python side, we need someone to own concrete objects (we can't just provide classes that wrap device_async_resource_ref since that is non-owning).

Probably should have this discussion somewhere else? (Opened #1662 in discussions) I was planning for this to be a C++-only PR, replacing #1634.

abellina · 2024-08-27T18:40:14Z

include/rmm/mr/device/per_device_resource.hpp

+  // Note: even though set_per_device_resource() and set_per_device_resource_ref() are not
+  // interchangeable, we call the latter from the former to maintain resource_ref
+  // state consistent with the resource pointer state. This is necessary because the
+  // Python API still uses the raw pointer API. Once the Python API is updated to use


cudf JNI also uses set_per_device_resource and get_current_device_resource today, so we'll also be needing to change our code to use _ref.

include/rmm/mr/device/aligned_resource_adaptor.hpp

wence-

Looks good to me, with a question about the (now-redundant) class template parameter on all the adaptor MRs.

wence- · 2024-08-28T09:51:16Z

include/rmm/mr/device/aligned_resource_adaptor.hpp

+   * @param alignment_threshold Only allocations with a size larger than or equal to this threshold
+   * are aligned.
+   */
+  explicit aligned_resource_adaptor(Upstream* upstream,


I'm not sure how to handle this migration without making a breaking change, but this class no longer needs to be a templated class, but rather this constructor should be templated.

And hence, question: does the constructor even need to exist, if there is transparent conversion from Upstream * to device_async_resource_ref?

That is, what doesn't work if the only constructor is:

aligned_resource_adaptor(device_async_resource_ref upstream, ...);

Applies mutatis mutandis to the other adaptor MR changes as well, I think.

Right, the plan is to remove the template parameter and the Upstream* constructors, once we add the resource_ref constructors and convert all of RAPIDS to use them. But we can't do it yet.

See #1457

include/rmm/mr/device/aligned_resource_adaptor.hpp

include/rmm/mr/device/detail/arena.hpp

tests/mr/device/arena_mr_tests.cpp

…tructors.

…nversion to refs

harrism · 2024-08-29T01:28:02Z

Let's not merge until I can do some more downstream testing, including JNI.

bdice

I got about 2/3 through the changes. I’ll finish tomorrow.

benchmarks/device_uvector/device_uvector_bench.cu

include/rmm/mr/device/aligned_resource_adaptor.hpp

include/rmm/mr/device/pool_memory_resource.hpp

include/rmm/resource_ref.hpp

bdice

I have a couple more comments but overall I think this is a good direction and this PR is of sufficient quality that I feel comfortable approving it. Thanks for answering my questions.

bdice · 2024-08-29T22:57:16Z

tests/CMakeLists.txt

@@ -40,6 +40,7 @@ function(ConfigureTestInternal TEST_NAME)
                             PUBLIC "SPDLOG_ACTIVE_LEVEL=SPDLOG_LEVEL_${RMM_LOGGING_LEVEL}")
  target_compile_options(${TEST_NAME} PUBLIC $<$<COMPILE_LANG_AND_ID:CXX,GNU,Clang>:-Wall -Werror
                                             -Wno-error=deprecated-declarations>)
+  target_compile_options(${TEST_NAME} PUBLIC "$<$<CONFIG:Debug>:-O0>")


Is this temporary / for testing?

Suggested change

target_compile_options(${TEST_NAME} PUBLIC "$<$<CONFIG:Debug>:-O0>")

I would like it to stay, at least until this fixes the problem: . It only applies to Debug builds. rapidsai/rapids-cmake#634 (comment)

bdice · 2024-08-29T23:02:36Z

tests/mr/host/mr_ref_tests.cpp

@@ -233,14 +233,17 @@ TYPED_TEST(MRRefTest, UnsupportedAlignmentTest)
  for (std::size_t num_trials = 0; num_trials < NUM_TRIALS; ++num_trials) {
    for (std::size_t alignment = MinTestedAlignment; alignment <= MaxTestedAlignment;
         alignment *= TestedAlignmentMultiplier) {
+#ifdef NDEBUG


Maybe this is related to the CMakeLists.txt change I noted above? What's the rationale for this change?

In debug mode, the deallocation will assert, crashing the test. I was debugging when getting this PR working, and discovered this.

CI does not test our debug builds, so these things creep in.

harrism · 2024-09-09T23:39:37Z

/merge

Merge after rapidsai/rmm#1661 Creates and uses CUDF internal wrappers around RMM `current_device_resource` functions. I've marked this PR as breaking because it breaks the ABI, however the API is compatible. For reviewers, the most substantial additions are in the new file `<cudf/utilities/memory_resource.hpp>`, and in the `DEVELOPER_GUIDE.md` and `*.rst` docs. The rest are all replacements of an include and all calls to `rmm::get_current_device_resource()` with `cudf::get_current_device_resource_ref()`. Closes #16676 Authors: - Mark Harris (https://github.com/harrism) Approvers: - Nghia Truong (https://github.com/ttnghia) - GALI PREM SAGAR (https://github.com/galipremsagar) - https://github.com/nvdbaranec - David Wendt (https://github.com/davidwendt) URL: #16679

harrism and others added 13 commits August 20, 2024 05:10

Add _ref versions of get/set_current_device_resource and friends.

a9555c5

doc: explain that set_current_device_resource is in transition

d893698

Update/add tests for get/set_current_device_resource_ref

b068ded

Apply suggestions from code review

c7d443d

Co-authored-by: Lawrence Mitchell <[email protected]>

RMM_EXPORT ref_map_lock()

3cd6917

Co-authored-by: Lawrence Mitchell <[email protected]>

Apply suggestions from code review

adbf7f6

Co-authored-by: Bradley Dice <[email protected]>

Merge branch 'fea-get_current_device_resource_ref' of github.com:harr…

db53c65

…ism/rmm into fea-get_current_device_resource_ref

Explicit RMM_EXPORT on all functions that return a function-local sta…

b5a831a

…tic.

Doc: ref and ptr versions are not interchangeable.

812d08d

Use resource_ref in MR adaptors and related tests.

1c43d63

Use resource_ref in torch_allocator and polymorphic_allocator

bd70c42

Make the raw pointer version of set_per_device_resource update the re…

00bd3ac

…source_ref state as well.

Update containers and tests to use get/set_current_device_resource_ref()

d7df60e

github-actions bot added Python Related to RMM Python API cpp Pertains to C++ code labels Aug 22, 2024

abellina reviewed Aug 27, 2024

View reviewed changes

This was referenced Aug 27, 2024

[FEA][JNI] Move to RMM apis get_current_device_resource_ref and set_current_device_resource_ref rapidsai/cudf#16673

Open

[FEA] Move to RMM apis get_current_device_resource_ref NVIDIA/spark-rapids-jni#2358

Closed

harrism added 5 commits August 28, 2024 01:24

Merge branch 'branch-24.10' into fea-resource_ref-upstream

81ba5a0

Remove an utterance of device_memory_resource* in test.

8d2bc62

Remove another unnecessary device_memory_resource*

3d2ef08

Copyrights

13172c5

Merge branch 'branch-24.10' into fea-resource_ref-upstream

efee997

harrism added feature request New feature or request non-breaking Non-breaking change labels Aug 28, 2024

harrism marked this pull request as ready for review August 28, 2024 08:18

harrism requested review from a team as code owners August 28, 2024 08:18

harrism requested a review from wence- August 28, 2024 08:18

miscco approved these changes Aug 28, 2024

View reviewed changes

include/rmm/mr/device/aligned_resource_adaptor.hpp Outdated Show resolved Hide resolved

wence- approved these changes Aug 28, 2024

View reviewed changes

rongou approved these changes Aug 28, 2024

View reviewed changes

harrism added 3 commits August 29, 2024 01:20

Add to_device_async_resource_ref_checked() to simplify adaptor cons…

4884f53

…tructors.

Use to_device_async_resource_ref_checked() and more adaptor tidy / co…

1bcb81a

…nversion to refs

Make test debug builds use -O0

56e5eaa

harrism requested a review from a team as a code owner August 29, 2024 01:25

Fix quoted include.

545323b

github-actions bot added the CMake label Aug 29, 2024

harrism added the 5 - DO NOT MERGE Hold off on merging; see PR for details label Aug 29, 2024

cmake style

5598840

bdice reviewed Aug 29, 2024

View reviewed changes

bdice approved these changes Aug 29, 2024

View reviewed changes

wence- mentioned this pull request Sep 2, 2024

Fallback resource adaptor #1665

Draft

4 tasks

harrism added 4 commits September 4, 2024 07:43

Don't take the address when passing MR to resource_ref.

c1d7b0d

Fix yoda code.

4c11382

Use resource_ref upstreams in binning_mr

6d2f09d

Remove incorrect @throws documentation

a367af6

harrism removed the 5 - DO NOT MERGE Hold off on merging; see PR for details label Sep 5, 2024

harrism mentioned this pull request Sep 5, 2024

Use RMM get_current_device_resource_ref() rapidsai/cuspatial#1454

Open

3 tasks

rapids-bot bot merged commit 6729def into rapidsai:branch-24.10 Sep 9, 2024
49 checks passed

harrism mentioned this pull request Sep 9, 2024

Switch rmm::get_current_device_resource to use references via rmm::get_current_device_resource_ref NVIDIA/spark-rapids-jni#2372

Merged

harrism mentioned this pull request Sep 12, 2024

Rewrite rmm::device_check_resource_adaptor to take a cuda::mr::async_resource_ref instead of a device_memory_resource* #1541

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Update RMM adaptors, containers and tests to use get/set_current_device_resource_ref() #1661

Update RMM adaptors, containers and tests to use get/set_current_device_resource_ref() #1661

harrism commented Aug 22, 2024

wence- commented Aug 23, 2024

harrism commented Aug 26, 2024 •

edited

Loading

abellina Aug 27, 2024

harrism Aug 27, 2024

wence- left a comment

wence- Aug 28, 2024

harrism Aug 29, 2024

harrism commented Aug 29, 2024

bdice left a comment

bdice left a comment

bdice Aug 29, 2024

harrism Sep 2, 2024

bdice Aug 29, 2024

harrism Sep 2, 2024 •

edited

Loading

harrism commented Sep 9, 2024

Update RMM adaptors, containers and tests to use get/set_current_device_resource_ref() #1661

Update RMM adaptors, containers and tests to use get/set_current_device_resource_ref() #1661

Conversation

harrism commented Aug 22, 2024

Description

Checklist

wence- commented Aug 23, 2024

harrism commented Aug 26, 2024 • edited Loading

abellina Aug 27, 2024

Choose a reason for hiding this comment

harrism Aug 27, 2024

Choose a reason for hiding this comment

wence- left a comment

Choose a reason for hiding this comment

wence- Aug 28, 2024

Choose a reason for hiding this comment

harrism Aug 29, 2024

Choose a reason for hiding this comment

harrism commented Aug 29, 2024

bdice left a comment

Choose a reason for hiding this comment

bdice left a comment

Choose a reason for hiding this comment

bdice Aug 29, 2024

Choose a reason for hiding this comment

harrism Sep 2, 2024

Choose a reason for hiding this comment

bdice Aug 29, 2024

Choose a reason for hiding this comment

harrism Sep 2, 2024 • edited Loading

Choose a reason for hiding this comment

harrism commented Sep 9, 2024

harrism commented Aug 26, 2024 •

edited

Loading

harrism Sep 2, 2024 •

edited

Loading