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