From 22c6153a3ba1dd2abc4699fbddb93bd8ac505e07 Mon Sep 17 00:00:00 2001 From: Jash Date: Mon, 15 Jun 2026 09:08:48 +0530 Subject: [PATCH 1/2] Add SuperLUDIST sparse factorization support --- Project.toml | 3 + ext/LinearSolveSuperLUDISTExt.jl | 196 +++++++++++++++++++++++++++++++ src/LinearSolve.jl | 3 +- src/extension_algs.jl | 81 +++++++++++++ test/runtests.jl | 7 ++ test/superludist/Project.toml | 17 +++ test/superludist/superludist.jl | 80 +++++++++++++ test/test_groups.toml | 5 + 8 files changed, 391 insertions(+), 1 deletion(-) create mode 100644 ext/LinearSolveSuperLUDISTExt.jl create mode 100644 test/superludist/Project.toml create mode 100644 test/superludist/superludist.jl diff --git a/Project.toml b/Project.toml index 4b7422b89..2bec19673 100644 --- a/Project.toml +++ b/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" diff --git a/ext/LinearSolveSuperLUDISTExt.jl b/ext/LinearSolveSuperLUDISTExt.jl new file mode 100644 index 000000000..811762cfc --- /dev/null +++ b/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 diff --git a/src/LinearSolve.jl b/src/LinearSolve.jl index 5098c0877..292fbd468 100644 --- a/src/LinearSolve.jl +++ b/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 diff --git a/src/extension_algs.jl b/src/extension_algs.jl index d080da6b7..16f908943 100644 --- a/src/extension_algs.jl +++ b/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...)` diff --git a/test/runtests.jl b/test/runtests.jl index 09484917b..031ca48e0 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -194,6 +194,13 @@ else end return nothing end, + "LinearSolveSuperLUDIST" => function () + if Base.Sys.islinux() && HAS_EXTENSIONS + activate_group_env(joinpath(@__DIR__, "LinearSolveSuperLUDIST")) + @time @safetestset "LinearSolveSuperLUDIST" include("LinearSolveSuperLUDIST/superludist.jl") + end + return nothing + end, "Trim" => function () if VERSION >= v"1.12.0" activate_group_env(joinpath(@__DIR__, "Trim")) diff --git a/test/superludist/Project.toml b/test/superludist/Project.toml new file mode 100644 index 000000000..d17fa3323 --- /dev/null +++ b/test/superludist/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" diff --git a/test/superludist/superludist.jl b/test/superludist/superludist.jl new file mode 100644 index 000000000..c08a0b231 --- /dev/null +++ b/test/superludist/superludist.jl @@ -0,0 +1,80 @@ +using LinearSolve +using MPI +using SparseArrays +using SuperLUDIST +using Test + +MPI.Initialized() || MPI.Init() +const SuperLUDISTExt = Base.get_extension(LinearSolve, :LinearSolveSuperLUDISTExt) + +function residual_ok(A, x, b; atol = 1.0e-8, rtol = 1.0e-8) + return isapprox(A * x, b; atol = atol, rtol = rtol) +end + +@testset "SuperLUDISTFactorization sparse solve and reuse" begin + A = sparse( + [ + 4.0 1.0 0.0 + 2.0 3.0 1.0 + 0.0 1.0 2.0 + ] + ) + b1 = [1.0, 2.0, 3.0] + b2 = [3.0, 2.0, 1.0] + A2 = sparse( + [ + 5.0 1.0 0.0 + 2.0 4.0 1.0 + 0.0 1.0 3.0 + ] + ) + + alg = SuperLUDISTFactorization(; comm = MPI.COMM_SELF) + sol = solve(LinearProblem(A, b1), alg) + @test sol.retcode == ReturnCode.Success + @test residual_ok(A, sol.u, b1) + SuperLUDISTExt.cleanup_superludist_cache!(sol) + + cache = init(LinearProblem(A, b1), alg) + sol1 = solve!(cache) + @test sol1.retcode == ReturnCode.Success + @test residual_ok(A, sol1.u, b1) + + cache.b = b2 + sol2 = solve!(cache) + @test sol2.retcode == ReturnCode.Success + @test residual_ok(A, sol2.u, b2) + + cache.A = A2 + sol3 = solve!(cache) + @test sol3.retcode == ReturnCode.Success + @test residual_ok(A2, sol3.u, b2) + SuperLUDISTExt.cleanup_superludist_cache!(cache) +end + +@testset "SuperLUDISTFactorization rejects unsupported complex inputs" begin + A = sparse( + ComplexF64[ + 3.0 + 0im 1.0 - 1im + 2.0 + 1im 4.0 + 0im + ] + ) + x = ComplexF64[1.0 - 1im, 2.0 + 0.5im] + b = A * x + + sol = solve(LinearProblem(A, b), SuperLUDISTFactorization()) + @test sol.retcode == ReturnCode.Failure +end + +@testset "SuperLUDISTFactorization rejects unsupported scalar types" begin + A = sparse( + BigFloat[ + 4 1 + 2 3 + ] + ) + b = BigFloat[1, 2] + + sol = solve(LinearProblem(A, b), SuperLUDISTFactorization()) + @test sol.retcode == ReturnCode.Failure +end diff --git a/test/test_groups.toml b/test/test_groups.toml index 058d9b955..1aacac66c 100644 --- a/test/test_groups.toml +++ b/test/test_groups.toml @@ -63,6 +63,11 @@ versions = ["lts", "1"] runner = "self-hosted" num_threads = 2 +[LinearSolveSuperLUDIST] +versions = ["lts", "1"] +runner = "self-hosted" +num_threads = 2 + [LinearSolveMUMPS] versions = ["lts", "1"] runner = "self-hosted" From 3448fd32a5803020916881f08516488dc81ef17a Mon Sep 17 00:00:00 2001 From: Jash Date: Mon, 15 Jun 2026 09:15:47 +0530 Subject: [PATCH 2/2] Align SuperLUDIST test group path with test harness --- test/{superludist => LinearSolveSuperLUDIST}/Project.toml | 0 test/{superludist => LinearSolveSuperLUDIST}/superludist.jl | 0 2 files changed, 0 insertions(+), 0 deletions(-) rename test/{superludist => LinearSolveSuperLUDIST}/Project.toml (100%) rename test/{superludist => LinearSolveSuperLUDIST}/superludist.jl (100%) diff --git a/test/superludist/Project.toml b/test/LinearSolveSuperLUDIST/Project.toml similarity index 100% rename from test/superludist/Project.toml rename to test/LinearSolveSuperLUDIST/Project.toml diff --git a/test/superludist/superludist.jl b/test/LinearSolveSuperLUDIST/superludist.jl similarity index 100% rename from test/superludist/superludist.jl rename to test/LinearSolveSuperLUDIST/superludist.jl