add CuInteriorPointSolver constructor (#185)

MadNLP · Jul 1, 2022 · ad48437 · ad48437
1 parent b7ceb81
commit ad48437
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 9 deletions.
diff --git a/lib/MadNLPGPU/src/MadNLPGPU.jl b/lib/MadNLPGPU/src/MadNLPGPU.jl
@@ -2,7 +2,7 @@ module MadNLPGPU
 
 import LinearAlgebra
 # CUDA
-import CUDA: CUBLAS, CUSOLVER, CuVector, CuMatrix, CuArray, toolkit_version, R_64F, has_cuda, @allowscalar, runtime_version
+import CUDA: CUBLAS, CUSOLVER, CuVector, CuMatrix, CuArray, R_64F, has_cuda, @allowscalar, runtime_version
 # Kernels
 import KernelAbstractions: @kernel, @index, wait, Event
 import CUDAKernels: CUDADevice
@@ -24,5 +24,6 @@ if has_cuda()
     include("lapackgpu.jl")
     export LapackGPUSolver
 end
+include("interface.jl")
 
 end # module
diff --git a/lib/MadNLPGPU/src/interface.jl b/lib/MadNLPGPU/src/interface.jl
@@ -0,0 +1,22 @@
+
+function CuInteriorPointSolver(nlp::AbstractNLPModel;
+    option_dict::Dict{Symbol,Any}=Dict{Symbol,Any}(), kwargs...
+)
+    opt = MadNLP.Options(linear_solver=LapackGPUSolver)
+    MadNLP.set_options!(opt,option_dict,kwargs)
+    MadNLP.check_option_sanity(opt)
+
+    KKTSystem = if (opt.kkt_system == MadNLP.SPARSE_KKT_SYSTEM) || (opt.kkt_system == MadNLP.SPARSE_UNREDUCED_KKT_SYSTEM)
+        error("Sparse KKT system are currently not supported on CUDA GPU.\n" *
+              "Please use `DENSE_KKT_SYSTEM` or `DENSE_CONDENSED_KKT_SYSTEM` instead.")
+    elseif opt.kkt_system == MadNLP.DENSE_KKT_SYSTEM
+        MT = CuMatrix{Float64}
+        VT = CuVector{Float64}
+        MadNLP.DenseKKTSystem{Float64, VT, MT}
+    elseif opt.kkt_system == MadNLP.DENSE_CONDENSED_KKT_SYSTEM
+        MT = CuMatrix{Float64}
+        VT = CuVector{Float64}
+        MadNLP.DenseCondensedKKTSystem{Float64, VT, MT}
+    end
+    return MadNLP.InteriorPointSolver{KKTSystem}(nlp, opt; option_linear_solver=option_dict)
+end
diff --git a/lib/MadNLPGPU/test/densekkt_gpu.jl b/lib/MadNLPGPU/test/densekkt_gpu.jl
@@ -2,26 +2,34 @@
 using CUDA
 using MadNLPTests
 
-function _compare_gpu_with_cpu(kkt_system, n, m, ind_fixed)
+function _compare_gpu_with_cpu(KKTSystem, n, m, ind_fixed)
+
+    opt_kkt = if (KKTSystem == MadNLP.DenseKKTSystem)
+        MadNLP.DENSE_KKT_SYSTEM
+    elseif (KKTSystem == MadNLP.DenseCondensedKKTSystem)
+        MadNLP.DENSE_CONDENSED_KKT_SYSTEM
+    end
+    # Define options
     madnlp_options = Dict{Symbol, Any}(
-        :kkt_system=>MadNLP.DENSE_CONDENSED_KKT_SYSTEM,
+        :kkt_system=>opt_kkt,
         :linear_solver=>LapackGPUSolver,
         :print_level=>MadNLP.ERROR,
     )
 
     nlp = MadNLPTests.DenseDummyQP(; n=n, m=m, fixed_variables=ind_fixed)
 
+    # Solve on CPU
     h_ips = MadNLP.InteriorPointSolver(nlp; option_dict=copy(madnlp_options))
     MadNLP.optimize!(h_ips)
 
-
-    # Init KKT on the GPU
-    TKKTGPU = kkt_system{Float64, CuVector{Float64}, CuMatrix{Float64}}
-    opt = MadNLP.Options(; madnlp_options...)
-    # Instantiate Solver with KKT on the GPU
-    d_ips = MadNLP.InteriorPointSolver{TKKTGPU}(nlp, opt; option_linear_solver=copy(madnlp_options))
+    # Solve on GPU
+    d_ips = MadNLPGPU.CuInteriorPointSolver(nlp; option_dict=copy(madnlp_options))
     MadNLP.optimize!(d_ips)
 
+    T = Float64
+    VT = CuVector{T}
+    MT = CuMatrix{T}
+    @test isa(d_ips.kkt, KKTSystem{T, VT, MT})
     # # Check that both results match exactly
     @test h_ips.cnt.k == d_ips.cnt.k
     @test h_ips.obj_val ≈ d_ips.obj_val atol=1e-10