Blog

• Article suivant
• Page de blog

Booster un réseau de neurones en C grâce à l'assembleur

Et voilà, je me suis finalement mis à l'Intelligence Artificielle, et c'est arrivé bien plus tôt que prévu !

Si une voyante m'avait annoncé le mois dernier que j'intégrerais du réseau de neurone dans mon code avant la mi-mai, j'aurais détourné les yeux de mon XTerm sous XWM pour la dévisager avec un air chiffonné, pour ensuite de me rouler au sol en riant durant de longues minutes… Et pourtant elle aurait vu juste, cela s'est produit : Tout est parti d'un besoin de créer un générateur d'images pour une bibliothèque d'icônes il y a une paire de semaines, tandis que comme tout développeur moyen j'entamais tête baissée la bonne cinquantaine de fonctions chronophages à implémenter, je réalisais soudain qu'au 21^ème siècle ce genre de besogne devait impérativement être déléguée aux réseaux de neurones.

J'avais toutefois deux obstacles titanesques à franchir :

1. Je n'y connaissais rien.
2. Bienvenue chez Python.

Or si apprendre un nouveau domaine a toujours été un plaisir pour moi, il n'était pas question d'imposer à mes clients l'installation de lourds environnements Python pour exécuter leurs logiciels, ni d'abandonner mon code C++ et mon terminal Unix. Mais alors, comment se doter de l'IA quand on est un programmeur à l'ancienne, est-on bon pour la casse ? J'ai bien essayé de passer quelques heures à utiliser le très célèbre TensorFlow, mais j'étais incapable de créer le moindre graphe qui fonctionne, d'ailleurs personne ne semble savoir le faire en langage C sur les forums d'experts. Heureusement, il existe au moins un grand Seigneur Nerd en ce monde, celui-ci a eu l'audace de construire une petite API de Machine Learning en C, nommée « nn.h », et d'en faire un petit projet à porté éducative. Grâce à lui on peut enfin comprendre le principe des réseaux de neurones, et en fabriquer soi-même sur n'importe quel ordinateur sans rien installer de plus qu'un simple compilateur. Mes amis, voici venu le temps de considérer l'IA comme un outil facile et rapide à intégrer à nos petites moulinettes du quotidien ! Cependant, l'auteur insiste bien sur le fait que son code est inadapté pour un usage en production. Ce n'est pas gênant pour la partie apprentissage, mais comment optimiser le résultat d'un entraînement pour l'intégrer à un vrai projet professionnel ?

Nous pouvons facilement intervenir sur trois aspects : Le poids des neurones en octets grâce à la quantification, la dépendance à la bibliothèque <math.h> pour calculer la fonction sigmoïde grâce à une approximation, et la vitesse de calcul grâce à une implémentation de la multiplication matricielle avec une gestion fine des instructions SIMD. Et pour deux de ces points, nous aurons besoin d'appeler le bon vieux langage assembleur à la rescousse !

Réduire la taille par quantification

La quatification, dont vous entendez peut-être plus souvent parler en tant que « quantization », consiste à tronquer des données. Dans le cas des réseaux de neurones, le but est de mettre en place une compression avec perte, en réduisant le nombre de bits utilisés pour stocker les paramètres.

Pour rappel, chaque neurone se définit par une liste de valeurs appelées paramètres, à raison d'une par connexion entrante, et d'une valeur additionnelle pour appliquer un biais aux calculs. Dans le cas de nn.h tous les nombres sont de type « float », ils occupent 32 bits de mémoire. Si le cerveau artificiel est grand, ou si plusieurs réseaux sont intégrés, cela peut facilement occuper des dizaines de kilo-octets.

Une pratique répandue consiste à utiliser un autre type de valeur dite « brain float » qui code les nombres à virgule sur 16 bits, et dont on attribue souvent la découverte à Google. En fait, ce format consiste simplement à se débarasser des deux derniers octets des floats réguliers puis à considérer tous les bits retirés comme étant égaux à zéro, ce qui apporte un gain de place considérable puisqu'il divise par deux la taille occupée par l'ensemble des paramètres. Le problème est que les standards du C ne nous permettent pas d'implémenter cette idée tout en ayant la conscience tranquile. Certes, en pratique il devrait suffire de définir une union d'un nombre IEEE 754 et d'un int de même taille pour y parvenir, mais sur le papier rien ne garantit que les alignements en mémoire et les formats de stockage le permettent. Pour une fois, nous nous retrouvons contraints de mettre de côté le C et d'opter pour l'assembleur, que nous allons néanmoins interfacer avec lui.

