Faster 3-argument dot

benchmark/bench_dot3.jl

@@ -0,0 +1,73 @@
+module BenchDot3
+
+# Three-argument `dot(x, A, y)` compared to the generic method from LinearAlgebra
+# (the implementation used before StaticArrays had a specialized method) and to
+# `dot(x, A * y)`. Run this file directly to print a table that flags sizes where the
+# specialized method is slower than either alternative:
+#
+#     julia --project=benchmark benchmark/bench_dot3.jl
+
+import Random
+using BenchmarkTools
+using LinearAlgebra
+using Printf
+using StaticArrays
+
+const suite = BenchmarkGroup()
+
+const shapes = if haskey(ENV, "GITHUB_EVENT_PATH")
+    [(2, 2), (3, 3), (4, 4), (8, 8), (15, 15)]
+else
+    [[(n, n) for n in 1:15]; [(3, 8), (8, 3), (4, 12), (15, 2), (2, 15)]]
+end
+
+const eltypes = (Float64, Float32, ComplexF64, ComplexF32)
+
+generic_dot(x, A, y) = invoke(dot, Tuple{AbstractVector, AbstractMatrix, AbstractVector}, x, A, y)
+matvec_dot(x, A, y) = dot(x, A * y)
+
+const implementations = (("dot(x,A,y)", dot), ("generic", generic_dot), ("dot(x,A*y)", matvec_dot))
+
+Random.seed!(1234)
+
+for T in eltypes
+    s1 = suite[string(T)] = BenchmarkGroup()
+    for (M, N) in shapes
+        x = @SVector rand(T, M)
+        A = @SMatrix rand(T, M, N)
+        y = @SVector rand(T, N)
+        s2 = s1[string(M, pad=2) * "x" * string(N, pad=2)] = BenchmarkGroup()
+        for (label, f) in implementations
+            s2[label] = @benchmarkable $f($x, $A, $y)
+        end
+    end
+end
+
+# Print minimum times and mark sizes where `dot(x,A,y)` is slower than an alternative by
+# more than the relative tolerance `rtol` plus an absolute tolerance `atol` (in ns).
+function compare(; rtol = 0.05, atol = 0.3, seconds = 0.5)
+    tune!(suite)
+    results = run(suite; seconds = seconds)
+    nregressions = 0
+    for T in eltypes
+        println(T)
+        r1 = results[string(T)]
+        for key in sort(collect(keys(r1)))
+            times = [time(minimum(r1[key][label])) for (label, _) in implementations]
+            slower = times[1] > (1 + rtol) * minimum(times[2:end]) + atol
+            nregressions += slower
+            @printf("  %s  dot(x,A,y) %8.2f ns  generic %8.2f ns  dot(x,A*y) %8.2f ns%s\n",
+                key, times..., slower ? "  <-- regression" : "")
+        end
+    end
+    println(nregressions == 0 ? "No regressions." : "$nregressions regression(s).")
+    return nregressions
+end
+
+end  # module
+
+if abspath(PROGRAM_FILE) == @__FILE__
+    exit(BenchDot3.compare() == 0 ? 0 : 1)
+end
+
+BenchDot3.suite

src/linalg.jl

@@ -276,6 +276,67 @@ end
     return ret
 end
 
