at-dispatch-v2

AT_DISPATCH 到 AT_DISPATCH_V2 转换器

此技能可帮助将 PyTorch 的旧版 AT_DISPATCH 宏转换为 aten/src/ATen/Dispatch_v2.h 中定义的新格式 AT_DISPATCH_V2。

何时使用此技能

在以下情况下使用此技能：

• 将 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", [&]() {
  // lambda body
});

新格式：

AT_DISPATCH_V2(dtype, "kernel_name", AT_WRAP([&]() {
  // lambda body
}), 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. 添加头文件：在其他 Dispatch 头文件附近添加 #include <ATen/Dispatch_v2.h>

说明

步骤 1：添加 Dispatch_v2.h 头文件

在现有的 #include <ATen/Dispatch.h> 附近添加 v2 头文件：

#include <ATen/Dispatch.h>
#include <ATen/Dispatch_v2.h>

暂时保留旧的 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 类型组
ALL_TYPES	AT_EXPAND(AT_ALL_TYPES)
FLOATING_TYPES	AT_EXPAND(AT_FLOATING_TYPES)
INTEGRAL_TYPES	AT_EXPAND(AT_INTEGRAL_TYPES)
COMPLEX_TYPES	AT_EXPAND(AT_COMPLEX_TYPES)
ALL_TYPES_AND_COMPLEX	AT_EXPAND(AT_ALL_TYPES_AND_COMPLEX)

对于组合模式，请使用多个 AT_EXPAND() 条目：

// Old: AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND2(...)
// New: 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)
                                             ^^^^^^^^^^^^^^^^^^^^^^^^
                                             Individual types from AND3

步骤 5：转换为 AT_DISPATCH_V2

应用以下转换：

模式：

AT_DISPATCH_V2(
  scalar_type,           // 1st: The dtype expression
  "name",                // 2nd: The debug string
  AT_WRAP(lambda),       // 3rd: The lambda wrapped in AT_WRAP
  type_groups,           // 4th+: Type groups with AT_EXPAND()
  individual_types       // Last: Individual types
)

转换示例：

// BEFORE
AT_DISPATCH_ALL_TYPES_AND3(
    kBFloat16, kHalf, kBool,
    iter.dtype(),
    "min_values_cuda",
    [&]() {
      min_values_kernel_cuda_impl<scalar_t>(iter);
    }
);

// AFTER
AT_DISPATCH_V2(
    iter.dtype(),
    "min_values_cuda",
    AT_WRAP([&]() {
      min_values_kernel_cuda_impl<scalar_t>(iter);
    }),
    AT_EXPAND(AT_ALL_TYPES),
    kBFloat16, kHalf, kBool
);

步骤 6：处理多行 lambda

对于包含内部逗号或复杂表达式的 lambda，AT_WRAP 至关重要：

AT_DISPATCH_V2(
    dtype,
    "complex_kernel",
    AT_WRAP([&]() {
      gpu_reduce_kernel<scalar_t, scalar_t>(
        iter,
        MinOps<scalar_t>{},
        thrust::pair<scalar_t, int64_t>(upper_bound(), 0)  // Commas inside!
      );
    }),
    AT_EXPAND(AT_ALL_TYPES)
);

步骤 7：验证转换

检查以下各项：

• AT_WRAP() 包装了整个 lambda
• 类型组使用 AT_EXPAND()
• 单独的类型不使用 AT_EXPAND()（只写 kBFloat16，而不是 AT_EXPAND(kBFloat16)）
• 参数顺序为：scalar_type、name、lambda、类型
• 已添加头文件：#include <ATen/Dispatch_v2.h>

类型组参考

可用的类型组宏（需与 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 + unsigned types
AT_BAREBONES_UNSIGNED_TYPES  // kUInt16, kUInt32, kUInt64
AT_FLOAT8_TYPES       // Float8 variants

常见模式

模式：AT_DISPATCH_ALL_TYPES_AND2

// Before
AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBFloat16, dtype, "op", [&]() {
  kernel<scalar_t>(data);
});

