Added size-inferable matrix-vector products

Can now do a fast StaticMatrix * AbstractVector and
related `A_mul_B!` calls.

Author: Andy Ferris

src/matrix_multiply.jl

@@ -99,6 +99,7 @@ end
 
     s = (sA[1],)
     T = promote_op(matprod, TA, Tb)
+    #println(T)
 
     if sb[1] != sA[2]
         error("Dimension mismatch")
@@ -130,6 +131,37 @@ end
     end
 end
 
+# This happens to be size-inferrable from A
+@generated function *(A::StaticMatrix, b::AbstractVector)
+    TA = eltype(A)
+    Tb = eltype(b)
+    sA = size(A)
+    #sb = size(b)
+
+    s = (sA[1],)
+    T = promote_op(matprod, TA, Tb)
+
+    if T == Tb
+        newtype = similar_type(A, s)
+    else
+        newtype = similar_type(A, T, s)
+    end
+
+    if sA[2] != 0
+        exprs = [reduce((ex1,ex2) -> :(+($ex1,$ex2)), [:(A[$(sub2ind(sA, k, j))]*b[$j]) for j = 1:sA[2]]) for k = 1:sA[1]]
+    else
+        exprs = [zero(T) for k = 1:sA[1]]
+    end
+
+    return quote
+        $(Expr(:meta,:inline))
+        if length(b) != $(sA[2])
+            error("Dimension mismatch")
+        end
+        @inbounds return $(Expr(:call, newtype, Expr(:tuple, exprs...)))
+    end
+end
+
 @generated function *(a::StaticVector, B::StaticMatrix)
     Ta = eltype(a)
     TB = eltype(B)
@@ -402,7 +434,7 @@ end
 
 # The idea here is to get pointers to stack variables and call BLAS.
 # This saves an aweful lot of time compared to copying SArray's to Ref{SArray{...}}
-# and should be fastest for (very) large SArrays
+# and using BLAS should be fastest for (very) large SArrays
 
 # Here is an LLVM function that gets the pointer to its input, %x
 # After this we would make the ccall above.
@@ -413,6 +445,48 @@ end
 #     ret i32* %1
 # }
 
+@generated function A_mul_B!(c::StaticVector, A::StaticMatrix, b::StaticVector)
+    sA = size(A)
+    sb = size(b)
+    s = size(c)
+    T = eltype(c)
+
+    if sb[1] != sA[2] || s[1] != sA[1]
+        error("Dimension mismatch")
+    end
+
+    if sA[2] != 0
+        exprs = [:(c[$k] = $(reduce((ex1,ex2) -> :(+($ex1,$ex2)), [:(A[$(sub2ind(sA, k, j))]*b[$j]) for j = 1:sA[2]]))) for k = 1:sA[1]]
+    else
+        exprs = [:(c[$k] = $(zero(T))) for k = 1:sA[1]]
+    end
+
+    return quote
+        $(Expr(:meta,:inline))
+        @inbounds $(Expr(:block, exprs...))
+    end
+end
+
+# The unrolled code is inferrable from the size of A
+@generated function A_mul_B!(c::AbstractVector, A::StaticMatrix, b::AbstractVector)
+    sA = size(A)
+    T = eltype(c)
+
+    if sA[2] != 0
+        exprs = [:(c[$k] = $(reduce((ex1,ex2) -> :(+($ex1,$ex2)), [:(A[$(sub2ind(sA, k, j))]*b[$j]) for j = 1:sA[2]]))) for k = 1:sA[1]]
+    else
+        exprs = [:(c[$k] = $(zero(T))) for k = 1:sA[1]]
+    end
+
+    return quote
+        $(Expr(:meta,:inline))
+        if length(b) != $(sA[2]) || length(c) != $(sA[1])
+            error("Dimension mismatch")
+        end
+        @inbounds $(Expr(:block, exprs...))
+    end
+end
+
 
 @generated function A_mul_B!(C::StaticMatrix, A::StaticMatrix, B::StaticMatrix)
     if isbits(C)

test/matrix_multiply.jl

@@ -4,22 +4,25 @@
         v = @SVector [1, 2]
         @test m*v === @SVector [5, 11]
         # More complicated eltype inference
-        v = @SVector [CartesianIndex((1,3)), CartesianIndex((3,1))]
-        x = @inferred(m*v)
+        v2 = @SVector [CartesianIndex((1,3)), CartesianIndex((3,1))]
+        x = @inferred(m*v2)
         @test isa(x, SVector{2,CartesianIndex{2}})
         @test x == @SVector [CartesianIndex((7,5)), CartesianIndex((15,13))]
 
-        m = @MMatrix [1 2; 3 4]
-        v = @MVector [1, 2]
-        @test (m*v)::MVector == @MVector [5, 11]
+        v3 = [1, 2]
+        @test m*v3 === @SVector [5, 11]
 
-        m = @SArray [1 2; 3 4]
-        v = @SArray [1, 2]
-        @test m*v === @SArray [5, 11]
+        m2 = @MMatrix [1 2; 3 4]
+        v4 = @MVector [1, 2]
+        @test (m2*v4)::MVector == @MVector [5, 11]
 
-        m = @MArray [1 2; 3 4]
-        v = @MArray [1, 2]
-        @test (m*v)::MArray == @MArray [5, 11]
+        m3 = @SArray [1 2; 3 4]
+        v5 = @SArray [1, 2]
+        @test m3*v5 === @SArray [5, 11]
+
+        m4 = @MArray [1 2; 3 4]
+        v6 = @MArray [1, 2]
+        @test (m4*v6)::MArray == @MArray [5, 11]
     end
 
     @testset "Vector-matrix" begin
@@ -117,9 +120,18 @@
     end
 
     @testset "A_mul_B!" begin
+        v = @SVector [2, 4]
+        v2 = [2, 4]
         m = @SMatrix [1 2; 3 4]
         n = @SMatrix [2 3; 4 5]
 
+        outvec = MVector{2,Int}()
+        A_mul_B!(outvec, m, v)
+        @test outvec == @MVector [10,22]
+        outvec2 = Vector{Int}(2)
+        A_mul_B!(outvec2, m, v2)
+        @test outvec2 == [10,22]
+
         a = MMatrix{2,2,Int,4}()
         A_mul_B!(a, m, n)
         @test a::MMatrix{2,2,Int,4} == @MMatrix [10 13; 22 29]