Added support of C++20 for OSS TF for Linux platform

PiperOrigin-RevId: 613370839

.bazelrc

@@ -380,8 +380,8 @@ build:android --cxxopt=-std=c++17
 build:android --host_cxxopt=-std=c++17
 build:ios --cxxopt=-std=c++17
 build:ios --host_cxxopt=-std=c++17
-build:linux --cxxopt=-std=c++17
-build:linux --host_cxxopt=-std=c++17
+build:linux --cxxopt=-std=c++20
+build:linux --host_cxxopt=-std=c++20
 build:macos --cxxopt=-std=c++17
 build:macos --host_cxxopt=-std=c++17
 build:windows --cxxopt=/std:c++17

ci/official/containers/linux_arm64/Dockerfile

@@ -1,5 +1,5 @@
 ################################################################################
-FROM ubuntu:20.04 as builder
+FROM ubuntu:22.04 as builder
 ################################################################################
 
 # Install devtoolset build dependencies

tensorflow/cc/saved_model/image_format/internal_api.cc

@@ -120,11 +120,6 @@ absl::StatusOr<std::tuple<std::string, bool>> WriteSavedModelToString(
 #if !IS_OSS
 // TODO(b/311769337): Define the function unconditionally after tf oss
 // dependency is updated to protobuf v22.x.
-absl::StatusOr<std::tuple<absl::Cord, bool>> WriteSavedModelToCord(
-    SavedModel* saved_model_proto) {
-  tools::proto_splitter::SavedModelSplitter splitter(saved_model_proto);
-  return splitter.WriteToCord();
-}
 #endif
 
 absl::Status WriteSavedModel(SavedModel* saved_model_proto,

tensorflow/compiler/aot/tests/BUILD

@@ -97,8 +97,10 @@ genrule(
     testonly = 1,
     outs = [
         "test_graph_tfadd.pb",
-        "test_graph_tfadd_with_ckpt.ckpt",
-        "test_graph_tfadd_with_ckpt.pb",
+        # copybara:uncomment_begin
+        # "test_graph_tfadd_with_ckpt.ckpt",
+        # "test_graph_tfadd_with_ckpt.pb",
+        # copybara:uncomment_end
         "test_graph_tfadd_with_ckpt_saver.ckpt",
         "test_graph_tfadd_with_ckpt_saver.pb",
         "test_graph_tfadd_with_ckpt_saver.saver",
@@ -152,19 +154,21 @@ tfcompile_test_dep_configs = [
                 "no_mac",  # TODO(b/228273415)
             ],
         ),
-        tf_library(
-            name = "test_graph_tfadd_with_ckpt" + suffix,
-            testonly = 1,
-            config = "test_graph_tfadd_with_ckpt.config.pbtxt",
-            cpp_class = "AddWithCkptComp",
-            freeze_checkpoint = "test_graph_tfadd_with_ckpt.ckpt",
-            graph = "test_graph_tfadd_with_ckpt.pb",
-            mlir_components = mlir_component,
-            tags = [
-                "manual",
-                "no_mac",  # TODO(b/228273415)
-            ],
-        ),
+        # copybara:uncomment_begin
+        # tf_library(
+        # name = "test_graph_tfadd_with_ckpt" + suffix,
+        # testonly = 1,
+        # config = "test_graph_tfadd_with_ckpt.config.pbtxt",
+        # cpp_class = "AddWithCkptComp",
+        # freeze_checkpoint = "test_graph_tfadd_with_ckpt.ckpt",
+        # graph = "test_graph_tfadd_with_ckpt.pb",
+        # mlir_components = mlir_component,
+        # tags = [
+        # "manual",
+        # "no_mac",  # TODO(b/228273415)
+        # ],
+        # ),
+        # copybara:uncomment_end
         tf_library(
             name = "test_graph_tfadd_with_ckpt_saver" + suffix,
             testonly = 1,

tensorflow/compiler/aot/tests/make_test_graphs.py

@@ -49,6 +49,7 @@ def tfadd(_):
   math_ops.add(x, y, name='x_y_sum')
 
 
+# copybara:comment_begin
 def tfadd_with_ckpt(out_dir):
   x = array_ops.placeholder(dtypes.int32, name='x_hold')
   y = variable_v1.VariableV1(constant_op.constant([0]), name='y_saved')
@@ -64,6 +65,9 @@ def tfadd_with_ckpt(out_dir):
     saver.save(sess, ckpt)
 
 
+# copybara:comment_end
+
+
 def tfadd_with_ckpt_saver(out_dir):
   x = array_ops.placeholder(dtypes.int32, name='x_hold')
   y = variable_v1.VariableV1(constant_op.constant([0]), name='y_saved')
@@ -87,7 +91,8 @@ def tfassert_eq(_):
   x = array_ops.placeholder(dtypes.int32, name='x_hold')
   y = array_ops.placeholder(dtypes.int32, name='y_hold')
   control_flow_assert.Assert(
-      math_ops.equal(x, y), ['Expected x == y.'], name='assert_eq')
+      math_ops.equal(x, y), ['Expected x == y.'], name='assert_eq'
+  )
   math_ops.add(x, math_ops.negative(y), name='x_y_diff')
 
 
@@ -120,7 +125,6 @@ def tfmatmulandadd(_):
 
 
 def tffunction(_):
-
   @function.Defun(dtypes.int32, dtypes.int32)
   def test_func(a, b):
     return a + b
@@ -201,7 +205,7 @@ def write_graph(build_graph, out_dir):
 def main(_):
   control_flow_util.enable_control_flow_v2()
   write_graph(tfadd, FLAGS.out_dir)
-  write_graph(tfadd_with_ckpt, FLAGS.out_dir)
+  write_graph(tfadd_with_ckpt, FLAGS.out_dir)  # copybara:comment
   write_graph(tfadd_with_ckpt_saver, FLAGS.out_dir)
   write_graph(tfassert_eq, FLAGS.out_dir)
   write_graph(tfcond, FLAGS.out_dir)
@@ -223,6 +227,7 @@ def main(_):
       '--out_dir',
       type=str,
       default='',
-      help='Output directory for graphs, checkpoints and savers.')
+      help='Output directory for graphs, checkpoints and savers.',
+  )
   FLAGS, unparsed = parser.parse_known_args()
   app.run(main=main, argv=[sys.argv[0]] + unparsed)

tensorflow/compiler/mlir/tensorflow/utils/eval_util.cc

@@ -102,7 +102,7 @@ mlir::LogicalResult EvaluateOperation(
     RETURN_FAILURE_IF_ERROR(status);
   }
 
-  VLOG(1) << "Start to evaluate node: " << *node_def;
+  VLOG(1) << "Start to evaluate node: " << (*node_def).name();
 
   // Adds inputs to the TF operation.
   for (const auto operand : operands) {

tensorflow/compiler/mlir/tfrt/transforms/ifrt/tf2hlo.cc

@@ -90,7 +90,7 @@ absl::StatusOr<tensorflow::tpu::TPUCompileMetadataProto> GetCompileMetadata(
         absl::StrCat("Missing ", kMetadataTextAttrName));
   }
 
-  VLOG(3) << "TpuCompileMetadata before shape is populated " << metadata;
+  // VLOG(3) << "TpuCompileMetadata before shape is populated " << metadata;
   if (metadata.num_replicas() < 1 || metadata.num_cores_per_replica() < 1) {
     return absl::InternalError(
         absl::StrCat("Number of replicas ", metadata.num_replicas(),
@@ -177,7 +177,7 @@ absl::StatusOr<Tf2HloResult> CompileTfToHlo(
   TF_ASSIGN_OR_RETURN(tensorflow::tpu::TPUCompileMetadataProto compile_metadata,
                       GetCompileMetadata(entry_fn, inputs, ifrt_client));
 
-  VLOG(1) << "Compilation metadata: " << compile_metadata;
+  // VLOG(1) << "Compilation metadata: " << compile_metadata;
 
   std::vector<TensorShape> arg_shapes;
   for (const auto& input : inputs) {

tensorflow/core/common_runtime/next_pluggable_device/plugin_coordination_service_agent_helper.h

@@ -16,6 +16,7 @@ limitations under the License.
 #ifndef TENSORFLOW_CORE_COMMON_RUNTIME_NEXT_PLUGGABLE_DEVICE_PLUGIN_COORDINATION_SERVICE_AGENT_HELPER_H_
 #define TENSORFLOW_CORE_COMMON_RUNTIME_NEXT_PLUGGABLE_DEVICE_PLUGIN_COORDINATION_SERVICE_AGENT_HELPER_H_
 
+#include "absl/flags/flag.h"
 #include "tensorflow/c/kernels.h"
 #include "tensorflow/c/tf_status_helper.h"
 #include "tensorflow/core/common_runtime/next_pluggable_device/c_plugin_coordination_service_agent.h"

tensorflow/core/common_runtime/next_pluggable_device/plugin_op_kernel_helper.h

@@ -16,6 +16,7 @@ limitations under the License.
 #ifndef TENSORFLOW_CORE_COMMON_RUNTIME_NEXT_PLUGGABLE_DEVICE_PLUGIN_OP_KERNEL_HELPER_H_
 #define TENSORFLOW_CORE_COMMON_RUNTIME_NEXT_PLUGGABLE_DEVICE_PLUGIN_OP_KERNEL_HELPER_H_
 
+#include "absl/flags/flag.h"
 #include "tensorflow/c/kernels.h"
 #include "tensorflow/c/tf_status_helper.h"
 #include "tensorflow/core/common_runtime/next_pluggable_device/c_plugin_op_kernel.h"

tensorflow/core/tfrt/mlrt/kernel/ifrt_ops_kernel.cc

@@ -58,7 +58,7 @@ absl::StatusOr<tsl::RCReference<xla::ifrt::Array>> LoadIfrtVariable(
   tensorflow::ifrt_serving::VariableDeviceShardingConfigProto sharding_config;
 
   if (!tensorflow::protobuf::TextFormat::ParseFromString(
-          sharding_config_proto_text, &sharding_config)) {
+          std::string(sharding_config_proto_text).c_str(), &sharding_config)) {
     return absl::InvalidArgumentError(absl::StrCat(
         "Attribute: ", sharding_config_proto_text, " cannot be parsed"));
   }

tensorflow/java/BUILD

@@ -1,8 +1,6 @@
 # Description:
 # TensorFlow Java API.
 
-load(":build_defs.bzl", "JAVACOPTS")
-load(":src/gen/gen_ops.bzl", "tf_java_op_gen_srcjar")
 load(
     "//tensorflow:tensorflow.bzl",
     "VERSION",
@@ -13,6 +11,8 @@ load(
     "tf_custom_op_library",
     "tf_java_test",
 )
+load(":build_defs.bzl", "JAVACOPTS")
+load(":src/gen/gen_ops.bzl", "tf_java_op_gen_srcjar")
 
 package(
     # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
@@ -438,8 +438,8 @@ tf_cc_binary(
         "//conditions:default": [
             "-z defs",
             "-s",
-            "-Wl,--version-script,$(location {})".format(LINKER_VERSION_SCRIPT),
             # copybara:uncomment_begin(google-only)
+            # "-Wl,--version-script,$(location {})".format(LINKER_VERSION_SCRIPT),
             # "-Wl,--undefined-version",
             # copybara:uncomment_end
         ],

tensorflow/workspace2.bzl

@@ -292,6 +292,9 @@ def _tf_repositories():
         name = "com_github_googlecloudplatform_google_cloud_cpp",
         sha256 = "ff82045b9491f0d880fc8e5c83fd9542eafb156dcac9ff8c6209ced66ed2a7f0",
         strip_prefix = "google-cloud-cpp-1.17.1",
+        # copybara:comment_begin(oss-only)
+        patch_file = ["//third_party/systemlibs:google_cloud_cpp_int_types_fix.patch"],
+        # copybara:comment_end
         repo_mapping = {
             "@com_github_curl_curl": "@curl",
             "@com_github_nlohmann_json": "@nlohmann_json_lib",
@@ -434,6 +437,9 @@ def _tf_repositories():
         patch_file = [
             "//third_party/grpc:generate_cc_env_fix.patch",
             "//third_party/grpc:register_go_toolchain.patch",
+            # copybara:comment_begin(oss-only)
+            "//third_party/grpc:cast_memory_order_to_int.patch",
+            # copybara:comment_end
         ],
         system_link_files = {
             "//third_party/systemlibs:BUILD": "bazel/BUILD",

third_party/grpc/cast_memory_order_to_int.patch

@@ -0,0 +1,23 @@
+diff --git a/src/core/lib/gprpp/atomic.h b/src/core/lib/gprpp/atomic.h
+index 095ebf1..4a53d2c 100644
+--- a/src/core/lib/gprpp/atomic.h
++++ b/src/core/lib/gprpp/atomic.h
+@@ -28,12 +28,12 @@
+ namespace grpc_core {
+
+ enum class MemoryOrder {
+-  RELAXED = std::memory_order_relaxed,
+-  CONSUME = std::memory_order_consume,
+-  ACQUIRE = std::memory_order_acquire,
+-  RELEASE = std::memory_order_release,
+-  ACQ_REL = std::memory_order_acq_rel,
+-  SEQ_CST = std::memory_order_seq_cst
++  RELAXED = static_cast<int>(std::memory_order_relaxed),
++  CONSUME = static_cast<int>(std::memory_order_consume),
++  ACQUIRE = static_cast<int>(std::memory_order_acquire),
++  RELEASE = static_cast<int>(std::memory_order_release),
++  ACQ_REL = static_cast<int>(std::memory_order_acq_rel),
++  SEQ_CST = static_cast<int>(std::memory_order_seq_cst)
+ };
+
+ template <typename T>

third_party/systemlibs/google_cloud_cpp_int_types_fix.patch

@@ -0,0 +1,26 @@
+diff --git a/google/cloud/iam_policy.h b/google/cloud/iam_policy.h
+index 45051d8..9004ebd 100644
+--- a/google/cloud/iam_policy.h
++++ b/google/cloud/iam_policy.h
+@@ -34,7 +34,7 @@ inline namespace GOOGLE_CLOUD_CPP_NS {
+  *     about ETags.
+  */
+ struct IamPolicy {
+-  std::int32_t version;
++  int32_t version;
+   IamBindings bindings;
+   std::string etag;
+ };
+diff --git a/google/cloud/storage/internal/hash_validator_impl.h b/google/cloud/storage/internal/hash_validator_impl.h
+index dbb387e..07d9d2d 100644
+--- a/google/cloud/storage/internal/hash_validator_impl.h
++++ b/google/cloud/storage/internal/hash_validator_impl.h
+@@ -63,7 +63,7 @@ class Crc32cHashValidator : public HashValidator {
+   Result Finish() && override;
+
+  private:
+-  std::uint32_t current_{0};
++  uint32_t current_{0};
+   std::string received_hash_;
+ };
+

third_party/xla/.bazelrc

@@ -380,8 +380,8 @@ build:android --cxxopt=-std=c++17
 build:android --host_cxxopt=-std=c++17
 build:ios --cxxopt=-std=c++17
 build:ios --host_cxxopt=-std=c++17
-build:linux --cxxopt=-std=c++17
-build:linux --host_cxxopt=-std=c++17
+build:linux --cxxopt=-std=c++20
+build:linux --host_cxxopt=-std=c++20
 build:macos --cxxopt=-std=c++17
 build:macos --host_cxxopt=-std=c++17
 build:windows --cxxopt=/std:c++17

third_party/xla/third_party/tsl/.bazelrc

@@ -380,8 +380,8 @@ build:android --cxxopt=-std=c++17
 build:android --host_cxxopt=-std=c++17
 build:ios --cxxopt=-std=c++17
 build:ios --host_cxxopt=-std=c++17
-build:linux --cxxopt=-std=c++17
-build:linux --host_cxxopt=-std=c++17
+build:linux --cxxopt=-std=c++20
+build:linux --host_cxxopt=-std=c++20
 build:macos --cxxopt=-std=c++17
 build:macos --host_cxxopt=-std=c++17
 build:windows --cxxopt=/std:c++17

third_party/xla/third_party/tsl/third_party/grpc/cast_memory_order_to_int.patch

@@ -0,0 +1,23 @@
+diff --git a/src/core/lib/gprpp/atomic.h b/src/core/lib/gprpp/atomic.h
+index 095ebf1..4a53d2c 100644
+--- a/src/core/lib/gprpp/atomic.h
++++ b/src/core/lib/gprpp/atomic.h
+@@ -28,12 +28,12 @@
+ namespace grpc_core {
+
+ enum class MemoryOrder {
+-  RELAXED = std::memory_order_relaxed,
+-  CONSUME = std::memory_order_consume,
+-  ACQUIRE = std::memory_order_acquire,
+-  RELEASE = std::memory_order_release,
+-  ACQ_REL = std::memory_order_acq_rel,
+-  SEQ_CST = std::memory_order_seq_cst
++  RELAXED = static_cast<int>(std::memory_order_relaxed),
++  CONSUME = static_cast<int>(std::memory_order_consume),
++  ACQUIRE = static_cast<int>(std::memory_order_acquire),
++  RELEASE = static_cast<int>(std::memory_order_release),
++  ACQ_REL = static_cast<int>(std::memory_order_acq_rel),
++  SEQ_CST = static_cast<int>(std::memory_order_seq_cst)
+ };
+
+ template <typename T>

third_party/xla/third_party/tsl/third_party/systemlibs/google_cloud_cpp_int_types_fix.patch

@@ -0,0 +1,26 @@
+diff --git a/google/cloud/iam_policy.h b/google/cloud/iam_policy.h
+index 45051d8..9004ebd 100644
+--- a/google/cloud/iam_policy.h
++++ b/google/cloud/iam_policy.h
+@@ -34,7 +34,7 @@ inline namespace GOOGLE_CLOUD_CPP_NS {
+  *     about ETags.
+  */
+ struct IamPolicy {
+-  std::int32_t version;
++  int32_t version;
+   IamBindings bindings;
+   std::string etag;
+ };
+diff --git a/google/cloud/storage/internal/hash_validator_impl.h b/google/cloud/storage/internal/hash_validator_impl.h
+index dbb387e..07d9d2d 100644
+--- a/google/cloud/storage/internal/hash_validator_impl.h
++++ b/google/cloud/storage/internal/hash_validator_impl.h
+@@ -63,7 +63,7 @@ class Crc32cHashValidator : public HashValidator {
+   Result Finish() && override;
+
+  private:
+-  std::uint32_t current_{0};
++  uint32_t current_{0};
+   std::string received_hash_;
+ };
+