at-dispatch-v2

intent

convert legacy pytorch at_dispatch macros to the new at_dispatch_v2 format defined in aten/src/ATen/Dispatch_v2.h. use this skill when porting aten kernels, cuda kernels, or native operator implementations that rely on at_dispatch_all_types_and*, at_dispatch_floating_types*, or similar dispatch macros. the v2 api removes arity from macro names (no more and2, and3, and4 variants) and allows composable type group expansion via at_expand().

inputs

• source file path: c++ file containing one or more legacy at_dispatch_* macro calls (typically in aten/src/ATen/native/, aten/src/ATen/native/cuda/, or similar)
• dispatch_v2.h reference: knowledge of aten/src/ATen/Dispatch_v2.h containing the new macro definitions and type group references (at_all_types, at_floating_types, at_complex_types, at_integral_types, etc.)
• scalar_type expression: the runtime dtype value (e.g., iter.dtype(), tensor.scalar_type(), dtype) passed to the old macro
• lambda body: the kernel implementation code, which may contain commas, function calls, or nested braces
• type list: the individual types appended to the old macro (e.g., khalf, kbfloat16, kbool extracted from at_dispatch_all_types_and3(khalf, kbfloat16, kbool, ...))

procedure

step 1: locate all at_dispatch_* calls in the target file

scan the file for patterns matching at_dispatch_* (case-insensitive in practice). note the exact macro name (e.g., at_dispatch_all_types_and2, at_dispatch_floating_types_and3, at_dispatch_all_types_and_complex_and2).

output of step 1: a list of line numbers and macro names to convert.

step 2: add the dispatch_v2.h include

if not already present, add #include <ATen/Dispatch_v2.h> immediately after or near the existing #include <ATen/Dispatch.h> line. keep the old include in place; other code may depend on it.

input: file location and existing includes section. output: file with the new include directive inserted.

step 3: for each legacy at_dispatch call, identify the components

extract:

• macro base name: e.g., all_types, floating_types, integral_types, all_types_and_complex
• arity suffix: the number appended (and2, and3, and4, etc.) or absence thereof
• individual types: the concrete types passed before scalar_type (e.g., khalf, kbfloat16, kbool)
• scalar_type argument: the dtype expression
• name string: the kernel/operation name as a string literal
• lambda body: the entire bracketed code

input: a single at_dispatch_* call (multi-line or single-line). output: structured breakdown of the five components above.

step 4: map the base macro name to at_expand() groups

use the reference table below to determine which type group(s) to use:

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), at_expand(at_complex_types)
floating_and_complex_types	at_expand(at_floating_types), at_expand(at_complex_types)

for compound macros (e.g., at_dispatch_all_types_and_complex_and2), include both the at_all_types and at_complex_types groups.

input: the base macro name from step 3. output: one or more at_expand(group) tokens.

step 5: construct the at_dispatch_v2 call

assemble the new call with this argument order:

AT_DISPATCH_V2(
  scalar_type,
  "name",
  AT_WRAP(lambda),
  type_groups_with_at_expand,
  individual_types_without_at_expand
)

• position 1: the scalar_type expression (unchanged from old call)
• position 2: the name string (unchanged from old call)
• position 3: the lambda wrapped in at_wrap(...)
• positions 4+: one or more at_expand(group) calls, then individual types (khalf, kbfloat16, etc.)

input: components from steps 3 and 4. output: a single at_dispatch_v2(...) call formatted across multiple lines for readability.

step 6: wrap the lambda in at_wrap()

enclose the entire lambda (from [...]() { to the final }) in at_wrap(...). this prevents internal commas or complex expressions from breaking macro parsing.

input: the lambda body from step 3. output: at_wrap(...{ ... }).

step 7: apply the transformation to the file

replace the old at_dispatch_* call with the new at_dispatch_v2 call using text substitution. preserve all whitespace, indentation, and surrounding code.

input: file path, old call text, new call text. output: file with the macro call replaced.

step 8: repeat steps 3-7 for each legacy dispatch call

iterate through the list from step 1 until all calls are converted.

input: remaining at_dispatch_* calls. output: fully converted file.

step 9: verify the converted file

check:

• all at_dispatch_* calls have been replaced with at_dispatch_v2 calls
• each at_dispatch_v2 call has at_wrap() around the lambda
• type groups use at_expand(); individual types do not
• argument order is: scalar_type, "name", at_wrap(lambda), type_groups, individual_types
• #include <ATen/Dispatch_v2.h> is present
• no syntax errors (unmatched brackets, missing commas, etc.)

input: converted file. output: confirmation of correctness, or a list of issues to fix.

decision points

if the old macro has an arity suffix (and2, and3, etc.): extract all the individual types from before the scalar_type argument and add them as trailing arguments to at_dispatch_v2 after the type group(s). each type becomes a separate argument; do not wrap individual types in at_expand().

if the old macro has no arity suffix (e.g., at_dispatch_all_types): there are no individual types to extract. proceed directly to constructing at_dispatch_v2 with only type groups. example:

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

if the lambda contains commas (function calls, nested structures, etc.): mandatory use of at_wrap(). without it, the preprocessor will misparse the commas as macro argument separators.

if the base macro name includes both "and_complex" and an arity suffix (e.g., at_dispatch_all_types_and_complex_and2): include multiple at_expand() calls for both at_all_types and at_complex_types, then list the individual types. example:

// before
AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND2(kComplexHalf, kHalf, dtype, "op", [&]() {});

// after
AT_DISPATCH_V2(dtype, "op", AT_WRAP([&]() {}), AT_EXPAND(AT_ALL_TYPES), AT_EXPAND(AT_COMPLEX_TYPES), kComplexHalf, kHalf);

if the target file does not already include dispatch_v2.h: add it near the existing dispatch.h include. do not remove dispatch.h, as other code in the file may depend on it.

if a dispatch call spans multiple lines with inconsistent indentation: preserve the indentation pattern of the surrounding file. reformat the new at_dispatch_v2 call to match the file's style.

output contract

the converted file must have the following properties:

• location: same file path as the input
• format: valid c++ preprocessor code, compatible with the pytorch build system
• includes: #include <ATen/Dispatch_v2.h> present; #include <ATen/Dispatch.h> retained
• macro calls: all legacy at_dispatch_* macros replaced with at_dispatch_v2(...) calls
• lambda wrapping: every lambda argument wrapped in at_wrap()
• type groups: every type group macro wrapped in at_expand(); individual types not wrapped
• argument order: (scalar_type, "name", at_wrap(lambda), type_groups..., individual_types...)
• syntax: no unmatched brackets, parentheses, or quotes; commas correctly placed
• semantics: the behavior and template instantiation of the kernel remain unchanged; scalar_t is still available inside the lambda

outcome signal

the user will know the skill worked if:

1. the file compiles without errors related to macro expansion or type mismatches (verified by running cd <pytorch_root> && python setup.py build or the local build command, though this skill does not perform compilation)
2. all at_dispatch_* macros have been replaced; a grep for at_dispatch_all_types|at_dispatch_floating_types|at_dispatch_integral_types returns no matches in the converted file (except in comments)
3. the user can inspect the converted file and see clear at_dispatch_v2 calls with consistent formatting, one call per kernel
4. the kernel still receives scalar_t as a template parameter inside the lambda, and all type-dependent operations work as before
5. the user can cross-reference the file against aten/src/ATen/Dispatch_v2.h and confirm that all type groups and individual types are valid