ATen

Active contributors: ezyang, mruberry, albanD

Purpose

ATen ("A Tensor Library") is the C++ tensor library that exposes everything torch.* does at the math level. It owns the operator definitions, the dispatcher, the in-tree native kernels, and the per-backend specializations. If you call a + b in Python, the Python binding eventually calls at::add(a, b) from ATen.

ATen lives at aten/src/ATen/. The namespace is at::.

Directory layout

Path	Contents
aten/src/ATen/	Public ATen headers and core sources
aten/src/ATen/core/	The dispatcher, op registry, schema parser, IValue, Library
aten/src/ATen/native/	The vast majority of CPU op implementations (and a few cross-cutting ones)
aten/src/ATen/native/native_functions.yaml	The op registry — the single source of truth for op schemas
aten/src/ATen/native/cpu/	Vectorized CPU kernels (AVX2, AVX512, NEON, …)
aten/src/ATen/native/cuda/	CUDA kernels
aten/src/ATen/native/cudnn/	cuDNN bindings (convolutions, RNNs)
aten/src/ATen/native/mkldnn/	oneDNN/MKL-DNN integration
aten/src/ATen/native/mps/	Apple Metal Performance Shaders kernels
aten/src/ATen/native/xpu/	Intel XPU stubs (real kernels live in intel-xpu submodule)
aten/src/ATen/native/sparse/	Sparse tensor ops
aten/src/ATen/native/quantized/	Quantized ops
aten/src/ATen/native/nested/	Nested tensor ops
aten/src/ATen/cuda/	CUDA runtime helpers (Blas, Solver, etc.) shared across kernels
aten/src/ATen/templates/	Mustache templates that torchgen renders
aten/src/ATen/test/	Gtest-based C++ unit tests

The directory is roughly 21K files and far too large to itemize.

Key abstractions

Type	File	Purpose
at::Tensor	aten/src/ATen/templates/TensorBody.h	Public C++ tensor type (intrusive_ptr)
at::TensorIterator	aten/src/ATen/TensorIterator.h	The shared loop driver for elementwise/reduction ops
at::Context	aten/src/ATen/Context.h	Global context: backend toggles, deterministic mode, BLAS choice
c10::Dispatcher	aten/src/ATen/core/dispatch/Dispatcher.h	The op routing table
at::OperatorHandle	aten/src/ATen/core/dispatch/Dispatcher.h	Handle to a registered operator
c10::FunctionSchema	aten/src/ATen/core/function_schema.h	Parsed op schema
at::native::*	aten/src/ATen/native/	Where most leaf op CPU implementations live
at::AutogradMeta	aten/src/ATen/core/Variable.h	Per-tensor autograd metadata (grad_fn, hooks, version counter)

The Context singleton (at::globalContext()) holds toggles like userEnabledCuDNN(), deterministicAlgorithms(), blasPreferredBackend(), etc.; this is the C++ side of torch.backends.*.

How it works

ATen has three logical pieces: the dispatcher (covered in Dispatcher), the op declaration system (in native_functions.yaml and core/), and the kernel implementations (native/).

Op declaration

Every operator is described by a YAML entry like:

- func: add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor
  device_check: NoCheck
  structured_delegate: add.out
  variants: function, method
  dispatch:
    SparseCPU, SparseCUDA: add_sparse
    SparseCsrCPU, SparseCsrCUDA: add_sparse_csr
    MkldnnCPU: mkldnn_add
    ZeroTensor: add_zerotensor
  tags: [pointwise, canonical]

Reading top-down: the schema, behavioural flags, the structured "delegate" that sub-kernels share, where it appears (function and/or method), the per-key dispatch table, and tags. torchgen consumes these and emits:

• a header in build/aten/src/ATen/Operators.h declaring at::add(...);
• generated registrations in RegisterCPU.cpp, RegisterCUDA.cpp, etc. that wire each dispatch: entry into the dispatcher;
• generated stubs for autograd's derivatives.yaml entries.

For "structured" kernels (the modern style for elementwise ops), torchgen also emits a meta function (shape inference) and an iter-driven base; the kernel author only writes the inner loop. See aten/src/ATen/native/README.md for the canonical guide.

TensorIterator

