fix: make hot loop alloc-free and test it

.github/workflows/Test.yml

@@ -6,7 +6,7 @@ on:
       - main
   pull_request:
 
-# needed to allow julia-actions/cache to delete old caches that it has created
+    # needed to allow julia-actions/cache to delete old caches that it has created
 permissions:
   actions: write
   contents: read
@@ -22,13 +22,18 @@ jobs:
     name: Julia ${{ matrix.version }} - ${{ matrix.group }}
     runs-on: ubuntu-latest
     strategy:
+      fail-fast: false
       matrix:
         version:
           - 'lts'
           - '1'
         group:
           - 'Core'
           - 'MOI'
+          - 'Perf'
+        exclude:
+          - version: 'lts'
+            group: 'Perf'
     env:
       COOLPDLP_TEST_GROUP: ${{ matrix.group }}
     steps:
@@ -39,6 +44,8 @@ jobs:
       - uses: julia-actions/cache@v3
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-runtest@v1
+        with:
+          coverage: ${{ matrix.group != 'Perf' }}
       - uses: julia-actions/julia-processcoverage@v1
       - uses: codecov/codecov-action@v6
         with:

.gitignore

@@ -39,3 +39,4 @@ playground*.jl
 .DS_Store
 *.trace
 *.gputrace
+*.pb.gz

src/algorithms/common.jl

@@ -134,9 +134,10 @@ abstract type AbstractState{T, V} end
 function prog_showvalues(state::AbstractState)
     err = state.stats.err
     (; primal, primal_scale, dual, dual_scale, gap, gap_scale) = err
