add specialized methods for Triangular*Triangular

Author: kleinschmidt

src/triangular.jl

@@ -17,9 +17,12 @@
 @inline Base.:*(A::LinearAlgebra.AbstractTriangular{<:Any,<:StaticMatrix}, B::Transpose{<:Any,<:StaticVecOrMat}) =
     transpose(transpose(B) * transpose(A))
 
+const StaticULT = Union{UpperTriangular{<:Any,<:StaticMatrix},LowerTriangular{<:Any,<:StaticMatrix}}
+
 @inline Base.:*(A::LinearAlgebra.AbstractTriangular{<:Any,<:StaticMatrix}, B::StaticVecOrMat) = _A_mul_B(Size(A), Size(B), A, B)
 @inline Base.:*(A::StaticVecOrMat, B::LinearAlgebra.AbstractTriangular{<:Any,<:StaticMatrix}) = _A_mul_B(Size(A), Size(B), A, B)
-@inline Base.:\(A::Union{UpperTriangular{<:Any,<:StaticMatrix},LowerTriangular{<:Any,<:StaticMatrix}}, B::StaticVecOrMat) = _A_ldiv_B(Size(A), Size(B), A, B)
+@inline Base.:*(A::StaticULT, B::StaticULT) = _A_mul_B(Size(A), Size(B), A, B)
+@inline Base.:\(A::StaticULT, B::StaticVecOrMat) = _A_ldiv_B(Size(A), Size(B), A, B)
 
 
 @generated function _A_mul_B(::Size{sa}, ::Size{sb}, A::UpperTriangular{TA,<:StaticMatrix}, B::StaticVecOrMat{TB}) where {sa,sb,TA,TB}
@@ -559,3 +562,129 @@ end
         @inbounds return similar_type(B, TAB)(tuple($(X...)))
     end
 end
+
+@generated function _A_mul_B(::Size{sa}, ::Size{sb}, A::UpperTriangular{<:TA,<:StaticMatrix}, B::UpperTriangular{<:TB,<:StaticMatrix}) where {sa,sb,TA,TB}
+    n = sa[1]
+    if n != sb[1]
+        throw(DimensionMismatch("left and right-hand must have same sizes, got $(n) and $(sb[1])"))
+    end
+
+    X = [Symbol("X_$(i)_$(j)") for i = 1:n, j = 1:n]
+
+    TAB = promote_op(*, eltype(TA), eltype(TB))
+    z = zero(TAB)
+
+    code = quote end
+    for j = 1:n
+        for i = 1:n
+            if i > j
+                push!(code.args, :($(X[i,j]) = $z))
+            else
+                ex = :(A.data[$(LinearIndices(sa)[i,i])] * B.data[$(LinearIndices(sb)[i,j])])
+                for k = i+1:j
+                    ex = :($ex + A.data[$(LinearIndices(sa)[i,k])] * B.data[$(LinearIndices(sb)[k,j])])
+                end
+                push!(code.args, :($(X[i,j]) = $ex))
+            end
+        end
+    end
+
+    return quote
+        @_inline_meta
+        @inbounds $code
+        return UpperTriangular(similar_type(B.data, $TAB)(tuple($(X...))))
+    end
+
+end
+
+@generated function _A_mul_B(::Size{sa}, ::Size{sb}, A::LowerTriangular{<:TA,<:StaticMatrix}, B::LowerTriangular{<:TB,<:StaticMatrix}) where {sa,sb,TA,TB}
+    n = sa[1]
+    if n != sb[1]
+        throw(DimensionMismatch("left and right-hand must have same sizes, got $(n) and $(sb[1])"))
+    end
+
+    X = [Symbol("X_$(i)_$(j)") for i = 1:n, j = 1:n]
+
+    TAB = promote_op(*, eltype(TA), eltype(TB))
+    z = zero(TAB)
+
+    code = quote end
+    for j = 1:n
+        for i = 1:n
+            if i < j
+                push!(code.args, :($(X[i,j]) = $z))
+            else
+                ex = :(A.data[$(LinearIndices(sa)[i,j])] * B.data[$(LinearIndices(sb)[j,j])])
+                for k = j+1:i
+                    ex = :($ex + A.data[$(LinearIndices(sa)[i,k])] * B.data[$(LinearIndices(sb)[k,j])])
+                end
+                push!(code.args, :($(X[i,j]) = $ex))
+            end
+        end
+    end
+
+    return quote
+        @_inline_meta
+        @inbounds $code
+        return LowerTriangular(similar_type(B.data, $TAB)(tuple($(X...))))
+    end
+
+end
+
+
+@generated function _A_mul_B(::Size{sa}, ::Size{sb}, A::UpperTriangular{<:TA,<:StaticMatrix}, B::LowerTriangular{<:TB,<:StaticMatrix}) where {sa,sb,TA,TB}
+    n = sa[1]
+    if n != sb[1]
+        throw(DimensionMismatch("left and right-hand must have same sizes, got $(n) and $(sb[1])"))
+    end
+
+    X = [Symbol("X_$(i)_$(j)") for i = 1:n, j = 1:n]
+
+    code = quote end
+    for j = 1:n
+        for i = 1:n
+            k1 = max(i,j)
+            ex = :(A.data[$(LinearIndices(sa)[i,k1])] * B.data[$(LinearIndices(sb)[k1,j])])
+            for k = k1+1:n
+                ex = :($ex + A.data[$(LinearIndices(sa)[i,k])] * B.data[$(LinearIndices(sb)[k,j])])
+            end
+            push!(code.args, :($(X[i,j]) = $ex))
+        end
+    end
+
+    return quote
+        @_inline_meta
+        @inbounds $code
+        TAB = promote_op(*, eltype(TA), eltype(TB))
+        return similar_type(B.data, TAB)(tuple($(X...)))
+    end
+
+end
+
+@generated function _A_mul_B(::Size{sa}, ::Size{sb}, A::LowerTriangular{<:TA,<:StaticMatrix}, B::UpperTriangular{<:TB,<:StaticMatrix}) where {sa,sb,TA,TB}
+    n = sa[1]
+    if n != sb[1]
+        throw(DimensionMismatch("left and right-hand must have same sizes, got $(n) and $(sb[1])"))
+    end
+
+    X = [Symbol("X_$(i)_$(j)") for i = 1:n, j = 1:n]
+
+    code = quote end
+    for j = 1:n
+        for i = 1:n
+            ex = :(A.data[$(LinearIndices(sa)[i,1])] * B.data[$(LinearIndices(sb)[1,j])])
+            for k = 2:min(i,j)
+                ex = :($ex + A.data[$(LinearIndices(sa)[i,k])] * B.data[$(LinearIndices(sb)[k,j])])
+            end
+            push!(code.args, :($(X[i,j]) = $ex))
+        end
+    end
+
+    return quote
+        @_inline_meta
+        @inbounds $code
+        TAB = promote_op(*, eltype(TA), eltype(TB))
+        return similar_type(B.data, TAB)(tuple($(X...)))
+    end
+
+end