CallGraph.h

Go to the documentation of this file.

 
#pragma once
 
#include <tvm/api.h>
 
#include <tvm/graph/internal/AbstractNode.h>
#include <tvm/graph/internal/DependencyGraph.h>
 
namespace tvm
{
 
namespace graph
{
 
class TVM_DLLAPI CallGraph
{
public:
  void add(std::shared_ptr<internal::Inputs> inputs);
 
  void update();
 
  inline void execute() const { plan_.execute(); }
 
  void clear();
 
protected:
  struct Call
  {
    internal::AbstractNode * node;
    int id;
    inline void operator()() const { node->update(id); }
  };
 
  struct CompareCall
  {
    bool operator()(const Call & c1, const Call & c2) const
    { return (c1.node < c2.node) || (c1.node == c2.node && c1.id < c2.id); }
  };
 
  struct Plan
  {
    void build(const CallGraph & graph);
 
    void clear();
 
    inline void execute() const
    {
      for(auto & c : plan_)
      {
        c();
      }
    }
 
  private:
    std::vector<Call> plan_;
  };
 
protected:
  std::vector<int> addOutput(abstract::Outputs * source, int output);
 
  int addCall(Call c);
 
  std::vector<std::shared_ptr<internal::Inputs>> inputs_;
  std::map<Call, int, CompareCall> callId_;
  std::vector<Call> calls_;
  internal::DependencyGraph dependencyGraph_;
 
  std::map<std::intptr_t, std::map<int, std::vector<int>>> visited_;
 
  Plan plan_;
};
 
} // namespace graph
 
} // namespace tvm