From 56c65da33764ee5a74305594d7365cd28ad8b9bb Mon Sep 17 00:00:00 2001
From: Jack Betteridge <J.D.Betteridge@bath.ac.uk>
Date: Fri, 27 May 2022 17:14:10 +0100
Subject: [PATCH] Don't add py-cpuinfo

---
 pyop2/__init__.py      |  5 +++++
 pyop2/compilation.py   | 21 +++++++++++++++++++++
 pyop2/configuration.py | 23 +----------------------
 requirements-ext.txt   |  1 -
 4 files changed, 27 insertions(+), 23 deletions(-)

diff --git a/pyop2/__init__.py b/pyop2/__init__.py
index f0deef2e1..864b55af9 100644
--- a/pyop2/__init__.py
+++ b/pyop2/__init__.py
@@ -7,3 +7,8 @@
 from pyop2._version import get_versions
 __version__ = get_versions()['version']
 del get_versions
+
+from pyop2.configuration import configuration
+from pyop2.compilation import max_simd_width
+if configuration["vectorization_strategy"]:
+    configuration["simd_width"] = max_simd_width()
diff --git a/pyop2/compilation.py b/pyop2/compilation.py
index 1f62b2c9b..2749ddde8 100644
--- a/pyop2/compilation.py
+++ b/pyop2/compilation.py
@@ -46,6 +46,7 @@
 
 from pyop2.mpi import MPI, collective, COMM_WORLD
 from pyop2.mpi import dup_comm, get_compilation_comm, set_compilation_comm
+from pyop2.caching import cached
 from pyop2.configuration import configuration
 from pyop2.logger import warning, debug, progress, INFO
 from pyop2.exceptions import CompilationError
@@ -696,3 +697,23 @@ def clear_cache(prompt=False):
         shutil.rmtree(cachedir)
     else:
         print("Not removing cached libraries")
+
+
+@cached(cache={})
+def max_simd_width():
+    prg_str = '''#include <stdio.h>
+
+int get_simd_width(){
+    return __builtin_cpu_supports("avx512f") ? 8:
+        __builtin_cpu_supports("avx") ? 4:
+        __builtin_cpu_supports("sse") ? 2:
+        1;
+}
+'''
+    try:
+        simd_width = load(prg_str, "c", "get_simd_width", restype=ctypes.c_int)
+        width = simd_width()
+    except (OSError, CompilationError):
+        warning("Cannot sniff SIMD width, using default of 4 doubles")
+        width = 4
+    return width
diff --git a/pyop2/configuration.py b/pyop2/configuration.py
index efd78071c..b859a973a 100644
--- a/pyop2/configuration.py
+++ b/pyop2/configuration.py
@@ -40,25 +40,6 @@
 from pyop2.exceptions import ConfigurationError
 
 
-def default_simd_width():
-    from cpuinfo import get_cpu_info
-    avx_to_width = {'avx': 2, 'avx1': 2, 'avx128': 2, 'avx2': 4,
-                    'avx256': 4, 'avx3': 8, 'avx512': 8}
-    longest_simd_extension = [t for t in get_cpu_info()["flags"] if t.startswith('avx')][-1]
-    if longest_simd_extension in avx_to_width.keys():
-        return avx_to_width[longest_simd_extension]
-    elif longest_simd_extension[:6] in avx_to_width.keys():
-        return avx_to_width[longest_simd_extension[:6]]
-    elif longest_simd_extension[:4] in avx_to_width.keys():
-        return avx_to_width[longest_simd_extension[:4]]
-    else:
-        raise ConfigurationError(f"The vector extension of your architecture is unknown.\
-                                   Must be one of {str(avx_to_width.keys())}.\
-                                   We advise to disable vectorisation \
-                                   with export PYOP2_VECT_STRATEGY=""."
-                                 )
-
-
 class Configuration(dict):
     r"""PyOP2 configuration parameters
 
@@ -117,7 +98,7 @@ class Configuration(dict):
         "ldflags":
             ("PYOP2_LDFLAGS", str, ""),
         "simd_width":
-            ("PYOP2_SIMD_WIDTH", int, 1),
+            ("PYOP2_SIMD_WIDTH", int, 4),
         "vectorization_strategy":
             ("PYOP2_VECT_STRATEGY", str, "cross-element"),
         "alignment":
@@ -193,7 +174,5 @@ def __setitem__(self, key, value):
 
 
 configuration = Configuration()
-if configuration["vectorization_strategy"]:
-    configuration["simd_width"] = default_simd_width()
 
 target = CWithGNULibcTarget()
diff --git a/requirements-ext.txt b/requirements-ext.txt
index 2121947f8..75adb64e3 100644
--- a/requirements-ext.txt
+++ b/requirements-ext.txt
@@ -4,7 +4,6 @@ pytest>=2.3
 flake8>=2.1.0
 pycparser>=2.10
 mpi4py>=1.3.1
-py-cpuinfo
 decorator<=4.4.2
 dataclasses
 cachetools