Fix queue lock algorithm for ParallelPSOKernel

Author: utkarsh530

Project.toml

@@ -5,7 +5,6 @@ version = "1.0.0-DEV"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
-CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 DiffEqGPU = "071ae1c0-96b5-11e9-1965-c90190d839ea"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"

src/kernels.jl

@@ -47,26 +47,33 @@ end
 
     particle = @private SPSOParticle{T1, T2} 1
 
-    # Initialize cost to be Inf
     @inbounds particle[1] = gpu_particles[i]
-    best_queue[tidx] = SPSOGBest(particle[1].best_position,
-                                 convert(T2, Inf))
-    queue_num[1] = UInt32(0)
+    @inbounds gbest = gbest_ref[1]
+
+    # Initialize cost to be Inf
+    if tidx == 1
+        fill!(best_queue, SPSOGBest(gbest.position, convert(typeof(gbest.cost), Inf)))
+        queue_num[1] = UInt32(0)
+    end
 
     @synchronize
 
-    @inbounds particle[1] = gpu_particles[i]
-    gbest = @inbounds gbest_ref[1]
-    @inbounds particle[1] = update_particle_state(particle[1], prob, gbest, w, c1, c2, i, opt)
+    @inbounds particle[1] = update_particle_state(particle[1],
+        prob,
+        gbest,
+        w,
+        c1,
+        c2,
+        i,
+        opt)
 
     @synchronize
 
-    @inbounds particle[1] = gpu_particles[i]
     gbest = @inbounds gbest_ref[1]
     if particle[1].best_cost < gbest.cost
         queue_idx = @atomic queue_num[1] += UInt32(1)
         @inbounds best_queue[queue_idx] = SPSOGBest(particle[1].best_position,
-                                                    particle[1].best_cost)
+            particle[1].best_cost)
     end
 
     @synchronize
@@ -98,6 +105,7 @@ end
             @atomicswap lock[1] = 0
         end
     end
+    @inbounds gpu_particles[i] = particle[1]
 end
 
 @kernel function update_particle_states!(prob,

test/Project.toml

@@ -1,4 +1,5 @@
 [deps]
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

test/gpu.jl

@@ -31,11 +31,15 @@ end
 
     sol = solve(prob, ParallelPSOKernel(n_particles; backend), maxiters = 500)
 
+    @test prob.f(prob.u0, prob.p) > sol.objective
+
+    @test sol.objective < 6e-4
+
     @test sol.retcode == ReturnCode.Default
 
     sol = solve(prob,
         ParallelPSOKernel(n_particles; backend, global_update = false),
-        maxiters = 500)
+        maxiters = 1000)
 
     @test prob.f(prob.u0, prob.p) > sol.objective
 
@@ -47,5 +51,5 @@ end
 
     @test prob.f(prob.u0, prob.p) > sol.objective
 
-    @test sol.objective < 6e-1
+    @test sol.objective < 6e-4
 end