improve float64 performance by sse2

1 files changed, 45 insertions, 1 deletions

diff --git a/benchmarks/benchmarks/bench_linalg.py b/benchmarks/benchmarks/bench_linalg.py
index 3abbe3670..d5f2a0f8b 100644
--- a/benchmarks/benchmarks/bench_linalg.py
+++ b/benchmarks/benchmarks/bench_linalg.py
@@ -2,7 +2,6 @@ from .common import Benchmark, get_squares_, get_indexes_rand, TYPES1
 
 import numpy as np
 
-
 class Eindot(Benchmark):
     def setup(self):
         self.a = np.arange(60000.0).reshape(150, 400)
@@ -105,3 +104,48 @@ class Lstsq(Benchmark):
 
     def time_numpy_linalg_lstsq_a__b_float64(self):
         np.linalg.lstsq(self.a, self.b, rcond=-1)
+
+class Einsum(Benchmark):
+    param_names = ['dtype']
+    params = [[np.float32, np.float64]]
+    def setup(self, dtype):
+        self.a = np.arange(3000, dtype=dtype)
+        self.b = np.arange(2990, dtype=dtype)
+        self.c = np.arange(24000, dtype=dtype).reshape(20, 30, 40)
+        self.c1 = np.arange(1200, dtype=dtype).reshape(30, 40)
+        self.c2 = np.arange(40, dtype=dtype)
+        self.c3 = np.arange(30000, dtype=dtype).reshape(30, 20, 50)
+        self.d = np.arange(2*1000, dtype=dtype).reshape(2, 1000)
+        self.e = np.arange(100*100, dtype=dtype).reshape(100, 100)
+
+    #outer(a,b)
+    def time_einsum_outer(self, dtype):
+        np.einsum("i,j", self.a, self.b, optimize=True)
+
+    #inner(a,b)
+    def time_einsum_inner(self, dtype):
+        np.einsum("...i, ...i", self.c, self.c2, optimize=True)
+
+    # swap axes
+    def time_einsum_swap(self, dtype):
+        np.einsum("ijk->jik", self.c, optimize=True)
+    
+    # sum(a, axis=0)
+    def time_einsum_sum(self, dtype):
+        np.einsum("i...->...", self.d, optimize=True)
+    
+    # trace(a)
+    def time_einsum_trace(self, dtype):
+        np.einsum("ii", self.e, optimize=True)
+
+    # multiply(a, b)
+    def time_einsum_multiply(self, dtype):
+        np.einsum("..., ...", self.c1, self.c , optimize=True)
+    
+    # tensordot(a, b)
+    def time_einsum_tensordot(self, dtype):
+        np.einsum("ijk, jil -> kl", self.c, self.c3 , optimize=True)
+
+    # a.dot(b)
+    def time_einsum_matmat(self, dtype):
+        np.einsum("ij,jk", self.e, self.e , optimize=True)