# AT_DISPATCH 到 AT_DISPATCH_V2 转换器此技能有助于将 PyTorch 的旧 AT_DISPATCH 宏转换为新的 AT_DISPATCH_V2 格式，如 `aten/src/ATen/Dispatch_v2.h` 中定义。 ## 什么时候使用这个技能在以下情况下使用此技能： * 将 AT_DISPATCH_* 宏转换为 AT_DISPATCH_V2 * 移植 ATen 内核以使用新的调度 API * 处理 `aten/src/ATen/native/` 中使用调度宏的文件 * 用户提及“AT_DISPATCH”、“dispatch v2”、“Dispatch_v2.h”或宏转换 ## 快速参考 **旧格式：** AT_DISPATCH_ALL_TYPES_AND3(kBFloat16, kHalf, kBool, dtype, "kernel_name", [&]() { // 拉姆达主体 }); **新格式：** AT_DISPATCH_V2(dtype, "kernel_name", AT_WRAP([&]() { // 拉姆达主体 }), AT_EXPAND(AT_ALL_TYPES), kBFloat16, kHalf, kBool); ## 关键转变 1. **重新排序参数**：首先是 `scalar_type` 和 `name`，然后是 lambda，然后是类型 2. **包裹 lambda** : 使用 `AT_WRAP(lambda)` 处理内部逗号 3. **扩展类型组** ：使用`AT_EXPAND(AT_ALL_TYPES)`代替隐式扩展 4. **列出单个类型** ：在扩展组后添加额外类型（kHalf、kBFloat16等） 5. **添加 include** : `#include ` 靠近其他 Dispatch include ## 说明 ### 第 1 步：添加 Dispatch_v2.h 包含文件在现有的 `#include ` 附近添加 v2 标头： #include #include 暂时保留旧的 Dispatch.h 包含文件（其他代码可能仍然需要它）。 ### 第 2 步：识别旧的调度模式要转换的常见模式： * `AT_DISPATCH_ALL_TYPES_AND{2,3,4}(type1, type2, ..., scalar_type, name, lambda)` * `AT_DISPATCH_FLOATING_TYPES_AND{2,3}(type1, type2, ..., scalar_type, name, lambda)` * `AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND{2,3}(type1, ..., scalar_type, name, lambda)` * `AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES_AND{2,3}(type1, ..., scalar_type, name, lambda)` ### 步骤 3：将旧宏映射到类型组确定哪个类型组宏对应于基本类型：老宏基地| AT_DISPATCH_V2类型组 ---|--- `所有类型` | `AT_EXPAND(AT_ALL_TYPES)` `浮动类型` | `AT_EXPAND(AT_FLOATING_TYPES)` `INTEGRAL_TYPES` | `AT_EXPAND(AT_INTEGRAL_TYPES)` `复杂_类型` | `AT_EXPAND(AT_COMPLEX_TYPES)` `所有类型和复杂` | `AT_EXPAND(AT_ALL_TYPES_AND_COMPLEX)` 对于组合模式，请使用多个“AT_EXPAND()”条目： // 旧：AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND2(...) // 新增：AT_EXPAND(AT_ALL_TYPES)、AT_EXPAND(AT_COMPLEX_TYPES)、type1、type2 ### 步骤 4：提取各个类型从“AT_DISPATCH_*_AND2(type1, type2, ...)”或“AT_DISPATCH_*_AND3(type1, type2, type3, ...)”中提取各个类型（type1、type2 等）。这些成为类型组之后的尾随参数： AT_DISPATCH_V2(..., AT_EXPAND(AT_ALL_TYPES), kBFloat16, kHalf, kBool) ^^^^^^^^^^^^^^^^^^^^^^^^^^ AND3 中的各个类型 ### 步骤5：转换为AT_DISPATCH_V2 应用变换： **图案：** AT_DISPATCH_V2( scalar_type, // 第一个：dtype 表达式 "name", // 第二个：调试字符串 AT_WRAP(lambda), // 第三个：包裹在 AT_WRAP 中的 lambda type_groups, // 第 4 个以上：使用 AT_EXPAND() 键入组 individual_types // 最后：个体类型） **转换示例：** // 之前 AT_DISPATCH_ALL_TYPES_AND3( kBFloat16、kHalf、kBool、迭代器.dtype(), “最小_值_cuda”， [&]() { min_values_kernel_cuda_impl(iter); } ）； // 之后 AT_DISPATCH_V2( 迭代器.dtype(), “最小_值_cuda”， AT_WRAP([&]() { min_values_kernel_cuda_impl(iter); }), AT_EXPAND（AT_ALL_TYPES）， kBFloat16、kHalf、kBool ）； ### 步骤 6：处理多行 lambda 对于带有内部逗号或复杂表达式的 lambda，AT_WRAP 是必不可少的： AT_DISPATCH_V2( 数据类型， “复杂的内核”， AT_WRAP([&]() { gpu_reduce_kernel( 迭代器， MinOps{}, Thrust::pair(upper_bound(), 0) // 里面有逗号！）； }), AT_EXPAND（AT_ALL_TYPES））； ### 步骤 7：验证转换检查： * `AT_WRAP()` 包装整个 lambda * 类型组使用`AT_EXPAND()` * 个别类型没有`AT_EXPAND()`（只有`kBFloat16`，而不是`AT_EXPAND(kBFloat16)`） * 参数顺序为：scalar_type、名称、lambda、类型 * 包含添加：`#include ` ## 输入组引用可用的类型组宏（与`AT_EXPAND()`一起使用）： AT_INTEGRAL_TYPES // kByte、kChar、kInt、kLong、kShort AT_FLOATING_TYPES // kDouble, kFloat AT_COMPLEX_TYPES // kComplexDouble, kComplexFloat AT_QINT_TYPES // kQInt8, kQUInt8, kQInt32 AT_ALL_TYPES // INTEGRAL_TYPES + FLOATING_TYPES AT_ALL_TYPES_AND_COMPLEX // ALL_TYPES + COMPLEX_TYPES AT_INTEGRAL_TYPES_V2 // INTEGRAL_TYPES + 无符号类型 AT_BAREBONES_UNSIGNED_TYPES // kUInt16、kUInt32、kUInt64 AT_FLOAT8_TYPES // Float8 变体 ## 常见模式 ### 模式：AT_DISPATCH_ALL_TYPES_AND2 // 之前 AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBFloat16, dtype, "op", [&]() { 内核（数据）； }); // 之后 AT_DISPATCH_V2(dtype, "op", AT_WRAP([&]() { 内核（数据）； }), AT_EXPAND(AT_ALL_TYPES), kHalf, kBFloat16); ### 模式：AT_DISPATCH_FLOATING_TYPES_AND3 // 之前 AT_DISPATCH_FLOATING_TYPES_AND3(kHalf, kBFloat16, kFloat8_e4m3fn, 张量.scalar_type(), "float_op", [&] { 过程（张量）； }); // 之后 AT_DISPATCH_V2(tensor.scalar_type(), "float_op", AT_WRAP([&] { 过程（张量）； }), AT_EXPAND(AT_FLOATING_TYPES), kHalf, kBFloat16, kFloat8_e4m3fn); ### 模式：AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND2 // 之前 AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND2( kComplexHalf, kHalf, self.scalar_type(), “复杂操作”， [&] { 结果=计算(self); } ）； // 之后 AT_DISPATCH_V2( self.scalar_type(), “复杂操作”， AT_WRAP([&] { 结果=计算(self); }), AT_EXPAND（AT_ALL_TYPES）， AT_EXPAND（AT_COMPLEX_TYPES）， kComplexHalf, k半）； ## 边缘情况 ### 情况 1：没有额外的类型（罕见） // 之前 AT_DISPATCH_ALL_TYPES(dtype, "op", [&]() { kernel(); }); // 之后 AT_DISPATCH_V2(dtype, "op", AT_WRAP([&]() { 内核(); }), AT_EXPAND(AT_ALL_TYPES)); ### 情况 2：许多个体类型（AND4、AND5 等） // 之前 AT_DISPATCH_FLOATING_TYPES_AND4(kHalf, kBFloat16, kFloat8_e4m3fn, kFloat8_e5m2, dtype, "float8_op", [&]() { kernel(); }); // 之后 AT_DISPATCH_V2(dtype, "float8_op", AT_WRAP([&]() { 内核(); }), AT_EXPAND(AT_FLOATING_TYPES), kHalf, kBFloat16, kFloat8_e4m3fn, kFloat8_e5m2); ### 案例 3：没有捕获的 Lambda // 之前 AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBool, dtype, "op", []() { static_kernel(); }); // 之后 AT_DISPATCH_V2(dtype, "op", AT_WRAP([]() { static_kernel(); }), AT_EXPAND(AT_ALL_TYPES), kHalf, kBool); ## AT_DISPATCH_V2 的好处 1. **宏名称中没有参数** : AND2, AND3, AND4 不需要不同的宏 2. **可组合类型集** ：使用“AT_EXPAND()”混合和匹配类型组 3. **可扩展**：轻松添加更多类型而不会触及宏限制 4. **更清晰**：类型组是显式的，而不是隐含在宏名称中 ## 重要提示 * 保留 `#include ` \ - 其他代码可能需要它 * `AT_WRAP()` 是强制性的 - 防止 lambda 中出现逗号解析问题 * 类型组需要`AT_EXPAND()`，单个类型不需要 * v2 API 位于 `aten/src/ATen/Dispatch_v2.h` 中 - 请参阅它以获取完整文档 * 请参阅Python脚本的头文件来重新生成宏实现 ## 工作流程当要求转换 AT_DISPATCH 宏时： 1.读取文件识别所有AT_DISPATCH的使用 2. 添加 `#include `（如果不存在） 3. 对于每个调度宏： * 识别模式并提取成分 * 映射基本类型组 * 提取单个类型 * 构造AT_DISPATCH_V2调用 * 使用编辑工具应用 4.向用户显示完整的转换后的文件 5. 解释更改的内容不要编译或测试代码 - 仅关注准确的转换。

在调度 V2

包含技能