// After
AT_DISPATCH_V2(dtype, "op", AT_WRAP([&]() {
  kernel<scalar_t>(data);
}), AT_EXPAND(AT_ALL_TYPES), kHalf, kBFloat16);

模式：AT_DISPATCH_FLOATING_TYPES_AND3

// Before
AT_DISPATCH_FLOATING_TYPES_AND3(kHalf, kBFloat16, kFloat8_e4m3fn,
    tensor.scalar_type(), "float_op", [&] {
  process<scalar_t>(tensor);
});

// After
AT_DISPATCH_V2(tensor.scalar_type(), "float_op", AT_WRAP([&] {
  process<scalar_t>(tensor);
}), AT_EXPAND(AT_FLOATING_TYPES), kHalf, kBFloat16, kFloat8_e4m3fn);

模式：AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND2

// Before
AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND2(
    kComplexHalf, kHalf,
    self.scalar_type(),
    "complex_op",
    [&] {
      result = compute<scalar_t>(self);
    }
);

// After
AT_DISPATCH_V2(
    self.scalar_type(),
    "complex_op",
    AT_WRAP([&] {
      result = compute<scalar_t>(self);
    }),
    AT_EXPAND(AT_ALL_TYPES),
    AT_EXPAND(AT_COMPLEX_TYPES),
    kComplexHalf,
    kHalf
);

边缘情况

情况 1：没有额外类型（少见）

// Before
AT_DISPATCH_ALL_TYPES(dtype, "op", [&]() { kernel<scalar_t>(); });

// After
AT_DISPATCH_V2(dtype, "op", AT_WRAP([&]() {
  kernel<scalar_t>();
}), AT_EXPAND(AT_ALL_TYPES));

情况 2：有多个单独的类型（AND4、AND5 等）

// Before
AT_DISPATCH_FLOATING_TYPES_AND4(kHalf, kBFloat16, kFloat8_e4m3fn, kFloat8_e5m2,
    dtype, "float8_op", [&]() { kernel<scalar_t>(); });

// After
AT_DISPATCH_V2(dtype, "float8_op", AT_WRAP([&]() {
  kernel<scalar_t>();
}), AT_EXPAND(AT_FLOATING_TYPES), kHalf, kBFloat16, kFloat8_e4m3fn, kFloat8_e5m2);

情况 3：lambda 不捕获任何变量

// Before
AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBool, dtype, "op", []() {
  static_kernel<scalar_t>();
});

// After
AT_DISPATCH_V2(dtype, "op", AT_WRAP([]() {
  static_kernel<scalar_t>();
}), AT_EXPAND(AT_ALL_TYPES), kHalf, kBool);

AT_DISPATCH_V2 的优势

1. 宏名称中不含参数数量：无需为 AND2、AND3、AND4 使用不同的宏
2. 类型集合可组合：使用 AT_EXPAND() 混合搭配类型组
3. 易于扩展：添加更多类型很方便，不会受宏数量限制
4. 更清晰：类型组明确列出，不再隐含于宏名称中

重要说明

• 保留 #include <ATen/Dispatch.h>——其他代码可能需要它
• AT_WRAP() 是必需的——可避免 lambda 中的逗号解析问题
• 类型组需要使用 AT_EXPAND()，单个类型不需要
• v2 API 位于 aten/src/ATen/Dispatch_v2.h——请参阅该文件了解完整文档
• 查看头文件，了解用于重新生成宏实现的 Python 脚本

工作流

当被要求转换 AT_DISPATCH 宏时：

1. 阅读文件，找出所有 AT_DISPATCH 的使用位置
2. 如果不存在，则添加 #include <ATen/Dispatch_v2.h>
3. 对每个调度宏：
   • 识别模式并提取各个组成部分
   • 映射基础类型组
   • 提取单个类型
   • 构造 AT_DISPATCH_V2 调用
   • 使用 Edit 工具进行修改
4. 向用户展示转换后的完整文件
5. 说明做了哪些更改

不要编译或测试代码——只需专注于准确转换。