用途:把 PyTorch ATen C++ 代码中的旧式 AT_DISPATCH 宏转换为 Dispatch_v2.h 中定义的 AT_DISPATCH_V2 格式,适用于将 ATen 内核迁移到新调度 API、处理 aten/src/ATen/native/ 下使用调度宏的文件。
主要能力:识别 AT_DISPATCH_ALL_TYPES_AND{2,3,4}、AT_DISPATCH_FLOATING_TYPES_AND{2,3}、AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND{2,3}、AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES_AND{2,3} 等模式并重写。规则包括:参数重新排序为 scalar_type、name、lambda、类型;用 AT_WRAP(lambda) 包裹 lambda 以处理内部逗号;类型组用 AT_EXPAND() 展开,如 AT_EXPAND(AT_ALL_TYPES);个别类型直接列出,如 kBFloat16、kHalf;并添加 #include <ATen/Dispatch_v2.h>,同时保留原有 Dispatch.h。
用法:先读取文件定位所有 AT_DISPATCH 用法,补上头文件;对每个宏识别模式、映射基础类型组(如 ALL_TYPES→AT_EXPAND(AT_ALL_TYPES),COMPLEX_TYPES→AT_EXPAND(AT_COMPLEX_TYPES))、抽出个别类型、构造 AT_DISPATCH_V2 调用并用 Edit 工具应用,最后展示转换后文件并说明改动。
限制:该技能只做准确转换,不编译或测试代码;旧 Dispatch.h include 需保留,AT_WRAP 为必需,类型组需 AT_EXPAND 而个别类型不需要。