add-uint-support
更新 PyTorch AT_DISPATCH_V2 宏,为算子和内核启用 uint16、uint32 及 uint64 支持。
简介
此技能旨在自动化将 PyTorch 算子类型覆盖范围扩展至包含无符号整数类型 (uint16、uint32 与 uint64) 的过程。它专为需要更新 C++ 派发逻辑以确保内核能够处理无符号输入的 PyTorch 开发人员与库维护者所设计。通过管理 AT_DISPATCH_V2 宏的转换与扩展,本技能确保了 CPU 与 CUDA 实现之间的一致性类型安全与内核执行性能。
此技能会识别当前的派发结构,判断是否有必要转换为 V2 宏格式,并应用适当的类型组扩展,例如 AT_BAREBONES_UNSIGNED_TYPES 或 AT_INTEGRAL_TYPES_V2。它处理了诸如从旧版 AT_DISPATCH 升级、将无符号支持整合至浮点数类型旁,以及处理单个源文件中多个派发站点等常见情境。这种操作支持能减少手动编写代码的工作量、降低类型不匹配的错误,并确保代码符合 PyTorch 不断演进的派发需求。
-
启用 PyTorch 算子与内核实现对 uint16、uint32 与 uint64 的支持。
-
标准化 AT_DISPATCH_V2 宏的使用,以提升类型派发效率与可维护性。
-
自动将 AT_BAREBONES_UNSIGNED_TYPES 与 AT_INTEGRAL_TYPES_V2 注入算子派发列表中。
-
简化从旧版派发宏迁移至建议的 V2 架构的路径。
-
支持多站点更新,确保单个文件内所有 CPU 与 CUDA 派发定义具备一致的覆盖范围。
-
在应用之前,请先确认现有的代码是否已使用 AT_DISPATCH_V2;若否,请先使用专用的迁移逻辑来升级宏格式。
-
在适用的情况下优先采用方法二 (替换为 AT_INTEGRAL_TYPES_V2),这提供了一种更简洁且包含无符号整数的完整覆盖方式。
-
确保所有派发站点 (如 CPU、CUDA 与内部的内核实现) 皆获得统一更新,以防止运行时的类型错误。
-
所有类型组皆应使用 AT_EXPAND(),以确保编译期间能正确执行宏展开。
-
本技能假设开发者可访问标准 ATen 派发头文件,并与 PyTorch 目前基于宏的类型派发系统兼容。
仓库统计
- Star 数
- 99,495
- Fork 数
- 27,610
- Open Issue 数
- 18,545
- 主要语言
- Python
- 默认分支
- main
- 同步状态
- 空闲
- 最近同步时间
- 2026年4月28日 12:39