Support multi-pass profiling in OpenXLA FunctionalHloRunner.

tsl/profiler/lib/BUILD

@@ -109,10 +109,12 @@ cc_library(
         "//learning/brain/tfrc/executor/api:__pkg__",
     ]),
     deps = [
+        ":multi_pass_profiler_controller",
         ":profiler_controller",
         ":profiler_interface",
         "//tsl/profiler/protobuf:profiler_options_proto_cc",
         "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/log",
         "@com_google_absl//absl/synchronization",
     ],
     alwayslink = True,
@@ -125,10 +127,12 @@ tsl_cc_test(
         ":profiler_factory",
         ":profiler_factory_impl",
         ":profiler_interface",
+        ":profiler_lock",
+        ":profiler_passes",
         "//tsl/profiler/protobuf:profiler_options_proto_cc",
         "//tsl/profiler/protobuf:xplane_proto_cc",
-        "@com_google_absl//absl/memory",
         "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings:string_view",
         "@com_google_googletest//:gtest_main",
         "@xla//xla/tsl/platform:macros",
         "@xla//xla/tsl/platform:test",
@@ -150,6 +154,7 @@ cc_library(
         "//tsl/profiler/protobuf:xplane_proto_cc",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
     ],
 )
 
@@ -200,6 +205,52 @@ tsl_cc_test(
     ],
 )
 