Certains matériels commencent à être capables de travailler directement sur des valeurs ayant ce nouveau type tronqué, je pense notamment aux processuers Intel dotés de l'unité de multiplication de matrices TMUL, mais ils sont encore très marginaux sur le marché. Le mieux en attendant une adoption plus large est de continuer à utiliser les nombres au format historique de 32 bits dans nos calculs, et d'utiliser les nouveaux brain floats sous forme d'entiers bruts de 16 bits pour le stockage. Il nous faut donc deux fonctions, une pour convertir un large réel en petit entier contenant sa représentation, et une pour faire l'inverse. Nous les nommerons respectivement « f2w() » et « w2f() ». Pour pouvoir les appeler depuis un code en C ou en C++ il suffira de les déclarer de façon tout à fait ordinaire :

#ifdef __cplusplus
extern "C" {
#endif

short f2w(float f);
float w2f(short w);

#ifdef __cplusplus
}
#endif

En revanche, leur implémentation se fera dans un fichier ayant l'extension « .asm », et non pas « .c ». Celui-ci devra contenir les lignes :

section .text

; f2w(f)

global f2w

; Argument (Linux64)
; -> f: xmm0[0]

f2w:
movd eax, xmm0 ;\.
shr eax, 16    ; | return f >> 16
ret            ;/

; w2f(w)

global w2f

; Argument (Linux64)
; -> w: rdi

w2f:
shl edi, 16    ;\.
movd xmm0, edi ; | return w << 16
ret            ;/

Ce code peut être assemblé en un « .o » sous GNU/Linux avec NASM grâce à la commande "nasm -f elf64" suivie d'un espace et du nom de fichier ASM. En voici une petite description pour ceux qui débutent dans le domaine et écarquillent les yeux face aux instructions x86-64 : Notre fichier commence par déclarer et définir le corps de la fonction f2w, qui reçoit son paramètre dans un registre spécial de SIMD nommé XMM0, elle effectue ensuite un décalage de 16 bits sur la droite en appelant l'instruction SHR, puis retourne la valeur dans AX. La deuxième fonction w2f réalise l'inverse, elle reçoit son paramètre dans EDI, décale sa valeur à gauche de 16 bits en appelant l'instruction SHL, puis la retourne dans le registre XMM0.

Une fois notre objet binaire généré, il ne nous reste plus qu'à appeler f2w en C sur la totalité des paramètres d'un cerveau artificiel et stocker le résultat dans un fichier de sauvegarde, ou mieux, dans la section « .rdata » du programme compilé afin d'éviter la dépendance à un fichier de données extérieur, les utilisateurs n'en seront que plus heureux. Ensuite lors du rechargement, il suffira de faire passer les valeurs stockées dans w2f pour retrouver des paramètres exploitables.

La grande question est maintenant de savoir si cette compression altère les capacités de nos IA. La réponse est oui, mais de façon assez limitée, je n'ai constaté aucune grande défaillance pour la génération d'images, bien que dans certains cas cela puisse provoquer des imperfections notables. C'est le cas avec cette image d'avertissement anxiogène dessinée dans un style Windows 95, qui a été générée par un réseau de 55 neurones, avec et sans l'utilisation du code assembleur donné ci-dessus :

Notez le bord droit du panneau qui se retrouve légèrement abîmé dans la version tronquée.

Remplacer la fonction sigmoïde

Après la compression de données, passons à l'indépendance vis-à-vis des bibliothèques externes. Pour rappel, le passage d'un signal dans une couche neuronale doit déclencher des fonctions d'activation. En général il s'agit soit d'un calcul de sigmoïde, soit d'une fonction dite « ReLU ».