TensorIterator (TensorIterator.h/.cpp, ~100K lines combined) is the heart of pointwise and reduction kernels. It computes broadcasting, dtype promotion, output allocation, and an iteration order that yields contiguous strided loops. CPU and CUDA both build a TensorIterator and then call vectorised loops on it; the bulk of CPU kernels in aten/src/ATen/native/cpu/*Kernel.cpp are short loops over an iter.

Op call sequence

A simplified eager-mode at::add(a, b):

sequenceDiagram
    participant Caller
    participant OpsH as Operators.h<br/>(generated)
    participant Disp as Dispatcher
    participant Reg as RegisterCUDA.cpp<br/>(generated)
    participant Kernel as native::add_kernel_cuda

    Caller->>OpsH: at::add(a, b)
    OpsH->>Disp: redispatch using a.key_set() | b.key_set()
    Disp->>Disp: pick highest-priority key with kernel
    Disp->>Reg: jump to wrapper for AutogradCUDA / CUDA
    Reg->>Kernel: native::add_kernel_cuda
    Kernel-->>Caller: returns Tensor

Backend layering

aten/src/ATen/native/ is where almost all in-tree kernels live. Backend-specific kernels nest under cpu/, cuda/, mps/, cudnn/, mkldnn/, quantized/, sparse/, nested/, nnapi/, vulkan/. Files in aten/src/ATen/cuda/, aten/src/ATen/mps/, etc. provide runtime helpers (BLAS handles, kernel launch utilities) but not op kernels.

External backends (XLA, MTIA, lazy) register their own kernels through TORCH_LIBRARY_IMPL from out-of-tree code.

Integration points

• Generated code consumers. Operators.h, Functions.h, RedispatchFunctions.h, the RegisterX.cpp files, and Python bindings under torch/csrc/autograd/generated/python_*.cpp are all produced from native_functions.yaml. See torchgen.
• Autograd. Forward kernels are wrapped by autograd's per-op codegen; see Autograd.
• Compile. AOT autograd traces through ATen ops symbolically (using FakeTensor); decompositions in torch/_decomp/ rewrite ATen ops in terms of simpler ops; see Tensor subclasses and AOT Autograd.
• JIT / TorchScript. Reuses the same dispatcher and operator registry. JIT IR ops are ATen ops.
• C API stable. torch/csrc/stable/ exposes a small ABI-stable subset of ATen for out-of-tree extensions; see C API stable.

Entry points for modification

To add an op:

1. Add a func: entry in aten/src/ATen/native/native_functions.yaml with a schema and dispatch table.
2. Implement the kernel(s) under aten/src/ATen/native/. For CPU pointwise ops put the inner loop in aten/src/ATen/native/cpu/<your>Kernel.cpp and a TensorIterator setup in aten/src/ATen/native/<your>.cpp.
3. If differentiable, add a derivative to tools/autograd/derivatives.yaml.
4. Add a Python binding entry if you need a non-default name (otherwise it's automatic).
5. Write tests in test/test_ops.py via OpInfo (torch/testing/_internal/common_methods_invocations.py) to get the suite of correctness tests for free.

To change an existing op's behaviour, find its entry in native_functions.yaml, follow the dispatch: table to the implementation, and remember to update structured kernel meta functions if the shape rule changes.

Key source files

File	Purpose
aten/src/ATen/native/native_functions.yaml	The op schema registry (single source of truth for ATen ops)
aten/src/ATen/native/README.md	Canonical guide to writing native functions
aten/src/ATen/TensorIterator.h, TensorIterator.cpp	Loop driver for elementwise/reduction ops
aten/src/ATen/Context.h, Context.cpp	Global toggles and backend selection
aten/src/ATen/core/dispatch/Dispatcher.h	The dispatcher
aten/src/ATen/core/op_registration/op_registration.h	Programmatic op registration
aten/src/ATen/templates/TensorBody.h	Public Tensor C++ class
aten/src/ATen/FunctionalTensorWrapper.h	Functionalization wrapper
aten/src/ATen/autocast_mode.h	Autocast (mixed-precision) dispatch
aten/src/ATen/native/cuda/	CUDA kernels
aten/src/ATen/native/cudnn/	cuDNN-backed convolutions, RNNs, attention