Inputs.h

Go to the documentation of this file.

 
#pragma once
 
#include <tvm/graph/abstract/Outputs.h>
 
#include <memory>
#include <set>
#include <unordered_map>
#include <unordered_set>
 
namespace tvm
{
 
namespace graph
{
 
class CallGraph;
 
namespace internal
{
 
class TVM_DLLAPI Inputs
{
public:
  friend class tvm::graph::CallGraph;
 
  virtual ~Inputs() = default;
 
  using inputs_t = std::unordered_map<abstract::Outputs *, std::set<int>>;
  using store_t = std::unordered_set<std::shared_ptr<abstract::Outputs>>;
 
  struct TVM_DLLAPI Iterator : public inputs_t::iterator
  {
    Iterator(inputs_t::iterator it, inputs_t::iterator end);
    operator bool();
 
  private:
    inputs_t::iterator end_;
  };
 
  template<typename T, typename EnumI, typename... Args>
  void addInput(std::shared_ptr<T> source, EnumI i, Args... args);
 
  template<typename T,
           typename EnumI,
           typename... Args,
           typename std::enable_if<std::is_base_of<abstract::Outputs, T>::value, int>::type = 0>
  void addInput(T & source, EnumI i, Args... args);
  template<typename T>
  void removeInput(T * source);
  template<typename T, typename... Args>
  void removeInput(T * source, Args... args);
  template<typename T>
  Iterator getInput(T * source);
  template<typename T>
  Iterator getInput(const std::shared_ptr<T> & source);
 
private:
  void removeInput(Iterator it, abstract::Outputs * source);
  inputs_t inputs_;
  store_t store_;
  template<typename T, typename EnumI>
  void addInput(T * source, EnumI i);
  template<typename T>
  void addInput(T &)
  {}
};
 
} // namespace internal
 
} // namespace graph
 
} // namespace tvm
 
#include "Inputs.hpp"