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

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
+277
View File
@@ -0,0 +1,277 @@
load("//tensorflow:tensorflow.bzl", "tf_cc_test", "tf_cuda_library")
load("//tensorflow:tensorflow.default.bzl", "filegroup")
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
licenses = ["notice"],
)
filegroup(
name = "pywrap_required_hdrs",
srcs = [
"devices.h",
"grappler_item.h",
"grappler_item_builder.h",
],
visibility = [
"//tensorflow/python/grappler:__pkg__",
],
)
cc_library(
name = "op_types",
srcs = ["op_types.cc"],
hdrs = ["op_types.h"],
visibility = ["//visibility:public"],
deps = [
":utils",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/strings",
],
)
cc_library(
name = "utils",
srcs = ["utils.cc"],
hdrs = ["utils.h"],
visibility = ["//visibility:public"],
deps = [
"//tensorflow/core:framework",
"//tensorflow/core:graph",
"//tensorflow/core:lib",
"//tensorflow/core:lib_internal",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/container:node_hash_map",
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/synchronization",
"@com_google_absl//absl/time",
"@com_google_absl//absl/types:span",
"@eigen_archive//:eigen3",
],
)
tf_cc_test(
name = "utils_test",
srcs = ["utils_test.cc"],
deps = [
":grappler_item",
":utils",
"//tensorflow/cc:cc_ops",
"//tensorflow/core:all_kernels",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core:testlib",
"//tensorflow/core/framework:tensor_testutil",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/synchronization",
],
)
tf_cuda_library(
name = "devices",
srcs = ["devices.cc"],
hdrs = ["devices.h"],
cuda_deps = [
"@xla//xla/stream_executor/gpu:gpu_init",
"//tensorflow/core/platform:stream_executor",
],
visibility = ["//visibility:public"],
deps = [
"//tensorflow/core:lib",
"//tensorflow/core:lib_internal",
"@com_google_absl//absl/log",
],
)
cc_library(
name = "graph_topology_view",
srcs = ["graph_topology_view.cc"],
hdrs = ["graph_topology_view.h"],
visibility = ["//visibility:public"],
deps = [
":graph_view",
"//tensorflow/core:graph",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:inlined_vector",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/types:optional",
"@com_google_absl//absl/types:span",
],
)
tf_cc_test(
name = "graph_topology_view_test",
srcs = ["graph_topology_view_test.cc"],
deps = [
":graph_topology_view",
":graph_view",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
],
)
cc_library(
name = "graph_view",
srcs = ["graph_view.cc"],
hdrs = ["graph_view.h"],
visibility = ["//visibility:public"],
deps = [
":utils",
"//tensorflow/core:framework",
"//tensorflow/core:graph",
"//tensorflow/core:lib",
"//tensorflow/core:lib_internal",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/hash",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
],
)
tf_cc_test(
name = "graph_view_test",
srcs = ["graph_view_test.cc"],
deps = [
":graph_view",
":grappler_item",
"//tensorflow/cc:cc_ops",
"//tensorflow/core:lib",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core:testlib",
"//tensorflow/core/grappler/inputs:trivial_test_graph_input_yielder",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/strings",
],
)
cc_library(
name = "grappler_item",
srcs = [
"grappler_item.cc",
],
hdrs = ["grappler_item.h"],
visibility = ["//visibility:public"],
deps = [
":op_types",
":utils",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/grappler/utils:transitive_fanin",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/strings",
"@tsl//tsl/platform:platform_port",
],
)
cc_library(
name = "grappler_item_builder",
srcs = [
"grappler_item_builder.cc",
],
hdrs = ["grappler_item_builder.h"],
visibility = ["//visibility:public"],
deps = [
":grappler_item",
":op_types",
":utils",
"//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:protos_all_cc",
"//tensorflow/core/grappler/inputs:utils",
"//tensorflow/core/grappler/optimizers:model_pruner",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
],
)
tf_cc_test(
name = "grappler_item_test",
srcs = ["grappler_item_test.cc"],
deps = [
":grappler_item",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core:testlib",
"//tensorflow/core/grappler/inputs:trivial_test_graph_input_yielder",
],
)
tf_cc_test(
name = "grappler_item_builder_test",
srcs = ["grappler_item_builder_test.cc"],
deps = [
":grappler_item_builder",
"//tensorflow/cc:cc_ops",
"//tensorflow/cc:functional_ops",
"//tensorflow/cc:grad_testutil",
"//tensorflow/cc:gradients",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core:testlib",
"//tensorflow/core/grappler/inputs:trivial_test_graph_input_yielder",
],
)
cc_library(
name = "mutable_graph_view",
srcs = [
"mutable_graph_view.cc",
],
hdrs = ["mutable_graph_view.h"],
visibility = ["//visibility:public"],
deps = [
":graph_view",
":op_types",
":utils",
"//tensorflow/core:framework",
"//tensorflow/core:graph",
"//tensorflow/core:lib",
"//tensorflow/core:lib_internal",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/types:span",
],
)
tf_cc_test(
name = "mutable_graph_view_test",
srcs = ["mutable_graph_view_test.cc"],
deps = [
":grappler_item",
":mutable_graph_view",
":utils",
"//tensorflow/cc:cc_ops",
"//tensorflow/core:graph",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core:testlib",
"//tensorflow/core/grappler/inputs:trivial_test_graph_input_yielder",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/types:span",
],
)
+175
View File
@@ -0,0 +1,175 @@
load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda")
load("//tensorflow:tensorflow.bzl", "tf_cc_test", "tf_cuda_library")
load("//tensorflow:tensorflow.default.bzl", "filegroup")
load(
"//tensorflow/core/platform:build_config_root.bzl",
"tf_cuda_tests_tags",
)
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
licenses = ["notice"],
)
tf_cuda_library(
name = "utils",
srcs = ["utils.cc"],
hdrs = ["utils.h"],
cuda_deps = [
"@local_config_cuda//cuda:cudnn_header",
],
visibility = ["//visibility:public"],
deps = [
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/common_runtime/gpu:gpu_id",
"@eigen_archive//:eigen3",
],
)
tf_cc_test(
name = "utils_test",
srcs = ["utils_test.cc"],
linkstatic = if_cuda(1, 0),
tags = tf_cuda_tests_tags(),
deps = [
":utils",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/common_runtime/gpu:gpu_id",
],
)
filegroup(
name = "pywrap_required_hdrs",
srcs = [
"cluster.h",
"single_machine.h",
"utils.h",
"virtual_cluster.h",
],
visibility = [
"//tensorflow/python/grappler:__pkg__",
],
)
cc_library(
name = "cluster",
srcs = ["cluster.cc"],
hdrs = [
"cluster.h",
],
visibility = ["//visibility:public"],
deps = [
"//tensorflow/core:core_cpu_lib",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:session_options",
"//tensorflow/core/grappler:grappler_item",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
],
)
cc_library(
name = "virtual_cluster",
srcs = ["virtual_cluster.cc"],
hdrs = [
"virtual_cluster.h",
],
visibility = ["//visibility:public"],
deps = [
":cluster",
":utils",
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/grappler/costs:analytical_cost_estimator",
"//tensorflow/core/grappler/costs:op_level_cost_estimator",
"//tensorflow/core/grappler/costs:virtual_scheduler",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
],
)
tf_cc_test(
name = "virtual_cluster_test",
srcs = ["virtual_cluster_test.cc"],
deps = [
":virtual_cluster",
"//tensorflow/cc:cc_ops",
"//tensorflow/cc:scope",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/grappler:grappler_item",
"//tensorflow/core/grappler/inputs:trivial_test_graph_input_yielder",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/status",
"@xla//xla/tsl/protobuf:error_codes_proto_impl_cc",
],
)
cc_library(
name = "single_machine",
srcs = ["single_machine.cc"],
hdrs = [
"single_machine.h",
],
visibility = ["//visibility:public"],
deps = [
":cluster",
":utils",
"//tensorflow/cc:coordinator",
"//tensorflow/cc:queue_runner",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core/common_runtime:core_cpu",
"//tensorflow/core/common_runtime:core_cpu_lib",
"//tensorflow/core/common_runtime:direct_session_internal",
"//tensorflow/core/common_runtime/gpu:gpu_id",
"//tensorflow/core/framework:cost_graph_proto_cc",
"//tensorflow/core/framework:graph_proto_cc",
"//tensorflow/core/grappler:utils",
"//tensorflow/core/kernels:ops_util",
"//tensorflow/core/protobuf:for_core_protos_cc",
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/synchronization",
"@com_google_absl//absl/time",
"@com_google_absl//absl/types:optional",
],
alwayslink = 1,
)
tf_cc_test(
name = "single_machine_test",
srcs = ["single_machine_test.cc"],
args = ["--heap_check="], # The GPU tracer leaks memory. TODO(b/185483595): use a dependency instead of a flag
tags = [
"no_cuda_on_cpu_tap",
"no_gpu",
"nomsan", # TODO(b/160921160): broken by NOAUTOROLLBACK CL
],
deps = [
":single_machine",
"//tensorflow/cc:cc_ops",
"//tensorflow/cc:resource_variable_ops",
"//tensorflow/cc:scope",
"//tensorflow/core:core_cpu_internal",
"//tensorflow/core:lib_proto_parsing",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/grappler:grappler_item",
"//tensorflow/core/grappler:utils",
"//tensorflow/core/grappler/inputs:trivial_test_graph_input_yielder",
],
)
@@ -0,0 +1,115 @@
/* 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/grappler/clusters/cluster.h"
#include <algorithm>
#include <string>
#include <vector>
#include "tensorflow/core/protobuf/config.pb.h"
#include "tensorflow/core/protobuf/rewriter_config.pb.h"
namespace tensorflow {
namespace grappler {
Cluster::Cluster(int timeout_s) : timeout_s_(timeout_s) {
DisableDetailedStats(false);
}
Cluster::~Cluster() {}
void Cluster::AllowSoftPlacement(bool soft_placement_state) {
options_.config.set_allow_soft_placement(soft_placement_state);
}
void Cluster::SetNumInterOpThreads(int num_threads) {
for (int i = 0; i < options_.config.session_inter_op_thread_pool_size();
++i) {
options_.config.mutable_session_inter_op_thread_pool(i)->set_num_threads(
num_threads);
}
}
void Cluster::SetNumWarmupSteps(int num_steps) {
options_.config.mutable_graph_options()->set_build_cost_model_after(
num_steps);
}
// Set executor type to instantiate
void Cluster::SetExecutorType(const std::string* executor_type) {
options_.config.mutable_experimental()->set_executor_type(*executor_type);
}
int Cluster::NumWarmupSteps() const {
return options_.config.graph_options().build_cost_model_after();
}
void Cluster::DisableDetailedStats(bool disable) {
if (disable) {
options_.config.mutable_graph_options()->set_build_cost_model(0);
run_options_.set_trace_level(RunOptions::NO_TRACE);
} else {
options_.config.mutable_graph_options()->set_build_cost_model(1);
run_options_.set_trace_level(RunOptions::HARDWARE_TRACE);
}
}
bool Cluster::DetailedStatsEnabled() const {
return options_.config.graph_options().build_cost_model() != 0;
}
void Cluster::DisableOptimizer(bool disable) {
OptimizerOptions* options =
options_.config.mutable_graph_options()->mutable_optimizer_options();
if (disable) {
options->set_opt_level(OptimizerOptions::L0);
// Disable Grappler optimizations.
auto rewriter_config =
options_.config.mutable_graph_options()->mutable_rewrite_options();
rewriter_config->set_layout_optimizer(RewriterConfig::OFF);
rewriter_config->set_disable_model_pruning(true);
rewriter_config->set_function_optimization(RewriterConfig::OFF);
rewriter_config->set_arithmetic_optimization(RewriterConfig::OFF);
rewriter_config->set_loop_optimization(RewriterConfig::OFF);
rewriter_config->set_dependency_optimization(RewriterConfig::OFF);
rewriter_config->set_constant_folding(RewriterConfig::OFF);
rewriter_config->set_memory_optimization(RewriterConfig::NO_MEM_OPT);
rewriter_config->set_shape_optimization(RewriterConfig::OFF);
rewriter_config->set_remapping(RewriterConfig::OFF);
rewriter_config->set_pin_to_host_optimization(RewriterConfig::OFF);
rewriter_config->mutable_auto_parallel()->set_enable(false);
rewriter_config->clear_optimizers();
} else {
options->set_opt_level(OptimizerOptions::L1);
auto rewriter_config =
options_.config.mutable_graph_options()->mutable_rewrite_options();
rewriter_config->set_constant_folding(RewriterConfig::DEFAULT);
rewriter_config->set_memory_optimization(RewriterConfig::DEFAULT_MEM_OPT);
}
}
const std::vector<std::string> Cluster::GetDeviceNames() const {
std::vector<std::string> device_names;
device_names.reserve(devices_.size());
for (const auto& device : devices_) {
device_names.push_back(device.first);
}
std::sort(device_names.begin(), device_names.end());
return device_names;
}
} // end namespace grappler
} // end namespace tensorflow
+150
View File
@@ -0,0 +1,150 @@
/* 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_GRAPPLER_CLUSTERS_CLUSTER_H_
#define TENSORFLOW_CORE_GRAPPLER_CLUSTERS_CLUSTER_H_
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "tensorflow/core/common_runtime/device_set.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/grappler/grappler_item.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/strings/strcat.h"
#include "tensorflow/core/protobuf/config.pb.h"
#include "tensorflow/core/protobuf/device_properties.pb.h"
#include "tensorflow/core/public/session_options.h"
namespace tensorflow {
namespace grappler {
// A cluster represents of collection of hardware resources available to run
// the TensorFlow model.
// A process can only create a single cluster at a time.
class Cluster {
public:
explicit Cluster(int timeout_s);
virtual ~Cluster();
// Returns a string that represent the type of cluster that was instantiated.
virtual std::string type() const = 0;
// Provision the hardware resources needed to run TensorFlow and start a
// TensorFlow session that can take advantage of these resources.
// The actual resources that are leveraged depend on the type of cluster
// instantiated.
// Returns OK iff all the requested resources could be reserved and a
// TensorFlow session successfully created. Returns an error otherwise.
// There is no graceful degradation to handle the case where only a subset
// of the requested resources are available.
virtual absl::Status Provision() = 0;
// Attempts to shutdown the cluster.
// Returns OK iff there are no pending calls to the Run() method and all the
// resources used by the cluster could be released. Returns an error
// otherwise.
virtual absl::Status Shutdown() { return absl::OkStatus(); }
// Whether soft placement is allowed. If allow_soft_placement is true,
// an op will be placed on CPU if there's no GPU implementation for the OP
// or if no GPU devices are known or registered or if we need to co-locate
// with reftype input(s) which are from CPU.
void AllowSoftPlacement(bool soft_placement_state);
// Update the number of inter-op threads for each per-session threadpool
void SetNumInterOpThreads(int num_threads);
// Set the number of steps required to warmup TensorFlow. Must be called
// before Provision().
void SetNumWarmupSteps(int num_steps);
// Set executor type to instantiate
void SetExecutorType(const std::string* executor_type);
// Returns the number of warmup steps.
int NumWarmupSteps() const;
// Disable the collection of detailed statistics. Must be called
// before Provision().
void DisableDetailedStats(bool disable);
// Returns true iff the collection of detailed statistics is enabled.
bool DetailedStatsEnabled() const;
// Disable the TensorFlow optimizer. This ensures that the graph that TF
// executes is similar to the input graph. Must be called before Provision().
void DisableOptimizer(bool disable);
// Return the list of TensorFlow devices that are available to execute a
// graph. This is empty until provision() is called.
const std::unordered_map<std::string, DeviceProperties>& GetDevices() const {
return devices_;
}
// Convenience method that returns the set of device names. These names are
// sorted alphabetically.
const std::vector<std::string> GetDeviceNames() const;
// The DeviceSet is not always available, but when it is it contains a
// superset of the devices listed in GetDevices/GetDeviceNames().
virtual const DeviceSet* GetDeviceSet() const { return nullptr; }
// Enables collecting the allocator stats. If called, must be called before
// Provision().
virtual absl::Status EnablePeakMemoryStats() {
return absl::UnimplementedError(absl::StrCat(
"Peak Memory Stats are not supported on ", type(), " clusters"));
}
// Returns peak memory of all devices during the session creation and session
// runs.
virtual absl::Status GetPeakMemoryUsage(
std::unordered_map<std::string, uint64_t>* device_peak_memory) const {
return absl::UnimplementedError(
"GetPeakMemoryUsage is not implemented for this type of cluster.");
}
// Prepare the session to run the specified grappler item. This include
// initializing all the model variables.
virtual absl::Status Initialize(const GrapplerItem& item) = 0;
// Run the specified graph_def and return the corresponding metadata.
virtual absl::Status Run(
const GraphDef& graph_def,
const std::vector<std::pair<std::string, Tensor>>& feed,
const std::vector<std::string>& fetch, RunMetadata* metadata) = 0;
// Run the specified GrapplerItem and return the corresponding metadata.
virtual absl::Status Run(const GrapplerItem& item, RunMetadata* metadata) {
return Run(item.graph, item.feed, item.fetch, metadata);
}
protected:
std::unordered_map<std::string, DeviceProperties> devices_;
const int timeout_s_;
SessionOptions options_;
RunOptions run_options_;
};
} // end namespace grappler
} // end namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPPLER_CLUSTERS_CLUSTER_H_
@@ -0,0 +1,492 @@
/* 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/grappler/clusters/single_machine.h"
#include <atomic>
#include <cstddef>
#include <cstdint>
#include <memory>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>
#include "absl/log/check.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "absl/synchronization/notification.h"
#include "absl/time/time.h"
#include "absl/types/optional.h"
#include "tensorflow/cc/training/queue_runner.h"
#include "tensorflow/core/common_runtime/device.h"
#include "tensorflow/core/common_runtime/device_mgr.h"
#include "tensorflow/core/common_runtime/gpu/gpu_id.h"
#include "tensorflow/core/common_runtime/gpu/gpu_id_manager.h"
#include "tensorflow/core/framework/cost_graph.pb.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/grappler/clusters/utils.h"
#include "tensorflow/core/grappler/utils.h"
#include "tensorflow/core/kernels/ops_util.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/strings/strcat.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/notification.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/protobuf/config.pb.h"
#include "tensorflow/core/protobuf/device_properties.pb.h"
#include "tensorflow/core/public/session.h"
namespace tensorflow {
namespace grappler {
static std::atomic<bool> already_provisioned(false);
SingleMachine::SingleMachine(int timeout_s, int num_cpu_cores, int num_gpus)
: Cluster(timeout_s), expected_init_time_s_(0), closing_(false) {
VLOG(1) << "Number of CPU cores: " << num_cpu_cores
<< " Number of GPUs: " << num_gpus;
thread_pool_ = std::make_unique<thread::ThreadPool>(
Env::Default(), SanitizeThreadSuffix("single_machine"), 2);
(*options_.config.mutable_device_count())["CPU"] = 1;
if (num_gpus > 0) {
(*options_.config.mutable_device_count())["GPU"] = num_gpus;
}
CHECK_GE(num_cpu_cores, 1);
options_.config.set_intra_op_parallelism_threads(num_cpu_cores);
// Create a session specific thread pool to ensure the threads are reset when
// the session is reset.
options_.config.add_session_inter_op_thread_pool()->set_num_threads(
num_cpu_cores);
if (timeout_s > 0) {
options_.config.set_operation_timeout_in_ms(timeout_s * 1000);
}
}
SingleMachine::~SingleMachine() {
CloseSession(false /*use_timeout*/).IgnoreError();
// Reset the thread-pool so that there are no outstanding Session::Run(...)s
// when we delete the session.
thread_pool_.reset();
}
absl::Status SingleMachine::Provision() {
// This is really ugly: to avoid leaking variables, we need to reset the tf
// session every time we're done processing a grappler item. However,
// variables are global, and therefore we can't have more than 1 session alive
// at a time. This check detects when more that one cluster is provisioned.
if (already_provisioned) {
return absl::UnavailableError(
"Can't provision more than one single cluster at a time");
}
TF_RETURN_IF_ERROR(ResetSession());
std::vector<DeviceAttributes> devices;
TF_RETURN_IF_ERROR(session_->ListDevices(&devices));
for (const auto& dev : devices) {
DeviceProperties attr;
if (dev.device_type() == "CPU") {
attr = GetLocalCPUInfo();
} else if (dev.device_type() == "GPU") {
DeviceNameUtils::ParsedName parsed;
if (!DeviceNameUtils::ParseFullName(dev.name(), &parsed)) {
return absl::InvalidArgumentError(
absl::StrCat("Not able to parse GPU device name: ", dev.name()));
}
TfDeviceId tf_device_id(parsed.id);
PlatformDeviceId platform_device_id;
absl::Status s =
GpuIdManager::TfToPlatformDeviceId(tf_device_id, &platform_device_id);
if (!s.ok()) {
return absl::UnavailableError(
absl::StrCat("Unknown TF GPU device with id ", tf_device_id.value(),
": ", s.message()));
}
attr = GetLocalGPUInfo(platform_device_id);
} else if (dev.device_type().find("XLA") == std::string::npos) {
// Filter out the fake XLA devices to avoid double counting the actual
// hardware resources that are available.
attr.set_type(dev.device_type());
}
// Overwrite the memory size since users might have requested to use only a
// fraction of the available device memory.
attr.set_memory_size(dev.memory_limit());
devices_[dev.name()] = attr;
}
already_provisioned = true;
// Clear highmark stats of all local allocators.
if (cpu_allocator_stats_enabled_) {
TF_RETURN_IF_ERROR(ClearAllocatorStats());
}
return absl::OkStatus();
}
absl::Status SingleMachine::Initialize(const GrapplerItem& item) {
mutex_lock l(this->last_graph_mu_);
if (last_graph_ != &item.graph || last_graph_id_ != item.id) {
init_ops_ = item.init_ops;
expected_init_time_s_ = item.expected_init_time;
last_graph_ = nullptr;
queue_runner_defs_ = item.queue_runners;
last_graph_id_ = item.id;
}
return absl::OkStatus();
}
absl::Status SingleMachine::Shutdown() {
TF_RETURN_IF_ERROR(ShutdownSession());
mutex_lock l(this->last_graph_mu_);
last_graph_ = nullptr;
already_provisioned = false;
return absl::OkStatus();
}
absl::Status SingleMachine::Run(
const GraphDef& graph_def,
const std::vector<std::pair<std::string, Tensor>>& feed,
const std::vector<std::string>& fetch, RunMetadata* metadata) {
mutex_lock l(this->last_graph_mu_);
if (last_graph_ != &graph_def) {
TF_RETURN_IF_ERROR(ResetSession());
TF_RETURN_IF_ERROR(session_->Create(graph_def));
if (!init_ops_.empty()) {
init_metadata_ = RunMetadata();
int64_t timeout_s = timeout_s_ + expected_init_time_s_;
TF_RETURN_IF_ERROR(
RunWithTimeout({}, init_ops_, &init_metadata_, timeout_s));
// The compute cost for init ops is likely to be pessimistic since init
// ops are run only once before warmup. Therefore we only keep their
// memory costs.
for (auto node : *init_metadata_.mutable_cost_graph()->mutable_node()) {
node.clear_compute_cost();
}
// Also clear the timeline to save memory
init_metadata_.clear_step_stats();
}
// We can have at most one hardware trace. Use it for the main graph, and
// downgrade tracing of the queue runners to a software trace.
RunOptions queue_options = run_options_;
if (queue_options.trace_level() >= RunOptions::HARDWARE_TRACE) {
queue_options.set_trace_level(RunOptions::SOFTWARE_TRACE);
}
for (size_t i = 0; i < queue_runner_defs_.size(); ++i) {
std::unique_ptr<QueueRunner> queue_runner;
TF_RETURN_IF_ERROR(QueueRunner::New(queue_runner_defs_[i],
coordinator_.get(), &queue_runner));
TF_RETURN_IF_ERROR(queue_runner->StartAndCollectCostGraph(session_.get(),
queue_options));
TF_RETURN_IF_ERROR(coordinator_->RegisterRunner(std::move(queue_runner)));
TF_RETURN_IF_ERROR(coordinator_->GetStatus());
}
// Warmup TensorFlow if needed
for (int i = 0; i < NumWarmupSteps(); ++i) {
TF_RETURN_IF_ERROR(RunWithTimeout(feed, fetch, nullptr));
}
}
if (metadata) {
TF_RETURN_IF_ERROR(RunWithTimeout(feed, fetch, metadata));
// Merge the costs of the initialization and the queue runners.
CostGraphDef queue_costs;
TF_RETURN_IF_ERROR(coordinator_->ExportCostGraph(&queue_costs));
MergeCosts(metadata->mutable_cost_graph(), init_metadata_.cost_graph(),
queue_costs);
} else {
TF_RETURN_IF_ERROR(RunWithTimeout(feed, fetch, nullptr));
}
last_graph_ = &graph_def;
return absl::OkStatus();
}
absl::Status SingleMachine::EnablePeakMemoryStats() {
EnableCPUAllocatorStats();
cpu_allocator_stats_enabled_ = true;
// No need to enable GPU allocator stats since its stats are always collected.
return absl::OkStatus();
}
absl::Status SingleMachine::GetPeakMemoryUsage(
std::unordered_map<std::string, uint64_t>* device_peak_memory) const {
// Cpu_allocator->TracksAllocationSizes() returns true doesn't always mean the
// the AllocatorStats would be collected.
if (!cpu_allocator_stats_enabled_) {
return absl::Status(absl::StatusCode::kInvalidArgument,
"Tracking allocation for CPU is not enabled.");
}
const DeviceMgr* device_mgr;
TF_RETURN_IF_ERROR(session_->LocalDeviceManager(&device_mgr));
std::vector<Device*> devices = device_mgr->ListDevices();
device_peak_memory->clear();
for (Device* device : devices) {
auto* allocator = device->GetAllocator(AllocatorAttributes());
if (!allocator->TracksAllocationSizes()) {
return absl::Status(absl::StatusCode::kInvalidArgument,
"Tracking allocation is not enabled.");
}
absl::optional<AllocatorStats> stats = allocator->GetStats();
(*device_peak_memory)[device->name()] =
(stats ? stats->peak_bytes_in_use : 0);
}
return absl::OkStatus();
}
absl::Status SingleMachine::RunWithTimeout(
const std::vector<std::pair<std::string, Tensor>>& feed,
const std::vector<std::string>& fetch, RunMetadata* run_metadata) {
return RunWithTimeout(feed, fetch, run_metadata, timeout_s_);
}
absl::Status SingleMachine::RunWithTimeout(
const std::vector<std::pair<std::string, Tensor>>& feed,
const std::vector<std::string>& fetch, RunMetadata* run_metadata,
int64_t timeout_s) {
// We shouldn't be running or closing the session at this point.
{
mutex_lock l(close_mu_);
CHECK(!closing_);
}
auto status = std::make_shared<absl::Status>();
auto local_metadata = std::make_shared<RunMetadata>();
const bool executed_in_time = ExecuteWithTimeout(
[this, status, local_metadata, feed, fetch]() {
*status = session_->Run(run_options_, feed, {}, fetch, nullptr,
local_metadata.get());
},
timeout_s * 1000, thread_pool_.get());
if (!executed_in_time) {
return absl::DeadlineExceededError(absl::StrCat(
"Failed to run the graph after ", timeout_s, " seconds, aborting"));
} else if (run_metadata && status->ok()) {
*run_metadata = *local_metadata;
}
return *status;
}
absl::Status SingleMachine::CloseSession(bool use_timeout) {
if (!session_ || !thread_pool_) {
return absl::OkStatus();
}
{
mutex_lock l(close_mu_);
if (!closing_) {
closing_ = true;
}
}
const bool executed_in_time = ExecuteWithTimeout(
[&]() {
if (this->coordinator_) {
this->coordinator_->RequestStop().IgnoreError();
// Wait for all the runners to have closed their queues.
while (!this->coordinator_->AllRunnersStopped()) {
Env::Default()->SleepForMicroseconds(1000000);
}
// Now we can close the session. This should cancel any pending I/O
// operation.
this->session_->Close().IgnoreError();
// Last but not least, we can delete the coordinator.
this->coordinator_.reset();
} else {
this->session_->Close().IgnoreError();
}
mutex_lock l2(close_mu_);
closing_ = false;
},
use_timeout ? timeout_s_ * 1000 : -1, thread_pool_.get());
if (!executed_in_time) {
// Let the caller know that we can't shutdown the session, and therefore
// can't process any further.
return absl::UnavailableError(
absl::StrCat("Failed to close the previous session after ", timeout_s_,
" seconds, aborting"));
}
return absl::OkStatus();
}
absl::Status SingleMachine::ShutdownSession() {
TF_RETURN_IF_ERROR(CloseSession(true /*use_timeout*/));
// Delete the threadpool: this ensures that all the pending closures complete
// before we return. Note that if TF deadlocked on us, the closures will
// never complete, and the call to thread_pool_.reset() will never return:
// therefore we need to delete the threadpool with the background thread.
// That thread itself will also never complete, so the user should
// abort the process to avoid leaking too many resources.
auto n = std::make_shared<absl::Notification>();
Env::Default()->SchedClosure([this, n]() {
thread_pool_.reset();
n->Notify();
});
const bool notified =
n->WaitForNotificationWithTimeout(absl::Seconds(timeout_s_));
if (!notified) {
// Let the caller know that we can't shutdown the session properly since
// there are calls to Session::Run() still running.
return absl::UnavailableError(absl::StrCat(
"The session is still running graphs after ", timeout_s_, " seconds"));
}
return absl::OkStatus();
}
absl::Status SingleMachine::ResetSession() {
if (session_) {
LOG(INFO) << "Cleaning up previous session";
// Make sure the session is properly closed
TF_RETURN_IF_ERROR(ShutdownSession());
// Destroying the object deletes all its variables as well. This is only
// true for DirectSession.
session_.reset();
}
LOG(INFO) << "Starting new session";
// Create a new threadpool
thread_pool_ = std::make_unique<thread::ThreadPool>(
Env::Default(), SanitizeThreadSuffix("single_machine"), 2);
session_.reset(NewSession(options_));
if (!session_) {
return absl::UnknownError("Failed to create session");
}
coordinator_ = std::make_unique<Coordinator>();
// Build the DeviceSet.
device_set_ = std::make_unique<DeviceSet>();
const DeviceMgr* device_mgr;
TF_RETURN_IF_ERROR(session_->LocalDeviceManager(&device_mgr));
for (auto d : device_mgr->ListDevices()) {
device_set_->AddDevice(d);
// We currently don't care about the client device.
}
return absl::OkStatus();
}
void SingleMachine::MergeCosts(CostGraphDef* graph_costs,
const CostGraphDef& init_costs,
const CostGraphDef& queue_costs) {
graph_costs->mutable_node()->Reserve(graph_costs->node_size() +
init_costs.node_size() +
queue_costs.node_size());
std::unordered_set<std::string> nodes_seen;
int queue_costs_id_offset = graph_costs->node_size();
for (const auto& node : graph_costs->node()) {
nodes_seen.insert(node.name());
if (node.id() >= queue_costs_id_offset) {
queue_costs_id_offset = node.id() + 1;
}
}
int init_costs_id_offset = queue_costs_id_offset + queue_costs.node_size();
// The costs obtained by running the main graph could be more stable than
// the one we get from the queue runners since the queue runners run
// asynchronously.
for (const auto& node : queue_costs.node()) {
if (nodes_seen.find(node.name()) != nodes_seen.end()) {
continue;
}
auto* new_node = graph_costs->add_node();
new_node->MergeFrom(node);
new_node->set_id(node.id() + queue_costs_id_offset);
if (new_node->id() >= init_costs_id_offset) {
init_costs_id_offset = new_node->id() + 1;
}
for (auto& input_info : *new_node->mutable_input_info()) {
input_info.set_preceding_node(input_info.preceding_node() +
queue_costs_id_offset);
}
for (auto& control_input : *new_node->mutable_control_input()) {
control_input += queue_costs_id_offset;
}
}
// Don't overwrite the costs with that generated during initialization since
// these are possibly outdated.
for (const auto& node : init_costs.node()) {
if (nodes_seen.find(node.name()) != nodes_seen.end()) {
continue;
}
auto* new_node = graph_costs->add_node();
new_node->MergeFrom(node);
new_node->set_id(node.id() + init_costs_id_offset);
for (auto& input_info : *new_node->mutable_input_info()) {
input_info.set_preceding_node(input_info.preceding_node() +
init_costs_id_offset);
}
for (auto& control_input : *new_node->mutable_control_input()) {
control_input += init_costs_id_offset;
}
}
}
absl::Status SingleMachine::ClearAllocatorStats() const {
// Cpu_allocator->TracksAllocationSizes() returns true doesn't always mean the
// the AllocatorStats would be collected.
if (!cpu_allocator_stats_enabled_) {
return absl::Status(absl::StatusCode::kInvalidArgument,
"Tracking allocation for CPU is not enabled.");
}
const DeviceMgr* device_mgr;
TF_RETURN_IF_ERROR(session_->LocalDeviceManager(&device_mgr));
std::vector<Device*> devices = device_mgr->ListDevices();
for (Device* device : devices) {
auto* allocator = device->GetAllocator(AllocatorAttributes());
if (!allocator->TracksAllocationSizes()) {
return absl::Status(absl::StatusCode::kInvalidArgument,
"Tracking allocation is not enabled.");
}
if (!allocator->ClearStats()) {
return absl::Status(
absl::StatusCode::kInvalidArgument,
absl::StrCat("Clearing allocation stats is not supported for ",
device->name()));
}
}
return absl::OkStatus();
}
} // namespace grappler
} // namespace tensorflow
@@ -0,0 +1,104 @@
/* 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_GRAPPLER_CLUSTERS_SINGLE_MACHINE_H_
#define TENSORFLOW_CORE_GRAPPLER_CLUSTERS_SINGLE_MACHINE_H_
#include <cstdint>
#include <memory>
#include <unordered_map>
#include <utility>
#include <vector>
#include "absl/status/status.h"
#include "tensorflow/cc/training/coordinator.h"
#include "tensorflow/core/framework/allocator.h"
#include "tensorflow/core/framework/cost_graph.pb.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/grappler/clusters/cluster.h"
#include "tensorflow/core/lib/core/threadpool.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/protobuf/config.pb.h"
#include "tensorflow/core/protobuf/queue_runner.pb.h"
#include "tensorflow/core/public/session.h"
namespace tensorflow {
namespace grappler {
// Create a simple cluster that makes available to grappler a subset of the
// nodes available on a single local computer.
class SingleMachine : public Cluster {
public:
SingleMachine(int timeout_s, int num_cpu_cores, int num_gpus);
~SingleMachine() override;
std::string type() const override { return "single_machine"; }
absl::Status Provision() override;
absl::Status Shutdown() override;
absl::Status Initialize(const GrapplerItem& item) override;
absl::Status Run(const GraphDef& item,
const std::vector<std::pair<std::string, Tensor>>& feed,
const std::vector<std::string>& fetch,
RunMetadata* metadata) override;
const DeviceSet* GetDeviceSet() const override { return device_set_.get(); }
absl::Status EnablePeakMemoryStats() override;
// It requires EnableAllocatorStats(true) be called before Provision().
absl::Status GetPeakMemoryUsage(std::unordered_map<std::string, uint64_t>*
device_peak_memory) const override;
private:
absl::Status RunWithTimeout(
const std::vector<std::pair<std::string, Tensor>>& feed,
const std::vector<std::string>& fetch, RunMetadata* run_metadata);
absl::Status RunWithTimeout(
const std::vector<std::pair<std::string, Tensor>>& feed,
const std::vector<std::string>& fetch, RunMetadata* run_metadata,
int64_t timeout_s);
absl::Status ResetSession();
absl::Status CloseSession(bool use_timeout);
absl::Status ShutdownSession();
void MergeCosts(CostGraphDef* graph_costs, const CostGraphDef& init_costs,
const CostGraphDef& queue_costs);
absl::Status ClearAllocatorStats() const;
std::unique_ptr<Session> session_;
std::vector<QueueRunnerDef> queue_runner_defs_;
std::string last_graph_id_;
mutex last_graph_mu_;
const GraphDef* last_graph_ TF_GUARDED_BY(last_graph_mu_) = nullptr;
std::vector<std::string> init_ops_;
int64_t expected_init_time_s_;
std::unique_ptr<Coordinator> coordinator_;
std::unique_ptr<thread::ThreadPool> thread_pool_;
std::unique_ptr<DeviceSet> device_set_;
RunMetadata init_metadata_;
mutex close_mu_;
bool closing_ TF_GUARDED_BY(close_mu_);
bool cpu_allocator_stats_enabled_ = false;
};
} // end namespace grappler
} // end namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPPLER_CLUSTERS_SINGLE_MACHINE_H_
@@ -0,0 +1,643 @@
/* 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/grappler/clusters/single_machine.h"
#include <memory>
#include "tensorflow/cc/framework/scope.h"
#include "tensorflow/cc/ops/resource_variable_ops.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/core/common_runtime/device.h"
#include "tensorflow/core/common_runtime/device_factory.h"
#include "tensorflow/core/framework/cost_graph.pb.h"
#include "tensorflow/core/framework/tensor_shape.pb.h"
#include "tensorflow/core/grappler/grappler_item.h"
#include "tensorflow/core/grappler/inputs/trivial_test_graph_input_yielder.h"
#include "tensorflow/core/grappler/utils.h"
#include "tensorflow/core/platform/protobuf.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/protobuf/queue_runner.pb.h"
namespace tensorflow {
namespace grappler {
namespace {
class SingleMachineTest : public ::testing::Test {
public:
void SetUp() override {
// Provision a single machine with 3 cpu cores, and a short timeout of 5
// seconds: since there isn't much work to process a test graph that should
// be plenty.
#if TENSORFLOW_USE_ROCM
// ROCm takes longer to start up
int timeout_s = 10;
#else
int timeout_s = 5;
#endif
#ifdef THREAD_SANITIZER
timeout_s *= 5;
#endif
cluster_ = std::make_unique<SingleMachine>(timeout_s, 3 /* num_cpu_cores */,
0 /* num_gpus */);
#if !defined(__APPLE__)
TF_CHECK_OK(cluster_->EnablePeakMemoryStats());
#endif
TF_CHECK_OK(cluster_->Provision());
}
void TearDown() override {
if (cluster_) {
TF_CHECK_OK(cluster_->Shutdown());
}
cluster_.reset();
}
protected:
std::unique_ptr<SingleMachine> cluster_;
};
TEST_F(SingleMachineTest, ClusterType) {
CHECK_EQ("single_machine", cluster_->type());
}
TEST_F(SingleMachineTest, CostModel) {
TrivialTestGraphInputYielder fake_input(4, 1, 10, false,
cluster_->GetDeviceNames());
GrapplerItem item;
CHECK(fake_input.NextItem(&item));
TF_CHECK_OK(cluster_->Initialize(item));
RunMetadata metadata;
const int64_t start_micros = Env::Default()->NowMicros();
TF_CHECK_OK(cluster_->Run(item.graph, item.feed, item.fetch, &metadata));
const int64_t run_duration_micros =
Env::Default()->NowMicros() - start_micros;
// There should be at least 4 nodes corresponding to the 4 stages we created
// in the fake input.
EXPECT_LE(4, metadata.cost_graph().node_size());
for (const auto& node : metadata.cost_graph().node()) {
// Skip the special nodes inserted by TF: these are prefixed with an
// underscore.
if (node.name()[0] == '_' || node.name().find("/_") != std::string::npos) {
continue;
}
#ifndef INTEL_MKL
// The output size of MKL op is 2, and cannot filter out the MKL op
// with the OP name (no op name here), so just disable this check in
// TF_MKL build.
EXPECT_EQ(1, node.output_info_size());
#endif // !INTEL_MKL
EXPECT_LE(8, node.output_info(0).size());
const TensorShapeProto& shape = node.output_info(0).shape();
EXPECT_EQ(2, shape.dim_size());
EXPECT_EQ(10, shape.dim(0).size());
EXPECT_EQ(1, shape.dim(1).size());
EXPECT_LE(0, node.compute_cost());
EXPECT_GE(run_duration_micros, node.compute_cost());
}
}
TEST_F(SingleMachineTest, Queue) {
TrivialTestGraphInputYielder fake_input(4, 1, 10, true,
cluster_->GetDeviceNames());
GrapplerItem item;
CHECK(fake_input.NextItem(&item));
TF_CHECK_OK(cluster_->Initialize(item));
RunMetadata metadata;
TF_CHECK_OK(cluster_->Run(item.graph, item.feed, item.fetch, &metadata));
}
TEST_F(SingleMachineTest, MultipleItems) {
TrivialTestGraphInputYielder fake_input(4, 1, 10, false,
cluster_->GetDeviceNames());
for (int i = 0; i < 3; ++i) {
GrapplerItem item;
CHECK(fake_input.NextItem(&item));
TF_CHECK_OK(cluster_->Initialize(item));
RunMetadata metadata1;
TF_CHECK_OK(cluster_->Run(item.graph, item.feed, item.fetch, &metadata1));
RunMetadata metadata2;
TF_CHECK_OK(cluster_->Run(item.graph, item.feed, item.fetch, &metadata2));
// There should be at least 4 nodes corresponding to the 4 stages we created
// in the fake input, plus 1 enqueue and 1 dequeue node.
EXPECT_LE(6, metadata1.cost_graph().node_size());
for (const auto& node : metadata1.cost_graph().node()) {
if (node.name()[0] == '_' ||
node.name().find("/_") != std::string::npos ||
node.name() == "queue") {
continue;
}
#ifndef INTEL_MKL
EXPECT_EQ(1, node.output_info_size());
#endif // !INTEL_MKL
const TensorShapeProto& shape = node.output_info(0).shape();
EXPECT_EQ(2, shape.dim_size());
EXPECT_EQ(10, shape.dim(0).size());
EXPECT_EQ(1, shape.dim(1).size());
}
for (int i = 0; i < metadata1.cost_graph().node_size(); ++i) {
metadata1.mutable_cost_graph()->mutable_node(i)->set_compute_cost(0);
metadata1.clear_step_stats();
}
for (int i = 0; i < metadata2.cost_graph().node_size(); ++i) {
metadata2.mutable_cost_graph()->mutable_node(i)->set_compute_cost(0);
metadata2.clear_step_stats();
}
std::string s1;
::tensorflow::protobuf::TextFormat::PrintToString(metadata1, &s1);
std::string s2;
::tensorflow::protobuf::TextFormat::PrintToString(metadata2, &s2);
EXPECT_EQ(s1, s2);
}
}
TEST_F(SingleMachineTest, GraphOptimizations) {
// Create a graph that can be fully precomputed
tensorflow::Scope root = tensorflow::Scope::NewRootScope();
auto zero = ops::Const(root.WithOpName("zero"), 0.0f, {2, 3});
auto one = ops::Const(root.WithOpName("one"), 1.0f, {2, 3});
auto add = ops::Add(root.WithOpName("add"), zero, one);
auto square = ops::Square(root.WithOpName("square"), add);
auto new_shape = ops::Const(root.WithOpName("new_shape"), {3, -1}, {2});
auto reshaped = ops::Reshape(root.WithOpName("reshaped"), square, new_shape);
auto final_shape = ops::Shape(root.WithOpName("final_shape"), reshaped);
auto expected_shape =
ops::Const(root.WithOpName("expected_shape"), {3, 2}, {2});
auto valid =
ops::Equal(root.WithOpName("valid"), final_shape, expected_shape);
auto all_dims = ops::Const(root.WithOpName("all_dims"), {0}, {1});
auto all_valid = ops::All(root.WithOpName("all_valid"), valid, all_dims);
auto assert_valid = ops::Assert(root.WithOpName("assert_valid"), all_valid,
{final_shape.output});
GrapplerItem item;
TF_CHECK_OK(root.ToGraphDef(&item.graph));
item.fetch.push_back("assert_valid");
// Force the placement of all the nodes on CPU since TF attempts to use a GPU
// when possible event though we created the session to have a single CPU !.
for (auto& node : *item.graph.mutable_node()) {
node.set_device("/cpu:0");
}
// With optimizations turned on, some nodes could have been optimized away,
// and the cost model could be partial. Restart the cluster with optimizations
// disabled and make sure we have all the information we're looking for.
TF_CHECK_OK(cluster_->Shutdown());
cluster_->DisableOptimizer(true);
TF_CHECK_OK(cluster_->Provision());
RunMetadata metadata;
TF_CHECK_OK(cluster_->Initialize(item));
TF_CHECK_OK(cluster_->Run(item.graph, item.feed, item.fetch, &metadata));
std::set<std::string> cost_nodes;
for (const auto& node : metadata.cost_graph().node()) {
#ifdef INTEL_MKL
// Skip the special nodes inserted by TF (and MKL): these are either
// prefixed with an underscore or contain "/_".
if (node.name()[0] == '_' || node.name().find("/_") != string::npos) {
continue;
}
cost_nodes.insert(node.name());
#else
// Skip nodes added by TF internally.
if (node.name()[0] != '_') {
cost_nodes.insert(node.name());
}
#endif
}
const std::set<std::string> expected_cost_nodes = {
"zero", "one", "add", "square",
"new_shape", "reshaped", "final_shape", "expected_shape",
"valid", "all_dims", "all_valid", "assert_valid"};
EXPECT_EQ(expected_cost_nodes, cost_nodes);
}
TEST_F(SingleMachineTest, TimeOuts) {
// Create a graph that will block forever: Just try to dequeue data from a
// queue that is never fed.
tensorflow::Scope root = tensorflow::Scope::NewRootScope();
auto q = ops::FIFOQueue(root.WithOpName("queue"), {DataType::DT_INT32});
auto dequeue =
ops::QueueDequeue(root.WithOpName("dequeue"), q, {DataType::DT_INT32});
GrapplerItem item;
TF_CHECK_OK(root.ToGraphDef(&item.graph));
item.fetch.push_back("dequeue");
TF_CHECK_OK(cluster_->Initialize(item));
RunMetadata metadata;
absl::Status s1 = cluster_->Run(item.graph, item.feed, item.fetch, &metadata);
EXPECT_TRUE(absl::IsDeadlineExceeded(s1));
absl::Status s2 = cluster_->Run(item.graph, item.feed, item.fetch, &metadata);
EXPECT_TRUE(absl::IsDeadlineExceeded(s2));
}
static void RunInfiniteTFLoop() {
// Create a while(true) loop
GrapplerItem item;
NodeDef* shp = item.graph.add_node();
shp->set_name("shape");
shp->set_op("Const");
(*shp->mutable_attr())["dtype"].set_type(DT_INT32);
Tensor shp_tensor(DT_INT32, TensorShape({1}));
shp_tensor.flat<int32_t>()(0) = 1;
shp_tensor.AsProtoTensorContent(
(*shp->mutable_attr())["value"].mutable_tensor());
NodeDef* r = item.graph.add_node();
r->set_name("random");
r->set_op("RandomUniform");
(*r->mutable_attr())["dtype"].set_type(DT_FLOAT);
(*r->mutable_attr())["T"].set_type(DT_INT32);
*r->add_input() = "shape";
NodeDef* e = item.graph.add_node();
e->set_name("while/Enter");
e->set_op("Enter");
(*e->mutable_attr())["T"].set_type(DT_FLOAT);
(*e->mutable_attr())["frame_name"].set_s("while/while/");
*e->add_input() = "random";
NodeDef* m = item.graph.add_node();
m->set_name("while/Merge");
m->set_op("Merge");
(*m->mutable_attr())["T"].set_type(DT_FLOAT);
(*m->mutable_attr())["N"].set_i(2);
*m->add_input() = "while/Enter";
*m->add_input() = "while/NextIteration";
NodeDef* t = item.graph.add_node();
t->set_name("always_true");
t->set_op("Const");
(*t->mutable_attr())["dtype"].set_type(DT_BOOL);
*t->add_input() = "^while/Merge";
Tensor true_tensor(DT_BOOL, TensorShape());
true_tensor.flat<bool>()(0) = true;
true_tensor.AsProtoTensorContent(
(*t->mutable_attr())["value"].mutable_tensor());
NodeDef* c = item.graph.add_node();
c->set_name("while/LoopCond");
c->set_op("LoopCond");
*c->add_input() = "always_true";
NodeDef* s = item.graph.add_node();
s->set_name("while/Switch");
(*s->mutable_attr())["T"].set_type(DT_FLOAT);
s->set_op("Switch");
*s->add_input() = "while/Merge";
*s->add_input() = "while/LoopCond";
NodeDef* i = item.graph.add_node();
i->set_name("while/Identity");
i->set_op("Identity");
(*i->mutable_attr())["T"].set_type(DT_FLOAT);
*i->add_input() = "while/Switch:1";
NodeDef* n = item.graph.add_node();
n->set_name("while/NextIteration");
n->set_op("NextIteration");
(*n->mutable_attr())["T"].set_type(DT_FLOAT);
*n->add_input() = "while/Identity";
NodeDef* x = item.graph.add_node();
x->set_name("while/Exit");
x->set_op("Exit");
(*x->mutable_attr())["T"].set_type(DT_FLOAT);
*x->add_input() = "while/Switch";
item.fetch.push_back("while/Exit");
// Create our own cluster to run it
SingleMachine cluster(5, 3, 0);
TF_CHECK_OK(cluster.Provision());
TF_CHECK_OK(cluster.Initialize(item));
absl::Status s1 = cluster.Run(item.graph, item.feed, item.fetch, nullptr);
if (!absl::IsDeadlineExceeded(s1)) {
LOG(ERROR) << "Expected 'deadline exceeded' error, got " << s1;
// Exit to break the infinite loop
_exit(1);
}
// Attempt to shutdown the cluster and make sure we get the proper error code.
absl::Status s2 = cluster.Shutdown();
if (!absl::IsUnavailable(s2)) {
LOG(ERROR) << "Expected 'unavailable' error, got " << s2;
// Exit to break the infinite loop
_exit(2);
}
// The isn't much we can do at this point. Exit with the error code 0 to
// indicate everything went according to plan.
_exit(0);
}
TEST_F(SingleMachineTest, InfiniteLoops) {
#if !(TENSORFLOW_USE_ROCM) // fails with ROCm (investigate)
// The RunInfiniteTFLoop function creates its own cluster.
TF_CHECK_OK(cluster_->Shutdown());
EXPECT_EXIT(RunInfiniteTFLoop(), ::testing::ExitedWithCode(0), ".*");
#endif
}
TEST_F(SingleMachineTest, InitializationMemory) {
// Build a variable and its initialization graph.
tensorflow::Scope s = tensorflow::Scope::NewRootScope();
int batch_size = 10;
Output x =
ops::RandomNormal(s.WithOpName("x"), {batch_size, 1}, DataType::DT_FLOAT);
Output v = ops::Variable(s.WithOpName("v"), TensorShape({batch_size, 1}),
DataType::DT_FLOAT);
Output init = ops::Assign(s.WithOpName("init"), v, x);
GrapplerItem item;
TF_CHECK_OK(s.ToGraphDef(&item.graph));
item.init_ops.push_back(init.name());
item.fetch.push_back(v.name());
TF_CHECK_OK(cluster_->Initialize(item));
RunMetadata metadata;
TF_CHECK_OK(cluster_->Run(item.graph, item.feed, item.fetch, &metadata));
// Check that the initialization op is present in the cost model.
bool found = false;
for (const auto& node : metadata.cost_graph().node()) {
found |= (node.name() == NodeName(init.name()));
}
EXPECT_TRUE(found);
}
namespace {
template <class T>
inline void SetNodeAttr(const std::string& key, const T& value, NodeDef* node) {
AttrValue attr_value;
SetAttrValue(value, &attr_value);
auto* attr_map = node->mutable_attr();
(*attr_map)[key] = attr_value;
}
template <>
inline void SetNodeAttr(const std::string& key, const Tensor& tensor,
NodeDef* node) {
TensorProto tensor_proto;
tensor.AsProtoTensorContent(&tensor_proto);
SetNodeAttr(key, tensor_proto, node);
}
} // namespace
TEST_F(SingleMachineTest, PersistentMemory) {
// Build a hashtable and its initialization graph.
GrapplerItem item;
const DataType key_dtype = DT_INT64;
const DataType data_dtype = DT_INT64;
NodeDef* hashtable_node = item.graph.add_node();
hashtable_node->set_op("HashTable");
hashtable_node->set_name("hash_table");
SetNodeAttr("key_dtype", key_dtype, hashtable_node);
SetNodeAttr("value_dtype", data_dtype, hashtable_node);
// Initial hashtable keys and values
NodeDef* keys_node = item.graph.add_node();
keys_node->set_op("Const");
keys_node->set_name("table_keys");
SetNodeAttr("dtype", key_dtype, keys_node);
Tensor keys(key_dtype, TensorShape{2});
keys.vec<int64_t>()(0) = 123;
keys.vec<int64_t>()(1) = 321;
SetNodeAttr("value", keys, keys_node);
NodeDef* values_node = item.graph.add_node();
values_node->set_op("Const");
values_node->set_name("table_values");
SetNodeAttr("dtype", data_dtype, values_node);
Tensor values(data_dtype, TensorShape{2});
values.vec<int64_t>()(0) = 789;
values.vec<int64_t>()(1) = 987;
SetNodeAttr("value", values, values_node);
// InitializeTable node
NodeDef* init_table_node = item.graph.add_node();
init_table_node->set_op("InitializeTable");
init_table_node->set_name("initialize_table");
SetNodeAttr("Tkey", key_dtype, init_table_node);
SetNodeAttr("Tval", data_dtype, init_table_node);
*init_table_node->add_input() = "hash_table";
*init_table_node->add_input() = "table_keys";
*init_table_node->add_input() = "table_values";
item.init_ops.push_back(init_table_node->name());
// Key to lookup
NodeDef* query_node = item.graph.add_node();
query_node->set_op("Const");
query_node->set_name("query");
SetNodeAttr("dtype", key_dtype, query_node);
Tensor query(key_dtype, TensorShape({}));
query.flat<int64_t>()(0) = 0;
SetNodeAttr("value", query, query_node);
// Default return value of hashtable lookup
NodeDef* default_value_node = item.graph.add_node();
default_value_node->set_op("Const");
default_value_node->set_name("default_table_value");
SetNodeAttr("dtype", data_dtype, default_value_node);
Tensor dflt(data_dtype, TensorShape({}));
dflt.flat<int64_t>()(0) = 456;
SetNodeAttr("value", dflt, default_value_node);
// HashTable lookup node
NodeDef* lookup_node = item.graph.add_node();
lookup_node->set_op("LookupTableFind");
lookup_node->set_name("table_lookup");
SetNodeAttr("Tin", key_dtype, lookup_node);
SetNodeAttr("Tout", data_dtype, lookup_node);
*lookup_node->add_input() = "hash_table";
*lookup_node->add_input() = "query";
*lookup_node->add_input() = "default_table_value";
item.fetch.push_back(lookup_node->name());
// Run the graph
TF_CHECK_OK(cluster_->Initialize(item));
RunMetadata metadata;
TF_CHECK_OK(cluster_->Run(item.graph, item.feed, item.fetch, &metadata));
// Check the cost model.
bool found_table_init = false;
bool found_hashtable = false;
for (const auto& node : metadata.cost_graph().node()) {
if (node.name() == "hash_table") {
found_hashtable = true;
// Persistent memory usage should be 0 since it's recorded as part of the
// initialize_table op.
EXPECT_EQ(0, node.persistent_memory_size());
} else if (node.name() == "initialize_table") {
found_table_init = true;
// Persistent memory should hold 2 keys and 2 values.
EXPECT_LE(4 * sizeof(int64_t), node.persistent_memory_size());
}
}
EXPECT_TRUE(found_table_init);
EXPECT_TRUE(found_hashtable);
}
GrapplerItem CreateGrapplerItemWithResourceMemory() {
tensorflow::Scope s = tensorflow::Scope::NewRootScope();
// Add a variable and initializer.
Output a = ops::Variable(s.WithOpName("a"), TensorShape({128, 256}),
DataType::DT_FLOAT);
Output a_init =
ops::RandomNormal(s.WithOpName("a/init"), {128, 256}, DataType::DT_FLOAT);
Output a_init_assign = ops::Assign(s.WithOpName("a/init/assign"), a, a_init);
// Add a resource variable.
Output b =
ops::VarHandleOp(s.WithOpName("b"), DataType::DT_FLOAT, {256, 512});
Output b_read =
ops::ReadVariableOp(s.WithOpName("b/read"), b, DataType::DT_FLOAT);
Output b_init =
ops::RandomNormal(s.WithOpName("b/init"), {256, 512}, DataType::DT_FLOAT);
auto b_init_assign =
ops::AssignVariableOp(s.WithOpName("b/init/assign"), b, b_init);
// Add a queue.
ops::FIFOQueue queue(s.WithOpName("queue"), {DataType::DT_STRING});
Output some_string =
ops::Const(s.WithOpName("some_string"), std::string("nothing"));
ops::QueueEnqueue enqueue(s.WithOpName("enqueue"), queue, {some_string});
ops::QueueDequeue dequeue(s.WithOpName("dequeue"), queue,
{DataType::DT_STRING});
// Add a IdentityReader.
ops::IdentityReader reader(s.WithOpName("identity_reader"));
ops::ReaderRead read(s.WithOpName("read_from_queue"), reader, queue);
Output var_mul = ops::MatMul(s.WithOpName("var_matmul"), a, b_read);
GrapplerItem item;
TF_CHECK_OK(s.ToGraphDef(&item.graph));
QueueRunnerDef queue_runner;
queue_runner.set_queue_name("queue");
*queue_runner.add_enqueue_op_name() = "enqueue";
item.queue_runners.push_back(queue_runner);
item.init_ops.push_back("a/init/assign");
item.init_ops.push_back("b/init/assign");
item.fetch.push_back("var_matmul");
item.fetch.push_back("dequeue");
return item;
}
#if defined(PLATFORM_GOOGLE)
TEST_F(SingleMachineTest, ReleaseMemoryAfterDestruction) {
GrapplerItem item = CreateGrapplerItemWithResourceMemory();
TF_CHECK_OK(cluster_->Initialize(item));
std::unordered_map<std::string, uint64_t> device_peak_memory_before;
TF_CHECK_OK(cluster_->GetPeakMemoryUsage(&device_peak_memory_before));
EXPECT_EQ(device_peak_memory_before.size(), 1);
// There might be a bit memory used before session's running anything.
EXPECT_LT(device_peak_memory_before.begin()->second, 400);
RunMetadata metadata;
TF_CHECK_OK(cluster_->Run(item.graph, item.feed, item.fetch, &metadata));
// Check there is memory that is not released.
std::unordered_map<std::string, uint64_t> device_peak_memory;
TF_CHECK_OK(cluster_->GetPeakMemoryUsage(&device_peak_memory));
EXPECT_EQ(device_peak_memory.size(), 1);
EXPECT_GT(device_peak_memory.begin()->second, 0);
// Reprovisioning the cluster would release all memory.
TF_CHECK_OK(cluster_->Shutdown());
TF_CHECK_OK(cluster_->Provision());
std::unordered_map<std::string, uint64_t> device_peak_memory_after;
TF_CHECK_OK(cluster_->GetPeakMemoryUsage(&device_peak_memory_after));
TF_CHECK_OK(cluster_->Shutdown());
// Check memory used by resources are released after cluster destruction.
EXPECT_EQ(device_peak_memory_before.size(), 1);
EXPECT_EQ(device_peak_memory_after.size(), 1);
EXPECT_LT(device_peak_memory_before.begin()->second, 400);
EXPECT_LT(device_peak_memory_after.begin()->second, 400);
}
TEST_F(SingleMachineTest, PeakMemory) {
GrapplerItem item = CreateGrapplerItemWithResourceMemory();
TF_CHECK_OK(cluster_->Initialize(item));
RunMetadata metadata;
TF_CHECK_OK(cluster_->Run(item.graph, item.feed, item.fetch, &metadata));
std::unordered_map<std::string, uint64_t> device_peak_memory;
TF_CHECK_OK(cluster_->GetPeakMemoryUsage(&device_peak_memory));
ASSERT_NE(
device_peak_memory.find("/job:localhost/replica:0/task:0/device:CPU:0"),
device_peak_memory.end());
uint64_t cpu_memory =
device_peak_memory["/job:localhost/replica:0/task:0/device:CPU:0"];
EXPECT_GT(cpu_memory, 0);
TF_CHECK_OK(cluster_->Shutdown());
TF_CHECK_OK(cluster_->Provision());
device_peak_memory.clear();
TF_CHECK_OK(cluster_->GetPeakMemoryUsage(&device_peak_memory));
TF_CHECK_OK(cluster_->Shutdown());
ASSERT_NE(
device_peak_memory.find("/job:localhost/replica:0/task:0/device:CPU:0"),
device_peak_memory.end());
cpu_memory =
device_peak_memory["/job:localhost/replica:0/task:0/device:CPU:0"];
EXPECT_LT(cpu_memory, 200);
}
TEST_F(SingleMachineTest, PeakMemoryStatsNotEnabled) {
GrapplerItem item = CreateGrapplerItemWithResourceMemory();
TF_CHECK_OK(cluster_->Shutdown());
cluster_.reset();
SingleMachine cluster(60 /* timeout_s */, 3 /* num_cpu_cores */,
0 /* num_gpus */);
TF_CHECK_OK(cluster.Provision());
TF_CHECK_OK(cluster.Initialize(item));
std::unordered_map<std::string, uint64_t> device_peak_memory;
absl::Status s = cluster.GetPeakMemoryUsage(&device_peak_memory);
TF_CHECK_OK(cluster.Shutdown());
ASSERT_FALSE(s.ok());
EXPECT_TRUE(absl::IsInvalidArgument(s));
}
#endif
} // namespace
} // namespace grappler
} // namespace tensorflow
+193
View File
@@ -0,0 +1,193 @@
/* 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/grappler/clusters/utils.h"
#include "Eigen/Core" // from @eigen_archive
#if GOOGLE_CUDA
#include "third_party/gpus/cuda/include/cuda.h"
#include "third_party/gpus/cuda/include/cuda_runtime_api.h"
#include "third_party/gpus/cudnn/cudnn.h"
#endif
#if TENSORFLOW_USE_ROCM
#include "rocm/include/hip/hip_runtime.h"
#endif
#include "tensorflow/core/common_runtime/gpu/gpu_id.h"
#include "tensorflow/core/common_runtime/gpu/gpu_id_manager.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/strings/numbers.h"
#include "tensorflow/core/lib/strings/strcat.h"
#include "tensorflow/core/platform/byte_order.h"
#include "tensorflow/core/platform/cpu_info.h"
#include "tensorflow/core/platform/mem.h"
namespace tensorflow {
namespace grappler {
DeviceProperties GetLocalCPUInfo() {
DeviceProperties device;
device.set_type("CPU");
device.set_vendor(port::CPUVendorIDString());
// Combine cpu family and model into the model string.
device.set_model(
absl::StrCat((port::CPUFamily() << 4) + port::CPUModelNum()));
device.set_frequency(port::NominalCPUFrequency() * 1e-6);
device.set_num_cores(port::NumSchedulableCPUs());
device.set_l1_cache_size(Eigen::l1CacheSize());
device.set_l2_cache_size(Eigen::l2CacheSize());
device.set_l3_cache_size(Eigen::l3CacheSize());
int64_t free_mem = port::AvailableRam();
if (free_mem < INT64_MAX) {
device.set_memory_size(free_mem);
}
(*device.mutable_environment())["cpu_instruction_set"] =
Eigen::SimdInstructionSetsInUse();
(*device.mutable_environment())["eigen"] = strings::StrCat(
EIGEN_WORLD_VERSION, ".", EIGEN_MAJOR_VERSION, ".", EIGEN_MINOR_VERSION);
return device;
}
DeviceProperties GetLocalGPUInfo(PlatformDeviceId platform_device_id) {
DeviceProperties device;
device.set_type("GPU");
#if GOOGLE_CUDA
cudaDeviceProp properties;
cudaError_t error =
cudaGetDeviceProperties(&properties, platform_device_id.value());
if (error != cudaSuccess) {
device.set_type("UNKNOWN");
LOG(ERROR) << "Failed to get device properties, error code: " << error;
return device;
}
device.set_vendor("NVIDIA");
device.set_model(properties.name);
int clockRate;
error = cudaDeviceGetAttribute(&clockRate, cudaDevAttrClockRate,
platform_device_id.value());
if (error != cudaSuccess) {
device.set_type("ERROR");
LOG(ERROR) << "Failed to get device attribute clockRate, error code: "
<< error;
return device;
}
device.set_frequency(clockRate * 1e-3);
device.set_num_cores(properties.multiProcessorCount);
device.set_num_registers(properties.regsPerMultiprocessor);
// For compute capability less than 5, l1 cache size is configurable to
// either 16 KB or 48 KB. We use the initial configuration 16 KB here. For
// compute capability larger or equal to 5, l1 cache (unified with texture
// cache) size is 24 KB. This number may need to be updated for future
// compute capabilities.
device.set_l1_cache_size((properties.major < 5) ? 16 * 1024 : 24 * 1024);
device.set_l2_cache_size(properties.l2CacheSize);
device.set_l3_cache_size(0);
device.set_shared_memory_size_per_multiprocessor(
properties.sharedMemPerMultiprocessor);
device.set_memory_size(properties.totalGlobalMem);
// 8 is the number of bits per byte. 2 is accounted for
// double data rate (DDR).
int memoryClockRate;
#if CUDART_VERSION >= 13000
error = cudaDeviceGetAttribute(&memoryClockRate, cudaDevAttrMemoryClockRate,
platform_device_id.value());
if (error != cudaSuccess) {
device.set_type("ERROR");
LOG(ERROR) << "Failed to get device attribute memoryClockRate, error code: "
<< error;
return device;
}
#else
memoryClockRate = properties.memoryClockRate;
#endif
device.set_bandwidth(properties.memoryBusWidth / 8 * memoryClockRate * 2ULL);
(*device.mutable_environment())["architecture"] =
absl::StrCat(properties.major, ".", properties.minor);
(*device.mutable_environment())["cuda"] = absl::StrCat(CUDA_VERSION);
(*device.mutable_environment())["cudnn"] = strings::StrCat(CUDNN_VERSION);
#elif TENSORFLOW_USE_ROCM
hipDeviceProp_t properties;
hipError_t error =
hipGetDeviceProperties(&properties, platform_device_id.value());
if (error != hipSuccess) {
device.set_type("UNKNOWN");
LOG(ERROR) << "Failed to get device properties, error code: " << error;
return device;
}
// ROCM TODO review if numbers here are valid
device.set_vendor("Advanced Micro Devices, Inc");
device.set_model(properties.name);
device.set_frequency(properties.clockRate * 1e-3);
device.set_num_cores(properties.multiProcessorCount);
device.set_num_registers(properties.regsPerBlock);
device.set_l1_cache_size(16 * 1024);
device.set_l2_cache_size(properties.l2CacheSize);
device.set_l3_cache_size(0);
device.set_shared_memory_size_per_multiprocessor(
properties.maxSharedMemoryPerMultiProcessor);
device.set_memory_size(properties.totalGlobalMem);
// 8 is the number of bits per byte. 2 is accounted for
// double data rate (DDR).
device.set_bandwidth(properties.memoryBusWidth / 8 *
properties.memoryClockRate * 2ULL);
(*device.mutable_environment())["architecture"] =
strings::StrCat("gfx", properties.gcnArchName);
#endif
return device;
}
DeviceProperties GetDeviceInfo(const DeviceNameUtils::ParsedName& device) {
DeviceProperties unknown;
unknown.set_type("UNKNOWN");
if (device.type == "CPU") {
return GetLocalCPUInfo();
} else if (device.type == "GPU") {
if (device.has_id) {
TfDeviceId tf_device_id(device.id);
PlatformDeviceId platform_device_id;
absl::Status s =
GpuIdManager::TfToPlatformDeviceId(tf_device_id, &platform_device_id);
if (!s.ok()) {
LOG(ERROR) << s;
return unknown;
}
return GetLocalGPUInfo(platform_device_id);
} else {
return GetLocalGPUInfo(PlatformDeviceId(0));
}
}
return unknown;
}
} // end namespace grappler
} // end namespace tensorflow
+39
View File
@@ -0,0 +1,39 @@
/* 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_GRAPPLER_CLUSTERS_UTILS_H_
#define TENSORFLOW_CORE_GRAPPLER_CLUSTERS_UTILS_H_
#include "tensorflow/core/common_runtime/gpu/gpu_id.h"
#include "tensorflow/core/protobuf/device_properties.pb.h"
#include "tensorflow/core/util/device_name_utils.h"
namespace tensorflow {
namespace grappler {
// Returns the DeviceProperties of the CPU on which grappler is running.
DeviceProperties GetLocalCPUInfo();
// Returns the DeviceProperties for the specified GPU attached to the server on
// which grappler is running.
DeviceProperties GetLocalGPUInfo(PlatformDeviceId platform_device_id);
// Returns the DeviceProperties of the specified device
DeviceProperties GetDeviceInfo(const DeviceNameUtils::ParsedName& device);
} // end namespace grappler
} // end namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPPLER_CLUSTERS_UTILS_H_
@@ -0,0 +1,121 @@
/* 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/grappler/clusters/utils.h"
#include "tensorflow/core/common_runtime/gpu/gpu_id.h"
#include "tensorflow/core/common_runtime/gpu/gpu_id_manager.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/protobuf/device_properties.pb.h"
namespace tensorflow {
namespace grappler {
namespace {
TEST(UtilsTest, GetLocalGPUInfo) {
GpuIdManager::TestOnlyReset();
#if GOOGLE_CUDA
LOG(INFO) << "CUDA is enabled.";
DeviceProperties properties;
// Invalid platform GPU ID.
properties = GetLocalGPUInfo(PlatformDeviceId(100));
EXPECT_EQ("UNKNOWN", properties.type());
// Succeed when a valid platform GPU id was inserted.
properties = GetLocalGPUInfo(PlatformDeviceId(0));
EXPECT_EQ("GPU", properties.type());
EXPECT_EQ("NVIDIA", properties.vendor());
#elif TENSORFLOW_USE_ROCM
LOG(INFO) << "ROCm is enabled.";
DeviceProperties properties;
// Invalid platform GPU ID.
properties = GetLocalGPUInfo(PlatformDeviceId(100));
EXPECT_EQ("UNKNOWN", properties.type());
// Succeed when a valid platform GPU id was inserted.
properties = GetLocalGPUInfo(PlatformDeviceId(0));
EXPECT_EQ("GPU", properties.type());
EXPECT_EQ("Advanced Micro Devices, Inc", properties.vendor());
#else
LOG(INFO) << "CUDA is not enabled.";
DeviceProperties properties;
properties = GetLocalGPUInfo(PlatformDeviceId(0));
EXPECT_EQ("GPU", properties.type());
properties = GetLocalGPUInfo(PlatformDeviceId(100));
EXPECT_EQ("GPU", properties.type());
#endif
}
TEST(UtilsTest, GetDeviceInfo) {
GpuIdManager::TestOnlyReset();
DeviceNameUtils::ParsedName device;
DeviceProperties properties;
// Invalid type.
properties = GetDeviceInfo(device);
EXPECT_EQ("UNKNOWN", properties.type());
// Cpu info.
device.type = "CPU";
properties = GetDeviceInfo(device);
EXPECT_EQ("CPU", properties.type());
// No TF GPU id provided.
device.type = "GPU";
device.has_id = false;
properties = GetDeviceInfo(device);
EXPECT_EQ("GPU", properties.type());
#if GOOGLE_CUDA
EXPECT_EQ("NVIDIA", properties.vendor());
#elif TENSORFLOW_USE_ROCM
EXPECT_EQ("Advanced Micro Devices, Inc", properties.vendor());
#endif
// TF to platform GPU id mapping entry doesn't exist.
device.has_id = true;
device.id = 0;
properties = GetDeviceInfo(device);
EXPECT_EQ("UNKNOWN", properties.type());
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
// Invalid platform GPU id.
TF_ASSERT_OK(GpuIdManager::InsertTfPlatformDeviceIdPair(
TfDeviceId(0), PlatformDeviceId(100)));
properties = GetDeviceInfo(device);
EXPECT_EQ("UNKNOWN", properties.type());
// Valid platform GPU id.
TF_ASSERT_OK(GpuIdManager::InsertTfPlatformDeviceIdPair(TfDeviceId(1),
PlatformDeviceId(0)));
device.id = 1;
properties = GetDeviceInfo(device);
EXPECT_EQ("GPU", properties.type());
#if GOOGLE_CUDA
EXPECT_EQ("NVIDIA", properties.vendor());
#elif TENSORFLOW_USE_ROCM
EXPECT_EQ("Advanced Micro Devices, Inc", properties.vendor());
#endif
#endif
}
} // namespace
} // namespace grappler
} // namespace tensorflow
@@ -0,0 +1,135 @@
/* 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/grappler/clusters/virtual_cluster.h"
#include <cstdint>
#include <memory>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "tensorflow/core/framework/cost_graph.pb.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/tensor_shape.pb.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/grappler/clusters/utils.h"
#include "tensorflow/core/grappler/costs/op_level_cost_estimator.h"
#include "tensorflow/core/protobuf/config.pb.h"
#include "tensorflow/core/protobuf/device_properties.pb.h"
namespace tensorflow {
namespace grappler {
VirtualCluster::VirtualCluster(
const std::unordered_map<std::string, DeviceProperties>& devices)
: VirtualCluster(devices, std::make_unique<OpLevelCostEstimator>(),
ReadyNodeManagerFactory("FirstReady")) {}
VirtualCluster::VirtualCluster(
const std::unordered_map<std::string, DeviceProperties>& devices,
std::unique_ptr<OpLevelCostEstimator> node_estimator,
std::unique_ptr<ReadyNodeManager> node_manager)
: Cluster(0) {
devices_ = devices;
// Note that we do not use aggressive shape inference to preserve unknown
// shapes from the input graph.
estimator_ = std::make_unique<AnalyticalCostEstimator>(
this, std::move(node_estimator), std::move(node_manager),
/*use_static_shapes=*/true, /*use_aggressive_shape_inference=*/false);
}
VirtualCluster::VirtualCluster(const DeviceSet* device_set)
: VirtualCluster(std::unordered_map<std::string, DeviceProperties>()) {
device_set_ = device_set;
for (const auto& device : device_set_->devices()) {
DeviceProperties props = GetDeviceInfo(device->parsed_name());
if (props.type() == "UNKNOWN") continue;
auto attrs = device->attributes();
props.set_memory_size(attrs.memory_limit());
devices_[device->name()] = props;
}
}
VirtualCluster::~VirtualCluster() {}
absl::Status VirtualCluster::Provision() { return absl::OkStatus(); }
absl::Status VirtualCluster::Initialize(const GrapplerItem& item) {
return absl::OkStatus();
}
absl::Status VirtualCluster::Run(
const GraphDef& graph,
const std::vector<std::pair<std::string, Tensor>>& feed,
const std::vector<std::string>& fetch, RunMetadata* metadata) {
GrapplerItem item;
item.graph = graph;
item.feed = feed;
item.fetch = fetch;
return Run(item, metadata);
}
absl::Status VirtualCluster::Run(const GrapplerItem& item,
RunMetadata* metadata) {
// Initializes an analytical cost estimator to estimate the graph cost. Makes
// sure to use static shape inference to prevent the virtual scheduler from
// calling the Run method on the cluster and creating an infinite loop.
if (metadata) {
metadata->clear_step_stats();
metadata->clear_cost_graph();
metadata->clear_partition_graphs();
}
TF_RETURN_IF_ERROR(estimator_->Initialize(item));
TF_RETURN_IF_ERROR(
estimator_->PredictCosts(item.graph, metadata, /*cost=*/nullptr));
const std::unordered_map<std::string, DeviceProperties>& device =
GetDevices();
std::unordered_map<std::string, int64_t> peak_mem_usage =
estimator_->GetScheduler()->GetPeakMemoryUsage();
for (const auto& mem_usage : peak_mem_usage) {
const std::string& device_name = mem_usage.first;
auto it = device.find(device_name);
if (it == device.end()) {
// It's probably the fake send/recv device. Eventually we'll need to
// remove this fake device to ensure proper memory accounting for
// multi-device settings.
continue;
}
const DeviceProperties& dev = it->second;
if (dev.memory_size() <= 0) {
// Available device memory unknown
continue;
}
int64_t peak_mem = mem_usage.second;
if (peak_mem >= dev.memory_size()) {
return absl::ResourceExhaustedError(absl::StrCat(
"Graph requires ", peak_mem, " bytes of memory on device ",
device_name, " to run ", " but device only has ", dev.memory_size(),
" available."));
}
}
return absl::OkStatus();
}
} // namespace grappler
} // namespace tensorflow
@@ -0,0 +1,70 @@
/* 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_GRAPPLER_CLUSTERS_VIRTUAL_CLUSTER_H_
#define TENSORFLOW_CORE_GRAPPLER_CLUSTERS_VIRTUAL_CLUSTER_H_
#include <memory>
#include <unordered_map>
#include <utility>
#include <vector>
#include "absl/status/status.h"
#include "tensorflow/core/common_runtime/device_set.h"
#include "tensorflow/core/grappler/clusters/cluster.h"
#include "tensorflow/core/grappler/costs/analytical_cost_estimator.h"
#include "tensorflow/core/grappler/costs/op_level_cost_estimator.h"
#include "tensorflow/core/grappler/costs/virtual_scheduler.h"
#include "tensorflow/core/protobuf/config.pb.h"
#include "tensorflow/core/protobuf/device_properties.pb.h"
namespace tensorflow {
namespace grappler {
// Create a simple cluster that lists the devices (and their properties)
// available in a TensorFlow session. This cluster simulates the execution of
// actual graphs.
class VirtualCluster : public Cluster {
public:
explicit VirtualCluster(
const std::unordered_map<std::string, DeviceProperties>& devices);
VirtualCluster(
const std::unordered_map<std::string, DeviceProperties>& devices,
std::unique_ptr<OpLevelCostEstimator> node_estimator,
std::unique_ptr<ReadyNodeManager> node_manager);
explicit VirtualCluster(const DeviceSet* device_set);
~VirtualCluster() override;
std::string type() const override { return "virtual"; }
absl::Status Provision() override;
absl::Status Initialize(const GrapplerItem& item) override;
absl::Status Run(const GraphDef& graph,
const std::vector<std::pair<std::string, Tensor>>& feed,
const std::vector<std::string>& fetch,
RunMetadata* metadata) override;
absl::Status Run(const GrapplerItem& item, RunMetadata* metadata) override;
const DeviceSet* GetDeviceSet() const override { return device_set_; }
private:
std::unique_ptr<AnalyticalCostEstimator> estimator_;
const DeviceSet* device_set_ = nullptr;
};
} // end namespace grappler
} // end namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPPLER_CLUSTERS_VIRTUAL_CLUSTER_H_
@@ -0,0 +1,133 @@
/* 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/grappler/clusters/virtual_cluster.h"
#include <memory>
#include <string>
#include <unordered_map>
#include "absl/log/check.h"
#include "absl/status/status.h"
#include "tensorflow/cc/framework/scope.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "xla/tsl/protobuf/error_codes.pb.h"
#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/framework/types.pb.h"
#include "tensorflow/core/grappler/grappler_item.h"
#include "tensorflow/core/grappler/inputs/trivial_test_graph_input_yielder.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/protobuf/config.pb.h"
#include "tensorflow/core/protobuf/device_properties.pb.h"
#include "tensorflow/core/protobuf/error_codes.pb.h"
namespace tensorflow {
namespace grappler {
namespace {
class VirtualClusterTest : public ::testing::Test {
public:
void SetUp() override {
// Invent a CPU so that predictions remain the same from machine to machine.
DeviceProperties cpu_device;
cpu_device.set_type("CPU");
cpu_device.set_frequency(1000);
cpu_device.set_num_cores(4);
cpu_device.set_bandwidth(32);
cpu_device.set_l1_cache_size(32 * 1024);
cpu_device.set_l2_cache_size(256 * 1024);
cpu_device.set_l3_cache_size(4 * 1024 * 1024);
cpu_device.set_memory_size(1024 * 1024);
std::unordered_map<std::string, DeviceProperties> devices;
devices["/job:localhost/replica:0/task:0/cpu:0"] = cpu_device;
cluster_ = std::make_unique<VirtualCluster>(devices);
TF_CHECK_OK(cluster_->Provision());
}
void TearDown() override {
TF_CHECK_OK(cluster_->Shutdown());
cluster_.reset();
}
protected:
std::unique_ptr<VirtualCluster> cluster_;
};
TEST_F(VirtualClusterTest, ClusterType) {
CHECK_EQ("virtual", cluster_->type());
}
TEST_F(VirtualClusterTest, CostModel) {
TrivialTestGraphInputYielder fake_input(4, 1, 10, false,
cluster_->GetDeviceNames());
GrapplerItem item;
CHECK(fake_input.NextItem(&item));
TF_CHECK_OK(cluster_->Initialize(item));
RunMetadata metadata;
TF_CHECK_OK(cluster_->Run(item.graph, item.feed, item.fetch, &metadata));
// There should be at least 4 nodes corresponding to the 4 stages we created
// in the fake input.
EXPECT_LE(4, metadata.cost_graph().node_size());
for (const auto& node : metadata.cost_graph().node()) {
// Skip the constant node that configures the random number generator.
if (node.name().find("Const/Const") != std::string::npos) {
continue;
}
EXPECT_EQ(1, node.output_info_size());
EXPECT_EQ(40, node.output_info(0).size());
const TensorShapeProto& shape = node.output_info(0).shape();
EXPECT_EQ(2, shape.dim_size());
EXPECT_EQ(10, shape.dim(0).size());
EXPECT_EQ(1, shape.dim(1).size());
if (node.name() == "x") {
EXPECT_EQ(1500, node.compute_cost());
} else {
EXPECT_EQ(2500, node.compute_cost());
}
}
for (const auto& dev_stat : metadata.step_stats().dev_stats()) {
EXPECT_EQ("/job:localhost/replica:0/task:0/cpu:0", dev_stat.device());
for (const auto& node : dev_stat.node_stats()) {
if (node.node_name() == "AddN") {
EXPECT_EQ(2500, node.op_end_rel_micros());
}
}
}
}
TEST_F(VirtualClusterTest, OutOfMemory) {
tensorflow::Scope root = tensorflow::Scope::NewRootScope();
// Create a large variable that can't fit in memory.
auto zero = ops::Variable(root.WithOpName("zero"), {1024, 1024}, DT_FLOAT);
auto identity = ops::Identity(root.WithOpName("i"), zero);
auto identity2 = ops::Identity(root.WithOpName("i2"), identity);
GrapplerItem item;
TF_CHECK_OK(root.ToGraphDef(&item.graph));
item.fetch.push_back("i2");
TF_CHECK_OK(cluster_->Initialize(item));
absl::Status s = cluster_->Run(item.graph, item.feed, item.fetch, nullptr);
EXPECT_EQ(error::RESOURCE_EXHAUSTED, s.code());
}
} // namespace
} // namespace grappler
} // namespace tensorflow
+423
View File
@@ -0,0 +1,423 @@
load("@rules_python//python:proto.bzl", "py_proto_library")
load(
"//tensorflow:tensorflow.bzl",
"tf_cc_test",
"tf_cuda_library",
)
load("//tensorflow:tensorflow.default.bzl", "filegroup")
load(
"//tensorflow/core/platform:build_config.bzl",
"tf_proto_library",
"tf_protos_grappler",
"tf_pyclif_proto_library",
)
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
licenses = ["notice"],
)
alias(
name = "graph_properties_testdata",
actual = "//tensorflow/core/grappler/costs/graph_properties_testdata:graph_properties_testdata",
visibility = ["//visibility:public"],
)
tf_proto_library(
name = "op_performance_data",
srcs = ["op_performance_data.proto"],
make_default_target_header_only = True,
protodeps = [
"//tensorflow/core/framework:attr_value_proto",
"//tensorflow/core/framework:resource_handle_proto",
"//tensorflow/core/framework:tensor_proto",
"//tensorflow/core/framework:tensor_shape_proto",
"//tensorflow/core/framework:types_proto",
"//tensorflow/core/protobuf:for_core_protos",
],
visibility = ["//visibility:public"],
)
tf_pyclif_proto_library(
name = "op_performance_data_pyclif",
proto_lib = ":op_performance_data",
proto_srcfile = "op_performance_data.proto",
visibility = ["//visibility:public"],
)
filegroup(
name = "pywrap_required_hdrs",
srcs = [
"analytical_cost_estimator.h",
"cost_estimator.h",
"graph_memory.h",
"graph_properties.h",
"measuring_cost_estimator.h",
"op_context.h",
"op_level_cost_estimator.h",
"utils.h",
"virtual_placer.h",
"virtual_scheduler.h",
],
visibility = [
"//tensorflow/python/grappler:__pkg__",
],
)
cc_library(
name = "graph_properties",
srcs = ["graph_properties.cc"],
hdrs = ["graph_properties.h"],
visibility = ["//visibility:public"],
deps = [
":utils",
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/grappler:grappler_item",
"//tensorflow/core/grappler:mutable_graph_view",
"//tensorflow/core/grappler:op_types",
"//tensorflow/core/grappler:utils",
"//tensorflow/core/grappler/clusters:cluster",
"//tensorflow/core/grappler/optimizers:evaluation_utils",
"//tensorflow/core/grappler/utils:functions",
"//tensorflow/core/grappler/utils:topological_sort",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/hash",
"@com_google_absl//absl/types:optional",
] + tf_protos_grappler(),
)
tf_cc_test(
name = "graph_properties_test",
srcs = ["graph_properties_test.cc"],
args = ["--heap_check="], # The GPU tracer leaks memory. TODO(b/185483595): use a dependency instead of a flag
data = [":graph_properties_testdata"],
tags = [
"nomsan", # TODO(b/160921160): broken by NOAUTOROLLBACK CL
],
deps = [
":graph_properties",
"//tensorflow/cc:cc_ops",
"//tensorflow/cc:functional_ops",
"//tensorflow/cc:scope",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:lib_proto_parsing",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/framework:tensor_testutil",
"//tensorflow/core/graph:mkl_graph_util",
"//tensorflow/core/grappler:grappler_item",
"//tensorflow/core/grappler/clusters:single_machine",
"//tensorflow/core/grappler/inputs:trivial_test_graph_input_yielder",
"//tensorflow/core/grappler/inputs:utils",
],
)
cc_library(
name = "graph_memory",
srcs = ["graph_memory.cc"],
hdrs = ["graph_memory.h"],
visibility = ["//visibility:public"],
deps = [
":cost_estimator",
":graph_properties",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/grappler:grappler_item",
"//tensorflow/core/grappler:utils",
"//tensorflow/core/grappler/clusters:cluster",
"//tensorflow/core/grappler/clusters:virtual_cluster",
],
)
tf_cc_test(
name = "graph_memory_test",
srcs = ["graph_memory_test.cc"],
args = ["--heap_check="], # The GPU tracer leaks memory. TODO(b/185483595): use a dependency instead of a flag
deps = [
":graph_memory",
"//tensorflow/cc:cc_ops",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/grappler:grappler_item",
"//tensorflow/core/grappler/inputs:trivial_test_graph_input_yielder",
],
)
cc_library(
name = "robust_stats",
srcs = ["robust_stats.cc"],
hdrs = ["robust_stats.h"],
visibility = ["//visibility:public"],
)
tf_cc_test(
name = "robust_stats_test",
srcs = ["robust_stats_test.cc"],
deps = [
":robust_stats",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
],
)
tf_cuda_library(
name = "utils",
srcs = ["utils.cc"],
hdrs = ["utils.h"],
visibility = ["//visibility:public"],
deps = [
":cost_estimator",
"//tensorflow/core:framework",
"//tensorflow/core:graph",
"//tensorflow/core:lib",
"//tensorflow/core:lib_proto_parsing",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/common_runtime/gpu:gpu_id",
"//tensorflow/core/grappler:utils",
"//tensorflow/core/grappler/clusters:utils",
"//tensorflow/core/util:overflow",
"@com_google_absl//absl/container:node_hash_map",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:str_format",
"@eigen_archive//:eigen3",
] + tf_protos_grappler(),
)
tf_cc_test(
name = "utils_test",
srcs = ["utils_test.cc"],
visibility = ["//visibility:public"],
deps = [
":utils",
"//tensorflow/cc:cc_ops",
"//tensorflow/core:all_kernels",
"//tensorflow/core:core_cpu_internal",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core:testlib",
"//tensorflow/core/framework:tensor_testutil",
],
)
cc_library(
name = "cost_estimator",
srcs = ["cost_estimator.cc"],
hdrs = ["cost_estimator.h"],
visibility = ["//visibility:public"],
deps = [
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:flat_hash_set",
],
)
tf_cc_test(
name = "cost_estimator_test",
srcs = ["cost_estimator_test.cc"],
deps = [
":cost_estimator",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"@com_google_googletest//:gtest",
],
)
cc_library(
name = "virtual_placer",
srcs = ["virtual_placer.cc"],
hdrs = ["virtual_placer.h"],
visibility = ["//visibility:public"],
deps = [
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/grappler:devices",
"//tensorflow/core/grappler/clusters:cluster",
],
)
tf_cc_test(
name = "virtual_placer_test",
srcs = ["virtual_placer_test.cc"],
deps = [
":virtual_placer",
"//tensorflow/core:core_cpu",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/grappler/clusters:virtual_cluster",
],
)
cc_library(
name = "op_context",
hdrs = ["op_context.h"],
visibility = ["//visibility:public"],
deps = [
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/container:flat_hash_map",
] + tf_protos_grappler(),
)
cc_library(
name = "virtual_scheduler",
srcs = ["virtual_scheduler.cc"],
hdrs = ["virtual_scheduler.h"],
visibility = ["//visibility:public"],
deps = [
":cost_estimator",
":graph_properties",
":op_context",
":utils",
":virtual_placer",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/grappler:grappler_item",
"//tensorflow/core/grappler:op_types",
"//tensorflow/core/grappler:utils",
"//tensorflow/core/grappler/clusters:utils",
"//tensorflow/core/grappler/utils:transitive_fanin",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:str_format",
],
)
tf_cc_test(
name = "virtual_scheduler_test",
srcs = ["virtual_scheduler_test.cc"],
deps = [
":graph_properties",
":utils",
":virtual_placer",
":virtual_scheduler",
"//tensorflow/cc:cc_ops",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:tensorflow",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/grappler/clusters:virtual_cluster",
"@com_google_absl//absl/strings",
],
)
cc_library(
name = "measuring_cost_estimator",
srcs = ["measuring_cost_estimator.cc"],
hdrs = ["measuring_cost_estimator.h"],
visibility = ["//visibility:public"],
deps = [
":robust_stats",
"//tensorflow/core:core_cpu",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:lib_internal",
"//tensorflow/core:lib_proto_parsing",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/grappler:grappler_item",
"//tensorflow/core/grappler/clusters:cluster",
"//tensorflow/core/grappler/costs:cost_estimator",
"//tensorflow/core/kernels:ops_util",
],
alwayslink = 1,
)
cc_library(
name = "op_level_cost_estimator",
srcs = ["op_level_cost_estimator.cc"],
hdrs = ["op_level_cost_estimator.h"],
visibility = ["//visibility:public"],
deps = [
":cost_estimator",
":op_context",
":utils",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/grappler/clusters:utils",
"//tensorflow/core/util:overflow",
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@eigen_archive//:eigen3",
"@tsl//tsl/platform:statusor",
] + tf_protos_grappler(),
)
tf_cc_test(
name = "op_level_cost_estimator_test",
srcs = ["op_level_cost_estimator_test.cc"],
tags = [
"no_oss",
"not_run:arm",
],
deps = [
":op_level_cost_estimator",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/platform:status_matchers",
],
)
cc_library(
name = "analytical_cost_estimator",
srcs = ["analytical_cost_estimator.cc"],
hdrs = ["analytical_cost_estimator.h"],
visibility = ["//visibility:public"],
deps = [
":cost_estimator",
":graph_properties",
":op_level_cost_estimator",
":utils",
":virtual_placer",
":virtual_scheduler",
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:graph",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/grappler:grappler_item",
"//tensorflow/core/grappler:op_types",
"//tensorflow/core/grappler:utils",
"//tensorflow/core/util:overflow",
] + tf_protos_grappler(),
alwayslink = 1,
)
tf_cc_test(
name = "analytical_cost_estimator_test",
srcs = ["analytical_cost_estimator_test.cc"],
deps = [
":analytical_cost_estimator",
":virtual_scheduler",
"//tensorflow/cc:cc_ops",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:tensorflow",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/grappler/clusters:virtual_cluster",
],
)
# copybara:uncomment_begin(google-only)
# py_proto_library(
# name = "op_performance_data_py_pb2",
# has_services = 0,
# visibility = ["//visibility:public"],
# deps = [":op_performance_data"],
# )
# copybara:uncomment_end
@@ -0,0 +1,251 @@
/* 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/grappler/costs/analytical_cost_estimator.h"
#include <limits>
#include <unordered_map>
#include "tensorflow/core/framework/tensor.pb.h" // NOLINT
#include "tensorflow/core/framework/tensor_shape.pb.h"
#include "tensorflow/core/graph/types.h"
#include "tensorflow/core/grappler/costs/graph_properties.h"
#include "tensorflow/core/grappler/costs/op_performance_data.pb.h"
#include "tensorflow/core/grappler/costs/utils.h"
#include "tensorflow/core/grappler/costs/virtual_placer.h"
#include "tensorflow/core/grappler/costs/virtual_scheduler.h"
#include "tensorflow/core/grappler/grappler_item.h"
#include "tensorflow/core/grappler/op_types.h"
#include "tensorflow/core/grappler/utils.h"
#include "tensorflow/core/util/overflow.h"
namespace tensorflow {
namespace grappler {
namespace {
// Helper function in PredictCosts() to add cost node to cost_graph.
absl::Status AddCostNode(
ReadyNodeManager* node_manager, const OpContext& op_context, int node_id,
const Costs& node_costs,
gtl::FlatMap<std::string, CostGraphDef::Node*>* name_to_cost_node,
gtl::FlatMap<std::string, int>* name_to_id, CostGraphDef* cost_graph) {
const std::string& op_name = op_context.name;
auto it = name_to_cost_node->find(op_name);
CostGraphDef::Node* node;
if (it != name_to_cost_node->end()) {
node = it->second;
node->clear_input_info();
node->clear_output_info();
} else {
node = cost_graph->add_node();
(*name_to_cost_node)[op_name] = node;
node->set_name(op_name);
node->set_id(node_id);
(*name_to_id)[node->name()] = node->id();
}
// For nodes we have seen before (e.g. Merge nodes are executed twice by
// VirtualScheduler), the following fields will be overwritten/updated.
node->set_device(op_context.device_name);
node->set_compute_cost(node_costs.execution_time.asMicroSeconds().count());
node->set_compute_time(node_costs.compute_time.asMicroSeconds().count());
node->set_memory_time(node_costs.memory_time.asMicroSeconds().count());
node->set_temporary_memory_size(node_costs.temporary_memory);
node->set_persistent_memory_size(node_costs.persistent_memory);
node->set_inaccurate(node_costs.inaccurate);
for (const std::string& input : node_manager->GetCurrNode()->input()) {
int input_port;
std::string input_name = ParseNodeName(input, &input_port);
// All inputs should have been seen already unless this is a Merge node.
if (name_to_id->find(input_name) == name_to_id->end()) {
if (!IsMerge(*node_manager->GetCurrNode()))
VLOG(1) << "input: " << input
<< " not found for non-Merge node: " << op_name;
// For Merge node, some of inputs may not be seen before
// For example, for a typical while loop in tensorflow, Merge node
// will be executed twice by VirtualScheduler (one for Enter, the
// other for NextIteration), so eventually both inputs will be added.
continue;
}
if (IsControlInput(input)) {
node->add_control_input(name_to_id->at(input_name));
} else {
auto* input_info = node->add_input_info();
input_info->set_preceding_node(name_to_id->at(input_name));
input_info->set_preceding_port(input_port);
}
}
for (const auto& output : op_context.op_info.outputs()) {
auto output_info = node->add_output_info();
output_info->set_alias_input_port(-1);
output_info->set_dtype(output.dtype());
*output_info->mutable_shape() = output.shape();
int64_t size = DataTypeSize(output.dtype());
for (const auto& dim : output.shape().dim()) {
size = MultiplyWithoutOverflow(size, std::max<int64_t>(1, dim.size()));
if (size < 0) {
return absl::InvalidArgumentError(
"Integer overflow encountered in dimension size.");
}
}
output_info->set_size(size);
}
return absl::OkStatus();
}
} // namespace
AnalyticalCostEstimator::AnalyticalCostEstimator(
Cluster* cluster, bool use_static_shapes,
bool use_aggressive_shape_inference)
: AnalyticalCostEstimator(
cluster, std::make_unique<OpLevelCostEstimator>(),
ReadyNodeManagerFactory("FirstReady"), use_static_shapes,
use_aggressive_shape_inference) {}
AnalyticalCostEstimator::AnalyticalCostEstimator(
Cluster* cluster, std::unique_ptr<OpLevelCostEstimator> node_estimator,
std::unique_ptr<ReadyNodeManager> node_manager, bool use_static_shapes,
bool use_aggressive_shape_inference)
: node_estimator_(std::move(node_estimator)),
node_manager_(std::move(node_manager)),
use_static_shapes_(use_static_shapes),
use_aggressive_shape_inference_(use_aggressive_shape_inference) {
scheduler_ = std::make_unique<VirtualScheduler>(
use_static_shapes_, use_aggressive_shape_inference_, cluster,
node_manager_.get(),
std::make_unique<VirtualPlacer>(cluster->GetDevices()));
}
AnalyticalCostEstimator::AnalyticalCostEstimator(
Cluster* cluster, std::unique_ptr<OpLevelCostEstimator> node_estimator,
std::unique_ptr<ReadyNodeManager> node_manager,
std::unique_ptr<VirtualPlacer> placer, bool use_static_shapes,
bool use_aggressive_shape_inference)
: node_estimator_(std::move(node_estimator)),
node_manager_(std::move(node_manager)),
use_static_shapes_(use_static_shapes),
use_aggressive_shape_inference_(use_aggressive_shape_inference) {
scheduler_ = std::make_unique<VirtualScheduler>(
use_static_shapes_, use_aggressive_shape_inference_, cluster,
node_manager_.get(), std::move(placer));
}
absl::Status AnalyticalCostEstimator::Initialize(const GrapplerItem& item) {
item_ = &item;
return absl::OkStatus();
}
absl::Status AnalyticalCostEstimator::PredictCosts(
const GraphDef& optimized_graph, RunMetadata* run_metadata,
Costs* costs) const {
std::unique_ptr<GrapplerItem> item_storage;
const GrapplerItem* item;
// Many callers to PredictCosts() pass the same optimized_graph as was used
// to initialize the estimator.
if (&optimized_graph == &item_->graph) {
item = item_;
} else {
GraphDef graph_copy = optimized_graph;
item_storage = std::make_unique<GrapplerItem>(
item_->WithGraph(std::move(graph_copy)));
item = item_storage.get();
}
auto status = scheduler_->Init(item);
if (!status.ok()) {
if (costs) {
costs->execution_time = Costs::Duration::max();
}
return status;
}
gtl::FlatMap<std::string, CostGraphDef::Node*> name_to_cost_node;
CostGraphDef* cost_graph = nullptr;
if (run_metadata) {
cost_graph = run_metadata->mutable_cost_graph();
// TODO(pcma): Clear nodes in cost_graph after we make sure we always pass
// in an empty cost_graph (a non-empty but incomplete cost_graph will cause
// problems, e.g., no node_id in cost_graph).
for (auto& node : *cost_graph->mutable_node()) {
name_to_cost_node[node.name()] = &node;
}
}
std::vector<std::string> inaccurate_nodes;
int nodes_executed = 0;
int node_id = 0;
gtl::FlatMap<std::string, int> name_to_id;
Costs node_costs;
do {
++nodes_executed;
OpContext op_context = scheduler_->GetCurrNode();
node_costs = node_estimator_->PredictCosts(op_context);
if (node_costs.inaccurate) {
inaccurate_nodes.push_back(op_context.name);
if (node_costs.num_ops_with_unknown_shapes > 0)
VLOG(4) << op_context.name << " has "
<< node_costs.num_ops_with_unknown_shapes << " unknown shapes";
}
// TODO(pcma): Add unit tests for generating CostGraphDef.
if (cost_graph) {
absl::Status s =
AddCostNode(node_manager_.get(), op_context, node_id++, node_costs,
&name_to_cost_node, &name_to_id, cost_graph);
if (!s.ok()) {
return s;
}
}
} while (scheduler_->MarkCurrNodeExecuted(node_costs));
VLOG(1) << inaccurate_nodes.size() << " out of " << nodes_executed
<< " nodes have inaccurate time estimation";
if (VLOG_IS_ON(3)) {
for (const auto& node : inaccurate_nodes) {
VLOG(4) << "Node with inaccurate time estimation: " << node;
}
}
// run_metadata gets step_stats and partition_graphs from Summary.
if (costs) {
*costs = scheduler_->Summary(run_metadata);
} else if (run_metadata) {
scheduler_->GenerateRunMetadata(run_metadata);
}
if (VLOG_IS_ON(1)) {
bool verbose = VLOG_IS_ON(2);
if (run_metadata) {
VLOG(1) << GetStatsStringFromRunMetadata(*run_metadata, verbose);
} else {
RunMetadata run_metadata;
scheduler_->GenerateRunMetadata(&run_metadata);
VLOG(1) << GetStatsStringFromRunMetadata(run_metadata, verbose);
}
}
return absl::OkStatus();
}
} // end namespace grappler
} // end namespace tensorflow
@@ -0,0 +1,83 @@
/* 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_GRAPPLER_COSTS_ANALYTICAL_COST_ESTIMATOR_H_
#define TENSORFLOW_CORE_GRAPPLER_COSTS_ANALYTICAL_COST_ESTIMATOR_H_
#include "tensorflow/core/grappler/costs/cost_estimator.h"
#include "tensorflow/core/grappler/costs/op_level_cost_estimator.h"
#include "tensorflow/core/grappler/costs/virtual_scheduler.h"
#include "tensorflow/core/grappler/grappler_item.h"
#include "tensorflow/core/lib/core/status.h"
namespace tensorflow {
class CostGraphDef;
class GraphDef;
} // namespace tensorflow
namespace tensorflow {
namespace grappler {
class Cluster;
struct GrapplerItem;
// Estimate the cost of running a Grappler item based on the theoretical
// performance of the hardware that will run the model. Note that this
// internally uses static shape inference. An option for aggressive shape
// inference is provided to minimize unknown shapes, and this is only applicable
// with static shape inference.
class AnalyticalCostEstimator : public CostEstimator {
public:
AnalyticalCostEstimator(Cluster* cluster, bool use_static_shapes,
bool use_aggressive_shape_inference);
AnalyticalCostEstimator(Cluster* cluster,
std::unique_ptr<OpLevelCostEstimator> node_estimator,
std::unique_ptr<ReadyNodeManager> node_manager,
bool use_static_shapes,
bool use_aggressive_shape_inference);
AnalyticalCostEstimator(Cluster* cluster,
std::unique_ptr<OpLevelCostEstimator> node_estimator,
std::unique_ptr<ReadyNodeManager> node_manager,
std::unique_ptr<VirtualPlacer> placer,
bool use_static_shapes,
bool use_aggressive_shape_inference);
~AnalyticalCostEstimator() override {}
// This implementation always returns OK.
absl::Status Initialize(const GrapplerItem& item) override;
// Predict the performance of each node of the optimized graph and annotate
// the RunMetadata with the corresponding estimates. Also returns the
// expected cost for the whole graph.
absl::Status PredictCosts(const GraphDef& optimized_graph,
RunMetadata* run_metadata,
Costs* cost) const override;
const VirtualScheduler* GetScheduler() const { return scheduler_.get(); }
private:
const GrapplerItem* item_;
std::unique_ptr<OpLevelCostEstimator> node_estimator_;
std::unique_ptr<ReadyNodeManager> node_manager_;
std::unique_ptr<VirtualScheduler> scheduler_;
bool use_static_shapes_;
bool use_aggressive_shape_inference_;
};
} // end namespace grappler
} // end namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPPLER_COSTS_ANALYTICAL_COST_ESTIMATOR_H_
@@ -0,0 +1,119 @@
/* 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/grappler/costs/virtual_scheduler.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/core/grappler/clusters/virtual_cluster.h"
#include "tensorflow/core/grappler/costs/analytical_cost_estimator.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace grappler {
class AnalyticalCostEstimatorTest : public ::testing::Test {
protected:
void SetUp() override {
// Initializes cluster_ and placer_.
std::unordered_map<std::string, DeviceProperties> devices;
DeviceProperties cpu_device;
cpu_device.set_type("CPU");
cpu_device.set_num_cores(4);
cpu_device.set_frequency(2600);
cpu_device.set_bandwidth(24 * 1024 * 1024);
devices["/job:localhost/replica:0/task:0/cpu:0"] = cpu_device;
DeviceProperties gpu_device;
gpu_device.set_type("GPU");
gpu_device.set_num_cores(12);
gpu_device.set_frequency(1100);
gpu_device.set_bandwidth(180 * 1024 * 1024);
(*gpu_device.mutable_environment())["architecture"] = "6";
devices["/job:localhost/replica:0/task:0/device:GPU:0"] = gpu_device;
cluster_.reset(new VirtualCluster(devices));
}
GrapplerItem CreateMiniGraph() {
const int batch = 1;
const int width = 28;
const int height = 28;
const int num_channels = 1;
const int num_labels = 10;
const int kernel_size = 3;
const int conv_filters = 32;
Scope s = Scope::NewRootScope();
auto images = ops::RandomUniform(
s.WithOpName("image"), {batch, width, height, num_channels}, DT_FLOAT);
auto labels = ops::RandomUniform(s.WithOpName("label"), {batch, num_labels},
DT_FLOAT);
auto w = ops::Variable(
s.WithOpName("W"),
{kernel_size, kernel_size, num_channels, conv_filters}, DT_FLOAT);
auto b = ops::Variable(s.WithOpName("B"), {conv_filters}, DT_FLOAT);
auto conv =
ops::Conv2D(s.WithOpName("conv"), images, w, {1, 1, 1, 1}, "SAME");
auto bias = ops::Add(s.WithOpName("bias"), conv, b);
auto relu = ops::Relu(s.WithOpName("relu"), bias);
auto flat_shape = ops::Const(s.WithOpName("flat_shape"),
{batch, width * height * conv_filters});
auto flat = ops::Reshape(s.WithOpName("flat"), relu, flat_shape);
auto w2 =
ops::Variable(s.WithOpName("W2"),
{width * height * conv_filters, num_labels}, DT_FLOAT);
auto b2 = ops::Variable(s.WithOpName("B2"), {num_labels}, DT_FLOAT);
auto matmul = ops::MatMul(s.WithOpName("matmul"), flat, w2);
auto logits = ops::Add(s.WithOpName("logits"), matmul, b2);
auto softmax = ops::Softmax(s.WithOpName("softmax"), logits);
auto lsm = ops::Log(s.WithOpName("lsm"), softmax);
GrapplerItem item;
item.fetch.push_back("lsm");
TF_CHECK_OK(s.ToGraphDef(&item.graph));
return item;
}
std::unique_ptr<VirtualCluster> cluster_;
};
TEST_F(AnalyticalCostEstimatorTest, SimpleTest) {
GrapplerItem item = CreateMiniGraph();
AnalyticalCostEstimator estimator(cluster_.get(), /*use_static_shapes=*/true,
/*use_aggressive_shape_inference=*/true);
TF_ASSERT_OK(estimator.Initialize(item));
RunMetadata run_metadata;
Costs summary;
TF_ASSERT_OK(estimator.PredictCosts(item.graph, &run_metadata, &summary));
EXPECT_EQ(Costs::NanoSeconds(9158), summary.execution_time);
// Note there are totally 17 nodes (RandomUniform creates 2 nodes), but
// grappler will not process "label", therefore we have 15 here instead
EXPECT_EQ(15, summary.num_ops_total);
// Make this estimate accurate:
// TODO(http://b/70031255): Accurate estimator for RandomUniform op needed
//
// Change to EXPECT_FALSE when the above TODOs are done:
EXPECT_TRUE(summary.inaccurate);
EXPECT_EQ(0, summary.num_ops_with_unknown_shapes);
}
} // end namespace grappler
} // end namespace tensorflow
@@ -0,0 +1,88 @@
/* 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/grappler/costs/cost_estimator.h"
namespace tensorflow {
namespace grappler {
Costs CombineCosts(const Costs& left, const Costs& right) {
CHECK_NE(left.max_memory, kMemoryUnknown);
CHECK_NE(left.max_per_op_buffers, kMemoryUnknown);
CHECK_NE(left.max_per_op_streaming, kMemoryUnknown);
Costs result = left;
result.execution_time += right.execution_time;
result.compute_time += right.compute_time;
result.memory_time += right.memory_time;
result.network_time += right.network_time;
result.intermediate_memory_time += right.intermediate_memory_time;
result.intermediate_memory_read_time += right.intermediate_memory_read_time;
result.intermediate_memory_write_time += right.intermediate_memory_write_time;
result.hbm_read_time += right.hbm_read_time;
result.hbm_write_time += right.hbm_write_time;
result.hbm_read_time_noderate += right.hbm_read_time_noderate;
result.hbm_write_time_noderate += right.hbm_write_time_noderate;
if (right.max_per_op_buffers != kMemoryUnknown) {
result.max_per_op_buffers =
std::max(left.max_per_op_buffers, right.max_per_op_buffers);
}
if (right.max_per_op_streaming != kMemoryUnknown) {
result.max_per_op_streaming =
std::max(left.max_per_op_streaming, right.max_per_op_streaming);
}
result.num_ops_total += right.num_ops_total;
if (right.inaccurate) {
result.inaccurate = true;
}
result.num_ops_with_unknown_shapes += right.num_ops_with_unknown_shapes;
if (right.max_memory != kMemoryUnknown) {
result.max_memory += right.max_memory;
}
return result;
}
// Multiplies Costs by a scalar.
// Equivalent to applying CombineCosts "multiplier" times.
// Note the field regarding num_ops and max_memory are not multiplied.
Costs MultiplyCosts(const Costs& costs, int multiplier) {
CHECK_GE(multiplier, 0);
if (multiplier == 0) {
return Costs::ZeroCosts();
}
if (multiplier == 1) {
return costs;
}
Costs result = costs;
result.execution_time *= multiplier;
result.compute_time *= multiplier;
result.memory_time *= multiplier;
result.hbm_read_time *= multiplier;
result.hbm_write_time *= multiplier;
result.hbm_read_time_noderate *= multiplier;
result.hbm_write_time_noderate *= multiplier;
result.network_time *= multiplier;
result.intermediate_memory_time *= multiplier;
result.intermediate_memory_read_time *= multiplier;
result.intermediate_memory_write_time *= multiplier;
return result;
}
} // end namespace grappler
} // end namespace tensorflow
@@ -0,0 +1,272 @@
/* 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_GRAPPLER_COSTS_COST_ESTIMATOR_H_
#define TENSORFLOW_CORE_GRAPPLER_COSTS_COST_ESTIMATOR_H_
#include <cmath>
#include <limits>
#include <string>
#include <unordered_map>
#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/protobuf/config.pb.h"
namespace tensorflow {
class GraphDef;
class CostGraphDef;
namespace grappler {
struct GrapplerItem;
constexpr uint64_t kMemoryUnknown = std::numeric_limits<uint64_t>::max();
constexpr uint64_t kZeroMemory = 0ULL;
struct DeviceInfo {
// Billions of operations executed per second.
double gigaops;
// Bandwidth to main memory in GB per second.
double gb_per_sec;
// Read bandwidth to intermediate memory in GB per second.
double intermediate_read_gb_per_sec;
// Write bandwidth to intermediate memory in GB per second.
double intermediate_write_gb_per_sec;
DeviceInfo()
: gigaops(INFINITY),
gb_per_sec(INFINITY),
intermediate_read_gb_per_sec(INFINITY),
intermediate_write_gb_per_sec(INFINITY) {}
DeviceInfo(const DeviceInfo& input)
: gigaops(input.gigaops),
gb_per_sec(input.gb_per_sec),
intermediate_read_gb_per_sec(input.intermediate_read_gb_per_sec),
intermediate_write_gb_per_sec(input.intermediate_write_gb_per_sec) {}
DeviceInfo(double gigaops, double gb_per_sec,
double intermediate_read_gb_per_sec = INFINITY,
double intermediate_write_gb_per_sec = INFINITY)
: gigaops(gigaops),
gb_per_sec(gb_per_sec),
intermediate_read_gb_per_sec(intermediate_read_gb_per_sec),
intermediate_write_gb_per_sec(intermediate_write_gb_per_sec) {}
};
// Holds the set of things we might want to estimate or measure in Grappler.
// Always produce execution time. Other fields are optional depending on the
// estimator being used.
struct Costs {
// Returns a Costs structure with default values for all of the fields.
inline Costs();
// Builds a Costs structure with all zero values, rather than unknowns.
static inline Costs ZeroCosts(bool inaccurate = false);
struct MilliSeconds : std::chrono::milliseconds {
MilliSeconds() : std::chrono::milliseconds(0) {}
MilliSeconds(double d)
: std::chrono::milliseconds(static_cast<int64_t>(d)) {}
MilliSeconds(const std::chrono::milliseconds& d)
: std::chrono::milliseconds(d) {}
MilliSeconds& operator=(const std::chrono::milliseconds& d) {
std::chrono::milliseconds::operator=(d);
return *this;
}
};
struct MicroSeconds : std::chrono::microseconds {
MicroSeconds() : std::chrono::microseconds(0) {}
MicroSeconds(double d)
: std::chrono::microseconds(static_cast<int64_t>(d)) {}
MicroSeconds(const std::chrono::microseconds& d)
: std::chrono::microseconds(d) {}
MicroSeconds& operator=(const std::chrono::microseconds& d) {
std::chrono::microseconds::operator=(d);
return *this;
}
MilliSeconds asMilliSeconds() const {
return std::chrono::duration_cast<std::chrono::milliseconds>(*this);
}
};
struct NanoSeconds : std::chrono::nanoseconds {
NanoSeconds() : std::chrono::nanoseconds(0) {}
NanoSeconds(double d) : std::chrono::nanoseconds(static_cast<int64_t>(d)) {}
NanoSeconds(const std::chrono::nanoseconds& d)
: std::chrono::nanoseconds(d) {}
NanoSeconds& operator=(const std::chrono::nanoseconds& d) {
std::chrono::nanoseconds::operator=(d);
return *this;
}
MicroSeconds asMicroSeconds() const {
return std::chrono::duration_cast<std::chrono::microseconds>(*this);
}
MilliSeconds asMilliSeconds() const {
return std::chrono::duration_cast<std::chrono::milliseconds>(*this);
}
static NanoSeconds infinity() {
return NanoSeconds(std::chrono::nanoseconds::max());
}
};
// We store all our times in nanoseconds. If needs be, we can always switch to
// picoseconds in the future by updating this typedef.
typedef NanoSeconds Duration;
// Overall cost of running the graph; latency.
Duration execution_time;
// Computation cost of running the graph.
Duration compute_time;
// Memory access cost of running the graph.
Duration memory_time;
// HBM read and write cost of running the graph. These are the achieved time
// where the HBM bw might haven been derated (from the peak bw). For small
// DMAs there is usually high derating, i.e. the achieved HBM bw is smaller
// than the peak bw and the achieved time is much larger than the theoretical
// roofline time.
Duration hbm_read_time;
Duration hbm_write_time;
// HBM read and write cost of running the graph without HBM bw derating. These
// are the theoretical roofline HBM time if the peak HBM bw is achieved.
Duration hbm_read_time_noderate;
Duration hbm_write_time_noderate;
// Intermediate memory access cost of running the graph
Duration intermediate_memory_time;
Duration intermediate_memory_read_time; // Intermediate memory read cost.
Duration intermediate_memory_write_time; // Intermediate memory write cost.
// Network time (colelctived ops - all gather, all reduce, etc.)
Duration network_time;
// This field can be a very pessimistic estimate of the main memory
// requirements of a graph. For example, it might assume that all activations
// are live for all of a graph's execution.
uint64_t max_memory; // Max main memory requirement in bytes over all ops.
uint64_t persistent_memory;
uint64_t temporary_memory;
// Output memory usage per port.
absl::flat_hash_map<int32_t, int64_t> output_tensor_size_bytes;
// Track persistent versus temporary memory.
absl::flat_hash_set<int32_t> persistent_output_ports;
// These fields are used for TPU-related estimations. They are per-op
// maximums, so each op is evaluated independently, but we want the maximum of
// the value over all ops.
int64_t max_per_op_buffers; // Sum of all buffers used by the ops.
int64_t max_per_op_streaming; // Ignore largest input buffer, assuming it
// streams from main memory.
// Number of ops included in this Costs in total.
// Default initialized to be one.
int64_t num_ops_total = 1;
// If the time estimation is inaccurate.
bool inaccurate = false;
// Number of ops that are estimated with unknown shapes.
int64_t num_ops_with_unknown_shapes = 0;
// TODO(pcma): include a counter for total inaccurate ops and counters for
// other reasons causing the inaccuracy
// Max possible memory usage per device.
std::unordered_map<std::string, uint64_t> estimated_max_memory_per_device;
};
inline std::ostream& operator<<(std::ostream& os, const Costs::MilliSeconds d) {
os << d.count() << "ms";
return os;
}
inline std::ostream& operator<<(std::ostream& os, const Costs::MicroSeconds d) {
os << d.count() << "us";
return os;
}
inline std::ostream& operator<<(std::ostream& os, const Costs::NanoSeconds d) {
os << d.count() << "ns";
return os;
}
Costs::Costs() {
execution_time = Duration::zero();
compute_time = Duration::zero();
memory_time = Duration::zero();
intermediate_memory_time = Duration::zero();
network_time = Duration::zero();
max_memory = kMemoryUnknown;
persistent_memory = kMemoryUnknown;
temporary_memory = kMemoryUnknown;
max_per_op_buffers = kMemoryUnknown;
max_per_op_streaming = kMemoryUnknown;
}
Costs Costs::ZeroCosts(bool inaccurate) {
Costs costs;
costs.execution_time = Duration::zero();
costs.compute_time = Duration::zero();
costs.memory_time = Duration::zero();
costs.intermediate_memory_time = Duration::zero();
costs.network_time = Duration::zero();
costs.max_memory = kZeroMemory;
costs.persistent_memory = kZeroMemory;
costs.temporary_memory = kZeroMemory;
costs.max_per_op_buffers = kZeroMemory;
costs.max_per_op_streaming = kZeroMemory;
costs.inaccurate = inaccurate;
return costs;
}
Costs CombineCosts(const Costs& left, const Costs& right);
// Multiplies Costs by a scalar.
// Equivalent to applying CombineCosts "multiplier" times.
Costs MultiplyCosts(const Costs& costs, int multiplier);
// Given a GrapperItem and an optimized implementation of the corresponding
// TensorFlow graph, the CostEstimator attempts to predicts the actual cost of
// running the graph.
class CostEstimator {
public:
virtual ~CostEstimator() {}
// Initializes the estimator for the specified grappler item.
// The estimator shouldn't be used if this function returns any status other
// that OK.
virtual absl::Status Initialize(const GrapplerItem& item) = 0;
// Predicts the cost of running the given optimized version of the grappler
// item.
// If a RunMetadata is passed, it will be populated with detailed information
// about the cost of running each operation of the optimized graph.
// if a double value is passed, it will be set to a value that reflects the
// overall cost of running the graph (e.g. the latency of the computation).
// Returns a status that indicate is the performance could be estimated or
// not.
virtual absl::Status PredictCosts(const GraphDef& optimized_graph,
RunMetadata* run_metadata,
Costs* cost) const = 0;
};
} // end namespace grappler
} // end namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPPLER_COSTS_COST_ESTIMATOR_H_
@@ -0,0 +1,101 @@
/* 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/grappler/costs/cost_estimator.h"
#include <gtest/gtest.h>
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace grappler {
namespace {
TEST(CostEstimatorTest, CombineCosts) {
Costs c = Costs::ZeroCosts();
c.execution_time = Costs::NanoSeconds(1);
c.compute_time = Costs::NanoSeconds(2);
c.memory_time = Costs::NanoSeconds(3);
c.intermediate_memory_time = Costs::NanoSeconds(4);
c.intermediate_memory_read_time = Costs::NanoSeconds(5);
c.intermediate_memory_write_time = Costs::NanoSeconds(6);
c.hbm_read_time = Costs::NanoSeconds(7);
c.hbm_write_time = Costs::NanoSeconds(8);
c.hbm_read_time_noderate = Costs::NanoSeconds(9);
c.hbm_write_time_noderate = Costs::NanoSeconds(10);
c.max_memory = 1;
c.max_per_op_buffers = 2;
c.max_per_op_streaming = 3;
c.num_ops_total = 1;
c.inaccurate = false;
c.num_ops_with_unknown_shapes = 0;
Costs sum = CombineCosts(c, c);
EXPECT_EQ(sum.execution_time, Costs::NanoSeconds(2));
EXPECT_EQ(sum.compute_time, Costs::NanoSeconds(4));
EXPECT_EQ(sum.memory_time, Costs::NanoSeconds(6));
EXPECT_EQ(sum.intermediate_memory_time, Costs::NanoSeconds(8));
EXPECT_EQ(sum.intermediate_memory_read_time, Costs::NanoSeconds(10));
EXPECT_EQ(sum.intermediate_memory_write_time, Costs::NanoSeconds(12));
EXPECT_EQ(sum.hbm_read_time, Costs::NanoSeconds(14));
EXPECT_EQ(sum.hbm_write_time, Costs::NanoSeconds(16));
EXPECT_EQ(sum.hbm_read_time_noderate, Costs::NanoSeconds(18));
EXPECT_EQ(sum.hbm_write_time_noderate, Costs::NanoSeconds(20));
EXPECT_EQ(sum.max_memory, 2);
EXPECT_EQ(sum.max_per_op_buffers, 2);
EXPECT_EQ(sum.max_per_op_streaming, 3);
EXPECT_EQ(sum.num_ops_total, 2);
EXPECT_FALSE(sum.inaccurate);
EXPECT_EQ(sum.num_ops_with_unknown_shapes, 0);
}
TEST(CostEstimatorTest, MultiplyCosts) {
Costs c = Costs::ZeroCosts();
c.execution_time = Costs::NanoSeconds(1);
c.compute_time = Costs::NanoSeconds(2);
c.memory_time = Costs::NanoSeconds(3);
c.intermediate_memory_time = Costs::NanoSeconds(4);
c.intermediate_memory_read_time = Costs::NanoSeconds(5);
c.intermediate_memory_write_time = Costs::NanoSeconds(6);
c.hbm_read_time_noderate = Costs::NanoSeconds(7);
c.hbm_write_time_noderate = Costs::NanoSeconds(8);
c.max_memory = 1;
c.max_per_op_buffers = 2;
c.max_per_op_streaming = 3;
c.num_ops_total = 1;
c.inaccurate = false;
c.num_ops_with_unknown_shapes = 0;
Costs product = MultiplyCosts(c, 10);
EXPECT_EQ(product.execution_time, Costs::NanoSeconds(10));
EXPECT_EQ(product.compute_time, Costs::NanoSeconds(20));
EXPECT_EQ(product.memory_time, Costs::NanoSeconds(30));
EXPECT_EQ(product.intermediate_memory_time, Costs::NanoSeconds(40));
EXPECT_EQ(product.intermediate_memory_read_time, Costs::NanoSeconds(50));
EXPECT_EQ(product.intermediate_memory_write_time, Costs::NanoSeconds(60));
EXPECT_EQ(product.hbm_read_time_noderate, Costs::NanoSeconds(70));
EXPECT_EQ(product.hbm_write_time_noderate, Costs::NanoSeconds(80));
EXPECT_EQ(product.max_memory, 1);
EXPECT_EQ(product.max_per_op_buffers, 2);
EXPECT_EQ(product.max_per_op_streaming, 3);
EXPECT_EQ(product.num_ops_total, 1);
EXPECT_FALSE(product.inaccurate);
EXPECT_EQ(product.num_ops_with_unknown_shapes, 0);
}
} // namespace
} // namespace grappler
} // namespace tensorflow
@@ -0,0 +1,294 @@
/* 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/grappler/costs/graph_memory.h"
#include <deque>
#include "tensorflow/core/framework/allocation_description.pb.h"
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/step_stats.pb.h"
#include "tensorflow/core/framework/tensor.pb.h" // NOLINT
#include "tensorflow/core/framework/tensor_description.pb.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/tensor_shape.pb.h"
#include "tensorflow/core/grappler/clusters/virtual_cluster.h"
#include "tensorflow/core/grappler/costs/graph_properties.h"
#include "tensorflow/core/grappler/utils.h"
namespace tensorflow {
namespace grappler {
absl::Status GraphMemory::InferStatically(
const std::unordered_map<std::string, DeviceProperties>& devices) {
VirtualCluster cluster(devices);
TF_RETURN_IF_ERROR(cluster.Provision());
TF_RETURN_IF_ERROR(cluster.Initialize(item_));
RunMetadata metadata;
absl::Status s = cluster.Run(item_, &metadata);
// The virtual cluster returns the RESOURCE_EXHAUSTED error when it detects
// that the model would run out of memory. We still get the metadata we need
// out of the simulation, so we just ignore this error.
if (!s.ok() && s.code() != error::RESOURCE_EXHAUSTED) {
return s;
}
InferFromTrace(metadata.step_stats());
return absl::OkStatus();
}
absl::Status GraphMemory::InferDynamically(Cluster* cluster) {
if (!cluster->DetailedStatsEnabled()) {
return absl::UnavailableError("Detailed stats collection must be enabled");
}
TF_RETURN_IF_ERROR(cluster->Initialize(item_));
RunMetadata metadata;
TF_RETURN_IF_ERROR(cluster->Run(item_, &metadata));
InferFromTrace(metadata.step_stats());
return absl::OkStatus();
}
int64_t GraphMemory::GetWorstCaseMemoryUsage() const {
int64_t worst_case = -1;
for (const auto& peak_usage : peak_usage_) {
worst_case = std::max(worst_case, peak_usage.second.used_memory);
}
return worst_case;
}
void GraphMemory::InferMemUsageForNodes(
const std::vector<const NodeDef*>& nodes, GraphProperties* properties,
int64_t* worst_case_memory_usage, int64_t* best_case_memory_usage) const {
// TODO(bsteiner) refine this: we should consider the multidevice case.
*worst_case_memory_usage = 0;
*best_case_memory_usage = 0;
for (const auto& node : item_.graph.node()) {
// Estimate the memory required to store the tensors generated by the node.
std::vector<OpInfo::TensorProperties> outputs =
properties->GetOutputProperties(node.name());
int64_t node_memory_usage = InferMemUsageForNeighbors(outputs);
// Worst case memory usage corresponds to the case where all the nodes are
// alive.
*worst_case_memory_usage += node_memory_usage;
// Estimate the memory required to store the input tensors needed by the
// node.
std::vector<OpInfo::TensorProperties> inputs =
properties->GetInputProperties(node.name());
node_memory_usage += InferMemUsageForNeighbors(inputs);
*best_case_memory_usage =
std::max(*best_case_memory_usage, node_memory_usage);
}
}
int64_t GraphMemory::InferMemUsageForNeighbors(
const std::vector<OpInfo::TensorProperties>& props) const {
int64_t neighbors_memory_usage = 0;
for (const auto& prop : props) {
DataType dtype = prop.dtype();
int size = DataTypeSize(dtype);
TensorShapeProto shape = prop.shape();
if (shape.unknown_rank()) {
// Can't infer the size if the rank is unknown, just skip.
continue;
}
// If one of the dimensions is unknown statically, assume it's one.
for (int i = 0; i < shape.dim_size(); ++i) {
if (shape.dim(i).size() < 0) {
shape.mutable_dim(i)->set_size(1);
}
}
int num_elems = TensorShape(shape).num_elements();
neighbors_memory_usage += num_elems * size;
}
return neighbors_memory_usage;
}
static GraphMemory::LiveTensor* FindOrCreateLiveTensor(
const std::string& node_name, int output_id,
std::unordered_map<std::string, GraphMemory::LiveTensor*>* live_tensors,
std::deque<GraphMemory::LiveTensor>* device_tensors) {
std::string name = absl::StrCat(node_name, ":", output_id);
GraphMemory::LiveTensor* live;
auto it = live_tensors->find(name);
if (it == live_tensors->end()) {
GraphMemory::LiveTensor temp;
temp.node = node_name;
temp.output_id = output_id;
temp.allocation_time = 0;
temp.deallocation_time = 0;
device_tensors->push_front(temp);
live = &device_tensors->front();
(*live_tensors)[name] = live;
} else {
live = it->second;
}
return live;
}
namespace {
struct Event {
Event(int64_t _timestamp, bool _allocated,
const GraphMemory::LiveTensor* _tensor)
: timestamp(_timestamp), allocated(_allocated), tensor(_tensor) {}
int64_t timestamp;
bool allocated;
const GraphMemory::LiveTensor* tensor;
bool operator<(const Event& other) const {
return timestamp < other.timestamp;
}
};
} // namespace
void GraphMemory::InferFromTrace(const StepStats& timeline) {
std::unordered_map<std::string, std::string> node_placement;
for (const auto& dev_stats : timeline.dev_stats()) {
for (const auto& node_stats : dev_stats.node_stats()) {
node_placement[node_stats.node_name()] = dev_stats.device();
}
}
std::unordered_map<std::string, LiveTensor*> live_tensors;
std::unordered_map<std::string, std::deque<LiveTensor>>
live_tensors_per_device;
std::unordered_map<std::string, const NodeDef*> node_map;
for (const NodeDef& node : item_.graph.node()) {
node_map[node.name()] = &node;
}
for (const auto& dev_stats : timeline.dev_stats()) {
const std::string& device_name = dev_stats.device();
const bool is_gpu = (device_name.find("GPU:") || device_name.find("gpu:"));
std::deque<LiveTensor>& device_tensors =
live_tensors_per_device[dev_stats.device()];
for (const auto& node_stats : dev_stats.node_stats()) {
for (int i = 0; i < node_stats.output_size(); ++i) {
const auto& output = node_stats.output(i);
LiveTensor* live = FindOrCreateLiveTensor(
node_stats.node_name(), i, &live_tensors, &device_tensors);
live->memory_used = output.tensor_description()
.allocation_description()
.allocated_bytes();
// Allocations typically take place at the very beginning of the op
// execution.
live->allocation_time =
Costs::MicroSeconds(node_stats.all_start_micros());
// Add one nanosecond to the completion time of the ops to account for
// TF overhead that slightly delays deallocations.
live->deallocation_time = std::max<Costs::Duration>(
live->deallocation_time,
Costs::NanoSeconds(1) +
Costs::MicroSeconds(node_stats.all_start_micros() +
node_stats.op_end_rel_micros()));
}
auto it = node_map.find(node_stats.node_name());
if (it == node_map.end()) {
// Skip nodes inserted by TF since they don't exist in the original
// graph (e.g _Send/_Recv nodes).
continue;
}
const NodeDef* node = it->second;
std::unordered_set<int> swapped_inputs;
if (is_gpu) {
auto it = node->attr().find("_swap_to_host");
if (it != node->attr().end()) {
const AttrValue& val = it->second;
for (int port_id : val.list().i()) {
swapped_inputs.insert(port_id);
}
}
}
for (int i = 0; i < node->input_size(); ++i) {
if (swapped_inputs.find(i) != swapped_inputs.end()) {
// The memory of swapped inputs will be released as early as possible:
// therefore ignore this input when determining the deallocation time
// of the tensor.
continue;
}
const std::string& input = node->input(i);
int position;
std::string input_node = ParseNodeName(input, &position);
if (position < 0) {
// Skip control dependencies
continue;
}
LiveTensor* live = FindOrCreateLiveTensor(
input_node, position, &live_tensors,
&live_tensors_per_device[node_placement[input_node]]);
live->deallocation_time = std::max<Costs::Duration>(
live->deallocation_time,
Costs::NanoSeconds(1) +
Costs::MicroSeconds(node_stats.all_start_micros() +
node_stats.op_end_rel_micros()));
}
}
}
for (const auto& live_per_device : live_tensors_per_device) {
std::vector<Event> events;
events.reserve(2 * live_per_device.second.size());
for (const auto& live : live_per_device.second) {
events.emplace_back(static_cast<int64_t>(live.allocation_time.count()),
true, &live);
events.emplace_back(static_cast<int64_t>(live.deallocation_time.count()),
false, &live);
}
std::stable_sort(events.begin(), events.end());
size_t peak = 0;
std::unordered_set<const LiveTensor*> live_at_peak;
size_t current = 0;
std::unordered_set<const LiveTensor*> currently_live;
int events_size = events.size();
for (int i = 0; i < events_size; ++i) {
const auto& event = events[i];
if (event.allocated) {
VLOG(1) << "At time " << event.timestamp << " allocated "
<< event.tensor->memory_used << " for tensor "
<< event.tensor->node << ":" << event.tensor->output_id;
current += event.tensor->memory_used;
currently_live.insert(event.tensor);
} else {
VLOG(1) << "At time " << event.timestamp << " deallocated "
<< event.tensor->memory_used << " for tensor "
<< event.tensor->node << ":" << event.tensor->output_id;
current -= event.tensor->memory_used;
currently_live.erase(event.tensor);
}
if (i + 1 == events_size || event.timestamp != events[i + 1].timestamp) {
if (current > peak) {
peak = current;
live_at_peak = currently_live;
}
}
}
MemoryUsage& peak_mem_usage = peak_usage_[live_per_device.first];
peak_mem_usage.used_memory = peak;
peak_mem_usage.live_tensors.clear();
peak_mem_usage.live_tensors.reserve(live_at_peak.size());
for (const auto& live : live_at_peak) {
peak_mem_usage.live_tensors.push_back(*live);
}
}
}
} // end namespace grappler
} // end namespace tensorflow
@@ -0,0 +1,81 @@
/* 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_GRAPPLER_COSTS_GRAPH_MEMORY_H_
#define TENSORFLOW_CORE_GRAPPLER_COSTS_GRAPH_MEMORY_H_
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/grappler/clusters/cluster.h"
#include "tensorflow/core/grappler/costs/cost_estimator.h"
#include "tensorflow/core/grappler/costs/graph_properties.h"
#include "tensorflow/core/grappler/grappler_item.h"
namespace tensorflow {
namespace grappler {
// Infer the worst case memory usage for a given grappler item.
class GraphMemory {
public:
struct LiveTensor {
std::string node;
int output_id;
size_t memory_used;
Costs::Duration allocation_time;
Costs::Duration deallocation_time;
};
struct MemoryUsage {
int64_t used_memory;
std::vector<LiveTensor> live_tensors;
};
explicit GraphMemory(const GrapplerItem& item)
: item_(item), unknown_usage_({-1, {}}) {}
absl::Status InferStatically(
const std::unordered_map<std::string, DeviceProperties>& devices);
absl::Status InferDynamically(Cluster* cluster);
// Worst case memory usage in bytes, or -1 if the usage is unknown. If there
// are multiple devices, returns the highest per device memory usage.
int64_t GetWorstCaseMemoryUsage() const;
// Returns the peak memory usage for the specified device.
const MemoryUsage& GetPeakMemoryUsage(const std::string& device) const {
auto it = peak_usage_.find(device);
if (it == peak_usage_.end()) {
return unknown_usage_;
}
return it->second;
}
private:
void InferMemUsageForNodes(const std::vector<const NodeDef*>& nodes,
GraphProperties* properties, int64_t* worst_case,
int64_t* best_case) const;
int64_t InferMemUsageForNeighbors(
const std::vector<OpInfo::TensorProperties>& props) const;
void InferFromTrace(const StepStats& timeline);
const GrapplerItem& item_;
std::unordered_map<std::string, int64_t> worst_case_memory_usage_;
std::unordered_map<std::string, MemoryUsage> peak_usage_;
const MemoryUsage unknown_usage_;
};
} // end namespace grappler
} // end namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPPLER_COSTS_GRAPH_MEMORY_H_
@@ -0,0 +1,228 @@
/* 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/grappler/costs/graph_memory.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/core/grappler/grappler_item.h"
#include "tensorflow/core/grappler/inputs/trivial_test_graph_input_yielder.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace grappler {
namespace {
class GraphMemoryTest : public ::testing::Test {
protected:
std::unordered_map<std::string, DeviceProperties> devices_;
public:
GraphMemoryTest() {
devices_["/CPU:0"].set_type("CPU");
devices_["/CPU:0"].set_num_cores(1);
devices_["/CPU:0"].set_frequency(1);
devices_["/CPU:0"].set_bandwidth(1);
devices_["/GPU:0"].set_type("GPU");
devices_["/GPU:0"].set_num_cores(1);
devices_["/GPU:0"].set_frequency(1);
devices_["/CPU:0"].set_bandwidth(1);
(*devices_["/GPU:0"].mutable_environment())["architecture"] = "3";
}
};
TEST_F(GraphMemoryTest, Basic) {
TrivialTestGraphInputYielder fake_input(4, 1, 10, false, {"/CPU:0"});
GrapplerItem item;
CHECK(fake_input.NextItem(&item));
item.feed.clear();
GraphMemory memory(item);
absl::Status s = memory.InferStatically(devices_);
TF_CHECK_OK(s);
const GraphMemory::MemoryUsage& mem_usage =
memory.GetPeakMemoryUsage("/CPU:0");
EXPECT_EQ(120, mem_usage.used_memory);
std::set<std::string> tensors;
for (const auto& t : mem_usage.live_tensors) {
tensors.insert(absl::StrCat(t.node, ":", t.output_id));
}
// When the execution of the 'Sign' node completes, TF can start executing
// 'Sign_1' and release the memory used by 'x'. Since we can't be sure of
// the order in which this takes place, in the worst case the 3 tensors are in
// memory.
std::set<std::string> expected;
expected.insert("Sign:0");
expected.insert("Sign_1:0");
expected.insert("x:0");
EXPECT_EQ(expected, tensors);
}
TEST_F(GraphMemoryTest, UnknownBatchSize) {
TrivialTestGraphInputYielder fake_input(4, 1, -1, false, {"/CPU:0"});
GrapplerItem item;
CHECK(fake_input.NextItem(&item));
item.feed.clear();
GraphMemory memory(item);
absl::Status s = memory.InferStatically(devices_);
TF_CHECK_OK(s);
// Same maths as before, except that batch size is unknown and therefore
// assumed to be one.
const GraphMemory::MemoryUsage& mem_usage =
memory.GetPeakMemoryUsage("/CPU:0");
EXPECT_EQ(16, mem_usage.used_memory);
std::set<std::string> tensors;
for (const auto& t : mem_usage.live_tensors) {
tensors.insert(absl::StrCat(t.node, ":", t.output_id));
}
std::set<std::string> expected;
expected.insert("Const/Const:0");
expected.insert("Sign:0");
expected.insert("x:0");
EXPECT_EQ(expected, tensors);
}
TEST_F(GraphMemoryTest, MultiDevice) {
TrivialTestGraphInputYielder fake_input(4, 2, 1024 * 1024, false,
{"/CPU:0", "/GPU:0"});
GrapplerItem item;
CHECK(fake_input.NextItem(&item));
item.feed.clear();
GraphMemory memory(item);
absl::Status s = memory.InferStatically(devices_);
TF_CHECK_OK(s);
const GraphMemory::MemoryUsage& cpu_mem = memory.GetPeakMemoryUsage("/CPU:0");
EXPECT_EQ(16777216, cpu_mem.used_memory);
std::set<std::string> cpu_tensors;
for (const auto& t : cpu_mem.live_tensors) {
cpu_tensors.insert(absl::StrCat(t.node, ":", t.output_id));
}
std::set<std::string> cpu_expected;
cpu_expected.insert("Recv_Sign_1_0_on_/CPU_0:0");
cpu_expected.insert("Sign:0");
cpu_expected.insert("x:0");
cpu_expected.insert("AddN:0");
EXPECT_EQ(cpu_expected, cpu_tensors);
const GraphMemory::MemoryUsage& gpu_mem = memory.GetPeakMemoryUsage("/GPU:0");
EXPECT_EQ(16777216, gpu_mem.used_memory);
std::set<std::string> gpu_tensors;
for (const auto& t : gpu_mem.live_tensors) {
gpu_tensors.insert(absl::StrCat(t.node, ":", t.output_id));
}
std::set<std::string> gpu_expected;
gpu_expected.insert("Recv_AddN_0_on_/GPU_0:0");
gpu_expected.insert("Sign_1:0");
gpu_expected.insert("AddN_1:0");
gpu_expected.insert("AddN_3:0");
EXPECT_EQ(gpu_expected, gpu_tensors);
}
TEST_F(GraphMemoryTest, GpuSwapping) {
TrivialTestGraphInputYielder fake_input(4, 2, 1024 * 1024, false, {"/GPU:0"});
GrapplerItem item;
CHECK(fake_input.NextItem(&item));
item.feed.clear();
{
// Estimate the max memory usage for the graph.
GraphMemory memory(item);
absl::Status s = memory.InferStatically(devices_);
TF_CHECK_OK(s);
const GraphMemory::MemoryUsage& gpu_mem =
memory.GetPeakMemoryUsage("/GPU:0");
EXPECT_EQ(20971520, gpu_mem.used_memory);
std::set<std::string> gpu_tensors;
for (const auto& t : gpu_mem.live_tensors) {
gpu_tensors.insert(absl::StrCat(t.node, ":", t.output_id));
}
std::set<std::string> gpu_expected;
gpu_expected.insert("Sign:0");
gpu_expected.insert("Sign_1:0");
gpu_expected.insert("AddN:0");
gpu_expected.insert("AddN_1:0");
gpu_expected.insert("AddN_2:0");
EXPECT_EQ(gpu_expected, gpu_tensors);
}
{
// Swap the first input to node AddN_1: its fanin (the square nodes) should
// not appear in the max cut anymore.
for (auto& node : *item.graph.mutable_node()) {
if (node.name() == "AddN_1") {
(*node.mutable_attr())["_swap_to_host"].mutable_list()->add_i(0);
}
}
GraphMemory memory(item);
absl::Status s = memory.InferStatically(devices_);
TF_CHECK_OK(s);
const GraphMemory::MemoryUsage& new_gpu_mem =
memory.GetPeakMemoryUsage("/GPU:0");
EXPECT_EQ(20971520, new_gpu_mem.used_memory);
std::set<std::string> new_gpu_tensors;
for (const auto& t : new_gpu_mem.live_tensors) {
new_gpu_tensors.insert(absl::StrCat(t.node, ":", t.output_id));
}
std::set<std::string> new_gpu_expected;
new_gpu_expected.insert("AddN:0");
new_gpu_expected.insert("AddN_1:0");
new_gpu_expected.insert("AddN_2:0");
new_gpu_expected.insert("AddN_3:0");
new_gpu_expected.insert("AddN_4:0");
EXPECT_EQ(new_gpu_expected, new_gpu_tensors);
}
}
TEST_F(GraphMemoryTest, CtrlDependencies) {
// Build a simple graph with a control dependency.
Scope s = Scope::NewRootScope();
Output a = ops::Const(s.WithOpName("a").WithDevice("/CPU:0"), 10.0f, {3});
Output v =
ops::Variable(s.WithOpName("v").WithDevice("/CPU:0"), {3}, DT_FLOAT);
Output assign =
ops::Assign(s.WithOpName("assign").WithDevice("/CPU:0"), v, a);
ops::NoOp init(
s.WithOpName("init").WithDevice("/CPU:0").WithControlDependencies(
assign));
GrapplerItem item;
item.fetch.push_back("init");
TF_CHECK_OK(s.ToGraphDef(&item.graph));
GraphMemory memory(item);
absl::Status status = memory.InferStatically(devices_);
TF_CHECK_OK(status);
const GraphMemory::MemoryUsage& mem = memory.GetPeakMemoryUsage("/CPU:0");
EXPECT_EQ(36, mem.used_memory);
std::set<std::string> tensors;
for (const auto& t : mem.live_tensors) {
tensors.insert(absl::StrCat(t.node, ":", t.output_id));
}
std::set<std::string> expected;
expected.insert("a:0");
expected.insert("v:0");
expected.insert("assign:0");
EXPECT_EQ(expected, tensors);
}
} // namespace
} // namespace grappler
} // namespace tensorflow
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,226 @@
/* 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_GRAPPLER_COSTS_GRAPH_PROPERTIES_H_
#define TENSORFLOW_CORE_GRAPPLER_COSTS_GRAPH_PROPERTIES_H_
#include <unordered_map>
#include <unordered_set>
#include <vector>
#include "absl/container/flat_hash_map.h"
#include "tensorflow/core/framework/shape_inference.h"
#include "tensorflow/core/grappler/clusters/cluster.h"
#include "tensorflow/core/grappler/costs/op_performance_data.pb.h"
#include "tensorflow/core/grappler/grappler_item.h"
namespace tensorflow {
namespace grappler {
// Optional attributes that tell about node output information.
// We use these side information, if provided, for static shape inference
// and VirtualScheduler scheduling.
// Switch op attribute as a vector of int that tells which branch the
// Switch output is taken on every round of execution.
// Used for scheduling ops after Switch correctly (e.g., While loop).
ABSL_CONST_INIT const char kOutputSlots[] = "_output_slot_vector";
// Example:
// Assume a node has two outputs and iterated for three times. Then it has:
// _execution_count = 3
// _output_sizes_vector = [2, 2, 2]
// _output_dtype_vector.size = 6
// _output_shape_vector.size = 6
// If all the iterations have same output shapes, then
// _execution_count = 3
// _same_output_for_iterations = true
// _output_sizes_vector = [2]
// _output_dtype_vector.size = 2
// _output_shape_vector.size = 2
// How many times this node has been executed.
ABSL_CONST_INIT const char kExecutionCount[] = "_execution_count";
// Records the output sizes for each round of execution.
ABSL_CONST_INIT const char kOutputSizes[] = "_output_sizes_vector";
// The node has been scheduled multiple times with outputs that have the same
// shape.
ABSL_CONST_INIT const char kOutputSame[] = "_same_output_for_iterations";
// Outputs DataType vector.
ABSL_CONST_INIT const char kOutputTypes[] = "_output_dtype_vector";
// Outputs TensorShapeProto vector.
ABSL_CONST_INIT const char kOutputShapes[] = "_output_shape_vector";
class SymbolicShapeRefiner;
class TopoQueue;
// Infer OpInfo::TensorProperties for graph nodes inputs/outputs.
//
// Typical use case, is to infer tensor properties from a graph, before doing
// optimization pass. Nodes modified during optimization pass have to be
// invalidated, to prevent further incorrect optimizations based on wrong shape
// and data type properties.
class GraphProperties {
public:
// The item must outlive the properties
explicit GraphProperties(const GrapplerItem& item) : item_(item) {}
// Infer the shapes through abstract interpretation. Feed information can be
// incorrect so it should be discarded to ensure correctness of the analysis.
// However, it can help infer shapes in the fanout of fed nodes (even though
// the correctness of these shapes can't be guaranteed), so in some cases
// (such as simulation or scheduling) it makes sense of keep these shapes.
// aggressive_shape_inference option executes nodes on the host to identify
// output values when possible and does other aggressive strategies.
// Similar to assuming_valid_feeds, this may cause incorrectness in graph
// analyses, but is useful for simulation or scheduling.
// If include_input_tensor_values is true, the values of constant tensors
// will included in the input properties.
// If include_output_tensor_values is true, the values of constant tensors
// will be included in the output properties.
absl::Status InferStatically(bool assume_valid_feeds,
bool aggressive_shape_inference,
bool include_input_tensor_values,
bool include_output_tensor_values);
absl::Status InferStatically(bool assume_valid_feeds,
bool aggressive_shape_inference,
bool include_tensor_values) {
return InferStatically(
assume_valid_feeds,
/*aggressive_shape_inference=*/aggressive_shape_inference,
/*include_input_tensor_values=*/include_tensor_values,
/*include_output_tensor_values=*/include_tensor_values);
}
absl::Status InferStatically(bool assume_valid_feeds) {
return InferStatically(assume_valid_feeds,
/*aggressive_shape_inference=*/false,
/*include_tensor_values=*/true);
}
// Infer the shape by running the graph on the specified cluster and recording
// the shapes of the processed tensors.
absl::Status InferDynamically(Cluster* cluster);
// Extract the properties from a cost graph. For testing only since there is
// no way to ensure that the cost graph match the item.
absl::Status InferFromCostGraph(const CostGraphDef& cost_graph);
// Stores `item_.graph` with the inferred output shapes to `output_graph_def`.
absl::Status AnnotateOutputShapes(GraphDef* output_graph_def) const;
// Return the properties of node inputs/outputs, including data types and
// shapes. Note that the dimensions in the shapes can be negative. We use the
// -1 value to denote that we don't know anything about a dimension. We use
// values strictly less than -1 to encode symbolic dimensions: although we
// don't know the actual value of the symbolic dimension, we know that all the
// dimensions denoted by the same negative value are the equal.
bool HasInputProperties(const std::string& node_name) const;
bool HasOutputProperties(const std::string& node_name) const;
const std::vector<OpInfo::TensorProperties>& GetInputProperties(
const std::string& node_name) const;
const std::vector<OpInfo::TensorProperties>& GetOutputProperties(
const std::string& node_name) const;
// Invalidate input/output properties for nodes modified during graph
// optimization pass, to prevent potential optimizations, based on incorrect
// shape information.
void ClearInputProperties(const std::string& node_name);
void ClearOutputProperties(const std::string& node_name);
// Returns true if we have *any* properties.
bool has_properties() const {
return !input_properties_.empty() || !output_properties_.empty();
}
bool CheckShapeIncompatible(const std::string& node_name) const {
return incompatible_shape_nodes_.find(node_name) !=
incompatible_shape_nodes_.end();
}
// Clear all infered properties.
void Clear() {
input_properties_.clear();
output_properties_.clear();
}
private:
// Relaxes shapes <shapes_and_types>, determined from an EnqueueV2 node, into
// <*queue_shapes_and_types>.
static absl::Status RelaxEnqueueShapesAndMergeTypes(
SymbolicShapeRefiner* shape_refiner, const NodeDef* qnode,
const std::vector<shape_inference::ShapeAndType>& shapes_and_types,
std::vector<shape_inference::ShapeAndType>* queue_shapes_and_types);
// Update the shapes of the enqueue node, port them over to the corresponding
// queue, and schedule the reprocessing of the queue if needed.
static absl::Status UpdateEnqueue(
const NodeDef* enqueue_node,
const absl::flat_hash_map<const NodeDef*, const NodeDef*>&
resource_handles,
SymbolicShapeRefiner* shape_refiner, bool* new_shapes);
// Update the shapes and types of the Queue node, if not set by Enqueue node.
static absl::Status UpdateQueue(const NodeDef* queue_node,
SymbolicShapeRefiner* shape_refiner,
bool* new_shapes);
// Update the output shapes of a Merge node, and enqueue its fanout in
// new_shapes if needed.
absl::Status UpdateMerge(SymbolicShapeRefiner* shape_refiner,
const NodeDef* node, bool* new_shapes) const;
// Process the Enter node, and enqueue its fanout in new_shapes if needed.
static absl::Status UpdateEnter(SymbolicShapeRefiner* shape_refiner,
const NodeDef* node, bool* new_shapes);
// Update the shapes for node 'n'. If output shapes for n have changed,
// enqueue its fanout in 'new_shapes'.
absl::Status UpdateShapes(
SymbolicShapeRefiner* shape_refiner,
const absl::flat_hash_map<const NodeDef*, const NodeDef*>&
resource_handles,
const NodeDef* n, bool* new_shapes) const;
// Propagate the shapes for the nodes enqueued in new_shapes and their
// transitive fanout until a fixed point is reached.
absl::Status PropagateShapes(
SymbolicShapeRefiner* shape_refiner, TopoQueue* new_shapes,
const absl::flat_hash_map<const NodeDef*, const NodeDef*>&
resource_handles,
int num_loops) const;
// Data members
const GrapplerItem& item_;
absl::flat_hash_map<std::string, std::vector<OpInfo::TensorProperties>>
input_properties_;
absl::flat_hash_map<std::string, std::vector<OpInfo::TensorProperties>>
output_properties_;
const std::vector<OpInfo::TensorProperties> missing_properties_;
// Nodes with output shape incompatible between shape inference and
// annotation.
std::unordered_set<std::string> incompatible_shape_nodes_;
};
// Helper function for GraphProperties.
bool IsShapeFullyDefinedIntegerVectorOrScalar(
shape_inference::InferenceContext* ic,
const shape_inference::ShapeHandle& shape,
const shape_inference::ShapeHandle& tensor_as_shape, const DataType& dtype);
} // end namespace grappler
} // end namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPPLER_COSTS_GRAPH_PROPERTIES_H_
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,18 @@
# Description:
# Graph properties test data.
load("//tensorflow:tensorflow.default.bzl", "filegroup")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
licenses = ["notice"],
)
filegroup(
name = "graph_properties_testdata",
srcs = glob([
"*.pbtxt",
"*.pbtxt.html",
]),
visibility = ["//visibility:public"],
)
@@ -0,0 +1,117 @@
node {
name: "Const"
op: "Const"
attr {
key: "dtype"
value {
type: DT_FLOAT
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_FLOAT
tensor_shape {
dim {
size: 1
}
dim {
size: 2
}
}
float_val: 2.0
}
}
}
}
node {
name: "Const_1"
op: "Const"
attr {
key: "dtype"
value {
type: DT_FLOAT
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_FLOAT
tensor_shape {
dim {
size: 1
}
dim {
size: 2
}
}
float_val: 2.0
}
}
}
}
node {
name: "MyAdd_yabA4wXEdM4"
op: "MyAdd_yabA4wXEdM4"
input: "Const"
input: "Const_1"
}
library {
function {
signature {
name: "MyAdd_yabA4wXEdM4"
input_arg {
name: "x"
type: DT_FLOAT
}
input_arg {
name: "y"
type: DT_FLOAT
}
output_arg {
name: "add_1"
type: DT_FLOAT
}
}
node_def {
name: "Add"
op: "Add"
input: "x"
input: "Add:z:0"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
node_def {
name: "Add_1"
op: "Add"
input: "Add:z:0"
input: "y"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
ret {
key: "add_1"
value: "Add_1:z:0"
}
attr {
key: "_noinline"
value {
b: true
}
}
}
}
versions {
producer: 26
min_consumer: 12
}
@@ -0,0 +1,246 @@
node {
name: "MyFunc_AenMyWWx1Us"
op: "MyFunc_AenMyWWx1Us"
}
library {
function {
signature {
name: "MyFunc_AenMyWWx1Us"
output_arg {
name: "while"
type: DT_INT32
}
output_arg {
name: "while_0"
type: DT_FLOAT
}
is_stateful: true
}
node_def {
name: "Const"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 10
}
}
}
}
node_def {
name: "While/input_1"
op: "Const"
attr {
key: "dtype"
value {
type: DT_FLOAT
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_FLOAT
tensor_shape {
}
float_val: 0.0
}
}
}
}
node_def {
name: "While"
op: "While"
input: "Const:output:0"
input: "While/input_1:output:0"
attr {
key: "T"
value {
list {
type: DT_INT32
type: DT_FLOAT
}
}
}
attr {
key: "body"
value {
func {
name: "Body_8GOMGeZeK5c"
}
}
}
attr {
key: "cond"
value {
func {
name: "Cond_Xf5ttAHgUCg"
}
}
}
attr {
key: "output_shapes"
value {
list {
}
}
}
}
ret {
key: "while"
value: "While:output:0"
}
ret {
key: "while_0"
value: "While:output:1"
}
attr {
key: "_noinline"
value {
b: true
}
}
}
function {
signature {
name: "Body_8GOMGeZeK5c"
input_arg {
name: "n"
type: DT_FLOAT
}
input_arg {
name: "x"
type: DT_FLOAT
}
output_arg {
name: "sub"
type: DT_FLOAT
}
output_arg {
name: "add"
type: DT_FLOAT
}
}
node_def {
name: "sub/y"
op: "Const"
attr {
key: "dtype"
value {
type: DT_FLOAT
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_FLOAT
tensor_shape {
}
float_val: 1.0
}
}
}
}
node_def {
name: "sub_0"
op: "Sub"
input: "n"
input: "sub/y:output:0"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
node_def {
name: "add_0"
op: "Add"
input: "x"
input: "n"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
ret {
key: "add"
value: "add_0:z:0"
}
ret {
key: "sub"
value: "sub_0:z:0"
}
}
function {
signature {
name: "Cond_Xf5ttAHgUCg"
input_arg {
name: "n"
type: DT_FLOAT
}
input_arg {
name: "unused_x"
type: DT_FLOAT
}
output_arg {
name: "greater"
type: DT_BOOL
}
}
node_def {
name: "Greater/y"
op: "Const"
attr {
key: "dtype"
value {
type: DT_FLOAT
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_FLOAT
tensor_shape {
}
float_val: 0.0
}
}
}
}
node_def {
name: "Greater"
op: "Greater"
input: "n"
input: "Greater/y:output:0"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
ret {
key: "greater"
value: "Greater:z:0"
}
}
}
versions {
producer: 26
min_consumer: 12
}
@@ -0,0 +1,251 @@
node {
name: "Const"
op: "Const"
attr {
key: "dtype"
value {
type: DT_FLOAT
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_FLOAT
tensor_shape {
dim {
size: 1
}
dim {
size: 2
}
}
float_val: 2.0
}
}
}
}
node {
name: "Less/x"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 0
}
}
}
}
node {
name: "Less/y"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 1
}
}
}
}
node {
name: "Less"
op: "Less"
input: "Less/x"
input: "Less/y"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "case/cond/Switch"
op: "Switch"
input: "Less"
input: "Less"
attr {
key: "T"
value {
type: DT_BOOL
}
}
}
node {
name: "case/cond/switch_t"
op: "Identity"
input: "case/cond/Switch:1"
attr {
key: "T"
value {
type: DT_BOOL
}
}
}
node {
name: "case/cond/switch_f"
op: "Identity"
input: "case/cond/Switch"
attr {
key: "T"
value {
type: DT_BOOL
}
}
}
node {
name: "case/cond/pred_id"
op: "Identity"
input: "Less"
attr {
key: "T"
value {
type: DT_BOOL
}
}
}
node {
name: "case/cond/Const"
op: "Const"
input: "^case/cond/switch_t"
attr {
key: "dtype"
value {
type: DT_FLOAT
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_FLOAT
tensor_shape {
dim {
size: 1
}
dim {
size: 2
}
}
float_val: 2.0
}
}
}
}
node {
name: "case/cond/Const_1"
op: "Const"
input: "^case/cond/switch_f"
attr {
key: "dtype"
value {
type: DT_FLOAT
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_FLOAT
tensor_shape {
dim {
size: 1
}
dim {
size: 2
}
}
float_val: 2.0
}
}
}
}
node {
name: "case/cond/Merge"
op: "Merge"
input: "case/cond/Const_1"
input: "case/cond/Const"
attr {
key: "N"
value {
i: 2
}
}
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
node {
name: "MyAdd_MPaeanipb7o"
op: "MyAdd_MPaeanipb7o"
input: "case/cond/Merge"
input: "Const"
}
library {
function {
signature {
name: "MyAdd_MPaeanipb7o"
input_arg {
name: "x"
type: DT_FLOAT
}
input_arg {
name: "y"
type: DT_FLOAT
}
output_arg {
name: "Add"
type: DT_FLOAT
}
}
node_def {
name: "Add"
op: "Add"
input: "x"
input: "y"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
ret {
key: "Add"
value: "Add:z:0"
}
attr {
key: "_noinline"
value {
b: true
}
}
}
}
versions {
producer: 26
min_consumer: 12
}
@@ -0,0 +1,251 @@
node {
name: "Const"
op: "Const"
attr {
key: "dtype"
value {
type: DT_FLOAT
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_FLOAT
tensor_shape {
dim {
size: 1
}
dim {
size: 2
}
}
float_val: 2.0
}
}
}
}
node {
name: "Less/x"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 1
}
}
}
}
node {
name: "Less/y"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 0
}
}
}
}
node {
name: "Less"
op: "Less"
input: "Less/x"
input: "Less/y"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "case/cond/Switch"
op: "Switch"
input: "Less"
input: "Less"
attr {
key: "T"
value {
type: DT_BOOL
}
}
}
node {
name: "case/cond/switch_t"
op: "Identity"
input: "case/cond/Switch:1"
attr {
key: "T"
value {
type: DT_BOOL
}
}
}
node {
name: "case/cond/switch_f"
op: "Identity"
input: "case/cond/Switch"
attr {
key: "T"
value {
type: DT_BOOL
}
}
}
node {
name: "case/cond/pred_id"
op: "Identity"
input: "Less"
attr {
key: "T"
value {
type: DT_BOOL
}
}
}
node {
name: "case/cond/Const"
op: "Const"
input: "^case/cond/switch_t"
attr {
key: "dtype"
value {
type: DT_FLOAT
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_FLOAT
tensor_shape {
dim {
size: 1
}
dim {
size: 2
}
}
float_val: 2.0
}
}
}
}
node {
name: "case/cond/Const_1"
op: "Const"
input: "^case/cond/switch_f"
attr {
key: "dtype"
value {
type: DT_FLOAT
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_FLOAT
tensor_shape {
dim {
size: 1
}
dim {
size: 2
}
}
float_val: 2.0
}
}
}
}
node {
name: "case/cond/Merge"
op: "Merge"
input: "case/cond/Const_1"
input: "case/cond/Const"
attr {
key: "N"
value {
i: 2
}
}
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
node {
name: "MyAdd_MPaeanipb7o"
op: "MyAdd_MPaeanipb7o"
input: "case/cond/Merge"
input: "Const"
}
library {
function {
signature {
name: "MyAdd_MPaeanipb7o"
input_arg {
name: "x"
type: DT_FLOAT
}
input_arg {
name: "y"
type: DT_FLOAT
}
output_arg {
name: "Add"
type: DT_FLOAT
}
}
node_def {
name: "Add"
op: "Add"
input: "x"
input: "y"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
ret {
key: "Add"
value: "Add:z:0"
}
attr {
key: "_noinline"
value {
b: true
}
}
}
}
versions {
producer: 26
min_consumer: 12
}
@@ -0,0 +1,317 @@
node {
name: "Const"
op: "Const"
attr {
key: "dtype"
value {
type: DT_FLOAT
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_FLOAT
tensor_shape {
dim {
size: 1
}
dim {
size: 3
}
}
float_val: 2.0
}
}
}
}
node {
name: "Less/x"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 1
}
}
}
}
node {
name: "Less/y"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 0
}
}
}
}
node {
name: "Less"
op: "Less"
input: "Less/x"
input: "Less/y"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "case/cond/Switch"
op: "Switch"
input: "Less"
input: "Less"
attr {
key: "T"
value {
type: DT_BOOL
}
}
}
node {
name: "case/cond/switch_t"
op: "Identity"
input: "case/cond/Switch:1"
attr {
key: "T"
value {
type: DT_BOOL
}
}
}
node {
name: "case/cond/switch_f"
op: "Identity"
input: "case/cond/Switch"
attr {
key: "T"
value {
type: DT_BOOL
}
}
}
node {
name: "case/cond/pred_id"
op: "Identity"
input: "Less"
attr {
key: "T"
value {
type: DT_BOOL
}
}
}
node {
name: "case/cond/Const"
op: "Const"
input: "^case/cond/switch_t"
attr {
key: "dtype"
value {
type: DT_FLOAT
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_FLOAT
tensor_shape {
dim {
size: 1
}
dim {
size: 2
}
}
float_val: 2.0
}
}
}
}
node {
name: "case/cond/Const_1"
op: "Const"
input: "^case/cond/switch_f"
attr {
key: "dtype"
value {
type: DT_FLOAT
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_FLOAT
tensor_shape {
dim {
size: 1
}
dim {
size: 2
}
}
float_val: 2.0
}
}
}
}
node {
name: "case/cond/Merge"
op: "Merge"
input: "case/cond/Const_1"
input: "case/cond/Const"
attr {
key: "N"
value {
i: 2
}
}
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
node {
name: "MyAdd_lEKAAnIwI5I"
op: "MyAdd_lEKAAnIwI5I"
input: "case/cond/Merge"
input: "Const"
}
library {
function {
signature {
name: "MyAdd_lEKAAnIwI5I"
input_arg {
name: "x"
type: DT_FLOAT
}
input_arg {
name: "y"
type: DT_FLOAT
}
output_arg {
name: "Add"
type: DT_FLOAT
}
}
node_def {
name: "Const"
op: "Const"
attr {
key: "dtype"
value {
type: DT_FLOAT
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_FLOAT
tensor_shape {
dim {
size: 1
}
dim {
size: 2
}
}
float_val: 2.0
}
}
}
}
node_def {
name: "Const_1"
op: "Const"
attr {
key: "dtype"
value {
type: DT_FLOAT
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_FLOAT
tensor_shape {
dim {
size: 1
}
dim {
size: 3
}
}
float_val: 2.0
}
}
}
}
node_def {
name: "Add"
op: "Add"
input: "x"
input: "Const:output:0"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
node_def {
name: "Add_1"
op: "Add"
input: "y"
input: "Const_1:output:0"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
ret {
key: "Add"
value: "Add:z:0"
}
attr {
key: "_noinline"
value {
b: true
}
}
}
}
versions {
producer: 26
min_consumer: 12
}
@@ -0,0 +1,130 @@
node {
name: "x"
op: "Const"
attr {
key: "dtype"
value {
type: DT_FLOAT
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_FLOAT
tensor_shape {
dim {
size: 1
}
dim {
size: 3
}
}
float_val: 2.0
}
}
}
}
node {
name: "y"
op: "_Arg"
attr {
key: "T"
value {
type: DT_RESOURCE
}
}
attr {
key: "_handle_dtypes"
value {
list {
type: DT_FLOAT
}
}
}
attr {
key: "_handle_shapes"
value {
list {
shape {
dim {
size: 1
}
dim {
size: 3
}
}
}
}
}
attr {
key: "index"
value {
i: 0
}
}
}
node {
name: "FunctionWithDtResourceInput"
op: "FunctionWithDtResourceInput"
input: "x"
input: "y"
}
library {
function {
signature {
name: "FunctionWithDtResourceInput"
input_arg {
name: "x"
type: DT_FLOAT
}
input_arg {
name: "y"
type: DT_RESOURCE
}
output_arg {
name: "z1"
type: DT_FLOAT
}
output_arg {
name: "z2"
type: DT_FLOAT
}
}
node_def {
name: "y1"
op: "ReadVariableOp"
input: "y"
attr {
key: "dtype"
value {
type: DT_FLOAT
}
}
}
node_def {
name: "Add"
op: "Add"
input: "x"
input: "y1:value:0"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
ret {
key: "z1"
value: "Add:z:0"
}
ret {
key: "z2"
value: "x"
}
}
}
versions {
producer: 26
min_consumer: 12
}
@@ -0,0 +1,613 @@
node {
name: "Const/Const"
op: "Const"
device: "/cpu:0"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
dim {
size: 1
}
}
int_val: 64
}
}
}
}
node {
name: "input_0_0"
op: "RandomUniform"
input: "Const/Const"
device: "/cpu:0"
attr {
key: "T"
value {
type: DT_INT32
}
}
attr {
key: "dtype"
value {
type: DT_FLOAT
}
}
attr {
key: "seed"
value {
i: 0
}
}
attr {
key: "seed2"
value {
i: 0
}
}
}
node {
name: "Const_1/Const"
op: "Const"
device: "/cpu:0"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
dim {
size: 4
}
}
int_val: 1
int_val: 1
int_val: 24
int_val: 64
}
}
}
}
node {
name: "input_1_0"
op: "RandomUniform"
input: "Const_1/Const"
device: "/cpu:0"
attr {
key: "T"
value {
type: DT_INT32
}
}
attr {
key: "dtype"
value {
type: DT_FLOAT
}
}
attr {
key: "seed"
value {
i: 0
}
}
attr {
key: "seed2"
value {
i: 0
}
}
}
node {
name: "Const_2/Const"
op: "Const"
device: "/cpu:0"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
dim {
size: 4
}
}
int_val: 128
int_val: 224
int_val: 224
int_val: 3
}
}
}
}
node {
name: "input_2_0"
op: "RandomUniform"
input: "Const_2/Const"
device: "/cpu:0"
attr {
key: "T"
value {
type: DT_INT32
}
}
attr {
key: "dtype"
value {
type: DT_FLOAT
}
}
attr {
key: "seed"
value {
i: 0
}
}
attr {
key: "seed2"
value {
i: 0
}
}
}
node {
name: "Const_3/Const"
op: "Const"
device: "/cpu:0"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
dim {
size: 4
}
}
int_val: 7
int_val: 7
int_val: 3
int_val: 8
}
}
}
}
node {
name: "input_3_0"
op: "RandomUniform"
input: "Const_3/Const"
device: "/cpu:0"
attr {
key: "T"
value {
type: DT_INT32
}
}
attr {
key: "dtype"
value {
type: DT_FLOAT
}
}
attr {
key: "seed"
value {
i: 0
}
}
attr {
key: "seed2"
value {
i: 0
}
}
}
node {
name: "y0"
op: "BiasAddx1_Conv2Dx1_DepthwiseConv2dNativex1_Relux1_95"
input: "input_0_0"
input: "input_1_0"
input: "input_2_0"
input: "input_3_0"
input: "^input_0_0"
input: "^input_1_0"
input: "^input_2_0"
input: "^input_3_0"
device: "/cpu:0"
}
node {
name: "shape"
op: "Shape"
input: "y0"
input: "^input_0_0"
input: "^input_1_0"
input: "^input_2_0"
input: "^input_3_0"
device: "/cpu:0"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
attr {
key: "out_type"
value {
type: DT_INT32
}
}
}
node {
name: "zeros"
op: "ZerosLike"
input: "shape"
input: "^input_0_0"
input: "^input_1_0"
input: "^input_2_0"
input: "^input_3_0"
device: "/cpu:0"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "ones"
op: "OnesLike"
input: "shape"
input: "^input_0_0"
input: "^input_1_0"
input: "^input_2_0"
input: "^input_3_0"
device: "/cpu:0"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "slice_0"
op: "Slice"
input: "y0"
input: "zeros"
input: "ones"
input: "^input_0_0"
input: "^input_1_0"
input: "^input_2_0"
input: "^input_3_0"
device: "/cpu:0"
attr {
key: "Index"
value {
type: DT_INT32
}
}
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
node {
name: "identity_0"
op: "Identity"
input: "slice_0"
input: "^input_0_0"
input: "^input_1_0"
input: "^input_2_0"
input: "^input_3_0"
device: "/cpu:0"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
node {
name: "shape_1"
op: "Shape"
input: "y0:1"
input: "^input_0_0"
input: "^input_1_0"
input: "^input_2_0"
input: "^input_3_0"
device: "/cpu:0"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
attr {
key: "out_type"
value {
type: DT_INT32
}
}
}
node {
name: "zeros_1"
op: "ZerosLike"
input: "shape_1"
input: "^input_0_0"
input: "^input_1_0"
input: "^input_2_0"
input: "^input_3_0"
device: "/cpu:0"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "ones_1"
op: "OnesLike"
input: "shape_1"
input: "^input_0_0"
input: "^input_1_0"
input: "^input_2_0"
input: "^input_3_0"
device: "/cpu:0"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "slice_1"
op: "Slice"
input: "y0:1"
input: "zeros_1"
input: "ones_1"
input: "^input_0_0"
input: "^input_1_0"
input: "^input_2_0"
input: "^input_3_0"
device: "/cpu:0"
attr {
key: "Index"
value {
type: DT_INT32
}
}
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
node {
name: "identity_1"
op: "Identity"
input: "slice_1"
input: "^input_0_0"
input: "^input_1_0"
input: "^input_2_0"
input: "^input_3_0"
device: "/cpu:0"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
library {
function {
signature {
name: "BiasAddx1_Conv2Dx1_DepthwiseConv2dNativex1_Relux1_95"
input_arg {
name: "InceptionV2/Conv2d_1a_7x7/biases/read"
type: DT_FLOAT
}
input_arg {
name: "InceptionV2/Conv2d_1a_7x7/pointwise_weights/read"
type: DT_FLOAT
}
input_arg {
name: "random_uniform"
type: DT_FLOAT
}
input_arg {
name: "InceptionV2/Conv2d_1a_7x7/depthwise_weights/read"
type: DT_FLOAT
}
output_arg {
name: "InceptionV2/InceptionV2/Conv2d_1a_7x7/Relu"
type: DT_FLOAT
}
output_arg {
name: "InceptionV2/InceptionV2/Conv2d_1a_7x7/separable_conv2d/depthwise"
type: DT_FLOAT
}
}
node_def {
name: "InceptionV2/InceptionV2/Conv2d_1a_7x7/BiasAdd"
op: "BiasAdd"
input: "InceptionV2/InceptionV2/Conv2d_1a_7x7/separable_conv2d:output:0"
input: "InceptionV2/Conv2d_1a_7x7/biases/read"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
attr {
key: "data_format"
value {
s: "NHWC"
}
}
}
node_def {
name: "InceptionV2/InceptionV2/Conv2d_1a_7x7/Relu"
op: "Relu"
input: "InceptionV2/InceptionV2/Conv2d_1a_7x7/BiasAdd:output:0"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
node_def {
name: "InceptionV2/InceptionV2/Conv2d_1a_7x7/separable_conv2d"
op: "Conv2D"
input: "InceptionV2/InceptionV2/Conv2d_1a_7x7/separable_conv2d/depthwise:output:0"
input: "InceptionV2/Conv2d_1a_7x7/pointwise_weights/read"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
attr {
key: "data_format"
value {
s: "NHWC"
}
}
attr {
key: "dilations"
value {
list {
i: 1
i: 1
i: 1
i: 1
}
}
}
attr {
key: "padding"
value {
s: "VALID"
}
}
attr {
key: "explicit_paddings"
value {
list {
}
}
}
attr {
key: "strides"
value {
list {
i: 1
i: 1
i: 1
i: 1
}
}
}
attr {
key: "use_cudnn_on_gpu"
value {
b: true
}
}
}
node_def {
name: "InceptionV2/InceptionV2/Conv2d_1a_7x7/separable_conv2d/depthwise"
op: "DepthwiseConv2dNative"
input: "random_uniform"
input: "InceptionV2/Conv2d_1a_7x7/depthwise_weights/read"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
attr {
key: "data_format"
value {
s: "NHWC"
}
}
attr {
key: "dilations"
value {
list {
i: 1
i: 1
i: 1
i: 1
}
}
}
attr {
key: "padding"
value {
s: "SAME"
}
}
attr {
key: "strides"
value {
list {
i: 1
i: 2
i: 2
i: 1
}
}
}
}
ret {
key: "InceptionV2/InceptionV2/Conv2d_1a_7x7/Relu"
value: "InceptionV2/InceptionV2/Conv2d_1a_7x7/Relu:activations:0"
}
ret {
key: "InceptionV2/InceptionV2/Conv2d_1a_7x7/separable_conv2d/depthwise"
value: "InceptionV2/InceptionV2/Conv2d_1a_7x7/separable_conv2d/depthwise:output:0"
}
attr {
key: "_noinline"
value {
b: true
}
}
}
}
versions {
producer: 26
min_consumer: 12
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,822 @@
node {
name: "Const"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 0
}
}
}
}
node {
name: "fifo_queue"
op: "FIFOQueueV2"
attr {
key: "capacity"
value {
i: 1
}
}
attr {
key: "component_types"
value {
list {
type: DT_FLOAT
}
}
}
attr {
key: "container"
value {
s: ""
}
}
attr {
key: "shapes"
value {
list {
}
}
}
attr {
key: "shared_name"
value {
s: ""
}
}
}
node {
name: "ones"
op: "Const"
attr {
key: "dtype"
value {
type: DT_FLOAT
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_FLOAT
tensor_shape {
dim {
size: 1
}
dim {
size: 1
}
dim {
size: 1
}
}
float_val: 1.0
}
}
}
}
node {
name: "while/Enter"
op: "Enter"
input: "Const"
attr {
key: "T"
value {
type: DT_INT32
}
}
attr {
key: "frame_name"
value {
s: "while/while/"
}
}
attr {
key: "is_constant"
value {
b: false
}
}
attr {
key: "parallel_iterations"
value {
i: 10
}
}
}
node {
name: "while/Enter_1"
op: "Enter"
input: "ones"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
attr {
key: "frame_name"
value {
s: "while/while/"
}
}
attr {
key: "is_constant"
value {
b: false
}
}
attr {
key: "parallel_iterations"
value {
i: 10
}
}
}
node {
name: "while/Merge"
op: "Merge"
input: "while/Enter"
input: "while/NextIteration"
attr {
key: "N"
value {
i: 2
}
}
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/Merge_1"
op: "Merge"
input: "while/Enter_1"
input: "while/NextIteration_1"
attr {
key: "N"
value {
i: 2
}
}
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
node {
name: "while/Less/y"
op: "Const"
input: "^while/Merge"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 3
}
}
}
}
node {
name: "while/Less"
op: "Less"
input: "while/Merge"
input: "while/Less/y"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/LoopCond"
op: "LoopCond"
input: "while/Less"
}
node {
name: "while/Switch"
op: "Switch"
input: "while/Merge"
input: "while/LoopCond"
attr {
key: "T"
value {
type: DT_INT32
}
}
attr {
key: "_class"
value {
list {
s: "loc:@while/Merge"
}
}
}
}
node {
name: "while/Switch_1"
op: "Switch"
input: "while/Merge_1"
input: "while/LoopCond"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
attr {
key: "_class"
value {
list {
s: "loc:@while/Merge_1"
}
}
}
}
node {
name: "while/Identity"
op: "Identity"
input: "while/Switch:1"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/Identity_1"
op: "Identity"
input: "while/Switch_1:1"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
node {
name: "while/fifo_queue_enqueue/Enter"
op: "Enter"
input: "fifo_queue"
attr {
key: "T"
value {
type: DT_RESOURCE
}
}
attr {
key: "frame_name"
value {
s: "while/while/"
}
}
attr {
key: "is_constant"
value {
b: true
}
}
attr {
key: "parallel_iterations"
value {
i: 10
}
}
}
node {
name: "while/fifo_queue_enqueue"
op: "QueueEnqueueV2"
input: "while/fifo_queue_enqueue/Enter"
input: "while/Identity_1"
attr {
key: "Tcomponents"
value {
list {
type: DT_FLOAT
}
}
}
attr {
key: "timeout_ms"
value {
i: -1
}
}
}
node {
name: "while/concat/axis"
op: "Const"
input: "^while/Identity"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 2
}
}
}
}
node {
name: "while/concat"
op: "ConcatV2"
input: "while/Identity_1"
input: "while/Identity_1"
input: "while/concat/axis"
attr {
key: "N"
value {
i: 2
}
}
attr {
key: "T"
value {
type: DT_FLOAT
}
}
attr {
key: "Tidx"
value {
type: DT_INT32
}
}
}
node {
name: "while/fifo_queue_Dequeue"
op: "QueueDequeueV2"
input: "while/fifo_queue_enqueue/Enter"
input: "^while/Identity"
attr {
key: "component_types"
value {
list {
type: DT_FLOAT
}
}
}
attr {
key: "timeout_ms"
value {
i: -1
}
}
}
node {
name: "while/while/Const"
op: "Const"
input: "^while/Identity"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 0
}
}
}
}
node {
name: "while/while/Enter"
op: "Enter"
input: "while/while/Const"
attr {
key: "T"
value {
type: DT_INT32
}
}
attr {
key: "frame_name"
value {
s: "while/while/while/while/"
}
}
attr {
key: "is_constant"
value {
b: false
}
}
attr {
key: "parallel_iterations"
value {
i: 10
}
}
}
node {
name: "while/while/Enter_1"
op: "Enter"
input: "while/fifo_queue_Dequeue"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
attr {
key: "frame_name"
value {
s: "while/while/while/while/"
}
}
attr {
key: "is_constant"
value {
b: false
}
}
attr {
key: "parallel_iterations"
value {
i: 10
}
}
}
node {
name: "while/while/Merge"
op: "Merge"
input: "while/while/Enter"
input: "while/while/NextIteration"
attr {
key: "N"
value {
i: 2
}
}
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/while/Merge_1"
op: "Merge"
input: "while/while/Enter_1"
input: "while/while/NextIteration_1"
attr {
key: "N"
value {
i: 2
}
}
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
node {
name: "while/while/Less/y"
op: "Const"
input: "^while/while/Merge"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 3
}
}
}
}
node {
name: "while/while/Less"
op: "Less"
input: "while/while/Merge"
input: "while/while/Less/y"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/while/LoopCond"
op: "LoopCond"
input: "while/while/Less"
}
node {
name: "while/while/Switch"
op: "Switch"
input: "while/while/Merge"
input: "while/while/LoopCond"
attr {
key: "T"
value {
type: DT_INT32
}
}
attr {
key: "_class"
value {
list {
s: "loc:@while/while/Merge"
}
}
}
}
node {
name: "while/while/Switch_1"
op: "Switch"
input: "while/while/Merge_1"
input: "while/while/LoopCond"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
attr {
key: "_class"
value {
list {
s: "loc:@while/while/Merge_1"
}
}
}
}
node {
name: "while/while/Identity"
op: "Identity"
input: "while/while/Switch:1"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/while/Identity_1"
op: "Identity"
input: "while/while/Switch_1:1"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
node {
name: "while/while/add/y"
op: "Const"
input: "^while/while/Identity"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 1
}
}
}
}
node {
name: "while/while/add"
op: "Add"
input: "while/while/Identity"
input: "while/while/add/y"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/while/concat/axis"
op: "Const"
input: "^while/while/Identity"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 0
}
}
}
}
node {
name: "while/while/concat"
op: "ConcatV2"
input: "while/while/Identity_1"
input: "while/while/Identity_1"
input: "while/while/concat/axis"
attr {
key: "N"
value {
i: 2
}
}
attr {
key: "T"
value {
type: DT_FLOAT
}
}
attr {
key: "Tidx"
value {
type: DT_INT32
}
}
}
node {
name: "while/while/NextIteration"
op: "NextIteration"
input: "while/while/add"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/while/NextIteration_1"
op: "NextIteration"
input: "while/while/concat"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
node {
name: "while/while/Exit"
op: "Exit"
input: "while/while/Switch"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/while/Exit_1"
op: "Exit"
input: "while/while/Switch_1"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
node {
name: "while/add/y"
op: "Const"
input: "^while/Identity"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 1
}
}
}
}
node {
name: "while/add"
op: "Add"
input: "while/Identity"
input: "while/add/y"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/NextIteration"
op: "NextIteration"
input: "while/add"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/NextIteration_1"
op: "NextIteration"
input: "while/concat"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
node {
name: "while/Exit"
op: "Exit"
input: "while/Switch"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/Exit_1"
op: "Exit"
input: "while/Switch_1"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
versions {
producer: 21
}
@@ -0,0 +1,762 @@
node {
name: "Const"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 0
}
}
}
}
node {
name: "loop_var"
op: "VarHandleOp"
attr {
key: "container"
value {
s: ""
}
}
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "shape"
value {
shape {
}
}
}
attr {
key: "shared_name"
value {
s: "loop_var"
}
}
}
node {
name: "loop_var/IsInitialized/VarIsInitializedOp"
op: "VarIsInitializedOp"
input: "loop_var"
}
node {
name: "loop_var/Assign"
op: "AssignVariableOp"
input: "loop_var"
input: "Const"
attr {
key: "_class"
value {
list {
s: "loc:@loop_var"
}
}
}
attr {
key: "dtype"
value {
type: DT_INT32
}
}
}
node {
name: "loop_var/Read/ReadVariableOp"
op: "ReadVariableOp"
input: "loop_var"
attr {
key: "_class"
value {
list {
s: "loc:@loop_var"
}
}
}
attr {
key: "dtype"
value {
type: DT_INT32
}
}
}
node {
name: "Const_1"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 1
}
}
}
}
node {
name: "while/ReadVariableOp"
op: "ReadVariableOp"
input: "loop_var"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
}
node {
name: "while/Enter"
op: "Enter"
input: "while/ReadVariableOp"
attr {
key: "T"
value {
type: DT_INT32
}
}
attr {
key: "frame_name"
value {
s: "while/while_context"
}
}
attr {
key: "is_constant"
value {
b: false
}
}
attr {
key: "parallel_iterations"
value {
i: 10
}
}
}
node {
name: "while/Enter_1"
op: "Enter"
input: "Const_1"
attr {
key: "T"
value {
type: DT_INT32
}
}
attr {
key: "frame_name"
value {
s: "while/while_context"
}
}
attr {
key: "is_constant"
value {
b: false
}
}
attr {
key: "parallel_iterations"
value {
i: 10
}
}
}
node {
name: "while/Merge"
op: "Merge"
input: "while/Enter"
input: "while/NextIteration"
attr {
key: "N"
value {
i: 2
}
}
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/Merge_1"
op: "Merge"
input: "while/Enter_1"
input: "while/NextIteration_1"
attr {
key: "N"
value {
i: 2
}
}
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/Less/y"
op: "Const"
input: "^while/Merge"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 3
}
}
}
}
node {
name: "while/Less"
op: "Less"
input: "while/Merge"
input: "while/Less/y"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/LoopCond"
op: "LoopCond"
input: "while/Less"
}
node {
name: "while/Switch"
op: "Switch"
input: "while/Merge"
input: "while/LoopCond"
attr {
key: "T"
value {
type: DT_INT32
}
}
attr {
key: "_class"
value {
list {
s: "loc:@while/Merge"
}
}
}
}
node {
name: "while/Switch_1"
op: "Switch"
input: "while/Merge_1"
input: "while/LoopCond"
attr {
key: "T"
value {
type: DT_INT32
}
}
attr {
key: "_class"
value {
list {
s: "loc:@while/Merge_1"
}
}
}
}
node {
name: "while/Identity"
op: "Identity"
input: "while/Switch:1"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/Identity_1"
op: "Identity"
input: "while/Switch_1:1"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/add"
op: "Add"
input: "while/Identity_1"
input: "while/Identity_1"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/while/Const"
op: "Const"
input: "^while/Identity"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 0
}
}
}
}
node {
name: "while/while/ReadVariableOp/Enter"
op: "Enter"
input: "loop_var"
attr {
key: "T"
value {
type: DT_RESOURCE
}
}
attr {
key: "frame_name"
value {
s: "while/while_context"
}
}
attr {
key: "is_constant"
value {
b: true
}
}
attr {
key: "parallel_iterations"
value {
i: 10
}
}
}
node {
name: "while/while/ReadVariableOp"
op: "ReadVariableOp"
input: "while/while/ReadVariableOp/Enter"
input: "^while/Identity"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
}
node {
name: "while/while/Enter"
op: "Enter"
input: "while/while/Const"
attr {
key: "T"
value {
type: DT_INT32
}
}
attr {
key: "frame_name"
value {
s: "while/while/while_context"
}
}
attr {
key: "is_constant"
value {
b: false
}
}
attr {
key: "parallel_iterations"
value {
i: 10
}
}
}
node {
name: "while/while/Enter_1"
op: "Enter"
input: "while/while/ReadVariableOp"
attr {
key: "T"
value {
type: DT_INT32
}
}
attr {
key: "frame_name"
value {
s: "while/while/while_context"
}
}
attr {
key: "is_constant"
value {
b: false
}
}
attr {
key: "parallel_iterations"
value {
i: 10
}
}
}
node {
name: "while/while/Merge"
op: "Merge"
input: "while/while/Enter"
input: "while/while/NextIteration"
attr {
key: "N"
value {
i: 2
}
}
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/while/Merge_1"
op: "Merge"
input: "while/while/Enter_1"
input: "while/while/NextIteration_1"
attr {
key: "N"
value {
i: 2
}
}
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/while/Less/y"
op: "Const"
input: "^while/while/Merge"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 3
}
}
}
}
node {
name: "while/while/Less"
op: "Less"
input: "while/while/Merge"
input: "while/while/Less/y"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/while/LoopCond"
op: "LoopCond"
input: "while/while/Less"
}
node {
name: "while/while/Switch"
op: "Switch"
input: "while/while/Merge"
input: "while/while/LoopCond"
attr {
key: "T"
value {
type: DT_INT32
}
}
attr {
key: "_class"
value {
list {
s: "loc:@while/while/Merge"
}
}
}
}
node {
name: "while/while/Switch_1"
op: "Switch"
input: "while/while/Merge_1"
input: "while/while/LoopCond"
attr {
key: "T"
value {
type: DT_INT32
}
}
attr {
key: "_class"
value {
list {
s: "loc:@while/while/Merge_1"
}
}
}
}
node {
name: "while/while/Identity"
op: "Identity"
input: "while/while/Switch:1"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/while/Identity_1"
op: "Identity"
input: "while/while/Switch_1:1"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/while/add/y"
op: "Const"
input: "^while/while/Identity"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 1
}
}
}
}
node {
name: "while/while/add"
op: "Add"
input: "while/while/Identity"
input: "while/while/add/y"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/while/add_1"
op: "Add"
input: "while/while/Identity_1"
input: "while/while/Identity_1"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/while/NextIteration"
op: "NextIteration"
input: "while/while/add"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/while/NextIteration_1"
op: "NextIteration"
input: "while/while/add_1"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/while/Exit"
op: "Exit"
input: "while/while/Switch"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/while/Exit_1"
op: "Exit"
input: "while/while/Switch_1"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/add_1/y"
op: "Const"
input: "^while/Identity"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 1
}
}
}
}
node {
name: "while/add_1"
op: "Add"
input: "while/Identity"
input: "while/add_1/y"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/NextIteration"
op: "NextIteration"
input: "while/add_1"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/NextIteration_1"
op: "NextIteration"
input: "while/add"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/Exit"
op: "Exit"
input: "while/Switch"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/Exit_1"
op: "Exit"
input: "while/Switch_1"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
versions {
producer: 24
}
@@ -0,0 +1,330 @@
node {
name: "Const"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 7
}
}
}
}
node {
name: "Const_1"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 5
}
}
}
}
node {
name: "ones"
op: "Const"
attr {
key: "dtype"
value {
type: DT_FLOAT
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_FLOAT
tensor_shape {
dim {
size: 1
}
dim {
size: 1
}
dim {
size: 1
}
}
float_val: 1.0
}
}
}
}
node {
name: "Less"
op: "Less"
input: "Const"
input: "Const_1"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "cond/Switch"
op: "Switch"
input: "Less"
input: "Less"
attr {
key: "T"
value {
type: DT_BOOL
}
}
}
node {
name: "cond/switch_t"
op: "Identity"
input: "cond/Switch:1"
attr {
key: "T"
value {
type: DT_BOOL
}
}
}
node {
name: "cond/switch_f"
op: "Identity"
input: "cond/Switch"
attr {
key: "T"
value {
type: DT_BOOL
}
}
}
node {
name: "cond/pred_id"
op: "Identity"
input: "Less"
attr {
key: "T"
value {
type: DT_BOOL
}
}
}
node {
name: "cond/concat/axis"
op: "Const"
input: "^cond/switch_t"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 0
}
}
}
}
node {
name: "cond/concat/Switch"
op: "Switch"
input: "ones"
input: "cond/pred_id"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
attr {
key: "_class"
value {
list {
s: "loc:@ones"
}
}
}
}
node {
name: "cond/concat"
op: "ConcatV2"
input: "cond/concat/Switch:1"
input: "cond/concat/Switch:1"
input: "cond/concat/axis"
attr {
key: "N"
value {
i: 2
}
}
attr {
key: "T"
value {
type: DT_FLOAT
}
}
attr {
key: "Tidx"
value {
type: DT_INT32
}
}
}
node {
name: "cond/concat_1/axis"
op: "Const"
input: "^cond/switch_f"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 1
}
}
}
}
node {
name: "cond/concat_1/Switch"
op: "Switch"
input: "ones"
input: "cond/pred_id"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
attr {
key: "_class"
value {
list {
s: "loc:@ones"
}
}
}
}
node {
name: "cond/concat_1"
op: "ConcatV2"
input: "cond/concat_1/Switch"
input: "cond/concat_1/Switch"
input: "cond/concat_1/axis"
attr {
key: "N"
value {
i: 2
}
}
attr {
key: "T"
value {
type: DT_FLOAT
}
}
attr {
key: "Tidx"
value {
type: DT_INT32
}
}
}
node {
name: "cond/Merge"
op: "Merge"
input: "cond/concat"
input: "cond/concat_1"
attr {
key: "N"
value {
i: 2
}
}
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
node {
name: "concat/axis"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 2
}
}
}
}
node {
name: "concat"
op: "ConcatV2"
input: "cond/Merge"
input: "cond/Merge"
input: "concat/axis"
attr {
key: "N"
value {
i: 2
}
}
attr {
key: "T"
value {
type: DT_FLOAT
}
}
attr {
key: "Tidx"
value {
type: DT_INT32
}
}
}
versions {
producer: 21
}
@@ -0,0 +1,716 @@
node {
name: "Const"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 0
}
}
}
}
node {
name: "ones"
op: "Const"
attr {
key: "dtype"
value {
type: DT_FLOAT
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_FLOAT
tensor_shape {
dim {
size: 1
}
dim {
size: 1
}
dim {
size: 1
}
}
float_val: 1.0
}
}
}
}
node {
name: "while/Enter"
op: "Enter"
input: "Const"
attr {
key: "T"
value {
type: DT_INT32
}
}
attr {
key: "frame_name"
value {
s: "while/while/"
}
}
attr {
key: "is_constant"
value {
b: false
}
}
attr {
key: "parallel_iterations"
value {
i: 10
}
}
}
node {
name: "while/Enter_1"
op: "Enter"
input: "ones"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
attr {
key: "frame_name"
value {
s: "while/while/"
}
}
attr {
key: "is_constant"
value {
b: false
}
}
attr {
key: "parallel_iterations"
value {
i: 10
}
}
}
node {
name: "while/Merge"
op: "Merge"
input: "while/Enter"
input: "while/NextIteration"
attr {
key: "N"
value {
i: 2
}
}
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/Merge_1"
op: "Merge"
input: "while/Enter_1"
input: "while/NextIteration_1"
attr {
key: "N"
value {
i: 2
}
}
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
node {
name: "while/Less/y"
op: "Const"
input: "^while/Merge"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 3
}
}
}
}
node {
name: "while/Less"
op: "Less"
input: "while/Merge"
input: "while/Less/y"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/LoopCond"
op: "LoopCond"
input: "while/Less"
}
node {
name: "while/Switch"
op: "Switch"
input: "while/Merge"
input: "while/LoopCond"
attr {
key: "T"
value {
type: DT_INT32
}
}
attr {
key: "_class"
value {
list {
s: "loc:@while/Merge"
}
}
}
}
node {
name: "while/Switch_1"
op: "Switch"
input: "while/Merge_1"
input: "while/LoopCond"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
attr {
key: "_class"
value {
list {
s: "loc:@while/Merge_1"
}
}
}
}
node {
name: "while/Identity"
op: "Identity"
input: "while/Switch:1"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/Identity_1"
op: "Identity"
input: "while/Switch_1:1"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
node {
name: "while/while/Const"
op: "Const"
input: "^while/Identity"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 0
}
}
}
}
node {
name: "while/while/Enter"
op: "Enter"
input: "while/while/Const"
attr {
key: "T"
value {
type: DT_INT32
}
}
attr {
key: "frame_name"
value {
s: "while/while/while/while/"
}
}
attr {
key: "is_constant"
value {
b: false
}
}
attr {
key: "parallel_iterations"
value {
i: 10
}
}
}
node {
name: "while/while/Enter_1"
op: "Enter"
input: "while/Identity_1"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
attr {
key: "frame_name"
value {
s: "while/while/while/while/"
}
}
attr {
key: "is_constant"
value {
b: false
}
}
attr {
key: "parallel_iterations"
value {
i: 10
}
}
}
node {
name: "while/while/Merge"
op: "Merge"
input: "while/while/Enter"
input: "while/while/NextIteration"
attr {
key: "N"
value {
i: 2
}
}
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/while/Merge_1"
op: "Merge"
input: "while/while/Enter_1"
input: "while/while/NextIteration_1"
attr {
key: "N"
value {
i: 2
}
}
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
node {
name: "while/while/Less/y"
op: "Const"
input: "^while/while/Merge"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 3
}
}
}
}
node {
name: "while/while/Less"
op: "Less"
input: "while/while/Merge"
input: "while/while/Less/y"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/while/LoopCond"
op: "LoopCond"
input: "while/while/Less"
}
node {
name: "while/while/Switch"
op: "Switch"
input: "while/while/Merge"
input: "while/while/LoopCond"
attr {
key: "T"
value {
type: DT_INT32
}
}
attr {
key: "_class"
value {
list {
s: "loc:@while/while/Merge"
}
}
}
}
node {
name: "while/while/Switch_1"
op: "Switch"
input: "while/while/Merge_1"
input: "while/while/LoopCond"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
attr {
key: "_class"
value {
list {
s: "loc:@while/while/Merge_1"
}
}
}
}
node {
name: "while/while/Identity"
op: "Identity"
input: "while/while/Switch:1"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/while/Identity_1"
op: "Identity"
input: "while/while/Switch_1:1"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
node {
name: "while/while/add/y"
op: "Const"
input: "^while/while/Identity"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 1
}
}
}
}
node {
name: "while/while/add"
op: "Add"
input: "while/while/Identity"
input: "while/while/add/y"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/while/concat/axis"
op: "Const"
input: "^while/while/Identity"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 2
}
}
}
}
node {
name: "while/while/concat"
op: "ConcatV2"
input: "while/while/Identity_1"
input: "while/while/Identity_1"
input: "while/while/concat/axis"
attr {
key: "N"
value {
i: 2
}
}
attr {
key: "T"
value {
type: DT_FLOAT
}
}
attr {
key: "Tidx"
value {
type: DT_INT32
}
}
}
node {
name: "while/while/NextIteration"
op: "NextIteration"
input: "while/while/add"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/while/NextIteration_1"
op: "NextIteration"
input: "while/while/concat"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
node {
name: "while/while/Exit"
op: "Exit"
input: "while/while/Switch"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/while/Exit_1"
op: "Exit"
input: "while/while/Switch_1"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
node {
name: "while/add/y"
op: "Const"
input: "^while/Identity"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 1
}
}
}
}
node {
name: "while/add"
op: "Add"
input: "while/Identity"
input: "while/add/y"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/concat/axis"
op: "Const"
input: "^while/Identity"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 0
}
}
}
}
node {
name: "while/concat"
op: "ConcatV2"
input: "while/Identity_1"
input: "while/Identity_1"
input: "while/concat/axis"
attr {
key: "N"
value {
i: 2
}
}
attr {
key: "T"
value {
type: DT_FLOAT
}
}
attr {
key: "Tidx"
value {
type: DT_INT32
}
}
}
node {
name: "while/NextIteration"
op: "NextIteration"
input: "while/add"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/NextIteration_1"
op: "NextIteration"
input: "while/concat"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
node {
name: "while/Exit"
op: "Exit"
input: "while/Switch"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/Exit_1"
op: "Exit"
input: "while/Switch_1"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
versions {
producer: 21
}
@@ -0,0 +1,569 @@
node {
name: "Const"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 0
}
}
}
}
node {
name: "fifo_queue"
op: "FIFOQueueV2"
attr {
key: "capacity"
value {
i: 1
}
}
attr {
key: "component_types"
value {
list {
type: DT_FLOAT
}
}
}
attr {
key: "container"
value {
s: ""
}
}
attr {
key: "shapes"
value {
list {
}
}
}
attr {
key: "shared_name"
value {
s: ""
}
}
}
node {
name: "ones"
op: "Const"
attr {
key: "dtype"
value {
type: DT_FLOAT
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_FLOAT
tensor_shape {
dim {
size: 2
}
dim {
size: 2
}
}
float_val: 1.0
}
}
}
}
node {
name: "fifo_queue_enqueue"
op: "QueueEnqueueV2"
input: "fifo_queue"
input: "ones"
attr {
key: "Tcomponents"
value {
list {
type: DT_FLOAT
}
}
}
attr {
key: "timeout_ms"
value {
i: -1
}
}
}
node {
name: "fifo_queue_1"
op: "FIFOQueueV2"
attr {
key: "capacity"
value {
i: 1
}
}
attr {
key: "component_types"
value {
list {
type: DT_FLOAT
}
}
}
attr {
key: "container"
value {
s: ""
}
}
attr {
key: "shapes"
value {
list {
}
}
}
attr {
key: "shared_name"
value {
s: ""
}
}
}
node {
name: "fifo_queue_Dequeue"
op: "QueueDequeueV2"
input: "fifo_queue"
attr {
key: "component_types"
value {
list {
type: DT_FLOAT
}
}
}
attr {
key: "timeout_ms"
value {
i: -1
}
}
}
node {
name: "while/Enter"
op: "Enter"
input: "Const"
attr {
key: "T"
value {
type: DT_INT32
}
}
attr {
key: "frame_name"
value {
s: "while/while/"
}
}
attr {
key: "is_constant"
value {
b: false
}
}
attr {
key: "parallel_iterations"
value {
i: 10
}
}
}
node {
name: "while/Enter_1"
op: "Enter"
input: "fifo_queue_Dequeue"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
attr {
key: "frame_name"
value {
s: "while/while/"
}
}
attr {
key: "is_constant"
value {
b: false
}
}
attr {
key: "parallel_iterations"
value {
i: 10
}
}
}
node {
name: "while/Merge"
op: "Merge"
input: "while/Enter"
input: "while/NextIteration"
attr {
key: "N"
value {
i: 2
}
}
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/Merge_1"
op: "Merge"
input: "while/Enter_1"
input: "while/NextIteration_1"
attr {
key: "N"
value {
i: 2
}
}
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
node {
name: "while/Less/y"
op: "Const"
input: "^while/Merge"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 10
}
}
}
}
node {
name: "while/Less"
op: "Less"
input: "while/Merge"
input: "while/Less/y"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/LoopCond"
op: "LoopCond"
input: "while/Less"
}
node {
name: "while/Switch"
op: "Switch"
input: "while/Merge"
input: "while/LoopCond"
attr {
key: "T"
value {
type: DT_INT32
}
}
attr {
key: "_class"
value {
list {
s: "loc:@while/Merge"
}
}
}
}
node {
name: "while/Switch_1"
op: "Switch"
input: "while/Merge_1"
input: "while/LoopCond"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
attr {
key: "_class"
value {
list {
s: "loc:@while/Merge_1"
}
}
}
}
node {
name: "while/Identity"
op: "Identity"
input: "while/Switch:1"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/Identity_1"
op: "Identity"
input: "while/Switch_1:1"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
node {
name: "while/add/y"
op: "Const"
input: "^while/Identity"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 1
}
}
}
}
node {
name: "while/add"
op: "Add"
input: "while/Identity"
input: "while/add/y"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/concat/axis"
op: "Const"
input: "^while/Identity"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 0
}
}
}
}
node {
name: "while/concat"
op: "ConcatV2"
input: "while/Identity_1"
input: "while/Identity_1"
input: "while/concat/axis"
attr {
key: "N"
value {
i: 2
}
}
attr {
key: "T"
value {
type: DT_FLOAT
}
}
attr {
key: "Tidx"
value {
type: DT_INT32
}
}
}
node {
name: "while/NextIteration"
op: "NextIteration"
input: "while/add"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/NextIteration_1"
op: "NextIteration"
input: "while/concat"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
node {
name: "while/Exit"
op: "Exit"
input: "while/Switch"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/Exit_1"
op: "Exit"
input: "while/Switch_1"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
node {
name: "fifo_queue_1_enqueue"
op: "QueueEnqueueV2"
input: "fifo_queue_1"
input: "while/Exit_1"
attr {
key: "Tcomponents"
value {
list {
type: DT_FLOAT
}
}
}
attr {
key: "timeout_ms"
value {
i: -1
}
}
}
node {
name: "fifo_queue_1_Dequeue"
op: "QueueDequeueV2"
input: "fifo_queue_1"
attr {
key: "component_types"
value {
list {
type: DT_FLOAT
}
}
}
attr {
key: "timeout_ms"
value {
i: -1
}
}
}
node {
name: "concat/axis"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 1
}
}
}
}
node {
name: "concat"
op: "ConcatV2"
input: "fifo_queue_1_Dequeue"
input: "fifo_queue_1_Dequeue"
input: "concat/axis"
attr {
key: "N"
value {
i: 2
}
}
attr {
key: "T"
value {
type: DT_FLOAT
}
}
attr {
key: "Tidx"
value {
type: DT_INT32
}
}
}
versions {
producer: 21
}
@@ -0,0 +1,111 @@
node {
name: "Const"
op: "Const"
attr {
key: "dtype"
value {
type: DT_FLOAT
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_FLOAT
tensor_shape {
dim {
size: 1
}
dim {
size: 2
}
}
float_val: 2.0
}
}
}
}
node {
name: "Const_1"
op: "Const"
attr {
key: "dtype"
value {
type: DT_FLOAT
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_FLOAT
tensor_shape {
dim {
size: 1
}
dim {
size: 2
}
}
float_val: 2.0
}
}
}
}
node {
name: "MyAdd_55e046a8"
op: "MyAdd_55e046a8"
input: "Const"
input: "Const_1"
}
node {
name: "MyAdd_55e046a8_1"
op: "MyAdd_55e046a8"
input: "Const"
input: "MyAdd_55e046a8"
}
library {
function {
signature {
name: "MyAdd_55e046a8"
input_arg {
name: "x"
type: DT_FLOAT
}
input_arg {
name: "y"
type: DT_FLOAT
}
output_arg {
name: "Add"
type: DT_FLOAT
}
}
node_def {
name: "Add"
op: "Add"
input: "x"
input: "y"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
ret {
key: "Add"
value: "Add:z:0"
}
attr {
key: "_noinline"
value {
b: true
}
}
}
}
versions {
producer: 24
min_consumer: 12
}
@@ -0,0 +1,367 @@
node {
name: "Const"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 0
}
}
}
}
node {
name: "ones"
op: "PlaceholderV2"
attr {
key: "dtype"
value {
type: DT_FLOAT
}
}
attr {
key: "shape"
value {
shape {
dim {
size: -1
}
dim {
size: 2
}
}
}
}
}
node {
name: "while/Enter"
op: "Enter"
input: "Const"
attr {
key: "T"
value {
type: DT_INT32
}
}
attr {
key: "frame_name"
value {
s: "while/while/"
}
}
attr {
key: "is_constant"
value {
b: false
}
}
attr {
key: "parallel_iterations"
value {
i: 10
}
}
}
node {
name: "while/Enter_1"
op: "Enter"
input: "ones"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
attr {
key: "frame_name"
value {
s: "while/while/"
}
}
attr {
key: "is_constant"
value {
b: false
}
}
attr {
key: "parallel_iterations"
value {
i: 10
}
}
}
node {
name: "while/Merge"
op: "Merge"
input: "while/Enter"
input: "while/NextIteration"
attr {
key: "N"
value {
i: 2
}
}
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/Merge_1"
op: "Merge"
input: "while/Enter_1"
input: "while/NextIteration_1"
attr {
key: "N"
value {
i: 2
}
}
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
node {
name: "while/Less/y"
op: "Const"
input: "^while/Merge"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 10
}
}
}
}
node {
name: "while/Less"
op: "Less"
input: "while/Merge"
input: "while/Less/y"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/LoopCond"
op: "LoopCond"
input: "while/Less"
}
node {
name: "while/Switch"
op: "Switch"
input: "while/Merge"
input: "while/LoopCond"
attr {
key: "T"
value {
type: DT_INT32
}
}
attr {
key: "_class"
value {
list {
s: "loc:@while/Merge"
}
}
}
}
node {
name: "while/Switch_1"
op: "Switch"
input: "while/Merge_1"
input: "while/LoopCond"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
attr {
key: "_class"
value {
list {
s: "loc:@while/Merge_1"
}
}
}
}
node {
name: "while/Identity"
op: "Identity"
input: "while/Switch:1"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/Identity_1"
op: "Identity"
input: "while/Switch_1:1"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
node {
name: "while/add/y"
op: "Const"
input: "^while/Identity"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 1
}
}
}
}
node {
name: "while/add"
op: "Add"
input: "while/Identity"
input: "while/add/y"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/concat/axis"
op: "Const"
input: "^while/Identity"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 0
}
}
}
}
node {
name: "while/concat"
op: "ConcatV2"
input: "while/Identity_1"
input: "while/Identity_1"
input: "while/concat/axis"
attr {
key: "N"
value {
i: 2
}
}
attr {
key: "T"
value {
type: DT_FLOAT
}
}
attr {
key: "Tidx"
value {
type: DT_INT32
}
}
}
node {
name: "while/NextIteration"
op: "NextIteration"
input: "while/add"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/NextIteration_1"
op: "NextIteration"
input: "while/concat"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
node {
name: "while/Exit"
op: "Exit"
input: "while/Switch"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "while/Exit_1"
op: "Exit"
input: "while/Switch_1"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
versions {
producer: 21
}
@@ -0,0 +1,291 @@
node {
name: "Const"
op: "Const"
attr {
key: "dtype"
value { type: DT_INT32 }
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {}
int_val: 0
}
}
}
}
node {
name: "loop_var"
op: "VarHandleOp"
attr {
key: "_class"
value { list { s: "loc:@loop_var" } }
}
attr {
key: "container"
value { s: "" }
}
attr {
key: "dtype"
value { type: DT_INT32 }
}
attr {
key: "shape"
value { shape {} }
}
attr {
key: "shared_name"
value { s: "loop_var" }
}
}
node {
name: "Const_1"
op: "Const"
attr {
key: "dtype"
value { type: DT_INT32 }
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {}
int_val: 1
}
}
}
}
node {
name: "while/Enter"
op: "Enter"
input: "loop_var"
attr {
key: "T"
value { type: DT_RESOURCE }
}
attr {
key: "frame_name"
value { s: "while/while_context" }
}
attr {
key: "is_constant"
value { b: false }
}
attr {
key: "parallel_iterations"
value { i: 10 }
}
}
node {
name: "while/Enter_1"
op: "Enter"
input: "Const_1"
attr {
key: "T"
value { type: DT_INT32 }
}
attr {
key: "frame_name"
value { s: "while/while_context" }
}
attr {
key: "is_constant"
value { b: false }
}
attr {
key: "parallel_iterations"
value { i: 10 }
}
}
node {
name: "while/Merge"
op: "Merge"
input: "while/Enter"
input: "while/NextIteration"
attr {
key: "N"
value { i: 2 }
}
attr {
key: "T"
value { type: DT_RESOURCE }
}
}
node {
name: "while/Merge_1"
op: "Merge"
input: "while/Enter_1"
input: "while/NextIteration_1"
attr {
key: "N"
value { i: 2 }
}
attr {
key: "T"
value { type: DT_INT32 }
}
}
node {
name: "while/Less/y"
op: "Const"
input: "^while/Merge"
attr {
key: "dtype"
value { type: DT_INT32 }
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {}
int_val: 3
}
}
}
}
node {
name: "while/Less"
op: "Less"
input: "while/Merge"
input: "while/Less/y"
attr {
key: "T"
value { type: DT_INT32 }
}
}
node { name: "while/LoopCond" op: "LoopCond" input: "while/Less" }
node {
name: "while/Switch"
op: "Switch"
input: "while/Merge"
input: "while/LoopCond"
attr {
key: "T"
value { type: DT_RESOURCE }
}
attr {
key: "_class"
value { list { s: "loc:@while/Merge" } }
}
}
node {
name: "while/Switch_1"
op: "Switch"
input: "while/Merge_1"
input: "while/LoopCond"
attr {
key: "T"
value { type: DT_INT32 }
}
attr {
key: "_class"
value { list { s: "loc:@while/Merge_1" } }
}
}
node {
name: "while/Identity"
op: "Identity"
input: "while/Switch:1"
attr {
key: "T"
value { type: DT_RESOURCE }
}
}
node {
name: "while/ReadVariableOp"
op: "ReadVariableOp"
input: "while/Identity"
attr {
key: "dtype"
value { type: DT_INT32 }
}
}
node {
name: "while/Identity_1"
op: "Identity"
input: "while/Switch_1:1"
attr {
key: "T"
value { type: DT_INT32 }
}
}
node {
name: "while/add/y"
op: "Const"
input: "^while/ReadVariableOp"
attr {
key: "dtype"
value { type: DT_INT32 }
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {}
int_val: 1
}
}
}
}
node {
name: "while/add"
op: "Add"
input: "while/ReadVariableOp"
input: "while/add/y"
attr {
key: "T"
value { type: DT_INT32 }
}
}
node {
name: "while/add_1"
op: "Add"
input: "while/Identity_1"
input: "while/Identity_1"
attr {
key: "T"
value { type: DT_INT32 }
}
}
node {
name: "while/NextIteration"
op: "NextIteration"
input: "while/Identity"
attr {
key: "T"
value { type: DT_RESOURCE }
}
}
node {
name: "while/NextIteration_1"
op: "NextIteration"
input: "while/add_1"
attr {
key: "T"
value { type: DT_INT32 }
}
}
node {
name: "while/Exit"
op: "Exit"
input: "while/Switch"
attr {
key: "T"
value { type: DT_RESOURCE }
}
}
node {
name: "while/Exit_1"
op: "Exit"
input: "while/Switch_1"
attr {
key: "T"
value { type: DT_INT32 }
}
}
versions { producer: 27 }
@@ -0,0 +1,150 @@
/* 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/grappler/costs/measuring_cost_estimator.h"
#include <limits>
#include "tensorflow/core/framework/cost_graph.pb.h"
#include "tensorflow/core/framework/step_stats.pb.h"
#include "tensorflow/core/grappler/clusters/cluster.h"
#include "tensorflow/core/grappler/costs/robust_stats.h"
#include "tensorflow/core/grappler/grappler_item.h"
#include "tensorflow/core/kernels/ops_util.h"
#include "tensorflow/core/platform/blocking_counter.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/public/session.h"
namespace tensorflow {
namespace grappler {
MeasuringCostEstimator::MeasuringCostEstimator(Cluster* cluster,
int measurement_steps,
int measurement_threads)
: measurement_steps_(measurement_steps),
measurement_threads_(measurement_threads) {
CHECK_GE(measurement_steps, 1);
if (measurement_threads > 0) {
thread_pool_.reset(new thread::ThreadPool(
Env::Default(), SanitizeThreadSuffix("measurements"),
measurement_threads));
}
cluster_ = cluster;
}
absl::Status MeasuringCostEstimator::Initialize(const GrapplerItem& item) {
feed_ = item.feed;
fetch_ = item.fetch;
return cluster_->Initialize(item);
}
absl::Status MeasuringCostEstimator::PredictCosts(
const GraphDef& optimized_graph, RunMetadata* run_metadata,
Costs* costs) const {
CostGraphDef* cost_graph = nullptr;
if (run_metadata) {
cost_graph = run_metadata->mutable_cost_graph();
}
const bool running_simulation = (cluster_->type() == "virtual");
std::vector<double> times(measurement_steps_);
BlockingCounter barrier(measurement_steps_);
mutex status_mu;
absl::Status status;
auto measurement_fn = [&](const int step) {
const Costs::MicroSeconds start = Env::Default()->NowMicros();
RunMetadata metadata;
const absl::Status local_status =
cluster_->Run(optimized_graph, feed_, fetch_, &metadata);
{
mutex_lock lock(status_mu);
status.Update(local_status);
}
if (step < 0) {
// Discard the first iteration as it triggers the warmup, and therefore
// takes much longer than a normal step.
return;
}
if (!local_status.ok()) {
// Discard the data if the run wasn't successful.
barrier.DecrementCount();
return;
}
const Costs::MicroSeconds finish = Env::Default()->NowMicros();
if (running_simulation) {
// When running simulation, return the estimated runtime, not the time it
// takes to run the simulation.
double time = 0.0;
for (const DeviceStepStats& stepstats :
metadata.step_stats().dev_stats()) {
for (const NodeExecStats& node_stats : stepstats.node_stats()) {
const double completion_time =
node_stats.all_end_rel_micros() + node_stats.all_start_micros();
time = std::max(time, completion_time * 1e3);
}
}
times[step] = time;
} else {
const double time = (finish - start).count() * 1e3;
times[step] = time;
}
if (cost_graph && (step + 1 == measurement_steps_)) {
metadata.mutable_cost_graph()->Swap(cost_graph);
}
barrier.DecrementCount();
};
// Initialize the computation and warm up TensorFlow.
measurement_fn(-1);
if (!status.ok()) {
LOG(ERROR) << "Failed to run start measurements: " << status.message();
costs->execution_time = Costs::Duration::max();
return status;
}
// Run "measurement_steps_" and measure the time.
VLOG(1) << "Number of measurement steps: " << measurement_steps_;
if (measurement_threads_ > 0) {
for (int i = 0; i < measurement_steps_; ++i) {
thread_pool_->Schedule([i, &measurement_fn]() { measurement_fn(i); });
}
barrier.Wait();
} else {
for (int i = 0; i < measurement_steps_ && status.ok(); ++i) {
measurement_fn(i);
}
}
if (!status.ok()) {
LOG(ERROR) << "Failed to measure graph performance: " << status.message();
costs->execution_time = Costs::Duration::max();
return status;
}
// Compute the average time of the measure steps. Use Huber statistics
// to filter out outliers.
RobustStats stats(times);
costs->execution_time = Costs::Duration(stats.mean());
return absl::OkStatus();
}
} // end namespace grappler
} // end namespace tensorflow
@@ -0,0 +1,77 @@
/* 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_GRAPPLER_COSTS_MEASURING_COST_ESTIMATOR_H_
#define TENSORFLOW_CORE_GRAPPLER_COSTS_MEASURING_COST_ESTIMATOR_H_
#include <string>
#include <utility>
#include <vector>
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/grappler/costs/cost_estimator.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/core/threadpool.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
class CostGraphDef;
class GraphDef;
} // namespace tensorflow
namespace tensorflow {
namespace grappler {
class Cluster;
struct GrapplerItem;
// Estimate the cost of running a Grappler item by actually running the
// corresponding TensorFlow graph on the specified cluster and measuring the
// runtimes.
class MeasuringCostEstimator : public CostEstimator {
public:
// Run the model for measurement_steps to measure its average cost.
// When measurement_threads is greater than 0, use a threadpool of as many
// threads to run the measurements; otherwise, run them serially. Does not
// take ownership of cluster.
explicit MeasuringCostEstimator(Cluster* cluster, int measurement_steps,
int measurement_threads);
~MeasuringCostEstimator() override {}
// Initializes the estimator for the specified grappler item.
// This implementation always returns OK.
absl::Status Initialize(const GrapplerItem& item) override;
// Runs the optimized version of the graph on the cluster, measures
// the runtimes of each operation, and annotates the CostGraphDef of
// RunMetadata with the corresponding measurements.
// Returns the average latency for the whole graph.
absl::Status PredictCosts(const GraphDef& optimized_graph,
RunMetadata* run_metadata,
Costs* cost) const override;
private:
Cluster* cluster_; // Not owned.
int measurement_steps_;
int measurement_threads_;
std::vector<std::pair<std::string, Tensor>> feed_;
std::vector<std::string> fetch_;
std::unique_ptr<thread::ThreadPool> thread_pool_;
};
} // end namespace grappler
} // end namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPPLER_COSTS_MEASURING_COST_ESTIMATOR_H_
@@ -0,0 +1,44 @@
/* 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_GRAPPLER_COSTS_OP_CONTEXT_H_
#define TENSORFLOW_CORE_GRAPPLER_COSTS_OP_CONTEXT_H_
#include "absl/container/flat_hash_map.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/grappler/costs/op_performance_data.pb.h"
namespace tensorflow {
namespace grappler {
// A structure to keep the context of op execution, including its shape,
// execution context, and other relevant information.
struct OpContext {
std::string name;
std::string device_name;
OpInfo op_info;
const FunctionDefLibrary* function_library; // Not owned.
// This map is used to stash meta attributes so that they may be
// communicated, for instance, from the scheduler that creates them to the
// CostEstimator or EventCostManager that uses them.
absl::flat_hash_map<std::string, absl::variant<int64_t, std::string>>
op_meta_attributes;
OpContext() { function_library = nullptr; }
};
} // end namespace grappler
} // end namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPPLER_COSTS_OP_CONTEXT_H_
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,346 @@
/* 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_GRAPPLER_COSTS_OP_LEVEL_COST_ESTIMATOR_H_
#define TENSORFLOW_CORE_GRAPPLER_COSTS_OP_LEVEL_COST_ESTIMATOR_H_
#include <cstdint>
#include <functional>
#include <map>
#include <numeric>
#include <set>
#include <vector>
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "tensorflow/core/grappler/costs/cost_estimator.h"
#include "tensorflow/core/grappler/costs/op_context.h"
#include "tensorflow/core/grappler/costs/op_performance_data.pb.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/util/padding.h"
namespace tensorflow {
namespace grappler {
bool GetTensorShapeProtoFromTensorProto(const TensorProto& tensor_proto,
TensorShapeProto* tensor_shape_proto);
std::vector<int64_t> MaybeGetMinimumShape(
const TensorShapeProto& original_shape, int rank,
bool* found_unknown_shapes);
// Node costs; an intermediate structure used within op level cost estimator.
struct NodeCosts {
// If this FLAG is true, override calculated compute time with a minimum
// value, instead of calculating it from num_compute_ops and compute ops/sec.
// For example, PredictIdentity, PredictVariable, PredictMetadata set this
// FLAG.
bool minimum_cost_op = false;
// Compute ops.
int64_t num_compute_ops = 0;
// Memory bytes accessed; note that these may be different to the size of
// tensors.
std::vector<int64_t> num_input_bytes_accessed; // ordered by input tensors.
std::vector<int64_t> num_output_bytes_accessed; // ordered by output ports.
int64_t internal_read_bytes = 0;
int64_t internal_write_bytes = 0;
// Convenience functions.
int64_t num_total_input_bytes() const {
return std::accumulate(num_input_bytes_accessed.begin(),
num_input_bytes_accessed.end(), 0LL);
}
int64_t num_total_read_bytes() const {
return num_total_input_bytes() + internal_read_bytes;
}
int64_t num_total_output_bytes() const {
return std::accumulate(num_output_bytes_accessed.begin(),
num_output_bytes_accessed.end(), 0LL);
}
int64_t num_total_write_bytes() const {
return num_total_output_bytes() + internal_write_bytes;
}
int64_t num_bytes_accessed() const {
return num_total_read_bytes() + num_total_write_bytes();
}
// Memory usage.
int64_t max_memory = 0;
int64_t persistent_memory = 0;
int64_t temporary_memory = 0;
// Stats.
int64_t num_nodes = 1;
int64_t num_nodes_with_unknown_shapes = 0;
int64_t num_nodes_with_unknown_op_type = 0;
int64_t num_nodes_with_pure_memory_op = 0;
bool inaccurate = false;
// TODO(dyoon): this is added for compatibility; some old code is hard to
// migrate; hence, using these as a backup. Once we clean up, we'll delete
// these fields. New code should not use these.
bool has_costs = false;
Costs costs;
};
class OpLevelCostEstimator {
public:
OpLevelCostEstimator();
virtual ~OpLevelCostEstimator() {}
virtual Costs PredictCosts(const OpContext& op_context) const;
// Returns basic device performance info.
virtual DeviceInfo GetDeviceInfo(const DeviceProperties& device) const;
protected:
// TODO(dyoon): Consider to remove PredictOpCountBasedCosts() with OpInfo.
// Naive cost estimate based on the given operations count and total
// input/output tensor sizes of the given op_info combined.
Costs PredictOpCountBasedCost(double operations, const OpInfo& op_info) const;
// Naive cost estimate based on the given operations count and the given total
// io size in bytes. Sizes of op_info inputs and outputs are not taken into
// consideration.
Costs PredictOpCountBasedCost(double operations, double input_io_bytes,
double output_io_bytes,
const OpInfo& op_info) const;
// Top-level method cost function (PredictCosts calls this method to get
// NodeCosts, and then converts it to Costs). PredictNodeCosts() calls other
// Predict methods depending on op types.
absl::Status PredictNodeCosts(const OpContext& op_context,
NodeCosts* node_costs) const;
// Predict cost of an op for which no accurate estimator is defined.
absl::Status PredictCostOfAnUnknownOp(const OpContext& op_context,
NodeCosts* node_costs) const;
// This family of routines predicts the costs to
// perform the specified TensorFlow Op on the
// device represented by a subclass. The default
// implementation just divides the operations to
// perform the op (from the "Count" routines,
// above) by the device peak operations per
// second.
// Implementation of costs other than
// execution_time is optional, depending on the
// device.
absl::Status PredictNaryOp(const OpContext& op_context,
NodeCosts* node_costs) const;
absl::Status PredictConv2D(const OpContext& op_context,
NodeCosts* node_costs) const;
absl::Status PredictCwiseOp(const OpContext& op_context,
NodeCosts* node_costs) const;
absl::Status PredictConv2DBackpropInput(const OpContext& op_context,
NodeCosts* node_costs) const;
absl::Status PredictConv2DBackpropFilter(const OpContext& op_context,
NodeCosts* node_costs) const;
absl::Status PredictFusedConv2DBiasActivation(const OpContext& op_context,
NodeCosts* node_costs) const;
absl::Status PredictMatMul(const OpContext& op_context,
NodeCosts* node_costs) const;
absl::Status PredictSparseTensorDenseMatMul(const OpContext& op_context,
NodeCosts* node_costs) const;
absl::Status PredictNoOp(const OpContext& op_context,
NodeCosts* node_costs) const;
absl::Status PredictIdentity(const OpContext& op_context,
NodeCosts* node_costs) const;
absl::Status PredictVariable(const OpContext& op_context,
NodeCosts* node_costs) const;
absl::Status PredictBatchMatMul(const OpContext& op_context,
NodeCosts* node_costs) const;
absl::Status PredictMetadata(const OpContext& op_context,
NodeCosts* node_costs) const;
absl::Status PredictGatherOrSlice(const OpContext& op_context,
NodeCosts* node_costs) const;
absl::Status PredictScatter(const OpContext& op_context,
NodeCosts* node_costs) const;
absl::Status PredictMaxPool(const OpContext& op_context,
NodeCosts* node_costs) const;
absl::Status PredictMaxPoolGrad(const OpContext& op_context,
NodeCosts* node_costs) const;
absl::Status PredictAvgPool(const OpContext& op_context,
NodeCosts* node_costs) const;
absl::Status PredictAvgPoolGrad(const OpContext& op_context,
NodeCosts* node_costs) const;
absl::Status PredictFusedBatchNorm(const OpContext& op_context,
NodeCosts* node_costs) const;
absl::Status PredictFusedBatchNormGrad(const OpContext& op_context,
NodeCosts* node_costs) const;
absl::Status PredictEinsum(const OpContext& op_context,
NodeCosts* node_costs) const;
absl::Status PredictAssignVariableOps(const OpContext& op_context,
NodeCosts* node_costs) const;
absl::Status PredictPureMemoryOp(const OpContext& op_context,
NodeCosts* node_costs) const;
absl::Status PredictSoftmax(const OpContext& op_context,
NodeCosts* node_costs) const;
absl::Status PredictResizeBilinear(const OpContext& op_context,
NodeCosts* node_costs) const;
absl::Status PredictCropAndResize(const OpContext& op_context,
NodeCosts* node_costs) const;
int64_t GetSoftmaxComputeOps(const OpContext& op_context) const;
// Generic cost prediction method for fused operations.
absl::Status PredictFusedOp(const OpContext& op_context,
const std::vector<OpContext>& fused_op_contexts,
NodeCosts* node_costs) const;
// Utility function for safe division. Returns 0
// if rhs is 0 or negative.
static double SafeDiv(const double lhs, const double rhs) {
if (rhs > 0) {
return lhs / rhs;
} else {
return 0.0;
}
}
// This family of routines counts the number of operations to perform the
// specified TensorFlow Op.
struct MatMulDimensions {
int m;
int n;
int k;
};
struct BatchMatMulDimensions {
std::vector<int> batch_dims;
MatMulDimensions matmul_dims;
};
struct ConvolutionDimensions {
int64_t batch; // Batch size.
int64_t ix; // Input size x.
int64_t iy; // Input size y.
int64_t iz; // Input depth.
int64_t kx; // Kernel x.
int64_t ky; // Kernel y.
int64_t kz; // Kernel depth (in case of group convolution, this will be
// smaller than input depth).
int64_t oz; // Output depth.
int64_t ox; // Output size x.
int64_t oy; // Output size y.
int64_t sx; // Stride x.
int64_t sy; // Stride y.
Padding padding; // SAME or VALID.
};
static int64_t CountConv2DOperations(const OpInfo& op_info,
bool* found_unknown_shapes);
static int64_t CountConv2DOperations(const OpInfo& op_info,
ConvolutionDimensions* conv_info,
bool* found_unknown_shapes);
static int64_t CountMatMulOperations(const OpInfo& op_info,
bool* found_unknown_shapes);
static int64_t CountMatMulOperations(const OpInfo& op_info,
MatMulDimensions* mat_mul,
bool* found_unknown_shapes);
static int64_t CountMatMulOperations(const OpInfo& op_info, bool transpose_a,
bool transpose_b,
MatMulDimensions* mat_mul,
bool* found_unknown_shapes);
bool GenerateBatchMatmulContextFromEinsum(const OpContext& einsum_context,
OpContext* batch_matmul_context,
bool* found_unknown_shapes) const;
static int64_t CountBatchMatMulOperations(const OpInfo& op_info,
bool* found_unknown_shapes);
static int64_t CountBatchMatMulOperations(
const OpInfo& op_info, BatchMatMulDimensions* batch_mat_mul,
bool* found_unknown_shapes);
static int64_t CountConv2DBackpropInputOperations(
const OpInfo& op_info, ConvolutionDimensions* returned_conv_dims,
bool* found_unknown_shapes);
static int64_t CountConv2DBackpropFilterOperations(
const OpInfo& op_info, ConvolutionDimensions* returned_conv_dims,
bool* found_unknown_shapes);
// Calculate the element count of an input/output tensor.
static int64_t CalculateTensorElementCount(
const OpInfo::TensorProperties& tensor, bool* found_unknown_shapes);
// Calculate the total size in bytes of an input/output tensor.
static int64_t CalculateTensorSize(const OpInfo::TensorProperties& tensor,
bool* found_unknown_shapes);
// Calculate the element count of the largest
// input of specified TensorFlow op.
static int64_t CalculateLargestInputCount(const OpInfo& op_info,
bool* found_unknown_shapes);
// Calculate the total size in bytes of the all
// the inputs of specified TensorFlow op.
static int64_t CalculateInputSize(const OpInfo& op_info,
bool* found_unknown_shapes);
// Same, but a vector format: one for each input.
static std::vector<int64_t> CalculateInputTensorSize(
const OpInfo& op_info, bool* found_unknown_shapes);
// Calculate the total size in bytes of the all
// the outputs of specified TensorFlow op.
static int64_t CalculateOutputSize(const OpInfo& op_info,
bool* found_unknown_shapes);
// Same, but a vector format: one for each output.
static std::vector<int64_t> CalculateOutputTensorSize(
const OpInfo& op_info, bool* found_unknown_shapes);
// For convolution and its grad ops.
static ConvolutionDimensions ConvolutionDimensionsFromInputs(
const TensorShapeProto& original_image_shape,
const TensorShapeProto& original_filter_shape, const OpInfo& op_info,
bool* found_unknown_shapes);
// For Pooling, FusedBatchNorm, and their grad ops.
static absl::StatusOr<ConvolutionDimensions> OpDimensionsFromInputs(
const TensorShapeProto& original_image_shape, const OpInfo& op_info,
bool* found_unknown_shapes);
// Helper to construct child operation contexts for the component operations
// of fused ops.
static OpContext FusedChildContext(
const OpContext& parent, const std::string& op_name,
const OpInfo::TensorProperties& output,
const std::vector<OpInfo::TensorProperties>& inputs);
// Helper to construct tensor shapes.
static OpInfo::TensorProperties DescribeTensor(
DataType type, const std::vector<int64_t>& dims);
// Helper method for building common case NodeCosts.
static absl::Status PredictDefaultNodeCosts(int64_t num_compute_ops,
const OpContext& op_context,
bool* found_unknown_shapes,
NodeCosts* node_costs);
protected:
std::map<std::string, int> elementwise_ops_;
typedef std::function<absl::Status(const OpContext& op_context, NodeCosts*)>
CostImpl;
std::map<std::string, CostImpl> device_cost_impl_;
// If true, assume compute and memory overlap; hence, the op cost is max of
// compute_time and memory_time, instead of sum of those two.
bool compute_memory_overlap_;
std::set<std::string> persistent_ops_;
private:
friend class OpLevelCostEstimatorTest;
};
} // end namespace grappler
} // end namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPPLER_COSTS_OP_LEVEL_COST_ESTIMATOR_H_
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,123 @@
/* 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.
==============================================================================*/
syntax = "proto3";
package tensorflow;
option cc_enable_arenas = true;
import "tensorflow/core/framework/tensor.proto";
import "tensorflow/core/framework/tensor_shape.proto";
import "tensorflow/core/framework/types.proto";
import "tensorflow/core/framework/attr_value.proto";
import "tensorflow/core/protobuf/device_properties.proto";
// Description of the session when an op is run.
message SessionInfo {
int64 intra_op_parallelism = 1;
}
// Description of an operation as well as the parameters expected to impact its
// performance.
message OpInfo {
// The operation name. There may be custom parameters in attrs.
string op = 1;
// Custom parameters impacting the behavior of the op.
map<string, AttrValue> attr = 2;
// Input data types, shapes and values if known.
message TensorProperties {
DataType dtype = 1;
TensorShapeProto shape = 2;
TensorProto value = 3;
};
repeated TensorProperties inputs = 3;
// Optional description of the op outputs
repeated TensorProperties outputs = 5;
// Device on which the operation is run.
DeviceProperties device = 4;
// Information about the session configs.
SessionInfo session_info = 6;
}
message NormalDistribution {
double mu = 1;
double sigma = 2;
}
message LogNormalDistribution {
double mu = 1;
double sigma = 2;
}
// Performance data for tensorflow operations
message OpPerformance {
// The op
OpInfo op = 1;
// Information about the session configs.
SessionInfo session_info = 12 [deprecated = true];
// The node name (optional). Makes it easier to associate the performance data
// with a specific graph node.
string node = 5;
// Temporary memory used by this node (in bytes).
int64 temporary_memory_size = 2;
// Time it takes to run the op (in nanoseconds).
int64 compute_cost = 3;
// Analytical compute cost (in nanoseconds).
int64 compute_time = 6;
// Analytical memory access cost (in nanoseconds).
int64 memory_time = 7;
// Percentage of theoretical compute performance.
double compute_efficiency = 4;
// Percentage of theoretical memory performance.
double memory_efficiency = 8;
// Expected execution time, modeled using one of 2 possible distributions.
oneof execution_time {
NormalDistribution execution_time_normal = 10;
LogNormalDistribution execution_time_log_normal = 11;
};
// Memory usage data for a tensorflow operation.
message OpMemory {
// The output information may have memory usage and output shapes.
repeated int64 output_memory = 1;
// Temp and persistent memory allocated by this node.
int64 temp_memory = 2;
int64 persistent_memory = 4;
int64 device_temp_memory = 3 [deprecated = true];
int64 device_persistent_memory = 5 [deprecated = true];
}
OpMemory op_memory = 9;
}
// A collection of OpPerformance data points.
message OpPerformanceList {
repeated OpPerformance op_performance = 1;
}
@@ -0,0 +1,153 @@
/* 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/grappler/costs/robust_stats.h"
#include <algorithm>
#include <cmath>
#include <utility>
namespace tensorflow {
namespace grappler {
// Given a sorted vector of values, calculate the median.
// Returns 0 for an empty vector. Does not verify sortedness.
static double SortedMedian(const std::vector<double> &values) {
const int n = values.size();
if (n == 0) return 0.0;
if (n & 1) {
return values[n / 2];
} else {
return (values[n / 2] + values[n / 2 - 1]) / 2.0;
}
}
// Given a vector of values (sorted or not), calculate the median.
static double Median(std::vector<double> &&values) {
const size_t n = values.size();
if (n == 0) return 0;
const auto middle = values.begin() + (n / 2);
// Put the middle value in its place.
std::nth_element(values.begin(), middle, values.end());
if (n & 1) {
return *middle;
}
// Return the average of the two elements, the max_element lower than
// *middle is found between begin and middle as a post-cond of
// nth_element.
const auto lower_middle = std::max_element(values.begin(), middle);
// Preventing overflow. We know that '*lower_middle <= *middle'.
// If both are on opposite sides of zero, the sum won't overflow, otherwise
// the difference won't overflow.
if (*lower_middle <= 0 && *middle >= 0) {
return (*lower_middle + *middle) / 2;
}
return *lower_middle + (*middle - *lower_middle) / 2;
}
// Given a set of values, calculates the scaled Median Absolute Deviation (a
// robust approximation to the standard deviation). This is calculated as the
// median of the absolute deviations from the median, scaled by 1.4826. Its
// advantage over the standard deviation is that it is not (as) affected by
// outlier values. Returns a pair<median, mad>.
static std::pair<double, double> ScaledMedianAbsoluteDeviation(
const std::vector<double> &sorted_values) {
double median = SortedMedian(sorted_values);
// Next, we calculate the absolute deviations from the median,
// find the median of the resulting data, and scale by 1.4826.
std::vector<double> deviations;
deviations.reserve(sorted_values.size());
for (double d : sorted_values) {
deviations.push_back(std::abs(d - median));
}
double mad = Median(std::move(deviations)) * 1.4826;
return std::pair<double, double>(median, mad);
}
RobustStats::RobustStats(const std::vector<double> &values)
: RobustStats(std::vector<double>(values)) {}
RobustStats::RobustStats(std::vector<double> &&values) {
std::sort(values.begin(), values.end());
lo_ = values[0];
hi_ = values.back();
HuberMAD(values);
}
// Computes an updated mean using Huber's weighting function (values beyond
// the margin are weighted by margin / abs(value - mean).
double UpdateHuberMean(const std::vector<double> &sorted_values, double mean,
double margin) {
int num_within = 0;
double sum = 0.0;
for (double d : sorted_values) {
if (d < mean - margin) {
sum -= margin;
} else if (d > mean + margin) {
sum += margin;
} else {
sum += d;
++num_within;
}
}
// It is possible, for a set with an interquartile distance of 0, i.e., with
// more than half of the values at the median, to encounter the case where
// the Huber mean drifts slightly off the median and there are no values
// within the margin. In that case, just return the old mean, and the caller
// will quit.
if (num_within > 0) {
return sum / num_within;
} else {
return mean;
}
}
// Given a list of values, this approximates the stddev using the MAD and then
// uses it to compute a Huber robust mean (sandwich mean). A margin of
// c*stddev is defined around the current mean, and values are weighted by
// margin / abs(value - mean) if outside the margin, or 1 if inside. This
// computes the mean iteratively, because each time it changes the margin
// shifts a bit. It typically settles very quickly, but it's possible for it
// to be unstable. We limit it to 10 iterations.
//
void RobustStats::HuberMAD(const std::vector<double> &sorted_values) {
const std::pair<double, double> median_mad =
ScaledMedianAbsoluteDeviation(sorted_values);
mean_ = median_mad.first;
stddev_ = median_mad.second;
// c = 1.345 is the commonly used cutoff with 95% efficiency at the normal.
// We're using c = 1.5 to be a little more conservative, and because that's
// the default in S-plus.
// TODO(dehnert): Specialize Stats for integral types so we don't implement
// methods that don't make sense.
const double c = 1.5;
const double margin = c * stddev_;
// Iterate 10 times, or until the Huber mean stabilizes.
// If the margin is zero, we don't want mean to drift from the median.
if (margin > 0.0) {
for (int k = 0; k < 10; ++k) {
double old_mean = mean_;
mean_ = UpdateHuberMean(sorted_values, mean_, margin);
if (mean_ == old_mean) break;
}
}
}
} // namespace grappler
} // namespace tensorflow
@@ -0,0 +1,42 @@
/* 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_GRAPPLER_COSTS_ROBUST_STATS_H_
#define TENSORFLOW_CORE_GRAPPLER_COSTS_ROBUST_STATS_H_
#include <vector>
namespace tensorflow {
namespace grappler {
class RobustStats {
public:
explicit RobustStats(const std::vector<double>& values);
explicit RobustStats(std::vector<double>&& values);
double lo() const { return lo_; }
double hi() const { return hi_; }
double mean() const { return mean_; }
private:
void HuberMAD(const std::vector<double>& values);
double lo_;
double hi_;
double mean_;
double stddev_;
};
} // namespace grappler
} // namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPPLER_COSTS_ROBUST_STATS_H_
@@ -0,0 +1,63 @@
/* 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/grappler/costs/robust_stats.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace grappler {
namespace {
class RobustStatsTest : public ::testing::Test {
public:
void SetUp() override {
for (double d = 1.0; d <= 5.0; d += 1.0) {
values1_.push_back(5.0 - d);
values1_.push_back(5.0 + d);
values2_.push_back(25.0 - 2 * d);
values2_.push_back(25.0 + 2 * d);
values3_.push_back(-3.0 - d);
values3_.push_back(-3.0 + d);
}
values1_.push_back(5.0); // Odd # elements, mean is 5.0
values3_.push_back(197.0);
values3_.push_back(-203.0); // Even # elements, mean is -3.0
}
std::vector<double> values1_;
std::vector<double> values2_;
std::vector<double> values3_;
};
TEST_F(RobustStatsTest, Simple) {
RobustStats s1(values1_);
EXPECT_EQ(5.0, s1.mean());
EXPECT_EQ(0.0, s1.lo());
EXPECT_EQ(10.0, s1.hi());
RobustStats s2(values2_);
EXPECT_EQ(25.0, s2.mean());
EXPECT_EQ(15.0, s2.lo());
EXPECT_EQ(35.0, s2.hi());
RobustStats s3(values3_);
EXPECT_EQ(-3.0, s3.mean());
EXPECT_EQ(-203.0, s3.lo());
EXPECT_EQ(197.0, s3.hi());
}
} // namespace
} // namespace grappler
} // namespace tensorflow
+512
View File
@@ -0,0 +1,512 @@
/* 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/grappler/costs/utils.h"
#include <stddef.h>
#include <utility>
#include "absl/strings/str_cat.h"
#include "absl/strings/str_format.h"
#include "Eigen/Core" // from @eigen_archive
#include "tensorflow/core/common_runtime/gpu/gpu_id.h"
#include "tensorflow/core/common_runtime/gpu/gpu_id_manager.h"
#include "tensorflow/core/framework/allocation_description.pb.h"
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_def.pb.h"
#include "tensorflow/core/framework/step_stats.pb.h"
#include "tensorflow/core/framework/tensor.pb.h"
#include "tensorflow/core/framework/tensor_description.pb.h"
#include "tensorflow/core/framework/tensor_shape.pb.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/graph/tensor_id.h"
#include "tensorflow/core/grappler/clusters/utils.h"
#include "tensorflow/core/grappler/utils.h"
#include "tensorflow/core/lib/core/bits.h"
#include "tensorflow/core/lib/strings/numbers.h"
#include "tensorflow/core/platform/byte_order.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/protobuf.h"
#include "tensorflow/core/protobuf/config.pb.h"
#include "tensorflow/core/util/device_name_utils.h"
#include "tensorflow/core/util/overflow.h"
namespace tensorflow {
namespace grappler {
static OpInfo::TensorProperties UnknownInput() {
OpInfo::TensorProperties input;
input.set_dtype(DataType::DT_INVALID);
input.mutable_shape()->set_unknown_rank(true);
return input;
}
static std::vector<TensorProto> ExtractTensors(const AttrValue& attr_value) {
std::vector<TensorProto> tensors;
switch (attr_value.value_case()) {
case AttrValue::kTensor: {
tensors.push_back(attr_value.tensor());
break;
}
case AttrValue::kList: {
for (const auto& tensor_proto : attr_value.list().tensor()) {
tensors.push_back(tensor_proto);
}
break;
}
default: {
}
}
return tensors;
}
// Annotate the op_info inputs with extra information when possible (e.g. the
// input value if it's known statically).
static void ExtractExtraProperties(
const NodeDef& node,
const std::unordered_map<std::string, const NodeDef*>& name_to_node,
OpInfo* op_info) {
OpRegistry* op_registry = OpRegistry::Global();
const OpDef* op_def = nullptr;
auto s = op_registry->LookUpOpDef(node.op(), &op_def);
if (!s.ok()) {
op_def = nullptr;
}
for (int i = 0; i < node.input_size(); ++i) {
const std::string input_name = node.input(i);
CHECK(!input_name.empty());
if (IsControlInput(input_name)) {
continue;
}
TensorId input_tensor_id = ParseTensorName(input_name);
const std::string input_node_name(input_tensor_id.first);
auto iter = name_to_node.find(input_node_name);
if (iter == name_to_node.end()) continue;
const NodeDef* input_node = iter->second;
if (i >= op_info->inputs_size()) {
LOG(ERROR) << "OpInfo's inputs doesn't match the graph! OpInfo: "
<< op_info->DebugString()
<< "\nCurrent node: " << node.DebugString()
<< "\nInput node: " << input_node->DebugString();
}
// The value attribute in Const input is useful for cost prediction.
if (input_node->op() == "Const" && i < op_info->inputs_size()) {
auto it = input_node->attr().find("value");
if (it == input_node->attr().end()) continue;
const AttrValue& attr_value = it->second;
std::vector<TensorProto> tensors = ExtractTensors(attr_value);
if (tensors.empty()) continue;
const TensorProto& t = tensors[0];
OpInfo::TensorProperties* input = op_info->mutable_inputs(i);
*(input->mutable_value()) = t;
// For filename input, the file size can also be useful.
if (op_def && i < op_def->input_arg_size() &&
op_def->input_arg(i).name().find("filename") != std::string::npos) {
Tensor tensor;
if (!tensor.FromProto(t)) {
continue;
}
if (tensor.NumElements() != 1) {
continue;
}
const std::string& filename = tensor.scalar<tstring>()();
Env* env = Env::Default();
FileStatistics stat;
absl::Status s = env->Stat(filename, &stat);
if (!s.ok()) {
continue;
}
AttrValue attr;
attr.set_i(stat.length);
std::string attr_key = absl::StrCat("input_", i, "_filesize");
(*op_info->mutable_attr())[attr_key] = attr;
}
}
// When the input is a handle (e.g. look up table handle), the information
// in the op itself is not sufficient to predict the op memory.
if (op_def && i < op_def->input_arg_size() &&
op_def->input_arg(i).name().find("handle") != std::string::npos) {
std::string new_key = absl::StrCat("parent_", i, "_op");
AttrValue attr;
attr.set_s(input_node->op());
(*op_info->mutable_attr())[new_key] = attr;
// TODO(yuefengz): Only parent node's op name is copied. Copy inputs
// and attributes when necessary.
}
}
}
std::vector<OpInfo::TensorProperties> FindInputFeatures(
const NodeDef& node,
const std::unordered_map<std::string, const CostGraphDef::Node*>&
name_to_cost,
const std::unordered_map<std::string, const NodeDef*>& name_to_node) {
std::vector<OpInfo::TensorProperties> inputs;
for (const auto& input_name : node.input()) {
CHECK(!input_name.empty());
TensorId input_tensor_id = ParseTensorName(input_name);
const std::string input_node_name(input_tensor_id.first);
const int output_index = input_tensor_id.second;
// Skip control inputs.
if (output_index == Graph::kControlSlot) {
continue;
}
auto it = name_to_cost.find(input_node_name);
if (it == name_to_cost.end() || output_index < 0) {
inputs.push_back(UnknownInput());
} else {
const CostGraphDef::Node* input_cost = it->second;
if (input_cost->output_info_size() == 0) {
inputs.push_back(UnknownInput());
} else {
const CostGraphDef::Node::OutputInfo& output =
input_cost->output_info(output_index);
OpInfo::TensorProperties input;
input.set_dtype(output.dtype());
*input.mutable_shape() = output.shape();
inputs.push_back(input);
}
}
}
return inputs;
}
int64_t CalculateTensorSize(const OpInfo::TensorProperties& prop) {
int64_t size = DataTypeSize(BaseType(prop.dtype()));
TensorShapeProto shape = prop.shape();
// Can't infer the size if the rank is unknown. It has to be at least a
// scalar though.
if (shape.unknown_rank()) {
VLOG(2) << "CalculateTensorSize() -- unknown rank";
return size;
}
// If one of the dimensions is unknown statically, assume it's at least one.
for (int i = 0; i < shape.dim_size(); ++i) {
if (shape.dim(i).size() < 0) {
shape.mutable_dim(i)->set_size(1);
VLOG(2) << "CalculateTensorSize() -- unknown dim: " << i;
}
}
int64_t num_elems = TensorShape(shape).num_elements();
int64_t tensor_size = MultiplyWithoutOverflow(num_elems, size);
if (tensor_size < 0) {
VLOG(1) << "Overflow encountered when computing tensor size, multiplying "
<< num_elems << " with " << size;
return -1;
}
return tensor_size;
}
int64_t CalculateOutputSize(
const std::vector<OpInfo::TensorProperties>& output_properties,
const int port_num) {
if (port_num < 0) return 4; // 4B for control dependency.
if (port_num >= output_properties.size()) {
LOG(ERROR) << "CalculateOutputSize() -- port_num: " << port_num
<< " >= output_properties.size(): " << output_properties.size();
return 0;
}
return CalculateTensorSize(output_properties[port_num]);
}
DeviceProperties GetDeviceInfo(const std::string& device_str) {
DeviceProperties unknown;
unknown.set_type("UNKNOWN");
DeviceNameUtils::ParsedName parsed;
if (DeviceNameUtils::ParseFullName(device_str, &parsed)) {
if (parsed.type == "GPU") {
TfDeviceId tf_device_id(parsed.id);
PlatformDeviceId platform_device_id;
absl::Status s =
GpuIdManager::TfToPlatformDeviceId(tf_device_id, &platform_device_id);
if (!s.ok()) {
// We are probably running simulation without linking cuda libraries.
platform_device_id = PlatformDeviceId(parsed.id);
}
return GetLocalGPUInfo(platform_device_id);
} else if (parsed.type == "CPU") {
return GetLocalCPUInfo();
}
}
return unknown;
}
DeviceProperties GetDeviceInfo(const CostGraphDef::Node& node) {
return GetDeviceInfo(node.device());
}
OpInfo BuildOpInfoWithoutDevice(
const NodeDef& node,
const std::unordered_map<std::string, const NodeDef*>& name_to_node,
const std::vector<OpInfo::TensorProperties>& inputs) {
OpInfo op_info;
op_info.set_op(node.op());
*op_info.mutable_attr() = node.attr();
for (auto& input : inputs) {
*op_info.add_inputs() = input;
}
ExtractExtraProperties(node, name_to_node, &op_info);
return op_info;
}
std::string GetOpDescription(const OpInfo& op_info) {
std::string description = "[";
description += "Op=" + op_info.op() + ", ";
description += "input_shapes=[";
for (auto const& input : op_info.inputs()) {
description += PartialTensorShape::DebugString(input.shape());
}
description += "]";
return description;
}
OpPerformanceList CostGraphToOpPerformanceData(const CostGraphDef& cost_graph,
const GraphDef& graph) {
OpPerformanceList ret;
std::unordered_map<std::string, const CostGraphDef::Node*> name_to_cost;
std::unordered_map<std::string, const NodeDef*> name_to_node;
for (auto& node : cost_graph.node()) {
name_to_cost[node.name()] = &node;
}
for (auto& node : graph.node()) {
name_to_node[node.name()] = &node;
}
for (const auto& node : graph.node()) {
// Skip the nodes that are not in the cost graph: these are nodes that
// aren't run, because they aren't in the intersection of transitive
// fan-in of a fetch node and the transitive fan-out of an input, or nodes
// that were optimized away by the optimizer. Since they don't contribute
// to the execution time we simply discard them.
auto it = name_to_cost.find(node.name());
if (it == name_to_cost.end()) {
continue;
}
const CostGraphDef::Node* cost_node = it->second;
OpPerformance* perf = ret.add_op_performance();
perf->set_node(node.name());
std::vector<OpInfo::TensorProperties> inputs =
FindInputFeatures(node, name_to_cost, name_to_node);
*perf->mutable_op() = BuildOpInfoWithoutDevice(node, name_to_node, inputs);
*perf->mutable_op()->mutable_device() = GetDeviceInfo(cost_node->device());
perf->set_temporary_memory_size(cost_node->temporary_memory_size());
// Note that CostGraphDef::Node::compute_cost is microseconds, while
// OpPerformance.compute_cost is nanoseconds.
perf->set_compute_cost(cost_node->compute_cost() * 1000);
perf->set_compute_time(cost_node->compute_time() * 1000);
perf->set_memory_time(cost_node->memory_time() * 1000);
for (const auto& output_info : cost_node->output_info()) {
perf->mutable_op_memory()->add_output_memory(output_info.size());
}
perf->mutable_op_memory()->set_temp_memory(
cost_node->temporary_memory_size());
perf->mutable_op_memory()->set_persistent_memory(
cost_node->persistent_memory_size());
}
return ret;
}
void TensorSizeHistogram::Add(const uint64_t value) {
num_elem_++;
sum_elem_ += value;
min_ = std::min(min_, value);
max_ = std::max(max_, value);
buckets_[Index(value)]++;
}
void TensorSizeHistogram::Merge(const TensorSizeHistogram& src) {
num_elem_ += src.num_elem_;
sum_elem_ += src.sum_elem_;
min_ = std::min(min_, src.min_);
max_ = std::max(max_, src.max_);
std::transform(buckets_.begin(), buckets_.end(), src.buckets_.begin(),
buckets_.begin(), std::plus<uint64_t>());
}
std::string TensorSizeHistogram::ToString() const {
std::string r = absl::StrFormat(
"Count: %lld, Average: %s, Min: %s, Max: %s"
"\n------------------------------------------------------\n",
num_elem_, strings::HumanReadableNumBytes(Average()),
strings::HumanReadableNumBytes(min_),
strings::HumanReadableNumBytes(max_));
const double mult = num_elem_ > 0 ? 100.0 / num_elem_ : 0.0;
uint64_t cumul_sum = 0;
for (int i = 0; i < buckets_.size(); i++) {
if (buckets_[i] == 0) continue;
cumul_sum += buckets_[i];
uint64_t left = i == 0 ? 0ULL : 1ULL << (i - 1);
uint64_t right = 1ULL << i;
absl::StrAppendFormat(&r, "[ %12s, %12s) %7d %7.3f%% %7.3f%% ",
strings::HumanReadableNumBytes(left),
strings::HumanReadableNumBytes(right),
buckets_[i], // count
mult * buckets_[i], // percentage
mult * cumul_sum); // cumulative percentage
// Add hash marks based on percentage; 40 marks for 100%.
auto marks = static_cast<int>(
(static_cast<double>(40 * buckets_[i] + (num_elem_ >> 1)) / num_elem_));
absl::StrAppendFormat(&r, "%s\n", std::string(marks, '#'));
}
return r;
}
const int TensorSizeHistogram::Index(const uint64_t value) const {
// Log2Floor64 returns -1 for 0, 0 for 1, 1 for 2-3, 2 for 4-7, ...
const auto index = Log2Floor64(value) + 1;
return std::min(index, kMaxBuckets - 1);
}
std::string GetDeviceClassForNonChannelDevice(const std::string& device_name) {
DeviceNameUtils::ParsedName parsed_name;
bool parsed = DeviceNameUtils::ParseFullName(device_name, &parsed_name);
if (!parsed) {
std::string name =
str_util::StringReplace(device_name, "/job_", "/job:", true);
name = str_util::StringReplace(name, "/replica_", "/replica:", true);
name = str_util::StringReplace(name, "/task_", "/task:", true);
name = str_util::StringReplace(name, "/device_", "/device:", true);
name = str_util::StringReplace(name, "GPU_", "GPU:", true);
name = str_util::StringReplace(name, "CPU_", "CPU:", true);
name = str_util::StringReplace(name, "gpu_", "gpu:", true);
name = str_util::StringReplace(name, "cpu_", "cpu:", true);
parsed = DeviceNameUtils::ParseFullName(name, &parsed_name);
}
if (parsed) {
const std::string jobname = parsed_name.has_job ? parsed_name.job : "";
return absl::StrCat("/", jobname, "/", parsed_name.type);
} else {
return "Unclassified";
}
}
std::string GetDeviceClass(const std::string& device_name) {
// TODO(dyoon): channel device name follows the convention we currently have
// in VirtualScheduler. This should be revised with VirtualScheduler as well
// as VirtualPlacer in the future.
if (device_name.find("Channel") != std::string::npos) {
const std::string from = "_from_";
const std::string to = "_to_";
const auto from_loc = device_name.find(from);
const auto to_loc = device_name.find(to);
const auto src_device_full = device_name.substr(
from_loc + from.size(), to_loc - (from_loc + from.size()));
const auto dst_device_full = device_name.substr(to_loc + to.size());
return absl::StrCat(
"Channel", ": ", GetDeviceClassForNonChannelDevice(src_device_full),
" -> ", GetDeviceClassForNonChannelDevice(dst_device_full));
} else {
return GetDeviceClassForNonChannelDevice(device_name);
}
}
std::string GetStatsStringFromRunMetadata(const RunMetadata& run_metadata,
bool verbosity) {
// TODO(dyoon): print out other stats as needed.
std::ostringstream output;
// Tensor size histogram:
// if verbosity, it outputs per-device histogram,
// otherwise, only per-class histogram.
std::unordered_map<std::string, TensorSizeHistogram> device_to_hist_map;
const auto& step_stats = run_metadata.step_stats();
for (const auto& dev_stat : step_stats.dev_stats()) {
const auto& device_name = dev_stat.device();
auto& hist = device_to_hist_map[device_name];
for (const auto& node_stat : dev_stat.node_stats()) {
for (const auto& node_output : node_stat.output()) {
// TODO(dyoon): Calculate tensor size from tensor_description's dtype
// and shape, instead of using optional allocation_description.
const auto size = node_output.tensor_description()
.allocation_description()
.allocated_bytes();
hist.Add(size);
}
}
}
if (verbosity) {
output << "\n";
output << "Per device tensor size histogram.\n";
}
std::unordered_map<std::string, TensorSizeHistogram> device_class_to_hist_map;
for (const auto& device_hist : device_to_hist_map) {
const auto& device_name = device_hist.first;
const auto& hist = device_hist.second;
if (verbosity) {
output << "Device: " << device_name << "\n" << hist.ToString() << "\n";
}
const auto device_class = GetDeviceClass(device_name);
auto it = device_class_to_hist_map.find(device_class);
if (it == device_class_to_hist_map.end()) {
device_class_to_hist_map.emplace(device_class, TensorSizeHistogram(hist));
} else {
it->second.Merge(hist);
}
}
output << "\n";
output << "Aggregated per device / channel type tensor size histogram:\n";
for (const auto& device_hist : device_class_to_hist_map) {
const auto& device_name = device_hist.first;
const auto& hist = device_hist.second;
output << "Device: " << device_name << "\n" << hist.ToString() << "\n";
}
output << "\n";
return output.str();
}
void CombineCostsAndUpdateExecutionTime(bool compute_memory_overlap,
Costs* costs) {
if (compute_memory_overlap) {
costs->execution_time =
std::max(costs->intermediate_memory_time,
std::max(costs->compute_time, costs->memory_time));
} else {
costs->execution_time = costs->compute_time + costs->memory_time +
costs->intermediate_memory_time;
}
}
} // end namespace grappler
} // end namespace tensorflow
+135
View File
@@ -0,0 +1,135 @@
/* 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_GRAPPLER_COSTS_UTILS_H_
#define TENSORFLOW_CORE_GRAPPLER_COSTS_UTILS_H_
#include <cstdint>
#include <limits>
#include <string>
#include <unordered_map>
#include <vector>
#include "tensorflow/core/framework/cost_graph.pb.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/op_def.pb.h"
#include "tensorflow/core/graph/types.h"
#include "tensorflow/core/grappler/costs/cost_estimator.h"
#include "tensorflow/core/grappler/costs/op_performance_data.pb.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/protobuf/config.pb.h"
#include "tensorflow/core/protobuf/device_properties.pb.h"
namespace tensorflow {
namespace grappler {
// Returns a vector of InputProperties for 'node'. The vector will contain one
// entry for each input of 'node'.
// For each node in the graph, the 'name_to_cost' map stores a pointer to the
// corresponding cost graph node indexed by node name. The 'name_to_node' maps a
// node name to its node definition.
std::vector<OpInfo::TensorProperties> FindInputFeatures(
const NodeDef& node,
const std::unordered_map<std::string, const CostGraphDef::Node*>&
name_to_cost,
const std::unordered_map<std::string, const NodeDef*>& name_to_node);
// Returns the size of tensor (unit: bytes). For tensor shape with unknown rank,
// it assumes the tensor to be scalar. For any unknown dimension, it assumes
// size one.
int64_t CalculateTensorSize(const OpInfo::TensorProperties& prop);
// Returns the size of output at port_num (unit: bytes). A special case is
// port_num -1, which is for control dependency and assumed to be 4 bytes.
int64_t CalculateOutputSize(
const std::vector<OpInfo::TensorProperties>& output_properties,
int port_num);
// Returns the DeviceProperties of the device on which 'node' runs.
DeviceProperties GetDeviceInfo(const CostGraphDef::Node& node);
DeviceProperties GetDeviceInfo(const std::string& device_str);
// Return a string describing a node given a nodeinfo.
std::string GetOpDescription(const OpInfo& op_info);
// Builds the OpInfo for node without filling its device information, given all
// nodes in the graph and its input properties.
OpInfo BuildOpInfoWithoutDevice(
const NodeDef& node,
const std::unordered_map<std::string, const NodeDef*>& name_to_node,
const std::vector<OpInfo::TensorProperties>& inputs);
// Gather performance data from a cost graph.
OpPerformanceList CostGraphToOpPerformanceData(const CostGraphDef& cost_graph,
const GraphDef& graph);
// Simple histogram for profiling Tensor size; histogram uses logarithmic
// buckets.
class TensorSizeHistogram {
public:
TensorSizeHistogram() : buckets_(kMaxBuckets, 0) {}
void Add(const uint64_t value);
void Merge(const TensorSizeHistogram& src);
double Average() const {
if (num_elem_ > 0) {
return static_cast<double>(sum_elem_) / num_elem_;
} else {
return 0.0;
}
}
uint64_t Min() const { return min_; }
uint64_t Max() const { return max_; }
uint64_t NumElem() const { return num_elem_; }
uint64_t SumElem() const { return sum_elem_; }
std::string ToString() const;
protected:
const int Index(const uint64_t value) const;
const std::vector<uint64_t>& GetBuckets() const { return buckets_; }
private:
const int kMaxBuckets = 64;
uint64_t num_elem_ = 0;
uint64_t sum_elem_ = 0;
// min_ and max_ are initialized to a very large value and zero, respectively,
// so that any value added can replace the initial min_ and max_.
uint64_t min_ = std::numeric_limits<uint64_t>::max();
uint64_t max_ = 0;
// Buckets are logarithmic:
// 0B, 1B, 2-3B, 4-7B, 8-15B, ..., 2^N - 2^(N+1)-1B, ...
std::vector<uint64_t> buckets_;
};
// Helper functions for aggregating per-device stats into per-device-class
// stats.
std::string GetDeviceClassForNonChannelDevice(const std::string& device_name);
std::string GetDeviceClass(const std::string& device_name);
// Get stats in string format from RunMetadata.
std::string GetStatsStringFromRunMetadata(const RunMetadata& run_metadata,
bool verbosity);
// This method calculates the execution time depending on whether IO can
// overlap with computation. It assumes the memory and the compute times have
// already been calculated.
void CombineCostsAndUpdateExecutionTime(bool compute_memory_overlap,
Costs* costs);
} // end namespace grappler
} // end namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPPLER_COSTS_UTILS_H_
@@ -0,0 +1,347 @@
/* 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/grappler/costs/utils.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def_builder.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace grappler {
namespace {
void CreateConstOp(const std::string& name, std::initializer_list<int64_t> dims,
NodeDef* node) {
Tensor tensor(DT_FLOAT, TensorShape(dims));
for (int64_t i = 0; i < tensor.NumElements(); ++i)
tensor.flat<float>()(i) = i / 10.0f;
TF_CHECK_OK(NodeDefBuilder(name, "Const")
.Attr("dtype", DT_FLOAT)
.Attr("value", tensor)
.Finalize(node));
}
void CreateConstSizesOp(const std::string& name,
const std::vector<int32_t>& sizes, NodeDef* node) {
TensorShape shape;
shape.AddDim(sizes.size());
Tensor tensor(DT_INT32, shape);
for (int64_t i = 0; i < tensor.NumElements(); ++i)
tensor.flat<int32_t>()(i) = sizes[i];
TF_CHECK_OK(NodeDefBuilder(name, "Const")
.Attr("dtype", DT_INT32)
.Attr("value", tensor)
.Finalize(node));
}
// Helper method for converting shapes vector to TensorProperty.
OpInfo::TensorProperties ShapeToTensorProperty(const std::vector<int>& shapes,
const DataType& data_type) {
OpInfo::TensorProperties prop;
prop.set_dtype(data_type);
for (int shape : shapes) prop.mutable_shape()->add_dim()->set_size(shape);
return prop;
}
TEST(UtilsTest, ConvOpInfo) {
int batch = 32;
int rows = 7;
int cols = 9;
int filter_rows = 3;
int filter_cols = 3;
int out_rows = 7;
int out_cols = 9;
int in_depth = 3;
int out_depth = 5;
int stride = 1;
std::unordered_map<std::string, const NodeDef*> name_to_node;
GraphDef graph;
NodeDef* input = graph.add_node();
name_to_node["input"] = input;
CreateConstOp("input", {batch, rows, cols, in_depth}, input);
NodeDef* filter = graph.add_node();
name_to_node["filter"] = filter;
CreateConstOp("filter", {filter_rows, filter_cols, in_depth, out_depth},
filter);
NodeDef* output_backprop = graph.add_node();
name_to_node["output_backprop"] = output_backprop;
CreateConstOp("output_backprop", {batch, out_rows, out_cols, out_depth},
output_backprop);
NodeDef* input_sizes = graph.add_node();
name_to_node["input_sizes"] = input;
CreateConstSizesOp("input_sizes",
std::vector<int32_t>({batch, rows, cols, in_depth}),
input_sizes);
NodeDef* filter_sizes = graph.add_node();
name_to_node["filter_sizes"] = filter_sizes;
CreateConstSizesOp(
"filter_sizes",
std::vector<int32_t>({filter_rows, filter_cols, in_depth, out_depth}),
filter_sizes);
TensorShape paddings_shape({4, 2});
Tensor paddings_tensor(DT_INT32, paddings_shape);
for (int64_t i = 0; i < paddings_tensor.NumElements(); ++i) {
paddings_tensor.flat<int32_t>()(i) = 0;
}
TF_CHECK_OK(NodeDefBuilder("paddings", "Const")
.Attr("dtype", DT_INT32)
.Attr("value", paddings_tensor)
.Finalize(graph.add_node()));
// Now add the convolution op
NodeDef* conv = graph.add_node();
TF_CHECK_OK(NodeDefBuilder("conv2d", "Conv2D")
.Input("input", 0, DT_FLOAT)
.Input("filter", 0, DT_FLOAT)
.Attr("strides", {1, stride, stride, 1})
.Attr("padding", "SAME")
.Finalize(conv));
NodeDef* conv_bp_in = graph.add_node();
TF_CHECK_OK(NodeDefBuilder("conv2d_bp_in", "Conv2DBackpropInput")
.Input("input_sizes", 0, DT_INT32)
.Input("filter", 0, DT_FLOAT)
.Input("output_backprop", 0, DT_FLOAT)
.Attr("strides", {1, stride, stride, 1})
.Attr("padding", "SAME")
.Finalize(conv_bp_in));
NodeDef* conv_bp_filter = graph.add_node();
TF_CHECK_OK(NodeDefBuilder("conv2d_bp_filter", "Conv2DBackpropFilter")
.Input("input", 0, DT_FLOAT)
.Input("filter_sizes", 0, DT_INT32)
.Input("output_backprop", 0, DT_FLOAT)
.Attr("strides", {1, stride, stride, 1})
.Attr("padding", "SAME")
.Finalize(conv_bp_filter));
for (const auto& node : graph.node()) {
if (node.name().find("conv2d") != 0) {
continue;
}
std::vector<OpInfo::TensorProperties> inputs;
inputs.resize(node.input_size());
OpInfo info = BuildOpInfoWithoutDevice(node, name_to_node, inputs);
if (node.name() == "conv2d") {
EXPECT_EQ(2, info.inputs_size());
} else if (node.name() == "conv2dbp_in") {
EXPECT_EQ(3, info.inputs_size());
} else if (node.name() == "conv2d_bp_filter") {
EXPECT_EQ(3, info.inputs_size());
}
}
}
TEST(UtilsTest, TestSkipControlInput) {
GraphDef graph;
TF_CHECK_OK(NodeDefBuilder("constant", "Const")
.Attr("dtype", DT_INT32)
.Finalize(graph.add_node()));
TF_CHECK_OK(NodeDefBuilder("constfold", "NoOp")
.ControlInput("constant")
.Finalize(graph.add_node()));
std::unordered_map<std::string, const NodeDef*> name_to_node;
for (const auto& node : graph.node()) {
name_to_node[node.name()] = &node;
}
bool node_found = false;
for (const auto& node : graph.node()) {
if (node.name() == "constfold") {
std::vector<OpInfo::TensorProperties> inputs;
OpInfo info = BuildOpInfoWithoutDevice(node, name_to_node, inputs);
node_found = true;
EXPECT_EQ(0, info.inputs_size());
}
}
EXPECT_TRUE(node_found);
}
TEST(UtilsTest, CalculateTensorSize) {
// Test normal usage.
EXPECT_EQ(DataTypeSize(DT_FLOAT) * 1,
CalculateTensorSize(ShapeToTensorProperty({1}, DT_FLOAT)));
EXPECT_EQ(DataTypeSize(DT_FLOAT) * 4 * 4,
CalculateTensorSize(ShapeToTensorProperty({4, 4}, DT_FLOAT)));
EXPECT_EQ(DataTypeSize(DT_HALF) * 10 * 10 * 10,
CalculateTensorSize(ShapeToTensorProperty({10, 10, 10}, DT_HALF)));
EXPECT_EQ(
DataTypeSize(DT_FLOAT) * 100 * 7 * 8 * 99,
CalculateTensorSize(ShapeToTensorProperty({100, 7, 8, 99}, DT_FLOAT)));
// Test unknown rank: assumes the tensor to be a scalar.
OpInfo::TensorProperties t = ShapeToTensorProperty({100, 7, 8, 99}, DT_FLOAT);
t.mutable_shape()->set_unknown_rank(true);
EXPECT_EQ(DataTypeSize(DT_FLOAT) * 1, CalculateTensorSize(t));
// Test unknown shape: assumes unknown shape (-1) to have size 1.
EXPECT_EQ(
DataTypeSize(DT_FLOAT) * 1 * 7 * 8 * 99,
CalculateTensorSize(ShapeToTensorProperty({-1, 7, 8, 99}, DT_FLOAT)));
EXPECT_EQ(
DataTypeSize(DT_FLOAT) * 1 * 7 * 1 * 99,
CalculateTensorSize(ShapeToTensorProperty({-1, 7, -1, 99}, DT_FLOAT)));
// Test overflow
EXPECT_EQ(-1, CalculateTensorSize(ShapeToTensorProperty(
{4096, 4096, 4096, 33554432}, DT_FLOAT)));
}
TEST(UtilsTest, CalculateOutputSize) {
// Create a set of tensor properties.
std::vector<OpInfo::TensorProperties> output = {
ShapeToTensorProperty({4, 4}, DT_FLOAT), // 0
ShapeToTensorProperty({-1, 7, -1, 99}, DT_FLOAT) // 1
};
// Test valid outputs.
EXPECT_EQ(DataTypeSize(DT_FLOAT) * 4 * 4, CalculateOutputSize(output, 0));
EXPECT_EQ(DataTypeSize(DT_FLOAT) * 1 * 7 * 1 * 99,
CalculateOutputSize(output, 1));
// port_num -1 is for control dependency: hard coded 4B.
EXPECT_EQ(4, CalculateOutputSize(output, -1));
// Invalid port_num (though it may be an error) shall yield zero
// output size.
EXPECT_EQ(0, CalculateOutputSize(output, 2));
}
// Class for testing TensorSizeHistogram.
class TestTensorSizeHistogram : public TensorSizeHistogram {
public:
FRIEND_TEST(TensorSizeHistogramTest, Constructor);
FRIEND_TEST(TensorSizeHistogramTest, Index);
FRIEND_TEST(TensorSizeHistogramTest, Add);
FRIEND_TEST(TensorSizeHistogramTest, Merge);
};
TEST(TensorSizeHistogramTest, Constructor) {
TestTensorSizeHistogram hist;
EXPECT_EQ(0, hist.NumElem());
EXPECT_EQ(0, hist.SumElem());
EXPECT_LT(1000000000, hist.Min()); // Initially, min_ is a very large value.
EXPECT_EQ(0, hist.Max());
EXPECT_EQ(0.0, hist.Average());
const auto& buckets = hist.GetBuckets();
for (const auto& bucket : buckets) {
EXPECT_EQ(0, bucket);
}
}
TEST(TensorSizeHistogramTest, Index) {
TestTensorSizeHistogram hist;
EXPECT_EQ(0, hist.Index(0));
EXPECT_EQ(1, hist.Index(1));
EXPECT_EQ(2, hist.Index(2));
EXPECT_EQ(2, hist.Index(3));
EXPECT_EQ(3, hist.Index(4));
EXPECT_EQ(3, hist.Index(5));
EXPECT_EQ(3, hist.Index(6));
EXPECT_EQ(3, hist.Index(7));
EXPECT_EQ(4, hist.Index(8));
EXPECT_EQ(4, hist.Index(15));
EXPECT_EQ(5, hist.Index(16));
EXPECT_EQ(5, hist.Index(31));
EXPECT_EQ(6, hist.Index(32));
EXPECT_EQ(11, hist.Index(1025));
}
TEST(TensorSizeHistogramTest, Add) {
TestTensorSizeHistogram hist;
hist.Add(1037);
hist.Add(1038);
hist.Add(1039);
const auto& buckets = hist.GetBuckets();
EXPECT_EQ(3, hist.NumElem());
EXPECT_EQ(1037 + 1038 + 1039, hist.SumElem());
EXPECT_DOUBLE_EQ(1038.0, hist.Average());
EXPECT_EQ(1037, hist.Min());
EXPECT_EQ(1039, hist.Max());
EXPECT_EQ(3, buckets.at(11));
}
TEST(TensorSizeHistogramTest, Merge) {
TestTensorSizeHistogram hist1;
const auto& buckets = hist1.GetBuckets();
hist1.Add(1037);
hist1.Add(1038);
hist1.Add(1039);
TestTensorSizeHistogram hist2(hist1);
hist1.Merge(hist2);
EXPECT_EQ(6, hist1.NumElem());
EXPECT_EQ(2 * (1037 + 1038 + 1039), hist1.SumElem());
EXPECT_DOUBLE_EQ(1038.0, hist1.Average());
EXPECT_EQ(1037, hist1.Min());
EXPECT_EQ(1039, hist1.Max());
EXPECT_EQ(6, buckets.at(11));
TestTensorSizeHistogram hist3;
hist3.Add(1);
hist3.Add(2);
hist3.Add(4);
hist1.Merge(hist3);
EXPECT_EQ(9, hist1.NumElem());
EXPECT_EQ(2 * (1037 + 1038 + 1039) + 1 + 2 + 4, hist1.SumElem());
EXPECT_DOUBLE_EQ((2 * (1037 + 1038 + 1039) + 1 + 2 + 4) / 9.0,
hist1.Average());
EXPECT_EQ(1, hist1.Min());
EXPECT_EQ(1039, hist1.Max());
EXPECT_EQ(1, buckets.at(1));
EXPECT_EQ(1, buckets.at(2));
EXPECT_EQ(1, buckets.at(3));
EXPECT_EQ(6, buckets.at(11));
}
TEST(DeviceClassTest, GetDeviceClass) {
EXPECT_EQ(
"Channel: /ps/CPU -> /worker/GPU",
GetDeviceClass("Channel_from_/job_ps/replica_0/task_0/device_CPU_0_to_"
"/job_worker/replica_7/task_0/device_GPU_7"));
EXPECT_EQ(
"Channel: /worker_train/CPU -> /ps/GPU",
GetDeviceClass(
"Channel_from_/job_worker_train/replica_0/task_0/device_CPU_0_to_"
"/job_ps/replica_7/task_0/device_GPU_7"));
}
TEST(DeviceClassTest, GetDeviceClassForNonChannelDevice) {
EXPECT_EQ("Unclassified",
GetDeviceClassForNonChannelDevice("SOMETHING_WEIRD_DEVICE_NAME"));
EXPECT_EQ("/worker/GPU", GetDeviceClassForNonChannelDevice(
"/job:worker/replica:0/task:0/device:GPU:0"));
EXPECT_EQ("/worker/CPU", GetDeviceClassForNonChannelDevice(
"/job:worker/replica:0/task:0/device:CPU:0"));
EXPECT_EQ("/worker_train/CPU", GetDeviceClassForNonChannelDevice(
"/job:worker_train/replica:7/CPU:0"));
EXPECT_EQ("//GPU", GetDeviceClassForNonChannelDevice("/device:GPU:7"));
}
} // namespace
} // end namespace grappler
} // end namespace tensorflow
@@ -0,0 +1,180 @@
/* 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/grappler/costs/virtual_placer.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/grappler/clusters/cluster.h"
#include "tensorflow/core/grappler/devices.h"
#include "tensorflow/core/lib/strings/str_util.h"
#include "tensorflow/core/util/device_name_utils.h"
namespace tensorflow {
namespace grappler {
VirtualPlacer::VirtualPlacer(
const std::unordered_map<std::string, DeviceProperties>& devices)
: devices_(devices),
// Default job name for canonical device name. Needs to be set before the
// first call to to_lfqn_or_empty()
default_job_name_lowercase_("localhost") {
lfqn_map_.reserve(devices_.size());
for (const auto& kv : devices_) {
const auto lfqn = to_lfqn_or_empty(kv.first);
if (lfqn.empty()) {
LOG(ERROR) << "VirtualPlacer couldn't parse device name from cluster: "
<< kv.first;
} else {
lfqn_map_[lfqn] = kv.first;
}
}
if (devices_.empty()) {
// If there are no devices in the cluster, add a single device, "UNKNOWN" to
// the cluster.
default_device_name_ = "UNKNOWN";
DeviceProperties& prop = devices_["UNKNOWN"];
prop.set_type("UNKNOWN");
} else if (devices_.size() == 1) {
// If there is only one device in the cluster, use it as default device,
// whatever it is.
default_device_name_ = devices_.begin()->first;
} else {
// Default device is set from the devices in the cluster in the following
// priority: /gpu:0, /cpu:0, or any device.
// TODO(dyoon): This logic assumes single machine with CPU and GPU devices.
// Make it more general to support multiple machines, job types, and devices
// other than CPU and GPU.
std::map<int, std::string>
cpu_devices; // CPU device map: id -> device name.
std::map<int, std::string>
gpu_devices; // GPU device map: id -> device name.
for (const auto& kv : lfqn_map_) {
const auto& lfqn = kv.first;
const auto& cluster_device_name = kv.second;
DeviceNameUtils::ParsedName parsed_name;
bool parsed = DeviceNameUtils::ParseFullName(lfqn, &parsed_name);
if (parsed) {
// Parsed devices are stored to cpu_devices or gpu_devices map,
// addressed (and ordered) by device id.
const auto type = absl::AsciiStrToLower(parsed_name.type);
if (type == "gpu") {
gpu_devices[parsed_name.id] = cluster_device_name;
} else if (type == "cpu") {
cpu_devices[parsed_name.id] = cluster_device_name;
}
}
}
if (!gpu_devices.empty()) {
// GPU:0 (or GPU with smallest device id).
default_device_name_ = gpu_devices.begin()->second;
} else if (!cpu_devices.empty()) {
// CPU:0 (or CPU with smallest device id).
default_device_name_ = cpu_devices.begin()->second;
} else {
default_device_name_ = devices_.begin()->first; // Any device.
}
}
VLOG(3) << "default device name: " << default_device_name_;
// Scan the device names from the cluster, and if there is one job name used,
// use it for canonical device name.
std::unordered_set<std::string> job_names_from_cluster;
for (const auto& device : lfqn_map_) {
const auto& lfqn = device.first;
DeviceNameUtils::ParsedName parsed_name;
bool parsed = DeviceNameUtils::ParseFullName(lfqn, &parsed_name);
if (parsed && !parsed_name.job.empty()) {
job_names_from_cluster.insert(parsed_name.job);
if (job_names_from_cluster.size() > 1) {
break;
}
}
}
// If there is only one type of job name in all the devices in the cluster,
// use that one as default job name; otherwise, use localhost.
// TODO(dyoon): this should be improved, especially when the cluster is
// composed of multiple worker, PS, and other types of jobs.
if (job_names_from_cluster.size() == 1) {
auto it = job_names_from_cluster.begin();
default_job_name_lowercase_ = *it;
}
VLOG(3) << "default job name: " << default_job_name_lowercase_;
}
const DeviceProperties& VirtualPlacer::get_device(const NodeDef& node) const {
std::string device = get_canonical_device_name(node);
VLOG(3) << "node.name=" << node.name() << " node.device=" << node.device()
<< " is placed on: " << device;
auto it = devices_.find(device);
DCHECK(it != devices_.end());
return it->second;
}
std::string VirtualPlacer::get_canonical_device_name(
const NodeDef& node) const {
if (node.device().empty()) {
return default_device_name_;
}
const auto lfqn = to_lfqn_or_empty(node.device());
if (lfqn.empty()) {
return default_device_name_;
}
const auto it = lfqn_map_.find(lfqn);
if (it != lfqn_map_.end()) {
return it->second;
}
return default_device_name_;
}
std::string VirtualPlacer::to_lfqn_or_empty(
const std::string& device_name) const {
DeviceNameUtils::ParsedName parsed_name;
const auto lowercase_name = absl::AsciiStrToLower(device_name);
bool parsed = DeviceNameUtils::ParseFullName(lowercase_name, &parsed_name);
if (!parsed) {
parsed = DeviceNameUtils::ParseLocalName(lowercase_name, &parsed_name);
parsed_name.job = "localhost";
}
if (!parsed) {
if (lowercase_name == "gpu" || lowercase_name == "cpu") {
parsed_name.job = "localhost";
parsed_name.type = lowercase_name;
parsed = true;
}
}
if (!parsed) {
return {};
}
if (parsed_name.job.empty()) {
parsed_name.job = default_job_name_lowercase_;
}
// Have to do this, because parser returns uppercase types for CPU and GPU.
parsed_name.type = absl::AsciiStrToLower(parsed_name.type);
std::string lfqn = strings::StrCat(
"/job:", parsed_name.job, "/replica:", parsed_name.replica,
"/task:", parsed_name.task, "/device:", parsed_name.type, ":",
parsed_name.id);
return lfqn;
}
} // end namespace grappler
} // end namespace tensorflow
@@ -0,0 +1,62 @@
/* 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_GRAPPLER_COSTS_VIRTUAL_PLACER_H_
#define TENSORFLOW_CORE_GRAPPLER_COSTS_VIRTUAL_PLACER_H_
#include <unordered_map>
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/protobuf/device_properties.pb.h"
namespace tensorflow {
class NodeDef;
namespace grappler {
class Cluster;
// The virtual placer emulates the behavior of the TF placer.
class VirtualPlacer {
public:
explicit VirtualPlacer(
const std::unordered_map<std::string, DeviceProperties>& devices);
const DeviceProperties& get_device(const NodeDef& node) const;
// Returns device name from cluster, which best matches the node.device()
// specification. Returns default device if no match was found or the
// node.device() could not be parsed.
std::string get_canonical_device_name(const NodeDef& node) const;
private:
// Converts given device name to Lowercase Fully-Qualified Name (LFQN) string.
// This helps us disambiguate device names internally and simplify matching.
// If device_name couldn't be parsed successfully, returns empty string.
std::string to_lfqn_or_empty(const std::string& device_name) const;
// Map based on the cluster info: cluster device name -> device properties.
std::unordered_map<std::string, DeviceProperties> devices_;
// Maps LFQN to original device name as it was declared in cluster.
std::unordered_map<std::string, std::string> lfqn_map_;
std::string default_device_name_;
std::string default_job_name_lowercase_;
};
} // namespace grappler
} // end namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPPLER_COSTS_VIRTUAL_PLACER_H_
@@ -0,0 +1,331 @@
/* 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/grappler/costs/virtual_placer.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/grappler/clusters/virtual_cluster.h"
#include "tensorflow/core/lib/strings/strcat.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/protobuf/device_properties.pb.h"
namespace tensorflow {
namespace grappler {
TEST(VirtualPlacerTest, LocalDevices) {
// Create a virtual cluster with a local CPU and a local GPU
std::unordered_map<std::string, DeviceProperties> devices;
DeviceProperties cpu_device;
cpu_device.set_type("CPU");
devices["/job:localhost/replica:0/task:0/cpu:0"] = cpu_device;
DeviceProperties gpu_device;
gpu_device.set_type("GPU");
devices["/job:localhost/replica:0/task:0/device:GPU:0"] = gpu_device;
VirtualCluster cluster(devices);
VirtualPlacer placer(devices);
NodeDef node;
node.set_op("Conv2D");
// node.device() is empty, but GPU is default device if there is.
EXPECT_EQ("GPU", placer.get_device(node).type());
EXPECT_EQ("/job:localhost/replica:0/task:0/device:GPU:0",
placer.get_canonical_device_name(node));
node.set_device("CPU");
EXPECT_EQ("CPU", placer.get_device(node).type());
EXPECT_EQ("/job:localhost/replica:0/task:0/cpu:0",
placer.get_canonical_device_name(node));
node.set_device("GPU:0");
EXPECT_EQ("GPU", placer.get_device(node).type());
EXPECT_EQ("/job:localhost/replica:0/task:0/device:GPU:0",
placer.get_canonical_device_name(node));
}
TEST(VirtualPlacerTest, ShortNames) {
// Create a virtual cluster with a local CPU and a local GPU
std::unordered_map<std::string, DeviceProperties> devices;
DeviceProperties cpu_device;
cpu_device.set_type("CPU");
devices["/CPU:0"] = cpu_device;
DeviceProperties gpu_device;
gpu_device.set_type("GPU");
devices["/GPU:0"] = gpu_device;
VirtualCluster cluster(devices);
VirtualPlacer placer(devices);
NodeDef node;
node.set_op("Conv2D");
// node.device() is empty, but GPU is default device if there is.
EXPECT_EQ("GPU", placer.get_device(node).type());
EXPECT_EQ("/GPU:0", placer.get_canonical_device_name(node));
node.set_device("CPU");
EXPECT_EQ("CPU", placer.get_device(node).type());
EXPECT_EQ("/CPU:0", placer.get_canonical_device_name(node));
node.set_device("GPU:0");
EXPECT_EQ("GPU", placer.get_device(node).type());
EXPECT_EQ("/GPU:0", placer.get_canonical_device_name(node));
}
TEST(VirtualPlacerTest, PlacementOnNonDefaultDevice) {
// Create a virtual cluster with a CPU and a device:TPU
// Test that placement on TPU works
// In contrast with GPU, TPU is not selected as default device at the moment.
std::unordered_map<std::string, DeviceProperties> devices;
DeviceProperties cpu_device;
cpu_device.set_type("CPU");
devices["/job:localhost/replica:0/task:0/cpu:0"] = cpu_device;
DeviceProperties tpu_device;
tpu_device.set_type("TPU");
devices["/job:localhost/replica:0/task:0/device:TPU:0"] = tpu_device;
VirtualCluster cluster(devices);
VirtualPlacer placer(devices);
NodeDef node;
node.set_op("Conv2D");
// node.device() is empty, and CPU is default device.
EXPECT_EQ("CPU", placer.get_device(node).type());
EXPECT_EQ("/job:localhost/replica:0/task:0/cpu:0",
placer.get_canonical_device_name(node));
node.set_device("/device:TPU:0");
EXPECT_EQ("TPU", placer.get_device(node).type());
EXPECT_EQ("/job:localhost/replica:0/task:0/device:TPU:0",
placer.get_canonical_device_name(node));
}
TEST(VirtualPlacerTest, EmptyJobName) {
// Virtual placer choose job name from the devices in cluster if a device name
// of an op is empty. In case there are more than one kind of job name
// or job names are missing in the devices in cluster, we use local_host.
for (const std::string& job_name : {"localhost", "worker", "worker_train"}) {
std::unordered_map<std::string, DeviceProperties> devices;
DeviceProperties cpu_device;
cpu_device.set_type("CPU");
devices[absl::StrCat("/job:", job_name, "/replica:0/task:0/cpu:0")] =
cpu_device;
DeviceProperties gpu_device;
gpu_device.set_type("GPU");
devices[absl::StrCat("/job:", job_name, "/replica:0/task:0/device:GPU:0")] =
gpu_device;
VirtualCluster cluster(devices);
VirtualPlacer placer(devices);
NodeDef node;
node.set_op("Conv2D");
node.set_device("/device:CPU:0");
EXPECT_EQ(absl::StrCat("/job:", job_name, "/replica:0/task:0/cpu:0"),
placer.get_canonical_device_name(node));
node.set_device("/device:GPU:0");
EXPECT_EQ(absl::StrCat("/job:", job_name, "/replica:0/task:0/device:GPU:0"),
placer.get_canonical_device_name(node));
}
// When more than one job names are used, we use default "localhost"
// This may be improved later.
std::unordered_map<std::string, DeviceProperties> devices;
DeviceProperties cpu_device;
cpu_device.set_type("CPU");
devices["/job:localhost/replica:0/task:0/cpu:0"] = cpu_device;
devices["/job:ps/replica:0/task:0/cpu:0"] = cpu_device;
devices["/job:worker/replica:0/task:0/cpu:0"] = cpu_device;
VirtualCluster cluster(devices);
VirtualPlacer placer(devices);
NodeDef node;
node.set_op("Conv2D");
node.set_device("/device:CPU:0");
EXPECT_EQ("/job:localhost/replica:0/task:0/cpu:0",
placer.get_canonical_device_name(node));
}
std::string GetDefaultDeviceName(
const std::unordered_map<std::string, DeviceProperties>& devices) {
VirtualCluster cluster(devices);
VirtualPlacer placer(devices);
NodeDef node;
node.set_op("Conv2D");
// Device is not set to the node, so get_canonical_device_name() will return
// the default_device_.
return placer.get_canonical_device_name(node);
}
TEST(VirtualPlacerTest, DefaultDevice) {
std::unordered_map<std::string, DeviceProperties> devices;
DeviceProperties cpu_device;
cpu_device.set_type("CPU");
devices["/job:worker/replica:0/task:0/cpu:0"] = cpu_device;
// CPU is default when there is only CPU.
EXPECT_EQ("/job:worker/replica:0/task:0/cpu:0",
GetDefaultDeviceName(devices));
DeviceProperties gpu_device;
gpu_device.set_type("GPU");
// If there is any GPU, then gpu:0 is default device.
for (int i = 0; i < 8; i++) {
devices[absl::StrCat("/job:worker/replica:0/task:0/gpu:", i)] = gpu_device;
EXPECT_EQ("/job:worker/replica:0/task:0/gpu:0",
GetDefaultDeviceName(devices));
}
}
TEST(VirtualPlacerTest, MultiReplica) {
// Create a cluster with 8 workers, each with 8 GPUs.
std::unordered_map<std::string, DeviceProperties> devices;
DeviceProperties cpu_device;
cpu_device.set_type("CPU");
DeviceProperties gpu_device;
gpu_device.set_type("GPU");
for (int i = 0; i < 8; i++) {
devices[absl::StrCat("/job:worker/replica:", i, "/task:0/cpu:0")] =
cpu_device;
for (int j = 0; j < 8; j++) {
devices[absl::StrCat("/job:worker/replica:", i, "/task:0/gpu:", j)] =
gpu_device;
}
}
std::unique_ptr<VirtualCluster> cluster(new VirtualCluster(devices));
std::unique_ptr<VirtualPlacer> placer(new VirtualPlacer(devices));
auto get_device_name = [&placer](const std::string& device) -> std::string {
NodeDef node;
node.set_op("Conv2D");
node.set_device(device);
return placer->get_canonical_device_name(node);
};
// Validate device name is correct when we pass only replica ID and device
// name.
EXPECT_EQ("/job:worker/replica:0/task:0/cpu:0",
get_device_name("/replica:0/cpu:0"));
EXPECT_EQ("/job:worker/replica:2/task:0/cpu:0",
get_device_name("/replica:2/cpu:0"));
EXPECT_EQ("/job:worker/replica:7/task:0/cpu:0",
get_device_name("/replica:7/cpu:0"));
EXPECT_EQ("/job:worker/replica:3/task:0/gpu:0",
get_device_name("/replica:3/gpu:0"));
EXPECT_EQ("/job:worker/replica:5/task:0/gpu:3",
get_device_name("/replica:5/gpu:3"));
EXPECT_EQ("/job:worker/replica:4/task:0/gpu:7",
get_device_name("/replica:4/gpu:7"));
// Now add PS replicas; with multiple job names present in the cluster,
// device names in nodes should specify job names correctly.
for (int i = 0; i < 4; i++) {
devices[absl::StrCat("/job:ps/replica:", i, "/task:0/cpu:0")] = cpu_device;
}
cluster.reset(new VirtualCluster(devices));
placer.reset(new VirtualPlacer(cluster->GetDevices()));
EXPECT_EQ("/job:worker/replica:0/task:0/cpu:0",
get_device_name("/job:worker/replica:0/cpu:0"));
EXPECT_EQ("/job:worker/replica:7/task:0/gpu:3",
get_device_name("/job:worker/replica:7/gpu:3"));
EXPECT_EQ("/job:ps/replica:0/task:0/cpu:0",
get_device_name("/job:ps/replica:0/cpu:0"));
EXPECT_EQ("/job:ps/replica:1/task:0/cpu:0",
get_device_name("/job:ps/replica:1/cpu:0"));
EXPECT_EQ("/job:ps/replica:2/task:0/cpu:0",
get_device_name("/job:ps/replica:2/cpu:0"));
EXPECT_EQ("/job:ps/replica:3/task:0/cpu:0",
get_device_name("/job:ps/replica:3/cpu:0"));
}
TEST(VirtualPlacerTest, FallBackUnknown) {
// Virtual placer falls back to "UNKNOWN" only if there are no devices in the
// cluster.
std::unordered_map<std::string, DeviceProperties> devices;
VirtualCluster cluster(devices);
VirtualPlacer placer(devices);
NodeDef node;
node.set_op("Conv2D");
// Device falls back to UNKNOWN since the cluster has no devices.
EXPECT_EQ("UNKNOWN", placer.get_device(node).type());
EXPECT_EQ("UNKNOWN", placer.get_canonical_device_name(node));
}
TEST(VirtualPlacerTest, FallBackCPU) {
std::unordered_map<std::string, DeviceProperties> devices;
DeviceProperties cpu_device;
cpu_device.set_type("CPU");
devices["/job:my_job/replica:0/task:0/cpu:0"] = cpu_device;
VirtualCluster cluster(devices);
VirtualPlacer placer(devices);
NodeDef node;
node.set_op("Conv2D");
// Device falls back to CPU since there is no GPU.
EXPECT_EQ("CPU", placer.get_device(node).type());
EXPECT_EQ("/job:my_job/replica:0/task:0/cpu:0",
placer.get_canonical_device_name(node));
}
TEST(VirtualPlacerTest, RemoteDevices) {
std::unordered_map<std::string, DeviceProperties> devices;
DeviceProperties cpu_device;
cpu_device.set_type("CPU");
devices["/job:my_job/replica:0/task:0/cpu:0"] = cpu_device;
DeviceProperties gpu_device;
gpu_device.set_type("GPU");
devices["/job:my_job/replica:0/task:0/device:GPU:0"] = gpu_device;
VirtualCluster cluster(devices);
VirtualPlacer placer(devices);
NodeDef node;
node.set_op("Conv2D");
// Device falls back to GPU.
EXPECT_EQ("GPU", placer.get_device(node).type());
EXPECT_EQ("/job:my_job/replica:0/task:0/device:GPU:0",
placer.get_canonical_device_name(node));
node.set_device("/job:my_job/replica:0/task:0/cpu:0");
EXPECT_EQ("CPU", placer.get_device(node).type());
EXPECT_EQ("/job:my_job/replica:0/task:0/cpu:0",
placer.get_canonical_device_name(node));
node.set_device("/job:my_job/replica:0/task:0/device:GPU:0");
EXPECT_EQ("GPU", placer.get_device(node).type());
EXPECT_EQ("/job:my_job/replica:0/task:0/device:GPU:0",
placer.get_canonical_device_name(node));
// There is no local cpu available. Device falls back to GPU.
node.set_device("CPU");
EXPECT_EQ("GPU", placer.get_device(node).type());
EXPECT_EQ("/job:my_job/replica:0/task:0/device:GPU:0",
placer.get_canonical_device_name(node));
node.set_device("GPU:0");
// There is no local GPU available. Fall back to default GPU.
EXPECT_EQ("GPU", placer.get_device(node).type());
EXPECT_EQ("/job:my_job/replica:0/task:0/device:GPU:0",
placer.get_canonical_device_name(node));
// This isn't a valid name. Fall back to GPU.
node.set_device("/job:my_job/replica:0/task:0");
EXPECT_EQ("GPU", placer.get_device(node).type());
EXPECT_EQ("/job:my_job/replica:0/task:0/device:GPU:0",
placer.get_canonical_device_name(node));
}
} // end namespace grappler
} // end namespace tensorflow
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,545 @@
/* 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_GRAPPLER_COSTS_VIRTUAL_SCHEDULER_H_
#define TENSORFLOW_CORE_GRAPPLER_COSTS_VIRTUAL_SCHEDULER_H_
#include <functional>
#include <list>
#include <memory>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/step_stats.pb.h"
#include "tensorflow/core/grappler/costs/cost_estimator.h"
#include "tensorflow/core/grappler/costs/graph_properties.h"
#include "tensorflow/core/grappler/costs/op_context.h"
#include "tensorflow/core/grappler/costs/virtual_placer.h"
#include "tensorflow/core/grappler/grappler_item.h"
namespace tensorflow {
namespace grappler {
ABSL_CONST_INIT extern const char kAttrInputSrc[];
ABSL_CONST_INIT extern const char kAttrSrcDevice[];
ABSL_CONST_INIT extern const char kAttrDstDevice[];
ABSL_CONST_INIT extern const char kAttrTensorName[];
ABSL_CONST_INIT extern const char kChannelDevice[];
ABSL_CONST_INIT extern const char kStreaming[];
struct NodeState {
// A node (i.e., an op) takes a set of input:port pairs and produces
// a set of output ports.
// Cross references to input and output nodes from graphdef.
std::vector<std::pair<const NodeDef*, int>> inputs; // Input, port pairs.
// List of output nodes (a list of nodes that takes this output port as input)
// keyed by port_num. Note that port_num -1 is used for control dependency.
std::unordered_map<int, std::vector<const NodeDef*>> outputs;
// Info from GraphProperties.
std::vector<OpInfo::TensorProperties> input_properties;
std::vector<OpInfo::TensorProperties> output_properties;
// Canonical device name used within VirtualScheduler.
std::string device_name;
// States updated as scheduling nodes.
int num_inputs_ready;
std::unordered_map<int, int> num_outputs_executed;
Costs::Duration time_ready;
Costs::Duration time_scheduled;
Costs::Duration time_finished;
// Time that all the consumers are executed (hence, no need to keep this
// output in memory), keyed by port_num.
std::unordered_map<int, Costs::Duration> time_no_references;
// Note that a node may have multiple output ports. The length of outputs,
// num_outputs_executed, and time_no_references should be
// identical when a NodeState is fully initialized.
// They should be 1 + output_properties.size() as we add [-1] for control
// dependency.
// Node will be ready to be executed at time_ready, scheduled at
// time_scheduled, and finishes execution at time_finished.
// Each output port uses up memory space from time_scheduled to its
// time_no_references.
Costs node_costs; // Node costs per execution
Costs TotalNodeCosts() const {
return MultiplyCosts(node_costs, execution_count);
}
// How many times this node has been executed, e.g. in a while loop.
int execution_count;
// Output shape incompatible between shape annotation and shape inference.
bool shape_incompatible;
NodeState() {
num_inputs_ready = 0;
time_ready = Costs::Duration::max();
time_scheduled = Costs::Duration::max();
time_finished = Costs::Duration::max();
execution_count = 0;
shape_incompatible = false;
// Note that num_outputs_executed and time_no_references are not initialized
// here, since we don't know the size (i.e., # outputs for this node).
}
};
struct DeviceState {
// Nodes executed on this device in execution order.
std::vector<const NodeDef*> nodes_executed;
struct NodePairHash {
public:
const std::size_t operator()(
const std::pair<const NodeDef*, int>& element) const {
return std::hash<const NodeDef*>()(element.first);
}
};
// Nodes currently allocated in memory: set of NodeDef* and port_num pairs
// so that we can track which output of the node is in memory.
std::unordered_set<std::pair<const NodeDef*, int>, NodePairHash>
nodes_in_memory;
// Nodes allocated in memory persistently: e.g., Variables.
std::unordered_set<std::pair<const NodeDef*, int>, NodePairHash>
persistent_nodes;
// Snapshot of nodes_in_memory, when memory usage is at peak.
// Same to nodes_in_memory, it's a set of NodeDef* and port_num pairs.
std::unordered_set<std::pair<const NodeDef*, int>, NodePairHash>
mem_usage_snapshot_at_peak;
// Vector of temporary memory usage trace in execution order.
// Each pair represents the current node name and current (accumulated)
// temporary memory usage of the device when the node is scheduled.
// Only enabled when mem_usage_tracking is enabled.
// Note: CPU uses an inter-op threadpool, so the execution order on CPU may
// not be deterministic.
std::vector<std::pair<std::string, int64_t>> temporary_memory_usage_trace;
Costs device_costs;
std::map<std::string, Costs> op_to_cost; // Per-op cost.
int64_t memory_usage; // Current temporary memory usage
int64_t max_memory_usage; // Max temporary memory usage
// Shape annotation statistics.
struct ShapeAnnotationStats {
// Number of ops with shape annotated.
int64_t num_ops_annotated = 0;
// Number of ops executed multiple times (e.g. in a loop).
int64_t num_ops_executed_more_than_once = 0;
// Number of ops executed: account for execution count.
int64_t num_ops_executed = 0;
// Number of ops with dynamic shapes (e.g. shape changes in a loop).
int64_t num_ops_with_dynamic_shapes = 0;
// Number of ops with incompatible shapes between annotation and shape
// inference.
int64_t num_ops_with_incompatible_shapes = 0;
} shape_annotation_stats;
DeviceState() {
device_costs = Costs::ZeroCosts();
device_costs.num_ops_total = 0;
memory_usage = 0;
max_memory_usage = 0;
}
Costs::Duration GetCurrTime() const { return device_costs.execution_time; }
};
// ReadyNodeManager (abstract class):
// Keeps ready nodes and picks the best one to be scheduled.
class ReadyNodeManager {
public:
ReadyNodeManager() {}
virtual ~ReadyNodeManager() {}
virtual absl::Status Init(
const std::unordered_map<const NodeDef*, NodeState>* node_map) {
return absl::OkStatus();
}
virtual void AddNode(const NodeDef* node) = 0;
virtual const NodeDef* GetCurrNode() = 0;
virtual void RemoveCurrNode() = 0;
virtual bool Empty() const = 0;
};
class FIFOManager : public ReadyNodeManager {
public:
FIFOManager() : ReadyNodeManager() {}
~FIFOManager() override {}
void AddNode(const NodeDef* node) override { nodes_.push_back(node); }
const NodeDef* GetCurrNode() override {
CHECK(!nodes_.empty()) << "GetCurrNode(), but there's no ready node";
return nodes_.front();
}
void RemoveCurrNode() override { nodes_.pop_front(); }
bool Empty() const override { return nodes_.empty(); }
private:
std::list<const NodeDef*> nodes_;
};
// The LIFOManager schedules nodes by returning the last one added to the
// scheduler. A node is executed and then its ready outputs are newly added to
// the scheduler, so the LIFOManager will return outputs to a node following
// that node's execution.
class LIFOManager : public ReadyNodeManager {
public:
LIFOManager() : ReadyNodeManager() {}
~LIFOManager() override {}
void AddNode(const NodeDef* node) override;
const NodeDef* GetCurrNode() override;
void RemoveCurrNode() override;
bool Empty() const override { return nodes_.empty(); }
private:
std::list<const NodeDef*> nodes_;
// Keep track of the current node being executed by saving its position.
// Necessary because nodes may be added to the end of the list while a node is
// executing, and we want to remove the correct node (the one that is
// executing) rather than the new ones being added.
std::list<const NodeDef*>::iterator curr_pos_ = nodes_.end();
};
// Abstract class that maintains a heap/priority queue for scheduling ready
// nodes. Derived class needs to implement the Greater() function which returns
// the comparator for the heap.
class HeapReadyManager : public ReadyNodeManager {
public:
HeapReadyManager();
absl::Status Init(
const std::unordered_map<const NodeDef*, NodeState>* node_map) override;
~HeapReadyManager() override {}
void AddNode(const NodeDef* node) override;
const NodeDef* GetCurrNode() override;
void RemoveCurrNode() override;
bool Empty() const override;
protected:
virtual std::function<bool(const NodeDef*, const NodeDef*)> Greater() = 0;
// nodes_ is the main queue, where we construct heap, and the front is the
// current node.
std::vector<const NodeDef*> nodes_;
// Comparator functor for heap; stl heap is max heap, so we use "greater than"
// functor for keeping the smallest time_ready node at the front of heap.
std::function<bool(const NodeDef*, const NodeDef*)> greater_;
// NodeState structure from SchedulerState to get time_ready of ready nodes.
// Not owned by FirstReadyManager.
const std::unordered_map<const NodeDef*, NodeState>* node_map_;
// Cached curr node. Set back to nullptr from RemoveCurrNode().
const NodeDef* curr_node_;
};
// FirstReadyManager picks a node with the minimum time_ready value.
// Behavior is deterministic when there are more than one nodes with the minimum
// time_ready value with unique node names as the tie-breaker.
class FirstReadyManager : public HeapReadyManager {
public:
FirstReadyManager() : HeapReadyManager() {}
~FirstReadyManager() override {}
protected:
std::function<bool(const NodeDef*, const NodeDef*)> Greater() override;
};
// PriorityReadyManager uses the given node priorities when picking up next node
// from all the ready nodes.
class PriorityReadyManager : public HeapReadyManager {
public:
PriorityReadyManager() : HeapReadyManager() {}
~PriorityReadyManager() override {}
void AddNode(const NodeDef* node) override;
// Note this should be called after Init().
absl::Status SetPriority(
const std::unordered_map<std::string, int>& node_priority);
protected:
std::function<bool(const NodeDef*, const NodeDef*)> Greater() override;
private:
// A map from unique node name to priority. Lower number means higher
// priority.
std::unordered_map<std::string, int> node_priority_;
};
// CompositeNodeManager has a few other NodeManagers: per-device LIFO for normal
// ops (neither _Send nor _Recv) and FirstReadyManagers for _Send ops and _Recv
// ops, and then it chooses FirstReady among the ops chosen from each
// internal NodeManagers. The objective is to maximize producer-consumer
// locality within device, while processing nodes across devices, including
// _Send and _Recv, fairly, in terms of their time_ready.
class CompositeNodeManager : public ReadyNodeManager {
public:
CompositeNodeManager();
~CompositeNodeManager() override {}
absl::Status Init(
const std::unordered_map<const NodeDef*, NodeState>* node_map) override;
void AddNode(const NodeDef* node) override;
const NodeDef* GetCurrNode() override;
void RemoveCurrNode() override;
bool Empty() const override;
private:
// Internal ready node managers:
// LIFO for normal ops to maximize producer consumer locality.
// One LIFO per device.
std::unordered_map<std::string, LIFOManager> ops_lifo_map_;
// FirstReady for send and recv. Handle send and recv separately ensures that
// send and recv do not block previously read ops with LIFO schedule.
FirstReadyManager send_manager_;
FirstReadyManager recv_manager_;
// NodeState structure from SchedulerState to get time_ready of ready nodes.
// Not owned by CompositeReadyManager.
const std::unordered_map<const NodeDef*, NodeState>* node_map_;
// Cached curr node. Set back to nullptr from RemoveCurrNode().
const NodeDef* curr_node_;
};
// Constructs a ready node manager from the given string.
std::unique_ptr<ReadyNodeManager> ReadyNodeManagerFactory(
const std::string& ready_node_manager);
// Encapsulates all of the various pieces uses to track state of a scheduler;
// enables reuse of all scheduler state-related utilities across different
// scheduler implementations.
class SchedulerState {
public:
SchedulerState(const bool use_static_shapes,
const bool use_aggressive_shape_inference, Cluster* cluster,
std::unique_ptr<VirtualPlacer> placer);
// Move constructor. Explicitly defined because it otherwise gets implicitly
// deleted. SchedulerState is a move-only class, as we have a <unique_ptr>
// for it in VirtualScheduler. A derivative of VirtualScheduler can move a
// <unique_ptr> SchedulerState to VirtualScheduler when it is constructed,
// which is where this move constructor is needed.
SchedulerState(SchedulerState&& arg) = default;
// We explicitly delete assinment and copy operators, this is done implicitly,
// but we state it here explicitly for clarity.
SchedulerState& operator=(SchedulerState&& arg) = delete;
SchedulerState(const SchedulerState&) = delete;
SchedulerState& operator=(const SchedulerState&) = delete;
// Destructor. Must be defined such that a derivative class can override it
// and allow proper desctruction of the derivative class. If this is not done
// properly, memory leaks can occur.
virtual ~SchedulerState();
// Sets up the graph while also performing some necessary transformations
// initial_nodes is the set of nodes (primary inputs) discovered by Init()
// which may be added by a ReadyNodeManager (or related/derivative scheduler)
// to begin node schedule and graph simulation.
absl::Status Init(const GrapplerItem* item,
std::vector<const NodeDef*>* initial_nodes,
bool create_explicit_channel_device = true);
virtual Costs Summary() const;
// Like the above, but writes detailed stats to RunMetadata.
// If metadata is nullptr, then just calls and return Summary().
virtual Costs Summary(RunMetadata* metadata);
// Generates RunMetadata's step_stats and partition_graphs fields from results
// of the virtual execution of the graph.
// TODO(rdegruijl) See if we can make this function and caller Summary()
// const.
void GenerateRunMetadata(RunMetadata* metadata);
// Returns per device memory usage.
const std::unordered_map<std::string, int64_t> GetPeakMemoryUsage() const;
const std::unordered_map<std::string, int64_t> GetPersistentMemoryUsage()
const;
void enable_mem_usage_tracking() { track_mem_usage_snapshot_ = true; }
// Returns (read only) device and node states.
const std::unordered_map<std::string, DeviceState>* GetDeviceStates() const {
return &device_;
}
const std::unordered_map<const NodeDef*, NodeState>* GetNodeStates() const {
return &node_map_;
}
virtual OpContext CreateOpContext(const NodeDef* node) const;
std::vector<const NodeDef*> MarkNodeExecuted(
const NodeDef* node, const Costs& node_costs, const OpContext& op_context,
bool extract_execution_count_attr = true,
const std::string& override_device_name = "");
// Some getter functions.
const GrapplerItem* GetGrapplerItem() { return grappler_item_; }
Costs GetGraphCost() { return graph_costs_; }
Cluster* GetCluster() { return cluster_; }
bool GetUseStaticShape() { return use_static_shapes_; }
bool GetUseAggressiveShapeInference() {
return use_aggressive_shape_inference_;
}
const std::unordered_map<const NodeDef*, NodeState>& GetNodeMap() {
return node_map_;
}
protected:
// Assigns the time_scheduled in the NodeState of node to the current
// execution_time of the device executing this node.
void SetNodeStateTimeScheduled(const NodeDef* node);
// This method can be used by a class derived from SchedulerState to
// access the device state map.
std::unordered_map<std::string, DeviceState>* GetMutableDeviceState() {
return &device_;
}
private:
// Methods called from Init(). Fails if initialize_ is set.
void MaybeUpdateInputOutput(const NodeDef* node);
NodeState& GetNodeStateOrCreateIt(const NodeDef* node);
// Creates a Send_ and Recv_ pair between from and to. The argument
// create_channel_device tells the function to create an explicit device for
// the channel.
std::pair<const NodeDef*, const NodeDef*> CreateSendRecv(
const NodeDef* from, const NodeDef* to, const NodeDef* input_node,
const std::string& input_name, bool create_channel_device);
std::string DeviceName(const NodeDef* node) const;
std::string SanitizedDeviceName(const NodeDef* node) const;
std::string ChannelDeviceName(const NodeDef* from, const NodeDef* to) const;
// Helper methods.
void GetOutputNodes(const NodeDef* node, const Costs::Duration& curr_time,
std::vector<const NodeDef*>* output_nodes);
// Retrieves output size from node_cost at a port_num. If the output size has
// not been set, defaults back to CalculateOutputSize.
int64_t GetOrCalculateOutputSize(const NodeState& node_state,
int port_num) const;
std::unordered_map<const NodeDef*, NodeState> node_map_;
std::unordered_map<std::string, DeviceState> device_;
// Pool of NodeDefs for SendRecv and Identity ops created.
std::vector<std::unique_ptr<NodeDef>> additional_nodes_;
// Stats:
// Op counts with key with input shape.
// Example key: "[Op=AssignSub, input_shapes=[[7,1,160,160][7,1,160,160]]"
std::map<std::string, int> op_counts_;
// Individual op costs with key with input shape.
// Integer field for execution time in micro seconds.
// Boolean field for whether the cost is accurate.
std::map<std::string, std::pair<int, bool>> op_costs_;
Costs graph_costs_; // Graph cost.
std::map<std::string, Costs> op_to_cost_; // Per-op cost.
// Auxiliary data structures for constructing NodeState and DeviceState.
std::unique_ptr<GraphProperties> graph_properties_; // Initialized in Init().
Cluster* cluster_; // Not owned.
const GrapplerItem* grappler_item_; // Not owned.
bool use_static_shapes_;
bool initialized_;
bool track_mem_usage_snapshot_;
const bool use_aggressive_shape_inference_;
std::unique_ptr<VirtualPlacer> placer_;
};
// The virtual scheduler emulates execution of nodes in a graph, considering
// dependencies, device, etc.
class VirtualScheduler {
public:
// Does not take ownership of cluster or ready_nodes.
VirtualScheduler(const bool use_static_shapes,
const bool use_aggressive_shape_inference, Cluster* cluster,
ReadyNodeManager* ready_nodes,
std::unique_ptr<VirtualPlacer> placer);
// This constructor can be called by a derivative of VirtualScheduler to
// construct the base class. It lets VirtualScheduler take ownership of
// a new SchedulerState or a derivative thereof.
// Note that this constructor does not set a VirtualPlacer, in this
// constructor the VirtialPlacer is passed as a member of the SchedulerState
// that is passed as an argument.
VirtualScheduler(ReadyNodeManager* ready_nodes,
std::unique_ptr<SchedulerState> scheduler_state);
virtual ~VirtualScheduler();
// Initializes the scheduler for the specific grappler item.
// Should be called immediately after the c'tor or when the scheduler will be
// reused for a new grappler item. All internal states of the scheduler
// related to the previous grappler item will be reset/cleared.
//
// This function should be called at least once after the scheduler is
// constructed. An uninitialized or failed-to-initialize scheduler will cause
// undefined behavior.
virtual absl::Status Init(const GrapplerItem* item);
// Gets the current scheduled node for execution; the caller of this function
// can accordingly simulate the execution of the current scheduled node.
virtual OpContext GetCurrNode();
// Marks the current scheduled node as executed. Note that we should call this
// function only after the execution of the node has been simulated;
// node_costs_ capture the simulated costs of the node.
// Returns true if there is any node to be scheduled.
virtual bool MarkCurrNodeExecuted(const Costs& node_costs);
// Prints out summary of execution (timing, memory usage, etc.)
Costs Summary() const { return scheduler_state_->Summary(); }
// Like the above, but writes detailed stats to RunMetadata.
// If metadata is nullptr, then just calls and return Summary().
Costs Summary(RunMetadata* metadata) {
return scheduler_state_->Summary(metadata);
}
// Generates RunMetadata's step_stats and partition_graphs fields from results
// of the virtual execution of the graph.
void GenerateRunMetadata(RunMetadata* metadata) {
scheduler_state_->GenerateRunMetadata(metadata);
}
// Returns per device memory usage.
const std::unordered_map<std::string, int64_t> GetPeakMemoryUsage() const {
return scheduler_state_->GetPeakMemoryUsage();
}
const std::unordered_map<std::string, int64_t> GetPersistentMemoryUsage()
const {
return scheduler_state_->GetPersistentMemoryUsage();
}
// Returns VirtualScheduler (read only) device and node states.
const std::unordered_map<std::string, DeviceState>* GetDeviceStates() const {
return scheduler_state_->GetDeviceStates();
}
const std::unordered_map<const NodeDef*, NodeState>* GetNodeStates() const {
return scheduler_state_->GetNodeStates();
}
void enable_mem_usage_tracking() {
scheduler_state_->enable_mem_usage_tracking();
}
protected:
// The state of the scheduler and the execution of the graph is encapsulated
// by the scheduler_state_ object.
std::unique_ptr<SchedulerState> scheduler_state_;
// ready_nodes_ is responsible for ordering the traversal of the graph.
ReadyNodeManager* ready_nodes_; // Not owned.
};
} // namespace grappler
} // end namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPPLER_COSTS_VIRTUAL_SCHEDULER_H_
File diff suppressed because it is too large Load Diff
+103
View File
@@ -0,0 +1,103 @@
/* 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/grappler/devices.h"
#include <memory>
#include "absl/log/log.h"
#include "tensorflow/core/platform/cpu_info.h"
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#include "xla/stream_executor/gpu/gpu_init.h"
#include "tensorflow/core/platform/stream_executor.h"
#endif // GOOGLE_CUDA
namespace tensorflow {
namespace grappler {
int GetNumAvailableGPUs(
const std::pair<int, int>& min_cuda_compute_capability) {
int num_eligible_gpus = 0;
#if TENSORFLOW_USE_ROCM
if (min_cuda_compute_capability.first != 0 ||
min_cuda_compute_capability.second != 0) {
LOG(ERROR) << "GetNumAvailableGPUs() should receive zero "
"min_cuda_compute_capability";
return 0;
}
#endif
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
if (se::ValidateGPUMachineManager().ok()) {
se::Platform* gpu_manager = se::GPUMachineManager();
if (gpu_manager != nullptr) {
int num_gpus = gpu_manager->VisibleDeviceCount();
for (int i = 0; i < num_gpus; i++) {
#if GOOGLE_CUDA
auto desc = gpu_manager->DescriptionForDevice(i);
if (desc.ok()) {
int min_gpu_core_count = 8;
if ((*desc)->core_count() >= min_gpu_core_count &&
(*desc)->cuda_compute_capability().IsAtLeast(
min_cuda_compute_capability.first,
min_cuda_compute_capability.second)) {
num_eligible_gpus++;
}
}
#else
num_eligible_gpus++;
#endif
}
}
}
#if GOOGLE_CUDA
LOG(INFO)
<< "Number of eligible GPUs (core count >= 8, compute capability >= "
<< min_cuda_compute_capability.first << "."
<< min_cuda_compute_capability.second << "): " << num_eligible_gpus;
#else
LOG(INFO) << "Number of eligible GPUs: " << num_eligible_gpus;
#endif
#else // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
LOG(INFO)
<< "Number of eligible GPUs (core count >= 8, compute capability >= "
<< min_cuda_compute_capability.first << "."
<< min_cuda_compute_capability.second << "): " << num_eligible_gpus
<< " (Note: TensorFlow was not compiled with CUDA or ROCm support)";
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
return num_eligible_gpus;
}
int64_t AvailableGPUMemory(int gpu_id) {
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
// Look up the device, to see its attributes.
se::Platform* gpu_platform = se::GPUMachineManager();
CHECK_LT(gpu_id, gpu_platform->VisibleDeviceCount());
se::StreamExecutor* se = gpu_platform->ExecutorForDevice(gpu_id).value();
int64_t total_memory, available_memory;
CHECK(se->DeviceMemoryUsage(&available_memory, &total_memory));
return available_memory;
#else
return 0;
#endif
}
int GetNumAvailableLogicalCPUCores() { return port::NumSchedulableCPUs(); }
} // end namespace grappler
} // end namespace tensorflow
+46
View File
@@ -0,0 +1,46 @@
/* 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_GRAPPLER_DEVICES_H_
#define TENSORFLOW_CORE_GRAPPLER_DEVICES_H_
#include <cstdint>
#include <functional>
#include <utility>
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/core/threadpool.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace grappler {
// Get the number of available GPUs whose number of multiprocessors is no less
// than 8 and whose CUDA compute capability is no less than
// min_cuda_compute_capability.
int GetNumAvailableGPUs(
const std::pair<int, int>& min_cuda_compute_capability = {0, 0});
// Maximum amount of gpu memory available per gpu. gpu_id must be in the range
// [0, num_available_gpu)
int64_t AvailableGPUMemory(int gpu_id);
// Get the number of logical CPU cores (aka hyperthreads) available.
int GetNumAvailableLogicalCPUCores();
} // end namespace grappler
} // end namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPPLER_DEVICES_H_
@@ -0,0 +1,163 @@
load("//tensorflow:tensorflow.bzl", "tf_cc_test")
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
licenses = ["notice"],
)
cc_library(
name = "graph_analyzer_lib",
srcs = [
"gen_node.cc",
"graph_analyzer.cc",
"sig_node.cc",
"subgraph.cc",
],
hdrs = [
"gen_node.h",
"graph_analyzer.h",
"hash_tools.h",
"map_tools.h",
"sig_node.h",
"subgraph.h",
],
visibility = ["//visibility:public"],
deps = [
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:tensorflow",
"//tensorflow/core/grappler:op_types",
"//tensorflow/core/grappler:utils",
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:str_format",
],
)
cc_library(
name = "graph_analyzer_tool",
srcs = ["graph_analyzer_tool.cc"],
hdrs = ["graph_analyzer_tool.h"],
visibility = ["//visibility:public"],
deps = [
":graph_analyzer_lib",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:tensorflow",
"//tensorflow/core/grappler/utils:transitive_fanin",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
],
)
cc_library(
name = "test_tools_lib",
testonly = 1,
srcs = [
"test_tools.cc",
],
hdrs = [
"test_tools.h",
],
visibility = ["//visibility:public"],
deps = [
":graph_analyzer_lib",
"//tensorflow/core:framework",
"//tensorflow/core/framework:graph_proto_cc",
"//tensorflow/core/framework:node_def_proto_cc",
"//tensorflow/core/grappler:op_types",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:str_format",
],
)
tf_cc_test(
name = "hash_tools_test",
testonly = 1,
srcs = [
"hash_tools_test.cc",
],
deps = [
":graph_analyzer_lib",
"@com_google_googletest//:gtest_main",
],
)
tf_cc_test(
name = "gen_node_test",
testonly = 1,
srcs = [
"gen_node_test.cc",
],
deps = [
":graph_analyzer_lib",
":test_tools_lib",
"//tensorflow/core/framework:graph_proto_cc",
"//tensorflow/core/framework:node_def_proto_cc",
"//tensorflow/core/framework:op_def_proto_cc",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_googletest//:gtest_main",
],
)
tf_cc_test(
name = "sig_node_test",
testonly = 1,
srcs = [
"sig_node_test.cc",
],
deps = [
":graph_analyzer_lib",
":test_tools_lib",
"//tensorflow/core/framework:graph_proto_cc",
"//tensorflow/core/framework:node_def_proto_cc",
"//tensorflow/core/grappler:utils",
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings:str_format",
"@com_google_googletest//:gtest_main",
],
)
tf_cc_test(
name = "graph_analyzer_test",
testonly = 1,
srcs = [
"graph_analyzer_test.cc",
],
deps = [
":graph_analyzer_lib",
":test_tools_lib",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/status",
"@com_google_googletest//:gtest_main",
],
)
tf_cc_test(
name = "subgraph_test",
testonly = 1,
srcs = [
"subgraph_test.cc",
],
deps = [
":graph_analyzer_lib",
":test_tools_lib",
"//tensorflow/core/framework:graph_proto_cc",
"//tensorflow/core/framework:node_def_proto_cc",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings:str_format",
"@com_google_googletest//:gtest_main",
],
)
@@ -0,0 +1,154 @@
/* 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/grappler/graph_analyzer/gen_node.h"
#include <memory>
#include <string>
#include "absl/status/status.h"
#include "absl/strings/str_format.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/grappler/graph_analyzer/hash_tools.h"
#include "tensorflow/core/grappler/op_types.h"
#include "tensorflow/core/grappler/utils.h"
namespace tensorflow {
namespace grappler {
namespace graph_analyzer {
GenNode::GenNode(const NodeDef* node) : node_(node), op_(nullptr) {}
absl::Status GenNode::BuildGraphInMap(const GraphDef& source, GenNodeMap* map) {
for (const auto& n : source.node()) {
const std::string& name = n.name();
if (map->find(name) != map->end()) {
// This error code looks more meaningful than ALREADY_EXISTS.
return absl::Status(absl::StatusCode::kInvalidArgument,
"Duplicate node name '" + name + "'.");
}
(*map)[name] = std::make_unique<GenNode>(&n);
}
// Now parse the links.
for (const auto& mapit : *map) {
absl::Status st = mapit.second->ParseInputs(map);
if (!st.ok()) {
return st;
}
}
return absl::OkStatus();
}
absl::Status GenNode::ParseInputs(const GenNodeMap* map) {
all_inputs_or_none_ = false;
absl::Status st = OpRegistry::Global()->LookUpOpDef(opcode(), &op_);
if (!st.ok()) {
return absl::Status(
absl::StatusCode::kInvalidArgument,
absl::StrFormat("Node '%s' contains an undefined operation '%s': %s",
name(), opcode(), st.message()));
}
int n_inputs = node_->input_size();
int n_named_inputs = op_->input_arg_size();
int n_multi_inputs = 0;
for (const auto& inarg : op_->input_arg()) {
if (!inarg.number_attr().empty() || !inarg.type_list_attr().empty()) {
++n_multi_inputs;
}
}
bool is_commutative = grappler::IsCommutative(*node_);
if (n_multi_inputs > 1 || (n_multi_inputs > 0 && n_named_inputs > 1)) {
// Can't handle more than one multi-input at a time.
// And can't handle the commutativeness of only some arguments
// rather than all of them.
is_commutative = false;
}
if (is_commutative) {
// If truly commutative, can treat all the inputs as one multi-input.
// It's possible to just treat the commutative nodes as AllInputsOrNone
// but (1) this way is a bit more efficient and (2) I want to preserve this
// more efficient code path that does all-or-none by a single input and
// perhaps extend its use in the future.
n_named_inputs = 1;
all_inputs_or_none_ = false;
} else if (n_multi_inputs > 0) {
all_inputs_or_none_ = true;
}
for (int i = 0; i < n_inputs; ++i) {
int other_position;
std::string other_name = ParseNodeName(node_->input(i), &other_position);
auto other_it = map->find(other_name);
if (other_it == map->end()) {
return absl::Status(
absl::StatusCode::kInvalidArgument,
absl::StrFormat(
"Node '%s' input %d refers to a non-existing node '%s'.", name(),
i, other_name));
}
GenNode* other_node = other_it->second.get();
int this_position = other_position < 0 ? -1 : (is_commutative ? 0 : i);
if (this_position >= 0 && n_multi_inputs == 0 &&
this_position >= n_named_inputs) {
return absl::Status(
absl::StatusCode::kInvalidArgument,
absl::StrFormat(
"Node '%s' has a non-control input from '%s' at index %d but its "
"operation '%s' defines only %d inputs.",
name(), other_name, this_position, op_->name(), n_named_inputs));
}
Port this_port(/*inbound=*/true, this_position);
Port other_port(/*inbound=*/false, other_position);
links_[this_port].emplace_back(LinkTarget(other_node, other_port));
other_node->links_[other_port].emplace_back(LinkTarget(this, this_port));
}
return absl::OkStatus();
}
bool GenNode::IsMultiInput(Port port) const {
if (!port.IsInbound()) {
return false;
}
auto it = links_.find(port);
if (it == links_.end()) {
return false; // Shouldn't happen.
}
return (it->second.size() > 1);
}
GenNode::Port::operator std::string() const {
std::string result = this->IsInbound() ? "i" : "o";
if (this->IsControl()) {
result.append("C");
} else {
result.append(absl::StrFormat("%d", this->Id()));
}
return result;
}
} // end namespace graph_analyzer
} // end namespace grappler
} // end namespace tensorflow
@@ -0,0 +1,169 @@
/* 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_GRAPPLER_GRAPH_ANALYZER_GEN_NODE_H_
#define TENSORFLOW_CORE_GRAPPLER_GRAPH_ANALYZER_GEN_NODE_H_
#include <map>
#include <memory>
#include <unordered_map>
#include <vector>
#include "absl/status/status.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/op_def.pb.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
namespace tensorflow {
namespace grappler {
namespace graph_analyzer {
class GenNode;
// To find nodes by name.
using GenNodeMap = std::unordered_map<std::string, std::unique_ptr<GenNode>>;
// One node in the graph, in the form convenient for traversal and generation of
// subgraphs. It refers to the original NodeDef protobuf for most information
// and adds the extra enrichment.
//
// The graph building is 2-stage: first match a GenNode with each NodeDef and
// collect them into a map that finds them by name, then process the map,
// deep-parse the underlying NodeDefs and connect the GenNodes together.
class GenNode {
public:
// Will keep the pointer, so the underlying object must not be deleted while
// GenNode is alive.
explicit GenNode(const NodeDef* node);
// Access wrappers.
const std::string& name() const { return node_->name(); }
const std::string& opcode() const { return node_->op(); }
const NodeDef* node_def() const { return node_; }
// Parse the inputs of this node and update the map accordingly, creating the
// links (i.e. edges, connections between nodes) in itself and in the nodes
// it's linked to (the map itself is unchanged, only the nodes in it are
// updated).
absl::Status ParseInputs(const GenNodeMap* map);
// Does the full 2-stage build of the graph. The map should be initially
// empty. The map keeps pointers to the nodes in source, so the source must
// not be destroyed before the map.
static absl::Status BuildGraphInMap(const GraphDef& source, GenNodeMap* map);
// The enrichment that constitutes the point of this class.
// Representation of a connection on a node.
class Port {
public:
// A port may be inbound or outbound.
// Negative ids (canonically -1) mean a control port.
Port(bool inbound, int32_t id) : value_(id << 1) {
if (inbound) {
value_ |= 1;
}
}
Port(const Port&) = default;
Port& operator=(const Port&) = default;
bool IsInbound() const { return (value_ & 0x1); }
bool IsControl() const { return (value_ < 0); }
int32_t Id() const {
// Arithmetic shift preserves the sign.
return (value_ >> 1);
}
// Integer type used to represent the encoded port value.
using IntPort = int32_t;
// Returns the encoded form of this port, so that it can be used
// as various map indexes.
IntPort Encoded() const { return value_; }
static Port Decode(IntPort encoded) { return Port(encoded); }
bool operator==(const Port& other) const { return value_ == other.value_; }
bool operator<(const Port& other) const { return value_ < other.value_; }
struct Hasher {
size_t operator()(const Port& port) const noexcept {
return hasher(port.Encoded());
}
std::hash<int32_t> hasher;
};
// Convenient for printing. I've really wanted it to be implicit but
// ClangTidy insists on making it explicit.
explicit operator std::string() const;
private:
explicit Port(IntPort value) : value_(value) {}
IntPort value_;
};
struct LinkTarget {
GenNode* node; // Node where this link points.
Port port; // Port on the remote side of this link.
LinkTarget(GenNode* a_node, Port a_port) : node(a_node), port(a_port) {}
};
// All the links that are connected to the same port of this node
// are collected in one vector. A link is an edge of the graph that connects
// 2 nodes. Each of the connected nodes has its own perspective on the link,
// seeing its local port, remote port and the remote node. The direction of
// the link is encoded in the ports, one port is always incoming and another
// one outgoing.
using LinkTargetVector = std::vector<LinkTarget>;
// Both inputs and outputs are stored in the same map.
using LinkMap = std::unordered_map<Port, LinkTargetVector, Port::Hasher>;
// Access to the link map.
const LinkMap& links() const { return links_; }
// Check whether the port is an input (including the controls) with multiple
// connections. Such inputs get handled in a special way when building the
// subgraphs, in an "all or nothing" fashion.
bool IsMultiInput(Port port) const;
// When building the subgraphs, must include either all non-control inputs of
// this node into the subgraph or none of them. This happens when at least one
// of the inputs is a multi-input (or if the opcode is commutative, thus
// treating all the inputs as one multi-input).
bool AllInputsOrNone() const { return all_inputs_or_none_; }
private:
const NodeDef* node_;
// Becomes valid only after ParseInputs().
const OpDef* op_;
// The opcode has a complicated structure of input args, with multi-input args
// that are not commutative. This means that to make sense, the subgraphs that
// include this node must also include either all its inputs or none of them.
bool all_inputs_or_none_ = false;
LinkMap links_;
};
} // end namespace graph_analyzer
} // end namespace grappler
} // end namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPPLER_GRAPH_ANALYZER_GEN_NODE_H_
@@ -0,0 +1,498 @@
/* 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/grappler/graph_analyzer/gen_node.h"
#include <memory>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/op_def.pb.h"
#include "tensorflow/core/grappler/graph_analyzer/test_tools.h"
namespace tensorflow {
namespace grappler {
namespace graph_analyzer {
namespace test {
namespace {
using ::testing::ElementsAre;
using ::testing::Eq;
using ::testing::Ne;
TEST(GenNodeTest, Port) {
{
GenNode::Port p(true, 100);
EXPECT_THAT(p.IsInbound(), Eq(true));
EXPECT_THAT(p.IsControl(), Eq(false));
EXPECT_THAT(p.Id(), Eq(100));
GenNode::Port p2 = GenNode::Port::Decode(p.Encoded());
EXPECT_THAT(p2.IsInbound(), Eq(true));
EXPECT_THAT(p2.IsControl(), Eq(false));
EXPECT_THAT(p2.Id(), Eq(100));
}
{
GenNode::Port p(false, 0);
EXPECT_THAT(p.IsInbound(), Eq(false));
EXPECT_THAT(p.IsControl(), Eq(false));
EXPECT_THAT(p.Id(), Eq(0));
GenNode::Port p2 = GenNode::Port::Decode(p.Encoded());
EXPECT_THAT(p2.IsInbound(), Eq(false));
EXPECT_THAT(p2.IsControl(), Eq(false));
EXPECT_THAT(p2.Id(), Eq(0));
}
{
GenNode::Port p(true, -100);
EXPECT_THAT(p.IsInbound(), Eq(true));
EXPECT_THAT(p.IsControl(), Eq(true));
EXPECT_THAT(p.Id(), Eq(-100));
GenNode::Port p2 = GenNode::Port::Decode(p.Encoded());
EXPECT_THAT(p2.IsInbound(), Eq(true));
EXPECT_THAT(p2.IsControl(), Eq(true));
EXPECT_THAT(p2.Id(), Eq(-100));
}
{
GenNode::Port p(false, -1);
EXPECT_THAT(p.IsInbound(), Eq(false));
EXPECT_THAT(p.IsControl(), Eq(true));
EXPECT_THAT(p.Id(), Eq(-1));
GenNode::Port p2 = GenNode::Port::Decode(p.Encoded());
EXPECT_THAT(p2.IsInbound(), Eq(false));
EXPECT_THAT(p2.IsControl(), Eq(true));
EXPECT_THAT(p2.Id(), Eq(-1));
}
}
TEST(GenNodeTest, ParseNodeNoInputs) {
GenNodeMap map;
NodeDef node1 = MakeNodeConst("node1");
map["node1"] = std::make_unique<GenNode>(&node1);
auto gn1 = map["node1"].get();
ASSERT_THAT(gn1->ParseInputs(&map), Eq(absl::OkStatus()));
EXPECT_THAT(DumpLinkMap(gn1->links()), ElementsAre());
}
// A general operation, and a control link.
TEST(GenNodeTest, ParseNodeWithControl) {
GenNodeMap map;
NodeDef node1 = MakeNodeConst("node1");
map["node1"] = std::make_unique<GenNode>(&node1);
NodeDef node2 = MakeNodeConst("node2");
map["node2"] = std::make_unique<GenNode>(&node2);
NodeDef node3 = MakeNodeSub("node3", "node1", "node2");
node3.add_input("^node1"); // The control link.
node3.add_input("^node2"); // The control link.
map["node3"] = std::make_unique<GenNode>(&node3);
auto gn1 = map["node1"].get();
auto gn2 = map["node2"].get();
auto gn3 = map["node3"].get();
ASSERT_THAT(gn3->ParseInputs(&map), Eq(absl::OkStatus()));
// clang-format off
EXPECT_THAT(DumpLinkMap(gn1->links()), ElementsAre(
"o0: node3[i0]",
"oC: node3[iC]"
));
EXPECT_THAT(DumpLinkMap(gn2->links()), ElementsAre(
"o0: node3[i1]",
"oC: node3[iC]"
));
EXPECT_THAT(DumpLinkMap(gn3->links()), ElementsAre(
"i0: node1[o0]",
"i1: node2[o0]",
"iC: node1[oC], node2[oC]"
));
// clang-format on
EXPECT_THAT(gn3->IsMultiInput(GenNode::Port(true, 0)), Eq(false));
// This is a multi-control-input.
EXPECT_THAT(gn3->IsMultiInput(GenNode::Port(true, -1)), Eq(true));
EXPECT_FALSE(gn1->AllInputsOrNone());
EXPECT_FALSE(gn2->AllInputsOrNone());
EXPECT_FALSE(gn3->AllInputsOrNone());
}
// Commutative nodes are treated as having a single input,
// because their inputs are equivalent.
TEST(GenNodeTest, ParseNodeCommutative) {
GenNodeMap map;
NodeDef node1 = MakeNodeConst("node1");
map["node1"] = std::make_unique<GenNode>(&node1);
NodeDef node2 = MakeNodeConst("node2");
map["node2"] = std::make_unique<GenNode>(&node2);
// TODO(babkin): grappler::IsCommutative() should return true for Add but
// apparently doesn't. So use Mul in the meantime.
NodeDef node3 = MakeNodeMul("node3", "node1", "node2");
map["node3"] = std::make_unique<GenNode>(&node3);
auto gn1 = map["node1"].get();
auto gn2 = map["node2"].get();
auto gn3 = map["node3"].get();
ASSERT_THAT(gn3->ParseInputs(&map), Eq(absl::OkStatus()));
// clang-format off
EXPECT_THAT(DumpLinkMap(gn1->links()), ElementsAre(
"o0: node3[i0]"
));
EXPECT_THAT(DumpLinkMap(gn2->links()), ElementsAre(
"o0: node3[i0]"
));
EXPECT_THAT(DumpLinkMap(gn3->links()), ElementsAre(
"i0: node1[o0], node2[o0]"
));
// clang-format on
EXPECT_THAT(gn3->IsMultiInput(GenNode::Port(true, 0)), Eq(true));
EXPECT_FALSE(gn3->AllInputsOrNone());
}
TEST(GenNodeTest, ParseNodeMultiInputCommutative) {
GenNodeMap map;
NodeDef node1 = MakeNodeConst("node1");
map["node1"] = std::make_unique<GenNode>(&node1);
NodeDef node2 = MakeNodeConst("node2");
map["node2"] = std::make_unique<GenNode>(&node2);
NodeDef node3 = MakeNodeAddN("node3", "node1", "node2");
map["node3"] = std::make_unique<GenNode>(&node3);
auto gn1 = map["node1"].get();
auto gn2 = map["node2"].get();
auto gn3 = map["node3"].get();
ASSERT_THAT(gn3->ParseInputs(&map), Eq(absl::OkStatus()));
// clang-format off
EXPECT_THAT(DumpLinkMap(gn1->links()), ElementsAre(
"o0: node3[i0]"
));
EXPECT_THAT(DumpLinkMap(gn2->links()), ElementsAre(
"o0: node3[i0]"
));
EXPECT_THAT(DumpLinkMap(gn3->links()), ElementsAre(
"i0: node1[o0], node2[o0]"
));
// clang-format on
// This is a multi-output.
EXPECT_THAT(gn2->IsMultiInput(GenNode::Port(false, 0)), Eq(false));
// This is a multi-input.
EXPECT_THAT(gn3->IsMultiInput(GenNode::Port(true, 0)), Eq(true));
EXPECT_FALSE(gn3->AllInputsOrNone());
}
TEST(GenNodeTest, ParseNodeMultiInputNotCommutative) {
GenNodeMap map;
NodeDef node1 = MakeNodeConst("node1");
map["node1"] = std::make_unique<GenNode>(&node1);
NodeDef node2 = MakeNodeConst("node2");
map["node2"] = std::make_unique<GenNode>(&node2);
NodeDef node3 = MakeNodeShapeN("node3", "node1", "node2");
map["node3"] = std::make_unique<GenNode>(&node3);
auto gn1 = map["node1"].get();
auto gn2 = map["node2"].get();
auto gn3 = map["node3"].get();
ASSERT_THAT(gn3->ParseInputs(&map), Eq(absl::OkStatus()));
// clang-format off
EXPECT_THAT(DumpLinkMap(gn1->links()), ElementsAre(
"o0: node3[i0]"
));
EXPECT_THAT(DumpLinkMap(gn2->links()), ElementsAre(
"o0: node3[i1]"
));
EXPECT_THAT(DumpLinkMap(gn3->links()), ElementsAre(
"i0: node1[o0]",
"i1: node2[o0]"
));
// clang-format on
// Non-commutative multi-input doesn't count.
EXPECT_THAT(gn3->IsMultiInput(GenNode::Port(true, 0)), Eq(false));
EXPECT_TRUE(gn3->AllInputsOrNone());
}
TEST(GenNodeTest, ParseNodeMultiInputList) {
GenNodeMap map;
NodeDef node1 = MakeNodeConst("node1");
map["node1"] = std::make_unique<GenNode>(&node1);
NodeDef node2 = MakeNodeConst("node2");
map["node2"] = std::make_unique<GenNode>(&node2);
NodeDef node3 = MakeNodeIdentityN("node3", "node1", "node2");
map["node3"] = std::make_unique<GenNode>(&node3);
auto gn1 = map["node1"].get();
auto gn2 = map["node2"].get();
auto gn3 = map["node3"].get();
ASSERT_THAT(gn3->ParseInputs(&map), Eq(absl::OkStatus()));
// clang-format off
EXPECT_THAT(DumpLinkMap(gn1->links()), ElementsAre(
"o0: node3[i0]"
));
EXPECT_THAT(DumpLinkMap(gn2->links()), ElementsAre(
"o0: node3[i1]"
));
EXPECT_THAT(DumpLinkMap(gn3->links()), ElementsAre(
"i0: node1[o0]",
"i1: node2[o0]"
));
// clang-format on
// Non-commutative multi-input doesn't count.
EXPECT_THAT(gn3->IsMultiInput(GenNode::Port(true, 0)), Eq(false));
EXPECT_TRUE(gn3->AllInputsOrNone());
}
TEST(GenNodeTest, ParseNodeMultiMultiInput) {
GenNodeMap map;
NodeDef node1 = MakeNodeConst("node1");
map["node1"] = std::make_unique<GenNode>(&node1);
NodeDef node2 = MakeNodeConst("node2");
map["node2"] = std::make_unique<GenNode>(&node2);
NodeDef node3 = MakeNodeConst("node3");
map["node3"] = std::make_unique<GenNode>(&node3);
NodeDef node4 = MakeNodeConst("node4");
map["node4"] = std::make_unique<GenNode>(&node4);
NodeDef node5 =
MakeNodeQuantizedConcat("node5", "node1", "node2", "node3", "node4");
map["node5"] = std::make_unique<GenNode>(&node5);
auto gn1 = map["node1"].get();
auto gn2 = map["node2"].get();
auto gn3 = map["node3"].get();
auto gn4 = map["node4"].get();
auto gn5 = map["node5"].get();
ASSERT_THAT(gn5->ParseInputs(&map), Eq(absl::OkStatus()));
// clang-format off
EXPECT_THAT(DumpLinkMap(gn1->links()), ElementsAre(
"o0: node5[i0]"
));
EXPECT_THAT(DumpLinkMap(gn2->links()), ElementsAre(
"o0: node5[i1]"
));
EXPECT_THAT(DumpLinkMap(gn3->links()), ElementsAre(
"o0: node5[i2]"
));
EXPECT_THAT(DumpLinkMap(gn4->links()), ElementsAre(
"o0: node5[i3]"
));
EXPECT_THAT(DumpLinkMap(gn5->links()), ElementsAre(
"i0: node1[o0]",
"i1: node2[o0]",
"i2: node3[o0]",
"i3: node4[o0]"
));
// clang-format on
// Non-commutative multi-input doesn't count.
EXPECT_THAT(gn5->IsMultiInput(GenNode::Port(true, 1)), Eq(false));
EXPECT_THAT(gn5->IsMultiInput(GenNode::Port(true, 2)), Eq(false));
EXPECT_TRUE(gn5->AllInputsOrNone());
}
TEST(GenNodeTest, ParseNodeMultiOutput) {
GenNodeMap map;
NodeDef node1 = MakeNodeConst("node1");
map["node1"] = std::make_unique<GenNode>(&node1);
NodeDef node2 = MakeNodeConst("node2");
map["node2"] = std::make_unique<GenNode>(&node2);
NodeDef node3 = MakeNodeBroadcastGradientArgs("node3", "node1", "node2");
map["node3"] = std::make_unique<GenNode>(&node3);
NodeDef node4 = MakeNodeSub("node4", "node3:1", "node3:0");
map["node4"] = std::make_unique<GenNode>(&node4);
auto gn4 = map["node4"].get();
ASSERT_THAT(gn4->ParseInputs(&map), Eq(absl::OkStatus()));
// clang-format off
EXPECT_THAT(DumpLinkMap(gn4->links()), ElementsAre(
"i0: node3[o1]",
"i1: node3[o0]"
));
// clang-format on
}
TEST(GenNodeTest, ParseNodeUndefinedOp) {
GenNodeMap map;
NodeDef node1;
node1.set_name("node1");
node1.set_op("Zzzx");
map["node1"] = std::make_unique<GenNode>(&node1);
const OpDef* opdef;
absl::Status nested_error = OpRegistry::Global()->LookUpOpDef("Zzzx", &opdef);
auto gn = map["node1"].get();
ASSERT_THAT(
gn->ParseInputs(&map),
Eq(absl::Status(
absl::StatusCode::kInvalidArgument,
absl::StrCat("Node 'node1' contains an undefined operation 'Zzzx': ",
nested_error.message()))));
}
TEST(GenNodeTest, ParseNodeUnexpectedInputs) {
GenNodeMap map;
NodeDef node1 = MakeNodeConst("node1");
map["node1"] = std::make_unique<GenNode>(&node1);
node1.add_input("node1");
auto gn1 = map["node1"].get();
EXPECT_THAT(gn1->ParseInputs(&map),
Eq(absl::Status(absl::StatusCode::kInvalidArgument,
"Node 'node1' has a non-control "
"input from 'node1' at index 0 but its operation "
"'Const' defines only 0 inputs.")));
NodeDef node2 = MakeNodeConst("node2");
map["node2"] = std::make_unique<GenNode>(&node2);
NodeDef node3 = MakeNodeSub("node3", "node1", "node2");
map["node3"] = std::make_unique<GenNode>(&node3);
node3.add_input("node1");
auto gn3 = map["node3"].get();
EXPECT_THAT(gn3->ParseInputs(&map),
Eq(absl::Status(absl::StatusCode::kInvalidArgument,
"Node 'node3' has a non-control "
"input from 'node1' at index 2 but its operation "
"'Sub' defines only 2 inputs.")));
}
// Even if an opcode defines no inputs, the node may still accept the control
// inputs.
TEST(GenNodeTest, ParseNodeControlInputsAlwaysOk) {
GenNodeMap map;
NodeDef node1 = MakeNodeConst("node1");
map["node1"] = std::make_unique<GenNode>(&node1);
node1.add_input("^node1");
auto gn1 = map["node1"].get();
ASSERT_THAT(gn1->ParseInputs(&map), Eq(absl::OkStatus()));
// clang-format off
EXPECT_THAT(DumpLinkMap(gn1->links()), ElementsAre(
"iC: node1[oC]",
"oC: node1[iC]"
));
// clang-format on
}
TEST(GenNodeTest, ParseNodeInvalidInput) {
GenNodeMap map;
NodeDef node1 = MakeNodeAddN("node1", "node2", "node3");
map["node1"] = std::make_unique<GenNode>(&node1);
node1.add_input("node1");
auto gn1 = map["node1"].get();
ASSERT_THAT(
gn1->ParseInputs(&map),
Eq(absl::Status(
absl::StatusCode::kInvalidArgument,
"Node 'node1' input 0 refers to a non-existing node 'node2'.")));
}
TEST(GenNodeTest, BuildGraphInMap) {
GraphDef graph;
// A topology with a loop.
(*graph.add_node()) = MakeNodeConst("node1");
(*graph.add_node()) = MakeNodeSub("node2", "node3:1", "node3:0");
(*graph.add_node()) =
MakeNodeBroadcastGradientArgs("node3", "node1", "node2");
GenNodeMap map;
ASSERT_THAT(GenNode::BuildGraphInMap(graph, &map), Eq(absl::OkStatus()));
ASSERT_THAT(map.find("node1"), Ne(map.end()));
ASSERT_THAT(map.find("node2"), Ne(map.end()));
ASSERT_THAT(map.find("node3"), Ne(map.end()));
EXPECT_THAT(map["node1"]->name(), Eq("node1"));
EXPECT_THAT(map["node2"]->name(), Eq("node2"));
EXPECT_THAT(map["node3"]->name(), Eq("node3"));
// clang-format off
EXPECT_THAT(DumpLinkMap(map["node1"]->links()), ElementsAre(
"o0: node3[i0]"
));
EXPECT_THAT(DumpLinkMap(map["node2"]->links()), ElementsAre(
"i0: node3[o1]",
"i1: node3[o0]",
"o0: node3[i1]"
));
EXPECT_THAT(DumpLinkMap(map["node3"]->links()), ElementsAre(
"i0: node1[o0]",
"i1: node2[o0]",
"o0: node2[i1]",
"o1: node2[i0]"
));
// clang-format on
}
TEST(GenNodeTest, BuildGraphInMapDuplicateNode) {
GraphDef graph;
(*graph.add_node()) = MakeNodeConst("node1");
(*graph.add_node()) = MakeNodeConst("node1");
GenNodeMap map;
ASSERT_THAT(GenNode::BuildGraphInMap(graph, &map),
Eq(absl::Status(absl::StatusCode::kInvalidArgument,
"Duplicate node name 'node1'.")));
}
TEST(GenNodeTest, BuildGraphInMapParseError) {
GraphDef graph;
// A topology with a loop.
(*graph.add_node()) = MakeNodeConst("node1");
(*graph.add_node()) = MakeNodeSub("node2", "node3:1", "node3:0");
GenNodeMap map;
ASSERT_THAT(
GenNode::BuildGraphInMap(graph, &map),
Eq(absl::Status(
absl::StatusCode::kInvalidArgument,
"Node 'node2' input 0 refers to a non-existing node 'node3'.")));
}
} // end namespace
} // end namespace test
} // end namespace graph_analyzer
} // end namespace grappler
} // end namespace tensorflow
@@ -0,0 +1,353 @@
/* 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/grappler/graph_analyzer/graph_analyzer.h"
#include <cstddef>
#include <deque>
#include <iostream>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "absl/status/status.h"
#include "absl/strings/str_format.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/grappler/graph_analyzer/gen_node.h"
#include "tensorflow/core/grappler/graph_analyzer/sig_node.h"
namespace tensorflow {
namespace grappler {
namespace graph_analyzer {
GraphAnalyzer::GraphAnalyzer(const GraphDef& graph, int subgraph_size)
: graph_(graph), subgraph_size_(subgraph_size) {}
GraphAnalyzer::~GraphAnalyzer() {}
absl::Status GraphAnalyzer::Run() {
// The signature computation code would detect this too, but better
// to report it up front than spend time computing all the graphs first.
if (subgraph_size_ > Signature::kMaxGraphSize) {
return absl::Status(
absl::StatusCode::kInvalidArgument,
absl::StrFormat("Subgraphs of %d nodes are not supported, "
"the maximal supported node count is %d.",
subgraph_size_, Signature::kMaxGraphSize));
}
absl::Status st = BuildMap();
if (!st.ok()) {
return st;
}
FindSubgraphs();
DropInvalidSubgraphs();
st = CollateResult();
if (!st.ok()) {
return st;
}
return absl::OkStatus();
}
absl::Status GraphAnalyzer::BuildMap() {
nodes_.clear();
return GenNode::BuildGraphInMap(graph_, &nodes_);
}
void GraphAnalyzer::FindSubgraphs() {
result_.clear();
if (subgraph_size_ < 1) {
return;
}
partial_.clear();
todo_.clear(); // Just in case.
// Start with all subgraphs of size 1.
const Subgraph::Identity empty_parent;
for (const auto& node : nodes_) {
if (subgraph_size_ == 1) {
result_.ExtendParent(empty_parent, node.second.get());
} else {
// At this point ExtendParent() is guaranteed to not return nullptr.
todo_.push_back(partial_.ExtendParent(empty_parent, node.second.get()));
}
}
// Then extend the subgraphs until no more extensions are possible.
while (!todo_.empty()) {
ExtendSubgraph(todo_.front());
todo_.pop_front();
}
partial_.clear();
}
void GraphAnalyzer::ExtendSubgraph(Subgraph* parent) {
const int next_parent_id = parent->id().size() + 1;
bool will_complete = (next_parent_id == subgraph_size_);
SubgraphPtrSet& sg_set = will_complete ? result_ : partial_;
const GenNode* last_all_or_none_node = nullptr;
for (SubgraphIterator sit(parent); !sit.AtEnd(); sit.Next()) {
const GenNode* node = sit.GetNode();
GenNode::Port port = sit.GetPort();
const GenNode::LinkTarget& neighbor = sit.GetNeighbor();
if (node->AllInputsOrNone() && port.IsInbound() && !port.IsControl()) {
if (node != last_all_or_none_node) {
ExtendSubgraphAllOrNone(parent, node);
last_all_or_none_node = node;
}
sit.SkipPort();
} else if (neighbor.node->AllInputsOrNone() && !port.IsInbound() &&
!port.IsControl()) {
if (parent->id().find(neighbor.node) == parent->id().end()) {
// Not added yet.
ExtendSubgraphAllOrNone(parent, neighbor.node);
}
} else if (node->IsMultiInput(port)) {
ExtendSubgraphPortAllOrNone(parent, node, port);
sit.SkipPort();
} else if (neighbor.node->IsMultiInput(neighbor.port)) {
// Would need to add all inputs of the neighbor node at this port at
// once.
if (parent->id().find(neighbor.node) != parent->id().end()) {
continue; // Already added.
}
ExtendSubgraphPortAllOrNone(parent, neighbor.node, neighbor.port);
} else {
Subgraph* sg = sg_set.ExtendParent(parent->id(), neighbor.node);
if (!will_complete && sg != nullptr) {
todo_.push_back(sg);
}
}
}
}
void GraphAnalyzer::ExtendSubgraphAllOrNone(Subgraph* parent,
const GenNode* node) {
Subgraph::Identity id = parent->id();
id.insert(node);
auto range_end = node->links().end();
for (auto nbit = node->links().begin(); nbit != range_end; ++nbit) {
auto port = nbit->first;
if (!port.IsInbound() || port.IsControl()) {
continue;
}
// Since there might be multiple links to the same nodes,
// have to add all links one-by-one to check whether the subgraph
// would grow too large. But if it does grow too large, there is no
// point in growing it more, can just skip over the rest of the links.
for (const auto& link : nbit->second) {
id.insert(link.node);
const int id_size = id.size();
if (id_size > subgraph_size_) {
return; // Too big.
}
}
}
AddExtendedSubgraph(parent, id);
}
void GraphAnalyzer::ExtendSubgraphPortAllOrNone(Subgraph* parent,
const GenNode* node,
GenNode::Port port) {
auto nbit = node->links().find(port);
if (nbit == node->links().end()) {
return; // Should never happen.
}
Subgraph::Identity id = parent->id();
id.insert(node);
// Since there might be multiple links to the same nodes,
// have to add all links one-by-one to check whether the subgraph
// would grow too large. But if it does grow too large, there is no
// point in growing it more, can just skip over the rest of the links.
for (const auto& link : nbit->second) {
id.insert(link.node);
const int id_size = id.size();
if (id_size > subgraph_size_) {
return; // Too big.
}
}
AddExtendedSubgraph(parent, id);
}
void GraphAnalyzer::AddExtendedSubgraph(Subgraph* parent,
const Subgraph::Identity& id) {
if (id.size() == parent->id().size()) {
return; // Nothing new was added.
}
auto sg = std::make_unique<Subgraph>(id);
SubgraphPtrSet& spec_sg_set =
(id.size() == subgraph_size_) ? result_ : partial_;
if (spec_sg_set.find(sg) != spec_sg_set.end()) {
// This subgraph was already found by extending from a different path.
return;
}
const int id_size = id.size();
if (id_size != subgraph_size_) {
todo_.push_back(sg.get());
}
spec_sg_set.insert(std::move(sg));
}
void GraphAnalyzer::DropInvalidSubgraphs() {
auto resit = result_.begin();
while (resit != result_.end()) {
if (HasInvalidMultiInputs(resit->get())) {
auto delit = resit;
++resit;
result_.erase(delit);
} else {
++resit;
}
}
}
bool GraphAnalyzer::HasInvalidMultiInputs(Subgraph* sg) {
// Do the all-or-none-input nodes.
for (auto const& node : sg->id()) {
if (!node->AllInputsOrNone()) {
continue;
}
bool anyIn = false;
bool anyOut = false;
auto range_end = node->links().end();
for (auto nbit = node->links().begin(); nbit != range_end; ++nbit) {
auto port = nbit->first;
if (!port.IsInbound() || port.IsControl()) {
continue;
}
// Since there might be multiple links to the same nodes,
// have to add all links one-by-one to check whether the subgraph
// would grow too large. But if it does grow too large, there is no
// point in growing it more, can just skip over the rest of the links.
for (const auto& link : nbit->second) {
if (sg->id().find(link.node) == sg->id().end()) {
anyOut = true;
} else {
anyIn = true;
}
}
}
if (anyIn && anyOut) {
return true;
}
}
// Do the multi-input ports.
for (SubgraphIterator sit(sg); !sit.AtEnd(); sit.Next()) {
if (sit.GetNode()->IsMultiInput(sit.GetPort())) {
bool anyIn = false;
bool anyOut = false;
do {
GenNode* peer = sit.GetNeighbor().node;
if (sg->id().find(peer) == sg->id().end()) {
anyOut = true;
} else {
anyIn = true;
}
} while (sit.NextIfSamePort());
if (anyIn && anyOut) {
return true;
}
}
}
return false;
}
absl::Status GraphAnalyzer::CollateResult() {
ordered_collation_.clear();
collation_map_.clear();
// Collate by the signatures of the graphs.
for (const auto& it : result_) {
auto sig = std::make_unique<Signature>();
it->ExtractForSignature(&sig->map);
absl::Status status = sig->Compute();
if (!status.ok()) {
return status;
}
auto& coll_entry = collation_map_[sig.get()];
if (coll_entry.sig == nullptr) {
coll_entry.sig = std::move(sig);
}
++coll_entry.count;
}
// Then order them by the count.
for (auto& entry : collation_map_) {
ordered_collation_.insert(&entry.second);
}
result_.clear(); // Not needed after collation.
return absl::OkStatus();
}
std::vector<std::string> GraphAnalyzer::DumpRawSubgraphs() {
std::vector<std::string> result;
for (const auto& it : result_) {
result.emplace_back(it->Dump());
}
return result;
}
std::vector<std::string> GraphAnalyzer::DumpSubgraphs() {
std::vector<std::string> result;
for (auto ptr : ordered_collation_) {
result.emplace_back(
absl::StrFormat("%d %s", ptr->count, ptr->sig->ToString()));
}
return result;
}
absl::Status GraphAnalyzer::OutputSubgraphs() {
size_t total = 0;
for (auto ptr : ordered_collation_) {
std::cout << ptr->count << ' ' << ptr->sig->ToString() << '\n';
total += ptr->count;
}
std::cout << "Total: " << total << '\n';
if (std::cout.fail()) {
return absl::Status(absl::StatusCode::kDataLoss,
"Failed to write to stdout");
} else {
return absl::OkStatus();
}
}
} // end namespace graph_analyzer
} // end namespace grappler
} // end namespace tensorflow
@@ -0,0 +1,155 @@
/* 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_GRAPPLER_GRAPH_ANALYZER_GRAPH_ANALYZER_H_
#define TENSORFLOW_CORE_GRAPPLER_GRAPH_ANALYZER_GRAPH_ANALYZER_H_
#include <deque>
#include <vector>
#include "absl/status/status.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/grappler/graph_analyzer/map_tools.h"
#include "tensorflow/core/grappler/graph_analyzer/sig_node.h"
#include "tensorflow/core/grappler/graph_analyzer/subgraph.h"
#include "tensorflow/core/lib/core/status.h"
namespace tensorflow {
namespace grappler {
namespace graph_analyzer {
namespace test {
class GraphAnalyzerTest;
} // end namespace test
// Finds all the subgraphs of a given size and groups them by equivalence.
class GraphAnalyzer {
public:
// Makes a copy of the graph.
GraphAnalyzer(const GraphDef& graph, int subgraph_size);
virtual ~GraphAnalyzer();
// Performs the analysis and collects the subgraphs.
absl::Status Run();
// Returns the subgraphs found in Run() printed to text.
std::vector<std::string> DumpSubgraphs();
// Prints the subgraphs found in Run() to stdout.
absl::Status OutputSubgraphs();
// TODO(babkin): add a way to extract the subgraphs as direct data
// structures and as protobufs, and to write protobufs to a RecordIO.
private:
GraphAnalyzer() = delete;
GraphAnalyzer(const GraphAnalyzer&) = delete;
void operator=(const GraphAnalyzer&) = delete;
friend class tensorflow::grappler::graph_analyzer::test::GraphAnalyzerTest;
// Builds the map of nodes from the original graph definition.
absl::Status BuildMap();
// Using nodes_, finds all the subgraphs of size subgraph_size_ and places
// them into result_.
void FindSubgraphs();
// Deletes from result_ the unacceptable subgraphs. Those include the
// subgraphs where not all the inputs at a multi-input port are included (this
// could happen if some of these inputs were reached and included through
// different paths).
void DropInvalidSubgraphs();
// Deletes from result_ duplicate entries of equivalent topology.
absl::Status CollateResult();
// Returns the raw subgraphs found in FindSubgraphs() printed to text.
std::vector<std::string> DumpRawSubgraphs();
// Finds and adds appropriately to either partial_ or result_ all the
// subgraphs that can be created by extending the parent subgraph by one node.
// Ignores the duplicates.
void ExtendSubgraph(Subgraph* parent);
// Extends the parent subgraph by adding another node (if it wasn't already
// added) and all its non-control inputs in the link map range at once.
// If the subgraph would grow over subgraph_size_, it gets ignored.
void ExtendSubgraphAllOrNone(Subgraph* parent, const GenNode* node);
// Same but adds one specific inbound port (even control) all-or-none.
void ExtendSubgraphPortAllOrNone(Subgraph* parent, const GenNode* node,
GenNode::Port port);
// The common final step called by ExtendSubgraph*AllOrNone() methods.
void AddExtendedSubgraph(Subgraph* parent, const Subgraph::Identity& id);
// Returns true if this subgraph has any multi-inputs that aren't all-in or
// all-out.
bool HasInvalidMultiInputs(Subgraph* sg);
// Graph to run the analysis on.
GraphDef graph_;
int subgraph_size_;
// The enriched graph of parsed nodes and connections.
GenNodeMap nodes_;
// The resulting set of subgraphs.
SubgraphPtrSet result_;
// The subgraphs of partial size, stored while finding the result.
SubgraphPtrSet partial_;
// The subgraphs of partial size (stored in partial_) that are still waiting
// to be extended.
//
// TODO(babkin): This is rather simple-minded, each subgraph is examined from
// scratch, which means that all its internal links get iterated too. But it's
// OK for the small subgraphs. This can be improved by keeping not just
// subgraphs but iterators on the list, each of them having the list not-yet
// examined nodes (and the link position of the next link to be examined for
// the first node). This would add extra constant overhead, so the break-even
// subgraph size is not clear yet.
std::deque<Subgraph*> todo_;
// The collation map by signature is designed to allow the removal of entries
// and moving of the signature references from the keys of this map to the
// outside world. Must be careful at inserting and removal: make sure that
// when a new entry is inserted, its signature reference gets populated with
// the same data as the key of the map, and that if a reference is moved out,
// the map entry gets removed before that reference gets destroyed.
struct CollationEntry {
std::shared_ptr<Signature> sig;
size_t count = 0;
};
using CollationMap =
std::unordered_map<Signature*, CollationEntry, HashAtPtr<Signature*>,
EqAtPtr<Signature*> >;
CollationMap collation_map_;
// The entries are owned by collation_map_, so must be removed from
// ordered_collation_ before removing them from collation_map_.
struct ReverseLessByCount {
bool operator()(CollationEntry* left, CollationEntry* right) const {
return left->count > right->count; // Reverse order.
}
};
using CollationOrderByCount =
std::multiset<CollationEntry*, ReverseLessByCount>;
CollationOrderByCount ordered_collation_;
};
} // end namespace graph_analyzer
} // end namespace grappler
} // end namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPPLER_GRAPH_ANALYZER_GRAPH_ANALYZER_H_
@@ -0,0 +1,576 @@
/* 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/grappler/graph_analyzer/graph_analyzer.h"
#include <deque>
#include <memory>
#include <string>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/status/status.h"
#include "tensorflow/core/grappler/graph_analyzer/test_tools.h"
namespace tensorflow {
namespace grappler {
namespace graph_analyzer {
namespace test {
using ::testing::ElementsAre;
using ::testing::Eq;
using ::testing::Ne;
using ::testing::SizeIs;
using ::testing::UnorderedElementsAre;
class GraphAnalyzerTest : public ::testing::Test, protected TestGraphs {
protected:
absl::Status BuildMap() { return gran_->BuildMap(); }
void FindSubgraphs() { gran_->FindSubgraphs(); }
void DropInvalidSubgraphs() { gran_->DropInvalidSubgraphs(); }
absl::Status CollateResult() { return gran_->CollateResult(); }
void ExtendSubgraph(Subgraph* parent) { gran_->ExtendSubgraph(parent); }
void ExtendSubgraphPortAllOrNone(Subgraph* parent, GenNode* node,
GenNode::Port port) {
gran_->ExtendSubgraphPortAllOrNone(parent, node, port);
}
void ExtendSubgraphAllOrNone(Subgraph* parent, GenNode* node) {
gran_->ExtendSubgraphAllOrNone(parent, node);
}
std::vector<std::string> DumpRawSubgraphs() {
return gran_->DumpRawSubgraphs();
}
std::vector<std::string> DumpPartials() {
std::vector<std::string> result;
for (const auto& it : gran_->partial_) {
result.emplace_back(it->Dump());
}
return result;
}
const GenNodeMap& GetNodes() { return gran_->nodes_; }
GenNode* GetNode(const std::string& name) {
return gran_->nodes_.at(name).get();
}
SubgraphPtrSet& GetResult() { return gran_->result_; }
SubgraphPtrSet& GetPartial() { return gran_->partial_; }
std::deque<Subgraph*>& GetTodo() { return gran_->todo_; }
// Gets initialized by a particular test from a suitable GraphDef.
std::unique_ptr<GraphAnalyzer> gran_;
};
TEST_F(GraphAnalyzerTest, BuildMap) {
gran_ = std::make_unique<GraphAnalyzer>(graph_3n_self_control_, 1);
absl::Status st = BuildMap();
EXPECT_THAT(st, Eq(absl::OkStatus()));
auto& map = GetNodes();
EXPECT_THAT(map.find("node1"), Ne(map.end()));
EXPECT_THAT(map.find("node2"), Ne(map.end()));
EXPECT_THAT(map.find("node3"), Ne(map.end()));
}
TEST_F(GraphAnalyzerTest, BuildMapError) {
// A duplicate node.
(*graph_3n_self_control_.add_node()) = MakeNodeConst("node1");
gran_ = std::make_unique<GraphAnalyzer>(graph_3n_self_control_, 1);
absl::Status st = BuildMap();
ASSERT_THAT(st, Eq(absl::Status(absl::StatusCode::kInvalidArgument,
"Duplicate node name 'node1'.")));
}
TEST_F(GraphAnalyzerTest, FindSubgraphs0) {
gran_ = std::make_unique<GraphAnalyzer>(graph_3n_self_control_, 0);
absl::Status st = BuildMap();
ASSERT_THAT(st, Eq(absl::OkStatus()));
FindSubgraphs();
auto& subgraphs = GetResult();
EXPECT_THAT(subgraphs, SizeIs(0));
EXPECT_THAT(DumpRawSubgraphs(), ElementsAre());
EXPECT_THAT(DumpPartials(), UnorderedElementsAre());
EXPECT_THAT(GetTodo(), SizeIs(0));
}
TEST_F(GraphAnalyzerTest, FindSubgraphs1) {
gran_ = std::make_unique<GraphAnalyzer>(graph_3n_self_control_, 1);
absl::Status st = BuildMap();
ASSERT_THAT(st, Eq(absl::OkStatus()));
FindSubgraphs();
auto& subgraphs = GetResult();
EXPECT_THAT(subgraphs, SizeIs(3));
// clang-format off
EXPECT_THAT(DumpRawSubgraphs(), UnorderedElementsAre(
"1: BroadcastGradientArgs(node3)",
"1: Const(node1)",
"1: Sub(node2)"
));
// clang-format on
EXPECT_THAT(DumpPartials(), UnorderedElementsAre());
EXPECT_THAT(GetTodo(), SizeIs(0));
}
// The required subgraphs are larger than the graph.
TEST_F(GraphAnalyzerTest, FindSubgraphsTooLarge) {
gran_ = std::make_unique<GraphAnalyzer>(graph_3n_self_control_, 4);
absl::Status st = BuildMap();
ASSERT_THAT(st, Eq(absl::OkStatus()));
FindSubgraphs();
EXPECT_THAT(DumpRawSubgraphs(), ElementsAre());
EXPECT_THAT(DumpPartials(), UnorderedElementsAre());
EXPECT_THAT(GetTodo(), SizeIs(0));
}
//===
// Successfully propagate backwards through a multi-input link,
// with the base (currently-extending) node already in the graph.
TEST_F(GraphAnalyzerTest, MultiInputSuccessBackwardsBaseIn) {
gran_ = std::make_unique<GraphAnalyzer>(graph_multi_input_, 4);
absl::Status st = BuildMap();
ASSERT_THAT(st, Eq(absl::OkStatus()));
auto root =
std::make_unique<Subgraph>(Subgraph::Identity({GetNode("add2")}));
ExtendSubgraphPortAllOrNone(root.get(), GetNode("add2"),
GenNode::Port(true, 0));
// clang-format off
EXPECT_THAT(DumpRawSubgraphs(), UnorderedElementsAre(
"1: AddN(add2), Const(const2_1), Const(const2_2), Const(const2_3)"
));
// clang-format on
EXPECT_THAT(DumpPartials(), UnorderedElementsAre());
EXPECT_THAT(GetTodo(), SizeIs(0));
}
// Successfully propagate backwards through a multi-input link,
// with the base (currently-extending) node not in the graph yet.
TEST_F(GraphAnalyzerTest, MultiInputSuccessBackwardsBaseOut) {
gran_ = std::make_unique<GraphAnalyzer>(graph_multi_input_, 4);
absl::Status st = BuildMap();
ASSERT_THAT(st, Eq(absl::OkStatus()));
auto parent = std::make_unique<Subgraph>(Subgraph::Identity());
auto root =
std::make_unique<Subgraph>(Subgraph::Identity({GetNode("add2")}));
ExtendSubgraphPortAllOrNone(parent.get(), GetNode("add2"),
GenNode::Port(true, 0));
// clang-format off
EXPECT_THAT(DumpRawSubgraphs(), UnorderedElementsAre(
"1: AddN(add2), Const(const2_1), Const(const2_2), Const(const2_3)"
));
// clang-format on
EXPECT_THAT(DumpPartials(), UnorderedElementsAre());
EXPECT_THAT(GetTodo(), SizeIs(0));
}
// Successfully propagate backwards through a multi-input link,
// where the target subgraph size is larger.
TEST_F(GraphAnalyzerTest, MultiInputSuccessBackwardsIncomplete) {
gran_ = std::make_unique<GraphAnalyzer>(graph_multi_input_, 5);
absl::Status st = BuildMap();
ASSERT_THAT(st, Eq(absl::OkStatus()));
auto root =
std::make_unique<Subgraph>(Subgraph::Identity({GetNode("add2")}));
ExtendSubgraphPortAllOrNone(root.get(), GetNode("add2"),
GenNode::Port(true, 0));
EXPECT_THAT(DumpRawSubgraphs(), UnorderedElementsAre());
// clang-format off
EXPECT_THAT(DumpPartials(), UnorderedElementsAre(
"1: AddN(add2), Const(const2_1), Const(const2_2), Const(const2_3)"
));
// clang-format on
EXPECT_THAT(GetTodo(), SizeIs(1));
}
// Propagate backwards through a multi-input link, finding that the
// resulting subgraph would be too large.
TEST_F(GraphAnalyzerTest, MultiInputTooLargeBackwards) {
gran_ = std::make_unique<GraphAnalyzer>(graph_multi_input_, 3);
absl::Status st = BuildMap();
ASSERT_THAT(st, Eq(absl::OkStatus()));
auto root =
std::make_unique<Subgraph>(Subgraph::Identity({GetNode("add2")}));
ExtendSubgraphPortAllOrNone(root.get(), GetNode("add2"),
GenNode::Port(true, 0));
EXPECT_THAT(DumpRawSubgraphs(), UnorderedElementsAre());
EXPECT_THAT(DumpPartials(), UnorderedElementsAre());
EXPECT_THAT(GetTodo(), SizeIs(0));
}
// Propagate backwards through a multi-input link, finding that nothing
// would be added to the parent subgraph.
TEST_F(GraphAnalyzerTest, MultiInputNothingAddedBackwards) {
gran_ = std::make_unique<GraphAnalyzer>(graph_multi_input_, 4);
absl::Status st = BuildMap();
ASSERT_THAT(st, Eq(absl::OkStatus()));
auto root = std::make_unique<Subgraph>(
Subgraph::Identity({GetNode("add2"), GetNode("const2_1"),
GetNode("const2_2"), GetNode("const2_3")}));
ExtendSubgraphPortAllOrNone(root.get(), GetNode("add2"),
GenNode::Port(true, 0));
EXPECT_THAT(DumpRawSubgraphs(), UnorderedElementsAre());
EXPECT_THAT(DumpPartials(), UnorderedElementsAre());
EXPECT_THAT(GetTodo(), SizeIs(0));
}
// Successfully propagate forwards through a multi-input link,
// with the base (currently-extending) node not in the subgraph yet.
TEST_F(GraphAnalyzerTest, MultiInputSuccessForwardsBaseOut) {
gran_ = std::make_unique<GraphAnalyzer>(graph_multi_input_, 4);
absl::Status st = BuildMap();
ASSERT_THAT(st, Eq(absl::OkStatus()));
auto root =
std::make_unique<Subgraph>(Subgraph::Identity({GetNode("const2_1")}));
ExtendSubgraphPortAllOrNone(root.get(), GetNode("add2"),
GenNode::Port(true, 0));
// clang-format off
EXPECT_THAT(DumpRawSubgraphs(), UnorderedElementsAre(
"1: AddN(add2), Const(const2_1), Const(const2_2), Const(const2_3)"
));
// clang-format on
EXPECT_THAT(DumpPartials(), UnorderedElementsAre());
EXPECT_THAT(GetTodo(), SizeIs(0));
}
// Successfully propagate backwards through a multi-input link.
TEST_F(GraphAnalyzerTest, MultiInputSuccessBackwardsFull) {
gran_ = std::make_unique<GraphAnalyzer>(graph_multi_input_, 4);
absl::Status st = BuildMap();
ASSERT_THAT(st, Eq(absl::OkStatus()));
auto root =
std::make_unique<Subgraph>(Subgraph::Identity({GetNode("add2")}));
ExtendSubgraph(root.get());
// clang-format off
EXPECT_THAT(DumpRawSubgraphs(), UnorderedElementsAre(
"1: AddN(add2), Const(const2_1), Const(const2_2), Const(const2_3)"
));
EXPECT_THAT(DumpPartials(), UnorderedElementsAre(
"1: AddN(add2), Sub(sub)"
));
// clang-format on
EXPECT_THAT(GetTodo(), SizeIs(1));
}
// Successfully propagate forwards through a multi-input link.
TEST_F(GraphAnalyzerTest, MultiInputSuccessForwardsFull) {
gran_ = std::make_unique<GraphAnalyzer>(graph_multi_input_, 4);
absl::Status st = BuildMap();
ASSERT_THAT(st, Eq(absl::OkStatus()));
auto root =
std::make_unique<Subgraph>(Subgraph::Identity({GetNode("const2_1")}));
ExtendSubgraph(root.get());
// clang-format off
EXPECT_THAT(DumpRawSubgraphs(), UnorderedElementsAre(
"1: AddN(add2), Const(const2_1), Const(const2_2), Const(const2_3)"
));
// clang-format on
EXPECT_THAT(DumpPartials(), UnorderedElementsAre());
EXPECT_THAT(GetTodo(), SizeIs(0));
}
TEST_F(GraphAnalyzerTest, DropInvalidSubgraphsMulti) {
gran_ = std::make_unique<GraphAnalyzer>(graph_multi_input_, 3);
absl::Status st = BuildMap();
ASSERT_THAT(st, Eq(absl::OkStatus()));
// A good one, multi-input is all-in.
GetResult().insert(std::make_unique<Subgraph>(Subgraph::Identity({
GetNode("const1_1"),
GetNode("const1_2"),
GetNode("add1"),
})));
// A good one, multi-input is all-out
GetResult().insert(std::make_unique<Subgraph>(Subgraph::Identity({
GetNode("add1"),
GetNode("add2"),
GetNode("sub"),
})));
// A bad one, multi-input is partially in.
GetResult().insert(std::make_unique<Subgraph>(Subgraph::Identity({
GetNode("const1_1"),
GetNode("add1"),
GetNode("sub"),
})));
// A bad one, multi-input is partially in.
GetResult().insert(std::make_unique<Subgraph>(Subgraph::Identity({
GetNode("add2"),
GetNode("const2_1"),
GetNode("const2_2"),
})));
DropInvalidSubgraphs();
// clang-format off
EXPECT_THAT(DumpRawSubgraphs(), UnorderedElementsAre(
"1: AddN(add1), AddN(add2), Sub(sub)",
"1: AddN(add1), Const(const1_1), Const(const1_2)"
));
// clang-format on
EXPECT_THAT(DumpPartials(), UnorderedElementsAre());
EXPECT_THAT(GetTodo(), SizeIs(0));
}
//===
// Successfully propagate backwards through a multi-input link,
// with the base (currently-extending) node already in the graph.
TEST_F(GraphAnalyzerTest, AllOrNoneInputSuccessBackwards) {
gran_ = std::make_unique<GraphAnalyzer>(graph_all_or_none_, 4);
absl::Status st = BuildMap();
ASSERT_THAT(st, Eq(absl::OkStatus()));
auto root =
std::make_unique<Subgraph>(Subgraph::Identity({GetNode("pass2")}));
ExtendSubgraphAllOrNone(root.get(), GetNode("pass2"));
// clang-format off
EXPECT_THAT(DumpRawSubgraphs(), UnorderedElementsAre(
"1: Const(const2_1), Const(const2_2), Const(const2_3), IdentityN(pass2)"
));
EXPECT_THAT(DumpPartials(), UnorderedElementsAre());
// clang-format on
EXPECT_THAT(GetTodo(), SizeIs(0));
}
// Successfully propagate backwards through a multi-input link,
// but no control links propagate. It also tests the situation
// where the target subgraph size is larger.
TEST_F(GraphAnalyzerTest, AllOrNoneInputSuccessBackwardsNoControl) {
gran_ = std::make_unique<GraphAnalyzer>(graph_all_or_none_, 5);
absl::Status st = BuildMap();
ASSERT_THAT(st, Eq(absl::OkStatus()));
auto root =
std::make_unique<Subgraph>(Subgraph::Identity({GetNode("pass1")}));
ExtendSubgraphAllOrNone(root.get(), GetNode("pass1"));
// clang-format off
EXPECT_THAT(DumpRawSubgraphs(), UnorderedElementsAre());
EXPECT_THAT(DumpPartials(), UnorderedElementsAre(
"1: Const(const1_1), Const(const1_2), IdentityN(pass1)"
));
// clang-format on
EXPECT_THAT(GetTodo(), SizeIs(1));
}
// The control links propagate separately as all-or-none, even on the nodes
// that are all-or-none for the normal inputs.
TEST_F(GraphAnalyzerTest, AllOrNoneInputSeparateControl) {
gran_ = std::make_unique<GraphAnalyzer>(graph_all_or_none_, 5);
absl::Status st = BuildMap();
ASSERT_THAT(st, Eq(absl::OkStatus()));
auto root =
std::make_unique<Subgraph>(Subgraph::Identity({GetNode("pass1")}));
ExtendSubgraphPortAllOrNone(root.get(), GetNode("pass1"),
GenNode::Port(true, -1));
// clang-format off
EXPECT_THAT(DumpRawSubgraphs(), UnorderedElementsAre());
EXPECT_THAT(DumpPartials(), UnorderedElementsAre(
"1: Const(const2_1), Const(const2_2), Const(const2_3), IdentityN(pass1)"
));
// clang-format on
EXPECT_THAT(GetTodo(), SizeIs(1));
}
// Propagate backwards from all-or-none-input node, finding that the
// resulting subgraph would be too large.
TEST_F(GraphAnalyzerTest, AllOrNoneInputTooLargeBackwards) {
gran_ = std::make_unique<GraphAnalyzer>(graph_all_or_none_, 3);
absl::Status st = BuildMap();
ASSERT_THAT(st, Eq(absl::OkStatus()));
auto root =
std::make_unique<Subgraph>(Subgraph::Identity({GetNode("pass2")}));
ExtendSubgraphAllOrNone(root.get(), GetNode("pass2"));
EXPECT_THAT(DumpRawSubgraphs(), UnorderedElementsAre());
EXPECT_THAT(DumpPartials(), UnorderedElementsAre());
EXPECT_THAT(GetTodo(), SizeIs(0));
}
// Propagate backwards from all-or-none-input node, finding that nothing
// would be added to the parent subgraph.
TEST_F(GraphAnalyzerTest, AllOrNoneInputNothingAddedBackwards) {
gran_ = std::make_unique<GraphAnalyzer>(graph_all_or_none_, 4);
absl::Status st = BuildMap();
ASSERT_THAT(st, Eq(absl::OkStatus()));
auto root = std::make_unique<Subgraph>(
Subgraph::Identity({GetNode("pass2"), GetNode("const2_1"),
GetNode("const2_2"), GetNode("const2_3")}));
ExtendSubgraphAllOrNone(root.get(), GetNode("pass2"));
EXPECT_THAT(DumpRawSubgraphs(), UnorderedElementsAre());
EXPECT_THAT(DumpPartials(), UnorderedElementsAre());
EXPECT_THAT(GetTodo(), SizeIs(0));
}
// Successfully propagate forwards to all-or-none-input node,
// with the base (currently-extending) node not in the subgraph yet.
TEST_F(GraphAnalyzerTest, AllOrNoneInputSuccessForwardsBaseOut) {
gran_ = std::make_unique<GraphAnalyzer>(graph_all_or_none_, 4);
absl::Status st = BuildMap();
ASSERT_THAT(st, Eq(absl::OkStatus()));
auto root =
std::make_unique<Subgraph>(Subgraph::Identity({GetNode("const2_1")}));
ExtendSubgraphAllOrNone(root.get(), GetNode("pass2"));
// clang-format off
EXPECT_THAT(DumpRawSubgraphs(), UnorderedElementsAre(
"1: Const(const2_1), Const(const2_2), Const(const2_3), IdentityN(pass2)"
));
EXPECT_THAT(DumpPartials(), UnorderedElementsAre());
// clang-format on
EXPECT_THAT(GetTodo(), SizeIs(0));
}
// Successfully propagate backwards from all-or-none-input node.
TEST_F(GraphAnalyzerTest, AllOrNoneInputSuccessBackwardsFull) {
gran_ = std::make_unique<GraphAnalyzer>(graph_all_or_none_, 4);
absl::Status st = BuildMap();
ASSERT_THAT(st, Eq(absl::OkStatus()));
auto root =
std::make_unique<Subgraph>(Subgraph::Identity({GetNode("pass2")}));
ExtendSubgraph(root.get());
// clang-format off
EXPECT_THAT(DumpRawSubgraphs(), UnorderedElementsAre(
"1: Const(const2_1), Const(const2_2), Const(const2_3), IdentityN(pass2)"
));
EXPECT_THAT(DumpPartials(), UnorderedElementsAre(
"1: IdentityN(pass2), Sub(sub)"
));
// clang-format on
EXPECT_THAT(GetTodo(), SizeIs(1));
}
// Successfully propagate forwards to all-or-none-input node. This includes
// both all-or-none-input for the normal inputs, and multi-input by the
// control path.
TEST_F(GraphAnalyzerTest, AllOrNoneInputSuccessForwardsFull) {
gran_ = std::make_unique<GraphAnalyzer>(graph_all_or_none_, 4);
absl::Status st = BuildMap();
ASSERT_THAT(st, Eq(absl::OkStatus()));
auto root =
std::make_unique<Subgraph>(Subgraph::Identity({GetNode("const2_1")}));
ExtendSubgraph(root.get());
// clang-format off
EXPECT_THAT(DumpRawSubgraphs(), UnorderedElementsAre(
"1: Const(const2_1), Const(const2_2), Const(const2_3), IdentityN(pass2)",
"1: Const(const2_1), Const(const2_2), Const(const2_3), IdentityN(pass1)"
));
EXPECT_THAT(DumpPartials(), UnorderedElementsAre());
// clang-format on
EXPECT_THAT(GetTodo(), SizeIs(0));
}
TEST_F(GraphAnalyzerTest, DropInvalidSubgraphsAllOrNone) {
gran_ = std::make_unique<GraphAnalyzer>(graph_all_or_none_, 3);
absl::Status st = BuildMap();
ASSERT_THAT(st, Eq(absl::OkStatus()));
// A good one, all-or-none is all-in.
GetResult().insert(std::make_unique<Subgraph>(Subgraph::Identity({
GetNode("const1_1"),
GetNode("const1_2"),
GetNode("pass1"),
})));
// A good one, all-or-none is all-out
GetResult().insert(std::make_unique<Subgraph>(Subgraph::Identity({
GetNode("pass1"),
GetNode("pass2"),
GetNode("sub"),
})));
// A bad one, all-or-none is partially in.
GetResult().insert(std::make_unique<Subgraph>(Subgraph::Identity({
GetNode("const1_1"),
GetNode("pass1"),
GetNode("sub"),
})));
// A bad one, all-or-none is partially in.
GetResult().insert(std::make_unique<Subgraph>(Subgraph::Identity({
GetNode("pass2"),
GetNode("const2_1"),
GetNode("const2_2"),
})));
DropInvalidSubgraphs();
// clang-format off
EXPECT_THAT(DumpRawSubgraphs(), UnorderedElementsAre(
"1: IdentityN(pass1), IdentityN(pass2), Sub(sub)",
"1: Const(const1_1), Const(const1_2), IdentityN(pass1)"
));
// clang-format on
EXPECT_THAT(DumpPartials(), UnorderedElementsAre());
EXPECT_THAT(GetTodo(), SizeIs(0));
}
} // end namespace test
} // end namespace graph_analyzer
} // end namespace grappler
} // end namespace tensorflow
@@ -0,0 +1,103 @@
/* 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 <string>
#include <vector>
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/grappler/graph_analyzer/graph_analyzer.h"
#include "tensorflow/core/grappler/utils/transitive_fanin.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
namespace tensorflow {
namespace grappler {
namespace graph_analyzer {
// Dies on failure.
static void LoadModel(const std::string& filename,
tensorflow::MetaGraphDef* metagraph) {
LOG(INFO) << "Loading model from " << filename;
absl::Status st;
st = ReadBinaryProto(Env::Default(), filename, metagraph);
if (!st.ok()) {
LOG(WARNING) << "Failed to read a binary metagraph: " << st;
st = ReadTextProto(Env::Default(), filename, metagraph);
if (!st.ok()) {
LOG(FATAL) << "Failed to read a text metagraph: " << st;
}
}
}
// Prune the graph to only keep the transitive fanin part with respect to a set
// of train ops (if provided).
void MaybePruneGraph(const tensorflow::MetaGraphDef& metagraph,
tensorflow::GraphDef* graph) {
std::vector<std::string> fetch_nodes;
for (const auto& fetch :
metagraph.collection_def().at("train_op").node_list().value()) {
LOG(INFO) << "Fetch node: " << fetch;
fetch_nodes.push_back(fetch);
}
if (fetch_nodes.empty()) {
*graph = metagraph.graph_def();
} else {
std::vector<const tensorflow::NodeDef*> fanin_nodes;
TF_CHECK_OK(tensorflow::grappler::ComputeTransitiveFanin(
metagraph.graph_def(), fetch_nodes, &fanin_nodes));
for (const tensorflow::NodeDef* node : fanin_nodes) {
*(graph->add_node()) = *node;
}
LOG(INFO) << "Pruned "
<< metagraph.graph_def().node_size() - graph->node_size()
<< " nodes. Original graph size: "
<< metagraph.graph_def().node_size()
<< ". New graph size: " << graph->node_size() << ".";
}
}
void GraphAnalyzerTool(const std::string& file_name, int n) {
if (n < 1) {
LOG(FATAL) << "Invalid subgraph size " << n << ", must be at least 1";
}
tensorflow::MetaGraphDef metagraph;
LoadModel(file_name, &metagraph);
tensorflow::GraphDef graph;
MaybePruneGraph(metagraph, &graph);
tensorflow::grappler::graph_analyzer::GraphAnalyzer analyzer(graph, n);
LOG(INFO) << "Running the analysis";
absl::Status st = analyzer.Run();
if (!st.ok()) {
LOG(FATAL) << "Analysis failed: " << st;
}
LOG(INFO) << "Printing the result";
st = analyzer.OutputSubgraphs();
if (!st.ok()) {
LOG(FATAL) << "Failed to print the result: " << st;
}
LOG(INFO) << "Completed";
}
} // end namespace graph_analyzer
} // end namespace grappler
} // end namespace tensorflow
@@ -0,0 +1,31 @@
/* 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_GRAPPLER_GRAPH_ANALYZER_GRAPH_ANALYZER_TOOL_H_
#define TENSORFLOW_CORE_GRAPPLER_GRAPH_ANALYZER_GRAPH_ANALYZER_TOOL_H_
#include "tensorflow/core/lib/strings/str_util.h"
namespace tensorflow {
namespace grappler {
namespace graph_analyzer {
void GraphAnalyzerTool(const std::string& file_name, int n);
} // end namespace graph_analyzer
} // end namespace grappler
} // end namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPPLER_GRAPH_ANALYZER_GRAPH_ANALYZER_TOOL_H_
@@ -0,0 +1,47 @@
/* 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_GRAPPLER_GRAPH_ANALYZER_HASH_TOOLS_H_
#define TENSORFLOW_CORE_GRAPPLER_GRAPH_ANALYZER_HASH_TOOLS_H_
#include <cstddef>
namespace tensorflow {
namespace grappler {
namespace graph_analyzer {
// Unfortunately, std::hash provides no way to combine hashes, so everyone
// is copying boost::hash_combine. This is a version that follows Google's
// guidelines on the arguments, and contains only the combination, without
// hashing.
inline void CombineHash(size_t from, size_t* to) {
*to ^= from + 0x9e3779b9 + (*to << 6) + (*to >> 2);
}
// Combine two hashes in such a way that the order of combination doesn't matter
// (so it's really both commutative and associative). The result is not a very
// high-quality hash but can be used in case if the order of sub-elements must
// not matter in the following comparison. An alternative would be to sort the
// hashes of the sub-elements and then combine them normally in the sorted
// order.
inline void CombineHashCommutative(size_t from, size_t* to) {
*to = *to + from + 0x9e3779b9;
}
} // end namespace graph_analyzer
} // end namespace grappler
} // end namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPPLER_GRAPH_ANALYZER_HASH_TOOLS_H_
@@ -0,0 +1,48 @@
/* 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/grappler/graph_analyzer/hash_tools.h"
#include <cstddef>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
namespace tensorflow {
namespace grappler {
namespace graph_analyzer {
namespace test {
namespace {
using ::testing::Eq;
TEST(HashToolsTest, CombineHashCommutative) {
size_t a = 0;
size_t b = 999;
size_t c = a;
CombineHashCommutative(b, &c);
size_t d = b;
CombineHashCommutative(a, &d);
EXPECT_THAT(c, Eq(d));
}
} // namespace
} // end namespace test
} // end namespace graph_analyzer
} // end namespace grappler
} // end namespace tensorflow
@@ -0,0 +1,46 @@
/* 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_GRAPPLER_GRAPH_ANALYZER_MAP_TOOLS_H_
#define TENSORFLOW_CORE_GRAPPLER_GRAPH_ANALYZER_MAP_TOOLS_H_
#include <functional>
namespace tensorflow {
namespace grappler {
namespace graph_analyzer {
// Helpers for building maps of pointers.
template <typename Ptr>
struct LessAtPtr : std::function<bool(Ptr, Ptr)> {
bool operator()(const Ptr& x, const Ptr& y) const { return *x < *y; }
};
template <typename Ptr>
struct EqAtPtr : std::function<bool(Ptr, Ptr)> {
bool operator()(const Ptr& x, const Ptr& y) const { return *x == *y; }
};
template <typename Ptr>
struct HashAtPtr : std::function<size_t(Ptr)> {
size_t operator()(const Ptr& x) const { return x->Hash(); }
};
} // end namespace graph_analyzer
} // end namespace grappler
} // end namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPPLER_GRAPH_ANALYZER_MAP_TOOLS_H_
@@ -0,0 +1,466 @@
/* 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/grappler/graph_analyzer/sig_node.h"
#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <functional>
#include <ios>
#include <map>
#include <string>
#include "absl/log/check.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/strings/str_format.h"
#include "tensorflow/core/framework/node_def.pb.h"
namespace tensorflow {
namespace grappler {
namespace graph_analyzer {
static constexpr bool debug = false;
//=== SigNode
SigNode::SigNode(const NodeDef* node) : node_(node) {}
void SigNode::CopyLinks(const GenNode& from, const TranslationMap& map) {
hash_to_link_.clear();
hashed_peers_.clear();
std::map<LinkTag, Link> link_map;
CopyLinksPass1(from, map, &link_map);
CopyLinksPass2(&link_map);
}
void SigNode::CopyLinksPass1(const GenNode& from, const TranslationMap& map,
std::map<LinkTag, Link>* link_map) {
LinkTag::Hasher link_hasher;
for (const auto& entry : from.links()) {
for (const auto& target : entry.second) {
auto nodeit = map.find(target.node);
if (nodeit == map.end()) {
// Node is not in the subgraph, ignore.
continue;
}
LinkTag tag(entry.first, target.port);
size_t hval = link_hasher(tag);
// This instantiates the entry if it was not present.
Link& map_entry = (*link_map)[tag];
if (map_entry.peers.empty()) {
map_entry.tag = tag;
map_entry.unique_hash = hval;
}
map_entry.peers.push_back(nodeit->second);
}
}
}
void SigNode::CopyLinksPass2(std::map<LinkTag, Link>* link_map) {
for (auto& entry : *link_map) {
Link* hl_entry_ptr = &hash_to_link_[entry.second.unique_hash];
// In case of a conflict, rehash. This should almost never happen.
// Because the order of iteration is predictable, the rehashed values
// will also be predictable.
while (!hl_entry_ptr->peers.empty()) {
CombineHash(1, &entry.second.unique_hash);
hl_entry_ptr = &hash_to_link_[entry.second.unique_hash];
}
for (const auto& peer : entry.second.peers) {
hashed_peers_.emplace_back(HashedPeer(entry.second.unique_hash, peer));
}
hl_entry_ptr->tag = entry.second.tag;
hl_entry_ptr->unique_hash = entry.second.unique_hash;
hl_entry_ptr->peers.swap(entry.second.peers);
}
}
void SigNode::ComputeTopoHash0() {
topo_hash_.clear();
last_hashed_nodes_ = next_hashed_nodes_ = node_mask_;
// TODO(babkin): include the attributes too, as an option.
size_t hval = std::hash<std::string>()(opcode());
// Getting the topology of the links in to the hash early should get more
// conflicts resolved early.
for (const auto& entry : hashed_peers_) {
CombineHash(entry.link_hash, &hval);
}
topo_hash_.push_back(hval);
}
void SigNode::ComputeTopoHash(int distance) {
// The new starting point.
next_hashed_nodes_ = last_hashed_nodes_;
if (debug) {
LOG(INFO) << "DEBUG node " << name() << " mask=" << std::hex
<< next_hashed_nodes_;
}
if (hash_is_final_) {
return;
}
const int64_t topo_hash_size = topo_hash_.size();
CHECK(topo_hash_size == distance);
int prev = distance - 1;
// Start with own's local topology hash. This value is stable, so
// if the hashes of the surrounding nodes don't change on the following
// distances, the hash of this node won't change either.
size_t hval = topo_hash_[0];
if (!hashed_peers_.empty()) {
size_t last_link_hash = hashed_peers_[0].link_hash;
size_t comm_hash = 0;
for (const auto& entry : hashed_peers_) {
if (entry.link_hash != last_link_hash) {
CombineHash(last_link_hash, &hval);
CombineHash(comm_hash, &hval);
comm_hash = 0;
last_link_hash = entry.link_hash;
}
// The links in the same vector are commutative, so combine their
// hashes in a commutative way.
CombineHashCommutative(entry.peer->GetTopoHash(prev), &comm_hash);
next_hashed_nodes_ |= entry.peer->last_hashed_nodes_;
if (debug) {
LOG(INFO) << "DEBUG node " << name() << " += " << entry.peer->name()
<< " mask=" << std::hex << next_hashed_nodes_;
}
}
// The last commutative group.
CombineHash(last_link_hash, &hval);
CombineHash(comm_hash, &hval);
}
topo_hash_.push_back(hval);
}
size_t SigNode::GetTopoHash(int distance) const {
CHECK(!topo_hash_.empty());
const int64_t topo_hash_size = topo_hash_.size();
if (distance >= topo_hash_size) {
CHECK(hash_is_final_);
return topo_hash_.back();
} else {
return topo_hash_[distance];
}
}
bool SigNode::operator==(const SigNode& other) const {
// TODO(babkin): add attributes too.
if (opcode() != other.opcode()) {
return false;
}
// Normally the caller is expected to compare the nodes
// at the same rank in different graphs, but just in case...
if (unique_rank_ != other.unique_rank_) {
return false;
}
if (hashed_peers_.size() != other.hashed_peers_.size()) {
return false;
}
for (auto it1 = hashed_peers_.begin(), it2 = other.hashed_peers_.begin();
it1 != hashed_peers_.end(); ++it1, ++it2) {
// TODO(babkin): might compare the actual values too
// but the hash is probably just as good.
if (it1->link_hash != it2->link_hash) {
return false;
}
if (it1->peer->unique_rank_ != it2->peer->unique_rank_) {
return false;
}
}
return true;
}
//=== Signature
constexpr int Signature::kMaxGraphSize;
std::string Signature::ToString() const {
std::string result;
for (size_t n = 0; n < nodes.size(); ++n) {
// TODO(babkin): add attributes too.
result += absl::StrFormat("%d:%s", n, nodes[n]->opcode());
for (const auto& entry : nodes[n]->hashed_peers_) {
const auto& link = nodes[n]->hash_to_link_[entry.link_hash];
// The link entries are already sorted, by tags and then by the
// node ranks.
if (link.tag.local.IsInbound()) {
result += absl::StrFormat("[%s:%s:%d]", std::string(link.tag.local),
std::string(link.tag.remote),
entry.peer->unique_rank_);
}
}
result.push_back(',');
}
return result;
}
absl::Status Signature::Compute() {
if (map.size() > kMaxGraphSize) {
return absl::Status(
absl::StatusCode::kInvalidArgument,
absl::StrFormat(
"A graph of %d nodes is too big for signature computation, "
"the maximal supported node count is %d.",
map.size(), kMaxGraphSize));
}
// The value that will be assigned next as the unique node id.
// This also means that all the entries in nodes at indexes less than this
// have been finalized and don't need to be touched any more.
size_t next_node_id = 0;
sig_short = 0;
sig_full.resize(0); // Keep the storage.
// The main signature generation.
PrepareNodes();
FindUniqueHashes(&next_node_id);
while (next_node_id < map.size()) {
ComputeOneRound(next_node_id);
FindUniqueHashes(&next_node_id);
}
OrderLinks();
return absl::OkStatus();
}
void Signature::PrepareNodes() {
nodes.resize(0); // Keep the storage.
// Initialize the nodes.
int64_t mask = 1;
for (const auto& entry : map) {
SigNode* node = entry.second.get();
node->last_hashed_nodes_ = node->node_mask_ = mask;
mask <<= 1;
node->unique_rank_ = ~0;
node->hash_is_final_ = false;
node->ComputeTopoHash0();
if (node->GetHighTopoHash() <= map.size()) {
// Would conflict with one of the reserved values.
node->ReHighTopoHash();
}
// The initial order is random.
nodes.emplace_back(node);
}
}
void Signature::FindUniqueHashes(size_t* next_node_id_p) {
// Start by sorting by the hash value.
std::stable_sort(nodes.begin() + *next_node_id_p, nodes.end(),
SigNode::NodeOrderLess());
// At each call, if no nodes have unique hashes, one node that has a
// non-unique (shared) hash can be made unique by assigning a unique id.
// This node gets picked predictably by taking the last node.
// TODO(babkin): Technically, more than one node can be unshared,
// as long as their last_hashed_nodes_ overlap only by the nodes that
// already had the assigned ids before the current round. But it's not clear
// yet, how often would this beneficial, because it looks like for many
// subgraphs unsharing one node should be enough to untangle them. This
// would need more measurement before implementing.
bool found_unique = false;
for (size_t n = *next_node_id_p; n < nodes.size(); ++n) {
size_t cur_hash = nodes[n]->GetHighTopoHash();
if (n + 1 < nodes.size() && nodes[n + 1]->GetHighTopoHash() == cur_hash) {
// A sequence of nodes sharing the same hash. Skip over it.
// TODO(babkin): check here for the arbitrary hash conflicts and resolve
// them.
for (++n;
n + 1 < nodes.size() && nodes[n + 1]->GetHighTopoHash() == cur_hash;
++n) {
}
if (found_unique || n != nodes.size() - 1) {
// Either some unique nodes have already been found, or this is
// not the last chance, keep trying to find the unique nodes.
continue;
}
// Here we're at the last node and haven't found any unique ones.
// So fall through and make this last node unique.
}
found_unique = true;
size_t id = (*next_node_id_p)++;
nodes[n]->unique_rank_ = id;
size_t last_hash = nodes[n]->GetHighTopoHash();
CombineHash(last_hash, &sig_short);
sig_full.push_back(last_hash);
// Take the hash at 0 and mix the unique rank into it. After that it will
// stay fixed.
nodes[n]->topo_hash_.resize(1);
nodes[n]->topo_hash_[0] = id + 1; // Avoid the value of 0.
nodes[n]->hash_is_final_ = true;
nodes[n]->last_hashed_nodes_ = nodes[n]->node_mask_;
if (n != id) {
std::swap(nodes[id], nodes[n]);
}
}
}
void Signature::ComputeOneRound(size_t next_node_id) {
// Reset the state of the nodes.
int debug_i = 0;
for (auto it = nodes.begin() + next_node_id; it != nodes.end(); ++it) {
auto node = *it;
// The hash at distance 0 never changes, so preserve it.
node->topo_hash_.resize(1);
node->last_hashed_nodes_ = node->node_mask_;
node->hash_is_final_ = false;
if (debug) {
LOG(INFO) << "DEBUG distance=" << 0 << " node " << debug_i++ << " "
<< node->name() << " mask=" << std::hex
<< node->last_hashed_nodes_;
}
}
bool stop = false;
// The distance can reach up to nodes.size()+1, to include not only all the
// nodes but also all the redundant paths.
for (int distance = 1; !stop; ++distance) {
for (auto it = nodes.begin() + next_node_id; it != nodes.end(); ++it) {
auto node = *it;
if (node->hash_is_final_) {
continue;
}
node->ComputeTopoHash(distance);
if (node->GetHighTopoHash() <= nodes.size()) {
// Would conflict with one of the reserved values.
node->ReHighTopoHash();
}
}
// Will be looking for the indications to not stop.
stop = true;
debug_i = 0;
// The bitmasks get moved after all the hash computations are done.
for (auto it = nodes.begin() + next_node_id; it != nodes.end(); ++it) {
auto node = *it;
if (debug) {
LOG(INFO) << "DEBUG distance=" << distance << " node " << debug_i++
<< " " << node->name() << " oldmask=" << std::hex
<< node->last_hashed_nodes_ << " mask=" << std::hex
<< node->next_hashed_nodes_;
}
if (node->last_hashed_nodes_ == node->next_hashed_nodes_) {
// Stopped growing, this part of the graph must be fully
// surrounded by nodes that already have the unique ids.
node->hash_is_final_ = true;
} else {
node->last_hashed_nodes_ = node->next_hashed_nodes_;
stop = false;
}
}
}
}
void Signature::OrderLinks() {
for (const auto& node : nodes) {
if (node->hashed_peers_.empty()) {
continue;
}
size_t cur_link_hash = node->hashed_peers_[0].link_hash + 1;
int first_idx = -1;
int idx;
for (idx = 0; idx < static_cast<int64_t>(node->hashed_peers_.size());
++idx) {
auto& entry = node->hashed_peers_[idx];
if (entry.link_hash == cur_link_hash) {
continue;
}
if (idx - first_idx > 1) {
// Need to sort.
std::sort(node->hashed_peers_.begin() + first_idx,
node->hashed_peers_.begin() + idx,
SigNode::HashedPeer::LessByRank());
}
cur_link_hash = entry.link_hash;
first_idx = idx;
}
if (idx - first_idx > 1) {
// Sort the last bunch.
std::sort(node->hashed_peers_.begin() + first_idx,
node->hashed_peers_.begin() + idx,
SigNode::HashedPeer::LessByRank());
}
}
}
bool Signature::operator==(const Signature& other) const {
// Tries to find the differences as early as possible by
// comparing the hashes first.
if (sig_short != other.sig_short) {
return false;
}
if (sig_full.size() != other.sig_full.size()) {
return false;
}
for (auto it1 = sig_full.begin(), it2 = other.sig_full.begin();
it1 != sig_full.end(); ++it1, ++it2) {
if (*it1 != *it2) {
return false;
}
}
if (nodes.size() != other.nodes.size()) {
return false;
}
for (auto it1 = nodes.begin(), it2 = other.nodes.begin(); it1 != nodes.end();
++it1, ++it2) {
if (**it1 != **it2) {
return false;
}
}
return true;
}
} // end namespace graph_analyzer
} // end namespace grappler
} // end namespace tensorflow
@@ -0,0 +1,306 @@
/* 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_GRAPPLER_GRAPH_ANALYZER_SIG_NODE_H_
#define TENSORFLOW_CORE_GRAPPLER_GRAPH_ANALYZER_SIG_NODE_H_
#include <map>
#include <memory>
#include <vector>
#include "absl/log/check.h"
#include "absl/status/status.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/grappler/graph_analyzer/gen_node.h"
#include "tensorflow/core/grappler/graph_analyzer/hash_tools.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
namespace tensorflow {
namespace grappler {
namespace graph_analyzer {
namespace test {
class SigBaseTest;
} // end namespace test
class SigNode;
// To find nodes by name. Having the map ordered makes the tests easier,
// and it isn't used in production code often enough to get any win from
// using an unordered map.
using SigNodeMap = std::map<std::string, std::unique_ptr<SigNode>>;
// One node in the graph, in the form convenient for generation of the signature
// of the graph, and comparison of two (sub)graphs for equivalence. It refers to
// the original NodeDef protobuf for most information and adds the extra
// enrichment.
//
// The graph building is 2-stage: first match a SigNode with each NodeDef and
// collect them into a map that finds them by name, then process the map,
// deep-parse the underlying NodeDefs and connect the SigNodes together.
class SigNode {
public:
friend struct Signature;
// Will keep the pointer to the underlying NodeDef, so that
// underlying object must not be deleted while SigNode is alive.
explicit SigNode(const NodeDef* node);
// Access wrappers.
const std::string& name() const { return node_->name(); }
const std::string& opcode() const { return node_->op(); }
const NodeDef* node_def() const { return node_; }
// For extraction of subgraphs into a separate SigNodeMap, copies the links
// that point inside the subgraph from a full-graph SigNode to a subgraph
// SigNode. The translation map defines the subgraph and gives the mapping
// from the nodes in the full graph to the matching nodes in subgraph.
using TranslationMap =
std::unordered_map<const GenNode* /*full_graph*/, SigNode* /*subgraph*/>;
void CopyLinks(const GenNode& from, const TranslationMap& map);
// A link is an edge of the graph that connects 2 nodes. Each of the connected
// nodes has its own perspective on the link, seeing its local port, remote
// port and the remote node. The direction of the link is encoded in the
// ports, one port is always incoming and another one outgoing.
//
// The link tag here contains both ports of the link viewed from the
// perspective of this node; consisting of both the local port (i.e. at this
// node) and remote port (i.e. on the other node), the local one going first.
struct LinkTag {
struct Hasher {
size_t operator()(const LinkTag& tag) const noexcept {
size_t hval = port_hasher(tag.local);
CombineHash(port_hasher(tag.remote), &hval);
return hval;
}
GenNode::Port::Hasher port_hasher;
};
LinkTag(GenNode::Port a_local, GenNode::Port a_remote)
: local(a_local), remote(a_remote) {}
// The default constructor is used for the default values in maps.
// (false, 99) is an arbitrary value that makes the uninitialized
// links easy to tell when debugging (they should never happen).
LinkTag() : local(false, 99), remote(false, 99) {}
// Port of the link on the local node.
GenNode::Port local;
// Port of the link on the remote node.
GenNode::Port remote;
bool operator==(const LinkTag& other) const {
return local == other.local && remote == other.remote;
}
bool operator<(const LinkTag& other) const {
return local < other.local ||
(local == other.local && remote < other.remote);
}
};
// Since the signature logic doesn't differentiate between the links
// with the same tag (other than by the "peer" nodes on their other ends),
// all the links with the same tag are grouped into a single structure.
struct Link {
LinkTag tag;
size_t unique_hash; // Hash of the tag after conflict resolution.
// The remote node(s) on the other side on the link(s).
using PeerVector = std::vector<SigNode*>;
PeerVector peers;
};
// A way to look up the link description by its hash.
using LinkHashMap = std::map<size_t, Link>;
const LinkHashMap& hash_to_link() const { return hash_to_link_; }
// The enumeration of all the peer nodes in a predictable order.
// Before the signature generation, only the link values determine the
// order, after the signature generation the entries at the same
// links get further sorted by their peer node ranks.
struct HashedPeer {
HashedPeer(size_t l, SigNode* p) : link_hash(l), peer(p) {}
struct LessByRank {
bool operator()(const SigNode::HashedPeer& left,
const SigNode::HashedPeer& right) {
return left.peer->unique_rank_ < right.peer->unique_rank_;
}
};
size_t link_hash;
SigNode* peer;
};
using HashedPeerVector = std::vector<HashedPeer>;
const HashedPeerVector& hashed_peers() const { return hashed_peers_; }
// Compares two nodes in two different graphs for equivalence (two nodes in
// the same graph would never be equivalent). Expects that the signatures of
// the graphs have already been computed, so unique_rank_ is filled in and
// the hashed_peers_ properly ordered.
bool operator==(const SigNode& other) const;
bool operator!=(const SigNode& other) const { return !(*this == other); }
private:
friend class test::SigBaseTest;
// The CopyLinks code is split into 2 parts for testability.
// The first pass builds a map ordered by LinkTag for predictability.
void CopyLinksPass1(const GenNode& from, const TranslationMap& map,
std::map<LinkTag, Link>* link_map);
// The second pass converts to the map by hash value,
// resolves any hash conflicts, and builds the hashed peer vector.
void CopyLinksPass2(std::map<LinkTag, Link>* link_map);
// Computes the topological hash at distance 0. Resets the topo_hash_ vector
// and hashed_nodes_;
void ComputeTopoHash0();
// Compute the topological has at the given distance. The hashes for all the
// lower distances must be already computed for all the nodes in the graph.
// Also computes next_hashed_nodes_ from last_hashed_nodes_.
void ComputeTopoHash(int distance);
// Get the hash value for a particular distance. It must be previously
// computed.
size_t GetTopoHash(int distance) const;
// The hash value for the highest computed distance. It must be previously
// computed.
size_t GetHighTopoHash() const {
CHECK(!topo_hash_.empty());
return topo_hash_.back();
}
// Rehash the topmost hash, to avoid conflicts.
void ReHighTopoHash() {
CHECK(!topo_hash_.empty());
CombineHash(1, &topo_hash_.back());
}
// Ordering by node order and highest available hash (it must be
// previously computed).
struct NodeOrderLess {
bool operator()(const SigNode* left, const SigNode* right) {
return left->topo_hash_.back() < right->topo_hash_.back();
}
};
private:
const NodeDef* node_;
// The bitmap mask with 1 bit set that represents this node in the set
// during the computation of the signature.
uint64_t node_mask_ = 0;
// The code that populates this map makes sure that there are no hash
// conflicts, rehashing if necessary.
LinkHashMap hash_to_link_;
// The enumeration of all the direct peers in the predictable order (which
// happens to be the order ot their link tags, but the order of the hashes
// would do too). It is used for the quick enumeration during the signature
// computation. After the signature building is completed, the entries that
// have the same link tag get further sorted in the order of the ranks of
// their nodes.
HashedPeerVector hashed_peers_;
// The unique rank represents the order in which the node will be included
// into the signature. It gets assigned in order either when the topo_hash_ of
// this node becomes unique in the graph, or when the nodes are completely
// equivalent, one of them is picked at random to assign the next rank, and
// then the rest of the nodes attempt to disambiguate based on that
// information.
size_t unique_rank_ = ~0;
// When hash_is_final_ is set, the topo_has_ vector stops growing, and the
// last value from it is used for all the further hashes.
bool hash_is_final_ = false;
// The hashes that include the topology of the nodes up to the distance N. The
// hash for distance 0 is produced from the attributes of this node itself and
// its general connectivity properties but no information about the
// neighboring nodes. The hash for distance D+1 is build from hashes at level
// D of this node and of all its immediate neighbors. The neighbors that are
// connected by equivalent links are included in a commutative way.
std::vector<size_t> topo_hash_;
// The set of nodes that got included into the computation of the
// last topo_hash_ entry.
uint64_t last_hashed_nodes_ = 0;
// The next set of nodes that gets used for the current topo_hash entry.
uint64_t next_hashed_nodes_ = 0;
};
// Signature of a graph. The computation is intertwined with the private methods
// of SigNode, so keeping both in the same file looks more convenient.
struct Signature {
friend class test::SigBaseTest;
// Maximal size of the graphs for which the signature can be computed.
// Changing this constant won't magically add the support for a larger size,
// the rest of implementation would have to be extended. The value of 64 is
// driven by the size of a bitset in an uint64_t, and should be enough for our
// purposes, while having a high efficiency of implementation.
static constexpr int kMaxGraphSize = 64;
// Using the map, computes the rest of the fields of a signature.
// Returns an error is the graph is too big.
absl::Status Compute();
// Convert the computed signature to a string representation.
std::string ToString() const;
SigNodeMap map; // The nodes in the graph, accessible by name.
size_t sig_short = 0; // Hash of the signature, for the quick equality check.
// The full signature: hashes of the nodes in a predictable order.
std::vector<size_t> sig_full;
// The nodes in the same order as they go in the signature.
std::vector<SigNode*> nodes;
// For building the unordered maps.
size_t Hash() const { return sig_short; }
// Returns true if the graphs are equivalent. The signature must be already
// computed.
bool operator==(const Signature& other) const;
private:
// Populates the nodes vector from the map and initializes the state of the
// nodes for the signature computation.
void PrepareNodes();
// Finds the nodes with the hashes that are unique and assigns the unique ids
// to them. If there are nodes with non-unique hashes, exactly one node from
// the first such sequence (in the order of hash values) will be picked and
// assigned a unique id. Assumes that the nodes[0...(next_node_id-1)] have
// been already assigned the unique ids. Advances next_node_id by at least 1.
void FindUniqueHashes(size_t* next_node_id_p);
// One round of the signature computation. Assumes that the
// nodes[0...(next_node_id-1)] have been already assigned the fixed
// positions, and thus computes the hashes only for the remaining nodes.
void ComputeOneRound(size_t next_node_id);
// Additional ordering of the hashed_peers_ links in the nodes, so that they
// can be compared and printed in a predictable order.
void OrderLinks();
};
} // end namespace graph_analyzer
} // end namespace grappler
} // end namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPPLER_GRAPH_ANALYZER_SIG_NODE_H_
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,243 @@
/* 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/grappler/graph_analyzer/subgraph.h"
#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <functional>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "absl/strings/str_format.h"
#include "absl/strings/str_join.h"
#include "tensorflow/core/grappler/graph_analyzer/hash_tools.h"
namespace tensorflow {
namespace grappler {
namespace graph_analyzer {
//=== Subgraph::Identity
Subgraph::Identity::Identity(InitializerList init) {
for (auto element : init) {
insert(element);
}
}
bool Subgraph::Identity::operator<(const Identity& other) const {
// Shorter sets go first.
if (this->size() < other.size()) {
return true;
}
if (this->size() > other.size()) {
return false;
}
for (auto lit = this->begin(), rit = other.begin(); lit != this->end();
++lit, ++rit) {
if (*lit < *rit) {
return true;
}
if (*lit > *rit) {
return false;
}
}
return false; // Equal.
}
bool Subgraph::Identity::operator==(const Identity& other) const {
if (this->size() != other.size()) {
return false;
}
for (auto lit = this->begin(), rit = other.begin(); lit != this->end();
++lit, ++rit) {
if (*lit != *rit) {
return false;
}
}
return true; // Equal.
}
size_t Subgraph::Identity::Hash() const {
std::hash<const GenNode*> hasher;
size_t result = 0;
for (auto ptr : *this) {
CombineHash(hasher(ptr), &result);
}
return result;
}
std::string Subgraph::Dump() {
// TODO(babkin): this is simplified for now.
std::vector<std::string> nodes;
for (const auto& n : id_) {
if (specific_) {
nodes.emplace_back(absl::StrFormat("%s(%s)", n->opcode(), n->name()));
} else {
nodes.emplace_back(n->opcode());
}
}
std::sort(nodes.begin(), nodes.end());
return absl::StrFormat("%d: ", collation_count_) + absl::StrJoin(nodes, ", ");
}
void Subgraph::ExtractForSignature(SigNodeMap* result) {
// Mapping of nodes from the original graph to the new one.
SigNode::TranslationMap full_to_new;
for (auto node : id_) {
auto newnode_ref = std::make_unique<SigNode>(node->node_def());
auto newnode = newnode_ref.get();
(*result)[node->name()] = std::move(newnode_ref);
full_to_new[node] = newnode;
}
for (const auto& mapping : full_to_new) {
mapping.second->CopyLinks(*mapping.first, full_to_new);
}
}
//=== Subgraph
Subgraph::Subgraph(const Identity& parent_id, GenNode* add_node)
: id_(parent_id) {
id_.insert(add_node);
hash_ = id_.Hash();
}
//=== SubgraphIterator
SubgraphIterator::SubgraphIterator(const Subgraph::Identity* id)
: id_(id), id_it_(id_->begin()) {
if (!id_->empty()) {
link_map_it_ = (*id_it_)->links().begin();
// In case if the node has no links.
while (link_map_it_ == (*id_it_)->links().end()) {
if (++id_it_ == id_->end()) {
return;
}
link_map_it_ = (*id_it_)->links().begin();
}
link_idx_ = 0;
// The LinkTargetVector should never be empty but just in case safeguard
// against that too.
PropagateNext();
}
}
bool SubgraphIterator::Next() {
if (AtEnd()) {
return false;
}
++link_idx_;
return PropagateNext();
}
bool SubgraphIterator::NextIfSamePort() {
if (AtEnd()) {
return false;
}
const int64_t link_map_it_second_size = link_map_it_->second.size();
if (link_idx_ + 1 < link_map_it_second_size) {
++link_idx_;
return true;
} else {
return false;
}
}
void SubgraphIterator::SkipPort() {
if (AtEnd()) {
return;
}
link_idx_ = link_map_it_->second.size() - 1;
}
void SubgraphIterator::SkipNode() {
if (AtEnd()) {
return;
}
for (auto next = link_map_it_; next != (*id_it_)->links().end(); ++next) {
link_map_it_ = next;
}
link_idx_ = link_map_it_->second.size() - 1;
}
bool SubgraphIterator::PropagateNext() {
// Loops are used to skip over the empty entries.
const int64_t link_map_it_second_size = link_map_it_->second.size();
while (link_idx_ >= link_map_it_second_size) {
++link_map_it_;
while (link_map_it_ == (*id_it_)->links().end()) {
if (++id_it_ == id_->end()) {
return false;
}
link_map_it_ = (*id_it_)->links().begin();
}
link_idx_ = 0;
}
return true;
}
bool SubgraphIterator::operator==(const SubgraphIterator& other) const {
if (id_ != other.id_) {
return false;
}
if (id_it_ != other.id_it_) {
return false;
}
// When AtEnd(), the rest of the fields are not valid.
if (AtEnd()) {
return true;
}
if (link_map_it_ != other.link_map_it_) {
return false;
}
if (link_idx_ != other.link_idx_) {
return false;
}
return true;
}
//=== SubgraphPtrSet
Subgraph* SubgraphPtrSet::ExtendParent(const Subgraph::Identity& parent_id,
GenNode* node) {
if (parent_id.find(node) != parent_id.end()) {
// This was another link to the node that is already in the parent.
return nullptr;
}
// Constructing an object just to check that an equivalent one is already
// present is kind of ugly but storing the references rather than the objects
// in the set avoids the need to make the object copyable.
auto sg = std::make_unique<Subgraph>(parent_id, node);
if (find(sg) != end()) {
// This subgraph was already found by extending from a different path.
return nullptr;
}
Subgraph* ptr = sg.get();
insert(std::move(sg));
return ptr;
}
} // end namespace graph_analyzer
} // end namespace grappler
} // end namespace tensorflow
@@ -0,0 +1,190 @@
/* 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_GRAPPLER_GRAPH_ANALYZER_SUBGRAPH_H_
#define TENSORFLOW_CORE_GRAPPLER_GRAPH_ANALYZER_SUBGRAPH_H_
#include <initializer_list>
#include <set>
#include <unordered_set>
#include "tensorflow/core/grappler/graph_analyzer/gen_node.h"
#include "tensorflow/core/grappler/graph_analyzer/map_tools.h"
#include "tensorflow/core/grappler/graph_analyzer/sig_node.h"
#include "tensorflow/core/lib/gtl/flatset.h"
namespace tensorflow {
namespace grappler {
namespace graph_analyzer {
// The description of a single subgraph for processing.
class Subgraph {
public:
// Identity of a single subgraph as a set of nodes.
class Identity : public gtl::FlatSet<const GenNode*> {
public:
using InitializerList = std::initializer_list<GenNode*>;
Identity() = default;
Identity(InitializerList init);
bool operator<(const Identity& other) const;
bool operator==(const Identity& other) const;
// Compute the hash.
size_t Hash() const;
};
explicit Subgraph(Identity id) : id_(std::move(id)), hash_(id_.Hash()) {}
// Construct by extending the parent identity with an extra node.
Subgraph(const Identity& parent_id, GenNode* add_node);
Subgraph() = delete;
Subgraph(const Subgraph& other) = delete;
void operator=(const Subgraph& other) = delete;
// Order for building sets of subgraphs.
bool operator<(const Subgraph& other) const { return this->id_ < other.id_; }
// Support for hashed sets.
bool operator==(const Subgraph& other) const {
return this->id_ == other.id_;
}
size_t Hash() const { return hash_; }
// Dump the subgraph information to a string.
std::string Dump();
// Extract this subgraph into a separate graph representation for signature
// building, that includes only the links between the nodes in the subgraph
// and drops all the external links. The result map should be clear before the
// call.
void ExtractForSignature(SigNodeMap* result);
const Identity& id() const { return id_; }
bool specific() const { return specific_; }
void SetSpecific(bool value) { specific_ = value; }
int32_t collation_count() const { return collation_count_; }
void AddCollation(int32_t n = 1) { collation_count_ += n; }
void ResetCollation() { collation_count_ = 1; }
void MergeCollation(const Subgraph& other) {
collation_count_ += other.collation_count_;
}
private:
// Identity also serves as the list of nodes. It never changes throughout the
// life of subgraph.
Identity id_;
size_t hash_; // Cached from the identity.
// Whether the dump should include the specific names of the nodes. The
// non-specific (i.e. generic) subgraphs represent a collation of multiple
// subgraphs.
bool specific_ = true;
// How many collated subgraphs are represented by this subgraph.
int32_t collation_count_ = 1;
};
// Iteration of all links in a subgraph. This is more like Java iterators than
// the normal C++ iterators. It's simpler this way and there seems to be no
// major reason to make it a proper C++ iterator.
class SubgraphIterator {
public:
// Obviously an iterator is valid only until the original object
// gets destroyed.
explicit SubgraphIterator(const Subgraph::Identity* id);
explicit SubgraphIterator(const Subgraph* sg) : SubgraphIterator(&sg->id()) {}
// Check whether the built-in iterator is at the end.
bool AtEnd() const { return id_it_ == id_->end(); }
// Get the neighbor at the current iterator.
// MUST NOT be called when AtEnd();
const GenNode::LinkTarget& GetNeighbor() const {
return link_map_it_->second[link_idx_];
}
// Get the node at the current iterator.
// MUST NOT be called when AtEnd();
const GenNode* GetNode() const { return *id_it_; }
// Get the port leading to the neighbor at the current iterator.
// MUST NOT be called when AtEnd();
GenNode::Port GetPort() const { return link_map_it_->first; }
// Increases the iterator.
// Returns true if NOT AtEnd() after increasing the iterator.
// Safe to call if already AtEnd().
bool Next();
// If there are more links at the same port, increases the iterator and
// returns true. Otherwise leaves the iterator unchanged and returns false.
bool NextIfSamePort();
// Increases the iterator directly to the last position on the current port
// (or if already there then doesn't increase). Equivalent to calling
// NextIfSamePort() while it returns true, but faster.
// Safe to call if already AtEnd().
void SkipPort();
// Increases the iterator directly to the last position on the current node.
// Safe to call if already AtEnd().
void SkipNode();
// Returns true if the iterators are exactly the same.
bool operator==(const SubgraphIterator& other) const;
bool operator!=(const SubgraphIterator& other) const {
return !(*this == other);
}
private:
// After link_idx_ has been increased, make sure that it points to the
// next valid element (or end) by increasing the higher levels of iteration if
// needed.
// Returns true if NOT AtEnd() after increasing the iterator.
// NOT safe to call if already AtEnd().
bool PropagateNext();
// Identity of the subgraph being iterated over.
const Subgraph::Identity* id_;
// The current position, allowing to iterate through the links (see the
// reasoning for it in the public section).
//
// (1) Iterator of the nodes in the subgraph.
Subgraph::Identity::const_iterator id_it_;
// (2) Iterator in the link map of the node.
GenNode::LinkMap::const_iterator link_map_it_;
// (3) Index in the vector of the links.
int32_t link_idx_;
};
// A convenient way to store subgraphs: in a set of unique_ptrs. This way the
// addresses of subgraph objects will stay stable, and the objects themselves
// won't be copied.
class SubgraphPtrSet
: public std::unordered_set<std::unique_ptr<Subgraph>,
HashAtPtr<std::unique_ptr<Subgraph>>,
EqAtPtr<std::unique_ptr<Subgraph>>> {
public:
// Attempts to extend the set by adding a new subgraph that gets created by
// adding one node to the parent subgraph. If such a subgraph already exists,
// returns nullptr, otherwise returns the pointer to the new subgraph.
Subgraph* ExtendParent(const Subgraph::Identity& parent_id, GenNode* node);
};
} // end namespace graph_analyzer
} // end namespace grappler
} // end namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPPLER_GRAPH_ANALYZER_SUBGRAPH_H_
@@ -0,0 +1,351 @@
/* 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/grappler/graph_analyzer/subgraph.h"
#include <algorithm>
#include <memory>
#include <string>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/status/status.h"
#include "absl/strings/str_format.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/grappler/graph_analyzer/test_tools.h"
namespace tensorflow {
namespace grappler {
namespace graph_analyzer {
namespace test {
namespace {
using ::testing::ElementsAre;
using ::testing::Eq;
using ::testing::Ne;
TEST(SubgraphTest, Comparison) {
GraphDef graph;
// A topology with a loop.
(*graph.add_node()) = MakeNodeConst("node1");
(*graph.add_node()) = MakeNodeConst("node2");
GenNodeMap map;
ASSERT_THAT(GenNode::BuildGraphInMap(graph, &map), Eq(absl::OkStatus()));
auto gn1 = map["node1"].get();
auto gn2 = map["node2"].get();
ASSERT_THAT(gn1, Ne(nullptr));
ASSERT_THAT(gn2, Ne(nullptr));
Subgraph::Identity id1;
Subgraph::Identity id2;
id1.insert(gn1);
id2.insert(gn2);
Subgraph sg1(id1);
Subgraph sg2(id2);
EXPECT_TRUE(id1 == sg1.id());
EXPECT_TRUE(id2 == sg2.id());
EXPECT_THAT(sg1 < sg2, Eq(id1 < id2));
}
TEST(SubgraphTest, EmptyIteration) {
NodeDef node1 = MakeNodeConst("node1");
auto gn1 = std::make_unique<GenNode>(&node1);
Subgraph::Identity id1;
id1.insert(gn1.get());
Subgraph sg1(id1);
SubgraphIterator sit(&sg1);
EXPECT_TRUE(sit.AtEnd());
EXPECT_FALSE(sit.Next());
EXPECT_TRUE(sit.AtEnd());
SubgraphIterator sit2(&sg1);
EXPECT_TRUE(sit == sit2);
}
TEST(SubgraphTest, Iteration) {
GraphDef graph;
// A topology with a loop.
(*graph.add_node()) = MakeNodeConst("node1");
(*graph.add_node()) = MakeNodeSub("node2", "node3:1", "node3:0");
auto node3 = graph.add_node();
*node3 = MakeNodeBroadcastGradientArgs("node3", "node1", "node2");
node3->add_input("^node3"); // The control link goes back to self.
GenNodeMap map;
ASSERT_THAT(GenNode::BuildGraphInMap(graph, &map), Eq(absl::OkStatus()));
ASSERT_THAT(map.find("node3"), Ne(map.end()));
Subgraph::Identity id;
id.insert(map["node3"].get());
Subgraph sg(id);
// node3 has 2 incoming data links, 2 outgoing data , 1 control incoming, 1
// control outgoing = total of 6
{
SubgraphIterator sit(&sg);
EXPECT_FALSE(sit.AtEnd()); // 1
EXPECT_TRUE(sit.Next());
EXPECT_FALSE(sit.AtEnd()); // 2
EXPECT_TRUE(sit.Next());
EXPECT_FALSE(sit.AtEnd()); // 3
EXPECT_TRUE(sit.Next());
EXPECT_FALSE(sit.AtEnd()); // 4
EXPECT_TRUE(sit.Next());
EXPECT_FALSE(sit.AtEnd()); // 5
EXPECT_TRUE(sit.Next());
EXPECT_FALSE(sit.AtEnd()); // 6
EXPECT_FALSE(sit.Next());
EXPECT_TRUE(sit.AtEnd());
}
// Now get the values out. And more equality testing along the way.
{
SubgraphIterator sit(&sg);
SubgraphIterator sit2(&sg);
std::vector<std::string> links;
for (; !sit.AtEnd(); sit.Next()) {
EXPECT_TRUE(sit == sit2);
sit2.Next();
EXPECT_FALSE(sit == sit2);
links.push_back(absl::StrFormat("[%s,%s,%s]", std::string(sit.GetPort()),
sit.GetNeighbor().node->name(),
std::string(sit.GetNeighbor().port)));
}
EXPECT_TRUE(sit == sit2);
std::sort(links.begin(), links.end());
// clang-format off
EXPECT_THAT(links, ElementsAre(
"[i0,node1,o0]",
"[i1,node2,o0]",
"[iC,node3,oC]",
"[o0,node2,i1]",
"[o1,node2,i0]",
"[oC,node3,iC]"
));
// clang-format on
}
}
TEST(SubgraphTest, IterationSamePort) {
GraphDef graph;
(*graph.add_node()) = MakeNodeConst("node1");
(*graph.add_node()) = MakeNodeSub("node2", "node3", "node3");
(*graph.add_node()) = MakeNodeAddN("node3", "node1", "node2");
GenNodeMap map;
ASSERT_THAT(GenNode::BuildGraphInMap(graph, &map), Eq(absl::OkStatus()));
ASSERT_THAT(map.find("node3"), Ne(map.end()));
Subgraph::Identity id;
id.insert(map["node3"].get());
Subgraph sg(id);
int total_links = 0;
for (SubgraphIterator sit(&sg); !sit.AtEnd(); sit.Next()) {
++total_links;
}
// Initialize the port as control, which doesn't occur in this graph.
GenNode::Port last_port(false, -1);
int steps_total_same_port = 0;
int steps_with_same_port = 0;
for (SubgraphIterator sit(&sg); !sit.AtEnd(); sit.Next()) {
GenNode::Port new_port = sit.GetPort();
EXPECT_THAT(last_port.Encoded(), Ne(new_port.Encoded()))
<< "At step " << steps_total_same_port;
last_port = new_port;
++steps_total_same_port;
SubgraphIterator sit2(sit);
sit2.SkipPort();
while (sit.NextIfSamePort()) {
new_port = sit.GetPort();
EXPECT_THAT(last_port.Encoded(), Eq(new_port.Encoded()))
<< "At step " << steps_total_same_port;
++steps_total_same_port;
++steps_with_same_port;
}
EXPECT_TRUE(sit == sit2);
}
EXPECT_THAT(steps_total_same_port, Eq(total_links));
// There is one 2-way input and one 2-way output.
EXPECT_THAT(steps_with_same_port, Eq(2));
}
TEST(SubgraphTest, IterationSameNode) {
GraphDef graph;
(*graph.add_node()) = MakeNodeConst("node1");
(*graph.add_node()) = MakeNodeSub("node2", "node3", "node3");
(*graph.add_node()) = MakeNodeAddN("node3", "node1", "node2");
GenNodeMap map;
ASSERT_THAT(GenNode::BuildGraphInMap(graph, &map), Eq(absl::OkStatus()));
ASSERT_THAT(map.find("node3"), Ne(map.end()));
Subgraph::Identity id;
id.insert(map["node3"].get());
Subgraph sg(id);
const GenNode* last_node = nullptr;
SubgraphIterator sit(&sg);
while (!sit.AtEnd()) {
const GenNode* new_node = sit.GetNode();
EXPECT_THAT(new_node, Ne(last_node)) << "At node " << new_node->name();
SubgraphIterator sit2(sit);
sit2.SkipNode();
ASSERT_FALSE(sit2.AtEnd());
EXPECT_THAT(sit2.GetNode(), Eq(new_node))
<< "At expected node " << new_node->name() << ", got "
<< sit2.GetNode()->name();
while (sit != sit2 && !sit.AtEnd()) {
sit.Next();
}
ASSERT_FALSE(sit.AtEnd());
EXPECT_THAT(sit.GetNode(), Eq(new_node))
<< "At expected node " << new_node->name() << ", got "
<< sit2.GetNode()->name();
sit.Next();
last_node = new_node;
}
// Check that it doesn't fail if already at end.
sit.SkipNode();
EXPECT_TRUE(sit.AtEnd());
}
TEST(SubgraphTest, ExtendSet) {
GraphDef graph;
// A topology with a loop.
(*graph.add_node()) = MakeNodeConst("node1");
(*graph.add_node()) = MakeNodeSub("node2", "node3:1", "node3:0");
auto node3 = graph.add_node();
*node3 = MakeNodeBroadcastGradientArgs("node3", "node1", "node2");
node3->add_input("^node3"); // The control link goes back to self.
GenNodeMap map;
ASSERT_THAT(GenNode::BuildGraphInMap(graph, &map), Eq(absl::OkStatus()));
ASSERT_THAT(map.find("node2"), Ne(map.end()));
ASSERT_THAT(map.find("node3"), Ne(map.end()));
Subgraph::Identity id_empty;
Subgraph::Identity id3;
id3.insert(map["node3"].get());
Subgraph::Identity id23 = id3;
id23.insert(map["node2"].get());
Subgraph* sg;
SubgraphPtrSet set;
// Extend an empty identity.
sg = set.ExtendParent(id_empty, map["node3"].get());
EXPECT_THAT(set.size(), Eq(1));
ASSERT_THAT(sg, Ne(nullptr));
EXPECT_TRUE(sg->id() == id3);
// Extend with a node that is already in the parent.
sg = set.ExtendParent(id3, map["node3"].get());
EXPECT_THAT(set.size(), Eq(1));
EXPECT_THAT(sg, Eq(nullptr));
// Extend to a 2-node subgraph.
sg = set.ExtendParent(id3, map["node2"].get());
EXPECT_THAT(set.size(), Eq(2));
ASSERT_THAT(sg, Ne(nullptr));
EXPECT_TRUE(sg->id() == id23);
// The second insert of the same node gets ignored.
sg = set.ExtendParent(id3, map["node2"].get());
EXPECT_THAT(set.size(), Eq(2));
EXPECT_THAT(sg, Eq(nullptr));
}
TEST(SubgraphTest, ExtractForSignature) {
GraphDef graph;
(*graph.add_node()) = MakeNodeConst("node1");
(*graph.add_node()) = MakeNodeSub("node2", "node3:1", "node3:0");
auto node3 = graph.add_node();
*node3 = MakeNodeBroadcastGradientArgs("node3", "node1", "node2");
node3->add_input("^node1");
node3->add_input("^node2");
node3->add_input("^node3"); // The control link goes back to self.
GenNodeMap map;
ASSERT_THAT(GenNode::BuildGraphInMap(graph, &map), Eq(absl::OkStatus()));
ASSERT_THAT(map.find("node1"), Ne(map.end()));
ASSERT_THAT(map.find("node2"), Ne(map.end()));
ASSERT_THAT(map.find("node3"), Ne(map.end()));
Subgraph::Identity id;
id.insert(map["node1"].get());
id.insert(map["node3"].get());
Subgraph sg(id);
SigNodeMap map2;
sg.ExtractForSignature(&map2);
ASSERT_THAT(map2.find("node1"), Ne(map2.end()));
ASSERT_THAT(map2.find("node2"), Eq(map2.end()));
ASSERT_THAT(map2.find("node3"), Ne(map2.end()));
// clang-format off
EXPECT_THAT(DumpLinkHashMap(map2["node1"]->hash_to_link()), ElementsAre(
"oC:iC: node3",
"o0:i0: node3"
));
EXPECT_THAT(DumpHashedPeerVector(map2["node1"]->hashed_peers()), ElementsAre(
"node3",
"node3"
));
EXPECT_THAT(DumpLinkHashMap(map2["node3"]->hash_to_link()), ElementsAre(
"oC:iC: node3",
"iC:oC: node1, node3",
"i0:o0: node1"
));
EXPECT_THAT(DumpHashedPeerVector(map2["node3"]->hashed_peers()), ElementsAre(
"node3",
"node1",
"node3",
"node1"
));
// clang-format on
}
} // end namespace
} // end namespace test
} // end namespace graph_analyzer
} // end namespace grappler
} // end namespace tensorflow
@@ -0,0 +1,309 @@
/* 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/grappler/graph_analyzer/test_tools.h"
#include <algorithm>
#include <cstddef>
#include <map>
#include <set>
#include <string>
#include <vector>
#include "absl/strings/str_format.h"
#include "absl/strings/str_join.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
namespace tensorflow {
namespace grappler {
namespace graph_analyzer {
namespace test {
//=== Helper methods to construct the nodes.
NodeDef MakeNodeConst(const std::string& name) {
NodeDef n;
n.set_name(name);
n.set_op("Const");
return n;
}
NodeDef MakeNode2Arg(const std::string& name, const std::string& opcode,
const std::string& arg1, const std::string& arg2) {
NodeDef n;
n.set_name(name);
n.set_op(opcode);
n.add_input(arg1);
n.add_input(arg2);
return n;
}
NodeDef MakeNode4Arg(const std::string& name, const std::string& opcode,
const std::string& arg1, const std::string& arg2,
const std::string& arg3, const std::string& arg4) {
NodeDef n;
n.set_name(name);
n.set_op(opcode);
n.add_input(arg1);
n.add_input(arg2);
n.add_input(arg3);
n.add_input(arg4);
return n;
}
// Not really a 2-argument but convenient to construct.
NodeDef MakeNodeShapeN(const std::string& name, const std::string& arg1,
const std::string& arg2) {
// This opcode is multi-input but not commutative.
return MakeNode2Arg(name, "ShapeN", arg1, arg2);
}
// Not really a 2-argument but convenient to construct.
NodeDef MakeNodeIdentityN(const std::string& name, const std::string& arg1,
const std::string& arg2) {
// The argument is of a list type.
return MakeNode2Arg(name, "IdentityN", arg1, arg2);
}
NodeDef MakeNodeQuantizedConcat(const std::string& name,
const std::string& arg1,
const std::string& arg2,
const std::string& arg3,
const std::string& arg4) {
// This opcode has multiple multi-inputs.
return MakeNode4Arg(name, "QuantizedConcat", arg1, arg2, arg3, arg4);
}
//=== Helper methods for analysing the structures.
std::vector<std::string> DumpLinkMap(const GenNode::LinkMap& link_map) {
// This will order the entries first.
std::map<std::string, std::string> ordered;
for (const auto& link : link_map) {
std::string key = std::string(link.first);
// Order the other sides too. They may be repeating, so store them
// in a multiset.
std::multiset<std::string> others;
for (const auto& other : link.second) {
others.emplace(absl::StrFormat("%s[%s]", other.node->name(),
std::string(other.port)));
}
ordered[key] = absl::StrJoin(others, ", ");
}
// Now dump the result in a predictable order.
std::vector<std::string> result;
result.reserve(ordered.size());
for (const auto& link : ordered) {
result.emplace_back(link.first + ": " + link.second);
}
return result;
}
std::vector<std::string> DumpLinkHashMap(
const SigNode::LinkHashMap& link_hash_map) {
// The entries in this map are ordered by hash value which might change
// at any point. Re-order them by the link tag.
std::map<SigNode::LinkTag, size_t> tags;
for (const auto& entry : link_hash_map) {
tags[entry.second.tag] = entry.first;
}
std::vector<std::string> result;
for (const auto& id : tags) {
// For predictability, the nodes need to be sorted.
std::vector<std::string> nodes;
for (const auto& peer : link_hash_map.at(id.second).peers) {
nodes.emplace_back(peer->name());
}
std::sort(nodes.begin(), nodes.end());
result.emplace_back(std::string(id.first.local) + ":" +
std::string(id.first.remote) + ": " +
absl::StrJoin(nodes, ", "));
}
return result;
}
std::vector<std::string> DumpHashedPeerVector(
const SigNode::HashedPeerVector& hashed_peers) {
std::vector<std::string> result;
// Each subset of nodes with the same hash has to be sorted by name.
// Other than that, the vector is already ordered by full tags.
size_t last_hash = 0;
// Index, since iterators may get invalidated on append.
size_t subset_start = 0;
for (const auto& entry : hashed_peers) {
if (entry.link_hash != last_hash) {
std::sort(result.begin() + subset_start, result.end());
subset_start = result.size();
}
result.emplace_back(entry.peer->name());
}
std::sort(result.begin() + subset_start, result.end());
return result;
}
TestGraphs::TestGraphs() {
{
GraphDef& graph = graph_3n_self_control_;
// The topology includes a loop and a link to self.
(*graph.add_node()) = MakeNodeConst("node1");
(*graph.add_node()) = MakeNodeSub("node2", "node3:1", "node3:0");
auto node3 = graph.add_node();
*node3 = MakeNodeBroadcastGradientArgs("node3", "node1", "node2");
node3->add_input("^node3"); // The control link goes back to self.
}
{
GraphDef& graph = graph_multi_input_;
// The topology includes a loop and a link to self.
(*graph.add_node()) = MakeNodeConst("const1_1");
(*graph.add_node()) = MakeNodeConst("const1_2");
(*graph.add_node()) = MakeNodeAddN("add1", "const1_1", "const1_2");
(*graph.add_node()) = MakeNodeConst("const2_1");
(*graph.add_node()) = MakeNodeConst("const2_2");
(*graph.add_node()) = MakeNodeConst("const2_3");
auto add2 = graph.add_node();
*add2 = MakeNodeAddN("add2", "const2_1", "const2_2");
// The 3rd node is connected twice, to 4 links total.
add2->add_input("const2_3");
add2->add_input("const2_3");
(*graph.add_node()) = MakeNodeSub("sub", "add1", "add2");
}
{
GraphDef& graph = graph_all_or_none_;
// The topology includes a loop and a link to self.
(*graph.add_node()) = MakeNodeConst("const1_1");
(*graph.add_node()) = MakeNodeConst("const1_2");
auto pass1 = graph.add_node();
*pass1 = MakeNodeIdentityN("pass1", "const1_1", "const1_2");
(*graph.add_node()) = MakeNodeConst("const2_1");
(*graph.add_node()) = MakeNodeConst("const2_2");
(*graph.add_node()) = MakeNodeConst("const2_3");
auto pass2 = graph.add_node();
*pass2 = MakeNodeIdentityN("pass2", "const2_1", "const2_2");
// The 3rd node is connected twice, to 4 links total.
pass2->add_input("const2_3");
pass2->add_input("const2_3");
// Add the control links, they get handled separately than the normal
// links.
pass1->add_input("^const2_1");
pass1->add_input("^const2_2");
pass1->add_input("^const2_3");
(*graph.add_node()) = MakeNodeSub("sub", "pass1", "pass2");
}
{
GraphDef& graph = graph_circular_onedir_;
(*graph.add_node()) = MakeNodeMul("node1", "node5", "node5");
(*graph.add_node()) = MakeNodeMul("node2", "node1", "node1");
(*graph.add_node()) = MakeNodeMul("node3", "node2", "node2");
(*graph.add_node()) = MakeNodeMul("node4", "node3", "node3");
(*graph.add_node()) = MakeNodeMul("node5", "node4", "node4");
}
{
GraphDef& graph = graph_circular_bidir_;
// The left and right links are intentionally mixed up.
(*graph.add_node()) = MakeNodeMul("node1", "node5", "node2");
(*graph.add_node()) = MakeNodeMul("node2", "node3", "node1");
(*graph.add_node()) = MakeNodeMul("node3", "node2", "node4");
(*graph.add_node()) = MakeNodeMul("node4", "node5", "node3");
(*graph.add_node()) = MakeNodeMul("node5", "node4", "node1");
}
{
GraphDef& graph = graph_linear_;
(*graph.add_node()) = MakeNodeConst("node1");
(*graph.add_node()) = MakeNodeMul("node2", "node1", "node1");
(*graph.add_node()) = MakeNodeMul("node3", "node2", "node2");
(*graph.add_node()) = MakeNodeMul("node4", "node3", "node3");
(*graph.add_node()) = MakeNodeMul("node5", "node4", "node4");
}
{
GraphDef& graph = graph_cross_;
(*graph.add_node()) = MakeNodeConst("node1");
(*graph.add_node()) = MakeNodeMul("node2", "node1", "node1");
(*graph.add_node()) = MakeNodeConst("node3");
(*graph.add_node()) = MakeNodeMul("node4", "node3", "node3");
(*graph.add_node()) = MakeNodeConst("node5");
(*graph.add_node()) = MakeNodeMul("node6", "node5", "node5");
(*graph.add_node()) = MakeNodeConst("node7");
(*graph.add_node()) = MakeNodeMul("node8", "node7", "node7");
auto center = graph.add_node();
*center = MakeNodeMul("node9", "node2", "node4");
center->add_input("node6");
center->add_input("node8");
}
{
GraphDef& graph = graph_small_cross_;
(*graph.add_node()) = MakeNodeConst("node1");
(*graph.add_node()) = MakeNodeConst("node2");
(*graph.add_node()) = MakeNodeConst("node3");
(*graph.add_node()) = MakeNodeConst("node4");
auto center = graph.add_node();
*center = MakeNodeMul("node5", "node1", "node2");
center->add_input("node3");
center->add_input("node4");
}
{
GraphDef& graph = graph_for_link_order_;
(*graph.add_node()) = MakeNodeConst("node1");
(*graph.add_node()) = MakeNodeConst("node2");
(*graph.add_node()) = MakeNodeConst("node3");
(*graph.add_node()) = MakeNodeConst("node4");
// One group of equivalent links.
auto center = graph.add_node();
*center = MakeNodeMul("node5", "node1", "node2");
center->add_input("node3");
center->add_input("node4");
// Multiple groups, separated by unique links.
auto center2 = graph.add_node();
*center2 = MakeNodeMul("node6", "node1", "node2");
center2->add_input("node2:1");
center2->add_input("node3:2");
center2->add_input("node4:2");
center2->add_input("node4:3");
}
{
GraphDef& graph = graph_sun_;
(*graph.add_node()) = MakeNodeConst("node1");
(*graph.add_node()) = MakeNodeConst("node2");
(*graph.add_node()) = MakeNodeConst("node3");
(*graph.add_node()) = MakeNodeConst("node4");
(*graph.add_node()) = MakeNodeConst("node5");
(*graph.add_node()) = MakeNodeSub("node6", "node1", "node10");
(*graph.add_node()) = MakeNodeSub("node7", "node2", "node6");
(*graph.add_node()) = MakeNodeSub("node8", "node3", "node7");
(*graph.add_node()) = MakeNodeSub("node9", "node4", "node8");
(*graph.add_node()) = MakeNodeSub("node10", "node5", "node9");
}
}
} // end namespace test
} // end namespace graph_analyzer
} // end namespace grappler
} // end namespace tensorflow
@@ -0,0 +1,125 @@
/* 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_GRAPPLER_GRAPH_ANALYZER_TEST_TOOLS_H_
#define TENSORFLOW_CORE_GRAPPLER_GRAPH_ANALYZER_TEST_TOOLS_H_
#include <string>
#include <vector>
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/grappler/graph_analyzer/gen_node.h"
#include "tensorflow/core/grappler/graph_analyzer/sig_node.h"
#include "tensorflow/core/grappler/op_types.h"
namespace tensorflow {
namespace grappler {
namespace graph_analyzer {
namespace test {
//=== Helper methods to construct the nodes.
NodeDef MakeNodeConst(const std::string& name);
NodeDef MakeNode2Arg(const std::string& name, const std::string& opcode,
const std::string& arg1, const std::string& arg2);
NodeDef MakeNode4Arg(const std::string& name, const std::string& opcode,
const std::string& arg1, const std::string& arg2,
const std::string& arg3, const std::string& arg4);
inline NodeDef MakeNodeMul(const std::string& name, const std::string& arg1,
const std::string& arg2) {
return MakeNode2Arg(name, "Mul", arg1, arg2);
}
// Not really a 2-argument but convenient to construct.
inline NodeDef MakeNodeAddN(const std::string& name, const std::string& arg1,
const std::string& arg2) {
return MakeNode2Arg(name, "AddN", arg1, arg2);
}
inline NodeDef MakeNodeSub(const std::string& name, const std::string& arg1,
const std::string& arg2) {
return MakeNode2Arg(name, "Sub", arg1, arg2);
}
// Has 2 honest outputs.
inline NodeDef MakeNodeBroadcastGradientArgs(const std::string& name,
const std::string& arg1,
const std::string& arg2) {
return MakeNode2Arg(name, "BroadcastGradientArgs", arg1, arg2);
}
NodeDef MakeNodeShapeN(const std::string& name, const std::string& arg1,
const std::string& arg2);
NodeDef MakeNodeIdentityN(const std::string& name, const std::string& arg1,
const std::string& arg2);
NodeDef MakeNodeQuantizedConcat(const std::string& name,
const std::string& arg1,
const std::string& arg2,
const std::string& arg3,
const std::string& arg4);
//=== A container of pre-constructed graphs.
class TestGraphs {
public:
TestGraphs();
// Graph with 3 nodes and a control link to self (which is not valid in
// reality but adds excitement to the tests).
GraphDef graph_3n_self_control_;
// Graph that has the multi-input links.
GraphDef graph_multi_input_;
// Graph that has the all-or-none nodes.
GraphDef graph_all_or_none_;
// All the nodes are connected in a circle that goes in one direction.
GraphDef graph_circular_onedir_;
// All the nodes are connected in a circle that goes in both directions.
GraphDef graph_circular_bidir_;
// The nodes are connected in a line.
GraphDef graph_linear_;
// The nodes are connected in a cross shape.
GraphDef graph_cross_;
GraphDef graph_small_cross_;
// For testing the ordering of links at the end of signature generation,
// a variation of a cross.
GraphDef graph_for_link_order_;
// Sun-shaped, a ring with "rays".
GraphDef graph_sun_;
};
//=== Helper methods for analysing the structures.
std::vector<std::string> DumpLinkMap(const GenNode::LinkMap& link_map);
// Also checks for the consistency of hash values.
std::vector<std::string> DumpLinkHashMap(
const SigNode::LinkHashMap& link_hash_map);
std::vector<std::string> DumpHashedPeerVector(
const SigNode::HashedPeerVector& hashed_peers);
} // end namespace test
} // end namespace graph_analyzer
} // end namespace grappler
} // end namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPPLER_GRAPH_ANALYZER_TEST_TOOLS_H_
@@ -0,0 +1,212 @@
/* 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/grappler/graph_topology_view.h"
#include <algorithm>
#include "absl/container/flat_hash_map.h"
#include "absl/container/inlined_vector.h"
#include "absl/strings/string_view.h"
#include "absl/types/optional.h"
#include "absl/types/span.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
namespace tensorflow {
namespace grappler {
namespace {
template <typename T>
inline void SortAndRemoveDuplicates(T* v) {
std::sort(v->begin(), v->end());
v->erase(std::unique(v->begin(), v->end()), v->end());
}
} // namespace
absl::Status GraphTopologyView::InitializeFromGraph(
const GraphDef& graph,
const absl::Span<const GraphView::Edge> ephemeral_edges,
bool ignore_control_edges) {
if (graph_ != nullptr) {
return absl::InvalidArgumentError(
"GraphTopologyView is already initialized.");
}
graph_ = &graph;
num_nodes_ = graph.node_size();
index_to_node_name_.resize(num_nodes_);
node_name_to_index_.rehash(num_nodes_);
fanins_.resize(num_nodes_);
fanouts_.resize(num_nodes_);
// Build map from name to index and vice versa.
for (int node_idx = 0; node_idx < num_nodes_; ++node_idx) {
const NodeDef& node = graph.node(node_idx);
node_name_to_index_.emplace(node.name(), node_idx);
index_to_node_name_.emplace_back(node.name());
}
// 1. Add ephemeral edges to the adjacency lists.
for (const GraphView::Edge& edge : ephemeral_edges) {
const auto src = node_name_to_index_.find(edge.src.node->name());
const bool valid_src = src != node_name_to_index_.end();
if (!valid_src) {
const std::string error_message =
absl::StrCat("Non-existent src node: ", edge.src.node->name());
if (skip_invalid_edges_) {
VLOG(0) << "Skip error: " << error_message;
} else {
return absl::InvalidArgumentError(error_message);
}
}
const auto dst = node_name_to_index_.find(edge.dst.node->name());
const bool valid_dst = dst != node_name_to_index_.end();
if (!valid_dst) {
const std::string error_message =
absl::StrCat("Non-existent dst node: ", edge.dst.node->name());
if (skip_invalid_edges_) {
VLOG(0) << "Skip error: " << error_message;
} else {
return absl::InvalidArgumentError(error_message);
}
}
if (valid_dst && valid_src) {
const int src_idx = src->second;
const int dst_idx = dst->second;
if (ignore_control_edges && (src_idx < 0 || dst_idx < 0)) {
continue;
}
fanins_[dst_idx].push_back(src_idx);
fanouts_[src_idx].push_back(dst_idx);
}
}
// 2. Add graph edges to the adjacency lists.
for (int node_idx = 0; node_idx < num_nodes_; ++node_idx) {
const NodeDef& node = graph.node(node_idx);
fanins_[node_idx].reserve(node.input_size());
for (const std::string& input : node.input()) {
TensorId tensor = ParseTensorName(input);
if (ignore_control_edges && IsTensorIdControl(tensor)) {
continue;
}
const auto it = node_name_to_index_.find(tensor.node());
const bool valid_input = it != node_name_to_index_.end();
if (!valid_input) {
const std::string error_message = absl::StrCat(
"Non-existent input ", input, " in node ", node.name());
if (skip_invalid_edges_) {
VLOG(3) << "Skip error: " << error_message;
} else {
return absl::InvalidArgumentError(error_message);
}
}
if (valid_input) {
const int input_idx = it->second;
fanins_[node_idx].push_back(input_idx);
fanouts_[input_idx].push_back(node_idx);
}
}
// Dedup the input list while it's still hot in cache.
SortAndRemoveDuplicates(&fanins_[node_idx]);
}
// Dedup outputs for all the graph nodes.
for (int node_idx = 0; node_idx < num_nodes_; ++node_idx) {
SortAndRemoveDuplicates(&fanouts_[node_idx]);
}
return absl::OkStatus();
}
absl::Status GraphTopologyView::InitializeFromGraph(
const GraphDef& graph,
const absl::Span<const GraphView::Edge> ephemeral_edges) {
return InitializeFromGraph(graph, ephemeral_edges,
/*ignore_control_edges=*/false);
}
absl::Status GraphTopologyView::InitializeFromGraph(const GraphDef& graph,
bool ignore_control_edges) {
return InitializeFromGraph(graph, absl::Span<GraphView::Edge>(),
ignore_control_edges);
}
absl::Status GraphTopologyView::InitializeFromGraph(const GraphDef& graph) {
return InitializeFromGraph(graph, absl::Span<GraphView::Edge>(),
/*ignore_control_edges*/ false);
}
bool GraphTopologyView::HasNode(const absl::string_view node_name) const {
DCHECK(is_initialized()) << "GraphTopologyView is not initialized";
const auto it = node_name_to_index_.find(node_name);
return it != node_name_to_index_.end();
}
const NodeDef* GraphTopologyView::GetNode(
const absl::string_view node_name) const {
DCHECK(is_initialized()) << "GraphTopologyView is not initialized";
const auto it = node_name_to_index_.find(node_name);
return it == node_name_to_index_.end() ? nullptr : &graph_->node(it->second);
}
const NodeDef* GraphTopologyView::GetNode(int node_idx) const {
DCHECK(is_initialized()) << "GraphTopologyView is not initialized";
DCHECK(node_idx >= 0 && node_idx < num_nodes_) << "node_idx is out of range";
return &graph_->node(node_idx);
}
const absl::optional<int> GraphTopologyView::GetNodeIndex(
const absl::string_view node_name) const {
DCHECK(is_initialized()) << "GraphTopologyView is not initialized";
const auto it = node_name_to_index_.find(node_name);
DCHECK(it != node_name_to_index_.end()) << "Node doesn't exist in a graph";
return it == node_name_to_index_.end() ? std::nullopt
: std::make_optional(it->second);
}
const absl::optional<int> GraphTopologyView::GetNodeIndex(
const NodeDef& node) const {
return GetNodeIndex(node.name());
}
const absl::InlinedVector<int, 4>& GraphTopologyView::GetFanin(
int node_idx) const {
DCHECK(is_initialized()) << "GraphTopologyView is not initialized";
const bool is_valid_node_idx = node_idx >= 0 && node_idx < num_nodes_;
DCHECK(is_valid_node_idx) << "node_idx is out of range";
return is_valid_node_idx ? fanins_[node_idx] : empty_fanin_;
}
const absl::InlinedVector<int, 2>& GraphTopologyView::GetFanout(
int node_idx) const {
DCHECK(is_initialized()) << "GraphTopologyView is not initialized";
const bool is_valid_node_idx = node_idx >= 0 && node_idx < num_nodes_;
DCHECK(is_valid_node_idx) << "node_idx is out of range";
return is_valid_node_idx ? fanouts_[node_idx] : empty_fanout_;
}
} // end namespace grappler
} // end namespace tensorflow
@@ -0,0 +1,116 @@
/* 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_GRAPPLER_GRAPH_TOPOLOGY_VIEW_H_
#define TENSORFLOW_CORE_GRAPPLER_GRAPH_TOPOLOGY_VIEW_H_
#include "absl/container/flat_hash_map.h"
#include "absl/container/inlined_vector.h"
#include "absl/strings/string_view.h"
#include "absl/types/optional.h"
#include "absl/types/span.h"
#include "tensorflow/core/graph/tensor_id.h"
#include "tensorflow/core/grappler/graph_view.h"
namespace tensorflow {
namespace grappler {
// GraphTopologyView is a helper class to simplify `node-to-node` connectivity
// traversals. Regular `GraphView` simplifies `tensor-to-tensor` traversals:
// connections between output tensors and inputs of a consumer nodes. For the
// topology view we are focused on nodes connected to nodes, and it's irrelevant
// if this connection is formed by one or multiple individual tensors.
//
// Example:
// a = Placeholder(..)
// b = Placeholder(..)
// c = AddN([a, a, b])
//
// GraphView edges: [a:0 -> c:0, a:0 -> c:1, b:0 -> c:2]
// GraphTopologyView edges: [a -> c, b -> c]
//
// GraphView is used for exploring single node fanins and fanouts, and
// GraphTopologyView is focused on efficient full graph traversals (computing
// graph node properties from transitive fanouts, etc...).
class GraphTopologyView {
public:
GraphTopologyView() = default;
explicit GraphTopologyView(bool skip_invalid_edges)
: skip_invalid_edges_(skip_invalid_edges) {}
// Initialize graph topology view from the graph. It's possible to pass
// additional edges that do not exist in a graph, but must be respected when
// computing graph topology. Example: Tensorflow runtime allows concurrent
// execution of dequeue/enqueue ops from the same queue resource, but we might
// want to enforce ordering between them for the purpose of graph analysis.
absl::Status InitializeFromGraph(
const GraphDef& graph, absl::Span<const GraphView::Edge> ephemeral_edges,
bool ignore_control_edges);
absl::Status InitializeFromGraph(
const GraphDef& graph, absl::Span<const GraphView::Edge> ephemeral_edges);
absl::Status InitializeFromGraph(const GraphDef& graph,
bool ignore_control_edges);
absl::Status InitializeFromGraph(const GraphDef& graph);
bool is_initialized() const { return graph_ != nullptr; }
int num_nodes() const { return num_nodes_; }
const GraphDef* graph() const { return graph_; }
// Returns true iff the node exists in the underlying graph.
bool HasNode(absl::string_view node_name) const;
// Finds a node by name or returns `nullptr` if it's not in the graph.
const NodeDef* GetNode(absl::string_view node_name) const;
// Returns a node corresponding to the given node index.
const NodeDef* GetNode(int node_idx) const;
// Returns a node index for the given node name, if the name exists in the
// underlying graph. Otherwise returns empty optional.
const absl::optional<int> GetNodeIndex(absl::string_view node_name) const;
// Returns a node index for the given node, if the node belongs to the
// underlying graph. Otherwise returns empty optional.
const absl::optional<int> GetNodeIndex(const NodeDef& node) const;
// Returns all the node indexes that are in the direct fanin of the given
// node. If the `node_idx` is outside of [0, num_nodes_) returns empty vector.
const absl::InlinedVector<int, 4>& GetFanin(int node_idx) const;
// Returns all the node indexes that are in the direct fanout of the given
// node. If the `node_idx` is outside of [0, num_nodes_) returns empty vector.
const absl::InlinedVector<int, 2>& GetFanout(int node_idx) const;
private:
// If true, all invalid edges and inputs (srd, dst or input node not found in
// a graph) will be skipped, otherwise initialization will fail with error.
bool skip_invalid_edges_ = false;
// WARN: `graph_` must outlive this object and graph nodes must not be
// destructed, because node names captured with absl::string_view.
const GraphDef* graph_ = nullptr; // do not own
int num_nodes_ = 0;
std::vector<absl::string_view> index_to_node_name_;
absl::flat_hash_map<absl::string_view, int> node_name_to_index_;
std::vector<absl::InlinedVector<int, 4>> fanins_; // node_idx->input nodes
std::vector<absl::InlinedVector<int, 2>> fanouts_; // node_idx->output nodes
// We need a valid reference to return from GetFanin/GetFanout if the
// `node_idx` argument is outside of the [0, num_nodes_) range.
absl::InlinedVector<int, 4> empty_fanin_;
absl::InlinedVector<int, 2> empty_fanout_;
};
} // end namespace grappler
} // end namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPPLER_GRAPH_TOPOLOGY_VIEW_H_
@@ -0,0 +1,153 @@
/* 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/grappler/graph_topology_view.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace grappler {
class GraphTopologyViewTest : public ::testing::Test {
protected:
using NodeConfig = std::pair<std::string, std::vector<std::string>>;
static GraphDef CreateGraph(const std::vector<NodeConfig>& nodes) {
GraphDef graph;
for (const NodeConfig& node : nodes) {
const auto& node_name = node.first;
const auto& node_inputs = node.second;
NodeDef node_def;
node_def.set_name(node_name);
for (const std::string& input : node_inputs) {
node_def.add_input(input);
}
*graph.add_node() = std::move(node_def);
}
return graph;
}
};
TEST_F(GraphTopologyViewTest, SimpleGraph) {
const GraphDef graph = CreateGraph({
{"a", {}}, // idx: 0
{"b", {}}, // idx: 1
{"c", {"a", "b"}}, // idx: 2
{"d", {"a", "c"}}, // idx: 3
});
GraphTopologyView graph_view;
TF_CHECK_OK(graph_view.InitializeFromGraph(graph));
EXPECT_TRUE(graph_view.is_initialized());
const NodeDef* a_by_name = graph_view.GetNode("a");
const NodeDef* a_by_idx = graph_view.GetNode(0);
ASSERT_TRUE(a_by_name);
ASSERT_TRUE(a_by_idx);
EXPECT_EQ(a_by_name, a_by_idx);
const NodeDef* b_by_name = graph_view.GetNode("b");
const NodeDef* b_by_idx = graph_view.GetNode(1);
ASSERT_TRUE(b_by_name);
ASSERT_TRUE(b_by_idx);
EXPECT_EQ(b_by_name, b_by_idx);
const absl::optional<int> b_idx = graph_view.GetNodeIndex(*b_by_name);
ASSERT_TRUE(b_idx.has_value());
EXPECT_EQ(b_idx.value(), 1);
const absl::optional<int> c_idx = graph_view.GetNodeIndex("c");
ASSERT_TRUE(c_idx.has_value());
EXPECT_EQ(c_idx.value(), 2);
using Fanin = absl::InlinedVector<int, 4>;
EXPECT_EQ(graph_view.GetFanin(0), Fanin());
EXPECT_EQ(graph_view.GetFanin(1), Fanin());
EXPECT_EQ(graph_view.GetFanin(2), Fanin({0, 1}));
EXPECT_EQ(graph_view.GetFanin(3), Fanin({0, 2}));
using Fanout = absl::InlinedVector<int, 2>;
EXPECT_EQ(graph_view.GetFanout(0), Fanout({2, 3}));
EXPECT_EQ(graph_view.GetFanout(1), Fanout({2}));
EXPECT_EQ(graph_view.GetFanout(2), Fanout({3}));
EXPECT_EQ(graph_view.GetFanout(3), Fanout());
}
TEST_F(GraphTopologyViewTest, GraphWithALoop) {
const GraphDef graph = CreateGraph({
{"a", {}}, // idx: 0
{"b", {}}, // idx: 1
{"c", {"a", "b", "d"}}, // idx: 2 <<<--- 'c' and 'd' have a loop
{"d", {"a", "c"}}, // idx: 3
});
GraphTopologyView graph_view;
TF_CHECK_OK(graph_view.InitializeFromGraph(graph));
EXPECT_TRUE(graph_view.is_initialized());
using Fanin = absl::InlinedVector<int, 4>;
EXPECT_EQ(graph_view.GetFanin(2), Fanin({0, 1, 3}));
EXPECT_EQ(graph_view.GetFanin(3), Fanin({0, 2}));
using Fanout = absl::InlinedVector<int, 2>;
EXPECT_EQ(graph_view.GetFanout(2), Fanout({3}));
EXPECT_EQ(graph_view.GetFanout(3), Fanout({2}));
}
TEST_F(GraphTopologyViewTest, GraphWithControls) {
const GraphDef graph = CreateGraph({
{"a", {}}, // idx: 0
{"b", {}}, // idx: 1
{"c", {"a", "b", "^d"}}, // idx: 2
{"d", {"a", "c"}}, // idx: 3
});
{
GraphTopologyView graph_view;
TF_CHECK_OK(graph_view.InitializeFromGraph(graph));
EXPECT_TRUE(graph_view.is_initialized());
using Fanin = absl::InlinedVector<int, 4>;
EXPECT_EQ(graph_view.GetFanin(2), Fanin({0, 1, 3}));
EXPECT_EQ(graph_view.GetFanin(3), Fanin({0, 2}));
using Fanout = absl::InlinedVector<int, 2>;
EXPECT_EQ(graph_view.GetFanout(2), Fanout({3}));
EXPECT_EQ(graph_view.GetFanout(3), Fanout({2}));
}
{
GraphTopologyView graph_view;
TF_CHECK_OK(
graph_view.InitializeFromGraph(graph, /*ignore_controls*/ true));
EXPECT_TRUE(graph_view.is_initialized());
using Fanin = absl::InlinedVector<int, 4>;
EXPECT_EQ(graph_view.GetFanin(2), Fanin({0, 1}));
EXPECT_EQ(graph_view.GetFanin(3), Fanin({0, 2}));
using Fanout = absl::InlinedVector<int, 2>;
EXPECT_EQ(graph_view.GetFanout(2), Fanout({3}));
EXPECT_EQ(graph_view.GetFanout(3), Fanout({}));
}
}
} // namespace grappler
} // namespace tensorflow
+60
View File
@@ -0,0 +1,60 @@
/* 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/grappler/graph_view.h"
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/grappler/utils.h"
namespace tensorflow {
namespace grappler {
int OpOutputPortIdToArgId(const NodeDef& node, const OpDef& op, int port_id) {
return OpPortIdToArgId(node, op.output_arg(), port_id);
}
int OpInputPortIdToArgId(const NodeDef& node, const OpDef& op, int port_id) {
return OpPortIdToArgId(node, op.input_arg(), port_id);
}
bool HasSingleFanoutNode(const GraphView& graph_view, const NodeDef* node,
int port) {
const auto output = GraphView::OutputPort(node, port);
return graph_view.GetFanout(output).size() <= 1;
}
bool HasFanouts(const GraphView& graph_view, const NodeDef* node, int port) {
const auto output = GraphView::OutputPort(node, port);
return !graph_view.GetFanout(output).empty();
}
bool HasControlFanin(const GraphView& graph_view, const NodeDef* node) {
const auto control_port = GraphView::InputPort(node, Graph::kControlSlot);
return !graph_view.GetFanin(control_port).empty();
}
bool HasControlFanout(const GraphView& graph_view, const NodeDef* node) {
const auto control_port = GraphView::OutputPort(node, Graph::kControlSlot);
return !graph_view.GetFanout(control_port).empty();
}
bool HasControlFaninOrFanout(const GraphView& graph_view, const NodeDef* node) {
return HasControlFanin(graph_view, node) ||
HasControlFanout(graph_view, node);
}
} // end namespace grappler
} // end namespace tensorflow
+428
View File
@@ -0,0 +1,428 @@
/* 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_GRAPPLER_GRAPH_VIEW_H_
#define TENSORFLOW_CORE_GRAPPLER_GRAPH_VIEW_H_
#include <unordered_map>
#include <unordered_set>
#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "absl/hash/hash.h"
#include "absl/status/status.h"
#include "absl/strings/string_view.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/op_def.pb.h"
#include "tensorflow/core/graph/tensor_id.h"
#include "tensorflow/core/grappler/utils.h"
#include "tensorflow/core/lib/gtl/map_util.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace grappler {
// Map a node/op's input/output port_id to arg_id.
//
// The port_id refers to the n-th tensor of the node, while the arg_id refers to
// the n-th arg of the op. These two can be different if an op's arg is a list
// of tensors.
//
// We return -1 for any invalid port_id (i.e., no corresponding arg_id).
int OpOutputPortIdToArgId(const NodeDef& node, const OpDef& op, int port_id);
int OpInputPortIdToArgId(const NodeDef& node, const OpDef& op, int port_id);
namespace internal {
// GraphViewInternal is a helper class to simplify graph traversal. It creates
// an immutable view of the nodes and edges represented by a GraphDef protocol
// buffer.
//
// There are two public classes implementing GraphViewInternal:
//
// - GraphView: constructed from the `const GraphDef` and doesn't allow
// to mutate underlying graph via input/output ports lookup functions (ports
// have const pointers to nodes).
//
// - MutableGraphView: constructed from the 'GraphDef` and allows to mutate
// the graph via input/output ports lookup functions (ports have non-const
// pointers to nodes), and also have couple additional functions to
// add/remove/replace nodes in the graph.
//
// --------------------------- !!! WARNING !!! ---------------------------------
// Removing nodes from the graph outside of MutableGraphView will
// lead to segfaults! Guaranteed by absl::string_view!
// -----------------------------------------------------------------------------
//
template <typename GraphDefT, typename NodeDefT>
class GraphViewInternal {
public:
struct Port {
Port() : node(nullptr), port_id(0) {}
Port(NodeDefT* n, int port) : node(n), port_id(port) {}
bool operator==(const Port& other) const {
return node == other.node && port_id == other.port_id;
}
template <typename H>
friend H AbslHashValue(H h, const Port& p) {
return H::combine(std::move(h), p.node, p.port_id);
}
NodeDefT* node;
int port_id;
};
struct InputPort : public Port {
using Port::Port;
};
struct OutputPort : public Port {
using Port::Port;
};
struct Edge {
Edge(OutputPort s, InputPort d) : src(s), dst(d) {}
bool operator==(const Edge& other) const {
return src == other.src && dst == other.dst;
}
template <typename H>
friend H AbslHashValue(H h, const Edge& e) {
return H::combine(std::move(h), e.src, e.dst);
}
OutputPort src;
InputPort dst;
};
GraphDefT* graph() const { return graph_; }
// Finds a node by name or return `nullptr` if it's not in the graph view.
NodeDefT* GetNode(absl::string_view node_name) const {
return gtl::FindWithDefault(nodes_, node_name, nullptr);
}
// Checks if a node by name is in the graph view.
bool HasNode(absl::string_view node_name) const {
return GetNode(node_name) != nullptr;
}
// Gets the specified input port. Note that the special '-1' port_id can be
// used to access the controlling nodes (i.e. the nodes connected to node_name
// through an incoming control dependency).
InputPort GetInputPort(absl::string_view node_name, int port_id) const {
return InputPort(GetNode(node_name), port_id);
}
// Gets the specified output port. Note that the special '-1' port_id can be
// used to access the controlled nodes (i.e. the nodes connected to node_name
// through an outgoing control dependency).
OutputPort GetOutputPort(absl::string_view node_name, int port_id) const {
return OutputPort(GetNode(node_name), port_id);
}
// Gets the input port(s) in the immediate fanout of an output port.
const absl::flat_hash_set<InputPort>& GetFanout(
const OutputPort& port) const {
return gtl::FindWithDefault(fanouts_, port, fanout_not_found_value_);
}
// Gets the output port(s) in the immediate fanin of an input port.
absl::flat_hash_set<OutputPort> GetFanin(const InputPort& port) const {
if (port.port_id >= 0) {
OutputPort regular_fanin = GetRegularFanin(port);
if (regular_fanin.node == nullptr) {
return {};
}
return {regular_fanin};
}
// Collect fanin for the control input.
absl::flat_hash_set<OutputPort> result;
const int first_control_port =
gtl::FindWithDefault(max_regular_input_port_, port.node, -1) + 1;
for (int i = first_control_port; i < port.node->input_size(); ++i) {
TensorId tensor_id = ParseTensorName(port.node->input(i));
auto it = nodes_.find(tensor_id.node());
if (it != nodes_.end()) result.emplace(it->second, tensor_id.index());
}
return result;
}
// Special case: regular (i.e. non-control) input ports can only have one
// fanin. If port.port_id is out of range or is a control dependency, then an
// empty OutputPort is returned.
const OutputPort GetRegularFanin(const InputPort& port) const {
if (port.port_id < 0 ||
port.port_id >
gtl::FindWithDefault(max_regular_input_port_, port.node, -1)) {
return OutputPort();
}
TensorId tensor_id = ParseTensorName(port.node->input(port.port_id));
return GetOutputPort(tensor_id.node(), tensor_id.index());
}
// Checks if a tensor id is a fanin of the node.
bool HasFanin(const NodeDefT& node, const TensorId& fanin) const {
int end = node.input_size();
if (end == 0 || fanin.index() < -1) {
return false;
}
const int num_regular_fanins =
gtl::FindWithDefault(max_regular_input_port_, &node, -1) + 1;
int start = 0;
if (fanin.index() > -1) {
end = num_regular_fanins;
} else {
start = num_regular_fanins;
}
for (int i = start; i < end; ++i) {
if (ParseTensorName(node.input(i)) == fanin) {
return true;
}
}
return false;
}
// Gets all the input ports in the immediate fanout of a node. Include the
// controlled nodes iff include_controlled_nodes is true.
absl::flat_hash_set<InputPort> GetFanouts(
const NodeDefT& node, bool include_controlled_nodes) const {
absl::flat_hash_set<InputPort> result;
OutputPort port;
port.node = const_cast<NodeDefT*>(&node);
const int first_port_id = include_controlled_nodes ? -1 : 0;
const int last_port_id =
gtl::FindWithDefault(max_regular_output_port_, &node, -1);
for (int i = first_port_id; i <= last_port_id; ++i) {
port.port_id = i;
auto it = fanouts_.find(port);
if (it != fanouts_.end()) {
result.insert(it->second.begin(), it->second.end());
}
}
return result;
}
// Gets all the output ports in the immediate fanin of a node. Include the
// controlling nodes iff include_controlling_nodes is true.
absl::flat_hash_set<OutputPort> GetFanins(
const NodeDefT& node, bool include_controlling_nodes) const {
absl::flat_hash_set<OutputPort> result;
const int max_input_port =
include_controlling_nodes
? node.input_size() - 1
: gtl::FindWithDefault(max_regular_input_port_, &node, -1);
for (int i = 0; i <= max_input_port; ++i) {
TensorId tensor_id = ParseTensorName(node.input(i));
auto it = nodes_.find(tensor_id.node());
if (it != nodes_.end()) result.emplace(it->second, tensor_id.index());
}
return result;
}
// Gets the number of ports in the immediate fanin of a node. Count the
// controlling nodes iff include_controlling_nodes is true.
int NumFanins(const NodeDefT& node, bool include_controlling_nodes) const {
if (include_controlling_nodes) {
return node.input_size();
}
return gtl::FindWithDefault(max_regular_input_port_, &node, -1) + 1;
}
// Gets the number of ports in the immediate fanout of a node. Count the
// controlled nodes iff include_controlled_nodes is true.
int NumFanouts(const NodeDefT& node, bool include_controlled_nodes) const {
int count = 0;
OutputPort port;
port.node = const_cast<NodeDefT*>(&node);
const int first_port_id = include_controlled_nodes ? -1 : 0;
const int last_port_id =
gtl::FindWithDefault(max_regular_output_port_, &node, -1);
for (int i = first_port_id; i <= last_port_id; ++i) {
port.port_id = i;
auto it = fanouts_.find(port);
if (it != fanouts_.end()) count += it->second.size();
}
return count;
}
// Gets all the edges in the immediate fanout of a node. Include the
// controlled edges iff include_controlled_edges is true.
absl::flat_hash_set<Edge> GetFanoutEdges(
const NodeDefT& node, bool include_controlled_edges) const {
absl::flat_hash_set<Edge> result;
OutputPort port;
port.node = const_cast<NodeDefT*>(&node);
const int first_port_id = include_controlled_edges ? -1 : 0;
const int last_port_id =
gtl::FindWithDefault(max_regular_output_port_, &node, -1);
for (int i = first_port_id; i <= last_port_id; ++i) {
port.port_id = i;
auto it = fanouts_.find(port);
if (it != fanouts_.end()) {
for (auto itr = it->second.begin(); itr != it->second.end(); ++itr) {
result.emplace(/*src=*/port, /*dst=*/*itr);
}
}
}
return result;
}
// Gets all the edges in the immediate fanin of a node. Include the
// controlling edges iff include_controlling_edges is true.
absl::flat_hash_set<Edge> GetFaninEdges(
const NodeDefT& node, bool include_controlling_edges) const {
absl::flat_hash_set<Edge> result;
const int max_input_port =
include_controlling_edges
? node.input_size() - 1
: gtl::FindWithDefault(max_regular_input_port_, &node, -1);
for (int i = 0; i <= max_input_port; ++i) {
TensorId tensor_id = ParseTensorName(node.input(i));
auto it = nodes_.find(tensor_id.node());
if (it != nodes_.end()) {
result.emplace(/*src=*/OutputPort(it->second, tensor_id.index()),
/*dst=*/InputPort(const_cast<NodeDefT*>(&node), i));
}
}
return result;
}
protected:
explicit GraphViewInternal(GraphDefT* graph) : graph_(graph) {}
absl::Status AddUniqueNode(NodeDefT* node) {
auto inserted = nodes_.emplace(node->name(), node);
return inserted.second
? absl::OkStatus()
: absl::InvalidArgumentError(absl::StrCat(
"Non unique node name detected: ", node->name()));
}
// TODO(ezhulenev): Remove this function.
void AddUniqueNodeOrDie(NodeDefT* node) {
absl::Status st = AddUniqueNode(node);
CHECK(st.ok()) << st.message();
}
// TODO(lyandy): Checks for self loops, Switch control dependencies, fanins
// exist, and all regular fanins come before controlling fanins.
void AddFanouts(NodeDefT* node) {
int max_input_port = -1;
for (int i = 0; i < node->input_size(); ++i) {
TensorId tensor_id = ParseTensorName(node->input(i));
OutputPort output(nodes_[tensor_id.node()], tensor_id.index());
if (output.port_id < 0) {
fanouts_[output].emplace(node, -1);
} else {
max_input_port = i;
int& max_regular_output_port = max_regular_output_port_[output.node];
max_regular_output_port =
std::max(max_regular_output_port, output.port_id);
fanouts_[output].emplace(node, i);
}
}
if (max_input_port > -1) {
max_regular_input_port_[node] = max_input_port;
}
}
// Access to the mutable internal state for MutableGraphView.
absl::flat_hash_map<absl::string_view, NodeDefT*>& nodes() { return nodes_; }
absl::flat_hash_map<OutputPort, absl::flat_hash_set<InputPort>>& fanouts() {
return fanouts_;
}
absl::flat_hash_map<const NodeDefT*, int>& max_regular_input_port() {
return max_regular_input_port_;
}
absl::flat_hash_map<const NodeDefT*, int>& max_regular_output_port() {
return max_regular_output_port_;
}
private:
GraphDefT* graph_; // must outlive the graph view
// A mapping from the node name to the node itself.
absl::flat_hash_map<absl::string_view, NodeDefT*> nodes_;
// A mapping from the output port to all inputs that read from it.
absl::flat_hash_map<OutputPort, absl::flat_hash_set<InputPort>> fanouts_;
// Keep a maximum index of input tensors of the node.
absl::flat_hash_map<const NodeDefT*, int> max_regular_input_port_;
// Keep a maximum index of tensor fetched from the node. It doesn't guarantee
// that all tensors in the [0, max_regular_output_port] range are actually
// fetched by other nodes.
absl::flat_hash_map<const NodeDefT*, int> max_regular_output_port_;
// If the node has no fanouts at given output port (output tensor consumers)
// we return a reference to this set from `GetFanout` (we can't construct new
// empty set every time, because we need a non-dangling reference).
absl::flat_hash_set<InputPort> fanout_not_found_value_;
};
} // namespace internal
// Immutable GraphView that keeps the constness of the GraphDef. If you need to
// mutate the graph or the nodes via the graph view lookup functions, see
// MutableGraphView.
class GraphView
: public internal::GraphViewInternal<const GraphDef, const NodeDef> {
public:
explicit GraphView(const GraphDef* graph) : GraphViewInternal(graph) {
for (const NodeDef& node : graph->node()) AddUniqueNodeOrDie(&node);
for (const NodeDef& node : graph->node()) AddFanouts(&node);
}
};
// Returns true if node has one (or zero) fanout nodes at given output port.
bool HasSingleFanoutNode(const GraphView& graph_view, const NodeDef* node,
int port = 0);
// Returns true if node has at least one fanout node at given output port.
bool HasFanouts(const GraphView& graph_view, const NodeDef* node, int port = 0);
// Returns true if the node has at least one input control dependency.
bool HasControlFanin(const GraphView& graph_view, const NodeDef* node);
// Returns true if the node has at least one output control dependency.
bool HasControlFanout(const GraphView& graph_view, const NodeDef* node);
// Returns true if the node has at least one input or output control dependency.
bool HasControlFaninOrFanout(const GraphView& graph_view, const NodeDef* node);
} // end namespace grappler
} // end namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPPLER_GRAPH_VIEW_H_
+507
View File
@@ -0,0 +1,507 @@
/* 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/grappler/graph_view.h"
#include "absl/container/flat_hash_set.h"
#include "absl/strings/str_cat.h"
#include "tensorflow/cc/ops/parsing_ops.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/core/graph/benchmark_testlib.h"
#include "tensorflow/core/grappler/grappler_item.h"
#include "tensorflow/core/grappler/inputs/trivial_test_graph_input_yielder.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/test_benchmark.h"
namespace tensorflow {
namespace grappler {
namespace {
class GraphViewTest : public ::testing::Test {};
TEST_F(GraphViewTest, OpPortIdToArgIdShapeN) {
tensorflow::Scope s = tensorflow::Scope::NewRootScope();
Output a = ops::Const(s.WithOpName("a"), 0.0f, {10, 10});
ops::ShapeN b(s.WithOpName("b"), {a, a, a});
GraphDef graph_def;
TF_CHECK_OK(s.ToGraphDef(&graph_def));
GraphView graph_view(&graph_def);
const NodeDef& a_node_def = *graph_view.GetNode("a");
const NodeDef& b_node_def = *graph_view.GetNode("b");
const OpDef* a_op_def = nullptr;
const OpDef* b_op_def = nullptr;
TF_EXPECT_OK(OpRegistry::Global()->LookUpOpDef(a_node_def.op(), &a_op_def));
TF_EXPECT_OK(OpRegistry::Global()->LookUpOpDef(b_node_def.op(), &b_op_def));
// Const has 0 inputs, 1 output.
EXPECT_EQ(OpInputPortIdToArgId(a_node_def, *a_op_def, 0), -1);
EXPECT_EQ(OpOutputPortIdToArgId(a_node_def, *a_op_def, 0), 0);
EXPECT_EQ(OpOutputPortIdToArgId(a_node_def, *a_op_def, 1), -1);
// ShapeN has N=3 inputs and outputs.
EXPECT_EQ(OpInputPortIdToArgId(b_node_def, *b_op_def, 0), 0);
EXPECT_EQ(OpInputPortIdToArgId(b_node_def, *b_op_def, 1), 0);
EXPECT_EQ(OpInputPortIdToArgId(b_node_def, *b_op_def, 2), 0);
EXPECT_EQ(OpInputPortIdToArgId(b_node_def, *b_op_def, 3), -1);
EXPECT_EQ(OpOutputPortIdToArgId(b_node_def, *b_op_def, 0), 0);
EXPECT_EQ(OpOutputPortIdToArgId(b_node_def, *b_op_def, 1), 0);
EXPECT_EQ(OpOutputPortIdToArgId(b_node_def, *b_op_def, 2), 0);
EXPECT_EQ(OpOutputPortIdToArgId(b_node_def, *b_op_def, 3), -1);
EXPECT_EQ(OpOutputPortIdToArgId(b_node_def, *b_op_def, 4), -1);
}
TEST_F(GraphViewTest, OpPortIdToArgIdSparseSplit) {
for (int num_splits : {1, 2}) {
tensorflow::Scope s = tensorflow::Scope::NewRootScope();
Output a = ops::Const<int64_t>(s.WithOpName("a"), 1, {10, 10});
ops::SparseSplit b(s.WithOpName("b"), a, a, a, a, num_splits);
GraphDef graph_def;
TF_CHECK_OK(s.ToGraphDef(&graph_def));
GraphView graph_view(&graph_def);
const NodeDef& b_node_def = *graph_view.GetNode("b");
const OpDef* b_op_def = nullptr;
TF_EXPECT_OK(OpRegistry::Global()->LookUpOpDef(b_node_def.op(), &b_op_def));
// We have 4 inputs.
EXPECT_EQ(OpInputPortIdToArgId(b_node_def, *b_op_def, 0), 0);
EXPECT_EQ(OpInputPortIdToArgId(b_node_def, *b_op_def, 1), 1);
EXPECT_EQ(OpInputPortIdToArgId(b_node_def, *b_op_def, 2), 2);
EXPECT_EQ(OpInputPortIdToArgId(b_node_def, *b_op_def, 3), 3);
EXPECT_EQ(OpInputPortIdToArgId(b_node_def, *b_op_def, 4), -1);
for (int port_id = 0; port_id <= num_splits * 3; ++port_id) {
int arg_id = -1;
if (port_id < num_splits * 3) {
arg_id = port_id / num_splits;
}
EXPECT_EQ(OpOutputPortIdToArgId(b_node_def, *b_op_def, port_id), arg_id);
}
}
}
TEST_F(GraphViewTest, ParseSingleExample) {
tensorflow::Scope s = tensorflow::Scope::NewRootScope();
Output a = ops::Const<tstring>(s.WithOpName("a"), "", {});
Output b = ops::Const<int64_t>(s.WithOpName("b"), 1, {1, 1});
ops::ParseSingleExample c(s.WithOpName("c"), a, {b, b}, 2, {"w", "x"},
{"y", "z"}, {DT_INT64, DT_INT64}, {{1}, {1}});
GraphDef graph_def;
TF_CHECK_OK(s.ToGraphDef(&graph_def));
GraphView graph_view(&graph_def);
const NodeDef& c_node_def = *graph_view.GetNode("c");
const OpDef* c_op_def = nullptr;
TF_EXPECT_OK(OpRegistry::Global()->LookUpOpDef(c_node_def.op(), &c_op_def));
EXPECT_EQ(OpOutputPortIdToArgId(c_node_def, *c_op_def, 0), 0);
EXPECT_EQ(OpOutputPortIdToArgId(c_node_def, *c_op_def, 1), 0);
EXPECT_EQ(OpOutputPortIdToArgId(c_node_def, *c_op_def, 2), 1);
EXPECT_EQ(OpOutputPortIdToArgId(c_node_def, *c_op_def, 3), 1);
EXPECT_EQ(OpOutputPortIdToArgId(c_node_def, *c_op_def, 4), 2);
EXPECT_EQ(OpOutputPortIdToArgId(c_node_def, *c_op_def, 5), 2);
EXPECT_EQ(OpOutputPortIdToArgId(c_node_def, *c_op_def, 6), 3);
EXPECT_EQ(OpOutputPortIdToArgId(c_node_def, *c_op_def, 7), 3);
EXPECT_EQ(OpOutputPortIdToArgId(c_node_def, *c_op_def, 8), -1);
}
TEST_F(GraphViewTest, BasicGraph) {
TrivialTestGraphInputYielder fake_input(4, 2, 2, false, {"/CPU:0", "/GPU:0"});
GrapplerItem item;
CHECK(fake_input.NextItem(&item));
GraphView graph(&item.graph);
GraphView::InputPort input = graph.GetInputPort("AddN", 0);
EXPECT_EQ(input.node->name(), "AddN");
EXPECT_EQ(input.port_id, 0);
GraphView::OutputPort fanin = graph.GetRegularFanin(input);
EXPECT_EQ(fanin.node->name(), "Sign");
EXPECT_EQ(fanin.port_id, 0);
input = graph.GetInputPort("AddN", 1);
EXPECT_EQ(input.node->name(), "AddN");
EXPECT_EQ(input.port_id, 1);
fanin = graph.GetRegularFanin(input);
EXPECT_EQ(fanin.node->name(), "Sign_1");
EXPECT_EQ(fanin.port_id, 0);
GraphView::OutputPort output = graph.GetOutputPort("AddN", 0);
EXPECT_EQ(output.node->name(), "AddN");
EXPECT_EQ(output.port_id, 0);
EXPECT_EQ(graph.GetFanout(output).size(), 2);
for (auto fanout : graph.GetFanout(output)) {
if (fanout.node->name() == "AddN_2" || fanout.node->name() == "AddN_3") {
EXPECT_EQ(fanout.port_id, 0);
} else {
// Invalid fanout
EXPECT_FALSE(true);
}
}
const NodeDef* add_node = graph.GetNode("AddN");
EXPECT_NE(add_node, nullptr);
absl::flat_hash_set<std::string> fanouts;
absl::flat_hash_set<std::string> expected_fanouts = {"AddN_2:0", "AddN_3:0"};
for (const auto& fo : graph.GetFanouts(*add_node, false)) {
fanouts.insert(absl::StrCat(fo.node->name(), ":", fo.port_id));
}
EXPECT_EQ(graph.NumFanouts(*add_node, false), 2);
EXPECT_EQ(fanouts, expected_fanouts);
absl::flat_hash_set<std::string> fanins;
absl::flat_hash_set<std::string> expected_fanins = {"Sign_1:0", "Sign:0"};
for (const auto& fi : graph.GetFanins(*add_node, false)) {
fanins.insert(absl::StrCat(fi.node->name(), ":", fi.port_id));
}
EXPECT_EQ(graph.NumFanins(*add_node, false), 2);
EXPECT_EQ(fanins, expected_fanins);
}
TEST_F(GraphViewTest, ControlDependencies) {
tensorflow::Scope s = tensorflow::Scope::NewRootScope();
Output a = ops::Const(s.WithOpName("a"), 0.0f, {10, 10});
Output b = ops::Square(s.WithOpName("b"), {a});
Output c = ops::Sqrt(s.WithOpName("c"), {b});
Output d = ops::AddN(s.WithOpName("d").WithControlDependencies(a), {b, c});
GrapplerItem item;
TF_CHECK_OK(s.ToGraphDef(&item.graph));
GraphView graph(&item.graph);
GraphView::OutputPort output = graph.GetOutputPort("a", -1);
EXPECT_EQ(output.node->name(), "a");
EXPECT_EQ(output.port_id, -1);
auto fanout = graph.GetFanout(output);
EXPECT_EQ(fanout.size(), 1);
EXPECT_EQ((*fanout.begin()).node->name(), "d");
EXPECT_EQ((*fanout.begin()).port_id, -1);
output = graph.GetOutputPort("a", 0);
EXPECT_EQ(output.node->name(), "a");
EXPECT_EQ(output.port_id, 0);
fanout = graph.GetFanout(output);
EXPECT_EQ(fanout.size(), 1);
EXPECT_EQ((*fanout.begin()).node->name(), "b");
EXPECT_EQ((*fanout.begin()).port_id, 0);
GraphView::InputPort input = graph.GetInputPort("d", -1);
EXPECT_EQ(input.node->name(), "d");
EXPECT_EQ(input.port_id, -1);
auto fanin = graph.GetFanin(input);
EXPECT_EQ(fanin.size(), 1);
EXPECT_EQ((*fanin.begin()).node->name(), "a");
EXPECT_EQ((*fanin.begin()).port_id, -1);
input = graph.GetInputPort("d", 0);
EXPECT_EQ(input.node->name(), "d");
EXPECT_EQ(input.port_id, 0);
fanin = graph.GetFanin(input);
EXPECT_EQ(fanin.size(), 1);
EXPECT_EQ((*fanin.begin()).node->name(), "b");
EXPECT_EQ((*fanin.begin()).port_id, 0);
input = graph.GetInputPort("d", 1);
EXPECT_EQ(input.node->name(), "d");
EXPECT_EQ(input.port_id, 1);
fanin = graph.GetFanin(input);
EXPECT_EQ(fanin.size(), 1);
EXPECT_EQ((*fanin.begin()).node->name(), "c");
EXPECT_EQ((*fanin.begin()).port_id, 0);
}
TEST_F(GraphViewTest, HasNode) {
tensorflow::Scope s = tensorflow::Scope::NewRootScope();
Output a = ops::Const(s.WithOpName("a"), 0.0f, {10, 10});
GrapplerItem item;
TF_CHECK_OK(s.ToGraphDef(&item.graph));
GraphView graph(&item.graph);
EXPECT_EQ(graph.HasNode("a"), true);
EXPECT_EQ(graph.HasNode("b"), false);
}
TEST_F(GraphViewTest, HasFanin) {
tensorflow::Scope s = tensorflow::Scope::NewRootScope();
Output a = ops::Const(s.WithOpName("a"), 0.0f, {10, 10});
Output b = ops::Square(s.WithOpName("b"), {a});
Output c = ops::Sqrt(s.WithOpName("c"), {b});
Output d = ops::AddN(s.WithOpName("d").WithControlDependencies(a), {b, c});
GrapplerItem item;
TF_CHECK_OK(s.ToGraphDef(&item.graph));
GraphView graph(&item.graph);
const NodeDef* d_node = graph.GetNode("d");
EXPECT_NE(d_node, nullptr);
EXPECT_EQ(graph.HasFanin(*d_node, {"a", Graph::kControlSlot}), true);
EXPECT_EQ(graph.HasFanin(*d_node, {"a", 0}), false);
EXPECT_EQ(graph.HasFanin(*d_node, {"b", 0}), true);
EXPECT_EQ(graph.HasFanin(*d_node, {"b", Graph::kControlSlot}), false);
EXPECT_EQ(graph.HasFanin(*d_node, {"c", 0}), true);
EXPECT_EQ(graph.HasFanin(*d_node, {"c", Graph::kControlSlot}), false);
}
TEST_F(GraphViewTest, GetRegularFaninPortOutOfBounds) {
tensorflow::Scope s = tensorflow::Scope::NewRootScope();
Output a = ops::Const(s.WithOpName("a"), 0.0f, {10, 10});
Output b = ops::Square(s.WithOpName("b"), {});
Output c = ops::Sqrt(s.WithOpName("c"), {b});
Output d = ops::AddN(s.WithOpName("d").WithControlDependencies(a), {b, c});
GrapplerItem item;
TF_CHECK_OK(s.ToGraphDef(&item.graph));
GraphView graph(&item.graph);
const NodeDef* b_node = graph.GetNode("b");
EXPECT_NE(b_node, nullptr);
const NodeDef* c_node = graph.GetNode("c");
EXPECT_NE(c_node, nullptr);
const NodeDef* d_node = graph.GetNode("d");
EXPECT_NE(d_node, nullptr);
auto d_output_0 = graph.GetRegularFanin({d_node, 0});
EXPECT_EQ(d_output_0, GraphView::OutputPort(b_node, 0));
auto d_output_1 = graph.GetRegularFanin({d_node, 1});
EXPECT_EQ(d_output_1, GraphView::OutputPort(c_node, 0));
auto d_output_2 = graph.GetRegularFanin({d_node, 2});
EXPECT_EQ(d_output_2, GraphView::OutputPort());
auto d_output_control = graph.GetRegularFanin({d_node, Graph::kControlSlot});
EXPECT_EQ(d_output_control, GraphView::OutputPort());
}
void BM_GraphViewConstruction(::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);
for (auto s : state) {
GraphView graph_view(&graph_def);
}
}
BENCHMARK(BM_GraphViewConstruction)
->ArgPair(10, 2)
->ArgPair(100, 2)
->ArgPair(1000, 2)
->ArgPair(10000, 2)
->ArgPair(25000, 2)
->ArgPair(50000, 2)
->ArgPair(100000, 2)
->ArgPair(10, 4)
->ArgPair(100, 4)
->ArgPair(1000, 4)
->ArgPair(10000, 4)
->ArgPair(25000, 4)
->ArgPair(50000, 4)
->ArgPair(100000, 4)
->ArgPair(10, 8)
->ArgPair(100, 8)
->ArgPair(1000, 8)
->ArgPair(10000, 8)
->ArgPair(25000, 8)
->ArgPair(50000, 8)
->ArgPair(100000, 8)
->ArgPair(10, 16)
->ArgPair(100, 16)
->ArgPair(1000, 16)
->ArgPair(10000, 16)
->ArgPair(25000, 16)
->ArgPair(50000, 16)
->ArgPair(100000, 16);
void BM_GraphViewGetNode(::testing::benchmark::State& state) {
const int num_nodes = state.range(0);
const GraphDef graph_def =
test::CreateGraphDef(num_nodes, /*num_edges_per_node=*/16);
GraphView graph_view(&graph_def);
for (auto s : state) {
graph_view.GetNode("out");
}
}
BENCHMARK(BM_GraphViewGetNode)
->Arg(10)
->Arg(100)
->Arg(1000)
->Arg(10000)
->Arg(25000)
->Arg(50000)
->Arg(100000);
#define RUN_FANIN_FANOUT_BENCHMARK(name) \
BENCHMARK(name) \
->ArgPair(10, 10) \
->ArgPair(10, 100) \
->ArgPair(10, 1000) \
->ArgPair(10, 10000) \
->ArgPair(10, 100000) \
->ArgPair(100, 10) \
->ArgPair(100, 100) \
->ArgPair(100, 1000) \
->ArgPair(100, 10000) \
->ArgPair(100, 100000) \
->ArgPair(1000, 10) \
->ArgPair(1000, 100) \
->ArgPair(1000, 1000) \
->ArgPair(1000, 10000) \
->ArgPair(1000, 100000) \
->ArgPair(10000, 10) \
->ArgPair(10000, 100) \
->ArgPair(10000, 1000) \
->ArgPair(10000, 10000) \
->ArgPair(10000, 100000) \
->ArgPair(100000, 10) \
->ArgPair(100000, 100) \
->ArgPair(100000, 1000) \
->ArgPair(100000, 10000) \
->ArgPair(100000, 100000);
void BM_GraphViewGetFanout(::testing::benchmark::State& state) {
const int num_fanins = state.range(0);
const int num_fanouts = state.range(1);
const GraphDef graph_def = test::CreateFaninFanoutNodeGraph(
num_fanins, num_fanouts, num_fanins, num_fanouts,
/*fanout_unique_index=*/true);
GraphView graph_view(&graph_def);
for (auto s : state) {
const NodeDef* node = graph_view.GetNode("node");
graph_view.GetFanout({node, 0});
}
}
RUN_FANIN_FANOUT_BENCHMARK(BM_GraphViewGetFanout);
void BM_GraphViewGetFanin(::testing::benchmark::State& state) {
const int num_fanins = state.range(0);
const int num_fanouts = state.range(1);
const GraphDef graph_def = test::CreateFaninFanoutNodeGraph(
num_fanins, num_fanouts, num_fanins, num_fanouts,
/*fanout_unique_index=*/true);
GraphView graph_view(&graph_def);
for (auto s : state) {
const NodeDef* node = graph_view.GetNode("node");
graph_view.GetFanin({node, 0});
}
}
RUN_FANIN_FANOUT_BENCHMARK(BM_GraphViewGetFanin);
void BM_GraphViewGetRegularFanin(::testing::benchmark::State& state) {
const int num_fanins = state.range(0);
const int num_fanouts = state.range(1);
const GraphDef graph_def = test::CreateFaninFanoutNodeGraph(
num_fanins, num_fanouts, num_fanins, num_fanouts,
/*fanout_unique_index=*/true);
GraphView graph_view(&graph_def);
for (auto s : state) {
const NodeDef* node = graph_view.GetNode("node");
graph_view.GetRegularFanin({node, 0});
}
}
RUN_FANIN_FANOUT_BENCHMARK(BM_GraphViewGetRegularFanin);
void BM_GraphViewGetFanouts(::testing::benchmark::State& state) {
const int num_fanins = state.range(0);
const int num_fanouts = state.range(1);
const GraphDef graph_def = test::CreateFaninFanoutNodeGraph(
num_fanins, num_fanouts, num_fanins, num_fanouts,
/*fanout_unique_index=*/true);
GraphView graph_view(&graph_def);
for (auto s : state) {
const NodeDef* node = graph_view.GetNode("node");
graph_view.GetFanouts(*node, /*include_controlled_nodes=*/false);
}
}
RUN_FANIN_FANOUT_BENCHMARK(BM_GraphViewGetFanouts);
void BM_GraphViewGetFanins(::testing::benchmark::State& state) {
const int num_fanins = state.range(0);
const int num_fanouts = state.range(1);
const GraphDef graph_def = test::CreateFaninFanoutNodeGraph(
num_fanins, num_fanouts, num_fanins, num_fanouts,
/*fanout_unique_index=*/true);
GraphView graph_view(&graph_def);
for (auto s : state) {
const NodeDef* node = graph_view.GetNode("node");
graph_view.GetFanins(*node, /*include_controlling_nodes=*/false);
}
}
RUN_FANIN_FANOUT_BENCHMARK(BM_GraphViewGetFanins);
void BM_GraphViewGetFanoutEdges(::testing::benchmark::State& state) {
const int num_fanins = state.range(0);
const int num_fanouts = state.range(1);
const GraphDef graph_def = test::CreateFaninFanoutNodeGraph(
num_fanins, num_fanouts, num_fanins, num_fanouts,
/*fanout_unique_index=*/true);
GraphView graph_view(&graph_def);
for (auto s : state) {
const NodeDef* node = graph_view.GetNode("node");
graph_view.GetFanoutEdges(*node, /*include_controlled_edges=*/false);
}
}
RUN_FANIN_FANOUT_BENCHMARK(BM_GraphViewGetFanoutEdges);
void BM_GraphViewGetFaninEdges(::testing::benchmark::State& state) {
const int num_fanins = state.range(0);
const int num_fanouts = state.range(1);
const GraphDef graph_def = test::CreateFaninFanoutNodeGraph(
num_fanins, num_fanouts, num_fanins, num_fanouts,
/*fanout_unique_index=*/true);
GraphView graph_view(&graph_def);
for (auto s : state) {
const NodeDef* node = graph_view.GetNode("node");
graph_view.GetFaninEdges(*node, /*include_controlling_edges=*/false);
}
}
RUN_FANIN_FANOUT_BENCHMARK(BM_GraphViewGetFaninEdges);
} // namespace
} // namespace grappler
} // namespace tensorflow
+230
View File
@@ -0,0 +1,230 @@
/* 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/grappler/grappler_item.h"
#include <unordered_map>
#include <unordered_set>
#include <vector>
#include "absl/container/flat_hash_set.h"
#include "absl/strings/str_join.h"
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/grappler/op_types.h"
#include "tensorflow/core/grappler/utils.h"
#include "tensorflow/core/grappler/utils/transitive_fanin.h"
#include "tensorflow/core/util/device_name_utils.h"
namespace tensorflow {
namespace grappler {
GrapplerItem::OptimizationOptions CreateOptOptionsForEager() {
GrapplerItem::OptimizationOptions optimization_options;
// Tensorflow 2.0 in eager mode with automatic control dependencies will
// prune all nodes that are not in the transitive fanin of the fetch nodes.
// However because the function will be executed via FunctionLibraryRuntime,
// and current function implementation does not prune stateful and dataset
// ops, we rely on Grappler to do the correct graph pruning.
optimization_options.allow_pruning_stateful_and_dataset_ops = true;
optimization_options.is_eager_mode = true;
// All the nested function calls will be executed and optimized via
// PartitionedCallOp, there is no need to optimize functions now.
optimization_options.optimize_function_library = false;
return optimization_options;
}
GrapplerItem GrapplerItem::WithGraph(GraphDef&& graph_def) const {
GrapplerItem item;
item.id = id;
item.feed = feed;
item.fetch = fetch;
item.init_ops = init_ops;
item.keep_ops = keep_ops;
item.expected_init_time = expected_init_time;
item.save_op = save_op;
item.restore_op = restore_op;
item.save_restore_loc_tensor = save_restore_loc_tensor;
item.queue_runners = queue_runners;
item.devices_ = devices_;
item.optimization_options_ = optimization_options_;
item.graph.Swap(&graph_def);
return item;
}
std::vector<const NodeDef*> GrapplerItem::MainOpsFanin() const {
std::vector<const NodeDef*> fanin_nodes;
TF_CHECK_OK(ComputeTransitiveFanin(graph, fetch, &fanin_nodes));
return fanin_nodes;
}
std::vector<const NodeDef*> GrapplerItem::EnqueueOpsFanin() const {
std::vector<std::string> enqueue_ops;
for (const auto& queue_runner : queue_runners) {
for (const std::string& enqueue_op : queue_runner.enqueue_op_name()) {
enqueue_ops.push_back(enqueue_op);
}
}
std::vector<const NodeDef*> fanin_nodes;
TF_CHECK_OK(ComputeTransitiveFanin(graph, fetch, &fanin_nodes));
return fanin_nodes;
}
std::vector<const NodeDef*> GrapplerItem::InitOpsFanin() const {
std::vector<const NodeDef*> fanin_nodes;
TF_CHECK_OK(ComputeTransitiveFanin(graph, init_ops, &fanin_nodes));
return fanin_nodes;
}
std::vector<const NodeDef*> GrapplerItem::MainVariables() const {
std::vector<const NodeDef*> fanin;
TF_CHECK_OK(ComputeTransitiveFanin(graph, init_ops, &fanin));
std::vector<const NodeDef*> vars;
for (const NodeDef* node : fanin) {
if (IsVariable(*node)) {
vars.push_back(node);
}
}
return vars;
}
std::unordered_set<std::string> GrapplerItem::NodesToPreserve() const {
std::unordered_set<std::string> result;
for (const std::string& f : fetch) {
VLOG(1) << "Add fetch " << f;
result.insert(NodeName(f));
}
for (const auto& f : feed) {
VLOG(1) << "Add feed " << f.first;
result.insert(NodeName(f.first));
}
for (const auto& node : init_ops) {
result.insert(NodeName(node));
}
for (const auto& node : keep_ops) {
result.insert(NodeName(node));
}
if (!save_op.empty()) {
result.insert(NodeName(save_op));
}
if (!restore_op.empty()) {
result.insert(NodeName(restore_op));
}
if (!save_restore_loc_tensor.empty()) {
result.insert(NodeName(save_restore_loc_tensor));
}
for (const auto& queue_runner : queue_runners) {
for (const std::string& enqueue_op : queue_runner.enqueue_op_name()) {
result.insert(NodeName(enqueue_op));
}
if (!queue_runner.close_op_name().empty()) {
result.insert(NodeName(queue_runner.close_op_name()));
}
if (!queue_runner.cancel_op_name().empty()) {
result.insert(NodeName(queue_runner.cancel_op_name()));
}
}
absl::optional<FunctionLibraryDefinition> fn_library;
if (!optimization_options_.allow_pruning_stateful_and_dataset_ops) {
fn_library.emplace(OpRegistry::Global(), graph.library());
}
for (const NodeDef& node : graph.node()) {
const auto attrs = AttrSlice(&node.attr());
// Tensorflow functions do not prune stateful or dataset-output ops from
// the function body (see PruneFunctionBody in common_runtime/function.cc).
if (!optimization_options_.allow_pruning_stateful_and_dataset_ops &&
(IsStateful(node, &*fn_library) || IsDataset(node))) {
result.insert(node.name());
}
// Do not remove ops with attribute _grappler_do_not_remove. This is useful
// for debugging.
bool do_not_remove;
if (TryGetNodeAttr(attrs, "_grappler_do_not_remove", &do_not_remove) &&
do_not_remove) {
result.insert(node.name());
}
}
return result;
}
const std::unordered_set<std::string>& GrapplerItem::devices() const {
return devices_;
}
absl::Status GrapplerItem::AddDevice(const std::string& device) {
DeviceNameUtils::ParsedName name;
if (!DeviceNameUtils::ParseFullName(device, &name)) {
return absl::InvalidArgumentError(
absl::StrCat("Invalid device name: device=", device));
} else if (!name.has_job || !name.has_replica || !name.has_task ||
!name.has_type || !name.has_id) {
return absl::InvalidArgumentError(
absl::StrCat("Not a fully defined device name: device=", device));
}
devices_.insert(DeviceNameUtils::ParsedNameToString(name));
return absl::OkStatus();
}
absl::Status GrapplerItem::AddDevices(const GrapplerItem& other) {
std::vector<absl::string_view> invalid_devices;
for (const std::string& device : other.devices()) {
absl::Status added = AddDevice(device);
if (!added.ok()) invalid_devices.emplace_back(device);
}
return invalid_devices.empty()
? absl::OkStatus()
: absl::InvalidArgumentError(
absl::StrCat("Skipped invalid devices: [",
absl::StrJoin(invalid_devices, ", "), "]"));
}
absl::Status GrapplerItem::InferDevicesFromGraph() {
absl::flat_hash_set<absl::string_view> invalid_devices;
for (const NodeDef& node : graph.node()) {
absl::Status added = AddDevice(node.device());
if (!added.ok()) invalid_devices.insert(node.device());
}
VLOG(2) << "Inferred device set: [" << absl::StrJoin(devices_, ", ") << "]";
return invalid_devices.empty()
? absl::OkStatus()
: absl::InvalidArgumentError(
absl::StrCat("Skipped invalid devices: [",
absl::StrJoin(invalid_devices, ", "), "]"));
}
void GrapplerItem::ClearDevices() { devices_.clear(); }
const GrapplerItem::OptimizationOptions& GrapplerItem::optimization_options()
const {
return optimization_options_;
}
GrapplerItem::OptimizationOptions& GrapplerItem::optimization_options() {
return optimization_options_;
}
} // end namespace grappler
} // end namespace tensorflow
+145
View File
@@ -0,0 +1,145 @@
/* 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_GRAPPLER_GRAPPLER_ITEM_H_
#define TENSORFLOW_CORE_GRAPPLER_GRAPPLER_ITEM_H_
#include <memory>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/variable.pb.h"
#include "tensorflow/core/protobuf/queue_runner.pb.h"
#include "tsl/platform/cpu_info.h"
namespace tensorflow {
namespace grappler {
// A TensorFlow model to optimize.
// Models are represented by the combination of a graph, one of more fetch
// nodes, and potentially a set of nodes to feed.
struct GrapplerItem {
GrapplerItem() = default;
GrapplerItem(const GrapplerItem& other) = default;
GrapplerItem(GrapplerItem&& other) = default;
GrapplerItem& operator=(const GrapplerItem& other) = default;
GrapplerItem& operator=(GrapplerItem&& other) = default;
virtual ~GrapplerItem() = default;
// Create a copy of this GrapplerItem with graph swapped with the argument.
GrapplerItem WithGraph(GraphDef&& graph) const;
std::string id; // A unique id for this item
// Inputs
GraphDef graph;
std::vector<std::pair<std::string, Tensor>> feed;
std::vector<std::string> fetch;
// Initialization op(s).
std::vector<std::string> init_ops;
// Expected initialization time in seconds, or 0 if unknown
int64_t expected_init_time = 0;
// Save/restore ops (if any)
std::string save_op;
std::string restore_op;
std::string save_restore_loc_tensor;
// Queue runner(s) required to run the queue(s) of this model.
std::vector<QueueRunnerDef> queue_runners;
// List of op names to keep in the graph. This includes nodes that are
// referenced in various collections, and therefore must be preserved to
// ensure that the optimized metagraph can still be loaded.
std::vector<std::string> keep_ops;
// Return the set of node evaluated during a regular train/inference step.
std::vector<const NodeDef*> MainOpsFanin() const;
// Return the set of node run to populate the queues (if any).
std::vector<const NodeDef*> EnqueueOpsFanin() const;
// Return the set nodes used by TensorFlow to initialize the graph.
std::vector<const NodeDef*> InitOpsFanin() const;
// Return the set of variables accessed during a regular train/inference step.
std::vector<const NodeDef*> MainVariables() const;
// Return a set of node names that must be preserved. This includes feed and
// fetch nodes, keep_ops, init_ops.
std::unordered_set<std::string> NodesToPreserve() const;
struct OptimizationOptions {
// Is it allowed to add nodes to the graph that do not have registered
// gradient function.
bool allow_non_differentiable_rewrites = true;
// Tensorflow function execution semantics is slightly different from the
// main Tensorflow graph, and we need to make sure that we do not change it
// by running Grappler optimizer passes. One main difference is that
// functions do not prune ops with side-effects and dataset-output ops (see
// PruneFunctionBody in common_runtime/function.cc).
bool allow_pruning_stateful_and_dataset_ops = true;
// If true Grappler will optimize the main graph, and also all functions in
// the graph function library (function can't be polymorphic, it can't have
// undefined type parameters in the function signature, or placeholder
// attributes in the function body).
bool optimize_function_library = true;
// Mark the grapper optimization run in eager mode or not.
bool is_eager_mode = false;
// Number of intra threads used to run operation.
int intra_op_parallelism_threads = tsl::port::MaxParallelism();
};
const std::unordered_set<std::string>& devices() const;
// Adds a device to a set of available devices, only if it's a valid fully
// defined device name. Returns `OkStatus()` if successfully added a device,
// and an error otherwise.
absl::Status AddDevice(const std::string& device);
// Adds all valid devices from the other Grappler item to the device set.
absl::Status AddDevices(const GrapplerItem& other);
// Adds all valid devices from the nodes of the graph to the device set.
// Returns `OkStatus()` if all device annotations found in a graph are valid
// fully defined device names, and an error otherwise.
absl::Status InferDevicesFromGraph();
// Clears a set of available devices.
void ClearDevices();
const OptimizationOptions& optimization_options() const;
OptimizationOptions& optimization_options();
private:
// TODO(ezhulenev) Make GrapplerItem a class and hide all public data members.
// TODO(ezhulenev): Migrate all unordered collections to absl.
// A set of fully defined device names that can be used to place the nodes of
// the `graph`.
// Example of a fully defined name: "/job:work/replica:1/task:1/device:CPU:0"
std::unordered_set<std::string> devices_;
OptimizationOptions optimization_options_;
};
GrapplerItem::OptimizationOptions CreateOptOptionsForEager();
} // end namespace grappler
} // end namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPPLER_GRAPPLER_ITEM_H_
@@ -0,0 +1,695 @@
/* 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/grappler/grappler_item_builder.h"
#include <type_traits>
#include <unordered_map>
#include <unordered_set>
#include <vector>
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "tensorflow/core/common_runtime/device.h"
#include "tensorflow/core/common_runtime/device_factory.h"
#include "tensorflow/core/common_runtime/device_mgr.h"
#include "tensorflow/core/common_runtime/function.h"
#include "tensorflow/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/common_runtime/graph_optimizer.h"
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/framework/graph_def_util.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/tensor.pb.h"
#include "tensorflow/core/framework/tensor_shape.pb.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/framework/variable.pb.h"
#include "tensorflow/core/framework/versions.pb.h"
#include "tensorflow/core/grappler/inputs/utils.h"
#include "tensorflow/core/grappler/op_types.h"
#include "tensorflow/core/grappler/optimizers/model_pruner.h"
#include "tensorflow/core/grappler/utils.h"
#include "tensorflow/core/lib/gtl/map_util.h"
#include "tensorflow/core/lib/io/path.h"
#include "tensorflow/core/platform/protobuf_internal.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
#include "tensorflow/core/protobuf/saver.pb.h"
#include "tensorflow/core/public/session_options.h"
namespace tensorflow {
namespace grappler {
namespace {
void InitializeTensor(DataType type, Tensor* tensor) {
const int period = 7;
if (type == DT_FLOAT) {
auto flat = tensor->flat<float>();
// Populate numbers 0, 0.1, 0.2, ..., 0.5, 0.6, 0, 0.1, 0.2, ...
for (int i = 0; i < flat.size(); i++) {
flat(i) = static_cast<float>(i % period) / 10.0f;
}
} else if (type == DT_INT64) {
auto flat = tensor->flat<int64_t>();
// Populate numbers 0, 1, 2, ..., 5, 6, 0, 1, 2, ...
for (int i = 0; i < flat.size(); i++) {
flat(i) = i % period;
}
} else if (type != DT_STRING && type != DT_RESOURCE && type != DT_VARIANT) {
// DT_STRING, DT_RESOURCE and DT_VARIANT are not simple types according to
// is_simple_type<> in tensorflow/core/framework/type_traits.h, and
// Allocator will run non-trivial constructor/destructor for a Tensor with
// one of these types, so we should not memset its buffer.
memset(const_cast<char*>(tensor->tensor_data().data()), 0,
tensor->tensor_data().size());
}
}
// Applies the same graph pruning logic to the graph as Session.Run in TF.
// If the returned status is not OK, item state may be inconsistent.
absl::Status PruneGraph(GrapplerItem* item) {
ModelPruner pruner;
GraphDef pruned_graph;
Cluster* cluster = nullptr; // ModelPruner doesn't check cluster.
TF_RETURN_IF_ERROR(pruner.Optimize(cluster, *item, &pruned_graph));
item->graph = std::move(pruned_graph);
return absl::OkStatus();
}
// Replace any unknown dimensions in a shape with
// cfg.placeholder_unknown_output_shape_dim if it is no less than 0.
absl::Status ReplaceUnknownShapeDim(const ItemConfig& cfg,
const TensorShapeProto& shape_pb_in,
TensorShapeProto* shape_pb_out,
TensorShape* shape_out) {
std::vector<int32_t> dims;
for (const auto& dim_proto : shape_pb_in.dim()) {
if (cfg.placeholder_unknown_output_shape_dim >= 0 &&
dim_proto.size() == -1) {
dims.push_back(cfg.placeholder_unknown_output_shape_dim);
shape_pb_out->add_dim()->set_size(
cfg.placeholder_unknown_output_shape_dim);
} else {
dims.push_back(std::max<int32_t>(1, dim_proto.size()));
shape_pb_out->add_dim()->set_size(dim_proto.size());
}
}
return TensorShapeUtils::MakeShape(dims.data(), dims.size(), shape_out);
}
// Replace unknown dimensions in Placeholder shape if
// cfg.placeholder_unknown_output_shape_dim is set or
// the Placeholder node has _output_shapes.
// Otherwise keep it intact to keep compatible with shape annotation
// (b/134092018).
absl::Status UpdatePlaceholderShape(
const ItemConfig& cfg,
const std::unordered_set<std::string>& signature_feed_nodes,
GrapplerItem* new_item, NodeDef* node) {
if (node->attr().count("dtype") == 0) {
return absl::InternalError(absl::StrCat("Unknown type for placeholder ",
node->name(),
", skipping this input"));
}
DataType type = node->attr().at("dtype").type();
// TODO(andiryxu): Consider cfg.placeholder_unknown_output_shape_dim >= 0 and
// _output_shapes is present case.
if (node->attr().count("shape") == 0) {
return absl::InternalError(absl::StrCat("Unknown shape for placeholder ",
node->name(),
", skipping this input"));
}
// Replace all unknown dimensions in the placeholder's tensorshape proto
// with cfg.placeholder_unknown_output_shape_dim and create a tensorshape
// from it. We do this because in newer protos, the input placeholder
// shape is not empty if the shape is partially defined.
TensorShape shape;
TensorShapeProto shape_proto;
absl::Status make_shape_status = ReplaceUnknownShapeDim(
cfg, node->attr().at("shape").shape(), &shape_proto, &shape);
if (!make_shape_status.ok()) {
return absl::InternalError(
absl::StrCat("Invalid shape for placeholder ", node->name(), ": ",
make_shape_status.ToString(), ", skipping this input"));
}
// Some placeholder nodes have a mismatch between the node
// attribute "shape" and a different node attribute "_output_shapes".
// Specifically, a shape with shape.dims() == 0 could indicate either
// a scalar or an unknown shape. In those cases, we check _output_shapes
// for additional information.
// This case is observed in the bnmt graphs. Have not observed any
// cases where there was more than 1 _output_shapes, so limit it
// to cases where there is only 1 _output_shapes.
// We only do this if cfg.placeholder_unknown_output_shape_dim has
// been set to avoid crashing non-BNMT graphs.
// TODO(andiryxu): Investigate if this is a bug in BNMT graph.
if ((cfg.placeholder_unknown_output_shape_dim >= 0) && (shape.dims() == 0) &&
(node->attr().count("_output_shapes") == 1)) {
const auto& output_shapes =
node->attr().at("_output_shapes").list().shape(0);
if (output_shapes.dim_size() != 0) {
shape.Clear();
shape_proto.clear_dim();
for (const auto& dim : output_shapes.dim()) {
auto size = dim.size();
if (size == -1) size = cfg.placeholder_unknown_output_shape_dim;
TF_RETURN_IF_ERROR(shape.AddDimWithStatus(size));
shape_proto.add_dim()->set_size(size);
}
}
}
Tensor fake_input(type, shape);
InitializeTensor(type, &fake_input);
if (cfg.feed_nodes.empty()) {
// No specific feed nodes were given. Assume all placeholders are fed.
if (signature_feed_nodes.count(node->name()) == 0) {
new_item->feed.emplace_back(node->name(), fake_input);
}
} else if (cfg.feed_nodes.count(node->name()) > 0) {
// If specific feed nodes were given, only update their tensors.
auto it = find_if(new_item->feed.begin(), new_item->feed.end(),
[&node](std::pair<std::string, Tensor>& f) {
return f.first == node->name();
});
DCHECK(it != new_item->feed.end());
it->second = fake_input;
}
// Set the shape of the node in the graph. This is needed for statically
// inferring shapes and is a no-op when dynamically inferring shapes as
// the Placeholder shape will match the shape passed from new_item->feed.
// Only replace node shape with known shape. For unknown shape keep it intact
// (b/134092018).
if (!shape_proto.dim().empty())
*(node->mutable_attr()->at("shape").mutable_shape()) = shape_proto;
return absl::OkStatus();
}
} // namespace
absl::Status RuntimeGraphOptimizer(const GraphDef& graph_def_arg,
GraphDef* output_graph_def,
const ItemConfig& cfg) {
// This is a temporary change that optimizes the graph in context of a single
// gpu machine. Down the line, we may want to make grappler_item_builder aware
// of the cluster type (E.g: single cpu, multiple gpu, etc) being simulated
// in order to get the correct session options and environment, and performing
// the correct optimizations.
// Return input as is if no graph-modifying config is set.
if (!cfg.apply_optimizations && !cfg.inline_functions &&
!cfg.erase_noinline_attributes) {
if (output_graph_def != &graph_def_arg) {
*output_graph_def = graph_def_arg;
}
return absl::OkStatus();
}
// Create a session option for a single GPU device.
SessionOptions options;
// Make a local copy of graph def, because we need to change some things.
GraphDef graph_def(graph_def_arg);
if (cfg.erase_noinline_attributes) {
// TF optimizer doesn't inline functions with "_noinline" attribute,
// so let's go over the function library and erase it.
for (auto& func : *graph_def.mutable_library()->mutable_function()) {
func.mutable_attr()->erase("_noinline");
}
}
// Instantiate all variables for function library runtime creation.
std::vector<std::unique_ptr<Device>> devices;
// Only CPU device is used so instead of calling DeviceFactory::AddDevices()
// with dummy session config, which will conflict with user defined options
// and create unwanted devices, call cpu_factory->CreateDevices() to get CPU
// only devices.
DeviceFactory* cpu_factory = DeviceFactory::GetFactory("CPU");
TF_RETURN_IF_ERROR(cpu_factory->CreateDevices(
options, "/job:localhost/replica:0/task:0", &devices));
Device* cpu_device = devices[0].get();
auto dvc_mgr = std::make_unique<StaticDeviceMgr>(std::move(devices));
FunctionLibraryDefinition function_library(OpRegistry::Global(),
graph_def.library());
Env* env = Env::Default();
// Optimizer options: L1 and inlining. L1 is default.
OptimizerOptions* optimizer_opts =
options.config.mutable_graph_options()->mutable_optimizer_options();
if (cfg.apply_optimizations) {
optimizer_opts->set_opt_level(::tensorflow::OptimizerOptions::L1);
} else {
optimizer_opts->set_opt_level(::tensorflow::OptimizerOptions::L0);
}
optimizer_opts->set_do_function_inlining(cfg.inline_functions);
// Create the function library runtime.
std::unique_ptr<ProcessFunctionLibraryRuntime> pflr(
new ProcessFunctionLibraryRuntime(dvc_mgr.get(), env, &options.config,
graph_def.versions().producer(),
&function_library, *optimizer_opts));
FunctionLibraryRuntime* flr = pflr->GetFLR(cpu_device->name());
// Create the GraphOptimizer to optimize the graph def.
GraphConstructorOptions graph_ctor_opts;
graph_ctor_opts.allow_internal_ops = true;
graph_ctor_opts.expect_device_spec = false;
std::unique_ptr<Graph> graphptr(new Graph(function_library));
TF_RETURN_IF_ERROR(ConvertGraphDefToGraph(
graph_ctor_opts, std::move(graph_def), graphptr.get()));
// Optimize the graph.
::tensorflow::GraphOptimizer optimizer(*optimizer_opts);
optimizer.Optimize(flr, env, cpu_device, &graphptr,
tensorflow::GraphOptimizer::Options());
graphptr->ToGraphDef(output_graph_def);
// The default values of attributes might have been stripped by the optimizer.
// Add them back.
return AddDefaultAttrsToGraphDef(output_graph_def, *graphptr->op_registry(),
0, true);
}
std::unique_ptr<GrapplerItem> GrapplerItemFromMetaGraphDef(
const std::string& id, const MetaGraphDef& meta_graph,
const ItemConfig& cfg) {
if (id.empty()) {
LOG(ERROR) << "id must be non-empty.";
return nullptr;
}
std::unique_ptr<GrapplerItem> new_item(new GrapplerItem());
new_item->id = id;
new_item->graph = meta_graph.graph_def();
// Fill in feed nodes from config, if any provided.
for (const auto& feed_node : cfg.feed_nodes) {
const std::string feed_name = NodeName(feed_node);
new_item->feed.emplace_back(feed_name, Tensor());
}
for (const auto& fetch_node : cfg.fetch_nodes) {
new_item->fetch.emplace_back(NodeName(fetch_node));
}
// Attempt to detect the fetch node(s) if they were not set explicitly.
if (new_item->fetch.empty() &&
meta_graph.collection_def().count("train_op") > 0) {
const CollectionDef& nodes = meta_graph.collection_def().at("train_op");
if (nodes.has_node_list()) {
for (const auto& node : nodes.node_list().value()) {
new_item->fetch.push_back(NodeName(node));
}
}
}
// Detect feed and fetch nodes from signature defs. Signatures may share same
// inputs or outputs.
std::unordered_set<std::string> signature_feed_nodes;
std::unordered_set<std::string> signature_fetch_nodes;
for (const auto& name_and_signature : meta_graph.signature_def()) {
for (const auto& name_and_input : name_and_signature.second.inputs()) {
const TensorInfo& input = name_and_input.second;
if (input.has_coo_sparse()) {
// Define the shapes following the comment of CooSparse.
// TODO(yuefengz): we probably want to use different dim values for the
// three tensors of a SparseTensor.
int64_t dim = std::max(1, cfg.placeholder_unknown_output_shape_dim);
TensorShape shape_1d({dim});
TensorShape shape_2d({dim, dim});
if (gtl::InsertIfNotPresent(
&signature_feed_nodes,
NodeName(input.coo_sparse().values_tensor_name()))) {
Tensor value_tensor(input.dtype(), shape_1d);
InitializeTensor(input.dtype(), &value_tensor);
new_item->feed.emplace_back(
NodeName(input.coo_sparse().values_tensor_name()), value_tensor);
}
if (gtl::InsertIfNotPresent(
&signature_feed_nodes,
NodeName(input.coo_sparse().indices_tensor_name()))) {
Tensor indices_tensor(DT_INT64, shape_2d);
InitializeTensor(input.dtype(), &indices_tensor);
new_item->feed.emplace_back(
NodeName(input.coo_sparse().indices_tensor_name()),
indices_tensor);
}
if (gtl::InsertIfNotPresent(
&signature_feed_nodes,
NodeName(input.coo_sparse().dense_shape_tensor_name()))) {
Tensor dense_shape_tensor(DT_INT64, shape_1d);
InitializeTensor(input.dtype(), &dense_shape_tensor);
new_item->feed.emplace_back(
NodeName(input.coo_sparse().dense_shape_tensor_name()),
dense_shape_tensor);
}
} else {
if (gtl::InsertIfNotPresent(&signature_feed_nodes,
NodeName(input.name()))) {
TensorShape shape;
TensorShapeProto shape_proto;
absl::Status s = ReplaceUnknownShapeDim(cfg, input.tensor_shape(),
&shape_proto, &shape);
if (!s.ok()) {
LOG(ERROR) << "Invalid shape for signature input " << input.name()
<< ": " << s << ", skipping this input";
return nullptr;
}
Tensor fake_input(input.dtype(), shape);
InitializeTensor(input.dtype(), &fake_input);
new_item->feed.emplace_back(NodeName(input.name()), fake_input);
}
}
}
for (const auto& name_and_output : name_and_signature.second.outputs()) {
const TensorInfo& output = name_and_output.second;
if (output.has_coo_sparse()) {
if (gtl::InsertIfNotPresent(
&signature_fetch_nodes,
NodeName(output.coo_sparse().values_tensor_name()))) {
new_item->fetch.push_back(
NodeName(output.coo_sparse().values_tensor_name()));
}
if (gtl::InsertIfNotPresent(
&signature_fetch_nodes,
NodeName(output.coo_sparse().indices_tensor_name()))) {
new_item->fetch.push_back(
NodeName(output.coo_sparse().indices_tensor_name()));
}
if (gtl::InsertIfNotPresent(
&signature_fetch_nodes,
NodeName(output.coo_sparse().dense_shape_tensor_name()))) {
new_item->fetch.push_back(
NodeName(output.coo_sparse().dense_shape_tensor_name()));
}
} else {
if (gtl::InsertIfNotPresent(&signature_fetch_nodes,
NodeName(output.name()))) {
new_item->fetch.push_back(NodeName(output.name()));
}
}
}
}
for (const auto& feed : new_item->feed) {
if (feed.first.empty()) {
LOG(ERROR) << "Invalid feed node name skipping this input";
return nullptr;
} else {
VLOG(1) << "Will use feed node " << feed.first;
}
}
for (const auto& fetch : new_item->fetch) {
if (fetch.empty()) {
LOG(ERROR) << "Invalid fetch node name skipping this input";
return nullptr;
} else {
VLOG(1) << "Will use fetch node " << fetch;
}
}
if (new_item->fetch.empty()) {
LOG(ERROR) << "Failed to detect the fetch node(s), skipping this input";
return nullptr;
}
// TODO(yuefengz): consider handling saved_model_main_op and legacy_init_op.
// The reason why they are difficult to handle is because they may not intend
// to initialize all variables that are required to run fetch nodes. We may
// have to run restore op first.
// Try to find initializers from variables and tables as init ops.
for (const std::string& var_collection :
{"variables", "local_variables", "model_variables",
"trainable_variables"}) {
if (meta_graph.collection_def().count(var_collection) == 0) {
continue;
}
const CollectionDef& vars = meta_graph.collection_def().at(var_collection);
for (const auto& raw_var : vars.bytes_list().value()) {
VariableDef var;
var.ParseFromString(raw_var);
if (!var.initializer_name().empty()) {
new_item->init_ops.push_back(NodeName(var.initializer_name()));
}
}
}
if (meta_graph.collection_def().count("table_initializer") > 0) {
const CollectionDef& inits =
meta_graph.collection_def().at("table_initializer");
if (inits.has_node_list()) {
for (const auto& node : inits.node_list().value()) {
new_item->init_ops.push_back(NodeName(node));
// Tables are initialized from files, which can take a long time. Add
// 30 minutes to the initialization time for each table to avoid
// timing out.
// TODO(bsteiner): adjust the timeout based on the file size.
new_item->expected_init_time += 30 * 60;
}
}
}
// We keep the mapping from asset node to asset files. This should have been
// used as feed but since asset node is usually a constant node, we will fill
// the values of these constant nodes with their actual asset file paths.
std::unordered_map<std::string, std::string> asset_node_to_value;
// Assets file may have changed their directory, we assemble their new paths
// if assets_directory_override is set. We also make sure we still can
// access these asset files.
if (!cfg.assets_directory_override.empty()) {
if (meta_graph.collection_def().count("saved_model_assets") > 0) {
const CollectionDef& collection =
meta_graph.collection_def().at("saved_model_assets");
const auto& any_assets = collection.any_list().value();
if (!any_assets.empty()) {
if (std::is_base_of<protobuf::Message, AssetFileDef>()) {
for (const auto& any_asset : any_assets) {
AssetFileDef asset_file_def;
if (!ParseAny(any_asset, &asset_file_def, "tensorflow.AssetFileDef")
.ok()) {
LOG(ERROR) << "Failed to parse AssetFile.";
continue;
}
std::string asset_filepath = io::JoinPath(
cfg.assets_directory_override, asset_file_def.filename());
if (!FilesExist({asset_filepath}, nullptr)) {
LOG(ERROR) << "Can't access one or more of the asset files "
<< asset_filepath << ", skipping this input";
return nullptr;
}
asset_node_to_value[NodeName(asset_file_def.tensor_info().name())] =
asset_filepath;
}
} else {
LOG(ERROR) << "Can't parse AssetFileDef when using lite protos.";
return nullptr;
}
}
}
} else if (meta_graph.collection_def().count("asset_filepaths") > 0) {
const CollectionDef& file_paths =
meta_graph.collection_def().at("asset_filepaths");
std::vector<std::string> paths;
for (const auto& raw_path : file_paths.bytes_list().value()) {
paths.push_back(raw_path);
}
if (!FilesExist(paths, nullptr)) {
LOG(ERROR) << "Can't access one or more of the asset files, skipping "
"this input";
return nullptr;
}
}
if (meta_graph.collection_def().count("queue_runners") > 0) {
const CollectionDef& vars = meta_graph.collection_def().at("queue_runners");
for (const auto& raw : vars.bytes_list().value()) {
QueueRunnerDef queue_runner;
if (!queue_runner.ParseFromString(raw)) {
LOG(ERROR) << "Could not parse queue_runners, skipping this input";
return nullptr;
}
if (queue_runner.cancel_op_name().empty()) {
LOG(ERROR) << "Queue without a cancel op, skipping this input";
return nullptr;
}
new_item->queue_runners.push_back(queue_runner);
}
}
// Add each node referenced in a collection to the list of nodes to keep.
for (const auto& col : meta_graph.collection_def()) {
const CollectionDef& collection = col.second;
for (const std::string& node : collection.node_list().value()) {
new_item->keep_ops.push_back(NodeName(node));
}
}
for (auto& node : *new_item->graph.mutable_node()) {
if (IsPlaceholder(node) && node.op() != "PlaceholderWithDefault") {
absl::Status s = UpdatePlaceholderShape(cfg, signature_feed_nodes,
new_item.get(), &node);
if (!s.ok()) return nullptr;
} else if (IsConstant(node)) {
auto it = asset_node_to_value.find(node.name());
if (it != asset_node_to_value.end()) {
auto iter = node.mutable_attr()->find("value");
if (iter == node.attr().end()) {
LOG(ERROR) << "Value attribute expected in const op for asset files";
return nullptr;
}
if (!iter->second.has_tensor() ||
iter->second.tensor().string_val_size() != 1) {
LOG(INFO) << "Unexpected AttrValue proto: "
<< iter->second.DebugString();
return nullptr;
}
LOG(INFO) << "Using asset file " << it->second << " for node "
<< node.name();
*(iter->second.mutable_tensor()->mutable_string_val(0)) = it->second;
}
}
// Erase the recorded result of any previous shape inference to start again
// from scratch.
node.mutable_attr()->erase("_output_shapes");
// Delete user specified placement if requested.
if (cfg.ignore_user_placement) {
node.clear_device();
}
// Delete colocation constraints if requested.
if (cfg.ignore_colocation) {
auto attr = node.mutable_attr();
auto it = attr->find("_class");
if (it != attr->end()) {
attr->erase(it);
}
}
}
if (meta_graph.collection_def().count("savers") > 0) {
const CollectionDef& savers = meta_graph.collection_def().at("savers");
for (const auto& raw : savers.bytes_list().value()) {
SaverDef saver;
// Skip bad savers since we don't need saves/restores to be able to run a
// graph.
if (!saver.ParseFromString(raw)) {
continue;
}
if (saver.filename_tensor_name().empty()) {
continue;
}
new_item->save_op = saver.save_tensor_name();
new_item->restore_op = saver.restore_op_name();
new_item->save_restore_loc_tensor = saver.filename_tensor_name();
// Only use the first saver since it's not clear what to do if there's
// more than one.
break;
}
} else {
const SaverDef& saver = meta_graph.saver_def();
new_item->save_op = saver.save_tensor_name();
new_item->restore_op = saver.restore_op_name();
new_item->save_restore_loc_tensor = saver.filename_tensor_name();
}
// Instantiate all the missing attributes with their default values.
absl::Status attr_status = AddDefaultAttrsToGraphDef(
&new_item->graph,
FunctionLibraryDefinition(OpRegistry::Global(),
new_item->graph.library()),
0, true);
if (!attr_status.ok()) {
LOG(ERROR) << "Failed to instantiate default attribute values: "
<< attr_status.message();
return nullptr;
}
// Optimize the graph (function inlining, l1 optimizations, etc).
VLOG(1) << "Number of nodes in graph before RuntimeGraphOptimizer: "
<< new_item->graph.node_size();
absl::Status optimize_status =
RuntimeGraphOptimizer(new_item->graph, &new_item->graph, cfg);
if (!optimize_status.ok()) {
LOG(ERROR) << "Graph preprocessing failed: " << optimize_status;
return nullptr;
}
VLOG(1) << "Number of nodes in graph after RuntimeGraphOptimizer: "
<< new_item->graph.node_size();
if (cfg.prune_graph) {
VLOG(1) << "Pruning graph...";
auto status = PruneGraph(new_item.get());
if (!status.ok()) {
LOG(ERROR) << "Pruning failed: " << status.message();
return nullptr;
}
VLOG(1) << "Number of nodes in graph after pruning: "
<< new_item->graph.node_size();
}
// Validate feed, fetch and init nodes
std::unordered_set<std::string> nodes;
for (const auto& node : new_item->graph.node()) {
nodes.insert(node.name());
}
for (const auto& feed : new_item->feed) {
if (nodes.find(feed.first) == nodes.end()) {
LOG(ERROR) << "Feed node " << feed.first << " doesn't exist in graph";
return nullptr;
}
}
for (const auto& fetch : new_item->fetch) {
if (nodes.find(fetch) == nodes.end()) {
LOG(ERROR) << "Fetch node " << fetch << " doesn't exist in graph";
return nullptr;
}
}
for (const auto& init : new_item->init_ops) {
if (nodes.find(init) == nodes.end()) {
LOG(ERROR) << "Init node " << init << " doesn't exist in graph";
return nullptr;
}
}
return new_item;
}
std::unique_ptr<GrapplerItem> GrapplerItemFromMetaGraphDefFile(
const std::string& id, const std::string& meta_graph_file,
const ItemConfig& cfg) {
MetaGraphDef meta_graph;
if (!ReadMetaGraphDefFromFile(meta_graph_file, &meta_graph).ok()) {
LOG(ERROR) << "Failed to read " << meta_graph_file;
return nullptr;
}
return GrapplerItemFromMetaGraphDef(id, meta_graph, cfg);
}
} // end namespace grappler
} // end namespace tensorflow
@@ -0,0 +1,87 @@
/* 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_GRAPPLER_GRAPPLER_ITEM_BUILDER_H_
#define TENSORFLOW_CORE_GRAPPLER_GRAPPLER_ITEM_BUILDER_H_
#include <memory>
#include <set>
#include <string>
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/grappler/grappler_item.h"
namespace tensorflow {
class MetaGraphDef;
namespace grappler {
struct ItemConfig {
ItemConfig() {}
// If true, ignore all user specified node placement.
bool ignore_user_placement = true;
// If true, ignore all user specified colocation attributes.
bool ignore_colocation = true;
// Dimension to use if a placeholder node has an _output_shapes attribute with
// a dimension of -1.
int placeholder_unknown_output_shape_dim = -1;
// If true, erases all "_noinline" attributes from user-defined functions.
// Has no effect if "inline_functions" is disabled.
bool erase_noinline_attributes = false;
// If non-empty, override the directory of asset paths.
std::string assets_directory_override;
// If true, runs ModelPruner on the graph.
bool prune_graph = false;
// Override feed nodes list.
std::set<std::string> feed_nodes;
// Override fetch nodes list.
std::set<std::string> fetch_nodes;
// Configs for graph optimizations from common_runtime. This is NOT Grappler
// function optimizer. When Grappler is invoked at runtime, it is typically
// running after common_runtime pass.
//
// If true, does L1 optimizations.
bool apply_optimizations = false;
// If true, does function inlining.
bool inline_functions = false;
};
// Method for optimizing the graph def (including function inlining and other
// optimizations). This is optimizations from common_runtime, NOT Grappler
// function optimizer.
absl::Status RuntimeGraphOptimizer(const GraphDef& graph_def_arg,
GraphDef* output_graph_def,
const ItemConfig& cfg);
// Factory method for creating a GrapplerItem from a MetaGraphDef.
// Returns nullptr if the given meta_graph cannot be converted.
std::unique_ptr<GrapplerItem> GrapplerItemFromMetaGraphDef(
const std::string& id, const MetaGraphDef& meta_graph,
const ItemConfig& cfg);
// Factory method for creating a GrapplerItem from a file
// containing a MetaGraphDef in either binary or text format.
// Returns nullptr if the given meta_graph cannot be converted.
std::unique_ptr<GrapplerItem> GrapplerItemFromMetaGraphDefFile(
const std::string& id, const std::string& meta_graph_file,
const ItemConfig& cfg);
} // end namespace grappler
} // end namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPPLER_GRAPPLER_ITEM_BUILDER_H_
@@ -0,0 +1,484 @@
/* 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/grappler/grappler_item_builder.h"
#include "google/protobuf/any.pb.h"
#include "tensorflow/cc/framework/gradients.h"
#include "tensorflow/cc/gradients/grad_testutil.h"
#include "tensorflow/cc/ops/functional_ops.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/core/framework/function_testlib.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/grappler/inputs/trivial_test_graph_input_yielder.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/lib/io/path.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
namespace tensorflow {
namespace grappler {
namespace {
class GrapplerItemBuilderTest : public ::testing::Test {};
TEST_F(GrapplerItemBuilderTest, AssetFilepathOverrideTest) {
MetaGraphDef meta_graph;
tensorflow::Scope s = tensorflow::Scope::NewRootScope();
Output var =
ops::Variable(s.WithOpName("var"), TensorShape(), DataType::DT_FLOAT);
Output filename_node =
ops::Const(s.WithOpName("filename"), std::string("model"), TensorShape());
Output tensor_name =
ops::Const(s.WithOpName("tensorname"), std::string("var"), TensorShape());
Output restore = ops::Restore(s.WithOpName("restore"), filename_node,
tensor_name, DataType::DT_FLOAT);
Output assign = ops::Assign(s.WithOpName("assign"), var, restore);
TF_CHECK_OK(s.ToGraphDef(meta_graph.mutable_graph_def()));
std::string temp_dir = testing::TmpDir();
Env *env = Env::Default();
std::string filename =
io::JoinPath(temp_dir, "grappler_item_builder_test_filename");
env->DeleteFile(filename).IgnoreError();
std::unique_ptr<WritableFile> file_to_write;
TF_CHECK_OK(env->NewWritableFile(filename, &file_to_write));
TF_CHECK_OK(file_to_write->Close());
TF_CHECK_OK(env->FileExists(filename));
LOG(INFO) << filename;
AssetFileDef asset_file_def;
*asset_file_def.mutable_tensor_info()->mutable_name() = "filename";
*asset_file_def.mutable_filename() = "grappler_item_builder_test_filename";
(*meta_graph.mutable_collection_def())["saved_model_assets"]
.mutable_any_list()
->add_value()
->PackFrom(asset_file_def);
*((*meta_graph.mutable_collection_def())["train_op"]
.mutable_node_list()
->add_value()) = "assign";
ItemConfig cfg;
cfg.assets_directory_override = temp_dir;
std::unique_ptr<GrapplerItem> item =
GrapplerItemFromMetaGraphDef("0", meta_graph, cfg);
ASSERT_TRUE(item != nullptr);
for (const NodeDef &node : item->graph.node()) {
if (node.name() == "filename") {
const auto iter = node.attr().find("value");
ASSERT_TRUE(iter != node.attr().end());
ASSERT_TRUE(iter->second.has_tensor());
ASSERT_EQ(1, iter->second.tensor().string_val_size());
std::string tensor_string_val = iter->second.tensor().string_val(0);
EXPECT_EQ(tensor_string_val, filename);
}
}
}
TEST_F(GrapplerItemBuilderTest, AssetFilepathOverrideTest_FileNotAccessible) {
MetaGraphDef meta_graph;
tensorflow::Scope s = tensorflow::Scope::NewRootScope();
Output var =
ops::Variable(s.WithOpName("var"), TensorShape(), DataType::DT_FLOAT);
Output filename_node1 = ops::Const(s.WithOpName("filename1"),
std::string("model1"), TensorShape());
Output filename_node2 = ops::Const(s.WithOpName("filename2"),
std::string("model2"), TensorShape());
Output tensor_name =
ops::Const(s.WithOpName("tensorname"), std::string("var"), TensorShape());
Output restore1 = ops::Restore(s.WithOpName("restore1"), filename_node1,
tensor_name, DataType::DT_FLOAT);
Output restore2 = ops::Restore(s.WithOpName("restore2"), filename_node1,
tensor_name, DataType::DT_FLOAT);
Output assign1 = ops::Assign(s.WithOpName("assign1"), var, restore1);
Output assign2 = ops::Assign(s.WithOpName("assign2"), var, restore2);
TF_CHECK_OK(s.ToGraphDef(meta_graph.mutable_graph_def()));
std::string temp_dir = testing::TmpDir();
// Create the first AssetFileDef that has a valid file.
Env *env = Env::Default();
std::string filename1 =
io::JoinPath(temp_dir, "grappler_item_builder_test_filename1");
env->DeleteFile(filename1).IgnoreError();
std::unique_ptr<WritableFile> file_to_write;
TF_CHECK_OK(env->NewWritableFile(filename1, &file_to_write));
TF_CHECK_OK(file_to_write->Close());
TF_CHECK_OK(env->FileExists(filename1));
AssetFileDef asset_file_def1;
*asset_file_def1.mutable_tensor_info()->mutable_name() = "filename1";
*asset_file_def1.mutable_filename() = "grappler_item_builder_test_filename1";
// Create the second AssetFileDef that has not a valid file.
std::string filename2 =
io::JoinPath(temp_dir, "grappler_item_builder_test_filename1");
env->DeleteFile(filename2).IgnoreError();
EXPECT_FALSE(env->FileExists(filename2).ok());
AssetFileDef asset_file_def2;
*asset_file_def2.mutable_tensor_info()->mutable_name() = "filename2";
*asset_file_def2.mutable_filename() = "grappler_item_builder_test_filename2";
(*meta_graph.mutable_collection_def())["saved_model_assets"]
.mutable_any_list()
->add_value()
->PackFrom(asset_file_def1);
(*meta_graph.mutable_collection_def())["saved_model_assets"]
.mutable_any_list()
->add_value()
->PackFrom(asset_file_def2);
*((*meta_graph.mutable_collection_def())["train_op"]
.mutable_node_list()
->add_value()) = "assign1";
*((*meta_graph.mutable_collection_def())["train_op"]
.mutable_node_list()
->add_value()) = "assign2";
ItemConfig cfg;
cfg.assets_directory_override = temp_dir;
std::unique_ptr<GrapplerItem> item =
GrapplerItemFromMetaGraphDef("0", meta_graph, cfg);
ASSERT_TRUE(item == nullptr);
}
TEST_F(GrapplerItemBuilderTest, GraphWithFunctions) {
MetaGraphDef meta_graph;
// y = XTimesTwo(x)
constexpr char device[] = "/cpu:0";
*meta_graph.mutable_graph_def() = test::function::GDef(
{test::function::NDef("x", "Const", {}, {{"dtype", DT_FLOAT}}, device),
test::function::NDef("y", "XTimesTwo", {"x"}, {{"T", DT_FLOAT}},
device)},
// FunctionLib
{
test::function::XTimesTwo(),
});
CollectionDef train_op;
train_op.mutable_node_list()->add_value("y");
(*meta_graph.mutable_collection_def())["train_op"] = train_op;
ItemConfig cfg;
std::unique_ptr<GrapplerItem> item =
GrapplerItemFromMetaGraphDef("0", meta_graph, cfg);
ASSERT_TRUE(item != nullptr);
}
TEST_F(GrapplerItemBuilderTest, GraphWithCustomOps) {
MetaGraphDef meta_graph;
// y = XTimesTwo(x)
constexpr char device[] = "/cpu:0";
*meta_graph.mutable_graph_def() = test::function::GDef(
{test::function::NDef("x", "Const", {}, {{"dtype", DT_FLOAT}}, device),
test::function::NDef("y", "CustomOp", {"x"}, {{"T", DT_FLOAT}}, device)},
{});
CollectionDef train_op;
train_op.mutable_node_list()->add_value("y");
(*meta_graph.mutable_collection_def())["train_op"] = train_op;
ItemConfig cfg;
std::unique_ptr<GrapplerItem> item =
GrapplerItemFromMetaGraphDef("0", meta_graph, cfg);
ASSERT_TRUE(item != nullptr);
}
TEST_F(GrapplerItemBuilderTest, FromGraphWithSignatureDef) {
tensorflow::Scope s = tensorflow::Scope::NewRootScope();
auto x = ops::Const(s.WithOpName("x"), 0);
auto y = ops::Const(s.WithOpName("y"), 1);
auto z = ops::Add(s.WithOpName("z"), x, y);
MetaGraphDef meta_graph;
TF_CHECK_OK(s.ToGraphDef(meta_graph.mutable_graph_def()));
TensorInfo input, output;
input.set_name("x");
input.set_dtype(DT_FLOAT);
output.set_name("z");
SignatureDef serving_signature;
(*serving_signature.mutable_inputs())["input"] = input;
(*serving_signature.mutable_outputs())["output"] = output;
(*meta_graph.mutable_signature_def())["serving"] = serving_signature;
// It should be able to dedup the input and output with same names.
TensorInfo input2, output2;
input.set_name("x");
input.set_dtype(DT_FLOAT);
output.set_name("z");
SignatureDef serving_signature2;
(*serving_signature.mutable_inputs())["input2"] = input2;
(*serving_signature.mutable_outputs())["output2"] = output2;
(*meta_graph.mutable_signature_def())["serving2"] = serving_signature2;
std::unique_ptr<GrapplerItem> item =
GrapplerItemFromMetaGraphDef("0", meta_graph, ItemConfig());
ASSERT_TRUE(item != nullptr);
EXPECT_EQ(item->feed.size(), 1);
EXPECT_EQ(item->fetch.size(), 1);
EXPECT_EQ(item->feed[0].first, "x");
EXPECT_EQ(item->fetch[0], "z");
}
TEST_F(GrapplerItemBuilderTest, FromGraphWithIncompleteSignatureDef) {
tensorflow::Scope s = tensorflow::Scope::NewRootScope();
auto x = ops::Const(s.WithOpName("x"), 0);
auto y = ops::Const(s.WithOpName("y"), 1);
MetaGraphDef meta_graph;
TF_CHECK_OK(s.ToGraphDef(meta_graph.mutable_graph_def()));
CollectionDef train_op;
train_op.mutable_node_list()->add_value("y");
(*meta_graph.mutable_collection_def())["train_op"] = train_op;
TensorInfo input, output;
input.set_name("x");
input.set_dtype(DT_FLOAT);
// Its coo_sparse proto is incomplete.
output.mutable_coo_sparse()->set_values_tensor_name("z");
SignatureDef serving_signature;
(*serving_signature.mutable_inputs())["input"] = input;
(*serving_signature.mutable_outputs())["output"] = output;
(*meta_graph.mutable_signature_def())["serving"] = serving_signature;
std::unique_ptr<GrapplerItem> item =
GrapplerItemFromMetaGraphDef("0", meta_graph, ItemConfig());
ASSERT_TRUE(item == nullptr);
}
TEST_F(GrapplerItemBuilderTest, FromGraphWithUnknownDimInSignatureInput) {
tensorflow::Scope s = tensorflow::Scope::NewRootScope();
auto shape_1d = PartialTensorShape({-1});
auto x = ops::Placeholder(s.WithOpName("x"), DT_FLOAT,
ops::Placeholder::Shape(shape_1d));
auto y = ops::Const(s.WithOpName("y"), static_cast<float>(1.0));
auto z = ops::Add(s.WithOpName("z"), x, y);
MetaGraphDef meta_graph;
TF_CHECK_OK(s.ToGraphDef(meta_graph.mutable_graph_def()));
TensorInfo input, output;
input.set_name("x");
input.set_dtype(DT_FLOAT);
shape_1d.AsProto(input.mutable_tensor_shape());
output.set_name("z");
SignatureDef serving_signature;
(*serving_signature.mutable_inputs())["input"] = input;
(*serving_signature.mutable_outputs())["output"] = output;
(*meta_graph.mutable_signature_def())["serving"] = serving_signature;
ItemConfig cfg;
cfg.placeholder_unknown_output_shape_dim = 64;
std::unique_ptr<GrapplerItem> item1 =
GrapplerItemFromMetaGraphDef("0", meta_graph, cfg);
ASSERT_TRUE(item1 != nullptr);
ASSERT_EQ(item1->feed.size(), 1);
EXPECT_EQ(item1->feed[0].second.NumElements(), 64);
std::unique_ptr<GrapplerItem> item2 =
GrapplerItemFromMetaGraphDef("0", meta_graph, ItemConfig());
ASSERT_TRUE(item2 != nullptr);
ASSERT_EQ(item2->feed.size(), 1);
EXPECT_EQ(item2->feed[0].second.NumElements(), 1);
}
TEST_F(GrapplerItemBuilderTest, ExplicitFeedAndFetch) {
tensorflow::Scope s = tensorflow::Scope::NewRootScope();
auto x = ops::Const(s.WithOpName("x"), 0);
auto y = ops::Const(s.WithOpName("y"), 1);
auto z = ops::Add(s.WithOpName("z"), x, y);
MetaGraphDef meta_graph;
TF_CHECK_OK(s.ToGraphDef(meta_graph.mutable_graph_def()));
ItemConfig config;
config.feed_nodes.insert("x");
config.fetch_nodes.insert("z");
std::unique_ptr<GrapplerItem> item =
GrapplerItemFromMetaGraphDef("0", meta_graph, config);
ASSERT_TRUE(item != nullptr);
EXPECT_EQ(item->feed.size(), 1);
EXPECT_EQ(item->fetch.size(), 1);
EXPECT_EQ(item->feed[0].first, "x");
EXPECT_EQ(item->fetch[0], "z");
}
TEST_F(GrapplerItemBuilderTest, UnknownRankPlaceholderTest) {
MetaGraphDef meta_graph;
const char* text_proto = R"EOF(
graph_def {
node {
name: "x"
op: "Placeholder"
attr { key: "dtype" value { type: DT_FLOAT } }
attr { key: "shape" value { shape { unknown_rank: true } } }
}
versions {
producer: 51
}
}
collection_def {
key: "train_op"
value {
node_list {
value: "x:0"
}
}
}
)EOF";
CHECK(protobuf::TextFormat::ParseFromString(text_proto, &meta_graph));
ItemConfig cfg;
std::unique_ptr<GrapplerItem> item =
GrapplerItemFromMetaGraphDef("0", meta_graph, cfg);
ASSERT_TRUE(item != nullptr);
const NodeDef& node = item->graph.node(0);
const auto iter = node.attr().find("shape");
ASSERT_TRUE(iter != node.attr().end());
ASSERT_TRUE(iter->second.has_shape());
const auto& shape = iter->second.shape();
// Do not update unknown shape.
EXPECT_TRUE(shape.unknown_rank());
}
TEST_F(GrapplerItemBuilderTest, ConfigPlaceholderTest) {
MetaGraphDef meta_graph;
const char* text_proto = R"EOF(
graph_def {
node {
name: "x"
op: "Placeholder"
attr { key: "dtype" value { type: DT_FLOAT } }
attr { key: "shape" value {
shape {
dim {
size: -1
}
dim {
size: -1
}
}
} }
}
versions {
producer: 51
}
}
collection_def {
key: "train_op"
value {
node_list {
value: "x:0"
}
}
}
)EOF";
CHECK(protobuf::TextFormat::ParseFromString(text_proto, &meta_graph));
ItemConfig cfg;
cfg.placeholder_unknown_output_shape_dim = 64;
std::unique_ptr<GrapplerItem> item =
GrapplerItemFromMetaGraphDef("0", meta_graph, cfg);
ASSERT_TRUE(item != nullptr);
const NodeDef& node = item->graph.node(0);
const auto iter = node.attr().find("shape");
ASSERT_TRUE(iter != node.attr().end());
ASSERT_TRUE(iter->second.has_shape());
const auto& shape = iter->second.shape();
EXPECT_EQ(shape.dim_size(), 2);
// Shape updated with placeholder_unknown_output_shape_dim.
EXPECT_EQ(shape.dim(0).size(), 64);
EXPECT_EQ(shape.dim(1).size(), 64);
}
TEST_F(GrapplerItemBuilderTest, OutputShapePlaceholderTest) {
MetaGraphDef meta_graph;
const char* text_proto = R"EOF(
graph_def {
node {
name: "x"
op: "Placeholder"
attr { key: "dtype" value { type: DT_FLOAT } }
attr { key: "shape" value { shape { unknown_rank: true } } }
attr { key: "_output_shapes" value { list {
shape {
dim {
size: -1
}
dim {
size: 32
}
}
} } }
}
versions {
producer: 51
}
}
collection_def {
key: "train_op"
value {
node_list {
value: "x:0"
}
}
}
)EOF";
CHECK(protobuf::TextFormat::ParseFromString(text_proto, &meta_graph));
ItemConfig cfg;
cfg.placeholder_unknown_output_shape_dim = 64;
std::unique_ptr<GrapplerItem> item =
GrapplerItemFromMetaGraphDef("0", meta_graph, cfg);
ASSERT_TRUE(item != nullptr);
const NodeDef& node = item->graph.node(0);
const auto iter = node.attr().find("shape");
ASSERT_TRUE(iter != node.attr().end());
ASSERT_TRUE(iter->second.has_shape());
const auto& shape = iter->second.shape();
EXPECT_EQ(shape.dim_size(), 2);
// Shape updated with placeholder_unknown_output_shape_dim
// and _output_shapes attr.
EXPECT_EQ(shape.dim(0).size(), 64);
EXPECT_EQ(shape.dim(1).size(), 32);
}
} // namespace
} // namespace grappler
} // namespace tensorflow
@@ -0,0 +1,77 @@
/* 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/grappler/grappler_item.h"
#include "tensorflow/core/framework/function_testlib.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/grappler/inputs/trivial_test_graph_input_yielder.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace grappler {
namespace {
class GrapplerItemTest : public ::testing::Test {};
TEST_F(GrapplerItemTest, Basic) {
TrivialTestGraphInputYielder fake_input(4, 1, 10, false, {{"CPU:0"}});
GrapplerItem item;
CHECK(fake_input.NextItem(&item));
EXPECT_TRUE(item.InitOpsFanin().empty());
std::vector<std::string> graph_nodes;
for (const auto& node : item.graph.node()) {
graph_nodes.push_back(node.name());
}
std::vector<std::string> main_ops;
for (const auto& node : item.MainOpsFanin()) {
main_ops.push_back(node->name());
}
std::sort(graph_nodes.begin(), graph_nodes.end());
std::sort(main_ops.begin(), main_ops.end());
EXPECT_EQ(main_ops, graph_nodes);
}
TEST_F(GrapplerItemTest, InferDevices) {
using test::function::NDef;
const std::string cpu0 = "/job:work/replica:1/task:1/device:CPU:0";
const std::string cpu1 = "/job:work/replica:1/task:1/device:CPU:1";
const std::string cpu2 = "/device:CPU:2";
GrapplerItem item;
item.graph = test::function::GDef(
{
NDef("a", "Placeholder", {}, {{"dtype", DT_FLOAT}}, cpu0),
NDef("b", "Placeholder", {}, {{"dtype", DT_FLOAT}}, cpu1),
NDef("c", "Placeholder", {}, {{"dtype", DT_FLOAT}}, cpu2),
},
{} /* Empty function library */);
ASSERT_FALSE(item.InferDevicesFromGraph().ok());
EXPECT_EQ(item.devices().size(), 2);
EXPECT_NE(item.devices().find(cpu0), item.devices().end());
EXPECT_NE(item.devices().find(cpu1), item.devices().end());
item.ClearDevices();
EXPECT_EQ(item.devices().size(), 0);
}
} // namespace
} // namespace grappler
} // namespace tensorflow
+96
View File
@@ -0,0 +1,96 @@
load("//tensorflow:tensorflow.bzl", "tf_cc_test")
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
licenses = ["notice"],
)
cc_library(
name = "utils",
srcs = [
"utils.cc",
],
hdrs = [
"utils.h",
],
visibility = ["//visibility:public"],
deps = [
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/status",
],
)
tf_cc_test(
name = "utils_test",
srcs = [
"utils_test.cc",
],
deps = [
":utils",
"//tensorflow/core:lib",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/framework:graph_proto_cc",
"//tensorflow/core/protobuf:for_core_protos_cc",
"@com_google_absl//absl/status",
"@tsl//tsl/platform:status",
],
)
cc_library(
name = "input_yielder",
hdrs = [
"input_yielder.h",
],
visibility = ["//visibility:public"],
deps = [],
)
cc_library(
name = "trivial_test_graph_input_yielder",
srcs = ["trivial_test_graph_input_yielder.cc"],
hdrs = [
"trivial_test_graph_input_yielder.h",
],
visibility = ["//visibility:public"],
deps = [
":input_yielder",
"//tensorflow/cc:cc_ops",
"//tensorflow/cc:ops",
"//tensorflow/cc:scope",
"//tensorflow/core:all_kernels",
"//tensorflow/core:framework",
"//tensorflow/core:functional_ops_op_lib",
"//tensorflow/core:lib",
"//tensorflow/core:math_ops_op_lib",
"//tensorflow/core:no_op_op_lib",
"//tensorflow/core:portable_gif_internal",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:sendrecv_ops_op_lib",
"//tensorflow/core/grappler:grappler_item",
"//tensorflow/core/kernels:aggregate_ops",
"@tsl//tsl/platform:status",
],
)
cc_library(
name = "file_input_yielder",
srcs = ["file_input_yielder.cc"],
hdrs = [
"file_input_yielder.h",
],
visibility = ["//visibility:public"],
deps = [
":input_yielder",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/grappler:grappler_item",
"//tensorflow/core/grappler:grappler_item_builder",
"//tensorflow/core/grappler:utils",
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/strings",
],
)
@@ -0,0 +1,141 @@
/* 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/grappler/inputs/file_input_yielder.h"
#include <cstddef>
#include <memory>
#include <string>
#include <unordered_set>
#include <utility>
#include <vector>
#include "absl/log/check.h"
#include "absl/log/log.h"
#include "absl/strings/str_cat.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/grappler/grappler_item.h"
#include "tensorflow/core/grappler/grappler_item_builder.h"
#include "tensorflow/core/grappler/utils.h"
#include "tensorflow/core/lib/strings/strcat.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/fingerprint.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
namespace tensorflow {
namespace grappler {
FileInputYielder::FileInputYielder(const std::vector<std::string>& filenames,
size_t max_iterations)
: filenames_(filenames),
current_file_(0),
current_iteration_(0),
max_iterations_(max_iterations),
bad_inputs_(0) {
CHECK_GT(filenames.size(), 0) << "List of filenames is empty.";
}
bool FileInputYielder::NextItem(GrapplerItem* item) {
if (filenames_.size() == bad_inputs_) {
// All the input files are bad, give up.
return false;
}
if (current_file_ >= filenames_.size()) {
if (current_iteration_ >= max_iterations_) {
return false;
} else {
++current_iteration_;
current_file_ = 0;
bad_inputs_ = 0;
}
}
const std::string& filename = filenames_[current_file_];
++current_file_;
if (!Env::Default()->FileExists(filename).ok()) {
LOG(WARNING) << "Skipping non existent file " << filename;
// Attempt to process the next item on the list
bad_inputs_ += 1;
return NextItem(item);
}
LOG(INFO) << "Loading model from " << filename;
MetaGraphDef metagraph;
absl::Status s = ReadBinaryProto(Env::Default(), filename, &metagraph);
if (!s.ok()) {
s = ReadTextProto(Env::Default(), filename, &metagraph);
}
if (!s.ok()) {
LOG(WARNING) << "Failed to read MetaGraphDef from " << filename << ": "
<< s;
// Attempt to process the next item on the list
bad_inputs_ += 1;
return NextItem(item);
}
if (metagraph.collection_def().count("train_op") == 0 ||
!metagraph.collection_def().at("train_op").has_node_list() ||
metagraph.collection_def().at("train_op").node_list().value_size() == 0) {
LOG(ERROR) << "No train op specified";
bad_inputs_ += 1;
metagraph = MetaGraphDef();
return NextItem(item);
} else {
std::unordered_set<std::string> train_ops;
for (const std::string& val :
metagraph.collection_def().at("train_op").node_list().value()) {
train_ops.insert(NodeName(val));
}
std::unordered_set<std::string> train_ops_found;
for (auto& node : metagraph.graph_def().node()) {
if (train_ops.find(node.name()) != train_ops.end()) {
train_ops_found.insert(node.name());
}
}
if (train_ops_found.size() != train_ops.size()) {
for (const auto& train_op : train_ops) {
if (train_ops_found.find(train_op) != train_ops_found.end()) {
LOG(ERROR) << "Non existent train op specified: " << train_op;
}
}
bad_inputs_ += 1;
metagraph = MetaGraphDef();
return NextItem(item);
}
}
const std::string id =
absl::StrCat(Fingerprint64(metagraph.SerializeAsString()));
ItemConfig cfg;
std::unique_ptr<GrapplerItem> new_item =
GrapplerItemFromMetaGraphDef(id, metagraph, cfg);
if (new_item == nullptr) {
bad_inputs_ += 1;
metagraph = MetaGraphDef();
return NextItem(item);
}
*item = std::move(*new_item);
return true;
}
} // end namespace grappler
} // end namespace tensorflow
@@ -0,0 +1,56 @@
/* 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.
==============================================================================*/
// The file input provides a mechanism to feed grappler with existing TensorFlow
// graphs stored in TensorFlow checkpoints. Note that at this point the weights
// that may be stored in the checkpoint are not restored in order to speedup the
// initialization.
#ifndef TENSORFLOW_CORE_GRAPPLER_INPUTS_FILE_INPUT_YIELDER_H_
#define TENSORFLOW_CORE_GRAPPLER_INPUTS_FILE_INPUT_YIELDER_H_
#include <stddef.h>
#include <limits>
#include <vector>
#include "tensorflow/core/grappler/inputs/input_yielder.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace grappler {
class GrapplerItem;
class FileInputYielder : public InputYielder {
public:
// Iterates over the files specified in the list of 'filename' up to
// 'max_iterations' times.
explicit FileInputYielder(
const std::vector<std::string>& filenames,
size_t max_iterations = std::numeric_limits<size_t>::max());
bool NextItem(GrapplerItem* item) override;
private:
const std::vector<std::string> filenames_;
size_t current_file_;
size_t current_iteration_;
size_t max_iterations_;
size_t bad_inputs_;
};
} // end namespace grappler
} // end namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPPLER_INPUTS_FILE_INPUT_YIELDER_H_
@@ -0,0 +1,35 @@
/* 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_GRAPPLER_INPUTS_INPUT_YIELDER_H_
#define TENSORFLOW_CORE_GRAPPLER_INPUTS_INPUT_YIELDER_H_
namespace tensorflow {
namespace grappler {
struct GrapplerItem;
// Abstract interface for yielding graphs that we want to optimize.
class InputYielder {
public:
virtual ~InputYielder() = default;
virtual bool NextItem(GrapplerItem* item) = 0;
};
} // end namespace grappler
} // end namespace tensorflow
#endif // TENSORFLOW_CORE_GRAPPLER_INPUTS_INPUT_YIELDER_H_

Some files were not shown because too many files have changed in this diff Show More