spatialrust_tensor/
dlpack.rs

1//! Audited DLPack major-version 1 CPU ownership boundary.
2
3use std::{
4    ffi::c_void,
5    mem::ManuallyDrop,
6    ptr::{self, NonNull},
7    slice,
8};
9
10use crate::{
11    DataType, DataTypeCode, Device, DeviceKind, TensorBuffer, TensorDescriptor, TensorError,
12    TensorStorage, TensorView,
13};
14
15/// DLPack ABI major version supported by this boundary.
16pub const DLPACK_MAJOR: u32 = 1;
17/// Baseline DLPack minor version emitted by this boundary.
18pub const DLPACK_MINOR: u32 = 0;
19
20const FLAG_READ_ONLY: u64 = 1;
21const MAX_RANK: usize = 64;
22
23/// Errors raised while validating a DLPack managed tensor.
24#[derive(Debug, thiserror::Error)]
25pub enum DlpackError {
26    /// A null managed-tensor pointer was supplied.
27    #[error("DLPack managed tensor pointer is null")]
28    NullManagedTensor,
29    /// The ABI major version is incompatible.
30    #[error("unsupported DLPack ABI version {major}.{minor}; expected major {DLPACK_MAJOR}")]
31    IncompatibleVersion {
32        /// Producer major version.
33        major: u32,
34        /// Producer minor version.
35        minor: u32,
36    },
37    /// Rank is negative or exceeds the defensive limit.
38    #[error("invalid DLPack tensor rank {0}")]
39    InvalidRank(i32),
40    /// A ranked tensor has no shape pointer.
41    #[error("DLPack tensor shape pointer is null for non-zero rank")]
42    NullShape,
43    /// A dimension is negative or cannot be represented by `usize`.
44    #[error("invalid DLPack shape dimension {0}")]
45    InvalidDimension(i64),
46    /// The data type code is unsupported or not byte-addressable.
47    #[error("unsupported DLPack dtype code={code}, bits={bits}, lanes={lanes}")]
48    UnsupportedDataType {
49        /// Raw DLPack code.
50        code: u8,
51        /// Bits per lane.
52        bits: u8,
53        /// Vector lanes.
54        lanes: u16,
55    },
56    /// The device code is unknown to this DLPack minor implementation.
57    #[error("unsupported DLPack device type {0}")]
58    UnsupportedDevice(i32),
59    /// A non-empty host tensor has no data pointer.
60    #[error("non-empty DLPack host tensor has a null data pointer")]
61    NullData,
62    /// A DLPack integer does not fit the host representation.
63    #[error("DLPack metadata does not fit the host integer representation")]
64    IntegerConversion,
65    /// The decoded tensor layout is invalid.
66    #[error(transparent)]
67    Tensor(#[from] TensorError),
68}
69
70#[repr(C)]
71#[derive(Clone, Copy)]
72struct RawVersion {
73    major: u32,
74    minor: u32,
75}
76
77#[repr(C)]
78#[derive(Clone, Copy)]
79struct RawDevice {
80    device_type: i32,
81    device_id: i32,
82}
83
84#[repr(C)]
85#[derive(Clone, Copy)]
86struct RawDataType {
87    code: u8,
88    bits: u8,
89    lanes: u16,
90}
91
92#[repr(C)]
93struct RawTensor {
94    data: *mut c_void,
95    device: RawDevice,
96    ndim: i32,
97    dtype: RawDataType,
98    shape: *mut i64,
99    strides: *mut i64,
100    byte_offset: u64,
101}
102
103#[repr(C)]
104struct RawManagedTensorVersioned {
105    version: RawVersion,
106    manager_ctx: *mut c_void,
107    deleter: Option<unsafe extern "C" fn(*mut RawManagedTensorVersioned)>,
108    flags: u64,
109    dl_tensor: RawTensor,
110}
111
112#[repr(C)]
113struct RawManagedTensorLegacy {
114    dl_tensor: RawTensor,
115    manager_ctx: *mut c_void,
116    deleter: Option<unsafe extern "C" fn(*mut RawManagedTensorLegacy)>,
117}
118
119struct ExportContext {
120    _allocation: TensorStorage,
121    _shape: Box<[i64]>,
122    _strides: Box<[i64]>,
123}
124
125/// Owner for a versioned DLPack managed tensor exported without copying CPU data.
126///
127/// Dropping this value calls the DLPack deleter. [`Self::into_raw`] transfers
128/// that responsibility to a capsule or another consumer.
129pub struct DlpackExport {
130    raw: NonNull<RawManagedTensorVersioned>,
131}
132
133/// Owner for a legacy DLPack managed tensor exported without copying CPU data.
134///
135/// This exists for consumers implementing the pre-1.0 capsule protocol.
136pub struct DlpackLegacyExport {
137    raw: NonNull<RawManagedTensorLegacy>,
138}
139
140/// Calls the producer deleter for a raw pointer previously returned by
141/// [`DlpackExport::into_raw`].
142///
143/// # Safety
144///
145/// `raw` must still carry exclusive deleter responsibility for a live
146/// `DLManagedTensorVersioned*`. It must not be used after this call.
147pub unsafe fn release_dlpack_raw(raw: *mut c_void) {
148    // SAFETY: forwarded from the public ownership contract above.
149    unsafe { call_deleter(raw.cast()) };
150}
151
152/// Calls the producer deleter for a legacy raw pointer.
153///
154/// # Safety
155///
156/// `raw` must carry exclusive deleter responsibility for a live
157/// legacy `DLManagedTensor*`. It must not be used after this call.
158pub unsafe fn release_dlpack_legacy_raw(raw: *mut c_void) {
159    // SAFETY: forwarded from the public ownership contract above.
160    unsafe { call_legacy_deleter(raw.cast()) };
161}
162
163impl DlpackExport {
164    /// Shares an owned CPU tensor allocation with a DLPack consumer without copying.
165    pub fn from_tensor(tensor: &TensorBuffer) -> Result<Self, DlpackError> {
166        let descriptor = tensor.descriptor();
167        if !descriptor.device().is_host_accessible() {
168            return Err(TensorError::DeviceNotHostAccessible(descriptor.device()).into());
169        }
170        let shape = descriptor
171            .shape()
172            .iter()
173            .map(|&dimension| i64::try_from(dimension).map_err(|_| DlpackError::IntegerConversion))
174            .collect::<Result<Vec<_>, _>>()?
175            .into_boxed_slice();
176        let strides = match descriptor.strides() {
177            Some(values) => values
178                .iter()
179                .map(|&stride| i64::try_from(stride).map_err(|_| DlpackError::IntegerConversion))
180                .collect::<Result<Vec<_>, _>>()?,
181            None => compact_strides(descriptor.shape())?,
182        }
183        .into_boxed_slice();
184        let ndim =
185            i32::try_from(descriptor.shape().len()).map_err(|_| DlpackError::IntegerConversion)?;
186        let byte_offset =
187            u64::try_from(descriptor.byte_offset()).map_err(|_| DlpackError::IntegerConversion)?;
188        let allocation = tensor.shared_allocation();
189        let data = if allocation.is_empty() {
190            ptr::null_mut()
191        } else {
192            allocation.as_ptr().cast_mut().cast()
193        };
194        let shape_ptr = if shape.is_empty() { ptr::null_mut() } else { shape.as_ptr().cast_mut() };
195        let strides_ptr =
196            if strides.is_empty() { ptr::null_mut() } else { strides.as_ptr().cast_mut() };
197        let context =
198            Box::new(ExportContext { _allocation: allocation, _shape: shape, _strides: strides });
199        let manager_ctx = Box::into_raw(context).cast();
200        let raw = Box::new(RawManagedTensorVersioned {
201            version: RawVersion { major: DLPACK_MAJOR, minor: DLPACK_MINOR },
202            manager_ctx,
203            deleter: Some(delete_export),
204            flags: FLAG_READ_ONLY,
205            dl_tensor: RawTensor {
206                data,
207                device: encode_device(descriptor.device()),
208                ndim,
209                dtype: encode_dtype(descriptor.dtype()),
210                shape: shape_ptr,
211                strides: strides_ptr,
212                byte_offset,
213            },
214        });
215        Ok(Self { raw: NonNull::from(Box::leak(raw)) })
216    }
217
218    /// Returns the opaque managed-tensor pointer without transferring ownership.
219    pub fn as_raw(&self) -> *mut c_void {
220        self.raw.as_ptr().cast()
221    }
222
223    /// Transfers deleter responsibility to an external DLPack consumer.
224    pub fn into_raw(self) -> *mut c_void {
225        let this = ManuallyDrop::new(self);
226        this.raw.as_ptr().cast()
227    }
228}
229
230impl Drop for DlpackExport {
231    fn drop(&mut self) {
232        // SAFETY: `DlpackExport` uniquely owns deleter responsibility until `into_raw`.
233        unsafe { call_deleter(self.raw.as_ptr()) };
234    }
235}
236
237impl DlpackLegacyExport {
238    /// Shares an owned CPU tensor allocation with a legacy DLPack consumer.
239    pub fn from_tensor(tensor: &TensorBuffer) -> Result<Self, DlpackError> {
240        let (context, tensor) = build_export(tensor)?;
241        let raw = Box::new(RawManagedTensorLegacy {
242            dl_tensor: tensor,
243            manager_ctx: Box::into_raw(context).cast(),
244            deleter: Some(delete_legacy_export),
245        });
246        Ok(Self { raw: NonNull::from(Box::leak(raw)) })
247    }
248
249    /// Transfers deleter responsibility to an external DLPack consumer.
250    pub fn into_raw(self) -> *mut c_void {
251        let this = ManuallyDrop::new(self);
252        this.raw.as_ptr().cast()
253    }
254}
255
256impl Drop for DlpackLegacyExport {
257    fn drop(&mut self) {
258        // SAFETY: this export uniquely owns deleter responsibility.
259        unsafe { call_legacy_deleter(self.raw.as_ptr()) };
260    }
261}
262
263/// Validated owner of a DLPack producer's managed tensor.
264#[derive(Debug)]
265pub struct DlpackImport {
266    raw: ImportedRaw,
267    descriptor: TensorDescriptor,
268    allocation_len: usize,
269    data: *const u8,
270    version: (u32, u32),
271    flags: u64,
272}
273
274#[derive(Debug)]
275enum ImportedRaw {
276    Versioned(NonNull<RawManagedTensorVersioned>),
277    Legacy(NonNull<RawManagedTensorLegacy>),
278}
279
280impl DlpackImport {
281    /// Takes deleter ownership of a DLPack managed tensor and validates its host view.
282    ///
283    /// # Safety
284    ///
285    /// `raw` must be a live, exclusively transferred `DLManagedTensorVersioned*`
286    /// produced according to DLPack. The caller must not use or delete it after
287    /// this call, whether validation succeeds or fails.
288    pub unsafe fn from_raw(raw: *mut c_void) -> Result<Self, DlpackError> {
289        let raw = NonNull::new(raw.cast::<RawManagedTensorVersioned>())
290            .ok_or(DlpackError::NullManagedTensor)?;
291        let guard = IncomingGuard { raw: Some(raw) };
292
293        // SAFETY: the caller promises a live versioned header. DLPack guarantees
294        // version and deleter positions remain accessible across major mismatch.
295        let managed = unsafe { raw.as_ref() };
296        let version = (managed.version.major, managed.version.minor);
297        if version.0 != DLPACK_MAJOR {
298            return Err(DlpackError::IncompatibleVersion { major: version.0, minor: version.1 });
299        }
300        let tensor = &managed.dl_tensor;
301        if tensor.ndim < 0 || tensor.ndim as usize > MAX_RANK {
302            return Err(DlpackError::InvalidRank(tensor.ndim));
303        }
304        let rank = tensor.ndim as usize;
305        if rank != 0 && tensor.shape.is_null() {
306            return Err(DlpackError::NullShape);
307        }
308        let shape_values = if rank == 0 {
309            &[][..]
310        } else {
311            // SAFETY: producer contract supplies `ndim` readable shape entries.
312            unsafe { slice::from_raw_parts(tensor.shape, rank) }
313        };
314        let shape = shape_values
315            .iter()
316            .map(|&dimension| {
317                usize::try_from(dimension).map_err(|_| DlpackError::InvalidDimension(dimension))
318            })
319            .collect::<Result<Vec<_>, _>>()?;
320        let strides = if tensor.strides.is_null() {
321            None
322        } else {
323            // SAFETY: producer contract supplies `ndim` readable stride entries.
324            let values = unsafe { slice::from_raw_parts(tensor.strides, rank) };
325            Some(
326                values
327                    .iter()
328                    .map(|&stride| {
329                        isize::try_from(stride).map_err(|_| DlpackError::IntegerConversion)
330                    })
331                    .collect::<Result<Vec<_>, _>>()?,
332            )
333        };
334        let dtype = decode_dtype(tensor.dtype)?;
335        let device = decode_device(tensor.device)?;
336        let byte_offset =
337            usize::try_from(tensor.byte_offset).map_err(|_| DlpackError::IntegerConversion)?;
338        let descriptor = match strides {
339            Some(strides) => {
340                TensorDescriptor::try_strided(dtype, shape, strides, byte_offset, device)?
341            }
342            None => {
343                let mut descriptor = TensorDescriptor::contiguous(dtype, shape, device);
344                if byte_offset != 0 {
345                    descriptor = TensorDescriptor::try_strided(
346                        dtype,
347                        descriptor.shape().to_vec(),
348                        compact_strides_isize(descriptor.shape())?,
349                        byte_offset,
350                        device,
351                    )?;
352                }
353                descriptor
354            }
355        };
356        let range = descriptor.required_byte_range()?;
357        if range.end != 0 && tensor.data.is_null() {
358            return Err(DlpackError::NullData);
359        }
360        let raw = ImportedRaw::Versioned(guard.disarm());
361        Ok(Self {
362            raw,
363            descriptor,
364            allocation_len: range.end,
365            data: tensor.data.cast(),
366            version,
367            flags: managed.flags,
368        })
369    }
370
371    /// Takes deleter ownership of a legacy DLPack managed tensor.
372    ///
373    /// # Safety
374    ///
375    /// `raw` must be a live, exclusively transferred legacy `DLManagedTensor*`.
376    pub unsafe fn from_legacy_raw(raw: *mut c_void) -> Result<Self, DlpackError> {
377        let raw = NonNull::new(raw.cast::<RawManagedTensorLegacy>())
378            .ok_or(DlpackError::NullManagedTensor)?;
379        let guard = IncomingLegacyGuard { raw: Some(raw) };
380        // SAFETY: the caller promises a live legacy managed tensor.
381        let managed = unsafe { raw.as_ref() };
382        let (descriptor, allocation_len, data) = validate_tensor(&managed.dl_tensor)?;
383        let raw = ImportedRaw::Legacy(guard.disarm());
384        Ok(Self { raw, descriptor, allocation_len, data, version: (0, 0), flags: 0 })
385    }
386
387    /// Returns producer ABI major and minor versions.
388    pub const fn version(&self) -> (u32, u32) {
389        self.version
390    }
391
392    /// Returns raw DLPack flags.
393    pub const fn flags(&self) -> u64 {
394        self.flags
395    }
396
397    /// Returns validated tensor metadata.
398    pub const fn descriptor(&self) -> &TensorDescriptor {
399        &self.descriptor
400    }
401
402    /// Borrows the imported host allocation without copying it.
403    pub fn view(&self) -> Result<TensorView<'_>, DlpackError> {
404        let bytes = if self.allocation_len == 0 {
405            &[][..]
406        } else {
407            // SAFETY: `from_raw` validated the DLPack producer contract and the
408            // managed tensor remains alive until this owner is dropped.
409            unsafe { slice::from_raw_parts(self.data, self.allocation_len) }
410        };
411        Ok(TensorView::try_new(bytes, self.descriptor.clone())?)
412    }
413}
414
415impl Drop for DlpackImport {
416    fn drop(&mut self) {
417        // SAFETY: this owner received exclusive deleter responsibility.
418        unsafe {
419            match self.raw {
420                ImportedRaw::Versioned(raw) => call_deleter(raw.as_ptr()),
421                ImportedRaw::Legacy(raw) => call_legacy_deleter(raw.as_ptr()),
422            }
423        };
424    }
425}
426
427struct IncomingGuard {
428    raw: Option<NonNull<RawManagedTensorVersioned>>,
429}
430
431impl IncomingGuard {
432    fn disarm(mut self) -> NonNull<RawManagedTensorVersioned> {
433        self.raw.take().expect("incoming pointer is present")
434    }
435}
436
437impl Drop for IncomingGuard {
438    fn drop(&mut self) {
439        if let Some(raw) = self.raw {
440            // SAFETY: the guard owns the transferred pointer on all error paths.
441            unsafe { call_deleter(raw.as_ptr()) };
442        }
443    }
444}
445
446struct IncomingLegacyGuard {
447    raw: Option<NonNull<RawManagedTensorLegacy>>,
448}
449
450impl IncomingLegacyGuard {
451    fn disarm(mut self) -> NonNull<RawManagedTensorLegacy> {
452        self.raw.take().expect("incoming pointer is present")
453    }
454}
455
456impl Drop for IncomingLegacyGuard {
457    fn drop(&mut self) {
458        if let Some(raw) = self.raw {
459            // SAFETY: the guard owns the transferred pointer on error paths.
460            unsafe { call_legacy_deleter(raw.as_ptr()) };
461        }
462    }
463}
464
465unsafe extern "C" fn delete_export(raw: *mut RawManagedTensorVersioned) {
466    if raw.is_null() {
467        return;
468    }
469    // SAFETY: this function is installed only for allocations built by
470    // `DlpackExport::from_tensor` and is called exactly once by ownership contract.
471    let managed = unsafe { Box::from_raw(raw) };
472    if !managed.manager_ctx.is_null() {
473        // SAFETY: manager_ctx was created with Box::into_raw for ExportContext.
474        drop(unsafe { Box::from_raw(managed.manager_ctx.cast::<ExportContext>()) });
475    }
476}
477
478unsafe extern "C" fn delete_legacy_export(raw: *mut RawManagedTensorLegacy) {
479    if raw.is_null() {
480        return;
481    }
482    // SAFETY: installed only for allocations built by `DlpackLegacyExport`.
483    let managed = unsafe { Box::from_raw(raw) };
484    if !managed.manager_ctx.is_null() {
485        // SAFETY: manager_ctx was created with Box::into_raw for ExportContext.
486        drop(unsafe { Box::from_raw(managed.manager_ctx.cast::<ExportContext>()) });
487    }
488}
489
490unsafe fn call_deleter(raw: *mut RawManagedTensorVersioned) {
491    if raw.is_null() {
492        return;
493    }
494    // SAFETY: caller owns a live managed-tensor pointer.
495    if let Some(deleter) = unsafe { (*raw).deleter } {
496        // SAFETY: deleter belongs to this exact managed tensor.
497        unsafe { deleter(raw) };
498    }
499}
500
501unsafe fn call_legacy_deleter(raw: *mut RawManagedTensorLegacy) {
502    if raw.is_null() {
503        return;
504    }
505    // SAFETY: caller owns a live legacy managed-tensor pointer.
506    if let Some(deleter) = unsafe { (*raw).deleter } {
507        // SAFETY: deleter belongs to this exact managed tensor.
508        unsafe { deleter(raw) };
509    }
510}
511
512fn build_export(tensor: &TensorBuffer) -> Result<(Box<ExportContext>, RawTensor), DlpackError> {
513    let descriptor = tensor.descriptor();
514    if !descriptor.device().is_host_accessible() {
515        return Err(TensorError::DeviceNotHostAccessible(descriptor.device()).into());
516    }
517    let shape = descriptor
518        .shape()
519        .iter()
520        .map(|&dimension| i64::try_from(dimension).map_err(|_| DlpackError::IntegerConversion))
521        .collect::<Result<Vec<_>, _>>()?
522        .into_boxed_slice();
523    let strides = match descriptor.strides() {
524        Some(values) => values
525            .iter()
526            .map(|&stride| i64::try_from(stride).map_err(|_| DlpackError::IntegerConversion))
527            .collect::<Result<Vec<_>, _>>()?,
528        None => compact_strides(descriptor.shape())?,
529    }
530    .into_boxed_slice();
531    let ndim =
532        i32::try_from(descriptor.shape().len()).map_err(|_| DlpackError::IntegerConversion)?;
533    let byte_offset =
534        u64::try_from(descriptor.byte_offset()).map_err(|_| DlpackError::IntegerConversion)?;
535    let allocation = tensor.shared_allocation();
536    let data =
537        if allocation.is_empty() { ptr::null_mut() } else { allocation.as_ptr().cast_mut().cast() };
538    let shape_ptr = if shape.is_empty() { ptr::null_mut() } else { shape.as_ptr().cast_mut() };
539    let strides_ptr =
540        if strides.is_empty() { ptr::null_mut() } else { strides.as_ptr().cast_mut() };
541    let context =
542        Box::new(ExportContext { _allocation: allocation, _shape: shape, _strides: strides });
543    let tensor = RawTensor {
544        data,
545        device: encode_device(descriptor.device()),
546        ndim,
547        dtype: encode_dtype(descriptor.dtype()),
548        shape: shape_ptr,
549        strides: strides_ptr,
550        byte_offset,
551    };
552    Ok((context, tensor))
553}
554
555fn validate_tensor(
556    tensor: &RawTensor,
557) -> Result<(TensorDescriptor, usize, *const u8), DlpackError> {
558    if tensor.ndim < 0 || tensor.ndim as usize > MAX_RANK {
559        return Err(DlpackError::InvalidRank(tensor.ndim));
560    }
561    let rank = tensor.ndim as usize;
562    if rank != 0 && tensor.shape.is_null() {
563        return Err(DlpackError::NullShape);
564    }
565    let shape_values = if rank == 0 {
566        &[][..]
567    } else {
568        // SAFETY: the producer contract supplies `ndim` readable shape entries.
569        unsafe { slice::from_raw_parts(tensor.shape, rank) }
570    };
571    let shape = shape_values
572        .iter()
573        .map(|&dimension| {
574            usize::try_from(dimension).map_err(|_| DlpackError::InvalidDimension(dimension))
575        })
576        .collect::<Result<Vec<_>, _>>()?;
577    let strides = if tensor.strides.is_null() {
578        None
579    } else {
580        // SAFETY: the producer contract supplies `ndim` readable stride entries.
581        let values = unsafe { slice::from_raw_parts(tensor.strides, rank) };
582        Some(
583            values
584                .iter()
585                .map(|&stride| isize::try_from(stride).map_err(|_| DlpackError::IntegerConversion))
586                .collect::<Result<Vec<_>, _>>()?,
587        )
588    };
589    let dtype = decode_dtype(tensor.dtype)?;
590    let device = decode_device(tensor.device)?;
591    let byte_offset =
592        usize::try_from(tensor.byte_offset).map_err(|_| DlpackError::IntegerConversion)?;
593    let descriptor = match strides {
594        Some(strides) => TensorDescriptor::try_strided(dtype, shape, strides, byte_offset, device)?,
595        None => {
596            let mut descriptor = TensorDescriptor::contiguous(dtype, shape, device);
597            if byte_offset != 0 {
598                descriptor = TensorDescriptor::try_strided(
599                    dtype,
600                    descriptor.shape().to_vec(),
601                    compact_strides_isize(descriptor.shape())?,
602                    byte_offset,
603                    device,
604                )?;
605            }
606            descriptor
607        }
608    };
609    let range = descriptor.required_byte_range()?;
610    if range.end != 0 && tensor.data.is_null() {
611        return Err(DlpackError::NullData);
612    }
613    Ok((descriptor, range.end, tensor.data.cast()))
614}
615
616fn compact_strides(shape: &[usize]) -> Result<Vec<i64>, DlpackError> {
617    let mut output = vec![0; shape.len()];
618    let mut stride = 1_i64;
619    for (index, &dimension) in shape.iter().enumerate().rev() {
620        output[index] = stride;
621        stride = stride
622            .checked_mul(
623                i64::try_from(dimension.max(1)).map_err(|_| DlpackError::IntegerConversion)?,
624            )
625            .ok_or(DlpackError::IntegerConversion)?;
626    }
627    Ok(output)
628}
629
630fn compact_strides_isize(shape: &[usize]) -> Result<Vec<isize>, DlpackError> {
631    compact_strides(shape)?
632        .into_iter()
633        .map(|stride| isize::try_from(stride).map_err(|_| DlpackError::IntegerConversion))
634        .collect()
635}
636
637fn encode_dtype(dtype: DataType) -> RawDataType {
638    RawDataType { code: dtype.code() as u8, bits: dtype.bits(), lanes: dtype.lanes() }
639}
640
641fn decode_dtype(raw: RawDataType) -> Result<DataType, DlpackError> {
642    let code = match raw.code {
643        0 => DataTypeCode::Int,
644        1 => DataTypeCode::UInt,
645        2 => DataTypeCode::Float,
646        4 => DataTypeCode::BFloat,
647        5 => DataTypeCode::Complex,
648        6 => DataTypeCode::Bool,
649        _ => {
650            return Err(DlpackError::UnsupportedDataType {
651                code: raw.code,
652                bits: raw.bits,
653                lanes: raw.lanes,
654            })
655        }
656    };
657    DataType::try_new(code, raw.bits, raw.lanes).map_err(|_| DlpackError::UnsupportedDataType {
658        code: raw.code,
659        bits: raw.bits,
660        lanes: raw.lanes,
661    })
662}
663
664fn encode_device(device: Device) -> RawDevice {
665    let device_type = match device.kind {
666        DeviceKind::Cpu => 1,
667        DeviceKind::Cuda => 2,
668        DeviceKind::CudaHost => 3,
669        DeviceKind::OpenCl => 4,
670        DeviceKind::Vulkan => 7,
671        DeviceKind::Metal => 8,
672        DeviceKind::Vpi => 9,
673        DeviceKind::Rocm => 10,
674        DeviceKind::RocmHost => 11,
675        DeviceKind::External => 12,
676        DeviceKind::CudaManaged => 13,
677        DeviceKind::OneApi => 14,
678        DeviceKind::WebGpu => 15,
679        DeviceKind::Hexagon => 16,
680        DeviceKind::Maia => 17,
681        DeviceKind::Trainium => 18,
682        DeviceKind::Tpu => 19,
683        DeviceKind::TpuHost => 20,
684    };
685    RawDevice { device_type, device_id: device.id }
686}
687
688fn decode_device(raw: RawDevice) -> Result<Device, DlpackError> {
689    let kind = match raw.device_type {
690        1 => DeviceKind::Cpu,
691        2 => DeviceKind::Cuda,
692        3 => DeviceKind::CudaHost,
693        4 => DeviceKind::OpenCl,
694        7 => DeviceKind::Vulkan,
695        8 => DeviceKind::Metal,
696        9 => DeviceKind::Vpi,
697        10 => DeviceKind::Rocm,
698        11 => DeviceKind::RocmHost,
699        12 => DeviceKind::External,
700        13 => DeviceKind::CudaManaged,
701        14 => DeviceKind::OneApi,
702        15 => DeviceKind::WebGpu,
703        16 => DeviceKind::Hexagon,
704        17 => DeviceKind::Maia,
705        18 => DeviceKind::Trainium,
706        19 => DeviceKind::Tpu,
707        20 => DeviceKind::TpuHost,
708        other => return Err(DlpackError::UnsupportedDevice(other)),
709    };
710    Ok(Device { kind, id: raw.device_id })
711}
712
713#[cfg(test)]
714mod tests {
715    use super::{
716        DlpackError, DlpackExport, DlpackImport, DlpackLegacyExport, DLPACK_MAJOR, DLPACK_MINOR,
717    };
718    use crate::{DataType, Device, TensorBuffer, TensorDescriptor};
719
720    #[test]
721    fn owned_cpu_roundtrip_is_zero_copy_and_versioned() {
722        let tensor = TensorBuffer::try_new(
723            (0_u8..24).collect(),
724            TensorDescriptor::contiguous(DataType::F32, vec![2, 3], Device::CPU),
725        )
726        .unwrap();
727        let original = tensor.allocation_bytes().as_ptr();
728        let allocation = tensor.shared_allocation();
729        let export = DlpackExport::from_tensor(&tensor).unwrap();
730        assert_eq!(allocation.strong_count(), 3);
731        // SAFETY: into_raw transfers the live export exactly once.
732        let imported = unsafe { DlpackImport::from_raw(export.into_raw()) }.unwrap();
733        assert_eq!(imported.version(), (DLPACK_MAJOR, DLPACK_MINOR));
734        let view = imported.view().unwrap();
735        assert_eq!(view.descriptor().shape(), &[2, 3]);
736        assert_eq!(view.allocation_bytes().as_ptr(), original);
737        drop(imported);
738        assert_eq!(allocation.strong_count(), 2);
739    }
740
741    #[test]
742    fn negative_stride_and_byte_offset_roundtrip() {
743        let descriptor =
744            TensorDescriptor::try_strided(DataType::U8, vec![4], vec![-1], 3, Device::CPU).unwrap();
745        let tensor = TensorBuffer::try_new(vec![10, 20, 30, 40], descriptor).unwrap();
746        let export = DlpackExport::from_tensor(&tensor).unwrap();
747        // SAFETY: the export pointer is transferred exactly once.
748        let imported = unsafe { DlpackImport::from_raw(export.into_raw()) }.unwrap();
749        let view = imported.view().unwrap();
750        assert_eq!(view.descriptor().strides(), Some(&[-1][..]));
751        assert_eq!(view.descriptor().byte_offset(), 3);
752        assert_eq!(view.allocation_bytes(), &[10, 20, 30, 40]);
753    }
754
755    #[test]
756    fn legacy_cpu_roundtrip_is_zero_copy() {
757        let tensor = TensorBuffer::try_new(
758            (0_u8..12).collect(),
759            TensorDescriptor::contiguous(DataType::U8, vec![3, 4], Device::CPU),
760        )
761        .unwrap();
762        let original = tensor.allocation_bytes().as_ptr();
763        let allocation = tensor.shared_allocation();
764        let export = DlpackLegacyExport::from_tensor(&tensor).unwrap();
765        assert_eq!(allocation.strong_count(), 3);
766        // SAFETY: into_raw transfers the live legacy export exactly once.
767        let imported = unsafe { DlpackImport::from_legacy_raw(export.into_raw()) }.unwrap();
768        assert_eq!(imported.version(), (0, 0));
769        assert_eq!(imported.flags(), 0);
770        let view = imported.view().unwrap();
771        assert_eq!(view.descriptor().shape(), &[3, 4]);
772        assert_eq!(view.allocation_bytes().as_ptr(), original);
773        drop(imported);
774        assert_eq!(allocation.strong_count(), 2);
775    }
776
777    #[test]
778    fn major_mismatch_is_rejected_and_deleted() {
779        let tensor = TensorBuffer::try_new(
780            vec![1],
781            TensorDescriptor::contiguous(DataType::U8, vec![1], Device::CPU),
782        )
783        .unwrap();
784        let allocation = tensor.shared_allocation();
785        let export = DlpackExport::from_tensor(&tensor).unwrap();
786        let raw = export.into_raw().cast::<super::RawManagedTensorVersioned>();
787        // SAFETY: test exclusively owns this live export and only changes its version header.
788        unsafe { (*raw).version.major = 99 };
789        // SAFETY: the pointer is still live and transferred exactly once.
790        let error = unsafe { DlpackImport::from_raw(raw.cast()) }.unwrap_err();
791        assert!(matches!(error, DlpackError::IncompatibleVersion { major: 99, .. }));
792        assert_eq!(allocation.strong_count(), 2);
793    }
794
795    #[test]
796    fn malformed_dtype_is_rejected_and_deleted() {
797        let tensor = TensorBuffer::try_new(
798            vec![1],
799            TensorDescriptor::contiguous(DataType::U8, vec![1], Device::CPU),
800        )
801        .unwrap();
802        let allocation = tensor.shared_allocation();
803        let raw = DlpackExport::from_tensor(&tensor)
804            .unwrap()
805            .into_raw()
806            .cast::<super::RawManagedTensorVersioned>();
807        // SAFETY: test owns the live export and corrupts only metadata under test.
808        unsafe { (*raw).dl_tensor.dtype.code = 255 };
809        // SAFETY: the corrupted but live pointer is transferred exactly once.
810        let error = unsafe { DlpackImport::from_raw(raw.cast()) }.unwrap_err();
811        assert!(matches!(error, DlpackError::UnsupportedDataType { code: 255, .. }));
812        assert_eq!(allocation.strong_count(), 2);
813    }
814
815    #[test]
816    fn null_nonempty_data_is_rejected_and_deleted() {
817        let tensor = TensorBuffer::try_new(
818            vec![1, 2],
819            TensorDescriptor::contiguous(DataType::U8, vec![2], Device::CPU),
820        )
821        .unwrap();
822        let allocation = tensor.shared_allocation();
823        let raw = DlpackExport::from_tensor(&tensor)
824            .unwrap()
825            .into_raw()
826            .cast::<super::RawManagedTensorVersioned>();
827        // SAFETY: test owns the live export and corrupts only metadata under test.
828        unsafe { (*raw).dl_tensor.data = std::ptr::null_mut() };
829        // SAFETY: the corrupted but live pointer is transferred exactly once.
830        let error = unsafe { DlpackImport::from_raw(raw.cast()) }.unwrap_err();
831        assert!(matches!(error, DlpackError::NullData));
832        assert_eq!(allocation.strong_count(), 2);
833    }
834
835    #[test]
836    fn negative_shape_is_rejected_and_deleted() {
837        let tensor = TensorBuffer::try_new(
838            vec![1],
839            TensorDescriptor::contiguous(DataType::U8, vec![1], Device::CPU),
840        )
841        .unwrap();
842        let allocation = tensor.shared_allocation();
843        let raw = DlpackExport::from_tensor(&tensor)
844            .unwrap()
845            .into_raw()
846            .cast::<super::RawManagedTensorVersioned>();
847        // SAFETY: exported rank is one and shape points to one writable context entry.
848        unsafe { *(*raw).dl_tensor.shape = -1 };
849        // SAFETY: the corrupted but live pointer is transferred exactly once.
850        let error = unsafe { DlpackImport::from_raw(raw.cast()) }.unwrap_err();
851        assert!(matches!(error, DlpackError::InvalidDimension(-1)));
852        assert_eq!(allocation.strong_count(), 2);
853    }
854}
spatialrust_tensor/dlpack.rs

spatialrust_tensor/
dlpack.rs
