chore: import upstream snapshot with attribution

tensorflow/core/util/activation_mode.cc

@@ -0,0 +1,50 @@
+/* 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/util/activation_mode.h"
+
+#include <string>
+
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "tensorflow/core/framework/node_def_util.h"
+#include "tensorflow/core/lib/core/errors.h"
+
+namespace tensorflow {
+
+absl::Status GetActivationModeFromString(const std::string& str_value,
+                                         ActivationMode* value) {
+  if (str_value == "None") {
+    *value = NONE;
+  } else if (str_value == "Sigmoid") {
+    *value = SIGMOID;
+  } else if (str_value == "Relu") {
+    *value = RELU;
+  } else if (str_value == "Relu6") {
+    *value = RELU6;
+  } else if (str_value == "ReluX") {
+    *value = RELUX;
+  } else if (str_value == "Tanh") {
+    *value = TANH;
+  } else if (str_value == "BandPass") {
+    *value = BANDPASS;
+  } else {
+    return absl::NotFoundError(
+        absl::StrCat(str_value, " is not an allowed activation mode"));
+  }
+  return absl::OkStatus();
+}
+
+}  // end namespace tensorflow

tensorflow/core/util/activation_mode.h

@@ -0,0 +1,67 @@
+/* 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_UTIL_ACTIVATION_MODE_H_
+#define TENSORFLOW_CORE_UTIL_ACTIVATION_MODE_H_
+
+// This file contains helper routines to deal with activation mode in various
+// ops and kernels.
+
+#include <string>
+
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/lib/core/status.h"
+
+namespace tensorflow {
+
+// ActivationMode: the activation function we apply to the input tensor:
+enum ActivationMode {
+  NONE = 0,
+  SIGMOID = 1,
+  RELU = 2,
+  RELU6 = 3,
+  RELUX = 4,
+  TANH = 5,
+  BANDPASS = 6,
+};
+
+// Specialization to parse an attribute directly into a ActivationMode enum.
+absl::Status GetActivationModeFromString(const std::string& str_value,
+                                         ActivationMode* value);
+
+inline absl::string_view ToString(ActivationMode mode) {
+  switch (mode) {
+    case NONE:
+      return "NONE";
+    case SIGMOID:
+      return "SIGMOID";
+    case RELU:
+      return "RELU";
+    case RELU6:
+      return "RELU6";
+    case RELUX:
+      return "RELUX";
+    case TANH:
+      return "TANH";
+    case BANDPASS:
+      return "BANDPASS";
+  }
+}
+
+}  // end namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_UTIL_ACTIVATION_MODE_H_

tensorflow/core/util/autotune_maps/BUILD

@@ -0,0 +1,220 @@
+# Defines data structures that store autotuning results.
+# The autotuning results can be serialized as a string, allowing you to save
+# and later restore them.
+
+load("@rules_python//python:proto.bzl", "py_proto_library")
+load(
+    "//tensorflow:tensorflow.bzl",
+    "tf_cc_test",
+    "tf_cuda_library",
+    "tf_cuda_only_cc_test",
+)
+
+# TODO(ruochengw): Currently only supports contrib's fused_conv2d_bias_activation_op.
+# We plan to add more ops and move fused_conv2d_bias_activation_op back into core library.
+load(
+    "//tensorflow/core/platform:build_config.bzl",
+    "tf_proto_library",
+    # copybara:uncomment "tf_pyclif_proto_library",
+)
+load(
+    "//tensorflow/core/platform:rules_cc.bzl",
+    "cc_library",
+)
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//tensorflow:__subpackages__"],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "conv_autotune_maps",
+    hdrs = [
+        "conv_autotune_maps.h",
+    ],
+    deps = [
+        "//tensorflow/core/kernels:gpu_util_hdrs",
+    ],
+)
+
+cc_library(
+    name = "conv_map_wrapper",
+    srcs = ["conv_map_wrapper.cc"],
+    hdrs = ["conv_map_wrapper.h"],
+    visibility = [
+        "//tensorflow:__subpackages__",
+        "//waymo/ml/deploy:__subpackages__",
+    ],
+    deps = [
+        ":autotune_map_proto_cc",
+        ":conv_parameters_proto_cc",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@xla//xla/tsl/lib/strings:proto_serialization",
+        "@xla//xla/tsl/protobuf:dnn_proto_cc",
+    ],
+)
+
+tf_cc_test(
+    name = "conv_map_wrapper_test",
+    srcs = ["conv_map_wrapper_test.cc"],
+    deps = [
+        ":autotune_map_proto_cc",
+        ":conv_map_wrapper",
+        ":conv_parameters_proto_cc",
+        "@com_google_googletest//:gtest_main",
+        "@xla//xla/tsl/platform:statusor",
+        "@xla//xla/tsl/protobuf:dnn_proto_cc",
+    ],
+)
+
+tf_proto_library(
+    name = "conv_parameters_proto",
+    srcs = [
+        "conv_parameters.proto",
+    ],
+    protodeps = [
+        "//tensorflow/core/framework:types_proto",
+        "@xla//xla/tsl/protobuf:dnn_proto",
+    ],
+)
+
+# This ensures the compilation of custom op fused_conv2d_bias_activation_op.
+# TODO(ruochengw): Remove this target once we move fused_conv2d_bias_activation_op back in core library.
+cc_library(
+    name = "conv_parameters_hdrs",
+    hdrs = ["conv_parameters.h"],
+    deps = [
+        ":conv_parameters_proto_cc",
+    ],
+)
+
+# For a more maintainable build this target should not exist and the headers
+# should  be split into the existing cc_library targets, but this change was
+# automatically  done so that we can remove long standing issues and complexity
+# in the build system. It's up to the OWNERS of this package to get rid of it or
+# not. The use of the textual_hdrs attribute is discouraged, use hdrs instead.
+# Here it is used to avoid header parsing errors in packages where the feature
+# parse_headers was enabled since loose headers were not being parsed. See
+# go/loose-lsc-one-target-approach for more details.
+cc_library(
+    name = "loose_headers",
+    tags = ["avoid_dep"],
+    textual_hdrs = [
+        "conv_autotune_maps.h",
+        "conv_parameters.h",
+    ],
+    visibility = ["//visibility:private"],
+)
+
+tf_cuda_library(
+    name = "conv_parameters",
+    srcs = ["conv_parameters.cc"],
+    hdrs = [
+        "conv_parameters.h",
+    ],
+    cuda_deps = [
+        "@xla//xla/tsl/lib/strings:proto_serialization",
+    ],
+    deps = [
+        ":conv_parameters_proto_cc",
+        "//tensorflow/core/platform:hash",
+        "//tensorflow/core/platform:protobuf",
+        "//tensorflow/core/platform:stream_executor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:str_format",
+        "@com_google_absl//absl/types:optional",
+        "@tsl//tsl/platform:regexp",
+    ],
+)
+
+tf_proto_library(
+    name = "autotune_map_proto",
+    srcs = [
+        "autotune_map.proto",
+    ],
+    protodeps = [
+        "//tensorflow/core/util/autotune_maps:conv_parameters_proto",
+        "@xla//xla/tsl/protobuf:dnn_proto",
+    ],
+    visibility = [
+        "//waymo/ml/deploy/benchmark:__subpackages__",
+        "//waymo/ml/deploy/system/autotuning:__subpackages__",
+    ],
+)
+
+# copybara:uncomment_begin(google-only)
+# py_proto_library(
+#     name = "autotune_map_py_pb2",
+#     has_services = 0,
+#     visibility = ["//waymo/ml/deploy/system/autotuning:__subpackages__"],
+#     deps = [":autotune_map_proto"],
+# )
+#
+# tf_pyclif_proto_library(
+#     name = "autotune_map_pyclif",
+#     proto_lib = ":autotune_map_proto",
+#     visibility = ["//waymo/ml/deploy/system/autotuning:__subpackages__"],
+# )
+# copybara:uncomment_end
+
+tf_cuda_library(
+    name = "autotune_serialize",
+    srcs = [
+        "autotune_serialize.cc",
+    ],
+    hdrs = [
+        "autotune_serialize.h",
+    ],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":autotune_map_proto_cc",
+        ":conv_autotune_maps",
+        ":conv_parameters",
+        ":conv_parameters_proto_cc",
+        "//tensorflow/core:framework",
+        "//tensorflow/core/platform:status",
+        "//tensorflow/core/platform:str_util",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings:string_view",
+        "@xla//xla:status_macros",
+        "@xla//xla/stream_executor:dnn",
+        "@xla//xla/stream_executor:platform_manager",
+        "@xla//xla/stream_executor/gpu:gpu_init",
+        "@xla//xla/tsl/lib/strings:proto_serialization",
+        "@xla//xla/tsl/protobuf:dnn_proto_cc",
+    ],
+)
+
+tf_cuda_only_cc_test(
+    name = "autotune_serialize_test",
+    size = "small",
+    srcs = ["autotune_serialize_test.cc"],
+    features = ["-layering_check"],
+    tags = ["no_rocm"],
+    deps = [
+        ":autotune_serialize",
+        ":conv_autotune_maps",
+        ":conv_parameters",
+        ":conv_parameters_proto_cc",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core/platform:status_matchers",
+        "@xla//xla/stream_executor/gpu:gpu_init",
+    ],
+)
+
+tf_cuda_only_cc_test(
+    name = "conv_parameters_test",
+    size = "small",
+    srcs = ["conv_parameters_test.cc"],
+    features = ["-layering_check"],
+    tags = ["no_rocm"],
+    deps = [
+        ":conv_parameters",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+    ],
+)

tensorflow/core/util/autotune_maps/autotune_map.proto

@@ -0,0 +1,40 @@
+/* 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.
+==============================================================================*/
+
+// Protocol messages for describing autotuning maps. This enables us to easily
+// serialize/deserialize existing autotune maps in the TF runtime.
+//
+// For Google-internal use only.
+syntax = "proto3";
+
+package tensorflow;
+
+import "xla/tsl/protobuf/dnn.proto";
+import "tensorflow/core/util/autotune_maps/conv_parameters.proto";
+
+message ConvMapProto {
+  message Entry {
+    tensorflow.ConvParametersProto key = 1;
+    stream_executor.dnn.AlgorithmConfigProto value = 2;
+  }
+
+  repeated Entry kv_pairs = 1;
+}
+
+// TODO(b/189530096): Support autotune maps for more ops.
+message AutotuneMapsProto {
+  ConvMapProto conv_map = 2;
+  ConvMapProto fused_conv_map = 3;
+}

tensorflow/core/util/autotune_maps/autotune_serialize.cc

@@ -0,0 +1,214 @@
+/* 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.
+==============================================================================*/
+
+// For Google-internal use only.
+#include "tensorflow/core/util/autotune_maps/autotune_serialize.h"
+
+#include <map>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include "xla/status_macros.h"
+#include "xla/stream_executor/dnn.h"
+#include "xla/stream_executor/gpu/gpu_init.h"
+#include "xla/stream_executor/platform_manager.h"
+#include "xla/tsl/lib/strings/proto_serialization.h"
+#include "xla/tsl/protobuf/dnn.pb.h"
+#include "tensorflow/core/platform/str_util.h"
+#include "tensorflow/core/util/activation_mode.h"
+#include "tensorflow/core/util/autotune_maps/autotune_map.pb.h"
+#include "tensorflow/core/util/autotune_maps/conv_autotune_maps.h"
+#include "tensorflow/core/util/autotune_maps/conv_parameters.h"
+#include "tensorflow/core/util/autotune_maps/conv_parameters.pb.h"
+
+namespace tensorflow {
+
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+namespace {
+using stream_executor::dnn::AlgorithmConfig;
+using stream_executor::dnn::AlgorithmConfigProto;
+using stream_executor::dnn::AlgorithmDesc;
+using stream_executor::dnn::AlgorithmProto;
+
+template <typename Op>
+absl::StatusOr<ConvMapProto> ConvMapToProto(
+    const AutotuneMap<ConvParameters, AutotuneEntry<Op>>& autotune_map) {
+  ConvMapProto proto;
+
+  // Deterministically sort the entries in autotune maps
+  // according to the serialized string of ConvParametersProto in order to
+  // enable deterministic serialization. The actual order is meaningless.
+  //
+  // This step also filters out duplicate entries (only device_id's are
+  // different) in the autotune maps. So that there is only one entry for a
+  // convolution operation with a specific GPU device type.
+  std::map<std::string, ConvMapProto::Entry> sorted_map;
+
+  for (auto const &p : autotune_map.GetMap()) {
+    const ConvParameters &params = p.first;
+    const ConvParametersProto &params_proto = params.proto();
+    VLOG(1) << "Reading: " << params.ToString();
+
+    ConvMapProto::Entry kv;
+    *kv.mutable_key() = params_proto;
+
+    if (p.second.is_algorithm_config()) {
+      *kv.mutable_value() = p.second.GetAlgorithmConfig().ToProto();
+    } else {
+      const auto &runners = p.second.GetOpRunners();
+      *kv.mutable_value()->mutable_algorithm() =
+          runners.primary->ToAlgorithmDesc().ToProto();
+      if (runners.no_scratch_fallback) {
+        *kv.mutable_value()->mutable_algorithm_no_scratch() =
+            runners.no_scratch_fallback->ToAlgorithmDesc().ToProto();
+      }
+    }
+
+    std::string serialized_params;
+    TF_RET_CHECK(
+        tsl::SerializeToStringDeterministic(params_proto, &serialized_params));
+    sorted_map.insert(std::make_pair(std::move(serialized_params), kv));
+  }
+
+  for (auto const &p : sorted_map) {
+    ConvMapProto::Entry *kv = proto.add_kv_pairs();
+    *kv = p.second;
+  }
+  return proto;
+}
+
+template <typename Op>
+absl::Status PopulateConvMap(
+    const ConvMapProto& m,
+    AutotuneMap<ConvParameters, AutotuneEntry<Op>>* autotune_map) {
+  if (m.kv_pairs().size() == 0) {
+    return absl::OkStatus();
+  }
+
+  // Get the list of all GPU StreamExecutors.
+  TF_ASSIGN_OR_RETURN(
+      se::Platform * platform,
+      se::PlatformManager::PlatformWithName(se::GpuPlatformName()));
+  std::vector<std::string> device_descs;
+  for (int i = 0; i < platform->VisibleDeviceCount(); i++) {
+    TF_ASSIGN_OR_RETURN(std::unique_ptr<se::DeviceDescription> device_desc,
+                        platform->DescriptionForDevice(i));
+    device_descs.push_back(
+        DeviceIdentifierForAutotuning(device_desc->model_str()));
+  }
+
+  std::set<std::string> unmatched_device_descs;
+  for (const ConvMapProto::Entry &kv : m.kv_pairs()) {
+    const ConvParametersProto &params_proto = kv.key();
+    // Abort the loading process whenever there is an entry whose version number
+    // doesn't match runtime version because the autotune results may be
+    // incorrect.
+    if (params_proto.version() != ConvParameters::kVersion) {
+      VLOG(1) << "ConvParametersProto with the incompatible version:"
+              << params_proto.DebugString();
+      return absl::AbortedError(absl::StrCat(
+          "Aborted because the loaded autotune results for convolution "
+          "operations have a version different "
+          "from runtime's version. Expected version: ",
+          ConvParameters::kVersion,
+          ". Actual version: ", params_proto.version()));
+    }
+
+    const AlgorithmConfigProto &algorithm_config_proto = kv.value();
+    const AlgorithmDesc primary(algorithm_config_proto.algorithm());
+    const std::optional<AlgorithmDesc> fallback =
+        algorithm_config_proto.has_algorithm_no_scratch()
+            ? std::optional<AlgorithmDesc>(
+                  AlgorithmDesc(algorithm_config_proto.algorithm_no_scratch()))
+            : std::nullopt;
+
+    bool devices_matched = false;
+    for (int ordinal = 0; ordinal < device_descs.size(); ordinal++) {
+      const std::string &desc_str = device_descs[ordinal];
+      if (desc_str != params_proto.device_identifier()) {
+        continue;
+      }
+      devices_matched = true;
+
+      AutotuneEntry<Op> entry;
+#if TENSORFLOW_USE_ROCM
+      // ROCm doesn't yet support the OpRunner-based API, so for the time being
+      // we still need legacy AlgorithmDesc entries in the autotune map.
+      // Long-term, this should be folded into the next case.
+      entry = AutotuneEntry<Op>(AlgorithmConfig(algorithm_config_proto));
+#else
+      entry = AutotuneEntry<Op>(primary, fallback);
+#endif
+
+      autotune_map->Insert(ConvParameters(ordinal, params_proto), entry);
+    }
+
+    if (!devices_matched) {
+      unmatched_device_descs.insert(params_proto.device_identifier());
+    }
+  }
+
+  if (!unmatched_device_descs.empty()) {
+    LOG(WARNING) << "Unmatched device id's from AoT autotuning data: "
+                 << absl::StrJoin(unmatched_device_descs, ", ")
+                 << "; existing devices: " << absl::StrJoin(device_descs, ", ");
+  }
+
+  return absl::OkStatus();
+}
+
+}  // namespace
+#endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+
+absl::Status SerializeAutotuneMaps(std::string *output) {
+  AutotuneMapsProto proto;
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+  TF_ASSIGN_OR_RETURN(*proto.mutable_conv_map(),
+                      ConvMapToProto(*ConvAutotuneMap::GetInstance()));
+  TF_ASSIGN_OR_RETURN(*proto.mutable_fused_conv_map(),
+                      ConvMapToProto(*FusedConvAutotuneMap::GetInstance()));
+#endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+  TF_RET_CHECK(tsl::SerializeToStringDeterministic(proto, output));
+  return absl::OkStatus();
+}
+
+absl::Status LoadSerializedAutotuneMaps(absl::string_view s) {
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+  AutotuneMapsProto proto;
+  // The explicit string conversion here is a workaround for
+  // resolving the issue that OSS proto library's ParseFromString only accepts
+  // std::string.
+  if (!proto.ParseFromString(s)) {
+    return absl::InvalidArgumentError(
+        "Failed to parse the autotune maps from string.");
+  }
+  TF_RETURN_IF_ERROR(
+      PopulateConvMap(proto.conv_map(), ConvAutotuneMap::GetInstance()));
+  TF_RETURN_IF_ERROR(PopulateConvMap(proto.fused_conv_map(),
+                                     FusedConvAutotuneMap::GetInstance()));
+  // TODO(b/189530096): Populate autotune maps for more ops.
+#endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+  return absl::OkStatus();
+}
+
+void ResetAutotuneMaps() {
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+  ConvAutotuneMap::GetInstance()->ClearMap();
+  FusedConvAutotuneMap::GetInstance()->ClearMap();
+#endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+}
+
+}  // namespace tensorflow

tensorflow/core/util/autotune_maps/autotune_serialize.h

@@ -0,0 +1,50 @@
+/* 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.
+==============================================================================*/
+
+// For Google-internal use only.
+//
+// Supports serializing the autotune maps to string
+// (SerializeAutotuneMaps), as well as deserializing them from
+// string and injecting them into TF runtime
+// (LoadSerializedAutotuneMaps).
+//
+// Aims to speed up the warmup time of neural nets.
+
+#ifndef TENSORFLOW_CORE_UTIL_AUTOTUNE_MAPS_AUTOTUNE_SERIALIZE_H_
+#define TENSORFLOW_CORE_UTIL_AUTOTUNE_MAPS_AUTOTUNE_SERIALIZE_H_
+
+#include <string>
+
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/core/platform/status.h"
+
+namespace tensorflow {
+
+// TODO(b/189530096) Support autotune maps for more ops.
+// Loads autotune maps from string output by SerializeAutotuneMaps and uses
+// them to update the runtime autotune maps.
+absl::Status LoadSerializedAutotuneMaps(absl::string_view s);
+
+// Serializes all the autotune maps into a string that can be decoded by
+// LoadSerializedAutotuneMaps.
+absl::Status SerializeAutotuneMaps(std::string* output);
+
+// Resets all autotune maps. For test use only.
+void ResetAutotuneMaps();
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_UTIL_AUTOTUNE_MAPS_AUTOTUNE_SERIALIZE_H_

tensorflow/core/util/autotune_maps/autotune_serialize_test.cc

@@ -0,0 +1,181 @@
+/* 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.
+==============================================================================*/
+
+// For Google-internal use only.
+#include "xla/stream_executor/platform_manager.h"
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+
+#include "xla/stream_executor/gpu/gpu_init.h"
+#include "tensorflow/core/platform/status_matchers.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/util/autotune_maps/autotune_serialize.h"
+#include "tensorflow/core/util/autotune_maps/conv_autotune_maps.h"
+#include "tensorflow/core/util/autotune_maps/conv_parameters.h"
+#include "tensorflow/core/util/autotune_maps/conv_parameters.pb.h"
+#include "tensorflow/core/util/tensor_format.h"
+
+namespace tensorflow {
+namespace {
+using stream_executor::dnn::AlgorithmConfig;
+using stream_executor::dnn::AlgorithmDesc;
+using ::tensorflow::testing::StatusIs;
+using ::testing::HasSubstr;
+
+// Gets a GPU StreamExecutor instance.  Any one will do.
+se::StreamExecutor* GetStreamExec() {
+  se::Platform* platform =
+      se::PlatformManager::PlatformWithName(se::GpuPlatformName()).value();
+  CHECK_GT(platform->VisibleDeviceCount(), 0);
+  return platform->ExecutorForDevice(0).value();
+}
+
+// Tests when there is no entry in the autotune maps.
+TEST(AutotuneSerializeTest, Empty) {
+  ResetAutotuneMaps();
+  std::string output;
+  TF_CHECK_OK(SerializeAutotuneMaps(&output));
+  TF_CHECK_OK(LoadSerializedAutotuneMaps(output));
+  EXPECT_EQ(ConvAutotuneMap::GetInstance()->GetMap().size(), 0);
+}
+
+// Tests the consistency of SerializeAutotuneMaps and LoadSerializedAutotuneMaps
+// by:
+// 1. Insert predefined entries into the autotune maps.
+// 2. Serialize it to string using SerializeAutotuneMaps.
+// 3. Reset autotune maps.
+// 4. Use MergeFromstring to load the entries from string to autotune maps.
+// 5. Check if entries in autotune maps are equal to the predefined ones.
+TEST(AutotuneSerializeTest, Consistency) {
+  ResetAutotuneMaps();
+  ConvParameters conv_params_example_a = {
+      GetStreamExec(),
+      /*batch=*/1,
+      /*in_depths=*/1,
+      /*in=*/{{1, 1}},
+      /*data_format=*/TensorFormat::FORMAT_NCHW,
+      /*out_depths=*/1,
+      /*filter=*/{{1, 1}},
+      /*dilation=*/{{1, 1}},
+      /*stride=*/{{1, 1}},
+      /*padding=*/{{1, 1}},
+      /*dtype=*/DataType::DT_INT8,
+      /*group_count=*/1};
+  ConvParameters fused_params_example_a = {
+      GetStreamExec(),
+      /*batch=*/1,
+      /*in_depths=*/1,
+      /*in=*/{{1, 1}},
+      /*data_format=*/TensorFormat::FORMAT_NCHW,
+      /*out_depths=*/1,
+      /*filter=*/{{1, 1}},
+      /*dilation=*/{{1, 1}},
+      /*stride=*/{{1, 1}},
+      /*padding=*/{{1, 1}},
+      /*dtype=*/DataType::DT_INT8,
+      /*group_count=*/1,
+      ConvParameters::FusionInfo{1.0, 0., 0.,
+                                 /*activation_mode=*/
+                                 se::dnn::ActivationMode::kNone,
+                                 /*is_contrib=*/false},
+  };
+  ConvParameters contrib_fused_params_example_a = {
+      GetStreamExec(),
+      /*batch=*/1,
+      /*in_depths=*/1,
+      /*in=*/{{1, 1}},
+      /*data_format=*/TensorFormat::FORMAT_NCHW,
+      /*out_depths=*/1,
+      /*filter=*/{{1, 1}},
+      /*dilation=*/{{1, 1}},
+      /*stride=*/{{1, 1}},
+      /*padding=*/{{1, 1}},
+      /*dtype=*/DataType::DT_INT8,
+      /*group_count=*/1,
+      ConvParameters::FusionInfo{1.0, 0., 0.,
+                                 /*activation_mode=*/
+                                 se::dnn::ActivationMode::kRelu,
+                                 /*is_contrib=*/true}};
+
+  AlgorithmDesc algorithm(/*algo_id=*/1, /*use_tensor_ops=*/true);
+  AlgorithmDesc algorithm_no_scratch(/*algo_id=*/1, /*use_tensor_ops=*/true);
+  AutotuneEntry<se::dnn::ConvOp> example_a(algorithm, algorithm_no_scratch);
+  ConvAutotuneMap::GetInstance()->Insert(conv_params_example_a, example_a);
+  ConvAutotuneMap::GetInstance()->Insert(fused_params_example_a, example_a);
+  ConvAutotuneMap::GetInstance()->Insert(contrib_fused_params_example_a,
+                                         example_a);
+  std::string serialized_string;
+  TF_CHECK_OK(SerializeAutotuneMaps(&serialized_string));
+  ResetAutotuneMaps();
+  TF_CHECK_OK(LoadSerializedAutotuneMaps(serialized_string));
+  EXPECT_EQ(ConvAutotuneMap::GetInstance()->GetMap().size(), 3);
+
+  AutotuneEntry<se::dnn::ConvOp> entry;
+  EXPECT_TRUE(
+      ConvAutotuneMap::GetInstance()->Find(conv_params_example_a, &entry));
+  EXPECT_EQ(entry, example_a);
+  EXPECT_TRUE(
+      ConvAutotuneMap::GetInstance()->Find(fused_params_example_a, &entry));
+  EXPECT_EQ(entry, example_a);
+  EXPECT_TRUE(ConvAutotuneMap::GetInstance()->Find(
+      contrib_fused_params_example_a, &entry));
+  EXPECT_EQ(entry, example_a);
+}
+
+// Test that LoadSerializedAutotuneMaps will reject entries with incompatible
+// version.
+TEST(AutotuneSerializeTest, VersionControl) {
+  ResetAutotuneMaps();
+
+  ConvParameters fused_params_example_a = {
+      GetStreamExec(),
+      /*batch=*/1,
+      /*in_depths=*/1,
+      /*in=*/{{1, 1}},
+      /*data_format=*/TensorFormat::FORMAT_NCHW,
+      /*out_depths=*/1,
+      /*filter=*/{{1, 1}},
+      /*dilation=*/{{1, 1}},
+      /*stride=*/{{1, 1}},
+      /*padding=*/{{1, 1}},
+      /*dtype=*/DataType::DT_INT8,
+      /*group_count=*/1,
+      ConvParameters::FusionInfo{1.0, 0., 0.,
+                                 /*activation_mode=*/
+                                 se::dnn::ActivationMode::kNone,
+                                 /*is_contrib=*/false},
+      /*version=*/ConvParameters::kVersion - 1};
+
+  AlgorithmDesc algorithm(/*algo_id=*/1, /*use_tensor_ops=*/true);
+  AlgorithmDesc algorithm_no_scratch(/*algo_id=*/1, /*use_tensor_ops=*/true);
+  AlgorithmConfig algorithm_config_example_a(algorithm, /*scratch_size=*/1,
+                                             algorithm_no_scratch);
+
+  ConvAutotuneMap::GetInstance()->Insert(
+      fused_params_example_a,
+      AutotuneEntry<se::dnn::ConvOp>(algorithm_config_example_a));
+
+  std::string serialized_string;
+  TF_CHECK_OK(SerializeAutotuneMaps(&serialized_string));
+
+  ResetAutotuneMaps();
+  EXPECT_THAT(LoadSerializedAutotuneMaps(serialized_string),
+              absl_testing::StatusIs(
+                  error::ABORTED,
+                  HasSubstr("Aborted because the loaded autotune results")));
+  EXPECT_EQ(ConvAutotuneMap::GetInstance()->GetMap().size(), 0);
+}
+}  // namespace
+}  // namespace tensorflow
+#endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM

tensorflow/core/util/autotune_maps/conv_autotune_maps.h

@@ -0,0 +1,60 @@
+/* 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.
+==============================================================================*/
+
+// For Google-internal use only.
+//
+// This file defines the map data structure for storing autotuning results for
+// fused_conv2d_bias_activation_op_kernels.
+//
+// The key of the map uniquely identifies a convolution operation that runs on a
+// particular device model while the value might be the autotuned algorithm we
+// choose for the conv.
+//
+// This map will be merged after fused_conv2d_bias_activation_op_kernels is
+// merged into conv_ops_fused_impl.h (b/177365158, b/189530096)
+
+#ifndef TENSORFLOW_CORE_UTIL_AUTOTUNE_MAPS_CONV_AUTOTUNE_MAPS_H_
+#define TENSORFLOW_CORE_UTIL_AUTOTUNE_MAPS_CONV_AUTOTUNE_MAPS_H_
+
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+#include <string>
+
+#include "tensorflow/core/kernels/gpu_utils.h"
+#include "tensorflow/core/platform/stream_executor.h"
+#include "tensorflow/core/util/autotune_maps/conv_parameters.h"
+
+namespace tensorflow {
+
+// A dummy type to group forward convolution autotune results together.
+struct ConvAutotuneGroup {
+  static std::string name() { return "Conv"; }
+};
+
+using ConvAutotuneMap = AutotuneSingleton<ConvAutotuneGroup, ConvParameters,
+                                          AutotuneEntry<se::dnn::ConvOp>>;
+
+// A dummy type to group fused convolution autotune results together.
+struct ConvFusedAutotuneGroup {
+  static std::string name() { return "FusedConv"; }
+};
+
+using FusedConvAutotuneMap =
+    AutotuneSingleton<ConvAutotuneGroup, ConvParameters,
+                      AutotuneEntry<se::dnn::FusedConvOp>>;
+
+}  // namespace tensorflow
+#endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+
+#endif  // TENSORFLOW_CORE_UTIL_AUTOTUNE_MAPS_CONV_AUTOTUNE_MAPS_H_

tensorflow/core/util/autotune_maps/conv_map_wrapper.cc

@@ -0,0 +1,84 @@
+/* Copyright 2024 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/util/autotune_maps/conv_map_wrapper.h"
+
+#include <vector>
+
+#include "absl/log/check.h"
+#include "absl/status/status.h"
+#include "xla/tsl/lib/strings/proto_serialization.h"
+#include "xla/tsl/protobuf/dnn.pb.h"
+#include "tensorflow/core/util/autotune_maps/autotune_map.pb.h"
+#include "tensorflow/core/util/autotune_maps/conv_parameters.pb.h"
+
+namespace tensorflow {
+
+/*static*/ absl::StatusOr<ConvMapWrapper> ConvMapWrapper::FromKeyAndValue(
+    OpaqueKey key, OpaqueValue value) {
+  ConvMapProto::Entry key_proto;
+  if (!key_proto.ParseFromString(key)) {
+    return absl::Status(absl::StatusCode::kInvalidArgument,
+                        "Could not parse the provided key");
+  }
+
+  ConvMapProto::Entry value_proto;
+  if (!value_proto.ParseFromString(value)) {
+    return absl::Status(absl::StatusCode::kInvalidArgument,
+                        "Could not parse the provided value");
+  }
+
+  ConvMapProto::Entry full_entry;
+  *full_entry.mutable_key() = key_proto.key();
+  *full_entry.mutable_value() = value_proto.value();
+  return ConvMapWrapper(full_entry);
+}
+
+ConvMapWrapper::OpaqueKey ConvMapWrapper::Key() const {
+  ConvMapProto::Entry entry;
+  *entry.mutable_key() = conv_map_entry_.key();
+  OpaqueKey serialized;
+  CHECK(tsl::SerializeToStringDeterministic(entry, &serialized));  // Crash OK
+  return serialized;
+}
+
+ConvMapWrapper::OpaqueValue ConvMapWrapper::Value() const {
+  ConvMapProto::Entry entry;
+  *entry.mutable_value() = conv_map_entry_.value();
+  OpaqueValue serialized;
+  CHECK(tsl::SerializeToStringDeterministic(entry, &serialized));  // Crash OK
+  return serialized;
+}
+
+/*static*/ std::vector<ConvMapWrapper> ConvMapWrapper::ConvMapToWrappers(
+    const ConvMapProto& autotune_results) {
+  std::vector<ConvMapWrapper> wrappers;
+  wrappers.reserve(autotune_results.kv_pairs_size());
+  for (const auto& entry : autotune_results.kv_pairs()) {
+    wrappers.push_back(ConvMapWrapper(entry));
+  }
+  return wrappers;
+}
+
+/*static*/ absl::StatusOr<ConvMapProto> ConvMapWrapper::ConvMapFromWrappers(
+    const std::vector<ConvMapWrapper>& wrappers) {
+  ConvMapProto conv_map_proto;
+  for (const auto& wrapper : wrappers) {
+    *conv_map_proto.add_kv_pairs() = wrapper.conv_map_entry_;
+  }
+  return conv_map_proto;
+}
+
+}  // namespace tensorflow

tensorflow/core/util/autotune_maps/conv_map_wrapper.h

@@ -0,0 +1,66 @@
+/* Copyright 2024 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_UTIL_AUTOTUNE_MAPS_CONV_MAP_WRAPPER_H_
+#define TENSORFLOW_CORE_UTIL_AUTOTUNE_MAPS_CONV_MAP_WRAPPER_H_
+
+#include <string>
+#include <vector>
+
+#include "absl/status/statusor.h"
+#include "tensorflow/core/util/autotune_maps/autotune_map.pb.h"
+
+namespace tensorflow {
+
+// This class is a thin wrapper around `ConvMapProto::Entry`. It is used to
+// provide opaque accessors to an entry's key and value without exposing the
+// internal structure of the entry.
+class ConvMapWrapper {
+ public:
+  using OpaqueKey = std::string;
+  using OpaqueValue = std::string;
+
+  // Creates an `ConvMapWrapper` from a key and value. The provided key and
+  // value must be ones that were previously returned by calls to `Key()` and
+  // `Value()`.
+  static absl::StatusOr<ConvMapWrapper> FromKeyAndValue(OpaqueKey key,
+                                                        OpaqueValue value);
+
+  // An opaque string that can be used as a key for this autotuning result.
+  // Do not rely on the format of this string.
+  OpaqueKey Key() const;
+
+  // An opaque string that encodes the autotuning result.
+  // Do not rely on the format of this string.
+  OpaqueValue Value() const;
+
+  static std::vector<ConvMapWrapper> ConvMapToWrappers(
+      const ConvMapProto& autotune_results);
+
+  // Returns the `ConvMapProto` proto that corresponds to the provided
+  // wrappers.
+  static absl::StatusOr<ConvMapProto> ConvMapFromWrappers(
+      const std::vector<ConvMapWrapper>& wrappers);
+
+ private:
+  explicit ConvMapWrapper(const ConvMapProto::Entry& entry)
+      : conv_map_entry_(entry) {}
+
+  ConvMapProto::Entry conv_map_entry_;
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_UTIL_AUTOTUNE_MAPS_CONV_MAP_WRAPPER_H_

tensorflow/core/util/autotune_maps/conv_map_wrapper_test.cc

@@ -0,0 +1,114 @@
+/* Copyright 2024 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/util/autotune_maps/conv_map_wrapper.h"
+
+#include <utility>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "xla/tsl/platform/statusor.h"
+#include "xla/tsl/protobuf/dnn.pb.h"
+#include "tensorflow/core/util/autotune_maps/autotune_map.pb.h"
+#include "tensorflow/core/util/autotune_maps/conv_parameters.pb.h"
+
+namespace tensorflow {
+namespace {
+
+ConvMapProto ThreeConvMapEntries() {
+  ConvMapProto proto;
+
+  auto r1 = proto.add_kv_pairs();
+  r1->mutable_key()->set_batch(1);
+  r1->mutable_key()->set_in_depths(2);
+  r1->mutable_key()->set_out_depths(3);
+  r1->mutable_value()->mutable_algorithm()->set_algo_id(4);
+
+  auto r2 = proto.add_kv_pairs();
+  r2->mutable_key()->set_batch(5);
+  r2->mutable_key()->set_in_depths(6);
+  r2->mutable_key()->set_out_depths(7);
+  r2->mutable_value()->mutable_algorithm()->set_algo_id(8);
+
+  auto r3 = proto.add_kv_pairs();
+  r3->mutable_key()->set_batch(9);
+  r3->mutable_key()->set_in_depths(10);
+  r3->mutable_key()->set_out_depths(11);
+  r3->mutable_value()->mutable_algorithm()->set_algo_id(12);
+
+  return proto;
+}
+
+TEST(ConvMapWrapperTest, FullRoundTrip) {
+  std::vector<ConvMapWrapper> wrappers =
+      ConvMapWrapper::ConvMapToWrappers(ThreeConvMapEntries());
+
+  std::vector<std::pair<ConvMapWrapper::OpaqueKey, ConvMapWrapper::OpaqueValue>>
+      key_value_pairs;
+  for (const auto& wrapper : wrappers) {
+    key_value_pairs.emplace_back(wrapper.Key(), wrapper.Value());
+  }
+
+  std::vector<ConvMapWrapper> new_wrappers;
+  for (const auto& [key, value] : key_value_pairs) {
+    TF_ASSERT_OK_AND_ASSIGN(ConvMapWrapper wrapper,
+                            ConvMapWrapper::FromKeyAndValue(key, value));
+    new_wrappers.push_back(wrapper);
+  }
+
+  TF_ASSERT_OK_AND_ASSIGN(ConvMapProto round_tripped,
+                          ConvMapWrapper::ConvMapFromWrappers(new_wrappers));
+  EXPECT_EQ(round_tripped.kv_pairs_size(), 3);
+  EXPECT_EQ(round_tripped.kv_pairs(0).key().batch(), 1);
+  EXPECT_EQ(round_tripped.kv_pairs(0).key().in_depths(), 2);
+  EXPECT_EQ(round_tripped.kv_pairs(0).key().out_depths(), 3);
+  EXPECT_EQ(round_tripped.kv_pairs(0).value().algorithm().algo_id(), 4);
+  EXPECT_EQ(round_tripped.kv_pairs(1).key().batch(), 5);
+  EXPECT_EQ(round_tripped.kv_pairs(1).key().in_depths(), 6);
+  EXPECT_EQ(round_tripped.kv_pairs(1).key().out_depths(), 7);
+  EXPECT_EQ(round_tripped.kv_pairs(1).value().algorithm().algo_id(), 8);
+  EXPECT_EQ(round_tripped.kv_pairs(2).key().batch(), 9);
+  EXPECT_EQ(round_tripped.kv_pairs(2).key().in_depths(), 10);
+  EXPECT_EQ(round_tripped.kv_pairs(2).key().out_depths(), 11);
+  EXPECT_EQ(round_tripped.kv_pairs(2).value().algorithm().algo_id(), 12);
+}
+
+TEST(ConvMapWrapperTest, DeterministicSerialization) {
+  std::vector<ConvMapWrapper> wrappers =
+      ConvMapWrapper::ConvMapToWrappers(ThreeConvMapEntries());
+  std::vector<ConvMapWrapper::OpaqueKey> keys;
+  std::vector<ConvMapWrapper::OpaqueValue> values;
+  for (const auto& wrapper : wrappers) {
+    keys.push_back(wrapper.Key());
+    values.push_back(wrapper.Value());
+  }
+
+  const int kNumIterations = 100;
+  for (int i = 0; i < kNumIterations; ++i) {
+    std::vector<ConvMapWrapper> test_wrappers =
+        ConvMapWrapper::ConvMapToWrappers(ThreeConvMapEntries());
+    std::vector<ConvMapWrapper::OpaqueKey> test_keys;
+    std::vector<ConvMapWrapper::OpaqueValue> test_values;
+    for (const auto& test_wrapper : test_wrappers) {
+      test_keys.push_back(test_wrapper.Key());
+      test_values.push_back(test_wrapper.Value());
+    }
+    EXPECT_EQ(keys, test_keys);
+    EXPECT_EQ(values, test_values);
+  }
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/core/util/autotune_maps/conv_parameters.cc

@@ -0,0 +1,144 @@
+/* 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.
+==============================================================================*/
+
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+#include "tensorflow/core/util/autotune_maps/conv_parameters.h"
+
+#include <cstddef>
+#include <vector>
+
+#include "absl/strings/str_format.h"
+#include "absl/strings/str_replace.h"
+#include "xla/tsl/lib/strings/proto_serialization.h"
+#include "tensorflow/core/platform/hash.h"
+#include "tensorflow/core/util/autotune_maps/conv_parameters.pb.h"
+#include "tsl/platform/regexp.h"
+
+namespace tensorflow {
+
+namespace {
+using ::tsl::protobuf::util::MessageDifferencer;
+
+uint64_t ComputeHash(int device_id, const ConvParametersProto& proto) {
+  return Hash64Combine(device_id, tsl::DeterministicProtoHash64(proto));
+}
+
+uint64_t ComputeHash(int device_id, const MatmulParametersProto& proto) {
+  return Hash64Combine(device_id, tsl::DeterministicProtoHash64(proto));
+}
+}  // namespace
+
+std::string DeviceIdentifierForAutotuning(absl::string_view device_identifier) {
+  static const LazyRE2 kDevicePattern = {
+      R"regexp(^sm_[^ ]+ with ([\d]+)B RAM, [\d]+ cores, [\d]+KHz clock, [\d]+KHz mem clock, [\d]+B L2\$$)regexp"};
+
+  uint64_t ram_bytes;
+  if (!RE2::FullMatch(device_identifier, *kDevicePattern, &ram_bytes)) {
+    return std::string(device_identifier);
+  }
+
+  // Round up RAM to GB with 1 decimal place.
+  double ram_gb = static_cast<double>(ram_bytes) / 1e9;
+  return absl::StrReplaceAll(device_identifier,
+                             {{absl::StrCat(ram_bytes, "B RAM"),
+                               absl::StrFormat("%.1fGB RAM", ram_gb)}});
+}
+
+ConvParameters::ConvParameters(
+    se::StreamExecutor* stream_exec, int64_t batch, int64_t in_depths,
+    const absl::Span<const int64_t> in, int data_format, int64_t out_depths,
+    const absl::Span<const int64_t> filter,
+    const absl::Span<const int64_t> dilation,
+    const absl::Span<const int64_t> stride,
+    const absl::Span<const int64_t> padding, DataType dtype, int group_count,
+    std::optional<ConvParameters::FusionInfo> fusion_info, int version)
+    : device_id_(stream_exec->device_ordinal()) {
+  proto_.set_batch(batch);
+  proto_.set_in_depths(in_depths);
+  *proto_.mutable_in() = {in.begin(), in.end()};
+  proto_.set_data_format(static_cast<int>(data_format));
+  proto_.set_out_depths(out_depths);
+  *proto_.mutable_filter() = {filter.begin(), filter.end()};
+  *proto_.mutable_dilation() = {dilation.begin(), dilation.end()};
+  *proto_.mutable_stride() = {stride.begin(), stride.end()};
+  *proto_.mutable_padding() = {padding.begin(), padding.end()};
+  proto_.set_dtype(dtype);
+  proto_.set_group_count(group_count);
+  if (fusion_info.has_value()) {
+    ConvParametersProto::Fusion fusion_proto;
+    fusion_proto.set_conv_scale(fusion_info.value().conv_scale);
+    fusion_proto.set_side_input_scale(fusion_info.value().side_input_scale);
+    fusion_proto.set_activation_mode(fusion_info.value().activation_mode);
+    fusion_proto.set_is_contrib(fusion_info.value().is_contrib);
+    *proto_.mutable_fusion() = fusion_proto;
+  }
+  proto_.set_device_identifier(DeviceIdentifierForAutotuning(
+      stream_exec->GetDeviceDescription().model_str()));
+  proto_.set_version(version);
+  hash_code_ = ComputeHash(device_id_, proto_);
+}
+ConvParameters::ConvParameters(int device_id, const ConvParametersProto& proto)
+    : device_id_(device_id),
+      proto_(proto),
+      hash_code_(ComputeHash(device_id_, proto_)) {}
+
+bool ConvParameters::operator==(const ConvParameters& other) const {
+  return device_id_ == other.device_id_ &&
+         MessageDifferencer::Equals(this->proto_, other.proto_);
+}
+
+std::string ConvParameters::ToString() const { return proto_.DebugString(); }
+
+MatmulParameters::MatmulParameters(
+    se::StreamExecutor* stream_exec, DataType ab_dtype, DataType c_dtype,
+    bool trans_a, bool trans_b, uint64_t m, uint64_t n, uint64_t k, int64_t lda,
+    int64_t ldb, int64_t ldc,
+    stream_executor::dnn::ActivationMode activation_mode, int version)
+    : device_id_(stream_exec->device_ordinal()) {
+  proto_.set_ab_dtype(ab_dtype);
+  proto_.set_c_dtype(c_dtype);
+
+  proto_.set_trans_a(trans_a);
+  proto_.set_trans_b(trans_b);
+  proto_.set_m(m);
+  proto_.set_n(n);
+  proto_.set_k(k);
+  proto_.set_lda(lda);
+  proto_.set_ldb(ldb);
+  proto_.set_ldc(ldc);
+  proto_.set_activation_mode(activation_mode);
+
+  proto_.set_device_identifier(DeviceIdentifierForAutotuning(
+      stream_exec->GetDeviceDescription().model_str()));
+  proto_.set_version(version);
+  hash_code_ = ComputeHash(device_id_, proto_);
+}
+
+MatmulParameters::MatmulParameters(se::StreamExecutor* stream_exec,
+                                   const MatmulParametersProto& proto)
+    : device_id_(stream_exec->device_ordinal()),
+      proto_(proto),
+      hash_code_(ComputeHash(device_id_, proto_)) {}
+
+bool MatmulParameters::operator==(const MatmulParameters& other) const {
+  return device_id_ == other.device_id_ &&
+         MessageDifferencer::Equals(this->proto_, other.proto_);
+}
+
+std::string MatmulParameters::ToString() const { return proto_.DebugString(); }
+
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM

tensorflow/core/util/autotune_maps/conv_parameters.h

@@ -0,0 +1,148 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_UTIL_AUTOTUNE_MAPS_CONV_PARAMETERS_H_
+#define TENSORFLOW_CORE_UTIL_AUTOTUNE_MAPS_CONV_PARAMETERS_H_
+
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+#include "absl/types/optional.h"
+#include "tensorflow/core/platform/stream_executor.h"
+#include "tensorflow/core/util/autotune_maps/conv_parameters.pb.h"
+
+namespace tensorflow {
+
+// Returns a string that uniquely identifies the device model based on the
+// device identifier string from the stream executor. There are cases where
+// the same device model may have different identifiers (e.g. different RAM).
+// This function normalizes the identifier to make it comparable to other
+// autotuning results.
+std::string DeviceIdentifierForAutotuning(absl::string_view device_identifier);
+
+// Uniquely identifies a convolution operation that runs on a particular device
+// model.
+//
+// This can serve as a hashtable key, where the value might be the autotuned
+// algorithm we choose for the conv.
+//
+// All of the data in this class other than the device_id is stored in the
+// ConvParametersProto, so it can be easily serialized (for the purposes of
+// ahead-of-time autotuning).
+//
+// When using the cudnn frontend API, two autotuning results for two different
+// GPUs of the same model are not interchangeable, because an autotuning result
+// includes a cudnn execution plan, which is tied to the GPU.  As a result, we
+// need to create separate ConvParameters objects for them.
+class ConvParameters {
+ public:
+  struct FusionInfo {
+    // For some implementations (e.g. cuDNN new backend) these scales are part
+    // of the algorithm, not part of the parameters an algorithm take. They need
+    // to be used to distinguish different algorithms.
+    double conv_scale;
+    double side_input_scale;
+    double leakyrelu_alpha;
+    stream_executor::dnn::ActivationMode activation_mode;
+    bool is_contrib;
+  };
+
+  // LINT.IfChange(conv_parameters_version)
+  // A positive number that denotes the version of this class. Should be
+  // incremented everytime this class or ConvParametersProto are updated in a
+  // way that may invalidate autotune results.
+  static constexpr int kVersion = 3;
+  // LINT.ThenChange()
+
+  // We have three kinds of convolutions today.  Vanilla unfused convolutions,
+  // fused convolutions, and fused convolutions as implemented in the `contrib`
+  // directory.  The two fused convolutions ultimately correspond to the same
+  // cudnn calls, but have slightly different semantics (e.g. they interpret
+  // padding differently).
+  ConvParameters(se::StreamExecutor* stream_exec, int64_t batch,
+                 int64_t in_depths, absl::Span<const int64_t> in,
+                 int data_format, int64_t out_depths,
+                 absl::Span<const int64_t> filter,
+                 absl::Span<const int64_t> dilation,
+                 absl::Span<const int64_t> stride,
+                 absl::Span<const int64_t> padding, DataType dtype,
+                 int group_count,
+                 std::optional<ConvParameters::FusionInfo> fusion_info =
+                     std::optional<ConvParameters::FusionInfo>(),
+                 // This argument should be set only for test use.
+                 int version = kVersion);
+
+  ConvParameters(int device_id, const ConvParametersProto& proto);
+
+  ConvParameters(se::StreamExecutor* stream_exec,
+                 const ConvParametersProto& proto)
+      : ConvParameters(stream_exec->device_ordinal(), proto) {}
+
+  bool operator==(const ConvParameters& other) const;
+
+  bool operator!=(const ConvParameters& other) const {
+    return !(*this == other);
+  }
+  uint64_t hash() const { return hash_code_; }
+
+  std::string ToString() const;
+
+  const ConvParametersProto& proto() const { return proto_; }
+
+ private:
+  int device_id_;
+  ConvParametersProto proto_;
+  uint64_t hash_code_;
+};
+
+class MatmulParameters {
+ public:
+  // LINT.IfChange(matmul_parameters_version)
+  // A positive number that denotes the version of this class. Should be
+  // incremented everytime this class or ConvParametersProto are updated in a
+  // way that may invalidate autotune results.
+  static constexpr int kVersion = 2;
+  // LINT.ThenChange()
+
+  MatmulParameters(se::StreamExecutor* stream_exec, DataType ab_dtype,
+                   DataType c_dtype, bool trans_a, bool trans_b, uint64_t m,
+                   uint64_t n, uint64_t k, int64_t lda, int64_t ldb,
+                   int64_t ldc,
+                   stream_executor::dnn::ActivationMode activation_mode,
+                   // This argument should be set only for test use.
+                   int version = kVersion);
+
+  MatmulParameters(se::StreamExecutor* stream_exec,
+                   const MatmulParametersProto& proto);
+
+  bool operator==(const MatmulParameters& other) const;
+
+  bool operator!=(const MatmulParameters& other) const {
+    return !(*this == other);
+  }
+  uint64_t hash() const { return hash_code_; }
+
+  std::string ToString() const;
+
+  const MatmulParametersProto& proto() const { return proto_; }
+
+ private:
+  int device_id_;
+  MatmulParametersProto proto_;
+  uint64_t hash_code_;
+};
+
+}  // namespace tensorflow
+#endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+
+#endif  // TENSORFLOW_CORE_UTIL_AUTOTUNE_MAPS_CONV_PARAMETERS_H_

tensorflow/core/util/autotune_maps/conv_parameters.proto

@@ -0,0 +1,102 @@
+// LINT: LEGACY_NAMES
+/* 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.
+==============================================================================*/
+
+// Protocol messages for describing parameters of convolution-related
+// operations.
+
+// For Google-internal use only.
+syntax = "proto3";
+
+package tensorflow;
+
+import "xla/tsl/protobuf/dnn.proto";
+import "tensorflow/core/framework/types.proto";
+
+// LINT.IfChange
+
+// This is the underlying data structure of class ConvParameters, which are used
+// as the keys in cuDNN autotuning maps for retrieving corresponding cuDNN
+// algorithms. This is used as a serialization format for saving/loading
+// autotuning databases.
+message ConvParametersProto {
+  // This stores the information for fused convolution operations where an
+  // activation and a side input might follow the convolution.
+  message Fusion {
+    // If true, this proto corresponds to a FusedConvBiasActivation operation
+    // implemented in the contrib library, otherwise this proto corresponds to
+    // the FusedConv operation implemented in the core library.
+    // Compared with FusedConv, FusedConvBiasActivation supports more types of
+    // activation function (including no activation) as well as the side_input.
+    // For now they have same type of keys in autotune maps, but the semantics
+    // of some fields (like padding) are different. So we add this field to
+    // distinguish them.
+    // TODO(b/177365158) Remove this field once these two operations are merged.
+    bool is_contrib = 1;
+    stream_executor.dnn.ActivationMode activation_mode = 2;
+    double conv_scale = 3;
+    double side_input_scale = 4;
+  }
+  int64 batch = 1;
+  int64 in_depths = 2;
+  int64 out_depths = 3;
+  repeated int64 in = 4;
+  // data_format corresponds to type TensorFormat in
+  // third_party/tensorflow/core/util/tensor_format.h.
+  int32 data_format = 5;
+  repeated int64 filter = 6;
+  repeated int64 dilation = 7;
+  repeated int64 stride = 8;
+  repeated int64 padding = 9;
+  tensorflow.DataType dtype = 10;
+  int32 group_count = 11;
+  // A string uniquely identifying a particular GPU model, e.g. V100 vs RTX
+  // 2080.
+  string device_identifier = 12;
+  Fusion fusion = 13;
+  // The version number of ConvParameters class. Offline autotune results
+  // whose version number is different from the runtime's version number
+  // (defined in ConvParameters::kVersion) will be rejected and ignored by
+  // LoadSerializedAutotuneMaps. This ensures that we will not load out-of-date
+  // autotune results.
+  int32 version = 14;
+}
+
+//  LINT.ThenChange(
+//  "conv_parameters.h:conv_parameters_version"
+//  )
+
+// LINT.IfChange
+
+message MatmulParametersProto {
+  tensorflow.DataType ab_dtype = 1;
+  tensorflow.DataType c_dtype = 2;
+  bool trans_a = 3;
+  bool trans_b = 4;
+  uint64 m = 5;
+  uint64 n = 6;
+  uint64 k = 7;
+  int64 lda = 8;
+  int64 ldb = 9;
+  int64 ldc = 10;
+  stream_executor.dnn.ActivationMode activation_mode = 11;
+
+  string device_identifier = 12;
+  int32 version = 14;
+}
+
+//  LINT.ThenChange(
+//  "conv_parameters.h:matmul_parameters_version"
+//  )

tensorflow/core/util/autotune_maps/conv_parameters_test.cc

@@ -0,0 +1,46 @@
+/* Copyright 2025 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/util/autotune_maps/conv_parameters.h"
+
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+
+namespace {
+
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+
+using ::testing::StrEq;
+
+TEST(DeviceIdentifierForAutotuning, RoundsUpRamToGb) {
+  EXPECT_THAT(
+      DeviceIdentifierForAutotuning(
+          "sm_7.0 with 98765432100B RAM, 97 cores, 777KHz clock, 204KHz "
+          "mem clock, 888B L2$"),
+      StrEq("sm_7.0 with 98.8GB RAM, 97 cores, 777KHz clock, 204KHz mem clock, "
+            "888B L2$"));
+}
+
+TEST(DeviceIdentifierForAutotuning, NoChangeIfRegexDoesNotMatch) {
+  EXPECT_THAT(DeviceIdentifierForAutotuning("expect no change 123456B RAM"),
+              StrEq("expect no change 123456B RAM"));
+}
+
+#endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+
+}  // namespace
+
+}  // namespace tensorflow

tensorflow/core/util/bad_indices_policy.cc

@@ -0,0 +1,39 @@
+/* Copyright 2024 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/util/bad_indices_policy.h"
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+
+namespace tensorflow {
+
+constexpr char kDefault[] = "DEFAULT";
+constexpr char kErrorStr[] = "ERROR";
+constexpr char kIgnoreStr[] = "IGNORE";
+
+absl::StatusOr<BadIndicesPolicy> BadIndicesPolicyFromString(
+    absl::string_view str) {
+  if (str.empty()) return BadIndicesPolicy::kDefault;
+  if (str == kDefault) return BadIndicesPolicy::kDefault;
+  if (str == kErrorStr) return BadIndicesPolicy::kError;
+  if (str == kIgnoreStr) return BadIndicesPolicy::kIgnore;
+  return absl::InvalidArgumentError(
+      absl::StrCat("Unknown bad indices handling attribute: ", str));
+}
+
+}  // namespace tensorflow

tensorflow/core/util/bad_indices_policy.h

@@ -0,0 +1,39 @@
+/* Copyright 2024 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_UTIL_BAD_INDICES_POLICY_H_
+#define TENSORFLOW_CORE_UTIL_BAD_INDICES_POLICY_H_
+
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+
+namespace tensorflow {
+enum class BadIndicesPolicy {
+  // Default behavior: return an error on CPU and ignore on GPU. This is because
+  // we handle bad indices differently on CPU and GPU before this policy is
+  // introduced.
+  kDefault,
+  // Return an error.
+  kError,
+  // Ignore bad indices.
+  kIgnore,
+};
+
+absl::StatusOr<BadIndicesPolicy> BadIndicesPolicyFromString(
+    absl::string_view str);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_UTIL_BAD_INDICES_POLICY_H_

tensorflow/core/util/bad_indices_policy_test.cc

@@ -0,0 +1,70 @@
+/* Copyright 2024 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/util/bad_indices_policy.h"
+
+#include <gmock/gmock.h>
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+
+namespace {  // Anonymous namespace to avoid name conflicts
+
+constexpr absl::string_view kDefault = "DEFAULT";
+constexpr absl::string_view kErrorStr = "ERROR";
+constexpr absl::string_view kIgnoreStr = "IGNORE";
+
+// Unit test class using Google Test framework
+class BadIndicesPolicyFromStringTest : public ::testing::Test {
+ protected:
+  // Reusable function to test valid inputs
+  void TestValidInput(absl::string_view input, BadIndicesPolicy expected) {
+    absl::StatusOr<BadIndicesPolicy> result = BadIndicesPolicyFromString(input);
+    ASSERT_TRUE(result.ok());             // Check for success
+    EXPECT_EQ(result.value(), expected);  // Verify the policy value
+  }
+};
+
+// Test cases covering valid inputs
+TEST_F(BadIndicesPolicyFromStringTest, EmptyString) {
+  TestValidInput("", BadIndicesPolicy::kDefault);
+}
+
+TEST_F(BadIndicesPolicyFromStringTest, DefaultKeyword) {
+  TestValidInput(kDefault, BadIndicesPolicy::kDefault);
+}
+
+TEST_F(BadIndicesPolicyFromStringTest, ErrorKeyword) {
+  TestValidInput(kErrorStr, BadIndicesPolicy::kError);
+}
+
+TEST_F(BadIndicesPolicyFromStringTest, IgnoreKeyword) {
+  TestValidInput(kIgnoreStr, BadIndicesPolicy::kIgnore);
+}
+
+// Test case for invalid input
+TEST_F(BadIndicesPolicyFromStringTest, InvalidInput) {
+  absl::StatusOr<BadIndicesPolicy> result =
+      BadIndicesPolicyFromString("unknown");
+  ASSERT_FALSE(result.ok());  // Check for failure
+  EXPECT_THAT(result.status().message(),
+              ::testing::HasSubstr("Unknown bad indices handling attribute"));
+}
+
+}  // namespace
+
+}  // namespace tensorflow

tensorflow/core/util/batch_util.cc

@@ -0,0 +1,503 @@
+/* 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/util/batch_util.h"
+
+#include <algorithm>
+#include <utility>
+
+#include "tensorflow/core/framework/full_type.pb.h"
+#include "tensorflow/core/framework/register_types.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/tstring.h"
+
+#define TF_CALL_DATASET_TYPES(m) TF_CALL_ALL_TYPES(m) TF_CALL_QUANTIZED_TYPES(m)
+
+namespace tensorflow {
+namespace batch_util {
+
+namespace {
+
+absl::Status ValidateInput(const Tensor& parent, const Tensor& element,
+                           int64_t index) {
+  DCHECK_NE(parent.dim_size(0), 0);
+  DCHECK_GE(index, 0);
+  if (element.NumElements() != (parent.NumElements() / parent.dim_size(0))) {
+    TensorShape chip_shape = parent.shape();
+    chip_shape.RemoveDim(0);
+    return absl::InternalError(absl::StrCat(
+        "ValidateInput Cannot perform copy: number of elements does not match. "
+        " Shapes are: [element]: ",
+        element.shape().DebugString(),
+        ", [parent slice]: ", chip_shape.DebugString()));
+  }
+  return absl::OkStatus();
+}
+
+template <typename T>
+absl::Status HandleElementToSlice(const Tensor& /* element */, T* src, T* dest,
+                                  int64_t num_values) {
+  static_assert(tsl::is_simple_type<T>::value,
+                "Memcpy requires a simple type.");
+  memcpy(dest, src, num_values * sizeof(T));
+  return absl::OkStatus();
+}
+
+template <>
+absl::Status HandleElementToSlice<tstring>(const Tensor& element, tstring* src,
+                                           tstring* dest, int64_t num_values) {
+  if (element.RefCountIsOne()) {
+    for (int64_t i = 0; i < num_values; ++i) {
+      *dest++ = std::move(*src++);
+    }
+  } else {
+    std::copy_n(src, num_values, dest);
+  }
+  return absl::OkStatus();
+}
+
+template <>
+absl::Status HandleElementToSlice<Variant>(const Tensor& element, Variant* src,
+                                           Variant* dest, int64_t num_values) {
+  if (element.RefCountIsOne()) {
+    for (int64_t i = 0; i < num_values; ++i) {
+      *dest++ = std::move(*src++);
+    }
+  } else {
+    std::copy_n(src, num_values, dest);
+  }
+  return absl::OkStatus();
+}
+
+template <>
+absl::Status HandleElementToSlice<ResourceHandle>(const Tensor& /* element */,
+                                                  ResourceHandle* src,
+                                                  ResourceHandle* dest,
+                                                  int64_t num_values) {
+  std::copy_n(src, num_values, dest);
+  return absl::OkStatus();
+}
+
+template <>
+absl::Status HandleElementToSlice<Eigen::half>(const Tensor& /* element */,
+                                               Eigen::half* src,
+                                               Eigen::half* dest,
+                                               int64_t num_values) {
+  std::copy_n(src, num_values, dest);
+  return absl::OkStatus();
+}
+
+template <typename T>
+void HandleSliceToElement(const T* src, T* dest, int64_t num_values) {
+  static_assert(tsl::is_simple_type<T>::value,
+                "Memcpy requires a simple type.");
+  memcpy(dest, src, num_values * sizeof(T));
+}
+
+template <>
+void HandleSliceToElement<tstring>(const tstring* src, tstring* dest,
+                                   int64_t num_values) {
+  std::copy_n(src, num_values, dest);
+}
+
+template <>
+void HandleSliceToElement<Variant>(const Variant* src, Variant* dest,
+                                   int64_t num_values) {
+  std::copy_n(src, num_values, dest);
+}
+
+template <>
+void HandleSliceToElement<ResourceHandle>(const ResourceHandle* src,
+                                          ResourceHandle* dest,
+                                          int64_t num_values) {
+  std::copy_n(src, num_values, dest);
+}
+
+template <>
+void HandleSliceToElement<Eigen::half>(const Eigen::half* src,
+                                       Eigen::half* dest, int64_t num_values) {
+  std::copy_n(src, num_values, dest);
+}
+
+template <typename T>
+void HandleSliceToElement(Tensor* parent, T* src, T* dest, int64_t num_values) {
+  static_assert(tsl::is_simple_type<T>::value,
+                "Memcpy requires a simple type.");
+  memcpy(dest, src, num_values * sizeof(T));
+}
+
+template <>
+void HandleSliceToElement<tstring>(Tensor* parent, tstring* src, tstring* dest,
+                                   int64_t num_values) {
+  if (parent->RefCountIsOne()) {
+    for (int64_t i = 0; i < num_values; ++i) {
+      dest[i] = std::move(src[i]);
+    }
+  } else {
+    std::copy_n(src, num_values, dest);
+  }
+}
+
+template <>
+void HandleSliceToElement<Variant>(Tensor* parent, Variant* src, Variant* dest,
+                                   int64_t num_values) {
+  if (parent->RefCountIsOne()) {
+    for (int64_t i = 0; i < num_values; ++i) {
+      dest[i] = std::move(src[i]);
+    }
+  } else {
+    std::copy_n(src, num_values, dest);
+  }
+}
+
+template <>
+void HandleSliceToElement<ResourceHandle>(Tensor* parent, ResourceHandle* src,
+                                          ResourceHandle* dest,
+                                          int64_t num_values) {
+  std::copy_n(src, num_values, dest);
+}
+
+template <>
+void HandleSliceToElement<Eigen::half>(Tensor* parent, Eigen::half* src,
+                                       Eigen::half* dest, int64_t num_values) {
+  std::copy_n(src, num_values, dest);
+}
+
+}  // namespace
+
+// Copies element into the index^th slice of parent (in the 0th dimension).
+absl::Status CopyElementToSlice(const Tensor& element, Tensor* parent,
+                                int64_t index) {
+  TF_RETURN_IF_ERROR(ValidateInput(*parent, element, index));
+  const int64_t num_values = element.NumElements();
+#define HANDLE_TYPE(T)                                              \
+  case DataTypeToEnum<T>::value: {                                  \
+    T* src = element.base<T>();                                     \
+    T* dest = parent->base<T>() + (num_values * index);             \
+    return HandleElementToSlice<T>(element, src, dest, num_values); \
+  }
+
+  switch (element.dtype()) {
+    TF_CALL_ALL_TYPES(HANDLE_TYPE);
+    TF_CALL_QUANTIZED_TYPES(HANDLE_TYPE);
+#undef HANDLE_TYPE
+    default:
+      return absl::UnimplementedError(absl::StrCat(
+          "CopyElementToSlice Unhandled data type: ", element.dtype()));
+  }
+}
+
+// Copies the index^th slice of parent (in the 0th dimension) into element.
+absl::Status CopySliceToElement(const Tensor& parent, Tensor* element,
+                                int64_t index) {
+  TF_RETURN_IF_ERROR(ValidateInput(parent, *element, index));
+  const int64_t num_values = element->NumElements();
+
+#define HANDLE_TYPE(T)                                      \
+  case DataTypeToEnum<T>::value: {                          \
+    const T* src = parent.base<T>() + (num_values * index); \
+    T* dest = element->base<T>();                           \
+    HandleSliceToElement<T>(src, dest, num_values);         \
+    return OkStatus();                                      \
+  }
+
+  switch (parent.dtype()) {
+    TF_CALL_ALL_TYPES(HANDLE_TYPE);
+    TF_CALL_QUANTIZED_TYPES(HANDLE_TYPE);
+#undef HANDLE_TYPE
+    default:
+      return absl::UnimplementedError(absl::StrCat(
+          "CopySliceToElement Unhandled data type: ", element->dtype()));
+  }
+}
+
+// Does the same thing as `CopyContiguousSlices` except it might move
+// the underlying data from `src` to `dst` when possible.
+absl::Status MaybeMoveContiguousSlices(Tensor& src, int64_t src_offset,
+                                       int64_t dst_offset, int64_t num_slices,
+                                       Tensor* dst) {
+  if (src.dtype() != dst->dtype()) {
+    return absl::FailedPreconditionError(absl::StrCat(
+        "MaybeMoveContiguousSlices cannot perform copy: src and dst have "
+        "different "
+        "dtypes. Source dtype: ",
+        src.dtype(), " dstination dtype: ", dst->dtype(), "."));
+  }
+  if (src.dims() < 1) {
+    return absl::FailedPreconditionError(absl::StrCat(
+        "MaybeMoveContiguousSlices cannot perform copy: src has to be a tensor "
+        "with "
+        "rank >= 1. Source shape: ",
+        src.shape().DebugString()));
+  }
+  if (dst->dims() < 1) {
+    return absl::FailedPreconditionError(absl::StrCat(
+        "MaybeMoveContiguousSlices cannot perform copy: dst has to be a tensor "
+        "with rank >= 1. Dest shape: ",
+        dst->shape().DebugString()));
+  }
+
+  const int64_t src_dim0 = src.dim_size(0);
+  const int64_t dst_dim0 = dst->dim_size(0);
+  int64_t src_chip_size = 1;
+  int64_t dst_chip_size = 1;
+  for (int i = 1; i < src.dims(); ++i) {
+    src_chip_size *= src.dim_size(i);
+  }
+  for (int i = 1; i < dst->dims(); ++i) {
+    dst_chip_size *= dst->dim_size(i);
+  }
+
+  if (src_chip_size != dst_chip_size) {
+    return absl::FailedPreconditionError(absl::StrCat(
+        "MaybeMoveContiguousSlices cannot perform copy: source and dst shapes "
+        "are"
+        "not compatible. Source shape: ",
+        src.shape().DebugString(),
+        ", dst shape: ", dst->shape().DebugString()));
+  }
+  if (src_chip_size == 0 && dst_chip_size == 0) {
+    return absl::OkStatus();
+  }
+  if (src_offset < 0 || src_offset + num_slices > src_dim0 || dst_offset < 0 ||
+      dst_offset + num_slices > dst_dim0) {
+    return absl::FailedPreconditionError(absl::StrCat(
+        "MaybeMoveContiguousSlices cannot perform copy: index out of range. "
+        "src_offset: ",
+        src_offset, ", num_slices: ", num_slices, ", src_dim0: ", src_dim0,
+        ", dst_offset: ", dst_offset, ", dst_dim0: ", dst_dim0, "."));
+  }
+
+#define HANDLE_TYPE(T)                                                       \
+  case DataTypeToEnum<T>::value: {                                           \
+    T* src_p = src.base<T>() + (src_chip_size * src_offset);                 \
+    T* dst_p = dst->base<T>() + (dst_chip_size * dst_offset);                \
+    HandleSliceToElement<T>(&src, src_p, dst_p, src_chip_size * num_slices); \
+    return OkStatus();                                                       \
+  }
+
+  switch (src.dtype()) {
+    TF_CALL_ALL_TYPES(HANDLE_TYPE);
+    TF_CALL_QUANTIZED_TYPES(HANDLE_TYPE);
+#undef HANDLE_TYPE
+    default:
+      return absl::FailedPreconditionError(absl::StrCat(
+          "MaybeMoveContiguousSlices unhandled data type: ", src.dtype()));
+  }
+}
+
+absl::Status CopyContiguousSlices(const Tensor& src, int64_t src_offset,
+                                  int64_t dst_offset, int64_t num_slices,
+                                  Tensor* dst) {
+  if (src.dtype() != dst->dtype()) {
+    return absl::FailedPreconditionError(absl::StrCat(
+        "CopyContiguousSlices cannot perform copy: src and dst have different "
+        "dtypes. Source dtype: ",
+        src.dtype(), " dstination dtype: ", dst->dtype(), "."));
+  }
+  if (src.dims() < 1) {
+    return absl::FailedPreconditionError(absl::StrCat(
+        "CopyContiguousSlices cannot perform copy: src has to be a tensor with "
+        "rank >= 1. Source shape: ",
+        src.shape().DebugString()));
+  }
+
+  if (dst->dims() < 1) {
+    return absl::FailedPreconditionError(absl::StrCat(
+        "CopyContiguousSlices cannot perform copy: dst has to be a tensor "
+        "with rank >= 1. Dest shape: ",
+        dst->shape().DebugString()));
+  }
+
+  const int64_t src_dim0 = src.dim_size(0);
+  const int64_t dst_dim0 = dst->dim_size(0);
+  int64_t src_chip_size = 1;
+  int64_t dst_chip_size = 1;
+  for (int i = 1; i < src.dims(); ++i) {
+    src_chip_size *= src.dim_size(i);
+  }
+  for (int i = 1; i < dst->dims(); ++i) {
+    dst_chip_size *= dst->dim_size(i);
+  }
+
+  if (src_chip_size != dst_chip_size) {
+    return absl::FailedPreconditionError(absl::StrCat(
+        "CopyContiguousSlices cannot perform copy: source and dst shapes are"
+        "not compatible. Source shape: ",
+        src.shape().DebugString(),
+        ", dst shape: ", dst->shape().DebugString()));
+  }
+
+  if (src_chip_size == 0 && dst_chip_size == 0) {
+    return absl::OkStatus();
+  }
+
+  if (src_offset < 0 || src_offset + num_slices > src_dim0 || dst_offset < 0 ||
+      dst_offset + num_slices > dst_dim0) {
+    return absl::FailedPreconditionError(absl::StrCat(
+        "CopyContiguousSlices cannot perform copy: index out of range. "
+        "src_offset: ",
+        src_offset, ", num_slices: ", num_slices, ", src_dim0: ", src_dim0,
+        ", dst_offset: ", dst_offset, ", dst_dim0: ", dst_dim0, "."));
+  }
+
+#define HANDLE_TYPE(T)                                                 \
+  case DataTypeToEnum<T>::value: {                                     \
+    const T* src_p = src.base<T>() + (src_chip_size * src_offset);     \
+    T* dst_p = dst->base<T>() + (dst_chip_size * dst_offset);          \
+    HandleSliceToElement<T>(src_p, dst_p, src_chip_size * num_slices); \
+    return OkStatus();                                                 \
+  }
+
+  switch (src.dtype()) {
+    TF_CALL_ALL_TYPES(HANDLE_TYPE);
+    TF_CALL_QUANTIZED_TYPES(HANDLE_TYPE);
+#undef HANDLE_TYPE
+    default:
+      return absl::UnimplementedError(absl::StrCat(
+          "CopyContiguousSlices unhandled data type: ", src.dtype()));
+  }
+}
+
+// Copies the index^th slice of parent (in the 0th dimension) into element.
+//
+// NOTE(mrry): The implementation may be able to optimize the copy to a move.
+// This is particularly important for DT_STRING tensors.
+absl::Status MaybeMoveSliceToElement(Tensor* parent, Tensor* element,
+                                     int64_t index) {
+  TF_RETURN_IF_ERROR(ValidateInput(*parent, *element, index));
+  const int64_t num_values = element->NumElements();
+
+#define HANDLE_TYPE(T)                                      \
+  case DataTypeToEnum<T>::value: {                          \
+    T* src = parent->base<T>() + (num_values * index);      \
+    T* dest = element->base<T>();                           \
+    HandleSliceToElement<T>(parent, src, dest, num_values); \
+    return OkStatus();                                      \
+  }
+
+  switch (parent->dtype()) {
+    TF_CALL_ALL_TYPES(HANDLE_TYPE);
+    TF_CALL_QUANTIZED_TYPES(HANDLE_TYPE);
+#undef HANDLE_TYPE
+    default:
+      return absl::UnimplementedError(absl::StrCat(
+          "MaybeMoveSliceToElement Unhandled data type: ", element->dtype()));
+  }
+}
+
+// The following five functions are copied from padding_fifo_queue.cc.
+// TODO(mrry): Reconcile these functions with the similar methods in the
+// queue implementation.
+absl::Status ValidateElementToLargerSlice(const Tensor& element,
+                                          Tensor* parent) {
+  DCHECK_NE(parent->dim_size(0), 0);
+  if (element.NumElements() > (parent->NumElements() / parent->dim_size(0))) {
+    TensorShape chip_shape = parent->shape();
+    chip_shape.RemoveDim(0);
+    return absl::InternalError(absl::StrCat(
+        "HandleElementToLargerSlice Cannot copy slice: number of entries in "
+        "element is greater than number of elements in parent slice.  ",
+        "Shapes are: [element]: ", element.shape().DebugString(),
+        ", [parent slice]: ", chip_shape.DebugString()));
+  }
+  return absl::OkStatus();
+}
+
+template <typename T, int NDIMS>
+absl::Status HandleElementToLargerSlice(const Tensor& element, Tensor* parent,
+                                        int index) {
+  TF_RETURN_IF_ERROR(ValidateElementToLargerSlice(element, parent));
+  if (element.NumElements() == 0) {
+    return absl::OkStatus();
+  }
+  auto element_t = element.tensor<T, NDIMS>();
+  auto parent_t = parent->tensor<T, NDIMS + 1>();
+  Eigen::DSizes<Eigen::DenseIndex, NDIMS + 1> slice_indices;
+  slice_indices[0] = index;
+  Eigen::DSizes<Eigen::DenseIndex, NDIMS + 1> slice_size;
+  slice_size[0] = 1;
+  for (size_t i = 1; i < slice_size.size(); ++i) {
+    slice_size[i] = element_t.dimension(i - 1);
+  }
+  parent_t.slice(slice_indices, slice_size) = element_t.reshape(slice_size);
+  return absl::OkStatus();
+}
+
+template <int NDIMS>
+absl::Status HandleElementToLargerSliceWithRank(const Tensor& element,
+                                                Tensor* parent, int index) {
+#define HANDLE_TYPE(T)                                                   \
+  case DataTypeToEnum<T>::value: {                                       \
+    return HandleElementToLargerSlice<T, NDIMS>(element, parent, index); \
+  }
+
+  switch (element.dtype()) {
+    TF_CALL_DATASET_TYPES(HANDLE_TYPE);
+#undef HANDLE_TYPE
+    default:
+      return absl::UnimplementedError(absl::StrCat(
+          "HandleElementToLargerSliceWithRank Unhandled data type: ",
+          element.dtype()));
+  }
+}
+
+absl::Status CopyElementToLargerSlice(const Tensor& element, Tensor* parent,
+                                      int index) {
+  if (parent->dims() != element.dims() + 1) {
+    return absl::InternalError(absl::StrCat(
+        "Mismatched ranks.  Element's rank is: ", element.dims(),
+        " but element is meant to be a slice in output Tensor having rank: ",
+        parent->dims(), " (should be: ", element.dims() + 1, ")"));
+  }
+
+#define HANDLE_DIMS(NDIMS)                                                  \
+  case NDIMS: {                                                             \
+    TF_RETURN_IF_ERROR(                                                     \
+        HandleElementToLargerSliceWithRank<NDIMS>(element, parent, index)); \
+    return OkStatus();                                                      \
+  }
+
+  switch (element.dims()) {
+    HANDLE_DIMS(0);
+    HANDLE_DIMS(1);
+    HANDLE_DIMS(2);
+    HANDLE_DIMS(3);
+    HANDLE_DIMS(4);
+    HANDLE_DIMS(5);
+#undef HANDLE_DIMS
+    default:
+      return absl::UnimplementedError(absl::StrCat(
+          "CopyElementToLargerSlice Unhandled rank: ", element.dims()));
+  }
+}
+
+absl::Status SetElementZero(Tensor* element, const Tensor& padding) {
+#define HANDLE_TYPE(T)                                     \
+  if (element->dtype() == DataTypeToEnum<T>::value) {      \
+    element->flat<T>().setConstant(padding.scalar<T>()()); \
+    return OkStatus();                                     \
+  }
+  TF_CALL_DATASET_TYPES(HANDLE_TYPE);
+#undef HANDLE_TYPE
+  return absl::UnimplementedError(
+      absl::StrCat("SetElementZero Unhandled data type: ", element->dtype()));
+}
+
+}  // namespace batch_util
+}  // namespace tensorflow

tensorflow/core/util/batch_util.h

@@ -0,0 +1,75 @@
+/* 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_UTIL_BATCH_UTIL_H_
+#define TENSORFLOW_CORE_UTIL_BATCH_UTIL_H_
+
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/lib/core/status.h"
+
+namespace tensorflow {
+namespace batch_util {
+
+// Copies element into the index^th slice of parent (in the 0th dimension).
+//
+// NOTE(mrry): The `element` argument is taken by value. Use `std::move()`
+// to move the `element` argument into this function, and the implementation
+// may be able to optimize the copy to a move. This is particularly important
+// for DT_STRING tensors.
+absl::Status CopyElementToSlice(const Tensor& element, Tensor* parent,
+                                int64_t index);
+
+// Copies the index^th slice of parent (in the 0th dimension) into element.
+absl::Status CopySliceToElement(const Tensor& parent, Tensor* element,
+                                int64_t index);
+
+// Copies 'num_slices' contiguous slices from 'src' tensor starting from index
+// 'src_offset' into target tensor 'dst', and places them into slices
+// starting from 'dst_offset'.
+//
+// This function requires 'src' and 'dst' to have compatible shapes. That is it
+// requires cum_prod(src.shape[1:] == cum_prod(dst->shape[1:]). For example if
+// source is of shape [x, 2, 1] and dst is a tensor of shape [y, 1, 2], this
+// function can still proceed successfully.
+absl::Status CopyContiguousSlices(const Tensor& src, int64_t src_offset,
+                                  int64_t dst_offset, int64_t num_slices,
+                                  Tensor* dst);
+
+// Copies the index^th slice of parent (in the 0th dimension) into element.
+//
+// NOTE(mrry): The implementation may be able to optimize the copy to a move.
+// This is particularly important for DT_STRING tensors.
+absl::Status MaybeMoveSliceToElement(Tensor* parent, Tensor* element,
+                                     int64_t index);
+
+// Moves `src` Tensor's data in [src_offset, src_offset+num_slices) along
+// the first dimension if possible. Otherwise, copy them into `dst`.
+absl::Status MaybeMoveContiguousSlices(Tensor& src, int64_t src_offset,
+                                       int64_t dst_offset, int64_t num_slices,
+                                       Tensor* dst);
+
+// Zero-initializes the tensor `element` using the scalar stored in `padding`.
+// Both `element` and `padding` must have matching `dtype`.
+absl::Status SetElementZero(Tensor* element, const Tensor& padding);
+
+// Copies `element` into a (0th dimension) slice of `parent`, assuming
+// the shape of `element` is strictly not larger along any axis than a
+// slice.
+absl::Status CopyElementToLargerSlice(const Tensor& element, Tensor* parent,
+                                      int index);
+
+}  // namespace batch_util
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_UTIL_BATCH_UTIL_H_

tensorflow/core/util/bcast.cc

@@ -0,0 +1,35 @@
+/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/util/bcast.h"
+
+#include "tensorflow/core/platform/logging.h"
+namespace tensorflow {
+
+BCast::Vec BCast::FromShape(const TensorShape& shape) {
+  const int N = shape.dims();
+  BCastList::Vec ret(N);
+  for (int i = 0; i < N; ++i) {
+    ret[i] = shape.dim_size(i);
+  }
+  return ret;
+}
+
+TensorShape BCast::ToShape(const BCastList::Vec& vec) {
+  TensorShape shape(vec);
+  return shape;
+}
+
+}  // end namespace tensorflow

tensorflow/core/util/bcast.h

@@ -0,0 +1,430 @@
+/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_UTIL_BCAST_H_
+#define TENSORFLOW_CORE_UTIL_BCAST_H_
+
+#include <algorithm>
+#include <vector>
+
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/lib/gtl/inlined_vector.h"
+#include "tensorflow/core/platform/macros.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/util/overflow.h"
+
+namespace tensorflow {
+
+// Returns the mapping from the output batch indices to the corresponding
+// input's batch indices, given the input's "reshape" and "bcast" shapes as
+// returned by the BCastList helper class. The i'th element denotes the
+// (flattened) batch index of the input that must be used to compute the i'th
+// batch output.
+//
+inline void ComputeBatchIndices(
+    const int64_t output_batch_size,
+    const absl::InlinedVector<int64_t, 4UL>& reshape,
+    const absl::InlinedVector<int64_t, 4UL>& bcast,
+    std::vector<int64_t>* out_indices) {
+  // Populates the mapping in out_indices. This algorithm is identical to
+  // the following steps:
+  //  - Reshape {0, 1, ..., input_batch_size - 1} to the input shape.
+  //  - Broadcast to the output shape.
+  //  - Reshape back to a flat 1D vector.
+  out_indices->resize(output_batch_size);
+  int64_t num_output_elements = 1;
+  int64_t num_input_elements = 1;
+  for (int64_t i = reshape.size() - 1; i >= 0; --i) {
+    // Replicate the already populated mapping an additional (dim - 1) times.
+    // If we are broadcasting, just copy the existing mapping.
+    // Otherwise, add another dimension from the input shape.
+    const int64_t dim = std::max(reshape[i], bcast[i]);
+    const int64_t incr = bcast[i] > 1 ? 0 : num_input_elements;
+    for (int64_t k = 0; k < (dim - 1) * num_output_elements; ++k) {
+      (*out_indices)[num_output_elements + k] = (*out_indices)[k] + incr;
+    }
+    num_output_elements *= dim;
+    num_input_elements *= reshape[i];
+  }
+}
+
+template <int N>
+class BCastList {
+ public:
+  // A vector of int64 representing the shape of tensor. The 0-th
+  // element is the outer-most dimension and the last element is the
+  // inner-most dimension. Note that we do not use TensorShape since
+  // it's more convenient to manipulate Vec directly for this module.
+  typedef absl::InlinedVector<int64_t, 4UL> Vec;
+
+  // Constructs all helper shapes, following the aforementioned rules.
+  //
+  // If "fewer_dims_optimization" is set to true (the default), the
+  // implementation tries to reduce intermediate dimensions needed to be more
+  // efficient.  This is transparent to the caller.
+  //
+  // If false, all intermediate shapes (except for grad_{x,y}_reduce_idx()) have
+  // the same number of dimensions as the larger of the two inputs.
+  //
+  // If return_flattened_batch_indices is true, the implementation will compute
+  // for each output member of the flattened output, which batch indices of
+  // each input correspond to it. This is disabled by default.
+  explicit BCastList(const Vec (&x)[N], bool fewer_dims_optimization = true,
+                     bool return_flattened_batch_indices = false);
+  ~BCastList() = default;
+
+  // Returns true iff two operands are compatible according to the
+  // broadcasting rule.
+  bool IsValid() const { return valid_; }
+  bool IsBroadcastingRequired() const { return broadcasting_required_; }
+
+  // If and only if IsValid(), the following fields can be used in
+  // implementing a broadcasted binary tensor operation according to
+  // the broadcasting rule.
+  const Vec& reshape(int i) const { return reshape_[i]; }
+  const Vec& bcast(int i) const { return bcast_[i]; }
+  const Vec& result_shape() const { return result_; }
+  const Vec& output_shape() const { return output_; }
+  const Vec& grad_reduce_idx(int i) const { return grad_reduce_idx_[i]; }
+  int64_t output_batch_size() const { return output_batch_size_; }
+
+  // Returns the mapping from the flattened output batch indices to x's
+  // flattened batch indices. The result is a vector of length
+  // output_batch_size(). To compute the i'th batch output, a binary matmul-like
+  // operation should use the `x_batch_indices()[i]`th batch index of `x`.
+  // Note: Returns an empty vector if broadcasting is not required. Callers
+  // should only use this when IsBroadcastingRequired() returns true.
+  const std::vector<int64_t>& batch_indices(int i) const {
+    return batch_indices_[i];
+  }
+
+ protected:
+  bool valid_ = true;
+  bool broadcasting_required_ = true;
+  Vec reshape_[N];
+  Vec bcast_[N];
+  Vec result_;
+  Vec output_;
+  Vec grad_reduce_idx_[N];
+
+  int64_t output_batch_size_;
+  std::vector<int64_t> batch_indices_[N];
+
+  static void Reverse(Vec* shape) {
+    std::reverse(shape->begin(), shape->end());
+  }
+
+  BCastList(const BCastList&) = delete;
+  void operator=(const BCastList&) = delete;
+};
+
+template <int N>
+BCastList<N>::BCastList(const BCastList::Vec (&x)[N],
+                        const bool fewer_dims_optimization,
+                        const bool return_flattened_batch_indices) {
+  typedef BCastList::Vec Vec;
+
+  // Safely multiplies dimensions taking into account symbolic shapes.
+  auto mul_dims = [](int64_t dim1, int64_t dim2) -> int64_t {
+    if (dim1 == 0 || dim2 == 0) return 0;
+    int64_t res = MultiplyWithoutOverflow(dim1, dim2);
+    return res < 0 ? -1 : res;
+  };
+
+  bool all_equal = true;
+  size_t largest_rank = 0;
+  output_batch_size_ = 1;
+  for (int i = 0; i < N; ++i) {
+    if (x[i] != x[0]) {
+      all_equal = false;
+    }
+    if (x[i].size() > largest_rank) {
+      largest_rank = x[i].size();
+    }
+  }
+  if (all_equal) {
+    broadcasting_required_ = false;
+  }
+  if (all_equal && TF_PREDICT_TRUE(fewer_dims_optimization)) {
+    // Fast path for common case of identical shapes.
+    int64_t elements = 1;
+    const int rank = x[0].size();
+    output_.resize(rank);
+    for (int i = 0; i < rank; i++) {
+      const int64_t dim = x[0][i];
+      elements = mul_dims(elements, dim);
+      output_[i] = dim;
+    }
+    result_.push_back(elements);
+    output_batch_size_ = elements;
+    for (int i = 0; i < N; ++i) {
+      reshape_[i].push_back(elements);
+      bcast_[i].push_back(1);
+    }
+    // grad_reduce_ is left as empty
+    return;
+  }
+
+  // Reverse all the shapes for convenience
+  // After the reverse, 0-th is the inner-most dimension.
+  Vec copy[N];
+  for (int i = 0; i < N; ++i) {
+    copy[i] = x[i];
+    Reverse(&copy[i]);
+  }
+
+  // 1-extend and align all vectors.
+  for (int i = 0; i < N; ++i) {
+    if (copy[i].size() < largest_rank) {
+      copy[i].resize(largest_rank, 1);
+    }
+  }
+  // Going through each dimension starting from the inner-most
+  // dimension, compares dimension of x and y. They are compatible if
+  // they are equal or either is 1.
+
+  // indices of j-th component of each input.
+  bool prev_is_one[N];
+  bool current_is_one[N];
+  for (int i = 0; i < N; ++i) {
+    prev_is_one[i] = false;
+    current_is_one[i] = false;
+  }
+  bool output_dim_set = false;
+  int64_t output_dim = -1;
+  bool none_is_one = true;
+  bool set_one = false;
+  for (int j = 0; j < largest_rank; ++j) {
+    output_dim = -1;
+    output_dim_set = false;
+    none_is_one = true;
+    // Find which indices are 1.
+    for (int i = 0; i < N; ++i) {
+      // Keep track of which indices are 1.
+      if (copy[i][j] == 1) {
+        current_is_one[i] = true;
+        none_is_one = false;
+      } else {
+        current_is_one[i] = false;
+        if (!output_dim_set || copy[i][j] == output_dim) {
+          output_dim = copy[i][j];
+          output_dim_set = true;
+        } else {
+          valid_ = false;
+          return;
+        }
+      }
+    }
+    output_.push_back(output_dim_set ? output_dim : 1);
+    output_batch_size_ = mul_dims(output_batch_size_, output_.back());
+    // All dimensions are 1.
+    if (!output_dim_set) {
+      if (!TF_PREDICT_TRUE(fewer_dims_optimization)) {
+        for (int i = 0; i < N; ++i) {
+          bcast_[i].push_back(1);
+          reshape_[i].push_back(1);
+        }
+        result_.push_back(1);
+      }
+      for (int i = 0; i < N; ++i) {
+        grad_reduce_idx_[i].push_back(largest_rank - 1 - j);
+      }
+      // This will skip updating the previous state to the current one. We'll
+      // explain why this is safe below.
+      // Consider the previous state P, current state C and the next state N.
+      // In the case where N also is all ones (N == C), we'll do the same
+      // optimization here (push back one dimensions if we need to), which is
+      // safe and is expected.
+      //
+      // When N != C, we'll continue as usual. However, we might trigger the
+      // next block if N == P (because we didn't update the previous state).
+      // We trigger the next block if `fewer_dims_optimization` is true.
+      // This means that we did not modify and broadcast / reshapes in this
+      // block (we skipped updating, since the one dimensions can be ignored).
+      // In essence, we only need to check whether the previous non-one state is
+      // equal to the current non-one state.
+
+      continue;
+    } else if (TF_PREDICT_TRUE(fewer_dims_optimization) &&
+               std::equal(current_is_one, current_is_one + N, prev_is_one) &&
+               set_one) {
+      // It is a run of the same broadcasting case as last time.
+      // We can reshape the input so that fewer dimensions
+      // are involved in the intermediate computation.
+      result_.back() = mul_dims(result_.back(), output_dim);
+      for (int i = 0; i < N; ++i) {
+        reshape_[i].back() = mul_dims(reshape_[i].back(), copy[i][j]);
+        bcast_[i].back() =
+            mul_dims(bcast_[i].back(), current_is_one[i] ? output_dim : 1);
+        if (current_is_one[i] && !none_is_one) {
+          grad_reduce_idx_[i].push_back(largest_rank - 1 - j);
+        }
+      }
+    } else {
+      result_.push_back(output_dim);
+      for (int i = 0; i < N; ++i) {
+        reshape_[i].push_back(copy[i][j]);
+        bcast_[i].push_back(current_is_one[i] ? output_dim : 1);
+        if (current_is_one[i] && !none_is_one) {
+          grad_reduce_idx_[i].push_back(largest_rank - 1 - j);
+        }
+      }
+    }
+    set_one = true;
+    for (int i = 0; i < N; ++i) {
+      prev_is_one[i] = current_is_one[i];
+    }
+  }
+  if (result_.empty()) {
+    result_.push_back(1);
+    for (int i = 0; i < N; ++i) {
+      reshape_[i].push_back(1);
+      bcast_[i].push_back(1);
+    }
+  }
+  // Do something about batches.
+  for (int i = 0; i < N; ++i) {
+    Reverse(&reshape_[i]);
+    Reverse(&bcast_[i]);
+    Reverse(&grad_reduce_idx_[i]);
+  }
+  Reverse(&result_);
+  Reverse(&output_);
+  // Only compute batch indices when we need broadcasting, and we aren't doing
+  // needless work (when the output size is 0 or the
+  // return_flattened_batch_indices isn't enabled).
+  if (return_flattened_batch_indices && broadcasting_required_ &&
+      output_batch_size_ > 0) {
+    for (int i = 0; i < N; ++i) {
+      ComputeBatchIndices(output_batch_size_, reshape_[i], bcast_[i],
+                          &batch_indices_[i]);
+    }
+  }
+}
+
+// BCast is a helper for broadcasting binary tensor operation.
+// TensorFlow's broadcasting rule follows that of numpy (See
+// http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html).
+//
+// The rule has the following properties:
+//
+//   1. suffix matching: the rule starts with the right-most
+//      dimension, and works towards the left-most dimension. Since
+//      TensorFlow is row-major, the right-most dimension (the last
+//      element in the shape of a tensor) is the inner-most, a.k.a.
+//      the fastest changing, dimension.
+//
+//   2. Two dimensions are compatible for broadcasting if both are the
+//      same or either is 1.
+//
+// BCast takes the shape of two tensors and computes a few vectors of
+// int32 that are useful for the caller to reshape the tensors, apply
+// the right broadcasts to them, compute the broadcasted operation,
+// and possibly the gradients. In a nutshell, the caller is expected
+// to compute the broadcasted operation as following:
+//
+//   BCast b(x.shape(), y.shape());
+//   output = x.reshape(b.x_reshape()).broadcast(b.x_bcast())
+//            _op_
+//            y.reshape(b.y_reshape()).broadcast(b.y_bcast())
+//
+// For the gradient computation,
+//   grad_x = sum(grad * backprop_x(x, y), grad_x_reduce_idx)
+//            .reshape(x.shape())
+//   grad_y = sum(grad * backprop_y(x, y), grad_y_reduce_idx)
+//            .reshape(y.shape())
+// backprop_x and backprop_y are functionals of the binary function "op",
+// e.g.,
+//   for +, backprop_x(x, y) = backprop_y(x, y) = 1;
+//   for *, backprop_x(x, y) =  y, backprop_y(x, y) = x;
+//   for /, backprop_x(x, y) = 1/y, backprop_y(x, y) = -x/y^2;
+//
+// The multiplication in the grad * backprop_x itself is also
+// broadcasting following the same rule.
+class BCast : public BCastList<2> {
+ public:
+  // Constructs all helper shapes, following the aforementioned rules.
+  //
+  // If "fewer_dims_optimization" is set to true (the default), the
+  // implementation tries to reduce intermediate dimensions needed to be more
+  // efficient.  This is transparent to the caller.
+  //
+  // If false, all intermediate shapes (except for grad_{x,y}_reduce_idx()) have
+  // the same number of dimensions as the larger of the two inputs.
+  typedef absl::InlinedVector<int64_t, 4UL> Vec;
+
+  BCast(const Vec& x, const Vec& y, const bool fewer_dims_optimization = true,
+        const bool return_flattened_batch_indices = false)
+      : BCastList<2>({x, y}, fewer_dims_optimization,
+                     return_flattened_batch_indices) {}
+
+  ~BCast() = default;
+
+  // If and only if IsValid(), the following fields can be used in
+  // implementing a broadcasted binary tensor operation according to
+  // the broadcasting rule.
+  const Vec& x_reshape() const { return reshape_[0]; }
+  const Vec& x_bcast() const { return bcast_[0]; }
+  const Vec& y_reshape() const { return reshape_[1]; }
+  const Vec& y_bcast() const { return bcast_[1]; }
+  const Vec& result_shape() const { return result_; }
+  const Vec& output_shape() const { return output_; }
+  const Vec& grad_x_reduce_idx() const { return grad_reduce_idx_[0]; }
+  const Vec& grad_y_reduce_idx() const { return grad_reduce_idx_[1]; }
+
+  // Returns the mapping from the flattened output batch indices to x's
+  // flattened batch indices. The result is a vector of length
+  // output_batch_size(). To compute the i'th batch output, a binary matmul-like
+  // operation should use the `x_batch_indices()[i]`th batch index of `x`.
+  // Note: Returns an empty vector if broadcasting is not required. Callers
+  // should only use this when IsBroadcastingRequired() returns true.
+  const std::vector<int64_t>& x_batch_indices() const {
+    return batch_indices_[0];
+  }
+  // Returns the mapping from the flattened output batch indices to y's
+  // flattened batch indices. Similar to x_batch_indices().
+  // Note: Returns an empty vector if broadcasting is not required. Callers
+  // should only use this when IsBroadcastingRequired() returns true.
+  const std::vector<int64_t>& y_batch_indices() const {
+    return batch_indices_[1];
+  }
+
+  template <typename IndexType, int NDIMS>
+  static Eigen::array<IndexType, NDIMS> ToIndexArrayType(
+      const BCast::Vec& vec) {
+    CHECK_EQ(vec.size(), NDIMS);
+    Eigen::array<IndexType, NDIMS> ret;
+    for (int i = 0; i < NDIMS; ++i) ret[i] = vec[i];
+    return ret;
+  }
+
+  template <int NDIMS>
+  static Eigen::array<Eigen::DenseIndex, NDIMS> ToIndexArray(
+      const BCast::Vec& vec) {
+    return ToIndexArrayType<Eigen::DenseIndex, NDIMS>(vec);
+  }
+
+  // Static helpers.
+  static Vec FromShape(const TensorShape& shape);
+  static TensorShape ToShape(const Vec& vec);
+
+ private:
+  BCast(const BCast&) = delete;
+  void operator=(const BCast&) = delete;
+};
+
+}  // end namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_UTIL_BCAST_H_

tensorflow/core/util/bcast_test.cc

@@ -0,0 +1,727 @@
+/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/util/bcast.h"
+
+#include <cstdint>
+
+#include "tensorflow/core/lib/strings/str_util.h"
+#include "tensorflow/core/lib/strings/strcat.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/platform/test_benchmark.h"
+
+namespace tensorflow {
+namespace {
+
+std::string BCast(const tensorflow::BCast::Vec& x,
+                  const tensorflow::BCast::Vec& y,
+                  const bool fewer_dims_optimization = true) {
+  tensorflow::BCast b(x, y, fewer_dims_optimization);
+  if (!b.IsValid()) {
+    return "invalid";
+  }
+  std::string ret;
+  absl::StrAppend(&ret, "[", absl::StrJoin(b.x_reshape(), ","), "]");
+  absl::StrAppend(&ret, "[", absl::StrJoin(b.x_bcast(), ","), "]");
+  absl::StrAppend(&ret, "[", absl::StrJoin(b.y_reshape(), ","), "]");
+  absl::StrAppend(&ret, "[", absl::StrJoin(b.y_bcast(), ","), "]");
+  absl::StrAppend(&ret, "[", absl::StrJoin(b.result_shape(), ","), "]");
+  absl::StrAppend(&ret, "[", absl::StrJoin(b.output_shape(), ","), "]");
+  absl::StrAppend(&ret, "[", absl::StrJoin(b.grad_x_reduce_idx(), ","), "]");
+  absl::StrAppend(&ret, "[", absl::StrJoin(b.grad_y_reduce_idx(), ","), "]");
+  return ret;
+}
+
+std::string BCastBatchIndices(const tensorflow::BCast::Vec& x,
+                              const tensorflow::BCast::Vec& y,
+                              const bool fewer_dims_optimization = true) {
+  tensorflow::BCast b(x, y, fewer_dims_optimization,
+                      /*return_flattened_batch_indices=*/true);
+  std::string ret;
+  absl::StrAppend(&ret, "[", absl::StrJoin(b.x_batch_indices(), ","), "]");
+  absl::StrAppend(&ret, "[", absl::StrJoin(b.y_batch_indices(), ","), "]");
+  return ret;
+}
+
+std::string BCastList3(const tensorflow::BCast::Vec& x,
+                       const tensorflow::BCast::Vec& y,
+                       const tensorflow::BCast::Vec& z,
+                       const bool fewer_dims_optimization = true) {
+  tensorflow::BCastList<3> b({x, y, z}, fewer_dims_optimization);
+  if (!b.IsValid()) {
+    return "invalid";
+  }
+  std::string ret;
+  absl::StrAppend(&ret, "[", absl::StrJoin(b.reshape(0), ","), "]");
+  absl::StrAppend(&ret, "[", absl::StrJoin(b.bcast(0), ","), "]");
+  absl::StrAppend(&ret, "[", absl::StrJoin(b.reshape(1), ","), "]");
+  absl::StrAppend(&ret, "[", absl::StrJoin(b.bcast(1), ","), "]");
+  absl::StrAppend(&ret, "[", absl::StrJoin(b.reshape(2), ","), "]");
+  absl::StrAppend(&ret, "[", absl::StrJoin(b.bcast(2), ","), "]");
+  absl::StrAppend(&ret, "[", absl::StrJoin(b.result_shape(), ","), "]");
+  absl::StrAppend(&ret, "[", absl::StrJoin(b.output_shape(), ","), "]");
+  absl::StrAppend(&ret, "[", absl::StrJoin(b.grad_reduce_idx(0), ","), "]");
+  absl::StrAppend(&ret, "[", absl::StrJoin(b.grad_reduce_idx(1), ","), "]");
+  absl::StrAppend(&ret, "[", absl::StrJoin(b.grad_reduce_idx(2), ","), "]");
+  return ret;
+}
+
+TEST(BCastTest, Invalid) {
+  for (const bool use_optimization : {true, false}) {
+    EXPECT_EQ("invalid", BCast({5, 3, 2}, {3}, use_optimization));
+    EXPECT_EQ("invalid", BCast({5, 3, 2}, {2, 2}, use_optimization));
+    EXPECT_EQ("invalid", BCast({5, 3, 2}, {10, 1, 1}, use_optimization));
+    EXPECT_EQ("invalid",
+              BCast({1, 2, 1, 2, 1, 2}, {2, 4, 2, 1, 2, 1}, use_optimization));
+  }
+}
+
+TEST(BCastListTest, Invalid) {
+  for (const bool use_optimization : {true, false}) {
+    EXPECT_EQ("invalid", BCastList3({5, 3, 2}, {3}, {1}, use_optimization));
+    EXPECT_EQ("invalid", BCastList3({5, 3, 2}, {2, 2}, {1}, use_optimization));
+    EXPECT_EQ("invalid",
+              BCastList3({5, 3, 2}, {10, 1, 1}, {1}, use_optimization));
+    EXPECT_EQ("invalid", BCastList3({1, 2, 1, 2, 1, 2}, {2, 4, 2, 1, 2, 1}, {1},
+                                    use_optimization));
+    EXPECT_EQ("invalid", BCastList3({5, 3, 2}, {1}, {3}, use_optimization));
+    EXPECT_EQ("invalid", BCastList3({5, 3, 2}, {1}, {2, 2}, use_optimization));
+    EXPECT_EQ("invalid",
+              BCastList3({5, 3, 2}, {1}, {10, 1, 1}, use_optimization));
+
+    EXPECT_EQ("invalid", BCastList3({1}, {5, 3, 2}, {3}, use_optimization));
+    EXPECT_EQ("invalid", BCastList3({1}, {5, 3, 2}, {2, 2}, use_optimization));
+    EXPECT_EQ("invalid",
+              BCastList3({1}, {5, 3, 2}, {10, 1, 1}, use_optimization));
+  }
+}
+
+TEST(BCastTest, Basic_SameShape) {
+  // Effectively no broadcast needed.
+  EXPECT_EQ(BCast({11, 7, 5, 3, 2}, {11, 7, 5, 3, 2}),
+            "[2310][1][2310][1]"
+            "[2310]"
+            "[11,7,5,3,2]"
+            "[][]");
+
+  EXPECT_EQ(BCast({11, 7, 5, 3, 2}, {11, 7, 5, 3, 2}, false),
+            "[11,7,5,3,2][1,1,1,1,1][11,7,5,3,2][1,1,1,1,1]"
+            "[11,7,5,3,2]"
+            "[11,7,5,3,2]"
+            "[][]");
+}
+
+TEST(BCastListTest, Basic_SameShape) {
+  // Effectively no broadcast needed.
+  EXPECT_EQ(BCastList3({11, 7, 5, 3, 2}, {11, 7, 5, 3, 2}, {11, 7, 5, 3, 2}),
+            "[2310][1][2310][1][2310][1]"
+            "[2310]"
+            "[11,7,5,3,2]"
+            "[][][]");
+
+  EXPECT_EQ(
+      BCastList3({11, 7, 5, 3, 2}, {11, 7, 5, 3, 2}, {11, 7, 5, 3, 2}, false),
+      "[11,7,5,3,2][1,1,1,1,1][11,7,5,3,2][1,1,1,1,1][11,7,5,3,2][1,1,1,1,1]"
+      "[11,7,5,3,2]"
+      "[11,7,5,3,2]"
+      "[][][]");
+}
+
+TEST(BCastTest, Basic_SameShapeWithZeroDim) {
+  // Effectively no broadcast needed.
+  EXPECT_EQ(BCast({11, 7, 0, 3, 2}, {11, 7, 0, 3, 2}),
+            "[0][1][0][1]"
+            "[0]"
+            "[11,7,0,3,2]"
+            "[][]");
+
+  EXPECT_EQ(BCast({11, 7, 0, 3, 2}, {11, 7, 0, 3, 2}, false),
+            "[11,7,0,3,2][1,1,1,1,1][11,7,0,3,2][1,1,1,1,1]"
+            "[11,7,0,3,2]"
+            "[11,7,0,3,2]"
+            "[][]");
+}
+
+TEST(BCastListTest, Basic_SameShapeWithZeroDim) {
+  // Effectively no broadcast needed.
+  EXPECT_EQ(BCastList3({11, 7, 0, 3, 2}, {11, 7, 0, 3, 2}, {11, 7, 0, 3, 2}),
+            "[0][1][0][1][0][1]"
+            "[0]"
+            "[11,7,0,3,2]"
+            "[][][]");
+
+  EXPECT_EQ(
+      BCastList3({11, 7, 0, 3, 2}, {11, 7, 0, 3, 2}, {11, 7, 0, 3, 2}, false),
+      "[11,7,0,3,2][1,1,1,1,1][11,7,0,3,2][1,1,1,1,1][11,7,0,3,2][1,1,1,1,1]"
+      "[11,7,0,3,2]"
+      "[11,7,0,3,2]"
+      "[][][]");
+}
+
+TEST(BCastTest, Basic_Scalar_Scalar) {
+  // Effectively it's a scalar and a scalar.
+  // [1, 1] [1]
+  //
+  EXPECT_EQ(BCast({1, 1}, {}),
+            "[1][1][1][1]"
+            "[1]"
+            "[1,1]"
+            "[0,1][0,1]");
+
+  EXPECT_EQ(BCast({1, 1}, {1}),
+            "[1][1][1][1]"
+            "[1]"
+            "[1,1]"
+            "[0,1][0,1]");
+
+  EXPECT_EQ(BCast({1, 1}, {1}, false),
+            "[1,1][1,1][1,1][1,1]"
+            "[1,1]"
+            "[1,1]"
+            "[0,1][0,1]");
+
+  // [1] [1, 1]
+  EXPECT_EQ(BCast({1}, {1, 1}),
+            "[1][1][1][1]"
+            "[1]"
+            "[1,1]"
+            "[0,1][0,1]");
+
+  EXPECT_EQ(BCast({1}, {1, 1}, false),
+            "[1,1][1,1][1,1][1,1]"
+            "[1,1]"
+            "[1,1]"
+            "[0,1][0,1]");
+}
+
+TEST(BCastTest, Basic_TrueScalar_Scalar) {
+  // [] []
+  EXPECT_EQ(BCast({}, {}),
+            "[1][1][1][1]"
+            "[1]"
+            "[]"
+            "[][]");
+
+  // [] [1]
+  EXPECT_EQ(BCast({}, {1}),
+            "[1][1][1][1]"
+            "[1]"
+            "[1]"
+            "[0][0]");
+
+  EXPECT_EQ(BCast({}, {1}, false),
+            "[1][1][1][1]"
+            "[1]"
+            "[1]"
+            "[0][0]");
+
+  // [] [1, 1]
+  EXPECT_EQ(BCast({}, {1, 1}),
+            "[1][1][1][1]"
+            "[1]"
+            "[1,1]"
+            "[0,1][0,1]");
+
+  EXPECT_EQ(BCast({}, {1, 1}, false),
+            "[1,1][1,1][1,1][1,1]"
+            "[1,1]"
+            "[1,1]"
+            "[0,1][0,1]");
+
+  // [1] []
+  EXPECT_EQ(BCast({1}, {}),
+            "[1][1][1][1]"
+            "[1]"
+            "[1]"
+            "[0][0]");
+
+  EXPECT_EQ(BCast({1}, {}, false),
+            "[1][1][1][1]"
+            "[1]"
+            "[1]"
+            "[0][0]");
+
+  // [1, 1] []
+  EXPECT_EQ(BCast({1, 1}, {}),
+            "[1][1][1][1]"
+            "[1]"
+            "[1,1]"
+            "[0,1][0,1]");
+
+  EXPECT_EQ(BCast({1, 1}, {}, false),
+            "[1,1][1,1][1,1][1,1]"
+            "[1,1]"
+            "[1,1]"
+            "[0,1][0,1]");
+}
+
+TEST(BCastListTest, Basic_Scalar_Scalar_Scalar) {
+  // Effectively it's a scalar and a scalar.
+  // [1, 1] [1] [1]
+  EXPECT_EQ(BCastList3({1, 1}, {1}, {1}),
+            "[1][1][1][1][1][1]"
+            "[1]"
+            "[1,1]"
+            "[0,1][0,1][0,1]");
+
+  EXPECT_EQ(BCastList3({1, 1}, {1}, {1}, false),
+            "[1,1][1,1][1,1][1,1][1,1][1,1]"
+            "[1,1]"
+            "[1,1]"
+            "[0,1][0,1][0,1]");
+
+  // [1] [1, 1] [1]
+  EXPECT_EQ(BCastList3({1}, {1, 1}, {1}),
+            "[1][1][1][1][1][1]"
+            "[1]"
+            "[1,1]"
+            "[0,1][0,1][0,1]");
+
+  EXPECT_EQ(BCastList3({1}, {1, 1}, {1}, false),
+            "[1,1][1,1][1,1][1,1][1,1][1,1]"
+            "[1,1]"
+            "[1,1]"
+            "[0,1][0,1][0,1]");
+
+  // [1] [1] [1, 1]
+  EXPECT_EQ(BCastList3({1}, {1}, {1, 1}),
+            "[1][1][1][1][1][1]"
+            "[1]"
+            "[1,1]"
+            "[0,1][0,1][0,1]");
+
+  EXPECT_EQ(BCastList3({1}, {1}, {1, 1}, false),
+            "[1,1][1,1][1,1][1,1][1,1][1,1]"
+            "[1,1]"
+            "[1,1]"
+            "[0,1][0,1][0,1]");
+}
+
+TEST(BCastListTest, Basic_TrueScalar_Scalar_Scalar) {
+  // Effectively it's a scalar and a scalar.
+  // [1, 1] [1] []
+  EXPECT_EQ(BCastList3({1, 1}, {1}, {}),
+            "[1][1][1][1][1][1]"
+            "[1]"
+            "[1,1]"
+            "[0,1][0,1][0,1]");
+
+  EXPECT_EQ(BCastList3({1, 1}, {1}, {}, false),
+            "[1,1][1,1][1,1][1,1][1,1][1,1]"
+            "[1,1]"
+            "[1,1]"
+            "[0,1][0,1][0,1]");
+
+  // [] [1, 1] [1]
+  EXPECT_EQ(BCastList3({}, {1, 1}, {1}),
+            "[1][1][1][1][1][1]"
+            "[1]"
+            "[1,1]"
+            "[0,1][0,1][0,1]");
+
+  EXPECT_EQ(BCastList3({}, {1, 1}, {1}, false),
+            "[1,1][1,1][1,1][1,1][1,1][1,1]"
+            "[1,1]"
+            "[1,1]"
+            "[0,1][0,1][0,1]");
+
+  // [1] [] [1, 1]
+  EXPECT_EQ(BCastList3({1}, {}, {1, 1}),
+            "[1][1][1][1][1][1]"
+            "[1]"
+            "[1,1]"
+            "[0,1][0,1][0,1]");
+
+  EXPECT_EQ(BCastList3({1}, {}, {1, 1}, false),
+            "[1,1][1,1][1,1][1,1][1,1][1,1]"
+            "[1,1]"
+            "[1,1]"
+            "[0,1][0,1][0,1]");
+}
+
+TEST(BCastTest, Basic_Tensor_Scalar) {
+  // Effectively it's a tensor and a scalar.
+  // [11, 7, 5, 3, 2] [1]
+  EXPECT_EQ(BCast({11, 7, 5, 3, 2}, {1}),
+            "[2310][1][1][2310]"
+            "[2310]"
+            "[11,7,5,3,2]"
+            "[][0,1,2,3,4]");
+
+  EXPECT_EQ(BCast({11, 7, 5, 3, 2}, {1}, false),
+            "[11,7,5,3,2][1,1,1,1,1][1,1,1,1,1][11,7,5,3,2]"
+            "[11,7,5,3,2]"
+            "[11,7,5,3,2]"
+            "[][0,1,2,3,4]");
+
+  // [1] [11, 7, 5, 3, 2]
+  EXPECT_EQ(BCast({1}, {11, 7, 5, 3, 2}),
+            "[1][2310][2310][1]"
+            "[2310]"
+            "[11,7,5,3,2]"
+            "[0,1,2,3,4][]");
+
+  EXPECT_EQ(BCast({1}, {11, 7, 5, 3, 2}, false),
+            "[1,1,1,1,1][11,7,5,3,2][11,7,5,3,2][1,1,1,1,1]"
+            "[11,7,5,3,2]"
+            "[11,7,5,3,2]"
+            "[0,1,2,3,4][]");
+
+  // int32 edge-case:
+  EXPECT_EQ(BCast({1, 2147483648}, {1}),
+            "[2147483648][1][1][2147483648]"
+            "[2147483648]"
+            "[1,2147483648]"
+            "[0][0,1]");
+}
+
+TEST(BCastTest, Basic_Tensor_With_DimSize_1_Scalar) {
+  // Effectively it's a tensor and a scalar.
+  // [11, 7, 5, 3, 2, 1] [1]
+  EXPECT_EQ(BCast({11, 7, 5, 3, 2, 1}, {1}),
+            "[2310][1][1][2310]"
+            "[2310]"
+            "[11,7,5,3,2,1]"
+            "[5][0,1,2,3,4,5]");
+
+  EXPECT_EQ(BCast({11, 7, 5, 3, 2, 1}, {1}, false),
+            "[11,7,5,3,2,1][1,1,1,1,1,1][1,1,1,1,1,1][11,7,5,3,2,1]"
+            "[11,7,5,3,2,1]"
+            "[11,7,5,3,2,1]"
+            "[5][0,1,2,3,4,5]");
+
+  // [1] [11, 7, 5, 3, 2, 1]
+  EXPECT_EQ(BCast({1}, {11, 7, 5, 3, 2, 1}),
+            "[1][2310][2310][1]"
+            "[2310]"
+            "[11,7,5,3,2,1]"
+            "[0,1,2,3,4,5][5]");
+
+  EXPECT_EQ(BCast({1}, {11, 7, 5, 3, 2, 1}, false),
+            "[1,1,1,1,1,1][11,7,5,3,2,1][11,7,5,3,2,1][1,1,1,1,1,1]"
+            "[11,7,5,3,2,1]"
+            "[11,7,5,3,2,1]"
+            "[0,1,2,3,4,5][5]");
+
+  // Effectively it's a tensor and a scalar.
+  // [11, 7, 5, 1, 1, 3, 2, 1] [1]
+  EXPECT_EQ(BCast({11, 7, 5, 1, 1, 3, 2, 1, 1}, {1}),
+            "[2310][1][1][2310]"
+            "[2310]"
+            "[11,7,5,1,1,3,2,1,1]"
+            "[3,4,7,8][0,1,2,3,4,5,6,7,8]");
+
+  EXPECT_EQ(BCast({11, 7, 5, 1, 1, 3, 2, 1, 1}, {1}, false),
+            "[11,7,5,1,1,3,2,1,1][1,1,1,1,1,1,1,1,1]"  // x_reshape(), x_bcast()
+            "[1,1,1,1,1,1,1,1,1][11,7,5,1,1,3,2,1,1]"  // y_reshape(), y_bcast()
+            "[11,7,5,1,1,3,2,1,1]"
+            "[11,7,5,1,1,3,2,1,1]"
+            "[3,4,7,8][0,1,2,3,4,5,6,7,8]");
+
+  // [1] [11, 7, 5, 1, 1, 3, 2, 1]
+  EXPECT_EQ(BCast({1}, {11, 7, 5, 1, 1, 3, 2, 1, 1}),
+            "[1][2310][2310][1]"
+            "[2310]"
+            "[11,7,5,1,1,3,2,1,1]"
+            "[0,1,2,3,4,5,6,7,8][3,4,7,8]");
+
+  EXPECT_EQ(BCast({1}, {11, 7, 5, 1, 1, 3, 2, 1, 1}, false),
+            "[1,1,1,1,1,1,1,1,1][11,7,5,1,1,3,2,1,1]"  // x_reshape(), x_bcast()
+            "[11,7,5,1,1,3,2,1,1][1,1,1,1,1,1,1,1,1]"  // y_reshape(), y_bcast()
+            "[11,7,5,1,1,3,2,1,1]"
+            "[11,7,5,1,1,3,2,1,1]"
+            "[0,1,2,3,4,5,6,7,8][3,4,7,8]");
+}
+
+TEST(BCastTest, Basic_Tensor_Vector) {
+  // [11, 7, 5, 3, 2] [2]
+  EXPECT_EQ(BCast({11, 7, 5, 3, 2}, {2}),
+            "[1155,2][1,1][1,2][1155,1]"
+            "[1155,2]"
+            "[11,7,5,3,2]"
+            "[][0,1,2,3]");
+
+  EXPECT_EQ(BCast({11, 7, 5, 3, 2}, {2}, false),
+            "[11,7,5,3,2][1,1,1,1,1][1,1,1,1,2][11,7,5,3,1]"
+            "[11,7,5,3,2]"
+            "[11,7,5,3,2]"
+            "[][0,1,2,3]");
+
+  // [2] [11, 7, 5, 3, 2]
+  EXPECT_EQ(BCast({2}, {11, 7, 5, 3, 2}),
+            "[1,2][1155,1][1155,2][1,1]"
+            "[1155,2]"
+            "[11,7,5,3,2]"
+            "[0,1,2,3][]");
+
+  EXPECT_EQ(BCast({2}, {11, 7, 5, 3, 2}, false),
+            "[1,1,1,1,2][11,7,5,3,1][11,7,5,3,2][1,1,1,1,1]"
+            "[11,7,5,3,2]"
+            "[11,7,5,3,2]"
+            "[0,1,2,3][]");
+}
+
+TEST(BCastTest, Basic_Tensor_Matrix) {
+  // [11, 7, 5, 3, 2] [3, 2]
+  EXPECT_EQ(BCast({11, 7, 5, 3, 2}, {3, 2}),
+            "[385,6][1,1][1,6][385,1]"
+            "[385,6]"
+            "[11,7,5,3,2]"
+            "[][0,1,2]");
+
+  EXPECT_EQ(BCast({11, 7, 5, 3, 2}, {3, 2}, false),
+            "[11,7,5,3,2][1,1,1,1,1][1,1,1,3,2][11,7,5,1,1]"
+            "[11,7,5,3,2]"
+            "[11,7,5,3,2]"
+            "[][0,1,2]");
+
+  // [3, 2] [11, 7, 5, 3, 2]
+  EXPECT_EQ(BCast({3, 2}, {11, 7, 5, 3, 2}),
+            "[1,6][385,1][385,6][1,1]"
+            "[385,6]"
+            "[11,7,5,3,2]"
+            "[0,1,2][]");
+
+  EXPECT_EQ(BCast({3, 2}, {11, 7, 5, 3, 2}, false),
+            "[1,1,1,3,2][11,7,5,1,1][11,7,5,3,2][1,1,1,1,1]"
+            "[11,7,5,3,2]"
+            "[11,7,5,3,2]"
+            "[0,1,2][]");
+}
+
+TEST(BCastTest, Basic_Tensor_Matrix_Column) {
+  // [11, 7, 5, 3, 2] [3, 1]
+  EXPECT_EQ(BCast({11, 7, 5, 3, 2}, {3, 1}),
+            "[385,3,2][1,1,1][1,3,1][385,1,2]"
+            "[385,3,2]"
+            "[11,7,5,3,2]"
+            "[][0,1,2,4]");
+
+  EXPECT_EQ(BCast({11, 7, 5, 3, 2}, {3, 1}, false),
+            "[11,7,5,3,2][1,1,1,1,1][1,1,1,3,1][11,7,5,1,2]"
+            "[11,7,5,3,2]"
+            "[11,7,5,3,2]"
+            "[][0,1,2,4]");
+
+  // [3, 1] [11, 7, 5, 3, 2]
+  EXPECT_EQ(BCast({3, 1}, {11, 7, 5, 3, 2}),
+            "[1,3,1][385,1,2][385,3,2][1,1,1]"
+            "[385,3,2]"
+            "[11,7,5,3,2]"
+            "[0,1,2,4][]");
+
+  EXPECT_EQ(BCast({3, 1}, {11, 7, 5, 3, 2}, false),
+            "[1,1,1,3,1][11,7,5,1,2][11,7,5,3,2][1,1,1,1,1]"
+            "[11,7,5,3,2]"
+            "[11,7,5,3,2]"
+            "[0,1,2,4][]");
+}
+
+TEST(BCastTest, Basic_Tensor_Matrix_As_Tensor) {
+  // [11, 7, 5, 3, 2] [7, 5, 1, 1]
+  EXPECT_EQ(BCast({11, 7, 5, 3, 2}, {7, 5, 1, 1}),
+            "[11,35,6][1,1,1][1,35,1][11,1,6]"
+            "[11,35,6]"
+            "[11,7,5,3,2]"
+            "[][0,3,4]");
+
+  EXPECT_EQ(BCast({11, 7, 5, 3, 2}, {7, 5, 1, 1}, false),
+            "[11,7,5,3,2][1,1,1,1,1][1,7,5,1,1][11,1,1,3,2]"
+            "[11,7,5,3,2]"
+            "[11,7,5,3,2]"
+            "[][0,3,4]");
+
+  // [7, 5, 1, 1] [11, 7, 5, 3, 2]
+  EXPECT_EQ(BCast({7, 5, 1, 1}, {11, 7, 5, 3, 2}),
+            "[1,35,1][11,1,6][11,35,6][1,1,1]"
+            "[11,35,6]"
+            "[11,7,5,3,2]"
+            "[0,3,4][]");
+
+  EXPECT_EQ(BCast({7, 5, 1, 1}, {11, 7, 5, 3, 2}, false),
+            "[1,7,5,1,1][11,1,1,3,2][11,7,5,3,2][1,1,1,1,1]"
+            "[11,7,5,3,2][11,7,5,3,2]"
+            "[0,3,4][]");
+}
+
+TEST(BCastTest, NegativeShape) {
+  const tensorflow::BCast bcast({10, -1, 1}, {10, 1, 1}, false);
+  EXPECT_TRUE(bcast.IsValid());
+  EXPECT_EQ(bcast.output_batch_size(), -1);
+  EXPECT_EQ(bcast.output_shape(), tensorflow::BCast::Vec({10, -1, 1}));
+}
+
+TEST(BCastTest, Overflow) {
+  const int64_t large_dim = 1LL << 62;
+  const tensorflow::BCast bcast({large_dim, 2}, {1, 1}, false);
+  EXPECT_TRUE(bcast.IsValid());
+  EXPECT_EQ(bcast.output_batch_size(), -1);
+}
+
+TEST(BCastTest, Basic_SymbolicShape) {
+  constexpr int64_t kSymDim1 = -10'000'000'000;
+  constexpr int64_t kSymDim2 = -10'000'000'001;
+
+  const tensorflow::BCast bcast({10, kSymDim1, kSymDim2}, {10, 1, 1}, false);
+  EXPECT_TRUE(bcast.IsValid());
+  EXPECT_EQ(bcast.output_batch_size(), -1);
+}
+
+TEST(BCastTest, Complex_BCast_To_Each_Other) {
+  // Rare cases. x and y broadcast to each other.  x and y are of
+  // different ranks.
+  // Can be verified in numpy as:
+  //   import numpy as np
+  //   x = np.arange(0,110).reshape([11,1,5,1,2])
+  //   y = np.arange(0,21).reshape([7,1,3,1])
+  //   np.shape(x + y)
+  //   Out[.]: (11, 7, 5, 3, 2)
+  std::string truth =
+      "[11,1,5,1,2][1,7,1,3,1][1,7,1,3,1][11,1,5,1,2]"
+      "[11,7,5,3,2]"
+      "[11,7,5,3,2]"
+      "[1,3][0,2,4]";
+
+  EXPECT_EQ(BCast({11, 1, 5, 1, 2}, {7, 1, 3, 1}), truth);
+  EXPECT_EQ(BCast({11, 1, 5, 1, 2}, {7, 1, 3, 1}, false), truth);
+}
+
+TEST(BCastListTest, Complex_BCast_To_Each_Other) {
+  // Rare cases. x, y and z broadcast to each other. x,y and z are of
+  // different ranks.
+  // Can be verified in numpy as:
+  //   import numpy as np
+  //   x = np.arange(0,22).reshape([11,1,1,1,2])
+  //   y = np.arange(0,21).reshape([7,1,3,1])
+  //   z = np.arange(0,5).reshape([5,1,1])
+  //   np.shape(x + y + z)
+  //   Out[.]: (11, 7, 5, 3, 2)
+  //
+  std::string truth =
+      "[11,1,1,1,2][1,7,5,3,1]"
+      "[1,7,1,3,1][11,1,5,1,2]"
+      "[1,1,5,1,1][11,7,1,3,2]"
+      "[11,7,5,3,2]"
+      "[11,7,5,3,2]"
+      "[1,2,3][0,2,4][0,1,3,4]";
+
+  EXPECT_EQ(BCastList3({11, 1, 1, 1, 2}, {7, 1, 3, 1}, {5, 1, 1}), truth);
+  EXPECT_EQ(BCastList3({11, 1, 1, 1, 2}, {7, 1, 3, 1}, {5, 1, 1}, false),
+            truth);
+}
+
+TEST(BCastTest, TestZeroDimensionShape) {
+  // (2,0,5) and (5) in both orders
+  EXPECT_EQ(BCast({2, 0, 5}, {5}),
+            "[0,5][1,1][1,5][0,1]"
+            "[0,5]"
+            "[2,0,5]"
+            "[][0,1]");
+  EXPECT_EQ(BCast({5}, {2, 0, 5}),
+            "[1,5][0,1][0,5][1,1]"
+            "[0,5]"
+            "[2,0,5]"
+            "[0,1][]");
+
+  EXPECT_EQ(BCast({2, 0, 5}, {5}, false),
+            "[2,0,5][1,1,1][1,1,5][2,0,1]"
+            "[2,0,5]"
+            "[2,0,5]"
+            "[][0,1]");
+  EXPECT_EQ(BCast({5}, {2, 0, 5}, false),
+            "[1,1,5][2,0,1][2,0,5][1,1,1]"
+            "[2,0,5]"
+            "[2,0,5]"
+            "[0,1][]");
+
+  // (2,0,3,0,5) and (5) in both orders
+  EXPECT_EQ(BCast({2, 0, 3, 0, 5}, {5}),
+            "[0,5][1,1][1,5][0,1]"
+            "[0,5]"
+            "[2,0,3,0,5]"
+            "[][0,1,2,3]");
+  EXPECT_EQ(BCast({5}, {2, 0, 3, 0, 5}),
+            "[1,5][0,1][0,5][1,1]"
+            "[0,5]"
+            "[2,0,3,0,5]"
+            "[0,1,2,3][]");
+
+  EXPECT_EQ(BCast({2, 0, 3, 0, 5}, {5}, false),
+            "[2,0,3,0,5][1,1,1,1,1][1,1,1,1,5][2,0,3,0,1]"
+            "[2,0,3,0,5]"
+            "[2,0,3,0,5]"
+            "[][0,1,2,3]");
+  EXPECT_EQ(BCast({5}, {2, 0, 3, 0, 5}, false),
+            "[1,1,1,1,5][2,0,3,0,1][2,0,3,0,5][1,1,1,1,1]"
+            "[2,0,3,0,5]"
+            "[2,0,3,0,5]"
+            "[0,1,2,3][]");
+
+  // (2,0,3,0,5) and (3,1,5) in both orders
+  EXPECT_EQ(BCast({2, 0, 3, 0, 5}, {3, 1, 5}),
+            "[0,3,0,5][1,1,1,1][1,3,1,5][0,1,0,1]"
+            "[0,3,0,5]"
+            "[2,0,3,0,5]"
+            "[][0,1,3]");
+  EXPECT_EQ(BCast({3, 1, 5}, {2, 0, 3, 0, 5}),
+            "[1,3,1,5][0,1,0,1][0,3,0,5][1,1,1,1]"
+            "[0,3,0,5]"
+            "[2,0,3,0,5]"
+            "[0,1,3][]");
+
+  EXPECT_EQ(BCast({2, 0, 3, 0, 5}, {3, 1, 5}, false),
+            "[2,0,3,0,5][1,1,1,1,1][1,1,3,1,5][2,0,1,0,1]"
+            "[2,0,3,0,5]"
+            "[2,0,3,0,5]"
+            "[][0,1,3]");
+  EXPECT_EQ(BCast({3, 1, 5}, {2, 0, 3, 0, 5}, false),
+            "[1,1,3,1,5][2,0,1,0,1][2,0,3,0,5][1,1,1,1,1]"
+            "[2,0,3,0,5]"
+            "[2,0,3,0,5]"
+            "[0,1,3][]");
+}
+
+TEST(BCastTest, BatchIndices) {
+  EXPECT_EQ("[0,0,0,0][0,1,2,3]", BCastBatchIndices({1}, {4}));
+  // Invalid broadcast.
+  EXPECT_EQ("[][]", BCastBatchIndices({5}, {7}));
+  // Same shape, no batch indices.
+  EXPECT_EQ("[][]", BCastBatchIndices({2, 4, 6}, {2, 4, 6}));
+  // More complicated broadcasts.
+  EXPECT_EQ("[0,0,0,0,1,1,1,1,2,2,2,2][0,1,2,3,0,1,2,3,0,1,2,3]",
+            BCastBatchIndices({3, 1}, {1, 4}));
+  EXPECT_EQ("[0,0,1,1,2,2,0,0,1,1,2,2][0,1,0,1,0,1,2,3,2,3,2,3]",
+            BCastBatchIndices({3, 1}, {2, 1, 2}));
+}
+
+void BM_BCastSetup(::testing::benchmark::State& state) {
+  const int same_shape = state.range(0);
+
+  if (same_shape) {
+    state.SetLabel("same_shapes");
+    for (auto s : state) {
+      class BCast b({1000, 100}, {1000, 100});
+    }
+  } else {
+    state.SetLabel("different_shapes");
+    for (auto s : state) {
+      class BCast b({3, 1, 5}, {2, 0, 3, 0, 5});
+    }
+  }
+}
+BENCHMARK(BM_BCastSetup)->Arg(0)->Arg(1);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/core/util/command_line_flags.h

@@ -0,0 +1,31 @@
+/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_UTIL_COMMAND_LINE_FLAGS_H_
+#define TENSORFLOW_CORE_UTIL_COMMAND_LINE_FLAGS_H_
+
+#include <functional>
+#include <string>
+#include <vector>
+
+#include "xla/tsl/util/command_line_flags.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+using tsl::Flag;   // NOLINT
+using tsl::Flags;  // NOLINT
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_UTIL_COMMAND_LINE_FLAGS_H_

tensorflow/core/util/command_line_flags_test.cc

@@ -0,0 +1,366 @@
+/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/util/command_line_flags.h"
+
+#include <ctype.h>
+
+#include <vector>
+
+#include "absl/strings/ascii.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace {
+// The returned array is only valid for the lifetime of the input vector.
+// We're using const casting because we need to pass in an argv-style array of
+// char* pointers for the API, even though we know they won't be altered.
+std::vector<char*> CharPointerVectorFromStrings(
+    const std::vector<std::string>& strings) {
+  std::vector<char *> result;
+  result.reserve(strings.size());
+  for (const std::string& string : strings) {
+    result.push_back(const_cast<char *>(string.c_str()));
+  }
+  return result;
+}
+}  // namespace
+
+TEST(CommandLineFlagsTest, BasicUsage) {
+  int some_int32_set_directly = 10;
+  int some_int32_set_via_hook = 20;
+  int64_t some_int64_set_directly = 21474836470;  // max int32 is 2147483647
+  int64_t some_int64_set_via_hook = 21474836479;  // max int32 is 2147483647
+  bool some_switch_set_directly = false;
+  bool some_switch_set_via_hook = true;
+  bool some_switch_set_capitalized = false;
+  bool some_switch_set_by_number = false;
+  std::string some_name_set_directly = "something_a";
+  std::string some_name_set_via_hook = "something_b";
+  float some_float_set_directly = -23.23f;
+  float some_float_set_via_hook = -25.23f;
+  std::vector<std::string> argv_strings = {
+      "program_name",
+      "--some_int32_set_directly=20",
+      "--some_int32_set_via_hook=50",
+      "--some_int64_set_directly=214748364700",
+      "--some_int64_set_via_hook=214748364710",
+      "--some_switch_set_directly",
+      "--some_switch_set_via_hook=false",
+      "--some_switch_set_capitalized=True",
+      "--some_switch_set_by_number=1",
+      "--some_name_set_directly=somethingelse",
+      "--some_name_set_via_hook=anythingelse",
+      "--some_float_set_directly=42.0",
+      "--some_float_set_via_hook=43.0"};
+  int argc = argv_strings.size();
+  std::vector<char *> argv_array = CharPointerVectorFromStrings(argv_strings);
+  bool parsed_ok = Flags::Parse(
+      &argc, argv_array.data(),
+      {
+          Flag("some_int32_set_directly", &some_int32_set_directly,
+               "some int32 set directly"),
+          Flag(
+              "some_int32_set_via_hook",
+              [&](int32_t value) {
+                some_int32_set_via_hook = value;
+                return true;
+              },
+              some_int32_set_via_hook, "some int32 set via hook"),
+          Flag("some_int64_set_directly", &some_int64_set_directly,
+               "some int64 set directly"),
+          Flag(
+              "some_int64_set_via_hook",
+              [&](int64_t value) {
+                some_int64_set_via_hook = value;
+                return true;
+              },
+              some_int64_set_via_hook, "some int64 set via hook"),
+          Flag("some_switch_set_directly", &some_switch_set_directly,
+               "some switch set directly"),
+          Flag(
+              "some_switch_set_via_hook",
+              [&](bool value) {
+                some_switch_set_via_hook = value;
+                return true;
+              },
+              some_switch_set_via_hook, "some switch set via hook"),
+          Flag("some_switch_set_capitalized", &some_switch_set_capitalized,
+               "some switch set capitalized"),
+          Flag("some_switch_set_by_number", &some_switch_set_by_number,
+               "some switch set by number"),
+          Flag("some_name_set_directly", &some_name_set_directly,
+               "some name set directly"),
+          Flag(
+              "some_name_set_via_hook",
+              [&](std::string value) {
+                some_name_set_via_hook = std::move(value);
+                return true;
+              },
+              some_name_set_via_hook, "some name set via hook"),
+          Flag("some_float_set_directly", &some_float_set_directly,
+               "some float set directly"),
+          Flag(
+              "some_float_set_via_hook",
+              [&](float value) {
+                some_float_set_via_hook = value;
+                return true;
+              },
+              some_float_set_via_hook, "some float set via hook"),
+      });
+
+  EXPECT_EQ(true, parsed_ok);
+  EXPECT_EQ(20, some_int32_set_directly);
+  EXPECT_EQ(50, some_int32_set_via_hook);
+  EXPECT_EQ(214748364700, some_int64_set_directly);
+  EXPECT_EQ(214748364710, some_int64_set_via_hook);
+  EXPECT_EQ(true, some_switch_set_directly);
+  EXPECT_EQ(false, some_switch_set_via_hook);
+  EXPECT_EQ(true, some_switch_set_capitalized);
+  EXPECT_EQ(true, some_switch_set_by_number);
+  EXPECT_EQ("somethingelse", some_name_set_directly);
+  EXPECT_EQ("anythingelse", some_name_set_via_hook);
+  EXPECT_NEAR(42.0f, some_float_set_directly, 1e-5f);
+  EXPECT_NEAR(43.0f, some_float_set_via_hook, 1e-5f);
+  EXPECT_EQ(argc, 1);
+}
+
+TEST(CommandLineFlagsTest, BadIntValue) {
+  int some_int = 10;
+  int argc = 2;
+  std::vector<std::string> argv_strings = {"program_name",
+                                           "--some_int=notanumber"};
+  std::vector<char *> argv_array = CharPointerVectorFromStrings(argv_strings);
+  bool parsed_ok = Flags::Parse(&argc, argv_array.data(),
+                                {Flag("some_int", &some_int, "some int")});
+
+  EXPECT_EQ(false, parsed_ok);
+  EXPECT_EQ(10, some_int);
+  EXPECT_EQ(argc, 1);
+}
+
+TEST(CommandLineFlagsTest, BadBoolValue) {
+  bool some_switch = false;
+  int argc = 2;
+  std::vector<std::string> argv_strings = {"program_name",
+                                           "--some_switch=notabool"};
+  std::vector<char *> argv_array = CharPointerVectorFromStrings(argv_strings);
+  bool parsed_ok =
+      Flags::Parse(&argc, argv_array.data(),
+                   {Flag("some_switch", &some_switch, "some switch")});
+
+  EXPECT_EQ(false, parsed_ok);
+  EXPECT_EQ(false, some_switch);
+  EXPECT_EQ(argc, 1);
+}
+
+TEST(CommandLineFlagsTest, BadFloatValue) {
+  float some_float = -23.23f;
+  int argc = 2;
+  std::vector<std::string> argv_strings = {"program_name",
+                                           "--some_float=notanumber"};
+  std::vector<char *> argv_array = CharPointerVectorFromStrings(argv_strings);
+  bool parsed_ok =
+      Flags::Parse(&argc, argv_array.data(),
+                   {Flag("some_float", &some_float, "some float")});
+
+  EXPECT_EQ(false, parsed_ok);
+  EXPECT_NEAR(-23.23f, some_float, 1e-5f);
+  EXPECT_EQ(argc, 1);
+}
+
+TEST(CommandLineFlagsTest, FailedInt32Hook) {
+  int argc = 2;
+  std::vector<std::string> argv_strings = {"program_name", "--some_int32=200"};
+  std::vector<char *> argv_array = CharPointerVectorFromStrings(argv_strings);
+  bool parsed_ok =
+      Flags::Parse(&argc, argv_array.data(),
+                   {Flag(
+                       "some_int32", [](int32_t value) { return false; }, 30,
+                       "some int32")});
+
+  EXPECT_EQ(false, parsed_ok);
+  EXPECT_EQ(argc, 1);
+}
+
+TEST(CommandLineFlagsTest, FailedInt64Hook) {
+  int argc = 2;
+  std::vector<std::string> argv_strings = {"program_name", "--some_int64=200"};
+  std::vector<char *> argv_array = CharPointerVectorFromStrings(argv_strings);
+  bool parsed_ok =
+      Flags::Parse(&argc, argv_array.data(),
+                   {Flag(
+                       "some_int64", [](int64_t value) { return false; }, 30,
+                       "some int64")});
+
+  EXPECT_EQ(false, parsed_ok);
+  EXPECT_EQ(argc, 1);
+}
+
+TEST(CommandLineFlagsTest, FailedFloatHook) {
+  int argc = 2;
+  std::vector<std::string> argv_strings = {"program_name",
+                                           "--some_float=200.0"};
+  std::vector<char *> argv_array = CharPointerVectorFromStrings(argv_strings);
+  bool parsed_ok =
+      Flags::Parse(&argc, argv_array.data(),
+                   {Flag("some_float", [](float value) { return false; }, 30.0f,
+                         "some float")});
+
+  EXPECT_EQ(false, parsed_ok);
+  EXPECT_EQ(argc, 1);
+}
+
+TEST(CommandLineFlagsTest, FailedBoolHook) {
+  int argc = 2;
+  std::vector<std::string> argv_strings = {"program_name",
+                                           "--some_switch=true"};
+  std::vector<char *> argv_array = CharPointerVectorFromStrings(argv_strings);
+  bool parsed_ok =
+      Flags::Parse(&argc, argv_array.data(),
+                   {Flag("some_switch", [](bool value) { return false; }, false,
+                         "some switch")});
+
+  EXPECT_EQ(false, parsed_ok);
+  EXPECT_EQ(argc, 1);
+}
+
+TEST(CommandLineFlagsTest, FailedStringHook) {
+  int argc = 2;
+  std::vector<std::string> argv_strings = {"program_name", "--some_name=true"};
+  std::vector<char *> argv_array = CharPointerVectorFromStrings(argv_strings);
+  bool parsed_ok =
+      Flags::Parse(&argc, argv_array.data(),
+                   {Flag(
+                       "some_name", [](std::string value) { return false; }, "",
+                       "some name")});
+
+  EXPECT_EQ(false, parsed_ok);
+  EXPECT_EQ(argc, 1);
+}
+
+TEST(CommandLineFlagsTest, RepeatedStringHook) {
+  int argc = 3;
+  std::vector<std::string> argv_strings = {"program_name", "--some_name=this",
+                                           "--some_name=that"};
+  std::vector<char *> argv_array = CharPointerVectorFromStrings(argv_strings);
+  int call_count = 0;
+  bool parsed_ok = Flags::Parse(&argc, argv_array.data(),
+                                {Flag(
+                                    "some_name",
+                                    [&call_count](std::string value) {
+                                      call_count++;
+                                      return true;
+                                    },
+                                    "", "some name")});
+
+  EXPECT_EQ(true, parsed_ok);
+  EXPECT_EQ(argc, 1);
+  EXPECT_EQ(call_count, 2);
+}
+
+// Return whether str==pat, but allowing any whitespace in pat
+// to match zero or more whitespace characters in str.
+static bool MatchWithAnyWhitespace(const std::string& str,
+                                   const std::string& pat) {
+  bool matching = true;
+  int pat_i = 0;
+  for (int str_i = 0; str_i != str.size() && matching; str_i++) {
+    if (absl::ascii_isspace(str[str_i])) {
+      matching = (pat_i != pat.size() && absl::ascii_isspace(pat[pat_i]));
+    } else {
+      while (pat_i != pat.size() && absl::ascii_isspace(pat[pat_i])) {
+        pat_i++;
+      }
+      matching = (pat_i != pat.size() && str[str_i] == pat[pat_i++]);
+    }
+  }
+  while (pat_i != pat.size() && absl::ascii_isspace(pat[pat_i])) {
+    pat_i++;
+  }
+  return (matching && pat_i == pat.size());
+}
+
+TEST(CommandLineFlagsTest, UsageString) {
+  int some_int = 10;
+  int64_t some_int64 = 21474836470;  // max int32 is 2147483647
+  bool some_switch = false;
+  std::string some_name = "something";
+  // Don't test float in this case, because precision is hard to predict and
+  // match against, and we don't want a franky test.
+  const std::string tool_name = "some_tool_name";
+  std::string usage = Flags::Usage(
+      tool_name + "<flags>", {Flag("some_int", &some_int, "some int"),
+                              Flag("some_int64", &some_int64, "some int64"),
+                              Flag("some_switch", &some_switch, "some switch"),
+                              Flag("some_name", &some_name, "some name")});
+  // Match the usage message, being sloppy about whitespace.
+  const char *expected_usage =
+      " usage: some_tool_name <flags>\n"
+      "Flags:\n"
+      "--some_int=10 int32 some int\n"
+      "--some_int64=21474836470 int64 some int64\n"
+      "--some_switch=false bool some switch\n"
+      "--some_name=\"something\" string some name\n";
+  ASSERT_EQ(MatchWithAnyWhitespace(usage, expected_usage), true);
+
+  // Again but with no flags.
+  usage = Flags::Usage(tool_name, {});
+  ASSERT_EQ(MatchWithAnyWhitespace(usage, " usage: some_tool_name\n"), true);
+}
+
+namespace {
+template <typename T, typename ExpectationFun>
+void PrefixTestTempl(ExpectationFun expectation_fun, const T& value0,
+                     const T& value1, std::string str0, std::string str1) {
+  int argc = 3;
+  std::vector<std::string> argv_strings = {
+      "program_name",
+      "--hello" + str0,
+      "--hello_world" + str1,
+  };
+  std::vector<char *> argv_array = CharPointerVectorFromStrings(argv_strings);
+
+  T hello{};
+  T hello_world{};
+  bool parsed_ok = Flags::Parse(
+      &argc, argv_array.data(),
+      {
+          Flag("hello", &hello, "usage of hello"),
+          Flag("hello_world", &hello_world, "usage of hello world"),
+      });
+
+  EXPECT_EQ(true, parsed_ok);
+  expectation_fun(value0, hello);
+  expectation_fun(value1, hello_world);
+  EXPECT_EQ(argc, 1);
+}
+}  // namespace
+
+TEST(CommandLineFlagsTest, OneArgumentIsAPrefixOfAnother) {
+  auto expect_eq = [](auto a, auto b) { EXPECT_EQ(a, b); };
+  auto expect_near = [](auto a, auto b) { EXPECT_NEAR(a, b, 1e-5f); };
+
+  PrefixTestTempl<int32_t>(expect_eq, 1, 2, "=1", "=2");
+  PrefixTestTempl<int64_t>(expect_eq, 1, 2, "=1", "=2");
+  PrefixTestTempl<bool>(expect_eq, false, true, "=false", "=true");
+  PrefixTestTempl<bool>(expect_eq, false, true, "=false", "");
+  PrefixTestTempl<bool>(expect_eq, true, false, "=true", "=false");
+  PrefixTestTempl<bool>(expect_eq, true, false, "", "=false");
+  PrefixTestTempl<std::string>(expect_eq, "a", "b", "=a", "=b");
+  PrefixTestTempl<float>(expect_near, 0.1f, 0.2f, "=0.1", "=0.2");
+}
+
+}  // namespace tensorflow

tensorflow/core/util/ctc/BUILD

@@ -0,0 +1,116 @@
+# Description: CTC, Connectionist Temporal Classification,
+# is a type of seq2seq loss.  The libraries in this directory
+# implement the CTC loss and a number of CTC decoders.
+
+load("//tensorflow:tensorflow.bzl", "tf_cc_tests")
+load("//tensorflow:tensorflow.default.bzl", "filegroup")
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//visibility:public"],
+    licenses = ["notice"],
+)
+
+filegroup(
+    name = "mobile_hdrs",
+    srcs = [
+        "ctc_beam_entry.h",
+        "ctc_beam_scorer.h",
+        "ctc_beam_search.h",
+        "ctc_decoder.h",
+        "ctc_loss_util.h",
+    ],
+)
+
+alias(
+    name = "android_srcs",
+    actual = ":mobile_hdrs",
+)
+
+cc_library(
+    name = "ctc",
+    deps = [
+        ":ctc_beam_search_lib",
+        ":ctc_loss_calculator_lib",
+    ],
+)
+
+cc_library(
+    name = "ctc_beam_search_lib",
+    srcs = [
+        "ctc_beam_entry.h",
+        "ctc_beam_scorer.h",
+        "ctc_beam_search.h",
+        "ctc_decoder.h",
+    ],
+    hdrs = [
+        "ctc_beam_entry.h",
+        "ctc_beam_scorer.h",
+        "ctc_beam_search.h",
+        "ctc_decoder.h",
+    ],
+    deps = [
+        ":ctc_loss_util_lib",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:lib_internal",
+        "@eigen_archive//:eigen3",
+    ],
+)
+
+tf_cc_tests(
+    name = "ctc_beam_search_test",
+    size = "small",
+    srcs = [
+        "ctc_beam_search_test.cc",
+    ],
+    deps = [
+        ":ctc_beam_search_lib",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "@com_google_absl//absl/log",
+    ],
+)
+
+cc_library(
+    name = "ctc_loss_calculator_lib",
+    srcs = [
+        "ctc_loss_calculator.cc",
+    ],
+    hdrs = [
+        "ctc_loss_calculator.h",
+    ],
+    deps = [
+        ":ctc_loss_util_lib",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@eigen_archive//:eigen3",
+    ],
+)
+
+cc_library(
+    name = "ctc_loss_util_lib",
+    hdrs = [
+        "ctc_loss_util.h",
+    ],
+)
+
+# For a more maintainable build this target should not exist and the headers
+# should  be split into the existing cc_library targets, but this change was
+# automatically  done so that we can remove long standing issues and complexity
+# in the build system. It's up to the OWNERS of this package to get rid of it or
+# not. The use of the textual_hdrs attribute is discouraged, use hdrs instead.
+# Here it is used to avoid header parsing errors in packages where the feature
+# parse_headers was enabled since loose headers were not being parsed. See
+# go/loose-lsc-one-target-approach for more details.
+cc_library(
+    name = "loose_headers",
+    tags = ["avoid_dep"],
+    textual_hdrs = ["ctc_beam_search.h"],
+    visibility = ["//visibility:public"],
+)

tensorflow/core/util/ctc/ctc_beam_entry.h

@@ -0,0 +1,154 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+// LINT.IfChange
+
+#ifndef TENSORFLOW_CORE_UTIL_CTC_CTC_BEAM_ENTRY_H_
+#define TENSORFLOW_CORE_UTIL_CTC_CTC_BEAM_ENTRY_H_
+
+#include <algorithm>
+#include <memory>
+#include <vector>
+
+#include "Eigen/Core"  // from @eigen_archive
+#include "tensorflow/core/lib/gtl/flatmap.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/macros.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/util/ctc/ctc_loss_util.h"
+
+namespace tensorflow {
+namespace ctc {
+
+// The ctc_beam_search namespace holds several classes meant to be accessed only
+// in case of extending the CTCBeamSearch decoder to allow custom scoring
+// functions.
+//
+// BeamEntry is exposed through template arguments BeamScorer and BeamComparer
+// of CTCBeamSearch (ctc_beam_search.h).
+namespace ctc_beam_search {
+
+struct EmptyBeamState {};
+
+template <typename T>
+struct BeamProbability {
+  BeamProbability()
+      : total(kLogZero<T>()), blank(kLogZero<T>()), label(kLogZero<T>()) {}
+  void Reset() {
+    total = kLogZero<T>();
+    blank = kLogZero<T>();
+    label = kLogZero<T>();
+  }
+  T total;
+  T blank;
+  T label;
+};
+
+template <class T, class CTCBeamState>
+class BeamRoot;
+
+template <class T, class CTCBeamState = EmptyBeamState>
+struct BeamEntry {
+  // BeamRoot<CTCBeamState>::AddEntry() serves as the factory method.
+  friend BeamEntry<T, CTCBeamState>* BeamRoot<T, CTCBeamState>::AddEntry(
+      BeamEntry<T, CTCBeamState>* p, int l);
+  inline bool Active() const { return newp.total != kLogZero<T>(); }
+  // Return the child at the given index, or construct a new one in-place if
+  // none was found.
+  BeamEntry<T, CTCBeamState>& GetChild(int ind) {
+    auto entry = children.emplace(ind, nullptr);
+    auto& child_entry = entry.first->second;
+    // If this is a new child, populate the BeamEntry<CTCBeamState>*.
+    if (entry.second) {
+      child_entry = beam_root->AddEntry(this, ind);
+    }
+    return *child_entry;
+  }
+  std::vector<int> LabelSeq(bool merge_repeated) const {
+    std::vector<int> labels;
+    int prev_label = -1;
+    const BeamEntry<T, CTCBeamState>* c = this;
+    while (c->parent != nullptr) {  // Checking c->parent to skip root leaf.
+      if (!merge_repeated || c->label != prev_label) {
+        labels.push_back(c->label);
+      }
+      prev_label = c->label;
+      c = c->parent;
+    }
+    std::reverse(labels.begin(), labels.end());
+    return labels;
+  }
+
+  BeamEntry<T, CTCBeamState>* parent;
+  int label;
+  // All instances of child BeamEntry are owned by *beam_root.
+  gtl::FlatMap<int, BeamEntry<T, CTCBeamState>*> children;
+  BeamProbability<T> oldp;
+  BeamProbability<T> newp;
+  CTCBeamState state;
+
+ private:
+  // Constructor giving parent, label, and the beam_root.
+  // The object pointed to by p cannot be copied and should not be moved,
+  // otherwise parent will become invalid.
+  // This private constructor is only called through the factory method
+  // BeamRoot<CTCBeamState>::AddEntry().
+  BeamEntry(BeamEntry* p, int l, BeamRoot<T, CTCBeamState>* beam_root)
+      : parent(p), label(l), beam_root(beam_root) {}
+  BeamRoot<T, CTCBeamState>* beam_root;
+  BeamEntry(const BeamEntry&) = delete;
+  void operator=(const BeamEntry&) = delete;
+};
+
+// This class owns all instances of BeamEntry.  This is used to avoid recursive
+// destructor call during destruction.
+template <class T, class CTCBeamState = EmptyBeamState>
+class BeamRoot {
+ public:
+  BeamRoot(BeamEntry<T, CTCBeamState>* p, int l) {
+    root_entry_ = AddEntry(p, l);
+  }
+  BeamRoot(const BeamRoot&) = delete;
+  BeamRoot& operator=(const BeamRoot&) = delete;
+
+  BeamEntry<T, CTCBeamState>* AddEntry(BeamEntry<T, CTCBeamState>* p, int l) {
+    auto* new_entry = new BeamEntry<T, CTCBeamState>(p, l, this);
+    beam_entries_.emplace_back(new_entry);
+    return new_entry;
+  }
+  BeamEntry<T, CTCBeamState>* RootEntry() const { return root_entry_; }
+
+ private:
+  BeamEntry<T, CTCBeamState>* root_entry_ = nullptr;
+  std::vector<std::unique_ptr<BeamEntry<T, CTCBeamState>>> beam_entries_;
+};
+
+// BeamComparer is the default beam comparer provided in CTCBeamSearch.
+template <class T, class CTCBeamState = EmptyBeamState>
+class BeamComparer {
+ public:
+  virtual ~BeamComparer() {}
+  virtual bool inline operator()(const BeamEntry<T, CTCBeamState>* a,
+                                 const BeamEntry<T, CTCBeamState>* b) const {
+    return a->newp.total > b->newp.total;
+  }
+};
+
+}  // namespace ctc_beam_search
+
+}  // namespace ctc
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_UTIL_CTC_CTC_BEAM_ENTRY_H_
+// LINT.ThenChange(//tensorflow/lite/kernels/ctc/ctc_beam_entry.h)

tensorflow/core/util/ctc/ctc_beam_scorer.h

@@ -0,0 +1,77 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+// LINT.IfChange
+
+// Collection of scoring classes that can be extended and provided to the
+// CTCBeamSearchDecoder to incorporate additional scoring logic (such as a
+// language model).
+//
+// To build a custom scorer extend and implement the pure virtual methods from
+// BeamScorerInterface. The default CTC decoding behavior is implemented
+// through BaseBeamScorer.
+
+#ifndef TENSORFLOW_CORE_UTIL_CTC_CTC_BEAM_SCORER_H_
+#define TENSORFLOW_CORE_UTIL_CTC_CTC_BEAM_SCORER_H_
+
+#include "tensorflow/core/util/ctc/ctc_beam_entry.h"
+
+namespace tensorflow {
+namespace ctc {
+
+// Base implementation of a beam scorer used by default by the decoder that can
+// be subclassed and provided as an argument to CTCBeamSearchDecoder, if complex
+// scoring is required. Its main purpose is to provide a thin layer for
+// integrating language model scoring easily.
+template <typename T, typename CTCBeamState>
+class BaseBeamScorer {
+ public:
+  virtual ~BaseBeamScorer() {}
+  // State initialization.
+  virtual void InitializeState(CTCBeamState* root) const {}
+  // ExpandState is called when expanding a beam to one of its children.
+  // Called at most once per child beam. In the simplest case, no state
+  // expansion is done.
+  virtual void ExpandState(const CTCBeamState& from_state, int from_label,
+                           CTCBeamState* to_state, int to_label) const {}
+  // ExpandStateEnd is called after decoding has finished. Its purpose is to
+  // allow a final scoring of the beam in its current state, before resorting
+  // and retrieving the TopN requested candidates. Called at most once per beam.
+  virtual void ExpandStateEnd(CTCBeamState* state) const {}
+  // GetStateExpansionScore should be an inexpensive method to retrieve the
+  // (cached) expansion score computed within ExpandState. The score is
+  // multiplied (log-addition) with the input score at the current step from
+  // the network.
+  //
+  // The score returned should be a log-probability. In the simplest case, as
+  // there's no state expansion logic, the expansion score is zero.
+  virtual T GetStateExpansionScore(const CTCBeamState& state,
+                                   T previous_score) const {
+    return previous_score;
+  }
+  // GetStateEndExpansionScore should be an inexpensive method to retrieve the
+  // (cached) expansion score computed within ExpandStateEnd. The score is
+  // multiplied (log-addition) with the final probability of the beam.
+  //
+  // The score returned should be a log-probability.
+  virtual T GetStateEndExpansionScore(const CTCBeamState& state) const {
+    return T(0);
+  }
+};
+
+}  // namespace ctc
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_UTIL_CTC_CTC_BEAM_SCORER_H_
+// LINT.ThenChange(//tensorflow/lite/kernels/ctc/ctc_beam_scorer.h)

tensorflow/core/util/ctc/ctc_beam_search.h

@@ -0,0 +1,438 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+// LINT.IfChange
+
+#ifndef TENSORFLOW_CORE_UTIL_CTC_CTC_BEAM_SEARCH_H_
+#define TENSORFLOW_CORE_UTIL_CTC_CTC_BEAM_SEARCH_H_
+
+#include <algorithm>
+#include <cmath>
+#include <limits>
+#include <memory>
+#include <vector>
+
+#include "Eigen/Core"  // from @eigen_archive
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/lib/gtl/top_n.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/macros.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/util/ctc/ctc_beam_entry.h"
+#include "tensorflow/core/util/ctc/ctc_beam_scorer.h"
+#include "tensorflow/core/util/ctc/ctc_decoder.h"
+#include "tensorflow/core/util/ctc/ctc_loss_util.h"
+
+namespace tensorflow {
+namespace ctc {
+
+template <typename T, typename CTCBeamState = ctc_beam_search::EmptyBeamState,
+          typename CTCBeamComparer =
+              ctc_beam_search::BeamComparer<T, CTCBeamState>>
+class CTCBeamSearchDecoder : public CTCDecoder<T> {
+  // Beam Search
+  //
+  // Example (GravesTh Fig. 7.5):
+  //         a    -
+  //  P = [ 0.3  0.7 ]  t = 0
+  //      [ 0.4  0.6 ]  t = 1
+  //
+  // Then P(l = -) = P(--) = 0.7 * 0.6 = 0.42
+  //      P(l = a) = P(a-) + P(aa) + P(-a) = 0.3*0.4 + ... = 0.58
+  //
+  // In this case, Best Path decoding is suboptimal.
+  //
+  // For Beam Search, we use the following main recurrence relations:
+  //
+  // Relation 1:
+  // ---------------------------------------------------------- Eq. 1
+  //      P(l=abcd @ t=7) = P(l=abc  @ t=6) * P(d @ 7)
+  //                      + P(l=abcd @ t=6) * (P(d @ 7) + P(- @ 7))
+  // where P(l=? @ t=7), ? = a, ab, abc, abcd are all stored and
+  // updated recursively in the beam entry.
+  //
+  // Relation 2:
+  // ---------------------------------------------------------- Eq. 2
+  //      P(l=abc? @ t=3) = P(l=abc @ t=2) * P(? @ 3)
+  // for ? in a, b, d, ..., (not including c or the blank index),
+  // and the recurrence starts from the beam entry for P(l=abc @ t=2).
+  //
+  // For this case, the length of the new sequence equals t+1 (t
+  // starts at 0).  This special case can be calculated as:
+  //   P(l=abc? @ t=3) = P(a @ 0)*P(b @ 1)*P(c @ 2)*P(? @ 3)
+  // but we calculate it recursively for speed purposes.
+  typedef ctc_beam_search::BeamEntry<T, CTCBeamState> BeamEntry;
+  typedef ctc_beam_search::BeamRoot<T, CTCBeamState> BeamRoot;
+  typedef ctc_beam_search::BeamProbability<T> BeamProbability;
+
+ public:
+  typedef BaseBeamScorer<T, CTCBeamState> DefaultBeamScorer;
+
+  // The beam search decoder is constructed specifying the beam_width (number of
+  // candidates to keep at each decoding timestep) and a beam scorer (used for
+  // custom scoring, for example enabling the use of a language model).
+  // The ownership of the scorer remains with the caller. The default
+  // implementation, CTCBeamSearchDecoder<>::DefaultBeamScorer, generates the
+  // standard beam search.
+  CTCBeamSearchDecoder(int num_classes, int beam_width,
+                       BaseBeamScorer<T, CTCBeamState>* scorer,
+                       int batch_size = 1, bool merge_repeated = false)
+      : CTCDecoder<T>(num_classes, batch_size, merge_repeated),
+        beam_width_(beam_width),
+        leaves_(beam_width),
+        beam_scorer_(CHECK_NOTNULL(scorer)) {
+    Reset();
+  }
+
+  ~CTCBeamSearchDecoder() override {}
+
+  // Run the hibernating beam search algorithm on the given input.
+  absl::Status Decode(const typename CTCDecoder<T>::SequenceLength& seq_len,
+                      const std::vector<typename CTCDecoder<T>::Input>& input,
+                      std::vector<typename CTCDecoder<T>::Output>* output,
+                      typename CTCDecoder<T>::ScoreOutput* scores) override;
+
+  // Calculate the next step of the beam search and update the internal state.
+  template <typename Vector>
+  void Step(const Vector& log_input_t);
+
+  template <typename Vector>
+  T GetTopK(const int K, const Vector& input, std::vector<T>* top_k_logits,
+            std::vector<int>* top_k_indices);
+
+  // Retrieve the beam scorer instance used during decoding.
+  BaseBeamScorer<T, CTCBeamState>* GetBeamScorer() const {
+    return beam_scorer_;
+  }
+
+  // Set label selection parameters for faster decoding.
+  // See comments for label_selection_size_ and label_selection_margin_.
+  void SetLabelSelectionParameters(int label_selection_size,
+                                   T label_selection_margin) {
+    label_selection_size_ = label_selection_size;
+    label_selection_margin_ = label_selection_margin;
+  }
+
+  // Reset the beam search
+  void Reset();
+
+  // Extract the top n paths at current time step
+  absl::Status TopPaths(int n, std::vector<std::vector<int>>* paths,
+                        std::vector<T>* log_probs, bool merge_repeated) const;
+
+ private:
+  int beam_width_;
+
+  // Label selection is designed to avoid possibly very expensive scorer calls,
+  // by pruning the hypotheses based on the input alone.
+  // Label selection size controls how many items in each beam are passed
+  // through to the beam scorer. Only items with top N input scores are
+  // considered.
+  // Label selection margin controls the difference between minimal input score
+  // (versus the best scoring label) for an item to be passed to the beam
+  // scorer. This margin is expressed in terms of log-probability.
+  // Default is to do no label selection.
+  // For more detail: https://research.google.com/pubs/pub44823.html
+  int label_selection_size_ = 0;       // zero means unlimited
+  T label_selection_margin_ = -1;      // -1 means unlimited.
+
+  gtl::TopN<BeamEntry*, CTCBeamComparer> leaves_;
+  std::unique_ptr<BeamRoot> beam_root_;
+  BaseBeamScorer<T, CTCBeamState>* beam_scorer_;
+
+  CTCBeamSearchDecoder(const CTCBeamSearchDecoder&) = delete;
+  void operator=(const CTCBeamSearchDecoder&) = delete;
+};
+
+template <typename T, typename CTCBeamState, typename CTCBeamComparer>
+absl::Status CTCBeamSearchDecoder<T, CTCBeamState, CTCBeamComparer>::Decode(
+    const typename CTCDecoder<T>::SequenceLength& seq_len,
+    const std::vector<typename CTCDecoder<T>::Input>& input,
+    std::vector<typename CTCDecoder<T>::Output>* output,
+    typename CTCDecoder<T>::ScoreOutput* scores) {
+  // Storage for top paths.
+  std::vector<std::vector<int>> beams;
+  std::vector<T> beam_log_probabilities;
+  int top_n = output->size();
+  if (std::any_of(output->begin(), output->end(),
+                  [this](const typename CTCDecoder<T>::Output& output) -> bool {
+                    return output.size() < this->batch_size_;
+                  })) {
+    return absl::InvalidArgumentError(
+        "output needs to be of size at least (top_n, batch_size).");
+  }
+  if (scores->rows() < this->batch_size_ || scores->cols() < top_n) {
+    return absl::InvalidArgumentError(
+        "scores needs to be of size at least (batch_size, top_n).");
+  }
+
+  for (int b = 0; b < this->batch_size_; ++b) {
+    int seq_len_b = seq_len[b];
+    Reset();
+
+    for (int t = 0; t < seq_len_b; ++t) {
+      // Pass log-probabilities for this example + time.
+      Step(input[t].row(b));
+    }  // for (int t...
+
+    // O(n * log(n))
+    std::unique_ptr<std::vector<BeamEntry*>> branches(leaves_.Extract());
+    leaves_.Reset();
+    for (int i = 0; i < branches->size(); ++i) {
+      BeamEntry* entry = (*branches)[i];
+      beam_scorer_->ExpandStateEnd(&entry->state);
+      entry->newp.total +=
+          beam_scorer_->GetStateEndExpansionScore(entry->state);
+      leaves_.push(entry);
+    }
+
+    absl::Status status =
+        TopPaths(top_n, &beams, &beam_log_probabilities, this->merge_repeated_);
+    if (!status.ok()) {
+      return status;
+    }
+
+    CHECK_EQ(top_n, beam_log_probabilities.size());
+    CHECK_EQ(beams.size(), beam_log_probabilities.size());
+
+    for (int i = 0; i < top_n; ++i) {
+      // Copy output to the correct beam + batch
+      (*output)[i][b].swap(beams[i]);
+      (*scores)(b, i) = -beam_log_probabilities[i];
+    }
+  }  // for (int b...
+  return absl::OkStatus();
+}
+
+template <typename T, typename CTCBeamState, typename CTCBeamComparer>
+template <typename Vector>
+T CTCBeamSearchDecoder<T, CTCBeamState, CTCBeamComparer>::GetTopK(
+    const int K, const Vector& input, std::vector<T>* top_k_logits,
+    std::vector<int>* top_k_indices) {
+  // Find Top K choices, complexity nk in worst case. The array input is read
+  // just once.
+  CHECK_EQ(this->num_classes_, input.size());
+  top_k_logits->clear();
+  top_k_indices->clear();
+  top_k_logits->resize(K, -INFINITY);
+  top_k_indices->resize(K, -1);
+  for (int j = 0; j < this->num_classes_ - 1; ++j) {
+    const T logit = input(j);
+    if (logit > (*top_k_logits)[K - 1]) {
+      int k = K - 1;
+      while (k > 0 && logit > (*top_k_logits)[k - 1]) {
+        (*top_k_logits)[k] = (*top_k_logits)[k - 1];
+        (*top_k_indices)[k] = (*top_k_indices)[k - 1];
+        k--;
+      }
+      (*top_k_logits)[k] = logit;
+      (*top_k_indices)[k] = j;
+    }
+  }
+  // Return max value which is in 0th index or blank character logit
+  return std::max((*top_k_logits)[0], input(this->num_classes_ - 1));
+}
+
+template <typename T, typename CTCBeamState, typename CTCBeamComparer>
+template <typename Vector>
+void CTCBeamSearchDecoder<T, CTCBeamState, CTCBeamComparer>::Step(
+    const Vector& raw_input) {
+  std::vector<T> top_k_logits;
+  std::vector<int> top_k_indices;
+  const bool top_k =
+      (label_selection_size_ > 0 && label_selection_size_ < raw_input.size());
+  // Number of character classes to consider in each step.
+  const int max_classes =
+      top_k ? label_selection_size_ : (this->num_classes_ - 1);
+  // Get max coefficient and remove it from raw_input later.
+  T max_coeff;
+  if (top_k) {
+    max_coeff = GetTopK(label_selection_size_, raw_input, &top_k_logits,
+                        &top_k_indices);
+  } else {
+    max_coeff = raw_input.maxCoeff();
+  }
+  // Get normalization term of softmax: log(sum(exp(logit[j]-max_coeff))).
+  T logsumexp = T(0.0);
+  for (int j = 0; j < raw_input.size(); ++j) {
+    logsumexp += Eigen::numext::exp(raw_input(j) - max_coeff);
+  }
+  logsumexp = Eigen::numext::log(logsumexp);
+  // Final normalization offset to get correct log probabilities.
+  T norm_offset = max_coeff + logsumexp;
+
+  const T label_selection_input_min =
+      (label_selection_margin_ >= 0) ? (max_coeff - label_selection_margin_)
+                                     : -std::numeric_limits<T>::infinity();
+
+  // Extract the beams sorted in decreasing new probability
+  CHECK_EQ(this->num_classes_, raw_input.size());
+
+  std::unique_ptr<std::vector<BeamEntry*>> branches(leaves_.Extract());
+  leaves_.Reset();
+
+  for (BeamEntry* b : *branches) {
+    // P(.. @ t) becomes the new P(.. @ t-1)
+    b->oldp = b->newp;
+  }
+
+  for (BeamEntry* b : *branches) {
+    if (b->parent != nullptr) {  // if not the root
+      if (b->parent->Active()) {
+        // If last two sequence characters are identical:
+        //   Plabel(l=acc @ t=6) = (Plabel(l=acc @ t=5)
+        //                          + Pblank(l=ac @ t=5))
+        // else:
+        //   Plabel(l=abc @ t=6) = (Plabel(l=abc @ t=5)
+        //                          + P(l=ab @ t=5))
+        T previous = (b->label == b->parent->label) ? b->parent->oldp.blank
+                                                    : b->parent->oldp.total;
+        b->newp.label =
+            LogSumExp(b->newp.label,
+                      beam_scorer_->GetStateExpansionScore(b->state, previous));
+      }
+      // Plabel(l=abc @ t=6) *= P(c @ 6)
+      b->newp.label += raw_input(b->label) - norm_offset;
+    }
+    // Pblank(l=abc @ t=6) = P(l=abc @ t=5) * P(- @ 6)
+    b->newp.blank = b->oldp.total + raw_input(this->blank_index_) - norm_offset;
+    // P(l=abc @ t=6) = Plabel(l=abc @ t=6) + Pblank(l=abc @ t=6)
+    b->newp.total = LogSumExp(b->newp.blank, b->newp.label);
+
+    // Push the entry back to the top paths list.
+    // Note, this will always fill leaves back up in sorted order.
+    leaves_.push(b);
+  }
+
+  // we need to resort branches in descending oldp order.
+
+  // branches is in descending oldp order because it was
+  // originally in descending newp order and we copied newp to oldp.
+
+  // Grow new leaves
+  for (BeamEntry* b : *branches) {
+    // A new leaf (represented by its BeamProbability) is a candidate
+    // iff its total probability is nonzero and either the beam list
+    // isn't full, or the lowest probability entry in the beam has a
+    // lower probability than the leaf.
+    auto is_candidate = [this](const BeamProbability& prob) {
+      return (prob.total > kLogZero<T>() &&
+              (leaves_.size() < beam_width_ ||
+               prob.total > leaves_.peek_bottom()->newp.total));
+    };
+
+    if (!is_candidate(b->oldp)) {
+      continue;
+    }
+
+    for (int ind = 0; ind < max_classes; ind++) {
+      const int label = top_k ? top_k_indices[ind] : ind;
+      const T logit = top_k ? top_k_logits[ind] : raw_input(ind);
+      // Perform label selection: if input for this label looks very
+      // unpromising, never evaluate it with a scorer.
+      // We may compare logits instead of log probabilities, 
+      // since the difference is the same in both cases.
+      if (logit < label_selection_input_min) {
+        continue;
+      }
+      BeamEntry& c = b->GetChild(label);
+      if (!c.Active()) {
+        //   Pblank(l=abcd @ t=6) = 0
+        c.newp.blank = kLogZero<T>();
+        // If new child label is identical to beam label:
+        //   Plabel(l=abcc @ t=6) = Pblank(l=abc @ t=5) * P(c @ 6)
+        // Otherwise:
+        //   Plabel(l=abcd @ t=6) = P(l=abc @ t=5) * P(d @ 6)
+        beam_scorer_->ExpandState(b->state, b->label, &c.state, c.label);
+        T previous = (c.label == b->label) ? b->oldp.blank : b->oldp.total;
+        c.newp.label = logit - norm_offset +
+                       beam_scorer_->GetStateExpansionScore(c.state, previous);
+        // P(l=abcd @ t=6) = Plabel(l=abcd @ t=6)
+        c.newp.total = c.newp.label;
+
+        if (is_candidate(c.newp)) {
+          // Before adding the new node to the beam, check if the beam
+          // is already at maximum width.
+          if (leaves_.size() == beam_width_) {
+            // Bottom is no longer in the beam search.  Reset
+            // its probability; signal it's no longer in the beam search.
+            BeamEntry* bottom = leaves_.peek_bottom();
+            bottom->newp.Reset();
+          }
+          leaves_.push(&c);
+        } else {
+          // Deactivate child.
+          c.oldp.Reset();
+          c.newp.Reset();
+        }
+      }
+    }
+  }  // for (BeamEntry* b...
+}
+
+template <typename T, typename CTCBeamState, typename CTCBeamComparer>
+void CTCBeamSearchDecoder<T, CTCBeamState, CTCBeamComparer>::Reset() {
+  leaves_.Reset();
+
+  // This beam root, and all of its children, will be in memory until
+  // the next reset.
+  beam_root_.reset(new BeamRoot(nullptr, -1));
+  beam_root_->RootEntry()->newp.total = T(0.0);  // ln(1)
+  beam_root_->RootEntry()->newp.blank = T(0.0);  // ln(1)
+
+  // Add the root as the initial leaf.
+  leaves_.push(beam_root_->RootEntry());
+
+  // Call initialize state on the root object.
+  beam_scorer_->InitializeState(&beam_root_->RootEntry()->state);
+}
+
+template <typename T, typename CTCBeamState, typename CTCBeamComparer>
+absl::Status CTCBeamSearchDecoder<T, CTCBeamState, CTCBeamComparer>::TopPaths(
+    int n, std::vector<std::vector<int>>* paths, std::vector<T>* log_probs,
+    bool merge_repeated) const {
+  CHECK_NOTNULL(paths)->clear();
+  CHECK_NOTNULL(log_probs)->clear();
+  if (n > beam_width_) {
+    return absl::InvalidArgumentError(
+        "requested more paths than the beam width.");
+  }
+  if (n > leaves_.size()) {
+    return absl::InvalidArgumentError(
+        "Less leaves in the beam search than requested.");
+  }
+
+  gtl::TopN<BeamEntry*, CTCBeamComparer> top_branches(n);
+
+  // O(beam_width_ * log(n)), space complexity is O(n)
+  for (auto it = leaves_.unsorted_begin(); it != leaves_.unsorted_end(); ++it) {
+    top_branches.push(*it);
+  }
+  // O(n * log(n))
+  std::unique_ptr<std::vector<BeamEntry*>> branches(top_branches.Extract());
+
+  for (int i = 0; i < n; ++i) {
+    BeamEntry* e((*branches)[i]);
+    paths->push_back(e->LabelSeq(merge_repeated));
+    log_probs->push_back(e->newp.total);
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace ctc
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_UTIL_CTC_CTC_BEAM_SEARCH_H_
+// LINT.ThenChange(//tensorflow/lite/kernels/ctc/ctc_beam_search.h)

tensorflow/core/util/ctc/ctc_beam_search_test.cc

@@ -0,0 +1,401 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// This test illustrates how to make use of the CTCBeamSearchDecoder using a
+// custom BeamScorer and BeamState based on a dictionary with a few artificial
+// words.
+#include "tensorflow/core/util/ctc/ctc_beam_search.h"
+
+#include <algorithm>
+#include <cmath>
+#include <vector>
+
+#include "absl/log/log.h"
+#include "tensorflow/core/lib/strings/strcat.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace {
+
+template <class T>
+using TestData = std::vector<std::vector<std::vector<T>>>;
+
+// The HistoryBeamState is used to keep track of the current candidate and
+// caches the expansion score (needed by the scorer).
+template <class T>
+struct HistoryBeamState {
+  T score;
+  std::vector<int> labels;
+};
+
+// DictionaryBeamScorer essentially favors candidates that can still become
+// dictionary words. As soon as a beam candidate is not a dictionary word or
+// a prefix of a dictionary word it gets a low probability at each step.
+//
+// The dictionary itself is hard-coded a static const variable of the class.
+template <class T, class BeamState>
+class DictionaryBeamScorer
+    : public tensorflow::ctc::BaseBeamScorer<T, BeamState> {
+ public:
+  DictionaryBeamScorer()
+      : tensorflow::ctc::BaseBeamScorer<T, BeamState>(),
+        dictionary_({{3}, {3, 1}}) {}
+
+  void InitializeState(BeamState* root) const override { root->score = 0; }
+
+  void ExpandState(const BeamState& from_state, int from_label,
+                   BeamState* to_state, int to_label) const override {
+    // Keep track of the current complete candidate by storing the labels along
+    // the expansion path in the beam state.
+    to_state->labels.push_back(to_label);
+    SetStateScoreAccordingToDict(to_state);
+  }
+
+  void ExpandStateEnd(BeamState* state) const override {
+    SetStateScoreAccordingToDict(state);
+  }
+
+  T GetStateExpansionScore(const BeamState& state,
+                           T previous_score) const override {
+    return previous_score + state.score;
+  }
+
+  T GetStateEndExpansionScore(const BeamState& state) const override {
+    return state.score;
+  }
+
+  // Simple dictionary used when scoring the beams to check if they are prefixes
+  // of dictionary words (see SetStateScoreAccordingToDict below).
+  const std::vector<std::vector<int>> dictionary_;
+
+ private:
+  void SetStateScoreAccordingToDict(BeamState* state) const;
+};
+
+template <class T, class BeamState>
+void DictionaryBeamScorer<T, BeamState>::SetStateScoreAccordingToDict(
+    BeamState* state) const {
+  // Check if the beam can still be a dictionary word (e.g. prefix of one).
+  const std::vector<int>& candidate = state->labels;
+  for (int w = 0; w < dictionary_.size(); ++w) {
+    const std::vector<int>& word = dictionary_[w];
+    // If the length of the current beam is already larger, skip.
+    if (candidate.size() > word.size()) {
+      continue;
+    }
+    if (std::equal(word.begin(), word.begin() + candidate.size(),
+                   candidate.begin())) {
+      state->score = std::log(T(1.0));
+      return;
+    }
+  }
+  // At this point, the candidate certainly can't be in the dictionary.
+  state->score = std::log(T(0.01));
+}
+
+template <class T>
+void ctc_beam_search_decoding_with_and_without_dictionary() {
+  const int batch_size = 1;
+  const int timesteps = 5;
+  const int top_paths = 3;
+  const int num_classes = 6;
+
+  // Plain decoder using hibernating beam search algorithm.
+  typename tensorflow::ctc::CTCBeamSearchDecoder<T>::DefaultBeamScorer
+      default_scorer;
+  tensorflow::ctc::CTCBeamSearchDecoder<T> decoder(num_classes, 10 * top_paths,
+                                                   &default_scorer);
+
+  // Dictionary decoder, allowing only two dictionary words : {3}, {3, 1}.
+  DictionaryBeamScorer<T, HistoryBeamState<T>> dictionary_scorer;
+  tensorflow::ctc::CTCBeamSearchDecoder<T, HistoryBeamState<T>>
+      dictionary_decoder(num_classes, top_paths, &dictionary_scorer);
+
+  // Raw data containers (arrays of floats64, ints, etc.).
+  int sequence_lengths[batch_size] = {timesteps};
+  T input_data_mat[timesteps][batch_size][num_classes] = {
+      {{0, 0.6, 0, 0.4, 0, 0}},
+      {{0, 0.5, 0, 0.5, 0, 0}},
+      {{0, 0.4, 0, 0.6, 0, 0}},
+      {{0, 0.4, 0, 0.6, 0, 0}},
+      {{0, 0.4, 0, 0.6, 0, 0}}};
+
+  // The CTCDecoder works with log-probs.
+  for (int t = 0; t < timesteps; ++t) {
+    for (int b = 0; b < batch_size; ++b) {
+      for (int c = 0; c < num_classes; ++c) {
+        input_data_mat[t][b][c] = std::log(input_data_mat[t][b][c]);
+      }
+    }
+  }
+
+  // Plain output, without any additional scoring.
+  std::vector<typename tensorflow::ctc::CTCDecoder<T>::Output> expected_output =
+      {
+          {{1, 3}, {1, 3, 1}, {3, 1, 3}},
+      };
+
+  // Dictionary outputs: preference for dictionary candidates. The
+  // second-candidate is there, despite it not being a dictionary word, due to
+  // stronger probability in the input to the decoder.
+  std::vector<typename tensorflow::ctc::CTCDecoder<T>::Output>
+      expected_dict_output = {
+          {{3}, {1, 3}, {3, 1}},
+      };
+
+  // Convert data containers to the format accepted by the decoder, simply
+  // mapping the memory from the container to an Eigen::ArrayXi,::MatrixXf,
+  // using Eigen::Map.
+  Eigen::Map<const Eigen::ArrayXi> seq_len(&sequence_lengths[0], batch_size);
+  std::vector<
+      Eigen::Map<const Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>>>
+      inputs;
+  inputs.reserve(timesteps);
+  for (int t = 0; t < timesteps; ++t) {
+    inputs.emplace_back(&input_data_mat[t][0][0], batch_size, num_classes);
+  }
+
+  // Prepare containers for output and scores.
+  std::vector<typename tensorflow::ctc::CTCDecoder<T>::Output> outputs(
+      top_paths);
+  for (typename tensorflow::ctc::CTCDecoder<T>::Output& output : outputs) {
+    output.resize(batch_size);
+  }
+  T score[batch_size][top_paths] = {{0.0}};
+  Eigen::Map<Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>> scores(
+      &score[0][0], batch_size, top_paths);
+
+  EXPECT_TRUE(decoder.Decode(seq_len, inputs, &outputs, &scores).ok());
+  for (int path = 0; path < top_paths; ++path) {
+    EXPECT_EQ(outputs[path][0], expected_output[0][path]);
+  }
+
+  // Prepare dictionary outputs.
+  std::vector<typename tensorflow::ctc::CTCDecoder<T>::Output> dict_outputs(
+      top_paths);
+  for (typename tensorflow::ctc::CTCDecoder<T>::Output& output : dict_outputs) {
+    output.resize(batch_size);
+  }
+  EXPECT_TRUE(
+      dictionary_decoder.Decode(seq_len, inputs, &dict_outputs, &scores).ok());
+  for (int path = 0; path < top_paths; ++path) {
+    EXPECT_EQ(dict_outputs[path][0], expected_dict_output[0][path]);
+  }
+}
+
+template <class T>
+void ctc_beam_search_decoding_all_beam_elements_have_finite_scores() {
+  const int batch_size = 1;
+  const int timesteps = 1;
+  const int top_paths = 3;
+  const int num_classes = 6;
+
+  // Plain decoder using hibernating beam search algorithm.
+  typename tensorflow::ctc::CTCBeamSearchDecoder<T>::DefaultBeamScorer
+      default_scorer;
+  tensorflow::ctc::CTCBeamSearchDecoder<T> decoder(num_classes, top_paths,
+                                                   &default_scorer);
+
+  // Raw data containers (arrays of floats64, ints, etc.).
+  int sequence_lengths[batch_size] = {timesteps};
+  T input_data_mat[timesteps][batch_size][num_classes] = {
+      {{0.4, 0.3, 0, 0, 0, 0.5}}};
+
+  // Convert data containers to the format accepted by the decoder, simply
+  // mapping the memory from the container to an Eigen::ArrayXi,::MatrixXf,
+  // using Eigen::Map.
+  Eigen::Map<const Eigen::ArrayXi> seq_len(&sequence_lengths[0], batch_size);
+  std::vector<
+      Eigen::Map<const Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>>>
+      inputs;
+  inputs.reserve(timesteps);
+  for (int t = 0; t < timesteps; ++t) {
+    inputs.emplace_back(&input_data_mat[t][0][0], batch_size, num_classes);
+  }
+
+  // Prepare containers for output and scores.
+  std::vector<typename tensorflow::ctc::CTCDecoder<T>::Output> outputs(
+      top_paths);
+  for (typename tensorflow::ctc::CTCDecoder<T>::Output& output : outputs) {
+    output.resize(batch_size);
+  }
+  T score[batch_size][top_paths] = {{0.0}};
+  Eigen::Map<Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>> scores(
+      &score[0][0], batch_size, top_paths);
+
+  EXPECT_TRUE(decoder.Decode(seq_len, inputs, &outputs, &scores).ok());
+  // Make sure all scores are finite.
+  for (int path = 0; path < top_paths; ++path) {
+    LOG(INFO) << "path " << path;
+    EXPECT_FALSE(std::isinf(score[0][path]));
+  }
+}
+
+// A beam decoder to test label selection. It simply models N labels with
+// rapidly dropping off log-probability.
+
+typedef int LabelState;  // The state is simply the final label.
+
+template <class T>
+class RapidlyDroppingLabelScorer
+    : public tensorflow::ctc::BaseBeamScorer<T, LabelState> {
+ public:
+  void InitializeState(LabelState* root) const override {}
+
+  void ExpandState(const LabelState& from_state, int from_label,
+                   LabelState* to_state, int to_label) const override {
+    *to_state = to_label;
+  }
+
+  void ExpandStateEnd(LabelState* state) const override {}
+
+  T GetStateExpansionScore(const LabelState& state,
+                           T previous_score) const override {
+    // Drop off rapidly for later labels.
+    const T kRapidly = 100;
+    return previous_score - kRapidly * state;
+  }
+
+  T GetStateEndExpansionScore(const LabelState& state) const override {
+    return T(0);
+  }
+};
+
+template <class T>
+void ctc_beam_search_label_selection() {
+  const int batch_size = 1;
+  const int timesteps = 3;
+  const int top_paths = 5;
+  const int num_classes = 6;
+
+  // Decoder which drops off log-probabilities for labels 0 >> 1 >> 2 >> 3.
+  RapidlyDroppingLabelScorer<T> scorer;
+  tensorflow::ctc::CTCBeamSearchDecoder<T, LabelState> decoder(
+      num_classes, top_paths, &scorer);
+
+  // Raw data containers (arrays of floats64, ints, etc.).
+  int sequence_lengths[batch_size] = {timesteps};
+  // Log probabilities, slightly preferring later labels, this decision
+  // should be overridden by the scorer which strongly prefers earlier labels.
+  // The last one is empty label, and for simplicity  we give it an extremely
+  // high cost to ignore it. We also use the first label to break up the
+  // repeated label sequence.
+  T input_data_mat[timesteps][batch_size][num_classes] = {
+      {{-1e6, 1, 2, 3, 4, -1e6}},
+      {{1e6, 0, 0, 0, 0, -1e6}},  // force label 0 to break up repeated
+      {{-1e6, 1.1, 2.2, 3.3, 4.4, -1e6}},
+  };
+
+  // Expected output without label selection
+  std::vector<typename tensorflow::ctc::CTCDecoder<T>::Output>
+      expected_default_output = {
+          {{1, 0, 1}, {1, 0, 2}, {2, 0, 1}, {1, 0, 3}, {2, 0, 2}},
+      };
+
+  // Expected output with label selection limiting to 2 items
+  // this is suboptimal because only labels 3 and 4 were allowed to be seen.
+  std::vector<typename tensorflow::ctc::CTCDecoder<T>::Output>
+      expected_output_size2 = {
+          {{3, 0, 3}, {3, 0, 4}, {4, 0, 3}, {4, 0, 4}, {3}},
+      };
+
+  // Expected output with label width of 2.0. This would permit three labels at
+  // the first timestep, but only two at the last.
+  std::vector<typename tensorflow::ctc::CTCDecoder<T>::Output>
+      expected_output_width2 = {
+          {{2, 0, 3}, {2, 0, 4}, {3, 0, 3}, {3, 0, 4}, {4, 0, 3}},
+      };
+
+  // Convert data containers to the format accepted by the decoder, simply
+  // mapping the memory from the container to an Eigen::ArrayXi,::MatrixXf,
+  // using Eigen::Map.
+  Eigen::Map<const Eigen::ArrayXi> seq_len(&sequence_lengths[0], batch_size);
+  std::vector<
+      Eigen::Map<const Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>>>
+      inputs;
+  inputs.reserve(timesteps);
+  for (int t = 0; t < timesteps; ++t) {
+    inputs.emplace_back(&input_data_mat[t][0][0], batch_size, num_classes);
+  }
+
+  // Prepare containers for output and scores.
+  std::vector<typename tensorflow::ctc::CTCDecoder<T>::Output> outputs(
+      top_paths);
+  for (typename tensorflow::ctc::CTCDecoder<T>::Output& output : outputs) {
+    output.resize(batch_size);
+  }
+  T score[batch_size][top_paths] = {{0.0}};
+  Eigen::Map<Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>> scores(
+      &score[0][0], batch_size, top_paths);
+
+  EXPECT_TRUE(decoder.Decode(seq_len, inputs, &outputs, &scores).ok());
+  for (int path = 0; path < top_paths; ++path) {
+    EXPECT_EQ(outputs[path][0], expected_default_output[0][path]);
+  }
+
+  // Try label selection size 2
+  decoder.SetLabelSelectionParameters(2, T(-1));
+  EXPECT_TRUE(decoder.Decode(seq_len, inputs, &outputs, &scores).ok());
+  for (int path = 0; path < top_paths; ++path) {
+    EXPECT_EQ(outputs[path][0], expected_output_size2[0][path]);
+  }
+
+  // Try label selection width 2.0
+  decoder.SetLabelSelectionParameters(0, T(2.0));
+  EXPECT_TRUE(decoder.Decode(seq_len, inputs, &outputs, &scores).ok());
+  for (int path = 0; path < top_paths; ++path) {
+    EXPECT_EQ(outputs[path][0], expected_output_width2[0][path]);
+  }
+
+  // Try both size 2 and width 2.0: the former is more constraining, so
+  // it's equivalent to that.
+  decoder.SetLabelSelectionParameters(2, T(2.0));
+  EXPECT_TRUE(decoder.Decode(seq_len, inputs, &outputs, &scores).ok());
+  for (int path = 0; path < top_paths; ++path) {
+    EXPECT_EQ(outputs[path][0], expected_output_size2[0][path]);
+  }
+
+  // Size 4 and width > 3.3 are equivalent to no label selection
+  decoder.SetLabelSelectionParameters(4, T(3.3001));
+  EXPECT_TRUE(decoder.Decode(seq_len, inputs, &outputs, &scores).ok());
+  for (int path = 0; path < top_paths; ++path) {
+    EXPECT_EQ(outputs[path][0], expected_default_output[0][path]);
+  }
+}
+
+TEST(CtcBeamSearch, FloatDecodingWithAndWithoutDictionary) {
+  ctc_beam_search_decoding_with_and_without_dictionary<float>();
+}
+
+TEST(CtcBeamSearch, DoubleDecodingWithAndWithoutDictionary) {
+  ctc_beam_search_decoding_with_and_without_dictionary<double>();
+}
+
+TEST(CtcBeamSearch, FloatAllBeamElementsHaveFiniteScores) {
+  ctc_beam_search_decoding_all_beam_elements_have_finite_scores<float>();
+}
+
+TEST(CtcBeamSearch, DoubleAllBeamElementsHaveFiniteScores) {
+  ctc_beam_search_decoding_all_beam_elements_have_finite_scores<double>();
+}
+
+TEST(CtcBeamSearch, FloatLabelSelection) {
+  ctc_beam_search_label_selection<float>();
+}
+
+TEST(CtcBeamSearch, DoubleLabelSelection) {
+  ctc_beam_search_label_selection<double>();
+}
+
+}  // namespace

tensorflow/core/util/ctc/ctc_decoder.h

@@ -0,0 +1,122 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+// LINT.IfChange
+
+#ifndef TENSORFLOW_CORE_UTIL_CTC_CTC_DECODER_H_
+#define TENSORFLOW_CORE_UTIL_CTC_CTC_DECODER_H_
+
+#include <memory>
+#include <vector>
+
+#include "Eigen/Core"  // from @eigen_archive
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/core/status.h"
+
+namespace tensorflow {
+namespace ctc {
+
+// The CTCDecoder is an abstract interface to be implemented when providing a
+// decoding method on the timestep output of a RNN trained with CTC loss.
+//
+// The two types of decoding available are:
+//   - greedy path, through the CTCGreedyDecoder
+//   - beam search, through the CTCBeamSearchDecoder
+template <class T>
+class CTCDecoder {
+ public:
+  typedef Eigen::Map<const Eigen::ArrayXi> SequenceLength;
+  typedef Eigen::Map<const Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>>
+      Input;
+  typedef std::vector<std::vector<int>> Output;
+  typedef Eigen::Map<Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>>
+      ScoreOutput;
+
+  CTCDecoder(int num_classes, int batch_size, bool merge_repeated)
+      : num_classes_(num_classes),
+        blank_index_(num_classes - 1),
+        batch_size_(batch_size),
+        merge_repeated_(merge_repeated) {}
+
+  virtual ~CTCDecoder() {}
+
+  // Dimensionality of the input/output is expected to be:
+  //  - seq_len[b] - b = 0 to batch_size_
+  //  - input[t].rows(b) - t = 0 to timesteps; b = 0 t batch_size_
+  //  - output.size() specifies the number of beams to be returned.
+  //  - scores(b, i) - b = 0 to batch_size; i = 0 to output.size()
+  virtual absl::Status Decode(const SequenceLength& seq_len,
+                              const std::vector<Input>& input,
+                              std::vector<Output>* output,
+                              ScoreOutput* scores) = 0;
+
+  int batch_size() { return batch_size_; }
+  int num_classes() { return num_classes_; }
+
+ protected:
+  int num_classes_;
+  int blank_index_;
+  int batch_size_;
+  bool merge_repeated_;
+};
+
+// CTCGreedyDecoder is an implementation of the simple best path decoding
+// algorithm, selecting at each timestep the most likely class at each timestep.
+template <class T>
+class CTCGreedyDecoder : public CTCDecoder<T> {
+ public:
+  typedef CTCDecoder<T> Decoder;
+  CTCGreedyDecoder(int num_classes, int batch_size, bool merge_repeated)
+      : CTCDecoder<T>(num_classes, batch_size, merge_repeated) {}
+
+  absl::Status Decode(const typename CTCDecoder<T>::SequenceLength& seq_len,
+                      const std::vector<typename CTCDecoder<T>::Input>& input,
+                      std::vector<typename CTCDecoder<T>::Output>* output,
+                      typename CTCDecoder<T>::ScoreOutput* scores) override {
+    if (output->empty() || (*output)[0].size() < Decoder::batch_size_) {
+      return absl::InvalidArgumentError(
+          "output needs to be of size at least (1, batch_size).");
+    }
+    if (scores->rows() < Decoder::batch_size_ || scores->cols() == 0) {
+      return absl::InvalidArgumentError(
+          "scores needs to be of size at least (batch_size, 1).");
+    }
+    // For each batch entry, identify the transitions
+    for (int b = 0; b < Decoder::batch_size_; ++b) {
+      int seq_len_b = seq_len[b];
+      // Only writing to beam 0
+      std::vector<int>& output_b = (*output)[0][b];
+
+      int prev_class_ix = -1;
+      (*scores)(b, 0) = 0;
+      for (int t = 0; t < seq_len_b; ++t) {
+        auto row = input[t].row(b);
+        int max_class_ix;
+        (*scores)(b, 0) += -row.maxCoeff(&max_class_ix);
+        if (max_class_ix != Decoder::blank_index_ &&
+            !(Decoder::merge_repeated_ && max_class_ix == prev_class_ix)) {
+          output_b.push_back(max_class_ix);
+        }
+        prev_class_ix = max_class_ix;
+      }
+    }
+    return absl::OkStatus();
+  }
+};
+
+}  // namespace ctc
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_UTIL_CTC_CTC_DECODER_H_
+// LINT.ThenChange(//tensorflow/lite/kernels/ctc/ctc_decoder.h)

tensorflow/core/util/ctc/ctc_loss_calculator.cc

@@ -0,0 +1,20 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/util/ctc/ctc_loss_calculator.h"
+
+namespace tensorflow {
+namespace ctc {}  // namespace ctc
+}  // namespace tensorflow

tensorflow/core/util/ctc/ctc_loss_calculator.h

@@ -0,0 +1,553 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_UTIL_CTC_CTC_LOSS_CALCULATOR_H_
+#define TENSORFLOW_CORE_UTIL_CTC_CTC_LOSS_CALCULATOR_H_
+
+#include <algorithm>
+#include <cstddef>
+#include <cstdint>
+#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/strings/str_join.h"
+#include "Eigen/Core"  // from @eigen_archive
+#include "tensorflow/core/framework/device_base.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/lib/strings/str_util.h"
+#include "tensorflow/core/lib/strings/strcat.h"
+#include "tensorflow/core/util/ctc/ctc_loss_util.h"
+#include "tensorflow/core/util/work_sharder.h"
+
+namespace tensorflow {
+namespace ctc {
+
+template <class T>
+class CTCLossCalculator {
+  // Connectionist Temporal Classification Loss
+  //
+  // Implementation by kanishkarao@, posenhuang@, and ebrevdo@.
+  //
+  // The CTC Loss layer learns a *transition* probability value for each
+  // input time step.  The transitions are on the class alphabet
+  //   {0, 1, ..., N-2}
+  // where N is the depth of the input layer (the size of the alphabet is N-1).
+  // Note: The token N-1 is reserved for the "no transition" output, so
+  // make sure that your input layer has a depth that's one larger than
+  // the set of classes you're training on.  Also make sure that your
+  // training labels do not have a class value of N-1, as training will skip
+  // these examples.
+  //
+  // Reference materials:
+  //  GravesTh: Alex Graves, "Supervised Sequence Labeling with Recurrent
+  //    Neural Networks" (PhD Thesis), Technische Universit¨at M¨unchen.
+ public:
+  typedef std::vector<std::vector<int>> LabelSequences;
+  using Matrix = Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>;
+  // typedef Eigen::MatrixXd Matrix;
+  using Array = Eigen::Array<T, Eigen::Dynamic, 1>;
+  // typedef Eigen::ArrayXd Array;
+  using InputMap = Eigen::Map<const Matrix>;
+  // typedef Eigen::Map<const Eigen::MatrixXd> InputMap;
+  using OutputMap = Eigen::Map<Matrix>;
+  // typedef Eigen::Map<Eigen::MatrixXd> OutputMap;
+
+  CTCLossCalculator(int blank_index, int output_delay)
+      : blank_index_(blank_index), output_delay_(output_delay) {}
+
+  template <typename VectorIn, typename VectorOut, typename MatrixIn,
+            typename MatrixOut>
+  absl::Status CalculateLoss(
+      const VectorIn& seq_len, const LabelSequences& labels,
+      const std::vector<MatrixIn>& inputs, bool preprocess_collapse_repeated,
+      bool ctc_merge_repeated, bool ignore_longer_outputs_than_inputs,
+      VectorOut* loss, std::vector<MatrixOut>* gradients,
+      DeviceBase::CpuWorkerThreads* workers = nullptr) const;
+
+ private:
+  void CalculateForwardVariables(const std::vector<int>& l_prime,
+                                 const Matrix& y, bool ctc_merge_repeated,
+                                 Matrix* log_alpha) const;
+
+  void CalculateBackwardVariables(const std::vector<int>& l_prime,
+                                  const Matrix& y, bool ctc_merge_repeated,
+                                  Matrix* log_beta) const;
+
+  void CalculateGradient(const std::vector<int>& l_prime, const Matrix& y,
+                         const Matrix& log_alpha, const Matrix& log_beta,
+                         T log_p_z_x, Matrix* dy) const;
+
+  void GetLPrimeIndices(const std::vector<int>& l,
+                        std::vector<int>* l_prime) const;
+
+  // Helper function that calculates the l_prime indices for all
+  // batches at the same time, and identifies errors for any given
+  // batch.  Return value:
+  //    max_{b in batch_size} l_primes[b].size()
+  template <typename Vector>
+  absl::Status PopulateLPrimes(bool preprocess_collapse_repeated,
+                               bool ignore_longer_outputs_than_inputs,
+                               int batch_size, int num_classes,
+                               const Vector& seq_len,
+                               const LabelSequences& labels,
+                               size_t* max_u_prime,
+                               LabelSequences* l_primes) const;
+
+  // Utility indices for the CTC algorithm.
+  int blank_index_;
+
+  // Delay for target labels in time steps.
+  // The delay in time steps before the output sequence.
+  const int output_delay_;
+};
+
+template <class T>
+template <typename VectorIn, typename VectorOut, typename MatrixIn,
+          typename MatrixOut>
+absl::Status CTCLossCalculator<T>::CalculateLoss(
+    const VectorIn& seq_len, const LabelSequences& labels,
+    const std::vector<MatrixIn>& inputs, bool preprocess_collapse_repeated,
+    bool ctc_merge_repeated, bool ignore_longer_outputs_than_inputs,
+    VectorOut* loss, std::vector<MatrixOut>* gradients,
+    DeviceBase::CpuWorkerThreads* workers) const {
+  using Eigen::numext::log;
+
+  auto num_time_steps = inputs.size();
+
+  if (loss == nullptr) {
+    return absl::InvalidArgumentError("loss == nullptr");
+  }
+
+  bool requires_backprop = (gradients != nullptr);
+
+  auto batch_size = inputs[0].rows();
+  auto num_classes = inputs[0].cols();
+
+  if (loss->size() != batch_size) {
+    return absl::InvalidArgumentError("loss.size() != batch_size");
+  }
+  loss->setZero();
+
+  for (int t = 1; t < num_time_steps; ++t) {
+    if (inputs[t].rows() != batch_size) {
+      return errors::InvalidArgument("Expected batch size at t: ", t,
+                                     " to be: ", batch_size,
+                                     " but got: ", inputs[t].rows());
+    }
+    if (inputs[t].cols() != num_classes) {
+      return errors::InvalidArgument("Expected class count at t: ", t,
+                                     " to be: ", num_classes,
+                                     " but got: ", inputs[t].cols());
+    }
+  }
+
+  // Check validity of sequence_length array values.
+  auto max_seq_len = seq_len(0);
+  for (int b = 0; b < batch_size; b++) {
+    if (seq_len(b) < 0) {
+      return absl::InvalidArgumentError(absl::StrCat("seq_len(", b, ") < 0"));
+    }
+    if (seq_len(b) > num_time_steps) {
+      return absl::InvalidArgumentError(
+          absl::StrCat("seq_len(", b, ") > num_time_steps"));
+    }
+    max_seq_len = std::max(seq_len(b), max_seq_len);
+  }
+
+  // Calculate the modified label sequence l' for each batch element,
+  // and calculate the maximum necessary allocation size.
+  LabelSequences l_primes(batch_size);
+  size_t max_u_prime = 0;
+  absl::Status l_p_ret = PopulateLPrimes(
+      preprocess_collapse_repeated, ignore_longer_outputs_than_inputs,
+      batch_size, num_classes, seq_len, labels, &max_u_prime, &l_primes);
+  if (!l_p_ret.ok()) {
+    return l_p_ret;
+  }
+
+  // Process each item in a batch in parallel, using at most kMaxThreads.
+  auto ComputeLossAndGradients = [this, num_classes, &labels, &l_primes,
+                                  &seq_len, &inputs, requires_backprop,
+                                  ctc_merge_repeated,
+                                  ignore_longer_outputs_than_inputs, &loss,
+                                  &gradients](int64_t start_row,
+                                              int64_t limit_row) {
+    for (int b = start_row; b < limit_row; b++) {
+      // Return zero gradient for empty sequences or sequences with labels
+      // longer than input, which is not supported by CTC.
+      if (seq_len(b) == 0 ||
+          (ignore_longer_outputs_than_inputs &&
+           labels[b].size() > seq_len(b) - this->output_delay_)) {
+        VLOG(1) << "The sequence length is either zero or shorter than the "
+                   "target output (CTC works only with shorter target sequence "
+                   "than input sequence). You can turn this into a warning by "
+                   "using the flag ignore_longer_outputs_than_inputs - "
+                << b << ": " << absl::StrJoin(labels[b], " ");
+        continue;
+      }
+
+      // For each batch element, log(alpha) and log(beta).
+      //   row size is: u_prime == l_prime.size()
+      //   col size is: seq_len[b] - output_delay_
+      const std::vector<int>& l_prime = l_primes[b];
+
+      Matrix log_alpha_b(l_prime.size(), seq_len(b) - this->output_delay_);
+      Matrix log_beta_b(l_prime.size(), seq_len(b) - this->output_delay_);
+
+      // Work matrices, pre-allocated to the size required by this batch item.
+      Matrix y(num_classes, seq_len(b));
+      Matrix dy;
+      if (requires_backprop) {
+        dy = Matrix::Zero(y.rows(), y.cols());
+      }
+
+      // For this batch, we'll only work with this shortened sequence_length.
+      Matrix y_b = y.leftCols(seq_len(b));
+
+      // Convert label from DistBelief
+      // y, prob are in num_classes x seq_len(b)
+      // Output activations.
+      Array y_b_col;
+      for (int t = 0; t < seq_len(b); t++) {
+        // Calculate the softmax of y_b.  Use original precision
+        // arithmetic for the sum.
+        T max_coeff = inputs[t].row(b).maxCoeff();
+        y_b_col = (inputs[t].row(b).array() - max_coeff).exp();
+        y_b.col(t) = y_b_col / y_b_col.sum();
+      }
+
+      // Compute forward, backward.
+      // Forward variables.
+      CalculateForwardVariables(l_prime, y_b, ctc_merge_repeated, &log_alpha_b);
+      // Backward variables.
+      CalculateBackwardVariables(l_prime, y_b, ctc_merge_repeated, &log_beta_b);
+
+      // The loss is computed as the log(p(z|x)) between the target and
+      // prediction. Do lazy evaluation of log_prob here.
+      T log_p_z_x = kLogZero<T>();
+      for (int u = 0; u < l_prime.size(); ++u) {
+        // (GravesTh) Eq 7.26, sum over all paths for t = 0.
+        log_p_z_x = LogSumExp(log_p_z_x, log_alpha_b(u, 0) + log_beta_b(u, 0));
+      }
+
+      (*loss)(b) = -log_p_z_x;  // Use negative log loss for display.
+
+      // We compute the derivative if needed.
+      if (requires_backprop) {
+        // Gradients with respect to input activations.
+        // Calculate gradient.
+        dy.setZero();
+        CalculateGradient(l_prime, y_b, log_alpha_b, log_beta_b, log_p_z_x,
+                          &dy);
+
+        // Convert gradient for current sample to DistBelief.
+        for (int t = 0; t < seq_len(b); t++) {
+          (*gradients)[t].row(b).array() = dy.col(t);
+        }
+      }
+    }  // for (int b = ...
+  };
+  if (workers) {
+    // *Rough* estimate of the cost for one item in the batch.
+    // Forward, Backward: O(T * U (= 2L + 1)), Gradients: O(T * (U + L)).
+    //
+    // softmax: T * L * (Cost(Exp) + Cost(Div))softmax +
+    // fwd,bwd: T * 2 * (2*L + 1) * (Cost(LogSumExp) + Cost(Log)) +
+    // grad: T * ((2L + 1) * Cost(LogSumExp) + L * (Cost(Expf) + Cost(Add)).
+    const int64_t cost_exp = Eigen::internal::functor_traits<
+        Eigen::internal::scalar_exp_op<T>>::Cost;
+    const int64_t cost_log = Eigen::internal::functor_traits<
+        Eigen::internal::scalar_log_op<T>>::Cost;
+    const int64_t cost_log_sum_exp =
+        Eigen::TensorOpCost::AddCost<T>() + cost_exp + cost_log;
+    const int64_t cost =
+        max_seq_len * num_classes *
+            (cost_exp + Eigen::TensorOpCost::DivCost<T>()) +
+        max_seq_len * 2 * (2 * num_classes + 1) *
+            (cost_log_sum_exp + cost_log) +
+        max_seq_len *
+            ((2 * num_classes + 1) * cost_log_sum_exp +
+             num_classes * (cost_exp + Eigen::TensorOpCost::AddCost<T>()));
+    Shard(workers->num_threads, workers->workers, batch_size, cost,
+          ComputeLossAndGradients);
+  } else {
+    ComputeLossAndGradients(0, batch_size);
+  }
+  return absl::OkStatus();
+}
+
+template <class T>
+template <typename Vector>
+absl::Status CTCLossCalculator<T>::PopulateLPrimes(
+    bool preprocess_collapse_repeated, bool ignore_longer_outputs_than_inputs,
+    int batch_size, int num_classes, const Vector& seq_len,
+    const LabelSequences& labels, size_t* max_u_prime,
+    LabelSequences* l_primes) const {
+  // labels is a Label array of size batch_size
+  if (labels.size() != batch_size) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "labels.size() != batch_size: ", labels.size(), " vs. ", batch_size));
+  }
+
+  *max_u_prime = 0;  // keep track of longest l' modified label sequence.
+  for (int b = 0; b < batch_size; b++) {
+    // Assume label is in Label proto
+    const std::vector<int>& label = labels[b];
+    if (label.size() == 0) {
+      return absl::InvalidArgumentError(
+          absl::StrCat("Labels length is zero in batch ", b));
+    }
+
+    // If debugging: output the labels coming into training.
+    //
+    VLOG(2) << "label for batch: " << b << ": " << absl::StrJoin(label, " ");
+
+    // Target indices, length = U.
+    std::vector<int> l;
+
+    // Convert label from DistBelief
+    bool finished_sequence = false;
+    for (int i = 0; i < label.size(); ++i) {
+      if (i == 0 || !preprocess_collapse_repeated || label[i] != label[i - 1]) {
+        if (label[i] >= num_classes - 1) {
+          finished_sequence = true;
+        } else {
+          if (finished_sequence) {
+            // Saw an invalid sequence with non-null following null
+            // labels.
+            return absl::InvalidArgumentError(absl::StrCat(
+                "Saw a non-null label (index >= num_classes - 1) "
+                "following a ",
+                "null label, batch: ", b, " num_classes: ", num_classes,
+                " labels: ", absl::StrJoin(label, ","),
+                " labels seen so far: ", absl::StrJoin(l, ",")));
+          }
+          l.push_back(label[i]);
+        }
+      }
+    }
+
+    for (int l_i : l) {
+      if (l_i < 0) {
+        return absl::InvalidArgumentError(
+            absl::StrCat("All labels must be nonnegative integers, batch: ", b,
+                         " labels: ", absl::StrJoin(l, ",")));
+      } else if (l_i >= num_classes) {
+        return absl::InvalidArgumentError(absl::StrCat(
+            "No label may be greater than num_classes. ", "num_classes: ",
+            num_classes, ", batch: ", b, " labels: ", absl::StrJoin(l, ",")));
+      }
+    }
+    if (!ignore_longer_outputs_than_inputs) {
+      // Make sure there is enough time to output the target indices.
+      int time = seq_len(b) - output_delay_;
+      int required_time = label.size();
+      if (required_time > time) {
+        return absl::InvalidArgumentError(absl::StrCat(
+            "Not enough time for target transition sequence ("
+            "required: ",
+            required_time, ", available: ", time, ")", b,
+            "You can turn this error into a warning by using the flag "
+            "ignore_longer_outputs_than_inputs"));
+      }
+    }
+    // Target indices with blanks before each index and a blank at the end.
+    // Length U' = 2U + 1.
+    // Convert l to l_prime
+    GetLPrimeIndices(l, &l_primes->at(b));
+    *max_u_prime = std::max(*max_u_prime, l_primes->at(b).size());
+  }
+  return absl::OkStatus();
+}
+
+// Calculates the alpha(t, u) as described in (GravesTh) Section 7.3.
+// Starting with t = 0 instead of t = 1 used in the text.
+// Based on Kanishka's CTC.
+template <typename TT>
+void CTCLossCalculator<TT>::CalculateForwardVariables(
+    const std::vector<int>& l_prime, const Matrix& y, bool ctc_merge_repeated,
+    Matrix* log_alpha) const {
+  using Eigen::numext::log;
+
+  // Number of cols is the number of time steps = number of cols in target
+  // after the output delay.
+  log_alpha->setConstant(kLogZero<TT>());
+
+  int U = l_prime.size();
+  int T = log_alpha->cols();
+
+  CHECK_EQ(U, log_alpha->rows());
+
+  // Initial alpha values in (GravesTh) Eq 7.5 and Eq 7.6.
+  log_alpha->coeffRef(0, 0) = log(y(blank_index_, output_delay_));
+  // Below, l_prime[1] == labels[0]
+  auto label_0 = (l_prime.size() > 1) ? l_prime[1] : blank_index_;
+  log_alpha->coeffRef(1, 0) = log(y(label_0, output_delay_));
+
+  for (int t = 1; t < T; ++t) {
+    // If there is not enough time to output the remaining labels or
+    // some labels have been skipped, then let log_alpha(u, t) continue to
+    // be kLogZero.
+    for (int u = std::max(0, U - (2 * (T - t))); u < std::min(U, 2 * (t + 1));
+         ++u) {
+      // Begin (GravesTh) Eq 7.9
+      // Add in the u, t - 1 term.
+      auto sum_log_alpha = kLogZero<TT>();
+      if (ctc_merge_repeated || l_prime[u] == blank_index_) {
+        sum_log_alpha = log_alpha->coeff(u, t - 1);
+      }
+
+      // Add in the u - 1, t - 1 term.
+      if (u > 0) {
+        sum_log_alpha =
+            LogSumExp(sum_log_alpha, log_alpha->coeff(u - 1, t - 1));
+      }
+
+      // Add in the u - 2, t - 1 term if l_prime(u) != blank or l_prime(u-2).
+      if (u > 1) {
+        const bool matching_labels_merge =
+            ctc_merge_repeated && (l_prime[u] == l_prime[u - 2]);
+        if (l_prime[u] != blank_index_ && !matching_labels_merge) {
+          sum_log_alpha =
+              LogSumExp(sum_log_alpha, log_alpha->coeff(u - 2, t - 1));
+        }
+      }
+      // Multiply the summed alphas with the activation log probability.
+      log_alpha->coeffRef(u, t) =
+          log(y(l_prime[u], output_delay_ + t)) + sum_log_alpha;
+    }  // End (GravesTh) Eq 7.9.
+  }
+}
+
+// Calculates the beta(t, u) as described in (GravesTh) Section 7.3.
+template <class TT>
+void CTCLossCalculator<TT>::CalculateBackwardVariables(
+    const std::vector<int>& l_prime, const Matrix& y, bool ctc_merge_repeated,
+    Matrix* log_beta) const {
+  // Number of cols is the number of time steps =  number of cols in target.
+  // Matrix log_beta =
+  //    Matrix::Constant(l_prime.size(), y.cols() - output_delay_,
+  // kLogZero);
+  using Eigen::numext::log;
+
+  log_beta->setConstant(kLogZero<TT>());
+  int T = log_beta->cols();
+  int U = l_prime.size();
+  CHECK_EQ(U, log_beta->rows());
+
+  // Initial beta values in (GravesTh) Eq 7.13: log of probability 1.
+  for (int u = U - 2; u < U; ++u) log_beta->coeffRef(u, T - 1) = 0;
+
+  for (int t = T - 1 - 1; t >= 0; --t) {
+    // If there is not enough time to output the remaining labels or
+    // some labels have been skipped, then let log_beta(u, t) continue to
+    // be kLogZero.
+    for (int u = std::max(0, U - (2 * (T - t))); u < std::min(U, 2 * (t + 1));
+         ++u) {
+      // Begin (GravesTh) Eq 7.15
+      // Add in the u, t + 1 term.
+      if (ctc_merge_repeated || l_prime[u] == blank_index_) {
+        log_beta->coeffRef(u, t) =
+            LogSumExp(log_beta->coeff(u, t),
+                      log_beta->coeff(u, t + 1) +
+                          log(y(l_prime[u], output_delay_ + t + 1)));
+      }
+
+      // Add in the u + 1, t + 1 term.
+      if (u + 1 < U) {
+        log_beta->coeffRef(u, t) =
+            LogSumExp(log_beta->coeff(u, t),
+                      log_beta->coeff(u + 1, t + 1) +
+                          log(y(l_prime[u + 1], output_delay_ + t + 1)));
+      }
+
+      // Add in the u + 2, t + 1 term if l_prime(u) != blank or l_prime(u+2).
+      if (u + 2 < U) {
+        const bool matching_labels_merge =
+            ctc_merge_repeated && (l_prime[u] == l_prime[u + 2]);
+        if (l_prime[u] != blank_index_ && !matching_labels_merge) {
+          // Add in u + 2 term.
+          log_beta->coeffRef(u, t) =
+              LogSumExp(log_beta->coeff(u, t),
+                        log_beta->coeff(u + 2, t + 1) +
+                            log(y(l_prime[u + 2], output_delay_ + t + 1)));
+        }
+      }  // End (GravesTh) Eq. 7.15
+    }
+  }
+}
+
+// Using (GravesTh) Eq 7.26 & 7.34.
+template <typename TT>
+void CTCLossCalculator<TT>::CalculateGradient(const std::vector<int>& l_prime,
+                                              const Matrix& y,
+                                              const Matrix& log_alpha,
+                                              const Matrix& log_beta,
+                                              TT log_p_z_x, Matrix* dy) const {
+  // Only working with the leftmost part of dy for this batch element.
+  auto dy_b = dy->leftCols(y.cols());
+
+  // It is possible that no valid path is found if the activations for the
+  // targets are zero.
+  if (log_p_z_x == kLogZero<TT>()) {
+    LOG(WARNING) << "No valid path found.";
+    dy_b = y;
+    return;
+  }
+
+  int L = y.rows();
+  int T = y.cols();
+  int U = l_prime.size();
+
+  for (int t = 0; t < T - output_delay_; ++t) {
+    Array prob_sum(L);
+    prob_sum.setConstant(kLogZero<TT>());
+
+    for (int u = 0; u < U; ++u) {
+      int l = l_prime[u];
+      prob_sum[l] = LogSumExp(prob_sum[l], log_alpha(u, t) + log_beta(u, t));
+    }
+
+    for (int l = 0; l < L; ++l) {
+      // Negative term in (GravesTh) Eq 7.28.
+      auto negative_term = expf(prob_sum[l] - log_p_z_x);
+
+      dy_b(l, output_delay_ + t) = y(l, output_delay_ + t) - negative_term;
+    }
+  }
+}
+
+template <class TT>
+void CTCLossCalculator<TT>::GetLPrimeIndices(const std::vector<int>& l,
+                                             std::vector<int>* l_prime) const {
+  // Assumption is that l_prime is empty.
+  l_prime->reserve(2 * l.size() + 1);
+
+  for (auto label : l) {
+    l_prime->push_back(blank_index_);
+    l_prime->push_back(label);
+  }
+  // Add final blank to l'.
+  l_prime->push_back(blank_index_);
+}
+
+}  // namespace ctc
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_UTIL_CTC_CTC_LOSS_CALCULATOR_H_

tensorflow/core/util/ctc/ctc_loss_util.h

@@ -0,0 +1,55 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+// LINT.IfChange
+
+#ifndef TENSORFLOW_CORE_UTIL_CTC_CTC_LOSS_UTIL_H_
+#define TENSORFLOW_CORE_UTIL_CTC_CTC_LOSS_UTIL_H_
+
+#include <cmath>
+#include <limits>
+
+namespace tensorflow {
+namespace ctc {
+
+template <class T>
+constexpr T kLogZero() {
+  return -std::numeric_limits<T>::infinity();  // NOLINT
+}
+
+// Add logarithmic probabilities using:
+// ln(a + b) = ln(a) + ln(1 + exp(ln(b) - ln(a)))
+// The two inputs are assumed to be log probabilities.
+// (GravesTh) Eq. 7.18
+template <typename T>
+inline T LogSumExp(T log_prob_1, T log_prob_2) {
+  // const T kLogZero = -std::numeric_limits<T>::infinity();
+  // Always have 'b' be the smaller number to avoid the exponential from
+  // blowing up.
+  if (log_prob_1 == kLogZero<T>()) {
+    return log_prob_2;
+  } else if (log_prob_2 == kLogZero<T>()) {
+    return log_prob_1;
+  } else {
+    return (log_prob_1 > log_prob_2)
+               ? log_prob_1 + log1pf(expf(log_prob_2 - log_prob_1))
+               : log_prob_2 + log1pf(expf(log_prob_1 - log_prob_2));
+  }
+}
+
+}  // namespace ctc
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_UTIL_CTC_CTC_LOSS_UTIL_H_
+// LINT.ThenChange(//tensorflow/lite/kernels/ctc/ctc_loss_util.h)

tensorflow/core/util/cuda_sparse.h

@@ -0,0 +1,728 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_UTIL_CUDA_SPARSE_H_
+#define TENSORFLOW_CORE_UTIL_CUDA_SPARSE_H_
+
+// This header declares the class GpuSparse, which contains wrappers of
+// cuSparse libraries for use in TensorFlow kernels.
+
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+
+#include <functional>
+#include <vector>
+
+#if GOOGLE_CUDA
+
+#include "third_party/gpus/cuda/include/cuda.h"
+#include "third_party/gpus/cuda/include/cusparse.h"
+
+using gpusparseStatus_t = cusparseStatus_t;
+using gpusparseOperation_t = cusparseOperation_t;
+using gpusparseMatDescr_t = cusparseMatDescr_t;
+using gpusparseAction_t = cusparseAction_t;
+using gpusparseHandle_t = cusparseHandle_t;
+using gpuStream_t = cudaStream_t;
+#if CUDA_VERSION >= 10020
+using gpusparseDnMatDescr_t = cusparseDnMatDescr_t;
+using gpusparseSpMatDescr_t = cusparseSpMatDescr_t;
+using gpusparseSpMMAlg_t = cusparseSpMMAlg_t;
+#endif
+
+#define GPUSPARSE(postfix) CUSPARSE_##postfix
+#define gpusparse(postfix) cusparse##postfix
+
+#elif TENSORFLOW_USE_ROCM
+
+#include "rocm/rocm_config.h"
+#include "xla/stream_executor/rocm/hipsparse_wrapper.h"
+
+using gpusparseStatus_t = hipsparseStatus_t;
+using gpusparseOperation_t = hipsparseOperation_t;
+using gpusparseMatDescr_t = hipsparseMatDescr_t;
+using gpusparseAction_t = hipsparseAction_t;
+using gpusparseHandle_t = hipsparseHandle_t;
+using gpuStream_t = hipStream_t;
+#if TF_ROCM_VERSION >= 40200
+using gpusparseDnMatDescr_t = hipsparseDnMatDescr_t;
+using gpusparseSpMatDescr_t = hipsparseSpMatDescr_t;
+using gpusparseSpMMAlg_t = hipsparseSpMMAlg_t;
+#endif
+#define GPUSPARSE(postfix) HIPSPARSE_##postfix
+#define gpusparse(postfix) hipsparse##postfix
+
+#endif
+
+#include "xla/stream_executor/data_type.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_types.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/platform/stream_executor.h"
+#include "tensorflow/core/public/version.h"
+
+#if GOOGLE_CUDA
+#include "xla/stream_executor/cuda/cuda_blas_utils.h"
+#endif
+
+// Macro that specializes a sparse method for all 4 standard
+// numeric types.
+// TODO: reuse with cuda_solvers
+#define TF_CALL_LAPACK_TYPES(m) \
+  m(float, S) m(double, D) m(std::complex<float>, C) m(std::complex<double>, Z)
+
+namespace tensorflow {
+
+inline std::string ConvertGPUSparseErrorToString(
+    const gpusparseStatus_t status) {
+  switch (status) {
+#define STRINGIZE(q) #q
+#define RETURN_IF_STATUS(err) \
+  case err:                   \
+    return STRINGIZE(err);
+
+#if GOOGLE_CUDA
+
+    RETURN_IF_STATUS(CUSPARSE_STATUS_SUCCESS)
+    RETURN_IF_STATUS(CUSPARSE_STATUS_NOT_INITIALIZED)
+    RETURN_IF_STATUS(CUSPARSE_STATUS_ALLOC_FAILED)
+    RETURN_IF_STATUS(CUSPARSE_STATUS_INVALID_VALUE)
+    RETURN_IF_STATUS(CUSPARSE_STATUS_ARCH_MISMATCH)
+    RETURN_IF_STATUS(CUSPARSE_STATUS_MAPPING_ERROR)
+    RETURN_IF_STATUS(CUSPARSE_STATUS_EXECUTION_FAILED)
+    RETURN_IF_STATUS(CUSPARSE_STATUS_INTERNAL_ERROR)
+    RETURN_IF_STATUS(CUSPARSE_STATUS_MATRIX_TYPE_NOT_SUPPORTED)
+
+    default:
+      return absl::StrCat("Unknown CUSPARSE error: ", static_cast<int>(status));
+#elif TENSORFLOW_USE_ROCM
+
+    RETURN_IF_STATUS(HIPSPARSE_STATUS_SUCCESS)
+    RETURN_IF_STATUS(HIPSPARSE_STATUS_NOT_INITIALIZED)
+    RETURN_IF_STATUS(HIPSPARSE_STATUS_ALLOC_FAILED)
+    RETURN_IF_STATUS(HIPSPARSE_STATUS_INVALID_VALUE)
+    RETURN_IF_STATUS(HIPSPARSE_STATUS_ARCH_MISMATCH)
+    RETURN_IF_STATUS(HIPSPARSE_STATUS_MAPPING_ERROR)
+    RETURN_IF_STATUS(HIPSPARSE_STATUS_EXECUTION_FAILED)
+    RETURN_IF_STATUS(HIPSPARSE_STATUS_INTERNAL_ERROR)
+    RETURN_IF_STATUS(HIPSPARSE_STATUS_MATRIX_TYPE_NOT_SUPPORTED)
+    RETURN_IF_STATUS(HIPSPARSE_STATUS_ZERO_PIVOT)
+
+    default:
+      return strings::StrCat("Unknown hipSPARSE error: ",
+                             static_cast<int>(status));
+#endif
+
+#undef RETURN_IF_STATUS
+#undef STRINGIZE
+  }
+}
+
+#if GOOGLE_CUDA
+
+#define TF_RETURN_IF_GPUSPARSE_ERROR(expr)                                 \
+  do {                                                                     \
+    auto status = (expr);                                                  \
+    if (TF_PREDICT_FALSE(status != CUSPARSE_STATUS_SUCCESS)) {             \
+      return errors::Internal(__FILE__, ":", __LINE__, " (", TF_STR(expr), \
+                              "): cuSparse call failed with status ",      \
+                              ConvertGPUSparseErrorToString(status));      \
+    }                                                                      \
+  } while (0)
+
+#elif TENSORFLOW_USE_ROCM
+
+#define TF_RETURN_IF_GPUSPARSE_ERROR(expr)                                 \
+  do {                                                                     \
+    auto status = (expr);                                                  \
+    if (TF_PREDICT_FALSE(status != HIPSPARSE_STATUS_SUCCESS)) {            \
+      return errors::Internal(__FILE__, ":", __LINE__, " (", TF_STR(expr), \
+                              "): hipSPARSE call failed with status ",     \
+                              ConvertGPUSparseErrorToString(status));      \
+    }                                                                      \
+  } while (0)
+
+#endif
+
+inline gpusparseOperation_t TransposeAndConjugateToGpuSparseOp(
+    bool transpose, bool conjugate, absl::Status* status) {
+#if GOOGLE_CUDA
+  if (transpose) {
+    return conjugate ? CUSPARSE_OPERATION_CONJUGATE_TRANSPOSE
+                     : CUSPARSE_OPERATION_TRANSPOSE;
+  } else {
+    if (conjugate) {
+      DCHECK(status != nullptr);
+      *status = absl::InvalidArgumentError(
+          "Conjugate == True and transpose == False is not supported.");
+    }
+    return CUSPARSE_OPERATION_NON_TRANSPOSE;
+  }
+#elif TENSORFLOW_USE_ROCM
+  if (transpose) {
+    return conjugate ? HIPSPARSE_OPERATION_CONJUGATE_TRANSPOSE
+                     : HIPSPARSE_OPERATION_TRANSPOSE;
+  } else {
+    if (conjugate) {
+      DCHECK(status != nullptr);
+      *status = errors::InvalidArgument(
+          "Conjugate == True and transpose == False is not supported.");
+    }
+    return HIPSPARSE_OPERATION_NON_TRANSPOSE;
+  }
+#endif
+}
+
+#if GOOGLE_CUDA && (CUDA_VERSION >= 12000)
+
+template <typename T>
+struct ToGpuSparseIndexType;
+template <>
+struct ToGpuSparseIndexType<int> {
+  static constexpr cusparseIndexType_t value = CUSPARSE_INDEX_32I;
+};
+template <>
+struct ToGpuSparseIndexType<int64_t> {
+  static constexpr cusparseIndexType_t value = CUSPARSE_INDEX_64I;
+};
+
+class GpuSparseSpGEMMDescr {
+ public:
+  GpuSparseSpGEMMDescr() : initialized_(false) {}
+  ~GpuSparseSpGEMMDescr() {
+    if (initialized_) {
+      cusparseSpGEMM_destroyDescr(descr_);
+    }
+  }
+  absl::Status Initialize() {
+    if (initialized_) {
+      return absl::InternalError(
+          "Double initializion of GpuSparseSpGEMMDescr.");
+    }
+    TF_RETURN_IF_GPUSPARSE_ERROR(cusparseSpGEMM_createDescr(&descr_));
+    initialized_ = true;
+    return absl::OkStatus();
+  }
+  cusparseSpGEMMDescr_t& get() { return descr_; }
+
+ private:
+  bool initialized_;
+  cusparseSpGEMMDescr_t descr_;
+
+  GpuSparseSpGEMMDescr(const GpuSparseSpGEMMDescr&) = delete;
+  void operator=(const GpuSparseSpGEMMDescr&) = delete;
+};
+
+class GpuSparseSpMatDescr {
+ public:
+  GpuSparseSpMatDescr() : initialized_(false) {}
+  ~GpuSparseSpMatDescr() {
+    if (initialized_) {
+      cusparseDestroySpMat(descr_);
+    }
+  }
+  template <typename IndexType, typename FloatType>
+  absl::Status InitializeCsr(int64_t rows, int64_t cols, int64_t nnz,
+                             IndexType* csrRowOffsets, IndexType* csrColInd,
+                             FloatType* csrValues) {
+    if (initialized_) {
+      return absl::InternalError("Double initializion of gpusparseSpMatDescr.");
+    }
+    using stream_executor::cuda::AsCudaDataType;
+    using stream_executor::dnn::ToDataType;
+    TF_RETURN_IF_GPUSPARSE_ERROR(cusparseCreateCsr(
+        &descr_, rows, cols, nnz, csrRowOffsets, csrColInd, csrValues,
+        ToGpuSparseIndexType<IndexType>::value,
+        ToGpuSparseIndexType<IndexType>::value, CUSPARSE_INDEX_BASE_ZERO,
+        AsCudaDataType(ToDataType<FloatType>::value)));
+    initialized_ = true;
+    return absl::OkStatus();
+  }
+  gpusparseSpMatDescr_t& get() { return descr_; }
+
+ private:
+  bool initialized_;
+  cusparseSpMatDescr_t descr_;
+  GpuSparseSpMatDescr(const GpuSparseSpMatDescr&) = delete;
+  void operator=(const GpuSparseSpMatDescr&) = delete;
+};
+
+class GpuSparseConstSpMatDescr {
+ public:
+  GpuSparseConstSpMatDescr() : initialized_(false) {}
+  ~GpuSparseConstSpMatDescr() {
+    if (initialized_) {
+      cusparseDestroySpMat(descr_);
+    }
+  }
+  template <typename IndexType, typename FloatType>
+  absl::Status InitializeCsr(int64_t rows, int64_t cols, int64_t nnz,
+                             const IndexType* csrRowOffsets,
+                             const IndexType* csrColInd,
+                             const FloatType* csrValues) {
+    if (initialized_) {
+      return absl::InternalError("Double initializion of gpusparseSpMatDescr.");
+    }
+    using stream_executor::cuda::AsCudaDataType;
+    using stream_executor::dnn::ToDataType;
+    TF_RETURN_IF_GPUSPARSE_ERROR(cusparseCreateConstCsr(
+        &descr_, rows, cols, nnz, csrRowOffsets, csrColInd, csrValues,
+        ToGpuSparseIndexType<IndexType>::value,
+        ToGpuSparseIndexType<IndexType>::value, CUSPARSE_INDEX_BASE_ZERO,
+        AsCudaDataType(ToDataType<FloatType>::value)));
+    initialized_ = true;
+    return absl::OkStatus();
+  }
+  cusparseConstSpMatDescr_t& get() { return descr_; }
+
+ private:
+  bool initialized_;
+  cusparseConstSpMatDescr_t descr_;
+  GpuSparseConstSpMatDescr(const GpuSparseConstSpMatDescr&) = delete;
+  void operator=(const GpuSparseConstSpMatDescr&) = delete;
+};
+
+#endif
+
+// The GpuSparse class provides a simplified templated API for cuSparse
+// (http://docs.nvidia.com/cuda/cusparse/index.html).
+// An object of this class wraps static cuSparse instances,
+// and will launch Cuda kernels on the stream wrapped by the GPU device
+// in the OpKernelContext provided to the constructor.
+//
+// Notice: All the computational member functions are asynchronous and simply
+// launch one or more Cuda kernels on the Cuda stream wrapped by the GpuSparse
+// object.
+
+class GpuSparse {
+ public:
+  // This object stores a pointer to context, which must outlive it.
+  explicit GpuSparse(OpKernelContext* context);
+  virtual ~GpuSparse() {}
+
+  // This initializes the GpuSparse class if it hasn't
+  // been initialized yet.  All following public methods require the
+  // class has been initialized.  Can be run multiple times; all
+  // subsequent calls after the first have no effect.
+  absl::Status Initialize();  // Move to constructor?
+
+  // ====================================================================
+  // Wrappers for cuSparse start here.
+  //
+
+  // Solves tridiagonal system of equations.
+  // See: https://docs.nvidia.com/cuda/cusparse/index.html#gtsv2
+  template <typename Scalar>
+  absl::Status Gtsv2(int m, int n, const Scalar* dl, const Scalar* d,
+                     const Scalar* du, Scalar* B, int ldb, void* pBuffer) const;
+
+  // Computes the size of a temporary buffer used by Gtsv2.
+  // See: https://docs.nvidia.com/cuda/cusparse/index.html#gtsv2_bufferSize
+  template <typename Scalar>
+  absl::Status Gtsv2BufferSizeExt(int m, int n, const Scalar* dl,
+                                  const Scalar* d, const Scalar* du,
+                                  const Scalar* B, int ldb,
+                                  size_t* bufferSizeInBytes) const;
+
+  // Solves tridiagonal system of equations without partial pivoting.
+  // See: https://docs.nvidia.com/cuda/cusparse/index.html#gtsv2_nopivot
+  template <typename Scalar>
+  absl::Status Gtsv2NoPivot(int m, int n, const Scalar* dl, const Scalar* d,
+                            const Scalar* du, Scalar* B, int ldb,
+                            void* pBuffer) const;
+
+  // Computes the size of a temporary buffer used by Gtsv2NoPivot.
+  // See:
+  // https://docs.nvidia.com/cuda/cusparse/index.html#gtsv2_nopivot_bufferSize
+  template <typename Scalar>
+  absl::Status Gtsv2NoPivotBufferSizeExt(int m, int n, const Scalar* dl,
+                                         const Scalar* d, const Scalar* du,
+                                         const Scalar* B, int ldb,
+                                         size_t* bufferSizeInBytes) const;
+
+  // Solves a batch of tridiagonal systems of equations. Doesn't support
+  // multiple right-hand sides per each system. Doesn't do pivoting.
+  // See: https://docs.nvidia.com/cuda/cusparse/index.html#gtsv2stridedbatch
+  template <typename Scalar>
+  absl::Status Gtsv2StridedBatch(int m, const Scalar* dl, const Scalar* d,
+                                 const Scalar* du, Scalar* x, int batchCount,
+                                 int batchStride, void* pBuffer) const;
+
+  // Computes the size of a temporary buffer used by Gtsv2StridedBatch.
+  // See:
+  // https://docs.nvidia.com/cuda/cusparse/index.html#gtsv2stridedbatch_bufferSize
+  template <typename Scalar>
+  absl::Status Gtsv2StridedBatchBufferSizeExt(int m, const Scalar* dl,
+                                              const Scalar* d, const Scalar* du,
+                                              const Scalar* x, int batchCount,
+                                              int batchStride,
+                                              size_t* bufferSizeInBytes) const;
+
+  // Compresses the indices of rows or columns. It can be interpreted as a
+  // conversion from COO to CSR sparse storage format. See:
+  // http://docs.nvidia.com/cuda/cusparse/index.html#cusparse-lt-t-gt-csr2coo.
+  absl::Status Csr2coo(const int* CsrRowPtr, int nnz, int m,
+                       int* cooRowInd) const;
+
+  // Uncompresses the indices of rows or columns. It can be interpreted as a
+  // conversion from CSR to COO sparse storage format. See:
+  // http://docs.nvidia.com/cuda/cusparse/index.html#cusparse-lt-t-gt-coo2csr.
+  absl::Status Coo2csr(const int* cooRowInd, int nnz, int m,
+                       int* csrRowPtr) const;
+
+#if (GOOGLE_CUDA && (CUDA_VERSION < 10020)) || \
+    (TENSORFLOW_USE_ROCM && TF_ROCM_VERSION < 40200)
+  // Sparse-dense matrix multiplication C = alpha * op(A) * op(B)  + beta * C,
+  // where A is a sparse matrix in CSR format, B and C are dense tall
+  // matrices.  This routine allows transposition of matrix B, which
+  // may improve performance.  See:
+  // http://docs.nvidia.com/cuda/cusparse/index.html#cusparse-lt-t-gt-csrmm2
+  //
+  // **NOTE** Matrices B and C are expected to be in column-major
+  // order; to make them consistent with TensorFlow they
+  // must be transposed (or the matmul op's pre/post-processing must take this
+  // into account).
+  //
+  // **NOTE** This is an in-place operation for data in C.
+  template <typename Scalar>
+  Status Csrmm(gpusparseOperation_t transA, gpusparseOperation_t transB, int m,
+               int n, int k, int nnz, const Scalar* alpha_host,
+               const gpusparseMatDescr_t descrA, const Scalar* csrSortedValA,
+               const int* csrSortedRowPtrA, const int* csrSortedColIndA,
+               const Scalar* B, int ldb, const Scalar* beta_host, Scalar* C,
+               int ldc) const;
+#else  // CUDA_VERSION >=10200 || TF_ROCM_VERSION >= 40200
+  // Workspace size query for sparse-dense matrix multiplication. Helper
+  // function for SpMM which computes y = alpha * op(A) * op(B) + beta * C,
+  // where A is a sparse matrix in CSR format, B and C are dense matricies in
+  // column-major format. Returns needed workspace size in bytes.
+  template <typename Scalar>
+  absl::Status SpMMBufferSize(gpusparseOperation_t transA,
+                              gpusparseOperation_t transB, const Scalar* alpha,
+                              const gpusparseSpMatDescr_t matA,
+                              const gpusparseDnMatDescr_t matB,
+                              const Scalar* beta, gpusparseDnMatDescr_t matC,
+                              gpusparseSpMMAlg_t alg, size_t* bufferSize) const;
+
+  // Sparse-dense matrix multiplication y = alpha * op(A) * op(B) + beta * C,
+  // where A is a sparse matrix in CSR format, B and C are dense matricies in
+  // column-major format. Buffer is assumed to be at least as large as the
+  // workspace size returned by SpMMBufferSize().
+  //
+  // **NOTE** This is an in-place operation for data in C.
+  template <typename Scalar>
+  absl::Status SpMM(gpusparseOperation_t transA, gpusparseOperation_t transB,
+                    const Scalar* alpha, const gpusparseSpMatDescr_t matA,
+                    const gpusparseDnMatDescr_t matB, const Scalar* beta,
+                    gpusparseDnMatDescr_t matC, gpusparseSpMMAlg_t alg,
+                    int8_t* buffer) const;
+#endif
+
+  // Sparse-dense vector multiplication y = alpha * op(A) * x  + beta * y,
+  // where A is a sparse matrix in CSR format, x and y are dense vectors. See:
+  // http://docs.nvidia.com/cuda/cusparse/index.html#cusparse-lt-t-gt-csrmv_mergepath
+  //
+  // **NOTE** This is an in-place operation for data in y.
+#if (GOOGLE_CUDA && (CUDA_VERSION < 10020)) || TENSORFLOW_USE_ROCM
+  template <typename Scalar>
+  Status Csrmv(gpusparseOperation_t transA, int m, int n, int nnz,
+               const Scalar* alpha_host, const gpusparseMatDescr_t descrA,
+               const Scalar* csrSortedValA, const int* csrSortedRowPtrA,
+               const int* csrSortedColIndA, const Scalar* x,
+               const Scalar* beta_host, Scalar* y) const;
+#else
+  template <typename Scalar>
+  absl::Status Csrmv(gpusparseOperation_t transA, int m, int n, int nnz,
+                     const Scalar* alpha_host, const Scalar* csrSortedValA,
+                     const int* csrSortedRowPtrA, const int* csrSortedColIndA,
+                     const Scalar* x, const Scalar* beta_host, Scalar* y) const;
+#endif  // CUDA_VERSION < 10020
+
+  // Computes workspace size for sparse - sparse matrix addition of matrices
+  // stored in CSR format.
+  template <typename Scalar>
+  absl::Status CsrgeamBufferSizeExt(
+      int m, int n, const Scalar* alpha, const gpusparseMatDescr_t descrA,
+      int nnzA, const Scalar* csrSortedValA, const int* csrSortedRowPtrA,
+      const int* csrSortedColIndA, const Scalar* beta,
+      const gpusparseMatDescr_t descrB, int nnzB, const Scalar* csrSortedValB,
+      const int* csrSortedRowPtrB, const int* csrSortedColIndB,
+      const gpusparseMatDescr_t descrC, Scalar* csrSortedValC,
+      int* csrSortedRowPtrC, int* csrSortedColIndC, size_t* bufferSize);
+
+  // Computes sparse-sparse matrix addition of matrices
+  // stored in CSR format.  This is part one: calculate nnz of the
+  // output.  csrSortedRowPtrC must be preallocated on device with
+  // m + 1 entries.  See:
+  // http://docs.nvidia.com/cuda/cusparse/index.html#cusparse-lt-t-gt-csrgeam.
+  absl::Status CsrgeamNnz(
+      int m, int n, const gpusparseMatDescr_t descrA, int nnzA,
+      const int* csrSortedRowPtrA, const int* csrSortedColIndA,
+      const gpusparseMatDescr_t descrB, int nnzB, const int* csrSortedRowPtrB,
+      const int* csrSortedColIndB, const gpusparseMatDescr_t descrC,
+      int* csrSortedRowPtrC, int* nnzTotalDevHostPtr, void* workspace);
+
+  // Computes sparse - sparse matrix addition of matrices
+  // stored in CSR format.  This is part two: perform sparse-sparse
+  // addition.  csrValC and csrColIndC must be allocated on the device
+  // with nnzTotalDevHostPtr entries (as calculated by CsrgeamNnz).  See:
+  // http://docs.nvidia.com/cuda/cusparse/index.html#cusparse-lt-t-gt-csrgeam.
+  template <typename Scalar>
+  absl::Status Csrgeam(int m, int n, const Scalar* alpha,
+                       const gpusparseMatDescr_t descrA, int nnzA,
+                       const Scalar* csrSortedValA, const int* csrSortedRowPtrA,
+                       const int* csrSortedColIndA, const Scalar* beta,
+                       const gpusparseMatDescr_t descrB, int nnzB,
+                       const Scalar* csrSortedValB, const int* csrSortedRowPtrB,
+                       const int* csrSortedColIndB,
+                       const gpusparseMatDescr_t descrC, Scalar* csrSortedValC,
+                       int* csrSortedRowPtrC, int* csrSortedColIndC,
+                       void* workspace);
+
+  // Computes sparse-sparse matrix multiplication of matrices
+  // stored in CSR format.
+#if TENSORFLOW_USE_ROCM
+  // Part one: calculate nnz of the output.
+  // csrSortedRowPtrC must be preallocated on device with m + 1 entries.
+  Status CsrgemmNnz(gpusparseOperation_t transA, gpusparseOperation_t transB,
+                    int m, int k, int n, const gpusparseMatDescr_t descrA,
+                    int nnzA, const int* csrSortedRowPtrA,
+                    const int* csrSortedColIndA,
+                    const gpusparseMatDescr_t descrB, int nnzB,
+                    const int* csrSortedRowPtrB, const int* csrSortedColIndB,
+                    const gpusparseMatDescr_t descrC, int* csrSortedRowPtrC,
+                    int* nnzTotalDevHostPtr);
+  // Part two: perform sparse-sparse matmul.
+  // csrValC and csrColIndC must be allocated on the device with
+  // nnzTotalDevHostPtr entries (as calculated by CsrgemmNnz).
+  template <typename Scalar>
+  Status Csrgemm(gpusparseOperation_t transA, gpusparseOperation_t transB,
+                 int m, int k, int n, const gpusparseMatDescr_t descrA,
+                 int nnzA, const Scalar* csrSortedValA,
+                 const int* csrSortedRowPtrA, const int* csrSortedColIndA,
+                 const gpusparseMatDescr_t descrB, int nnzB,
+                 const Scalar* csrSortedValB, const int* csrSortedRowPtrB,
+                 const int* csrSortedColIndB, const gpusparseMatDescr_t descrC,
+                 Scalar* csrSortedValC, int* csrSortedRowPtrC,
+                 int* csrSortedColIndC);
+#elif CUDA_VERSION < 12000
+  // Part zero: calculate required workspace size.
+  template <typename Scalar>
+  Status CsrgemmBufferSize(
+      int m, int n, int k, const gpusparseMatDescr_t descrA, int nnzA,
+      const int* csrSortedRowPtrA, const int* csrSortedColIndA,
+      const gpusparseMatDescr_t descrB, int nnzB, const int* csrSortedRowPtrB,
+      const int* csrSortedColIndB, csrgemm2Info_t info, size_t* workspaceBytes);
+  // Part one: calculate nnz of the output.
+  // csrSortedRowPtrC must be preallocated on device with m + 1 entries.
+  Status CsrgemmNnz(int m, int n, int k, const gpusparseMatDescr_t descrA,
+                    int nnzA, const int* csrSortedRowPtrA,
+                    const int* csrSortedColIndA,
+                    const gpusparseMatDescr_t descrB, int nnzB,
+                    const int* csrSortedRowPtrB, const int* csrSortedColIndB,
+                    const gpusparseMatDescr_t descrC, int* csrSortedRowPtrC,
+                    int* nnzTotalDevHostPtr, csrgemm2Info_t info,
+                    void* workspace);
+  // Part two: perform sparse-sparse matmul.
+  // csrValC and csrColIndC must be allocated on the device with
+  // nnzTotalDevHostPtr entries (as calculated by CsrgemmNnz).
+  template <typename Scalar>
+  Status Csrgemm(int m, int n, int k, const gpusparseMatDescr_t descrA,
+                 int nnzA, const Scalar* csrSortedValA,
+                 const int* csrSortedRowPtrA, const int* csrSortedColIndA,
+                 const gpusparseMatDescr_t descrB, int nnzB,
+                 const Scalar* csrSortedValB, const int* csrSortedRowPtrB,
+                 const int* csrSortedColIndB, const gpusparseMatDescr_t descrC,
+                 Scalar* csrSortedValC, int* csrSortedRowPtrC,
+                 int* csrSortedColIndC, const csrgemm2Info_t info,
+                 void* workspace);
+#else  // CUDA_VERSION >= 12000
+  template <typename Scalar>
+  absl::Status SpGEMM_workEstimation(GpuSparseConstSpMatDescr& matA,
+                                     GpuSparseConstSpMatDescr& matB,
+                                     GpuSparseSpMatDescr& matC,
+                                     GpuSparseSpGEMMDescr& spgemmDescr,
+                                     size_t* bufferSize1,
+                                     void* externalBuffer1);
+  template <typename Scalar>
+  absl::Status SpGEMM_compute(GpuSparseConstSpMatDescr& matA,
+                              GpuSparseConstSpMatDescr& matB,
+                              GpuSparseSpMatDescr& matC,
+                              GpuSparseSpGEMMDescr& spgemmDescr,
+                              size_t* bufferSize2, void* externalBuffer2);
+  template <typename Scalar>
+  absl::Status SpGEMM_copy(GpuSparseConstSpMatDescr& matA,
+                           GpuSparseConstSpMatDescr& matB,
+                           GpuSparseSpMatDescr& matC,
+                           GpuSparseSpGEMMDescr& spgemmDescr);
+#endif
+
+  // In-place reordering of unsorted CSR to sorted CSR.
+  // http://docs.nvidia.com/cuda/cusparse/index.html#cusparse-lt-t-gt-csru2csr
+  template <typename Scalar>
+  absl::Status Csru2csr(int m, int n, int nnz, const gpusparseMatDescr_t descrA,
+                        Scalar* csrVal, const int* csrRowPtr, int* csrColInd);
+
+  // Converts from CSR to CSC format (equivalently, transpose).
+  // http://docs.nvidia.com/cuda/cusparse/index.html#cusparse-csr2cscEx
+  template <typename Scalar>
+  absl::Status Csr2csc(int m, int n, int nnz, const Scalar* csrVal,
+                       const int* csrRowPtr, const int* csrColInd,
+                       Scalar* cscVal, int* cscRowInd, int* cscColPtr,
+                       const gpusparseAction_t copyValues);
+
+ private:
+  bool initialized_;
+  OpKernelContext* context_;  // not owned.
+  gpuStream_t gpu_stream_;
+  gpusparseHandle_t* gpusparse_handle_;  // not owned.
+
+  GpuSparse(const GpuSparse&) = delete;
+  void operator=(const GpuSparse&) = delete;
+};
+
+// A wrapper class to ensure that a CUDA sparse matrix descriptor is initialized
+// only once. For more details on the descriptor (gpusparseMatDescr_t), see:
+// https://docs.nvidia.com/cuda/cusparse/index.html#cusparsematdescrt
+class GpuSparseMatrixDescriptor {
+ public:
+  explicit GpuSparseMatrixDescriptor() : initialized_(false) {}
+
+  GpuSparseMatrixDescriptor(GpuSparseMatrixDescriptor&& rhs)
+      : initialized_(rhs.initialized_), descr_(std::move(rhs.descr_)) {
+    rhs.initialized_ = false;
+  }
+
+  GpuSparseMatrixDescriptor& operator=(GpuSparseMatrixDescriptor&& rhs) {
+    if (this == &rhs) return *this;
+    Release();
+    initialized_ = rhs.initialized_;
+    descr_ = std::move(rhs.descr_);
+    rhs.initialized_ = false;
+    return *this;
+  }
+
+  ~GpuSparseMatrixDescriptor() { Release(); }
+
+  // Initializes the underlying descriptor.  Will fail on the second call if
+  // called more than once.
+  absl::Status Initialize() {
+    DCHECK(!initialized_);
+#if GOOGLE_CUDA
+    TF_RETURN_IF_GPUSPARSE_ERROR(cusparseCreateMatDescr(&descr_));
+#elif TENSORFLOW_USE_ROCM
+    TF_RETURN_IF_GPUSPARSE_ERROR(se::wrap::hipsparseCreateMatDescr(&descr_));
+#endif
+    initialized_ = true;
+    return absl::OkStatus();
+  }
+
+  gpusparseMatDescr_t& descr() {
+    DCHECK(initialized_);
+    return descr_;
+  }
+
+  const gpusparseMatDescr_t& descr() const {
+    DCHECK(initialized_);
+    return descr_;
+  }
+
+ private:
+  void Release() {
+    if (initialized_) {
+#if GOOGLE_CUDA
+      cusparseDestroyMatDescr(descr_);
+#elif TENSORFLOW_USE_ROCM
+      se::wrap::hipsparseDestroyMatDescr(descr_);
+#endif
+      initialized_ = false;
+    }
+  }
+
+  bool initialized_;
+  gpusparseMatDescr_t descr_;
+
+  GpuSparseMatrixDescriptor(const GpuSparseMatrixDescriptor&) = delete;
+  void operator=(const GpuSparseMatrixDescriptor&) = delete;
+};
+
+#if GOOGLE_CUDA
+
+// A wrapper class to ensure that an unsorted/sorted CSR conversion information
+// struct (csru2csrInfo_t) is initialized only once. See:
+// https://docs.nvidia.com/cuda/cusparse/index.html#csru2csr
+class GpuSparseCsrSortingConversionInfo {
+ public:
+  explicit GpuSparseCsrSortingConversionInfo() : initialized_(false) {}
+
+  GpuSparseCsrSortingConversionInfo(GpuSparseCsrSortingConversionInfo&& rhs)
+      : initialized_(rhs.initialized_), info_(std::move(rhs.info_)) {
+    rhs.initialized_ = false;
+  }
+
+  GpuSparseCsrSortingConversionInfo& operator=(
+      GpuSparseCsrSortingConversionInfo&& rhs) {
+    if (this == &rhs) return *this;
+    Release();
+    initialized_ = rhs.initialized_;
+    info_ = std::move(rhs.info_);
+    rhs.initialized_ = false;
+    return *this;
+  }
+
+  ~GpuSparseCsrSortingConversionInfo() { Release(); }
+
+  // Initializes the underlying info. Will fail on the second call if called
+  // more than once.
+  absl::Status Initialize() {
+    DCHECK(!initialized_);
+    TF_RETURN_IF_GPUSPARSE_ERROR(cusparseCreateCsru2csrInfo(&info_));
+    initialized_ = true;
+    return absl::OkStatus();
+  }
+
+  csru2csrInfo_t& info() {
+    DCHECK(initialized_);
+    return info_;
+  }
+
+  const csru2csrInfo_t& info() const {
+    DCHECK(initialized_);
+    return info_;
+  }
+
+ private:
+  void Release() {
+    if (initialized_) {
+      cusparseDestroyCsru2csrInfo(info_);
+      initialized_ = false;
+    }
+  }
+
+  bool initialized_;
+  csru2csrInfo_t info_;
+
+  GpuSparseCsrSortingConversionInfo(const GpuSparseCsrSortingConversionInfo&) =
+      delete;
+  void operator=(const GpuSparseCsrSortingConversionInfo&) = delete;
+};
+
+#endif  // GOOGLE_CUDA
+
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+
+#endif  // TENSORFLOW_CORE_UTIL_CUDA_SPARSE_H_

tensorflow/core/util/debug_data_dumper.cc

@@ -0,0 +1,174 @@
+/* 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/util/debug_data_dumper.h"
+
+#include <optional>
+#include <set>
+#include <string>
+#include <vector>
+
+#include "absl/strings/str_format.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/path.h"
+#include "tensorflow/core/util/dump_graph.h"
+
+namespace tensorflow {
+
+DebugDataDumper* DebugDataDumper::Global() {
+  static DebugDataDumper* global_instance_ = new DebugDataDumper();
+  return global_instance_;
+}
+
+DebugDataDumper::DebugDataDumper() { LoadEnvvars(); }
+
+void DebugDataDumper::LoadEnvvars() {
+  // Load TF_DUMP_GRAPH_WRAPPED.
+  const char* dump_wrapped = getenv("TF_DUMP_GRAPH_WRAPPED");
+  dump_wrapped_ = static_cast<bool>(dump_wrapped);
+
+  // Load the name filter. Default value is null.
+  const char* name_filter = getenv("TF_DUMP_GRAPH_NAME_FILTER");
+  name_filter_ =
+      name_filter ? std::optional<std::string>{name_filter} : std::nullopt;
+
+  // Load the groups filter. Default value is "main".
+  const char* groups_filter = getenv("TF_DUMP_GRAPH_GROUPS");
+  groups_filter_ =
+      groups_filter ? std::set<std::string>(absl::StrSplit(groups_filter, ','))
+                    : std::set<std::string>({kDebugGroupMain});
+}
+
+bool DebugDataDumper::ShouldDump(const std::string& name,
+                                 const std::string& group) const {
+  // Skip dumping wrapped functions if needed.
+  if (!dump_wrapped_ && absl::StartsWith(name, "__wrapped__")) return false;
+
+  // Check the name filter.
+  if (name_filter_ == std::nullopt) {
+    VLOG(1) << "Skip dumping graph '" << name
+            << "', because TF_DUMP_GRAPH_NAME_FILTER is not set";
+    return false;
+  }
+
+  // If name_filter is not '*' or name doesn't contain the name_filter,
+  // skip the dump.
+  if (!absl::EqualsIgnoreCase(*name_filter_, "*") &&
+      !absl::StrContains(name, *name_filter_)) {
+    VLOG(1) << "Skip dumping graph '" << name
+            << "', because TF_DUMP_GRAPH_NAME_FILTER is not '*' and "
+            << "it is not contained by the graph name";
+    return false;
+  }
+
+  // Check the group filter.
+  if (groups_filter_.find(group) == groups_filter_.end() &&
+      groups_filter_.find("*") == groups_filter_.end())
+    return false;
+
+  // If all conditions are met, return true to allow the dump.
+  return true;
+}
+
+void DebugDataDumper::DumpOpCreationStackTraces(const std::string& name,
+                                                const std::string& group,
+                                                const std::string& tag,
+                                                const Graph* graph) {
+  // Check if we should take the dump.
+  if (!ShouldDump(name, group)) return;
+
+  // Construct the dump filename.
+  std::string dump_filename = GetDumpFilename(name, group, tag);
+
+  DumpToFile(dump_filename, "", ".csv", "StackTrace",
+             [graph, &dump_filename](WritableFile* file) {
+               auto status = file->Append("node_id,node_name,stackframes\n");
+               if (!status.ok()) {
+                 LOG(WARNING) << "error writing to file to " << dump_filename
+                              << ": " << status.message();
+                 return status;
+               }
+
+               for (Node* node : graph->nodes()) {
+                 auto stack_trace = node->GetStackTrace();
+                 if (stack_trace == nullptr) continue;
+
+                 int node_id = node->id();
+                 const std::string& node_name = node->name();
+                 std::vector<std::string> stackframes;
+                 stackframes.reserve(stack_trace->ToFrames().size());
+
+                 for (auto& frame : stack_trace->ToFrames()) {
+                   stackframes.push_back(
+                       absl::StrFormat("%s(%d): %s", frame.file_name,
+                                       frame.line_number, frame.function_name));
+                 }
+
+                 status = file->Append(
+                     absl::StrFormat("%d,%s,%s\n", node_id, node_name,
+                                     absl::StrJoin(stackframes, ";")));
+
+                 if (!status.ok()) {
+                   LOG(WARNING) << "error writing to file to " << dump_filename
+                                << ": " << status.message();
+                   return status;
+                 }
+               }
+
+               return file->Close();
+             });
+}
+
+void DebugDataDumper::DumpGraph(const std::string& name,
+                                const std::string& group,
+                                const std::string& tag, const Graph* graph,
+                                const FunctionLibraryDefinition* func_lib_def,
+                                bool bypass_filter) {
+  if (!ShouldDump(name, group) && !bypass_filter) return;
+
+  // Construct the dump filename.
+  std::string dump_filename = GetDumpFilename(name, group, tag);
+
+  // Make sure the dump filename is not longer than 255,
+  // because Linux won't take filename that long.
+  if (dump_filename.size() > 255) {
+    LOG(WARNING) << "Failed to dump graph " << dump_filename << " to "
+                 << ", because the file name is longer than 255";
+    return;
+  }
+
+  // Construct a graph def.
+  GraphDef graph_def;
+  graph->ToGraphDef(&graph_def);
+
+  if (func_lib_def) {
+    FunctionLibraryDefinition reachable_lib_def =
+        func_lib_def->ReachableDefinitions(graph_def);
+    *graph_def.mutable_library() = reachable_lib_def.ToProto();
+  }
+
+  // Now dump the graph into the target file.
+  DumpGraphDefToFile(dump_filename, graph_def);
+}
+
+std::string DebugDataDumper::GetDumpFilename(const std::string& name,
+                                             const std::string& group,
+                                             const std::string& tag) {
+  std::string dump_name = name.empty() ? "unknown_graph" : name;
+  return absl::StrFormat("%s.%04d.%s.%s", dump_name, GetNextDumpId(name), group,
+                         tag);
+}
+
+}  // namespace tensorflow

tensorflow/core/util/debug_data_dumper.h

@@ -0,0 +1,138 @@
+/* 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_UTIL_DEBUG_DATA_DUMPER_H_
+#define TENSORFLOW_CORE_UTIL_DEBUG_DATA_DUMPER_H_
+
+#include <optional>
+#include <set>
+#include <string>
+
+#include "absl/container/flat_hash_map.h"
+#include "tensorflow/core/platform/mutex.h"
+
+#define DEBUG_DATA_DUMPER() ::tensorflow::DebugDataDumper::Global()
+
+inline constexpr const char* kDebugGroupMain = "main";
+inline constexpr const char* kDebugGroupOpStacktrace = "op_stacktrace";
+inline constexpr const char* kDebugGroupGraphOptPass = "graph_opt_pass";
+inline constexpr const char* kDebugGroupBridgePhase1Clustering =
+    "bridge_phase1_clustering";
+inline constexpr const char* kDebugGroupRuntimeLowering = "runtime_lowering";
+inline constexpr const char* kDebugGroupBridgePhase1ExecutorExport =
+    "bridge_phase1_executor_export";
+inline constexpr const char* kDebugGroupBridgePhase2 = "bridge_phase2";
+inline constexpr const char* kDebugGroupDTensorMlir = "dtensor_mlir";
+inline constexpr const char* kDebugGroupDTensorGraph = "dtensor_graph";
+inline constexpr const char* kDebugGroupDTensorLayout = "dtensor_layout";
+
+namespace tensorflow {
+
+class FunctionLibraryDefinition;
+class Graph;
+
+////////////////////////////////////////////////////////////////////////////////
+// This class is responsible for dumping debugging data (e.g., GraphDef, MLIR).
+//
+// To dump GraphDef/MLIRs, take the following steps:
+// * Set envvar TF_DUMP_GRAPH_PREFIX to your target dump directory.
+// * Set envvar TF_DUMP_GRAPH_NAME_FILTER to '*' to dump all graphs,
+//   or a name filter to dump graphs with a name containing it.
+// * Set envvar TF_DUMP_GRAPH_GROUPS to your dump groups (comma-separated).
+//
+// The dumped graphs then can be found in your target dump directory.
+// The filename of the dump looks like this:
+// <name>.<order-id>.<group>.<tag>
+//
+// This is what each field means:
+// * <name>     : The name of your dump.
+// * <order-id> : The order of dumps of a specific name.
+//                Lower orders are executed before higher orders.
+// * <group>    : The group of your dump, e.g., main.
+// * <tag>      : The tag of your dump, e.g., your pass name.
+//
+// Example dump files are:
+// __inference_train_step_441.0.main.before_pre_placement_passes.pbtxt
+// __inference_train_step_441.1.main.before_placer.pbtxt
+// __inference_train_step_441.2.main.before_post_placement_passes.pbtxt
+// __inference_train_step_441.3.main.before_graph_optimization.pbtxt
+// __inference_train_step_441.4.main.after_graph_optimization.pbtxt
+// __inference_train_step_441.5.main.before_post_rewrite_for_exec_passes.pbtxt
+////////////////////////////////////////////////////////////////////////////////
+class DebugDataDumper {
+ public:
+  // Get the singleton instance.
+  static DebugDataDumper* Global();
+
+  // Initialize the debug data dumper.
+  void LoadEnvvars();
+
+  // Check if we should dump debug data.
+  // We should dump debug data only if the followings are true:
+  // 1. Envvar TF_DUMP_GRAPH_PREFIX is set to your target dump directory.
+  // 2. This condition is true if one of the followings is true.
+  //    2.1. TF_DUMP_GRAPH_NAME_FILTER is set to '*'
+  //    2.2. TF_DUMP_GRAPH_NAME_FILTER is set to a name filter
+  //         which is a substr of name.
+  // 3. The group is defined in TF_DUMP_GRAPH_GROUPS.
+  bool ShouldDump(const std::string& name, const std::string& group) const;
+
+  // Dump op creation callstacks, if ShouldDump returns true.
+  void DumpOpCreationStackTraces(const std::string& name,
+                                 const std::string& group,
+                                 const std::string& tag, const Graph* graph);
+
+  // Dump a graph, if ShouldDump returns true.
+  void DumpGraph(const std::string& name, const std::string& group,
+                 const std::string& tag, const Graph* graph,
+                 const FunctionLibraryDefinition* func_lib_def,
+                 bool bypass_filter = false);
+
+  // Get the dump file basename. Dump file basenames are in this format:
+  // <name>.<order-id>.<group>.<tag>
+  //
+  // What each field means is explained on the class level comment.
+  std::string GetDumpFilename(const std::string& name, const std::string& group,
+                              const std::string& tag);
+
+ private:
+  DebugDataDumper();
+
+  // Get next dump id for a name.
+  int GetNextDumpId(const std::string& name) {
+    // Use a lock to make sure this is thread safe.
+    const mutex_lock lock(lock_);
+    return dump_order_ids_[name]++;
+  }
+
+  // A dict to maintain the mapping from dump name to its current dump id.
+  absl::flat_hash_map<std::string, int> dump_order_ids_;
+
+  // A mutex to make sure this is thread safe.
+  tensorflow::mutex lock_;
+
+  // The name filter.
+  std::optional<std::string> name_filter_;
+
+  // The groups filter.
+  std::set<std::string> groups_filter_;
+
+  // A flag indicating whether to dump wrapped graphs.
+  bool dump_wrapped_;
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_UTIL_DEBUG_DATA_DUMPER_H_

tensorflow/core/util/debug_data_dumper_test.cc

@@ -0,0 +1,178 @@
+/* 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/util/debug_data_dumper.h"
+
+#include <string>
+
+#include "absl/strings/str_format.h"
+#include "tensorflow/core/graph/graph.h"
+#include "tensorflow/core/graph/node_builder.h"
+#include "tensorflow/core/lib/io/path.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace {
+
+TEST(DebugDataDumper, NoPrefixTest) {
+  EXPECT_EQ(false, DEBUG_DATA_DUMPER()->ShouldDump("DumpGraphToFileTest",
+                                                   kDebugGroupMain));
+}
+
+TEST(DebugDataDumper, NoNameFilterTest) {
+  std::string dir = testing::TmpDir();
+  setenv("TF_DUMP_GRAPH_PREFIX", dir.c_str(), 1);
+  DEBUG_DATA_DUMPER()->LoadEnvvars();
+
+  EXPECT_EQ(false, DEBUG_DATA_DUMPER()->ShouldDump("DumpGraphToFileTest",
+                                                   kDebugGroupMain));
+}
+
+TEST(DebugDataDumper, ShouldDumpTest) {
+  std::string dir = testing::TmpDir();
+  setenv("TF_DUMP_GRAPH_PREFIX", dir.c_str(), 1);
+  setenv("TF_DUMP_GRAPH_NAME_FILTER", "*", 1);
+  DEBUG_DATA_DUMPER()->LoadEnvvars();
+  EXPECT_EQ(true, DEBUG_DATA_DUMPER()->ShouldDump("DumpGraphToFileTest",
+                                                  kDebugGroupMain));
+
+  setenv("TF_DUMP_GRAPH_NAME_FILTER", "DumpGraph", 1);
+  DEBUG_DATA_DUMPER()->LoadEnvvars();
+  EXPECT_EQ(true, DEBUG_DATA_DUMPER()->ShouldDump("DumpGraphToFileTest",
+                                                  kDebugGroupMain));
+
+  setenv("TF_DUMP_GRAPH_NAME_FILTER", "DoNotDumpGraph", 1);
+  DEBUG_DATA_DUMPER()->LoadEnvvars();
+  EXPECT_EQ(false, DEBUG_DATA_DUMPER()->ShouldDump("DumpGraphToFileTest",
+                                                   kDebugGroupMain));
+
+  setenv("TF_DUMP_GRAPH_NAME_FILTER", "*", 1);
+  DEBUG_DATA_DUMPER()->LoadEnvvars();
+  EXPECT_EQ(false,
+            DEBUG_DATA_DUMPER()->ShouldDump("DumpGraphToFileTest",
+                                            kDebugGroupBridgePhase1Clustering));
+
+  setenv("TF_DUMP_GRAPH_GROUPS", "main,bridge_phase1_clustering", 1);
+  DEBUG_DATA_DUMPER()->LoadEnvvars();
+  EXPECT_EQ(true,
+            DEBUG_DATA_DUMPER()->ShouldDump("DumpGraphToFileTest",
+                                            kDebugGroupBridgePhase1Clustering));
+
+  DEBUG_DATA_DUMPER()->LoadEnvvars();
+  EXPECT_EQ(false, DEBUG_DATA_DUMPER()->ShouldDump(
+                       "__wrapped__DumpGraphToFileTest", kDebugGroupMain));
+
+  setenv("TF_DUMP_GRAPH_WRAPPED", "true", 1);
+  DEBUG_DATA_DUMPER()->LoadEnvvars();
+  EXPECT_EQ(true, DEBUG_DATA_DUMPER()->ShouldDump(
+                      "__wrapped__DumpGraphToFileTest", kDebugGroupMain));
+}
+
+TEST(DebugDataDumper, DumpFileBasenameTest) {
+  // For the same name, the order id should increment for each new dump file
+  // name.
+  EXPECT_EQ("DumpFileBasenameTest1.0000.main.tag1",
+            DEBUG_DATA_DUMPER()->GetDumpFilename("DumpFileBasenameTest1",
+                                                 kDebugGroupMain, "tag1"));
+  EXPECT_EQ("DumpFileBasenameTest1.0001.main.tag2",
+            DEBUG_DATA_DUMPER()->GetDumpFilename("DumpFileBasenameTest1",
+                                                 kDebugGroupMain, "tag2"));
+
+  // For other names, the order id should restart from 0.
+  EXPECT_EQ("DumpFileBasenameTest2.0000.main.tag1",
+            DEBUG_DATA_DUMPER()->GetDumpFilename("DumpFileBasenameTest2",
+                                                 kDebugGroupMain, "tag1"));
+}
+
+TEST(DebugDataDumper, DumpGraphToFileTest) {
+  Graph graph(OpRegistry::Global());
+  Node* node;
+  TF_CHECK_OK(NodeBuilder("A", "NoOp").Finalize(&graph, &node));
+
+  std::string dir = testing::TmpDir();
+  setenv("TF_DUMP_GRAPH_PREFIX", dir.c_str(), 1);
+  setenv("TF_DUMP_GRAPH_NAME_FILTER", "*", 1);
+  DEBUG_DATA_DUMPER()->LoadEnvvars();
+
+  DEBUG_DATA_DUMPER()->DumpGraph("DumpGraphToFileTest", kDebugGroupMain, "tag",
+                                 &graph, nullptr, false);
+
+  std::string dumpFilename =
+      io::JoinPath(dir, "DumpGraphToFileTest.0000.main.tag.pbtxt");
+  EXPECT_EQ(absl::OkStatus(), Env::Default()->FileExists(dumpFilename));
+}
+
+TEST(DebugDataDumper, DumpGraphLongFileNameCrashTest) {
+  Graph graph(OpRegistry::Global());
+  Node* node;
+  TF_CHECK_OK(NodeBuilder("A", "NoOp").Finalize(&graph, &node));
+
+  std::string dir = testing::TmpDir();
+  setenv("TF_DUMP_GRAPH_PREFIX", dir.c_str(), 1);
+  setenv("TF_DUMP_GRAPH_NAME_FILTER", "*", 1);
+  DEBUG_DATA_DUMPER()->LoadEnvvars();
+
+  // Make sure long file name does not crash.
+  std::string name = std::string(256, 'x');
+  DEBUG_DATA_DUMPER()->DumpGraph(name, kDebugGroupMain, "tag", &graph, nullptr,
+                                 false);
+
+  std::string dumpFilename = io::JoinPath(
+      dir, absl::StrFormat("%s.0000.main.tag.pbtxt", name.c_str()));
+  EXPECT_EQ(absl::StatusCode::kNotFound,
+            Env::Default()->FileExists(dumpFilename).code());
+}
+
+TEST(DebugDataDumper, DumpOpCreationStacktracesTest) {
+  Graph graph(OpRegistry::Global());
+  Node* node;
+  TF_CHECK_OK(NodeBuilder("A", "NoOp").Finalize(&graph, &node));
+
+  std::string dir = testing::TmpDir();
+  setenv("TF_DUMP_GRAPH_PREFIX", dir.c_str(), 1);
+  setenv("TF_DUMP_GRAPH_NAME_FILTER", "*", 1);
+  setenv("TF_DUMP_OP_CREATION_STACKTRACES", "1", 1);
+  DEBUG_DATA_DUMPER()->LoadEnvvars();
+
+  DEBUG_DATA_DUMPER()->DumpOpCreationStackTraces(
+      "DumpOpCreationStacktracesTest", kDebugGroupMain, "test", &graph);
+
+  std::string dumpFilename =
+      io::JoinPath(dir, "DumpOpCreationStacktracesTest.0000.main.test.csv");
+  EXPECT_EQ(absl::OkStatus(), Env::Default()->FileExists(dumpFilename));
+}
+
+TEST(DebugDataDumper, NoDumpOpCreationStacktracesTest) {
+  Graph graph(OpRegistry::Global());
+  Node* node;
+  TF_CHECK_OK(NodeBuilder("A", "NoOp").Finalize(&graph, &node));
+
+  std::string dir = testing::TmpDir();
+  setenv("TF_DUMP_GRAPH_PREFIX", dir.c_str(), 1);
+  setenv("TF_DUMP_GRAPH_NAME_FILTER", "*", 1);
+  DEBUG_DATA_DUMPER()->LoadEnvvars();
+
+  DEBUG_DATA_DUMPER()->DumpOpCreationStackTraces(
+      "DumpOpCreationStacktracesTest", kDebugGroupMain, "test", &graph);
+
+  std::string dumpFilename =
+      io::JoinPath(dir, "DumpOpCreationStacktracesTest.0000.main.test.json");
+  EXPECT_EQ(absl::StatusCode::kNotFound,
+            Env::Default()->FileExists(dumpFilename).code());
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/core/util/debug_events_writer.cc

@@ -0,0 +1,589 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/util/debug_events_writer.h"
+
+#include <deque>
+#include <memory>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include "tensorflow/core/lib/io/path.h"
+#include "tensorflow/core/lib/strings/strcat.h"
+#include "tensorflow/core/lib/strings/stringprintf.h"
+#include "tensorflow/core/platform/host_info.h"
+#include "tensorflow/core/public/release_version.h"
+
+namespace tensorflow {
+namespace tfdbg {
+
+namespace {
+void MaybeSetDebugEventTimestamp(DebugEvent* debug_event, Env* env) {
+  if (debug_event->wall_time() == 0) {
+    debug_event->set_wall_time(env->NowMicros() / 1e6);
+  }
+}
+}  // namespace
+
+SingleDebugEventFileWriter::SingleDebugEventFileWriter(
+    const std::string& file_path)
+    : env_(Env::Default()),
+      file_path_(file_path),
+      num_outstanding_events_(0),
+      writer_mu_() {}
+
+absl::Status SingleDebugEventFileWriter::Init() {
+  if (record_writer_ != nullptr) {
+    // TODO(cais): We currently don't check for file deletion. When the need
+    // arises, check and fix it.
+    return absl::OkStatus();
+  }
+
+  // Reset recordio_writer (which has a reference to writable_file_) so final
+  // Flush() and Close() call have access to writable_file_.
+  record_writer_.reset();
+
+  TF_RETURN_WITH_CONTEXT_IF_ERROR(
+      env_->NewWritableFile(file_path_, &writable_file_),
+      "Creating writable file ", file_path_);
+  record_writer_ = std::make_unique<io::RecordWriter>(writable_file_.get());
+  if (record_writer_ == nullptr) {
+    return absl::UnknownError(
+        absl::StrCat("Could not create record writer at path: ", file_path_));
+  }
+  num_outstanding_events_.store(0);
+  VLOG(1) << "Successfully opened debug events file: " << file_path_;
+  return absl::OkStatus();
+}
+
+void SingleDebugEventFileWriter::WriteSerializedDebugEvent(
+    absl::string_view debug_event_str) {
+  if (record_writer_ == nullptr) {
+    if (!Init().ok()) {
+      LOG(ERROR) << "Write failed because file could not be opened.";
+      return;
+    }
+  }
+  num_outstanding_events_.fetch_add(1);
+  {
+    mutex_lock l(writer_mu_);
+    record_writer_->WriteRecord(debug_event_str).IgnoreError();
+  }
+}
+
+absl::Status SingleDebugEventFileWriter::Flush() {
+  const int num_outstanding = num_outstanding_events_.load();
+  if (num_outstanding == 0) {
+    return absl::OkStatus();
+  }
+  if (writable_file_ == nullptr) {
+    return absl::UnknownError(
+        absl::StrCat("Unexpected NULL file for path: ", file_path_));
+  }
+
+  {
+    mutex_lock l(writer_mu_);
+    TF_RETURN_WITH_CONTEXT_IF_ERROR(record_writer_->Flush(), "Failed to flush ",
+                                    num_outstanding, " debug events to ",
+                                    file_path_);
+  }
+
+  TF_RETURN_WITH_CONTEXT_IF_ERROR(writable_file_->Sync(), "Failed to sync ",
+                                  num_outstanding, " debug events to ",
+                                  file_path_);
+  num_outstanding_events_.store(0);
+  return absl::OkStatus();
+}
+
+absl::Status SingleDebugEventFileWriter::Close() {
+  absl::Status status = Flush();
+  if (writable_file_ != nullptr) {
+    absl::Status close_status = writable_file_->Close();
+    if (!close_status.ok()) {
+      status = close_status;
+    }
+    record_writer_.reset(nullptr);
+    writable_file_.reset(nullptr);
+  }
+  num_outstanding_events_ = 0;
+  return status;
+}
+
+const std::string SingleDebugEventFileWriter::FileName() { return file_path_; }
+
+mutex DebugEventsWriter::factory_mu_(LINKER_INITIALIZED);
+
+DebugEventsWriter::~DebugEventsWriter() { Close().IgnoreError(); }
+
+// static
+DebugEventsWriter* DebugEventsWriter::GetDebugEventsWriter(
+    const std::string& dump_root, const std::string& tfdbg_run_id,
+    int64_t circular_buffer_size) {
+  mutex_lock l(DebugEventsWriter::factory_mu_);
+  std::unordered_map<std::string, std::unique_ptr<DebugEventsWriter>>*
+      writer_pool = DebugEventsWriter::GetDebugEventsWriterMap();
+  if (writer_pool->find(dump_root) == writer_pool->end()) {
+    std::unique_ptr<DebugEventsWriter> writer(
+        new DebugEventsWriter(dump_root, tfdbg_run_id, circular_buffer_size));
+    writer_pool->insert(std::make_pair(dump_root, std::move(writer)));
+  }
+  return (*writer_pool)[dump_root].get();
+}
+
+// static
+absl::Status DebugEventsWriter::LookUpDebugEventsWriter(
+    const std::string& dump_root, DebugEventsWriter** debug_events_writer) {
+  mutex_lock l(DebugEventsWriter::factory_mu_);
+  std::unordered_map<std::string, std::unique_ptr<DebugEventsWriter>>*
+      writer_pool = DebugEventsWriter::GetDebugEventsWriterMap();
+  if (writer_pool->find(dump_root) == writer_pool->end()) {
+    return absl::FailedPreconditionError(absl::StrCat(
+        "No DebugEventsWriter has been created at dump root ", dump_root));
+  }
+  *debug_events_writer = (*writer_pool)[dump_root].get();
+  return absl::OkStatus();
+}
+
+absl::Status DebugEventsWriter::Init() {
+  mutex_lock l(initialization_mu_);
+
+  // TODO(cais): We currently don't check for file deletion. When the need
+  // arises, check and fix file deletion.
+  if (is_initialized_) {
+    return absl::OkStatus();
+  }
+
+  if (!env_->IsDirectory(dump_root_).ok()) {
+    TF_RETURN_WITH_CONTEXT_IF_ERROR(env_->RecursivelyCreateDir(dump_root_),
+                                    "Failed to create directory ", dump_root_);
+  }
+
+  int64_t time_in_seconds = env_->NowMicros() / 1e6;
+  file_prefix_ = io::JoinPath(
+      dump_root_, absl::StrFormat("%s.%010lld.%s", kFileNamePrefix,
+                                  static_cast<long long>(time_in_seconds),
+                                  port::Hostname().c_str()));
+  TF_RETURN_IF_ERROR(InitNonMetadataFile(SOURCE_FILES));
+  TF_RETURN_IF_ERROR(InitNonMetadataFile(STACK_FRAMES));
+  TF_RETURN_IF_ERROR(InitNonMetadataFile(GRAPHS));
+
+  // In case there is one left over from before.
+  metadata_writer_.reset();
+
+  // The metadata file should be created.
+  std::string metadata_filename = GetFileNameInternal(METADATA);
+  metadata_writer_ =
+      std::make_unique<SingleDebugEventFileWriter>(metadata_filename);
+  if (metadata_writer_ == nullptr) {
+    return absl::UnknownError(
+        "Could not create debug event metadata file writer");
+  }
+
+  DebugEvent debug_event;
+  DebugMetadata* metadata = debug_event.mutable_debug_metadata();
+  metadata->set_tensorflow_version(TF_VERSION_STRING);
+  metadata->set_file_version(
+      absl::StrFormat("%s%d", kVersionPrefix, kCurrentFormatVersion));
+  metadata->set_tfdbg_run_id(tfdbg_run_id_);
+  TF_RETURN_IF_ERROR(SerializeAndWriteDebugEvent(&debug_event, METADATA));
+  TF_RETURN_WITH_CONTEXT_IF_ERROR(
+      metadata_writer_->Flush(), "Failed to flush debug event metadata writer");
+
+  TF_RETURN_IF_ERROR(InitNonMetadataFile(EXECUTION));
+  TF_RETURN_IF_ERROR(InitNonMetadataFile(GRAPH_EXECUTION_TRACES));
+  is_initialized_ = true;
+  return absl::OkStatus();
+}
+
+absl::Status DebugEventsWriter::WriteSourceFile(SourceFile* source_file) {
+  DebugEvent debug_event;
+  debug_event.set_allocated_source_file(source_file);
+  return SerializeAndWriteDebugEvent(&debug_event, SOURCE_FILES);
+}
+
+absl::Status DebugEventsWriter::WriteStackFrameWithId(
+    StackFrameWithId* stack_frame_with_id) {
+  DebugEvent debug_event;
+  debug_event.set_allocated_stack_frame_with_id(stack_frame_with_id);
+  return SerializeAndWriteDebugEvent(&debug_event, STACK_FRAMES);
+}
+
+absl::Status DebugEventsWriter::WriteGraphOpCreation(
+    GraphOpCreation* graph_op_creation) {
+  DebugEvent debug_event;
+  debug_event.set_allocated_graph_op_creation(graph_op_creation);
+  return SerializeAndWriteDebugEvent(&debug_event, GRAPHS);
+}
+
+absl::Status DebugEventsWriter::WriteDebuggedGraph(
+    DebuggedGraph* debugged_graph) {
+  DebugEvent debug_event;
+  debug_event.set_allocated_debugged_graph(debugged_graph);
+  return SerializeAndWriteDebugEvent(&debug_event, GRAPHS);
+}
+
+absl::Status DebugEventsWriter::WriteExecution(Execution* execution) {
+  if (circular_buffer_size_ <= 0) {
+    // No cyclic-buffer behavior.
+    DebugEvent debug_event;
+    debug_event.set_allocated_execution(execution);
+    return SerializeAndWriteDebugEvent(&debug_event, EXECUTION);
+  } else {
+    // Circular buffer behavior.
+    DebugEvent debug_event;
+    MaybeSetDebugEventTimestamp(&debug_event, env_);
+    debug_event.set_allocated_execution(execution);
+    std::string serialized;
+    debug_event.SerializeToString(&serialized);
+
+    mutex_lock l(execution_buffer_mu_);
+    execution_buffer_.emplace_back(std::move(serialized));
+    if (execution_buffer_.size() > circular_buffer_size_) {
+      execution_buffer_.pop_front();
+    }
+    return absl::OkStatus();
+  }
+}
+
+absl::Status DebugEventsWriter::WriteGraphExecutionTrace(
+    GraphExecutionTrace* graph_execution_trace) {
+  TF_RETURN_IF_ERROR(Init());
+  if (circular_buffer_size_ <= 0) {
+    // No cyclic-buffer behavior.
+    DebugEvent debug_event;
+    debug_event.set_allocated_graph_execution_trace(graph_execution_trace);
+    return SerializeAndWriteDebugEvent(&debug_event, GRAPH_EXECUTION_TRACES);
+  } else {
+    // Circular buffer behavior.
+    DebugEvent debug_event;
+    MaybeSetDebugEventTimestamp(&debug_event, env_);
+    debug_event.set_allocated_graph_execution_trace(graph_execution_trace);
+    std::string serialized;
+    debug_event.SerializeToString(&serialized);
+
+    mutex_lock l(graph_execution_trace_buffer_mu_);
+    graph_execution_trace_buffer_.emplace_back(std::move(serialized));
+    if (graph_execution_trace_buffer_.size() > circular_buffer_size_) {
+      graph_execution_trace_buffer_.pop_front();
+    }
+    return absl::OkStatus();
+  }
+}
+
+absl::Status DebugEventsWriter::WriteGraphExecutionTrace(
+    const std::string& tfdbg_context_id, const std::string& device_name,
+    const std::string& op_name, int32_t output_slot, int32_t tensor_debug_mode,
+    const Tensor& tensor_value) {
+  std::unique_ptr<GraphExecutionTrace> trace(new GraphExecutionTrace());
+  trace->set_tfdbg_context_id(tfdbg_context_id);
+  if (!op_name.empty()) {
+    trace->set_op_name(op_name);
+  }
+  if (output_slot > 0) {
+    trace->set_output_slot(output_slot);
+  }
+  if (tensor_debug_mode > 0) {
+    trace->set_tensor_debug_mode(TensorDebugMode(tensor_debug_mode));
+  }
+  trace->set_device_name(device_name);
+  tensor_value.AsProtoTensorContent(trace->mutable_tensor_proto());
+  return WriteGraphExecutionTrace(trace.release());
+}
+
+void DebugEventsWriter::WriteSerializedNonExecutionDebugEvent(
+    const std::string& debug_event_str, DebugEventFileType type) {
+  std::unique_ptr<SingleDebugEventFileWriter>* writer = nullptr;
+  SelectWriter(type, &writer);
+  (*writer)->WriteSerializedDebugEvent(debug_event_str);
+}
+
+void DebugEventsWriter::WriteSerializedExecutionDebugEvent(
+    const std::string& debug_event_str, DebugEventFileType type) {
+  const std::unique_ptr<SingleDebugEventFileWriter>* writer = nullptr;
+  std::deque<std::string>* buffer = nullptr;
+  mutex* mu = nullptr;
+  switch (type) {
+    case EXECUTION:
+      writer = &execution_writer_;
+      buffer = &execution_buffer_;
+      mu = &execution_buffer_mu_;
+      break;
+    case GRAPH_EXECUTION_TRACES:
+      writer = &graph_execution_traces_writer_;
+      buffer = &graph_execution_trace_buffer_;
+      mu = &graph_execution_trace_buffer_mu_;
+      break;
+    default:
+      return;
+  }
+
+  if (circular_buffer_size_ <= 0) {
+    // No cyclic-buffer behavior.
+    (*writer)->WriteSerializedDebugEvent(debug_event_str);
+  } else {
+    // Circular buffer behavior.
+    mutex_lock l(*mu);
+    buffer->push_back(debug_event_str);
+    if (buffer->size() > circular_buffer_size_) {
+      buffer->pop_front();
+    }
+  }
+}
+
+int DebugEventsWriter::RegisterDeviceAndGetId(const std::string& device_name) {
+  mutex_lock l(device_mu_);
+  int& device_id = device_name_to_id_[device_name];
+  if (device_id == 0) {
+    device_id = device_name_to_id_.size();
+    DebugEvent debug_event;
+    MaybeSetDebugEventTimestamp(&debug_event, env_);
+    DebuggedDevice* debugged_device = debug_event.mutable_debugged_device();
+    debugged_device->set_device_name(device_name);
+    debugged_device->set_device_id(device_id);
+    std::string serialized;
+    debug_event.SerializeToString(&serialized);
+    graphs_writer_->WriteSerializedDebugEvent(serialized);
+  }
+  return device_id;
+}
+
+absl::Status DebugEventsWriter::FlushNonExecutionFiles() {
+  TF_RETURN_IF_ERROR(Init());
+  if (source_files_writer_ != nullptr) {
+    TF_RETURN_IF_ERROR(source_files_writer_->Flush());
+  }
+  if (stack_frames_writer_ != nullptr) {
+    TF_RETURN_IF_ERROR(stack_frames_writer_->Flush());
+  }
+  if (graphs_writer_ != nullptr) {
+    TF_RETURN_IF_ERROR(graphs_writer_->Flush());
+  }
+  return absl::OkStatus();
+}
+
+absl::Status DebugEventsWriter::FlushExecutionFiles() {
+  TF_RETURN_IF_ERROR(Init());
+
+  if (execution_writer_ != nullptr) {
+    if (circular_buffer_size_ > 0) {
+      // Write out all the content in the circular buffers.
+      mutex_lock l(execution_buffer_mu_);
+      while (!execution_buffer_.empty()) {
+        execution_writer_->WriteSerializedDebugEvent(execution_buffer_.front());
+        // SerializeAndWriteDebugEvent(&execution_buffer_.front());
+        execution_buffer_.pop_front();
+      }
+    }
+    TF_RETURN_IF_ERROR(execution_writer_->Flush());
+  }
+
+  if (graph_execution_traces_writer_ != nullptr) {
+    if (circular_buffer_size_ > 0) {
+      // Write out all the content in the circular buffers.
+      mutex_lock l(graph_execution_trace_buffer_mu_);
+      while (!graph_execution_trace_buffer_.empty()) {
+        graph_execution_traces_writer_->WriteSerializedDebugEvent(
+            graph_execution_trace_buffer_.front());
+        graph_execution_trace_buffer_.pop_front();
+      }
+    }
+    TF_RETURN_IF_ERROR(graph_execution_traces_writer_->Flush());
+  }
+
+  return absl::OkStatus();
+}
+
+std::string DebugEventsWriter::FileName(DebugEventFileType type) {
+  if (file_prefix_.empty()) {
+    Init().IgnoreError();
+  }
+  return GetFileNameInternal(type);
+}
+
+absl::Status DebugEventsWriter::Close() {
+  {
+    mutex_lock l(initialization_mu_);
+    if (!is_initialized_) {
+      return absl::OkStatus();
+    }
+  }
+
+  std::vector<std::string> failed_to_close_files;
+
+  if (metadata_writer_ != nullptr) {
+    if (!metadata_writer_->Close().ok()) {
+      failed_to_close_files.push_back(metadata_writer_->FileName());
+    }
+    metadata_writer_.reset(nullptr);
+  }
+
+  TF_RETURN_IF_ERROR(FlushNonExecutionFiles());
+  if (source_files_writer_ != nullptr) {
+    if (!source_files_writer_->Close().ok()) {
+      failed_to_close_files.push_back(source_files_writer_->FileName());
+    }
+    source_files_writer_.reset(nullptr);
+  }
+  if (stack_frames_writer_ != nullptr) {
+    if (!stack_frames_writer_->Close().ok()) {
+      failed_to_close_files.push_back(stack_frames_writer_->FileName());
+    }
+    stack_frames_writer_.reset(nullptr);
+  }
+  if (graphs_writer_ != nullptr) {
+    if (!graphs_writer_->Close().ok()) {
+      failed_to_close_files.push_back(graphs_writer_->FileName());
+    }
+    graphs_writer_.reset(nullptr);
+  }
+
+  TF_RETURN_IF_ERROR(FlushExecutionFiles());
+  if (execution_writer_ != nullptr) {
+    if (!execution_writer_->Close().ok()) {
+      failed_to_close_files.push_back(execution_writer_->FileName());
+    }
+    execution_writer_.reset(nullptr);
+  }
+  if (graph_execution_traces_writer_ != nullptr) {
+    if (!graph_execution_traces_writer_->Close().ok()) {
+      failed_to_close_files.push_back(
+          graph_execution_traces_writer_->FileName());
+    }
+    graph_execution_traces_writer_.reset(nullptr);
+  }
+
+  if (failed_to_close_files.empty()) {
+    return absl::OkStatus();
+  } else {
+    return absl::FailedPreconditionError(absl::StrCat(
+        "Failed to close %d debug-events files associated with tfdbg",
+        failed_to_close_files.size()));
+  }
+}
+
+// static
+std::unordered_map<std::string, std::unique_ptr<DebugEventsWriter>>*
+DebugEventsWriter::GetDebugEventsWriterMap() {
+  static std::unordered_map<std::string,
+                            std::unique_ptr<DebugEventsWriter>>* writer_pool =
+      new std::unordered_map<std::string, std::unique_ptr<DebugEventsWriter>>();
+  return writer_pool;
+}
+
+DebugEventsWriter::DebugEventsWriter(const std::string& dump_root,
+                                     const std::string& tfdbg_run_id,
+                                     int64_t circular_buffer_size)
+    : env_(Env::Default()),
+      dump_root_(dump_root),
+      tfdbg_run_id_(tfdbg_run_id),
+      is_initialized_(false),
+      initialization_mu_(),
+      circular_buffer_size_(circular_buffer_size),
+      execution_buffer_(),
+      execution_buffer_mu_(),
+      graph_execution_trace_buffer_(),
+      graph_execution_trace_buffer_mu_(),
+      device_name_to_id_(),
+      device_mu_() {}
+
+absl::Status DebugEventsWriter::InitNonMetadataFile(DebugEventFileType type) {
+  std::unique_ptr<SingleDebugEventFileWriter>* writer = nullptr;
+  SelectWriter(type, &writer);
+  const std::string filename = GetFileNameInternal(type);
+  writer->reset();
+
+  *writer = std::make_unique<SingleDebugEventFileWriter>(filename);
+  if (*writer == nullptr) {
+    return absl::UnknownError(absl::StrCat(
+        "Could not create debug event file writer for ", filename));
+  }
+  TF_RETURN_WITH_CONTEXT_IF_ERROR(
+      (*writer)->Init(), "Initializing debug event writer at path ", filename);
+  VLOG(1) << "Successfully opened debug event file: " << filename;
+
+  return absl::OkStatus();
+}
+
+absl::Status DebugEventsWriter::SerializeAndWriteDebugEvent(
+    DebugEvent* debug_event, DebugEventFileType type) {
+  std::unique_ptr<SingleDebugEventFileWriter>* writer = nullptr;
+  SelectWriter(type, &writer);
+  if (writer != nullptr) {
+    // Timestamp is in seconds, with double precision.
+    MaybeSetDebugEventTimestamp(debug_event, env_);
+    std::string str;
+    debug_event->AppendToString(&str);
+    (*writer)->WriteSerializedDebugEvent(str);
+    return absl::OkStatus();
+  } else {
+    return absl::InternalError(absl::StrCat(
+        "Unable to find debug events file writer for DebugEventsFileType ",
+        type));
+  }
+}
+
+void DebugEventsWriter::SelectWriter(
+    DebugEventFileType type,
+    std::unique_ptr<SingleDebugEventFileWriter>** writer) {
+  switch (type) {
+    case METADATA:
+      *writer = &metadata_writer_;
+      break;
+    case SOURCE_FILES:
+      *writer = &source_files_writer_;
+      break;
+    case STACK_FRAMES:
+      *writer = &stack_frames_writer_;
+      break;
+    case GRAPHS:
+      *writer = &graphs_writer_;
+      break;
+    case EXECUTION:
+      *writer = &execution_writer_;
+      break;
+    case GRAPH_EXECUTION_TRACES:
+      *writer = &graph_execution_traces_writer_;
+      break;
+  }
+}
+
+const std::string DebugEventsWriter::GetSuffix(DebugEventFileType type) {
+  switch (type) {
+    case METADATA:
+      return kMetadataSuffix;
+    case SOURCE_FILES:
+      return kSourceFilesSuffix;
+    case STACK_FRAMES:
+      return kStackFramesSuffix;
+    case GRAPHS:
+      return kGraphsSuffix;
+    case EXECUTION:
+      return kExecutionSuffix;
+    case GRAPH_EXECUTION_TRACES:
+      return kGraphExecutionTracesSuffix;
+    default:
+      std::string suffix;
+      return suffix;
+  }
+}
+
+std::string DebugEventsWriter::GetFileNameInternal(DebugEventFileType type) {
+  const std::string suffix = GetSuffix(type);
+  return absl::StrCat(file_prefix_, ".", suffix);
+}
+
+}  // namespace tfdbg
+}  // namespace tensorflow

tensorflow/core/util/debug_events_writer.h

@@ -0,0 +1,279 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_UTIL_DEBUG_EVENTS_WRITER_H_
+#define TENSORFLOW_CORE_UTIL_DEBUG_EVENTS_WRITER_H_
+
+#include <atomic>
+#include <deque>
+#include <memory>
+#include <unordered_map>
+
+#include "absl/container/flat_hash_map.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/lib/io/record_writer.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/macros.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/protobuf/debug_event.pb.h"
+
+namespace tensorflow {
+namespace tfdbg {
+
+// The set of files generated by a debugged TensorFlow program.
+enum DebugEventFileType {
+  METADATA,
+  SOURCE_FILES,
+  STACK_FRAMES,
+  GRAPHS,
+  EXECUTION,
+  GRAPH_EXECUTION_TRACES,
+};
+
+// Helper class for DebugEventsWriter.
+// This class manages the writing of data to a single TFRecord file.
+// Each object of the DebugEventsWriter class below involves multiple
+// TFRecord files, and hence utilizes multiple objects of this helper class.
+class SingleDebugEventFileWriter {
+ public:
+  explicit SingleDebugEventFileWriter(const std::string& file_path);
+
+  absl::Status Init();
+
+  void WriteSerializedDebugEvent(absl::string_view debug_event_str);
+
+  absl::Status Flush();
+  absl::Status Close();
+
+  const std::string FileName();
+
+ private:
+  Env* env_;
+  const std::string file_path_;
+  std::atomic_int_fast32_t num_outstanding_events_;
+
+  std::unique_ptr<WritableFile> writable_file_;
+  std::unique_ptr<io::RecordWriter> record_writer_ TF_PT_GUARDED_BY(writer_mu_);
+  mutex writer_mu_;
+};
+
+// The DebugEvents writer class.
+class DebugEventsWriter {
+ public:
+#ifndef SWIG
+  // Prefix of version string present in the first entry of every event file.
+  // Default size of each circular buffer (unit: number of DebugEvent protos).
+  static constexpr const int64_t kDefaultCyclicBufferSize = 1000;
+
+  static constexpr const char* kFileNamePrefix = "tfdbg_events";
+  static constexpr const char* kMetadataSuffix = "metadata";
+  static constexpr const char* kSourceFilesSuffix = "source_files";
+  static constexpr const char* kStackFramesSuffix = "stack_frames";
+  static constexpr const char* kGraphsSuffix = "graphs";
+  static constexpr const char* kExecutionSuffix = "execution";
+  static constexpr const char* kGraphExecutionTracesSuffix =
+      "graph_execution_traces";
+
+  static constexpr const char* kVersionPrefix = "debug.Event:";
+  static constexpr const int kCurrentFormatVersion = 1;
+#endif
+
+  // Get the DebugEventsWriter for the given dump_root.
+  // For a given dump_root value, it is a singleton. tfdbg event files come in
+  // sets of six. The singleton pattern avoids storing multiple sets in a single
+  // folder, which might cause confusion.
+  //
+  // If an instance of DebugEventsWriter has already been created at a
+  // `dump_root`, calling this method with the same `dump_root` will return
+  // the existing instance.
+  //
+  // Args:
+  //   dump_root: Dump root directory. If it doesn't exist, will be created.
+  //   tfdbg_run_id: Debugging run ID of the writer.
+  //   circular_buffer_size: Circular buffer size (in number of DebugEvent
+  //     protos). If set to a value <=0, will abolish the circular-buffer
+  //     behavior.
+  // Returns:
+  //   A pointer to a DebugEventsWriter object: a per-dump_root singleton.
+  static DebugEventsWriter* GetDebugEventsWriter(
+      const std::string& dump_root, const std::string& tfdbg_run_id,
+      int64_t circular_buffer_size);
+  // Look up existing events writer by dump_root.
+  // If no DebugEventsWriter has been created at the dump_root, a non-OK
+  // Status will be returned. Else an OK status will be returned, with
+  // the pointer to the existing instance provided by reference.
+  static absl::Status LookUpDebugEventsWriter(
+      const std::string& dump_root, DebugEventsWriter** debug_events_writer);
+  ~DebugEventsWriter();
+
+  // Sets the debug event filenames and opens file for writing.
+  // All files (see the DebugEventFileType enum) share the same prefix and
+  // differ only in their suffixes. If not called by user, will be invoked
+  // automatically by a call to FileName() or any of the Write*() methods().
+  // Idempotent: if the metadata file exists and is open, this is a no-op.
+  // If on the other hand the file was opened, but has since disappeared (e.g.
+  // deleted by another process), this will open a new file.
+  absl::Status Init();
+
+  // The four DebugEvent fields below are written _without_ the circular
+  // buffer. Source file contents are written to the *.source_files file.
+  // Takes ownership of source_file.
+  absl::Status WriteSourceFile(SourceFile* source_file);
+  // Stack frames are written to the *.code_locations file.
+  // Takes ownership of stack_frame_with_id.
+  absl::Status WriteStackFrameWithId(StackFrameWithId* stack_frame_with_id);
+  // Graph op creation events are written to the *.graphs file.
+  // Takes ownership of graph_op_creation.
+  absl::Status WriteGraphOpCreation(GraphOpCreation* graph_op_creation);
+  // Debugged graphs are written to the *.graphs file.
+  // Takes ownership of debugged_graph.
+  absl::Status WriteDebuggedGraph(DebuggedGraph* debugged_graph);
+
+  // The two DebugEvent fields below are written to the circular buffer
+  // and saved to disk only at the FlushExecutionFiles() call.
+  // Execution events (eager execution of an op or a tf.function) are written
+  // to the *.execution file. Takes ownership of execution.
+  absl::Status WriteExecution(Execution* execution);
+  // Graph execution traces (graph-internal tensor values or their summaries)
+  // are written to the *.graph_execution_traces file.
+  // Takes ownership of graph_execution_trace.
+  absl::Status WriteGraphExecutionTrace(
+      GraphExecutionTrace* graph_execution_trace);
+
+  // Write a graph execution trace without using a protocol buffer.
+  // Instead, pass the raw values related to the graph execution trace.
+  // Args:
+  //   tfdbg_context_id: A unique ID for the context of interest, e.g., a
+  //   concreted compiled tf.function that the op of interest belongs to.
+  //   op_name: Name of the op that this graph execution trace is concerned
+  //     with. Applicable only to the single-tensor trace case. For cases in
+  //     which the trace concerns multiple tensors, this is an empty string.
+  //   output_slot: Output slot index of the op that this trace is concerned
+  //     with.
+  //   tensor_debug_mode: An integer that represents the tensor-debug mode
+  //   enum. tensor_value: The value of the tensor that describes the
+  //   tensor(s)
+  //     that this trace is concerned with. The semantics of this tensor value
+  //     depends on the value of `tensor_debug_mode`.
+  absl::Status WriteGraphExecutionTrace(const std::string& tfdbg_context_id,
+                                        const std::string& device_name,
+                                        const std::string& op_name,
+                                        int32_t output_slot,
+                                        int32_t tensor_debug_mode,
+                                        const Tensor& tensor_value);
+
+  // Writes a serialized DebugEvent to one of the debug-events files
+  // concerned with the non-execution events: the SOURCE_FILES, STACK_FRAMES
+  // and GRAPHS files.
+  // NOTE: Actually used in the Python binding, to avoid overhead of
+  // serializing and parsing protos at the language interface.
+  void WriteSerializedNonExecutionDebugEvent(const std::string& debug_event_str,
+                                             DebugEventFileType type);
+
+  // Writes a serialized DebugEvent to one of the debug-events files
+  // concerned with the execution-related events: the EXECUTION and
+  // GRAPH_EXECUTION_TRACES files. This involves the cyclic-buffer behavior if
+  // circular_buffer_size is configured to be >0.
+  // NOTE: Actually used in the Python binding, to avoid overhead of
+  // serializing and parsing protos at the language interface.
+  void WriteSerializedExecutionDebugEvent(const std::string& debug_event_str,
+                                          DebugEventFileType type);
+
+  // Given name of the device, retrieve a unique integer ID. As a side effect,
+  // if this is the first time this object encounters the device name,
+  // writes a DebuggedDevice proto to the .graphs file in the file set.
+  int RegisterDeviceAndGetId(const std::string& device_name);
+
+  // EventWriter automatically flushes and closes on destruction, but
+  // this method is provided for users who want to write to disk sooner
+  // and/or check for success.
+  // FlushNonExecutionFiles() pushes outstanding DebugEvents not written
+  // events to the circular buffer to their respective files.
+  absl::Status FlushNonExecutionFiles();
+
+  // Writes current contents of the circular buffers to their respective
+  // debug event files and clears the circular buffers.
+  absl::Status FlushExecutionFiles();
+
+  // Close() calls FlushNonExecutionFiles() and FlushExecutionFiles()
+  // and then closes the current debug events files.
+  absl::Status Close();
+
+ private:
+  static std::unordered_map<std::string, std::unique_ptr<DebugEventsWriter>>*
+
+  // Get a static map from dump-root path to DebugEventsWriter objects.
+  // This helps the per-dump-root singletone pattern.
+  GetDebugEventsWriterMap();
+
+  // Guards calls to the GetDebugEventsWriter() method.
+  static mutex factory_mu_;
+
+  DebugEventsWriter(const std::string& dump_root,
+                    const std::string& tfdbg_run_id,
+                    int64_t circular_buffer_size);
+
+  // Get the path prefix. The same for all files, which differ only in the
+  // suffix.
+  std::string FileName(DebugEventFileType type);
+
+  // Initialize the TFRecord writer for non-metadata file type.
+  absl::Status InitNonMetadataFile(DebugEventFileType type);
+
+  absl::Status SerializeAndWriteDebugEvent(DebugEvent* debug_event,
+                                           DebugEventFileType type);
+
+  void SelectWriter(DebugEventFileType type,
+                    std::unique_ptr<SingleDebugEventFileWriter>** writer);
+  const std::string GetSuffix(DebugEventFileType type);
+  std::string GetFileNameInternal(DebugEventFileType type);
+
+  Env* env_;
+  const std::string dump_root_;
+  const std::string tfdbg_run_id_;
+
+  std::string file_prefix_;
+  bool is_initialized_ TF_GUARDED_BY(initialization_mu_);
+  mutex initialization_mu_;
+
+  const int64_t circular_buffer_size_;
+  std::deque<std::string> execution_buffer_ TF_GUARDED_BY(execution_buffer_mu_);
+  mutex execution_buffer_mu_;
+  std::deque<std::string> graph_execution_trace_buffer_
+      TF_GUARDED_BY(graph_execution_trace_buffer_mu_);
+  mutex graph_execution_trace_buffer_mu_;
+
+  absl::flat_hash_map<std::string, int> device_name_to_id_
+      TF_GUARDED_BY(device_mu_);
+  mutex device_mu_;
+
+  std::unique_ptr<SingleDebugEventFileWriter> metadata_writer_;
+  std::unique_ptr<SingleDebugEventFileWriter> source_files_writer_;
+  std::unique_ptr<SingleDebugEventFileWriter> stack_frames_writer_;
+  std::unique_ptr<SingleDebugEventFileWriter> graphs_writer_;
+  std::unique_ptr<SingleDebugEventFileWriter> execution_writer_;
+  std::unique_ptr<SingleDebugEventFileWriter> graph_execution_traces_writer_;
+
+  DebugEventsWriter(const DebugEventsWriter&) = delete;
+  void operator=(const DebugEventsWriter&) = delete;
+
+  friend class DebugEventsWriterTest;
+};
+
+}  // namespace tfdbg
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_UTIL_DEBUG_EVENTS_WRITER_H_

tensorflow/core/util/debug_events_writer_test.cc

@@ -0,0 +1,890 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/util/debug_events_writer.h"
+
+#include <algorithm>
+#include <atomic>
+#include <memory>
+#include <vector>
+
+#include "absl/container/flat_hash_set.h"
+#include "tensorflow/core/framework/graph_debug_info.pb.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/lib/core/threadpool.h"
+#include "tensorflow/core/lib/io/path.h"
+#include "tensorflow/core/lib/io/record_reader.h"
+#include "tensorflow/core/lib/strings/stringprintf.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace tfdbg {
+
+// shorthand
+Env* env() { return Env::Default(); }
+
+class DebugEventsWriterTest : public ::testing::Test {
+ public:
+  static std::string GetDebugEventFileName(DebugEventsWriter* writer,
+                                           DebugEventFileType type) {
+    return writer->FileName(type);
+  }
+
+  static void ReadDebugEventProtos(DebugEventsWriter* writer,
+                                   DebugEventFileType type,
+                                   std::vector<DebugEvent>* protos) {
+    protos->clear();
+    const std::string filename = writer->FileName(type);
+    std::unique_ptr<RandomAccessFile> debug_events_file;
+    TF_CHECK_OK(env()->NewRandomAccessFile(filename, &debug_events_file));
+    io::RecordReader* reader = new io::RecordReader(debug_events_file.get());
+
+    uint64_t offset = 0;
+    DebugEvent actual;
+    while (ReadDebugEventProto(reader, &offset, &actual)) {
+      protos->push_back(actual);
+    }
+
+    delete reader;
+  }
+
+  static bool ReadDebugEventProto(io::RecordReader* reader, uint64_t* offset,
+                                  DebugEvent* proto) {
+    tstring record;
+    absl::Status s = reader->ReadRecord(offset, &record);
+    if (!s.ok()) {
+      return false;
+    }
+    return ParseProtoUnlimited(proto, record);
+  }
+
+  void SetUp() override {
+    dump_root_ = io::JoinPath(
+        testing::TmpDir(),
+        absl::StrFormat("%010lld", static_cast<long long>(env()->NowMicros())));
+    tfdbg_run_id_ = "test_tfdbg_run_id";
+  }
+
+  void TearDown() override {
+    if (env()->IsDirectory(dump_root_).ok()) {
+      int64_t undeleted_files = 0;
+      int64_t undeleted_dirs = 0;
+      TF_ASSERT_OK(env()->DeleteRecursively(dump_root_, &undeleted_files,
+                                            &undeleted_dirs));
+      ASSERT_EQ(0, undeleted_files);
+      ASSERT_EQ(0, undeleted_dirs);
+    }
+  }
+
+  std::string dump_root_;
+  std::string tfdbg_run_id_;
+};
+
+TEST_F(DebugEventsWriterTest, GetDebugEventsWriterSameRootGivesSameObject) {
+  // Test the per-dump_root_ singleton pattern.
+  DebugEventsWriter* writer_1 = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, DebugEventsWriter::kDefaultCyclicBufferSize);
+  DebugEventsWriter* writer_2 = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, DebugEventsWriter::kDefaultCyclicBufferSize);
+  EXPECT_EQ(writer_1, writer_2);
+}
+
+TEST_F(DebugEventsWriterTest, ConcurrentGetDebugEventsWriterSameDumpRoot) {
+  thread::ThreadPool* thread_pool =
+      new thread::ThreadPool(Env::Default(), "test_pool", 4);
+
+  std::vector<DebugEventsWriter*> writers;
+  mutex mu;
+  auto fn = [this, &writers, &mu]() {
+    DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
+        dump_root_, tfdbg_run_id_, DebugEventsWriter::kDefaultCyclicBufferSize);
+    {
+      mutex_lock l(mu);
+      writers.push_back(writer);
+    }
+  };
+  for (size_t i = 0; i < 4; ++i) {
+    thread_pool->Schedule(fn);
+  }
+  delete thread_pool;
+
+  EXPECT_EQ(writers.size(), 4);
+  EXPECT_EQ(writers[0], writers[1]);
+  EXPECT_EQ(writers[1], writers[2]);
+  EXPECT_EQ(writers[2], writers[3]);
+}
+
+TEST_F(DebugEventsWriterTest, ConcurrentGetDebugEventsWriterDiffDumpRoots) {
+  thread::ThreadPool* thread_pool =
+      new thread::ThreadPool(Env::Default(), "test_pool", 3);
+
+  std::atomic_int_fast64_t counter(0);
+  std::vector<DebugEventsWriter*> writers;
+  mutex mu;
+  auto fn = [this, &counter, &writers, &mu]() {
+    const std::string new_dump_root =
+        io::JoinPath(dump_root_, absl::StrFormat("%ld", counter.fetch_add(1)));
+    DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
+        new_dump_root, tfdbg_run_id_,
+        DebugEventsWriter::kDefaultCyclicBufferSize);
+    {
+      mutex_lock l(mu);
+      writers.push_back(writer);
+    }
+  };
+  for (size_t i = 0; i < 3; ++i) {
+    thread_pool->Schedule(fn);
+  }
+  delete thread_pool;
+
+  EXPECT_EQ(writers.size(), 3);
+  EXPECT_NE(writers[0], writers[1]);
+  EXPECT_NE(writers[0], writers[2]);
+  EXPECT_NE(writers[1], writers[2]);
+}
+
+TEST_F(DebugEventsWriterTest, GetDebugEventsWriterDifferentRoots) {
+  // Test the DebugEventsWriters for different directories are different.
+  DebugEventsWriter* writer_1 = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, DebugEventsWriter::kDefaultCyclicBufferSize);
+  const std::string dump_root_2 = io::JoinPath(dump_root_, "subdirectory");
+  DebugEventsWriter* writer_2 = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_2, tfdbg_run_id_, DebugEventsWriter::kDefaultCyclicBufferSize);
+  EXPECT_NE(writer_1, writer_2);
+}
+
+TEST_F(DebugEventsWriterTest, GetAndInitDebugEventsWriter) {
+  DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, DebugEventsWriter::kDefaultCyclicBufferSize);
+  TF_ASSERT_OK(writer->Init());
+  TF_ASSERT_OK(writer->Close());
+
+  // Verify the metadata file's content.
+  std::vector<DebugEvent> actuals;
+  ReadDebugEventProtos(writer, DebugEventFileType::METADATA, &actuals);
+  EXPECT_EQ(actuals.size(), 1);
+  EXPECT_GT(actuals[0].debug_metadata().tensorflow_version().length(), 0);
+  // Check the content of the file version string.
+  const std::string file_version = actuals[0].debug_metadata().file_version();
+  EXPECT_EQ(file_version.find(DebugEventsWriter::kVersionPrefix), 0);
+  EXPECT_GT(file_version.size(), strlen(DebugEventsWriter::kVersionPrefix));
+  // Check the tfdbg run ID.
+  EXPECT_EQ(actuals[0].debug_metadata().tfdbg_run_id(), "test_tfdbg_run_id");
+
+  // Verify that the .source_files file has been created and is empty.
+  ReadDebugEventProtos(writer, DebugEventFileType::SOURCE_FILES, &actuals);
+  // Verify that the .stack_frames file has been created and is empty.
+  ReadDebugEventProtos(writer, DebugEventFileType::STACK_FRAMES, &actuals);
+}
+
+TEST_F(DebugEventsWriterTest, CallingCloseWithoutInitIsOkay) {
+  DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, DebugEventsWriter::kDefaultCyclicBufferSize);
+  TF_ASSERT_OK(writer->Close());
+}
+
+TEST_F(DebugEventsWriterTest, CallingCloseTwiceIsOkay) {
+  DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, DebugEventsWriter::kDefaultCyclicBufferSize);
+  TF_ASSERT_OK(writer->Close());
+  TF_ASSERT_OK(writer->Close());
+}
+
+TEST_F(DebugEventsWriterTest, ConcurrentInitCalls) {
+  // Test that concurrent calls to Init() works correctly.
+  DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, DebugEventsWriter::kDefaultCyclicBufferSize);
+
+  thread::ThreadPool* thread_pool =
+      new thread::ThreadPool(Env::Default(), "test_pool", 4);
+  auto fn = [&writer]() { TF_ASSERT_OK(writer->Init()); };
+  for (size_t i = 0; i < 3; ++i) {
+    thread_pool->Schedule(fn);
+  }
+  delete thread_pool;
+
+  TF_ASSERT_OK(writer->Close());
+
+  // Verify the metadata file's content.
+  std::vector<DebugEvent> actuals;
+  ReadDebugEventProtos(writer, DebugEventFileType::METADATA, &actuals);
+  EXPECT_EQ(actuals.size(), 1);
+  EXPECT_GT(actuals[0].debug_metadata().tensorflow_version().length(), 0);
+  // Check the content of the file version string.
+  const std::string file_version = actuals[0].debug_metadata().file_version();
+  EXPECT_EQ(file_version.find(DebugEventsWriter::kVersionPrefix), 0);
+  EXPECT_GT(file_version.size(), strlen(DebugEventsWriter::kVersionPrefix));
+  EXPECT_EQ(actuals[0].debug_metadata().tfdbg_run_id(), "test_tfdbg_run_id");
+
+  // Verify that the .source_files file has been created and is empty.
+  ReadDebugEventProtos(writer, DebugEventFileType::SOURCE_FILES, &actuals);
+  // Verify that the .stack_frames file has been created and is empty.
+  ReadDebugEventProtos(writer, DebugEventFileType::STACK_FRAMES, &actuals);
+}
+
+TEST_F(DebugEventsWriterTest, InitTwiceDoesNotCreateNewMetadataFile) {
+  // Test that Init() is idempotent.
+  DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, DebugEventsWriter::kDefaultCyclicBufferSize);
+  TF_ASSERT_OK(writer->Init());
+
+  std::vector<DebugEvent> actuals;
+  ReadDebugEventProtos(writer, DebugEventFileType::METADATA, &actuals);
+  EXPECT_EQ(actuals.size(), 1);
+  EXPECT_GT(actuals[0].debug_metadata().tensorflow_version().length(), 0);
+  EXPECT_EQ(actuals[0].debug_metadata().tfdbg_run_id(), "test_tfdbg_run_id");
+  EXPECT_GE(actuals[0].debug_metadata().file_version().size(), 0);
+
+  std::string metadata_path_1 =
+      GetDebugEventFileName(writer, DebugEventFileType::METADATA);
+  TF_ASSERT_OK(writer->Init());
+  EXPECT_EQ(GetDebugEventFileName(writer, DebugEventFileType::METADATA),
+            metadata_path_1);
+  TF_ASSERT_OK(writer->Close());
+
+  // Verify the metadata file's content.
+  ReadDebugEventProtos(writer, DebugEventFileType::METADATA, &actuals);
+  EXPECT_EQ(actuals.size(), 1);
+  EXPECT_GT(actuals[0].debug_metadata().tensorflow_version().length(), 0);
+  EXPECT_EQ(actuals[0].debug_metadata().tfdbg_run_id(), "test_tfdbg_run_id");
+  EXPECT_GE(actuals[0].debug_metadata().file_version().size(), 0);
+}
+
+TEST_F(DebugEventsWriterTest, WriteSourceFile) {
+  DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, DebugEventsWriter::kDefaultCyclicBufferSize);
+  TF_ASSERT_OK(writer->Init());
+
+  SourceFile* source_file_1 = new SourceFile();
+  source_file_1->set_file_path("/home/tf_programs/main.py");
+  source_file_1->set_host_name("localhost.localdomain");
+  source_file_1->add_lines("import tensorflow as tf");
+  source_file_1->add_lines("");
+  source_file_1->add_lines("print(tf.constant([42.0]))");
+  source_file_1->add_lines("");
+  TF_ASSERT_OK(writer->WriteSourceFile(source_file_1));
+
+  SourceFile* source_file_2 = new SourceFile();
+  source_file_2->set_file_path("/home/tf_programs/train.py");
+  source_file_2->set_host_name("localhost.localdomain");
+  source_file_2->add_lines("import tensorflow.keras as keras");
+  source_file_2->add_lines("");
+  source_file_2->add_lines("model = keras.Sequential()");
+  TF_ASSERT_OK(writer->WriteSourceFile(source_file_2));
+
+  TF_ASSERT_OK(writer->FlushNonExecutionFiles());
+  TF_ASSERT_OK(writer->Close());
+
+  std::vector<DebugEvent> actuals;
+  ReadDebugEventProtos(writer, DebugEventFileType::SOURCE_FILES, &actuals);
+  EXPECT_EQ(actuals.size(), 2);
+  EXPECT_GT(actuals[0].wall_time(), 0);
+  EXPECT_GT(actuals[1].wall_time(), actuals[0].wall_time());
+
+  SourceFile actual_source_file_1 = actuals[0].source_file();
+  EXPECT_EQ(actual_source_file_1.file_path(), "/home/tf_programs/main.py");
+  EXPECT_EQ(actual_source_file_1.host_name(), "localhost.localdomain");
+  EXPECT_EQ(actual_source_file_1.lines().size(), 4);
+  EXPECT_EQ(actual_source_file_1.lines()[0], "import tensorflow as tf");
+  EXPECT_EQ(actual_source_file_1.lines()[1], "");
+  EXPECT_EQ(actual_source_file_1.lines()[2], "print(tf.constant([42.0]))");
+  EXPECT_EQ(actual_source_file_1.lines()[3], "");
+
+  SourceFile actual_source_file_2 = actuals[1].source_file();
+  EXPECT_EQ(actual_source_file_2.file_path(), "/home/tf_programs/train.py");
+  EXPECT_EQ(actual_source_file_2.host_name(), "localhost.localdomain");
+  EXPECT_EQ(actual_source_file_2.lines().size(), 3);
+  EXPECT_EQ(actual_source_file_2.lines()[0],
+            "import tensorflow.keras as keras");
+  EXPECT_EQ(actual_source_file_2.lines()[1], "");
+  EXPECT_EQ(actual_source_file_2.lines()[2], "model = keras.Sequential()");
+
+  // Verify no cross talk in the other non-execution debug-event files.
+  ReadDebugEventProtos(writer, DebugEventFileType::STACK_FRAMES, &actuals);
+  EXPECT_EQ(actuals.size(), 0);
+  ReadDebugEventProtos(writer, DebugEventFileType::GRAPHS, &actuals);
+  EXPECT_EQ(actuals.size(), 0);
+  ReadDebugEventProtos(writer, DebugEventFileType::EXECUTION, &actuals);
+  EXPECT_EQ(actuals.size(), 0);
+  ReadDebugEventProtos(writer, DebugEventFileType::GRAPH_EXECUTION_TRACES,
+                       &actuals);
+  EXPECT_EQ(actuals.size(), 0);
+}
+
+TEST_F(DebugEventsWriterTest, WriteStackFramesFile) {
+  DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, DebugEventsWriter::kDefaultCyclicBufferSize);
+  TF_ASSERT_OK(writer->Init());
+
+  StackFrameWithId* stack_frame_1 = new StackFrameWithId();
+  stack_frame_1->set_id("deadbeaf");
+  GraphDebugInfo::FileLineCol* file_line_col =
+      stack_frame_1->mutable_file_line_col();
+  file_line_col->set_file_index(12);
+  file_line_col->set_line(20);
+  file_line_col->set_col(2);
+  file_line_col->set_func("my_func");
+  file_line_col->set_code("  x = y + z");
+
+  StackFrameWithId* stack_frame_2 = new StackFrameWithId();
+  stack_frame_2->set_id("eeeeeeec");
+  file_line_col = stack_frame_2->mutable_file_line_col();
+  file_line_col->set_file_index(12);
+  file_line_col->set_line(21);
+  file_line_col->set_col(4);
+  file_line_col->set_func("my_func");
+  file_line_col->set_code("  x = x ** 2.0");
+
+  TF_ASSERT_OK(writer->WriteStackFrameWithId(stack_frame_1));
+  TF_ASSERT_OK(writer->WriteStackFrameWithId(stack_frame_2));
+  TF_ASSERT_OK(writer->FlushNonExecutionFiles());
+  TF_ASSERT_OK(writer->Close());
+
+  std::vector<DebugEvent> actuals;
+  ReadDebugEventProtos(writer, DebugEventFileType::STACK_FRAMES, &actuals);
+  EXPECT_EQ(actuals.size(), 2);
+  EXPECT_GT(actuals[0].wall_time(), 0);
+  EXPECT_GT(actuals[1].wall_time(), actuals[0].wall_time());
+
+  StackFrameWithId actual_stack_frame_1 = actuals[0].stack_frame_with_id();
+  EXPECT_EQ(actual_stack_frame_1.id(), "deadbeaf");
+  GraphDebugInfo::FileLineCol file_line_col_1 =
+      actual_stack_frame_1.file_line_col();
+  EXPECT_EQ(file_line_col_1.file_index(), 12);
+  EXPECT_EQ(file_line_col_1.line(), 20);
+  EXPECT_EQ(file_line_col_1.col(), 2);
+  EXPECT_EQ(file_line_col_1.func(), "my_func");
+  EXPECT_EQ(file_line_col_1.code(), "  x = y + z");
+
+  StackFrameWithId actual_stack_frame_2 = actuals[1].stack_frame_with_id();
+  EXPECT_EQ(actual_stack_frame_2.id(), "eeeeeeec");
+  GraphDebugInfo::FileLineCol file_line_col_2 =
+      actual_stack_frame_2.file_line_col();
+  EXPECT_EQ(file_line_col_2.file_index(), 12);
+  EXPECT_EQ(file_line_col_2.line(), 21);
+  EXPECT_EQ(file_line_col_2.col(), 4);
+  EXPECT_EQ(file_line_col_2.func(), "my_func");
+  EXPECT_EQ(file_line_col_2.code(), "  x = x ** 2.0");
+
+  // Verify no cross talk in the other non-execution debug-event files.
+  ReadDebugEventProtos(writer, DebugEventFileType::SOURCE_FILES, &actuals);
+  EXPECT_EQ(actuals.size(), 0);
+  ReadDebugEventProtos(writer, DebugEventFileType::GRAPHS, &actuals);
+  EXPECT_EQ(actuals.size(), 0);
+}
+
+TEST_F(DebugEventsWriterTest, WriteGraphOpCreationAndDebuggedGraph) {
+  DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, DebugEventsWriter::kDefaultCyclicBufferSize);
+  TF_ASSERT_OK(writer->Init());
+
+  GraphOpCreation* graph_op_creation = new GraphOpCreation();
+  graph_op_creation->set_op_type("MatMul");
+  graph_op_creation->set_op_name("Dense_1/MatMul");
+  TF_ASSERT_OK(writer->WriteGraphOpCreation(graph_op_creation));
+
+  DebuggedGraph* debugged_graph = new DebuggedGraph();
+  debugged_graph->set_graph_id("deadbeaf");
+  debugged_graph->set_graph_name("my_func_graph");
+  TF_ASSERT_OK(writer->WriteDebuggedGraph(debugged_graph));
+
+  TF_ASSERT_OK(writer->FlushNonExecutionFiles());
+  TF_ASSERT_OK(writer->Close());
+
+  std::vector<DebugEvent> actuals;
+  ReadDebugEventProtos(writer, DebugEventFileType::GRAPHS, &actuals);
+  EXPECT_EQ(actuals.size(), 2);
+  EXPECT_GT(actuals[0].wall_time(), 0);
+  EXPECT_GT(actuals[1].wall_time(), actuals[0].wall_time());
+
+  GraphOpCreation actual_op_creation = actuals[0].graph_op_creation();
+  EXPECT_EQ(actual_op_creation.op_type(), "MatMul");
+  EXPECT_EQ(actual_op_creation.op_name(), "Dense_1/MatMul");
+
+  DebuggedGraph actual_debugged_graph = actuals[1].debugged_graph();
+  EXPECT_EQ(actual_debugged_graph.graph_id(), "deadbeaf");
+  EXPECT_EQ(actual_debugged_graph.graph_name(), "my_func_graph");
+
+  // Verify no cross talk in the other non-execution debug-event files.
+  ReadDebugEventProtos(writer, DebugEventFileType::SOURCE_FILES, &actuals);
+  EXPECT_EQ(actuals.size(), 0);
+  ReadDebugEventProtos(writer, DebugEventFileType::STACK_FRAMES, &actuals);
+  EXPECT_EQ(actuals.size(), 0);
+}
+
+TEST_F(DebugEventsWriterTest, ConcurrentWriteCallsToTheSameFile) {
+  const size_t kConcurrentWrites = 100;
+  DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, DebugEventsWriter::kDefaultCyclicBufferSize);
+  TF_ASSERT_OK(writer->Init());
+
+  thread::ThreadPool* thread_pool =
+      new thread::ThreadPool(Env::Default(), "test_pool", 8);
+  std::atomic_int_fast64_t counter(0);
+  auto fn = [&writer, &counter]() {
+    const std::string file_path = absl::StrFormat(
+        "/home/tf_programs/program_%.3ld.py", counter.fetch_add(1));
+    SourceFile* source_file = new SourceFile();
+    source_file->set_file_path(file_path);
+    source_file->set_host_name("localhost.localdomain");
+    TF_ASSERT_OK(writer->WriteSourceFile(source_file));
+  };
+  for (size_t i = 0; i < kConcurrentWrites; ++i) {
+    thread_pool->Schedule(fn);
+  }
+  delete thread_pool;
+
+  TF_ASSERT_OK(writer->Close());
+
+  std::vector<DebugEvent> actuals;
+  ReadDebugEventProtos(writer, DebugEventFileType::SOURCE_FILES, &actuals);
+  EXPECT_EQ(actuals.size(), kConcurrentWrites);
+  std::vector<std::string> file_paths;
+  std::vector<std::string> host_names;
+  for (size_t i = 0; i < kConcurrentWrites; ++i) {
+    file_paths.push_back(actuals[i].source_file().file_path());
+    host_names.push_back(actuals[i].source_file().host_name());
+  }
+  std::sort(file_paths.begin(), file_paths.end());
+  for (size_t i = 0; i < kConcurrentWrites; ++i) {
+    EXPECT_EQ(file_paths[i],
+              absl::StrFormat("/home/tf_programs/program_%.3ld.py", i));
+    EXPECT_EQ(host_names[i], "localhost.localdomain");
+  }
+}
+
+TEST_F(DebugEventsWriterTest, ConcurrentWriteAndFlushCallsToTheSameFile) {
+  const size_t kConcurrentWrites = 100;
+  DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, DebugEventsWriter::kDefaultCyclicBufferSize);
+  TF_ASSERT_OK(writer->Init());
+
+  thread::ThreadPool* thread_pool =
+      new thread::ThreadPool(Env::Default(), "test_pool", 8);
+  std::atomic_int_fast64_t counter(0);
+  auto fn = [&writer, &counter]() {
+    const std::string file_path = absl::StrFormat(
+        "/home/tf_programs/program_%.3ld.py", counter.fetch_add(1));
+    SourceFile* source_file = new SourceFile();
+    source_file->set_file_path(file_path);
+    source_file->set_host_name("localhost.localdomain");
+    TF_ASSERT_OK(writer->WriteSourceFile(source_file));
+    TF_ASSERT_OK(writer->FlushNonExecutionFiles());
+  };
+  for (size_t i = 0; i < kConcurrentWrites; ++i) {
+    thread_pool->Schedule(fn);
+  }
+  delete thread_pool;
+
+  TF_ASSERT_OK(writer->Close());
+
+  std::vector<DebugEvent> actuals;
+  ReadDebugEventProtos(writer, DebugEventFileType::SOURCE_FILES, &actuals);
+  EXPECT_EQ(actuals.size(), kConcurrentWrites);
+  std::vector<std::string> file_paths;
+  std::vector<std::string> host_names;
+  for (size_t i = 0; i < kConcurrentWrites; ++i) {
+    file_paths.push_back(actuals[i].source_file().file_path());
+    host_names.push_back(actuals[i].source_file().host_name());
+  }
+  std::sort(file_paths.begin(), file_paths.end());
+  for (size_t i = 0; i < kConcurrentWrites; ++i) {
+    EXPECT_EQ(file_paths[i],
+              absl::StrFormat("/home/tf_programs/program_%.3ld.py", i));
+    EXPECT_EQ(host_names[i], "localhost.localdomain");
+  }
+}
+
+TEST_F(DebugEventsWriterTest, ConcurrentWriteCallsToTheDifferentFiles) {
+  const int32_t kConcurrentWrites = 30;
+  DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, DebugEventsWriter::kDefaultCyclicBufferSize);
+  TF_ASSERT_OK(writer->Init());
+
+  thread::ThreadPool* thread_pool =
+      new thread::ThreadPool(Env::Default(), "test_pool", 10);
+  std::atomic_int_fast32_t counter(0);
+  auto fn = [&writer, &counter]() {
+    const int32_t index = counter.fetch_add(1);
+    if (index % 3 == 0) {
+      SourceFile* source_file = new SourceFile();
+      source_file->set_file_path(
+          absl::StrFormat("/home/tf_programs/program_%.2d.py", index));
+      source_file->set_host_name("localhost.localdomain");
+      TF_ASSERT_OK(writer->WriteSourceFile(source_file));
+    } else if (index % 3 == 1) {
+      StackFrameWithId* stack_frame = new StackFrameWithId();
+      stack_frame->set_id(absl::StrFormat("e%.2d", index));
+      TF_ASSERT_OK(writer->WriteStackFrameWithId(stack_frame));
+    } else {
+      GraphOpCreation* op_creation = new GraphOpCreation();
+      op_creation->set_op_type("Log");
+      op_creation->set_op_name(absl::StrFormat("Log_%.2d", index));
+      TF_ASSERT_OK(writer->WriteGraphOpCreation(op_creation));
+    }
+  };
+  for (size_t i = 0; i < kConcurrentWrites; ++i) {
+    thread_pool->Schedule(fn);
+  }
+  delete thread_pool;
+
+  TF_ASSERT_OK(writer->Close());
+
+  std::vector<DebugEvent> actuals;
+  ReadDebugEventProtos(writer, DebugEventFileType::SOURCE_FILES, &actuals);
+  EXPECT_EQ(actuals.size(), kConcurrentWrites / 3);
+  std::vector<std::string> file_paths;
+  std::vector<std::string> host_names;
+  for (int32_t i = 0; i < kConcurrentWrites / 3; ++i) {
+    file_paths.push_back(actuals[i].source_file().file_path());
+    host_names.push_back(actuals[i].source_file().host_name());
+  }
+  std::sort(file_paths.begin(), file_paths.end());
+  for (int32_t i = 0; i < kConcurrentWrites / 3; ++i) {
+    EXPECT_EQ(file_paths[i],
+              absl::StrFormat("/home/tf_programs/program_%.2d.py", i * 3));
+    EXPECT_EQ(host_names[i], "localhost.localdomain");
+  }
+
+  ReadDebugEventProtos(writer, DebugEventFileType::STACK_FRAMES, &actuals);
+  EXPECT_EQ(actuals.size(), kConcurrentWrites / 3);
+  std::vector<std::string> stack_frame_ids;
+  for (int32_t i = 0; i < kConcurrentWrites / 3; ++i) {
+    stack_frame_ids.push_back(actuals[i].stack_frame_with_id().id());
+  }
+  std::sort(stack_frame_ids.begin(), stack_frame_ids.end());
+  for (int32_t i = 0; i < kConcurrentWrites / 3; ++i) {
+    EXPECT_EQ(stack_frame_ids[i], absl::StrFormat("e%.2d", i * 3 + 1));
+  }
+
+  ReadDebugEventProtos(writer, DebugEventFileType::GRAPHS, &actuals);
+  EXPECT_EQ(actuals.size(), kConcurrentWrites / 3);
+  std::vector<std::string> op_types;
+  std::vector<std::string> op_names;
+  for (int32_t i = 0; i < kConcurrentWrites / 3; ++i) {
+    op_types.push_back(actuals[i].graph_op_creation().op_type());
+    op_names.push_back(actuals[i].graph_op_creation().op_name());
+  }
+  std::sort(op_names.begin(), op_names.end());
+  for (int32_t i = 0; i < kConcurrentWrites / 3; ++i) {
+    EXPECT_EQ(op_types[i], "Log");
+    EXPECT_EQ(op_names[i], absl::StrFormat("Log_%.2d", i * 3 + 2));
+  }
+}
+
+TEST_F(DebugEventsWriterTest, WriteExecutionWithCyclicBufferNoFlush) {
+  // Verify that no writing to disk happens until the flushing method is called.
+  const size_t kCyclicBufferSize = 10;
+  DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, kCyclicBufferSize);
+  TF_ASSERT_OK(writer->Init());
+
+  // First, try writing and flushing more debug events than the capacity
+  // of the circular buffer, in a serial fashion.
+  for (size_t i = 0; i < kCyclicBufferSize * 2; ++i) {
+    Execution* execution = new Execution();
+    execution->set_op_type("Log");
+    execution->add_input_tensor_ids(i);
+    TF_ASSERT_OK(writer->WriteExecution(execution));
+  }
+
+  std::vector<DebugEvent> actuals;
+  // Before FlushExecutionFiles() is called, the file should be empty.
+  ReadDebugEventProtos(writer, DebugEventFileType::EXECUTION, &actuals);
+  EXPECT_EQ(actuals.size(), 0);
+
+  // Close the writer so the files can be safely deleted.
+  TF_ASSERT_OK(writer->Close());
+}
+
+TEST_F(DebugEventsWriterTest, WriteExecutionWithCyclicBufferFlush) {
+  // Verify that writing to disk happens when the flushing method is called.
+  const size_t kCyclicBufferSize = 10;
+  DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, kCyclicBufferSize);
+  TF_ASSERT_OK(writer->Init());
+
+  // First, try writing and flushing more debug events than the capacity
+  // of the circular buffer, in a serial fashion.
+  for (size_t i = 0; i < kCyclicBufferSize * 2; ++i) {
+    Execution* execution = new Execution();
+    execution->set_op_type("Log");
+    execution->add_input_tensor_ids(i);
+    TF_ASSERT_OK(writer->WriteExecution(execution));
+  }
+
+  TF_ASSERT_OK(writer->FlushExecutionFiles());
+
+  std::vector<DebugEvent> actuals;
+  // Expect there to be only the last kCyclicBufferSize debug events,
+  // and the order should be correct.
+  ReadDebugEventProtos(writer, DebugEventFileType::EXECUTION, &actuals);
+  EXPECT_EQ(actuals.size(), kCyclicBufferSize);
+  for (size_t i = 0; i < kCyclicBufferSize; ++i) {
+    EXPECT_EQ(actuals[i].execution().op_type(), "Log");
+    EXPECT_EQ(actuals[i].execution().input_tensor_ids().size(), 1);
+    EXPECT_EQ(actuals[i].execution().input_tensor_ids()[0],
+              kCyclicBufferSize + i);
+  }
+
+  // Second, write more than the capacity of the circular buffer,
+  // in a concurrent fashion.
+  thread::ThreadPool* thread_pool =
+      new thread::ThreadPool(Env::Default(), "test_pool", 8);
+  std::atomic_int_fast64_t counter(0);
+  auto fn = [&writer, &counter]() {
+    Execution* execution = new Execution();
+    execution->set_op_type("Abs");
+    execution->add_input_tensor_ids(counter.fetch_add(1));
+    TF_ASSERT_OK(writer->WriteExecution(execution));
+  };
+  for (size_t i = 0; i < kCyclicBufferSize * 2; ++i) {
+    thread_pool->Schedule(fn);
+  }
+  delete thread_pool;
+  TF_ASSERT_OK(writer->Close());
+
+  ReadDebugEventProtos(writer, DebugEventFileType::EXECUTION, &actuals);
+  // NOTE: This includes the files from the first stage above, because the
+  // .execution file hasn't changed.
+  EXPECT_EQ(actuals.size(), kCyclicBufferSize * 2);
+  for (size_t i = 0; i < kCyclicBufferSize; ++i) {
+    const size_t index = i + kCyclicBufferSize;
+    EXPECT_EQ(actuals[index].execution().op_type(), "Abs");
+    EXPECT_EQ(actuals[index].execution().input_tensor_ids().size(), 1);
+    EXPECT_GE(actuals[index].execution().input_tensor_ids()[0], 0);
+    EXPECT_LE(actuals[index].execution().input_tensor_ids()[0],
+              kCyclicBufferSize * 2);
+  }
+
+  // Verify no cross-talk.
+  ReadDebugEventProtos(writer, DebugEventFileType::SOURCE_FILES, &actuals);
+  EXPECT_EQ(actuals.size(), 0);
+  ReadDebugEventProtos(writer, DebugEventFileType::STACK_FRAMES, &actuals);
+  EXPECT_EQ(actuals.size(), 0);
+  ReadDebugEventProtos(writer, DebugEventFileType::GRAPHS, &actuals);
+  EXPECT_EQ(actuals.size(), 0);
+  ReadDebugEventProtos(writer, DebugEventFileType::GRAPH_EXECUTION_TRACES,
+                       &actuals);
+  EXPECT_EQ(actuals.size(), 0);
+}
+
+TEST_F(DebugEventsWriterTest, WriteGrahExecutionTraceWithCyclicBufferNoFlush) {
+  // Check no writing to disk happens before the flushing method is called.
+  const size_t kCyclicBufferSize = 10;
+  DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, kCyclicBufferSize);
+  TF_ASSERT_OK(writer->Init());
+
+  // First, try writing and flushing more debug events than the capacity
+  // of the circular buffer, in a serial fashion.
+  for (size_t i = 0; i < kCyclicBufferSize * 2; ++i) {
+    GraphExecutionTrace* trace = new GraphExecutionTrace();
+    trace->set_tfdbg_context_id(absl::StrFormat("graph_%.2ld", i));
+    TF_ASSERT_OK(writer->WriteGraphExecutionTrace(trace));
+  }
+
+  std::vector<DebugEvent> actuals;
+  // Before FlushExecutionFiles() is called, the file should be empty.
+  ReadDebugEventProtos(writer, DebugEventFileType::GRAPH_EXECUTION_TRACES,
+                       &actuals);
+  EXPECT_EQ(actuals.size(), 0);
+
+  // Close the writer so the files can be safely deleted.
+  TF_ASSERT_OK(writer->Close());
+}
+
+TEST_F(DebugEventsWriterTest, WriteGrahExecutionTraceWithoutPreviousInitCall) {
+  const size_t kCyclicBufferSize = -1;
+  DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, kCyclicBufferSize);
+  // NOTE(cais): `writer->Init()` is not called here before
+  // WriteGraphExecutionTrace() is called. This test checks that this is okay
+  // and the `GraphExecutionTrace` gets written correctly even without `Init()`
+  // being called first. This scenario can happen when a TF Graph with tfdbg
+  // debug ops are executed on a remote TF server.
+
+  GraphExecutionTrace* trace = new GraphExecutionTrace();
+  trace->set_tfdbg_context_id(absl::StrFormat("graph_0"));
+  TF_ASSERT_OK(writer->WriteGraphExecutionTrace(trace));
+  TF_ASSERT_OK(writer->FlushExecutionFiles());
+
+  std::vector<DebugEvent> actuals;
+  ReadDebugEventProtos(writer, DebugEventFileType::GRAPH_EXECUTION_TRACES,
+                       &actuals);
+  EXPECT_EQ(actuals.size(), 1);
+  EXPECT_EQ(actuals[0].graph_execution_trace().tfdbg_context_id(), "graph_0");
+
+  // Close the writer so the files can be safely deleted.
+  TF_ASSERT_OK(writer->Close());
+}
+
+TEST_F(DebugEventsWriterTest, WriteGrahExecutionTraceWithCyclicBufferFlush) {
+  const size_t kCyclicBufferSize = 10;
+  DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, kCyclicBufferSize);
+  TF_ASSERT_OK(writer->Init());
+
+  // First, try writing and flushing more debug events than the capacity
+  // of the circular buffer, in a serial fashion.
+  for (size_t i = 0; i < kCyclicBufferSize * 2; ++i) {
+    GraphExecutionTrace* trace = new GraphExecutionTrace();
+    trace->set_tfdbg_context_id(absl::StrFormat("graph_%.2ld", i));
+    TF_ASSERT_OK(writer->WriteGraphExecutionTrace(trace));
+  }
+
+  TF_ASSERT_OK(writer->FlushExecutionFiles());
+
+  std::vector<DebugEvent> actuals;
+  // Expect there to be only the last kCyclicBufferSize debug events,
+  // and the order should be correct.
+  ReadDebugEventProtos(writer, DebugEventFileType::GRAPH_EXECUTION_TRACES,
+                       &actuals);
+  EXPECT_EQ(actuals.size(), kCyclicBufferSize);
+  for (size_t i = 0; i < kCyclicBufferSize; ++i) {
+    EXPECT_EQ(actuals[i].graph_execution_trace().tfdbg_context_id(),
+              absl::StrFormat("graph_%.2ld", i + kCyclicBufferSize));
+  }
+
+  // Second, write more than the capacity of the circular buffer,
+  // in a concurrent fashion.
+  thread::ThreadPool* thread_pool =
+      new thread::ThreadPool(Env::Default(), "test_pool", 8);
+  std::atomic_int_fast64_t counter(0);
+  auto fn = [&writer, &counter]() {
+    GraphExecutionTrace* trace = new GraphExecutionTrace();
+    trace->set_tfdbg_context_id(
+        absl::StrFormat("new_graph_%.2ld", counter.fetch_add(1)));
+    TF_ASSERT_OK(writer->WriteGraphExecutionTrace(trace));
+  };
+  for (size_t i = 0; i < kCyclicBufferSize * 2; ++i) {
+    thread_pool->Schedule(fn);
+  }
+  delete thread_pool;
+  TF_ASSERT_OK(writer->Close());
+
+  ReadDebugEventProtos(writer, DebugEventFileType::GRAPH_EXECUTION_TRACES,
+                       &actuals);
+  // NOTE: This includes the files from the first stage above, because the
+  // .graph_execution_traces file hasn't changed.
+  EXPECT_EQ(actuals.size(), kCyclicBufferSize * 2);
+  for (size_t i = 0; i < kCyclicBufferSize; ++i) {
+    const size_t index = i + kCyclicBufferSize;
+    EXPECT_EQ(actuals[index].graph_execution_trace().tfdbg_context_id().find(
+                  "new_graph_"),
+              0);
+  }
+
+  // Verify no cross-talk.
+  ReadDebugEventProtos(writer, DebugEventFileType::SOURCE_FILES, &actuals);
+  EXPECT_EQ(actuals.size(), 0);
+  ReadDebugEventProtos(writer, DebugEventFileType::STACK_FRAMES, &actuals);
+  EXPECT_EQ(actuals.size(), 0);
+  ReadDebugEventProtos(writer, DebugEventFileType::GRAPHS, &actuals);
+  EXPECT_EQ(actuals.size(), 0);
+  ReadDebugEventProtos(writer, DebugEventFileType::EXECUTION, &actuals);
+  EXPECT_EQ(actuals.size(), 0);
+}
+
+TEST_F(DebugEventsWriterTest, RegisterDeviceAndGetIdTrace) {
+  DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, DebugEventsWriter::kDefaultCyclicBufferSize);
+  TF_ASSERT_OK(writer->Init());
+
+  // Register and get some device IDs in a concurrent fashion.
+  thread::ThreadPool* thread_pool =
+      new thread::ThreadPool(Env::Default(), "test_pool", 8);
+  int device_ids[8];
+  for (int i = 0; i < 8; ++i) {
+    thread_pool->Schedule([i, &writer, &device_ids]() {
+      const std::string device_name = absl::StrFormat(
+          "/job:localhost/replica:0/task:0/device:GPU:%d", i % 4);
+      device_ids[i] = writer->RegisterDeviceAndGetId(device_name);
+    });
+  }
+  delete thread_pool;
+  TF_ASSERT_OK(writer->FlushNonExecutionFiles());
+  TF_ASSERT_OK(writer->Close());
+
+  // There should be only 4 unique device IDs, because there are only 4 unique
+  // device names.
+  EXPECT_EQ(device_ids[0], device_ids[4]);
+  EXPECT_EQ(device_ids[1], device_ids[5]);
+  EXPECT_EQ(device_ids[2], device_ids[6]);
+  EXPECT_EQ(device_ids[3], device_ids[7]);
+  // Assert that the four device IDs are all unique.
+  EXPECT_EQ(absl::flat_hash_set<int>(device_ids, device_ids + 8).size(), 4);
+
+  std::vector<DebugEvent> actuals;
+  ReadDebugEventProtos(writer, DebugEventFileType::GRAPHS, &actuals);
+  // Due to the `% 4`, there are only 4 unique device names, even though there
+  // are 8 threads each calling `RegisterDeviceAndGetId`.
+  EXPECT_EQ(actuals.size(), 4);
+  for (const DebugEvent& actual : actuals) {
+    const std::string& device_name = actual.debugged_device().device_name();
+    int device_index = -1;
+    CHECK(absl::SimpleAtoi(device_name.substr(strlen(
+                               "/job:localhost/replica:0/task:0/device:GPU:")),
+                           &device_index));
+    EXPECT_EQ(actual.debugged_device().device_id(), device_ids[device_index]);
+  }
+}
+
+TEST_F(DebugEventsWriterTest, DisableCyclicBufferBehavior) {
+  const size_t kCyclicBufferSize = 0;  // A value <= 0 disables cyclic behavior.
+  DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, kCyclicBufferSize);
+  TF_ASSERT_OK(writer->Init());
+
+  const size_t kNumEvents = 20;
+
+  for (size_t i = 0; i < kNumEvents; ++i) {
+    Execution* execution = new Execution();
+    execution->set_op_type("Log");
+    execution->add_input_tensor_ids(i);
+    TF_ASSERT_OK(writer->WriteExecution(execution));
+  }
+  TF_ASSERT_OK(writer->FlushExecutionFiles());
+
+  std::vector<DebugEvent> actuals;
+  ReadDebugEventProtos(writer, DebugEventFileType::EXECUTION, &actuals);
+  EXPECT_EQ(actuals.size(), kNumEvents);
+  for (size_t i = 0; i < kNumEvents; ++i) {
+    EXPECT_EQ(actuals[i].execution().op_type(), "Log");
+    EXPECT_EQ(actuals[i].execution().input_tensor_ids().size(), 1);
+    EXPECT_EQ(actuals[i].execution().input_tensor_ids()[0], i);
+  }
+
+  for (size_t i = 0; i < kNumEvents; ++i) {
+    GraphExecutionTrace* trace = new GraphExecutionTrace();
+    trace->set_tfdbg_context_id(absl::StrFormat("graph_%.2ld", i));
+    TF_ASSERT_OK(writer->WriteGraphExecutionTrace(trace));
+  }
+  TF_ASSERT_OK(writer->FlushExecutionFiles());
+
+  ReadDebugEventProtos(writer, DebugEventFileType::GRAPH_EXECUTION_TRACES,
+                       &actuals);
+  EXPECT_EQ(actuals.size(), kNumEvents);
+  for (size_t i = 0; i < kNumEvents; ++i) {
+    EXPECT_EQ(actuals[i].graph_execution_trace().tfdbg_context_id(),
+              absl::StrFormat("graph_%.2ld", i));
+  }
+
+  // Close the writer so the files can be safely deleted.
+  TF_ASSERT_OK(writer->Close());
+}
+
+}  // namespace tfdbg
+}  // namespace tensorflow

tensorflow/core/util/determinism.h

@@ -0,0 +1,29 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_UTIL_DETERMINISM_H_
+#define TENSORFLOW_CORE_UTIL_DETERMINISM_H_
+
+#include "xla/tsl/util/determinism.h"
+
+namespace tensorflow {
+
+using tsl::EnableOpDeterminism;
+using tsl::OpDeterminismRequired;
+using tsl::OpOrderDeterminismRequired;
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_UTIL_DETERMINISM_H_

tensorflow/core/util/device_name_utils.h

@@ -0,0 +1,27 @@
+/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_UTIL_DEVICE_NAME_UTILS_H_
+#define TENSORFLOW_CORE_UTIL_DEVICE_NAME_UTILS_H_
+
+#include "xla/tsl/util/device_name_utils.h"
+
+namespace tensorflow {
+// NOLINTBEGIN(misc-unused-using-decls)
+using tsl::DeviceNameUtils;
+// NOLINTEND(misc-unused-using-decls)
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_UTIL_DEVICE_NAME_UTILS_H_

tensorflow/core/util/dump_graph.cc

@@ -0,0 +1,321 @@
+/* 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.
+==============================================================================*/
+
+// Helper functions for dumping Graphs, GraphDefs, and FunctionDefs to files for
+// debugging.
+
+#include "tensorflow/core/util/dump_graph.h"
+
+#include <functional>
+#include <memory>
+#include <unordered_map>
+
+#include "absl/strings/match.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "xla/tsl/util/env_var.h"
+#include "tensorflow/core/lib/strings/proto_serialization.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/file_system.h"
+#include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/platform/path.h"
+
+namespace tensorflow {
+
+namespace {
+using strings::StrCat;
+
+struct NameCounts {
+  mutex counts_mutex;
+  std::unordered_map<std::string, int> counts;
+};
+
+std::string MakeUniqueFilename(std::string name,
+                               const std::string& suffix = ".pbtxt") {
+  static NameCounts& instance = *new NameCounts;
+
+  // Remove illegal characters from `name`.
+  for (int i = 0; i < name.size(); ++i) {
+    char ch = name[i];
+    if (ch == '/' || ch == '[' || ch == ']' || ch == '*' || ch == '?' ||
+        ch == '\\') {
+      name[i] = '_';
+    }
+  }
+
+  int count;
+  {
+    mutex_lock lock(instance.counts_mutex);
+    count = instance.counts[name]++;
+  }
+
+  std::string filename = name;
+  if (count > 0) {
+    absl::StrAppend(&filename, "_", count);
+  }
+  absl::StrAppend(&filename, suffix);
+  return filename;
+}
+
+struct GraphDumperConfig {
+  mutex mu;
+
+  // The dumper and suffix configured.
+  struct Config {
+    bool IsSet() const { return dumper != nullptr; }
+    std::function<absl::Status(const Graph& graph,
+                               const FunctionLibraryDefinition* flib_def,
+                               WritableFile*)>
+        dumper = nullptr;
+    std::string suffix = ".pbtxt";
+  } config TF_GUARDED_BY(mu);
+
+  // Returns whether a custom dumper is set.
+  bool IsSet() TF_LOCKS_EXCLUDED(mu) {
+    mutex_lock lock(mu);
+    return config.IsSet();
+  }
+};
+
+GraphDumperConfig& GetGraphDumperConfig() {
+  static GraphDumperConfig config;
+  return config;
+}
+
+std::string GetDumpGraphFormatLowerCase() {
+  std::string fmt;
+  absl::Status status =
+      tsl::ReadStringFromEnvVar("TF_DUMP_GRAPH_FMT", "TXT", &fmt);
+  if (!status.ok()) {
+    LOG(WARNING) << "Failed to read TF_DUMP_GRAPH_FMT: " << status;
+    return "txt";
+  }
+  fmt = absl::AsciiStrToLower(fmt);
+  return fmt;
+}
+
+std::string GetDumpGraphSuffix() {
+  std::string fmt = GetDumpGraphFormatLowerCase();
+  if (fmt == "txt") {
+    return ".pbtxt";
+  } else if (fmt == "bin") {
+    return ".pb";
+  } else {
+    return ".pbtxt";
+  }
+}
+
+// WritableFile that simply prints to stderr.
+class StderrWritableFile : public WritableFile {
+ public:
+  StderrWritableFile() = default;
+
+  absl::Status Append(absl::string_view data) override {
+    fprintf(stderr, "%.*s", static_cast<int>(data.size()), data.data());
+    return absl::OkStatus();
+  }
+
+  absl::Status Close() override { return absl::OkStatus(); }
+
+  absl::Status Flush() override {
+    fflush(stderr);
+    return absl::OkStatus();
+  }
+
+  absl::Status Name(absl::string_view* result) const override {
+    *result = "stderr";
+    return absl::OkStatus();
+  }
+
+  absl::Status Sync() override { return absl::OkStatus(); }
+
+  absl::Status Tell(int64_t* position) override {
+    return absl::UnimplementedError("Stream not seekable");
+  }
+};
+
+absl::Status CreateWritableFile(Env* env, const std::string& dirname,
+                                const std::string& name,
+                                const std::string& suffix,
+                                std::string* filepath,
+                                std::unique_ptr<WritableFile>* file) {
+  std::string dir;
+  if (!dirname.empty()) {
+    dir = dirname;
+  } else {
+    const char* prefix = getenv("TF_DUMP_GRAPH_PREFIX");
+    if (prefix != nullptr) dir = prefix;
+  }
+  if (dir.empty()) {
+    LOG(WARNING)
+        << "Failed to dump " << name << " because dump location is not "
+        << " specified through either TF_DUMP_GRAPH_PREFIX environment "
+        << "variable or function argument.";
+    return absl::InvalidArgumentError("TF_DUMP_GRAPH_PREFIX not specified");
+  }
+
+  if (absl::EqualsIgnoreCase(dir, "sponge") ||
+      absl::EqualsIgnoreCase(dir, "test_undeclared_outputs_dir")) {
+    if (!io::GetTestUndeclaredOutputsDir(&dir)) {
+      LOG(WARNING) << "TF_DUMP_GRAPH_PREFIX=sponge, but "
+                      "TEST_UNDECLARED_OUTPUT_DIRS is not set, dumping to log";
+      dir = "-";
+    }
+  }
+
+  *filepath = "NULL";
+  if (dir == "-") {
+    *file = std::make_unique<StderrWritableFile>();
+    *filepath = "(stderr)";
+    return absl::OkStatus();
+  }
+
+  TF_RETURN_IF_ERROR(env->RecursivelyCreateDir(dir));
+  *filepath = io::JoinPath(dir, MakeUniqueFilename(name, suffix));
+  return env->NewWritableFile(*filepath, file);
+}
+
+absl::Status WriteProtoToUniqueFile(const tensorflow::protobuf::Message& proto,
+                                    WritableFile* file) {
+  std::string s;
+  std::string format = GetDumpGraphFormatLowerCase();
+  if (format == "txt" &&
+      !::tensorflow::protobuf::TextFormat::PrintToString(proto, &s)) {
+    return absl::FailedPreconditionError("Unable to convert proto to text.");
+  } else if (format == "bin" && !SerializeToStringDeterministic(proto, &s)) {
+    return absl::FailedPreconditionError(
+        "Failed to serialize proto to string.");
+  } else if (format != "txt" && format != "bin") {
+    return absl::FailedPreconditionError(
+        absl ::StrCat("Unknown format: ", format));
+  }
+  TF_RETURN_IF_ERROR(file->Append(s));
+  absl::string_view name;
+  TF_RETURN_IF_ERROR(file->Name(&name));
+  VLOG(5) << name;
+  VLOG(5) << s;
+  return file->Close();
+}
+
+absl::Status WriteProtoToUniqueFile(
+    const tensorflow::protobuf::MessageLite& proto, WritableFile* file) {
+  std::string s;
+  if (!SerializeToStringDeterministic(proto, &s)) {
+    return absl::InternalError("Failed to serialize proto to string.");
+  }
+  absl::string_view name;
+  TF_RETURN_IF_ERROR(file->Name(&name));
+  VLOG(5) << name;
+  VLOG(5) << s;
+  TF_RETURN_IF_ERROR(file->Append(s));
+  return file->Close();
+}
+
+}  // anonymous namespace
+
+std::string DumpToFile(const std::string& name, const std::string& dirname,
+                       const std::string& suffix, absl::string_view type_name,
+                       std::function<absl::Status(WritableFile*)> dumper) {
+  std::string filepath;
+  std::unique_ptr<WritableFile> file;
+  absl::Status status = CreateWritableFile(Env::Default(), dirname, name,
+                                           suffix, &filepath, &file);
+  if (!status.ok()) {
+    return absl::StrCat("(failed to create writable file: ", status.ToString(),
+                        ")");
+  }
+
+  status = dumper(file.get());
+  if (!status.ok()) {
+    return absl::StrCat("(failed to dump ", type_name, " to '", filepath,
+                        "': ", status.ToString(), ")");
+  }
+  LOG(INFO) << "Dumped " << type_name << " to " << filepath;
+  return filepath;
+}
+
+void SetGraphDumper(
+    std::function<absl::Status(const Graph& graph,
+                               const FunctionLibraryDefinition* flib_def,
+                               WritableFile*)>
+        dumper,
+    std::string suffix) {
+  GraphDumperConfig& dumper_config = GetGraphDumperConfig();
+  mutex_lock lock(dumper_config.mu);
+  dumper_config.config.dumper = dumper;
+  dumper_config.config.suffix = suffix;
+}
+
+std::string DumpGraphDefToFile(const std::string& name,
+                               GraphDef const& graph_def,
+                               const std::string& dirname) {
+  return DumpToFile(name, dirname, GetDumpGraphSuffix(), "Graph",
+                    [&](WritableFile* file) {
+                      return WriteProtoToUniqueFile(graph_def, file);
+                    });
+}
+
+std::string DumpCostGraphDefToFile(const std::string& name,
+                                   CostGraphDef const& graph_def,
+                                   const std::string& dirname) {
+  return DumpToFile(name, dirname, GetDumpGraphSuffix(), "Graph",
+                    [&](WritableFile* file) {
+                      return WriteProtoToUniqueFile(graph_def, file);
+                    });
+}
+
+std::string DumpGraphToFile(const std::string& name, Graph const& graph,
+                            const FunctionLibraryDefinition* flib_def,
+                            const std::string& dirname) {
+  auto& dumper_config = GetGraphDumperConfig();
+  if (dumper_config.IsSet()) {
+    GraphDumperConfig::Config config;
+    {
+      mutex_lock lock(dumper_config.mu);
+      config = dumper_config.config;
+    }
+    if (config.IsSet()) {
+      return DumpToFile(name, dirname, config.suffix, "Graph",
+                        [&](WritableFile* file) {
+                          return config.dumper(graph, flib_def, file);
+                        });
+    }
+  }
+
+  GraphDef graph_def;
+  graph.ToGraphDef(&graph_def);
+  if (flib_def) {
+    *graph_def.mutable_library() = flib_def->ToProto();
+  }
+  return DumpGraphDefToFile(name, graph_def, dirname);
+}
+
+std::string DumpFunctionDefToFile(const std::string& name,
+                                  FunctionDef const& fdef,
+                                  const std::string& dirname) {
+  return DumpToFile(
+      name, dirname, GetDumpGraphSuffix(), "FunctionDef",
+      [&](WritableFile* file) { return WriteProtoToUniqueFile(fdef, file); });
+}
+
+std::string DumpProtoToFile(const std::string& name,
+                            tensorflow::protobuf::Message const& proto,
+                            const std::string& dirname) {
+  return DumpToFile(
+      name, dirname, GetDumpGraphSuffix(), proto.GetTypeName(),
+      [&](WritableFile* file) { return WriteProtoToUniqueFile(proto, file); });
+}
+
+}  // namespace tensorflow

tensorflow/core/util/dump_graph.h

@@ -0,0 +1,91 @@
+/* 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.
+==============================================================================*/
+
+// Helper functions for dumping Graphs, GraphDefs, and FunctionDefs to files for
+// debugging.
+
+#ifndef TENSORFLOW_CORE_UTIL_DUMP_GRAPH_H_
+#define TENSORFLOW_CORE_UTIL_DUMP_GRAPH_H_
+
+#include <functional>
+#include <string>
+
+#include "absl/strings/string_view.h"
+#include "tensorflow/core/framework/cost_graph.pb.h"
+#include "tensorflow/core/framework/function.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/graph/graph.h"
+
+namespace tensorflow {
+
+// Dumps 'graph_def' to a file, as a GraphDef text or binary proto. Returns the
+// file name chosen. The format is determined by the TF_DUMP_GRAPH_FMT
+// environment variable (TXT or BIN).
+//
+// If the TF_DUMP_GRAPH_PREFIX environment variable is "-", then instead the
+// GraphDef will be logged (using the LOG() macro).
+//
+// Automatically picks a file name. Prefixes 'name' with the value of the
+// TF_DUMP_GRAPH_PREFIX environment variable if 'dirname' is empty, and suffixes
+// 'name' with '.pbtxt' or '.pb'. If a graph has already been dumped by
+// this process with the same name, suffixes with "_n.pb(txt)", where 'n' is a
+// sequence number.
+std::string DumpGraphDefToFile(const std::string& name,
+                               GraphDef const& graph_def,
+                               const std::string& dirname = "");
+
+// Similar to DumpGraphDefToFile, use CostGraphDef instead of GraphDef.
+std::string DumpCostGraphDefToFile(const std::string& name,
+                                   CostGraphDef const& graph_def,
+                                   const std::string& dirname = "");
+
+// Similar to DumpGraphDefToFile, but builds the GraphDef to dump from a 'graph'
+// and an optional function library 'flib_def'. Returns the file name chosen.
+std::string DumpGraphToFile(const std::string& name, Graph const& graph,
+                            const FunctionLibraryDefinition* flib_def = nullptr,
+                            const std::string& dirname = "");
+
+// Similar to DumpGraphDefToFile, but dumps a function as a FunctionDef text
+// proto. Returns the file name chosen.
+std::string DumpFunctionDefToFile(const std::string& name,
+                                  FunctionDef const& fdef,
+                                  const std::string& dirname = "");
+
+// Similar to DumpGraphDefToFile, but dumps a proto of any type. Returns the
+// file name chosen.
+std::string DumpProtoToFile(const std::string& name,
+                            tensorflow::protobuf::Message const& proto,
+                            const std::string& dirname = "");
+
+// Sets a custom Graph dumper. If set, this dumper will be used to dump graphs
+// instead via DumpGraphToFile. As the custom dumper may not produce protobufs,
+// allow specifying a file suffix/extension too.
+void SetGraphDumper(
+    std::function<absl::Status(const Graph& graph,
+                               const FunctionLibraryDefinition* flib_def,
+                               WritableFile*)>
+        dumper,
+    std::string suffix = ".pbtxt");
+
+// Dump data to a file.
+// This function will create a WritableFile and pass it to the dumper.
+// The dumper callback will be responsible for writing data to the file.
+std::string DumpToFile(const std::string& name, const std::string& dirname,
+                       const std::string& suffix, absl::string_view type_name,
+                       std::function<absl::Status(WritableFile*)> dumper);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_UTIL_DUMP_GRAPH_H_

tensorflow/core/util/dump_graph_test.cc

@@ -0,0 +1,96 @@
+/* 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/util/dump_graph.h"
+
+#include "absl/strings/match.h"
+#include "xla/tsl/lib/core/status_test_util.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/graph/graph.h"
+#include "tensorflow/core/graph/node_builder.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/path.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/platform/types.h"
+#include "tsl/platform/status.h"
+
+namespace tensorflow {
+namespace {
+
+TEST(DumpGraph, DumpGraphToFileSuccess) {
+  Graph graph(OpRegistry::Global());
+  Node* node;
+  TF_CHECK_OK(NodeBuilder("A", "NoOp").Finalize(&graph, &node));
+
+  setenv("TF_DUMP_GRAPH_PREFIX", testing::TmpDir().c_str(), 1);
+  std::string ret = DumpGraphToFile("graph", graph);
+  EXPECT_EQ(ret, io::JoinPath(testing::TmpDir(), "graph.pbtxt"));
+  ret = DumpGraphToFile("graph", graph);
+  EXPECT_EQ(ret, io::JoinPath(testing::TmpDir(), "graph_1.pbtxt"));
+
+  GraphDef gdef;
+  TF_ASSERT_OK(ReadTextProto(
+      Env::Default(), io::JoinPath(testing::TmpDir(), "graph.pbtxt"), &gdef));
+  std::string read, written;
+  gdef.AppendToString(&read);
+  graph.ToGraphDefDebug().AppendToString(&written);
+  EXPECT_EQ(read, written);
+}
+
+TEST(DumpGraph, DumpGraphToFileNoEnvPrefix) {
+  Graph graph(OpRegistry::Global());
+  unsetenv("TF_DUMP_GRAPH_PREFIX");
+  std::string ret = DumpGraphToFile("graph", graph);
+  EXPECT_TRUE(absl::StrContains(ret, "TF_DUMP_GRAPH_PREFIX not specified"));
+}
+
+TEST(DumpGraph, DumpFunctionDefToFileSuccess) {
+  FunctionDef fdef;
+  setenv("TF_DUMP_GRAPH_PREFIX", testing::TmpDir().c_str(), 1);
+  std::string ret = DumpFunctionDefToFile("function", fdef);
+  EXPECT_EQ(ret, io::JoinPath(testing::TmpDir(), "function.pbtxt"));
+}
+
+TEST(DumpGraph, DumpProtoToFileSuccess) {
+  NodeDef ndef_in;
+  ndef_in.set_name("foo");
+  ndef_in.set_op("bar");
+  ndef_in.add_input("baz");
+  ndef_in.set_device("cpu:0");
+
+  setenv("TF_DUMP_GRAPH_PREFIX", testing::TmpDir().c_str(), 1);
+  setenv("TF_DUMP_GRAPH_FMT", "TXT", 1);
+  std::string expected_filepath =
+      io::JoinPath(testing::TmpDir(), "node_def.pbtxt");
+  std::string actual_filepath = DumpProtoToFile("node_def", ndef_in);
+  EXPECT_EQ(expected_filepath, actual_filepath);
+
+  NodeDef ndef_out;
+  TF_ASSERT_OK(ReadTextProto(Env::Default(), expected_filepath, &ndef_out));
+  EXPECT_EQ(ndef_in.DebugString(), ndef_out.DebugString());
+
+  setenv("TF_DUMP_GRAPH_FMT", "BIN", 1);
+  std::string ret = DumpProtoToFile("node_def", ndef_in);
+  EXPECT_EQ(ret, io::JoinPath(testing::TmpDir(), "node_def_1.pb"));
+  TF_ASSERT_OK(ReadBinaryProto(Env::Default(), ret, &ndef_out));
+  EXPECT_EQ(ndef_out.DebugString(), ndef_in.DebugString());
+
+  setenv("TF_DUMP_GRAPH_FMT", "unknown", 1);
+  ret = DumpProtoToFile("node_def", ndef_in);
+  EXPECT_TRUE(absl::StrContains(ret, "Unknown format"));
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/core/util/einsum_op_util.cc

@@ -0,0 +1,141 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/util/einsum_op_util.h"
+
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/strings/str_split.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/lib/gtl/inlined_vector.h"
+
+namespace tensorflow {
+
+absl::Status ValidateEinsumEquation(
+    const std::string& equation,
+    absl::InlinedVector<std::string, 2UL>* input_subscripts,
+    std::string* output_subscript) {
+  absl::InlinedVector<std::string, 2UL> inputs_and_output_subscripts =
+      absl::StrSplit(equation, "->");
+  if (inputs_and_output_subscripts.size() != 2) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "Expecting exactly one '->' in einsum equation: ", equation));
+  }
+  *output_subscript = std::move(inputs_and_output_subscripts[1]);
+  *input_subscripts =
+      absl::StrSplit(std::move(inputs_and_output_subscripts[0]), ',');
+  if (input_subscripts->size() != 1 && input_subscripts->size() != 2) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("Expecting 1 or 2 input subscripts in equation '",
+                     equation, "' but got: ", input_subscripts->size()));
+  }
+  return absl::OkStatus();
+}
+
+// Returns the EinsumDimensionType given whether the corresponding label is
+// present in exactly one input subscript (is_unique) and whether it is absent
+// from the output subscripts (is_removed). Does not handle broadcasting
+// dimensions.
+EinsumDimensionType GetDimensionType(bool is_removed, bool is_unique) {
+  if (!is_removed && !is_unique)
+    return kBatch;
+  else if (!is_removed && is_unique)
+    return kFree;
+  else if (is_removed && !is_unique)
+    return kContract;
+  else  // is_removed && is_unique
+    return kReduce;
+}
+
+// Maps the character labels to consecutive integers.
+void MapToLabels(const std::string& subscript, Labels* labels,
+                 absl::flat_hash_map<char, int>* label_mapping) {
+  for (int i = 0; i < subscript.size(); ++i) {
+    const char label_char = subscript[i];
+    if (label_char == '.') {
+      labels->push_back(kEllipsisLabel);
+      i += 2;  // Skip next 2 characters as well.
+      continue;
+    }
+    if (!label_mapping->contains(label_char)) {
+      const int next_label = label_mapping->size();
+      (*label_mapping)[label_char] = next_label;
+    }
+    const int mapped_label = (*label_mapping)[label_char];
+    labels->push_back(mapped_label);
+  }
+}
+
+absl::Status ParseEinsumEquation(
+    const std::string& equation, OperandLabels* input_labels,
+    Labels* output_labels, std::vector<EinsumDimensionType>* label_types,
+    OperandLabelCounts* input_label_counts, LabelCounts* output_label_counts,
+    absl::InlinedVector<bool, 2UL>* input_has_ellipsis,
+    bool* output_has_ellipsis) {
+  absl::InlinedVector<std::string, 2UL> input_str;
+  std::string output_str;
+  TF_RETURN_IF_ERROR(ValidateEinsumEquation(equation, &input_str, &output_str));
+
+  // Temporary map from single character labels to (consecutive) integer labels.
+  absl::flat_hash_map<char, int> label_mapping;
+  int num_inputs = input_str.size();
+  input_labels->resize(num_inputs);
+
+  // Map from single characters to integer labels.
+  for (int i = 0; i < num_inputs; ++i) {
+    MapToLabels(input_str[i], &input_labels->at(i), &label_mapping);
+  }
+  MapToLabels(output_str, output_labels, &label_mapping);
+
+  // Compute counts for input and output labels.
+  int num_labels = label_mapping.size();
+  input_label_counts->resize(num_inputs);
+  input_has_ellipsis->resize(num_inputs);
+  for (int i = 0; i < num_inputs; ++i) {
+    input_label_counts->at(i).resize(num_labels);
+    input_has_ellipsis->at(i) = false;
+    for (const int label : input_labels->at(i)) {
+      if (label != kEllipsisLabel)
+        input_label_counts->at(i)[label] += 1;
+      else
+        input_has_ellipsis->at(i) = true;
+    }
+  }
+  output_label_counts->resize(num_labels);
+  *output_has_ellipsis = false;
+  for (const int label : *output_labels) {
+    if (label != kEllipsisLabel)
+      output_label_counts->at(label) += 1;
+    else
+      *output_has_ellipsis = true;
+  }
+
+  // Map each label to a unique EinsumDimensionType.
+  label_types->resize(num_labels);
+  for (int label = 0; label < num_labels; ++label) {
+    if (label == kEllipsisLabel) continue;
+    bool removed = (*output_label_counts)[label] == 0;
+    bool unique = num_inputs == 1 || (*input_label_counts)[0][label] == 0 ||
+                  (*input_label_counts)[1][label] == 0;
+    (*label_types)[label] = GetDimensionType(removed, unique);
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace tensorflow

tensorflow/core/util/einsum_op_util.h

@@ -0,0 +1,73 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_CORE_UTIL_EINSUM_OP_UTIL_H_
+#define TENSORFLOW_CORE_UTIL_EINSUM_OP_UTIL_H_
+
+#include <vector>
+
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/lib/gtl/inlined_vector.h"
+
+namespace tensorflow {
+
+using Labels = absl::InlinedVector<int, 8UL>;
+using OperandLabels = absl::InlinedVector<Labels, 2UL>;
+using LabelCounts = absl::InlinedVector<int, 8UL>;
+using OperandLabelCounts = absl::InlinedVector<LabelCounts, 2UL>;
+
+// Dummy axis label used to denote an ellipsis in an input or output subscript.
+constexpr int kEllipsisLabel = -1;
+
+// Each dimension is categorized into exactly one of five types based on
+// whether its corresponding label is present in the input and/or the output
+// subscripts.
+enum EinsumDimensionType {
+  // Batch dimensions are those present in two inputs as well as the output.
+  // They are part of the batch dimensions during Tensor contraction. Such
+  // dimensions may be broadcasting dimensions (those mapping to ellipsis)
+  // or explicit batch dimensions corresponding to named axis labels.
+  kBroadcasting = 0,
+  kBatch = 1,
+  // Free dimensions are present in exactly one of the inputs, and also the
+  // output. These are non-contracted axes in the Tensor contraction.
+  kFree = 2,
+  // Contract dimensions are present in two inputs, but not the output. These
+  // dimensions are contracted in Tensor contraction.
+  kContract = 3,
+  // Reduce dimensions are present in exactly one input; and not in the output
+  // and are summed over prior to Tensor contraction.
+  kReduce = 4,
+};
+
+// Parses and validates an einsum equation in explicit form.
+absl::Status ValidateEinsumEquation(
+    const std::string& equation,
+    absl::InlinedVector<std::string, 2UL>* input_subscripts,
+    std::string* output_subscript);
+
+// Parses and validates the equation and the input shapes. Single character
+// labels are integerized and we populate input and output label subscripts
+// and corresponding counts. Also create the mapping from (named) labels to
+// their EinsumDimensionType.
+absl::Status ParseEinsumEquation(
+    const std::string& equation, OperandLabels* input_labels,
+    Labels* output_labels, std::vector<EinsumDimensionType>* label_types,
+    OperandLabelCounts* input_label_counts, LabelCounts* output_label_counts,
+    absl::InlinedVector<bool, 2UL>* input_has_ellipsis,
+    bool* output_has_ellipsis);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_UTIL_EINSUM_OP_UTIL_H_

tensorflow/core/util/env_var.h

@@ -0,0 +1,34 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_UTIL_ENV_VAR_H_
+#define TENSORFLOW_CORE_UTIL_ENV_VAR_H_
+
+#include "xla/tsl/util/env_var.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/stringpiece.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+
+using tsl::ReadBoolFromEnvVar;
+using tsl::ReadFloatFromEnvVar;
+using tsl::ReadInt64FromEnvVar;
+using tsl::ReadStringFromEnvVar;
+using tsl::ReadStringsFromEnvVar;
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_UTIL_ENV_VAR_H_

tensorflow/core/util/equal_graph_def.cc

@@ -0,0 +1,282 @@
+/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/util/equal_graph_def.h"
+
+#include <map>
+#include <set>
+#include <unordered_map>
+#include <unordered_set>
+
+#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/attr_value_util.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/framework/node_def.pb.h"
+#include "tensorflow/core/framework/node_def_util.h"
+#include "tensorflow/core/lib/hash/hash.h"
+#include "tensorflow/core/lib/strings/str_util.h"
+#include "tensorflow/core/lib/strings/strcat.h"
+#include "tensorflow/core/platform/protobuf.h"
+
+namespace tensorflow {
+
+bool EqualGraphDef(const GraphDef& actual, const GraphDef& expected,
+                   std::string* diff, const EqualGraphDefOptions& options) {
+  // Intentionally do not check that versions match so that this routine can
+  // be used for less brittle golden file tests.
+  return EqualRepeatedNodeDef(actual.node(), expected.node(), diff, options);
+}
+
+uint64_t GraphDefHash(const GraphDef& gdef,
+                      const EqualGraphDefOptions& options) {
+  return RepeatedNodeDefHash(gdef.node(), options);
+}
+
+bool EqualRepeatedNodeDef(const protobuf::RepeatedPtrField<NodeDef>& actual,
+                          const protobuf::RepeatedPtrField<NodeDef>& expected,
+                          std::string* diff,
+                          const EqualGraphDefOptions& options) {
+  std::unordered_map<std::string, const NodeDef*> actual_index;
+  for (const NodeDef& node : actual) {
+    actual_index[node.name()] = &node;
+  }
+
+  for (const NodeDef& expected_node : expected) {
+    auto actual_iter = actual_index.find(expected_node.name());
+    if (actual_iter == actual_index.end()) {
+      if (diff != nullptr) {
+        *diff = absl::StrCat("Did not find expected node '",
+                             SummarizeNodeDef(expected_node), "'");
+      }
+      return false;
+    }
+
+    if (!EqualNodeDef(*actual_iter->second, expected_node, diff, options)) {
+      return false;
+    }
+
+    actual_index.erase(actual_iter);
+  }
+
+  if (!actual_index.empty()) {
+    if (diff != nullptr) {
+      *diff =
+          absl::StrCat("Found unexpected node '",
+                       SummarizeNodeDef(*actual_index.begin()->second), "'");
+    }
+    return false;
+  }
+
+  return true;
+}
+
+uint64_t RepeatedNodeDefHash(const protobuf::RepeatedPtrField<NodeDef>& ndefs,
+                             const EqualGraphDefOptions& options) {
+  uint64_t h = 0xDECAFCAFFE;
+  // Insert NodeDefs into map to deterministically sort by name
+  std::map<std::string, const NodeDef*> nodes;
+  for (const NodeDef& node : ndefs) {
+    nodes[node.name()] = &node;
+  }
+  for (const auto& pair : nodes) {
+    h = Hash64(pair.first.data(), pair.first.size(), h);
+    h = Hash64Combine(NodeDefHash(*pair.second, options), h);
+  }
+  return h;
+}
+
+namespace {
+
+std::string JoinStringField(const protobuf::RepeatedPtrField<std::string>& f) {
+  std::string ret;
+  for (int i = 0; i < f.size(); ++i) {
+    if (i > 0) absl::StrAppend(&ret, ", ");
+    absl::StrAppend(&ret, f.Get(i));
+  }
+  return ret;
+}
+
+}  // namespace
+
+bool EqualNodeDef(const NodeDef& actual, const NodeDef& expected,
+                  std::string* diff, const EqualGraphDefOptions& options) {
+  if (actual.name() != expected.name()) {
+    if (diff != nullptr) {
+      *diff = strings::StrCat("Actual node name '", actual.name(),
+                              "' is not expected '", expected.name(), "'");
+    }
+    return false;
+  }
+
+  if (actual.op() != expected.op()) {
+    if (diff != nullptr) {
+      *diff = strings::StrCat("Node named '", actual.name(), "' has op '",
+                              actual.op(), "' that is not expected '",
+                              expected.op(), "'");
+    }
+    return false;
+  }
+
+  if (actual.device() != expected.device()) {
+    if (diff != nullptr) {
+      *diff = strings::StrCat("Node named '", actual.name(), "' has device '",
+                              actual.device(), "' that is not expected '",
+                              expected.device(), "'");
+    }
+    return false;
+  }
+
+  if (actual.input_size() != expected.input_size()) {
+    if (diff != nullptr) {
+      *diff = strings::StrCat("Node named '", actual.name(), "' has inputs '",
+                              JoinStringField(actual.input()),
+                              "' that don't match expected '",
+                              JoinStringField(expected.input()), "'");
+    }
+    return false;
+  }
+
+  int first_control_input = actual.input_size();
+  for (int i = 0; i < actual.input_size(); ++i) {
+    if (absl::StartsWith(actual.input(i), "^")) {
+      first_control_input = i;
+      break;
+    }
+    // Special case for inputs: "tensor" is equivalent to "tensor:0"
+    if (actual.input(i) != expected.input(i) &&
+        actual.input(i) != absl::StrCat(expected.input(i), ":0") &&
+        absl::StrCat(actual.input(i), ":0") != expected.input(i)) {
+      if (diff != nullptr) {
+        *diff = strings::StrCat("Node named '", actual.name(), "' has input ",
+                                i, " '", actual.input(i),
+                                "' that doesn't match expected '",
+                                expected.input(i), "'");
+      }
+      return false;
+    }
+  }
+
+  std::unordered_set<std::string> actual_control;
+  std::unordered_set<std::string> expected_control;
+  for (int i = first_control_input; i < actual.input_size(); ++i) {
+    actual_control.insert(actual.input(i));
+    expected_control.insert(expected.input(i));
+  }
+  for (const auto& e : expected_control) {
+    if (actual_control.erase(e) == 0) {
+      if (diff != nullptr) {
+        *diff = strings::StrCat("Node named '", actual.name(),
+                                "' missing expected control input '", e, "'");
+      }
+      return false;
+    }
+  }
+  if (!actual_control.empty()) {
+    if (diff != nullptr) {
+      *diff = strings::StrCat("Node named '", actual.name(),
+                              "' has unexpected control input '",
+                              *actual_control.begin(), "'");
+    }
+    return false;
+  }
+
+  std::unordered_set<std::string> actual_attr;
+  for (const auto& a : actual.attr()) {
+    if (options.ignore_internal_attrs && !a.first.empty() &&
+        a.first[0] == '_') {
+      continue;
+    }
+    actual_attr.insert(a.first);
+  }
+  for (const auto& e : expected.attr()) {
+    if (options.ignore_internal_attrs && !e.first.empty() &&
+        e.first[0] == '_') {
+      continue;
+    }
+
+    if (actual_attr.erase(e.first) == 0) {
+      if (diff != nullptr) {
+        *diff = strings::StrCat("Node named '", actual.name(),
+                                "' missing expected attr '", e.first,
+                                "' with value: ", SummarizeAttrValue(e.second));
+      }
+      return false;
+    }
+    auto iter = actual.attr().find(e.first);
+    if (!AreAttrValuesEqual(e.second, iter->second)) {
+      if (diff != nullptr) {
+        *diff = strings::StrCat(
+            "Node named '", actual.name(), "' has attr '", e.first,
+            "' with value: ", SummarizeAttrValue(iter->second),
+            " that does not match expected: ", SummarizeAttrValue(e.second));
+      }
+      return false;
+    }
+  }
+  if (!actual_attr.empty()) {
+    if (diff != nullptr) {
+      *diff = strings::StrCat(
+          "Node named '", actual.name(), "' has unexpected attr '",
+          *actual_attr.begin(), "' with value: ",
+          SummarizeAttrValue(actual.attr().find(*actual_attr.begin())->second));
+    }
+    return false;
+  }
+
+  return true;
+}
+
+uint64_t NodeDefHash(const NodeDef& ndef, const EqualGraphDefOptions& options) {
+  uint64_t h = Hash64(ndef.name());
+  h = Hash64(ndef.op().data(), ndef.op().size(), h);
+  h = Hash64(ndef.device().data(), ndef.device().size(), h);
+
+  // Normal inputs. Order important.
+  int first_control_input = ndef.input_size();
+  for (int i = 0; i < ndef.input_size(); ++i) {
+    if (absl::StartsWith(ndef.input(i), "^")) {
+      first_control_input = i;
+      break;
+    }
+    h = Hash64(ndef.input(i).data(), ndef.input(i).size(), h);
+  }
+
+  // Control inputs. Order irrelevant.
+  std::set<std::string> ndef_control;
+  for (int i = first_control_input; i < ndef.input_size(); ++i) {
+    ndef_control.insert(ndef.input(i));
+  }
+  for (const std::string& s : ndef_control) {
+    h = Hash64(s.data(), s.size(), h);
+  }
+
+  // Attributes
+  std::map<std::string, AttrValue> ndef_attr;
+  for (const auto& a : ndef.attr()) {
+    if (options.ignore_internal_attrs && !a.first.empty() &&
+        a.first[0] == '_') {
+      continue;
+    }
+    ndef_attr[a.first] = a.second;
+  }
+  for (const auto& a : ndef_attr) {
+    h = Hash64(a.first.data(), a.first.size(), h);
+    h = Hash64Combine(AttrValueHash(a.second), h);
+  }
+
+  return h;
+}
+
+}  // namespace tensorflow

tensorflow/core/util/equal_graph_def.h

@@ -0,0 +1,100 @@
+/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_UTIL_EQUAL_GRAPH_DEF_H_
+#define TENSORFLOW_CORE_UTIL_EQUAL_GRAPH_DEF_H_
+
+#include "tensorflow/core/framework/graph_def_util.h"
+#include "tensorflow/core/platform/protobuf.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+
+class GraphDef;
+class NodeDef;
+
+struct EqualGraphDefOptions {
+  // Should internal attributes (attribute names that start with '_') be
+  // ignored?
+  bool ignore_internal_attrs = true;
+};
+
+// Determines if actual and expected are equal, ignoring versions and ordering
+// of nodes, attrs, and control inputs.  If the GraphDefs are different and
+// diff != nullptr, *diff is set to an explanation of the difference.  Note that
+// we use node names to match up nodes between the graphs, and so the naming of
+// nodes must be consistent.
+bool EqualGraphDef(const GraphDef& actual, const GraphDef& expected,
+                   std::string* diff, const EqualGraphDefOptions& options = {});
+
+// Returns a hash of `gdef` that is consistent with EqualGraphDef. In other
+// words, if two graph defs compare equal according to EqualGraphDef,
+// GraphDefHash will return the same value for both of them when called
+// with the same `options` that was used in the call to EqualGraphDef.
+// Similarly to protobuf deterministic serialization, hash value is
+// guaranteed to be stable only for a given binary. In particular, one should
+// probably not persist the returned value.
+uint64_t GraphDefHash(const GraphDef& gdef,
+                      const EqualGraphDefOptions& options = {});
+
+// Determines if actual and expected are equal, ignoring: ordering of
+// attrs, internal attributes (if set in `options`), and control inputs.
+//
+// If the NodeDefs are different and
+// diff != nullptr, *diff is set to an explanation of the difference.
+bool EqualNodeDef(const NodeDef& actual, const NodeDef& expected,
+                  std::string* diff, const EqualGraphDefOptions& options = {});
+
+// Returns a hash of `ndef` that is consistent with EqualNodeDef. In other
+// words, if two node defs compare equal according to EqualNodeDef, NodeDefHash
+// will return the same value for both of them when called with the same
+// `options` that was used in the call to EqualNodeDef.
+// Similarly to protobuf deterministic serialization, hash value is
+// guaranteed to be stable only for a given binary. In particular, one should
+// probably not persist the returned value.
+uint64_t NodeDefHash(const NodeDef& ndef,
+                     const EqualGraphDefOptions& options = {});
+
+// Determines if actual and expected are equal, ignoring ordering. If they're
+// different and diff != nullptr, *diff is set to an explanation of the
+// difference.
+bool EqualRepeatedNodeDef(const protobuf::RepeatedPtrField<NodeDef>& actual,
+                          const protobuf::RepeatedPtrField<NodeDef>& expected,
+                          std::string* diff,
+                          const EqualGraphDefOptions& options = {});
+
+// Returns a hash of `ndefs` that is consistent with EqualRepeatedNodeDef.
+// In other words, if two ndefs compare equal according to
+// EqualRepeatedNodeDef, RepeatedNodeDefHash will return the same value for
+// both of them when called with the same `options` that was used in
+// the call to EqualRepeatedNodeDef.
+// Similarly to protobuf deterministic serialization, hash value is
+// guaranteed to be stable only for a given binary. In particular, one should
+// probably not persist the returned value.
+uint64_t RepeatedNodeDefHash(const protobuf::RepeatedPtrField<NodeDef>& ndefs,
+                             const EqualGraphDefOptions& options = {});
+
+#define TF_EXPECT_GRAPH_EQ(expected, actual)            \
+  do {                                                  \
+    string diff;                                        \
+    EXPECT_TRUE(EqualGraphDef(actual, expected, &diff)) \
+        << diff << "\nExpected:\n"                      \
+        << SummarizeGraphDef(expected) << "\nActual:\n" \
+        << SummarizeGraphDef(actual);                   \
+  } while (false)
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_UTIL_EQUAL_GRAPH_DEF_H_

tensorflow/core/util/equal_graph_def_test.cc

@@ -0,0 +1,323 @@
+/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <utility>
+
+#include "tensorflow/core/util/equal_graph_def.h"
+
+#include "tensorflow/core/framework/node_def_util.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/graph/graph_def_builder.h"
+#include "tensorflow/core/kernels/ops_util.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/public/version.h"
+
+namespace tensorflow {
+namespace {
+
+REGISTER_OP("Input").Output("o: float");
+REGISTER_OP("Alternate").Output("o: float");
+REGISTER_OP("Combine").Input("a: float").Input("b: float").Output("o: float");
+
+Node* Input(const GraphDefBuilder::Options& opts) {
+  return ops::SourceOp("Input", opts);
+}
+
+Node* Alternate(const GraphDefBuilder::Options& opts) {
+  return ops::SourceOp("Alternate", opts);
+}
+
+Node* Combine(ops::NodeOut a, ops::NodeOut b,
+              const GraphDefBuilder::Options& opts) {
+  return ops::BinaryOp("Combine", std::move(a), std::move(b), opts);
+}
+
+class EqualGraphDefTest : public ::testing::Test {
+ protected:
+  EqualGraphDefTest()
+      : e_(GraphDefBuilder::kFailImmediately),
+        a_(GraphDefBuilder::kFailImmediately) {}
+
+  bool Match() {
+    GraphDef expected;
+    TF_EXPECT_OK(e_.ToGraphDef(&expected));
+    GraphDef actual;
+    TF_EXPECT_OK(a_.ToGraphDef(&actual));
+    bool match = EqualGraphDef(actual, expected, &diff_);
+    if (match) {
+      EXPECT_EQ(GraphDefHash(expected), GraphDefHash(actual));
+    } else {
+      // While, strictly speaking, this does not have to be the case,
+      // we want to check that our hash is more than "return 0;".
+      // If, in the extremely unlikely case, some different graphs
+      // legitimately produce equal hash values in this test, we can always
+      // tweak them a little to produce different hash values.
+      EXPECT_NE(GraphDefHash(expected), GraphDefHash(actual));
+    }
+    return match;
+  }
+
+  GraphDefBuilder e_;
+  GraphDefBuilder a_;
+  std::string diff_;
+};
+
+TEST_F(EqualGraphDefTest, Match) {
+  Input(e_.opts().WithName("A"));
+  Input(a_.opts().WithName("A"));
+  EXPECT_TRUE(Match()) << diff_;
+}
+
+TEST_F(EqualGraphDefTest, NoMatch) {
+  Input(e_.opts().WithName("A"));
+  Input(a_.opts().WithName("B"));
+  EXPECT_FALSE(Match());
+  EXPECT_EQ("Did not find expected node '{{node A}} = Input[]()'", diff_);
+}
+
+TEST_F(EqualGraphDefTest, MissingNode) {
+  Input(e_.opts().WithName("A"));
+  Input(e_.opts().WithName("B"));
+  Input(a_.opts().WithName("A"));
+  EXPECT_FALSE(Match());
+  EXPECT_EQ("Did not find expected node '{{node B}} = Input[]()'", diff_);
+}
+
+TEST_F(EqualGraphDefTest, ExtraNode) {
+  Input(e_.opts().WithName("A"));
+  Input(a_.opts().WithName("A"));
+  Input(a_.opts().WithName("B"));
+  EXPECT_FALSE(Match());
+  EXPECT_EQ("Found unexpected node '{{node B}} = Input[]()'", diff_);
+}
+
+TEST_F(EqualGraphDefTest, NodeOrder) {
+  Node* a = Input(e_.opts().WithName("A"));
+  Node* b = Input(e_.opts().WithName("B"));
+  Combine(a, b, e_.opts().WithName("C"));
+
+  b = Input(a_.opts().WithName("B"));
+  a = Input(a_.opts().WithName("A"));
+  Combine(a, b, a_.opts().WithName("C"));
+  EXPECT_TRUE(Match()) << diff_;
+}
+
+TEST_F(EqualGraphDefTest, NameMismatch) {
+  Node* a = Input(e_.opts().WithName("A"));
+  Node* b = Input(e_.opts().WithName("B"));
+  // Have to call EqualNodeDef() directly here, since EqualGraphDef()
+  // only calls EqualNodeDef() with nodes that have matching names.
+  EXPECT_FALSE(EqualNodeDef(a->def(), b->def(), &diff_));
+  EXPECT_EQ("Actual node name 'A' is not expected 'B'", diff_);
+}
+
+TEST_F(EqualGraphDefTest, OpMismatch) {
+  Input(e_.opts().WithName("A"));
+  Alternate(a_.opts().WithName("A"));
+  EXPECT_FALSE(Match());
+  EXPECT_EQ("Node named 'A' has op 'Alternate' that is not expected 'Input'",
+            diff_);
+}
+
+TEST_F(EqualGraphDefTest, DeviceMatch) {
+  Input(e_.opts().WithName("A").WithDevice("/cpu:0"));
+  Input(a_.opts().WithName("A").WithDevice("/cpu:0"));
+  EXPECT_TRUE(Match()) << diff_;
+}
+
+TEST_F(EqualGraphDefTest, DeviceMismatch) {
+  Input(e_.opts().WithName("A").WithDevice("/cpu:0"));
+  Input(a_.opts().WithName("A").WithDevice("/cpu:1"));
+  EXPECT_FALSE(Match());
+  EXPECT_EQ("Node named 'A' has device '/cpu:1' that is not expected '/cpu:0'",
+            diff_);
+}
+
+TEST_F(EqualGraphDefTest, InputMismatch) {
+  Node* a = Input(e_.opts().WithName("A"));
+  Node* b = Input(e_.opts().WithName("B"));
+  Combine(a, a, e_.opts().WithName("C"));
+
+  a = Input(a_.opts().WithName("A"));
+  b = Input(a_.opts().WithName("B"));
+  Combine(b, b, a_.opts().WithName("C"));
+  EXPECT_FALSE(Match());
+  EXPECT_EQ("Node named 'C' has input 0 'B' that doesn't match expected 'A'",
+            diff_);
+}
+
+TEST_F(EqualGraphDefTest, InputOrderMismatch) {
+  Node* a = Input(e_.opts().WithName("A"));
+  Node* b = Input(e_.opts().WithName("B"));
+  Combine(a, b, e_.opts().WithName("C"));
+
+  a = Input(a_.opts().WithName("A"));
+  b = Input(a_.opts().WithName("B"));
+  Combine(b, a, a_.opts().WithName("C"));
+  EXPECT_FALSE(Match());
+  EXPECT_EQ("Node named 'C' has input 0 'B' that doesn't match expected 'A'",
+            diff_);
+}
+
+TEST_F(EqualGraphDefTest, ControlInputOrder) {
+  Node* a = Input(e_.opts().WithName("A"));
+  Node* b = Input(e_.opts().WithName("B"));
+  Node* c = Input(e_.opts().WithName("C"));
+  Node* d = Input(e_.opts().WithName("D"));
+  Combine(a, a,
+          e_.opts()
+              .WithName("E")
+              .WithControlInput(b)
+              .WithControlInput(c)
+              .WithControlInput(d));
+
+  a = Input(a_.opts().WithName("A"));
+  b = Input(a_.opts().WithName("B"));
+  c = Input(a_.opts().WithName("C"));
+  d = Input(a_.opts().WithName("D"));
+  Combine(a, a,
+          a_.opts()
+              .WithName("E")
+              .WithControlInput(c)
+              .WithControlInput(d)
+              .WithControlInput(b));
+  EXPECT_TRUE(Match()) << diff_;
+}
+
+TEST_F(EqualGraphDefTest, ControlInputMismatch) {
+  Node* a = Input(e_.opts().WithName("A"));
+  Node* b = Input(e_.opts().WithName("B"));
+  Node* c = Input(e_.opts().WithName("C"));
+  Node* d = Input(e_.opts().WithName("D"));
+  Combine(a, a,
+          e_.opts().WithName("E").WithControlInput(b).WithControlInput(c));
+
+  a = Input(a_.opts().WithName("A"));
+  b = Input(a_.opts().WithName("B"));
+  c = Input(a_.opts().WithName("C"));
+  d = Input(a_.opts().WithName("D"));
+  Combine(a, a,
+          a_.opts().WithName("E").WithControlInput(b).WithControlInput(d));
+  EXPECT_FALSE(Match());
+  EXPECT_EQ("Node named 'E' missing expected control input '^C'", diff_);
+}
+
+TEST_F(EqualGraphDefTest, ControlInputAdded) {
+  Node* a = Input(e_.opts().WithName("A"));
+  Node* b = Input(e_.opts().WithName("B"));
+  Node* c = Input(e_.opts().WithName("C"));
+  Combine(a, a, e_.opts().WithName("D").WithControlInput(b));
+
+  a = Input(a_.opts().WithName("A"));
+  b = Input(a_.opts().WithName("B"));
+  c = Input(a_.opts().WithName("C"));
+  Combine(a, a,
+          a_.opts().WithName("D").WithControlInput(b).WithControlInput(c));
+  EXPECT_FALSE(Match());
+  EXPECT_EQ(
+      "Node named 'D' has inputs 'A, A, ^B, ^C' that don't match "
+      "expected 'A, A, ^B'",
+      diff_);
+}
+
+TEST_F(EqualGraphDefTest, ControlInputRemoved) {
+  Node* a = Input(e_.opts().WithName("A"));
+  Node* b = Input(e_.opts().WithName("B"));
+  Node* c = Input(e_.opts().WithName("C"));
+  Combine(a, a,
+          e_.opts().WithName("D").WithControlInput(b).WithControlInput(c));
+
+  a = Input(a_.opts().WithName("A"));
+  b = Input(a_.opts().WithName("B"));
+  c = Input(a_.opts().WithName("C"));
+  Combine(a, a, a_.opts().WithName("D").WithControlInput(b));
+  EXPECT_FALSE(Match());
+  EXPECT_EQ(
+      "Node named 'D' has inputs 'A, A, ^B' that don't match "
+      "expected 'A, A, ^B, ^C'",
+      diff_);
+}
+
+TEST_F(EqualGraphDefTest, Attr) {
+  Node* a = Input(e_.opts().WithName("A"));
+  NodeDef same(a->def());
+  AddNodeAttr("foo", "bar", &same);
+  EXPECT_TRUE(EqualNodeDef(same, same, &diff_)) << diff_;
+}
+
+TEST_F(EqualGraphDefTest, AttrAdded) {
+  Node* a = Input(e_.opts().WithName("A"));
+  NodeDef actual(a->def());
+  AddNodeAttr("foo", "bar", &actual);
+  EXPECT_FALSE(EqualNodeDef(actual, a->def(), &diff_));
+  EXPECT_EQ("Node named 'A' has unexpected attr 'foo' with value: \"bar\"",
+            diff_);
+}
+
+TEST_F(EqualGraphDefTest, AttrRemoved) {
+  Node* a = Input(e_.opts().WithName("A"));
+  NodeDef expected(a->def());
+  AddNodeAttr("foo", "bar", &expected);
+  EXPECT_FALSE(EqualNodeDef(a->def(), expected, &diff_));
+  EXPECT_EQ("Node named 'A' missing expected attr 'foo' with value: \"bar\"",
+            diff_);
+}
+
+TEST_F(EqualGraphDefTest, AttrOrder) {
+  Node* a = Input(e_.opts().WithName("A"));
+  NodeDef actual(a->def());
+  AddNodeAttr("foo", "bar", &actual);
+  AddNodeAttr("baz", 42, &actual);
+
+  NodeDef expected(a->def());
+  AddNodeAttr("baz", 42, &expected);
+  AddNodeAttr("foo", "bar", &expected);
+
+  EXPECT_TRUE(EqualNodeDef(actual, expected, &diff_)) << diff_;
+}
+
+TEST_F(EqualGraphDefTest, AttrMismatch) {
+  Node* a = Input(e_.opts().WithName("A"));
+  NodeDef actual(a->def());
+  AddNodeAttr("foo", "bar", &actual);
+  AddNodeAttr("baz", 5, &actual);
+
+  NodeDef expected(a->def());
+  AddNodeAttr("baz", 42, &expected);
+  AddNodeAttr("foo", "bar", &expected);
+
+  EXPECT_FALSE(EqualNodeDef(actual, expected, &diff_));
+  EXPECT_EQ(
+      "Node named 'A' has attr 'baz' with value: 5 that does not match "
+      "expected: 42",
+      diff_);
+}
+
+TEST_F(EqualGraphDefTest, IgnoreInternalAttrs) {
+  Node* a = Input(e_.opts().WithName("A"));
+  NodeDef actual(a->def());
+  AddNodeAttr("foo", "bar", &actual);
+  // Internal attrs are ignored.
+  AddNodeAttr("_class", 5, &actual);
+
+  NodeDef expected(a->def());
+  AddNodeAttr("foo", "bar", &expected);
+  AddNodeAttr("_kernel", "eigen", &actual);
+  EXPECT_TRUE(EqualNodeDef(actual, expected, &diff_));
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/core/util/event.proto

@@ -0,0 +1,158 @@
+// Copyright 2026 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;
+
+import "tensorflow/core/framework/summary.proto";
+
+option cc_enable_arenas = true;
+option java_outer_classname = "EventProtos";
+option java_multiple_files = true;
+option java_package = "org.tensorflow.util";
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/util/event_go_proto";
+
+// Protocol buffer representing an event that happened during
+// the execution of a Brain model.
+message Event {
+  // Timestamp of the event.
+  double wall_time = 1;
+
+  // Global step of the event.
+  int64 step = 2;
+
+  oneof what {
+    // An event file was started, with the specified version.
+    // This is use to identify the contents of the record IO files
+    // easily.  Current version is "brain.Event:2".  All versions
+    // start with "brain.Event:".
+    string file_version = 3;
+    // An encoded version of a GraphDef.
+    bytes graph_def = 4;
+    // A summary was generated.
+    Summary summary = 5;
+    // The user output a log message. This was theoretically used by the defunct
+    // tensorboard_logging module, which has since been removed; this field is
+    // now deprecated and should not be used.
+    LogMessage log_message = 6 [deprecated = true];
+    // The state of the session which can be used for restarting after crashes.
+    SessionLog session_log = 7;
+    // The metadata returned by running a session.run() call.
+    TaggedRunMetadata tagged_run_metadata = 8;
+    // An encoded version of a MetaGraphDef.
+    bytes meta_graph_def = 9;
+  }
+
+  // Information of the source that writes the events, this is only logged in
+  // the very first event along with the `file_version` field.
+  SourceMetadata source_metadata = 10;
+}
+
+// Holds the information of the source that writes the events.
+message SourceMetadata {
+  // Low level name of the summary writer, such as
+  // `tensorflow.core.util.events_writer`.
+  string writer = 1;
+}
+
+// Protocol buffer used for logging messages to the events file.
+//
+// This was theoretically used by the defunct tensorboard_logging module, which
+// has been removed; this message is now deprecated and should not be used.
+message LogMessage {
+  option deprecated = true;
+
+  enum Level {
+    option deprecated = true;
+
+    UNKNOWN = 0;
+    // Note: The logging level 10 cannot be named DEBUG. Some software
+    // projects compile their C/C++ code with -DDEBUG in debug builds. So the
+    // C++ code generated from this file should not have an identifier named
+    // DEBUG.
+    DEBUGGING = 10;
+    INFO = 20;
+    WARN = 30;
+    ERROR = 40;
+    FATAL = 50;
+  }
+  Level level = 1;
+  string message = 2;
+}
+
+// Protocol buffer used for logging session state.
+message SessionLog {
+  enum SessionStatus {
+    STATUS_UNSPECIFIED = 0;
+    START = 1;
+    STOP = 2;
+    CHECKPOINT = 3;
+  }
+
+  SessionStatus status = 1;
+  // This checkpoint_path contains both the path and filename.
+  string checkpoint_path = 2;
+  string msg = 3;
+}
+
+// For logging the metadata output for a single session.run() call.
+message TaggedRunMetadata {
+  // Tag name associated with this metadata.
+  string tag = 1;
+  // Byte-encoded version of the `RunMetadata` proto in order to allow lazy
+  // deserialization.
+  bytes run_metadata = 2;
+}
+
+// Worker heartbeat messages.  Support for these operations is currently
+// internal and expected to change.
+
+// Current health status of a worker.
+enum WorkerHealth {
+  OK = 0;  // By default a worker is healthy.
+  RECEIVED_SHUTDOWN_SIGNAL = 1;
+  INTERNAL_ERROR = 2;
+  SHUTTING_DOWN = 3;  // Worker has been instructed to shutdown after a timeout.
+}
+
+// Indicates the behavior of the worker when an internal error or shutdown
+// signal is received.
+enum WorkerShutdownMode {
+  DEFAULT = 0;
+  NOT_CONFIGURED = 1;
+  WAIT_FOR_COORDINATOR = 2;
+  SHUTDOWN_AFTER_TIMEOUT = 3;
+}
+
+message WatchdogConfig {
+  int64 timeout_ms = 1;
+}
+
+message RequestedExitCode {
+  int32 exit_code = 1;
+}
+
+message WorkerHeartbeatRequest {
+  WorkerShutdownMode shutdown_mode = 1;
+  WatchdogConfig watchdog_config = 2;
+  RequestedExitCode exit_code = 3;
+}
+
+message WorkerHeartbeatResponse {
+  WorkerHealth health_status = 1;
+  repeated Event worker_log = 2;
+  string hostname = 3;
+}

tensorflow/core/util/events_writer.cc

@@ -0,0 +1,167 @@
+/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/util/events_writer.h"
+
+#include <stddef.h>  // for NULL
+
+#include <memory>
+
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/lib/io/path.h"
+#include "tensorflow/core/lib/strings/strcat.h"
+#include "tensorflow/core/lib/strings/stringprintf.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/host_info.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/util/event.pb.h"
+
+namespace tensorflow {
+
+EventsWriter::EventsWriter(const std::string& file_prefix)
+    // TODO(jeff,sanjay): Pass in env and use that here instead of Env::Default
+    : env_(Env::Default()),
+      file_prefix_(file_prefix),
+      num_outstanding_events_(0) {}
+
+EventsWriter::~EventsWriter() {
+  Close().IgnoreError();  // Autoclose in destructor.
+}
+
+absl::Status EventsWriter::Init() { return InitWithSuffix(""); }
+
+absl::Status EventsWriter::InitWithSuffix(const std::string& suffix) {
+  file_suffix_ = suffix;
+  return InitIfNeeded();
+}
+
+absl::Status EventsWriter::InitIfNeeded() {
+  if (recordio_writer_ != nullptr) {
+    CHECK(!filename_.empty());
+    if (!FileStillExists().ok()) {
+      // Warn user of data loss and let .reset() below do basic cleanup.
+      if (num_outstanding_events_ > 0) {
+        LOG(WARNING) << "Re-initialization, attempting to open a new file, "
+                     << num_outstanding_events_ << " events will be lost.";
+      }
+    } else {
+      // No-op: File is present and writer is initialized.
+      return absl::OkStatus();
+    }
+  }
+
+  int64_t time_in_seconds = env_->NowMicros() / 1000000;
+
+  filename_ =
+      absl::StrFormat("%s.out.tfevents.%010lld.%s%s", file_prefix_.c_str(),
+                      static_cast<long long>(time_in_seconds),
+                      port::Hostname().c_str(), file_suffix_.c_str());
+
+  // Reset recordio_writer (which has a reference to recordio_file_) so final
+  // Flush() and Close() call have access to recordio_file_.
+  recordio_writer_.reset();
+
+  TF_RETURN_WITH_CONTEXT_IF_ERROR(
+      env_->NewWritableFile(filename_, &recordio_file_),
+      "Creating writable file ", filename_);
+  recordio_writer_ = std::make_unique<io::RecordWriter>(recordio_file_.get());
+  if (recordio_writer_ == nullptr) {
+    return absl::UnknownError("Could not create record writer");
+  }
+  num_outstanding_events_ = 0;
+  VLOG(1) << "Successfully opened events file: " << filename_;
+  {
+    // Write the first event with the current version, and flush
+    // right away so the file contents will be easily determined.
+
+    Event event;
+    event.set_wall_time(time_in_seconds);
+    event.set_file_version(absl::StrCat(kVersionPrefix, kCurrentVersion));
+    SourceMetadata* source_metadata = event.mutable_source_metadata();
+    source_metadata->set_writer(kWriterSourceMetadata);
+    WriteEvent(event);
+    TF_RETURN_WITH_CONTEXT_IF_ERROR(Flush(), "Flushing first event.");
+  }
+  return absl::OkStatus();
+}
+
+std::string EventsWriter::FileName() {
+  if (filename_.empty()) {
+    InitIfNeeded().IgnoreError();
+  }
+  return filename_;
+}
+
+void EventsWriter::WriteSerializedEvent(absl::string_view event_str) {
+  if (recordio_writer_ == nullptr) {
+    if (!InitIfNeeded().ok()) {
+      LOG(ERROR) << "Write failed because file could not be opened.";
+      return;
+    }
+  }
+  num_outstanding_events_++;
+  recordio_writer_->WriteRecord(event_str).IgnoreError();
+}
+
+// NOTE(touts); This is NOT the function called by the Python code.
+// Python calls WriteSerializedEvent(), see events_writer.i.
+void EventsWriter::WriteEvent(const Event& event) {
+  std::string record;
+  event.AppendToString(&record);
+  WriteSerializedEvent(record);
+}
+
+absl::Status EventsWriter::Flush() {
+  if (num_outstanding_events_ == 0) return absl::OkStatus();
+  CHECK(recordio_file_ != nullptr) << "Unexpected NULL file";
+
+  TF_RETURN_WITH_CONTEXT_IF_ERROR(recordio_writer_->Flush(), "Failed to flush ",
+                                  num_outstanding_events_, " events to ",
+                                  filename_);
+  TF_RETURN_WITH_CONTEXT_IF_ERROR(recordio_file_->Sync(), "Failed to sync ",
+                                  num_outstanding_events_, " events to ",
+                                  filename_);
+  VLOG(1) << "Wrote " << num_outstanding_events_ << " events to disk.";
+  num_outstanding_events_ = 0;
+  return absl::OkStatus();
+}
+
+absl::Status EventsWriter::Close() {
+  absl::Status status = Flush();
+  if (recordio_file_ != nullptr) {
+    absl::Status close_status = recordio_file_->Close();
+    if (!close_status.ok()) {
+      status = close_status;
+    }
+    recordio_writer_.reset(nullptr);
+    recordio_file_.reset(nullptr);
+  }
+  num_outstanding_events_ = 0;
+  return status;
+}
+
+absl::Status EventsWriter::FileStillExists() {
+  if (env_->FileExists(filename_).ok()) {
+    return absl::OkStatus();
+  }
+  // This can happen even with non-null recordio_writer_ if some other
+  // process has removed the file.
+  return absl::UnknownError(
+      absl::StrCat("The events file ", filename_, " has disappeared."));
+}
+
+}  // namespace tensorflow

tensorflow/core/util/events_writer.h

@@ -0,0 +1,103 @@
+/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_UTIL_EVENTS_WRITER_H_
+#define TENSORFLOW_CORE_UTIL_EVENTS_WRITER_H_
+
+#include <memory>
+#include <string>
+
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/lib/io/record_writer.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/macros.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/util/event.pb.h"
+
+namespace tensorflow {
+
+class EventsWriter {
+ public:
+#ifndef SWIG
+  // Prefix of version string present in the first entry of every event file.
+  static constexpr const char* kVersionPrefix = "brain.Event:";
+  static constexpr const int kCurrentVersion = 2;
+  static constexpr const char* kWriterSourceMetadata =
+      "tensorflow.core.util.events_writer";
+#endif
+
+  // Events files typically have a name of the form
+  //   '/some/file/path/my.file.out.events.[timestamp].[hostname][suffix]'
+  // To create and EventWriter, the user should provide file_prefix =
+  //   '/some/file/path/my.file'
+  // The EventsWriter will append '.out.events.[timestamp].[hostname][suffix]'
+  // to the ultimate filename once Init() is called.
+  // Note that it is not recommended to simultaneously have two
+  // EventWriters writing to the same file_prefix.
+  explicit EventsWriter(const std::string& file_prefix);
+  ~EventsWriter();
+
+  // Sets the event file filename and opens file for writing.  If not called by
+  // user, will be invoked automatically by a call to FileName() or Write*().
+  // Returns false if the file could not be opened.  Idempotent: if file exists
+  // and is open this is a no-op.  If on the other hand the file was opened,
+  // but has since disappeared (e.g. deleted by another process), this will open
+  // a new file with a new timestamp in its filename.
+  absl::Status Init();
+  absl::Status InitWithSuffix(const std::string& suffix);
+
+  // Returns the filename for the current events file:
+  // filename_ = [file_prefix_].out.events.[timestamp].[hostname][suffix]
+  std::string FileName();
+
+  // Append "event" to the file.  The "tensorflow::" part is for swig happiness.
+  void WriteEvent(const tensorflow::Event& event);
+
+  // Append "event_str", a serialized Event, to the file.
+  // Note that this function does NOT check that de-serializing event_str
+  // results in a valid Event proto.  The tensorflow:: bit makes SWIG happy.
+  void WriteSerializedEvent(absl::string_view event_str);
+
+  // EventWriter automatically flushes and closes on destruction, but
+  // these two methods are provided for users who want to write to disk sooner
+  // and/or check for success.
+  //   Flush() pushes outstanding events to disk.  Returns false if the
+  // events file could not be created, or if the file exists but could not
+  // be written too.
+  //   Close() calls Flush() and then closes the current events file.
+  // Returns true only if both the flush and the closure were successful.
+  absl::Status Flush();
+  absl::Status Close();
+
+ private:
+  absl::Status FileStillExists();  // OK if event_file_path_ exists.
+  absl::Status InitIfNeeded();
+
+  Env* env_;
+  const std::string file_prefix_;
+  std::string file_suffix_;
+  std::string filename_;
+  std::unique_ptr<WritableFile> recordio_file_;
+  std::unique_ptr<io::RecordWriter> recordio_writer_;
+  int num_outstanding_events_;
+#ifndef SWIG
+  EventsWriter(const EventsWriter&) = delete;
+  void operator=(const EventsWriter&) = delete;
+#endif
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_UTIL_EVENTS_WRITER_H_

tensorflow/core/util/events_writer_test.cc

@@ -0,0 +1,212 @@
+/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/util/events_writer.h"
+
+#include <math.h>
+
+#include <memory>
+
+#include "tensorflow/core/framework/summary.pb.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/lib/io/path.h"
+#include "tensorflow/core/lib/io/record_reader.h"
+#include "tensorflow/core/lib/strings/strcat.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/protobuf.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/util/event.pb.h"
+
+namespace tensorflow {
+namespace {
+
+// shorthand
+Env* env() { return Env::Default(); }
+
+void WriteSimpleValue(EventsWriter* writer, double wall_time, int64_t step,
+                      const std::string& tag, float simple_value) {
+  Event event;
+  event.set_wall_time(wall_time);
+  event.set_step(step);
+  Summary::Value* summ_val = event.mutable_summary()->add_value();
+  summ_val->set_tag(tag);
+  summ_val->set_simple_value(simple_value);
+  writer->WriteEvent(event);
+}
+
+void WriteFile(EventsWriter* writer) {
+  WriteSimpleValue(writer, 1234, 34, "foo", 3.14159);
+  WriteSimpleValue(writer, 2345, 35, "bar", -42);
+}
+
+static bool ReadEventProto(io::RecordReader* reader, uint64_t* offset,
+                           Event* proto) {
+  tstring record;
+  absl::Status s = reader->ReadRecord(offset, &record);
+  if (!s.ok()) {
+    return false;
+  }
+  return ParseProtoUnlimited(proto, record);
+}
+
+void VerifyFile(const std::string& filename) {
+  TF_CHECK_OK(env()->FileExists(filename));
+  std::unique_ptr<RandomAccessFile> event_file;
+  TF_CHECK_OK(env()->NewRandomAccessFile(filename, &event_file));
+  io::RecordReader* reader = new io::RecordReader(event_file.get());
+
+  uint64_t offset = 0;
+
+  Event actual;
+  CHECK(ReadEventProto(reader, &offset, &actual));
+  VLOG(1) << actual.ShortDebugString();
+  // Wall time should be within 5s of now.
+
+  double current_time = env()->NowMicros() / 1000000.0;
+  EXPECT_LT(fabs(actual.wall_time() - current_time), 5);
+  // Should have the current version number.
+  EXPECT_EQ(actual.file_version(), absl::StrCat(EventsWriter::kVersionPrefix,
+                                                EventsWriter::kCurrentVersion));
+  // Should have the current source metadata.
+  EXPECT_EQ(actual.source_metadata().writer(),
+            EventsWriter::kWriterSourceMetadata);
+
+  Event expected;
+  CHECK(ReadEventProto(reader, &offset, &actual));
+  VLOG(1) << actual.ShortDebugString();
+  ASSERT_TRUE(protobuf::TextFormat::ParseFromString(
+      "wall_time: 1234 step: 34 "
+      "summary { value { tag: 'foo' simple_value: 3.14159 } }",
+      &expected));
+  // TODO(keveman): Enable this check
+  // EXPECT_THAT(expected, EqualsProto(actual));
+
+  CHECK(ReadEventProto(reader, &offset, &actual));
+  VLOG(1) << actual.ShortDebugString();
+  ASSERT_TRUE(protobuf::TextFormat::ParseFromString(
+      "wall_time: 2345 step: 35 "
+      "summary { value { tag: 'bar' simple_value: -42 } }",
+      &expected));
+  // TODO(keveman): Enable this check
+  // EXPECT_THAT(expected, EqualsProto(actual));
+
+  TF_CHECK_OK(env()->DeleteFile(filename));
+  delete reader;
+}
+
+std::string GetDirName(const std::string& suffix) {
+  return io::JoinPath(testing::TmpDir(), suffix);
+}
+
+TEST(EventWriter, WriteFlush) {
+  std::string file_prefix = GetDirName("/writeflush_test");
+  EventsWriter writer(file_prefix);
+  WriteFile(&writer);
+  TF_EXPECT_OK(writer.Flush());
+  std::string filename = writer.FileName();
+  VerifyFile(filename);
+}
+
+TEST(EventWriter, WriteClose) {
+  std::string file_prefix = GetDirName("/writeclose_test");
+  EventsWriter writer(file_prefix);
+  WriteFile(&writer);
+  TF_EXPECT_OK(writer.Close());
+  std::string filename = writer.FileName();
+  VerifyFile(filename);
+}
+
+TEST(EventWriter, WriteDelete) {
+  std::string file_prefix = GetDirName("/writedelete_test");
+  EventsWriter* writer = new EventsWriter(file_prefix);
+  WriteFile(writer);
+  std::string filename = writer->FileName();
+  delete writer;
+  VerifyFile(filename);
+}
+
+TEST(EventWriter, FailFlush) {
+  std::string file_prefix = GetDirName("/failflush_test");
+  EventsWriter writer(file_prefix);
+  std::string filename = writer.FileName();
+  WriteFile(&writer);
+  TF_EXPECT_OK(env()->FileExists(filename));
+  TF_ASSERT_OK(env()->DeleteFile(filename));
+  EXPECT_TRUE(writer.Flush().ok());
+}
+
+TEST(EventWriter, FailClose) {
+  std::string file_prefix = GetDirName("/failclose_test");
+  EventsWriter writer(file_prefix);
+  std::string filename = writer.FileName();
+  WriteFile(&writer);
+  TF_EXPECT_OK(env()->FileExists(filename));
+  TF_ASSERT_OK(env()->DeleteFile(filename));
+  EXPECT_TRUE(writer.Close().ok());
+}
+
+TEST(EventWriter, InitWriteClose) {
+  std::string file_prefix = GetDirName("/initwriteclose_test");
+  EventsWriter writer(file_prefix);
+  TF_EXPECT_OK(writer.Init());
+  std::string filename0 = writer.FileName();
+  TF_EXPECT_OK(env()->FileExists(filename0));
+  WriteFile(&writer);
+  TF_EXPECT_OK(writer.Close());
+  std::string filename1 = writer.FileName();
+  EXPECT_EQ(filename0, filename1);
+  VerifyFile(filename1);
+}
+
+TEST(EventWriter, NameWriteClose) {
+  std::string file_prefix = GetDirName("/namewriteclose_test");
+  EventsWriter writer(file_prefix);
+  std::string filename = writer.FileName();
+  TF_EXPECT_OK(env()->FileExists(filename));
+  WriteFile(&writer);
+  TF_EXPECT_OK(writer.Close());
+  VerifyFile(filename);
+}
+
+TEST(EventWriter, NameClose) {
+  std::string file_prefix = GetDirName("/nameclose_test");
+  EventsWriter writer(file_prefix);
+  std::string filename = writer.FileName();
+  TF_EXPECT_OK(writer.Close());
+  TF_EXPECT_OK(env()->FileExists(filename));
+  TF_ASSERT_OK(env()->DeleteFile(filename));
+}
+
+TEST(EventWriter, FileDeletionBeforeWriting) {
+  std::string file_prefix = GetDirName("/fdbw_test");
+  EventsWriter writer(file_prefix);
+  std::string filename0 = writer.FileName();
+  TF_EXPECT_OK(env()->FileExists(filename0));
+  env()->SleepForMicroseconds(
+      2000000);  // To make sure timestamp part of filename will differ.
+  TF_ASSERT_OK(env()->DeleteFile(filename0));
+  TF_EXPECT_OK(writer.Init());  // Init should reopen file.
+  WriteFile(&writer);
+  TF_EXPECT_OK(writer.Flush());
+  std::string filename1 = writer.FileName();
+  EXPECT_NE(filename0, filename1);
+  VerifyFile(filename1);
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/core/util/example_proto_fast_parsing.h

@@ -0,0 +1,172 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_UTIL_EXAMPLE_PROTO_FAST_PARSING_H_
+#define TENSORFLOW_CORE_UTIL_EXAMPLE_PROTO_FAST_PARSING_H_
+
+#include <string>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include "tensorflow/core/example/example.pb.h"
+#include "tensorflow/core/framework/allocator.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/partial_tensor_shape.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/lib/gtl/array_slice.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/util/sparse/sparse_tensor.h"
+
+namespace tensorflow {
+namespace example {
+
+// FastParseExampleConfig defines how to parse features in Example.
+// Each sub-config is responsible for one feature identified with feature_name.
+// FastParseExampleConfig can't have two sub-configs with the same feature_name.
+// dtype identifies the type of output vector and the kind of Feature expected
+// in Example.
+struct FastParseExampleConfig {
+  struct Dense {
+    Dense(absl::string_view feature_name, DataType dtype,
+          PartialTensorShape shape, Tensor default_value, bool variable_length,
+          std::size_t elements_per_stride)
+        : feature_name(feature_name),  // TODO(mrry): Switch to preallocated
+                                       // tstring when this is available.
+          dtype(dtype),
+          shape(std::move(shape)),
+          default_value(std::move(default_value)),
+          variable_length(variable_length),
+          elements_per_stride(elements_per_stride) {}
+    Dense() = default;
+
+    tstring feature_name;
+    DataType dtype;
+    // These 2 fields correspond exactly to dense_shapes and dense_defaults in
+    // ParseExample op.
+    // Documentation is available in: tensorflow/core/ops/parsing_ops.cc
+    PartialTensorShape shape;
+    Tensor default_value;
+    bool variable_length;
+    std::size_t elements_per_stride;
+  };
+
+  struct Sparse {
+    Sparse(absl::string_view feature_name, DataType dtype)
+        : feature_name(feature_name),  // TODO(mrry): Switch to preallocated
+                                       // tstring when this is available.
+          dtype(dtype) {}
+    Sparse() = default;
+
+    tstring feature_name;
+    DataType dtype;
+  };
+
+  struct Ragged {
+    Ragged(absl::string_view feature_name, DataType dtype,
+           DataType splits_dtype)
+        : feature_name(feature_name),  // TODO(mrry): Switch to preallocated
+                                       // tstring when this is available.
+          dtype(dtype),
+          splits_dtype(splits_dtype) {}
+    Ragged() = default;
+
+    tstring feature_name;
+    DataType dtype;
+    DataType splits_dtype;
+  };
+
+  std::vector<Dense> dense;
+  std::vector<Sparse> sparse;
+  std::vector<Ragged> ragged;
+
+  // If `true`, `Result::feature_stats` will contain one
+  // `PerExampleFeatureStats` for each serialized example in the input.
+  bool collect_feature_stats = false;
+};
+
+// Statistics about the features in each example passed to
+// `FastParse[Single]Example()`.
+//
+// TODO(b/111553342): The gathered statistics currently have two limitations:
+// * Feature names that appear more than once will be counted multiple times.
+// * The feature values count only represents the counts for features that were
+//   requested in the `FastParseExampleConfig`.
+// These could be addressed with additional work at runtime.
+struct PerExampleFeatureStats {
+  // The number of feature names in an example.
+  size_t features_count = 0;
+
+  // The sum of the number of values in each feature that is parsed.
+  size_t feature_values_count = 0;
+};
+
+// This is exactly the output of TF's ParseExample Op.
+// Documentation is available in: tensorflow/core/ops/parsing_ops.cc
+struct Result {
+  std::vector<Tensor> sparse_indices;
+  std::vector<Tensor> sparse_values;
+  std::vector<Tensor> sparse_shapes;
+  std::vector<Tensor> dense_values;
+  std::vector<Tensor> ragged_values;
+  std::vector<Tensor> ragged_splits;
+  std::vector<Tensor> ragged_outer_splits;  // For SequenceExamples
+
+  // This vector will be populated with one element per example if
+  // `FastParseExampleConfig::collect_feature_stats` is set to `true`.
+  std::vector<PerExampleFeatureStats> feature_stats;
+};
+
+// Parses a batch of serialized Example protos and converts them into result
+// according to given config.
+// Given example names have to either be empty or the same size as serialized.
+// example_names are used only for error messages.
+absl::Status FastParseExample(const FastParseExampleConfig& config,
+                              absl::Span<const tstring> serialized,
+                              absl::Span<const tstring> example_names,
+                              thread::ThreadPool* thread_pool, Result* result);
+
+// TODO(mrry): Move the hash table construction into the config object.
+typedef FastParseExampleConfig FastParseSingleExampleConfig;
+
+absl::Status FastParseSingleExample(const FastParseSingleExampleConfig& config,
+                                    absl::string_view serialized,
+                                    Result* result);
+
+// Parses a batch of serialized SequenceExample protos and converts them into
+// result according to given config.
+// Given example names have to either be empty or the same size as serialized.
+// example_names are used only for error messages.
+// (If batch=true, then this parses a single SequenceExample.)
+absl::Status FastParseSequenceExample(
+    const example::FastParseExampleConfig& context_config,
+    const example::FastParseExampleConfig& sequence_config,
+    absl::Span<const tstring> serialized,
+    absl::Span<const tstring> example_names, thread::ThreadPool* thread_pool,
+    example::Result* context_result, example::Result* sequence_result,
+    std::vector<Tensor>* dense_feature_lengths, bool is_batch = true);
+
+// This function parses serialized Example and populates given example.
+// It uses the same specialized parser as FastParseExample which is efficient.
+// But then constructs Example which is relatively slow.
+// It is exported here as a convenient API to test parser part separately.
+bool TestFastParse(const std::string& serialized, Example* example);
+
+}  // namespace example
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_UTIL_EXAMPLE_PROTO_FAST_PARSING_H_

tensorflow/core/util/example_proto_fast_parsing_test.cc

@@ -0,0 +1,624 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/util/example_proto_fast_parsing.h"
+
+#include <cstdint>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "tensorflow/core/example/example.pb.h"
+#include "tensorflow/core/example/feature.pb.h"
+#include "tensorflow/core/lib/random/philox_random.h"
+#include "tensorflow/core/lib/random/simple_philox.h"
+#include "tensorflow/core/platform/protobuf.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/platform/test_benchmark.h"
+#include "tensorflow/core/util/example_proto_fast_parsing_test.pb.h"
+
+namespace tensorflow {
+namespace example {
+namespace {
+
+constexpr char kDenseInt64Key[] = "dense_int64";
+constexpr char kDenseFloatKey[] = "dense_float";
+constexpr char kDenseStringKey[] = "dense_string";
+
+constexpr char kSparseInt64Key[] = "sparse_int64";
+constexpr char kSparseFloatKey[] = "sparse_float";
+constexpr char kSparseStringKey[] = "sparse_string";
+
+std::string SerializedToReadable(std::string serialized) {
+  std::string result;
+  result += '"';
+  for (char c : serialized)
+    absl::StrAppend(&result, "\\x", absl::Hex(c, absl::kZeroPad2));
+  result += '"';
+  return result;
+}
+
+template <class T>
+std::string Serialize(const T& example) {
+  std::string serialized;
+  example.SerializeToString(&serialized);
+  return serialized;
+}
+
+// Tests that serialized gets parsed identically by TestFastParse(..)
+// and the regular Example.ParseFromString(..).
+void TestCorrectness(const std::string& serialized) {
+  Example example;
+  Example fast_example;
+  EXPECT_TRUE(example.ParseFromString(serialized));
+  example.DiscardUnknownFields();
+  EXPECT_TRUE(TestFastParse(serialized, &fast_example));
+  EXPECT_EQ(example.DebugString(), fast_example.DebugString());
+  if (example.DebugString() != fast_example.DebugString()) {
+    LOG(ERROR) << "Bad serialized: " << SerializedToReadable(serialized);
+  }
+}
+
+// Fast parsing does not differentiate between EmptyExample and EmptyFeatures
+// TEST(FastParse, EmptyExample) {
+//   Example example;
+//   TestCorrectness(example);
+// }
+
+TEST(FastParse, IgnoresPrecedingUnknownTopLevelFields) {
+  ExampleWithExtras example;
+  (*example.mutable_features()->mutable_feature())["age"]
+      .mutable_int64_list()
+      ->add_value(13);
+  example.set_extra1("some_str");
+  example.set_extra2(123);
+  example.set_extra3(234);
+  example.set_extra4(345);
+  example.set_extra5(4.56);
+  example.add_extra6(5.67);
+  example.add_extra6(6.78);
+  (*example.mutable_extra7()->mutable_feature())["extra7"]
+      .mutable_int64_list()
+      ->add_value(1337);
+
+  Example context;
+  (*context.mutable_features()->mutable_feature())["zipcode"]
+      .mutable_int64_list()
+      ->add_value(94043);
+
+  TestCorrectness(absl::StrCat(Serialize(example), Serialize(context)));
+}
+
+TEST(FastParse, IgnoresTrailingUnknownTopLevelFields) {
+  Example example;
+  (*example.mutable_features()->mutable_feature())["age"]
+      .mutable_int64_list()
+      ->add_value(13);
+
+  ExampleWithExtras context;
+  (*context.mutable_features()->mutable_feature())["zipcode"]
+      .mutable_int64_list()
+      ->add_value(94043);
+  context.set_extra1("some_str");
+  context.set_extra2(123);
+  context.set_extra3(234);
+  context.set_extra4(345);
+  context.set_extra5(4.56);
+  context.add_extra6(5.67);
+  context.add_extra6(6.78);
+  (*context.mutable_extra7()->mutable_feature())["extra7"]
+      .mutable_int64_list()
+      ->add_value(1337);
+
+  TestCorrectness(absl::StrCat(Serialize(example), Serialize(context)));
+}
+
+TEST(FastParse, SingleInt64WithContext) {
+  Example example;
+  (*example.mutable_features()->mutable_feature())["age"]
+      .mutable_int64_list()
+      ->add_value(13);
+
+  Example context;
+  (*context.mutable_features()->mutable_feature())["zipcode"]
+      .mutable_int64_list()
+      ->add_value(94043);
+
+  TestCorrectness(absl::StrCat(Serialize(example), Serialize(context)));
+}
+
+TEST(FastParse, DenseInt64WithContext) {
+  Example example;
+  (*example.mutable_features()->mutable_feature())["age"]
+      .mutable_int64_list()
+      ->add_value(0);
+
+  Example context;
+  (*context.mutable_features()->mutable_feature())["age"]
+      .mutable_int64_list()
+      ->add_value(15);
+
+  std::string serialized = Serialize(example) + Serialize(context);
+
+  {
+    Example deserialized;
+    EXPECT_TRUE(deserialized.ParseFromString(serialized));
+    EXPECT_EQ(deserialized.DebugString(), context.DebugString());
+    // Whoa! Last EQ is very surprising, but standard deserialization is what it
+    // is and Servo team requested to replicate this 'feature'.
+    // In future we should return error.
+  }
+  TestCorrectness(serialized);
+}
+
+TEST(FastParse, NonPacked) {
+  TestCorrectness(
+      "\x0a\x0e\x0a\x0c\x0a\x03\x61\x67\x65\x12\x05\x1a\x03\x0a\x01\x0d");
+}
+
+TEST(FastParse, Packed) {
+  TestCorrectness(
+      "\x0a\x0d\x0a\x0b\x0a\x03\x61\x67\x65\x12\x04\x1a\x02\x08\x0d");
+}
+
+TEST(FastParse, ValueBeforeKeyInMap) {
+  TestCorrectness("\x0a\x12\x0a\x10\x12\x09\x0a\x07\x0a\x05value\x0a\x03key");
+}
+
+TEST(FastParse, EmptyFeatures) {
+  Example example;
+  example.mutable_features();
+  TestCorrectness(Serialize(example));
+}
+
+void TestCorrectnessJson(const std::string& json) {
+  auto resolver = protobuf::util::NewTypeResolverForDescriptorPool(
+      "type.googleapis.com", protobuf::DescriptorPool::generated_pool());
+  std::string serialized;
+  auto s = protobuf::util::JsonToBinaryString(
+      resolver, "type.googleapis.com/tensorflow.Example", json, &serialized);
+  EXPECT_TRUE(s.ok()) << s;
+  delete resolver;
+  TestCorrectness(serialized);
+}
+
+TEST(FastParse, JsonUnivalent) {
+  TestCorrectnessJson(
+      "{'features': {"
+      "  'feature': {'age': {'int64_list': {'value': [0]} }}, "
+      "  'feature': {'flo': {'float_list': {'value': [1.1]} }}, "
+      "  'feature': {'byt': {'bytes_list': {'value': ['WW8='] }}}"
+      "}}");
+}
+
+TEST(FastParse, JsonMultivalent) {
+  TestCorrectnessJson(
+      "{'features': {"
+      "  'feature': {'age': {'int64_list': {'value': [0, 13, 23]} }}, "
+      "  'feature': {'flo': {'float_list': {'value': [1.1, 1.2, 1.3]} }}, "
+      "  'feature': {'byt': {'bytes_list': {'value': ['WW8=', 'WW8K'] }}}"
+      "}}");
+}
+
+TEST(FastParse, SingleInt64) {
+  Example example;
+  (*example.mutable_features()->mutable_feature())["age"]
+      .mutable_int64_list()
+      ->add_value(13);
+  TestCorrectness(Serialize(example));
+}
+
+static std::string ExampleWithSomeFeatures() {
+  Example example;
+
+  (*example.mutable_features()->mutable_feature())[""];
+
+  (*example.mutable_features()->mutable_feature())["empty_bytes_list"]
+      .mutable_bytes_list();
+  (*example.mutable_features()->mutable_feature())["empty_float_list"]
+      .mutable_float_list();
+  (*example.mutable_features()->mutable_feature())["empty_int64_list"]
+      .mutable_int64_list();
+
+  BytesList* bytes_list =
+      (*example.mutable_features()->mutable_feature())["bytes_list"]
+          .mutable_bytes_list();
+  bytes_list->add_value("bytes1");
+  bytes_list->add_value("bytes2");
+
+  FloatList* float_list =
+      (*example.mutable_features()->mutable_feature())["float_list"]
+          .mutable_float_list();
+  float_list->add_value(1.0);
+  float_list->add_value(2.0);
+
+  Int64List* int64_list =
+      (*example.mutable_features()->mutable_feature())["int64_list"]
+          .mutable_int64_list();
+  int64_list->add_value(3);
+  int64_list->add_value(270);
+  int64_list->add_value(86942);
+
+  return Serialize(example);
+}
+
+TEST(FastParse, SomeFeatures) { TestCorrectness(ExampleWithSomeFeatures()); }
+
+static void AddDenseFeature(const char* feature_name, DataType dtype,
+                            PartialTensorShape shape, bool variable_length,
+                            size_t elements_per_stride,
+                            FastParseExampleConfig* out_config) {
+  out_config->dense.emplace_back();
+  auto& new_feature = out_config->dense.back();
+  new_feature.feature_name = feature_name;
+  new_feature.dtype = dtype;
+  new_feature.shape = std::move(shape);
+  new_feature.default_value = Tensor(dtype, {});
+  new_feature.variable_length = variable_length;
+  new_feature.elements_per_stride = elements_per_stride;
+}
+
+static void AddSparseFeature(const char* feature_name, DataType dtype,
+                             FastParseExampleConfig* out_config) {
+  out_config->sparse.emplace_back();
+  auto& new_feature = out_config->sparse.back();
+  new_feature.feature_name = feature_name;
+  new_feature.dtype = dtype;
+}
+
+TEST(FastParse, StatsCollection) {
+  const size_t kNumExamples = 13;
+  std::vector<tstring> serialized(kNumExamples, ExampleWithSomeFeatures());
+
+  FastParseExampleConfig config_dense;
+  AddDenseFeature("bytes_list", DT_STRING, {2}, false, 2, &config_dense);
+  AddDenseFeature("float_list", DT_FLOAT, {2}, false, 2, &config_dense);
+  AddDenseFeature("int64_list", DT_INT64, {3}, false, 3, &config_dense);
+  config_dense.collect_feature_stats = true;
+
+  FastParseExampleConfig config_varlen;
+  AddDenseFeature("bytes_list", DT_STRING, {-1}, true, 1, &config_varlen);
+  AddDenseFeature("float_list", DT_FLOAT, {-1}, true, 1, &config_varlen);
+  AddDenseFeature("int64_list", DT_INT64, {-1}, true, 1, &config_varlen);
+  config_varlen.collect_feature_stats = true;
+
+  FastParseExampleConfig config_sparse;
+  AddSparseFeature("bytes_list", DT_STRING, &config_sparse);
+  AddSparseFeature("float_list", DT_FLOAT, &config_sparse);
+  AddSparseFeature("int64_list", DT_INT64, &config_sparse);
+  config_sparse.collect_feature_stats = true;
+
+  FastParseExampleConfig config_mixed;
+  AddDenseFeature("bytes_list", DT_STRING, {2}, false, 2, &config_mixed);
+  AddDenseFeature("float_list", DT_FLOAT, {-1}, true, 1, &config_mixed);
+  AddSparseFeature("int64_list", DT_INT64, &config_mixed);
+  config_mixed.collect_feature_stats = true;
+
+  for (const FastParseExampleConfig& config :
+       {config_dense, config_varlen, config_sparse, config_mixed}) {
+    {
+      Result result;
+      TF_CHECK_OK(FastParseExample(config, serialized, {}, nullptr, &result));
+      EXPECT_EQ(kNumExamples, result.feature_stats.size());
+      for (const PerExampleFeatureStats& stats : result.feature_stats) {
+        EXPECT_EQ(7, stats.features_count);
+        EXPECT_EQ(7, stats.feature_values_count);
+      }
+    }
+
+    {
+      Result result;
+      TF_CHECK_OK(FastParseSingleExample(config, serialized[0], &result));
+      EXPECT_EQ(1, result.feature_stats.size());
+      EXPECT_EQ(7, result.feature_stats[0].features_count);
+      EXPECT_EQ(7, result.feature_stats[0].feature_values_count);
+    }
+  }
+}
+
+std::string RandStr(random::SimplePhilox* rng) {
+  static const char key_char_lookup[] =
+      "0123456789{}~`!@#$%^&*()"
+      "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
+      "abcdefghijklmnopqrstuvwxyz";
+  auto len = 1 + rng->Rand32() % 200;
+  std::string str;
+  str.reserve(len);
+  while (len-- > 0) {
+    str.push_back(
+        key_char_lookup[rng->Rand32() % (sizeof(key_char_lookup) /
+                                         sizeof(key_char_lookup[0]))]);
+  }
+  return str;
+}
+
+void Fuzz(random::SimplePhilox* rng) {
+  // Generate keys.
+  auto num_keys = 1 + rng->Rand32() % 100;
+  std::unordered_set<std::string> unique_keys;
+  for (auto i = 0; i < num_keys; ++i) {
+    unique_keys.emplace(RandStr(rng));
+  }
+
+  // Generate serialized example.
+  Example example;
+  std::string serialized_example;
+  auto num_concats = 1 + rng->Rand32() % 4;
+  std::vector<Feature::KindCase> feat_types(
+      {Feature::kBytesList, Feature::kFloatList, Feature::kInt64List});
+  std::vector<std::string> all_keys(unique_keys.begin(), unique_keys.end());
+  while (num_concats--) {
+    example.Clear();
+    auto num_active_keys = 1 + rng->Rand32() % all_keys.size();
+
+    // Generate features.
+    for (auto i = 0; i < num_active_keys; ++i) {
+      auto fkey = all_keys[rng->Rand32() % all_keys.size()];
+      auto ftype_idx = rng->Rand32() % feat_types.size();
+      auto num_features = 1 + rng->Rand32() % 5;
+      switch (static_cast<Feature::KindCase>(feat_types[ftype_idx])) {
+        case Feature::kBytesList: {
+          BytesList* bytes_list =
+              (*example.mutable_features()->mutable_feature())[fkey]
+                  .mutable_bytes_list();
+          while (num_features--) {
+            bytes_list->add_value(RandStr(rng));
+          }
+          break;
+        }
+        case Feature::kFloatList: {
+          FloatList* float_list =
+              (*example.mutable_features()->mutable_feature())[fkey]
+                  .mutable_float_list();
+          while (num_features--) {
+            float_list->add_value(rng->RandFloat());
+          }
+          break;
+        }
+        case Feature::kInt64List: {
+          Int64List* int64_list =
+              (*example.mutable_features()->mutable_feature())[fkey]
+                  .mutable_int64_list();
+          while (num_features--) {
+            int64_list->add_value(rng->Rand64());
+          }
+          break;
+        }
+        default: {
+          LOG(QFATAL);
+          break;
+        }
+      }
+    }
+    serialized_example += example.SerializeAsString();
+  }
+
+  // Test correctness.
+  TestCorrectness(serialized_example);
+}
+
+TEST(FastParse, FuzzTest) {
+  const uint64_t seed = 1337;
+  random::PhiloxRandom philox(seed);
+  random::SimplePhilox rng(&philox);
+  auto num_runs = 200;
+  while (num_runs--) {
+    LOG(INFO) << "runs left: " << num_runs;
+    Fuzz(&rng);
+  }
+}
+
+TEST(TestFastParseExample, Empty) {
+  Result result;
+  FastParseExampleConfig config;
+  config.sparse.push_back({"test", DT_STRING});
+  absl::Status status =
+      FastParseExample(config, absl::Span<const tstring>(),
+                       absl::Span<const tstring>(), nullptr, &result);
+  EXPECT_TRUE(status.ok()) << status;
+}
+
+TEST(FastParse, OOB_Write_Vulnerability_NonPacked_FloatList) {
+  FastParseExampleConfig config;
+  AddDenseFeature("f", DT_FLOAT, {1}, false, 1, &config);
+
+  auto encode_varint = [](uint32_t v, std::string* out) {
+    while (v >= 0x80) {
+      out->push_back((v & 0x7f) | 0x80);
+      v >>= 7;
+    }
+    out->push_back(v);
+  };
+
+  std::string float_list_data;
+  int num_elements = 10000;  // Large number to force crash
+  for (int i = 0; i < num_elements; ++i) {
+    float_list_data.push_back(13);  // kFixed32Tag(1)
+    float v = 1.0f;
+    const char* p = reinterpret_cast<const char*>(&v);
+    float_list_data.append(p, 4);
+  }
+
+  std::string serialized_feature;
+  serialized_feature.push_back(18);  // kDelimitedTag(2) for float_list
+  encode_varint(float_list_data.size(), &serialized_feature);
+  serialized_feature.append(float_list_data);
+
+  std::string map_entry;
+  map_entry.push_back(10);  // kDelimitedTag(1) for key
+  map_entry.push_back(1);
+  map_entry.push_back('f');
+  map_entry.push_back(18);  // kDelimitedTag(2) for value
+  encode_varint(serialized_feature.size(), &map_entry);
+  map_entry.append(serialized_feature);
+
+  std::string features_msg;
+  features_msg.push_back(10);  // kDelimitedTag(1) for map entry
+  encode_varint(map_entry.size(), &features_msg);
+  features_msg.append(map_entry);
+
+  std::string serialized_example;
+  serialized_example.push_back(10);  // kDelimitedTag(1) for features
+  encode_varint(features_msg.size(), &serialized_example);
+  serialized_example.append(features_msg);
+
+  Result result;
+  std::vector<tstring> serialized_vec = {tstring(serialized_example)};
+  absl::Status parse_status =
+      FastParseExample(config, serialized_vec, {}, nullptr, &result);
+
+  // We expect this to fail with INVALID_ARGUMENT due to size mismatch,
+  // but WITHOUT crashing.
+  EXPECT_FALSE(parse_status.ok());
+  EXPECT_TRUE(absl::IsInvalidArgument(parse_status));
+}
+
+TEST(FastParse, DenseFloat_TooManyElements_ReportsError) {
+  FastParseExampleConfig config;
+  AddDenseFeature("f", DT_FLOAT, {1}, false, 1, &config);
+
+  auto encode_varint = [](uint32_t v, std::string* out) {
+    while (v >= 0x80) {
+      out->push_back((v & 0x7f) | 0x80);
+      v >>= 7;
+    }
+    out->push_back(v);
+  };
+
+  std::string float_list_data;
+  int num_elements = 5;  // Expecting 1, but providing 5
+  for (int i = 0; i < num_elements; ++i) {
+    float_list_data.push_back(13);  // kFixed32Tag(1)
+    float v = 1.0f;
+    const char* p = reinterpret_cast<const char*>(&v);
+    float_list_data.append(p, 4);
+  }
+
+  std::string serialized_feature;
+  serialized_feature.push_back(18);  // kDelimitedTag(2) for float_list
+  encode_varint(float_list_data.size(), &serialized_feature);
+  serialized_feature.append(float_list_data);
+
+  std::string map_entry;
+  map_entry.push_back(10);  // kDelimitedTag(1) for key
+  map_entry.push_back(1);
+  map_entry.push_back('f');
+  map_entry.push_back(18);  // kDelimitedTag(2) for value
+  encode_varint(serialized_feature.size(), &map_entry);
+  map_entry.append(serialized_feature);
+
+  std::string features_msg;
+  features_msg.push_back(10);  // kDelimitedTag(1) for map entry
+  encode_varint(map_entry.size(), &features_msg);
+  features_msg.append(map_entry);
+
+  std::string serialized_example;
+  serialized_example.push_back(10);  // kDelimitedTag(1) for features
+  encode_varint(features_msg.size(), &serialized_example);
+  serialized_example.append(features_msg);
+
+  Result result;
+  std::vector<tstring> serialized_vec = {tstring(serialized_example)};
+  absl::Status parse_status =
+      FastParseExample(config, serialized_vec, {}, nullptr, &result);
+
+  EXPECT_FALSE(parse_status.ok());
+  EXPECT_TRUE(absl::IsInvalidArgument(parse_status));
+  EXPECT_NE(parse_status.ToString().find("Number of float values != expected"),
+            std::string::npos);
+}
+
+TEST(FastParse, DenseFloat_TooFewElements_ReportsError) {
+  FastParseExampleConfig config;
+  // Expecting 3 elements per stride
+  AddDenseFeature("f", DT_FLOAT, {3}, false, 3, &config);
+
+  auto encode_varint = [](uint32_t v, std::string* out) {
+    while (v >= 0x80) {
+      out->push_back((v & 0x7f) | 0x80);
+      v >>= 7;
+    }
+    out->push_back(v);
+  };
+
+  std::string float_list_data;
+  int num_elements = 1;  // Providing only 1
+  for (int i = 0; i < num_elements; ++i) {
+    float_list_data.push_back(13);  // kFixed32Tag(1)
+    float v = 1.0f;
+    const char* p = reinterpret_cast<const char*>(&v);
+    float_list_data.append(p, 4);
+  }
+
+  std::string serialized_feature;
+  serialized_feature.push_back(18);  // kDelimitedTag(2) for float_list
+  encode_varint(float_list_data.size(), &serialized_feature);
+  serialized_feature.append(float_list_data);
+
+  std::string map_entry;
+  map_entry.push_back(10);  // kDelimitedTag(1) for key
+  map_entry.push_back(1);
+  map_entry.push_back('f');
+  map_entry.push_back(18);  // kDelimitedTag(2) for value
+  encode_varint(serialized_feature.size(), &map_entry);
+  map_entry.append(serialized_feature);
+
+  std::string features_msg;
+  features_msg.push_back(10);  // kDelimitedTag(1) for map entry
+  encode_varint(map_entry.size(), &features_msg);
+  features_msg.append(map_entry);
+
+  std::string serialized_example;
+  serialized_example.push_back(10);  // kDelimitedTag(1) for features
+  encode_varint(features_msg.size(), &serialized_example);
+  serialized_example.append(features_msg);
+
+  Result result;
+  std::vector<tstring> serialized_vec = {tstring(serialized_example)};
+  absl::Status parse_status =
+      FastParseExample(config, serialized_vec, {}, nullptr, &result);
+
+  EXPECT_FALSE(parse_status.ok());
+  EXPECT_TRUE(absl::IsInvalidArgument(parse_status));
+  EXPECT_NE(parse_status.ToString().find("Number of float values != expected"),
+            std::string::npos);
+}
+
+TEST(FastParse, OobWriteVulnerabilityMalformedSparseSequenceExample) {
+  FastParseExampleConfig context_config;
+  FastParseExampleConfig sequence_config;
+  AddSparseFeature("s", DT_STRING, &sequence_config);
+
+  // Serialized SequenceExample with a malformed bytes list in feature list "s".
+  // The bytes list value declares a length of 100 bytes (0x64), but the stream
+  // terminates immediately, causing CodedInputStream::Skip to fail.
+  const std::vector<tstring> serialized = {
+      tstring("\x12\x0d\x0a\x0b\x0a\x01s\x12\x06\x0a\x04\x0a\x02\x0a\x64", 15)};
+
+  Result context_result, sequence_result;
+  std::vector<Tensor> dense_lengths;
+  const absl::Status status = FastParseSequenceExample(
+      context_config, sequence_config, serialized, {}, nullptr, &context_result,
+      &sequence_result, &dense_lengths);
+
+  EXPECT_FALSE(status.ok());
+  EXPECT_TRUE(absl::IsInvalidArgument(status));
+}
+
+}  // namespace
+}  // namespace example
+}  // namespace tensorflow

tensorflow/core/util/example_proto_fast_parsing_test.proto

@@ -0,0 +1,37 @@
+// Copyright 2026 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.
+// ==============================================================================
+
+// Protocol message for the fast Example parse unit test.
+syntax = "proto3";
+
+package tensorflow;
+
+import "tensorflow/core/example/feature.proto";
+
+option cc_enable_arenas = true;
+
+// This message is parallel to Example, but with additional fields to test
+// unknown fields handling in example_proto_fast_parsing_test.cc.
+message ExampleWithExtras {
+  Features features = 1;
+
+  string extra1 = 1337;
+  int64 extra2 = 1338;
+  fixed32 extra3 = 1339;
+  fixed64 extra4 = 1340;
+  double extra5 = 1341;
+  repeated float extra6 = 1342;
+  Features extra7 = 1343;
+}

tensorflow/core/util/example_proto_helper.cc

@@ -0,0 +1,693 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/core/util/example_proto_helper.h"
+
+#include <algorithm>
+#include <cstddef>
+#include <limits>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/strings/str_format.h"
+#include "tensorflow/core/example/example.pb.h"
+#include "tensorflow/core/example/feature.pb.h"
+#include "tensorflow/core/framework/numeric_op.h"
+#include "tensorflow/core/framework/register_types.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/protobuf.h"
+#include "tensorflow/core/util/sparse/sparse_tensor.h"
+
+namespace tensorflow {
+
+absl::Status CheckValidType(const DataType& dtype) {
+  switch (dtype) {
+    case DT_INT64:
+    case DT_FLOAT:
+    case DT_STRING:
+      return absl::OkStatus();
+    default:
+      return absl::InvalidArgumentError(
+          absl::StrCat("Received input dtype: ", DataTypeString(dtype)));
+  }
+}
+
+absl::Status CheckTypesMatch(const Feature& feature, const DataType& dtype,
+                             bool* match) {
+  switch (dtype) {
+    case DT_INT64:
+      *match = (feature.kind_case() == Feature::kInt64List);
+      break;
+    case DT_FLOAT:
+      *match = (feature.kind_case() == Feature::kFloatList);
+      break;
+    case DT_STRING:
+      *match = (feature.kind_case() == Feature::kBytesList);
+      break;
+    default:
+      return absl::InvalidArgumentError(
+          absl::StrCat("Invalid input dtype: ", DataTypeString(dtype)));
+  }
+  return absl::OkStatus();
+}
+
+absl::Status FeatureDenseCopy(const std::size_t out_index,
+                              const std::string& name, const std::string& key,
+                              const DataType& dtype, const TensorShape& shape,
+                              const Feature& feature, Tensor* out) {
+  const std::size_t num_elements = shape.num_elements();
+  const std::size_t offset = out_index * num_elements;
+
+  switch (dtype) {
+    case DT_INT64: {
+      const Int64List& values = feature.int64_list();
+      if (static_cast<size_t>(values.value_size()) != num_elements) {
+        return absl::InvalidArgumentError(absl::StrCat(
+            "Name: ", name, ", Key: ", key, ", Index: ", out_index,
+            ".  Number of int64 values != expected.  "
+            "values size: ",
+            values.value_size(), " but output shape: ", shape.DebugString()));
+      }
+      auto out_p = out->flat<int64_t>().data() + offset;
+      std::copy_n(values.value().data(), num_elements, out_p);
+      return absl::OkStatus();
+    }
+    case DT_FLOAT: {
+      const FloatList& values = feature.float_list();
+      if (static_cast<size_t>(values.value_size()) != num_elements) {
+        return absl::InvalidArgumentError(absl::StrCat(
+            "Name: ", name, ", Key: ", key, ", Index: ", out_index,
+            ".  Number of float values != expected.  "
+            "values size: ",
+            values.value_size(), " but output shape: ", shape.DebugString()));
+      }
+      auto out_p = out->flat<float>().data() + offset;
+      std::copy_n(values.value().data(), num_elements, out_p);
+      return absl::OkStatus();
+    }
+    case DT_STRING: {
+      const BytesList& values = feature.bytes_list();
+      if (static_cast<size_t>(values.value_size()) != num_elements) {
+        return absl::InvalidArgumentError(absl::StrCat(
+            "Name: ", name, ", Key ", key, ", Index: ", out_index,
+            ".  Number of bytes values != expected.  "
+            "Values size: ",
+            values.value_size(), " but output shape: ", shape.DebugString()));
+      }
+      auto out_p = out->flat<tstring>().data() + offset;
+      std::transform(values.value().data(),
+                     values.value().data() + num_elements, out_p,
+                     [](const std::string* s) { return *s; });
+      return absl::OkStatus();
+    }
+    default:
+      return absl::InvalidArgumentError(
+          absl::StrCat("Invalid input dtype: ", DataTypeString(dtype)));
+  }
+}
+
+Tensor FeatureSparseCopy(const std::size_t batch, const std::string& key,
+                         const DataType& dtype, const Feature& feature) {
+  switch (dtype) {
+    case DT_INT64: {
+      const Int64List& values = feature.int64_list();
+      const int64_t num_elements = values.value_size();
+      Tensor out(dtype, TensorShape({num_elements}));
+      auto out_p = out.flat<int64_t>().data();
+      std::copy_n(values.value().data(), num_elements, out_p);
+      return out;
+    }
+    case DT_FLOAT: {
+      const FloatList& values = feature.float_list();
+      const int64_t num_elements = values.value_size();
+      Tensor out(dtype, TensorShape({num_elements}));
+      auto out_p = out.flat<float>().data();
+      std::copy_n(values.value().data(), num_elements, out_p);
+      return out;
+    }
+    case DT_STRING: {
+      const BytesList& values = feature.bytes_list();
+      const int64_t num_elements = values.value_size();
+      Tensor out(dtype, TensorShape({num_elements}));
+      auto out_p = out.flat<tstring>().data();
+      std::transform(values.value().data(),
+                     values.value().data() + num_elements, out_p,
+                     [](const std::string* s) { return *s; });
+      return out;
+    }
+    default:
+      LOG(FATAL) << "not supposed to be here.  dtype requested: " << dtype;
+  }
+}
+
+int64_t CopyIntoSparseTensor(const Tensor& in, const int batch,
+                             const int64_t offset, Tensor* indices,
+                             Tensor* values) {
+  const int64_t num_elements = in.shape().num_elements();
+  const DataType& dtype = in.dtype();
+  CHECK_EQ(dtype, values->dtype());
+
+  // Update indices.
+  if (num_elements > 0) {
+    auto ix_t = indices->matrix<int64_t>();
+    int64_t* ix_p = &ix_t(offset, 0);
+    for (int64_t i = 0; i < num_elements; ++i, ix_p += 2) {
+      *ix_p = batch;    // Column 0 stores the batch entry
+      *(ix_p + 1) = i;  // Column 1 stores the index in the batch
+    }
+  }
+
+  // Copy values over.
+  switch (dtype) {
+    case DT_INT64: {
+      std::copy_n(in.flat<int64_t>().data(), num_elements,
+                  values->flat<int64_t>().data() + offset);
+      break;
+    }
+    case DT_FLOAT: {
+      std::copy_n(in.flat<float>().data(), num_elements,
+                  values->flat<float>().data() + offset);
+      break;
+    }
+    case DT_STRING: {
+      std::copy_n(in.flat<tstring>().data(), num_elements,
+                  values->flat<tstring>().data() + offset);
+      break;
+    }
+    default:
+      LOG(FATAL) << "Not supposed to be here.  Saw dtype: " << dtype;
+  }
+
+  return num_elements;
+}
+
+void RowDenseCopy(const std::size_t& out_index, const DataType& dtype,
+                  const Tensor& in, Tensor* out) {
+  const std::size_t num_elements = in.shape().num_elements();
+  const std::size_t offset = out_index * num_elements;
+
+  switch (dtype) {
+    case DT_INT64: {
+      std::copy_n(in.flat<int64_t>().data(), num_elements,
+                  out->flat<int64_t>().data() + offset);
+      break;
+    }
+    case DT_FLOAT: {
+      std::copy_n(in.flat<float>().data(), num_elements,
+                  out->flat<float>().data() + offset);
+      break;
+    }
+    case DT_STRING: {
+      // TODO(dero): verify.
+      std::copy_n(in.flat<tstring>().data(), num_elements,
+                  out->flat<tstring>().data() + offset);
+      break;
+    }
+    default:
+      LOG(FATAL) << "Not supposed to be here.  Saw dtype: " << dtype;
+  }
+}
+
+absl::Status SingleExampleProtoToTensors(
+    const Example& example, const std::string& example_name,
+    const int batch_index,
+    const std::vector<FixedLenFeature>& fixed_len_features,
+    const std::vector<VarLenFeature>& var_len_features,
+    std::vector<Tensor*>* output_dense_values_tensor,
+    std::vector<std::vector<Tensor>>* output_sparse_values_tmp) {
+  const Features& features = example.features();
+  const auto& feature_dict = features.feature();
+
+  // Handle dense features.
+  for (size_t d = 0; d < fixed_len_features.size(); ++d) {
+    const FixedLenFeature& feature_config = fixed_len_features[d];
+    const std::string& key = feature_config.key;
+    const DataType& dtype = feature_config.dtype;
+    const TensorShape& shape = feature_config.shape;
+    const Tensor& default_value = feature_config.default_value;
+    bool required = (default_value.NumElements() == 0);
+    const auto& feature_found = feature_dict.find(key);
+    const bool feature_has_data =  // Found key & data type is set
+        (feature_found != feature_dict.end() &&
+         (feature_found->second.kind_case() != Feature::KIND_NOT_SET));
+
+    const bool required_ok = feature_has_data || !required;
+    if (!required_ok) {
+      return absl::InvalidArgumentError(
+          absl::StrCat("Name: ", example_name, ", Feature: ", key,
+                       " is required but could not be found."));
+    }
+
+    // Perform the FeatureDenseCopy into the output dense_values tensor (if
+    // the value is present).
+    if (feature_has_data) {
+      const Feature& f = feature_found->second;
+      bool types_match;
+      TF_RETURN_IF_ERROR(CheckTypesMatch(f, dtype, &types_match));
+      if (!types_match) {
+        return absl::InvalidArgumentError(absl::StrCat(
+            "Name: ", example_name, ", Feature: ", key,
+            ".  Data types don't match. ", "Expected type: ",
+            DataTypeString(dtype), "  Feature is: ", f.DebugString()));
+      }
+      TF_RETURN_IF_ERROR(FeatureDenseCopy(batch_index, example_name, key, dtype,
+                                          shape, f,
+                                          (*output_dense_values_tensor)[d]));
+    } else {
+      // If the value is missing, RowDenseCopy the default value.
+      RowDenseCopy(batch_index, dtype, default_value,
+                   (*output_dense_values_tensor)[d]);
+    }
+  }
+
+  // Handle sparse features.
+  for (size_t d = 0; d < var_len_features.size(); ++d) {
+    const VarLenFeature& feature_config = var_len_features[d];
+    const std::string& key = feature_config.key;
+    const DataType& dtype = feature_config.dtype;
+    const auto& feature_found = feature_dict.find(key);
+
+    const bool feature_has_data =  // Found key & data type is set
+        (feature_found != feature_dict.end() &&
+         (feature_found->second.kind_case() != Feature::KIND_NOT_SET));
+
+    if (feature_has_data) {
+      const Feature& f = feature_found->second;
+      bool types_match;
+      TF_RETURN_IF_ERROR(CheckTypesMatch(f, dtype, &types_match));
+      if (!types_match) {
+        return absl::InvalidArgumentError(absl::StrCat(
+            "Name: ", example_name, ", Feature: ", key,
+            ".  Data types don't match. ", "Expected type: ",
+            DataTypeString(dtype), "  Feature is: ", f.DebugString()));
+      }
+      (*output_sparse_values_tmp)[d][batch_index] =
+          FeatureSparseCopy(batch_index, key, dtype, f);
+    } else {
+      (*output_sparse_values_tmp)[d][batch_index] =
+          Tensor(dtype, TensorShape({0}));
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status GetSparseTensorShapes(const VarLenFeature& var_len_feature,
+                                   const std::vector<Tensor>& sparse_values_tmp,
+                                   const int batch_size,
+                                   VarLenFeatureBatchShapes* output_shapes) {
+  int64_t total_num_features = 0;
+  int64_t max_num_features = 0;
+  for (int b = 0; b < batch_size; ++b) {
+    const Tensor& t = sparse_values_tmp[b];
+    const int64_t num_elements = t.shape().num_elements();
+    total_num_features += num_elements;
+    max_num_features = std::max(max_num_features, num_elements);
+  }
+  output_shapes->indices_shape.AddDim(total_num_features);
+  output_shapes->indices_shape.AddDim(2);
+  output_shapes->values_shape.AddDim(total_num_features);
+  output_shapes->max_num_features = max_num_features;
+  return absl::OkStatus();
+}
+
+absl::Status BatchExampleProtoToTensors(
+    const std::vector<const Example*>& examples,
+    const std::vector<std::string>& names,
+    const std::vector<FixedLenFeature>& fixed_len_features,
+    const std::vector<VarLenFeature>& var_len_features, Allocator* allocator,
+    std::vector<Tensor>* output_dense_values_tensor,
+    std::vector<Tensor>* output_sparse_indices_tensor,
+    std::vector<Tensor>* output_sparse_values_tensor,
+    std::vector<Tensor>* output_sparse_shapes_tensor) {
+  const int batch_size = examples.size();
+
+  const bool has_names = (!names.empty());
+  if (has_names) {
+    if (names.size() != examples.size()) {
+      return absl::InvalidArgumentError(absl::StrCat(
+          "Expected len(names) == len(examples), but got: ", names.size(),
+          " vs. ", examples.size()));
+    }
+  }
+
+  // We also need a map of Tensor pointers for the SingleExampleProtoToTensors
+  // call. (Is there a better solution here?)
+  std::vector<Tensor*> output_dense_values_tensor_ptrs(
+      fixed_len_features.size());
+
+  // Preallocate dense_values, since we know their sizes.
+  for (size_t d = 0; d < fixed_len_features.size(); ++d) {
+    const FixedLenFeature& config = fixed_len_features[d];
+    TensorShape out_shape;
+    out_shape.AddDim(batch_size);
+    const TensorShape& shape = config.shape;
+    const DataType& dtype = config.dtype;
+    for (const int dim : shape.dim_sizes()) out_shape.AddDim(dim);
+    (*output_dense_values_tensor)[d] = Tensor(allocator, dtype, out_shape);
+    output_dense_values_tensor_ptrs[d] = &(*output_dense_values_tensor)[d];
+  }
+
+  // Temporary vector to hold sparse values.
+  std::vector<std::vector<Tensor>> sparse_values_tmp(var_len_features.size());
+
+  for (size_t d = 0; d < var_len_features.size(); ++d) {
+    sparse_values_tmp[d] = std::vector<Tensor>(batch_size);
+  }
+
+  for (size_t b = 0; b < examples.size(); ++b) {
+    const Example& ex = *(examples[b]);
+    const std::string& example_name = (has_names) ? names[b] : "<unknown>";
+    TF_RETURN_IF_ERROR(SingleExampleProtoToTensors(
+        ex, example_name, b, fixed_len_features, var_len_features,
+        &output_dense_values_tensor_ptrs, &sparse_values_tmp));
+  }
+
+  for (size_t d = 0; d < var_len_features.size(); ++d) {
+    const VarLenFeature& feature_config = var_len_features[d];
+    const DataType& dtype = feature_config.dtype;
+    const std::vector<Tensor>& sparse_values_tensor = sparse_values_tmp[d];
+
+    VarLenFeatureBatchShapes sparse_tensor_batch_shapes;
+    TF_RETURN_IF_ERROR(GetSparseTensorShapes(feature_config,
+                                             sparse_values_tensor, batch_size,
+                                             &sparse_tensor_batch_shapes));
+    const TensorShape& indices_shape = sparse_tensor_batch_shapes.indices_shape;
+    const TensorShape& values_shape = sparse_tensor_batch_shapes.values_shape;
+
+    // Allocate the sparse indices here.
+    (*output_sparse_indices_tensor)[d] =
+        Tensor(allocator, DT_INT64, indices_shape);
+    (*output_sparse_values_tensor)[d] = Tensor(allocator, dtype, values_shape);
+    (*output_sparse_shapes_tensor)[d] =
+        Tensor(allocator, DT_INT64, TensorShape({2}));
+
+    auto shape_t = (*output_sparse_shapes_tensor)[d].vec<int64_t>();
+    shape_t(0) = batch_size;
+    shape_t(1) = sparse_tensor_batch_shapes.max_num_features;
+
+    Tensor* sp_indices_d = &(*output_sparse_indices_tensor)[d];
+    Tensor* sp_values_d = &(*output_sparse_values_tensor)[d];
+
+    int64_t offset = 0;
+    for (int b = 0; b < batch_size; ++b) {
+      const int64_t num_elements = CopyIntoSparseTensor(
+          sparse_values_tensor[b], b, offset, sp_indices_d, sp_values_d);
+      offset += num_elements;
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status ParseExampleAttrs::UpdateDenseShapes(
+    const std::vector<size_t>& got_dims) {
+  if (got_dims.size() != dense_shapes.size()) {
+    return absl::InvalidArgumentError(absl::StrFormat(
+        "got_dims.size() (%d) must match dense_shapes.size() (%d)",
+        got_dims.size(), dense_shapes.size()));
+  }
+  for (size_t d = 0; d < dense_shapes.size(); ++d) {
+    dense_shapes[d].set_dim(0, got_dims[d]);
+  }
+  // Recalculate relative fields.
+  variable_length.clear();
+  elements_per_stride.clear();
+  return GetDenseShapes(dense_shapes, &variable_length, &elements_per_stride);
+}
+
+absl::Status ParseExampleAttrs::FinishInit(int op_version) {
+  switch (op_version) {
+    case 1:
+      num_ragged = 0;
+      break;
+    case 2:
+      num_dense = dense_types.size();
+      num_ragged = ragged_value_types.size();
+      break;
+    default:
+      return absl::InvalidArgumentError(
+          absl::StrCat("Unexpected op_version", op_version));
+  }
+  if (static_cast<size_t>(num_sparse) != sparse_types.size()) {
+    return absl::InvalidArgumentError("len(sparse_keys) != len(sparse_types)");
+  }
+  if (static_cast<size_t>(num_dense) != dense_types.size()) {
+    return absl::InvalidArgumentError("len(dense_keys) != len(dense_types)");
+  }
+  if (static_cast<size_t>(num_dense) != dense_shapes.size()) {
+    return absl::InvalidArgumentError("len(dense_keys) != len(dense_shapes)");
+  }
+  if (static_cast<size_t>(num_ragged) != ragged_value_types.size()) {
+    return absl::InvalidArgumentError(
+        "len(ragged_keys) != len(ragged_value_types)");
+  }
+  if (static_cast<size_t>(num_ragged) != ragged_split_types.size()) {
+    return absl::InvalidArgumentError(
+        "len(ragged_keys) != len(ragged_split_types)");
+  }
+  if (num_dense > std::numeric_limits<int32_t>::max()) {
+    return absl::InvalidArgumentError("num_dense_ too large");
+  }
+  for (const DataType& type : dense_types) {
+    TF_RETURN_IF_ERROR(CheckValidType(type));
+  }
+  for (const DataType& type : sparse_types) {
+    TF_RETURN_IF_ERROR(CheckValidType(type));
+  }
+  for (const DataType& type : ragged_value_types) {
+    TF_RETURN_IF_ERROR(CheckValidType(type));
+  }
+  for (const DataType& type : ragged_split_types) {
+    if (!(type == DT_INT64 || type == DT_INT32)) {
+      return absl::InvalidArgumentError(
+          absl::StrCat("Invalid ragged_split_type: ", DataTypeString(type)));
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status ParseSingleExampleAttrs::FinishInit() {
+  if (sparse_keys.size() != sparse_types.size()) {
+    return absl::InvalidArgumentError("len(sparse_keys) != len(sparse_types)");
+  }
+  if (dense_keys.size() != dense_types.size()) {
+    return absl::InvalidArgumentError("len(dense_keys) != len(dense_types)");
+  }
+  if (dense_keys.size() != dense_shapes.size()) {
+    return absl::InvalidArgumentError("len(dense_keys) != len(dense_shapes)");
+  }
+  for (const DataType& type : dense_types) {
+    TF_RETURN_IF_ERROR(CheckValidType(type));
+  }
+  for (const DataType& type : sparse_types) {
+    TF_RETURN_IF_ERROR(CheckValidType(type));
+  }
+  return absl::OkStatus();
+}
+
+absl::Status ParseSequenceExampleAttrs::FinishInit(int op_version) {
+  switch (op_version) {
+    case 1:
+      num_context_ragged = 0;
+      num_feature_list_ragged = 0;
+      if (num_context_sparse != context_sparse_keys.size()) {
+        return absl::InvalidArgumentError(
+            absl::StrCat("num_context_sparse (", num_context_sparse,
+                         ") must match the size of context_sparse_keys (",
+                         context_sparse_keys.size(), ")"));
+      }
+      if (num_context_dense != context_dense_keys.size()) {
+        return absl::InvalidArgumentError(
+            absl::StrCat("num_context_dense (", num_context_dense,
+                         ") must match the size of context_dense_keys (",
+                         context_dense_keys.size(), ")"));
+      }
+      if (num_feature_list_sparse != feature_list_sparse_keys.size()) {
+        return absl::InvalidArgumentError(
+            absl::StrCat("num_feature_list_sparse (", num_feature_list_sparse,
+                         ") must match the size of feature_list_sparse_keys (",
+                         feature_list_sparse_keys.size(), ")"));
+      }
+      if (num_feature_list_dense != feature_list_dense_keys.size()) {
+        return absl::InvalidArgumentError(
+            absl::StrCat("num_feature_list_dense (", num_feature_list_dense,
+                         ") must match the size of feature_list_dense_keys (",
+                         feature_list_dense_keys.size(), ")"));
+      }
+      break;
+    case 2:
+      num_context_dense = context_dense_types.size();
+      num_context_ragged = context_ragged_value_types.size();
+      num_feature_list_ragged = feature_list_ragged_value_types.size();
+      break;
+    default:
+      return absl::InvalidArgumentError(
+          absl::StrCat("Unexpected op_version", op_version));
+  }
+  if (num_context_sparse != context_sparse_types.size()) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("num_context_sparse (", num_context_sparse,
+                     ") must match the size of context_sparse_types (",
+                     context_sparse_types.size(), ")"));
+  }
+  if (num_context_dense != context_dense_types.size() ||
+      num_context_dense != context_dense_shapes.size()) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("num_context_dense (", num_context_dense,
+                     ") must match the size of context_dense_types (",
+                     context_dense_types.size(), ") and context_dense_shapes (",
+                     context_dense_shapes.size(), ")"));
+  }
+  if ((num_context_ragged != context_ragged_value_types.size()) ||
+      (num_context_ragged != context_ragged_split_types.size())) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "num_context_ragged (", num_context_ragged,
+        ") must match the size of context_ragged_value_types (",
+        context_ragged_value_types.size(), ") and context_ragged_split_types (",
+        context_ragged_split_types.size(), ")"));
+  }
+  if (num_feature_list_sparse != feature_list_sparse_types.size()) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("num_feature_list_sparse (", num_feature_list_sparse,
+                     ") must match the size of feature_list_sparse_types (",
+                     feature_list_sparse_types.size(), ")"));
+  }
+  if (num_feature_list_dense != feature_list_dense_types.size() ||
+      num_feature_list_dense != feature_list_dense_shapes.size()) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "num_feature_list_dense (", num_feature_list_dense,
+        ") must match the size of feature_list_dense_types (",
+        feature_list_dense_types.size(), ") and feature_list_dense_shapes (",
+        feature_list_dense_shapes.size(), ")"));
+  }
+  if ((num_feature_list_ragged != feature_list_ragged_value_types.size()) ||
+      (num_feature_list_ragged != feature_list_ragged_split_types.size())) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "num_feature_list_ragged (", num_feature_list_ragged,
+        ") must match the size of feature_list_ragged_value_types (",
+        feature_list_ragged_value_types.size(),
+        ") and feature_list_ragged_split_types (",
+        feature_list_ragged_split_types.size(), ")"));
+  }
+  for (const DataType& type : context_dense_types) {
+    TF_RETURN_IF_ERROR(CheckValidType(type));
+  }
+  for (const DataType& type : context_sparse_types) {
+    TF_RETURN_IF_ERROR(CheckValidType(type));
+  }
+  for (const DataType& type : feature_list_dense_types) {
+    TF_RETURN_IF_ERROR(CheckValidType(type));
+  }
+  for (const DataType& type : feature_list_sparse_types) {
+    TF_RETURN_IF_ERROR(CheckValidType(type));
+  }
+  for (const DataType& type : context_ragged_value_types) {
+    TF_RETURN_IF_ERROR(CheckValidType(type));
+  }
+  for (const DataType& type : context_ragged_split_types) {
+    if (!(type == DT_INT64 || type == DT_INT32)) {
+      return absl::InvalidArgumentError(absl::StrCat(
+          "Invalid context_ragged_split_type: ", DataTypeString(type)));
+    }
+  }
+  for (const DataType& type : feature_list_ragged_value_types) {
+    TF_RETURN_IF_ERROR(CheckValidType(type));
+  }
+  for (const DataType& type : feature_list_ragged_split_types) {
+    if (!(type == DT_INT64 || type == DT_INT32)) {
+      return absl::InvalidArgumentError(absl::StrCat(
+          "Invalid feature_list_ragged_split_type: ", DataTypeString(type)));
+    }
+  }
+
+  return absl::OkStatus();
+}
+
+absl::Status ParseSingleSequenceExampleAttrs::FinishInit() {
+  if (static_cast<size_t>(num_context_sparse) != context_sparse_types.size()) {
+    return absl::InvalidArgumentError(
+        "len(context_sparse_keys) != len(context_sparse_types)");
+  }
+  if (static_cast<size_t>(num_context_dense) != context_dense_types.size()) {
+    return absl::InvalidArgumentError(
+        "len(context_dense_keys) != len(context_dense_types)");
+  }
+  if (static_cast<size_t>(num_context_dense) != context_dense_shapes.size()) {
+    return absl::InvalidArgumentError(
+        "len(context_dense_keys) != len(context_dense_shapes)");
+  }
+  if (static_cast<size_t>(num_feature_list_sparse) !=
+      feature_list_sparse_types.size()) {
+    return absl::InvalidArgumentError(
+        "len(feature_list_sparse_keys) != len(feature_list_sparse_types)");
+  }
+  if (static_cast<size_t>(num_feature_list_dense) !=
+      feature_list_dense_types.size()) {
+    return absl::InvalidArgumentError(
+        "len(feature_list_dense_keys) != "
+        "len(feature_list_dense_types)");
+  }
+  for (const DataType& type : context_dense_types) {
+    TF_RETURN_IF_ERROR(CheckValidType(type));
+  }
+  for (const DataType& type : context_sparse_types) {
+    TF_RETURN_IF_ERROR(CheckValidType(type));
+  }
+  for (const DataType& type : feature_list_dense_types) {
+    TF_RETURN_IF_ERROR(CheckValidType(type));
+  }
+  for (const DataType& type : feature_list_sparse_types) {
+    TF_RETURN_IF_ERROR(CheckValidType(type));
+  }
+  return absl::OkStatus();
+}
+
+absl::Status GetDenseShapes(const std::vector<PartialTensorShape>& dense_shapes,
+                            std::vector<bool>* variable_length,
+                            std::vector<std::size_t>* elements_per_stride) {
+  // Temporary check until we start allowing a variable length outer
+  // dimension.
+  for (int i = 0; i < dense_shapes.size(); ++i) {
+    bool shape_ok = true;
+    if (dense_shapes[i].dims() == -1) {
+      shape_ok = false;
+    } else {
+      for (int d = 1; d < dense_shapes[i].dims(); ++d) {
+        if (dense_shapes[i].dim_size(d) == -1) {
+          shape_ok = false;
+        }
+      }
+    }
+    if (!shape_ok) {
+      return absl::InvalidArgumentError(
+          absl::StrCat("dense_shapes[", i,
+                       "] has unknown rank or unknown inner dimensions: ",
+                       dense_shapes[i].DebugString()));
+    }
+    TensorShape dense_shape;
+    if (dense_shapes[i].dims() > 0 && dense_shapes[i].dim_size(0) == -1) {
+      variable_length->push_back(true);
+      for (int d = 1; d < dense_shapes[i].dims(); ++d) {
+        dense_shape.AddDim(dense_shapes[i].dim_size(d));
+      }
+    } else {
+      variable_length->push_back(false);
+      dense_shapes[i].AsTensorShape(&dense_shape);
+    }
+    elements_per_stride->push_back(dense_shape.num_elements());
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace tensorflow

tensorflow/core/util/example_proto_helper.h

@@ -0,0 +1,376 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_UTIL_EXAMPLE_PROTO_HELPER_H_
+#define TENSORFLOW_CORE_UTIL_EXAMPLE_PROTO_HELPER_H_
+
+#include <cstddef>
+#include <cstdint>
+#include <unordered_set>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "xla/tsl/platform/errors.h"
+#include "tensorflow/core/example/example.pb.h"
+#include "tensorflow/core/example/feature.pb.h"
+#include "tensorflow/core/framework/allocator.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/platform/tstring.h"
+#include "tensorflow/core/platform/types.h"
+
+// This is a set of helper methods that will make it possible to share
+// tensorflow::Example proto Tensor conversion code inside the ExampleParserOp
+// OpKernel as well as in external code.
+namespace tensorflow {
+
+// "Dense" feature configuration.
+struct FixedLenFeature {
+  std::string key;
+  DataType dtype;
+  TensorShape shape;
+  Tensor default_value;
+  std::string values_output_tensor_name;
+};
+
+// "Sparse" feature configuration.
+struct VarLenFeature {
+  std::string key;
+  DataType dtype;
+  std::string values_output_tensor_name;
+  std::string indices_output_tensor_name;
+  std::string shapes_output_tensor_name;
+};
+
+// Given a single tensorflow::Example, with an optional example name
+// at a particular index within a batch, and dense and sparse feature
+// configurations from fixed_len_features, var_len_features, this method
+// updates the dense value tensor and the sparse values temporary vector
+// of tensors. The indexing of the output vectors correspond 1:1 to the
+// indexing of the feature configuration vectors.
+//
+// The fixed_len_features and var_len_features maps are assume to be
+// have disjoint key fields from the Feature map in the tensorflow.Example
+// proto.
+//
+// For each sparse feature, the sparse values temporary vector holds a
+// tensor for each Example. Each tensor is either empty or filled, depending
+// on if the sparse feature value is set for the Example. This
+// temporary structure is needed because we need to know the total number
+// of filled elements in the batch to get the proper final sparse tensor
+// shapes allocated.  After the entire batch is processed,
+// GetSparseTensorShape can be used to calculate the final shapes and
+// CopyIntoSparseTensor can be used to copy from the temporary vector
+// into the final allocated tensors.
+absl::Status SingleExampleProtoToTensors(
+    const Example& example, const std::string& name, int batch_index,
+    const std::vector<FixedLenFeature>& fixed_len_features,
+    const std::vector<VarLenFeature>& var_len_features,
+    std::vector<Tensor*>* output_dense_values_tensor,
+    std::vector<std::vector<Tensor>>* output_sparse_values_tmp);
+
+// The shape of the indices and values tensors associated with a SparseTensor
+// are dependent on the contents of the batch.
+struct VarLenFeatureBatchShapes {
+  TensorShape indices_shape;
+  TensorShape values_shape;
+  int max_num_features;
+};
+
+// Get the shape of the sparse values and indices tensors for the batch,
+// given how many of the tensors in the temporary sparse values vector
+// are actually filled.
+absl::Status GetSparseTensorShapes(const VarLenFeature& var_len_feature,
+                                   const std::vector<Tensor>& sparse_values_tmp,
+                                   int batch_size,
+                                   VarLenFeatureBatchShapes* output_shapes);
+
+// A method to convert a batch of tensorflow::Example protos into output
+// tensors. This method is useful if there already is a batch of deserialized
+// Example protos in memory (such as a serving use-case) and we do not wish
+// to incur an extraneous serialize/deserialize.  It is intended
+// as an outside of OpKernel compatible replacement for the functionality of
+// ExampleParserOp. In a serving setting, this method could be used to produce
+// a feed_dict of Tensors that could bypass the ExampleParserOp.
+//
+// Note that unlike SingleExampleProtoToTensors, output tensors are
+// allocated using a provided Allocator within this method.
+absl::Status BatchExampleProtoToTensors(
+    const std::vector<const Example*>& examples,
+    const std::vector<std::string>& names,
+    const std::vector<FixedLenFeature>& fixed_len_features,
+    const std::vector<VarLenFeature>& var_len_features, Allocator* allocator,
+    std::vector<Tensor>* output_dense_values_tensor,
+    std::vector<Tensor>* output_sparse_indices_tensor,
+    std::vector<Tensor>* output_sparse_values_tensor,
+    std::vector<Tensor>* output_sparse_shapes_tensor);
+
+// Check that the given dtype is one that is compatible with
+// tensorflow::Example protocol buffer feature values.
+absl::Status CheckValidType(const DataType& dtype);
+
+// Check that the provided Feature proto message's oneof value
+// matches that of the provided dtype.
+absl::Status CheckTypesMatch(const Feature& feature, const DataType& dtype,
+                             bool* match);
+
+// For a single Example, copy a dense feature value into an output
+// dense value tensor Out at the provided out_index offset.
+absl::Status FeatureDenseCopy(std::size_t out_index, const std::string& name,
+                              const std::string& key, const DataType& dtype,
+                              const TensorShape& shape, const Feature& feature,
+                              Tensor* out);
+
+// Copy the value a provided Tensor into an output dense_value tensor Out
+// at the provided out_index offset.
+void RowDenseCopy(const std::size_t& out_index, const DataType& dtype,
+                  const Tensor& in, Tensor* out);
+
+// For a single Example, and given sparse feature return a temporary output
+// Tensor suitable for being collected in the temporary sparse value vector.
+Tensor FeatureSparseCopy(std::size_t batch, const std::string& key,
+                         const DataType& dtype, const Feature& feature);
+
+// Copy a temporary Tensor into the final sparse indices and values
+// tensor at a given batch index and element offset. This method
+// assumes that the indices/values Tensors have been properly allocated
+// for the batch.
+int64_t CopyIntoSparseTensor(const Tensor& in, int batch, int64_t offset,
+                             Tensor* indices, Tensor* values);
+
+// Check that each dense_shape has known rank and inner dimensions; and
+// update variable_length (whether the outer dimension is None) and
+// elements_per_stride for each denes_shape.
+absl::Status GetDenseShapes(const std::vector<PartialTensorShape>& dense_shapes,
+                            std::vector<bool>* variable_length,
+                            std::vector<std::size_t>* elements_per_stride);
+
+// Parses the attributes passed to ParseExample.
+// REQUIRES: Init must be called after construction.
+struct ParseExampleAttrs {
+ public:
+  template <typename ContextType>
+  absl::Status Init(ContextType* ctx, int op_version = 1) {
+    TF_RETURN_IF_ERROR(ctx->GetAttr("sparse_types", &sparse_types));
+    TF_RETURN_IF_ERROR(ctx->GetAttr("Tdense", &dense_types));
+    TF_RETURN_IF_ERROR(ctx->GetAttr("dense_shapes", &dense_shapes));
+    TF_RETURN_IF_ERROR(
+        GetDenseShapes(dense_shapes, &variable_length, &elements_per_stride));
+    switch (op_version) {
+      case 1:
+        TF_RETURN_IF_ERROR(ctx->GetAttr("Nsparse", &num_sparse));
+        TF_RETURN_IF_ERROR(ctx->GetAttr("Ndense", &num_dense));
+        break;
+      case 2:
+        TF_RETURN_IF_ERROR(
+            ctx->GetAttr("ragged_value_types", &ragged_value_types));
+        TF_RETURN_IF_ERROR(ctx->GetAttr("num_sparse", &num_sparse));
+        TF_RETURN_IF_ERROR(
+            ctx->GetAttr("ragged_split_types", &ragged_split_types));
+        break;
+      default:
+        return absl::InvalidArgumentError(
+            absl::StrCat("Unexpected op_version", op_version));
+    }
+    return FinishInit(op_version);
+  }
+
+  absl::Status UpdateDenseShapes(const std::vector<size_t>& got_dims);
+
+  int64_t num_sparse;
+  int64_t num_dense;
+  int64_t num_ragged;
+  std::vector<DataType> sparse_types;
+  std::vector<DataType> dense_types;
+  std::vector<DataType> ragged_value_types;
+  std::vector<DataType> ragged_split_types;
+  std::vector<PartialTensorShape> dense_shapes;
+  std::vector<bool> variable_length;
+  std::vector<std::size_t> elements_per_stride;
+
+ private:
+  absl::Status FinishInit(
+      int op_version);  // for context-independent parts of Init.
+};
+
+// Parses the attributes passed to ParseSingleExample.
+// REQUIRES: Init must be called after construction.
+struct ParseSingleExampleAttrs {
+ public:
+  template <typename ContextType>
+  absl::Status Init(ContextType* ctx) {
+    TF_RETURN_IF_ERROR(ctx->GetAttr("sparse_keys", &sparse_keys));
+    TF_RETURN_IF_ERROR(ctx->GetAttr("sparse_types", &sparse_types));
+    TF_RETURN_IF_ERROR(ctx->GetAttr("dense_keys", &dense_keys));
+    TF_RETURN_IF_ERROR(ctx->GetAttr("Tdense", &dense_types));
+    TF_RETURN_IF_ERROR(ctx->GetAttr("dense_shapes", &dense_shapes));
+
+    int num_sparse;
+    TF_RETURN_IF_ERROR(ctx->GetAttr("num_sparse", &num_sparse));
+    if (num_sparse != sparse_keys.size() || num_sparse != sparse_types.size()) {
+      return absl::InvalidArgumentError(absl::StrCat(
+          "num_sparse (", num_sparse, ") must match the size of sparse_keys (",
+          sparse_keys.size(), ") and sparse_types (", sparse_types.size(),
+          ")"));
+    }
+
+    TF_RETURN_IF_ERROR(
+        GetDenseShapes(dense_shapes, &variable_length, &elements_per_stride));
+    return FinishInit();
+  }
+
+  std::vector<tstring> sparse_keys;
+  std::vector<DataType> sparse_types;
+  std::vector<tstring> dense_keys;
+  std::vector<DataType> dense_types;
+  std::vector<PartialTensorShape> dense_shapes;
+  std::vector<bool> variable_length;
+  std::vector<std::size_t> elements_per_stride;
+
+ private:
+  absl::Status FinishInit();  // for context-independent parts of Init.
+};
+
+// Parses the attributes passed to ParseSequenceExample.
+// REQUIRES: Init must be called after construction.
+struct ParseSequenceExampleAttrs {
+ public:
+  template <typename ContextType>
+  absl::Status Init(ContextType* ctx, int op_version = 1) {
+    switch (op_version) {
+      case 1: {
+        std::vector<std::string> missing_empty_vector;
+        TF_RETURN_IF_ERROR(ctx->GetAttr(
+            "feature_list_dense_missing_assumed_empty", &missing_empty_vector));
+        for (const std::string& feature : missing_empty_vector) {
+          feature_list_dense_missing_assumed_empty.insert(feature);
+        }
+      }
+        TF_RETURN_IF_ERROR(
+            ctx->GetAttr("context_sparse_keys", &context_sparse_keys));
+        TF_RETURN_IF_ERROR(
+            ctx->GetAttr("context_dense_keys", &context_dense_keys));
+        TF_RETURN_IF_ERROR(ctx->GetAttr("feature_list_sparse_keys",
+                                        &feature_list_sparse_keys));
+        TF_RETURN_IF_ERROR(
+            ctx->GetAttr("feature_list_dense_keys", &feature_list_dense_keys));
+        TF_RETURN_IF_ERROR(ctx->GetAttr("Ncontext_dense", &num_context_dense));
+        break;
+      case 2:
+        TF_RETURN_IF_ERROR(ctx->GetAttr("context_ragged_value_types",
+                                        &context_ragged_value_types));
+        TF_RETURN_IF_ERROR(ctx->GetAttr("context_ragged_split_types",
+                                        &context_ragged_split_types));
+        TF_RETURN_IF_ERROR(ctx->GetAttr("feature_list_ragged_value_types",
+                                        &feature_list_ragged_value_types));
+        TF_RETURN_IF_ERROR(ctx->GetAttr("feature_list_ragged_split_types",
+                                        &feature_list_ragged_split_types));
+        break;
+      default:
+        return absl::InvalidArgumentError(
+            absl::StrCat("Unexpected op_version", op_version));
+    }
+    TF_RETURN_IF_ERROR(
+        ctx->GetAttr("context_sparse_types", &context_sparse_types));
+    TF_RETURN_IF_ERROR(
+        ctx->GetAttr("Nfeature_list_dense", &num_feature_list_dense));
+    TF_RETURN_IF_ERROR(ctx->GetAttr("Ncontext_sparse", &num_context_sparse));
+    TF_RETURN_IF_ERROR(ctx->GetAttr("Tcontext_dense", &context_dense_types));
+    TF_RETURN_IF_ERROR(
+        ctx->GetAttr("feature_list_sparse_types", &feature_list_sparse_types));
+    TF_RETURN_IF_ERROR(
+        ctx->GetAttr("feature_list_dense_types", &feature_list_dense_types));
+    TF_RETURN_IF_ERROR(
+        ctx->GetAttr("Nfeature_list_sparse", &num_feature_list_sparse));
+    TF_RETURN_IF_ERROR(
+        ctx->GetAttr("context_dense_shapes", &context_dense_shapes));
+    TF_RETURN_IF_ERROR(
+        ctx->GetAttr("feature_list_dense_shapes", &feature_list_dense_shapes));
+    return FinishInit(op_version);
+  }
+
+  std::unordered_set<std::string> feature_list_dense_missing_assumed_empty;
+  int64_t num_context_sparse;
+  int64_t num_context_dense;
+  int64_t num_context_ragged;
+  int64_t num_feature_list_sparse;
+  int64_t num_feature_list_dense;
+  int64_t num_feature_list_ragged;
+  std::vector<tstring> context_sparse_keys;
+  std::vector<tstring> context_dense_keys;
+  std::vector<tstring> feature_list_sparse_keys;
+  std::vector<tstring> feature_list_dense_keys;
+  std::vector<DataType> context_sparse_types;
+  std::vector<DataType> context_dense_types;
+  std::vector<TensorShape> context_dense_shapes;
+  std::vector<DataType> feature_list_sparse_types;
+  std::vector<DataType> feature_list_dense_types;
+  std::vector<TensorShape> feature_list_dense_shapes;
+  std::vector<DataType> context_ragged_value_types;
+  std::vector<DataType> context_ragged_split_types;
+  std::vector<DataType> feature_list_ragged_value_types;
+  std::vector<DataType> feature_list_ragged_split_types;
+
+ private:
+  absl::Status FinishInit(
+      int op_version);  // for context-independent parts of Init.
+};
+
+// Parses the attributes passed to ParseSingleSequenceExample.
+// REQUIRES: Init must be called after construction.
+struct ParseSingleSequenceExampleAttrs {
+ public:
+  template <typename ContextType>
+  absl::Status Init(ContextType* ctx) {
+    TF_RETURN_IF_ERROR(
+        ctx->GetAttr("context_sparse_types", &context_sparse_types));
+    TF_RETURN_IF_ERROR(ctx->GetAttr("Ncontext_dense", &num_context_dense));
+    TF_RETURN_IF_ERROR(
+        ctx->GetAttr("Nfeature_list_dense", &num_feature_list_dense));
+    TF_RETURN_IF_ERROR(ctx->GetAttr("Ncontext_sparse", &num_context_sparse));
+    TF_RETURN_IF_ERROR(ctx->GetAttr("Tcontext_dense", &context_dense_types));
+    TF_RETURN_IF_ERROR(
+        ctx->GetAttr("feature_list_sparse_types", &feature_list_sparse_types));
+    TF_RETURN_IF_ERROR(
+        ctx->GetAttr("feature_list_dense_types", &feature_list_dense_types));
+    TF_RETURN_IF_ERROR(
+        ctx->GetAttr("Nfeature_list_sparse", &num_feature_list_sparse));
+    TF_RETURN_IF_ERROR(
+        ctx->GetAttr("context_dense_shapes", &context_dense_shapes));
+    TF_RETURN_IF_ERROR(
+        ctx->GetAttr("feature_list_dense_shapes", &feature_list_dense_shapes));
+    return FinishInit();
+  }
+
+  int64_t num_context_sparse;
+  int64_t num_context_dense;
+  int64_t num_feature_list_sparse;
+  int64_t num_feature_list_dense;
+  std::vector<DataType> context_sparse_types;
+  std::vector<DataType> context_dense_types;
+  std::vector<TensorShape> context_dense_shapes;
+  std::vector<DataType> feature_list_sparse_types;
+  std::vector<DataType> feature_list_dense_types;
+  std::vector<TensorShape> feature_list_dense_shapes;
+
+ private:
+  absl::Status FinishInit();  // for context-independent parts of Init.
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_UTIL_EXAMPLE_PROTO_HELPER_H_

tensorflow/core/util/example_proto_helper_test.cc

@@ -0,0 +1,261 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/core/util/example_proto_helper.h"
+
+#include <cstdint>
+#include <vector>
+
+#include "tensorflow/core/example/example.pb.h"
+#include "tensorflow/core/example/feature.pb.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/platform/tstring.h"
+
+namespace tensorflow {
+namespace {
+
+TEST(CopyIntoSparseTensorTest, String) {
+  Tensor in_tensor(DT_STRING, TensorShape({2}));
+  in_tensor.flat<tstring>()(0) = "hello";
+  in_tensor.flat<tstring>()(1) = "world";
+  int n_values = 5;
+  Tensor ix_tensor(DT_INT64, TensorShape({n_values, 2}));
+  auto ix_matrix = ix_tensor.matrix<int64_t>();
+  for (int i = 0; i < n_values; ++i) {
+    for (int j = 0; j < 2; ++j) {
+      ix_matrix(i, j) = 0;
+    }
+  }
+  Tensor value_tensor(DT_STRING, TensorShape({n_values}));
+  int batch = 67;
+  int64_t offset = 1;
+  auto n_elems =
+      CopyIntoSparseTensor(in_tensor, batch, offset, &ix_tensor, &value_tensor);
+  EXPECT_EQ(2, n_elems);
+
+  EXPECT_EQ(0, ix_matrix(0, 0));
+  EXPECT_EQ(0, ix_matrix(0, 1));
+  EXPECT_EQ(batch, ix_matrix(1, 0));
+  EXPECT_EQ(0, ix_matrix(1, 1));
+  EXPECT_EQ(batch, ix_matrix(2, 0));
+  EXPECT_EQ(1, ix_matrix(2, 1));
+  EXPECT_EQ(0, ix_matrix(3, 0));
+  EXPECT_EQ(0, ix_matrix(3, 1));
+  EXPECT_EQ(0, ix_matrix(4, 0));
+  EXPECT_EQ(0, ix_matrix(4, 1));
+  auto values = value_tensor.flat<tstring>();
+  EXPECT_EQ("", values(0));
+  EXPECT_EQ("hello", values(1));
+  EXPECT_EQ("world", values(2));
+  EXPECT_EQ("", values(3));
+  EXPECT_EQ("", values(4));
+}
+
+constexpr char kDenseInt64Key[] = "dense_int64";
+constexpr char kDenseFloatKey[] = "dense_float";
+constexpr char kDenseStringKey[] = "dense_string";
+
+constexpr char kSparseInt64Key[] = "sparse_int64";
+constexpr char kSparseFloatKey[] = "sparse_float";
+constexpr char kSparseStringKey[] = "sparse_string";
+
+// Note that this method is also extensively tested by the python unit test:
+// tensorflow/python/kernel_tests/parsing_ops_test.py
+class SingleExampleProtoToTensorsTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    // Setup dense feature configuration.
+    FixedLenFeature int64_dense_config;
+    int64_dense_config.key = kDenseInt64Key;
+    int64_dense_config.dtype = DT_INT64;
+    int64_dense_config.shape = TensorShape({1});
+    int64_dense_config.default_value = Tensor(DT_INT64, TensorShape({1}));
+    int64_dense_config.default_value.scalar<int64_t>()() = 0;
+    dense_vec_.push_back(int64_dense_config);
+
+    FixedLenFeature float_dense_config;
+    float_dense_config.key = kDenseFloatKey;
+    float_dense_config.dtype = DT_FLOAT;
+    float_dense_config.shape = TensorShape({1});
+    float_dense_config.default_value = Tensor(DT_FLOAT, TensorShape({1}));
+    float_dense_config.default_value.scalar<float>()() = 0.0;
+    dense_vec_.push_back(float_dense_config);
+
+    FixedLenFeature string_dense_config;
+    string_dense_config.key = kDenseStringKey;
+    string_dense_config.dtype = DT_STRING;
+    string_dense_config.shape = TensorShape({1});
+    string_dense_config.default_value = Tensor(DT_STRING, TensorShape({1}));
+    string_dense_config.default_value.scalar<tstring>()() = "default";
+    dense_vec_.push_back(string_dense_config);
+
+    // Setup sparse feature configuration.
+    VarLenFeature int64_sparse_config;
+    int64_sparse_config.key = kSparseInt64Key;
+    int64_sparse_config.dtype = DT_INT64;
+    sparse_vec_.push_back(int64_sparse_config);
+
+    VarLenFeature float_sparse_config;
+    float_sparse_config.key = kSparseFloatKey;
+    float_sparse_config.dtype = DT_FLOAT;
+    sparse_vec_.push_back(float_sparse_config);
+
+    VarLenFeature string_sparse_config;
+    string_sparse_config.key = kSparseStringKey;
+    string_sparse_config.dtype = DT_STRING;
+    sparse_vec_.push_back(string_sparse_config);
+  }
+
+  std::vector<FixedLenFeature> dense_vec_;
+  std::vector<VarLenFeature> sparse_vec_;
+};
+
+TEST_F(SingleExampleProtoToTensorsTest, SparseOnlyTrivial) {
+  Example ex;
+  // Set up a feature for each of our supported types.
+  (*ex.mutable_features()->mutable_feature())[kSparseInt64Key]
+      .mutable_int64_list()
+      ->add_value(42);
+  (*ex.mutable_features()->mutable_feature())[kSparseFloatKey]
+      .mutable_float_list()
+      ->add_value(4.2);
+  (*ex.mutable_features()->mutable_feature())[kSparseStringKey]
+      .mutable_bytes_list()
+      ->add_value("forty-two");
+
+  std::vector<Tensor*> output_dense_values(0);
+  std::vector<std::vector<Tensor>> output_sparse_values_tmp(3);
+  for (int i = 0; i < 3; ++i) {
+    output_sparse_values_tmp[i] = std::vector<Tensor>(1);
+  }
+
+  std::vector<FixedLenFeature> empty_dense_vec;
+  TF_EXPECT_OK(SingleExampleProtoToTensors(ex, "", 0, empty_dense_vec,
+                                           sparse_vec_, &output_dense_values,
+                                           &output_sparse_values_tmp));
+
+  const std::vector<Tensor>& int64_tensor_vec = output_sparse_values_tmp[0];
+  EXPECT_EQ(1, int64_tensor_vec.size());
+  EXPECT_EQ(42, int64_tensor_vec[0].vec<int64_t>()(0));
+
+  const std::vector<Tensor>& float_tensor_vec = output_sparse_values_tmp[1];
+  EXPECT_EQ(1, float_tensor_vec.size());
+  EXPECT_NEAR(4.2, float_tensor_vec[0].vec<float>()(0), 0.001);
+
+  const std::vector<Tensor>& string_tensor_vec = output_sparse_values_tmp[2];
+  EXPECT_EQ(1, string_tensor_vec.size());
+  EXPECT_EQ("forty-two", string_tensor_vec[0].vec<tstring>()(0));
+}
+
+TEST_F(SingleExampleProtoToTensorsTest, SparseOnlyEmpty) {
+  Example empty;
+  std::vector<Tensor*> output_dense_values(0);
+  std::vector<std::vector<Tensor>> output_sparse_values_tmp(3);
+  for (int i = 0; i < 3; ++i) {
+    output_sparse_values_tmp[i] = std::vector<Tensor>(1);
+  }
+
+  std::vector<FixedLenFeature> empty_dense_vec;
+  TF_EXPECT_OK(SingleExampleProtoToTensors(empty, "", 0, empty_dense_vec,
+                                           sparse_vec_, &output_dense_values,
+                                           &output_sparse_values_tmp));
+
+  // Each feature will still have a tensor vector, however the tensor
+  // in the vector will be empty.
+  const std::vector<Tensor>& int64_tensor_vec = output_sparse_values_tmp[0];
+  EXPECT_EQ(1, int64_tensor_vec.size());
+  EXPECT_EQ(0, int64_tensor_vec[0].vec<int64_t>().size());
+
+  const std::vector<Tensor>& float_tensor_vec = output_sparse_values_tmp[1];
+  EXPECT_EQ(1, float_tensor_vec.size());
+  EXPECT_EQ(0, float_tensor_vec[0].vec<float>().size());
+
+  const std::vector<Tensor>& string_tensor_vec = output_sparse_values_tmp[2];
+  EXPECT_EQ(1, string_tensor_vec.size());
+  EXPECT_EQ(0, string_tensor_vec[0].vec<tstring>().size());
+}
+
+TEST_F(SingleExampleProtoToTensorsTest, DenseOnlyTrivial) {
+  Example ex;
+  // Set up a feature for each of our supported types.
+  (*ex.mutable_features()->mutable_feature())[kDenseInt64Key]
+      .mutable_int64_list()
+      ->add_value(42);
+  (*ex.mutable_features()->mutable_feature())[kDenseFloatKey]
+      .mutable_float_list()
+      ->add_value(4.2);
+  (*ex.mutable_features()->mutable_feature())[kDenseStringKey]
+      .mutable_bytes_list()
+      ->add_value("forty-two");
+
+  std::vector<Tensor*> output_dense_values(3);
+  Tensor int64_dense_output(DT_INT64, TensorShape({1, 1}));
+  output_dense_values[0] = &int64_dense_output;
+
+  Tensor float_dense_output(DT_FLOAT, TensorShape({1, 1}));
+  output_dense_values[1] = &float_dense_output;
+
+  Tensor str_dense_output(DT_STRING, TensorShape({1, 1}));
+  output_dense_values[2] = &str_dense_output;
+
+  std::vector<VarLenFeature> empty_sparse_vec;
+  std::vector<std::vector<Tensor>> output_sparse_values_tmp;
+  TF_EXPECT_OK(SingleExampleProtoToTensors(
+      ex, "", 0, dense_vec_, empty_sparse_vec, &output_dense_values,
+      &output_sparse_values_tmp));
+  EXPECT_TRUE(output_sparse_values_tmp.empty());
+
+  EXPECT_EQ(1, int64_dense_output.matrix<int64_t>().size());
+  EXPECT_EQ(42, int64_dense_output.matrix<int64_t>()(0, 0));
+
+  EXPECT_EQ(1, float_dense_output.matrix<float>().size());
+  EXPECT_NEAR(4.2, float_dense_output.matrix<float>()(0, 0), 0.001);
+
+  EXPECT_EQ(1, str_dense_output.matrix<tstring>().size());
+  EXPECT_EQ("forty-two", str_dense_output.matrix<tstring>()(0, 0));
+}
+
+TEST_F(SingleExampleProtoToTensorsTest, DenseOnlyDefaults) {
+  std::vector<Tensor*> output_dense_values(3);
+  Tensor int64_dense_output(DT_INT64, TensorShape({1, 1}));
+  output_dense_values[0] = &int64_dense_output;
+
+  Tensor float_dense_output(DT_FLOAT, TensorShape({1, 1}));
+  output_dense_values[1] = &float_dense_output;
+
+  Tensor str_dense_output(DT_STRING, TensorShape({1, 1}));
+  output_dense_values[2] = &str_dense_output;
+
+  Example empty;
+
+  std::vector<VarLenFeature> empty_sparse_vec;
+  std::vector<std::vector<Tensor>> output_sparse_values_tmp;
+  TF_EXPECT_OK(SingleExampleProtoToTensors(
+      empty, "", 0, dense_vec_, empty_sparse_vec, &output_dense_values,
+      &output_sparse_values_tmp));
+
+  EXPECT_EQ(1, int64_dense_output.matrix<int64_t>().size());
+  EXPECT_EQ(0, int64_dense_output.matrix<int64_t>()(0, 0));
+
+  EXPECT_EQ(1, float_dense_output.matrix<float>().size());
+  EXPECT_NEAR(0.0, float_dense_output.matrix<float>()(0, 0), 0.001);
+
+  EXPECT_EQ(1, str_dense_output.matrix<tstring>().size());
+  EXPECT_EQ("default", str_dense_output.matrix<tstring>()(0, 0));
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/core/util/exec_on_stall.h

@@ -0,0 +1,89 @@
+/* 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_UTIL_EXEC_ON_STALL_H_
+#define TENSORFLOW_CORE_UTIL_EXEC_ON_STALL_H_
+
+#include <functional>
+
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/mutex.h"
+
+namespace tensorflow {
+
+// An object that executes a particular function only if it
+// is not deleted within the allotted number of seconds.
+//
+// This can be useful in diagnosing deadlocks, stalls and memory leaks
+// without logging too aggressively.
+class ExecuteOnStall {
+ public:
+  // delay_secs: If the object still exists after this many seconds,
+  //     execute f.
+  // f: The function to be executed, for example a detailed log of the
+  //    the state of an object to which this is attached.
+  // poll_microseconds: The spawned thread will wake and test whether
+  //    the destructor has been invoked this frequently.
+  ExecuteOnStall(int delay_secs, std::function<void()> f,
+                 int32_t poll_microseconds = 100)
+      : disabled_(false),
+        joined_(false),
+        env_(Env::Default()),
+        f_(f),
+        poll_microseconds_(poll_microseconds) {
+    deadline_ = env_->NowMicros() + 1000000 * delay_secs;
+    env_->SchedClosure([this]() {
+      while (env_->NowMicros() < deadline_) {
+        {
+          mutex_lock l(mu_);
+          if (disabled_) {
+            break;
+          }
+        }
+        env_->SleepForMicroseconds(poll_microseconds_);
+      }
+      {
+        mutex_lock l(mu_);
+        if (!disabled_) {
+          f_();
+        }
+        joined_ = true;
+        cond_var_.notify_all();
+      }
+    });
+  }
+
+  ~ExecuteOnStall() {
+    // Wait for spawned thread to terminate.
+    mutex_lock l(mu_);
+    disabled_ = true;
+    if (!joined_) {
+      cond_var_.wait(l);
+    }
+  }
+
+ private:
+  mutex mu_;
+  condition_variable cond_var_;
+  bool disabled_ TF_GUARDED_BY(mu_);
+  bool joined_ TF_GUARDED_BY(mu_);
+  Env* env_;
+  std::function<void()> f_;
+  int64_t deadline_;
+  int32_t poll_microseconds_;
+};
+
+}  // namespace tensorflow
+#endif  // TENSORFLOW_CORE_UTIL_EXEC_ON_STALL_H_

tensorflow/core/util/exec_on_stall_test.cc

@@ -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.
+==============================================================================*/
+
+#include "tensorflow/core/util/exec_on_stall.h"
+
+#include <functional>
+#include <memory>
+#include <utility>
+
+#include "tensorflow/core/platform/macros.h"
+#include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace {
+
+struct Chunk {
+  std::unique_ptr<ExecuteOnStall> stall_closure;
+};
+
+Chunk* NewChunk(int stall_seconds, std::function<void()> f) {
+  Chunk* c = new Chunk;
+  c->stall_closure =
+      std::make_unique<ExecuteOnStall>(stall_seconds, std::move(f));
+  return c;
+}
+
+TEST(ExecuteOnStallTest, BothWays) {
+  mutex mu;
+  bool a_triggered(false);
+  bool b_triggered(false);
+  Chunk* a = NewChunk(1, [&mu, &a_triggered]() {
+    mutex_lock l(mu);
+    a_triggered = true;
+  });
+  Chunk* b = NewChunk(1, [&mu, &b_triggered]() {
+    mutex_lock l(mu);
+    b_triggered = true;
+  });
+  delete a;
+  Env::Default()->SleepForMicroseconds(2000000);
+  {
+    mutex_lock l(mu);
+    EXPECT_FALSE(a_triggered);
+    EXPECT_TRUE(b_triggered);
+  }
+  delete b;
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/core/util/fake_clock_env.cc

@@ -0,0 +1,38 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/util/fake_clock_env.h"
+
+#include <string>
+
+namespace tensorflow {
+
+FakeClockEnv::FakeClockEnv(Env* wrapped) : EnvWrapper(wrapped) {}
+
+void FakeClockEnv::AdvanceByMicroseconds(int64_t micros) {
+  {
+    mutex_lock l(mu_);
+    current_time_ += micros;
+  }
+}
+
+uint64_t FakeClockEnv::NowMicros() const {
+  {
+    mutex_lock l(mu_);
+    return current_time_;
+  }
+}
+
+}  // namespace tensorflow

tensorflow/core/util/fake_clock_env.h

@@ -0,0 +1,58 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_UTIL_FAKE_CLOCK_ENV_H_
+#define TENSORFLOW_CORE_UTIL_FAKE_CLOCK_ENV_H_
+
+#include <functional>
+#include <string>
+#include <vector>
+
+#include "tensorflow/core/lib/core/notification.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/macros.h"
+#include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/platform/thread_annotations.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+
+// An Env implementation with a fake clock for NowMicros().
+// The clock doesn't advance on its own. It advances
+// via an explicit AdvanceByMicroseconds() method. All other Env virtual methods
+// pass through to a wrapped Env.
+class FakeClockEnv : public EnvWrapper {
+ public:
+  explicit FakeClockEnv(Env* wrapped);
+  ~FakeClockEnv() override = default;
+
+  // Advance the clock by a certain number of microseconds.
+  void AdvanceByMicroseconds(int64_t micros);
+
+  // Returns the current time of FakeClockEnv in microseconds.
+  uint64_t NowMicros() const override;
+
+ private:
+  mutable mutex mu_;
+  uint64_t current_time_ TF_GUARDED_BY(mu_) = 0;
+
+  FakeClockEnv(const FakeClockEnv&) = delete;
+  void operator=(const FakeClockEnv&) = delete;
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_UTIL_FAKE_CLOCK_ENV_H_

tensorflow/core/util/fake_clock_env_test.cc

@@ -0,0 +1,65 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/util/fake_clock_env.h"
+
+#include <memory>
+
+#include <gtest/gtest.h>
+#include "tensorflow/core/platform/env.h"
+
+namespace tensorflow {
+namespace {
+
+class FakeClockEnvTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    fake_clock_env_ = std::make_unique<FakeClockEnv>(Env::Default());
+  }
+
+  void TearDown() override { fake_clock_env_.reset(); }
+
+  std::unique_ptr<FakeClockEnv> fake_clock_env_;
+};
+
+TEST_F(FakeClockEnvTest, TimeInitializedToZero) {
+  EXPECT_EQ(0, fake_clock_env_->NowMicros());
+}
+
+TEST_F(FakeClockEnvTest, AdvanceTimeByMicroseconds) {
+  int current_time = fake_clock_env_->NowMicros();
+
+  // Advance current time and fake clock by equal duration.
+  int64_t duration = 100;
+  current_time += duration;
+  fake_clock_env_->AdvanceByMicroseconds(duration);
+  EXPECT_EQ(current_time, fake_clock_env_->NowMicros());
+
+  // Multiple advancements of current time and fake clock.
+  for (int i = 0; i < 5; ++i) {
+    fake_clock_env_->AdvanceByMicroseconds(100);
+    current_time += 100;
+  }
+  EXPECT_EQ(current_time, fake_clock_env_->NowMicros());
+
+  // Advance current time and fake clock by unequal durations.
+  current_time += duration;
+  duration = 200;
+  fake_clock_env_->AdvanceByMicroseconds(duration);
+  EXPECT_NE(current_time, fake_clock_env_->NowMicros());
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/core/util/gpu_cuda_alias.h

@@ -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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_UTIL_GPU_CUDA_ALIAS_H_
+#define TENSORFLOW_CORE_UTIL_GPU_CUDA_ALIAS_H_
+
+// Several forwarding macros are defined in this file to serve for backward
+// compatibility usage as we migrating from CUDA prefixed function to GPU
+// prefixed functions. Both Cuda and ROCm can unify under the new GPU prefix
+// naming scheme. In the migration period, we provide equivalent CUDA* and GPU*
+// function. Over time, all CUDA* functions will be deprecated.
+
+namespace tensorflow {
+
+// CREATE_CUDA_HOST_FUNCTION_ALIAS forward the host function to its CUDA Alias.
+#ifndef TENSORFLOW_USE_ROCM
+#define CREATE_CUDA_HOST_FUNCTION_ALIAS(func, cuda_alias) \
+  template <typename... Args>                             \
+  auto cuda_alias(Args&&... args)                         \
+      ->decltype(func(std::forward<Args>(args)...)) {     \
+    return func(std::forward<Args>(args)...);             \
+  }
+#else
+#define CREATE_CUDA_HOST_FUNCTION_ALIAS(func, cuda_alias)
+#endif
+
+// CREATE_CUDA_DEVICE_FUNCTION_ALIAS forward the device function to its CUDA
+// Alias.
+#ifndef TENSORFLOW_USE_ROCM
+#define CREATE_CUDA_DEVICE_FUNCTION_ALIAS(func, cuda_alias) \
+  template <typename... Args>                               \
+  __device__ auto cuda_alias(Args&&... args)                \
+      ->decltype(func(std::forward<Args>(args)...)) {       \
+    return func(std::forward<Args>(args)...);               \
+  }
+#else
+#define CREATE_CUDA_DEVICE_FUNCTION_ALIAS(func, cuda_alias)
+#endif
+
+// CREATE_CUDA_TYPE_ALIAS forward the type to its CUDA Alias.
+#ifndef TENSORFLOW_USE_ROCM
+#define CREATE_CUDA_TYPE_ALIAS(type, cuda_alias) using cuda_alias = type;
+#else
+#define CREATE_CUDA_TYPE_ALIAS(type, cuda_alias)
+#endif
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_UTIL_GPU_CUDA_ALIAS_H_

tensorflow/core/util/gpu_device_functions.h

@@ -0,0 +1,995 @@
+/* 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_UTIL_GPU_DEVICE_FUNCTIONS_H_
+#define TENSORFLOW_CORE_UTIL_GPU_DEVICE_FUNCTIONS_H_
+
+/**
+ * Wrappers and helpers for CUDA device code.
+ *
+ * Wraps the warp-cooperative intrinsics introduced in CUDA 9 to provide
+ * backwards compatibility, see go/volta-porting for details.
+ * Provides atomic operations on types that aren't natively supported.
+ * Defines a number of macros and types providing a shared interface
+ * to either CUDA or ROCm APIs, depending on the build.
+ */
+
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+
+#include <algorithm>
+#include <complex>
+
+#include "unsupported/Eigen/CXX11/Tensor"  // from @eigen_archive
+#if GOOGLE_CUDA
+#include "third_party/gpus/cuda/include/cuda.h"
+#else
+#include "rocm/include/hip/hip_bf16.h"
+#include "rocm/include/hip/hip_complex.h"
+#include "rocm/include/hip/hip_fp16.h"
+#endif
+
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/util/gpu_cuda_alias.h"
+
+#if GOOGLE_CUDA
+using gpuStream_t = cudaStream_t;
+using gpuEvent_t = cudaEvent_t;
+#define gpuEventRecord cudaEventRecord
+#define gpuEventSynchronize cudaEventSynchronize
+#define gpuEventDestroy cudaEventDestroy
+#define gpuEventCreate cudaEventCreate
+#define gpuEventCreateWithFlags cudaEventCreateWithFlags
+#define gpuEventDisableTiming cudaEventDisableTiming
+#define gpuDeviceSynchronize cudaDeviceSynchronize
+#define gpuFree cudaFree
+#elif TENSORFLOW_USE_ROCM
+using gpuStream_t = hipStream_t;
+using gpuEvent_t = hipEvent_t;
+using cudaError = int;
+using cudaError_t = int;
+#define cudaSuccess 0
+#define cudaGetLastError hipGetLastError
+#define gpuEventRecord hipEventRecord
+#define gpuEventDestroy hipEventDestroy
+#define gpuEventSynchronize hipEventSynchronize
+#define gpuEventCreate hipEventCreate
+#define gpuEventCreateWithFlags hipEventCreateWithFlags
+#define gpuEventDisableTiming hipEventDisableTiming
+#define gpuDeviceSynchronize hipDeviceSynchronize
+#define gpuFree hipFree
+static std::string cudaGetErrorString(int err) { return std::to_string(err); }
+#endif
+
+#define TF_RETURN_IF_CUDA_ERROR(result)                                       \
+  do {                                                                        \
+    cudaError_t error(result);                                                \
+    if (!TF_PREDICT_TRUE(error == cudaSuccess)) {                             \
+      return absl::InternalError(                                             \
+          absl::StrCat("Cuda call failed with ", cudaGetErrorString(error))); \
+    }                                                                         \
+  } while (0)
+
+#define TF_OP_REQUIRES_CUDA_SUCCESS(context, result)                          \
+  do {                                                                        \
+    cudaError_t error(result);                                                \
+    if (!TF_PREDICT_TRUE(error == cudaSuccess)) {                             \
+      context->SetStatus(absl::InternalError(                                 \
+          absl::StrCat("Cuda call failed with", cudaGetErrorString(error)))); \
+      return;                                                                 \
+    }                                                                         \
+  } while (0)
+
+namespace tensorflow {
+// According to HIP developer guide at
+// https://github.com/ROCm-Developer-Tools/HIP/blob/master/docs/markdown/hip_kernel_language.md#assert
+// assert is not supported by HIP. While we are waiting for assert support in
+// hip kernels, the assert call should be macroed to NOP so that it does not
+// block us from creating a debug build
+#if TENSORFLOW_USE_ROCM
+#undef assert
+#define assert(x) \
+  {}
+#endif
+
+namespace detail {
+
+// Helper for range-based for loop using 'delta' increments.
+// Usage: see GpuGridRange?() functions below.
+template <typename T>
+class GpuGridRange {
+  struct Iterator {
+    __device__ Iterator(T index, T delta) : index_(index), delta_(delta) {}
+    __device__ T operator*() const { return index_; }
+    __device__ Iterator& operator++() {
+      index_ += delta_;
+      return *this;
+    }
+    __device__ bool operator!=(const Iterator& other) const {
+      bool greater = index_ > other.index_;
+      bool less = index_ < other.index_;
+      // Anything past an end iterator (delta_ == 0) is equal.
+      // In range-based for loops, this optimizes to 'return less'.
+      if (!other.delta_) {
+        return less;
+      }
+      if (!delta_) {
+        return greater;
+      }
+      return less || greater;
+    }
+
+   private:
+    T index_;
+    const T delta_;
+  };
+
+ public:
+  __device__ GpuGridRange(T begin, T delta, T end)
+      : begin_(begin), delta_(delta), end_(end) {}
+
+  __device__ Iterator begin() const { return Iterator{begin_, delta_}; }
+  __device__ Iterator end() const { return Iterator{end_, 0}; }
+
+ private:
+  T begin_;
+  T delta_;
+  T end_;
+};
+
+#ifndef TENSORFLOW_USE_ROCM
+template <typename... T>
+using CudaGridRange = GpuGridRange<T...>;
+#endif
+}  // namespace detail
+
+// Helper to visit indices in the range 0 <= i < count, using the x-coordinate
+// of the global thread index. That is, each index i is visited by all threads
+// with the same x-coordinate.
+// Usage: for(int i : GpuGridRangeX(count)) { visit(i); }
+template <typename T>
+__device__ detail::GpuGridRange<T> GpuGridRangeX(T count) {
+  return detail::GpuGridRange<T>(
+      /*begin=*/static_cast<T>(blockIdx.x) * blockDim.x + threadIdx.x,
+      /*delta=*/static_cast<T>(gridDim.x) * blockDim.x, /*end=*/count);
+}
+CREATE_CUDA_DEVICE_FUNCTION_ALIAS(GpuGridRangeX, CudaGridRangeX);
+
+// Helper to visit indices in the range 0 <= i < count using the y-coordinate.
+// Usage: for(int i : GpuGridRangeY(count)) { visit(i); }
+template <typename T>
+__device__ detail::GpuGridRange<T> GpuGridRangeY(T count) {
+  return detail::GpuGridRange<T>(
+      /*begin=*/static_cast<T>(blockIdx.y) * blockDim.y + threadIdx.y,
+      /*delta=*/static_cast<T>(gridDim.y) * blockDim.y, /*end=*/count);
+}
+CREATE_CUDA_DEVICE_FUNCTION_ALIAS(GpuGridRangeY, CudaGridRangeY);
+
+// Helper to visit indices in the range 0 <= i < count using the z-coordinate.
+// Usage: for(int i : GpuGridRangeZ(count)) { visit(i); }
+template <typename T>
+__device__ detail::GpuGridRange<T> GpuGridRangeZ(T count) {
+  return detail::GpuGridRange<T>(
+      /*begin=*/static_cast<T>(blockIdx.z) * blockDim.z + threadIdx.z,
+      /*delta=*/static_cast<T>(gridDim.z) * blockDim.z, /*end=*/count);
+}
+CREATE_CUDA_DEVICE_FUNCTION_ALIAS(GpuGridRangeZ, CudaGridRangeZ);
+
+// Mask for all 32 threads in a warp.
+__device__ const unsigned kCudaWarpAll = 0xffffffff;
+// ROCM TODO add ROCM implementation
+// Mask for all 64 threads in a wavefront.
+__device__ const unsigned kGpuWarpAll = 0xffffffff;
+
+// Returns the warp lane ID of the calling thread
+__device__ inline unsigned GpuLaneId() {
+  unsigned int lane_id;
+#if GOOGLE_CUDA
+#if __clang__ && !__NVCC__
+  return __nvvm_read_ptx_sreg_laneid();
+#else   // __clang__
+  asm("mov.u32 %0, %%laneid;" : "=r"(lane_id));
+#endif  // __clang__
+#elif TENSORFLOW_USE_ROCM
+  lane_id = __lane_id();
+#endif
+  return lane_id;
+}
+CREATE_CUDA_DEVICE_FUNCTION_ALIAS(GpuLaneId, CudaLaneId);
+
+namespace detail {
+// Returns true if mask is a valid parameter for __shfl*sync to return a well
+// defined value, assuming the calling lane will read from src_lane as part of
+// the shuffle operation.
+//
+// Specifically, returns true iff mask has the calling lane bit and the src_lane
+// bit set, and the src_lane calls this function with the same mask value
+// (required for the two threads to wait for each other).
+//
+// On Volta, for some invalid masks, this function hangs or returns false
+// positives, because the implementation shuffles with the same mask that
+// we are validating. Run on Pascal if you suspect that the mask is incorrect.
+__device__ inline bool GpuValidateShuffleSyncMask(unsigned mask,
+                                                  unsigned src_lane) {
+  unsigned src_dst_mask = 1u << GpuLaneId() | 1u << src_lane;
+#if GOOGLE_CUDA
+  unsigned src_lane_mask = __shfl_sync(mask, mask, src_lane);
+#else  // TENSORFLOW_USE_ROCM
+  unsigned src_lane_mask =
+      __shfl(static_cast<int>(mask), static_cast<int>(src_lane));
+#endif
+  return (src_dst_mask & ~mask) == 0 && src_lane_mask == mask;
+}
+CREATE_CUDA_DEVICE_FUNCTION_ALIAS(GpuValidateShuffleSyncMask,
+                                  CudaValidateShuffleSyncMask);
+
+// Returns the actual source lane for shuffle.
+__device__ inline unsigned GpuShuffleGetSrcLane(int src_lane, int width) {
+  int lane_id = GpuLaneId();
+  int lane_base = lane_id & ~width + 1;
+  int lane_offset = src_lane & width - 1;
+  return lane_base + lane_offset;
+}
+CREATE_CUDA_DEVICE_FUNCTION_ALIAS(GpuShuffleGetSrcLane, CudaShuffleGetSrcLane);
+
+// Returns the source lane for shuffle up.
+__device__ inline unsigned GpuShuffleUpGetSrcLane(unsigned delta, int width) {
+  unsigned lane_id = GpuLaneId();
+  if ((lane_id & width - 1) < delta) {
+    return lane_id;
+  }
+  return lane_id - delta;
+}
+CREATE_CUDA_DEVICE_FUNCTION_ALIAS(GpuShuffleUpGetSrcLane,
+                                  CudaShuffleUpGetSrcLane);
+
+// Returns the source lane for shuffle down.
+__device__ inline unsigned GpuShuffleDownGetSrcLane(unsigned delta, int width) {
+  unsigned lane_id = GpuLaneId();
+  if ((lane_id & width - 1) + delta >= width) {
+    return lane_id;
+  }
+  return lane_id + delta;
+}
+CREATE_CUDA_DEVICE_FUNCTION_ALIAS(GpuShuffleDownGetSrcLane,
+                                  CudaShuffleDownGetSrcLane);
+
+// Returns the source lane for shuffle xor.
+__device__ inline unsigned GpuShuffleXorGetSrcLane(int lane_mask, int width) {
+  int lane_id = GpuLaneId();
+  int src_lane = lane_id ^ lane_mask;
+  if (src_lane > (lane_id | width - 1)) {
+    return lane_id;
+  }
+  return src_lane;
+}
+CREATE_CUDA_DEVICE_FUNCTION_ALIAS(GpuShuffleXorGetSrcLane,
+                                  CudaShuffleXorGetSrcLane);
+}  // namespace detail
+
+// For all *_sync wrappers below, it is illegal to synchronize threads from
+// different program locations, because that is not supported before sm_70.
+// In other words, all threads in 'mask' must call the functions in convergence.
+// Code that requires sm_70 (and CUDA 9) may use the intrinsic directly.
+//
+// It is also illegal to shuffle with a mask that produces an undefined result
+// for any of the threads. Specifically, all source threads of the shuffle
+// must have their corresponding bit in 'mask' set.
+
+// Wrapper for __syncwarp. No-op for CUDA 8 and earlier.
+__device__ inline void GpuSyncWarp(unsigned mask = kCudaWarpAll) {
+  assert(mask & 1u << GpuLaneId());
+#if GOOGLE_CUDA
+  __syncwarp(mask);
+#endif
+}
+CREATE_CUDA_DEVICE_FUNCTION_ALIAS(GpuSyncWarp, CudaSyncWarp);
+
+// Wrapper for __ballot_sync. All threads in 'mask' must call this function in
+// convergence, see comment above for details.
+__device__ inline unsigned GpuBallotSync(unsigned mask, int pred) {
+  assert(mask & 1u << GpuLaneId());
+#if GOOGLE_CUDA
+  return __ballot_sync(mask, pred);
+#else  // TENSORFLOW_USE_ROCM
+  return __ballot(pred) & mask;  // Apply mask to match __ballot_sync's spec.
+#endif
+}
+CREATE_CUDA_DEVICE_FUNCTION_ALIAS(GpuBallotSync, CudaBallotSync);
+
+// Wrapper for __any_sync. All threads in 'mask' must call this function in
+// convergence, see comment above for details.
+__device__ inline int GpuAnySync(unsigned mask, int pred) {
+  assert(mask & 1u << GpuLaneId());
+#if GOOGLE_CUDA
+  return __any_sync(mask, pred);
+#else  // TENSORFLOW_USE_ROCM
+  return __any(pred);
+#endif
+}
+CREATE_CUDA_DEVICE_FUNCTION_ALIAS(GpuAnySync, CudaAnySync);
+
+// Wrapper for __all_sync. All threads in 'mask' must call this function in
+// convergence, see comment above for details.
+__device__ inline int GpuAllSync(unsigned mask, int pred) {
+  assert(mask & 1u << GpuLaneId());
+#if GOOGLE_CUDA
+  return __all_sync(mask, pred);
+#else  // TENSORFLOW_USE_ROCM
+  return __all(pred);
+#endif
+}
+CREATE_CUDA_DEVICE_FUNCTION_ALIAS(GpuAllSync, CudaAllSync);
+
+// Wrapper for __shfl_sync. All threads in 'mask' must call this function in
+// convergence, see comment above for details.
+template <typename T>
+__device__ T GpuShuffleSync(unsigned mask, T value, int src_lane,
+                            int width = warpSize) {
+  assert(!(width & width - 1));
+  assert(detail::GpuValidateShuffleSyncMask(
+      mask, detail::GpuShuffleGetSrcLane(src_lane, width)));
+#if GOOGLE_CUDA
+  return __shfl_sync(mask, value, src_lane, width);
+#else  // TENSORFLOW_USE_ROCM
+  return __shfl(value, src_lane, width);
+#endif
+}
+
+// Variant of the (undocumented) version from the CUDA SDK, but using unsigned
+// instead of float for lo and hi (which is incorrect with ftz, for example).
+// See b/69446944.
+__device__ inline double GpuShuffleSync(unsigned mask, double value,
+                                        int src_lane, int width = warpSize) {
+  auto tmp = __double_as_longlong(value);
+  auto lo = static_cast<unsigned>(tmp);
+  auto hi = static_cast<unsigned>(tmp >> 32);
+  hi = GpuShuffleSync(mask, hi, src_lane, width);
+  lo = GpuShuffleSync(mask, lo, src_lane, width);
+  return __longlong_as_double(static_cast<uint64_t>(hi) << 32 | lo);
+}
+CREATE_CUDA_DEVICE_FUNCTION_ALIAS(GpuShuffleSync, CudaShuffleSync);
+
+// Wrapper for __shfl_up_sync. All threads in 'mask' must call this function in
+// convergence, see comment above for details.
+template <typename T>
+__device__ inline T GpuShuffleUpSync(unsigned mask, T value, unsigned delta,
+                                     int width = warpSize) {
+  assert(!(width & width - 1));
+  assert(detail::GpuValidateShuffleSyncMask(
+      mask, detail::GpuShuffleUpGetSrcLane(delta, width)));
+#if GOOGLE_CUDA
+  return __shfl_up_sync(mask, value, delta, width);
+#else  // TENSORFLOW_USE_ROCM
+  return __shfl_up(value, delta, width);
+#endif
+}
+
+// Variant of the (undocumented) version from the CUDA SDK, but using unsigned
+// instead of float for lo and hi (which is incorrect with ftz, for example).
+// See b/69446944.
+__device__ inline double GpuShuffleUpSync(unsigned mask, double value,
+                                          unsigned delta,
+                                          int width = warpSize) {
+  auto tmp = __double_as_longlong(value);
+  auto lo = static_cast<unsigned>(tmp);
+  auto hi = static_cast<unsigned>(tmp >> 32);
+  hi = GpuShuffleUpSync(mask, hi, delta, width);
+  lo = GpuShuffleUpSync(mask, lo, delta, width);
+  return __longlong_as_double(static_cast<uint64_t>(hi) << 32 | lo);
+}
+CREATE_CUDA_DEVICE_FUNCTION_ALIAS(GpuShuffleUpSync, CudaShuffleUpSync);
+
+// Wrapper for __shfl_down_sync. All threads in 'mask' must call this function
+// in convergence, see comment above for details.
+template <typename T>
+__device__ inline T GpuShuffleDownSync(unsigned mask, T value, unsigned delta,
+                                       int width = warpSize) {
+  assert(!(width & width - 1));
+  assert(detail::GpuValidateShuffleSyncMask(
+      mask, detail::GpuShuffleDownGetSrcLane(delta, width)));
+#if GOOGLE_CUDA
+  return __shfl_down_sync(mask, value, delta, width);
+#else  // TENSORFLOW_USE_ROCM
+  return __shfl_down(value, delta, width);
+#endif
+}
+
+// Variant of the (undocumented) version from the CUDA SDK, but using unsigned
+// instead of float for lo and hi (which is incorrect with ftz, for example).
+// See b/69446944.
+__device__ inline double GpuShuffleDownSync(unsigned mask, double value,
+                                            unsigned delta,
+                                            int width = warpSize) {
+  auto tmp = __double_as_longlong(value);
+  auto lo = static_cast<unsigned>(tmp);
+  auto hi = static_cast<unsigned>(tmp >> 32);
+  hi = GpuShuffleDownSync(mask, hi, delta, width);
+  lo = GpuShuffleDownSync(mask, lo, delta, width);
+  return __longlong_as_double(static_cast<uint64_t>(hi) << 32 | lo);
+}
+CREATE_CUDA_DEVICE_FUNCTION_ALIAS(GpuShuffleDownSync, CudaShuffleDownSync);
+
+// Wrapper for __shfl_xor_sync. All threads in 'mask' must call this function in
+// convergence, see comment above for details.
+template <typename T>
+__device__ T GpuShuffleXorSync(unsigned mask, T value, int lane_mask,
+                               int width = warpSize) {
+  assert(!(width & width - 1));
+  assert(detail::GpuValidateShuffleSyncMask(
+      mask, detail::GpuShuffleXorGetSrcLane(lane_mask, width)));
+#if GOOGLE_CUDA
+  return __shfl_xor_sync(mask, value, lane_mask, width);
+#elif TENSORFLOW_USE_ROCM
+  // ROCM TODO: check if HIP should be changed to cope with more types
+  return __shfl_xor(static_cast<int>(value), lane_mask, width);
+#endif
+}
+
+// Variant of the (undocumented) version from the CUDA SDK, but using unsigned
+// instead of float for lo and hi (which is incorrect with ftz, for example).
+// See b/69446944.
+__device__ inline double GpuShuffleXorSync(unsigned mask, double value,
+                                           int lane_mask,
+                                           int width = warpSize) {
+  auto tmp = __double_as_longlong(value);
+  auto lo = static_cast<unsigned>(tmp);
+  auto hi = static_cast<unsigned>(tmp >> 32);
+  hi = GpuShuffleXorSync(mask, hi, lane_mask, width);
+  lo = GpuShuffleXorSync(mask, lo, lane_mask, width);
+  return __longlong_as_double(static_cast<uint64_t>(hi) << 32 | lo);
+}
+CREATE_CUDA_DEVICE_FUNCTION_ALIAS(GpuShuffleXorSync, CudaShuffleXorSync);
+
+// Wrapper for __ldg.
+template <typename T>
+__host__ __device__ T GpuLdg(const T* address) {
+#if __CUDA_ARCH__ >= 350
+  return __ldg(address);
+#else
+  return *address;
+#endif
+}
+
+__host__ __device__ inline bool GpuLdg(const bool* address) {
+  return GpuLdg(reinterpret_cast<const char*>(address)) != 0;
+}
+
+__host__ __device__ inline std::complex<float> GpuLdg(
+    const std::complex<float>* address) {
+#if __CUDA_ARCH__ >= 350
+  float2 mem = __ldg(reinterpret_cast<const float2*>(address));
+  return std::complex<float>(mem.x, mem.y);
+#else
+  return *address;
+#endif
+}
+
+__host__ __device__ inline std::complex<double> GpuLdg(
+    const std::complex<double>* address) {
+#if __CUDA_ARCH__ >= 350
+  double2 mem = __ldg(reinterpret_cast<const double2*>(address));
+  return std::complex<double>(mem.x, mem.y);
+#else
+  return *address;
+#endif
+}
+CREATE_CUDA_DEVICE_FUNCTION_ALIAS(GpuLdg, CudaLdg);
+
+// Zeroes count elements starting at ptr using all threads of a 1-D grid.
+// Note: this function does not synchronize, and therefore the memory range is
+// not guaranteed to be zero until the next kernel launch.
+template <typename T>
+__global__ void SetZero(const int64_t count, T* __restrict__ ptr) {
+  // Check that the grid is one dimensional and index doesn't overflow.
+  assert(blockDim.y == 1);
+  assert(blockDim.z == 1);
+  assert(blockDim.x * gridDim.x / blockDim.x == gridDim.x);
+  for (int64_t i : GpuGridRangeX(count)) {
+    ptr[i] = T(0);
+  }
+}
+
+// Helper to set all tensor entries to a specific value.
+template <typename T, typename Tvalue = T>
+__global__ void SetToValue(const int64_t count, T* __restrict__ ptr,
+                           Tvalue value) {
+  // Check that the grid is one dimensional and index doesn't overflow.
+  assert(blockDim.y == 1);
+  assert(blockDim.z == 1);
+  assert(blockDim.x * gridDim.x / blockDim.x == gridDim.x);
+  for (int64_t i : GpuGridRangeX(count)) {
+    ptr[i] = static_cast<T>(value);
+  }
+}
+
+namespace detail {
+
+template <int N, typename T>
+__device__ T* AddressSpaceHint(T* ptr) {
+#if defined(TENSORFLOW_USE_ROCM)
+  using AS = __attribute__((address_space(N))) T*;
+  auto ptr_ = reinterpret_cast<AS>(reinterpret_cast<uintptr_t>(ptr));
+  return (T*)(ptr_);
+#else
+  return ptr;  // NOOP
+#endif
+}
+
+// Helper function for atomic accumulation implemented as CAS.
+template <typename T, typename F>
+__device__ T GpuAtomicCasHelper(T* ptr, F accumulate) {
+  ptr = detail::AddressSpaceHint<1>(ptr);
+  T old = *ptr;
+  T assumed;
+  do {
+    assumed = old;
+    old = atomicCAS(ptr, assumed, accumulate(assumed));
+  } while (assumed != old);
+  return old;
+}
+CREATE_CUDA_DEVICE_FUNCTION_ALIAS(GpuAtomicCasHelper, CudaAtomicCasHelper);
+
+// Overload for floating point (using integer comparison to handle NaN
+// correctly).
+template <typename F>
+__device__ float GpuAtomicCasHelper(float* ptr, F accumulate) {
+  return __int_as_float(GpuAtomicCasHelper(
+      reinterpret_cast<int32_t*>(ptr), [accumulate](int32_t a) {
+        return __float_as_int(accumulate(__int_as_float(a)));
+      }));
+}
+template <typename F>
+__device__ double GpuAtomicCasHelper(double* ptr, F accumulate) {
+  return __longlong_as_double(GpuAtomicCasHelper(
+      reinterpret_cast<unsigned long long*>(ptr), [accumulate](uint64_t a) {
+        return __double_as_longlong(accumulate(__longlong_as_double(a)));
+      }));
+}
+
+// Overload of above function for half. Note that we don't have
+// atomicCAS() for anything less than 32 bits, so we need to include the
+// other 16 bits in the operation.
+//
+// This version is going to be very slow
+// under high concurrency, since most threads will be spinning on failing
+// their compare-and-swap tests. (The fact that we get false sharing on the
+// neighboring fp16 makes this even worse.) If you are doing a large reduction,
+// you are much better off with doing the intermediate steps in fp32 and then
+// switching to fp16 as late as you can in the calculations.
+//
+// Note: Assumes little endian.
+template <typename F>
+__device__ Eigen::half GpuAtomicCasHelper(Eigen::half* ptr, F accumulate) {
+#if defined(__BYTE_ORDER__) && defined(__ORDER_LITTLE_ENDIAN__)
+  static_assert(__BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__, "Not little endian");
+#endif
+  uintptr_t intptr = reinterpret_cast<uintptr_t>(ptr);
+  uint32_t shift = (intptr & 0x2) * 8U;
+  uint32_t mask = 0xFFFF0000U >> shift;
+
+  assert(!(intptr & 0x1));  // should be 2-aligned.
+  uint32_t* address = reinterpret_cast<uint32_t*>(intptr & ~0x3);
+  uint32_t result =
+      GpuAtomicCasHelper(address, [accumulate, shift, mask](uint32_t arg) {
+        uint16_t high = static_cast<uint16_t>(arg >> shift);
+        Eigen::half acc =
+            accumulate(Eigen::numext::bit_cast<Eigen::half>(high));
+        return (static_cast<uint32_t>(Eigen::numext::bit_cast<uint16_t>(acc))
+                << shift) |
+               (arg & mask);
+      });
+  return Eigen::numext::bit_cast<Eigen::half>(
+      static_cast<uint16_t>(result >> shift));
+}
+
+template <typename F>
+__device__ Eigen::bfloat16 GpuAtomicCasHelper(Eigen::bfloat16* ptr,
+                                              F accumulate) {
+  Eigen::half ret = detail::GpuAtomicCasHelper(
+      reinterpret_cast<Eigen::half*>(ptr), [accumulate](Eigen::half a) {
+        Eigen::bfloat16 acc =
+            accumulate(Eigen::numext::bit_cast<Eigen::bfloat16>(a));
+        return Eigen::numext::bit_cast<Eigen::half>(acc);
+      });
+  return Eigen::numext::bit_cast<Eigen::bfloat16>(ret);
+}
+
+template <typename F>
+__device__ long long GpuAtomicCasHelper(long long* ptr, F accumulate) {
+  return static_cast<long long>(
+      GpuAtomicCasHelper(reinterpret_cast<unsigned long long*>(ptr),
+                         [accumulate](unsigned long long a) {
+                           return static_cast<unsigned long long>(
+                               accumulate(static_cast<long long>(a)));
+                         }));
+}
+
+template <typename From, typename To>
+using ToTypeIfConvertible =
+    typename std::enable_if<std::is_convertible<From, To>::value, To>::type;
+
+template <typename T>
+struct CudaSupportedTypeImpl {
+  using type = T;
+};
+
+template <>
+struct CudaSupportedTypeImpl<long long> {
+  using type = unsigned long long;
+};
+
+template <>
+struct CudaSupportedTypeImpl<unsigned long> {
+  using type =
+      typename std::conditional<sizeof(unsigned long) == sizeof(unsigned int),
+                                unsigned int, unsigned long long>::type;
+};
+
+template <>
+struct CudaSupportedTypeImpl<long> {
+  // This cast should be safe since module-2 addition should work fine. However,
+  // signed overflow is not handled correctly since it's undefined behavior.
+  using type = typename CudaSupportedTypeImpl<unsigned long>::type;
+};
+
+template <typename T>
+using CudaSupportedType = typename CudaSupportedTypeImpl<T>::type;
+
+template <typename T>
+__device__ CudaSupportedType<T>* ToCudaSupportedPtr(T* ptr) {
+  return reinterpret_cast<CudaSupportedType<T>*>(ptr);
+}
+
+}  // namespace detail
+
+// CUDA provides atomic ops, but not for all types.  We provide wrappers
+// for some ops and provide implementation for all reasonable types.
+
+template <typename T, typename U>
+__device__ detail::ToTypeIfConvertible<U, T> GpuAtomicAdd(T* ptr, U value) {
+  return atomicAdd(detail::ToCudaSupportedPtr(detail::AddressSpaceHint<1>(ptr)),
+                   value);
+}
+
+#if GOOGLE_CUDA
+__device__ inline Eigen::bfloat16 GpuAtomicAdd(Eigen::bfloat16* ptr,
+                                               Eigen::bfloat16 value) {
+  return detail::GpuAtomicCasHelper(
+      ptr, [value](Eigen::bfloat16 a) { return a + value; });
+}
+
+__device__ inline Eigen::half GpuAtomicAdd(Eigen::half* ptr,
+                                           Eigen::half value) {
+  return detail::GpuAtomicCasHelper(
+      ptr, [value](Eigen::half a) { return a + value; });
+}
+#endif
+
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 600)
+__device__ inline double GpuAtomicAdd(double* ptr, double value) {
+  return detail::GpuAtomicCasHelper(ptr,
+                                    [value](double a) { return a + value; });
+}
+#endif
+
+#if TENSORFLOW_USE_ROCM
+namespace detail {
+
+template <typename P, typename T, typename F>
+__device__ inline T GpuAtomicAddHalfHelper(T* ptr, T value, F add) {
+  typedef P __attribute__((ext_vector_type(2))) P2;
+  auto ptr2 = (__attribute__((address_space(1)))
+               P2*)(reinterpret_cast<uintptr_t>(ptr) & ~0x3);
+  uintptr_t shift = ((reinterpret_cast<uintptr_t>(ptr) & 0x2) * 8);
+  // Eigen::numext::bit_cast on ext_vector produces redudant inlined memcpy.
+  // Use union instead.
+  union {
+    P2 v2;
+    uint32_t i;
+  } u;
+
+  u.i = static_cast<uint32_t>(Eigen::numext::bit_cast<uint16_t>(value))
+        << shift;
+
+  // Performs + (T)0 on adjacent location, so this is not idempotent with
+  // regards to its bit pattern. Should be fine as long as that locations is
+  // used to hold T.
+  u.v2 = add(ptr2, u.v2);
+  return Eigen::numext::bit_cast<T>(static_cast<uint16_t>(u.i >> shift));
+}
+
+}  // namespace detail
+
+__device__ inline Eigen::bfloat16 GpuAtomicAdd(Eigen::bfloat16* ptr,
+                                               Eigen::bfloat16 value) {
+#if __has_builtin(__builtin_amdgcn_global_atomic_fadd_v2bf16)
+  return detail::GpuAtomicAddHalfHelper<short>(ptr, value, [](auto p, auto v) {
+    return __builtin_amdgcn_global_atomic_fadd_v2bf16(p, v);
+  });
+#else
+  return detail::GpuAtomicCasHelper(
+      ptr, [value](Eigen::bfloat16 a) { return a + value; });
+#endif
+}
+
+__device__ inline Eigen::half GpuAtomicAdd(Eigen::half* ptr,
+                                           Eigen::half value) {
+#if __has_builtin(__builtin_amdgcn_global_atomic_fadd_v2f16)
+  return detail::GpuAtomicAddHalfHelper<_Float16>(
+      ptr, value, [](auto p, auto v) {
+        return __builtin_amdgcn_global_atomic_fadd_v2f16(p, v);
+      });
+#else
+  return detail::GpuAtomicCasHelper(
+      ptr, [value](Eigen::half a) { return a + value; });
+#endif
+}
+#endif
+
+// GpuAtomicAdd
+// Specializations of GpuAtomicAdd for complex types, which GpuAtomicAdd does
+// not support. We treat a std::complex<T>* as a T* (the C++ standard section
+// 26.4.4 allows this explicitly) and atomic add the real and imaginary
+// components individually. The operation as a whole is not atomic, but we can
+// safely treat the components independently for the purpose of accumulating.
+
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+__device__ inline std::complex<float> GpuAtomicAdd(std::complex<float>* ptr,
+                                                   std::complex<float> value) {
+  auto ptr_scalar = reinterpret_cast<float*>(ptr);
+  return std::complex<float>(GpuAtomicAdd(ptr_scalar, value.real()),
+                             GpuAtomicAdd(ptr_scalar + 1, value.imag()));
+}
+
+__device__ inline std::complex<double> GpuAtomicAdd(
+    std::complex<double>* ptr, std::complex<double> value) {
+  auto ptr_scalar = reinterpret_cast<double*>(ptr);
+  return std::complex<double>(GpuAtomicAdd(ptr_scalar, value.real()),
+                              GpuAtomicAdd(ptr_scalar + 1, value.imag()));
+}
+#endif
+CREATE_CUDA_DEVICE_FUNCTION_ALIAS(GpuAtomicAdd, CudaAtomicAdd);
+
+// GpuAtomicSub
+template <typename T, typename U>
+__device__ detail::ToTypeIfConvertible<U, T> GpuAtomicSub(T* ptr, U value) {
+  return atomicSub(detail::AddressSpaceHint<1>(ptr), value);
+}
+
+// Specializations of substraction which add the negative value.
+__device__ inline float GpuAtomicSub(float* ptr, float value) {
+  return GpuAtomicAdd(ptr, -value);
+}
+
+__device__ inline double GpuAtomicSub(double* ptr, double value) {
+  return GpuAtomicAdd(ptr, -value);
+}
+
+__device__ inline int64_t GpuAtomicSub(int64_t* ptr, int64_t value) {
+  return GpuAtomicAdd(ptr, -value);
+}
+
+__device__ inline uint64_t GpuAtomicSub(uint64_t* ptr, uint64_t value) {
+  return GpuAtomicAdd(ptr, -static_cast<int64_t>(value));
+}
+
+__device__ inline Eigen::half GpuAtomicSub(Eigen::half* ptr,
+                                           Eigen::half value) {
+  return GpuAtomicAdd(ptr, -value);
+}
+
+__device__ inline Eigen::bfloat16 GpuAtomicSub(Eigen::bfloat16* ptr,
+                                               Eigen::bfloat16 value) {
+  return GpuAtomicAdd(ptr, -value);
+}
+
+CREATE_CUDA_DEVICE_FUNCTION_ALIAS(GpuAtomicSub, CudaAtomicSub);
+
+// GpuAtomicMax
+template <typename T, typename U>
+__device__ detail::ToTypeIfConvertible<U, T> GpuAtomicMax(T* ptr, U value) {
+  return atomicMax(detail::ToCudaSupportedPtr(detail::AddressSpaceHint<1>(ptr)),
+                   value);
+}
+
+#if TENSORFLOW_USE_ROCM
+
+/*
+ * CUDA runtime headers have the following defined
+ *   __device__  int max(int, int)
+ *   __device__  float max(float, float)
+ *   __device__  double max(double, double)
+ *
+ * and many others, where as HIP runtime headers only have the "int" version
+ *
+ * Therefore need to special case ROCm version to call the correct underlying
+ * routines for float and double types.
+ *
+ */
+
+__device__ inline float GpuAtomicMax(float* ptr, float value) {
+  return detail::GpuAtomicCasHelper(
+      ptr, [value](float a) { return fmaxf(a, value); });
+}
+
+__device__ inline double GpuAtomicMax(double* ptr, double value) {
+  return detail::GpuAtomicCasHelper(
+      ptr, [value](double a) { return fmax(a, value); });
+}
+
+#else
+
+__device__ inline float GpuAtomicMax(float* ptr, float value) {
+  return detail::GpuAtomicCasHelper(ptr,
+                                    [value](float a) { return max(a, value); });
+}
+
+__device__ inline double GpuAtomicMax(double* ptr, double value) {
+  return detail::GpuAtomicCasHelper(
+      ptr, [value](double a) { return max(a, value); });
+}
+
+#endif
+
+__device__ inline Eigen::half GpuAtomicMax(Eigen::half* ptr,
+                                           Eigen::half value) {
+  return detail::GpuAtomicCasHelper(
+      ptr, [value](Eigen::half a) { return max(a, value); });
+}
+
+__device__ inline Eigen::bfloat16 GpuAtomicMax(Eigen::bfloat16* ptr,
+                                               Eigen::bfloat16 value) {
+  return detail::GpuAtomicCasHelper(
+      ptr, [value](Eigen::bfloat16 a) { return max(a, value); });
+}
+
+#if TENSORFLOW_USE_ROCM || (__CUDA_ARCH__ < 320)
+__device__ inline uint64_t GpuAtomicMax(uint64_t* ptr, uint64_t value) {
+  return detail::GpuAtomicCasHelper(
+      detail::ToCudaSupportedPtr(ptr),
+      [value](uint64_t a) { return max(a, value); });
+}
+
+__device__ inline int64_t GpuAtomicMax(int64_t* ptr, int64_t value) {
+  return detail::GpuAtomicCasHelper(
+      detail::ToCudaSupportedPtr(ptr),
+      [value](int64_t a) { return max(a, value); });
+}
+#endif
+CREATE_CUDA_DEVICE_FUNCTION_ALIAS(GpuAtomicMax, CudaAtomicMax);
+
+// GpuAtomicMin
+template <typename T, typename U>
+__device__ detail::ToTypeIfConvertible<U, T> GpuAtomicMin(T* ptr, U value) {
+  return atomicMin(detail::ToCudaSupportedPtr(detail::AddressSpaceHint<1>(ptr)),
+                   value);
+}
+
+#if TENSORFLOW_USE_ROCM
+
+/*
+ * CUDA runtime headers have the following defined
+ *   __device__  int min(int, int)
+ *   __device__  float min(float, float)
+ *   __device__  double min(double, double)
+ *
+ * and many others, where as HIP runtime headers only have the "int" version
+ *
+ * Therefore need to special case ROCm version to call the correct underlying
+ * routines for float and double types.
+ *
+ */
+
+__device__ inline float GpuAtomicMin(float* ptr, float value) {
+  return detail::GpuAtomicCasHelper(
+      ptr, [value](float a) { return fminf(a, value); });
+}
+
+__device__ inline double GpuAtomicMin(double* ptr, double value) {
+  return detail::GpuAtomicCasHelper(
+      ptr, [value](double a) { return fmin(a, value); });
+}
+
+#else
+
+__device__ inline float GpuAtomicMin(float* ptr, float value) {
+  return detail::GpuAtomicCasHelper(ptr,
+                                    [value](float a) { return min(a, value); });
+}
+
+__device__ inline double GpuAtomicMin(double* ptr, double value) {
+  return detail::GpuAtomicCasHelper(
+      ptr, [value](double a) { return min(a, value); });
+}
+
+#endif
+
+__device__ inline Eigen::half GpuAtomicMin(Eigen::half* ptr,
+                                           Eigen::half value) {
+  return detail::GpuAtomicCasHelper(
+      ptr, [value](Eigen::half a) { return min(a, value); });
+}
+
+__device__ inline Eigen::bfloat16 GpuAtomicMin(Eigen::bfloat16* ptr,
+                                               Eigen::bfloat16 value) {
+  return detail::GpuAtomicCasHelper(
+      ptr, [value](Eigen::bfloat16 a) { return min(a, value); });
+}
+
+#if TENSORFLOW_USE_ROCM || (__CUDA_ARCH__ < 320)
+__device__ inline uint64_t GpuAtomicMin(uint64_t* ptr, uint64_t value) {
+  return detail::GpuAtomicCasHelper(
+      detail::ToCudaSupportedPtr(ptr),
+      [value](uint64_t a) { return min(a, value); });
+}
+
+__device__ inline int64_t GpuAtomicMin(int64_t* ptr, int64_t value) {
+  return detail::GpuAtomicCasHelper(
+      detail::ToCudaSupportedPtr(ptr),
+      [value](int64_t a) { return min(a, value); });
+}
+#endif
+CREATE_CUDA_DEVICE_FUNCTION_ALIAS(GpuAtomicMin, CudaAtomicMin);
+
+// GpuAtomicMul
+template <typename T, typename U>
+__device__ detail::ToTypeIfConvertible<U, T> GpuAtomicMul(T* ptr, U value) {
+  return detail::GpuAtomicCasHelper(ptr, [value](T a) { return a * value; });
+}
+CREATE_CUDA_DEVICE_FUNCTION_ALIAS(GpuAtomicMul, CudaAtomicMul);
+
+// GpuAtomicDiv
+template <typename T, typename U>
+__device__ detail::ToTypeIfConvertible<U, T> GpuAtomicDiv(T* ptr, U value) {
+  return detail::GpuAtomicCasHelper(ptr, [value](T a) { return a / value; });
+}
+CREATE_CUDA_DEVICE_FUNCTION_ALIAS(GpuAtomicDiv, CudaAtomicDiv);
+
+// Import all specialized std::complex device operators in namespace tensorflow.
+#if GOOGLE_CUDA && defined(EIGEN_USING_STD_COMPLEX_OPERATORS)
+EIGEN_USING_STD_COMPLEX_OPERATORS
+#endif  // GOOGLE_CUDA
+
+namespace functor {
+// Import all specialized std::complex device operators in namespace functor.
+#if GOOGLE_CUDA && defined(EIGEN_USING_STD_COMPLEX_OPERATORS)
+EIGEN_USING_STD_COMPLEX_OPERATORS
+#endif  // GOOGLE_CUDA
+
+// ROCm hcc(clang) has severe difficulties dealing with std::complex directly
+// due to a header issue. This template assists in casting std::complex into the
+// corresponding internal ROCm types.
+template <class T>
+struct MapComplexToHipComplex {
+  typedef T TM;
+};
+
+#if TENSORFLOW_USE_ROCM
+template <>
+struct MapComplexToHipComplex<std::complex<float> > {
+  typedef hipFloatComplex TM;
+};
+
+template <>
+struct MapComplexToHipComplex<std::complex<double> > {
+  typedef hipDoubleComplex TM;
+};
+#endif
+};  // namespace functor
+
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+#endif  // TENSORFLOW_CORE_UTIL_GPU_DEVICE_FUNCTIONS_H_

tensorflow/core/util/gpu_kernel_helper.h

@@ -0,0 +1,524 @@
+/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_UTIL_GPU_KERNEL_HELPER_H_
+#define TENSORFLOW_CORE_UTIL_GPU_KERNEL_HELPER_H_
+
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+
+#include <type_traits>
+
+#if GOOGLE_CUDA
+#include "third_party/gpus/cuda/include/cuda_fp16.h"
+#endif
+#include "tensorflow/core/util/gpu_cuda_alias.h"
+#include "tensorflow/core/util/gpu_device_functions.h"
+#include "tensorflow/core/util/gpu_launch_config.h"
+
+#if GOOGLE_CUDA
+#define TF_RED_WARPSIZE 32
+#elif TENSORFLOW_USE_ROCM
+#define TF_RED_WARPSIZE 64
+#endif
+
+// Deprecated, use 'for(int i : GpuGridRangeX(n))' instead.
+#define GPU_1D_KERNEL_LOOP(i, n) \
+  for (int i : ::tensorflow::GpuGridRangeX<int>(n))
+#define CUDA_1D_KERNEL_LOOP(i, n) \
+  for (int i : ::tensorflow::GpuGridRangeX<int>(n))
+
+// Deprecated, use 'for(int i : GpuGridRange?(n))' instead.
+#define GPU_AXIS_KERNEL_LOOP(i, n, axis) \
+  for (int i : ::tensorflow::GpuGridRange##axis<int>(n))
+#define CUDA_AXIS_KERNEL_LOOP(i, n, axis) \
+  for (int i : ::tensorflow::GpuGridRange##axis<int>(n))
+
+#if GOOGLE_CUDA
+#define gpuSuccess cudaSuccess
+using gpuStream_t = cudaStream_t;
+using gpuError_t = cudaError_t;
+#elif TENSORFLOW_USE_ROCM
+#define gpuSuccess hipSuccess
+using gpuStream_t = hipStream_t;
+using gpuError_t = hipError_t;
+#endif
+
+// macro wrapper to declare dynamic shared memory
+#if GOOGLE_CUDA
+
+#define GPU_DYNAMIC_SHARED_MEM_DECL(ALIGN, TYPE, NAME) \
+  extern __shared__ __align__(ALIGN)                   \
+  TYPE NAME[]
+
+#elif TENSORFLOW_USE_ROCM
+
+#define GPU_DYNAMIC_SHARED_MEM_DECL(ALIGN, TYPE, NAME) \
+  HIP_DYNAMIC_SHARED(TYPE, NAME)
+
+#endif
+
+namespace tensorflow {
+
+#if GOOGLE_CUDA
+// cudaGetErrorString is available to both host and device
+__host__ __device__ inline const char* GpuGetErrorString(cudaError_t error) {
+  return cudaGetErrorString(error);
+}
+#elif TENSORFLOW_USE_ROCM
+// hipGetErrorString is available on host side only
+inline const char* GpuGetErrorString(hipError_t error) {
+  return hipGetErrorString(error);
+}
+#endif
+
+// Returns a raw reference to the current cuda stream. Required by a
+// number of kernel calls (for which StreamInterface* does not work),
+// i.e. CUB and certain cublas primitives.
+inline gpuStream_t GetGpuStream(OpKernelContext* context) {
+  void* opaque_stream = CHECK_NOTNULL(context->op_device_context()
+                                          ->stream()
+                                          ->platform_specific_handle()
+                                          .stream);
+  return reinterpret_cast<gpuStream_t>(opaque_stream);
+}
+
+// Launches a GPU kernel through cudaLaunchKernel in CUDA environment, or
+// hipLaunchKernel in ROCm environment with the given arguments.
+//
+// The kernel parameters 'Ts' must be constructible from the arguments 'Args'.
+template <typename... Ts, typename... Args>
+absl::Status GpuLaunchKernel(void (*function)(Ts...), dim3 grid_dim,
+                             dim3 block_dim, size_t shared_memory_size_bytes,
+                             gpuStream_t stream, Args... arguments) {
+  static_assert(detail::NoneIsReference<Ts...>(),
+                "Kernels with reference arguments have undefined behaviour.");
+  if (grid_dim.x * grid_dim.y * grid_dim.z > 0 &&
+      block_dim.x * block_dim.y * block_dim.z > 0) {
+#if GOOGLE_CUDA
+    auto func_ptr = absl::bit_cast<const void*>(function);
+    // Cast arguments and forward them as an array of pointers.
+    auto args_tuple = std::tuple<Ts...>(arguments...);
+    auto arg_ptrs = detail::GetArrayOfElementPointers(&args_tuple);
+    auto result =
+        cudaLaunchKernel(func_ptr, grid_dim, block_dim, arg_ptrs.data(),
+                         shared_memory_size_bytes, stream);
+    if (result != cudaSuccess) {
+      return errors::Internal(cudaGetErrorString(result));
+    }
+#elif TENSORFLOW_USE_ROCM
+    hipLaunchKernelGGL(function, grid_dim, block_dim, shared_memory_size_bytes,
+                       stream, std::forward<Args>(arguments)...);
+    TF_RETURN_IF_CUDA_ERROR(hipGetLastError());
+#endif
+  }
+  return absl::OkStatus();
+}
+
+// Perfect forwarding to make CudaLaunchKernel available to both ROCm and CUDA
+// builds
+template <typename... Args>
+auto CudaLaunchKernel(Args&&... args)
+    -> decltype(GpuLaunchKernel(std::forward<Args>(args)...)) {
+  return GpuLaunchKernel(std::forward<Args>(args)...);
+}
+
+__host__ __device__ inline tensorflow::bfloat16 GpuLdg(
+    const tensorflow::bfloat16* address) {
+  return Eigen::numext::bit_cast<tensorflow::bfloat16>(
+      GpuLdg(reinterpret_cast<const uint16_t*>(address)));
+}
+// Already aliased in gpu_device_functions.h
+
+template <typename T>
+__host__ __device__ inline T ldg(const T* ptr) {
+  return GpuLdg(ptr);
+}
+
+template <typename T>
+__host__ __device__ inline const T& tf_min(const T& x, const T& y) {
+  return x < y ? x : y;
+}
+
+template <typename T>
+__host__ __device__ inline const T& tf_max(const T& x, const T& y) {
+  return x < y ? y : x;
+}
+
+// Overloads of the above functions for float and double.
+__host__ __device__ inline float tf_min(float x, float y) {
+  return fminf(x, y);
+}
+__host__ __device__ inline double tf_min(double x, double y) {
+  return fmin(x, y);
+}
+__host__ __device__ inline float tf_max(float x, float y) {
+  return fmaxf(x, y);
+}
+__host__ __device__ inline double tf_max(double x, double y) {
+  return fmax(x, y);
+}
+
+#ifdef _MSC_VER
+#if _MSC_VER >= 1930
+using std::max;
+using std::min;
+__host__ __device__ inline int tf_min(int x, int y) { return min(x, y); }
+__host__ __device__ inline int tf_max(int x, int y) { return max(x, y); }
+#endif
+#endif
+
+// ROCM TODO re-enable them after adding fp16 support logic
+#if GOOGLE_CUDA
+__device__ inline Eigen::half GpuShuffleSync(unsigned mask, Eigen::half value,
+                                             int src_lane,
+                                             int width = warpSize) {
+  return Eigen::half(
+      GpuShuffleSync(mask, static_cast<uint16_t>(value), src_lane, width));
+}
+// Aliased in gpu_device_functions.h
+
+__device__ EIGEN_ALWAYS_INLINE Eigen::half GpuShuffleUpSync(
+    unsigned mask, Eigen::half value, int delta, int width = warpSize) {
+  return Eigen::half(
+      GpuShuffleUpSync(mask, static_cast<uint16_t>(value), delta, width));
+}
+// Aliased in gpu_device_functions.h
+
+__device__ EIGEN_ALWAYS_INLINE Eigen::half GpuShuffleDownSync(
+    unsigned mask, Eigen::half value, int delta, int width = warpSize) {
+  return Eigen::half(
+      GpuShuffleDownSync(mask, static_cast<uint16_t>(value), delta, width));
+}
+// Aliased in gpu_device_functions.h
+
+__device__ EIGEN_ALWAYS_INLINE Eigen::half GpuShuffleXorSync(
+    unsigned mask, Eigen::half value, int lane_mask, int width = warpSize) {
+  return Eigen::half(
+      GpuShuffleXorSync(mask, static_cast<uint16_t>(value), lane_mask, width));
+}
+// Aliased in gpu_device_functions.h
+#endif
+
+#ifdef __CUDA_ARCH__
+#define UNROLL_ON_DEVICE _Pragma("unroll")
+#else
+#define UNROLL_ON_DEVICE
+#endif
+
+// Represents an aligned array of N elements of T. Data pointers can be
+// reinterpreted as this type to generate vectorized loads/stores in a kernel.
+template <typename T, int N>
+class alignas(alignof(T) * N) AlignedVector {
+ public:
+  typedef T value_type;
+  static constexpr const int kSize = N;
+
+  AlignedVector() = default;
+
+  // Uniform initialization.
+  __host__ __device__ explicit AlignedVector(value_type uniform) {
+    UNROLL_ON_DEVICE for (int i = 0; i < kSize; ++i) { values_[i] = uniform; }
+  }
+  // Uniform initialization with explicit conversion.
+  // Note: This is required for T=Eigen::half because it only supports explicit
+  // conversions from other types and its template constructor is too relaxed
+  // to be able to use std::is_constructible.
+  template <typename U, typename std::enable_if<std::is_arithmetic<U>::value,
+                                                int>::type = 0>
+  __host__ __device__ explicit AlignedVector(U uniform_u) {
+    value_type uniform(uniform_u);
+    UNROLL_ON_DEVICE for (int i = 0; i < kSize; ++i) { values_[i] = uniform; }
+  }
+  // Implicit conversion.
+  template <typename U, typename std::enable_if<
+                            std::is_convertible<U, T>::value, int>::type = 0>
+  __host__ __device__ AlignedVector(const AlignedVector<U, N>& other) {
+    UNROLL_ON_DEVICE for (int i = 0; i < kSize; ++i) { values_[i] = other[i]; }
+  }
+  // Explicit conversion.
+  template <typename U,
+            typename std::enable_if<!std::is_convertible<U, T>::value &&
+                                        std::is_constructible<T, U>::value,
+                                    int>::type = 0>
+  __host__ __device__ explicit AlignedVector(const AlignedVector<U, N>& other) {
+    UNROLL_ON_DEVICE for (int i = 0; i < kSize; ++i) {
+      values_[i] = T(other[i]);
+    }
+  }
+
+  __host__ __device__ value_type& operator[](int i) { return values_[i]; }
+  __host__ __device__ const value_type& operator[](int i) const {
+    return values_[i];
+  }
+
+#define DEFINE_BINARY_UPDATE_OPERATOR(op)                                      \
+  __host__ __device__ AlignedVector& operator op(const AlignedVector& rhs) {   \
+    UNROLL_ON_DEVICE for (int i = 0; i < kSize; ++i) { values_[i] op rhs[i]; } \
+    return *this;                                                              \
+  }
+  DEFINE_BINARY_UPDATE_OPERATOR(+=)
+  DEFINE_BINARY_UPDATE_OPERATOR(-=)
+  DEFINE_BINARY_UPDATE_OPERATOR(*=)
+  DEFINE_BINARY_UPDATE_OPERATOR(/=)
+#undef DEFINE_BINARY_UPDATE_OPERATOR
+
+#define DEFINE_BINARY_OPERATOR(op)                          \
+  friend __host__ __device__ AlignedVector operator op(     \
+      const AlignedVector& lhs, const AlignedVector& rhs) { \
+    AlignedVector ret;                                      \
+    UNROLL_ON_DEVICE for (int i = 0; i < kSize; ++i) {      \
+      ret[i] = lhs[i] op rhs[i];                            \
+    }                                                       \
+    return ret;                                             \
+  }
+  DEFINE_BINARY_OPERATOR(+)
+  DEFINE_BINARY_OPERATOR(-)
+  DEFINE_BINARY_OPERATOR(*)
+  DEFINE_BINARY_OPERATOR(/)
+#undef DEFINE_BINARY_OPERATOR
+
+#define DEFINE_BINARY_FUNCTION(func)                                        \
+  friend __host__ __device__ AlignedVector func(const AlignedVector& lhs,   \
+                                                const AlignedVector& rhs) { \
+    AlignedVector ret;                                                      \
+    UNROLL_ON_DEVICE for (int i = 0; i < kSize; ++i) {                      \
+      ret[i] = func(lhs[i], rhs[i]);                                        \
+    }                                                                       \
+    return ret;                                                             \
+  }
+  DEFINE_BINARY_FUNCTION(min)
+  DEFINE_BINARY_FUNCTION(max)
+#undef DEFINE_BINARY_FUNCTION
+
+ private:
+  value_type values_[N];
+};
+
+#undef UNROLL_ON_DEVICE
+
+// Returns the maximum power-of-two alignment (in units of elements, not bytes)
+// of a stride or pointer value.
+inline int64_t alignment_of(int64_t element_stride) {
+  // A zero/nullptr value means that the stride/pointer is not used, so it
+  // effectively has infinite alignment.
+  constexpr int64_t kMaxAlignment = 512;
+  if (element_stride == 0) return kMaxAlignment;
+  return element_stride & -element_stride;
+}
+
+template <typename T>
+inline int64_t alignment_of(T* ptr) {
+  const intptr_t ptr_val = reinterpret_cast<std::uintptr_t>(ptr);
+  // Pointers should always be aligned to sizeof(T) bytes.
+  DCHECK_EQ(ptr_val % sizeof(T), 0);
+  // Note that we want the alignment in elements, not bytes.
+  return alignment_of(ptr_val / sizeof(T));
+}
+
+template <typename... Args>
+int64_t MinAlignmentOf(Args... args) {
+  return std::min({alignment_of(args)...});
+}
+
+namespace detail {
+
+template <int64_t VecSize, template <int vec_size> class Functor>
+struct DispatchToVectorizedHelper {
+  template <typename... Args>
+  absl::Status operator()(int64_t max_vec_size, Args&&... args) const {
+    if (max_vec_size >= VecSize) {
+      return Functor<VecSize>()(std::forward<Args>(args)...);
+    }
+    return DispatchToVectorizedHelper<VecSize / 2, Functor>()(
+        max_vec_size, std::forward<Args>(args)...);
+  }
+};
+template <template <int vec_size> class Functor>
+struct DispatchToVectorizedHelper<1, Functor> {
+  template <typename... Args>
+  absl::Status operator()(int64_t max_vec_size, Args&&... args) const {
+    return Functor<1>()(std::forward<Args>(args)...);
+  }
+};
+
+}  // namespace detail
+
+// Calls Functor<vec_size>()(args...) with vec_size set to the optimal GPU
+// vector instruction size for type T that is <= max_vec_size. The max_vec_size
+// argument should be set to the minimum alignment of all relevant parameters.
+// Requires sizeof(T) to be a power of 2.
+template <typename T, template <int vec_size> class Functor, typename... Args>
+absl::Status DispatchToVectorized(int64_t max_vec_size, Args&&... args) {
+  static_assert((sizeof(T) & (sizeof(T) - 1)) == 0,
+                "sizeof(T) must be a power of 2");
+  if (max_vec_size <= 0) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("DispatchToVectorized: max_vec_size (", max_vec_size,
+                     ") must be greater than zero."));
+  }
+  constexpr const int kOptimalVecSizeBytes = 16;
+  // The optimal number of (aligned) elements of T to load/store in a
+  // single instruction inside a kernel.
+  constexpr const int optimal_vec_size =
+      (kOptimalVecSizeBytes - 1) / sizeof(T) + 1;
+  return detail::DispatchToVectorizedHelper<optimal_vec_size, Functor>()(
+      max_vec_size, std::forward<Args>(args)...);
+}
+
+// Similar to std::upper_bound, this returns the index of the first element in
+// [first, first + count) that is greater than `val`, or `count` if no such
+// element is found. Assumes [first, first + count) is sorted.
+namespace gpu_helper {
+template <typename T, typename OutType = int32_t, typename Iterator = const T*>
+__device__ OutType upper_bound(Iterator first, OutType count, T val) {
+  Iterator orig = first;
+  OutType step = 0;
+  while (count > 0) {
+    Iterator it = first;
+    step = count / 2;
+    it += step;
+    if (!(val < *it)) {
+      first = ++it;
+      count -= step + 1;
+    } else {
+      count = step;
+    }
+  }
+
+  return first - orig;
+}
+
+// Similar to std::lower_bound, this returns the index of the first element in
+// [first, first + count) that is not less than `val`, or `count` if no such
+// element is found. Assumes [first, first + count) is sorted.
+template <typename T, typename OutType = int32_t, typename Iterator = const T*>
+__device__ OutType lower_bound(Iterator first, OutType count, T val) {
+  Iterator orig = first;
+  OutType step = 0;
+  while (count > 0) {
+    Iterator it = first;
+    step = count / 2;
+    it += step;
+    if (*it < val) {
+      first = ++it;
+      count -= step + 1;
+    } else {
+      count = step;
+    }
+  }
+
+  return first - orig;
+}
+
+}  // namespace gpu_helper
+
+#ifndef TENSORFLOW_USE_ROCM
+namespace cuda_helper = gpu_helper;
+#endif
+
+// For int division, we can substitute the fast multiplication for slow
+// division. For detailed information see:
+//   https://ridiculousfish.com/blog/posts/labor-of-division-episode-i.html
+//
+// Warning: This implementation only works when the divisor is [1, INT32_MAX]
+//          and the numerator has to be [0, INT32_MAX]. This is enough for our
+//          purpose for computing integer indices.
+// Basics: the typical int division can be written as:
+//   n / d = (m * n) / 2^(32 + s)
+// where 'n' is the numerator and 'd' is the divisor. For a given 'd', we
+// need to find a magic number 'm' and a shift 's'. See update_magic().
+struct FastDividerUint32 {
+  EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC FastDividerUint32(uint32_t d)
+      : divisor(d) {
+    assert(divisor >= 1 && divisor <= INT32_MAX);
+    update_magic();
+  }
+
+  EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC void update_magic() {
+    // (1). The shift 's' is calculated by log2ceil(d).
+#if defined(__CUDA_ARCH__)
+    shift = 32 - __clz(divisor - 1);
+#else
+    for (shift = 0; shift < 32; shift++) {
+      if ((1U << shift) >= divisor) break;
+    }
+#endif
+
+    // (2). The magic number 'm' is calculated by:
+    //   m = 2^(32 + s) / d + 1
+    // Note, the digit '1' is to round up 'm * n', which will be rounded down
+    // later by dividing two. In practice, 'm' is a 33-bit value. To fit the
+    // 32-bit range, we introduce:
+    //   magic = m - 2^32
+    //         = 2^(32 + s) / d - 2^32 + 1
+    //         = 2^32 * 2^s / d - 2^32 * d / d + 1
+    //         = (2^32 * (2^s - d)) / d + 1, where 'magic' will be in 32-bit.
+    uint64_t m = (0x100000000ull * ((0x1ull << shift) - divisor)) / divisor + 1;
+    magic = static_cast<uint32_t>(m);
+  }
+
+  EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC FastDividerUint32& operator=(
+      uint32_t d) {
+    assert(divisor >= 1 && divisor <= INT32_MAX);
+    this->divisor = d;
+    update_magic();
+    return *this;
+  }
+
+  EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC operator uint32_t() const {
+    return divisor;
+  }
+
+  uint32_t divisor;
+  uint32_t magic;
+  uint32_t shift;
+};
+
+EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC uint32_t
+operator/(const uint32_t n, const FastDividerUint32& fdiv) {
+  // (3). We use the 32-bit 'magic' instead of 'm' in the formula:
+  //   n / d = (m * n) / 2^(32 + s)
+  //         = (magic + 2^32) * n / 2^(32 + s)
+  //         = (magic * n) / 2^(32 + s) + n / 2^s
+  //         = (magic * n) / 2^32 / 2^s + n / 2^s
+  //         = (magic * n / 2^32 + n) / 2^s
+#if defined(__CUDA_ARCH__)
+  uint32_t q = __umulhi(n, fdiv.magic);
+#else
+  uint32_t q =
+      static_cast<uint32_t>((static_cast<uint64_t>(n) * fdiv.magic) >> 32);
+#endif
+  return (n + q) >> fdiv.shift;
+}
+
+EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC uint32_t
+operator%(const uint32_t n, const FastDividerUint32& fdiv) {
+  return n - (n / fdiv) * fdiv.divisor;
+}
+
+EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC uint32_t
+operator/(const int n, const FastDividerUint32& fdiv) {
+  return static_cast<uint32_t>(n) / fdiv;
+}
+
+EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC uint32_t
+operator%(const int n, const FastDividerUint32& fdiv) {
+  return static_cast<uint32_t>(n) % fdiv;
+}
+
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+#endif  // TENSORFLOW_CORE_UTIL_GPU_KERNEL_HELPER_H_

tensorflow/core/util/gpu_kernel_helper_test.cu.cc

@@ -0,0 +1,369 @@
+/* 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.
+==============================================================================*/
+
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+#define EIGEN_USE_GPU
+
+#include <time.h>
+
+#include <numeric>
+
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/util/gpu_kernel_helper.h"
+#include "tensorflow/core/util/gpu_launch_config.h"
+
+#define CUDA_EXPECT_SUCCESS                                 \
+  {                                                         \
+    gpuDeviceSynchronize();                                 \
+    cudaError_t err = cudaGetLastError();                   \
+    EXPECT_EQ(cudaSuccess, err) << cudaGetErrorString(err); \
+  }
+
+#define CUDA_ASSERT_SUCCESS                                 \
+  {                                                         \
+    gpuDeviceSynchronize();                                 \
+    cudaError_t err = cudaGetLastError();                   \
+    ASSERT_EQ(cudaSuccess, err) << cudaGetErrorString(err); \
+  }
+
+namespace tensorflow {
+
+namespace {
+
+__global__ void SetOutbufZero(GpuLaunchConfig config,
+                              int* __restrict__ outbuf) {
+  GPU_1D_KERNEL_LOOP(x, config.virtual_thread_count) { outbuf[x] = 0; }
+}
+
+// counting number of jobs by using atomic +1
+__global__ void Count1D(GpuLaunchConfig config, int bufsize,
+                        int* __restrict__ outbuf) {
+  GPU_1D_KERNEL_LOOP(x, config.virtual_thread_count) {
+    if (x < 0) {  // x might overflow when testing extreme case
+      break;
+    }
+    atomicAdd(&outbuf[x % bufsize], 1);
+  }
+}
+__global__ void Count2D(Gpu2DLaunchConfig config, int bufsize,
+                        int* __restrict__ outbuf) {
+  GPU_AXIS_KERNEL_LOOP(x, config.virtual_thread_count.x, X) {
+    if (x < 0) {  // x might overflow when testing extreme case
+      break;
+    }
+    GPU_AXIS_KERNEL_LOOP(y, config.virtual_thread_count.y, Y) {
+      if (y < 0) {  // y might overflow when testing extreme case
+        break;
+      }
+      int idx = x * config.virtual_thread_count.y + y;
+      atomicAdd(&outbuf[idx % bufsize], 1);
+    }
+  }
+}
+__global__ void Count3D(Gpu3DLaunchConfig config, int bufsize,
+                        int* __restrict__ outbuf) {
+  GPU_AXIS_KERNEL_LOOP(x, config.virtual_thread_count.x, X) {
+    if (x < 0) {  // x might overflow when testing extreme case
+      break;
+    }
+    GPU_AXIS_KERNEL_LOOP(y, config.virtual_thread_count.y, Y) {
+      if (y < 0) {  // y might overflow when testing extreme case
+        break;
+      }
+      GPU_AXIS_KERNEL_LOOP(z, config.virtual_thread_count.z, Z) {
+        if (z < 0) {  // z might overflow when testing extreme case
+          break;
+        }
+        int idx =
+            x * config.virtual_thread_count.y * config.virtual_thread_count.z +
+            y * config.virtual_thread_count.z + z;
+        atomicAdd(&outbuf[idx % bufsize], 1);
+      }
+    }
+  }
+}
+
+__global__ void GpuShuffleGetSrcLaneTest(unsigned* __restrict__ failure_count) {
+  unsigned lane_id = GpuLaneId();
+  for (int width = warpSize; width > 1; width /= 2) {
+    auto check_result = [&](const char* op_name, int param, unsigned actual,
+                            unsigned expected) {
+      if (actual != expected) {
+        printf("Cuda%sGetSrcLane(%d, %d) for lane %d returned %d, not %d\n",
+               op_name, param, width, lane_id, actual, expected);
+        GpuAtomicAdd(failure_count, 1);
+      }
+    };
+
+    for (int src_lane = -warpSize; src_lane <= warpSize; ++src_lane) {
+#if TENSORFLOW_USE_ROCM
+      if (src_lane < 0 || src_lane >= width) continue;
+#endif
+      unsigned actual_lane = detail::GpuShuffleGetSrcLane(src_lane, width);
+      unsigned expect_lane =
+          GpuShuffleSync(kCudaWarpAll, lane_id, src_lane, width);
+      check_result("Shuffle", src_lane, actual_lane, expect_lane);
+    }
+
+    for (unsigned delta = 0; delta <= warpSize; ++delta) {
+      unsigned actual_lane = detail::GpuShuffleUpGetSrcLane(delta, width);
+      unsigned expect_lane =
+          GpuShuffleUpSync(kCudaWarpAll, lane_id, delta, width);
+      check_result("ShuffleUp", delta, actual_lane, expect_lane);
+    }
+
+    for (unsigned delta = 0; delta <= warpSize; ++delta) {
+      unsigned actual_lane = detail::GpuShuffleDownGetSrcLane(delta, width);
+      unsigned expect_lane =
+          GpuShuffleDownSync(kCudaWarpAll, lane_id, delta, width);
+      check_result("ShuffleDown", delta, actual_lane, expect_lane);
+    }
+
+    for (int lane_lane = warpSize; lane_lane > 0; lane_lane /= 2) {
+      unsigned actual_lane = detail::GpuShuffleXorGetSrcLane(lane_lane, width);
+      unsigned expect_lane =
+          GpuShuffleXorSync(kCudaWarpAll, lane_id, lane_lane, width);
+      check_result("ShuffleXor", lane_lane, actual_lane, expect_lane);
+    }
+  }
+}
+
+}  // namespace
+
+class GpuLaunchConfigTest : public ::testing::Test {
+ protected:
+  static const int bufsize = 1024;
+  int* outbuf = nullptr;
+  int* outbuf_host = nullptr;
+  int hostbuf[bufsize];
+  Eigen::GpuStreamDevice stream;
+  Eigen::GpuDevice d = Eigen::GpuDevice(&stream);
+
+  void copyToHost() {
+#if TENSORFLOW_USE_ROCM
+    hipError_t err = hipMemcpy(hostbuf, outbuf, sizeof(int) * bufsize,
+                               hipMemcpyDeviceToHost);
+    ASSERT_EQ(cudaSuccess, err) << cudaGetErrorString(err);
+#endif
+  }
+  virtual void SetUp() {
+#if GOOGLE_CUDA
+    cudaError_t err = cudaMallocManaged(&outbuf, sizeof(int) * bufsize);
+    outbuf_host = outbuf;
+#else
+    hipError_t err = hipMalloc(&outbuf, sizeof(int) * bufsize);
+    outbuf_host = hostbuf;
+#endif
+    ASSERT_EQ(cudaSuccess, err) << cudaGetErrorString(err);
+  }
+
+  virtual void TearDown() {
+    ASSERT_EQ(gpuDeviceSynchronize(), cudaSuccess);
+    ASSERT_EQ(gpuFree(outbuf), cudaSuccess);
+    outbuf = nullptr;
+  }
+};
+
+TEST_F(GpuLaunchConfigTest, GetGpuLaunchConfig) {
+  GpuLaunchConfig cfg;
+
+// test valid inputs
+#define TEST_LAUNCH_PARAMETER(work_element_count)                             \
+  cfg = GetGpuLaunchConfig(bufsize, d);                                       \
+  TF_CHECK_OK(GpuLaunchKernel(SetOutbufZero, cfg.block_count,                 \
+                              cfg.thread_per_block, 0, d.stream(), cfg,       \
+                              outbuf));                                       \
+  CUDA_ASSERT_SUCCESS                                                         \
+  cfg = GetGpuLaunchConfig(work_element_count, d);                            \
+  TF_CHECK_OK(GpuLaunchKernel(Count1D, cfg.block_count, cfg.thread_per_block, \
+                              0, d.stream(), cfg, bufsize, outbuf));          \
+  CUDA_EXPECT_SUCCESS                                                         \
+  copyToHost();                                                               \
+  EXPECT_EQ(work_element_count,                                               \
+            std::accumulate(outbuf_host, outbuf_host + bufsize, 0));          \
+                                                                              \
+  cfg = GetGpuLaunchConfig(bufsize, d, SetOutbufZero, 0, 0);                  \
+  TF_CHECK_OK(GpuLaunchKernel(SetOutbufZero, cfg.block_count,                 \
+                              cfg.thread_per_block, 0, d.stream(), cfg,       \
+                              outbuf));                                       \
+  CUDA_ASSERT_SUCCESS                                                         \
+  cfg = GetGpuLaunchConfig(work_element_count, d, Count1D, 0, 0);             \
+  TF_CHECK_OK(GpuLaunchKernel(Count1D, cfg.block_count, cfg.thread_per_block, \
+                              0, d.stream(), cfg, bufsize, outbuf));          \
+  CUDA_EXPECT_SUCCESS                                                         \
+  copyToHost();                                                               \
+  EXPECT_EQ(work_element_count,                                               \
+            std::accumulate(outbuf_host, outbuf_host + bufsize, 0));
+
+  TEST_LAUNCH_PARAMETER(128);
+  TEST_LAUNCH_PARAMETER(129);
+  TEST_LAUNCH_PARAMETER(511);
+  TEST_LAUNCH_PARAMETER(512);
+  TEST_LAUNCH_PARAMETER(2048);
+  TEST_LAUNCH_PARAMETER(2049);
+  TEST_LAUNCH_PARAMETER(8191);
+  TEST_LAUNCH_PARAMETER(8192);
+  TEST_LAUNCH_PARAMETER(123456);
+  TEST_LAUNCH_PARAMETER(1 << 30);
+#undef TEST_LAUNCH_PARAMETER
+}
+
+bool operator==(const Gpu2DLaunchConfig& a, const Gpu2DLaunchConfig& b) {
+  return a.thread_per_block.x == b.thread_per_block.x &&
+         a.thread_per_block.y == b.thread_per_block.y &&
+         a.thread_per_block.z == b.thread_per_block.z &&
+         a.block_count.x == b.block_count.x &&
+         a.block_count.y == b.block_count.y &&
+         a.block_count.z == b.block_count.z &&
+         a.thread_per_block.x == b.thread_per_block.x &&
+         a.thread_per_block.y == b.thread_per_block.y &&
+         a.thread_per_block.z == b.thread_per_block.z;
+}
+
+TEST_F(GpuLaunchConfigTest, GetGpu2DLaunchConfig) {
+  Gpu2DLaunchConfig cfg;
+  GpuLaunchConfig cfg1d;
+
+// test valid inputs
+#define TEST_LAUNCH_PARAMETER(dimx, dimy)                                      \
+  cfg1d = GetGpuLaunchConfig(bufsize, d);                                      \
+  TF_EXPECT_OK(GpuLaunchKernel(SetOutbufZero, cfg1d.block_count,               \
+                               cfg1d.thread_per_block, 0, d.stream(), cfg1d,   \
+                               outbuf));                                       \
+  CUDA_ASSERT_SUCCESS                                                          \
+  cfg = GetGpu2DLaunchConfig(dimx, dimy, d);                                   \
+  TF_EXPECT_OK(GpuLaunchKernel(Count2D, cfg.block_count, cfg.thread_per_block, \
+                               0, d.stream(), cfg, bufsize, outbuf));          \
+  CUDA_EXPECT_SUCCESS                                                          \
+  copyToHost();                                                                \
+  EXPECT_EQ(dimx* dimy,                                                        \
+            std::accumulate(outbuf_host, outbuf_host + bufsize, 0));           \
+                                                                               \
+  cfg1d = GetGpuLaunchConfig(bufsize, d, SetOutbufZero, 0, 0);                 \
+  TF_EXPECT_OK(GpuLaunchKernel(SetOutbufZero, cfg1d.block_count,               \
+                               cfg1d.thread_per_block, 0, d.stream(), cfg1d,   \
+                               outbuf));                                       \
+  CUDA_ASSERT_SUCCESS                                                          \
+  cfg = GetGpu2DLaunchConfig(dimx, dimy, d, Count2D, 0, 0);                    \
+  TF_EXPECT_OK(GpuLaunchKernel(Count2D, cfg.block_count, cfg.thread_per_block, \
+                               0, d.stream(), cfg, bufsize, outbuf));          \
+  CUDA_EXPECT_SUCCESS                                                          \
+  copyToHost();                                                                \
+  EXPECT_EQ(dimx* dimy, std::accumulate(outbuf_host, outbuf_host + bufsize, 0))
+
+  TEST_LAUNCH_PARAMETER(128, 128);
+  TEST_LAUNCH_PARAMETER(129, 64);
+  TEST_LAUNCH_PARAMETER(511, 2048);
+  TEST_LAUNCH_PARAMETER(512, 512);
+  TEST_LAUNCH_PARAMETER(2048, 1024);
+  TEST_LAUNCH_PARAMETER(2049, 32);
+  TEST_LAUNCH_PARAMETER(8191, 1);
+  TEST_LAUNCH_PARAMETER(8192, 10);
+  TEST_LAUNCH_PARAMETER(123456, 12);
+  TEST_LAUNCH_PARAMETER(1, 1 << 30);
+  TEST_LAUNCH_PARAMETER(1 << 30, 1);
+#undef TEST_LAUNCH_PARAMETER
+}
+
+TEST_F(GpuLaunchConfigTest, GetGpu3DLaunchConfig) {
+  Gpu3DLaunchConfig cfg;
+  GpuLaunchConfig cfg1d;
+
+// test valid inputs
+#define TEST_LAUNCH_PARAMETER(dimx, dimy, dimz)                                \
+  cfg1d = GetGpuLaunchConfig(bufsize, d, SetOutbufZero, 0, 0);                 \
+  TF_EXPECT_OK(GpuLaunchKernel(SetOutbufZero, cfg1d.block_count,               \
+                               cfg1d.thread_per_block, 0, d.stream(), cfg1d,   \
+                               outbuf));                                       \
+  CUDA_ASSERT_SUCCESS                                                          \
+  cfg = GetGpu3DLaunchConfig(dimx, dimy, dimz, d, Count3D, 0, 0);              \
+  TF_EXPECT_OK(GpuLaunchKernel(Count3D, cfg.block_count, cfg.thread_per_block, \
+                               0, d.stream(), cfg, bufsize, outbuf));          \
+  CUDA_EXPECT_SUCCESS                                                          \
+  copyToHost();                                                                \
+  EXPECT_EQ(dimx* dimy* dimz,                                                  \
+            std::accumulate(outbuf_host, outbuf_host + bufsize, 0))
+
+  TEST_LAUNCH_PARAMETER(128, 128, 128);
+  TEST_LAUNCH_PARAMETER(129, 64, 1024);
+  TEST_LAUNCH_PARAMETER(511, 2048, 128);
+  TEST_LAUNCH_PARAMETER(512, 512, 64);
+  TEST_LAUNCH_PARAMETER(2048, 1024, 128);
+  TEST_LAUNCH_PARAMETER(2049, 32, 1024);
+  TEST_LAUNCH_PARAMETER(8191, 1, 1024);
+  TEST_LAUNCH_PARAMETER(8192, 10, 32);
+  TEST_LAUNCH_PARAMETER(123456, 12, 21);
+  TEST_LAUNCH_PARAMETER(1, 1, 1 << 30);
+  TEST_LAUNCH_PARAMETER(1, 1 << 30, 1);
+  TEST_LAUNCH_PARAMETER(1 << 30, 1, 1);
+#undef TEST_LAUNCH_PARAMETER
+}
+
+__global__ void TestGpuAtomicAddDouble(double* ptr, double value) {
+  GpuAtomicAdd(ptr, value);
+}
+
+TEST_F(GpuLaunchConfigTest, GpuAtomicAddDoubleTest) {
+  double* d_ptr;
+#if GOOGLE_CUDA
+  cudaError_t err = cudaMallocManaged(&d_ptr, sizeof(double));
+#else
+  hipError_t err = hipMalloc(&d_ptr, sizeof(double));
+#endif
+  ASSERT_EQ(cudaSuccess, err) << cudaGetErrorString(err);
+
+  double init_val = 10.5;
+#if GOOGLE_CUDA
+  *d_ptr = init_val;
+#else
+  ASSERT_EQ(hipMemcpy(d_ptr, &init_val, sizeof(double), hipMemcpyHostToDevice),
+            cudaSuccess);
+#endif
+
+  TF_EXPECT_OK(GpuLaunchKernel(TestGpuAtomicAddDouble, 1, 1, 0, d.stream(),
+                               d_ptr, 5.25));
+  CUDA_EXPECT_SUCCESS
+
+  double final_val;
+#if GOOGLE_CUDA
+  final_val = *d_ptr;
+#else
+  ASSERT_EQ(hipMemcpy(&final_val, d_ptr, sizeof(double), hipMemcpyDeviceToHost),
+            cudaSuccess);
+#endif
+  EXPECT_DOUBLE_EQ(final_val, 15.75);
+
+  ASSERT_EQ(gpuFree(d_ptr), cudaSuccess);
+}
+
+TEST(CudaDeviceFunctionsTest, ShuffleGetSrcLane) {
+  unsigned* failure_count;
+#if GOOGLE_CUDA
+  ASSERT_EQ(cudaMallocManaged(&failure_count, sizeof(unsigned)), cudaSuccess);
+#else
+  ASSERT_EQ(hipHostMalloc(&failure_count, sizeof(unsigned), 0), cudaSuccess);
+#endif
+  *failure_count = 0;
+  TF_EXPECT_OK(GpuLaunchKernel(GpuShuffleGetSrcLaneTest, 1, TF_RED_WARPSIZE, 0,
+                               nullptr, failure_count));
+  ASSERT_EQ(gpuDeviceSynchronize(), cudaSuccess);
+  ASSERT_EQ(*failure_count, 0);
+  ASSERT_EQ(gpuFree(failure_count), cudaSuccess);
+}
+
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM

tensorflow/core/util/gpu_launch_config.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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_UTIL_GPU_LAUNCH_CONFIG_H_
+#define TENSORFLOW_CORE_UTIL_GPU_LAUNCH_CONFIG_H_
+
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+
+#include <algorithm>
+#include <cstdint>
+#include <limits>
+
+#include "absl/base/casts.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "Eigen/Core"  // from @eigen_archive
+#include "unsupported/Eigen/CXX11/Tensor"  // from @eigen_archive
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/stream_executor.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/util/gpu_cuda_alias.h"
+
+// Usage of GetGpuLaunchConfig, GetGpu2DLaunchConfig, and
+// GetGpu3DLaunchConfig:
+//
+// There are two versions of GetGpuLaunchConfig and GetGpu2DLaunchConfig, one
+// version uses heuristics without any knowledge of the device kernel, the other
+// version uses cudaOccupancyMaxPotentialBlockSize to determine the theoretical
+// launch parameters that maximize occupancy. Currently, only the maximum
+// occupancy version of GetGpu3DLaunchConfig is available.
+//
+// For large number of work elements, the convention is that each kernel would
+// iterate through its assigned range. The return value of GetGpuLaunchConfig
+// is struct GpuLaunchConfig, which contains all the information needed for the
+// kernel launch, including: virtual number of threads, the number of threads
+// per block and number of threads per block used inside <<< >>> of a kernel
+// launch. GetGpu2DLaunchConfig and GetGpu3DLaunchConfig does the same thing
+// as GpuLaunchConfig. The only difference is the dimension. The macros
+// GPU_1D_KERNEL_LOOP and GPU_AXIS_KERNEL_LOOP might be used to do inner loop.
+//
+/* Sample code:
+
+__global__ void MyKernel1D(GpuLaunchConfig config, other_args...) {
+  GPU_1D_KERNEL_LOOP(x, config.virtual_thread_count) {
+    do_your_job_here;
+  }
+}
+
+__global__ void MyKernel2D(Gpu2DLaunchConfig config, other_args...) {
+  GPU_AXIS_KERNEL_LOOP(x, config.virtual_thread_count, x) {
+    GPU_AXIS_KERNEL_LOOP(y, config.virtual_thread_count, y) {
+      do_your_job_here;
+    }
+  }
+}
+
+__global__ void MyKernel3D(Gpu3DLaunchConfig config, other_args...) {
+  GPU_AXIS_KERNEL_LOOP(x, config.virtual_thread_count, x) {
+    GPU_AXIS_KERNEL_LOOP(y, config.virtual_thread_count, y) {
+      GPU_AXIS_KERNEL_LOOP(z, config.virtual_thread_count, z) {
+        do_your_job_here;
+      }
+    }
+  }
+}
+
+void MyDriverFunc(const Eigen::GpuDevice &d) {
+  // use heuristics
+  GpuLaunchConfig cfg1 = GetGpuLaunchConfig(10240, d);
+  MyKernel1D <<<config.block_count,
+                config.thread_per_block, 0, d.stream()>>> (cfg1, other_args...);
+  Gpu2DLaunchConfig cfg2 = GetGpu2DLaunchConfig(10240, 10240, d);
+  MyKernel2D <<<config.block_count,
+                config.thread_per_block, 0, d.stream()>>> (cfg2, other_args...);
+  Gpu3DLaunchConfig cfg3 = GetGpu3DLaunchConfig(4096, 4096, 100, d);
+  MyKernel3D <<<config.block_count,
+                config.thread_per_block, 0, d.stream()>>> (cfg3, other_args...);
+
+  // maximize occupancy
+  GpuLaunchConfig cfg4 = GetGpuLaunchConfig(10240, d, MyKernel1D, 0, 0 );
+  MyKernel1D <<<config.block_count,
+                config.thread_per_block, 0, d.stream()>>> (cfg4, other_args...);
+  Gpu2DLaunchConfig cfg5 = GetGpu2DLaunchConfig(10240, 10240, d,
+                                                  MyKernel1D, 0, 0);
+  MyKernel2D <<<config.block_count,
+                config.thread_per_block, 0, d.stream()>>> (cfg5, other_args...);
+  Gpu3DLaunchConfig cfg6 = GetGpu3DLaunchConfig(4096, 4096, 100, d,
+                                                  MyKernel1D, 0, 0);
+  MyKernel3D <<<config.block_count,
+                config.thread_per_block, 0, d.stream()>>> (cfg6, other_args...);
+}
+
+// See the test for this for more example:
+//
+https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/util/gpu_kernel_helper_test.cu.cc
+
+*/
+
+namespace tensorflow {
+
+using Eigen::numext::div_ceil;
+
+[[deprecated("Use Eigen::numext::div_ceil instead.")]]
+inline int DivUp(int a, int b) {
+  return div_ceil(a, b);
+}
+
+struct GpuLaunchConfig {
+  // Logical number of thread that works on the elements. If each logical
+  // thread works on exactly a single element, this is the same as the working
+  // element count.
+  int virtual_thread_count = -1;
+  // Number of threads per block.
+  int thread_per_block = -1;
+  // Number of blocks for GPU kernel launch.
+  int block_count = -1;
+};
+CREATE_CUDA_TYPE_ALIAS(GpuLaunchConfig, CudaLaunchConfig);
+
+// GPU launch configuration for 64-bit work element counts.
+struct GpuLaunchConfig64 {
+  int64_t virtual_thread_count = -1;
+  int thread_per_block = -1;
+  int block_count = -1;
+};
+
+// Calculate the GPU launch config we should use for a kernel launch.
+// This is assuming the kernel is quite simple and will largely be
+// memory-limited.
+// REQUIRES: work_element_count >= 0.
+inline GpuLaunchConfig GetGpuLaunchConfig(int work_element_count,
+                                          const Eigen::GpuDevice& d) {
+  CHECK_GE(work_element_count, 0);
+  GpuLaunchConfig config;
+  const int virtual_thread_count = work_element_count;
+  const int physical_thread_count = std::min(
+      d.getNumGpuMultiProcessors() * d.maxGpuThreadsPerMultiProcessor(),
+      virtual_thread_count);
+  const int thread_per_block = std::min(1024, d.maxGpuThreadsPerBlock());
+  const int block_count =
+      std::min(div_ceil(physical_thread_count, thread_per_block),
+               d.getNumGpuMultiProcessors());
+
+  config.virtual_thread_count = virtual_thread_count;
+  config.thread_per_block = thread_per_block;
+  config.block_count = block_count;
+  return config;
+}
+#ifndef TENSORFLOW_USE_ROCM
+inline CudaLaunchConfig GetCudaLaunchConfig(int work_element_count,
+                                            const Eigen::GpuDevice& d) {
+  return GetGpuLaunchConfig(work_element_count, d);
+}
+#endif
+
+// Calculate the GPU launch config for 64-bit work element counts.
+inline absl::StatusOr<GpuLaunchConfig64> GetGpuLaunchConfig64(
+    int64_t work_element_count, const Eigen::GpuDevice& d) {
+  if (work_element_count < 0) {
+    return absl::InvalidArgumentError("work_element_count must be >= 0");
+  }
+
+  const int64_t virtual_thread_count = work_element_count;
+  const int64_t physical_thread_count =
+      std::min(static_cast<int64_t>(d.getNumGpuMultiProcessors()) *
+                   d.maxGpuThreadsPerMultiProcessor(),
+               virtual_thread_count);
+  const int64_t thread_per_block =
+      std::min(1024L, static_cast<int64_t>(d.maxGpuThreadsPerBlock()));
+  const int64_t block_count =
+      std::min(div_ceil(physical_thread_count, thread_per_block),
+               static_cast<int64_t>(d.getNumGpuMultiProcessors()));
+
+  if (thread_per_block > std::numeric_limits<int>::max()) {
+    return absl::InternalError("thread_per_block exceeds int limit");
+  }
+  if (block_count > std::numeric_limits<int>::max()) {
+    return absl::InternalError("block_count exceeds int limit");
+  }
+
+  GpuLaunchConfig64 config;
+  config.virtual_thread_count = virtual_thread_count;
+  config.thread_per_block = static_cast<int>(thread_per_block);
+  config.block_count = static_cast<int>(block_count);
+  return config;
+}
+
+// Calculate the GPU launch config we should use for a kernel launch. This
+// variant takes the resource limits of func into account to maximize occupancy.
+// REQUIRES: work_element_count >= 0.
+template <typename DeviceFunc>
+GpuLaunchConfig GetGpuLaunchConfig(int work_element_count,
+                                   const Eigen::GpuDevice& d, DeviceFunc func,
+                                   size_t dynamic_shared_memory_size,
+                                   int block_size_limit) {
+  CHECK_GE(work_element_count, 0);
+  GpuLaunchConfig config;
+  int block_count = 0;
+  int thread_per_block = 0;
+
+#if GOOGLE_CUDA
+  cudaError_t err = cudaOccupancyMaxPotentialBlockSize(
+      &block_count, &thread_per_block, func, dynamic_shared_memory_size,
+      block_size_limit);
+  CHECK_EQ(err, cudaSuccess);
+#elif TENSORFLOW_USE_ROCM
+  hipError_t err = hipOccupancyMaxPotentialBlockSize(
+      &block_count, &thread_per_block, func, dynamic_shared_memory_size,
+      block_size_limit);
+  CHECK_EQ(err, hipSuccess);
+#endif
+
+  block_count =
+      std::min(block_count, div_ceil(work_element_count, thread_per_block));
+
+  config.virtual_thread_count = work_element_count;
+  config.thread_per_block = thread_per_block;
+  config.block_count = block_count;
+  return config;
+}
+CREATE_CUDA_HOST_FUNCTION_ALIAS(GetGpuLaunchConfig, GetCudaLaunchConfig);
+
+// Calculate the GPU launch config for 64-bit work element counts using
+// occupancy API.
+template <typename DeviceFunc>
+absl::StatusOr<GpuLaunchConfig64> GetGpuLaunchConfig64(
+    int64_t work_element_count, const Eigen::GpuDevice& d, DeviceFunc func,
+    size_t dynamic_shared_memory_size, int block_size_limit) {
+  if (work_element_count < 0) {
+    return absl::InvalidArgumentError("work_element_count must be >= 0");
+  }
+  int block_count = 0;
+  int thread_per_block = 0;
+
+#if GOOGLE_CUDA
+  cudaError_t err = cudaOccupancyMaxPotentialBlockSize(
+      &block_count, &thread_per_block, func, dynamic_shared_memory_size,
+      block_size_limit);
+  if (err != cudaSuccess) {
+    return absl::InternalError("cudaOccupancyMaxPotentialBlockSize failed");
+  }
+#elif TENSORFLOW_USE_ROCM
+  hipError_t err = hipOccupancyMaxPotentialBlockSize(
+      &block_count, &thread_per_block, func, dynamic_shared_memory_size,
+      block_size_limit);
+  if (err != hipSuccess) {
+    return absl::InternalError("hipOccupancyMaxPotentialBlockSize failed");
+  }
+#endif
+
+  int64_t block_count_int64 = std::min(
+      static_cast<int64_t>(block_count),
+      div_ceil(work_element_count, static_cast<int64_t>(thread_per_block)));
+  if (block_count_int64 > std::numeric_limits<int>::max()) {
+    return absl::InternalError("block_count exceeds int limit");
+  }
+
+  GpuLaunchConfig64 config;
+  config.virtual_thread_count = work_element_count;
+  config.thread_per_block = thread_per_block;
+  config.block_count = static_cast<int>(block_count_int64);
+  return config;
+}
+
+// Calculate the GPU launch config we should use for a kernel launch. This
+// variant takes the resource limits of func into account to maximize occupancy.
+// The returned launch config has thread_per_block set to fixed_block_size.
+// REQUIRES: work_element_count >= 0.
+template <typename DeviceFunc>
+GpuLaunchConfig GetGpuLaunchConfigFixedBlockSize(
+    int work_element_count, const Eigen::GpuDevice& d, DeviceFunc func,
+    size_t dynamic_shared_memory_size, int fixed_block_size) {
+  CHECK_GE(work_element_count, 0);
+  GpuLaunchConfig config;
+  int block_count = 0;
+
+#if GOOGLE_CUDA
+  cudaError_t err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
+      &block_count, func, fixed_block_size, dynamic_shared_memory_size);
+  CHECK_EQ(err, cudaSuccess);
+#elif TENSORFLOW_USE_ROCM
+  hipError_t err = hipOccupancyMaxActiveBlocksPerMultiprocessor(
+      &block_count, func, fixed_block_size, dynamic_shared_memory_size);
+  CHECK_EQ(err, hipSuccess);
+#endif
+  block_count = std::min(block_count * d.getNumGpuMultiProcessors(),
+                         div_ceil(work_element_count, fixed_block_size));
+
+  config.virtual_thread_count = work_element_count;
+  config.thread_per_block = fixed_block_size;
+  config.block_count = block_count;
+  return config;
+}
+CREATE_CUDA_HOST_FUNCTION_ALIAS(GetGpuLaunchConfigFixedBlockSize,
+                                GetCudaLaunchConfigFixedBlockSize);
+
+struct Gpu2DLaunchConfig {
+  dim3 virtual_thread_count = dim3(0, 0, 0);
+  dim3 thread_per_block = dim3(0, 0, 0);
+  dim3 block_count = dim3(0, 0, 0);
+};
+CREATE_CUDA_TYPE_ALIAS(Gpu2DLaunchConfig, Cuda2DLaunchConfig);
+
+inline Gpu2DLaunchConfig GetGpu2DLaunchConfig(int xdim, int ydim,
+                                              const Eigen::GpuDevice& d) {
+  Gpu2DLaunchConfig config;
+
+  if (xdim <= 0 || ydim <= 0) {
+    return config;
+  }
+
+  const int kThreadsPerBlock = 256;
+  int block_cols = std::min(xdim, kThreadsPerBlock);
+  // ok to round down here and just do more loops in the kernel
+  int block_rows = std::max(kThreadsPerBlock / block_cols, 1);
+
+  const int physical_thread_count =
+      d.getNumGpuMultiProcessors() * d.maxGpuThreadsPerMultiProcessor();
+
+  const int max_blocks = std::max(physical_thread_count / kThreadsPerBlock, 1);
+
+  config.virtual_thread_count = dim3(xdim, ydim, 1);
+  config.thread_per_block = dim3(block_cols, block_rows, 1);
+
+  int grid_x = std::min(div_ceil(xdim, block_cols), max_blocks);
+
+  config.block_count = dim3(
+      grid_x, std::min(max_blocks / grid_x, std::max(ydim / block_rows, 1)), 1);
+  return config;
+}
+#ifndef TENSORFLOW_USE_ROCM
+inline Cuda2DLaunchConfig GetCuda2DLaunchConfig(int xdim, int ydim,
+                                                const Eigen::GpuDevice& d) {
+  return GetGpu2DLaunchConfig(xdim, ydim, d);
+}
+#endif
+
+// Calculate the GPU 2D and 3D launch config we should use for a kernel launch.
+// This variant takes the resource limits of func into account to maximize
+// occupancy.
+using Gpu3DLaunchConfig = Gpu2DLaunchConfig;
+CREATE_CUDA_TYPE_ALIAS(Gpu3DLaunchConfig, Cuda3DLaunchConfig);
+
+template <typename DeviceFunc>
+Gpu3DLaunchConfig GetGpu3DLaunchConfig(int xdim, int ydim, int zdim,
+                                       const Eigen::GpuDevice& d,
+                                       DeviceFunc func,
+                                       size_t dynamic_shared_memory_size,
+                                       int block_size_limit) {
+  Gpu3DLaunchConfig config;
+
+  if (xdim <= 0 || ydim <= 0 || zdim <= 0) {
+    return config;
+  }
+
+  int dev;
+  int xthreadlimit = 0, ythreadlimit = 0, zthreadlimit = 0;
+  int xgridlimit = 0, ygridlimit = 0, zgridlimit = 0;
+#if GOOGLE_CUDA
+  cudaGetDevice(&dev);
+  cudaDeviceGetAttribute(&xthreadlimit, cudaDevAttrMaxBlockDimX, dev);
+  cudaDeviceGetAttribute(&ythreadlimit, cudaDevAttrMaxBlockDimY, dev);
+  cudaDeviceGetAttribute(&zthreadlimit, cudaDevAttrMaxBlockDimZ, dev);
+  cudaDeviceGetAttribute(&xgridlimit, cudaDevAttrMaxGridDimX, dev);
+  cudaDeviceGetAttribute(&ygridlimit, cudaDevAttrMaxGridDimY, dev);
+  cudaDeviceGetAttribute(&zgridlimit, cudaDevAttrMaxGridDimZ, dev);
+#elif TENSORFLOW_USE_ROCM
+  hipError_t err = hipGetDevice(&dev);
+  CHECK_EQ(err, hipSuccess);
+  err =
+      hipDeviceGetAttribute(&xthreadlimit, hipDeviceAttributeMaxBlockDimX, dev);
+  CHECK_EQ(err, hipSuccess);
+  err =
+      hipDeviceGetAttribute(&ythreadlimit, hipDeviceAttributeMaxBlockDimY, dev);
+  CHECK_EQ(err, hipSuccess);
+  err =
+      hipDeviceGetAttribute(&zthreadlimit, hipDeviceAttributeMaxBlockDimZ, dev);
+  CHECK_EQ(err, hipSuccess);
+  err = hipDeviceGetAttribute(&xgridlimit, hipDeviceAttributeMaxGridDimX, dev);
+  CHECK_EQ(err, hipSuccess);
+  err = hipDeviceGetAttribute(&ygridlimit, hipDeviceAttributeMaxGridDimY, dev);
+  CHECK_EQ(err, hipSuccess);
+  err = hipDeviceGetAttribute(&zgridlimit, hipDeviceAttributeMaxGridDimZ, dev);
+  CHECK_EQ(err, hipSuccess);
+#endif
+
+  int block_count = 0;
+  int thread_per_block = 0;
+
+#if GOOGLE_CUDA
+  cudaError_t err = cudaOccupancyMaxPotentialBlockSize(
+      &block_count, &thread_per_block, func, dynamic_shared_memory_size,
+      block_size_limit);
+  CHECK_EQ(err, cudaSuccess);
+#elif TENSORFLOW_USE_ROCM
+  // ROCM TODO re-enable this after hipOccupancyMaxPotentialBlockSize is
+  // implemented
+  // hipError_t err = hipOccupancyMaxPotentialBlockSize(
+  //    &block_count, &thread_per_block, func, dynamic_shared_memory_size,
+  //    block_size_limit);
+  // CHECK_EQ(err, hipSuccess);
+
+  const int physical_thread_count =
+      d.getNumGpuMultiProcessors() * d.maxGpuThreadsPerMultiProcessor();
+  thread_per_block = std::min(1024, d.maxGpuThreadsPerBlock());
+  block_count = std::min(div_ceil(physical_thread_count, thread_per_block),
+                         d.getNumGpuMultiProcessors());
+#endif
+
+  int threadsx = std::min({xdim, thread_per_block, xthreadlimit});
+  int threadsy =
+      std::min({ydim, std::max(thread_per_block / threadsx, 1), ythreadlimit});
+  int threadsz =
+      std::min({zdim, std::max(thread_per_block / (threadsx * threadsy), 1),
+                zthreadlimit});
+
+  int blocksx = std::min({block_count, div_ceil(xdim, threadsx), xgridlimit});
+  int blocksy = std::min(
+      {div_ceil(block_count, blocksx), div_ceil(ydim, threadsy), ygridlimit});
+  int blocksz = std::min({div_ceil(block_count, (blocksx * blocksy)),
+                          div_ceil(zdim, threadsz), zgridlimit});
+
+  config.virtual_thread_count = dim3(xdim, ydim, zdim);
+  config.thread_per_block = dim3(threadsx, threadsy, threadsz);
+  config.block_count = dim3(blocksx, blocksy, blocksz);
+  return config;
+}
+CREATE_CUDA_HOST_FUNCTION_ALIAS(GetGpu3DLaunchConfig, GetCuda3DLaunchConfig);
+
+template <typename DeviceFunc>
+Gpu2DLaunchConfig GetGpu2DLaunchConfig(int xdim, int ydim,
+                                       const Eigen::GpuDevice& d,
+                                       DeviceFunc func,
+                                       size_t dynamic_shared_memory_size,
+                                       int block_size_limit) {
+  return GetGpu3DLaunchConfig(xdim, ydim, 1, d, func,
+                              dynamic_shared_memory_size, block_size_limit);
+}
+CREATE_CUDA_HOST_FUNCTION_ALIAS(GetGpu2DLaunchConfig, GetCuda2DLaunchConfig);
+
+#if GOOGLE_CUDA
+template <typename DeviceFunc>
+Cuda2DLaunchConfig GetCuda2DLaunchConfig(int xdim, int ydim,
+                                         const Eigen::GpuDevice& d,
+                                         DeviceFunc func,
+                                         size_t dynamic_shared_memory_size,
+                                         int block_size_limit) {
+  return GetGpu2DLaunchConfig(xdim, ydim, d, func, dynamic_shared_memory_size,
+                              block_size_limit);
+}
+#endif  // GOOGLE_CUDA
+
+namespace detail {
+template <typename... Ts, size_t... Is>
+std::array<void*, sizeof...(Ts)> GetArrayOfElementPointersImpl(
+    std::tuple<Ts...>* tuple, absl::index_sequence<Is...>) {
+  return {{&std::get<Is>(*tuple)...}};
+}
+// Returns an array of void pointers to the elements of the given tuple.
+template <typename... Ts>
+std::array<void*, sizeof...(Ts)> GetArrayOfElementPointers(
+    std::tuple<Ts...>* tuple) {
+  return GetArrayOfElementPointersImpl(tuple,
+                                       absl::index_sequence_for<Ts...>{});
+}
+
+template <bool...>
+struct BoolPack;
+template <bool... Bs>
+using NoneTrue = std::is_same<BoolPack<Bs..., false>, BoolPack<false, Bs...>>;
+// Returns whether none of the types in Ts is a reference.
+template <typename... Ts>
+constexpr bool NoneIsReference() {
+  return NoneTrue<(std::is_reference<Ts>::value)...>::value;
+}
+}  // namespace detail
+}  // namespace tensorflow
+#endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+#endif  // TENSORFLOW_CORE_UTIL_GPU_LAUNCH_CONFIG_H_

tensorflow/core/util/gpu_solvers.h

@@ -0,0 +1,711 @@
+/* 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_UTIL_GPU_SOLVERS_H_
+#define TENSORFLOW_CORE_UTIL_GPU_SOLVERS_H_
+
+// This header declares the class GpuSolver, which contains wrappers of linear
+// algebra solvers in the cuBlas/cuSolverDN or rocmSolver libraries for use in
+// TensorFlow kernels.
+
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+
+#include <functional>
+#include <vector>
+
+#if GOOGLE_CUDA
+#include "third_party/gpus/cuda/include/cublas_v2.h"
+#include "third_party/gpus/cuda/include/cuda.h"
+#include "third_party/gpus/cuda/include/cusolverDn.h"
+#else
+#include "rocm/include/hip/hip_complex.h"
+#include "rocm/include/rocblas.h"
+#include "rocm/rocm_config.h"
+#if TF_ROCM_VERSION >= 40500
+#include "xla/stream_executor/rocm/hipsolver_wrapper.h"
+#endif
+#include "xla/stream_executor/rocm/rocsolver_wrapper.h"
+#endif
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_reference.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/platform/stream_executor.h"
+
+namespace tensorflow {
+
+#if GOOGLE_CUDA
+// Type traits to get CUDA complex types from std::complex<T>.
+template <typename T>
+struct CUDAComplexT {
+  typedef T type;
+};
+template <>
+struct CUDAComplexT<std::complex<float>> {
+  typedef cuComplex type;
+};
+template <>
+struct CUDAComplexT<std::complex<double>> {
+  typedef cuDoubleComplex type;
+};
+// Converts pointers of std::complex<> to pointers of
+// cuComplex/cuDoubleComplex. No type conversion for non-complex types.
+template <typename T>
+inline const typename CUDAComplexT<T>::type* CUDAComplex(const T* p) {
+  return reinterpret_cast<const typename CUDAComplexT<T>::type*>(p);
+}
+template <typename T>
+inline typename CUDAComplexT<T>::type* CUDAComplex(T* p) {
+  return reinterpret_cast<typename CUDAComplexT<T>::type*>(p);
+}
+
+// Template to give the Cublas adjoint operation for real and complex types.
+template <typename T>
+cublasOperation_t CublasAdjointOp() {
+  return Eigen::NumTraits<T>::IsComplex ? CUBLAS_OP_C : CUBLAS_OP_T;
+}
+#else  // TENSORFLOW_USE_ROCM
+// Type traits to get ROCm complex types from std::complex<T>.
+template <typename T>
+struct ROCmComplexT {
+  typedef T type;
+};
+template <>
+struct ROCmComplexT<std::complex<float>> {
+  typedef rocblas_float_complex type;
+};
+template <>
+struct ROCmComplexT<std::complex<double>> {
+  typedef rocblas_double_complex type;
+};
+// Converts pointers of std::complex<> to pointers of
+// ROCmComplex/ROCmDoubleComplex. No type conversion for non-complex types.
+template <typename T>
+inline const typename ROCmComplexT<T>::type* ROCmComplex(const T* p) {
+  return reinterpret_cast<const typename ROCmComplexT<T>::type*>(p);
+}
+template <typename T>
+inline typename ROCmComplexT<T>::type* ROCmComplex(T* p) {
+  return reinterpret_cast<typename ROCmComplexT<T>::type*>(p);
+}
+
+// Type traits to get HIP complex types from std::complex<>
+
+template <typename T>
+struct HipComplexT {
+  typedef T type;
+};
+
+template <>
+struct HipComplexT<std::complex<float>> {
+  typedef hipFloatComplex type;
+};
+
+template <>
+struct HipComplexT<std::complex<double>> {
+  typedef hipDoubleComplex type;
+};
+
+// Convert pointers of std::complex<> to pointers of
+// hipFloatComplex/hipDoubleComplex. No type conversion for non-complex types.
+template <typename T>
+inline const typename HipComplexT<T>::type* AsHipComplex(const T* p) {
+  return reinterpret_cast<const typename HipComplexT<T>::type*>(p);
+}
+
+template <typename T>
+inline typename HipComplexT<T>::type* AsHipComplex(T* p) {
+  return reinterpret_cast<typename HipComplexT<T>::type*>(p);
+}
+// Template to give the Rocblas adjoint operation for real and complex types.
+template <typename T>
+rocblas_operation RocblasAdjointOp() {
+  return Eigen::NumTraits<T>::IsComplex ? rocblas_operation_conjugate_transpose
+                                        : rocblas_operation_transpose;
+}
+
+#if TF_ROCM_VERSION >= 40500
+using gpuSolverOp_t = hipsolverOperation_t;
+using gpuSolverFill_t = hipsolverFillMode_t;
+using gpuSolverSide_t = hipsolverSideMode_t;
+#else
+using gpuSolverOp_t = rocblas_operation;
+using gpuSolverFill_t = rocblas_fill;
+using gpuSolverSide_t = rocblas_side;
+#endif
+#endif
+
+// Container of LAPACK info data (an array of int) generated on-device by
+// a GpuSolver call. One or more such objects can be passed to
+// GpuSolver::CopyLapackInfoToHostAsync() along with a callback to
+// check the LAPACK info data after the corresponding kernels
+// finish and LAPACK info has been copied from the device to the host.
+class DeviceLapackInfo;
+
+// Host-side copy of LAPACK info.
+class HostLapackInfo;
+
+// The GpuSolver class provides a simplified templated API for the dense linear
+// solvers implemented in cuSolverDN (http://docs.nvidia.com/cuda/cusolver) and
+// cuBlas (http://docs.nvidia.com/cuda/cublas/#blas-like-extension/).
+// An object of this class wraps static cuSolver and cuBlas instances,
+// and will launch Cuda kernels on the stream wrapped by the GPU device
+// in the OpKernelContext provided to the constructor.
+//
+// Notice: All the computational member functions are asynchronous and simply
+// launch one or more Cuda kernels on the Cuda stream wrapped by the GpuSolver
+// object. To check the final status of the kernels run, call
+// CopyLapackInfoToHostAsync() on the GpuSolver object to set a callback that
+// will be invoked with the status of the kernels launched thus far as
+// arguments.
+//
+// Example of an asynchronous TensorFlow kernel using GpuSolver:
+//
+// template <typename Scalar>
+// class SymmetricPositiveDefiniteSolveOpGpu : public AsyncOpKernel {
+//  public:
+//   explicit SymmetricPositiveDefiniteSolveOpGpu(OpKernelConstruction* context)
+//       : AsyncOpKernel(context) { }
+//   void ComputeAsync(OpKernelContext* context, DoneCallback done) final {
+//     // 1. Set up input and output device ptrs. See, e.g.,
+//     // matrix_inverse_op.cc for a full example.
+//     ...
+//
+//     // 2. Initialize the solver object.
+//     std::unique_ptr<GpuSolver> solver(new GpuSolver(context));
+//
+//     // 3. Launch the two compute kernels back to back on the stream without
+//     // synchronizing.
+//     std::vector<DeviceLapackInfo> dev_info;
+//     const int batch_size = 1;
+//     dev_info.push_back(solver->GetDeviceLapackInfo(batch_size, "potrf");
+//     // Compute the Cholesky decomposition of the input matrix.
+//     OP_REQUIRES_OK_ASYNC(context,
+//                          solver->Potrf(uplo, n, dev_matrix_ptrs, n,
+//                                        dev_info.back().mutable_data()),
+//                          done);
+//     dev_info.push_back(solver->GetDeviceLapackInfo(batch_size, "potrs");
+//     // Use the Cholesky decomposition of the input matrix to solve A X = RHS.
+//     OP_REQUIRES_OK_ASYNC(context,
+//                          solver->Potrs(uplo, n, nrhs, dev_matrix_ptrs, n,
+//                                        dev_output_ptrs, ldrhs,
+//                                        dev_info.back().mutable_data()),
+//                          done);
+//
+//     // 4. Check the status after the computation finishes and call done.
+//     solver.CheckLapackInfoAndDeleteSolverAsync(std::move(solver), dev_info,
+//                                                std::move(done));
+//   }
+// };
+
+template <typename Scalar>
+class ScratchSpace;
+
+class GpuSolver {
+ public:
+  // This object stores a pointer to context, which must outlive it.
+  explicit GpuSolver(OpKernelContext* context);
+  virtual ~GpuSolver();
+
+  // Launches a memcpy of solver status data specified by dev_lapack_info from
+  // device to the host, and asynchronously invokes the given callback when the
+  // copy is complete. The first Status argument to the callback will be
+  // Status::OK if all lapack infos retrieved are zero, otherwise an error
+  // status is given. The second argument contains a host-side copy of the
+  // entire set of infos retrieved, and can be used for generating detailed
+  // error messages.
+  // `info_checker_callback` must call the DoneCallback of any asynchronous
+  // OpKernel within which `solver` is used.
+  static void CheckLapackInfoAndDeleteSolverAsync(
+      std::unique_ptr<GpuSolver> solver,
+      const std::vector<DeviceLapackInfo>& dev_lapack_info,
+      std::function<void(const absl::Status&,
+                         const std::vector<HostLapackInfo>&)>
+          info_checker_callback);
+
+  // Simpler version to use if no special error checking / messages are needed
+  // apart from checking that the Status of all calls was Status::OK.
+  // `done` may be nullptr.
+  static void CheckLapackInfoAndDeleteSolverAsync(
+      std::unique_ptr<GpuSolver> solver,
+      const std::vector<DeviceLapackInfo>& dev_lapack_info,
+      AsyncOpKernel::DoneCallback done);
+
+  // Returns a ScratchSpace. The GpuSolver object maintains a TensorReference
+  // to the underlying Tensor to prevent it from being deallocated prematurely.
+  template <typename Scalar>
+  ScratchSpace<Scalar> GetScratchSpace(const TensorShape& shape,
+                                       const std::string& debug_info,
+                                       bool on_host);
+  template <typename Scalar>
+  ScratchSpace<Scalar> GetScratchSpace(int64_t size,
+                                       const std::string& debug_info,
+                                       bool on_host);
+  // Returns a DeviceLapackInfo that will live for the duration of the
+  // GpuSolver object.
+  inline DeviceLapackInfo GetDeviceLapackInfo(int64_t size,
+                                              const std::string& debug_info);
+
+  // Allocates a temporary tensor that will live for the duration of the
+  // GpuSolver object.
+  absl::Status allocate_scoped_tensor(DataType type, const TensorShape& shape,
+                                      Tensor* scoped_tensor);
+  absl::Status forward_input_or_allocate_scoped_tensor(
+      absl::Span<const int> candidate_input_indices, DataType type,
+      const TensorShape& shape, Tensor* input_alias_or_new_scoped_tensor);
+
+  OpKernelContext* context() { return context_; }
+
+#if TENSORFLOW_USE_ROCM
+  // ====================================================================
+  // Wrappers for ROCSolver start here
+  //
+  // The method names below
+  // map to those in ROCSolver/Hipsolver, which follow the naming
+  // convention in LAPACK. See rocm_solvers.cc for a mapping of
+  // GpuSolverMethod to library API
+
+  // LU factorization.
+  // Computes LU factorization with partial pivoting P * A = L * U.
+  template <typename Scalar>
+  Status Getrf(int m, int n, Scalar* dev_A, int lda, int* dev_pivots,
+               int* info);
+
+  // Uses LU factorization to solve A * X = B.
+  template <typename Scalar>
+  Status Getrs(const gpuSolverOp_t trans, int n, int nrhs, Scalar* A, int lda,
+               int* dev_pivots, Scalar* B, int ldb, int* dev_lapack_info);
+
+  template <typename Scalar>
+  Status GetrfBatched(int n, Scalar** dev_A, int lda, int* dev_pivots,
+                      DeviceLapackInfo* info, const int batch_count);
+
+  template <typename Scalar>
+  Status GetrsBatched(const rocblas_operation trans, int n, int nrhs,
+                      Scalar** A, int lda, int* dev_pivots, Scalar** B,
+                      const int ldb, int* lapack_info, const int batch_count);
+
+  // Computes the Cholesky factorization A = L * L^H for a single matrix.
+  template <typename Scalar>
+  Status Potrf(gpuSolverFill_t uplo, int n, Scalar* dev_A, int lda,
+               int* dev_lapack_info);
+
+  // Computes matrix inverses for a batch of small matrices. Uses the outputs
+  // from GetrfBatched. No HipSolver implementation yet
+  template <typename Scalar>
+  Status GetriBatched(int n, const Scalar* const host_a_dev_ptrs[], int lda,
+                      const int* dev_pivots,
+                      const Scalar* const host_a_inverse_dev_ptrs[], int ldainv,
+                      DeviceLapackInfo* dev_lapack_info, int batch_size);
+
+  // Computes matrix inverses for a batch of small matrices with size n < 32.
+  // Returns OkStatus() if the kernel was launched successfully. Uses
+  // GetrfBatched and GetriBatched
+  template <typename Scalar>
+  Status MatInvBatched(int n, const Scalar* const host_a_dev_ptrs[], int lda,
+                       const Scalar* const host_a_inverse_dev_ptrs[],
+                       int ldainv, DeviceLapackInfo* dev_lapack_info,
+                       int batch_size);
+
+  // Cholesky factorization
+  // Computes the Cholesky factorization A = L * L^H for a batch of small
+  // matrices.
+  template <typename Scalar>
+  Status PotrfBatched(gpuSolverFill_t uplo, int n,
+                      const Scalar* const host_a_dev_ptrs[], int lda,
+                      DeviceLapackInfo* dev_lapack_info, int batch_size);
+
+  // See
+  // https://rocblas.readthedocs.io/en/latest/API_Reference_Guide.html#trsm_batched
+  // trsm_batched performs the following batched operation:
+  // op(A_i)*X_i = alpha*B_i or
+  // X_i*op(A_i) = alpha*B_i, for i = 1, ..., batch_count,
+  // where alpha is a scalar, X and B are batched m by n matrices,
+  // A is triangular batched matrix and op(A) is one of
+  // op( A ) = A   or
+  // op( A ) = A^T   or
+  // op( A ) = A^H.
+  // Each matrix X_i is overwritten on B_i for i = 1, ..., batch_count.
+  template <typename Scalar>
+  Status TrsmBatched(rocblas_side side, rocblas_fill uplo,
+                     rocblas_operation trans, rocblas_diagonal diag, int m,
+                     int n, const Scalar* alpha,
+                     const Scalar* const dev_Aarray[], int lda,
+                     Scalar* dev_Barray[], int ldb, int batch_size);
+
+  template <typename Scalar>
+  Status Trsv(rocblas_fill uplo, rocblas_operation trans, rocblas_diagonal diag,
+              int n, const Scalar* A, int lda, Scalar* x, int intcx);
+
+  template <typename Scalar>
+  Status Trsm(rocblas_side side, rocblas_fill uplo, rocblas_operation trans,
+              rocblas_diagonal diag, int m, int n, const Scalar* alpha,
+              const Scalar* A, int lda, Scalar* B, int ldb);
+
+  // Singular value decomposition.
+  // See: https://hipsolver.readthedocs.io/en/latest/api_lapackfunc.html#svds
+  // No GesvdjBatched yet.
+  template <typename Scalar>
+  Status Gesvd(signed char jobu, signed char jobvt, int m, int n, Scalar* dev_A,
+               int lda, Scalar* dev_S, Scalar* dev_U, int ldu, Scalar* dev_VT,
+               int ldvt, int* dev_lapack_info);
+
+  // QR factorization.
+  // Computes QR factorization A = Q * R.
+  template <typename Scalar>
+  Status Geqrf(int m, int n, Scalar* dev_A, int lda, Scalar* dev_tau,
+               int* dev_lapack_info);
+
+  // This function performs the matrix-matrix addition/transposition
+  //   C = alpha * op(A) + beta * op(B).
+  template <typename Scalar>
+  Status Geam(rocblas_operation transa, rocblas_operation transb, int m, int n,
+              const Scalar* alpha, /* host or device pointer */
+              const Scalar* A, int lda,
+              const Scalar* beta, /* host or device pointer */
+              const Scalar* B, int ldb, Scalar* C, int ldc);
+
+  // Overwrite matrix C by product of C and the unitary Householder matrix Q.
+  // The Householder matrix Q is represented by the output from Geqrf in dev_a
+  // and dev_tau.
+  template <typename Scalar>
+  Status Unmqr(gpuSolverSide_t side, gpuSolverOp_t trans, int m, int n, int k,
+               const Scalar* dev_a, int lda, const Scalar* dev_tau,
+               Scalar* dev_c, int ldc, int* dev_lapack_info);
+
+  // Overwrites QR factorization produced by Geqrf by the unitary Householder
+  // matrix Q. On input, the Householder matrix Q is represented by the output
+  // from Geqrf in dev_a and dev_tau. On output, dev_a is overwritten with the
+  // first n columns of Q. Requires m >= n >= 0.
+  template <typename Scalar>
+  Status Ungqr(int m, int n, int k, Scalar* dev_a, int lda,
+               const Scalar* dev_tau, int* dev_lapack_info);
+
+#if TF_ROCM_VERSION >= 40500
+  // Hermitian (Symmetric) Eigen decomposition.
+  template <typename Scalar>
+  Status Heevd(hipsolverEigMode_t jobz, gpuSolverFill_t uplo, int n,
+               Scalar* dev_A, int lda,
+               typename Eigen::NumTraits<Scalar>::Real* dev_W,
+               int* dev_lapack_info);
+#endif
+
+#else  // GOOGLE_CUDA
+  // ====================================================================
+  // Wrappers for cuSolverDN and cuBlas solvers start here.
+  //
+  // Apart from capitalization of the first letter, the method names below
+  // map to those in cuSolverDN and cuBlas, which follow the naming
+  // convention in LAPACK see, e.g.,
+  // http://docs.nvidia.com/cuda/cusolver/#naming-convention
+
+  // This function performs the matrix-matrix addition/transposition
+  //   C = alpha * op(A) + beta * op(B).
+  // Returns OkStatus() if the kernel was launched successfully.  See:
+  // http://docs.nvidia.com/cuda/cublas/index.html#cublas-lt-t-gt-geam
+  // NOTE(ebrevdo): Does not support in-place transpose of non-square
+  // matrices.
+
+  template <typename Scalar>
+  absl::Status Geam(cublasOperation_t transa, cublasOperation_t transb, int m,
+                    int n, const Scalar* alpha, /* host or device pointer */
+                    const Scalar* A, int lda,
+                    const Scalar* beta, /* host or device pointer */
+                    const Scalar* B, int ldb, Scalar* C, int ldc) const;
+
+  // Computes the Cholesky factorization A = L * L^H for a single matrix.
+  // Returns OkStatus() if the kernel was launched successfully. See:
+  // http://docs.nvidia.com/cuda/cusolver/#cuds-lt-t-gt-potrf
+  template <typename Scalar>
+  absl::Status Potrf(cublasFillMode_t uplo, int n, Scalar* dev_A, int lda,
+                     int* dev_lapack_info);
+
+  // Computes the Cholesky factorization A = L * L^H for a batch of small
+  // matrices.
+  // Returns OkStatus() if the kernel was launched successfully. See:
+  // http://docs.nvidia.com/cuda/cusolver/index.html#cuds-lt-t-gt-potrfBatched
+  template <typename Scalar>
+  absl::Status PotrfBatched(cublasFillMode_t uplo, int n,
+                            const Scalar* const host_a_dev_ptrs[], int lda,
+                            DeviceLapackInfo* dev_lapack_info, int batch_size);
+  // LU factorization.
+  // Computes LU factorization with partial pivoting P * A = L * U.
+  // See: http://docs.nvidia.com/cuda/cusolver/#cuds-lt-t-gt-getrf
+  template <typename Scalar>
+  absl::Status Getrf(int m, int n, Scalar* dev_A, int lda, int* dev_pivots,
+                     int* dev_lapack_info);
+
+  // Uses LU factorization to solve A * X = B.
+  // See: http://docs.nvidia.com/cuda/cusolver/#cuds-lt-t-gt-getrs
+  template <typename Scalar>
+  absl::Status Getrs(cublasOperation_t trans, int n, int nrhs, const Scalar* A,
+                     int lda, const int* pivots, Scalar* B, int ldb,
+                     int* dev_lapack_info) const;
+
+  // Computes partially pivoted LU factorizations for a batch of small matrices.
+  // Returns OkStatus() if the kernel was launched successfully. See:
+  // http://docs.nvidia.com/cuda/cublas/index.html#cublas-lt-t-gt-getrfbatched
+  template <typename Scalar>
+  absl::Status GetrfBatched(int n, const Scalar* const host_a_dev_ptrs[],
+                            int lda, int* dev_pivots,
+                            DeviceLapackInfo* dev_lapack_info, int batch_size);
+
+  // Batched linear solver using LU factorization from getrfBatched.
+  // Notice that lapack_info is returned on the host, as opposed to
+  // most of the other functions that return it on the device. See:
+  // http://docs.nvidia.com/cuda/cublas/index.html#cublas-lt-t-gt-getrsbatched
+  template <typename Scalar>
+  absl::Status GetrsBatched(cublasOperation_t trans, int n, int nrhs,
+                            const Scalar* const dev_Aarray[], int lda,
+                            const int* devIpiv,
+                            const Scalar* const dev_Barray[], int ldb,
+                            int* host_lapack_info, int batch_size);
+
+  // Computes matrix inverses for a batch of small matrices. Uses the outputs
+  // from GetrfBatched. Returns OkStatus() if the kernel was launched
+  // successfully. See:
+  // http://docs.nvidia.com/cuda/cublas/index.html#cublas-lt-t-gt-getribatched
+  template <typename Scalar>
+  absl::Status GetriBatched(int n, const Scalar* const host_a_dev_ptrs[],
+                            int lda, const int* dev_pivots,
+                            const Scalar* const host_a_inverse_dev_ptrs[],
+                            int ldainv, DeviceLapackInfo* dev_lapack_info,
+                            int batch_size);
+
+  // Computes matrix inverses for a batch of small matrices with size n < 32.
+  // Returns OkStatus() if the kernel was launched successfully. See:
+  // http://docs.nvidia.com/cuda/cublas/index.html#cublas-lt-t-gt-matinvbatched
+  template <typename Scalar>
+  absl::Status MatInvBatched(int n, const Scalar* const host_a_dev_ptrs[],
+                             int lda,
+                             const Scalar* const host_a_inverse_dev_ptrs[],
+                             int ldainv, DeviceLapackInfo* dev_lapack_info,
+                             int batch_size);
+
+  // QR factorization.
+  // Computes QR factorization A = Q * R.
+  // Returns OkStatus() if the kernel was launched successfully.
+  // See: http://docs.nvidia.com/cuda/cusolver/#cuds-lt-t-gt-geqrf
+  template <typename Scalar>
+  absl::Status Geqrf(int m, int n, Scalar* dev_A, int lda, Scalar* dev_tau,
+                     int* dev_lapack_info);
+
+  // Overwrite matrix C by product of C and the unitary Householder matrix Q.
+  // The Householder matrix Q is represented by the output from Geqrf in dev_a
+  // and dev_tau.
+  // Notice: If Scalar is real, only trans=CUBLAS_OP_N or trans=CUBLAS_OP_T is
+  // supported. If Scalar is complex, trans=CUBLAS_OP_N or trans=CUBLAS_OP_C is
+  // supported.
+  // Returns OkStatus() if the kernel was launched successfully.
+  // See: http://docs.nvidia.com/cuda/cusolver/#cuds-lt-t-gt-ormqr
+  template <typename Scalar>
+  absl::Status Unmqr(cublasSideMode_t side, cublasOperation_t trans, int m,
+                     int n, int k, const Scalar* dev_a, int lda,
+                     const Scalar* dev_tau, Scalar* dev_c, int ldc,
+                     int* dev_lapack_info);
+
+  // Overwrites QR factorization produced by Geqrf by the unitary Householder
+  // matrix Q. On input, the Householder matrix Q is represented by the output
+  // from Geqrf in dev_a and dev_tau. On output, dev_a is overwritten with the
+  // first n columns of Q. Requires m >= n >= 0.
+  // Returns OkStatus() if the kernel was launched successfully.
+  // See: http://docs.nvidia.com/cuda/cusolver/#cuds-lt-t-gt-orgqr
+  template <typename Scalar>
+  absl::Status Ungqr(int m, int n, int k, Scalar* dev_a, int lda,
+                     const Scalar* dev_tau, int* dev_lapack_info);
+
+  // Hermitian (Symmetric) Eigen decomposition.
+  // See: http://docs.nvidia.com/cuda/cusolver/#cuds-lt-t-gt-syevd
+  template <typename Scalar>
+  absl::Status Heevd(cusolverEigMode_t jobz, cublasFillMode_t uplo, int n,
+                     Scalar* dev_A, int lda,
+                     typename Eigen::NumTraits<Scalar>::Real* dev_W,
+                     int* dev_lapack_info);
+
+  // Singular value decomposition.
+  // Returns OkStatus() if the kernel was launched successfully.
+  // TODO(rmlarsen, volunteers): Add support for complex types.
+  // See: http://docs.nvidia.com/cuda/cusolver/#cuds-lt-t-gt-gesvd
+  template <typename Scalar>
+  absl::Status Gesvd(signed char jobu, signed char jobvt, int m, int n,
+                     Scalar* dev_A, int lda, Scalar* dev_S, Scalar* dev_U,
+                     int ldu, Scalar* dev_VT, int ldvt, int* dev_lapack_info);
+  template <typename Scalar>
+  absl::Status GesvdjBatched(cusolverEigMode_t jobz, int m, int n,
+                             Scalar* dev_A, int lda, Scalar* dev_S,
+                             Scalar* dev_U, int ldu, Scalar* dev_V, int ldv,
+                             int* dev_lapack_info, int batch_size);
+
+  // Triangular solve
+  // Returns OkStatus() if the kernel was launched successfully.
+  // See https://docs.nvidia.com/cuda/cublas/index.html#cublas-lt-t-gt-trsm
+  template <typename Scalar>
+  absl::Status Trsm(cublasSideMode_t side, cublasFillMode_t uplo,
+                    cublasOperation_t trans, cublasDiagType_t diag, int m,
+                    int n, const Scalar* alpha, const Scalar* A, int lda,
+                    Scalar* B, int ldb);
+
+  template <typename Scalar>
+  absl::Status Trsv(cublasFillMode_t uplo, cublasOperation_t trans,
+                    cublasDiagType_t diag, int n, const Scalar* A, int lda,
+                    Scalar* x, int intcx);
+
+  // See
+  // https://docs.nvidia.com/cuda/cublas/index.html#cublas-lt-t-gt-trsmbatched
+  template <typename Scalar>
+  absl::Status TrsmBatched(cublasSideMode_t side, cublasFillMode_t uplo,
+                           cublasOperation_t trans, cublasDiagType_t diag,
+                           int m, int n, const Scalar* alpha,
+                           const Scalar* const dev_Aarray[], int lda,
+                           Scalar* dev_Barray[], int ldb, int batch_size);
+#endif
+
+ private:
+  OpKernelContext* context_;  // not owned.
+#if GOOGLE_CUDA
+  cudaStream_t cuda_stream_;
+  cusolverDnHandle_t cusolver_dn_handle_;
+  cublasHandle_t cublas_handle_;
+#else  // TENSORFLOW_USE_ROCM
+  hipStream_t hip_stream_;
+  rocblas_handle rocm_blas_handle_;
+#if TF_ROCM_VERSION >= 40500
+  hipsolverHandle_t hipsolver_handle_;
+#endif
+#endif
+
+  std::vector<TensorReference> scratch_tensor_refs_;
+
+  GpuSolver(const GpuSolver&) = delete;
+  void operator=(const GpuSolver&) = delete;
+};
+
+// Helper class to allocate scratch memory and keep track of debug info.
+// Mostly a thin wrapper around Tensor & allocate_temp.
+template <typename Scalar>
+class ScratchSpace {
+ public:
+  ScratchSpace(OpKernelContext* context, int64_t size, bool on_host)
+      : ScratchSpace(context, TensorShape({size}), "", on_host) {}
+
+  ScratchSpace(OpKernelContext* context, int64_t size,
+               const std::string& debug_info, bool on_host)
+      : ScratchSpace(context, TensorShape({size}), debug_info, on_host) {}
+
+  ScratchSpace(OpKernelContext* context, const TensorShape& shape,
+               const std::string& debug_info, bool on_host)
+      : context_(context), debug_info_(debug_info), on_host_(on_host) {
+    AllocatorAttributes alloc_attr;
+    if (on_host) {
+      // Allocate pinned memory on the host to avoid unnecessary
+      // synchronization.
+      alloc_attr.set_on_host(true);
+      alloc_attr.set_gpu_compatible(true);
+    }
+    TF_CHECK_OK(context->allocate_temp(DataTypeToEnum<Scalar>::value, shape,
+                                       &scratch_tensor_, alloc_attr));
+  }
+
+  virtual ~ScratchSpace() {}
+
+  Scalar* mutable_data() {
+    return scratch_tensor_.template flat<Scalar>().data();
+  }
+  const Scalar* data() const {
+    return scratch_tensor_.template flat<Scalar>().data();
+  }
+  Scalar& operator()(int64_t i) {
+    return scratch_tensor_.template flat<Scalar>()(i);
+  }
+  const Scalar& operator()(int64_t i) const {
+    return scratch_tensor_.template flat<Scalar>()(i);
+  }
+  int64_t bytes() const { return scratch_tensor_.TotalBytes(); }
+  int64_t size() const { return scratch_tensor_.NumElements(); }
+  const std::string& debug_info() const { return debug_info_; }
+
+  Tensor& tensor() { return scratch_tensor_; }
+  const Tensor& tensor() const { return scratch_tensor_; }
+
+  // Returns true if this ScratchSpace is in host memory.
+  bool on_host() const { return on_host_; }
+
+ protected:
+  OpKernelContext* context() const { return context_; }
+
+ private:
+  OpKernelContext* context_;  // not owned
+  const std::string debug_info_;
+  const bool on_host_;
+  Tensor scratch_tensor_;
+};
+
+class HostLapackInfo : public ScratchSpace<int> {
+ public:
+  HostLapackInfo(OpKernelContext* context, int64_t size,
+                 const std::string& debug_info)
+      : ScratchSpace<int>(context, size, debug_info, /* on_host */ true) {}
+};
+
+class DeviceLapackInfo : public ScratchSpace<int> {
+ public:
+  DeviceLapackInfo(OpKernelContext* context, int64_t size,
+                   const std::string& debug_info)
+      : ScratchSpace<int>(context, size, debug_info, /* on_host */ false) {}
+
+  // Allocates a new scratch space on the host and launches a copy of the
+  // contents of *this to the new scratch space. Sets success to true if
+  // the copy kernel was launched successfully.
+  HostLapackInfo CopyToHost(bool* success) const {
+    CHECK(success != nullptr);
+    HostLapackInfo copy(context(), size(), debug_info());
+    auto stream = context()->op_device_context()->stream();
+    stream_executor::DeviceAddressBase wrapped_src(
+        static_cast<void*>(const_cast<int*>(this->data())));
+    *success =
+        stream->Memcpy(copy.mutable_data(), wrapped_src, this->bytes()).ok();
+    return copy;
+  }
+};
+
+template <typename Scalar>
+ScratchSpace<Scalar> GpuSolver::GetScratchSpace(const TensorShape& shape,
+                                                const std::string& debug_info,
+                                                bool on_host) {
+  ScratchSpace<Scalar> new_scratch_space(context_, shape, debug_info, on_host);
+  scratch_tensor_refs_.emplace_back(new_scratch_space.tensor());
+  return std::move(new_scratch_space);
+}
+
+template <typename Scalar>
+ScratchSpace<Scalar> GpuSolver::GetScratchSpace(int64_t size,
+                                                const std::string& debug_info,
+                                                bool on_host) {
+  return GetScratchSpace<Scalar>(TensorShape({size}), debug_info, on_host);
+}
+
+inline DeviceLapackInfo GpuSolver::GetDeviceLapackInfo(
+    int64_t size, const std::string& debug_info) {
+  DeviceLapackInfo new_dev_info(context_, size, debug_info);
+  scratch_tensor_refs_.emplace_back(new_dev_info.tensor());
+  return new_dev_info;
+}
+
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+
+#endif  // TENSORFLOW_CORE_UTIL_GPU_SOLVERS_H_

tensorflow/core/util/guarded_philox_random.cc

@@ -0,0 +1,69 @@
+/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/util/guarded_philox_random.h"
+
+#include "tensorflow/core/lib/random/random.h"
+#include "tensorflow/core/util/determinism.h"
+
+namespace tensorflow {
+
+absl::Status GuardedPhiloxRandom::Init(OpKernelConstruction* context) {
+  // Grab seed Attrs.
+  int64_t seed, seed2;
+  auto status = context->GetAttr("seed", &seed);
+  if (!status.ok()) return status;
+  status = context->GetAttr("seed2", &seed2);
+  if (!status.ok()) return status;
+  if (seed == 0 && seed2 == 0 && OpDeterminismRequired()) {
+    return absl::InvalidArgumentError(
+        "When determinism is enabled, random ops "
+        "must have a seed specified.");
+  }
+
+  // Initialize with the given seeds
+  Init(seed, seed2);
+  return absl::OkStatus();
+}
+
+void GuardedPhiloxRandom::Init(int64_t seed, int64_t seed2) {
+  CHECK(!initialized_);
+  if (seed == 0 && seed2 == 0) {
+    // If both seeds are unspecified, use completely random seeds.
+    seed = random::New64();
+    seed2 = random::New64();
+  }
+  mutex_lock lock(mu_);
+  generator_ = random::PhiloxRandom(seed, seed2);
+  initialized_ = true;
+}
+
+void GuardedPhiloxRandom::Init(random::PhiloxRandom::ResultType counter,
+                               random::PhiloxRandom::Key key) {
+  CHECK(!initialized_);
+  mutex_lock lock(mu_);
+  generator_ = random::PhiloxRandom(counter, key);
+  initialized_ = true;
+}
+
+random::PhiloxRandom GuardedPhiloxRandom::ReserveSamples128(int64_t samples) {
+  CHECK(initialized_);
+  mutex_lock lock(mu_);
+  auto local = generator_;
+  generator_.Skip(samples);
+  return local;
+}
+
+}  // namespace tensorflow

tensorflow/core/util/guarded_philox_random.h

@@ -0,0 +1,83 @@
+/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_UTIL_GUARDED_PHILOX_RANDOM_H_
+#define TENSORFLOW_CORE_UTIL_GUARDED_PHILOX_RANDOM_H_
+
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/lib/random/philox_random.h"
+#include "tensorflow/core/platform/macros.h"
+#include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+
+// A thread safe wrapper around a Philox generator.  Example usage:
+//
+//   GuardedRandomPhilox generator;
+//   generator.Init(context);
+//
+//   // In thread safe code
+//   const int samples = ...;
+//   auto local_generator = generator.ReserveSamples128(samples);
+//   for (int i = 0; i < samples; i++)
+//     Array<uint32, 4> sample = local_generator();
+//     // Use sample
+//   }
+//
+class GuardedPhiloxRandom {
+ public:
+  // Must call Init to finish initialization
+  GuardedPhiloxRandom() : initialized_(false) {}
+
+  // Initialize the generator from attributes "seed" and "seed2".
+  // If both seeds are unspecified, use random seeds.
+  // Must be called exactly once.
+  absl::Status Init(OpKernelConstruction* context);
+
+  // Initialize with given seeds.
+  void Init(int64_t seed, int64_t seed2);
+  void Init(random::PhiloxRandom::ResultType counter,
+            random::PhiloxRandom::Key key);
+
+  // Reserve a certain number of 128-bit samples.
+  // This function is thread safe.  The returned generator is valid for the
+  // given number of samples, and can be used without a lock.
+  random::PhiloxRandom ReserveSamples128(int64_t samples);
+
+  // Reserve a certain number of 32-bit samples.
+  random::PhiloxRandom ReserveSamples32(int64_t samples) {
+    return ReserveSamples128((samples + 3) / 4);
+  }
+
+  // Reserve enough random samples in the generator for the given output count.
+  random::PhiloxRandom ReserveRandomOutputs(int64_t output_count,
+                                            int multiplier) {
+    int64_t conservative_sample_count = output_count * multiplier;
+    return ReserveSamples128(conservative_sample_count);
+  }
+
+ private:
+  mutex mu_;
+  random::PhiloxRandom generator_ TF_GUARDED_BY(mu_);
+  bool initialized_;
+
+  GuardedPhiloxRandom(const GuardedPhiloxRandom&) = delete;
+  void operator=(const GuardedPhiloxRandom&) = delete;
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_UTIL_GUARDED_PHILOX_RANDOM_H_

tensorflow/core/util/image_resizer_state.h

@@ -0,0 +1,267 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// This is a helper struct to package up the input and output
+// parameters of an image resizer (the height, widths, etc.).  To
+// reduce code duplication and ensure consistency across the different
+// resizers, it performs the input validation.
+
+#ifndef TENSORFLOW_CORE_UTIL_IMAGE_RESIZER_STATE_H_
+#define TENSORFLOW_CORE_UTIL_IMAGE_RESIZER_STATE_H_
+
+#define EIGEN_USE_THREADS
+#include <math.h>
+
+#include <algorithm>
+#include <array>
+
+#include "unsupported/Eigen/CXX11/Tensor"  // from @eigen_archive
+#include "tensorflow/core/framework/bounds_check.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/register_types.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.h"
+
+namespace tensorflow {
+
+// CalculateResizeScale determines the float scaling factor.
+inline float CalculateResizeScale(int64_t in_size, int64_t out_size,
+                                  bool align_corners) {
+  return (align_corners && in_size > 1 && out_size > 1)
+             ? (in_size - 1) / static_cast<float>(out_size - 1)
+             : in_size / static_cast<float>(out_size);
+}
+
+// Half pixel scaler scales assuming that the pixel centers are at 0.5, i.e. the
+// floating point coordinates of the top,left pixel is 0.5,0.5.
+struct HalfPixelScaler {
+  HalfPixelScaler() = default;
+  inline float operator()(const int x, const float scale) const {
+    // Note that we subtract 0.5 from the return value, as the existing bilinear
+    // sampling code etc assumes pixels are in the old coordinate system.
+    return (static_cast<float>(x) + 0.5f) * scale - 0.5f;
+  }
+};
+
+// Older incorrect scaling method that causes all resizes to have a slight
+// translation leading to inconsistent results. For example, a flip then a
+// resize gives different results then a resize then a flip.
+struct LegacyScaler {
+  LegacyScaler() = default;
+  inline float operator()(const int x, const float scale) const {
+    return static_cast<float>(x) * scale;
+  }
+};
+
+struct ImageResizerState {
+  explicit ImageResizerState(bool align_corners, bool half_pixel_centers)
+      : align_corners_(align_corners),
+        half_pixel_centers_(half_pixel_centers) {}
+
+  // ValidateAndCalculateOutputSize checks the bounds on the input tensors
+  // and requested size, sets up some of the resizing state such as the
+  // height_scale and width_scale, and calculates the output size.
+  // If any of these operations fails, it sets an error status in
+  // the context, which the caller must check.
+  void ValidateAndCalculateOutputSize(OpKernelContext* context) {
+    OP_REQUIRES(
+        context,
+        !half_pixel_centers_ || (half_pixel_centers_ && !align_corners_),
+        absl::InvalidArgumentError("If half_pixel_centers is True, "
+                                   "align_corners must be False."));
+
+    const TensorShape& input_shape = context->input(0).shape();
+    OP_REQUIRES(context, input_shape.dims() == 4,
+                absl::InvalidArgumentError(absl::StrCat(
+                    "input must be 4-dimensional", input_shape.DebugString())));
+    batch_size = input_shape.dim_size(0);
+    channels = input_shape.dim_size(3);
+    OP_REQUIRES(
+        context, channels > 0,
+        absl::InvalidArgumentError("image must have at least one channel"));
+
+    // Verify and assign `in_height` and `in_width`.
+    OP_REQUIRES(
+        context, input_shape.dim_size(1) > 0 && input_shape.dim_size(2) > 0,
+        absl::InvalidArgumentError("input image must be of non-zero size"));
+    OP_REQUIRES(context,
+                FastBoundsCheck(input_shape.dim_size(1),
+                                std::numeric_limits<int32_t>::max()) &&
+                    FastBoundsCheck(input_shape.dim_size(2),
+                                    std::numeric_limits<int32_t>::max()),
+                absl::InvalidArgumentError(
+                    "input sizes must be between 0 and max int32"));
+    in_height = static_cast<int32_t>(input_shape.dim_size(1));
+    in_width = static_cast<int32_t>(input_shape.dim_size(2));
+
+    // Verify the output tensor's shape.
+    const Tensor& shape_t = context->input(1);
+    OP_REQUIRES(
+        context, shape_t.dims() == 1,
+        absl::InvalidArgumentError(absl::StrCat(
+            "shape_t must be 1-dimensional", shape_t.shape().DebugString())));
+    OP_REQUIRES(
+        context, shape_t.NumElements() == 2,
+        absl::InvalidArgumentError(absl::StrCat(
+            "shape_t must have two elements", shape_t.shape().DebugString())));
+
+    // Verify and assign `out_height` and `out_width`.
+    auto Svec = shape_t.vec<int32_t>();
+    out_height = internal::SubtleMustCopy(Svec(0));
+    out_width = internal::SubtleMustCopy(Svec(1));
+    OP_REQUIRES(
+        context, out_height > 0 && out_width > 0,
+        absl::InvalidArgumentError("output dimensions must be positive"));
+
+    height_scale = CalculateResizeScale(in_height, out_height, align_corners_);
+    width_scale = CalculateResizeScale(in_width, out_width, align_corners_);
+
+    // Guard against overflows
+    OP_REQUIRES(context,
+                ceilf((out_height - 1) * height_scale) <=
+                    static_cast<float>(std::numeric_limits<int64_t>::max()),
+                absl::InvalidArgumentError(
+                    "input image height scale would cause an overflow"));
+    OP_REQUIRES(
+        context,
+        ceilf((out_width - 1) * width_scale) <= static_cast<float>(INT_MAX),
+        absl::InvalidArgumentError(
+            "input image width scale would cause an overflow"));
+  }
+
+  // Calculates all the required variables, and allocates the output.
+  void ValidateAndCreateOutput(OpKernelContext* context) {
+    ValidateAndCalculateOutputSize(context);
+    if (!context->status().ok()) return;
+
+    TensorShape shape;
+    // Guard against shape overflow
+    OP_REQUIRES_OK(context, shape.AddDimWithStatus(batch_size));
+    OP_REQUIRES_OK(context, shape.AddDimWithStatus(out_height));
+    OP_REQUIRES_OK(context, shape.AddDimWithStatus(out_width));
+    OP_REQUIRES_OK(context, shape.AddDimWithStatus(channels));
+
+    OP_REQUIRES_OK(context, context->allocate_output(0, shape, &output));
+  }
+
+  int64_t batch_size;
+  int64_t out_height;
+  int64_t out_width;
+  int64_t in_height;
+  int64_t in_width;
+  int64_t channels;
+  float height_scale;
+  float width_scale;
+  Tensor* output = nullptr;
+
+ private:
+  bool align_corners_;
+  bool half_pixel_centers_;
+};
+
+struct ImageResizerGradientState {
+  explicit ImageResizerGradientState(bool align_corners,
+                                     bool half_pixel_centers)
+      : align_corners_(align_corners),
+        half_pixel_centers_(half_pixel_centers) {}
+
+  void ValidateAndCreateOutput(OpKernelContext* context) {
+    OP_REQUIRES(
+        context,
+        !half_pixel_centers_ || (half_pixel_centers_ && !align_corners_),
+        absl::InvalidArgumentError("If half_pixel_centers is True, "
+                                   "align_corners must be False."));
+
+    const Tensor& input = context->input(0);
+    OP_REQUIRES(
+        context, input.dims() == 4,
+        absl::InvalidArgumentError(absl::StrCat(
+            "input_grad must be 4-dimensional", input.shape().DebugString())));
+
+    // Resizers always produce float images, so input gradient must
+    // always be a float.
+    OP_REQUIRES(context, input.dtype() == DT_FLOAT,
+                absl::InvalidArgumentError(
+                    absl::StrCat("input_grad must be of type float",
+                                 DataTypeString(input.dtype()))));
+
+    batch_size = input.dim_size(0);
+    channels = input.dim_size(3);
+
+    resized_height = input.dim_size(1);
+    resized_width = input.dim_size(2);
+
+    // The following check is also carried out for the forward op. It is added
+    // here to prevent a divide-by-zero exception when either height_scale or
+    // width_scale is being calculated.
+    OP_REQUIRES(
+        context, resized_height > 0 && resized_width > 0,
+        absl::InvalidArgumentError("resized dimensions must be positive"));
+
+    const TensorShape& output_shape = context->input(1).shape();
+    OP_REQUIRES(context, output_shape.dims() == 4,
+                absl::InvalidArgumentError(
+                    absl::StrCat("original_image must be 4-dimensional",
+                                 output_shape.DebugString())));
+    original_height = output_shape.dim_size(1);
+    original_width = output_shape.dim_size(2);
+
+    // The following check is also carried out for the forward op. It is added
+    // here to prevent either height_scale or width_scale from being set to
+    // zero, which would cause a divide-by-zero exception in the deterministic
+    // back-prop path.
+    OP_REQUIRES(
+        context, original_height > 0 && original_width > 0,
+        absl::InvalidArgumentError("original dimensions must be positive"));
+
+    OP_REQUIRES(
+        context,
+        FastBoundsCheck(original_height, std::numeric_limits<int32_t>::max()) &&
+            FastBoundsCheck(original_width,
+                            std::numeric_limits<int32_t>::max()),
+        absl::InvalidArgumentError(
+            "original sizes must be between 0 and max int32"));
+
+    height_scale =
+        CalculateResizeScale(original_height, resized_height, align_corners_);
+    width_scale =
+        CalculateResizeScale(original_width, resized_width, align_corners_);
+
+    OP_REQUIRES_OK(context, context->allocate_output(
+                                0,
+                                TensorShape({batch_size, original_height,
+                                             original_width, channels}),
+                                &output));
+  }
+
+  int64_t batch_size;
+  int64_t channels;
+  int64_t resized_height;
+  int64_t resized_width;
+  int64_t original_height;
+  int64_t original_width;
+  float height_scale;
+  float width_scale;
+  Tensor* output = nullptr;
+
+ private:
+  bool align_corners_;
+  bool half_pixel_centers_;
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_UTIL_IMAGE_RESIZER_STATE_H_

tensorflow/core/util/incremental_barrier.cc

@@ -0,0 +1,65 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/util/incremental_barrier.h"
+
+#include <atomic>
+#include <functional>
+#include <utility>
+
+#include "absl/functional/bind_front.h"
+#include "tensorflow/core/platform/logging.h"
+
+namespace tensorflow {
+
+class InternalIncrementalBarrier {
+ public:
+  explicit InternalIncrementalBarrier(IncrementalBarrier::DoneCallback callback)
+      : left_(1), done_callback_(std::move(callback)) {}
+
+  void operator()() {
+    DCHECK_GE(left_.load(std::memory_order_relaxed), 0);
+
+    if (left_.fetch_sub(1, std::memory_order_acq_rel) - 1 == 0) {
+      IncrementalBarrier::DoneCallback done_callback =
+          std::move(done_callback_);
+      delete this;
+      done_callback();
+    }
+  }
+
+  IncrementalBarrier::BarrierCallback Inc() {
+    left_.fetch_add(1, std::memory_order_acq_rel);
+
+    // std::bind_front is only available ever since C++20.
+    return absl::bind_front(&InternalIncrementalBarrier::operator(), this);
+  }
+
+ private:
+  std::atomic<int> left_;
+  IncrementalBarrier::DoneCallback done_callback_;
+};
+
+IncrementalBarrier::IncrementalBarrier(DoneCallback done_callback)
+    : internal_barrier_(
+          new InternalIncrementalBarrier(std::move(done_callback))) {}
+
+IncrementalBarrier::~IncrementalBarrier() { (*internal_barrier_)(); }
+
+IncrementalBarrier::BarrierCallback IncrementalBarrier::Inc() {
+  return internal_barrier_->Inc();
+}
+
+}  // namespace tensorflow

tensorflow/core/util/incremental_barrier.h

@@ -0,0 +1,81 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_UTIL_INCREMENTAL_BARRIER_H_
+#define TENSORFLOW_CORE_UTIL_INCREMENTAL_BARRIER_H_
+
+#include <atomic>
+#include <functional>
+
+namespace tensorflow {
+
+class InternalIncrementalBarrier;
+
+// BarrierClosure (see
+// https://github.com/chromium/chromium/blob/master/base/barrier_closure.h)
+// executes a callback after it has been invoked |num_closures| times.
+// Plus, `BarrierClosure` is a continuation-passing style abstraction and self-
+// deleting.
+
+// IncrementalBarrier is a convenience class to be used in place of a barrier
+// closure, which is particularly helpful (e.g. simplify code) because callers
+// don't need to calculate the |num_closures| beforehand.
+//
+// Example Usage:
+//   void MakeCalls() {
+//     typedef std::function<void()> Callback;
+//     typedef std::function<void(Callback)> OtherCallback;
+//     Callback done_callback = ...
+//     OtherCallback cb1 = ...
+//     OtherCallback cb2 = ...
+//     std::thread threads[2];
+//     {
+//         IncrementalBarrier barrier(done_callback);
+//         threads[0] = std::thread(cb1(barrier.Inc());
+//         threads[1] = std::thread(cb2(barrier.Inc());
+//         ... at this moment, `barrier` is incremented twice, and then
+//         destructed....
+//     }
+//     threads[0].join();
+//     threads[1].join();
+//   }
+//
+//  `done_callback` will be called when both conditions are true:
+//  1) after `barrier` is destructed.
+//  2) Each `BarrierCallback` returned by `Inc` is called.
+// This class is thread-safe.
+class IncrementalBarrier {
+ public:
+  typedef std::function<void()> DoneCallback;
+  typedef std::function<void()> BarrierCallback;
+  explicit IncrementalBarrier(DoneCallback callback);
+
+  ~IncrementalBarrier();
+
+  // Returns a BarrierCallback (std::function) that individual task call to
+  // signal its completeness.
+  // The returned BarrierCallback outlives this `IncrementalBarrier` instance.
+  // Furthermore, each task should eventually call the returned function, or
+  // else done_callback wouldn't be called.
+  BarrierCallback Inc();
+
+ private:
+  // self-deleting, thereby not owned by 'IncrementalBarrier'.
+  InternalIncrementalBarrier* internal_barrier_;
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_UTIL_INCREMENTAL_BARRIER_H_

tensorflow/core/util/incremental_barrier_test.cc

@@ -0,0 +1,133 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/util/incremental_barrier.h"
+
+#include <atomic>
+
+#include "absl/functional/bind_front.h"
+#include "absl/time/time.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/platform/platform.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/platform/test_benchmark.h"
+#include "tensorflow/core/platform/thread_annotations.h"
+#include "tensorflow/core/platform/threadpool.h"
+
+namespace tensorflow {
+namespace {
+
+// A thread-safe counter class.
+class Counter {
+ public:
+  void Increment() TF_LOCKS_EXCLUDED(mu_) {
+    mutex_lock l(mu_);
+    ++count_;
+  }
+
+  int GetCount() TF_LOCKS_EXCLUDED(mu_) {
+    mutex_lock l(mu_);
+    return count_;
+  }
+
+ private:
+  mutex mu_;
+  int count_ = 0;
+};
+
+TEST(IncrementalBarrierTest, RunInstantlyWhenZeroClosure) {
+  Counter counter;
+  EXPECT_EQ(counter.GetCount(), 0);
+  {
+    IncrementalBarrier::DoneCallback done_callback =
+        absl::bind_front(&Counter::Increment, &counter);
+    IncrementalBarrier barrier(done_callback);
+    EXPECT_EQ(counter.GetCount(), 0);
+  }
+  EXPECT_EQ(counter.GetCount(), 1);
+}
+
+TEST(IncrementalBarrierTest, RunAfterNumClosuresOneNowTwoLater) {
+  Counter counter;
+
+  IncrementalBarrier::BarrierCallback bc1, bc2;
+  {
+    IncrementalBarrier::DoneCallback done_callback =
+        absl::bind_front(&Counter::Increment, &counter);
+    IncrementalBarrier barrier(done_callback);
+
+    CHECK_EQ(counter.GetCount(), 0);
+
+    bc1 = barrier.Inc();
+    bc2 = barrier.Inc();
+
+    IncrementalBarrier::BarrierCallback bc3 = barrier.Inc();
+    bc3();
+
+    CHECK_EQ(counter.GetCount(), 0);
+  }
+
+  CHECK_EQ(counter.GetCount(), 0);
+  bc1();
+  CHECK_EQ(counter.GetCount(), 0);
+  bc2();
+  CHECK_EQ(counter.GetCount(), 1);
+}
+
+TEST(IncrementalBarrierTest, RunAfterNumClosuresConcurrency) {
+  const int num_closure = 100, num_thread = 2;
+  std::atomic<int> schedule_count{0};
+  Counter counter;
+
+  {
+    IncrementalBarrier::DoneCallback done_callback =
+        absl::bind_front(&Counter::Increment, &counter);
+    IncrementalBarrier barrier(done_callback);
+
+    CHECK_EQ(counter.GetCount(), 0);
+
+    tensorflow::thread::ThreadPool pool(tensorflow::Env::Default(),
+                                        "BarrierClosure", num_thread);
+    for (int i = 0; i < num_closure; ++i) {
+      pool.Schedule([&barrier, &schedule_count]() {
+        schedule_count.fetch_add(1);
+        IncrementalBarrier::BarrierCallback bc = barrier.Inc();
+
+        Env::Default()->SleepForMicroseconds(100);
+        bc();
+      });
+    }
+
+    CHECK_EQ(counter.GetCount(), 0);
+  }
+
+  CHECK_EQ(schedule_count.load(std::memory_order_relaxed), 100);
+  CHECK_EQ(counter.GetCount(), 1);
+}
+
+#if defined(PLATFORM_GOOGLE)
+void BM_FunctionInc(benchmark::State& state) {
+  IncrementalBarrier barrier([] {});
+  for (auto _ : state) {
+    barrier.Inc()();
+  }
+}
+
+BENCHMARK(BM_FunctionInc);
+#endif  // PLATFORM_GOOGLE
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/core/util/managed_stack_trace.h

@@ -0,0 +1,121 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_UTIL_MANAGED_STACK_TRACE_H_
+#define TENSORFLOW_CORE_UTIL_MANAGED_STACK_TRACE_H_
+
+#include <functional>
+#include <memory>
+#include <string>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include "absl/strings/match.h"
+#include "absl/strings/str_cat.h"
+#include "absl/types/optional.h"
+#include "tensorflow/core/platform/stack_frame.h"
+
+namespace tensorflow {
+
+// Returns "true" on filenames which should be skipped.
+using StackTraceFilter = std::function<bool(const char*)>;
+
+using SourceLoc = std::pair<std::string, int>;
+
+// Using absl::Hash breaks NVCC under Windows :P
+struct PairHash {
+  template <class T1, class T2>
+  std::size_t operator()(const std::pair<T1, T2>& pair) const {
+    std::size_t h1 = std::hash<T1>()(pair.first);
+    std::size_t h2 = std::hash<T2>()(pair.second);
+    return h1 + 0x9e3779b9 + (h2 << 6) + (h2 >> 2);
+  }
+};
+
+// Maps filename/line_no combination into a stack frame.
+using SourceMap = std::unordered_map<SourceLoc, StackFrame, PairHash>;
+
+using ToStackFramesFunctor = std::vector<StackFrame>(int, const SourceMap&,
+                                                     const StackTraceFilter&,
+                                                     bool, int);
+
+// Returns whether the given frame is internal to TF.
+inline bool IsInternalFrameForFilename(absl::string_view file_name) {
+  // Use a simple heuristic for now.
+  // TODO(cheshire): Build a more sophisticated mechanism, rely on @tf.export.
+  return (absl::StrContains(file_name, "tensorflow/python") ||
+          absl::StrContains(file_name, "tensorflow\\python")) &&
+         !absl::StrContains(file_name, "keras") &&
+         !absl::StrContains(file_name, "test.py");
+}
+
+class CapturedStackTrace {
+ public:
+  virtual ~CapturedStackTrace() = default;
+
+  std::vector<StackFrame> ToStackFrames(const SourceMap& source_map,
+                                        const StackTraceFilter& filtered) {
+    return ToStackFrames(source_map, filtered, /*reverse_traversal=*/false,
+                         /*limit=*/-1);
+  }
+  virtual std::vector<StackFrame> ToStackFrames(
+      const SourceMap& source_map, const StackTraceFilter& filtered,
+      bool reverse_traversal, int limit) const = 0;
+};
+
+// Kept for backwards compatibility with existing users, this simply wraps an
+// underlying stack trace pointer.
+class ManagedStackTrace : public CapturedStackTrace {
+ public:
+  explicit ManagedStackTrace(std::shared_ptr<CapturedStackTrace> trace)
+      : trace_(trace) {}
+
+  ~ManagedStackTrace() override { trace_.reset(); }
+
+  // Returns stack trace as a vector of `StackFrame`s.
+  std::vector<StackFrame> ToStackFrames(const SourceMap& source_map,
+                                        const StackTraceFilter& filtered,
+                                        bool reverse_traversal,
+                                        int limit) const override {
+    return trace_->ToStackFrames(source_map, filtered, reverse_traversal,
+                                 limit);
+  }
+
+ private:
+  std::shared_ptr<CapturedStackTrace> trace_;
+};
+
+// Generates a message with a definition location based on a provided stack
+// trace, or an empty one if the stack trace is empty.
+inline std::string DefinitionLocationMsg(
+    const std::optional<ManagedStackTrace>& stack_trace) {
+  if (stack_trace.has_value()) {
+    std::vector<StackFrame> stack_frames =
+        stack_trace->ToStackFrames({}, IsInternalFrameForFilename,
+                                   /*reverse_traversal=*/true,
+                                   /*limit=*/1);
+    if (!stack_frames.empty()) {
+      const StackFrame& last_frame = stack_frames[0];
+      return absl::StrCat(" (defined @ ", last_frame.file_name, ":",
+                          last_frame.line_number, ")");
+    }
+  }
+  return "";
+}
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_UTIL_MANAGED_STACK_TRACE_H_

tensorflow/core/util/matmul_autotune.cc

@@ -0,0 +1,51 @@
+/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/util/matmul_autotune.h"
+
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/lib/core/stringpiece.h"
+#include "tensorflow/core/util/env_var.h"
+
+namespace tensorflow {
+bool MatmulAutotuneEnable() {
+  bool value;
+  absl::Status status =
+      ReadBoolFromEnvVar("TF_MATMUL_AUTOTUNE_ENABLE", false, &value);
+  if (!status.ok()) {
+    LOG(ERROR) << status.message();
+  }
+  return value;
+}
+
+bool MatmulDoFP32ComputationFP16Input() {
+  bool value;
+  // Feedback from NVIDIA: the "true floating point 16" compute capability is
+  // absent from compute capability SM 5.2. The native 16 bit floating point
+  // computation was introduced in SM 5.3 and higher compute capability. So
+  // for compatibility, set this to be true by default for now.
+  // TODO(yangzihao): In the future, we need to return three possibilities:
+  // user-set-true, user-set-false, user-no-setting. In the calling sites,
+  // check the compatibilities. Note that user-set-false with compute
+  // capability <= 5.2 will cause an error in the later cublasGemmEx() call.
+  absl::Status status =
+      ReadBoolFromEnvVar("TF_FP16_MATMUL_USE_FP32_COMPUTE", true, &value);
+  if (!status.ok()) {
+    LOG(ERROR) << status.message();
+  }
+  return value;
+}
+
+}  // namespace tensorflow

tensorflow/core/util/matmul_autotune.h

@@ -0,0 +1,28 @@
+/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// The utility to check matmul autotune related flags.
+
+#ifndef TENSORFLOW_CORE_UTIL_MATMUL_AUTOTUNE_H_
+#define TENSORFLOW_CORE_UTIL_MATMUL_AUTOTUNE_H_
+
+namespace tensorflow {
+
+bool MatmulAutotuneEnable();
+bool MatmulDoFP32ComputationFP16Input();
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_UTIL_MATMUL_AUTOTUNE_H_

tensorflow/core/util/matmul_bcast.h

@@ -0,0 +1,107 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_UTIL_MATMUL_BCAST_H_
+#define TENSORFLOW_CORE_UTIL_MATMUL_BCAST_H_
+
+#include <algorithm>
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/lib/gtl/inlined_vector.h"
+#include "tensorflow/core/util/bcast.h"
+
+namespace tensorflow {
+
+// Simple wrapper over BCast specialized for MatMul.
+// Provides utilities for broadcasting across batch dimensions for binary
+// MatMul-like operations. If neither argument has batch dimensions (rank <= 2)
+// then no broadcasting is needed and the operation MatMul operation is
+// considered valid.
+class MatMulBCast {
+ public:
+  using Vec = BCast::Vec;
+
+  MatMulBCast(const Vec& x, const Vec& y) {
+    if (std::max(x.size(), y.size()) == 2) return;
+    const Vec x_resized(x.begin(), x.end() - 2);
+    const Vec y_resized(y.begin(), y.end() - 2);
+
+    batch_bcast_ =
+        std::make_unique<BCast>(std::move(x_resized), std::move(y_resized));
+    if (!batch_bcast_->IsValid()) {
+      // Set broadcasting_required_ to true to make IsValid() return false;
+      broadcasting_required_ = true;
+      return;
+    }
+
+    x_batch_size_ = TensorShape(batch_bcast_->x_reshape()).num_elements();
+    y_batch_size_ = TensorShape(batch_bcast_->y_reshape()).num_elements();
+    output_batch_shape_ = TensorShape(batch_bcast_->output_shape());
+    output_batch_size_ = output_batch_shape_.num_elements();
+    broadcasting_required_ =
+        std::min(x_batch_size_, y_batch_size_) != output_batch_size_;
+
+    if (broadcasting_required_) {
+      ComputeBatchIndices(output_batch_size_, batch_bcast_->x_reshape(),
+                          batch_bcast_->x_bcast(), &x_batch_indices_);
+      ComputeBatchIndices(output_batch_size_, batch_bcast_->y_reshape(),
+                          batch_bcast_->y_bcast(), &y_batch_indices_);
+    }
+  }
+
+  bool IsValid() const {
+    return !broadcasting_required_ || (batch_bcast_ && batch_bcast_->IsValid());
+  }
+  bool IsBroadcastingRequired() const { return broadcasting_required_; }
+
+  int64_t output_batch_size() const { return output_batch_size_; }
+  int64_t x_batch_size() const { return x_batch_size_; }
+  int64_t y_batch_size() const { return y_batch_size_; }
+  const TensorShape& output_batch_shape() const { return output_batch_shape_; }
+
+  // Returns the mapping from the flattened output batch indices to x's
+  // flattened batch indices. The result is a vector of length
+  // output_batch_size(). To compute the i'th batch output, a binary matmul-like
+  // operation should use the `x_batch_indices()[i]`th batch index of `x`.
+  // Note: Returns an empty vector if broadcasting is not required. Callers
+  // should only use this when IsBroadcastingRequired() returns true.
+  const std::vector<int64_t>& x_batch_indices() const {
+    return x_batch_indices_;
+  }
+  // Returns the mapping from the flattened output batch indices to y's
+  // flattened batch indices. Similar to x_batch_indices().
+  // Note: Returns an empty vector if broadcasting is not required. Callers
+  // should only use this when IsBroadcastingRequired() returns true.
+  const std::vector<int64_t>& y_batch_indices() const {
+    return y_batch_indices_;
+  }
+
+ private:
+  std::unique_ptr<BCast> batch_bcast_;
+  bool broadcasting_required_ = false;
+  int64_t x_batch_size_ = 1;
+  int64_t y_batch_size_ = 1;
+  TensorShape output_batch_shape_;
+  int64_t output_batch_size_ = 1;
+  std::vector<int64_t> x_batch_indices_;
+  std::vector<int64_t> y_batch_indices_;
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_UTIL_MATMUL_BCAST_H_

tensorflow/core/util/matmul_bcast_test.cc

@@ -0,0 +1,142 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/util/matmul_bcast.h"
+
+#include "tensorflow/core/lib/strings/str_util.h"
+#include "tensorflow/core/lib/strings/strcat.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace {
+
+std::string MatMulBCastToStr(const MatMulBCast& b) {
+  if (!b.IsValid()) {
+    return "invalid";
+  }
+  std::string ret;
+  absl::StrAppend(&ret, "[",
+                  absl::StrJoin(b.output_batch_shape().dim_sizes(), ","), "]");
+  absl::StrAppend(&ret, "[", absl::StrJoin(b.x_batch_indices(), ","), "]");
+  absl::StrAppend(&ret, "[", absl::StrJoin(b.y_batch_indices(), ","), "]");
+  return ret;
+}
+
+TEST(MatMulBCastTest, SimpleBroadcast) {
+  MatMulBCast bcast({1, 5, 3}, {4, 3, 7});
+
+  EXPECT_TRUE(bcast.IsValid());
+  EXPECT_TRUE(bcast.IsBroadcastingRequired());
+
+  EXPECT_EQ(1, bcast.x_batch_size());
+  EXPECT_EQ(4, bcast.y_batch_size());
+  EXPECT_EQ(4, bcast.output_batch_size());
+
+  EXPECT_EQ("[4][0,0,0,0][0,1,2,3]", MatMulBCastToStr(bcast));
+}
+
+TEST(MatMulBCastTest, EmptyBatchBroadcast) {
+  MatMulBCast bcast({5, 3}, {3, 7});
+
+  EXPECT_TRUE(bcast.IsValid());
+  EXPECT_FALSE(bcast.IsBroadcastingRequired());
+
+  EXPECT_EQ(1, bcast.x_batch_size());
+  EXPECT_EQ(1, bcast.y_batch_size());
+  EXPECT_EQ(1, bcast.output_batch_size());
+
+  EXPECT_EQ("[][][]", MatMulBCastToStr(bcast));
+}
+
+TEST(MatMulBCastTest, BroadcastingNotRequired) {
+  MatMulBCast bcast({2, 4, 6, 5, 3}, {2, 4, 6, 3, 7});
+
+  EXPECT_TRUE(bcast.IsValid());
+  EXPECT_FALSE(bcast.IsBroadcastingRequired());
+
+  EXPECT_EQ(48, bcast.x_batch_size());
+  EXPECT_EQ(48, bcast.y_batch_size());
+  EXPECT_EQ(48, bcast.output_batch_size());
+
+  EXPECT_EQ("[2,4,6][][]", MatMulBCastToStr(bcast));
+}
+
+TEST(MatMulBCastTest, EmptyWithNonEmptyBatchBroadcast) {
+  MatMulBCast bcast1({5, 3}, {6, 3, 7});
+
+  EXPECT_TRUE(bcast1.IsValid());
+  EXPECT_TRUE(bcast1.IsBroadcastingRequired());
+
+  EXPECT_EQ(1, bcast1.x_batch_size());
+  EXPECT_EQ(6, bcast1.y_batch_size());
+  EXPECT_EQ(6, bcast1.output_batch_size());
+  EXPECT_EQ("[6][0,0,0,0,0,0][0,1,2,3,4,5]", MatMulBCastToStr(bcast1));
+
+  MatMulBCast bcast2({2, 5, 3}, {3, 7});
+  EXPECT_TRUE(bcast2.IsValid());
+  EXPECT_TRUE(bcast2.IsBroadcastingRequired());
+
+  EXPECT_EQ(2, bcast2.x_batch_size());
+  EXPECT_EQ(1, bcast2.y_batch_size());
+  EXPECT_EQ(2, bcast2.output_batch_size());
+  EXPECT_EQ("[2][0,1][0,0]", MatMulBCastToStr(bcast2));
+}
+
+TEST(MatMulBCastTest, NoBathcDimensions) {
+  MatMulBCast bcast1({3, 3}, {3});
+  EXPECT_TRUE(bcast1.IsValid());
+
+  MatMulBCast bcast2({3}, {3, 3});
+  EXPECT_TRUE(bcast2.IsValid());
+
+  MatMulBCast bcast3({3, 3}, {3, 3});
+  EXPECT_TRUE(bcast3.IsValid());
+}
+
+TEST(MatMulBCastTest, InvalidDimensions) {
+  // Batch dimensions not broadcastable.
+  MatMulBCast bcast3({4, 5, 3}, {2, 3, 7});
+  EXPECT_FALSE(bcast3.IsValid());
+
+  MatMulBCast bcast4({2, 1, 5, 3}, {1, 3, 1, 3, 7});
+  EXPECT_FALSE(bcast4.IsValid());
+}
+
+TEST(MatMulBCastTest, BroadcastBothOperands) {
+  MatMulBCast bcast({3, 1, 5, 3}, {1, 4, 3, 7});
+  EXPECT_TRUE(bcast.IsValid());
+
+  EXPECT_EQ(3, bcast.x_batch_size());
+  EXPECT_EQ(4, bcast.y_batch_size());
+  EXPECT_EQ(12, bcast.output_batch_size());
+
+  EXPECT_EQ("[3,4][0,0,0,0,1,1,1,1,2,2,2,2][0,1,2,3,0,1,2,3,0,1,2,3]",
+            MatMulBCastToStr(bcast));
+}
+
+TEST(MatMulBCastTest, DifferentRanks) {
+  MatMulBCast bcast({3, 1, 5, 3}, {2, 1, 2, 3, 7});
+  EXPECT_TRUE(bcast.IsValid());
+
+  EXPECT_EQ(3, bcast.x_batch_size());
+  EXPECT_EQ(4, bcast.y_batch_size());
+  EXPECT_EQ(12, bcast.output_batch_size());
+
+  EXPECT_EQ("[2,3,2][0,0,1,1,2,2,0,0,1,1,2,2][0,1,0,1,0,1,2,3,2,3,2,3]",
+            MatMulBCastToStr(bcast));
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/core/util/memmapped_file_system.cc

@@ -0,0 +1,328 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/core/util/memmapped_file_system.h"
+
+#include <algorithm>
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/strings/str_util.h"
+#include "tensorflow/core/platform/file_system.h"
+#include "tensorflow/core/platform/protobuf.h"
+#include "tensorflow/core/util/memmapped_file_system.pb.h"
+
+namespace tensorflow {
+
+namespace {
+
+uint64_t DecodeUint64LittleEndian(const uint8_t* buffer) {
+  uint64_t result = 0;
+  for (int i = 0; i < static_cast<int>(sizeof(uint64_t)); ++i) {
+    result |= static_cast<uint64_t>(buffer[i]) << (8 * i);
+  }
+  return result;
+}
+
+}  // namespace
+
+namespace {
+
+class ReadOnlyMemoryRegionFromMemmapped : public ReadOnlyMemoryRegion {
+ public:
+  ReadOnlyMemoryRegionFromMemmapped(const void* data, uint64_t length)
+      : data_(data), length_(length) {}
+  ~ReadOnlyMemoryRegionFromMemmapped() override = default;
+  const void* data() override { return data_; }
+  uint64_t length() override { return length_; }
+
+ private:
+  const void* const data_;
+  const uint64_t length_;
+  // intentionally copyable
+};
+
+class RandomAccessFileFromMemmapped : public RandomAccessFile {
+ public:
+  RandomAccessFileFromMemmapped(const void* data, uint64_t length)
+      : data_(data), length_(length) {}
+
+  ~RandomAccessFileFromMemmapped() override = default;
+
+  absl::Status Name(absl::string_view* result) const override {
+    return absl::UnimplementedError(
+        "RandomAccessFileFromMemmapped does not support Name()");
+  }
+
+  absl::Status Read(uint64_t offset, size_t to_read, absl::string_view* result,
+                    char* scratch) const override {
+    if (offset >= length_) {
+      *result = absl::string_view(scratch, 0);
+      return absl::Status(absl::StatusCode::kOutOfRange, "Read after file end");
+    }
+    const uint64_t region_left =
+        std::min(length_ - offset, static_cast<uint64_t>(to_read));
+    *result = absl::string_view(reinterpret_cast<const char*>(data_) + offset,
+                                region_left);
+    return (region_left == to_read)
+               ? absl::OkStatus()
+               : absl::Status(absl::StatusCode::kOutOfRange,
+                              "Read less bytes than requested");
+  }
+
+ private:
+  const void* const data_;
+  const uint64_t length_;
+  // intentionally copyable
+};
+
+}  // namespace
+
+MemmappedFileSystem::MemmappedFileSystem() = default;
+
+absl::Status MemmappedFileSystem::FileExists(absl::string_view fname) {
+  if (!mapped_memory_) {
+    return absl::FailedPreconditionError("MemmappedEnv is not initialized");
+  }
+  const auto dir_element = directory_.find(fname);
+  if (dir_element != directory_.end()) {
+    return absl::OkStatus();
+  }
+  return absl::NotFoundError(absl::StrCat(fname, " not found"));
+}
+
+absl::Status MemmappedFileSystem::NewRandomAccessFile(
+    const std::string& filename, std::unique_ptr<RandomAccessFile>* result) {
+  if (!mapped_memory_) {
+    return absl::FailedPreconditionError("MemmappedEnv is not initialized");
+  }
+  const auto dir_element = directory_.find(filename);
+  if (dir_element == directory_.end()) {
+    return absl::NotFoundError(
+        absl::StrCat("Region ", filename, " is not found"));
+  }
+  *result = std::make_unique<RandomAccessFileFromMemmapped>(
+      GetMemoryWithOffset(dir_element->second.offset),
+      dir_element->second.length);
+  return absl::OkStatus();
+}
+
+absl::Status MemmappedFileSystem::NewReadOnlyMemoryRegionFromFile(
+    const std::string& filename,
+    std::unique_ptr<ReadOnlyMemoryRegion>* result) {
+  if (!mapped_memory_) {
+    return absl::FailedPreconditionError("MemmappedEnv is not initialized");
+  }
+  const auto dir_element = directory_.find(filename);
+  if (dir_element == directory_.end()) {
+    return absl::NotFoundError(
+        absl::StrCat("Region ", filename, " is not found"));
+  }
+  *result = std::make_unique<ReadOnlyMemoryRegionFromMemmapped>(
+      GetMemoryWithOffset(dir_element->second.offset),
+      dir_element->second.length);
+  return absl::OkStatus();
+}
+
+absl::Status MemmappedFileSystem::GetFileSize(const std::string& filename,
+                                              uint64_t* size) {
+  if (!mapped_memory_) {
+    return absl::FailedPreconditionError("MemmappedEnv is not initialized");
+  }
+  const auto dir_element = directory_.find(filename);
+  if (dir_element == directory_.end()) {
+    return absl::NotFoundError(
+        absl::StrCat("Region ", filename, " is not found"));
+  }
+  *size = dir_element->second.length;
+  return absl::OkStatus();
+}
+
+absl::Status MemmappedFileSystem::Stat(const std::string& fname,
+                                       FileStatistics* stat) {
+  uint64_t size;
+  auto status = GetFileSize(fname, &size);
+  if (status.ok()) {
+    stat->length = size;
+  }
+  return status;
+}
+
+absl::Status MemmappedFileSystem::NewWritableFile(
+    const std::string& filename, std::unique_ptr<WritableFile>* wf) {
+  return absl::UnimplementedError("memmapped format doesn't support writing");
+}
+
+absl::Status MemmappedFileSystem::NewAppendableFile(
+    const std::string& filename, std::unique_ptr<WritableFile>* result) {
+  return absl::UnimplementedError("memmapped format doesn't support writing");
+}
+
+absl::Status MemmappedFileSystem::GetChildren(
+    const std::string& filename, std::vector<std::string>* strings) {
+  return absl::UnimplementedError(
+      "memmapped format doesn't support GetChildren");
+}
+
+absl::Status MemmappedFileSystem::GetMatchingPaths(
+    const std::string& pattern, std::vector<std::string>* results) {
+  return absl::UnimplementedError(
+      "memmapped format doesn't support GetMatchingPaths");
+}
+
+absl::Status MemmappedFileSystem::DeleteFile(const std::string& filename) {
+  return absl::UnimplementedError(
+      "memmapped format doesn't support DeleteFile");
+}
+
+absl::Status MemmappedFileSystem::CreateDir(const std::string& dirname) {
+  return absl::UnimplementedError("memmapped format doesn't support CreateDir");
+}
+
+absl::Status MemmappedFileSystem::DeleteDir(const std::string& dirname) {
+  return absl::UnimplementedError("memmapped format doesn't support DeleteDir");
+}
+
+absl::Status MemmappedFileSystem::RenameFile(const std::string& filename_from,
+                                             const std::string& filename_to) {
+  return absl::UnimplementedError(
+      "memmapped format doesn't support RenameFile");
+}
+
+const void* MemmappedFileSystem::GetMemoryWithOffset(uint64_t offset) const {
+  return reinterpret_cast<const uint8_t*>(mapped_memory_->data()) + offset;
+}
+
+constexpr const char MemmappedFileSystem::kMemmappedPackagePrefix[];
+constexpr const char MemmappedFileSystem::kMemmappedPackageDefaultGraphDef[];
+
+absl::Status MemmappedFileSystem::InitializeFromFile(
+    Env* env, const std::string& filename) {
+  TF_RETURN_IF_ERROR(
+      env->NewReadOnlyMemoryRegionFromFile(filename, &mapped_memory_));
+  directory_.clear();
+  if (mapped_memory_->length() <= sizeof(uint64_t)) {
+    return absl::DataLossError(absl::StrCat(
+        "Corrupted memmapped model file: ", filename, " Invalid package size"));
+  }
+  const auto memory_start =
+      reinterpret_cast<const uint8_t*>(mapped_memory_->data());
+  const uint64_t directory_offset = DecodeUint64LittleEndian(
+      memory_start + mapped_memory_->length() - sizeof(uint64_t));
+  if (directory_offset > mapped_memory_->length() - sizeof(uint64_t)) {
+    return absl::DataLossError(
+        absl::StrCat("Corrupted memmapped model file: ", filename,
+                     " Invalid directory offset"));
+  }
+  MemmappedFileSystemDirectory proto_directory;
+  if (!ParseProtoUnlimited(
+          &proto_directory, memory_start + directory_offset,
+          mapped_memory_->length() - directory_offset - sizeof(uint64_t))) {
+    return absl::DataLossError(
+        absl::StrCat("Corrupted memmapped model file: ", filename,
+                     " Can't parse its internal directory"));
+  }
+
+  // Iterating in reverse order to get lengths of elements;
+  uint64_t prev_element_offset = directory_offset;
+  for (auto element_iter = proto_directory.element().rbegin();
+       element_iter != proto_directory.element().rend(); ++element_iter) {
+    // Check that the element offset and length are in the right range.
+    if (element_iter->offset() >= prev_element_offset ||
+        element_iter->length() > prev_element_offset - element_iter->offset()) {
+      return absl::DataLossError(
+          absl::StrCat("Corrupted memmapped model file: ", filename,
+                       " Invalid offset or length of internal component"));
+    }
+    if (!directory_
+             .insert(std::make_pair(
+                 element_iter->name(),
+                 FileRegion(element_iter->offset(), element_iter->length())))
+             .second) {
+      return absl::DataLossError(absl::StrCat(
+          "Corrupted memmapped model file: ", filename,
+          " Duplicate name of internal component ", element_iter->name()));
+    }
+    prev_element_offset = element_iter->offset();
+  }
+  return absl::OkStatus();
+}
+
+bool MemmappedFileSystem::IsMemmappedPackageFilename(
+    absl::string_view filename) {
+  return absl::StartsWith(filename, kMemmappedPackagePrefix);
+}
+
+namespace {
+bool IsValidRegionChar(char c) {
+  return absl::ascii_isalnum(c) || c == '_' || c == '.';
+}
+}  // namespace
+
+bool MemmappedFileSystem::IsWellFormedMemmappedPackageFilename(
+    const std::string& filename) {
+  if (!IsMemmappedPackageFilename(filename)) {
+    return false;
+  }
+  for (char c :
+       filename.substr(strlen(kMemmappedPackagePrefix),
+                       filename.length() - strlen(kMemmappedPackagePrefix))) {
+    if (!IsValidRegionChar(c)) {
+      return false;
+    }
+  }
+  return true;
+}
+
+MemmappedEnv::MemmappedEnv(Env* env) : EnvWrapper(env) {}
+
+absl::Status MemmappedEnv::GetFileSystemForFile(absl::string_view fname,
+                                                FileSystem** result) {
+  if (MemmappedFileSystem::IsMemmappedPackageFilename(fname)) {
+    if (!memmapped_file_system_) {
+      return absl::FailedPreconditionError(
+          "MemmappedEnv is not initialized from a file.");
+    }
+    *result = memmapped_file_system_.get();
+    return absl::OkStatus();
+  }
+  return EnvWrapper::GetFileSystemForFile(fname, result);
+}
+
+absl::Status MemmappedEnv::GetRegisteredFileSystemSchemes(
+    std::vector<std::string>* schemes) {
+  const auto status = EnvWrapper::GetRegisteredFileSystemSchemes(schemes);
+  if (status.ok()) {
+    schemes->emplace_back(MemmappedFileSystem::kMemmappedPackagePrefix);
+  }
+  return status;
+}
+
+absl::Status MemmappedEnv::InitializeFromFile(
+    const std::string& package_filename) {
+  std::unique_ptr<MemmappedFileSystem> file_system_ptr(new MemmappedFileSystem);
+  const auto status =
+      file_system_ptr->InitializeFromFile(target(), package_filename);
+  if (status.ok()) {
+    memmapped_file_system_ = std::move(file_system_ptr);
+  }
+  return status;
+}
+
+}  // namespace tensorflow

tensorflow/core/util/memmapped_file_system.h

@@ -0,0 +1,136 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_UTIL_MEMMAPPED_FILE_SYSTEM_H_
+#define TENSORFLOW_CORE_UTIL_MEMMAPPED_FILE_SYSTEM_H_
+
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "absl/container/flat_hash_map.h"
+#include "tensorflow/core/platform/env.h"
+
+namespace tensorflow {
+
+// A file system that uses a graph saved in memmapped format by
+// MemmappedEnvWriter as a file system.
+//
+// The format supports saved tensors and protos. Tensors are saved at aligned
+// offsets.
+//
+// Format specification:
+// - last 8 bytes of a package is encoded offset to the directory. The encoding
+// is always little endian, independently from the platform, done by functions
+// EncodeUint64LittleEndian/DecodeUint64LittleEndian
+// - the directory starts from the encoded offset and is saved proto
+// MemmappedFileSystemDirectory with names and offsets to the regions.
+// - at the offsets in the directory the file regions are stored. Tensor regions
+// are aligned such way that when the package mapped to RAM they have the right
+// offset to be used by ImmutableConst operator.
+//
+// Region naming:
+// Region naming is up to the application, all of them starts from
+// kMemmappedPackagePrefix. The default graph usually has name
+// kMemmappedPackageDefaultGraphDef;
+//
+// A "frozen" GraphDef can be converted into this format using
+// tensorflow/contrib/util/convert_graphdef_memmapped_format
+class MemmappedFileSystem : public FileSystem {
+ public:
+  // Memmapped regions use this prefix to distinguish from
+  // the filesystem.
+  static constexpr const char kMemmappedPackagePrefix[] =
+      "memmapped_package://";
+
+  // The default graphdef in the package.
+  static constexpr const char kMemmappedPackageDefaultGraphDef[] =
+      "memmapped_package://.";
+
+  MemmappedFileSystem();
+  ~MemmappedFileSystem() override = default;
+
+  absl::Status FileExists(absl::string_view fname) override;
+  absl::Status NewRandomAccessFile(
+      const std::string& filename,
+      std::unique_ptr<RandomAccessFile>* result) override;
+  absl::Status NewReadOnlyMemoryRegionFromFile(
+      const std::string& filename,
+      std::unique_ptr<ReadOnlyMemoryRegion>* result) override;
+
+  // All these functions return Unimplemented error, the memmapped storage is
+  // read only.
+  absl::Status NewWritableFile(const std::string& fname,
+                               std::unique_ptr<WritableFile>* result) override;
+  absl::Status NewAppendableFile(
+      const std::string& fname, std::unique_ptr<WritableFile>* result) override;
+  absl::Status GetChildren(const std::string& dir,
+                           std::vector<std::string>* r) override;
+  absl::Status GetMatchingPaths(const std::string& pattern,
+                                std::vector<std::string>* results) override;
+  absl::Status DeleteFile(const std::string& f) override;
+  absl::Status CreateDir(const std::string& d) override;
+  absl::Status DeleteDir(const std::string& d) override;
+  absl::Status RenameFile(const std::string& s, const std::string& t) override;
+
+  // These functions are implemented.
+  absl::Status GetFileSize(const std::string& f, uint64_t* s) override;
+  // Currently just returns size.
+  absl::Status Stat(const std::string& fname, FileStatistics* stat) override;
+
+  // Initializes filesystem from a file in memmapped format.
+  absl::Status InitializeFromFile(Env* env, const std::string& filename);
+
+  // Checks if the filename has a correct prefix.
+  static bool IsMemmappedPackageFilename(absl::string_view filename);
+
+  static bool IsWellFormedMemmappedPackageFilename(const std::string& filename);
+
+ private:
+  struct FileRegion {
+    FileRegion(uint64_t o, uint64_t l) : offset(o), length(l) {}
+
+    uint64_t offset;  // Offset from the beginning of the file.
+    uint64_t length;  // Length of the region.
+  };
+
+  using DirectoryType = absl::flat_hash_map<std::string, FileRegion>;
+
+  const void* GetMemoryWithOffset(uint64_t offset) const;
+
+  std::unique_ptr<ReadOnlyMemoryRegion> mapped_memory_;
+  DirectoryType directory_;
+
+  MemmappedFileSystem(const MemmappedFileSystem&) = delete;
+  void operator=(const MemmappedFileSystem&) = delete;
+};
+
+class MemmappedEnv : public EnvWrapper {
+ public:
+  explicit MemmappedEnv(Env* env);
+  ~MemmappedEnv() override = default;
+  absl::Status GetFileSystemForFile(absl::string_view fname,
+                                    FileSystem** result) override;
+  absl::Status GetRegisteredFileSystemSchemes(
+      std::vector<std::string>* schemes) override;
+  absl::Status InitializeFromFile(const std::string& filename);
+
+ protected:
+  std::unique_ptr<MemmappedFileSystem> memmapped_file_system_;
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_UTIL_MEMMAPPED_FILE_SYSTEM_H_

tensorflow/core/util/memmapped_file_system.proto

@@ -0,0 +1,31 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+syntax = "proto3";
+
+package tensorflow;
+
+option cc_enable_arenas = true;
+
+// A message that describes one region of memmapped file.
+message MemmappedFileSystemDirectoryElement {
+  uint64 offset = 1;
+  string name = 2;
+  uint64 length = 3;
+}
+
+// A directory of regions in a memmapped file.
+message MemmappedFileSystemDirectory {
+  repeated MemmappedFileSystemDirectoryElement element = 1;
+}

tensorflow/core/util/memmapped_file_system_test.cc

@@ -0,0 +1,232 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/core/util/memmapped_file_system.h"
+
+#include <cstdint>
+#include <memory>
+
+#include "absl/base/internal/endian.h"
+#include "absl/status/status.h"
+#include "absl/strings/match.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/framework/versions.pb.h"
+#include "tensorflow/core/graph/graph_def_builder.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/lib/io/path.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/file_system.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/util/memmapped_file_system_writer.h"
+
+#ifdef PLATFORM_WINDOWS
+#undef DeleteFile
+#endif
+
+namespace tensorflow {
+
+namespace {
+
+// Names of files in memmapped environment.
+constexpr char kTensor1FileName[] = "memmapped_package://t1";
+constexpr char kTensor2FileName[] = "memmapped_package://t2";
+constexpr char kProtoFileName[] = "memmapped_package://b";
+constexpr int kTestGraphDefVersion = 666;
+
+absl::Status CreateMemmappedFileSystemFile(const std::string& filename,
+                                           bool corrupted,
+                                           Tensor* test_tensor) {
+  Env* env = Env::Default();
+  MemmappedFileSystemWriter writer;
+  TF_RETURN_IF_ERROR(writer.InitializeToFile(env, filename));
+
+  // Try to write a tensor and proto.
+  test::FillFn<float>(test_tensor,
+                      [](int i) { return static_cast<float>(i * i); });
+
+  TF_RETURN_IF_ERROR(writer.SaveTensor(*test_tensor, kTensor1FileName));
+
+  // Create a proto with some fields.
+  GraphDef graph_def;
+  graph_def.mutable_versions()->set_producer(kTestGraphDefVersion);
+  graph_def.mutable_versions()->set_min_consumer(kTestGraphDefVersion);
+  TF_RETURN_IF_ERROR(writer.SaveProtobuf(graph_def, kProtoFileName));
+
+  // Save a tensor after the proto to check that alignment works.
+  test::FillFn<float>(test_tensor,
+                      [](int i) { return static_cast<float>(i) * i * i; });
+  TF_RETURN_IF_ERROR(writer.SaveTensor(*test_tensor, kTensor2FileName));
+
+  if (!corrupted) {
+    // Flush and close the file.
+    TF_RETURN_IF_ERROR(writer.FlushAndClose());
+  }
+  return absl::OkStatus();
+}
+
+TEST(MemmappedFileSystemTest, SimpleTest) {
+  const TensorShape test_tensor_shape = {10, 200};
+  Tensor test_tensor(DT_FLOAT, test_tensor_shape);
+  const std::string dir = testing::TmpDir();
+  const std::string filename = io::JoinPath(dir, "memmapped_env_test");
+  TF_ASSERT_OK(CreateMemmappedFileSystemFile(filename, false, &test_tensor));
+
+  // Check that we can memmap the created file.
+  MemmappedEnv memmapped_env(Env::Default());
+  TF_ASSERT_OK(memmapped_env.InitializeFromFile(filename));
+  // Try to load a proto from the file.
+  GraphDef test_graph_def;
+  TF_EXPECT_OK(
+      ReadBinaryProto(&memmapped_env, kProtoFileName, &test_graph_def));
+  EXPECT_EQ(kTestGraphDefVersion, test_graph_def.versions().producer());
+  EXPECT_EQ(kTestGraphDefVersion, test_graph_def.versions().min_consumer());
+  // Check that we can correctly get a tensor memory.
+  std::unique_ptr<ReadOnlyMemoryRegion> memory_region;
+  TF_ASSERT_OK(memmapped_env.NewReadOnlyMemoryRegionFromFile(kTensor2FileName,
+                                                             &memory_region));
+
+  // The memory region can be bigger but not less than Tensor size.
+  ASSERT_GE(memory_region->length(), test_tensor.TotalBytes());
+  EXPECT_EQ(test_tensor.tensor_data(),
+            absl::string_view(static_cast<const char*>(memory_region->data()),
+                              test_tensor.TotalBytes()));
+  // Check that GetFileSize works.
+  uint64_t file_size = 0;
+  TF_ASSERT_OK(memmapped_env.GetFileSize(kTensor2FileName, &file_size));
+  EXPECT_EQ(test_tensor.TotalBytes(), file_size);
+
+  // Check that Stat works.
+  FileStatistics stat;
+  TF_ASSERT_OK(memmapped_env.Stat(kTensor2FileName, &stat));
+  EXPECT_EQ(test_tensor.TotalBytes(), stat.length);
+
+  // Check that if file not found correct error message returned.
+  EXPECT_EQ(
+      error::NOT_FOUND,
+      memmapped_env.NewReadOnlyMemoryRegionFromFile("bla-bla", &memory_region)
+          .code());
+
+  // Check FileExists.
+  TF_EXPECT_OK(memmapped_env.FileExists(kTensor2FileName));
+  EXPECT_EQ(error::Code::NOT_FOUND,
+            memmapped_env.FileExists("bla-bla-bla").code());
+}
+
+TEST(MemmappedFileSystemTest, NotInitialized) {
+  MemmappedEnv memmapped_env(Env::Default());
+  std::unique_ptr<ReadOnlyMemoryRegion> memory_region;
+  EXPECT_EQ(
+      error::FAILED_PRECONDITION,
+      memmapped_env
+          .NewReadOnlyMemoryRegionFromFile(kTensor1FileName, &memory_region)
+          .code());
+  std::unique_ptr<RandomAccessFile> file;
+  EXPECT_EQ(error::FAILED_PRECONDITION,
+            memmapped_env.NewRandomAccessFile(kProtoFileName, &file).code());
+}
+
+TEST(MemmappedFileSystemTest, Corrupted) {
+  // Create a corrupted file (it is not closed it properly).
+  const TensorShape test_tensor_shape = {100, 200};
+  Tensor test_tensor(DT_FLOAT, test_tensor_shape);
+  const std::string dir = testing::TmpDir();
+  const std::string filename =
+      io::JoinPath(dir, "memmapped_env_corrupted_test");
+  TF_ASSERT_OK(CreateMemmappedFileSystemFile(filename, true, &test_tensor));
+  MemmappedFileSystem memmapped_env;
+  ASSERT_NE(memmapped_env.InitializeFromFile(Env::Default(), filename),
+            absl::OkStatus());
+}
+
+TEST(MemmappedFileSystemTest, CorruptedLength) {
+  const std::string dir = testing::TmpDir();
+  const std::string filename =
+      io::JoinPath(dir, "memmapped_env_corrupted_length_test");
+
+  // Create a properly formatted memmapped file first
+  const TensorShape test_tensor_shape = {10, 20};
+  Tensor test_tensor(DT_FLOAT, test_tensor_shape);
+  TF_ASSERT_OK(CreateMemmappedFileSystemFile(filename, false, &test_tensor));
+
+  // Corrupt the length of an element in the directory to be too large
+  Env* env = Env::Default();
+  std::string file_content;
+  TF_ASSERT_OK(ReadFileToString(env, filename, &file_content));
+
+  // Re-parse the directory to find where it is
+  const uint64_t directory_offset = absl::little_endian::Load64(
+      reinterpret_cast<const uint8_t*>(file_content.data()) +
+      file_content.length() - sizeof(uint64_t));
+
+  MemmappedFileSystemDirectory proto_directory;
+  ASSERT_TRUE(proto_directory.ParseFromString(absl::string_view(
+      file_content.data() + directory_offset,
+      file_content.length() - directory_offset - sizeof(uint64_t))));
+
+  // Modify the first element to have an incredibly large length
+  proto_directory.mutable_element(0)->set_length(file_content.length() * 2);
+
+  std::string new_directory_content = proto_directory.SerializeAsString();
+
+  // Create a new corrupted file
+  std::unique_ptr<WritableFile> corrupted_file;
+  TF_ASSERT_OK(env->NewWritableFile(filename, &corrupted_file));
+  TF_ASSERT_OK(corrupted_file->Append(
+      absl::string_view(file_content.data(), directory_offset)));
+  TF_ASSERT_OK(corrupted_file->Append(new_directory_content));
+
+  // Write offset in little endian format
+  uint8_t offset_buffer[sizeof(uint64_t)];
+  absl::little_endian::Store64(offset_buffer, directory_offset);
+  TF_ASSERT_OK(corrupted_file->Append(absl::string_view(
+      reinterpret_cast<char*>(offset_buffer), sizeof(uint64_t))));
+  TF_ASSERT_OK(corrupted_file->Close());
+
+  MemmappedFileSystem memmapped_env;
+  auto status = memmapped_env.InitializeFromFile(Env::Default(), filename);
+  EXPECT_EQ(status.code(), absl::StatusCode::kDataLoss);
+  EXPECT_TRUE(absl::StrContains(
+      status.message(), "Invalid offset or length of internal component"));
+}
+
+TEST(MemmappedFileSystemTest, ProxyToDefault) {
+  MemmappedEnv memmapped_env(Env::Default());
+  const std::string dir = testing::TmpDir();
+  const std::string filename = io::JoinPath(dir, "test_file");
+  // Check that we can create write and read ordinary file.
+  std::unique_ptr<WritableFile> writable_file_temp;
+  TF_ASSERT_OK(memmapped_env.NewAppendableFile(filename, &writable_file_temp));
+  // Making sure to clean up after the test finishes.
+  const auto adh = [&memmapped_env, &filename](WritableFile* f) {
+    delete f;
+    TF_CHECK_OK(memmapped_env.DeleteFile(filename));
+  };
+  std::unique_ptr<WritableFile, decltype(adh)> writable_file(
+      writable_file_temp.release(), adh);
+  const std::string test_string = "bla-bla-bla";
+  TF_ASSERT_OK(writable_file->Append(test_string));
+  TF_ASSERT_OK(writable_file->Close());
+  uint64_t file_length = 0;
+  TF_EXPECT_OK(memmapped_env.GetFileSize(filename, &file_length));
+  EXPECT_EQ(test_string.length(), file_length);
+  FileStatistics stat;
+  TF_EXPECT_OK(memmapped_env.Stat(filename, &stat));
+  EXPECT_EQ(test_string.length(), stat.length);
+  std::unique_ptr<RandomAccessFile> random_access_file;
+  TF_ASSERT_OK(
+      memmapped_env.NewRandomAccessFile(filename, &random_access_file));
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/core/util/memmapped_file_system_writer.cc

@@ -0,0 +1,136 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/core/util/memmapped_file_system_writer.h"
+
+#include <algorithm>
+
+namespace tensorflow {
+
+absl::Status MemmappedFileSystemWriter::InitializeToFile(
+    Env* env, const std::string& filename) {
+  auto status = env->NewWritableFile(filename, &output_file_);
+  if (status.ok()) {
+    output_file_offset_ = 0;
+  }
+  return status;
+}
+
+absl::Status MemmappedFileSystemWriter::SaveTensor(
+    const Tensor& tensor, const std::string& element_name) {
+  if (!output_file_) {
+    return absl::FailedPreconditionError(
+        "MemmappedEnvWritter: saving tensor into not opened file");
+  }
+  if (!MemmappedFileSystem::IsWellFormedMemmappedPackageFilename(
+          element_name)) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "MemmappedEnvWritter: element_name is invalid: must have memmapped ",
+        "package prefix ", MemmappedFileSystem::kMemmappedPackagePrefix,
+        " and include [A-Za-z0-9_.]"));
+  }
+  const auto tensor_data = tensor.tensor_data();
+  if (tensor_data.empty()) {
+    return absl::InvalidArgumentError(
+        "MemmappedEnvWritter: saving tensor with 0 size");
+  }
+  // Adds pad for correct alignment after memmapping.
+  TF_RETURN_IF_ERROR(AdjustAlignment(Allocator::kAllocatorAlignment));
+  AddToDirectoryElement(element_name, tensor_data.size());
+  const auto result = output_file_->Append(tensor_data);
+  if (result.ok()) {
+    output_file_offset_ += tensor_data.size();
+  }
+  return result;
+}
+
+absl::Status MemmappedFileSystemWriter::SaveProtobuf(
+    const protobuf::MessageLite& message, const std::string& element_name) {
+  if (!output_file_) {
+    return absl::FailedPreconditionError(
+        "MemmappedEnvWritter: saving protobuf into not opened file");
+  }
+  if (!MemmappedFileSystem::IsWellFormedMemmappedPackageFilename(
+          element_name)) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "MemmappedEnvWritter: element_name is invalid: must have memmapped "
+        "package prefix ",
+        MemmappedFileSystem::kMemmappedPackagePrefix,
+        " and include [A-Za-z0-9_.]"));
+  }
+  const std::string encoded = message.SerializeAsString();
+  AddToDirectoryElement(element_name, encoded.size());
+  const auto res = output_file_->Append(encoded);
+  if (res.ok()) {
+    output_file_offset_ += encoded.size();
+  }
+  return res;
+}
+
+namespace {
+
+absl::string_view EncodeUint64LittleEndian(uint64_t val, char* output_buffer) {
+  for (unsigned int i = 0; i < sizeof(uint64_t); ++i) {
+    output_buffer[i] = (val >> i * 8);
+  }
+  return {output_buffer, sizeof(uint64_t)};
+}
+
+}  // namespace
+
+absl::Status MemmappedFileSystemWriter::FlushAndClose() {
+  if (!output_file_) {
+    return absl::FailedPreconditionError(
+        "MemmappedEnvWritter: flushing into not opened file");
+  }
+  const std::string dir = directory_.SerializeAsString();
+  TF_RETURN_IF_ERROR(output_file_->Append(dir));
+
+  // Write the directory offset.
+  char buffer[sizeof(uint64_t)];
+  TF_RETURN_IF_ERROR(output_file_->Append(
+      EncodeUint64LittleEndian(output_file_offset_, buffer)));
+
+  // Flush and close the file.
+  TF_RETURN_IF_ERROR(output_file_->Flush());
+  TF_RETURN_IF_ERROR(output_file_->Close());
+  output_file_.reset();
+  return absl::OkStatus();
+}
+
+absl::Status MemmappedFileSystemWriter::AdjustAlignment(uint64_t alignment) {
+  const uint64_t alignment_rest = output_file_offset_ % alignment;
+  const uint64_t to_write_for_alignment =
+      (alignment_rest == 0) ? 0 : alignment - (output_file_offset_ % alignment);
+  static constexpr uint64_t kFillerBufferSize = 16;
+  const char kFillerBuffer[kFillerBufferSize] = {};
+  for (uint64_t rest = to_write_for_alignment; rest > 0;) {
+    absl::string_view sp(kFillerBuffer, std::min(rest, kFillerBufferSize));
+    TF_RETURN_IF_ERROR(output_file_->Append(sp));
+    rest -= sp.size();
+    output_file_offset_ += sp.size();
+  }
+  return absl::OkStatus();
+}
+
+void MemmappedFileSystemWriter::AddToDirectoryElement(const std::string& name,
+                                                      uint64_t length) {
+  MemmappedFileSystemDirectoryElement* new_directory_element =
+      directory_.add_element();
+  new_directory_element->set_offset(output_file_offset_);
+  new_directory_element->set_name(name);
+  new_directory_element->set_length(length);
+}
+
+}  // namespace tensorflow

tensorflow/core/util/memmapped_file_system_writer.h

@@ -0,0 +1,55 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_CORE_UTIL_MEMMAPPED_FILE_SYSTEM_WRITER_H_
+#define TENSORFLOW_CORE_UTIL_MEMMAPPED_FILE_SYSTEM_WRITER_H_
+
+#include <memory>
+#include <vector>
+
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/util/memmapped_file_system.h"
+#include "tensorflow/core/util/memmapped_file_system.pb.h"
+
+namespace tensorflow {
+
+// A class for saving into the memmapped format that can be read by
+// MemmappedFileSystem.
+class MemmappedFileSystemWriter {
+ public:
+  MemmappedFileSystemWriter() = default;
+  ~MemmappedFileSystemWriter() = default;
+  absl::Status InitializeToFile(Env* env, const std::string& filename);
+  absl::Status SaveTensor(const Tensor& tensor,
+                          const std::string& element_name);
+  absl::Status SaveProtobuf(const protobuf::MessageLite& message,
+                            const std::string& element_name);
+  // Writes out the directory of regions and closes the output file.
+  absl::Status FlushAndClose();
+
+ private:
+  absl::Status AdjustAlignment(uint64_t alignment);
+  void AddToDirectoryElement(const std::string& element_name, uint64_t length);
+  MemmappedFileSystemDirectory directory_;
+  // The current offset in the file, to support alignment.
+  uint64_t output_file_offset_ = 0;
+  std::unique_ptr<WritableFile> output_file_;
+  MemmappedFileSystemWriter(const MemmappedFileSystemWriter&) = delete;
+  void operator=(const MemmappedFileSystemWriter&) = delete;
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_UTIL_MEMMAPPED_FILE_SYSTEM_WRITER_H_

tensorflow/core/util/mirror_pad_mode.cc

@@ -0,0 +1,44 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/util/mirror_pad_mode.h"
+
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/framework/node_def_util.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/core/stringpiece.h"
+
+namespace tensorflow {
+
+absl::Status GetNodeAttr(const NodeDef& node_def, absl::string_view attr_name,
+                         MirrorPadMode* value) {
+  std::string str_value;
+  TF_RETURN_IF_ERROR(GetNodeAttr(node_def, attr_name, &str_value));
+  if (str_value == "REFLECT") {
+    *value = MirrorPadMode::REFLECT;
+  } else if (str_value == "SYMMETRIC") {
+    *value = MirrorPadMode::SYMMETRIC;
+  } else {
+    return absl::NotFoundError(
+        absl::StrCat(str_value, " is not an allowed padding mode."));
+  }
+  return absl::OkStatus();
+}
+
+std::string GetMirrorPadModeAttrString() {
+  return "mode: {'REFLECT', 'SYMMETRIC'}";
+}
+
+}  // end namespace tensorflow

tensorflow/core/util/mirror_pad_mode.h

@@ -0,0 +1,53 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_UTIL_MIRROR_PAD_MODE_H_
+#define TENSORFLOW_CORE_UTIL_MIRROR_PAD_MODE_H_
+
+// This file contains helper routines to deal with padding in various ops and
+// kernels.
+
+#include <string>
+
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/lib/core/stringpiece.h"
+
+namespace tensorflow {
+
+// REFLECT: Border elements are not mirrored to padded regions.
+// SYMMETRIC: Border elements are mirrored to padded regions.
+//
+// For example, if two elements are padded to the right of an array [1, 2, 3],
+// then the result is [1, 2, 3, 2, 1] for REFLECT mode, and is [1, 2, 3, 3, 2]
+// for SYMMETRIC mode.
+enum class MirrorPadMode {
+  REFLECT = 1,
+  SYMMETRIC = 2,
+};
+
+// Return the string containing the list of valid padding modes, that can be
+// used as an Attr() in REGISTER_OP.
+std::string GetMirrorPadModeAttrString();
+
+// Forward declaration to avoid including core/framework/graph.proto.
+class NodeDef;
+
+// Specialization to parse an attribute directly into a MirrorPadMode enum.
+absl::Status GetNodeAttr(const NodeDef& node_def, absl::string_view attr_name,
+                         MirrorPadMode* value);
+
+}  // end namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_UTIL_MIRROR_PAD_MODE_H_

tensorflow/core/util/mkl_heuristics.h

@@ -0,0 +1,194 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// This file contains heuristics data and methods that are used to
+// decide whether to rewrite node to use oneDNN kernels
+
+#ifndef TENSORFLOW_CORE_UTIL_MKL_HEURISTICS_H_
+#define TENSORFLOW_CORE_UTIL_MKL_HEURISTICS_H_
+#ifdef INTEL_MKL
+
+#include <vector>
+
+#include "tensorflow/core/framework/node_def_util.h"
+#include "tensorflow/core/graph/graph.h"
+#include "tsl/platform/cpu_info.h"
+
+namespace tensorflow {
+
+struct RewriteThreshold {
+  std::string op;
+  int cpu_family;
+  int cpu_model_num;
+  // The model that is used to decide whether it is worth
+  // accelerating operations using oneDNN is:
+  //
+  // threshold = thread_synchronisation * thread_num + framework_tax
+  //
+  // This finds threshold when framework overhead and thread synchronisations
+  // are amortized with amount of computation that has to be performed.
+  // If we are below this threshold then we will not rewrite the operation to
+  // to be run using oneDNN primitive.
+  struct PerformanceParameters {
+    double thread_sync_cost;
+    double framework_cost;
+  } params;
+};
+
+// Table storing thread synchronization and framework overhead costs on each CPU
+// architecture for each oneNN-eligible operation. Our heuristics use these
+// costs to determine whether we should rewrite the operation to use oneDNN.
+static const RewriteThreshold rewrite_thresholds[] = {
+#ifdef DNNL_AARCH64_USE_ACL
+    {"Conv2D", 0x41, 0xd40, {0.9349, 22.603}},
+    {"_FusedConv2D", 0x41, 0xd40, {0.9349, 22.603}},
+    {"FusedBatchNormV3", 0x41, 0xd40, {0.3223, -0.8822}},
+    {"Sigmoid", 0x41, 0xd40, {0.0, 0.064736}},
+#endif  // DNNL_AARCH64_USE_ACL
+    {"", 0x0, 0x0, {0, 0}}};
+
+static double FindRewriteThreshold(const string node_name, int threads) {
+  int cpu_family_ = tsl::port::CPUFamily();
+  int cpu_model_num_ = tsl::port::CPUModelNum();
+
+  if (threads == 0) {
+    // if we do not have information how many threads are used
+    // to parallelise operation we revert to the old behaviour
+    return 0;
+  }
+
+  for (const RewriteThreshold* i = rewrite_thresholds;
+       i->op != "" && threads > 0; i++) {
+    if (node_name == i->op && cpu_family_ == i->cpu_family &&
+        cpu_model_num_ == i->cpu_model_num) {
+      return i->params.thread_sync_cost * threads + i->params.framework_cost;
+    }
+  }
+
+  return 0;
+}
+
+static double CalculateNodeMFlops(const AttrSlice& attrs,
+                                  const string node_name) {
+  // Check if we can obtained dimensions for this node.
+  std::vector<const TensorShapeProto*> shape_attrs;
+  if (!TryGetNodeAttr(attrs, "_input_shapes", &shape_attrs)) {
+    // We can't obtain shape so we will revert to default behaviour
+    // to rewrite node.
+    return -1;
+  }
+
+  if ((node_name == "Conv2D" || node_name == "_FusedConv2D") &&
+      shape_attrs.size() == 2) {
+    TensorShape input_shape, filter_shape;
+    if (TensorShape::BuildTensorShape(*shape_attrs[0], &input_shape) !=
+        tsl::OkStatus()) {
+      return -1;
+    }
+    if (TensorShape::BuildTensorShape(*shape_attrs[1], &filter_shape) !=
+        tsl::OkStatus()) {
+      return -1;
+    }
+
+    // MFLOPS = N * H * W * C * FH * FW * FC / 1e6.
+    return input_shape.dim_size(0) * input_shape.dim_size(1) *
+           input_shape.dim_size(2) * input_shape.dim_size(3) *
+           filter_shape.dim_size(0) * filter_shape.dim_size(1) *
+           filter_shape.dim_size(3) / (double)1e6;
+  } else if ((node_name == "FusedBatchNormV3" || node_name == "Sigmoid") &&
+             shape_attrs.size() >= 1) {
+    TensorShape input_shape;
+    if (TensorShape::BuildTensorShape(*shape_attrs[0], &input_shape) !=
+        tsl::OkStatus()) {
+      return -1;
+    }
+    return input_shape.dim_size(0) * input_shape.dim_size(1) *
+           input_shape.dim_size(2) * input_shape.dim_size(3) / (double)1e6;
+  }
+
+  return -1;
+}
+
+// MatMulHeuristic returns true to rewrite the node with oneDNN
+// false to execute the node in Eigen
+static bool MatMulHeuristic(const Node* n) {
+  // Run heuristic only if CPU is ARM_NEOVERSE_V1
+  if (!tsl::port::TestAarch64CPU(tsl::port::Aarch64CPU::ARM_NEOVERSE_V1)) {
+    return true;
+  }
+  // Check if we can obtain dimensions for this node.
+  std::vector<const TensorShapeProto*> shape_attrs;
+  if (!TryGetNodeAttr(n->attrs(), "_input_shapes", &shape_attrs)) {
+    // We can't obtain shape so we will revert to default behaviour
+    // to rewrite node.
+    return true;
+  }
+
+  if ((n->type_string() == "MatMul" || n->type_string() == "_FusedMatMul")) {
+    TensorShape lhs_shape, rhs_shape;
+    if (TensorShape::BuildTensorShape(*shape_attrs[0], &lhs_shape) !=
+        tsl::OkStatus()) {
+      return true;
+    }
+    if (TensorShape::BuildTensorShape(*shape_attrs[1], &rhs_shape) !=
+        tsl::OkStatus()) {
+      return true;
+    }
+
+    auto M = lhs_shape.dim_size(0);
+    auto K = lhs_shape.dim_size(1);
+    auto N = rhs_shape.dim_size(1);
+    auto ops = M * N * K;
+    std::array<int, 3> n_threshold = {7560, 250, 1536};
+    std::array<int, 2> m_threshold = {378, 80};
+    std::array<int, 2> ops_threshold = {5242880, 1090519040};
+
+    if (N <= n_threshold.at(0)) {
+      if (ops <= ops_threshold.at(0)) {
+        if (M <= m_threshold.at(0)) {
+          return false;
+        } else {
+          if (N <= n_threshold.at(1)) {
+            return false;
+          } else {
+            return true;
+          }
+        }
+      } else {
+        if (M <= m_threshold.at(1)) {
+          if (N <= n_threshold.at(2)) {
+            return true;
+          } else {
+            return false;
+          }
+        } else {
+          if (ops <= ops_threshold.at(1)) {
+            return true;
+          } else {
+            return false;
+          }
+        }
+      }
+    } else {
+      return false;
+    }
+  }
+  return true;
+}
+
+}  // namespace tensorflow
+
+#endif  // INTEL_MKL
+#endif  // TENSORFLOW_CORE_UTIL_MKL_HEURISTICS_H_

tensorflow/core/util/mkl_heuristics_test.cc

@@ -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.
+==============================================================================*/
+
+#ifdef INTEL_MKL
+#define EIGEN_USE_THREADS
+
+#include "tensorflow/core/util/mkl_heuristics.h"
+
+#include "tensorflow/core/kernels/ops_testutil.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+
+namespace {
+
+TEST(MklHeuristicsTest, MklCalculateMFlops) {
+  int batch = 8;
+  int width = 32;
+  int height = 32;
+  int in_depth = 3;
+
+  int filter_h = 3;
+  int filter_w = 3;
+  int out_depth = 64;
+
+  // Test calculation for number of MFLOPs for convolution
+  AttrValue attr_input_shape;
+  TensorShapeProto* proto = attr_input_shape.mutable_list()->add_shape();
+  proto->add_dim()->set_size(batch);
+  proto->add_dim()->set_size(width);
+  proto->add_dim()->set_size(height);
+  proto->add_dim()->set_size(in_depth);
+  proto = attr_input_shape.mutable_list()->add_shape();
+  proto->add_dim()->set_size(filter_h);
+  proto->add_dim()->set_size(filter_w);
+  proto->add_dim()->set_size(in_depth);
+  proto->add_dim()->set_size(out_depth);
+
+  NodeDef ndef;
+
+  // If node doesn't have any _input_shapes it should return -1
+  double calculated_empty_mflops =
+      CalculateNodeMFlops(AttrSlice(ndef), "Conv2D");
+  EXPECT_EQ(calculated_empty_mflops, -1);
+
+  (*ndef.mutable_attr())["_input_shapes"] = attr_input_shape;
+
+  double conv_calculated_mflops =
+      CalculateNodeMFlops(AttrSlice(ndef), "Conv2D");
+  double expected_conv_mflops = batch * width * height * in_depth * filter_h *
+                                filter_w * out_depth / static_cast<double>(1e6);
+  EXPECT_EQ(conv_calculated_mflops, expected_conv_mflops);
+
+  // We should get the same calculation for fused convolution too
+  double fused_calculated_mflops =
+      CalculateNodeMFlops(AttrSlice(ndef), "_FusedConv2D");
+  EXPECT_EQ(conv_calculated_mflops, expected_conv_mflops);
+
+  // Finally calculate for sigmoid number of MFLOPS
+  double sigmoid_calculated_mflops =
+      CalculateNodeMFlops(AttrSlice(ndef), "Sigmoid");
+  double expected_sigmoid_mflops =
+      batch * width * height * in_depth / static_cast<double>(1e6);
+  EXPECT_EQ(sigmoid_calculated_mflops, expected_sigmoid_mflops);
+}
+
+#ifdef DNNL_AARCH64_USE_ACL
+TEST(MklHeuristicsTest, MklThresholds) {
+  int cpu_family = tsl::port::CPUFamily();
+  int cpu_model_num = tsl::port::CPUModelNum();
+
+  int neoverse_v1_family = 0x41;
+  int neoverse_v1_model = 0xd40;
+
+  string op_type = "Conv2D";
+
+  if (neoverse_v1_family == cpu_family && neoverse_v1_model == cpu_model_num) {
+    double thread_sync_cost = -1;
+    double framework_cost = -1;
+    for (const RewriteThreshold* i = rewrite_thresholds; i->op != ""; i++) {
+      if (i->op == op_type) {
+        thread_sync_cost = i->params.thread_sync_cost;
+        framework_cost = i->params.framework_cost;
+        break;
+      }
+    }
+
+    EXPECT_NE(thread_sync_cost, -1);
+    EXPECT_NE(thread_sync_cost, -1);
+
+    int no_threads = 0;
+    double calculated_threshold_zero_threads =
+        FindRewriteThreshold(op_type, no_threads);
+    EXPECT_EQ(calculated_threshold_zero_threads, 0);
+
+    int threads = 8;
+    double calculated_threshold = FindRewriteThreshold(op_type, threads);
+    double expected_threshold = threads * thread_sync_cost + framework_cost;
+    EXPECT_EQ(expected_threshold, calculated_threshold);
+  }
+}
+#endif  // DNNL_AARCG64_USE_ACL
+
+}  // namespace
+}  // namespace tensorflow
+
+#endif  // INTEL_MKL

tensorflow/core/util/mkl_util_test.cc

@@ -0,0 +1,94 @@
+/* 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.
+==============================================================================*/
+
+#ifdef INTEL_MKL
+
+#include "tensorflow/core/util/mkl_util.h"
+
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace {
+
+TEST(MklUtilTest, MklDnnTfShape) {
+  auto cpu_engine = engine(engine::kind::cpu, 0);
+  MklDnnData<float> a(&cpu_engine);
+
+  const int N = 1, C = 2, H = 3, W = 4;
+  memory::dims a_dims = {N, C, H, W};
+  MklDnnShape a_mkldnn_shape;
+  a_mkldnn_shape.SetMklTensor(true);
+  // Create TF layout in NCHW.
+  a_mkldnn_shape.SetTfLayout(a_dims.size(), a_dims,
+                             MklTensorFormat::FORMAT_NCHW);
+  TensorShape a_tf_shape_nchw({N, C, H, W});
+  TensorShape a_tf_shape_nhwc({N, H, W, C});
+  TensorShape a_mkldnn_tf_shape = a_mkldnn_shape.GetTfShape();
+  // Check that returned shape is in NCHW format.
+  EXPECT_EQ(a_tf_shape_nchw, a_mkldnn_tf_shape);
+  EXPECT_NE(a_tf_shape_nhwc, a_mkldnn_tf_shape);
+
+  memory::dims b_dims = {N, C, H, W};
+  MklDnnShape b_mkldnn_shape;
+  b_mkldnn_shape.SetMklTensor(true);
+  // Create TF layout in NHWC.
+  b_mkldnn_shape.SetTfLayout(b_dims.size(), b_dims,
+                             MklTensorFormat::FORMAT_NHWC);
+  TensorShape b_tf_shape_nhwc({N, H, W, C});
+  TensorShape b_tf_shape_nchw({N, C, H, W});
+  TensorShape b_mkldnn_tf_shape = b_mkldnn_shape.GetTfShape();
+  // Check that returned shape is in NHWC format.
+  EXPECT_EQ(b_tf_shape_nhwc, b_mkldnn_tf_shape);
+  EXPECT_NE(b_tf_shape_nchw, b_mkldnn_tf_shape);
+}
+
+TEST(MklUtilTest, LRUCacheTest) {
+  // The cached objects are of type int*
+  size_t capacity = 100;
+  size_t num_objects = capacity + 10;
+  LRUCache<int> lru_cache(capacity);
+
+  // Test SetOp: be able to set more ops than the capacity
+  for (int k = 0; k < num_objects; k++) {
+    lru_cache.SetOp(std::to_string(k), new int(k));
+  }
+
+  // Test GetOp and capacity:
+  // Least recently accessed objects should not be in cache any more.
+  for (int k = 0; k < num_objects - capacity; ++k) {
+    EXPECT_EQ(nullptr, lru_cache.GetOp(std::to_string(k)));
+  }
+
+  // Test GetOp and capacity:
+  // Most recently accessed objects should still be in cache.
+  for (int k = num_objects - capacity; k < num_objects; ++k) {
+    int* int_ptr = lru_cache.GetOp(std::to_string(k));
+    EXPECT_NE(nullptr, int_ptr);
+    EXPECT_EQ(*int_ptr, k);
+  }
+
+  // Clean up the cache
+  lru_cache.Clear();
+
+  // After clean up, there should be no cached object.
+  for (int k = 0; k < num_objects; ++k) {
+    EXPECT_EQ(nullptr, lru_cache.GetOp(std::to_string(k)));
+  }
+}
+
+}  // namespace
+}  // namespace tensorflow
+
+#endif  // INTEL_MKL