Skip to content

include/tenzr/common/trace.h

Namespaces

Name
tenzr
tenzr::profiler

Classes

Name
struct tenzr::profiler::TraceEvent
One Chrome Trace Event entry recorded by the profiler.
struct tenzr::profiler::TopOp
Aggregated timing statistics for one operation in a profile.
struct tenzr::profiler::BottleneckReport
Heuristic analysis of the dominant performance bottleneck.
struct tenzr::profiler::ProfileSummary
Aggregated profile statistics for a recorded trace.
class tenzr::profiler::Profiler
Process-wide collector of profile events with Chrome Trace export.
class tenzr::profiler::ScopedTrace
RAII helper that records a duration for its enclosing scope.
class tenzr::profiler::TraceSession
RAII helper that starts profiling on construction and exports on destruction.
class tenzr::profiler::StepMarker
RAII helper that marks the start and end of a training step.

Defines

Name
TENZR_TRACE_SCOPE(name)
TENZR_TRACE_SCOPE_CAT(cat, name)

Macros Documentation

define TENZR_TRACE_SCOPE

#define TENZR_TRACE_SCOPE(
    name
)
::tenzr::profiler::ScopedTraceTENZR_CONCAT(_tenzr_trace_, __LINE__)(name, "runtime")

define TENZR_TRACE_SCOPE_CAT

#define TENZR_TRACE_SCOPE_CAT(
    cat,
    name
)
::tenzr::profiler::ScopedTraceTENZR_CONCAT(_tenzr_trace_, __LINE__)(name, cat)

Source code

// Copyright 2026 The Tenzr Project Authors
//
// SPDX-License-Identifier: Apache-2.0
//
// 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
//
//      https://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 TENZR_COMMON_TRACE_H_
#define TENZR_COMMON_TRACE_H_

#include <atomic>
#include <chrono>
#include <concepts>
#include <cstdint>
#include <functional>
#include <memory>
#include <mutex>
#include <string>
#include <unordered_map>
#include <vector>

#include "tenzr/common/macros.h"

namespace tenzr::profiler {

  struct TraceEvent {
    std::string name;
    std::string cat;
    char ph{'X'};  // 'X'=Complete, 'i'=Instant, 'C'=Counter, 'M'=Metadata
    int64_t ts_us{0};
    int64_t dur_us{0};
    uint32_t pid{1};
    uint64_t tid{0};
    char s{'t'};  // Instant event scope ('g'=global, 'p'=process, 't'=thread)
    std::unordered_map<std::string, std::string> args;
  };

  struct TopOp {
    std::string name;
    std::string cat;
    int64_t total_dur_us{0};
    size_t count{0};
    double avg_dur_us{0.0};
  };

  struct BottleneckReport {
    std::string primary_bottleneck;
    double compute_to_comm_ratio{0.0};
    double memory_transfer_pct{0.0};
    double collective_comm_pct{0.0};
    int64_t total_device_kernel_us{0};
    int64_t total_memory_transfer_us{0};
    int64_t total_collective_comm_us{0};
    std::vector<std::string> recommendations;

    [[nodiscard]] std::string Format() const;
  };

  struct ProfileSummary {
    int64_t total_duration_us{0};
    size_t total_events{0};
    std::unordered_map<std::string, int64_t> category_durations_us;
    std::vector<TopOp> top_ops;

    [[nodiscard]] std::string FormatReport() const;
    [[nodiscard]] BottleneckReport IdentifyBottlenecks() const;
    [[nodiscard]] bool CheckLatencyRegression(int64_t baseline_us, double threshold_pct = 10.0, std::string* message = nullptr) const;
  };

  class Profiler {
   public:
    static Profiler& Instance();

    void Start();
    void Stop();
    void Reset();
    [[nodiscard]] bool IsEnabled() const { return enabled_.load(std::memory_order_relaxed); }

    void RecordEvent(TraceEvent event);
    void RecordInstant(std::string name, std::string cat = "general", char scope = 't');
    void RecordCounter(std::string name, int64_t value, std::string cat = "memory");
    void RecordMetadata(std::string name, std::string value);

    // Hardware-accelerated and stream profiling
    void RecordStreamKernel(std::string name, int64_t stream_id, int64_t dur_us, std::string cat = "gpu_kernel",
                            std::unordered_map<std::string, std::string> args = {});
    void RecordMemoryTransfer(std::string name, size_t bytes, bool h2d, int64_t dur_us);
    void RecordCollectiveComm(std::string name, size_t bytes, int64_t dur_us);

    // Exports Chrome Trace Event format (loadable by Perfetto, chrome://tracing, and OpenXLA xprof)
    bool ExportChromeTrace(const std::string& path) const;

    [[nodiscard]] ProfileSummary AnalyzeTrace() const;
    [[nodiscard]] size_t EventCount() const;

   private:
    Profiler();
    ~Profiler();

