PyTorch is a tensor library that became the dominant ML training
framework through one architectural bet: define-by-run
autodiff. There is no graph-building phase separate from
computation โ every eager a * b call is graph construction. The
call flows from Python through a generated C++ binding into a
central dispatcher, which routes to an autograd kernel whose job is
twofold: run the real kernel for the backend, and leave a receipt
(a Node) on the output tensor describing how to invert the
operation. loss.backward() is then a tiny DAG scheduler โ ready
queues, in-degree countdowns, one worker thread per device โ
walking those receipts in reverse topological order. The lesson
this codebase teaches better than any other: how to layer a thin
Python frontend over a C++ core through one dispatch table, how
to generate that entire surface (bindings, kernels, autograd
nodes, signatures) from two YAML schemas, and what it costs to
make every single operation carry the weight of an implicit,
user-invisible tracing system.
๐ง Explain like I'm 10
A restaurant where every cook writes a receipt for what they
made, and at closing time a cleanup crew walks the receipts
backwards to pay everyone back. When you order (call torch.mul
in Python), the receptionist (the dispatcher, in C++) looks at
where your ingredients came from (each tensor carries a tag โ the
dispatch key set โ saying CPU or GPU, needs-grad or not) and routes
the order to the right kitchen (the kernel). As the cook works,
the receptionist also staples a receipt (the autograd Node) onto
the finished dish, and each receipt points back to the receipts of
the dishes it was made from (next_edges). At closing time
(loss.backward()), the cleanup crew (the autograd engine) has
one clerk per room (one worker thread per device), each with a
pile of receipts to process (the ReadyQueue); a clerk may only
process a receipt once every dish that contributed to it has
delivered its share, because contributions going to the same cook
must be added together first (the InputBuffer). If someone
sneaks back and rewrites a saved ingredient after its receipt was
stapled, the crew refuses to close the books (the version counter
check). And because receipts are generated for every operation
in advance from a recipe book (the derivatives.yaml schema), no
cook ever has to write one by hand.
Why the grown-ups did it this way: when the computation graph depends on Python control flow (loop lengths, data-dependent branches), the only graph builder that can capture it faithfully is the execution itself โ so the design records while it runs, and pays for that choice on every single op call.
1. Architecture overview
flowchart TD
PY["Python: torch package<br/>(torch/__init__.py:2635)<br/>torch.autograd"]
subgraph csrc["torch/csrc โ Python boundary"]
GEN["generated bindings<br/>(python_torch_functions.cpp etc.)<br/>+ PythonArgParser :212"]
THPV["THPVariable<br/>wrap / unpack<br/>(python_variable.cpp:408)"]
PYENG["THPEngine / PythonEngine<br/>(python_engine.cpp:120)"]
end
subgraph disp["c10 dispatcher โ aten/src/ATen/core/dispatch"]
DCALL["Dispatcher::call :784"]
OE["OperatorEntry::lookup :182<br/>flat array indexed by key"]
DKS["DispatchKeySet<br/>(c10/core/DispatchKeySet.h:50)"]
end
subgraph grad["autograd layer โ torch/csrc/autograd"]
VT["generated VariableType kernels<br/>(Autograd dispatch key)<br/>record + redispatch"]
NODES["generated Node subclasses<br/>(MulBackward0, ...)"]
ENG["Engine + GraphTask<br/>+ ReadyQueue (engine.cpp)"]
AG["AccumulateGrad"]
end
subgraph aten["ATen โ backend kernels"]
SK["structured kernel mul_out<br/>(BinaryOps.cpp:441)"]
TI["TensorIterator<br/>(TensorIterator.cpp:1467)"]
STUB["DispatchStub: AVX2/AVX512<br/>(cpu/BinaryOpsKernel.cpp:126)"]
CUDAK["cuda/BinaryMulKernel.cu:22"]
end
subgraph c10core["c10 โ tensor core"]
TI2["TensorImpl :569<br/>+ Storage, intrusive_ptr"]
AM["AutogradMeta (nullable)<br/>(variable.h:231)"]
end
Y1["native_functions.yaml:4246 mul.Tensor"]
Y2["derivatives.yaml:1251 mul.Tensor"]
CG["codegen: torchgen/ + tools/autograd/<br/>generates bindings, VariableType,<br/>Node classes, redispatch fns"]
PY --> GEN --> THPV --> DCALL
DCALL --> OE --> VT
VT -->|"record"| NODES
VT -->|"redispatch below Autograd"| SK
SK --> TI --> STUB
SK --> CUDAK
SK --> TI2
TI2 --- AM
PY -->|"backward()"| PYENG --> ENG --> NODES --> AG --> AM
Y1 & Y2 --> CG
CG -.->|"generates"| GEN & VT & NODES
DKS -.-> DCALL
Diagram source โ enable JavaScript to render the diagram.- Python is a skin, not the body.
torch/__init__.py:2635doesfrom torch._C._VariableFunctions import *; those functions are generated C bindings (tools/autograd/templates/ python_torch_functions.cpp), one staticPyObject*function per operator, each fronting a cachedPythonArgParser(torch/csrc/utils/python_arg_parser.h:212). The result wraps back into Python viaTHPVariable_WrapWithType(torch/csrc/autograd/python_variable.cpp:408), which allocates (or finds โ see the PyObject slot below) the Python wrapper for aTensorImpl. - One dispatcher is the entire language boundary. Every op โ
Python-called, C++-called, custom, or compiled โ goes through
Dispatcher::call(Dispatcher.h:784): compute aDispatchKeySetfrom the arguments' tensors (each tensor carries its key set in itsTensorImpl), thenOperatorEntry::lookup(OperatorEntry.h:182) โ a single array index into a per-op table of kernels, sized to stay cache-resident. What looks like "the C++ API" is the same table reached from the other side. - The Autograd layer is a dispatch key, not a wrapper class.
Old frameworks wrapped tensors in differentiable shells. PyTorch
registers generated autograd kernels at the
Autograddispatch key (tools/autograd/gen_variable_type.py:912emits theTORCH_LIBRARY_IMPL(aten, Autograd, ...)registrations). Each kernel runs the backward-recipe (fromderivatives.yaml), saves what backward needs, redispatches to the backend key below autograd, then staples theNodeonto the output (set_history,torch/csrc/autograd/functions/utils.h:67). - ATen kernels never know autograd exists. A structured op
like
mul.outimplements oneTORCH_IMPL_FUNC(mul_out)(aten/src/ATen/native/BinaryOps.cpp:441) that hands a preparedTensorIteratorto aDispatchStub(aten/src/ATen/native/ DispatchStub.h) โ which picks the AVX2/AVX512/NEON-compiled version ofmul_kernel(cpu/BinaryOpsKernel.cpp:126) or the CUDA one (cuda/BinaryMulKernel.cu:22).TensorIterator::build(TensorIterator.cpp:1467) did the hard part: broadcasting, dtype promotion, dimension reordering, dimension coalescing โ so the kernel body is a flat, vectorizable loop. - The spine is one op call end-to-end: Python
torch.mulโ generated binding โDispatcher::callโ generated Autograd kernel (record) โ redispatch โmul_outโTensorIteratorโ vectorized CPU loop โ wrapped output with a freshgrad_fnโ and later,backward()walking the same receipts in reverse. The Walkthrough traces exactly this.
Walkthrough: one multiply recorded, then walked backwards
sequenceDiagram
autonumber
participant P as Python
participant B as generated binding + Parser
participant D as Dispatcher::call
participant V as VariableType kernel (Autograd key)
participant K as mul_out โ TensorIterator
participant E as Engine (backward)
participant G as AccumulateGrad (leaf .grad)
P->>B: torch.mul(a, b) both require grad
B->>D: parse args โ at::mul(a, b)
D->>V: keyset = CPU + Autograd โ lookup
V->>V: make MulBackward0 node
V->>V: next_edges โ a and b grad_fn
V->>V: save a and b into the node
V->>K: redispatch with Autograd excluded
K-->>V: result tensor
V->>V: set_history(result, node)
V-->>P: THPVariable_Wrap(result)
Note over P: training loop ... then loss.backward()
P->>E: _engine_run_backward โ Engine::execute
E->>E: compute_dependencies (in-degree count)
E->>E: seed ReadyQueue with root
E->>E: thread_main โ MulBackward0::apply
E->>G: InputBuffer โ countdown โ AccumulateGrad
G-->>P: a.grad and b.grad set
Diagram source โ enable JavaScript to render the diagram.The order enters reception
torch.mul(a, b) lands
in a generated binding (shaped by tools/autograd/ gen_python_functions.py:835): a static PythonArgParser
(python_arg_parser.h:212) compiled once per process matches the
arguments against the overload signatures generated from
native_functions.yaml โ mul.Tensor before mul.Scalar, zero-dim
tensors binding to Scalar unless they require grad (the parser's
header comment documents these rules). The tensors are unwrapped
from their Python wrappers (THPVariable_Unpack) โ no copies, the
Python object and the C++ tensor share one TensorImpl.
Dispatch by tag, not by type
Dispatcher::call
(Dispatcher.h:784) extracts each tensor argument's
DispatchKeySet and ORs them; a and b on CPU with
requires_grad=True yield a set containing AutogradCPU. One
array lookup (OperatorEntry::lookup, OperatorEntry.h:182)
selects the highest-priority functionality's kernel โ the
generated VariableType kernel registered at the Autograd key.
The receipt is prepared before the dish exists
The
kernel (generated from tools/autograd/derivatives.yaml:1251,
the mul.Tensor entry whose formula is grad * other and
grad * self) checks whether any input requires grad and grad
mode is on; if so it creates a MulBackward0 Node
(tools/autograd/gen_variable_type.py:637 ASSIGN_GRAD_FN,
:704 SET_HISTORY templates show the generated shape), whose
next_edges come from collect_next_edges(a, b)
(torch/csrc/autograd/function.h:71) โ i.e. the inputs'
existing grad_fns (or their AccumulateGrad accumulators, for
leaves). The operands the backward formula needs are saved into
the node (SavedVariable), each pinned at its current version
counter.
Compute, below autograd
The kernel calls the
generated at::redispatch::mul_ with the autograd keys excluded
(gen_variable_type.py:664 CALL_REDISPATCH): the dispatcher
lands on the CPU backend kernel TORCH_IMPL_FUNC(mul_out)
(BinaryOps.cpp:441) โ mul_stub โ mul_kernel
(cpu/BinaryOpsKernel.cpp:126). TensorIterator::build
(TensorIterator.cpp:1467) has broadcast shapes, promoted dtypes,
reordered and coalesced dimensions; mul_kernel then runs
cpu_kernel_vec with a scalar lambda and a Vectorized<T> lambda
side by side โ with special cases for bool (logical and) and a
fast path when one operand is a Python scalar. The result
TensorImpl shares nothing with the inputs' storage.
Staple and return
Back in the autograd kernel,
set_history (functions/utils.h:67) calls
impl::set_gradient_edge(result, {node, output_nr}) โ the output's
AutogradMeta::grad_fn_ now points at MulBackward0. The
generated binding wraps the tensor (THPVariable_WrapWithType,
python_variable.cpp:408); the PyObject and the TensorImpl keep
each other alive (an impl::PyObjectSlot on the TensorImpl,
c10/core/TensorImpl.h:2269). The forward pass was the graph
builder. Nothing else was recorded anywhere.
Close the books
loss.backward() goes
torch/autograd/graph.py:1100 _engine_run_backward โ
THPEngine_run_backward (python_engine.cpp:179) โ
PythonEngine::execute (python_engine.cpp:120, which
TORCH_CHECKs the GIL is released โ expensive C++ work must
not serialize Python) โ Engine::execute
(engine.cpp:1428). It builds a GraphTask
(graph_task.h:18), wraps the roots in a GraphRoot, and
compute_dependencies (engine.cpp:1360) walks next_edges
counting each node's in-degree into dependencies_ โ a plain
unordered_map<Node*, int>; no topological sort, no locks per
edge.
Walk the receipts
The calling thread itself becomes
the CPU worker (thread_main, engine.cpp:575); GPU-device nodes
are executed by one persistent thread per device created lazily by
start_device_threads (engine.cpp:1725, each named
pt_autograd_N, engine.cpp:344 thread_init). Each
evaluate_function (engine.cpp:1168) runs the node's apply
via call_function (engine.cpp:1114 โ pre-hooks, the node call,
post-hooks, output shape validation), then for each output edge
decrements the target's dependency count; a target whose count
reaches zero graduates from its staging InputBuffer (partial
gradient sums accumulated in not_ready_, input_buffer.h:41)
into a NodeTask pushed on the target device's ReadyQueue
(engine.h:87 โ a priority queue ordered by sequence_nr, which
is why AccumulateGrad constructs itself with
sequence_nr = UINT64_MAX to always sort first (set in its
constructor at functions/accumulate_grad.cpp:87, applied in
AccumulateGrad::apply at :92).
Leaves hold their gradients, and the tape is spent
Gradient paths ending at leaves reach AccumulateGrad::apply
(accumulate_grad.cpp:92): under the node's mutex it either
steals the incoming gradient (when the engine's refcount protocol
says nobody else holds it โ call_function bumps the expected
count exactly when post-hooks exist) or accumulates into
variable.grad_ in place. When outstanding_tasks_ โ incremented
by every ReadyQueue::push (engine.cpp:234) and decremented on
every completed task โ hits zero, the GraphTask's future
completes; Engine::execute's fut->wait() returns the captured
gradients. With keep_graph=False each node released its saved
tensors as it ran: the tape is single-use by default.
2. Core data structures โ and why these
TensorImplโ one refcounted body, three optional souls. What:c10/core/TensorImpl.h:569โ storage pointer, dtype, device, sizes/strides (possibly symbolic), aDispatchKeySet, a nullableAutogradMetapointer, aPyObjectslot, and a sharedVariableVersioncounter. Owned byintrusive_ptr; the Pythontorch.Tensorand the C++Tensorare both thin handles to it. Why: Python refcounting, C++ refcounting, and autograd state each get to exist or not exist independently โ a plain float tensor pays zero autograd bytes (autograd_meta_stays null, see the invariant note atTensorImpl.h:3003), and a C++-only workload (libtorch) never materializes a PyObject. The descriptor/physical split of llama.cpp'sggml_tensor, taken further: here even language identity is an optional slot.DispatchKeySetโ dispatch as a uint64 bitmap. What:c10/core/DispatchKeySet.h:50โ a 64-bit word of ordered "functionality" bits (Autograd above backend) plus backend bits; every tensor carries one, and per-thread include/exclude masks (guards liketorch.no_gradclear the Autograd bit) refine it at call time (DispatchKeyExtractor,aten/src/ATen/core/dispatch/DispatchKeyExtractor.h:39). Why: multiple dispatch over N arguments collapses into an OR of bitsets plus onecountr_zero-style priority pick โ the dispatch step is O(1), branch-predictable, and needs no per-op virtual call. The cost is a fixed bit budget (see the "No More Than 16 Backends" note in the header) and aliasing rules that the file documents at length โ a fair price versus a type-tower.EdgeandNodeโ the tape as ownership graph. What:torch/csrc/autograd/edge.hโ anEdgeis exactly(intrusive_ptr<Node>, input_nr); aNode(torch/csrc/autograd/node.h) owns itsnext_edges_via refcounted pointers, plussequence_nr(creation order โ queue priority),topological_nr(an internal DAG property used to prune traversal), input metadata for shape validation, and amutex_for concurrentapply(node.h:682documents exactly when users must lock it). Why: the tape is not a separate data structure the framework stores somewhere โ it is literally the ownership chain of small objects hanging off tensors'AutogradMeta::grad_fn_. Free the tensor, and if no other output keeps the node alive, the subgraph below it deallocates itself. Liveness and graph extent are one fact.GraphTaskโ the run-state of one backward. What:torch/csrc/autograd/graph_task.h:18โ atomicoutstanding_tasks_, thedependencies_in-degree map,not_ready_stagedInputBuffers,exec_info_(which paths to actually execute wheninputs=is passed), captured leaf streams for CUDA, and the completion future. Why: everything a lock-light dataflow scheduler needs, and nothing else: no global pass over the graph at completion (the counter is the barrier), no scheduler thread (workers pull), no per-node promise. Compare the countdown pattern with Kubernetes controllers โ same "count dependent edges" idea, here paid per-node in nanoseconds.TensorIteratorโ the great leveler of elementwise kernels. What:aten/src/ATen/TensorIterator.cpp:1467 buildโ fills operand slots, checks memory overlap, computes the broadcast shape, promotes a common dtype, permutes dimensions for coalescing, coalesces adjacent dims, allocates outputs; afterwards the kernel sees flat pointers/strides and (usually) a contiguous 1-D loop. Why: one C++ class absorbs the combinatorial horror of NumPy semantics (broadcast + cast + layout) so that ~200 pointwise ops are each a dozen lines โmul_kernelis literallya * btwice, scalar and vector. The ATen "structured kernel" split (mul.Tensordelegates tomul.outviastructured_delegatein the yaml) means the out-variant owns the iterator setup for all three variants.SavedVariable+VariableVersionโ snapshots that can complain. What: saving for backward pins the tensor's data pointer and records its version counter (saved_variable.h);unpack(saved_variable.cpp:166) refuses ifsaved_version_ != current_versionโ the error users know as "one of the variables needed for gradient computation has been modified by an inplace operation." Why: in-place ops bump the shared counter (TensorImpl.h'sbump_version, shared between a view and its base), so correctness of the tape doesn't require copying saved inputs โ just detecting staleness cheaply. A validation budget of one atomic counter, checked once per save.
3. Concurrency and failure model
Two parallelism worlds, deliberately separate. Forward execution is single-threaded per op call; parallelism lives inside kernels (CPU thread pools, CUDA streams). Backward is parallel across the DAG:
Engineruns one worker thread per device (created on first backward touching that device,engine.cpp:1725) plus whatever thread calledbackward()as the CPU worker (engine.cpp:1585). No lock protects "the graph" โ the graph is immutable during execution; coordination happens through per-queue mutexes, the in-degree map, andInputBufferaccumulation under the GraphTask's own mutex (engine.cpp:1223).Readiness is a countdown, not a scheduler. A node runs when its dependency count hits zero (
engine.cpp:1295), because all its incoming gradients have arrived and been summed into itsInputBufferโ an implicit addition node (the file comment ininput_buffer.hcalls it exactly that). Priorities inside eachReadyQueueare(reentrant depth, sequence_nr)(engine.h:87), so within a device, work is approximately topological, which keeps peak memory of intermediate gradients bounded without any global ordering pass.CUDA correctness is streams-and-events bookkeeping. Each produced gradient records the stream it was computed on;
InputBuffer::add(input_buffer.cpp:350) records events andevaluate_function(engine.cpp:1203) makes the consumer's stream wait on the producer's event when they differ. The GraphTask collectsleaf_streamsandexec_post_processingsyncs them with the caller's streams โ the semantics documented publicly as "stream semantics of backward passes." The Nodemutex_(node.h:682) serializes concurrentapplyon the same node, which is what makes the whole scheme safe without per-gradient locks.Failure: poison the task, not the thread. A throwing node's exception is captured into the GraphTask (
graph_task.h:set_exception),has_error_flips, other workers drain the queue without executing further nodes, and the exception is rethrown to the caller when the future is joined inEngine::execute. A Python exception raised in a Python-side node crosses this boundary throughPythonEngine::execute'spython_errorrestore path (python_engine.cpp:120). CUDA kernels failing asynchronously surface at the next synchronization point โ by design, the engine never polls the GPU.Reentrant backwards are a bounded recursion. A
Node::applythat itself callsbackward()would recurse throughEngine::execute; beyondMAX_DEPTH = 60levels (engine.h:38) the engine stops recursing and hands the new GraphTask to a small standby thread pool (Note [Reentrant backwards],engine.cpp:146). The constant is honest about why: TSAN deadlocks past ~65 held mutexes.
4. API and extension design
The public surface is schema-first. An operator exists once in
aten/src/ATen/native/native_functions.yaml(with aliasing and mutation annotations likeTensor(a!)), and codegen emits: C++ declarations, Python bindings,out=variants, structured kernel scaffolding, meta (shape) kernels, the dispatch registrations โ and, paired with aderivatives.yamlentry, the autograd kernel and theNodeclass. The yaml is the API reference; the Redis commands.json pattern at 2,000-op scale.Five extension seams, all funnels into one table. (1)
torch.library/TORCH_LIBRARY(torch/library.h:974): register a new operator and kernels at any dispatch key โ Python'storch._library.custom_ops.py:272 CustomOpDefwraps it with schema inference and a pt2-compliance tag. (2) Customautograd.Functionโ a PythonNodesubclass the engine runs with the GIL held just for it. (3)PrivateUse1โ a reserved dispatch key so out-of-tree backends (TPUs, NPUs) plug in without forking. (4) TorchDispatch/TorchFunction modes โ thread-local interceptors that see every op call (how functionalization, functorch, and profiling are built). (5)torch.compilebackends โ a registry of graph handlers receiving the FX graph Dynamo extracted. All five converge on the same dispatcher; there is no second system.torch.compile is the eager design's honest sequel, not a replacement. Dynamo installs a CPython frame-evaluation hook (PEP 523;
torch/csrc/dynamo/eval_frame.c:329 enable_eval_frame_shimโ_PyInterpreterState_SetEvalFrameFunc), rewrites bytecode to trace tensor ops into an FX graph guarded on shapes/properties (torch/_dynamo/convert_frame.py:1703 _compile), and lowers through AOTAutograd and Inductor's generated Triton/C++ (torch/_inductor/compile_fx.py:3130). Because ops were already funneled through one dispatcher with schema and meta kernels, capture did not require a new runtime โ the same seams that gave custom backends gave the compiler.Honest warts. The generated layer is invisible in a fresh clone:
torch/csrc/autograd/generated/exists only after a build, so reading the "real" autograd kernel means reading Python generators โ and the codebase knows it (the generatedVariableTypeEverything.cppexists "for ease of inspection... only for convenience; not actually used in the build," per the sharding note intools/autograd/templates/VariableType.cpp). Schema-string overload sprawl (mul.Tensor/mul.Scalar/mul.out/mul_.Tensor...) is the price of uniformity. And the compile stack's depth (Dynamo โ AOTAutograd โ Inductor) makes failure diagnosis a genuinely hard skill โ the flexibility that won research is exactly what capture must fight to recover.
5. The defining trade-offs
Trade-off 1: define-by-run recording over pre-declared static graphs
- Tension: ML frameworks before PyTorch made you declare the graph, then run it โ enabling whole-graph optimization but breaking Python control flow, debugging, and shape dynamism.
- Choice: make execution the graph builder. The generated
Autograd kernel records while computing; the "graph" is never
materialized anywhere else but the live tensors'
grad_fnchain. - Evidence:
tools/autograd/gen_variable_type.py:664CALL_REDISPATCHand:704SET_HISTORYtemplates โ the generated kernel first computes with the backend kernel, then attaches the node;tools/autograd/derivatives.yaml:1251supplies the backward formula as data;functions/utils.h:67 set_historyis the staple. The bill for the choice is paid later: no eager pass can fuse ops, so the codebase eventually built a tracing compiler on top (eval_frame.c:329) to recover the static graphs it deliberately never had. - Given up: whole-graph optimization in the default path (fused kernels, memory planning), plus per-op dispatch cost โ both recovered only partially, and ten years later, by torch.compile.
- Verdict: right for the research workload that made PyTorch
dominant โ models whose structure changes with every run, where
the value of printing/branching/debugging mid-graph exceeds
fused-kernel throughput. Wrong when your model is fixed and
shapes are static โ ship the compiled artifact (the same
codebase's
torch.compile/torch.exportexist precisely for that phase transition).
Trade-off 2: thin generated bindings over a C++ dispatcher core
- Tension: the frontend everyone wanted was Python, but the kernels needed C++ speed and vendors needed a Python-free path; per-call overhead of crossing CPython is real.
- Choice: Python bindings exist only as generated stubs over
at::functions; all real logic sits behind one dispatcher reachable identically from C++ (libtorch) and Python. - Evidence: the cost side is visible at
gen_python_functions.py:835(every call parses Python objects against compiled signature strings) andDispatcher.h:784 callโOperatorEntry.h:182 lookup(keyset computation and one more indirect call per op); the payoff side is thatautograd_not_implemented_fallback.cpp, TorchDispatch modes, profiling, and torch.compile all hook the same single point โ and libtorch ships the whole engine with no CPython linked. - Given up: a hand-tuned Python-first fast path; the per-op constant factor that frameworks like llama.cpp eliminate by having no frontend at all โ PyTorch pays it on every call and eventually attacks it with a bytecode-level JIT.
- Verdict: right when many frontends and many backends must meet at one semantic core, and ops are heavy enough (real matmuls) that a few hundred ns of dispatch vanish in the kernel. Wrong for tiny-op-dominated workloads โ which is a large part of why torch.compile exists.
Trade-off 3: generate everything from two YAML schemas
- Tension: 2,000+ operators ร 4โ6 artifacts each (C++ decl, Python binding, dispatch registration, autograd kernel, Node class, meta kernel) โ hand-maintaining them guarantees drift.
- Choice:
native_functions.yaml+derivatives.yamlare the single source;torchgen/andtools/autograd/generate everything (theVariableType.cpptemplate even splits itself into compile-parallel shards because the unsplit file was 36,000 lines โ the template's own sharding note). - Evidence:
native_functions.yaml:4246declares all fourmuloverloads in 40 lines, including sparse/Mkldnn dispatch map entries;tools/autograd/gen_autograd.py's header states the build destination (torch/csrc/autograd/generated/);torchgen/gen.py:2581-2646writesRegisterBackendSelect.cpp,RegisterSchema.cpp,Operators.cpp,Functions.cppfrom the same parse. - Given up: a readable, greppable implementation โ the "real"
autograd code for
muldoes not exist as a committed file, and reviewing a PR means reading generator templates plus imagined output. Build complexity and a Python-based build step are permanent tenants. - Verdict: right at this scale โ the alternative is the pre-2017 framework world where bindings, docs, and autograd drifted apart per-op. Wrong below a few dozen ops, where codegen is ceremony (curl hand-writes its surface and stays coherent). This is PATTERNS.md's "generate the surface from data," taken to its maximum credible extent.
6. Reading path
torch/csrc/autograd/variable.h:53โ the Variable doc block: leaf vs. interior,grad_fnvs.grad_accumulator, versions, views. The densest design document in the codebase; do not skip it.torch/csrc/autograd/edge.hthentorch/csrc/autograd/node.hโEdgeis(Node, input_nr); readNode'snext_edges_,sequence_nr,topological_nr, themutex_contract at:682, andoperator()at:172.aten/src/ATen/native/native_functions.yaml:4246โ the fourmuloverloads; noticeTensor(a!)mutation annotations,structured: Trueonmul.out, and per-backend dispatch entries. Thentools/autograd/derivatives.yaml:1251โ the backward formulas as data.tools/autograd/gen_variable_type.py:614-726โ the templates that are the autograd kernels (SETUP_DERIVATIVE,CALL_REDISPATCHat:664,SET_HISTORYat:704); then:912for theTORCH_LIBRARY_IMPL(aten, Autograd)wrapper generation. Optionally runpython -m tools.autograd.gen_autograd(its docstring shows the invocation) to materializetorch/csrc/autograd/generated/and read the realVariableTypeEverything.cpp.torch/csrc/autograd/functions/utils.h:67 set_historyโ the staple:add_input_metadata+set_gradient_edge. Three lines that make every differentiable op "record."torch/csrc/autograd/python_variable.cpp:408 THPVariable_WrapWithTypeandtorch/csrc/utils/ python_arg_parser.h:212 PythonArgParserโ the CPython boundary in two structures; note the PyObject/TensorImpl co-ownership and the overload-matching gotchas in the header comment.aten/src/ATen/core/dispatch/Dispatcher.h:784 call+OperatorEntry.h:182 lookup+c10/core/DispatchKeySet.h:50โ the O(1) heart: bitset OR, one array index. ThenDispatchKeyExtractor.h:39for how TLS guards (no_grad, inference mode) mask keys in and out.aten/src/ATen/native/BinaryOps.cpp:441 TORCH_IMPL_FUNC (mul_out)โaten/src/ATen/native/cpu/BinaryOpsKernel.cpp:126 mul_kernelโaten/src/ATen/native/DispatchStub.hโ the kernel side: structured op โ arch-dispatched stub โ a two-lambda loop body. Readcpu/Loops.h'sbasic_loopnext.aten/src/ATen/TensorIterator.cpp:1467 buildโ walk one call:compute_shape,compute_types,reorder_dimensionsat:226,coalesce_dimensionsat:602. This is where NumPy semantics are actually implemented.torch/autograd/graph.py:1100 _engine_run_backwardโtorch/csrc/autograd/python_engine.cpp:179 THPEngine_run_backwardand:120 PythonEngine::executeโ backward's front door, and the GIL discipline around it.torch/csrc/autograd/engine.cppin this order::1428 execute(GraphTask, GraphRoot, dependencies) โ:1360 compute_dependenciesโ:234 ReadyQueue::pushโ:575 thread_mainโ:1168 evaluate_function(the countdown, the InputBuffer staging, the stream guards) โ:1114 call_function. This file is the best-written ~2,000 lines of scheduling code you will read this year.torch/csrc/autograd/functions/accumulate_grad.cpp:92 AccumulateGrad::applyโ steal-vs-clone refcount logic, thesequence_nr = UINT64_MAXtrick (constructor at:87), the mutex; thentorch/csrc/autograd/saved_variable.cpp:166โ the version check every "modified by an inplace operation" error comes from.torch/csrc/dynamo/eval_frame.c:329 enable_eval_frame_shimandtorch/_dynamo/convert_frame.py:1703 _compileโ the sequel: how the codebase later recovered static graphs without breaking the eager contract from steps 1โ12.
You have now traced one multiply from a Python call through the generated binding, the dispatcher, a recorded autograd node, a vectorized CPU kernel, and back โ and then the entire backward pass that the recorded tape made possible.
7. What to steal
- Make execution the recorder โ when users can't or won't
describe the structure up front (control-flow-dependent work,
debugging workflows), build the trace as a side effect of the
real path, one receipt at a time. Seen at
gen_variable_type.py's SET_HISTORY template: three lines after every compute. - Nullable capability fields on hot objects โ a tensor that
doesn't need autograd stores null
AutogradMetaand pays nothing; a tensor Python never touches stores no PyObject. Zero-cost abstractions via optional slots, not subclass trees (TensorImpl.h:569). - Schedule DAGs with counters and buffers, not a scheduler โ
in-degree map + staged input buffers + per-worker queues +
one atomic outstanding counter as the barrier: no scheduler
thread, no graph locks, naturally parallel across devices
(
engine.cpp:1168 evaluate_function). - Version counters to make cheap snapshots safe โ share one
atomic counter between an object and its views; snapshots pin
the version and refuse to run if it moved (
saved_variable.cpp:166). - Generate the API surface from schema data โ one yaml
entry becomes bindings, kernels, autograd, docs, and compile
metadata; drift between them becomes structurally impossible
(
native_functions.yaml, the Redis pattern at 100ร scale). - Put every extension point at the same dispatch table โ
custom ops, custom backends, tracing modes, and compilers all
register at one seam (
torch/library.h:974), so the ecosystem never forks the core โ the placement decision that made PyTorch a platform rather than a library.
8. Difficulty and prerequisites
- Difficulty:
demanding. Realistic budget: 9โ12 hours for this path. The C++ is modern and heavily commented, but two genuinely hard things await: the generated-code indirection (much of what runs does not exist as committed source โ budget an hour to run the codegen or read generators fluently) and the engine's synchronization reasoning (read it twice). - Prerequisites: what a gradient and the chain rule are;
comfortable C++ (templates,
intrusive_ptrrefcounting); Python C API at the level of "aPyObject*wraps our object"; DAG vocabulary (topological order, in-degree). CUDA streams only for section 3's depth โ the CPU path tells the full story. CPython is the natural companion dive; bring llama.cpp as the deliberately opposite design (no frontend, no dispatcher) for contrast.
9. Further reading
- The upstream Autograd mechanics note and The dispatcher notes on docs.pytorch.org โ the public version of what section 1 walks.
- In the clone:
GLOSSARY.md(operation/kernel/node vocabulary),aten/src/ATen/core/dispatch/README.md, and the notes referenced inengine.cpp(Note [Reentrant backwards], Note [Streaming backwards]). - PyTorch: An Imperative Style, High-Performance Deep Learning Library (Paszke et al., NeurIPS 2019) โ the design rationale in the authors' words, including the eager-over-static argument this dive's Trade-off 1 reconstructs.
- Automatic Differentiation in Machine Learning: a Survey (Baydin et al., 2018) โ where "define-by-run" vs "define-and-run" got their names.
- Deep Learning Compiler PyTorch 2.0 (Ansel et al., 2023, arXiv 2309.06180) โ what torch.compile (section 4) does with the tape-less eager design once graphs are recovered.