chore: import upstream snapshot with attribution

tensorflow/tools/optimization/BUILD

@@ -0,0 +1,54 @@
+# Description:
+#   Utilities that perform useful transformations on graphs
+
+load(
+    "//tensorflow:tensorflow.bzl",
+    "tf_cc_binary",
+    "tf_cuda_library",
+)
+load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//visibility:public"],
+    licenses = ["notice"],
+)
+
+tf_cuda_library(
+    name = "optimization_pass_runner_lib",
+    srcs = ["optimization_pass_runner.cc"],
+    hdrs = ["optimization_pass_runner.h"],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        "//tensorflow/core:core_cpu",
+        "//tensorflow/core:core_cpu_base",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:framework_internal",
+        "//tensorflow/core:framework_lite",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core:tensorflow",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:string_view",
+    ],
+)
+
+tf_cc_binary(
+    name = "gpu_optimization_pass_runner",
+    srcs = ["gpu_optimization_pass_runner_main.cc"],
+    deps = [
+        ":optimization_pass_runner_lib",
+        "//tensorflow/compiler/jit:xla_cpu_jit",
+        "//tensorflow/compiler/jit:xla_gpu_jit",
+        "//tensorflow/compiler/tf2xla:xla_compiler",
+        "//tensorflow/core:core_cpu_base",
+        "//tensorflow/core:framework_internal",
+        "//tensorflow/core:framework_lite",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core:tensorflow",
+        "@com_google_absl//absl/status",
+        "@xla//xla/tsl/platform:status",
+    ],
+)

tensorflow/tools/optimization/gpu_optimization_pass_runner_main.cc

@@ -0,0 +1,93 @@
+/* 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.
+==============================================================================*/
+// This file creates a binary that can run any registered optimization pass.
+// ./xla_gpu_opt  --input_file_path=/tmp/input.pbtxt
+// --output_file_path=/tmp/output.pbtxt
+// --optimization_pass=NameOfGraphOptimizationPass
+
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "xla/tsl/platform/errors.h"
+#include "xla/tsl/platform/status.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/init_main.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/protobuf/config.pb.h"
+#include "tensorflow/core/util/command_line_flags.h"
+#include "tensorflow/tools/optimization/optimization_pass_runner.h"
+
+namespace tensorflow {
+namespace {
+absl::Status RealMain(int argc, char** argv) {
+  std::string input_file_path;
+  std::string output_file_path;
+  std::string optimization_pass;
+
+  const std::vector<Flag> flag_list = {
+      Flag("input_file_path", &input_file_path, "Location of the input graph."),
+      Flag("output_file_path", &output_file_path,
+           "Location to write the resulting graph."),
+      // For now only a single optimization pass can be run.
+      Flag("optimization_pass", &optimization_pass,
+           "Which optimization pass to run."),
+  };
+  if (!Flags::Parse(&argc, argv, flag_list)) {
+    return absl::FailedPreconditionError("Invalid flags passed");
+  }
+  port::InitMain(argv[0], &argc, &argv);
+
+  if (input_file_path.empty()) {
+    return absl::FailedPreconditionError("input_file_path is a required flag.");
+  }
+  if (output_file_path.empty()) {
+    return absl::FailedPreconditionError(
+        "output_file_path is a required flag.");
+  }
+  if (optimization_pass.empty()) {
+    return absl::FailedPreconditionError(
+        "optimization_pass is a required flag.");
+  }
+
+  GraphDef graphdef_input;
+  TF_RETURN_IF_ERROR(
+      ReadTextProto(Env::Default(), input_file_path, &graphdef_input));
+
+  tensorflow::OptimizationPassRunner runner;
+
+  // Most machines in our servers currently use 8 gpus. There is nothing special
+  // about this number and it can be decreased or increased to test other
+  // configurations.
+  TF_RETURN_IF_ERROR(runner.AddCpus(8));
+  TF_RETURN_IF_ERROR(runner.AddGpus(8));
+
+  // This binary is used to test TF:XLA behavior, so turn on auto_jit.
+  TF_RETURN_IF_ERROR(runner.SetJitLevel(tensorflow::OptimizerOptions::ON_2));
+  GraphDef graphdef_output;
+  TF_RETURN_IF_ERROR(runner.Run(optimization_pass, std::move(graphdef_input),
+                                &graphdef_output));
+  return WriteTextProto(Env::Default(), output_file_path, graphdef_output);
+}
+}  // namespace
+}  // namespace tensorflow
+
+int main(int argc, char** argv) {
+  TF_CHECK_OK(tensorflow::RealMain(argc, argv));
+  return 0;
+}

