small patch to avx

BlaiseMuhirwa · Nov 30, 2023 · f75ad3c · f75ad3c
1 parent 5fcc8e8
commit f75ad3c
Show file tree

Hide file tree

Showing 5 changed files with 47 additions and 17 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -34,8 +34,7 @@ set(CMAKE_CXX_FLAGS
     -ffast-math \
     -funroll-loops \
     -mavx \
-    -mavx512f \
-    -ftree-vectorize")
+    -mavx512f")
 
 option(CMAKE_BUILD_TYPE "Build type" Release)
 if(CMAKE_BUILD_TYPE STREQUAL "Debug")
@@ -155,9 +154,11 @@ message(STATUS "Building tests: ${BUILD_TESTS}")
 message(STATUS "Building examples: ${BUILD_EXAMPLES}")
 message(STATUS "Building benchmarks: ${BUILD_BENCHMARKS}")
 
+# Enable auto-vectorization if we are not using SIMD.
 if(NO_MANUAL_VECTORIZATION)
   message(STATUS "Disabling manual vectorization (SIMD)")
   add_definitions(-DNO_MANUAL_VECTORIZATION)
+  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -ftree-vectorize")
 endif()
 
 if(BUILD_BENCHMARKS)

diff --git a/flatnav/distances/InnerProductDistance.h b/flatnav/distances/InnerProductDistance.h
@@ -90,7 +90,7 @@ class InnerProductDistance : public DistanceInterface<InnerProductDistance> {
       _distance_computer = distanceImplInnerProductSIMD16ExtAVX;
     }
 #endif
-    if (!_dimension % 16 == 0) {
+    if (!(_dimension % 16 == 0)) {
       if (_dimension % 4 == 0) {
         _distance_computer = distanceImplInnerProductSIMD4ExtSSE;
       } else if (_dimension > 16) {

diff --git a/flatnav/distances/SquaredL2Distance.h b/flatnav/distances/SquaredL2Distance.h
@@ -88,7 +88,7 @@ class SquaredL2Distance : public DistanceInterface<SquaredL2Distance> {
       _distance_computer = distanceImplSquaredL2SIMD16ExtAVX;
     }
 #endif
-    if (!_dimension % 16 == 0) {
+    if (!(_dimension % 16 == 0)) {
       if (_dimension % 4 == 0) {
         _distance_computer = distanceImplSquaredL2SIMD4Ext;
       } else if (_dimension > 16) {

diff --git a/flatnav/util/SIMDDistanceSpecializations.h b/flatnav/util/SIMDDistanceSpecializations.h
@@ -628,4 +628,12 @@ static float distanceImplSquaredL2SIMD16ExtResiduals(const void *x,
 }
 #endif
 
-// namespace flatnav::util
+// namespace flatnav::util
+
+// with manual vectorization: [INFO] Mean Recall: 0.999931, Duration:2.7517
+// milliseconds without manual vectorization: [INFO] Mean Recall: 0.999931,
+// Duration:2.55 milliseconds
+
+// without the -ftree-vectorize option and no manual vectorization:
+// duration: 2.5767 without the -ftree-vectorize option but with manual
+// vectorization: duration:
diff --git a/flatnav_python/setup.py b/flatnav_python/setup.py
@@ -1,14 +1,30 @@
 import os
 from pybind11.setup_helpers import Pybind11Extension, build_ext
 from setuptools import setup
-import sys
+import sys, subprocess
 
 __version__ = "0.0.1"
 
 CURRENT_DIR = os.getcwd()
 SOURCE_PATH = os.path.join(CURRENT_DIR, "python_bindings.cpp")
 
 
+def platform_has_avx_support():
+    """
+    TODO: Make this check more robust across compilers and supported platforms.
+    """
+    if sys.platform in ["linux", "linux2", "darwin"]:
+        try:
+            output = subprocess.check_output("cat /proc/cpuinfo", shell=True)
+            decoded_output = output.decode()
+            return "avx" in decoded_output or "avx512" in decoded_output
+        except Exception:
+            return False
+
+    else:
+        return False
+
+
 INCLUDE_DIRS = [
     os.path.join(CURRENT_DIR, ".."),
     os.path.join(CURRENT_DIR, "..", "external", "cereal", "include"),
@@ -24,24 +40,29 @@
     EXTRA_LINK_ARGS.extend(["-fopenmp"])
 
 
+EXTRA_COMPILE_ARGS = [
+    omp_flag,  # Enable OpenMP
+    "-Ofast",  # Use the fastest optimization
+    "-fpic",  # Position-independent code
+    "-w",  # Suppress all warnings (note: this overrides -Wall)
+    "-ffast-math",  # Enable fast math optimizations
+    "-funroll-loops",  # Unroll loops
+    "-mavx",  # Enable AVX instructions
+    "-mavx512f",  # Enable AVX-512 instructions
+]
+
+if not platform_has_avx_support():
+    EXTRA_COMPILE_ARGS.append("-ftree-vectorize")
+
+
 ext_modules = [
     Pybind11Extension(
         "flatnav",
         [SOURCE_PATH],
         define_macros=[("VERSION_INFO", __version__)],
         cxx_std=17,
         include_dirs=INCLUDE_DIRS,
-        extra_compile_args=[
-            omp_flag,  # Enable OpenMP
-            "-Ofast",  # Use the fastest optimization
-            "-fpic",  # Position-independent code
-            "-w",  # Suppress all warnings (note: this overrides -Wall)
-            "-ffast-math",  # Enable fast math optimizations
-            "-funroll-loops",  # Unroll loops
-            "-ftree-vectorize",  # Vectorize where possible
-            "-mavx",  # Enable AVX instructions
-            "-mavx512f",  # Enable AVX-512 instructions
-        ],
+        extra_compile_args=EXTRA_COMPILE_ARGS,
         extra_link_args=EXTRA_LINK_ARGS,  # Link OpenMP when linking the extension
     )
 ]