Merge pull request symengine#291 from richardotis/master

LambdaDouble/LLVMDouble: Add cdef nogil function wrappers for fast calling
isuruf · Aug 4, 2019 · 45d649e · 45d649e
2 parents d21ec88 + bb5b40c
commit 45d649e
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 3 deletions.
diff --git a/symengine/lib/symengine_wrapper.pxd b/symengine/lib/symengine_wrapper.pxd
@@ -39,9 +39,11 @@ cdef class _Lambdify(object):
 
     cdef _init(self, symengine.vec_basic& args_, symengine.vec_basic& outs_, cppbool cse)
     cdef _load(self, const string &s)
+    cdef void unsafe_real_ptr(self, double *inp, double *out) nogil
     cpdef unsafe_real(self,
                       double[::1] inp, double[::1] out,
                       int inp_offset=*, int out_offset=*)
+    cdef void unsafe_complex_ptr(self, double complex *inp, double complex *out) nogil
     cpdef unsafe_complex(self, double complex[::1] inp, double complex[::1] out,
                          int inp_offset=*, int out_offset=*)
     cpdef eval_real(self, inp, out)
@@ -51,7 +53,9 @@ cdef class LambdaDouble(_Lambdify):
     cdef vector[symengine.LambdaRealDoubleVisitor] lambda_double
     cdef vector[symengine.LambdaComplexDoubleVisitor] lambda_double_complex
     cdef _init(self, symengine.vec_basic& args_, symengine.vec_basic& outs_, cppbool cse)
+    cdef void unsafe_real_ptr(self, double *inp, double *out) nogil
     cpdef unsafe_real(self, double[::1] inp, double[::1] out, int inp_offset=*, int out_offset=*)
+    cdef void unsafe_complex_ptr(self, double complex *inp, double complex *out) nogil
     cpdef unsafe_complex(self, double complex[::1] inp, double complex[::1] out, int inp_offset=*, int out_offset=*)
     cpdef as_scipy_low_level_callable(self)
 
@@ -60,5 +64,6 @@ IF HAVE_SYMENGINE_LLVM:
         cdef vector[symengine.LLVMDoubleVisitor] lambda_double
         cdef _init(self, symengine.vec_basic& args_, symengine.vec_basic& outs_, cppbool cse)
         cdef _load(self, const string &s)
+        cdef void unsafe_real_ptr(self, double *inp, double *out) nogil
         cpdef unsafe_real(self, double[::1] inp, double[::1] out, int inp_offset=*, int out_offset=*)
         cpdef as_scipy_low_level_callable(self)
diff --git a/symengine/lib/symengine_wrapper.pyx b/symengine/lib/symengine_wrapper.pyx
@@ -4507,11 +4507,19 @@ cdef class _Lambdify(object):
     cdef _load(self, const string &s):
         raise ValueError("Not supported")
 
+    cdef void unsafe_real_ptr(self, double *inp, double *out) nogil:
+        with gil:
+            raise ValueError("Not supported")
+
     cpdef unsafe_real(self,
                       double[::1] inp, double[::1] out,
                       int inp_offset=0, int out_offset=0):
         raise ValueError("Not supported")
 
+    cdef void unsafe_complex_ptr(self, double complex *inp, double complex *out) nogil:
+        with gil:
+            raise ValueError("Not supported")
+
     cpdef unsafe_complex(self, double complex[::1] inp, double complex[::1] out,
                          int inp_offset=0, int out_offset=0):
         raise ValueError("Not supported")
@@ -4690,11 +4698,17 @@ cdef class LambdaDouble(_Lambdify):
             self.lambda_double_complex.resize(1)
             self.lambda_double_complex[0].init(args_, outs_, cse)
 
+    cdef void unsafe_real_ptr(self, double *inp, double *out) nogil:
+        self.lambda_double[0].call(out, inp)
+
     cpdef unsafe_real(self, double[::1] inp, double[::1] out, int inp_offset=0, int out_offset=0):
-        self.lambda_double[0].call(&out[out_offset], &inp[inp_offset])
+        self.unsafe_real_ptr(&inp[inp_offset], &out[out_offset])
+
+    cdef void unsafe_complex_ptr(self, double complex *inp, double complex *out) nogil:
+        self.lambda_double_complex[0].call(out, inp)
 
     cpdef unsafe_complex(self, double complex[::1] inp, double complex[::1] out, int inp_offset=0, int out_offset=0):
-        self.lambda_double_complex[0].call(&out[out_offset], &inp[inp_offset])
+        self.unsafe_complex_ptr(&inp[inp_offset], &out[out_offset])
 
     cpdef as_scipy_low_level_callable(self):
         from ctypes import c_double, c_void_p, c_int, cast, POINTER, CFUNCTYPE
@@ -4726,8 +4740,11 @@ IF HAVE_SYMENGINE_LLVM:
             return llvm_loading_func, (self.args_size, self.tot_out_size, self.out_shapes, self.real, \
                 self.n_exprs, self.order, self.accum_out_sizes, self.numpy_dtype, s)
 
+        cdef void unsafe_real_ptr(self, double *inp, double *out) nogil:
+            self.lambda_double[0].call(out, inp)
+
         cpdef unsafe_real(self, double[::1] inp, double[::1] out, int inp_offset=0, int out_offset=0):
-            self.lambda_double[0].call(&out[out_offset], &inp[inp_offset])
+            self.unsafe_real_ptr(&inp[inp_offset], &out[out_offset])
 
         cpdef as_scipy_low_level_callable(self):
             from ctypes import c_double, c_void_p, c_int, cast, POINTER, CFUNCTYPE