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: 799665233

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/debug/BUILD

@@ -66,10 +66,8 @@ xla_cc_test(
     deps = [
         ":unstable_reduction_detector",
         "//xla/hlo/parser:hlo_parser",
-        "//xla/tsl/platform:logging",
         "//xla/tsl/platform:statusor",
         "@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:status_matchers",

xla/service/debug/unstable_reduction_detector_test.cc

@@ -15,18 +15,14 @@ limitations under the License.
 
 #include "xla/service/debug/unstable_reduction_detector.h"
 
-#include <type_traits>
-
 #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/status_matchers.h"
 #include "absl/strings/string_view.h"
 #include "xla/hlo/parser/hlo_parser.h"
-#include "xla/tsl/platform/logging.h"
 #include "xla/tsl/platform/statusor.h"
 
 namespace xla {
@@ -80,18 +76,16 @@ TEST(UnstableReductionDetectorTest, FailOnUnstableReductions) {
           DebugOptions::UNSTABLE_REDUCTION_DETECTION_MODE_FAIL);
   UnstableReductionDetector detector;
   ::absl::ScopedMockLog log;
-  if constexpr (std::is_same_v<absl::LogSink, tsl::TFLogSink>) {
-    EXPECT_CALL(
-        log,
-        Log(LogSeverity::kWarning, _,
-            HasSubstr("1 unstable reductions found in module 'module_main'")));
-    EXPECT_CALL(log,
-                Log(LogSeverity::kWarning, _,
-                    "Unstable reduction: %red.1 = bf16[] reduce(%p0.1, %init), "
-                    "dimensions={0}, to_apply=%red, "
-                    "metadata={op_name=\"op_name\" "
-                    "source_file=\"source_file.py\" source_line=42}"));
-  }
+  EXPECT_CALL(
+      log,
+      Log(LogSeverity::kWarning, _,
+          HasSubstr("1 unstable reductions found in module 'module_main'")));
+  EXPECT_CALL(log,
+              Log(LogSeverity::kWarning, _,
+                  "Unstable reduction: %red.1 = bf16[] reduce(%p0.1, %init), "
+                  "dimensions={0}, to_apply=%red, "
+                  "metadata={op_name=\"op_name\" "
+                  "source_file=\"source_file.py\" source_line=42}"));
   log.StartCapturingLogs();
   EXPECT_THAT(
       detector.Run(module.get(), /*execution_threads=*/{}),
@@ -110,16 +104,13 @@ TEST(UnstableReductionDetectorTest, WarningOnUnstableReduction) {
           DebugOptions::UNSTABLE_REDUCTION_DETECTION_MODE_WARNING);
   UnstableReductionDetector detector;
   ::absl::ScopedMockLog log;
-  if constexpr (std::is_same_v<absl::LogSink, tsl::TFLogSink>) {
-    EXPECT_CALL(log,
-                Log(LogSeverity::kWarning, _,
-                    "1 unstable reductions found in module 'module_main'"));
-    EXPECT_CALL(log, Log(LogSeverity::kWarning, _,
-                         "Unstable reduction: %red.1 = bf16[] reduce(%p0.1, "
-                         "%init), dimensions={0}, to_apply=%red, "
-                         "metadata={op_name=\"op_name\" "
-                         "source_file=\"source_file.py\" source_line=42}"));
-  }
+  EXPECT_CALL(log, Log(LogSeverity::kWarning, _,
+                       "1 unstable reductions found in module 'module_main'"));
+  EXPECT_CALL(log, Log(LogSeverity::kWarning, _,
+                       "Unstable reduction: %red.1 = bf16[] reduce(%p0.1, "
+                       "%init), dimensions={0}, to_apply=%red, "
+                       "metadata={op_name=\"op_name\" "
+                       "source_file=\"source_file.py\" source_line=42}"));
   log.StartCapturingLogs();
   EXPECT_THAT(detector.Run(module.get(), /*execution_threads=*/{}),
               IsOkAndHolds(false));
@@ -135,14 +126,11 @@ TEST(UnstableReductionDetectorTest, FailOnUnstableReductionNoMetadata) {
           DebugOptions::UNSTABLE_REDUCTION_DETECTION_MODE_FAIL);
   UnstableReductionDetector detector;
   ::absl::ScopedMockLog log;
-  if constexpr (std::is_same_v<absl::LogSink, tsl::TFLogSink>) {
-    EXPECT_CALL(log,
-                Log(LogSeverity::kWarning, _,
-                    "1 unstable reductions found in module 'module_main'"));
-    EXPECT_CALL(log, Log(LogSeverity::kWarning, _,
-                         "Unstable reduction: %red.1 = bf16[] reduce(%p0.1, "
-                         "%init), dimensions={0}, to_apply=%red"));
-  }
+  EXPECT_CALL(log, Log(LogSeverity::kWarning, _,
+                       "1 unstable reductions found in module 'module_main'"));
+  EXPECT_CALL(log, Log(LogSeverity::kWarning, _,
+                       "Unstable reduction: %red.1 = bf16[] reduce(%p0.1, "
+                       "%init), dimensions={0}, to_apply=%red"));
   log.StartCapturingLogs();
   EXPECT_THAT(detector.Run(module.get(), /*execution_threads=*/{}),
               StatusIs(absl::StatusCode::kFailedPrecondition,

xla/service/gpu/BUILD

@@ -1762,7 +1762,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:status_matchers",
         "@com_google_absl//absl/status:statusor",
@@ -2281,7 +2280,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/status_matchers.h"
 #include "absl/status/statusor.h"
@@ -1908,15 +1906,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