Aller au contenu

3 · La mémoire partagée

Le chapitre qui transforme un programmeur GPU débutant en programmeur GPU. Tout ce qui suit dans ce cours — GEMM, FlashAttention, megakernels — est une variation sur le motif présenté ici.

3.1 Le problème que ça résout

Reprenons la multiplication matricielle naïve \(\mathbf{C} = \mathbf{A}\mathbf{B}\) avec des matrices \(N \times N\) :

__global__ void matmul_naif(const float* A, const float* B, float* C, int N) {
    int lig = blockIdx.y * blockDim.y + threadIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    if (lig >= N || col >= N) return;

    float acc = 0.0f;
    for (int k = 0; k < N; ++k) {
        acc += A[lig * N + k] * B[k * N + col];
    }
    C[lig * N + col] = acc;
}

Comptons.

Opérations : chaque thread fait \(N\) multiplications et \(N\) additions, soit \(2N\) opérations. Pour \(N^2\) threads : \(2N^3\) opérations au total.

Octets lus depuis la mémoire globale : chaque thread lit \(2N\) float, soit \(8N\) octets. Total : \(8N^3\) octets.

Intensité arithmétique :

\[ I = \frac{2N^3}{8N^3} = \frac{1}{4} \text{ opération par octet} \]

Sur un H100, dont le ratio machine en FP32 est \(67 \times 10^{12} / 3{,}35 \times 10^{12} = 20\) opérations par octet, on est 80 fois en dessous du seuil. Le noyau est massivement limité par la mémoire, et il plafonnera à :

\[ \frac{3{,}35 \times 10^{12} \text{ o/s}}{4} = 0{,}84 \text{ TFLOPS} \]

soit 1,2 % du pic FP32. Et ce calcul est optimiste : il suppose une coalescence parfaite, alors que B[k * N + col] est lue avec un pas de \(N\).

Le vrai gaspillage

L'élément A[lig][k] est lu par les \(N\) threads de la ligne lig. L'élément B[k][col] est lu par les \(N\) threads de la colonne col. Chaque élément est donc lu \(N\) fois depuis la HBM. Toute la question est : comment ne le lire qu'une fois ?

3.2 Le motif fondamental

┌─────────────────────────────────────────────┐
│ 1. Chaque thread charge un élément de la    │
│    tuile depuis la mémoire globale          │  ← accès coalescés
│    vers la mémoire partagée                 │
├─────────────────────────────────────────────┤
│ 2. __syncthreads()                          │  ← tout le monde a fini
├─────────────────────────────────────────────┤
│ 3. Chaque thread lit plusieurs fois dans    │
│    la tuile partagée et accumule            │  ← accès rapides
├─────────────────────────────────────────────┤
│ 4. __syncthreads()                          │  ← avant d'écraser
├─────────────────────────────────────────────┤
│ 5. Tuile suivante                           │
└─────────────────────────────────────────────┘

Les deux __syncthreads() sont tous les deux obligatoires et pour des raisons différentes :

  • le premier garantit que la tuile est complètement écrite avant qu'on la lise ;
  • le second garantit que tout le monde a fini de lire avant qu'on écrase la tuile à l'itération suivante.

Oublier le second est un bug classique qui ne se manifeste que sous charge.

3.3 La GEMM par tuiles

#define T 32   // taille de tuile

__global__ void matmul_tuiles(const float* A, const float* B, float* C, int N) {
    __shared__ float As[T][T];
    __shared__ float Bs[T][T];

    int tx = threadIdx.x, ty = threadIdx.y;
    int lig = blockIdx.y * T + ty;
    int col = blockIdx.x * T + tx;

    float acc = 0.0f;

    for (int t = 0; t < N / T; ++t) {
        // 1. Chargement coopératif
        As[ty][tx] = A[lig * N + (t * T + tx)];
        Bs[ty][tx] = B[(t * T + ty) * N + col];

        __syncthreads();                      // 2.

        // 3. Calcul sur la tuile
        for (int k = 0; k < T; ++k) {
            acc += As[ty][k] * Bs[k][tx];
        }

        __syncthreads();                      // 4.
    }

    C[lig * N + col] = acc;
}

Recomptons l'intensité arithmétique.

Par bloc : on charge \(2 T^2\) float par itération, soit \(8T^2\) octets, et on effectue \(2T^3\) opérations. Avec \(N/T\) itérations :

  • octets : \(8T^2 \times (N/T) = 8TN\) par bloc, et il y a \((N/T)^2\) blocs, soit \(8N^3/T\) octets au total ;
  • opérations : \(2N^3\) inchangé.
\[ I = \frac{2N^3}{8N^3/T} = \frac{T}{4} \]

Pour \(T = 32\) : \(I = 8\) opérations par octet. On a gagné un facteur 32.

La leçon générale