-    rel_primal = @sprintf("%.3e", primal / primal_scale)
-    rel_dual = @sprintf("%.3e", dual / dual_scale)
-    rel_gap = @sprintf("%.3e", gap / gap_scale)
+    # @sprintf induces string formatting overhead in hot loops
+    rel_primal = primal / primal_scale
+    rel_dual = dual / dual_scale
+    rel_gap = gap / gap_scale
     return (
         ("primal", rel_primal),
         ("dual", rel_dual),
@@ -192,7 +193,7 @@ function solve(
     milp, sol = preprocess(milp_init_cpu, sol_init_cpu, algo)
     state = initialize(milp, sol, algo; starting_time)
     if nbcons(milp) == 0 && all(iszero, milp.c) # early exit for 0 obj/no cons
-        @. sol.x = proj_box(zero(eltype(milp.lv)), milp.lv, milp.uv)
+        @. sol.x = clamp(zero(eltype(milp.lv)), milp.lv, milp.uv)
         state.stats.termination_status = OPTIMAL
         return get_solution(state, milp), state.stats
     end

src/algorithms/pdhg.jl

@@ -87,14 +87,14 @@ function step!(
 
     τ, σ = η / ω, η * ω
 
-    # xp = proj_box.(x - τ * (c - At * y), lv, uv)
+    # xp = clamp.(x - τ * (c - At * y), lv, uv)
     At_y = mul!(scratch.x, At, y)
-    @. sol.x = proj_box(x - τ * (c - At_y), lv, uv)
+    @. sol.x = clamp(x - τ * (c - At_y), lv, uv)
     xdiff = @. scratch.x = 2sol.x - x
 
-    # yp = y - σ * A * (2xp - x) - σ * proj_box.(inv(σ) * y - A * (2xp - x), -uc, -lc)
+    # yp = y - σ * A * (2xp - x) - σ * clamp.(inv(σ) * y - A * (2xp - x), -uc, -lc)
     A_xdiff = mul!(scratch.y, A, xdiff)
-    @. sol.y = y - σ * A_xdiff - σ * proj_box(inv(σ) * y - A_xdiff, -uc, -lc)
+    @. sol.y = y - σ * A_xdiff - σ * clamp(inv(σ) * y - A_xdiff, -uc, -lc)
 
     # other updates
     state.stats.kkt_passes += 1

src/algorithms/pdlp.jl

@@ -102,14 +102,14 @@ function step!(
 
     τ, σ = η / ω, η * ω
 
-    # xp = proj_box.(x - τ * (c - At * y), lv, uv)
+    # xp = clamp.(x - τ * (c - At * y), lv, uv)
     At_y = mul!(scratch.x, At, y)
-    @. sol.x = proj_box(x - τ * (c - At_y), lv, uv)
+    @. sol.x = clamp(x - τ * (c - At_y), lv, uv)
     xdiff = @. scratch.x = 2sol.x - x
 
-    # yp = y - σ * A * (2xp - x) - σ * proj_box.(inv(σ) * y - A * (2xp - x), -uc, -lc)
+    # yp = y - σ * A * (2xp - x) - σ * clamp.(inv(σ) * y - A * (2xp - x), -uc, -lc)
     A_xdiff = mul!(scratch.y, A, xdiff)
-    @. sol.y = y - σ * A_xdiff - σ * proj_box(inv(σ) * y - A_xdiff, -uc, -lc)
+    @. sol.y = y - σ * A_xdiff - σ * clamp(inv(σ) * y - A_xdiff, -uc, -lc)
 
     # other updates
     state.stats.kkt_passes += 1

src/components/errors.jl

@@ -74,7 +74,7 @@ function kkt_errors!(
     At_y = mul!(scratch.x, At, y)
     r = @. scratch.r = proj_multiplier(c - At_y, lv, uv)
 
-    primal_diff = @. scratch.y = inv(D1.diag) * (A_x - proj_box(A_x, lc, uc))
+    primal_diff = @. scratch.y = inv(D1.diag) * (A_x - clamp(A_x, lc, uc))
     primal = norm(primal_diff)
     rescaled_combined_bounds = @. scratch.y = inv(D1.diag) * combine(lc, uc)
     primal_scale = one(T) + norm(rescaled_combined_bounds)

src/problems/solution.jl

@@ -92,3 +92,7 @@ end
 function Base.isapprox(sol1::PrimalDualSolution{T, V}, sol2::PrimalDualSolution{T, V}; kwargs...) where {T, V}
     return isapprox(sol1.x, sol2.x; kwargs...) && isapprox(sol1.y, sol2.y; kwargs...)
 end
+
+function PrimalDualSolution(milp::MILP)
+    return PrimalDualSolution(zero(milp.lv), zero(milp.lc))
+end

src/utils/linalg.jl

@@ -25,8 +25,6 @@ end
 
 @inline safeprod_left(left, right) = ifelse(isinf(left), right, left * right)
 
-@inline proj_box(x::Number, l::Number, u::Number) = min(u, max(l, x))
-
 """
     proj_multiplier(λ, l, u)
 

src/utils/test.jl

@@ -44,7 +44,7 @@ function random_milp_and_sol(m::Int, n::Int, p::Float64)
         end
     end
     int_var = rand(Bool, length(c))
-    x = proj_box.(randn(n), lv, uv)
+    x = clamp.(randn(n), lv, uv)
     y = proj_multiplier.(randn(m), lc, uc)
     return MILP(; c, lv, uv, A, lc, uc, int_var), PrimalDualSolution(x, y)
 end

test/Project.toml

@@ -1,6 +1,7 @@
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
+Chairmarks = "0ca39b1e-fe0b-4e98-acfc-b1656634c4de"
 CoolPDLP = "9e90bdc5-3073-4e0f-a275-e19f28ac1b53"
 ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7"
 GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
@@ -14,6 +15,7 @@ MathOptBenchmarkInstances = "f7f8d0a1-fd34-491e-a7ac-a4cf52f91fe5"
 MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Preferences = "21216c6a-2e73-6563-6e65-726566657250"
+ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SCS = "c946c3f1-0d1f-5ce8-9dea-7daa1f7e2d13"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"

test/components/errors.jl

@@ -25,7 +25,7 @@ err = CoolPDLP.kkt_errors!(scratch, sol, milp)
 err_p = CoolPDLP.kkt_errors!(scratch, sol_p, milp_p)
 
 @testset "Correct KKT errors" begin
-    @test err.primal ≈ norm(A * x - CoolPDLP.proj_box.(A * x, lc, uc))
+    @test err.primal ≈ norm(A * x - CoolPDLP.clamp.(A * x, lc, uc))
     @test err.dual ≈ norm(c - At * y - r)
     @test err.gap ≈ abs(dot(c, x) + p(-y, lc, uc) + p(-r, lv, uv))
     @test err.primal_scale ≈ 1 + norm(CoolPDLP.combine.(lc, uc))

test/components/preconditioning.jl

@@ -70,6 +70,6 @@ end
 
     @test objective_value(sol.x, milp) ≈ objective_value(sol_p.x, milp_p)
     @test dot(sol.y, milp.A, sol.x) ≈ dot(sol_p.y, milp_p.A, sol_p.x)
-    @test CoolPDLP.proj_box.(sol.x, milp.lv, milp.uv) ≈ prec.D2 * CoolPDLP.proj_box.(sol_p.x, milp_p.lv, milp_p.uv)
-    @test CoolPDLP.proj_box.(milp.A * sol.x, milp.lc, milp.uc) ≈ prec.D1 \ CoolPDLP.proj_box.(milp_p.A * sol_p.x, milp_p.lc, milp_p.uc)
+    @test CoolPDLP.clamp.(sol.x, milp.lv, milp.uv) ≈ prec.D2 * CoolPDLP.clamp.(sol_p.x, milp_p.lv, milp_p.uv)
+    @test CoolPDLP.clamp.(milp.A * sol.x, milp.lc, milp.uc) ≈ prec.D1 \ CoolPDLP.clamp.(milp_p.A * sol_p.x, milp_p.lc, milp_p.uc)
 end

test/perf.jl

@@ -0,0 +1,25 @@
+using Chairmarks
+using CoolPDLP
+using MathOptBenchmarkInstances
+using ProgressMeter
+using SparseArrays
+using Test
+
+prepstate(milp, algo) = initialize(
+    milp, PrimalDualSolution(milp), algo; starting_time = time()
+)
+
+@testset verbose = true "Allocation-free `solve!`" begin
+    milp = MILP(read_instance(Netlib, first(list_instances(Netlib)))[1])
+    @testset "$(typeof(algo))" for algo in [
+            PDHG(time_limit = 1.0, record_error_history = false)
+            PDLP(time_limit = 1.0, record_error_history = false)
+        ]
+        milp = MILP(read_instance(Netlib, first(list_instances(Netlib)))[1])
+        algo = PDHG(time_limit = 1.0, record_error_history = false)
+        solve!(prepstate(milp, algo), milp, algo)
+        result = @b prepstate(milp, algo) solve!(_, milp, algo) seconds = 5
+        result_nosolve = @b ProgressUnknown(; desc = "placeholder")
+        @test result.allocs == result_nosolve.allocs
+    end
+end

test/runtests.jl

@@ -31,7 +31,14 @@ GROUP = get(ENV, "COOLPDLP_TEST_GROUP", nothing)
             include("moi.jl")
         end
     end
-    if GROUP == "CUDA"  # don't test this if GROUP is not specified
+    # don't test this if GROUP is not specified
+    if GROUP == "Perf"
+        # test separately in CI to avoid Codecov noise
+        @testset "Performance" begin
+            include("perf.jl")
+        end
+    end
+    if GROUP == "CUDA"
         Pkg.add("CUDA")
         @testset verbose = true "CUDA" begin
             include("cuda/runtests.jl")

test/utils/linalg.jl

@@ -9,16 +9,6 @@ using Random: Xoshiro
         @test x ≈ CoolPDLP.positive_part(x) - CoolPDLP.negative_part(x)
         @test CoolPDLP.positive_part(x) >= 0
         @test CoolPDLP.negative_part(x) >= 0
-        a = randn()
-        l, u = a - rand(), a + rand()
-        @test l <= CoolPDLP.proj_box(x, l, u) <= u
-        if l <= x <= u
-            @test CoolPDLP.proj_box(x, l, u) == x
-        elseif x >= u
-            @test CoolPDLP.proj_box(x, l, u) == u
-        elseif x <= l
-            @test CoolPDLP.proj_box(x, l, u) == l
-        end
     end
 end