+@inline dot(x::StaticVector, A::StaticMatrix, y::StaticVector) = _dot3(Size(x), Size(A), Size(y), x, A, y)
+
+# Maximum size in bytes of `A` for which `_dot3` is fully unrolled.
+const DOT3_UNROLL_BUDGET = 8192
+
+# Fully unrolled `dot(x, A*y)`: the accumulators `t_i` of `A*y` are updated column by
+# column with independent operations, a pattern that LLVM's SLP vectorizer turns into packed
+# loads and multiply-adds. Unlike the generic method, columns where `y[j]` is zero are not
+# skipped. Arrays of non-numbers use the generic method.
+#
+# The choice between `muladd` and `+`/`*` and of the final reduction was made empirically:
+# on x86-64 `muladd` of complex numbers with 64-bit parts prevents vectorization, and for
+# non-real numbers reducing with the `@simd` loop of the two-argument `dot` is faster.
+# On other architectures (e.g. Apple M5) `+`/`*` is slower than the `muladd`s of `A * y`.
+@generated function _dot3(::Size{sx}, ::Size{sA}, ::Size{sy}, x::StaticVector{<:Any, Tx}, A::StaticMatrix{<:Any, <:Any, TA}, y::StaticVector{<:Any, Ty}) where {sx, sA, sy, Tx, TA, Ty}
+    M, N = sA
+    if sx[1] != M || sy[1] != N
+        return :(throw(DimensionMismatch("dimensions of x $sx, A $sA and y $sy are incompatible")))
+    end
+    if M == 0 || N == 0
+        # outermost zero call to avoid spurious sign ambiguity (like 0.0 - 0.0im)
+        return :(zero(dot(zero(eltype(x)), zero(eltype(A)), zero(eltype(y)))))
+    end
+    if !(Tx <: Number && TA <: Number && Ty <: Number)
+        return :(invoke(dot, Tuple{AbstractVector, AbstractMatrix, AbstractVector}, x, A, y))
+    end
+    if M * N > 256 || !isbitstype(TA) || M * N * sizeof(TA) > DOT3_UNROLL_BUDGET
+        return :(dot(x, A * y))
+    end
+
+    is_real = Tx <: Real && TA <: Real && Ty <: Real
+    use_muladd = is_real || all(T -> T <: Union{Real, Complex} && sizeof(real(T)) <= 4, (TA, Ty))
+    if !use_muladd && Sys.ARCH !== :x86_64
+        return :(dot(x, A * y))
+    end
+    t = [Symbol("t_", i) for i = 1:M]
+    exprs = [:($(t[i]) = A[$i] * y[1]) for i = 1:M]
+    for j = 2:N, i = 1:M
+        a = :(A[$(LinearIndices(sA)[i, j])])
+        rhs = use_muladd ? :(muladd($a, y[$j], $(t[i]))) : :($(t[i]) + $a * y[$j])
+        push!(exprs, :($(t[i]) = $rhs))
+    end
+    if is_real
+        s = :(dot(x[1], $(t[1])))
+        for i = 2:M
+            s = :($s + dot(x[$i], $(t[i])))
+        end
+    else
+        s = :(dot(x, SVector{$M}($(t...))))
+    end
+
+    return quote
+        @_inline_meta
+        @inbounds begin
+            $(exprs...)
+            return $s
+        end
+    end
+end
+
+
 #--------------------------------------------------
 # Norms
 _inner_eltype(v::AbstractArray) = isempty(v) ? eltype(v) : _inner_eltype(first(v))

test/linalg.jl

@@ -266,6 +266,66 @@ end
         @test @inferred(dot(m1, m2)) === dot(v1, v2)
     end
 
