add SparseBandedMatrices and associated implementation

Author: Shreyas-Ekanathan

Project.toml

@@ -8,6 +8,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890"
 Polyester = "f517fe37-dbe3-4b94-8317-1923a5111588"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 StrideArraysCore = "7792a7ef-975c-4747-a70f-980b88e8d1da"
 TriangularSolve = "d5829a12-d9aa-46ab-831f-fb7c9ab06edf"
 VectorizedRNG = "33b4df10-0173-11e9-2a0c-851a7edac40e"
@@ -17,6 +18,7 @@ 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"
+SparseArrays = "1.11.0"
 StrideArraysCore = "0.5.5"
 TriangularSolve = "0.2"
 VectorizedRNG = "0.2.25"

src/butterflylu.jl

@@ -2,6 +2,193 @@ using VectorizedRNG
 using LinearAlgebra: Diagonal, I
 using LoopVectorization
 using RecursiveFactorization
+using SparseArrays
+
+struct SparseBandedMatrix{T} <: AbstractMatrix{T}
+    size :: Tuple{Int, Int}
+    indices :: Vector{Int}
+    diags :: Vector{Vector{T}}
+    function SparseBandedMatrix{T}(::UndefInitializer, N, M) where T
+        size = (N, M)
+        indices = Int[]
+        diags = Vector{T}[]
+        new(size, indices, diags)
+    end
+    function SparseBandedMatrix{T}(ind_vals, diag_vals, N, M) where T
+        size = (N, M)
+        perm = sortperm(ind_vals)
+        indices = ind_vals[perm]
+        for i in 1 : length(indices) - 1
+            @assert indices[i] != indices[i + 1]
+        end
+        diags = diag_vals[perm]
+        new(size, indices, diags)
+    end
+end
+
+function Base.size(M :: SparseBandedMatrix) 
+    M.size
+end
+
+function Base.getindex(M :: SparseBandedMatrix{T}, i :: Int, j :: Int, I :: Int...) where T
+    @boundscheck checkbounds(M, i, j, I...)
+    rows, cols = size(M)
+    wanted_ind = rows - i + j
+    ind = searchsortedfirst(M.indices, wanted_ind)
+    if (ind <= length(M.indices) && M.indices[ind] == wanted_ind)
+        if (i > j)
+            return M.diags[ind][j]
+        else
+            return M.diags[ind][i]
+        end
+    end
+    zero(T)
+end
+
+function Base.setindex!(M :: SparseBandedMatrix{T}, val, i :: Int, j :: Int, I :: Int...) where T #TODO IF VAL ISNT OF TYPE T
+    @boundscheck checkbounds(M, i, j, I...) 
+    rows, cols = size(M)
+    wanted_ind = rows - i + j
+    ind = searchsortedfirst(M.indices, wanted_ind)
+    if (ind > length(M.indices) || M.indices[ind] != wanted_ind)
+        insert!(M.indices, ind, wanted_ind)
+        insert!(M.diags, ind, zeros(T, rows - abs(wanted_ind - rows)))
+    end
+    if (i > j)
+        M.diags[ind][j] = val isa T ? val : convert(T, val)::T
+    else
+        M.diags[ind][i] = val isa T ? val : convert(T, val)::T
+    end
+    val
+ end
+
+ function setdiagonal!(M :: SparseBandedMatrix{T}, diagvals, lower :: Bool) where T
+    rows, cols = size(M) 
+    if length(diagvals) > rows
+        error("size of diagonal is too big for the matrix")
+    end
+    if lower
+        wanted_ind = length(diagvals)
+    else
+        wanted_ind = 2 * rows - length(diagvals)
+    end
+
+    ind = searchsortedfirst(M.indices, wanted_ind)
+    if (ind > length(M.indices) || M.indices[ind] != wanted_ind)
+        insert!(M.indices, ind, wanted_ind)
+        insert!(M.diags, ind, diagvals isa Vector{T} ? diagvals : convert(Vector{T}, diagvals)::Vector{T}) 
+    else
+        for i in 1 : eachindex(diagvals)
+            M.diags[ind][i] = diagvals[i] isa T ? diagvals[i] : convert(T, diagvals[i])::T
+        end
+    end
+    diagvals
+end
+
+using LinearAlgebra
+using .Threads
+
+# C = Cb + aAB
+function LinearAlgebra.mul!(C :: Matrix{T}, A:: SparseBandedMatrix{T}, B :: Matrix{T}, a :: Number, b :: Number) where T
+    @assert size(A, 2) == size(B, 1)
+    @assert size(A, 1) == size(C, 1)
+    @assert size(B, 2) == size(C, 2)
+    C.*=b
+
+    rows, cols = size(A)
+    @inbounds for (ind, location) in enumerate(A.indices)
+        @threads for i in 1:length(A.diags[ind])
+            # value: diag[i]
+            # index in array: 
+            #       if ind < rows(A), then index = (rows - loc + i, i)
+            #       else index = (i, loc - cols + i)
+            val = A.diags[ind][i] * a
+            if location < rows 
+                index_i = rows - location + i 
+                index_j = i 
+            else
+                index_i = i 
+                index_j = location - cols + i 
+            end
+            #A[index_i, index_j] * B[index_j, j] = C[index_i, j]
+            @simd for j in 1 : size(B, 2)
+                C[index_i, j] = fma(val, B[index_j, j], C[index_i, j])
+            end
+        end
+    end
+    C
+end
+
+# C = Cb + aBA
+function LinearAlgebra.mul!(C :: Matrix{T}, A:: Matrix{T}, B :: SparseBandedMatrix{T}, a :: Number, b :: Number) where T
+    @assert size(A, 2) == size(B, 1)
+    @assert size(A, 1) == size(C, 1)
+    @assert size(B, 2) == size(C, 2)
+
+    C.*=b
+
+    rows, cols = size(B)
+    @inbounds for (ind, location) in enumerate(B.indices)
+        @threads for i in eachindex(B.diags[ind])
+                val = B.diags[ind][i] * a
+                if location < rows 
+                    index_i = rows - location + i 
+                    index_j = i 
+                else
+                    index_i = i 
+                    index_j = location - cols + i 
+                end
+            @simd for j in 1 : size(A, 1)
+                C[j, index_j] = fma(val, A[j, index_i], C[j, index_j])
+            end
+        end
+    end
+    C
+end
+
+function LinearAlgebra.mul!(C :: Matrix{T}, A:: SparseBandedMatrix{T}, B :: SparseBandedMatrix{T}, a :: Number, b :: Number) where T
+    @assert size(A, 2) == size(B, 1)
+    @assert size(A, 1) == size(C, 1)
+    @assert size(B, 2) == size(C, 2)
+
+    C.*=b
+
+    rows_a, cols_a = size(A)
+    rows_b, cols_b = size(B)
+    @inbounds for (ind_a, location_a) in enumerate(A.indices)
+        @threads for i in eachindex(A.diags[ind_a])
+            val_a = A.diags[ind_a][i] * a
+            if location_a < rows_a 
+                index_ia = rows_a - location_a + i 
+                index_ja = i 
+            else
+                index_ia = i 
+                index_ja = location_a - cols_a + i 
+            end
+            min_loc = rows_b - index_ja + 1
+            max_loc = 2 * rows_b - index_ja
+            for (ind_b, location_b) in enumerate(B.indices)
+                #index_ib = index_ja
+                #       if ind < rows(A), then index = (rows - loc + i, i)
+                #rows - loc + j = index_ja, j = index_ja - rows + loc
+                #       else index = (i, loc - cols + i)
+                # if location < rows(B), then 
+                if location_b <= rows_b && location_b >= min_loc
+                    j = index_ja - rows_b + location_b
+                    index_jb = j
+                    val_b = B.diags[ind_b][j]
+                    C[index_ia, index_jb] = muladd(val_a, val_b, C[index_ia, index_jb])         
+                elseif location_b > rows_b && location_b <= max_loc
+                    j = index_ja
+                    index_jb = location_b - cols_b + j 
+                    val_b = B.diags[ind_b][j]
+                    C[index_ia, index_jb] = muladd(val_a, val_b, C[index_ia, index_jb])         
+                end           
+            end
+        end
+    end
+    C
+end
 
 @inline exphalf(x) = exp(x) * oftype(x, 0.5)
 function 🦋!(wv, ::Val{SEED} = Val(888)) where {SEED}
