adding Diagonal to the new matrix multiplication scheme

Author: mateuszbaran

benchmark/bench_mat_mul.jl

@@ -14,7 +14,8 @@ mul_wrappers = [
     (m -> UnitUpperTriangular(m), "uup-tri"),
     (m -> UnitLowerTriangular(m), "ulo-tri"),
     (m -> Adjoint(m), "adjoint"),
-    (m -> Transpose(m), "transpo")]
+    (m -> Transpose(m), "transpo"),
+    (m -> Diagonal(m), "diag   ")]
 
 for N in [2, 4, 8, 10, 16]
 

src/SDiagonal.jl

@@ -18,8 +18,6 @@ size(::Type{<:SDiagonal{N}}) where {N} = (N,N)
 size(::Type{<:SDiagonal{N}}, d::Int) where {N} = d > 2 ? 1 : N
 
 # define specific methods to avoid allocating mutable arrays
-*(A::StaticMatrix, D::SDiagonal) = A .* transpose(D.diag)
-*(D::SDiagonal, A::StaticMatrix) = D.diag .* A
 \(D::SDiagonal, b::AbstractVector) = D.diag .\ b
 \(D::SDiagonal, b::StaticVector) = D.diag .\ b # catch ambiguity
 

src/matrix_multiply.jl

@@ -66,6 +66,9 @@ end
 function gen_by_access(expr_gen, a::Type{<:Adjoint{<:Any, <:StaticVecOrMat}}, asym = :wrapped_a)
     return expr_gen(:adjoint)
 end
+function gen_by_access(expr_gen, a::Type{<:SDiagonal}, asym = :wrapped_a)
+    return expr_gen(:diagonal)
+end
 """
     gen_by_access(expr_gen, a::Type{<:AbstractArray}, b::Type{<:AbstractArray})
 
@@ -148,6 +151,13 @@ function gen_by_access(expr_gen, a::Type{<:Adjoint{<:Any, <:StaticMatrix}}, b::T
         end)
     end
 end
+function gen_by_access(expr_gen, a::Type{<:SDiagonal}, b::Type)
+    return quote
+        return $(gen_by_access(b, :wrapped_b) do access_b
+            expr_gen(:diagonal, access_b)
+        end)
+    end
+end
 
 """
     mul_result_structure(a::Type, b::Type)
@@ -164,6 +174,21 @@ end
 function mul_result_structure(::LowerTriangular{<:Any, <:StaticMatrix}, ::LowerTriangular{<:Any, <:StaticMatrix})
     return LowerTriangular
 end
+function mul_result_structure(::UpperTriangular{<:Any, <:StaticMatrix}, ::SDiagonal)
+    return UpperTriangular
+end
+function mul_result_structure(::LowerTriangular{<:Any, <:StaticMatrix}, ::SDiagonal)
+    return LowerTriangular
+end
+function mul_result_structure(::SDiagonal, ::UpperTriangular{<:Any, <:StaticMatrix})
+    return UpperTriangular
+end
+function mul_result_structure(::SDiagonal, ::LowerTriangular{<:Any, <:StaticMatrix})
+    return LowerTriangular
+end
+function mul_result_structure(::SDiagonal, ::SDiagonal)
+    return Diagonal
+end
 
 """
     uplo_access(sa, asym, k, j, uplo)
@@ -247,6 +272,12 @@ function uplo_access(sa, asym, k, j, uplo)
         return :(transpose($asym[$(LinearIndices(reverse(sa))[j, k])]))
     elseif uplo == :adjoint
         return :(adjoint($asym[$(LinearIndices(reverse(sa))[j, k])]))
+    elseif uplo == :diagonal
+        if k == j
+            return :($asym[$k])
+        else
+            return :(zero($TAsym))
+        end
     else
         error("Unknown uplo: $uplo")
     end
@@ -347,12 +378,12 @@ end
 
 @generated function _mul(Sa::Size{sa}, Sb::Size{sb}, a::StaticMatMulLike{<:Any, <:Any, Ta}, b::StaticMatMulLike{<:Any, <:Any, Tb}) where {sa, sb, Ta, Tb}
     # Heuristic choice for amount of codegen
-    if sa[1]*sa[2]*sb[2] <= 8*8*8 || !(a <: StaticMatrix) || !(b <: StaticMatrix)
+    if sa[1]*sa[2]*sb[2] <= 8*8*8 || a <: Diagonal || b <: Diagonal
         return quote
             @_inline_meta
             return mul_unrolled(Sa, Sb, a, b)
         end
-    elseif sa[1] <= 14 && sa[2] <= 14 && sb[2] <= 14
+    elseif (sa[1] <= 14 && sa[2] <= 14 && sb[2] <= 14) || !(a <: StaticMatrix) || !(b <: StaticMatrix)
         return quote
             @_inline_meta
             return mul_unrolled_chunks(Sa, Sb, a, b)
@@ -436,7 +467,7 @@ end
     tmp_type_out = :(SVector{$(sa[1]), T})
 
     retexpr = gen_by_access(wrapped_a, wrapped_b) do access_a, access_b
-        vect_exprs = [:($(Symbol("tmp_$k2"))::$tmp_type_out = partly_unrolled_multiply($(Size{sa}()), $(Size{(sb[1],)}()),
+        vect_exprs = [:($(Symbol("tmp_$k2")) = partly_unrolled_multiply($(Size{sa}()), $(Size{(sb[1],)}()),
             a, $(Expr(:call, tmp_type_in, [uplo_access(sb, :b, i, k2, access_b) for i = 1:sb[1]]...)), $(Val(access_a)))::$tmp_type_out) for k2 = 1:sb[2]]
 
         exprs = [:($(Symbol("tmp_$k2"))[$k1]) for k1 = 1:sa[1], k2 = 1:sb[2]]

src/matrix_multiply_add.jl

@@ -27,7 +27,8 @@ const StaticMatMulLike{s1, s2, T} = Union{
     UnitLowerTriangular{T, <:StaticMatrix{s1, s2, T}},
     UnitUpperTriangular{T, <:StaticMatrix{s1, s2, T}},
     Adjoint{T, <:StaticMatrix{s1, s2, T}},
-    Transpose{T, <:StaticMatrix{s1, s2, T}}}
+    Transpose{T, <:StaticMatrix{s1, s2, T}},
+    SDiagonal{s1, T}}
 
 
 """ Size that stores whether a Matrix is a Transpose
