Add workspace to Factor to keep 3-arg API. Redesign CholmodWorkspace similar to UmfpackWS inside UmfpackLU

Author: lorenzogambichler

src/solvers/cholmod.jl

@@ -25,7 +25,6 @@ import LinearAlgebra: (\), AdjointFactorization,
 using SparseArrays
 using SparseArrays: getcolptr, AbstractSparseVecOrMat
 export
-    CholmodWorkspace,
     Dense,
     Factor,
     Sparse
@@ -259,8 +258,50 @@ function Sparse(p::Ptr{cholmod_sparse})
     Sparse{jlxtype(s.xtype, s.dtype)}(p)
 end
 
+"""
+    CholmodWorkspace()
+
+Reusable workspace for allocation-free CHOLMOD solves.
+
+`cholmod_solve2` allocates two temporary dense matrices (Y and E) on every solve.
+A `CholmodWorkspace` holds those buffers across calls so that CHOLMOD reuses them
+instead of reallocating. It also pre-allocates the `cholmod_dense_struct` wrappers
+for `x` and `b`, eliminating all Julia-side heap allocations per solve.
+
+Each [`Factor`](@ref) embeds its own workspace and uses it automatically in
+`ldiv!(x, F, b)`. For multi-thread solves with a single factor, use `copy(F)`
+for every thread so each copy gets its own independent workspace and lock.
+
+The workspace is finalizer-protected. The Y and E buffers are freed automatically
+when the workspace is garbage-collected.
+"""
+mutable struct CholmodWorkspace
+    dense_x::cholmod_dense_struct
+    dense_b::cholmod_dense_struct
+    X::Ref{Ptr{cholmod_dense_struct}}
+    Y::Ref{Ptr{cholmod_dense_struct}}
+    E::Ref{Ptr{cholmod_dense_struct}}
+
+    function CholmodWorkspace()
+        ws = new(
+            cholmod_dense_struct(), # all fields will be written in ldiv! before use
+            cholmod_dense_struct(),
+            Ref(Ptr{cholmod_dense_struct}(C_NULL)),
+            Ref(Ptr{cholmod_dense_struct}(C_NULL)),
+            Ref(Ptr{cholmod_dense_struct}(C_NULL))
+        )
+        finalizer(ws) do w
+            w.Y[] != C_NULL && free!(w.Y[])
+            w.E[] != C_NULL && free!(w.E[])
+        end
+        ws
+    end
+end
+
 mutable struct Factor{Tv<:VTypes, Ti<:ITypes} <: Factorization{Tv}
     ptr::Ptr{cholmod_factor}
