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
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
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! {
/// Flags carried by `DLManagedTensorVersioned`.
#[repr(transparent)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct DlpackFlags: u64 {
/// Consumers must not modify the tensor data.
const READ_ONLY = 1 << 0;
/// The export owns an unaliased copy of the tensor data.
const IS_COPIED = 1 << 1;
/// Packed sub-byte data has per-element padding.
const IS_SUBBYTE_TYPE_PADDED = 1 << 2;
}
}
impl DlpackFlags {
/// Whether setting `self` as the new flags, given the tensor's `current`
/// flags, would newly assert [`DlpackFlags::IS_COPIED`] — turn it on
/// when it wasn't already set.
///
/// Turning it on is the risky direction: `ManagedBox::cpu_data_slice_mut`
/// and `array_view_from_dlpack_mut` trust it unconditionally to skip
/// aliasing checks. Leaving an already-set `IS_COPIED` on asserts nothing
/// new, so that case is not flagged.
pub(crate) fn newly_asserts_is_copied(self, current: DlpackFlags) -> bool {
self.contains(DlpackFlags::IS_COPIED) && !current.contains(DlpackFlags::IS_COPIED)
}
}
/// Common operations implemented by the legacy and versioned managed tensor
/// ABIs.
///
/// This trait lets [`crate::Builder`] and [`crate::ManagedBox`] operate
/// generically while preserving the concrete C layout selected by the caller.
pub trait ManagedTensorBase {
/// Constructs a managed tensor from its embedded tensor and ownership
/// fields.
fn from_parts(
tensor: DLTensor,
manager_ctx: *mut c_void,
deleter: Option<unsafe extern "C" fn(self_: *mut Self)>,
) -> Self;
/// Returns the embedded tensor descriptor.
fn tensor(&self) -> &DLTensor;
/// Returns mutable access used while initializing the descriptor.
fn tensor_mut(&mut self) -> &mut DLTensor;
/// Returns the producer-owned opaque context pointer.
fn manager_ctx(&self) -> *mut c_void;
/// Returns the managed tensor deleter.
fn deleter(&self) -> Option<unsafe extern "C" fn(self_: *mut Self)>;
/// Returns versioned flags, or empty flags for the legacy ABI.
fn flags(&self) -> DlpackFlags {
DlpackFlags::empty()
}
/// Applies DLPack flags to this managed tensor verbatim, including
/// [`DlpackFlags::IS_COPIED`].
///
/// 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.
///
/// # Safety
///
/// If `flags` includes `IS_COPIED`, the caller must ensure that no other
/// reference to the tensor's data exists: `ManagedBox::cpu_data_slice_mut`
/// and `array_view_from_dlpack_mut` trust that bit unconditionally and
/// skip aliasing checks accordingly.
unsafe fn set_flags_unchecked(&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
}
unsafe fn set_flags_unchecked(&mut self, flags: crate::DlpackFlags) {
self.flags = flags;
}
fn flags(&self) -> DlpackFlags {
self.flags
}
}