/add-uint-support

name: add-uint-support
description: Add unsigned integer (uint) type support to PyTorch operators by updating AT_DISPATCH macros. Use when adding support for uint16, uint32, uint64 types to operators, kernels, or when user mentions enabling unsigned types, barebones unsigned types, or 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:**

// 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:**

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:

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:**

// 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:**

// 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:

// 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

// 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

// 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

// 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