adjoint and transpose wrappers for multiplication; more documentation

Author: mateuszbaran

src/matrix_multiply.jl

@@ -4,12 +4,20 @@ import LinearAlgebra: BlasFloat, matprod, mul!
 # Manage dispatch of * and mul!
 # TODO Adjoint? (Inner product?)
 
+"""
+    StaticMatMulLike
+
+Static wrappers used for multiplication dispatch.
+"""
 const StaticMatMulLike{s1, s2, T} = Union{
     StaticMatrix{s1, s2, T},
     Symmetric{T, <:StaticMatrix{s1, s2, T}},
     Hermitian{T, <:StaticMatrix{s1, s2, T}},
     LowerTriangular{T, <:StaticMatrix{s1, s2, T}},
-    UpperTriangular{T, <:StaticMatrix{s1, s2, T}}}
+    UpperTriangular{T, <:StaticMatrix{s1, s2, T}},
+    Adjoint{T, <:StaticMatrix{s1, s2, T}},
+    Transpose{T, <:StaticMatrix{s1, s2, T}}}
+
 
 @inline *(A::StaticMatMulLike, B::AbstractVector) = _mul(Size(A), A, B)
 @inline *(A::StaticMatMulLike, B::StaticVector) = _mul(Size(A), Size(B), A, B)
