Add lightweight, zero-overhead tracing macros in tpu_sync/common/trace.h supporting both TraceMe and Perfetto backends.

MODULE.bazel

@@ -89,6 +89,12 @@ bazel_dep(name = "boringssl", version = "0.20240913.0")
 llvm = use_extension("@xla//third_party/extensions:llvm.bzl", "llvm_extension")
 use_repo(llvm, "llvm-project")
 
+tsl_extension = use_extension("@xla//third_party/extensions:tsl.bzl", "tsl_extension")
+use_repo(
+    tsl_extension,
+    "tsl",
+)
+
 jax_python_wheel_repository = use_repo_rule(
     "@jax//jaxlib:jax_python_wheel.bzl",
     "jax_python_wheel_repository",
@@ -102,6 +108,7 @@ bazel_dep(name = "googletest", version = "1.17.0.bcr.2", repo_name = "com_google
 bazel_dep(name = "flatbuffers", version = "25.12.19", repo_name = "com_github_google_flatbuffers")
 bazel_dep(name = "highway", version = "1.2.0")
 bazel_dep(name = "prometheus-cpp", version = "1.3.0", repo_name = "com_github_jupp0r_prometheus_cpp")
+bazel_dep(name = "perfetto")
 
 pybind11_internal_configure = use_extension(
     "@pybind11_bazel//:internal_configure.bzl",
@@ -143,6 +150,16 @@ git_override(
     remote = "https://github.com/google/nanobind_abseil.git",
 )
 
+git_override(
+    module_name = "perfetto",
+    commit = "d0d8c6fa5f7411b59fccacf99233e79bcf94e3fc",
+    patch_strip = 1,
+    patches = [
+        "//:third_party/perfetto/no_android_deps.patch",
+    ],
+    remote = "https://github.com/google/perfetto.git",
+)
+
 # Single-sources the wheel version from pyproject.toml [project].version and
 # appends WHEEL_VERSION_EXTRAS (e.g. ".dev<timestamp>"). Consumed by
 # //ci/wheel:raiden_wheel via @raiden_version//:version.bzl.

third_party/perfetto/no_android_deps.patch

@@ -0,0 +1,117 @@
+diff --git a/MODULE.bazel b/MODULE.bazel
+index 10941b7f..ea7db590 100644
+--- a/MODULE.bazel
++++ b/MODULE.bazel
+@@ -31,50 +31,9 @@ bazel_dep(name = "platforms", version = "0.0.10")
+ bazel_dep(name = "protobuf", version = "31.1", repo_name = "com_google_protobuf")
+ bazel_dep(name = "re2", version = "2024-07-02.bcr.1")
+ bazel_dep(name = "rules_python", version = "1.0.0")
+-bazel_dep(name = "rules_android", version = "0.6.6")
+ bazel_dep(name = "rules_cc", version = "0.2.19")
+ bazel_dep(name = "zstd", version = "1.5.7.bcr.1")
+
+-remote_android_extensions = use_extension(
+-    "@rules_android//bzlmod_extensions:android_extensions.bzl",
+-    "remote_android_tools_extensions",
+-)
+-use_repo(remote_android_extensions, "android_tools")
+-
+-android_sdk_repository_extension = use_extension(
+-    "@rules_android//rules/android_sdk_repository:rule.bzl",
+-    "android_sdk_repository_extension",
+-)
+-
+-# When built using 'tools/bazel', 'ANDROID_HOME' environment variable points
+-# to the hermetic Android SDK installation, that should be downloaded first
+-# with 'tools/install-build-deps --android'
+-android_sdk_repository_extension.configure(
+-    api_level = 35,
+-    build_tools_version = "35.0.1",
+-)
+-use_repo(android_sdk_repository_extension, "androidsdk")
+-
+-register_toolchains("@androidsdk//:sdk-toolchain", "@androidsdk//:all")
+-
+-bazel_dep(name = "rules_jvm_external", version = "6.9")
+-bazel_dep(name = "rules_android_ndk", version = "0.1.3")
+-
+-android_ndk_repository_extension = use_extension(
+-    "@rules_android_ndk//:extension.bzl",
+-    "android_ndk_repository_extension",
+-)
+-
+-# When built using 'tools/bazel', 'ANDROID_NDK_HOME' environment variable points
+-# to the hermetic Android NDK installation, that should be downloaded first
+-# with 'tools/install-build-deps --android'
+-android_ndk_repository_extension.configure(
+-    api_level = 26,
+-)
+-use_repo(android_ndk_repository_extension, "androidndk")
+-
+-register_toolchains("@androidndk//:all")
+-
+ # Perfetto configuration repository extension.
+ # This creates @perfetto_cfg which provides PERFETTO_CONFIG struct.
+ perfetto_cfg_ext = use_extension(
+@@ -100,27 +59,3 @@ use_repo(
+     "perfetto_dep_zlib",
+ )
+
+-# Maven dependencies for Android instrumentation tests.
+-# Use a unique name to avoid conflicts with other modules (e.g., bazel_worker_java).
+-maven = use_extension(
+-    "@rules_jvm_external//:extensions.bzl",
+-    "maven",
+-)
+-maven.install(
+-    name = "perfetto_maven",
+-    # Use rules_android's aar_import to avoid toolchain type mismatch.
+-    aar_import_bzl_label = "@rules_android//rules:rules.bzl",
+-    artifacts = [
+-        "androidx.test:runner:1.6.2",
+-        "androidx.test:monitor:1.7.2",
+-        "com.google.truth:truth:1.4.4",
+-        "junit:junit:4.13.2",
+-        "androidx.test.ext:junit:1.2.1",
+-        "com.google.errorprone:error_prone_annotations:2.36.0",
+-    ],
+-    repositories = [
+-        "https://maven.google.com",
+-        "https://repo1.maven.org/maven2",
+-    ],
+-)
+-use_repo(maven, "perfetto_maven")
+diff --git a/bazel/rules.bzl b/bazel/rules.bzl
+index 958e9bc0..f5990adf 100644
+--- a/bazel/rules.bzl
++++ b/bazel/rules.bzl
+@@ -13,9 +13,16 @@
+ # limitations under the License.
+
+ load("@perfetto//bazel:proto_gen.bzl", "proto_descriptor_gen", "proto_gen")
+-load("@perfetto//bazel:run_ait_with_adb.bzl", "android_instrumentation_test")
+ load("@perfetto_cfg//:perfetto_cfg.bzl", "PERFETTO_CONFIG")
+-load("@rules_android//android:rules.bzl", "android_binary", "android_library")
++
++def android_binary(**kwargs):
++    pass
++
++def android_library(**kwargs):
++    pass
++
++def android_instrumentation_test(**kwargs):
++    pass
+
+ # +----------------------------------------------------------------------------+
+ # | Base C++ rules.                                                            |
+diff --git a/bazel/standalone/perfetto_cfg.bzl b/bazel/standalone/perfetto_cfg.bzl
+--- a/bazel/standalone/perfetto_cfg.bzl
++++ b/bazel/standalone/perfetto_cfg.bzl
+@@ -53,3 +53,6 @@ PERFETTO_CONFIG = struct(
+         protoc = ["@com_google_protobuf//:protoc"],
+-        protoc_lib = ["@com_google_protobuf//:protoc_lib"],
++        protoc_lib = [
++            "@com_google_protobuf//src/google/protobuf/compiler:code_generator",
++            "@com_google_protobuf//src/google/protobuf/compiler:importer",
++        ],
+         protobuf_lite = ["@com_google_protobuf//:protobuf_lite"],

tpu_sync/common/BUILD

@@ -48,6 +48,34 @@ cc_test(
     ],
 )
 
+cc_library(
+    name = "trace",
+    srcs = ["trace.cc"],
+    hdrs = ["trace.h"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "@com_google_absl//absl/base",
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/synchronization",
+        "@perfetto//:libperfetto_client_experimental",
+        "@tsl//tsl/profiler/lib:traceme",
+        "@tsl//tsl/profiler/lib:traceme_encode",
+    ],
+)
+
+cc_test(
+    name = "trace_test",
+    srcs = ["trace_test.cc"],
+    deps = [
+        ":trace",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/time",
+        "@com_google_googletest//:gtest_main",
+        "@perfetto//:libperfetto_client_experimental",
+    ],
+)
+
 nanobind_extension(
     name = "_raiden_id",
     srcs = ["raiden_id_binding.cc"],