Make manager context pinnable with safer and more ergonomic bi-directional conversions

Author: Ehsan M. Kermani

examples/sample/src/main.rs

@@ -1,8 +1,8 @@
-#![allow(clippy::drop_copy)]
+#![allow(clippy::drop_ref)]
 
-use std::{ffi::c_void, mem, ptr};
+use std::ffi::c_void;
 
-use dlpackrs::{ffi, DataType, Device, ManagedTensor as DLManagedTensor, Tensor as DLTensor};
+use dlpackrs::{DataType, Device, ManagedTensor as DLManagedTensor, Tensor as DLTensor};
 use ndarray::{Array, ArrayD, RawArrayViewMut};
 
 #[derive(Debug)]
@@ -32,19 +32,20 @@ impl<'tensor> From<&'tensor mut Tensor<'tensor>> for ArrayD<f32> {
     }
 }
 
+// The context holds DLManagedTensor
 #[derive(Debug)]
-pub struct ManagedTensor<'tensor, 'ctx>(DLManagedTensor<'tensor, 'ctx>);
+pub struct ManagedContext<'tensor, C>(DLManagedTensor<'tensor, C>);
 
-impl<'tensor, 'ctx> From<&'tensor mut ArrayD<f32>> for ManagedTensor<'tensor, 'ctx> {
+impl<'tensor, C> From<&'tensor mut ArrayD<f32>> for ManagedContext<'tensor, C> {
     fn from(t: &'tensor mut ArrayD<f32>) -> Self {
         let dlt: Tensor<'tensor> = Tensor::from(t);
-        let inner = DLManagedTensor::new(dlt.0, ptr::null_mut());
-        ManagedTensor(inner)
+        let inner = DLManagedTensor::new(dlt.0, None);
+        ManagedContext(inner)
     }
 }
 