@@ -32,7 +219,6 @@ end
 const butterfly_workspace = 🦋workspace;
 
 function 🦋mul_level!(A, u, v)
-    # for now, assume... 
     M, N = size(A)
     Ml = M >>> 1
     Nl = N >>> 1
@@ -97,8 +283,53 @@ function diagnegbottom(x)
     end
     Diagonal(y), Diagonal(z)
 end
-🦋(A, B) = [A B
-           A -B]
+
+#🦋(A, B) = [A B
+#            A -B]
+            
+    #Bu2 = [🦋(U₁u, U₁l) 0*I
+    #       0*I 🦋(U₂u, U₂l)]
+# U1u U1l 0 0
+# U1u -U1l 0 0
+#=
+function 🦋!(C, A, B)
+    A1, A2 = size(A)
+    B1, B2 = size(B)
+    @assert A1 == B1
+    for j in 1 : A2, i in 1 : A1
+        C[i, j] = A[i, j]
+        C[i + A1, j] = A[i, j]
+    end
+    for j in A2 + 1 : A2 + B2, i in 1 : A1
+        C[i, j] = B[i, j - A2]
+        C[i + A1, j] = -B[i, j - A2]
+    end
+    C
+end
+=#
+function 🦋!(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 🦋!(C::SparseBandedMatrix, A::Diagonal, B::Diagonal)
+    @assert size(A) == size(B)
+
+    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)
@@ -112,13 +343,31 @@ function materializeUV(A, (uv,))
     Uu, Ul = diagnegbottom(@view(uv[(1 + 2 * Mh + 2 * Nh):(2 * Mh + 2 * Nh + M)]))
     Vu, Vl = diagnegbottom(@view(uv[(1 + 2 * Mh + 2 * Nh + M):(2 * Mh + 2 * Nh + M + N)]))
 
