[Linear Solver] Added undocumented cholesky solver (MadNLP#292)

* undocumented cholesky solver added

* Added testing and made CuCholesky default

lib/MadNLPGPU/Project.toml

@@ -1,25 +1,27 @@
 name = "MadNLPGPU"
 uuid = "d72a61cc-809d-412f-99be-fd81f4b8a598"
-version = "0.6"
+version = "0.6.0"
 
 [deps]
+AMD = "14f7f29c-3bd6-536c-9a0b-7339e30b5a3e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 CUSOLVERRF = "a8cc9031-bad2-4722-94f5-40deabb4245c"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MadNLP = "2621e9c9-9eb4-46b1-8089-e8c72242dfb6"
 
 [compat]
-CUDA = "4, 5"
+AMD = "0.5"
+CUDA = "5"
 CUSOLVERRF = "0.2"
 KernelAbstractions = "0.9"
 MadNLP = "0.7"
 MadNLPTests = "0.3, 0.4"
 julia = "1.7"
 
 [extras]
-MadNLPTests = "b52a2a03-04ab-4a5f-9698-6a2deff93217"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+MadNLPTests = "b52a2a03-04ab-4a5f-9698-6a2deff93217"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]

lib/MadNLPGPU/src/MadNLPGPU.jl

@@ -22,6 +22,9 @@ import MadNLP:
     introduce, factorize!, solve!, improve!, is_inertia, inertia, tril_to_full!,
     LapackOptions, input_type, is_supported, default_options, symul!
 
+# AMD
+import AMD
+
 symul!(y, A, x::CuVector{T}, α = 1., β = 0.) where T = CUBLAS.symv!('L', T(α), A, x, T(β), y)
 MadNLP._ger!(alpha::Number, x::CuVector{T}, y::CuVector{T}, A::CuMatrix{T}) where T = CUBLAS.ger!(alpha, x, y, A)
 function MadNLP._madnlp_unsafe_wrap(vec::VT, n, shift=1) where {T, VT <: CuVector{T}}
