Largeur d'index CUDA dans PyTorch

Utilisez cette skill quand un kernel CUDA déborde de l'indexation int, échoue près de 2^31 éléments, ou nécessite un examen de la mathématique d'index int vs int64_t.

Commencer par le site du débordement

Ne convertissez pas aveuglément toutes les variables d'index en int64_t. Trouvez l'expression qui peut dépasser 32 bits et classifiez où elle s'exécute :

Configuration unique par CTA ou par groupe d'éléments de sortie : préférez un cast 64-bit local à la première multiplication.
Indexation linéaire per-élément chaude : template le kernel sur index_t et dispatchez int vs int64_t.
Algorithme non supporté avec hypothèses 32-bit uniquement : ajoutez un TORCH_CHECK(canUse32BitIndexMath(...)) clair au lieu de déborder silencieusement.
Débordement de dimension de grille : changer le type arithmétique ne suffit pas ; ajoutez du tiling/striding sur cette dimension.

Utilitaires canoniques

Incluez/utilisez les utilitaires PyTorch existants au lieu de vérifications ad hoc :

#include <ATen/native/CanUse32BitIndexMath.h>
#include <ATen/cuda/detail/KernelUtils.h>

at::native::canUse32BitIndexMath(tensor, INT_MAX) vérifie à la fois numel et l'offset de stockage maximal.
Pour les fichiers CUDA qui incluent déjà des wrappers de détails cuda, at::cuda::detail::canUse32BitIndexMath peut être disponible en tant qu'alias.
AT_DISPATCH_INDEX_TYPES(cond ? ScalarType::Int : ScalarType::Long, "name", [&] { ... }); fournit index_t.
CUDA_KERNEL_LOOP_TYPE(index, nthreads, index_t) garde les boucles grid-stride correctes pour l'une ou l'autre largeur.

Vérifiez tous les tenseurs dont les offsets sont calculés avec le type d'index sélectionné, pas seulement le tenseur de sortie.

Motifs de correction

1. Débordement d'offset de base localisé

Si seul un offset de pointeur de base peut déborder et qu'il est calculé en dehors de la boucle chaude, gardez le kernel par ailleurs inchangé :

int64_t plane = blockIdx.x;
input = input + plane * strideD;
output = output + plane * osizeH * osizeW;

Cela évite de doubler les instantiations de kernel et garde l'arithmétique de boucle interne en 32-bit. Utilisez ceci quand les dimensions à l'intérieur de la tile tiennent toujours dans int.

2. Kernels linéaires grid-stride

Si l'index de boucle, la décomposition modulo/division, ou l'accès final data[index] peut dépasser 32 bits, template le kernel :

template <typename scalar_t, typename index_t>
__global__ void kernel(index_t n, const scalar_t* in, scalar_t* out) {
  CUDA_KERNEL_LOOP_TYPE(index, n, index_t) {
    out[index] = in[index];
  }
}

AT_DISPATCH_INDEX_TYPES(
    canUse32BitIndexMath(out, INT_MAX) && canUse32BitIndexMath(in, INT_MAX)
        ? ScalarType::Int
        : ScalarType::Long,
    "kernel_index_type",
    [&] {
      kernel<scalar_t, index_t><<<blocks, threads, 0, stream>>>(n, in, out);
      C10_CUDA_KERNEL_LAUNCH_CHECK();
    });

Préférez ceci à un changement inconditionnel de l'index de boucle en int64_t, car la division/modulo 64-bit dans une boucle CUDA chaude peut être mesurable.

3. TensorInfo/accessor ou kernels strided

Quand les offsets sont calculés à partir de tailles/strides, dispatchez sur un type d'index seulement si tous les tenseurs participants passent canUse32BitIndexMath pour ce type. Rappelez-vous qu'un tenseur petit numel() peut toujours avoir besoin d'offsets 64-bit s'il s'agit d'une grande vue strided.

4. Kernels 32-bit uniquement

Si supporter l'indexation 64-bit nécessiterait une réécriture d'algorithme plus grande ou dépasserait les limites de lancement CUDA, échouez tôt :

TORCH_CHECK(
    canUse32BitIndexMath(input) && canUse32BitIndexMath(output),
    "op_name: tensors must fit into 32-bit index math");

Utilisez ceci seulement quand l'opérateur a déjà une limitation de taille documentée ou acceptée ; ne transformez pas un bug de correction signalé en une limitation inutile.

Compromis taille binaire et performance

Templater sur index_t duplique chaque kernel affecté pour chaque dtype scalaire et spécialisation de format mémoire. Avant d'ajouter un dispatch d'index à plusieurs kernels, demandez-vous si le débordement se trouve dans un chemin chaud ou seulement dans une expression de configuration.

Candidats A/B pour réparer quand le choix n'est pas évident :

Construisez chaque candidat à partir d'un diff propre en utilisant le même environnement de construction.

Enregistrez les tailles d'objet CUDA et de bibliothèque modifiées :

stat -c '%s %n' build/aten/src/ATen/CMakeFiles/torch_cuda.dir/native/cuda/<file>.cu.o torch/lib/libtorch_cuda.so

Vérifiez la multiplication de symboles pour le nom du kernel :

nm -S --size-sort -C torch/lib/libtorch_cuda.so | rg '<kernel_name>|index_t|long|int'

Si l'arithmétique de boucle chaude a changé, faites un benchmark de petits et grands tenseurs représentatifs ; ne signalez pas les performances de runs sanitizer.

Décision par défaut :

Une ou deux multiplications 64-bit de configuration : préférez le cast local.
L'indexation per-élément peut dépasser 32 bits : préférez le dispatch index_t.
La famille de kernels existante dispatche déjà index_t : étendez le motif existant.
Changer de nombreux kernels spécialisés scalaires : considérez la taille binaire avant de templater tous.

Tests pour les correctifs d'index large

Ajoutez une régression qui traverse la limite exacte qui a échoué :

Utilisez le plus petit dtype et la taille de sortie qui exercent toujours le débordement.
Assertez tensor.numel() > torch.iinfo(torch.int32).max ou assertez la limite d'offset spécifique.
Échantillonnez des valeurs à la fois en dessous et au-dessus de la limite ; évitez les comparaisons CPU de tenseur entier pour les tenseurs énormes.

Exécutez le repro original sous compute-sanitizer pour les bugs mémoire :

CUDA_LAUNCH_BLOCKING=1 PYTORCH_NO_CUDA_MEMORY_CACHING=1 compute-sanitizer --tool memcheck --error-exitcode=99 <python> repro.py

Pour les tests PyTorch, préférez ajouter la régression près des tests de pooling/indexing liés et gardez les cas coûteux avec @largeTensorTest et le décorateur de device pertinent.

Checklist de révision

[ ] L'expression débordante exacte est identifiée.
[ ] La mathématique 64-bit est limitée aux expressions qui en ont besoin, ou le kernel est templateé quand l'index chaud en a besoin.
[ ] canUse32BitIndexMath considère tous les tenseurs dont les offsets utilisent le type sélectionné.
[ ] Les dimensions de lancement CUDA tiennent toujours dans les limites du matériel.
[ ] Le test échoue avant le correctif et passe après, ou le rapport explique pourquoi l'échec pré-correctif n'a pas été réexécuté.
[ ] L'impact taille binaire/performance est mentionné si le nouveau dispatch index_t duplique les kernels.

cuda-index-width