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

eigenpy
Author(s): Justin Carpentier, Nicolas Mansard
autogenerated on Fri Jun 2 2023 02:10:26