Exemple d’optimisation du code en R

Un billet pour expliquer les étapes conduisant à une amélioration des performances d’un programme assez simple : réalisant un bootstrap résiduel dans un modèle linéaire général. Cette méthode statistique est adaptée aux écarts assez sévères à la normalité, fragilisant le théorème central limite et limitant la validité de la méthode de Wald. Cette méthode ne résout en rien les défauts d’additivité ou de linéarité du modèle ou les problèmes de corrélation entre observations. Dans le scenario à deux groupes, la méthode est équivalente au parabootstrap avec variances égales. On peut la résumer à la procédure suivante pour un échantillon de taille N:

  1. Tirer au hasard N valeurs parmi les N résidus, avec remise.
  2. Ajouter ces N nouveaux résidus aux N valeurs ajustées (fitted values), créant ainsi un nouveau vecteur réponse
  3. Réajuster (refitting) le modèle à ce nouveau vecteur réponse, en conservant la matrice de modèle (model matrix) inchangée
  4. Répéter les étapes (1) à (3) un grand nombre de fois (p.e. dix mille fois) tout en enregistrant les coefficients du modèle à chaque refitting
  5. Calculer les intervalles de confiance des coefficients comme les quantiles des coefficients de l’ensemble des expériences enregistrées en (4)

L’échantillon de test était issu d’une vraie étude, et comportait 924 observations avec 11 covariables qualitatives conduisant à 38 coefficients distincts après recodage des variables qualitatives à k modalités en k-1 variables binaires et ajout d’une ordonnée à l’origine (intercept).

La première approche est directe, basée sur le package ‘boot’ et sur la fonction ‘lm’:

resid_boot1=function(mod, R1=10, R2=1000) {
	f=fitted(mod)
	res=resid(mod)
	boot(data=model.frame(mod), statistic=function(data, idx) {
		datax=data;
		# on change la variable réponse
		datax[,1]=f+res[idx]
		# puis on réajuste (refitting)
		coef(lm(data=datax, formula=formula(mod)))
	}, R=R1*R2)
}

Le paramètre mod représente un modèle ajusté (fitted) sur le jeu de données original et R1 et R2 sont deux paramètres représentant le nombre de simulations. Habituellement, on utilise un seul paramètre R représentant le nombre total de simulations, mais pour avoir une homogénéité d’interface, j’ai ici utilisé deux paramètres dont le produit est égal au nombre total de simulations. On pourra remarquer que pour réajuster (refitting) le modèle, je n’ai pas fait appel à la fonction update() de R. Cette fonction est extrêmement dangereuse et ne devrait jamais être utilisée parce qu’elle ne réajuste pas le modèle dans l’environnement dans lequel il avait été ajusté la première fois.

Une fois cette fonction exécutée, il suffit d’appeler boot.ci() pour chacun des coefficients du modèle afin de calculer les intervalles de confiance:

bt = resid_boot1(original_model, R1=10,R2=100)

confidence_intervals = t(sapply(seq_along(bt$t0),
function (i) {
  boot.ci(bt, type="perc",index=i)$perc[4:5]
}))

La fonction resid_boot1 est malheureusement lente pour R version 4.0.3. Sur un PC assez moderne (AMD Ryzen 7 1700, 3 Ghz, 8 coeurs, 24 Go de RAM, bien ventilé, tournant en turbo à 3.3 Ghz en continu) sous Windows 10 64 bits il faut 4,8 secondes pour mille simulations (R1=10, R2=100).

Le premier constat qui permet d’améliorer les performances, c’est que l’appel à la fonction lm() refait itérativement des calculs inutiles. Cette fonction va commencer par recoder les variables explicatives qualitatives en variables binaires, mais comme ces variables explicatives ne changent pas d’une simulation à l’autre, le recodage est redondant. Pour gagner en temps, on peut utiliser la fonction lm.fit() qui prend directement en entrée, une matrice numérique dite matrice de modèle (model matrix) qui a déjà toutes les colonnes recodées.

Ainsi, la deuxième implémentation du boostrap résiduel se base sur lm.fit:

resid_boot2=function(mod, R1=10, R2=1000) {
	f=fitted(mod)
	res=resid(mod)
	boot(data=model.matrix(mod), statistic=function(data, idx) {
		lm.fit(data, f+res[idx])$coefficients
	}, R=R1*R2)
}
system.time(bt2<-resid_boot2(modT, R1=10, R2=100))

