Skip to content
Permalink
Browse files

issue #2433: low-overhead function calls in python

  • Loading branch information...
jgillis committed Jul 21, 2019
1 parent 92dcf7e commit 3e9919e9195e8f057d09cd9bcc4f7fa59e6fc104
@@ -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);
}
@@ -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();
@@ -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_);
}
@@ -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
@@ -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
@@ -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"
@@ -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

@@ -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")

0 comments on commit 3e9919e

Please sign in to comment.
You can’t perform that action at this time.