"Fixed iteration" DiffEq solvers on the GPU?

Hey,

I am currently trying to (re)write some old hobby project into Julia (related to raytracing) and want to rely on the DiffEq solver interface, ideally inside a KernelAbstractions.jl kernel (as I am performing a lot of operations during and after solving that can not easily be expressed through the ensemble or callback API (without losing out on performance), but still want to stay on the GPU).

After some initial benchmarking, a problem I noticed was that ray bundles were having warp divergence issues (not unexpected!). Playing around with rel/abstol I also noticed that I can partially get rid of this and virtually none of the final rendered pixel values changing, so in theory I could get away with a fixed number of iterations inside a solver (i.e. hardcode/unroll the number of say, TSit5() iterations instead of a while loop, so there is zero warp divergence, and still being able to mark rays that exceeded error tolerances). This would get me some adaptivity in time-stepping, i.e. would be much better than just the naive fixed \Delta t approaches.

Does anyone know if such solvers exist anywhere? Would they even be compatible with the existing DiffEq.jl API or would I just be hacking my way into it?

[I know this is not really related to “real” DE solving since we in effect accept a possibly unbounded error. However, for my application (rays governed by a DE, landing on surface, followed by a discrete pixel being colored - as long the error is on the order of a single pixel’s angular resolution, its acceptable and I can just run unit tests vs a “real” solver) this would be a huge leap.]