Optique physique sur le GPU Nvidia

Intro

Pendant la majeure partie de ma carrière en optique, j'ai simulé l'imagerie de modulateurs spatiaux de lumière, et plus particulièrement de SLMs illuminés par des sources partiellement cohérentes. En général, la FFT fait un travail assez rapide, au moins pour une seule image. Parfois, cependant, il est nécessaire d'effectuer un grand nombre de simulations. C'est le cas par exemple lorsque l'on travaille sur les exigences optiques d'un système optique. Une approche qui s'est avérée utile est l'utilisation de Échantillonnage hypercube latin. Pour un système optique, on inclurait probablement jusqu'aux aberrations de cinquième ordre, voire un peu plus, et cela deviendrait une simulation assez importante. Peut-être 1000 ou 2000 images. Ensuite, vous pouvez examiner les données. Il se peut que l'espace sondé (la portée) soit trop petit ou trop grand, et que l'on veuille refaire ces simulations plusieurs fois.

Bien que ce soit l'ordinateur qui le fasse, il n'est pas toujours facile de programmer une autre tâche, et même si c'est le cas, il n'est pas très productif de changer de tâche toutes les demi-heures ou toutes les heures (voire plus). C'est pourquoi il est important de pouvoir calculer rapidement des FFT bidimensionnelles, et c'est même très important.

La structure d'une simulation d'image

L'approche des simulations d'images partiellement cohérentes dépend généralement du fait que vous prévoyez de calculer un grand nombre d'images en utilisant une distribution spécifique de sources, alors votre approche pourrait être SOCS (Sum Of Coherent Systems), ou si vous voulez rechercher une bonne distribution d'intensité, le surcoût de SOCS n'en vaut pas la peine et l'approche consistera simplement à calculer les transformées inverses des amplitudes réfléchies par votre masque ou, comme c'était généralement le cas pour moi, SLM. Chacune de ces transformées inverses représente une image dont l'intensité doit être additionnée afin de trouver l'exposition complète.

En fonction de la taille du domaine périodique et du facteur de remplissage de la pupille de l'illumination (généralement appelé sigma), nous pourrions facilement nous retrouver à devoir effectuer plusieurs centaines de transformations inverses. À ce stade, on peut décider de procéder à un échantillonnage plus grossier de la source lumineuse, mais je ne descendrais pas en dessous de 100 points d'illumination. Si nous descendons à 25 (par exemple), il y aura certainement des contributions visibles de la périodicité de la fonction de cohérence. Même 50 est un peu juste, alors en règle générale, 100 n'est pas un mauvais chiffre. Parfois, la densité d'échantillonnage est donnée par un homogénéisateur, il n'y a donc pas grand-chose à ajouter.

Pour résumer, la simulation d'image de base dont je parlerai ici consiste en une transformation directe pour calculer l'amplitude qui illumine la pupille du système, et une série de transformations inverses qui ne sont que des translations triviales d'une transformation, pour calculer les images partielles. Il y a un filtrage de la pupille entre les deux, mais il s'agit d'une opération très peu coûteuse sur le plan informatique.

API CUDA

Dans un premier temps, le API CUDA peut sembler un peu intimidant, et même avec le temps considérable que j'ai passé à faire ces choses, j'étais un peu inquiet de la douleur que j'allais devoir endurer avant de sortir de la courbe d'apprentissage. Ce qui a quelque peu compliqué les choses, c'est que je voulais un code qui puisse être utilisé avec Octave ou Matlab et je le voulais en C. Ce code doit également renvoyer des résultats identiques lorsqu'il est compilé pour un utilisateur qui n'a pas accès à CUDA.

À ma grande surprise, CUDA a été beaucoup plus facile à appréhender que je ne l'avais prévu, et les choses que je craignais, comme devoir redémarrer un nombre infini de fois parce que j'avais réussi à bloquer ma carte graphique, ne se sont jamais produites.

Quelques choix stratégiques

The first decision point regarding the API was, should I use the CUDA FFTW-wrapper library or back-off a little and use the cuFFT library directly. If you have a stand-alone application based on the FFTW-library, then maybe it is a choice, even if there are some obvious things that this wrapper just cannot do, such as allocating both device and host buffers for your data. In the CUDA lingo, “device” memory is on the GPU and host memory is bound to you CPU.

With Octave/Matlab, there’s one more thing, which is that using the CUDA FFTW-wrapper library will create name conflicts between FFTW and identically named functions in the wrapper library. Even if we decide to provide implementations to all Octave/Matlab FFT variants, we still won’t know if there isn’t some binary somewhere expecting to get a plan from fftw_plan_dft_2d, and while the wrapper library provides such a function, the libraries are not binary compatible. Therefore, using the conversion guide to cuFFT is a good idea.

