diff --git a/Project.toml b/Project.toml index 7923060c4..d7835f9f4 100644 --- a/Project.toml +++ b/Project.toml @@ -65,6 +65,7 @@ PartitionedArrays = "5a9dfac6-5c52-46f7-8278-5e2210713be9" PartitionedSolvers = "11b65f7f-80ac-401b-9ef2-3db765482d62" PureUMFPACK = "b7e1f0a2-3c4d-4e5f-9a0b-1c2d3e4f5a6b" RecursiveFactorization = "f2c3362d-daeb-58d1-803e-2bc74f2840b4" +TriangularSolve = "d5829a12-d9aa-46ab-831f-fb7c9ab06edf" STRUMPACK_jll = "86fbd0b9-476f-557c-b766-62c724b42d8c" SparseMatricesCSR = "a0a7dd2c-ebf4-11e9-1f05-cf50bc540ca1" Sparspak = "e56a9233-b9d6-4f03-8d0f-1825330902ac" @@ -106,7 +107,7 @@ LinearSolvePartitionedSolversExt = ["PartitionedArrays", "PartitionedSolvers"] LinearSolveParUExt = ["ParU_jll", "SparseArrays"] LinearSolvePureUMFPACKExt = ["PureUMFPACK", "SparseArrays"] LinearSolvePardisoExt = ["Pardiso", "SparseArrays"] -LinearSolveRecursiveFactorizationExt = "RecursiveFactorization" +LinearSolveRecursiveFactorizationExt = ["RecursiveFactorization", "TriangularSolve"] LinearSolveSuperLUDISTExt = ["SparseArrays", "SuperLUDIST"] LinearSolveSTRUMPACKExt = ["SparseArrays", "STRUMPACK_jll"] LinearSolveSparseArraysExt = "SparseArrays" @@ -174,6 +175,7 @@ PureUMFPACK = "0.1" Random = "1.10" RecursiveArrayTools = "3.37, 4" RecursiveFactorization = "0.2.26" +TriangularSolve = "0.2" Reexport = "1.2.2" STRUMPACK_jll = "8" SafeTestsets = "0.1" diff --git a/ext/LinearSolveRecursiveFactorizationExt.jl b/ext/LinearSolveRecursiveFactorizationExt.jl index e387be470..20e9eaff2 100644 --- a/ext/LinearSolveRecursiveFactorizationExt.jl +++ b/ext/LinearSolveRecursiveFactorizationExt.jl @@ -4,6 +4,7 @@ using LinearSolve: LinearSolve, userecursivefactorization, LinearCache, @get_cac RFLUFactorization, ButterflyFactorization, RF32MixedLUFactorization, default_alias_A, default_alias_b, LinearVerbosity using LinearSolve.LinearAlgebra, LinearSolve.ArrayInterface, RecursiveFactorization +using TriangularSolve: TriangularSolve using SciMLBase: SciMLBase, ReturnCode using SciMLLogging: @SciMLMessage @@ -176,4 +177,31 @@ function LinearSolve._custom_adjoint_factorization_solve( return solution[1:n] end +# ---- SupernodalLU panel triangular solves --------------------------------- +# The vendored supernodal sparse LU (src/SupernodalLU) applies two BLAS-3 +# trsms per supernode against its just-factored diagonal block: the L21 panel +# on the right by U11, and the U12 panel on the left by unit-L11. Route them +# through TriangularSolve, which RecursiveFactorization already depends on +# and uses for its own trsms — so when RFLU is the dense default, the sparse +# solver's panel work runs on the same kernels. Measured: recovers the +# 2D-mesh refactorization gap left by the stdlib trsms. +const SNLU = LinearSolve.SupernodalLU +const SNLUTypes = Union{Float32, Float64} + +function SNLU._panel_rdiv!(W::Matrix{Tv}, np::Int, len::Int) where {Tv <: SNLUTypes} + len > np || return nothing + TriangularSolve.rdiv!( + view(W, (np + 1):len, 1:np), UpperTriangular(view(W, 1:np, 1:np)), Val(false) + ) + return nothing +end + +function SNLU._panel_ldiv!(W::Matrix{Tv}, np::Int, Z::Matrix{Tv}) where {Tv <: SNLUTypes} + isempty(Z) && return nothing + TriangularSolve.ldiv!( + UnitLowerTriangular(view(W, 1:np, 1:np)), Z, Val(false) + ) + return nothing +end + end diff --git a/src/SupernodalLU/numeric.jl b/src/SupernodalLU/numeric.jl index d29bd3919..053d9ecbe 100644 --- a/src/SupernodalLU/numeric.jl +++ b/src/SupernodalLU/numeric.jl @@ -267,6 +267,23 @@ function _cache_lu!( return true end +# Panel triangular solves (L21 := L21·U11⁻¹ and U12 := L11⁻¹·U12). These are +# BLAS-3 trsms against the just-factored diagonal block, not linear solves in +# the LinearSolve sense, so they are plain overridable hooks: the defaults +# below use the stdlib triangular solves, and +# LinearSolveRecursiveFactorizationExt routes them through TriangularSolve — +# the same library RecursiveFactorization (and hence the RFLU dense default) +# already uses internally for its own trsms. +function _panel_rdiv!(W::Matrix{Tv}, np::Int, len::Int) where {Tv} + rdiv!(view(W, (np + 1):len, 1:np), UpperTriangular(view(W, 1:np, 1:np))) + return nothing +end + +function _panel_ldiv!(W::Matrix{Tv}, np::Int, Z::Matrix{Tv}) where {Tv} + ldiv!(UnitLowerTriangular(view(W, 1:np, 1:np)), Z) + return nothing +end + _lu_from_cacheval(cv::LinearAlgebra.LU) = cv _lu_from_cacheval(cv::Tuple) = _lu_from_cacheval(first(cv)) @@ -490,8 +507,8 @@ function _process_supernode!( end if nu > 0 - rdiv!(view(Ws, (np + 1):len, 1:np), UpperTriangular(view(Ws, 1:np, 1:np))) - ldiv!(UnitLowerTriangular(view(Ws, 1:np, 1:np)), Zs) + _panel_rdiv!(Ws, np, len) + _panel_ldiv!(Ws, np, Zs) end for a in 1:np