GraphProbe.h

Go to the documentation of this file.

 
#pragma once
 
#include <tvm/api.h>
 
#include <tvm/Variable.h>
#include <tvm/diagnostic/internal/probe.h>
#include <tvm/diagnostic/internal/traits.h>
#include <tvm/diagnostic/matrix.h>
#include <tvm/graph/CallGraph.h>
#include <tvm/internal/MatrixWithProperties.h>
 
#include <mpark/variant.hpp>
 
#include <Eigen/Core>
 
#include <iosfwd>
#include <map>
#include <memory>
#include <optional>
#include <tuple>
#include <vector>
 
namespace tvm::diagnostic
{
class TVM_DLLAPI GraphProbe
{
public:
  using Output = graph::internal::Log::Output;
  using Update = graph::internal::Log::Update;
  using OutputVal = std::tuple<Output, VariablePtr, Eigen::MatrixXd>;
 
  struct TVM_DLLAPI ProbeNode
  {
    ProbeNode(const Output & o) : val(o) {}
    ProbeNode(const Update & u) : val(u) {}
    mpark::variant<Output, Update> val;
    std::vector<std::unique_ptr<ProbeNode>> children;
 
    ProbeNode & operator=(const ProbeNode &) = delete;
    ProbeNode(const ProbeNode &) = delete;
  };
 
  static constexpr auto absInRange = [](double rmin, double rmax) {
    return [rmin, rmax](const Eigen::MatrixXd & M) { return hasElemAbsInRange(M, rmin, rmax); };
  };
 
  static constexpr auto hasNan = [](const Eigen::MatrixXd & M) { return M.hasNaN(); };
 
  GraphProbe(const graph::internal::Log & log = tvm::graph::internal::Logger::logger().log());
 
  template<typename T, typename MethodT, typename EnumOutput, typename ConvertT = std::nullopt_t>
  void registerAccessor(EnumOutput o, MethodT method, ConvertT convert = std::nullopt);
 
  template<typename T>
  void registerTVMFunction();
 
  template<typename T>
  void registerTVMConstraint();
 
  template<typename T>
  void registerTVMTaskDynamics();
 
  std::vector<OutputVal> listOutputVal(const graph::CallGraph * const g, bool verbose = false) const;
 
  std::vector<OutputVal> listOutputVal(const Output & o, bool verbose = false) const;
 
  std::unique_ptr<ProbeNode> followUp(
      const Output & o,
      const std::function<bool(const Eigen::MatrixXd &)> & select = [](const Eigen::MatrixXd &) { return true; }) const;
 
  std::vector<std::unique_ptr<ProbeNode>> followUp(
      const graph::CallGraph * const g,
      const std::function<bool(const Eigen::MatrixXd &)> & select = [](const Eigen::MatrixXd &) { return true; }) const;
 
  void print(std::ostream & os, const std::unique_ptr<ProbeNode> & root) const;
  void print(std::ostream & os, const std::vector<std::unique_ptr<ProbeNode>> & roots) const;
 
private:
  struct Processed
  {
    Processed() = default;
    Processed(const graph::internal::Log & log);
 
    std::map<Output, bool> outputs;
    std::map<Update, bool> updates;
    std::vector<Output> allOutputs;
  };
 
  bool addOutputVal(std::vector<OutputVal> & ov, const Output & o) const;
  std::vector<OutputVal> outputVal(const Output & o) const;
  std::vector<OutputVal> listOutputVal(const std::vector<graph::internal::Log::Output> & o, bool verbose) const;
  std::unique_ptr<ProbeNode> followUp(const Output & o,
                                      const std::function<bool(const Eigen::MatrixXd &)> & select,
                                      Processed & processed) const;
  std::unique_ptr<ProbeNode> followUp(const Update & u,
                                      const std::function<bool(const Eigen::MatrixXd &)> & select,
                                      Processed & processed) const;
 
  void print(std::ostream & os, const std::unique_ptr<ProbeNode> & node, int depth) const;
 
