add-uint-support

为运算符添加无符号整数（uint）支持

此技能可帮助通过更新 AT_DISPATCH 宏，为 PyTorch 运算符添加对无符号整数类型（uint16、uint32、uint64）的支持。

何时使用此技能

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

• 为运算符添加对 uint16、uint32 或 uint64 的支持
• 用户提到“无符号类型”“uint 支持”或“基础无符号类型”
• 在内核中启用对 kUInt16、kUInt32、kUInt64 的支持
• 处理需要扩大类型覆盖范围的运算符实现

快速参考

为现有的分派添加无符号类型：

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

// After (method 1: add unsigned types explicitly)
AT_DISPATCH_V2(dtype, "op", AT_WRAP([&]() {
  kernel<scalar_t>();
}), AT_EXPAND(AT_ALL_TYPES), AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES));

// After (method 2: use V2 integral types if AT_INTEGRAL_TYPES present)
AT_DISPATCH_V2(dtype, "op", AT_WRAP([&]() {
  kernel<scalar_t>();
}), AT_EXPAND(AT_INTEGRAL_TYPES_V2), AT_EXPAND(AT_FLOATING_TYPES));

类型组参考

无符号类型组：

• AT_BAREBONES_UNSIGNED_TYPES：kUInt16、kUInt32、kUInt64
• AT_INTEGRAL_TYPES_V2：AT_INTEGRAL_TYPES + AT_BAREBONES_UNSIGNED_TYPES

关系：

AT_INTEGRAL_TYPES          // kByte, kChar, kInt, kLong, kShort
AT_BAREBONES_UNSIGNED_TYPES  // kUInt16, kUInt32, kUInt64
AT_INTEGRAL_TYPES_V2       // INTEGRAL_TYPES + BAREBONES_UNSIGNED_TYPES

说明

步骤 1：确定是否需要转换为 V2

检查文件是否使用 AT_DISPATCH_V2：

如果使用旧版 AT_DISPATCH：

• 先使用 at-dispatch-v2 技能转换为 AT_DISPATCH_V2
• 然后继续添加 uint 支持

如果已使用 AT_DISPATCH_V2：

• 直接继续步骤 2

步骤 2：分析当前的分派宏

识别当前使用的类型组：

AT_DISPATCH_V2(dtype, "op", AT_WRAP([&]() {
  // body
}), AT_EXPAND(AT_ALL_TYPES), kHalf, kBFloat16);
    ^^^^^^^^^^^^^^^^^^^^^^^^^
    Current type coverage

常见模式：

• AT_EXPAND(AT_ALL_TYPES) → 包含 AT_INTEGRAL_TYPES + AT_FLOATING_TYPES
• AT_EXPAND(AT_INTEGRAL_TYPES) → 仅有符号整数
• AT_EXPAND(AT_FLOATING_TYPES) → 浮点类型

步骤 3：选择添加 uint 的方法

有两种方法：

方法 1：显式添加 AT_BAREBONES_UNSIGNED_TYPES

• 适用情况：希望明确添加 uint 支持
• 将 AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES) 添加到类型列表

方法 2：将 AT_INTEGRAL_TYPES 替换为 AT_INTEGRAL_TYPES_V2

• 适用情况：分派已使用 AT_EXPAND(AT_INTEGRAL_TYPES)
• 更简洁：用其超集替换一个类型组
• 仅适用于存在 AT_INTEGRAL_TYPES 的情况

步骤 4：应用转换

方法 1 示例：

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

// After (add unsigned types)
AT_DISPATCH_V2(
    dtype,
    "min_values_cuda",
    AT_WRAP([&]() {
      kernel_impl<scalar_t>(iter);
    }),
    AT_EXPAND(AT_ALL_TYPES),
    AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES),
    kBFloat16, kHalf, kBool
);

方法 2 示例：

// Before
AT_DISPATCH_V2(
    dtype,
    "integral_op",
    AT_WRAP([&]() {
      kernel<scalar_t>();
    }),
    AT_EXPAND(AT_INTEGRAL_TYPES)
);

// After (substitute with V2)
AT_DISPATCH_V2(
    dtype,
    "integral_op",
    AT_WRAP([&]() {
      kernel<scalar_t>();
    }),
    AT_EXPAND(AT_INTEGRAL_TYPES_V2)
);

步骤 5：处理 AT_ALL_TYPES 与单独的类型组

