FisSal2011: Implement faster get/set basis for Gurobi

Project.toml

@@ -6,6 +6,7 @@ version = "0.4.2"
 [deps]
 Conda = "8f4d0f93-b110-5947-807f-2305c1781a2d"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
+Gurobi = "2e9cd046-0924-5485-92f1-d5272153d98b"
 HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
 HiGHS = "87dc4568-4c63-4d18-b0c0-bb2238e4078b"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
@@ -28,6 +29,7 @@ TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 [compat]
 Conda = "1"
 DataStructures = "0.18"
+Gurobi = "1.7.5"
 HDF5 = "0.16"
 HiGHS = "1"
 JLD2 = "0.4"
@@ -36,10 +38,10 @@ JuMP = "1"
 KLU = "0.4"
 MathOptInterface = "1"
 OrderedCollections = "1"
+PrecompileTools = "1"
 PyCall = "1"
 Requires = "1"
+SCIP = "0.12"
 Statistics = "1"
 TimerOutputs = "0.5"
 julia = "1"
-PrecompileTools = "1"
-SCIP = "0.12"

src/Cuts/tableau/gmi_dual.jl

@@ -332,6 +332,7 @@ function collect_gmi_FisSal2011(
         pool_cut_hashes = Set{UInt64}()
         pool_size_mb = 0
         tableau_density::Float32 = 0.05
+        basis_cache = nothing
         λ, Δ = 0, 0
         μ = 10
     end
@@ -417,8 +418,10 @@ function collect_gmi_FisSal2011(
         end
 
         @timeit "Optimize LP (lagrangian)" begin
+            basis_cache === nothing || set_basis(model_s, basis_cache)
             set_silent(model_s)
             optimize!(model_s)
+            basis_cache = get_basis(model_s)
             status = termination_status(model_s)
             if status != MOI.OPTIMAL
                 error("Non-optimal termination status: $status")
@@ -473,8 +476,10 @@ function collect_gmi_FisSal2011(
         end
 
         if mod(round - 1, interval_read_tableau) == 0
-            @timeit "Select tableau rows" begin
+            @timeit "Get basis" begin
                 basis = get_basis(model_s)
+            end
+            @timeit "Select tableau rows" begin
                 selected_rows =
                     select_gmi_rows(data_s, basis, sol_frac, max_rows = max_cuts_per_round)
             end

src/Cuts/tableau/tableau.jl

@@ -4,48 +4,161 @@
 
 using KLU
 using TimerOutputs
+using Gurobi
 
 function get_basis(model::JuMP.Model)::Basis
-    var_basic = Int[]
-    var_nonbasic = Int[]
-    constr_basic = Int[]
-    constr_nonbasic = Int[]
-
-    # Variables
-    for (i, var) in enumerate(all_variables(model))
-        bstatus = MOI.get(model, MOI.VariableBasisStatus(), var)
-        if bstatus == MOI.BASIC
-            push!(var_basic, i)
-        elseif bstatus == MOI.NONBASIC_AT_LOWER
-            push!(var_nonbasic, i)
-        else
-            error("Unknown basis status: $bstatus")
+    if isa(unsafe_backend(model), Gurobi.Optimizer)
+        return get_basis_gurobi(model)
+    end
+
+    @timeit "Initialization" begin
+        var_basic = Int[]
+        var_nonbasic = Int[]
+        constr_basic = Int[]
+        constr_nonbasic = Int[]
+        nvars = num_variables(model)
+        sizehint!(var_basic, nvars)
+        sizehint!(var_nonbasic, nvars)
+    end
+
+    @timeit "Query variables" begin
+        for (i, var) in enumerate(all_variables(model))
+            bstatus = MOI.get(model, MOI.VariableBasisStatus(), var)
+            if bstatus == MOI.BASIC
+                push!(var_basic, i)
+            elseif bstatus == MOI.NONBASIC_AT_LOWER
+                push!(var_nonbasic, i)
+            else
+                error("Unknown basis status: $bstatus")
+            end
         end
     end
 
-    # Constraints
-    constr_index = 1
-    for (ftype, stype) in list_of_constraint_types(model)
-        for constr in all_constraints(model, ftype, stype)
-            if ftype == VariableRef
-                # nop
-            elseif ftype == AffExpr
-                bstatus = MOI.get(model, MOI.ConstraintBasisStatus(), constr)
-                if bstatus == MOI.BASIC
-                    push!(constr_basic, constr_index)
-                elseif bstatus == MOI.NONBASIC
-                    push!(constr_nonbasic, constr_index)
+    @timeit "Query constraints" begin
+        constr_index = 1
+        for (ftype, stype) in list_of_constraint_types(model)
+            for constr in all_constraints(model, ftype, stype)
+                if ftype == VariableRef
+                    # nop
+                elseif ftype == AffExpr
+                    bstatus = MOI.get(model, MOI.ConstraintBasisStatus(), constr)
+                    if bstatus == MOI.BASIC
+                        push!(constr_basic, constr_index)
+                    elseif bstatus == MOI.NONBASIC
+                        push!(constr_nonbasic, constr_index)
+                    else
+                        error("Unknown basis status: $bstatus")
+                    end
+                    constr_index += 1
                 else
-                    error("Unknown basis status: $bstatus")
+                    error("Unsupported constraint type: ($ftype, $stype)")
                 end
-                constr_index += 1
+            end
+        end
+    end
+
+    @timeit "Build basis struct" begin
+        basis = Basis(; var_basic, var_nonbasic, constr_basic, constr_nonbasic)
+    end
+
+    return basis
+end
+
+function set_basis(model::JuMP.Model, basis::Basis)
+    if isa(unsafe_backend(model), Gurobi.Optimizer)
+        # NOP
+        return
+    end
+
+    @timeit "Initialization" begin
+        nvars = num_variables(model)
+        gurobi_model = unsafe_backend(model).inner
+    end
+
+    @timeit "Set variable basis" begin
+        var_basis_statuses = Vector{Cint}(undef, nvars)
+        fill!(var_basis_statuses, -1)  # Default to GRB_NONBASIC_LOWER
+        
+        for var_idx in basis.var_basic
+            var_basis_statuses[var_idx] = 0  # GRB_BASIC
+        end
+        
+        ret = GRBsetintattrarray(gurobi_model, "VBasis", 0, nvars, var_basis_statuses)
+        if ret != 0
+            error("Failed to set variable basis statuses in Gurobi: error code $ret")
+        end
+    end
+
+    @timeit "Set constraint basis" begin
+        nconstr = num_constraints(model, AffExpr, MOI.EqualTo{Float64})
+        constr_basis_statuses = Vector{Cint}(undef, nconstr)
+        fill!(constr_basis_statuses, -1)  # Default to GRB_NONBASIC
+        
+        for constr_idx in basis.constr_basic
+            constr_basis_statuses[constr_idx] = 0  # GRB_BASIC
+        end
+        
+        ret = GRBsetintattrarray(gurobi_model, "CBasis", 0, nconstr, constr_basis_statuses)
+        if ret != 0
+            error("Failed to set constraint basis statuses in Gurobi: error code $ret")
+        end
+    end
+
+    return nothing
+end
+
+function get_basis_gurobi(model::JuMP.Model)::Basis
+    @timeit "Initialization" begin
+        var_basic = Int[]
+        var_nonbasic = Int[]
+        constr_basic = Int[]
+        constr_nonbasic = Int[]
+        nvars = num_variables(model)
+        sizehint!(var_basic, nvars)
+        sizehint!(var_nonbasic, nvars)
+        gurobi_model = unsafe_backend(model).inner
+    end
+
+    @timeit "Query variables" begin
+        var_basis_statuses = Vector{Cint}(undef, nvars)
+        ret = GRBgetintattrarray(gurobi_model, "VBasis", 0, nvars, var_basis_statuses)
+        if ret != 0
+            error("Failed to get variable basis statuses from Gurobi: error code $ret")
+        end
+        for i in 1:nvars
+            if var_basis_statuses[i] == 0  # GRB_BASIC
+                push!(var_basic, i)
+            elseif var_basis_statuses[i] == -1  # GRB_NONBASIC_LOWER
+                push!(var_nonbasic, i)
+            else
+                error("Unknown variable basis status: $(var_basis_statuses[i])")
+            end
+        end
+    end
+
+    @timeit "Query constraints" begin
+        nconstr = num_constraints(model, AffExpr, MOI.EqualTo{Float64})
+        constr_basis_statuses = Vector{Cint}(undef, nconstr)
+        ret = GRBgetintattrarray(gurobi_model, "CBasis", 0, nconstr, constr_basis_statuses)
+        if ret != 0
+            error("Failed to get constraint basis statuses from Gurobi: error code $ret")
+        end
+        for i in 1:nconstr
+            if constr_basis_statuses[i] == 0  # GRB_BASIC
+                push!(constr_basic, i)
+            elseif constr_basis_statuses[i] == -1  # GRB_NONBASIC
+                push!(constr_nonbasic, i)
             else
-                error("Unsupported constraint type: ($ftype, $stype)")
+                error("Unknown constraint basis status: $(constr_basis_statuses[i])")
             end
         end
     end
 
-    return Basis(; var_basic, var_nonbasic, constr_basic, constr_nonbasic)
+    @timeit "Build basis struct" begin
+        basis = Basis(; var_basic, var_nonbasic, constr_basic, constr_nonbasic)
+    end
+
+    return basis
 end
 
 function get_x(model::JuMP.Model)
@@ -150,4 +263,4 @@ function compute_tableau(
     return Tableau(obj = tableau_obj, lhs = tableau_lhs, rhs = tableau_rhs, z = z)
 end
 
-export get_basis, get_x, compute_tableau
+export get_basis, get_basis_gurobi, set_basis, get_x, compute_tableau