一、ATen 是什么在 PyTorch 中扮演什么角色当你在 Python 中写下torch.add(a, b)或a.cos()时最终执行计算的不是 Python而是一套 C 库——这就是ATenA Tensor Library。ATen 在 PyTorch 整体架构中的位置如下┌─────────────────────────────────────────────┐ │ Python API 层torch/, torch.nn, … │ ├─────────────────────────────────────────────┤ │ Python-C 绑定层torch/csrc/ │ ├─────────────────────────────────────────────┤ │ ★ ATenaten/src/ATen/ │ ← 本文分析目标 │ ├─ 算子声明native_functions.yaml │ │ ├─ 算子实现native/ │ │ ├─ 核心抽象Tensor, TensorIterator, … │ │ ├─ 设备后端cpu/, cuda/, mps/, … │ │ └─ 代码生成模板templates/ │ ├─────────────────────────────────────────────┤ │ c10 核心库c10/ │ │ ├─ Storage, TensorImpl, DispatchKey │ │ └─ Allocator, Device, ScalarType │ └─────────────────────────────────────────────┘一句话总结c10 提供最底层的张量存储和分发基础设施ATen 在其之上实现了1000 个算子的具体逻辑而 Python 层只是这些 C 实现的薄封装。二、aten/ 目录全景aten/ ├── CMakeLists.txt # 构建入口 ├── tools/ # 代码生成辅助工具 └── src/ ├── README.md # 历史说明TH/THC 的演变 ├── THC/ # 遗留的 CUDA 辅助代码 └── ATen/ # ★ 核心所在 ├── native/ # ★★★ 算子实现最大的子目录 ├── core/ # 核心抽象移动端友好 ├── cpu/ # CPU 向量化支持 ├── cuda/ # CUDA 运行时支持 ├── detail/ # 设备钩子接口 ├── functorch/ # vmap/grad 批处理规则 ├── ops/ # 代码生成的算子绑定 ├── templates/ # 代码生成模板 ├── quantized/ # 量化相关 ├── cudnn/ # cuDNN 封装 ├── miopen/ # MIOpen (AMD) 封装 ├── mps/ # Apple Metal 封装 ├── mkl/ # Intel MKL 封装 ├── vulkan/ # Vulkan 移动端后端 ├── metal/ # Metal 移动端后端 ├── xpu/ # Intel GPU 后端 ├── test/ # C 单元测试 ├── benchmarks/ # 性能基准 └── *.h / *.cpp # 顶层核心头文件与实现三、核心中的核心native_functions.yaml如果说 ATen 有一个上帝文件那就是aten/src/ATen/native/native_functions.yaml。PyTorch 的每一个原生算子都在这个 YAML 文件中声明。3.1 格式解析一个典型的算子声明长这样-func:add.Tensor(Tensor self,Tensor other,*,Scalar alpha1)-Tensorvariants:function,methoddispatch:CPU:add_cpuCUDA:add_cudaSparseCPU:add_sparseSparseCUDA:add_sparse_cudaMPS:add_mps它声明了以下信息字段含义func算子签名名称、参数类型、默认值、返回类型variants生成at::add()函数和Tensor::add()方法dispatch不同设备/后端对应的 C 实现函数名3.2 代码生成流程当你修改了native_functions.yamlPyTorch 的构建系统torchgen/会native_functions.yaml │ ▼ torchgen/ 代码生成器 │ ├─→ aten/src/ATen/ops/ # 生成的 C 绑定 │ ├─ Functions.h # at::add() 声明 │ ├─ TensorBody.h # Tensor::add() 方法 │ ├─ Operators.cpp # 分发桩代码 │ └─ RegisterSchema.cpp # Schema 注册 │ ├─→ torch/csrc/autograd/ # 生成的 Python 绑定 │ └─→ build/aten/src/ATen/ # 构建产物这意味着开发者增加新算子的第一步不是写 C 代码而是先在 YAML 中声明签名。3.3 三种 Dispatch 关键字dispatch字段的值决定了算子的实现策略这是 ATen 设计哲学的精髓Dispatch 关键字语义适用场景CompositeImplicitAutograd由其他算子组合而成自动推导梯度可分解的高层算子如addcmulCompositeExplicitAutograd由其他算子组合但需手写梯度在derivatives.yaml需要自定义反向传播的委托函数CPU/CUDA/MPS/ …设备特化实现需要针对特定硬件优化的算子设计哲学如果一个算子能用已有算子组合实现就用CompositeImplicitAutograd让 autograd 自动推导梯度。只有性能或数值精度不满足时才下降到设备特化实现。四、native/ 目录算子实现的主战场aten/src/ATen/native/是整个 ATen 中最庞大的目录包含数百个算子的 C 实现。4.1 顶层分类native/ 下的文件按算子类别组织文件/文件组涵盖算子UnaryOps.cppabs, cos, sin, exp, log, sqrt, …BinaryOps.cppadd, sub, mul, div, …ReduceOps.cppsum, mean, prod, max, min, …TensorShape.cppreshape, view, permute, transpose, cat, stack, …TensorFactories.cppzeros, ones, empty, rand, randn, arange, …TensorCompare.cppeq, ne, lt, gt, where, …TensorAdvancedIndexing.cppindex, index_put, gather, scatter, …LinearAlgebra.cppmm, bmm, addmm, svd, qr, …Convolution.cppconv1d, conv2d, conv3d 路由逻辑Normalization.cppbatch_norm, layer_norm, group_norm, …Activation.cpprelu, gelu, silu, hardswish, …Loss.cpp/LossCTC.cppcross_entropy, nll_loss, ctc_loss, …SoftMax.cppsoftmax, log_softmaxPooling.cppmax_pool, avg_pool, adaptive_pool, …Embedding.cppembedding, embedding_bagDistributions.cppnormal_, uniform_, bernoulli_, …RNN.cpplstm, gru 计算内核Copy.cpptensor 间的拷贝逻辑Resize.cpp张量尺寸调整4.2 设备后端子目录native/ 下有多个设备特化的子目录native/ ├── cpu/ # CPU 向量化内核AVX2/AVX512/NEON/SVE 多版本编译 ├── cuda/ # CUDA 内核.cu 文件约 300 个 ├── mps/ # Apple Metal 内核.mm 文件 ├── mkldnn/ # Intel oneDNN 加速内核 ├── cudnn/ # cuDNN 加速卷积、BN、RNN ├── miopen/ # AMD MIOpen 加速 ├── xnnpack/ # ARM CPU XNNPACK 加速移动端 ├── sparse/ # 稀疏张量操作 ├── nested/ # 嵌套张量操作 ├── quantized/ # 量化算子 ├── transformers/ # Attention 实现FlashAttention/高效注意力 ├── vulkan/ # Vulkan 移动端 GPU └── kleidiai/ # ARM KleidiAI 加速4.3 CPU 内核的多版本编译机制native/cpu/目录下的文件有一个非常特殊的编译策略每个.cpp文件会被编译多次每次使用不同的指令集标志如-mavx2、-mavx512f。BinaryOpsKernel.cpp │ ├─ 编译为 BinaryOpsKernel_DEFAULT.o (基础 SSE) ├─ 编译为 BinaryOpsKernel_AVX2.o (-mavx2 -mfma) └─ 编译为 BinaryOpsKernel_AVX512.o (-mavx512f -mavx512bw)运行时通过DispatchStub机制根据 CPU 特性自动选择最优版本// 在 native/BinaryOps.h 中声明DECLARE_DISPATCH(add_fn,add_stub);// 在 native/BinaryOps.cpp 中定义DEFINE_DISPATCH(add_stub);// 在 native/cpu/BinaryOpsKernel.cpp 中注册REGISTER_DISPATCH(add_stub,add_kernel);关键约束来自 READMEnative/cpu/中的所有实现必须放在匿名命名空间中否则不同指令集版本的符号会在链接时冲突。只有需要利用 SIMD 指令的计算密集型内核才应放在native/cpu/。4.4 CUDA 内核的组织native/cuda/包含约300 个.cu文件是 ATen 中第二大的代码区。组织特点native/cuda/ ├── UnaryOpsKernel.cu # 一元操作 ├── BinaryMulKernel.cu # 二元操作按类型拆分 ├── BinaryDivTrueKernel.cu ├── ReduceSumProdKernel.cu # 归约操作 ├── SortStable.cu # 排序 ├── Indexing.cu # 索引 ├── Embedding.cu # 嵌入 ├── SpectralOps.cu # FFT ├── fused_adam_impl.cu # 融合优化器 ├── fused_adamw_impl.cu ├── int4mm.cu # 量化矩阵乘 ├── int8mm.cu ├── RowwiseScaledMM.cu # 逐行缩放 GEMM ├── ScaledGroupMM.cu # 缩放分组 GEMM ├── cutlass_extensions/ # CUTLASS 扩展模板 ├── linalg/ # CUDA 线性代数 ├── tunable/ # 自动调优框架 └── jit_utils.cpp # CUDA JIT 编译工具每个 CUDA 内核通常使用以下模式之一TensorIterator gpu_kernel用于逐元素操作自动处理广播和类型提升。手写 CUDA kernel用于无法用 TensorIterator 表达的复杂操作如 sort、scatter。调用外部库cuBLASGEMM、cuDNN卷积/BN、cuSOLVERSVD/QR等。五、核心基础设施5.1 TensorIterator — 逐元素操作的瑞士军刀TensorIterator声明在ATen/TensorIterator.h是 ATen 中使用最广泛的基础设施灵感来自 NumPy 的NpyIter。它自动处理广播broadcasting类型提升type promotion内存连续性优化多线程并行使用模式// 配置autoiterTensorIteratorConfig().add_output(output).add_input(input_a).add_input(input_b).build();// CPU 内核cpu_kernel(iter,[](floata,floatb)-float{returnab;});// CUDA 内核gpu_kernel(iter,[]GPU_LAMBDA(floata,floatb)-float{returnab;});为什么重要绝大多数逐元素和归约算子add、mul、cos、sum……都通过 TensorIterator 实现。它将算子业务逻辑和内存遍历/并行策略彻底解耦。5.2 AT_DISPATCH 宏 — 类型分发PyTorch 支持 20 种数据类型float16/32/64, int8/16/32/64, bfloat16, complex64/128, …。C 模板不能直接用 runtime 的ScalarType实例化因此 ATen 使用AT_DISPATCH_*宏族做类型分发AT_DISPATCH_ALL_TYPES_AND2(kHalf,kBFloat16,input.scalar_type(),my_kernel,[](){// scalar_t 在这里被绑定为具体类型cpu_kernel(iter,[](scalar_t a)-scalar_t{returna*a;});});常用宏包括AT_DISPATCH_ALL_TYPES— float/double/int8/16/32/64AT_DISPATCH_FLOATING_TYPES— float/doubleAT_DISPATCH_ALL_TYPES_AND(kHalf, ...)— 增加 float16AT_DISPATCH_ALL_TYPES_AND_COMPLEX— 包含 complex64/1285.3 DispatchStub — CPU 指令集分发如前所述DispatchStub 是 CPU 多版本编译的分发机制DispatchStubfn_type, impl │ ├─ DEFAULT → 基础实现 ├─ AVX2 → AVX2 优化实现 └─ AVX512 → AVX512 优化实现支持的设备类型不限于 CPU还包括 CUDA、HIP、MPS、XPU 等。运行时根据cpuinfo报告的 CPU 特性自动选择最优路径。5.4 向量化库 cpu/vec/ATen/cpu/vec/提供了跨平台的 SIMD 向量化抽象cpu/vec/ ├── vec.h # 统一入口 ├── vec_base.h # 标量回退实现 ├── vec256/ # 256-bit 向量AVX2/NEON │ ├── vec256_float.h │ ├── vec256_int.h │ └── ... ├── vec512/ # 512-bit 向量AVX512 ├── vec128/ # 128-bit 向量SSE/NEON ├── sve/ # ARM SVE 支持 ├── functional.h # map/reduce 等函数式接口 ├── vec_convert.h # 类型转换 └── vec_mask.h # 掩码操作VectorizedT模板类封装了不同 SIMD 指令集的寄存器操作让算子开发者可以写出一次编写、多指令集运行的代码#includeATen/cpu/vec/vec.husingnamespaceat::vec;voidmy_fast_cos(float*out,constfloat*in,int64_tn){int64_td0;for(;dn-(n%Vectorizedfloat::size());dVectorizedfloat::size()){autoxVectorizedfloat::loadu(ind);x.cos().store(outd);}for(;dn;d){out[d]std::cos(in[d]);}}六、core/ 目录移动端友好的核心抽象ATen/core/包含 ATen 的最小内核子集经过精心设计以控制二进制体积适合移动端部署。core/ ├── Tensor.h / Tensor.cpp # Tensor 核心定义 ├── TensorBase.h # Tensor 基类无 autograd ├── Generator.h # 随机数生成器抽象 ├── Dimname.h # 命名维度 ├── ivalue.h # IValueJIT 系统的通用值类型 ├── function_schema.h # 算子签名描述 ├── op_registration/ # 算子注册基础设施 ├── dispatch/ # 分发器核心 ├── boxing/ # 装箱/拆箱泛型调用 ├── class_type.h / custom_class.h # TorchScript 自定义类 ├── List.h / Dict.h # 泛型容器 └── PythonFallbackKernel.cpp # Python 回调内核关键设计约束来自 README二进制大小是此目录中文件的重要约束。这意味着 core/ 中的代码不能随意添加重型依赖。七、CUDA 运行时支持cuda/ 目录ATen/cuda/不是算子实现那在native/cuda/而是CUDA 运行时基础设施ATen/cuda/ ├── CUDAContext.h/.cpp # CUDA 上下文管理当前设备、stream ├── CUDABlas.h/.cpp # cuBLAS 封装GEMM、BatchedGEMM ├── CUDASparseBlas.h/.cpp # cuSPARSE 封装 ├── CUDAGraph.h/.cpp/.cu # CUDA Graph 录制与回放 ├── CUDAGeneratorImpl.h/.cpp # CUDA 随机数生成器 ├── CUDAEvent.h # CUDA Event 封装 ├── CachingHostAllocator.h # 页锁定内存缓存分配器 ├── MemPool.h/.cpp # 内存池管理 ├── PeerToPeerAccess.h/.cpp # 多 GPU P2P 访问 ├── Atomic.cuh # 原子操作 ├── cub.cuh # CUB 库封装排序、扫描 ├── PhiloxCudaState.h # Philox 随机数状态 ├── tunable/ # 自动调优框架 ├── detail/ # 内部实现细节 └── nvrtc_stub/ # NVRTC JIT 编译桩CUDABlas是一个重要的性能热点封装它包装了 cuBLAS 的 GEMM 调用并在内部处理Tensor Core 的自动使用workspace 管理混合精度支持FP16/BF16 GEMM FP32 累加八、detail/ 目录设备钩子架构ATen/detail/定义了 PyTorch 的设备钩子接口Hooks Interface这是支持多后端可扩展性的关键detail/ ├── CUDAHooksInterface.h # CUDA 后端钩子 ├── HIPHooksInterface.h # ROCm/HIP 后端钩子 ├── MPSHooksInterface.h # Apple Metal 后端钩子 ├── XPUHooksInterface.h # Intel GPU 后端钩子 ├── MTIAHooksInterface.h # Meta AI 加速器钩子 ├── HPUHooksInterface.h # Intel Gaudi 钩子 ├── PrivateUse1HooksInterface.h# 自定义设备钩子 ├── AcceleratorHooksInterface.h# 加速器通用接口 ├── XLAHooksInterface.h # XLA 后端钩子 └── CPUGuardImpl.cpp # CPU 设备守卫设计模式每个XxxHooksInterface.h定义了一组虚函数如getDeviceCount()、getStream()对应后端在自己的库中提供具体实现通过动态链接注册到 ATen。这使得核心 ATen 库不需要链接任何设备特定库即可编译设备支持是插件式的。九、代码生成系统templates/ 与 ops/9.1 templates/ — 代码模板templates/ ├── TensorBody.h # Tensor 类方法声明模板 ├── TensorMethods.cpp # Tensor 类方法实现模板 ├── Functions.h # at:: 命名空间函数声明 ├── Functions.cpp # 函数实现 ├── Operators.h/.cpp # 分发桩 ├── RegisterSchema.cpp # Schema 注册 ├── RegisterDispatchKey.cpp # 各 DispatchKey 注册 ├── NativeFunctions.h # Native 函数声明 ├── DispatchKeyFunctions.h # 按 key 的函数声明 └── RegisterFunctionalization.cpp # 功能化变换注册这些模板通过torchgen/代码生成器填充产出到build/aten/src/ATen/目录。9.2 ops/ — 生成产物仓库中只有少量文件ops/ ├── tensor.h # 简便头文件 └── from_blob.h # from_blob 工厂函数构建后ops/目录会包含每个算子的独立头文件如ops/add.h是面向用户的 C API 入口。十、专项领域模块10.1 native/transformers/ — 注意力机制native/transformers/ ├── attention.h/.cpp # 统一注意力入口 ├── sdp_utils_cpp.h/.cpp # Scaled Dot Product Attention 工具 ├── transformer.cpp # 整体 Transformer 组件 ├── cuda/ # FlashAttention / Memory-Efficient Attention ├── hip/ # AMD GPU 注意力实现 └── xpu/ # Intel GPU 注意力实现这是近年来 PyTorch 性能提升最显著的模块之一封装了 FlashAttention v1/v2、Memory-Efficient Attention 等高效实现。10.2 native/sparse/ — 稀疏张量native/sparse/ ├── SparseTensor.cpp # COO 格式稀疏张量操作 ├── SparseTensorMath.cpp # 稀疏数学运算 ├── SparseCsrTensor.cpp # CSR/CSC/BSR/BSC 格式 ├── SparseFactories.cpp # 稀疏张量工厂函数 ├── SparseBlas.cpp # 稀疏 BLAS ├── cuda/ # CUDA 稀疏实现 ├── mps/ # Metal 稀疏实现 └── eigen/ # Eigen 库方案10.3 native/nested/ — 嵌套张量native/nested/ ├── NestedTensorUtils.h/.cpp # 核心工具 ├── NestedTensorMath.cpp # 数学操作 ├── NestedTensorMatmul.cpp # 矩阵乘 ├── NestedTensorTransformerFunctions.cpp # Transformer 适配 ├── NestedTensorFactories.cpp # 工厂函数 └── cuda/ # CUDA 实现嵌套张量Nested Tensor是 PyTorch 对变长序列批处理的原生支持避免了 padding 带来的计算浪费。10.4 native/quantized/ — 量化量化模块支持 INT8/INT4 推理包括仿射量化器AffineQuantizer伪量化FakeQuantAffine用于量化感知训练量化卷积/线性层打包通过 FBGEMM/QNNPACKCUDA 量化内核int4mm, int8mm10.5 functorch/ — 函数变换的批处理规则ATen/functorch/ ├── BatchedTensorImpl.h/.cpp # 批处理张量实现 ├── DynamicLayer.h/.cpp # 变换层栈管理 ├── Interpreter.h/.cpp # 变换解释器 ├── BatchRules*.cpp # ★ 各算子的 vmap 批处理规则约 20 个文件 │ ├── BatchRulesBinaryOps.cpp │ ├── BatchRulesActivation.cpp │ ├── BatchRulesLinearAlgebra.cpp │ └── ... ├── TensorWrapper.h/.cpp # Tensor 包装器 └── VmapInterpreter.h/.cpp # vmap 解释器当用户调用torch.vmap(fn)时functorch 需要知道每个算子在增加一个批维度后应该如何行为——这就是BatchRules*.cpp中定义的规则。十一、数据流全景一次 torch.add() 的完整旅程ctorch.add(a,b)在 C 侧的完整调用路径1. at::add(a, b) // 生成的函数ops/Functions.h │ ▼ 2. Dispatcher::call(aten::add.Tensor) // 分发器查找 │ ├─ 检查 DispatchKey 优先级 │ Autograd FuncTorch Backend │ ▼ 3. AutogradCPU::add() // autograd 包装 │ 记录 grad_fnAddBackward到计算图 │ ▼ 4. CPU::add() (或 CUDA::add()) // 设备特化实现 │ ▼ 5. add_stub(kCPU, ...) // DispatchStub 选择指令集 │ ├─ 如果 CPU 支持 AVX512 → add_kernelAVX512 ├─ 如果 CPU 支持 AVX2 → add_kernelAVX2 └─ 否则 → add_kernelDEFAULT │ ▼ 6. TensorIterator 构建 // 自动广播 类型提升 │ ▼ 7. cpu_kernel(iter, [](scalar_t a, scalar_t b) { return a b; // 最终的计算逻辑 }); │ ▼ 8. 内部并行化OpenMP / 线程池 Vectorizedfloat::loadu → vadd → storeu关键洞察用户写的一行 Python 代码在 C 侧经过了分发器路由 → autograd 记录 → 设备/指令集选择 → 迭代器构建 → 向量化执行至少 6 层抽象。十二、文件规模与热点分布基于目录遍历的粗略统计子目录文件数约核心职责native/顶层~200 个 .cpp/.h算子公共实现native/cuda/~300 个 .cu/.cppCUDA 内核native/cpu/~90 个 .cppCPU 向量化内核native/sparse/~25 个稀疏张量native/quantized/~20 个量化native/nested/~15 个嵌套张量native/transformers/~10 个注意力机制core/~80 个核心抽象cuda/~60 个CUDA 运行时functorch/~40 个vmap 批处理规则顶层ATen/*.h/*.cpp~100 个公共 API 与基础设施热点改动最频繁的区域native/cuda/— 性能优化的永恒战场native/transformers/— LLM 时代的核心模块native_functions.yaml— 每个新算子必经之路十三、目录速查表路径一句话定位native_functions.yaml上帝文件所有算子的声明中心native/*.cpp算子的通用 / CompositeImplicitAutograd 实现native/cpu/*.cppCPU 多指令集内核AVX2/512/NEONnative/cuda/*.cuCUDA GPU 内核native/mps/Apple Metal 内核native/mkldnn/Intel oneDNN 加速native/cudnn/cuDNN 封装卷积/BN/RNNnative/sparse/稀疏张量操作native/nested/嵌套张量变长批处理native/transformers/FlashAttention 等注意力实现native/quantized/INT8/INT4 量化算子TensorIterator.h逐元素操作的统一迭代框架Dispatch.hAT_DISPATCH 类型分发宏native/DispatchStub.hCPU 指令集运行时分发cpu/vec/SIMD 向量化抽象库core/移动端友好的最小内核cuda/CUDA 运行时cuBLAS/Context/Graphdetail/设备钩子接口插件式后端functorch/vmap/grad 批处理规则templates/代码生成模板ops/生成的 C 算子绑定十四、阅读建议如果你想添加新算子先读native/README.md本文第四节的源文件在native_functions.yaml中添加声明在native/下编写 CompositeImplicitAutograd 实现如需 CUDA 特化在native/cuda/添加.cu文件在tools/autograd/derivatives.yaml添加梯度公式如果你想理解性能瓶颈TensorIterator.h— 理解广播和并行策略cpu/vec/— 理解 SIMD 向量化cuda/CUDABlas.cpp— 理解 GEMM 调用链native/cuda/Reduce.cuh— 理解 GPU 归约模式如果你想支持新硬件后端detail/*HooksInterface.h— 实现设备钩子native/DispatchStub.h— 注册新设备的 stubnative_functions.yaml— 添加新的 dispatch key总结ATen 是 PyTorch 的C 算子引擎其架构可以用三层来理解声明层native_functions.yaml统一声明所有算子的签名和分发策略。基础设施层TensorIterator迭代、DispatchStubCPU 指令集分发、AT_DISPATCH类型分发、cpu/vecSIMD 向量化构成了算子实现的四大支柱。实现层native/下按算子类别和设备后端组织的数百个 C/CUDA 文件。理解 ATen你就掌握了 PyTorch 性能优化的底牌——因为不管编译器栈多么花哨Dynamo、Inductor、Triton最终的单算子执行路径仍然落在 ATen 的 C 内核上。