Imputer l'énergie GPU au noyau qui l'a dépensée

J'ai passé un été à Oak Ridge à travailler là-dessus, et la question de départ est courte : le noyau qui a pris le plus de temps réel dans mon exécution n'était pas celui qui a coûté le plus d'énergie. Le profileur classait tout par le temps, avec assurance, et ce classement était tout simplement le mauvais dès lors que ce qu'on vous demande de réduire, c'est la facture électrique. Les clusters tournent de plus en plus sous un plafond de puissance plutôt que sous une cible de fréquence, donc l'énergie par solution devient le chiffre qui compte, et presque rien dans un flux HPC normal ne la rapporte par noyau. J'ai donc construit un outil qui le fait : un connecteur Kokkos Tools qui impute les joules à chaque région profilée sans toucher à l'application qu'il mesure.

Voici le résultat qu'il a produit sur une exécution, puis comment il y parvient.

$ export KOKKOS_TOOLS_LIBS=/opt/kp/libkp_gpu_energy.so
$ ./solver --mesh big.h5

region            calls     time(s)   energy(J)   J/call   avg(W)
-----------------------------------------------------------------
gemm_apply        12480       8.42      1936.1     0.155      230
spmv_matvec       49920      11.07      1421.8     0.028      128
halo_exchange     49920       3.91       402.7     0.008      103
-----------------------------------------------------------------
device idle baseline ~ 61 W  (subtracted for the J/call column)

Le produit matrice-vecteur creux a tourné le plus longtemps, onze secondes contre huit pour le bloc dense, et a tout de même coûté un tiers d'énergie en moins, parce qu'il est limité par la mémoire et laisse le GPU tirer environ la moitié de la puissance. Le temps me disait d'optimiser spmv_matvec. L'énergie me disait de regarder gemm_apply d'abord. Ce sont des consignes différentes, et jusqu'à l'existence de ce connecteur je ne voyais que la première.

Kokkos Tools, ou de l'instrumentation que vous n'avez pas à compiler

Kokkos annonce déjà ce qu'il fait. Chaque parallel_for, parallel_reduce et parallel_scan déclenche un callback de début avant de se lancer et un callback de fin une fois terminé, et on peut entourer des portions de code arbitraires de régions nommées avec des marqueurs push et pop. Un connecteur Kokkos Tools n'est qu'une bibliothèque partagée qui implémente ces callbacks, et on l'attache en pointant une variable d'environnement vers elle. Aucune recompilation de l'application, aucune annotation dans ses sources, aucun fork du code. On met KOKKOS_TOOLS_LIBS sur le chemin de la bibliothèque et le runtime la charge.

Cette propriété est toute la raison d'être de l'approche. Je pouvais prendre un solveur que je n'avais pas écrit, que personne ne veut me laisser modifier, et apprendre le coût énergétique de chacun de ses noyaux en chargeant une bibliothèque de plus à côté. Le connecteur écoute les événements que Kokkos émet déjà. La mesure suit le mouvement.

On ne peut pas lire l'énergie, seulement observer la puissance

La première version évidente lit le capteur de puissance au callback de début, le relit à la fin, et rapporte la moyenne fois la durée. Ça ne marche pas, et la raison pour laquelle ça ne marche pas est au cœur du problème. La bibliothèque de gestion de NVIDIA, NVML, expose nvmlDeviceGetPowerUsage, qui renvoie la puissance instantanée de la carte en milliwatts. Le piège est double. Ce capteur se met à jour à un rythme modeste, de l'ordre de la dizaine de hertz, et un noyau GPU peut facilement être plus court que l'intervalle entre deux mises à jour : début et fin renvoient alors souvent la même valeur périmée et la durée ne dit rien. Et même quand le noyau est assez long pour couvrir plusieurs mises à jour, deux lectures ponctuelles ne peuvent pas décrire une courbe qui monte et descend tout au long de la vie du noyau.

Le problème de fond, c'est que la puissance est la mauvaise grandeur à échantillonner aux bornes. La puissance est instantanée, des watts, un débit. Ce que vous payez, c'est de l'énergie, des joules, et l'énergie est l'intégrale de la puissance dans le temps. Deux lectures vous donnent deux hauteurs d'une courbe. La facture, c'est l'aire en dessous. Ma première version racontait n'importe quoi sur les noyaux courts, parfois même un écart négatif quand les deux lectures tombaient de part et d'autre d'une mise à jour du capteur, et c'était le signal pour cesser d'échantillonner au rythme du noyau et commencer à échantillonner au rythme de l'horloge.

Un thread dédié, une cadence fixe, et un trapèze

