Optimisation de l'algorithme MUSIC avec BLAS et extensions Python/C++

Contexte et problématique

Dans le cadre du développement d'une librairie de localisation par interférométrie, j'ai procédé à un profilage approfondi du code pour identifier les goulots d'étranglement de performance. Cette analyse a révélé que 80% du temps d'exécution des simulations était consacré au calcul du spectre MUSIC (MUltiple SIgnal Classification).

Pour plus de détails sur l'algorithme MUSIC, consultez mon article dessus.

Cette observation m'a conduit à réécrire la fonction critique en C++ via une extension Python. Ensuite, j'ai tiré parti de la bibliothèque BLAS pour accélérer les opérations d'algèbre linéaire.

Extensions Python/C++ et modules natifs

Les extensions Python/C++ permettent d'intégrer du code natif directement dans l'écosystème Python, offrant ainsi le meilleur des deux mondes : la simplicité et la richesse de Python avec les performances du code compilé. On peut tout particulièrement écrire des fonctions utilisant l'API NumPy qui bénéficient d'une excellente compatibilité avec les ndarrays et héritent automatiquement des mécanismes de broadcasting et de gestion mémoire optimisée de NumPy.

Les ufuncs et gufuncs : des fonctions optimisées pour NumPy

NumPy propose deux types de fonctions particulièrement efficaces :

📊 Les ufuncs (Universal Functions)

Des fonctions qui s'appliquent élément par élément sur des tableaux NumPy. Exemples : np.sin(), np.add(), np.multiply(). Elles gèrent automatiquement le broadcasting et sont vectorisées en C.

🔧 Les gufuncs (Generalized Universal Functions)

Version avancée des ufuncs qui peut traiter des blocs de données plutôt que juste des éléments individuels. Parfait pour l'algèbre linéaire : produits matriciels, décompositions, etc. Mon algorithme MUSIC en bénéficie directement !

L'avantage des gufuncs pour mon projet MUSIC ? Elles me permettent de définir des opérations complexes (comme le calcul du spectre) qui s'exécutent nativement en C++ tout en conservant la simplicité d'utilisation des fonctions NumPy classiques.

💻 Exemple d'implémentation d'une gufunc

1. Fichiers d'en-tête

#define PY_SSIZE_T_CLEAN
#include <Python.h>
#include "numpy/arrayobject.h"
#include "numpy/ufuncobject.h"

2. Définition du noyau

static void music_kernel(char **args, npy_intp *dimensions,
                         npy_intp* steps, void* data) {
    // args : pointeurs vers les tableaux d'entrée/sortie
    // dimensions : longueurs des core dimensions
    double *input = (double*) args[0];
    double *output = (double*) args[1];
    
    // Implémentation du calcul MUSIC optimisé
    for (npy_intp i = 0; i < dimensions[0]; ++i) {
        // Calculs vectorisés ici
        output[i] = /* calcul MUSIC */ input[i];
    }
}

3. Enregistrement

PyObject *uf = PyUFunc_FromFuncAndData(
    kernels,           // pointeurs vers les kernels
    data,              // pointeurs vers les données
    typecodes,         // types des arguments
    1,                 // nombre de kernels
    1,                 // nombre d'entrées
    1,                 // nombre de sorties
    PyUFunc_None,      // identité
    "music_spectrum",  // nom exposé
    "Calcul spectre MUSIC optimisé",
    0
);

4. Compilation avec setuptools

from setuptools import setup, Extension
import numpy

ext = Extension(
    "music_module",
    sources=["music_module.cpp"],
    include_dirs=[numpy.get_include()],
    language="c++",
)

setup(
    name="music_optimization",
    ext_modules=[ext],
)

Optimisation avec BLAS

Pour pousser l'optimisation encore plus loin, j'ai intégré OpenBLAS, une implémentation haute performance de la Basic Linear Algebra Subprograms (BLAS).

Pourquoi BLAS ?

BLAS est une spécification standardisée qui définit un ensemble de routines d'algèbre linéaire de bas niveau. L'utilisation d'OpenBLAS apporte des gains substantiels grâce à :

🚀 Vectorisation SIMD

Exploitation automatique des instructions SSE/AVX pour traiter plusieurs éléments simultanément

🧠 Optimisations cache

Algorithmes adaptés à la hiérarchie mémoire pour maximiser la réutilisation des données

⚡ Parallélisation

Distribution automatique sur tous les cœurs CPU (contrôlable via OPENBLAS_NUM_THREADS)

🎯 Assembleur optimisé

Routines spécialisées pour chaque architecture de processeur (x86, ARM, etc.)

Décomposition du spectre MUSIC

Le spectre MUSIC pour un vecteur de direction e se calcule selon la formule :

\[P_{\mathrm{MUSIC}}(\theta) = \frac{1}{\mathbf{e}(\theta)^\dagger \mathbf{U}_n \mathbf{U}_n^\dagger \mathbf{e}(\theta)}\]

où 𝐞(θ) est le vecteur de direction, 𝐔ₙ la matrice des vecteurs propres du sous-espace bruit, et † désigne la transposée conjuguée

Cette expression peut être efficacement décomposée en opérations BLAS élémentaires :

Décomposition algorithmique :

BLAS Level 2 : tmp = 𝐔ₙ† × 𝐞 (produit matrice-vecteur via cblas_zgemv)
BLAS Level 1 : result = tmp† × tmp (produit scalaire conjugué via cblas_zdotc)
Opération scalaire : spectrum = 1 / result

Cette décomposition permet d'exploiter les routines BLAS hautement optimisées pour les nombres complexes, spécifiquement conçues pour tirer parti des architectures modernes.

Conclusion

Cette optimisation progressive de l'algorithme MUSIC, depuis l'extension Python/C++ jusqu'à l'intégration d'OpenBLAS, a permis de diviser par 5 le temps de calcul tout en préservant la compatibilité avec l'écosystème NumPy. L'approche hybride a ainsi rendu possible l'exécution de simulations complexes en temps réel.

Un avantage particulièrement appréciable de l'utilisation des gufuncs est la conservation du mécanisme de broadcasting NumPy. Cette compatibilité a facilité la réalisation d'études paramétriques complexes, permettant de lancer de nombreux calculs MUSIC en parallèle de manière vectorisée. Les simulations multi-paramètres qui nécessitaient auparavant des boucles explicites s'exécutent désormais automatiquement sur l'ensemble des configurations en une seule opération optimisée.