3 ms·
This is precisely correct. You just have to make sure in the thousands of lines of code that you only use higher level building blocks and make no assumptions g
by ChrisRackauckas 7y ago
This is precisely correct. You just have to make sure in the thousands of lines of code that you only use higher level building blocks and make no assumptions greater than what a GPU-based array allows. GPU-based arrays don't even want to allow you to do scalar indexing for example, since every scalar indexing operation is a data transfer back to the CPU. So, with those rules in mind, getting the stiff ODE solvers to automatically work with GPU-based arrays amounted to:
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/init.jl#L104-L110 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!).