introduced a MaxVirtualStaticRows compile time constant and realized as many static rows specific virtual reverseAD methods in the CallRecord interface to speedup the Jacobian evaluatio.
HannesSommer
parent
4966f5a942
commit
fb24ab586e
|
|
@ -72,27 +72,119 @@ public:
|
|||
};
|
||||
|
||||
//-----------------------------------------------------------------------------
|
||||
/**
|
||||
* MaxVirtualStaticRows defines how many separate virtual reverseAD with specific
|
||||
* static rows (1..MaxVirtualStaticRows) methods will be part of the CallRecord interface.
|
||||
*/
|
||||
const int MaxVirtualStaticRows = 4;
|
||||
|
||||
namespace internal {
|
||||
/**
|
||||
* ConvertToDynamicIf converts to a dense matrix with dynamic rows iff ConvertToDynamicRows (colums stay as they are) otherwise
|
||||
* it just passes dense Eigen matrices through.
|
||||
*/
|
||||
template <bool ConvertToDynamicRows>
|
||||
struct ConvertToDynamicRowsIf {
|
||||
template <typename Derived>
|
||||
static Eigen::Matrix<double, Eigen::Dynamic, Derived::ColsAtCompileTime> convert(const Eigen::MatrixBase<Derived> & x){
|
||||
return x;
|
||||
}
|
||||
};
|
||||
template <>
|
||||
struct ConvertToDynamicRowsIf<false> {
|
||||
template <int Rows, int Cols>
|
||||
static const Eigen::Matrix<double, Rows, Cols> & convert(const Eigen::Matrix<double, Rows, Cols> & x){
|
||||
return x;
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* Recursive definition of an interface having several purely
|
||||
* virtual _reverseAD(const Eigen::Matrix<double, Rows, Cols> &, JacobianMap&)
|
||||
* with Rows in 1..MaxSupportedStaticRows
|
||||
*/
|
||||
template<int MaxSupportedStaticRows, int Cols>
|
||||
struct ReverseADInterface : public ReverseADInterface < MaxSupportedStaticRows - 1, Cols> {
|
||||
protected:
|
||||
using ReverseADInterface < MaxSupportedStaticRows - 1, Cols>::_reverseAD;
|
||||
virtual void _reverseAD(const Eigen::Matrix<double, MaxSupportedStaticRows, Cols> & dFdT, JacobianMap& jacobians) const = 0;
|
||||
};
|
||||
|
||||
template<int Cols>
|
||||
struct ReverseADInterface<0, Cols> {
|
||||
protected:
|
||||
void _reverseAD(){} //dummy to allow the using directive in the template without failing for MaxSupportedStaticRows == 1.
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* The CallRecord class stores the Jacobians of applying a function
|
||||
* with respect to each of its arguments. It also stores an executation trace
|
||||
* (defined below) for each of its arguments.
|
||||
*
|
||||
* It is sub-classed in the function-style ExpressionNode sub-classes below.
|
||||
* It is implemented in the function-style ExpressionNode's nested Record class below.
|
||||
*/
|
||||
template<int COLS>
|
||||
struct CallRecord {
|
||||
static size_t const N = 0;
|
||||
virtual void print(const std::string& indent) const {
|
||||
template<int Cols>
|
||||
struct CallRecord : private internal::ReverseADInterface<MaxVirtualStaticRows, Cols> {
|
||||
inline void print(const std::string& indent) const {
|
||||
_print(indent);
|
||||
}
|
||||
virtual void startReverseAD(JacobianMap& jacobians) const {
|
||||
inline void startReverseAD(JacobianMap& jacobians) const {
|
||||
_startReverseAD(jacobians);
|
||||
}
|
||||
virtual void reverseAD(const Matrix& dFdT, JacobianMap& jacobians) const {
|
||||
template <int Rows>
|
||||
inline void reverseAD(const Eigen::Matrix<double, Rows, Cols> & dFdT, JacobianMap& jacobians) const{
|
||||
_reverseAD(internal::ConvertToDynamicRowsIf<(Rows > MaxVirtualStaticRows)>::convert(dFdT), jacobians);
|
||||
}
|
||||
typedef Eigen::Matrix<double, 2, COLS> Jacobian2T;
|
||||
virtual void reverseAD2(const Jacobian2T& dFdT,
|
||||
JacobianMap& jacobians) const {
|
||||
virtual ~CallRecord() {
|
||||
}
|
||||
private:
|
||||
using internal::ReverseADInterface<MaxVirtualStaticRows, Cols>::_reverseAD;
|
||||
virtual void _print(const std::string& indent) const = 0;
|
||||
virtual void _startReverseAD(JacobianMap& jacobians) const = 0;
|
||||
virtual void _reverseAD(const Eigen::Matrix<double, Eigen::Dynamic, Cols> & dFdT, JacobianMap& jacobians) const = 0;
|
||||
virtual void _reverseAD(const Matrix & dFdT, JacobianMap& jacobians) const = 0;
|
||||
};
|
||||
|
||||
namespace internal {
|
||||
/**
|
||||
* ReverseADImplementor is a utility class used by CallRecordImplementor to implementing the recursive CallRecord interface.
|
||||
*/
|
||||
template <typename Derived, int MaxSupportedStaticRows, int Cols>
|
||||
struct ReverseADImplementor : ReverseADImplementor<Derived, MaxSupportedStaticRows - 1, Cols> {
|
||||
protected:
|
||||
virtual void _reverseAD(const Eigen::Matrix<double, MaxSupportedStaticRows, Cols> & dFdT, JacobianMap& jacobians) const {
|
||||
static_cast<const Derived *>(this)->reverseAD(dFdT, jacobians);
|
||||
}
|
||||
};
|
||||
template<typename Derived, int Cols>
|
||||
struct ReverseADImplementor<Derived, 0, Cols> : CallRecord<Cols> {
|
||||
};
|
||||
|
||||
/**
|
||||
* The CallRecordImplementor implements the CallRecord interface for a Derived class by
|
||||
* delegating to its corresponding (templated) non-virtual methods.
|
||||
*/
|
||||
template<typename Derived, int Cols>
|
||||
struct CallRecordImplementor : public ReverseADImplementor<Derived, MaxVirtualStaticRows, Cols> {
|
||||
private:
|
||||
const Derived & derived() const {
|
||||
return static_cast<const Derived&>(*this);
|
||||
}
|
||||
virtual void _print(const std::string& indent) const {
|
||||
derived().print(indent);
|
||||
}
|
||||
virtual void _startReverseAD(JacobianMap& jacobians) const {
|
||||
derived().startReverseAD(jacobians);
|
||||
}
|
||||
virtual void _reverseAD(const Eigen::Matrix<double, Eigen::Dynamic, Cols> & dFdT, JacobianMap& jacobians) const {
|
||||
derived().reverseAD(dFdT, jacobians);
|
||||
}
|
||||
virtual void _reverseAD(const Matrix & dFdT, JacobianMap& jacobians) const {
|
||||
derived().reverseAD(dFdT, jacobians);
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
//-----------------------------------------------------------------------------
|
||||
/// Handle Leaf Case: reverseAD ends here, by writing a matrix into Jacobians
|
||||
|
|
@ -101,10 +193,10 @@ void handleLeafCase(const Eigen::Matrix<double, ROWS, COLS>& dTdA,
|
|||
JacobianMap& jacobians, Key key) {
|
||||
jacobians(key).block<ROWS, COLS>(0, 0) += dTdA; // block makes HUGE difference
|
||||
}
|
||||
/// Handle Leaf Case for Dynamic Matrix type (slower)
|
||||
template<>
|
||||
/// Handle Leaf Case for Dynamic ROWS Matrix type (slower)
|
||||
template<int COLS>
|
||||
inline void handleLeafCase(
|
||||
const Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic>& dTdA,
|
||||
const Eigen::Matrix<double, Eigen::Dynamic, COLS>& dTdA,
|
||||
JacobianMap& jacobians, Key key) {
|
||||
jacobians(key) += dTdA;
|
||||
}
|
||||
|
|
@ -193,45 +285,18 @@ public:
|
|||
content.ptr->startReverseAD(jacobians);
|
||||
}
|
||||
// Either add to Jacobians (Leaf) or propagate (Function)
|
||||
void reverseAD(const Matrix& dTdA, JacobianMap& jacobians) const {
|
||||
template <typename DerivedMatrix>
|
||||
void reverseAD(const Eigen::MatrixBase<DerivedMatrix> & dTdA, JacobianMap& jacobians) const {
|
||||
if (kind == Leaf)
|
||||
handleLeafCase(dTdA, jacobians, content.key);
|
||||
handleLeafCase(dTdA.eval(), jacobians, content.key);
|
||||
else if (kind == Function)
|
||||
content.ptr->reverseAD(dTdA, jacobians);
|
||||
}
|
||||
// Either add to Jacobians (Leaf) or propagate (Function)
|
||||
typedef Eigen::Matrix<double, 2, Dim> Jacobian2T;
|
||||
void reverseAD2(const Jacobian2T& dTdA, JacobianMap& jacobians) const {
|
||||
if (kind == Leaf)
|
||||
handleLeafCase(dTdA, jacobians, content.key);
|
||||
else if (kind == Function)
|
||||
content.ptr->reverseAD2(dTdA, jacobians);
|
||||
content.ptr->reverseAD(dTdA.eval(), jacobians);
|
||||
}
|
||||
|
||||
/// Define type so we can apply it as a meta-function
|
||||
typedef ExecutionTrace<T> type;
|
||||
};
|
||||
|
||||
/// Primary template calls the generic Matrix reverseAD pipeline
|
||||
template<size_t ROWS, class A>
|
||||
struct Select {
|
||||
typedef Eigen::Matrix<double, ROWS, traits::dimension<A>::value> Jacobian;
|
||||
static void reverseAD(const ExecutionTrace<A>& trace, const Jacobian& dTdA,
|
||||
JacobianMap& jacobians) {
|
||||
trace.reverseAD(dTdA, jacobians);
|
||||
}
|
||||
};
|
||||
|
||||
/// Partially specialized template calls the 2-dimensional output version
|
||||
template<class A>
|
||||
struct Select<2, A> {
|
||||
typedef Eigen::Matrix<double, 2, traits::dimension<A>::value> Jacobian;
|
||||
static void reverseAD(const ExecutionTrace<A>& trace, const Jacobian& dTdA,
|
||||
JacobianMap& jacobians) {
|
||||
trace.reverseAD2(dTdA, jacobians);
|
||||
}
|
||||
};
|
||||
|
||||
//-----------------------------------------------------------------------------
|
||||
/**
|
||||
* Expression node. The superclass for objects that do the heavy lifting
|
||||
|
|
@ -479,8 +544,15 @@ template<class T>
|
|||
struct FunctionalBase: ExpressionNode<T> {
|
||||
static size_t const N = 0; // number of arguments
|
||||
|
||||
typedef CallRecord<traits::dimension<T>::value> Record;
|
||||
|
||||
struct Record {
|
||||
void print(const std::string& indent) const {
|
||||
}
|
||||
void startReverseAD(JacobianMap& jacobians) const {
|
||||
}
|
||||
template <typename SomeMatrix>
|
||||
void reverseAD(const SomeMatrix & dFdT, JacobianMap& jacobians) const {
|
||||
}
|
||||
};
|
||||
/// Construct an execution trace for reverse AD
|
||||
void trace(const Values& values, Record* record, char*& raw) const {
|
||||
// base case: does not do anything
|
||||
|
|
@ -539,7 +611,7 @@ struct GenerateFunctionalNode: Argument<T, A, Base::N + 1>, Base {
|
|||
typedef JacobianTrace<T, A, N> This;
|
||||
|
||||
/// Print to std::cout
|
||||
virtual void print(const std::string& indent) const {
|
||||
void print(const std::string& indent) const {
|
||||
Base::Record::print(indent);
|
||||
static const Eigen::IOFormat matlab(0, 1, " ", "; ", "", "", "[", "]");
|
||||
std::cout << This::dTdA.format(matlab) << std::endl;
|
||||
|
|
@ -547,25 +619,17 @@ struct GenerateFunctionalNode: Argument<T, A, Base::N + 1>, Base {
|
|||
}
|
||||
|
||||
/// Start the reverse AD process
|
||||
virtual void startReverseAD(JacobianMap& jacobians) const {
|
||||
void startReverseAD(JacobianMap& jacobians) const {
|
||||
Base::Record::startReverseAD(jacobians);
|
||||
Select<traits::dimension<T>::value, A>::reverseAD(This::trace, This::dTdA,
|
||||
jacobians);
|
||||
This::trace.reverseAD(This::dTdA, jacobians);
|
||||
}
|
||||
|
||||
/// Given df/dT, multiply in dT/dA and continue reverse AD process
|
||||
virtual void reverseAD(const Matrix& dFdT, JacobianMap& jacobians) const {
|
||||
template <int Rows, int Cols>
|
||||
void reverseAD(const Eigen::Matrix<double, Rows, Cols> & dFdT, JacobianMap& jacobians) const {
|
||||
Base::Record::reverseAD(dFdT, jacobians);
|
||||
This::trace.reverseAD(dFdT * This::dTdA, jacobians);
|
||||
}
|
||||
|
||||
/// Version specialized to 2-dimensional output
|
||||
typedef Eigen::Matrix<double, 2, traits::dimension<T>::value> Jacobian2T;
|
||||
virtual void reverseAD2(const Jacobian2T& dFdT,
|
||||
JacobianMap& jacobians) const {
|
||||
Base::Record::reverseAD2(dFdT, jacobians);
|
||||
This::trace.reverseAD2(dFdT * This::dTdA, jacobians);
|
||||
}
|
||||
};
|
||||
|
||||
/// Construct an execution trace for reverse AD
|
||||
|
|
@ -619,8 +683,16 @@ struct FunctionalNode {
|
|||
return static_cast<Argument<T, A, N> const &>(*this).expression;
|
||||
}
|
||||
|
||||
/// Provide convenience access to Record storage
|
||||
struct Record: public Base::Record {
|
||||
/// Provide convenience access to Record storage and implement
|
||||
/// the virtual function based interface of CallRecord using the CallRecordImplementor
|
||||
struct Record:
|
||||
public internal::CallRecordImplementor<Record, traits::dimension<T>::value>,
|
||||
public Base::Record {
|
||||
using Base::Record::print;
|
||||
using Base::Record::startReverseAD;
|
||||
using Base::Record::reverseAD;
|
||||
|
||||
virtual ~Record(){}
|
||||
|
||||
/// Access Value
|
||||
template<class A, size_t N>
|
||||
|
|
@ -633,7 +705,6 @@ struct FunctionalNode {
|
|||
typename Jacobian<T, A>::type& jacobian() {
|
||||
return static_cast<JacobianTrace<T, A, N>&>(*this).dTdA;
|
||||
}
|
||||
|
||||
};
|
||||
|
||||
/// Construct an execution trace for reverse AD
|
||||
|
|
|
|||
Loading…
Reference in New Issue