Pour le même nombre de simulations (mille), il ne faut plus que 1,42 secondes soit une multiplication par 3,4 de la vitesse d’exécution.

Cette fonction lm.fit fait quelques vérifications avant d’appeler .lm.fit, qui fonctionne presque de la même manière mais devrait être un peu plus rapide; avec cette dernière fonction le temps d’exécution passe à 1,36 seconde; amélioration minime.

L’implémentation de .lm.fit de R n’étant pas optimale, il existe une fonction lmfit dans le package Rfast, dont l’interface est semblable mais qui devrait être un peu plus rapide. Le temps d’exécution passe à 0,97 seconde avec Rfast::lmfit.

L’étape suivante d’optimisation repose sur une propriété importante du modèle linéaire général estimé par les moindres carrés: conditionnellement à la matrice de modèle (c’est-à-dire aux variables explicatives), chacun des coefficients du modèle n’est qu’une combinaison linéaire (c’est-à-dire, une moyenne pondérée) du vecteur réponse. Cela se voit très bien dans la formule classique d’estimation du modèle linéaire général. Étant donné une matrice de modèle X à n lignes (représentant les n observations) et k colonnes (représentant les k variables explicatives après recodage) et un vecteur réponse Y à n valeurs (la variable à expliquer), les coefficients du modèle s’estiment comme (X’ X)-1X’ Y où X’ représente la matrice transposée de X. Ainsi, on commence par faire le produit de la transposée de X avec X elle même, puis on en prend l’inverse; ensuite on remultiplie par la transposée de X; enfin, on multiple par le vecteur Y. Seule cette dernière étape va varier dans le boostrap résiduel, puisque X est constante. On peut donc gagner une inversion de matrice et trois produits de matrice en précalculant (X’ X)-1X’.

C’est donc l’étape d’optimisation suivante:

resid_boot5=function(mod, R1=10, R2=1000) {
	f=fitted(mod)
	res=resid(mod)
	x = model.matrix(mod)
	linmat = solve(t(x) %*% x) %*% t(x)
	boot(data=model.matrix(mod), statistic=function(data, idx) {
		(linmat %*% (f+res[idx]))[,1]
	}, R=R1*R2)
}

Le temps d’exécution passe à 0,12 seconde pour mille itérations, soit environ huit fois plus rapide que la version précédente; afin de pouvoir calculer ce chiffre avec précision, il faut multiplier par dix le nombre de simulations : 1,20 seconde pour dix mille simulations.

Les optimisations suivantes sont des « micro optimisations » qui feront quand même gagner une vistesse non négligeable (vous verrez). Afin de mieux calculer les timings et de se défaire de quelques propriétés inimportunes du package boot sur la manière de générer les nombres aléatoires (en un énorme bloc si simple=FALSE ou en trop petits blocs si simple=TRUE), on peut réécrire la même fonction de manière plus brute. Au passage, on peut transposer la matrice avant la multiplication de matrice et utiliser crossprod qui est un peu plus rapide. En même temps, on va réajuster (refitter) R1=10 simulations à chaque itération, afin de réduire les coûts intertiels (overheads) du produit matriciel et du tirage au sort des observations.

resid_boot6=function(mod, R1=10, R2=1000) {
	f=fitted(mod)
	res=resid(mod)
	ff = matrix(fitted(mod), nrow=nrow(x), ncol=R1)
	x = model.matrix(mod)
	tlinmat = t(solve(t(x) %*% x) %*% t(x))
	N=nrow(x)
	a = replicate(R2, crossprod(tlinmat,
                                    ff+sample(res,N*R1,replace=TRUE)))
	dim(a)=c(length(coef(mod)),R1*R2)
	list(t=t(as.matrix(a)), t0=coef(mod))
}

Pour dix mille itérations, on passe à 0,97 seconde, ce qui est un maigre gain par rapport aux 1,20 secondes de l’implémentation précédente.

L’étape suivante a consisté à évaluer si le goulot d’étranglement (bottleneck) se situait dans le produit matriciel ou l’échantillonnage des résidus, en décomposant la fonction. On pouvait constater que l’échantillonnage des résidus était le facteur limitant, comme si le générateur pseudo-aléatoire était trop lent. En essayant tous les générateurs pseudo-aléatoires proposés par R (cf fonction set.seed), j’observais une vitesse d’exécution presque identique, mais lente, pour tous les générateurs. Même en se basant sur la fonction runif(), qui est la plus directe des fonctions, on ne dépassait pas les 15,3 millions de nombres aléatoires générés par seconde, comme le montre la ligne de code suivante:

