Skip to content

Clean up ButterflyFactorization #103

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ChrisRackauckas merged 4 commits into from
Oct 9, 2025
Merged

Clean up ButterflyFactorization #103

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
0e635a2
refactor butterfly into new struct
Shreyas-Ekanathan Oct 5, 2025
9917b57
small edit
Shreyas-Ekanathan Oct 5, 2025
280b169
fixes
Shreyas-Ekanathan Oct 8, 2025
28feecc
undo error
Shreyas-Ekanathan Oct 8, 2025
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
126 changes: 61 additions & 65 deletions src/butterflylu.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,34 +5,49 @@ using RecursiveFactorization
using SparseBandedMatrices

@inline exphalf(x) = exp(x) * oftype(x, 0.5)
function 🦋!(wv, ::Val{SEED} = Val(888)) where {SEED}
function generate_rand_butterfly_vals!(wv, ::Val{SEED} = Val(888)) where {SEED}
T = eltype(wv)
mrng = VectorizedRNG.MutableXoshift(SEED)
GC.@preserve mrng begin rand!(exphalf, VectorizedRNG.Xoshift(mrng), wv, static(0),
T(-0.05), T(0.1)) end
end

function 🦋generate_random!(A, ::Val{SEED} = Val(888)) where {SEED}
Usz = 2 * size(A, 1)
Vsz = 2 * size(A, 2)
uv = similar(A, Usz + Vsz)
🦋!(uv, Val(SEED))
(uv,)
uv = similar(A, 4 * size(A, 1))
generate_rand_butterfly_vals!(uv, Val(SEED))
uv
end

function 🦋workspace(A, b, B::Matrix{T}, U::Adjoint{T, Matrix{T}}, V::Matrix{T}, thread, ::Val{SEED} = Val(888)) where {T, SEED}
M = size(A, 1)
if (M % 4 != 0)
A = pad!(A)
struct 🦋workspace{T}
A::Matrix{T}
b::Vector{T}
B::Matrix{T}
ws::Vector{T}
U::Matrix{T}
V::Matrix{T}
out::Vector{T}
function 🦋workspace(A, b, ::Val{SEED} = Val(888)) where {SEED}
M = size(A, 1)
out = similar(b, M)
if (M % 4 != 0)
A = pad!(A)
xn = 4 - M % 4
b = [b; rand(xn)]
end
B = similar(A)
U, V = (similar(A), similar(A))
ws = 🦋generate_random!(A)
new{eltype(A)}(A, b, B, ws, U, V, out)
end
B = similar(A)
ws = 🦋generate_random!(copyto!(B, A))
🦋mul!(copyto!(B, A), ws)
U, V = materializeUV(B, ws)
F = RecursiveFactorization.lu!(B, thread)
out = similar(b, M)
end

U, V, F, out
function 🦋lu!(workspace::🦋workspace, M, thread)
Copy link
Copy Markdown
Contributor

@oscardssmith oscardssmith Oct 9, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@Shreyas-Ekanathan sorry for not mentioning this earlier, but this function has the wrong name. It doesn't return a factorization. It performs a linear solve. It also shouldn't take M as an argument. That should be stored in the workspace.

