add-uint-support

# 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:** ```cpp // 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:** ```cpp 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: ```cpp 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:** ```cpp // 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:** ```cpp // 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: ```cpp // 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 ```cpp // 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 ```cpp // 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 ```cpp // 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: ```cpp 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: ```cpp // 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: ```cpp 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: 1. Read the target file 2. Check if using AT_DISPATCH_V2: - If not → use at-dispatch-v2 skill first 3. Identify all dispatch macro sites 4. For each dispatch: - Analyze current type groups - Choose method (add BAREBONES_UNSIGNED or upgrade to V2) - Apply transformation with Edit tool 5. Show the user the changes 6. 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.