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, kUInt64AT_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_TYPESAT_EXPAND(AT_INTEGRAL_TYPES)→ entiers signés uniquementAT_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_V2est utilisé → support uint déjà présent - Si
AT_BAREBONES_UNSIGNED_TYPESest 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 :
- Lisez le fichier cible
- Vérifiez si AT_DISPATCH_V2 est utilisé :
- Si non → utilisez d'abord la skill at-dispatch-v2
- Identifiez tous les sites de macro dispatch
- 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
- Montrez les modifications à l'utilisateur
- 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.