wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 6 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tensorflow/core/graph/BUILD
+302
View File
@@ -0,0 +1,302 @@
load(
"@xla//xla/tsl/mkl:build_defs.bzl",
"if_mkl",
)
load(
"//tensorflow:tensorflow.bzl",
"tf_cc_test",
"tf_cc_tests",
)
load("//tensorflow:tensorflow.default.bzl", "filegroup")
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
"//tensorflow/core:__subpackages__",
],
licenses = ["notice"],
)
# TODO(bmzhao): This target a holdover from tensorflow/core/BUILD. We
# will add proper dependencies once tf/core/graph/BUILD has granular
# targets added in a subsequent changes.
cc_library(
name = "mkl_graph_util",
hdrs = ["mkl_graph_util.h"],
deps = [
"@com_google_absl//absl/base",
"@com_google_absl//absl/container:flat_hash_map",
],
)
cc_library(
name = "zen_graph_util",
hdrs = ["zen_graph_util.h"],
deps = [
"@com_google_absl//absl/container:flat_hash_map",
],
)
tf_cc_test(
name = "collective_order_test",
size = "small",
srcs = [
"collective_order_test.cc",
],
deps = [
"//tensorflow/core",
"//tensorflow/core:core_cpu",
"//tensorflow/core:core_cpu_internal",
"//tensorflow/core:framework",
"//tensorflow/core:framework_internal",
"//tensorflow/core:lib",
"//tensorflow/core:lib_internal",
"//tensorflow/core:ops",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"@com_google_googletest//:gtest_main",
],
)
filegroup(
name = "core_cpu_headers",
srcs = [
"algorithm.h",
"default_device.h",
"graph.h",
"graph_debug_info_builder.h",
"graph_def_builder.h",
"graph_node_util.h",
"node_builder.h",
"validate.h",
"while_context.h",
],
)
filegroup(
name = "framework_internal_private_headers",
srcs = [
"edgeset.h",
"graph.h",
"graph_debug_info_builder.h",
"graph_def_builder.h",
"graph_node_util.h",
"node_builder.h",
"tensor_id.h",
"types.h",
],
)
filegroup(
name = "framework_internal_impl_srcs",
srcs = [
"edgeset.cc",
"graph.cc",
"graph_debug_info_builder.cc",
"graph_def_builder.cc",
"graph_node_util.cc",
"node_builder.cc",
"tensor_id.cc",
"while_context.cc",
"while_context.h",
],
)
# Note(bmzhao): This target is a holdover from the GRAPH_HDRS array
# in tensorflow/core/BUILD. This target contains all '.h' files under
# tensorflow/core/graph, except for the following:
# 'benchmark_testlib.h'.
filegroup(
name = "graph_headers",
srcs = [
"algorithm.h",
"collective_order.h",
"colors.h",
"control_flow.h",
"costmodel.h",
"default_device.h",
"edgeset.h",
"graph.h",
"graph_debug_info_builder.h",
"graph_def_builder.h",
"graph_node_util.h",
"graph_partition.h",
"node_builder.h",
"optimizer_cse.h",
"subgraph.h",
"tensor_id.h",
"testlib.h",
"types.h",
"validate.h",
"while_context.h",
],
)
filegroup(
name = "graph_srcs",
srcs = [
"algorithm.cc",
"collective_order.cc",
"colors.cc",
"control_flow.cc",
"costmodel.cc",
"graph_partition.cc",
"optimizer_cse.cc",
"subgraph.cc",
"validate.cc",
],
)
filegroup(
name = "testlib_headers",
srcs = [
"benchmark_testlib.h",
"testlib.h",
] + if_mkl(["mkl_testlib.h"]),
)
filegroup(
name = "testlib_srcs",
srcs = [
"testlib.cc",
] + if_mkl(["mkl_testlib.cc"]),
)
filegroup(
name = "mkl_graph_util_header",
srcs = [
"mkl_graph_util.h",
],
)
filegroup(
name = "mobile_srcs_only_runtime",
srcs = [
"algorithm.cc",
"collective_order.cc",
"colors.cc",
"control_flow.cc",
"costmodel.cc",
"edgeset.cc",
"graph.cc",
"graph_debug_info_builder.cc",
"graph_def_builder.cc",
"graph_node_util.cc",
"graph_partition.cc",
"node_builder.cc",
"optimizer_cse.cc",
"subgraph.cc",
"tensor_id.cc",
"validate.cc",
"while_context.cc",
],
)
filegroup(
name = "mobile_hdrs_only_runtime",
srcs = [
"algorithm.h",
"benchmark_testlib.h",
"collective_order.h",
"colors.h",
"control_flow.h",
"costmodel.h",
"default_device.h",
"edgeset.h",
"graph.h",
"graph_debug_info_builder.h",
"graph_def_builder.h",
"graph_node_util.h",
"graph_partition.h",
"mkl_graph_util.h",
"node_builder.h",
"optimizer_cse.h",
"subgraph.h",
"tensor_id.h",
"testlib.h",
"types.h",
"validate.h",
"while_context.h",
],
)
# Note(bmzhao): Ideally we would use a filegroup to represent these tests instead.
# However, that causes tf_cc_tests to link all of these tests into a single object
# file. This breaks tensorflow/core:core_higher_level_tests, because some of these
# tests redefine the same symbol. This will be fixed by having granular tests
# instead, after phase 4 of the tensorflow's build refactoring:
# https://github.com/tensorflow/community/pull/179
exports_files(
srcs = [
"algorithm_test.cc",
"control_flow_test.cc",
"costmodel_test.cc",
"edgeset_test.cc",
"graph_debug_info_builder_test.cc",
"graph_def_builder_test.cc",
"graph_partition_test.cc",
"graph_test.cc",
"node_builder_test.cc",
"optimizer_cse_test.cc",
"subgraph_test.cc",
"tensor_id_test.cc",
"validate_test.cc",
],
visibility = ["//tensorflow/core:__pkg__"],
)
tf_cc_tests(
name = "higher_level_tests",
size = "small",
srcs = [
"algorithm_test.cc",
"control_flow_test.cc",
"costmodel_test.cc",
"edgeset_test.cc",
"graph_debug_info_builder_test.cc",
"graph_def_builder_test.cc",
"graph_partition_test.cc",
"graph_test.cc",
"node_builder_test.cc",
"optimizer_cse_test.cc",
"subgraph_test.cc",
"tensor_id_test.cc",
"validate_test.cc",
],
linkopts = select({
"//tensorflow:macos": ["-headerpad_max_install_names"],
"//conditions:default": [],
}),
deps = [
"//tensorflow/cc:cc_ops",
"//tensorflow/cc:cc_ops_internal",
"//tensorflow/cc:function_ops",
"//tensorflow/cc:ops",
"//tensorflow/cc:scope",
"//tensorflow/cc:sendrecv_ops",
"//tensorflow/cc:while_loop",
"//tensorflow/core",
"//tensorflow/core:core_cpu",
"//tensorflow/core:core_cpu_internal",
"//tensorflow/core:direct_session_internal",
"//tensorflow/core:framework",
"//tensorflow/core:framework_internal",
"//tensorflow/core:lib",
"//tensorflow/core:lib_internal",
"//tensorflow/core:ops",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core:testlib",
"//tensorflow/core/kernels:ops_util",
"//tensorflow/core/platform:regexp",
"//tensorflow/core/platform:status_matchers",
"//tensorflow/core/util:protos_test_cc",
"@com_google_absl//absl/base",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/strings",
"@com_google_googletest//:gtest",
"@eigen_archive//:eigen3",
],
)
tensorflow/core/graph/algorithm.cc
+323
View File
@@ -0,0 +1,323 @@
/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#include "tensorflow/core/graph/algorithm.h"
#include <algorithm>
#include <deque>
#include <vector>
#include "tensorflow/core/platform/logging.h"
namespace tensorflow {
namespace {
template <typename T>
void DFSFromHelper(const Graph& g, gtl::ArraySlice<T> start,
const std::function<void(T)>& enter,
const std::function<void(T)>& leave,
const NodeComparator& stable_comparator,
const EdgeFilter& edge_filter) {
// Stack of work to do.
struct Work {
T node;
bool leave; // Are we entering or leaving n?
};
std::vector<Work> stack(start.size());
for (int i = 0; i < start.size(); ++i) {
stack[i] = Work{start[i], false};
}
std::vector<bool> visited(g.num_node_ids(), false);
while (!stack.empty()) {
Work w = stack.back();
stack.pop_back();
T n = w.node;
if (w.leave) {
leave(n);
continue;
}
if (visited[n->id()]) continue;
visited[n->id()] = true;
if (enter) enter(n);
// Arrange to call leave(n) when all done with descendants.
if (leave) stack.push_back(Work{n, true});
auto add_work = [&visited, &stack](Node* out) {
if (!visited[out->id()]) {
// Note; we must not mark as visited until we actually process it.
stack.push_back(Work{out, false});
}
};
if (stable_comparator) {
std::vector<Node*> nodes_sorted;
for (const Edge* out_edge : n->out_edges()) {
if (!edge_filter || edge_filter(*out_edge)) {
nodes_sorted.emplace_back(out_edge->dst());
}
}
std::sort(nodes_sorted.begin(), nodes_sorted.end(), stable_comparator);
for (Node* out : nodes_sorted) {
add_work(out);
}
} else {
for (const Edge* out_edge : n->out_edges()) {
if (!edge_filter || edge_filter(*out_edge)) {
add_work(out_edge->dst());
}
}
}
}
}
} // namespace
void DFS(const Graph& g, const std::function<void(Node*)>& enter,
const std::function<void(Node*)>& leave,
const NodeComparator& stable_comparator,
const EdgeFilter& edge_filter) {
DFSFromHelper(g, {g.source_node()}, enter, leave, stable_comparator,
edge_filter);
}
void DFSFrom(const Graph& g, absl::Span<Node* const> start,
const std::function<void(Node*)>& enter,
const std::function<void(Node*)>& leave,
const NodeComparator& stable_comparator,
const EdgeFilter& edge_filter) {
DFSFromHelper(g, start, enter, leave, stable_comparator, edge_filter);
}
void DFSFrom(const Graph& g, absl::Span<const Node* const> start,
const std::function<void(const Node*)>& enter,
const std::function<void(const Node*)>& leave,
const NodeComparator& stable_comparator,
const EdgeFilter& edge_filter) {
DFSFromHelper(g, start, enter, leave, stable_comparator, edge_filter);
}
void ReverseDFS(const Graph& g, const std::function<void(Node*)>& enter,
const std::function<void(Node*)>& leave,
const NodeComparator& stable_comparator,
const EdgeFilter& edge_filter) {
ReverseDFSFrom(g, {g.sink_node()}, enter, leave, stable_comparator,
edge_filter);
}
namespace {
template <typename T>
void ReverseDFSFromHelper(const Graph& g, gtl::ArraySlice<T> start,
const std::function<void(T)>& enter,
const std::function<void(T)>& leave,
const NodeComparator& stable_comparator,
const EdgeFilter& edge_filter) {
// Stack of work to do.
struct Work {
T node;
bool leave; // Are we entering or leaving n?
};
std::vector<Work> stack(start.size());
for (int i = 0; i < start.size(); ++i) {
stack[i] = Work{start[i], false};
}
std::vector<bool> visited(g.num_node_ids(), false);
while (!stack.empty()) {
Work w = stack.back();
stack.pop_back();
T n = w.node;
if (w.leave) {
leave(n);
continue;
}
if (visited[n->id()]) continue;
visited[n->id()] = true;
if (enter) enter(n);
// Arrange to call leave(n) when all done with descendants.
if (leave) stack.push_back(Work{n, true});
auto add_work = [&visited, &stack](T out) {
if (!visited[out->id()]) {
// Note; we must not mark as visited until we actually process it.
stack.push_back(Work{out, false});
}
};
if (stable_comparator) {
std::vector<T> nodes_sorted;
for (const Edge* in_edge : n->in_edges()) {
if (!edge_filter || edge_filter(*in_edge)) {
nodes_sorted.emplace_back(in_edge->src());
}
}
std::sort(nodes_sorted.begin(), nodes_sorted.end(), stable_comparator);
for (T in : nodes_sorted) {
add_work(in);
}
} else {
for (const Edge* in_edge : n->in_edges()) {
if (!edge_filter || edge_filter(*in_edge)) {
add_work(in_edge->src());
}
}
}
}
}
} // namespace
void ReverseDFSFrom(const Graph& g, absl::Span<const Node* const> start,
const std::function<void(const Node*)>& enter,
const std::function<void(const Node*)>& leave,
const NodeComparator& stable_comparator,
const EdgeFilter& edge_filter) {
ReverseDFSFromHelper(g, start, enter, leave, stable_comparator, edge_filter);
}
void ReverseDFSFrom(const Graph& g, absl::Span<Node* const> start,
const std::function<void(Node*)>& enter,
const std::function<void(Node*)>& leave,
const NodeComparator& stable_comparator,
const EdgeFilter& edge_filter) {
ReverseDFSFromHelper(g, start, enter, leave, stable_comparator, edge_filter);
}
void GetPostOrder(const Graph& g, std::vector<Node*>* order,
const NodeComparator& stable_comparator,
const EdgeFilter& edge_filter) {
order->clear();
DFS(
g, nullptr, [order](Node* n) { order->push_back(n); }, stable_comparator,
edge_filter);
}
void GetReversePostOrder(const Graph& g, std::vector<Node*>* order,
const NodeComparator& stable_comparator,
const EdgeFilter& edge_filter) {
GetPostOrder(g, order, stable_comparator, edge_filter);
std::reverse(order->begin(), order->end());
}
bool PruneForReverseReachability(Graph* g,
std::unordered_set<const Node*> start) {
// Compute set of nodes that we need to traverse in order to reach
// the nodes in "start" by performing a breadth-first search from those
// nodes, and accumulating the visited nodes.
std::vector<bool> visited(g->num_node_ids());
for (auto node : start) {
visited[node->id()] = true;
}
std::deque<const Node*> queue(start.begin(), start.end());
while (!queue.empty()) {
const Node* n = queue.front();
queue.pop_front();
for (const Node* in : n->in_nodes()) {
if (!visited[in->id()]) {
visited[in->id()] = true;
queue.push_back(in);
VLOG(2) << "Reverse reach : " << n->name() << " from " << in->name();
}
}
}
// Make a pass over the graph to remove nodes not in "visited".
bool any_removed = false;
for (int i = 0; i < visited.size(); ++i) {
if (!visited[i]) {
Node* n = g->FindNodeId(i);
if (n != nullptr && !n->IsSource() && !n->IsSink()) {
g->RemoveNode(n);
any_removed = true;
}
}
}
return any_removed;
}
bool FixupSourceAndSinkEdges(Graph* g) {
// Connect all nodes with no incoming edges to source.
// Connect all nodes with no outgoing edges to sink.
bool changed = false;
for (Node* n : g->nodes()) {
if (!n->IsSource() && n->in_edges().empty()) {
g->AddControlEdge(g->source_node(), n,
true /* skip test for duplicates */);
changed = true;
}
if (!n->IsSink() && n->out_edges().empty()) {
g->AddControlEdge(n, g->sink_node(), true /* skip test for duplicates */);
changed = true;
}
}
return changed;
}
namespace {
template <class T>
void BreadthFirstTraversalHelper(const Graph& g, gtl::ArraySlice<T> start,
const std::function<void(T)>& visit,
NodeComparator stable_comparator) {
std::deque<T> stack;
if (start.empty()) {
for (T n : g.nodes()) {
if (n->in_edges().empty()) {
stack.push_back(n);
}
}
}
std::vector<bool> seen(g.num_node_ids(), false);
while (!stack.empty()) {
T n = stack.front();
stack.pop_front();
seen[n->id()] = true;
visit(n);
std::vector<T> nodes_sorted;
for (const Edge* out_edge : n->out_edges()) {
if (!seen[out_edge->dst()->id()]) {
seen[out_edge->dst()->id()] = true;
nodes_sorted.emplace_back(out_edge->dst());
}
}
std::sort(nodes_sorted.begin(), nodes_sorted.end(), stable_comparator);
for (T out : nodes_sorted) {
stack.push_back(out);
}
}
}
} // namespace
void BreadthFirstTraversal(const Graph& g, absl::Span<const Node* const> start,
const std::function<void(const Node*)>& visit,
NodeComparator stable_comparator) {
return BreadthFirstTraversalHelper<const Node*>(g, start, visit,
stable_comparator);
}
void BreadthFirstTraversal(Graph& g, absl::Span<Node* const> start,
const std::function<void(Node*)>& visit,
NodeComparator stable_comparator) {
return BreadthFirstTraversalHelper<Node*>(g, start, visit, stable_comparator);
}
} // namespace tensorflow
tensorflow/core/graph/algorithm.h
+154
View File
@@ -0,0 +1,154 @@
/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
Licensed 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 TENSORFLOW_CORE_GRAPH_ALGORITHM_H_
#define TENSORFLOW_CORE_GRAPH_ALGORITHM_H_
#include <functional>
#include <unordered_set>
#include <vector>
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/lib/gtl/array_slice.h"
namespace tensorflow {
// Comparator for two nodes. This is used in order to get a stable ording.
using NodeComparator = std::function<bool(const Node*, const Node*)>;
using EdgeFilter = std::function<bool(const Edge&)>;
// Compares two node based on their ids.
struct NodeComparatorID {
bool operator()(const Node* n1, const Node* n2) const {
return n1->id() < n2->id();
}
};
// Compare two nodes based on their names.
struct NodeComparatorName {
bool operator()(const Node* n1, const Node* n2) const {
return n1->name() < n2->name();
}
};
// Perform a depth-first-search on g starting at the source node.
// If enter is not empty, calls enter(n) before visiting any children of n.
// If leave is not empty, calls leave(n) after visiting all children of n.
// If stable_comparator is set, a stable ordering of visit is achieved by
// sorting a node's neighbors first before visiting them.
// If edge_filter is set then ignores edges for which edge_filter returns false.
void DFS(const Graph& g, const std::function<void(Node*)>& enter,
const std::function<void(Node*)>& leave,
const NodeComparator& stable_comparator = {},
const EdgeFilter& edge_filter = {});
// Perform a depth-first-search on g starting at the 'start' nodes.
// If enter is not empty, calls enter(n) before visiting any children of n.
// If leave is not empty, calls leave(n) after visiting all children of n.
// If stable_comparator is set, a stable ordering of visit is achieved by
// sorting a node's neighbors first before visiting them.
// If edge_filter is set then ignores edges for which edge_filter returns false.
void DFSFrom(const Graph& g, absl::Span<Node* const> start,
const std::function<void(Node*)>& enter,
const std::function<void(Node*)>& leave,
const NodeComparator& stable_comparator = {},
const EdgeFilter& edge_filter = {});
void DFSFrom(const Graph& g, absl::Span<const Node* const> start,
const std::function<void(const Node*)>& enter,
const std::function<void(const Node*)>& leave,
const NodeComparator& stable_comparator = {},
const EdgeFilter& edge_filter = {});
// Perform a reverse depth-first-search on g starting at the sink node.
// If enter is not empty, calls enter(n) before visiting any parents of n.
// If leave is not empty, calls leave(n) after visiting all parents of n.
// If stable_comparator is set, a stable ordering of visit is achieved by
// sorting a node's neighbors first before visiting them.
// If edge_filter is set then ignores edges for which edge_filter returns false.
void ReverseDFS(const Graph& g, const std::function<void(Node*)>& enter,
const std::function<void(Node*)>& leave,
const NodeComparator& stable_comparator = {},
const EdgeFilter& edge_filter = {});
// Perform a reverse depth-first-search on g starting at the 'start' nodes.
// If enter is not empty, calls enter(n) before visiting any parents of n.
// If leave is not empty, calls leave(n) after visiting all parents of n.
// If stable_comparator is set, a stable ordering of visit is achieved by
// sorting a node's neighbors first before visiting them.
// If edge_filter is set then ignores edges for which edge_filter returns false.
void ReverseDFSFrom(const Graph& g, absl::Span<Node* const> start,
const std::function<void(Node*)>& enter,
const std::function<void(Node*)>& leave,
const NodeComparator& stable_comparator = {},
const EdgeFilter& edge_filter = {});
void ReverseDFSFrom(const Graph& g, absl::Span<const Node* const> start,
const std::function<void(const Node*)>& enter,
const std::function<void(const Node*)>& leave,
const NodeComparator& stable_comparator = {},
const EdgeFilter& edge_filter = {});
void BreadthFirstTraversal(
const Graph& g, absl::Span<const Node* const> start,
const std::function<void(const Node*)>& visit,
NodeComparator stable_comparator = NodeComparatorID());
void BreadthFirstTraversal(
Graph& g, absl::Span<Node* const> start,
const std::function<void(Node*)>& visit,
NodeComparator stable_comparator = NodeComparatorID());
// Stores in *order the post-order numbering of all nodes
// in graph found via a depth first search starting at the source node.
//
// Note that this is equivalent to reverse topological sorting when the
// graph does not have cycles.
//
// If stable_comparator is set, a stable ordering of visit is achieved by
// sorting a node's neighbors first before visiting them.
//
// If edge_filter is set then ignores edges for which edge_filter returns
// false.
//
// REQUIRES: order is not NULL.
void GetPostOrder(const Graph& g, std::vector<Node*>* order,
const NodeComparator& stable_comparator = {},
const EdgeFilter& edge_filter = {});
// Stores in *order the reverse post-order numbering of all nodes
// If stable_comparator is set, a stable ordering of visit is achieved by
// sorting a node's neighbors first before visiting them.
//
// If edge_filter is set then ignores edges for which edge_filter returns
// false.
void GetReversePostOrder(const Graph& g, std::vector<Node*>* order,
const NodeComparator& stable_comparator = {},
const EdgeFilter& edge_filter = {});
// Prune nodes in "g" that are not in some path from the source node
// to any node in 'nodes'. Returns true if changes were made to the graph.
// Does not fix up source and sink edges.
bool PruneForReverseReachability(Graph* g,
std::unordered_set<const Node*> nodes);
// Connect all nodes with no incoming edges to source.
// Connect all nodes with no outgoing edges to sink.
//
// Returns true if and only if 'g' is mutated.
bool FixupSourceAndSinkEdges(Graph* g);
} // namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPH_ALGORITHM_H_
tensorflow/core/graph/algorithm_test.cc
+246
View File
@@ -0,0 +1,246 @@
/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#include "tensorflow/core/graph/algorithm.h"
#include <string>
#include <vector>
#include "tensorflow/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/common_runtime/graph_def_builder_util.h"
#include "tensorflow/core/graph/benchmark_testlib.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/graph/graph_def_builder.h"
#include "tensorflow/core/graph/subgraph.h"
#include "tensorflow/core/kernels/ops_util.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/test_benchmark.h"
// TODO(josh11b): Test setting the "device" field of a NodeDef.
// TODO(josh11b): Test that feeding won't prune targets.
namespace tensorflow {
namespace {
REGISTER_OP("TestParams").Output("o: float");
REGISTER_OP("TestInput").Output("a: float").Output("b: float");
REGISTER_OP("TestMul").Input("a: float").Input("b: float").Output("o: float");
REGISTER_OP("TestUnary").Input("a: float").Output("o: float");
REGISTER_OP("TestBinary")
.Input("a: float")
.Input("b: float")
.Output("o: float");
// Compares that the order of nodes in 'inputs' respects the
// pair orders described in 'ordered_pairs'.
bool ExpectBefore(
const std::vector<std::pair<std::string, std::string>>& ordered_pairs,
const std::vector<Node*>& inputs, std::string* error) {
for (const std::pair<std::string, std::string>& pair : ordered_pairs) {
const std::string& before_node = pair.first;
const std::string& after_node = pair.second;
bool seen_before = false;
bool seen_both = false;
for (const Node* node : inputs) {
if (!seen_before && after_node == node->name()) {
*error = absl::StrCat("Saw ", after_node, " before ", before_node);
return false;
}
if (before_node == node->name()) {
seen_before = true;
} else if (after_node == node->name()) {
seen_both = seen_before;
break;
}
}
if (!seen_both) {
*error =
absl::StrCat("didn't see either ", before_node, " or ", after_node);
return false;
}
}
return true;
}
TEST(AlgorithmTest, ReversePostOrder) {
GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
using namespace ::tensorflow::ops; // NOLINT(build/namespaces)
Node* w1 = SourceOp("TestParams", b.opts().WithName("W1"));
Node* w2 = SourceOp("TestParams", b.opts().WithName("W2"));
Node* input =
SourceOp("TestInput", b.opts().WithName("input").WithControlInput(w1));
Node* t1 = BinaryOp("TestMul", w1, {input, 1}, b.opts().WithName("t1"));
BinaryOp("TestMul", w1, {input, 1},
b.opts().WithName("t2").WithControlInput(t1));
BinaryOp("TestMul", w2, {input, 1}, b.opts().WithName("t3"));
Graph g(OpRegistry::Global());
TF_ASSERT_OK(GraphDefBuilderToGraph(b, &g));
std::vector<Node*> order;
// Test reverse post order:
GetReversePostOrder(g, &order);
// Check that the order respects the dependencies correctly.
std::vector<std::pair<std::string, std::string>> reverse_orders = {
{"W1", "input"}, {"W1", "t1"}, {"W1", "t2"}, {"W1", "t3"},
{"input", "t1"}, {"input", "t3"}, {"t1", "t2"}, {"W2", "t3"}};
std::string error;
EXPECT_TRUE(ExpectBefore(reverse_orders, order, &error)) << error;
// A false ordering should fail the check.
reverse_orders = {{"input", "W1"}};
EXPECT_FALSE(ExpectBefore(reverse_orders, order, &error));
// Test post order:
GetPostOrder(g, &order);
// Check that the order respects the dependencies correctly.
std::vector<std::pair<std::string, std::string>> orders = {
{"input", "W1"}, {"t1", "W1"}, {"t2", "W1"}, {"t3", "W1"},
{"t1", "input"}, {"t3", "input"}, {"t2", "t1"}, {"t3", "W2"}};
EXPECT_TRUE(ExpectBefore(orders, order, &error)) << error;
// A false ordering should fail the check.
orders = {{"W1", "t3"}};
EXPECT_FALSE(ExpectBefore(orders, order, &error));
}
TEST(AlgorithmTest, ReversePostOrderStable) {
int64_t run_count = 100;
using namespace ::tensorflow::ops; // NOLINT(build/namespaces)
for (int64_t i = 0; i < run_count; ++i) {
// One source of nondeterminism comes from unordered set with key of a
// pointer type, for example the order of FlatSet<Node*> depends on the
// raw pointer value of Node. Stable post order suppose to remove this
// nondeterminism by enforcing an ordering based on node ids.
GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
std::string error;
Node* w1 = SourceOp("TestParams", b.opts().WithName("W1"));
Node* input =
SourceOp("TestInput", b.opts().WithName("input").WithControlInput(w1));
BinaryOp("TestMul", w1, {input, 1}, b.opts().WithName("t2"));
// Insert different number of nodes between the allocation of t2 and t3,
// this creates enough entropy in the memory distance between t2 and t3 thus
// forces them to have randomized ordering had stable DFS was not
// implemented correctly.
for (int64_t j = 0; j < i; ++j) {
BinaryOp("TestMul", w1, {input, 1},
b.opts().WithName(absl::StrCat("internal", j)));
}
BinaryOp("TestMul", w1, {input, 1}, b.opts().WithName("t3"));
Graph g(OpRegistry::Global());
TF_ASSERT_OK(GraphDefBuilderToGraph(b, &g));
std::vector<Node*> order;
// Test reverse post order generates expected ordering.
GetReversePostOrder(g, &order, /*stable_comparator=*/NodeComparatorName());
EXPECT_TRUE(ExpectBefore({{"t2", "t3"}}, order, &error));
}
}
TEST(AlgorithmTest, PostOrderWithEdgeFilter) {
GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
Node* n0 = ops::SourceOp("TestParams", b.opts().WithName("n0"));
Node* n1 = ops::UnaryOp("TestUnary", n0, b.opts().WithName("n1"));
Node* n2 = ops::UnaryOp("TestUnary", n1, b.opts().WithName("n2"));
Node* n3 = ops::BinaryOp("TestBinary", n2, n0, b.opts().WithName("n3"));
Graph g(OpRegistry::Global());
TF_ASSERT_OK(GraphDefBuilderToGraph(b, &g));
g.AddEdge(g.FindNodeId(n3->id()), 0, g.FindNodeId(n1->id()), 1);
std::vector<Node*> post_order;
auto edge_filter = [&](const Edge& e) {
return !(e.src()->id() == n3->id() && e.dst()->id() == n1->id());
};
std::vector<Node*> expected_post_order = {
g.sink_node(), g.FindNodeId(n3->id()), g.FindNodeId(n2->id()),
g.FindNodeId(n1->id()), g.FindNodeId(n0->id()), g.source_node()};
std::vector<Node*> expected_reverse_post_order = expected_post_order;
std::reverse(expected_reverse_post_order.begin(),
expected_reverse_post_order.end());
GetPostOrder(g, &post_order, /*stable_comparator=*/{},
/*edge_filter=*/edge_filter);
ASSERT_EQ(expected_post_order.size(), post_order.size());
for (int i = 0; i < post_order.size(); i++) {
CHECK_EQ(post_order[i], expected_post_order[i])
<< post_order[i]->name() << " vs. " << expected_post_order[i]->name();
}
std::vector<Node*> reverse_post_order;
GetReversePostOrder(g, &reverse_post_order, /*stable_comparator=*/{},
/*edge_filter=*/edge_filter);
ASSERT_EQ(expected_reverse_post_order.size(), reverse_post_order.size());
for (int i = 0; i < reverse_post_order.size(); i++) {
CHECK_EQ(reverse_post_order[i], expected_reverse_post_order[i])
<< reverse_post_order[i]->name() << " vs. "
<< expected_reverse_post_order[i]->name();
}
}
void BM_PruneForReverseReachability(::testing::benchmark::State& state) {
const int num_nodes = state.range(0);
const int num_edges_per_node = state.range(1);
const GraphDef graph_def =
test::CreateGraphDef(num_nodes, num_edges_per_node);
const auto registry = OpRegistry::Global();
GraphConstructorOptions opts;
for (auto s : state) {
state.PauseTiming();
Graph graph(registry);
TF_CHECK_OK(ConvertGraphDefToGraph(opts, graph_def, &graph));
std::unordered_set<const Node*> visited;
visited.insert(graph.FindNodeId(graph.num_nodes() - 1));
state.ResumeTiming();
PruneForReverseReachability(&graph, std::move(visited));
}
}
BENCHMARK(BM_PruneForReverseReachability)->ArgPair(10, 2);
BENCHMARK(BM_PruneForReverseReachability)->ArgPair(1 << 6, 2);
BENCHMARK(BM_PruneForReverseReachability)->ArgPair(1 << 9, 2);
BENCHMARK(BM_PruneForReverseReachability)->ArgPair(1 << 12, 2);
BENCHMARK(BM_PruneForReverseReachability)->ArgPair(1 << 15, 2);
BENCHMARK(BM_PruneForReverseReachability)->ArgPair(10, 4);
BENCHMARK(BM_PruneForReverseReachability)->ArgPair(1 << 6, 4);
BENCHMARK(BM_PruneForReverseReachability)->ArgPair(1 << 9, 4);
BENCHMARK(BM_PruneForReverseReachability)->ArgPair(1 << 12, 4);
BENCHMARK(BM_PruneForReverseReachability)->ArgPair(1 << 15, 4);
BENCHMARK(BM_PruneForReverseReachability)->ArgPair(10, 8);
BENCHMARK(BM_PruneForReverseReachability)->ArgPair(1 << 6, 8);
BENCHMARK(BM_PruneForReverseReachability)->ArgPair(1 << 9, 8);
BENCHMARK(BM_PruneForReverseReachability)->ArgPair(1 << 12, 8);
BENCHMARK(BM_PruneForReverseReachability)->ArgPair(1 << 15, 8);
BENCHMARK(BM_PruneForReverseReachability)->ArgPair(10, 16);
BENCHMARK(BM_PruneForReverseReachability)->ArgPair(1 << 6, 16);
BENCHMARK(BM_PruneForReverseReachability)->ArgPair(1 << 9, 16);
BENCHMARK(BM_PruneForReverseReachability)->ArgPair(1 << 12, 16);
BENCHMARK(BM_PruneForReverseReachability)->ArgPair(1 << 15, 16);
} // namespace
} // namespace tensorflow
tensorflow/core/graph/benchmark_testlib.h
+191
View File
@@ -0,0 +1,191 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed 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 TENSORFLOW_CORE_GRAPH_BENCHMARK_TESTLIB_H_
#define TENSORFLOW_CORE_GRAPH_BENCHMARK_TESTLIB_H_
#include "absl/strings/str_cat.h"
#include "absl/strings/str_format.h"
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/lib/random/philox_random.h"
#include "tensorflow/core/lib/random/simple_philox.h"
namespace tensorflow {
namespace test {
REGISTER_OP("Input").Output("y: float");
REGISTER_OP("Output")
.Input("x: N * float")
.Attr("N: int >= 1")
.Output("y: float");
REGISTER_OP("In2Out1").Input("a: float").Input("b: float").Output("y: float");
REGISTER_OP("In4Out1")
.Input("a: float")
.Input("b: float")
.Input("c: float")
.Input("d: float")
.Output("y: float");
REGISTER_OP("In8Out1")
.Input("a: float")
.Input("b: float")
.Input("c: float")
.Input("d: float")
.Input("e: float")
.Input("f: float")
.Input("g: float")
.Input("h: float")
.Output("y: float");
REGISTER_OP("In16Out1")
.Input("a: float")
.Input("b: float")
.Input("c: float")
.Input("d: float")
.Input("e: float")
.Input("f: float")
.Input("g: float")
.Input("h: float")
.Input("i: float")
.Input("j: float")
.Input("k: float")
.Input("l: float")
.Input("m: float")
.Input("n: float")
.Input("o: float")
.Input("p: float")
.Output("y: float");
inline GraphDef CreateGraphDef(int num_nodes, int num_edges_per_node) {
const int kNumInNodes = 10 * num_edges_per_node;
GraphDef graph_def;
auto create_node = [](const std::string& name, const std::string& op) {
NodeDef node;
node.set_name(name);
node.set_op(op);
return node;
};
NodeDef node;
for (int in = 0; in < kNumInNodes; ++in) {
node = create_node(/*name=*/absl::StrFormat("in%04d", in), /*op=*/"Input");
*graph_def.add_node() = std::move(node);
}
random::PhiloxRandom philox(301, 17);
random::SimplePhilox rnd(&philox);
for (int op = 0; op < num_nodes; ++op) {
node = create_node(/*name=*/absl::StrFormat("op%05d", op),
/*op=*/absl::StrFormat("In%dOut1", num_edges_per_node));
for (int edge = 0; edge < num_edges_per_node; ++edge) {
node.add_input(absl::StrFormat("in%04d", rnd.Uniform(kNumInNodes)));
}
*graph_def.add_node() = std::move(node);
}
// Add a single sink node. Otherwise a lot of time is spent in
// FixupSourceAndSinkEdges().
node = create_node(/*name=*/"out", /*op=*/"Output");
for (int op = 0; op < num_nodes; ++op) {
node.add_input(absl::StrFormat("op%05d", op));
}
AttrValue attr;
attr.set_i(num_nodes);
node.mutable_attr()->insert({"N", std::move(attr)});
*graph_def.add_node() = std::move(node);
return graph_def;
}
inline GraphDef CreateRandomGraph(int size) {
random::PhiloxRandom philox(0x12345);
random::SimplePhilox rnd(&philox);
std::string prefix = "long_node_name_prefix_to_measure_string_copy_overhead";
GraphDef graph;
for (int i = 0; i < size; ++i) {
const std::string name = absl::StrCat(prefix, i);
const uint32_t num_inputs = rnd.Uniform(std::min(i, 5));
NodeDef node;
node.set_name(name);
for (int n = 0; n < num_inputs; ++n) {
const uint32_t input_node = rnd.Uniform(i);
node.add_input(absl::StrCat(prefix, input_node));
}
*graph.add_node() = std::move(node);
}
return graph;
}
inline GraphDef CreateFaninFanoutNodeGraph(int num_regular_fanins,
int num_regular_fanouts,
int num_controlling_fanins,
int num_controlled_fanouts,
bool fanout_unique_index) {
GraphDef graph;
auto create_node = [](const std::string& name) {
NodeDef node;
node.set_name(name);
return node;
};
NodeDef node = create_node(/*name=*/"node");
for (int i = 0; i < num_regular_fanins; ++i) {
const std::string input_node_name = absl::StrFormat("in%05d", i);
NodeDef input_node = create_node(/*name=*/input_node_name);
*graph.add_node() = std::move(input_node);
node.add_input(input_node_name);
}
for (int i = 0; i < num_controlling_fanins; ++i) {
const std::string input_node_name = absl::StrFormat("control_in%05d", i);
NodeDef input_node = create_node(/*name=*/input_node_name);
*graph.add_node() = std::move(input_node);
node.add_input(absl::StrCat("^", input_node_name));
}
for (int i = 0; i < num_regular_fanouts; ++i) {
NodeDef output_node = create_node(/*name=*/absl::StrFormat("out%05d", i));
const std::string input_node_index =
fanout_unique_index ? absl::StrCat(node.name(), ":", i) : node.name();
output_node.add_input(input_node_index);
*graph.add_node() = std::move(output_node);
}
const std::string controlled_fanout_input = absl::StrCat("^", node.name());
for (int i = 0; i < num_controlled_fanouts; ++i) {
NodeDef output_node =
create_node(/*name=*/absl::StrFormat("control_out%05d", i));
output_node.add_input(controlled_fanout_input);
*graph.add_node() = std::move(output_node);
}
*graph.add_node() = std::move(node);
return graph;
}
} // namespace test
} // namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPH_BENCHMARK_TESTLIB_H_
tensorflow/core/graph/collective_order.cc
+207
View File
@@ -0,0 +1,207 @@
/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#include "tensorflow/core/graph/collective_order.h"
#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "tensorflow/core/graph/algorithm.h"
namespace tensorflow {
namespace {
// Find all CollectiveReduce nodes and the existing data dependencies between
// them.
absl::Status DiscoverDataDependencies(
const Graph* graph, std::vector<Node*>* collective_nodes,
std::vector<int32_t>* instance_keys,
absl::flat_hash_map<Node*, absl::flat_hash_set<int32_t>>*
data_dependencies) {
absl::Status s;
// Algorithm: do Reverse DFS starting at sink. `node_leave` is called when
// all parents of `node` have been visited. At that point,
// `data_dependencies[node]` is a list containing `instance_key` of every
// `CollectiveReduce` on which `node` has a data dependency.
// For this node's children, add all these instance keys. Also, if this node
// is collective, add as a dependency for the children.
auto node_leave = [collective_nodes, instance_keys, data_dependencies,
&s](Node* node) {
int32_t instance_key;
bool enter_node =
node->IsCollective() && node->type_string() == "CollectiveReduce";
if (enter_node) {
absl::Status get_attr_status =
GetNodeAttr(node->attrs(), "instance_key", &instance_key);
s.Update(get_attr_status);
collective_nodes->push_back(node);
instance_keys->push_back(instance_key);
VLOG(2) << "collective node " << node->DebugString();
}
// Avoid reference invalidation of `node_deps`.
data_dependencies->reserve(data_dependencies->size() + 1 +
node->out_edges().size());
const auto& node_deps = (*data_dependencies)[node];
for (const Edge* out_edge : node->out_edges()) {
auto& child_deps = (*data_dependencies)[out_edge->dst()];
child_deps.insert(node_deps.begin(), node_deps.end());
if (enter_node && s.ok()) {
child_deps.insert(instance_key);
}
}
};
ReverseDFS(*graph, nullptr, node_leave);
return s;
}
// Given a list of `collective_nodes` and `data_dependencies` between the
// collective nodes, create control dependencies between concurrent collectives
// and store in `dependency_edges`.
// If there exists an edge a -> b then `dependency_edges[a]` contains `b`
absl::Status CreateControlDependencies(
const std::vector<Node*>& collective_nodes,
const std::vector<int32_t>& instance_keys,
absl::flat_hash_map<Node*, absl::flat_hash_set<int32_t>>* data_dependencies,
absl::flat_hash_map<Node*, absl::flat_hash_set<Node*>>* dependency_edges) {
// If there exists some path a -> ... -> b then `all_paths[a]` contains `b`
absl::flat_hash_map<Node*, absl::flat_hash_set<Node*>> all_paths;
for (int i = 0; i < collective_nodes.size() - 1; i++) {
if (!collective_nodes[i]->IsCollective() ||
collective_nodes[i]->type_string() != "CollectiveReduce") {
return absl::InternalError(
absl::StrCat("Unexpected node ", collective_nodes[i]->DebugString()));
}
const auto& deps_i = (*data_dependencies)[collective_nodes[i]];
for (int j = i + 1; j < collective_nodes.size(); j++) {
if (collective_nodes[i]->requested_device() !=
collective_nodes[j]->requested_device()) {
continue;
}
if (instance_keys[i] == instance_keys[j]) {
return absl::InternalError(
absl::StrCat("Unexpected same instance_key ", instance_keys[i],
" on 2 nodes with the same device ",
collective_nodes[i]->requested_device()));
}
const auto& deps_j = (*data_dependencies)[collective_nodes[j]];
if (deps_i.find(instance_keys[j]) == deps_i.end() &&
deps_j.find(instance_keys[i]) == deps_j.end()) {
int src_idx = instance_keys[i] > instance_keys[j] ? i : j;
int dst_idx = instance_keys[i] > instance_keys[j] ? j : i;
Node* src_node = collective_nodes[src_idx];
Node* dst_node = collective_nodes[dst_idx];
VLOG(1) << "Adding control dependency from node " << src_node->name()
<< " instance " << instance_keys[src_idx] << " to node "
<< dst_node->name() << " instance " << instance_keys[dst_idx];
(*dependency_edges)[src_node].insert(dst_node);
auto& src_paths = all_paths[src_node];
src_paths.insert(dst_node);
for (Node* downstream_node : all_paths[dst_node]) {
src_paths.insert(downstream_node);
}
}
}
}
// Prune dependency edges so that if there are edges a -> b, b -> c, and a ->
// c, then remove a -> c. This dependency would be handled naturally during
// op scheduling.
for (int i = 0; i < collective_nodes.size(); ++i) {
Node* node = collective_nodes[i];
auto& neighbor_set = (*dependency_edges)[node];
std::vector<Node*> neighbor_list(neighbor_set.begin(), neighbor_set.end());
// For all n1, n2 in `neighbor_list` if there is a path from n1 -> n2 then
// eliminate n2 from `neighbor_set` and `neighbor_list`. We remove from
// `neighbor_list` by replacing with a `nullptr`, hence the `nullptr` checks
// below.
for (int j = 0; j < neighbor_list.size(); ++j) {
Node* n1 = neighbor_list[j];
if (n1 == nullptr) continue;
auto& n1_paths = all_paths[n1];
for (int k = 0; k < neighbor_list.size(); ++k) {
Node* n2 = neighbor_list[k];
if (j == k || n2 == nullptr) continue;
if (n1_paths.find(n2) != n1_paths.end()) {
neighbor_set.erase(n2);
neighbor_list[k] = nullptr;
}
}
}
}
return absl::OkStatus();
}
// Insert control dependencies defined by `dependency_edges` in `graph`. If
// `order_type` is `kEdges`, insert explicit control edges, else if `order_type`
// is `kAttrs`, encode dependencies as an attribute on collective node.
absl::Status InsertControlDependencies(
Graph* graph, GraphCollectiveOrder order_type,
const absl::flat_hash_map<Node*, absl::flat_hash_set<Node*>>&
dependency_edges) {
if (order_type == GraphCollectiveOrder::kEdges) {
for (const auto& pair : dependency_edges) {
Node* src_node = pair.first;
for (Node* dst_node : pair.second) {
graph->AddControlEdge(src_node, dst_node);
}
}
} else if (order_type == GraphCollectiveOrder::kAttrs) {
// `wait_for` is the inverse of `dependency_edges`, i.e. `wait_for[node]`
// contains the list of instance keys for which `node` must wait.
absl::flat_hash_map<Node*, absl::flat_hash_set<int32_t>> wait_for;
for (const auto& pair : dependency_edges) {
int32_t src_instance;
TF_RETURN_IF_ERROR(
GetNodeAttr(pair.first->attrs(), "instance_key", &src_instance));
for (Node* dst_node : pair.second) {
wait_for[dst_node].insert(src_instance);
}
}
for (const auto& pair : wait_for) {
std::vector<int32_t> wait_for_list(pair.second.begin(),
pair.second.end());
pair.first->ClearAttr("wait_for");
pair.first->AddAttr("wait_for", wait_for_list);
}
} else {
return absl::InternalError(absl::StrCat(
"Unexpected GraphCollectiveOrder type ", static_cast<int>(order_type)));
}
return absl::OkStatus();
}
} // namespace
absl::Status OrderCollectives(Graph* graph, GraphCollectiveOrder order_type) {
// `instance_keys[i]` corresponds to `collective_nodes[i]`
std::vector<Node*> collective_nodes;
std::vector<int32_t> instance_keys;
// node -> set of collectives on which node depends.
absl::flat_hash_map<Node*, absl::flat_hash_set<int32_t>> data_dependencies;
TF_RETURN_IF_ERROR(DiscoverDataDependencies(
graph, &collective_nodes, &instance_keys, &data_dependencies));
if (collective_nodes.empty()) return absl::OkStatus();
absl::flat_hash_map<Node*, absl::flat_hash_set<Node*>> dependency_edges;
// For all pairs of collective nodes n1 and n2 on the same device, if n1 does
// not depend on n2 and n2 does not depend on n1, then they are potentially
// concurrent. Create an arbitrary, deterministic ordering between them.
TF_RETURN_IF_ERROR(CreateControlDependencies(
collective_nodes, instance_keys, &data_dependencies, &dependency_edges));
return InsertControlDependencies(graph, order_type, dependency_edges);
}
} // namespace tensorflow
tensorflow/core/graph/collective_order.h
+36
View File
@@ -0,0 +1,36 @@
/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
Licensed 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 TENSORFLOW_CORE_GRAPH_COLLECTIVE_ORDER_H_
#define TENSORFLOW_CORE_GRAPH_COLLECTIVE_ORDER_H_
#include "tensorflow/core/graph/graph.h"
namespace tensorflow {
enum class GraphCollectiveOrder { kNone, kEdges, kAttrs };
// Introduces a deterministic execution order between potentially concurrent
// CollectiveOps. This may be used to execute collectives in the same order
// across all workers in a distributed execution, if all workers are executing
// the same graph.
// If `order_type` is `kEdges`, introduce the ordering in the form of explicit
// control edges between collective graph nodes. If `order_type` is `kAttrs`,
// add an attribute to the node which may be used by collective executor to
// ensure the required ordering.
absl::Status OrderCollectives(Graph* graph, GraphCollectiveOrder order_type);
} // namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPH_COLLECTIVE_ORDER_H_
tensorflow/core/graph/collective_order_test.cc
+236
View File
@@ -0,0 +1,236 @@
/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#include "tensorflow/core/graph/collective_order.h"
#include <gmock/gmock.h>
#include "tensorflow/core/common_runtime/graph_def_builder_util.h"
#include "tensorflow/core/framework/node_def_builder.h"
#include "tensorflow/core/graph/graph_def_builder.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace {
using ::testing::UnorderedElementsAreArray;
REGISTER_OP("TestParams").Output("o: float");
// Verifies that the list of collective nodes in `graph` matches
// `expected_collective_nodes`, and that the list of control edges between these
// collective nodes matches `expected_collective_control_edges`.
void VerifyGraph(const Graph& graph,
const std::vector<std::string>& expected_collective_nodes,
const std::vector<std::pair<std::string, std::string>>&
expected_collective_control_edges) {
std::vector<std::string> actual_collective_nodes;
std::vector<std::pair<std::string, std::string>>
actual_collective_control_edges;
for (const Node* src : graph.nodes()) {
if (!src->IsCollective()) {
continue;
}
actual_collective_nodes.push_back(src->name());
for (const Edge* edge : src->out_edges()) {
VLOG(2) << "collective edge " << edge->src()->name() << " -> "
<< edge->dst()->name();
// Add all control edges found except those to `_SINK`.
if (!edge->IsControlEdge() || edge->dst()->name() == "_SINK") {
continue;
}
actual_collective_control_edges.emplace_back(src->name(),
edge->dst()->name());
}
}
EXPECT_THAT(actual_collective_nodes,
UnorderedElementsAreArray(expected_collective_nodes));
EXPECT_THAT(actual_collective_control_edges,
UnorderedElementsAreArray(expected_collective_control_edges));
}
// Verifies that the `wait_for` attribute on collective nodes matches
// `wait_for_map`.
void VerifyAttrs(
const Graph& graph,
const std::unordered_map<std::string, std::vector<int32_t>> wait_for_map) {
for (const Node* node : graph.nodes()) {
if (node->IsCollective() ||
wait_for_map.find(node->name()) == wait_for_map.end()) {
continue;
}
std::vector<int32_t> wait_for_actual;
TF_EXPECT_OK(GetNodeAttr(node->attrs(), "wait_for", &wait_for_actual));
auto wait_for_expected = wait_for_map.at(node->name());
EXPECT_THAT(wait_for_actual, UnorderedElementsAreArray(wait_for_expected));
}
}
Node* CollectiveReduceNode(GraphDefBuilder* builder, Node* input,
const std::string& name, const std::string& device,
int instance_key) {
Node* collective_node =
ops::UnaryOp("CollectiveReduce", input,
builder->opts()
.WithName(name)
.WithDevice(device)
.WithAttr("T", DT_FLOAT)
.WithAttr("group_size", 2)
.WithAttr("group_key", 1)
.WithAttr("instance_key", instance_key)
.WithAttr("merge_op", "Add")
.WithAttr("final_op", "Id")
.WithAttr("subdiv_offsets", {1}));
return collective_node;
}
// Initialize the following graph:
//
// (cpu0) (cpu1)
// a b
// | |
// c1 c1
// | |
// id id
// / \ / \
// c2 c3 c2 c3
//
// Here ci denotes a collective node with `instance_key` i. `a` and `b` are
// inputs, `id` is identity node.
std::unique_ptr<Graph> InitGraph() {
GraphDefBuilder builder(GraphDefBuilder::kFailImmediately);
const std::string dev0 = "/job:localhost/replica:0/task:0/device:CPU:0";
const std::string dev1 = "/job:localhost/replica:0/task:0/device:CPU:1";
Node* a = ops::SourceOp("TestParams",
builder.opts().WithName("a").WithDevice(dev0));
Node* b = ops::SourceOp("TestParams",
builder.opts().WithName("b").WithDevice(dev1));
Node* c1_0 = CollectiveReduceNode(&builder, a, "c1_0", dev0, 1);
Node* c1_1 = CollectiveReduceNode(&builder, b, "c1_1", dev1, 1);
Node* id0 = ops::UnaryOp(
"Identity", c1_0,
builder.opts().WithName("id0").WithDevice(dev0).WithAttr("T", DT_FLOAT));
Node* id1 = ops::UnaryOp(
"Identity", c1_1,
builder.opts().WithName("id1").WithDevice(dev1).WithAttr("T", DT_FLOAT));
CollectiveReduceNode(&builder, id0, "c2_0", dev0, 2);
CollectiveReduceNode(&builder, id1, "c2_1", dev1, 2);
CollectiveReduceNode(&builder, id0, "c3_0", dev0, 3);
CollectiveReduceNode(&builder, id1, "c3_1", dev1, 3);
std::unique_ptr<Graph> graph = absl::make_unique<Graph>(OpRegistry::Global());
absl::Status s = GraphDefBuilderToGraph(builder, graph.get());
if (!s.ok()) {
LOG(FATAL) << "Error building graph " << s;
}
return graph;
}
// Tests that in the graph created by `InitGraph`, exactly 2 control edges are
// added after calling `OrderCollectives`: c3_0 -> c2_0 and c3_1 -> c2_1.
TEST(CollectiveOrderTest, SimpleOrder) {
std::unique_ptr<Graph> graph = InitGraph();
TF_EXPECT_OK(OrderCollectives(graph.get(), GraphCollectiveOrder::kEdges));
VerifyGraph(*graph, {"c1_0", "c1_1", "c2_0", "c2_1", "c3_0", "c3_1"},
{{"c3_0", "c2_0"}, {"c3_1", "c2_1"}});
}
TEST(CollectiveOrderTest, SimpleOrderAttr) {
std::unique_ptr<Graph> graph = InitGraph();
TF_EXPECT_OK(OrderCollectives(graph.get(), GraphCollectiveOrder::kAttrs));
VerifyAttrs(*graph, {{"c2_0", {3}}, {"c2_1", {3}}});
}
// Initialize the following graph:
//
// a
// |
// c1
// / \
// c4 id
// / \
// c2 c3
//
// Here ci denotes a collective node with `instance_key` i. `a` is an input,
// `id` is identity node.
std::unique_ptr<Graph> InitGraph2() {
GraphDefBuilder builder(GraphDefBuilder::kFailImmediately);
const std::string dev0 = "/job:localhost/replica:0/task:0/device:CPU:0";
Node* a = ops::SourceOp("TestParams",
builder.opts().WithName("a").WithDevice(dev0));
Node* c1 = CollectiveReduceNode(&builder, a, "c1", dev0, 1);
CollectiveReduceNode(&builder, c1, "c4", dev0, 4);
Node* id = ops::UnaryOp(
"Identity", c1,
builder.opts().WithName("id").WithDevice(dev0).WithAttr("T", DT_FLOAT));
CollectiveReduceNode(&builder, id, "c2", dev0, 2);
CollectiveReduceNode(&builder, id, "c3", dev0, 3);
std::unique_ptr<Graph> graph = absl::make_unique<Graph>(OpRegistry::Global());
absl::Status s = GraphDefBuilderToGraph(builder, graph.get());
if (!s.ok()) {
LOG(FATAL) << "Error building graph " << s;
}
return graph;
}
// Tests that in the graph created by `InitGraph2`, we add the following control
// edges after calling `OrderCollectives`: c4 -> c3, c3 -> c2. c4->c2 is
// pruned because it follows from the other two edges.
TEST(CollectiveOrderTest, SimpleOrder2) {
std::unique_ptr<Graph> graph = InitGraph2();
TF_EXPECT_OK(OrderCollectives(graph.get(), GraphCollectiveOrder::kEdges));
VerifyGraph(*graph, {"c1", "c2", "c3", "c4"}, {{"c4", "c3"}, {"c3", "c2"}});
}
// Initialize the following graph:
//
// w x y z
// | | | |
// c1 c2 c3 c4
//
std::unique_ptr<Graph> InitGraphForPruning() {
GraphDefBuilder builder(GraphDefBuilder::kFailImmediately);
const std::string dev0 = "/job:localhost/replica:0/task:0/device:CPU:0";
Node* w = ops::SourceOp("TestParams",
builder.opts().WithName("w").WithDevice(dev0));
Node* x = ops::SourceOp("TestParams",
builder.opts().WithName("x").WithDevice(dev0));
Node* y = ops::SourceOp("TestParams",
builder.opts().WithName("y").WithDevice(dev0));
Node* z = ops::SourceOp("TestParams",
builder.opts().WithName("z").WithDevice(dev0));
CollectiveReduceNode(&builder, w, "c1", dev0, 1);
CollectiveReduceNode(&builder, x, "c2", dev0, 2);
CollectiveReduceNode(&builder, y, "c3", dev0, 3);
CollectiveReduceNode(&builder, z, "c4", dev0, 4);
std::unique_ptr<Graph> graph = absl::make_unique<Graph>(OpRegistry::Global());
absl::Status s = GraphDefBuilderToGraph(builder, graph.get());
if (!s.ok()) {
LOG(FATAL) << "Error building graph " << s;
}
return graph;
}
// Tests that in the graph created by `InitGraphForPruning`, we only add c4 ->
// c3, c3 -> c2, c2 -> c1, and other edges are pruned away.
TEST(CollectiveOrderTest, Pruning) {
std::unique_ptr<Graph> graph = InitGraphForPruning();
TF_EXPECT_OK(OrderCollectives(graph.get(), GraphCollectiveOrder::kAttrs));
VerifyAttrs(*graph, {{"c3", {4}}, {"c2", {3}}, {"c1", {2}}});
}
} // namespace
} // namespace tensorflow
tensorflow/core/graph/colors.cc
+41
View File
@@ -0,0 +1,41 @@
/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#include "tensorflow/core/graph/colors.h"
#include "tensorflow/core/platform/macros.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
// Color palette
// http://www.mulinblog.com/a-color-palette-optimized-for-data-visualization/
static const char* kColors[] = {
"#F15854", // red
"#5DA5DA", // blue
"#FAA43A", // orange
"#60BD68", // green
"#F17CB0", // pink
"#B2912F", // brown
"#B276B2", // purple
"#DECF3F", // yellow
"#4D4D4D", // gray
};
const char* ColorFor(int dindex) {
return kColors[dindex % TF_ARRAYSIZE(kColors)];
}
} // namespace tensorflow
tensorflow/core/graph/colors.h
+29
View File
@@ -0,0 +1,29 @@
/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
Licensed 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 TENSORFLOW_CORE_GRAPH_COLORS_H_
#define TENSORFLOW_CORE_GRAPH_COLORS_H_
namespace tensorflow {
// Return a color drawn from a palette to represent an entity
// identified by "i". The return value has the form "#RRGGBB" Note
// that the palette has a limited set of colors and therefore colors
// will be reused eventually.
const char* ColorFor(int dindex);
} // namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPH_COLORS_H_
tensorflow/core/graph/control_flow.cc
+188
View File
@@ -0,0 +1,188 @@
/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#include "tensorflow/core/graph/control_flow.h"
#include <deque>
#include <vector>
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/graph/node_builder.h"
#include "tensorflow/core/lib/core/errors.h"
namespace tensorflow {
namespace {
// Information about a loop frame structure.
struct Frame {
std::string name;
// Pointer to the parent frame. The root frame has a pointer to itself.
Frame* parent = nullptr;
// The loop condition of the loop. There should be exactly one loop condition
// in every loop.
const Node* loop_cond = nullptr;
};
// Verify that the ControlFlowInfo of the graph has valid loop structure.
absl::Status ValidateControlFlowInfo(
const Graph* graph, const std::vector<ControlFlowInfo>& cf_info) {
std::unordered_map<std::string, Frame> frames;
for (const Node* node : graph->op_nodes()) {
const ControlFlowInfo& cf = cf_info[node->id()];
if (!cf.frame || !cf.parent_frame) {
// Skip nodes unreachable from the source node. They might be pruned
// later.
continue;
}
Frame& frame = frames[cf.frame_name];
Frame* parent = &frames[cf_info[cf.parent_frame->id()].frame_name];
if (frame.parent == nullptr) {
frame.parent = parent;
frame.name = cf.frame_name;
} else if (frame.parent != parent) {
return absl::InternalError(absl::StrCat(
"Invalid loop structure: Mismatched parent frames for \"",
cf.frame_name, "\": \"", parent->name, "\" vs \"", frame.parent->name,
"\". The node giving this error: ", FormatNodeForError(*node),
". This is an internal bug, please file a bug report with "
"instructions on how to reproduce the error."));
}
if (IsLoopCond(node)) {
// ForwardLoopCounter runs in the same frame as the forward loop and
// BackPropLoopCounter runs in the same frame as the backprop loop. They
// are the only cases that multiple loops share the same frame.
if (frame.loop_cond &&
!absl::StrContains(frame.loop_cond->name(), "LoopCounter") &&
!absl::StrContains(node->name(), "LoopCounter")) {
return absl::InvalidArgumentError(absl::StrCat(
"Invalid loop structure: Loop \"", cf.frame_name,
"\" has more than one LoopCond node: ", FormatNodeForError(*node),
" and ", FormatNodeForError(*frame.loop_cond),
". This is an internal bug, please file a bug report with "
"instructions on how to reproduce the error."));
}
frame.loop_cond = node;
}
}
return absl::OkStatus();
}
} // namespace
absl::Status BuildControlFlowInfo(const Graph* g,
std::vector<ControlFlowInfo>* info,
std::vector<std::string>* unreachable_nodes) {
info->clear();
info->resize(g->num_node_ids());
std::vector<const Node*> parent_nodes;
parent_nodes.resize(g->num_node_ids());
const Node* src_node = g->source_node();
ControlFlowInfo& src_info = (*info)[src_node->id()];
src_info.frame = src_node;
src_info.parent_frame = src_node;
std::string frame_name;
std::deque<const Node*> ready;
ready.push_back(src_node);
while (!ready.empty()) {
const Node* curr_node = ready.front();
ready.pop_front();
const ControlFlowInfo& curr_info = (*info)[curr_node->id()];
const Node* frame = curr_info.frame;
const Node* parent = curr_info.parent_frame;
frame_name = curr_info.frame_name;
if (IsExit(curr_node)) {
// Exit to the parent frame.
const ControlFlowInfo& parent_info = (*info)[parent->id()];
frame = parent_info.frame;
parent = parent_info.parent_frame;
frame_name = parent_info.frame_name;
}
for (const Edge* out_edge : curr_node->out_edges()) {
const Node* out = out_edge->dst();
int out_id = out->id();
ControlFlowInfo* out_info = &(*info)[out_id];
const Node* out_parent = out_info->parent_frame;
bool is_visited = (parent_nodes[out_id] != nullptr);
// Skip Sink/Source nodes.
if (!out->IsOp()) continue;
// Add to ready queue if not seen.
if (!is_visited) {
parent_nodes[out->id()] = curr_node;
ready.push_back(out);
}
// Process the node 'out'.
if (IsEnter(out)) {
if (is_visited) {
const std::string& parent_frame =
(*info)[out_parent->id()].frame_name;
if (parent_frame != frame_name) {
return absl::InvalidArgumentError(absl::StrCat(
FormatNodeForError(*out),
" has inputs from different frames. The input ",
FormatNodeForError(*curr_node), " is in frame '", frame_name,
"'. The input ", FormatNodeForError(*parent_nodes[out->id()]),
" is in frame '", parent_frame, "'."));
}
} else {
out_info->frame = out;
out_info->parent_frame = frame;
TF_RETURN_IF_ERROR(
GetNodeAttr(out->attrs(), "frame_name", &out_info->frame_name));
if (out_info->frame_name.empty()) {
return absl::InvalidArgumentError(
absl::StrCat("The Enter ", FormatNodeForError(*out),
" must have a frame name."));
}
}
} else {
if (is_visited) {
if (out_info->frame_name != frame_name) {
return absl::InvalidArgumentError(absl::StrCat(
FormatNodeForError(*out),
" has inputs from different frames. The input ",
FormatNodeForError(*curr_node), " is in frame '", frame_name,
"'. The input ", FormatNodeForError(*parent_nodes[out->id()]),
" is in frame '", out_info->frame_name, "'."));
}
} else {
out_info->frame = frame;
out_info->parent_frame = parent;
out_info->frame_name = frame_name;
}
}
}
}
if (unreachable_nodes) {
for (const Node* node : g->op_nodes()) {
if (!parent_nodes[node->id()]) {
unreachable_nodes->push_back(node->name());
}
}
}
TF_RETURN_IF_ERROR(ValidateControlFlowInfo(g, *info));
return absl::OkStatus();
}
} // namespace tensorflow
tensorflow/core/graph/control_flow.h
+61
View File
@@ -0,0 +1,61 @@
/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
Licensed 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 TENSORFLOW_CORE_GRAPH_CONTROL_FLOW_H_
#define TENSORFLOW_CORE_GRAPH_CONTROL_FLOW_H_
#include <vector>
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/lib/core/status.h"
namespace tensorflow {
// Control flow info for a graph node.
struct ControlFlowInfo {
// 'frame' and 'parent_frame' are pointers to:
//
// a) One of the Enter nodes corresponding to the loop body, if the node
// executes inside a loop. If multiple tensors enter the while loop, it's
// undefined which Enter node will be used.
//
// b) SOURCE node (node.id() == Graph::kSourceId), if the node is not inside
// any of the while loops.
const Node* frame = nullptr; // frame of a node
const Node* parent_frame = nullptr; // parent frame of a node
std::string frame_name; // frame name of a node
};
// Clear and populate `info` with each node's frame and the level it belongs to.
// We check the well-formedness of the graph:
// 1) All inputs to a node must come from the same frame and have the same
// "static" iteration level.
// 2) Each frame has at most one LoopCond node.
// 3) Each frame has a single parent frame.
// If `unreachable_nodes` is set, return names of nodes unreachable from the
// source node. We cannot build ControlFlowInfo for such nodes. They might be
// pruned later.
//
// NOTE(yuanbyu): For now, we require all sends/recvs have iteration level 0.
// This essentially means there can't be multiple serial Nexts in an iteration,
// which all sane front-ends should satisfy.
absl::Status BuildControlFlowInfo(
const Graph* g, std::vector<ControlFlowInfo>* info,
std::vector<std::string>* unreachable_nodes = nullptr);
} // namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPH_CONTROL_FLOW_H_
tensorflow/core/graph/control_flow_test.cc
+146
View File
@@ -0,0 +1,146 @@
/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#include "tensorflow/core/graph/control_flow.h"
#include <string>
#include <vector>
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/cc/ops/while_loop.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/lib/strings/str_util.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace {
absl::Status LessThanTenCond(const Scope& scope,
const std::vector<Output>& inputs,
Output* output) {
*output = ops::Less(scope, inputs[0], 10);
return scope.status();
}
absl::Status AddOneBody(const Scope& scope, const std::vector<Output>& inputs,
std::vector<Output>* outputs) {
outputs->push_back(ops::AddN(scope, {inputs[0], 1}));
return scope.status();
}
absl::Status NestedLoopBody(const Scope& scope,
const std::vector<Output>& inputs,
std::vector<Output>* outputs) {
return ops::BuildWhileLoop(scope.NewSubScope("inner"), inputs,
LessThanTenCond, AddOneBody, "inner_loop",
outputs);
}
TEST(ValidateControlFlowTest, InputsFromDifferentFrames) {
Scope scope = Scope::NewRootScope().ExitOnError();
std::vector<Output> inputs;
inputs.push_back(ops::Placeholder(scope, DT_INT32));
std::vector<Output> outputs;
TF_ASSERT_OK(ops::BuildWhileLoop(scope.NewSubScope("outer"), inputs,
LessThanTenCond, NestedLoopBody,
"outer_loop", &outputs));
std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
TF_ASSERT_OK(scope.ToGraph(graph.get()));
// {inner/Enter', 'outer/Switch'} --> 'inner/Merge'. 'inner/Enter' is in frame
// 'inner_loop'. 'outer/Switch' is in frame 'outer_loop'.
std::vector<ControlFlowInfo> info;
absl::Status status = BuildControlFlowInfo(graph.get(), &info);
EXPECT_FALSE(status.ok());
EXPECT_TRUE(
absl::StrContains(status.message(), "has inputs from different frames"))
<< status.message();
EXPECT_TRUE(
absl::StrContains(status.message(), "{{node outer/body/inner/Merge}}"))
<< status.message();
EXPECT_TRUE(
absl::StrContains(status.message(), "{{node outer/body/inner/Enter}}"))
<< status.message();
EXPECT_TRUE(absl::StrContains(status.message(), "{{node outer/Switch}}"))
<< status.message();
}
TEST(ValidateControlFlowTest, MismatchedParentFrames) {
Scope scope = Scope::NewRootScope().ExitOnError();
std::vector<Output> inputs;
inputs.push_back(ops::Placeholder(scope, DT_INT32));
std::vector<Output> outputs;
TF_ASSERT_OK(ops::BuildWhileLoop(scope, inputs, LessThanTenCond, AddOneBody,
"test_loop", &outputs));
std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
TF_ASSERT_OK(scope.ToGraph(graph.get()));
Node* enter_1 = nullptr;
for (Node* node : graph->op_nodes()) {
if (IsEnter(node)) {
enter_1 = node;
}
}
ASSERT_TRUE(enter_1 != nullptr);
NodeDef enter;
enter.set_name("Enter2");
enter.set_op("Enter");
(*enter.mutable_attr())["T"].set_type(DT_INT32);
(*enter.mutable_attr())["frame_name"].set_s("test_loop");
*enter.add_input() = "Enter";
absl::Status status;
Node* enter_2 = graph->AddNode(enter, &status);
TF_ASSERT_OK(status);
graph->AddControlEdge(enter_1, enter_2);
// SOURCE("") --> Enter("test_loop") --> Enter2("test_loop")
// For node 'Enter', the parent frame of "test_loop" is empty.
// For node 'Enter2', the parent frame of "test_loop" is "test_loop".
std::vector<ControlFlowInfo> info;
status = BuildControlFlowInfo(graph.get(), &info);
EXPECT_FALSE(status.ok());
EXPECT_TRUE(absl::StrContains(status.message(), "Mismatched parent frames"))
<< status.message();
EXPECT_TRUE(absl::StrContains(status.message(), "{{node Enter2}}"))
<< status.message();
}
TEST(ValidateControlFlowTest, TwoLoopCond) {
// Test that one frame has at most one LoopCond node. This is necessary for
// functionalize control flow.
Scope scope = Scope::NewRootScope().ExitOnError();
std::vector<Output> inputs;
inputs.push_back(ops::Placeholder(scope, DT_INT32));
std::vector<Output> outputs;
TF_ASSERT_OK(ops::BuildWhileLoop(scope, inputs, LessThanTenCond, AddOneBody,
"test_loop", &outputs));
outputs.clear();
TF_ASSERT_OK(ops::BuildWhileLoop(scope.NewSubScope("sub"), inputs,
LessThanTenCond, AddOneBody, "test_loop",
&outputs, false));
std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
TF_ASSERT_OK(scope.ToGraph(graph.get()));
std::vector<ControlFlowInfo> info;
absl::Status status = BuildControlFlowInfo(graph.get(), &info);
EXPECT_FALSE(status.ok());
EXPECT_TRUE(
absl::StrContains(status.message(), "more than one LoopCond node"))
<< status.message();
EXPECT_TRUE(absl::StrContains(status.message(), "{{node sub/LoopCond}}"))
<< status.message();
EXPECT_TRUE(absl::StrContains(status.message(), "{{node LoopCond}}"))
<< status.message();
}
} // namespace
} // namespace tensorflow
tensorflow/core/graph/costmodel.cc
+586
View File
@@ -0,0 +1,586 @@
/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#include "tensorflow/core/graph/costmodel.h"
#include <algorithm>
#include <vector>
#include "tensorflow/core/framework/allocation_description.pb.h"
#include "tensorflow/core/framework/cost_graph.pb.h"
#include "tensorflow/core/framework/step_stats.pb.h"
#include "tensorflow/core/framework/tensor_description.pb.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/platform/logging.h"
namespace tensorflow {
namespace {
const Microseconds kDefaultTimeEstimate(1);
const Microseconds kMinTimeEstimate(1);
} // namespace
void CostModel::SuppressInfrequent() {
// Find the median of the non-zero counts, and use half of its value
// as the cutoff for a "normal" execution mode node.
if (count_.empty()) return;
std::vector<int32_t> non_zero;
for (auto v : count_) {
if (v > 0) non_zero.push_back(v);
}
const size_t sz = non_zero.size();
if (sz > 0) {
std::nth_element(non_zero.begin(), non_zero.begin() + sz / 2,
non_zero.end());
int32_t median_value = non_zero[sz / 2];
min_count_ = median_value / 2;
VLOG(1) << "num non_zero vals: " << non_zero.size() << " median_value "
<< median_value;
} else {
min_count_ = 1;
}
}
void CostModel::MergeFromLocal(const Graph& g, const CostModel& cm) {
CHECK(is_global_);
CHECK(!cm.is_global());
for (const Node* n : g.nodes()) {
const int local_id = cm.Id(n);
const int global_id = Id(n);
if (local_id < 0 || global_id < 0) continue;
int num_slots = cm.slot_bytes_[local_id].size();
Ensure(global_id, num_slots);
count_[global_id] += cm.count_[local_id];
time_[global_id] += cm.time_[local_id];
if (num_slots > 0) {
if (slot_bytes_[global_id].empty()) {
slot_bytes_[global_id].resize(num_slots);
} else {
CHECK_EQ(num_slots, slot_bytes_[global_id].size());
}
for (int s = 0; s < num_slots; ++s) {
auto& current_v = slot_bytes_[global_id][s];
auto other_v = cm.slot_bytes_[local_id][s];
if (current_v < 0) {
current_v = other_v;
} else if (other_v > 0) {
current_v += other_v;
}
}
}
}
}
void CostModel::MergeFromGlobal(const CostModel& cm) {
CHECK(is_global_);
CHECK_EQ(true, cm.is_global());
const int num_nodes = cm.count_.size();
for (int i = num_nodes - 1; i >= 0; --i) {
int num_slots = cm.slot_bytes_[i].size();
Ensure(i, num_slots);
count_[i] += cm.count_[i];
time_[i] += cm.time_[i];
if (num_slots > 0) {
if (slot_bytes_[i].empty()) {
slot_bytes_[i].resize(num_slots);
} else {
CHECK_EQ(num_slots, slot_bytes_[i].size());
}
for (int s = 0; s < num_slots; ++s) {
auto& current_v = slot_bytes_[i][s];
auto other_v = cm.slot_bytes_[i][s];
if (current_v < 0) {
current_v = other_v;
} else if (other_v > 0) {
current_v += other_v;
}
}
}
}
}
void CostModel::MergeFromStats(const NodeNameToCostIdMap& map,
const StepStats& ss) {
CHECK(is_global_);
for (auto& ds : ss.dev_stats()) {
for (auto& ns : ds.node_stats()) {
NodeNameToCostIdMap::const_iterator iter = map.find(ns.node_name());
// We don't keep stats for nodes not in the global graph, i.e.
// copy/send/recv nodes, feed/fetch, etc.
if (iter == map.end()) continue;
int32_t global_id = iter->second;
Ensure(global_id, ns.output_size());
int64_t elapsed_micros =
ns.op_end_rel_micros() - ns.op_start_rel_micros();
count_[global_id]++;
time_[global_id] += elapsed_micros;
for (auto& no : ns.output()) {
int si = no.slot();
if (static_cast<size_t>(si) >= slot_bytes_[global_id].size()) {
slot_bytes_[global_id].resize(1 + si);
}
auto& current_v = slot_bytes_[global_id][si];
auto other_v =
no.tensor_description().allocation_description().requested_bytes();
if (current_v < 0) {
current_v = other_v;
} else if (other_v > 0) {
current_v += other_v;
}
}
}
}
}
void CostModel::Ensure(int id, int num_outputs) {
if (slot_bytes_.size() <= static_cast<size_t>(id)) {
slot_bytes_.resize(id + 1);
count_.resize(id + 1);
time_.resize(id + 1);
max_mem_usage_.resize(id + 1);
max_exec_time_.resize(id + 1);
output_port_alloc_ids_.resize(id + 1);
}
if (num_outputs > 0) {
auto perslot = &slot_bytes_[id];
auto output_port_alloc_ids = &output_port_alloc_ids_[id];
auto max_mem_usage = &max_mem_usage_[id];
CHECK_LE(perslot->size(), num_outputs);
DCHECK_EQ(output_port_alloc_ids->size(), perslot->size());
DCHECK_EQ(max_mem_usage->output_port_mem.size(), perslot->size());
DCHECK_EQ(max_mem_usage->output_port_shape.size(), perslot->size());
DCHECK_EQ(max_mem_usage->output_port_type.size(), perslot->size());
perslot->resize(num_outputs, Bytes(-1));
output_port_alloc_ids->resize(num_outputs, -1);
max_mem_usage->output_port_mem.resize(num_outputs, Bytes(-1));
max_mem_usage->output_port_shape.resize(num_outputs, unknown_shape_);
max_mem_usage->output_port_type.resize(num_outputs, DT_INVALID);
}
}
void CostModel::SetNumOutputs(const Node* node, int num_outputs) {
const int id = Id(node);
if (id < 0) return;
// Do not resize the number of slots before checking its existing number of
// slots.
Ensure(id, 0);
auto perslot = &slot_bytes_[id];
if (!perslot->empty()) {
CHECK_EQ(num_outputs, perslot->size())
<< "Cannot resize slot_bytes, node=" << node->name();
}
Ensure(id, num_outputs);
}
void CostModel::RecordCount(const Node* node, int count) {
const int id = Id(node);
if (id < 0) return;
CHECK_LT(id, slot_bytes_.size());
count_[id] += count;
}
int32_t CostModel::TotalCount(const Node* node) const {
const int id = Id(node);
if (id < 0) return 0;
return (static_cast<size_t>(id) < slot_bytes_.size()) ? count_[id] : 0;
}
void CostModel::RecordSize(const Node* node, int slot, Bytes bytes) {
const int id = Id(node);
if (id < 0) return;
CHECK_LT(id, slot_bytes_.size());
auto perslot = &slot_bytes_[id];
CHECK_LT(slot, perslot->size());
auto v = &(*perslot)[slot];
if (*v >= 0) {
*v += bytes;
} else {
*v = bytes;
}
}
Bytes CostModel::TotalBytes(const Node* node, int slot) const {
const int id = Id(node);
if (id < 0 || static_cast<size_t>(id) >= slot_bytes_.size() ||
slot_bytes_[id].size() <= static_cast<size_t>(slot)) {
return Bytes(0);
}
return slot_bytes_[id][slot];
}
Bytes CostModel::SizeEstimate(const Node* node, int slot) const {
int32_t count = TotalCount(node);
if (count < min_count_) return Bytes(0);
return TotalBytes(node, slot) / std::max(1, TotalCount(node));
}
void CostModel::RecordTime(const Node* node, Microseconds time) {
const int id = Id(node);
if (id < 0) return;
DCHECK(node->IsOp()) << node->DebugString();
Ensure(id, node->num_outputs());
time_[id] += time;
}
Microseconds CostModel::TotalTime(const Node* node) const {
DCHECK(node->IsOp()) << node->DebugString();
const int id = Id(node);
if (id < 0 || static_cast<size_t>(id) >= time_.size() ||
time_[id] < Microseconds(0)) {
return Microseconds(0);
}
return time_[id];
}
Microseconds CostModel::TimeEstimate(const Node* node) const {
int32_t count = TotalCount(node);
if (count <= min_count_) return kMinTimeEstimate;
return std::max(kMinTimeEstimate, TotalTime(node) / std::max(1, count));
}
void CostModel::CheckInitialized(const Graph& graph) const {
for (const Node* n : graph.op_nodes()) {
CHECK(static_cast<size_t>(n->id()) < time_.size() &&
time_[n->id()] >= Microseconds(0))
<< ": no time estimate for " << n->DebugString();
CHECK(static_cast<size_t>(n->id()) < slot_bytes_.size())
<< ": no size estimate for " << n->DebugString();
const auto& perslot = slot_bytes_[n->id()];
for (size_t i = 0; i < perslot.size(); i++) {
CHECK_GE(perslot[i], Bytes(0)) << ": no size estimate for output# " << i
<< " of " << n->DebugString();
}
}
}
void CostModel::RecordMaxMemorySize(const Node* node, int output_slot,
Bytes bytes,
const TensorShapeProto& tensor_shape,
const DataType& dtype) {
const int id = Id(node);
if (id < 0) return;
if (output_slot >= node->num_outputs()) {
LOG(ERROR) << "Unexpected output slot for node " << node->DebugString()
<< ". Got " << output_slot << " but its num_outputs is "
<< node->num_outputs();
return;
}
Ensure(id, node->num_outputs());
auto& current_max = max_mem_usage_[id].output_port_mem[output_slot];
// If the memory allocator doesn't track memory usage, let's infer a lower
// bound from the tensor shape and its data type.
if (bytes.value() < 0) {
bytes = MinTensorMemoryUsage(tensor_shape, dtype);
}
if (bytes.value() > current_max.value()) {
current_max = bytes.value();
max_mem_usage_[id].output_port_shape[output_slot] = tensor_shape;
max_mem_usage_[id].output_port_type[output_slot] = dtype;
}
}
Bytes CostModel::MaxMemorySize(const Node* node, int slot) const {
const int id = Id(node);
if (id < 0 || static_cast<size_t>(id) >= max_mem_usage_.size() ||
max_mem_usage_[id].output_port_mem.size() <= static_cast<size_t>(slot)) {
return Bytes(0);
}
return max_mem_usage_[id].output_port_mem[slot];
}
const TensorShapeProto& CostModel::MaxMemoryShape(const Node* node,
int slot) const {
const int id = Id(node);
if (id < 0 || static_cast<size_t>(id) >= max_mem_usage_.size() ||
max_mem_usage_[id].output_port_shape.size() <=
static_cast<size_t>(slot)) {
return unknown_shape_;
}
return max_mem_usage_[id].output_port_shape[slot];
}
DataType CostModel::MaxMemoryType(const Node* node, int slot) const {
const int id = Id(node);
if (id < 0 || static_cast<size_t>(id) >= max_mem_usage_.size() ||
max_mem_usage_[id].output_port_type.size() <= static_cast<size_t>(slot)) {
return DT_INVALID;
}
return max_mem_usage_[id].output_port_type[slot];
}
Bytes CostModel::TempMemorySize(const Node* node) const {
const int id = Id(node);
if (id < 0 || static_cast<size_t>(id) >= max_mem_usage_.size()) {
return Bytes(0);
}
return max_mem_usage_[id].temp_memory_size;
}
Bytes CostModel::PersistentMemorySize(const Node* node) const {
const int id = Id(node);
if (id < 0 || static_cast<size_t>(id) >= max_mem_usage_.size()) {
return Bytes(0);
}
return max_mem_usage_[id].persistent_memory_size;
}
void CostModel::RecordMemoryStats(const Node* node,
const MemoryStats& memory_stats) {
const int id = Id(node);
if (id < 0) return;
Ensure(id, node->num_outputs());
max_mem_usage_[id].temp_memory_size = memory_stats.temp_memory_size();
max_mem_usage_[id].persistent_memory_size =
memory_stats.persistent_memory_size();
for (int64_t alloc_id : memory_stats.persistent_tensor_alloc_ids()) {
if (alloc_id > 0) {
persistent_alloc_ids_.insert(alloc_id);
}
}
}
void CostModel::RecordMaxExecutionTime(const Node* node, Microseconds time) {
const int id = Id(node);
if (id < 0) return;
Ensure(id, node->num_outputs());
max_exec_time_[id] = std::max(max_exec_time_[id], time);
}
Microseconds CostModel::MaxExecutionTime(const Node* node) const {
const int id = Id(node);
if (id < 0 || static_cast<size_t>(id) >= max_exec_time_.size()) {
return Microseconds(0);
}
return max_exec_time_[id];
}
void CostModel::RecordAllocationId(const Node* node, int output_slot,
int64_t alloc_id) {
const int id = Id(node);
if (id < 0) return;
Ensure(id, node->num_outputs());
output_port_alloc_ids_[id][output_slot] = alloc_id;
}
int64_t CostModel::AllocationId(const Node* node, int slot) const {
const int id = Id(node);
if (id < 0 || static_cast<size_t>(id) >= output_port_alloc_ids_.size() ||
output_port_alloc_ids_[id].size() <= static_cast<size_t>(slot)) {
return -1;
}
return output_port_alloc_ids_[id][slot];
}
bool CostModel::IsPersistentTensor(const Node* node, int64_t alloc_id) const {
if (persistent_alloc_ids_.count(alloc_id) > 0) {
return true;
}
return false;
}
Microseconds CostModel::CopyTimeEstimate(Bytes b, double network_latency_millis,
double estimated_gbps) {
// TODO(jeff,sanjay): estimate cost based on bandwidth along the
// communication path and the type of transport we are using between
// devices.
//
// We assume the copy time follows a linear model:
// copy_time = copy_bytes / rate + min_time
int64_t copy_bytes = b.value();
const double bytes_per_usec = estimated_gbps * 1000.0 / 8;
const double min_micros = network_latency_millis * 1000.0;
return Microseconds(
static_cast<int64_t>(copy_bytes / bytes_per_usec + min_micros));
}
Microseconds CostModel::ComputationTimeEstimate(int64_t math_ops) {
// TODO(jeff,sanjay): Eventually we should pass in the type of device
// (GPU vs. CPU) and use that to affect the estimate.
// We estimate the microseconds using that value. We divide
// by 1000 to convert the madd number into microseconds (assuming
// roughly 1000 madds per microsecond (~1 GHz for one core)).
return Microseconds(math_ops / 1000);
}
void CostModel::IncrementUpdateTimes() { update_times_++; }
int32_t CostModel::GetUpdateTimes() const { return update_times_; }
// ----------------------------------------------------------------------------
// InitCostModel
// ----------------------------------------------------------------------------
namespace {
static void AddNodesToCostModel(const Graph& g, CostModel* cost_model) {
for (Node* n : g.nodes()) {
const int num_outputs = n->num_outputs();
cost_model->SetNumOutputs(n, num_outputs);
for (int output = 0; output < num_outputs; output++) {
// Set up an initial bogus estimate for the node's outputs
cost_model->RecordSize(n, output, Bytes(1));
}
}
}
static void AssignSizes(const Graph& g, CostModel* cost_model) {
for (const Edge* e : g.edges()) {
// Skip if it is a control edge.
if (e->IsControlEdge()) {
continue;
}
const Node* src = e->src();
// TODO(josh11b): Get an estimate from the Op
Bytes size(1);
cost_model->RecordSize(src, e->src_output(), size);
}
}
// This generates an extremely simple initial guess for the
// computation cost of each node. For ordinary Ops, its value should quickly
// be wiped out by the real runtime measurements. For other Ops we don't
// actually generate measurements, so suppression of infrequent Ops ends up
// giving them 0 costs. So, this is not of much consequence except perhaps
// in tests.
static Microseconds TimeEstimateForNode(CostModel* cost_model, Node* n) {
CHECK(n->IsOp());
VLOG(2) << "Node " << n->id() << ": " << n->name()
<< " type_string: " << n->type_string();
if (IsConstant(n) || IsVariable(n)) {
return Microseconds(0);
}
return kDefaultTimeEstimate;
}
static void EstimateComputationCosts(const Graph& g, CostModel* cost_model) {
for (Node* n : g.nodes()) {
if (!n->IsOp()) continue;
cost_model->RecordTime(n, TimeEstimateForNode(cost_model, n));
}
}
} // namespace
void CostModel::InitFromGraph(const Graph& g) {
const int num_node_ids = g.num_node_ids();
slot_bytes_.reserve(num_node_ids);
count_.reserve(num_node_ids);
time_.reserve(num_node_ids);
max_mem_usage_.reserve(num_node_ids);
max_exec_time_.reserve(num_node_ids);
output_port_alloc_ids_.reserve(num_node_ids);
AddNodesToCostModel(g, this);
AssignSizes(g, this);
EstimateComputationCosts(g, this);
CheckInitialized(g);
}
void CostModel::AddToCostGraphDef(const Graph* graph,
CostGraphDef* cost_graph) const {
std::vector<const Edge*> inputs;
std::vector<const Edge*> control_inputs;
int offset = cost_graph->node_size();
for (const Node* n : graph->nodes()) {
CostGraphDef::Node* cnode = cost_graph->add_node();
cnode->set_name(n->name());
cnode->set_device(n->assigned_device_name());
cnode->set_id(GlobalId(n, offset));
inputs.clear();
inputs.resize(n->num_inputs(), nullptr);
control_inputs.clear();
for (const Edge* e : n->in_edges()) {
if (e->IsControlEdge()) {
control_inputs.push_back(e);
} else {
inputs[e->dst_input()] = e;
}
}
std::sort(control_inputs.begin(), control_inputs.end(),
[this](Edge const* a, Edge const* b) {
return Id(a->src()) < Id(b->src());
});
for (const Edge* e : inputs) {
CostGraphDef::Node::InputInfo* input_info = cnode->add_input_info();
input_info->set_preceding_node(GlobalId(e->src(), offset));
input_info->set_preceding_port(e->src_output());
}
for (int i = 0; i < n->num_outputs(); i++) {
CostGraphDef::Node::OutputInfo* output_info = cnode->add_output_info();
int64_t alloc_id = AllocationId(n, i);
int64_t alias_to_input = -1;
for (const Edge* e : inputs) {
int64_t input_alloc_id = AllocationId(e->src(), e->src_output());
if (input_alloc_id == alloc_id) {
alias_to_input = e->dst_input();
break;
}
}
output_info->set_alias_input_port(alias_to_input);
output_info->set_dtype(MaxMemoryType(n, i));
*output_info->mutable_shape() = MaxMemoryShape(n, i);
if (alias_to_input < 0 && IsPersistentTensor(n, alloc_id)) {
output_info->set_size(0);
} else {
output_info->set_size(MaxMemorySize(n, i).value());
}
}
for (const Edge* e : control_inputs) {
cnode->add_control_input(GlobalId(e->src(), offset));
}
cnode->set_temporary_memory_size(TempMemorySize(n).value());
cnode->set_persistent_memory_size(PersistentMemorySize(n).value());
cnode->set_compute_cost(MaxExecutionTime(n).value());
// For now we treat all send nodes as final.
// TODO(yuanbyu): Send nodes for fetches shouldn't be treated as final.
cnode->set_is_final(n->IsSend());
}
}
void CostModel::WriteSummaryToLog() const {
LOG(INFO) << " min_count_=" << min_count_;
for (size_t i = 0; i < count_.size(); ++i) {
LOG(INFO) << "Node " << i << " count " << count_[i] << " total time "
<< time_[i] << " avg time "
<< (time_[i] / (std::max(1, count_[i])));
}
}
Bytes CostModel::MinTensorMemoryUsage(const TensorShapeProto& tensor_shape,
const DataType& dtype) {
if (tensor_shape.unknown_rank()) {
return Bytes(-1);
}
size_t num_coefficients = 1;
for (const TensorShapeProto::Dim& dim : tensor_shape.dim()) {
// If the dimension is unknown, it has to be at least 1
num_coefficients *= std::max<size_t>(dim.size(), 1);
}
return Bytes(num_coefficients * DataTypeSize(dtype));
}
} // namespace tensorflow
tensorflow/core/graph/costmodel.h
+241
View File
@@ -0,0 +1,241 @@
/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
Licensed 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 TENSORFLOW_CORE_GRAPH_COSTMODEL_H_
#define TENSORFLOW_CORE_GRAPH_COSTMODEL_H_
#include <set>
#include <unordered_map>
#include <vector>
#include "tensorflow/core/framework/cost_graph.pb.h"
#include "tensorflow/core/framework/step_stats.pb.h"
#include "tensorflow/core/framework/tensor_shape.pb.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/graph/types.h"
#include "tensorflow/core/lib/core/stringpiece.h"
#include "tensorflow/core/lib/gtl/array_slice.h"
#include "tensorflow/core/platform/macros.h"
#include "tensorflow/core/platform/protobuf.h"
namespace tensorflow {
typedef std::unordered_map<absl::string_view, int32_t, StringPieceHasher>
NodeNameToCostIdMap;
class StepStats;
// CostModel keeps track of the following runtime statistics for nodes
// of a single Graph:
// * The total number of times a node has executed.
// * The accumulated execution time (in microseconds) of a node.
// * The accumulated size (in bytes) of each node's output.
//
// This class is NOT thread-safe.
class CostModel {
public:
// If "global" is true, maintains costs based on Node::cost_id, otherwise
// maintains costs based on Node::id.
explicit CostModel(bool is_global) : is_global_(is_global) {
unknown_shape_.set_unknown_rank(true);
}
// Assigns min_count_ as a function of the median count for a Node.
// This value is then used for suppressing the time/size costs of
// infrequent operations.
// NOTE(tucker): Maybe this should move to a subclass of CostModel.
void SuppressInfrequent();
bool is_global() const { return is_global_; }
int Id(const Node* n) const {
if (is_global_) {
return n->cost_id();
} else {
return n->id();
}
}
int GlobalId(const Node* n, int offset) const {
if (is_global_) {
return n->cost_id();
} else {
return n->id() + offset;
}
}
// Initializes cost model for 'g'.
void InitFromGraph(const Graph& g);
// Merges costs from cm.
// REQUIRES: is_global_ is true for this and for "cm"
void MergeFromGlobal(const CostModel& cm);
// Merges costs from "cm", which has been computed relative to "g".
// REQUIRES: is_global_ is true for this, and false for "cm".
void MergeFromLocal(const Graph& g, const CostModel& cm);
void MergeFromStats(const NodeNameToCostIdMap& map, const StepStats& ss);
// Sets the number of outputs of "node".
void SetNumOutputs(const Node* node, int num_outputs);
// Records that "node" has executed "num_count" more times.
void RecordCount(const Node* node, int num_count);
// Returns how many times "node" has been executed.
int32_t TotalCount(const Node* node) const;
// Records that "output_slot" of "node" has produced tensors of
// aggregated "bytes".
void RecordSize(const Node* node, int output_slot, Bytes bytes);
// Returns total bytes of tensors produced by "node"s output slot.
Bytes TotalBytes(const Node* node, int output_slot) const;
// Returns a prediction for the size of the tensor at the
// output_slot produced by one execution of "node".
Bytes SizeEstimate(const Node* node, int output_slot) const;
// Records that Executions of "node" have taken "time" microseconds.
void RecordTime(const Node* node, Microseconds time);
// Returns the total execution time for "node".
Microseconds TotalTime(const Node* node) const;
// Returns a prediction for one execution of "node".
Microseconds TimeEstimate(const Node* node) const;
// Check that an estimate is available for every OP node in graph.
void CheckInitialized(const Graph& graph) const;
// Records the maximum size in bytes and optionally the corresponding shape of
// the tensor generated by "output_slot" of "node". If
void RecordMaxMemorySize(const Node* node, int output_slot, Bytes bytes,
const TensorShapeProto& tensor_shape,
const DataType& dtype);
// Returns the maximum size in bytes of the tensor generated by "output_slot"
// of "node".
Bytes MaxMemorySize(const Node* node, int output_slot) const;
// Returns the shape corresponding to the largest memory size of the tensor
// generated by "output_slot" of "node".
const TensorShapeProto& MaxMemoryShape(const Node* node,
int output_slot) const;
// Returns the shape corresponding to the largest memory size of the tensor
// generated by "output_slot" of "node".
DataType MaxMemoryType(const Node* node, int output_slot) const;
// Returns the size in bytes of temporary memory consumed by "node".
Bytes TempMemorySize(const Node* node) const;
// Returns the size of persistent memory allocated by "node".
Bytes PersistentMemorySize(const Node* node) const;
// Records memory stats such as temp momory and persistent memory.
void RecordMemoryStats(const Node* node, const MemoryStats& memory_stats);
// Records the maximum execution time (in microseconds) of "node".
void RecordMaxExecutionTime(const Node* node, Microseconds time);
// Returns the maximum execution time (in microseconds) of "node".
Microseconds MaxExecutionTime(const Node* node) const;
// Record the unique id of the tensor generated by "output_slot" of "node".
// Any other tensor sharing the same id will be an alias, i.e. it will share
// the same underlying memory storage area.
void RecordAllocationId(const Node* node, int output_slot, int64_t alloc_id);
// Return the unique id of the tensor generated by "output_slot" of "node".
int64_t AllocationId(const Node* node, int output_slot) const;
bool IsPersistentTensor(const Node* node, int64_t alloc_id) const;
// Helper routines to encapsulate static estimation heuristics
// Compute an estimate of the time to copy "b" bytes over the network,
// given a fixed cost of "network_latency_millis" milliseconds and
// an estimated bandwidth of "estimated_gbps" gigabits per second (note that
// this value is in gigabits, not gigabytes).
static Microseconds CopyTimeEstimate(Bytes b, double network_latency_millis,
double estimated_gbps);
static Microseconds ComputationTimeEstimate(int64_t mathops);
// Add this CostModel into the CostGraphDef.
void AddToCostGraphDef(const Graph* graph, CostGraphDef* cost_graph) const;
// Write the contents of the CostModel to the INFO log.
void WriteSummaryToLog() const;
// Increment the times that the cost model is updated.
void IncrementUpdateTimes();
// Get the times that the cost model is updated.
int32_t GetUpdateTimes() const;
private:
static Bytes MinTensorMemoryUsage(const TensorShapeProto& tensor_shape,
const DataType& dtype);
const bool is_global_;
// Resizes vectors so that they are large enough for "id" and id's outputs.
void Ensure(int id, int num_outputs);
// Nodes and Edges whose count is < this value
// get type/byte estimates of 0.
int32_t min_count_ = 0;
// The number of times the cost model is updated.
int32_t update_times_ = 0;
// Number of times each Node has been executed.
std::vector<int32_t> count_;
// Cumulative execution time.
std::vector<Microseconds> time_;
// Cumulative Bytes output on each channel.
std::vector<absl::InlinedVector<Bytes, 2UL>> slot_bytes_;
// Maximum execution time
std::vector<Microseconds> max_exec_time_;
// Maximum memory usage
struct MemUsage {
MemUsage() : temp_memory_size(0), persistent_memory_size(0) {}
// TODO(yuefengz): temp_memory_size is not being used, remove it.
Bytes temp_memory_size;
Bytes persistent_memory_size;
absl::InlinedVector<Bytes, 2UL> output_port_mem;
absl::InlinedVector<TensorShapeProto, 2UL> output_port_shape;
absl::InlinedVector<DataType, 2UL> output_port_type;
};
std::vector<MemUsage> max_mem_usage_;
std::vector<absl::InlinedVector<int64_t, 2UL>> output_port_alloc_ids_;
std::set<int64_t> persistent_alloc_ids_;
TensorShapeProto unknown_shape_;
CostModel(const CostModel&) = delete;
void operator=(const CostModel&) = delete;
};
} // namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPH_COSTMODEL_H_
tensorflow/core/graph/costmodel_test.cc
+665
View File
@@ -0,0 +1,665 @@
/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#include "tensorflow/core/graph/costmodel.h"
#include <memory>
#include <string>
#include <unordered_map>
#include "tensorflow/cc/framework/scope.h"
#include "tensorflow/core/common_runtime/costmodel_manager.h"
#include "tensorflow/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/common_runtime/step_stats_collector.h"
#include "tensorflow/core/framework/allocation_description.pb.h"
#include "tensorflow/core/framework/cost_graph.pb.h"
#include "tensorflow/core/framework/node_def_builder.h"
#include "tensorflow/core/framework/step_stats.pb.h"
#include "tensorflow/core/framework/tensor_description.pb.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/graph/types.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/protobuf.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/util/dump_graph.h"
namespace tensorflow {
namespace {
using ::testing::Not;
// Work-around for lack of ShapeProtoEquals in OSS.
MATCHER_P(ShapeProtoEquals, other, "") {
if (arg.unknown_rank()) {
return other.unknown_rank();
}
if (arg.dim_size() != other.dim_size()) {
return false;
}
for (int i = 0; i < arg.dim_size(); ++i) {
if (arg.dim(i).size() != other.dim(i).size()) {
return false;
}
}
return true;
}
static void InitGraph(const std::string& s, Graph* graph) {
GraphDef graph_def;
auto parser = protobuf::TextFormat::Parser();
CHECK(parser.MergeFromString(s, &graph_def)) << s;
GraphConstructorOptions opts;
TF_CHECK_OK(ConvertGraphDefToGraph(opts, graph_def, graph));
}
static void InitModelFromGraph(const Graph& graph, CostModel& cm) {
// This adjusts the model to include all of the graph's nodes.
// Unlike CostModel::InitFromGraph(), this method does not add
// default estimates for sizes or times.
for (const auto& node : graph.nodes()) {
cm.SetNumOutputs(node, node->num_outputs());
}
}
// Creates a graph with two multiply nodes.
static std::unique_ptr<Graph> CreateBasicTestGraph() {
auto graph = std::make_unique<Graph>(OpRegistry::Global());
InitGraph(
"node { name: 'A' op: 'Input'}"
"node { name: 'B' op: 'Input'}"
"node { name: 'C' op: 'Mul' attr { key: 'T' value { type: DT_FLOAT } }"
" input: ['A', 'B'] }"
"node { name: 'D' op: 'Mul' attr { key: 'T' value { type: DT_FLOAT } }"
" input: ['A', 'B'] }",
graph.get());
return graph;
}
Node* FindNode(const Graph& graph, std::string name) {
for (const auto& node : graph.nodes()) {
if (node->name() == name) {
return node;
}
}
return nullptr;
}
Node* AddNode(Graph& graph, const std::string& name,
const std::string& node_type, int num_inputs) {
auto builder = NodeDefBuilder(name, node_type);
for (int i = 0; i < num_inputs; ++i) {
builder = builder.Input(absl::StrCat("node_", i), i, DT_FLOAT);
}
NodeDef node_def;
TF_CHECK_OK(builder.Finalize(&node_def));
absl::Status s;
Node* node = graph.AddNode(node_def, &s);
TF_CHECK_OK(s);
return node;
}
static void GenerateStepStats(Graph* graph, StepStats* step_stats,
const std::string& device_name) {
// Fill RunMetadata's step_stats and partition_graphs fields.
DeviceStepStats* device_stepstats = step_stats->add_dev_stats();
device_stepstats->set_device(device_name);
for (const auto& node_def : graph->nodes()) {
NodeExecStats* node_stats = device_stepstats->add_node_stats();
node_stats->set_node_name(node_def->name());
}
}
REGISTER_OP("Input").Output("o: float").SetIsStateful();
TEST(CostModelTest, WorksWithManager) {
Scope scope = Scope::NewRootScope().ExitOnError();
auto graph1 = std::make_unique<Graph>(OpRegistry::Global());
auto graph2 = std::make_unique<Graph>(OpRegistry::Global());
InitGraph(
"node { name: 'A1' op: 'Input'}"
"node { name: 'B1' op: 'Input'}"
"node { name: 'C1' op: 'Mul' attr { key: 'T' value { type: DT_FLOAT } }"
" input: ['A1', 'B1'] }"
"node { name: 'D1' op: 'Mul' attr { key: 'T' value { type: DT_FLOAT } }"
" input: ['A1', 'B1'] }",
graph1.get());
InitGraph(
"node { name: 'A2' op: 'Input'}"
"node { name: 'B2' op: 'Input'}"
"node { name: 'C2' op: 'Mul' attr { key: 'T' value { type: DT_FLOAT } }"
" input: ['A2', 'B2'] }"
"node { name: 'D2' op: 'Mul' attr { key: 'T' value { type: DT_FLOAT } }"
" input: ['A2', 'B2'] }",
graph2.get());
StepStats step_stats;
GenerateStepStats(graph1.get(), &step_stats, "DummyDevice1");
GenerateStepStats(graph2.get(), &step_stats, "DummyDevice2");
StepStatsCollector collector(&step_stats);
std::unordered_map<std::string, const Graph*> device_map;
device_map["DummyDevice1"] = graph1.get();
device_map["DummyDevice2"] = graph2.get();
CostModelManager cost_model_manager;
collector.BuildCostModel(&cost_model_manager, device_map);
CostGraphDef cost_graph_def;
TF_ASSERT_OK(
cost_model_manager.AddToCostGraphDef(graph1.get(), &cost_graph_def));
TF_ASSERT_OK(
cost_model_manager.AddToCostGraphDef(graph2.get(), &cost_graph_def));
ASSERT_EQ(cost_graph_def.node_size(), 12);
absl::flat_hash_map<int32_t, const CostGraphDef::Node> ids;
for (auto node : cost_graph_def.node()) {
int32_t index = node.id();
auto result = ids.insert({index, node});
EXPECT_TRUE(result.second);
}
}
TEST(CostModelTest, GlobalId) {
auto graph = CreateBasicTestGraph();
CostModel cm_local(/*is_global=*/false);
CostModel cm_global(/*is_global=*/true);
constexpr int kOffset = 7;
for (const auto& node : graph->nodes()) {
// Local cost models use the local id and offset.
EXPECT_EQ(cm_local.GlobalId(node, kOffset), node->id() + kOffset);
// Global cost modesl use the cost id.
EXPECT_EQ(cm_global.GlobalId(node, kOffset), node->cost_id());
}
}
TEST(CostModelTest, RecordTime) {
auto graph = CreateBasicTestGraph();
CostModel cm(/*is_global=*/false);
InitModelFromGraph(*graph, cm);
constexpr int kIters = 100;
constexpr int kMicrosPerIter = 1000;
for (int i = 0; i < kIters; ++i) {
for (const auto& node : graph->op_nodes()) {
cm.RecordTime(node, node->id() * Microseconds(kMicrosPerIter));
}
}
for (const auto& node : graph->op_nodes()) {
EXPECT_EQ(cm.TotalTime(node),
Microseconds(node->id() * kIters * kMicrosPerIter));
}
// Total time for unrecorded node is 0.
Node* E = AddNode(*graph, "E", "Mul", 2);
EXPECT_EQ(cm.TotalTime(E), Microseconds(0));
}
TEST(CostModelTest, RecordCount) {
auto graph = CreateBasicTestGraph();
CostModel cm(/*is_global=*/false);
InitModelFromGraph(*graph, cm);
constexpr int kIters = 100;
constexpr int kCountPerIter = 4;
for (int i = 0; i < kIters; ++i) {
for (const auto& node : graph->op_nodes()) {
cm.RecordCount(node, node->id() * kCountPerIter);
}
}
for (const auto& node : graph->op_nodes()) {
EXPECT_EQ(cm.TotalCount(node), node->id() * kIters * kCountPerIter);
}
// Total count for unrecorded node is 0.
Node* E = AddNode(*graph, "E", "Mul", 2);
EXPECT_EQ(cm.TotalCount(E), 0);
}
TEST(CostModelTest, RecordSize) {
auto graph = CreateBasicTestGraph();
CostModel cm(/*is_global=*/false);
InitModelFromGraph(*graph, cm);
constexpr int kIters = 100;
constexpr int kBytesPerIter = 4;
for (int i = 0; i < kIters; ++i) {
for (const auto& node : graph->op_nodes()) {
for (int slot = 0; slot < node->num_outputs(); ++slot) {
cm.RecordSize(node, slot, Bytes((node->id() + slot) * kBytesPerIter));
}
}
}
for (const auto& node : graph->op_nodes()) {
for (int slot = 0; slot < node->num_outputs(); ++slot) {
EXPECT_EQ(cm.TotalBytes(node, slot),
Bytes((node->id() + slot) * kIters * kBytesPerIter));
}
}
// Total size for unrecorded node is 0.
Node* E = AddNode(*graph, "E", "Mul", 2);
EXPECT_EQ(cm.TotalBytes(E, 0), Bytes(0));
}
TEST(CostModelTest, SizeEstimate) {
auto graph = CreateBasicTestGraph();
CostModel cm(/*is_global=*/false);
InitModelFromGraph(*graph, cm);
Node* C = FindNode(*graph, "C");
// Size estimate should be total bytes / total count.
constexpr int kBytesPerCount = 31;
constexpr int kCount = 17;
cm.RecordCount(C, kCount);
cm.RecordSize(C, 0, Bytes(kCount * kBytesPerCount));
EXPECT_EQ(cm.SizeEstimate(C, 0), Bytes(kBytesPerCount));
}
TEST(CostModelTest, TimeEstimate) {
auto graph = CreateBasicTestGraph();
CostModel cm(/*is_global=*/false);
InitModelFromGraph(*graph, cm);
Node* C = FindNode(*graph, "C");
// Time estimate should be total time / total count.
constexpr int kMicrosPerCount = 31;
constexpr int kCount = 17;
cm.RecordCount(C, kCount);
cm.RecordTime(C, Microseconds(kCount * kMicrosPerCount));
EXPECT_EQ(cm.TimeEstimate(C), Microseconds(kMicrosPerCount));
}
TensorShapeProto CreateTensorShapeProto(absl::Span<const int64_t> dims) {
TensorShapeProto shape;
for (int i = 0; i < dims.size(); ++i) {
shape.add_dim()->set_size(dims[i]);
}
return shape;
}
int64_t Count(const TensorShapeProto& shape) {
int64_t count = 1;
for (int i = 0; i < shape.dim_size(); ++i) {
count *= shape.dim(i).size();
}
return count;
}
TEST(CostModelTest, RecordMaxMemorySize) {
auto graph = CreateBasicTestGraph();
CostModel cm(/*is_global=*/false);
Node* C = FindNode(*graph, "C");
InitModelFromGraph(*graph, cm);
EXPECT_EQ(cm.MaxMemorySize(C, 0), Bytes(-1));
{
const TensorShapeProto shape = CreateTensorShapeProto({2, 5, 10});
const DataType dtype = DataType::DT_FLOAT;
const Bytes bytes = Bytes(Count(shape) * sizeof(float));
cm.RecordMaxMemorySize(C, 0, bytes, shape, dtype);
EXPECT_EQ(cm.MaxMemorySize(C, 0), bytes);
EXPECT_EQ(cm.MaxMemoryType(C, 0), dtype);
EXPECT_THAT(cm.MaxMemoryShape(C, 0), ShapeProtoEquals(shape));
}
// Records higher memory value.
{
const TensorShapeProto shape = CreateTensorShapeProto({3, 6, 11});
const DataType dtype = DataType::DT_DOUBLE;
const Bytes bytes = Bytes(Count(shape) * sizeof(double));
cm.RecordMaxMemorySize(C, 0, bytes, shape, dtype);
EXPECT_EQ(cm.MaxMemorySize(C, 0), bytes);
EXPECT_EQ(cm.MaxMemoryType(C, 0), dtype);
EXPECT_THAT(cm.MaxMemoryShape(C, 0), ShapeProtoEquals(shape));
}
// Lower memory value ignored.
{
const TensorShapeProto shape = CreateTensorShapeProto({1, 1, 1});
const DataType dtype = DataType::DT_BFLOAT16;
const Bytes bytes = Bytes(Count(shape) * sizeof(double));
cm.RecordMaxMemorySize(C, 0, bytes, shape, dtype);
EXPECT_GT(cm.MaxMemorySize(C, 0), bytes);
EXPECT_NE(cm.MaxMemoryType(C, 0), dtype);
EXPECT_THAT(cm.MaxMemoryShape(C, 0), Not(ShapeProtoEquals(shape)));
}
// Bytes computed from shape/dtype.
{
const TensorShapeProto shape = CreateTensorShapeProto({100, 100, 100});
const DataType dtype = DataType::DT_BFLOAT16;
cm.RecordMaxMemorySize(C, 0, Bytes(-1), shape, dtype);
EXPECT_EQ(cm.MaxMemorySize(C, 0), Bytes(Count(shape) * sizeof(bfloat16)));
EXPECT_EQ(cm.MaxMemoryType(C, 0), dtype);
EXPECT_THAT(cm.MaxMemoryShape(C, 0), ShapeProtoEquals(shape));
}
// Max memory size for unrecorded node is 0.
Node* E = AddNode(*graph, "E", "Mul", 2);
EXPECT_EQ(cm.MaxMemorySize(E, 0), Bytes(0));
EXPECT_THAT(cm.MaxMemoryType(E, 0), DataType::DT_INVALID);
TensorShapeProto unknown;
unknown.set_unknown_rank(true);
EXPECT_THAT(cm.MaxMemoryShape(E, 0), ShapeProtoEquals(unknown));
}
TEST(CostModelTest, RecordMaxExecutionTime) {
auto graph = CreateBasicTestGraph();
CostModel cm(/*is_global=*/false);
InitModelFromGraph(*graph, cm);
Node* C = FindNode(*graph, "C");
EXPECT_EQ(cm.MaxExecutionTime(C), Microseconds(0));
cm.RecordMaxExecutionTime(C, Microseconds(13));
EXPECT_EQ(cm.MaxExecutionTime(C), Microseconds(13));
cm.RecordMaxExecutionTime(C, Microseconds(27));
EXPECT_EQ(cm.MaxExecutionTime(C), Microseconds(27));
cm.RecordMaxExecutionTime(C, Microseconds(9));
EXPECT_EQ(cm.MaxExecutionTime(C), Microseconds(27));
// Max execution time for unrecorded node is 0.
Node* E = AddNode(*graph, "E", "Mul", 2);
EXPECT_EQ(cm.MaxExecutionTime(E), Microseconds(0));
}
TEST(CostModelTest, RecordMemoryStats) {
auto graph = CreateBasicTestGraph();
CostModel cm(/*is_global=*/false);
InitModelFromGraph(*graph, cm);
Node* C = FindNode(*graph, "C");
MemoryStats stats;
stats.set_temp_memory_size(256);
stats.set_persistent_memory_size(16);
stats.add_persistent_tensor_alloc_ids(1);
stats.add_persistent_tensor_alloc_ids(3);
stats.add_persistent_tensor_alloc_ids(5);
stats.add_persistent_tensor_alloc_ids(5); // Intentional duplicate.
cm.RecordMemoryStats(C, stats);
EXPECT_EQ(cm.TempMemorySize(C), stats.temp_memory_size());
EXPECT_EQ(cm.PersistentMemorySize(C), stats.persistent_memory_size());
EXPECT_TRUE(cm.IsPersistentTensor(C, 1));
EXPECT_TRUE(cm.IsPersistentTensor(C, 3));
EXPECT_TRUE(cm.IsPersistentTensor(C, 5));
EXPECT_FALSE(cm.IsPersistentTensor(C, 31));
// Info for unrecorded node is 0.
Node* E = AddNode(*graph, "E", "Mul", 2);
EXPECT_EQ(cm.TempMemorySize(E), Bytes(0));
EXPECT_EQ(cm.PersistentMemorySize(E), Bytes(0));
}
TEST(CostModelTest, RecordAllocationId) {
auto graph = CreateBasicTestGraph();
CostModel cm(/*is_global=*/false);
InitModelFromGraph(*graph, cm);
Node* C = FindNode(*graph, "C");
cm.RecordAllocationId(C, /*output_slot=*/0, /*alloc_id=*/13);
EXPECT_EQ(cm.AllocationId(C, /*output_slot=*/0), 13);
// Invalid slot returns -1.
EXPECT_EQ(cm.AllocationId(C, /*output_slot=*/7), -1);
// Unrecorded node returns -1.
Node* E = AddNode(*graph, "E", "Mul", 2);
EXPECT_EQ(cm.AllocationId(E, /*output_slot=*/0), -1);
}
TEST(CostModelTest, CopyTimeEstimate) {
// Current estimate is a linear model bytes / rate + latency.
int64_t bytes = 32568;
double latency_ms = 10.2;
double gbps = 2.2;
double bytes_per_usec = gbps * 1000 / 8;
double cost_usecs = (bytes / bytes_per_usec + latency_ms * 1000);
EXPECT_EQ(CostModel::CopyTimeEstimate(Bytes(bytes), latency_ms, gbps),
Microseconds(static_cast<uint64_t>(cost_usecs)));
}
TEST(CostModelTest, ComputationTimeEstimate) {
// Current estimate is 1000 math ops per microsecond.
constexpr int64_t kNumMathOps = 32150;
EXPECT_EQ(CostModel::ComputationTimeEstimate(kNumMathOps),
Microseconds(kNumMathOps / 1000));
}
TEST(CostModel, UpdateTimes) {
CostModel cm(/*is_global=*/false);
EXPECT_EQ(cm.GetUpdateTimes(), 0);
constexpr int kNumUpdates = 111;
for (int i = 0; i < kNumUpdates; ++i) {
cm.IncrementUpdateTimes();
}
EXPECT_EQ(cm.GetUpdateTimes(), kNumUpdates);
}
TEST(CostModel, SuppressInfrequent) {
// Infrequent count is used in the size and time estimates.
CostModel cm(/*is_global=*/false);
auto graph = std::make_unique<Graph>(OpRegistry::Global());
Node* A = AddNode(*graph, "A", "Mul", 2);
Node* B = AddNode(*graph, "B", "Mul", 2);
Node* C = AddNode(*graph, "B", "Mul", 2);
InitModelFromGraph(*graph, cm);
// A and B are frequent, C is not.
cm.RecordCount(A, 1000);
cm.RecordSize(A, 0, Bytes(8 * 1000));
cm.RecordTime(A, Microseconds(8 * 1000));
cm.RecordCount(B, 2000);
cm.RecordSize(B, 0, Bytes(2000 * 10));
cm.RecordTime(B, Microseconds(2000 * 10));
cm.RecordCount(C, 17);
cm.RecordSize(C, 0, Bytes(32 * 17));
cm.RecordTime(C, Microseconds(32 * 17));
// Estimate size and time without suppression.
EXPECT_EQ(cm.SizeEstimate(A, 0), Bytes(8));
EXPECT_EQ(cm.TimeEstimate(A), Microseconds(8));
EXPECT_EQ(cm.SizeEstimate(B, 0), Bytes(10));
EXPECT_EQ(cm.TimeEstimate(B), Microseconds(10));
EXPECT_EQ(cm.SizeEstimate(C, 0), Bytes(32));
EXPECT_EQ(cm.TimeEstimate(C), Microseconds(32));
cm.SuppressInfrequent();
// Sizes and times suppressed for C but not A, B.
EXPECT_EQ(cm.SizeEstimate(A, 0), Bytes(8));
EXPECT_EQ(cm.TimeEstimate(A), Microseconds(8));
EXPECT_EQ(cm.SizeEstimate(B, 0), Bytes(10));
EXPECT_EQ(cm.TimeEstimate(B), Microseconds(10));
EXPECT_EQ(cm.SizeEstimate(C, 0), Bytes(0));
EXPECT_EQ(cm.TimeEstimate(C), Microseconds(1)); // kMinTimeEstimate.
}
TEST(CostModelTest, MergeFromLocal) {
CostModel cm_global(/*is_global=*/true);
CostModel cm_local(/*is_global=*/false);
auto graph = CreateBasicTestGraph();
InitModelFromGraph(*graph, cm_global);
// Populate global model.
Node* C = FindNode(*graph, "C");
Node* D = FindNode(*graph, "D");
cm_global.RecordCount(C, 23);
cm_global.RecordSize(C, 0, Bytes(23));
cm_global.RecordTime(C, Microseconds(123));
cm_global.RecordCount(D, 17);
cm_global.RecordSize(D, 0, Bytes(17));
cm_global.RecordTime(D, Microseconds(117));
// Add new nodes and add cost to a local model.
Node* E = AddNode(*graph, "E", "Mul", 2);
graph->AddEdge(C, 0, E, 0);
graph->AddEdge(D, 0, E, 1);
Node* F = AddNode(*graph, "F", "Mul", 2);
graph->AddEdge(E, 0, F, 0);
graph->AddEdge(D, 0, F, 1);
InitModelFromGraph(*graph, cm_local);
cm_local.RecordCount(E, 37);
cm_local.RecordSize(E, 0, Bytes(37));
cm_local.RecordTime(E, Microseconds(137));
cm_local.RecordCount(F, 41);
cm_local.RecordSize(F, 0, Bytes(41));
cm_local.RecordTime(F, Microseconds(141));
// Add existing node to check stats are added.
cm_local.RecordCount(C, 1);
cm_local.RecordSize(C, 0, Bytes(1));
cm_local.RecordTime(C, Microseconds(100));
// Merge and check that stats from local are now in global.
cm_global.MergeFromLocal(*graph, cm_local);
EXPECT_EQ(cm_global.TotalCount(E), cm_local.TotalCount(E));
EXPECT_EQ(cm_global.TotalBytes(E, 0), cm_local.TotalBytes(E, 0));
EXPECT_EQ(cm_global.TotalTime(E), cm_local.TotalTime(E));
EXPECT_EQ(cm_global.TotalCount(F), cm_local.TotalCount(F));
EXPECT_EQ(cm_global.TotalBytes(F, 0), cm_local.TotalBytes(F, 0));
EXPECT_EQ(cm_global.TotalTime(F), cm_local.TotalTime(F));
// Stats for C are added.
EXPECT_EQ(cm_global.TotalCount(C), Microseconds(24));
EXPECT_EQ(cm_global.TotalBytes(C, 0), Bytes(24));
EXPECT_EQ(cm_global.TotalTime(C), Microseconds(223));
}
TEST(CostModelTest, MergeFromGlobal) {
CostModel cm1(/*is_global=*/true);
CostModel cm2(/*is_global=*/true);
auto graph = CreateBasicTestGraph();
InitModelFromGraph(*graph, cm1);
// Populate global model.
Node* C = FindNode(*graph, "C");
Node* D = FindNode(*graph, "D");
cm1.RecordCount(C, 23);
cm1.RecordSize(C, 0, Bytes(23));
cm1.RecordTime(C, Microseconds(123));
cm1.RecordCount(D, 17);
cm1.RecordSize(D, 0, Bytes(17));
cm1.RecordTime(D, Microseconds(117));
// Add new nodes and add cost to a local model.
Node* E = AddNode(*graph, "E", "Mul", 2);
graph->AddEdge(C, 0, E, 0);
graph->AddEdge(D, 0, E, 1);
Node* F = AddNode(*graph, "F", "Mul", 2);
graph->AddEdge(E, 0, F, 0);
graph->AddEdge(D, 0, F, 1);
InitModelFromGraph(*graph, cm2);
cm2.RecordCount(E, 37);
cm2.RecordSize(E, 0, Bytes(37));
cm2.RecordTime(E, Microseconds(137));
cm2.RecordCount(F, 41);
cm2.RecordSize(F, 0, Bytes(41));
cm2.RecordTime(F, Microseconds(141));
// Add existing node to check stats are added.
cm2.RecordCount(C, 1);
cm2.RecordSize(C, 0, Bytes(1));
cm2.RecordTime(C, Microseconds(100));
// Merge and check that stats are merged.
cm1.MergeFromGlobal(cm2);
EXPECT_EQ(cm1.TotalCount(E), cm2.TotalCount(E));
EXPECT_EQ(cm1.TotalBytes(E, 0), cm2.TotalBytes(E, 0));
EXPECT_EQ(cm1.TotalTime(E), cm2.TotalTime(E));
EXPECT_EQ(cm1.TotalCount(F), cm2.TotalCount(F));
EXPECT_EQ(cm1.TotalBytes(F, 0), cm2.TotalBytes(F, 0));
EXPECT_EQ(cm1.TotalTime(F), cm2.TotalTime(F));
// Stats for C are added.
EXPECT_EQ(cm1.TotalCount(C), Microseconds(24));
EXPECT_EQ(cm1.TotalBytes(C, 0), Bytes(24));
EXPECT_EQ(cm1.TotalTime(C), Microseconds(223));
}
NodeExecStats CreateNodeExecStats(const Node* node, int64_t time,
int64_t bytes) {
NodeExecStats stats;
stats.set_node_name(node->name());
stats.set_op_start_rel_micros(10);
stats.set_op_end_rel_micros(10 + time);
for (int i = 0; i < node->num_outputs(); ++i) {
NodeOutput* no = stats.add_output();
no->set_slot(i);
no->mutable_tensor_description()
->mutable_allocation_description()
->set_requested_bytes(bytes);
}
return stats;
}
TEST(CostModelTest, MergeFromStats) {
CostModel cm(/*is_global=*/true);
auto graph = CreateBasicTestGraph();
InitModelFromGraph(*graph, cm);
// Populate global model.
Node* C = FindNode(*graph, "C");
Node* D = FindNode(*graph, "D");
cm.RecordCount(C, 23);
cm.RecordTime(C, Microseconds(123));
cm.RecordCount(D, 17);
cm.RecordTime(D, Microseconds(117));
// Add new nodes and create stats.
Node* E = AddNode(*graph, "E", "Mul", 2);
graph->AddEdge(C, 0, E, 0);
graph->AddEdge(D, 0, E, 1);
Node* F = AddNode(*graph, "F", "Mul", 2);
graph->AddEdge(E, 0, F, 0);
graph->AddEdge(D, 0, F, 1);
StepStats stats;
DeviceStepStats* dstats = stats.add_dev_stats();
*(dstats->add_node_stats()) = CreateNodeExecStats(C, 10, 10);
*(dstats->add_node_stats()) = CreateNodeExecStats(D, 10, 10);
*(dstats->add_node_stats()) = CreateNodeExecStats(E, 20, 20);
*(dstats->add_node_stats()) = CreateNodeExecStats(E, 20, 20);
*(dstats->add_node_stats()) = CreateNodeExecStats(F, 30, 30);
*(dstats->add_node_stats()) = CreateNodeExecStats(F, 30, 30);
NodeNameToCostIdMap id_map;
for (const auto& node : graph->nodes()) {
id_map.emplace(node->name(), node->cost_id());
}
cm.MergeFromStats(id_map, stats);
// Stats for C/D are added to existing.
EXPECT_EQ(cm.TotalCount(C), 24);
EXPECT_EQ(cm.TotalTime(C), Microseconds(133));
EXPECT_EQ(cm.TotalBytes(C, 0), Bytes(10));
EXPECT_EQ(cm.TotalCount(D), 18);
EXPECT_EQ(cm.TotalTime(D), Microseconds(127));
EXPECT_EQ(cm.TotalBytes(D, 0), Bytes(10));
// Stats for E/F are accumulated.
EXPECT_EQ(cm.TotalCount(E), 2);
EXPECT_EQ(cm.TotalTime(E), Microseconds(40));
EXPECT_EQ(cm.TotalBytes(E, 0), Bytes(40));
EXPECT_EQ(cm.TotalCount(F), 2);
EXPECT_EQ(cm.TotalTime(F), Microseconds(60));
EXPECT_EQ(cm.TotalBytes(F, 0), Bytes(60));
}
} // namespace
} // namespace tensorflow
tensorflow/core/graph/default_device.h
+41
View File
@@ -0,0 +1,41 @@
/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
Licensed 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 TENSORFLOW_CORE_GRAPH_DEFAULT_DEVICE_H_
#define TENSORFLOW_CORE_GRAPH_DEFAULT_DEVICE_H_
#include <string>
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
namespace tensorflow {
namespace graph {
// Sets the default device for all nodes in graph_def to "device",
// only if not already set.
inline void SetDefaultDevice(const std::string& device, GraphDef* graph_def) {
for (int i = 0; i < graph_def->node_size(); ++i) {
auto node = graph_def->mutable_node(i);
if (node->device().empty()) {
node->set_device(device);
}
}
}
} // namespace graph
} // namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPH_DEFAULT_DEVICE_H_
tensorflow/core/graph/edgeset.cc
+70
View File
@@ -0,0 +1,70 @@
/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#include "tensorflow/core/graph/edgeset.h"
namespace tensorflow {
std::pair<EdgeSet::const_iterator, bool> EdgeSet::insert(value_type value) {
RegisterMutation();
const_iterator ci;
ci.Init(this);
auto s = get_set();
if (!s) {
for (int i = 0; i < kInline; i++) {
if (ptrs_[i] == value) {
ci.array_iter_ = &ptrs_[i];
return std::make_pair(ci, false);
}
}
for (int i = 0; i < kInline; i++) {
if (ptrs_[i] == nullptr) {
ptrs_[i] = value;
ci.array_iter_ = &ptrs_[i];
return std::make_pair(ci, true);
}
}
// array is full. convert to set.
s = new gtl::FlatSet<const Edge*>;
s->insert(reinterpret_cast<const Edge**>(std::begin(ptrs_)),
reinterpret_cast<const Edge**>(std::end(ptrs_)));
ptrs_[0] = this;
ptrs_[1] = s;
// fall through.
}
auto p = s->insert(value);
ci.tree_iter_ = p.first;
return std::make_pair(ci, p.second);
}
EdgeSet::size_type EdgeSet::erase(key_type key) {
RegisterMutation();
auto s = get_set();
if (!s) {
for (int i = 0; i < kInline; i++) {
if (ptrs_[i] == key) {
size_t n = size();
ptrs_[i] = ptrs_[n - 1];
ptrs_[n - 1] = nullptr;
return 1;
}
}
return 0;
} else {
return s->erase(key);
}
}
} // namespace tensorflow
tensorflow/core/graph/edgeset.h
+246
View File
@@ -0,0 +1,246 @@
/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
Licensed 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 TENSORFLOW_CORE_GRAPH_EDGESET_H_
#define TENSORFLOW_CORE_GRAPH_EDGESET_H_
#include <stddef.h>
#include "tensorflow/core/lib/gtl/flatset.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/macros.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
class Edge;
// An unordered set of edges. Uses very little memory for small sets.
// Unlike gtl::FlatSet, EdgeSet does NOT allow mutations during
// iteration.
class EdgeSet {
public:
EdgeSet();
~EdgeSet();
typedef const Edge* key_type;
typedef const Edge* value_type;
typedef size_t size_type;
typedef ptrdiff_t difference_type;
class const_iterator;
typedef const_iterator iterator;
bool empty() const;
size_type size() const;
void clear();
std::pair<iterator, bool> insert(value_type value);
size_type erase(key_type key);
void reserve(size_type new_size) {
if (new_size > kInline) {
auto s = new gtl::FlatSet<const Edge*>(new_size);
s->insert(reinterpret_cast<const Edge**>(std::begin(ptrs_)),
reinterpret_cast<const Edge**>(&ptrs_[0] + size()));
ptrs_[0] = this;
ptrs_[1] = s;
}
}
// Caller is not allowed to mutate the EdgeSet while iterating.
const_iterator begin() const;
const_iterator end() const;
private:
// Up to kInline elements are stored directly in ptrs_ (nullptr means none).
// If ptrs_[0] == this then ptrs_[1] points to a set<const Edge*>.
// kInline must be >= 2, and is chosen such that ptrs_ fills a 64 byte
// cacheline.
static constexpr int kInline = 64 / sizeof(const void*);
const void* ptrs_[kInline];
gtl::FlatSet<const Edge*>* get_set() const {
if (ptrs_[0] == this) {
return static_cast<gtl::FlatSet<const Edge*>*>(
const_cast<void*>(ptrs_[1]));
} else {
return nullptr;
}
}
// To detect mutations while iterating.
#ifdef NDEBUG
void RegisterMutation() {}
#else
uint32_t mutations_ = 0;
void RegisterMutation() { mutations_++; }
#endif
EdgeSet(const EdgeSet&) = delete;
void operator=(const EdgeSet&) = delete;
};
class EdgeSet::const_iterator {
public:
typedef typename EdgeSet::value_type value_type;
typedef const typename EdgeSet::value_type& reference;
typedef const typename EdgeSet::value_type* pointer;
typedef typename EdgeSet::difference_type difference_type;
typedef std::forward_iterator_tag iterator_category;
const_iterator() {}
const_iterator& operator++();
const_iterator operator++(int /*unused*/);
const value_type* operator->() const;
value_type operator*() const;
bool operator==(const const_iterator& other) const;
bool operator!=(const const_iterator& other) const {
return !(*this == other);
}
private:
friend class EdgeSet;
void const* const* array_iter_ = nullptr;
typename gtl::FlatSet<const Edge*>::const_iterator tree_iter_;
#ifdef NDEBUG
inline void Init(const EdgeSet* e) {}
inline void CheckNoMutations() const {}
#else
void Init(const EdgeSet* e) {
owner_ = e;
init_mutations_ = e->mutations_;
}
void CheckNoMutations() const {
CHECK_EQ(init_mutations_, owner_->mutations_);
}
const EdgeSet* owner_ = nullptr;
uint32_t init_mutations_ = 0;
#endif
};
inline EdgeSet::EdgeSet() {
for (int i = 0; i < kInline; i++) {
ptrs_[i] = nullptr;
}
}
inline EdgeSet::~EdgeSet() { delete get_set(); }
inline bool EdgeSet::empty() const { return size() == 0; }
inline EdgeSet::size_type EdgeSet::size() const {
auto s = get_set();
if (s) {
return s->size();
} else {
size_t result = 0;
for (int i = 0; i < kInline; i++) {
if (ptrs_[i]) result++;
}
return result;
}
}
inline void EdgeSet::clear() {
RegisterMutation();
delete get_set();
for (int i = 0; i < kInline; i++) {
ptrs_[i] = nullptr;
}
}
inline EdgeSet::const_iterator EdgeSet::begin() const {
const_iterator ci;
ci.Init(this);
auto s = get_set();
if (s) {
ci.tree_iter_ = s->begin();
} else {
ci.array_iter_ = &ptrs_[0];
}
return ci;
}
inline EdgeSet::const_iterator EdgeSet::end() const {
const_iterator ci;
ci.Init(this);
auto s = get_set();
if (s) {
ci.tree_iter_ = s->end();
} else {
ci.array_iter_ = &ptrs_[size()];
}
return ci;
}
inline EdgeSet::const_iterator& EdgeSet::const_iterator::operator++() {
CheckNoMutations();
if (array_iter_ != nullptr) {
++array_iter_;
} else {
++tree_iter_;
}
return *this;
}
inline EdgeSet::const_iterator EdgeSet::const_iterator::operator++(
int /*unused*/) {
CheckNoMutations();
const_iterator tmp = *this;
operator++();
return tmp;
}
// gcc's set and multiset always use const_iterator since it will otherwise
// allow modification of keys.
inline const EdgeSet::const_iterator::value_type* EdgeSet::const_iterator::
operator->() const {
CheckNoMutations();
if (array_iter_ != nullptr) {
return reinterpret_cast<const value_type*>(array_iter_);
} else {
return tree_iter_.operator->();
}
}
// gcc's set and multiset always use const_iterator since it will otherwise
// allow modification of keys.
inline EdgeSet::const_iterator::value_type EdgeSet::const_iterator::operator*()
const {
CheckNoMutations();
if (array_iter_ != nullptr) {
return static_cast<value_type>(*array_iter_);
} else {
return *tree_iter_;
}
}
inline bool EdgeSet::const_iterator::operator==(
const const_iterator& other) const {
DCHECK((array_iter_ == nullptr) == (other.array_iter_ == nullptr))
<< "Iterators being compared must be from same set that has not "
<< "been modified since the iterator was constructed";
CheckNoMutations();
if (array_iter_ != nullptr) {
return array_iter_ == other.array_iter_;
} else {
return other.array_iter_ == nullptr && tree_iter_ == other.tree_iter_;
}
}
} // namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPH_EDGESET_H_
tensorflow/core/graph/edgeset_test.cc
+109
View File
@@ -0,0 +1,109 @@
/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#include "tensorflow/core/graph/edgeset.h"
#include <set>
#include <vector>
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
class EdgeSetTest : public ::testing::Test {
public:
EdgeSetTest() : edges_(nullptr) {}
~EdgeSetTest() override { delete[] edges_; }
void MakeEdgeSet(int n) {
if (edges_) {
delete[] edges_;
}
edges_ = new Edge[n];
eset_.clear();
model_.clear();
for (int i = 0; i < n; i++) {
eset_.insert(&edges_[i]);
model_.insert(&edges_[i]);
}
}
void CheckSame() {
EXPECT_EQ(model_.size(), eset_.size());
EXPECT_EQ(model_.empty(), eset_.empty());
std::vector<const Edge*> modelv(model_.begin(), model_.end());
std::vector<const Edge*> esetv(eset_.begin(), eset_.end());
std::sort(modelv.begin(), modelv.end());
std::sort(esetv.begin(), esetv.end());
EXPECT_EQ(modelv.size(), esetv.size());
for (size_t i = 0; i < modelv.size(); i++) {
EXPECT_EQ(modelv[i], esetv[i]) << i;
}
}
static constexpr int kInline = 64 / sizeof(const void*);
Edge nonexistent_;
Edge* edges_;
EdgeSet eset_;
std::set<const Edge*> model_;
};
namespace {
TEST_F(EdgeSetTest, Ops) {
for (int n : {0, 1, 2, kInline + 1}) {
MakeEdgeSet(n);
CheckSame();
EXPECT_EQ((n == 0), eset_.empty());
EXPECT_EQ(n, eset_.size());
eset_.clear();
model_.clear();
CheckSame();
eset_.insert(&edges_[0]);
model_.insert(&edges_[0]);
CheckSame();
}
}
// Try insert/erase of existing elements at different positions.
TEST_F(EdgeSetTest, Exists) {
for (int n : {0, 1, 2, kInline + 1}) {
MakeEdgeSet(n);
for (int pos = 0; pos < n; pos++) {
auto p = eset_.insert(&edges_[pos]);
EXPECT_FALSE(p.second);
EXPECT_EQ(&edges_[pos], *p.first);
EXPECT_EQ(1, eset_.erase(&edges_[pos]));
model_.erase(&edges_[pos]);
CheckSame();
}
}
}
// Try insert/erase of non-existent element.
TEST_F(EdgeSetTest, DoesNotExist) {
for (int n : {0, 1, 2, kInline + 1}) {
MakeEdgeSet(n);
EXPECT_EQ(0, eset_.erase(&nonexistent_));
auto p = eset_.insert(&nonexistent_);
EXPECT_TRUE(p.second);
EXPECT_EQ(&nonexistent_, *p.first);
}
}
} // namespace
} // namespace tensorflow
tensorflow/core/graph/graph.cc
+1100
View File
File diff suppressed because it is too large Load Diff
tensorflow/core/graph/graph.h
+1116
View File
File diff suppressed because it is too large Load Diff
tensorflow/core/graph/graph_debug_info_builder.cc
+299
View File
@@ -0,0 +1,299 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#include "tensorflow/core/graph/graph_debug_info_builder.h"
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "absl/container/flat_hash_map.h"
#include "absl/hash/hash.h"
#include "absl/status/status.h"
#include "absl/strings/str_format.h"
#include "absl/types/span.h"
#include "tensorflow/core/config/flag_defs.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/graph_debug_info.pb.h"
#include "tensorflow/core/framework/logging.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/stack_frame.h"
#include "tsl/platform/path.h"
namespace tensorflow {
// Ignore the frames containing this substring for common prefix calculation.
static const char* kFilenameToIgnorePrefix = "<embedded";
// Converts the given stack frame to a string.
std::string StackFrameToString(const StackFrame& frame,
int shared_prefix_length) {
std::string out = absl::StrFormat(
"File \"%s\", line %d, in %s",
absl::StrContains(frame.file_name, kFilenameToIgnorePrefix)
? frame.file_name
: frame.file_name.substr(shared_prefix_length),
frame.line_number, frame.function_name);
return out;
}
std::string ToStringHelper(absl::Span<const StackFrame> stack_frames,
int shared_prefix_length) {
return absl::StrJoin(
stack_frames, "\n", [&](std::string* out, const StackFrame& frame) {
absl::StrAppend(out, StackFrameToString(frame, shared_prefix_length));
});
}
FrozenStackTrace::FrozenStackTrace(absl::Span<StackFrame const> frames,
absl::Span<StackFrame const> user_frames)
: frames_(frames.begin(), frames.end()),
user_frames_(user_frames.begin(), user_frames.end()) {
if (user_frames.empty()) {
user_frames_ = frames_;
}
}
FrozenStackTrace::FrozenStackTrace(
const GraphDebugInfo::StackTrace& stack_trace,
const GraphDebugInfo& debug_info) {
auto push_frame = [this,
&debug_info](const GraphDebugInfo::FileLineCol& frame) {
int file_index = frame.file_index();
std::string file_name =
(file_index >= 0 && file_index < debug_info.files_size())
? debug_info.files(file_index)
: "<UNKNOWN_FILE_NAME>";
frames_.push_back(StackFrame(file_name, frame.line(), frame.func()));
};
if (!stack_trace.file_line_cols().empty()) {
for (const GraphDebugInfo::FileLineCol& frame :
stack_trace.file_line_cols()) {
push_frame(frame);
}
} else {
for (const uint64_t frame_id : stack_trace.frame_id()) {
if (debug_info.frames_by_id().contains(frame_id)) {
push_frame(debug_info.frames_by_id().at(frame_id));
} else {
LOG_FIRST_N(ERROR, 5) << "No matching frame for id:" << frame_id;
}
}
}
}
absl::Span<StackFrame const> FrozenStackTrace::ToFrames() const {
return frames_;
}
std::vector<StackFrame> FrozenStackTrace::ToUncachedFrames() const {
return frames_;
}
StackFrame FrozenStackTrace::LastUserFrame() const { return frames_.back(); }
std::vector<StackFrame> FrozenStackTrace::GetUserFrames(int limit) const {
std::vector<StackFrame> result;
if (limit < 0 || limit > user_frames_.size()) {
limit = user_frames_.size();
}
result.reserve(limit);
for (int i = 0; i < limit; ++i) {
result.push_back(user_frames_[i]);
}
return result;
}
std::string FrozenStackTrace::ToString(const TracePrintingOptions& opts) const {
int shared_prefix_length = 0;
if (opts.filter_common_prefix) {
std::vector<std::string> prefix_file_names;
for (const StackFrame& frame : frames_) {
if (!absl::StrContains(frame.file_name, kFilenameToIgnorePrefix)) {
prefix_file_names.push_back(frame.file_name);
}
}
shared_prefix_length = tsl::io::CommonPathPrefix(prefix_file_names).size();
}
if (!opts.drop_internal_frames) {
return ToStringHelper(frames_, shared_prefix_length);
}
std::vector<StackFrame> non_internal_frames;
for (const StackFrame& frame : frames_) {
if (!IsInternalFrameForFilename(frame.file_name)) {
non_internal_frames.push_back(frame);
}
}
return ToStringHelper(non_internal_frames, shared_prefix_length);
}
GraphDebugInfoBuilder::GraphDebugInfoBuilder()
: debug_info_(std::make_unique<GraphDebugInfo>()) {}
void GraphDebugInfoBuilder::AccumulateStackTracesMap(
const StackTracesMap& stack_traces_map, absl::string_view key_suffix,
const GraphDebugInfoBuilder::Options& options) {
trace_to_index_.reserve(trace_to_index_.size() + stack_traces_map.size());
for (const auto& [node_name, stack_trace] : stack_traces_map) {
if (stack_trace == nullptr) continue;
std::string trace_key = absl::StrCat(node_name, key_suffix);
AccumulateStackTrace(stack_trace, trace_key, options);
}
}
void GraphDebugInfoBuilder::AccumulateStackTrace(
std::shared_ptr<AbstractStackTrace> trace, absl::string_view traces_key,
const GraphDebugInfoBuilder::Options& options) {
int trace_index = 0;
StackTracePointer p{trace};
auto found = trace_to_index_.find(p);
if (found != trace_to_index_.end()) {
trace_index = found->second;
} else {
trace_index = debug_info_->traces_by_id().size();
trace_to_index_[p] = trace_index;
GraphDebugInfo::StackTrace& stack_trace_proto =
(*debug_info_->mutable_traces_by_id())[trace_index];
if (options.user_frames) {
frame_to_index_.reserve(
frame_to_index_.size() +
trace->GetUserFrames(options.user_frames_limit).size());
for (const auto& stack_frame :
trace->GetUserFrames(options.user_frames_limit)) {
AppendToStackTraceProto(stack_frame, stack_trace_proto);
}
} else {
if (flags::Global()
.enable_graph_debug_info_caching_for_stack_frames.value()) {
frame_to_index_.reserve(frame_to_index_.size() +
trace->ToFrames().size());
for (const auto& stack_frame : trace->ToFrames()) {
AppendToStackTraceProto(stack_frame, stack_trace_proto);
}
} else {
frame_to_index_.reserve(frame_to_index_.size() +
trace->ToUncachedFrames().size());
for (const auto& stack_frame : trace->ToUncachedFrames()) {
AppendToStackTraceProto(stack_frame, stack_trace_proto);
}
}
}
}
(*debug_info_->mutable_name_to_trace_id())[traces_key] = trace_index;
}
void GraphDebugInfoBuilder::AppendToStackTraceProto(
const StackFrame& stack_frame,
GraphDebugInfo::StackTrace& stack_trace_proto) {
int frame_index = 0;
auto found = frame_to_index_.find(stack_frame);
if (found != frame_to_index_.end()) {
frame_index = found->second;
} else {
frame_index = debug_info_->frames_by_id().size();
frame_to_index_[stack_frame] = frame_index;
GraphDebugInfo::FileLineCol& frame =
(*debug_info_->mutable_frames_by_id())[frame_index];
auto file_index = file_name_to_index_.find(stack_frame.file_name);
if (file_index != file_name_to_index_.end()) {
frame.set_file_index(file_index->second);
} else {
frame.set_file_index(new_name_index_);
file_name_to_index_[stack_frame.file_name] = new_name_index_;
*debug_info_->add_files() = stack_frame.file_name;
new_name_index_++;
}
frame.set_line(stack_frame.line_number);
frame.set_func(stack_frame.function_name);
}
stack_trace_proto.add_frame_id(frame_index);
}
void GraphDebugInfoBuilder::AppendGraphDebugInfo(
absl::string_view prefix, const GraphDebugInfo& new_info) {
for (const auto& pair : new_info.name_to_trace_id()) {
auto trace = new_info.traces_by_id().at(pair.second);
auto frozen = std::make_shared<FrozenStackTrace>(trace, new_info);
std::string key =
prefix.empty() ? pair.first : absl::StrCat(pair.first, "@", prefix);
AccumulateStackTrace(frozen, key, GraphDebugInfoBuilder::Options{});
}
}
GraphDebugInfo GraphDebugInfoBuilder::Build() const { return *debug_info_; }
absl::Status GraphDebugInfoBuilder::AppendGraphDebugInfoStr(
absl::string_view prefix, absl::string_view new_info_str) {
GraphDebugInfo debug_info;
if (!debug_info.ParseFromString(new_info_str)) {
return absl::InvalidArgumentError("Failed to parse GraphDebugInfo proto.");
}
AppendGraphDebugInfo(prefix, debug_info);
return absl::OkStatus();
}
std::string GraphDebugInfoBuilder::ToGraphDebugInfoStr() const {
return Build().SerializeAsString();
}
StackTracesMap LoadTracesFromDebugInfo(const GraphDebugInfo& debug_info) {
StackTracesMap traces;
absl::flat_hash_map<uint64_t, std::shared_ptr<AbstractStackTrace>>
traces_by_id;
traces_by_id.reserve(debug_info.traces_by_id_size());
for (const auto& [id, frames] : debug_info.traces_by_id()) {
traces_by_id[id] = std::make_shared<FrozenStackTrace>(frames, debug_info);
}
traces.reserve(debug_info.name_to_trace_id_size() + debug_info.traces_size());
for (const auto& [name, trace_id] : debug_info.name_to_trace_id()) {
if (!traces_by_id.contains(trace_id)) {
LOG_FIRST_N(ERROR, 5) << "No matching trace for id:" << trace_id;
continue;
}
traces[name] = traces_by_id[trace_id];
}
for (const auto& [name, frames] : debug_info.traces()) {
traces[name] = std::make_shared<FrozenStackTrace>(frames, debug_info);
}
return traces;
}
absl::StatusOr<StackTracesMap> LoadTracesFromDebugInfoStr(
absl::string_view debug_info_str) {
GraphDebugInfo debug_info;
if (!debug_info.ParseFromString(debug_info_str)) {
return absl::InvalidArgumentError("Failed to parse GraphDebugInfo proto.");
}
return LoadTracesFromDebugInfo(debug_info);
}
GraphDebugInfo StackTracesMapToGraphDebugInfo(const StackTracesMap& map,
bool user_frames) {
GraphDebugInfoBuilder builder;
GraphDebugInfoBuilder::Options options;
options.user_frames = user_frames;
options.user_frames_limit = -1;
builder.AccumulateStackTracesMap(map, "", options);
return builder.Build();
}
} // namespace tensorflow
tensorflow/core/graph/graph_debug_info_builder.h
+210
View File
@@ -0,0 +1,210 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed 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 TENSORFLOW_CORE_GRAPH_GRAPH_DEBUG_INFO_BUILDER_H_
#define TENSORFLOW_CORE_GRAPH_GRAPH_DEBUG_INFO_BUILDER_H_
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "absl/container/flat_hash_map.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "absl/types/span.h"
#include "tensorflow/core/framework/graph_debug_info.pb.h"
#include "tensorflow/core/platform/stack_frame.h"
#include "tsl/platform/macros.h"
namespace tensorflow {
// Language agnostic stack traces.
class AbstractStackTrace {
public:
struct TracePrintingOptions {
// Show inline the contents of each stack line.
bool show_line_contents = false;
// Drop the common largest prefix of all filenames in stack frames.
bool filter_common_prefix = false;
// Do not show internal frames.
bool drop_internal_frames = false;
};
virtual ~AbstractStackTrace() = default;
// The returned span is alive as long as the AbstractStackTrace is alive.
virtual absl::Span<StackFrame const> ToFrames() const = 0;
// Returns the stack frames without caching any generated data.
virtual std::vector<StackFrame> ToUncachedFrames() const = 0;
// Returns the last stack frame from user code, attempting to ignore the
// framework code. Returns an empty frame if no such stack frame was found.
virtual StackFrame LastUserFrame() const = 0;
// Returns stack trace from user code (instead of op creation ones returned in
// ToFrames).
virtual std::vector<StackFrame> GetUserFrames(int limit) const = 0;
virtual std::string ToString(const TracePrintingOptions& opts) const = 0;
};
// A frozen sequence of StackFrames; an adapter for a span of StackFrames that
// conforms to the AbstractStackTrace contract.
class FrozenStackTrace : public AbstractStackTrace {
public:
// Constructs a FrozenStackTrace from a span of StackFrames by making a copy
// of each stack frame.
explicit FrozenStackTrace(absl::Span<StackFrame const> frames,
absl::Span<StackFrame const> user_frames = {});
explicit FrozenStackTrace(std::vector<StackFrame>&& frames)
: frames_(std::move(frames)), user_frames_({}) {}
FrozenStackTrace(FrozenStackTrace&&) = default;
// Constructs a FrozenStackTrace from serialized proto data.
FrozenStackTrace(const GraphDebugInfo::StackTrace& stack_trace,
const GraphDebugInfo& debug_info);
~FrozenStackTrace() override = default;
absl::Span<StackFrame const> ToFrames() const override;
std::vector<StackFrame> ToUncachedFrames() const override;
StackFrame LastUserFrame() const override;
std::vector<StackFrame> GetUserFrames(int limit) const override;
std::string ToString(const TracePrintingOptions& opts) const override;
private:
std::vector<StackFrame> frames_;
std::vector<StackFrame> user_frames_;
};
// Holder type to use `AbstractStackTrace` as a key.
struct StackTracePointer {
std::shared_ptr<AbstractStackTrace> trace;
template <class H>
friend H AbslHashValue(H h, const StackTracePointer& p) {
for (const auto& frame : p.trace->ToFrames()) {
h = H::combine(std::move(h), frame);
}
return h;
}
bool operator==(const StackTracePointer& other) const {
absl::Span<StackFrame const> other_frames = other.trace->ToFrames();
absl::Span<StackFrame const> frames = trace->ToFrames();
return frames == other_frames;
}
};
using StackTracesMap =
absl::flat_hash_map<std::string,
std::shared_ptr<tensorflow::AbstractStackTrace>>;
// Load all stack traces from `debug_info`.
StackTracesMap LoadTracesFromDebugInfo(const GraphDebugInfo& debug_info);
absl::StatusOr<StackTracesMap> LoadTracesFromDebugInfoStr(
absl::string_view debug_info_str);
// Generates a GraphDebugInfo proto from a StackTracesMap object. Returns user
// frames by default. If `user_frames` is false, returns all frames.
GraphDebugInfo StackTracesMapToGraphDebugInfo(const StackTracesMap& map,
bool user_frames = true);
// Builder for GraphDebugInfo protos from either an existing map of string keys
// to stack traces, or individual stack traces, or both. All stack traces in a
// GraphDebugInfo are stored with a string key in the `traces` field. In the
// case of an existing map, its keys are used, appended with a key suffix,
// which may be empty. If it is not empty, it is conventionally of the form
// "@function_name", although this class doesn't care. In the case of an
// individual stack trace, a key for `traces` must be provided.
//
// This builder will create a list of the unique file names across all stack
// traces and store it in the `files` field. When storing stack traces into the
// proto, file names are replaced by their index into `files`.
//
// Typical usage is to call one or both of the accumulate methods one or more
// times and then to call the Build().
class GraphDebugInfoBuilder {
public:
struct Options {
// Call the AbstractTraceMap GetUserFrames method rather than ToFrames
bool user_frames;
// Value of `limit` to pass to GetUserFrames if `user_frames` is true,
// otherwise ignored
int user_frames_limit;
};
GraphDebugInfoBuilder();
virtual ~GraphDebugInfoBuilder() = default;
// Adds a map of stack traces to the GraphDebugInfo proto. For each key (node
// id) and stack traces entry in `stack_traces_map`, combine the key with
// `key_suffix` to form a new key and use that to add the stack traces to the
// `traces` field of the proto. If not empty, the suffix is typically of the
// form "@function_name", although this function doesn't care.
void AccumulateStackTracesMap(const StackTracesMap& stack_traces_map,
absl::string_view key_suffix = "",
const GraphDebugInfoBuilder::Options& options =
GraphDebugInfoBuilder::Options());
// Adds one stack trace to the GraphDebugInfo proto, using `traces_key` as the
// key for the `traces` field of the proto.
void AccumulateStackTrace(std::shared_ptr<AbstractStackTrace> trace,
absl::string_view traces_key,
const GraphDebugInfoBuilder::Options& options =
GraphDebugInfoBuilder::Options());
void AppendGraphDebugInfo(absl::string_view prefix,
const GraphDebugInfo& new_info);
// These string methods are used in the Python bindings to avoid symbol
// resolution errors with pybind on Windows.
absl::Status AppendGraphDebugInfoStr(absl::string_view prefix,
absl::string_view new_info_str);
std::string ToGraphDebugInfoStr() const;
// Returns the GraphDebugInfo proto.
GraphDebugInfo Build() const;
private:
void AppendToStackTraceProto(const StackFrame& stack_frame,
GraphDebugInfo::StackTrace& stack_trace_proto);
std::unique_ptr<GraphDebugInfo> debug_info_;
absl::flat_hash_map<std::string, int> file_name_to_index_;
absl::flat_hash_map<StackTracePointer, int> trace_to_index_;
absl::flat_hash_map<StackFrame, int> frame_to_index_;
int new_name_index_ = 0;
GraphDebugInfoBuilder(const GraphDebugInfoBuilder&) = delete;
void operator=(const GraphDebugInfoBuilder&) = delete;
};
} // namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPH_GRAPH_DEBUG_INFO_BUILDER_H_
tensorflow/core/graph/graph_debug_info_builder_test.cc
+239
View File
@@ -0,0 +1,239 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#include "tensorflow/core/graph/graph_debug_info_builder.h"
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include <gmock/gmock.h>
#include "absl/strings/str_cat.h"
#include "tensorflow/core/framework/graph_debug_info.pb.h"
#include "tensorflow/core/platform/stack_frame.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace {
using ::testing::Eq;
using ::testing::Ne;
using ::testing::UnorderedElementsAre;
class TestStackTrace : public AbstractStackTrace {
public:
explicit TestStackTrace(const std::vector<StackFrame> frames)
: frames_(std::move(frames)) {}
absl::Span<StackFrame const> ToFrames() const override { return frames_; }
std::vector<StackFrame> ToUncachedFrames() const override { return frames_; }
std::vector<StackFrame> GetUserFrames(int limit) const override {
return frames_;
}
StackFrame LastUserFrame() const override { return frames_.back(); }
std::string ToString(const TracePrintingOptions& opts) const override {
auto frame = LastUserFrame();
return absl::StrCat(frame.file_name, ":", frame.line_number, ":",
frame.function_name);
}
std::vector<StackFrame> frames_;
};
TEST(GraphDebugInfoBuilderTest, AccumulateStackTrace) {
auto stack_trace = std::make_shared<TestStackTrace>(
std::vector<StackFrame>{{"dummy_file_alpha.cc", 20, "function_bar"},
{"dummy_file_beta.cc", 30, "function_sop"}});
GraphDebugInfoBuilder builder;
builder.AccumulateStackTrace(stack_trace, "alpha_beta");
GraphDebugInfo debug_info = builder.Build();
EXPECT_THAT(debug_info.files(), UnorderedElementsAre("dummy_file_alpha.cc",
"dummy_file_beta.cc"));
EXPECT_THAT(debug_info.traces_by_id_size(), Eq(1));
EXPECT_THAT(debug_info.name_to_trace_id().find("alpha_beta"),
Ne(debug_info.name_to_trace_id().end()));
auto actual_stack_trace = debug_info.traces_by_id().at(
debug_info.name_to_trace_id().at("alpha_beta"));
EXPECT_THAT(actual_stack_trace.frame_id_size(), Eq(2))
<< debug_info.DebugString();
}
TEST(GraphDebugInfoBuilderTest, AccumulateStackTracesMap) {
StackTracesMap stack_traces;
stack_traces["two"] = std::make_shared<TestStackTrace>(
std::vector<StackFrame>{{"dummy_file_alpha.cc", 20, "function_bar"},
{"dummy_file_beta.cc", 30, "function_sop"}});
stack_traces["scale"] =
std::make_shared<TestStackTrace>(std::vector<StackFrame>{
{"dummy_file_alpha.cc", 10, "function_foo"},
{"dummy_file_beta.cc", 30, "function_sop"},
});
stack_traces["y"] = std::make_shared<TestStackTrace>(std::vector<StackFrame>{
{"dummy_file_alpha.cc", 15, "function_flex"},
{"dummy_file_alpha.cc", 20, "function_bar"},
{"dummy_file_beta.cc", 30, "function_sop"},
});
GraphDebugInfoBuilder builder;
builder.AccumulateStackTracesMap(stack_traces, "@func");
GraphDebugInfo debug_info = builder.Build();
EXPECT_THAT(debug_info.files(), UnorderedElementsAre("dummy_file_alpha.cc",
"dummy_file_beta.cc"));
EXPECT_THAT(debug_info.name_to_trace_id_size(), Eq(3));
// Examine one of the three stack traces in detail.
EXPECT_THAT(debug_info.name_to_trace_id().find("scale@func"),
Ne(debug_info.name_to_trace_id().end()));
auto stack_trace = debug_info.traces_by_id().at(
debug_info.name_to_trace_id().at("scale@func"));
EXPECT_THAT(stack_trace.frame_id_size(), Eq(2));
std::vector<GraphDebugInfo::FileLineCol> file_line_cols;
for (auto& frame_id : stack_trace.frame_id()) {
file_line_cols.push_back(debug_info.frames_by_id().at(frame_id));
}
// `FileLineCol.file_index` is non-deterministic because the GraphDebugInfo is
// built by accumulating all file names into a set, and then storing that in
// the `files` field in an arbitrary order.
auto file_line_col_0 = file_line_cols[0];
auto file_line_col_1 = file_line_cols[1];
EXPECT_THAT(std::vector<int>(
{file_line_col_0.file_index(), file_line_col_1.file_index()}),
UnorderedElementsAre(0, 1));
EXPECT_THAT(file_line_col_0.line(), Eq(10));
EXPECT_THAT(file_line_col_0.func(), Eq("function_foo"));
EXPECT_THAT(file_line_col_1.line(), Eq(30));
EXPECT_THAT(file_line_col_1.func(), Eq("function_sop"));
}
TEST(GraphDebugInfoBuilderTest, AppendGraphDebugInfo) {
GraphDebugInfo a;
// Function stack traces are commonly returned without a prefix.
// Validate that we can accumulate these correctly.
{
GraphDebugInfoBuilder builder;
StackTracesMap stack_traces;
stack_traces["two"] = std::make_shared<TestStackTrace>(
std::vector<StackFrame>{{"dummy_file_alpha.cc", 20, "function_bar"}});
stack_traces["scale"] = std::make_shared<TestStackTrace>(
std::vector<StackFrame>{{"dummy_file_alpha.cc", 10, "function_foo"}});
builder.AccumulateStackTracesMap(stack_traces, "");
a = builder.Build();
}
GraphDebugInfo b;
{
GraphDebugInfoBuilder builder;
StackTracesMap stack_traces;
stack_traces["y"] =
std::make_shared<TestStackTrace>(std::vector<StackFrame>{
{"dummy_file_alpha.cc", 15, "function_flex"},
});
builder.AccumulateStackTracesMap(stack_traces, "");
b = builder.Build();
}
// With builtin prefix
GraphDebugInfo c;
{
GraphDebugInfoBuilder builder;
StackTracesMap stack_traces;
stack_traces["z"] =
std::make_shared<TestStackTrace>(std::vector<StackFrame>{
{"dummy_file_alpha.cc", 15, "function_flex"},
});
builder.AccumulateStackTracesMap(stack_traces, "@func3");
c = builder.Build();
}
GraphDebugInfoBuilder builder;
builder.AppendGraphDebugInfo("func1", a);
builder.AppendGraphDebugInfo("func2", b);
builder.AppendGraphDebugInfo("", c);
GraphDebugInfo combined = builder.Build();
EXPECT_EQ(combined.name_to_trace_id().size(), 4);
std::vector<std::string> keys{"two@func1", "scale@func1", "y@func2",
"z@func3"};
for (const auto& key : keys) {
EXPECT_THAT(combined.name_to_trace_id().find(key),
Ne(combined.name_to_trace_id().end()));
}
}
TEST(StackTracesMapToGraphDebugInfoTest, EmptyMap) {
StackTracesMap map;
GraphDebugInfo generated = StackTracesMapToGraphDebugInfo(map);
EXPECT_EQ(generated.files_size(), 0);
EXPECT_EQ(generated.traces_size(), 0);
}
TEST(StackTracesMapToGraphDebugInfoTest, EmptyFrames) {
StackTracesMap map;
std::vector<StackFrame> frames;
auto stack_trace = std::make_shared<FrozenStackTrace>(frames);
map.insert({"dummy_name", stack_trace});
GraphDebugInfo generated = StackTracesMapToGraphDebugInfo(map);
EXPECT_EQ(generated.files_size(), 0);
EXPECT_EQ(generated.traces_by_id_size(), 1);
EXPECT_TRUE(generated.name_to_trace_id().contains("dummy_name"));
}
TEST(StackTracesMapToGraphDebugInfoTest, RoundTripStackTraces) {
StackTracesMap map;
std::vector<StackFrame> frames = {
StackFrame({"dummy_file_name", 10, "dummy_function_name"}),
StackFrame({"dummy_file_name", 20, "other_function_name"})};
auto stack_trace = std::make_shared<FrozenStackTrace>(frames);
map.insert({"dummy_name", stack_trace});
GraphDebugInfo generated = StackTracesMapToGraphDebugInfo(map);
StackTracesMap output = LoadTracesFromDebugInfo(generated);
for (auto [name, trace] : output) {
auto orig_trace = map[name];
EXPECT_NE(orig_trace, nullptr);
EXPECT_EQ(orig_trace->ToFrames(), trace->ToFrames());
}
}
TEST(StackTracesTest, ToFrames) {
StackTracesMap map;
std::vector<StackFrame> frames = {
StackFrame({"dummy_file_name", 10, "dummy_function_name"}),
StackFrame({"other_file_name", 20, "other_function_name"})};
auto stack_trace = TestStackTrace(frames);
EXPECT_EQ(stack_trace.ToFrames().size(), 2);
auto uncached_frames = stack_trace.ToUncachedFrames();
EXPECT_EQ(uncached_frames.size(), 2);
EXPECT_EQ(frames[0], uncached_frames[0]);
EXPECT_EQ(frames[1], uncached_frames[1]);
}
} // namespace
} // namespace tensorflow
tensorflow/core/graph/graph_def_builder.cc
+138
View File
@@ -0,0 +1,138 @@
/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#include "tensorflow/core/graph/graph_def_builder.h"
#include <utility>
#include "tensorflow/core/graph/tensor_id.h"
#include "tensorflow/core/lib/core/errors.h"
namespace tensorflow {
GraphDefBuilder::Options::Options(Graph* graph, absl::Status* status)
: graph_(graph), status_(status) {}
GraphDefBuilder::Options::~Options() {}
GraphDefBuilder::Options GraphDefBuilder::Options::WithName(
absl::string_view name) const {
return Options(*this).WithNameImpl(name);
}
GraphDefBuilder::Options GraphDefBuilder::Options::WithDevice(
absl::string_view device) const {
return Options(*this).WithDeviceImpl(device);
}
GraphDefBuilder::Options GraphDefBuilder::Options::WithControlInput(
Node* control_input) const {
return Options(*this).WithControlInputImpl(control_input);
}
GraphDefBuilder::Options GraphDefBuilder::Options::WithControlInputs(
absl::Span<Node* const> control_inputs) const {
return Options(*this).WithControlInputsImpl(control_inputs);
}
GraphDefBuilder::Options GraphDefBuilder::Options::WithNameImpl(
absl::string_view name) {
name_ = std::string(name);
return *this;
}
GraphDefBuilder::Options GraphDefBuilder::Options::WithDeviceImpl(
absl::string_view device) {
device_ = std::string(device);
return *this;
}
GraphDefBuilder::Options GraphDefBuilder::Options::WithControlInputImpl(
Node* control_input) {
control_inputs_.push_back(control_input);
return *this;
}
GraphDefBuilder::Options GraphDefBuilder::Options::WithControlInputsImpl(
absl::Span<Node* const> control_inputs) {
control_inputs_.insert(control_inputs_.end(), control_inputs.begin(),
control_inputs.end());
return *this;
}
absl::Status GraphDefBuilder::ToGraphDef(GraphDef* graph_def) const {
if (status_.ok()) {
graph_.ToGraphDef(graph_def);
*graph_def->mutable_library() = flib_def_.ToProto();
}
return status_;
}
std::string GraphDefBuilder::Options::GetNameForOp(absl::string_view op) const {
if (name_.empty()) return graph_->NewName(op);
return name_;
}
Node* GraphDefBuilder::Options::FinalizeBuilder(NodeBuilder* builder) const {
builder->ControlInputs(control_inputs_);
if (!device_.empty()) builder->Device(device_);
for (const auto& attr : attrs_) {
builder->Attr(attr.first, attr.second);
}
Node* returned_node;
UpdateStatus(builder->Finalize(graph_, &returned_node));
return returned_node;
}
void GraphDefBuilder::Options::UpdateStatus(const absl::Status& status) const {
if (status_ == nullptr) {
TF_CHECK_OK(status);
} else {
status_->Update(status);
}
}
namespace ops {
Node* SourceOp(const std::string& op_name,
const GraphDefBuilder::Options& opts) {
if (opts.HaveError()) return nullptr;
NodeBuilder node_builder(opts.GetNameForOp(op_name), op_name,
opts.op_registry());
return opts.FinalizeBuilder(&node_builder);
}
Node* UnaryOp(const std::string& op_name, NodeOut input,
const GraphDefBuilder::Options& opts) {
if (opts.HaveError()) return nullptr;
NodeBuilder node_builder(opts.GetNameForOp(op_name), op_name,
opts.op_registry());
node_builder.Input(std::move(input));
return opts.FinalizeBuilder(&node_builder);
}
Node* BinaryOp(const std::string& op_name, NodeOut a, NodeOut b,
const GraphDefBuilder::Options& opts) {
if (opts.HaveError()) return nullptr;
NodeBuilder node_builder(opts.GetNameForOp(op_name), op_name,
opts.op_registry());
node_builder.Input(std::move(a)).Input(std::move(b));
return opts.FinalizeBuilder(&node_builder);
}
Node* TernaryOp(const std::string& op_name, NodeOut a, NodeOut b, NodeOut c,
const GraphDefBuilder::Options& opts) {
if (opts.HaveError()) return nullptr;
NodeBuilder node_builder(opts.GetNameForOp(op_name), op_name,
opts.op_registry());
node_builder.Input(std::move(a)).Input(std::move(b)).Input(std::move(c));
return opts.FinalizeBuilder(&node_builder);
}
} // end namespace ops
} // end namespace tensorflow
tensorflow/core/graph/graph_def_builder.h
+217
View File
@@ -0,0 +1,217 @@
/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
Licensed 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 TENSORFLOW_CORE_GRAPH_GRAPH_DEF_BUILDER_H_
#define TENSORFLOW_CORE_GRAPH_GRAPH_DEF_BUILDER_H_
#include <string>
#include <vector>
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/graph/node_builder.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/core/stringpiece.h"
#include "tensorflow/core/lib/gtl/array_slice.h"
namespace tensorflow {
// Given a function like:
// namespace ops {
// Node* Identity(NodeOut input, const GraphDefBuilder::Options& opts) {
// if (opts.HaveError()) return nullptr;
// static const string kOpName = "Identity";
// NodeBuilder node_builder(opts.GetNameForOp(kOpName), kOpName,
// opts.op_registry());
// node_builder.Input(input);
// return opts.FinalizeBuilder(&node_builder);
// }
// } // namespace ops
//
// // Or, alternatively:
// namespace ops {
// Node* Identity(NodeOut input, const GraphDefBuilder::Options& opts) {
// static const string kOpName = "Identity";
// return UnaryOp(kOpName, input, opts);
// }
// } // namespace ops
//
// You call it like:
// GraphDefBuilder b;
// using namespace ::tensorflow::ops; // NOLINT(build/namespaces)
// Node* na = Const(7, b.opts());
// // Note: WithName() returns a copy, opts is unchanged.
// Node* nb = Const(5, b.opts().WithName("control-input"));
// Node* nc = Identity(na, b.opts().WithControlInput(nb));
// GraphDef graph_def;
// Status status = b.ToGraphDef(&graph_def);
// if (!status.ok()) { /* Handle error */ }
//
// In tests you can skip the status handling via:
// GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
// ...
// b.ToGraphDef(&graph_def);
class GraphDefBuilder {
public:
// Options for adding a Node to a Graph.
class Options {
public:
// Sets the Graph (that Nodes will be added to) and the status. The
// status may be set to nullptr, in which case errors cause CHECK
// failures. The graph and status must outlive *this.
Options(Graph* graph, absl::Status* status);
~Options();
// Methods for setting options. These are const methods: they
// return a copy of *this with the option set.
Options WithName(absl::string_view name) const;
Options WithDevice(absl::string_view device) const;
Options WithControlInput(Node* control_input) const;
Options WithControlInputs(absl::Span<Node* const> control_inputs) const;
// Override the default value for an optional attr.
template <class T>
Options WithAttr(absl::string_view attr_name, T&& value) const {
return Options(*this).WithAttrImpl(attr_name, std::forward<T>(value));
}
// Note: overload needed to allow {...} expressions for value.
template <class T>
Options WithAttr(absl::string_view attr_name,
std::initializer_list<T> value) const {
return WithAttr<std::initializer_list<T>>(attr_name, std::move(value));
}
// Methods for using options from a function that creates a Node.
// Returns true if the status associated with *this has an error.
// Use this to skip processing that may depend on prior results.
bool HaveError() const { return status_ != nullptr && !status_->ok(); }
// Returns a string representation of the status associated with *this.
// Returns the string `"OK"` if the status doesn't have any error.
std::string StatusToString() const {
return status_->ok() ? "OK" : std::string(status_->message());
}
// Given the Op type name, return a name for a node of that type.
// Uses the value set in WithName() if that has been called. Otherwise,
// returns a name built out of the Op type name.
std::string GetNameForOp(absl::string_view op) const;
// Sets the device, adds control inputs, adds attrs, and calls Finalize().
// If Finalize returns an error, it is saved and this function returns
// nullptr.
Node* FinalizeBuilder(NodeBuilder* builder) const;
// Updates the associated status, if any, or calls TF_CHECK_OK if none.
void UpdateStatus(const absl::Status& status) const;
// Accessor
const OpRegistryInterface* op_registry() const {
return graph_->op_registry();
}
private:
Options WithNameImpl(absl::string_view name);
Options WithDeviceImpl(absl::string_view device);
Options WithControlInputImpl(Node* control_input);
Options WithControlInputsImpl(absl::Span<Node* const> control_inputs);
template <class T>
Options WithAttrImpl(absl::string_view name, T&& value) {
attrs_.emplace_back(std::string(name), AttrValue());
SetAttrValue(std::forward<T>(value), &attrs_.back().second);
return *this;
}
Graph* const graph_;
absl::Status* const status_;
std::string name_;
std::string device_;
std::vector<Node*> control_inputs_;
std::vector<std::pair<std::string, AttrValue>> attrs_;
};
// Start building a new graph.
explicit GraphDefBuilder(
const OpRegistryInterface* op_registry = OpRegistry::Global())
: graph_(op_registry), flib_def_(op_registry), opts_(&graph_, &status_) {}
// For use in tests, where you want to fail immediately on error instead
// of checking the status at the end.
enum TestFailImmediatelyType { kFailImmediately };
explicit GraphDefBuilder(
TestFailImmediatelyType,
const OpRegistryInterface* op_registry = OpRegistry::Global())
: graph_(op_registry), flib_def_(op_registry), opts_(&graph_, nullptr) {}
// Gets the Options with the associated Graph and Status.
const Options& opts() const { return opts_; }
// Once all the nodes have been added, call this to get whether it was
// successful, and if so fill *graph_def.
absl::Status ToGraphDef(GraphDef* graph_def) const;
// Adds the function and gradient definitions in `fdef_lib` to this graph's op
// registry. Ignores duplicate functions, and returns a bad status if an
// imported function differs from an existing function or op with the same
// name.
absl::Status AddFunctionLibrary(const FunctionDefLibrary& fdef_lib) {
return flib_def_.AddLibrary(fdef_lib);
}
// Returns whether a user-defined function with `name` already exists in the
// graph.
bool HasFunction(const std::string& name) {
return flib_def_.Find(name) != nullptr;
}
private:
Graph graph_;
FunctionLibraryDefinition flib_def_;
absl::Status status_;
Options opts_;
};
namespace ops {
// A NodeOut may either be a regular input or back input. Regular
// inputs are specified via either a Node* or a Node* and an output
// index. Back inputs are specified by a node name, output index, and
// output type.
typedef NodeBuilder::NodeOut NodeOut;
// For adding an Op with no inputs to a GraphDefBuilder.
Node* SourceOp(const std::string& op_name,
const GraphDefBuilder::Options& opts);
// For adding an Op with one input to a GraphDefBuilder.
Node* UnaryOp(const std::string& op_name, NodeOut input,
const GraphDefBuilder::Options& opts);
// For adding an Op with two inputs to a GraphDefBuilder.
Node* BinaryOp(const std::string& op_name, NodeOut a, NodeOut b,
const GraphDefBuilder::Options& opts);
// For adding an Op with three inputs to a GraphDefBuilder.
Node* TernaryOp(const std::string& op_name, NodeOut a, NodeOut b, NodeOut c,
const GraphDefBuilder::Options& opts);
} // namespace ops
} // namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPH_GRAPH_DEF_BUILDER_H_
tensorflow/core/graph/graph_def_builder_test.cc
+51
View File
@@ -0,0 +1,51 @@
/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#include "tensorflow/core/graph/graph_def_builder.h"
#include "tensorflow/core/common_runtime/graph_def_builder_util.h"
#include "tensorflow/core/framework/versions.pb.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/kernels/ops_util.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/public/version.h"
namespace tensorflow {
namespace {
TEST(GraphDefBuilderTest, Version) {
// Verify that our assertions will be nontrivial
ASSERT_LT(0, TF_GRAPH_DEF_VERSION);
// Newly built graphs should use the current version
GraphDefBuilder builder(GraphDefBuilder::kFailImmediately);
// Check version when we convert to a Graph
Graph graph(OpRegistry::Global());
TF_EXPECT_OK(GraphDefBuilderToGraph(builder, &graph));
ASSERT_EQ(graph.versions().producer(), TF_GRAPH_DEF_VERSION);
ASSERT_EQ(graph.versions().min_consumer(), TF_GRAPH_DEF_VERSION_MIN_CONSUMER);
// Check version when we convert to a GraphDef
GraphDef graph_def;
TF_EXPECT_OK(builder.ToGraphDef(&graph_def));
ASSERT_EQ(graph_def.versions().producer(), TF_GRAPH_DEF_VERSION);
ASSERT_EQ(graph_def.versions().min_consumer(),
TF_GRAPH_DEF_VERSION_MIN_CONSUMER);
}
} // namespace
} // namespace tensorflow
tensorflow/core/graph/graph_node_util.cc
+92
View File
@@ -0,0 +1,92 @@
/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#include "tensorflow/core/graph/graph_node_util.h"
#include <vector>
#include "absl/container/btree_set.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/core/stringpiece.h"
namespace tensorflow {
std::string SummarizeNode(const Node& node) {
return SummarizeNodeDef(node.def());
}
std::string FormatNodeForError(const Node& node) {
return FormatNodeDefForError(node.def());
}
absl::Status NameRangesForNode(const Node& node, const OpDef& op_def,
NameRangeMap* inputs, NameRangeMap* outputs) {
return NameRangesForNode(node.def(), op_def, inputs, outputs);
}
absl::Status AttachDef(const absl::Status& status, const Node& node,
bool allow_multiple_formatted_node) {
return AttachDef(status, node.def(), allow_multiple_formatted_node);
}
absl::btree_set<std::string> GetMergedNames(
const std::vector<std::string>& from_names,
const std::vector<std::string>& to_names) {
absl::btree_set<std::string> merged_names;
merged_names.insert(from_names.begin(), from_names.end());
merged_names.insert(to_names.begin(), to_names.end());
return merged_names;
}
void MergeDebugInfo(const NodeDebugInfo& from, Node* to_node) {
NodeDebugInfo to = NodeDebugInfo(*to_node);
if (!from.original_node_names.empty()) {
auto node_names =
GetMergedNames(from.original_node_names, to.original_node_names);
to_node->set_original_node_names({node_names.begin(), node_names.end()});
}
if (!from.original_func_names.empty()) {
auto func_names =
GetMergedNames(from.original_func_names, to.original_func_names);
to_node->set_original_func_names({func_names.begin(), func_names.end()});
}
}
void MergeDebugInfo(const NodeDebugInfo& from, NodeDef* to_node_def) {
NodeDebugInfo to = NodeDebugInfo(*to_node_def);
if (!from.original_node_names.empty()) {
auto node_names =
GetMergedNames(from.original_node_names, to.original_node_names);
to_node_def->mutable_experimental_debug_info()->clear_original_node_names();
*to_node_def->mutable_experimental_debug_info()
->mutable_original_node_names() = {node_names.begin(),
node_names.end()};
}
if (!from.original_func_names.empty()) {
auto func_names =
GetMergedNames(from.original_func_names, to.original_func_names);
to_node_def->mutable_experimental_debug_info()->clear_original_func_names();
*to_node_def->mutable_experimental_debug_info()
->mutable_original_func_names() = {func_names.begin(),
func_names.end()};
}
}
void MergeDebugInfo(const NodeDef& from_node_def, NodeDef* to_node_def) {
MergeDebugInfo(NodeDebugInfo(from_node_def), to_node_def);
}
} // namespace tensorflow
tensorflow/core/graph/graph_node_util.h
+64
View File
@@ -0,0 +1,64 @@
/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
Licensed 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 TENSORFLOW_CORE_GRAPH_GRAPH_NODE_UTIL_H_
#define TENSORFLOW_CORE_GRAPH_GRAPH_NODE_UTIL_H_
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/lib/core/stringpiece.h"
#include "tensorflow/core/platform/status.h"
namespace tensorflow {
class Node;
struct NodeDebugInfo;
// We forward declare protos so that kernels don't need to depend on them
class NodeDef;
class OpDef;
// Produce a human-readable version of a Node or NodeDef that is more concise
// than a text-format proto.
std::string SummarizeNode(const Node& node);
// Produces a formatted string pattern from the node which can uniquely identify
// this node upstream to produce an informative error message. The pattern
// followed is: {{node <node_name>}}
std::string FormatNodeForError(const Node& node);
// Merges the original node names from the debug information of 'from' to the
// debug information of 'to'.
void MergeDebugInfo(const NodeDebugInfo& from, Node* to);
void MergeDebugInfo(const NodeDebugInfo& from, NodeDef* to);
void MergeDebugInfo(const NodeDef& from, NodeDef* to);
// Computes the mapping from input/output argument name to the
// corresponding input/output index range. For example,
// input "foo" corresponds to input indices
// [ (*inputs)["foo"].first, (*inputs)["foo"].second ).
// NOTE(mrry): To reduce allocations when the map is used and save
// space, the returned `NameRangeMap` objects borrow the input/output
// argument names from `op_def`. The `op_def` must outlive the
// returned `NameRangeMap` objects.
absl::Status NameRangesForNode(const Node& node, const OpDef& op_def,
NameRangeMap* inputs, NameRangeMap* outputs);
// Returns "status" with formatted Node attached as additional text
// in the error message. If 'allow_multiple_formatted_node' is false and there
// is already a formatted Node present in 'status', we simply attach the name
// of the Node instead of the formatted string.
absl::Status AttachDef(const absl::Status& status, const Node& node,
bool allow_multiple_formatted_node = false);
} // namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPH_GRAPH_NODE_UTIL_H_
tensorflow/core/graph/graph_partition.cc
+1315
View File
File diff suppressed because it is too large Load Diff
tensorflow/core/graph/graph_partition.h
+110
View File
@@ -0,0 +1,110 @@
/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
Licensed 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 TENSORFLOW_CORE_GRAPH_GRAPH_PARTITION_H_
#define TENSORFLOW_CORE_GRAPH_GRAPH_PARTITION_H_
#include <functional>
#include <string>
#include <unordered_map>
#include <vector>
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/graph/costmodel.h"
#include "tensorflow/core/graph/graph.h"
namespace tensorflow {
struct PartitionOptions {
// A function that returns a location for the execution of a given
// Node.
typedef std::function<std::string(const Node*)> NodeToLocFunc;
NodeToLocFunc node_to_loc = nullptr;
// A function that returns a unique graph node name with the given
// prefix.
typedef std::function<std::string(const std::string&)> NewNameFunc;
NewNameFunc new_name = nullptr;
// A function that returns the incarnation of a device given the
// device's fullname. If not found, GetIncarnationFunc should return
// kIllegalIncarnation.
static constexpr uint64_t kIllegalIncarnation = 0;
typedef std::function<uint64_t(const std::string&)> GetIncarnationFunc;
GetIncarnationFunc get_incarnation = nullptr;
// If specified, flib_def defines a function library that should be
// partitioned and replicated into each resulting partition graphs.
const FunctionLibraryDefinition* flib_def = nullptr;
// True if all the control flow "code" has already been added. The
// control flow code needs to be added when we still have the entire
// graph before any partitioning. So this flag should be false for
// the first partitioning but true for all subsequent partitioning.
//
// TODO(yuanbyu): We could also make the addition of the control
// flow code incremental based on 'node_to_loc'. This makes the
// communication a broadcast tree, which could be more efficient when
// the number of participating devices is large.
bool control_flow_added = false;
// A function that returns the data type into which the tensor
// should be cast before sent over the wire.
typedef std::function<DataType(const Edge*)> ShouldCastFunc;
ShouldCastFunc should_cast = nullptr;
// Schedule the execution of the recvs based on their start times
// computed by some scheduling algorithm. The recvs are divided into
// epochs based on their start times. A recv is enabled only when
// execution reaches its epoch - N for some predefined N.
bool scheduling_for_recvs = false;
// The start time for each node in the graph computed by some scheduling
// algorithm. If 'need_to_record_start_times' is true, we record them
// in the graph as a node attribute.
bool need_to_record_start_times = false;
std::vector<Microseconds> start_times;
// Optional customized function to compute the "tensor_name" attr value of
// Send/Recv ops inserted during partitioning.
std::function<std::string(const Edge*)> get_tensor_name_attr = nullptr;
// If true, the `Partition()` function can make destructive changes to the
// passed-in `Graph`.
//
// TODO(b/327983931): Add wrapper functions for partitioning that clearly
// signal this intent by taking a `Graph` or `Graph&&`.
bool can_make_destructive_changes = false;
};
// Partition "input" graph into a set of graphs, one per location.
// The location for node n is derived by calling opts.node_to_loc(n).
// New nodes added by Partition use "opts.new_name(old_name)" to
// generate node names.
//
// Stores the partitions in *partitions.
absl::Status Partition(const PartitionOptions& opts, Graph* input,
std::unordered_map<std::string, GraphDef>* partitions);
// Add control edges to the partitions to control the ordering
// and timing of the recv nodes based on the start times calculated
// using some scheduling algorithm.
absl::Status AddControlEdges(
const PartitionOptions& opts,
std::unordered_map<std::string, GraphDef>* partitions);
} // namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPH_GRAPH_PARTITION_H_
tensorflow/core/graph/graph_partition_test.cc
+817
View File
@@ -0,0 +1,817 @@
/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#include "tensorflow/core/graph/graph_partition.h"
#include <memory>
#include <string>
#include <unordered_map>
#include <utility>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/strings/str_cat.h"
#include "tensorflow/cc/ops/array_ops.h"
#include "tensorflow/cc/ops/const_op.h"
#include "tensorflow/cc/ops/control_flow_ops.h"
#include "tensorflow/cc/ops/control_flow_ops_internal.h"
#include "tensorflow/cc/ops/math_ops.h"
#include "tensorflow/cc/ops/random_ops.h"
#include "tensorflow/cc/ops/sendrecv_ops.h"
#include "tensorflow/cc/ops/while_loop.h"
#include "tensorflow/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/framework/common_shape_fns.h"
#include "tensorflow/core/framework/function_testlib.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/versions.pb.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/graph/graph_debug_info_builder.h"
#include "tensorflow/core/graph/graph_def_builder.h"
#include "tensorflow/core/kernels/ops_util.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/lib/strings/str_util.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/protobuf.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/public/version.h"
#include "tensorflow/core/util/equal_graph_def.h"
namespace tensorflow {
// from graph_partition.cc
extern absl::Status TopologicalSortNodesWithTimePriority(
const GraphDef* gdef,
std::vector<std::pair<const NodeDef*, int64_t>>* nodes,
std::unordered_map<const NodeDef*, int64_t>* node_to_start_time_out);
namespace {
using ops::_Recv;
using ops::_Send;
using ops::Const;
using ops::Identity;
using ops::LoopCond;
using ops::NextIteration;
using ::testing::Ne;
const char gpu_device[] = "/job:a/replica:0/task:0/device:GPU:0";
std::string SplitByDevice(const Node* node) {
return node->assigned_device_name();
}
std::string DeviceName(const Node* node) {
char first = node->name()[0];
if (first == 'G') {
return gpu_device;
} else {
const std::string cpu_prefix = "/job:a/replica:0/task:0/cpu:";
int index = first - 'A';
return absl::StrCat(cpu_prefix, index);
}
}
void Partition(const GraphDef& graph_def,
std::unordered_map<std::string, GraphDef>* partitions) {
Graph g(OpRegistry::Global());
GraphConstructorOptions opts;
TF_CHECK_OK(ConvertGraphDefToGraph(opts, graph_def, &g));
// Assigns devices to each node. Uses 1st letter of the node name as the
// device index if no device is specified.
for (Node* node : g.nodes()) {
std::string device_name = !node->requested_device().empty()
? node->requested_device()
: DeviceName(node);
node->set_assigned_device_name(device_name);
}
PartitionOptions popts;
popts.node_to_loc = SplitByDevice;
popts.new_name = [&g](const std::string& prefix) {
return g.NewName(prefix);
};
popts.get_incarnation = [](const std::string& name) {
return (name[0] - 'A') + 100;
};
absl::Status s = Partition(popts, &g, partitions);
CHECK(s.ok()) << s;
// Check versions.
EXPECT_EQ(graph_def.versions().producer(), TF_GRAPH_DEF_VERSION);
// Partitions must inherit the versions of the original graph.
for (auto& it : *partitions) {
EXPECT_EQ(graph_def.versions().producer(), it.second.versions().producer());
EXPECT_EQ(graph_def.versions().min_consumer(),
it.second.versions().min_consumer());
}
}
void CheckLoopConstruction(const GraphDef& graph_def) {
std::unordered_map<std::string, GraphDef> partitions;
Partition(graph_def, &partitions);
for (const auto& kv : partitions) {
const GraphDef& gdef = kv.second;
bool has_control_enter = false;
bool has_control_merge = false;
bool has_control_switch = false;
bool has_control_next = false;
for (const NodeDef& ndef : gdef.node()) {
// _recvs must have a control input
if (ndef.op() == "_Recv") {
bool has_control = false;
for (const std::string& input_name : ndef.input()) {
if (absl::StartsWith(input_name, "^")) {
has_control = true;
break;
}
}
EXPECT_TRUE(has_control);
}
// Must have a control loop
if (absl::StartsWith(ndef.name(), "_cloop")) {
if (ndef.op() == "Enter") {
has_control_enter = true;
}
if (ndef.op() == "Merge") {
has_control_merge = true;
}
if (ndef.op() == "Switch") {
has_control_switch = true;
}
if (ndef.op() == "NextIteration") {
has_control_next = true;
}
}
}
EXPECT_TRUE(has_control_enter);
EXPECT_TRUE(has_control_merge);
EXPECT_TRUE(has_control_switch);
EXPECT_TRUE(has_control_next);
}
}
REGISTER_OP("FloatInput")
.Output("o: float")
.SetShapeFn(shape_inference::UnknownShape);
REGISTER_OP("BoolInput")
.Output("o: bool")
.SetShapeFn(shape_inference::UnknownShape);
REGISTER_OP("Combine")
.Input("a: float")
.Input("b: float")
.Output("o: float")
.SetShapeFn(shape_inference::UnknownShape);
Output ConstructOp(const Scope& scope, const std::string& op_type,
const absl::Span<const Input> inputs) {
if (!scope.ok()) return Output();
const std::string unique_name = scope.GetUniqueNameForOp(op_type);
auto builder =
NodeBuilder(unique_name, op_type, scope.graph()->op_registry());
for (auto const& input : inputs) {
builder.Input(ops::NodeOut(input.node(), input.index()));
}
scope.UpdateBuilder(&builder);
Node* ret;
scope.UpdateStatus(builder.Finalize(scope.graph(), &ret));
if (!scope.ok()) return Output();
scope.UpdateStatus(scope.DoShapeInference(ret));
if (!scope.ok()) return Output();
return Output(ret);
}
Output FloatInput(const Scope& scope) {
return ConstructOp(scope, "FloatInput", {});
}
Output BoolInput(const Scope& scope) {
return ConstructOp(scope, "BoolInput", {});
}
Output Combine(const Scope& scope, Input a, Input b) {
return ConstructOp(scope, "Combine", {std::move(a), std::move(b)});
}
std::string FormatStackTrace(const GraphDebugInfo::StackTrace& stack_trace,
const GraphDebugInfo& debug_info) {
std::string result;
for (const GraphDebugInfo::FileLineCol& file_line_col :
stack_trace.file_line_cols()) {
const std::string& file = debug_info.files(file_line_col.file_index());
absl::StrAppend(&result, file_line_col.func(), "@", file, ":",
file_line_col.line(), ".", file_line_col.col(), "\n");
}
return result;
}
class GraphPartitionTest : public ::testing::Test {
protected:
GraphPartitionTest()
: in_(Scope::NewRootScope().ExitOnError()),
scope_a_(Scope::NewRootScope().ExitOnError().WithDevice(
"/job:a/replica:0/task:0/cpu:0")),
scope_b_(Scope::NewRootScope().ExitOnError().WithDevice(
"/job:a/replica:0/task:0/cpu:1")) {}
const GraphDef& ToGraphDef(bool include_debug_info = false) {
TF_EXPECT_OK(in_.ToGraphDef(&in_graph_def_, include_debug_info));
return in_graph_def_;
}
void ExpectMatchA() {
GraphDef graph_def;
TF_EXPECT_OK(scope_a_.ToGraphDef(&graph_def));
std::string a = "/job:a/replica:0/task:0/cpu:0";
TF_EXPECT_GRAPH_EQ(graph_def, partitions_[a]);
}
void ExpectMatchB() {
GraphDef graph_def;
TF_EXPECT_OK(scope_b_.ToGraphDef(&graph_def));
std::string b = "/job:a/replica:0/task:0/cpu:1";
TF_EXPECT_GRAPH_EQ(graph_def, partitions_[b]);
}
void ExpectFunctions(const FunctionDefLibrary& library,
const std::set<std::string>& expected_names) {
std::set<std::string> actual_names;
for (const FunctionDef& fdef : library.function()) {
actual_names.insert(fdef.signature().name());
}
EXPECT_EQ(actual_names, expected_names);
}
Scope in_;
GraphDef in_graph_def_;
Scope scope_a_;
Scope scope_b_;
std::unordered_map<std::string, GraphDef> partitions_;
};
TEST_F(GraphPartitionTest, SingleDevice) {
auto a1 = FloatInput(in_.WithOpName("A1"));
Combine(in_.WithOpName("A2"), a1, a1);
Partition(ToGraphDef(), &partitions_);
EXPECT_EQ(1, partitions_.size());
a1 = FloatInput(scope_a_.WithOpName("A1"));
Combine(scope_a_.WithOpName("A2"), a1, a1);
ExpectMatchA();
}
TEST_F(GraphPartitionTest, CrossDeviceData) {
auto a1 = FloatInput(in_.WithOpName("A1"));
auto b1 = FloatInput(in_.WithOpName("B1"));
Combine(in_.WithOpName("B2"), a1, b1);
Partition(ToGraphDef(), &partitions_);
EXPECT_EQ(2, partitions_.size());
std::string a = "/job:a/replica:0/task:0/cpu:0";
std::string b = "/job:a/replica:0/task:0/cpu:1";
a1 = FloatInput(scope_a_.WithOpName("A1"));
_Send(scope_a_.WithOpName("A1/_0"), a1, "edge_1_A1", a, 82, b);
ExpectMatchA();
b1 = FloatInput(scope_b_.WithOpName("B1"));
auto recv =
_Recv(scope_b_.WithOpName("A1/_1"), DT_FLOAT, "edge_1_A1", a, 82, b);
Combine(scope_b_.WithOpName("B2"), recv, b1);
ExpectMatchB();
}
TEST_F(GraphPartitionTest, CrossDeviceControl) {
auto a1 = FloatInput(in_.WithOpName("A1"));
auto b1 = FloatInput(in_.WithOpName("B1"));
Combine(in_.WithOpName("B2").WithControlDependencies(a1), b1, b1);
Partition(ToGraphDef(), &partitions_);
EXPECT_EQ(2, partitions_.size());
std::string a = "/job:a/replica:0/task:0/cpu:0";
std::string b = "/job:a/replica:0/task:0/cpu:1";
a1 = FloatInput(scope_a_.WithOpName("A1"));
auto c =
Const(scope_a_.WithOpName("A1/ctrl/_0").WithControlDependencies(a1), {});
_Send(scope_a_.WithOpName("A1/_1"), c, "edge_3_A1", a, 82, b);
ExpectMatchA();
auto recv =
_Recv(scope_b_.WithOpName("A1/_2"), DT_FLOAT, "edge_3_A1", a, 82, b);
auto id = Identity(scope_b_.WithOpName("A1/_3"), recv);
b1 = FloatInput(scope_b_.WithOpName("B1"));
Combine(scope_b_.WithOpName("B2").WithControlDependencies(id), b1, b1);
ExpectMatchB();
}
TEST_F(GraphPartitionTest, CrossDeviceData_MultiUse) {
auto a1 = FloatInput(in_.WithOpName("A1"));
auto b1 = FloatInput(in_.WithOpName("B1"));
Combine(in_.WithOpName("B2"), a1, b1);
Combine(in_.WithOpName("B3"), a1, a1);
Partition(ToGraphDef(), &partitions_);
EXPECT_EQ(2, partitions_.size());
std::string a = "/job:a/replica:0/task:0/cpu:0";
std::string b = "/job:a/replica:0/task:0/cpu:1";
a1 = FloatInput(scope_a_.WithOpName("A1"));
_Send(scope_a_.WithOpName("A1/_0"), a1, "edge_1_A1", a, 82, b);
ExpectMatchA();
auto recv =
_Recv(scope_b_.WithOpName("A1/_1"), DT_FLOAT, "edge_1_A1", a, 82, b);
b1 = FloatInput(scope_b_.WithOpName("B1"));
Combine(scope_b_.WithOpName("B2"), recv, b1);
Combine(scope_b_.WithOpName("B3"), recv, recv);
ExpectMatchB();
}
TEST_F(GraphPartitionTest, CrossDeviceControl_MultiUse) {
auto a1 = FloatInput(in_.WithOpName("A1"));
auto b1 = FloatInput(in_.WithOpName("B1"));
Combine(in_.WithOpName("B2").WithControlDependencies(a1), b1, b1);
FloatInput(in_.WithOpName("B3").WithControlDependencies(a1));
Partition(ToGraphDef(), &partitions_);
EXPECT_EQ(2, partitions_.size());
std::string a = "/job:a/replica:0/task:0/cpu:0";
std::string b = "/job:a/replica:0/task:0/cpu:1";
a1 = FloatInput(scope_a_.WithOpName("A1"));
auto c =
Const(scope_a_.WithOpName("A1/ctrl/_0").WithControlDependencies(a1), {});
_Send(scope_a_.WithOpName("A1/_1"), c, "edge_3_A1", a, 82, b);
ExpectMatchA();
auto recv =
_Recv(scope_b_.WithOpName("A1/_2"), DT_FLOAT, "edge_3_A1", a, 82, b);
auto id = Identity(scope_b_.WithOpName("A1/_3"), recv);
b1 = FloatInput(scope_b_.WithOpName("B1"));
Combine(scope_b_.WithOpName("B2").WithControlDependencies(id), b1, b1);
FloatInput(scope_b_.WithOpName("B3").WithControlDependencies(id));
ExpectMatchB();
}
TEST_F(GraphPartitionTest, CrossDevice_DataControl) {
auto a1 = FloatInput(in_.WithOpName("A1"));
auto b1 = FloatInput(in_.WithOpName("B1"));
Combine(in_.WithOpName("B2"), a1, b1);
FloatInput(in_.WithOpName("B3").WithControlDependencies(a1));
Partition(ToGraphDef(), &partitions_);
EXPECT_EQ(2, partitions_.size());
std::string a = "/job:a/replica:0/task:0/cpu:0";
std::string b = "/job:a/replica:0/task:0/cpu:1";
a1 = FloatInput(scope_a_.WithOpName("A1"));
_Send(scope_a_.WithOpName("A1/_0"), a1, "edge_1_A1", a, 82, b);
auto c =
Const(scope_a_.WithOpName("A1/ctrl/_2").WithControlDependencies(a1), {});
// NOTE: Send 0 A1/_1 -> A1/_2 is not necessarily needed. We could
// use A1/_0 -> A1/_4 as the control as a minor optimization.
_Send(scope_a_.WithOpName("A1/_3"), c, "edge_3_A1", a, 82, b);
ExpectMatchA();
auto recv1 =
_Recv(scope_b_.WithOpName("A1/_4"), DT_FLOAT, "edge_3_A1", a, 82, b);
auto id1 = Identity(scope_b_.WithOpName("A1/_5"), recv1);
auto recv2 =
_Recv(scope_b_.WithOpName("A1/_1"), DT_FLOAT, "edge_1_A1", a, 82, b);
b1 = FloatInput(scope_b_.WithOpName("B1"));
Combine(scope_b_.WithOpName("B2"), recv2, b1);
FloatInput(scope_b_.WithOpName("B3").WithControlDependencies(id1));
ExpectMatchB();
}
TEST_F(GraphPartitionTest, CrossDeviceLoopSimple) {
auto a1 = BoolInput(in_.WithOpName("A1"));
auto a2 = ::tensorflow::ops::internal::Enter(in_.WithOpName("A2"), a1, "foo");
auto a3 = ::tensorflow::ops::Merge(in_.WithOpName("A3"),
{a2, Input("A5", 0, DT_BOOL)})
.output;
LoopCond(in_.WithOpName("A4"), a3);
auto b1 = Identity(in_.WithOpName("B1"), a3);
NextIteration(in_.WithOpName("A5"), b1);
CheckLoopConstruction(ToGraphDef());
}
TEST_F(GraphPartitionTest, CrossDeviceLoopSimple1) {
auto a1 = BoolInput(in_.WithOpName("A1"));
auto a2 = ::tensorflow::ops::internal::Enter(in_.WithOpName("B2"), a1, "foo");
auto a3 = ::tensorflow::ops::Merge(in_.WithOpName("A3"),
{a2, Input("B5", 0, DT_BOOL)})
.output;
LoopCond(in_.WithOpName("A4"), a3);
auto b1 = Identity(in_.WithOpName("B1"), a3);
NextIteration(in_.WithOpName("B5"), b1);
std::unordered_map<std::string, GraphDef> partitions;
Partition(ToGraphDef(), &partitions);
for (const auto& kv : partitions) {
const GraphDef& gdef = kv.second;
for (const NodeDef& ndef : gdef.node()) {
if (ndef.name() == "A3") {
// A3, B2, and B5 are on the same device.
EXPECT_EQ(ndef.input(0), "B2");
EXPECT_EQ(ndef.input(1), "B5");
}
}
}
}
TEST_F(GraphPartitionTest, CrossDeviceLoopFull) {
Scope cpu0 = in_.WithDevice("/job:a/replica:0/task:0/cpu:0");
auto p1 = ops::Placeholder(cpu0, DT_INT32);
auto p2 = ops::Placeholder(cpu0, DT_INT32);
OutputList outputs;
// while i1 < 10: i1 += i2
TF_ASSERT_OK(ops::BuildWhileLoop(
cpu0, {p1, p2},
[](const Scope& s, const std::vector<Output>& inputs, Output* output) {
*output = ops::Less(s, inputs[0], 10);
return s.status();
},
[](const Scope& s, const std::vector<Output>& inputs,
std::vector<Output>* outputs) {
Scope cpu1 = s.WithDevice("/job:a/replica:0/task:0/cpu:1");
outputs->push_back(ops::AddN(cpu1, {inputs[0], inputs[1]}));
outputs->push_back(inputs[1]);
return s.status();
},
"test_loop", &outputs));
CheckLoopConstruction(ToGraphDef());
}
TEST_F(GraphPartitionTest, PartitionIncompleteGraph) {
NodeDef ndef;
Graph g(OpRegistry::Global());
// Invalid graph since the Combine node requires an input.
bool parsed = protobuf::TextFormat::ParseFromString(
R"EOF(
name: "N"
op: "Combine"
)EOF",
&ndef);
ASSERT_TRUE(parsed);
absl::Status status;
g.AddNode(ndef, &status);
TF_ASSERT_OK(status);
PartitionOptions popts;
popts.node_to_loc = SplitByDevice;
popts.new_name = [&g](const std::string& prefix) {
return g.NewName(prefix);
};
popts.get_incarnation = [](const std::string&) { return 1; };
std::unordered_map<std::string, GraphDef> partitions;
status = Partition(popts, &g, &partitions);
// Partitioning should fail, but not crash like it did before the
// changes that accompanied the addition of this test.
EXPECT_EQ(error::INVALID_ARGUMENT, status.code()) << status;
}
TEST_F(GraphPartitionTest, Functions) {
FunctionDefLibrary fdef_lib;
*fdef_lib.add_function() = test::function::XTimesTwo();
*fdef_lib.add_function() = test::function::XTimesFour();
TF_ASSERT_OK(in_.graph()->AddFunctionLibrary(fdef_lib));
auto a1 = FloatInput(in_.WithOpName("A1"));
auto b1 = FloatInput(in_.WithOpName("B1"));
ConstructOp(in_.WithOpName("A2"), "XTimesTwo", {a1});
ConstructOp(in_.WithOpName("B2"), "XTimesFour", {b1});
// The `Partition()` helper function uses the first letter of the op name ('A'
// or 'B') to choose a device for each node.
Partition(ToGraphDef(), &partitions_);
EXPECT_EQ(2, partitions_.size());
// Test that partition graphs inherit function library from original graph.
std::string a = "/job:a/replica:0/task:0/cpu:0";
std::string b = "/job:a/replica:0/task:0/cpu:1";
// Node "A2" is placed in part `a`, and uses only "XTimesTwo".
ExpectFunctions(partitions_[a].library(), {"XTimesTwo"});
// Node "B2" is placed in part `b`, and uses both "XTimesFour" directly,
// and "XTimesTwo" in the body of "XTimesFour".
ExpectFunctions(partitions_[b].library(), {"XTimesTwo", "XTimesFour"});
}
TEST_F(GraphPartitionTest, SetIncarnation) {
GraphDef gdef;
const char* const kSendRecvAttrs = R"pb(
attr {
key: 'T'
value { type: DT_FLOAT }
}
attr {
key: 'client_terminated'
value { b: false }
}
attr {
key: 'recv_device'
value { s: 'B' }
}
attr {
key: 'send_device'
value { s: 'A' }
}
attr {
key: 'send_device_incarnation'
value { i: 0 }
}
attr {
key: 'tensor_name'
value { s: 'test' }
}
)pb";
CHECK(protobuf::TextFormat::ParseFromString(
strings::StrCat(
"node { name: 'A/Pi' op: 'Const' ",
" attr { key: 'dtype' value { type: DT_FLOAT } } ",
" attr { key: 'value' value { tensor { ",
" dtype: DT_FLOAT tensor_shape {} float_val: 3.14 } } } }",
"node { name: 'A' op: '_Send' input: 'A/Pi' ", kSendRecvAttrs, "}",
"node { name: 'B' op: '_Recv' ", kSendRecvAttrs,
" attr { key: 'tensor_type' value { type:DT_FLOAT}}}"),
&gdef));
gdef.mutable_versions()->set_producer(TF_GRAPH_DEF_VERSION);
Partition(gdef, &partitions_);
EXPECT_EQ(2, partitions_.size());
for (const auto& kv : partitions_) {
const GraphDef& gdef = kv.second;
for (const NodeDef& ndef : gdef.node()) {
if (ndef.name() == "A" || ndef.name() == "B") {
int64_t val;
TF_CHECK_OK(GetNodeAttr(ndef, "send_device_incarnation", &val));
EXPECT_EQ(val, 100); // Send device is "A".
}
}
}
}
TEST_F(GraphPartitionTest, GraphDebugInfo) {
GraphDef graph_def;
Output a1 = FloatInput(in_.WithOpName("A1"));
Output b1 = FloatInput(in_.WithOpName("B1"));
Combine(in_.WithOpName("B2"), a1, b1);
Node *a1_node = nullptr, *b1_node = nullptr, *b2_node = nullptr;
for (Node* node : in_.graph()->op_nodes()) {
if (node->name() == "A1") {
a1_node = node;
} else if (node->name() == "B1") {
b1_node = node;
} else if (node->name() == "B2") {
b2_node = node;
}
}
EXPECT_NE(a1_node, nullptr);
EXPECT_NE(b1_node, nullptr);
EXPECT_NE(b2_node, nullptr);
std::vector<StackFrame> a1_stack_trace{{"main.cc", 20, "x"},
{"alpha.cc", 30, "a1"}};
std::vector<StackFrame> b1_stack_trace{{"window.cc", 21, "y"},
{"beta.cc", 35, "b1"}};
std::vector<StackFrame> b2_stack_trace{{"cache.cc", 22, "bar"},
{"beta.cc", 39, "b2"}};
a1_node->SetStackTrace(std::make_shared<FrozenStackTrace>(a1_stack_trace));
b1_node->SetStackTrace(std::make_shared<FrozenStackTrace>(b1_stack_trace));
b2_node->SetStackTrace(std::make_shared<FrozenStackTrace>(b2_stack_trace));
TF_EXPECT_OK(in_.ToGraphDef(&graph_def, /*include_debug_info=*/true));
// `Partition()` uses the first letter of the op name ('A' or 'B') to choose a
// device for each node. It calls the function under test, also named
// `Partition()`, to do the actual partitioning.
Partition(ToGraphDef(/*include_debug_info=*/true), &partitions_);
EXPECT_EQ(2, partitions_.size());
// Expect each partitioned graph to contain the stack traces for its nodes.
// A stack trace for A1 should be in the A partition (".../cpu:0").
std::string a = "/job:a/replica:0/task:0/cpu:0";
const GraphDebugInfo& a_debug_info = partitions_[a].debug_info();
StackTracesMap traces = LoadTracesFromDebugInfo(a_debug_info);
const auto& a_it = traces.find("A1");
EXPECT_THAT(a_it, Ne(traces.end()));
EXPECT_THAT(a_it->second->ToString({}),
::testing::ContainsRegex("alpha.cc.*30"));
// Stack traces for B1 and B2 should be in the B partition (".../cpu:1").
std::string b = "/job:a/replica:0/task:0/cpu:1";
const GraphDebugInfo& b_debug_info = partitions_[b].debug_info();
traces = LoadTracesFromDebugInfo(b_debug_info);
const auto& b1_it = traces.find("B1");
const auto& b2_it = traces.find("B2");
EXPECT_THAT(b1_it, Ne(traces.end()));
EXPECT_THAT(b2_it, Ne(traces.end()));
EXPECT_THAT(b1_it->second->ToString({}),
::testing::ContainsRegex("beta.cc.*35"));
EXPECT_THAT(b2_it->second->ToString({}),
::testing::ContainsRegex("beta.cc.*39"));
}
TEST(TopologicalSortNodesWithTimePriorityTest, NoDependencies) {
// Create placeholders, shuffle them so the order in the graph is not strictly
// increasing.
Scope root = Scope::NewRootScope().ExitOnError();
std::vector<int> indexes;
for (int i = 0; i < 20; ++i) {
indexes.push_back((i + 2001) % 20);
}
std::vector<ops::Placeholder> placeholders;
for (int i : indexes) {
placeholders.emplace_back(root.WithOpName(absl::StrCat("p", i)), DT_FLOAT);
placeholders.back().node()->AddAttr("_start_time", i + 1);
}
GraphDef gdef;
TF_EXPECT_OK(root.ToGraphDef(&gdef));
std::vector<std::pair<const NodeDef*, int64_t>> nodes;
std::unordered_map<const NodeDef*, int64_t> node_to_start_time;
TF_CHECK_OK(
TopologicalSortNodesWithTimePriority(&gdef, &nodes, &node_to_start_time));
ASSERT_EQ(nodes.size(), 20);
for (int i = 0; i < nodes.size(); ++i) {
EXPECT_EQ(absl::StrCat("p", i), nodes[i].first->name());
EXPECT_EQ(i + 1, nodes[i].second);
}
}
TEST(TopologicalSortNodesWithTimePriority, Dependencies) {
// Create placeholders, shuffle them so the order in the graph is not strictly
// increasing.
Scope root = Scope::NewRootScope().ExitOnError();
std::vector<int> indexes;
std::vector<ops::Placeholder> placeholders_in_order;
const int num_leaves = 20;
for (int i = 0; i < num_leaves; ++i) {
indexes.push_back((i + 2001) % num_leaves);
placeholders_in_order.emplace_back(root.WithOpName(absl::StrCat("p", i)),
DT_FLOAT);
placeholders_in_order.back().node()->AddAttr("_start_time", i + 1);
}
std::vector<ops::Placeholder> placeholders;
for (int i : indexes) {
placeholders.push_back(placeholders_in_order[i]);
}
// Create ops that depend on the placeholders. We give start times to these
// that are in descending order (e.g., the op that depends on the first
// placeholder runs last).
std::vector<ops::Square> squares;
for (int i : indexes) {
squares.emplace_back(root.WithOpName(absl::StrCat("s", i)),
placeholders[i]);
squares.back().node()->AddAttr("_start_time", 50 - (i + 1));
}
// Create addn to sum all squares.
std::vector<Input> inputs;
for (const auto& s : squares) inputs.push_back(s);
ops::AddN addn =
ops::AddN(root.WithOpName("addn"), absl::Span<const Input>(inputs));
// Start times is actually listed earlier than the nodes it depends on.
// But because of dependency ordering, it is last in the list.
addn.node()->AddAttr("_start_time", 1);
GraphDef gdef;
TF_EXPECT_OK(root.ToGraphDef(&gdef));
std::vector<std::pair<const NodeDef*, int64_t>> nodes;
std::unordered_map<const NodeDef*, int64_t> node_to_start_time;
TF_CHECK_OK(
TopologicalSortNodesWithTimePriority(&gdef, &nodes, &node_to_start_time));
ASSERT_EQ(1 + squares.size() + placeholders.size(), nodes.size());
for (int i = 0; i < placeholders.size(); ++i) {
const NodeDef* node = nodes[i].first;
EXPECT_EQ(absl::StrCat("p", i), node->name());
EXPECT_EQ(i + 1, nodes[i].second);
EXPECT_EQ(i + 1, node_to_start_time[node]);
}
for (int i = 0; i < squares.size(); ++i) {
int node_index = placeholders.size() + i;
int square_index = num_leaves - 1 - i;
const NodeDef* node = nodes[node_index].first;
EXPECT_EQ(absl::StrCat("s", square_index), node->name());
EXPECT_EQ(50 - (square_index + 1), nodes[node_index].second);
EXPECT_EQ(50 - (square_index + 1), node_to_start_time[node]);
}
EXPECT_EQ("addn", nodes.back().first->name());
EXPECT_EQ(50, nodes.back().second);
EXPECT_EQ(50, node_to_start_time[nodes.back().first]);
}
TEST(TopologicalSortNodesWithTimePriority, WhileLoop) {
using namespace ::tensorflow::ops; // NOLINT(build/namespaces)
using namespace ::tensorflow::ops::internal; // NOLINT(build/namespaces)
// Create placeholders.
Scope root = Scope::NewRootScope().ExitOnError();
std::vector<int> indexes;
std::vector<Placeholder> placeholders_in_order;
const int num_leaves = 20;
for (int i = 0; i < num_leaves; ++i) {
indexes.push_back((i + 2001) % num_leaves);
placeholders_in_order.emplace_back(root.WithOpName(absl::StrCat("p", i)),
DT_FLOAT);
placeholders_in_order.back().node()->AddAttr("_start_time", i + 1);
}
std::vector<Placeholder> placeholders;
placeholders.reserve(indexes.size());
for (int i : indexes) {
placeholders.push_back(placeholders_in_order[i]);
}
// Add a while loop above each placeholder.
std::vector<Exit> while_exits;
const int nodes_per_loop = 8;
for (int i : indexes) {
Scope scope = root.NewSubScope(absl::StrCat("while", i));
auto dummy = Placeholder(scope, DT_FLOAT);
Enter enter(scope, placeholders[i], absl::StrCat("frame", i));
Merge merge(scope, std::initializer_list<Input>{enter, dummy});
auto cv = Const(scope.WithControlDependencies({merge.output}), false);
LoopCond loop_cond(scope, cv);
Switch switch_node(scope, merge.output, loop_cond);
Identity identity(scope, switch_node.output_true);
NextIteration next_iteration(scope, identity);
while_exits.emplace_back(scope.WithOpName("exit"),
switch_node.output_false);
// Complete loop by removing dummy node and attaching NextIteration to
// that input of the merge node.
scope.graph()->RemoveNode(dummy.node());
scope.graph()->AddEdge(next_iteration.node(), 0, merge.output.node(), 1);
int base_start_time = i * 10 + 100;
for (const auto& op : std::initializer_list<Output>{
enter, merge.output, cv, loop_cond, switch_node.output_false,
identity, next_iteration, while_exits.back()}) {
op.node()->AddAttr("_start_time", base_start_time++);
}
}
// Create ops that depend on the loop exits.
std::vector<Square> squares;
squares.reserve(indexes.size());
for (int i : indexes) {
squares.emplace_back(root.WithOpName(absl::StrCat("s", i)), while_exits[i]);
squares.back().node()->AddAttr("_start_time", 500 - (i + 1));
}
GraphDef gdef;
TF_EXPECT_OK(root.ToGraphDef(&gdef));
// Run the sort. The while loop nodes do not appear in the output <nodes>.
std::vector<std::pair<const NodeDef*, int64_t>> nodes;
std::unordered_map<const NodeDef*, int64_t> node_to_start_time;
TF_CHECK_OK(
TopologicalSortNodesWithTimePriority(&gdef, &nodes, &node_to_start_time));
ASSERT_LT(while_exits.size() + squares.size() + placeholders.size(),
nodes.size());
int node_index = 0;
for (int i = 0; i < placeholders.size(); ++i, ++node_index) {
const NodeDef* node = nodes[i].first;
EXPECT_EQ(absl::StrCat("p", i), node->name());
EXPECT_EQ(i + 1, nodes[i].second);
EXPECT_EQ(i + 1, node_to_start_time[node]);
}
for (int i = 0; i < while_exits.size(); ++i, node_index += nodes_per_loop) {
const NodeDef* node = nodes[node_index].first;
EXPECT_EQ(absl::StrCat("while", i, "/Enter"), node->name());
EXPECT_EQ(100 + i * 10, nodes[node_index].second);
EXPECT_EQ(100 + i * 10, node_to_start_time[node]);
}
for (int i = 0; i < squares.size(); ++i, ++node_index) {
int square_index = num_leaves - 1 - i;
const NodeDef* node = nodes[node_index].first;
EXPECT_EQ(absl::StrCat("s", square_index), node->name());
EXPECT_EQ(500 - (square_index + 1), nodes[node_index].second);
EXPECT_EQ(500 - (square_index + 1), node_to_start_time[node]);
}
}
} // namespace
} // namespace tensorflow
tensorflow/core/graph/graph_test.cc
+1021
View File
File diff suppressed because it is too large Load Diff
tensorflow/core/graph/mkl_graph_util.h
+284
View File
@@ -0,0 +1,284 @@
/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
Licensed 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 TENSORFLOW_CORE_GRAPH_MKL_GRAPH_UTIL_H_
#define TENSORFLOW_CORE_GRAPH_MKL_GRAPH_UTIL_H_
#ifdef INTEL_MKL
#include "absl/base/call_once.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/util/env_var.h"
#include "tensorflow/core/util/util.h"
namespace tensorflow {
// Since our ops are going to produce and also consume N addition tensors
// (Mkl) for N Tensorflow tensors, we can have following different
// orderings among these 2N tensors.
//
// E.g., for Tensorflow tensors A, B, and C, our ops will produce and
// consume A_m, B_m, and C_m additionally.
//
// INTERLEAVED: in this case 2N tensors are interleaved. So for above
// example, the ordering looks like: A, A_m, B, B_m, C, C_m.
//
// CONTIGUOUS: in thi case N Tensorflow tensors are contiguous followed
// by N Mkl tensors. So for above example, the ordering looks
// like: A, B, C, A_m, B_m, C_m
//
// Following APIs map index of original Tensorflow tensors to their
// appropriate position based on selected ordering. For contiguous ordering,
// we need to know the total number of tensors (parameter total).
//
typedef enum { TENSORS_INTERLEAVED, TENSORS_CONTIGUOUS } MklTfTensorOrdering;
// NOTE: Currently, we use contiguous ordering. If you change this, then you
// would need to change Mkl op definitions in nn_ops.cc.
static const MklTfTensorOrdering kTensorOrdering = TENSORS_CONTIGUOUS;
// Get index of MetaData tensor from index 'n' of Data tensor.
inline int DataIndexToMetaDataIndex(int n, int total_tensors) {
if (kTensorOrdering == MklTfTensorOrdering::TENSORS_INTERLEAVED) {
// For interleaved ordering, Mkl tensor follows immediately after
// Tensorflow tensor.
return n + 1;
} else {
CHECK_EQ(kTensorOrdering, MklTfTensorOrdering::TENSORS_CONTIGUOUS);
// For contiguous ordering, Mkl tensor is n+total_tensors / 2 away.
return n + total_tensors / 2;
}
}
int inline GetTensorDataIndex(int n, int total_tensors) {
if (kTensorOrdering == MklTfTensorOrdering::TENSORS_INTERLEAVED) {
return 2 * n; // index corresponding to nth input/output tensor
} else {
CHECK_EQ(kTensorOrdering, MklTfTensorOrdering::TENSORS_CONTIGUOUS);
return n;
}
}
int inline GetTensorMetaDataIndex(int n, int total_tensors) {
// Get index for TensorData first and then use mapping function
// to get TensorMetaData index from TensorData index.
int tidx = GetTensorDataIndex(n, total_tensors);
return DataIndexToMetaDataIndex(tidx, total_tensors);
}
// check if the control between src and dst nodes already exists
bool inline DoesControlEdgeExist(const Node* src, const Node* dst) {
for (const Edge* edge : src->out_edges()) {
if (edge->IsControlEdge() && edge->dst() == dst) {
return true;
}
}
return false;
}
// In TF 2.8, oneDNN blocked format will not be supported.
// TODO(intel_tf): Cleanup shall be done in future:
// (1) Remove this method;
// (2) Update related code wherever it is called.
bool inline NativeFormatEnabled() { return true; }
// Check if the data_format attribute in the node def represents 5D tensor
bool inline Check5DFormat(const NodeDef& ndef) {
string data_format;
TF_CHECK_OK(GetNodeAttr(ndef, "data_format", &data_format));
if (data_format.compare("NCDHW") == 0 || data_format.compare("NDHWC") == 0) {
return true;
}
return false;
}
namespace mkl_op_registry {
// MKL operators whose kernels are registered with 'MklLayoutDependentOp' label
// (e.g., MklConv2D) understand input tensors in MKL layout. These operators
// get additional meta-tensors for actual input tensors.
static const char* kMklLayoutDependentOpLabel = "MklLayoutDependentOp";
static const char* kMklLayoutDependentOpLabelPattern =
"label='MklLayoutDependentOp'";
// MKL operators whose kernels are registered with 'MklNameChangeOp' label
// (e.g., MklMatMul, MklTranspose) do not understand input tensors in MKL
// layout. These operators do not get additional meta-tensors. The signatures of
// these operators are the same as the original TensorFlow operators that they
// correspond to. So these ops just go through a name change during graph
// rewrite pass.
static const char* kMklNameChangeOpLabel = "MklNameChangeOp";
static const char* kMklNameChangeOpLabelPattern = "label='MklNameChangeOp'";
static const char* kMklQuantizedOpLabel = "QuantizedMklOp";
static const char* kMklQuantizedOpLabelPattern = "label='QuantizedMklOp'";
// Prefix that we add to Tensorflow op name to construct Mkl op name.
static const char* const kMklOpPrefix = "_Mkl";
// TODO(intel-tf): PR review feedback (penpornk)
// Can we add eager_mode (or is_eager) as an op attribute instead?
// This way we don't need to rename the op just to pass eager_mode
// through template parameter.
static const char* const kMklEagerOpPrefix = "_MklEager";
// Prefix that we add to TF op name to construct MKL op that does not
// depend on layout propagation. It will be used in both Eager and graph
// modes unless there is a reason to have additional op name with
// _MklEager prefix.
static const char* const kMklNativeOpPrefix = "_MklNative";
// Get the name of Mkl Native (does not depend on layout propagation) op
// from original TensorFlow op.
inline string GetMklNativeOpName(const string& name) {
// There are few operators that don't depend on layout propagation but are
// prefixed with _Mkl instead of _MklNative.
bool result =
(0 == name.compare("ConjugateTranspose") ||
0 == name.compare("SparseTensorDenseMatMul") ||
0 == name.compare("BatchMatMul") || 0 == name.compare("BatchMatMulV2") ||
0 == name.compare("Einsum") || 0 == name.compare("MatMul") ||
0 == name.compare("Transpose") || 0 == name.compare("QuantizeV2") ||
0 == name.compare("Dequantize") || 0 == name.compare("Softmax") ||
0 == name.rfind("Quantized", 0));
if (result) {
return string(kMklOpPrefix) + name;
} else {
return string(kMklNativeOpPrefix) + name;
}
}
// Get the name of Mkl op from original TensorFlow op
// We prefix the original op with _Mkl or _MklNative to get Mkl op.
inline string GetMklOpName(const string& name) {
if (!NativeFormatEnabled()) {
return string(kMklOpPrefix) + name;
} else {
return GetMklNativeOpName(name);
}
}
// Get the name of Mkl Eager op from original TensorFlow op
// We prefix 'MklEager' to the original op to get Mkl Eager op.
inline string GetMklEagerOpName(const string& name) {
return string(kMklEagerOpPrefix) + name;
}
// Check whether opname with type T is registered as MKL operator
// that will go through name change or layout change pass.
//
// @input: name of the op
// @input: T datatype to be used for checking op
// @return: true if opname is registered as MKL op that will go through name
// change or layout change pass; false otherwise
static inline bool IsMklOp(const string& op_name, DataType T,
bool is_native_op) {
string label = is_native_op ? kMklNameChangeOpLabelPattern
: kMklLayoutDependentOpLabelPattern;
string registered_kernels_key = op_name + label + std::to_string(T);
thread_local static auto registered_kernels_map =
std::make_unique<absl::flat_hash_map<string, bool>>();
auto kernel_element = registered_kernels_map->find(registered_kernels_key);
bool kernel_registered = false;
if (kernel_element == registered_kernels_map->end()) {
string registered_kernels = KernelsRegisteredForOp(op_name);
// String returned by KernelsRegisteredForOp looks like below:
//
// Op = _MklMatMul, kernels =
// device='CPU'; label='MklNameChangeOp'; T in [DT_COMPLEX128]
// device='CPU'; label='MklNameChangeOp'; T in [DT_COMPLEX64]
// device='CPU'; label='MklNameChangeOp'; T in [DT_DOUBLE]
// device='CPU'; label='MklNameChangeOp'; T in [DT_FLOAT]
if (is_native_op &&
registered_kernels.find(kMklQuantizedOpLabelPattern) != string::npos) {
// Restrict quantized ops to QUINT8, QINT8 and DT_QINT32
kernel_registered = (T == DT_QUINT8 || T == DT_QINT8 || T == DT_QINT32);
}
// Now we just construct a search string to match what we are looking for.
string search_string =
label + string("; T in [") + DataType_Name(T) + string("]");
if (registered_kernels.find(search_string) != string::npos) {
kernel_registered = is_native_op
? (T == DT_COMPLEX128 || T == DT_COMPLEX64 ||
T == DT_DOUBLE || T == DT_FLOAT)
: T == DT_FLOAT;
if (!kernel_registered) {
if ((T == DT_BFLOAT16 || T == DT_HALF) &&
IsDataTypeSupportedByOneDNNOnThisCPU(T)) {
kernel_registered = true;
} else {
DataTypeUnsupportedWarning(T);
}
}
}
registered_kernels_map->insert(
std::make_pair(registered_kernels_key, kernel_registered));
} else {
// Kernel is visited at least once. Return stored registration result.
kernel_registered = kernel_element->second;
}
return kernel_registered;
}
// TODO(intel-tf): QuantizedConv2D is registered with input: QUINT8
// filter:QINT8 for oneDNN integration. First a dummy kernel is created
// and then it is replaced by an actual kernel.
static inline bool IsMklQuantizedOp(const string& op_name, DataType Tinput,
DataType Tfilter) {
// Restrict quantized ops to QUINT8 and QINT8 for now
if (IsMklOp(op_name, Tinput, kMklQuantizedOpLabelPattern)) {
return (Tfilter == DT_QINT8);
}
return false;
}
// Check if the operator with 'op_name' and type 'T' is an MKL operator that
// will either understand input tensors in MKL layout or will go through name
// rewrite that some operators go through.
static inline bool IsMklOp(const string& op_name, DataType T) {
return IsMklOp(op_name, T, true) || IsMklOp(op_name, T, false);
}
static inline bool IsMklOp(const Node* n) {
DataType T;
return GetNodeAttr(n->def(), "T", &T).ok() && IsMklOp(n->type_string(), T);
}
// Check whether opname with type T is registered as MKL-compliant and
// is element-wise.
//
// @input: name of the op
// @input: T datatype to be used for checking op
// @return: true if opname is registered as element-wise Mkl op;
// false otherwise
static inline bool IsMklElementWiseOp(const string& op_name, DataType T) {
if (!IsMklOp(op_name, T)) {
return false;
}
bool result = (0 == op_name.compare(GetMklOpName("Add")) ||
0 == op_name.compare(GetMklOpName("AddV2")) ||
0 == op_name.compare(GetMklOpName("Sub")) ||
0 == op_name.compare(GetMklOpName("Mul")) ||
0 == op_name.compare(GetMklOpName("Maximum")) ||
0 == op_name.compare(GetMklOpName("Sigmoid")) ||
0 == op_name.compare(GetMklOpName("SquaredDifference")));
return result;
}
} // namespace mkl_op_registry
} // namespace tensorflow
#endif // INTEL_MKL
#endif // TENSORFLOW_CORE_GRAPH_MKL_GRAPH_UTIL_H_
tensorflow/core/graph/mkl_testlib.cc
+52
View File
@@ -0,0 +1,52 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#ifdef INTEL_MKL
#include "tensorflow/core/graph/mkl_testlib.h"
#include "tensorflow/core/graph/mkl_graph_util.h"
#include "tensorflow/core/graph/node_builder.h"
namespace tensorflow {
namespace test {
namespace graph {
Node* oneDNNSoftmax(Graph* g, Node* input) {
Node* ret = nullptr;
TF_CHECK_OK(NodeBuilder(g->NewName("n"), "_MklSoftmax")
.Input(input)
.Attr("_kernel", mkl_op_registry::kMklNameChangeOpLabel)
.Finalize(g, &ret));
return ret;
}
#ifdef ENABLE_ONEDNN_V3
Node* oneDNNSparseCSRMatmul(Graph* g, Node* csr_matrix_t, Node* b) {
Node* ret = nullptr;
TF_CHECK_OK(NodeBuilder(g->NewName("n"), "_MklNativeSparseMatrixMatMul")
.Input(csr_matrix_t)
.Input(b)
.Attr("T", DT_FLOAT)
.Attr("_kernel", mkl_op_registry::kMklNameChangeOpLabel)
.Finalize(g, &ret));
return ret;
}
#endif // ENABLE_ONEDNN_V3
} // namespace graph
} // namespace test
} // namespace tensorflow
#endif // INTEL_MKL
tensorflow/core/graph/mkl_testlib.h
+38
View File
@@ -0,0 +1,38 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed 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 TENSORFLOW_CORE_GRAPH_MKL_TESTLIB_H_
#define TENSORFLOW_CORE_GRAPH_MKL_TESTLIB_H_
#ifdef INTEL_MKL
#include "tensorflow/core/graph/graph.h"
namespace tensorflow {
namespace test {
namespace graph {
Node* oneDNNSoftmax(Graph* g, Node* input);
#ifdef ENABLE_ONEDNN_V3
Node* oneDNNSparseCSRMatmul(Graph* g, Node* csr_matrix_t, Node* b);
#endif // ENABLE_ONEDNN_V3
} // namespace graph
} // namespace test
} // namespace tensorflow
#endif // INTEL_MKL
#endif // TENSORFLOW_CORE_GRAPH_MKL_TESTLIB_H_
tensorflow/core/graph/node_builder.cc
+180
View File
@@ -0,0 +1,180 @@
/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#include "tensorflow/core/graph/node_builder.h"
#include <unordered_map>
#include <vector>
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/framework/versions.pb.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/protobuf/error_codes.pb.h"
namespace tensorflow {
NodeBuilder::NodeOut::NodeOut(Node* n, int32_t i) // NOLINT(runtime/explicit)
: node(n),
error(false),
name(node != nullptr ? node->name() : (error = true, "")),
index(i),
dt(SafeGetOutput(node, i, &error)) {}
NodeBuilder::NodeOut::NodeOut(OutputTensor t) : NodeOut(t.node, t.index) {}
NodeBuilder::NodeOut::NodeOut(absl::string_view n, int32_t i, DataType t)
: node(nullptr), error(false), name(n), index(i), dt(t) {}
NodeBuilder::NodeOut::NodeOut()
: node(nullptr), error(true), index(0), dt(DT_FLOAT) {}
NodeBuilder::NodeBuilder(absl::string_view name, absl::string_view op_name,
const OpRegistryInterface* op_registry,
const NodeDebugInfo* debug)
: def_builder_(name, op_name, op_registry, debug) {}
NodeBuilder::NodeBuilder(absl::string_view name, const OpDef* op_def)
: def_builder_(name, op_def) {}
NodeBuilder::NodeBuilder(const NodeDefBuilder& def_builder)
: def_builder_(def_builder) {}
NodeBuilder& NodeBuilder::Input(Node* src_node, int src_index) {
inputs_.emplace_back(src_node, src_index);
DataType dt;
if (GetOutputType(src_node, src_index, &dt)) {
def_builder_.Input(src_node->name(), src_index, dt);
}
return *this;
}
NodeBuilder& NodeBuilder::Input(NodeOut src) {
if (src.error) {
AddIndexError(src.node, src.index);
} else {
inputs_.emplace_back(src.node, src.index);
def_builder_.Input(src.name, src.index, src.dt);
}
return *this;
}
NodeBuilder& NodeBuilder::Input(absl::Span<const NodeOut> src_list) {
std::vector<NodeDefBuilder::NodeOut> srcs;
srcs.reserve(src_list.size());
for (const auto& node_out : src_list) {
if (node_out.error) {
AddIndexError(node_out.node, node_out.index);
} else {
srcs.emplace_back(node_out.name, node_out.index, node_out.dt);
inputs_.emplace_back(node_out.node, node_out.index);
}
}
def_builder_.Input(absl::Span<const NodeDefBuilder::NodeOut>(srcs));
return *this;
}
NodeBuilder& NodeBuilder::ControlInput(Node* src_node) {
control_inputs_.emplace_back(src_node);
def_builder_.ControlInput(src_node->name());
return *this;
}
NodeBuilder& NodeBuilder::ControlInputs(absl::Span<Node* const> src_nodes) {
control_inputs_.insert(control_inputs_.end(), src_nodes.begin(),
src_nodes.end());
for (const Node* src_node : src_nodes) {
def_builder_.ControlInput(src_node->name());
}
return *this;
}
NodeBuilder& NodeBuilder::Device(absl::string_view device_spec) {
def_builder_.Device(device_spec);
return *this;
}
NodeBuilder& NodeBuilder::AssignedDevice(absl::string_view device) {
assigned_device_ = std::string(device);
return *this;
}
NodeBuilder& NodeBuilder::XlaCluster(absl::string_view xla_cluster) {
def_builder_.Attr("_XlaCluster", xla_cluster);
return *this;
}
absl::StatusOr<Node*> NodeBuilder::Finalize(Graph* graph, bool consume) {
Node* out;
TF_RETURN_IF_ERROR(Finalize(graph, &out, consume));
return out;
}
absl::Status NodeBuilder::Finalize(Graph* graph, Node** created_node,
bool consume) {
// In case of error, set *created_node to nullptr.
if (created_node != nullptr) {
*created_node = nullptr;
}
if (!errors_.empty()) {
return absl::InvalidArgumentError(absl::StrJoin(errors_, "\n"));
}
NodeDef node_def;
TF_RETURN_IF_ERROR(def_builder_.Finalize(&node_def, consume));
TF_RETURN_IF_ERROR(ValidateNodeDef(node_def, def_builder_.op_def()));
TF_RETURN_IF_ERROR(
CheckOpDeprecation(def_builder_.op_def(), graph->versions().producer()));
TF_ASSIGN_OR_RETURN(Node * node, graph->AddNode(std::move(node_def)));
node->set_assigned_device_name(assigned_device_);
for (size_t i = 0; i < inputs_.size(); ++i) {
if (inputs_[i].node != nullptr) { // Skip back edges.
graph->AddEdge(inputs_[i].node, inputs_[i].index, node, i);
}
}
for (Node* control_input : control_inputs_) {
graph->AddControlEdge(control_input, node);
}
if (created_node != nullptr) *created_node = node;
return absl::OkStatus();
}
void NodeBuilder::AddIndexError(const Node* node, int i) {
if (node == nullptr) {
errors_.emplace_back(
absl::StrCat("Attempt to add nullptr Node to node with type ",
def_builder_.op_def().name()));
} else {
errors_.emplace_back(strings::StrCat(
"Attempt to add output ", i, " of ", node->name(), " not in range [0, ",
node->num_outputs(), ") to node with type ",
def_builder_.op_def().name(), ". Node: ", FormatNodeForError(*node)));
}
}
bool NodeBuilder::GetOutputType(const Node* node, int i, DataType* dt) {
bool error;
*dt = SafeGetOutput(node, i, &error);
if (error) AddIndexError(node, i);
return !error;
}
} // namespace tensorflow
tensorflow/core/graph/node_builder.h
+181
View File
@@ -0,0 +1,181 @@
/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
Licensed 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 TENSORFLOW_CORE_GRAPH_NODE_BUILDER_H_
#define TENSORFLOW_CORE_GRAPH_NODE_BUILDER_H_
#include <vector>
#include "tensorflow/core/framework/node_def_builder.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_def.pb.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/core/stringpiece.h"
#include "tensorflow/core/lib/gtl/array_slice.h"
namespace tensorflow {
// This is a helper for creating a Node and adding it to a Graph.
// Internally, it uses a NodeDefBuilder to automatically set attrs
// that can be inferred from the inputs, and use default values
// (where they exist) for unspecified attrs. Example usage:
//
// Node* node;
// Status status = NodeBuilder(node_name, op_name)
// .Input(...)
// .Attr(...)
// .Finalize(&graph, &node);
// if (!status.ok()) return status;
// // Use node here.
class NodeBuilder {
public:
// For specifying the output of a Node to provide to one of the Input()
// functions below. It supports both regular inputs (where you are
// connecting to an existing Node*), and inputs from outside the graph
// (or haven't been added to the graph yet, like back edges, where
// you don't have a Node*). Both types can be mixed, e.g. in an
// ArraySlice.
struct NodeOut {
// For referencing an existing Node.
NodeOut(Node* n, int32_t i = 0);
NodeOut(OutputTensor t);
// For referencing Nodes not in the graph being built. It is
// useful when preparing a graph for ExtendSession or creating a
// back edge to a node that hasn't been added to the graph yet,
// but will be.
NodeOut(absl::string_view name, int32_t i, DataType t);
// Default constructor for std::vector<NodeOut>.
NodeOut();
Node* node;
// error is set to true if:
// * the NodeOut was default constructed and never overwritten,
// * a nullptr Node* was passed to the NodeOut constructor, or
// * an out-of-range index was passed to the NodeOut constructor.
bool error;
std::string name;
int32_t index;
DataType dt;
};
// Specify the name and the Op (either via an OpDef or the name of
// the Op plus a registry) for the Node. Other fields are
// specified by calling the methods below.
// REQUIRES: The OpDef must satisfy ValidateOpDef().
NodeBuilder(absl::string_view name, absl::string_view op_name,
const OpRegistryInterface* op_registry = OpRegistry::Global(),
const NodeDebugInfo* debug = nullptr);
NodeBuilder(absl::string_view name, const OpDef* op_def);
// Create a NodeBuilder from an existing NodeDefBuilder.
NodeBuilder(const NodeDefBuilder& def_builder);
// You must call one Input() function per input_arg in the Op,
// *and in the same order as the input_args appear in the OpDef.*
// For inputs that take a single tensor.
NodeBuilder& Input(Node* src_node, int src_index = 0);
NodeBuilder& Input(NodeOut src);
// For inputs that take a list of tensors.
NodeBuilder& Input(absl::Span<const NodeOut> src_list);
// Require that this node run after src_node(s).
NodeBuilder& ControlInput(Node* src_node);
NodeBuilder& ControlInputs(absl::Span<Node* const> src_nodes);
// Sets the "requested device spec" in the NodeDef (not the
// "assigned device" in the Node).
NodeBuilder& Device(absl::string_view device_spec);
// Sets the device name in the "assigned device" field in tensorflow::Node.
NodeBuilder& AssignedDevice(absl::string_view device);
// Sets the _XlaCluster attribute in created node to `xla_cluster`.
NodeBuilder& XlaCluster(absl::string_view xla_cluster);
// Set the value of an attr. attr_name must match the name of one of
// attrs defined by the Op, and value must have the corresponding type
// (see SetAttrValue() in ../framework/attr_value_util.h for legal
// types for value). Note that attrs will be set automatically if
// they can be determined by the inputs.
template <class T>
NodeBuilder& Attr(absl::string_view attr_name, T&& value);
template <class T>
NodeBuilder& Attr(absl::string_view attr_name,
std::initializer_list<T> value);
// Validates the described node and adds it to *graph, adding edges
// for all (non-back) inputs. If created_node is not nullptr,
// *created_node will be set to the new node (or nullptr on error).
// If `consume` is true, the builder state will be moved into `node_def`,
// and the builder will be left in an undefined state.
absl::Status Finalize(Graph* graph, Node** created_node,
bool consume = false);
// Same as `Finalize` above, but using StatusOr to return value. Preferred
// form.
absl::StatusOr<Node*> Finalize(Graph* graph, bool consume = false);
// Accessors for the values set in the constructor.
const std::string& node_name() const { return def_builder_.node_name(); }
const OpDef& op_def() const { return def_builder_.op_def(); }
private:
static DataType SafeGetOutput(const Node* node, int i, bool* error) {
if (node != nullptr && i >= 0 && i < node->num_outputs()) {
*error = false;
return node->output_type(i);
} else {
*error = true;
return DT_FLOAT;
}
}
// If SafeGetOutput indicates a range error, add it to errors_.
void AddIndexError(const Node* node, int i);
// Set *dt and returns true if i is in range. Combines
// SafeGetOutput() and AddIndexError().
bool GetOutputType(const Node* node, int i, DataType* dt);
NodeDefBuilder def_builder_;
const OpRegistryInterface* op_registry_;
std::vector<NodeOut> inputs_;
std::vector<Node*> control_inputs_;
std::vector<std::string> errors_;
std::string assigned_device_;
};
// IMPLEMENTATION -------------------------------------------------------------
template <class T>
NodeBuilder& NodeBuilder::Attr(absl::string_view attr_name, T&& value) {
def_builder_.Attr(attr_name, std::forward<T>(value));
return *this;
}
template <class T>
NodeBuilder& NodeBuilder::Attr(absl::string_view attr_name,
std::initializer_list<T> value) {
def_builder_.Attr(attr_name, value);
return *this;
}
} // namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPH_NODE_BUILDER_H_
tensorflow/core/graph/node_builder_test.cc
+131
View File
@@ -0,0 +1,131 @@
/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#include "tensorflow/core/graph/node_builder.h"
#include <string>
#include "tensorflow/core/framework/full_type.pb.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_def_builder.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/kernels/ops_util.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace {
REGISTER_OP("Source").Output("o: out_types").Attr("out_types: list(type)");
REGISTER_OP("Sink").Input("i: T").Attr("T: type");
TEST(NodeBuilderTest, Simple) {
Graph graph(OpRegistry::Global());
Node* source_node;
TF_EXPECT_OK(NodeBuilder("source_op", "Source")
.Attr("out_types", {DT_INT32, DT_STRING})
.Finalize(&graph, &source_node));
ASSERT_TRUE(source_node != nullptr);
// Try connecting to each of source_node's outputs.
TF_EXPECT_OK(NodeBuilder("sink1", "Sink")
.Input(source_node)
.Finalize(&graph, nullptr));
TF_EXPECT_OK(NodeBuilder("sink2", "Sink")
.Input(source_node, 1)
.Finalize(&graph, nullptr));
// Generate an error if the index is out of range.
EXPECT_FALSE(NodeBuilder("sink3", "Sink")
.Input(source_node, 2)
.Finalize(&graph, nullptr)
.ok());
EXPECT_FALSE(NodeBuilder("sink4", "Sink")
.Input(source_node, -1)
.Finalize(&graph, nullptr)
.ok());
EXPECT_FALSE(NodeBuilder("sink5", "Sink")
.Input({source_node, -1})
.Finalize(&graph, nullptr)
.ok());
// Generate an error if the node is nullptr. This can happen when using
// GraphDefBuilder if there was an error creating the input node.
EXPECT_FALSE(NodeBuilder("sink6", "Sink")
.Input(nullptr)
.Finalize(&graph, nullptr)
.ok());
EXPECT_FALSE(NodeBuilder("sink7", "Sink")
.Input(NodeBuilder::NodeOut(nullptr, 0))
.Finalize(&graph, nullptr)
.ok());
}
REGISTER_OP("FullTypeOpBasicType")
.Output("o1: out_type")
.Attr("out_type: type")
.SetTypeConstructor([](OpDef* op_def) {
FullTypeDef* tdef =
op_def->mutable_output_arg(0)->mutable_experimental_full_type();
tdef->set_type_id(TFT_ARRAY);
FullTypeDef* arg = tdef->add_args();
arg->set_type_id(TFT_VAR);
arg->set_s("out_type");
return absl::OkStatus();
});
TEST(NodeBuilderTest, TypeConstructorBasicType) {
Graph graph(OpRegistry::Global());
Node* node;
TF_EXPECT_OK(NodeBuilder("op", "FullTypeOpBasicType")
.Attr("out_type", DT_FLOAT)
.Finalize(&graph, &node));
ASSERT_TRUE(node->def().has_experimental_type());
const FullTypeDef& ft = node->def().experimental_type();
ASSERT_EQ(ft.type_id(), TFT_PRODUCT);
ASSERT_EQ(ft.args_size(), 1);
auto ot = ft.args(0);
ASSERT_EQ(ot.type_id(), TFT_ARRAY);
ASSERT_EQ(ot.args(0).type_id(), TFT_FLOAT);
ASSERT_EQ(ot.args(0).args().size(), 0);
}
REGISTER_OP("FullTypeOpListType")
.Output("o1: out_types")
.Attr("out_types: list(type)")
.SetTypeConstructor([](OpDef* op_def) {
FullTypeDef* tdef =
op_def->mutable_output_arg(0)->mutable_experimental_full_type();
tdef->set_type_id(TFT_ARRAY);
FullTypeDef* arg = tdef->add_args();
arg->set_type_id(TFT_VAR);
arg->set_s("out_types");
return absl::OkStatus();
});
TEST(NodeBuilderTest, TypeConstructorListType) {
Graph graph(OpRegistry::Global());
Node* node;
ASSERT_FALSE(NodeBuilder("op", "FullTypeOpListType")
.Attr("out_types", {DT_FLOAT, DT_INT32})
.Finalize(&graph, &node)
.ok());
}
} // namespace
} // namespace tensorflow
tensorflow/core/graph/optimizer_cse.cc
+338
View File
@@ -0,0 +1,338 @@
/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
// This module implements a common subexpression elimination pass. We
// process the nodes in the graph in reverse postorder
// (i.e. inputs before their downstream dependencies). The rough algorithm is
// as follows:
//
// std::unordered_map<size_t, Node*> available
// for each node n in forward topological order:
// h = NodeHash(n)
// if available[h] exists and Equivalent(available(h), h)
// redirect downstream uses of outputs of n to available[h]
// remove n from graph
// else
// if available[h] does not exist
// available[h] = n
//
// This is similar to the global value number algorithm describe in this
// paper:
//
// "Global code motion/global value numbering", Cliff Click, PLDI '95
// Proceedings of the ACM SIGPLAN 1995 conference on Programming
// language design and implementation, Pages 246-257
// http://dl.acm.org/citation.cfm?id=207154
#include "tensorflow/core/graph/optimizer_cse.h"
#include <iostream>
#include <unordered_map>
#include <utility>
#include <vector>
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/graph/algorithm.h"
#include "tensorflow/core/graph/graph_node_util.h"
#include "tensorflow/core/lib/gtl/map_util.h"
#include "tensorflow/core/lib/hash/hash.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/protobuf.h"
namespace tensorflow {
class OptimizerCSE {
public:
explicit OptimizerCSE(Graph* g) : g_(g) {}
bool Optimize(const std::function<bool(const Node*)>& consider_fn);
private:
static size_t NodeHash(const Node* n);
static bool Equivalent(const Node* a, const Node* b,
AttrSlice::Scratch* scratch);
Graph* g_;
};
static void FillInputs(
const Node* n, absl::InlinedVector<const Node*, 4UL>* control_edges,
absl::InlinedVector<std::pair<const Node*, int>, 4UL>* in) {
DCHECK_EQ(in->size(), n->num_inputs());
control_edges->clear();
for (const Edge* e : n->in_edges()) {
if (e->IsControlEdge()) {
control_edges->push_back(e->src());
} else {
(*in)[e->dst_input()] = std::make_pair(e->src(), e->src_output());
}
}
std::sort(control_edges->begin(), control_edges->end());
if (n->op_def().is_commutative()) {
// For commutative inputs, we sort the input by the input Node*
// to get a canonical ordering (so that add(a,b) and add(b, a) will
// hash to the same value if is_commutative is true for 'add').
std::sort(in->begin(), in->end());
}
}
static size_t kIllegalNodeHash = 0;
class Hasher {
public:
uint64_t hash() { return h_ == kIllegalNodeHash ? kIllegalNodeHash + 1 : h_; }
void MixString(const std::string& s) { h_ = Hash64(s.data(), s.size(), h_); }
void MixInteger(size_t z) { h_ = Hash64Combine(h_, z); }
void MixProto(const protobuf::MessageLite& msg) {
msg.ByteSizeLong(); // Ensure sizes are cached accurately.
HashingOutputStream hasher;
{
// CodedOutputStream doesn't call BackUp until it's destroyed, so we need
// it to be destroyed before we call hasher.hash().
protobuf::io::CodedOutputStream stream(&hasher);
stream.EnableAliasing(true);
stream.SetSerializationDeterministic(true);
msg.SerializeWithCachedSizes(&stream);
}
h_ = Hash64Combine(h_, hasher.hash());
}
private:
// HashingOutputStream produces the same exact hash as if you serialized the
// proto and hashed it sequentially in kBufSize chunks, except it doesn't
// manifest the entire proto into memory at any point.
class HashingOutputStream : public protobuf::io::ZeroCopyOutputStream {
public:
// This kBufSize makes sizeof(HashingOutputStream) == 256. It's not chosen
// for any particular reason except it's a nice even number of cache lines.
static constexpr size_t kBufSize = 228;
static constexpr uint64_t kDefaultSeed = 2570847921467975139ULL;
bool Next(void** data, int* size) override {
if (i_ == kBufSize) {
// Mix the chunk in.
Mix(buf_, kBufSize);
*data = buf_;
*size = kBufSize;
} else {
*data = buf_ + i_;
*size = kBufSize - i_;
}
// We always set i_ to be past the end, since we've given the rest of buf_
// out.
i_ = kBufSize;
return true;
}
void BackUp(int count) override { i_ -= count; }
int64_t ByteCount() const override { return byte_count_; }
bool WriteAliasedRaw(const void* void_data, int size) override {
// We can't do math on void*.
const char* data = static_cast<const char*>(void_data);
const auto remaining = kBufSize - i_;
if (remaining > 0) {
if (size < remaining) {
memcpy(buf_ + i_, data, size);
i_ += size;
return true;
}
memcpy(buf_ + i_, data, remaining);
i_ = kBufSize;
data += remaining;
size -= remaining;
}
if (i_ == kBufSize) {
Mix(buf_, kBufSize);
i_ = 0;
}
while (size >= kBufSize) {
Mix(data, kBufSize);
data += kBufSize;
size -= kBufSize;
}
memcpy(buf_, data, size);
i_ = size;
return true;
}
bool AllowsAliasing() const override { return true; }
uint64_t hash() {
if (i_ != 0) {
Mix(buf_, i_);
i_ = 0;
}
return h_;
}
private:
void Mix(const char* p, size_t n) {
byte_count_ += n;
h_ = Hash64(p, n, h_);
}
char buf_[kBufSize];
int i_ = 0;
int64_t byte_count_ = 0;
uint64_t h_ = kDefaultSeed;
};
uint64_t h_ = HashingOutputStream::kDefaultSeed;
};
size_t OptimizerCSE::NodeHash(const Node* n) {
Hasher hasher;
hasher.MixString(n->type_string());
hasher.MixInteger(n->output_types().size());
for (DataType dt : n->output_types()) {
hasher.MixInteger(dt);
}
hasher.MixInteger(n->num_inputs());
absl::InlinedVector<const Node*, 4UL> control_edges;
absl::InlinedVector<std::pair<const Node*, int>, 4UL> in(n->num_inputs());
FillInputs(n, &control_edges, &in);
for (const auto& edge : in) {
hasher.MixInteger(edge.first->id());
hasher.MixInteger(edge.second);
}
#if !defined(__ANDROID__)
// Hash the attrs. For example, this makes sure different constants
// end up in different hash buckets.
size_t attr_hashes = 0;
for (const auto& attr : n->attrs()) {
Hasher h;
h.MixString(attr.first);
h.MixProto(attr.second);
attr_hashes = Hash64CombineUnordered(attr_hashes, h.hash());
}
hasher.MixInteger(attr_hashes);
#endif
return hasher.hash();
}
static bool HasRefInput(const Node* n) {
for (auto dt : n->input_types()) {
if (IsRefType(dt)) return true;
}
return false;
}
bool OptimizerCSE::Equivalent(const Node* a, const Node* b,
AttrSlice::Scratch* scratch) {
// Different op names are different
if (a->type_string() != b->type_string()) return false;
// Never consider stateful nodes (such as non-const inputs) equivalent.
if (a->op_def().is_stateful()) return false;
// For now, we consider any node that takes a ref input to not be
// equivalent to any other node.
if (HasRefInput(a) || HasRefInput(b)) return false;
// Compare attrs. Note that equal attrs implies equal input and
// output types.
if (!a->attrs().EqualAttrs(b->attrs(), scratch)) return false;
// Compare input sources
if (a->num_inputs() != b->num_inputs()) return false;
const int N_in = a->num_inputs();
absl::InlinedVector<const Node*, 4UL> a_control_edges;
absl::InlinedVector<const Node*, 4UL> b_control_edges;
absl::InlinedVector<std::pair<const Node*, int>, 4UL> a_in(N_in);
absl::InlinedVector<std::pair<const Node*, int>, 4UL> b_in(N_in);
FillInputs(a, &a_control_edges, &a_in);
FillInputs(b, &b_control_edges, &b_in);
if (a_in != b_in) return false;
if (a_control_edges != b_control_edges) return false;
return true;
}
bool OptimizerCSE::Optimize(
const std::function<bool(const Node*)>& consider_fn) {
// This very simple implementation works if the whole graph is one
// giant basic block (because we just traverse nodes in a
// topological order). This simple implementation works well
// with control flow/loops/etc. But we need to be careful about
// control flow if we want to add more sophisticated CSE optimizations.
// TODO(jeff): We need to handle Update nodes specially, but dealing
// with more general control flow will also solve this issue, and for
// now, our updates are almost always the most downstream nodes in
// the graph.
std::vector<Node*> order;
GetReversePostOrder(*g_, &order, NodeComparatorID());
// Our value is just a single Node*, meaning we keep just a single
// candidate for a given node hash value. This may cause us to
// (rarely) lose some optimization opportunities if there are
// hash collisions, but it allows us to avoid having the value
// be a set<Node*> (or equivalent).
std::unordered_map<size_t, Node*> available;
// Scratch space for Equivalent calls. Allocated here and passed in to
// Equivalent to avoid allocation inside the loop below.
bool changed = false;
AttrSlice::Scratch scratch;
for (Node* n : order) {
if (!n->IsOp()) continue;
// Don't prune placeholder nodes.
if (n->type_string() == "Placeholder" ||
n->type_string() == "PlaceholderV2" ||
n->type_string() == "PlaceholderWithDefault") {
continue;
}
// See if we should consider this node at all
if (consider_fn != nullptr && !consider_fn(n)) continue;
size_t h = NodeHash(n);
Node** candidate = &available[h];
if (*candidate == nullptr) {
// No existing match: insert "n" into the hash table under "h"
*candidate = n;
} else if (Equivalent(*candidate, n, &scratch)) {
VLOG(1) << "CSE: equivalent: " << (*candidate)->name() << " and "
<< n->name();
// *candidate and n are equivalent. Therefore, we can replace
// n with *candidate by fixing up outgoing edges from "n" to instead
// come from "*candidate", and then delete n from the graph
for (const Edge* e : n->out_edges()) {
g_->AddEdge(*candidate, e->src_output(), e->dst(), e->dst_input());
}
MergeDebugInfo(NodeDebugInfo(*n), *candidate);
g_->RemoveNode(n);
changed = true;
}
}
return changed;
}
bool OptimizeCSE(Graph* g,
const std::function<bool(const Node*)>& consider_fn) {
OptimizerCSE opt(g);
return opt.Optimize(consider_fn);
}
} // namespace tensorflow
tensorflow/core/graph/optimizer_cse.h
+37
View File
@@ -0,0 +1,37 @@
/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
// An optimization pass that performs common subexpression elimination.
#ifndef TENSORFLOW_CORE_GRAPH_OPTIMIZER_CSE_H_
#define TENSORFLOW_CORE_GRAPH_OPTIMIZER_CSE_H_
#include <sys/types.h>
#include "tensorflow/core/graph/graph.h"
namespace tensorflow {
// Perform common-subexpression elimination on the graph "*g". If
// "consider_fn" is not nullptr, then only nodes for which
// consider_fn(node) returns true will be considered for combining
// during the common subexpression elimination.
//
// Returns true if and only if 'g' is mutated.
extern bool OptimizeCSE(Graph* g,
const std::function<bool(const Node*)>& consider_fn);
} // namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPH_OPTIMIZER_CSE_H_
tensorflow/core/graph/optimizer_cse_test.cc
+389
View File
@@ -0,0 +1,389 @@
/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#include "tensorflow/core/graph/optimizer_cse.h"
#include <utility>
#include <vector>
#include "tensorflow/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/graph/testlib.h"
#include "tensorflow/core/kernels/ops_util.h"
#include "tensorflow/core/lib/random/simple_philox.h"
#include "tensorflow/core/lib/strings/str_util.h"
#include "tensorflow/core/lib/strings/stringprintf.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/protobuf.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/test_benchmark.h"
#include "tsl/platform/protobuf.h"
namespace tensorflow {
namespace {
static void InitGraph(const std::string& s, Graph* graph) {
GraphDef graph_def;
auto parser = protobuf::TextFormat::Parser();
// parser.AllowRelaxedWhitespace(true);
CHECK(parser.MergeFromString(s, &graph_def)) << s;
GraphConstructorOptions opts;
TF_CHECK_OK(ConvertGraphDefToGraph(opts, graph_def, graph));
}
class OptimizerCSETest : public ::testing::Test {
public:
OptimizerCSETest() : graph_(OpRegistry::Global()) {}
void InitGraph(const std::string& s) {
::tensorflow::InitGraph(s, &graph_);
original_ = CanonicalGraphString(&graph_);
}
static bool IncludeNode(const Node* n) { return n->IsOp(); }
static std::string EdgeId(const Node* n, int index) {
if (index == 0) {
return n->name();
} else if (index == Graph::kControlSlot) {
return absl::StrCat(n->name(), ":control");
} else {
return absl::StrCat(n->name(), ":", index);
}
}
std::string CanonicalGraphString(Graph* g) {
std::vector<std::string> nodes;
std::vector<std::string> edges;
for (const Node* n : g->nodes()) {
if (IncludeNode(n)) {
nodes.push_back(absl::StrCat(n->name(), "(", n->type_string(), ")"));
}
}
for (const Edge* e : g->edges()) {
if (IncludeNode(e->src()) && IncludeNode(e->dst())) {
edges.push_back(absl::StrCat(EdgeId(e->src(), e->src_output()), "->",
EdgeId(e->dst(), e->dst_input())));
}
}
// Canonicalize
std::sort(nodes.begin(), nodes.end());
std::sort(edges.begin(), edges.end());
return absl::StrCat(absl::StrJoin(nodes, ";"), "|",
absl::StrJoin(edges, ";"));
}
std::string DoCSE(
const std::function<bool(const Node*)>& consider_fn = nullptr) {
std::string before = CanonicalGraphString(&graph_);
LOG(ERROR) << "Before rewrites: " << before;
OptimizeCSE(&graph_, consider_fn);
std::string result = CanonicalGraphString(&graph_);
LOG(ERROR) << "After rewrites: " << result;
return result;
}
const std::string& OriginalGraph() const { return original_; }
Graph graph_;
std::string original_;
};
REGISTER_OP("Input").Output("o: float").SetIsStateful();
// Note that the "rules" in these tests are not meant to be logically correct
TEST_F(OptimizerCSETest, Simple) {
InitGraph(
"node { name: 'A' op: 'Input'}"
"node { name: 'B' op: 'Input'}"
"node { name: 'C' op: 'Mul' attr { key: 'T' value { type: DT_FLOAT } }"
" input: ['A', 'B'] }"
"node { name: 'D' op: 'Mul' attr { key: 'T' value { type: DT_FLOAT } }"
" input: ['A', 'B'] }");
EXPECT_EQ(DoCSE(),
"A(Input);B(Input);C(Mul)|"
"A->C;B->C:1");
}
TEST_F(OptimizerCSETest, Simple_ThreeEquivalent) {
InitGraph(
"node { name: 'A' op: 'Input'}"
"node { name: 'B' op: 'Input'}"
"node { name: 'C' op: 'Mul' attr { key: 'T' value { type: DT_FLOAT } }"
" input: ['A', 'B'] }"
"node { name: 'D' op: 'Mul' attr { key: 'T' value { type: DT_FLOAT } }"
" input: ['A', 'B'] }"
"node { name: 'E' op: 'Mul' attr { key: 'T' value { type: DT_FLOAT } }"
" input: ['A', 'B'] }");
EXPECT_EQ(DoCSE(),
"A(Input);B(Input);C(Mul)|"
"A->C;B->C:1");
}
TEST_F(OptimizerCSETest, Simple_WithFixups) {
InitGraph(
"node { name: 'A' op: 'Input'}"
"node { name: 'B' op: 'Input'}"
"node { name: 'C' op: 'Mul' attr { key: 'T' value { type: DT_FLOAT } }"
" input: ['A', 'B'] }"
"node { name: 'D' op: 'Mul' attr { key: 'T' value { type: DT_FLOAT } }"
" input: ['A', 'B'] }"
"node { name: 'E' op: 'Mul' attr { key: 'T' value { type: DT_FLOAT } }"
" input: ['C', 'D'] }");
EXPECT_EQ(DoCSE(),
"A(Input);B(Input);C(Mul);E(Mul)|"
"A->C;B->C:1;C->E;C->E:1");
}
TEST_F(OptimizerCSETest, Simple_Commutative) {
InitGraph(
"node { name: 'A' op: 'Input'}"
"node { name: 'B' op: 'Input'}"
"node { name: 'C' op: 'Mul' attr { key: 'T' value { type: DT_FLOAT } }"
" input: ['A', 'B'] }"
"node { name: 'D' op: 'Mul' attr { key: 'T' value { type: DT_FLOAT } }"
" input: ['B', 'A'] }");
EXPECT_EQ(DoCSE(),
"A(Input);B(Input);C(Mul)|"
"A->C;B->C:1");
}
static bool IsNotMultiply(const Node* n) { return n->type_string() != "Mul"; }
// Like Simple_Commutative,
TEST_F(OptimizerCSETest, Simple_Filtered) {
InitGraph(
"node { name: 'A' op: 'Input'}"
"node { name: 'B' op: 'Input'}"
"node { name: 'C' op: 'Mul' attr { key: 'T' value { type: DT_FLOAT } }"
" input: ['A', 'B'] }"
"node { name: 'D' op: 'Mul' attr { key: 'T' value { type: DT_FLOAT } }"
" input: ['B', 'A'] }");
EXPECT_EQ(DoCSE(IsNotMultiply), OriginalGraph());
}
TEST_F(OptimizerCSETest, Simple_NotCommutative) {
InitGraph(
"node { name: 'A' op: 'Input'}"
"node { name: 'B' op: 'Input'}"
"node { name: 'C' op: 'Sub' attr { key: 'T' value { type: DT_FLOAT } }"
" input: ['A', 'B'] }"
"node { name: 'D' op: 'Sub' attr { key: 'T' value { type: DT_FLOAT } }"
" input: ['B', 'A'] }");
EXPECT_EQ(DoCSE(), OriginalGraph());
}
TEST_F(OptimizerCSETest, NotEquivalent_Ops) {
InitGraph(
"node { name: 'A' op: 'Input'}"
"node { name: 'B' op: 'Input'}"
"node { name: 'C' op: 'Mul' attr { key: 'T' value { type: DT_FLOAT } }"
" input: ['A', 'B'] }"
"node { name: 'D' op: 'Sub' attr { key: 'T' value { type: DT_FLOAT } }"
" input: ['A', 'B'] }");
EXPECT_EQ(DoCSE(), OriginalGraph());
}
TEST_F(OptimizerCSETest, Simple_SameOps_SameAttrs1) {
// Should still do CSE for ops with attrs if they match.
InitGraph(
"node { name: 'A' op: 'Input'}"
"node { name: 'B' op: 'Input'}"
"node { name: 'C' op: 'Mul' attr { key: 'T' value { type: DT_FLOAT } }"
" input: ['A', 'B'] attr { key: 'shape'"
" value { shape: { dim: { size: 37 name: 'SAME_NAME' } } } } }"
"node { name: 'D' op: 'Mul' attr { key: 'T' value { type: DT_FLOAT } }"
" input: ['A', 'B'] attr { key: 'shape'"
" value { shape: { dim: { size: 37 name: 'SAME_NAME' } } } } }");
EXPECT_EQ(DoCSE(),
"A(Input);B(Input);C(Mul)|"
"A->C;B->C:1");
}
TEST_F(OptimizerCSETest, Simple_SameOps_SameAttrs2) {
// Should still do CSE for ops with attrs if they match, even if they
// are not in the same order.
InitGraph(
"node { name: 'A' op: 'Input'}"
"node { name: 'B' op: 'Input'}"
"node { name: 'C' op: 'Mul' attr { key: 'T' value { type: DT_FLOAT } }"
" input: ['A', 'B']"
" attr { key: 'a' value { i: 3 } }"
" attr { key: 't' value { type: DT_INT32 } } }"
"node { name: 'D' op: 'Mul' attr { key: 'T' value { type: DT_FLOAT } }"
" input: ['A', 'B']"
" attr { key: 't' value { type: DT_INT32 } }"
" attr { key: 'a' value { i: 3 } } }");
EXPECT_EQ(DoCSE(),
"A(Input);B(Input);C(Mul)|"
"A->C;B->C:1");
}
TEST_F(OptimizerCSETest, SameConstants) {
// Should still do CSE for ops with constants if the values are identical
InitGraph(
"node { name: 'A' op: 'Const' "
" attr { key: 'dtype' value { type: DT_INT32 } }"
" attr { key: 'value' value {"
" tensor { dtype: DT_INT32 tensor_shape { dim { size: 1 } } "
" int_val: 0 } } } }"
"node { name: 'B' op: 'Const' "
" attr { key: 'dtype' value { type: DT_INT32 } }"
" attr { key: 'value' value {"
" tensor { dtype: DT_INT32 tensor_shape { dim { size: 1 } } "
" int_val: 0 } } } }"
"node { name: 'D' op: 'Mul' attr { key: 'T' value { type: DT_INT32 } }"
" input: ['A', 'B'] }");
EXPECT_EQ(DoCSE(),
"A(Const);D(Mul)|"
"A->D;A->D:1");
}
TEST_F(OptimizerCSETest, DifferentConstants) {
// Should still do CSE for ops with extensions if the extensions are identical
InitGraph(
"node { name: 'A' op: 'Const' "
" attr { key: 'dtype' value { type: DT_INT32 } }"
" attr { key: 'value' value {"
" tensor { dtype: DT_INT32 tensor_shape { dim { size: 1 } } "
" int_val: 0 } } } }"
"node { name: 'B' op: 'Const' "
" attr { key: 'dtype' value { type: DT_INT32 } }"
" attr { key: 'value' value {"
" tensor { dtype: DT_INT32 tensor_shape { dim { size: 1 } } "
" int_val: 100000 } } } }"
"node { name: 'D' op: 'Mul' attr { key: 'T' value { type: DT_INT32 } }"
" input: ['A', 'B'] }");
EXPECT_EQ(DoCSE(),
"A(Const);B(Const);D(Mul)|"
"A->D;B->D:1");
}
TEST_F(OptimizerCSETest, SameOps_DifferentAttrs1) {
InitGraph(
"node { name: 'A' op: 'Input'}"
"node { name: 'B' op: 'Input'}"
"node { name: 'C' op: 'Mul' attr { key: 'T' value { type: DT_FLOAT } }"
" input: ['A', 'B']"
" attr { key: 'a' value { i: 3 } }"
" attr { key: 't' value { type: DT_INT32 } } }"
"node { name: 'D' op: 'Mul' attr { key: 'T' value { type: DT_FLOAT } }"
" input: ['A', 'B']"
" attr { key: 't' value { type: DT_INT32 } }"
" attr { key: 'a' value { i: 4 } } }");
EXPECT_EQ(DoCSE(), OriginalGraph());
}
TEST_F(OptimizerCSETest, SameOps_DifferentAttrs2) {
InitGraph(
"node { name: 'A' op: 'Input'}"
"node { name: 'B' op: 'Input'}"
"node { name: 'C' op: 'Mul' attr { key: 'T' value { type: DT_FLOAT } }"
" input: ['A', 'B']"
" attr { key: 'a' value { i: 3 } }"
" attr { key: 't' value { type: DT_FLOAT } } }"
"node { name: 'D' op: 'Mul' attr { key: 'T' value { type: DT_FLOAT } }"
" input: ['A', 'B']"
" attr { key: 't' value { type: DT_INT32 } }"
" attr { key: 'a' value { i: 3 } } }");
EXPECT_EQ(DoCSE(), OriginalGraph());
}
TEST_F(OptimizerCSETest, NotEquivalent_Inputs) {
InitGraph(
"node { name: 'A' op: 'Input'}"
"node { name: 'B' op: 'Input'}"
"node { name: 'C' op: 'Input'}"
"node { name: 'D' op: 'Mul' attr { key: 'T' value { type: DT_FLOAT } }"
" input: ['A', 'B'] }"
"node { name: 'E' op: 'Mul' attr { key: 'T' value { type: DT_FLOAT } }"
" input: ['A', 'C'] }");
EXPECT_EQ(DoCSE(), OriginalGraph());
}
TEST_F(OptimizerCSETest, Constant_Dedup) {
Tensor a(DT_FLOAT, TensorShape({1}));
a.flat<float>()(0) = 1.0;
Tensor b(DT_DOUBLE, TensorShape({1})); // Different type
b.flat<double>()(0) = 1.0;
Tensor c(DT_FLOAT, TensorShape({1, 1})); // Different shape
c.flat<float>()(0) = 1.0;
Tensor d(DT_FLOAT, TensorShape({1})); // Different value
d.flat<float>()(0) = 2.0;
// A graph contains a bunch of constants.
Graph g(OpRegistry::Global());
for (const auto& val : {a, b, c, d, d, c, b, a}) {
test::graph::Constant(&g, val); // Node name is n/_0, n/_1, ...
}
GraphDef gdef;
test::graph::ToGraphDef(&g, &gdef);
InitGraph(tsl::LegacyUnredactedDebugString(gdef));
EXPECT_EQ(OriginalGraph(),
"n/_0(Const);n/_1(Const);n/_2(Const);n/_3(Const);"
"n/_4(Const);n/_5(Const);n/_6(Const);n/_7(Const)|");
std::vector<std::string> nodes = str_util::Split(DoCSE(), ";|");
std::set<std::string> node_set(nodes.begin(), nodes.end());
// Expect exactly one of each type of node to be retained after CSE.
EXPECT_EQ(node_set.count("n/_0(Const)") + node_set.count("n/_7(Const)"), 1);
EXPECT_EQ(node_set.count("n/_1(Const)") + node_set.count("n/_6(Const)"), 1);
EXPECT_EQ(node_set.count("n/_2(Const)") + node_set.count("n/_5(Const)"), 1);
EXPECT_EQ(node_set.count("n/_3(Const)") + node_set.count("n/_4(Const)"), 1);
}
void BM_CSE(::testing::benchmark::State& state) {
const int op_nodes = state.range(0);
std::string s;
for (int in = 0; in < 10; in++) {
s += absl::StrFormat("node { name: 'in%04d' op: 'Input'}", in);
}
random::PhiloxRandom philox(301, 17);
random::SimplePhilox rnd(&philox);
for (int op = 0; op < op_nodes; op++) {
s += absl::StrFormat(
"node { name: 'op%04d' op: 'Mul' attr { key: 'T' value { "
"type: DT_FLOAT } } input: ['in%04d', 'in%04d' ] }",
op, rnd.Uniform(10), rnd.Uniform(10));
}
bool first = true;
for (auto i : state) {
state.PauseTiming();
Graph* graph = new Graph(OpRegistry::Global());
InitGraph(s, graph);
int N = graph->num_node_ids();
if (first) {
state.SetLabel(absl::StrCat("Per graph node. Nodes: ", N));
first = false;
}
{
state.ResumeTiming();
OptimizeCSE(graph, nullptr);
state.PauseTiming();
}
delete graph;
state.ResumeTiming();
}
}
BENCHMARK(BM_CSE)->Arg(1000)->Arg(10000);
} // namespace
} // namespace tensorflow
tensorflow/core/graph/regularization/BUILD
+88
View File
@@ -0,0 +1,88 @@
# Description:
# A TensorFlow Graph Regularization Library
load(
"//tensorflow:tensorflow.bzl",
"tf_cc_test",
)
load("//tensorflow:tensorflow.default.bzl", "filegroup")
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
"//tensorflow:__subpackages__",
],
licenses = ["notice"],
)
filegroup(
name = "test_files",
srcs = glob([
"testdata/bert2/**",
"testdata/bert1/**",
]),
)
cc_library(
name = "util",
srcs = ["util.cc"],
hdrs = ["util.h"],
deps = [
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/platform:numbers",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
],
)
tf_cc_test(
name = "util_test",
size = "small",
srcs = [
"util_test.cc",
],
deps = [
":util",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"@com_google_googletest//:gtest_main",
],
)
cc_library(
name = "simple_delete",
srcs = ["simple_delete.cc"],
hdrs = ["simple_delete.h"],
deps = [
":util",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/grappler:op_types",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
],
)
tf_cc_test(
name = "simple_delete_test",
size = "small",
srcs = [
"simple_delete_test.cc",
],
data = [
":test_files",
],
deps = [
":simple_delete",
":util",
"//tensorflow/core:lib",
"//tensorflow/core:test",
"//tensorflow/core/framework:node_def_proto_cc",
"//tensorflow/core/platform:path",
"//tensorflow/core/protobuf:for_core_protos_cc",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings:string_view",
"@com_google_googletest//:gtest_main",
],
)
tensorflow/core/graph/regularization/README.md
+37
View File
@@ -0,0 +1,37 @@
# GraphDef Regularization
This directory contains the code for TensorFlow GraphDef regularization,
sometimes referred to as "canonicalization".
## What does it mean to "regularize" a GraphDef?
The TensorFlow GraphDef is the representation of TensorFlow programs. It shares
a lot of the characteristics of an
[intermediate representation](https://en.wikipedia.org/wiki/Intermediate_representation)
or IR. A single TensorFlow program can produce different GraphDefs, depending on
the device, platform, TF version, runtime state, etc.
"Regularization" is the process of removing this non-determinism from the
GraphDef.
## Interesting Problems
GraphDef regularization helps us answer a variety of interesting questions:
- [Graph Isomorphism](https://en.wikipedia.org/wiki/Graph_isomorphism): Do two
GraphDefs describe the same program?
- [Graph Fingerprinting](https://github.com/tensorflow/community/pull/415): How
can we can uniquely identify a graph using a much shorter fingerprint?
## Algorithms
- `simple_delete`: An algorithm that deletes parts of the GraphDef that are not
deterministic.
## Testing
- TODO(b/239046865): Complete this section.
## Contributions Welcome
If you would like to contribute to the GraphDef regularization library, please
send us a pull request. We welcome collaboration!
tensorflow/core/graph/regularization/simple_delete.cc
+87
View File
@@ -0,0 +1,87 @@
/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#include "tensorflow/core/graph/regularization/simple_delete.h"
#include <cstdint>
#include <string>
#include "absl/status/statusor.h"
#include "absl/strings/strip.h"
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/framework/versions.pb.h"
#include "tensorflow/core/graph/regularization/util.h"
#include "tensorflow/core/grappler/op_types.h"
namespace tensorflow::graph_regularization {
namespace {
// This function mutates `graph_def`, changing the names and config_proto's
// of the Function nodes.
void RegularizeNodes(GraphDef* graph_def) {
for (NodeDef& node : *graph_def->mutable_node()) {
// Check if this is a function call.
if (grappler::IsPartitionedCall(node) ||
grappler::IsStatefulPartitionedCall(node)) {
// Regularize "f" attribute, the function name for PartitionedCall and
// and StatefulPartitionedCall ops, by stripping the suffix UID if it
// has one.
auto f_attr = node.mutable_attr()->find("f");
if (f_attr != node.mutable_attr()->end() && f_attr->second.has_func()) {
std::string function_name = f_attr->second.func().name();
absl::StatusOr<int64_t> uid = GetSuffixUID(function_name);
if (uid.ok()) {
f_attr->second.mutable_func()->set_name(
absl::StripSuffix(function_name, std::to_string(*uid)));
}
}
// Erase the "config_proto" attribute which contains device-specific
// information.
auto node_config_proto = node.mutable_attr()->find("config_proto");
if (node_config_proto != node.attr().end()) {
node_config_proto->second.mutable_s()->erase();
}
}
// Erase the value of string constants, which can vary based on platform.
if (grappler::IsConstant(node)) {
auto dtype_attr = node.attr().find("dtype");
if (dtype_attr != node.attr().end() &&
dtype_attr->second.type() == DT_STRING) {
auto value_attr = node.mutable_attr()->find("value");
if (value_attr != node.mutable_attr()->end()) {
value_attr->second.clear_value();
}
}
}
}
}
} // namespace
void SimpleDelete(GraphDef& graph_def) {
// The GraphDef contains two main sections: a list of nodes and the
// FunctionDefLibrary. Regularization treats these two sections separately.
RegularizeNodes(&graph_def);
// TODO(b/240173815): Complete canonicalization of the FunctionDefLibrary.
// For now, we just completely clear the FunctionDefLibrary.
graph_def.mutable_library()->Clear();
graph_def.mutable_versions()->Clear();
}
} // namespace tensorflow::graph_regularization
tensorflow/core/graph/regularization/simple_delete.h
+28
View File
@@ -0,0 +1,28 @@
/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Licensed 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 TENSORFLOW_CORE_GRAPH_REGULARIZATION_SIMPLE_DELETE_H_
#define TENSORFLOW_CORE_GRAPH_REGULARIZATION_SIMPLE_DELETE_H_
#include "tensorflow/core/framework/graph.pb.h"
namespace tensorflow::graph_regularization {
// Regularizes the graph_def by deleting non-deterministic sections.
void SimpleDelete(GraphDef& graph_def);
} // namespace tensorflow::graph_regularization
#endif // TENSORFLOW_CORE_GRAPH_REGULARIZATION_SIMPLE_DELETE_H_
tensorflow/core/graph/regularization/simple_delete_test.cc
+89
View File
@@ -0,0 +1,89 @@
/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#include "tensorflow/core/graph/regularization/simple_delete.h"
#include <cstdint>
#include <string>
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/graph/regularization/util.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/path.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
#include "tensorflow/core/protobuf/saved_model.pb.h"
#include "tsl/platform/statusor.h"
namespace tensorflow::graph_regularization {
namespace {
absl::StatusOr<SavedModel> ReadSavedModel(absl::string_view file_dir) {
std::string file_path = io::JoinPath(file_dir, "saved_model.pb");
std::string serialized_saved_model;
auto status =
ReadFileToString(Env::Default(), file_path, &serialized_saved_model);
if (!status.ok()) {
return status;
}
SavedModel saved_model_pb;
saved_model_pb.ParseFromString(serialized_saved_model);
return saved_model_pb;
}
TEST(SimpleDeleteTest, MissingFAttributeDoesNotCrash) {
GraphDef graph_def;
NodeDef* node = graph_def.add_node();
node->set_op("PartitionedCall");
node->set_name("my_node");
// Do not add the "f" attribute.
// This should not crash.
SimpleDelete(graph_def);
}
// Test that SimpleDelete algorithm returns the same hash for two models saved
// by calling `tf.saved_model.save` twice in a row in the same program.
TEST(SimpleDeleteTest, TestSimpleDeleteModelSavedTwice) {
const std::string export_dir =
io::JoinPath(testing::TensorFlowSrcRoot(),
"core/graph/regularization/testdata", "bert1");
TF_ASSERT_OK_AND_ASSIGN(SavedModel saved_model_pb,
ReadSavedModel(export_dir));
MetaGraphDef* metagraph = saved_model_pb.mutable_meta_graphs(0);
GraphDef* graph_def = metagraph->mutable_graph_def();
SimpleDelete(*graph_def);
uint64_t hash1 = ComputeHash(*graph_def);
const std::string export_dir2 =
io::JoinPath(testing::TensorFlowSrcRoot(),
"core/graph/regularization/testdata", "bert2");
TF_ASSERT_OK_AND_ASSIGN(SavedModel saved_model_pb2,
ReadSavedModel(export_dir2));
const MetaGraphDef& metagraph2 = saved_model_pb2.meta_graphs(0);
GraphDef graph_def2 = metagraph2.graph_def();
SimpleDelete(graph_def2);
uint64_t hash2 = ComputeHash(graph_def2);
EXPECT_EQ(hash1, hash2);
}
} // namespace
} // namespace tensorflow::graph_regularization
tensorflow/core/graph/regularization/testdata/bert1/saved_model.pb
Vendored
BIN
View File
Binary file not shown.
tensorflow/core/graph/regularization/testdata/bert2/saved_model.pb
Vendored
BIN
View File
Binary file not shown.
tensorflow/core/graph/regularization/util.cc
+53
View File
@@ -0,0 +1,53 @@
/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#include "tensorflow/core/graph/regularization/util.h"
#include <cstdint>
#include <string>
#include <vector>
#include "absl/status/statusor.h"
#include "absl/strings/numbers.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_split.h"
#include "absl/strings/string_view.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/lib/strings/proto_serialization.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/fingerprint.h"
#include "tensorflow/core/platform/numbers.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow::graph_regularization {
uint64_t ComputeHash(const GraphDef& graph_def) {
std::string graph_def_string;
SerializeToStringDeterministic(graph_def, &graph_def_string);
return tensorflow::Fingerprint64(graph_def_string);
}
absl::StatusOr<int64_t> GetSuffixUID(absl::string_view function_name) {
std::vector<absl::string_view> v = absl::StrSplit(function_name, '_');
int64_t uid;
if (!absl::SimpleAtoi(v.back(), &uid)) {
return absl::InvalidArgumentError(absl::StrCat(
"Function name: `", function_name, "` does not end in an integer."));
}
return uid;
}
} // namespace tensorflow::graph_regularization
tensorflow/core/graph/regularization/util.h
+37
View File
@@ -0,0 +1,37 @@
/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Licensed 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 TENSORFLOW_CORE_GRAPH_REGULARIZATION_UTIL_H_
#define TENSORFLOW_CORE_GRAPH_REGULARIZATION_UTIL_H_
#include <cstdint>
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow::graph_regularization {
// Computes the Fingerprint64 hash of the GraphDef.
uint64_t ComputeHash(const GraphDef& graph_def);
// Returns the suffix UID of `function_name`, returns an error if there is none.
absl::StatusOr<int64_t> GetSuffixUID(absl::string_view function_name);
} // namespace tensorflow::graph_regularization
#endif // TENSORFLOW_CORE_GRAPH_REGULARIZATION_UTIL_H_
tensorflow/core/graph/regularization/util_test.cc
+52
View File
@@ -0,0 +1,52 @@
/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#include "tensorflow/core/graph/regularization/util.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow::graph_regularization {
namespace {
GraphDef CreateTestGraph() {
GraphDef graph_def;
NodeDef* node = graph_def.add_node();
node->set_name("name1");
node->set_op("op1");
node = graph_def.add_node();
node->set_name("name2");
node->set_op("op2");
return graph_def;
}
TEST(UtilTest, TestGetSuffixUID) { EXPECT_EQ(*GetSuffixUID("foo_32"), 32); }
TEST(UtilTest, TestGetSuffixUID64Bit) {
EXPECT_EQ(*GetSuffixUID("foo_2209431640"), 2209431640);
}
TEST(UtilTest, TestGetSuffixUIDInvalid) {
EXPECT_FALSE(GetSuffixUID("foo").ok());
}
TEST(FingerprintingTest, TestComputeHash) {
GraphDef graph_def = CreateTestGraph();
EXPECT_EQ(ComputeHash(graph_def), 4870331646167591885);
}
} // namespace
} // namespace tensorflow::graph_regularization
tensorflow/core/graph/subgraph.cc
+408
View File
@@ -0,0 +1,408 @@
/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#include "tensorflow/core/graph/subgraph.h"
#include <algorithm>
#include <deque>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/graph/algorithm.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/graph/tensor_id.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/platform/logging.h"
namespace tensorflow {
namespace subgraph {
// ----------------------------------------------------------------------------
// Subgraph construction-related routines
// ----------------------------------------------------------------------------
// TODO(vrv): Profile the unordered_set and unordered_map use in this file to
// see if we should use an alternative implementation.
namespace {
typedef std::unordered_map<absl::string_view, Node*, StringPieceHasher>
NameIndex;
// Rewrite graph by replacing the output tensors specified in
// "fed_outputs" with special feed nodes for each specified output
// tensor, and removing any nodes that are now disconnected from the
// part of the graph that reaches the sink node. The set of special
// feed nodes added to the graph are returned in "*feed_nodes".
//
// Return true on success. On error, return false and sets *error to
// an appropriate error message (and *g is left in an indeterminate
// state).
absl::Status FeedInputs(
Graph* g, const std::vector<std::unique_ptr<PruneRewrite>>& feed_rewrites,
NameIndex* name_index, DataTypeVector* out_feed_types) {
out_feed_types->clear();
out_feed_types->reserve(feed_rewrites.size());
for (size_t i = 0; i < feed_rewrites.size(); ++i) {
const std::string& t = feed_rewrites[i]->endpoint_name();
TensorId id(ParseTensorName(t));
auto iter = name_index->find(id.first);
if (iter == name_index->end()) {
return absl::NotFoundError(
absl::StrCat("FeedInputs: unable to find feed output ", t));
}
Node* n = iter->second;
DCHECK_EQ(n->name(), id.first);
if (id.second < 0 || id.second >= n->num_outputs()) {
return absl::InvalidArgumentError(
absl::StrCat("FeedInputs: ", t, " should have output index < ",
n->num_outputs(), " and >= 0"));
}
Node* feed_node;
TF_RETURN_IF_ERROR(
feed_rewrites[i]->AddNode(g, {n, id.second}, &feed_node));
// Update name_index
(*name_index)[feed_node->name()] = feed_node;
// Duplicate control edges aren't allowed, but feed_node was *just* created
// so there's no need to check for a duplicate.
g->AddControlEdge(g->source_node(), feed_node, true);
// Look through edges coming out of "n" for edges whose src_output() index
// matches "output_index". If found, replace the edges with a connection
// from the special feed node.
std::vector<const Edge*> to_remove;
for (const Edge* e : n->out_edges()) {
if (e->src_output() == id.second) {
to_remove.emplace_back(e);
} else if (e->src_output() == Graph::kControlSlot &&
(n->type_string() == "Placeholder" ||
n->type_string() == "PlaceholderV2")) {
// When feeding a Placeholder node, any outgoing control edges
// will be replaced with a control edge from the replacement
// feed_node.
// TODO(josh11b,mrry): Come up with a more elegant way of addressing
// the general version of this problem.
to_remove.emplace_back(e);
}
}
for (const Edge* e : to_remove) {
if (e->src_output() == id.second) {
g->AddEdge(feed_node, 0, e->dst(), e->dst_input());
} else {
CHECK_EQ(Graph::kControlSlot, e->src_output());
// Duplicate control edges aren't allowed, but feed_node was *just*
// created so there's no need to check for a duplicate.
g->AddControlEdge(feed_node, e->dst(), true);
}
g->RemoveEdge(e);
}
out_feed_types->push_back(BaseType(n->output_type(id.second)));
}
return absl::OkStatus();
}
absl::Status FetchOutputs(
Graph* g, const std::vector<std::unique_ptr<PruneRewrite>>& fetch_rewrites,
NameIndex* name_index, std::vector<Node*>* out_fetch_nodes,
DataTypeVector* out_fetch_types) {
out_fetch_nodes->clear();
out_fetch_nodes->reserve(fetch_rewrites.size());
for (size_t i = 0; i < fetch_rewrites.size(); ++i) {
const std::string& t = fetch_rewrites[i]->endpoint_name();
// Parse t into node_name and output_index.
TensorId id(ParseTensorName(t));
// Find node in graph with that name.
auto iter = name_index->find(id.first);
if (iter == name_index->end()) {
return absl::NotFoundError(
absl::StrCat("FetchOutputs node ", t, ": not found"));
}
Node* n = iter->second;
DCHECK_EQ(n->name(), id.first);
VLOG(2) << "Found fetch node for " << t;
// Validate output_index
if (n->num_outputs() == 0) {
return absl::InvalidArgumentError(absl::StrCat(
"Tried to fetch data for '", t,
"', which produces no output. To run to a node but not fetch any "
"data, pass '",
t,
"' as an argument to the 'target_node_names' argument of the "
"Session::Run API."));
} else if (id.second < 0 || id.second >= n->num_outputs()) {
return absl::InvalidArgumentError(
absl::StrCat("FetchOutputs ", t, ": output index must be < ",
n->num_outputs(), " and >= 0"));
}
// Create the fetch Node and connect it up
Node* fetch_node;
TF_RETURN_IF_ERROR(
fetch_rewrites[i]->AddNode(g, {n, id.second}, &fetch_node));
// Update the index.
(*name_index)[fetch_node->name()] = fetch_node;
// Duplicate control edges aren't allowed, but fetch_node was *just* created
// so there's no need to check for a duplicate.
g->AddControlEdge(fetch_node, g->sink_node(), true);
out_fetch_nodes->push_back(fetch_node);
out_fetch_types->push_back(BaseType(n->output_type(id.second)));
}
return absl::OkStatus();
}
bool AddNodeToTargets(const std::string& node_or_tensor_name,
const NameIndex& name_index,
std::unordered_set<const Node*>* targets) {
TensorId id = ParseTensorName(node_or_tensor_name);
auto iter = name_index.find(id.first);
if (iter == name_index.end()) {
return false;
}
const Node* n = iter->second;
CHECK_EQ(n->name(), id.first);
targets->insert(n);
return true;
}
absl::Status PruneForTargets(
Graph* g, const NameIndex& name_index,
const std::vector<Node*>& fetch_nodes,
const absl::Span<const std::string>& target_nodes) {
std::string not_found;
std::unordered_set<const Node*> targets;
for (Node* n : fetch_nodes) {
if (!AddNodeToTargets(n->name(), name_index, &targets)) {
absl::StrAppend(&not_found, n->name(), " ");
}
}
for (const std::string& s : target_nodes) {
if (!AddNodeToTargets(s, name_index, &targets)) {
absl::StrAppend(&not_found, s, " ");
}
}
if (!not_found.empty()) {
return absl::NotFoundError(absl::StrCat(
"PruneForTargets: Some target nodes not found: ", not_found));
}
PruneForReverseReachability(g, std::move(targets));
// Reconnect nodes with no outgoing edges to the sink node
FixupSourceAndSinkEdges(g);
return absl::OkStatus();
}
} // namespace
absl::Status ArgFeedRewrite::AddNode(Graph* g, NodeBuilder::NodeOut feed_tensor,
Node** out_node) {
// NOTE(mrry): We must include the index as part of the node
// name, because _Arg is a "stateful" kernel and therefore
// its name must uniquely identify a kernel instance across all
// graphs in the same session.
TF_RETURN_IF_ERROR(
NodeBuilder(strings::StrCat("_arg_", feed_tensor.node->name(), "_",
feed_tensor.index, "_", arg_index_),
"_Arg")
.Attr("T", BaseType(feed_tensor.node->output_type(feed_tensor.index)))
.Attr("index", arg_index_)
.Finalize(g, out_node, /*consume=*/true));
(*out_node)->set_assigned_device_name(device_info().name());
return absl::OkStatus();
}
absl::Status RecvFeedRewrite::AddNode(Graph* g,
NodeBuilder::NodeOut feed_tensor,
Node** out_node) {
TF_RETURN_IF_ERROR(
NodeBuilder(absl::StrCat("_recv_", feed_tensor.node->name(), "_",
feed_tensor.index),
"_Recv")
.Attr("tensor_type",
BaseType(feed_tensor.node->output_type(feed_tensor.index)))
.Attr("tensor_name", endpoint_name())
.Attr("send_device", device_info().name())
.Attr("recv_device", device_info().name())
.Attr("send_device_incarnation",
static_cast<int64_t>(device_info().incarnation()))
.Attr("client_terminated", true)
.Finalize(g, out_node, /*consume=*/true));
(*out_node)->set_assigned_device_name(device_info().name());
return absl::OkStatus();
}
absl::Status RetvalFetchRewrite::AddNode(Graph* g,
NodeBuilder::NodeOut fetch_tensor,
Node** out_node) {
// NOTE(mrry): We must include the index as part of the node
// name, because _Retval is a "stateful" kernel and therefore
// its name must uniquely identify a kernel instance across all
// graphs in the same session.
TF_RETURN_IF_ERROR(
NodeBuilder(strings::StrCat("_retval_", fetch_tensor.node->name(), "_",
fetch_tensor.index, "_", retval_index_),
"_Retval")
.Input(fetch_tensor.node, fetch_tensor.index)
.Attr("T",
BaseType(fetch_tensor.node->output_type(fetch_tensor.index)))
.Attr("index", retval_index_)
.Finalize(g, out_node, /*consume=*/true));
(*out_node)->set_assigned_device_name(device_info().name());
return absl::OkStatus();
}
absl::Status SendFetchRewrite::AddNode(Graph* g,
NodeBuilder::NodeOut fetch_tensor,
Node** out_node) {
TF_RETURN_IF_ERROR(
NodeBuilder(absl::StrCat("_send_", fetch_tensor.node->name(), "_",
fetch_tensor.index),
"_Send")
.Input(fetch_tensor.node, fetch_tensor.index)
.Attr("tensor_name", endpoint_name())
.Attr("send_device", device_info().name())
.Attr("recv_device", device_info().name())
.Attr("send_device_incarnation",
static_cast<int64_t>(device_info().incarnation()))
.Attr("client_terminated", true)
.Finalize(g, out_node, /*consume=*/true));
(*out_node)->set_assigned_device_name(device_info().name());
return absl::OkStatus();
}
absl::Status RewriteGraphForExecution(
Graph* g, const absl::Span<const std::string>& fed_outputs,
const absl::Span<const std::string>& fetch_outputs,
const absl::Span<const std::string>& target_node_names,
const DeviceAttributes& device_info, bool use_function_convention,
RewriteGraphMetadata* out_metadata) {
std::vector<std::unique_ptr<PruneRewrite>> feed_rewrites;
feed_rewrites.reserve(fed_outputs.size());
if (use_function_convention) {
for (size_t i = 0; i < fed_outputs.size(); ++i) {
feed_rewrites.emplace_back(new ArgFeedRewrite(
&fed_outputs[i], &device_info, static_cast<int32_t>(i)));
}
} else {
for (const std::string& fed_output : fed_outputs) {
feed_rewrites.emplace_back(
new RecvFeedRewrite(&fed_output, &device_info));
}
}
std::vector<std::unique_ptr<PruneRewrite>> fetch_rewrites;
fetch_rewrites.reserve(fetch_outputs.size());
if (use_function_convention) {
for (size_t i = 0; i < fetch_outputs.size(); ++i) {
fetch_rewrites.emplace_back(new RetvalFetchRewrite(
&fetch_outputs[i], &device_info, static_cast<int32_t>(i)));
}
} else {
for (const std::string& fetch_output : fetch_outputs) {
fetch_rewrites.emplace_back(
new SendFetchRewrite(&fetch_output, &device_info));
}
}
return RewriteGraphForExecution(g, feed_rewrites, fetch_rewrites,
target_node_names, out_metadata);
}
namespace {
template <typename StringContainer>
std::vector<std::string> ConvertToVector(StringContainer field) {
return std::vector<std::string>(field.begin(), field.end());
}
} // namespace
absl::Status RewriteGraphForExecution(
Graph* g, const std::vector<std::unique_ptr<PruneRewrite>>& feed_rewrites,
const std::vector<std::unique_ptr<PruneRewrite>>& fetch_rewrites,
const absl::Span<const std::string>& target_node_names,
RewriteGraphMetadata* out_metadata) {
if (fetch_rewrites.empty() && target_node_names.empty()) {
return absl::InvalidArgumentError(
"Must specify at least one target to fetch or execute.");
}
std::unordered_set<std::string> endpoints;
for (const auto& feed_rewrite : feed_rewrites) {
auto result = endpoints.insert(feed_rewrite->endpoint_name());
if (!result.second) {
return absl::InvalidArgumentError(
absl::StrCat("Endpoint \"", feed_rewrite->endpoint_name(),
"\" fed more than once."));
}
}
for (const auto& fetch_rewrite : fetch_rewrites) {
if (endpoints.count(fetch_rewrite->endpoint_name()) > 0) {
return absl::InvalidArgumentError(absl::StrCat(
fetch_rewrite->endpoint_name(), " is both fed and fetched."));
}
}
// A separate index mapping name to Node*, for use by FeedInputs,
// FetchOutputs, and PruneForTargets
NameIndex name_index;
name_index.reserve(g->num_nodes());
for (Node* n : g->nodes()) {
name_index[n->name()] = n;
}
// Add the feeds. This may replace nodes in the graph, including the nodes
// currently listed in "fetch_rewrites". We pass "name_index" so the index is
// kept up to date.
if (!feed_rewrites.empty()) {
TF_RETURN_IF_ERROR(
FeedInputs(g, feed_rewrites, &name_index, &out_metadata->feed_types));
}
// Add the fetch nodes, also updating "name_index".
std::vector<Node*> fetch_nodes;
if (!fetch_rewrites.empty()) {
TF_RETURN_IF_ERROR(FetchOutputs(g, fetch_rewrites, &name_index,
&fetch_nodes, &out_metadata->fetch_types));
}
// Prune the graph to only compute what is needed for the fetch nodes and the
// target nodes.
if (!fetch_nodes.empty() || !target_node_names.empty()) {
TF_RETURN_IF_ERROR(
PruneForTargets(g, name_index, fetch_nodes, target_node_names));
}
return absl::OkStatus();
}
} // namespace subgraph
} // namespace tensorflow
tensorflow/core/graph/subgraph.h
+166
View File
@@ -0,0 +1,166 @@
/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
Licensed 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 TENSORFLOW_CORE_GRAPH_SUBGRAPH_H_
#define TENSORFLOW_CORE_GRAPH_SUBGRAPH_H_
#include <string>
#include "tensorflow/core/framework/device_attributes.pb.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/graph/node_builder.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/gtl/array_slice.h"
#include "tensorflow/core/protobuf/config.pb.h"
namespace tensorflow {
namespace subgraph {
// Information about a graph rewritten by `RewriteGraphForExecution()`.
struct RewriteGraphMetadata {
// The element type of each tensor fed to this subgraph. The order
// of types corresponds to the order of tensor names in
// `fed_outputs` when calling `RewriteGraphForExecution()`.
DataTypeVector feed_types;
// The element type of each tensor fetched from this subgraph. The
// order of types corresponds to the order of tensor names in
// `fetch_outputs` when calling `RewriteGraphForExecution()`.
DataTypeVector fetch_types;
};
// Describes the action to take on a particular tensor endpoint (described by
// a "<node_name>:<output_index>" pair) when pruning the graph.
//
// The `AddNode()` method must be overridden to describe this action. The method
// will be invoked once during `RewriteGraphForExecution()` with tensor endpoint
// named by `endpoint_name`, and it may either create a single new node, or fail
// with an error if the resulting graph would be invalid.
class PruneRewrite {
public:
// `endpoint_name` and `device_info` must outlive this object.
PruneRewrite(const std::string* endpoint_name,
const DeviceAttributes* device_info)
: endpoint_name_(endpoint_name), device_info_(device_info) {}
virtual ~PruneRewrite() {}
// Creates a new node whose output replaces the given `tensor` in graph `g`.
// The node will be assigned to the device named in `device_info`.
virtual absl::Status AddNode(Graph* g, NodeBuilder::NodeOut tensor,
Node** out_node) = 0;
// Returns the name of the tensor to which this rewrite applies.
const std::string& endpoint_name() { return *endpoint_name_; }
protected:
// The device on which the new node will be created.
const DeviceAttributes& device_info() { return *device_info_; }
private:
const std::string* const endpoint_name_; // Not owned.
const DeviceAttributes* const device_info_; // Not owned.
};
// Rewrite the graph structure of "*g" to deal with feeding node
// outputs, fetching node outputs, and only running a subset of the
// graph. "fed_outputs" and "fetch_outputs" are both lists of
// output tensor identifiers in the form of
// "<name>[:<optional_output_index>]", and "target_nodes_str" is a
// lists of target node names in "*g" "g".
//
// In the resulting graph "*g", output edges in "fed_outputs" have
// been redirected to special "_recv" nodes introduced into the graph.
// If these fed nodes are not needed in order to compute the effects
// of the nodes in "target_node_names" and "fetch_outputs", then these may
// be omitted from the graph.
//
// In the resulting graph "*g", additional "_send" nodes are connected
// to every output in "fetch_outputs". These "_send" nodes are set up
// to execute on the device described by device_info.
//
// On success, returns OK, and sets "*g" to a version of "*g"
// that represents the portions of the graph necessary for producing
// the output of all nodes listed in "target_node_names" and fetching the
// specific node outputs specified in "fetch_outputs".
//
// On failure, returns the error status. Possible errors include:
// - fed output "node:output_index" does not exist in "*g"
// - fetch output "node:output_index" does not exist in "*g"
// - target node "node" does not exist in "*g"
absl::Status RewriteGraphForExecution(
Graph* g, const absl::Span<const std::string>& fed_outputs,
const absl::Span<const std::string>& fetch_outputs,
const absl::Span<const std::string>& target_node_names,
const DeviceAttributes& device_info, bool use_function_convention,
RewriteGraphMetadata* out_metadata);
// A more general version of the above function that supports
// customizable rewriting actions for each fed and fetched tensor.
absl::Status RewriteGraphForExecution(
Graph* g, const std::vector<std::unique_ptr<PruneRewrite>>& feed_rewrites,
const std::vector<std::unique_ptr<PruneRewrite>>& fetch_rewrites,
const absl::Span<const std::string>& target_node_names,
RewriteGraphMetadata* out_metadata);
/////////////////////////////////////////////////////////
// Custom rewrite actions for fed and fetched tensors. //
/////////////////////////////////////////////////////////
// A rewrite action that adds an _Arg node for a fed tensor.
class ArgFeedRewrite : public PruneRewrite {
public:
ArgFeedRewrite(const std::string* endpoint_name,
const DeviceAttributes* device_info, int32_t arg_index)
: PruneRewrite(endpoint_name, device_info), arg_index_(arg_index) {}
absl::Status AddNode(Graph* g, NodeBuilder::NodeOut feed_tensor,
Node** out_node) override;
private:
const int32_t arg_index_;
};
// A rewrite action that adds a client-terminated _Recv node for a fed tensor.
class RecvFeedRewrite : public PruneRewrite {
public:
using PruneRewrite::PruneRewrite;
absl::Status AddNode(Graph* g, NodeBuilder::NodeOut feed_tensor,
Node** out_node) override;
};
// A rewrite action that adds a _Retval node for a fetched tensor.
class RetvalFetchRewrite : public PruneRewrite {
public:
RetvalFetchRewrite(const std::string* endpoint_name,
const DeviceAttributes* device_info, int32_t retval_index)
: PruneRewrite(endpoint_name, device_info), retval_index_(retval_index) {}
absl::Status AddNode(Graph* g, NodeBuilder::NodeOut fetch_tensor,
Node** out_node) override;
private:
const int32_t retval_index_;
};
// A rewrite action that adds a client-terminated _Send node for a
// fetched tensor.
class SendFetchRewrite : public PruneRewrite {
public:
using PruneRewrite::PruneRewrite;
absl::Status AddNode(Graph* g, NodeBuilder::NodeOut fetch_tensor,
Node** out_node) override;
};
} // namespace subgraph
} // namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPH_SUBGRAPH_H_
tensorflow/core/graph/subgraph_test.cc
+401
View File
@@ -0,0 +1,401 @@
/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#include "tensorflow/core/graph/subgraph.h"
#include <string>
#include <vector>
#include "tensorflow/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/common_runtime/graph_def_builder_util.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/partial_tensor_shape.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/graph/graph_def_builder.h"
#include "tensorflow/core/kernels/ops_util.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/lib/strings/str_util.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/protobuf.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/test_benchmark.h"
// TODO(josh11b): Test setting the "device" field of a NodeDef.
// TODO(josh11b): Test that feeding won't prune targets.
namespace tensorflow {
namespace {
class SubgraphTest : public ::testing::Test {
protected:
SubgraphTest() : g_(new Graph(OpRegistry::Global())) {
device_info_.set_name("/job:a/replica:0/task:0/cpu:0");
device_info_.set_device_type(DeviceType(DEVICE_CPU).type());
device_info_.set_incarnation(0);
}
~SubgraphTest() override {}
void ExpectOK(const std::string& gdef_ascii) {
CHECK(protobuf::TextFormat::ParseFromString(gdef_ascii, &gdef_));
GraphConstructorOptions opts;
TF_CHECK_OK(ConvertGraphDefToGraph(opts, gdef_, g_.get()));
}
Node* FindNode(const std::string& name) {
for (Node* n : g_->nodes()) {
if (n->name() == name) return n;
}
return nullptr;
}
bool HasNode(const std::string& name) { return FindNode(name) != nullptr; }
void ExpectNodes(const std::string& nodes) {
int count = 0;
std::vector<std::string> actual_nodes;
for (Node* n : g_->nodes()) {
if (n->IsOp()) {
count++;
actual_nodes.push_back(n->name());
}
}
std::sort(actual_nodes.begin(), actual_nodes.end());
LOG(INFO) << "Nodes present: " << absl::StrJoin(actual_nodes, " ");
std::vector<std::string> expected_nodes = str_util::Split(nodes, ',');
std::sort(expected_nodes.begin(), expected_nodes.end());
for (const std::string& s : expected_nodes) {
Node* n = FindNode(s);
EXPECT_TRUE(n != nullptr) << s;
if (n->type_string() == "_Send" || n->type_string() == "_Recv") {
EXPECT_EQ(device_info_.name(), n->assigned_device_name()) << s;
}
}
EXPECT_TRUE(actual_nodes.size() == expected_nodes.size())
<< "\nActual: " << absl::StrJoin(actual_nodes, ",")
<< "\nExpected: " << absl::StrJoin(expected_nodes, ",");
}
bool HasEdge(const std::string& src, int src_out, const std::string& dst,
int dst_in) {
for (const Edge* e : g_->edges()) {
if (e->src()->name() == src && e->src_output() == src_out &&
e->dst()->name() == dst && e->dst_input() == dst_in)
return true;
}
return false;
}
bool HasControlEdge(const std::string& src, const std::string& dst) {
return HasEdge(src, Graph::kControlSlot, dst, Graph::kControlSlot);
}
std::string Subgraph(const std::string& fed_str, const std::string& fetch_str,
const std::string& targets_str,
bool use_function_convention = false) {
Graph* subgraph = new Graph(OpRegistry::Global());
CopyGraph(*g_, subgraph);
std::vector<std::string> fed =
str_util::Split(fed_str, ',', str_util::SkipEmpty());
std::vector<std::string> fetch =
str_util::Split(fetch_str, ',', str_util::SkipEmpty());
std::vector<std::string> targets =
str_util::Split(targets_str, ',', str_util::SkipEmpty());
subgraph::RewriteGraphMetadata metadata;
absl::Status s = subgraph::RewriteGraphForExecution(
subgraph, fed, fetch, targets, device_info_, use_function_convention,
&metadata);
if (!s.ok()) {
delete subgraph;
return s.ToString();
}
EXPECT_EQ(fed.size(), metadata.feed_types.size());
EXPECT_EQ(fetch.size(), metadata.fetch_types.size());
// Replace the graph with the subgraph for the rest of the display program
g_.reset(subgraph);
return "OK";
}
Graph* graph() { return g_.get(); }
private:
GraphDef gdef_;
std::unique_ptr<Graph> g_;
DeviceAttributes device_info_;
};
REGISTER_OP("TestParams").Output("o: float");
REGISTER_OP("TestInput").Output("a: float").Output("b: float");
REGISTER_OP("TestRelu").Input("i: float").Output("o: float");
REGISTER_OP("TestMul").Input("a: float").Input("b: float").Output("o: float");
TEST_F(SubgraphTest, Targets1) {
ExpectOK(
"node { name: 'W1' op: 'TestParams' }"
"node { name: 'W2' op: 'TestParams' }"
"node { name: 'input' op: 'TestInput' }"
"node { name: 't1' op: 'TestMul' input: [ 'W1', 'input:1' ] }"
"node { name: 't2' op: 'TestMul' input: [ 'W2', 't1' ] }"
"node { name: 't3_a' op: 'TestRelu' input: 't2' }"
"node { name: 't3_b' op: 'TestRelu' input: 't2' }");
EXPECT_EQ("OK", Subgraph("", "", "t1"));
ExpectNodes("W1,input,t1");
}
TEST_F(SubgraphTest, Targets2) {
ExpectOK(
"node { name: 'W1' op: 'TestParams' }"
"node { name: 'W2' op: 'TestParams' }"
"node { name: 'input' op: 'TestInput' }"
"node { name: 't1' op: 'TestMul' input: 'W1' input: 'input:1' }"
"node { name: 't2' op: 'TestMul' input: 'W2' input: 't1' }"
"node { name: 't3_a' op: 'TestRelu' input: 't2' }"
"node { name: 't3_b' op: 'TestRelu' input: 't2' }");
EXPECT_EQ("OK", Subgraph("", "", "t2,t3_a"));
ExpectNodes("W1,W2,input,t1,t2,t3_a");
}
TEST_F(SubgraphTest, FedOutputs1) {
ExpectOK(
"node { name: 'W1' op: 'TestParams' }"
"node { name: 'W2' op: 'TestParams' }"
"node { name: 'input' op: 'TestInput' }"
"node { name: 't1' op: 'TestMul' input: [ 'W1', 'input:1' ] }"
"node { name: 't2' op: 'TestMul' input: [ 'W2', 't1' ] }"
"node { name: 't3_a' op: 'TestRelu' input: 't2' }"
"node { name: 't3_b' op: 'TestRelu' input: 't2' }");
EXPECT_EQ("OK", Subgraph("input:1", "", "t2"));
ExpectNodes("W1,W2,_recv_input_1,t1,t2");
}
TEST_F(SubgraphTest, FedOutputs1_FunctionConvention) {
ExpectOK(
"node { name: 'W1' op: 'TestParams' }"
"node { name: 'W2' op: 'TestParams' }"
"node { name: 'input' op: 'TestInput' }"
"node { name: 't1' op: 'TestMul' input: [ 'W1', 'input:1' ] }"
"node { name: 't2' op: 'TestMul' input: [ 'W2', 't1' ] }"
"node { name: 't3_a' op: 'TestRelu' input: 't2' }"
"node { name: 't3_b' op: 'TestRelu' input: 't2' }");
EXPECT_EQ("OK",
Subgraph("input:1", "", "t2", true /* use_function_convention */));
ExpectNodes("W1,W2,_arg_input_1_0,t1,t2");
}
TEST_F(SubgraphTest, FedRefNode) {
ExpectOK(
"node { name: 'W1' op: 'TestParams' }"
"node { name: 'W2' op: 'TestParams' }"
"node { name: 't1' op: 'TestMul' input: [ 'W2', 'W1' ] }");
EXPECT_EQ("OK", Subgraph("W1:0", "", "t1"));
ExpectNodes("_recv_W1_0,W2,t1");
Node* n = FindNode("_recv_W1_0");
EXPECT_FALSE(IsRefType(CHECK_NOTNULL(n)->output_type(0)));
}
TEST_F(SubgraphTest, FedRefNode_FunctionConvention) {
ExpectOK(
"node { name: 'W1' op: 'TestParams' }"
"node { name: 'W2' op: 'TestParams' }"
"node { name: 't1' op: 'TestMul' input: [ 'W2', 'W1' ] }");
EXPECT_EQ("OK",
Subgraph("W1:0", "", "t1", true /* use_function_convention */));
ExpectNodes("_arg_W1_0_0,W2,t1");
Node* n = FindNode("_arg_W1_0_0");
EXPECT_FALSE(IsRefType(CHECK_NOTNULL(n)->output_type(0)));
}
TEST_F(SubgraphTest, FedOutputs2_FunctionConvention) {
ExpectOK(
"node { name: 'W1' op: 'TestParams' }"
"node { name: 'W2' op: 'TestParams' }"
"node { name: 'input' op: 'TestInput' }"
"node { name: 't1' op: 'TestMul' input: [ 'W1', 'input:1' ] }"
"node { name: 't2' op: 'TestMul' input: [ 'W2', 't1' ] }"
"node { name: 't3_a' op: 'TestRelu' input: 't2' }"
"node { name: 't3_b' op: 'TestRelu' input: 't2' }");
// We feed input:1, but nothing connects to it, so the _recv(input:1)
// node also disappears.
EXPECT_EQ("OK", Subgraph("input:1,t1,W2", "", "t2",
true /* use_function_convention */));
ExpectNodes("_arg_t1_0_1,_arg_W2_0_2,t2");
}
TEST_F(SubgraphTest, FetchOutputs1) {
ExpectOK(
"node { name: 'W1' op: 'TestParams' }"
"node { name: 'W2' op: 'TestParams' }"
"node { name: 'input' op: 'TestInput' }"
"node { name: 't1' op: 'TestMul' input: [ 'W1', 'input:1' ] }"
"node { name: 't2' op: 'TestMul' input: [ 'W2', 't1' ] }"
"node { name: 't3_a' op: 'TestRelu' input: 't2' }"
"node { name: 't3_b' op: 'TestRelu' input: 't2' }");
EXPECT_EQ("OK", Subgraph("", "W2,input:1,t1,t2", "t2"));
ExpectNodes(
"W1,W2,input,t1,t2,_send_W2_0,_send_input_1,_send_t1_0,_send_t2_0");
}
TEST_F(SubgraphTest, FetchOutputs1_FunctionConvention) {
ExpectOK(
"node { name: 'W1' op: 'TestParams' }"
"node { name: 'W2' op: 'TestParams' }"
"node { name: 'input' op: 'TestInput' }"
"node { name: 't1' op: 'TestMul' input: [ 'W1', 'input:1' ] }"
"node { name: 't2' op: 'TestMul' input: [ 'W2', 't1' ] }"
"node { name: 't3_a' op: 'TestRelu' input: 't2' }"
"node { name: 't3_b' op: 'TestRelu' input: 't2' }");
EXPECT_EQ("OK", Subgraph("", "W2,input:1,t1,t2", "t2",
true /* use_function_convention */));
ExpectNodes(
"W1,W2,input,t1,t2,_retval_W2_0_0,_retval_input_1_1,_retval_t1_0_2,_"
"retval_t2_0_3");
}
TEST_F(SubgraphTest, FetchOutputs2) {
ExpectOK(
"node { name: 'W1' op: 'TestParams' }"
"node { name: 'W2' op: 'TestParams' }"
"node { name: 'input' op: 'TestInput' }"
"node { name: 't1' op: 'TestMul' input: [ 'W1', 'input:1' ] }"
"node { name: 't2' op: 'TestMul' input: [ 'W2', 't1' ] }"
"node { name: 't3_a' op: 'TestRelu' input: 't2' }"
"node { name: 't3_b' op: 'TestRelu' input: 't2' }");
EXPECT_EQ("OK", Subgraph("", "t3_a", "t2"));
ExpectNodes("W1,W2,input,t1,t2,t3_a,_send_t3_a_0");
}
TEST_F(SubgraphTest, FetchOutputs2_FunctionConvention) {
ExpectOK(
"node { name: 'W1' op: 'TestParams' }"
"node { name: 'W2' op: 'TestParams' }"
"node { name: 'input' op: 'TestInput' }"
"node { name: 't1' op: 'TestMul' input: [ 'W1', 'input:1' ] }"
"node { name: 't2' op: 'TestMul' input: [ 'W2', 't1' ] }"
"node { name: 't3_a' op: 'TestRelu' input: 't2' }"
"node { name: 't3_b' op: 'TestRelu' input: 't2' }");
EXPECT_EQ("OK",
Subgraph("", "t3_a", "t2", true /* use_function_convention */));
ExpectNodes("W1,W2,input,t1,t2,t3_a,_retval_t3_a_0_0");
}
TEST_F(SubgraphTest, ChainOfFools) {
ExpectOK(
"node { name: 'a' op: 'TestParams' }"
"node { name: 'b' op: 'TestRelu' input: 'a'}"
"node { name: 'c' op: 'TestRelu' input: 'b'}"
"node { name: 'd' op: 'TestRelu' input: 'c'}"
"node { name: 'e' op: 'TestRelu' input: 'd'}"
"node { name: 'f' op: 'TestRelu' input: 'e'}");
EXPECT_EQ("OK", Subgraph("c:0", "b:0,e:0", ""));
ExpectNodes("a,b,_send_b_0,_recv_c_0,d,e,_send_e_0");
EXPECT_TRUE(HasEdge("a", 0, "b", 0));
EXPECT_TRUE(HasEdge("b", 0, "_send_b_0", 0));
EXPECT_TRUE(HasEdge("_recv_c_0", 0, "d", 0));
EXPECT_TRUE(HasEdge("d", 0, "e", 0));
EXPECT_TRUE(HasEdge("e", 0, "_send_e_0", 0));
}
static bool HasSubstr(absl::string_view base, absl::string_view substr) {
bool ok = absl::StrContains(base, substr);
EXPECT_TRUE(ok) << base << ", expected substring " << substr;
return ok;
}
TEST_F(SubgraphTest, Errors) {
ExpectOK(
"node { name: 'a' op: 'TestParams' }"
"node { name: 'b' op: 'TestRelu' input: 'a'}"
"node { name: 'c' op: 'TestRelu' input: 'b'}"
"node { name: 'd' op: 'TestRelu' input: 'c'}"
"node { name: 'e' op: 'TestRelu' input: 'd'}"
"node { name: 'f' op: 'TestRelu' input: 'e'}");
// Duplicated feed and fetch
EXPECT_TRUE(
HasSubstr(Subgraph("c:0", "b:0,c:0", ""), "both fed and fetched"));
// Feed not found.
EXPECT_TRUE(HasSubstr(Subgraph("foo:0", "c:0", ""), "unable to find"));
// Fetch not found.
EXPECT_TRUE(HasSubstr(Subgraph("", "foo:0", ""), "not found"));
// Target not found.
EXPECT_TRUE(HasSubstr(Subgraph("", "", "foo"), "not found"));
// No targets specified.
EXPECT_TRUE(HasSubstr(Subgraph("", "", ""), "at least one target"));
}
REGISTER_OP("In").Output("o: float");
REGISTER_OP("Op").Input("i: float").Output("o: float");
void BM_SubgraphHelper(::testing::benchmark::State& state,
bool use_function_convention) {
const int num_nodes = state.range(0);
DeviceAttributes device_info;
device_info.set_name("/job:a/replica:0/task:0/cpu:0");
device_info.set_device_type(DeviceType(DEVICE_CPU).type());
device_info.set_incarnation(0);
Graph g(OpRegistry::Global());
{ // Scope for temporary variables used to construct g.
GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
Node* last_node = nullptr;
for (int i = 0; i < num_nodes; i++) {
std::string name = absl::StrCat("N", i);
if (i > 0) {
last_node = ops::UnaryOp("Op", last_node, b.opts().WithName(name));
} else {
last_node = ops::SourceOp("In", b.opts().WithName(name));
}
}
TF_CHECK_OK(GraphDefBuilderToGraph(b, &g));
}
std::vector<std::string> fed;
if (num_nodes > 1000) {
fed.push_back(absl::StrCat("N", num_nodes - 1000));
}
std::vector<std::string> fetch;
std::vector<std::string> targets = {absl::StrCat("N", num_nodes - 1)};
for (auto s : state) {
Graph* subgraph = new Graph(OpRegistry::Global());
CopyGraph(g, subgraph);
subgraph::RewriteGraphMetadata metadata;
TF_CHECK_OK(subgraph::RewriteGraphForExecution(
subgraph, fed, fetch, targets, device_info, use_function_convention,
&metadata));
delete subgraph;
}
}
void BM_Subgraph(::testing::benchmark::State& state) {
BM_SubgraphHelper(state, false /* use_function_convention */);
}
void BM_SubgraphFunctionConvention(::testing::benchmark::State& state) {
BM_SubgraphHelper(state, true /* use_function_convention */);
}
BENCHMARK(BM_Subgraph)->Arg(100)->Arg(1000)->Arg(10000)->Arg(100000);
BENCHMARK(BM_SubgraphFunctionConvention)
->Arg(100)
->Arg(1000)
->Arg(10000)
->Arg(100000);
} // namespace
} // namespace tensorflow
tensorflow/core/graph/tensor_id.cc
+62
View File
@@ -0,0 +1,62 @@
/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#include "tensorflow/core/graph/tensor_id.h"
#include <cstddef>
#include <cstdint>
#include <limits>
#include <string>
#include "absl/strings/string_view.h"
#include "absl/strings/strip.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/platform/str_util.h"
namespace tensorflow {
TensorId::TensorId(const SafeTensorId& id) : TensorId(id.first, id.second) {}
SafeTensorId::SafeTensorId(const TensorId& id)
: SafeTensorId(std::string(id.first), id.second) {}
TensorId ParseTensorName(absl::string_view name) {
// Parse either a name, ^name, or name:digits. To do so, we go backwards from
// the end of the string, skipping over a run of digits. If we hit a ':'
// character, then we know we are in the 'name:digits' regime. Otherwise, we
// see if the name starts with '^', indicating a control edge. If we find
// neither ':' nor '^' characters, the output index is implicitly 0, and the
// whole name string forms the first part of the tensor name.
size_t colon_pos = name.rfind(':');
if (colon_pos != absl::string_view::npos) {
absl::string_view prefix = name.substr(0, colon_pos);
absl::string_view suffix = name.substr(colon_pos + 1);
uint64_t index;
if (str_util::ConsumeLeadingDigits(&suffix, &index) && suffix.empty() &&
index <= std::numeric_limits<int>::max()) {
return TensorId(prefix, index);
}
}
if (absl::ConsumePrefix(&name, "^")) {
return TensorId(name, Graph::kControlSlot);
}
return TensorId(name, 0);
}
bool IsTensorIdControl(const TensorId& tensor_id) {
return tensor_id.index() == Graph::kControlSlot;
}
} // namespace tensorflow
tensorflow/core/graph/tensor_id.h
+93
View File
@@ -0,0 +1,93 @@
/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
Licensed 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 TENSORFLOW_CORE_GRAPH_TENSOR_ID_H_
#define TENSORFLOW_CORE_GRAPH_TENSOR_ID_H_
#include <string>
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/lib/core/stringpiece.h"
#include "tensorflow/core/lib/hash/hash.h"
#include "tensorflow/core/lib/strings/strcat.h"
namespace tensorflow {
struct SafeTensorId;
// Identifier for a tensor within a step.
// first == operation_name, second == output_index
// Note: does not own backing storage for name.
struct TensorId : public std::pair<absl::string_view, int> {
typedef std::pair<absl::string_view, int> Base;
// Inherit the set of constructors.
using Base::pair;
// NOTE(skyewm): this is required on some platforms. I'm not sure why the
// using statement above isn't always sufficient.
TensorId() : Base() {}
TensorId(const SafeTensorId& id);
absl::string_view node() const { return first; }
int index() const { return second; }
std::string ToString() const {
if (second == Graph::kControlSlot) return absl::StrCat("^", first);
return absl::StrCat(first, ":", second);
}
struct Hasher {
public:
std::size_t operator()(const TensorId& x) const {
return Hash32(x.first.data(), x.first.size(), x.second);
}
};
};
TensorId ParseTensorName(absl::string_view name);
bool IsTensorIdControl(const TensorId& tensor_id);
// Same as TensorId, except owns the backing storage for the op name. This makes
// the memory management simpler at the expense of a copy.
struct SafeTensorId : public std::pair<std::string, int> {
typedef std::pair<std::string, int> Base;
// NOTE(skyewm): this is required on some platforms. I'm not sure why the
// using "using Base::pair;" isn't always sufficient.
SafeTensorId() : Base() {}
SafeTensorId(const std::string& str, int idx) : Base(str, idx) {}
SafeTensorId(const TensorId& id);
const std::string& node() const { return first; }
int index() const { return second; }
std::string ToString() const {
if (second == Graph::kControlSlot) return absl::StrCat("^", first);
return absl::StrCat(first, ":", second);
}
struct Hasher {
public:
std::size_t operator()(const TensorId& x) const {
return Hash32(x.first.data(), x.first.size(), x.second);
}
};
};
} // namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPH_TENSOR_ID_H_
tensorflow/core/graph/tensor_id_test.cc
+123
View File
@@ -0,0 +1,123 @@
/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#include "tensorflow/core/graph/tensor_id.h"
#include <vector>
#include "tensorflow/core/lib/random/simple_philox.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/test_benchmark.h"
namespace tensorflow {
namespace {
std::string ParseHelper(const std::string& n) {
return ParseTensorName(n).ToString();
}
TEST(TensorIdTest, ParseTensorName) {
EXPECT_EQ(ParseHelper("W1"), "W1:0");
EXPECT_EQ(ParseHelper("W1:0"), "W1:0");
EXPECT_EQ(ParseHelper("weights:0"), "weights:0");
EXPECT_EQ(ParseHelper("W1:1"), "W1:1");
EXPECT_EQ(ParseHelper("W1:17"), "W1:17");
EXPECT_EQ(ParseHelper("xyz1_17"), "xyz1_17:0");
EXPECT_EQ(ParseHelper("^foo"), "^foo");
// Index out of range for int32.
EXPECT_EQ(ParseHelper("W1:2147483648"), "W1:2147483648:0");
}
uint32_t Skewed(random::SimplePhilox* rnd, int max_log) {
const uint32_t space = 1 << (rnd->Rand32() % (max_log + 1));
return rnd->Rand32() % space;
}
void BM_ParseTensorName(::testing::benchmark::State& state) {
const int arg = state.range(0);
random::PhiloxRandom philox(301, 17);
random::SimplePhilox rnd(&philox);
std::vector<std::string> names;
for (int i = 0; i < 100; i++) {
std::string name;
switch (arg) {
case 0: { // Generate random names
size_t len = Skewed(&rnd, 4);
while (name.size() < len) {
name += rnd.OneIn(4) ? '0' : 'a';
}
if (rnd.OneIn(3)) {
absl::StrAppend(&name, ":", rnd.Uniform(12));
}
break;
}
case 1:
name = "W1";
break;
case 2:
name = "t0003";
break;
case 3:
name = "weights";
break;
case 4:
name = "weights:17";
break;
case 5:
name = "^weights";
break;
default:
LOG(FATAL) << "Unexpected arg";
break;
}
names.push_back(name);
}
TensorId id;
int index = 0;
int sum = 0;
for (auto s : state) {
id = ParseTensorName(names[index++ % names.size()]);
sum += id.second;
}
VLOG(2) << sum; // Prevent compiler from eliminating loop body
}
BENCHMARK(BM_ParseTensorName)->Arg(0)->Arg(1)->Arg(2)->Arg(3)->Arg(4)->Arg(5);
TEST(TensorIdTest, IsTensorIdControl) {
std::string input = "^foo";
TensorId tensor_id = ParseTensorName(input);
EXPECT_TRUE(IsTensorIdControl(tensor_id));
input = "foo";
tensor_id = ParseTensorName(input);
EXPECT_FALSE(IsTensorIdControl(tensor_id));
input = "foo:2";
tensor_id = ParseTensorName(input);
EXPECT_FALSE(IsTensorIdControl(tensor_id));
}
TEST(TensorIdTest, PortZero) {
for (std::string input : {"foo", "foo:0"}) {
TensorId tensor_id = ParseTensorName(input);
EXPECT_EQ("foo", tensor_id.node());
EXPECT_EQ(0, tensor_id.index());
}
}
} // namespace
} // namespace tensorflow
tensorflow/core/graph/testlib.cc
+532
View File
@@ -0,0 +1,532 @@
/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#include "tensorflow/core/graph/testlib.h"
#include <vector>
#include "tensorflow/core/framework/common_shape_fns.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def_builder.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/graph/node_builder.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/platform/logging.h"
namespace tensorflow {
namespace test {
namespace graph {
Node* Send(Graph* g, Node* input, const std::string& tensor,
const std::string& sender, const uint64_t sender_incarnation,
const std::string& receiver) {
Node* ret;
TF_CHECK_OK(NodeBuilder(g->NewName("n"), "_Send")
.Input(input, 0)
.Attr("tensor_name", tensor)
.Attr("send_device", sender)
.Attr("send_device_incarnation",
static_cast<int64_t>(sender_incarnation))
.Attr("recv_device", receiver)
.Finalize(g, &ret));
return ret;
}
Node* Recv(Graph* g, const std::string& tensor, const std::string& type,
const std::string& sender, const uint64_t sender_incarnation,
const std::string& receiver) {
Node* ret;
DataType dtype;
CHECK(DataTypeFromString(type, &dtype));
TF_CHECK_OK(NodeBuilder(g->NewName("n"), "_Recv")
.Attr("tensor_type", dtype)
.Attr("tensor_name", tensor)
.Attr("send_device", sender)
.Attr("send_device_incarnation",
static_cast<int64_t>(sender_incarnation))
.Attr("recv_device", receiver)
.Finalize(g, &ret));
return ret;
}
Node* Constant(Graph* g, const Tensor& tensor) {
Node* ret;
TF_CHECK_OK(NodeBuilder(g->NewName("n"), "Const")
.Attr("dtype", tensor.dtype())
.Attr("value", tensor)
.Finalize(g, &ret));
return ret;
}
Node* Constant(Graph* g, const Tensor& tensor, const std::string& name) {
Node* ret;
TF_CHECK_OK(NodeBuilder(name, "Const")
.Attr("dtype", tensor.dtype())
.Attr("value", tensor)
.Finalize(g, &ret));
return ret;
}
Node* HostConstant(Graph* g, const Tensor& tensor) {
return HostConstant(g, tensor, g->NewName("n"));
}
Node* HostConstant(Graph* g, const Tensor& tensor, const std::string& name) {
Node* ret;
TF_CHECK_OK(NodeBuilder(name, "HostConst")
.Attr("dtype", tensor.dtype())
.Attr("value", tensor)
.Finalize(g, &ret));
return ret;
}
Node* Var(Graph* g, const DataType dtype, const TensorShape& shape) {
Node* ret;
TF_CHECK_OK(NodeBuilder(g->NewName("n"), "Variable")
.Attr("dtype", dtype)
.Attr("shape", shape)
.Finalize(g, &ret));
return ret;
}
Node* Var(Graph* g, const DataType dtype, const TensorShape& shape,
const std::string& name) {
Node* ret;
TF_CHECK_OK(NodeBuilder(name, "Variable")
.Attr("dtype", dtype)
.Attr("shape", shape)
.Finalize(g, &ret));
return ret;
}
Node* Assign(Graph* g, Node* var, Node* val) {
Node* ret;
TF_CHECK_OK(NodeBuilder(g->NewName("n"), "Assign")
.Input(var)
.Input(val)
.Attr("use_locking", true)
.Finalize(g, &ret));
return ret;
}
Node* Cumsum(Graph* g, Node* data, Node* axes, bool exclusive, bool reverse) {
Node* ret;
TF_CHECK_OK(NodeBuilder(g->NewName("n"), "Cumsum")
.Input(data)
.Input(axes)
.Attr("exclusive", exclusive)
.Attr("reverse", reverse)
.Finalize(g, &ret));
return ret;
}
Node* Reduce(Graph* g, const std::string& reduce, Node* data, Node* axes,
bool keep_dims) {
Node* ret;
TF_CHECK_OK(NodeBuilder(g->NewName("n"), reduce, g->op_registry())
.Input(data)
.Input(axes)
.Attr("keep_dims", keep_dims)
.Finalize(g, &ret));
return ret;
}
Node* QuantizeToUINT8(Graph* g, Node* data) {
Node* ret;
TF_CHECK_OK(NodeBuilder(g->NewName("n"), "Quantize")
.Input(data)
.Attr("T", DT_QUINT8)
.Attr("max_range", 1.0f)
.Attr("min_range", -1.0f)
.Finalize(g, &ret));
return ret;
}
Node* Matmul(Graph* g, Node* in0, Node* in1, bool transpose_a,
bool transpose_b) {
Node* ret;
TF_CHECK_OK(NodeBuilder(g->NewName("n"), "MatMul")
.Input(in0)
.Input(in1)
.Attr("transpose_a", transpose_a)
.Attr("transpose_b", transpose_b)
.Finalize(g, &ret));
return ret;
}
Node* BatchMatmul(Graph* g, Node* in0, Node* in1, bool adj_x, bool adj_y) {
Node* ret;
TF_CHECK_OK(NodeBuilder(g->NewName("n"), "BatchMatMul")
.Input(in0)
.Input(in1)
.Attr("adj_x", adj_x)
.Attr("adj_y", adj_y)
.Finalize(g, &ret));
return ret;
}
Node* RandomNumberGenerator(const std::string& op, Graph* g, Node* input,
DataType dtype) {
Node* ret;
TF_CHECK_OK(NodeBuilder(g->NewName("n"), op, g->op_registry())
.Input(input)
.Attr("dtype", dtype)
.Attr("seed", 0)
.Finalize(g, &ret));
return ret;
}
Node* RandomUniform(Graph* g, Node* input, DataType dtype) {
return RandomNumberGenerator("RandomUniform", g, input, dtype);
}
Node* RandomGaussian(Graph* g, Node* input, DataType dtype) {
return RandomNumberGenerator("RandomStandardNormal", g, input, dtype);
}
Node* TruncatedNormal(Graph* g, Node* input, DataType dtype) {
return RandomNumberGenerator("TruncatedNormal", g, input, dtype);
}
Node* RandomGamma(Graph* g, Node* shape, Node* alpha) {
Node* ret;
TF_CHECK_OK(NodeBuilder(g->NewName("n"), "RandomGamma")
.Input(shape)
.Input(alpha)
.Attr("seed", 0)
.Finalize(g, &ret));
return ret;
}
Node* RandomPoisson(Graph* g, Node* shape, Node* lam) {
Node* ret;
TF_CHECK_OK(NodeBuilder(g->NewName("n"), "RandomPoisson")
.Input(shape)
.Input(lam)
.Attr("seed", 0)
.Finalize(g, &ret));
return ret;
}
Node* Unary(Graph* g, const std::string& func, Node* input, int index) {
Node* ret;
TF_CHECK_OK(NodeBuilder(g->NewName("n"), func, g->op_registry())
.Input(input, index)
.Finalize(g, &ret));
return ret;
}
Node* Binary(Graph* g, const std::string& func, Node* in0, Node* in1) {
Node* ret;
TF_CHECK_OK(NodeBuilder(g->NewName("n"), func, g->op_registry())
.Input(in0)
.Input(in1)
.Finalize(g, &ret));
return ret;
}
Node* Multi(Graph* g, const std::string& func, absl::Span<Node* const> ins) {
Node* ret;
auto b = NodeBuilder(g->NewName("n"), func, g->op_registry());
for (Node* n : ins) b = b.Input(n);
TF_CHECK_OK(b.Finalize(g, &ret));
return ret;
}
Node* Identity(Graph* g, Node* input, int index) {
Node* ret;
TF_CHECK_OK(NodeBuilder(g->NewName("n"), "Identity")
.Input(input, index)
.Finalize(g, &ret));
return ret;
}
Node* Add(Graph* g, Node* in0, Node* in1) { return Binary(g, "Add", in0, in1); }
Node* Reverse(Graph* g, Node* tensor, Node* axis) {
return Binary(g, "ReverseV2", tensor, axis);
}
Node* Roll(Graph* g, Node* input, Node* shift, Node* axis) {
Node* ret;
TF_CHECK_OK(NodeBuilder(g->NewName("n"), "Roll", g->op_registry())
.Input(input)
.Input(shift)
.Input(axis)
.Finalize(g, &ret));
return ret;
}
Node* Error(Graph* g, Node* input, const std::string& errmsg, bool log_error) {
Node* ret;
TF_CHECK_OK(NodeBuilder(g->NewName("n"), "Error")
.Input(input)
.Attr("message", errmsg)
.Attr("log_error", log_error)
.Finalize(g, &ret));
return ret;
}
Node* InvalidRefType(Graph* g, DataType out_type, DataType invalid_type) {
DCHECK(out_type != invalid_type);
Node* ret;
TF_CHECK_OK(NodeBuilder(g->NewName("n"), "InvalidRefType")
.Attr("TIn", out_type)
.Attr("TOut", invalid_type)
.Finalize(g, &ret));
return ret;
}
Node* Delay(Graph* g, Node* input, Microseconds delay_micros) {
Node* ret;
TF_CHECK_OK(NodeBuilder(g->NewName("n"), "Delay")
.Input(input)
.Attr("micros", delay_micros.value())
.Finalize(g, &ret));
return ret;
}
Node* NoOp(Graph* g, const std::vector<Node*>& control_inputs) {
Node* ret;
TF_CHECK_OK(NodeBuilder(g->NewName("n"), "NoOp")
.ControlInputs(control_inputs)
.Finalize(g, &ret));
return ret;
}
Node* Switch(Graph* g, Node* in0, Node* in1) {
Node* ret;
TF_CHECK_OK(NodeBuilder(g->NewName("n"), "Switch")
.Input(in0)
.Input(in1)
.Finalize(g, &ret));
return ret;
}
Node* Enter(Graph* g, Node* input, const std::string& frame_name) {
Node* ret;
TF_CHECK_OK(NodeBuilder(g->NewName("n"), "Enter")
.Input(input)
.Attr("frame_name", frame_name)
.Finalize(g, &ret));
return ret;
}
Node* Exit(Graph* g, Node* input) {
Node* ret;
TF_CHECK_OK(
NodeBuilder(g->NewName("n"), "Exit").Input(input).Finalize(g, &ret));
return ret;
}
Node* Merge(Graph* g, Node* in0, Node* in1) {
Node* ret;
TF_CHECK_OK(NodeBuilder(g->NewName("n"), "Merge")
.Input({in0, in1})
.Finalize(g, &ret));
return ret;
}
Node* Merge(Graph* g, Node* in0, absl::Span<const std::string> remaining_in) {
std::vector<NodeBuilder::NodeOut> inputs;
inputs.reserve(remaining_in.size() + 1);
inputs.emplace_back(in0);
for (const std::string& in_name : remaining_in) {
inputs.emplace_back(in_name, 0, inputs[0].dt);
}
Node* ret;
TF_CHECK_OK(
NodeBuilder(g->NewName("n"), "Merge").Input(inputs).Finalize(g, &ret));
return ret;
}
Node* Concat(Graph* g, Node* concat_dim, absl::Span<Node* const> tensors) {
std::vector<NodeBuilder::NodeOut> nodeouts;
nodeouts.reserve(tensors.size());
for (auto const t : tensors) {
nodeouts.emplace_back(t);
}
Node* ret;
TF_CHECK_OK(NodeBuilder(g->NewName("n"), "Concat")
.Input(concat_dim)
.Input(nodeouts)
.Finalize(g, &ret));
return ret;
}
Node* ConcatV2(Graph* g, absl::Span<Node* const> tensors, Node* concat_dim) {
std::vector<NodeBuilder::NodeOut> nodeouts;
nodeouts.reserve(tensors.size());
for (auto const t : tensors) {
nodeouts.emplace_back(t);
}
Node* ret;
TF_CHECK_OK(NodeBuilder(g->NewName("n"), "ConcatV2")
.Input(nodeouts)
.Input(concat_dim)
.Finalize(g, &ret));
return ret;
}
Node* Next(Graph* g, const std::string& name, Node* input) {
Node* ret;
TF_CHECK_OK(
NodeBuilder(name, "NextIteration").Input(input).Finalize(g, &ret));
return ret;
}
Node* LoopCond(Graph* g, Node* input) {
Node* ret;
TF_CHECK_OK(
NodeBuilder(g->NewName("n"), "LoopCond").Input(input).Finalize(g, &ret));
return ret;
}
Node* Less(Graph* g, Node* in0, Node* in1) {
return Binary(g, "Less", in0, in1);
}
Node* Select(Graph* g, Node* c, Node* inx, Node* iny) {
Node* ret;
TF_CHECK_OK(NodeBuilder(g->NewName("n"), "Select")
.Input(c)
.Input(inx)
.Input(iny)
.Finalize(g, &ret));
return ret;
}
Node* Cast(Graph* g, Node* in, DataType dst) {
Node* ret;
TF_CHECK_OK(NodeBuilder(g->NewName("n"), "Cast")
.Input(in)
.Attr("DstT", dst)
.Finalize(g, &ret));
return ret;
}
Node* Gather(Graph* g, Node* in0, Node* in1, Node* axis) {
Node* ret;
TF_CHECK_OK(NodeBuilder(g->NewName("n"), "GatherV2")
.Input(in0)
.Input(in1)
.Input(axis)
.Finalize(g, &ret));
return ret;
}
Node* GetSessionTensor(Graph* g, Node* in) {
Node* ret;
TF_CHECK_OK(NodeBuilder(g->NewName("n"), "GetSessionTensor")
.Input(in, 0)
.Attr("dtype", DT_FLOAT)
.Finalize(g, &ret));
return ret;
}
Node* Relu(Graph* g, Node* in) {
Node* ret;
TF_CHECK_OK(NodeBuilder(g->NewName("n"), "Relu")
.Input(in, 0)
.Attr("T", DT_FLOAT)
.Finalize(g, &ret));
return ret;
}
Node* Relu6(Graph* g, Node* in) {
Node* ret;
TF_CHECK_OK(NodeBuilder(g->NewName("n"), "Relu6")
.Input(in, 0)
.Attr("T", DT_FLOAT)
.Finalize(g, &ret));
return ret;
}
Node* BiasAdd(Graph* g, Node* value, Node* bias) {
Node* ret;
TF_CHECK_OK(NodeBuilder(g->NewName("n"), "BiasAdd")
.Input(value)
.Input(bias)
.Attr("T", DT_FLOAT)
.Finalize(g, &ret));
return ret;
}
Node* Conv2D(Graph* g, Node* in0, Node* in1) {
Node* ret;
TF_CHECK_OK(NodeBuilder(g->NewName("n"), "Conv2D")
.Input(in0)
.Input(in1)
.Attr("T", DT_FLOAT)
.Attr("strides", {1, 1, 1, 1})
.Attr("padding", "SAME")
.Finalize(g, &ret));
return ret;
}
Node* Diag(Graph* g, Node* in, DataType type) {
Node* ret;
TF_CHECK_OK(NodeBuilder(g->NewName("n"), "Diag")
.Input(in)
.Attr("T", type)
.Finalize(g, &ret));
return ret;
}
Node* DiagPart(Graph* g, Node* in, DataType type) {
Node* ret;
TF_CHECK_OK(NodeBuilder(g->NewName("n"), "DiagPart")
.Input(in)
.Attr("T", type)
.Finalize(g, &ret));
return ret;
}
Node* CheckNumerics(Graph* g, Node* in, const std::string& message) {
Node* ret;
TF_CHECK_OK(NodeBuilder(g->NewName("n"), "CheckNumerics")
.Input(in)
.Attr("message", message)
.Finalize(g, &ret));
return ret;
}
Node* Arg(Graph* g, int64_t index, DataType type) {
Node* ret;
TF_CHECK_OK(NodeBuilder(g->NewName("n"), "_Arg")
.Attr("T", type)
.Attr("index", index)
.Finalize(g, &ret));
return ret;
}
Node* Retval(Graph* g, int64_t index, Node* in, int64_t in_index) {
Node* ret;
TF_CHECK_OK(NodeBuilder(g->NewName("n"), "_Retval")
.Input(in, in_index)
.Attr("index", index)
.Finalize(g, &ret));
return ret;
}
void ToGraphDef(Graph* g, GraphDef* gdef) { g->ToGraphDef(gdef); }
} // end namespace graph
} // end namespace test
} // end namespace tensorflow
tensorflow/core/graph/testlib.h
+231
View File
@@ -0,0 +1,231 @@
/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
// DEPRECATED: Use the C++ API defined in tensorflow/cc instead.
#ifndef TENSORFLOW_CORE_GRAPH_TESTLIB_H_
#define TENSORFLOW_CORE_GRAPH_TESTLIB_H_
#include <string>
#include <vector>
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/graph/types.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace test {
namespace graph {
// Converts "g" into its corresponding GraphDef "def".
ABSL_DEPRECATED("Call g->ToGraphDef(def) instead.")
void ToGraphDef(Graph* g, GraphDef* def);
// A few helpers to construct a graph.
// Adds a node in "g" producing a constant "tensor".
Node* Constant(Graph* g, const Tensor& tensor);
Node* Constant(Graph* g, const Tensor& tensor, const std::string& name);
// Adds a node in "g" producing a constant "tensor" on the host.
// The given node which, unlike the regular Constant above, always
// stores its output on the host. This is necessary for use
// in GPU tests where the test Op in question runs on the device
// but requires some arguments to be pinned to the host.
Node* HostConstant(Graph* g, const Tensor& tensor);
Node* HostConstant(Graph* g, const Tensor& tensor, const std::string& name);
// Adds a variable in "g" of the given "shape" and "dtype".
Node* Var(Graph* g, DataType dtype, const TensorShape& shape);
Node* Var(Graph* g, DataType dtype, const TensorShape& shape,
const std::string& name);
// Adds an assign node in "g" which assigns "val" into "var".
Node* Assign(Graph* g, Node* var, Node* val);
// Adds a send node "g" sending "input" as a named "tensor" from
// "sender" to "receiver".
Node* Send(Graph* g, Node* input, const std::string& tensor,
const std::string& sender, uint64_t sender_incarnation,
const std::string& receiver);
// Adds a recv node in "g" receiving a named "tensor" from "sender"
// to "receiver".
Node* Recv(Graph* g, const std::string& tensor, const std::string& type,
const std::string& sender, uint64_t sender_incarnation,
const std::string& receiver);
// Adds a cumsum "node" in "g" doing cumsum(data, axes).
Node* Cumsum(Graph* g, Node* data, Node* axes, bool exclusive = false,
bool reverse = false);
// Adds a reduction "node" in "g" doing sum(data, axes). "reduce" is
// a reduction, e.g., Sum, Max, Min, Mean, etc.
Node* Reduce(Graph* g, const std::string& reduce, Node* data, Node* axes,
bool keep_dims = false);
// Adds a Matmul node in g doing in0.contract(in1).
Node* Matmul(Graph* g, Node* in0, Node* in1, bool transpose_a,
bool transpose_b);
// Adds a Matmul node in g doing in0.contract(in1).
Node* BatchMatmul(Graph* g, Node* in0, Node* in1, bool adj_x, bool adj_y);
// Adds a Quantize node into g that quantize floats into QUINT8. The range of
// the input float tensor is assumed to be [-1, 1].
Node* QuantizeToUINT8(Graph* g, Node* data);
// Adds a unary function "func" "node" in "g" taking "input".
Node* Unary(Graph* g, const std::string& func, Node* input, int index = 0);
// Adds an identity node in "g" taking "input" and producing an
// identity copy.
Node* Identity(Graph* g, Node* input, int index = 0);
// Adds a binary function "func" node in "g" taking "in0" and "in1".
Node* Binary(Graph* g, const std::string& func, Node* in0, Node* in1);
// Adds a function "func" node in "g" taking inputs "ins".
Node* Multi(Graph* g, const std::string& func, absl::Span<Node* const> ins);
// Adds a binary add node in "g" doing in0 + in1.
Node* Add(Graph* g, Node* in0, Node* in1);
// Reverses <axis> dimensions of <tensor>>
Node* Reverse(Graph* g, Node* tensor, Node* axis);
// Generates random unit uniform distribution of the input shape.
Node* RandomUniform(Graph* g, Node* input, DataType dtype);
// Generates random unit normal distribution of the input shape.
Node* RandomGaussian(Graph* g, Node* input, DataType dtype);
// Generates random gamma distribution with the given shape and alpha[s].
// Output dtype determined by alpha.
Node* RandomGamma(Graph* g, Node* shape, Node* alpha);
// Generates random poisson distribution with the given shape and lam[s].
// Output dtype determined by lam.
Node* RandomPoisson(Graph* g, Node* shape, Node* lam);
// Rolls tensor by an offset of <shift> along the corresponding
// <axis> dimensions.
Node* Roll(Graph* g, Node* input, Node* shift, Node* axis);
// Generates random parameters from the truncated standard normal distribution
// of the input shape
Node* TruncatedNormal(Graph* g, Node* input, DataType dtype);
// Adds an error node in "g". The node's computation always
// generates an error with the given error message "errmsg".
Node* Error(Graph* g, Node* input, const std::string& errmsg,
bool log_error = false);
// Adds a node that generates a invalid ref output.
Node* InvalidRefType(Graph* g, DataType out_type, DataType invalid_type);
// Adds a node in "g". Its Compute() sleeps a while and outputs the
// input (i.e., same as identity).
Node* Delay(Graph* g, Node* input, Microseconds delay_micros);
// Adds a no-op "node" in "g", with control inputs from all nodes in
// control_inputs vector.
Node* NoOp(Graph* g, const std::vector<Node*>& control_inputs);
// Adds a Switch node in "g". If "in1" is true, it forwards "in0" to
// output 1. Otherwise, it forwards "in0" to output 0.
Node* Switch(Graph* g, Node* in0, Node* in1);
// Adds an Enter node in "g", which enters a new frame.
Node* Enter(Graph* g, Node* input, const std::string& frame_name);
// Adds an Exit node in "g", which exits a frame.
Node* Exit(Graph* g, Node* input);
// Adds a Merge node in "g" with two inputs "in0" and "in1".
Node* Merge(Graph* g, Node* in0, Node* in1);
// Adds a Merge node in "g". The first input is "in0", the remaining
// inputs are only given by their names in remaining_in.
Node* Merge(Graph* g, Node* in0, absl::Span<const std::string> remaining_in);
// Adds a NextIteration node in "g", which makes its input available
// to the next iteration.
Node* Next(Graph* g, const std::string& name, Node* input);
// Adds a LoopCond node in "g", representing the "pivot" termination
// condition of a loop.
Node* LoopCond(Graph* g, Node* input);
// Adds a less node in "g", which returns true iff "in0" < "in1".
Node* Less(Graph* g, Node* in0, Node* in1);
// Adds a select node in "g", which outputs either "inx" or "iny"
// depending on the boolean value of "c".
Node* Select(Graph* g, Node* c, Node* inx, Node* iny);
// Casts "in" into data type "dst".
Node* Cast(Graph* g, Node* in, DataType dst);
// Perform gather op on params "in0" with indices "in1" and axis "axis".
Node* Gather(Graph* g, Node* in0, Node* in1, Node* axis);
// Gets a tensor stored in the session state.
Node* GetSessionTensor(Graph* g, Node* in);
// Adds a Concat node in "g". The first input is "concat_dim", the
// dimension to concatenate on, and the tensors to concatenate are
// given in "tensors".
Node* Concat(Graph* g, Node* concat_dim, absl::Span<Node* const> tensors);
// Adds a ConcatV2 node in "g". The last input is "concat_dim", the
// dimension to concatenate on, and the tensors to concatenate are
// given in "tensors".
Node* ConcatV2(Graph* g, absl::Span<Node* const> tensors, Node* concat_dim);
// Add a Relu node in "g".
Node* Relu(Graph* g, Node* in);
// Add a Relu6 node in "g".
Node* Relu6(Graph* g, Node* in);
// Add a BiasAdd node in "g".
Node* BiasAdd(Graph* g, Node* value, Node* bias);
// Add a Conv2D node in "g".
Node* Conv2D(Graph* g, Node* in0, Node* in1);
// Add a Diag node in "g".
Node* Diag(Graph* g, Node* in, DataType type);
// Add a DiagPart node in "g".
Node* DiagPart(Graph* g, Node* in, DataType type);
// Add a CheckNumerics node in "g".
Node* CheckNumerics(Graph* g, Node* in, const std::string& message);
// Add an _Arg node in "g".
Node* Arg(Graph* g, int64_t index, DataType type);
// Add a _Retval node in "g".
Node* Retval(Graph* g, int64_t index, Node* in, int64_t in_index = 0);
} // end namespace graph
} // end namespace test
} // end namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPH_TESTLIB_H_
tensorflow/core/graph/types.h
+35
View File
@@ -0,0 +1,35 @@
/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
Licensed 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 TENSORFLOW_CORE_GRAPH_TYPES_H_
#define TENSORFLOW_CORE_GRAPH_TYPES_H_
#include "tensorflow/core/lib/gtl/int_type.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
// We model running time in microseconds.
TSL_LIB_GTL_DEFINE_INT_TYPE(Microseconds, int64_t);
// We can also model running time in nanoseconds for more accuracy.
TSL_LIB_GTL_DEFINE_INT_TYPE(Nanoseconds, int64_t);
// We model size in bytes.
TSL_LIB_GTL_DEFINE_INT_TYPE(Bytes, int64_t);
} // namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPH_TYPES_H_
tensorflow/core/graph/validate.cc
+133
View File
@@ -0,0 +1,133 @@
/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#include "tensorflow/core/graph/validate.h"
#include "absl/container/flat_hash_set.h"
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "tensorflow/core/framework/graph_def_util.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/framework/op_def_util.h"
#include "tensorflow/core/framework/versions.pb.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace graph {
absl::Status ValidateGraphDef(const GraphDef& graph_def,
const OpRegistryInterface& op_registry) {
absl::Status s;
const int version = graph_def.versions().producer();
for (const NodeDef& node_def : graph_def.node()) {
// Look up the OpDef for the node_def's op name.
const OpDef* op_def;
TF_RETURN_IF_ERROR(op_registry.LookUpOpDef(node_def.op(), &op_def));
TF_RETURN_IF_ERROR(ValidateNodeDef(node_def, *op_def));
TF_RETURN_IF_ERROR(CheckOpDeprecation(*op_def, version));
}
return s;
}
absl::Status ValidateGraphDefAgainstOpRegistry(
const GraphDef& graph_def, const OpRegistryInterface& op_registry) {
GraphDef copy(graph_def);
TF_RETURN_IF_ERROR(AddDefaultAttrsToGraphDef(&copy, op_registry, 0));
return ValidateGraphDef(copy, op_registry);
}
absl::Status ValidateGraphDefAgainstOpList(const GraphDef& graph_def,
const OpList& op_list) {
OpListOpRegistry registry(&op_list);
return ValidateGraphDefAgainstOpRegistry(graph_def, registry);
}
void GetOpListForValidation(OpList* op_list, const OpRegistry& op_registry) {
op_registry.Export(false, op_list);
RemoveDescriptionsFromOpList(op_list);
}
absl::Status ValidateGraphHasNoCycle(const Graph& graph) {
// A node is ready when all of its inputs have been visited.
std::vector<const Node*> ready;
std::vector<int> pending_count(graph.num_node_ids(), 0);
for (int i = 0; i < graph.num_node_ids(); ++i) {
const Node* n = graph.FindNodeId(i);
if (n == nullptr) continue;
pending_count[i] = n->in_edges().size();
if (n->IsMerge()) {
// While-loop cycles are legal cycles so we manually adjust the
// pending_count to make sure that the loop is visited.
for (const Edge* e : n->in_edges()) {
if (!e->IsControlEdge() && e->src()->IsNextIteration()) {
pending_count[i]--;
}
}
}
if (pending_count[i] == 0) {
ready.push_back(n);
}
}
int processed = 0;
while (!ready.empty()) {
const Node* node = ready.back();
ready.pop_back();
++processed;
for (const Edge* out : node->out_edges()) {
const int output_id = out->dst()->id();
pending_count[output_id]--;
if (pending_count[output_id] == 0) {
ready.push_back(out->dst());
}
}
}
if (processed < graph.num_nodes()) {
std::vector<std::string> nodes_in_cycle;
for (int i = 0; i < pending_count.size() && nodes_in_cycle.size() < 3;
++i) {
if (pending_count[i] != 0) {
nodes_in_cycle.push_back(graph.FindNodeId(i)->name());
}
}
return absl::InvalidArgumentError(absl::StrCat(
"Graph is invalid, contains a cycle with ",
graph.num_nodes() - processed,
" nodes, including: ", absl::StrJoin(nodes_in_cycle, ", ")));
}
return absl::OkStatus();
}
absl::Status VerifyNoDuplicateNodeNames(const GraphDef& graph) {
absl::flat_hash_set<absl::string_view> nodes;
for (const auto& node : graph.node()) {
if (nodes.contains(node.name())) {
return absl::AlreadyExistsError(
absl::StrCat("Node already exists: ", node.name()));
}
nodes.insert(node.name());
}
return absl::OkStatus();
}
} // namespace graph
} // namespace tensorflow
tensorflow/core/graph/validate.h
+68
View File
@@ -0,0 +1,68 @@
/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
Licensed 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 TENSORFLOW_CORE_GRAPH_VALIDATE_H_
#define TENSORFLOW_CORE_GRAPH_VALIDATE_H_
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/lib/core/status.h"
namespace tensorflow {
namespace graph {
// Returns OK if every NodeDef in `graph_def` is valid with respect to
// its corresponding OpDef (as defined by ValidateNodeDef()) as
// registered in `op_registry`. Also checks for deprecated ops.
//
// REQUIRES:
// * `op_registry` is not nullptr.
// * `graph_def` has default attrs filled in (see AddDefaultAttrsToGraphDef()).
absl::Status ValidateGraphDef(const GraphDef& graph_def,
const OpRegistryInterface& op_registry);
// Like ValidateGraphDef() except it makes a copy of `graph_def` and calls
// AddDefaultAttrsToGraphDef() on the copy, removing that requirement from the
// caller.
absl::Status ValidateGraphDefAgainstOpRegistry(
const GraphDef& graph_def, const OpRegistryInterface& op_registry);
// Like ValidateGraphDefAgainstOpRegistry() except it takes an OpList
// instead of an OpRegistryInterface. Note that the OpList need not
// have descriptions, which can be a big space savings, see
// GetOpListForValidation() below.
absl::Status ValidateGraphDefAgainstOpList(const GraphDef& graph_def,
const OpList& op_list);
// Get an OpList from `*op_registry` with all the descriptions removed.
void GetOpListForValidation(
OpList* op_list, const OpRegistry& op_registry = *OpRegistry::Global());
// Validate that the graph has no cycle except for legal while loop cycles.
// This traverses the specified nodes in topological order to verify there are
// no cycles. Starting with inputless nodes, it visits nodes whose inputs have
// all been visited, and counts the total number of visited nodes. If there is a
// cycle, nodes in the cycle will never be visited, and the visited count will
// be less than the total node count.
absl::Status ValidateGraphHasNoCycle(const Graph& graph);
// Returns OK if the graph has no duplicate node names.
absl::Status VerifyNoDuplicateNodeNames(const GraphDef& graph);
} // namespace graph
} // namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPH_VALIDATE_H_
tensorflow/core/graph/validate_test.cc
+258
View File
@@ -0,0 +1,258 @@
/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#include "tensorflow/core/graph/validate.h"
#include <string>
#include "tensorflow/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/graph_def_util.h"
#include "tensorflow/core/framework/op_def_builder.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/graph/graph_def_builder.h"
#include "tensorflow/core/graph/subgraph.h"
#include "tensorflow/core/kernels/ops_util.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/lib/strings/str_util.h"
#include "tensorflow/core/platform/status_matchers.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace {
REGISTER_OP("FloatInput").Output("o: float");
REGISTER_OP("Int32Input").Output("o: int32");
TEST(ValidateGraphDefTest, TestValidGraph) {
const std::string graph_def_str =
"node { name: 'A' op: 'FloatInput' }"
"node { name: 'B' op: 'FloatInput' }"
"node { name: 'C' op: 'Mul' attr { key: 'T' value { type: DT_FLOAT } }"
" input: ['A', 'B'] }";
GraphDef graph_def;
auto parser = protobuf::TextFormat::Parser();
CHECK(parser.MergeFromString(graph_def_str, &graph_def)) << graph_def_str;
TF_ASSERT_OK(graph::ValidateGraphDef(graph_def, *OpRegistry::Global()));
}
TEST(ValidateGraphDefTest, GraphWithUnspecifiedDefaultAttr) {
const std::string graph_def_str =
"node { name: 'A' op: 'FloatInput' }"
"node { name: 'B' op: 'Int32Input' }"
"node { "
" name: 'C' op: 'Sum' "
" attr { key: 'T' value { type: DT_FLOAT } }"
" input: ['A', 'B'] "
"}";
GraphDef graph_def;
auto parser = protobuf::TextFormat::Parser();
CHECK(parser.MergeFromString(graph_def_str, &graph_def)) << graph_def_str;
absl::Status s = graph::ValidateGraphDef(graph_def, *OpRegistry::Global());
EXPECT_FALSE(s.ok());
EXPECT_TRUE(absl::StrContains(s.ToString(), "NodeDef missing attr"));
// Add the defaults.
TF_ASSERT_OK(AddDefaultAttrsToGraphDef(&graph_def, *OpRegistry::Global(), 0));
// Validation should succeed.
TF_ASSERT_OK(graph::ValidateGraphDef(graph_def, *OpRegistry::Global()));
}
TEST(ValidateGraphDefTest, GraphWithUnspecifiedRequiredAttr) {
// "DstT" attribute is missing.
const std::string graph_def_str =
"node { name: 'A' op: 'FloatInput' }"
"node { "
" name: 'B' op: 'Cast' "
" attr { key: 'SrcT' value { type: DT_FLOAT } }"
" input: ['A'] "
"}";
GraphDef graph_def;
auto parser = protobuf::TextFormat::Parser();
CHECK(parser.MergeFromString(graph_def_str, &graph_def)) << graph_def_str;
absl::Status s = graph::ValidateGraphDef(graph_def, *OpRegistry::Global());
EXPECT_FALSE(s.ok());
EXPECT_TRUE(absl::StrContains(s.ToString(), "NodeDef missing attr"));
// Add the defaults.
TF_ASSERT_OK(AddDefaultAttrsToGraphDef(&graph_def, *OpRegistry::Global(), 0));
// Validation should still fail.
s = graph::ValidateGraphDef(graph_def, *OpRegistry::Global());
EXPECT_FALSE(s.ok());
EXPECT_TRUE(absl::StrContains(s.ToString(), "NodeDef missing attr"));
}
TEST(ValidateGraphDefAgainstOpListTest, GraphWithOpOnlyInOpList) {
OpRegistrationData op_reg_data;
TF_ASSERT_OK(OpDefBuilder("UniqueSnowflake").Finalize(&op_reg_data));
OpList op_list;
*op_list.add_op() = op_reg_data.op_def;
const std::string graph_def_str = "node { name: 'A' op: 'UniqueSnowflake' }";
GraphDef graph_def;
auto parser = protobuf::TextFormat::Parser();
CHECK(parser.MergeFromString(graph_def_str, &graph_def)) << graph_def_str;
TF_ASSERT_OK(graph::ValidateGraphDefAgainstOpList(graph_def, op_list));
}
TEST(ValidateGraphDefAgainstOpListTest, GraphWithGlobalOpNotInOpList) {
OpRegistrationData op_reg_data;
TF_ASSERT_OK(OpDefBuilder("NotAnywhere").Finalize(&op_reg_data));
OpList op_list;
*op_list.add_op() = op_reg_data.op_def;
const std::string graph_def_str = "node { name: 'A' op: 'FloatInput' }";
GraphDef graph_def;
auto parser = protobuf::TextFormat::Parser();
CHECK(parser.MergeFromString(graph_def_str, &graph_def)) << graph_def_str;
ASSERT_FALSE(graph::ValidateGraphDefAgainstOpList(graph_def, op_list).ok());
}
REGISTER_OP("HasDocs").Doc("This is in the summary.");
TEST(GetOpListForValidationTest, ShouldStripDocs) {
bool found_float = false;
bool found_int32 = false;
bool found_has_docs = false;
OpList op_list;
graph::GetOpListForValidation(&op_list);
for (const OpDef& op_def : op_list.op()) {
if (op_def.name() == "FloatInput") {
EXPECT_FALSE(found_float);
found_float = true;
}
if (op_def.name() == "Int32Input") {
EXPECT_FALSE(found_int32);
found_int32 = true;
}
if (op_def.name() == "HasDocs") {
EXPECT_FALSE(found_has_docs);
found_has_docs = true;
EXPECT_TRUE(op_def.summary().empty());
}
}
EXPECT_TRUE(found_float);
EXPECT_TRUE(found_int32);
EXPECT_TRUE(found_has_docs);
}
TEST(VerifyNoDuplicateNodeNames, NoDuplicateNodeNames) {
const std::string graph_def_str =
"node { name: 'A' op: 'FloatInput' }"
"node { name: 'B' op: 'Int32Input' }"
"node { "
" name: 'C' op: 'Sum' "
" attr { key: 'T' value { type: DT_FLOAT } }"
" input: ['A', 'B'] "
"}";
GraphDef graph_def;
auto parser = protobuf::TextFormat::Parser();
CHECK(parser.MergeFromString(graph_def_str, &graph_def)) << graph_def_str;
TF_ASSERT_OK(graph::VerifyNoDuplicateNodeNames(graph_def));
}
TEST(VerifyNoDuplicateNodeNames, DuplicateNodeNames) {
const std::string graph_def_str =
"node { name: 'A' op: 'FloatInput' }"
"node { name: 'A' op: 'Int32Input' }"
"node { "
" name: 'C' op: 'Sum' "
" attr { key: 'T' value { type: DT_FLOAT } }"
" input: ['A', 'A'] "
"}";
GraphDef graph_def;
auto parser = protobuf::TextFormat::Parser();
CHECK(parser.MergeFromString(graph_def_str, &graph_def)) << graph_def_str;
EXPECT_EQ(graph::VerifyNoDuplicateNodeNames(graph_def).code(),
tensorflow::error::ALREADY_EXISTS);
}
TEST(ValidateGraphHasNoCycleTest, NoCyclePasses) {
const std::string graph_def_str =
"node { name: 'A' op: 'FloatInput' }"
"node { name: 'B' op: 'FloatInput' }"
"node { name: 'C' op: 'Mul' attr { key: 'T' value { type: DT_FLOAT } }"
" input: ['A', 'B'] }";
GraphDef graph_def;
auto parser = protobuf::TextFormat::Parser();
CHECK(parser.MergeFromString(graph_def_str, &graph_def)) << graph_def_str;
Graph graph(OpRegistry::Global());
GraphConstructorOptions opts;
TF_ASSERT_OK(ConvertGraphDefToGraph(opts, graph_def, &graph));
TF_EXPECT_OK(graph::ValidateGraphHasNoCycle(graph));
}
TEST(ValidateGraphHasNoCycleTest, NoCycleWithMergePasses) {
const std::string graph_def_str =
R"EOF(
node { name: 'A' op: 'FloatInput' }
node { name: 'merge' op: 'Merge' input: [ 'A:0', 'next:0' ]
attr { key: "N" value: { i: 2 } }
attr { key: "T" value: { type: DT_FLOAT } } }
node { name: 'B' op: 'Mul'
attr { key: 'T' value { type: DT_FLOAT } }
input: [ 'merge:0', 'merge:0' ] }
node { name: 'next' op: 'NextIteration' input: ['B:0']
attr { key: "T" value: { type: DT_FLOAT } } }
)EOF";
GraphDef graph_def;
auto parser = protobuf::TextFormat::Parser();
CHECK(parser.MergeFromString(graph_def_str, &graph_def)) << graph_def_str;
Graph graph(OpRegistry::Global());
GraphConstructorOptions opts;
TF_ASSERT_OK(ConvertGraphDefToGraph(opts, graph_def, &graph));
TF_EXPECT_OK(graph::ValidateGraphHasNoCycle(graph));
}
Node* AddNodeFromNodeDef(Graph& graph, const std::string& name,
const std::string& node_type, int num_inputs) {
auto builder = NodeDefBuilder(name, node_type);
for (int i = 0; i < num_inputs; ++i) {
builder = builder.Input(absl::StrCat("node_", i), i, DT_FLOAT);
}
NodeDef node_def;
TF_CHECK_OK(builder.Finalize(&node_def));
absl::Status s;
Node* node = graph.AddNode(node_def, &s);
TF_CHECK_OK(s);
return node;
}
TEST(ValidateGraphHasNoCycleTest, CycleFails) {
// Need to construct graph explicitly, since GraphDefToGraph has its own
// cycle validation routine.
Graph graph(OpRegistry::Global());
Node* a = AddNodeFromNodeDef(graph, "A", "FloatInput", 0);
Node* c = AddNodeFromNodeDef(graph, "B", "Mul", 2);
graph.AddEdge(a, 0, c, 0);
graph.AddEdge(c, 0, c, 1); // Loop from C->C.
EXPECT_THAT(
graph::ValidateGraphHasNoCycle(graph),
absl_testing::StatusIs(
tsl::error::Code::INVALID_ARGUMENT,
::testing::ContainsRegex("Graph is invalid, contains a cycle")));
}
} // namespace
} // namespace tensorflow
tensorflow/core/graph/while_context.cc
+38
View File
@@ -0,0 +1,38 @@
/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#include "tensorflow/core/graph/while_context.h"
namespace tensorflow {
WhileContext::WhileContext(absl::string_view frame_name,
std::vector<Node*> enter_nodes,
std::vector<Node*> exit_nodes,
OutputTensor cond_output,
std::vector<OutputTensor> body_inputs,
std::vector<OutputTensor> body_outputs)
: frame_name_(frame_name),
enter_nodes_(std::move(enter_nodes)),
exit_nodes_(std::move(exit_nodes)),
cond_output_(cond_output),
body_inputs_(std::move(body_inputs)),
body_outputs_(std::move(body_outputs)) {
const size_t num_loop_vars = enter_nodes_.size();
DCHECK_EQ(exit_nodes_.size(), num_loop_vars);
DCHECK_EQ(body_inputs_.size(), num_loop_vars);
DCHECK_EQ(body_outputs_.size(), num_loop_vars);
}
} // namespace tensorflow
tensorflow/core/graph/while_context.h
+76
View File
@@ -0,0 +1,76 @@
/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
Licensed 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 TENSORFLOW_CORE_GRAPH_WHILE_CONTEXT_H_
#define TENSORFLOW_CORE_GRAPH_WHILE_CONTEXT_H_
#include "tensorflow/core/graph/graph.h"
namespace tensorflow {
// Information about a while loop. Every user-defined while loop has an
// associated WhileContext, i.e., there is a WhileContext for every execution
// frame. Created with the while loop and used during gradient
// construction. Note that the gradient graph of while loop contains while loops
// itself, but these do not generate separate WhileContexts.
//
// TODO(skyewm): this is currently insufficient to handle nested loops and
// conditionals (and possibly other requirements). This may change a lot in the
// future to support these features.
//
// TODO(skyewm): de/serialize in MetaGraphDef so imported while loops will be
// differentiable. Figure out backwards compatibility story.
class WhileContext {
public:
WhileContext(absl::string_view frame_name, std::vector<Node*> enter_nodes,
std::vector<Node*> exit_nodes, OutputTensor cond_output,
std::vector<OutputTensor> body_inputs,
std::vector<OutputTensor> body_outputs);
const std::string& frame_name() const { return frame_name_; }
const std::vector<Node*>& enter_nodes() const { return enter_nodes_; }
const std::vector<Node*>& exit_nodes() const { return exit_nodes_; }
const OutputTensor& cond_output() const { return cond_output_; }
const std::vector<OutputTensor>& body_inputs() const { return body_inputs_; }
const std::vector<OutputTensor>& body_outputs() const {
return body_outputs_;
}
private:
// Each user-defined while loop defines a new execution frame, which is
// uniquely identified by its frame name. Frames are used by the executor to
// manage the iterations of a loop. See the FrameState comment in
// core/common_runtime/executor.cc for more details.
const std::string frame_name_;
// The enter nodes defining the input loop variables to the while loop. This
// vector defines the order of the loop variables.
const std::vector<Node*> enter_nodes_;
// The exit nodes defining the outputs of the while loop. These are in loop
// variable order.
const std::vector<Node*> exit_nodes_;
// The boolean output of the loop predicate.
const OutputTensor cond_output_;
// The inputs and outputs to the loop body.
const std::vector<OutputTensor> body_inputs_;
const std::vector<OutputTensor> body_outputs_;
};
} // namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPH_WHILE_CONTEXT_H_
tensorflow/core/graph/zen_graph_util.h
+83
View File
@@ -0,0 +1,83 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed 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 TENSORFLOW_CORE_GRAPH_ZEN_GRAPH_UTIL_H_
#define TENSORFLOW_CORE_GRAPH_ZEN_GRAPH_UTIL_H_
#ifdef AMD_ZENDNN
#include <string>
#include <utility>
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/platform/cpu_info.h"
#include "tensorflow/core/util/env_var.h"
namespace tensorflow {
namespace zen_op_registry {
// Prefix that we add to Tensorflow op name to construct Zen op name.
static const char* const kZenNodePrefix = "_Zen";
// Get the name of Zen op from original TensorFlow op.
// We prefix the original op with "Zen" to get Zen op.
inline string GetZenOpName(const string& name) {
return string(kZenNodePrefix) + name;
}
// Check whether op name with type T is registered as Zen operator
// that will go through name change or layout change pass.
//
// @input op_name - name of the op.
// @input T - datatype to be used for checking op.
// @return true if op name is registered as Zen op that will go through name
// change or layout change pass; false otherwise.
static inline bool IsZenOpKernelRegistered(const string& op_name, DataType T) {
string registered_kernels_key = op_name + string(DataType_Name(T));
thread_local static auto* registered_kernels_map =
new absl::flat_hash_map<string, bool>();
auto kernel_element = registered_kernels_map->find(registered_kernels_key);
bool kernel_registered = false;
if (kernel_element == registered_kernels_map->end()) {
string registered_kernels = KernelsRegisteredForOp(op_name);
// String returned by KernelsRegisteredForOp looks like below:
//
// Op = ZenMatMul, kernels =
// device='CPU'; T in [DT_FLOAT]
// device='CPU'; T in [DT_DOUBLE]
// If we have multiple kernels registered for the op. We need to verify
// our datatype
if (registered_kernels.find(string(DataType_Name(T))) != string::npos) {
kernel_registered = true;
}
registered_kernels_map->insert(
std::make_pair(registered_kernels_key, kernel_registered));
} else {
// Kernel is visited at least once. Return stored registration result.
kernel_registered = kernel_element->second;
}
return kernel_registered;
}
} // namespace zen_op_registry
} // namespace tensorflow
#endif // AMD_ZENDNN
#endif // TENSORFLOW_CORE_GRAPH_ZEN_GRAPH_UTIL_H_