Óptica física en la GPU Nvidia

Introducción

Durante la mayor parte de mi carrera en óptica, he simulado imágenes de moduladores espaciales de luz y, sobre todo, de SLM iluminados con fuentes parcialmente coherentes. Normalmente, la FFT hace un trabajo bastante rápido, al menos para una sola imagen. A veces, sin embargo, hay que hacer muchas simulaciones. Un ejemplo de ello es cuando se trabaja en los requisitos ópticos de un sistema óptico. Un enfoque que ha resultado útil es utilizar Muestreo por hipercubos latinos. Para un sistema óptico, probablemente habría que incluir hasta aberraciones de 5º orden, quizá un poco más, y esto se convierte en una simulación bastante grande. Quizá 1000 o 2000 imágenes. Entonces, tal vez se examinen los datos. Tal vez el espacio sondeado (rango) era demasiado pequeño, o demasiado grande, y uno puede querer rehacer esas simulaciones un par de veces.

Aunque sea el ordenador el que lo hace, no siempre es fácil programar alguna otra tarea, e incluso si lo es, cambiar de tarea cada media hora u hora (o quizá más) no es muy productivo. Por lo tanto, poder calcular FFT bidimensionales con rapidez importa, e importa bastante.

La estructura de una simulación de imagen

La aproximación a las simulaciones de imágenes parcialmente coherentes normalmente depende de si planeas calcular muchas imágenes usando una distribución específica de fuentes, entonces quizás tu aproximación a esto sería SOCS (Suma De Sistemas Coherentes), o Si quieres buscar una buena distribución de intensidad, la sobrecarga con SOCS no merece la pena y la aproximación será simplemente calcular las transformadas inversas de las amplitudes reflejadas por tu máscara o, como era normalmente mi caso, SLM. Cada una de estas transformadas inversas representa una imagen cuya intensidad debe sumarse para hallar la exposición completa.

Dependiendo del tamaño del dominio periódico, y del factor de relleno de la pupila de la iluminación (normalmente llamado sigma), podríamos fácilmente acabar necesitando hacer varios cientos de transformaciones inversas. En este punto, uno puede decidir hacer un muestreo más grueso de la fuente de luz, pero yo no iría significativamente por debajo de 100 puntos de iluminación. Si bajamos a 25 (por ejemplo), sin duda habrá contribuciones visibles de la periodicidad de la función de coherencia. Incluso 50 es demasiado, así que, como regla general, 100 no es un mal número. A veces, la densidad de muestreo viene dada por un homogeneizador, por lo que no podemos añadir mucho más.

Resumiendo, la simulación de imagen básica de la que hablaré aquí consiste en una transformación directa para calcular la amplitud que ilumina la pupila del sistema, y una serie de transformaciones inversas, que no son más que traslaciones triviales de una transformación, para calcular las imágenes parciales. Entre medias se realiza un filtrado de la pupila, pero se trata de una operación muy barata desde el punto de vista computacional.

API CUDA

Al principio, el API CUDA puede parecer un poco intimidante, e incluso con mi bastante considerable tiempo haciendo estas cosas, estaba un poco ansioso por cuánto dolor tendría que hacerme pasar antes de salir de la curva de aprendizaje. Lo que complicaba un poco las cosas era el hecho de que quería un código que pudiera utilizarse junto con Octave o Matlab y lo quería en C. Este código también debe devolver resultados idénticos cuando se compila para usuarios que no tienen acceso a CUDA.

Para mi sorpresa, CUDA era mucho más fácil de usar de lo que esperaba, y las cosas que me preocupaban, como tener que reiniciar un sinfín de veces porque había conseguido bloquear mi tarjeta gráfica, nunca ocurrieron.

Algunas opciones estratégicas

El primer punto de decisión con respecto a la API fue si debía utilizar la biblioteca CUDA FFTW-wrapper o echarme un poco atrás y utilizar directamente la biblioteca cuFFT. Si tienes una aplicación independiente basada en la librería FFTW, entonces quizás sea una opción, incluso si hay algunas cosas obvias que esta envoltura simplemente no puede hacer, como asignar buffers tanto de dispositivo como de host para tus datos. En la jerga CUDA, la memoria "de dispositivo" se encuentra en la GPU y la memoria host está vinculada a la CPU.

Con Octave/Matlab, hay una cosa más, y es que el uso de la librería CUDA FFTW-wrapper creará conflictos de nombres entre FFTW y funciones con nombres idénticos en la librería wrapper. Incluso si decidimos proporcionar implementaciones a todas las variantes FFT de Octave/Matlab, seguiremos sin saber si no hay algún binario en algún lugar esperando obtener un plan de fftw_plan_dft_2d, y aunque la biblioteca envoltorio proporciona dicha función, las bibliotecas no son compatibles binariamente. Por lo tanto, utilizar la guía de conversión a cuFFT es una buena idea.

