--- name: at-dispatch-v2 description: Convert PyTorch AT_DISPATCH macros to AT_DISPATCH_V2 format in ATen C++ code. Use when porting AT_DISPATCH_ALL_TYPES_AND*, AT_DISPATCH_FLOATING_TYPES*, or other dispatch macros to the new v2 API. For ATen kernel files, CUDA kernels, and native operator implementations. --- # AT_DISPATCH to AT_DISPATCH_V2 Converter This skill helps convert PyTorch's legacy AT_DISPATCH macros to the new AT_DISPATCH_V2 format, as defined in `aten/src/ATen/Dispatch_v2.h`. ## When to use this skill Use this skill when: - Converting AT_DISPATCH_* macros to AT_DISPATCH_V2 - Porting ATen kernels to use the new dispatch API - Working with files in `aten/src/ATen/native/` that use dispatch macros - User mentions "AT_DISPATCH", "dispatch v2", "Dispatch_v2.h", or macro conversion ## Quick reference **Old format:** ```cpp AT_DISPATCH_ALL_TYPES_AND3(kBFloat16, kHalf, kBool, dtype, "kernel_name", [&]() { // lambda body }); ``` **New format:** ```cpp AT_DISPATCH_V2(dtype, "kernel_name", AT_WRAP([&]() { // lambda body }), AT_EXPAND(AT_ALL_TYPES), kBFloat16, kHalf, kBool); ``` ## Key transformations 1. **Reorder arguments**: `scalar_type` and `name` come first, then lambda, then types 2. **Wrap the lambda**: Use `AT_WRAP(lambda)` to handle internal commas 3. **Expand type groups**: Use `AT_EXPAND(AT_ALL_TYPES)` instead of implicit expansion 4. **List individual types**: Add extra types (kHalf, kBFloat16, etc.) after expanded groups 5. **Add include**: `#include ` near other Dispatch includes ## Instructions ### Step 1: Add the Dispatch_v2.h include Add the v2 header near the existing `#include `: ```cpp #include #include ``` Keep the old Dispatch.h include for now (other code may still need it). ### Step 2: Identify the old dispatch pattern Common patterns to convert: - `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)` ### Step 3: Map the old macro to type groups Identify which type group macro corresponds to the base types: | Old macro base | AT_DISPATCH_V2 type group | |----------------|---------------------------| | `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)` | For combined patterns, use multiple `AT_EXPAND()` entries: ```cpp // Old: AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND2(...) // New: AT_EXPAND(AT_ALL_TYPES), AT_EXPAND(AT_COMPLEX_TYPES), type1, type2 ``` ### Step 4: Extract the individual types From `AT_DISPATCH_*_AND2(type1, type2, ...)` or `AT_DISPATCH_*_AND3(type1, type2, type3, ...)`, extract the individual types (type1, type2, etc.). These become the trailing arguments after the type group: ```cpp AT_DISPATCH_V2(..., AT_EXPAND(AT_ALL_TYPES), kBFloat16, kHalf, kBool) ^^^^^^^^^^^^^^^^^^^^^^^^ Individual types from AND3 ``` ### Step 5: Transform to AT_DISPATCH_V2 Apply the transformation: **Pattern:** ```cpp 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 ) ``` **Example transformation:** ```cpp // BEFORE AT_DISPATCH_ALL_TYPES_AND3( kBFloat16, kHalf, kBool, iter.dtype(), "min_values_cuda", [&]() { min_values_kernel_cuda_impl(iter); } ); // AFTER AT_DISPATCH_V2( iter.dtype(), "min_values_cuda", AT_WRAP([&]() { min_values_kernel_cuda_impl(iter); }), AT_EXPAND(AT_ALL_TYPES), kBFloat16, kHalf, kBool ); ``` ### Step 6: Handle multi-line lambdas For lambdas with internal commas or complex expressions, AT_WRAP is essential: ```cpp AT_DISPATCH_V2( dtype, "complex_kernel", AT_WRAP([&]() { gpu_reduce_kernel( iter, MinOps{}, thrust::pair(upper_bound(), 0) // Commas inside! ); }), AT_EXPAND(AT_ALL_TYPES) ); ``` ### Step 7: Verify the conversion Check that: - [ ] `AT_WRAP()` wraps the entire lambda - [ ] Type groups use `AT_EXPAND()` - [ ] Individual types don't have `AT_EXPAND()` (just `kBFloat16`, not `AT_EXPAND(kBFloat16)`) - [ ] Argument order is: scalar_type, name, lambda, types - [ ] Include added: `#include ` ## Type group reference Available type group macros (use with `AT_EXPAND()`): ```cpp 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 ``` ## Common patterns ### Pattern: AT_DISPATCH_ALL_TYPES_AND2 ```cpp // Before AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBFloat16, dtype, "op", [&]() { kernel(data); }); // After AT_DISPATCH_V2(dtype, "op", AT_WRAP([&]() { kernel(data); }), AT_EXPAND(AT_ALL_TYPES), kHalf, kBFloat16); ``` ### Pattern: AT_DISPATCH_FLOATING_TYPES_AND3 ```cpp // Before AT_DISPATCH_FLOATING_TYPES_AND3(kHalf, kBFloat16, kFloat8_e4m3fn, tensor.scalar_type(), "float_op", [&] { process(tensor); }); // After AT_DISPATCH_V2(tensor.scalar_type(), "float_op", AT_WRAP([&] { process(tensor); }), AT_EXPAND(AT_FLOATING_TYPES), kHalf, kBFloat16, kFloat8_e4m3fn); ``` ### Pattern: AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND2 ```cpp // Before AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND2( kComplexHalf, kHalf, self.scalar_type(), "complex_op", [&] { result = compute(self); } ); // After AT_DISPATCH_V2( self.scalar_type(), "complex_op", AT_WRAP([&] { result = compute(self); }), AT_EXPAND(AT_ALL_TYPES), AT_EXPAND(AT_COMPLEX_TYPES), kComplexHalf, kHalf ); ``` ## Edge cases ### Case 1: No extra types (rare) ```cpp // Before AT_DISPATCH_ALL_TYPES(dtype, "op", [&]() { kernel(); }); // After AT_DISPATCH_V2(dtype, "op", AT_WRAP([&]() { kernel(); }), AT_EXPAND(AT_ALL_TYPES)); ``` ### Case 2: Many individual types (AND4, AND5, etc.) ```cpp // Before AT_DISPATCH_FLOATING_TYPES_AND4(kHalf, kBFloat16, kFloat8_e4m3fn, kFloat8_e5m2, dtype, "float8_op", [&]() { kernel(); }); // After AT_DISPATCH_V2(dtype, "float8_op", AT_WRAP([&]() { kernel(); }), AT_EXPAND(AT_FLOATING_TYPES), kHalf, kBFloat16, kFloat8_e4m3fn, kFloat8_e5m2); ``` ### Case 3: Lambda with no captures ```cpp // Before AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBool, dtype, "op", []() { static_kernel(); }); // After AT_DISPATCH_V2(dtype, "op", AT_WRAP([]() { static_kernel(); }), AT_EXPAND(AT_ALL_TYPES), kHalf, kBool); ``` ## Benefits of AT_DISPATCH_V2 1. **No arity in macro name**: Don't need different macros for AND2, AND3, AND4 2. **Composable type sets**: Mix and match type groups with `AT_EXPAND()` 3. **Extensible**: Easy to add more types without hitting macro limits 4. **Clearer**: Type groups are explicit, not implicit in macro name ## Important notes - Keep `#include ` - other code may need it - The `AT_WRAP()` is mandatory - prevents comma parsing issues in the lambda - Type groups need `AT_EXPAND()`, individual types don't - The v2 API is in `aten/src/ATen/Dispatch_v2.h` - refer to it for full docs - See the header file for the Python script to regenerate the macro implementation ## Workflow When asked to convert AT_DISPATCH macros: 1. Read the file to identify all AT_DISPATCH uses 2. Add `#include ` if not present 3. For each dispatch macro: - Identify the pattern and extract components - Map the base type group - Extract individual types - Construct the AT_DISPATCH_V2 call - Apply with Edit tool 4. Show the user the complete converted file 5. Explain what was changed Do NOT compile or test the code - focus on accurate conversion only.