Profiler

Active contributors: scotts, ryanzhang22

Purpose

torch.profiler is the modern profiler for PyTorch. It records CPU op times, CUDA kernel times (via Kineto + CUPTI), memory allocations, and Python stack traces, and exports them to Chrome trace format and TensorBoard. The legacy torch.autograd.profiler API is still present but is mostly a thin wrapper that delegates to the modern profiler.

Directory layout

Path	Contents
torch/profiler/	Public Python API
torch/profiler/profiler.py	profile, schedule, tensorboard_trace_handler
torch/profiler/_memory_profiler.py	Memory profiler
torch/profiler/_pattern_matcher.py	Performance anti-pattern detection
torch/profiler/python_tracer.py	Python frame recorder
torch/autograd/profiler.py	Legacy autograd profiler
torch/csrc/profiler/	C++ implementation
torch/csrc/profiler/collection.cpp/.h	Recording machinery
torch/csrc/profiler/orchestration/	Lifecycle: start/step/stop
torch/csrc/profiler/perf.cpp	perf_event hardware counters
torch/csrc/autograd/profiler*	Autograd-side hooks
third_party/kineto/	Kineto submodule (the GPU-trace library)

Key abstractions

Type	File	Purpose
torch.profiler.profile	torch/profiler/profiler.py	Context manager / ProfilerAction state machine
ProfilerActivity	torch/profiler/profiler.py	Enum: CPU, CUDA, XPU, MTIA
KinetoEvent	torch/csrc/profiler/collection.h	One recorded event
_MemoryProfile	torch/profiler/_memory_profiler.py	Allocation timeline
record_function	torch/autograd/profiler.py	Add a custom span

How it works

graph LR
    User[User code] -->|with profile()| Hook[Profiler hooks]
    Hook -->|RecordFunction callbacks| AT[ATen ops]
    Hook -->|CUPTI| CUDA[CUDA kernels]
    Hook -->|frame eval| Py[Python frames]
    Hook --> Coll[collection.cpp<br/>per-thread ring buffer]
    Coll -->|on stop| Tree[Build event tree]
    Tree -->|format| Chrome[Chrome trace JSON]
    Tree -->|format| TB[TensorBoard plugin]

Recording

with torch.profiler.profile(activities=[...], schedule=..., on_trace_ready=...) as prof: installs three kinds of hooks:

1. RecordFunction callbacks — every ATen op call enters a RecordFunction scope (defined in aten/src/ATen/record_function.h). When the profiler is active it pushes a span around the call.
2. CUPTI / Kineto — on CUDA, Kineto subscribes to CUPTI events to record kernel launches, memcpy, and HW counters.
3. Python tracing — python_tracer.py uses sys.setprofile to record Python call/return events.

Each thread keeps a ring buffer of KinetoEvents; on stop the buffers are merged, ordered by timestamp, and a parent/child tree is built using the call stack.

Schedule

The schedule argument lets users skip warmup, record a short window, repeat, etc.:

torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=2)

The state machine that drives this lives in torch/csrc/profiler/orchestration/.

Memory profiler

profile(profile_memory=True) records every allocation/free with a Python stack at allocation time. The post-processed timeline (_MemoryProfile) is used by torch.profiler.profile.export_memory_timeline().

There's a separate, lower-level memory snapshot API in torch.cuda.memory._record_memory_history() that records into a binary file usable by https://docs.pytorch.org/memory_viz.

Hardware counters

torch/csrc/profiler/perf.cpp exposes Linux perf_event_open for HW counters (cycles, instructions, cache misses) on CPU.

Pattern matcher

torch/profiler/_pattern_matcher.py runs after profiling completes and flags common anti-patterns: extra to() calls, broadcasted small tensors, autograd graphs not being detached, etc.

Integration points

• Kineto (third_party/kineto) is a hard runtime dep for GPU profiling.
• CUPTI is dynamically loaded for CUDA traces.
• TensorBoard plugin consumes the exported JSON; it lives in a separate repo (pytorch/kineto/tb_plugin).
• Autograd profiler is wired into the autograd engine for backward-time attribution.

Entry points for modification

• New event source → add a recorder in torch/csrc/profiler/. See python_tracer.py for a complete small example.
• New trace format → see the export functions in torch/profiler/profiler.py; the Chrome trace exporter is in torch/csrc/profiler/orchestration/python_tracer.cpp.
• New pattern → add to _pattern_matcher.py.

Key source files

File	Purpose
torch/profiler/profiler.py	Public API, schedule, state machine
torch/profiler/_memory_profiler.py	Memory profiler
torch/csrc/profiler/collection.cpp	Event collection
torch/csrc/profiler/orchestration/	Lifecycle
torch/csrc/profiler/perf.cpp	perf_event counters
torch/csrc/autograd/profiler_kineto.cpp	Kineto integration
aten/src/ATen/record_function.h	RecordFunction callback API