[MadNLPGPU] Upgrade CUDSS -- support iterative refinement and hybrid …

…mode (#329)
MadNLP · Jul 15, 2024 · dbfc0fc · dbfc0fc
1 parent faa17a8
commit dbfc0fc
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 1 deletion.
diff --git a/lib/MadNLPGPU/Project.toml b/lib/MadNLPGPU/Project.toml
@@ -17,7 +17,7 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 [compat]
 AMD = "0.5"
 CUDA = "5"
-CUDSS = "0.1"
+CUDSS = "0.3.1"
 CUSOLVERRF = "0.2"
 KernelAbstractions = "0.9"
 MadNLP = "0.7, 0.8"

diff --git a/lib/MadNLPGPU/src/LinearSolvers/cudss.jl b/lib/MadNLPGPU/src/LinearSolvers/cudss.jl
@@ -5,6 +5,8 @@ import CUDSS
     cudss_algorithm::MadNLP.LinearFactorization = MadNLP.LDL
     ordering::ORDERING = DEFAULT_ORDERING
     perm::Vector{Cint} = Cint[]
+    ir::Int = 0
+    hybrid::Bool = false
 end
 
 mutable struct CUDSSSolver{T} <: MadNLP.AbstractLinearSolver{T}
@@ -61,6 +63,8 @@ function CUDSSSolver(
         end
         CUDSS.cudss_set(solver, "user_perm", opt.perm)
     end
+    (opt.ir > 0) && CUDSS.cudss_set(solver, "ir_n_steps", opt.ir)
+    opt.hybrid && CUDSS.cudss_set(solver, "hybrid_mode", 1)
 
     x_gpu = CUDA.zeros(T, n)
     b_gpu = CUDA.zeros(T, n)

diff --git a/lib/MadNLPGPU/test/madnlpgpu_test.jl b/lib/MadNLPGPU/test/madnlpgpu_test.jl
@@ -26,6 +26,16 @@ testset = [
         ),
         [],
     ],
+    [
+        "CUDSS-HYBRID",
+        ()->MadNLP.Optimizer(
+            linear_solver=MadNLPGPU.CUDSSSolver,
+            print_level=MadNLP.ERROR,
+            hybrid=true,
+            ir=1,
+        ),
+        [],
+    ],
     [
         "CUSOLVERRF",
         ()->MadNLP.Optimizer(