Migrate RMM usage to CCCL MR design (#7951)

## Summary
- Migrate all raw RMM `allocate`/`deallocate` calls to the new CCCL 3-argument API that requires explicit alignment
- Replace removed `rmm.librmm.per_device_resource` Cython import with `rmm.pylibrmm.memory_resource` and use `make_any_device_resource` to obtain the resource for `device_buffer` construction

Depends on rapidsai/rmm#2361.
Depends on rapidsai/ucxx#636.
Depends on rapidsai/raft#2996.
Depends on rapidsai/cuvs#1990.

## Changes
- **`cpp/src/genetic/genetic.cu`**: Add explicit `alignof(node)` / `alignof(program)` to all `allocate` and `deallocate` calls in `parallel_evolve` and `symFit`; fix deallocation bug in `parallel_evolve` where `h_nextprogs[i].len` was incorrectly used instead of `tmp.len` to compute the buffer size being freed
- **`cpp/examples/symreg/symreg_example.cpp`**: Use `params.population_size * sizeof(cg::program)` and `alignof(cg::program)` for `allocate`/`deallocate` calls, fixing incorrect byte-size computation; remove unused `<rmm/aligned.hpp>` include
- **`cpp/tests/sg/genetic/evolution_test.cu`**: Add alignment arguments to allocate/deallocate in `SymReg` test
- **`cpp/tests/sg/genetic/program_test.cu`**: Add alignment arguments to `SetUp`/`TearDown` allocate/deallocate calls
- **`python/cuml/cuml/manifold/umap/umap.pyx`**: Replace `get_current_device_resource()` with `make_any_device_resource(get_current_device_resource().get_mr())` for `device_buffer` construction

Authors:
  - Bradley Dice (https://github.com/bdice)

Approvers:
  - Simon Adorf (https://github.com/csadorf)
  - Divye Gala (https://github.com/divyegala)
  - Victor Lafargue (https://github.com/viclafargue)

URL: #7951

Author: bdice

cpp/examples/symreg/symreg_example.cpp

@@ -235,7 +235,8 @@ int main(int argc, char* argv[])
 
   // Initialize AST
   auto curr_mr = rmm::mr::get_current_device_resource_ref();
-  d_finalprogs = static_cast<cg::program_t>(curr_mr.allocate(stream, params.population_size));
+  d_finalprogs = static_cast<cg::program_t>(
+    curr_mr.allocate(stream, params.population_size * sizeof(cg::program), alignof(cg::program)));
 
   std::vector<std::vector<cg::program>> history;
   history.reserve(params.generations);
@@ -327,7 +328,8 @@ int main(int argc, char* argv[])
 
   /* ======================= Reset data ======================= */
 
-  curr_mr.deallocate(stream, d_finalprogs, params.population_size);
+  curr_mr.deallocate(
+    stream, d_finalprogs, params.population_size * sizeof(cg::program), alignof(cg::program));
   CUDA_RT_CALL(cudaEventDestroy(start));
   CUDA_RT_CALL(cudaEventDestroy(stop));
   return 0;

cpp/src/genetic/genetic.cu

@@ -221,15 +221,15 @@ void parallel_evolve(const raft::handle_t& h,
 
     // Set current generation device nodes
     tmp.nodes = (node*)rmm::mr::get_current_device_resource_ref().allocate(
-      stream, h_nextprogs[i].len * sizeof(node));
+      stream, h_nextprogs[i].len * sizeof(node), alignof(node));
     raft::copy(tmp.nodes, h_nextprogs[i].nodes, h_nextprogs[i].len, stream);
     raft::copy(d_nextprogs + i, &tmp, 1, stream);
 
     if (generation > 1) {
       // Free device memory allocated to program nodes in previous generation
       raft::copy(&tmp, d_oldprogs + i, 1, stream);
       rmm::mr::get_current_device_resource_ref().deallocate(
-        stream, tmp.nodes, h_nextprogs[i].len * sizeof(node));
+        stream, tmp.nodes, tmp.len * sizeof(node), alignof(node));
     }
 
     tmp.nodes = nullptr;
@@ -399,7 +399,7 @@ void symFit(const raft::handle_t& handle,
 
   program_t d_currprogs;  // pointer to current programs
   d_currprogs = (program_t)rmm::mr::get_current_device_resource_ref().allocate(
-    stream, params.population_size * sizeof(program));
+    stream, params.population_size * sizeof(program), alignof(program));
   program_t d_nextprogs = final_progs;  // Reuse memory already allocated for final_progs
   final_progs           = nullptr;
 
@@ -481,7 +481,7 @@ void symFit(const raft::handle_t& handle,
 
   // Deallocate the previous generation device memory
   rmm::mr::get_current_device_resource_ref().deallocate(
-    stream, d_nextprogs, params.population_size * sizeof(program));
+    stream, d_nextprogs, params.population_size * sizeof(program), alignof(program));
   d_currprogs = nullptr;
   d_nextprogs = nullptr;
 }

cpp/tests/sg/genetic/evolution_test.cu

@@ -259,7 +259,7 @@ TEST_F(GeneticEvolutionTest, SymReg)
   MLCommon::CompareApprox<float> compApprox(tolerance);
   program_t final_progs;
   final_progs = (program_t)rmm::mr::get_current_device_resource_ref().allocate(
-    stream, hyper_params.population_size * sizeof(program));
+    stream, hyper_params.population_size * sizeof(program), alignof(program));
   std::vector<std::vector<program>> history;
   history.reserve(hyper_params.generations);
 
@@ -327,12 +327,12 @@ TEST_F(GeneticEvolutionTest, SymReg)
     program tmp = program();
     raft::copy(&tmp, final_progs + i, 1, stream);
     rmm::mr::get_current_device_resource_ref().deallocate(
-      stream, tmp.nodes, tmp.len * sizeof(node));
+      stream, tmp.nodes, tmp.len * sizeof(node), alignof(node));
     tmp.nodes = nullptr;
   }
   // deallocate the final programs from device memory
   rmm::mr::get_current_device_resource_ref().deallocate(
-    stream, final_progs, hyper_params.population_size * sizeof(program));
+    stream, final_progs, hyper_params.population_size * sizeof(program), alignof(program));
 
   ASSERT_TRUE(compApprox(history[n_gen - 1][best_idx].raw_fitness_, 0.0036f));
   std::cout << "Some Predicted test values:" << std::endl;

