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 に変換する場合
• 新しいディスパッチ API を使うように ATen カーネルを移植する場合
• ディスパッチマクロを使用している 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 を指定し、次にラムダ、最後に型を指定します
2. ラムダをラップする：内部のカンマを処理するために 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：複数行のラムダを処理する

内部にカンマが含まれるラムダや複雑な式を使う場合は、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() でラムダ全体をラップしている
• 型グループで AT_EXPAND() を使用している
• 個別の型には AT_EXPAND() を付けていない（AT_EXPAND(kBFloat16) ではなく kBFloat16）
• 引数の順序が scalar_type、name、ラムダ、型になっている
• インクルードを追加している：#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：キャプチャのないラムダ

// 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() は必須です。ラムダ内でカンマが誤って解析されるのを防ぎます。
• 型グループには 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. 変更内容を説明します。

コードのコンパイルやテストは行わないでください。正確な変換のみに集中してください。