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
merged 4 commits into from
Oct 9, 2025
Merged

Clean up ButterflyFactorization #103

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
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
172 changes: 83 additions & 89 deletions src/butterflylu.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,49 +5,62 @@ 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}
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
U, V = (similar(A), similar(A))
ws = 🦋generate_random!(A)
materializeUV(U, V, ws)
new{eltype(A)}(A, 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)

U, V, F, out
end

function 🦋lu!(workspace::🦋workspace, M, thread)
Copy link

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, ws, U, V, out) = workspace
🦋mul!(A, ws)
F = RecursiveFactorization.lu!(A, Val(false), thread)
sol = V * (F \ (U' * b))
out .= @view sol[1:M]
out
end

const butterfly_workspace = 🦋workspace;

function 🦋mul_level!(A, u, v)
M, N = size(A)
@assert M == length(u) && N == length(v)
Mh = M >>> 1
Nh = N >>> 1
@turbo for n in 1 : Nh
for m in 1 : Mh
M_half = M >>> 1
N_half = N >>> 1
@turbo for n in 1 : N_half
for m in 1 : M_half
A11 = A[m, n]
A21 = A[m + Mh, n]
A12 = A[m, n + Nh]
A22 = A[m + Mh, n + Nh]
A21 = A[m + M_half, n]
A12 = A[m, n + N_half]
A22 = A[m + M_half, n + N_half]

T1 = A11 + A12
T2 = A21 + A22
Expand All @@ -59,32 +72,32 @@ function 🦋mul_level!(A, u, v)
C22 = T3 - T4

u1 = u[m]
u2 = u[m + Mh]
u2 = u[m + M_half]
v1 = v[n]
v2 = v[n + Nh]
v2 = v[n + N_half]

A[m, n] = u1 * C11 * v1
A[m + Mh, n] = u2 * C21 * v1
A[m, n + Nh] = u1 * C12 * v2
A[m + Mh, n + Nh] = u2 * C22 * v2
A[m + M_half, n] = u2 * C21 * v1
A[m, n + N_half] = u1 * C12 * v2
A[m + M_half, n + N_half] = u2 * C22 * v2
end
end
end

function 🦋mul!(A, (uv,))
function 🦋mul!(A, uv)
M, N = size(A)
@assert M == N
Mh = M >>> 1
M_half = M >>> 1

U₁ = @view(uv[1:Mh])
V₁ = @view(uv[(Mh + 1):(M)])
U₂ = @view(uv[(1 + M):(M + Mh)])
V₂ = @view(uv[(1 + M + Mh):(2 * M)])
U₁ = @view(uv[1:M_half])
V₁ = @view(uv[(M_half + 1):(M)])
U₂ = @view(uv[(1 + M):(M + M_half)])
V₂ = @view(uv[(1 + M + M_half):(2 * M)])

🦋mul_level!(@view(A[1:Mh, 1:Mh]), U₁, V₁)
🦋mul_level!(@view(A[Mh + 1:M, 1:Mh]), U₂, V₁)
🦋mul_level!(@view(A[1:Mh, Mh + 1:M]), U₁, V₂)
🦋mul_level!(@view(A[Mh + 1:M, Mh + 1:M]), U₂, V₂)
🦋mul_level!(@view(A[1:M_half, 1:M_half]), U₁, V₁)
🦋mul_level!(@view(A[M_half + 1:M, 1:M_half]), U₂, V₁)
🦋mul_level!(@view(A[1:M_half, M_half + 1:M]), U₁, V₂)
🦋mul_level!(@view(A[M_half + 1:M, M_half + 1:M]), U₂, V₂)

U = @view(uv[(1 + 2 * M):(3 * M)])
V = @view(uv[(1 + 3 * M):(4 * M)])
Expand All @@ -106,7 +119,14 @@ function diagnegbottom(x)
Diagonal(y), Diagonal(z)
end

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

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
Contributor Author

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 🦋!(C, A::Diagonal, B::Diagonal)
@assert size(A) == size(B)
A1 = size(A, 1)

Expand All @@ -120,61 +140,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)
M_half = M >>> 1

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:M_half])) #M_half
U₂u, U₂l = diagnegbottom(@view(uv[(1 + 2 * M_half):(M + M_half)])) #M_half
V₁u, V₁l = diagnegbottom(@view(uv[(M_half + 1):(2 * M_half)])) #M_half
V₂u, V₂l = diagnegbottom(@view(uv[(1 + 3 * M_half):(2 * M_half + M)])) #M_half
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)
🦋!(view(Bu2, 1 : M_half, 1 : M_half), U₁u, U₁l)
🦋!(view(Bu2, M_half + 1: M, M_half + 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)
🦋!(view(Bv2, 1 : M_half, 1 : M_half), V₁u, V₁l)
🦋!(view(Bv2, M_half + 1: M, M_half + 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
16 changes: 5 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,9 @@ 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(copy(A), copy(b))
out = RecursiveFactorization.🦋lu!(ws, i, Val(true))
@test norm(A * out .- b) <= 1e-10
end
end
end

Loading