如果分派使用 AT_EXPAND(AT_ALL_TYPES)：

• AT_ALL_TYPES = AT_INTEGRAL_TYPES + AT_FLOATING_TYPES
• 添加 uint：将 AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES) 添加到列表

如果分派分别列出 INTEGRAL 和 FLOATING：

// Before
AT_EXPAND(AT_INTEGRAL_TYPES), AT_EXPAND(AT_FLOATING_TYPES)

// After (Method 2 preferred)
AT_EXPAND(AT_INTEGRAL_TYPES_V2), AT_EXPAND(AT_FLOATING_TYPES)

步骤 6：验证所有分派位置

检查文件中所有需要 uint 支持的分派宏：

• 某些运算符有多个分派位置（CPU、CUDA、不同函数）
• 在所有位置一致地应用转换
• 确保每处都更新为相同的类型覆盖范围

步骤 7：验证更改

检查以下各项：

• 使用 AT_DISPATCH_V2 格式（而非旧版 AT_DISPATCH）
• 通过两种方法之一添加无符号类型
• 更新文件中所有相关的分派位置
• 类型组使用 AT_EXPAND()
• 参数格式正确，并以逗号分隔

常见模式

模式 1：AT_ALL_TYPES + 额外类型

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

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

模式 2：分别列出 INTEGRAL + FLOATING

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

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

模式 3：先转换旧版分派

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

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

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

多个分派位置示例

对于包含多个函数的文件：

void min_values_kernel_cuda(TensorIterator& iter) {
  AT_DISPATCH_V2(iter.dtype(), "min_values_cuda", AT_WRAP([&]() {
    impl<scalar_t>(iter);
  }), AT_EXPAND(AT_ALL_TYPES), AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES), kBFloat16, kHalf);
  //                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  //                           Added uint support
}

void min_launch_kernel(TensorIterator &iter) {
  AT_DISPATCH_V2(iter.input_dtype(), "min_cuda", AT_WRAP([&]() {
    gpu_reduce_kernel<scalar_t>(iter);
  }), AT_EXPAND(AT_ALL_TYPES), AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES), kBFloat16, kHalf);
  //                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  //                           Added uint support here too
}

决策树

使用此决策树确定处理方法：

Is the file using AT_DISPATCH_V2?
├─ No → Use at-dispatch-v2 skill first, then continue
└─ Yes
   └─ Does it use AT_EXPAND(AT_INTEGRAL_TYPES)?
      ├─ Yes → Replace with AT_EXPAND(AT_INTEGRAL_TYPES_V2)
      └─ No → Add AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES) to type list

边缘情况

情况 1：分派仅包含浮点类型

如果运算符仅支持浮点类型，则不要添加 uint 支持：

// Leave as-is - floating point only operator
AT_DISPATCH_V2(dtype, "float_op", AT_WRAP([&]() {
  kernel<scalar_t>();
}), AT_EXPAND(AT_FLOATING_TYPES), kHalf);

情况 2：存在复数类型

无符号类型可以与复数类型一起使用：

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

情况 3：已支持 uint

检查是否已包含 uint 类型：

• 如果使用 AT_INTEGRAL_TYPES_V2 → 已支持 uint
• 如果列表中已包含 AT_BAREBONES_UNSIGNED_TYPES → 已支持 uint
• 如果已支持 uint，则跳过该文件

工作流

当被要求添加 uint 支持时：

1. 读取目标文件
2. 检查是否使用 AT_DISPATCH_V2：
   • 如果没有 → 先使用 at-dispatch-v2 技能
3. 找出所有分发宏位置
4. 对每个分发：
   • 分析当前类型组
   • 选择方法（添加 BAREBONES_UNSIGNED 或升级到 V2）
   • 使用 Edit 工具应用转换
5. 向用户展示更改
6. 说明修改内容

重要说明

• 始终先检查是否需要转换为 v2
• 在文件的所有分发位置一致地应用更改
• 方法 2（AT_INTEGRAL_TYPES_V2）在适用时更简洁
• 方法 1（显式 AT_BAREBONES_UNSIGNED_TYPES）更明确
• 无符号类型包括：kUInt16、kUInt32、kUInt64（不包括 kByte，后者是 uint8）
• 某些运算符在语义上可能不支持无符号类型，请酌情判断

测试

添加 uint 支持后，该运算符应能接受 uint16、uint32 和 uint64 张量。功能测试由用户负责。