graph.h

Go to the documentation of this file.

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
 
#ifndef NNVM_GRAPH_H_
#define NNVM_GRAPH_H_
 
#include <algorithm>
#include <memory>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>
 
#include "base.h"
#include "node.h"
#include "symbolic.h"
 
namespace nnvm {
 
class IndexedGraph;
 
class Graph {
 public:
  std::vector<NodeEntry> outputs;
  std::unordered_map<std::string, std::shared_ptr<any> > attrs;
  template <typename T>
  inline const T& GetAttr(const std::string& attr_name) const;
  inline bool HasAttr(const std::string& attr_name) const;
  template <typename T>
  inline T MoveCopyAttr(const std::string& attr_name);
  const IndexedGraph& indexed_graph() const;
 
 private:
  // internal structure of indexed graph
  mutable std::shared_ptr<const IndexedGraph> indexed_graph_;
};
 
class IndexedGraph {
 public:
  struct NodeEntry {
    uint32_t node_id;
    uint32_t index;
    uint32_t version;
  };
  struct Node {
    const nnvm::Node* source;
    array_view<NodeEntry> inputs;
    array_view<uint32_t> control_deps;
    std::weak_ptr<nnvm::Node> weak_ref;
  };
  inline size_t num_nodes() const { return nodes_.size(); }
  inline size_t num_node_entries() const { return entry_rptr_.back(); }
  inline uint32_t entry_id(uint32_t node_id, uint32_t index) const {
    return entry_rptr_[node_id] + index;
  }
  inline uint32_t entry_id(const NodeEntry& e) const { return entry_rptr_[e.node_id] + e.index; }
  inline uint32_t entry_id(const nnvm::NodeEntry& e) const {
    return entry_rptr_[node_id(e.node.get())] + e.index;
  }
  inline uint32_t node_id(const nnvm::Node* node) const { return node2index_.at(node); }
  inline const Node& operator[](uint32_t node_id) const { return nodes_[node_id]; }
  inline const Node& operator[](const nnvm::Node* node) const { return nodes_[node_id(node)]; }
  inline const std::vector<uint32_t>& input_nodes() const { return input_nodes_; }
  inline const std::unordered_set<uint32_t>& mutable_input_nodes() const {
    return mutable_input_nodes_;
  }
  inline const std::vector<NodeEntry>& outputs() const { return outputs_; }
 
  inline bool exist(const nnvm::Node* node) const { return node2index_.count(node); }
 
  // disalllow copy assign
  IndexedGraph(const IndexedGraph&) = delete;
 
 private:
  friend class Graph;
  explicit IndexedGraph(const Graph& other);
  // Node pointers in CSR structure.
  std::vector<Node> nodes_;
  // Index to all input nodes.
  std::vector<uint32_t> input_nodes_;
  // Index to all mutable input nodes.
  std::unordered_set<uint32_t> mutable_input_nodes_;
  // space to store the outputs entries
  std::vector<NodeEntry> outputs_;
  // mapping from node to index.
  std::unordered_map<const nnvm::Node*, uint32_t> node2index_;
  // CSR pointer of node entries
  std::vector<size_t> entry_rptr_;
  // space to store input entries of each
  std::vector<NodeEntry> input_entries_;
  // control flow dependencies
  std::vector<uint32_t> control_deps_;
};
 
template <typename FVisit>
inline void DFSVisit(const std::vector<NodeEntry>& heads, FVisit fvisit);
 
// inline function implementations
template <typename T>
inline const T& Graph::GetAttr(const std::string& attr_name) const {
  auto it = attrs.find(attr_name);
  CHECK(it != attrs.end()) << "Cannot find attribute " << attr_name << " in the graph";
  return nnvm::unsafe_get<T>(*it->second);
}
 
inline bool Graph::HasAttr(const std::string& attr_name) const {
  auto it = attrs.find(attr_name);
  return it != attrs.end();
}
 
template <typename T>
inline T Graph::MoveCopyAttr(const std::string& attr_name) {
  auto it = attrs.find(attr_name);
  CHECK(it != attrs.end()) << "Cannot find attribute " << attr_name << " in the graph";
  std::shared_ptr<any> sptr = it->second;
  attrs.erase(it);
  if (sptr.unique()) {
    return std::move(nnvm::get<T>(*sptr));
  } else {
    return nnvm::get<T>(*sptr);
  }
}
 
template <typename GNode, typename HashType, typename FVisit, typename HashFunc, typename InDegree,
          typename GetInput>
void PostOrderDFSVisit(const std::vector<GNode>& heads, FVisit fvisit, HashFunc hash,
                       InDegree indegree, GetInput getinput) {
  std::vector<std::pair<GNode, uint32_t> > stack;
  std::unordered_set<HashType> visited;
  for (auto& head : heads) {
    HashType head_hash = hash(head);
    if (visited.count(head_hash) == 0) {
      stack.push_back(std::make_pair(head, 0));
      visited.insert(head_hash);
    }
    while (!stack.empty()) {
      std::pair<GNode, uint32_t>& back = stack.back();
      if (back.second == indegree(back.first)) {
        fvisit(back.first);
        stack.pop_back();
      } else {
        const GNode& input = getinput(back.first, back.second++);
        HashType input_hash = hash(input);
        if (visited.count(input_hash) == 0) {
          stack.push_back(std::make_pair(input, 0));
          visited.insert(input_hash);
        }
      }
    }
  }
}
 
template <typename FVisit>
inline void DFSVisit(const std::vector<NodeEntry>& heads, FVisit fvisit) {
  typedef const ObjectPtr* GNode;
  std::vector<GNode> head_nodes(heads.size());
  std::transform(heads.begin(), heads.end(), head_nodes.begin(),
                 [](const NodeEntry& e) -> GNode { return &e.node; });
  PostOrderDFSVisit<GNode, Node*>(
      head_nodes, [fvisit](GNode n) { fvisit(*n); },  // FVisit
      [](GNode n) -> Node* { return n->get(); },      // HashFunc
      [](GNode n) -> uint32_t {                       // InDegree
        if (!(*n)) return 0;
        return (*n)->inputs.size() + (*n)->control_deps.size();
      },
      [](GNode n, uint32_t index) -> GNode {  // GetInput
        if (index < (*n)->inputs.size()) {
          return &(*n)->inputs.at(index).node;
        } else {
          return &(*n)->control_deps.at(index - (*n)->inputs.size());
        }
      });
}
 
}  // namespace nnvm
 
#endif  // NNVM_GRAPH_H_