at-dispatch-v2

Par pytorch · pytorch

Convertir les macros PyTorch `AT_DISPATCH` au format `AT_DISPATCH_V2` dans le code ATen C++. À utiliser lors du portage de `AT_DISPATCH_ALL_TYPES_AND*`, `AT_DISPATCH_FLOATING_TYPES*` ou d'autres macros de dispatch vers la nouvelle API v2. Pour les fichiers de kernels ATen, les kernels CUDA et les implémentations d'opérateurs natifs.

npx skills add https://github.com/pytorch/pytorch --skill at-dispatch-v2

Convertisseur AT_DISPATCH vers AT_DISPATCH_V2

Cette skill aide à convertir les macros AT_DISPATCH héritées de PyTorch vers le nouveau format AT_DISPATCH_V2, tel que défini dans aten/src/ATen/Dispatch_v2.h.

Quand utiliser cette skill

Utilisez cette skill quand :

  • Conversion de macros ATDISPATCH* vers AT_DISPATCH_V2
  • Portage de kernels ATen pour utiliser la nouvelle API de dispatch
  • Travail avec des fichiers dans aten/src/ATen/native/ qui utilisent des macros de dispatch
  • L'utilisateur mentionne « AT_DISPATCH », « dispatch v2 », « Dispatch_v2.h », ou conversion de macros

Référence rapide

Ancien format :

AT_DISPATCH_ALL_TYPES_AND3(kBFloat16, kHalf, kBool, dtype, "kernel_name", [&]() {
  // corps lambda
});

Nouveau format :

AT_DISPATCH_V2(dtype, "kernel_name", AT_WRAP([&]() {
  // corps lambda
}), AT_EXPAND(AT_ALL_TYPES), kBFloat16, kHalf, kBool);

Transformations clés

  1. Réordonner les arguments : scalar_type et name en premier, puis lambda, puis types
  2. Envelopper la lambda : Utiliser AT_WRAP(lambda) pour gérer les virgules internes
  3. Développer les groupes de types : Utiliser AT_EXPAND(AT_ALL_TYPES) au lieu du développement implicite
  4. Lister les types individuels : Ajouter les types supplémentaires (kHalf, kBFloat16, etc.) après les groupes développés
  5. Ajouter l'include : #include <ATen/Dispatch_v2.h> près des autres includes Dispatch

Instructions

Étape 1 : Ajouter l'include Dispatch_v2.h

Ajoutez l'en-tête v2 près du #include <ATen/Dispatch.h> existant :

#include <ATen/Dispatch.h>
#include <ATen/Dispatch_v2.h>

Conservez l'ancien include Dispatch.h pour l'instant (autre code peut en avoir besoin).

Étape 2 : Identifier l'ancien pattern de dispatch

Patterns courants à convertir :

  • 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)

Étape 3 : Mapper l'ancienne macro aux groupes de types

Identifiez quel macro de groupe de type correspond aux types de base :

Base de l'ancienne macro Groupe de type AT_DISPATCH_V2
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)

Pour les patterns combinés, utilisez plusieurs entrées AT_EXPAND() :

// Ancien : AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND2(...)
// Nouveau : AT_EXPAND(AT_ALL_TYPES), AT_EXPAND(AT_COMPLEX_TYPES), type1, type2

Étape 4 : Extraire les types individuels

De AT_DISPATCH_*_AND2(type1, type2, ...) ou AT_DISPATCH_*_AND3(type1, type2, type3, ...), extrayez les types individuels (type1, type2, etc.).

Ces types deviennent les arguments suivants après le groupe de type :

AT_DISPATCH_V2(..., AT_EXPAND(AT_ALL_TYPES), kBFloat16, kHalf, kBool)
                                             ^^^^^^^^^^^^^^^^^^^^^^^^
                                             Types individuels de AND3

Étape 5 : Transformer en AT_DISPATCH_V2

Appliquez la transformation :

Pattern :

AT_DISPATCH_V2(
  scalar_type,           // 1er : L'expression de dtype
  "name",                // 2e : La chaîne de debug
  AT_WRAP(lambda),       // 3e : La lambda enveloppée dans AT_WRAP
  type_groups,           // 4e+ : Groupes de type avec AT_EXPAND()
  individual_types       // Dernier : Types individuels
)

Exemple de transformation :

// AVANT
AT_DISPATCH_ALL_TYPES_AND3(
    kBFloat16, kHalf, kBool,
    iter.dtype(),
    "min_values_cuda",
    [&]() {
      min_values_kernel_cuda_impl<scalar_t>(iter);
    }
);

// APRÈS
AT_DISPATCH_V2(
    iter.dtype(),
    "min_values_cuda",
    AT_WRAP([&]() {
      min_values_kernel_cuda_impl<scalar_t>(iter);
    }),
    AT_EXPAND(AT_ALL_TYPES),
    kBFloat16, kHalf, kBool
);

Étape 6 : Gérer les lambdas multi-lignes

Pour les lambdas avec virgules internes ou expressions complexes, AT_WRAP est essentiel :

AT_DISPATCH_V2(
    dtype,
    "complex_kernel",
    AT_WRAP([&]() {
      gpu_reduce_kernel<scalar_t, scalar_t>(
        iter,
        MinOps<scalar_t>{},
        thrust::pair<scalar_t, int64_t>(upper_bound(), 0)  // Virgules à l'intérieur !
      );
    }),
    AT_EXPAND(AT_ALL_TYPES)
);

Étape 7 : Vérifier la conversion

Vérifiez que :

  • [ ] AT_WRAP() enveloppe la lambda entière
  • [ ] Les groupes de type utilisent AT_EXPAND()
  • [ ] Les types individuels n'ont pas AT_EXPAND() (juste kBFloat16, pas AT_EXPAND(kBFloat16))
  • [ ] L'ordre des arguments est : scalar_type, name, lambda, types
  • [ ] Include ajouté : #include <ATen/Dispatch_v2.h>