-    Bu2 = [🦋(U₁u, U₁l) 0*I
-           0*I 🦋(U₂u, U₂l)]
-    Bu1 = 🦋(Uu, Ul)
+    #WRITE OUT MERGINGS EXPLICITLY
+    #Bu2 = [🦋(U₁u, U₁l) 0*I
+    #       0*I 🦋(U₂u, U₂l)]
+    #show size(Bu2)[1] #808
+    #@show size(🦋(V₁u, V₁l))[1] #404
+
+    #Bu2 = spzeros(M, N)
+    Bu2 = SparseBandedMatrix{typeof(uv[1])}(undef, M, N)
+    
+    🦋!(view(Bu2, 1 : M ÷ 2, 1 : N ÷ 2), U₁u, U₁l)
+    🦋!(view(Bu2, M ÷ 2 + 1 : M, N ÷ 2 + 1 : N), U₂u, U₂l)
+
+    #Bu1 = spzeros(M, N)
+    Bu1 = SparseBandedMatrix{typeof(uv[1])}(undef, M, N)
+    🦋!(Bu1, Uu, Ul)
+
+    #Bv2 = spzeros(M, N)
+    Bv2 = SparseBandedMatrix{typeof(uv[1])}(undef, M, N)
+
+    🦋!(view(Bv2, 1 : M ÷ 2, 1 : N ÷ 2), V₁u, V₁l)
+    🦋!(view(Bv2, M ÷ 2 + 1 : M, N ÷ 2 + 1 : N), V₂u, V₂l)
 
-    Bv2 = [🦋(V₁u, V₁l) 0*I
-           0*I 🦋(V₂u, V₂l)]
-    Bv1 = 🦋(Vu, Vl)
+    #Bv1 = spzeros(M, N)
+    Bv1 = SparseBandedMatrix{typeof(uv[1])}(undef, M, N)
+    🦋!(Bv1, Vu, Vl)
 
     (Bu2 * Bu1)', Bv2 * Bv1, A
 end
@@ -136,7 +385,7 @@ function pad!(A)
         @inbounds A_new[j, i] = 0
     end
 
-    for i in M + 1 : M + xn, j in N + 1 : N + xn
+    for j in N + 1 : N + xn, i in M + 1 : M + xn
         @inbounds A_new[i,j] = i == j
     end
     A_new