sx_function.hpp

/*
 *    This file is part of CasADi.
 *
 *    CasADi -- A symbolic framework for dynamic optimization.
 *    Copyright (C) 2010-2023 Joel Andersson, Joris Gillis, Moritz Diehl,
 *                            KU Leuven. All rights reserved.
 *    Copyright (C) 2011-2014 Greg Horn
 *
 *    CasADi is free software; you can redistribute it and/or
 *    modify it under the terms of the GNU Lesser General Public
 *    License as published by the Free Software Foundation; either
 *    version 3 of the License, or (at your option) any later version.
 *
 *    CasADi is distributed in the hope that it will be useful,
 *    but WITHOUT ANY WARRANTY; without even the implied warranty of
 *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 *    Lesser General Public License for more details.
 *
 *    You should have received a copy of the GNU Lesser General Public
 *    License along with CasADi; if not, write to the Free Software
 *    Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
 *
 */
 
 
#ifndef CASADI_SX_FUNCTION_HPP
#define CASADI_SX_FUNCTION_HPP
 
#include "x_function.hpp"
 
 
namespace casadi {
  struct ScalarAtomic {
    int op;     
    int i0;
    union {
      double d;
      struct { int i1, i2; };
    };
  };
 
class CASADI_EXPORT SXFunction :
        public XFunction<SXFunction, Matrix<SXElem>, SXNode>{
  public:
    SXFunction(const std::string& name,
               const std::vector<Matrix<SXElem> >& inputv,
               const std::vector<Matrix<SXElem> >& outputv,
               const std::vector<std::string>& name_in,
               const std::vector<std::string>& name_out);
 
  ~SXFunction() override;
 
  int eval(const double** arg, double** res, casadi_int* iw, double* w, void* mem) const override;
 
  int eval_sx(const SXElem** arg, SXElem** res,
              casadi_int* iw, SXElem* w, void* mem,
              bool always_inline, bool never_inline) const override;
 
  void eval_mx(const MXVector& arg, MXVector& res,
                bool always_inline, bool never_inline) const override;
 
  bool should_inline(bool with_sx, bool always_inline, bool never_inline) const override;
 
  void ad_forward(const std::vector<std::vector<SX> >& fseed,
                            std::vector<std::vector<SX> >& fsens) const;
 
  void ad_reverse(const std::vector<std::vector<SX> >& aseed,
                            std::vector<std::vector<SX> >& asens) const;
 
  bool is_smooth() const;
 
  // print an element of an algorithm
  std::string print(const ScalarAtomic& a) const;
 
  // Print the input arguments of an instruction
  void print_arg(std::ostream &stream, casadi_int k, const ScalarAtomic& el,
    const double* w) const;
 
  // Print the input arguments of an instruction
  void print_arg(CodeGenerator& g, casadi_int k, const ScalarAtomic& el) const;
 
  // Print the output arguments of an instruction
  void print_res(std::ostream &stream, casadi_int k, const ScalarAtomic& el,
    const double* w) const;
 
  // Print the output arguments of an instruction
  void print_res(CodeGenerator& g, casadi_int k, const ScalarAtomic& el) const;
 
  void disp_more(std::ostream& stream) const override;
 
  std::string class_name() const override {return "SXFunction";}
 
  bool is_a(const std::string& type, bool recursive) const override;
 
 
  const SX sx_in(casadi_int ind) const override;
  const std::vector<SX> sx_in() const override;
 
  std::vector<SX> free_sx() const override {
    std::vector<SX> ret(free_vars_.size());
    std::copy(free_vars_.begin(), free_vars_.end(), ret.begin());
    return ret;
  }
 
  bool has_free() const override { return !free_vars_.empty();}
 
  std::vector<std::string> get_free() const override {
    std::vector<std::string> ret;
    for (auto&& e : free_vars_) ret.push_back(e.name());
    return ret;
  }
 
  SX hess(casadi_int iind=0, casadi_int oind=0);
 
  casadi_int n_instructions() const override { return algorithm_.size();}
 
  casadi_int instruction_id(casadi_int k) const override { return algorithm_.at(k).op;}
 
  std::vector<casadi_int> instruction_input(casadi_int k) const override {
    auto e = algorithm_.at(k);
    if (e.op==OP_CALL) {
      const ExtendedAlgEl& m = call_.el[e.i1];
      return vector_static_cast<casadi_int>(m.dep);
    } else if (casadi_math<double>::ndeps(e.op)==2 || e.op==OP_INPUT) {
      return {e.i1, e.i2};
    } else if (casadi_math<double>::ndeps(e.op)==1) {
      return {e.i1};
    } else {
      return {};
    }
  }
 
  double instruction_constant(casadi_int k) const override {
    return algorithm_.at(k).d;
  }
 
  std::vector<casadi_int> instruction_output(casadi_int k) const override {
    auto e = algorithm_.at(k);
    if (e.op==OP_CALL) {
      const ExtendedAlgEl& m = call_.el[e.i1];
      return vector_static_cast<casadi_int>(m.res);
    } else if (e.op==OP_OUTPUT) {
      return {e.i0, e.i2};
    } else {
      return {e.i0};
    }
  }
 
  casadi_int n_nodes() const override { return algorithm_.size() - nnz_out();}
 
  typedef ScalarAtomic AlgEl;
 
  template<typename T>
  struct TapeEl {
    T d[2];
  };
 
  std::vector<AlgEl> algorithm_;
 
  // Work vector size
  size_t worksize_;
 
  std::vector<SXElem> free_vars_;
 
  std::vector<SXElem> operations_;
 
  std::vector<SXElem> constants_;
 
  std::vector<double> default_in_;
 
  std::vector<bool> copy_elision_;
 
  bool print_instructions_;
 
  void serialize_body(SerializingStream &s) const override;
 
  // call node information that won't fit into AlgEl
  struct ExtendedAlgEl {
    ExtendedAlgEl(const Function& fun);
    Function f;
    // Work vector indices of the arguments (cfr AlgEl::arg)
    std::vector<int> dep;
    // Work vector indices of the results (cfr AlgEl::res)
    std::vector<int> res;
 
    std::vector<int> copy_elision_arg;
    std::vector<int> copy_elision_offset;
 
    // Following fields are redundant but will increase eval speed
    casadi_int n_dep;
    casadi_int n_res;
    casadi_int f_n_in;
    casadi_int f_n_out;
    std::vector<int> f_nnz_in;
    std::vector<int> f_nnz_out;
  };
 
  struct CallInfo {
    // Maximum memory requirements across all call nodes
    size_t sz_arg = 0, sz_res = 0, sz_iw = 0, sz_w = 0;
    size_t sz_w_arg = 0, sz_w_res = 0;
    std::vector<ExtendedAlgEl> el;
  } call_;
 
  static ProtoFunction* deserialize(DeserializingStream& s);
 
  static std::vector<SX> order(const std::vector<SX>& expr);
 
 
  static const Options options_;
  const Options& get_options() const override { return options_;}
 
  Dict generate_options(const std::string& target="clone") const override;
 
  void init(const Dict& opts) override;
 
  void init_copy_elision();
 
  size_t codegen_sz_w(const CodeGenerator& g) const override;
 
  void codegen_declarations(CodeGenerator& g) const override;
 
  void codegen_body(CodeGenerator& g) const override;
 
  int sp_forward(const bvec_t** arg, bvec_t** res,
                  casadi_int* iw, bvec_t* w, void* mem) const override;
 
  int sp_reverse(bvec_t** arg, bvec_t** res, casadi_int* iw, bvec_t* w, void* mem) const override;
 
  SX instructions_sx() const override;
 
  // Get all embedded functions, recursively
  void find(std::map<FunctionInternal*, Function>& all_fun, casadi_int max_depth) const override;
 
  double get_default_in(casadi_int ind) const override { return default_in_.at(ind);}
 
  void export_code_body(const std::string& lang,
    std::ostream &stream, const Dict& options) const override;
 
  bool just_in_time_opencl_;
 
  bool just_in_time_sparsity_;
 
  bool live_variables_;
 
protected:
  template<typename T>
  void call_fwd(const AlgEl& e, const T** arg, T** res, casadi_int* iw, T* w) const;
 
  template<typename T>
  void call_rev(const AlgEl& e, T** arg, T** res, casadi_int* iw, T* w) const;
 
  template<typename T, typename CT>
  void call_setup(const ExtendedAlgEl& m,
    CT*** call_arg, T*** call_res, casadi_int** call_iw, T** call_w, T** nz_in, T** nz_out) const;
 
  explicit SXFunction(DeserializingStream& s);
};
 
 
} // namespace casadi
 
#endif // CASADI_SX_FUNCTION_HPP