cpp/tests/sg/genetic/program_test.cu

@@ -90,10 +90,12 @@ class GeneticProgramTest : public ::testing::Test {
     d_lY.resize(250, stream);
     d_lunitW.resize(250, stream);
     d_lW.resize(250, stream);
-    d_nodes1 = (node*)rmm::mr::get_current_device_resource_ref().allocate(stream, 7 * sizeof(node));
-    d_nodes2 = (node*)rmm::mr::get_current_device_resource_ref().allocate(stream, 7 * sizeof(node));
-    d_progs =
-      (program_t)rmm::mr::get_current_device_resource_ref().allocate(stream, 2 * sizeof(program));
+    d_nodes1 = (node*)rmm::mr::get_current_device_resource_ref().allocate(
+      stream, 7 * sizeof(node), alignof(node));
+    d_nodes2 = (node*)rmm::mr::get_current_device_resource_ref().allocate(
+      stream, 7 * sizeof(node), alignof(node));
+    d_progs = (program_t)rmm::mr::get_current_device_resource_ref().allocate(
+      stream, 2 * sizeof(program), alignof(program));
 
     RAFT_CUDA_TRY(cudaMemcpyAsync(
       d_lYpred.data(), h_lYpred.data(), 500 * sizeof(float), cudaMemcpyHostToDevice, stream));
@@ -146,9 +148,12 @@ class GeneticProgramTest : public ::testing::Test {
 
   void TearDown() override
   {
-    rmm::mr::get_current_device_resource_ref().deallocate(stream, d_nodes1, 7 * sizeof(node));
-    rmm::mr::get_current_device_resource_ref().deallocate(stream, d_nodes2, 7 * sizeof(node));
-    rmm::mr::get_current_device_resource_ref().deallocate(stream, d_progs, 2 * sizeof(program));
+    rmm::mr::get_current_device_resource_ref().deallocate(
+      stream, d_nodes1, 7 * sizeof(node), alignof(node));
+    rmm::mr::get_current_device_resource_ref().deallocate(
+      stream, d_nodes2, 7 * sizeof(node), alignof(node));
+    rmm::mr::get_current_device_resource_ref().deallocate(
+      stream, d_progs, 2 * sizeof(program), alignof(program));
   }
 
   raft::handle_t handle;

python/cuml/cuml/manifold/umap/umap.pyx

@@ -42,8 +42,9 @@ from libcpp.memory cimport unique_ptr
 from libcpp.utility cimport move
 from pylibraft.common.handle cimport handle_t
 from rmm.librmm.device_buffer cimport device_buffer
-from rmm.librmm.per_device_resource cimport get_current_device_resource
+from rmm.librmm.memory_resource cimport make_any_device_resource
 from rmm.pylibrmm.device_buffer cimport DeviceBuffer
+from rmm.pylibrmm.memory_resource cimport get_current_device_resource
 
 cimport cuml.manifold.umap.lib as lib
 from cuml.metrics.distance_type cimport DistanceType
@@ -1286,7 +1287,7 @@ class UMAP(Base, InteropMixin, CMajorInputTagMixin, SparseInputTagMixin):
                     <const void*><uintptr_t>init_m.ptr,
                     init_m.size,
                     handle_.get_stream(),
-                    get_current_device_resource()
+                    make_any_device_resource(get_current_device_resource().get_mr())
                 )
             )