Pytorch backend - #540
Merged
Merged
Pytorch backend#540
Conversation
Copies the transitive header closure of ATen/record_function.h (605 files) from Aurora's `module load frameworks` PyTorch install (2.10.0a0), following the same vendoring pattern hip/ze/mpi use for their real API headers. Verified standalone compile against these headers with the POC tracer. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
pytorch_events.yaml declares op_entry/op_exit (cxi-style hand-written events, no API model needed since we only trace two RecordFunction callbacks). tracer_pytorch.cpp adapts the POC tracer to fire these as LTTng-UST tracepoints instead of printing to stdout, registering via at::addGlobalCallback at library load time so it works under LD_PRELOAD with no application changes. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Generates the lttng_ust_pytorch tracepoint provider from pytorch_events.yaml (cxi-style, via gen_custom_probes.rb), builds it into libpytorchtracepoints.la, and links libTracerPytorch.la against it plus two throwaway dummy libtorch_cpu.so/libc10.so (built with `-x c++ /dev/null`, matching the POC's build.sh) so the linker can resolve -ltorch_cpu/-lc10 without a real torch install at THAPI build time. The dummy libs are removed right after each link (all-local, install-exec-hook) and again by clean-local, using explicit file removal + rmdir rather than a recursive delete. Validated against the real THAPI Autotools build (autogen.sh, configure, make, make install, make clean) using the thapi-tracer-environment-config/-build module set on Aurora, and confirmed end-to-end: the installed libTracerPytorch.so, LD_PRELOADed onto the real `frameworks` module's PyTorch, records a genuine LTTng session that babeltrace2 reads back as well-formed op_entry/op_exit events. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Adds backends/pytorch to SUBDIRS and its Makefile/tracer script to AC_CONFIG_FILES, matching every other backend. Also adds tracer_pytorch.sh.in, the standalone manual-usage wrapper (omp/itt style): it creates its own LTTng session, enables lttng_ust_pytorch, locates the active python3's torch lib directory to satisfy libTracerPytorch.so's runtime dependency on libtorch_cpu.so/libc10.so, and LD_PRELOADs it. Validated with a full `autogen.sh && configure && make -j install` of the whole tree (all existing backends plus pytorch) on Aurora, then ran tracer_pytorch.sh directly against the `frameworks` module's real PyTorch: the model runs forward+backward and babeltrace2 reads back 1708 well-formed op_entry/op_exit events from the recorded trace. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Adds enable_events_pytorch (enables lttng_ust_pytorch:*, mirroring itt/hip) and an all_env_tracers block that, when `pytorch` is requested via --backends, probes `python3 -c "import torch..."` for the active interpreter's torch lib directory and sets LD_LIBRARY_PATH/LD_PRELOAD accordingly so libTracerPytorch.so resolves its libtorch_cpu.so/libc10.so dependency at runtime. When torch isn't importable, it warns and leaves the backend disabled rather than failing iprof. This completes the goal for this stage: `iprof --backends pytorch --no-analysis -- python model.py` now records a trace, and babeltrace2 reads it back as well-formed op_entry/op_exit events (verified: 1708 events from the POC model's forward+backward pass, on Aurora's real `frameworks` PyTorch). The default --backends path (pytorch not requested) is unaffected. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
tracer_pytorch.cpp only touches five names from <ATen/record_function.h>:
RecordScope, ObserverContext, RecordFunction::name(),
RecordFunctionCallback, and addGlobalCallback. Everything else in the
previously-vendored 605-file/4.3MB closure (ATen/core, ATen/ops, c10/*,
torch/headeronly) was pulled in transitively only because the real
record_function.h needs it to typecheck, not because our tracer calls
into it. Replaces backends/pytorch/include/ with a single ~90-line
ATen/record_function.h declaring only those five names.
Advantages:
- 605 files / 4.3MB -> 1 file / ~2.5KB: nothing to keep in sync with
an upstream torch checkout, no per-release re-extraction step.
- Self-documents the tracer's actual API surface: the whole contract
with libtorch is visible in one small file instead of implied by
a large, mostly-unused header tree.
- Verified equivalent at compile, link, and runtime: builds clean
under -Wall -Wextra -Werror, links against the same build-time
dummy libtorch_cpu.so/libc10.so stubs, and produces byte-identical
output when LD_PRELOADed onto the real `frameworks` module's
PyTorch -- babeltrace2 reads back the same 1708 op_entry/op_exit
events from the POC model's forward+backward pass as the
real-header build did.
Risks, and why they're judged acceptable here:
- RecordFunction is declared as a bare method signature with no
fields: safe, because we only ever receive `const RecordFunction&`
from the real library and call .name() against its real
out-of-line symbol -- we never construct or lay out this type
ourselves.
- RecordFunctionCallback is NOT safe by the same argument: we
construct it ourselves and pass it BY VALUE into addGlobalCallback,
so its four private fields (start_, end_, sampling_prob_, the
scopes_ bool array sized by RecordScope::NUM_SCOPES, three trailing
bools) must match upstream's layout exactly, field for field. They
were copied verbatim from the real header for this PyTorch version.
- If a future PyTorch release reorders, adds, or removes a field on
RecordFunctionCallback, or changes RecordScope's member count, this
header will still compile without any warning or error -- the
mismatch would silently corrupt which scopes are enabled or which
function pointers are read, not fail the build. This is a real
regression in safety compared to vendoring real headers, which
can't drift from whatever torch they were extracted from.
- Mitigation: this layout must be re-verified against the real
ATen/record_function.h (diff the RecordFunctionCallback class and
RecordScope enum) whenever the target PyTorch version changes, and
the change re-validated end-to-end (this commit's babeltrace2
event-count check) rather than trusted on compile success alone.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
at::addGlobalCallback and at::RecordFunction::name, the only two ATen symbols tracer_pytorch.cpp calls, are both defined in libtorch_cpu.so. Verified with nm against the real installed libs and confirmed the resulting libTracerPytorch.so still LD_PRELOADs correctly and produces an identical RecordFunction trace against the POC workload. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Shorten the dummy-lib comment to just state which symbols come from libtorch_cpu.so. Also drop the all-local target: no other backend Makefile.am uses one, and clean-local/install-exec-hook already remove the dummy libtorch_cpu.so. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
clang++ -Werror flags two things gcc doesn't: - -Wmissing-field-initializers on LTTng's own tracepoint.h aggregate inits (not our code) - -Wunused-private-field on RecordFunctionCallback's fields, which exist only to match PyTorch's real class layout for addGlobalCallback and are read by libtorch, never by us Reproduced both errors locally against the CI's clang toolchain and confirmed these two -Wno- flags clear them without touching gcc builds. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
| # include/ATen/record_function.h) -- the real libtorch reads them, we | ||
| # never do, by design. | ||
| libTracerPytorch_la_CXXFLAGS = -std=c++17 -Wall -Wextra -Wno-unused-parameter -Wno-missing-field-initializers -Wno-unused-private-field $(WERROR) $(LTTNG_UST_CFLAGS) | ||
| libTracerPytorch_la_LDFLAGS = $(LTTNG_UST_LIBS) -Ldummy_libs -ltorch_cpu -avoid-version -module |
Collaborator
There was a problem hiding this comment.
why libtorch_cpu and not just libtorch? libtorch
We should verify that it work with the GPU backend of pytorch too if not done already
Collaborator
Author
There was a problem hiding this comment.
LibTracerPytorch.so calls at::addGlobalCallback() and RecordFunction::name(), both are implemeted in libtorch_cpu.so
TApplencourt
approved these changes
Sep 14, 2026
Collaborator
|
Thanks. So next step |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds a pytorch backend to THAPI that traces PyTorch RecordFunction (forward + backward) ops. Enable it explicitly:
Pending follow-up work (not in this PR):
--backends pytorch, and requires--no-analysissince there's no analysis-side decoder for this backend yet. Making it a default backend depends on adding that.iprof -- python3 model.py(active torch env) works;iprof -- venv/bin/python3 module.py(non-active env) doesn't yet — needs env vars to let the target interpreter be specified explicitly.Build: suppress two clang-only warnings in the pytorch tracer
tracer_pytorch.cppis the only backend translation unit that calls LTTng'stracepoint()macro from C++ instead of C (every other backend does so from a.cfile). That exposes two warnings that only clang raises, both false positives.Before:
After:
-Wno-missing-field-initializers— LTTng'stracepoint.hinitializes some of its internal structs with a designated initializer that only names one field (e.g..struct_size = ...), leaving the rest to their default zero value — a standard, safe C idiom. C accepts this silently; C++ (under clang) does not, and treats it as a missing-initializer warning. Every other backend includes this same header from a.cfile compiled as C, so they never trigger it;tracer_pytorch.cppis the only backend source that pulls it into a C++ translation unit.-Wno-unused-private-field—RecordFunctionCallbackin our hand-writtenrecord_function.hdeclares private fields (start_,end_,sampling_prob_, etc.) that our own code never reads. That's intentional: this struct is passed by value into PyTorch's realaddGlobalCallback, so its memory layout must match PyTorch's real class byte-for-byte — the fields are read by libtorch on the other side of that call, not by us. Clang can't see that and flags them as dead; gcc doesn't have this check at all.