Add Butterfly Factorization

.github/workflows/ci.yml

@@ -12,14 +12,14 @@ jobs:
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
-        julia-version: ['1']
+        julia-version: ['lts','1','pre']
         threads:
           - '1'
           - '3'
         os: [ubuntu-latest, windows-latest, macOS-latest]
     steps:
       - uses: actions/checkout@v4
-      - uses: julia-actions/setup-julia@v1
+      - uses: julia-actions/setup-julia@v2
         with:
           version: ${{ matrix.julia-version }}
       - uses: actions/cache@v4

Project.toml

@@ -8,16 +8,20 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890"
 Polyester = "f517fe37-dbe3-4b94-8317-1923a5111588"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
+SparseBandedMatrices = "bd59d7e1-4699-4102-944e-d05209cb92aa"
 StrideArraysCore = "7792a7ef-975c-4747-a70f-980b88e8d1da"
 TriangularSolve = "d5829a12-d9aa-46ab-831f-fb7c9ab06edf"
+VectorizedRNG = "33b4df10-0173-11e9-2a0c-851a7edac40e"
 
 [compat]
 LinearAlgebra = "1.5"
 LoopVectorization = "0.10,0.11, 0.12"
 Polyester = "0.3.2,0.4.1, 0.5, 0.6, 0.7"
 PrecompileTools = "1"
+SparseBandedMatrices = "0.1.0"
 StrideArraysCore = "0.5.5"
 TriangularSolve = "0.2"
+VectorizedRNG = "0.2.25"
 julia = "1.10"
 
 [extras]

perf/lu.jl

@@ -78,4 +78,4 @@ xaxis!(plt, "size (N x N)")
 yaxis!(plt, "GFLOPS")
 savefig("lubench.png")
 savefig("lubench.pdf")
-=#
+=#

src/RecursiveFactorization.jl

@@ -4,6 +4,7 @@ if isdefined(Base, :Experimental) &&
     @eval Base.Experimental.@max_methods 1
 end
 include("./lu.jl")
+include("./butterflylu.jl")
 
 import PrecompileTools
 

src/butterflylu.jl

@@ -0,0 +1,197 @@
+using VectorizedRNG
+using LinearAlgebra: Diagonal, I
+using LoopVectorization
+using RecursiveFactorization
+using SparseBandedMatrices
+
+@inline exphalf(x) = exp(x) * oftype(x, 0.5)
+function 🦋!(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,)
+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)
+    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
+
+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
+            A11 = A[m, n]
+            A21 = A[m + Mh, n]
+            A12 = A[m, n + Nh]
+            A22 = A[m + Mh, n + Nh]
+
+            T1 = A11 + A12
+            T2 = A21 + A22
+            T3 = A11 - A12
+            T4 = A21 - A22
+            C11 = T1 + T2
+            C21 = T1 - T2
+            C12 = T3 + T4
+            C22 = T3 - T4
+
+            u1 = u[m]
+            u2 = u[m + Mh]
+            v1 = v[n]
+            v2 = v[n + Nh]
+
+            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
+        end
+    end 
+end
+
+function 🦋mul!(A, (uv,))
+    M, N = size(A)
+    @assert M == N
+    Mh = 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)]) 
+
+    🦋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₂) 
+
+    U = @view(uv[(1 + 2 * M):(3 * M)]) 
+    V = @view(uv[(1 + 3 * M):(4 * M)]) 
+
+    🦋mul_level!(@view(A[1:M, 1:N]), U, V)
+    A
+end
+
+function diagnegbottom(x)
+    N = length(x)
+    y = similar(x, N >>> 1)
+    z = similar(x, N >>> 1)
+    for n in 1:(N >>> 1)
+        y[n] = x[n]
+    end
+    for n in 1:(N >>> 1)
+        z[n] = x[n + (N >>> 1)]
+    end
+    Diagonal(y), Diagonal(z)
+end
+
+function 🦋2!(C, A::Diagonal, B::Diagonal)
+    @assert size(A) == size(B)
+    A1 = size(A, 1)
+
+    for i in 1:A1
+        C[i, i] = A[i, i]
+        C[i + A1, i] = A[i, i]
+        C[i, i + A1] = B[i, i]
+        C[i + A1, i + A1] = -B[i, i]
+    end
+
+    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
+
+    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
+
+    Bu2 = SparseBandedMatrix{typeof(uv[1])}(undef, M, N)
+
+    🦋2!(view(Bu2, 1 : Mh, 1 : Nh), U₁u, U₁l)
+    🦋2!(view(Bu2, Mh + 1: M, Nh + 1: N), U₂u, U₂l)
+
+    Bu1 = SparseBandedMatrix{typeof(uv[1])}(undef, M, N)
+    🦋!(A, Bu1, Uu, Ul)
+
+    Bv2 = SparseBandedMatrix{typeof(uv[1])}(undef, M, N)
+
+    🦋2!(view(Bv2, 1 : Mh, 1 : Nh), V₁u, V₁l)
+    🦋2!(view(Bv2, Mh + 1: M, Nh + 1: N), V₂u, V₂l)
+
+    Bv1 = SparseBandedMatrix{typeof(uv[1])}(undef, M, N)
+    🦋!(A, Bv1, Vu, Vl)
+
+    U = (Bu2 * Bu1)'
+    V = Bv2 * Bv1
+
+    U, V
+end
+
+function pad!(A)
+    M, N = size(A)
+    xn = 4 - M % 4
+    A_new = similar(A, M + xn, N + xn)
+    for j in 1 : N, i in 1 : M
+        @inbounds A_new[i, j] = A[i, j]
+    end
+
+    for j in M + 1 : M + xn, i in 1:M
+        @inbounds A_new[i, j] = 0
+        @inbounds A_new[j, i] = 0
+    end
+
+    for j in N + 1 : N + xn, i in M + 1 : M + xn
+        @inbounds A_new[i,j] = i == j
+    end
+    A_new
+end

