Parallelize initialization of particles

Author: utkarsh530

src/PSOGPU.jl

@@ -3,6 +3,7 @@ module PSOGPU
 using SciMLBase, StaticArrays, Setfield, KernelAbstractions
 using QuasiMonteCarlo, Optimization, SimpleNonlinearSolve, ForwardDiff
 import Adapt
+import Adapt: adapt
 import Enzyme: autodiff_deferred, Active, Reverse
 import KernelAbstractions: @atomic, @atomicreplace, @atomicswap
 using QuasiMonteCarlo

src/solve.jl

@@ -14,26 +14,49 @@ function SciMLBase.init(
     backend = opt.backend
     @assert prob.u0 isa SArray
 
-    ## initialize cache
-
     ## Bounds check
     lb, ub = check_init_bounds(prob)
     lb, ub = check_init_bounds(prob)
     prob = remake(prob; lb = lb, ub = ub)
 
-    init_gbest, particles = init_particles(prob, opt, typeof(prob.u0))
+    particles = KernelAbstractions.allocate(
+        backend, SPSOParticle{typeof(prob.u0), eltype(typeof(prob.u0))}, opt.num_particles)
+    kernel! = gpu_init_particles!(backend)
 
-    # TODO: Do the equivalent of cu()/roc()
-    particles_eltype = eltype(particles) === Float64 ? Float32 : eltype(particles)
-    gpu_particles = KernelAbstractions.allocate(backend,
-        particles_eltype,
-        size(particles))
-    copyto!(gpu_particles, particles)
-    gpu_init_gbest = KernelAbstractions.allocate(backend, typeof(init_gbest), (1,))
-    copyto!(gpu_init_gbest, [init_gbest])
+    kernel!(particles, prob, opt, typeof(prob.u0); ndrange = opt.num_particles)
+
+    best_particle = minimum(particles)
+    _init_gbest = SPSOGBest(best_particle.best_position, best_particle.best_cost)
+
+    init_gbest = KernelAbstractions.allocate(backend, typeof(_init_gbest), (1,))
+    copyto!(init_gbest, [_init_gbest])
     return PSOCache{
-        typeof(prob), typeof(opt), typeof(gpu_particles), typeof(gpu_init_gbest)}(
-        prob, opt, gpu_particles, gpu_init_gbest)
+        typeof(prob), typeof(opt), typeof(particles), typeof(init_gbest)}(
+        prob, opt, particles, init_gbest)
+end
+
+function SciMLBase.reinit!(cache::PSOCache; kwargs...)
+    reinit_cache!(cache, cache.alg)
+end
+
+function reinit_cache!(cache::PSOCache, opt::ParallelPSOKernel)
+    prob = cache.prob
+    backend = opt.backend
+    particles = cache.particles
+
+    kernel! = PSOGPU.gpu_init_particles!(backend)
+    kernel!(particles, prob, opt, typeof(prob.u0); ndrange = opt.num_particles)
+
+    best_particle = minimum(particles)
+    _init_gbest = SPSOGBest(best_particle.best_position, best_particle.best_cost)
+
+    copyto!(cache.gbest, [_init_gbest])
+
+    return nothing
+end
+
+function SciMLBase.solve!(cache, args...; maxiters = 100, kwargs...)
+    solve!(cache, cache.alg, args...; maxiters, kwargs...)
 end
 
 function SciMLBase.solve!(

src/utils.jl

@@ -11,6 +11,59 @@ function uniform(dim::Int, lb::AbstractArray{T}, ub::AbstractArray{T}) where {T}
     return arr
 end
 
+@kernel function gpu_init_particles!(particles, prob, opt, ::Type{T}) where {T <: SArray}
+    i = @index(Global, Linear)
+
+    dim = length(prob.u0)
+    lb = prob.lb
+    ub = prob.ub
+    cost_func = prob.f
+    p = prob.p
+    num_particles = opt.num_particles
+
+    if lb === nothing || (all(isinf, lb) && all(isinf, ub))
+        gbest_position = StaticArrays.sacollect(T,
+            ifelse(
+                abs(prob.u0[i]) > 0, prob.u0[i] + rand(eltype(prob.u0)) * abs(prob.u0[i]),
+                rand(eltype(prob.u0))) for i in 1:dim)
+    else
+        gbest_position = StaticArrays.sacollect(T, uniform_itr(dim, lb, ub))
+    end
+
+    gbest_position = convert(T, gbest_position)
+    gbest_cost = cost_func(gbest_position, p)
+    if !isnothing(prob.f.cons)
+        penalty = calc_penalty(gbest_position, prob, 1, opt.θ, opt.γ, opt.h)
+        gbest_cost = cost_func(gbest_position, p) + penalty
+    else
+        gbest_cost = cost_func(gbest_position, p)
+    end
+    gbest_cost = cost_func(gbest_position, p)
+
+    if lb === nothing || (all(isinf, lb) && all(isinf, ub))
+        position = StaticArrays.sacollect(T,
+            ifelse(abs(prob.u0[i]) > 0,
+                prob.u0[i] + rand(eltype(prob.u0)) * abs(prob.u0[i]),
+                rand(eltype(prob.u0))) for i in 1:dim)
+    else
+        position = StaticArrays.sacollect(T, uniform_itr(dim, lb, ub))
+    end
+
+    velocity = zero(T)
+
+    if !isnothing(prob.f.cons)
+        penalty = calc_penalty(position, prob, 1, opt.θ, opt.γ, opt.h)
+        cost = cost_func(position, p) + penalty
+    else
+        cost = cost_func(position, p)
+    end
+
+    best_position = position
+    best_cost = cost
+    @inbounds particles[i] = SPSOParticle(
+        position, velocity, cost, best_position, best_cost)
+end
+
 function init_particles!(particles, prob, opt, ::Type{T}) where {T <: SArray}
     dim = length(prob.u0)
     lb = prob.lb