Python binding SIMD vs Scalar dynamic dispatch

Author: ThePituLegend

bindings/CMakeLists.txt

@@ -7,13 +7,37 @@ target_link_libraries(quickedcpp PRIVATE quicked)
 find_package(Python3 REQUIRED COMPONENTS Interpreter Development.Module)
 add_subdirectory(python/pybind11)
 
-pybind11_add_module(pyquicked python/quicked.cpp cpp/quicked.cpp)
-set_target_properties(pyquicked PROPERTIES OUTPUT_NAME quicked ARCHIVE_OUTPUT_NAME pyquicked)
-target_include_directories(pyquicked PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/cpp)
-target_link_libraries(pyquicked PRIVATE quicked)
+pybind11_add_module(pyquicked_scalar python/quicked.cpp cpp/quicked.cpp)
+target_compile_definitions(pyquicked_scalar PRIVATE PY_MODULE_NAME=_quicked_scalar)
+set_target_properties(pyquicked_scalar PROPERTIES OUTPUT_NAME _quicked_scalar ARCHIVE_OUTPUT_NAME pyquicked_scalar)
+target_include_directories(pyquicked_scalar PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/cpp)
+target_link_libraries(pyquicked_scalar PRIVATE quicked_scalar)
 if (NOT MSVC)
-    target_compile_options(pyquicked PRIVATE -fvisibility=hidden)
+    target_compile_options(pyquicked_scalar PRIVATE -fvisibility=hidden)
+    target_compile_options(pyquicked_scalar PRIVATE -mno-sse4.1 -mno-avx2) # Disable SSE4.1 and AVX2 for the scalar Python binding
+else()
+    target_compile_options(pyquicked_scalar PRIVATE /arch:IA32) # Disable SSE4.1 and AVX2 for the scalar Python binding
+    target_compile_definitions(pyquicked_scalar PRIVATE _M_IX86)
 endif()
 
-install(TARGETS pyquicked DESTINATION . COMPONENT pyquicked EXCLUDE_FROM_ALL)
-install(FILES python/quicked.pyi DESTINATION . COMPONENT pyquicked EXCLUDE_FROM_ALL)
+
+pybind11_add_module(pyquicked_simd python/quicked.cpp cpp/quicked.cpp)
+target_compile_definitions(pyquicked_simd PRIVATE PY_MODULE_NAME=_quicked_simd)
+set_target_properties(pyquicked_simd PROPERTIES OUTPUT_NAME _quicked_simd ARCHIVE_OUTPUT_NAME pyquicked_simd)
+target_include_directories(pyquicked_simd PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/cpp)
+target_link_libraries(pyquicked_simd PRIVATE quicked_simd)
+if (NOT MSVC)
+    target_compile_options(pyquicked_simd PRIVATE -fvisibility=hidden)
+    target_compile_options(pyquicked_simd PRIVATE -mno-sse4.1 -mno-avx2) # Enable SSE4.1 and AVX2 for the SIMD Python binding
+else()
+    target_compile_options(pyquicked_simd PRIVATE /arch:AVX) # Enable SSE4.1 and AVX2 for the SIMD Python binding
+endif()
+
+add_custom_target(pyquicked ALL
+    DEPENDS pyquicked_scalar pyquicked_simd
+)
+
+install(TARGETS pyquicked_scalar DESTINATION quicked/backend COMPONENT pyquicked EXCLUDE_FROM_ALL)
+install(TARGETS pyquicked_simd DESTINATION quicked/backend COMPONENT pyquicked EXCLUDE_FROM_ALL)
+install(FILES python/__init__.py DESTINATION quicked COMPONENT pyquicked EXCLUDE_FROM_ALL)
+install(FILES python/quicked.pyi DESTINATION quicked COMPONENT pyquicked EXCLUDE_FROM_ALL)

bindings/python/__init__.py

