menu

An Analysis of the XLA Profiler Subsystem: Part 1

2025.07.16 :: #perf #profiling #xla #jax #tensorflow

1. Introduction: Purpose and Scope of the XLA Profiler§

The Accelerated Linear Algebra (XLA) compiler is a specialized infrastructure designed to optimize machine learning models for execution on high-performance hardware, including Graphics Processing Units (GPUs) and Tensor Processing Units (TPUs). Achieving maximum performance requires a detailed understanding of how time is spent during execution, the identification of performance bottlenecks, and a clear analysis of the compiled code's behavior on the target hardware. The XLA Profiler is the subsystem responsible for providing these essential insights. Its primary technical challenge is to capture performance data from multiple, disparate sources—such as the host Central Processing Unit (CPU), the Python interpreter, and the accelerator device—and to correlate this data accurately, especially within complex asynchronous execution environments.

The primary objectives of this subsystem are as follows:

  • Unified Data Presentation: To capture performance data from all constituent components and present it in a correlated, holistic manner.
  • Minimal Performance Perturbation: To ensure the process of profiling has a negligible impact on the performance characteristics of the program under analysis.
  • Extensibility and Platform Independence: Although the framework provides deep, hardware-specific insights, its core architecture is designed for extensibility to accommodate new platforms and devices. This modularity is exemplified by a factory design pattern, defined in tp/profiler/lib/profiler_factory.h, which facilitates the registration of different tracer implementations.
  • Remote Profiling Capabilities: To permit the profiling of applications deployed on remote servers or embedded systems, a common requirement in production and research contexts.

The profiler is not only a tool for post-mortem analysis; it also serves as a foundational component for other XLA systems. Of particular note is its integration with the Autotuner, which leverages the profiler to empirically evaluate and select the most performant kernel implementation for a given operation, as specified in autotuner/profiler.h.

2. Core Architecture: A Unified and Extensible System§

The architecture of the XLA Profiler can be understood as a layered system that ranges from low-level instrumentation to high-level data analysis and remote control. At its core is the ProfilerSession class, defined in tp/profiler/lib/profiler_session.h, which acts as the central controller for a profiling run. When a session is initiated, it activates a collection of registered "tracers." Each of these tracers implements the ProfilerInterface and is responsible for the collection of a distinct category of performance data.

2.1. Key Components§

  • ProfilerInterface (tp/profiler/lib/profiler_interface.h): This abstract base class is the contract for all data collectors (tracers). Any implementation must provide methods for Start, Stop, and CollectData. This standardized interface is fundamental to the profiler's extensibility, allowing for the integration of new data sources, such as new hardware accelerators, without requiring modifications to the core session management logic.
  • ProfilerSession (tp/profiler/lib/profiler_session.h): This class manages the lifecycle of a profiling session. Its principal responsibilities, as defined by its public methods, include the following: instantiating and registering ProfilerInterface implementations via a factory system (tp/profiler/lib/profiler_factory.cc); coordinating the starting and stopping of these tracers; collecting profiling data from all tracers into a unified XSpace data structure; and managing profiling configurations and options, with a structure defined by its use in profiler/utils/profiler_options_util.h.
  • XPlane and XSpace: The XSpace and XPlane constructs are the cornerstone of the profiler's data representation model. Although the Protocol Buffer definition is not located in this directory, its structure is clearly defined by its usage throughout the codebase (e.g., profiler/utils/xplane_builder.h). XSpace is the top-level container for all profiling data from a session, containing one or more XPlane instances. An XPlane typically represents all trace data from a single computational device (e.g., one for the host CPU, one for each GPU).
    • An XPlane is composed of XLines, which represent discrete timelines (e.g., a CPU thread or a GPU stream).
    • Each XLine contains a sequence of XEvents, which are the individual, timed events (e.g., a function invocation or a kernel execution).
    • Both XPlanes and XEvents can have XStats, which are key-value pairs containing metadata such as kernel specifications, memory allocation dimensions, or correlation identifiers. The schema for common statistics is formally defined in profiler/utils/xplane_schema.h, providing semantic meaning to the collected raw data.

