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

[harrism]

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.

[wence-]

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]

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]

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.

[harrism]

🙏

[wence-]

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

[wence-]

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.

[harrism]

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

[harrism]

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.

[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]

Is this temporary / for testing?

Suggested change

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

[harrism]

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]

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

[harrism]

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]

/merge