@@ -47,7 +50,7 @@ function MadNLP.MadNLPOptions(nlp::AbstractNLPModel{T,VT}) where {T, VT <: CuVec
     kkt_system = is_dense_callback ? MadNLP.DenseCondensedKKTSystem : MadNLP.SparseCondensedKKTSystem
 
     # if dense kkt system, we use a dense linear solver
-    linear_solver = is_dense_callback ? LapackGPUSolver : RFSolver
+    linear_solver = is_dense_callback ? LapackGPUSolver : CuCholeskySolver
 
     equality_treatment = is_dense_callback ? MadNLP.EnforceEquality : MadNLP.RelaxEquality
 

lib/MadNLPGPU/src/cusolverrf.jl

@@ -174,3 +174,128 @@ MadNLP.improve!(M::GLUSolver) = false
 MadNLP.is_supported(::Type{GLUSolver},::Type{Float32}) = true
 MadNLP.is_supported(::Type{GLUSolver},::Type{Float64}) = true
 MadNLP.introduce(M::GLUSolver) = "GLU"
+
+
+
+#=
+    Undocumented Cholesky Solver
+=#
+
+@kwdef mutable struct CuCholeskySolverOptions <: MadNLP.AbstractOptions
+end
+
+mutable struct CuCholeskySolver{T} <: MadNLP.AbstractLinearSolver{T}
+    inner::Union{Nothing,CUSOLVER.SparseCholesky}
+
+    tril::CUSPARSE.CuSparseMatrixCSC{T}
+    full::CUSPARSE.CuSparseMatrixCSR{T}
+    tril_to_full_view::CuSubVector{T}
+    buffer::CUDA.CuVector{T}
+
+    fullp::CUSPARSE.CuSparseMatrixCSR{T}
+    p::CUDA.CuVector{Int}
+    pnzval::CUDA.CuVector{Int}
+    rhs::CUDA.CuVector{T}
+
+    singularity::Bool
+
+    opt::CuCholeskySolverOptions
+    logger::MadNLP.MadNLPLogger
+end
+
+function CuCholeskySolver(
+    csc::CUSPARSE.CuSparseMatrixCSC;
+    opt=CuCholeskySolverOptions(),
+    logger=MadNLP.MadNLPLogger(),
+)
+    n, m = size(csc)
+    @assert n == m
+
+    full,tril_to_full_view = MadNLP.get_tril_to_full(csc)
+    buffer = similar(csc.nzVal,1)
+
+    full = CUSPARSE.CuSparseMatrixCSR(
+        full.colPtr,
+        full.rowVal,
+        full.nzVal,
+        full.dims
+    )
+
+    full_cpu = MadNLP.SparseMatrixCSC(full)
+    full_cpu_order = MadNLP.SparseMatrixCSC(
+        n, m,
+        full_cpu.colptr,
+        full_cpu.rowval,
+        Array(1:length(full_cpu.nzval)),
+    )
+    p = AMD.amd(full_cpu_order)
+    full_cpu_reorder = full_cpu_order[p,p]
+    pnzval = full_cpu_reorder.nzval
+
+    fullp = CUSPARSE.CuSparseMatrixCSR(
+        CuArray(full_cpu_reorder.colptr),
+        CuArray(full_cpu_reorder.rowval),
+        similar(full.nzVal),
+        (n, m),
+    )    
+
+    rhs = similar(csc.nzVal, n)
+
+    return CuCholeskySolver(
+        nothing, csc, full, tril_to_full_view, buffer,
+        fullp, CuArray{Int}(p), CuArray{Int}(pnzval), rhs, false,
+        opt, logger
+    )
+end
+
+function MadNLP.factorize!(M::CuCholeskySolver)
+
+    copyto!(M.full.nzVal, M.tril_to_full_view)
+    _copy_to_perm_2!(CUDABackend())(M.fullp.nzVal, M.pnzval, M.full.nzVal; ndrange= length(M.pnzval))
+    synchronize(CUDABackend())
+    if M.inner == nothing
+        M.inner = CUSOLVER.SparseCholesky(M.fullp)
+        CUSOLVER.spcholesky_buffer(M.inner, M.fullp)
+    end
+    try
+        CUSOLVER.spcholesky_factorise(M.inner, M.fullp, eltype(M.fullp.nzVal) == Float32 ? 1e-6 : 1e-12)
+        M.singularity = false
+    catch e
+        println(e)
+        M.singularity = true
+    end
+
+    return M
+end
+
+function MadNLP.solve!(M::CuCholeskySolver{T}, x) where T
+    _copy_to_perm_2!(CUDABackend())(M.rhs, M.p, x; ndrange=length(M.p))
+    synchronize(CUDABackend())
+    CUSOLVER.spcholesky_solve(M.inner, M.rhs, x)
+    _copy_to_perm!(CUDABackend())(M.rhs, M.p, x; ndrange=length(M.p))
+    synchronize(CUDABackend())
+    copyto!(x, M.rhs)
+    return x
+end
+
+@kernel function _copy_to_perm!(y,p,x)
+    i = @index(Global)
+    @inbounds y[p[i]] = x[i]
+end
+@kernel function _copy_to_perm_2!(y,p,x)
+    i = @index(Global)
+    @inbounds y[i] = x[p[i]]
+end
+function MadNLP.inertia(M::CuCholeskySolver{T}) where T
+    return !(M.singularity) ? (size(M.fullp,1),0,0) : (size(M.fullp,1) - 2,1,1)
+end
+
+MadNLP.input_type(::Type{CuCholeskySolver}) = :csc
+MadNLP.default_options(::Type{CuCholeskySolver}) = CuCholeskySolverOptions()
+MadNLP.is_inertia(M::CuCholeskySolver) = true
+MadNLP.improve!(M::CuCholeskySolver) = false
+MadNLP.is_supported(::Type{CuCholeskySolver},::Type{Float32}) = true
+MadNLP.is_supported(::Type{CuCholeskySolver},::Type{Float64}) = true
+MadNLP.introduce(M::CuCholeskySolver) = "cuSolverCholesky"
+
+export CuCholeskySolver

lib/MadNLPGPU/test/madnlpgpu_test.jl

@@ -8,6 +8,14 @@ testset = [
         ),
         [],
     ],
+    [
+        "LapackGPU-CUSOLVERRF",
+        ()->MadNLP.Optimizer(
+            linear_solver=MadNLPGPU.CuCholeskySolver,
+            print_level=MadNLP.ERROR
+        ),
+        [],
+    ],
     [
         "LapackGPU-CUSOLVERRF",
         ()->MadNLP.Optimizer(