numpy/npyffi/
ufunc.rs

1//! Low-Level binding for [UFunc API](https://numpy.org/doc/stable/reference/c-api/ufunc.html)
2
3use std::ffi::*;
4
5use pyo3::{ffi::PyObject, sync::PyOnceLock};
6
7use crate::npyffi::*;
8
9fn mod_name(py: Python<'_>) -> PyResult<&'static str> {
10    static MOD_NAME: PyOnceLock<String> = PyOnceLock::new();
11    MOD_NAME
12        .get_or_try_init(py, || {
13            let numpy_core = super::array::numpy_core_name(py)?;
14            Ok(format!("{numpy_core}.umath"))
15        })
16        .map(String::as_str)
17}
18
19const CAPSULE_NAME: &str = "_UFUNC_API";
20
21/// A global variable which stores a ['capsule'](https://docs.python.org/3/c-api/capsule.html)
22/// pointer to [Numpy UFunc API](https://numpy.org/doc/stable/reference/c-api/ufunc.html).
23pub static PY_UFUNC_API: PyUFuncAPI = PyUFuncAPI(PyOnceLock::new());
24
25pub struct PyUFuncAPI(PyOnceLock<NonNull<*const c_void>>);
26
27unsafe impl Send for PyUFuncAPI {}
28
29unsafe impl Sync for PyUFuncAPI {}
30
31impl PyUFuncAPI {
32    pub(super) unsafe fn get<'py>(&self, py: Python<'py>, offset: isize) -> NonNull<*const c_void> {
33        let api = self
34            .0
35            .get_or_try_init(py, || get_numpy_api(py, mod_name(py)?, CAPSULE_NAME))
36            .expect("Failed to access NumPy ufunc API capsule");
37
38        api.offset(offset)
39    }
40}
41
42// See PyArrayAPI for a reference of NumPy API mapping.
43impl PyUFuncAPI {
44    impl_api![1; PyUFunc_FromFuncAndData(func: *mut PyUFuncGenericFunction, data: *mut *mut c_void, types: *mut c_char, ntypes: c_int, nin: c_int, nout: c_int, identity: c_int, name: *const c_char, doc: *const c_char, unused: c_int) -> *mut PyObject];
45    impl_api![2; PyUFunc_RegisterLoopForType(ufunc: *mut PyUFuncObject, usertype: c_int, function: PyUFuncGenericFunction, arg_types: *mut c_int, data: *mut c_void) -> c_int];
46    // Unused slot 3, was `PyUFunc_GenericFunction`
47    impl_api![4; PyUFunc_f_f_As_d_d(args: *mut *mut c_char, dimensions: *mut npy_intp, steps: *mut npy_intp, func: *mut c_void)];
48    impl_api![5; PyUFunc_d_d(args: *mut *mut c_char, dimensions: *mut npy_intp, steps: *mut npy_intp, func: *mut c_void)];
49    impl_api![6; PyUFunc_f_f(args: *mut *mut c_char, dimensions: *mut npy_intp, steps: *mut npy_intp, func: *mut c_void)];
50    impl_api![7; PyUFunc_g_g(args: *mut *mut c_char, dimensions: *mut npy_intp, steps: *mut npy_intp, func: *mut c_void)];
51    impl_api![8; PyUFunc_F_F_As_D_D(args: *mut *mut c_char, dimensions: *mut npy_intp, steps: *mut npy_intp, func: *mut c_void)];
52    impl_api![9; PyUFunc_F_F(args: *mut *mut c_char, dimensions: *mut npy_intp, steps: *mut npy_intp, func: *mut c_void)];
53    impl_api![10; PyUFunc_D_D(args: *mut *mut c_char, dimensions: *mut npy_intp, steps: *mut npy_intp, func: *mut c_void)];
54    impl_api![11; PyUFunc_G_G(args: *mut *mut c_char, dimensions: *mut npy_intp, steps: *mut npy_intp, func: *mut c_void)];
55    impl_api![12; PyUFunc_O_O(args: *mut *mut c_char, dimensions: *mut npy_intp, steps: *mut npy_intp, func: *mut c_void)];
56    impl_api![13; PyUFunc_ff_f_As_dd_d(args: *mut *mut c_char, dimensions: *mut npy_intp, steps: *mut npy_intp, func: *mut c_void)];
57    impl_api![14; PyUFunc_ff_f(args: *mut *mut c_char, dimensions: *mut npy_intp, steps: *mut npy_intp, func: *mut c_void)];
58    impl_api![15; PyUFunc_dd_d(args: *mut *mut c_char, dimensions: *mut npy_intp, steps: *mut npy_intp, func: *mut c_void)];
59    impl_api![16; PyUFunc_gg_g(args: *mut *mut c_char, dimensions: *mut npy_intp, steps: *mut npy_intp, func: *mut c_void)];
60    impl_api![17; PyUFunc_FF_F_As_DD_D(args: *mut *mut c_char, dimensions: *mut npy_intp, steps: *mut npy_intp, func: *mut c_void)];
61    impl_api![18; PyUFunc_DD_D(args: *mut *mut c_char, dimensions: *mut npy_intp, steps: *mut npy_intp, func: *mut c_void)];
62    impl_api![19; PyUFunc_FF_F(args: *mut *mut c_char, dimensions: *mut npy_intp, steps: *mut npy_intp, func: *mut c_void)];
63    impl_api![20; PyUFunc_GG_G(args: *mut *mut c_char, dimensions: *mut npy_intp, steps: *mut npy_intp, func: *mut c_void)];
64    impl_api![21; PyUFunc_OO_O(args: *mut *mut c_char, dimensions: *mut npy_intp, steps: *mut npy_intp, func: *mut c_void)];
65    impl_api![22; PyUFunc_O_O_method(args: *mut *mut c_char, dimensions: *mut npy_intp, steps: *mut npy_intp, func: *mut c_void)];
66    impl_api![23; PyUFunc_OO_O_method(args: *mut *mut c_char, dimensions: *mut npy_intp, steps: *mut npy_intp, func: *mut c_void)];
67    impl_api![24; PyUFunc_On_Om(args: *mut *mut c_char, dimensions: *mut npy_intp, steps: *mut npy_intp, func: *mut c_void)];
68    // Unused slot 25, was `PyUFunc_GetPyValues`
69    // Unused slot 26, was `PyUFunc_checkfperr`
70    impl_api![27; PyUFunc_clearfperr()];
71    impl_api![28; PyUFunc_getfperr() -> c_int];
72    // Unused slot 29, was `PyUFunc_handlefperr`
73    impl_api![30; PyUFunc_ReplaceLoopBySignature(func: *mut PyUFuncObject, newfunc: PyUFuncGenericFunction, signature: *mut c_int, oldfunc: *mut PyUFuncGenericFunction) -> c_int];
74    impl_api![31; PyUFunc_FromFuncAndDataAndSignature(func: *mut PyUFuncGenericFunction, data: *mut *mut c_void, types: *mut c_char, ntypes: c_int, nin: c_int, nout: c_int, identity: c_int, name: *const c_char, doc: *const c_char, unused: c_int, signature: *const c_char) -> *mut PyObject];
75    impl_api![32; PyUFunc_SetUsesArraysAsData(data: *mut *mut c_void, i: usize) -> c_int];
76    impl_api![33; PyUFunc_e_e(args: *mut *mut c_char, dimensions: *mut npy_intp, steps: *mut npy_intp, func: *mut c_void)];
77    impl_api![34; PyUFunc_e_e_As_f_f(args: *mut *mut c_char, dimensions: *mut npy_intp, steps: *mut npy_intp, func: *mut c_void)];
78    impl_api![35; PyUFunc_e_e_As_d_d(args: *mut *mut c_char, dimensions: *mut npy_intp, steps: *mut npy_intp, func: *mut c_void)];
79    impl_api![36; PyUFunc_ee_e(args: *mut *mut c_char, dimensions: *mut npy_intp, steps: *mut npy_intp, func: *mut c_void)];
80    impl_api![37; PyUFunc_ee_e_As_ff_f(args: *mut *mut c_char, dimensions: *mut npy_intp, steps: *mut npy_intp, func: *mut c_void)];
81    impl_api![38; PyUFunc_ee_e_As_dd_d(args: *mut *mut c_char, dimensions: *mut npy_intp, steps: *mut npy_intp, func: *mut c_void)];
82    impl_api![39; PyUFunc_DefaultTypeResolver(ufunc: *mut PyUFuncObject, casting: NPY_CASTING, operands: *mut *mut PyArrayObject, type_tup: *mut PyObject, out_dtypes: *mut *mut PyArray_Descr) -> c_int];
83    impl_api![40; PyUFunc_ValidateCasting(ufunc: *mut PyUFuncObject, casting: NPY_CASTING, operands: *mut *mut PyArrayObject, dtypes: *mut *mut PyArray_Descr) -> c_int];
84    impl_api![41; PyUFunc_RegisterLoopForDescr(ufunc: *mut PyUFuncObject, user_dtype: *mut PyArray_Descr, function: PyUFuncGenericFunction, arg_dtypes: *mut *mut PyArray_Descr, data: *mut c_void) -> c_int];
85    // Min v1.16 impl_api![42; PyUFunc_FromFuncAndDataAndSignatureAndIdentity(
86    //     ufunc: *mut PyUFuncObject,
87    //     data: *mut *mut c_void,
88    //     types: *mut c_char,
89    //     ntypes: c_int,
90    //     nin: c_int,
91    //     nout: c_int,
92    //     identity: c_int,
93    //     name: *const c_char,
94    //     doc: *const c_char,
95    //     unused: c_int,
96    //     signature: *const c_char,
97    //     identity_value: *const c_char,
98    // ) -> c_int];
99    // Min v2.0 impl_api![43; PyUFunc_AddLoopFromSpec(ufunc: *mut PyObject, spec: *mut PyArrayMethod_Spec) -> c_int];
100    // Min v2.0 impl_api![44; PyUFunc_AddPromoter(ufunc: *mut PyObject, dtype_tuple: *mut PyObject, promoter: *mut PyObject) -> c_int];
101    // Min v2.0 impl_api![45; PyUFunc_AddWrappingLoop(
102    //     ufunc_obj: *mut PyObject,
103    //     new_dtypes: *mut *mut PyArray_DTypeMeta,
104    //     wrapped_dtypes: *mut *mut PyArray_DTypeMeta,
105    //     translate_given_descrs: *mut PyArrayMethod_TranslateGivenDescriptors,
106    //     translate_loop_descrs: *mut PyArrayMethod_TranslateLoopDescriptors) -> c_int];
107    // Min v2.0 impl_api![46; PyUFunc_GiveFloatingpointErrors(name: *mut c_char, fpe_errors: c_int) -> c_int];
108}
numpy/npyffi/ufunc.rs

numpy/npyffi/
ufunc.rs
