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
use super::{
Error,
allocation::allocate_borrowed,
raw::{checked_ndim, drop_borrowed, initialize},
};
use crate::{OpaqueContext, managed_tensor::ManagedTensorBase};
use std::ptr::NonNull;
/// Borrows fixed-rank metadata and allocates only the managed tensor.
#[derive(Debug, Clone, Copy)]
pub struct BorrowedArray<'a, const N: usize> {
shape: &'a [i64; N],
strides: &'a [i64; N],
}
impl<'a, const N: usize> BorrowedArray<'a, N> {
/// Creates fixed-rank metadata that points to caller-owned arrays.
#[inline]
pub fn new(shape: &'a [i64; N], strides: &'a [i64; N]) -> Self {
Self { shape, strides }
}
/// Allocates a managed tensor that borrows the shape and strides arrays.
///
/// # Safety
///
/// The borrowed arrays must outlive the resulting managed tensor. They
/// must not be mutated through the DLPack `shape`/`strides` pointers while
/// the managed tensor is alive; this API starts from shared Rust
/// references and exposes them through DLPack's mutable pointer fields.
#[inline]
pub unsafe fn allocate<C, M>(self, ctx: C) -> NonNull<M>
where
C: OpaqueContext,
M: ManagedTensorBase,
{
const { assert!(N <= i32::MAX as usize, "N must fit in i32") };
unsafe {
initialize(
allocate_borrowed(),
self.shape.as_ptr().cast_mut(),
self.strides.as_ptr().cast_mut(),
N as i32,
ctx,
drop_borrowed::<C, M>,
)
}
}
}
/// Borrows runtime-rank metadata and allocates only the managed tensor.
#[derive(Debug, Clone, Copy)]
pub struct BorrowedSlice<'a> {
shape: &'a [i64],
strides: &'a [i64],
}
impl<'a> BorrowedSlice<'a> {
/// Creates runtime-rank metadata that points to caller-owned slices.
#[inline]
pub fn new(shape: &'a [i64], strides: &'a [i64]) -> Self {
Self { shape, strides }
}
/// Allocates a managed tensor that borrows the shape and strides slices.
///
/// # Safety
///
/// The borrowed slices must outlive the resulting managed tensor. They
/// must not be mutated through the DLPack `shape`/`strides` pointers while
/// the managed tensor is alive; this API starts from shared Rust
/// references and exposes them through DLPack's mutable pointer fields.
#[inline]
pub unsafe fn allocate<C, M>(self, ctx: C) -> Result<NonNull<M>, Error>
where
C: OpaqueContext,
M: ManagedTensorBase,
{
let ndim = checked_ndim(self.shape.len(), self.strides.len())?;
Ok(unsafe {
initialize(
allocate_borrowed(),
self.shape.as_ptr().cast_mut(),
self.strides.as_ptr().cast_mut(),
ndim,
ctx,
drop_borrowed::<C, M>,
)
})
}
/// Allocates borrowed metadata storage without checking runtime invariants.
///
/// # Safety
///
/// The borrowed slices must outlive the resulting managed tensor, and
/// shape and strides must have the same length with `ndim` fitting in
/// `i32`. They must not be mutated through the DLPack `shape`/`strides`
/// pointers while the managed tensor is alive; this API starts from
/// shared Rust references and exposes them through DLPack's mutable
/// pointer fields.
#[inline]
pub unsafe fn allocate_unchecked<C, M>(self, ctx: C) -> NonNull<M>
where
C: OpaqueContext,
M: ManagedTensorBase,
{
let ndim = self.shape.len();
debug_assert_eq!(self.shape.len(), self.strides.len());
debug_assert!(ndim <= i32::MAX as usize);
unsafe {
initialize(
allocate_borrowed(),
self.shape.as_ptr().cast_mut(),
self.strides.as_ptr().cast_mut(),
ndim as i32,
ctx,
drop_borrowed::<C, M>,
)
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::ffi::DLManagedTensor;
#[test]
fn borrowed_array_allocates_header_only_and_reuses_pointers() {
let shape = [2i64, 3];
let strides = [3i64, 1];
let managed = unsafe {
BorrowedArray::new(&shape, &strides).allocate::<_, DLManagedTensor>(Box::new(()))
};
let tensor = unsafe { managed.as_ref().tensor() };
assert_eq!(tensor.ndim, 2);
assert_eq!(tensor.shape, shape.as_ptr().cast_mut());
assert_eq!(tensor.strides, strides.as_ptr().cast_mut());
unsafe { DLManagedTensor::drop_raw(managed.as_ptr()) };
}
}