blob: 4f93424c667a4920dd853351bc178034c419c800 [file] [log] [blame]
#include "caffe2/core/graph.h"
#include "caffe2/core/common.h"
#include "caffe2/core/logging.h"
#include "caffe2/core/net.h"
#include "caffe2/proto/caffe2_pb.h"
namespace caffe2 {
namespace transform {
Graph::Graph(const NetDef& net) : netdef_(net) {
nodes_.clear();
nodes_.resize(net.op_size());
// Copy over operators
for (int x = 0; x < net.op_size(); x++) {
node(x).op = net.op(x);
}
// For any blob, which operator was the last to write to it?
// In python, this is known as "versions".
std::unordered_map<string, int> edge_parent;
for (int i = 0; i < (int)nodes_.size(); i++) {
for (const string& blob : node(i).op.input()) {
auto it = edge_parent.find(blob);
if (it != edge_parent.end()) {
int j = it->second;
node(i).parents[j].push_back(blob);
node(j).children[i].push_back(blob);
} else {
external_input_.insert(blob);
}
}
for (const string& blob : node(i).op.output()) {
edge_parent[blob] = i;
}
}
// Traverse opposite direction to find external outputs
// For any blob, which operator was the last to read to from it?
std::unordered_map<string, int> edge_child;
for (int i = (int)nodes_.size() - 1; i >= 0; i--) {
for (const string& blob : node(i).op.output()) {
auto it = edge_child.find(blob);
if (it == edge_child.end()) {
external_output_.insert(blob);
}
}
for (const string& blob : node(i).op.input()) {
edge_child[blob] = i;
}
}
}
const std::vector<std::pair<string, int>> Graph::GetSubgraphInput(
const std::vector<int>& match) {
return GetSubgraphPerimeterHelper(true, match);
}
const std::vector<std::pair<string, int>> Graph::GetSubgraphOutput(
const std::vector<int>& match) {
return GetSubgraphPerimeterHelper(false, match);
}
// This helper function will either get:
// 1) a list for the blobs that write INTO a subgraph
// 2) a list of for the blobs that are written FROM a subgraph.
//
// The "from_children" flag determines if it is case 1 (true) or case 2 (false).
const std::vector<std::pair<string, int>> Graph::GetSubgraphPerimeterHelper(
bool from_children,
const std::vector<int>& match) {
std::vector<std::pair<string, int>> edge_list;
std::unordered_set<int> match_set(match.begin(), match.end());
for (int x = 0; x < (int)nodes_.size(); x++) {
if (!is_node_active(x)) {
continue;
}
if (!match_set.count(x)) { // x is not in subgraph
const auto& list = from_children ? node(x).children : node(x).parents;
for (const auto& edge : list) {
int parent = edge.first;
const auto& blobs = edge.second;
if (match_set.count(parent)) { // but has a parent that is in subgraph
for (const string& blob : blobs) {
// NOLINTNEXTLINE(modernize-use-emplace)
edge_list.push_back({blob, x});
}
}
}
}
}
// return the list in sorted order, to allow binary searching
std::sort(edge_list.begin(), edge_list.end());
return edge_list;
}
NetDef Graph::GetNetDef() {
std::vector<bool> visited(nodes_.size(), false);
// Copy over all the properties of the netdef we're based on
NetDef netdef = netdef_;
// But we're going to put in our own operators.
netdef.clear_op();
// Keeps track of the number of parents yet to be processed.
std::vector<int> unchecked_parent_count;
// We will perform a topological traversal on the nodes, but we will prefer
// nodes that come earlier in the execution order.
// This is a min-heap, which stores its elements in ascending order.
// This stores the nodes in the order we process them to be in.
// This guarantees the lowest lexicographical topological ordering.
// This also means the original nodes will be kept in their execution order.
// NOLINTNEXTLINE(modernize-use-transparent-functors)
std::priority_queue<int, std::vector<int>, std::greater<int>> q;
// In our graph, G, the nodes don't have a strict ordering. But in the netdef,
// they must (since nets are operators executed in some order).
// How do we make sure that the order of operators in our generated netdef
// is valid?
// 1) The ordering of the netdef must be topologically sorted, respect to G.
// If A -> B is an edge in the graph G, then A must come before B in the
// netdef's ordering.
// 2) No blob conflicts: If A -> B is an edge in the graph G, and A writes to
// blob X and B reads from blob X, then there cannot be an op that writes
// to blob X between A and B in the ordering.
//
// Perform a Topological Sort, to find an order for the Operators to be in.
// We will keep track of the number of parents each node has.
// We begin with an empty queue, and push in all nodes that do not have any
// parents. Then, we keep track of all unprocessed parents for each node.
// When a node has no more unprocessed parents, we can push it into the queue
// to be processed. This guarantees condition 1 is satisfied.
// TODO(benz): Currently, condition 2 is not guaranteed to be satisified.
// However, giving each blob unique names via SSA will satisfy this condition.
// Then, the resulting graph can be optimized with memonger.
for (int i = 0; i < (int)nodes_.size(); i++) {
unchecked_parent_count.push_back(node(i).parents.size());
if (node(i).parents.size() == 0 && is_node_active(i)) {
q.push(i);
visited[i] = true;
}
}
while (!q.empty()) {
int idx = q.top();
q.pop();
if (!is_node_active(idx)) {
continue;
}
// Creates a new OperatorDef in NetDef
auto& op = *(netdef.add_op());
// Sets it equal to the OperatorDef at node(idx)
op = node(idx).op;
for (const auto& edge : node(idx).children) {
int child = edge.first;
if (!visited[child] && is_node_active(child)) {
unchecked_parent_count[child]--;
if (unchecked_parent_count[child] == 0) {
q.push(child);
visited[child] = true;
}
}
}
}
return netdef;
}
void Graph::DeactivateSubgraph(std::vector<int> subgraph) {
for (int idx : subgraph) {
// remove all edges connected to inactive node
for (const auto& edge : node(idx).parents) {
int parent = edge.first;
node(parent).children.erase(idx);
}
for (const auto& edge : node(idx).children) {
int child = edge.first;
node(child).parents.erase(idx);
}
// actually mark flags as false
node(idx).active = false;
}
}
} // namespace transform
OperatorDef* AddOp(
NetDef* netdef_ptr,
string op_type,
std::vector<string> inputs,
std::vector<string> outputs) {
CHECK(netdef_ptr);
auto& netdef = *netdef_ptr;
auto op_ptr = netdef.add_op();
auto& op = *op_ptr;
op.set_type(op_type);
for (const string& inp : inputs) {
op.add_input(inp);
}
for (const string& outp : outputs) {
op.add_output(outp);
}
return op_ptr;
}
bool MatchStrings(string p, string s) {
if (p == "*") { // star accepts anything
return true;
}
// TODO(benz): memoize this. (high constant factor boost in performance)
vector<string> choices = split('|', p);
for (const string& candidate : choices) {
if (candidate == s) {
return true;
}
}
return false;
}
bool MatchArguments(const OperatorDef& p_op, const OperatorDef& g_op) {
for (const auto& p_arg : p_op.arg()) {
if (!p_arg.has_name()) {
continue;
}
bool found = false;
for (const auto& g_arg : g_op.arg()) {
if (p_arg.name() == g_arg.name()) {
found = true;
if (p_arg.has_f()) {
if (!g_arg.has_f() || p_arg.f() != g_arg.f()) {
return false;
}
}
if (p_arg.has_i()) {
if (!g_arg.has_i() || p_arg.i() != g_arg.i()) {
return false;
}
}
if (p_arg.has_s()) {
if (!g_arg.has_s() || !MatchStrings(p_arg.s(), g_arg.s())) {
return false;
}
}
if (p_arg.floats_size() != g_arg.floats_size()) {
return false;
}
for (int i = 0; i < p_arg.floats_size(); i++) {
if (p_arg.floats(i) != g_arg.floats(i)) {
return false;
}
}
if (p_arg.ints_size() != g_arg.ints_size()) {
return false;
}
for (int i = 0; i < p_arg.ints_size(); i++) {
if (p_arg.ints(i) != g_arg.ints(i)) {
return false;
}
}
if (p_arg.strings_size() != g_arg.strings_size()) {
return false;
}
for (int i = 0; i < p_arg.strings_size(); i++) {
if (!MatchStrings(p_arg.strings(i), g_arg.strings(i))) {
return false;
}
}
}
}
if (!found) {
return false;
}
}
return true;
}
} // namespace caffe2