Unify OSS and internal logging libraries.

This removes the duplicated implementation of `LOG`, `VLOG` and `CHECK` macros
and relies directly on the implementation provided by Abseil.

TSL logging is not consistently used in XLA. Sometime the Abseil logging headers are included instead. This is not a problem internally because TSL logging aliases Abseil logging. However in OSS, TSL logging is a custom logging library. Abseil logging doesn't recognise the `TF_CPP_MAX_VLOG_LEVEL` and `TF_CPP_VMODULE` environment variables while OSS TSL logging doesn't recognise the `--vmodule` and `--v` flags.

To preserve backward compatibility in OSS, we run static initialisation logic that initialises Abseil logging from the `TF_CPP_MIN_LOG_LEVEL`/`TF_CPP_MAX_VLOG_LEVEL`/`TF_CPP_VMODULE` env variables.

PiperOrigin-RevId: 780164601

Author: allanrenucci

xla/BUILD

@@ -264,7 +264,6 @@ xla_cc_test(
         "//xla/tsl/platform:statusor",
         "//xla/tsl/platform:test_main",
         "@com_google_absl//absl/base:log_severity",
-        "@com_google_absl//absl/log:log_sink",
         "@com_google_absl//absl/log:scoped_mock_log",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",

xla/service/gpu/BUILD

@@ -1738,7 +1738,6 @@ xla_test(
         "@com_google_absl//absl/base:log_severity",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
-        "@com_google_absl//absl/log:log_sink",
         "@com_google_absl//absl/log:scoped_mock_log",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
@@ -2252,7 +2251,6 @@ xla_test(
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/base:log_severity",
         "@com_google_absl//absl/log",
-        "@com_google_absl//absl/log:log_sink",
         "@com_google_absl//absl/log:scoped_mock_log",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings:string_view",

xla/service/gpu/gpu_compiler_test.cc

@@ -25,7 +25,6 @@ limitations under the License.
 #include <ostream>
 #include <string>
 #include <tuple>
-#include <type_traits>
 #include <utility>
 #include <variant>
 #include <vector>
@@ -35,7 +34,6 @@ limitations under the License.
 #include "absl/base/log_severity.h"
 #include "absl/log/check.h"
 #include "absl/log/log.h"
-#include "absl/log/log_sink.h"
 #include "absl/log/scoped_mock_log.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
@@ -1901,15 +1899,12 @@ ENTRY %main {
 
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                           ParseAndReturnVerifiedModule(kHlo, config));
-  // absl::ScopedMockLog only works if we're actually using ABSL logging, and
-  // TSL supports a homegrown logging implementation, so we should only check
-  // the log is emitted when ABSL logging is used.
+
   absl::ScopedMockLog mock_log(absl::MockLogDefault::kIgnoreUnexpected);
-  if constexpr (std::is_same_v<absl::LogSink, tsl::TFLogSink>) {
-    EXPECT_CALL(mock_log,
-                Log(absl::LogSeverity::kWarning, EndsWith("/gpu_compiler.cc"),
-                    StartsWith("Using fallback sort algorithm")));
-  }
+  EXPECT_CALL(mock_log,
+              Log(absl::LogSeverity::kWarning, EndsWith("/gpu_compiler.cc"),
+                  StartsWith("Using fallback sort algorithm")));
+
   // StartCapturingLogs has to be called even if we expect not to capture any
   // logs.
   mock_log.StartCapturingLogs();

xla/service/gpu/gpu_hlo_schedule_test.cc

@@ -22,15 +22,13 @@ limitations under the License.
 #include <optional>
 #include <string>
 #include <tuple>
-#include <type_traits>
 #include <vector>
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include "absl/algorithm/container.h"
 #include "absl/base/log_severity.h"
 #include "absl/log/log.h"
-#include "absl/log/log_sink.h"
 #include "absl/log/scoped_mock_log.h"
 #include "absl/status/status.h"
 #include "absl/strings/string_view.h"
@@ -1809,15 +1807,10 @@ TEST_F(GpuHloScheduleTest, LogAnErrorWhenArgumentSizeExceedsMemoryLimit) {
       auto module, ParseAndReturnVerifiedModule(kHloText, module_config));
 
   absl::ScopedMockLog mock_log(absl::MockLogDefault::kIgnoreUnexpected);
-  // absl::ScopedMockLog only works if we're actually using ABSL logging, and
-  // TSL supports a homegrown logging implementation, so we should only check
-  // the log is emitted when ABSL logging is used.
-  if constexpr (std::is_same_v<absl::LogSink, tsl::TFLogSink>) {
-    EXPECT_CALL(mock_log,
-                Log(absl::LogSeverity::kError, _,
-                    EndsWith("This indicates an error in the calculation!")))
-        .Times(1);
-  }
+  EXPECT_CALL(mock_log,
+              Log(absl::LogSeverity::kError, _,
+                  EndsWith("This indicates an error in the calculation!")))
+      .Times(1);
   mock_log.StartCapturingLogs();
   TF_ASSERT_OK_AND_ASSIGN(auto metadata, ScheduleGpuModule(module.get()));
   EXPECT_EQ(metadata.scheduler_mem_limit, 0);

xla/status_macros_test.cc

@@ -17,13 +17,11 @@ limitations under the License.
 
 #include <functional>
 #include <string>
-#include <type_traits>
 #include <utility>
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include "absl/base/log_severity.h"
-#include "absl/log/log_sink.h"
 #include "absl/log/scoped_mock_log.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
@@ -93,15 +91,10 @@ TEST(StatusMacros, RetCheckFailingWithExtraMessage) {
 }
 
 TEST(StatusMacros, RetCheckLogWarning) {
-  // absl::ScopedMockLog only works if we're actually using ABSL logging, and
-  // TSL supports a homegrown logging implementation, so we should only check
-  // the log is emitted when ABSL logging is used.
   absl::ScopedMockLog mock_log(absl::MockLogDefault::kIgnoreUnexpected);
   const std::string kExpectedRegex = "RET_CHECK.*1 == 2 extra message";
-  if constexpr (std::is_same_v<absl::LogSink, tsl::TFLogSink>) {
-    EXPECT_CALL(mock_log, Log(absl::LogSeverity::kWarning, ::testing::_,
-                              ::testing::ContainsRegex(kExpectedRegex)));
-  }
+  EXPECT_CALL(mock_log, Log(absl::LogSeverity::kWarning, ::testing::_,
+                            ::testing::ContainsRegex(kExpectedRegex)));
   mock_log.StartCapturingLogs();
   absl::Status status =
       RetCheckFailWithLogSeverity(absl::LogSeverity::kWarning);
@@ -112,10 +105,8 @@ TEST(StatusMacros, RetCheckLogWarning) {
 TEST(StatusMacros, RetCheckLogInfo) {
   absl::ScopedMockLog mock_log(absl::MockLogDefault::kIgnoreUnexpected);
   const std::string kExpectedRegex = "RET_CHECK.*1 == 2 extra message";
-  if constexpr (std::is_same_v<absl::LogSink, tsl::TFLogSink>) {
-    EXPECT_CALL(mock_log, Log(absl::LogSeverity::kInfo, ::testing::_,
-                              ::testing::ContainsRegex(kExpectedRegex)));
-  }
+  EXPECT_CALL(mock_log, Log(absl::LogSeverity::kInfo, ::testing::_,
+                            ::testing::ContainsRegex(kExpectedRegex)));
   mock_log.StartCapturingLogs();
   absl::Status status = RetCheckFailWithLogSeverity(absl::LogSeverity::kInfo);
   EXPECT_EQ(status.code(), tsl::error::INTERNAL);
@@ -184,10 +175,8 @@ TEST(StatusMacros, AssignOrReturnUnsuccessfully) {
 TEST(StatusMacros, XlaRetCheckFailLogWarning) {
   absl::ScopedMockLog mock_log(absl::MockLogDefault::kIgnoreUnexpected);
   const std::string kExpectedLog = "xla ret check fail message";
-  if constexpr (std::is_same_v<absl::LogSink, tsl::TFLogSink>) {
-    EXPECT_CALL(mock_log, Log(absl::LogSeverity::kWarning, ::testing::_,
-                              ::testing::HasSubstr(kExpectedLog)));
-  }
+  EXPECT_CALL(mock_log, Log(absl::LogSeverity::kWarning, ::testing::_,
+                            ::testing::HasSubstr(kExpectedLog)));
   mock_log.StartCapturingLogs();
   absl::Status status = XlaRetCheckFailLogWarning();
   EXPECT_EQ(status.code(), tsl::error::INTERNAL);

xla/tsl/platform/BUILD

@@ -5,13 +5,13 @@ load(
 load(
     "//xla/tsl:tsl.bzl",
     "if_not_fuchsia",
+    "if_oss",
     "internal_visibility",
     "tsl_copts",
 )
 load("//xla/tsl:tsl.default.bzl", "get_compatible_with_portable")
 load(
     "//xla/tsl/platform:build_config.bzl",
-    "tf_logging_deps",
     "tf_platform_alias",
     "tf_platform_deps",
     "tf_protobuf_compiler_deps",
@@ -51,7 +51,6 @@ exports_files(
         "rocm_rocdl_path.h",
         "resource_loader.h",
         "file_system.cc",
-        "logging.h",
         "file_system.h",
         "file_system_helper.cc",
         "file_system_helper.h",
@@ -305,6 +304,7 @@ filegroup(
     name = "xla_cpu_runtime_srcs",
     srcs = [
         "dynamic_annotations.h",
+        "env_time.h",
         "macros.h",
     ],
     compatible_with = get_compatible_with_portable(),
@@ -468,12 +468,31 @@ cc_library(
 
 cc_library(
     name = "logging",
+    srcs = ["logging.cc"],
+    hdrs = ["logging.h"],
     compatible_with = get_compatible_with_portable(),
-    textual_hdrs = ["logging.h"],
     visibility = [
         "//visibility:public",
     ],
-    deps = tf_logging_deps(),
+    deps = [
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/log:globals",
+        "@com_google_absl//absl/log:vlog_is_on",
+        "@com_google_absl//absl/strings:string_view",
+    ] + if_oss([":logging_initializer"]),
+)
+
+cc_library(
+    name = "logging_initializer",
+    srcs = ["logging_initializer.cc"],
+    deps = [
+        "@com_google_absl//absl/base:log_severity",
+        "@com_google_absl//absl/log:globals",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:string_view",
+    ],
+    alwayslink = True,
 )
 
 tsl_cc_test(
@@ -486,11 +505,11 @@ tsl_cc_test(
         ":logging",
         ":statusor",
         ":test",
-        "@com_google_absl//absl/base:log_severity",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings:str_format",
         "@com_google_absl//absl/strings:string_view",
+        "@com_google_googletest//:gtest",
         "@tsl//tsl/platform:path",
         "@tsl//tsl/platform:stacktrace_handler",
     ],
@@ -518,6 +537,9 @@ cc_library(
         "@com_google_absl//absl/functional:function_ref",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/log:log_entry",
+        "@com_google_absl//absl/log:log_sink",
+        "@com_google_absl//absl/log:log_sink_registry",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:cord",

xla/tsl/platform/build_config.bzl

@@ -23,7 +23,6 @@ load(
     _tf_google_mobile_srcs_only_runtime = "tf_google_mobile_srcs_only_runtime",
     _tf_jspb_proto_library = "tf_jspb_proto_library",
     _tf_lib_proto_parsing_deps = "tf_lib_proto_parsing_deps",
-    _tf_logging_deps = "tf_logging_deps",
     _tf_platform_alias = "tf_platform_alias",
     _tf_platform_deps = "tf_platform_deps",
     _tf_portable_deps_no_runtime = "tf_portable_deps_no_runtime",
@@ -62,7 +61,6 @@ tf_google_mobile_srcs_no_runtime = _tf_google_mobile_srcs_no_runtime
 tf_google_mobile_srcs_only_runtime = _tf_google_mobile_srcs_only_runtime
 tf_jspb_proto_library = _tf_jspb_proto_library
 tf_lib_proto_parsing_deps = _tf_lib_proto_parsing_deps
-tf_logging_deps = _tf_logging_deps
 tf_platform_alias = _tf_platform_alias
 tf_platform_deps = _tf_platform_deps
 tf_portable_proto_lib = _tf_portable_proto_lib

xla/tsl/platform/default/BUILD

@@ -285,32 +285,6 @@ cc_library(
     ]),
 )
 
-cc_library(
-    name = "logging",
-    srcs = ["logging.cc"],
-    hdrs = ["//xla/tsl/platform:logging.h"],
-    tags = [
-        "manual",
-        "no_oss",
-        "nobuilder",
-    ],
-    textual_hdrs = ["logging.h"],
-    deps = [
-        "//xla/tsl/platform:env_time",
-        "//xla/tsl/platform:macros",
-        "//xla/tsl/platform:types",
-        "@com_google_absl//absl/base",
-        "@com_google_absl//absl/base:core_headers",
-        "@com_google_absl//absl/base:log_severity",
-        "@com_google_absl//absl/container:flat_hash_map",
-        "@com_google_absl//absl/strings",
-        "@com_google_absl//absl/strings:str_format",
-        "@com_google_absl//absl/strings:string_view",
-        "@com_google_absl//absl/synchronization",
-        "@tsl//tsl/platform",
-    ],
-)
-
 filegroup(
     name = "xla_cpu_runtime_srcs",
     srcs = [
@@ -656,7 +630,6 @@ exports_files(
         ["*"],
         exclude = [
             "integral_types.h",
-            "logging.h",
             "test.cc",
         ],
     ),
@@ -666,7 +639,6 @@ exports_files(
 exports_files(
     srcs = [
         "integral_types.h",
-        "logging.h",
         "test.cc",
     ],
     visibility = internal_visibility([

xla/tsl/platform/default/build_config.bzl

@@ -318,7 +318,6 @@ def tf_additional_lib_hdrs():
         clean_dep("//xla/tsl/platform/default:context.h"),
         clean_dep("//xla/tsl/platform/default:criticality.h"),
         clean_dep("//xla/tsl/platform/default:integral_types.h"),
-        clean_dep("//xla/tsl/platform/default:logging.h"),
         clean_dep("//xla/tsl/platform/default:stacktrace.h"),
         clean_dep("//xla/tsl/platform/default:status.h"),
         clean_dep("//xla/tsl/platform/default:statusor.h"),
@@ -549,9 +548,6 @@ def tf_stream_executor_deps(name, platform_dir = "@xla//xla/tsl/platform/"):
 def tf_platform_alias(name, platform_dir = "@xla//xla/tsl/platform/"):
     return [platform_dir + "default:" + name]
 
-def tf_logging_deps():
-    return [clean_dep("//xla/tsl/platform/default:logging")]
-
 def tf_error_logging_deps():
     return [clean_dep("//xla/tsl/platform/default:error_logging")]