Profiler Integration

Created On: Dec 26, 2025 | Last Updated On: Jul 07, 2026

Background

PyTorch ships a device-agnostic profiler that instruments CPU-side operator dispatch, coordinates with accelerator collectors, captures Python stacks, and exports aggregated statistics or Chrome/Perfetto traces. For core architecture, see torch/csrc/profiler/README.md.

There are two primary integration paths for accelerators:

1. Legacy autograd profiler:

   • Can attach backend-specific hooks via ProfilerStubs to record device events and compute elapsed times.

   • Works without Kineto; suitable for PrivateUse1 backends that want a minimal, self-contained path.

2. Kineto IActivityProfiler plugin:

   • Registers a full activity profiler with Kineto via REGISTER_PRIVATEUSE1_PROFILER.

   • Wires Kineto sessions and correlation-ID plumbing; vendors extend this to emit kernel events, flow links, and Chrome/Perfetto trace compatibility.

   • Requires Kineto at backend build time (kineto_LIBRARY from find_package(Torch), guarded by USE_KINETO).

Path	Python API	Profiler State	What it provides
Legacy (1)	torch.autograd.profiler.profile(use_device="openreg") (default use_kineto=False)	KINETO_PRIVATEUSE1_FALLBACK	Operator-level timing via ProfilerStubs device events
Kineto plugin (2)	profile(activities=[ProfilerActivity.CPU, ProfilerActivity.PrivateUse1])	KINETO_PRIVATEUSE1	Kineto session + correlation-ID plumbing; vendors add kernel events and flow links

Both paths can coexist when the backend extension is built with Kineto available (kineto_LIBRARY from find_package(Torch)). The legacy stubs path always works; the Kineto plugin path requires USE_KINETO at backend build time. PyTorch core already exposes REGISTER_PRIVATEUSE1_PROFILER; vendors implement and register their IActivityProfiler in the backend extension.

Design

Architecture overview

Layer	Responsibility	Source
Python control plane	Owns profiler lifecycle (prepare → start → stop → step) and exposes user APIs such as torch.autograd.profiler.profile.	torch/autograd/profiler.py
Profiler stubs	Implements torch::profiler::impl::ProfilerStubs so the profiler can record device events, synchronize, iterate devices, and compute elapsed time.	torch/csrc/profiler/stubs/
Device runtime	Provides streams, events, and device guards used by the stubs; implementation is backend-specific.	Backend extension (vendor code)

This layering keeps PyTorch device-agnostic: Python brokers the session, ProfilerStubs translate profiler requests into backend runtime calls, and the runtime interacts with the accelerator.

Key contracts

• Record hooks: record() must capture (optional) device index, allocate a backend event, optionally stash a CPU timestamp, and enqueue the event on the active stream.

• Elapsed time: elapsed() is responsible for synchronizing individual events and returning durations in microseconds.

• Synchronization hooks: synchronize() and onEachDevice() guarantee phase transitions (e.g., warmup → active) are aligned across devices.

• Annotations: mark, rangePush, and rangePop can be implemented to enrich traces; otherwise they may be left as no-ops.

Implementation (Legacy way)

Here we use OpenReg (Open Registration) to illustrate the minimal set of hooks a PrivateUse1 accelerator needs to expose so the profiler can attribute ATen ops, record_function ranges, and user code to device activity. OpenReg keeps upstream code untouched by translating profiler requests into its runtime calls, mirroring what a production accelerator would implement inside an out-of-tree extension.

OpenReg supports both paths: the legacy autograd profiler (use_kineto=False, the default) for operator-level timing via stubs, and the modern torch.profiler.profile API (use_kineto=True) for the Kineto plugin path described below.

Profiler stubs (C++)

torch::profiler::impl::OpenRegMethods inherits from ProfilerStubs and wires the hooks described above:

Method	Purpose
record	Grabs the current OpenRegStream, creates an orEvent, captures an optional CPU timestamp via c10::getTime(), and records the event on the stream.
elapsed	Synchronizes both events, calls orEventElapsedTime, and converts milliseconds to microseconds for the profiler.
onEachDevice	Uses c10::DeviceGuard(DeviceType::PrivateUse1) to iterate over torch.openreg.device_count() so schedulers can run per-device setup or teardown.
synchronize	Calls orDeviceSynchronize() to align device work with CPU scheduling phases.
enabled and annotation shims	Report availability and provide placeholder implementations for mark/push/pop.

