From 0243630b13a375b06298aef620f01a72d1812b34 Mon Sep 17 00:00:00 2001
From: viclafargue <viclafargue@nvidia.com>
Date: Tue, 15 Mar 2022 19:35:48 +0100
Subject: [PATCH] Resolve remaining issues

---
 .../raft/random/detail/make_regression.cuh    |  1 -
 .../raft/stats/detail/homogeneity_score.cuh   |  1 -
 cpp/test/spatial/fused_l2_knn.cu              |  6 +--
 cpp/test/stats/information_criterion.cu       | 53 ++++++++++---------
 4 files changed, 30 insertions(+), 31 deletions(-)

diff --git a/cpp/include/raft/random/detail/make_regression.cuh b/cpp/include/raft/random/detail/make_regression.cuh
index 8bab85e485..42c1319889 100644
--- a/cpp/include/raft/random/detail/make_regression.cuh
+++ b/cpp/include/raft/random/detail/make_regression.cuh
@@ -30,7 +30,6 @@
 #include <raft/linalg/qr.cuh>
 #include <raft/linalg/transpose.cuh>
 #include <raft/matrix/matrix.cuh>
-#include <raft/mr/device/buffer.hpp>
 #include <raft/random/permute.cuh>
 #include <raft/random/rng.cuh>
 #include <rmm/device_uvector.hpp>
diff --git a/cpp/include/raft/stats/detail/homogeneity_score.cuh b/cpp/include/raft/stats/detail/homogeneity_score.cuh
index 4c78553258..e781b58875 100644
--- a/cpp/include/raft/stats/detail/homogeneity_score.cuh
+++ b/cpp/include/raft/stats/detail/homogeneity_score.cuh
@@ -22,7 +22,6 @@
 
 #pragma once
 
-#include <raft/mr/device/allocator.hpp>
 #include <raft/stats/entropy.cuh>
 #include <raft/stats/mutual_info_score.cuh>
 
diff --git a/cpp/test/spatial/fused_l2_knn.cu b/cpp/test/spatial/fused_l2_knn.cu
index 3be4b47256..2ec4e86d1f 100644
--- a/cpp/test/spatial/fused_l2_knn.cu
+++ b/cpp/test/spatial/fused_l2_knn.cu
@@ -114,9 +114,9 @@ class FusedL2KNNTest : public ::testing::TestWithParam<FusedL2KNNInputs> {
     RAFT_CUDA_TRY(
       cudaMemsetAsync(search_queries.data(), 0, search_queries.size() * sizeof(T), stream_));
     RAFT_CUDA_TRY(
-      cudaMemsetAsync(raft_indices_.data(), 0, raft_indices_.size() * sizeof(T), stream_));
-    RAFT_CUDA_TRY(cudaMemsetAsync(
-      raft_distances_.data(), 0, raft_distances_.size() * sizeof(int64_t), stream_));
+      cudaMemsetAsync(raft_indices_.data(), 0, raft_indices_.size() * sizeof(int64_t), stream_));
+    RAFT_CUDA_TRY(
+      cudaMemsetAsync(raft_distances_.data(), 0, raft_distances_.size() * sizeof(T), stream_));
     RAFT_CUDA_TRY(
       cudaMemsetAsync(faiss_indices_.data(), 0, faiss_indices_.size() * sizeof(int64_t), stream_));
     RAFT_CUDA_TRY(
diff --git a/cpp/test/stats/information_criterion.cu b/cpp/test/stats/information_criterion.cu
index 802e3fee23..d61f8591a5 100644
--- a/cpp/test/stats/information_criterion.cu
+++ b/cpp/test/stats/information_criterion.cu
@@ -19,7 +19,8 @@
 #include <raft/stats/information_criterion.cuh>
 
 #include <raft/cudart_utils.h>
-#include <raft/mr/device/allocator.hpp>
+#include <raft/handle.hpp>
+#include <rmm/device_uvector.hpp>
 
 #include <gtest/gtest.h>
 
@@ -59,21 +60,23 @@ struct BatchedICInputs {
 
 template <typename T>
 class BatchedICTest : public ::testing::TestWithParam<BatchedICInputs<T>> {
+ public:
+  BatchedICTest()
+    : params(::testing::TestWithParam<BatchedICInputs<T>>::GetParam()),
+      stream(handle.get_stream()),
+      res_d(sizeof(T) * params.batch_size, stream)
+  {
+  }
+
  protected:
   void SetUp() override
   {
     using std::vector;
-    params = ::testing::TestWithParam<BatchedICInputs<T>>::GetParam();
-
-    // Create stream and allocator
-    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;
@@ -83,11 +86,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_batched(res_d,
-                                  loglike_d,
+    information_criterion_batched(res_d.data(),
+                                  loglike_d.data(),
                                   params.ic_type,
                                   params.n_params,
                                   params.batch_size,
@@ -103,22 +106,14 @@ class BatchedICTest : public ::testing::TestWithParam<BatchedICInputs<T>> {
              params.n_samples);
 
     RAFT_CUDA_TRY(cudaStreamSynchronize(stream));
-
-    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;
+  raft::handle_t handle;
+  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)
@@ -133,13 +128,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));