Merge pull request numpy#20846 from seiko2plus/simd_argmaxmin

ENH, SIMD: improve argmax/argmin performance
Shettyvarsha · Feb 13, 2022 · 7fe5d7f · 7fe5d7f
2 parents 344a3cb + 52787cc
commit 7fe5d7f
Show file tree

Hide file tree

Showing 6 changed files with 601 additions and 114 deletions.
diff --git a/benchmarks/benchmarks/bench_reduce.py b/benchmarks/benchmarks/bench_reduce.py
@@ -73,7 +73,8 @@ def time_max(self, dtype):
         np.fmax.reduce(self.d)
 
 class ArgMax(Benchmark):
-    params = [np.float32, np.float64, bool]
+    params = [np.int8, np.uint8, np.int16, np.uint16, np.int32, np.uint32,
+              np.int64, np.uint64, np.float32, np.float64, bool]
     param_names = ['dtype']
 
     def setup(self, dtype):
@@ -82,6 +83,17 @@ def setup(self, dtype):
     def time_argmax(self, dtype):
         np.argmax(self.d)
 
+class ArgMin(Benchmark):
+    params = [np.int8, np.uint8, np.int16, np.uint16, np.int32, np.uint32,
+              np.int64, np.uint64, np.float32, np.float64, bool]
+    param_names = ['dtype']
+
+    def setup(self, dtype):
+        self.d = np.ones(200000, dtype=dtype)
+
+    def time_argmin(self, dtype):
+        np.argmin(self.d)
+
 class SmallReduction(Benchmark):
     def setup(self):
         self.d = np.ones(100, dtype=np.float32)

diff --git a/numpy/core/setup.py b/numpy/core/setup.py
@@ -830,7 +830,7 @@ def gl_if_msvc(build_cmd):
     multiarray_deps = [
             join('src', 'multiarray', 'abstractdtypes.h'),
             join('src', 'multiarray', 'arrayobject.h'),
-            join('src', 'multiarray', 'arraytypes.h'),
+            join('src', 'multiarray', 'arraytypes.h.src'),
             join('src', 'multiarray', 'arrayfunction_override.h'),
             join('src', 'multiarray', 'array_coercion.h'),
             join('src', 'multiarray', 'array_method.h'),
@@ -892,7 +892,9 @@ def gl_if_msvc(build_cmd):
             join('src', 'multiarray', 'abstractdtypes.c'),
             join('src', 'multiarray', 'alloc.c'),
             join('src', 'multiarray', 'arrayobject.c'),
+            join('src', 'multiarray', 'arraytypes.h.src'),
             join('src', 'multiarray', 'arraytypes.c.src'),
+            join('src', 'multiarray', 'argfunc.dispatch.c.src'),
             join('src', 'multiarray', 'array_coercion.c'),
             join('src', 'multiarray', 'array_method.c'),
             join('src', 'multiarray', 'array_assign_scalar.c'),