Remove RAFT memory management (2/2) (#4526)

Authors:
  - Victor Lafargue (https://github.com/viclafargue)

Approvers:
  - Corey J. Nolet (https://github.com/cjnolet)

URL: #4526

cpp/bench/common/ml_benchmark.hpp

@@ -22,7 +22,6 @@
 #include <cuml/common/utils.hpp>
 #include <memory>
 #include <raft/cudart_utils.h>
-#include <raft/mr/device/allocator.hpp>
 #include <sstream>
 #include <string>
 #include <vector>

cpp/bench/prims/permute.cu

@@ -16,7 +16,6 @@
 
 #include <common/ml_benchmark.hpp>
 #include <raft/cudart_utils.h>
-#include <raft/mr/device/allocator.hpp>
 #include <raft/random/permute.hpp>
 #include <raft/random/rng.hpp>
 

cpp/bench/sg/benchmark.cuh

@@ -24,7 +24,6 @@
 #include <cuml/common/logger.hpp>
 #include <raft/cudart_utils.h>
 #include <raft/handle.hpp>
-#include <raft/mr/device/allocator.hpp>
 
 namespace ML {
 namespace Bench {

cpp/examples/symreg/symreg_example.cpp

@@ -198,7 +198,6 @@ int main(int argc, char* argv[])
 
   /* ======================= Begin GPU memory allocation ======================= */
   std::cout << "***************************************" << std::endl;
-  std::shared_ptr<raft::mr::device::allocator> allocator(new raft::mr::device::default_allocator());
 
   cudaStream_t stream;
   raft::handle_t handle{stream};

cpp/src/common/cuml_api.cpp

@@ -20,64 +20,9 @@
 #include <cuml/cuml_api.h>
 
 #include <raft/cudart_utils.h>
-#include <raft/mr/device/allocator.hpp>
-#include <raft/mr/host/allocator.hpp>
 
 #include <cstddef>
 #include <functional>
-namespace ML {
-namespace detail {
-
-class hostAllocatorFunctionWrapper : public raft::mr::host::allocator {
- public:
-  hostAllocatorFunctionWrapper(cuml_allocate allocate_fn, cuml_deallocate deallocate_fn)
-    : _allocate_fn(allocate_fn), _deallocate_fn(deallocate_fn)
-  {
-  }
-
-  virtual void* allocate(std::size_t n, cudaStream_t stream)
-  {
-    void* ptr = 0;
-    RAFT_CUDA_TRY(_allocate_fn(&ptr, n, stream));
-    return ptr;
-  }
-
-  virtual void deallocate(void* p, std::size_t n, cudaStream_t stream)
-  {
-    RAFT_CUDA_TRY_NO_THROW(_deallocate_fn(p, n, stream));
-  }
-
- private:
-  const std::function<cudaError_t(void**, size_t, cudaStream_t)> _allocate_fn;
-  const std::function<cudaError_t(void*, size_t, cudaStream_t)> _deallocate_fn;
-};
-
-class deviceAllocatorFunctionWrapper : public raft::mr::device::default_allocator {
- public:
-  deviceAllocatorFunctionWrapper(cuml_allocate allocate_fn, cuml_deallocate deallocate_fn)
-    : _allocate_fn(allocate_fn), _deallocate_fn(deallocate_fn)
-  {
-  }
-
-  virtual void* allocate(std::size_t n, cudaStream_t stream)
-  {
-    void* ptr = 0;
-    RAFT_CUDA_TRY(_allocate_fn(&ptr, n, stream));
-    return ptr;
-  }
-
-  virtual void deallocate(void* p, std::size_t n, cudaStream_t stream)
-  {
-    RAFT_CUDA_TRY_NO_THROW(_deallocate_fn(p, n, stream));
-  }
-
- private:
-  const std::function<cudaError_t(void**, size_t, cudaStream_t)> _allocate_fn;
-  const std::function<cudaError_t(void*, size_t, cudaStream_t)> _deallocate_fn;
-};
-
-}  // end namespace detail
-}  // end namespace ML
 
 extern "C" const char* cumlGetErrorString(cumlError_t error)
 {

cpp/src/pca/sign_flip_mg.cu

@@ -18,8 +18,6 @@
 #include <thrust/device_vector.h>
 #include <thrust/execution_policy.h>
 
-#include <common/allocatorAdapter.hpp>
-
 #include <raft/comms/comms.hpp>
 #include <raft/cuda_utils.cuh>
 #include <raft/handle.hpp>

cpp/test/mg/knn_classify.cu

@@ -28,7 +28,6 @@ void generate_partitions(float* data,
                          int n_cols,
                          int n_clusters,
                          int my_rank,
-                         std::shared_ptr<raft::mr::device::allocator> allocator,
                          cudaStream_t stream)
 {
   Random::make_blobs<float, int>(data,

cpp/test/prims/batched/information_criterion.cu

@@ -19,7 +19,7 @@
 #include <metrics/batched/information_criterion.cuh>
 
 #include <raft/cudart_utils.h>
-#include <raft/mr/device/allocator.hpp>
+#include <rmm/device_uvector.hpp>
 
 #include <gtest/gtest.h>
 
@@ -61,20 +61,23 @@ struct BatchedICInputs {
 template <typename T>
 class BatchedICTest : public ::testing::TestWithParam<BatchedICInputs<T>> {
  protected:
+  BatchedICTest()
+    : params(::testing::TestWithParam<BatchedICInputs<T>>::GetParam()),
+      res_d(sizeof(T) * params.batch_size, stream)
+  {
+  }
+
   void SetUp() override
   {
     using std::vector;
-    params = ::testing::TestWithParam<BatchedICInputs<T>>::GetParam();
 
-    // Create stream and allocator
+    // Create stream
     RAFT_CUDA_TRY(cudaStreamCreate(&stream));
-    allocator = std::make_shared<raft::mr::device::default_allocator>();
 
     // Create arrays
     std::vector<T> loglike_h = std::vector<T>(params.batch_size);
     res_h.resize(params.batch_size);
-    T* loglike_d = (T*)allocator->allocate(sizeof(T) * params.batch_size, stream);
-    res_d        = (T*)allocator->allocate(sizeof(T) * params.batch_size, stream);
+    rmm::device_uvector<T> loglike_d(sizeof(T) * params.batch_size, stream);
 
     // Generate random data
     std::random_device rd;
@@ -84,11 +87,11 @@ class BatchedICTest : public ::testing::TestWithParam<BatchedICInputs<T>> {
       loglike_h[i] = std::log(udis(gen));
 
     // Copy the data to the device
-    raft::update_device(loglike_d, loglike_h.data(), params.batch_size, stream);
+    raft::update_device(loglike_d.data(), loglike_h.data(), params.batch_size, stream);
 
     // Compute the tested results
-    information_criterion(res_d,
-                          loglike_d,
+    information_criterion(res_d.data(),
+                          loglike_d.data(),
                           params.ic_type,
                           params.n_params,
                           params.batch_size,
@@ -102,22 +105,13 @@ class BatchedICTest : public ::testing::TestWithParam<BatchedICInputs<T>> {
              params.n_params,
              params.batch_size,
              params.n_samples);
-
-    allocator->deallocate(loglike_d, sizeof(T) * params.batch_size, stream);
-  }
-
-  void TearDown() override
-  {
-    allocator->deallocate(res_d, sizeof(T) * params.batch_size, stream);
-    RAFT_CUDA_TRY(cudaStreamDestroy(stream));
   }
 
  protected:
-  std::shared_ptr<raft::mr::device::default_allocator> allocator;
+  cudaStream_t stream = 0;
   BatchedICInputs<T> params;
-  T* res_d;
+  rmm::device_uvector<T> res_d;
   std::vector<T> res_h;
-  cudaStream_t stream = 0;
 };
 
 // Test parameters (op, n_batches, m, n, p, q, tolerance)
@@ -132,13 +126,19 @@ using BatchedICTestD = BatchedICTest<double>;
 using BatchedICTestF = BatchedICTest<float>;
 TEST_P(BatchedICTestD, Result)
 {
-  ASSERT_TRUE(devArrMatchHost(
-    res_h.data(), res_d, params.batch_size, raft::CompareApprox<double>(params.tolerance), stream));
+  ASSERT_TRUE(devArrMatchHost(res_h.data(),
+                              res_d.data(),
+                              params.batch_size,
+                              raft::CompareApprox<double>(params.tolerance),
+                              stream));
 }
 TEST_P(BatchedICTestF, Result)
 {
-  ASSERT_TRUE(devArrMatchHost(
-    res_h.data(), res_d, params.batch_size, raft::CompareApprox<float>(params.tolerance), stream));
+  ASSERT_TRUE(devArrMatchHost(res_h.data(),
+                              res_d.data(),
+                              params.batch_size,
+                              raft::CompareApprox<float>(params.tolerance),
+                              stream));
 }
 
 INSTANTIATE_TEST_CASE_P(BatchedICTests, BatchedICTestD, ::testing::ValuesIn(inputsd));