1/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. 2 3Licensed under the Apache License, Version 2.0 (the "License"); 4you may not use this file except in compliance with the License. 5You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9Unless required by applicable law or agreed to in writing, software 10distributed under the License is distributed on an "AS IS" BASIS, 11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12See the License for the specific language governing permissions and 13limitations under the License. 14==============================================================================*/ 15#ifndef TENSORFLOW_CONTRIB_LITE_TOCO_TENSORFLOW_GRAPH_MATCHING_RESOLVE_CLUSTER_H 16#define TENSORFLOW_CONTRIB_LITE_TOCO_TENSORFLOW_GRAPH_MATCHING_RESOLVE_CLUSTER_H 17 18#include <string> 19#include <unordered_map> 20#include <vector> 21 22#include "tensorflow/contrib/lite/toco/tensorflow_graph_matching/cluster.h" 23#include "tensorflow/contrib/lite/toco/tensorflow_graph_matching/resolve_svdf.h" 24#include "tensorflow/core/framework/graph.pb.h" 25#include "tensorflow/core/framework/node_def.pb.h" 26 27namespace toco { 28 29// Given a graph info and a list of cluster classes (cluster_factories), it 30// partitions the graph to clusters, and then collapses each cluster into their 31// corresponding composite ops. It generates a new graph using the newly 32// generated composite ops. Each cluster factory is responsible to recognize a 33// cluster of nodes into a cluster using a name-based pattern matching approach. 34std::unique_ptr<tensorflow::GraphDef> MaybeResolveClusters( 35 const tensorflow::GraphDef& graph_def, 36 const std::vector<ClusterFactoryInterface*>& cluster_factories); 37 38// Adds a node to a given graph. The added node will be a copy of a given source 39// node, except for the inputs. If the inputs are coming from a node which 40// belongs to another cluster, then those inputs are renamed to the source 41// cluster name. 42void AddNodeToGraph(const tensorflow::NodeDef& node, 43 const std::vector<string>& cluster_names, 44 tensorflow::GraphDef* graph); 45 46// Given a graph and a cluster class, it finds all the nodes which belong to a 47// given class factory, encapsulate them inside a cluster of the given type and 48// returns a vector of those clusters. It also labels the nodes in that graph if 49// they belong to the generated clusters. 50bool FindCluster(const ClusterFactoryInterface& cluster_factory, 51 const tensorflow::GraphDef& graph_def, 52 std::unordered_map<string, bool>* is_node_in_cluster, 53 std::vector<std::unique_ptr<Cluster>>* clusters); 54 55// Receives a graph and generates another graph by replacing the cluster of 56// nodes which matches a given composite op. Each composite op is represented 57// using a class factory. 58std::unique_ptr<tensorflow::GraphDef> MaybeReplaceCompositeSubgraph( 59 const tensorflow::GraphDef& tf_graph); 60 61} // end namespace toco 62 63#endif // CONTRIB_LITE_TOCO_TENSORFLOW_GRAPH_MATCHING_RESOLVE_CLUSTER_H 64