Options pour résoudre les systèmes ODE sur les GPU?

Je voudrais consolider la résolution de systèmes d'ODE sur des GPU, dans un cadre «trivialement parallélisable». Par exemple, faire une analyse de sensibilité avec 512 jeux de paramètres différents.

Idéalement, je veux faire une résolution ODE avec un solveur d'horodateur adaptatif intelligent comme CVODE, plutôt qu'un horodatage fixe comme Forward Euler, mais en l'exécutant sur un GPU NVIDIA au lieu d'un CPU.

Quelqu'un a-t-il fait cela? Existe-t-il des bibliothèques pour cela?

Vous voudrez peut-être examiner la bibliothèque d'odeint de Boost et Thrust . Ils peuvent être combinés comme discuté ici .

La bibliothèque DifferentialEquations.jl est une bibliothèque pour un langage de haut niveau (Julia) qui dispose d'outils pour transformer automatiquement le système ODE en une version optimisée pour une solution parallèle sur les GPU. Il existe deux formes de parallélisme qui peuvent être utilisées: le parallélisme basé sur des matrices pour les grands systèmes ODE et le parallélisme des paramètres pour les études de paramètres sur les systèmes ODE relativement petits (<100). Il prend en charge des méthodes implicites et explicites de haut niveau et surpasse régulièrement ou correspond à d'autres systèmes dans les benchmarks (à tout le moins, il encapsule les autres, il est donc facile de les vérifier et de les utiliser!)

Pour cette fonctionnalité spécifique, vous voudrez peut-être jeter un œil à DiffEqGPU.jl qui est le module de parallélisme de paramètres automatisé. La bibliothèque DifferentialEquations.jl a des fonctionnalités pour les études de paramètres parallèles , et ce module augmente les configurations existantes pour que l'étude se fasse automatiquement en parallèle. Ce que l'on fait, c'est transformer leur existant ODEProblem(ou autre DEProblemsimilaire SDEProblem) en un EnsembleProblemet spécifier avec un prob_funccomment les autres problèmes sont générés à partir du prototype. Ce qui suit résout 10 000 trajectoires de l'équation de Lorenz sur le GPU avec une méthode adaptative explicite d'ordre élevé:

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)

Notez que l'utilisateur n'a besoin d'écrire aucun code GPU, et avec un seul RTX 2080, cela se compare comme une amélioration 5x par rapport à l'utilisation d'une machine Xeon à 16 cœurs avec un parallélisme multithread. On peut ensuite consulter le README pour savoir comment faire des choses comme utiliser plusieurs GPU et faire du multiprocessing + GPU pour utiliser simultanément un cluster complet de GPU . Notez que le passage au multithreading au lieu des GPU est un changement de ligne: EnsembleThreads()au lieu de EnsembleGPUArray().

Ensuite, pour les solveurs implicites, la même interface est valable. Par exemple, ce qui suit utilise des méthodes Rosenbrock d'ordre élevé et des méthodes Runge-Kutta implicites:

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)

Bien que ce formulaire nécessite que vous donniez un jacobien afin d'être utilisé sur le GPU (actuellement, il sera bientôt corrigé), la documentation DifferentialEquations.jl montre comment effectuer des calculs jacobiens symboliques automatiques sur des fonctions définies numériquement , il n'y a donc toujours pas de manuel travailler ici. Je recommanderais fortement ces algorithmes car la logique de branchement d'une méthode comme CVODE provoque généralement la désynchronisation des threads et ne semble pas fonctionner aussi bien qu'une méthode Rosenbrock dans ces types de scénarios de toute façon.

En utilisant DifferentialEquations.jl, vous avez également accès à la bibliothèque complète, qui comprend des fonctionnalités telles que l'analyse de sensibilité globale qui peut utiliser cette accélération GPU. Il est également compatible avec les nombres doubles pour une analyse de sensibilité locale rapide . Le code basé sur GPU obtient toutes les fonctionnalités de DifferentialEquations.jl, comme la gestion des événements et le grand ensemble de solveurs ODE qui sont optimisés pour différents types de problèmes , ce qui signifie qu'il ne s'agit pas simplement d'un simple solveur ODE GPU unique, mais plutôt d'un partie d'un système complet qui a également un support GPU efficace.