ExecutionTrace.h

Go to the documentation of this file.

/* ----------------------------------------------------------------------------

 * GTSAM Copyright 2010, Georgia Tech Research Corporation,
 * Atlanta, Georgia 30332-0415
 * All Rights Reserved
 * Authors: Frank Dellaert, et al. (see THANKS for the full author list)

 * See LICENSE for the license information

 * -------------------------------------------------------------------------- */

#pragma once

#include <gtsam/config.h>      // Configuration from CMake
#include <gtsam/nonlinear/internal/JacobianMap.h>
#include <gtsam/inference/Key.h>
#include <gtsam/base/Manifold.h>

#include <Eigen/Core>

#include <iostream>
#include <optional>
#include <string>
#include <type_traits>

namespace gtsam {
namespace internal {

template<int T> struct CallRecord;

#ifdef _MSC_VER
// TODO(dellaert): this might lead to trouble if Eigen decides to use 32 on Windows.
static const unsigned TraceAlignment = 16; // 16 bytes max_align on Windows
#else
static const unsigned TraceAlignment = 32; // Alignment used by Eigen on some platforms.
#endif
// TODO(dellaert): we *should* be able to simplify the code using the pointer arithmetic from ExecutionTraceStorage.
typedef std::aligned_storage<1, TraceAlignment>::type ExecutionTraceStorage;

template<bool UseBlock, typename Derived>
struct UseBlockIf {
  static void addToJacobian(const Eigen::MatrixBase<Derived>& dTdA,
      JacobianMap& jacobians, Key key) {
    // block makes HUGE difference
    jacobians(key).block<Derived::RowsAtCompileTime, Derived::ColsAtCompileTime>(
        0, 0) += dTdA;
  }
};

template<typename Derived>
struct UseBlockIf<false, Derived> {
  static void addToJacobian(const Eigen::MatrixBase<Derived>& dTdA,
      JacobianMap& jacobians, Key key) {
    jacobians(key) += dTdA;
  }
};

template<typename Derived>
void handleLeafCase(const Eigen::MatrixBase<Derived>& dTdA,
    JacobianMap& jacobians, Key key) {
  UseBlockIf<
      Derived::RowsAtCompileTime != Eigen::Dynamic
          && Derived::ColsAtCompileTime != Eigen::Dynamic, Derived>::addToJacobian(
      dTdA, jacobians, key);
}

template<class T>
class ExecutionTrace {
  static const int Dim = traits<T>::dimension;
  typedef Eigen::Matrix<double, Dim, Dim> JacobianTT;
  enum {
    Constant, Leaf, Function
  } kind;
  union {
    Key key;
    CallRecord<Dim>* ptr;
  } content;

 public:

  ExecutionTrace() :
      kind(Constant) {
  }

  void setLeaf(Key key) {
    kind = Leaf;
    content.key = key;
  }

  void setFunction(CallRecord<Dim>* record) {
    kind = Function;
    content.ptr = record;
  }

  void print(const std::string& indent = "") const {
    if (kind == Constant) {
      std::cout << indent << "Constant" << std::endl;
    } else if (kind == Leaf) {
      std::cout << indent << "Leaf, key = " << content.key << std::endl;
    } else if (kind == Function) {
      content.ptr->print(indent + "  ");
    }
  }

  template<class Record>
  std::optional<Record*> record() {
    if (kind != Function) {
      return {};
    } else {
      Record* p = dynamic_cast<Record*>(content.ptr);
      return p ? std::optional<Record*>(p) : std::nullopt;
    }
  }

  void startReverseAD1(JacobianMap& jacobians) const {
    if (kind == Leaf) {
      // This branch will only be called on trivial Leaf expressions, i.e. Priors
      static const JacobianTT I = JacobianTT::Identity();
      handleLeafCase(I, jacobians, content.key);
    } else if (kind == Function) {
      // This is the more typical entry point, starting the AD pipeline
      // Inside startReverseAD2 the correctly dimensioned pipeline is chosen.
      content.ptr->startReverseAD2(jacobians);
    }
  }

  // This is a crucial method in selecting the pipeline dimension: since it is templated method
  // it will call the templated Record method reverseAD2. Hence, at this point in the pipeline
  // there are as many reverseAD1 and reverseAD2 versions as there are different Jacobian sizes.
  template<typename DerivedMatrix>
  void reverseAD1(const Eigen::MatrixBase<DerivedMatrix> & dTdA,
      JacobianMap& jacobians) const {
    if (kind == Leaf)
      handleLeafCase(dTdA, jacobians, content.key);
    else if (kind == Function)
      content.ptr->reverseAD2(dTdA, jacobians);
  }

  ~ExecutionTrace() {
    if (kind == Function) {
      content.ptr->~CallRecord<Dim>();
    }
  }

  typedef ExecutionTrace<T> type;
};

} // namespace internal
} // namespace gtsam