Qu'est-ce qu'un kernel, au juste ?

Deux nombres sur un GPU

Vous savez déjà faire c = a + b en Python. Vous l'avez fait des milliers de fois. Ici, nous allons en parler avec les tenseurs PyTorch. Un tenseur est simplement un tableau de nombres. Le placer sur un GPU signifie que ce tableau réside dans la mémoire du GPU au lieu de la mémoire CPU classique. Quand a et b sont deux tenseurs sur un GPU, cette ligne s'exécute suffisamment vite pour que vous n'y pensiez même pas.

Rétrécissons maintenant. Disons que a et b sont deux flottants uniques, tous deux situés sur le GPU. Même ligne. c = a + b. Qu'est-ce qui s'exécute réellement sur la puce ?

La réponse est un kernel. Dans ce contexte, un kernel est un petit programme que le GPU exécute sur un ensemble de données. Pas le kernel du système d'exploitation qui démarre votre ordinateur portable. Pas le kernel mathématique d'un manuel d'algèbre linéaire. Le mot est beaucoup réutilisé, et ce n'est pas de votre faute. Dans l'univers des GPU, un kernel signifie simplement : une petite fonction que le GPU est chargé d'exécuter, immédiatement, en parallèle, sur les données que vous lui avez fournies.

À la fin de cet article, vous serez capable de regarder un extrait de code PyTorch et de compter combien de kernels le GPU va exécuter. Cela peut sembler un petit truc, et ça l'est, mais c'est aussi le premier pas pour se débarrasser de l'impression que le GPU est une boîte noire, celle où votre modèle est lent et vous ne savez pas pourquoi. Toute question que vous pouvez vous poser sur les performances du GPU finit par revenir à « quels kernels se sont exécutés, et que faisaient-ils ? ». C'est donc par là que nous commençons.

Votre premier kernel

Grossissons un peu a et b. Cette fois, des tenseurs de longueur 8. Toujours une seule ligne PyTorch : c = a + b.

Lorsque vous exécutez ceci, votre CPU (la machine qui exécute le Python proprement dit) dit au GPU : hé, va exécuter ce programme sur ces données. Cette instruction s'appelle un **lancement\. Ce qui est lancé est un kernel : un programme, prêt à être exécuté. Les lancements eux-mêmes sont peu coûteux, de l'ordre de quelques microsecondes chacun. Ce qui se trouve autour du lancement (les données allant vers le GPU, les résultats en revenant) est là où se trouve le coût réel, et c'est ce que nous allons compter au fur et à mesure.

À l'intérieur du kernel, le travail réel est effectué par de minuscules ouvriers appelés threads. Un GPU en a des milliers à disposition. Pour notre addition de longueur 8, 8 threads prennent le travail en charge : le thread 0 gère l'élément 0, le thread 1 gère l'élément 1, et ainsi de suite jusqu'au thread 7. Chaque thread exécute le même petit programme : lire un élément de a, lire l'élément correspondant de b, les additionner, écrire le résultat dans c.

(En pratique, le GPU lance les threads en groupes de taille fixe appelés warps, toujours 32 threads sur les cartes NVIDIA, et masque les supplémentaires lorsque votre tableau ne se divise pas exactement. On peut ignorer cela pour l'instant.)

Nous avons donc une ligne PyTorch, un lancement, un kernel, 8 threads effectuant 8 additions. Comptons maintenant ce qui a réellement traversé la puce. Pour effectuer l'addition, chaque thread a besoin de son élément de a et de son élément de b. Cela fait 8 lectures de a et 8 lectures de b. Ensuite, chaque thread écrit son résultat dans c. Cela fait 8 écritures.

Ces lectures et écritures vont vers la grande mémoire située juste à côté de la puce du GPU. Sur les cartes pour centres de données (A100, H100), cette mémoire s'appelle HBM (mémoire à haute bande passante). Sur les cartes grand public (RTX 4090) et les T4 de type Colab sur lesquelles les gens sont susceptibles d'essayer, on l'appelle GDDR. Dans les deux cas, c'est une mémoire rapide située à côté de la puce, et nous l'appellerons simplement mémoire GPU. Elle est rapide, mais elle n'est pas gratuite, et chaque accès a un coût.

Un kernel = un lancement = un passage sur les données. Quoi que le kernel fasse à l'intérieur de son corps, les lectures et écritures à ses extrémités (les accès à la mémoire GPU pour récupérer les entrées, l'accès pour écrire les sorties) sont la partie qui coûte. C'est toute la structure.

