add-uint-support

⚠ Archivé — pas de mise à jour depuis 1 mois

Par pytorch · pytorch

Ajoutez la prise en charge des types entiers non signés (uint) aux opérateurs PyTorch en mettant à jour les macros `AT_DISPATCH`. À utiliser lors de l'ajout de la prise en charge des types `uint16`, `uint32`, `uint64` aux opérateurs et kernels, ou lorsque l'utilisateur mentionne l'activation des types non signés, des types non signés de base, ou le support uint.

npx skills add https://github.com/pytorch/pytorch --skill add-uint-support

Ajouter le support des entiers non signés (uint) aux opérateurs

Cette skill aide à ajouter le support des types entiers non signés (uint16, uint32, uint64) aux opérateurs PyTorch en mettant à jour leurs macros AT_DISPATCH.

Quand utiliser cette skill

Utilisez cette skill quand :

  • Vous ajoutez le support uint16, uint32 ou uint64 à un opérateur
  • L'utilisateur mentionne « unsigned types », « uint support », « barebones unsigned types »
  • Vous activez le support pour kUInt16, kUInt32, kUInt64 dans les kernels
  • Vous travaillez sur des implémentations d'opérateurs qui ont besoin d'une couverture de types étendue

Référence rapide

Ajouter les types non signés à une dispatch existante :

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

// Après (méthode 1 : ajouter les types non signés explicitement)
AT_DISPATCH_V2(dtype, "op", AT_WRAP([&]() {
  kernel<scalar_t>();
}), AT_EXPAND(AT_ALL_TYPES), AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES));

// Après (méthode 2 : utiliser les types intégraux V2 si AT_INTEGRAL_TYPES est présent)
AT_DISPATCH_V2(dtype, "op", AT_WRAP([&]() {
  kernel<scalar_t>();
}), AT_EXPAND(AT_INTEGRAL_TYPES_V2), AT_EXPAND(AT_FLOATING_TYPES));

Référence des groupes de types

Groupes de types non signés :

  • AT_BAREBONES_UNSIGNED_TYPES : kUInt16, kUInt32, kUInt64
  • AT_INTEGRAL_TYPES_V2 : AT_INTEGRAL_TYPES + AT_BAREBONES_UNSIGNED_TYPES

Relations :

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

Étape 1 : Déterminer si une conversion en V2 est nécessaire

Vérifiez si le fichier utilise AT_DISPATCH_V2 :

Si vous utilisez l'ancien AT_DISPATCH :

  • Convertissez d'abord en AT_DISPATCH_V2 à l'aide de la skill at-dispatch-v2
  • Puis procédez à l'ajout du support uint

Si vous utilisez déjà AT_DISPATCH_V2 :

  • Allez directement à l'étape 2

Étape 2 : Analyser la macro dispatch actuelle

Identifiez les groupes de types actuellement en usage :

AT_DISPATCH_V2(dtype, "op", AT_WRAP([&]() {
  // body
}), AT_EXPAND(AT_ALL_TYPES), kHalf, kBFloat16);
    ^^^^^^^^^^^^^^^^^^^^^^^^^
    Couverture de type actuelle

Motifs courants :

  • AT_EXPAND(AT_ALL_TYPES) → inclut AT_INTEGRAL_TYPES + AT_FLOATING_TYPES
  • AT_EXPAND(AT_INTEGRAL_TYPES) → entiers signés uniquement
  • AT_EXPAND(AT_FLOATING_TYPES) → types virgule flottante

Étape 3 : Choisir la méthode d'ajout de uint

Deux approches :

Méthode 1 : Ajouter AT_BAREBONES_UNSIGNED_TYPES explicitement

  • À utiliser quand : Vous voulez être explicite sur l'ajout du support uint
  • Ajoutez AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES) à la liste de types

Méthode 2 : Remplacer AT_INTEGRAL_TYPES par AT_INTEGRAL_TYPES_V2

  • À utiliser quand : La dispatch utilise déjà AT_EXPAND(AT_INTEGRAL_TYPES)
  • Plus concis : remplace un groupe de types par son sur-ensemble
  • Applicable uniquement si AT_INTEGRAL_TYPES est présent

Étape 4 : Appliquer la transformation

Exemple méthode 1 :

// Avant
AT_DISPATCH_V2(
    dtype,
    "min_values_cuda",
    AT_WRAP([&]() {
      kernel_impl<scalar_t>(iter);
    }),
    AT_EXPAND(AT_ALL_TYPES),
    kBFloat16, kHalf, kBool
);

// Après (ajouter les types non signés)
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
);

Exemple méthode 2 :

// Avant
AT_DISPATCH_V2(
    dtype,
    "integral_op",
    AT_WRAP([&]() {
      kernel<scalar_t>();
    }),
    AT_EXPAND(AT_INTEGRAL_TYPES)
);

// Après (remplacer par V2)
AT_DISPATCH_V2(
    dtype,
    "integral_op",
    AT_WRAP([&]() {
      kernel<scalar_t>();
    }),
    AT_EXPAND(AT_INTEGRAL_TYPES_V2)
);

Étape 5 : Gérer AT_ALL_TYPES vs groupes de types individuels

Si la dispatch utilise AT_EXPAND(AT_ALL_TYPES) :

  • AT_ALL_TYPES = AT_INTEGRAL_TYPES + AT_FLOATING_TYPES
  • Pour ajouter uint : ajoutez AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES) à la liste

Si la dispatch liste séparément INTEGRAL et FLOATING :

// Avant
AT_EXPAND(AT_INTEGRAL_TYPES), AT_EXPAND(AT_FLOATING_TYPES)

