optimize for Symmetric,Hermitian,UpperTriangular as well as StridedArray

Author: joshday

Project.toml

@@ -1,6 +1,6 @@
 name = "SweepOperator"
 uuid = "7522ee7d-7047-56d0-94d9-4bc626e7058d"
-version = "0.3.2"
+version = "0.3.3"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

src/SweepOperator.jl

@@ -39,19 +39,47 @@ function sweep_with_buffer!(akk::AVec{T}, A::AMat{T}, k::Integer, inv::Bool = fa
     1 ≤ k ≤ p || throw(BoundsError(A, k))
     p == length(akk) || throw(DimensionError("Incorrect buffer size."))
     @inbounds @views begin
-        d = one(T) / A[k, k]                            # pivot
-        copy!(akk, Symmetric(A, :U)[:, k])            # akk = A[:, k]
-        if A isa StridedMatrix{<:Union{LinearAlgebra.BlasFloat, LinearAlgebra.BlasComplex}}
-            BLAS.syrk!('U', 'N', -d, akk, one(T), A)    # everything not in col/row k
-        else
-            A .+= UpperTriangular(-d * akk * akk')
-        end
-        rmul!(akk, d * (-one(T)) ^ inv)                 # akk .* d (negated if inv=true)
-        copy!(A[1:k-1,k], akk[1:k-1])                 # col k
-        copy!(A[k, k+1:end], akk[k+1:end])            # row k
-        A[k, k] = -d                                    # pivot element
+        d = one(T) / A[k, k]                # pivot
+        copy!(akk, Symmetric(A, :U)[:, k])  # akk = A[:, k]
+        syrk!(A, -d, akk)                   # everything not in row/col k
+        rmul!(akk, d * (-one(T)) ^ inv)     # akk .* d (negated if inv=true)
+        setrowcol!(A, k, akk)
+        A[k, k] = -d                        # pivot element
     end
     return A
 end
 
+#-----------------------------------------------------------------------------# setrowcol!
+# Set upper triangle of: (A[k, :] = x; A[:, k] = x)
+function setrowcol!(A::StridedArray, k, x)
+    @views copy!(A[1:k-1,k], x[1:k-1])         # col k
+    @views copy!(A[k, k+1:end], x[k+1:end])    # row k
+end
+
+setrowcol!(A::Union{Hermitian,Symmetric,UpperTriangular}, k, x) = setrowcol!(A.data, k, x)
+
+#-----------------------------------------------------------------------------# syrk!
+const BlasNumber = Union{LinearAlgebra.BlasFloat, LinearAlgebra.BlasComplex}
+
+# In-place update of (the upper triangle of) A + α * x * x'
+function syrk!(A::StridedMatrix{T}, α::T, x::AbstractArray{<:T}) where {T<:BlasNumber}
+    BLAS.syrk!('U', 'N', α, x, one(T), A)
+end
+
+function syrk!(A::Hermitian{T, S}, α::T, x::AbstractArray{<:T}) where {T<:BlasNumber, S<:StridedMatrix{T}}
+    Hermitian(BLAS.syrk!('U', 'N', α, x, one(T), A.data))
+end
+
+function syrk!(A::Symmetric{T, S}, α::T, x::AbstractArray{<:T}) where {T<:BlasNumber, S<:StridedMatrix{T}}
+    Symmetric(BLAS.syrk!('U', 'N', α, x, one(T), A.data))
+end
+
+function syrk!(A, α, x) where {T}
+    p = checksquare(A)
+    for i in 1:p, j in i:p
+        @inbounds A[i,j] += α * x[i] * x[j]
+    end
+end
+
+
 end # module

test/runtests.jl

@@ -6,41 +6,51 @@ x = randn(n, p)
 xtx = x'x
 
 @testset "Sweep One By One" begin
-    A = deepcopy(xtx)
-    B = deepcopy(xtx)
+    A = copy(xtx)
+    B = copy(xtx)
     for j in 1:p
         sweep!(A, j)
         sweep!(A, j, true)
     end
-    @test A ≈ B
+    @test UpperTriangular(A) ≈ UpperTriangular(B)
 
-    A = deepcopy(xtx)
-    B = deepcopy(xtx)
+    A = copy(xtx)
+    B = copy(xtx)
     for j in 1:p
-        sweep!(A, j, true)
+        sweep!(A, j)
     end
     for j in 1:p
         sweep!(A, j, true)
     end
-    @test A ≈ B
+    @test UpperTriangular(A) ≈ UpperTriangular(B)
 end
 
 @testset "Sweep All" begin
-    A = deepcopy(xtx)
-    B = deepcopy(xtx)
+    A = copy(xtx)
+    B = copy(xtx)
     sweep!(A, 1:p)
     sweep!(A, 1:p, true)
     @test A ≈ B
 end
 
-@testset "Non-StridedArray" begin
-    A = Diagonal(deepcopy(xtx))
-    B = Diagonal(deepcopy(xtx))
+@testset "UpperTriangular" begin
+    A = UpperTriangular(copy(xtx))
+    B = UpperTriangular(copy(xtx))
     sweep!(A, 1:p)
     sweep!(A, 1:p, true)
     @test A ≈ B
 end
 
+@testset "Hermitian/Symmetric" begin
+    A = Hermitian(copy(xtx))
+    B = Symmetric(copy(xtx))
+    sweep!(A, 1:p)
+    sweep!(A, 1:p, true)
+    sweep!(B, 1:p)
+    sweep!(B, 1:p, true)
+    @test A ≈ B ≈ xtx
+end
+
 @testset "Linear Regression" begin
     y = x * collect(1.:p) + randn(n)
     xy = [x y]