Add SuperLUDIST sparse factorization support

Project.toml

@@ -66,6 +66,7 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SparseMatricesCSR = "a0a7dd2c-ebf4-11e9-1f05-cf50bc540ca1"
 Sparspak = "e56a9233-b9d6-4f03-8d0f-1825330902ac"
 SpecializingFactorizations = "fa08b7a1-13d3-4faf-875d-5cbc1520e3f3"
+SuperLUDIST = "4cd002a6-0da4-410d-a012-232df062f478"
 blis_jll = "6136c539-28a5-5bf0-87cc-b183200dce32"
 
 [extensions]
@@ -100,6 +101,7 @@ LinearSolveParUExt = ["ParU_jll", "SparseArrays"]
 LinearSolvePureUMFPACKExt = ["PureUMFPACK", "SparseArrays"]
 LinearSolvePardisoExt = ["Pardiso", "SparseArrays"]
 LinearSolveRecursiveFactorizationExt = "RecursiveFactorization"
+LinearSolveSuperLUDISTExt = ["SparseArrays", "SuperLUDIST"]
 LinearSolveSTRUMPACKExt = ["SparseArrays", "STRUMPACK_jll"]
 LinearSolveSparseArraysExt = "SparseArrays"
 LinearSolveSparspakExt = ["SparseArrays", "Sparspak"]
@@ -176,6 +178,7 @@ Sparspak = "0.3.9"
 SpecializingFactorizations = "0.1"
 StaticArrays = "1.9"
 StaticArraysCore = "1.4.3"
+SuperLUDIST = "1"
 Test = "1.10"
 Zygote = "0.7"
 blis_jll = "0.9.0"

ext/LinearSolveSuperLUDISTExt.jl