// Après (méthode 2 préférée)
AT_EXPAND(AT_INTEGRAL_TYPES_V2), AT_EXPAND(AT_FLOATING_TYPES)

Étape 6 : Vérifier tous les sites de dispatch

Vérifiez le fichier pour TOUTES les macros dispatch qui ont besoin du support uint :

  • Certains opérateurs ont plusieurs sites de dispatch (CPU, CUDA, différentes fonctions)
  • Appliquez la transformation de manière cohérente sur tous les sites
  • Assurez-vous que chacun obtient les mêmes mises à jour de couverture de type

Étape 7 : Valider les modifications

Vérifiez que :

  • [ ] Le format AT_DISPATCH_V2 est utilisé (pas l'ancien AT_DISPATCH)
  • [ ] Les types non signés sont ajoutés via l'une des deux méthodes
  • [ ] Tous les sites de dispatch pertinents du fichier sont mis à jour
  • [ ] Les groupes de types utilisent AT_EXPAND()
  • [ ] Les arguments sont correctement formatés et séparés par des virgules

Motifs courants

Motif 1 : AT_ALL_TYPES + extras

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

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

Motif 2 : INTEGRAL + FLOATING séparés

// Avant
AT_DISPATCH_V2(dtype, "op", AT_WRAP([&]() {
  kernel<scalar_t>();
}), AT_EXPAND(AT_INTEGRAL_TYPES), AT_EXPAND(AT_FLOATING_TYPES));

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

Motif 3 : L'ancienne dispatch a besoin d'une conversion d'abord

// Avant (nécessite d'abord une conversion v2)
AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBFloat16, dtype, "op", [&]() {
  kernel<scalar_t>();
});

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

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

Exemple avec plusieurs sites de dispatch

Pour un fichier avec plusieurs fonctions :

void min_values_kernel_cuda(TensorIterator& iter) {
  AT_DISPATCH_V2(iter.dtype(), "min_values_cuda", AT_WRAP([&]() {
    impl<scalar_t>(iter);
  }), AT_EXPAND(AT_ALL_TYPES), AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES), kBFloat16, kHalf);
  //                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  //                           Support uint ajouté
}

void min_launch_kernel(TensorIterator &iter) {
  AT_DISPATCH_V2(iter.input_dtype(), "min_cuda", AT_WRAP([&]() {
    gpu_reduce_kernel<scalar_t>(iter);
  }), AT_EXPAND(AT_ALL_TYPES), AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES), kBFloat16, kHalf);
  //                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  //                           Support uint ajouté ici aussi
}

Arbre de décision

Utilisez cet arbre de décision pour déterminer l'approche :

Le fichier utilise-t-il AT_DISPATCH_V2 ?
├─ Non → Utilisez d'abord la skill at-dispatch-v2, puis continuez
└─ Oui
   └─ Utilise-t-il AT_EXPAND(AT_INTEGRAL_TYPES) ?
      ├─ Oui → Remplacez par AT_EXPAND(AT_INTEGRAL_TYPES_V2)
      └─ Non → Ajoutez AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES) à la liste de types

Cas limites

Cas 1 : Dispatch avec uniquement les types virgule flottante

Si l'opérateur ne supporte que les types virgule flottante, n'ajoutez pas le support uint :

// Laisser tel quel - opérateur virgule flottante uniquement
AT_DISPATCH_V2(dtype, "float_op", AT_WRAP([&]() {
  kernel<scalar_t>();
}), AT_EXPAND(AT_FLOATING_TYPES), kHalf);

Cas 2 : Types complexes présents

Les types non signés fonctionnent aux côtés des types complexes :

AT_DISPATCH_V2(dtype, "op", AT_WRAP([&]() {
  kernel<scalar_t>();
}), AT_EXPAND(AT_ALL_TYPES),
    AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES),
    AT_EXPAND(AT_COMPLEX_TYPES),
    kHalf, kBFloat16);

Cas 3 : Support uint déjà présent

Vérifiez si les types uint sont déjà présents :

  • Si AT_INTEGRAL_TYPES_V2 est utilisé → support uint déjà présent
  • Si AT_BAREBONES_UNSIGNED_TYPES est déjà dans la liste → support uint déjà présent
  • Ignorez le fichier si le support uint est déjà présent

Flux de travail

Quand on vous demande d'ajouter le support uint :

  1. Lisez le fichier cible
  2. Vérifiez si AT_DISPATCH_V2 est utilisé :
    • Si non → utilisez d'abord la skill at-dispatch-v2
  3. Identifiez tous les sites de macro dispatch
  4. Pour chaque dispatch :
    • Analysez les groupes de types actuels
    • Choisissez la méthode (ajouter BAREBONES_UNSIGNED ou mettre à niveau en V2)
    • Appliquez la transformation avec l'outil Edit
  5. Montrez les modifications à l'utilisateur
  6. Expliquez ce qui a été modifié

Notes importantes

  • Vérifiez toujours si une conversion v2 est nécessaire d'abord
  • Appliquez les modifications de manière cohérente sur tous les sites de dispatch du fichier
  • La méthode 2 (AT_INTEGRAL_TYPES_V2) est plus propre quand applicable
  • La méthode 1 (AT_BAREBONES_UNSIGNED_TYPES explicite) est plus explicite
  • Les types non signés sont : kUInt16, kUInt32, kUInt64 (pas kByte qui est uint8)
  • Certains opérateurs peuvent ne pas supporter sémantiquement les types non signés - utilisez votre jugement

Tests

Après l'ajout du support uint, l'opérateur doit accepter les tenseurs uint16, uint32 et uint64. L'utilisateur est responsable des tests fonctionnels.

Skills similaires