This structured data format is highly flexible, enabling the rich representation of complex, multi-device execution traces. The XPlaneBuilder class provides a convenient API for the programmatic construction of these data structures.

// From profiler/utils/xplane_builder.h (conceptual example)
XPlaneBuilder host_plane_builder("Host CPU");
host_plane_builder.GetOrCreateLine(thread_id);
XEventBuilder event_builder = host_plane_builder.GetEventBuilder(line_builder);
event_builder.SetTimestampNs(start_time_ns);
event_builder.SetDurationNs(end_time_ns - start_time_ns);
event_builder.AddStat("My Stat", my_value);

3. Instrumentation, Collection, and Correlation§

The profiler is designed to collect data from multiple sources simultaneously and then establish relationships between the collected events.

3.1. Host Tracing§

Host-side tracing involves capturing activity on the CPU, which includes both user-level Python code and the C++ runtime environment.

  • TraceMe (tp/profiler/lib/traceme.h): The primary mechanism for instrumenting C++ code is TraceMe. It is a lightweight Resource Acquisition Is Initialization (RAII) object that records start and end timestamps. It is used extensively in performance-sensitive code sections.
    void MyFunction() {
  profiler::TraceMe trace("MyFunction");
  // ... work to be profiled ...
}
  • HostTracer (backends/profiler/cpu/host_tracer.h): The HostTracer is the implementation of the ProfilerInterface for host-side tracing. It is responsible for collecting trace data from all threads and aggregating metadata about the host environment, such as CPU specifications, which is gathered by the MetadataCollector (backends/profiler/cpu/metadata_collector.cc).
  • Python Tracing (backends/profiler/cpu/python_tracer.h): The Python tracing mechanism hooks into the Python interpreter to capture the execution of Python code. It uses the PyEval_SetProfile function from the Python C API, configured in py/profiler/internal/python_hooks.cc, to set a callback that is invoked on the entry and exit of Python functions.
The Producer-Consumer Model: TraceMeRecorder§

To minimize the performance overhead associated with TraceMe, the system uses a highly optimized producer-consumer model.

  • Producer: Instrumented code sections, acting as producers, generate Activity events.
  • Buffer: These events are written to a thread-local, lock-free ring buffer managed by the TraceMeRecorder class (profiler/backends/cpu/traceme_recorder.h). Using a thread-local buffer avoids costly inter-thread synchronization, which makes event production extremely efficient.
  • Consumer: The HostTracer acts as the consumer; at the end of a profiling session, it invokes the TraceMeRecorder::Consume method on each thread's recorder to process the buffered events and construct the final XEvent objects within the host's XPlane.

3.2. Device Tracing§

Device-side tracing is inherently platform-specific, requiring dedicated tracers for each distinct hardware architecture.

  • NVIDIA GPUs (backends/profiler/gpu/cupti_tracer.h): For NVIDIA GPUs, the implementation uses the CUDA Profiling Tools Interface (CUPTI) to collect a detailed log of all on-device activities, including kernel executions, memory transfers, and stream synchronizations.
  • AMD GPUs (backends/profiler/gpu/rocm_tracer.h): For AMD GPUs, it uses the ROCm profiler library (ROCProfiler) to collect device activity.
  • TPUs (backends/profiler/tpu/tpu_tracer.cc): The tracer for TPUs operates differently; it communicates with the TPU's intrinsic profiling runtime, requests a trace, receives a fully-formed XSpace protocol buffer, and merges it with the XSpace being collected on the host.

3.3. Event Correlation§

A raw trace is simply a collection of disconnected events. The profiler employs several distinct mechanisms to establish relational dependencies among these events, a process that is fundamental to understanding the end-to-end execution flow of a given computational operation.

3.3.1. CPU Parent-Child Relationships§

