Optionen zum Lösen von ODE-Systemen auf GPUs?

Ich möchte das Lösen von ODE-Systemen auf GPUs in einer "trivial parallelisierbaren" Umgebung durchführen. Führen Sie beispielsweise eine Sensitivitätsanalyse mit 512 verschiedenen Parametersätzen durch.

Idealerweise möchte ich ODE-Lösungen mit einem intelligenten adaptiven Zeitschrittlöser wie CVODE durchführen, anstatt mit einem festen Zeitschritt wie Forward Euler, aber auf einer NVIDIA-GPU anstelle einer CPU.

Hat jemand das getan? Gibt es Bibliotheken dafür?

Vielleicht möchten Sie einen Blick in Boosts Odeint-Bibliothek und Thrust werfen . Sie können wie hier beschrieben kombiniert werden .

Die Bibliothek DifferentialEquations.jl ist eine Bibliothek für eine Hochsprache (Julia), die Tools zum automatischen Umwandeln des ODE-Systems in eine optimierte Version für die parallele Lösung auf GPUs enthält. Es gibt zwei Formen der Parallelität, die verwendet werden können: Array-basierte Parallelität für große ODE-Systeme und Parameterparallelität für Parameterstudien an relativ kleinen (<100) ODE-Systemen. Es unterstützt implizite und explizite Methoden hoher Ordnung und übertrifft oder vergleicht routinemäßig andere Systeme in Benchmarks (zumindest umschließt es die anderen, sodass es einfach ist, sie zu überprüfen und zu verwenden!).

Für diese spezielle Funktionalität sollten Sie sich DiffEqGPU.jl ansehen , das Modul für die automatisierte Parameterparallelität. Die Bibliothek DifferentialEquations.jl bietet Funktionen für parallele Parameterstudien. Dieses Modul erweitert die vorhandenen Konfigurationen, damit die Studie automatisch parallel durchgeführt wird. Was man tut, ist, sein vorhandenes ODEProblem(oder DEProblemähnliches SDEProblem) in ein umzuwandeln EnsembleProblemund mit einem anzugeben, prob_funcwie die anderen Probleme aus dem Prototyp generiert werden. Das Folgende löst 10.000 Trajektorien der Lorenz-Gleichung auf der GPU mit einer expliziten adaptiven Methode hoher Ordnung:

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)

Beachten Sie, dass der Benutzer keinen GPU-Code schreiben muss. Bei einem einzelnen RTX 2080 ist dies eine 5-fache Verbesserung gegenüber der Verwendung eines 16-Kern-Xeon-Computers mit Multithread-Parallelität. In der README-Datei erfahren Sie, wie Sie mehrere GPUs verwenden und Multiprocessing + GPUs ausführen, um einen vollständigen Cluster von GPUs gleichzeitig zu verwenden . Beachten Sie, dass das Umschalten auf Multithreading anstelle von GPUs eine Zeilenänderung ist: EnsembleThreads()anstelle von EnsembleGPUArray().

Für implizite Löser gilt dann dieselbe Schnittstelle. Im Folgenden werden beispielsweise Rosenbrock-Methoden hoher Ordnung und implizite Runge-Kutta-Methoden verwendet:

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)

Während dieses Formular erfordert, dass Sie einen Jacobian angeben, um auf der GPU verwendet zu werden (wird derzeit in Kürze behoben), zeigt die Dokumentation zu DifferentialEquations.jl , wie automatische symbolische Jacobian-Berechnungen für numerisch definierte Funktionen durchgeführt werden. Daher gibt es noch kein Handbuch Arbeit hier. Ich würde diese Algorithmen wärmstens empfehlen, da die Verzweigungslogik einer Methode wie CVODE im Allgemeinen eine Thread-Desynchronisierung verursacht und in solchen Szenarien sowieso nicht so gut zu funktionieren scheint wie eine Rosenbrock-Methode.

Durch die Verwendung von DifferentialEquations.jl erhalten Sie auch Zugriff auf die vollständige Bibliothek, die Funktionen wie die globale Sensitivitätsanalyse enthält, die diese GPU-Beschleunigung nutzen können. Es ist auch mit zwei Zahlen für eine schnelle lokale Sensitivitätsanalyse kompatibel . Der GPU-basierte Code bietet alle Funktionen von DifferentialEquations.jl, wie die Ereignisbehandlung und die große Anzahl von ODE-Solvern, die für verschiedene Arten von Problemen optimiert sind. Dies bedeutet, dass es sich nicht nur um einen einfachen einmaligen GPU-ODE-Solver handelt, sondern um einen Teil eines voll funktionsfähigen Systems, das auch eine effiziente GPU-Unterstützung bietet.