(;A, b, B, ws, U, V, out) = workspace
🦋mul!(copyto!(B, A), ws)
materializeUV(U, V, ws)
F = RecursiveFactorization.lu!(B, Val(false), thread)
sol = V * (F \ (U' * b))
out .= @view sol[1:M]
out
end

const butterfly_workspace = 🦋workspace;
Expand Down Expand Up @@ -71,7 +86,7 @@ function 🦋mul_level!(A, u, v)
end
end

function 🦋mul!(A, (uv,))
function 🦋mul!(A, uv)
M, N = size(A)
@assert M == N
Mh = M >>> 1
Expand Down Expand Up @@ -106,6 +121,13 @@ function diagnegbottom(x)
Diagonal(y), Diagonal(z)
end

function 🦋!(C::SparseBandedMatrix, A::Diagonal, B::Diagonal)
setdiagonal!(C, [A.diag; -B.diag], true)
Copy link
Copy Markdown
Contributor

@oscardssmith oscardssmith Oct 5, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[A.diag; -B.diag] allocates. Any way to avoid?

Copy link
Copy Markdown
Contributor Author

@Shreyas-Ekanathan Shreyas-Ekanathan Oct 5, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

hmm not sure to be honest

setdiagonal!(C, A.diag, true)
setdiagonal!(C, B.diag, false)
C
end

function 🦋2!(C, A::Diagonal, B::Diagonal)
Copy link
Copy Markdown
Contributor

@oscardssmith oscardssmith Oct 5, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What's the difference between 🦋2! and 🦋!? Can we unify them?

Copy link
Copy Markdown
Contributor Author

@Shreyas-Ekanathan Shreyas-Ekanathan Oct 5, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yeah this was one thing I wanted to talk about, currently we have two functions because 🦋! can't accept submatrices generated by view calls, so as of right now we need the second 🦋2! function for those cases, but ideally we would combine those.

Copy link
Copy Markdown
Contributor

@oscardssmith oscardssmith Oct 6, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think the answer here might be the same as to the issue with 🦋! allocating, in that it might be worth allocating a butterfly matrix with the proper diagonals preallocated and fill in the entries manually.

Copy link
Copy Markdown
Contributor Author

@Shreyas-Ekanathan Shreyas-Ekanathan Oct 7, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

But this wouldn't allow us to pass submatrices into the function right? so we would still need 🦋2!?

Copy link
Copy Markdown
Contributor

@oscardssmith oscardssmith Oct 7, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it could just be a separate method of 🦋!.

@assert size(A) == size(B)
A1 = size(A, 1)
Expand All @@ -120,61 +142,35 @@ function 🦋2!(C, A::Diagonal, B::Diagonal)
C
end

function 🦋!(A::Matrix, C::SparseBandedMatrix, X::Diagonal, Y::Diagonal)
@assert size(X) == size(Y)
if (size(X, 1) + size(Y, 1) != size(A, 1))
x = size(A, 1) - size(X, 1) - size(Y, 1)
setdiagonal!(C, [X.diag; rand(x); -Y.diag], true)
setdiagonal!(C, X.diag, true)
setdiagonal!(C, Y.diag, false)
else
setdiagonal!(C, [X.diag; -Y.diag], true)
setdiagonal!(C, X.diag, true)
setdiagonal!(C, Y.diag, false)
end

C
end

function 🦋2!(C::SparseBandedMatrix, A::Diagonal, B::Diagonal)
setdiagonal!(C, [A.diag; -B.diag], true)
setdiagonal!(C, A.diag, true)
setdiagonal!(C, B.diag, false)
C
end

function materializeUV(A, (uv,))
M, N = size(A)
Mh = M >>> 1
Nh = N >>> 1
function materializeUV(U, V, uv)
M = size(U, 1)
Mh = M >>> 1
Copy link
Copy Markdown
Contributor

@oscardssmith oscardssmith Oct 5, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

rename to n and n_half?


U₁u, U₁l = diagnegbottom(@view(uv[1:Mh])) #Mh
U₂u, U₂l = diagnegbottom(@view(uv[(1 + Mh + Nh):(M + Nh)])) #M2
V₁u, V₁l = diagnegbottom(@view(uv[(Mh + 1):(Mh + Nh)])) #Nh
V₂u, V₂l = diagnegbottom(@view(uv[(1 + 2 * Mh + Nh):(2 * Mh + N)])) #N2
Uu, Ul = diagnegbottom(@view(uv[(1 + M + N):(2 * M + N)])) #M
Vu, Vl = diagnegbottom(@view(uv[(1 + 2 * M + N):(2 * M + 2 * N)])) #N
U₂u, U₂l = diagnegbottom(@view(uv[(1 + 2 * Mh):(M + Mh)])) #Mh
V₁u, V₁l = diagnegbottom(@view(uv[(Mh + 1):(2 * Mh)])) #Mh
V₂u, V₂l = diagnegbottom(@view(uv[(1 + 3 * Mh):(2 * Mh + M)])) #Mh
Uu, Ul = diagnegbottom(@view(uv[(1 + 2 * M):(3 * M)])) #M
Vu, Vl = diagnegbottom(@view(uv[(1 + 3 * M):(4 * M)])) #M

Bu2 = SparseBandedMatrix{typeof(uv[1])}(undef, M, N)
Bu2 = SparseBandedMatrix{typeof(uv[1])}(undef, M, M)

🦋2!(view(Bu2, 1 : Mh, 1 : Nh), U₁u, U₁l)
🦋2!(view(Bu2, Mh + 1: M, Nh + 1: N), U₂u, U₂l)
🦋2!(view(Bu2, 1 : Mh, 1 : Mh), U₁u, U₁l)
🦋2!(view(Bu2, Mh + 1: M, Mh + 1: M), U₂u, U₂l)

Bu1 = SparseBandedMatrix{typeof(uv[1])}(undef, M, N)
🦋!(A, Bu1, Uu, Ul)
Bu1 = SparseBandedMatrix{typeof(uv[1])}(undef, M, M)
🦋!(Bu1, Uu, Ul)

Bv2 = SparseBandedMatrix{typeof(uv[1])}(undef, M, N)
Bv2 = SparseBandedMatrix{typeof(uv[1])}(undef, M, M)

🦋2!(view(Bv2, 1 : Mh, 1 : Nh), V₁u, V₁l)
🦋2!(view(Bv2, Mh + 1: M, Nh + 1: N), V₂u, V₂l)
🦋2!(view(Bv2, 1 : Mh, 1 : Mh), V₁u, V₁l)
🦋2!(view(Bv2, Mh + 1: M, Mh + 1: M), V₂u, V₂l)

Bv1 = SparseBandedMatrix{typeof(uv[1])}(undef, M, N)
🦋!(A, Bv1, Vu, Vl)
Bv1 = SparseBandedMatrix{typeof(uv[1])}(undef, M, M)
🦋!(Bv1, Vu, Vl)

U = (Bu2 * Bu1)'
V = Bv2 * Bv1

U, V
mul!(U, Bu2, Bu1)
mul!(V, Bv2, Bv1)
end

function pad!(A)
Expand Down
15 changes: 4 additions & 11 deletions test/runtests.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -81,15 +81,8 @@ end
for i in 790 : 810
A = wilkinson(i)
b = rand(i)
U, V, F, out = RecursiveFactorization.🦋workspace(A, b, A, A', A, Val(true))
M = size(A, 1)
xn = 4 - M % 4
if (M % 4 != 0)
xn = 4 - M % 4
b = [b; rand(xn)]
end
sol = V * (F \ (U * b))
out .= @view sol[1:M]
@test norm(A * out .- b[1:M]) <= 1e-10
ws = RecursiveFactorization.🦋workspace(A, b)
out = RecursiveFactorization.🦋lu!(ws, i, Val(true))
@test norm(A * out .- b) <= 1e-10
end
end
end
Loading