Make BufferResource shared_ptr-managed, device_mr() now keeps BR alive

[madsbk]

Fixes the C++ side of #641 by making BufferResource shared-pointer-managed.

Lifetime safety is provided by OwningResourceAdaptor, a CCCL-compatible adaptor that wraps the device memory resource and carries a std::weak_ptr<BufferResource> to its parent.

The adaptor stored inside BufferResource itself holds only the weak reference, so there is no reference cycle. When CCCL deep-copies the adaptor, which happens when downstream code promotes the ref into an owning cuda::mr::any_resource<...> (for example rmm::device_buffer::_mr), the adaptor copy constructor promotes the weak reference to a std::shared_ptr<BufferResource>.

As a result, the copy stored inside the owning any_resource keeps the entire BufferResource alive, including its stream pool, spill manager, and statistics, for as long as the downstream object exists.

If the BufferResource has already been destroyed at copy time, the adaptor copy constructor throws std::bad_weak_ptr instead of creating a dangling reference.

The caller must still ensure that BufferResource outlives any direct use of a bare device_async_resource_ref. For example:

auto mr = br->device_mr();
mr.allocate(...);  // UB if br is already destroyed

This matches the existing RMM contract for non-owning resource refs. The new safety guarantees only apply once the ref has been captured into an owning container.

Context

The original sketch in #641 proposed making BufferResource itself a CCCL device memory resource, but that conflates roles. device_mr() is the actual device memory resource, while host_mr() and pinned_mr() expose separate host and pinned resources. What matters is that all three keep the parent BufferResource alive.

---

This PR only changes BufferResource::device_mr(). Extending the same ownership semantics to host_mr() and pinned_mr() is tracked separately in #1070.

Python-side lifetime gaps are tracked separately in #1074. Python-created BufferResources still borrow their stream pool, device MR, pinned MR, and statistics from Python-owned objects, so the new C++ ownership model does not yet fully protect those cases.

[nirandaperera]

@madsbk Let's back up a bit. Prior to #1061, we could have passed a plain device mr or a RmmResourceAdapter wrapped device mr to BufferResource. If former was passed, reservations were unlimited, and latter was explicitly passed in from_options method. So, there was a notion that, if I want to limit reservations, I should have a RmmResourceAdapter in my context. Callers would have likely used that mr for external operations that needed a memory resource (rmm, cudf etc)
Eg:

auto cuda_mr ... 
RmmResourceAdapter mr(cuda_mr);
auto br = BufferResource::from_options(mr, ...); 
auto stream_ref = br->stream_pool().get_stream();
...
rmm::device_buffer some_buf(..., mr, stream_ref); // issue here, stream/br could go out of scope without some_buf knowing 

But now, RmmResourceAdapter is internal to br. So, then in the call sites, the callers should explicitly take the adapter out from the br after construction.

auto cuda_mr ... 
auto br = BufferResource::from_options(cuda_mr, ...); 
auto mr = br->device_mr(); // use this everywhere from here on!!!! 
... 

It is very easy to miss this last line, and continue using cuda_mr everywhere else. So, we are loosing track of allocations outside of rapidsmpf.

So, my suggestion is, before fixing the lifetime guarantees, we need to accept a RmmResourceAdapter to br, instead of an any_device_resource.

[nirandaperera]

@madsbk Now, regarding ownership guarantees, we anyway need to store a back-reference of buffer resource to the provided device_mr. I agree with managing buffer resource from a shared_ptr. What if we attach an OwningResourceAdaptor after constructing the shared_ptr.
eg:

class BufferResource{
private:
 OwningResourceAdapter<RmmResourceAdapter, shared_ptr<BufferResource>> mr_owner_;

 BufferResource(RmmResourceAdapter mr, ...): mr_owner_{move(mr), nullptr},...{}

public:
shared_ptr create(RmmResourceAdapter mr, ...){
 auto ret = make_shared(move(mr), ...); // ret doesnt have br tied yet! 
 ret->mr_owner_->set_backref(ret); // hacky! but now, we can safely return ref/ any_resource without any compromise. 
 return ret;
}

any_device_resource device_mr() { return mr_owner_;}
device_async_resource_ref device_mr_ref() {return mr_owner_;}
}

[nirandaperera]

Or else, we can use a shared_ptr<void> with a custom deleter to reset the buffer resource shared_ptr.

