Diffeq.jl v6.4: Full GPU ODEs, Neural ODEs with Batching on GPUs, and More
juliadiffeq.org
juliadiffeq.org
Making differential equations 'just work' on a GPU is super cool. I feel like if this were almost any other language, it'd be impossible to get this level of code-reuse where a user can just bring their own differential equation that may not even know about GPUs and then have the differential equation library solve it on the GPU by just giving a GPUarray as the initial condition.
________________________________
Also, since the blog post doesn't seem to link to it, here's a link to the actual DifferentialEquations.jl repo: https://github.com/JuliaDiffEq/DifferentialEquations.jl
Otherwise, I would probably use it for everything now given the performance and the very good FFI story.
Good IDE support (juno). Julia is fast. Numpy is actually c, julia is just julia.
You have an amazing interop that just works. With python, matlab, c,c++,etc etc.
https://github.com/JuliaPy/PyCall.jl
Package manager just works.
The language is a lot more powerful. Diffeq is around 1000 lines of code which is nuts.
1. Making sure every operation in the solvers used a higher level building block (@..) that automatically built fused broadcast kernels, and CuArrays.jl then supplies the broadcast overloads. That makes all of the non-stiff methods work.
2. Change the default choice of Jacobian to "dense array of size NxN" to "zero'd outer product of input array defined via broadcast". This translates to `J = false .* z .* z'` as a trick to make a Jacobian match whatever array type the array type library thinks is best, so J will be a GPUArray while DiffEq knows nothing about GPUs (also works for things like MPI, but GPUs is easier for people to reason about these days). So a one line change to the stiff ODE solvers made them all GPU compatible! But...
3. There's a small bug, or at least missing feature, in the GPU libraries that their QR doesn't implement backsolve. Technically the user can work around that because we allow someone to pass a linear solver routine with the `linsolve` argument, but it's easier to just handle it for everyone in the defaults. So using Requires.jl, if you have CuArrays installed, then we add a linsolve routine to the defaults for you, defined by https://github.com/JuliaDiffEq/DiffEqBase.jl/blob/master/src... . Technically this can get upstreamed to CuArrays.jl and get deleted from DiffEqBase, but for now it'll live there so that everything is easier on users.
So yes, make it high level, avoid assuming things like the Jacobian live on the CPU, and add one missing method to the GPU array library, and now when the user passes `u0 = CuArray(...)` for the initial condition, the algorithm compiles a version which does all state operations on the GPU (and all logic on the CPU, so it's a pretty good mix!).
In particular: https://github.com/jpsamaroo/HSAArrays.jl
I don't know how far along it is (no documentation yet), but those are definite steps.
1. fresh install of juno (1.1.0)
2. add latest packages
3. something goes wrong. this time it's
p = destructure(model)
from example page givesUndefVarError: destructure not defined top-level scope at none:0
https://i.imgur.com/8fCpVrR.png
how the hell does anyone get anything when the codebase is shifting this fast!?
edit:
UndefVarError: JacVecOperator not defined
:(
also this
https://gist.github.com/makslevental/fe9144d196cade685d333cd...
i realize it's just a warning but it shows another undef error