@@ -188,7 +189,7 @@ end
     can_blas = Tc == Ta && Tc == Tb && Tc <: BlasFloat && a <: Union{StaticMatrix,Transpose} && b <: Union{StaticMatrix,Transpose}
 
     mult_dim = multiplied_dimension(a,b)
-    if mult_dim < 4*4*4
+    if mult_dim < 4*4*4 || a <: Diagonal || b <: Diagonal
         return quote
             @_inline_meta
             muladd_unrolled_all!(Sc, c, Sa, Sb, a, b, _add)

test/matrix_multiply.jl

@@ -11,7 +11,8 @@ mul_wrappers = [
     m -> UnitUpperTriangular(m),
     m -> UnitLowerTriangular(m),
     m -> Adjoint(m),
-    m -> Transpose(m)]
+    m -> Transpose(m),
+    m -> Diagonal(m)]
 
 @testset "Matrix multiplication" begin
     @testset "Matrix-vector" begin
@@ -172,6 +173,8 @@ mul_wrappers = [
                     LowerTriangular{Int,typeof(mm)}
                 elseif res_structure == UpperTriangular
                     UpperTriangular{Int,typeof(mm)}
+                elseif res_structure == Diagonal
+                    Diagonal{Int,<:SVector}
                 else
                     error("Unknown structure: ", res_structure)
                 end

test/matrix_multiply_add.jl

@@ -22,8 +22,9 @@ mul_add_wrappers = [
     m -> LowerTriangular(m),
     m -> UnitUpperTriangular(m),
     m -> UnitLowerTriangular(m),
-    m -> adjoint(m),
-    m -> transpose(m)]
+    m -> Adjoint(m),
+    m -> Transpose(m),
+    m -> Diagonal(m)]
 
 
 # check_dims
@@ -224,15 +225,15 @@ function test_wrappers_for_size(N, test_block)
         # matrix-vector
         for wrapper in mul_add_wrappers
             # LinearAlgebra can't handle these
-            if all(T -> !isa(wrapper([1 2; 3 4]), T), [Symmetric, Hermitian])
+            if all(T -> !isa(wrapper([1 2; 3 4]), T), [Symmetric, Hermitian, Diagonal])
                 mul!(Cv_block, wrapper(A_block), bv_block)
                 @test Cv_block == wrapper(Array(A_block))*Array(bv_block)
             end
         end
 
         # matrix-matrix
         for wrapper_a in mul_add_wrappers, wrapper_b in mul_add_wrappers
-            if all(T -> !isa(wrapper_a([1 2; 3 4]), T) && !isa(wrapper_b([1 2; 3 4]), T), [Symmetric, Hermitian])
+            if all(T -> !isa(wrapper_a([1 2; 3 4]), T) && !isa(wrapper_b([1 2; 3 4]), T), [Symmetric, Hermitian, Diagonal])
                 mul!(C_block, wrapper_a(A_block), wrapper_b(B_block))
                 @test C_block == wrapper_a(Array(A_block))*wrapper_b(Array(B_block))
             end