On a single CPU thread, parent-child relationships are captured in two main ways:

  1. Scoped Nesting: The RAII nature of TraceMe (tp/profiler/lib/traceme.h) naturally creates nested events. A TraceMe object for a child function created inside the scope of a parent's TraceMe will have its start and end times entirely within the parent's, which visualization tools display as a nested relationship.
  2. Explicit Annotation Stack: For more explicit call-stack-style annotation, the system uses an AnnotationStack (profiler/backends/cpu/annotation_stack.h). The ScopedAnnotation class (tp/profiler/lib/scoped_annotation.h) pushes and pops annotations on this stack to create a hierarchical name (e.g., Parent::Child) that is applied to all TraceMe events within that scope.
3.3.2. Asynchronous Event Correlation§

Linking an event on a host thread to an asynchronous operation that executes subsequently on a different thread or on a separate device is more complex. This is handled by two primary mechanisms, each designed for a specific correlation scenario.

  • Mechanism 1: ConnectedTraceMe (Host-to-Host Correlation) This general-purpose software pattern (tp/profiler/lib/connected_traceme.h) is used to correlate two related events that occur on different host threads. It uses a producer-consumer model with a shared context_id to link an event that dispatches work to a thread pool with the corresponding event on the worker thread that performs the execution.
  • Mechanism 2: Driver-Level Correlation (Host-to-Device Correlation) This is a more specialized mechanism for directly linking a host-side API call to a specific device execution. It uses a thread-local tag set on the host thread just before a driver call. The driver associates this tag with the launched work, and the device activity record contains the same tag, creating a direct link. This process is managed by RAII objects, such as CuptiTracer::Annotation (backends/profiler/gpu/cupti_tracer.h).

4. Illustrative Example: Conceptual Trace of a JAX Matrix Multiplication§

To illustrate how the previously described components interact, this section presents a conceptual trace of a JAX matrix multiplication operation. The following walkthrough is designed to synthesize the roles of the different profiler components by following the logic and code paths established in the source code, rather than representing an empirical log from a specific execution.

The Scenario: Consider a scenario where the following JAX code is executed while a profiler session is active.

import jax
import jax.numpy as jnp

@jax.jit
def matmul_op(x, y):
  return jnp.dot(x, y)

# Profiler is active
x = jnp.ones((2048, 2048))
y = jnp.ones((2048, 2048))
result = matmul_op(x, y).block_until_ready()

The Profiling Lifecycle:

  1. Python Frontend (PythonTracer): When activated through hooks in py/profiler/internal/python_hooks.cc, the PythonTracer is designed to observe the call to matmul_op and create a high-level XEvent within the host CPU's XPlane. This event, named something like JAX-Execute: matmul_op, would become the top-level parent in the trace hierarchy.
  2. C++ Host Runtime (TraceMe): The call would then descend through the PjRt C++ API into the XLA runtime, triggering a series of nested TraceMe scopes. XEvents for PjRtCpuClient::Compile, PjRtStreamExecutorClient::Execute, and ultimately GpuExecutable::Execute would be recorded on the same host thread's XLine, nested within the originating Python event.
  3. Thunk Execution and Correlation (GemmThunk): The GpuExecutable contains a sequence of "thunks" to be executed on the GPU stream. For the matrix multiplication operation, the GemmThunk::ExecuteOnStream method (backends/gpu/runtime/gemm_thunk.cc) would be invoked. This is a critical point for establishing host-to-device correlation.
    • Inside this function, a CuptiTracer::Annotation object would be created on the stack.
    • Its constructor is designed to call cuptiActivityPushCorrelationId, tagging the current host thread with a unique ID (e.g., 12345).
    • A TraceMe event for gemm::DoGemm is created, which the HostTracer then records on the host XPlane with an attached correlation_id stat.
    • The thunk then calls the underlying cuBLAS library to launch the GEMM kernel on the GPU stream. This call is asynchronous and returns immediately.
  4. Device Execution (CuptiTracer): At some later time, the GPU would execute the cuBLAS kernel. Because the launching thread was tagged with ID 12345, the CUPTI record received by the CuptiTracer would also contain the same correlation ID. The CuptiTracer then creates an XEvent on the GPU's XPlane (e.g., on the CUDA Stream #1 XLine) and attaches a matching correlation_id stat.
  5. Data Finalization and Post-Processing: When the session ends, the ProfilerSession finalizes the data.
    • Collection: It calls the CollectData method on each active tracer (e.g., HostTracer, CuptiTracer). Each tracer receives a pointer to a single, shared XSpace object, which it populates by adding its own XPlane(s).
    • Post-Processing: The populated XSpace, containing the raw trace, is then post-processed by functions like PostProcessSingleHostXPlane (profiler/convert/post_process_single_host_xplane.h). A key utility used here is GroupEvents (profiler/utils/group_events.h), which analyzes correlation IDs and other metadata to build a meaningful hierarchy that links parent events to their children across disparate threads and devices.
  6. Final Visualization: When such a processed trace is loaded into a visualization tool like TensorBoard, it provides a clear, hierarchical view. Expanding the JAX-Execute: matmul_op event reveals the C++ runtime calls. Critically, a directed edge would connect the gemm::DoGemm event on the host CPU timeline to the cuBLAS kernel execution event on the GPU timeline, providing an unambiguous illustration of the causal relationship of the asynchronous launch.

