Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions MANIFEST.in
Original file line number Diff line number Diff line change
@@ -1 +1,4 @@
include python/usearch/py.typed
recursive-include numkong/include *.h *.hpp
recursive-include numkong/c *.c *.h
recursive-include stringzilla/include *.h *.hpp
32 changes: 32 additions & 0 deletions python/numkong_bundle.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,32 @@
/**
* @file python/numkong_bundle.cpp
* @brief A single-file C++ wrapper for NumKong C files to ensure static linkage on Android.
*/

extern "C" {
#include "../numkong/c/numkong.c"
#include "../numkong/c/dispatch_bf16.c"
#include "../numkong/c/dispatch_bf16c.c"
#include "../numkong/c/dispatch_e2m3.c"
#include "../numkong/c/dispatch_e3m2.c"
#include "../numkong/c/dispatch_e4m3.c"
#include "../numkong/c/dispatch_e5m2.c"
#include "../numkong/c/dispatch_f16.c"
#include "../numkong/c/dispatch_f16c.c"
#include "../numkong/c/dispatch_f32.c"
#include "../numkong/c/dispatch_f32c.c"
#include "../numkong/c/dispatch_f64.c"
#include "../numkong/c/dispatch_f64c.c"
#include "../numkong/c/dispatch_i16.c"
#include "../numkong/c/dispatch_i32.c"
#include "../numkong/c/dispatch_i4.c"
#include "../numkong/c/dispatch_i64.c"
#include "../numkong/c/dispatch_i8.c"
#include "../numkong/c/dispatch_other.c"
#include "../numkong/c/dispatch_u1.c"
#include "../numkong/c/dispatch_u16.c"
#include "../numkong/c/dispatch_u32.c"
#include "../numkong/c/dispatch_u4.c"
#include "../numkong/c/dispatch_u64.c"
#include "../numkong/c/dispatch_u8.c"
}
4 changes: 4 additions & 0 deletions python/usearch/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,10 @@
# any bundled sibling DLLs, then load without `RTLD_GLOBAL`.
os.add_dll_directory(os.path.dirname(_numkong_path))
numkong_lib = ctypes.CDLL(_numkong_path)
elif sys.platform == "android":
# On Android, NumKong is statically bundled into the USearch extension
# because the Bionic linker does not support `RTLD_GLOBAL`.
pass
else:
# On Linux/macOS we need `RTLD_GLOBAL` so USearch's compiled module can
# resolve NumKong symbols at its own load time.
Expand Down
41 changes: 22 additions & 19 deletions setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ def get_bool_env_w_name(name: str, preference: bool) -> tuple:


# Check the environment variables
is_android: bool = sys.platform == "android"
is_linux: bool = sys.platform == "linux"
is_macos: bool = sys.platform == "darwin"
is_windows: bool = sys.platform == "win32"
Expand All @@ -28,7 +29,7 @@ def get_bool_env_w_name(name: str, preference: bool) -> tuple:

is_gcc = False
is_clang = False
if is_linux:
if is_linux or is_android:
cxx = os.environ.get("CXX")
if cxx:
try:
Expand All @@ -45,11 +46,16 @@ def get_bool_env_w_name(name: str, preference: bool) -> tuple:
# ? Using `ctypes.CDLL(numkong.__file__)` breaks the CI
# ? with "Windows fatal exception: access violation".
prefer_numkong: bool = not is_windows
prefer_openmp: bool = is_linux and is_gcc
prefer_openmp: bool = (is_linux or is_android) and is_gcc

use_numkong: bool = get_bool_env("USEARCH_USE_NUMKONG", prefer_numkong)
use_openmp: bool = get_bool_env("USEARCH_USE_OPENMP", prefer_openmp)

# Dynamic LOADING of a separate NumKong library often fails at runtime on Android
# due to the Bionic linker's behavior with RTLD_GLOBAL and Python extensions.
# In such cases, we bundle NumKong source files directly into the extension.
use_numkong_bundle: bool = use_numkong and is_android

# Common arguments for all platforms
macros_args.append(("USEARCH_USE_OPENMP", "1" if use_openmp else "0"))
macros_args.append(("USEARCH_USE_NUMKONG", "1" if use_numkong else "0"))
Expand All @@ -61,7 +67,7 @@ def get_bool_env_w_name(name: str, preference: bool) -> tuple:
macros_args.extend(
[
("NK_DYNAMIC_DISPATCH", "1" if use_numkong else "0"),
("NK_TARGET_NEON", "0"), # ? Hide-out all complex intrinsics
("NK_TARGET_NEON", "1" if is_android and ("arm" in machine or "aarch64" in machine) else "0"),
("NK_TARGET_NEONBFDOT", "0"), # ? Hide-out all complex intrinsics
("NK_TARGET_NEONHALF", "0"), # ? Hide-out all complex intrinsics
("NK_TARGET_NEONSDOT", "0"), # ? Hide-out all complex intrinsics
Expand Down Expand Up @@ -98,7 +104,7 @@ def get_bool_env_w_name(name: str, preference: bool) -> tuple:
]
)

if is_linux:
if is_linux or is_android:
compile_args.append("-std=c++17")
compile_args.append("-O3") # Maximize performance
compile_args.append("-ffast-math") # Maximize floating-point performance
Expand All @@ -110,7 +116,10 @@ def get_bool_env_w_name(name: str, preference: bool) -> tuple:

# Linking to NumKong
compile_args.append("-Wl,--unresolved-symbols=ignore-in-shared-libs")
link_args.append("-static-libstdc++")

# On Android we don't usually want to link libstdc++ statically
if is_linux:
link_args.append("-static-libstdc++")

if use_openmp:
compile_args.append("-fopenmp")
Expand All @@ -133,26 +142,17 @@ def get_bool_env_w_name(name: str, preference: bool) -> tuple:
link_args.append("-undefined")
link_args.append("dynamic_lookup")

# Linking OpenMP requires additional preparation in CIBuildWheel.
# We must install `brew install llvm` ahead of time.
# import subprocess as cli
# llvm_base = cli.check_output(["brew", "--prefix", "llvm"]).strip().decode("utf-8")
# if len(llvm_base):
# compile_args.append(f"-I{llvm_base}/include")
# compile_args.append("-Xpreprocessor -fopenmp")
# link_args.append(f"-L{llvm_base}/lib")
# link_args.append("-lomp")
# macros_args.append(("USEARCH_USE_OPENMP", "1"))

if is_windows:
compile_args.append("/std:c++17")
compile_args.append("/O2")
compile_args.append("/fp:fast") # Enable fast math for MSVC
compile_args.append("/O2") # Maximize performance
compile_args.append("/fp:fast") # Maximize floating-point performance
compile_args.append("/W1") # Reduce warnings verbosity
link_args.append("/FORCE") # Force linking with missing NumKong symbols


sources = ["python/lib.cpp"]
if use_numkong_bundle:
sources.append("python/numkong_bundle.cpp")

ext_modules = [
Pybind11Extension(
Expand Down Expand Up @@ -184,7 +184,10 @@ def get_bool_env_w_name(name: str, preference: bool) -> tuple:
]
if use_numkong:
include_dirs.append("numkong/include")
install_requires.append("numkong")
if use_numkong_bundle:
include_dirs.append("numkong/c")
else:
install_requires.append("numkong")


# With Clang, `setuptools` doesn't properly use the `language="c++"` argument we pass.
Expand Down