cgarling / StarFormationHistories.jl

Fitting astrophysical star formation histories via CMD modelling.
MIT License
3 stars 0 forks source link

GPU support #23

Closed cgarling closed 3 months ago

cgarling commented 7 months ago

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 the composite! 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.

cgarling commented 7 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.

cgarling commented 6 months ago

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.