Program Listing for File ufunc.hpp
↰ Return to documentation for file (include/eigenpy/ufunc.hpp)
//
// Copyright (c) 2020-2021 INRIA
// code aptapted from
// https://github.com/numpy/numpy/blob/41977b24ae011a51f64faa75cb524c7350fdedd9/numpy/core/src/umath/_rational_tests.c.src
//
#ifndef __eigenpy_ufunc_hpp__
#define __eigenpy_ufunc_hpp__
#include "eigenpy/register.hpp"
#include "eigenpy/user-type.hpp"
#include "eigenpy/utils/python-compat.hpp"
namespace eigenpy {
namespace internal {
#ifdef NPY_1_19_API_VERSION
#define EIGENPY_NPY_CONST_UFUNC_ARG const
#else
#define EIGENPY_NPY_CONST_UFUNC_ARG
#endif
template <typename T>
void matrix_multiply(char **args, npy_intp const *dimensions,
npy_intp const *steps) {
/* pointers to data for input and output arrays */
char *ip1 = args[0];
char *ip2 = args[1];
char *op = args[2];
/* lengths of core dimensions */
npy_intp dm = dimensions[0];
npy_intp dn = dimensions[1];
npy_intp dp = dimensions[2];
/* striding over core dimensions */
npy_intp is1_m = steps[0];
npy_intp is1_n = steps[1];
npy_intp is2_n = steps[2];
npy_intp is2_p = steps[3];
npy_intp os_m = steps[4];
npy_intp os_p = steps[5];
/* core dimensions counters */
npy_intp m, p;
/* calculate dot product for each row/column vector pair */
for (m = 0; m < dm; m++) {
for (p = 0; p < dp; p++) {
SpecialMethods<T>::dotfunc(ip1, is1_n, ip2, is2_n, op, dn, NULL);
/* advance to next column of 2nd input array and output array */
ip2 += is2_p;
op += os_p;
}
/* reset to first column of 2nd input array and output array */
ip2 -= is2_p * p;
op -= os_p * p;
/* advance to next row of 1st input array and output array */
ip1 += is1_m;
op += os_m;
}
}
template <typename T>
void gufunc_matrix_multiply(char **args,
npy_intp EIGENPY_NPY_CONST_UFUNC_ARG *dimensions,
npy_intp EIGENPY_NPY_CONST_UFUNC_ARG *steps,
void *NPY_UNUSED(func)) {
/* outer dimensions counter */
npy_intp N_;
/* length of flattened outer dimensions */
npy_intp dN = dimensions[0];
/* striding over flattened outer dimensions for input and output arrays */
npy_intp s0 = steps[0];
npy_intp s1 = steps[1];
npy_intp s2 = steps[2];
/*
* loop through outer dimensions, performing matrix multiply on
* core dimensions for each loop
*/
for (N_ = 0; N_ < dN; N_++, args[0] += s0, args[1] += s1, args[2] += s2) {
matrix_multiply<T>(args, dimensions + 1, steps + 3);
}
}
#define EIGENPY_REGISTER_BINARY_OPERATOR(name, op) \
template <typename T1, typename T2, typename R> \
void binary_op_##name( \
char **args, EIGENPY_NPY_CONST_UFUNC_ARG npy_intp *dimensions, \
EIGENPY_NPY_CONST_UFUNC_ARG npy_intp *steps, void * /*data*/) { \
npy_intp is0 = steps[0], is1 = steps[1], os = steps[2], n = *dimensions; \
char *i0 = args[0], *i1 = args[1], *o = args[2]; \
int k; \
for (k = 0; k < n; k++) { \
T1 &x = *static_cast<T1 *>(static_cast<void *>(i0)); \
T2 &y = *static_cast<T2 *>(static_cast<void *>(i1)); \
R &res = *static_cast<R *>(static_cast<void *>(o)); \
res = x op y; \
i0 += is0; \
i1 += is1; \
o += os; \
} \
} \
\
template <typename T> \
void binary_op_##name( \
char **args, EIGENPY_NPY_CONST_UFUNC_ARG npy_intp *dimensions, \
EIGENPY_NPY_CONST_UFUNC_ARG npy_intp *steps, void *data) { \
binary_op_##name<T, T, T>(args, dimensions, steps, data); \
}
EIGENPY_REGISTER_BINARY_OPERATOR(add, +)
EIGENPY_REGISTER_BINARY_OPERATOR(subtract, -)
EIGENPY_REGISTER_BINARY_OPERATOR(multiply, *)
EIGENPY_REGISTER_BINARY_OPERATOR(divide, /)
EIGENPY_REGISTER_BINARY_OPERATOR(equal, ==)
EIGENPY_REGISTER_BINARY_OPERATOR(not_equal, !=)
EIGENPY_REGISTER_BINARY_OPERATOR(less, <)
EIGENPY_REGISTER_BINARY_OPERATOR(greater, >)
EIGENPY_REGISTER_BINARY_OPERATOR(less_equal, <=)
EIGENPY_REGISTER_BINARY_OPERATOR(greater_equal, >=)
#define EIGENPY_REGISTER_UNARY_OPERATOR(name, op) \
template <typename T, typename R> \
void unary_op_##name( \
char **args, EIGENPY_NPY_CONST_UFUNC_ARG npy_intp *dimensions, \
EIGENPY_NPY_CONST_UFUNC_ARG npy_intp *steps, void * /*data*/) { \
npy_intp is = steps[0], os = steps[1], n = *dimensions; \
char *i = args[0], *o = args[1]; \
int k; \
for (k = 0; k < n; k++) { \
T &x = *static_cast<T *>(static_cast<void *>(i)); \
R &res = *static_cast<R *>(static_cast<void *>(o)); \
res = op x; \
i += is; \
o += os; \
} \
} \
\
template <typename T> \
void unary_op_##name( \
char **args, EIGENPY_NPY_CONST_UFUNC_ARG npy_intp *dimensions, \
EIGENPY_NPY_CONST_UFUNC_ARG npy_intp *steps, void *data) { \
unary_op_##name<T, T>(args, dimensions, steps, data); \
}
EIGENPY_REGISTER_UNARY_OPERATOR(negative, -)
} // namespace internal
#define EIGENPY_REGISTER_BINARY_UFUNC(name, code, T1, T2, R) \
{ \
PyUFuncObject *ufunc = \
(PyUFuncObject *)PyObject_GetAttrString(numpy, #name); \
int _types[3] = {Register::getTypeCode<T1>(), Register::getTypeCode<T2>(), \
Register::getTypeCode<R>()}; \
if (!ufunc) { \
/*goto fail; \*/ \
} \
if (sizeof(_types) / sizeof(int) != ufunc->nargs) { \
PyErr_Format(PyExc_AssertionError, \
"ufunc %s takes %d arguments, our loop takes %lu", #name, \
ufunc->nargs, \
(unsigned long)(sizeof(_types) / sizeof(int))); \
Py_DECREF(ufunc); \
} \
if (PyUFunc_RegisterLoopForType((PyUFuncObject *)ufunc, code, \
internal::binary_op_##name<T1, T2, R>, \
_types, 0) < 0) { \
/*Py_DECREF(ufunc);*/ \
/*goto fail; \*/ \
} \
Py_DECREF(ufunc); \
}
#define EIGENPY_REGISTER_UNARY_UFUNC(name, code, T, R) \
{ \
PyUFuncObject *ufunc = \
(PyUFuncObject *)PyObject_GetAttrString(numpy, #name); \
int _types[2] = {Register::getTypeCode<T>(), Register::getTypeCode<R>()}; \
if (!ufunc) { \
/*goto fail; \*/ \
} \
if (sizeof(_types) / sizeof(int) != ufunc->nargs) { \
PyErr_Format(PyExc_AssertionError, \
"ufunc %s takes %d arguments, our loop takes %lu", #name, \
ufunc->nargs, \
(unsigned long)(sizeof(_types) / sizeof(int))); \
Py_DECREF(ufunc); \
} \
if (PyUFunc_RegisterLoopForType((PyUFuncObject *)ufunc, code, \
internal::unary_op_##name<T, R>, _types, \
0) < 0) { \
/*Py_DECREF(ufunc);*/ \
/*goto fail; \*/ \
} \
Py_DECREF(ufunc); \
}
template <typename Scalar>
void registerCommonUfunc() {
const int type_code = Register::getTypeCode<Scalar>();
PyObject *numpy_str;
numpy_str = PyStr_FromString("numpy");
PyObject *numpy;
numpy = PyImport_Import(numpy_str);
Py_DECREF(numpy_str);
import_ufunc();
// Matrix multiply
{
int types[3] = {type_code, type_code, type_code};
std::stringstream ss;
ss << "return result of multiplying two matrices of ";
ss << bp::type_info(typeid(Scalar)).name();
PyUFuncObject *ufunc =
(PyUFuncObject *)PyObject_GetAttrString(numpy, "matmul");
if (!ufunc) {
std::stringstream ss;
ss << "Impossible to define matrix_multiply for given type "
<< bp::type_info(typeid(Scalar)).name() << std::endl;
eigenpy::Exception(ss.str());
}
if (PyUFunc_RegisterLoopForType((PyUFuncObject *)ufunc, type_code,
&internal::gufunc_matrix_multiply<Scalar>,
types, 0) < 0) {
std::stringstream ss;
ss << "Impossible to register matrix_multiply for given type "
<< bp::type_info(typeid(Scalar)).name() << std::endl;
eigenpy::Exception(ss.str());
}
Py_DECREF(ufunc);
}
// Binary operators
EIGENPY_REGISTER_BINARY_UFUNC(add, type_code, Scalar, Scalar, Scalar);
EIGENPY_REGISTER_BINARY_UFUNC(subtract, type_code, Scalar, Scalar, Scalar);
EIGENPY_REGISTER_BINARY_UFUNC(multiply, type_code, Scalar, Scalar, Scalar);
EIGENPY_REGISTER_BINARY_UFUNC(divide, type_code, Scalar, Scalar, Scalar);
// Comparison operators
EIGENPY_REGISTER_BINARY_UFUNC(equal, type_code, Scalar, Scalar, bool);
EIGENPY_REGISTER_BINARY_UFUNC(not_equal, type_code, Scalar, Scalar, bool);
EIGENPY_REGISTER_BINARY_UFUNC(greater, type_code, Scalar, Scalar, bool);
EIGENPY_REGISTER_BINARY_UFUNC(less, type_code, Scalar, Scalar, bool);
EIGENPY_REGISTER_BINARY_UFUNC(greater_equal, type_code, Scalar, Scalar, bool);
EIGENPY_REGISTER_BINARY_UFUNC(less_equal, type_code, Scalar, Scalar, bool);
// Unary operators
EIGENPY_REGISTER_UNARY_UFUNC(negative, type_code, Scalar, Scalar);
Py_DECREF(numpy);
}
} // namespace eigenpy
#endif // __eigenpy_ufunc_hpp__