La solución más básica

Empecemos por el principio. ¿Qué es lo mínimo que debemos hacer para ejecutar FFT en la GPU? En realidad, no hace falta mucho. Debemos

• Asignar memoria GPU - device_mem = cudaMalloc(n);
• Hacer un plan - cufftPlan2D(&plan, n1, n2, CUFFT_Z2Z);
• Copia los datos del host a la GPU - cudaMemcpy(device_mem, host_mem, n, cudaMemcopyDefault);
• Ejecutar el plan - cufftExecZ2Z(plan, device_mem_in, device_mem_out, CUFFT_FORWARD);
• Devuelve el resultado - cudaMemcpy(host_mem, device_mem, n, cudaMemcpyDefault);

Evidentemente, soy bastante chapucero en este aspecto, pero seguir esta línea nos permitiría realizar la FFT en la GPU, pero también sería un poco decepcionante. La mayor parte del tiempo se pasaría copiando datos hacia y desde el dispositivo. Lo más probable (si es que "bueno" es la palabra adecuada) es que no nos diéramos cuenta de que la GPU hace nada, salvo que calcula el resultado correcto.

Entonces, ¿qué teníamos que hacer con los datos que habíamos copiado de la GPU? Lo único que teníamos que hacer era calcular el producto de la transformada y su conjugado para cada muestra y acumular la suma hasta que termináramos de escanear las fuentes de iluminación. Si pudiéramos hacer eso en la GPU, al menos no tendríamos que leer el resultado más de una vez. Esto significa que debemos aprender a ejecutar algún código trivial en la GPU.

Un trozo de código que haría eso es:

__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<<(img, amp, factor, n);
}

El prototipo para addIntensity2Df tiene que ser exportado a una cabecera y puede ser usado desde código C normal. Una vez finalizado el bucle principal sobre las fuentes, tendremos que volver a leer la intensidad acumulada y liberar un montón de buffers.

Un tamaño de bloque de 256 funcionó bastante bien, pero esto es ajustable. Es una buena idea poner trozos de código como este en su propio archivo .cu y compilarlo con nvcc. Esto nos permitió compilar el código C estándar restante con sus optimizaciones habituales usando GCC o Clang. Y no hay ningún problema en enlazarlo todo a un archivo .mex y utilizarlo con Octave.

La API de CUDA ofrece muchas formas de abordar esta simulación para conseguir un mejor aprovechamiento de la GPU, como por ejemplo los streams, pero incluso con esta sencilla implementación podemos conseguir un aprovechamiento de la GPU de 70% para dominios de simulación grandes (por ejemplo, 2048×2048 puntos de muestreo). Sin embargo, no se gana mucho implementando una alternativa a FFTW para una única transformación. La sobrecarga consume casi por completo la ganancia de realizar el cálculo en la GPU.

¿Merece la pena el esfuerzo?

¿Ha merecido la pena? Mi GPU principal es una Ryzen 7 8700G. El código que suelo ejecutar está optimizado para este tipo de CPU desde hace bastante tiempo para sacar el máximo partido a los 8 núcleos. La GPU con la que quería compararla es una RTX 3070Ti, que no es un demonio de velocidad ni siquiera para los estándares de la época en que salió al mercado. Aun así, los resultados muestran que computa imágenes parcialmente coherentes unas 10 veces más rápido que esta CPU algo reciente. Es un buen aumento de velocidad para un primer intento. No es tan fácil encontrar una CPU 10 veces más rápida a cualquier precio.

Para que funcione a una velocidad decente, hay que reducir la cantidad de datos transferidos por bucle. Dado que se trataba de un esfuerzo mínimo basado en código optimizado para una CPU, era bastante reacio a convertir demasiado código para ejecutarlo en la GPU. Por suerte, sólo hay que copiar la parte de la transformación que atraviesa la pupila, que es una cantidad muy limitada de datos a copiar. Quizá en alguna revisión futura traslade más parte del código a la GPU y entonces se eliminen todas las copias por iteración. Puede que haya otra ganancia de velocidad de 60%-70% y, una vez eliminados todos los pequeños cuellos de botella, una 5070Ti podría brillar de verdad.

Volviendo la vista atrás, el cálculo de imágenes parcialmente coherentes encaja perfectamente en una GPU. El gran número de transformaciones que hay que hacer se puede guardar en la GPU y, además, acumular la intensidad es realmente sencillo, y la GPU hace un trabajo fantástico. No se me ocurre ninguna aplicación que ofrezca tantas mejoras con tan poco esfuerzo.