Skip to content
This repository was archived by the owner on Dec 16, 2024. It is now read-only.

Allocate umfpack's workspace on julia's side and reuse it #67

Closed
wants to merge 4 commits into from
Closed

Allocate umfpack's workspace on julia's side and reuse it #67

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
7093e4a
Allocate umfpack's workspace on julia's side and reuse it
SobhanMP May 27, 2022
3434903
simplify
rayegun May 28, 2022
679a298
Add alloc solve to the tests
SobhanMP May 29, 2022
8443aeb
add 4 arg ldiv!
SobhanMP May 29, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
130 changes: 81 additions & 49 deletions src/umfpack.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -137,6 +137,18 @@ function show_umf_info(level::Real = 2.0)
umfpack_dl_report_info(umf_ctrl, umf_info)
umf_ctrl[1] = old_prt
end
"""
Working space for Umfpack so ldiv! doesn't reallocate
"""
struct UmfpackWS{T<:UMFITypes}
Wi::Vector{T}
W::Vector{Float64}
end

Base.similar(w::UmfpackWS) = UmfpackWS(similar(w.Wi), similar(w.W))
# TODO, this actually doesn't need to be this big if iter refinement is off
UmfpackWS(S::SparseMatrixCSC{Float64, T}) where T= UmfpackWS{T}(Vector{T}(undef, size(S, 2)), Vector{Float64}(undef, 5 * size(S, 2)))
UmfpackWS(S::SparseMatrixCSC{ComplexF64, T}) where T = UmfpackWS{T}(Vector{T}(undef, size(S, 2)), Vector{Float64}(undef, 10 * size(S, 2)))

## Should this type be immutable?
mutable struct UmfpackLU{Tv<:UMFVTypes,Ti<:UMFITypes} <: Factorization{Tv}
Expand All @@ -148,11 +160,13 @@ mutable struct UmfpackLU{Tv<:UMFVTypes,Ti<:UMFITypes} <: Factorization{Tv}
rowval::Vector{Ti} # 0-based row indices
nzval::Vector{Tv}
status::Int
workspace::UmfpackWS
end

Base.adjoint(F::UmfpackLU) = Adjoint(F)
Base.transpose(F::UmfpackLU) = Transpose(F)