Rien de tout cela ne change lorsque les tenseurs deviennent plus grands. Même ligne PyTorch, même kernel unique, juste plus de threads. Si a et b ont un million d'éléments chacun, le GPU lance le même kernel avec un essaim plus important de threads. Le calcul augmente, les octets déplacés augmentent, le modèle mental ne change pas. Une ligne, un kernel.

Ce qui se passe entre deux opérations

1c = (a + b).relu()

Vous connaissez assez bien Python pour savoir qu'il s'agit de deux opérations, une addition puis une activation ReLU, chaînées. Dans un interpréteur, cela représente deux appels de fonction. Sur un GPU, en mode eager PyTorch, cela représente deux lancements de kernel : un pour l'addition, un pour la ReLU. Jusqu'ici, rien de surprenant.

Ce qui est réellement intéressant, c'est ce qui se passe entre les deux kernels.

Lorsque l'addition se termine, elle doit placer son résultat quelque part. Ce quelque part est la mémoire GPU. L'addition écrit un tableau intermédiaire complet (appelons-le tmp) en mémoire. Puis, un instant plus tard, la ReLU est lancée, et sa première tâche est de lire ce même tableau tmp depuis la mémoire. Elle lit l'intégralité du tableau, applique la ReLU à chaque élément, et écrit le résultat dans c.

Comptons le trafic mémoire pour ces deux kernels :

• L'addition : lit a, lit b, écrit tmp. Trois transferts de la taille d'un tableau.
• La ReLU : lit tmp, écrit c. Deux transferts supplémentaires.

Cinq transferts de la taille d'un tableau au total. Comparez cela à l'addition de longueur 8 seule de la section précédente, qui en nécessitait trois. Ajouter .relu() à la chaîne ne vous a pas seulement coûté le calcul de la ReLU. Cela vous a coûté un aller-retour complet du tableau à travers la mémoire GPU, car tmp a dû être écrit juste pour que le kernel suivant puisse le relire.

Rien n'a été mis en cache. tmp n'a pas pu rester dans un registre ou un cache local rapide. Il est allé dans la mémoire GPU (le type lent et éloigné) et en est revenu immédiatement. Les deux kernels sont des étrangers l'un pour l'autre. Ils doivent se passer les données via le seul support qu'ils connaissent tous les deux : la mémoire GPU.

Pourquoi PyTorch fait-il cela ? Parce qu'en mode eager, lorsque vous écrivez a + b, PyTorch l'exécute \immédiatement**. Il ne sait pas que vous alliez appeler .relu() ensuite. Chaque opération est dispatchée au moment où sa ligne Python s'exécute. Il n'y a pas de plan, pas d'anticipation. Chaque opération est indépendante, produit un tableau réel, et le transmet à l'opération suivante via la mémoire.

C'est le schéma à retenir. Chaque valeur intermédiaire dans votre code PyTorch est physiquement écrite dans la mémoire GPU et relue par l'opération suivante. Chacune. C'est ce que le « nombre de kernels » mesure réellement. Chaque kernel supplémentaire est un autre aller-retour que vos données doivent effectuer à travers la mémoire GPU.

Fusion : deux opérations, un seul kernel

