From 9a4cd70edb668afb458392ac66f30d5e82e78839 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Fri, 22 May 2026 13:25:37 -0700 Subject: [PATCH] Support multipass profiling in multihost_hlo_runner. PiperOrigin-RevId: 919831591 --- tsl/profiler/lib/BUILD | 54 +++++ .../lib/multi_pass_profiler_controller.cc | 176 +++++++++++++++ .../lib/multi_pass_profiler_controller.h | 65 ++++++ tsl/profiler/lib/profiler_factory.cc | 39 ++++ tsl/profiler/lib/profiler_factory.h | 12 + tsl/profiler/lib/profiler_factory_test.cc | 52 +++++ tsl/profiler/lib/profiler_interface.h | 23 ++ tsl/profiler/lib/profiler_passes.cc | 205 ++++++++++++++++++ tsl/profiler/lib/profiler_passes.h | 98 +++++++++ tsl/profiler/protobuf/profiler_options.proto | 5 +- 10 files changed, 728 insertions(+), 1 deletion(-) create mode 100644 tsl/profiler/lib/multi_pass_profiler_controller.cc create mode 100644 tsl/profiler/lib/multi_pass_profiler_controller.h create mode 100644 tsl/profiler/lib/profiler_passes.cc create mode 100644 tsl/profiler/lib/profiler_passes.h diff --git a/tsl/profiler/lib/BUILD b/tsl/profiler/lib/BUILD index 9293a2988..2908c59bd 100644 --- a/tsl/profiler/lib/BUILD +++ b/tsl/profiler/lib/BUILD @@ -109,10 +109,12 @@ cc_library( "//learning/brain/tfrc/executor/stream_executor:__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, @@ -129,6 +131,7 @@ tsl_cc_test( "//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", @@ -148,6 +151,8 @@ cc_library( ]), deps = [ "//tsl/profiler/protobuf:xplane_proto_cc", + "@com_google_absl//absl/status", + "@com_google_absl//absl/strings", "@xla//xla/tsl/platform:status", ], ) @@ -183,6 +188,54 @@ 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_collection", + ":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/status:status_macros", + "@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"], @@ -380,6 +433,7 @@ 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/stream_executor:__pkg__", "@xla//xla/backends/profiler/cpu:__pkg__", diff --git a/tsl/profiler/lib/multi_pass_profiler_controller.cc b/tsl/profiler/lib/multi_pass_profiler_controller.cc new file mode 100644 index 000000000..58f54400d --- /dev/null +++ b/tsl/profiler/lib/multi_pass_profiler_controller.cc @@ -0,0 +1,176 @@ +/* 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 +#include + +#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 profiler) + : profiler_(std::move(profiler)) {} + +MultiPassProfilerController::~MultiPassProfilerController() { + // Ensure a successfully started pass is stopped. + if (state_ == MultiPassProfilerState::kPassStarted && status_.ok()) { + profiler_->StopPass().IgnoreError(); + state_ = MultiPassProfilerState::kPassStopped; + } + 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); + } 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 (status_.ok()) { + status = status_ = profiler_->PopRange(); + } else { + 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::StopPass() { + absl::Status status; + if (state_ == MultiPassProfilerState::kPassStarted) { + if (status_.ok()) { + status = status_ = profiler_->StopPass(); + state_ = MultiPassProfilerState::kPassStopped; + } else { + 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 || + state_ == MultiPassProfilerState::kPassStopped || + state_ == MultiPassProfilerState::kStarted) { + state_ = MultiPassProfilerState::kStopped; + if (status_.ok()) { + status = status_ = profiler_->Stop(); + } else { + status = absl::AbortedError("Previous call returned an error."); + } + } else { + status = absl::AbortedError("Stop called in the wrong order"); + } + if (!status.ok()) LOG(ERROR) << status; + return status; +} + +} // namespace profiler + +} // namespace tsl diff --git a/tsl/profiler/lib/multi_pass_profiler_controller.h b/tsl/profiler/lib/multi_pass_profiler_controller.h new file mode 100644 index 000000000..0bdf3f0ef --- /dev/null +++ b/tsl/profiler/lib/multi_pass_profiler_controller.h @@ -0,0 +1,65 @@ +/* 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 + +#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 { + +class MultiPassProfilerController : public MultiPassProfilerInterface { + public: + explicit MultiPassProfilerController( + std::unique_ptr 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, + }; + + MultiPassProfilerState state_ = MultiPassProfilerState::kInit; + std::unique_ptr profiler_; + absl::Status status_; // result of calls to profiler_ +}; + +} // namespace profiler +} // namespace tsl + +#endif // TENSORFLOW_TSL_PROFILER_LIB_MULTI_PASS_PROFILER_CONTROLLER_H_ diff --git a/tsl/profiler/lib/profiler_factory.cc b/tsl/profiler/lib/profiler_factory.cc index a11a11b50..c18a10541 100644 --- a/tsl/profiler/lib/profiler_factory.cc +++ b/tsl/profiler/lib/profiler_factory.cc @@ -19,7 +19,9 @@ limitations under the License. #include #include "absl/base/const_init.h" +#include "absl/log/log.h" #include "absl/synchronization/mutex.h" +#include "tsl/profiler/lib/multi_pass_profiler_controller.h" #include "tsl/profiler/lib/profiler_controller.h" #include "tsl/profiler/lib/profiler_interface.h" #include "tsl/profiler/protobuf/profiler_options.pb.h" @@ -29,12 +31,18 @@ namespace profiler { namespace { absl::Mutex mu(absl::kConstInit); +absl::Mutex mu_multipass(absl::kConstInit); std::vector* GetFactories() { static auto factories = new std::vector(); return factories; } +std::vector* GetMultiPassFactories() { + static auto factories = new std::vector(); + return factories; +} + } // namespace void RegisterProfilerFactory(ProfilerFactory factory) { @@ -42,6 +50,11 @@ void RegisterProfilerFactory(ProfilerFactory factory) { GetFactories()->push_back(std::move(factory)); } +void RegisterMultiPassProfilerFactory(MultiPassProfilerFactory factory) { + absl::MutexLock lock(mu_multipass); + GetMultiPassFactories()->push_back(std::move(factory)); +} + std::vector> CreateProfilers( const tensorflow::ProfileOptions& options) { std::vector> result; @@ -56,10 +69,36 @@ std::vector> CreateProfilers( return result; } +std::vector> +CreateMultiPassProfilers(const tensorflow::ProfileOptions& options) { + std::vector> multipass_profilers; + { + absl::MutexLock lock(mu_multipass); + VLOG(3) << "=====Creating MultiPassProfilers() with " + << GetMultiPassFactories()->size() << " factories."; + for (const auto& factory : *GetMultiPassFactories()) { + auto profiler = factory(options); + if (profiler == nullptr) { + continue; + } + multipass_profilers.emplace_back( + std::make_unique(std::move(profiler))); + } + VLOG(3) << "=====Created " << multipass_profilers.size() + << " multipass profilers."; + } + return multipass_profilers; +} + void ClearRegisteredProfilersForTest() { absl::MutexLock lock(mu); GetFactories()->clear(); } +void ClearRegisteredMultiPassProfilersForTest() { + absl::MutexLock lock(mu_multipass); + GetMultiPassFactories()->clear(); +} + } // namespace profiler } // namespace tsl diff --git a/tsl/profiler/lib/profiler_factory.h b/tsl/profiler/lib/profiler_factory.h index 7e6c64c39..2cf1b07e1 100644 --- a/tsl/profiler/lib/profiler_factory.h +++ b/tsl/profiler/lib/profiler_factory.h @@ -30,16 +30,28 @@ namespace profiler { using ProfilerFactory = std::function( const tensorflow::ProfileOptions&)>; +using MultiPassProfilerFactory = + std::function( + const tensorflow::ProfileOptions&)>; + // Registers a profiler factory. Should be invoked at most once per factory. void RegisterProfilerFactory(ProfilerFactory factory); +void RegisterMultiPassProfilerFactory(MultiPassProfilerFactory factory); + // Invokes all registered profiler factories with the given options, and // returns the instantiated (non-null) profiler interfaces. std::vector> CreateProfilers( const tensorflow::ProfileOptions& options); +// Invokes all registered multi-pass profiler factories with the given options, +// and returns the instantiated (non-null) multi-pass profiler interfaces. +std::vector> +CreateMultiPassProfilers(const tensorflow::ProfileOptions& options); + // For testing only. void ClearRegisteredProfilersForTest(); +void ClearRegisteredMultiPassProfilersForTest(); } // namespace profiler } // namespace tsl diff --git a/tsl/profiler/lib/profiler_factory_test.cc b/tsl/profiler/lib/profiler_factory_test.cc index 97dc6f914..a05e0559a 100644 --- a/tsl/profiler/lib/profiler_factory_test.cc +++ b/tsl/profiler/lib/profiler_factory_test.cc @@ -19,6 +19,7 @@ limitations under the License. #include "absl/memory/memory.h" #include "absl/status/status.h" +#include "absl/strings/string_view.h" #include "xla/tsl/platform/macros.h" #include "xla/tsl/platform/test.h" #include "tsl/profiler/lib/profiler_interface.h" @@ -98,6 +99,57 @@ TEST(ProfilerFactoryTest, FactoryClassCapturedByLambda) { EXPECT_EQ(profilers.size(), 1); } +class TestMultiPassProfiler : public MultiPassProfilerInterface { + public: + bool NeedMorePasses() override { return false; } + absl::Status StartPass() override { return absl::OkStatus(); } + absl::Status PushRange(absl::string_view name) override { + return absl::OkStatus(); + } + absl::Status PopRange() override { return absl::OkStatus(); } + absl::Status StopPass() override { return absl::OkStatus(); } + + absl::Status Start() override { return absl::OkStatus(); } + absl::Status Stop() override { return absl::OkStatus(); } + absl::Status CollectData(tensorflow::profiler::XSpace*) override { + return absl::OkStatus(); + } +}; + +std::unique_ptr TestMultiPassFactoryFunction( + const tensorflow::ProfileOptions& options) { + return absl::make_unique(); +} + +TEST(ProfilerFactoryTest, MultiPassFactoryFunctionPointer) { + ClearRegisteredMultiPassProfilersForTest(); + RegisterMultiPassProfilerFactory(&TestMultiPassFactoryFunction); + auto profilers = CreateMultiPassProfilers(tensorflow::ProfileOptions()); + EXPECT_EQ(profilers.size(), 1); +} + +TEST(ProfilerFactoryTest, MultiPassFactoryLambda) { + ClearRegisteredMultiPassProfilersForTest(); + RegisterMultiPassProfilerFactory( + [](const tensorflow::ProfileOptions& options) { + return absl::make_unique(); + }); + auto profilers = CreateMultiPassProfilers(tensorflow::ProfileOptions()); + EXPECT_EQ(profilers.size(), 1); +} + +std::unique_ptr NullMultiPassFactoryFunction( + const tensorflow::ProfileOptions& options) { + return nullptr; +} + +TEST(ProfilerFactoryTest, MultiPassFactoryReturnsNull) { + ClearRegisteredMultiPassProfilersForTest(); + RegisterMultiPassProfilerFactory(&NullMultiPassFactoryFunction); + auto profilers = CreateMultiPassProfilers(tensorflow::ProfileOptions()); + EXPECT_TRUE(profilers.empty()); +} + } // namespace } // namespace profiler } // namespace tsl diff --git a/tsl/profiler/lib/profiler_interface.h b/tsl/profiler/lib/profiler_interface.h index 2b0b71242..011959910 100644 --- a/tsl/profiler/lib/profiler_interface.h +++ b/tsl/profiler/lib/profiler_interface.h @@ -15,6 +15,8 @@ limitations under the License. #ifndef TENSORFLOW_TSL_PROFILER_LIB_PROFILER_INTERFACE_H_ #define TENSORFLOW_TSL_PROFILER_LIB_PROFILER_INTERFACE_H_ +#include "absl/status/status.h" +#include "absl/strings/string_view.h" #include "xla/tsl/platform/status.h" #include "tsl/profiler/protobuf/xplane.pb.h" @@ -43,6 +45,27 @@ class ProfilerInterface { virtual absl::Status CollectData(tensorflow::profiler::XSpace* space) = 0; }; +// XProfMultipass manages multi-pass profiling. It plans the passes, +// configures the profiler for each pass, and aggregates the results. +class MultiPassProfilerInterface : public ProfilerInterface { + public: + // Returns true if there are more passes to profile. + virtual bool NeedMorePasses() = 0; + + // Starts a new pass. + virtual absl::Status StartPass() = 0; + + virtual absl::Status PushRange(absl::string_view name) = 0; + + virtual absl::Status PopRange() = 0; + + virtual absl::Status StopPass() = 0; + + absl::Status Start() override { return absl::OkStatus(); } + + absl::Status Stop() override { return absl::OkStatus(); } +}; + } // namespace profiler } // namespace tsl diff --git a/tsl/profiler/lib/profiler_passes.cc b/tsl/profiler/lib/profiler_passes.cc new file mode 100644 index 000000000..6311c2a5b --- /dev/null +++ b/tsl/profiler/lib/profiler_passes.cc @@ -0,0 +1,205 @@ +/* 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/profiler_passes.h" + +#include +#include + +#include "absl/memory/memory.h" +#include "absl/status/status.h" +#include "absl/status/status_macros.h" +#include "absl/strings/string_view.h" +#include "absl/synchronization/mutex.h" +#include "xla/tsl/platform/logging.h" +#include "xla/tsl/profiler/utils/xplane_builder.h" +#include "xla/tsl/profiler/utils/xplane_schema.h" +#include "xla/tsl/profiler/utils/xplane_utils.h" +#include "tsl/profiler/lib/profiler_interface.h" +#include "tsl/profiler/protobuf/profiler_options.pb.h" +#include "tsl/profiler/protobuf/xplane.pb.h" + +#if !defined(IS_MOBILE_PLATFORM) +#include "xla/tsl/platform/env.h" +#include "xla/tsl/profiler/convert/post_process_single_host_xplane.h" +#include "xla/tsl/profiler/utils/time_utils.h" +#include "tsl/platform/host_info.h" +#include "tsl/profiler/lib/profiler_factory.h" +#include "tsl/profiler/lib/profiler_lock.h" +#endif + +namespace tsl { +namespace { + +using tensorflow::ProfileOptions; +using tensorflow::profiler::XSpace; +using ::tsl::profiler::XPlaneBuilder; + +void SetProfileOptionsIntoSpace(const ProfileOptions& options, XSpace* space) { + XPlaneBuilder xplane(profiler::FindOrAddMutablePlaneWithName( + space, tsl::profiler::kTaskEnvPlaneName)); + xplane.AddStatValue( + *xplane.GetOrCreateStatMetadata(tsl::profiler::GetTaskEnvStatTypeStr( + tsl::profiler::kEnvProfileOptions)), + options); +} + +} // namespace + +std::unique_ptr ProfilerPasses::Create( + const ProfileOptions& options) { + return absl::WrapUnique(new ProfilerPasses(options)); +} + +absl::Status ProfilerPasses::Status() { + absl::MutexLock l(&mutex_); + return status_; +} + +bool ProfilerPasses::NeedMorePasses() { + absl::MutexLock l(&mutex_); + if (!status_.ok()) return false; + if (multi_pass_profiler_ == nullptr) { + return false; + } + return multi_pass_profiler_->NeedMorePasses(); +} + +absl::Status ProfilerPasses::StartPass() { + absl::MutexLock l(&mutex_); + if (!status_.ok()) { + return absl::InternalError("Previous operations failed"); + } + if (multi_pass_profiler_ == nullptr) { + return absl::OkStatus(); + } + + if (pass_count_++ == 0) { + first_pass_start_time_ns_ = profiler::GetCurrentTimeNanos(); + absl::Status status = multi_pass_profiler_->Start(); + if (options_.raise_error_on_start_failure()) { + status_ = status; + } else { + status.IgnoreError(); + } + RETURN_IF_ERROR(status_); + + status = multi_pass_profiler_->StartPass(); + if (options_.raise_error_on_start_failure()) { + status_ = status; + } else { + status.IgnoreError(); + } + return status_; + } + + // passes after first pass. + return status_ = multi_pass_profiler_->StartPass(); +} + +absl::Status ProfilerPasses::PushRange(absl::string_view name) { + absl::MutexLock l(&mutex_); + RETURN_IF_ERROR(status_); + if (multi_pass_profiler_ != nullptr) { + return status_ = multi_pass_profiler_->PushRange(name); + } + return absl::OkStatus(); +} + +absl::Status ProfilerPasses::PopRange() { + absl::MutexLock l(&mutex_); + RETURN_IF_ERROR(status_); + if (multi_pass_profiler_ != nullptr) { + return status_ = multi_pass_profiler_->PopRange(); + } + return absl::OkStatus(); +} + +absl::Status ProfilerPasses::StopPass() { + absl::MutexLock l(&mutex_); + if (pass_count_ == 1) { + first_pass_stop_time_ns_ = profiler::GetCurrentTimeNanos(); + } + if (multi_pass_profiler_ != nullptr) { + status_.Update(multi_pass_profiler_->StopPass()); + } + return status_; +} + +#if !defined(IS_MOBILE_PLATFORM) +absl::Status ProfilerPasses::CollectDataInternal(XSpace* space) { + absl::MutexLock l(&mutex_); + RETURN_IF_ERROR(status_); + VLOG(3) << "Profiler passes collecting data."; + if (multi_pass_profiler_ != nullptr) { + multi_pass_profiler_->Stop().IgnoreError(); + multi_pass_profiler_->CollectData(space).IgnoreError(); + multi_pass_profiler_.reset(); + } + // Allow another session to start. + profiler_lock_.ReleaseIfActive(); + return absl::OkStatus(); +} +#endif + +absl::Status ProfilerPasses::CollectData(XSpace* space) { +#if !defined(IS_MOBILE_PLATFORM) + space->add_hostnames(port::Hostname()); + RETURN_IF_ERROR(CollectDataInternal(space)); + profiler::SetXSpacePidIfNotSet(*space, tsl::Env::Default()->GetProcessId()); + profiler::PostProcessSingleHostXSpace(space, first_pass_start_time_ns_, + first_pass_stop_time_ns_); +#endif + SetProfileOptionsIntoSpace(options_, space); + return absl::OkStatus(); +} + +ProfilerPasses::ProfilerPasses(const ProfileOptions& options) +#if defined(IS_MOBILE_PLATFORM) + : status_(absl::UnimplementedError( + "Profiler is unimplemented for mobile platforms.")) { +#else + : options_(options) { + auto profiler_lock = profiler::ProfilerLock::Acquire(); + if (!profiler_lock.ok()) { + status_ = profiler_lock.status(); + return; + } + profiler_lock_ = *std::move(profiler_lock); + + DCHECK(profiler_lock_.Active()); + VLOG(3) << "======== Profiler passes initializing. options.enable_multipass =" + << options_.enable_multipass(); + auto multi_pass_profilers = profiler::CreateMultiPassProfilers(options_); + if (!multi_pass_profilers.empty()) { + multi_pass_profiler_ = std::move(multi_pass_profilers[0]); + if (multi_pass_profilers.size() > 1) { + LOG(ERROR) + << "More than one multipass-planner found, only use one of them."; + } + VLOG(3) << "======== Find one multipass planner!"; + } else { + VLOG(3) << "======== No multipass planner found!"; + } +#endif +} + +ProfilerPasses::~ProfilerPasses() { +#if !defined(IS_MOBILE_PLATFORM) + VLOG(3) << "Profiler passes tear down."; +#endif +} + +} // namespace tsl diff --git a/tsl/profiler/lib/profiler_passes.h b/tsl/profiler/lib/profiler_passes.h new file mode 100644 index 000000000..4d8970fa1 --- /dev/null +++ b/tsl/profiler/lib/profiler_passes.h @@ -0,0 +1,98 @@ +/* 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_PROFILER_PASSES_H_ +#define TENSORFLOW_TSL_PROFILER_LIB_PROFILER_PASSES_H_ + +#include +#include + +#include "absl/base/thread_annotations.h" +#include "absl/status/status.h" +#include "absl/strings/string_view.h" +#include "absl/synchronization/mutex.h" +#include "tsl/profiler/lib/profiler_collection.h" +#include "tsl/profiler/protobuf/profiler_options.pb.h" +#include "tsl/profiler/protobuf/xplane.pb.h" + +#if !defined(IS_MOBILE_PLATFORM) +#include "tsl/profiler/lib/profiler_interface.h" +#include "tsl/profiler/lib/profiler_lock.h" +#endif + +namespace tsl { + +class ProfilerPasses { + public: + // Creates a ProfilerPasses and starts profiling. + static std::unique_ptr Create( + const tensorflow::ProfileOptions& options); + + static tensorflow::ProfileOptions DefaultOptions() { + tensorflow::ProfileOptions options; + options.set_version(1); + options.set_device_tracer_level(1); + options.set_host_tracer_level(2); + options.set_device_type(tensorflow::ProfileOptions::UNSPECIFIED); + options.set_python_tracer_level(0); + options.set_enable_hlo_proto(true); + options.set_include_dataset_ops(true); + options.set_enable_multipass(true); + return options; + } + + ~ProfilerPasses(); + + absl::Status Status() ABSL_LOCKS_EXCLUDED(mutex_); + + bool NeedMorePasses() ABSL_LOCKS_EXCLUDED(mutex_); + + absl::Status StartPass() ABSL_LOCKS_EXCLUDED(mutex_); + + absl::Status PushRange(absl::string_view name) ABSL_LOCKS_EXCLUDED(mutex_); + + absl::Status PopRange() ABSL_LOCKS_EXCLUDED(mutex_); + + absl::Status StopPass() ABSL_LOCKS_EXCLUDED(mutex_); + + // Collects profile data into XSpace. + absl::Status CollectData(tensorflow::profiler::XSpace* space) + ABSL_LOCKS_EXCLUDED(mutex_); + + private: + explicit ProfilerPasses(const tensorflow::ProfileOptions& options); + + ProfilerPasses(const ProfilerPasses&) = delete; + ProfilerPasses& operator=(const ProfilerPasses&) = delete; + +#if !defined(IS_MOBILE_PLATFORM) + absl::Status CollectDataInternal(tensorflow::profiler::XSpace* space); + + profiler::ProfilerLock profiler_lock_ ABSL_GUARDED_BY(mutex_); + std::unique_ptr multi_pass_profiler_ + ABSL_GUARDED_BY(mutex_); +#endif + + absl::Mutex mutex_; + absl::Status status_ ABSL_GUARDED_BY(mutex_); + tensorflow::ProfileOptions options_; + int pass_count_ ABSL_GUARDED_BY(mutex_) = 0; + int64_t first_pass_start_time_ns_ = 0; + int64_t first_pass_stop_time_ns_ = 0; +}; + +} // namespace tsl + +#endif // TENSORFLOW_TSL_PROFILER_LIB_PROFILER_PASSES_H_ diff --git a/tsl/profiler/protobuf/profiler_options.proto b/tsl/profiler/protobuf/profiler_options.proto index 04c46af2d..42cab8d07 100644 --- a/tsl/profiler/protobuf/profiler_options.proto +++ b/tsl/profiler/protobuf/profiler_options.proto @@ -2,7 +2,7 @@ syntax = "proto3"; package tensorflow; -// Next ID: 16 +// Next ID: 17 message ProfileOptions { // Some default value of option are not proto3 default value. Use this version // to determine if we should use default option value instead of proto3 @@ -103,6 +103,9 @@ message ProfileOptions { // If set, this hostname will be used to name the profile file. string override_hostname = 15; + + // Programmatic replay multipass enabled or not. + bool enable_multipass = 16; } // Options for remote profiler session manager.