Closed cgarling closed 3 months ago
Did a short experiment with Metal.jl but didn't have much luck. Trouble with allocations made the GPU kernels slower than the CPU alternatives. That was on v0.5.2, v1.0.0 might be better. Low priority for now.
Now that we have implementations of composite!
and ∇loglikelihood!
that use matrix-vector products through LinearAlgebra.mul!
, this may be easier. When inputting GPU array types, mul!
should have an overload defined to perform the calculation on the GPU. As such, it may be much easier to get performant GPU evaluation with these methods than with our hand-written loops.
For large isochrone grids, the vast majority of the computation time is spent in the computational kernels in
fitting/fitting_base.jl
doing things like constructing the complex model from a linear combination of single-population models and computing the associated loglikelihood and gradient. The LL and grad contain slightly more complex logic, but thecomposite!
method is just linear algebra and could be gainfully accelerated on GPUs. Would be worth trying with Metal.jl, probably through a package extension: see dev docs, discussion on Julia <1.9.