system.time(for(i in 1:1e4) sample.int(1:nobs(model),1e4,replace=TRUE))
#   user  system elapsed 
#   6.55    0.00    6.55

Pour un processeur qui tourne en permanence en mode turbo à 3,3 Ghz, ça fait environ 216 cycles de CPU pour calculer un nombre aléatoire. Pour l’algorithme Mersenne-Twister par défaut, ça paraît beaucoup, puisque cet algorithme est connu pour être plutôt rapide, basé sur quelques xor et bit shiftings par nombre 32 bits généré. Afin de vérifier la vitesse optimale de génération des nombres aléatoires selon cet algorithme, j’ai téléchargé une version très basique de l’algorithme (https://github.com/ESultanik/mtwister) et l’ai compilée avec GCC 8.3.0 pour x86_64 et l’option d’optimisation -O2. Elle générait un milliard de nombres pour 3.9 seconde, soit à peine 13 cycles de CPU par itération. Cela correspondait globalement à une multiplication par 16,7 de la vitesse de génération des nombres aléatoires. C’est ainsi qu’il devenait rentable de réimplementer une fonction d’échantillonnage (fastsample) en C:

Le cœur de l’algorithme est donné ci-dessous en code C:

	for(R_xlen_t i=0; i<inobs; i++) {
		double unif=(((double)genRandLong(rangen))
                            * (1.0/(1+(double)0xFFFFFFFFL)));
		presult[i] = psamp[(size_t)(ssize*unif)];
	}

On pouvait alors remplacer sample() par fastsample() pour quasiment annihiler le temps de génération des nombres pseudo-aléatoires:

resid_boot7=function(mod, R1=10, R2=1000, seed=1337) {
	rgen=rangen(seed)
	f=fitted(mod)
	res=resid(mod)
	ff = matrix(fitted(mod), nrow=nrow(x), ncol=R1)
	x = model.matrix(mod)
	tlinmat = t(solve(t(x) %*% x) %*% t(x))
	N=nrow(x)
	a = replicate(R2, crossprod(tlinmat, ff+fastsample(rgen, res,N*R1)))
	dim(a)=c(length(coef(mod)),R1*R2)
	list(t=t(as.matrix(a)), t0=coef(mod))
}

Pour dix mille itérations, il ne fallait plus que 0,40 seconde. Sur ces 0,40 seconde, le fastsample représentait environ 0,04 seconde. Pour ce calcul, j’ai dû multiplier par dix le nombre d’itérations (soit cent mille itérations). Presque tout le temps de calcul restant provenait de la multiplication de matrice (crossprod). Cette multiplication, bien que censée être optimisée, est loin d’être optimale comme on peut le voir dans la fonction internal_crossprod du fichier src/main/array.c du code source de R. C’est une implémentation assez naïve, basée sur les bonnes vieilles FP80 du FPU x87. En utilisant les instructions vectorielles (SIMD) avec des FP64, on peut a priori gagner en performance. Pour bien utiliser les SIMD FP64 des processeurs x86, c’est-à-dire, les SSE2 et AVX/AVX2, il faut aligner les accès à la mémoire sur 128 bits (SSE2/AVX) ou 256 bits (AVX2). Cela oblige à arranger les matrices pour ces alignements. Par exemple, une matrice 7×3 (7 lignes et 3 colonnes) va devoir être transformée en une matrice 8×3, avec la dernière ligne à zéro pour que les accès mémoire soient bien alignés, sans compter l’allocation du vecteur de données qui devra être bien alignée. Cela va demander donc un petit temps de transformation de la matrice; heureusement, pour la matrice (X’X)-1X’, ce travail peut n’être fait qu’une seule fois. C’est ainsi, qu’il a été possible d’écrire deux fonctions en langage C en utilisant les Vector Extensions de GCC: une fonction plmat() qui prépare la matrice à gauche de la multiplication (transposition/alignement) et une fonction fastmatmul() qui multiplie cette matrice préparée par une nouvelle matrice. En utilisant les SSE2, la nouvelle fonction s’exécutait en 0,14 seconde pour dix mille itérations. En passant à cent mille itérations (R1=100, R2=1000), il fallait 1,24 seconde. Avec AVX2+FMA3, on passait à 1,18 seconde pour cent mille itérations. À ce stade, il ne sert à rien d’optimiser par rapport à un usage ponctuel du bootstrap résiduel. Cependant pour évaluer les propriétés statistiques (notamment les biais) du boostrap résiduel, on peut avoir besoin de faire des dizaines de milliers de bootstrap. En visant dix mille simulations par bootstrap, on peut facilement se retrouver avec des centaines de millions de calculs. C’est là qu’on peut gagner en performance en utilisant les multiples cœurs du processeur; cela se fait bien avec le package ‘snow’ sous Windows ou parallel::mclapply sous GNU/Linux. En pratique, avec huit cœurs on gagne sans difficulté un facteur huit; avec le SMT, on peut encore multiplier la vitesse par 1,15 ou 1,20 environ, mais au-delà il est difficile de gagner en performances sur un CPU. Sur GPGPU on peut espérer y gagner encore. Cependant mes résultats préliminaires avec OpenCL sur une GeForce GTX 1060 3Go sont extrêmement décevants, avec des performances pires que sur CPU avec OpenCL: cela semble s’expliquer par un usage trop intensif de nombres entiers. Il y a probablement des optimisations auxquelles je n’ai pas pensé et qui pourraient accélérer.

Le tableau ci-dessous récapitule les gains de performance successifs:

AlgorithmeTemps estimé pour 100 000 itérations (secondes)
1 : naïf avec lm()480
2 : amélioré avec lm.fit()142
3 : amélioré avec .lm.fit()136
4 : avec Rfast::lmfit97
5 : précalculant (X’X)-1X’11,9
6 : pré-transposition et calcul de 10 itérations simultanées9,7
7 : réécriture fastsample() en C4,0
8 : optimisation multiplication de matrices avec SSE20,124
9 : optimisation AVX2+FMA30,118
10 : OpenCL+Geforce GTX 1060 3 Go0,13
Performance des différents algorithmes de bootstrap résiduel du modèle linéaire général

La version N°9 est environ quatre mille fois plus rapide que la version N°1. Il n’y avait qu’une seule vraie optimisation algorithmique (méthode N°5), faisant gagner un facteur 8, les autres optimisations étant des petites bidouilles qui ne changent pas la complexité algorithmique.

Pseudo-aléatoire R

Pourquoi le générateur pseudo-aléatoire est-il aussi lent ? Même l’accès le plus direct à la source pseudo-aléatoire, la fonction runif() est trois fois plus lente que fastsample() bien qu’elle ne fasse pas l’indexation d’un vecteur. En fouillant un peu le code source de R, en appelant les différentes parties pour en évaluer les performances et en regardant le code assembleur généré par GCC, on découvre que le problème ne vient pas du générateur Mersenne-Twister lui-même (fonction MT_genrand dans src/main/RNG.c) ni de unif_rand dans src/main/RNG.c dont les inerties (overheads) sont faibles. Le problème vient de random2() dans src/library/stats/src/random.c. Soit dit au passage, le code de src/main/random.c est redondant avec celui de src/library/stats/src/random.c. Le premier code ne semble finalement jamais être appelé en situation ordinaire et il est plus performant que le second.

> runif
function (n, min = 0, max = 1) 
.Call(C_runif, n, min, max)
<bytecode: 0x0000000004706130>
<environment: namespace:stats>
> system.time(for(i in 1:1e4)runif(1e4))
   user  system elapsed 
   1.48    0.00    1.49
> system.time(for(i in 1:1e4).Internal(runif(1e4,0,1)))
   user  system elapsed 
   1.28    0.00    1.29

Sur l’AMD Ryzen 1700, on constate donc des performances un peu meilleures pour la fonction interne, implémentée dans src/main/random.c que la fonction de bibliothèque dynamique implémentée dans src/library/stats/src/random.c, même si les deux restent mauvaises. Concernant src/library/stats/src/random.c on comprend rapidement le problème en lisant le code source:

for (R_xlen_t i = i0; i < n; i++) {
    rx[i] = fn(ra[i % na], rb[i % nb]);
    if (ISNAN(rx[i])) naflag = TRUE;
}

Où la fonction fn() génère un nombre aléatoire selon une loi uniforme; les deux paramètres de fn représentent respectivement le minimum et le maximum de la loi uniforme.

Le problème de performance provient des deux divisions entières (i % na et i % nb) réalisées pour le cyclage des arguments min et max de runif pour gérer des appels tels que stats::runif(100, 1:5, c(50,100)). La division entière est une opération très lente sur de nombreux processeurs, nécessitant 14 à 45 cycles (reciprocal throughput) sur la microarchitecture Zen 1 (https://www.agner.org/optimize/instruction_tables.pdf). Sur processeur Intel, c’est encore pire; sur Coffe Lake (8ème génération d’Intel Core) la latence est de 42 à 95 cycles et le reciprocal throughput est de 24 à 90 cycles. Le problème est moins sévère avec le code de src/random qui utilise MOD_ITERATE2_CORE (src/include/R_ext/Itermacros.h), dont le code est montré ci-dessous:

define MOD_ITERATE2_CORE(n, n1, n2, i, i1, i2, loop_body) do { \
        for (; i < n;                                                   \
             i1 = (++i1 == n1) ? 0 : i1,                                \
                 i2 = (++i2 == n2) ? 0 : i2,                            \
                 ++i) {                                                 \
            loop_body                                                   \
                }                                                       \
    } while (0)

Le code assembleur généré utilise quand même des cmov suboptimaux et a du mal à tenir dans les registres. On peut accélérer énormément les choses en réécrivant le code pour le cas particulier où les arguments min et max à runif() sont des vecteurs de longueur 1. Le nouveau code de random2() dans src/main/random.c devient alors:

static Rboolean random2(double (*f) (double, double),
                        double *a, R_xlen_t na, double *b, R_xlen_t nb,
                        double *x, R_xlen_t n)
{
    double ai, bi;
    R_xlen_t i, ia, ib;
    Rboolean naflag = FALSE;
    errno = 0;
    if (na == 1 && nb == 1) {
        double ra=*a, rb=*b;
        double *end = x+n;
        double *x0=x;
        while(x < end) {
            *x = f(ra, rb);
            if (ISNAN(*x)) naflag = TRUE;
            x++;
        }
        *x0=(R_xlen_t)f;
    } else {
        MOD_ITERATE2(n, na, nb, i, ia, ib, {
            ai = a[ia];
            bi = b[ib];
            x[i] = f(ai, bi);
            if (ISNAN(x[i])) naflag = TRUE;
        });
    }
    return(naflag);
}

Par ailleurs la fonction runif dans src/nmath/runif.c qui est chargée de calculer un seul nombre aléatoire entre min et max, est sous-optimale dans le cas où min=0 et max=1, parce qu’elle va faire une addition, une soustraction et une multiplication inutiles pour fournir un résultat compris entre min et max. On peut alors gagner un petit peu à tester le scenario commun min=0 et max=1:

double runif(double a, double b)
{
    if (a == 0 && b == 1) { /* ajout pour accélérer */
        return unif_rand();
    }
    if (!R_FINITE(a) || !R_FINITE(b) || b < a)  ML_WARN_return_NAN;

    if (a == b)
        return a;
    else {
        double u;
        /* This is true of all builtin generators, but protect against
           user-supplied ones */
        do {u = unif_rand();} while (u <= 0 || u >= 1);
        return a + (b - a) * u;
    }
}

Moyennant ces modifications, les appels à runif() bénficièrent d’une accélération par un facteur trois environ. Il existe encore quatre ralentissements non négligeables que je pus évaluer en réécrivant complètement la fonction pour shunter certaines étapes : (1) l’usage d’une boucle générique pour tous les générateurs pseudo-aléatoires à deux paramètres (rnorm, rbinom, runif, etc.) engendre un niveau d’indirection (2) le test ISNAN(*x) répétitif n’est pas nécessaire pour une distribution uniforme et engendre un UCOMISD et un CMOVP à chaque itération (3) la comparaison a==0 && b == 1 dans runif(double a,double b) est exécutée à chaque itération (4) la fonction unif_rand() qui s’occupe de sélectionner le générateur pseudo-aléatoire ajoute encore un niveau d’indirection par rapport à un appel direct à MT_genrand() pour l’usage du générateur Mersenne-Twister. Cette dernière sélection du générateur est particulièrement peu performante, car basée sur un gros switch() après lecture d’une variable globale. Si on corrige ces 4 problèmes on gagne encore un facteur deux en performances. C’est ainsi qu’on peut multiplier par six les performances de runif.

En mettant ça au propre, je pense que je pourrais écrire un petit patch pour R.

Rejoindre la conversation

1 commentaire

Laisser un commentaire

Votre adresse e-mail ne sera pas publiée. Les champs obligatoires sont indiqués avec *