Reduce over scalar costs when picking the best particle

src/ParallelParticleSwarms.jl

@@ -76,6 +76,37 @@ function Base.isless(
     return a.cost < b.cost
 end
 
+# Reduce over scalar costs instead of whole particles: large-dimension particles
+# overflow GPU shared memory in `minimum`, giving a zero-thread launch.
+# Candidates are ranked by `isless` on (cost, index): NaN sorts after every number,
+# including Inf, so it only wins when all costs are NaN, and ties go to the lowest
+# index. An all-NaN population therefore yields the first particle.
+_cost(p::SPSOParticle) = p.best_cost
+_cost(g::SPSOGBest) = g.cost
+
+_min_cost_index(a, b) = isless(b, a) ? b : a
+
+# Identity for `_min_cost_index`: ranks after every (cost, index) candidate.
+_worst_cost(::Type{T}) where {T <: AbstractFloat} = T(NaN)
+_worst_cost(::Type{T}) where {T} = typemax(T)
+
+function _best_particle(ps)
+    costs = map(_cost, ps)
+    init = (_worst_cost(eltype(costs)), typemax(Int))
+    _, i = mapreduce(tuple, _min_cost_index, costs, eachindex(costs); init)
+    return Array(view(ps, i:i))[1]
+end
+
+function _best_particle(ps::Array)
+    best = firstindex(ps)
+    @inbounds for i in eachindex(ps)
+        if isless(_cost(ps[i]), _cost(ps[best]))
+            best = i
+        end
+    end
+    return @inbounds ps[best]
+end
+
 function Base.typemax(::Type{ParallelParticleSwarms.SPSOParticle{T1, T2}}) where {T1, T2}
     return ParallelParticleSwarms.SPSOParticle{T1, T2}(
         similar(T1),

src/init.jl

@@ -29,7 +29,7 @@ function __init!(
         particles, qmc_samples, prob, opt, typeof(prob.u0), T; ndrange = opt.num_particles
     )
 
-    best_particle = minimum(particles)
+    best_particle = _best_particle(particles)
     init_gbest = SPSOGBest(best_particle.best_position, best_particle.best_cost)
 
     return particles, init_gbest
@@ -46,7 +46,7 @@ function __init!(
 
     kernel!(particles, prob, opt, typeof(prob.u0), T; ndrange = opt.num_particles)
 
-    best_particle = minimum(particles)
+    best_particle = _best_particle(particles)
 
     init_gbest = SPSOGBest(best_particle.best_position, best_particle.best_cost)
 

src/kernels.jl

@@ -133,7 +133,10 @@ end
     costs = @localmem T2 (gs)
     idxs = @localmem Int32 (gs)
 
-    @inbounds costs[tidx] = convert(T2, Inf)
+    # Rank with `isless` so NaN sorts after every number. Padding (i > n) fills the
+    # block's highest slots with NaN, and a slot only takes a higher slot's candidate
+    # when it is strictly less, so padding never displaces a real particle.
+    @inbounds costs[tidx] = convert(T2, NaN)
     @inbounds idxs[tidx] = Int32(0)
 
     if i <= n
@@ -151,7 +154,7 @@ end
         half[1] = cld(nactive[1], 2)
         @synchronize
         if tidx <= nactive[1] - half[1]
-            @inbounds if costs[tidx + half[1]] < costs[tidx]
+            @inbounds if isless(costs[tidx + half[1]], costs[tidx])
                 costs[tidx] = costs[tidx + half[1]]
                 idxs[tidx] = idxs[tidx + half[1]]
             end

src/lowerlevel_solve.jl

@@ -32,7 +32,7 @@ function vectorized_solve!(
             w, opt;
             ndrange = padded_ndrange
         )
-        gbest = minimum(block_particles)
+        gbest = _best_particle(block_particles)
         w = w * wdamp
     end
 
@@ -60,7 +60,7 @@ function vectorized_solve!(
             w, opt;
             ndrange = length(gpu_particles)
         )
-        best_particle = minimum(gpu_particles)
+        best_particle = _best_particle(gpu_particles)
         gbest = SPSOGBest(best_particle.best_position, best_particle.best_cost)
         w = w * wdamp
     end
@@ -122,7 +122,7 @@ function vectorized_solve!(
         ndrange = padded_ndrange
     )
 
-    best_particle = minimum(gpu_particles)
+    best_particle = _best_particle(gpu_particles)
     return SPSOGBest(best_particle.best_position, best_particle.best_cost), gpu_particles
 end
 

src/ode_pso.jl

@@ -97,7 +97,7 @@ function parameter_estim_ode!(
 
         KernelAbstractions.synchronize(backend)
 
-        best_particle = minimum(gpu_particles)
+        best_particle = _best_particle(gpu_particles)
 
         KernelAbstractions.synchronize(backend)
 
@@ -156,7 +156,7 @@ function parameter_estim_ode!(
 
         KernelAbstractions.synchronize(backend)
 
-        best_particle = minimum(gpu_particles)
+        best_particle = _best_particle(gpu_particles)
 
         KernelAbstractions.synchronize(backend)
 

test/gpu.jl

@@ -60,7 +60,7 @@ end
 @testset "reduction is dimension-independent" begin
     Random.seed!(1234)
     n_particles = 512
-    for D in (3, 10, 30)
+    for D in (3, 10, 30, 128)
         lb = SVector{D, Float64}(ntuple(_ -> -5.0, Val(D)))
         ub = SVector{D, Float64}(ntuple(_ -> 5.0, Val(D)))
         x0 = zero(lb)
@@ -69,6 +69,8 @@ end
         for opt in (
                 ParallelSyncPSOKernel(n_particles; backend),
                 ParallelPSOKernel(n_particles; backend, global_update = true),
+                # Ragged population: the last workgroup is only partly filled.
+                ParallelPSOKernel(257; backend, global_update = false),
             )
             sol = solve(prob, opt; maxiters = 100)
             @test isfinite(sol.objective)
@@ -78,32 +80,98 @@ end
                 ParallelPSOKernel(
                     1024; backend, global_update = true, workgroupsize = 1024
                 ),
+                ParallelPSOKernel(
+                    1024; backend, global_update = false, workgroupsize = 1024
+                ),
             )
             sol = solve(prob, opt; maxiters = 20)
             @test isfinite(sol.objective)
         end
     end
 end
 
+@testset "best particle ranks NaN last and breaks ties by index" begin
+    # Each candidate's position encodes its index, so the winner is identifiable even
+    # when its cost is NaN.
+    cases = [
+        [NaN, Inf] => 2,
+        [Inf, NaN] => 1,
+        [NaN, 3.0, NaN, 1.0, 2.0] => 4,
+        [1.0, -Inf, NaN] => 2,
+        [NaN, NaN, NaN] => 1,
+        [9.0, 1.0, 1.0, 9.0] => 2,
+    ]
+    pos(i) = SVector(Float64(i), Float64(i))
+    for (costs, winner) in cases
+        particles = [
+            ParallelParticleSwarms.SPSOParticle(pos(i), pos(i), c, pos(i), c)
+                for (i, c) in enumerate(costs)
+        ]
+        gbests = [
+            ParallelParticleSwarms.SPSOGBest(pos(i), c) for (i, c) in enumerate(costs)
+        ]
+        best = ParallelParticleSwarms._best_particle(
+            ParallelParticleSwarms.adapt(backend, particles)
+        )
+        @test best.best_position == pos(winner)
+        @test isequal(best.best_cost, costs[winner])
+        best = ParallelParticleSwarms._best_particle(
+            ParallelParticleSwarms.adapt(backend, gbests)
+        )
+        @test best.position == pos(winner)
+        @test isequal(best.cost, costs[winner])
+    end
+end
+
+@testset "NaN costs never win the block reduction" begin
+    Random.seed!(1234)
+    D = 4
+    lb = SVector{D, Float64}(ntuple(_ -> -5.0, Val(D)))
+    ub = SVector{D, Float64}(ntuple(_ -> 5.0, Val(D)))
+    x0 = zero(lb)
+    # About half the initial particles start (and stay) at a NaN personal best.
+    optf = OptimizationFunction{false}(
+        (x, p) -> x[1] < 0 ? NaN : sum(abs2, x), SciMLBase.NoAD()
+    )
+    prob = OptimizationProblem{false}(optf, x0, nothing; lb, ub)
+    for opt in (
+            ParallelSyncPSOKernel(64; backend, workgroupsize = 64),
+            ParallelSyncPSOKernel(100; backend, workgroupsize = 32),
+            ParallelPSOKernel(100; backend, global_update = true, workgroupsize = 32),
+            ParallelPSOKernel(100; backend, global_update = false, workgroupsize = 32),
+        )
+        sol = solve(prob, opt; maxiters = 5)
+        @test isfinite(sol.objective)
+        @test sol.u[1] >= 0
+    end
+end
+
 @testset "block argmin matches brute force" begin
     Random.seed!(1234)
     D = 10
     lb = SVector{D, Float64}(ntuple(_ -> -5.0, Val(D)))
     ub = SVector{D, Float64}(ntuple(_ -> 5.0, Val(D)))
     x0 = zero(lb)
-    optf = OptimizationFunction{false}((x, p) -> sum(abs2, x), SciMLBase.NoAD())
-    prob = OptimizationProblem{false}(optf, x0, nothing; lb, ub)
-    for n in (10, 100, 5000),
+    # `floor(abs(x[1]))` is 0 across a whole slab, so many particles tie at the minimum.
+    for f in ((x, p) -> sum(abs2, x), (x, p) -> floor(abs(x[1]))),
+            n in (10, 100, 5000),
             opt_f in (
                 n -> ParallelSyncPSOKernel(n; backend),
                 n -> ParallelPSOKernel(n; backend, global_update = true),
+                n -> ParallelPSOKernel(n; backend, global_update = false),
             )
+        optf = OptimizationFunction{false}(f, SciMLBase.NoAD())
+        prob = OptimizationProblem{false}(optf, x0, nothing; lb, ub)
         cache = init(prob, opt_f(n))
         sol = solve!(cache; maxiters = 5)
-        best = argmin(p -> p.best_cost, Array(cache.particles))
-        @test sol.u == best.best_position
+        particles = Array(cache.particles)
+        best_cost = minimum(p -> p.best_cost, particles)
+        # The reductions need not pick the lowest-index minimizer: the block tree and
+        # the shared global best (which keeps its incumbent) may settle on any of them.
+        minimizers = [p.best_position for p in particles if p.best_cost == best_cost]
+        @test sol.u in minimizers
         # host f vs FMA-contracted kernel cost: each n-term sum of squares is within n*eps/2
-        @test isapprox(sol.objective, best.best_cost; rtol = length(sol.u) * eps())
+        @test isapprox(sol.objective, best_cost; rtol = length(sol.u) * eps())
     end
 end