chore: import upstream snapshot with attribution

tensorflow/tools/tfg_graph_transforms/BUILD

@@ -0,0 +1,90 @@
+# Description:
+#   Utilities to perform MLIR TFG graph transformations.
+
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load(
+    "//tensorflow:tensorflow.bzl",
+    "tf_cc_binary",
+)
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "utils",
+    srcs = [
+        "utils.cc",
+    ],
+    hdrs = [
+        "utils.h",
+    ],
+    deps = [
+        "//tensorflow/cc/saved_model/image_format:internal_api",
+        "//tensorflow/core:lib",
+        "//tensorflow/core/platform:errors",
+        "//tensorflow/core/platform:path",
+        "//tensorflow/core/platform:status",
+        "//tensorflow/core/protobuf:for_core_protos_cc",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings:string_view",
+    ],
+)
+
+TFG_GRAPH_TRANSFORM_DEPS = [
+    ":utils",
+    "@llvm-project//llvm:Support",
+    "@llvm-project//mlir:IR",
+    "@llvm-project//mlir:Pass",
+    "@llvm-project//mlir:Transforms",
+    "//tensorflow/compiler/mlir:init_mlir",
+    "//tensorflow/compiler/mlir/tensorflow",
+    "//tensorflow/compiler/mlir/tensorflow:error_util",
+    "//tensorflow/core:lib",
+    "//tensorflow/core/ir:Dialect",
+    "//tensorflow/core/protobuf:for_core_protos_cc",
+    "//tensorflow/core/transforms:PassRegistration",
+    "//tensorflow/compiler/tf2xla/ops:xla_ops",
+    "//tensorflow/core:ops",
+    "//tensorflow/core:protos_all_cc",
+    "//tensorflow/core/ir/importexport:graphdef_export",
+    "//tensorflow/core/ir/importexport:graphdef_import",
+    "//tensorflow/core/ir/importexport:savedmodel_export",
+    "//tensorflow/core/ir/importexport:savedmodel_import",
+    "//tensorflow/core/ir:tf_op_registry",
+]
+
+# Description:
+#   A tool that provides a mechanism to run TFG graph optimizations operating on the
+#   GraphDef from the supplied SavedModel as an input.
+tf_cc_binary(
+    name = "tfg_graph_transforms",
+    srcs = [
+        "tfg_graph_transforms_main.cc",
+    ],
+    deps = TFG_GRAPH_TRANSFORM_DEPS + [
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@llvm-project//mlir:Support",
+    ],
+)
+
+# Description:
+#    The tool wrapped into a library, so that it could be used
+#    when custom defined ops and passes need to be added.
+cc_library(
+    name = "tfg_graph_transforms_main",
+    srcs = [
+        "tfg_graph_transforms_main.cc",
+    ],
+    visibility = ["//visibility:public"],
+    deps = TFG_GRAPH_TRANSFORM_DEPS + [
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@llvm-project//mlir:Support",
+    ],
+)

tensorflow/tools/tfg_graph_transforms/tfg_graph_transforms_main.cc