@@ -0,0 +1,69 @@
+"""
+QuickEd - A high-performance exact sequence alignment based on the bound-and-align paradigm
+===========================================================================================
+
+QuickEd is a high-performance exact sequence alignment based on the bound-and-align paradigm.
+Currently, QuickEd focuses on DNA sequence alignment, using the edit distance (Levenshtein distance) metric.
+
+    >>> pattern = "ACGT"
+    >>> text = "ACTT"
+    >>> aligner = quicked.QuickedAligner()
+    >>> aligner.align(pattern, text)
+    >>> score = aligner.getScore()
+    >>> cigar = aligner.getCigar()
+
+QuickEd is actually a C library, and this package is it’s wrapper for Python.
+Check out QuickEd's GitHub for more code examples and more details on how QuickEd works.
+"""
+def __cpu_supports_simd():
+    try:
+        import cpuinfo
+
+        info = cpuinfo.get_cpu_info()
+        arch = info.get("arch", "").lower()
+        if "x86" not in arch:
+            return False
+        flags = info.get("flags", [])
+        return "avx2" in flags and "sse4_1" in flags
+    except Exception:
+        return False
+
+
+def __load_compiled_backend(name: str):
+    import os
+    import importlib.util
+    import importlib.machinery
+
+    base_dir = os.path.dirname(__file__)
+    backend_dir = os.path.join(base_dir, "backend")
+    for ext in importlib.machinery.EXTENSION_SUFFIXES:
+        candidate = os.path.join(backend_dir, name + ext)
+        if os.path.isfile(candidate):
+            loader = importlib.machinery.ExtensionFileLoader(name, candidate)
+            spec = importlib.util.spec_from_file_location(
+                name, candidate, loader=loader
+            )
+            mod = importlib.util.module_from_spec(spec)
+            spec.loader.exec_module(mod)
+            return mod
+    raise ImportError(f"Cannot find compiled backend: {name}")
+
+
+backend = "_quicked_simd" if __cpu_supports_simd() else "_quicked_scalar"
+"""
+The name of the backend module to load, determined based on CPU SIMD support.
+If the CPU supports AVX2 and SSE4.1, the SIMD-optimized backend ("_quicked_simd") is used.
+Otherwise, the scalar backend ("_quicked_scalar") is selected.
+"""
+
+__backend_module = __load_compiled_backend(backend)
+
+# Import public symbols
+__public_symbols = getattr(
+    __backend_module,
+    "__all__",
+    [k for k in dir(__backend_module) if not k.startswith("_")],
+)
+globals().update({k: getattr(__backend_module, k) for k in __public_symbols})
+__all__ = __public_symbols + ["backend"]
+

bindings/python/quicked.cpp

@@ -28,8 +28,12 @@
 namespace py = pybind11;
 using namespace py::literals;
 