@@ -0,0 +1,196 @@
+module LinearSolveSuperLUDISTExt
+
+using LinearAlgebra
+using LinearSolve
+using LinearSolve: LinearVerbosity, OperatorAssumptions
+using SparseArrays
+using SuperLUDIST
+import SciMLBase
+import SciMLBase: ReturnCode, LinearSolution
+using SciMLLogging: @SciMLMessage
+
+const MPI = SuperLUDIST.MPI
+const SparseBase = SuperLUDIST.SparseBase
+const CIndex = SuperLUDIST.CIndex
+
+mutable struct SuperLUDISTCache
+    factor::Any
+
+    function SuperLUDISTCache()
+        return new(nothing)
+    end
+end
+
+cleanup_superludist_cache!(cache::SuperLUDISTCache) = (cache.factor = nothing; cache)
+cleanup_superludist_cache!(cache::LinearSolve.LinearCache) = cleanup_superludist_cache!(cache.cacheval)
+cleanup_superludist_cache!(sol::LinearSolution) = cleanup_superludist_cache!(sol.cache.cacheval)
+
+LinearSolve.needs_concrete_A(::LinearSolve.SuperLUDISTFactorization) = true
+
+function LinearSolve.init_cacheval(
+        alg::LinearSolve.SuperLUDISTFactorization, A, b, u, Pl, Pr, maxiters::Int,
+        abstol, reltol, verbose::Union{LinearVerbosity, Bool}, assumptions::OperatorAssumptions
+    )
+    return SuperLUDISTCache()
+end
+
+_superlu_eltype(::Type{Float32}) = Float32
+_superlu_eltype(::Type{Float64}) = Float64
+function _superlu_eltype(::Type{T}) where {T}
+    throw(
+        ArgumentError(
+            "SuperLUDISTFactorization only supports Float32 and Float64 inputs; got element type $T"
+        )
+    )
+end
+
+function _superlu_index_type(A::SparseMatrixCSC)
+    limit = max(size(A, 1), size(A, 2), nnz(A) + 1)
+    return limit <= typemax(Int32) ? Int32 : Int64
+end
+
+function _grid_dims(nprocs::Integer, nprow::Integer, npcol::Integer)
+    if nprow == 0 && npcol == 0
+        root = floor(Int, sqrt(nprocs))
+        for r in root:-1:1
+            if nprocs % r == 0
+                return r, nprocs ÷ r
+            end
+        end
+    elseif nprow == 0
+        nprocs % npcol == 0 || error("npcol=$npcol does not divide communicator size $nprocs")
+        return nprocs ÷ npcol, npcol
+    elseif npcol == 0
+        nprocs % nprow == 0 || error("nprow=$nprow does not divide communicator size $nprocs")
+        return nprow, nprocs ÷ nprow
+    else
+        nprow * npcol == nprocs || error(
+            "nprow*npcol must equal communicator size ($nprow * $npcol != $nprocs)"
+        )
+        return nprow, npcol
+    end
+    error("could not determine a valid SuperLU_DIST process grid")
+end
+
+function _build_options(opt)
+    options = SuperLUDIST.Options()
+    if opt === nothing
+        return options
+    elseif opt isa NamedTuple
+        for (name, value) in pairs(opt)
+            setproperty!(options, name, value)
+        end
+        return options
+    elseif opt isa SuperLUDIST.Options
+        return deepcopy(opt)
+    else
+        throw(
+            ArgumentError(
+                "options must be nothing, a NamedTuple, or a SuperLUDIST.Options object; got $(typeof(opt))"
+            )
+        )
+    end
+end
+
+function _to_superlu_store(A::SparseMatrixCSC{T}, ::Type{I}) where {T, I}
+    ptr = CIndex{I}.(A.colptr)
+    idx = CIndex{I}.(A.rowval)
+    vals = Vector{T}(A.nzval)
+    return SparseBase.CSCStore(ptr, idx, vals, size(A))
+end
+
+function _to_superlu_matrix(A::SparseMatrixCSC)
+    T = _superlu_eltype(eltype(A))
+    I = _superlu_index_type(A)
+    Ac = SparseMatrixCSC{T, Int}(A)
+    store = _to_superlu_store(Ac, I)
+    return Ac, store, I
+end
+
+function _rhs_matrix(b, ::Type{T}) where {T}
+    if b isa AbstractVector
+        return reshape(T.(collect(b)), :, 1), true
+    else
+        return Matrix{T}(b), false
+    end
+end
+
+function _copy_solution!(u::AbstractVector, x::AbstractMatrix)
+    copyto!(u, vec(x))
+    return u
+end
+
+function _copy_solution!(u::AbstractMatrix, x::AbstractMatrix)
+    copyto!(u, x)
+    return u
+end
+
+function _solve_failed_solution(
+        alg::LinearSolve.SuperLUDISTFactorization,
+        cache::LinearSolve.LinearCache,
+        msg::AbstractString
+    )
+    @SciMLMessage(msg, cache.verbose, :solver_failure)
+    return SciMLBase.build_linear_solution(
+        alg, cache.u, nothing, cache; retcode = ReturnCode.Failure
+    )
+end
+
+function _build_factorization(
+        alg::LinearSolve.SuperLUDISTFactorization,
+        A::SparseMatrixCSC,
+        bmat,
+    )
+    Ac, store, I = _to_superlu_matrix(A)
+    comm = alg.comm === nothing ? MPI.COMM_SELF : alg.comm
+    nprocs = MPI.Comm_size(comm)
+    nprow, npcol = _grid_dims(nprocs, alg.nprow, alg.npcol)
+    alg.threads === nothing || SuperLUDIST.superlu_set_num_threads(I, alg.threads)
+    grid = SuperLUDIST.Grid{I}(nprow, npcol, comm)
+    Aslu = SuperLUDIST.ReplicatedSuperMatrix(store, grid)
+    options = _build_options(alg.options)
+    x, factor = SuperLUDIST.pgssvx!(Aslu, copy(bmat); options = options)
+    return x, factor, Ac
+end
+
+function SciMLBase.solve!(
+        cache::LinearSolve.LinearCache, alg::LinearSolve.SuperLUDISTFactorization;
+        kwargs...
+    )
+    A = convert(AbstractMatrix, cache.A)
+    A_sparse = A isa SparseMatrixCSC ? A : sparse(A)
+    T = try
+        _superlu_eltype(promote_type(eltype(A_sparse), eltype(cache.b)))
+    catch err
+        return _solve_failed_solution(alg, cache, sprint(showerror, err))
+    end
+    bmat, _ = _rhs_matrix(cache.b, T)
+    scache = LinearSolve.@get_cacheval(cache, :SuperLUDISTFactorization)
+
+    x = nothing
+    if cache.isfresh || scache.factor === nothing
+        cleanup_superludist_cache!(scache)
+        result = try
+            _build_factorization(alg, SparseMatrixCSC{T, Int}(A_sparse), bmat)
+        catch err
+            return _solve_failed_solution(alg, cache, sprint(showerror, err))
+        end
+        x, scache.factor = result[1], result[2]
+        cache.isfresh = false
+    else
+        x = try
+            y = copy(bmat)
+            ldiv!(scache.factor, y)
+            y
+        catch err
+            return _solve_failed_solution(alg, cache, sprint(showerror, err))
+        end
+    end
+
+    _copy_solution!(cache.u, x)
+    return SciMLBase.build_linear_solution(
+        alg, cache.u, nothing, cache; retcode = ReturnCode.Success
+    )
+end
+
+end

src/LinearSolve.jl

@@ -526,7 +526,7 @@ export LUFactorization, SVDFactorization, QRFactorization, GenericFactorization,
     SparspakFactorization, DiagonalFactorization, CholeskyFactorization,
     BunchKaufmanFactorization, CHOLMODFactorization, LDLtFactorization,
     CUSOLVERRFFactorization, CliqueTreesFactorization, ParUFactorization,
