pytorch/at-dispatch-v2

O que at-dispatch-v2 faz?

Converta macros AT_DISPATCH do PyTorch para o formato AT_DISPATCH_V2 em código C++ do ATen. Use ao portar AT_DISPATCH_ALL_TYPES_AND*, AT_DISPATCH_FLOATING_TYPES* ou outras macros de dispatch para a nova API v2. Para arquivos kernel do ATen, kernels CUDA e implementações nativas de operadores.

Descrição

Este skill converte automaticamente as chamadas de macro de dispatch do PyTorch, como AT_DISPATCH_ALL_TYPES_AND* e AT_DISPATCH_FLOATING_TYPES*, para o novo formato AT_DISPATCH_V2. Isso é necessário para migrar código C++ do ATen (incluindo kernels CUDA e implementações nativas) para a API v2 do PyTorch.

Uso

1. Selecione o bloco de código contendo macros AT_DISPATCH.
2. Execute o comando pytorch/at-dispatch-v2.
3. O código será convertido para o formato AT_DISPATCH_V2 equivalente.

Exemplo de Conversão

Antes:

AT_DISPATCH_ALL_TYPES_AND(kHalf, input, "my_op", [&]() {
  // ... código usando scalar_t ...
});

Depois:

AT_DISPATCH_ALL_TYPES_AND_V2(kHalf, input, "my_op", AT_EXPAND(AT_WRAP([&]() {
  // ... código usando scalar_t ...
})));

Macros Suportadas

• AT_DISPATCH_ALL_TYPES_AND → AT_DISPATCH_ALL_TYPES_AND_V2
• AT_DISPATCH_FLOATING_TYPES → AT_DISPATCH_FLOATING_TYPES_V2
• AT_DISPATCH_COMPLEX_TYPES → AT_DISPATCH_COMPLEX_TYPES_V2
• AT_DISPATCH_INTEGRAL_TYPES → AT_DISPATCH_INTEGRAL_TYPES_V2
• Outras macros similares

Detalhes Técnicos

• Preserva a indentação e estrutura do código.
• Mantém comentários e espaços em branco.
• Converte corretamente listas de tipos (ex.: kHalf, kBFloat16).
• Funciona em arquivos .cpp e .cu.

Limitações

• Não modifica a lógica interna das lambdas (apenas o invólucro da macro).
• Requer que as macros estejam escritas no formato padrão (sem customizações).

Contribuições

Relate problemas ou sugira melhorias no repositório oficial.