Profiler

Performance instrumentation and the on-device + Python profiler. Lives in tensorflow/core/profiler/ (C++) and tensorflow/python/profiler/ (Python).

Purpose

• Capture trace events (op start/stop, GPU kernels, memory allocations) during a TF program.
• Expose them through the tf.profiler.experimental.* Python API and a profiler RPC for remote capture.
• Render them in TensorBoard's "Profile" plugin or as Chrome trace JSON.

Directory layout

tensorflow/core/profiler/
├── lib/
│   ├── traceme.h                  # TraceMe macro
│   ├── traceme_encode.h
│   ├── profiler_session.{h,cc}    # Session-scoped capture
│   ├── profiler_factory.{h,cc}
│   ├── profiler_lock.{h,cc}
│   └── profiler_collection.{h,cc}
├── backends/                      # CPU, GPU (CUPTI/ROCm), TPU collection
├── convert/                       # Convert raw traces → various JSON dialects
├── protobuf/                      # XSpace, XPlane, XStat protos (the canonical trace format)
├── rpc/                           # Profiler gRPC service (remote capture)
├── utils/
└── ...

tensorflow/python/profiler/
├── profiler_v2.py                  # tf.profiler.experimental.start/stop
├── profiler_client.py              # Remote profiling client
├── trace.py                        # tf.profiler.experimental.Trace context manager
└── ...

Key abstractions

Type / macro	File	Purpose
TraceMe	tensorflow/core/profiler/lib/traceme.h	RAII macro that emits a (start, stop, attrs) trace event.
XPlane / XSpace	tensorflow/core/profiler/protobuf/xplane.proto	The canonical column-oriented trace format.
ProfilerSession	tensorflow/core/profiler/lib/profiler_session.h	Owns a single capture; dumps an XSpace.
ProfilerFactory	tensorflow/core/profiler/lib/profiler_factory.h	Plugs in CPU/GPU/TPU back-ends.
tf.profiler.experimental.start/stop	tensorflow/python/profiler/profiler_v2.py	Python entrypoint.
Profiler gRPC service	tensorflow/core/profiler/rpc/	Allows profiler_client to capture on remote workers.

How a capture works

sequenceDiagram
    participant User as Python user
    participant Session as ProfilerSession
    participant Backends as CPU/GPU/TPU back-ends
    participant File as TensorBoard logdir

    User->>Session: tf.profiler.experimental.start(logdir)
    Session->>Backends: Start tracing
    Note over Backends: TraceMe events,<br/>CUPTI GPU events,<br/>etc. accumulate
    User->>Session: tf.profiler.experimental.stop()
    Session->>Backends: Stop tracing
    Backends-->>Session: XSpace proto
    Session->>File: write profile.tf.gz, overview, ...
    User->>User: open in TensorBoard

TraceMe is sprinkled through hot paths (executor, kernels, eager runtime). Each TraceMe emits a start/stop event with optional metadata strings (e.g. shapes). When a ProfilerSession is active, the events are recorded into the shared TraceMeRecorder; when no session is active, the macro short-circuits to nothing.

Back-ends

• Host CPU — TraceMeRecorder collects events from the running threads.
• GPU — uses CUPTI (NVIDIA) / ROC-profiler to capture kernel launches and memcpy events.
• TPU — TPU runtime emits XPlanes directly.

The profiler driver sequences these and serialises them into the unified XSpace/XPlane format consumed by TensorBoard.

Python API surface

import tensorflow as tf

tf.profiler.experimental.start("logdir")
# ... model code ...
tf.profiler.experimental.stop()

# Or:
with tf.profiler.experimental.Profile("logdir"):
    train_one_epoch()

# Per-step trace inside a tf.function:
with tf.profiler.experimental.Trace("train", step_num=step):
    train_step(batch)

The remote-capture client tf.profiler.experimental.client.trace("worker:6009", "logdir", duration_ms=2000) uses the gRPC service in tensorflow/core/profiler/rpc/.

Integration points

• Used by: any subsystem that wraps a critical region in tensorflow::profiler::TraceMe.
• Surfaces in: TensorBoard's "Profile" plugin (separate repo tensorflow/profiler).
• DTensor / tf.distribute rely on the gRPC profiler service to capture across hosts.

Entry points for modification

• Add a TraceMe to a hot path — single-line addition; usually obvious from top -H or a slow trace.
• New profiler back-end — implement tensorflow::profiler::ProfilerInterface and register with ProfilerFactory.
• New trace event payload — extend XStat/XPlane schemas in tensorflow/core/profiler/protobuf/ (mind backward compatibility with TensorBoard).

Related

• core-runtime — much of the existing instrumentation lives in the executor and op kernels.
• features/distribution-strategy — multi-host profiling needs the gRPC service.