Merge branch 'main' into feat/dlpack-exchange-api

Author: HaoZeke

.github/workflows/tests.yml

@@ -15,7 +15,7 @@ jobs:
     strategy:
       matrix:
         toolchain:
-          - 1.74
+          - "1.74"
           - stable
           - nightly
     steps:
@@ -27,6 +27,19 @@ jobs:
       - name: Setup sccache-cache
         uses: mozilla-actions/sccache-action@v0.0.9
 
+      # MSRV backfill: pyo3 0.28 requires rustc >=1.83 and pyo3 0.27 is
+      # the last line with MSRV 1.74. dlpk accepts `>=0.27, <0.29`, so we
+      # pin the pyo3 stack to the 0.27 line for the 1.74 job only. Stable
+      # and nightly resolve to 0.28 naturally.
+      - name: Pin pyo3 to 0.27 on MSRV toolchain
+        if: matrix.toolchain == '1.74'
+        run: |
+          cargo update -p pyo3 --precise 0.27.2
+          cargo update -p pyo3-build-config --precise 0.27.2
+          cargo update -p pyo3-macros --precise 0.27.2
+          cargo update -p pyo3-macros-backend --precise 0.27.2
+          cargo update -p pyo3-ffi --precise 0.27.2
+
       - name: Run tests
         env:
           SCCACHE_GHA_ENABLED: "true"

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "dlpk"
-version = "0.1.5"
+version = "0.3.1"
 edition = "2021"
 rust-version = "1.74"
 license = "MIT OR Apache-2.0"
@@ -18,7 +18,11 @@ keywords = ["dlpack", "deep-learning", "machine-learning"]
 
 [dependencies]
 ndarray = { version = "0.17", optional = true }
-pyo3 = { version = "0.26", optional = true }
+# Accept either 0.27 or 0.28 so downstream crates on either line can take
+# dlpk as a dep without forcing a pyo3 version conflict. 0.28 has MSRV 1.83
+# while 0.27 keeps MSRV 1.74; the MSRV CI job uses `cargo update --precise`
+# to pin to 0.27.
+pyo3 = { version = ">=0.27, <0.29", optional = true }
 ouroboros = { version = "0.18", optional = true }
 
 # 2.3 has MSRV 1.7.0

src/data_types.rs

@@ -174,6 +174,7 @@ mod tests {
     }
 
     #[test]
+    #[cfg_attr(miri, ignore)]
     fn test_dlpack_pointer_cast() {
         let value: u32 = 42;
         let mut mock_data = value;

src/lib.rs

@@ -28,7 +28,7 @@
 #![allow(clippy::needless_return, clippy::redundant_field_names)]
 #![forbid(clippy::as_ptr_cast_mut, clippy::ptr_cast_constness)]
 
-use std::{ffi::c_void, ptr::NonNull};
+use std::ptr::NonNull;
 
 pub mod sys;
 pub use self::sys::{DLDataType, DLDataTypeCode, DLDevice, DLDeviceType, DLPackVersion};
@@ -252,63 +252,26 @@ impl Drop for DLPackTensor {
     }
 }
 