Imaginez un seul kernel qui fait tout cela en une seule fois : lit son élément de a, lit son élément de b, les additionne, applique la ReLU au résultat (tout cela à l'intérieur du kernel, sur un espace de travail temporaire minuscule par thread qui ne quitte jamais la puce), et seulement ensuite écrit la valeur finale dans c. L'intermédiaire (a + b) existe toujours, mais uniquement à l'intérieur du kernel, dans l'espace de travail privé de chaque thread. Il n'est jamais écrit dans la mémoire GPU. tmp, en tant que tableau réel, n'existe pas du tout.

Comptez les transferts maintenant. Lectures de a : 1 par élément. Lectures de b : 1 par élément. Écritures de c : 1 par élément. Trois transferts de la taille d'un tableau. Même calcul que la version à deux kernels, mais deux allers-retours en moins.

Pour une longueur de 8, c'est une erreur d'arrondi. Personne ne s'en soucie. Pour une longueur d'un million, ou de 100 millions, ces allers-retours supplémentaires deviennent une part importante du temps d'exécution, et l'horloge murale le reflète. \Pourquoi** le trafic mémoire finit par dominer ainsi est le sujet principal de l'article 2 de cette série, donc je vais laisser le « pourquoi » de côté ici. Le point immédiat est simplement : même calcul, moins de trajets, plus rapide en pratique.

Cette astuce, qui consiste à combiner des opérations qui auraient été des kernels séparés en un seul kernel afin que l'intermédiaire n'ait jamais à visiter la mémoire GPU, a un nom. Cela s'appelle la fusion. C'est tout le mot. C'est toute l'idée.

Maintenant, la partie délicate. Écrire ce kernel combiné à la main semble facile pour addition + ReLU. Deux opérations. Une ligne de « calcul » au milieu. Mais le code PyTorch réel a des dizaines d'opérations chaînées, chacune avec ses propres formes, types de données et règles de broadcasting. Écrire un kernel fusionné qui gère tout cela correctement est un véritable travail d'ingénierie. Vous n'écririez généralement pas ces kernels élémentaires courants à la main.

Bonne nouvelle : PyTorch est livré avec un outil qui effectue la réécriture pour vous, exactement pour ce type de cas, automatiquement. Il s'appelle torch.compile.

Vous avez presque certainement vu torch.compile utilisé comme ceci :

1model = torch.compile(model)

Une ligne. Quelqu'un sur Internet vous a dit que cela rendait les choses plus rapides. Voici ce qu'il fait réellement, en termes simples : au lieu d'exécuter vos opérations une par une comme le fait le mode eager, torch.compile capture les opérations sur les tenseurs que votre fonction effectue, cherche des opportunités de les combiner et génère du code optimisé. Les appels ultérieurs qui correspondent aux mêmes hypothèses peuvent réutiliser ce code.

La fusion que nous avons effectuée sur papier ci-dessus (addition et ReLU, partageant un seul kernel, tmp ne touchant jamais la mémoire) est exactement le genre de chose que torch.compile fera à votre code automatiquement, tant que les opérations sont suffisamment simples. Quand les gens disent que torch.compile « rend PyTorch plus rapide », c'est une grande partie de ce qu'ils veulent dire.

Pour les cas que torch.compile ne peut pas fusionner seul (opérations personnalisées qu'il ne reconnaît pas, réductions inhabituelles, dispositions mémoire étranges), quelqu'un doit encore écrire un kernel à la main. C'est à cela que servent des outils comme Triton et CUDA. Un article séparé.

Voyez-le par vous-même

Tout ce qui précède consistait à compter les kernels sur papier. Il est temps de les compter sur un vrai GPU. Si vous avez une machine avec un GPU CUDA à portée de main (une station de travail, un notebook Colab, une instance cloud), vous pouvez exécuter ceci vous-même en quelques minutes.

L'outil est torch.profiler. Il est intégré à PyTorch. Tout ce qu'il fait est d'enregistrer ce que le GPU a réellement fait pendant l'exécution de votre code, et de vous renvoyer un tableau lisible.

Étape 1 : la version eager.

Enveloppez la ligne à deux opérations dans une fonction pour avoir quelque chose à appeler :

1import torch
2from torch.profiler import profile, ProfilerActivity
3
4def add_relu(a, b):
5    return (a + b).relu()
6
7a = torch.randn(1_000_000, device="cuda")
8b = torch.randn(1_000_000, device="cuda")
9
10with profile(activities=[ProfilerActivity.CUDA]) as prof:
11    add_relu(a, b)
12    torch.cuda.synchronize()
13
14print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))