J'aurais voulu vivre dans le meilleur des mondes et utiliser ReLU qui est très simple a faire exécuter par un ordinateur. Mais pour d'obscures raisons, toutes mes tentatives de l'exploiter ont échoué, je suis donc condamné à rester sur une sigmoïde coûteuse qui nécessite de calculer une expentielle pour chaque neurone recevant une information. De plus, en C il est nécessaire de lier le code à la bibliothèque mathématique standard pour accéder à la fonction « exp() ». Traduction : Cela rend presque impossible la transmission d'un objet compilé à un utilisateur ayant une version différente de la libm sur son ordinateur. Pour les profanes de l'IA comme moi, nous pouvons tenter de remédier à ce problème en cherchant une approximation de la courbe par une suite de point, et avec un peu de tâtonnement nous finissons par retomber sur des valeurs donnant d'excellents résultats, comme celles-ci :

float not_sigmoid(float x) {

	if (x < 0.0f) return 1.0 - not_sigmoid(-x);

	if (x < 0.5f) return x*0.25f + 0.5f;
	if (x < 1.0f) return (x - 0.5f)*0.21875f + 0.625f;
	if (x < 1.5f) return (x - 1.0f)*0.171875f + 0.734375f;
	if (x < 2.0f) return (x - 1.5f)*0.125f + 0.8203125f;
	if (x < 2.5f) return (x - 2.0f)*0.078125f + 0.8828125f;
	if (x < 3.0f) return (x - 2.5f)*0.0625f + 0.921875f;
	if (x < 4.0f) return (x - 3.0f)*0.03125f + 0.953125f;
	if (x < 6.0f) return (x - 4.0f)*0.0078125f + 0.984375f;

	return 1.0f;

}

Il est à noter que les nombres choisis sont exclusivement sous la forme ∑2^a_i avec a_i∈ℤ pour minimiser les risques de troncature en cas de différence d'implémentation ou d'architecture. Cela devrait augmenter sérieusement les chances d'avoir des résultats identiques sur différentes plateformes.

Cette simplification fonctionne tellement bien qu'il est pratiquement impossible de distinguer à l'œil nu les différences qu'elle engendre dans les images par rapport à l'utilisation d'une véritable sigmoïde. Voici la comparaison d'une icône de programme générée par un même réseau de neurones avec et sans approximation :

Attention, il ne faut surtout pas utiliser la dérivée de cette fausse sigmoïde pendant l'exécution de la backpropagation d'un apprentissage, car avec ses deux coefficients directeurs nuls en deça de -1 et au-delà de 1, elle ne peut que générer des résultats catastrophiques. Le gain de l'approximation n'est en fait accessible qu'aux procédures appliquant un simple forwarding du réseau, ce qui inclut bien entendu les algorithmes finaux à intégrer dans les projets pouvant se destiner à des clients.

Il n'y a pas eu besoin d'assembleur cette fois. En réalité pour une fonction aussi simple le compilateur accomplit déjà un excellent travail d'optimisation, faire mieux que lui serait impossible ou demanderait d'y passer trop de temps, le rapport bénéfice/effort serait éminemment défavorable.

Améliorer le temps d'exécution en SIMD

Passons maintenant aux choses sérieuses, nous allons défier le compilateur sur une procédure un peu plus complexe, en vue de produire un forwarding plus rapide que le sien. Cela n'est pas si difficile car il existe un domaine dans lequel GCC et Clang ne sont pas encore tout à fait au top : Celui des produits matriciels.

Avant d'aller plus loin je vais devoir m'arrêter sur le principe du SSE, il s'agit d'une extension aujourd'hui accessible sur tous les microprocesseurs Intel. D'ailleurs, sa présence est requise par la spécification x86-64, ce qui nous permet d'en user et d'en abuser les yeux fermés, notre code fonctionnera sur tous les PC en 64 bits. Le SSE permet de réaliser des calculs en parallèle (principe du SIMD) grâce à des registres spéciaux pouvant contenir quatre valeurs de type float. Les instruction existent en deux modes, le premier dit « scalar » exécute des opérations sur un seul nombre, tandis que le deuxième dit « packed » lance quatre opérations à la fois, sur quatre nombre différents. Lors de l'implémentation d'une multiplication matricielle, les grands compilateurs semblent échouer à trouver comment employer le mode packed pour paralléliser les calculs, ils se limitent donc au scalar qui réalise les opérations une par une, même en cas d'optimisation maximum avec le passage du paramètre « -O3 ».

