From 7d5174a145d46a2ae6d1cbf3828915319223c213 Mon Sep 17 00:00:00 2001
From: Ben Frederickson <ben@benfrederickson.com>
Date: Tue, 20 Feb 2024 10:35:37 -0800
Subject: [PATCH] Add resources to to_host/to_device functions

eventually we'll want the stream arg on these
---
 rust/cuvs/examples/cagra.rs  | 10 +++++-----
 rust/cuvs/src/cagra/index.rs | 10 +++++-----
 rust/cuvs/src/dlpack.rs      |  4 +++-
 3 files changed, 13 insertions(+), 11 deletions(-)

diff --git a/rust/cuvs/examples/cagra.rs b/rust/cuvs/examples/cagra.rs
index 51e28d023..879fac787 100644
--- a/rust/cuvs/examples/cagra.rs
+++ b/rust/cuvs/examples/cagra.rs
@@ -47,20 +47,20 @@ fn cagra_example() -> Result<()> {
     // CAGRA search API requires queries and outputs to be on device memory
     // copy query data over, and allocate new device memory for the distances/ neighbors
     // outputs
-    let queries = ManagedTensor::from_ndarray(&queries).to_device()?;
+    let queries = ManagedTensor::from_ndarray(&queries).to_device(&res)?;
     let mut neighbors_host = ndarray::Array::<u32, _>::zeros((n_queries, k));
-    let neighbors = ManagedTensor::from_ndarray(&neighbors_host).to_device()?;
+    let neighbors = ManagedTensor::from_ndarray(&neighbors_host).to_device(&res)?;
 
     let mut distances_host = ndarray::Array::<f32, _>::zeros((n_queries, k));
-    let distances = ManagedTensor::from_ndarray(&distances_host).to_device()?;
+    let distances = ManagedTensor::from_ndarray(&distances_host).to_device(&res)?;
 
     let search_params = SearchParams::new()?;
 
     index.search(&res, &search_params, &queries, &neighbors, &distances)?;
 
     // Copy back to host memory
-    distances.to_host(&mut distances_host)?;
-    neighbors.to_host(&mut neighbors_host)?;
+    distances.to_host(&res, &mut distances_host)?;
+    neighbors.to_host(&res, &mut neighbors_host)?;
 
     // nearest neighbors should be themselves, since queries are from the
     // dataset
diff --git a/rust/cuvs/src/cagra/index.rs b/rust/cuvs/src/cagra/index.rs
index 7f2bc96fb..be1543a29 100644
--- a/rust/cuvs/src/cagra/index.rs
+++ b/rust/cuvs/src/cagra/index.rs
@@ -105,15 +105,15 @@ mod tests {
         // CAGRA search API requires queries and outputs to be on device memory
         // copy query data over, and allocate new device memory for the distances/ neighbors
         // outputs
-        let queries = ManagedTensor::from_ndarray(&queries).to_device().unwrap();
+        let queries = ManagedTensor::from_ndarray(&queries).to_device(&res).unwrap();
         let mut neighbors_host = ndarray::Array::<u32, _>::zeros((n_queries, k));
         let neighbors = ManagedTensor::from_ndarray(&neighbors_host)
-            .to_device()
+            .to_device(&res)
             .unwrap();
 
         let mut distances_host = ndarray::Array::<f32, _>::zeros((n_queries, k));
         let distances = ManagedTensor::from_ndarray(&distances_host)
-            .to_device()
+            .to_device(&res)
             .unwrap();
 
         let search_params = SearchParams::new().unwrap();
@@ -123,8 +123,8 @@ mod tests {
             .unwrap();
 
         // Copy back to host memory
-        distances.to_host(&mut distances_host).unwrap();
-        neighbors.to_host(&mut neighbors_host).unwrap();
+        distances.to_host(&res, &mut distances_host).unwrap();
+        neighbors.to_host(&res, &mut neighbors_host).unwrap();
 
         // nearest neighbors should be themselves, since queries are from the
         // dataset
diff --git a/rust/cuvs/src/dlpack.rs b/rust/cuvs/src/dlpack.rs
index b9d7d7e6d..26d618a4a 100644
--- a/rust/cuvs/src/dlpack.rs
+++ b/rust/cuvs/src/dlpack.rs
@@ -15,6 +15,7 @@
  */
 
 use crate::error::{check_cuda, Result};
+use crate::resources::Resources;
 
 #[derive(Debug)]
 pub struct ManagedTensor(ffi::DLManagedTensor);
@@ -66,7 +67,7 @@ impl ManagedTensor {
         bytes
     }
 
-    pub fn to_device(&self) -> Result<ManagedTensor> {
+    pub fn to_device(&self, _res: &Resources) -> Result<ManagedTensor> {
         unsafe {
             let bytes = self.bytes();
             let mut device_data: *mut std::ffi::c_void = std::ptr::null_mut();
@@ -95,6 +96,7 @@ impl ManagedTensor {
         D: ndarray::Dimension,
     >(
         &self,
+        _res: &Resources,
         arr: &mut ndarray::ArrayBase<S, D>,
     ) -> Result<()> {
         unsafe {