The constructor registers the methods once via registerPrivateUse1Methods(&methods);, making them discoverable whenever the profiler is enabled with use_device="openreg".

Python control plane

On the Python side, no new entrypoint is required—developers use the standard autograd profiler:

from torch.autograd.profiler import profile as autograd_profile
from torch.profiler import record_function

with autograd_profile(use_device="openreg", record_shapes=True) as prof:
    with record_function("matmul"):
        x = torch.randn(512, 512, device="openreg")
        y = torch.randn(512, 512, device="openreg")
        z = x @ y

print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=10))
prof.export_chrome_trace("openreg_trace.json")

Data capture flow

1. User code enters autograd_profile(use_device="openreg").

2. The profiler transitions to ProfilerState.KINETO_PRIVATEUSE1_FALLBACK.

3. The profiler asks the active backend to record() an event.

4. The OpenReg stubs allocate orEvent objects, attach them to the current stream, and stash CPU timestamps.

5. When events end, the profiler calls elapsed() to compute durations.

Implementation (Kineto Plugin)

Note

This section covers the Kineto IActivityProfiler plugin path for kernel-level tracing. It requires USE_KINETO at build time. All Kineto-dependent code must be guarded with #ifdef USE_KINETO.

The plugin path has two layers: a device library component (the CUPTI analog) and the PyTorch integration layer. OpenReg keeps these clearly separated.

Device library: correlation tracking

The device library provides openreg::profiler::OpenRegTracer (third_party/openreg/csrc/tracer.h/.cpp) — a singleton with a thread-local correlation-ID stack and an atomic enable/disable flag that the profiler session uses to control the recording window.

Kineto pushes/pops correlation IDs through the session. The session calls C-style activity APIs in openreg.h (mirroring CUPTI):

• orActivityEnableTracing() / orActivityDisableTracing() — control the recording window

• orActivityPushExternalCorrelationId() / orActivityPopExternalCorrelationId() — maintain the correlation stack

A real vendor’s equivalent would be their device tracing SDK (e.g., CUPTI for CUDA).

PyTorch integration: IActivityProfiler and IActivityProfilerSession

Implement the two Kineto interfaces from third_party/kineto/libkineto/include/IActivityProfiler.h. In OpenReg, these live in torch_openreg/csrc/profiler/ — the backend extension integration layer.

• IActivityProfiler — stateless factory. Two configure() overloads both create and return a session:

  • configure(activity_types, config) — synchronous overload; required by the interface. The OpenReg stub implements this as the core session-creation path.

  • configure(ts_ms, duration_ms, activity_types, config) — Kineto’s child-profiler path calls this overload for all traces (including on-demand), passing profileStartTime() epoch ms and profileDuration() ms. The OpenReg stub ignores scheduling and delegates to the first overload; vendors use ts_ms/duration_ms to defer device-SDK activation.

• IActivityProfilerSession — per-trace session. start()/stop() manage the profiling window and toggle activity tracing via orActivityEnableTracing()/orActivityDisableTracing(); getTraceBuffer() returns the buffer to Kineto.

  • Reference stub: processTrace() only sets the trace span (traceBuffer_.span = TraceSpan(startTs_, endTs_, "openreg")); it emits no kernel records.

  • Vendor extension: replace processTrace() to flush records from your device tracing SDK, emit GenericTraceActivity entries with timestamps (µs), correlation IDs, and flow links (flow.id = correlationId, flow.type = kLinkAsyncCpuGpu, flow.start = 0).

Registration and build

Register with one line: REGISTER_PRIVATEUSE1_PROFILER(OpenRegActivityProfiler). The macro (defined in torch/csrc/profiler/standalone/privateuse1_profiler.h) creates a static registration object that forwards a factory to Kineto at profiler init time.

Kineto is on by default in PyTorch; no special build flags are needed unless explicitly disabled it with USE_KINETO=0. In backend extension, find_package(Torch) sets kineto_LIBRARY; link against kineto and torch_cpu_library, and guard Kineto code with #ifdef USE_KINETO. Without Kineto, the plugin compiles as a no-op and only the legacy stubs path is available.

Usage

import torch
from torch.profiler import profile, ProfilerActivity

with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.PrivateUse1]) as prof:
    x = torch.randn(512, 512, device="openreg")
    y = torch.randn(512, 512, device="openreg")
    z = x @ y

prof.export_chrome_trace("kernel_trace.json")