+    @testset "three-argument dot" begin
+        generic_dot(x, A, y) = invoke(dot, Tuple{AbstractVector, AbstractMatrix, AbstractVector}, x, A, y)
+
+        x = @SVector [1, 2, 3]
+        A = @SMatrix [1 2; 3 4; 5 6]
+        y = @SVector [7, 8]
+        @test @inferred(dot(x, A, y)) === generic_dot(x, A, y) === dot(Vector(x), Matrix(A), Vector(y))
+        @test @inferred(dot(MVector(x), MMatrix(A), MVector(y))) === generic_dot(x, A, y)
+        @test @inferred(dot(SizedVector{3}([1, 2, 3]), SizedMatrix{3,2}(Matrix(A)), y)) === generic_dot(x, A, y)
+        @test_throws DimensionMismatch dot(y, A, y)
+        @test_throws DimensionMismatch dot(x, A, x)
+
+        # real and complex element types, including mixed ones, for which different
+        # code paths are generated
+        for (M, N) in ((1, 1), (2, 2), (3, 3), (4, 7), (7, 4), (15, 15), (16, 16), (17, 17), (3, 40))
+            for (Tx, TA, Ty) in ((Float64, Float64, Float64), (Float32, Float32, Float32),
+                                 (ComplexF64, ComplexF64, ComplexF64), (Float64, ComplexF64, Float64),
+                                 (ComplexF32, Float32, ComplexF64), (Int, Int, Int))
+                x = @SVector rand(Tx, M)
+                A = @SMatrix rand(TA, M, N)
+                y = @SVector rand(Ty, N)
+                r = generic_dot(x, A, y)
+                @test @inferred(dot(x, A, y)) isa typeof(r)
+                # the generic method accumulates `A'*x` instead of `A*y`, which matters for mixed precision
+                @test dot(x, A, y) ≈ r rtol=1e-5
+                @test dot(x, A, y) ≈ dot(x, A * y)
+            end
+        end
+
+        # complex conjugation of `x` only
+        x = @SVector [1.0im, 2.0]
+        A = @SMatrix [1.0 2.0im; 3.0 4.0]
+        y = @SVector [1.0im, 1.0]
+        @test dot(x, A, y) ≈ dot(Vector(x), Matrix(A), Vector(y))
+
+        # element type without a bits representation
+        x = SVector(big(1.0), big(2.0))
+        A = SMatrix{2,2}(big.([1.0 2.0; 3.0 4.0]))
+        y = SVector(big(1.0), big(2.0))
+        @test @inferred(dot(x, A, y)) == generic_dot(x, A, y) == 27
+
+        # recursive like the generic method
+        x = @SVector [[1, 2], [3, 4]]
+        A = SMatrix{2,2}(([1 2; 3 4], [9 10; 11 12], [5 6; 7 8], [13 14; 15 16]))
+        y = @SVector [[1, 0], [0, 1]]
+        @test @inferred(dot(x, A, y)) == dot(Vector(x), Matrix(A), Vector(y))
+
+        # empty arrays
+        @test @inferred(dot(SVector{0,Float64}(), SMatrix{0,2,Float64}(), SVector(1.0, 2.0))) === 0.0
+        @test @inferred(dot(SVector(1.0, 2.0), SMatrix{2,0,Float64}(), SVector{0,Float64}())) === 0.0
+        @test @inferred(dot(SVector{0,ComplexF64}(), SMatrix{0,0,ComplexF64}(), SVector{0,ComplexF64}())) === zero(ComplexF64)
+
+        alloc_dot(x, A, y) = @allocated dot(x, A, y)
+        x = @SVector rand(8)
+        A = @SMatrix rand(8, 8)
+        y = @SVector rand(8)
+        alloc_dot(x, A, y)
+        @test alloc_dot(x, A, y) == 0
+    end
+
     @testset "transpose() and conj()" begin
         @test @inferred(conj(SVector(1+im, 2+im))) === SVector(1-im, 2-im)
 

test/unitful.jl

@@ -11,4 +11,9 @@ using Unitful
 
     @test norm([SVector(1.0u"m", 2.0u"m")]) == norm([1.0u"m", 2.0u"m"])
     @test isapprox(SVector(1.0u"m", 2.0u"m"), SVector(1.1u"m", 2.1u"m"); atol=0.2u"m")
+
+    x = SVector(1.0u"m", 2.0u"m")
+    A = SMatrix{2,2}(1.0u"s", 3.0u"s", 2.0u"s", 4.0u"s")
+    y = SVector(1.0u"kg", 2.0u"kg")
+    @test @inferred(dot(x, A, y)) == dot(Vector(x), Matrix(A), Vector(y)) == 27.0u"m*s*kg"
 end