Complete refactor of CPU and GPU PSO codes

Author: utkarsh530

src/PSOGPU.jl

@@ -6,44 +6,51 @@ import DiffEqGPU: GPUTsit5, vectorized_asolve, make_prob_compatible
 
 ## Use lb and ub either as StaticArray or pass them separately as CuArrays
 ## Passing as CuArrays makes more sense, or maybe SArray? The based on no. of dimension
-struct PSOParticle{T1, T2 <: eltype(T1)}
+struct SPSOParticle{T1, T2 <: eltype(T1)}
     position::T1
     velocity::T1
     cost::T2
     best_position::T1
     best_cost::T2
 end
-struct PSOGBest{T1, T2 <: eltype(T1)}
+struct SPSOGBest{T1, T2 <: eltype(T1)}
     position::T1
     cost::T2
 end
 
+mutable struct MPSOParticle{T}
+    position::AbstractArray{T}
+    velocity::AbstractArray{T}
+    cost::T
+    best_position::AbstractArray{T}
+    best_cost::T
+end
+mutable struct MPSOGBest{T}
+    position::AbstractArray{T}
+    cost::T
+end
+
 ## required overloads for min or max computation on particles
-function Base.isless(a::PSOGPU.PSOParticle{T1, T2},
-        b::PSOGPU.PSOParticle{T1, T2}) where {T1, T2}
+function Base.isless(a::PSOGPU.SPSOParticle{T1, T2},
+        b::PSOGPU.SPSOParticle{T1, T2}) where {T1, T2}
     a.best_cost < b.best_cost
 end
 