class RmmResourceAdapter{
private:
 shared_ptr<void> owning{}; 

public: 
// might need templates 
void set_owned(shared_ptr<void> needs_owned){
 owning = move(needs_owned);
}
};

class BufferResource{
private:
 RmmResourceAdapter mr_owner_;

 BufferResource(RmmResourceAdapter mr, ...): mr_owner_{move(mr), nullptr},...{}

public:
shared_ptr create(RmmResourceAdapter mr, ...){
 auto ret = make_shared(move(mr), ...); // ret doesnt have br tied yet! 
 ret->mr_owner_->set_owned(std::shared_ptr<void>(nullptr, [ret](void*) {
    }));
 return ret;
}

any_device_resource device_mr() { return mr_owner_;}
device_async_resource_ref device_mr_ref() {return mr_owner_;}
}

[wence-]

Small changes I think, but otherwise I think this looks good.

[madsbk]

@madsbk Let's back up a bit. Prior to #1061, we could have passed a plain device mr or a RmmResourceAdapter wrapped device mr to BufferResource. If former was passed, reservations were unlimited, and latter was explicitly passed in from_options method. So, there was a notion that, if I want to limit reservations, I should have a RmmResourceAdapter in my context. Callers would have likely used that mr for external operations that needed a memory resource (rmm, cudf etc)
...
It is very easy to miss this last line, and continue using cuda_mr everywhere else. So, we are loosing track of allocations outside of rapidsmpf.

So, my suggestion is, before fixing the lifetime guarantees, we need to accept a RmmResourceAdapter to br, instead of an any_device_resource.

That's a valid API design concern, but I think it's orthogonal to the lifetime issue this PR is addressing. I've updated the docs to make the tracking semantics explicit: f2663dd.

I'm also not convinced we want to preserve the old semantic. My intuition is that we do want to force users through device_mr() so allocations become visible to BufferResource accounting, limits, and statistics.

Let's discuss that separately.

[pentschev]

Left a couple docstrings improvement suggestions, otherwise LGTM.

[pentschev]

This appears to overstate the current Python implementation. BufferResource.device_mr still returns self._device_mr, which is the original DeviceMemoryResource passed to the constructor, not the internal tracked adaptor. See the assignment at lines 212-214 and the property at lines 294-303.

Since Python-side lifetime/resource exposure is deferred to #1074, can we avoid documenting br.device_mr as the tracked wrapper here? Either this property should actually expose the tracked adaptor, or this Python doc should keep the weaker wording that the C++ BufferResource wraps the MR internally and that allocations made through the original Python MR are not tracked.

Same issue in line 119: tells Python users to use BufferResource.device_mr for tracked allocations, but the property currently returns the original unwrapped Python MR. This should be softened unless this PR also changes the property to expose the tracked adaptor.

[madsbk]

Very good point. It has been fixed in 19f50ef.
@pentschev and @wence- can you take a look before we merge?

[pentschev]

One final nit, otherwise LGTM. Thanks Mads!

[nirandaperera]

I wanted to float the idea of a mixin here. But dont want to block because @vyasr is waiting for this.
So, approving.

[nirandaperera]

@madsbk I feel like making this implement Resource concept is an overkill. How about converting this to a mixin, BackRefMixin?

template<typename BackRef>
class BackRefMixin{

  void set_backref(std::weak_ptr<BackRef> backref){...}
// might need the equality, copy constructors,.  overload
  private:
    std::weak_ptr<BackRef> weak_{nullptr};
    std::shared_ptr<BackRef> strong_{nullptr};
}

That way, we can mix it in with RmmResourceAdaptor or create a simple class,
RmmResourceAdapterWithBR: RmmResourceAdaptor, BackRefMixin

This will remove the Resource boilerplate

[madsbk]

Interesting idea, but I think it becomes cumbersome once we need to handle other resources like the pinned and host MRs. You would likely end up needing a separate XWithBR glue class for each adaptor/resource combination.

I also like that the current adaptor is generic enough that we could potentially upstream it to RMM at some point, whereas a mixin-based approach feels a bit more tied to our specific BufferResource use case.

[nirandaperera]

On the other hand, a mixin can be used for any class. We could easily add that to Pinned mr and host mr, and well as stream pool (will need a rapidsmpf stream pool in the end). Let me sketch something up, now that this PR is merged.

[nirandaperera]

@madsbk This is what I had in mind #1078
I feel its more generic.

[madsbk]

/merge