support apple darwin

CMakeLists.txt

@@ -25,7 +25,6 @@ set(CMAKE_CXX_STANDARD 17)
 set(CMAKE_CXX_FLAGS
     "${CMAKE_CXX_FLAGS} \
     -std=c++17 \
-    -Wall \
     -Ofast \
     -DHAVE_CXX0X \
     -DNDEBUG \
@@ -38,10 +37,9 @@ set(CMAKE_CXX_FLAGS
 
 option(CMAKE_BUILD_TYPE "Build type" Release)
 if(CMAKE_BUILD_TYPE STREQUAL "Debug")
-  # In Debug mode, we will also add the -g flag to enable debugging and the
-  # address sanitizer.
+  # Add debug compile flags
   message(STATUS "Building in Debug mode")
-  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -g -fsanitize=address")
+  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -g -Wall -fsanitize=address")
 endif()
 
 include(ExternalProject)

bin/build.sh

@@ -55,6 +55,8 @@ if [[ "$(uname)" == "Darwin" ]]; then
     echo "Using LLVM clang"
     export CC=/opt/homebrew/opt/llvm/bin/clang
     export CXX=/opt/homebrew/opt/llvm/bin/clang++
+    export LDFLAGS="-L/opt/homebrew/opt/libomp/lib"
+    export CPPFLAGS="-I/opt/homebrew/opt/libomp/include"
 elif [[ "$(uname)" == "Linux" ]]; then
     echo "Using system clang"
 else

flatnav_python/pyproject.toml

@@ -13,8 +13,8 @@ authors = [
 license = "Apache 2.0"
 
 [tool.poetry.dependencies]
-python = "3.10.12"
-pybind11 = "2.8.1"
+python = ">=3.9"
+pybind11 = "2.10.4"
 setuptools = "68.2.2"
 
 [tool.poetry.dev-dependencies]
@@ -27,5 +27,5 @@ requests = "^2.31.0"
 
 
 [build-system]
-requires = ["setuptools>=68", "wheel", "pybind11=2.8.1", "python=3.10.12"]
+requires = ["setuptools>=68", "wheel", "pybind11=2.10.4", "python>=3.9"]
 build-backend = "setuptools.build_meta"

flatnav_python/setup.py

@@ -1,33 +1,46 @@
 import os
 from pybind11.setup_helpers import Pybind11Extension, build_ext
 from setuptools import setup
+import sys
 
 __version__ = "0.0.1"
 
 CURRENT_DIR = os.getcwd()
 SOURCE_PATH = os.path.join(CURRENT_DIR, "python_bindings.cpp")
 
 
+INCLUDE_DIRS = [
+    os.path.join(CURRENT_DIR, ".."),
+    os.path.join(CURRENT_DIR, "..", "external", "cereal", "include"),
+]
+EXTRA_LINK_ARGS = []
+
+if sys.platform == "darwin":
+    omp_flag = "-Xclang -fopenmp"
+    INCLUDE_DIRS.extend(["/opt/homebrew/opt/libomp/include"])
+    EXTRA_LINK_ARGS.extend(["-lomp", "-L/opt/homebrew/opt/libomp/lib"])
+elif sys.platform() == "linux":
+    omp_flag = "-fopenmp"
+    EXTRA_LINK_ARGS.extend(["-fopenmp"])
+
+
 ext_modules = [
     Pybind11Extension(
         "flatnav",
         [SOURCE_PATH],
         define_macros=[("VERSION_INFO", __version__)],
         cxx_std=17,
-        include_dirs=[
-            os.path.join(CURRENT_DIR, ".."),
-            os.path.join(CURRENT_DIR, "..", "external", "cereal", "include"),
-        ],
+        include_dirs=INCLUDE_DIRS,
         extra_compile_args=[
-            "-fopenmp",  # Enable OpenMP
+            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
         ],
-        extra_link_args=["-fopenmp"],  # Link OpenMP when linking the extension
+        extra_link_args=EXTRA_LINK_ARGS,  # Link OpenMP when linking the extension
     )
 ]
 

flatnav_python/test_index.py

@@ -147,7 +147,7 @@ def test_flatnav_ip_index_random_dataset():
     dataset_to_index = generate_random_data(dataset_length=30_000, dim=225)
     queries = generate_random_data(dataset_length=5_000, dim=225)
     ground_truth = np.random.randint(low=0, high=50, size=(5_000, 100))
-    
+
     index = create_index(
         distance_type="angular",
         dim=dataset_to_index.shape[1],
@@ -166,8 +166,8 @@ def test_flatnav_ip_index_random_dataset():
         queries=queries,
         ground_truth=ground_truth,
     )
-    
-    
+
+
 def test_flatnav_index_with_reordering():
     training_set, queries, ground_truth, _ = get_ann_benchmark_dataset(
         dataset_name="mnist-784-euclidean"
@@ -193,7 +193,7 @@ def test_flatnav_index_with_reordering():
         assert_recall_threshold=True,
         recall_threshold=0.97,
         use_reordering=True,
-        reordering_algorithm="gorder"
+        reordering_algorithm="gorder",
     )
 
 

quantization/ProductQuantization.h

@@ -14,7 +14,11 @@
 #include <flatnav/distances/InnerProductDistance.h>
 #include <flatnav/distances/SquaredL2Distance.h>
 #include <memory>
+
+#ifdef _OPENMP
 #include <omp.h>
+#endif 
+
 #include <quantization/CentroidsGenerator.h>
 #include <quantization/Utils.h>
 #include <queue>