@@ -0,0 +1,295 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <cstdlib>
+#include <string>
+#include <utility>
+
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "llvm/ADT/Twine.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/LogicalResult.h"
+#include "mlir/IR/AsmState.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/Diagnostics.h"  // from @llvm-project
+#include "mlir/IR/Location.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "mlir/Pass/PassRegistry.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Transforms/Passes.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/init_mlir.h"
+#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/error_util.h"
+#include "xla/tsl/platform/errors.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/framework/graph_debug_info.pb.h"
+#include "tensorflow/core/ir/dialect.h"
+#include "tensorflow/core/ir/importexport/graphdef_export.h"
+#include "tensorflow/core/ir/importexport/graphdef_import.h"
+#include "tensorflow/core/ir/importexport/savedmodel_export.h"
+#include "tensorflow/core/ir/importexport/savedmodel_import.h"
+#include "tensorflow/core/ir/tf_op_registry.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/protobuf/saved_model.pb.h"
+#include "tensorflow/core/transforms/pass_registration.h"
+#include "tensorflow/tools/tfg_graph_transforms/utils.h"
+
+namespace {
+
+llvm::cl::OptionCategory tfg_graph_transform_category(
+    "TFG graph transform options");
+
+// NOLINTNEXTLINE
+llvm::cl::opt<std::string> input_file(
+    llvm::cl::Positional, llvm::cl::desc("<Input model>"),
+    llvm::cl::value_desc("Full path to the input model"),
+    llvm::cl::cat(tfg_graph_transform_category), llvm::cl::Required);
+
+// NOLINTNEXTLINE
+llvm::cl::opt<std::string> output_file(
+    "o", llvm::cl::desc("Output model"),
+    llvm::cl::value_desc("Full path to the output model"),
+    llvm::cl::cat(tfg_graph_transform_category), llvm::cl::Required);
+
+enum class DataFormat { SavedModel = 0, GraphDef = 1 };
+
+// NOLINTNEXTLINE
+llvm::cl::opt<DataFormat> data_format(
+    "data_format",
+    llvm::cl::desc(
+        "Data format for both input and output, e.g., SavedModel or GraphDef"),
+    values(clEnumValN(DataFormat::SavedModel, "savedmodel",
+                      "SavedModel format"),
+           clEnumValN(DataFormat::GraphDef, "graphdef", "GraphDef format")),
+    llvm::cl::init(DataFormat::SavedModel),
+    llvm::cl::cat(tfg_graph_transform_category));
+
+// NOLINTNEXTLINE
+llvm::cl::opt<bool> experimental_image_format(
+    "experimental_image_format",
+    llvm::cl::desc("Whether to expect use the experimental SavedModel image "
+                   "format. Only applies to SavedModel inputs and outputs. "
+                   "When enabled, the output filename may have a different "
+                   "extension than the one provided."),
+    llvm::cl::init(false));
+
+// NOLINTNEXTLINE
+llvm::cl::opt<int> experimental_image_format_max_proto_size(
+    "experimental_image_format_max_proto_size",
+    llvm::cl::desc(
+        "Sets the maximum chunk size in bytes allowed for protos (2GB by "
+        "default). This flag should only be used for testing purposes and can "
+        "be removed at any time."),
+    llvm::cl::init(0));
+
+// Validate CL options and returns false in case of an error.
+bool CheckCLParams() {
+  if (input_file == output_file) {
+    LOG(WARNING)
+        << "Input and output files are set to the same location. "
+           "The resulted model protobuf will overwrite the original one.\n";
+  }
+  if (!tensorflow::Env::Default()->FileExists(input_file).ok()) {
+    LOG(ERROR) << "Provided file or directory does not exist: '" << input_file
+               << "'\n";
+    return false;
+  }
+
+  if (tensorflow::Env::Default()->IsDirectory(input_file).ok()) {
+    LOG(ERROR)
+        << "Expected full path to the model protobuf file, given directory: '"
+        << input_file << "'\n";
+    return false;
+  }
+
+  return true;
+}
+
+void RegisterDialects(mlir::DialectRegistry& registry) {
+  // This potentially could be limited, for now keep all TF.
+  mlir::RegisterAllTensorFlowDialects(registry);
+
+  // Register the TF op registry interface so that passes can query it.
+  registry.addExtension(
+      +[](mlir::MLIRContext* ctx, mlir::tfg::TFGraphDialect* dialect) {
+        dialect->addInterfaces<mlir::tfg::TensorFlowOpRegistryInterface>();
+      });
+}
+
+absl::Status RunOptimizationPasses(
+    const mlir::PassPipelineCLParser& passPipeline, mlir::ModuleOp module,
+    mlir::MLIRContext* context) {
+  mlir::PassManager pm(context);
+  mlir::registerPassManagerCLOptions();
+  if (failed(mlir::applyPassManagerCLOptions(pm))) {
+    return absl::InvalidArgumentError(
+        "Could not initialize MLIR pass manager CL options");
+  }
+
+  auto error_handler = [&](const llvm::Twine& msg) {
+    emitError(mlir::UnknownLoc::get(pm.getContext())) << msg;
+    return mlir::failure();
+  };
+  if (failed(passPipeline.addToPipeline(pm, error_handler))) {
+    return absl::InvalidArgumentError("Pipeline initialization failed");
+  }
+
+  mlir::StatusScopedDiagnosticHandler diagnostics_handler(context);
+  if (failed(pm.run(module))) {
+    return diagnostics_handler.Combine(
+        absl::InvalidArgumentError("MLIR Pass Manager failure: "));
+  }
+
+  return diagnostics_handler.ConsumeStatus();
+}
+
+// Import model to the TFG MLIR module.
+absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> ImportModel(
+    DataFormat data_format, const std::string& input_file,
+    bool experimental_image_format, mlir::MLIRContext* mlir_context) {
+  tensorflow::GraphDebugInfo debug_info;
+
+  switch (data_format) {
+    case DataFormat::SavedModel: {
+      tensorflow::SavedModel saved_model;
+      if (experimental_image_format) {
+        TF_RETURN_IF_ERROR(
+            mlir::tfg::graph_transforms::ReadSavedModelImageFormat(
+                input_file, saved_model));
+      } else {
+        TF_RETURN_IF_ERROR(
+            mlir::tfg::graph_transforms::ReadModelProto<tensorflow::SavedModel>(
+                input_file, saved_model));
+      }
+      return mlir::tfg::ImportSavedModelToMlir(mlir_context, debug_info,
+                                               saved_model);
+    }
+    case DataFormat::GraphDef: {
+      tensorflow::GraphDef graph_def;
+      TF_RETURN_IF_ERROR(
+          mlir::tfg::graph_transforms::ReadModelProto<tensorflow::GraphDef>(
+              input_file, graph_def));
+      return mlir::tfg::ImportGraphDef(mlir_context, debug_info, graph_def);
+    }
+  }
+}
+
+absl::Status ExportTFGModule(mlir::ModuleOp module_op, DataFormat data_format,
+                             const std::string& input_file,
+                             const std::string& output_file,
+                             bool experimental_image_format,
+                             int experimental_image_format_max_size) {
+  switch (data_format) {
+    case DataFormat::SavedModel: {
+      tensorflow::SavedModel original_saved_model;
+      if (experimental_image_format) {
+        TF_RETURN_IF_ERROR(
+            mlir::tfg::graph_transforms::ReadSavedModelImageFormat(
+                input_file, original_saved_model));
+      } else {
+        TF_RETURN_IF_ERROR(
+            mlir::tfg::graph_transforms::ReadModelProto<tensorflow::SavedModel>(
+                input_file, original_saved_model));
+      }
+
+      tensorflow::SavedModel final_saved_model;
+
+      TF_RETURN_WITH_CONTEXT_IF_ERROR(
+          mlir::tfg::ExportMlirToSavedModel(module_op, original_saved_model,
+                                            &final_saved_model),
+          "while converting TFG to SavedModel");
+
+      if (experimental_image_format) {
+        VLOG(1) << "Serializing resulting SavedModel to " << output_file
+                << " (filename might not exactly match since "
+                   "`experimental_image_format` has been enabled)";
+        return mlir::tfg::graph_transforms::WriteSavedModelImageFormat(
+            &final_saved_model, output_file,
+            experimental_image_format_max_proto_size);
+      } else {
+        VLOG(1) << "Serializing resulting SavedModel to " << output_file;
+        return mlir::tfg::graph_transforms::SerializeProto<
+            tensorflow::SavedModel>(final_saved_model, output_file);
+      }
+    }
+    case DataFormat::GraphDef: {
+      tensorflow::GraphDef new_graphdef;
+      TF_RETURN_WITH_CONTEXT_IF_ERROR(
+          mlir::tfg::ConvertToGraphDef(module_op, &new_graphdef),
+          "while converting TFG to GraphDef");
+
+      VLOG(1) << "Serializing resulting GraphDef to " << output_file;
+      return mlir::tfg::graph_transforms::SerializeProto<tensorflow::GraphDef>(
+          new_graphdef, output_file);
+    }
+  }
+}
+
+}  // namespace
+
+int main(int argc, char** argv) {
+  tensorflow::InitMlir y(&argc, &argv);
+  mlir::registerAsmPrinterCLOptions();
+  mlir::registerMLIRContextCLOptions();
+  mlir::registerPassManagerCLOptions();
+  mlir::tfg::registerTFGraphPasses();
+  mlir::registerSymbolPrivatizePass();
+  mlir::registerSymbolDCEPass();
+
+  mlir::PassPipelineCLParser pass_pipeline("", "TFG passes to run");
+  llvm::cl::ParseCommandLineOptions(argc, argv, "TFG optimization tool\n");
+
+  if (!CheckCLParams()) {
+    LOG(QFATAL) << "Command line parameters are invalid";
+  }
+
+  mlir::DialectRegistry registry;
+  RegisterDialects(registry);
+  mlir::MLIRContext context(registry);
+
+  // Import model to the TFG MLIR module.
+  auto module_ref_status =
+      ImportModel(data_format, input_file, experimental_image_format, &context);
+
+  if (!module_ref_status.ok()) {
+    LOG(QFATAL) << "Model import failed: " << module_ref_status.status();
+  }
+  auto module_ref = std::move(module_ref_status.value());
+
+  // Parse the optimization pipeline configuration and run requested graph
+  // optimizations.
+  absl::Status pass_pipeline_status =
+      RunOptimizationPasses(pass_pipeline, *module_ref, &context);
+  if (!pass_pipeline_status.ok()) {
+    LOG(QFATAL) << pass_pipeline_status << "\n";
+  }
+
+  // Export MLIR TFG module to the resulting model proto.
+  absl::Status export_status = ExportTFGModule(
+      *module_ref, data_format, input_file, output_file,
+      experimental_image_format, experimental_image_format_max_proto_size);
+
+  if (!export_status.ok()) {
+    LOG(QFATAL) << "Export of TFG module failed: " << export_status << "\n";
+  }
+
+  return EXIT_SUCCESS;
+}

