issue #2433: low-overhead function calls in python

jgillis · jgillis · commit 3e9919e9195e · 2019-08-09T12:57:53.000+02:00
diff --git a/casadi/core/callback.cpp b/casadi/core/callback.cpp
@@ -52,6 +52,13 @@ namespace casadi {
     }
   }
 
+  int Callback::eval_buffer(const double **arg, const std::vector<casadi_int>& sizes_arg,
+                              double **res, const std::vector<casadi_int>& sizes_res) const {
+    casadi_error("eval_buffer not overloaded.");
+  }
+  bool Callback::has_eval_buffer() const {
+    return false;
+  }
   std::vector<DM> Callback::eval(const std::vector<DM>& arg) const {
     return (*this)->FunctionInternal::eval_dm(arg);
   }
diff --git a/casadi/core/callback.hpp b/casadi/core/callback.hpp
@@ -85,7 +85,15 @@ namespace casadi {
     /** \brief Evaluate numerically, temporary matrices and work vectors */
     virtual std::vector<DM> eval(const std::vector<DM>& arg) const;
 
-   /** \brief Get the number of inputs
+    /** \brief A copy-free low level interface 
+     * 
+     * In Python, you will be passed two tuples of memoryview objects
+    */
+    virtual int eval_buffer(const double **arg, const std::vector<casadi_int>& sizes_arg,
+                              double **res, const std::vector<casadi_int>& sizes_res) const;
+    virtual bool has_eval_buffer() const;
+
+    /** \brief Get the number of inputs
      * This function is called during construction.
      */
     virtual casadi_int get_n_in();
diff --git a/casadi/core/callback_internal.cpp b/casadi/core/callback_internal.cpp
@@ -39,7 +39,7 @@ namespace casadi {
 
   CallbackInternal::
   CallbackInternal(const std::string& name, Callback *self)
-    : FunctionInternal(name), self_(self) {
+    : FunctionInternal(name), self_(self), has_eval_buffer_(false) {
   }
 
   CallbackInternal::~CallbackInternal() {
@@ -70,6 +70,10 @@ namespace casadi {
     TRY_CALL(get_name_out, self_, i);
   }
 
+  bool CallbackInternal::has_eval_buffer() const {
+    TRY_CALL(has_eval_buffer, self_);
+  }
+
   void CallbackInternal::init(const Dict& opts) {
     // Initialize the base classes
     FunctionInternal::init(opts);
@@ -86,12 +90,35 @@ namespace casadi {
 
     // Finalize the base classes
     FunctionInternal::finalize();
+
+    has_eval_buffer_ = has_eval_buffer();
+
+    if (has_eval_buffer_) {
+      sizes_arg_.resize(n_in_);
+      for (casadi_int i=0;i<n_in_;++i) {
+        sizes_arg_[i] = nnz_in(i);
+      }
+      sizes_res_.resize(n_out_);
+      for (casadi_int i=0;i<n_out_;++i) {
+        sizes_res_[i] = nnz_out(i);
+      }
+    }
   }
 
   std::vector<DM> CallbackInternal::eval_dm(const std::vector<DM>& arg) const {
     TRY_CALL(eval, self_, arg);
   }
 
+  /** \brief  Evaluate numerically */
+  int CallbackInternal::eval(const double** arg, double** res,
+      casadi_int* iw, double* w, void* mem) const {
+    if (has_eval_dm()) {
+      return FunctionInternal::eval(arg, res, iw, w, mem);
+    } else {
+      TRY_CALL(eval_buffer, self_, arg, sizes_arg_, res, sizes_res_);
+    }
+  }
+
   bool CallbackInternal::uses_output() const {
     TRY_CALL(uses_output, self_);
   }
diff --git a/casadi/core/callback_internal.hpp b/casadi/core/callback_internal.hpp
@@ -68,11 +68,16 @@ namespace casadi {
     void finalize() override;
 
     ///@{
-    /** \brief Evaluate with DM matrices (NOTE: eval not defined) */
+    /** \brief Evaluate with DM matrices */
     std::vector<DM> eval_dm(const std::vector<DM>& arg) const override;
-    bool has_eval_dm() const override { return true;}
+    bool has_eval_dm() const override { return !has_eval_buffer_;}
     ///@}
 
+    /** \brief  Evaluate numerically */
+    virtual int eval(const double** arg, double** res,
+      casadi_int* iw, double* w, void* mem) const override;
+    bool has_eval_buffer() const;
+
     /** \brief Do the derivative functions need nondifferentiated outputs? */
     bool uses_output() const override;
 
@@ -105,6 +110,10 @@ namespace casadi {
 
     /** \brief Pointer to the public class */
     Callback* self_;
+
+    // For buffered evaluation
+    std::vector<casadi_int> sizes_arg_, sizes_res_;
+    bool has_eval_buffer_;
   };
 
 } // namespace casadi
diff --git a/casadi/core/function.cpp b/casadi/core/function.cpp
@@ -1694,5 +1694,67 @@ namespace casadi {
     return (*this)->info();
   }
 
+  FunctionBuffer::FunctionBuffer(const Function& f) : f_(f) {
+    w_.resize(f_.sz_w());
+    iw_.resize(f_.sz_iw());
+    arg_.resize(f_.sz_arg());
+    res_.resize(f_.sz_res());
+    mem_ = f_->checkout();
+    mem_internal_ = f.memory(mem_);
+    f_node_ = f.operator->();
+  }
+
+  FunctionBuffer::~FunctionBuffer() {
+    if (f_->release_) {
+      f_->release_(mem_);
+    } else {
+      f_.release(mem_);
+    }
+  }
+
+  FunctionBuffer::FunctionBuffer(const FunctionBuffer& f) : f_(f.f_) {
+    operator=(f);
+  }
+
+  FunctionBuffer& FunctionBuffer::operator=(const FunctionBuffer& f) {
+    f_ = f.f_;
+    w_ = f.w_; iw_ = f.iw_; arg_ = f.arg_; res_ = f.res_; f_node_ = f.f_node_;
+    // Checkout fresh memory
+    if (f_->checkout_) {
+      mem_ = f_->checkout_();
+    } else {
+      mem_ = f_.checkout();
+      mem_internal_ = f_.memory(mem_);
+    }
+
+    return *this;
+  }
+
+  void FunctionBuffer::set_arg(casadi_int i, const double* a, casadi_int size) {
+    casadi_assert(size>=f_.nnz_in(i)*sizeof(double),
+     "Buffer is not large enough. Needed " + str(f_.nnz_in(i)*sizeof(double)) +
+     " bytes, got " + str(size) + ".");
+    arg_.at(i) = a;
+  }
+  void FunctionBuffer::set_res(casadi_int i, double* a, casadi_int size) {
+    casadi_assert(size>=f_.nnz_out(i)*sizeof(double),
+     "Buffer is not large enough. Needed " + str(f_.nnz_out(i)*sizeof(double)) +
+     " bytes, got " + str(size) + ".");
+    res_.at(i) = a;
+  }
+  void FunctionBuffer::_eval() {
+    if (f_node_->eval_) {
+      ret_ = f_node_->eval_(get_ptr(arg_), get_ptr(res_), get_ptr(iw_), get_ptr(w_), mem_);
+    } else {
+      ret_ = f_node_->eval(get_ptr(arg_), get_ptr(res_), get_ptr(iw_), get_ptr(w_), mem_internal_);
+    }
+  }
+  int FunctionBuffer::ret() {
+    return ret_;
+  }
+
+  void CASADI_EXPORT _function_buffer_eval(void* raw) {
+    static_cast<FunctionBuffer*>(raw)->_eval();
+  }
 
 } // namespace casadi
diff --git a/casadi/core/function.hpp b/casadi/core/function.hpp
@@ -29,6 +29,7 @@
 #include "mx.hpp"
 #include "printable.hpp"
 #include <exception>
+#include <stack>
 
 namespace casadi {
 
@@ -1024,6 +1025,52 @@ namespace casadi {
 
   };
 
+
+/** \brief Class to achieve minimal overhead function evaluations
+*/
+class CASADI_EXPORT FunctionBuffer {
+  Function f_;
+  std::vector<double> w_;
+  std::vector<casadi_int> iw_;
+  std::vector<const double*> arg_;
+  std::vector<double*> res_;
+  FunctionInternal* f_node_;
+  casadi_int mem_;
+  void *mem_internal_;
+  int ret_;
+public:
+  /** \brief Main constructor */
+  FunctionBuffer(const Function& f);
+#ifndef SWIG
+  ~FunctionBuffer();
+  FunctionBuffer(const FunctionBuffer& f);
+  FunctionBuffer& operator=(const FunctionBuffer& f);
+#endif // SWIG
+
+  /** \brief Set input buffer for input i
+
+      mem.set_arg(0, memoryview(a))
+
+      Note that CasADi uses 'fortran' order: column-by-column
+  */
+  void set_arg(casadi_int i, const double* a, casadi_int size);
+
+  /** \brief Set output buffer for ouput i
+
+      mem.set_res(0, memoryview(a))
+
+      Note that CasADi uses 'fortran' order: column-by-column
+  */
+  void set_res(casadi_int i, double* a, casadi_int size);
+  /// Get last return value
+  int ret();
+  void _eval();
+  void* _self() { return this; }
+};
+
+void CASADI_EXPORT _function_buffer_eval(void* raw);
+
+
 } // namespace casadi
 
 #include "casadi_interrupt.hpp"
diff --git a/swig/casadi.i b/swig/casadi.i
@@ -23,6 +23,7 @@
  */
 
 
+
 %module(package="casadi",directors=1) casadi
 
 #ifdef CASADI_WITH_COPYSIGN_UNDEF
@@ -2294,6 +2295,59 @@ namespace std {
   #define L_STR "str"
 #endif
 
+#ifdef SWIGPYTHON
+%typemap(in, doc="memoryview(ro)", noblock=1, fragment="casadi_all") (const double * a, casadi_int size) (Py_buffer* buffer) {
+  if (!PyMemoryView_Check($input)) SWIG_exception_fail(SWIG_TypeError, "Must supply a MemoryView.");
+  buffer = PyMemoryView_GET_BUFFER($input);
+  $1 = static_cast<double*>(buffer->buf); // const double cast comes later
+  $2 = buffer->len;
+ }
+
+%typemap(in, doc="memoryview(rw)", noblock=1, fragment="casadi_all") (double * a, casadi_int size)  (Py_buffer* buffer) {
+  if (!PyMemoryView_Check($input)) SWIG_exception_fail(SWIG_TypeError, "Must supply a writable MemoryView.");
+  buffer = PyMemoryView_GET_BUFFER($input);
+  if (buffer->readonly) SWIG_exception_fail(SWIG_TypeError, "Must supply a writable MemoryView.");
+  $1 = static_cast<double*>(buffer->buf);
+  $2 = buffer->len;
+ }
+
+// Directorin typemap; as output
+%typemap(directorin, noblock=1, fragment="casadi_all") (const double** arg, const std::vector<casadi_int>& sizes_arg) (PyObject* my_tuple) {
+  PyObject * arg_tuple = PyTuple_New($2.size());
+  for (casadi_int i=0;i<$2.size();++i) {
+    
+#ifdef WITH_PYTHON3
+    PyObject* buf = $1[i] ? PyMemoryView_FromMemory(reinterpret_cast<char*>(const_cast<double*>($1[i])), $2[i]*sizeof(double), PyBUF_READ) : SWIG_Py_Void();
+#else
+    PyObject* buf = $1[i] ? PyBuffer_FromMemory(const_cast<double*>($1[i]), $2[i]*sizeof(double)) : SWIG_Py_Void();
+#endif
+    PyTuple_SET_ITEM(arg_tuple, i, buf);
+  }
+  $input = arg_tuple;
+}
+
+%typemap(directorin, noblock=1, fragment="casadi_all") (double** res, const std::vector<casadi_int>& sizes_res) {
+  PyObject* res_tuple = PyTuple_New($2.size());
+  for (casadi_int i=0;i<$2.size();++i) {
+#ifdef WITH_PYTHON3
+    PyObject* buf = $1[i] ? PyMemoryView_FromMemory(reinterpret_cast<char*>(const_cast<double*>($1[i])), $2[i]*sizeof(double), PyBUF_WRITE) : SWIG_Py_Void();
+#else
+    PyObject* buf = $1[i] ? PyBuffer_FromReadWriteMemory($1[i], $2[i]*sizeof(double)) : SWIG_Py_Void();
+#endif
+    PyTuple_SET_ITEM(res_tuple, i, buf);
+  }
+  $input = res_tuple;
+}
+
+%typemap(in, doc="void*", noblock=1, fragment="casadi_all") void* raw {
+  $1 = PyCapsule_GetPointer($input, NULL);
+}
+
+%typemap(out, doc="void*", noblock=1, fragment="casadi_all") void* {
+  $result = PyCapsule_New($1, NULL,NULL);
+}
+#endif
+
 %casadi_typemaps(L_STR, PREC_STRING, std::string)
 %casadi_template(LL L_STR LR, PREC_VECTOR, std::vector<std::string>)
 %casadi_typemaps("Sparsity", PREC_SPARSITY, casadi::Sparsity)
@@ -3950,6 +4004,11 @@ def PyFunction(name, obj, inputs, outputs, opts={}):
 %}
 #endif
 
+#ifndef SWIGPYTHON
+%ignore FunctionBuffer;
+%ignore _function_buffer_eval;
+#endif
+
 %include <casadi/core/function.hpp>
 #ifdef SWIGPYTHON
 namespace casadi{
@@ -3971,8 +4030,21 @@ namespace casadi{
       else:
         # Named inputs -> return dictionary
         return self.call(kwargs)
+
+    def buffer(self):
+      """
+      Create a FunctionBuffer object for evaluating with minimal overhead
+
+      """
+      import functools
+      fb = FunctionBuffer(self)
+      caller = functools.partial(_casadi._function_buffer_eval, fb._self())
+      return (fb, caller)
   %}
+
+
  }
+
 }
 #endif // SWIGPYTHON
 
diff --git a/swig/python/CMakeLists.txt b/swig/python/CMakeLists.txt
@@ -58,6 +58,8 @@ endif()
 set(PYTHONFLAG "")
 set(CMAKE_SWIG_FLAGS ${CMAKE_SWIG_FLAGS} "-DPy_USING_UNICODE")
 set(CMAKE_SWIG_FLAGS ${CMAKE_SWIG_FLAGS} "-noproxydel")
+set(CMAKE_SWIG_FLAGS ${CMAKE_SWIG_FLAGS} "-fastunpack")
+set(CMAKE_SWIG_FLAGS ${CMAKE_SWIG_FLAGS} "-modernargs")
 if("${PYTHON_VERSION_MAJOR}" STREQUAL "3")
 set(PYTHONFLAG "3")
 set(CMAKE_SWIG_FLAGS ${CMAKE_SWIG_FLAGS} "-py3")
diff --git a/test/python/function.py b/test/python/function.py

Original file line number	Diff line number	Diff line change
`@@ -52,6 +52,13 @@ namespace casadi {`
`52`	`52`	`}`
`53`	`53`	`}`
`54`	`54`
	`55`	`+ int Callback::eval_buffer(const double **arg, const std::vector<casadi_int>& sizes_arg,`
	`56`	`+ double **res, const std::vector<casadi_int>& sizes_res) const {`
	`57`	`+ casadi_error("eval_buffer not overloaded.");`
	`58`	`+ }`
	`59`	`+ bool Callback::has_eval_buffer() const {`
	`60`	`+ return false;`
	`61`	`+ }`
`55`	`62`	`std::vector<DM> Callback::eval(const std::vector<DM>& arg) const {`
`56`	`63`	`return (*this)->FunctionInternal::eval_dm(arg);`
`57`	`64`	`}`