tpu_sync: add POSIX shared-memory layout for telemetry

Multi-worker TPU training setups require aggregated host-level metrics across multiple isolated worker processes without inter-process lock contention or high CPU serialization overhead.

Define 64-byte aligned, lock-free POSIX shared memory layout (ShmHistogramSlot and ShmStoreLayout) for atomic counter, gauge, and histogram metric updates.

PiperOrigin-RevId: 970337171

Author: mrtarunarora-google

tpu_sync/telemetry/shm/BUILD

@@ -0,0 +1,36 @@
+# Copyright 2026 Google LLC.
+#
+# 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.
+
+load("@rules_cc//cc:defs.bzl", "cc_library", "cc_test")
+
+package(default_visibility = ["//visibility:public"])
+
+cc_library(
+    name = "shm_layout",
+    hdrs = ["shm_layout.h"],
+    deps = [
+        "//tpu_sync/telemetry:metrics_backend",
+        "@com_google_absl//absl/types:span",
+    ],
+)
+
+cc_test(
+    name = "shm_layout_test",
+    srcs = ["shm_layout_test.cc"],
+    deps = [
+        ":shm_layout",
+        "//tpu_sync/telemetry:metrics_backend",
+        "@com_google_googletest//:gtest_main",
+    ],
+)

tpu_sync/telemetry/shm/shm_layout.h

@@ -0,0 +1,130 @@
+// Copyright 2026 Google LLC.
+//
+// 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 THIRD_PARTY_TPU_RAIDEN_TPU_SYNC_TELEMETRY_SHM_SHM_LAYOUT_H_
+#define THIRD_PARTY_TPU_RAIDEN_TPU_SYNC_TELEMETRY_SHM_SHM_LAYOUT_H_
+
+#include <atomic>
+#include <cmath>
+#include <cstddef>
+#include <cstdint>
+#include <iterator>
+#include <type_traits>
+
+#include "tpu_sync/telemetry/metrics_backend.h"
+
+namespace tpu_raiden::telemetry {
+
+inline constexpr uint32_t kRaidenShmMagic = 0xABCD1234;
+inline constexpr uint32_t kSupportedVersion = 2;
+inline constexpr size_t kMaxTocEntries = 1024;
+inline constexpr size_t kMaxDataPoolBytes = 64 * 1024;  // 64 KB
+inline constexpr size_t kNumHistogramBuckets =
+    std::size(kDefaultHistogramBuckets);
+
+enum class TocEntryState : uint32_t {
+  kUninitialized = 0,
+  kWriting = 1,
+  kCommitted = 2,
+};
+
+// Table of Contents entry descriptor for a single metric stream.
+struct alignas(64) ShmTocEntry {
+  char metric_name[64];
+  char encoded_labels[128];
+  MetricType type;
+  uint32_t offset;
+  uint32_t size;
+  std::atomic<uint32_t> entry_state{
+      static_cast<uint32_t>(TocEntryState::kUninitialized)};
+  uint8_t padding[48]{};
+};
+
+// 64-byte aligned header for the shared-memory telemetry segment.
+struct alignas(64) ShmTocHeader {
+  std::atomic<uint32_t> magic{0};
+  uint32_t version{kSupportedVersion};
+  int64_t pid{0};
+  std::atomic<uint32_t> toc_entry_count{0};
+  uint32_t max_toc_entries{kMaxTocEntries};
+  uint32_t data_pool_offset{0};
+  std::atomic<uint32_t> data_pool_bytes{0};
+  uint32_t chunk_index{0};
+  uint8_t padding[28]{};
+};
+
+// Lock-free shared-memory slot for a single histogram metric stream.
+struct alignas(64) ShmHistogramSlot {
+  std::atomic<uint64_t> sample_count{0};
+  std::atomic<double> sample_sum{0.0};
+  std::atomic<uint64_t> bucket_counts[kNumHistogramBuckets + 1]{};
+  uint8_t padding[8]{};
+
+  void Observe(double value) {
+    if (!std::isfinite(value)) return;
+    sample_count.fetch_add(1, std::memory_order_relaxed);
+    double current_sum = sample_sum.load(std::memory_order_relaxed);
+    while (!sample_sum.compare_exchange_weak(current_sum, current_sum + value,
+                                             std::memory_order_relaxed,
+                                             std::memory_order_relaxed)) {
+    }
+    for (size_t i = 0; i < kNumHistogramBuckets; ++i) {
+      if (value <= kDefaultHistogramBuckets[i]) {
+        bucket_counts[i].fetch_add(1, std::memory_order_relaxed);
+      }
+    }
+    bucket_counts[kNumHistogramBuckets].fetch_add(1, std::memory_order_relaxed);
+  }
+};
+
+// Complete 64-byte aligned segment layout containing header and TOC entries.
+struct alignas(64) ShmSegmentLayout {
+  ShmTocHeader header;
+  ShmTocEntry toc[kMaxTocEntries];
+};
+
+constexpr size_t CalculateChunkFileSize(size_t max_toc_entries,
+                                        size_t data_pool_bytes) {
+  return sizeof(ShmTocHeader) + (sizeof(ShmTocEntry) * max_toc_entries) +
+         data_pool_bytes;
+}
+
+inline constexpr size_t kSegmentTotalFileSize =
+    CalculateChunkFileSize(kMaxTocEntries, kMaxDataPoolBytes);
+
+static_assert(std::is_standard_layout_v<ShmTocEntry> &&
+              std::is_standard_layout_v<ShmTocHeader> &&
+              std::is_standard_layout_v<ShmHistogramSlot> &&
+              std::is_standard_layout_v<ShmSegmentLayout>);
+
+static_assert(std::is_trivially_copyable_v<ShmTocEntry> &&
+              std::is_trivially_copyable_v<ShmTocHeader> &&
+              std::is_trivially_copyable_v<ShmHistogramSlot> &&
+              std::is_trivially_copyable_v<ShmSegmentLayout>);
+
+static_assert(alignof(ShmTocHeader) == 64 && sizeof(ShmTocHeader) == 64);
+static_assert(alignof(ShmTocEntry) == 64 && sizeof(ShmTocEntry) == 256);
+static_assert(alignof(ShmHistogramSlot) == 64 &&
+              sizeof(ShmHistogramSlot) == 192);
+static_assert(alignof(ShmSegmentLayout) == 64 &&
+              sizeof(ShmSegmentLayout) == 262208);
+
+static_assert(std::atomic<uint64_t>::is_always_lock_free &&
+              std::atomic<int64_t>::is_always_lock_free &&
+              std::atomic<double>::is_always_lock_free &&
+              std::atomic<uint32_t>::is_always_lock_free);
+
+}  // namespace tpu_raiden::telemetry
+
+#endif  // THIRD_PARTY_TPU_RAIDEN_TPU_SYNC_TELEMETRY_SHM_SHM_LAYOUT_H_