torch.cuda.synchronize() est simplement là pour s'assurer que le GPU a terminé avant que nous lisions les temps. Le travail du GPU s'exécute de manière asynchrone, et sans la synchronisation, vous mesureriez parfois la surcharge du lancement au lieu du travail réel du kernel.

Étape 2 : lisez la sortie.

La sortie réelle de votre profileur aura probablement plus de lignes que prévu. Un tas de lignes d'allocation mémoire et de comptabilité interne de PyTorch seront mélangées. Les lignes qui nous intéressent sont les kernels CUDA, les fonctions réelles que le GPU a exécutées. Recherchez les lignes avec kernel dans le nom. Les deux ressembleront à peu près à ceci :

1vectorized_elementwise_kernel<...BinaryFunctor_add...>    1    12us
2vectorized_elementwise_kernel<...threshold_kernel...>     1     9us

Les noms de modèles exacts changent entre les versions de PyTorch (la ReLU apparaît souvent sous le nom de threshold car c'est l'opération sous-jacente, et l'addition parfois sous le nom de CUDAFunctor_add). N'essayez pas d'analyser la totalité. Comptez simplement les lignes. Deux lignes. Deux kernels. Un pour l'addition, un pour la ReLU. Exactement ce que nous avions dit qu'il se passerait il y a une section.

Étape 3 : la version compilée.

Un changement d'une ligne. Enveloppez la fonction dans torch.compile :

1compiled = torch.compile(add_relu)

Voici un piège qu'il est bon de connaître avant de l'exécuter. Le \premier appel à une fonction `torch.compile` est lent, parfois étrangement lent, car c'est à ce moment-là que `torch.compile` fait son travail : analyser votre code, déterminer quoi fusionner, générer le kernel fusionné. Si vous profilez le tout premier appel, vous mesurez l'étape de compilation, pas le kernel. Donc le schéma est : appelez-la une fois pour la chauffer, jetez le résultat, puis** profilez.

1compiled(a, b)                # échauffement, à jeter
2torch.cuda.synchronize()
3
4with profile(activities=[ProfilerActivity.CUDA]) as prof:
5    compiled(a, b)
6    torch.cuda.synchronize()
7
8print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))

Étape 4 : relisez la sortie.

Cette fois-ci :

1triton_poi_fused_add_relu_0    1    14us

Une ligne. Un kernel. Le nom vous dit même ce qu'il a fait : addition et ReLU fusionnées. Même calcul qu'avant, un lancement au lieu de deux.

Vous venez de faire la chose dont parle cet article, en une phrase : vous avez demandé à PyTorch de combiner deux opérations en un seul kernel, vous avez regardé le profileur et vous avez confirmé que le nombre est passé de deux à un. La fusion, dans la nature, sur votre machine.

Si vous voulez voir cela de manière plus spectaculaire, essayez avec différentes tailles de tenseurs. Pour une longueur de 100, les deux versions s'exécutent si vite que la différence se perd dans le bruit. Pour une longueur de 10 millions ou 100 millions, la version compilée commence à prendre nettement l'avantage, car l'aller-retour que nous avons supprimé représente une part réelle du travail à cette échelle.

Compter les kernels n'est plus un conseil abstrait. Vous avez un moyen de vérifier.

Résumons tout cela

Voici l'essentiel en une seule fois.

Votre code PyTorch, lorsqu'il s'exécute sur un GPU, se transforme en une séquence de kernels. Chaque kernel est un lancement, un passage sur vos données, un aller-retour à travers la mémoire GPU pour récupérer les entrées et écrire les sorties. Les opérations simples deviennent un kernel. Les chaînes d'opérations deviennent par défaut un kernel \par opération**, avec des allers-retours des intermédiaires à travers la mémoire entre eux. torch.compile peut fusionner les chaînes simples pour vous afin que ces intermédiaires ne touchent jamais la mémoire. Moins de kernels signifie généralement moins de trafic mémoire. Et moins de trafic mémoire signifie généralement une exécution plus rapide.