[pull] master from tensorflow:master

tensorflow/compiler/aot/tfcompile.bzl

@@ -378,7 +378,8 @@ def _tf_library(
         "-e \"s|{{TFCOMPILE_HEADER}}|$(location " + header_file + ")|g\" " +
         "-e \"s|{{TFCOMPILE_CPP_CLASS}}|" + cpp_class + "|g\" " +
         "-e \"s|{{TFCOMPILE_NAME}}|" + no_ns_name + "|g\" " +
-        "-e \"s!bazel-out/[^/]*/(bin|genfiles)/!!g\" "
+        "-e \"s!bazel-out/[^/]*/bin/!!g\" " +
+        "-e \"s!bazel-out/[^/]*/genfiles/!!g\" "
     )
 
     if gen_test:

third_party/xla/tensorflow.bazelrc

@@ -292,20 +292,20 @@ common:rocm_clang_hermetic --@rules_ml_toolchain//common:enable_rocm=True
 common:rocm_clang_hermetic --@rules_ml_toolchain//common:enable_cuda=False
 common:rocm_clang_hermetic --@rules_ml_toolchain//common:enable_sycl=False
 common:rocm_clang_hermetic --@rules_ml_toolchain//common:enable_hermetic_cc=True
-common:rocm_clang_hermetic --@local_config_rocm//rocm:rocm_path_type=hermetic --config=workspace
+common:rocm_clang_hermetic --@local_config_rocm//rocm:rocm_path_type=hermetic
 common:rocm_clang_hermetic --strategy=CppLink=local
 common:rocm_clang_hermetic --@rules_ml_toolchain//common:static_libcxx=False
 
 common:rocm --config=rocm_clang_hermetic
 common:rocm_ci --config=rocm
-common:rocm_ci --@local_config_rocm//rocm:rocm_path_type=hermetic --config=workspace
+common:rocm_ci --@local_config_rocm//rocm:rocm_path_type=hermetic
 
 common:rocm_ci_hermetic --dynamic_mode=off
 common:rocm_ci_hermetic --config=rocm_clang_hermetic
 common:rocm_ci_hermetic --repo_env=TF_ROCM_AMDGPU_TARGETS="gfx908,gfx90a"
 common:rocm_ci_hermetic --repo_env=ROCM_DISTRO_VERSION="rocm_7.13.0_gfx90a"
 common:rocm_ci_hermetic --repo_env=SYSROOT_DIST=linux_glibc_2_31
-common:rocm_ci_hermetic --@local_config_rocm//rocm:rocm_path_type=hermetic --config=workspace
+common:rocm_ci_hermetic --@local_config_rocm//rocm:rocm_path_type=hermetic
 
 # This config option is used for SYCL as GPU backend.
 # SYCL Configuration (non-hermetic)

third_party/xla/third_party/gpus/rocm/BUILD.tpl

@@ -129,16 +129,16 @@ cc_library(
     name = "rocm_rpath",
     linkopts = select({
         ":build_hermetic": [
-            "-Wl,-rpath,external/%{rocm_repo_name}/rocm/%{rocm_root}/lib",
+            "-Wl,-rpath,../%{rocm_repo_name}/rocm/%{rocm_root}/lib",
         ],
         ":link_only": [
         ],
         ":multiple_rocm_paths": [
-            "-Wl,-rpath,external/%{rocm_repo_name}/rocm/%{rocm_root}/lib",
+            "-Wl,-rpath,../%{rocm_repo_name}/rocm/%{rocm_root}/lib",
             "-Wl,-rpath=%{rocm_lib_paths}",
         ],
         "//conditions:default": [
-            "-Wl,-rpath,external/%{rocm_repo_name}/rocm/%{rocm_root}/lib",
+            "-Wl,-rpath,../%{rocm_repo_name}/rocm/%{rocm_root}/lib",
             "-Wl,-rpath,/opt/rocm/lib",
         ],
     }),

third_party/xla/xla/backends/profiler/gpu/BUILD

@@ -963,6 +963,7 @@ xla_test(
         "h100",
         "b200",
     ],
+    shard_count = 3,
     tags = [
         "cuda-only",
         "no_mac",

third_party/xla/xla/backends/profiler/gpu/cupti_tracer.cc

@@ -22,7 +22,6 @@ limitations under the License.
 #include <cstdint>
 #include <ios>
 #include <limits>
-#include <list>
 #include <memory>
 #include <string>
 #include <string_view>
@@ -31,8 +30,10 @@ limitations under the License.
 #include <vector>
 
 #include "absl/algorithm/container.h"
+#include "absl/base/const_init.h"
 #include "absl/base/no_destructor.h"
 #include "absl/base/optimization.h"
+#include "absl/base/thread_annotations.h"
 #include "absl/cleanup/cleanup.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/container/flat_hash_set.h"
@@ -59,7 +60,6 @@ limitations under the License.
 #include "xla/backends/profiler/gpu/cupti_utils.h"
 #include "xla/tsl/platform/env.h"
 #include "xla/tsl/platform/errors.h"
-#include "xla/tsl/platform/statusor.h"
 #include "xla/tsl/profiler/backends/cpu/annotation_stack.h"
 #include "xla/tsl/profiler/utils/per_thread.h"
 #include "xla/tsl/profiler/utils/xplane_builder.h"
@@ -1120,13 +1120,6 @@ const char* GetCuptiErrorString(CuptiInterface* cupti_interface,
   return err_str;
 }
 
-bool& IsCuptiHardwareEventSystemEnabled() {
-  // This flag can not flip to true once per process. Once enabled, it will stay
-  // enabled until the process is terminated.
-  static bool is_enabled = false;
-  return is_enabled;
-}
-
 }  // namespace
 
 CuptiTracer::CuptiTracer(CuptiInterface* cupti_interface)
@@ -1679,31 +1672,6 @@ absl::Status CuptiTracer::EnableActivityTracing() {
                      << err;
       }
     }