+    workspace::CholmodWorkspace
+    lock::ReentrantLock
     function Factor{Tv, Ti}(ptr::Ptr{cholmod_factor}, register_finalizer = true) where {Tv, Ti}
         if ptr == C_NULL
             throw(ArgumentError("factorization construction failed for " *
@@ -277,7 +318,7 @@ mutable struct Factor{Tv<:VTypes, Ti<:ITypes} <: Factorization{Tv}
             free!(ptr, Ti)
             throw(CHOLMODException("dtype=$(dtyp(Tv)) not supported"))
         end
-        F = new(ptr)
+        F = new(ptr, CholmodWorkspace(), ReentrantLock())
         if register_finalizer
             finalizer(free!, F)
         end
@@ -312,6 +353,8 @@ Base.pointer(x::Dense{Tv}) where {Tv}  = Base.unsafe_convert(Ptr{cholmod_dense},
 Base.pointer(x::Sparse{Tv}) where {Tv} = Base.unsafe_convert(Ptr{cholmod_sparse}, x)
 Base.pointer(x::Factor{Tv}) where {Tv} = Base.unsafe_convert(Ptr{cholmod_factor}, x)
 
+getworkspace(F::Factor) = F.workspace
+
 function _factor_components(is_ll)
     is_ll ? (:L, :U, :PtL, :UP) : (:L, :U, :PtL, :UP, :D, :LD, :DU, :PtLD, :DUP)
 end
@@ -858,32 +901,6 @@ function Base.convert(::Type{Dense{Tnew}}, A::Dense{T}) where {Tnew, T}
 end
 Base.convert(::Type{Dense{T}}, A::Dense{T}) where T = A
 
-# Just calling Dense(x) or Dense(b) will allocate new
-# `cholmod_dense_struct`s in CHOLMOD. Instead, we want to reuse
-# the existing memory. We can do this by creating new
-# `cholmod_dense_struct`s and filling them manually.
-function wrap_dense_and_ptr(x::StridedVecOrMat{T}) where {T <: VTypes}
-    dense_x = cholmod_dense_struct()
-    dense_x.nrow = size(x, 1)
-    dense_x.ncol = size(x, 2)
-    dense_x.nzmax = length(x)
-    dense_x.d = stride(x, 2)
-    dense_x.x = pointer(x)
-    dense_x.z = C_NULL
-    dense_x.xtype = xtyp(eltype(x))
-    dense_x.dtype = dtyp(eltype(x))
-    return dense_x, pointer_from_objref(dense_x)
-end
-# We need to use a special handling for the case of `Dense`
-# input arrays since the `pointer` refers to the pointer to the
-# `cholmod_dense`, not to the array values themselves as for
-# standard arrays.
-function wrap_dense_and_ptr(x::Dense{T}) where {T <: VTypes}
-    dense_x_ptr = x.ptr
-    dense_x = unsafe_load(dense_x_ptr)
-    return dense_x, pointer_from_objref(dense_x)
-end
-
 # This constructor assumes zero based colptr and rowval
 function Sparse(m::Integer, n::Integer,
         colptr0::Vector{Ti}, rowval0::Vector{Ti},
@@ -1326,8 +1343,8 @@ end
 @inline function getproperty(F::Factor, sym::Symbol)
     if sym === :p
         return get_perm(F)
-    elseif sym === :ptr
-        return getfield(F, :ptr)
+    elseif sym === :ptr || sym === :workspace || sym === :lock # add workspace and lock
+        return getfield(F, sym)
     else
         return FactorComponent(F, sym)
     end
@@ -1434,11 +1451,11 @@ end
 
 function cholesky!(F::Factor{Tv}, A::Sparse{Tv};
                    shift::Real=0.0, check::Bool = true) where Tv
-    # Compute the numerical factorization
-    @cholmod_param final_ll = true begin
-        factorize_p!(A, shift, F)
+    @lock F.lock begin
+        @cholmod_param final_ll = true begin
+            factorize_p!(A, shift, F)
+        end
     end
-
     check && (issuccess(F) || throw(LinearAlgebra.PosDefException(1)))
     return F
 end
@@ -1609,12 +1626,10 @@ LinearAlgebra._cholesky(A::Union{SparseMatrixCSC{T}, SparseMatrixCSC{Complex{T}}
 
 function ldlt!(F::Factor{Tv}, A::Sparse{Tv};
                shift::Real=0.0, check::Bool = true) where Tv
-    # Makes it an LDLt
-    change_factor!(F, false, false, true, false)
-
-    # Compute the numerical factorization
-    factorize_p!(A, shift, F)
-
+    @lock F.lock begin
+        change_factor!(F, false, false, true, false)
+        factorize_p!(A, shift, F)
+    end
     check && (issuccess(F) || throw(LinearAlgebra.ZeroPivotException(1)))
     return F
 end
@@ -1916,118 +1931,53 @@ const AbstractSparseVecOrMatInclAdjAndTrans = Union{AbstractSparseVecOrMat, AdjO
     throw(ArgumentError("self-adjoint sparse system solve not implemented for sparse rhs B," *
         " consider to convert B to a dense array"))
 
-"""
-    CholmodWorkspace()
-
-Reusable workspace for allocation-free calls to `ldiv!(x, L, b, ws)`.
-
-`cholmod_solve2` allocates two temporary dense matrices (Y and E) on every solve.
-A `CholmodWorkspace` holds those buffers across calls so that CHOLMOD reuses them
-instead of reallocating. It also pre-allocates the `cholmod_dense_struct` wrappers
-for `x` and `b`.
-This eliminates all Julia-side heap allocations per solve.
-
-A finalizer frees Y adn E automatically when the worspace is garbage-collected.
-
-A single workspace can be reused across calls with different arrays or a different number of RHS columns since
-CHOLMOD will resize Y and E.
-
-# Example
-```julia
-F = cholesky(A)
-ws = CholmodWorkspace()
-for b in rhs_list
-    ldiv!(x, F, b, ws)
-end
-```
-"""
-
-mutable struct CholmodWorkspace
-    dense_x::cholmod_dense_struct
-    dense_b::cholmod_dense_struct
-    X::Ref{Ptr{cholmod_dense_struct}}
-    Y::Ref{Ptr{cholmod_dense_struct}}
-    E::Ref{Ptr{cholmod_dense_struct}}
+# Julia array -> fill ws.dense_b fields and return pointer to it
+# CHOLMOD Dense object -> struct is already set up by CHOLMOD, return b.ptr directly
+@inline function _setup_bptr(b::StridedVecOrMat{T}, ws::CholmodWorkspace) where T<:VTypes
+    ws.dense_b.nrow  = size(b, 1)
+    ws.dense_b.ncol  = size(b, 2)
+    ws.dense_b.nzmax = length(b)
+    ws.dense_b.d     = stride(b, 2)
+    ws.dense_b.x     = pointer(b)
+    ws.dense_b.z     = C_NULL
+    ws.dense_b.xtype = xtyp(T)
+    ws.dense_b.dtype = dtyp(T)
+    return Ptr{cholmod_dense_struct}(pointer_from_objref(ws.dense_b))
+end
+@inline _setup_bptr(b::Dense{<:VTypes}, ::CholmodWorkspace) = b.ptr
 
-    function CholmodWorkspace()
-        ws = new(
-            cholmod_dense_struct(), # all fields will be written in ldiv! before use
-            cholmod_dense_struct(),
-            Ref(Ptr{cholmod_dense_struct}(C_NULL)),
-            Ref(Ptr{cholmod_dense_struct}(C_NULL)),
-            Ref(Ptr{cholmod_dense_struct}(C_NULL))
-        )
-        finalizer(ws) do w
-            w.Y[] != C_NULL && free!(w.Y[])
-            w.E[] != C_NULL && free!(w.E[])
-        end
-        ws
-    end
-end
-
-# in-place ldiv! and optional allocation-free ldiv!
 for TI in IndexTypes
-    @eval function ldiv!(x::StridedVecOrMat{T},
-                         L::Factor{T, $TI},
-                         b::StridedVecOrMat{T}) where {T<:VTypes}
-        if x === b
-            throw(ArgumentError("output array must not be aliased with input array"))
-        end
-        if size(L, 1) != size(b, 1)
-            throw(DimensionMismatch("Factorization and RHS should have the same number of rows. " *
-                "Factorization has $(size(L, 2)) rows, but RHS has $(size(b, 1)) rows."))
-        end
-        if size(L, 2) != size(x, 1)
-            throw(DimensionMismatch("Factorization and solution should match sizes. " *
-                "Factorization has $(size(L, 1)) columns, but solution has $(size(x, 1)) rows."))
-        end
-        if size(x, 2) != size(b, 2)
-            throw(DimensionMismatch("Solution and RHS should have the same number of columns. " *
-                "Solution has $(size(x, 2)) columns, but RHS has $(size(b, 2)) columns."))
-        end
-        if !issuccess(L)
-            s = unsafe_load(pointer(L))
-            if s.is_ll == 1
-                throw(LinearAlgebra.PosDefException(s.minor))
-            else
-                throw(LinearAlgebra.ZeroPivotException(s.minor))
+    @eval function solve!(x::StridedVecOrMat{T},
+                          L::Factor{T, $TI},
+                          b::StridedVecOrMat{T};
+                          ws = getworkspace(L)) where {T<:VTypes}
+        @lock L.lock begin
+            ws.dense_x.nrow  = size(x, 1)
+            ws.dense_x.ncol  = size(x, 2)
+            ws.dense_x.nzmax = length(x)
+            ws.dense_x.d     = stride(x, 2)
+            ws.dense_x.x     = pointer(x)
+            ws.dense_x.z     = C_NULL
+            ws.dense_x.xtype = xtyp(T)
+            ws.dense_x.dtype = dtyp(T)
+
+            ws.X[] = Ptr{cholmod_dense_struct}(pointer_from_objref(ws.dense_x))
+            Bptr = _setup_bptr(b, ws)
+            status = GC.@preserve x b ws begin
+                $(cholname(:solve2, TI))(
+                    CHOLMOD_A, L,
+                    Bptr, C_NULL,
+                    ws.X, C_NULL,
+                    ws.Y, ws.E,
+                    getcommon($TI))
             end
+            @assert !iszero(status)
         end
-
-        # Just calling Dense(x) or Dense(b) will allocate new
-        # `cholmod_dense_struct`s in CHOLMOD. Instead, we want to reuse
-        # the existing memory. We can do this by creating new
-        # `cholmod_dense_struct`s and filling them manually.
-        dense_x, dense_x_ptr = wrap_dense_and_ptr(x)
-        dense_b, dense_b_ptr = wrap_dense_and_ptr(b)
-
-        X_Handle = Ptr{cholmod_dense_struct}(dense_x_ptr)
-        Y_Handle = Ptr{cholmod_dense_struct}(C_NULL)
-        E_Handle = Ptr{cholmod_dense_struct}(C_NULL)
-        status = GC.@preserve x dense_x b dense_b begin
-            $(cholname(:solve2, TI))(
-                CHOLMOD_A, L,
-                Ref(dense_b), C_NULL,
-                Ref(X_Handle), C_NULL,
-                Ref(Y_Handle),
-                Ref(E_Handle),
-                getcommon($TI))
-        end
-        if Y_Handle != C_NULL
-            free!(Y_Handle)
-        end
-        if E_Handle != C_NULL
-            free!(E_Handle)
-        end
-        @assert !iszero(status)
-
-        return x
     end
 
     @eval function ldiv!(x::StridedVecOrMat{T},
                          L::Factor{T, $TI},
-                         b::StridedVecOrMat{T},
-                         ws::CholmodWorkspace) where {T<:VTypes}
+                         b::StridedVecOrMat{T}) where {T<:VTypes}
         if x === b
             throw(ArgumentError("output array must not be aliased with input array"))
         end
@@ -2051,38 +2001,7 @@ for TI in IndexTypes
                 throw(LinearAlgebra.ZeroPivotException(s.minor))
             end
         end
-
-        # Update all cholmod_dense_structs fields on every call in case dimensions or T change
-        ws.dense_x.nrow  = size(x, 1)
-        ws.dense_x.ncol  = size(x, 2)
-        ws.dense_x.nzmax = length(x)
-        ws.dense_x.d     = stride(x, 2)
-        ws.dense_x.x     = pointer(x)
-        ws.dense_x.z     = C_NULL
-        ws.dense_x.xtype = xtyp(T)
-        ws.dense_x.dtype = dtyp(T)
-
-        ws.dense_b.nrow  = size(b, 1)
-        ws.dense_b.ncol  = size(b, 2)
-        ws.dense_b.nzmax = length(b)
-        ws.dense_b.d     = stride(b, 2)
-        ws.dense_b.x     = pointer(b)
-        ws.dense_b.z     = C_NULL
-        ws.dense_b.xtype = xtyp(T)
-        ws.dense_b.dtype = dtyp(T)
-
-        ws.X[] = Ptr{cholmod_dense_struct}(pointer_from_objref(ws.dense_x))
-        Bptr = Ptr{cholmod_dense_struct}(pointer_from_objref(ws.dense_b))
-        status = GC.@preserve x b ws begin
-            $(cholname(:solve2, TI))(
-                CHOLMOD_A, L,
-                Bptr, C_NULL,
-                ws.X, C_NULL,
-                ws.Y, ws.E,
-                getcommon($TI))
-        end
-        @assert !iszero(status)
-
+        solve!(x, L, b)
         return x
     end
 end