"""
lu(A::SparseMatrixCSC; check = true) -> F::UmfpackLU

Expand Down Expand Up @@ -197,12 +211,12 @@ See also [`lu!`](@ref)

[^ACM832]: Davis, Timothy A. (2004b). Algorithm 832: UMFPACK V4.3---an Unsymmetric-Pattern Multifrontal Method. ACM Trans. Math. Softw., 30(2), 196–199. [doi:10.1145/992200.992206](https://doi.org/10.1145/992200.992206)
"""
function lu(S::SparseMatrixCSC{<:UMFVTypes,<:UMFITypes}; check::Bool = true)
function lu(S::SparseMatrixCSC{<:UMFVTypes, Ti}; check::Bool = true) where {Ti <:UMFITypes}
zerobased = getcolptr(S)[1] == 0
res = UmfpackLU(C_NULL, C_NULL, size(S, 1), size(S, 2),
zerobased ? copy(getcolptr(S)) : decrement(getcolptr(S)),
zerobased ? copy(rowvals(S)) : decrement(rowvals(S)),
copy(nonzeros(S)), 0)
copy(nonzeros(S)), 0, UmfpackWS(S))
finalizer(umfpack_free_symbolic, res)
umfpack_numeric!(res)
check && (issuccess(res) || throw(LinearAlgebra.SingularException(0)))
Expand Down Expand Up @@ -267,12 +281,17 @@ julia> F \\ ones(2)
"""
function lu!(F::UmfpackLU, S::SparseMatrixCSC{<:UMFVTypes,<:UMFITypes}; check::Bool=true)
zerobased = getcolptr(S)[1] == 0
# resize workspace if needed
if F.n < size(S, 2)
F.workspace = UmfpackWS(S)
end

F.m = size(S, 1)
F.n = size(S, 2)
F.colptr = zerobased ? copy(getcolptr(S)) : decrement(getcolptr(S))
F.rowval = zerobased ? copy(rowvals(S)) : decrement(rowvals(S))
F.nzval = copy(nonzeros(S))

umfpack_numeric!(F, reuse_numeric = false)
check && (issuccess(F) || throw(LinearAlgebra.SingularException(0)))
return F
Expand Down Expand Up @@ -324,7 +343,8 @@ function deserialize(s::AbstractSerializer, t::Type{UmfpackLU{Tv,Ti}}) where {Tv
rowval = deserialize(s)
nzval = deserialize(s)
status = deserialize(s)
obj = UmfpackLU{Tv,Ti}(symbolic, numeric, m, n, colptr, rowval, nzval, status)
ws = deserialize(s)
obj = UmfpackLU{Tv,Ti}(symbolic, numeric, m, n, colptr, rowval, nzval, status, ws)

finalizer(umfpack_free_symbolic, obj)

Expand Down Expand Up @@ -363,6 +383,8 @@ for itype in UmfpackIndexTypes
num_c = Symbol(umf_nm("numeric", :ComplexF64, itype))
sol_r = Symbol(umf_nm("solve", :Float64, itype))
sol_c = Symbol(umf_nm("solve", :ComplexF64, itype))
wsol_r = Symbol(umf_nm("wsolve", :Float64, itype))
wsol_c = Symbol(umf_nm("wsolve", :ComplexF64, itype))
det_r = Symbol(umf_nm("get_determinant", :Float64, itype))
det_z = Symbol(umf_nm("get_determinant", :ComplexF64, itype))
lunz_r = Symbol(umf_nm("get_lunz", :Float64, itype))
Expand Down Expand Up @@ -412,32 +434,41 @@ for itype in UmfpackIndexTypes
U.numeric = tmp[1]
return U
end
function solve!(x::StridedVector{Float64}, lu::UmfpackLU{Float64,$itype}, b::StridedVector{Float64}, typ::Integer)
function solve!(
x::StridedVector{Float64}, lu::UmfpackLU{Float64,$itype}, b::StridedVector{Float64}, typ::Integer;
Wi::StridedVector{$itype} = lu.workspace.Wi, W::StridedVector{Float64} = lu.workspace.W)
if x === b
throw(ArgumentError("output array must not be aliased with input array"))
end
if stride(x, 1) != 1 || stride(b, 1) != 1
throw(ArgumentError("in and output vectors must have unit strides"))
end
if size(lu, 2) > length(Wi)
throw(ArgumentError("Wi should be at least larger than `size(Af, 2)`"))
end
umfpack_numeric!(lu)
(size(b,1) == lu.m) && (size(b) == size(x)) || throw(DimensionMismatch())
@isok $sol_r(typ, lu.colptr, lu.rowval, lu.nzval,
@isok $wsol_r(typ, lu.colptr, lu.rowval, lu.nzval,
x, b, lu.numeric, umf_ctrl,
umf_info)
umf_info, Wi, W)
return x
end
function solve!(x::StridedVector{ComplexF64}, lu::UmfpackLU{ComplexF64,$itype}, b::StridedVector{ComplexF64}, typ::Integer)
function solve!(
x::StridedVector{ComplexF64}, lu::UmfpackLU{ComplexF64,$itype}, b::StridedVector{ComplexF64}, typ::Integer;
Wi::StridedVector{$itype} = lu.workspace.Wi, W::StridedVector{Float64} = lu.workspace.W)
if x === b
throw(ArgumentError("output array must not be aliased with input array"))
end
if stride(x, 1) != 1 || stride(b, 1) != 1
throw(ArgumentError("in and output vectors must have unit strides"))
end
if size(lu, 2) > length(Wi)
throw(ArgumentError("Wi should be at least larger than `size(Af, 2)`"))
end
umfpack_numeric!(lu)
(size(b, 1) == lu.m) && (size(b) == size(x)) || throw(DimensionMismatch())
n = size(b, 1)
@isok $sol_c(typ, lu.colptr, lu.rowval, lu.nzval, C_NULL, x, C_NULL, b,
C_NULL, lu.numeric, umf_ctrl, umf_info)
@isok $wsol_c(typ, lu.colptr, lu.rowval, lu.nzval, C_NULL, x, C_NULL, b,
C_NULL, lu.numeric, umf_ctrl, umf_info, Wi, W)
return x
end
function det(lu::UmfpackLU{Float64,$itype})
Expand Down Expand Up @@ -680,63 +711,64 @@ LinearAlgebra.issuccess(lu::UmfpackLU) = lu.status == UMFPACK_OK

import LinearAlgebra.ldiv!

ldiv!(lu::UmfpackLU{T}, B::StridedVecOrMat{T}) where {T<:UMFVTypes} =
ldiv!(B, lu, copy(B))
ldiv!(translu::Transpose{T,<:UmfpackLU{T}}, B::StridedVecOrMat{T}) where {T<:UMFVTypes} =
(lu = translu.parent; ldiv!(B, transpose(lu), copy(B)))
ldiv!(adjlu::Adjoint{T,<:UmfpackLU{T}}, B::StridedVecOrMat{T}) where {T<:UMFVTypes} =
(lu = adjlu.parent; ldiv!(B, adjoint(lu), copy(B)))
ldiv!(lu::UmfpackLU{Float64}, B::StridedVecOrMat{<:Complex}) =
ldiv!(B, lu, copy(B))
ldiv!(translu::Transpose{Float64,<:UmfpackLU{Float64}}, B::StridedVecOrMat{<:Complex}) =
(lu = translu.parent; ldiv!(B, transpose(lu), copy(B)))
ldiv!(adjlu::Adjoint{Float64,<:UmfpackLU{Float64}}, B::StridedVecOrMat{<:Complex}) =
(lu = adjlu.parent; ldiv!(B, adjoint(lu), copy(B)))

ldiv!(X::StridedVecOrMat{T}, lu::UmfpackLU{T}, B::StridedVecOrMat{T}) where {T<:UMFVTypes} =
_Aq_ldiv_B!(X, lu, B, UMFPACK_A)
ldiv!(X::StridedVecOrMat{T}, translu::Transpose{T,<:UmfpackLU{T}}, B::StridedVecOrMat{T}) where {T<:UMFVTypes} =
(lu = translu.parent; _Aq_ldiv_B!(X, lu, B, UMFPACK_Aat))
ldiv!(X::StridedVecOrMat{T}, adjlu::Adjoint{T,<:UmfpackLU{T}}, B::StridedVecOrMat{T}) where {T<:UMFVTypes} =
(lu = adjlu.parent; _Aq_ldiv_B!(X, lu, B, UMFPACK_At))
ldiv!(X::StridedVecOrMat{Tb}, lu::UmfpackLU{Float64}, B::StridedVecOrMat{Tb}) where {Tb<:Complex} =
_Aq_ldiv_B!(X, lu, B, UMFPACK_A)
ldiv!(X::StridedVecOrMat{Tb}, translu::Transpose{Float64,<:UmfpackLU{Float64}}, B::StridedVecOrMat{Tb}) where {Tb<:Complex} =
(lu = translu.parent; _Aq_ldiv_B!(X, lu, B, UMFPACK_Aat))
ldiv!(X::StridedVecOrMat{Tb}, adjlu::Adjoint{Float64,<:UmfpackLU{Float64}}, B::StridedVecOrMat{Tb}) where {Tb<:Complex} =
(lu = adjlu.parent; _Aq_ldiv_B!(X, lu, B, UMFPACK_At))

function _Aq_ldiv_B!(X::StridedVecOrMat, lu::UmfpackLU, B::StridedVecOrMat, transposeoptype)
ldiv!(lu::UmfpackLU{T}, B::StridedVecOrMat{T}; workspace=lu.workspace) where {T<:UMFVTypes} =
ldiv!(B, lu, copy(B); workspace)
ldiv!(translu::Transpose{T,<:UmfpackLU{T}}, B::StridedVecOrMat{T}; workspace=translu.parent.workspace) where {T<:UMFVTypes} =
(lu = translu.parent; ldiv!(B, transpose(lu), copy(B); workspace))
ldiv!(adjlu::Adjoint{T,<:UmfpackLU{T}}, B::StridedVecOrMat{T}; workspace=adjlu.parent.workspace) where {T<:UMFVTypes} =
(lu = adjlu.parent; ldiv!(B, adjoint(lu), copy(B); workspace))
ldiv!(lu::UmfpackLU{Float64}, B::StridedVecOrMat{<:Complex}; workspace=lu.workspace) =
ldiv!(B, lu, copy(B); workspace)
ldiv!(translu::Transpose{Float64,<:UmfpackLU{Float64}}, B::StridedVecOrMat{<:Complex}; workspace=translu.parent.workspace) =
(lu = translu.parent; ldiv!(B, transpose(lu), copy(B); workspace))
ldiv!(adjlu::Adjoint{Float64,<:UmfpackLU{Float64}}, B::StridedVecOrMat{<:Complex}; workspace=adjlu.parent.workspace) =
(lu = adjlu.parent; ldiv!(B, adjoint(lu), copy(B); workspace))

ldiv!(X::StridedVecOrMat{T}, lu::UmfpackLU{T}, B::StridedVecOrMat{T}; workspace=lu.workspace) where {T<:UMFVTypes} =
_Aq_ldiv_B!(X, lu, B, UMFPACK_A, workspace)
ldiv!(X::StridedVecOrMat{T}, translu::Transpose{T,<:UmfpackLU{T}}, B::StridedVecOrMat{T}; workspace=translu.parent.workspace) where {T<:UMFVTypes} =
(lu = translu.parent; _Aq_ldiv_B!(X, lu, B, UMFPACK_Aat, workspace))
ldiv!(X::StridedVecOrMat{T}, adjlu::Adjoint{T,<:UmfpackLU{T}}, B::StridedVecOrMat{T}; workspace=adjlu.parent.workspace) where {T<:UMFVTypes} =
(lu = adjlu.parent; _Aq_ldiv_B!(X, lu, B, UMFPACK_At, workspace))
ldiv!(X::StridedVecOrMat{Tb}, lu::UmfpackLU{Float64}, B::StridedVecOrMat{Tb}; workspace=lu.workspace) where {Tb<:Complex} =
_Aq_ldiv_B!(X, lu, B, UMFPACK_A, workspace)
ldiv!(X::StridedVecOrMat{Tb}, translu::Transpose{Float64,<:UmfpackLU{Float64}}, B::StridedVecOrMat{Tb};
workspace=translu.parent.workspace) where {Tb<:Complex} =
(lu = translu.parent; _Aq_ldiv_B!(X, lu, B, UMFPACK_Aat, workspace))
ldiv!(X::StridedVecOrMat{Tb}, adjlu::Adjoint{Float64,<:UmfpackLU{Float64}}, B::StridedVecOrMat{Tb}; workspace=adjlu.parent.workspace) where {Tb<:Complex} =
(lu = adjlu.parent; _Aq_ldiv_B!(X, lu, B, UMFPACK_At, workspace))

function _Aq_ldiv_B!(X::StridedVecOrMat, lu::UmfpackLU, B::StridedVecOrMat, transposeoptype, workspace::UmfpackWS)
if size(X, 2) != size(B, 2)
throw(DimensionMismatch("input and output arrays must have same number of columns"))
end
_AqldivB_kernel!(X, lu, B, transposeoptype)
_AqldivB_kernel!(X, lu, B, transposeoptype, workspace)
return X
end
function _AqldivB_kernel!(x::StridedVector{T}, lu::UmfpackLU{T},
b::StridedVector{T}, transposeoptype) where T<:UMFVTypes
solve!(x, lu, b, transposeoptype)
b::StridedVector{T}, transposeoptype, workspace::UmfpackWS) where {T<:UMFVTypes}
solve!(x, lu, b, transposeoptype; Wi=workspace.Wi, W=workspace.W)
end
function _AqldivB_kernel!(X::StridedMatrix{T}, lu::UmfpackLU{T},
B::StridedMatrix{T}, transposeoptype) where T<:UMFVTypes
B::StridedMatrix{T}, transposeoptype, workspace::UmfpackWS) where {T<:UMFVTypes}
for col in 1:size(X, 2)
solve!(view(X, :, col), lu, view(B, :, col), transposeoptype)
solve!(view(X, :, col), lu, view(B, :, col), transposeoptype; Wi=workspace.Wi, W=workspace.W)
end
end
function _AqldivB_kernel!(x::StridedVector{Tb}, lu::UmfpackLU{Float64},
b::StridedVector{Tb}, transposeoptype) where Tb<:Complex
b::StridedVector{Tb}, transposeoptype, workspace::UmfpackWS) where Tb<:Complex
r, i = similar(b, Float64), similar(b, Float64)
solve!(r, lu, Vector{Float64}(real(b)), transposeoptype)
solve!(i, lu, Vector{Float64}(imag(b)), transposeoptype)
solve!(r, lu, Vector{Float64}(real(b)), transposeoptype; Wi=workspace.Wi, W=workspace.W)
solve!(i, lu, Vector{Float64}(imag(b)), transposeoptype; Wi=workspace.Wi, W=workspace.W)
map!(complex, x, r, i)
end
function _AqldivB_kernel!(X::StridedMatrix{Tb}, lu::UmfpackLU{Float64},
B::StridedMatrix{Tb}, transposeoptype) where Tb<:Complex
B::StridedMatrix{Tb}, transposeoptype, workspace::UmfpackWS) where Tb<:Complex
r = similar(B, Float64, size(B, 1))
i = similar(B, Float64, size(B, 1))
for j in 1:size(B, 2)
solve!(r, lu, Vector{Float64}(real(view(B, :, j))), transposeoptype)
solve!(i, lu, Vector{Float64}(imag(view(B, :, j))), transposeoptype)
solve!(r, lu, Vector{Float64}(real(view(B, :, j))), transposeoptype; Wi=workspace.Wi, W=workspace.W)
solve!(i, lu, Vector{Float64}(imag(view(B, :, j))), transposeoptype; Wi=workspace.Wi, W=workspace.W)
map!(complex, view(X, :, j), r, i)
end
end
Expand Down
79 changes: 79 additions & 0 deletions test/umfpack.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,85 @@ using Serialization
using LinearAlgebra:
I, det, issuccess, ldiv!, lu, lu!, Adjoint, Transpose, SingularException, Diagonal
using SparseArrays: nnz, sparse, sprand, sprandn, SparseMatrixCSC
for itype in UMFPACK.UmfpackIndexTypes
sol_r = Symbol(UMFPACK.umf_nm("solve", :Float64, itype))
sol_c = Symbol(UMFPACK.umf_nm("solve", :ComplexF64, itype))
@eval begin
function alloc_solve!(x::StridedVector{Float64}, lu::UmfpackLU{Float64,$itype}, b::StridedVector{Float64}, typ::Integer)
if x === b
throw(ArgumentError("output array must not be aliased with input array"))
end
if stride(x, 1) != 1 || stride(b, 1) != 1
throw(ArgumentError("in and output vectors must have unit strides"))
end
UMFPACK.umfpack_numeric!(lu)
(size(b,1) == lu.m) && (size(b) == size(x)) || throw(DimensionMismatch())
UMFPACK.@isok UMFPACK.$sol_r(typ, lu.colptr, lu.rowval, lu.nzval,
x, b, lu.numeric, UMFPACK.umf_ctrl,
UMFPACK.umf_info)
return x
end
function alloc_solve!(x::StridedVector{ComplexF64}, lu::UmfpackLU{ComplexF64,$itype}, b::StridedVector{ComplexF64}, typ::Integer)
if x === b
throw(ArgumentError("output array must not be aliased with input array"))
end
if stride(x, 1) != 1 || stride(b, 1) != 1
throw(ArgumentError("in and output vectors must have unit strides"))
end
UMFPACK.umfpack_numeric!(lu)
(size(b, 1) == lu.m) && (size(b) == size(x)) || throw(DimensionMismatch())
UMFPACK.@isok UMFPACK.$sol_c(typ, lu.colptr, lu.rowval, lu.nzval, C_NULL, x, C_NULL, b,
C_NULL, lu.numeric, UMFPACK.umf_ctrl, UMFPACK.umf_info)
return x
end
end
end

@testset "Workspace management" begin
A0 = I + sprandn(100, 100, 0.01)
b0 = randn(100)
bn0 = rand(100, 20)
@testset "Core functionality for $Tv elements" for Tv in (Float64, ComplexF64)
for Ti in Base.uniontypes(UMFPACK.UMFITypes)
A = convert(SparseMatrixCSC{Tv,Ti}, A0)
Af = lu(A)
b = convert(Vector{Tv}, b0)
x = alloc_solve!(
similar(b),
Af, b,
UMFPACK.UMFPACK_A)
@test A \ b == x
bn = convert(Matrix{Tv}, bn0)
xn = similar(bn)
for i in 1:20
xn[:, i] .= alloc_solve!(
similar(bn[:, i]),
Af, bn[:, i],
UMFPACK.UMFPACK_A)
end
@test A \ bn == xn
end
end
f = function(Tv, Ti)
A = convert(SparseMatrixCSC{Tv,Ti}, A0)
Af = lu(A)
b = convert(Vector{Tv}, b0)
x = similar(b)
ldiv!(x, Af, b)
aloc1 = @allocated ldiv!(x, Af, b)
bn = convert(Matrix{Tv}, bn0)
xn = similar(bn)
ldiv!(xn, Af, bn)
aloc2 = @allocated ldiv!(xn, Af, bn)
return aloc1 + aloc2
end
@testset "Allocations" begin
for Tv in Base.uniontypes(UMFPACK.UMFVTypes),
Ti in Base.uniontypes(UMFPACK.UMFITypes)
@test f(Tv, Ti) == 0
end
end
end

@testset "UMFPACK wrappers" begin
se33 = sparse(1.0I, 3, 3)
Expand Down