add-uint-supportスキルでPyTorchオペレーターにuint型サポートを追加する方法

PyTorchのオペレーターが符号なし整数型（uint16、uint32、uint64）に対応していないために、型エラーが発生した経験はありませんか？ATENカーネルの型ディスパッチを拡張する作業は、マクロの構造を正確に理解する必要があります。

この記事では、PyTorchリポジトリのadd-uint-supportスキルを使って、AT_DISPATCHマクロにuint型サポートを追加する方法を解説します。

このスキルは何をしてくれるのか

add-uint-supportは、PyTorchオペレーターのAT_DISPATCHマクロを更新して、uint16/uint32/uint64のサポートを追加するスキルです：

• AT_DISPATCH_V2形式への変換確認
• 2つの追加メソッド（明示的追加 or 型グループ置換）の選択ガイド
• 決定木による適切なアプローチの判定
• ファイル内の全ディスパッチサイトの一括更新

PyTorchコアの開発に携わり、オペレーターの型カバレッジを拡張する必要がある開発者に向いています。

インストール方法

前提条件

• Claude Codeがインストール済みであること
• PyTorchリポジトリのソースコードへのアクセス

インストールコマンド

Copyclaude mcp add github.com/pytorch/pytorch/tree/main/.claude/skills/add-uint-support

使い方

基本的な使い方

uint型サポートを追加したいファイルを指示します：

Copyこのファイルにuint型サポートを追加して

2つの追加メソッド

メソッド1: AT_BAREBONES_UNSIGNED_TYPESを明示的に追加

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

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

メソッド2: AT_INTEGRAL_TYPESをAT_INTEGRAL_TYPES_V2に置換

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

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

決定木

CopyAT_DISPATCH_V2を使用しているか？
├─ No → 先にat-dispatch-v2スキルで変換
└─ Yes
   └─ AT_INTEGRAL_TYPESを使用しているか？
      ├─ Yes → メソッド2（V2に置換）
      └─ No → メソッド1（明示的に追加）

知っておくべき注意点

V2形式が前提

レガシーのAT_DISPATCHマクロを使用している場合、先にat-dispatch-v2スキルでV2形式に変換する必要があります。

浮動小数点型のみのオペレーターは対象外

AT_FLOATING_TYPESのみを使用するオペレーターは、意味的にuint型をサポートする必要がないため、変更しないでください。

全ディスパッチサイトを一貫して更新

1つのファイル内にCPU用・CUDA用など複数のディスパッチサイトがある場合、すべてを一貫して更新する必要があります。

まとめ

add-uint-supportスキルを使うと、PyTorchオペレーターへのuint16/uint32/uint64サポート追加を、2つのメソッドから適切な方を選択して進められます。決定木に従い、V2変換の必要性も含めて判断できます。

詳細な仕様やオプションは、スキル詳細ページをご確認ください。

関連リンク

→ add-uint-support スキル詳細