La solution, c'est de séparer entièrement l'échantillonnage des noyaux. Un thread en arrière-plan interroge le capteur de puissance à intervalle fixe, espacé de quelques millisecondes, et horodate chaque lecture, construisant une trace continue de l'évolution de la consommation de la carte sur toute l'exécution. Les callbacks de début et de fin ne lisent plus du tout la puissance. Ils enregistrent une fenêtre en temps réel, le moment où la région s'est ouverte et celui où elle s'est refermée. Pour obtenir l'énergie d'une région, le connecteur intègre la trace de puissance sur cette fenêtre avec la règle des trapèzes, en sommant les petits trapèzes entre échantillons consécutifs qui tombent à l'intérieur. Comme la même région est franchie des milliers de fois, ses joules s'accumulent sur chaque appel : c'est la colonne energy(J) ci-dessus, et la seule façon honnête de parler d'un noyau qui s'exécute en quelques dizaines de microsecondes.

Échantillonner sur l'horloge plutôt que sur le noyau, c'est ce qui rend les noyaux courts mesurables. Un seul lancement peut être trop bref pour capter ne serait-ce qu'une lecture fraîche du capteur, mais dix mille lancements sous une trace interrogée régulièrement déposent assez d'échantillons pour que l'agrégat soit solide. Le compromis, c'est un peu de surcoût dû au thread d'interrogation et un plancher de résolution fixé par l'intervalle d'échantillonnage, tous deux petits et, surtout, bornés et connus.

Ce que le chiffre est, et ce qu'il n'est pas

Je préfère énoncer les limites clairement plutôt que laisser le tableau suggérer plus de précision qu'il n'en a. NVML rapporte la puissance de la carte entière, pas par multiprocesseur de flux, donc c'est une imputation à l'échelle du GPU entier. Si deux noyaux s'exécutent en même temps sur le même appareil, sur des flux distincts, la trace ne peut pas dire lequel a tiré quel watt, et l'énergie du recouvrement ne peut pas être répartie proprement entre eux. Le chiffre inclut aussi la consommation au repos de l'appareil, les dizaines de watts qu'un GPU allumé dépense à ne rien faire : pour le coût marginal d'un noyau, on mesure donc une référence de repos appareil au calme et on la soustrait, c'est la ligne sous le tableau et le plancher dans le schéma. Rien de tout cela ne rend la mesure fausse. Cela la rend à l'échelle de l'appareil entier et honnête sur sa résolution, ce qui, pour classer les noyaux par énergie, suffit exactement.

Deux backends, deux questions différentes

NVML répond à une question avec beaucoup de précision : qu'a tiré ce GPU NVIDIA. C'est par carte, à la résolution du milliwatt, NVIDIA seulement, et il ne voit rien en dehors de la carte. Le connecteur a donc un second backend bâti sur Variorum, indépendant du fournisseur, qui lit la puissance au niveau du nœud et du socket, y compris le CPU via RAPL, la DRAM, et certains GPU, sur du matériel qui n'est pas NVIDIA. Les deux ne sont pas redondants. NVML vous dit ce que le GPU a tiré. Variorum vous dit ce que le nœud a tiré. Un noyau qui paraît bon marché sur la carte peut tout de même brasser assez de données pour faire chauffer le CPU et les contrôleurs mémoire autour de lui, et seule la vue au niveau du nœud l'attrape. On se tourne vers NVML quand on règle un noyau GPU isolément, et vers Variorum quand on veut la facture énergétique que la salle machine voit réellement.

Ce que les joules par noyau vous apportent

Une fois l'énergie imputée au noyau qui l'a dépensée, on peut enfin optimiser la grandeur qu'on vous facture vraiment, au lieu d'utiliser le temps comme approximation en espérant que les deux concordent. Ils ne concordent pas toujours, et c'est tout l'enjeu : le noyau le plus rapide n'est souvent pas le plus efficace en énergie, parce qu'aller vite peut vouloir dire faire tourner le silicium à son plafond de puissance, et un noyau plus lent limité par la mémoire peut être le moins cher à exécuter un million de fois. Rien de tout cela n'est visible sur une frise temporelle, et tout le devient une fois les joules posés à côté du nombre d'appels.

Le connecteur existe sous forme d'une petite pile de PR ouvertes en amont sur kokkos/kokkos-tools, le backend NVML et celui Variorum, et sur mes propres machines les joules par noyau alimentent le même tableau de bord que le reste de mon travail GPU, si bien qu'une exécution montre l'énergie par solution à côté de l'utilisation, plutôt que dans un log séparé que personne n'ouvre. La prochaine chose que je veux, c'est pousser la résolution sous la carte entière, parce que l'imputation à l'appareil entier est honnête mais grossière, et les flux concurrents méritent une réponse plus propre que celle que je peux donner aujourd'hui. Si vous avez mesuré l'énergie par solution à une résolution plus fine que l'appareil entier, ou trouvé une façon saine de répartir la facture entre noyaux qui se recouvrent, j'aimerais sincèrement comparer nos notes. Je suis aussi ouvert à des postes où c'est le quotidien, à partir de janvier 2027.