From dd11f2d656feb894b27f8b56429e6cd197a2d82f Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Tue, 13 Oct 2026 03:25:08 +0530 Subject: [PATCH 1/3] Reduce over scalar costs when picking the best particle GPU `minimum` over whole particles/gbests stages each element in shared memory, so large-dimension positions overflow it and the reduction launches with zero threads. Pick the best via `argmin` over costs instead, mapping NaN costs to typemax so they never win. This covers the sync kernel's block reduction over SPSOGBest as well as the particle reductions in init, the async/CPU solvers, and ODE parameter estimation. Adds a 128-dimension case to the dimension-independence GPU test. --- src/ParallelParticleSwarms.jl | 14 ++++++++++++++ src/init.jl | 4 ++-- src/lowerlevel_solve.jl | 6 +++--- src/ode_pso.jl | 4 ++-- test/gpu.jl | 2 +- 5 files changed, 22 insertions(+), 8 deletions(-) diff --git a/src/ParallelParticleSwarms.jl b/src/ParallelParticleSwarms.jl index 4324a95..854af55 100644 --- a/src/ParallelParticleSwarms.jl +++ b/src/ParallelParticleSwarms.jl @@ -76,6 +76,20 @@ 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. +# NaN costs map to typemax so they never win (argmin would otherwise pick NaN). +_cost(p::SPSOParticle) = p.best_cost +_cost(g::SPSOGBest) = g.cost +_ranked_cost(p) = (c = _cost(p); isnan(c) ? typemax(c) : c) + +function _best_particle(ps) + i = argmin(map(_ranked_cost, ps)) + return Array(view(ps, i:i))[1] +end + +_best_particle(ps::Array) = argmin(_ranked_cost, ps) + 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/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 529b160..6270ebb 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) From 73dee70c9f918ee372cc3b003f9422a0123ff19d Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Tue, 13 Oct 2026 03:36:02 +0530 Subject: [PATCH 2/3] Compare winning position exactly in block argmin test sol.objective is prob.f recomputed on the host, so it can differ from the device-computed best_cost by an ulp on CUDA (also seen on main). Check the copied position exactly and the cost approximately. --- test/gpu.jl | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/test/gpu.jl b/test/gpu.jl index 6270ebb..dec37d3 100644 --- a/test/gpu.jl +++ b/test/gpu.jl @@ -100,8 +100,11 @@ end ) cache = init(prob, opt_f(n)) sol = solve!(cache; maxiters = 5) - best = minimum(p.best_cost for p in Array(cache.particles)) - @test sol.objective == best + best = argmin(p -> p.best_cost, Array(cache.particles)) + # The winning position is copied exactly; `sol.objective` is recomputed on the + # host and can differ from the device-computed `best_cost` by rounding. + @test sol.u == best.best_position + @test sol.objective ≈ best.best_cost end end From c78d6481655573b592adfddbbda18854fe73c4a5 Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Sat, 24 Oct 2026 14:25:28 +0530 Subject: [PATCH 3/3] Rank NaN costs last in best-particle selection and accept tied minimizers - `_best_particle` reduces over (cost, index) with `isless`, so NaN ranks after every number including Inf ([NaN, Inf] now picks Inf), ties go to the lowest index, and an all-NaN population yields the first particle. The CPU path is a plain loop over the same ordering. - The sync block kernel compares with `isless` and pads with NaN, so a NaN particle no longer displaces a finite one within a block. - Tests: helper regressions for SPSOParticle and SPSOGBest ([NaN, Inf], its reverse, mixed finite/NaN, -Inf, all NaN, ties); a NaN-objective solve for every kernel variant; global_update = false (including a ragged 257-particle population) in the D = 128 test; the brute-force check accepts any tied minimizer and adds a tie-heavy objective. --- src/ParallelParticleSwarms.jl | 25 +++++++++-- src/kernels.jl | 7 ++- test/gpu.jl | 82 ++++++++++++++++++++++++++++++++--- 3 files changed, 101 insertions(+), 13 deletions(-) diff --git a/src/ParallelParticleSwarms.jl b/src/ParallelParticleSwarms.jl index 854af55..ae704c7 100644 --- a/src/ParallelParticleSwarms.jl +++ b/src/ParallelParticleSwarms.jl @@ -78,17 +78,34 @@ end # Reduce over scalar costs instead of whole particles: large-dimension particles # overflow GPU shared memory in `minimum`, giving a zero-thread launch. -# NaN costs map to typemax so they never win (argmin would otherwise pick NaN). +# 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 -_ranked_cost(p) = (c = _cost(p); isnan(c) ? typemax(c) : c) + +_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) - i = argmin(map(_ranked_cost, 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 -_best_particle(ps::Array) = argmin(_ranked_cost, ps) +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}( 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/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