  using OutputKey = std::pair<std::size_t, graph::internal::Log::EnumValue>;
  using VarMatrixPair = std::pair<VariablePtr, Eigen::MatrixXd>;
  const graph::internal::Log & log_;
  std::unordered_map<OutputKey, std::function<Eigen::MatrixXd(uintptr_t)>, internal::PairHasher> outputAccessor_;
  std::unordered_map<OutputKey, std::function<std::vector<VarMatrixPair>(uintptr_t)>, internal::PairHasher>
      varDepOutputAccessor_;
};
 
template<typename T, typename MethodT, typename EnumOutput, typename ConvertT>
inline void GraphProbe::registerAccessor(EnumOutput o, MethodT method, ConvertT convertIn)
{
  using CheckAccessor = internal::CheckAccessor<T, MethodT>;
  if constexpr(CheckAccessor::isVoidAccessor)
  {
    using ReturnT = typename CheckAccessor::ReturnT;
    auto convert = internal::MakeConvert<ReturnT>(std::move(convertIn));
    OutputKey k{std::type_index(typeid(T)).hash_code(), tvm::graph::internal::Log::EnumValue(o)};
    outputAccessor_[k] = [method, convert](uintptr_t t) { return convert((reinterpret_cast<T *>(t)->*method)()); };
  }
  else if constexpr(CheckAccessor::isVariableAccessor)
  {
    using ReturnT = typename CheckAccessor::ReturnT;
    auto convert = internal::MakeConvert<ReturnT>(std::move(convertIn));
    OutputKey k{std::type_index(typeid(T)).hash_code(), tvm::graph::internal::Log::EnumValue(o)};
    varDepOutputAccessor_[k] = [method, convert](uintptr_t t) {
      std::vector<VarMatrixPair> ret;
      T * ptr = reinterpret_cast<T *>(t);
      for(const auto & v : ptr->variables())
        ret.emplace_back(v, convert((ptr->*method)(*v)));
      return ret;
    };
  }
  else
  {
    static_assert(!std::is_same_v<T, T>, "Provided method does not have a valid signature");
  }
}
 
template<typename T>
inline void GraphProbe::registerTVMFunction()
{
  using GetVectorT = const Eigen::VectorXd & (T::*)() const;
  using GetJacobianT = tvm::internal::MatrixConstRefWithProperties (T::*)(const Variable &) const;
  using GetJDotT = MatrixConstRef (T::*)(const Variable &) const;
  registerAccessor<T>(T::Output::Value, static_cast<GetVectorT>(&T::value));
  registerAccessor<T>(T::Output::Velocity, static_cast<GetVectorT>(&T::velocity));
  registerAccessor<T>(T::Output::Jacobian, static_cast<GetJacobianT>(&T::jacobian));
  registerAccessor<T>(T::Output::NormalAcceleration, static_cast<GetVectorT>(&T::normalAcceleration));
  registerAccessor<T>(T::Output::JDot, static_cast<GetJDotT>(&T::JDot));
}
 
template<typename T>
inline void GraphProbe::registerTVMConstraint()
{
  using GetVectorT = const Eigen::VectorXd & (T::*)() const;
  using GetJacobianT = tvm::internal::MatrixConstRefWithProperties (T::*)(const Variable &) const;
  registerAccessor<T>(T::Output::Value, static_cast<GetVectorT>(&T::value));
  registerAccessor<T>(T::Output::Jacobian, static_cast<GetJacobianT>(&T::jacobian));
  registerAccessor<T>(T::Output::L, static_cast<GetVectorT>(&T::l));
  registerAccessor<T>(T::Output::U, static_cast<GetVectorT>(&T::u));
  registerAccessor<T>(T::Output::E, static_cast<GetVectorT>(&T::e));
}
 
template<typename T>
inline void GraphProbe::registerTVMTaskDynamics()
{
  using GetVectorT = const Eigen::VectorXd & (T::Impl::*)() const;
  registerAccessor<typename T::Impl>(T::Impl::Output::Value, static_cast<GetVectorT>(&T::Impl::value));
}
} // namespace tvm::diagnostic
 
TVM_DLLAPI std::ostream & operator<<(std::ostream & os,
                                     const std::vector<tvm::diagnostic::GraphProbe::OutputVal> & vals);