graph.h

#pragma once

#include "caffe2/core/common.h"
#include "caffe2/proto/caffe2_pb.h"
#include "caffe2/utils/proto_utils.h"
#include "caffe2/utils/string_utils.h"

#include <algorithm>
#include <unordered_map>
#include <unordered_set>

namespace caffe2 {

namespace transform {

struct CAFFE2_API Node {
 public:
  // Empty constructor for resize
  Node() {}

  // Alternate constructor
  Node(
      const OperatorDef& op,
      bool active,
      std::map<int, std::vector<string>> parents,
      std::map<int, std::vector<string>> children)
      : op(op), active(active), parents(parents), children(children) {}

  // The OperatorDef which this node represents.
  OperatorDef op;

  // Keeps track of if an operator has been deleted through a transformation.
  bool active = true;

  // Stores a pair (idx, blob_list),
  //  idx = index of the child
  //  blob_list = a list of strings, containing the blobs that connect the nodes
  std::map<int, std::vector<string>> parents;
  std::map<int, std::vector<string>> children;
};

struct CAFFE2_API Graph {
 public:
  const std::vector<std::pair<string, int>> GetSubgraphInput(
      const std::vector<int>& subgraph);

  const std::vector<std::pair<string, int>> GetSubgraphOutput(
      const std::vector<int>& subgraph);

  explicit Graph(const NetDef& net_def);

  NetDef GetNetDef();

  void DeactivateSubgraph(std::vector<int> subgraph);

  size_t size() const {
    return nodes_.size();
  }

  void push_node(const Node& new_node) {
    return nodes_.push_back(new_node);
  }

  void resize_nodes(size_t new_size) {
    nodes_.resize(new_size);
  }

  // Index safe, less verbose way to access nodes
  inline const Node& node(size_t idx) const {
    return nodes_.at(idx);
  }

  inline Node& node(size_t idx) {
    return nodes_.at(idx);
  }

  inline bool is_node_active(size_t idx) {
    return node(idx).active;
  }

  inline const std::set<string>& external_input() const {
    return external_input_;
  }

  inline const std::set<string>& external_output() const {
    return external_output_;
  }

 private:
  const std::vector<std::pair<string, int>> GetSubgraphPerimeterHelper(
      bool from_children,
      const std::vector<int>& match);

  // Stores the netdef representation. Is updated upon calls to GetNetDef.
  NetDef netdef_;

  // Stores which blobs the graph reads from, and writes to.
  std::set<string> external_input_;
  std::set<string> external_output_;

  // Keeps track of all the Operators currently within graph, even if inactive.
  std::vector<Node> nodes_;
};

} // namespace transform

// Adds an operator def to a netdef.
// Returns the ptr, if you want to add anything extra (such as device_option)
CAFFE2_API OperatorDef* AddOp(
    NetDef* netdef_ptr,
    string op_type,
    std::vector<string> inputs,
    std::vector<string> outputs);

CAFFE2_API bool MatchStrings(string p, string s);

CAFFE2_API bool MatchArguments(const OperatorDef& p_op, const OperatorDef& g_op);

} // namespace caffe2