-    if (option_->enable_activity_hardware_tracing) {
-      if (IsCuptiHardwareEventSystemEnabled()) {
-        LOG(INFO) << "CUPTI activity HW trace already enabled.";
-      } else {
-        auto err = cupti_interface_->ActivityEnableHWTrace(true);
-        if (err == CUPTI_ERROR_NOT_SUPPORTED) {
-          LOG(INFO)
-              << "CUPTI activity HW trace not enabled due to not supported on "
-                 "this platform!";
-        } else if (err != CUPTI_SUCCESS) {
-          LOG(WARNING)
-              << "Fail to enable CUPTI activity HW trace, CUPTI ERROR CODE:"
-              << err << " (" << GetCuptiErrorString(cupti_interface_, err)
-              << ")";
-        } else {
-          LOG(INFO) << "CUPTI activity HW trace successfully enabled.";
-          IsCuptiHardwareEventSystemEnabled() = true;
-        }
-      }
-    } else {
-      if (IsCuptiHardwareEventSystemEnabled()) {
-        LOG(INFO)
-            << "CUPTI activity HW trace already enabled, continue with it.";
-      }
-    }
 
     if (using_v2_subscriber_api_) {
       RETURN_IF_CUPTI_ERROR(ActivityRegisterCallbacksV2(
@@ -2073,6 +2041,66 @@ absl::Status CuptiTracer::ProcessActivityBuffer(CUcontext context,
   return "";
 }
 
+/*static*/ absl::Status CuptiTracer::EnableHES() {
+  static absl::Mutex mu(absl::kConstInit);
+  static bool is_hes_enabled ABSL_GUARDED_BY(mu) = false;
+
+  absl::MutexLock lock(mu);
+  if (is_hes_enabled) {
+    LOG(INFO) << "CUPTI activity HW trace already enabled.";
+    return absl::OkStatus();
+  }
+
+  CUresult cu_err = cuInit(0);
+  if (cu_err != CUDA_SUCCESS) {
+    return absl::InternalError(absl::StrCat(
+        "cuInit(0) failed with error code: ", static_cast<int>(cu_err)));
+  }
+
+  CUcontext ctx = nullptr;
+  if (cuCtxGetCurrent(&ctx) == CUDA_SUCCESS && ctx != nullptr) {
+    return absl::FailedPreconditionError(
+        "Cannot enable HES: a CUDA context is already active on the current "
+        "thread.");
+  }
+
+  int gpu_count = NumGpus();
+  for (int i = 0; i < gpu_count; ++i) {
+    CUdevice dev;
+    if (cuDeviceGet(&dev, i) == CUDA_SUCCESS) {
+      unsigned int flags = 0;
+      int active = 0;
+      if (cuDevicePrimaryCtxGetState(dev, &flags, &active) == CUDA_SUCCESS &&
+          active) {
+        return absl::FailedPreconditionError(absl::StrCat(
+            "Cannot enable HES: active primary CUDA context found on device ",
+            i));
+      }
+    }
+  }
+
+  CuptiInterface* cupti_interface = GetCuptiInterface();
+  auto err = cupti_interface->ActivityEnableHWTrace(true);
+  if (err == CUPTI_ERROR_NOT_SUPPORTED) {
+    LOG(INFO)
+        << "CUPTI activity HW trace not enabled due to not supported on this "
+           "platform!";
+    return absl::UnimplementedError(
+        "CUPTI activity HW trace not supported on this platform.");
+  }
+  if (err != CUPTI_SUCCESS) {
+    LOG(WARNING) << "Fail to enable CUPTI activity HW trace, CUPTI ERROR CODE: "
+                 << err << " (" << GetCuptiErrorString(cupti_interface, err)
+                 << ")";
+    return absl::InternalError(
+        absl::StrCat("Fail to enable CUPTI activity HW trace: ",
+                     GetCuptiErrorString(cupti_interface, err)));
+  }
+  LOG(INFO) << "CUPTI activity HW trace successfully enabled.";
+  is_hes_enabled = true;
+  return absl::OkStatus();
+}
+
 std::vector<CallbackAnnotationsAndEvents>
 CuptiTracer::GatherCallbackAnnotationsAndEvents(bool stop_recording) {
   // Note that it is OK to call PerThread<T>'s StartRecording() multiple times

third_party/xla/xla/backends/profiler/gpu/cupti_tracer.h

@@ -158,6 +158,10 @@ class CuptiTracer {
   // Returns the error (if any) when using libcupti.
   static std::string ErrorIfAny();
 
+  // Enables activity hardware events tracing using HES (Hardware Event System).
+  // Once enabled, it stays enabled for the process lifetime.
+  static absl::Status EnableHES();
+
   // Returns true if the number of annotation strings is too large. The input
   // count is the per-thread count.
   bool TooManyAnnotationStrings(size_t count) const;

third_party/xla/xla/backends/profiler/gpu/device_tracer_cuda.cc

@@ -48,6 +48,18 @@ using tensorflow::ProfileOptions;
 using tensorflow::profiler::XSpace;
 using tsl::ReadBoolFromEnvVar;
 
+static void MaybeEnableHES() {
+  bool enable_hes = false;
+  tsl::ReadBoolFromEnvVar("TF_GPU_CUPTI_ENABLE_ACTIVITY_HW_TRACING", false,
+                          &enable_hes)
+      .IgnoreError();
+  if (enable_hes) {
+    if (auto status = CuptiTracer::EnableHES(); !status.ok()) {
+      LOG(WARNING) << "Failed to enable HES: " << status.message();
+    }
+  }
+}
+
 // GpuTracer for GPU.
 class GpuTracer : public tsl::profiler::ProfilerInterface {
  public:
@@ -103,11 +115,6 @@ absl::Status GpuTracer::DoStart() {
   options_.activities_selected.push_back(CUPTI_ACTIVITY_KIND_OVERHEAD);
   options_.activities_selected.push_back(CUPTI_ACTIVITY_KIND_MEMSET);
 
-  // TODO: Change default to true once we have more confidence in HES.
-  ReadBoolFromEnvVar("TF_GPU_CUPTI_ENABLE_ACTIVITY_HW_TRACING", false,
-                     &options_.enable_activity_hardware_tracing)
-      .IgnoreError();
-
 // CUDA/CUPTI 10 have issues (leaks and crashes) with cuptiFinalize.
 #if CUDA_VERSION >= 11000
   options_.cupti_finalize = true;
@@ -243,6 +250,7 @@ std::unique_ptr<tsl::profiler::ProfilerInterface> CreateGpuTracer(
 }
 
 auto register_gpu_tracer_factory = [] {
+  MaybeEnableHES();
   RegisterProfilerFactory(&CreateGpuTracer);
   return 0;
 }();

third_party/xla/xla/backends/profiler/gpu/profile_with_cuda_kernels_test.cc

@@ -110,6 +110,12 @@ void HandleRecords(PmSamples* samples) {
 
 void SimpleAddSubWithProfilerTest(bool enable_activity_hardware_tracing,
                                   bool enable_pm_sampling) {
+  if (enable_activity_hardware_tracing) {
+    if (auto status = CuptiTracer::EnableHES(); !status.ok()) {
+      LOG(WARNING) << "Failed to enable HES: " << status.message();
+    }
+  }
+
   uint32_t cupti_version = 0;
   cuptiGetVersion(&cupti_version);
   LOG(INFO) << "RUNTIME CUPTI version " << cupti_version

third_party/xla/xla/hlo/transforms/collectives/BUILD

@@ -572,6 +572,7 @@ cc_library(
         "//xla/hlo/transforms/simplifiers:tuple_simplifier",
         "//xla/service:call_graph",
         "//xla/service:tuple_util",
+        "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
@@ -591,6 +592,7 @@ xla_cc_test(
     srcs = ["infeed_token_propagation_test.cc"],
     deps = [
         ":infeed_token_propagation",
+        "//xla:shape_util",
         "//xla/hlo/analysis:hlo_ordering",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/testlib:hlo_hardware_independent_test_base",