src/lu.jl

@@ -21,6 +21,8 @@ function lu(A::AbstractMatrix, pivot = Val(true), thread = Val(false); kwargs...
 end
 
 const CUSTOMIZABLE_PIVOT = VERSION >= v"1.8.0-DEV.1507"
+# Julia 1.11+ uses negative info for NoPivot() failures
+const NOPIVOT_NEGATIVE_INFO = VERSION >= v"1.11.0-DEV"
 
 struct NotIPIV <: AbstractVector{BlasInt}
     len::Int
@@ -235,7 +237,14 @@ function reckernel!(A::AbstractMatrix{T}, pivot::Val{Pivot}, m, n, ipiv, info, b
         # A21 <- P2 A21
         Pivot && apply_permutation!(P2, A21, thread)
 
-        info != previnfo && (info += n1)
+        if info != previnfo
+            # Handle negative info for NoPivot (Julia 1.11+ convention)
+            if NOPIVOT_NEGATIVE_INFO && info < 0
+                info -= n1
+            else
+                info += n1
+            end
+        end
         if Pivot
             @turbo warn_check_args=false for i in 1:n2
                 P2[i] += n1
@@ -303,6 +312,10 @@ function _generic_lufact!(A, ::Val{Pivot}, ipiv, info) where {Pivot}
                 end
             elseif info == 0
                 info = k
+                # Julia 1.11+ convention: negative info for NoPivot
+                if !Pivot && NOPIVOT_NEGATIVE_INFO
+                    info = -info
+                end
             end
             k == minmn && break
             # Update the rest

test/runtests.jl

@@ -64,3 +64,32 @@ testlu(A::Union{Transpose, Adjoint}, MF, BF, p) = testlu(parent(A), parent(MF),
         end
     end
 end
+
+function wilkinson(N)
+    A = zeros(N, N)
+    A[1:(N+1):N*N] .= 1
+    A[:, end] .= 1
+    for n in 1:(N - 1)
+        for r in (n + 1):N
+            @inbounds A[r, n] = -1
+        end
+    end
+    A
+end
+
+@testset "🦋" begin
+    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
+    end
+end