PyTorch源码篇之源码架构总览【第一篇】
本文基于pytorch源码分析 PyTorch 的整体架构。PyTorch 是 Meta 开源的动态图深度学习框架,核心特性:
- 动态计算图(Eager Execution 即时执行): 代码运行时实时构建计算图,调试友好;
Python 优先,底层 C++ 高性能内核; 支持 CPU / CUDA / MPS / XPU / ROCm 多硬件后端; 两大核心场景:科研原型实验、工业大模型训练 / 推理; 配套生态: torch.compile编译栈、TorchServe、TorchRec、TorchVision、TorchAudio 等。
核心顶层能力:
- 张量 Tensor: 多维数组,自动微分载体,支持异构设备;
- Autograd 自动微分引擎: 自动构建反向传播图;
- ATen 算子库: CPU/GPU 底层数学算子集合;
- nn.Module 神经网络组件: 层封装、参数管理、优化器;
- Compiler 编译栈(Dynamo + Inductor): 动态图捕获、编译优化、算子融合;
- 分布式训练(c10d): DDP、FSDP、DTensor 分布式张量;
- 推理生态: TorchScript、AOTInductor、Export 模型导出。
两大执行模式:
- Eager Mode(默认即时执行): 运行一行执行一行,动态图,调试简单;
- Compiled Mode(torch.compile): Dynamo 拦截 Python 字节码,捕获计算图,编译优化静态子图。
用户视角: PyTorch 的核心入口是 torch、torch.nn、torch.autograd、torch.optim、torch.utils、torch.distributed、torch.compile 等模块。从源码视角看,更合适的主干分层是:用户 API 层、C++ 桥接层、张量与算子库层、c10 基础抽象层。在这条主干之外,自动求导、dispatcher、编译栈、设备后端、分布式和代码生成系统作为横切机制或扩展层贯穿其中。
一、整体架构和分层
1.1 分层框图
PyTorch 的源码架构可以先按主调用链分层理解:用户从 Python API 进入,经过 C++ 桥接层,下沉到 ATen 张量与算子库,再依赖 c10 基础抽象表达设备、内存、调度和 Tensor 元数据,最后由具体后端执行数值计算。自动求导、编译栈、分布式、代码生成和扩展机制横跨这些层。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 +--------------------------------------------------------------------+| 用户模型与训练代码 || model(x), loss.backward(), optimizer.step(), torch.compile(model) |+--------------------------------------------------------------------+ | v+--------------------------------------------------------------------+| 第一层:用户 API 层 || torch / torch.nn / torch.optim / torch.autograd / torch.distributed || torch.utils / torch.export / torch.compile / torch.fx |+--------------------------------------------------------------------+ | v+--------------------------------------------------------------------+| 第二层:C++ 桥接层 || torch._C / torch/csrc / Python binding / generated bindings || THPVariable / pybind / C extension / C++ frontend |+--------------------------------------------------------------------+ | v+--------------------------------------------------------------------+| 第三层:张量与算子库层 || ATen / at::Tensor / TensorBase / native_functions.yaml || c10::Dispatcher / OperatorEntry / native kernels / torchgen |+--------------------------------------------------------------------+ | v+--------------------------------------------------------------------+| 第四层:c10 基础抽象层 || TensorImpl / Storage / DataPtr / Allocator / Device / ScalarType || DispatchKey / DispatchKeySet / SymInt / intrusive_ptr |+--------------------------------------------------------------------+ | v+--------------------------------------------------------------------+| 第五层:后端执行层 || CPU / CUDA / HIP / XPU / MPS / Vulkan / Metal / MKL / cuDNN / NCCL || backend kernels / memory allocator / stream / event / libraries |+--------------------------------------------------------------------+横切机制层:Autograd / Dispatcher / torch.compile / FX / AOTAutograd / Inductor /Distributed / torchgen / custom op / custom backend / Tensor subclass 这张图表达的是主执行路径。以 torch.add(x, y) 为例,用户 API 层接收 Python 调用,C++ 桥接层把 Python 参数转换为 C++ Tensor 调用,ATen 根据算子 schema 进入 dispatcher,dispatcher 依据 DispatchKeySet 选择 CPU、CUDA、Meta、Autograd 或 Composite 等 kernel,kernel 再依赖 c10 中的 Tensor 元数据、设备、内存和调度抽象完成执行。
1.2 各层功能总表
torch/ | torch.Tensornn.Module、Optimizer、DataLoader、ProcessGroup、GraphModule | ||
torch/csrc/torch/_C | THPVariable | ||
aten/src/ATen/aten/src/ATen/native/、torchgen/ | at::TensorTensorBase、native_functions.yaml、c10::Dispatcher、OperatorEntry | ||
c10/ | TensorImplStorageImpl、DataPtr、Allocator、Device、DispatchKeySet、SymInt | ||
aten/src/ATen/cpucuda、hip、xpu、mps、c10/cuda 等 | |||
torch/autogradtorch/csrc/autograd、torch/_dynamo、torch/_inductor、torch/distributed、tools/autograd | AutogradMetaNode、Engine、FX graph、AOTAutograd、Inductor、DDP/FSDP、torchgen |
1.3 用户 API 层功能表
torch | torch.Tensor | |
torch.nn | ModuleParameter、各类 layer/loss | |
torch.optim | Optimizer | |
torch.autograd | Function | |
torch.utils | DatasetDataLoader、cpp extension | |
torch.distributed | ||
torch.compiletorch.export / torch.fx | GraphModule |
1.4 C++ 桥接层功能表
torch._C | ||
torch/csrc | ||
torch/csrc/autograd | THPVariable | Function 和 C++ autograd engine。 |
torch/csrc/api | ||
native_functions.yaml 中的算子转换成 Python 可调用入口。 | ||
1.5 张量与算子库层功能表
ATen/core | TensorBase | |
ATen/native | native_functions.yaml | |
ATen/ops | ||
c10::DispatcherOperatorEntry、KernelFunction | ||
torchgen | ||
1.6 c10 基础抽象层功能表
c10/core/TensorImpl.h | TensorImpl | |
c10/core/Storage.h | StorageStorageImpl | |
c10/core/Allocator.h | AllocatorDataPtr | |
c10/core/Device.h | DeviceDeviceType | |
c10/core/DispatchKey.h | DispatchKeyDispatchKeySet | |
c10/core/SymInt.h | SymIntSymBool | |
c10/util | intrusive_ptr |
1.7 后端执行层功能表
1.8 横切机制层功能表
torch/autogradtorch/csrc/autograd、tools/autograd | ||
ATen/core/dispatchc10/core/DispatchKey.h | ||
torch/_dynamotorch/fx、torch/_functorch、torch/_inductor | ||
torch/distributedtorch/csrc/distributed | ||
torchgentools/autograd | ||
torch.library |
二、源码目录结构详解
2.1 顶层目录
当前本地源码根目录为 pytorch/,主要目录如下。
torch/ | torch/csrc/、编译栈、分布式、设备模块等。 |
aten/ | |
c10/ | |
torchgen/ | native_functions.yaml、autograd 配置和模板生成 ATen、dispatcher、binding、autograd 等代码。 |
tools/ | |
test/ | |
caffe2/ | |
cmake/ | |
third_party/ | |
functorch/ | torch.func、AOTAutograd 等路径。 |
docs/ | |
benchmarks/ | |
android/ | |
binaries/ | |
scripts/ | |
mypy_plugins/ | |
.github/.ci/ |
2.2 torch/ 目录
torch/ 是 Python 包主体,也是用户最容易接触的源码目录。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 torch/├── nn/ # 神经网络层、loss、初始化、functional API、parallel wrapper├── autograd/ # Python 侧 autograd API├── csrc/ # C++ extension module 源码,连接 Python 和 C++ 核心│ ├── autograd/ # autograd engine、Node、Variable、SavedVariable、Function、binding│ ├── jit/ # TorchScript/JIT C++ 实现│ ├── distributed/ # 分布式 C++ 支撑│ ├── dynamo/ # Dynamo 与 CPython eval frame 相关 C++ 支撑│ └── inductor/ # Inductor C++ 支撑├── _dynamo/ # Python 字节码捕获和 guard 系统├── fx/ # FX 图 IR 和 symbolic tracing├── _inductor/ # 编译器后端、lowering、调度、代码生成、runtime cache├── _functorch/ # vmap、grad、functional transforms、AOTAutograd├── export/ # 导出 API├── _export/ # 导出内部实现和约束处理├── distributed/ # 分布式训练 Python API├── utils/ # DataLoader、C++ extension、benchmark、checkpoint、pytree├── cuda/ # CUDA Python API├── xpu/ # XPU Python API├── mps/ # MPS Python API├── _subclasses/ # FakeTensor、ProxyTensor、FunctionalTensor、Tensor subclass├── _decomp/ # 算子分解├── _refs/ # 参考实现└── _prims/ # primitive 算子体系 2.3 aten/ 目录
aten/ 是算子和 Tensor C++ API 的核心目录。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 aten/└── src/ATen/ ├── core/ # TensorBase、IValue、operator schema、dispatcher、boxing、TensorOptions ├── native/ # 大量 native 算子实现 │ └── native_functions.yaml # 算子声明中心 ├── ops/ # 生成的算子头文件/包装入口 ├── cpu/ # CPU 相关支持 ├── cuda/ # CUDA 相关支持 ├── cudnn/ # cuDNN 集成 ├── hip/ # HIP/ROCm 后端 ├── xpu/ # XPU 后端 ├── mps/ # MPS 后端 ├── metal/ # Metal 后端 ├── vulkan/ # Vulkan 后端 ├── quantized/ # 量化算子 ├── functorch/ # functorch 与 ATen 交互部分 ├── templates/ # 代码生成模板 └── test/ # ATen C++ 测试 2.4 c10/ 目录
c10/ 是 PyTorch C++ 基础库,许多文件不依赖 Python,也不依赖高层 ATen。
1 2 3 4 5 6 7 8 9 10 11 c10/├── core/ # Device、DispatchKey、TensorImpl、Storage、Allocator、ScalarType、SymInt、Layout├── util/ # Exception、ArrayRef、Optional、SmallVector、intrusive_ptr、typeid、数值工具├── cuda/ # CUDA stream、event、allocator、device guard 等基础封装├── hip/ # HIP/ROCm 基础封装├── xpu/ # XPU 基础封装├── metal/ # Metal 基础封装├── macros/ # 导出宏、编译器宏、平台宏├── mobile/ # 移动端相关基础设施├── test/ # c10 测试└── benchmark/ # c10 基准测试 2.5 torchgen/ 目录
torchgen/ 是理解 PyTorch 源码生成机制的关键。
1 2 3 4 5 6 7 8 9 torchgen/├── api/ # schema 到 C++ API / dispatcher API / native API / structured API 的翻译├── model.py # native function、schema、dispatch key、operator name 等数据模型├── gen.py # ATen 代码生成主入口之一├── dest/ # 生成目标文件逻辑├── decompositions/ # 算子分解相关生成├── aoti/ # AOTInductor 相关生成├── selective_build/ # 选择性构建,用于移动端和裁剪场景└── shape_functions/ # shape function 相关内容 PyTorch 大量 C++ 文件不是手写,而是由 yaml、模板和 torchgen 生成。读源码时需要区分“手写源码”和“生成源码”,否则容易找不到入口。
2.6 tools/ 目录
tools/ 包含开发和生成工具。
1 2 3 4 5 6 7 8 tools/├── autograd/ # 自动求导相关代码生成,包括 derivatives 配置处理├── codegen/ # 构建期代码生成工具或相关脚本├── linter/ # lint 工具├── code_coverage/ # 代码覆盖率工具├── coverage_plugins_package/ # 覆盖率插件包├── stats/ # 统计工具└── experimental/ # 实验性工具 2.7 test/ 目录
test/ 不只是验证目录,也是理解行为语义的重要入口。
1 2 3 4 5 6 7 8 9 10 test/├── test_*.py # 大量 Python API 行为测试├── cpp/ # C++ API、JIT、distributed、profiler、lazy、custom backend 测试├── distributed/ # 分布式测试├── inductor/ # Inductor 测试├── dynamo/ # Dynamo 测试├── functorch/ # functorch / 函数变换测试├── cpp_extensions/ # C++/CUDA 扩展测试├── mobile/ # 移动端测试└── onnx/ # ONNX 导出相关测试 三、用户 API 层
3.1 层功能说明
用户 API 层是 PyTorch 暴露给模型开发者的入口,主要位于 torch/。这一层负责把底层 Tensor、算子、自动求导、优化器、数据加载、分布式和编译能力组织成自然的 Python 编程体验。它本身不承担所有计算,更多承担 API 组织、对象管理、参数校验和用户语义表达,性能敏感路径会继续下沉到 C++ 桥接层、ATen 或编译栈。
torch | torch.Tensortorch.add、torch.save、torch.device | |
torch.nn | ModuleParameter、Linear、Conv2d、functional | |
torch.optim | OptimizerSGD、Adam、AdamW | |
torch.autograd | Functiongrad、no_grad、enable_grad | |
torch.utils | DatasetDataLoader、cpp_extension | |
torch.distributed | DDPFSDP、ProcessGroup、RPC | |
torch.compiletorch.export / torch.fx | compileGraphModule、ExportedProgram |
3.2 torch 顶层 API
torch 顶层模块是最基础的用户入口,提供 Tensor 构造、顶层算子、dtype/device 管理、随机数、序列化和全局配置。用户调用 torch.add、torch.matmul、torch.empty、torch.save 等函数时,Python 层通常只做轻量包装,随后进入 generated binding、ATen dispatcher 或设备相关实现。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 torch/├── __init__.py # 顶层命名空间,组织 torch.* API├── _tensor.py # Python 侧 Tensor 方法补充├── _C/__init__.pyi.in # torch._C 类型存根,描述 C++ 扩展暴露的接口├── serialization.py # torch.save / torch.load├── random.py # 随机数相关 Python API├── _prims/ # primitive 算子└── _refs/ # 参考实现核心对象:├── torch.Tensor # Python 用户看到的 Tensor,底层对应 C++ at::Tensor├── torch.dtype # dtype 描述├── torch.device # 设备描述└── torch.Generator # 随机数生成器 顶层 API 的价值是统一用户心智:不管 Tensor 在 CPU、CUDA、MPS、XPU 还是 Meta 上,用户都先面对同一组 torch.* 函数。后端差异由 dispatcher、dispatch key 和 backend kernel 处理。
3.3 torch.nn 模块体系
torch.nn 负责组织神经网络结构,核心类是 Module。Module 提供参数登记、buffer 登记、子模块管理、hook、state_dict、训练/推理模式切换等能力。Parameter 是特殊 Tensor,被赋值到 Module 属性上时会自动进入参数集合。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 torch/nn/├── modules/module.py # Module 基类:参数、buffer、子模块、hook、state_dict├── parameter.py # Parameter / Buffer├── modules/ # Linear、Conv2d、Norm、Activation、Loss 等模块├── functional.py # 无状态函数式接口├── init.py # 参数初始化└── parallel/ # DataParallel / DistributedDataParallel 相关封装核心类/方法:├── Module # 所有 nn 模块基类├── Parameter # 自动注册为参数的 Tensor 子类├── Buffer # 非参数状态张量├── Sequential / ModuleList / ModuleDict├── Module.__setattr__ # 参数、buffer、子模块自动登记入口├── register_parameter()├── register_buffer()├── state_dict()└── load_state_dict() torch.nn 中的 layer 例如 Linear、Conv2d、LayerNorm 通常负责保存权重和超参数,真正计算会调用 torch.nn.functional 或 ATen 算子。也就是说,nn 是模型结构组织层,不是底层 kernel 实现层。
3.4 torch.optim 优化器体系
torch.optim 以 Optimizer 为基类,负责维护参数组、学习率、优化器状态和 step() 逻辑。优化器不直接理解 Module 的层级结构,而是接收一组参数 Tensor,根据每个参数的 .grad 和内部状态进行原地更新。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 torch/optim/├── optimizer.py # Optimizer 基类、参数组、状态字典├── sgd.py # SGD├── adam.py / adamw.py # Adam / AdamW├── rmsprop.py / lbfgs.py # RMSprop / LBFGS 等优化器├── lr_scheduler.py # 学习率调度└── _functional.py # 函数式参数更新逻辑核心对象/方法:├── Optimizer├── Optimizer.param_groups # 参数组├── Optimizer.state # 参数状态,如 momentum、exp_avg├── zero_grad()├── step()└── state_dict() SGD、Adam、AdamW 等优化器的差异主要体现在状态变量和更新公式上。优化器层依赖 autograd 产生梯度,依赖 Tensor 原地操作更新参数,因此它连接了用户训练循环、autograd 和 ATen Tensor 操作。
3.5 torch.autograd 用户入口
torch.autograd 提供自动求导的用户可见接口,包括 backward、grad、Function、grad mode 和 profiler 等。用户只需要设置 requires_grad=True 并执行前向计算,PyTorch 会在底层构建反向图。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 torch/autograd/├── __init__.py # backward / grad 等用户入口├── function.py # Python 自定义 autograd.Function├── grad_mode.py # no_grad / enable_grad / inference_mode├── profiler.py # profiler 用户接口└── profiler_legacy.pytorch/csrc/autograd/ # C++ autograd engine 和 Python binding核心对象/函数:├── Function├── backward()├── grad()├── no_grad()├── enable_grad()└── inference_mode() Python 层的 torch.autograd.Function 用于自定义求导逻辑,用户实现 forward 和 backward,桥接层再把它连接到 C++ autograd engine。no_grad、enable_grad、inference_mode 这类上下文则通过线程局部状态影响 dispatcher 和 autograd 行为。
3.6 torch.utils 工具模块
torch.utils 提供训练工程中常用的辅助能力。Dataset 和 DataLoader 负责数据集抽象、批处理、多进程加载、pin memory 等;cpp_extension 支持构建 C++/CUDA 扩展;checkpoint、benchmark、pytree 等工具分别服务内存优化、性能测量和结构化参数处理。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 torch/utils/├── data/│ ├── dataset.py # Dataset / IterableDataset│ └── dataloader.py # DataLoader,多进程加载、pin memory、collate├── cpp_extension.py # C++/CUDA 扩展构建├── checkpoint.py # activation checkpoint├── benchmark/ # benchmark 工具└── _pytree.py # 嵌套结构 flatten / unflatten核心类:├── Dataset├── IterableDataset├── DataLoader├── Sampler└── BatchSampler 这一层的特点是贴近用户工程实践。它不定义 PyTorch 的核心 Tensor 语义,但大量训练程序离不开这些工具。
3.7 torch.distributed 分布式入口
torch.distributed 是分布式训练的用户入口,覆盖通信原语、进程组、DDP、FSDP、RPC、elastic、checkpoint 和 DTensor 等能力。它把底层通信后端包装成用户可用的训练抽象。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 torch/distributed/├── distributed_c10d.py # ProcessGroup 和集合通信 Python API├── device_mesh.py # DeviceMesh├── fsdp/│ ├── fully_sharded_data_parallel.py # FSDP 主类│ └── _flat_param.py # 参数 flatten / shard├── rpc/ # RPC├── elastic/ # 弹性训练├── checkpoint/ # 分布式 checkpoint└── tensor/ # DTensor 相关能力torch/csrc/distributed/ # 底层 C++ 通信实现核心对象:├── ProcessGroup├── DistributedDataParallel├── FullyShardedDataParallel├── DeviceMesh├── RPC agent└── checkpoint planner DDP 主要围绕梯度同步工作,FSDP 主要围绕参数、梯度和优化器状态分片工作。它们都与 Module、autograd hook、optimizer 和通信后端深度交互。
3.8 torch.compile、torch.export 与 torch.fx
torch.compile 提供动态图编译入口,目标是在尽量不改变用户代码的情况下捕获可优化区域。torch.export 更偏稳定导出语义,适合保存和下游消费。torch.fx 提供 Python 级图表示和图变换能力。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 torch/_dynamo/├── eval_frame.py # torch.compile 入口之一├── convert_frame.py # frame 转换├── symbolic_convert.py # 字节码符号执行└── guards.py # guard 生成和检查torch/fx/├── graph.py # Graph IR├── node.py # Node└── graph_module.py # GraphModuletorch/export/ and torch/_export/ # 导出 API 和内部实现torch/_inductor/├── compile_fx.py # FX 图编译入口├── graph.py # GraphLowering├── ir.py # Inductor IR├── scheduler.py # 调度和融合└── codegen/ # 后端代码生成核心对象:├── Graph / Node / GraphModule├── ExportedProgram├── Dynamo guard├── Inductor GraphLowering└── Inductor IR node 这些模块让用户 API 层不只支持 eager execution,也支持图捕获、编译优化和部署导出。底层会进一步进入 Dynamo、FX、AOTAutograd、Inductor 和后端代码生成。
四、C++ 桥接层
4.1 层功能说明
C++ 桥接层主要位于 torch/csrc/,编译后以 torch._C 的形式被 Python 侧导入。它是 Python API 和 C++ 核心之间的边界层,负责对象表示、参数转换、异常转换、GIL、引用计数、模块初始化和 generated binding 接入。
torch._C | ||
torch/csrc | ||
torch/csrc/autograd | THPVariable | |
torch/csrc/api | Module、optimizer 等 | |
torch/csrc/distributed | ||
torch/csrc/dynamoinductor |
4.2 torch._C 扩展模块
torch._C 是 Python 侧看到的 C extension module。它不是普通 Python 文件,而是编译产物,集中暴露 dispatcher、Tensor 类型、autograd、JIT、profiler、distributed 和设备管理等 C++ 能力。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 torch/_C/└── __init__.pyi.in # torch._C Python 类型存根torch/csrc/├── Module.cpp # 主扩展模块初始化,大量 C++ 能力注册├── Exceptions.cpp # C++ 异常到 Python 异常转换├── Device.cpp # Device Python binding├── Dtype.cpp # dtype Python binding├── Layout.cpp # layout Python binding└── Stream.cpp # stream Python binding核心对象:├── THPModule 初始化逻辑├── THPVariableClass├── torch._C._VariableFunctions├── Python type object└── pybind 注册对象 很多 Python API 表面上位于 torch/,实际执行会进入 torch._C。例如 Tensor 方法、部分设备 API、Dynamo eval frame hook、C++ profiler 和底层 distributed binding 都依赖这一模块。
4.3 torch/csrc 绑定源码
torch/csrc 是桥接层源码主体。它负责 Python 模块初始化、Python 类型定义、C++ 类型转换、异常转换和生命周期管理。这里的代码需要同时理解 CPython C API、C++ Tensor 语义和 PyTorch 自身的 intrusive pointer 机制。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 torch/csrc/├── python_headers.h # Python 头文件包含和隔离├── utils/│ ├── python_arg_parser.h # Python 参数解析│ └── python_arg_parser.cpp├── utils.cpp / utils.h # Tensor 包装、类型检查、通用转换工具├── tensor/python_tensor.cpp # Tensor Python 类型相关逻辑├── Storage.cpp # Storage Python binding└── StorageMethods.cpp # Storage 方法绑定核心函数:├── THPVariable_Wrap()├── THPVariable_Unpack()├── PythonArgParser└── THP*initModule() 典型调用路径是:
1 2 3 4 5 Python torch.Tensor method -> torch._C 暴露的 binding -> 参数解析和类型转换 -> C++ at::Tensor / dispatcher / autograd -> C++ 返回值包装为 Python 对象 4.4 torch/csrc/autograd 绑定
torch/csrc/autograd 连接 Python Tensor 和 C++ autograd engine。THPVariable 等结构负责让 Python 侧 Tensor 持有 C++ at::Tensor,同时暴露 .grad、.grad_fn、requires_grad 等求导属性。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 torch/csrc/autograd/├── python_variable.h # THPVariable / Wrap / Unpack├── python_variable.cpp # Python Tensor 类型行为├── python_variable_indexing.cpp # Tensor 索引├── python_function.cpp # Python autograd.Function 接入├── python_engine.cpp # C++ Engine 的 Python binding├── engine.h / engine.cpp # backward 执行引擎├── node.h # 反向图 Node├── function.h # autograd function 基础结构├── variable.h # Variable = at::Tensor 相关接口└── autograd_meta.cpp # AutogradMeta核心类:├── THPVariable # Python Tensor 壳├── torch::autograd::Engine├── Node / Edge├── AutogradMeta├── SavedVariable└── GraphTask Python 自定义 autograd.Function 也需要通过这一层接入 C++ autograd graph。用户在 Python 中定义 backward,底层仍然需要以 Node/Edge 的形式进入 autograd engine 调度。
4.5 torch/csrc/api C++ frontend
torch/csrc/api 提供 C++ frontend,使用户可以在 C++ 中使用类似 Python nn.Module、optimizer 和 Tensor API 的模型构建方式。它不是 Python API 的简单翻译,而是面向 C++ 用户的高层封装。
1 2 3 4 5 6 7 8 9 10 11 12 13 torch/csrc/api/├── include/torch/ # C++ frontend 公开头文件│ ├── torch.h # 常用总入口│ ├── nn/ # C++ 版模块体系│ ├── optim/ # C++ 版优化器│ └── data/ # C++ 版数据 API└── src/ # C++ frontend 实现核心类:├── torch::nn::Module├── torch::nn::Linear├── torch::optim::Optimizer└── torch::data::Dataset C++ frontend 依赖 ATen Tensor 和 dispatcher,但提供更接近深度学习框架的对象模型,例如 module、sequential、optimizer 和 serialization。
4.6 分布式和编译栈 C++ 支撑
torch/csrc/distributed 提供分布式 Python API 背后的 C++ 支撑,包括通信后端绑定、process group 对象和部分 runtime 能力。torch/csrc/dynamo 和 torch/csrc/inductor 则支撑现代编译栈,例如 CPython eval frame hook、编译 runtime 和低开销执行路径。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 torch/csrc/distributed/├── c10d/ # C++ 通信抽象和后端绑定│ ├── nccl/ # NCCL 后端│ ├── gloo/ # Gloo 后端│ ├── ucc/ # UCC 后端│ └── store/ # Store / rendezvous 支撑├── rpc/ # RPC C++ 支撑└── autograd/ # 分布式 autogradtorch/csrc/dynamo/ # Dynamo C++ 支撑,如 guards、compiled autogradtorch/csrc/inductor/ # Inductor runtime / launcher 支撑核心对象:├── ProcessGroup├── Store├── RPC agent├── distributed autograd context├── Dynamo guard manager└── Inductor launcher 这些模块说明桥接层不只是 Tensor binding,而是所有需要从 Python 高层进入 C++ runtime 的系统边界。
夜雨聆风