alt Hacker NewsVery interesting stuff. How would I get GPU offload working? I have a rather complex scientific code I'm working on with JAX. Most of it can be expressed well with JAX's programming model, but the last 10% really sucks. It's still worth it so I don't have to mess around with offload onto whatever XPU flavor of the week. But going to C++ would really make my life easier, as long as I could use e.g. Kokkos.
Replies
dionhaefner • today at 4:55 PM
No idea lol. I assume it’s possible since both Enzyme and GPU programming are pervasive in Julia. Let us know if you end up trying.
You can write the cuda kernels and directly differentiate them using Enzyme. Last time I checked KOKKOS was still WIP: https://github.com/EnzymeAD/Enzyme/issues?q=kokkos