feat: Reactant compatibility

.github/workflows/Test.yml

@@ -32,6 +32,7 @@ jobs:
           - 'MOI'
           - 'Perf'
           - 'OpenCL'
+          - 'Reactant'
         exclude:
           - version: 'lts'
             group: 'Perf'

Project.toml

@@ -20,17 +20,20 @@ Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
 QPSReader = "10f199a5-22af-520b-b891-7ce84a7b1bd0"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+ReactantCore = "a3311ec8-5e00-46d5-b541-4f83e724a433"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 
 [weakdeps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
-cuSPARSE = "b26da814-b3bc-49ef-b0ee-c816305aa060"
 GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
+Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
+cuSPARSE = "b26da814-b3bc-49ef-b0ee-c816305aa060"
 
 [extensions]
 CoolPDLPCUDAExt = ["CUDA", "cuSPARSE"]
 CoolPDLPGPUArraysExt = "GPUArrays"
+CoolPDLPReactantExt = "Reactant"
 
 [compat]
 Adapt = "4.4.0"
@@ -48,6 +51,8 @@ Printf = "1"
 ProgressMeter = "1.11.0"
 QPSReader = "0.2.1"
 Random = "1"
+Reactant = "0.2.279"
+ReactantCore = "0.1.21"
 SparseArrays = "1"
 StableRNGs = "1.0.3"
 cuSPARSE = "6.0.0"

ext/CoolPDLPReactantExt.jl

@@ -0,0 +1,63 @@
+module CoolPDLPReactantExt
+
+using CoolPDLP:
+    CoolPDLP,
+    ConvergenceStats,
+    KKTErrors,
+    MILP,
+    PDHGState,
+    PrimalDualSolution,
+    Scratch,
+    StepSizes,
+    custom_to_rarray
+using Reactant: ConcreteRArray, to_rarray
+
+function CoolPDLP.custom_to_rarray(milp::MILP; kwargs...)
+    return to_rarray(milp; kwargs...)
+end
+
+function CoolPDLP.custom_to_rarray(sol::PrimalDualSolution; kwargs...)
+    return to_rarray(sol; kwargs...)
+end
+
+function CoolPDLP.custom_to_rarray(scratch::Scratch; kwargs...)
+    (; x, y, z, b1, b2) = scratch
+    xr = to_rarray(x; kwargs...)
+    yr = to_rarray(y; kwargs...)
+    zr = to_rarray(z; kwargs...)
+    b1r = to_rarray(b1; kwargs...)
+    b2r = to_rarray(b2; kwargs...)
+    return Scratch(;
+        x = xr, y = yr, z = zr, b1 = b1r, b2 = b2r
+    )
+end
+
+function CoolPDLP.custom_to_rarray(stats::ConvergenceStats; kwargs...)
+    (; err, starting_time, time_elapsed, kkt_passes, termination_status, error_history) = stats
+    return ConvergenceStats(
+        to_rarray(err; kwargs...);
+        starting_time,
+        time_elapsed,
+        kkt_passes,
+        termination_status,
+        error_history = to_rarray(error_history; kwargs...)
+    )
+end
+
+function CoolPDLP.custom_to_rarray(state::PDHGState; kwargs...)
+    (; sol, sol_last, step_sizes, scratch, stats) = state
+    sol_r = to_rarray(sol; kwargs...)
+    sol_last_r = to_rarray(sol_last; kwargs...)
+    step_sizes_r = to_rarray(step_sizes; kwargs...)
+    scratch_r = custom_to_rarray(scratch; kwargs...)
+    stats_r = custom_to_rarray(stats; kwargs...)
+    return PDHGState(;
+        sol = sol_r,
+        sol_last = sol_last_r,
+        step_sizes = step_sizes_r,
+        scratch = scratch_r,
+        stats = stats_r,
+    )
+end
+
+end

src/CoolPDLP.jl

@@ -4,12 +4,13 @@ module CoolPDLP
 using Adapt: Adapt, adapt
 using Atomix: Atomix
 using BangBang: add!!, broadcast!!
-using DispatchDoctor: @stable
+using DispatchDoctor: @stable, @unstable
 using DocStringExtensions: TYPEDFIELDS
 using IterativeSolvers: powm!
 using KernelAbstractions: KernelAbstractions, Backend, CPU, @kernel, @index, allocate, get_backend
 import MathOptInterface as MOI
 using ProgressMeter: ProgressUnknown, finish!, next!
+using ReactantCore: @trace
 using QPSReader: QPSData, VTYPE_Binary, VTYPE_Integer
 using StableRNGs: StableRNG
 
@@ -21,34 +22,36 @@ using SparseArrays: SparseArrays, SparseMatrixCSC, AbstractSparseMatrix, findnz,
 
 include("public.jl")
 
-@stable begin
-    include("utils/device.jl")
-    include("utils/mat_coo.jl")
-    include("utils/mat_csr.jl")
-    include("utils/mat_ell.jl")
-    include("utils/linalg.jl")
-    include("utils/test.jl")
-    include("utils/batching.jl")
+# @stable begin
+include("utils/device.jl")
+include("utils/mat_coo.jl")
+include("utils/mat_csr.jl")
+include("utils/mat_ell.jl")
+include("utils/linalg.jl")
+include("utils/test.jl")
+include("utils/batching.jl")
 
-    include("problems/milp.jl")
-    include("problems/solution.jl")
-    include("problems/modify.jl")
+include("problems/milp.jl")
+include("problems/solution.jl")
+include("problems/modify.jl")
 
-    include("components/scratch.jl")
-    include("components/conversion.jl")
-    include("components/preconditioning.jl")
-    include("components/permutation.jl")
-    include("components/step_size.jl")
-    include("components/errors.jl")
-    include("components/iteration.jl")
-    include("components/restart.jl")
-    include("components/generic.jl")
-    include("components/termination.jl")
+include("components/scratch.jl")
+include("components/conversion.jl")
+include("components/preconditioning.jl")
+include("components/permutation.jl")
+include("components/step_size.jl")
+include("components/errors.jl")
+include("components/iteration.jl")
+include("components/restart.jl")
+include("components/generic.jl")
+include("components/termination.jl")
 
-    include("algorithms/common.jl")
-    include("algorithms/pdhg.jl")
-    include("algorithms/pdlp.jl")
-end
+include("algorithms/common.jl")
+include("algorithms/pdhg.jl")
+include("algorithms/pdlp.jl")
+
+include("extensions.jl")
+# end
 
 include("MOI_wrapper.jl")
 

src/algorithms/pdhg.jl

@@ -23,7 +23,7 @@ end
 $(TYPEDFIELDS)
 """
 @kwdef mutable struct PDHGState{
-        T <: Number, V <: AbstractVecOrMat{T}, S <: BatchedNumber{T},
+        T <: Number, V <: AbstractVecOrMat{T}, S <: BatchedNumber,
         Sc <: Scratch{T, V, S},
     } <: AbstractState{T, V}
     "current solution"
@@ -111,5 +111,5 @@ function step!(
 
     # other updates
     state.stats.kkt_passes += 1
-    return nothing
+    return state, milp
 end

src/algorithms/pdlp.jl

@@ -18,8 +18,8 @@ end
 $(TYPEDFIELDS)
 """
 @kwdef mutable struct PDLPState{
-        T <: Number, V <: AbstractVecOrMat{T}, S <: BatchedNumber{T},
-        B <: BatchedNumber{Bool}, Sc <: Scratch{T, V, S},
+        T <: Number, V <: AbstractVecOrMat{T}, S <: BatchedNumber,
+        B <: BatchedNumber, Sc <: Scratch{T, V, S},
     } <: AbstractState{T, V}
     "current solution"
     sol::PrimalDualSolution{T, V}

src/components/errors.jl

@@ -7,19 +7,19 @@ Mutable so that [`kkt_errors!`](@ref) can refill it without allocating.
 
 $(TYPEDFIELDS)
 """
-@kwdef mutable struct KKTErrors{T <: BatchedNumber}
+@kwdef mutable struct KKTErrors{T <: Number, B <: BatchedNumber{T}}
     "primal feasibility error"
-    primal::T
+    primal::B
     "characteristic scale of the primal constraint RHS"
-    primal_scale::T
+    primal_scale::B
     "dual feasibility error"
-    dual::T
+    dual::B
     "characteristic scale of the dual constraint RHS"
-    dual_scale::T
+    dual_scale::B
     "primal-dual gap"
-    gap::T
+    gap::B
     "characteristic scale of the gap"
-    gap_scale::T
+    gap_scale::B
 end
 
 format_error(e::Number) = @sprintf("%.3e", e)
@@ -77,7 +77,7 @@ Base.copy(err::KKTErrors) = KKTErrors(
 Fill `dest`, column by column, with the errors of `err_true` where `cond` holds and those of `err_false` elsewhere.
 """
 function select_errors!!(
-        dest::KKTErrors, cond::BatchedNumber{Bool},
+        dest::KKTErrors, cond::BatchedNumber,
         err_true::KKTErrors, err_false::KKTErrors,
     )
     dest.primal = broadcast!!(ifelse, dest.primal, cond, err_true.primal, err_false.primal)

src/components/restart.jl

@@ -28,7 +28,7 @@ end
 
 $(TYPEDFIELDS)
 """
-mutable struct RestartStats{T <: BatchedNumber, B <: BatchedNumber{Bool}}
+mutable struct RestartStats{T <: BatchedNumber, B <: BatchedNumber}
     "whether to restart from the average solution, column by column"
     restart_from_avg::B
     "KKT errors of the current solution"

src/components/scratch.jl

@@ -1,4 +1,4 @@
-@kwdef struct Scratch{T <: Number, V <: AbstractVecOrMat{T}, S <: BatchedNumber{T}}
+@kwdef struct Scratch{T <: Number, V <: AbstractVecOrMat{T}, S <: BatchedNumber}
     "primal scratch (length `nvar`)"
     x::V
     "dual scratch (length `ncons`)"

src/components/termination.jl

@@ -106,13 +106,25 @@ function termination_status!!(
     )
     (; err, time_elapsed, kkt_passes) = stats
     (; termination_reltol, time_limit, max_kkt_passes) = params
-    if batched_all(<=(termination_reltol), relative!!(dest, err))
-        return OPTIMAL
+    st = if batched_all(<=(termination_reltol), relative!!(dest, err))
+        OPTIMAL
     elseif time_elapsed >= time_limit
-        return TIME_LIMIT
+        TIME_LIMIT
     elseif kkt_passes >= max_kkt_passes
-        return ITERATION_LIMIT
+        ITERATION_LIMIT
     else
-        return STILL_RUNNING
+        STILL_RUNNING
     end
+    return st
+end
+
+function should_terminate!!(
+        dest::BatchedNumber, stats::ConvergenceStats, params::TerminationParameters
+    )
+    (; err, time_elapsed, kkt_passes) = stats
+    (; termination_reltol, time_limit, max_kkt_passes) = params
+    is_optimal = batched_all(<=(termination_reltol), relative!!(dest, err))
+    is_time_limit = time_elapsed >= time_limit
+    is_iteration_limit = kkt_passes >= max_kkt_passes
+    return is_optimal || is_time_limit || is_iteration_limit
 end

src/extensions.jl

@@ -0,0 +1,10 @@
+"""
+    custom_to_rarray
+
+Alternative to `Reactant.to_rarray` designed to trace structs inside the CoolPDLP package.
+
+# See also
+
+- https://github.com/EnzymeAD/Reactant.jl/issues/3184
+"""
+function custom_to_rarray end

src/problems/milp.jl

@@ -34,7 +34,7 @@ struct MILP{
         Dg2 <: Diagonal{T},
         M <: AbstractMatrix{T},
         Mt <: AbstractMatrix{T},
-        Vb <: DenseVector{Bool},
+        Vb <: DenseVector,
     }
     "objective vector"
     c::Vo
@@ -197,14 +197,14 @@ nbvar(milp::MILP) = size(milp.c, 1)
 
 Return the number of integer variables in `milp`.
 """
-nbvar_int(milp::MILP) = sum(milp.int_var)
+nbvar_int(milp::MILP) = sum(identity, milp.int_var)
 
 """
     nbvar_cont(milp)
 
 Return the number of continuous variables in `milp`.
 """
-nbvar_cont(milp::MILP) = nbvar(milp) - nbvar_int(milp)
+nbvar_cont(milp::MILP) = sum(!, milp.int_var)
 
 """
     nbcons(milp)

src/problems/solution.jl

@@ -115,7 +115,7 @@ end
 Overwrite the columns of `sol` for which `cond` holds with those of `sol_other`.
 """
 function batched_select!(
-        sol::PrimalDualSolution, cond::BatchedNumber{Bool},
+        sol::PrimalDualSolution, cond::BatchedNumber,
         sol_other::PrimalDualSolution,
     )
     condr = transpose(cond)

src/utils/linalg.jl

@@ -122,8 +122,8 @@ end
 Return the largest finite absolute value between the two bounds, or zero if neither is finite.
 """
 function combine(l::Number, u::Number)
-    ls = isfinite(l) ? abs(l) : zero(l)
-    us = isfinite(u) ? abs(u) : zero(u)
+    ls = ifelse(isfinite(l), abs(l), zero(l))
+    us = ifelse(isfinite(u), abs(u), zero(u))
     return max(zero(l), ls, us)
 end
 

test/Project.toml

@@ -18,6 +18,7 @@ Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Preferences = "21216c6a-2e73-6563-6e65-726566657250"
 ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
 SCS = "c946c3f1-0d1f-5ce8-9dea-7daa1f7e2d13"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"

test/gpu/reactant/runtests.jl

@@ -0,0 +1,31 @@
+using CoolPDLP
+using MathOptBenchmarkInstances
+using Reactant
+using Test
+
+dataset = Netlib
+list = list_instances(dataset);
+name = list[4]
+qps, path = read_instance(dataset, name);
+
+milp0 = MILP(qps; dataset, name, path);
+sol0 = PrimalDualSolution(milp0);
+
+algo = PDHG(
+    Float32,
+    Int32,
+    Matrix;
+    termination_reltol = 1.0f-6,
+    time_limit = 10.0,
+    record_error_history = false,
+    show_progress = false
+);
+
+milp, sol = preprocess(milp0, sol0, algo);
+state = initialize(milp, sol, algo; starting_time = time());
+
+milp_r = CoolPDLP.custom_to_rarray(milp);
+state_r = CoolPDLP.custom_to_rarray(state; track_numbers = true);
+
+compile_options = CompileOptions(; donated_args = :none)
+@test_nowarn compiled_step! = @compile compile_options = compile_options CoolPDLP.step!(state_r, milp_r)

test/runtests.jl

@@ -76,4 +76,9 @@ include("fixtures.jl")
             include("gpu/opencl/runtests.jl")
         end
     end
+    if GROUP == "Reactant"
+        @testset verbose = true "Reactant" begin
+            include("gpu/reactant/runtests.jl")
+        end
+    end
 end