diff --git a/Cargo.lock b/Cargo.lock index d0950d8..af72f49 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1382,8 +1382,8 @@ dependencies = [ [[package]] name = "dlpark" -version = "0.7.0" -source = "git+https://github.com/kylebarron/dlpark?rev=31c6f49c064e634326c97172d39a00acecd854b6#31c6f49c064e634326c97172d39a00acecd854b6" +version = "0.8.0" +source = "git+https://github.com/kylebarron/dlpark?rev=f6887bf8b367d98d079710be9659c4177ba899f8#f6887bf8b367d98d079710be9659c4177ba899f8" dependencies = [ "bitflags 2.13.0", "pyo3", diff --git a/Cargo.toml b/Cargo.toml index 2283a88..9d146f4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -41,7 +41,10 @@ arrow-buffer = "59" arrow-schema = "59" async-trait = { version = "0.1", optional = true } bytes = "1.12.0" -dlpark = { git = "https://github.com/kylebarron/dlpark", rev = "31c6f49c064e634326c97172d39a00acecd854b6", features = [ +# Fork of dlpark 0.8.0 that gates bindgen's 64-bit layout tests so the crate +# compiles for 32-bit targets (wasm32-unknown-emscripten). Upstreamed to +# SunDoge/dlpark; revert to the crates.io release once a fixed version ships. +dlpark = { git = "https://github.com/kylebarron/dlpark", rev = "f6887bf8b367d98d079710be9659c4177ba899f8", features = [ "pyo3", ] } icechunk = { version = "2.0.0", default-features = false, optional = true } diff --git a/src/data/tensor.rs b/src/data/tensor.rs index 8713fac..98f7b44 100644 --- a/src/data/tensor.rs +++ b/src/data/tensor.rs @@ -2,9 +2,9 @@ use std::ffi::{c_int, c_void}; use std::sync::Arc; use bytes::Bytes; -use dlpark::SafeManagedTensor; -use dlpark::ffi::{Device, DeviceType}; -use dlpark::traits::{RowMajorCompactLayout, TensorLike}; +use dlpark::ffi::{DLDataType, DLDataTypeCode, DLDevice, DLDeviceType}; +use dlpark::metadata::CopiedSlice; +use dlpark::{Builder, legacy}; use pyo3::exceptions::{PyNotImplementedError, PyValueError}; use pyo3::ffi; use pyo3::prelude::*; @@ -14,7 +14,7 @@ use zarrs::array::DataType; use crate::data::buffer_protocol::PyTensorBuffer; use crate::dtype::PyDataType; -use crate::error::{ZarristaError, ZarristaResult}; +use crate::error::ZarristaResult; /// Fixed-width, dense decoded data. /// @@ -118,13 +118,39 @@ impl PyTensor { fn __dlpack__<'py>( &self, _kwargs: Option>, - ) -> ZarristaResult { - SafeManagedTensor::new(self.clone()) + ) -> ZarristaResult { + let dtype = self.dlpack_data_type()?; + let shape = self + .shape + .iter() + .map(|s| i64::try_from(*s).expect("overflow converting shape to i64")) + .collect::>(); + let strides = row_major_compact_strides(&shape); + + // The boxed `Bytes` handed to the builder is what keeps the buffer alive: dlpark stores it + // as the managed tensor's `manager_ctx` and drops it from the deleter, which a consumer may + // run on any thread and long after this `PyTensor` is gone. + let data = self.bytes.as_ptr().cast::().cast_mut(); + let builder = Builder::new( + Box::new(self.bytes.clone()), + CopiedSlice::new(shape, strides), + ); + // SAFETY: `data` points at the start of the `Bytes` allocation moved into the context, so + // it stays valid until the deleter runs. The shape is the tensor's own shape, the strides + // are row-major compact, and `dtype` is the element type those bytes were decoded as, so + // together they describe exactly the initialized elements of the buffer. + let builder = unsafe { builder.data(data) }; + + Ok(builder + .device(DLDevice::CPU) + .dtype(dtype) + .try_build() + .map_err(|err| PyValueError::new_err(err.to_string()))?) } /// The DLPack device this data lives on: `(device_type, device_id)`. Always CPU. - fn __dlpack_device__(&self) -> (i32, i32) { - (DeviceType::Cpu as i32, 0) + fn __dlpack_device__(&self) -> (u32, i32) { + (DLDeviceType::CPU.0, 0) } /// Export as a PEP 3118 buffer: an N-dimensional, typed, read-only, @@ -148,64 +174,52 @@ impl PyTensor { unsafe fn __releasebuffer__(&self, _view: *mut ffi::Py_buffer) {} } -impl TensorLike for PyTensor { - type Error = ZarristaError; - - fn data_ptr(&self) -> *mut c_void { - self.bytes.as_ptr().cast::().cast_mut() - } - - fn memory_layout(&self) -> RowMajorCompactLayout { - let shape = self - .shape() - .iter() - .map(|s| i64::try_from(*s).expect("overflow converting shape to i64")) - .collect(); - RowMajorCompactLayout::new(shape) - } - - fn byte_offset(&self) -> u64 { - 0 - } - - fn device(&self) -> Result { - Ok(Device::CPU) - } - - fn data_type(&self) -> Result { +impl PyTensor { + /// The DLPack element descriptor for this tensor's zarr data type. + fn dlpack_data_type(&self) -> ZarristaResult { use zarrs::array::data_type::*; let dtype = &self.data_type; - if dtype.is::() { - Ok(dlpark::ffi::DataType::BOOL) + let (code, bits) = if dtype.is::() { + (DLDataTypeCode::BOOL, 8) } else if dtype.is::() { - Ok(dlpark::ffi::DataType::I8) + (DLDataTypeCode::INT, 8) } else if dtype.is::() { - Ok(dlpark::ffi::DataType::I16) + (DLDataTypeCode::INT, 16) } else if dtype.is::() { - Ok(dlpark::ffi::DataType::I32) + (DLDataTypeCode::INT, 32) } else if dtype.is::() { - Ok(dlpark::ffi::DataType::I64) + (DLDataTypeCode::INT, 64) } else if dtype.is::() { - Ok(dlpark::ffi::DataType::U8) + (DLDataTypeCode::UINT, 8) } else if dtype.is::() { - Ok(dlpark::ffi::DataType::U16) + (DLDataTypeCode::UINT, 16) } else if dtype.is::() { - Ok(dlpark::ffi::DataType::U32) + (DLDataTypeCode::UINT, 32) } else if dtype.is::() { - Ok(dlpark::ffi::DataType::U64) + (DLDataTypeCode::UINT, 64) } else if dtype.is::() { - Ok(dlpark::ffi::DataType::F16) + (DLDataTypeCode::FLOAT, 16) } else if dtype.is::() { - Ok(dlpark::ffi::DataType::F32) + (DLDataTypeCode::FLOAT, 32) } else if dtype.is::() { - Ok(dlpark::ffi::DataType::F64) + (DLDataTypeCode::FLOAT, 64) } else if dtype.is::() { - Ok(dlpark::ffi::DataType::BF16) + (DLDataTypeCode::BFLOAT, 16) } else { - Err(PyValueError::new_err("Unsupported data type in dlpack").into()) - } + return Err(PyValueError::new_err("Unsupported data type in dlpack").into()); + }; + Ok(DLDataType::scalar(code, bits)) + } +} + +/// Strides, in elements, for a row-major contiguous buffer of the given shape. +fn row_major_compact_strides(shape: &[i64]) -> Vec { + let mut strides = vec![1; shape.len()]; + for axis in (0..shape.len().saturating_sub(1)).rev() { + strides[axis] = strides[axis + 1] * shape[axis + 1]; } + strides } /// Fixed-width data with a validity mask. Skeleton.