-impl<'tensor, 'ctx> From<&mut ManagedTensor<'tensor, 'ctx>> for ArrayD<f32> {
-    fn from(mt: &mut ManagedTensor<'tensor, 'ctx>) -> Self {
+impl<'tensor, C> From<&mut ManagedContext<'tensor, C>> for ArrayD<f32> {
+    fn from(mt: &mut ManagedContext<'tensor, C>) -> Self {
         let dlt: DLTensor = mt.0.inner.dl_tensor.into();
         unsafe {
             let arr = RawArrayViewMut::from_shape_ptr(dlt.shape().unwrap(), dlt.data() as *mut f32);
@@ -70,16 +71,11 @@ fn main() {
     let pong = ArrayD::from(&mut tensor);
     println!("pong {:?}", pong);
     assert!(pong.into_dyn().abs_diff_eq(&ping, 1e-8f32));
-    // TODO: use dummy ctx holding the managed_tensor
-    let mut managed_tensor = ManagedTensor::from(&mut ping);
+    let mut managed_tensor: ManagedContext<f32> = (&mut ping).into();
     println!("managed tensor {:?}", managed_tensor);
-    let deleter = unsafe {
-        mem::transmute::<fn(&mut ManagedTensor), unsafe extern "C" fn(*mut ffi::DLManagedTensor)>(
-            |mt| {
-                println!("deleter is called!");
-                drop(mt as *mut _);
-            },
-        )
+    let deleter: fn(&mut DLManagedTensor<f32>) = |managed_tensor| {
+        println!("manager tensor deleter is called");
+        drop(managed_tensor);
     };
     managed_tensor.0.set_deleter(deleter);
     println!("managed tensor with deleter {:?}", managed_tensor);

src/lib.rs

@@ -25,7 +25,7 @@
 //!
 //! In this case, `ManagedTensor` is built from `ManagedTensorProxy` which is a safe proxy for the unsafe `ffi::DLManagedTensor`.
 //!
-//! ### Non-Memory Manged Tensor (View)
+//! ### Plain Not-Memory-Managed Tensor
 //!
 //! In this case, the (invariant) Rust wrapper `Tensor` can be used or if needed the unsafe `ffi::DLTensor`.
 //!
@@ -35,21 +35,21 @@
 //!
 //! When ownership is concerned, one can use the `ManagedTensor`. Here is an example on how the bi-directional conversion
 //!
-//! <div align="center">ndarray::ArrayD <--> ManagedTensor</div>
+//! <div align="center">ndarray::ArrayD <---> ManagedTensor</div>
 //!
 //! is done at zero-cost.
 //!
 //! ```no_run
-//! impl<'tensor, 'ctx> From<&'tensor mut ArrayD<f32>> for ManagedTensor<'tensor, 'ctx> {
+//! impl<'tensor, C> From<&'tensor mut ArrayD<f32>> for ManagedContext<'tensor, C> {
 //!     fn from(t: &'tensor mut ArrayD<f32>) -> Self {
 //!         let dlt: Tensor<'tensor> = Tensor::from(t);
-//!         let inner = DLManagedTensor::new(dlt.0, ptr::null_mut());
-//!         ManagedTensor(inner)
+//!         let inner = DLManagedTensor::new(dlt.0, None);
+//!         ManagedContext(inner)
 //!     }
 //! }
 //!
-//! impl<'tensor, 'ctx> From<&mut ManagedTensor<'tensor, 'ctx>> for ArrayD<f32> {
-//!     fn from(mt: &mut ManagedTensor<'tensor, 'ctx>) -> Self {
+//! impl<'tensor, C> From<&mut ManagedContext<'tensor, C>> for ArrayD<f32> {
+//!     fn from(mt: &mut ManagedContext<'tensor, C>) -> Self {
 //!         let dlt: DLTensor = mt.0.inner.dl_tensor.into();
 //!         unsafe {
 //!             let arr = RawArrayViewMut::from_shape_ptr(dlt.shape().unwrap(), dlt.data() as *mut f32);
@@ -63,7 +63,7 @@
 //!
 //! And when ownership is not concerned, one can use `Tensor` as a view. Here is an example on how the bi-directional converion
 //!
-//! <div align="center">ndarray::ArrayD <--> Tensor</div>
+//! <div align="center">ndarray::ArrayD <---> Tensor</div>
 //!
 //! is done at zero-cost.
 //!
@@ -104,12 +104,12 @@ pub mod ffi {
 
 pub mod datatype;
 pub mod device;
-mod errors;
+pub mod errors;
 pub mod tensor;
 
 pub use datatype::{DataType, DataTypeCode};
 pub use device::{Device, DeviceType};
-pub use tensor::{ManagedTensor, ManagedTensorProxy, Tensor};
+pub use tensor::{ManagedTensor, ManagedTensorProxy, ManagerContext, Tensor};
 
 pub fn version() -> u32 {
     ffi::DLPACK_VERSION

src/tensor.rs

@@ -1,7 +1,14 @@
 use pin_project::{pin_project, pinned_drop};
 
 use core::slice;
-use std::{marker::PhantomData, os::raw::c_void, pin::Pin};
+use std::{
+    fmt::Debug,
+    marker::{PhantomData, PhantomPinned},
+    mem::transmute,
+    os::raw::c_void,
+    pin::Pin,
+    ptr::{self, NonNull},
+};
 
 use crate::{
     datatype::DataType,
@@ -138,84 +145,137 @@ impl<'tensor> Tensor<'tensor> {
     }
 }
 
-/// Safe proxy to ffi::DLManagedTensor which is self-referential by design.
-/// See [DLManagedTensor](https://dmlc.github.io/dlpack/latest/c_api.html#_CPPv415DLManagedTensor)
+/// A typed ManagerContext type that is `!Unpin` i.e. pinnable for safety since it holds a pointer to the underlying DLTensor.
 #[derive(Debug)]
-#[pin_project(PinnedDrop)]
-pub struct ManagedTensorProxy<'ctx> {
-    pub dl_tensor: DLTensor,
-    #[pin]
-    /// Holds the underlying DLTensor.
-    pub manager_ctx: *mut c_void,
-    #[pin]
-    pub deleter: Option<unsafe extern "C" fn(*mut DLManagedTensor)>,
-    _marker: PhantomData<&'ctx ()>, // covariant wrt 'ctx
+#[repr(C)]
+pub struct ManagerContext<C> {
+    pub ptr: Option<NonNull<*mut c_void>>,
+    ty: PhantomData<C>,
+    _pin: PhantomPinned,
 }
 
-impl<'ctx> From<DLManagedTensor> for ManagedTensorProxy<'ctx> {
-    fn from(dlmt: DLManagedTensor) -> Self {
-        ManagedTensorProxy {
-            dl_tensor: dlmt.dl_tensor,
-            manager_ctx: dlmt.manager_ctx,
-            deleter: dlmt.deleter,
-            _marker: PhantomData,
+impl<C> ManagerContext<C> {
+    pub fn new(ptr: Option<NonNull<*mut c_void>>) -> Self {
+        Self {
+            ptr,
+            ty: PhantomData,
+            _pin: PhantomPinned,
         }
     }
 }
 
-impl<'ctx> From<ManagedTensorProxy<'ctx>> for DLManagedTensor {
-    fn from(pmt: ManagedTensorProxy<'ctx>) -> Self {
-        DLManagedTensor {
-            dl_tensor: pmt.dl_tensor,
-            manager_ctx: pmt.manager_ctx,
-            deleter: pmt.deleter,
-        }
-    }
+/// Safe proxy to ffi::DLManagedTensor which is self-referential by design.
+/// See [DLManagedTensor](https://dmlc.github.io/dlpack/latest/c_api.html#_CPPv415DLManagedTensor)
+#[pin_project(PinnedDrop)]
+#[repr(C)]
+pub struct ManagedTensorProxy<C> {
+    /// Holds the underlying tensor.
+    pub dl_tensor: DLTensor,
+    /// The context holding the underlying DLTensor.
+    #[pin]
+    pub manager_ctx: ManagerContext<C>, // safe typed wrapper for *mut c_void which is !Unpin i.e. pinnable
+    /// Deleter function pointer.
+    // TODO: should this be `#[pin]`?
+    pub deleter: Option<fn(&mut ManagedTensor<C>)>,
 }
 
-impl<'ctx> From<Pin<&mut ManagedTensorProxy<'ctx>>> for DLManagedTensor {
-    fn from(pmt: Pin<&mut ManagedTensorProxy<'ctx>>) -> Self {
-        DLManagedTensor {
-            dl_tensor: pmt.dl_tensor,
-            manager_ctx: pmt.manager_ctx,
-            deleter: pmt.deleter,
-        }
+impl<C: Debug> Debug for ManagedTensorProxy<C> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("ManagedTensorProxy")
+            .field("dl_tensor", &self.dl_tensor)
+            .field("manager_ctx", &self.manager_ctx)
+            .finish()
     }
 }
 
-impl<'ctx> ManagedTensorProxy<'ctx> {
+impl<C> ManagedTensorProxy<C> {
     pub fn dl_tensor(&self) -> DLTensor {
         self.dl_tensor
     }
 
-    pub fn manager_ctx(self: Pin<&mut Self>) -> *mut c_void {
-        let this = self.project();
-        *this.manager_ctx.get_mut()
+    pub fn manager_ctx(self: Pin<&mut Self>) -> Option<NonNull<*mut c_void>> {
+        let mut this = self.project();
+        this.manager_ctx.as_mut().ptr
     }
 
-    pub fn set_manager_ctx(self: Pin<&mut Self>, manager_ctx: *mut c_void) {
+    pub fn set_manager_ctx(self: Pin<&mut Self>, manager_ctx: NonNull<*mut c_void>) {
         let mut this = self.project();
-        *this.manager_ctx = manager_ctx;
+        let new = ManagerContext::new(Some(manager_ctx));
+        this.manager_ctx.set(new);
+    }
+}
+
+impl<C> From<DLManagedTensor> for ManagedTensorProxy<C> {
+    fn from(mut dlmt: DLManagedTensor) -> Self {
+        let ptr: Option<NonNull<*mut c_void>> = if dlmt.manager_ctx.is_null() {
+            None
+        } else {
+            unsafe { Some(NonNull::new_unchecked(&mut dlmt.manager_ctx as *mut _)) }
+        };
+        let manager_ctx = ManagerContext::new(ptr);
+        let deleter = dlmt.deleter.take().map(|del| unsafe {
+            transmute::<unsafe extern "C" fn(*mut DLManagedTensor), fn(&mut ManagedTensor<C>)>(del)
+        });
+        ManagedTensorProxy {
+            dl_tensor: dlmt.dl_tensor,
+            manager_ctx,
+            deleter,
+        }
+    }
+}
+
+impl<C> From<ManagedTensorProxy<C>> for DLManagedTensor {
+    fn from(pmt: ManagedTensorProxy<C>) -> Self {
+        let dl_tensor = pmt.dl_tensor;
+        let manager_ctx = match pmt.manager_ctx.ptr {
+            None => ptr::null_mut(),
+            Some(nnptr) => unsafe { *nnptr.as_ptr() },
+        };
+        let deleter = unsafe {
+            pmt.deleter.map(|del_fn| {
+                transmute::<fn(&mut ManagedTensor<C>), unsafe extern "C" fn(*mut DLManagedTensor)>(
+                    del_fn,
+                )
+            })
+        };
+        DLManagedTensor {
+            dl_tensor,
+            manager_ctx,
+            deleter,
+        }
     }
+}
 
-    pub fn deleter(
-        self: Pin<&mut Self>,
-    ) -> Option<unsafe extern "C" fn(self_: *mut DLManagedTensor)> {
-        let this = self.project();
-        *this.deleter.get_mut()
+impl<C> From<Pin<&mut ManagedTensorProxy<C>>> for DLManagedTensor {
+    fn from(pmt: Pin<&mut ManagedTensorProxy<C>>) -> Self {
+        let dl_tensor = pmt.dl_tensor;
+        let manager_ctx = match pmt.manager_ctx.ptr {
+            None => ptr::null_mut(),
+            Some(nnptr) => unsafe { *nnptr.as_ptr() },
+        };
+        let deleter = unsafe {
+            pmt.deleter.map(|del_fn| {
+                transmute::<fn(&mut ManagedTensor<C>), unsafe extern "C" fn(*mut DLManagedTensor)>(
+                    del_fn,
+                )
+            })
+        };
+        DLManagedTensor {
+            dl_tensor,
+            manager_ctx,
+            deleter,
+        }
     }
 }
 
 #[allow(clippy::needless_lifetimes)]
 #[pinned_drop]
-impl<'ctx> PinnedDrop for ManagedTensorProxy<'ctx> {
+impl<C> PinnedDrop for ManagedTensorProxy<C> {
     fn drop(mut self: Pin<&mut Self>) {
         let mut dlm: DLManagedTensor = self.as_mut().into();
-        if let Some(fptr) = self.deleter() {
+        if let Some(fptr) = self.deleter {
             unsafe {
-                let cfptr = std::mem::transmute::<*const (), unsafe fn(*mut DLManagedTensor)>(
-                    fptr as *const (),
-                );
+                let cfptr = transmute::<fn(&mut ManagedTensor<C>), fn(*mut DLManagedTensor)>(fptr);
                 cfptr(&mut dlm as *mut _);
             };
         }
@@ -227,18 +287,35 @@ impl<'ctx> PinnedDrop for ManagedTensorProxy<'ctx> {
 ///  See [DLManagedTensor](https://dmlc.github.io/dlpack/latest/c_api.html#_CPPv415DLManagedTensor)
 #[derive(Debug)]
 #[repr(transparent)]
-pub struct ManagedTensor<'tensor, 'ctx: 'tensor> {
-    pub inner: ManagedTensorProxy<'ctx>,
-    _marker: PhantomData<fn(&'tensor ()) -> &'tensor ()>,
+pub struct ManagedTensor<'tensor, C: 'tensor> {
+    pub inner: ManagedTensorProxy<C>,
+    _marker: PhantomData<fn(&'tensor ()) -> &'tensor ()>, // invariant wrt 'tensor
+}
+
+impl<'tensor, C> From<DLManagedTensor> for ManagedTensor<'tensor, C> {
+    fn from(dlm: DLManagedTensor) -> Self {
+        let proxy: ManagedTensorProxy<C> = dlm.into();
+        ManagedTensor {
+            inner: proxy,
+            _marker: PhantomData,
+        }
+    }
 }
 
-impl<'tensor, 'ctx> ManagedTensor<'tensor, 'ctx> {
-    pub fn new(tensor: Tensor<'tensor>, manager_ctx: *mut c_void) -> Self {
+impl<'tensor, C> From<ManagedTensor<'tensor, C>> for DLManagedTensor {
+    fn from(mt: ManagedTensor<'tensor, C>) -> Self {
+        mt.inner.into()
+    }
+}
+
+impl<'tensor, C: 'tensor> ManagedTensor<'tensor, C> {
+    /// Contructor.
+    pub fn new(tensor: Tensor<'tensor>, manager_ctx: Option<NonNull<*mut c_void>>) -> Self {
+        let manager_ctx = ManagerContext::new(manager_ctx);
         let inner = ManagedTensorProxy {
             dl_tensor: tensor.into_inner(),
             manager_ctx,
             deleter: None,
-            _marker: PhantomData,
         };
 
         ManagedTensor {
@@ -247,8 +324,8 @@ impl<'tensor, 'ctx> ManagedTensor<'tensor, 'ctx> {
         }
     }
 
-    /// Sets a deleter function pointer
-    pub fn set_deleter(&mut self, deleter: unsafe extern "C" fn(*mut DLManagedTensor)) {
+    /// Sets a deleter function pointer.
+    pub fn set_deleter(&mut self, deleter: fn(&mut ManagedTensor<C>)) {
         self.inner.deleter = Some(deleter);
     }
 
@@ -261,9 +338,8 @@ impl<'tensor, 'ctx> ManagedTensor<'tensor, 'ctx> {
     /// Returns a ManagedTensor instances from a raw pointer to DLManagedTensor.
     pub unsafe fn from_raw(ptr: *mut DLManagedTensor) -> Self {
         debug_assert!(!ptr.is_null());
-        let proxy = (*ptr).into();
         ManagedTensor {
-            inner: proxy,
+            inner: (*ptr).into(),
             _marker: PhantomData,
         }
     }