+#ifndef PY_MODULE_NAME
+#define PY_MODULE_NAME _quicked
+#endif
+
 namespace quicked {
-    PYBIND11_MODULE(quicked, m) {
+    PYBIND11_MODULE(PY_MODULE_NAME, m) {
         m.doc() = R"pbdoc(
 QuickEd - A high-performance exact sequence alignment based on the bound-and-align paradigm
 ===========================================================================================

bindings/python/quicked.pyi

@@ -1,72 +1,27 @@
-"""
-
-QuickEd - A high-performance exact sequence alignment based on the bound-and-align paradigm
-===========================================================================================
-
-QuickEd is a high-performance exact sequence alignment based on the bound-and-align paradigm.
-Currently, QuickEd focuses on DNA sequence alignment, using the edit distance (Levenshtein distance) metric.
-
-    >>> pattern = "ACGT"
-    >>> text = "ACTT"
-    >>> aligner = quicked.QuickedAligner()
-    >>> aligner.align(pattern, text)
-    >>> score = aligner.getScore()
-    >>> cigar = aligner.getCigar()
-
-QuickEd is actually a C library, and this package is it’s wrapper for Python.
-Check out QuickEd's GitHub for more code examples and more details on how QuickEd works.
-"""
 from __future__ import annotations
 import typing
 __all__ = ['QuickedAlgo', 'QuickedAligner', 'QuickedException']
 class QuickedAlgo:
     """
-
     QuickedAlgo
     ===========
 
-    Enumeration of alignment algorithms
+    Enumeration of alignment algorithms.
 
-    Members:
-      QUICKED : QuickEd algorithm
-      WINDOWED : WindowEd algorithm
-      BANDED : BandEd algorithm
-      HIRSCHBERG : Hirschberg algorithm
+    Attributes:
+        QUICKED (QuickedAlgo): QuickEd algorithm.
+        WINDOWED (QuickedAlgo): WindowEd algorithm.
+        BANDED (QuickedAlgo): BandEd algorithm.
+        HIRSCHBERG (QuickedAlgo): Hirschberg algorithm.
     """
     BANDED: typing.ClassVar[QuickedAlgo]  # value = <QuickedAlgo.BANDED: 2>
     HIRSCHBERG: typing.ClassVar[QuickedAlgo]  # value = <QuickedAlgo.HIRSCHBERG: 3>
     QUICKED: typing.ClassVar[QuickedAlgo]  # value = <QuickedAlgo.QUICKED: 0>
     WINDOWED: typing.ClassVar[QuickedAlgo]  # value = <QuickedAlgo.WINDOWED: 1>
     __members__: typing.ClassVar[dict[str, QuickedAlgo]]  # value = {'QUICKED': <QuickedAlgo.QUICKED: 0>, 'WINDOWED': <QuickedAlgo.WINDOWED: 1>, 'BANDED': <QuickedAlgo.BANDED: 2>, 'HIRSCHBERG': <QuickedAlgo.HIRSCHBERG: 3>}
-    def __eq__(self, other: typing.Any) -> bool:
-        ...
-    def __getstate__(self) -> int:
-        ...
-    def __hash__(self) -> int:
-        ...
-    def __index__(self) -> int:
-        ...
-    def __init__(self, value: int) -> None:
-        ...
-    def __int__(self) -> int:
-        ...
-    def __ne__(self, other: typing.Any) -> bool:
-        ...
-    def __repr__(self) -> str:
-        ...
-    def __setstate__(self, state: int) -> None:
-        ...
-    def __str__(self) -> str:
-        ...
-    @property
-    def name(self) -> str:
-        ...
-    @property
-    def value(self) -> int:
-        ...
+
 class QuickedAligner:
     """
-
     QuickedAligner
     ==============
 
@@ -120,5 +75,6 @@ class QuickedAligner:
         """
         Sets the window size for WindowEd alignment
         """
+
 class QuickedException(Exception):
     pass

examples/bindings/align_benchmark.py

@@ -3,7 +3,7 @@
 import sys
 import timeit
 
-from quicked import QuickedAlgo, QuickedAligner, QuickedException
+from quicked import QuickedAligner, QuickedException, backend
 
 if len(sys.argv) < 2:
     print("Usage: python3 test.py <input_file>")
@@ -41,6 +41,7 @@ def align_sequences():
             print(f"Error in pair {i // 2 + 1}: {e}")
 
 try:
+    print("Using backend:", backend)
     execution_time = timeit.timeit(align_sequences, number=1)
     print(f"Execution time: {execution_time:.2f} seconds")
 except QuickedException as e:

pyproject.toml

@@ -3,15 +3,15 @@ requires = ["scikit-build-core==0.11.1", "packaging>=24.2"]
 build-backend = "scikit_build_core.build"
 
 [tool.scikit-build]
-build.verbose = true
+cmake.version = "==3.22.1" # Pinning CMake version to avoid conflicts with pinned dependencies
 logging.level = "INFO"
 build.targets = ["pyquicked"]
 install.components = ["pyquicked"]
 
 [tool.cibuildwheel]
 environment = { CMAKE_ARGS = "-DQUICKED_NONATIVE=ON" }
 archs = ["auto64"]
-skip = ["pp*", "musllinux_*"] # Skip PyPy and musllinux builds to cut down on CI time
+skip = ["pp*", "*musllinux*"] # Skip PyPy and musllinux builds to cut down on CI time
 
 [project]
 name = "QuickEd"

quicked/CMakeLists.txt

@@ -34,4 +34,33 @@ install(FILES
     ${PROJECT_SOURCE_DIR}/quicked_utils/include/vector.h
     DESTINATION include/quicked_utils/include
     COMPONENT quicked
-)
+)
+
+# Shim targets for building the Python bindings
+add_library(quicked_scalar OBJECT ${QUICKED_SRCS} ${UTILS_SRCS})
+target_include_directories(quicked_scalar PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/include) # Library includes
+target_include_directories(quicked_scalar PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}) # Library root
+if(NOT (WIN_CLANG OR MSVC))
+    target_link_libraries(quicked_scalar PUBLIC m) # Link math library
+endif()
+if (NOT MSVC)
+    target_compile_options(quicked_scalar PRIVATE -mno-sse4.1 -mno-avx2) # Disable SSE4.1 and AVX2 for the scalar Python binding
+else()
+    target_compile_options(quicked_scalar PRIVATE /arch:IA32) # Disable SSE4.1 and AVX2 for the scalar Python binding
+    target_compile_definitions(quicked_scalar PRIVATE _M_IX86)
+endif()
+set_target_properties(quicked_scalar PROPERTIES EXCLUDE_FROM_ALL TRUE)
+
+
+add_library(quicked_simd OBJECT ${QUICKED_SRCS} ${UTILS_SRCS})
+target_include_directories(quicked_simd PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/include) # Library includes
+target_include_directories(quicked_simd PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}) # Library root
+if(NOT (WIN_CLANG OR MSVC))
+    target_link_libraries(quicked_simd PUBLIC m) # Link math library
+endif()
+if (NOT MSVC)
+    target_compile_options(quicked_simd PRIVATE -msse4.1 -mavx2) # Enable SSE4.1 and AVX2 for the SIMD Python binding
+else()
+    target_compile_options(quicked_simd PRIVATE /arch:AVX) # Enable SSE4.1 and AVX2 for the SIMD Python binding
+endif()
+set_target_properties(quicked_simd PROPERTIES EXCLUDE_FROM_ALL TRUE)