-struct RustBoxedManager {
-    // tensor: std::pin::Pin<sys::DLManagedTensorVersioned>,
-    original_ctx: *mut c_void,
-    original_deleter: Option<unsafe extern "C" fn(*mut sys::DLManagedTensorVersioned) -> ()>,
-}
-
-unsafe extern "C" fn rust_boxed_manager_deleter(tensor: *mut sys::DLManagedTensorVersioned) {
-    if tensor.is_null() {
-        return;
-    }
-
-    let manager = (*tensor).manager_ctx.cast::<RustBoxedManager>();
-    assert!(!manager.is_null());
-
-    (*tensor).manager_ctx = (*manager).original_ctx;
-    (*tensor).deleter = (*manager).original_deleter;
-
-    if let Some(deleter) = (*tensor).deleter {
-        deleter(tensor);
-    }
-
-    std::mem::drop(Box::from_raw(manager));
-    std::mem::drop(Box::from_raw(tensor));
-}
-
 impl DLPackTensor {
-    /// Create a `DLPackTensor` from a raw `DLManagedTensorVersioned`.
+    /// Create a `DLPackTensor` from a pointer to `DLManagedTensorVersioned`.
     ///
     /// # Safety
     ///
     /// The `DLManagedTensorVersioned` should have a valid `deleter` that can
     /// be called from Rust, or have the deleter set to `None`.
-    pub unsafe fn from_raw(mut tensor: sys::DLManagedTensorVersioned) -> DLPackTensor {
-        // we need to move the tensor to the heap, so we need to wrap the
-        // manager_ctx and deleter into another one that will also free the
-        // tensor from the heap.
-        let manager = Box::new(RustBoxedManager {
-            original_ctx: tensor.manager_ctx,
-            original_deleter: tensor.deleter,
-        });
-        tensor.manager_ctx = Box::into_raw(manager).cast();
-        tensor.deleter = Some(rust_boxed_manager_deleter);
-
-        let tensor = Box::new(tensor);
+    pub unsafe fn from_ptr(tensor: *mut sys::DLManagedTensorVersioned) -> DLPackTensor {
+        let tensor = NonNull::new(tensor).expect("DLManagedTensorVersioned pointer is null");
 
-        return DLPackTensor{
-            raw: NonNull::new_unchecked(Box::into_raw(tensor)),
-        };
+        return DLPackTensor::from_raw(tensor);
     }
 
-    /// Create a `DLPackTensor` from a non-null pointer to `DLManagedTensorVersioned`.
+    /// Create a `DLPackTensor` from a non-null pointer to
+    /// `DLManagedTensorVersioned`.
     ///
     /// # Safety
     ///
-    /// The `DLManagedTensorVersioned` should have a valid `deleter` that can
-    /// be called from Rust, or have the deleter set to `None`.
-    pub unsafe fn from_ptr(tensor: NonNull<sys::DLManagedTensorVersioned>) -> DLPackTensor {
+    /// The same safety requirements as `from_ptr` apply.
+    pub unsafe fn from_raw(tensor: NonNull<sys::DLManagedTensorVersioned>) -> DLPackTensor {
         if tensor.as_ref().version.major != sys::DLPACK_MAJOR_VERSION {
             // from the spec, we need to call the deleter here (and it is the
             // only thing we can do)
@@ -353,6 +316,14 @@ impl DLPackTensor {
         }
     }
 
+    /// Get the ABI version of this DLPack tensor.
+    pub fn version(&self) -> DLPackVersion {
+        let tensor_ref = unsafe {
+            self.raw.as_ref()
+        };
+        tensor_ref.version
+    }
+
     /// Get a pointer to data in this tensor. This pointer can be a device
     /// pointer according to [`DLPackTensor::device`].
     pub fn data_ptr<T>(&self) -> Result<*const T, CastError>  where T: DLPackPointerCast {
@@ -504,9 +475,13 @@ impl<'a> DLPackTensorRef<'a> {
 
     /// Get the shape of this tensor
     pub fn shape(&self) -> &[i64] {
-        assert!(!self.raw.shape.is_null());
-        unsafe {
-            return std::slice::from_raw_parts(self.raw.shape, self.n_dims());
+        if self.raw.shape.is_null() {
+            assert!(self.raw.ndim == 0, "Shape pointer is null but ndim is not 0");
+            return &[];
+        } else {
+            unsafe {
+                return std::slice::from_raw_parts(self.raw.shape, self.n_dims());
+            }
         }
     }
 
@@ -639,6 +614,14 @@ pub mod pyo3;
 #[cfg(feature = "sync")]
 pub mod sync;
 
+/// Small wrapper type to mark a tensor as read-only when there is a choice
+/// between a read-only and a read-write implementation.
+pub struct ReadOnly<T>(pub T);
+
+/// Small wrapper type to mark a tensor as read-write when there is a choice
+/// between a read-only and a read-write implementation.
+pub struct ReadWrite<T>(pub T);
+
 
 #[cfg(test)]
 mod tests {