5. Python Bindings and User-Facing API§

The C++ backend is made available to the Python environment using the pybind11 library. The primary bindings are defined in py/profiler.cc. This file creates the _profiler Python module and exposes essential functionalities to the Python interpreter.

  • Session Control: The module exposes start and stop functions that internally manage the lifecycle of a ProfilerSession. The start function takes a server address and profiling options as arguments, while the stop function finalizes the session and retrieves the collected XSpace data.
  • TraceMe for Python: To allow for custom events to be inserted from Python code, a TraceMe context manager is provided. The C++ implementation for this wrapper is located in py/profiler/internal/traceme_wrapper.h. This allows for the following usage pattern:
    from xla.python.profiler import TraceMe
with TraceMe("My custom Python block"):
  # ... code to be profiled ...

6. Data Visualization§

The final stage of the profiling pipeline is the conversion of the post-processed XSpace data structure into a visualizable format. The most important conversion is to the Trace Event Format, a JSON format understood by visualization tools like chrome://tracing and TensorBoard. This conversion is handled by the ConvertXSpaceToTraceEvents (profiler/convert/xplane_to_trace_events.cc) and TraceToJson (profiler/convert/trace_events_to_json.cc) functions.

7. Integration with the Autotuner§

The Autotuner's purpose is to find the most performant algorithm for a given computational operation. The profiler is the essential measurement tool that makes this possible, as specified in autotuner/profiler.h. The GPU-specific implementation, located at backends/gpu/autotuner/gpu_profiler.cc, uses device-specific timing events (e.g., CUDA events) to get high-precision measurements of kernel execution durations. These measurements are then used to populate a performance database.

8. Remote Profiling§

The profiler can be controlled remotely through a gRPC-based client-server architecture.

  • ProfilerService (profiler/rpc/profiler_service_impl.h): A gRPC service runs on the process being profiled; its Profile method starts the entire data collection, aggregation, and post-processing pipeline.
  • ProfilerClient (profiler/rpc/client/profiler_client.h): A client application, such as TensorBoard or a command-line utility, connects to the service to start and stop profiling and to retrieve the resulting data.
  • SessionManager (profiler/utils/session_manager.h): The SessionManager is a server-side component that manages concurrency by ensuring only one profiling session is active at a time, which prevents resource conflicts.

9. Conclusion§

The XLA Profiler is a sophisticated, multi-layered system that is essential for performance analysis and optimization. Its key strengths are its modular architecture, based on the ProfilerInterface; its unified data representation using the XPlane structure; and its robust remote capabilities. By combining host-side TraceMe instrumentation, detailed device-specific tracing, and several complementary correlation mechanisms, it provides a comprehensive picture of a model's execution. The final stages of data aggregation and event grouping transform this raw data into an understandable and actionable trace. This allows both developers and automated systems like the Autotuner to improve performance on modern hardware accelerators.