tensorflow/tools/optimization/optimization_pass_runner.cc

@@ -0,0 +1,156 @@
+/* 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.
+==============================================================================*/
+// This file creates a library that can run any registered optimization pass.
+// The binary that uses this will be run in a form similar to:
+// ./optimization_pass_runner  --input_file_path=/tmp/input.pbtxt
+// --output_file_path=/tmp/output.pbtxt
+// --optimization_pass=NameOfGraphOptimizationPass
+#include "tensorflow/tools/optimization/optimization_pass_runner.h"
+
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "xla/tsl/platform/errors.h"
+#include "tensorflow/core/common_runtime/device_set.h"
+#include "tensorflow/core/common_runtime/graph_constructor.h"
+#include "tensorflow/core/common_runtime/optimization_registry.h"
+#include "tensorflow/core/framework/device_attributes.pb.h"
+#include "tensorflow/core/framework/function.h"
+#include "tensorflow/core/framework/function.pb.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/graph/graph.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/protobuf/config.pb.h"
+#include "tensorflow/core/public/session_options.h"
+
+namespace tensorflow {
+namespace {
+// A fake device used to populate a DeviceSet.
+class FakeDevice : public Device {
+ private:
+  explicit FakeDevice(const DeviceAttributes& device_attributes)
+      : Device(nullptr, device_attributes) {}
+
+ public:
+  absl::Status Sync() override;
+  static std::unique_ptr<Device> Make(const std::string& name,
+                                      const std::string& type);
+};
+
+absl::Status FakeDevice::Sync() {
+  return absl::UnimplementedError("FakeDevice::Sync()");
+}
+
+std::unique_ptr<Device> FakeDevice::Make(const std::string& name,
+                                         const std::string& type) {
+  DeviceAttributes device_attributes;
+  device_attributes.set_name(name);
+  device_attributes.set_device_type(DeviceType(type).type());
+  return std::unique_ptr<Device>(new FakeDevice(device_attributes));
+}
+
+absl::Status FindPassWithName(absl::string_view name,
+                              GraphOptimizationPass** result) {
+  *result = nullptr;
+  // Run the optimization pass specified by the command line flag.
+  for (const auto& groups_and_passes :
+       OptimizationPassRegistry::Global()->groups()) {
+    for (const auto& phase_and_passes : groups_and_passes.second) {
+      for (const auto& pass : phase_and_passes.second) {
+        if (pass->name() == name) {
+          if (*result) {
+            return absl::InternalError(
+                absl::StrCat("Found more than one pass with name ", name));
+          }
+          *result = pass.get();
+        }
+      }
+    }
+  }
+
+  return *result == nullptr ? absl::InternalError(absl::StrCat(
+                                  "Could not find pass with name ", name))
+                            : absl::OkStatus();
+}
+}  // namespace
+
+absl::Status OptimizationPassRunner::Run(absl::string_view pass_to_run,
+                                         GraphDef input, GraphDef* result) {
+  auto session_options = std::make_unique<SessionOptions>();
+  session_options->config.mutable_graph_options()
+      ->mutable_optimizer_options()
+      ->set_global_jit_level(jit_level_);
+  FunctionDefLibrary flib;
+  std::unique_ptr<Graph> graph = std::make_unique<Graph>(OpRegistry::Global());
+
+  GraphOptimizationPassOptions options;
+  options.session_options = session_options.get();
+  options.graph = &graph;
+  std::unique_ptr<FunctionLibraryDefinition> flib_def(
+      new FunctionLibraryDefinition((*options.graph)->op_registry(), flib));
+  options.flib_def = flib_def.get();
+
+  // Grab the data
+  GraphConstructorOptions graph_opts;
+  graph_opts.expect_device_spec = true;
+  graph_opts.allow_internal_ops = true;
+  TF_RETURN_IF_ERROR(ConvertGraphDefToGraph(graph_opts, std::move(input),
+                                            options.graph->get()));
+
+  // Add all devices that were previously configured with AddDevice.
+  DeviceSet device_set;
+  for (auto& device : devices_) {
+    device_set.AddDevice(device.get());
+  }
+  options.device_set = &device_set;
+
+  GraphOptimizationPass* pass;
+  TF_RETURN_IF_ERROR(FindPassWithName(pass_to_run, &pass));
+  TF_RETURN_IF_ERROR(pass->Run(options));
+
+  options.graph->get()->ToGraphDef(result);
+  return absl::OkStatus();
+}
+
+absl::Status OptimizationPassRunner::SetJitLevel(
+    OptimizerOptions::GlobalJitLevel jit_level) {
+  jit_level_ = jit_level;
+  return absl::OkStatus();
+}
+
+absl::Status OptimizationPassRunner::AddDevices(absl::string_view type,
+                                                int count) {
+  for (int i = 0; i < count; i++) {
+    devices_.push_back(FakeDevice::Make(
+        absl::StrCat("/job:localhost/replica:0/task:0/device:", type, ":", i),
+        absl::StrCat(type)));
+    devices_.push_back(FakeDevice::Make(
+        absl::StrCat("/job:localhost/replica:0/task:0/device:XLA_", type, ":",
+                     i),
+        absl::StrCat(type)));
+  }
+
+  return absl::OkStatus();
+}
+}  // namespace tensorflow

