specialize mul_fast and add_fast to allow FMA (#988)

* specialize mul_fast and add_fast to allow FMA

* bump patch version number

Author: ranocha

Project.toml

@@ -1,6 +1,6 @@
 name = "StaticArrays"
 uuid = "90137ffa-7385-5640-81b9-e52037218182"
-version = "1.3.2"
+version = "1.3.3"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

src/linalg.jl

@@ -50,6 +50,15 @@ end
 @inline Base.muladd(scalar::Number, a::StaticArray, b::StaticArray) = map((ai, bi) -> muladd(scalar, ai, bi), a, b)
 @inline Base.muladd(a::StaticArray, scalar::Number, b::StaticArray) = map((ai, bi) -> muladd(ai, scalar, bi), a, b)
 
+
+# @fastmath operators
+@inline Base.FastMath.mul_fast(a::Number, b::StaticArray) = map(c -> Base.FastMath.mul_fast(a, c), b)
+@inline Base.FastMath.mul_fast(a::StaticArray, b::Number) = map(c -> Base.FastMath.mul_fast(c, b), a)
+
+@inline Base.FastMath.add_fast(a::StaticArray, b::StaticArray) = map(Base.FastMath.add_fast, a, b)
+@inline Base.FastMath.sub_fast(a::StaticArray, b::StaticArray) = map(Base.FastMath.sub_fast, a, b)
+
+
 #--------------------------------------------------
 # Matrix algebra
 

test/linalg.jl

@@ -69,6 +69,31 @@ StaticArrays.similar_type(::Union{RotMat2,Type{RotMat2}}) = SMatrix{2,2,Float64,
         end
     end
 
+    @testset "@fastmath operators" begin
+        for T in (Int, Float32, Float64)
+            s0 = convert(T, 2)
+            v1 = @SVector T[2, 4, 6, 8]
+            v2 = @SVector T[4, 3, 2, 1]
+            m1 = @SMatrix T[2 4; 6 8]
+            m2 = @SMatrix T[4 3; 2 1]
+
+            # Use that these small integers can be represetnted exactly
+            # as floating point numbers. In general, the comparison of
+            # floats should use `≈` instead of `===`.
+            # These should be turned into `vfmadd...` calls
+            @test @fastmath(@inferred(s0 * v1 + v2)) === @SVector T[8, 11, 14, 17]
+            @test @fastmath(@inferred(v1 * s0 + v2)) === @SVector T[8, 11, 14, 17]
+            @test @fastmath(@inferred(s0 * m1 + m2)) === @SMatrix T[8 11; 14 17]
+            @test @fastmath(@inferred(m1 * s0 + m2)) === @SMatrix T[8 11; 14 17]
+
+            # These should be turned into `vfmsub...` calls
+            @test @fastmath(@inferred(s0 * v1 - v2)) === @SVector T[0, 5, 10, 15]
+            @test @fastmath(@inferred(v1 * s0 - v2)) === @SVector T[0, 5, 10, 15]
+            @test @fastmath(@inferred(s0 * m1 - m2)) === @SMatrix T[0 5; 10 15]
+            @test @fastmath(@inferred(m1 * s0 - m2)) === @SMatrix T[0 5; 10 15]
+        end
+    end
+
     @testset "Interaction with `UniformScaling`" begin
         @test @inferred(@SMatrix([0 1; 2 3]) + I) === @SMatrix [1 1; 2 4]
         @test @inferred(I + @SMatrix([0 1; 2 3])) === @SMatrix [1 1; 2 4]