dlpark 0.8.0-alpha.1

dlpack Rust binding for Python
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123

use std::ffi::c_void;

use crate::ffi::{DLManagedTensor, DLManagedTensorVersioned, DLPackVersion, DLTensor};
use bitflags::bitflags;

impl Default for DLPackVersion {
    fn default() -> Self {
        Self {
            major: crate::ffi::DLPACK_MAJOR_VERSION,
            minor: crate::ffi::DLPACK_MINOR_VERSION,
        }
    }
}

bitflags! {
    #[repr(transparent)]
    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
    pub struct DlpackFlags: u64 {
        const READ_ONLY = 1 << 0;
        const IS_COPIED = 1 << 1;
        const IS_SUBBYTE_TYPE_PADDED = 1 << 2;
    }
}

pub trait ManagedTensorBase {
    fn from_parts(
        tensor: DLTensor,
        manager_ctx: *mut c_void,
        deleter: Option<unsafe extern "C" fn(self_: *mut Self)>,
    ) -> Self;

    fn tensor(&self) -> &DLTensor;
    fn tensor_mut(&mut self) -> &mut DLTensor;
    fn manager_ctx(&self) -> *mut c_void;
    fn deleter(&self) -> Option<unsafe extern "C" fn(self_: *mut Self)>;
    fn flags(&self) -> DlpackFlags {
        DlpackFlags::empty()
    }

    /// Applies DLPack flags to this managed tensor.
    ///
    /// Only `DLManagedTensorVersioned` carries a `flags` field; the legacy
    /// `DLManagedTensor` has none and inherits the default no-op, so callers
    /// can set flags generically over `M` without knowing which ABI it is.
    fn set_flags(&mut self, _flags: crate::DlpackFlags) {}

    /// Drops a raw managed tensor pointer through its DLPack deleter.
    ///
    /// # Safety
    ///
    /// The caller must ensure that `ptr` is a valid pointer to `Self` and has not been dropped/freed yet.
    unsafe fn drop_raw(ptr: *mut Self) {
        if let Some(deleter) = unsafe { (*ptr).deleter() } {
            unsafe { deleter(ptr) };
        }
    }
}

impl ManagedTensorBase for DLManagedTensor {
    fn from_parts(
        tensor: DLTensor,
        manager_ctx: *mut c_void,
        deleter: Option<unsafe extern "C" fn(self_: *mut Self)>,
    ) -> Self {
        Self {
            dl_tensor: tensor,
            manager_ctx,
            deleter,
        }
    }

    fn tensor(&self) -> &DLTensor {
        &self.dl_tensor
    }
    fn tensor_mut(&mut self) -> &mut DLTensor {
        &mut self.dl_tensor
    }
    fn manager_ctx(&self) -> *mut c_void {
        self.manager_ctx
    }

    fn deleter(&self) -> Option<unsafe extern "C" fn(self_: *mut Self)> {
        self.deleter
    }
}

impl ManagedTensorBase for DLManagedTensorVersioned {
    fn from_parts(
        tensor: DLTensor,
        manager_ctx: *mut c_void,
        deleter: Option<unsafe extern "C" fn(self_: *mut Self)>,
    ) -> Self {
        Self {
            version: DLPackVersion::default(),
            manager_ctx,
            deleter,
            flags: DlpackFlags::empty(),
            dl_tensor: tensor,
        }
    }

    fn tensor(&self) -> &DLTensor {
        &self.dl_tensor
    }
    fn tensor_mut(&mut self) -> &mut DLTensor {
        &mut self.dl_tensor
    }
    fn manager_ctx(&self) -> *mut c_void {
        self.manager_ctx
    }

    fn deleter(&self) -> Option<unsafe extern "C" fn(self_: *mut Self)> {
        self.deleter
    }

    fn set_flags(&mut self, flags: crate::DlpackFlags) {
        self.flags = flags;
    }

    fn flags(&self) -> DlpackFlags {
        self.flags
    }
}