-function Base.typemax(::Type{PSOGPU.PSOParticle{T1, T2}}) where {T1, T2}
-    PSOGPU.PSOParticle{T1, T2}(similar(T1),
+function Base.typemax(::Type{PSOGPU.SPSOParticle{T1, T2}}) where {T1, T2}
+    PSOGPU.SPSOParticle{T1, T2}(similar(T1),
         similar(T1),
         typemax(T2),
         similar(T1),
         typemax(T2))
 end
 
 include("./algorithms.jl")
-include("./pso_cpu.jl")
-include("./pso_gpu.jl")
-include("./pso_async_gpu.jl")
 include("./utils.jl")
-include("./pso_sync_gpu.jl")
 include("./ode_pso.jl")
 include("./kernels.jl")
 include("./lowerlevel_solve.jl")
 include("./solve.jl")
 
 export ParallelPSOKernel,
     ParallelSyncPSOKernel, ParallelPSOArray, SerialPSO, OptimizationProblem, solve
-
 end

src/algorithms.jl

@@ -87,8 +87,8 @@ struct SerialPSO <: PSOAlogrithm
 end
 
 function ParallelPSOKernel(num_particles::Int;
-        async = false, backend = CPU())
-    ParallelPSOKernel(num_particles, async, backend)
+        global_update = false, backend = CPU())
+    ParallelPSOKernel(num_particles, global_update, backend)
 end
 
 function ParallelSyncPSOKernel(num_particles::Int;

src/kernels.jl

@@ -56,3 +56,30 @@ end
 
     @inbounds gpu_particles[i] = particle
 end
+
+@kernel function update_particle_states_async!(prob,
+        gpu_particles,
+        gbest_ref,
+        w, wdamp, maxiters;
+        c1 = 1.4962f0,
+        c2 = 1.4962f0)
+    i = @index(Global, Linear)
+
+    gbest = gbest_ref[1]
+
+    ## Access the particle
+    @inbounds particle = gpu_particles[i]
+
+    ## Run all generations
+    for i in 1:maxiters
+        particle = update_particle_state(particle, prob, gbest, w, c1, c2)
+        if particle.best_cost < gbest.cost
+            @set! gbest.position = particle.best_position
+            @set! gbest.cost = particle.best_cost
+        end
+        w = w * wdamp
+    end
+
+    @inbounds gpu_particles[i] = particle
+    @inbounds gbest_ref[1] = gbest
+end

src/lowerlevel_solve.jl

@@ -16,7 +16,7 @@ function vectorized_solve!(prob,
             w, opt;
             ndrange = length(gpu_particles))
         best_particle = minimum(gpu_particles)
-        gbest = PSOGBest(best_particle.position, best_particle.best_cost)
+        gbest = SPSOGBest(best_particle.position, best_particle.best_cost)
         w = w * wdamp
     end
 
@@ -25,7 +25,7 @@ end
 
 function vectorized_solve!(prob,
         gbest,
-        gpu_particles, opt::ParallelPSOKernel;
+        gpu_particles, opt::ParallelPSOKernel, ::Val{true};
         maxiters = 100,
         w = 0.7298f0,
         wdamp = 1.0f0,
@@ -45,3 +45,97 @@ function vectorized_solve!(prob,
 
     return Array(gbest)[1]
 end
+
+function vectorized_solve!(prob,
+        gbest,
+        gpu_particles, opt::ParallelPSOKernel, ::Val{false};
+        maxiters = 100,
+        w = 0.7298f0,
+        wdamp = 1.0f0,
+        debug = false)
+    backend = get_backend(gpu_particles)
+
+    kernel = update_particle_states_async!(backend)
+    kernel(prob, gpu_particles, gbest, w, wdamp, maxiters; ndrange = length(gpu_particles))
+
+    best_particle = minimum(gpu_particles)
+    return SPSOGBest(best_particle.best_position, best_particle.best_cost)
+end
+
+function vectorized_solve!(prob, gbest,
+        particles, opt::ParallelPSOArray;
+        maxiters = 100,
+        w = 0.7298f0,
+        wdamp = 1.0f0,
+        c1 = 1.4962f0,
+        c2 = 1.4962f0,
+        verbose = false)
+    cost_func = prob.f
+    num_particles = length(particles)
+    rand_eltype = eltype(particles[1].velocity)
+    # main loop
+
+    for iter in 1:maxiters
+        Threads.@threads for i in 1:num_particles
+            particles[i].velocity .= w .* particles[i].velocity .+
+                                     c1 .* rand.(rand_eltype) .*
+                                     (particles[i].best_position .-
+                                      particles[i].position) .+
+                                     c2 .* rand.(rand_eltype) .*
+                                     (gbest.position .- particles[i].position)
+
+            particles[i].position .= particles[i].position .+ particles[i].velocity
+            particles[i].position .= max.(particles[i].position, prob.lb)
+            particles[i].position .= min.(particles[i].position, prob.ub)
+
+            particles[i].cost = cost_func(particles[i].position, prob.p)
+
+            if particles[i].cost < particles[i].best_cost
+                copy!(particles[i].best_position, particles[i].position)
+                particles[i].best_cost = particles[i].cost
+
+                ## Possible race condition here
+                if particles[i].best_cost < gbest.cost
+                    copy!(gbest.position, particles[i].best_position)
+                    gbest.cost = particles[i].best_cost
+                end
+            end
+        end
+        w = w * wdamp
+    end
+    gbest
+end
+
+function update_particle_states_cpu!(prob, particles, gbest_ref, w; c1 = 1.4962f0,
+        c2 = 1.4962f0)
+    gbest = gbest_ref[]
+
+    for i in eachindex(particles)
+        @inbounds particle = particles[i]
+        particle = update_particle_state(particle, prob, gbest, w, c1, c2)
+
+        if particle.best_cost < gbest.cost
+            @set! gbest.position = particle.best_position
+            @set! gbest.cost = particle.best_cost
+        end
+
+        particles[i] = particle
+    end
+    gbest_ref[] = gbest
+    return nothing
+end
+
+function vectorized_solve!(prob,
+        gbest,
+        particles, opt::SerialPSO;
+        maxiters = 100,
+        w = 0.7298f0,
+        wdamp = 1.0f0,
+        debug = false)
+    sol_ref = Ref(gbest)
+    for i in 1:maxiters
+        update_particle_states_cpu!(prob, particles, sol_ref, w)
+        w = w * wdamp
+    end
+    return sol_ref[]
+end

src/ode_pso.jl

@@ -83,13 +83,13 @@ function parameter_estim_ode!(prob::ODEProblem,
         update_costs!(losses, gpu_particles; ndrange = length(losses))
 
         best_particle = minimum(gpu_particles,
-            init = PSOGPU.PSOParticle(gbest.position,
+            init = PSOGPU.SPSOParticle(gbest.position,
                 gbest.position,
                 gbest.cost,
                 gbest.position,
                 gbest.cost))
 
-        gbest = PSOGPU.PSOGBest(best_particle.best_position, best_particle.best_cost)
+        gbest = PSOGPU.SPSOGBest(best_particle.best_position, best_particle.best_cost)
         w = w * wdamp
     end
     return gbest