Fix GPU crossover with CSR kernels and broadcast-safe helpers.

Run effective bounds on device via row kernels (CSR/ELL) with atomic bound
merging, keep CSC on CPU, and vectorize crossover_n_changed so docs and CI
GPU tutorial examples no longer hit scalar indexing errors.

Author: Charles Vielzeuf

src/components/crossover.jl

@@ -55,10 +55,16 @@ function crossover_kkt_acceptable(
     return true
 end
 
-"""True if coordinate `j` is on a finite box bound (within `atol`)."""
-function _crossover_at_box_bound(xj, lj, uj; atol::Real = 1.0e-12)
-    return (isfinite(lj) && abs(xj - lj) <= atol) ||
-        (isfinite(uj) && abs(xj - uj) <= atol)
+"""GPU-safe mask: coordinate `j` is on a finite box bound (within `atol`)."""
+function _crossover_at_box_mask(
+        x::AbstractVector,
+        lv::AbstractVector,
+        uv::AbstractVector;
+        atol::Real = 1.0e-12,
+    )
+    at_l = isfinite.(lv) .& (abs.(x .- lv) .<= atol)
+    at_u = isfinite.(uv) .& (abs.(x .- uv) .<= atol)
+    return at_l .| at_u
 end
 
 """
@@ -78,45 +84,257 @@ function crossover_effective_bounds(
     ) where {T}
     lv_eff = copy(milp.lv)
     uv_eff = copy(milp.uv)
-    m, n = size(milp.A)
-    at_box = [
-        _crossover_at_box_bound(x[j], milp.lv[j], milp.uv[j]; atol = bound_atol) for j in 1:n
-    ]
+    at_box = _crossover_at_box_mask(x, milp.lv, milp.uv; atol = bound_atol)
+    _crossover_effective_bounds!(
+        lv_eff,
+        uv_eff,
+        milp.A,
+        milp.lc,
+        milp.uc,
+        x,
+        at_box;
+        eq_atol,
+    )
+    return lv_eff, uv_eff
+end
+
+function _crossover_effective_bounds!(
+        lv_eff,
+        uv_eff,
+        A::SparseMatrixCSC{T},
+        lc,
+        uc,
+        x,
+        at_box;
+        eq_atol::Real = 1.0e-12,
+    ) where {T}
+    m, n = size(A)
+    at_box_cpu = Vector(at_box)
+    x_cpu = Vector(x)
+    lv_cpu = Vector(lv_eff)
+    uv_cpu = Vector(uv_eff)
+    lc_cpu = Vector(lc)
+    uc_cpu = Vector(uc)
     for i in 1:m
-        isapprox(milp.lc[i], milp.uc[i]; atol = eq_atol) || continue
+        isapprox(lc_cpu[i], uc_cpu[i]; atol = eq_atol) || continue
         free = Int[]
-        slack = milp.lc[i]
+        slack = lc_cpu[i]
         @inbounds for j in 1:n
-            aij = milp.A[i, j]
+            aij = A[i, j]
             aij == 0 && continue
-            if at_box[j]
-                slack -= aij * x[j]
+            if at_box_cpu[j]
+                slack -= aij * x_cpu[j]
             else
                 push!(free, j)
             end
         end
         length(free) == 1 || continue
         j = only(free)
-        aij = milp.A[i, j]
+        aij = A[i, j]
         if aij > 0
             implied = slack / aij
-            if !isfinite(uv_eff[j])
-                uv_eff[j] = implied
+            if !isfinite(uv_cpu[j])
+                uv_cpu[j] = implied
             else
-                uv_eff[j] = min(uv_eff[j], implied)
+                uv_cpu[j] = min(uv_cpu[j], implied)
             end
         elseif aij < 0
             implied = slack / aij
-            if !isfinite(lv_eff[j])
-                lv_eff[j] = implied
+            if !isfinite(lv_cpu[j])
+                lv_cpu[j] = implied
             else
-                lv_eff[j] = max(lv_eff[j], implied)
+                lv_cpu[j] = max(lv_cpu[j], implied)
+            end
+        end
+    end
+    lv_eff .= lv_cpu
+    uv_eff .= uv_cpu
+    return lv_eff, uv_eff
+end
+
+@kernel function crossover_effective_bounds_csr!(
+        lv_eff::DenseVector{T},
+        uv_eff::DenseVector{T},
+        A_rowptr::DenseVector{Ti},
+        A_colval::DenseVector{Ti},
+        A_nzval::DenseVector{T},
+        lc::DenseVector{T},
+        uc::DenseVector{T},
+        x::DenseVector{T},
+        at_box::DenseVector{Bool},
+        eq_atol::T,
+    ) where {T, Ti}
+    i = @index(Global, Linear)
+    if abs(lc[i] - uc[i]) <= eq_atol
+        slack = lc[i]
+        free_j = zero(Ti)
+        free_aij = zero(T)
+        n_free = zero(Ti)
+        for k in A_rowptr[i]:(A_rowptr[i + Ti(1)] - Ti(1))
+            j = A_colval[k]
+            aij = A_nzval[k]
+            if at_box[j]
+                slack -= aij * x[j]
+            else
+                n_free += Ti(1)
+                if n_free == Ti(1)
+                    free_j = j
+                    free_aij = aij
+                end
+            end
+        end
+        if n_free == Ti(1)
+            j = free_j
+            aij = free_aij
+            if aij > zero(T)
+                implied = slack / aij
+                old = Atomix.@atomic uv_eff[j]
+                while implied < old
+                    (old2, success) = Atomix.@atomicreplace uv_eff[j] old => implied
+                    success && break
+                    old = old2
+                end
+            elseif aij < zero(T)
+                implied = slack / aij
+                old = Atomix.@atomic lv_eff[j]
+                while implied > old
+                    (old2, success) = Atomix.@atomicreplace lv_eff[j] old => implied
+                    success && break
+                    old = old2
+                end
             end
         end
     end
+end
+
+function _crossover_effective_bounds!(
+        lv_eff,
+        uv_eff,
+        A::GPUSparseMatrixCSR,
+        lc,
+        uc,
+        x,
+        at_box;
+        eq_atol::Real = 1.0e-12,
+    )
+    backend = common_backend(lv_eff, uv_eff, A, lc, uc, x, at_box)
+    eq_tol = eltype(x)(eq_atol)
+    kernel! = crossover_effective_bounds_csr!(backend)
+    kernel!(
+        lv_eff,
+        uv_eff,
+        A.rowptr,
+        A.colval,
+        A.nzval,
+        lc,
+        uc,
+        x,
+        at_box,
+        eq_tol;
+        ndrange = size(A, 1),
+    )
     return lv_eff, uv_eff
 end
 
+@kernel function crossover_effective_bounds_ell!(
+        lv_eff::DenseVector{T},
+        uv_eff::DenseVector{T},
+        A_colval::AbstractMatrix{Ti},
+        A_nzval::AbstractMatrix{T},
+        lc::DenseVector{T},
+        uc::DenseVector{T},
+        x::DenseVector{T},
+        at_box::DenseVector{Bool},
+        eq_atol::T,
+    ) where {T, Ti}
+    i = @index(Global, Linear)
+    if abs(lc[i] - uc[i]) <= eq_atol
+        slack = lc[i]
+        free_j = zero(Ti)
+        free_aij = zero(T)
+        n_free = zero(Ti)
+        for k in axes(A_colval, 2)
+            j = A_colval[i, k]
+            if j != zero(Ti)
+                aij = A_nzval[i, k]
+                if at_box[j]
+                    slack -= aij * x[j]
+                else
+                    n_free += Ti(1)
+                    if n_free == Ti(1)
+                        free_j = j
+                        free_aij = aij
+                    end
+                end
+            end
+        end
+        if n_free == Ti(1)
+            j = free_j
+            aij = free_aij
+            if aij > zero(T)
+                implied = slack / aij
+                old = Atomix.@atomic uv_eff[j]
+                while implied < old
+                    (old2, success) = Atomix.@atomicreplace uv_eff[j] old => implied
+                    success && break
+                    old = old2
+                end
+            elseif aij < zero(T)
+                implied = slack / aij
+                old = Atomix.@atomic lv_eff[j]
+                while implied > old
+                    (old2, success) = Atomix.@atomicreplace lv_eff[j] old => implied
+                    success && break
+                    old = old2
+                end
+            end
+        end
+    end
+end
+
+function _crossover_effective_bounds!(
+        lv_eff,
+        uv_eff,
+        A::GPUSparseMatrixELL,
+        lc,
+        uc,
+        x,
+        at_box;
+        eq_atol::Real = 1.0e-12,
+    )
+    backend = common_backend(lv_eff, uv_eff, A, lc, uc, x, at_box)
+    eq_tol = eltype(x)(eq_atol)
+    kernel! = crossover_effective_bounds_ell!(backend)
+    kernel!(
+        lv_eff,
+        uv_eff,
+        A.colval,
+        A.nzval,
+        lc,
+        uc,
+        x,
+        at_box,
+        eq_tol;
+        ndrange = size(A, 1),
+    )
+    return lv_eff, uv_eff
+end
+
+function _crossover_effective_bounds!(
+        lv_eff,
+        uv_eff,
+        A::GPUSparseMatrixCOO,
+        lc,
+        uc,
+        x,
+        at_box;
+        eq_atol::Real = 1.0e-12,
+    )
+    backend = get_backend(x)
+    A_csr = adapt(backend, GPUSparseMatrixCSR(SparseMatrixCSC(A)))
+    return _crossover_effective_bounds!(lv_eff, uv_eff, A_csr, lc, uc, x, at_box; eq_atol)
+end
+
 """
     crossover_threshold!(x, lv, uv, params::CrossoverParameters)
 
@@ -167,7 +385,7 @@ end
 
 """Count coordinates changed by crossover (exact compare)."""
 function crossover_n_changed(x_after, x_before)
-    return count(i -> x_after[i] != x_before[i], eachindex(x_before))
+    return sum(x_after .!= x_before)
 end
 
 """

test/components/crossover.jl

@@ -43,6 +43,26 @@ end
     @test isinf(uv_eff[2])
 end
 
+@testset "crossover_effective_bounds GPU" begin
+    using Adapt
+    using JLArrays: JLBackend, JLArray
+
+    milp = MILP(;
+        c = [1.0, 2.0],
+        lv = [0.0, 0.0],
+        uv = [Inf, Inf],
+        A = sparse([1.0 1.0]),
+        lc = [1.0],
+        uc = [1.0],
+    )
+    milp_gpu = adapt(JLBackend(), CoolPDLP.set_matrix_type(CoolPDLP.GPUSparseMatrixCSR, milp))
+    x = adapt(JLBackend(), [1.0 - 1.0e-7, 0.0])
+    lv_eff, uv_eff = CoolPDLP.crossover_effective_bounds(milp_gpu, x)
+    @test lv_eff isa JLArray
+    @test isapprox(Array(uv_eff)[1], 1.0; atol = 1.0e-12)
+    @test isinf(Array(uv_eff)[2])
+end
+
 @testset "crossover improves fraction_at_bounds" begin
     milp = MILP(;
         c = [1.0, 2.0],
@@ -132,7 +152,7 @@ end
 end
 
 @testset "crossover rollback" begin
-  @testset "crossover_kkt_acceptable edge cases" begin
+    @testset "crossover_kkt_acceptable edge cases" begin
         err_ok = _kkt_err(1.0e-5, 1.0e-5, 1.0e-5)
         err_cert_fail = _kkt_err(2.0e-4, 1.0e-5, 1.0e-5)
         err_regress = _kkt_err(1.1e-5, 1.0e-5, 1.0e-5)