Opções para resolver sistemas ODE em GPUs?

Eu gostaria de descobrir como resolver sistemas de ODEs em GPUs, em uma configuração 'trivialmente paralelizável'. Por exemplo, fazendo uma análise de sensibilidade com 512 conjuntos de parâmetros diferentes.

Idealmente, eu quero resolver o ODE com um solucionador de tempo inteligente adaptável como CVODE, em vez de um tempo fixo fixo como o Forward Euler, mas executá-lo em uma GPU NVIDIA em vez de CPU.

Alguém já fez isso? Existem bibliotecas para isso?

Você pode procurar na biblioteca odeint do Boost e no Thrust . Eles podem ser combinados conforme discutido aqui .

A biblioteca DifferentialEquations.jl é uma biblioteca para uma linguagem de alto nível (Julia) que possui ferramentas para transformar automaticamente o sistema ODE em uma versão otimizada para solução paralela em GPUs. Existem duas formas de paralelismo que podem ser empregadas: paralelismo baseado em matriz para sistemas ODE grandes e paralelismo de parâmetros para estudos de parâmetros em sistemas ODE relativamente pequenos (<100). Ele suporta métodos implícitos e explícitos de alta ordem e supera rotineiramente ou combina com outros sistemas em benchmarks (no mínimo, envolve os outros, para que seja fácil verificar e usá-los!)

Para essa funcionalidade específica, você pode dar uma olhada no DiffEqGPU.jl, que é o módulo para o paralelismo automatizado de parâmetros. A biblioteca DifferentialEquations.jl possui funcionalidade para estudos de parâmetros paralelos , e este módulo aumenta as configurações existentes para que o estudo aconteça automaticamente em paralelo. O que se faz é transformar o existente ODEProblem(ou outro DEProblemsemelhante SDEProblem) em um EnsembleProbleme especificar com prob_funccomo os outros problemas são gerados a partir do protótipo. A seguir, são resolvidas 10.000 trajetórias da equação de Lorenz na GPU com um método adaptativo explícito de alta ordem:

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 o usuário não precisa escrever nenhum código da GPU e, com um único RTX 2080, esse benchmark é uma melhoria de 5x em relação ao uso de uma máquina Xeon de 16 núcleos com paralelismo multithread. Pode-se, então, verificar o README para saber como executar várias GPUs e executar multiprocessamento + GPUs para utilizar um cluster completo de GPUs simultaneamente . Observe que alternar para multithreading em vez de GPUs é uma alteração de linha: em EnsembleThreads()vez de EnsembleGPUArray().

Então, para solucionadores implícitos, a mesma interface é válida. Por exemplo, o seguinte usa os métodos Rosenbrock de alta ordem e Runge-Kutta implícito:

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)

Embora este formulário exija que você forneça um jacobiano para ser usado na GPU (atualmente, será corrigido em breve), a documentação DifferentialEquations.jl demonstra como fazer cálculos jacobianos simbólicos automáticos em funções numericamente definidas , portanto ainda não há manual trabalhe aqui. Eu recomendo esses algoritmos porque a lógica de ramificação de um método como o CVODE geralmente causa dessincronização de threads e parece não ter um desempenho tão bom quanto o método de Rosenbrock nesses tipos de cenários.

Ao usar DifferentialEquations.jl, você também obtém acesso à biblioteca completa, que inclui funcionalidades como a análise de sensibilidade global, que pode fazer uso dessa aceleração de GPU. Também é compatível com números duplos para análise de sensibilidade local rápida . O código baseado em GPU obtém todos os recursos do DifferentialEquations.jl, como a manipulação de eventos e o grande conjunto de solucionadores de ODE otimizados para diferentes tipos de problemas , o que significa que não é apenas um simples solucionador único de GPU ODE, mas sim um parte de um sistema completo que também possui suporte eficiente à GPU.