PyTorch 运算符添加无符号整数支持指南：AT_DISPATCH宏更新与uint16/32/64集成

add-uint-support by pytorch/pytorch

203 周安装量

98,500 GitHub Stars

GitHub

安装命令

npx skills add https://github.com/pytorch/pytorch --skill add-uint-support

AI/机器学习开发 C++

🇨🇳中文介绍

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

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

何时使用此技能

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

为运算符添加 uint16、uint32 或 uint64 支持时
用户提及"无符号类型"、"uint 支持"、"barebones unsigned types"时
在内核中启用 kUInt16、kUInt32、kUInt64 支持时
处理需要扩展类型覆盖范围的运算符实现时

快速参考

向现有调度添加无符号类型：

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

// After (方法 1：显式添加无符号类型)
AT_DISPATCH_V2(dtype, "op", AT_WRAP([&]() {
  kernel<scalar_t>();
}), AT_EXPAND(AT_ALL_TYPES), AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES));

// After (方法 2：如果存在 AT_INTEGRAL_TYPES，则使用 V2 整数类型)
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

广告位招租

在这里展示您的产品或服务

触达数万 AI 开发者，精准高效

联系我们

步骤 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);
        ^^^^^^^^^^^^^^^^^^^^^^^^^
        当前类型覆盖范围

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_V2 替换 AT_INTEGRAL_TYPES

使用时机：调度已经使用 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 (添加无符号类型)
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 (替换为 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 (首选方法 2)
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 (需要先进行 v2 转换)
AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBFloat16, dtype, "op", [&]() {
  kernel<scalar_t>();
});

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

// After 添加 uint 支持
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);
  //                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  //                           已添加 uint 支持
}

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);
  //                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  //                           此处也已添加 uint 支持
}

使用此决策树确定方法：

文件是否使用 AT_DISPATCH_V2？
├─ 否 → 先使用 at-dispatch-v2 技能，然后继续
└─ 是
   └─ 是否使用 AT_EXPAND(AT_INTEGRAL_TYPES)？
      ├─ 是 → 替换为 AT_EXPAND(AT_INTEGRAL_TYPES_V2)
      └─ 否 → 将 AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES) 添加到类型列表

情况 1：仅包含浮点类型的调度

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

// 保持原样 - 仅浮点运算符
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 支持时：

读取目标文件
检查是否使用 AT_DISPATCH_V2：
- 如果不是 → 先使用 at-dispatch-v2 技能
识别所有调度宏点
对于每个调度：
- 分析当前类型组
- 选择方法（添加 BAREBONES_UNSIGNED 或升级到 V2）
- 使用 Edit 工具应用转换
向用户展示更改
解释修改内容

始终首先检查是否需要 v2 转换
在文件中的所有调度点上一致地应用更改
方法 2（AT_INTEGRAL_TYPES_V2）在适用时更简洁
方法 1（显式 AT_BAREBONES_UNSIGNED_TYPES）更明确
无符号类型是：kUInt16、kUInt32、kUInt64（不是 kByte，它是 uint8）
某些运算符在语义上可能不支持无符号类型 - 请使用判断

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

🇺🇸English

Add Unsigned Integer (uint) Support to Operators

This skill helps add support for unsigned integer types (uint16, uint32, uint64) to PyTorch operators by updating their AT_DISPATCH macros.

When to use this skill

Use this skill when:

Adding uint16, uint32, or uint64 support to an operator
User mentions "unsigned types", "uint support", "barebones unsigned types"
Enabling support for kUInt16, kUInt32, kUInt64 in kernels
Working with operator implementations that need expanded type coverage

Quick reference

Add unsigned types to existing dispatch:

// 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));

Type group reference

Unsigned type groups:

AT_BAREBONES_UNSIGNED_TYPES: kUInt16, kUInt32, kUInt64
AT_INTEGRAL_TYPES_V2: AT_INTEGRAL_TYPES + AT_BAREBONES_UNSIGNED_TYPES

Relationship:

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

Instructions

Step 1: Determine if conversion to V2 is needed

Check if the file uses AT_DISPATCH_V2:

If using old AT_DISPATCH:

