From a1aa074ee542f68060d192d8710adaddeb3a3290 Mon Sep 17 00:00:00 2001
From: Michael Abbott <me@escbook>
Date: Thu, 8 Oct 2020 09:33:21 +0200
Subject: [PATCH] add simple 3-arg and 4-arg * methods

---
 src/matrix_multiply.jl  | 10 +++++++---
 test/matrix_multiply.jl |  3 +++
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/src/matrix_multiply.jl b/src/matrix_multiply.jl
index 2b08ca4a7..76d54eb93 100644
--- a/src/matrix_multiply.jl
+++ b/src/matrix_multiply.jl
@@ -15,11 +15,15 @@ import LinearAlgebra: BlasFloat, matprod, mul!
 @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
 
+# Avoid LinearAlgebra._quad_matmul's order calculation on equal sizes
+@inline *(A::StaticMatrix{N,N}, B::StaticMatrix{N,N}, C::StaticMatrix{N,N}) where {N} = (A*B)*C
+@inline *(A::StaticMatrix{N,N}, B::StaticMatrix{N,N}, C::StaticMatrix{N,N}, D::StaticMatrix{N,N}) where {N} = ((A*B)*C)*D
+
 """
     mul_result_structure(a::Type, b::Type)
 
 Get a structure wrapper that should be applied to the result of multiplication of matrices
-of given types (a*b). 
+of given types (a*b).
 """
 function mul_result_structure(a, b)
     return identity
@@ -119,7 +123,7 @@ end
     else
         exprs = [:(a[$i] * transpose(b[$j])) for i = 1:sa[1], j = 1:sb[2]]
     end
-    
+
     return quote
         @_inline_meta
         T = promote_op(*, Ta, Tb)
@@ -214,7 +218,7 @@ end
         while m < M
             mu = min(M, m + M_r)
             mrange = m+1:mu
-            
+
             atemps_init = [:($(atemps[k1]) = a[$k1]) for k1 = mrange]
             exprs_init = [:($(tmps[k1,k2])  = $(atemps[k1]) * b[$(1 + (k2-1) * sb[1])]) for k1 = mrange, k2 = nrange]
             atemps_loop_init = [:($(atemps[k1]) = a[$(k1-sa[1]) + $(sa[1])*j]) for k1 = mrange]
diff --git a/test/matrix_multiply.jl b/test/matrix_multiply.jl
index d9996ad26..794332153 100644
--- a/test/matrix_multiply.jl
+++ b/test/matrix_multiply.jl
@@ -152,6 +152,9 @@ mul_wrappers = [
         @test m*transpose(n) === @SMatrix [8 14; 18 32]
         @test transpose(m)*transpose(n) === @SMatrix [11 19; 16 28]
 
+        @test @inferred(m*n*m)  === @SMatrix [49 72; 109 160]
+        @test @inferred(m*n*m*n)  === @SMatrix [386 507; 858 1127]
+
         # check different sizes because there are multiple implementations for matrices of different sizes
         for (mm, nn) in [
             (m, n),