La solution la plus simple

So, let’s start lite. What is the bare minimum we must do to run FFTs on the GPU? Actually, that does not require much. We must

• Allocate GPU memory – device_mem = cudaMalloc(n);
• Make a plan – cufftPlan2D(&plan, n1, n2, CUFFT_Z2Z);
• Copy host data to the GPU – cudaMemcpy(device_mem, host_mem, n, cudaMemcopyDefault);
• Execute the plan – cufftExecZ2Z(plan, device_mem_in, device_mem_out, CUFFT_FORWARD);
• Bring the result back – cudaMemcpy(host_mem, device_mem, n, cudaMemcpyDefault);

Je suis évidemment un peu négligent, mais en suivant ces lignes, nous pourrons effectivement effectuer la FFT sur le GPU, mais ce sera aussi un peu décevant. La plupart du temps serait consacré à copier des données vers et depuis le périphérique. Il y a de bonnes chances (si le mot "bonnes" est approprié) que nous ne remarquions pas que le GPU fait quoi que ce soit, si ce n'est qu'il calcule le bon résultat.

Alors, que devons-nous faire avec les données que nous avons copiées depuis le GPU ? La seule chose que nous devions faire était de calculer le produit de la transformée et son conjugué pour chaque échantillon et d'accumuler la somme jusqu'à ce que nous ayons fini de scanner les sources d'illumination. Si nous pouvions le faire sur le GPU, nous n'aurions pas à relire le résultat plus d'une fois. Cela signifie que nous devons apprendre à exécuter un code trivial sur le GPU.

Voici un morceau de code qui permettrait de faire cela :

__global__ void _launchAddIntensity2Df(float *img, const cufftComplex *amp, float factor, int n)
{
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        float re = amp[idx].x;
        float im = amp[idx].y;
        img[idx] += (re * re + im * im) * factor;
    }
}

extern "C" void addIntensity2Df(float *img, const cufftComplex *amp, float factor, int n)  {
    int blockSize = 256;
    int gridSize = (n + blockSize - 1) / blockSize;
    _launchAddIntensity2Df<<<gridSize, blockSize>>>(img, amp, factor, n);
}

Le prototype de addIntensity2Df doit être exporté dans un en-tête et peut être utilisé à partir d'un code C normal. Une fois la boucle principale sur les sources terminée, nous devrons relire l'intensité accumulée et libérer un grand nombre de tampons.

A block size of 256 worked out quite well, but this is tunable. It’s a good idea to put pieces of code like this in it own .cu file and compile it with nvcc. This allowed us to compile the remaining standard C-code with it’s usual optimizations using GCC or Clang. And there is no problem then linking it all to a .mex file and use with Octave.

The CUDA API offers many ways to address this simulation in order to achieve better GPU utilization, such for example streams but even with this simple implementation, we can achieve 70% GPU utilization for large simulation domains (for example 2048×2048 sampling points). There is, however, little to be gained to implement an alternative to FFTW for a single transform. The overhead almost completely consumes the gain of doing the computation on the GPU.

Le jeu en vaut-il la chandelle ?

So, was this worth it? My main GPU is a Ryzen 7 8700G. The code I usually run has been optimized for this type of CPU for quite some time to get most out of the 8 cores. The GPU that I wanted to compare it with is an RTX 3070Ti, not a speed daemon even by the standards of the days when in was released. Still, the results show it computes partially coherent images about 10 times faster than this somewhat recent CPU. It’s a nice speed-up for a first attempt. Not so easy to find a 10 times faster CPU at any price.

Pour que cela fonctionne à une vitesse décente, il faut réduire la quantité de données transférées par boucle, ce qui est obsolète. Comme il s'agissait d'un effort minimal basé sur un code optimisé pour un CPU, j'étais plutôt réticent à convertir trop de code pour l'exécuter sur le GPU. Heureusement, il suffit de copier la partie de la transformation qui passe par la pupille, ce qui représente une quantité très limitée de données à copier. Peut-être que lors d'une prochaine révision, je transférerai une plus grande partie du code sur le GPU, ce qui permettra d'éliminer toutes les copies par itération. Il pourrait y avoir un autre gain de vitesse de 60%-70% et une fois que tous les petits goulots d'étranglement seront éliminés, un 5070Ti pourrait vraiment briller.

En y repensant, le calcul d'images partiellement cohérentes convient parfaitement à un GPU. Le grand nombre de transformations à effectuer peut être conservé sur le GPU et l'accumulation de l'intensité est vraiment simple, et le GPU fait un travail fantastique. De mémoire, je n'ai pas trouvé d'application qui apporte autant d'améliorations avec si peu d'efforts.