First convert to AT_DISPATCH_V2 using the at-dispatch-v2 skill
Then proceed with adding uint support

If already using AT_DISPATCH_V2:

Proceed directly to Step 2

Step 2: Analyze the current dispatch macro

Identify what type groups are currently in use:

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

Common patterns:

AT_EXPAND(AT_ALL_TYPES) → includes AT_INTEGRAL_TYPES + AT_FLOATING_TYPES
AT_EXPAND(AT_INTEGRAL_TYPES) → signed integers only
AT_EXPAND(AT_FLOATING_TYPES) → floating point types

Step 3: Choose the uint addition method

Two approaches:

Method 1: Add AT_BAREBONES_UNSIGNED_TYPES explicitly

Use when: You want to be explicit about adding uint support
Add AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES) to the type list

Method 2: Substitute AT_INTEGRAL_TYPES with AT_INTEGRAL_TYPES_V2

Use when: The dispatch already uses AT_EXPAND(AT_INTEGRAL_TYPES)
More concise: replaces one type group with its superset
Only applicable if AT_INTEGRAL_TYPES is present

Step 4: Apply the transformation

Method 1 example:

// 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
);

Method 2 example:

// 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)
);

Step 5: Handle AT_ALL_TYPES vs individual type groups

If the dispatch uses AT_EXPAND(AT_ALL_TYPES):

AT_ALL_TYPES = AT_INTEGRAL_TYPES + AT_FLOATING_TYPES
To add uint: add AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES) to the list

If the dispatch separately lists INTEGRAL and 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)

Step 6: Verify all dispatch sites

Check the file for ALL dispatch macros that need uint support:

Some operators have multiple dispatch sites (CPU, CUDA, different functions)
Apply the transformation consistently across all sites
Ensure each gets the same type coverage updates

Step 7: Validate the changes

Check that:

AT_DISPATCH_V2 format is used (not old AT_DISPATCH)
Unsigned types are added via one of the two methods
All relevant dispatch sites in the file are updated
Type groups use AT_EXPAND()
Arguments are properly formatted and comma-separated

Common patterns

Pattern 1: AT_ALL_TYPES + extras

// 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);

Pattern 2: Separate 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));

Pattern 3: Old dispatch needs conversion first

// 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);

Multiple dispatch sites example

For a file with multiple functions:

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
}

Decision tree

Use this decision tree to determine the approach:

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

Edge cases

Case 1: Dispatch with only floating types

If the operator only supports floating point types, don't add uint support:

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

Case 2: Complex types present

Unsigned types work alongside complex types:

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);

Case 3: Already has uint support

Check if uint types are already present:

If AT_INTEGRAL_TYPES_V2 is used → already has uint support
If AT_BAREBONES_UNSIGNED_TYPES is already in list → already has uint support
Skip the file if uint support is already present

Workflow

When asked to add uint support:

Read the target file
Check if using AT_DISPATCH_V2:
- If not → use at-dispatch-v2 skill first
Identify all dispatch macro sites
For each dispatch:
- Analyze current type groups
- Choose method (add BAREBONES_UNSIGNED or upgrade to V2)
- Apply transformation with Edit tool
Show the user the changes
Explain what was modified

Important notes

Always check if v2 conversion is needed first
Apply changes consistently across all dispatch sites in the file
Method 2 (AT_INTEGRAL_TYPES_V2) is cleaner when applicable
Method 1 (explicit AT_BAREBONES_UNSIGNED_TYPES) is more explicit
Unsigned types are: kUInt16, kUInt32, kUInt64 (not kByte which is uint8)
Some operators may not semantically support unsigned types - use judgment

Testing

After adding uint support, the operator should accept uint16, uint32, and uint64 tensors. The user is responsible for functional testing.

Weekly Installs

203

Repository

pytorch/pytorch

GitHub Stars

98.5K

First Seen

Jan 20, 2026

Security Audits

Gen Agent Trust HubPass SocketPass SnykPass

Installed on

claude-code191

opencode190

cursor186

codex184

gemini-cli183

github-copilot174

React 组合模式指南：Vercel 组件架构最佳实践，提升代码可维护性

111,800 周安装