dlpark 0.8.0

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
124
125
126
127
128
129

use super::{
    Error, MismatchedLengthSnafu, NdimOverflowSnafu,
    allocation::{copied_array_layout, copied_slice_layout},
};
use crate::{OpaqueContext, ffi::DLTensor, managed_tensor::ManagedTensorBase};
use snafu::{ResultExt, ensure};
use std::{alloc::Layout, ptr::NonNull};

#[inline]
pub(super) unsafe fn copy_i64_metadata(src: &[i64], dst: *mut i64) {
    unsafe { std::ptr::copy_nonoverlapping(src.as_ptr(), dst, src.len()) };
}

#[inline]
pub(super) unsafe fn copy_i64_metadata_n(src: &[i64], dst: *mut i64, len: usize) {
    unsafe { std::ptr::copy_nonoverlapping(src.as_ptr(), dst, len) };
}

#[inline]
pub(super) unsafe fn try_copy_generic_metadata<T: Copy + TryInto<i64>>(
    src: &[T],
    dst: *mut i64,
) -> Result<(), usize> {
    for (axis, &value) in src.iter().enumerate() {
        let value = value.try_into().map_err(|_| axis)?;
        unsafe { dst.add(axis).write(value) };
    }
    Ok(())
}

#[inline]
pub(super) unsafe fn copy_generic_metadata_unchecked<T: Copy + TryInto<i64>>(
    src: &[T],
    dst: *mut i64,
) {
    for (index, &value) in src.iter().enumerate() {
        let value = match value.try_into() {
            Ok(value) => value,
            Err(_) => unsafe { std::hint::unreachable_unchecked() },
        };
        unsafe { dst.add(index).write(value) };
    }
}

#[inline]
pub(super) unsafe fn initialize<C, M>(
    managed_tensor: NonNull<M>,
    shape: *mut i64,
    strides: *mut i64,
    ndim: i32,
    ctx: C,
    deleter: unsafe extern "C" fn(*mut M),
) -> NonNull<M>
where
    C: OpaqueContext,
    M: ManagedTensorBase,
{
    unsafe {
        managed_tensor.as_ptr().write(M::from_parts(
            DLTensor::from_parts(shape, strides, ndim),
            ctx.into_raw(),
            Some(deleter),
        ));
    }
    managed_tensor
}

#[inline]
pub(super) fn checked_ndim(shape_len: usize, strides_len: usize) -> Result<i32, Error> {
    ensure!(
        shape_len == strides_len,
        MismatchedLengthSnafu {
            shape_len,
            strides_len
        }
    );
    shape_len
        .try_into()
        .context(NdimOverflowSnafu { ndim: shape_len })
}

pub(super) unsafe extern "C" fn drop_copied_array<
    C: OpaqueContext,
    M: ManagedTensorBase,
    const N: usize,
>(
    managed: *mut M,
) {
    if managed.is_null() {
        return;
    }
    unsafe {
        C::drop_raw((*managed).manager_ctx());
        std::ptr::drop_in_place(managed);
        std::alloc::dealloc(managed.cast(), copied_array_layout::<M, N>().0);
    }
}

pub(super) unsafe extern "C" fn drop_copied_slice<C: OpaqueContext, M: ManagedTensorBase>(
    managed: *mut M,
) {
    if managed.is_null() {
        return;
    }
    unsafe {
        let allocation_ndim_offset = copied_slice_layout::<M>(0).1;
        let allocation_ndim = managed
            .cast::<u8>()
            .add(allocation_ndim_offset)
            .cast::<usize>()
            .read();
        C::drop_raw((*managed).manager_ctx());
        std::ptr::drop_in_place(managed);
        std::alloc::dealloc(managed.cast(), copied_slice_layout::<M>(allocation_ndim).0);
    }
}

pub(super) unsafe extern "C" fn drop_borrowed<C: OpaqueContext, M: ManagedTensorBase>(
    managed: *mut M,
) {
    if managed.is_null() {
        return;
    }
    unsafe {
        C::drop_raw((*managed).manager_ctx());
        std::ptr::drop_in_place(managed);
        std::alloc::dealloc(managed.cast(), Layout::new::<M>());
    }
}