Voici une fonction typique de produit matricielle en langage C, c'est elle que nous allons devoir reproduire et améliorer en assembleur :

void dotm(const float A[], const float B[], float C[], size_t m, size_t n, size_t p) {

	size_t i, j, k, l;

	if (n) {
		for (i = 0 ; i < m ; i++) {
			for (j = 0 ; j < p ; j++) {
				*C = 0.0f;
				l = j;
				for (k = 0 ; k < n ; k++) {
					*C += A[k]*B[l];
					l += p;
				}
				C++;
			}
			A += n;
		}
	}

}

Les six paramètres correspondent à la description de deux matrices Amn et Bnp dont la multiplication génère une troisième matrice Cmp. Comme nous sommes dans le contexte de l'IA, nous allons nous concentrer plus spécifiquement sur la capacité à multiplier rapidement un vecteur ligne par une grande matrice. Nous prendrons donc les éléments du vecteur un par un, de façon scalar, mais chaque valeur sera recopiée pour remplir tout un registre de quatre floats. Ensuite, nous ne travaillerons plus qu'en mode packed, pour prendre les valeurs de la matrice B par paquets, les multiplier à celles du vecteur, puis tout additionner dans un autre registre intermédiaire. Cette succession d'opérations sera répétée pour toutes les lignes de B, ce qui donnera quatre résultats à déposer tels quels dans la matrice C, et il ne restera plus qu'à passer aux colonnes suivantes de B.

Voici sans plus attendre la version ASM de cet algorithme. Compte tenu de sa longueur je ne vais pas pouvoir décrire le code en détail ici, mais en gros, les suffixes « SS » et « SP » indiquent si les instructions concernent un seul float ou plusieurs. Aussi, comme le nombre de colonnes de B n'est pas forcément un multiple de quatre, trois petites boucles succèdent la procédure principale. Elles servent à traiter les valeurs résiduelles dans les cas où le paramètre p serait de la forme « 4x + 1 », « 4x + 2 » ou « 4x + 3 », avec x∈ℕ :

section .text

; dotm(A, B, C, m, n, p)

global dotm

; Arguments (Linux64)
; -> A: rdi
; -> B: rsi
; -> C: rdx
; -> m: rcx
; -> n: r8
; -> p: r9

dotm_end:
ret

dotm:

test rcx, rcx ;\ if (!m) return
jz dotm_end   ;/

shl r8, 2   ; n *= sizeof(float)
jz dotm_end ; if (!n) return

mov r10, r9 ;\ sz = p*sizeof(float)
shl r10, 2  ;/
jz dotm_end ; if (!sz) return

foreach_a_rows:

push rsi ; _B = B

mov rax, rdi ; _A = A
add rdi, r8  ; A += n

push r9 ; i = p

cmp r9, 4      ;\ if (i < 4) goto last_b_cols
jb last_b_cols ;/

foreach_b_chunks:

push rax ; a = _A
push rsi ; b = B

xorps xmm0, xmm0 ; xmm0 = 0

foreach_a_cols:      ;\. do { // Note: (n != 0) => (a != A)
movss xmm1, [rax]    ; |   xmm1[0] = *a
shufps xmm1, xmm1, 0 ; |   xmm1[3] = xmm1[2] = xmm1[1] = xmm[0]
add rax, 4           ; |   a += sizeof(float)
movups xmm2, [rsi]   ; |   xmm2 = *b
add rsi, r10         ; |   b += sz
mulps xmm1, xmm2     ; |   xmm1 *= xmm2
addps xmm0, xmm1     ; |   xmm0 += xmm1
cmp rax, rdi         ; | }
jne foreach_a_cols   ;/  while (a != A)

movups [rdx], xmm0 ; *C = xmm0
add rdx, 16        ; C += sizeof(float[4])

pop rsi     ;\ B += sizeof(float[4])
add rsi, 16 ;/

pop rax

sub r9, 4            ;\.
cmp r9, 4            ; | if (--i >= 4) goto foreach_b_cols
jae foreach_b_chunks ;/
break_b_chunks:

test r9, r9    ;\ if (!i) goto skip_b_cols
jz skip_b_cols ;/