-    STRUMPACKFactorization, MUMPSFactorization,
+    STRUMPACKFactorization, MUMPSFactorization, SuperLUDISTFactorization,
     SpecializedLUFactorization, SpecializedQRFactorization,
     HSLMA57Factorization, HSLMA97Factorization
 
@@ -556,6 +556,7 @@ export MKLPardisoFactorize, MKLPardisoIterate
 export PanuaPardisoFactorize, PanuaPardisoIterate
 export PardisoJL
 export MUMPSFactorization
+export SuperLUDISTFactorization
 export MKLLUFactorization
 export OpenBLASLUFactorization
 export OpenBLAS32MixedLUFactorization

src/extension_algs.jl

@@ -1461,6 +1461,87 @@ struct MUMPSFactorization <: AbstractSparseFactorization
     end
 end
 
+"""
+`SuperLUDISTFactorization(; comm = nothing, nprow = 0, npcol = 0, options = nothing, threads = nothing)`
+
+A sparse direct solver wrapper around
+[SuperLUDIST.jl](https://github.com/JuliaSparse/SuperLUDIST.jl), backed by the
+`SuperLU_DIST_jll` artifact through that package.
+
+This wrapper targets ordinary replicated Julia sparse matrices
+(`SparseMatrixCSC`). Each participating rank builds the same Julia matrix and
+right-hand side locally, and `SuperLUDIST` performs the distributed solve on
+the requested communicator. The resulting factorization is cached inside the
+`LinearSolve` cache and reused across repeated solves with new right-hand
+sides.
+
+## Keyword Arguments
+
+  - `comm`: optional MPI communicator. `nothing` maps to `MPI.COMM_SELF`.
+  - `nprow`, `npcol`: process-grid dimensions for SuperLU_DIST. If both are
+    left as `0`, a near-square grid is chosen automatically from the
+    communicator size.
+  - `options`: either `nothing`, a `NamedTuple` of `SuperLUDIST.Options` field
+    updates, or a prebuilt `SuperLUDIST.Options` object.
+  - `threads`: optional OpenMP thread count passed through
+    `SuperLUDIST.superlu_set_num_threads`.
+
+!!! note
+
+    Using this solver requires loading `SuperLUDIST.jl` and `SparseArrays`, and
+    initializing MPI for multi-rank usage:
+    ```julia
+    using MPI, SuperLUDIST, SparseArrays
+    MPI.Init()
+    ```
+
+## Supported Element Types
+
+`Float32` and `Float64`.
+
+`ComplexF64` is intentionally rejected for now because the current upstream
+`SuperLUDIST.jl` replicated complex solve path crashes during finalization.
+
+## Example
+
+```julia
+using LinearSolve, SparseArrays, MPI, SuperLUDIST
+
+MPI.Init()
+A = sparse([4.0 1.0; 2.0 3.0])
+b = [1.0, 2.0]
+
+sol = solve(
+    LinearProblem(A, b),
+    SuperLUDISTFactorization(; comm = MPI.COMM_SELF)
+)
+```
+"""
+struct SuperLUDISTFactorization <: AbstractSparseFactorization
+    comm::Any
+    nprow::Int
+    npcol::Int
+    options::Any
+    threads::Union{Nothing, Int}
+
+    function SuperLUDISTFactorization(;
+            comm = nothing,
+            nprow::Integer = 0,
+            npcol::Integer = 0,
+            options = nothing,
+            threads::Union{Nothing, Integer} = nothing,
+        )
+        ext = Base.get_extension(@__MODULE__, :LinearSolveSuperLUDISTExt)
+        if ext === nothing
+            error("SuperLUDISTFactorization requires that SuperLUDIST and SparseArrays are loaded, i.e. `using MPI, SuperLUDIST, SparseArrays`")
+        end
+        nprow >= 0 || error("nprow must be nonnegative")
+        npcol >= 0 || error("npcol must be nonnegative")
+        threads === nothing || threads > 0 || error("threads must be positive")
+        return new(comm, Int(nprow), Int(npcol), options, threads === nothing ? nothing : Int(threads))
+    end
+end
+
 """
 `HSLMA57Factorization(; kwargs...)`
 

test/LinearSolveSuperLUDIST/Project.toml

@@ -0,0 +1,17 @@
+[deps]
+LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
+MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+SuperLUDIST = "4cd002a6-0da4-410d-a012-232df062f478"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[sources]
+LinearSolve = {path = "../.."}
+
+[compat]
+LinearSolve = "3"
+MPI = "0.20"
+SparseArrays = "1.10"
+SuperLUDIST = "1"
+Test = "1.10"
+julia = "1.10"