rapidsai / rmm

RAPIDS Memory Manager
https://docs.rapids.ai/api/rmm/stable/
Apache License 2.0
446 stars 188 forks source link

Add resource_ref versions of get/set_current_device_resource #1598

Open harrism opened 1 week ago

harrism commented 1 week ago

Description

Fixes #1597

Depends on #1589

Adds newget_per_device_resource_ref(), set_per_device_resource_ref() and current_device versions of these, intended to replace get_per_device_resource(). The new functions deal in device_async_resource_ref, while the old functions deal in device_memory_resource pointers. Tests are updated to use the new functions.

Note that I have also added reset_per_device_resource_ref() and reset_current_device_resource_ref() which are necessary because previously we implemented the resetting behavior by passing nullptr to set_current_device_resource(), which doesn't work with resource_ref because it can't refer to nullptr (no such thing as a null reference).

Also needed to update some adaptors (statistics, tracking) to take a resource_ref at construction. This is a preview of updating all adaptors to work this way. This was needed because they needed to work with get/set_current_resource_ref().

Checklist