Aller au contenu

4 · Les primitives de warp

Communiquer entre threads sans jamais toucher la mémoire. Quelques cycles au lieu de quelques dizaines, et c'est la base de toutes les réductions rapides.

4.1 Pourquoi ça existe

Les 32 threads d'un warp vivent dans la même partition d'un même SM. Leurs registres sont physiquement dans le même banc. Il n'y a donc aucune raison matérielle de passer par la mémoire partagée pour qu'ils échangent une valeur — il suffit d'un chemin de données entre voies de registres.

C'est ce que fournissent les instructions de shuffle.

Chemin Latence approximative
Registre → registre (__shfl_*) ~3-5 cycles
Via mémoire partagée (écriture + __syncwarp + lecture) ~40-60 cycles
Via mémoire globale ~500 cycles

4.2 Les shuffles

Toutes les variantes suivent la même forme :

T __shfl_sync(unsigned mask, T var, int srcLane, int width = 32);

Chaque thread fournit sa valeur var ; chaque thread reçoit la var du thread désigné.

Fonction Le thread l reçoit la valeur de…
__shfl_sync(m, v, s) thread s (diffusion si s est uniforme)
__shfl_up_sync(m, v, d) thread l − d (inchangé si l − d < 0)
__shfl_down_sync(m, v, d) thread l + d (inchangé si l + d ≥ 32)
__shfl_xor_sync(m, v, k) thread l XOR k — un « papillon »

Le masque

Le premier argument indique quels threads participent. Tous les threads listés dans le masque doivent atteindre l'instruction, sinon le comportement est indéfini.

__shfl_down_sync(0xffffffff, x, 16);   // les 32 threads participent
__shfl_down_sync(0x0000ffff, x,  8);   // seuls les 16 premiers

Le masque n'est pas décoratif

Depuis Volta, les threads d'un warp peuvent être à des points différents du programme. Le masque dit au matériel qui attendre. Un masque 0xffffffff dans un warp où seuls 20 threads sont actifs (à cause d'une garde de bord) produit un blocage ou un résultat faux.

Si vous ne pouvez pas garantir que le warp est plein, utilisez :

unsigned m = __activemask();

...mais sachez que c'est un pis-aller : la bonne réponse est de structurer le code pour que les warps soient toujours complets.

4.3 Réduction sur un warp

Le motif le plus utilisé du chapitre.

__device__ __forceinline__ float reduce_add_warp(float x) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        x += __shfl_down_sync(0xffffffff, x, offset);
    }
    return x;   // valide uniquement dans la lane 0
}

Cinq étapes (\(\log_2 32\)), aucun accès mémoire, aucun __syncthreads().

Variante butterfly, avec __shfl_xor_sync, qui laisse le résultat dans tous les threads :

__device__ __forceinline__ float allreduce_add_warp(float x) {
    #pragma unroll
    for (int k = 16; k > 0; k >>= 1) {
        x += __shfl_xor_sync(0xffffffff, x, k);
    }
    return x;   // valide dans les 32 lanes
}

C'est la version à préférer quand tous les threads ont besoin du résultat — par exemple pour normaliser après un softmax.

Depuis Ampere, il existe des versions matérielles pour les entiers :

int s = __reduce_add_sync(0xffffffff, x);     // int / unsigned uniquement
int m = __reduce_max_sync(0xffffffff, x);
int a = __reduce_and_sync(0xffffffff, x);

Une instruction au lieu de cinq. Elles ne couvrent malheureusement pas les flottants.

4.4 Les votes

unsigned __ballot_sync(unsigned mask, int pred);
int      __all_sync   (unsigned mask, int pred);
int      __any_sync   (unsigned mask, int pred);
unsigned __match_any_sync(unsigned mask, T valeur);

__ballot_sync renvoie un masque 32 bits où le bit \(l\) vaut 1 si le prédicat est vrai pour la lane \(l\). C'est la brique de la compaction de flux :

// Chaque thread a un élément ; on veut ne garder que ceux qui passent le test,
// compactés sans trou, dans un tableau de sortie.
__device__ void compacter(float valeur, bool garder,
                          float* sortie, int* compteur) {
    unsigned m    = __ballot_sync(0xffffffff, garder);
    int      rang = __popc(m & ((1u << (threadIdx.x & 31)) - 1));  // rang local
    int      total= __popc(m);

    int base;
    if ((threadIdx.x & 31) == 0)                 // un seul atomique par warp
        base = atomicAdd(compteur, total);
    base = __shfl_sync(0xffffffff, base, 0);     // diffusé à tout le warp

    if (garder) sortie[base + rang] = valeur;
}