last_b_cols:

xorps xmm0, xmm0 ; xmm0 = 0

cmp r9, 3      ;\ if (r9 == 3) goto last_3_cols
je last_3_cols ;/

cmp r9, 2      ;\ if (r9 == 2) goto last_2_cols
je last_2_cols ;/

foreach_1_col:    ;\. do {
movss xmm1, [rax] ; |   xmm1[0] = *a
add rax, 4        ; |   a += sizeof(float)
movss xmm2, [rsi] ; |   xmm2[0] = *B
add rsi, r10      ; |   b += sz
mulss xmm1, xmm2  ; |   xmm1[0] *= xmm2[0]
addss xmm0, xmm1  ; |   xmm0[0] += xmm1[0]
cmp rax, rdi      ; | }
jne foreach_1_col ;/  while (a != A)

movss [rdx], xmm0  ; *C = xmm0[0]
add rdx, 4         ; C += sizeof(float)

jmp skip_b_cols
last_2_cols:

foreach_2_cols:         ;\. do {
mov r11, rax            ; |   _a = a
add rax, 4              ; |   a += sizeof(float)
movlps xmm2, [rsi]      ; |   xmm2[0] = b[0], xmm2[1] = b[1]
cmp rax, rdi            ; |   if (a == A)
je process_last_col     ; |   goto process_last_col
add rsi, r10            ; |   b += sz
movlps xmm1, [r11]      ; |   xmm1[0] = *_a[0], xmm1[1] = *_a[1]
shufps xmm1, xmm1, 0x50 ; |   xmm1[2] = xmm1[3] = xmm1[1], xmm1[1] = xmm1[0]
add rax, 4              ; |   a += sizeof(float)
movhps xmm2, [rsi]      ; |   xmm2[2] = b[0], xmm2[3] = b[1]
add rsi, r10            ; |   b += sz
mulps xmm1, xmm2        ; |   xmm1 *= xmm2
addps xmm0, xmm1        ; |   xmm0 += xmm1
movhlps xmm1, xmm1      ; |   xmm1[0] = xmm1[2], xmm1[1] = xmm1[3]
addps xmm0, xmm1        ; |   xmm0 += xmm1
cmp rax, rdi            ; | }
jne foreach_2_cols      ;/  while (a != A)

jmp skip_last_col
process_last_col:
movss xmm1, [r11]    ; xmm1[0] = *_a
shufps xmm1, xmm1, 0 ; xmm1[1] = xmm1[0]
mulps xmm1, xmm2     ; xmm1 *= xmm2
addps xmm0, xmm1     ; xmm0 += xmm1
skip_last_col:

movlps [rdx], xmm0 ; C[0] = xmm0[0], C[1] = xmm0[1]
add rdx, 8         ; C += sizeof(float[2])

jmp skip_b_cols
last_3_cols:

foreach_3_cols:      ;\. do {
movss xmm1, [rax]    ; |   xmm1[0] = *a
shufps xmm1, xmm1, 0 ; |   xmm1[2] = xmm1[3] = xmm[0]
add rax, 4           ; |   a += sizeof(float)
movss xmm2, [rsi]    ; |   xmm2[0] = *b
movhps xmm2, [rsi+4] ; |   xmm2[2] = (b + 4)[0], xmm2[3] = (b + 4)[1]
add rsi, r10         ; |   b += sz
mulps xmm1, xmm2     ; |   xmm1 *= xmm2
addps xmm0, xmm1     ; |   xmm0 += xmm1
cmp rax, rdi         ; | }
jne foreach_3_cols   ;/  while (a != A)

movss [rdx], xmm0    ; *C = xmm0[0]
movhps [rdx+4], xmm0 ; (C + 4)[0] = xmm[2], (C + 4)[1] = xmm[3]
add rdx, 12          ; C += sizeof(float[3])

skip_b_cols:

pop r9
pop rsi ; B = _B

loop next_a_row    ;\.
ret                ; | if (--m) goto foreach_a_rows
next_a_row:        ; | else return
jmp foreach_a_rows ;/

