at-dispatch-v2スキルでPyTorchのディスパッチマクロをV2形式に変換する方法

PyTorchのATenカーネル開発で、レガシーなAT_DISPATCHマクロからAT_DISPATCH_V2への移行作業を行う場面はありませんか？マクロの引数の並び替え、ラムダのラップ、型グループの展開など、変換ルールが複数あるため、手作業ではミスが起きやすい作業です。

この記事では、at-dispatch-v2スキルを使ってAT_DISPATCHマクロを安全にV2形式に変換する方法を解説します。

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

at-dispatch-v2は、PyTorchのレガシーAT_DISPATCHマクロをAT_DISPATCH_V2形式に変換するスキルです：

• AT_DISPATCH_ALL_TYPES_AND*、AT_DISPATCH_FLOATING_TYPES_AND*等のマクロを対象
• 引数の並び替え（scalar_type, name → 先頭に移動）
• ラムダのAT_WRAP()によるラップ
• 型グループのAT_EXPAND()による展開
• 個別型（kHalf, kBFloat16等）の分離

ATenカーネルやCUDAカーネルの保守・移行作業を行うPyTorchコア開発者に向いています。

インストール方法

前提条件

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

インストールコマンド

Copyclaude mcp add github.com/pytorch/pytorch/tree/main/.claude/skills/at-dispatch-v2

使い方

基本的な使い方

変換したいファイルを指示します：

CopyこのファイルのAT_DISPATCHマクロをV2に変換して

変換ルール

旧形式：

CopyAT_DISPATCH_ALL_TYPES_AND3(kBFloat16, kHalf, kBool,
    dtype, "kernel_name", [&]() {
  kernel<scalar_t>();
});

新形式：

CopyAT_DISPATCH_V2(dtype, "kernel_name", AT_WRAP([&]() {
  kernel<scalar_t>();
}), AT_EXPAND(AT_ALL_TYPES), kBFloat16, kHalf, kBool);

型グループの対応表

知っておくべき注意点

AT_WRAP()は必須

ラムダ内にカンマが含まれる場合、マクロのパースが壊れます。AT_WRAP()は省略できません。

個別型にAT_EXPAND()は不要

kHalf、kBFloat16などの個別型はそのまま記述します。AT_EXPAND(kHalf)とは書きません。

Dispatch.hは残す

#include <ATen/Dispatch.h>は他のコードが依存している可能性があるため、#include <ATen/Dispatch_v2.h>を追加しても既存のインクルードは削除しないでください。

まとめ

at-dispatch-v2スキルを使うと、レガシーAT_DISPATCHマクロからV2形式への変換を、引数並び替え・ラムダラップ・型グループ展開のルールに沿って進められます。V2形式は拡張性が高く、AND2、AND3といったマクロ名の使い分けが不要になります。

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

関連リンク

→ at-dispatch-v2 スキル詳細