tensorflow/tools/tfg_graph_transforms/utils.cc

@@ -0,0 +1,69 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/tools/tfg_graph_transforms/utils.h"
+
+#include <string>
+
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/cc/saved_model/image_format/internal_api.h"
+#include "tensorflow/core/platform/path.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/stringpiece.h"
+#include "tensorflow/core/protobuf/saved_model.pb.h"
+#include "tsl/platform/stringpiece.h"
+
+namespace mlir {
+namespace tfg {
+namespace graph_transforms {
+
+namespace {
+
+absl::string_view GetNameWithoutExtension(absl::string_view filename) {
+  auto pos = filename.rfind('.');
+  if (pos == absl::string_view::npos) return filename;
+  return filename.substr(0, pos);
+}
+
+}  // namespace
+
+bool IsTextProto(const std::string& input_file) {
+  absl::string_view extension = tensorflow::io::Extension(input_file);
+  return !extension.compare("pbtxt");
+}
+
+absl::Status ReadSavedModelImageFormat(const std::string& input_file,
+                                       tensorflow::SavedModel& model_proto) {
+  std::string saved_model_prefix(GetNameWithoutExtension(input_file));
+  return tensorflow::image_format::ReadSavedModel(saved_model_prefix,
+                                                  &model_proto);
+}
+absl::Status WriteSavedModelImageFormat(tensorflow::SavedModel* model_proto,
+                                        const std::string& output_file,
+                                        int debug_max_size) {
+  std::string saved_model_prefix(GetNameWithoutExtension(output_file));
+  if (debug_max_size > 0) {
+    return tensorflow::image_format::WriteSavedModel(
+        model_proto, saved_model_prefix, debug_max_size);
+  } else {
+    return tensorflow::image_format::WriteSavedModel(model_proto,
+                                                     saved_model_prefix);
+  }
+}
+
+}  // namespace graph_transforms
+}  // namespace tfg
+}  // namespace mlir