+cc_library(
+    name = "multi_pass_profiler_controller",
+    srcs = ["multi_pass_profiler_controller.cc"],
+    hdrs = ["multi_pass_profiler_controller.h"],
+    visibility = internal_visibility([
+        "@xla//xla/tsl/profiler:internal",
+    ]),
+    deps = [
+        ":profiler_interface",
+        "//tsl/profiler/protobuf:xplane_proto_cc",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@xla//xla/tsl/platform:logging",
+    ],
+)
+
+cc_library(
+    name = "profiler_passes",
+    srcs = ["profiler_passes.cc"],
+    hdrs = ["profiler_passes.h"],
+    visibility = internal_visibility([
+        "@xla//xla/tsl/profiler:internal",
+        "@xla//xla/tsl/profiler:friends",
+    ]),
+    deps = [
+        ":profiler_factory",
+        ":profiler_interface",
+        ":profiler_lock",
+        "//tsl/platform:platform_port",
+        "//tsl/profiler/protobuf:profiler_options_proto_cc",
+        "//tsl/profiler/protobuf:xplane_proto_cc",
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/synchronization",
+        "@xla//xla/tsl/platform:env",
+        "@xla//xla/tsl/platform:logging",
+        "@xla//xla/tsl/profiler/convert:post_process_single_host_xplane",
+        "@xla//xla/tsl/profiler/utils:time_utils",
+        "@xla//xla/tsl/profiler/utils:xplane_builder",
+        "@xla//xla/tsl/profiler/utils:xplane_schema",
+        "@xla//xla/tsl/profiler/utils:xplane_utils",
+    ],
+)
+
 cc_library(
     name = "profiler_session",
     hdrs = ["profiler_session.h"],
@@ -399,10 +450,13 @@ cc_library(
     srcs = ["profiler_collection.cc"],
     hdrs = ["profiler_collection.h"],
     visibility = internal_visibility([
+        "@xla//xla/backends/profiler:__pkg__",
         "@xla//xla/backends/profiler/plugin:__pkg__",
         "//learning/brain/tfrc/executor/api:__pkg__",
         "@xla//xla/backends/profiler/cpu:__pkg__",
         "@xla//xla/backends/profiler/subprocess:__pkg__",
+        "@xla//xla/backends/profiler/tpu:__pkg__",
+        "@xla//xla/tsl/profiler:internal",
     ]),
     deps = [
         ":profiler_interface",

tsl/profiler/lib/multi_pass_profiler_controller.cc

@@ -0,0 +1,203 @@
+/* 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.
+==============================================================================*/
+
+#include "tsl/profiler/lib/multi_pass_profiler_controller.h"
+
+#include <memory>
+#include <utility>
+
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
+#include "xla/tsl/platform/logging.h"
+#include "tsl/profiler/lib/profiler_interface.h"
+#include "tsl/profiler/protobuf/xplane.pb.h"
+
+namespace tsl {
+namespace profiler {
+
+MultiPassProfilerController::MultiPassProfilerController(
+    std::unique_ptr<MultiPassProfilerInterface> profiler)
+    : profiler_(std::move(profiler)) {}
+
+MultiPassProfilerController::~MultiPassProfilerController() {
+  // Ensure an active pass is stopped.
+  if (state_ == MultiPassProfilerState::kPassStarted) {
+    StopPass().IgnoreError();
+  }
+  if (state_ == MultiPassProfilerState::kStarted ||
+      state_ == MultiPassProfilerState::kPassStopped) {
+    profiler_->Stop().IgnoreError();
+    state_ = MultiPassProfilerState::kStopped;
+  }
+}
+
+bool MultiPassProfilerController::NeedMorePasses() {
+  if (state_ == MultiPassProfilerState::kPassStopped ||
+      state_ == MultiPassProfilerState::kStarted ||
+      state_ == MultiPassProfilerState::kInit) {
+    if (!status_.ok()) return false;
+    return profiler_->NeedMorePasses();
+  }
+  return false;
+}
+
+absl::Status MultiPassProfilerController::StartPass() {
+  absl::Status status;
+  if (state_ == MultiPassProfilerState::kStarted ||
+      state_ == MultiPassProfilerState::kPassStopped) {
+    state_ = MultiPassProfilerState::kPassStarted;
+    if (status_.ok()) {
+      status = status_ = profiler_->StartPass();
+    } else {
+      status = absl::AbortedError("Previous call returned an error.");
+    }
+  } else {
+    status = absl::AbortedError("StartPass called in the wrong order");
+  }
+  if (!status.ok()) LOG(ERROR) << status;
+  return status;
+}
+
+absl::Status MultiPassProfilerController::PushRange(absl::string_view name) {
+  absl::Status status;
+  if (state_ == MultiPassProfilerState::kPassStarted) {
+    if (status_.ok()) {
+      status = status_ = profiler_->PushRange(name);
+      if (status_.ok()) {
+        ++active_range_count_;
+      }
+    } else {
+      status = absl::AbortedError("Previous call returned an error.");
+    }
+  } else {
+    status = absl::AbortedError("PushRange called in the wrong order");
+  }
+  if (!status.ok()) LOG(ERROR) << status;
+  return status;
+}
+
+absl::Status MultiPassProfilerController::PopRange() {
+  absl::Status status;
+  if (state_ == MultiPassProfilerState::kPassStarted) {
+    if (active_range_count_ <= 0) {
+      if (status_.ok()) {
+        status = status_ =
+            absl::InternalError("PopRange called with no active ranges.");
+      } else {
+        status = absl::AbortedError("Previous call returned an error.");
+      }
+    } else {
+      --active_range_count_;
+      if (status_.ok()) {
+        status = status_ = profiler_->PopRange();
+      } else {
+        profiler_->PopRange().IgnoreError();
+        status = absl::AbortedError("Previous call returned an error.");
+      }
+    }
+  } else {
+    status = absl::AbortedError("PopRange called in the wrong order");
+  }
+  if (!status.ok()) LOG(ERROR) << status;
+  return status;
+}
+
+absl::Status MultiPassProfilerController::PopAllRanges() {
+  absl::Status status;
+  while (active_range_count_ > 0) {
+    status.Update(PopRange());
+  }
+  return status;
+}
+
+absl::Status MultiPassProfilerController::StopPass() {
+  absl::Status status;
+  if (state_ == MultiPassProfilerState::kPassStarted) {
+    status = PopAllRanges();
+    state_ = MultiPassProfilerState::kPassStopped;
+    if (status_.ok()) {
+      status = status_ = profiler_->StopPass();
+    } else {
+      profiler_->StopPass().IgnoreError();
+      if (status.ok()) {
+        status = absl::AbortedError("Previous call returned an error.");
+      }
+    }
+  } else {
+    status = absl::AbortedError("StopPass called in the wrong order");
+  }
+  if (!status.ok()) LOG(ERROR) << status;
+  return status;
+}
+
+absl::Status MultiPassProfilerController::CollectData(
+    tensorflow::profiler::XSpace* space) {
+  absl::Status status;
+  if (state_ == MultiPassProfilerState::kInit) {
+    status = absl::OkStatus();
+    state_ = MultiPassProfilerState::kCollectData;
+  } else if (state_ == MultiPassProfilerState::kStopped) {
+    state_ = MultiPassProfilerState::kCollectData;
+    if (status_.ok()) {
+      status = status_ = profiler_->CollectData(space);
+    } else {
+      status = absl::AbortedError("Previous call returned an error.");
+    }
+  } else {
+    status = absl::AbortedError("CollectData called in the wrong order.");
+  }
+  if (!status.ok()) LOG(ERROR) << status;
+  return status;
+}
+
+absl::Status MultiPassProfilerController::Start() {
+  absl::Status status;
+  if (state_ == MultiPassProfilerState::kInit) {
+    state_ = MultiPassProfilerState::kStarted;
+    if (status_.ok()) {
+      status = status_ = profiler_->Start();
+    } else {
+      status = absl::AbortedError("Previous call returned an error.");
+    }
+  } else {
+    status = absl::AbortedError("Start called in the wrong order");
+  }
+  if (!status.ok()) LOG(ERROR) << status;
+  return status;
+}
+
+absl::Status MultiPassProfilerController::Stop() {
+  absl::Status status;
+  if (state_ == MultiPassProfilerState::kPassStarted) {
+    status = StopPass();
+  }
+  if (state_ == MultiPassProfilerState::kPassStopped ||
+      state_ == MultiPassProfilerState::kStarted) {
+    state_ = MultiPassProfilerState::kStopped;
+    if (status_.ok()) {
+      status = status_ = profiler_->Stop();
+    } else if (status.ok()) {
+      status = absl::AbortedError("Previous call returned an error.");
+    }
+  } else if (status.ok()) {
+    status = absl::AbortedError("Stop called in the wrong order");
+  }
+  if (!status.ok()) LOG(ERROR) << status;
+  return status;
+}
+
+}  // namespace profiler
+
+}  // namespace tsl

tsl/profiler/lib/multi_pass_profiler_controller.h

@@ -0,0 +1,77 @@
+/* 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_TSL_PROFILER_LIB_MULTI_PASS_PROFILER_CONTROLLER_H_
+#define TENSORFLOW_TSL_PROFILER_LIB_MULTI_PASS_PROFILER_CONTROLLER_H_
+
+#include <memory>
+
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
+#include "tsl/profiler/lib/profiler_interface.h"
+#include "tsl/profiler/protobuf/xplane.pb.h"
+
+namespace tsl {
+namespace profiler {
+
+// Decorator for multi-pass profiler plugins.
+//
+// Tracks that calls to the underlying multi-pass profiler interface functions
+// are made in the expected order: Start, StartPass, [PushRange/PopRange],
+// StopPass, ..., Stop, and CollectData. Making calls in an invalid order causes
+// them to be aborted.
+//
+// If a call to the decorated profiler fails, subsequent calls will be aborted
+// and no further calls will be forwarded to the underlying profiler.
+class MultiPassProfilerController : public MultiPassProfilerInterface {
+ public:
+  explicit MultiPassProfilerController(
+      std::unique_ptr<MultiPassProfilerInterface> profiler);
+  ~MultiPassProfilerController() override;
+
+  // MultiPassProfilerInterface methods
+  bool NeedMorePasses() override;
+  absl::Status StartPass() override;
+  absl::Status PushRange(absl::string_view name) override;
+  absl::Status PopRange() override;
+  absl::Status StopPass() override;
+
+  // ProfilerInterface methods
+  absl::Status Start() override;
+  absl::Status Stop() override;
+  absl::Status CollectData(tensorflow::profiler::XSpace* space) override;
+
+ private:
+  enum class MultiPassProfilerState {
+    kInit = 0,
+    kStarted = 1,
+    kPassStarted = 2,
+    kPassStopped = 3,
+    kStopped = 4,
+    kCollectData = 5,
+  };
+
+  absl::Status PopAllRanges();
+
+  MultiPassProfilerState state_ = MultiPassProfilerState::kInit;
+  std::unique_ptr<MultiPassProfilerInterface> profiler_;
+  absl::Status status_;  // result of calls to profiler_
+  int active_range_count_ = 0;
+};
+
+}  // namespace profiler
+}  // namespace tsl
+
+#endif  // TENSORFLOW_TSL_PROFILER_LIB_MULTI_PASS_PROFILER_CONTROLLER_H_