    std::atomic<bool> enabled_{false};
    std::chrono::high_resolution_clock::time_point start_time_;
    mutable std::mutex mutex_;
    std::vector<TraceEvent> events_;
    std::string default_output_path_;
  };

  inline bool IsEnabled() { return Profiler::Instance().IsEnabled(); }

  inline void Start() { Profiler::Instance().Start(); }

  inline void Stop() { Profiler::Instance().Stop(); }

  inline void Reset() { Profiler::Instance().Reset(); }

  inline void RecordCounter(std::string name, int64_t value, std::string cat = "memory") {
    Profiler::Instance().RecordCounter(std::move(name), value, std::move(cat));
  }

  inline void RecordInstant(std::string name, std::string cat = "general", char scope = 't') {
    Profiler::Instance().RecordInstant(std::move(name), std::move(cat), scope);
  }

  inline bool ExportChromeTrace(const std::string& path) { return Profiler::Instance().ExportChromeTrace(path); }

  inline ProfileSummary AnalyzeTrace() { return Profiler::Instance().AnalyzeTrace(); }

  inline void RecordStreamKernel(std::string name, int64_t stream_id, int64_t dur_us, std::string cat = "gpu_kernel",
                                 std::unordered_map<std::string, std::string> args = {}) {
    Profiler::Instance().RecordStreamKernel(std::move(name), stream_id, dur_us, std::move(cat), std::move(args));
  }

  inline void RecordMemoryTransfer(std::string name, size_t bytes, bool h2d, int64_t dur_us) {
    Profiler::Instance().RecordMemoryTransfer(std::move(name), bytes, h2d, dur_us);
  }

  inline void RecordCollectiveComm(std::string name, size_t bytes, int64_t dur_us) {
    Profiler::Instance().RecordCollectiveComm(std::move(name), bytes, dur_us);
  }

  class ScopedTrace {
   public:
    explicit ScopedTrace(std::string name, std::string cat = "op") : name_(std::move(name)), cat_(std::move(cat)), enabled_(IsEnabled()) {
      if (enabled_) {
        start_us_ = CurrentTimeUs();
      }
    }

    void AddArg(std::string key, std::string value) {
      if (enabled_) {
        args_[std::move(key)] = std::move(value);
      }
    }

    ~ScopedTrace() {
      if (enabled_) {
        int64_t end_us = CurrentTimeUs();
        TraceEvent ev;
        ev.name = std::move(name_);
        ev.cat = std::move(cat_);
        ev.ph = 'X';
        ev.ts_us = start_us_;
        ev.dur_us = end_us - start_us_;
        ev.pid = 1;
        ev.tid = GetTid();
        ev.args = std::move(args_);
        Profiler::Instance().RecordEvent(std::move(ev));
      }
    }

    static int64_t CurrentTimeUs();
    static uint64_t GetTid();

   private:
    std::string name_;
    std::string cat_;
    bool enabled_;
    int64_t start_us_{0};
    std::unordered_map<std::string, std::string> args_;
  };

  class TraceSession {
   public:
    explicit TraceSession(std::string path = "tenzr_profile.trace.json") : path_(std::move(path)) { Profiler::Instance().Start(); }

    ~TraceSession() { Stop(); }

    void Stop() {
      if (Profiler::Instance().IsEnabled()) {
        Profiler::Instance().Stop();
        Profiler::Instance().ExportChromeTrace(path_);
      }
    }

    TraceSession(const TraceSession&) = delete;
    TraceSession& operator=(const TraceSession&) = delete;
    TraceSession(TraceSession&&) = default;
    TraceSession& operator=(TraceSession&&) = default;

   private:
    std::string path_;
  };

  class StepMarker {
   public:
    explicit StepMarker(int64_t step, std::string name = "Step") : name_(std::move(name) + "_" + std::to_string(step)), scope_(name_, "step") {
      Profiler::Instance().RecordInstant(name_ + "_start", "step");
    }

    ~StepMarker() { Profiler::Instance().RecordInstant(name_ + "_end", "step"); }

    StepMarker(const StepMarker&) = delete;
    StepMarker& operator=(const StepMarker&) = delete;

   private:
    std::string name_;
    ScopedTrace scope_;
  };

  template <typename Fn, typename... Args> auto Trace(const std::string& path, Fn&& fn, Args&&... args) {
    TraceSession session(path);
    return std::invoke(std::forward<Fn>(fn), std::forward<Args>(args)...);
  }

}  // namespace tenzr::profiler

#define TENZR_TRACE_SCOPE(name) ::tenzr::profiler::ScopedTrace TENZR_CONCAT(_tenzr_trace_, __LINE__)(name, "runtime")

#define TENZR_TRACE_SCOPE_CAT(cat, name) ::tenzr::profiler::ScopedTrace TENZR_CONCAT(_tenzr_trace_, __LINE__)(name, cat)

#endif  // TENZR_COMMON_TRACE_H_