NVIDIA / spark-rapids-jni

RAPIDS Accelerator JNI For Apache Spark
Apache License 2.0
32 stars 59 forks source link

[FEA] Convert `rmm` resource parameters to `rmm::device_async_resource_ref` #1979

Open ttnghia opened 2 months ago

ttnghia commented 2 months ago

In order to move along with rmm code refactoring, similar to https://github.com/rapidsai/cudf/pull/15507, we need to perform a repository-wide conversion from rmm::mr::device_memory_resource * into rmm::device_async_resource_ref.

This needs to be done ASAP before rmm::mr::device_memory_resource * is removed from rmm.

pmattione-nvidia commented 1 month ago

Merged PR 2011 does much of this work; this conversion has now been performed at all of the kernel interfaces. However the resource allocator in SparkResourceAdaptor.jni still uses device_memory_resource and needs to be updated.

pmattione-nvidia commented 3 weeks ago

Further work is blocked until RMM updates its memory resource classes. The cuDF Java initialize() function ultimately creates the RMM C++ memory resources, but they all still inherit from device_memory_resource. The RMM team says that they still need to figure some things out before they can update them to use the new memory concepts.

In anticipation of the upcoming changes, the following draft PRs have been created, containing all of the work that can be done until then:

cuDF Java Code Draft PR spark-rapids-jni Draft PR

The remaining work that needs to be done before we can update our Java code is being tracked here:

Refactor RMM resources issue