diff --git a/thirdparty/faiss/.github/actions/build_cmake/action.yml b/thirdparty/faiss/.github/actions/build_cmake/action.yml index 556e6becf..b3bf27e0d 100644 --- a/thirdparty/faiss/.github/actions/build_cmake/action.yml +++ b/thirdparty/faiss/.github/actions/build_cmake/action.yml @@ -40,7 +40,7 @@ runs: uses: conda-incubator/setup-miniconda@v3 with: python-version: '3.12' - miniforge-version: latest # ensures conda-forge channel is used. + miniforge-version: latest channels: conda-forge conda-remove-defaults: 'true' # Set to aarch64 if we're on arm64 because there's no miniforge ARM64 package, just aarch64. @@ -50,18 +50,15 @@ runs: if: inputs.metal != 'ON' shell: bash run: | - # initialize Conda - conda config --set solver libmamba # Ensure starting packages are from conda-forge. conda list --show-channel-urls - conda install -y -q "conda<=25.07" echo "$CONDA/bin" >> $GITHUB_PATH conda install -y -q python=3.12 cmake=3.30.4 make=4.2 swig=4.0 "numpy>=2.0,<3.0" scipy=1.16 pytest=7.4 gflags=2.2 setuptools # install base packages for ARM64 if [ "${{ runner.arch }}" = "ARM64" ]; then - conda install -y -q -c conda-forge openblas=0.3.29 gxx_linux-aarch64=14.2 sysroot_linux-aarch64=2.17 + conda install -y -q -c conda-forge openblas=0.3.33 gxx_linux-aarch64=14.2 sysroot_linux-aarch64=2.17 fi # install base packages for X86_64 @@ -79,7 +76,11 @@ runs: conda install -y -q cuda-libraries-dev=12.6 cuda-nvcc=12.6 cuda-nvtx=12.6 cuda-cupti=12.6 cuda-cudart-dev=12.6 gxx_linux-64=12.4 -c "nvidia/label/cuda-12.6" # and CUDA from cuVS channel for cuVS builds elif [ "${{ inputs.cuvs }}" = "ON" ]; then - conda install -y -q libcuvs=26.02 'cuda-version=12.9' cuda-toolkit=12.9 sysroot_linux-64=2.34 -c rapidsai -c rapidsai-nightly -c conda-forge + conda install -y -q libcuvs=26.02 'cuda-version=12.9' sysroot_linux-64=2.34 -c rapidsai -c rapidsai-nightly -c conda-forge + # Clean index cache to prevent sqlite3 "database is locked" errors + # from conda-libmamba-solver's shards cache between consecutive installs. + # See: https://github.com/conda/conda-libmamba-solver/issues/667 + conda clean --index-cache 2>/dev/null || true fi # install SVS runtime for SVS builds @@ -241,6 +242,7 @@ runs: shell: bash run: | conda list --show-channel-urls + conda install -y -q "pip<26" pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm6.1 - name: Python tests (CPU only) if: inputs.gpu == 'OFF' && inputs.metal != 'ON' diff --git a/thirdparty/faiss/.github/actions/build_conda/action.yml b/thirdparty/faiss/.github/actions/build_conda/action.yml index 6406a7fb8..8a0c2b226 100644 --- a/thirdparty/faiss/.github/actions/build_conda/action.yml +++ b/thirdparty/faiss/.github/actions/build_conda/action.yml @@ -30,7 +30,7 @@ runs: uses: conda-incubator/setup-miniconda@v3 with: python-version: '3.12' - miniforge-version: latest # ensures conda-forge channel is used. + miniforge-version: latest channels: conda-forge conda-remove-defaults: 'true' activate-environment: 'base' @@ -81,7 +81,7 @@ runs: working-directory: conda run: | conda list --show-channel-urls - conda build faiss-gpu --variants '{ "cudatoolkit": "${{ inputs.cuda }}" }' \ + conda build faiss-gpu --python 3.12 --variants '{ "cudatoolkit": "${{ inputs.cuda }}" }' \ -c pytorch -c nvidia/label/cuda-${{ inputs.cuda }} -c nvidia - name: Conda build (GPU) w/ anaconda upload if: inputs.label != '' && inputs.cuda != '' && inputs.cuvs == '' @@ -99,7 +99,7 @@ runs: working-directory: conda run: | conda list --show-channel-urls - conda build faiss-gpu-cuvs --variants '{ "cudatoolkit": "${{ inputs.cuda }}" }' \ + conda build faiss-gpu-cuvs --python 3.12 --variants '{ "cudatoolkit": "${{ inputs.cuda }}" }' \ -c pytorch -c rapidsai -c rapidsai-nightly -c conda-forge -c nvidia - name: Conda build (GPU w/ cuVS) w/ anaconda upload if: inputs.label != '' && inputs.cuda != '' && inputs.cuvs != '' diff --git a/thirdparty/faiss/.github/workflows/build-pull-request.yml b/thirdparty/faiss/.github/workflows/build-pull-request.yml index 4b72208f9..43e267fc7 100644 --- a/thirdparty/faiss/.github/workflows/build-pull-request.yml +++ b/thirdparty/faiss/.github/workflows/build-pull-request.yml @@ -302,6 +302,109 @@ jobs: uses: ./.github/actions/build_cmake with: svs: ON + linux-riscv64-DD-cmake: + name: Linux riscv64 Dynamic Dispatch cross-compile (cmake) + needs: linux-x86_64-cmake + runs-on: ubuntu-latest + steps: + - name: Checkout + uses: actions/checkout@v4 + + - name: Set up riscv64 cross-compilation environment + run: | + sudo dpkg --add-architecture riscv64 + + # Restrict existing apt sources to amd64/i386 so that apt does not + # try (and fail) to fetch riscv64 packages from the main Ubuntu + # mirror, which does not carry riscv64. + if [ -f /etc/apt/sources.list.d/ubuntu.sources ]; then + # Ubuntu 24.04+ deb822 format: insert Architectures field. + sudo sed -i '/^Types: deb$/a Architectures: amd64 i386' \ + /etc/apt/sources.list.d/ubuntu.sources + fi + if grep -q '^deb http' /etc/apt/sources.list 2>/dev/null; then + # Ubuntu 22.04 traditional .list format. + sudo sed -i 's|^deb http|deb [arch=amd64,i386] http|g' \ + /etc/apt/sources.list + fi + + # Add Ubuntu Ports as the riscv64 package source. + CODENAME=$(. /etc/os-release && echo "$VERSION_CODENAME") + echo "deb [arch=riscv64] http://ports.ubuntu.com/ubuntu-ports ${CODENAME} main restricted universe" \ + | sudo tee /etc/apt/sources.list.d/riscv64-ports.list + echo "deb [arch=riscv64] http://ports.ubuntu.com/ubuntu-ports ${CODENAME}-updates main restricted universe" \ + | sudo tee -a /etc/apt/sources.list.d/riscv64-ports.list + + sudo apt-get update -qq + sudo apt-get install -y -qq \ + cmake \ + gcc-riscv64-linux-gnu \ + g++-riscv64-linux-gnu \ + libopenblas-dev:riscv64 \ + libgomp1:riscv64 \ + qemu-user-static + + # Bridge Ubuntu's multiarch library directory into the sysroot so + # that CMake's find_library() (ONLY mode, root=/usr/riscv64-linux-gnu) + # can reach packages installed as :riscv64. + # CMake searches ${root}/usr/lib// thanks to the + # compiler's -print-multiarch output (riscv64-linux-gnu). + sudo mkdir -p /usr/riscv64-linux-gnu/usr/lib + sudo ln -sfn /usr/lib/riscv64-linux-gnu \ + /usr/riscv64-linux-gnu/usr/lib/riscv64-linux-gnu + + - name: Setup ccache + uses: hendrikmuhs/ccache-action@v1 + with: + key: linux-riscv64-dd + max-size: 2G + update-package-index: true + + - name: Configure (cmake) + run: | + cmake -B build \ + -DCMAKE_TOOLCHAIN_FILE=cmake/toolchains/riscv64-linux-gnu.cmake \ + -DBUILD_TESTING=ON \ + -DBUILD_SHARED_LIBS=ON \ + -DFAISS_ENABLE_GPU=OFF \ + -DFAISS_OPT_LEVEL=dd \ + -DFAISS_ENABLE_PYTHON=OFF \ + -DFAISS_ENABLE_C_API=OFF \ + -DBLA_VENDOR=OpenBLAS \ + -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_C_COMPILER_LAUNCHER=ccache \ + -DCMAKE_CXX_COMPILER_LAUNCHER=ccache \ + -DCMAKE_BUILD_RPATH=/usr/lib/riscv64-linux-gnu \ + . + + - name: Build faiss_test (includes test_factory_tools) + run: cmake --build build --target faiss_test -j$(nproc) + + - name: Verify binary with ldd (via qemu) + run: | + file build/tests/faiss_test + file build/tests/faiss_test | grep -q "RISC-V" || \ + { echo "ERROR: not a RISC-V binary"; exit 1; } + + # Expose riscv64 libs inside the qemu sysroot so the riscv64 ld.so + # launched by qemu can resolve NEEDED entries. + sudo mkdir -p /usr/riscv64-linux-gnu/usr/lib + sudo ln -sfn /usr/lib/riscv64-linux-gnu \ + /usr/riscv64-linux-gnu/usr/lib/riscv64-linux-gnu + + # Invoke the riscv64 dynamic linker under qemu with --list, which + # is exactly what ldd does on native hardware. + RISCV64_LD=$(find /usr/riscv64-linux-gnu/lib \ + -name "ld-linux-riscv64-*.so.1" | head -1) + qemu-riscv64-static -L /usr/riscv64-linux-gnu \ + "$RISCV64_LD" --list build/tests/faiss_test + + - name: Run test_simd_levels (via QEMU) + run: | + qemu-riscv64-static -L /usr/riscv64-linux-gnu \ + build/tests/faiss_test \ + --gtest_filter="SIMDConfig.*:SIMDLevel.*:CompileOptions.*" + index-io-backward-compatibility: needs: linux-x86_64-cmake name: Index serialization backward compatibility diff --git a/thirdparty/faiss/.github/workflows/index-io-backward-compatibility.yml b/thirdparty/faiss/.github/workflows/index-io-backward-compatibility.yml index a24320e41..864f1e72e 100644 --- a/thirdparty/faiss/.github/workflows/index-io-backward-compatibility.yml +++ b/thirdparty/faiss/.github/workflows/index-io-backward-compatibility.yml @@ -52,7 +52,7 @@ jobs: eval "$(conda shell.bash hook)" conda create -n faiss_conda_read -y python=3.12 conda activate faiss_conda_read - conda install -y -c pytorch -c conda-forge faiss-cpu=1.14.1 + conda install -y -c pytorch -c conda-forge faiss-cpu=1.14.1 "mkl>=2024.2.2,<2026" conda list - name: Run Conda reader (read Faiss index and verify) @@ -89,7 +89,7 @@ jobs: eval "$(conda shell.bash hook)" conda create -n faiss_conda_write -y python=3.12 conda activate faiss_conda_write - conda install -y -c pytorch -c conda-forge faiss-cpu=1.14.1 + conda install -y -c pytorch -c conda-forge faiss-cpu=1.14.1 "mkl>=2024.2.2,<2026" conda list - name: Create shared data directory diff --git a/thirdparty/faiss/CMakeLists.txt b/thirdparty/faiss/CMakeLists.txt index eecc014c3..0eed58166 100644 --- a/thirdparty/faiss/CMakeLists.txt +++ b/thirdparty/faiss/CMakeLists.txt @@ -127,7 +127,7 @@ endif() include(CTest) if(BUILD_TESTING) add_subdirectory(tests) - if(NOT WIN32) + if(NOT WIN32 AND NOT CMAKE_CROSSCOMPILING) add_subdirectory(perf_tests) endif() if(FAISS_ENABLE_GPU) diff --git a/thirdparty/faiss/benchs/bench_fw/optimize.py b/thirdparty/faiss/benchs/bench_fw/optimize.py index 1357d556c..f97f55118 100644 --- a/thirdparty/faiss/benchs/bench_fw/optimize.py +++ b/thirdparty/faiss/benchs/bench_fw/optimize.py @@ -275,7 +275,7 @@ def optimize_codec( pareto_metric=ParetoMetric.TIME_SPACE, ) results = [ - factory[r] for r in set(v["factory"] for _, _, _, k, v in filtered) + factory[r] for r in {v["factory"] for _, _, _, k, v in filtered} ] return results diff --git a/thirdparty/faiss/benchs/bench_super_kmeans.py b/thirdparty/faiss/benchs/bench_super_kmeans.py new file mode 100644 index 000000000..d6ddbb4c3 --- /dev/null +++ b/thirdparty/faiss/benchs/bench_super_kmeans.py @@ -0,0 +1,112 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +"""Benchmark `faiss.SuperKMeans` against vanilla `faiss.Clustering` on the same +input. Inputs may be a .fvecs file or a synthetic Gaussian mixture. + +Examples: + # Synthetic data + python bench_super_kmeans.py --n 10000 --d 128 --k 64 --niter 10 + + # Real .fvecs (e.g., SIFT base) + python bench_super_kmeans.py --fvecs /path/to/sift_base.fvecs --k 1024 \\ + --niter 10 +""" + +import argparse +import sys +import time + +import faiss +import numpy as np +from datasets import fvecs_read + + +def gaussian_mixture(n, d, k, seed): + """k centers in [-1, 1]^d, ~n/k samples per center from N(center, 0.1).""" + rng = np.random.RandomState(seed) + centers = rng.uniform(-1.0, 1.0, size=(k, d)).astype("float32") + cluster = np.arange(n) % k + noise = rng.normal(0.0, 0.1, size=(n, d)).astype("float32") + return centers[cluster] + noise + + +def run_vanilla(x, d, k, niter, seed): + quant = faiss.IndexFlatL2(d) + cl = faiss.Clustering(d, k) + cl.seed = seed + cl.niter = niter + cl.verbose = False + t0 = time.perf_counter() + cl.train(x, quant) + dt = time.perf_counter() - t0 + final_obj = cl.iteration_stats.at(cl.iteration_stats.size() - 1).obj + return dt, final_obj, cl.iteration_stats.size() + + +def run_super(x, d, k, niter, seed): + p = faiss.SuperKMeansParameters() + p.seed = seed + p.niter = niter + p.verbose = False + sc = faiss.SuperKMeans(d, k, p) + t0 = time.perf_counter() + sc.train(x) + dt = time.perf_counter() - t0 + final_obj = sc.iteration_stats.at(sc.iteration_stats.size() - 1).obj + iter_sum = sum( + sc.iteration_stats.at(i).time for i in range(sc.iteration_stats.size()) + ) + setup = dt - iter_sum + rates = faiss.vector_to_array(sc.gemm_pruning_rates) + return dt, final_obj, sc.iteration_stats.size(), iter_sum, setup, rates + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + src = ap.add_mutually_exclusive_group() + src.add_argument("--fvecs", help="path to a .fvecs file") + ap.add_argument("--n", type=int, default=10000, help="synthetic n") + ap.add_argument("--d", type=int, default=128, help="synthetic d") + ap.add_argument("--k", type=int, default=64, help="number of centroids") + ap.add_argument("--niter", type=int, default=10, help="iterations") + ap.add_argument("--seed", type=int, default=42, help="RNG seed") + args = ap.parse_args() + + if args.fvecs: + x = fvecs_read(args.fvecs) + n, d = x.shape + print("=== bench_super_kmeans ===") + print( + f"fvecs={args.fvecs} n={n} d={d} k={args.k} " + f"niter={args.niter} seed={args.seed}" + ) + else: + n, d = args.n, args.d + x = gaussian_mixture(n, d, args.k, args.seed) + print("=== bench_super_kmeans ===") + print( + f"n={n} d={d} k={args.k} niter={args.niter} seed={args.seed}" + ) + + print("\n--- faiss.Clustering ---") + v_dt, v_obj, v_iters = run_vanilla(x, d, args.k, args.niter, args.seed) + print(f" wall_time={v_dt:.3f}s final_obj={v_obj:g} iters={v_iters}") + + print("\n--- faiss.SuperKMeans ---") + s_dt, s_obj, s_iters, iter_sum, setup, rates = run_super( + x, d, args.k, args.niter, args.seed + ) + print( + f" wall_time={s_dt:.3f}s final_obj={s_obj:g} iters={s_iters} " + f"iter_sum={iter_sum:.3f}s setup={setup:.3f}s" + ) + print(" gemm_pruning_rates:", " ".join(f"{r:.3f}" for r in rates)) + print() + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/thirdparty/faiss/cmake/toolchains/riscv64-linux-gnu.cmake b/thirdparty/faiss/cmake/toolchains/riscv64-linux-gnu.cmake new file mode 100644 index 000000000..e2655397e --- /dev/null +++ b/thirdparty/faiss/cmake/toolchains/riscv64-linux-gnu.cmake @@ -0,0 +1,33 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +# Cross-compilation toolchain for RISC-V 64-bit (lp64d ABI) on Ubuntu/Debian. +# +# Requires these packages installed on the build host: +# gcc-riscv64-linux-gnu g++-riscv64-linux-gnu +# +# Target libraries (e.g. libopenblas-dev:riscv64) are installed via apt +# multiarch to /usr/lib/riscv64-linux-gnu/. CMake's ONLY find-root mode +# searches ${CMAKE_FIND_ROOT_PATH}/usr/lib/riscv64-linux-gnu/ (via the +# compiler's multiarch tuple), so the CI script creates a symlink: +# /usr/riscv64-linux-gnu/usr/lib/riscv64-linux-gnu +# -> /usr/lib/riscv64-linux-gnu +# before invoking cmake. + +set(CMAKE_SYSTEM_NAME Linux) +set(CMAKE_SYSTEM_PROCESSOR riscv64) + +set(CMAKE_C_COMPILER riscv64-linux-gnu-gcc) +set(CMAKE_CXX_COMPILER riscv64-linux-gnu-g++) + +# Cross-compiler sysroot provided by gcc-riscv64-linux-gnu. +set(CMAKE_FIND_ROOT_PATH /usr/riscv64-linux-gnu) + +# Never look for host-side tools (cmake, python, …) inside the sysroot. +set(CMAKE_FIND_ROOT_PATH_MODE_PROGRAM NEVER) +# Look for target libraries/headers/packages only inside the sysroot. +set(CMAKE_FIND_ROOT_PATH_MODE_LIBRARY ONLY) +set(CMAKE_FIND_ROOT_PATH_MODE_INCLUDE ONLY) +set(CMAKE_FIND_ROOT_PATH_MODE_PACKAGE ONLY) diff --git a/thirdparty/faiss/conda/faiss-gpu-cuvs/meta.yaml b/thirdparty/faiss/conda/faiss-gpu-cuvs/meta.yaml index 8a9ac120d..1aa8f2f94 100644 --- a/thirdparty/faiss/conda/faiss-gpu-cuvs/meta.yaml +++ b/thirdparty/faiss/conda/faiss-gpu-cuvs/meta.yaml @@ -7,6 +7,7 @@ {% set suffix = "_nightly" if environ.get('PACKAGE_TYPE') == 'nightly' else "" %} {% set number = GIT_DESCRIBE_NUMBER %} {% set cuda_major = (cudatoolkit | default("12.0")).split('.')[0] | int %} +{% set cuda_minor = (cudatoolkit | default("12.0")).split('.')[1] | int %} {% set cuda_constraints=">=13.2,<13.3" %} {% set libcublas_constraints=">=13.3,<13.4" %} # libcublas 13.3 is actually for cuda 13.2 {% set cudart_constraints=">=13.2,<13.3" %} @@ -48,8 +49,8 @@ outputs: - cmake >=3.30.4 - make =4.2 # [not win] - _openmp_mutex =4.5=2_kmp_llvm # [x86_64] - - mkl >=2024.2.2 # [x86_64] - - mkl-devel >=2024.2.2 # [x86_64] + - mkl >=2024.2.2,<2026 # [x86_64] + - mkl-devel >=2024.2.2,<2026 # [x86_64] - cuda-toolkit {{ cudatoolkit }} - cuda-cudart {{ cudart_constraints }} - cuda-cudart-dev {{ cudart_constraints }} @@ -59,15 +60,15 @@ outputs: - cuda-cudart-static_linux-64 {{ cudart_constraints }} # [linux64] host: - _openmp_mutex =4.5=2_kmp_llvm # [x86_64] - - mkl >=2024.2.2 # [x86_64] - - openblas =0.3.32 # [not x86_64] + - mkl >=2024.2.2,<2026 # [x86_64] + - openblas =0.3.33 # [not x86_64] - libcuvs =26.02 - cuda-version {{ cuda_constraints }} - libsvs-runtime =0.3.0 # [x86_64 and linux] run: - _openmp_mutex =4.5=2_kmp_llvm # [x86_64] - - mkl >=2024.2.2 # [x86_64] - - openblas =0.3.32 # [not x86_64] + - mkl >=2024.2.2,<2026 # [x86_64] + - openblas =0.3.33 # [not x86_64] - cuda-cudart {{ cuda_constraints }} - libcublas {{ libcublas_constraints }} - libcuvs =26.02 @@ -98,17 +99,17 @@ outputs: - cmake >=3.30.4 - make =4.2 # [not win] - _openmp_mutex =4.5=2_kmp_llvm # [x86_64] - - mkl >=2024.2.2 # [x86_64] + - mkl >=2024.2.2,<2026 # [x86_64] - cuda-toolkit {{ cudatoolkit }} host: - - mkl >=2024.2.2 # [x86_64] + - mkl >=2024.2.2,<2026 # [x86_64] - _openmp_mutex =4.5=2_kmp_llvm # [x86_64] - python {{ python }} - numpy >=2.0,<2.3 - setuptools - {{ pin_subpackage('libfaiss', exact=True) }} run: - - mkl >=2024.2.2 # [x86_64] + - mkl >=2024.2.2,<2026 # [x86_64] - _openmp_mutex =4.5=2_kmp_llvm # [x86_64] - python {{ python }} - numpy >=2.0,<2.3 @@ -118,18 +119,19 @@ outputs: requires: - numpy >=2.0,<2.3 - scipy -# TODO: remove cuda_major guard when we move to PyPI (pytorch via PyPI works on CUDA 13) -{% if cuda_major < 13 %} +# pytorch-gpu on conda-forge requires CUDA >=12.9; skip for older CUDA 12.x. +# TODO: remove guard when we move to PyPI (pytorch via PyPI works on CUDA 13) +{% if cuda_major == 12 and cuda_minor >= 9 %} - pytorch-gpu >=2.7 {% endif %} commands: - python -X faulthandler -m unittest discover -v -s tests/ -p "test_*" -{% if cuda_major < 13 %} +{% if cuda_major == 12 and cuda_minor >= 9 %} - python -X faulthandler -m unittest discover -v -s tests/ -p "torch_*" {% endif %} - cp tests/common_faiss_tests.py faiss/gpu/test - python -X faulthandler -m unittest discover -v -s faiss/gpu/test/ -p "test_*" -{% if cuda_major < 13 %} +{% if cuda_major == 12 and cuda_minor >= 9 %} - python -X faulthandler -m unittest discover -v -s faiss/gpu/test/ -p "torch_*" {% endif %} - sh test_cpu_dispatch.sh # [linux64] diff --git a/thirdparty/faiss/conda/faiss-gpu/meta.yaml b/thirdparty/faiss/conda/faiss-gpu/meta.yaml index 5a21a3f93..bd8bb194f 100644 --- a/thirdparty/faiss/conda/faiss-gpu/meta.yaml +++ b/thirdparty/faiss/conda/faiss-gpu/meta.yaml @@ -7,6 +7,7 @@ {% set suffix = "_nightly" if environ.get('PACKAGE_TYPE') == 'nightly' else "" %} {% set number = GIT_DESCRIBE_NUMBER %} {% set cuda_major = (cudatoolkit | default("12.0")).split('.')[0] | int %} +{% set cuda_minor = (cudatoolkit | default("12.0")).split('.')[1] | int %} {% if cuda_major >= 13 %} {% set cuda_constraints=">=13.2,<13.3" %} {% set libcublas_constraints=">=13.3,<13.4" %} # libcublas 13.3 is actually for cuda 13.2 @@ -55,15 +56,15 @@ outputs: - cmake >=3.24.0 - make =4.2 # [not win and not (osx and arm64)] - make =4.4 # [osx and arm64] - - mkl-devel >=2024.2.2 # [x86_64] + - mkl-devel >=2024.2.2,<2026 # [x86_64] - cuda-toolkit {{ cudatoolkit }} host: - - mkl >=2024.2.2 # [x86_64] - - openblas =0.3.32 # [not x86_64] + - mkl >=2024.2.2,<2026 # [x86_64] + - openblas =0.3.33 # [not x86_64] - libsvs-runtime =0.3.0 # [x86_64 and linux] run: - - mkl >=2024.2.2 # [x86_64] - - openblas =0.3.32 # [not x86_64] + - mkl >=2024.2.2,<2026 # [x86_64] + - openblas =0.3.33 # [not x86_64] - cuda-cudart {{ cuda_constraints }} - libcublas {{ libcublas_constraints }} - libsvs-runtime =0.3.0 # [x86_64 and linux] @@ -93,16 +94,16 @@ outputs: - make =4.4 # [osx and arm64] - _openmp_mutex =4.5=2_kmp_llvm # [x86_64 and not win] - cuda-toolkit {{ cudatoolkit }} - - mkl-devel >=2024.2.2 # [x86_64] + - mkl-devel >=2024.2.2,<2026 # [x86_64] host: - - mkl >=2024.2.2 # [x86_64] + - mkl >=2024.2.2,<2026 # [x86_64] - python {{ python }} - numpy >=2.0,<3.0 - setuptools - _openmp_mutex =4.5=2_kmp_llvm # [x86_64 and not win] - {{ pin_subpackage('libfaiss', exact=True) }} run: - - mkl >=2024.2.2 # [x86_64] + - mkl >=2024.2.2,<2026 # [x86_64] - python {{ python }} - numpy >=2.0,<3.0 - packaging @@ -111,18 +112,19 @@ outputs: requires: - numpy >=2.0,<3.0 - scipy -# TODO: remove cuda_major guard when we move to PyPI (pytorch via PyPI works on CUDA 13) -{% if cuda_major < 13 %} +# pytorch-gpu on conda-forge requires CUDA >=12.9; skip for older CUDA 12.x. +# TODO: remove guard when we move to PyPI (pytorch via PyPI works on CUDA 13) +{% if cuda_major == 12 and cuda_minor >= 9 %} - pytorch-gpu >=2.7 {% endif %} commands: - python -X faulthandler -m unittest discover -v -s tests/ -p "test_*" -{% if cuda_major < 13 %} +{% if cuda_major == 12 and cuda_minor >= 9 %} - python -X faulthandler -m unittest discover -v -s tests/ -p "torch_*" {% endif %} - cp tests/common_faiss_tests.py faiss/gpu/test - python -X faulthandler -m unittest discover -v -s faiss/gpu/test/ -p "test_*" -{% if cuda_major < 13 %} +{% if cuda_major == 12 and cuda_minor >= 9 %} - python -X faulthandler -m unittest discover -v -s faiss/gpu/test/ -p "torch_*" {% endif %} - sh test_cpu_dispatch.sh # [linux64] diff --git a/thirdparty/faiss/conda/faiss/meta.yaml b/thirdparty/faiss/conda/faiss/meta.yaml index be09a017c..91bd78119 100644 --- a/thirdparty/faiss/conda/faiss/meta.yaml +++ b/thirdparty/faiss/conda/faiss/meta.yaml @@ -45,46 +45,46 @@ outputs: - make =4.2 # [not win and not (osx and arm64)] - make =4.4 # [osx and arm64] {% if PY_VER == '3.10' or PY_VER == '3.11' %} - - mkl-devel >=2024.2.2 # [x86_64] + - mkl-devel >=2024.2.2,<2026 # [x86_64] - python_abi <3.12 {% elif PY_VER == '3.12' %} - - mkl-devel >=2024.2.2 # [x86_64 and not win] - - mkl-devel >=2024.2.2 # [x86_64 and win] + - mkl-devel >=2024.2.2,<2026 # [x86_64 and not win] + - mkl-devel >=2024.2.2,<2026 # [x86_64 and win] - python_abi =3.12 {% else %} - - mkl-devel >=2024.2.2 # [x86_64 and not win] - - mkl-devel >=2024.2.2 # [x86_64 and win] + - mkl-devel >=2024.2.2,<2026 # [x86_64 and not win] + - mkl-devel >=2024.2.2,<2026 # [x86_64 and win] - python_abi {% endif %} - - libopenblas =0.3.32 # [osx] + - libopenblas =0.3.33 # [osx] host: - python {{ python }} - libcxx =20.1.1 # [osx and arm64] - libsvs-runtime =0.3.0 # [x86_64 and linux] {% if PY_VER == '3.10' or PY_VER == '3.11' %} - - mkl >=2024.2.2 # [x86_64] + - mkl >=2024.2.2,<2026 # [x86_64] - python_abi <3.12 {% elif PY_VER == '3.12' %} - - mkl >=2024.2.2 # [x86_64 and not win] - - mkl >=2024.2.2 # [x86_64 and win] + - mkl >=2024.2.2,<2026 # [x86_64 and not win] + - mkl >=2024.2.2,<2026 # [x86_64 and win] - python_abi =3.12 {% endif %} - - openblas =0.3.32 # [not x86_64] - - libopenblas =0.3.32 # [osx] + - openblas =0.3.33 # [not x86_64] + - libopenblas =0.3.33 # [osx] run: - python {{ python }} - libcxx =20.1.1 # [osx and arm64] - libsvs-runtime =0.3.0 # [x86_64 and linux] {% if PY_VER == '3.10' or PY_VER == '3.11' %} - - mkl >=2024.2.2 # [x86_64] + - mkl >=2024.2.2,<2026 # [x86_64] - python_abi <3.12 {% elif PY_VER == '3.12' %} - - mkl >=2024.2.2 # [x86_64 and not win] - - mkl >=2024.2.2 # [x86_64 and win] + - mkl >=2024.2.2,<2026 # [x86_64 and not win] + - mkl >=2024.2.2,<2026 # [x86_64 and win] - python_abi =3.12 {% endif %} - - openblas =0.3.32 # [not x86_64] - - libopenblas =0.3.32 # [osx] + - openblas =0.3.33 # [not x86_64] + - libopenblas =0.3.33 # [osx] test: requires: - conda-build =25.1.2 @@ -113,14 +113,14 @@ outputs: - make =4.4 # [osx and arm64] - _openmp_mutex =4.5=2_kmp_llvm # [x86_64 and not win] {% if PY_VER == '3.10' or PY_VER == '3.11' %} - - mkl >=2024.2.2 # [x86_64] + - mkl >=2024.2.2,<2026 # [x86_64] - python_abi <3.12 {% elif PY_VER == '3.12' %} - - mkl >=2024.2.2 # [x86_64 and not win] - - mkl >=2024.2.2 # [x86_64 and win] + - mkl >=2024.2.2,<2026 # [x86_64 and not win] + - mkl >=2024.2.2,<2026 # [x86_64 and win] - python_abi =3.12 {% endif %} - - libopenblas =0.3.32 # [osx] + - libopenblas =0.3.33 # [osx] host: - python {{ python }} - numpy >=2.0,<3.0 @@ -129,14 +129,14 @@ outputs: - libcxx =20.1.1 # [osx and arm64] - _openmp_mutex =4.5=2_kmp_llvm # [x86_64 and not win] {% if PY_VER == '3.10' or PY_VER == '3.11' %} - - mkl >=2024.2.2 # [x86_64] + - mkl >=2024.2.2,<2026 # [x86_64] - python_abi <3.12 {% elif PY_VER == '3.12' %} - - mkl >=2024.2.2 # [x86_64 and not win] - - mkl >=2024.2.2 # [x86_64 and win] + - mkl >=2024.2.2,<2026 # [x86_64 and not win] + - mkl >=2024.2.2,<2026 # [x86_64 and win] - python_abi =3.12 {% endif %} - - libopenblas =0.3.32 # [osx] + - libopenblas =0.3.33 # [osx] run: - python {{ python }} - numpy >=2.0,<3.0 @@ -144,14 +144,14 @@ outputs: - {{ pin_subpackage('libfaiss', exact=True) }} - libcxx =20.1.1 # [osx and arm64] {% if PY_VER == '3.10' or PY_VER == '3.11' %} - - mkl >=2024.2.2 # [x86_64] + - mkl >=2024.2.2,<2026 # [x86_64] - python_abi <3.12 {% elif PY_VER == '3.12' %} - - mkl >=2024.2.2 # [x86_64 and not win] - - mkl >=2024.2.2 # [x86_64 and win] + - mkl >=2024.2.2,<2026 # [x86_64 and not win] + - mkl >=2024.2.2,<2026 # [x86_64 and win] - python_abi =3.12 {% endif %} - - libopenblas =0.3.32 # [osx] + - libopenblas =0.3.33 # [osx] test: requires: - numpy >=2.0,<3.0 @@ -159,11 +159,11 @@ outputs: - pytorch-cpu >=2.7 # [not win] - pytorch >=2.7 # [win] {% if PY_VER == '3.10' or PY_VER == '3.11' %} - - mkl >=2024.2.2 # [x86_64] + - mkl >=2024.2.2,<2026 # [x86_64] - python_abi <3.12 {% elif PY_VER == '3.12' %} - - mkl >=2024.2.2 # [x86_64 and not win] - - mkl >=2024.2.2 # [x86_64 and win] + - mkl >=2024.2.2,<2026 # [x86_64 and not win] + - mkl >=2024.2.2,<2026 # [x86_64 and win] - python_abi =3.12 {% endif %} commands: diff --git a/thirdparty/faiss/contrib/ivf_tools.py b/thirdparty/faiss/contrib/ivf_tools.py index 2cfec86b5..8d155f7c3 100644 --- a/thirdparty/faiss/contrib/ivf_tools.py +++ b/thirdparty/faiss/contrib/ivf_tools.py @@ -72,7 +72,7 @@ def range_search_preassigned(index_ivf, x, radius, list_nos, coarse_dis=None): # the coarse distances are used in IVFPQ with L2 distance and # by_residual=True otherwise we provide dummy coarse_dis if coarse_dis is None: - coarse_dis = np.empty((n, index_ivf.nprobe), dtype=dis_type) + coarse_dis = np.zeros((n, index_ivf.nprobe), dtype=dis_type) else: assert coarse_dis.shape == (n, index_ivf.nprobe) diff --git a/thirdparty/faiss/faiss/CMakeLists.txt b/thirdparty/faiss/faiss/CMakeLists.txt index b5f406792..c5cb80626 100644 --- a/thirdparty/faiss/faiss/CMakeLists.txt +++ b/thirdparty/faiss/faiss/CMakeLists.txt @@ -51,6 +51,7 @@ set(FAISS_SIMD_SVE_SRC utils/simd_impl/distances_arm_sve.cpp ) set(FAISS_SIMD_RVV_SRC + impl/fast_scan/impl-riscv.cpp impl/pq_code_distance/rvv.cpp impl/scalar_quantizer/sq-rvv.cpp impl/binary_hamming/rvv.cpp @@ -73,6 +74,7 @@ endif() set(FAISS_SRC AutoTune.cpp Clustering.cpp + SuperKMeans.cpp IVFlib.cpp Index.cpp Index2Layer.cpp @@ -200,6 +202,7 @@ endif() set(FAISS_HEADERS AutoTune.h Clustering.h + SuperKMeans.h IVFlib.h Index.h Index2Layer.h diff --git a/thirdparty/faiss/faiss/IndexIDMap.cpp b/thirdparty/faiss/faiss/IndexIDMap.cpp index 5e978ef41..f2817ad4d 100644 --- a/thirdparty/faiss/faiss/IndexIDMap.cpp +++ b/thirdparty/faiss/faiss/IndexIDMap.cpp @@ -196,9 +196,10 @@ void IndexIDMapTemplate::search_ex( } index->search_ex(n, x, numeric_type, k, distances, labels, params); idx_t* li = labels; + const idx_t id_map_size = static_cast(id_map.size()); #pragma omp parallel for for (idx_t i = 0; i < n * k; i++) { - li[i] = li[i] < 0 ? li[i] : id_map[li[i]]; + li[i] = (li[i] < 0 || li[i] >= id_map_size) ? idx_t(-1) : id_map[li[i]]; } } @@ -237,10 +238,12 @@ void IndexIDMapTemplate::range_search( index->range_search(n, x, radius, result); } + const idx_t id_map_size = static_cast(id_map.size()); #pragma omp parallel for for (idx_t i = 0; i < static_cast(result->lims[result->nq]); i++) { - result->labels[i] = result->labels[i] < 0 ? result->labels[i] - : id_map[result->labels[i]]; + const idx_t label = result->labels[i]; + result->labels[i] = + (label < 0 || label >= id_map_size) ? idx_t(-1) : id_map[label]; } } diff --git a/thirdparty/faiss/faiss/IndexIVFPQ.cpp b/thirdparty/faiss/faiss/IndexIVFPQ.cpp index 7b98f3ab1..2f3ee0519 100644 --- a/thirdparty/faiss/faiss/IndexIVFPQ.cpp +++ b/thirdparty/faiss/faiss/IndexIVFPQ.cpp @@ -404,6 +404,9 @@ void initialize_IVFPQ_precomputed_table( return; } + const size_t m_ksub = + mul_no_overflow(pq.M, pq.ksub, "IVFPQ precomputed_table"); + if (use_precomputed_table == 0) { // then choose the type of table if (!(quantizer->metric_type == METRIC_L2 && by_residual)) { if (verbose) { @@ -418,7 +421,10 @@ void initialize_IVFPQ_precomputed_table( if (miq && pq.M % miq->pq.M == 0) { use_precomputed_table = 2; } else { - size_t table_size = pq.M * pq.ksub * nlist * sizeof(float); + size_t table_size = mul_no_overflow( + mul_no_overflow(m_ksub, nlist, "IVFPQ precomputed_table"), + sizeof(float), + "IVFPQ precomputed_table"); if (table_size > precomputed_table_max_bytes) { if (verbose) { printf("IndexIVFPQ::precompute_table: not precomputing table, " @@ -438,7 +444,7 @@ void initialize_IVFPQ_precomputed_table( } // squared norms of the PQ centroids - std::vector r_norms(pq.M * pq.ksub, NAN); + std::vector r_norms(m_ksub, NAN); for (size_t m = 0; m < pq.M; m++) { for (size_t j = 0; j < pq.ksub; j++) { r_norms[m * pq.ksub + j] = @@ -447,15 +453,16 @@ void initialize_IVFPQ_precomputed_table( } if (use_precomputed_table == 1) { - precomputed_table.resize(nlist * pq.M * pq.ksub); + precomputed_table.resize( + mul_no_overflow(nlist, m_ksub, "IVFPQ precomputed_table")); std::vector centroid(d); for (size_t i = 0; i < nlist; i++) { quantizer->reconstruct(i, centroid.data()); - float* tab = &precomputed_table[i * pq.M * pq.ksub]; + float* tab = &precomputed_table[i * m_ksub]; pq.compute_inner_prod_table(centroid.data(), tab); - fvec_madd_dispatch(pq.M * pq.ksub, r_norms.data(), 2.0, tab, tab); + fvec_madd_dispatch(m_ksub, r_norms.data(), 2.0, tab, tab); } } else if (use_precomputed_table == 2) { const MultiIndexQuantizer* miq = @@ -464,7 +471,8 @@ void initialize_IVFPQ_precomputed_table( const ProductQuantizer& cpq = miq->pq; FAISS_THROW_IF_NOT(pq.M % cpq.M == 0); - precomputed_table.resize(cpq.ksub * pq.M * pq.ksub); + precomputed_table.resize( + mul_no_overflow(cpq.ksub, m_ksub, "IVFPQ precomputed_table")); // reorder PQ centroid table std::vector centroids(d * cpq.ksub, NAN); @@ -481,8 +489,8 @@ void initialize_IVFPQ_precomputed_table( cpq.ksub, centroids.data(), precomputed_table.data()); for (size_t i = 0; i < cpq.ksub; i++) { - float* tab = &precomputed_table[i * pq.M * pq.ksub]; - fvec_madd_dispatch(pq.M * pq.ksub, r_norms.data(), 2.0, tab, tab); + float* tab = &precomputed_table[i * m_ksub]; + fvec_madd_dispatch(m_ksub, r_norms.data(), 2.0, tab, tab); } } } diff --git a/thirdparty/faiss/faiss/IndexRaBitQFastScan.h b/thirdparty/faiss/faiss/IndexRaBitQFastScan.h index 449c9a2b1..2aaf0ce81 100644 --- a/thirdparty/faiss/faiss/IndexRaBitQFastScan.h +++ b/thirdparty/faiss/faiss/IndexRaBitQFastScan.h @@ -134,6 +134,7 @@ template < struct RaBitQHeapHandler : simd_result_handlers::ResultHandlerCompare { using RHC = simd_result_handlers::ResultHandlerCompare; + using RHC::handle; using RHC::normalizers; static constexpr SIMDLevel SL256 = simd256_level_selector::value; using simd16uint16 = simd16uint16_tpl; @@ -187,7 +188,7 @@ struct RaBitQHeapHandler } } - void handle(size_t q, size_t b, simd16uint16 d0, simd16uint16 d1) override { + void handle(size_t q, size_t b, simd16uint16 d0, simd16uint16 d1) { ALIGNED(32) uint16_t d32tab[32]; d0.store(d32tab); d1.store(d32tab + 16); diff --git a/thirdparty/faiss/faiss/SuperKMeans.cpp b/thirdparty/faiss/faiss/SuperKMeans.cpp new file mode 100644 index 000000000..ace16aa71 --- /dev/null +++ b/thirdparty/faiss/faiss/SuperKMeans.cpp @@ -0,0 +1,656 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#ifndef FINTEGER +#define FINTEGER long +#endif + +extern "C" { + +/* declare BLAS functions, see http://www.netlib.org/clapack/cblas/ */ + +int sgemm_( + const char* transa, + const char* transb, + FINTEGER* m, + FINTEGER* n, + FINTEGER* k, + const float* alpha, + const float* a, + FINTEGER* lda, + const float* b, + FINTEGER* ldb, + float* beta, + float* c, + FINTEGER* ldc); +} + +namespace faiss { + +namespace { + +struct TrainState { + /// Orthogonal rotation. Train in rotated space (X_tilde = X * R); + /// un-rotate centroids before return. + faiss::RandomRotationMatrix R; + + std::vector X_tilde; // (n, d) row-major + int n = 0; + std::vector Y_tilde; // (k, d) row-major + + std::vector assignments; // size n + std::vector best_dists; // size n; tau per vector + + /// ||X_tilde[i, 0:d_prime]||^2; recomputed when d_prime changes. + std::vector x_norms_partial; + + int d_prime = 0; + + /// ADSampling threshold table; size d+1. + std::vector ad_coeff; + + /// PDX block layout for the trailing pruning sweep: block b covers + /// original dims [true_block_end[b] - block_dim[b], true_block_end[b]). + /// Recomputed when d_prime changes. + std::vector block_dim; + std::vector true_block_end; + + /// Counter for the verbose-mode "low pruning" warning. + int low_pruning_streak = 0; + bool low_pruning_warning_printed = false; + + explicit TrainState(int d) : R(d, d) {} +}; + +/// Rebuild state.block_dim and state.true_block_end from the current +/// state.d_prime and pdx_block_size. Call after any change to d_prime. +void rebuild_pdx_block_layout(int d, int pdx_block_size, TrainState& state) { + const int dp = state.d_prime; + const int d_trail = d - dp; + const int n_full_blocks = d_trail / pdx_block_size; + const int tail = d_trail % pdx_block_size; + const int n_blocks = n_full_blocks + (tail > 0 ? 1 : 0); + state.block_dim.assign(n_blocks, pdx_block_size); + state.true_block_end.resize(n_blocks); + if (n_blocks > 0) { + assert(!state.block_dim.empty()); + assert(!state.true_block_end.empty()); + for (int b = 0; b < n_full_blocks; ++b) { + state.true_block_end[b] = dp + (b + 1) * pdx_block_size; + } + if (tail > 0) { + state.block_dim[n_full_blocks] = tail; + state.true_block_end[n_full_blocks] = d; + } + } +} + +struct IterScratch { + std::vector partial_ip; // (bx_max, by_max) for the GEMM tile + std::vector Y_pdx; // PDX-laid-out trailing block + std::vector Y_trail; // row-major (k, d_trail) input to pdxify + std::vector y_norms_partial; // ||Y_tilde[j, 0:dp]||^2 + std::vector labels64; // size n; widened state.assignments + int prev_d_trail = -1; +}; + +/// Iter 0: full GEMM via knn_L2sqr (vanilla Lloyd's). Fills +/// state.assignments and state.best_dists. Returns objective. +double run_iter0_full_gemm(int d, int k, TrainState& state) { + std::vector labels(state.n); + std::vector distances(state.n); + knn_L2sqr( + state.X_tilde.data(), + state.Y_tilde.data(), + d, + state.n, + k, + /*k=*/1, + distances.data(), + labels.data(), + /*y_norm2=*/nullptr); + + assert(!state.assignments.empty()); + assert(!state.best_dists.empty()); + double objective = 0.0; + for (int i = 0; i < state.n; ++i) { + state.assignments[i] = static_cast(labels[i]); + state.best_dists[i] = distances[i]; + objective += distances[i]; + } + return objective; +} + +/// Iter 1+: partial GEMM over [0, d_prime) + ADSampling progressive +/// pruning over the PDX-laid-out trailing block. Updates +/// state.assignments and state.best_dists. Writes total_pairs and +/// pruned_at_gemm. Returns objective. +double run_iter_pruned( + int d, + int k, + const SuperKMeansParameters& cp, + TrainState& state, + IterScratch& scratch, + int64_t& total_pairs, + int64_t& pruned_at_gemm) { + const int dp = state.d_prime; + assert(dp >= 1); + assert(!state.ad_coeff.empty()); + assert(!scratch.partial_ip.empty()); + assert(!scratch.y_norms_partial.empty()); + const int d_trail = d - dp; + const int n_train = state.n; + assert(static_cast(state.best_dists.size()) >= n_train); + assert(static_cast(state.x_norms_partial.size()) >= n_train); + assert(static_cast(state.assignments.size()) >= n_train); + + if (d_trail != scratch.prev_d_trail) { + scratch.Y_pdx.resize(static_cast(k) * d_trail); + scratch.Y_trail.resize(static_cast(k) * d_trail); + scratch.prev_d_trail = d_trail; + } + for (int j = 0; j < k; ++j) { + std::memcpy( + scratch.Y_trail.data() + static_cast(j) * d_trail, + state.Y_tilde.data() + static_cast(j) * d + dp, + d_trail * sizeof(float)); + } + detail::pdxify( + scratch.Y_trail.data(), + k, + d_trail, + cp.pdx_block_size, + scratch.Y_pdx.data()); + + detail::compute_partial_norms( + state.Y_tilde.data(), k, d, dp, scratch.y_norms_partial.data()); + + const int n_blocks = static_cast(state.block_dim.size()); + + for (int xi = 0; xi < n_train; xi += cp.x_batch) { + const int bx = std::min(cp.x_batch, n_train - xi); + + // Refresh tau: recompute full-d L2 distance to the previously + // assigned centroid. This is intentionally over all d dims (not + // just d_prime) because tau must be an exact distance for the + // chi-squared pruning bound to be valid. Cost is O(bx * d) per + // x-batch, amortized across the y-batch tiles that follow. +#pragma omp parallel for + for (int i = 0; i < bx; ++i) { + const int j_prev = state.assignments[xi + i]; + const float* xrow = + state.X_tilde.data() + static_cast(xi + i) * d; + const float* yrow = + state.Y_tilde.data() + static_cast(j_prev) * d; + float tau = 0.0f; + for (int m = 0; m < d; ++m) { + const float diff = xrow[m] - yrow[m]; + tau += diff * diff; + } + state.best_dists[xi + i] = tau; + } + + for (int yj = 0; yj < k; yj += cp.y_batch) { + const int by = std::min(cp.y_batch, k - yj); + + // GEMM phase: column-major sgemm computes + // partial_ip[i*by + j] = . + { + FINTEGER M = by; + FINTEGER N_ = bx; + FINTEGER K_ = dp; + float alpha = 1.0f; + float beta = 0.0f; + FINTEGER lda_y = d; + FINTEGER lda_x = d; + FINTEGER ldc = by; + sgemm_("Transpose", + "Not transpose", + &M, + &N_, + &K_, + &alpha, + state.Y_tilde.data() + static_cast(yj) * d, + &lda_y, + state.X_tilde.data() + static_cast(xi) * d, + &lda_x, + &beta, + scratch.partial_ip.data(), + &ldc); + } + + // One SIMD dispatch per (xi, yj) tile — block_l2 below is + // a direct call (no per-call switch on SIMDConfig::level). + with_simd_level([&]() { + [[maybe_unused]] const int omp_chunk_local = cp.omp_chunk; + int64_t total_pairs_local = 0; + int64_t pruned_at_gemm_local = 0; +#pragma omp parallel for schedule(dynamic, omp_chunk_local) \ + reduction(+ : total_pairs_local) reduction(+ : pruned_at_gemm_local) + for (int i = 0; i < bx; ++i) { + // tau is the best full-d distance found so far for this + // point; tightened as closer centroids are found. + float tau = state.best_dists[xi + i]; + int best_j = state.assignments[xi + i]; + const float xnp_i = state.x_norms_partial[xi + i]; + const float* xrow = state.X_tilde.data() + + static_cast(xi + i) * d; + + for (int j = 0; j < by; ++j) { + ++total_pairs_local; + + // L2-from-IP; clamp to handle catastrophic + // cancellation when the true distance is ~0. + float pd = xnp_i + scratch.y_norms_partial[yj + j] - + 2.0f * + scratch.partial_ip + [static_cast(i) * by + + j]; + if (pd < 0.0f) { + pd = 0.0f; + } + + if (pd > state.ad_coeff[dp] * tau) { + ++pruned_at_gemm_local; + continue; + } + + // double accumulator mitigates float drift over many + // block additions. + double dist = pd; + bool keep = true; + + // Progressive pruning across PDX blocks. Per block: + // stride = k * block_dim[b] floats, column-major + // across centroids. + size_t pdx_offset = 0; + for (int b = 0; b < n_blocks; ++b) { + const int n_in_block = state.block_dim.at(b); + const int true_end = state.true_block_end.at(b); + const float* xblk = xrow + (true_end - n_in_block); + const float* yblk = scratch.Y_pdx.data() + + pdx_offset + + static_cast(yj + j) * n_in_block; + dist += faiss::detail::block_l2( + xblk, yblk, n_in_block); + pdx_offset += static_cast(k) * n_in_block; + + if (dist > + static_cast(state.ad_coeff[true_end]) * + tau) { + keep = false; + break; + } + } + + if (keep && dist < tau) { + tau = static_cast(dist); + best_j = yj + j; + } + } + + state.best_dists[xi + i] = tau; + state.assignments[xi + i] = best_j; + } + total_pairs += total_pairs_local; + pruned_at_gemm += pruned_at_gemm_local; + }); + } + } + + double objective = 0.0; + for (int i = 0; i < n_train; ++i) { + objective += state.best_dists[i]; + } + return objective; +} + +/// Post-iteration: update centroids and split empties. Returns nsplit. +int update_centroids_and_split( + int d, + int k, + TrainState& state, + IterScratch& scratch, + std::vector& hassign) { + std::fill(hassign.begin(), hassign.end(), 0.0f); + assert(!scratch.labels64.empty()); + assert(!state.assignments.empty()); + for (int i = 0; i < state.n; ++i) { + scratch.labels64[i] = static_cast(state.assignments[i]); + } + detail::compute_centroids( + d, + k, + state.n, + /*k_frozen=*/0, + reinterpret_cast(state.X_tilde.data()), + /*codec=*/nullptr, + scratch.labels64.data(), + /*weights=*/nullptr, + hassign.data(), + state.Y_tilde.data()); + if (state.n <= k) { + return 0; + } + return detail::split_clusters( + d, + k, + state.n, + /*k_frozen=*/0, + hassign.data(), + state.Y_tilde.data()); +} + +/// Stay-in-band controller: nudge state.d_prime based on observed +/// pruning rate. Recomputes x_norms_partial if d_prime changed. Returns +/// the observed pruning rate (0 when there were no pairs). +float adapt_d_prime( + int d, + const SuperKMeansParameters& cp, + TrainState& state, + int64_t total_pairs, + int64_t pruned_at_gemm) { + if (total_pairs == 0) { + return 0.0f; + } + const float pruning_rate = static_cast(pruned_at_gemm) / + static_cast(total_pairs); + int new_dp = state.d_prime; + if (pruning_rate > cp.pruning_target_high) { + new_dp = static_cast( + std::lround(state.d_prime * (1.0f - cp.d_prime_adjust))); + } else if (pruning_rate < cp.pruning_target_low) { + new_dp = static_cast( + std::lround(state.d_prime * (1.0f + cp.d_prime_adjust))); + } + new_dp = std::max(cp.d_prime_min, new_dp); + new_dp = std::min(d / 2, new_dp); + if (new_dp != state.d_prime) { + state.d_prime = new_dp; + detail::compute_partial_norms( + state.X_tilde.data(), + state.n, + d, + state.d_prime, + state.x_norms_partial.data()); + rebuild_pdx_block_layout(d, cp.pdx_block_size, state); + } + return pruning_rate; +} + +/// Pre-loop setup: subsample, rotate, Forgy init, build ADSampling table, +/// allocate scratch. Returned `sampled_x_owner` keeps the subsampled buffer +/// alive when subsampling occurred (otherwise empty). +std::unique_ptr setup_train_state( + TrainState& state, + IterScratch& scratch, + std::vector& hassign, + const SuperKMeansParameters& cp, + int d, + int k, + idx_t n, + const float* x) { + const size_t line_size = sizeof(float) * static_cast(d); + idx_t nx = n; + const uint8_t* x_bytes = reinterpret_cast(x); + std::unique_ptr sampled_x_owner; + if (static_cast(nx) > + static_cast(k) * cp.max_points_per_centroid) { + Clustering tmp_clus(d, k, cp); + uint8_t* x_new = nullptr; + float* w_unused = nullptr; + nx = detail::subsample_training_set( + tmp_clus, + nx, + x_bytes, + line_size, + /*weights=*/nullptr, + &x_new, + &w_unused); + FAISS_ASSERT(x_new != nullptr); + sampled_x_owner.reset(x_new); + x_bytes = x_new; + } + const float* x_sampled = reinterpret_cast(x_bytes); + + FAISS_THROW_IF_NOT_MSG( + nx <= static_cast(std::numeric_limits::max()), + "SuperKMeans: training set size exceeds INT_MAX after sampling"); + state.n = static_cast(nx); + + state.R.init(cp.seed); + + state.X_tilde.resize(static_cast(state.n) * d); + state.R.apply_noalloc(state.n, x_sampled, state.X_tilde.data()); + + // Forgy init: pick k random rows from the rotated pool as initial + // centroids. These remain in rotated space; un-rotation happens + // after the iteration loop. + state.Y_tilde.resize(static_cast(k) * d); + { + std::vector perm(state.n); + rand_perm(perm.data(), state.n, static_cast(cp.seed) + 1); + for (int j = 0; j < k; ++j) { + std::memcpy( + state.Y_tilde.data() + static_cast(j) * d, + state.X_tilde.data() + static_cast(perm[j]) * d, + sizeof(float) * d); + } + } + + state.d_prime = + std::max(cp.d_prime_min, static_cast(d * cp.d_prime_fraction)); + state.d_prime = std::min(state.d_prime, d / 2); + rebuild_pdx_block_layout(d, cp.pdx_block_size, state); + + // Iter 1+ uses L2-from-IP only over [0, d_prime), so full ||X[i]||^2 is + // never read; iter 0 routes through knn_L2sqr which carries its own. + state.x_norms_partial.resize(state.n); + detail::compute_partial_norms( + state.X_tilde.data(), + state.n, + d, + state.d_prime, + state.x_norms_partial.data()); + + const double epsilon = static_cast(cp.ad_epsilon_factor) / d; + state.ad_coeff = detail::precompute_ad_thresholds(d, epsilon); + FAISS_ASSERT_MSG( + state.ad_coeff.size() == static_cast(d + 1), + "ad_coeff size mismatch"); + + state.assignments.assign(state.n, 0); + state.best_dists.assign(state.n, std::numeric_limits::max()); + + hassign.assign(k, 0.0f); + + const int by_max = std::min(cp.y_batch, k); + const int bx_max = std::min(cp.x_batch, state.n); + scratch.partial_ip.resize(static_cast(bx_max) * by_max); + scratch.y_norms_partial.resize(k); + scratch.labels64.resize(state.n); + + return sampled_x_owner; +} + +/// Un-rotate centroids into output buffer. R orthogonal, so +/// reverse_transform applies R^T = R^-1. +void untransform_centroids( + std::vector& centroids, + const RandomRotationMatrix& R, + int d, + int k, + const float* Y_tilde) { + centroids.resize(static_cast(k) * d); + R.reverse_transform(k, Y_tilde, centroids.data()); +} + +} // namespace + +SuperKMeans::SuperKMeans(int d, int k, const SuperKMeansParameters& cp_in) + : cp(cp_in), d(d), k(k) { + FAISS_THROW_IF_NOT_MSG(d > 0, "SuperKMeans: d must be positive"); + FAISS_THROW_IF_NOT_MSG(k > 0, "SuperKMeans: k must be positive"); + FAISS_THROW_IF_NOT_MSG( + cp.d_prime_fraction > 0.0f && cp.d_prime_fraction <= 1.0f, + "SuperKMeans: d_prime_fraction must be in (0, 1]"); + FAISS_THROW_IF_NOT_MSG( + cp.d_prime_adjust >= 0.0f && cp.d_prime_adjust < 1.0f, + "SuperKMeans: d_prime_adjust must be in [0, 1)"); + // d >= 2 * d_prime_min keeps d_prime in the chi-squared validity + // range after both clamping steps (floor at d_prime_min, ceiling + // at d/2). See AdSampling.h on the p >= 16 contract. + FAISS_THROW_IF_NOT_FMT( + d >= 2 * cp.d_prime_min, + "SuperKMeans: d (%d) must be >= 2 * d_prime_min (%d)", + d, + cp.d_prime_min); + FAISS_THROW_IF_NOT_MSG( + cp.d_prime_min >= 16, + "SuperKMeans: d_prime_min must be >= 16 (chi-squared validity floor)"); + FAISS_THROW_IF_NOT_MSG( + cp.pdx_block_size > 0, "SuperKMeans: pdx_block_size must be > 0"); + FAISS_THROW_IF_NOT_MSG(cp.x_batch > 0, "SuperKMeans: x_batch must be > 0"); + FAISS_THROW_IF_NOT_MSG(cp.y_batch > 0, "SuperKMeans: y_batch must be > 0"); + FAISS_THROW_IF_NOT_MSG( + cp.pruning_target_low > 0.0f && + cp.pruning_target_low <= cp.pruning_target_high && + cp.pruning_target_high < 1.0f, + "SuperKMeans: require 0 < pruning_target_low <= pruning_target_high < 1"); + // epsilon = ad_epsilon_factor / d is the chi-squared significance + // level. precompute_ad_thresholds requires epsilon in (0, 1). + FAISS_THROW_IF_NOT_MSG( + cp.ad_epsilon_factor > 0.0f && + cp.ad_epsilon_factor < static_cast(d), + "SuperKMeans: ad_epsilon_factor must be in (0, d) " + "so epsilon = factor/d is in (0,1)"); + FAISS_THROW_IF_NOT_MSG( + cp.omp_chunk > 0, "SuperKMeans: omp_chunk must be > 0"); +} + +void SuperKMeans::train(idx_t n, const float* x) { + FAISS_THROW_IF_NOT_MSG(n > 0, "SuperKMeans: n must be positive"); + FAISS_THROW_IF_NOT_MSG(x != nullptr, "SuperKMeans: x must not be null"); + FAISS_THROW_IF_NOT_MSG( + n >= static_cast(k), "SuperKMeans: n must be >= k"); + if (cp.check_input_data_for_NaNs) { + for (size_t i = 0; i < static_cast(n) * d; i++) { + FAISS_THROW_IF_NOT_MSG( + std::isfinite(x[i]), + "SuperKMeans: input contains NaN's or Inf's"); + } + } + if (cp.verbose && n < static_cast(k) * cp.min_points_per_centroid) { + printf("WARNING: clustering %" PRId64 + " points to %d centroids: please provide at least " + "%" PRId64 " training points\n", + n, + k, + static_cast(k) * cp.min_points_per_centroid); + } + + TrainState state(d); + IterScratch scratch; + std::vector hassign; + [[maybe_unused]] auto sampled_x_owner = + setup_train_state(state, scratch, hassign, cp, d, k, n, x); + + iteration_stats.clear(); + iteration_stats.reserve(cp.niter); + gemm_pruning_rates.clear(); + gemm_pruning_rates.reserve(cp.niter); + + const double t_train_start = getmillisecs(); + + for (int iter = 0; iter < cp.niter; ++iter) { + const double t_iter_start = getmillisecs(); + double objective = 0.0; + int64_t total_pairs = 0; + int64_t pruned_at_gemm = 0; + + if (iter == 0) { + objective = run_iter0_full_gemm(d, k, state); + } else { + objective = run_iter_pruned( + d, k, cp, state, scratch, total_pairs, pruned_at_gemm); + } + + const int nsplit = + update_centroids_and_split(d, k, state, scratch, hassign); + const float pruning_rate = (iter == 0) + ? 0.0f + : adapt_d_prime(d, cp, state, total_pairs, pruned_at_gemm); + + ClusteringIterationStats stat{}; + stat.obj = static_cast(objective); + stat.time = (getmillisecs() - t_iter_start) / 1000.0; + stat.time_search = stat.time; + stat.imbalance_factor = std::numeric_limits::quiet_NaN(); + stat.nsplit = nsplit; + iteration_stats.push_back(stat); + gemm_pruning_rates.push_back(pruning_rate); + + if (iter > 0) { + if (pruning_rate < 0.85f) { + state.low_pruning_streak++; + } else { + state.low_pruning_streak = 0; + } + if (cp.verbose && state.low_pruning_streak >= 3 && + !state.low_pruning_warning_printed) { + fprintf(stderr, + "WARNING: SuperKMeans steady-state pruning < 0.85 for 3+ iters " + "(current=%.2f). Data may not be a good fit for ADSampling; " + "consider falling back to faiss::Clustering.\n", + pruning_rate); + state.low_pruning_warning_printed = true; + } + } + + if (cp.verbose) { + printf(" Iter %d: obj=%g time=%.3fs prune=%.4f dp=%d nsplit=%d\n", + iter, + stat.obj, + stat.time, + pruning_rate, + state.d_prime, + nsplit); + } + } + + if (cp.verbose) { + printf("Total training time: %.3fs\n", + (getmillisecs() - t_train_start) / 1000.0); + } + + untransform_centroids(centroids, state.R, d, k, state.Y_tilde.data()); +} + +} // namespace faiss diff --git a/thirdparty/faiss/faiss/SuperKMeans.h b/thirdparty/faiss/faiss/SuperKMeans.h new file mode 100644 index 000000000..033ec9bba --- /dev/null +++ b/thirdparty/faiss/faiss/SuperKMeans.h @@ -0,0 +1,97 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ + +// SuperKMeans — a faster k-means alternative to faiss::Clustering for IVF +// training and other large-k workloads. +// +// Based on: +// Kuffo, L., Hepkema, S., & Boncz, P. (2026). +// "A Super Fast K-means for Indexing Vector Embeddings." +// arXiv preprint arXiv:2603.20009. +// +// Use when: L2 metric, k >= 1024, d >= 128, dense float embeddings. +// Do not use for: IP/cosine (use Clustering with cp.spherical=true), small k, +// binary data (use IndexBinaryIVF), or near-unit-sphere embeddings with +// k < 4096 (chi-squared assumption breaks down). +// +// `cp`, `d`, `k` are public mutable; mutating them post-construction bypasses +// validation and may produce undefined behavior. + +#pragma once + +#include + +#include +#include + +namespace faiss { + +struct SuperKMeansParameters : public ClusteringParameters { + /// Initial d_prime as a fraction of d (GEMM/PRUNING split). + float d_prime_fraction = 0.125f; + + /// Matches block_l2 SIMD width. + int pdx_block_size = 64; + + /// ADSampling significance: epsilon = ad_epsilon_factor / d. + float ad_epsilon_factor = 1.0f; + + /// Adaptive d_prime stay-in-band controller. Per iteration: + /// pruning > high → shrink d_prime (over-pruning, cheaper threshold) + /// pruning < low → grow d_prime (under-pruning, more GEMM work) + /// else: hold. + float pruning_target_low = 0.95f; + float pruning_target_high = 0.97f; + + /// Relative step size for d_prime adjustments. + float d_prime_adjust = 0.20f; + + /// Floor on d_prime; below this the chi-squared bound is unvalidated. + int d_prime_min = 16; + + int x_batch = 4096; + int y_batch = 1024; + + /// OpenMP dynamic-schedule chunk size for the pruning loop. + int omp_chunk = 8; +}; + +/** Drop-in faster k-means: same interface as faiss::Clustering. Per iteration: + * iter 0: full GEMM over all d dims (vanilla Lloyd's). + * iter 1..niter-1: GEMM over front d_prime dims, then ADSampling + * progressive pruning over PDX-laid trailing dims. + * + * Trains in a randomly-rotated space; centroids are un-rotated before return. + */ +struct SuperKMeans { + SuperKMeansParameters cp; + int d; + int k; + + /// Output: k * d floats, row-major, un-rotated. + std::vector centroids; + + /// Per-iter stats. `obj`, `time`, and `nsplit` are populated faithfully. + /// `time_search` mirrors `time` (phases not timed separately) and + /// `imbalance_factor` is set to NaN (not computed). + std::vector iteration_stats; + + /// Per-iter fraction of (vector, centroid) pairs pruned at the d_prime + /// GEMM-boundary chi-squared check. Per-PDX-block early-exits are NOT + /// counted. iter 0 uses full GEMM, so gemm_pruning_rates[0] == 0.0f. + std::vector gemm_pruning_rates; + + SuperKMeans(int d, int k, const SuperKMeansParameters& cp = {}); + + /// Train on `n` row-major vectors of dimension `d`. Honors the applicable + /// ClusteringParameters fields (niter, seed, verbose, + /// max_points_per_centroid, use_faster_subsampling, + /// check_input_data_for_NaNs). + void train(idx_t n, const float* x); +}; + +} // namespace faiss diff --git a/thirdparty/faiss/faiss/clone_index.cpp b/thirdparty/faiss/faiss/clone_index.cpp index 8eb035f6b..4b3387cc0 100644 --- a/thirdparty/faiss/faiss/clone_index.cpp +++ b/thirdparty/faiss/faiss/clone_index.cpp @@ -377,6 +377,7 @@ Index* Cloner::clone_Index(const Index* index) { IndexRowwiseMinMaxBase* res = clone_IndexRowwiseMinMax(irmmb); res->own_fields = true; res->index = clone_Index(irmmb->index); + return res; } else if ( dynamic_cast(index) || dynamic_cast(index) || diff --git a/thirdparty/faiss/faiss/gpu/test/TestGpuFilterConvert.cu b/thirdparty/faiss/faiss/gpu/test/TestGpuFilterConvert.cu index 7751f28c1..9e7ad6d7b 100644 --- a/thirdparty/faiss/faiss/gpu/test/TestGpuFilterConvert.cu +++ b/thirdparty/faiss/faiss/gpu/test/TestGpuFilterConvert.cu @@ -78,12 +78,12 @@ void run_complex() { auto or_selector = faiss::IDSelectorOr(&range_selector, &array_selector); - auto bitmap_faiss_cpu = std::vector((spec.bitset_len + 8) / 8); + auto bitmap_faiss_cpu = std::vector((spec.bitset_len + 7) / 8); for (indexing_t i = 0; i < bitmap_faiss_cpu.size(); i++) { bitmap_faiss_cpu[i] = (uint8_t)faiss::gpu::randVal(0, 255); } - auto bitmap_selector = - faiss::IDSelectorBitmap(spec.bitset_len, bitmap_faiss_cpu.data()); + auto bitmap_selector = faiss::IDSelectorBitmap( + bitmap_faiss_cpu.size(), bitmap_faiss_cpu.data()); auto not_bitmap_selector = faiss::IDSelectorNot(&bitmap_selector); auto xor_selector = @@ -163,12 +163,12 @@ void run_bitmap() { const raft::device_resources& raft_handle = gpuRes->getRaftHandleCurrentDevice(); // generate random bitmap selector - auto bitmap_faiss_cpu = std::vector((spec.bitset_len + 8) / 8); + auto bitmap_faiss_cpu = std::vector((spec.bitset_len + 7) / 8); for (indexing_t i = 0; i < bitmap_faiss_cpu.size(); i++) { bitmap_faiss_cpu[i] = (uint8_t)faiss::gpu::randVal(0, 255); } - auto bitmap_selector = - faiss::IDSelectorBitmap(spec.bitset_len, bitmap_faiss_cpu.data()); + auto bitmap_selector = faiss::IDSelectorBitmap( + bitmap_faiss_cpu.size(), bitmap_faiss_cpu.data()); auto bitset = cuvs::core::bitset( raft_handle, spec.bitset_len, false); faiss::gpu::convert_to_bitset(gpuRes.get(), bitmap_selector, bitset.view()); @@ -232,6 +232,64 @@ void run_array() { } } +void run_bitmap_byte_convention() { + // Explicitly tests that n = byte count works correctly. + // This is the convention used by the Python wrapper: + // faiss.IDSelectorBitmap(numpy_uint8_array) + // -> IDSelectorBitmap(len(array), array.data()) + using bitset_t = uint32_t; + using indexing_t = int64_t; + + // Use a bit count that is NOT a multiple of 8 to stress edge cases + size_t bit_count = 100003; + size_t byte_count = (bit_count + 7) / 8; // 12501 + + faiss::gpu::StandardGpuResources res; + res.noTempMemory(); + auto gpuRes = res.getResources(); + const raft::device_resources& raft_handle = + gpuRes->getRaftHandleCurrentDevice(); + + // Generate random bitmap + auto bitmap_faiss_cpu = std::vector(byte_count); + for (size_t i = 0; i < bitmap_faiss_cpu.size(); i++) { + bitmap_faiss_cpu[i] = (uint8_t)faiss::gpu::randVal(0, 255); + } + + // Construct with n = byte_count (the only valid convention) + auto bitmap_selector = + faiss::IDSelectorBitmap(byte_count, bitmap_faiss_cpu.data()); + + // Create bitset with the actual number of bits + auto bitset = cuvs::core::bitset( + raft_handle, bit_count, false); + + faiss::gpu::convert_to_bitset(gpuRes.get(), bitmap_selector, bitset.view()); + + // Verify every bit matches + auto bitset_converted_cpu = + raft::make_host_vector(bitset.n_elements()); + raft::copy(raft_handle, bitset_converted_cpu.view(), bitset.to_mdspan()); + raft::resource::sync_stream(raft_handle); + auto bitset_converted_cpu_view = + cuvs::core::bitset_view( + bitset_converted_cpu.data_handle(), bit_count); + for (indexing_t i = 0; i < static_cast(bit_count); i++) { + if (bitset_converted_cpu_view.test(i) != bitmap_selector.is_member(i)) { + ASSERT_TRUE( + testing::AssertionFailure() + << "actual=" << bitset_converted_cpu_view.test(i) + << " != expected=" << bitmap_selector.is_member(i) << " @" + << i << " bit_count: " << bit_count + << " byte_count (n): " << byte_count); + } + } +} + +TEST(TestGpuFilterConvert, BitmapByteConvention) { + run_bitmap_byte_convention(); +} + TEST(TestGpuFilterConvert, Complex) { run_complex(); } diff --git a/thirdparty/faiss/faiss/gpu/test/TestGpuIndexBinaryCagra.cu b/thirdparty/faiss/faiss/gpu/test/TestGpuIndexBinaryCagra.cu index 93aa734a4..6b3c2488f 100644 --- a/thirdparty/faiss/faiss/gpu/test/TestGpuIndexBinaryCagra.cu +++ b/thirdparty/faiss/faiss/gpu/test/TestGpuIndexBinaryCagra.cu @@ -39,9 +39,9 @@ struct Options { Options() { - numTrain = 2 * faiss::gpu::randVal(2000, 5000); - dim = faiss::gpu::randVal(1, 20) * 8; - numAdd = faiss::gpu::randVal(1000, 3000); + numTrain = 2 * faiss::gpu::randVal(1000, 2000); + dim = faiss::gpu::randVal(1, 10) * 8; + numAdd = faiss::gpu::randVal(500, 1500); graphDegree = faiss::gpu::randSelect({32, 64}); intermediateGraphDegree = faiss::gpu::randSelect({64, 98}); @@ -75,7 +75,7 @@ struct Options { void queryTest( faiss::gpu::graph_build_algo build_algo, double expected_recall) { - for (int tries = 0; tries < 5; ++tries) { + for (int tries = 0; tries < 3; ++tries) { Options opt; auto trainVecs = faiss::gpu::randBinaryVecs(opt.numTrain, opt.dim); @@ -190,7 +190,7 @@ TEST(TestGpuIndexBinaryCagra, Query_ITERATIVE_SEARCH) { void copyToTest( faiss::gpu::graph_build_algo build_algo, double expected_recall) { - for (int tries = 0; tries < 5; ++tries) { + for (int tries = 0; tries < 3; ++tries) { Options opt; auto trainVecs = faiss::gpu::randBinaryVecs(opt.numTrain, opt.dim); @@ -324,7 +324,7 @@ TEST(TestGpuIndexBinaryCagra, CopyTo_ITERATIVE_SEARCH) { void copyFromTest( faiss::gpu::graph_build_algo build_algo, double expected_recall) { - for (int tries = 0; tries < 5; ++tries) { + for (int tries = 0; tries < 3; ++tries) { Options opt; auto trainVecs = faiss::gpu::randBinaryVecs(opt.numTrain, opt.dim); diff --git a/thirdparty/faiss/faiss/gpu/test/test_lut_dtype_binding.py b/thirdparty/faiss/faiss/gpu/test/test_lut_dtype_binding.py new file mode 100644 index 000000000..079ff53cb --- /dev/null +++ b/thirdparty/faiss/faiss/gpu/test/test_lut_dtype_binding.py @@ -0,0 +1,60 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import unittest + +import faiss + + +@unittest.skipIf( + "CUVS" not in faiss.get_compile_options(), + "only if cuVS is compiled in") +class TestIVFPQSearchCagraConfigDtypes(unittest.TestCase): + """Regression tests for the SWIG int typemap on cudaDataType_t. + + Before the typemap was added, IVFPQSearchCagraConfig.lut_dtype and + .internal_distance_dtype were exposed as SwigPyObject pointers that + Python users could neither construct nor assign to, and the + CUDA_R_* enum values were not exported. These tests pin the + behavior that they are now plain ints settable from Python. + """ + + def test_constants_exported(self): + self.assertEqual(faiss.CUDA_R_32F, 0) + self.assertEqual(faiss.CUDA_R_64F, 1) + self.assertEqual(faiss.CUDA_R_16F, 2) + self.assertEqual(faiss.CUDA_R_8I, 3) + self.assertEqual(faiss.CUDA_R_8U, 8) + + def test_default_lut_dtype_is_fp32(self): + c = faiss.IVFPQSearchCagraConfig() + self.assertEqual(c.lut_dtype, faiss.CUDA_R_32F) + self.assertEqual(c.internal_distance_dtype, faiss.CUDA_R_32F) + + def test_set_lut_dtype_via_constants(self): + c = faiss.IVFPQSearchCagraConfig() + for value in ( + faiss.CUDA_R_16F, + faiss.CUDA_R_8U, + faiss.CUDA_R_32F, + ): + c.lut_dtype = value + self.assertEqual(c.lut_dtype, value) + + def test_set_lut_dtype_via_raw_int(self): + c = faiss.IVFPQSearchCagraConfig() + c.lut_dtype = 2 # CUDA_R_16F + self.assertEqual(c.lut_dtype, 2) + + def test_dtype_fields_are_independent(self): + c = faiss.IVFPQSearchCagraConfig() + c.lut_dtype = faiss.CUDA_R_8U + c.internal_distance_dtype = faiss.CUDA_R_16F + self.assertEqual(c.lut_dtype, faiss.CUDA_R_8U) + self.assertEqual(c.internal_distance_dtype, faiss.CUDA_R_16F) + + +if __name__ == "__main__": + unittest.main() diff --git a/thirdparty/faiss/faiss/gpu/utils/CuvsFilterConvert.cu b/thirdparty/faiss/faiss/gpu/utils/CuvsFilterConvert.cu index 037ca8600..8a10aee66 100644 --- a/thirdparty/faiss/faiss/gpu/utils/CuvsFilterConvert.cu +++ b/thirdparty/faiss/faiss/gpu/utils/CuvsFilterConvert.cu @@ -145,25 +145,32 @@ void convert_to_bitset_bitmap( raft::resources const& res, const faiss::IDSelectorBitmap& selector, cuvs::core::bitset_view bitset) { - auto n = selector.n; + // IDSelectorBitmap.n is the byte count of the bitmap array (number of + // uint8_t elements). This matches the documented C++ API: + // @param n size of the bitmap array + // is_member(id) requires id / 8 < n + auto n = selector.n; // byte count of bitmap array auto bitset_original_nbits = bitset.get_original_nbits(); if (bitset_original_nbits == 0) { bitset_original_nbits = sizeof(uint32_t) * 8; } + auto bitset_bytes = (bitset.size() + 7) / 8; RAFT_EXPECTS( - bitset.size() == n, - "IDSelectorBitmap is out of range for the given bitset: %ld != %zu", - bitset.size(), - n); + static_cast(n) >= bitset_bytes, + "IDSelectorBitmap is too small for the given bitset: " + "n=%zu bytes, bitset needs %ld bytes (%ld bits)", + n, + bitset_bytes, + bitset.size()); auto stream = raft::resource::get_cuda_stream(res); - auto n_elements = (selector.n + 7) / 8; + auto n_elements = bitset_bytes; auto d_bitmap = raft::make_device_vector(res, n_elements); auto d_bitmap_ptr = d_bitmap.data_handle(); raft::copy( - res, - d_bitmap.view(), - raft::make_host_vector_view( - selector.bitmap, n_elements)); + d_bitmap_ptr, + reinterpret_cast(selector.bitmap), + n_elements, + stream); const int threads_per_block = 256; const int blocks = (n_elements + threads_per_block - 1) / threads_per_block; diff --git a/thirdparty/faiss/faiss/gpu_metal/CMakeLists.txt b/thirdparty/faiss/faiss/gpu_metal/CMakeLists.txt index ee73e87af..656bfe80d 100644 --- a/thirdparty/faiss/faiss/gpu_metal/CMakeLists.txt +++ b/thirdparty/faiss/faiss/gpu_metal/CMakeLists.txt @@ -11,6 +11,7 @@ set(FAISS_METAL_SRC MetalResources.mm MetalIndex.mm + MetalKernels.mm MetalFlatKernels.mm MetalIndexFlat.mm StandardMetalResources.mm @@ -44,3 +45,30 @@ set_target_properties(faiss_metal PROPERTIES OBJCXX_STANDARD 17 OBJCXX_STANDARD_REQUIRED ON ) + +# Pre-compile Metal shaders: .metal -> .air -> .metallib +find_program(XCRUN xcrun REQUIRED) +set(METAL_SOURCE ${CMAKE_CURRENT_SOURCE_DIR}/MetalDistance.metal) +set(METAL_AIR ${CMAKE_CURRENT_BINARY_DIR}/MetalDistance.air) +set(METAL_LIB ${CMAKE_CURRENT_BINARY_DIR}/MetalDistance.metallib) + +add_custom_command( + OUTPUT ${METAL_AIR} + COMMAND ${XCRUN} -sdk macosx metal -c ${METAL_SOURCE} -o ${METAL_AIR} + DEPENDS ${METAL_SOURCE} + COMMENT "Compiling MetalDistance.metal -> .air" +) + +add_custom_command( + OUTPUT ${METAL_LIB} + COMMAND ${XCRUN} -sdk macosx metallib ${METAL_AIR} -o ${METAL_LIB} + DEPENDS ${METAL_AIR} + COMMENT "Linking MetalDistance.air -> .metallib" +) + +add_custom_target(metal_shaders DEPENDS ${METAL_LIB}) +add_dependencies(faiss_metal metal_shaders) + +target_compile_definitions(faiss_metal PRIVATE + FAISS_METALLIB_PATH="${METAL_LIB}" +) diff --git a/thirdparty/faiss/faiss/gpu_metal/MetalDistance.metal b/thirdparty/faiss/faiss/gpu_metal/MetalDistance.metal new file mode 100644 index 000000000..cf9e829e7 --- /dev/null +++ b/thirdparty/faiss/faiss/gpu_metal/MetalDistance.metal @@ -0,0 +1,241 @@ +// @lint-ignore-every LICENSELINT +/** + * Copyright (c) Meta Platforms, Inc. and its affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + * + * Metal Shading Language kernels for distance computation and top-k selection. + * + * Kernel organization: + * - Distance kernels: l2_squared_matrix, ip_matrix (tiled GEMM-style) + * - Top-k selection: topk_threadgroup_K (parallel bitonic sort, K <= 256) + */ + +#include +using namespace metal; + +kernel void l2_squared_matrix( + device const float* queries [[buffer(0)]], + device const float* vectors [[buffer(1)]], + device float* distances [[buffer(2)]], + device const uint* params [[buffer(3)]], + uint2 tgid [[threadgroup_position_in_grid]], + uint2 ltid [[thread_position_in_threadgroup]] +) { + constexpr uint TILE_M = 32; + constexpr uint TILE_N = 32; + constexpr uint TILE_K = 16; + constexpr uint TG_THREADS = 256; + + uint nq = params[0], nb = params[1], d = params[2]; + uint row0 = tgid.y * TILE_M; + uint col0 = tgid.x * TILE_N; + uint ty = ltid.y, tx = ltid.x; + uint tid = ty * 16 + tx; + + float acc00 = 0.0f, acc01 = 0.0f, acc10 = 0.0f, acc11 = 0.0f; + + threadgroup float tgQ[TILE_M * TILE_K]; + threadgroup float tgV[TILE_N * TILE_K]; + + for (uint dk = 0; dk < d; dk += TILE_K) { + uint kLen = min(TILE_K, d - dk); + + for (uint i = tid; i < TILE_M * TILE_K; i += TG_THREADS) { + uint mr = i / TILE_K, mk = i % TILE_K; + uint gRow = row0 + mr; + tgQ[i] = (gRow < nq && mk < kLen) ? queries[gRow * d + dk + mk] : 0.0f; + } + for (uint i = tid; i < TILE_N * TILE_K; i += TG_THREADS) { + uint mr = i / TILE_K, mk = i % TILE_K; + uint gCol = col0 + mr; + tgV[i] = (gCol < nb && mk < kLen) ? vectors[gCol * d + dk + mk] : 0.0f; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint kk = 0; kk < TILE_K; kk++) { + float q0 = tgQ[(ty * 2) * TILE_K + kk]; + float q1 = tgQ[(ty * 2 + 1) * TILE_K + kk]; + float v0 = tgV[(tx * 2) * TILE_K + kk]; + float v1 = tgV[(tx * 2 + 1) * TILE_K + kk]; + float d00 = q0 - v0; acc00 += d00 * d00; + float d01 = q0 - v1; acc01 += d01 * d01; + float d10 = q1 - v0; acc10 += d10 * d10; + float d11 = q1 - v1; acc11 += d11 * d11; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + uint r0 = row0 + ty * 2, r1 = r0 + 1; + uint c0 = col0 + tx * 2, c1 = c0 + 1; + if (r0 < nq && c0 < nb) distances[r0 * nb + c0] = acc00; + if (r0 < nq && c1 < nb) distances[r0 * nb + c1] = acc01; + if (r1 < nq && c0 < nb) distances[r1 * nb + c0] = acc10; + if (r1 < nq && c1 < nb) distances[r1 * nb + c1] = acc11; +} + +kernel void ip_matrix( + device const float* queries [[buffer(0)]], + device const float* vectors [[buffer(1)]], + device float* distances [[buffer(2)]], + device const uint* params [[buffer(3)]], + uint2 tgid [[threadgroup_position_in_grid]], + uint2 ltid [[thread_position_in_threadgroup]] +) { + constexpr uint TILE_M = 32; + constexpr uint TILE_N = 32; + constexpr uint TILE_K = 16; + constexpr uint TG_THREADS = 256; + + uint nq = params[0], nb = params[1], d = params[2]; + uint row0 = tgid.y * TILE_M; + uint col0 = tgid.x * TILE_N; + uint ty = ltid.y, tx = ltid.x; + uint tid = ty * 16 + tx; + + float acc00 = 0.0f, acc01 = 0.0f, acc10 = 0.0f, acc11 = 0.0f; + + threadgroup float tgQ[TILE_M * TILE_K]; + threadgroup float tgV[TILE_N * TILE_K]; + + for (uint dk = 0; dk < d; dk += TILE_K) { + uint kLen = min(TILE_K, d - dk); + + for (uint i = tid; i < TILE_M * TILE_K; i += TG_THREADS) { + uint mr = i / TILE_K, mk = i % TILE_K; + uint gRow = row0 + mr; + tgQ[i] = (gRow < nq && mk < kLen) ? queries[gRow * d + dk + mk] : 0.0f; + } + for (uint i = tid; i < TILE_N * TILE_K; i += TG_THREADS) { + uint mr = i / TILE_K, mk = i % TILE_K; + uint gCol = col0 + mr; + tgV[i] = (gCol < nb && mk < kLen) ? vectors[gCol * d + dk + mk] : 0.0f; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint kk = 0; kk < TILE_K; kk++) { + float q0 = tgQ[(ty * 2) * TILE_K + kk]; + float q1 = tgQ[(ty * 2 + 1) * TILE_K + kk]; + float v0 = tgV[(tx * 2) * TILE_K + kk]; + float v1 = tgV[(tx * 2 + 1) * TILE_K + kk]; + acc00 += q0 * v0; acc01 += q0 * v1; + acc10 += q1 * v0; acc11 += q1 * v1; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + uint r0 = row0 + ty * 2, r1 = r0 + 1; + uint c0 = col0 + tx * 2, c1 = c0 + 1; + if (r0 < nq && c0 < nb) distances[r0 * nb + c0] = acc00; + if (r0 < nq && c1 < nb) distances[r0 * nb + c1] = acc01; + if (r1 < nq && c0 < nb) distances[r1 * nb + c0] = acc10; + if (r1 < nq && c1 < nb) distances[r1 * nb + c1] = acc11; +} + +// ============================================================ +// Parallel threadgroup-based top-k (bitonic sort) +// One threadgroup (256 threads) per query, 4 candidates per thread = 1024. +// ============================================================ + +#define TOPK_THREADGROUP_VARIANT(K) \ +kernel void topk_threadgroup_##K( \ + device const float* distances [[buffer(0)]], \ + device float* outDistances [[buffer(1)]], \ + device int* outIndices [[buffer(2)]], \ + device const uint* params [[buffer(3)]], \ + uint qi [[threadgroup_position_in_grid]], \ + uint tid [[thread_position_in_threadgroup]] \ +) { \ + constexpr uint TG_SIZE = 256; \ + constexpr uint R = 4; \ + constexpr uint CANDIDATES = TG_SIZE * R; \ + threadgroup float tgDist[CANDIDATES]; \ + threadgroup int tgIdx[CANDIDATES]; \ + uint nq = params[0], nb = params[1], k = params[2], want_min = params[3]; \ + if (qi >= nq || k == 0) return; \ + const device float* row = distances + qi * nb; \ + uint kk = min(k, nb); \ + uint K_out = min((uint)K, kk); \ + \ + float localDist[R]; \ + int localIdx[R]; \ + uint localCount = 0; \ + \ + for (uint j = tid; j < nb; j += TG_SIZE) { \ + float v = row[j]; \ + if (localCount < R) { \ + uint pos = localCount; \ + while (pos > 0 && ((want_min && v < localDist[pos-1]) || (!want_min && v > localDist[pos-1]))) { \ + localDist[pos] = localDist[pos-1]; \ + localIdx[pos] = localIdx[pos-1]; \ + pos--; \ + } \ + localDist[pos] = v; \ + localIdx[pos] = (int)j; \ + localCount++; \ + } else { \ + bool better = want_min ? (v < localDist[R-1]) : (v > localDist[R-1]); \ + if (better) { \ + uint pos = R - 1; \ + while (pos > 0 && ((want_min && v < localDist[pos-1]) || (!want_min && v > localDist[pos-1]))) { \ + localDist[pos] = localDist[pos-1]; \ + localIdx[pos] = localIdx[pos-1]; \ + pos--; \ + } \ + localDist[pos] = v; \ + localIdx[pos] = (int)j; \ + } \ + } \ + } \ + \ + for (uint i = 0; i < R; i++) { \ + uint idx = tid * R + i; \ + if (i < localCount) { \ + tgDist[idx] = localDist[i]; \ + tgIdx[idx] = localIdx[i]; \ + } else { \ + tgDist[idx] = want_min ? 1e38f : -1e38f; \ + tgIdx[idx] = -1; \ + } \ + } \ + threadgroup_barrier(mem_flags::mem_threadgroup); \ + \ + for (uint k2 = 2; k2 <= CANDIDATES; k2 *= 2) { \ + for (uint j = k2 >> 1; j > 0; j >>= 1) { \ + for (uint idx = tid; idx < CANDIDATES; idx += TG_SIZE) { \ + uint partner = idx ^ j; \ + if (partner < CANDIDATES && partner > idx) { \ + bool ascending = ((idx & k2) == 0); \ + bool partnerBetter = want_min \ + ? (tgDist[partner] < tgDist[idx] || (tgDist[partner] == tgDist[idx] && tgIdx[partner] < tgIdx[idx])) \ + : (tgDist[partner] > tgDist[idx] || (tgDist[partner] == tgDist[idx] && tgIdx[partner] < tgIdx[idx])); \ + bool idxBetter = want_min \ + ? (tgDist[idx] < tgDist[partner] || (tgDist[idx] == tgDist[partner] && tgIdx[idx] < tgIdx[partner])) \ + : (tgDist[idx] > tgDist[partner] || (tgDist[idx] == tgDist[partner] && tgIdx[idx] < tgIdx[partner])); \ + bool swap = ascending ? partnerBetter : idxBetter; \ + if (swap) { \ + float td = tgDist[idx]; tgDist[idx] = tgDist[partner]; tgDist[partner] = td; \ + int ti = tgIdx[idx]; tgIdx[idx] = tgIdx[partner]; tgIdx[partner] = ti; \ + } \ + } \ + } \ + threadgroup_barrier(mem_flags::mem_threadgroup); \ + } \ + } \ + \ + for (uint i = tid; i < K_out; i += TG_SIZE) { \ + outDistances[qi * k + i] = tgDist[i]; \ + outIndices[qi * k + i] = tgIdx[i]; \ + } \ + for (uint i = tid; i < k - K_out; i += TG_SIZE) { \ + outDistances[qi * k + K_out + i] = want_min ? 1e38f : -1e38f; \ + outIndices[qi * k + K_out + i] = -1; \ + } \ +} + +TOPK_THREADGROUP_VARIANT(32) +TOPK_THREADGROUP_VARIANT(64) +TOPK_THREADGROUP_VARIANT(128) +TOPK_THREADGROUP_VARIANT(256) +#undef TOPK_THREADGROUP_VARIANT diff --git a/thirdparty/faiss/faiss/gpu_metal/MetalFlatKernels.mm b/thirdparty/faiss/faiss/gpu_metal/MetalFlatKernels.mm index 4f3707536..18bbead31 100644 --- a/thirdparty/faiss/faiss/gpu_metal/MetalFlatKernels.mm +++ b/thirdparty/faiss/faiss/gpu_metal/MetalFlatKernels.mm @@ -5,126 +5,17 @@ * This source code is licensed under the MIT license found in the * LICENSE file in the root directory of this source tree. * - * MSL kernels: L2 squared / IP distance matrix, then top-k reduction. + * Flat search dispatch: distance matrix + top-k via MetalKernels. */ #import "MetalFlatKernels.h" -#include +#include +#import "MetalKernels.h" namespace faiss { namespace gpu_metal { -namespace { - -static const char* kMSLSource = R"msl( -#include -using namespace metal; - -kernel void l2_squared_matrix( - device const float* queries [[buffer(0)]], - device const float* vectors [[buffer(1)]], - device float* distances [[buffer(2)]], - device const uint* params [[buffer(3)]], // nq, nb, d - uint2 gid [[thread_position_in_grid]] -) { - uint nq = params[0], nb = params[1], d = params[2]; - uint i = gid.y; - uint j = gid.x; - if (i >= nq || j >= nb) return; - float sum = 0.0f; - for (uint t = 0; t < d; t++) { - float a = queries[i * d + t]; - float b = vectors[j * d + t]; - sum += (a - b) * (a - b); - } - distances[i * nb + j] = sum; -} - -kernel void ip_matrix( - device const float* queries [[buffer(0)]], - device const float* vectors [[buffer(1)]], - device float* distances [[buffer(2)]], - device const uint* params [[buffer(3)]], // nq, nb, d - uint2 gid [[thread_position_in_grid]] -) { - uint nq = params[0], nb = params[1], d = params[2]; - uint i = gid.y; - uint j = gid.x; - if (i >= nq || j >= nb) return; - float sum = 0.0f; - for (uint t = 0; t < d; t++) - sum += queries[i * d + t] * vectors[j * d + t]; - distances[i * nb + j] = sum; -} - -// want_min: 1 = k smallest (L2), 0 = k largest (IP). k <= 256. -kernel void topk( - device const float* distances [[buffer(0)]], - device float* outDistances [[buffer(1)]], - device int* outIndices [[buffer(2)]], - device const uint* params [[buffer(3)]], // nq, nb, k, want_min - uint qi [[thread_position_in_grid]] -) { - uint nq = params[0], nb = params[1], k = params[2], want_min = params[3]; - if (qi >= nq || k == 0) return; - const device float* row = distances + qi * nb; - float bestDist[256]; - int bestIdx[256]; - uint kk = min(k, (uint)256); - uint n = min(kk, nb); - for (uint i = 0; i < n; i++) { - bestDist[i] = row[i]; - bestIdx[i] = (int)i; - } - // Sort first n by distance (ascending for L2/smallest, descending for IP/largest) - for (uint i = 0; i < n; i++) { - for (uint j = i + 1; j < n; j++) { - bool swap = want_min ? (bestDist[j] < bestDist[i]) : (bestDist[j] > bestDist[i]); - if (swap) { - float td = bestDist[i]; bestDist[i] = bestDist[j]; bestDist[j] = td; - int ti = bestIdx[i]; bestIdx[i] = bestIdx[j]; bestIdx[j] = ti; - } - } - } - for (uint i = n; i < kk; i++) { - bestDist[i] = want_min ? 1e38f : -1e38f; - bestIdx[i] = -1; - } - for (uint j = n; j < nb; j++) { - float v = row[j]; - bool insert = want_min ? (v < bestDist[kk-1]) : (v > bestDist[kk-1]); - if (!insert) continue; - uint pos = kk - 1; - if (want_min) { - while (pos > 0 && v < bestDist[pos-1]) { - bestDist[pos] = bestDist[pos-1]; - bestIdx[pos] = bestIdx[pos-1]; - pos--; - } - } else { - while (pos > 0 && v > bestDist[pos-1]) { - bestDist[pos] = bestDist[pos-1]; - bestIdx[pos] = bestIdx[pos-1]; - pos--; - } - } - bestDist[pos] = v; - bestIdx[pos] = (int)j; - } - for (uint i = 0; i < kk; i++) { - outDistances[qi * k + i] = bestDist[i]; - outIndices[qi * k + i] = bestIdx[i]; - } - for (uint i = kk; i < k; i++) { - outDistances[qi * k + i] = want_min ? 1e38f : -1e38f; - outIndices[qi * k + i] = -1; - } -} -)msl"; - -static constexpr int kMaxK = 256; - -} // namespace +static constexpr int kMaxK = MetalKernels::kMaxK; bool runFlatSearchGPU( id device, @@ -146,66 +37,27 @@ bool runFlatSearchGPU( return false; } - NSError* err = nil; - id lib = [device newLibraryWithSource:@(kMSLSource) - options:nil - error:&err]; - if (!lib) { - return false; - } - - id fnDist = [lib - newFunctionWithName:isL2 ? @"l2_squared_matrix" : @"ip_matrix"]; - id fnTopK = [lib newFunctionWithName:@"topk"]; - if (!fnDist || !fnTopK) { - return false; - } - - id psDist = - [device newComputePipelineStateWithFunction:fnDist error:&err]; - id psTopK = - [device newComputePipelineStateWithFunction:fnTopK error:&err]; - if (!psDist || !psTopK) { + MetalKernels& kernels = getMetalKernels(device); + if (!kernels.isValid()) { return false; } - id cmdBuf = [queue commandBuffer]; - id enc = [cmdBuf computeCommandEncoder]; - - // Distance matrix: (nb x nq) threadgroups of (1,1) or tile for better - // occupancy - const NSUInteger w = 16; - const NSUInteger h = 16; - [enc setComputePipelineState:psDist]; - [enc setBuffer:queries offset:0 atIndex:0]; - [enc setBuffer:vectors offset:0 atIndex:1]; id distMatrix = [device newBufferWithLength:(size_t)nq * (size_t)nb * sizeof(float) options:MTLResourceStorageModeShared]; if (!distMatrix) { - [enc endEncoding]; return false; } - [enc setBuffer:distMatrix offset:0 atIndex:2]; - uint32_t distArgs[3] = {(uint32_t)nq, (uint32_t)nb, (uint32_t)d}; - [enc setBytes:distArgs length:sizeof(distArgs) atIndex:3]; - MTLSize tgSize = MTLSizeMake(w, h, 1); - MTLSize gridSize = MTLSizeMake((nb + w - 1) / w, (nq + h - 1) / h, 1); - [enc dispatchThreadgroups:gridSize threadsPerThreadgroup:tgSize]; - // Top-k: nq threads - [enc setComputePipelineState:psTopK]; - [enc setBuffer:distMatrix offset:0 atIndex:0]; - [enc setBuffer:outDistances offset:0 atIndex:1]; - [enc setBuffer:outIndices offset:0 atIndex:2]; - uint32_t topkArgs[4] = { - (uint32_t)nq, (uint32_t)nb, (uint32_t)k, isL2 ? 1u : 0u}; - [enc setBytes:topkArgs - length:sizeof(topkArgs) - atIndex:3]; // nq, nb, k, want_min - MTLSize gridTopK = MTLSizeMake(nq, 1, 1); - [enc dispatchThreadgroups:gridTopK - threadsPerThreadgroup:MTLSizeMake(1, 1, 1)]; + id cmdBuf = [queue commandBuffer]; + id enc = [cmdBuf computeCommandEncoder]; + + MetricType metric = isL2 ? METRIC_L2 : METRIC_INNER_PRODUCT; + kernels.encodeDistanceMatrix( + enc, queries, vectors, distMatrix, nq, nb, d, metric); + + kernels.encodeTopKThreadgroup( + enc, distMatrix, outDistances, outIndices, nq, nb, k, isL2); [enc endEncoding]; [cmdBuf commit]; diff --git a/thirdparty/faiss/faiss/gpu_metal/MetalKernels.h b/thirdparty/faiss/faiss/gpu_metal/MetalKernels.h new file mode 100644 index 000000000..586103b26 --- /dev/null +++ b/thirdparty/faiss/faiss/gpu_metal/MetalKernels.h @@ -0,0 +1,66 @@ +// @lint-ignore-every LICENSELINT +/** + * Copyright (c) Meta Platforms, Inc. and its affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + * + * MetalKernels: typed wrapper around Metal compute kernels. + * Owns library compilation, pipeline caching, and dispatch encoding. + */ + +#pragma once + +#import +#include +#include +#include + +namespace faiss { +namespace gpu_metal { + +class MetalKernels { + public: + explicit MetalKernels(id device); + ~MetalKernels(); + + bool isValid() const; + static constexpr int kMaxK = 256; + + void encodeDistanceMatrix( + id enc, + id queries, + id vectors, + id distances, + int nq, + int nb, + int d, + MetricType metric); + + void encodeTopKThreadgroup( + id enc, + id distances, + id outDist, + id outIdx, + int nq, + int nb, + int k, + bool wantMin); + + static int selectTopKVariantIndex(int k); + + private: + id pipeline(const char* name); + + id device_; + id library_; + std::unordered_map> cache_; + + static constexpr int kTopKVariantSizes[] = {32, 64, 128, 256}; + static constexpr int kNumTopKVariants = 4; +}; + +MetalKernels& getMetalKernels(id device); + +} // namespace gpu_metal +} // namespace faiss diff --git a/thirdparty/faiss/faiss/gpu_metal/MetalKernels.mm b/thirdparty/faiss/faiss/gpu_metal/MetalKernels.mm new file mode 100644 index 000000000..42270d51b --- /dev/null +++ b/thirdparty/faiss/faiss/gpu_metal/MetalKernels.mm @@ -0,0 +1,136 @@ +// @lint-ignore-every LICENSELINT +/** + * Copyright (c) Meta Platforms, Inc. and its affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + * + * MetalKernels implementation: library loading, pipeline caching, and + * encode methods for Metal compute kernels used by the Faiss Metal backend. + */ + +#import "MetalKernels.h" +#include + +#ifndef FAISS_METALLIB_PATH +#error "FAISS_METALLIB_PATH must be defined by CMake" +#endif + +namespace faiss { +namespace gpu_metal { + +namespace { + +static const char* kThreadgroupNames[] = { + "topk_threadgroup_32", + "topk_threadgroup_64", + "topk_threadgroup_128", + "topk_threadgroup_256"}; + +} // namespace + +constexpr int MetalKernels::kTopKVariantSizes[]; + +int MetalKernels::selectTopKVariantIndex(int k) { + for (int i = 0; i < kNumTopKVariants; i++) { + if (k <= kTopKVariantSizes[i]) + return i; + } + return kNumTopKVariants - 1; +} + +MetalKernels::MetalKernels(id device) + : device_(device), library_(nil) { + NSString* path = @(FAISS_METALLIB_PATH); + NSURL* url = [NSURL fileURLWithPath:path]; + NSError* err = nil; + library_ = [device_ newLibraryWithURL:url error:&err]; +} + +MetalKernels::~MetalKernels() = default; + +bool MetalKernels::isValid() const { + return library_ != nil; +} + +id MetalKernels::pipeline(const char* name) { + std::string key(name); + auto it = cache_.find(key); + if (it != cache_.end()) + return it->second; + id fn = [library_ newFunctionWithName:@(name)]; + if (!fn) + return nil; + NSError* err = nil; + id ps = + [device_ newComputePipelineStateWithFunction:fn error:&err]; + if (!ps) + return nil; + cache_[key] = ps; + return ps; +} + +void MetalKernels::encodeDistanceMatrix( + id enc, + id queries, + id vectors, + id distances, + int nq, + int nb, + int d, + MetricType metric) { + FAISS_THROW_IF_NOT_MSG( + metric == METRIC_L2 || metric == METRIC_INNER_PRODUCT, + "Metal distance matrix only supports L2 and inner product"); + const char* kernel = + metric == METRIC_L2 ? "l2_squared_matrix" : "ip_matrix"; + [enc setComputePipelineState:pipeline(kernel)]; + [enc setBuffer:queries offset:0 atIndex:0]; + [enc setBuffer:vectors offset:0 atIndex:1]; + [enc setBuffer:distances offset:0 atIndex:2]; + uint32_t args[3] = {(uint32_t)nq, (uint32_t)nb, (uint32_t)d}; + [enc setBytes:args length:sizeof(args) atIndex:3]; + const NSUInteger tileN = 32; + const NSUInteger tileM = 32; + MTLSize grid = MTLSizeMake( + ((NSUInteger)nb + tileN - 1) / tileN, + ((NSUInteger)nq + tileM - 1) / tileM, + 1); + [enc dispatchThreadgroups:grid + threadsPerThreadgroup:MTLSizeMake(16, 16, 1)]; +} + +void MetalKernels::encodeTopKThreadgroup( + id enc, + id distances, + id outDist, + id outIdx, + int nq, + int nb, + int k, + bool wantMin) { + int vi = selectTopKVariantIndex(k); + [enc setComputePipelineState:pipeline(kThreadgroupNames[vi])]; + [enc setBuffer:distances offset:0 atIndex:0]; + [enc setBuffer:outDist offset:0 atIndex:1]; + [enc setBuffer:outIdx offset:0 atIndex:2]; + uint32_t args[4] = { + (uint32_t)nq, (uint32_t)nb, (uint32_t)k, wantMin ? 1u : 0u}; + [enc setBytes:args length:sizeof(args) atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)nq, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; +} + +MetalKernels& getMetalKernels(id device) { + static std::mutex mu; + static std::unordered_map> map; + uintptr_t key = (uintptr_t)(__bridge void*)device; + std::lock_guard lock(mu); + auto& ptr = map[key]; + if (!ptr) + ptr = std::make_unique(device); + return *ptr; +} + +} // namespace gpu_metal +} // namespace faiss diff --git a/thirdparty/faiss/faiss/impl/ClusteringHelpers.cpp b/thirdparty/faiss/faiss/impl/ClusteringHelpers.cpp index b4704402b..669cdacf4 100644 --- a/thirdparty/faiss/faiss/impl/ClusteringHelpers.cpp +++ b/thirdparty/faiss/faiss/impl/ClusteringHelpers.cpp @@ -78,7 +78,7 @@ idx_t subsample_training_set( "subsample_training_set: perm size %zu < required nx %" PRId64, perm.size(), nx); - assert(!perm.empty()); // pattern for clang-tidy LocalUncheckedArrayBounds + assert(!perm.empty()); uint8_t* x_new = new uint8_t[nx * line_size]; *x_out = x_new; @@ -116,15 +116,6 @@ void compute_centroids( size_t line_size = codec ? codec->sa_code_size() : d * sizeof(float); - // NOTE: FAISS_THROW_IF_NOT inside this OMP region is formally UB per the - // OpenMP spec (exceptions cannot propagate out of a parallel region). - // In practice, most runtimes terminate the process, which is acceptable - // for a corrupted-assignment invariant violation. Inherited from - // Clustering.cpp; not worth refactoring without a broader FAISS effort. - // - // Thread safety of hassign[ci]: each thread owns the exclusive slice - // [c0, c1) and only writes hassign[ci] when ci falls in that slice - // (same guard as centroids[ci*d]). No two threads share an index. #pragma omp parallel { int nt = omp_get_num_threads(); diff --git a/thirdparty/faiss/faiss/impl/NSG.cpp b/thirdparty/faiss/faiss/impl/NSG.cpp index 2826863b9..74250e135 100644 --- a/thirdparty/faiss/faiss/impl/NSG.cpp +++ b/thirdparty/faiss/faiss/impl/NSG.cpp @@ -256,6 +256,8 @@ void NSG::search_on_graph( retset.resize(pool_size + 1); std::vector init_ids(pool_size); + vt.reserve(pool_size); + int num_ids = 0; std::vector neighbors(graph.K); size_t nneigh = graph.get_neighbors(ep, neighbors.data()); diff --git a/thirdparty/faiss/faiss/impl/ResultHandler.h b/thirdparty/faiss/faiss/impl/ResultHandler.h index 93e34d5c5..6beb9f470 100644 --- a/thirdparty/faiss/faiss/impl/ResultHandler.h +++ b/thirdparty/faiss/faiss/impl/ResultHandler.h @@ -377,8 +377,9 @@ struct HeapBlockResultHandler : TopkBlockResultHandler { /// series of results for queries i0..i1 is done void end_multiple() final { - // maybe parallel for - for (size_t i = i0; i < i1; i++) { +#pragma omp parallel for schedule(static) if ((i1 - i0) * k >= 1024) + for (int64_t i = static_cast(i0); i < static_cast(i1); + i++) { heap_reorder(k, this->dis_tab + i * k, this->ids_tab + i * k); } } @@ -568,9 +569,10 @@ struct ReservoirBlockResultHandler : TopkBlockResultHandler { /// series of results for queries i0..i1 is done void end_multiple() final { - // maybe parallel for - for (size_t i = i0; i < i1; i++) { - reservoirs[i - i0].to_result( +#pragma omp parallel for schedule(static) if ((i1 - i0) * this->k >= 1024) + for (int64_t i = static_cast(i0); i < static_cast(i1); + i++) { + reservoirs[i - static_cast(i0)].to_result( this->dis_tab + i * this->k, this->ids_tab + i * this->k); } } diff --git a/thirdparty/faiss/faiss/impl/VisitedTable.h b/thirdparty/faiss/faiss/impl/VisitedTable.h index f71aa7766..eb6f49b0b 100644 --- a/thirdparty/faiss/faiss/impl/VisitedTable.h +++ b/thirdparty/faiss/faiss/impl/VisitedTable.h @@ -45,6 +45,13 @@ struct VisitedTable { } } + /// pre-allocate bucket space to avoid rehashing during repeated set() calls + void reserve(size_t n) { + if (visno == 0) { + visited_set.reserve(n); + } + } + /// get flag #no bool get(size_t no) const { if (visno == 0) { diff --git a/thirdparty/faiss/faiss/impl/fast_scan/accumulate_loops.h b/thirdparty/faiss/faiss/impl/fast_scan/accumulate_loops.h index 692556dd5..34d5e174a 100644 --- a/thirdparty/faiss/faiss/impl/fast_scan/accumulate_loops.h +++ b/thirdparty/faiss/faiss/impl/fast_scan/accumulate_loops.h @@ -47,7 +47,12 @@ using namespace simd_result_handlers; * Search_1 path helpers (multi-BB kernel, bbs > 32) ***************************************************************/ -template +template < + int NQ, + int BB, + SIMDLevel KernelSL = SINGLE_SIMD_LEVEL, + class ResultHandler, + class Scaler> void accumulate_fixed_blocks( size_t nb, int nsq, @@ -58,13 +63,17 @@ void accumulate_fixed_blocks( size_t block_stride) { constexpr int bbs = 32 * BB; for_each_block(nb, codes, block_stride, res, [&](size_t) { - FixedStorageHandler res2; - kernel_accumulate_block(nsq, codes, LUT, res2, scaler); + FixedStorageHandler res2; + kernel_accumulate_block( + nsq, codes, LUT, res2, scaler); res2.to_other_handler(res); }); } -template +template < + SIMDLevel KernelSL = SINGLE_SIMD_LEVEL, + class ResultHandler, + class Scaler> void pq4_accumulate_loop_fixed_scaler( int nq, size_t nb, @@ -82,7 +91,7 @@ void pq4_accumulate_loop_fixed_scaler( #define FAISS_ACCLOOP_DISPATCH(NQ, BB) \ case NQ * 1000 + BB: \ - accumulate_fixed_blocks( \ + accumulate_fixed_blocks( \ nb, nsq, codes, LUT, res, scaler, block_stride); \ break @@ -106,7 +115,11 @@ void pq4_accumulate_loop_fixed_scaler( * QBS path helpers (bbs == 32, 256-bit kernel only) ***************************************************************/ -template +template < + int QBS, + SIMDLevel KernelSL = SINGLE_SIMD_LEVEL, + class ResultHandler, + class Scaler> void accumulate_q_4step_256( size_t ntotal2, int nsq, @@ -122,29 +135,32 @@ void accumulate_q_4step_256( constexpr int SQ = Q1 + Q2 + Q3 + Q4; for_each_block<32>(ntotal2, codes, block_stride, res, [&](size_t) { - FixedStorageHandler res2; + FixedStorageHandler res2; const uint8_t* LUT = LUT0; - pq4_kernel_qbs_256(nsq, codes, LUT, res2, scaler); + pq4_kernel_qbs_256(nsq, codes, LUT, res2, scaler); LUT += Q1 * nsq * 16; if (Q2 > 0) { res2.set_block_origin(Q1, 0); - pq4_kernel_qbs_256(nsq, codes, LUT, res2, scaler); + pq4_kernel_qbs_256(nsq, codes, LUT, res2, scaler); LUT += Q2 * nsq * 16; } if (Q3 > 0) { res2.set_block_origin(Q1 + Q2, 0); - pq4_kernel_qbs_256(nsq, codes, LUT, res2, scaler); + pq4_kernel_qbs_256(nsq, codes, LUT, res2, scaler); LUT += Q3 * nsq * 16; } if (Q4 > 0) { res2.set_block_origin(Q1 + Q2 + Q3, 0); - pq4_kernel_qbs_256(nsq, codes, LUT, res2, scaler); + pq4_kernel_qbs_256(nsq, codes, LUT, res2, scaler); } res2.to_other_handler(res); }); } -template +template < + SIMDLevel KernelSL = SINGLE_SIMD_LEVEL, + class ResultHandler, + class Scaler> void pq4_accumulate_loop_qbs_fixed_scaler_256( int qbs, size_t ntotal2, @@ -161,7 +177,7 @@ void pq4_accumulate_loop_qbs_fixed_scaler_256( switch (qbs) { #define FAISS_QBS256_DISPATCH(QBS) \ case QBS: \ - accumulate_q_4step_256( \ + accumulate_q_4step_256( \ ntotal2, nsq, codes, LUT0, res, scaler, block_stride); \ return; FAISS_QBS256_DISPATCH(0x3333); // 12 @@ -199,9 +215,9 @@ void pq4_accumulate_loop_qbs_fixed_scaler_256( int nq = qi & 15; qi >>= 4; res.set_block_origin(i0, j0); -#define FAISS_NQ256_DISPATCH(NQ) \ - case NQ: \ - pq4_kernel_qbs_256(nsq, codes, LUT, res, scaler); \ +#define FAISS_NQ256_DISPATCH(NQ) \ + case NQ: \ + pq4_kernel_qbs_256(nsq, codes, LUT, res, scaler); \ break switch (nq) { FAISS_NQ256_DISPATCH(1); diff --git a/thirdparty/faiss/faiss/impl/fast_scan/accumulate_loops_512.h b/thirdparty/faiss/faiss/impl/fast_scan/accumulate_loops_512.h index 15a4286cc..1b9d73b60 100644 --- a/thirdparty/faiss/faiss/impl/fast_scan/accumulate_loops_512.h +++ b/thirdparty/faiss/faiss/impl/fast_scan/accumulate_loops_512.h @@ -15,10 +15,9 @@ * (from kernels_simd512.h) instead of pq4_kernel_qbs_256. * * The 512-bit kernels produce simd16uint16_tpl results (via - * combine4x2). The virtual SIMDResultHandler::handle() expects - * simd16uint16_tpl in DD mode. FixedStorage512 bridges this gap: - * it stores AVX2-level results internally, then converts to the handler's - * level via storeu/load in to_other_handler(). + * combine4x2). FixedStorage512 stores these results without virtual + * dispatch, then forwards them to the outer handler via + * to_other_handler(). * * Only included from the AVX512 per-ISA TU (impl-avx512.cpp) via * dispatching.h's conditional include. @@ -46,8 +45,8 @@ using namespace simd_result_handlers; * 512-bit kernels produce simd16uint16_tpl. By avoiding * inheritance, handle() can accept AVX2-level types directly. * - * The conversion to the outer handler's type happens in - * to_other_handler() via a store-to-memory roundtrip. + * to_other_handler() passes AVX2-level values directly to the + * outer handler. ***************************************************************/ template @@ -72,17 +71,9 @@ struct FixedStorage512 { template void to_other_handler(OtherResultHandler& other) const { - using handler_simd16 = simd16uint16_tpl; for (int q = 0; q < NQ; q++) { for (int b = 0; b < BB; b += 2) { - // Convert AVX2 → handler level (NONE in DD mode) - ALIGNED(32) uint16_t buf0[16], buf1[16]; - dis[q][b].storeu(buf0); - dis[q][b + 1].storeu(buf1); - handler_simd16 h0, h1; - h0.loadu(buf0); - h1.loadu(buf1); - other.handle(q, b / 2, h0, h1); + other.handle(q, b / 2, dis[q][b], dis[q][b + 1]); } } } @@ -176,17 +167,16 @@ void pq4_accumulate_loop_qbs_fixed_scaler_512( #undef FAISS_QBS512_DISPATCH } - // Fallback for unknown QBS values: use 256-bit path with NONE-level - // scalers for type compatibility. This is rare — pq4_preferred_qbs() - // covers all values above. + // Fallback for unknown QBS values: use 256-bit path. + // This is rare — pq4_preferred_qbs() covers all values above. if constexpr (Scaler::nscale == 0) { - DummyScaler<> scaler_none; - pq4_accumulate_loop_qbs_fixed_scaler_256( - qbs, ntotal2, nsq, codes, LUT0, res, scaler_none, block_stride); + DummyScaler scaler_avx2; + pq4_accumulate_loop_qbs_fixed_scaler_256( + qbs, ntotal2, nsq, codes, LUT0, res, scaler_avx2, block_stride); } else { - NormTableScaler<> scaler_none(scaler.scale_int); - pq4_accumulate_loop_qbs_fixed_scaler_256( - qbs, ntotal2, nsq, codes, LUT0, res, scaler_none, block_stride); + NormTableScaler scaler_avx2(scaler.scale_int); + pq4_accumulate_loop_qbs_fixed_scaler_256( + qbs, ntotal2, nsq, codes, LUT0, res, scaler_avx2, block_stride); } } diff --git a/thirdparty/faiss/faiss/impl/fast_scan/decompose_qbs.h b/thirdparty/faiss/faiss/impl/fast_scan/decompose_qbs.h index fd3f23241..48ca40210 100644 --- a/thirdparty/faiss/faiss/impl/fast_scan/decompose_qbs.h +++ b/thirdparty/faiss/faiss/impl/fast_scan/decompose_qbs.h @@ -21,8 +21,16 @@ using namespace simd_result_handlers; /* * Unified kernel: selects 256-bit vs 512-bit path based on * compile-time __AVX512F__ guard. + * + * KernelSL selects the SIMD type width used for the inner accumulation + * loop. In DD mode the caller passes the dispatch level so the kernel + * uses real AVX2 types rather than the emulated-scalar fallback. */ -template +template < + int NQ, + SIMDLevel KernelSL = SINGLE_SIMD_LEVEL, + class ResultHandler, + class Scaler> void kernel_accumulate_block( int nsq, const uint8_t* codes, @@ -30,14 +38,24 @@ void kernel_accumulate_block( ResultHandler& res, const Scaler& scaler) { #ifdef __AVX512F__ - pq4_kernel_qbs_512(nsq, codes, LUT, res, scaler); + if constexpr ( + KernelSL == SIMDLevel::AVX512 || + KernelSL == SIMDLevel::AVX512_SPR) { + pq4_kernel_qbs_512(nsq, codes, LUT, res, scaler); + } else { + pq4_kernel_qbs_256(nsq, codes, LUT, res, scaler); + } #else - pq4_kernel_qbs_256(nsq, codes, LUT, res, scaler); + pq4_kernel_qbs_256(nsq, codes, LUT, res, scaler); #endif } // handle at most 4 blocks of queries -template +template < + int QBS, + SIMDLevel KernelSL = SINGLE_SIMD_LEVEL, + class ResultHandler, + class Scaler> void accumulate_q_4step( size_t ntotal2, int nsq, @@ -53,29 +71,36 @@ void accumulate_q_4step( constexpr int SQ = Q1 + Q2 + Q3 + Q4; for_each_block<32>(ntotal2, codes, block_stride, res, [&](size_t) { - FixedStorageHandler res2; + FixedStorageHandler res2; const uint8_t* LUT = LUT0; - kernel_accumulate_block(nsq, codes, LUT, res2, scaler); + kernel_accumulate_block(nsq, codes, LUT, res2, scaler); LUT += Q1 * nsq * 16; if (Q2 > 0) { res2.set_block_origin(Q1, 0); - kernel_accumulate_block(nsq, codes, LUT, res2, scaler); + kernel_accumulate_block( + nsq, codes, LUT, res2, scaler); LUT += Q2 * nsq * 16; } if (Q3 > 0) { res2.set_block_origin(Q1 + Q2, 0); - kernel_accumulate_block(nsq, codes, LUT, res2, scaler); + kernel_accumulate_block( + nsq, codes, LUT, res2, scaler); LUT += Q3 * nsq * 16; } if (Q4 > 0) { res2.set_block_origin(Q1 + Q2 + Q3, 0); - kernel_accumulate_block(nsq, codes, LUT, res2, scaler); + kernel_accumulate_block( + nsq, codes, LUT, res2, scaler); } res2.to_other_handler(res); }); } -template +template < + int NQ, + SIMDLevel KernelSL = SINGLE_SIMD_LEVEL, + class ResultHandler, + class Scaler> void kernel_accumulate_block_loop( size_t ntotal2, int nsq, @@ -85,13 +110,15 @@ void kernel_accumulate_block_loop( const Scaler& scaler, size_t block_stride) { for_each_block<32>(ntotal2, codes, block_stride, res, [&](size_t) { - kernel_accumulate_block( - nsq, codes, LUT, res, scaler); + kernel_accumulate_block(nsq, codes, LUT, res, scaler); }); } // non-template version of accumulate kernel -- dispatches dynamically -template +template < + SIMDLevel KernelSL = SINGLE_SIMD_LEVEL, + class ResultHandler, + class Scaler> void accumulate( int nq, size_t ntotal2, @@ -106,7 +133,7 @@ void accumulate( #define DISPATCH(NQ) \ case NQ: \ - kernel_accumulate_block_loop( \ + kernel_accumulate_block_loop( \ ntotal2, nsq, codes, LUT, res, scaler, block_stride); \ return @@ -121,7 +148,10 @@ void accumulate( #undef DISPATCH } -template +template < + SIMDLevel KernelSL = SINGLE_SIMD_LEVEL, + class ResultHandler, + class Scaler> void pq4_accumulate_loop_qbs_fixed_scaler( int qbs, size_t ntotal2, @@ -139,7 +169,7 @@ void pq4_accumulate_loop_qbs_fixed_scaler( switch (qbs) { #define DISPATCH(QBS) \ case QBS: \ - accumulate_q_4step( \ + accumulate_q_4step( \ ntotal2, nsq, codes, LUT0, res, scaler, block_stride); \ return; DISPATCH(0x3333); // 12 @@ -177,10 +207,9 @@ void pq4_accumulate_loop_qbs_fixed_scaler( int nq = qi & 15; qi >>= 4; res.set_block_origin(i0, j0); -#define DISPATCH(NQ) \ - case NQ: \ - kernel_accumulate_block( \ - nsq, codes, LUT, res, scaler); \ +#define DISPATCH(NQ) \ + case NQ: \ + kernel_accumulate_block(nsq, codes, LUT, res, scaler); \ break switch (nq) { DISPATCH(1); diff --git a/thirdparty/faiss/faiss/impl/fast_scan/dispatching.h b/thirdparty/faiss/faiss/impl/fast_scan/dispatching.h index abc1f7757..dde2698fd 100644 --- a/thirdparty/faiss/faiss/impl/fast_scan/dispatching.h +++ b/thirdparty/faiss/faiss/impl/fast_scan/dispatching.h @@ -70,8 +70,8 @@ struct ScannerMixIn : FastScanCodeScanner { int pq2x4_scale, size_t block_stride) override { if (pq2x4_scale) { - NormTableScaler<> scaler(pq2x4_scale); - pq4_accumulate_loop_fixed_scaler( + NormTableScaler scaler(pq2x4_scale); + pq4_accumulate_loop_fixed_scaler( nq, nb, bbs, @@ -82,8 +82,8 @@ struct ScannerMixIn : FastScanCodeScanner { scaler, block_stride); } else { - DummyScaler<> dummy; - pq4_accumulate_loop_fixed_scaler( + DummyScaler dummy; + pq4_accumulate_loop_fixed_scaler( nq, nb, bbs, @@ -138,8 +138,8 @@ struct ScannerMixIn : FastScanCodeScanner { } } else { if (pq2x4_scale) { - NormTableScaler<> scaler(pq2x4_scale); - pq4_accumulate_loop_qbs_fixed_scaler_256( + NormTableScaler scaler(pq2x4_scale); + pq4_accumulate_loop_qbs_fixed_scaler_256( qbs, nb, nsq, @@ -149,8 +149,8 @@ struct ScannerMixIn : FastScanCodeScanner { scaler, block_stride); } else { - DummyScaler<> dummy; - pq4_accumulate_loop_qbs_fixed_scaler_256( + DummyScaler dummy; + pq4_accumulate_loop_qbs_fixed_scaler_256( qbs, nb, nsq, @@ -188,15 +188,15 @@ std::unique_ptr make_fast_scan_scanner_impl< // Helper lambda: given comparator C and with_id_map W, select handler auto make = [&]() -> std::unique_ptr { if (k == 1) { - using H = SingleResultHandler; + using H = SingleResultHandler; return std::make_unique>( nq, ntotal, distances, ids, sel); } else if (impl % 2 == 0) { - using H = HeapHandler; + using H = HeapHandler; return std::make_unique>( nq, ntotal, k, distances, ids, sel); } else { - using H = ReservoirHandler; + using H = ReservoirHandler; return std::make_unique>( nq, ntotal, size_t(k), size_t(2 * k), distances, ids, sel); } @@ -231,11 +231,13 @@ std::unique_ptr make_range_scanner_impl< const IDSelector* sel) { if (is_max) { using C = CMax; - return std::make_unique>>( + return std::make_unique< + ScannerMixIn>>( rres, radius, ntotal, sel); } else { using C = CMin; - return std::make_unique>>( + return std::make_unique< + ScannerMixIn>>( rres, radius, ntotal, sel); } } @@ -252,11 +254,13 @@ std::unique_ptr make_partial_range_scanner_impl< const IDSelector* sel) { if (is_max) { using C = CMax; - return std::make_unique>>( + return std::make_unique>>( pres, radius, ntotal, q0, q1, sel); } else { using C = CMin; - return std::make_unique>>( + return std::make_unique>>( pres, radius, ntotal, q0, q1, sel); } } diff --git a/thirdparty/faiss/faiss/impl/fast_scan/fast_scan.cpp b/thirdparty/faiss/faiss/impl/fast_scan/fast_scan.cpp index bfee685d0..c469fb2df 100644 --- a/thirdparty/faiss/faiss/impl/fast_scan/fast_scan.cpp +++ b/thirdparty/faiss/faiss/impl/fast_scan/fast_scan.cpp @@ -385,20 +385,6 @@ int pq4_preferred_qbs(int n) { } } -void accumulate_to_mem( - int nq, - size_t ntotal2, - int nsq, - const uint8_t* codes, - const uint8_t* LUT, - uint16_t* accu) { - using namespace simd_result_handlers; - FAISS_THROW_IF_NOT(ntotal2 % 32 == 0); - StoreResultHandler<> handler(accu, ntotal2); - DummyScaler<> scaler; - accumulate(nq, ntotal2, nsq, codes, LUT, handler, scaler, 32 * nsq / 2); -} - } // namespace faiss /*************************************************************** @@ -416,81 +402,42 @@ void accumulate_to_mem( namespace faiss { -#ifdef COMPILE_SIMD_RISCV_RVV -template <> -std::unique_ptr make_fast_scan_scanner_impl< - SIMDLevel::RISCV_RVV>( - bool is_max, - int impl, - size_t nq, - size_t ntotal, - int64_t k, - float* distances, - int64_t* ids, - const IDSelector* sel, - bool with_id_map) { - return make_fast_scan_scanner_impl( - is_max, impl, nq, ntotal, k, distances, ids, sel, with_id_map); -} +using namespace simd_result_handlers; -template <> -std::unique_ptr make_range_scanner_impl< - SIMDLevel::RISCV_RVV>( - bool is_max, - RangeSearchResult& rres, - float radius, - size_t ntotal, - const IDSelector* sel) { - return make_range_scanner_impl( - is_max, rres, radius, ntotal, sel); -} +/*************************************************************** + * accumulate_to_mem: NONE specialization + runtime dispatch. + ***************************************************************/ template <> -std::unique_ptr make_partial_range_scanner_impl< - SIMDLevel::RISCV_RVV>( - bool is_max, - RangeSearchPartialResult& pres, - float radius, - size_t ntotal, - size_t q0, - size_t q1, - const IDSelector* sel) { - return make_partial_range_scanner_impl( - is_max, pres, radius, ntotal, q0, q1, sel); +void accumulate_to_mem_impl( + int nq, + size_t ntotal2, + int nsq, + const uint8_t* codes, + const uint8_t* LUT, + uint16_t* accu) { + StoreResultHandler handler(accu, ntotal2); + DummyScaler scaler; + accumulate( + nq, ntotal2, nsq, codes, LUT, handler, scaler, 32 * nsq / 2); } -template <> -std::unique_ptr rabitq_make_knn_scanner_impl< - SIMDLevel::RISCV_RVV>( - const IndexRaBitQFastScan* index, - bool is_max, - size_t nq, - int64_t k, - float* distances, - int64_t* ids, - const IDSelector* sel, - const FastScanDistancePostProcessing& context, - bool is_multi_bit) { - return rabitq_make_knn_scanner_impl( - index, is_max, nq, k, distances, ids, sel, context, is_multi_bit); +void accumulate_to_mem( + int nq, + size_t ntotal2, + int nsq, + const uint8_t* codes, + const uint8_t* LUT, + uint16_t* accu) { + FAISS_THROW_IF_NOT(ntotal2 % 32 == 0); + with_simd_level([&]() { + accumulate_to_mem_impl(nq, ntotal2, nsq, codes, LUT, accu); + }); } -template <> -std::unique_ptr rabitq_ivf_make_knn_scanner_impl< - SIMDLevel::RISCV_RVV>( - bool is_max, - const IndexIVFRaBitQFastScan* index, - size_t nq, - size_t k, - float* distances, - int64_t* ids, - const IDSelector* sel, - const FastScanDistancePostProcessing* context, - bool multi_bit) { - return rabitq_ivf_make_knn_scanner_impl( - is_max, index, nq, k, distances, ids, sel, context, multi_bit); -} -#endif // COMPILE_SIMD_RISCV_RVV +} // namespace faiss + +namespace faiss { std::unique_ptr make_fast_scan_knn_scanner( bool is_max, diff --git a/thirdparty/faiss/faiss/impl/fast_scan/fast_scan.h b/thirdparty/faiss/faiss/impl/fast_scan/fast_scan.h index f138ce224..0d260295f 100644 --- a/thirdparty/faiss/faiss/impl/fast_scan/fast_scan.h +++ b/thirdparty/faiss/faiss/impl/fast_scan/fast_scan.h @@ -168,6 +168,16 @@ void accumulate_to_mem( const uint8_t* LUT, uint16_t* accu); +/// Per-SIMD specialization of accumulate_to_mem (defined in per-SIMD TUs) +template +void accumulate_to_mem_impl( + int nq, + size_t ntotal2, + int nsq, + const uint8_t* codes, + const uint8_t* LUT, + uint16_t* accu); + /*************************************************************** * FastScanCodeScanner: virtual base that bundles handler + kernel * behind the SIMD dispatch boundary. Per-SIMD TUs instantiate this diff --git a/thirdparty/faiss/faiss/impl/fast_scan/impl-avx2.cpp b/thirdparty/faiss/faiss/impl/fast_scan/impl-avx2.cpp index 3dcac04f4..5dba4f9ad 100644 --- a/thirdparty/faiss/faiss/impl/fast_scan/impl-avx2.cpp +++ b/thirdparty/faiss/faiss/impl/fast_scan/impl-avx2.cpp @@ -11,4 +11,26 @@ #include // IWYU pragma: keep #include // IWYU pragma: keep +#include + +namespace faiss { + +using namespace simd_result_handlers; + +template <> +void accumulate_to_mem_impl( + int nq, + size_t ntotal2, + int nsq, + const uint8_t* codes, + const uint8_t* LUT, + uint16_t* accu) { + StoreResultHandler handler(accu, ntotal2); + DummyScaler scaler; + accumulate( + nq, ntotal2, nsq, codes, LUT, handler, scaler, 32 * nsq / 2); +} + +} // namespace faiss + #endif // COMPILE_SIMD_AVX2 diff --git a/thirdparty/faiss/faiss/impl/fast_scan/impl-avx512.cpp b/thirdparty/faiss/faiss/impl/fast_scan/impl-avx512.cpp index 2fac7920d..9f03766e8 100644 --- a/thirdparty/faiss/faiss/impl/fast_scan/impl-avx512.cpp +++ b/thirdparty/faiss/faiss/impl/fast_scan/impl-avx512.cpp @@ -11,4 +11,30 @@ #include // IWYU pragma: keep #include // IWYU pragma: keep +#include + +namespace faiss { + +using namespace simd_result_handlers; + +template <> +void accumulate_to_mem_impl( + int nq, + size_t ntotal2, + int nsq, + const uint8_t* codes, + const uint8_t* LUT, + uint16_t* accu) { + // Use AVX2-level handler (256-bit StoreResultHandler) since the 512-bit + // kernels reduce to AVX2-level simd16uint16 via FixedStorage512. + StoreResultHandler handler(accu, ntotal2); + DummyScaler scaler; + // kernel_accumulate_block in decompose_qbs.h selects pq4_kernel_qbs_512 + // via #ifdef __AVX512F__ (which is set for this TU). + accumulate( + nq, ntotal2, nsq, codes, LUT, handler, scaler, 32 * nsq / 2); +} + +} // namespace faiss + #endif // COMPILE_SIMD_AVX512 diff --git a/thirdparty/faiss/faiss/impl/fast_scan/impl-neon.cpp b/thirdparty/faiss/faiss/impl/fast_scan/impl-neon.cpp index af051b918..7c40e832d 100644 --- a/thirdparty/faiss/faiss/impl/fast_scan/impl-neon.cpp +++ b/thirdparty/faiss/faiss/impl/fast_scan/impl-neon.cpp @@ -11,6 +11,28 @@ #include // IWYU pragma: keep #include // IWYU pragma: keep +#include + +namespace faiss { + +using namespace simd_result_handlers; + +template <> +void accumulate_to_mem_impl( + int nq, + size_t ntotal2, + int nsq, + const uint8_t* codes, + const uint8_t* LUT, + uint16_t* accu) { + StoreResultHandler handler(accu, ntotal2); + DummyScaler scaler; + accumulate( + nq, ntotal2, nsq, codes, LUT, handler, scaler, 32 * nsq / 2); +} + +} // namespace faiss + // ARM_SVE: forward to ARM_NEON implementation until a dedicated SVE // specialization is written (same pattern as scalar_quantizer/sq-neon.cpp). #ifdef COMPILE_SIMD_ARM_SVE diff --git a/thirdparty/faiss/faiss/impl/fast_scan/impl-riscv.cpp b/thirdparty/faiss/faiss/impl/fast_scan/impl-riscv.cpp new file mode 100644 index 000000000..b765a016f --- /dev/null +++ b/thirdparty/faiss/faiss/impl/fast_scan/impl-riscv.cpp @@ -0,0 +1,104 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ + +// RISC-V RVV: forward all fast_scan specializations to NONE until +// dedicated RVV implementations are written. + +#ifdef COMPILE_SIMD_RISCV_RVV + +#include + +namespace faiss { + +template <> +void accumulate_to_mem_impl( + int nq, + size_t ntotal2, + int nsq, + const uint8_t* codes, + const uint8_t* LUT, + uint16_t* accu) { + accumulate_to_mem_impl(nq, ntotal2, nsq, codes, LUT, accu); +} + +template <> +std::unique_ptr make_fast_scan_scanner_impl< + SIMDLevel::RISCV_RVV>( + bool is_max, + int impl, + size_t nq, + size_t ntotal, + int64_t k, + float* distances, + int64_t* ids, + const IDSelector* sel, + bool with_id_map) { + return make_fast_scan_scanner_impl( + is_max, impl, nq, ntotal, k, distances, ids, sel, with_id_map); +} + +template <> +std::unique_ptr make_range_scanner_impl< + SIMDLevel::RISCV_RVV>( + bool is_max, + RangeSearchResult& rres, + float radius, + size_t ntotal, + const IDSelector* sel) { + return make_range_scanner_impl( + is_max, rres, radius, ntotal, sel); +} + +template <> +std::unique_ptr make_partial_range_scanner_impl< + SIMDLevel::RISCV_RVV>( + bool is_max, + RangeSearchPartialResult& pres, + float radius, + size_t ntotal, + size_t q0, + size_t q1, + const IDSelector* sel) { + return make_partial_range_scanner_impl( + is_max, pres, radius, ntotal, q0, q1, sel); +} + +template <> +std::unique_ptr rabitq_make_knn_scanner_impl< + SIMDLevel::RISCV_RVV>( + const IndexRaBitQFastScan* index, + bool is_max, + size_t nq, + int64_t k, + float* distances, + int64_t* ids, + const IDSelector* sel, + const FastScanDistancePostProcessing& context, + bool is_multi_bit) { + return rabitq_make_knn_scanner_impl( + index, is_max, nq, k, distances, ids, sel, context, is_multi_bit); +} + +template <> +std::unique_ptr rabitq_ivf_make_knn_scanner_impl< + SIMDLevel::RISCV_RVV>( + bool is_max, + const IndexIVFRaBitQFastScan* index, + size_t nq, + size_t k, + float* distances, + int64_t* ids, + const IDSelector* sel, + const FastScanDistancePostProcessing* context, + bool multi_bit) { + return rabitq_ivf_make_knn_scanner_impl( + is_max, index, nq, k, distances, ids, sel, context, multi_bit); +} + +} // namespace faiss + +#endif // COMPILE_SIMD_RISCV_RVV diff --git a/thirdparty/faiss/faiss/impl/fast_scan/kernels_simd256.h b/thirdparty/faiss/faiss/impl/fast_scan/kernels_simd256.h index cd889e5fd..71142832e 100644 --- a/thirdparty/faiss/faiss/impl/fast_scan/kernels_simd256.h +++ b/thirdparty/faiss/faiss/impl/fast_scan/kernels_simd256.h @@ -11,21 +11,31 @@ namespace faiss { -// Explicit SIMD-level aliases for this file (no global bare aliases). -using simd16uint16 = simd16uint16_tpl; -using simd32uint8 = simd32uint8_tpl; - /* * Multi-BB variant: accumulates NQ queries x BB*32 database elements. * Used by the search_1 path (bbs > 32). + * + * KernelSL selects the SIMD type width used for the inner accumulation + * loop. In DD mode the caller passes THE_LEVEL_TO_DISPATCH so the + * kernel uses real AVX2/AVX512 types rather than the emulated-scalar + * fallback that SINGLE_SIMD_LEVEL_256 would give. */ -template +template < + int NQ, + int BB, + SIMDLevel KernelSL = SINGLE_SIMD_LEVEL, + class ResultHandler, + class Scaler> void kernel_accumulate_block( int nsq, const uint8_t* codes, const uint8_t* LUT, ResultHandler& res, const Scaler& scaler) { + constexpr SIMDLevel SL256 = simd256_level_selector::value; + using simd16uint16 = simd16uint16_tpl; + using simd32uint8 = simd32uint8_tpl; + // distance accumulators simd16uint16 accu[NQ][BB][4]; @@ -112,13 +122,21 @@ void kernel_accumulate_block( * Single-BB QBS variant: accumulates NQ queries x 32 db elements (BB=1). * Used by the decompose_qbs layer for non-AVX512 paths. */ -template +template < + int NQ, + SIMDLevel KernelSL = SINGLE_SIMD_LEVEL, + class ResultHandler, + class Scaler> void pq4_kernel_qbs_256( int nsq, const uint8_t* codes, const uint8_t* LUT, ResultHandler& res, const Scaler& scaler) { + constexpr SIMDLevel SL256 = simd256_level_selector::value; + using simd16uint16 = simd16uint16_tpl; + using simd32uint8 = simd32uint8_tpl; + // dummy alloc to keep the windows compiler happy constexpr int NQA = NQ > 0 ? NQ : 1; // distance accumulators diff --git a/thirdparty/faiss/faiss/impl/fast_scan/rabitq_dispatching.h b/thirdparty/faiss/faiss/impl/fast_scan/rabitq_dispatching.h index 28d2a9cd0..5c0f83097 100644 --- a/thirdparty/faiss/faiss/impl/fast_scan/rabitq_dispatching.h +++ b/thirdparty/faiss/faiss/impl/fast_scan/rabitq_dispatching.h @@ -43,11 +43,17 @@ std::unique_ptr rabitq_make_knn_scanner_impl< const FastScanDistancePostProcessing& context, bool is_multi_bit) { if (is_max) { - using H = RaBitQHeapHandler, false>; + using H = RaBitQHeapHandler< + CMax, + false, + THE_LEVEL_TO_DISPATCH>; return std::make_unique>( index, nq, k, distances, ids, sel, &context, is_multi_bit); } else { - using H = RaBitQHeapHandler, false>; + using H = RaBitQHeapHandler< + CMin, + false, + THE_LEVEL_TO_DISPATCH>; return std::make_unique>( index, nq, k, distances, ids, sel, &context, is_multi_bit); } @@ -68,12 +74,14 @@ std::unique_ptr rabitq_ivf_make_knn_scanner_impl< bool multi_bit) { if (is_max) { using C = CMax; - using H = simd_result_handlers::IVFRaBitQHeapHandler; + using H = simd_result_handlers:: + IVFRaBitQHeapHandler; return std::make_unique>( index, nq, k, distances, ids, sel, context, multi_bit); } else { using C = CMin; - using H = simd_result_handlers::IVFRaBitQHeapHandler; + using H = simd_result_handlers:: + IVFRaBitQHeapHandler; return std::make_unique>( index, nq, k, distances, ids, sel, context, multi_bit); } diff --git a/thirdparty/faiss/faiss/impl/fast_scan/rabitq_result_handler.h b/thirdparty/faiss/faiss/impl/fast_scan/rabitq_result_handler.h index 59950e8ae..eee359b2a 100644 --- a/thirdparty/faiss/faiss/impl/fast_scan/rabitq_result_handler.h +++ b/thirdparty/faiss/faiss/impl/fast_scan/rabitq_result_handler.h @@ -38,6 +38,7 @@ using rabitq_utils::SignBitFactorsWithError; template struct IVFRaBitQHeapHandler : ResultHandlerCompare { using RHC = ResultHandlerCompare; + using RHC::handle; using typename RHC::simd16uint16; const IndexIVFRaBitQFastScan* index; @@ -82,7 +83,7 @@ struct IVFRaBitQHeapHandler : ResultHandlerCompare { const FastScanDistancePostProcessing* ctx = nullptr, bool multibit = false); - void handle(size_t q, size_t b, simd16uint16 d0, simd16uint16 d1) override; + void handle(size_t q, size_t b, simd16uint16 d0, simd16uint16 d1); void set_list_context(size_t list_no, const std::vector& probe_map) override; diff --git a/thirdparty/faiss/faiss/impl/fast_scan/simd_result_handlers.h b/thirdparty/faiss/faiss/impl/fast_scan/simd_result_handlers.h index e7b810350..32ce8c449 100644 --- a/thirdparty/faiss/faiss/impl/fast_scan/simd_result_handlers.h +++ b/thirdparty/faiss/faiss/impl/fast_scan/simd_result_handlers.h @@ -83,7 +83,8 @@ inline void for_each_block( } // namespace -// Explicit SIMD-level alias for the virtual interface below. +// SIMD-level alias for the virtual interface fallback (always NONE so the +// base-class signature is stable across translation units in DD mode). using simd16uint16 = simd16uint16_tpl; struct SIMDResultHandler { @@ -100,12 +101,19 @@ struct SIMDResultHandler { /** called when 32 distances are computed and provided in two * simd16uint16. (q, b) indicate which entry it is - * in the block. */ + * in the block. + * + * Non-pure: concrete handlers define their own handle() with the + * SIMD types matching their SL template parameter. When SL != NONE + * the handler's handle() hides this base version (different param + * types); when SL == NONE it overrides it. The kernel code always + * calls handle() through a concrete-type reference, so virtual + * dispatch is not on the hot path. */ virtual void handle( - size_t q, - size_t b, - simd16uint16 d0, - simd16uint16 d1) = 0; + size_t /* q */, + size_t /* b */, + simd16uint16 /* d0 */, + simd16uint16 /* d1 */) {} /// set the sub-matrix that is being computed virtual void set_block_origin(size_t i0, size_t j0) = 0; @@ -179,18 +187,17 @@ namespace simd_result_handlers { * (to avoid the computation to be optimized away) */ template struct DummyResultHandler : SIMDResultHandler { + using SIMDResultHandler::handle; static constexpr SIMDLevel SL256 = simd256_level_selector::value; using simd16uint16 = simd16uint16_tpl; size_t cs = 0; - void handle(size_t q, size_t b, simd16uint16 d0, simd16uint16 d1) final { + void handle(size_t q, size_t b, simd16uint16 d0, simd16uint16 d1) { cs += q * 123 + b * 789 + d0.get_scalar_0() + d1.get_scalar_0(); } void set_block_origin(size_t, size_t) final {} - - ~DummyResultHandler() {} }; /** memorize results in a nq-by-nb matrix. @@ -199,6 +206,7 @@ struct DummyResultHandler : SIMDResultHandler { */ template struct StoreResultHandler : SIMDResultHandler { + using SIMDResultHandler::handle; static constexpr SIMDLevel SL256 = simd256_level_selector::value; using simd16uint16 = simd16uint16_tpl; @@ -209,7 +217,7 @@ struct StoreResultHandler : SIMDResultHandler { StoreResultHandler(uint16_t* data, size_t ld) : data(data), ld(ld) {} - void handle(size_t q, size_t b, simd16uint16 d0, simd16uint16 d1) final { + void handle(size_t q, size_t b, simd16uint16 d0, simd16uint16 d1) { size_t ofs = (q + i0) * ld + j0 + b * 32; d0.storeu(data + ofs); d1.storeu(data + ofs + 16); @@ -224,13 +232,14 @@ struct StoreResultHandler : SIMDResultHandler { /** stores results in fixed-size matrix. */ template struct FixedStorageHandler : SIMDResultHandler { + using SIMDResultHandler::handle; static constexpr SIMDLevel SL256 = simd256_level_selector::value; using simd16uint16 = simd16uint16_tpl; simd16uint16 dis[NQ][BB]; int i0 = 0; - void handle(size_t q, size_t b, simd16uint16 d0, simd16uint16 d1) final { + void handle(size_t q, size_t b, simd16uint16 d0, simd16uint16 d1) { dis[q + i0][2 * b] = d0; dis[q + i0][2 * b + 1] = d1; } @@ -248,13 +257,12 @@ struct FixedStorageHandler : SIMDResultHandler { } } } - - virtual ~FixedStorageHandler() {} }; /** Result handler that compares distances to check if they need to be kept */ template struct ResultHandlerCompare : SIMDResultHandlerToFloat { + using SIMDResultHandler::handle; using TI = typename C::TI; static constexpr SIMDLevel SL256 = simd256_level_selector::value; using simd16uint16 = simd16uint16_tpl; @@ -332,8 +340,6 @@ struct ResultHandlerCompare : SIMDResultHandlerToFloat { } return lt_mask; } - - virtual ~ResultHandlerCompare() {} }; /** Special version for k=1 */ @@ -345,6 +351,7 @@ struct SingleResultHandler : ResultHandlerCompare { using T = typename C::T; using TI = typename C::TI; using RHC = ResultHandlerCompare; + using RHC::handle; using RHC::normalizers; using typename RHC::simd16uint16; @@ -365,7 +372,7 @@ struct SingleResultHandler : ResultHandlerCompare { } } - void handle(size_t q, size_t b, simd16uint16 d0, simd16uint16 d1) final { + void handle(size_t q, size_t b, simd16uint16 d0, simd16uint16 d1) { if (this->disable) { return; } @@ -438,6 +445,7 @@ struct HeapHandler : ResultHandlerCompare { using T = typename C::T; using TI = typename C::TI; using RHC = ResultHandlerCompare; + using RHC::handle; using RHC::normalizers; using typename RHC::simd16uint16; @@ -480,7 +488,7 @@ struct HeapHandler : ResultHandlerCompare { return C::neutral(); } - void handle(size_t q, size_t b, simd16uint16 d0, simd16uint16 d1) final { + void handle(size_t q, size_t b, simd16uint16 d0, simd16uint16 d1) { if (this->disable) { return; } @@ -580,6 +588,7 @@ struct ReservoirHandler : ResultHandlerCompare { using T = typename C::T; using TI = typename C::TI; using RHC = ResultHandlerCompare; + using RHC::handle; using RHC::normalizers; using typename RHC::simd16uint16; @@ -617,7 +626,7 @@ struct ReservoirHandler : ResultHandlerCompare { } } - void handle(size_t q, size_t b, simd16uint16 d0, simd16uint16 d1) final { + void handle(size_t q, size_t b, simd16uint16 d0, simd16uint16 d1) { if (this->disable) { return; } @@ -714,6 +723,7 @@ struct RangeHandler : ResultHandlerCompare { using T = typename C::T; using TI = typename C::TI; using RHC = ResultHandlerCompare; + using RHC::handle; using RHC::normalizers; using RHC::nq; using typename RHC::simd16uint16; @@ -751,7 +761,7 @@ struct RangeHandler : ResultHandlerCompare { } } - void handle(size_t q, size_t b, simd16uint16 d0, simd16uint16 d1) final { + void handle(size_t q, size_t b, simd16uint16 d0, simd16uint16 d1) { if (this->disable) { return; } @@ -902,6 +912,7 @@ struct SingleQueryResultCollectHandler using T = typename C::T; using TI = typename C::TI; using RHC = ResultHandlerCompare; + using RHC::handle; using RHC::normalizers; using simd16uint16 = typename RHC::simd16uint16; @@ -921,7 +932,7 @@ struct SingleQueryResultCollectHandler normalizers = norms; } - void handle(size_t q, size_t b, simd16uint16 d0, simd16uint16 d1) final { + void handle(size_t q, size_t b, simd16uint16 d0, simd16uint16 d1) { if (this->disable) { return; } diff --git a/thirdparty/faiss/faiss/impl/index_read.cpp b/thirdparty/faiss/faiss/impl/index_read.cpp index 1ee914b9e..574849114 100644 --- a/thirdparty/faiss/faiss/impl/index_read.cpp +++ b/thirdparty/faiss/faiss/impl/index_read.cpp @@ -86,6 +86,24 @@ namespace faiss { namespace { size_t deserialization_loop_limit_ = 0; size_t deserialization_vector_byte_limit_ = uint64_t{1} << 40; // 1 TB + +#ifdef FAISS_ENABLE_SVS +// Read and validate an SVSStorageKind from the stream. Centralizes the +// [0, SVS_count) range check so every SVS read site rejects out-of-range +// values uniformly at the deserialization boundary, instead of letting +// to_svs_storage_kind() surface the failure later from inside an SVS +// runtime load. +SVSStorageKind read_svs_storage_kind(IOReader* f) { + int sk; + READ1(sk); + FAISS_THROW_IF_NOT_FMT( + sk >= 0 && sk < static_cast(SVS_count), + "invalid SVS storage_kind=%d (must be in [0, %d))", + sk, + static_cast(SVS_count)); + return static_cast(sk); +} +#endif // FAISS_ENABLE_SVS } // namespace size_t get_deserialization_loop_limit() { @@ -712,6 +730,13 @@ void read_ProductQuantizer(ProductQuantizer* pq, IOReader* f) { FAISS_THROW_IF_NOT_MSG( n < get_deserialization_vector_byte_limit() / sizeof(float), "PQ centroids allocation would exceed deserialization byte limit"); + // Per-subquantizer tables (e.g. IVFPQ residual norms, search-time + // distance tables) are sized M * ksub. + size_t m_ksub = mul_no_overflow(pq->M, ksub, "PQ M*ksub"); + FAISS_THROW_IF_NOT_MSG( + m_ksub < + get_deserialization_vector_byte_limit() / sizeof(float), + "PQ M*ksub allocation would exceed deserialization byte limit"); } pq->set_derived_values(); READVECTOR(pq->centroids); @@ -2120,6 +2145,12 @@ std::unique_ptr read_index_up(IOReader* f, int io_flags) { ") does not match ntotal (%" PRId64 ")", int64_t(idxmap->id_map.size()), idxmap->ntotal); + FAISS_THROW_IF_NOT_FMT( + idxmap->index->ntotal == idxmap->ntotal, + "IndexIDMap inner index ntotal (%" PRId64 + ") does not match IndexIDMap ntotal (%" PRId64 ")", + idxmap->index->ntotal, + idxmap->ntotal); if (is_map2) { static_cast(idxmap.get())->construct_rev_map(); } @@ -2522,13 +2553,14 @@ std::unique_ptr read_index_up(IOReader* f, int io_flags) { } #ifdef FAISS_ENABLE_SVS else if ( - h == fourcc("ILVQ") || h == fourcc("ISVL") || h == fourcc("ISVD")) { + h == fourcc("ILVQ") || h == fourcc("ISVL") || h == fourcc("ISVD") || + h == fourcc("ISV2")) { std::unique_ptr svs; if (h == fourcc("ILVQ")) { svs = std::make_unique(); } else if (h == fourcc("ISVL")) { svs = std::make_unique(); - } else if (h == fourcc("ISVD")) { + } else if (h == fourcc("ISVD") || h == fourcc("ISV2")) { svs = std::make_unique(); } @@ -2542,14 +2574,7 @@ std::unique_ptr read_index_up(IOReader* f, int io_flags) { READ1(svs->prune_to); READ1(svs->use_full_search_history); - int sk; - READ1(sk); - FAISS_THROW_IF_NOT_FMT( - sk >= 0 && sk < static_cast(SVS_count), - "invalid SVS storage_kind=%d (must be in [0, %d))", - sk, - static_cast(SVS_count)); - svs->storage_kind = static_cast(sk); + svs->storage_kind = read_svs_storage_kind(f); if (h == fourcc("ISVL")) { auto* leanvec = dynamic_cast(svs.get()); @@ -2580,6 +2605,11 @@ std::unique_ptr read_index_up(IOReader* f, int io_flags) { leanvec->deserialize_training_data(is); } } + if (h == fourcc("ISV2")) { + READVECTOR(svs->stored_vectors); + } else { + svs->stored_vectors_valid = false; + } idx = std::move(svs); } else if (h == fourcc("ISVF")) { auto svs = std::make_unique(); @@ -2617,7 +2647,7 @@ std::unique_ptr read_index_up(IOReader* f, int io_flags) { READ1(svs_ivf->k_reorder); READ1(svs_ivf->num_threads); READ1(svs_ivf->intra_query_threads); - READ1(svs_ivf->storage_kind); + svs_ivf->storage_kind = read_svs_storage_kind(f); READ1(svs_ivf->is_static); if (h == fourcc("ISIL")) { auto* leanvec = dynamic_cast(svs_ivf.get()); @@ -3009,7 +3039,18 @@ std::unique_ptr read_index_binary_up(IOReader* f, int io_flags) { idxmap->index = read_index_binary(f, io_flags); idxmap->own_fields = true; READVECTOR(idxmap->id_map); - FAISS_THROW_IF_NOT(idxmap->id_map.size() == idxmap->ntotal); + FAISS_THROW_IF_NOT_FMT( + idxmap->id_map.size() == idxmap->ntotal, + "IndexBinaryIDMap id_map size (%" PRId64 + ") does not match ntotal (%" PRId64 ")", + int64_t(idxmap->id_map.size()), + idxmap->ntotal); + FAISS_THROW_IF_NOT_FMT( + idxmap->index->ntotal == idxmap->ntotal, + "IndexBinaryIDMap inner index ntotal (%" PRId64 + ") does not match IndexBinaryIDMap ntotal (%" PRId64 ")", + idxmap->index->ntotal, + idxmap->ntotal); if (is_map2) { static_cast(idxmap.get())->construct_rev_map(); } diff --git a/thirdparty/faiss/faiss/impl/index_write.cpp b/thirdparty/faiss/faiss/impl/index_write.cpp index c67a74c41..89142b899 100644 --- a/thirdparty/faiss/faiss/impl/index_write.cpp +++ b/thirdparty/faiss/faiss/impl/index_write.cpp @@ -1027,6 +1027,8 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) { h = fourcc("ILVQ"); // LVQ } else if (lean != nullptr) { h = fourcc("ISVL"); // LeanVec + } else if (svs->stored_vectors_valid && !svs->stored_vectors.empty()) { + h = fourcc("ISV2"); // uncompressed + stored_vectors } else { h = fourcc("ISVD"); // uncompressed } @@ -1069,6 +1071,10 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) { os.flush(); } } + + if (h == fourcc("ISV2")) { + WRITEVECTOR(svs->stored_vectors); + } } else if ( const IndexSVSFlat* svs = dynamic_cast(idx)) { uint32_t h = fourcc("ISVF"); diff --git a/thirdparty/faiss/faiss/index_factory.cpp b/thirdparty/faiss/faiss/index_factory.cpp index 877f56430..e0bb80733 100644 --- a/thirdparty/faiss/faiss/index_factory.cpp +++ b/thirdparty/faiss/faiss/index_factory.cpp @@ -130,6 +130,11 @@ char get_trains_alone(const Index* coarse_quantizer) { if (dynamic_cast(coarse_quantizer)) { return 2; } +#ifdef FAISS_ENABLE_SVS + if (dynamic_cast(coarse_quantizer)) { + return 2; + } +#endif return 2; // for complicated indexes, we assume they can't be used as a // kmeans index } @@ -299,6 +304,28 @@ Index* parse_coarse_quantizer( int R = std::stoi(sm[2]); return new IndexNSGFlat(d, R, mt); } +#ifdef FAISS_ENABLE_SVS + if (match("IVF([0-9]+[kM]?)_SVSVamana([0-9]*)(_.+)?")) { + nlist = parse_nlist(sm[1].str()); + int degree = sm[2].length() > 0 ? std::stoi(sm[2]) : 32; + SVSStorageKind storage = SVSStorageKind::SVS_FP32; + if (sm[3].matched) { + std::string s = sm[3].str().substr(1); + if (s == "SQI8") { + storage = SVSStorageKind::SVS_SQI8; + } else if (s == "FP16") { + storage = SVSStorageKind::SVS_FP16; + } else if (s == "FP32") { + storage = SVSStorageKind::SVS_FP32; + } else { + FAISS_THROW_FMT( + "unsupported SVSVamana coarse quantizer storage: %s", + s.c_str()); + } + } + return new IndexSVSVamana(d, degree, mt, storage); + } +#endif if (match("IVF([0-9]+[kM]?)\\(Index([0-9])\\)")) { nlist = parse_nlist(sm[1].str()); int no = std::stoi(sm[2].str()); diff --git a/thirdparty/faiss/faiss/python/__init__.py b/thirdparty/faiss/faiss/python/__init__.py index 05b376efc..f3aa4268a 100644 --- a/thirdparty/faiss/faiss/python/__init__.py +++ b/thirdparty/faiss/faiss/python/__init__.py @@ -22,8 +22,8 @@ from faiss.array_conversions import * from faiss.extra_wrappers import kmin, kmax, pairwise_distances, rand, randint, \ lrand, randn, rand_smooth_vectors, eval_intersection, normalize_L2, \ - ResultHeap, knn, Kmeans, checksum, matrix_bucket_sort_inplace, bucket_sort, \ - merge_knn_results, MapInt64ToInt64, knn_hamming, \ + ResultHeap, knn, Kmeans, SuperKmeans, checksum, matrix_bucket_sort_inplace, \ + bucket_sort, merge_knn_results, MapInt64ToInt64, knn_hamming, \ pack_bitstrings, unpack_bitstrings @@ -34,6 +34,7 @@ logger = logging.getLogger(__name__) class_wrappers.handle_Clustering(Clustering) +class_wrappers.handle_SuperKMeans(SuperKMeans) class_wrappers.handle_Clustering1D(Clustering1D) class_wrappers.handle_MatrixStats(MatrixStats) class_wrappers.handle_IOWriter(IOWriter) diff --git a/thirdparty/faiss/faiss/python/class_wrappers.py b/thirdparty/faiss/faiss/python/class_wrappers.py index 66552ee6c..683969b36 100644 --- a/thirdparty/faiss/faiss/python/class_wrappers.py +++ b/thirdparty/faiss/faiss/python/class_wrappers.py @@ -126,6 +126,24 @@ def replacement_train_encoded(self, x, codec, index, weights=None): replace_method(the_class, 'train_encoded', replacement_train_encoded) +def handle_SuperKMeans(the_class): + + def replacement_train(self, x): + """Perform SuperKMeans clustering on a set of vectors. + + Parameters + ---------- + x : array_like + Training vectors, shape (n, self.d). `dtype` must be float32. + """ + n, d = x.shape + assert d == self.d + x = np.ascontiguousarray(x, dtype='float32') + self.train_c(n, swig_ptr(x)) + + replace_method(the_class, 'train', replacement_train) + + def handle_Clustering1D(the_class): def replacement_train_exact(self, x): @@ -770,6 +788,8 @@ def replacement_search_preassigned(self, x, k, Iq, Dq, *, params=None, D=None, I if Dq is not None: Dq = np.ascontiguousarray(Dq, dtype='float32') assert Dq.shape == Iq.shape + else: + Dq = np.zeros(Iq.shape, dtype='float32') self.search_preassigned_c( n, swig_ptr(x), @@ -823,6 +843,8 @@ def replacement_range_search_preassigned(self, x, thresh, Iq, Dq, *, params=None if Dq is not None: Dq = np.ascontiguousarray(Dq, dtype='float32') assert Dq.shape == Iq.shape + else: + Dq = np.zeros(Iq.shape, dtype='float32') thresh = float(thresh) res = RangeSearchResult(n) @@ -1000,6 +1022,8 @@ def replacement_search_preassigned(self, x, k, Iq, Dq): if Dq is not None: Dq = np.ascontiguousarray(Dq, dtype='int32') assert Dq.shape == Iq.shape + else: + Dq = np.zeros(Iq.shape, dtype='int32') self.search_preassigned_c( n, swig_ptr(x), @@ -1035,6 +1059,8 @@ def replacement_range_search_preassigned(self, x, thresh, Iq, Dq, *, params=None if Dq is not None: Dq = np.ascontiguousarray(Dq, dtype='int32') assert Dq.shape == Iq.shape + else: + Dq = np.zeros(Iq.shape, dtype='int32') thresh = int(thresh) res = RangeSearchResult(n) @@ -1298,11 +1324,9 @@ def __exit__(self, *ignored): def handle_SearchParameters(the_class): - """ this wrapper is to enable initializations of the form - SearchParametersXX(a=3, b=SearchParamsYY) - This also requires the enclosing class to keep a reference on the - sub-object, since the C++ code assumes the object ownership is - handled externally. + """Protect SearchParameters from leaking SWIG-owned sub-objects assigned + via either kwargs construction (SearchParametersXX(sel=x)) or bare + attribute assignment (params.sel = x). """ the_class.original_init = the_class.__init__ @@ -1310,13 +1334,33 @@ def replacement_init(self, **args): self.original_init() for k, v in args.items(): assert hasattr(self, k) - with RememberSwigOwnership(v): - setattr(self, k, v) - if type(v) not in (int, float, bool, str): - add_to_referenced_objects(self, v) + setattr(self, k, v) the_class.__init__ = replacement_init + # Install __setattr__ once per hierarchy; subclasses inherit via MRO. + if getattr(the_class, "_protected_setattr", False): + return + parent_setattr = the_class.__setattr__ + + def replacement_setattr(self, k, v): + # Per-field ref dict. Reassigning the same field drops the prior + # ref instead of accumulating, so a long-lived SearchParameters + # with repeated `params.sel = ...` does not leak. + if v is not None and hasattr(v, "thisown"): + if not hasattr(self, "_sp_field_refs"): + parent_setattr(self, "_sp_field_refs", {}) + with RememberSwigOwnership(v): + parent_setattr(self, k, v) + self._sp_field_refs[k] = v + else: + parent_setattr(self, k, v) + if hasattr(self, "_sp_field_refs"): + self._sp_field_refs.pop(k, None) + + the_class.__setattr__ = replacement_setattr + the_class._protected_setattr = True + def handle_IDSelectorSubset(the_class, class_owns, force_int64=True): the_class.original_init = the_class.__init__ diff --git a/thirdparty/faiss/faiss/python/extra_wrappers.py b/thirdparty/faiss/faiss/python/extra_wrappers.py index b6fc1dbd4..77b3266b2 100644 --- a/thirdparty/faiss/faiss/python/extra_wrappers.py +++ b/thirdparty/faiss/faiss/python/extra_wrappers.py @@ -602,6 +602,56 @@ def assign(self, x): return D.ravel(), I.ravel() +class SuperKmeans(Kmeans): + """Drop-in replacement for `Kmeans` that runs `SuperKMeans` (ADSampling + + PDX progressive pruning) instead of `Clustering`. Same `centroids`, `obj`, + `iteration_stats`, and `assign()` surface; additionally exposes + `gemm_pruning_rates`. + + kwargs are forwarded to `SuperKMeansParameters`. Fields not present on it + (e.g. `spherical`, `int_centroids`, `nredo`, `frozen_centroids`, + `init_method`, `update_index`, `early_stop_threshold`, + `progressive_dim_steps`, `gpu`) raise `AttributeError`. + """ + + def __init__(self, d, k, **kwargs): + self.d = d + self.reset(k) + self.gpu = False + self.cp = SuperKMeansParameters() + for key, v in kwargs.items(): + getattr(self.cp, key) + setattr(self.cp, key, v) + self.set_index() + + def set_index(self): + self.index = IndexFlatL2(self.d) + + def train(self, x, weights=None, init_centroids=None): + assert weights is None, "SuperKmeans does not support weights" + assert init_centroids is None, \ + "SuperKmeans does not support init_centroids" + x = np.ascontiguousarray(x, dtype='float32') + n, d = x.shape + assert d == self.d + + sc = SuperKMeans(d, self.k, self.cp) + sc.train(x) + + centroids = faiss.vector_to_array(sc.centroids) + self.centroids = centroids.reshape(self.k, d) + stats = sc.iteration_stats + stats = [stats.at(i) for i in range(stats.size())] + self.obj = np.array([st.obj for st in stats]) + stat_fields = 'obj time time_search imbalance_factor nsplit'.split() + self.iteration_stats = [ + {field: getattr(st, field) for field in stat_fields} + for st in stats + ] + self.gemm_pruning_rates = faiss.vector_to_array(sc.gemm_pruning_rates) + return self.obj[-1] if self.obj.size > 0 else 0.0 + + ########################################### # Packing and unpacking bitstrings ########################################### diff --git a/thirdparty/faiss/faiss/python/swigfaiss.swig b/thirdparty/faiss/faiss/python/swigfaiss.swig index 8ab89f5e9..12be4c343 100644 --- a/thirdparty/faiss/faiss/python/swigfaiss.swig +++ b/thirdparty/faiss/faiss/python/swigfaiss.swig @@ -188,6 +188,7 @@ typedef uint64_t size_t; #endif // !_MSC_VER #include +#include #include #include @@ -521,6 +522,7 @@ void gpu_sync_all_devices() %include %include %include +%include %include @@ -725,6 +727,25 @@ void gpu_sync_all_devices() %include %include #ifdef FAISS_ENABLE_CUVS +// IVFPQSearchCagraConfig.lut_dtype and .internal_distance_dtype are typed as +// cudaDataType_t (a C enum from ). Without an int typemap +// SWIG generates accessors that read/write a cudaDataType_t* SwigPyObject, +// which Python users cannot construct, so the fields are effectively +// read-only from Python and the documented values (CUDA_R_16F, CUDA_R_8U) +// are not exported. Treat cudaDataType_t as int for SWIG and expose the +// relevant enum values so: +// +// c = faiss.IVFPQSearchCagraConfig() +// c.lut_dtype = faiss.CUDA_R_16F +// +// works as expected. Values match cudaDataType_t in . +typedef int cudaDataType_t; +%apply int { cudaDataType_t }; +%constant int CUDA_R_32F = 0; +%constant int CUDA_R_64F = 1; +%constant int CUDA_R_16F = 2; +%constant int CUDA_R_8I = 3; +%constant int CUDA_R_8U = 8; %include %include #endif diff --git a/thirdparty/faiss/faiss/svs/IndexSVSVamana.cpp b/thirdparty/faiss/faiss/svs/IndexSVSVamana.cpp index 64d1b2767..c17d73348 100644 --- a/thirdparty/faiss/faiss/svs/IndexSVSVamana.cpp +++ b/thirdparty/faiss/faiss/svs/IndexSVSVamana.cpp @@ -31,6 +31,7 @@ #include #include +#include #include #include #include @@ -114,13 +115,32 @@ void IndexSVSVamana::add(idx_t n, const float* x) { if (!status.ok()) { FAISS_THROW_MSG(status.message()); } + + if (stored_vectors_valid) { + size_t prev = static_cast(ntotal) * d; + stored_vectors.resize(prev + static_cast(n) * d); + std::memcpy(stored_vectors.data() + prev, x, sizeof(float) * n * d); + } ntotal += n; } +void IndexSVSVamana::reconstruct(idx_t key, float* recons) const { + FAISS_THROW_IF_NOT_MSG( + key >= 0 && key < ntotal, + "IndexSVSVamana::reconstruct: key out of range"); + FAISS_THROW_IF_NOT_MSG( + stored_vectors_valid && !stored_vectors.empty(), + "IndexSVSVamana::reconstruct: stored_vectors unavailable " + "(invalidated by remove_ids or not restored after deserialization)"); + std::memcpy(recons, stored_vectors.data() + key * d, sizeof(float) * d); +} + void IndexSVSVamana::reset() { if (impl) { impl->reset(); } + stored_vectors.clear(); + stored_vectors_valid = true; is_trained = false; ntotal = 0; } @@ -189,6 +209,8 @@ size_t IndexSVSVamana::remove_ids(const IDSelector& sel) { size_t removed = 0; auto Status = impl->remove_selected(&removed, id_filter); ntotal -= removed; + stored_vectors.clear(); + stored_vectors_valid = false; return removed; } diff --git a/thirdparty/faiss/faiss/svs/IndexSVSVamana.h b/thirdparty/faiss/faiss/svs/IndexSVSVamana.h index 162088d73..0bbe0f978 100644 --- a/thirdparty/faiss/faiss/svs/IndexSVSVamana.h +++ b/thirdparty/faiss/faiss/svs/IndexSVSVamana.h @@ -31,6 +31,7 @@ #include #include +#include namespace faiss { @@ -111,6 +112,8 @@ struct IndexSVSVamana : Index { void add(idx_t n, const float* x) override; + void reconstruct(idx_t key, float* recons) const override; + void search( idx_t n, const float* x, @@ -137,6 +140,12 @@ struct IndexSVSVamana : Index { /* The actual SVS implementation */ svs_runtime::DynamicVamanaIndex* impl{nullptr}; + // The SVS runtime API does not expose vector retrieval, so we keep a copy + // of added vectors to support reconstruct(). When used as a coarse + // quantizer this holds only nlist centroids. + std::vector stored_vectors; + bool stored_vectors_valid{true}; + protected: /* Initializes the implementation*/ virtual void create_impl(); diff --git a/thirdparty/faiss/faiss/svs/IndexSVSVamanaLeanVec.cpp b/thirdparty/faiss/faiss/svs/IndexSVSVamanaLeanVec.cpp index e6b99f9ff..dfcb76e96 100644 --- a/thirdparty/faiss/faiss/svs/IndexSVSVamanaLeanVec.cpp +++ b/thirdparty/faiss/faiss/svs/IndexSVSVamanaLeanVec.cpp @@ -56,7 +56,7 @@ IndexSVSVamanaLeanVec::~IndexSVSVamanaLeanVec() { FAISS_ASSERT(status.ok()); training_data = nullptr; } - IndexSVSVamana::~IndexSVSVamana(); + // Base class destructor handles impl cleanup } void IndexSVSVamanaLeanVec::add(idx_t n, const float* x) { diff --git a/thirdparty/faiss/faiss/utils/simd_impl/distances_arm_sve.cpp b/thirdparty/faiss/faiss/utils/simd_impl/distances_arm_sve.cpp index f7dd56746..0902db345 100644 --- a/thirdparty/faiss/faiss/utils/simd_impl/distances_arm_sve.cpp +++ b/thirdparty/faiss/faiss/utils/simd_impl/distances_arm_sve.cpp @@ -655,7 +655,10 @@ void exhaustive_L2sqr_blas_cmax( ip_block.get(), &nyi); } - for (int64_t i = i0; i < i1; i++) { +#pragma omp parallel for schedule(static) if ((i1 - i0) >= 16) + for (int64_t i = static_cast(i0); + i < static_cast(i1); + i++) { const size_t count = j1 - j0; float* ip_line = ip_block.get() + (i - i0) * count; diff --git a/thirdparty/faiss/faiss/utils/simd_impl/distances_avx2.cpp b/thirdparty/faiss/faiss/utils/simd_impl/distances_avx2.cpp index 5731fbb3e..7f524f6da 100644 --- a/thirdparty/faiss/faiss/utils/simd_impl/distances_avx2.cpp +++ b/thirdparty/faiss/faiss/utils/simd_impl/distances_avx2.cpp @@ -1276,7 +1276,10 @@ void exhaustive_L2sqr_blas_cmax( ip_block.get(), &nyi); } - for (int64_t i = i0; i < i1; i++) { +#pragma omp parallel for schedule(static) if ((i1 - i0) >= 16) + for (int64_t i = static_cast(i0); + i < static_cast(i1); + i++) { float* ip_line = ip_block.get() + (i - i0) * (j1 - j0); _mm_prefetch((const char*)ip_line, _MM_HINT_NTA); diff --git a/thirdparty/faiss/faiss/utils/simd_impl/super_kmeans_dispatch.h b/thirdparty/faiss/faiss/utils/simd_impl/super_kmeans_dispatch.h index 06d4072a7..9f6614ab8 100644 --- a/thirdparty/faiss/faiss/utils/simd_impl/super_kmeans_dispatch.h +++ b/thirdparty/faiss/faiss/utils/simd_impl/super_kmeans_dispatch.h @@ -13,7 +13,7 @@ // scalar primary template; adding NEON/SVE means just adding a new // specialization file alongside the AVX ones. // -// Known perf gap: aarch64 (NEON/SVE) specializations are not implemented in v1. +// Known perf gap: aarch64 (NEON/SVE) specializations are not implemented yet. // aarch64 falls through to the scalar primary template. Validating SVE requires // a Graviton-class host; deferred to a focused follow-up. diff --git a/thirdparty/faiss/faiss/utils/simd_impl/super_kmeans_kernels.h b/thirdparty/faiss/faiss/utils/simd_impl/super_kmeans_kernels.h index 0eec99b94..b36c05f5c 100644 --- a/thirdparty/faiss/faiss/utils/simd_impl/super_kmeans_kernels.h +++ b/thirdparty/faiss/faiss/utils/simd_impl/super_kmeans_kernels.h @@ -27,15 +27,17 @@ inline float block_l2(const float* x, const float* y, int n) { return s; } -// Mirrors the impl-file guards; turns off-arch direct calls into a -// compile-time error instead of a linker error. -#if defined(__x86_64__) +// COMPILE_SIMD_* is a build-system define (link-time promise that the +// specialization will be available). Mirrors the impl-file guards. +#ifdef COMPILE_SIMD_AVX2 template <> float block_l2(const float* x, const float* y, int n); +#endif +#ifdef COMPILE_SIMD_AVX512 template <> float block_l2(const float* x, const float* y, int n); -#endif // __x86_64__ +#endif } // namespace detail } // namespace faiss diff --git a/thirdparty/faiss/faiss/utils/simd_impl/super_kmeans_kernels_avx2.cpp b/thirdparty/faiss/faiss/utils/simd_impl/super_kmeans_kernels_avx2.cpp index 47cff327c..d5759a75a 100644 --- a/thirdparty/faiss/faiss/utils/simd_impl/super_kmeans_kernels_avx2.cpp +++ b/thirdparty/faiss/faiss/utils/simd_impl/super_kmeans_kernels_avx2.cpp @@ -5,7 +5,7 @@ * LICENSE file in the root directory of this source tree. */ -#if defined(__x86_64__) +#ifdef COMPILE_SIMD_AVX2 #include @@ -54,4 +54,4 @@ float block_l2(const float* x, const float* y, int n) { } // namespace detail } // namespace faiss -#endif // __x86_64__ +#endif // COMPILE_SIMD_AVX2 diff --git a/thirdparty/faiss/faiss/utils/simd_impl/super_kmeans_kernels_avx512.cpp b/thirdparty/faiss/faiss/utils/simd_impl/super_kmeans_kernels_avx512.cpp index 4fe67b1d1..dcb155bac 100644 --- a/thirdparty/faiss/faiss/utils/simd_impl/super_kmeans_kernels_avx512.cpp +++ b/thirdparty/faiss/faiss/utils/simd_impl/super_kmeans_kernels_avx512.cpp @@ -5,7 +5,7 @@ * LICENSE file in the root directory of this source tree. */ -#if defined(__x86_64__) && defined(__AVX512F__) +#ifdef COMPILE_SIMD_AVX512 #include @@ -42,4 +42,4 @@ float block_l2(const float* x, const float* y, int n) { } // namespace detail } // namespace faiss -#endif // __x86_64__ && __AVX512F__ +#endif // COMPILE_SIMD_AVX512 diff --git a/thirdparty/faiss/tests/CMakeLists.txt b/thirdparty/faiss/tests/CMakeLists.txt index 24f9df0e4..282aaf770 100644 --- a/thirdparty/faiss/tests/CMakeLists.txt +++ b/thirdparty/faiss/tests/CMakeLists.txt @@ -146,7 +146,11 @@ target_link_libraries(faiss_test PRIVATE # Defines `gtest_discover_tests()`. include(GoogleTest) -if(WIN32) +if(CMAKE_CROSSCOMPILING) + # Skip test discovery when cross-compiling: the binary cannot run on the + # host, so gtest_discover_tests() would fail. The build itself verifies + # that the binary compiles and links correctly. +elseif(WIN32) # On Windows, defer test discovery to ctest time (PRE_TEST) instead of # build time (POST_BUILD) because faiss.dll is not in PATH during build. gtest_discover_tests(faiss_test DISCOVERY_MODE PRE_TEST) diff --git a/thirdparty/faiss/tests/test_clone.py b/thirdparty/faiss/tests/test_clone.py index 7ed047d54..06f71f87e 100644 --- a/thirdparty/faiss/tests/test_clone.py +++ b/thirdparty/faiss/tests/test_clone.py @@ -86,3 +86,22 @@ def test_AdditiveQuantizerFastScan(self): self.do_test_clone("RQ3x4fs_32_Nlsq2x4") self.do_test_clone("PLSQ2x3x4fs_Nlsq2x4") self.do_test_clone("PRQ2x3x4fs_Nrq2x4") + + def test_RowwiseMinMax(self): + # clone_Index for IndexRowwiseMinMaxBase was missing return res, so + # clone_index returned nullptr and leaked the clone allocation plus + # the deep-copied inner index. + d = 32 + n = 200 + rng = np.random.default_rng(42) + x = rng.standard_normal((n, d)).astype(np.float32) + + for factory in ("MinMax,SQ8", "MinMaxFP16,SQ8"): + codec = faiss.index_factory(d, factory) + codec.train(x) + codes = codec.sa_encode(x) + + codec2 = faiss.clone_index(codec) + self.assertEqual(type(codec), type(codec2)) + codes2 = codec2.sa_encode(x) + np.testing.assert_array_equal(codes, codes2) diff --git a/thirdparty/faiss/tests/test_contrib.py b/thirdparty/faiss/tests/test_contrib.py index fbe0473b4..ed1e4e827 100644 --- a/thirdparty/faiss/tests/test_contrib.py +++ b/thirdparty/faiss/tests/test_contrib.py @@ -464,6 +464,29 @@ def test_binary(self): l0, l1 = lims[q], lims[q + 1] self.assertTrue(set(I[q]) <= set(IR[l0:l1])) + def test_range_search_preassigned_coarse_dis_none_matches_zeros(self): + """coarse_dis=None must match coarse_dis=zeros. Before the fix, + np.empty left uninitialized values that IVFPQ by_residual=True + reads for distance correction, producing nondeterministic results.""" + ds = datasets.SyntheticDataset(32, 1000, 500, 50) + index = faiss.index_factory(ds.d, "IVF10,PQ4x4") + index.train(ds.get_train()) + index.add(ds.get_database()) + index.nprobe = 5 + self.assertTrue(index.by_residual) + xq = ds.get_queries() + _, list_nos = index.quantizer.search(xq, index.nprobe) + zero_dis = np.zeros(list_nos.shape, dtype=np.float32) + lz, Dz, Iz = ivf_tools.range_search_preassigned( + index, xq, 10.0, list_nos, zero_dis + ) + ln, Dn, In = ivf_tools.range_search_preassigned( + index, xq, 10.0, list_nos + ) + np.testing.assert_array_equal(lz, ln) + np.testing.assert_array_equal(Iz, In) + np.testing.assert_array_equal(Dz, Dn) + class TestRangeSearchMaxResults(unittest.TestCase): diff --git a/thirdparty/faiss/tests/test_fast_scan.py b/thirdparty/faiss/tests/test_fast_scan.py index 2dfeea457..502061d5b 100644 --- a/thirdparty/faiss/tests/test_fast_scan.py +++ b/thirdparty/faiss/tests/test_fast_scan.py @@ -26,7 +26,7 @@ class TestSearch(unittest.TestCase): def test_PQ4_accuracy(self): - ds = datasets.SyntheticDataset(32, 2000, 5000, 1000) + ds = datasets.SyntheticDataset(32, 2000, 5000, 1000) index_gt = faiss.IndexFlatL2(32) index_gt.add(ds.get_database()) @@ -45,14 +45,17 @@ def test_PQ4_search_matches_none(self): """Fast-scan PQ search dispatches via the 256/512-bit QBS path. The integer QBS kernel is deterministic across SIMD levels *only if* - the upstream float PQ distance table is bit-identical. On some - toolchains (cmake aarch64 NEON/SVE) the float distance-table - computation reduces in a different order at NEON/SVE vs scalar, - producing a non-bit-identical LUT and thus different integer - distances. We diagnose which case applies, then assert accordingly: - - LUT identical -> assert (D, I) matches exactly (tie-tolerantly). - - LUT diverges -> fall back to a recall@1 check (the integer - kernel can't be bit-exact if its input isn't). + the upstream float PQ distance table is bit-identical AND the + handler's end() method (int-to-float conversion) compiles to + identical FP code. Two cases break this: + + 1. aarch64 NEON/SVE: float distance-table computation reduces + in a different order, producing non-bit-identical LUT. + 2. DD mode: per-SIMD TUs compile with different flags (-mfma + vs no -mfma), so handler end() may contract mul+add into + fma, producing 1-ULP float distance differences. + + In both cases, fall back to a recall@1 check. """ ds = datasets.SyntheticDataset(32, 2000, 5000, 200) index = faiss.index_factory(32, 'PQ16x4fs') @@ -75,62 +78,80 @@ def test_PQ4_search_matches_none(self): D_none, I_none = index.search(xq, 10) lut_diffs = int((tab != tab_none).sum()) - if lut_diffs == 0: - # Float LUT is bit-identical -> integer kernel must be too - # (modulo tied-index ordering). + dd_mode = "DD" in faiss.get_compile_options() + if lut_diffs == 0 and not dd_mode: compare_binary_result_lists(D, I, D_none, I_none) else: - # Float LUT diverges across SIMD levels (float reductions are - # not order-stable across vector widths). Integer distances - # cannot be expected to match; check result-list overlap. recall_at_1 = (I[:, 0] == I_none[:, 0]).mean() self.assertGreater( recall_at_1, 0.95, - f"LUT diverges ({lut_diffs} entries differ); " + f"LUT diverges ({lut_diffs} entries differ, DD={dd_mode}); " f"recall@1 vs NONE = {recall_at_1:.3f}") - # This is an experiment to see if we can catch performance - # regressions. It runs 2 codes, one should be faster than the - # other by a factor ~10 in opt mode. We check for a factor 5. - # hopefully the jitter in executtion time will not produce - # too many spurious test failures. Unoptimized timings are - # not exploitable, hence the flag test on that as well. - # DD mode uses virtual dispatch which, combined with ASAN overhead, - # makes the 4x speed threshold unreliable on dev machines. - @unittest.skipUnless( - ('AVX2' in faiss.get_compile_options() or - 'AVX512' in faiss.get_compile_options() or - 'NEON' in faiss.get_compile_options()) and - "OPTIMIZE" in faiss.get_compile_options() and - "DD" not in faiss.get_compile_options(), - "only test in static optimized mode with avx2/avx512/neon") - def test_PQ4_speed(self): - ds = datasets.SyntheticDataset(32, 2000, 5000, 1000) - xt = ds.get_train() - xb = ds.get_database() - xq = ds.get_queries() +class TestFastScanDDSpeed(unittest.TestCase): + """Verify DD-mode FastScan kernels use real SIMD.""" - index = faiss.index_factory(32, 'PQ16x4') - index.train(xt) + def _build_index(self): + d, n, nq, k = 128, 100_000, 1000, 10 + np.random.seed(123) + xb = np.random.randn(n, d).astype("float32") + xq = np.random.randn(nq, d).astype("float32") + index = faiss.index_factory(d, "PQ16x4fs") + index.train(xb) index.add(xb) + index.search(xq, k) + return index, xq, k + def _bench(self, index, xq, k, level, iters=50): + faiss.SIMDConfig.set_level(level) t0 = time.time() - D1, I1 = index.search(xq, 10) - t1 = time.time() - pq_t = t1 - t0 - print('PQ16x4 search time:', pq_t) + for _ in range(iters): + index.search(xq, k) + return time.time() - t0 + + def _check_speedup(self, none_time, simd_time, label): + speedup = none_time / simd_time + print( + f"FastScan DD: NONE={none_time:.3f}s, " + f"{label}={simd_time:.3f}s, {speedup:.1f}x" + ) + # 2x is conservative: scalar-to-AVX2 theoretical max is ~32x, + # observed 5-27x in isolation. 2x leaves margin for CI noise. + self.assertGreater( + speedup, 2.0, + f"{label} should be >2x faster than NONE, " + f"got {speedup:.1f}x", + ) - index2 = faiss.index_factory(32, 'PQ16x4fs') - index2.train(xt) - index2.add(xb) + @unittest.skipUnless( + "DD" in faiss.get_compile_options() + and "OPTIMIZE" in faiss.get_compile_options() + and faiss.SIMDConfig.is_simd_level_available(faiss.SIMDLevel_AVX2), + "DD-only, needs opt mode and AVX2", + ) + def test_dd_fastscan_avx2_faster_than_none(self): + index, xq, k = self._build_index() + none_t = self._bench(index, xq, k, faiss.SIMDLevel_NONE) + avx2_t = self._bench(index, xq, k, faiss.SIMDLevel_AVX2) + faiss.SIMDConfig.set_level( + faiss.SIMDConfig.get_dispatched_level()) + self._check_speedup(none_t, avx2_t, "AVX2") - t0 = time.time() - D2, I2 = index2.search(xq, 10) - t1 = time.time() - pqfs_t = t1 - t0 - print('PQ16x4fs search time:', pqfs_t) - self.assertLess(pqfs_t * 4, pq_t) + @unittest.skipUnless( + "DD" in faiss.get_compile_options() + and "OPTIMIZE" in faiss.get_compile_options() + and faiss.SIMDConfig.is_simd_level_available(faiss.SIMDLevel_AVX512), + "DD-only, needs opt mode and AVX512", + ) + def test_dd_fastscan_avx512_faster_than_none(self): + index, xq, k = self._build_index() + none_t = self._bench(index, xq, k, faiss.SIMDLevel_NONE) + avx512_t = self._bench( + index, xq, k, faiss.SIMDLevel_AVX512) + faiss.SIMDConfig.set_level( + faiss.SIMDConfig.get_dispatched_level()) + self._check_speedup(none_t, avx512_t, "AVX512") @for_all_simd_levels @@ -235,7 +256,7 @@ def reference_accu(codes, LUT): a = np.uint16(0) for sp in range(0, nsp, 2): c = codes[j, sp // 2] - a += LUT[i, sp , c & 15].astype('uint16') + a += LUT[i, sp, c & 15].astype('uint16') a += LUT[i, sp + 1, c >> 4].astype('uint16') accu[i, j] = a return accu @@ -300,12 +321,12 @@ def test_22(self): ########################################################## def verify_with_draws(testcase, Dref, Iref, Dnew, Inew): - """ verify a list of results where there are draws in the distances (because - they are integer). """ + """Verify a list of results where there are draws in the distances + (because they are integer).""" np.testing.assert_array_almost_equal(Dref, Dnew, decimal=5) # here we have to be careful because of draws for i in range(len(Iref)): - if np.all(Iref[i] == Inew[i]): # easy case + if np.all(Iref[i] == Inew[i]): continue # we can deduce nothing about the latest line skip_dis = Dref[i, -1] @@ -476,7 +497,7 @@ def do_test_add(self, d, bbs): ds = datasets.SyntheticDataset(d, 2000, 5000, 200) - index = faiss.index_factory(d, f'PQ{d//2}x4np') + index = faiss.index_factory(d, f'PQ{d // 2}x4np') index.train(ds.get_train()) xb = ds.get_database() @@ -506,7 +527,7 @@ def test_constructor(self): d = 32 ds = datasets.SyntheticDataset(d, 2000, 5000, 200) - index = faiss.index_factory(d, f'PQ{d//2}x4np') + index = faiss.index_factory(d, f'PQ{d // 2}x4np') index.train(ds.get_train()) index.add(ds.get_database()) Dref, Iref = index.search(ds.get_queries(), 10) @@ -764,7 +785,7 @@ def test_issue_2739(self): index.train(ds.get_train()) index.add(ds.get_database()) - np.testing.assert_array_equal( - index.pq.decode(index.pq.compute_codes(ds.get_database()))[0, ::100], - index.reconstruct(0)[::100] - ) + decoded = index.pq.decode( + index.pq.compute_codes(ds.get_database()) + )[0, ::100] + np.testing.assert_array_equal(decoded, index.reconstruct(0)[::100]) diff --git a/thirdparty/faiss/tests/test_fastscan_filter.py b/thirdparty/faiss/tests/test_fastscan_filter.py index 3c78a65e0..97f7d938d 100644 --- a/thirdparty/faiss/tests/test_fastscan_filter.py +++ b/thirdparty/faiss/tests/test_fastscan_filter.py @@ -65,13 +65,13 @@ def do_test_filter( ]) inner_sel = faiss.IDSelectorBatch(excluded) sel = faiss.IDSelectorNot(inner_sel) - allowed = set(i for i in range(nb) if i not in excluded) + allowed = {i for i in range(nb) if i not in excluded} elif id_selector_type == "partial_batch": # Exclude a few IDs within a single block (not a whole block) excluded = np.array([5, 10, 20, 31], dtype="int64") inner_sel = faiss.IDSelectorBatch(excluded) sel = faiss.IDSelectorNot(inner_sel) - allowed = set(i for i in range(nb) if i not in excluded) + allowed = {i for i in range(nb) if i not in excluded} elif id_selector_type == "empty": sel = faiss.IDSelectorBatch(np.array([], dtype="int64")) allowed = set() @@ -213,7 +213,7 @@ def test_blockskip_consistency_with_ivfpq(self): ]) inner_sel = faiss.IDSelectorBatch(excluded) sel = faiss.IDSelectorNot(inner_sel) - allowed = set(i for i in range(nb) if i not in excluded) + allowed = {i for i in range(nb) if i not in excluded} # FastScan index index_fs = faiss.index_factory(d, "IVF32,PQ4x4fs") diff --git a/thirdparty/faiss/tests/test_index_accuracy.py b/thirdparty/faiss/tests/test_index_accuracy.py index 21ce3d447..6507457f5 100644 --- a/thirdparty/faiss/tests/test_index_accuracy.py +++ b/thirdparty/faiss/tests/test_index_accuracy.py @@ -655,7 +655,10 @@ def test_OIVFPQ(self): # verify same on OIVFPQ for r in 1, 10, 100: - assert e_oivfpq[r] >= e_ivfpq[r] + assert e_oivfpq[r] >= e_ivfpq[r] - 0.005, ( + f"recall@{r}: OPQ+IVFPQ ({e_oivfpq[r]:.4f}) " + f"< IVFPQ ({e_ivfpq[r]:.4f})" + ) class TestRoundoff(unittest.TestCase): diff --git a/thirdparty/faiss/tests/test_read_index_deserialize.cpp b/thirdparty/faiss/tests/test_read_index_deserialize.cpp index 993529310..2e20afeb9 100644 --- a/thirdparty/faiss/tests/test_read_index_deserialize.cpp +++ b/thirdparty/faiss/tests/test_read_index_deserialize.cpp @@ -22,6 +22,7 @@ #include #include #include +#include #include #include #include @@ -3798,6 +3799,107 @@ TEST(ReadIndexDeserialize, IndexLatticeCodeSizeTooLarge) { expect_read_throws_with(buf, "code_size"); } +// ----------------------------------------------------------------------- +// Test: ProductQuantizer with M*ksub exceeding the deserialization byte +// limit is rejected. ProductQuantizer::set_derived_values only validates +// d % M == 0, which is satisfied by any M when d == 0, so a crafted PQ +// header could carry an enormous M alongside an empty centroids vector. +// Without the read-side bound, downstream allocations sized M*ksub (e.g. +// IVFPQ residual norms in initialize_IVFPQ_precomputed_table) reach +// std::vector::max_size() and abort the process via std::length_error. +// ----------------------------------------------------------------------- +TEST(ReadIndexDeserialize, ProductQuantizerMKsubExceedsByteLimit) { + const size_t old_limit = get_deserialization_vector_byte_limit(); + set_deserialization_vector_byte_limit(1 << 20); // 1 MB + + std::vector buf; + push_fourcc(buf, "IxPq"); + push_index_header(buf, /*d=*/0, /*ntotal=*/0); + // PQ: d=0 (so 0 % M == 0 holds), M=1<<20, nbits=8 -> ksub=256 + // M * ksub = 2^28 floats = 1 GB, well above the 1 MB byte limit. + push_pq(buf, + /*d=*/0, + /*M=*/(size_t{1} << 20), + /*nbits=*/8, + /*centroids=*/{}); + + expect_read_throws_with(buf, "M*ksub"); + + set_deserialization_vector_byte_limit(old_limit); +} + +// ----------------------------------------------------------------------- +// Test: initialize_IVFPQ_precomputed_table rejects parameter combinations +// whose nlist * pq.M * pq.ksub * sizeof(float) computation overflows +// size_t. Without overflow-checked multiplication, the wrapped value +// silently bypasses the precomputed_table_max_bytes guard and the +// subsequent r_norms allocation aborts the process. +// ----------------------------------------------------------------------- +TEST(ReadIndexDeserialize, InitializeIVFPQPrecomputedTableOverflowRejected) { + // Construct a PQ with M*ksub already saturating most of size_t. The + // nlist factor (taken from the quantizer's ntotal) then forces overflow + // when multiplied in. + ProductQuantizer pq; + pq.d = 0; + pq.M = size_t{1} << 40; + pq.nbits = 24; + pq.ksub = size_t{1} << pq.nbits; + pq.dsub = 0; + pq.code_size = 0; + // Skip set_derived_values (which would resize centroids); we are + // exercising initialize_IVFPQ_precomputed_table's arithmetic guard, not + // PQ training. + + IndexFlatL2 quantizer(/*d_in=*/0); + quantizer.ntotal = size_t{1} << 20; // nlist + AlignedTable precomputed_table; + int use_precomputed_table = 0; + + EXPECT_THROW( + initialize_IVFPQ_precomputed_table( + use_precomputed_table, + &quantizer, + pq, + precomputed_table, + /*by_residual=*/true, + /*verbose=*/false), + faiss::FaissException); +} + +// ----------------------------------------------------------------------- +// Test: initialize_IVFPQ_precomputed_table rejects an overflowing M*ksub +// even when the caller has pre-set use_precomputed_table to 1, which +// bypasses the in-function size-class branch. Protects in-memory IVFPQ +// callers (training, IndexHNSW, IndexIVFPQFastScan) and the +// IndexIVFIndependentQuantizer deserialization path that reads +// use_precomputed_table directly from the stream. +// ----------------------------------------------------------------------- +TEST(ReadIndexDeserialize, + InitializeIVFPQPrecomputedTableUserSetFlagOverflowRejected) { + ProductQuantizer pq; + pq.d = 0; + pq.M = size_t{1} << 40; + pq.nbits = 24; + pq.ksub = size_t{1} << pq.nbits; + pq.dsub = 0; + pq.code_size = 0; + + IndexFlatL2 quantizer(/*d_in=*/0); + quantizer.ntotal = 1; // nlist + AlignedTable precomputed_table; + int use_precomputed_table = 1; // caller pre-selects table type 1 + + EXPECT_THROW( + initialize_IVFPQ_precomputed_table( + use_precomputed_table, + &quantizer, + pq, + precomputed_table, + /*by_residual=*/true, + /*verbose=*/false), + faiss::FaissException); +} + // ============================================================ // SVS fourcc rejection / deserialization safety (Group F: T262015608) // ============================================================ @@ -3806,11 +3908,12 @@ TEST(ReadIndexDeserialize, IndexLatticeCodeSizeTooLarge) { #include -// An invalid storage_kind value should be rejected at deserialization time -// with a FaissException, not abort via FAISS_ASSERT in to_svs_storage_kind(). -TEST(ReadIndexDeserialize, SVSVamanaInvalidStorageKind) { - std::vector buf; - push_fourcc(buf, "ISVD"); +namespace { + +/// Append the IndexSVSVamana on-disk fields up to and including storage_kind. +/// Caller may then push initialized=true (to trigger the SVS deserialize_impl +/// path) or any other trailing fields needed to reach the validation site. +void push_svs_vamana_prefix(std::vector& buf, int storage_kind) { push_index_header(buf, 8, 0); push_val(buf, 32); // graph_max_degree push_val(buf, 1.2f); // alpha @@ -3820,10 +3923,85 @@ TEST(ReadIndexDeserialize, SVSVamanaInvalidStorageKind) { push_val(buf, 750); // max_candidate_pool_size push_val(buf, 28); // prune_to push_val(buf, false); // use_full_search_history - push_val( - buf, - static_cast(SVS_count)); // storage_kind — first invalid value - push_val(buf, true); // initialized + push_val(buf, storage_kind); +} + +/// Append the IndexSVSIVF on-disk fields up to and including storage_kind. +/// Used to exercise the validation guard at the ISIQ/ISIL/ISID read sites. +void push_svs_ivf_prefix(std::vector& buf, int storage_kind) { + push_index_header(buf, 8, 0); + push_val(buf, 8); // num_centroids + push_val(buf, 64); // minibatch_size + push_val(buf, 1); // num_iterations + push_val(buf, false); // is_hierarchical + push_val(buf, 0.1f); // training_fraction + push_val(buf, 0); // hierarchical_level1_clusters + push_val(buf, 0); // seed + push_val(buf, 1); // n_probes + push_val(buf, 1.0f); // k_reorder (float, not size_t) + push_val(buf, 1); // num_threads + push_val(buf, 1); // intra_query_threads + push_val(buf, storage_kind); +} + +} // namespace + +// An invalid storage_kind value should be rejected at deserialization time +// with a FaissException, not abort via FAISS_ASSERT in to_svs_storage_kind(). +TEST(ReadIndexDeserialize, SVSVamanaInvalidStorageKind) { + std::vector buf; + push_fourcc(buf, "ISVD"); + push_svs_vamana_prefix(buf, /*storage_kind=*/static_cast(SVS_count)); + push_val(buf, true); // initialized + + expect_read_throws_with(buf, "storage_kind"); +} + +// The Vamana validator must reject negative storage_kind values too — the +// shared helper takes a signed int because READ1 reads four bytes that could +// be either sign. +TEST(ReadIndexDeserialize, SVSVamanaNegativeStorageKind) { + std::vector buf; + push_fourcc(buf, "ISVD"); + push_svs_vamana_prefix(buf, /*storage_kind=*/-1); + push_val(buf, true); + + expect_read_throws_with(buf, "storage_kind"); +} + +// IVF Vamana flavours (ISIQ = IndexSVSIVFLVQ, ISIL = IndexSVSIVFLeanVec, +// ISID = IndexSVSIVF) all share the same storage_kind read site and must +// reject out-of-range values at the deserialization boundary, mirroring the +// Vamana branch above. Without the shared validator the bad value would +// propagate into IndexSVSIVF::deserialize_impl and only get rejected from +// to_svs_storage_kind() after several allocations and an SVS-runtime call. +TEST(ReadIndexDeserialize, SVSIVFInvalidStorageKind) { + std::vector buf; + push_fourcc(buf, "ISID"); + push_svs_ivf_prefix(buf, /*storage_kind=*/static_cast(SVS_count)); + push_val(buf, true); // is_static + push_val(buf, true); // initialized + + expect_read_throws_with(buf, "storage_kind"); +} + +TEST(ReadIndexDeserialize, SVSIVFLVQInvalidStorageKind) { + std::vector buf; + push_fourcc(buf, "ISIQ"); + push_svs_ivf_prefix(buf, /*storage_kind=*/-1); + push_val(buf, true); + push_val(buf, true); + + expect_read_throws_with(buf, "storage_kind"); +} + +TEST(ReadIndexDeserialize, SVSIVFLeanVecInvalidStorageKind) { + std::vector buf; + push_fourcc(buf, "ISIL"); + push_svs_ivf_prefix(buf, /*storage_kind=*/static_cast(SVS_count) + 7); + push_val(buf, true); // is_static + push_val(buf, 8); // leanvec_d (only on ISIL path) + push_val(buf, true); // initialized expect_read_throws_with(buf, "storage_kind"); } @@ -4052,3 +4230,127 @@ TEST(ReadIndexDeserialize, SVSVamanaFourccRejected) { } #endif // FAISS_ENABLE_SVS + +// ----------------------------------------------------------------------- +// IndexIDMap deserialization invariant: id_map.size(), idxmap.ntotal, +// and the inner index's ntotal must all agree. The outer IDMap header +// and the inner index header carry independent ntotal fields on disk, +// so deserialization must reject payloads where they disagree. +// ----------------------------------------------------------------------- + +namespace { + +/// Build a serialized IxMp/IxM2 payload. The inner IxF2 (IndexFlat) carries +/// codes sized to match its own ntotal, so it passes the inner deserializer's +/// codes-vs-ntotal check and the IDMap-vs-inner-ntotal check is exercised +/// in isolation. +std::vector make_idmap_payload( + const char fourcc[4], + int64_t outer_ntotal, + int64_t inner_ntotal, + size_t id_map_size) { + constexpr int d = 4; + std::vector buf; + push_fourcc(buf, fourcc); + push_index_header(buf, d, outer_ntotal); + push_fourcc(buf, "IxF2"); + push_index_header(buf, d, inner_ntotal); + // IndexFlat codes are written via WRITEXBVECTOR: a size_t prefix of + // (byte_count / 4) followed by byte_count bytes. For inner_ntotal + // float vectors of dimension d, byte_count = inner_ntotal * d * 4. + const size_t byte_count = + static_cast(inner_ntotal) * d * sizeof(float); + push_val(buf, byte_count / 4); + buf.insert(buf.end(), byte_count, uint8_t{0}); + std::vector id_map(id_map_size, 0); + push_vector(buf, id_map); + return buf; +} + +std::vector make_binary_idmap_payload( + const char fourcc[4], + int64_t outer_ntotal, + int64_t inner_ntotal, + size_t id_map_size) { + constexpr int d = 8; + constexpr int code_size = d / 8; + std::vector buf; + push_fourcc(buf, fourcc); + push_binary_index_header(buf, d, outer_ntotal); + push_fourcc(buf, "IBxF"); + push_binary_index_header(buf, d, inner_ntotal); + std::vector codes(inner_ntotal * code_size, 0); + push_vector(buf, codes); + std::vector id_map(id_map_size, 0); + push_vector(buf, id_map); + return buf; +} + +} // namespace + +TEST(ReadIndexDeserialize, IndexIDMapInnerNtotalMismatchRejected) { + // Outer IDMap ntotal=0 agrees with id_map.size()=0, but the inner + // IndexFlat reports ntotal=3. + auto buf = make_idmap_payload( + "IxMp", /*outer_ntotal=*/0, /*inner_ntotal=*/3, /*id_map_size=*/0); + expect_read_throws_with(buf, "inner index ntotal"); +} + +TEST(ReadIndexDeserialize, IndexIDMap2InnerNtotalMismatchRejected) { + auto buf = make_idmap_payload( + "IxM2", /*outer_ntotal=*/0, /*inner_ntotal=*/3, /*id_map_size=*/0); + expect_read_throws_with(buf, "inner index ntotal"); +} + +TEST(ReadIndexDeserialize, IndexBinaryIDMapInnerNtotalMismatchRejected) { + auto buf = make_binary_idmap_payload( + "IBMp", /*outer_ntotal=*/0, /*inner_ntotal=*/3, /*id_map_size=*/0); + expect_binary_read_throws_with(buf, "inner index ntotal"); +} + +TEST(ReadIndexDeserialize, IndexBinaryIDMap2InnerNtotalMismatchRejected) { + auto buf = make_binary_idmap_payload( + "IBM2", /*outer_ntotal=*/0, /*inner_ntotal=*/3, /*id_map_size=*/0); + expect_binary_read_throws_with(buf, "inner index ntotal"); +} + +// A well-formed IxMp payload with matching ntotals deserializes cleanly. +TEST(ReadIndexDeserialize, IndexIDMapMatchingNtotalsAccepted) { + auto buf = make_idmap_payload( + "IxMp", /*outer_ntotal=*/0, /*inner_ntotal=*/0, /*id_map_size=*/0); + VectorIOReader reader; + reader.data = buf; + EXPECT_NO_THROW({ + auto idx = read_index_up(&reader); + EXPECT_NE(idx, nullptr); + }); +} + +// ----------------------------------------------------------------------- +// IndexIDMap::search_ex maps each inner-index label through id_map. When +// inner.ntotal exceeds id_map.size() (e.g. mid-construction, or after a +// recoverable failure in add_with_ids_ex / add_sa_codes / merge_from +// leaves the IDMap layer behind the inner index), inner labels can land +// outside id_map. Such labels must be remapped to the -1 "no result" +// sentinel. +// ----------------------------------------------------------------------- +TEST(IndexIDMapSafety, SearchRemapsOutOfRangeLabelsToSentinel) { + constexpr int d = 4; + constexpr int nb = 3; + std::vector xb(nb * d, 0.0f); + for (int i = 0; i < nb; i++) { + xb[i * d] = static_cast(i); + } + + IndexFlat inner(d, METRIC_L2); + IndexIDMap idmap(&inner); + // inner.ntotal == nb, idmap.id_map is empty. + inner.add(nb, xb.data()); + + std::vector xq(d, 0.0f); + std::vector distances(1); + std::vector labels(1, 0); + EXPECT_NO_THROW( + idmap.search(1, xq.data(), 1, distances.data(), labels.data())); + EXPECT_EQ(labels[0], -1); +} diff --git a/thirdparty/faiss/tests/test_referenced_objects.py b/thirdparty/faiss/tests/test_referenced_objects.py index a7d697a16..ae8c7366c 100644 --- a/thirdparty/faiss/tests/test_referenced_objects.py +++ b/thirdparty/faiss/tests/test_referenced_objects.py @@ -84,6 +84,77 @@ def test_shards(self): gc.collect() index.search(xb, 10) + def test_SearchParameters_setattr_no_leak(self): + # Regression: assigning IDSelectorBatch to params.sel after + # construction must not leak. Before the fix, the SWIG property + # setter flipped IDSelectorBatch.thisown to 0 but + # SearchParameters never freed it -> ~80 MB orphaned per iter. + n_ids = 5_000_000 + n_iters = 10 + + def body(): + params = faiss.SearchParameters() + params.sel = faiss.IDSelectorBatch(np.arange(n_ids)) + del params + + body() # warmup: prime allocator + gc.collect() + start_kb = faiss.get_mem_usage_kb() + for _ in range(n_iters): + body() + gc.collect() + growth_mb = (faiss.get_mem_usage_kb() - start_kb) / 1024.0 + # One leaked IDSelectorBatch is ~80 MB; 10 iters would leak + # ~800 MB. 200 MB cap leaves generous headroom for allocator + # noise but is well below the leak signal. + self.assertLess(growth_mb, 200) + + def test_SearchParameters_subclass_no_double_wrap(self): + # Regression: SWIG does not generate per-class __setattr__, so + # the ownership-protecting override must be installed only on + # the base SearchParameters; otherwise subclasses inherit and + # re-wrap, doubling the refcount on the assigned selector. + params = faiss.SearchParametersIVF() + sel = faiss.IDSelectorBatch(np.arange(10)) + before = sys.getrefcount(sel) + params.sel = sel + delta = sys.getrefcount(sel) - before + self.assertEqual(delta, 1) + + def test_SearchParameters_reassignment_no_accumulation(self): + # Reassigning the same field on a long-lived SearchParameters + # must release the prior ref. Without per-field tracking the + # protection list grows unbounded - reproducing the original + # leak shape under a different access pattern. + params = faiss.SearchParameters() + sel1 = faiss.IDSelectorBatch(np.arange(10)) + sel2 = faiss.IDSelectorBatch(np.arange(10)) + before1 = sys.getrefcount(sel1) + params.sel = sel1 + params.sel = sel2 + # sel1 ref dropped when sel2 took its slot + self.assertEqual(sys.getrefcount(sel1), before1) + + def test_SearchParameters_setattr_else_branch(self): + # Else branch coverage: setting a field to None drops the prior + # SWIG ref; assigning two different fields keeps both alive + # independently so dropping one does not affect the other. + params = faiss.SearchParametersIVF() + sel = faiss.IDSelectorBatch(np.arange(10)) + quant = faiss.SearchParameters() + sel_before = sys.getrefcount(sel) + quant_before = sys.getrefcount(quant) + + params.sel = sel + params.quantizer_params = quant + self.assertEqual(sys.getrefcount(sel) - sel_before, 1) + self.assertEqual(sys.getrefcount(quant) - quant_before, 1) + + # Drop sel via None; quantizer_params untouched. + params.sel = None + self.assertEqual(sys.getrefcount(sel), sel_before) + self.assertEqual(sys.getrefcount(quant) - quant_before, 1) + dbin = 32 xtbin = np.random.randint(256, size=(100, int(dbin / 8))).astype('uint8') diff --git a/thirdparty/faiss/tests/test_super_kmeans.py b/thirdparty/faiss/tests/test_super_kmeans.py new file mode 100644 index 000000000..f0a339574 --- /dev/null +++ b/thirdparty/faiss/tests/test_super_kmeans.py @@ -0,0 +1,188 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import unittest + +import faiss +import numpy as np +from faiss.contrib.datasets import SyntheticDataset + + +class SuperKMeansTest(unittest.TestCase): + def test_train_rejects_too_few_points(self): + sc = faiss.SuperKMeans(128, 16) + x = SyntheticDataset(128, 10, 0, 0).get_train() # n=10 < k=16 + with self.assertRaises(RuntimeError): + sc.train(x) + + def test_objective_decreases_monotonically(self): + d, k, n = 32, 8, 1000 + x = SyntheticDataset(d, n, 0, 0).get_train() + + p = faiss.SuperKMeansParameters() + p.seed = 42 + p.niter = 10 + sc = faiss.SuperKMeans(d, k, p) + sc.train(x) + + stats = sc.iteration_stats + objs = [stats.at(i).obj for i in range(stats.size())] + for i in range(1, len(objs)): + self.assertLessEqual(objs[i], objs[i - 1] * 1.001) + + def test_pruned_iter_assignments_close_to_vanilla(self): + d, k, n = 64, 16, 2000 + x = SyntheticDataset(d, n, 0, 0).get_train() + + p = faiss.SuperKMeansParameters() + p.seed = 42 + p.niter = 10 + sc = faiss.SuperKMeans(d, k, p) + sc.train(x) + + quant = faiss.IndexFlatL2(d) + vanilla = faiss.Clustering(d, k) + vanilla.seed = 42 + vanilla.niter = 10 + vanilla.train(x, quant) + + sc_final = sc.iteration_stats.at(sc.iteration_stats.size() - 1).obj + v_stats = vanilla.iteration_stats + v_final = v_stats.at(v_stats.size() - 1).obj + self.assertLess(abs(sc_final - v_final) / v_final, 0.05) + + def test_pruning_rate_in_expected_range(self): + # The 10-d intrinsic manifold of SyntheticDataset gives ADSampling + # the assigned-vs-other distance gap it needs to prune effectively. + d, k, n = 64, 64, 5000 + x = SyntheticDataset(d, n, 0, 0).get_train() + + p = faiss.SuperKMeansParameters() + p.seed = 42 + p.niter = 5 + sc = faiss.SuperKMeans(d, k, p) + sc.train(x) + + rates = faiss.vector_to_array(sc.gemm_pruning_rates) + self.assertEqual(rates.shape, (5,)) + self.assertEqual(rates[0], 0.0) # iter 0: no pruning + # Lower bound 0.30: looser than the 0.93 design target because + # synthetic d=64 data lacks the cluster gap real embeddings have at + # d=768; the assertion still confirms meaningful pruning occurs. + for i in range(2, 5): + self.assertGreaterEqual(rates[i], 0.30) + + def test_pdx_offset_correct_for_unaligned_y_batch(self): + # PDX block layout is sized for the full k centroids, not per y-batch + # tile. The y-batch tile boundary must be invisible to PDX-offset + # arithmetic. k=70 with y_batch=32 forces a 6-wide tail tile; assert + # the resulting centroids are identical to running with y_batch=70 + # (single tile, no tail). + d, k, n = 64, 70, 5000 + x = SyntheticDataset(d, n, 0, 0).get_train() + + def train_with_y_batch(y_batch): + p = faiss.SuperKMeansParameters() + p.seed = 1234 + p.niter = 5 + p.y_batch = y_batch + sc = faiss.SuperKMeans(d, k, p) + sc.train(x) + return faiss.vector_to_array(sc.centroids) + + tiled = train_with_y_batch(32) # k % y_batch = 6 + single = train_with_y_batch(k) # k % y_batch = 0 + np.testing.assert_array_equal(tiled, single) + + def test_determinism(self): + # Bit-exact reproducibility requires MKL_CBWR=COMPATIBLE and a fixed + # OMP thread count. The BUCK target sets these in `env`; for OSS + # ctest the user must export them before running. + d, k, n = 128, 64, 5000 + x = SyntheticDataset(d, n, 0, 0).get_train() + + p = faiss.SuperKMeansParameters() + p.niter = 5 + p.seed = 1234 + + sc1 = faiss.SuperKMeans(d, k, p) + sc1.train(x) + sc2 = faiss.SuperKMeans(d, k, p) + sc2.train(x) + + np.testing.assert_array_equal( + faiss.vector_to_array(sc1.centroids), + faiss.vector_to_array(sc2.centroids), + ) + np.testing.assert_array_equal( + faiss.vector_to_array(sc1.gemm_pruning_rates), + faiss.vector_to_array(sc2.gemm_pruning_rates), + ) + + def test_rejects_nan_input(self): + # check_input_data_for_NaNs defaults to True. + sc = faiss.SuperKMeans(128, 64) + x = SyntheticDataset(128, 5000, 0, 0).get_train() + x.flat[42] = float("nan") + with self.assertRaises(RuntimeError): + sc.train(x) + + def test_cluster_splitting_with_many_duplicates(self): + # 80% duplicates force empty clusters after the first assignment; + # split_clusters must fire and produce distinct centroids. + d, k, n = 64, 16, 1000 + n_dup = int(n * 0.8) + + ds = SyntheticDataset(d, (n - n_dup) + 1, 0, 0).get_train() + x = np.empty((n, d), dtype="float32") + x[:n_dup] = ds[0] + x[n_dup:] = ds[1:] + + p = faiss.SuperKMeansParameters() + p.seed = 42 + p.niter = 5 + sc = faiss.SuperKMeans(d, k, p) + sc.train(x) + + stats = sc.iteration_stats + total_splits = sum(stats.at(i).nsplit for i in range(stats.size())) + self.assertGreater(total_splits, 0) + + centroids = faiss.vector_to_array(sc.centroids).reshape(k, d) + unique_centroids = {centroids[j].tobytes() for j in range(k)} + self.assertEqual(len(unique_centroids), k) + + +class SuperKmeansWrapperTest(unittest.TestCase): + def test_train_populates_kmeans_surface(self): + d, k, n = 64, 16, 2000 + x = SyntheticDataset(d, n, 0, 0).get_train() + + km = faiss.SuperKmeans(d, k, niter=5, seed=42) + final_obj = km.train(x) + + self.assertEqual(km.centroids.shape, (k, d)) + self.assertEqual(km.obj.shape, (5,)) + self.assertEqual(len(km.iteration_stats), 5) + self.assertEqual(final_obj, km.obj[-1]) + self.assertEqual(km.gemm_pruning_rates.shape, (5,)) + + D, I = km.assign(x) + self.assertEqual(I.shape, (n,)) + self.assertEqual(D.shape, (n,)) + self.assertTrue((I >= 0).all() and (I < k).all()) + + def test_final_obj_matches_vanilla_kmeans(self): + d, k, n = 64, 16, 2000 + x = SyntheticDataset(d, n, 0, 0).get_train() + + vanilla = faiss.Kmeans(d, k, niter=10, seed=42) + vanilla.train(x) + + sk = faiss.SuperKmeans(d, k, niter=10, seed=42) + sk.train(x) + + rel_diff = abs(sk.obj[-1] - vanilla.obj[-1]) / vanilla.obj[-1] + self.assertLess(rel_diff, 0.05) diff --git a/thirdparty/faiss/tests/test_super_kmeans_foundations.cpp b/thirdparty/faiss/tests/test_super_kmeans_foundations.cpp index 2e87aa287..4936237cd 100644 --- a/thirdparty/faiss/tests/test_super_kmeans_foundations.cpp +++ b/thirdparty/faiss/tests/test_super_kmeans_foundations.cpp @@ -171,35 +171,6 @@ TEST(BlockL2, ScalarMatchesReference) { } } -TEST(BlockL2, Avx512MatchesScalar) { -#if defined(__x86_64__) && defined(__AVX512F__) - constexpr int N = 64; - std::mt19937 rng(43); - std::uniform_real_distribution dist(-1.0f, 1.0f); - std::vector x(N), y(N); - for (int m = 0; m < N; ++m) { - x[m] = dist(rng); - y[m] = dist(rng); - } - - const float scalar = faiss::detail::block_l2( - x.data(), y.data(), N); - const float avx512 = faiss::detail::block_l2( - x.data(), y.data(), N); - EXPECT_NEAR(avx512, scalar, 1e-4f); - - for (int n = 1; n < N; ++n) { - const float s = faiss::detail::block_l2( - x.data(), y.data(), n); - const float a = faiss::detail::block_l2( - x.data(), y.data(), n); - EXPECT_NEAR(a, s, 1e-4f) << "n=" << n; - } -#else - GTEST_SKIP() << "AVX-512 test requires __AVX512F__"; -#endif -} - TEST(BlockL2, DispatchMatchesScalar) { constexpr int N = 64; std::mt19937 rng(44); diff --git a/thirdparty/faiss/tests/test_svs_py.py b/thirdparty/faiss/tests/test_svs_py.py index 6247debe9..f5e108041 100644 --- a/thirdparty/faiss/tests/test_svs_py.py +++ b/thirdparty/faiss/tests/test_svs_py.py @@ -308,6 +308,35 @@ def test_svs_factory_leanvec(self): self.assertEqual(index.leanvec_d, 64) +@unittest.skipIf(_SKIP_SVS, _SKIP_REASON) +class TestSVSIVFCoarseQuantizerFactory(unittest.TestCase): + """Test that IVF with SVSVamana coarse quantizer supports storage kinds""" + + def test_ivf_svsvamana_default(self): + index = faiss.index_factory(32, "IVF256_SVSVamana32,Flat") + self.assertEqual(index.d, 32) + self.assertEqual(index.nlist, 256) + + def test_ivf_svsvamana_sqi8(self): + index = faiss.index_factory(32, "IVF256_SVSVamana32_SQI8,Flat") + self.assertEqual(index.d, 32) + self.assertEqual(index.nlist, 256) + + def test_ivf_svsvamana_fp16(self): + index = faiss.index_factory(32, "IVF256_SVSVamana32_FP16,Flat") + self.assertEqual(index.d, 32) + self.assertEqual(index.nlist, 256) + + def test_ivf_svsvamana_fp32_explicit(self): + index = faiss.index_factory(32, "IVF256_SVSVamana32_FP32,Flat") + self.assertEqual(index.d, 32) + self.assertEqual(index.nlist, 256) + + def test_ivf_svsvamana_invalid_storage(self): + with self.assertRaises(RuntimeError): + faiss.index_factory(32, "IVF256_SVSVamana32_INVALID,Flat") + + @unittest.skipIf(_SKIP_SVS, _SKIP_REASON) class TestSVSAdapterFP16(TestSVSAdapter): """Repeat all tests for SVS Float16 variant""" @@ -644,7 +673,9 @@ def setUp(self): self.xq = np.random.random((self.nq, self.d)).astype('float32') def _create_instance(self): - return self.target_class(self.d, self.nlist) + idx = self.target_class(self.d, self.nlist) + idx.num_threads = 4 + return idx def test_ivf_construction(self): """Test construction and basic properties""" @@ -731,16 +762,20 @@ def test_ivf_reset(self): class TestSVSIVFAdapterFP16(TestSVSIVFAdapter): """Repeat IVF tests for FP16 variant""" def _create_instance(self): - return self.target_class(self.d, self.nlist, faiss.METRIC_L2, - faiss.SVS_FP16) + idx = self.target_class(self.d, self.nlist, faiss.METRIC_L2, + faiss.SVS_FP16) + idx.num_threads = 4 + return idx @unittest.skipIf(_SKIP_SVS, _SKIP_REASON) class TestSVSIVFAdapterSQI8(TestSVSIVFAdapter): """Repeat IVF tests for SQI8 variant""" def _create_instance(self): - return self.target_class(self.d, self.nlist, faiss.METRIC_L2, - faiss.SVS_SQI8) + idx = self.target_class(self.d, self.nlist, faiss.METRIC_L2, + faiss.SVS_SQI8) + idx.num_threads = 4 + return idx @unittest.skipIf(_SKIP_SVS_LL, _SKIP_SVS_LL_REASON) @@ -754,6 +789,7 @@ def setUpClass(cls): def _create_instance(self): idx = self.target_class(self.d, self.nlist) idx.storage_kind = faiss.SVS_LVQ4x4 + idx.num_threads = 4 return idx @@ -768,6 +804,7 @@ def setUpClass(cls): def _create_instance(self): idx = self.target_class(self.d, self.nlist) idx.storage_kind = faiss.SVS_LVQ4x8 + idx.num_threads = 4 return idx @@ -782,6 +819,7 @@ def setUpClass(cls): def _create_instance(self): idx = self.target_class(self.d, self.nlist) idx.storage_kind = faiss.SVS_LVQ8x0 + idx.num_threads = 4 return idx @@ -794,9 +832,11 @@ def setUpClass(cls): cls.target_class = faiss.IndexSVSIVFLeanVec def _create_instance(self): - return self.target_class( + idx = self.target_class( self.d, self.nlist, faiss.METRIC_L2, 0, faiss.SVS_LeanVec4x4) + idx.num_threads = 4 + return idx @unittest.skipIf(_SKIP_SVS, _SKIP_REASON) @@ -892,6 +932,7 @@ def test_ivf_parameter_serialization(self): index.num_iterations = 15 index.minibatch_size = 128 index.seed = 77 + index.num_threads = 4 # Train and add index.train(self.xb) @@ -1085,17 +1126,23 @@ def setUp(self): self.xb = np.random.random((self.nb, self.d)).astype('float32') self.xq = np.random.random((self.nq, self.d)).astype('float32') + def _create_static(self, storage=None): + if storage is None: + storage = faiss.SVS_FP32 + idx = faiss.IndexSVSIVF( + self.d, self.nlist, faiss.METRIC_L2, storage, True) + idx.num_threads = 4 + return idx + def test_static_ivf_construction(self): """Test construction sets is_static flag""" - index = faiss.IndexSVSIVF( - self.d, self.nlist, faiss.METRIC_L2, faiss.SVS_FP32, True) + index = self._create_static() self.assertTrue(index.is_static) self.assertFalse(index.is_trained) def test_static_ivf_train_and_search(self): """Train includes all data; search works normally""" - index = faiss.IndexSVSIVF( - self.d, self.nlist, faiss.METRIC_L2, faiss.SVS_FP32, True) + index = self._create_static() index.train(self.xb) self.assertTrue(index.is_trained) self.assertEqual(index.ntotal, self.nb) @@ -1107,8 +1154,7 @@ def test_static_ivf_train_and_search(self): def test_static_ivf_add_throws(self): """add() must raise on a static index""" - index = faiss.IndexSVSIVF( - self.d, self.nlist, faiss.METRIC_L2, faiss.SVS_FP32, True) + index = self._create_static() index.train(self.xb) extra = np.random.random((100, self.d)).astype('float32') with self.assertRaises(RuntimeError): @@ -1116,8 +1162,7 @@ def test_static_ivf_add_throws(self): def test_static_ivf_remove_throws(self): """remove_ids() must raise on a static index""" - index = faiss.IndexSVSIVF( - self.d, self.nlist, faiss.METRIC_L2, faiss.SVS_FP32, True) + index = self._create_static() index.train(self.xb) ids = np.arange(0, 10, dtype=np.int64) sel = faiss.IDSelectorArray(len(ids), faiss.swig_ptr(ids)) @@ -1126,8 +1171,7 @@ def test_static_ivf_remove_throws(self): def test_static_ivf_serialization(self): """Serialize/deserialize preserves is_static and search results""" - index = faiss.IndexSVSIVF( - self.d, self.nlist, faiss.METRIC_L2, faiss.SVS_FP32, True) + index = self._create_static() index.train(self.xb) D_before, I_before = index.search(self.xq, 4) @@ -1142,8 +1186,7 @@ def test_static_ivf_serialization(self): def test_static_ivf_fp16(self): """Static IVF with FP16 storage""" - index = faiss.IndexSVSIVF( - self.d, self.nlist, faiss.METRIC_L2, faiss.SVS_FP16, True) + index = self._create_static(faiss.SVS_FP16) index.train(self.xb) self.assertTrue(index.is_trained) self.assertTrue(index.is_static) @@ -1153,8 +1196,7 @@ def test_static_ivf_fp16(self): def test_static_ivf_reset(self): """reset() clears the static index so it can be retrained""" - index = faiss.IndexSVSIVF( - self.d, self.nlist, faiss.METRIC_L2, faiss.SVS_FP32, True) + index = self._create_static() index.train(self.xb) self.assertTrue(index.is_trained) @@ -1181,10 +1223,15 @@ def setUp(self): self.xb = np.random.random((self.nb, self.d)).astype('float32') self.xq = np.random.random((self.nq, self.d)).astype('float32') + def _create_static_lvq(self): + idx = faiss.IndexSVSIVFLVQ( + self.d, self.nlist, faiss.METRIC_L2, faiss.SVS_LVQ4x4, True) + idx.num_threads = 4 + return idx + def test_static_ivf_lvq_train_search(self): """Static IVF LVQ: train and search""" - index = faiss.IndexSVSIVFLVQ( - self.d, self.nlist, faiss.METRIC_L2, faiss.SVS_LVQ4x4, True) + index = self._create_static_lvq() index.train(self.xb) self.assertTrue(index.is_static) self.assertTrue(index.is_trained) @@ -1194,16 +1241,14 @@ def test_static_ivf_lvq_train_search(self): def test_static_ivf_lvq_add_throws(self): """Static IVF LVQ: add() must raise""" - index = faiss.IndexSVSIVFLVQ( - self.d, self.nlist, faiss.METRIC_L2, faiss.SVS_LVQ4x4, True) + index = self._create_static_lvq() index.train(self.xb) with self.assertRaises(RuntimeError): index.add(self.xb) def test_static_ivf_lvq_serialization(self): """Static IVF LVQ: serialization round-trip""" - index = faiss.IndexSVSIVFLVQ( - self.d, self.nlist, faiss.METRIC_L2, faiss.SVS_LVQ4x4, True) + index = self._create_static_lvq() index.train(self.xb) D_before, _ = index.search(self.xq, 4) @@ -1229,11 +1274,16 @@ def setUp(self): self.xb = np.random.random((self.nb, self.d)).astype('float32') self.xq = np.random.random((self.nq, self.d)).astype('float32') - def test_static_ivf_leanvec_train_search(self): - """Static IVF LeanVec: train and search""" - index = faiss.IndexSVSIVFLeanVec( + def _create_static_leanvec(self): + idx = faiss.IndexSVSIVFLeanVec( self.d, self.nlist, faiss.METRIC_L2, 0, faiss.SVS_LeanVec4x4, True) + idx.num_threads = 4 + return idx + + def test_static_ivf_leanvec_train_search(self): + """Static IVF LeanVec: train and search""" + index = self._create_static_leanvec() index.train(self.xb) self.assertTrue(index.is_static) self.assertTrue(index.is_trained) @@ -1243,18 +1293,14 @@ def test_static_ivf_leanvec_train_search(self): def test_static_ivf_leanvec_add_throws(self): """Static IVF LeanVec: add() must raise""" - index = faiss.IndexSVSIVFLeanVec( - self.d, self.nlist, faiss.METRIC_L2, 0, - faiss.SVS_LeanVec4x4, True) + index = self._create_static_leanvec() index.train(self.xb) with self.assertRaises(RuntimeError): index.add(self.xb) def test_static_ivf_leanvec_serialization(self): """Static IVF LeanVec: serialization round-trip""" - index = faiss.IndexSVSIVFLeanVec( - self.d, self.nlist, faiss.METRIC_L2, 0, - faiss.SVS_LeanVec4x4, True) + index = self._create_static_leanvec() index.train(self.xb) D_before, _ = index.search(self.xq, 4) @@ -1268,9 +1314,7 @@ def test_static_ivf_leanvec_serialization(self): def test_static_ivf_leanvec_ood_train_search(self): """Static IVF LeanVec OOD: train with queries and search""" - index = faiss.IndexSVSIVFLeanVec( - self.d, self.nlist, faiss.METRIC_L2, 0, - faiss.SVS_LeanVec4x4, True) + index = self._create_static_leanvec() index.train(self.xb, xq_train=self.xq) self.assertTrue(index.is_static) self.assertTrue(index.is_trained) diff --git a/thirdparty/faiss/tests/test_swig_wrapper.py b/thirdparty/faiss/tests/test_swig_wrapper.py index 44c9ab394..7e147f121 100644 --- a/thirdparty/faiss/tests/test_swig_wrapper.py +++ b/thirdparty/faiss/tests/test_swig_wrapper.py @@ -241,6 +241,43 @@ def test_sa_encode_preallocated_matches_auto_allocated(self): self.assertIs(codes_pre, buf) +class TestSearchPreassignedDqNone(unittest.TestCase): + """Dq=None must behave like Dq=zeros for search_preassigned and + range_search_preassigned (regression: swig_ptr(None) raised ValueError).""" + + def _check(self, index, xq, Iq, k, thresh, zero_dtype): + Dq_zero = np.zeros(Iq.shape, dtype=zero_dtype) + _, I_zero = index.search_preassigned(xq, k, Iq, Dq_zero) + _, I_none = index.search_preassigned(xq, k, Iq, None) + np.testing.assert_array_equal(I_zero, I_none) + lz, _, Iz = index.range_search_preassigned(xq, thresh, Iq, Dq_zero) + ln, _, In = index.range_search_preassigned(xq, thresh, Iq, None) + np.testing.assert_array_equal(lz, ln) + np.testing.assert_array_equal(Iz, In) + + def test_float_ivf(self): + rng = np.random.default_rng(42) + xb = rng.random((200, 16), dtype=np.float32) + xq = rng.random((10, 16), dtype=np.float32) + index = faiss.index_factory(16, "IVF4,Flat") + index.train(xb) + index.add(xb) + index.nprobe = 2 + _, Iq = index.quantizer.search(xq, 2) + self._check(index, xq, Iq, k=5, thresh=2.0, zero_dtype=np.float32) + + def test_binary_ivf(self): + rng = np.random.default_rng(42) + xb = rng.integers(0, 256, size=(200, 8), dtype=np.uint8) + xq = rng.integers(0, 256, size=(10, 8), dtype=np.uint8) + index = faiss.IndexBinaryIVF(faiss.IndexBinaryFlat(64), 64, 4) + index.train(xb) + index.add(xb) + index.nprobe = 2 + _, Iq = index.quantizer.search(xq, 2) + self._check(index, xq, Iq, k=5, thresh=20, zero_dtype=np.int32) + + @unittest.skipIf(faiss.swig_version() < 0x040000, "swig < 4 does not support Doxygen comments") class TestDoxygen(unittest.TestCase): diff --git a/thirdparty/faiss/tests/test_wheel_smoke.py b/thirdparty/faiss/tests/test_wheel_smoke.py index 5697fe236..2faca2f0f 100644 --- a/thirdparty/faiss/tests/test_wheel_smoke.py +++ b/thirdparty/faiss/tests/test_wheel_smoke.py @@ -101,6 +101,7 @@ def test_hnsw(self): np.random.seed(42) xb = np.random.random((n, d)).astype("float32") index = faiss.IndexHNSWFlat(d, 16) + index.hnsw.efSearch = 64 index.add(xb) assert index.ntotal == n D, I = index.search(xb[:5], 10)