New ISA dispatcher implementation (#133)

This ISA dispatcher allows detecting whether the running machine
supports none AVX, AVX2, or AVX512, and it dispatches to the appropriate
and optimized code path accordingly.

Author: Dian-Lun (Aaron) Lin

CMakeLists.txt

@@ -35,8 +35,10 @@ set(SVS_VERSION_MINOR ${PROJECT_VERSION_MINOR})
 set(SVS_VERSION_PATCH ${PROJECT_VERSION_PATCH})
 
 add_library(${SVS_LIB} INTERFACE)
-set_target_properties(${SVS_LIB} PROPERTIES EXPORT_NAME svs)
-add_library(svs::svs ALIAS ${SVS_LIB})
+add_library(svs_export INTERFACE)
+set_target_properties(svs_export PROPERTIES EXPORT_NAME svs)
+target_link_libraries(svs_export INTERFACE ${SVS_LIB})
+add_library(svs::svs ALIAS svs_export)
 
 # Runtime include directories are established in the installation logic.
 target_include_directories(
@@ -67,6 +69,13 @@ target_compile_options(
 
 include("cmake/options.cmake")
 
+add_library(svs_x86_options_base INTERFACE)
+add_library(svs::x86_options_base ALIAS svs_x86_options_base)
+if(CMAKE_SYSTEM_PROCESSOR MATCHES "(x86)|(X86)|(amd64)|(AMD64)")
+    target_compile_options(svs_x86_options_base INTERFACE -march=nehalem -mtune=nehalem)
+    include("cmake/multi-arch.cmake")
+endif()
+
 include("cmake/clang-tidy.cmake")
 include("cmake/eve.cmake")
 include("cmake/pthread.cmake")
@@ -112,7 +121,7 @@ set(LIB_CONFIG_INSTALL_DIR "${CMAKE_INSTALL_LIBDIR}/cmake/svs")
 
 # Install headers and target information.
 install(
-    TARGETS svs_devel svs_compile_options svs_native_options
+    TARGETS svs_export svs_devel svs_compile_options svs_native_options
     EXPORT svs-targets
     INCLUDES DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}
 )

benchmark/CMakeLists.txt

@@ -102,9 +102,9 @@ endif()
 target_link_libraries(
     svs_benchmark_library
     PUBLIC
-        ${SVS_LIB}
+        svs::svs
         svs_compile_options
-        svs_native_options
+        svs_x86_options_base
         fmt::fmt
 )
 

bindings/python/CMakeLists.txt

@@ -24,68 +24,9 @@ FetchContent_Declare(
 )
 FetchContent_MakeAvailable(pybind11)
 
-# Try to find the Python executable.
-#
-# If it's given as part of the Cmake arguments given by "scikit build", then use that.
-# Otherwise, fall back to using plain old "python".
-# If *THAT* doesn't work, give up.
-if(DEFINED PYTHON_EXECUTABLE)
-    set(SVS_PYTHON_EXECUTABLE "${PYTHON_EXECUTABLE}")
-else()
-    set(SVS_PYTHON_EXECUTABLE "python")
-endif()
-
-# The micro architectures to compile for.
-if(NOT DEFINED SVS_MICROARCHS)
-    set(SVS_MICROARCHS native)
-endif()
-
 # Include the SVS library directly.
 add_subdirectory("../.." "${CMAKE_CURRENT_BINARY_DIR}/svs")
 
-# Run the python script to get optimization flags for the desired back-ends.
-#
-# FLAGS_SCRIPT - Path to the Python script that will take the compiler, compiler version,
-#   and list of desired microarchitectures and generate optimization flags for each
-#   microarchitecture.
-#
-# FLAGS_TEXT_FILE - List of optimization flags for each architecture.
-#   Expected format:
-#       -march=arch1,-mtune=arch1
-#       -march=arch2,-mtune=arch2
-#       ...
-#       -march=archN,-mtune=archN
-#
-#   The number of lines should be equal to the number of microarchitectures.
-#   NOTE: The entries within each line are separated by a comma on purpose to allow CMake
-#   to read the whole file as a List and then use string replacement on the commas to turn
-#   each line into a list in its own right.
-#
-# TEMP_JSON - JSON Manifest file describing the generated binaries. This is meant to be
-#   included in the Python package to allow the Python code to reason about the packaged
-#   libraries and select the correct one for loading.
-#
-set(FLAGS_SCRIPT "${CMAKE_CURRENT_LIST_DIR}/microarch.py")
-set(FLAGS_TEXT_FILE "${CMAKE_CURRENT_BINARY_DIR}/optimization_flags.txt")
-set(FLAGS_MANIFEST_JSON "${CMAKE_CURRENT_BINARY_DIR}/flags_manifest.json")
-
-execute_process(
-    COMMAND
-        ${SVS_PYTHON_EXECUTABLE}
-        ${FLAGS_SCRIPT}
-        ${FLAGS_TEXT_FILE}
-        ${FLAGS_MANIFEST_JSON}
-        --compiler ${CMAKE_CXX_COMPILER_ID}
-        --compiler-version ${CMAKE_CXX_COMPILER_VERSION}
-        --microarchitectures ${SVS_MICROARCHS}
-    COMMAND_ERROR_IS_FATAL ANY
-)
-
-file(STRINGS "${FLAGS_TEXT_FILE}" OPTIMIZATION_FLAGS)
-message("Flags: ${OPTIMIZATION_FLAGS}")
-list(LENGTH OPTIMIZATION_FLAGS OPT_FLAGS_LENGTH)
-message("Length of flags: ${OPT_FLAGS_LENGTH}")
-
 # C++ files makind up the python bindings.
 set(CPP_FILES
     src/allocator.cpp
@@ -98,50 +39,26 @@ set(CPP_FILES
     src/svs_mkl.cpp
 )
 
-# Generate a shared library for each target microarchitecture.
-foreach(MICRO OPT_FLAGS IN ZIP_LISTS SVS_MICROARCHS OPTIMIZATION_FLAGS)
-    set(LIB_NAME "_svs_${MICRO}")
-
-    pybind11_add_module(${LIB_NAME} MODULE ${CPP_FILES})
-    target_link_libraries(${LIB_NAME} PUBLIC svs::svs)
-    # Dependency "fmt::fmt" obtained from "svs"
-    target_link_libraries(${LIB_NAME} PRIVATE svs::compile_options fmt::fmt)
-
-    string(REPLACE "," ";" OPT_FLAGS ${OPT_FLAGS})
-    message("OPT Flags: ${OPT_FLAGS}")
-    target_compile_options(${LIB_NAME} PRIVATE ${OPT_FLAGS})
-
-    # Header files.
-    target_include_directories(
+set(LIB_NAME "_svs")
+pybind11_add_module(${LIB_NAME} MODULE ${CPP_FILES})
+target_link_libraries(${LIB_NAME} PRIVATE pybind11::module)
+target_link_libraries(${LIB_NAME} PUBLIC svs::svs)
+# Dependency "fmt::fmt" obtained from "svs"
+target_link_libraries(${LIB_NAME} PRIVATE svs::compile_options fmt::fmt svs_x86_options_base)
+target_include_directories(
+    ${LIB_NAME}
+    PUBLIC $<BUILD_INTERFACE:${CMAKE_CURRENT_LIST_DIR}/include>
+)
+if(DEFINED SKBUILD)
+    install(TARGETS ${LIB_NAME} DESTINATION .)
+    set_target_properties(
         ${LIB_NAME}
-        PUBLIC $<BUILD_INTERFACE:${CMAKE_CURRENT_LIST_DIR}/include>
+        PROPERTIES
+        INSTALL_RPATH "$ORIGIN/${CMAKE_INSTALL_LIBDIR}"
     )
-
-    # Comunicate to the C++ library the desired name of the library
-    target_compile_options(${LIB_NAME} PRIVATE "-DSVS_MODULE_NAME=${LIB_NAME}")
-
-    # If scikit build is running the compilation process,
-    if(DEFINED SKBUILD)
-        install(TARGETS ${LIB_NAME} DESTINATION .)
-
-        # The extension module may need to load build or included libraries when loaded.
-
-        # Placing build depedencies in the package and using relative RPATHs that
-        # don't point outside of the package means that the built package is
-        # relocatable. This allows for safe binary redistribution.
-        set_target_properties(
-            ${LIB_NAME}
-            PROPERTIES
-                INSTALL_RPATH "$ORIGIN/${CMAKE_INSTALL_LIBDIR}"
-        )
-    endif()
-endforeach()
+endif()
 
 if(DEFINED SKBUILD)
-    # Install the manifest JSON file.
-    # This is kind of a hack to avoid the needing to explicitly move JSON file into the
-    # source folder of the python library.
-    install(FILES ${FLAGS_MANIFEST_JSON} DESTINATION .)
 
     # Install header files.
     install(

bindings/python/setup.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 
 from skbuild import setup
-import archspec.cpu as cpu
 import os
 
 # If building in a cibuildwheel context, compile multiple versions of the library for
@@ -25,27 +24,6 @@
     "-DCMAKE_EXPORT_COMPILE_COMMANDS=YES",
 ]
 
-# Utility to convert micro-architecture strings to
-def target(arch):
-    return cpu.TARGETS[arch]
-
-# N.B.: cibuildwheel must configure the multi-arch environment variable.
-# Also, the micro-architectures defined below should be in order of preference.
-if os.environ.get("SVS_MULTIARCH", None) is not None:
-    svs_microarchs = [
-        "cascadelake",
-        "x86_64_v3", # conservative base CPU for x86 CPUs.
-    ]
-
-    # Add the current host to the list of micro-architecture if it doesn't already exist.
-    last_target = target(svs_microarchs[-1])
-    host_name = cpu.host().name
-    if host_name not in svs_microarchs and target(host_name) < last_target:
-        svs_microarchs.append(host_name)
-
-    cmake_array = ";".join(svs_microarchs)
-    cmake_args.append(f"-DSVS_MICROARCHS={cmake_array}")
-
 # Determine the root of the repository
 base_dir = os.path.relpath(os.path.join(os.path.dirname(__file__), '..', '..'))
 

bindings/python/src/python_bindings.cpp

@@ -44,15 +44,6 @@
 #include <filesystem>
 #include <optional>
 
-// Get the expected name of the library
-// Make sure CMake stays up to date with defining this parameter.
-//
-// The variable allows us to customize the name of the python module to support
-// micro-architecture versioning.
-#if !defined(SVS_MODULE_NAME)
-#define SVS_MODULE_NAME _svs_native
-#endif
-
 namespace py = pybind11;
 
 namespace {
@@ -144,7 +135,7 @@ class ScopedModuleNameOverride {
 
 } // namespace
 
-PYBIND11_MODULE(SVS_MODULE_NAME, m) {
+PYBIND11_MODULE(_svs, m) {
     // Internall, the top level `__init__.py` imports everything from the C++ module named
     // `_svs`.
     //

bindings/python/src/svs/__init__.py

@@ -12,13 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Dynamic loading logic.
-from .loader import library, current_backend, available_backends
+import svs._svs as _svs
 
-# Reexport all public functions and structs from the inner module.
-lib = library()
 globals().update(
-    {k : v for (k, v) in lib.__dict__.items() if not k.startswith("__")}
+    {k : v for (k, v) in _svs.__dict__.items() if not k.startswith("__")}
 )
 
 # Misc types and functions

bindings/python/src/svs/common.py

@@ -18,8 +18,7 @@
 import os
 
 import numpy as np
-from .loader import library
-lib = library()
+import svs._svs as _svs
 
 def np_to_svs(nptype):
     """
@@ -34,29 +33,29 @@ def np_to_svs(nptype):
     """
     # Signed
     if nptype == np.int8:
-        return lib.int8
+        return _svs.int8
     if nptype == np.int16:
-        return lib.int16
+        return _svs.int16
     if nptype == np.int32:
-        return lib.int32
+        return _svs.int32
     if nptype == np.int64:
-        return lib.int64
+        return _svs.int64
     # Unsigned
     if nptype == np.uint8:
-        return lib.uint8
+        return _svs.uint8
     if nptype == np.uint16:
-        return lib.uint16
+        return _svs.uint16
     if nptype == np.uint32:
-        return lib.uint32
+        return _svs.uint32
     if nptype == np.uint64:
-        return lib.uint64
+        return _svs.uint64
     # Float
     if nptype == np.float16:
-        return lib.float16
+        return _svs.float16
     if nptype == np.float32:
-        return lib.float32
+        return _svs.float32
     if nptype == np.float64:
-        return lib.float64
+        return _svs.float64
 
     raise Exception(f"Could not convert {nptype} to a svs.DataType enum!");
 
@@ -221,7 +220,7 @@ def generate_test_dataset(
         query_seed = None,
         num_threads = 1,
         num_neighbors: int = 100,
-        distance = lib.DistanceType.L2
+        distance = _svs.DistanceType.L2
         ):
     """
     Generate a sample dataset consisting of the base data, queries, and groundtruth all in
@@ -259,7 +258,7 @@ def generate_test_dataset(
     write_vecs(queries, os.path.join(directory, "queries.fvecs"))
 
     print("Generating Groundtruth")
-    index = lib.Flat(data, distance, num_threads = num_threads)
+    index = _svs.Flat(data, distance, num_threads = num_threads)
     I, _ = index.search(queries, num_neighbors)
     write_vecs(I.astype(np.uint32), os.path.join(directory, "groundtruth.ivecs"))