L'intensité arithmétique d'un calcul par tuiles est proportionnelle à la taille de la tuile. C'est pour cela que toute l'ingénierie GPU moderne consiste à faire tenir la plus grande tuile possible en mémoire rapide — et pourquoi Blackwell a ajouté 256 Ko de TMEM.

Avec \(I = 8\) et un seuil machine de 20, on reste limité par la mémoire, mais plus que d'un facteur 2,5 au lieu de 80. Pour franchir le seuil, il faut soit des tuiles plus grandes (impossible : \(32\times32\) float × 2 = 8 Ko, on pourrait aller à \(T=96\) mais 1 024 threads maximum par bloc l'empêche), soit plusieurs éléments par thread.

3.4 Le registre comme troisième niveau de tuile

C'est l'astuce qui fait passer de 20 % à 90 % de cuBLAS. Chaque thread calcule non pas 1 élément de \(\mathbf{C}\), mais un petit bloc \(R \times R\), gardé en registres.

#define BM 128   // tuile en lignes de C, par bloc
#define BN 128   // tuile en colonnes de C, par bloc
#define BK   8   // profondeur de la tuile
#define TM   8   // lignes de C par thread
#define TN   8   // colonnes de C par thread
// 128*128 / (8*8) = 256 threads par bloc

__global__ void matmul_registres(const float* A, const float* B,
                                 float* C, int N) {
    __shared__ float As[BK][BM];    // stockée transposée
    __shared__ float Bs[BK][BN];

    float acc[TM][TN] = {0.0f};     // ← 64 registres par thread
    float regA[TM], regB[TN];

    for (int t = 0; t < N; t += BK) {
        // chargement coopératif de As et Bs (détails omis)
        charger_tuiles(A, B, As, Bs, t, N);
        __syncthreads();

        for (int k = 0; k < BK; ++k) {
            // charger une colonne de As et une ligne de Bs en registres
            for (int i = 0; i < TM; ++i) regA[i] = As[k][/* ... */ + i];
            for (int j = 0; j < TN; ++j) regB[j] = Bs[k][/* ... */ + j];

            // produit extérieur : TM*TN FMA pour TM+TN lectures
            for (int i = 0; i < TM; ++i)
                for (int j = 0; j < TN; ++j)
                    acc[i][j] += regA[i] * regB[j];
        }
        __syncthreads();
    }
    // écriture de acc dans C (détails omis)
}

Le cœur du gain est dans le produit extérieur : pour \(TM + TN = 16\) lectures en mémoire partagée, on effectue \(TM \times TN = 64\) FMA. Le ratio calcul/mémoire-partagée passe de 1:2 à 4:1.

L'intensité arithmétique vis-à-vis de la mémoire globale devient \(I \approx BM \cdot BN / (2(BM + BN)) \cdot 1/4 \approx 8\) pour cette configuration, mais le vrai gain est que la mémoire partagée n'est plus le goulot.

Le coût en registres