@@ -20,6 +28,18 @@ const StaticMatMulLike{s1, s2, T} = Union{
 @inline *(A::StaticArray{Tuple{N,1},<:Any,2}, B::Adjoint{<:Any,<:StaticVector}) where {N} = vec(A) * B
 @inline *(A::StaticArray{Tuple{N,1},<:Any,2}, B::Transpose{<:Any,<:StaticVector}) where {N} = vec(A) * B
 
+"""
+    gen_by_access(expr_gen, a::Type{<:AbstractArray}, asym = :a)
+
+Statically generate outer code for fully unrolled multiplication loops.
+Returned code does wrapper-specific tests (for example if a symmetric matrix view is
+`U` or `L`) and the body of the if expression is then generated by function `expr_gen`.
+The function `expr_gen` receives access pattern description symbol as its argument
+and this symbol is then consumed by uplo_access to generate the right code for matrix
+element access.
+
+The name of the matrix to test is indicated by `asym`.
+"""
 function gen_by_access(expr_gen, a::Type{<:StaticMatrix}, asym = :a)
     return expr_gen(:any)
 end
@@ -47,6 +67,19 @@ end
 function gen_by_access(expr_gen, a::Type{<:LowerTriangular{<:Any, <:StaticMatrix}}, asym = :a)
     return expr_gen(:lower_triangular)
 end
+function gen_by_access(expr_gen, a::Type{<:Transpose{<:Any, <:StaticMatrix}}, asym = :a)
+    return expr_gen(:transpose)
+end
+function gen_by_access(expr_gen, a::Type{<:Adjoint{<:Any, <:StaticMatrix}}, asym = :a)
+    return expr_gen(:adjoint)
+end
+"""
+    gen_by_access(expr_gen, a::Type{<:AbstractArray}, b::Type{<:AbstractArray})
+
+Simiar to gen_by_access with only one type argument. The difference is that tests for both
+arrays of type `a` and `b` are generated and `expr_gen` receives two access arguments,
+first for matrix `a` and the second for matrix `b`.
+"""
 function gen_by_access(expr_gen, a::Type{<:StaticMatrix}, b::Type)
     return quote
         return $(gen_by_access(b, :b) do access_b
@@ -94,6 +127,20 @@ function gen_by_access(expr_gen, a::Type{<:LowerTriangular{<:Any, <:StaticMatrix
         end)
     end
 end
+function gen_by_access(expr_gen, a::Type{<:Transpose{<:Any, <:StaticMatrix}}, b::Type)
+    return quote
+        return $(gen_by_access(b, :b) do access_b
+            expr_gen(:transpose, access_b)
+        end)
+    end
+end
+function gen_by_access(expr_gen, a::Type{<:Adjoint{<:Any, <:StaticMatrix}}, b::Type)
+    return quote
+        return $(gen_by_access(b, :b) do access_b
+            expr_gen(:adjoint, access_b)
+        end)
+    end
+end
 
 """
     mul_result_structure(a::Type, b::Type)
@@ -111,6 +158,14 @@ function mul_result_structure(::LowerTriangular{<:Any, <:StaticMatrix}, ::LowerT
     return LowerTriangular
 end
 
+"""
+    uplo_access(sa, asym, k, j, uplo)
+
+Generate code for matrix element access, for a matrix of size `sa` locally referred to
+as `asym` in the context where the result will be used. Both indices `k` and `j` need to be
+statically known for this function to work. `uplo` is the access pattern mode generated
+by the `gen_by_access` function.
+"""
 function uplo_access(sa, asym, k, j, uplo)
     if uplo == :any
         return :($asym[$(LinearIndices(sa)[k, j])])
@@ -150,6 +205,10 @@ function uplo_access(sa, asym, k, j, uplo)
         else
             return :(zero(T))
         end
+    elseif uplo == :transpose
+        return :($asym[$(LinearIndices(sa)[j, k])])
+    elseif uplo == :ajoint
+        return :(adjoint($asym[$(LinearIndices(sa)[j, k])]))
     end
 end
 
@@ -216,23 +275,35 @@ end
     end
 end
 
+_unstatic_array(::Type{TSA}) where {S, T, N, TSA<:StaticArray{S,T,N}} = AbstractArray{T,N}
+for TWR in [Adjoint, Transpose, Symmetric, Hermitian, LowerTriangular, UpperTriangular]
+    @eval _unstatic_array(::Type{$TWR{T,TSA}}) where {S, T, N, TSA<:StaticArray{S,T,N}} = $TWR{T,<:AbstractArray{T,N}}
+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)
         return quote
             @_inline_meta
             return mul_unrolled(Sa, Sb, a, b)
         end
-    elseif sa[1] <= 14 && sa[2] <= 14 && sb[2] <= 14
+    elseif a <: StaticMatrix && b <:StaticMatrix && sa[1] <= 14 && sa[2] <= 14 && sb[2] <= 14
         return quote
             @_inline_meta
             return mul_unrolled_chunks(Sa, Sb, a, b)
         end
-    else
+    elseif a <: StaticMatrix && b <:StaticMatrix
         return quote
             @_inline_meta
             return mul_loop(Sa, Sb, a, b)
         end
+    else
+        # we don't have any special code for handling this case so let's fall back to
+        # the generic implementation of matrix multiplication
+        return quote
+            @_inline_meta
+            return invoke(*, Tuple{$(_unstatic_array(a)),$(_unstatic_array(b))}, a, b)
+        end
     end
 end
 

src/triangular.jl

@@ -6,13 +6,13 @@
     LinearAlgebra.LowerTriangular(transpose(A.data))
 @inline adjoint(A::LinearAlgebra.UpperTriangular{<:Any,<:StaticMatrix}) =
     LinearAlgebra.LowerTriangular(adjoint(A.data))
-@inline Base.:*(A::Adjoint{<:Any,<:StaticVecOrMat}, B::LinearAlgebra.AbstractTriangular{<:Any,<:StaticMatrix}) =
+@inline Base.:*(A::Adjoint{<:Any,<:StaticVector}, B::LinearAlgebra.AbstractTriangular{<:Any,<:StaticMatrix}) =
     adjoint(adjoint(B) * adjoint(A))
-@inline Base.:*(A::Transpose{<:Any,<:StaticVecOrMat}, B::LinearAlgebra.AbstractTriangular{<:Any,<:StaticMatrix}) =
+@inline Base.:*(A::Transpose{<:Any,<:StaticVector}, B::LinearAlgebra.AbstractTriangular{<:Any,<:StaticMatrix}) =
     transpose(transpose(B) * transpose(A))
-@inline Base.:*(A::LinearAlgebra.AbstractTriangular{<:Any,<:StaticMatrix}, B::Adjoint{<:Any,<:StaticVecOrMat}) =
+@inline Base.:*(A::LinearAlgebra.AbstractTriangular{<:Any,<:StaticMatrix}, B::Adjoint{<:Any,<:StaticVector}) =
     adjoint(adjoint(B) * adjoint(A))
-@inline Base.:*(A::LinearAlgebra.AbstractTriangular{<:Any,<:StaticMatrix}, B::Transpose{<:Any,<:StaticVecOrMat}) =
+@inline Base.:*(A::LinearAlgebra.AbstractTriangular{<:Any,<:StaticMatrix}, B::Transpose{<:Any,<:StaticVector}) =
     transpose(transpose(B) * transpose(A))
 
 const StaticULT = Union{UpperTriangular{<:Any,<:StaticMatrix},LowerTriangular{<:Any,<:StaticMatrix}}

test/matrix_multiply.jl

@@ -7,7 +7,9 @@ mul_wrappers = [
     m -> Hermitian(m, :U),
     m -> Hermitian(m, :L),
     m -> UpperTriangular(m),
-    m -> LowerTriangular(m)]
+    m -> LowerTriangular(m),
+    m -> adjoint(m),
+    m -> transpose(m)]
 
 @testset "Matrix multiplication" begin
     @testset "Matrix-vector" begin
@@ -149,13 +151,19 @@ mul_wrappers = [
         # check different sizes because there are multiple implementations for matrices of different sizes
         for (mm, nn) in [
             (m, n),
-            #(SMatrix{10, 10}(collect(1:100)), SMatrix{10, 10}(collect(1:100))),
-            (SMatrix{15, 15}(collect(1:225)), SMatrix{15, 15}(collect(1:225)))]
+            (SMatrix{10, 10}(collect(1:100)), SMatrix{10, 10}(collect(1:100))),
+            (SMatrix{15, 15}(collect(1:225)), SMatrix{15, 15}(collect(1:225)))
+            ]
             for wrapper_m in mul_wrappers, wrapper_n in mul_wrappers
                 wm = wrapper_m(mm)
                 wn = wrapper_n(nn)
+                if length(mm) >= 100 && (!isa(wm, StaticArray) || !isa(wn, StaticArray))
+                    continue
+                end
                 res_structure = StaticArrays.mul_result_structure(wm, wn)
-                expected_type = if res_structure == identity
+                expected_type = if length(m) >= 100
+                    Matrix{Int}
+                elseif res_structure == identity
                     typeof(mm)
                 elseif res_structure == LowerTriangular
                     LowerTriangular{Int,typeof(mm)}