¿Opciones para resolver sistemas ODE en GPU?

Me gustaría desarrollar sistemas de resolución de EDO en GPU, en un entorno 'trivialmente paralelo'. Por ejemplo, haciendo un análisis de sensibilidad con 512 conjuntos de parámetros diferentes.

Idealmente, quiero resolver ODE con un solucionador de paso de tiempo inteligente y adaptable como CVODE, en lugar de un paso de tiempo fijo como Forward Euler, pero ejecutándolo en una GPU NVIDIA en lugar de CPU.

¿Alguien ha hecho esto? ¿Hay bibliotecas para ello?

Es posible que desee mirar en la biblioteca odeint de Boost y Thrust . Se pueden combinar como se discute aquí .

La biblioteca DifferentialEquations.jl es una biblioteca para un lenguaje de alto nivel (Julia) que tiene herramientas para transformar automáticamente el sistema ODE a una versión optimizada para una solución paralela en GPU. Se pueden emplear dos formas de paralelismo: paralelismo basado en matrices para sistemas ODE grandes y paralelismo de parámetros para estudios de parámetros en sistemas ODE relativamente pequeños (<100). Soporta métodos implícitos y explícitos de alto orden y rutinariamente supera o coincide con otros sistemas en los puntos de referencia (al menos, envuelve a los demás para que sea fácil verificarlos y usarlos)

Para esta funcionalidad específica, es posible que desee echar un vistazo a DiffEqGPU.jl, que es el módulo para el paralelismo automático de parámetros. La biblioteca DifferentialEquations.jl tiene funcionalidad para estudios de parámetros paralelos , y este módulo aumenta las configuraciones existentes para que el estudio se realice automáticamente en paralelo. Lo que uno hace es transformar su existente ODEProblem(u otro DEProblemsimilar SDEProblem) en un EnsembleProblemy especificar con un prob_funccómo se generan los otros problemas del prototipo. Lo siguiente resuelve 10,000 trayectorias de la ecuación de Lorenz en la GPU con un método adaptativo explícito de alto orden:

using OrdinaryDiffEq, DiffEqGPU
function lorenz(du,u,p,t)
 @inbounds begin
     du[1] = p[1]*(u[2]-u[1])
     du[2] = u[1]*(p[2]-u[3]) - u[2]
     du[3] = u[1]*u[2] - p[3]*u[3]
 end
 nothing
end

u0 = Float32[1.0;0.0;0.0]
tspan = (0.0f0,100.0f0)
p = (10.0f0,28.0f0,8/3f0)
prob = ODEProblem(lorenz,u0,tspan,p)
prob_func = (prob,i,repeat) -> remake(prob,p=rand(Float32,3).*p)
monteprob = EnsembleProblem(prob, prob_func = prob_func)
@time sol = solve(monteprob,Tsit5(),EnsembleGPUArray(),trajectories=10_000,saveat=1.0f0)

Observe que el usuario no necesita escribir ningún código de GPU, y con un solo RTX 2080, este punto de referencia es una mejora 5 veces mayor que el uso de una máquina Xeon de 16 núcleos con paralelismo multiproceso. Luego, puede consultar el archivo README para saber cómo hacer cosas como utilizar múltiples GPU y hacer multiprocesamiento + GPU para utilizar un grupo completo de GPU simultáneamente . Tenga en cuenta que cambiar a subprocesos múltiples en lugar de GPU es un cambio de línea: en EnsembleThreads()lugar de EnsembleGPUArray().

Luego, para los solucionadores implícitos, se mantiene la misma interfaz. Por ejemplo, lo siguiente utiliza Rosenbrock de alto orden y métodos implícitos de Runge-Kutta:

function lorenz_jac(J,u,p,t)
 @inbounds begin
     σ = p[1]
     ρ = p[2]
     β = p[3]
     x = u[1]
     y = u[2]
     z = u[3]
     J[1,1] = -σ
     J[2,1] = ρ - z
     J[3,1] = y
     J[1,2] = σ
     J[2,2] = -1
     J[3,2] = x
     J[1,3] = 0
     J[2,3] = -x
     J[3,3] = -β
 end
 nothing
end

function lorenz_tgrad(J,u,p,t)
 nothing
end

func = ODEFunction(lorenz,jac=lorenz_jac,tgrad=lorenz_tgrad)
prob_jac = ODEProblem(func,u0,tspan,p)
monteprob_jac = EnsembleProblem(prob_jac, prob_func = prob_func)

@time solve(monteprob_jac,Rodas5(linsolve=LinSolveGPUSplitFactorize()),EnsembleGPUArray(),dt=0.1,trajectories=10_000,saveat=1.0f0)
@time solve(monteprob_jac,TRBDF2(linsolve=LinSolveGPUSplitFactorize()),EnsembleGPUArray(),dt=0.1,trajectories=10_000,saveat=1.0f0)

Si bien este formulario requiere que le dé un jacobiano para usarlo en la GPU (actualmente, se solucionará pronto), la documentación de DifferentialEquations.jl muestra cómo hacer cálculos automáticos simbólicos jacobianos en funciones definidas numéricamente , por lo que todavía no hay manual trabajo aquí. Recomiendo encarecidamente estos algoritmos porque la lógica de ramificación de un método como CVODE generalmente causa la desincronización de subprocesos y, de todos modos, no parece funcionar tan bien como un método de Rosenbrock en este tipo de escenarios.

Al usar DifferentialEquations.jl, también obtiene acceso a la biblioteca completa, que incluye funcionalidades como el análisis de sensibilidad global que puede hacer uso de esta aceleración de GPU. También es compatible con números duales para un rápido análisis de sensibilidad local . El código basado en GPU obtiene todas las características de DifferentialEquations.jl, como el manejo de eventos y el gran conjunto de solucionadores ODE que están optimizados para diferentes tipos de problemas , lo que significa que no es solo un simple solucionador ODE GPU único sino un parte de un sistema con todas las funciones que también tiene un eficiente soporte de GPU.