dlpark 0.9.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

use super::{DLTENSOR, DLTENSOR_VERSIONED, USED_DLTENSOR, USED_DLTENSOR_VERSIONED};
use crate::{
    Managed,
    ffi::{DLManagedTensor, DLManagedTensorVersioned},
};
use pyo3::{Bound, PyAny, PyErr, conversion::IntoPyObject, exceptions::PyBufferError};
use std::ffi::CStr;

fn fetch_python_error() -> PyErr {
    unsafe { PyErr::fetch(pyo3::Python::assume_attached()) }
}

unsafe extern "C" fn dlpack_capsule_deleter(capsule: *mut pyo3::ffi::PyObject) {
    unsafe {
        if pyo3::ffi::PyCapsule_IsValid(capsule, USED_DLTENSOR.as_ptr()) == 1 {
            return;
        }

        let ptr = pyo3::ffi::PyCapsule_GetPointer(capsule, DLTENSOR.as_ptr());

        if ptr.is_null() {
            pyo3::ffi::PyErr_WriteUnraisable(capsule);
            return;
        }

        let _ = Managed::<DLManagedTensor>::from_raw_unchecked(ptr as *mut _);
    }
}

unsafe extern "C" fn dlpack_capsule_deleter_versioned(capsule: *mut pyo3::ffi::PyObject) {
    unsafe {
        if pyo3::ffi::PyCapsule_IsValid(capsule, USED_DLTENSOR_VERSIONED.as_ptr()) == 1 {
            return;
        }

        let ptr = pyo3::ffi::PyCapsule_GetPointer(capsule, DLTENSOR_VERSIONED.as_ptr());

        if ptr.is_null() {
            pyo3::ffi::PyErr_WriteUnraisable(capsule);
            return;
        }

        let _ = Managed::<DLManagedTensorVersioned>::from_raw_unchecked(ptr as *mut _);
    }
}

fn raw_dlpack_to_capsule(
    ptr: *mut std::ffi::c_void,
    name: &CStr,
    deleter: unsafe extern "C" fn(*mut pyo3::ffi::PyObject),
) -> pyo3::PyResult<*mut pyo3::ffi::PyObject> {
    if ptr.is_null() {
        return Err(PyBufferError::new_err(
            "cannot create a DLPack capsule from a null pointer",
        ));
    }
    let capsule = unsafe { pyo3::ffi::PyCapsule_New(ptr, name.as_ptr(), Some(deleter)) };
    if capsule.is_null() {
        Err(fetch_python_error())
    } else {
        Ok(capsule)
    }
}

impl<'py> IntoPyObject<'py> for Managed<DLManagedTensor> {
    type Target = PyAny;
    type Output = Bound<'py, PyAny>;
    type Error = pyo3::PyErr;

    fn into_pyobject(self, py: pyo3::Python<'py>) -> pyo3::PyResult<Self::Output> {
        unsafe {
            let raw = self.into_raw();
            let capsule =
                match raw_dlpack_to_capsule(raw as *mut _, DLTENSOR, dlpack_capsule_deleter) {
                    Ok(capsule) => capsule,
                    Err(err) => {
                        let _ = Managed::<DLManagedTensor>::from_raw_unchecked(raw);
                        return Err(err);
                    }
                };
            Bound::from_owned_ptr_or_err(py, capsule)
        }
    }
}

impl<'py> IntoPyObject<'py> for Managed<DLManagedTensorVersioned> {
    type Target = PyAny;
    type Output = Bound<'py, PyAny>;
    type Error = pyo3::PyErr;

    fn into_pyobject(self, py: pyo3::Python<'py>) -> pyo3::PyResult<Self::Output> {
        unsafe {
            let raw = self.into_raw();
            let capsule = match raw_dlpack_to_capsule(
                raw as *mut _,
                DLTENSOR_VERSIONED,
                dlpack_capsule_deleter_versioned,
            ) {
                Ok(capsule) => capsule,
                Err(err) => {
                    let _ = Managed::<DLManagedTensorVersioned>::from_raw_unchecked(raw);
                    return Err(err);
                }
            };
            Bound::from_owned_ptr_or_err(py, capsule)
        }
    }
}