Référence des groupes de type

Macros de groupe de type disponibles (à utiliser avec AT_EXPAND()) :

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 + types non-signés
AT_BAREBONES_UNSIGNED_TYPES  // kUInt16, kUInt32, kUInt64
AT_FLOAT8_TYPES       // Variantes Float8

Patterns courants

Pattern : AT_DISPATCH_ALL_TYPES_AND2

// Avant
AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBFloat16, dtype, "op", [&]() {
  kernel<scalar_t>(data);
});

// Après
AT_DISPATCH_V2(dtype, "op", AT_WRAP([&]() {
  kernel<scalar_t>(data);
}), AT_EXPAND(AT_ALL_TYPES), kHalf, kBFloat16);

Pattern : AT_DISPATCH_FLOATING_TYPES_AND3

// Avant
AT_DISPATCH_FLOATING_TYPES_AND3(kHalf, kBFloat16, kFloat8_e4m3fn,
    tensor.scalar_type(), "float_op", [&] {
  process<scalar_t>(tensor);
});

// Après
AT_DISPATCH_V2(tensor.scalar_type(), "float_op", AT_WRAP([&] {
  process<scalar_t>(tensor);
}), AT_EXPAND(AT_FLOATING_TYPES), kHalf, kBFloat16, kFloat8_e4m3fn);

Pattern : AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND2

// Avant
AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND2(
    kComplexHalf, kHalf,
    self.scalar_type(),
    "complex_op",
    [&] {
      result = compute<scalar_t>(self);
    }
);

// Après
AT_DISPATCH_V2(
    self.scalar_type(),
    "complex_op",
    AT_WRAP([&] {
      result = compute<scalar_t>(self);
    }),
    AT_EXPAND(AT_ALL_TYPES),
    AT_EXPAND(AT_COMPLEX_TYPES),
    kComplexHalf,
    kHalf
);

Cas limites

Cas 1 : Pas de types supplémentaires (rare)

// Avant
AT_DISPATCH_ALL_TYPES(dtype, "op", [&]() { kernel<scalar_t>(); });

// Après
AT_DISPATCH_V2(dtype, "op", AT_WRAP([&]() {
  kernel<scalar_t>();
}), AT_EXPAND(AT_ALL_TYPES));

Cas 2 : Nombreux types individuels (AND4, AND5, etc.)

// Avant
AT_DISPATCH_FLOATING_TYPES_AND4(kHalf, kBFloat16, kFloat8_e4m3fn, kFloat8_e5m2,
    dtype, "float8_op", [&]() { kernel<scalar_t>(); });

// Après
AT_DISPATCH_V2(dtype, "float8_op", AT_WRAP([&]() {
  kernel<scalar_t>();
}), AT_EXPAND(AT_FLOATING_TYPES), kHalf, kBFloat16, kFloat8_e4m3fn, kFloat8_e5m2);

Cas 3 : Lambda sans captures

// Avant
AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBool, dtype, "op", []() {
  static_kernel<scalar_t>();
});

// Après
AT_DISPATCH_V2(dtype, "op", AT_WRAP([]() {
  static_kernel<scalar_t>();
}), AT_EXPAND(AT_ALL_TYPES), kHalf, kBool);

Avantages de AT_DISPATCH_V2

  1. Pas d'arité dans le nom de macro : Pas besoin de macros différentes pour AND2, AND3, AND4
  2. Ensembles de type composables : Combinez et appariez des groupes de type avec AT_EXPAND()
  3. Extensible : Facile d'ajouter plus de types sans atteindre les limites des macros
  4. Plus clair : Les groupes de type sont explicites, pas implicites dans le nom de la macro

Notes importantes

  • Conservez #include <ATen/Dispatch.h> - autre code peut en avoir besoin
  • Le AT_WRAP() est obligatoire - empêche les problèmes d'analyse des virgules dans la lambda
  • Les groupes de type ont besoin de AT_EXPAND(), les types individuels non
  • L'API v2 se trouve dans aten/src/ATen/Dispatch_v2.h - consultez-la pour la documentation complète
  • Consultez le fichier d'en-tête pour le script Python afin de régénérer l'implémentation des macros

Workflow

Quand on vous demande de convertir des macros AT_DISPATCH :

  1. Lire le fichier pour identifier tous les usages AT_DISPATCH
  2. Ajouter #include <ATen/Dispatch_v2.h> s'il n'est pas présent
  3. Pour chaque macro de dispatch :
    • Identifier le pattern et extraire les composants
    • Mapper le groupe de type de base
    • Extraire les types individuels
    • Construire l'appel AT_DISPATCH_V2
    • Appliquer avec l'outil Edit
  4. Afficher le fichier converti complet à l'utilisateur
  5. Expliquer ce qui a changé

N'essayez PAS de compiler ou tester le code - concentrez-vous sur une conversion précise uniquement.

Skills similaires

at-dispatch-v2
pytorch / pytorch

Convertir les macros PyTorch ATDISPATCH legacy vers le format moderne ATDISPATCH_V2.

104 801
add-uint-support
pytorch / pytorch

Ajouter le support des entiers non signés (uint16/32/64) aux opérateurs PyTorch.

104 801
add-uint-support
pytorch / pytorch

Ajouter le support des entiers non signés (uint16/32/64) aux opérateurs PyTorch.

104 801
quality-playbook
github / awesome-copilot

Générer un système qualité complet et personnalisé pour tout projet logiciel.

40 074
metal-kernel
pytorch / pytorch

Implémenter des kernels Metal natifs pour opérateurs PyTorch sur Apple Silicon.

104 801