Ceci étant fait, vérifions les gains de performance. Voici la comparaison des temps d'exécution de 100000000 (cent millions) de produits matriciels aléatoires entre la fonction C compilée par GCC/Clang (avec optimisation « -O3 -ffast-math ») et notre version manuelle codée en assembleur. Le CPU utilisé pour les exécuter est un Intel Core-i5 quelconque, tandis que le texte dans les cellules du tableau correspond au champ real de la commande time :

• bleu : Assembleur plus rapide que le C.
• rouge : Assembleur plus lent que le C.

Ce benchmarking ne va pas plus loin qu'une matrice B ayant une taille 4×8, mais les gains semblent continuer d'augmenter avec des dimensions plus grandes. En tout cas les résultats sont au-delà de toute espérance, si avec une matrice ridiculement petite la fonction assembleur n'a rien de miraculeux à montrer, lorsque les objets à multiplier prennent un minimum de volume, elle devient deux à quatre fois plus rapide que ses concurrentes en C. Il est aussi amusant de noter que GCC surpasse de loin Clang pour les petits travaux, mais qu'il se fait doubler par ce dernier quand les tâches deviennent vraiment denses.

Avec ça et les deux autre optimisations vues dans les parties précédentes, nous voilà fin prêts à présenter quelque chose à nos clients !

Pour aller plus loin

On reproche souvent au code assembleur sa portabilité réduite : Il est vrai que les fonctions présentées dans cet article ne peuvent tourner que sous GNU/Linux avec un microprocesseur x86-64. Une toute première amélioration à faire serait alors de créer des conditions pour utiliser d'autres instructions sous Windows et avec l'architecture IA-32.

Le passage à une version Windows serait extrêmement simple, le seul changement à faire consisterait à récupérer les arguments depuis des registres différents, c'est à dire typiquement, réaliser les substitutions suivantes :

Cependant, pour une adaptation aux architectures IA-32 nous avons de pain sur la planche, car avec elles la présence de SSE n'est pas garantie. Il faudrait donc maintenir une liste de pointeurs vers différentes versions de chaque fonction, et exécuter l'instruction « cpuid » à l'initialisation de notre programme pour vérifier les capacités du CPU. Il se poserait alors la question de savoir comment optimiser le produit matriciel sans SSE, la seule autre extension SIMD capable de gérer des floats ayant existé étant « 3DNow! », et elle n'était présente que sur les processeurs AMD conçus au début de ce millénaire.

Puisque nous en sommes revenus à notre multiplication de matrices, profitons-en pour nous demander si nous pouvons encore améliorer les performances de notre code. Nous n'avons pas fait appel aux instructions de prefetch qui permettent de gérer manuellement le cache de la mémoire, il serait donc intéressant de voir à quel endroit nous pourrions les insérer pour en tirer parti. Par ailleurs, nous n'avons pas non plus utilisé les registres généraux EBX et EBP, alors que nous envoyons des valeurs temporaires sur la pile. Les réquisitionner pour en faire des registres de stockage constituerait une autre piste sérieuse à étudier.

Enfin, ne laissons pas sur le carreaux nos ancêtres UNIX, ils sont toujours vivants ! Nous autres jeunes fougueux avons le devoir de prendre soin d'eux, de leur témoigner la reconnaissance qu'ils méritent. Ajoutons donc à notre TODO liste le projet de voir si notre code assembleur peut être adapté à d'autres architectures notables, et mettons-y aussi de l'ARM, car il faut savoir se tourner vers l'avenir de temps à autre :

• Étudier la faisabilité pour l'AltiVec sur PowerPC.
• Étudier la faisabilité pour le VIS sur SPARC.
• Étudier la faisabilité pour le NEON sur ARM.

Sinon, nous pouvons également refaire surface dans le haut niveau, c'est à dire le langage C (n'est-ce pas !), pour évaluer la solution très prometteuse d'OpenGL. Je suis prêt à parier qu'un peu de GLSL bien placé nous permettrait d'accomplir des miracles…

Et si nous mettions en amont de notre code assembleur une condition pour vérifier la disponibilité d'OpenGL, afin d'utiliser ses fonctionnalités quand elles sont présentes et de renvoyer vers nos fonctions artisanales dans le cas contraire, ne nous retrouverions-nous pas avec la perfection ultime entre les mains ?