Update ParallelSyncPSOKernel

Author: utkarsh530

src/solve.jl

@@ -35,6 +35,56 @@ function SciMLBase.init(
         prob, opt, particles, init_gbest)
 end
 
+function SciMLBase.init(
+        prob::OptimizationProblem, opt::ParallelPSOKernel, args...; kwargs...)
+    backend = opt.backend
+    @assert prob.u0 isa SArray
+
+    ## Bounds check
+    lb, ub = check_init_bounds(prob)
+    lb, ub = check_init_bounds(prob)
+    prob = remake(prob; lb = lb, ub = ub)
+
+    particles = KernelAbstractions.allocate(
+        backend, SPSOParticle{typeof(prob.u0), eltype(typeof(prob.u0))}, opt.num_particles)
+    kernel! = 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)
+
+    init_gbest = KernelAbstractions.allocate(backend, typeof(_init_gbest), (1,))
+    copyto!(init_gbest, [_init_gbest])
+    return PSOCache{
+        typeof(prob), typeof(opt), typeof(particles), typeof(init_gbest)}(
+        prob, opt, particles, init_gbest)
+end
+
+function SciMLBase.init(
+        prob::OptimizationProblem, opt::ParallelSyncPSOKernel, args...; kwargs...)
+    backend = opt.backend
+    @assert prob.u0 isa SArray
+
+    ## Bounds check
+    lb, ub = check_init_bounds(prob)
+    lb, ub = check_init_bounds(prob)
+    prob = remake(prob; lb = lb, ub = ub)
+
+    particles = KernelAbstractions.allocate(
+        backend, SPSOParticle{typeof(prob.u0), eltype(typeof(prob.u0))}, opt.num_particles)
+    kernel! = 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)
+
+    return PSOCache{
+        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
@@ -55,6 +105,23 @@ function reinit_cache!(cache::PSOCache, opt::ParallelPSOKernel)
     return nothing
 end
 
+function reinit_cache!(cache::PSOCache, opt::ParallelSyncPSOKernel)
+    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)
+
+    cache.gbest = init_gbest
+
+    return nothing
+end
+
 function SciMLBase.solve!(cache, args...; maxiters = 100, kwargs...)
     solve!(cache, cache.alg, args...; maxiters, kwargs...)
 end
@@ -78,6 +145,24 @@ function SciMLBase.solve!(
         stats = Optimization.OptimizationStats(; time = t1 - t0))
 end
 
+function SciMLBase.solve!(
+        cache::PSOCache, opt::ParallelSyncPSOKernel, args...; maxiters = 100, kwargs...)
+    prob = cache.prob
+    t0 = time()
+    gbest, particles = vectorized_solve!(prob,
+        init_gbest,
+        gpu_particles,
+        opt,
+        args...;
+        kwargs...)
+    t1 = time()
+
+    particles_positions = get_pos.(particles)
+    SciMLBase.build_solution(SciMLBase.DefaultOptimizationCache(prob.f, prob.p), opt,
+        gbest.position, prob.f(gbest.position, prob.p), original = particles_positions,
+        stats = Optimization.OptimizationStats(; time = t1 - t0))
+end
+
 function SciMLBase.solve(prob::OptimizationProblem, opt::ParallelPSOKernel,
         args...; maxiters = 100, kwargs...)
     solve!(init(prob, opt, args...; maxiters, kwargs...), opt)
@@ -99,34 +184,6 @@ function SciMLBase.__solve(prob::OptimizationProblem,
         stats = Optimization.OptimizationStats(; time = solve_time))
 end
 
-function pso_solve(prob::OptimizationProblem,
-        opt::ParallelPSOKernel,
-        args...;
-        kwargs...)
-    backend = opt.backend
-    @assert prob.u0 isa SArray
-    init_gbest, particles = init_particles(prob, opt, typeof(prob.u0))
-    # 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])
-
-    t0 = time()
-    gbest, particles = vectorized_solve!(prob,
-        gpu_init_gbest,
-        gpu_particles,
-        opt,
-        Val(opt.global_update),
-        args...;
-        kwargs...)
-    t1 = time()
-    gbest, particles, t1 - t0
-end
-
 function pso_solve(prob::OptimizationProblem,
         opt::ParallelPSOArray,
         args...;
@@ -150,25 +207,3 @@ function pso_solve(prob::OptimizationProblem, opt::SerialPSO, args...; kwargs...
     t1 = time()
     gbest, particles, t1 - t0
 end
-
-function pso_solve(prob::OptimizationProblem,
-        opt::ParallelSyncPSOKernel,
-        args...;
-        kwargs...)
-    backend = opt.backend
-    init_gbest, particles = init_particles(prob, opt, typeof(prob.u0))
-    particles_eltype = eltype(particles) === Float64 ? Float32 : eltype(particles)
-    gpu_particles = KernelAbstractions.allocate(backend, particles_eltype, size(particles))
-    copyto!(gpu_particles, particles)
-    init_gbest = init_gbest
-
-    t0 = time()
-    gbest, particles = vectorized_solve!(prob,
-        init_gbest,
-        gpu_particles,
-        opt,
-        args...;
-        kwargs...)
-    t1 = time()
-    gbest, particles, t1 - t0
-end