acc[8][8] = 64 registres, plus regA, regB et les indices : on approche les 100 registres par thread, donc ~20 warps résidents par SM (31 % d'occupancy). C'est voulu. C'est exactement le régime « faible occupancy, fort parallélisme d'instructions » décrit par Volkov.

Le déroulement complet de cette optimisation, du noyau naïf à 95 % de cuBLAS, est détaillé dans IA · GEMM et dans le worklog de Simon Boehm, qui est la référence pédagogique du domaine.

3.5 Allocation statique et dynamique

Statique — la taille est connue à la compilation :

__shared__ float tuile[32][32];

Dynamique — la taille est passée au lancement :

extern __shared__ char partagee[];        // un seul tableau, typé char

__global__ void noyau(int n) {
    float* a = reinterpret_cast<float*>(partagee);
    float* b = a + n;                     // découpage manuel
    ...
}

// lancement
size_t octets = 2 * n * sizeof(float);
noyau<<<grille, bloc, octets>>>(n);

Au-delà de 48 Ko, il faut le demander

Par défaut, un bloc ne peut allouer que 48 Ko de mémoire partagée, quelle que soit la carte. Pour aller jusqu'aux 227 Ko de Hopper/Blackwell :

cudaFuncSetAttribute(noyau,
    cudaFuncAttributeMaxDynamicSharedMemorySize, 227 * 1024);

Sans cet appel, le lancement échoue avec cudaErrorInvalidValue — souvent ignoré si on ne vérifie pas les erreurs.

3.6 Les conflits de banc, en pratique

Rappel du chapitre 4 des fondations : 32 bancs de 4 octets, entrelacés.

Le cas typique dans une transposition :

__shared__ float tuile[32][32];

// Écriture : tuile[ty][tx], tx varie sur le warp → bancs 0..31 → OK
tuile[threadIdx.y][threadIdx.x] = entree[...];
__syncthreads();

// Lecture transposée : tuile[tx][ty], ty constant sur le warp
//   adresse = 32*tx + ty → banc = (32*tx + ty) mod 32 = ty
//   → les 32 threads visent LE MÊME banc → conflit à 32 voies
sortie[...] = tuile[threadIdx.x][threadIdx.y];

Le remède, une seule colonne de plus :

__shared__ float tuile[32][33];   // ← +33 au lieu de 32
// adresse = 33*tx + ty → banc = (33*tx + ty) mod 32 = (tx + ty) mod 32
// → tous distincts quand tx varie. Conflit résolu.

Coût : \(32 \times 4 = 128\) octets de mémoire partagée gaspillés. Gain : facteur 32 sur cet accès.

Le swizzling, l'alternative moderne

Le padding gaspille de la mémoire et casse l'alignement nécessaire à TMA et wgmma. Les bibliothèques modernes (CUTLASS, ThunderKittens) utilisent plutôt le swizzling : une permutation XOR des adresses qui élimine les conflits sans ajouter d'octets.

\[\text{adresse\_reelle} = \text{adresse} \oplus \big((\text{adresse} \gg s) \,\&\, m\big)\]

C'est plus complexe, et c'est exactement le genre de chose qu'on délègue à une bibliothèque. Voir Écosystème · CUTLASS.

3.7 La mémoire partagée comme ressource globale

Un point souvent négligé, et qui devient central dans les megakernels.

La mémoire partagée est une ressource par SM, pas par bloc. Si votre noyau demande 100 Ko, un SM Hopper (227 Ko) ne peut héberger que 2 blocs. Cela plafonne l'occupancy quelle que soit la consommation de registres.

\[ \text{blocs résidents} \le \left\lfloor \frac{\text{mém. partagée du SM}}{\text{mém. partagée par bloc}} \right\rfloor \]

Dans un megakernel, où un seul bloc persistant occupe un SM pendant toute la durée de l'exécution, la mémoire partagée devient un espace à gérer dynamiquement. C'est exactement ce que fait l'allocateur par pages de Hazy Research : les 213 Ko de mémoire partagée d'un H100 sont découpés en treize pages de 16 Ko, que les « instructions » demandent et libèrent. MPK utilise des pages de 32 Ko.

Retenez cette idée : au-delà d'un certain niveau de sophistication, la mémoire partagée cesse d'être une variable et devient un tas.

Résumé du chapitre

À retenir

  • Le motif universel : charger coopérativement → __syncthreads() → calculer → __syncthreads() → tuile suivante. Les deux barrières sont nécessaires.
  • L'intensité arithmétique d'un calcul par tuiles est proportionnelle à la taille de tuile : \(I = T/4\) pour une GEMM en float.
  • Le vrai gain vient du troisième niveau de tuile, en registres : un produit extérieur \(TM \times TN\) donne \(TM\cdot TN\) FMA pour \(TM+TN\) lectures partagées.
  • Au-delà de 48 Ko par bloc, cudaFuncSetAttribute est obligatoire.
  • Les conflits de banc se règlent par padding (simple, gaspilleur) ou par swizzling (complexe, gratuit, requis par TMA).
  • La mémoire partagée limite le nombre de blocs résidents. Dans un megakernel, elle devient un tas à gérer.

Vérifiez que vous avez compris

Que se passe-t-il si on supprime le second __syncthreads() de la GEMM par tuiles ?

Un thread rapide peut entamer l'itération \(t+1\) et écraser As[ty][tx] pendant qu'un thread lent lit encore cette case pour l'itération \(t\). Résultat : corruption silencieuse, non déterministe, qui n'apparaît souvent qu'avec de grandes matrices ou sur une autre carte.

C'est le type de bug qui coûte deux jours. compute-sanitizer --tool racecheck le détecte.

Pourquoi stocker As transposée (As[BK][BM] plutôt que As[BM][BK]) dans la version à registres ?

Parce que la boucle interne lit une colonne de As pour un k donné : As[k][i] pour i = 0..TM-1. Avec le stockage As[BK][BM], ces éléments sont contigus en mémoire partagée, donc dans des bancs différents : pas de conflit, et on peut même les lire en float4.

Avec As[BM][BK], les mêmes éléments seraient espacés de BK : conflits de banc garantis.

Un noyau utilise 100 Ko de mémoire partagée par bloc sur H100 et vous mesurez une occupancy de 12,5 %. D'où vient ce chiffre ?

227 Ko / 100 Ko = 2 blocs résidents par SM. Si chaque bloc fait 128 threads, cela donne 256 threads sur les 2 048 possibles, soit exactement 12,5 %.

Est-ce grave ? Pas forcément. Si chaque thread a beaucoup de parallélisme d'instructions et que les tuiles sont grandes, ce régime peut être optimal — c'est celui des noyaux CUTLASS. Ce qui compte est le débit obtenu, pas l'occupancy. Voir Performance · Occupancy.

Chapitre suivant : 4 · Les primitives de warp

Sources de ce chapitre