tensorflow/tools/optimization/optimization_pass_runner.h

@@ -0,0 +1,63 @@
+/* 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_TOOLS_OPTIMIZATION_OPTIMIZATION_PASS_RUNNER_H_
+#define TENSORFLOW_TOOLS_OPTIMIZATION_OPTIMIZATION_PASS_RUNNER_H_
+
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/core/common_runtime/device.h"
+#include "tensorflow/core/common_runtime/optimization_registry.h"
+#include "tensorflow/core/framework/device.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/protobuf/config.pb.h"
+
+namespace tensorflow {
+
+// OptimizationPassRunner can be initialized, populated with devices, then run
+// to test individual Tensorflow Optimization passes.
+class OptimizationPassRunner {
+ public:
+  explicit OptimizationPassRunner() : jit_level_(OptimizerOptions::DEFAULT) {}
+
+  // Increasing the Jit level will cause XLA to compile parts of the tensorflow
+  // graph that it is able to.
+  absl::Status SetJitLevel(OptimizerOptions::GlobalJitLevel jit_level);
+
+  absl::Status Run(absl::string_view pass_to_run, GraphDef input,
+                   GraphDef* result);
+
+  absl::Status AddCpus(int count) {
+    return AddDevices(tensorflow::DEVICE_CPU, count);
+  }
+
+  absl::Status AddGpus(int count) {
+    return AddDevices(tensorflow::DEVICE_GPU, count);
+  }
+
+ private:
+  absl::Status AddDevices(absl::string_view type, int count);
+
+  OptimizerOptions::GlobalJitLevel jit_level_;
+  std::vector<std::unique_ptr<Device>> devices_;
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_TOOLS_OPTIMIZATION_OPTIMIZATION_PASS_RUNNER_H_