diff --git a/src/ParallelParticleSwarms.jl b/src/ParallelParticleSwarms.jl index 4324a95..ae704c7 100644 --- a/src/ParallelParticleSwarms.jl +++ b/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), diff --git a/src/init.jl b/src/init.jl index c6c3c6e..8139661 100644 --- a/src/init.jl +++ b/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) diff --git a/src/kernels.jl b/src/kernels.jl index 1895e64..62b0ea7 100644 --- a/src/kernels.jl +++ b/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 diff --git a/src/lowerlevel_solve.jl b/src/lowerlevel_solve.jl index 6b5b580..23ab04c 100644 --- a/src/lowerlevel_solve.jl +++ b/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 diff --git a/src/ode_pso.jl b/src/ode_pso.jl index de58aee..1c8255b 100644 --- a/src/ode_pso.jl +++ b/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) diff --git a/test/gpu.jl b/test/gpu.jl index 89ee3f5..fc0996e 100644 --- a/test/gpu.jl +++ b/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,6 +80,9 @@ 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) @@ -85,25 +90,88 @@ 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