Ce code illustre une technique essentielle : l'agrégation d'atomiques au niveau du warp. Au lieu de 32 atomicAdd en conflit sur la même adresse, on en fait un seul. Sur des motifs à forte contention, le gain dépasse couramment un facteur 10.

__popc compte les bits à 1 (population count), en une instruction matérielle.

À retenir

Chaque fois que vous écrivez atomicAdd sur une adresse partagée par tout un warp, demandez-vous si vous pouvez agréger d'abord avec __ballot_sync + __popc, ou avec une réduction de warp. C'est presque toujours possible et presque toujours rentable.

4.5 Les groupes coopératifs

Introduits en CUDA 9, les cooperative groups offrent une abstraction typée au-dessus de tout cela.

#include <cooperative_groups.h>
#include <cooperative_groups/reduce.h>
namespace cg = cooperative_groups;

__global__ void exemple(float* donnees) {
    cg::thread_block bloc = cg::this_thread_block();
    cg::thread_block_tile<32> warp = cg::tiled_partition<32>(bloc);

    float x = donnees[bloc.thread_rank()];

    // Réduction sur le warp, portable et lisible
    float somme = cg::reduce(warp, x, cg::plus<float>());

    if (warp.thread_rank() == 0) { /* ... */ }

    bloc.sync();                    // équivaut à __syncthreads()
}

Avantages :

  • lisibilité : warp.shfl_down(x, 16) est plus clair que __shfl_down_sync(0xffffffff, x, 16) ;
  • composition : on peut partitionner en tuiles de 8 ou 16 threads (tiled_partition<8>) et écrire du code générique ;
  • le masque est géré par le type.

Le groupe le plus important pour la suite est le grid_group :

cg::grid_group grille = cg::this_grid();
grille.sync();     // ← barrière sur TOUS les blocs

C'est la seule synchronisation inter-blocs officiellement supportée par CUDA. Elle exige :

  • un lancement par cudaLaunchCooperativeKernel ;
  • que tous les blocs soient simultanément résidents, donc une grille dont la taille est bornée par cudaOccupancyMaxActiveBlocksPerMultiprocessor × nb_SM.

Le coût réel de grid.sync()

Une barrière de grille est implémentée par des atomiques en mémoire globale et coûte quelques microsecondes. Le monokernel Kog mesure sur MI300X une synchronisation de grille naïve à 7,59–7,88 µs, réduite à 0,80–0,93 µs avec une technique de sentinelle. Ils attribuent à la synchronisation de grille « environ 35 % du temps total de génération d'un jeton » dans l'implémentation de départ.

Autrement dit : grid.sync() résout le problème de la barrière globale, mais reste bien trop coûteux pour être appelé cent fois par jeton. C'est précisément le raisonnement qui conduit aux megakernels, où l'on remplace la barrière globale par des dépendances fines.

Depuis Hopper, il existe aussi le cluster :

cg::cluster_group cluster = cg::this_cluster();
cluster.sync();                      // barrière sur les blocs du cluster
float* voisin = cluster.map_shared_rank(ptr_partage, rang);  // DSMEM

Traité au chapitre CUDA moderne.

4.6 Un exemple complet : softmax sur une ligne

Assemblons. Le softmax numériquement stable d'un vecteur \(\mathbf{x}\) :

\[ \text{softmax}(\mathbf{x})_i = \frac{\exp(x_i - \max_j x_j)}{\sum_k \exp(x_k - \max_j x_j)} \]

La soustraction du maximum est ce qui empêche \(\exp\) de déborder. Il faut donc deux réductions : un max, puis une somme.

