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 の追加方法を選択する

方法は 2 つあります。

方法 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) を使用しているとき
• より簡潔です。1 つの型グループをそのスーパーセットで置き換えます
• 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 ではない）
• 2 つの方法のいずれかで符号なし型を追加している
• ファイル内の該当するすべてのディスパッチ箇所を更新している
• 型グループで 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 のテンソルを受け入れる必要があります。機能テストはユーザーが実施します。