ufunc.hpp
Go to the documentation of this file.
1 //
2 // Copyright (c) 2020-2025 INRIA
3 // code aptapted from
4 // https://github.com/numpy/numpy/blob/41977b24ae011a51f64faa75cb524c7350fdedd9/numpy/core/src/umath/_rational_tests.c.src
5 //
6 
7 #ifndef __eigenpy_ufunc_hpp__
8 #define __eigenpy_ufunc_hpp__
9 
10 #include "eigenpy/register.hpp"
11 #include "eigenpy/user-type.hpp"
12 #include "eigenpy/utils/python-compat.hpp"
13 
14 namespace eigenpy {
15 namespace internal {
16 
17 #ifdef NPY_1_19_API_VERSION
18 #define EIGENPY_NPY_CONST_UFUNC_ARG const
19 #else
20 #define EIGENPY_NPY_CONST_UFUNC_ARG
21 #endif
22 
23 template <typename T>
24 void matrix_multiply(char **args, npy_intp const *dimensions,
25  npy_intp const *steps) {
26  /* pointers to data for input and output arrays */
27  char *ip1 = args[0];
28  char *ip2 = args[1];
29  char *op = args[2];
30 
31  /* lengths of core dimensions */
32  npy_intp dm = dimensions[0];
33  npy_intp dn = dimensions[1];
34  npy_intp dp = dimensions[2];
35 
36  /* striding over core dimensions */
37  npy_intp is1_m = steps[0];
38  npy_intp is1_n = steps[1];
39  npy_intp is2_n = steps[2];
40  npy_intp is2_p = steps[3];
41  npy_intp os_m = steps[4];
42  npy_intp os_p = steps[5];
43 
44  /* core dimensions counters */
45  npy_intp m, p;
46 
47  /* calculate dot product for each row/column vector pair */
48  for (m = 0; m < dm; m++) {
49  for (p = 0; p < dp; p++) {
50  SpecialMethods<T>::dotfunc(ip1, is1_n, ip2, is2_n, op, dn, NULL);
51 
52  /* advance to next column of 2nd input array and output array */
53  ip2 += is2_p;
54  op += os_p;
55  }
56 
57  /* reset to first column of 2nd input array and output array */
58  ip2 -= is2_p * p;
59  op -= os_p * p;
60 
61  /* advance to next row of 1st input array and output array */
62  ip1 += is1_m;
63  op += os_m;
64  }
65 }
66 
67 template <typename T>
68 void gufunc_matrix_multiply(char **args,
69  npy_intp EIGENPY_NPY_CONST_UFUNC_ARG *dimensions,
70  npy_intp EIGENPY_NPY_CONST_UFUNC_ARG *steps,
71  void *NPY_UNUSED(func)) {
72  /* outer dimensions counter */
73  npy_intp N_;
74 
75  /* length of flattened outer dimensions */
76  npy_intp dN = dimensions[0];
77 
78  /* striding over flattened outer dimensions for input and output arrays */
79  npy_intp s0 = steps[0];
80  npy_intp s1 = steps[1];
81  npy_intp s2 = steps[2];
82 
83  /*
84  * loop through outer dimensions, performing matrix multiply on
85  * core dimensions for each loop
86  */
87  for (N_ = 0; N_ < dN; N_++, args[0] += s0, args[1] += s1, args[2] += s2) {
88  matrix_multiply<T>(args, dimensions + 1, steps + 3);
89  }
90 }
91 
92 #define EIGENPY_REGISTER_BINARY_OPERATOR(name, op) \
93  template <typename T1, typename T2, typename R> \
94  void binary_op_##name( \
95  char **args, EIGENPY_NPY_CONST_UFUNC_ARG npy_intp *dimensions, \
96  EIGENPY_NPY_CONST_UFUNC_ARG npy_intp *steps, void * /*data*/) { \
97  npy_intp is0 = steps[0], is1 = steps[1], os = steps[2], n = *dimensions; \
98  char *i0 = args[0], *i1 = args[1], *o = args[2]; \
99  int k; \
100  for (k = 0; k < n; k++) { \
101  T1 &x = *static_cast<T1 *>(static_cast<void *>(i0)); \
102  T2 &y = *static_cast<T2 *>(static_cast<void *>(i1)); \
103  R &res = *static_cast<R *>(static_cast<void *>(o)); \
104  res = x op y; \
105  i0 += is0; \
106  i1 += is1; \
107  o += os; \
108  } \
109  } \
110  \
111  template <typename T> \
112  void binary_op_##name( \
113  char **args, EIGENPY_NPY_CONST_UFUNC_ARG npy_intp *dimensions, \
114  EIGENPY_NPY_CONST_UFUNC_ARG npy_intp *steps, void *data) { \
115  binary_op_##name<T, T, T>(args, dimensions, steps, data); \
116  }
117 
118 EIGENPY_REGISTER_BINARY_OPERATOR(add, +)
119 EIGENPY_REGISTER_BINARY_OPERATOR(subtract, -)
120 EIGENPY_REGISTER_BINARY_OPERATOR(multiply, *)
121 EIGENPY_REGISTER_BINARY_OPERATOR(divide, /)
122 EIGENPY_REGISTER_BINARY_OPERATOR(equal, ==)
123 EIGENPY_REGISTER_BINARY_OPERATOR(not_equal, !=)
124 EIGENPY_REGISTER_BINARY_OPERATOR(less, <)
125 EIGENPY_REGISTER_BINARY_OPERATOR(greater, >)
126 EIGENPY_REGISTER_BINARY_OPERATOR(less_equal, <=)
127 EIGENPY_REGISTER_BINARY_OPERATOR(greater_equal, >=)
128 
129 #define EIGENPY_REGISTER_UNARY_OPERATOR(name, op) \
130  template <typename T, typename R> \
131  void unary_op_##name( \
132  char **args, EIGENPY_NPY_CONST_UFUNC_ARG npy_intp *dimensions, \
133  EIGENPY_NPY_CONST_UFUNC_ARG npy_intp *steps, void * /*data*/) { \
134  npy_intp is = steps[0], os = steps[1], n = *dimensions; \
135  char *i = args[0], *o = args[1]; \
136  int k; \
137  for (k = 0; k < n; k++) { \
138  T &x = *static_cast<T *>(static_cast<void *>(i)); \
139  R &res = *static_cast<R *>(static_cast<void *>(o)); \
140  res = op x; \
141  i += is; \
142  o += os; \
143  } \
144  } \
145  \
146  template <typename T> \
147  void unary_op_##name( \
148  char **args, EIGENPY_NPY_CONST_UFUNC_ARG npy_intp *dimensions, \
149  EIGENPY_NPY_CONST_UFUNC_ARG npy_intp *steps, void *data) { \
150  unary_op_##name<T, T>(args, dimensions, steps, data); \
151  }
152 
153 EIGENPY_REGISTER_UNARY_OPERATOR(negative, -)
154 EIGENPY_REGISTER_UNARY_OPERATOR(square, x *)
155 
156 } // namespace internal
157 
158 #define EIGENPY_REGISTER_BINARY_UFUNC(name, code, T1, T2, R) \
159  { \
160  PyUFuncObject *ufunc = \
161  (PyUFuncObject *)PyObject_GetAttrString(numpy, #name); \
162  int _types[3] = {Register::getTypeCode<T1>(), Register::getTypeCode<T2>(), \
163  Register::getTypeCode<R>()}; \
164  if (!ufunc) { \
165  /*goto fail; \*/ \
166  } \
167  if (sizeof(_types) / sizeof(int) != ufunc->nargs) { \
168  PyErr_Format(PyExc_AssertionError, \
169  "ufunc %s takes %d arguments, our loop takes %lu", #name, \
170  ufunc->nargs, \
171  (unsigned long)(sizeof(_types) / sizeof(int))); \
172  Py_DECREF(ufunc); \
173  } \
174  if (PyUFunc_RegisterLoopForType((PyUFuncObject *)ufunc, code, \
175  internal::binary_op_##name<T1, T2, R>, \
176  _types, 0) < 0) { \
177  /*Py_DECREF(ufunc);*/ \
178  /*goto fail; \*/ \
179  } \
180  Py_DECREF(ufunc); \
181  }
182 
183 #define EIGENPY_REGISTER_UNARY_UFUNC(name, code, T, R) \
184  { \
185  PyUFuncObject *ufunc = \
186  (PyUFuncObject *)PyObject_GetAttrString(numpy, #name); \
187  int _types[2] = {Register::getTypeCode<T>(), Register::getTypeCode<R>()}; \
188  if (!ufunc) { \
189  /*goto fail; \*/ \
190  } \
191  if (sizeof(_types) / sizeof(int) != ufunc->nargs) { \
192  PyErr_Format(PyExc_AssertionError, \
193  "ufunc %s takes %d arguments, our loop takes %lu", #name, \
194  ufunc->nargs, \
195  (unsigned long)(sizeof(_types) / sizeof(int))); \
196  Py_DECREF(ufunc); \
197  } \
198  if (PyUFunc_RegisterLoopForType((PyUFuncObject *)ufunc, code, \
199  internal::unary_op_##name<T, R>, _types, \
200  0) < 0) { \
201  /*Py_DECREF(ufunc);*/ \
202  /*goto fail; \*/ \
203  } \
204  Py_DECREF(ufunc); \
205  }
206 
207 template <typename Scalar>
208 void registerCommonUfunc() {
209  const int type_code = Register::getTypeCode<Scalar>();
210 
211  PyObject *numpy_str;
212  numpy_str = PyStr_FromString("numpy");
213  PyObject *numpy;
214  numpy = PyImport_Import(numpy_str);
215  Py_DECREF(numpy_str);
216 
217  import_ufunc();
218 
219  // Matrix multiply
220  {
221  int types[3] = {type_code, type_code, type_code};
222 
223  std::stringstream ss;
224  ss << "return result of multiplying two matrices of ";
225  ss << bp::type_info(typeid(Scalar)).name();
226  PyUFuncObject *ufunc =
227  (PyUFuncObject *)PyObject_GetAttrString(numpy, "matmul");
228  if (!ufunc) {
229  std::stringstream ss;
230  ss << "Impossible to define matrix_multiply for given type "
231  << bp::type_info(typeid(Scalar)).name() << std::endl;
232  eigenpy::Exception(ss.str());
233  }
234  if (PyUFunc_RegisterLoopForType((PyUFuncObject *)ufunc, type_code,
235  &internal::gufunc_matrix_multiply<Scalar>,
236  types, 0) < 0) {
237  std::stringstream ss;
238  ss << "Impossible to register matrix_multiply for given type "
239  << bp::type_info(typeid(Scalar)).name() << std::endl;
240  eigenpy::Exception(ss.str());
241  }
242 
243  Py_DECREF(ufunc);
244  }
245 
246  // Binary operators
247  EIGENPY_REGISTER_BINARY_UFUNC(add, type_code, Scalar, Scalar, Scalar);
248  EIGENPY_REGISTER_BINARY_UFUNC(subtract, type_code, Scalar, Scalar, Scalar);
249  EIGENPY_REGISTER_BINARY_UFUNC(multiply, type_code, Scalar, Scalar, Scalar);
250  EIGENPY_REGISTER_BINARY_UFUNC(divide, type_code, Scalar, Scalar, Scalar);
251 
252  // Comparison operators
253  EIGENPY_REGISTER_BINARY_UFUNC(equal, type_code, Scalar, Scalar, bool);
254  EIGENPY_REGISTER_BINARY_UFUNC(not_equal, type_code, Scalar, Scalar, bool);
255  EIGENPY_REGISTER_BINARY_UFUNC(greater, type_code, Scalar, Scalar, bool);
256  EIGENPY_REGISTER_BINARY_UFUNC(less, type_code, Scalar, Scalar, bool);
257  EIGENPY_REGISTER_BINARY_UFUNC(greater_equal, type_code, Scalar, Scalar, bool);
258  EIGENPY_REGISTER_BINARY_UFUNC(less_equal, type_code, Scalar, Scalar, bool);
259 
260  // Unary operators
261  EIGENPY_REGISTER_UNARY_UFUNC(negative, type_code, Scalar, Scalar);
262  EIGENPY_REGISTER_UNARY_UFUNC(square, type_code, Scalar, Scalar);
263 
264  Py_DECREF(numpy);
265 }
266 
267 } // namespace eigenpy
268 
269 #endif // __eigenpy_ufunc_hpp__
EIGENPY_REGISTER_UNARY_OPERATOR
#define EIGENPY_REGISTER_UNARY_OPERATOR(name, op)
Definition: ufunc.hpp:129
register.hpp
EIGENPY_REGISTER_UNARY_UFUNC
#define EIGENPY_REGISTER_UNARY_UFUNC(name, code, T, R)
Definition: ufunc.hpp:183
eigenpy::internal::gufunc_matrix_multiply
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))
Definition: ufunc.hpp:68
EIGENPY_REGISTER_BINARY_UFUNC
#define EIGENPY_REGISTER_BINARY_UFUNC(name, code, T1, T2, R)
Definition: ufunc.hpp:158
user-type.hpp
eigenpy
Definition: alignment.hpp:14
eigenpy::registerCommonUfunc
void registerCommonUfunc()
Definition: ufunc.hpp:208
eigenpy::type_info
boost::typeindex::type_index type_info(const T &value)
Definition: type_info.hpp:17
eigenpy::internal::matrix_multiply
void matrix_multiply(char **args, npy_intp const *dimensions, npy_intp const *steps)
Definition: ufunc.hpp:24
python-compat.hpp
test_user_struct.x
x
Definition: test_user_struct.py:4
eigenpy::Exception
Definition: exception.hpp:19
test_sparse_matrix.m
m
Definition: test_sparse_matrix.py:5
PyStr_FromString
#define PyStr_FromString
Definition: python-compat.hpp:19
EIGENPY_NPY_CONST_UFUNC_ARG
#define EIGENPY_NPY_CONST_UFUNC_ARG
Definition: ufunc.hpp:20
EIGENPY_REGISTER_BINARY_OPERATOR
#define EIGENPY_REGISTER_BINARY_OPERATOR(name, op)
Definition: ufunc.hpp:92
eigenpy::internal::SpecialMethods::dotfunc
static void dotfunc(void *, npy_intp, void *, npy_intp, void *, npy_intp, void *)

eigenpy
Author(s): Justin Carpentier, Nicolas Mansard
autogenerated on Fri Feb 14 2025 03:16:26