// Un warp traite une ligne de n <= 1024 éléments
__global__ void softmax_ligne(const float* entree, float* sortie, int n) {
    int ligne = blockIdx.x;
    int lane  = threadIdx.x;              // 0..31, un warp par bloc
    const float* x = entree + ligne * n;
    float*       y = sortie + ligne * n;

    // Passe 1 : maximum
    float m = -INFINITY;
    for (int i = lane; i < n; i += 32) m = fmaxf(m, x[i]);
    #pragma unroll
    for (int k = 16; k > 0; k >>= 1)
        m = fmaxf(m, __shfl_xor_sync(0xffffffff, m, k));
    // m est maintenant le max de la ligne, dans TOUS les threads

    // Passe 2 : somme des exponentielles
    float s = 0.0f;
    for (int i = lane; i < n; i += 32) s += __expf(x[i] - m);
    #pragma unroll
    for (int k = 16; k > 0; k >>= 1)
        s += __shfl_xor_sync(0xffffffff, s, k);

    // Passe 3 : normalisation
    float inv = 1.0f / s;
    for (int i = lane; i < n; i += 32) y[i] = __expf(x[i] - m) * inv;
}

Trois lectures de x depuis la mémoire globale. Si \(n\) est petit, on peut les garder en registres et n'en faire qu'une — c'est l'idée du softmax en ligne qui est au cœur de FlashAttention (IA · FlashAttention).

__expf contre expf

__expf est la version intrinsèque rapide, calculée par la SFU avec une précision d'environ 2 ULP. expf est la version conforme, plus lente d'un facteur 3 à 5. Pour un softmax en apprentissage profond, __expf suffit largement.

Fait notable : sur Blackwell, la SFU devient un goulot d'étranglement pour l'attention, parce que le débit des tensor cores a doublé alors que celui de la SFU n'a pas bougé. FlashAttention-4 calcule donc les exponentielles par approximation polynomiale sur les unités FMA plutôt que sur la SFU (arXiv:2603.05451). C'est un excellent exemple de co-conception algorithme/matériel.

Résumé du chapitre

À retenir

  • Les shuffles échangent des registres entre threads d'un warp en ~4 cycles, sans toucher la mémoire.
  • __shfl_down_sync pour une réduction vers la lane 0, __shfl_xor_sync pour une réduction visible par tous.
  • Le masque _sync est obligatoire depuis Volta et n'est pas décoratif.
  • __ballot_sync + __popc permet d'agréger les atomiques au niveau du warp : un atomicAdd au lieu de 32.
  • Les cooperative groups offrent la même chose en plus lisible et composable, plus la barrière de grille — qui existe mais coûte plusieurs microsecondes.
  • Le softmax stable exige deux réductions ; réduire ce coût est l'idée de départ de FlashAttention.

Vérifiez que vous avez compris

Pourquoi __shfl_down_sync(0xffffffff, x, 16) puis ..., 8) etc. et pas l'inverse (1, 2, 4, 8, 16) ?

Les deux fonctionnent pour une somme. La descente 16→1 place le résultat dans la lane 0 ; la montée 1→16 le placerait aussi en lane 0 mais avec un motif d'accès différent.

La convention 16→1 est préférée parce qu'elle correspond à l'arbre binaire classique et parce qu'à la première étape, la moitié des threads (16-31) n'ont plus de travail utile : le compilateur peut parfois en tirer parti. En pratique la différence est nulle ; c'est une convention.

Vous avez un bloc de 256 threads et voulez une réduction sur tout le bloc. Comment procéder ?

En deux niveaux :

__device__ float reduce_bloc(float x) {
    __shared__ float partiels[32];       // au plus 32 warps par bloc
    int lane = threadIdx.x & 31;
    int warp = threadIdx.x >> 5;

    x = reduce_add_warp(x);              // 1. réduction intra-warp
    if (lane == 0) partiels[warp] = x;   // 2. un partiel par warp
    __syncthreads();

    // 3. le premier warp réduit les partiels
    x = (threadIdx.x < blockDim.x / 32) ? partiels[lane] : 0.0f;
    if (warp == 0) x = reduce_add_warp(x);
    return x;                            // valide dans le thread 0
}

Pour 256 threads : 8 warps → 8 partiels → une dernière réduction de warp. Un seul __syncthreads() au lieu de \(\log_2 256 = 8\) dans une version naïve en mémoire partagée.

Dans la compaction de flux, pourquoi diffuser base avec __shfl_sync plutôt que de le stocker en mémoire partagée ?

Parce que c'est plus rapide (4 cycles contre ~40) et parce que cela ne consomme aucune mémoire partagée — qui est une ressource limitant l'occupancy.

Notez aussi qu'il n'y a pas besoin de __syncwarp() avant le shuffle : le _sync de __shfl_sync fait office de barrière au niveau du warp.

Chapitre suivant : 5 · Réductions et scans

Sources de ce chapitre