tensorflow/tools/tfg_graph_transforms/utils.h

@@ -0,0 +1,108 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_TOOLS_TFG_GRAPH_TRANSFORMS_UTILS_H_
+#define TENSORFLOW_TOOLS_TFG_GRAPH_TRANSFORMS_UTILS_H_
+
+#include <string>
+
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
+#include "xla/tsl/platform/errors.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/path.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/stringpiece.h"
+#include "tensorflow/core/protobuf/saved_model.pb.h"
+
+namespace mlir {
+namespace tfg {
+namespace graph_transforms {
+
+// Reads the model proto from `input_file`.
+// If the format of proto cannot be identified based on the file extension,
+// attempts to load in a binary format first and then in a text format.
+template <class T>
+absl::Status ReadModelProto(const std::string& input_file, T& model_proto) {
+  // Proto might be either in binary or text format.
+  absl::string_view extension = tensorflow::io::Extension(input_file);
+  bool binary_extenstion = !extension.compare("pb");
+  bool text_extension = !extension.compare("pbtxt");
+
+  if (!binary_extenstion && !text_extension) {
+    LOG(WARNING) << "Proto type cannot be identified based on the extension";
+    // Try load binary first.
+    auto status = tensorflow::ReadBinaryProto(tensorflow::Env::Default(),
+                                              input_file, &model_proto);
+    if (status.ok()) {
+      return status;
+    }
+
+    // Binary proto loading failed, attempt to load text proto.
+    return tensorflow::ReadTextProto(tensorflow::Env::Default(), input_file,
+                                     &model_proto);
+  }
+
+  if (binary_extenstion) {
+    return tensorflow::ReadBinaryProto(tensorflow::Env::Default(), input_file,
+                                       &model_proto);
+  }
+
+  if (text_extension) {
+    return tensorflow::ReadTextProto(tensorflow::Env::Default(), input_file,
+                                     &model_proto);
+  }
+
+  return absl::InvalidArgumentError("Expected either binary or text protobuf");
+}
+
+// Best effort to identify if the protobuf file `input_file` is
+// in a text or binary format.
+bool IsTextProto(const std::string& input_file);
+
+template <class T>
+absl::Status SerializeProto(T model_proto, const std::string& output_file) {
+  auto output_dir = tensorflow::io::Dirname(output_file);
+
+  TF_RETURN_IF_ERROR(tensorflow::Env::Default()->RecursivelyCreateDir(
+      {output_dir.data(), output_dir.length()}));
+  if (IsTextProto(output_file)) {
+    TF_RETURN_WITH_CONTEXT_IF_ERROR(
+        tensorflow::WriteTextProto(tensorflow::Env::Default(), output_file,
+                                   model_proto),
+        "Error while writing the resulting model proto");
+  } else {
+    TF_RETURN_WITH_CONTEXT_IF_ERROR(
+        tensorflow::WriteBinaryProto(tensorflow::Env::Default(), output_file,
+                                     model_proto),
+        "Error while writing the resulting model proto");
+  }
+  return absl::OkStatus();
+}
+
+// Read and write to the experimental SavedModel Image format.
+absl::Status ReadSavedModelImageFormat(const std::string& input_file,
+                                       tensorflow::SavedModel& model_proto);
+absl::Status WriteSavedModelImageFormat(tensorflow::SavedModel* model_proto,
+                                        const std::string& output_file,
+                                        int debug_max_size);
+
+}  // namespace graph_transforms
+}  // namespace tfg
+}  // namespace mlir
+
+#endif  // TENSORFLOW_TOOLS_TFG_GRAPH_TRANSFORMS_UTILS_H_