diff --git a/tsl/profiler/lib/BUILD b/tsl/profiler/lib/BUILD index 1071f08e2..2b4301f07 100644 --- a/tsl/profiler/lib/BUILD +++ b/tsl/profiler/lib/BUILD @@ -28,6 +28,18 @@ cc_library( visibility = ["//visibility:public"], ) +cc_library( + name = "traceme_global_flags", + srcs = ["traceme_global_flags.cc"], + hdrs = ["traceme_global_flags.h"], + copts = tf_profiler_copts(), + visibility = ["//visibility:public"], + deps = [ + "@xla//xla/tsl/platform:macros", + ], + alwayslink = True, +) + filegroup( name = "mobile_srcs_no_runtime", srcs = [ diff --git a/tsl/profiler/lib/traceme_global_flags.cc b/tsl/profiler/lib/traceme_global_flags.cc new file mode 100644 index 000000000..2e4dcacdb --- /dev/null +++ b/tsl/profiler/lib/traceme_global_flags.cc @@ -0,0 +1,33 @@ +/* Copyright 2026 The OpenXLA Authors. + +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/traceme_global_flags.h" + +#include + +namespace tsl { +namespace profiler { + +#ifdef _WIN32 +#define DECL_DLL_EXPORT __declspec(dllexport) +#else +#define DECL_DLL_EXPORT +#endif +// DLL imported variables cannot be initialized on Windows. This file is +// included only on DLL exports. +DECL_DLL_EXPORT std::atomic g_enable_source_location(true); + +} // namespace profiler +} // namespace tsl diff --git a/tsl/profiler/lib/traceme_global_flags.h b/tsl/profiler/lib/traceme_global_flags.h new file mode 100644 index 000000000..f3bc29130 --- /dev/null +++ b/tsl/profiler/lib/traceme_global_flags.h @@ -0,0 +1,44 @@ +/* Copyright 2026 The OpenXLA Authors. + +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_TRACEME_GLOBAL_FLAGS_H_ +#define TENSORFLOW_TSL_PROFILER_LIB_TRACEME_GLOBAL_FLAGS_H_ + +#include +#include + +#include + +#include "xla/tsl/platform/macros.h" + +// Flags that are global to the TraceMe implementation, which may be shared by +// TraceMe and TraceMeRecorder. +namespace tsl { +namespace profiler { + +TF_EXPORT extern std::atomic g_enable_source_location; + +class TraceMeGlobalFlags { + public: + // Returns whether source location is enabled. + static bool IsSourceLocationEnabled() { + return g_enable_source_location.load(std::memory_order_relaxed); + } +}; + +} // namespace profiler +} // namespace tsl + +#endif // TENSORFLOW_TSL_PROFILER_LIB_TRACEME_GLOBAL_FLAGS_H_