at-dispatch-v2
将 PyTorch ATen C++ 代码从旧版 AT_DISPATCH 宏转换为现代 AT_DISPATCH_V2 API 的自动化工具。
简介
此技能为将 PyTorch ATen C++ 算子实现从旧版分派宏迁移至现代化 AT_DISPATCH_V2 架构提供了系统化框架。专为 PyTorch 内核开发人员与系统工程师设计,通过 Dispatch_v2.h 头文件中定义的标准化转换规则、参数重排序逻辑与宏展开指引,简化移植流程。目标针对 aten/src/ATen/native 目录下的文件,确保旧有代码库能持续与演进中的 PyTorch 分派需求保持兼容。
-
自动化处理旧版模式(如 AT_DISPATCH_ALL_TYPES_AND* 与 AT_DISPATCH_FLOATING_TYPES_AND*)至统一 AT_DISPATCH_V2 格式的迁移。
-
提供类型组展开的清晰对应表,确保正确使用 AT_EXPAND() 处理如 AT_ALL_TYPES、AT_INTEGRAL_TYPES 与 AT_FLOATING_TYPES 等基础类型集。
-
使用 AT_WRAP 封装 Lambda 表达式,以处理包含内部逗号或多行逻辑的复杂 C++ 代码块。
-
将 #include <ATen/Dispatch_v2.h> 与现有分派依赖项整合,标准化头文件管理。
-
支持对类型列表的细粒度控制,允许开发人员在应用广泛类型组宏后,附加如 kBFloat16、kHalf 或 kBool 等特定类型。
-
在内核实现中识别旧版分派宏,包含 ATen 原生算子与 CUDA 内核。
-
通过重排序参数(将 scalar_type 与算子名称置于调用开头)来系统化替换宏。
-
使用 AT_WRAP 封装 Lambda 本体,以维持预处理期间正确的 C++ 宏解析。
-
针对涉及的特定类型类别应用正确的 AT_EXPAND 宏,例如 AT_ALL_TYPES_AND_COMPLEX 或 AT_BAREBONES_UNSIGNED_TYPES。
-
根据必要的签名验证最终转换结果:AT_DISPATCH_V2(scalar_type, name, lambda, type_groups, individual_types)。
-
确保个别类型遵循 API 规范,置于类型组之后且不被 AT_EXPAND 包裹。
仓库统计
- Star 数
- 181
- Fork 数
- 24
- Open Issue 数
- 4
- 主要语言
- Python
- 默认分支
- main
- 同步状态
- 空闲
- 最近同步时间
- 2026年4月29日 14:12