NVIDIA Warp: Nueva era en la simulación de física computacional para inteligencia artificial

La Ingeniería Asistida por Computadora (CAE) está evolucionando hacia flujos de trabajo impulsados por inteligencia artificial, incorporando modelos de física que se generalizan a través de diversas geometrías y condiciones operativas. A diferencia de los modelos de lenguaje, estos modelos requieren grandes volúmenes de datos de alta fidelidad y cumplimiento físico.

Recientes trabajos sobre surrogados en dinámica de fluidos computacionales (CFD) indican que los datos de entrenamiento generados por simulaciones suelen ser el costo limitante en la práctica. Esto genera exigencias sobre el simulador, que debe ser nativo de GPU, rápido y capaz de integrarse directamente en flujos de trabajo de aprendizaje automático.

NVIDIA Warp: Un marco para simulaciones aceleradas

NVIDIA Warp es un marco diseñado para la simulación acelerada, generación de datos y computación espacial que conecta CUDA y Python. Permite a los desarrolladores escribir núcleos de alto rendimiento como funciones de Python que se compilan en código eficiente para ejecutarse en la GPU. A diferencia de los marcos basados en tensores, donde los desarrolladores expresan cálculos como operaciones sobre matrices N-dimensionales enteras, en Warp los desarrolladores pueden crear núcleos flexibles que se ejecutan simultáneamente en todos los elementos de una red computacional.

Los núcleos de simulación a menudo se expresan en redes computacionales y dependen de flujos de control dependientes de datos, como condicionales y actualizaciones selectivas que varían por elemento. En los marcos tensoriales, estos patrones requieren la creación de máscaras booleanas que pueden volverse engorrosas y desperdiciar cómputo en elementos irrelevantes. En un núcleo de Warp, cada hilo puede ramificarse, omitir o salir de forma independiente, expresando esta lógica de manera natural sin necesidad de soluciones alternativas de enmascarado.

Creando un solucionador de Navier–Stokes en Warp

Este artículo guía sobre la construcción de un solucionador 2D de Navier–Stokes completamente en Warp. Se detalla cómo el modelo de programación de Warp se asigna a un solucionador de ecuaciones diferenciales parciales (EDP). Luego, se diferencia a través de la simulación para resolver un problema de perturbación óptima de forma integral. Finalmente, se presentan estudios de caso industriales que muestran lo que Warp puede habilitar en flujos de trabajo de producción.

Para mantener el enfoque en Warp en lugar de en métodos numéricos, se utiliza un ejemplo de turbulencia decaída 2D, descrito mediante la formulación de vorticidad-funcion de corriente de las ecuaciones incomprensibles de Navier-Stokes. La vorticidad (omega) evoluciona de acuerdo con la ecuación de transporte:

(frac{partial omega}{partial t} + frac{partial psi}{partial y}frac{partial omega}{partial x} – frac{partial psi}{partial x}frac{partial omega}{partial y} = frac{1}{text{Re}}nabla^2 omega tag{1})

y la función de corriente (psi) se recupera de la vorticidad a través de la ecuación de Poisson:

(nabla^2 psi = -omega tag{2})

Con condiciones de frontera periódicas, la ecuación anterior se reduce a una ecuación algebraica en el espacio de Fourier, evitando la necesidad de solucionadores iterativos:

(hat{psi}_{m,n} = frac{hat{omega}_{m,n}}{k_x^2 + k_y^2} tag{3})

donde ((k_x, k_y)) es el par de números de onda en el espacio de Fourier. El solucionador utiliza el algoritmo de Transformada Rápida de Fourier (FFT) para transformar de manera eficiente (omega) y (psi) al espacio de Fourier y viceversa.

Componentes clave del solucionador

Cada paso de tiempo tiene dos subcomponentes. Primero, la ecuación de transporte de vorticidad se discretiza en una red (N x N) sobre un dominio cuadrado (L x L). La solución se avanza en el tiempo utilizando un esquema de Runge-Kutta de tercera orden (RK3) para obtener (omega(t+Delta t)). En segundo lugar, la ecuación de Poisson se resuelve en el espacio de Fourier para obtener la actualización de (psi(t+Delta t)).

El step() en el ejemplo avanza la simulación mediante un solo llamado. El método self._solve_poisson() abstrae el pipeline de (omega(t+Delta t) xrightarrow{text{FFT}} hat{omega} xrightarrow{text{Eq.,3}} hat{psi} xrightarrow{text{IFFT}} psi(t+Delta t)) para modularidad.

Diferenciación a través del solucionador

Una vez construido el solucionador, surge la pregunta de cómo hacerlo diferenciable. La diferenciación automática (AD) calcula derivadas exactas aplicando la regla de la cadena a cada operación elemental en el gráfico computacional. A diferencia de las diferencias finitas, AD evita el ajuste de tamaño de paso y proporciona gradientes precisos hasta la precisión de la máquina.

Para ilustrar, se aborda el problema de perturbación óptima descrito en Prediction and Control of Two-Dimensional Decaying Turbulence Using Generative Adversarial Networks. En flujos turbulentos, pequeñas perturbaciones a las condiciones iniciales pueden amplificarse con el tiempo y alterar significativamente la trayectoria del flujo.

El sistema de diferenciación automática de Warp genera dos versiones de un programa en tiempo de compilación para una simulación diferenciable. Se pueden encontrar más detalles y ejemplos en el repositorio de GitHub de NVIDIA/warp.

Ilustración de un hombre mayor con auriculares y chaqueta