Skip to content

Graph-structured Indices for Scalable, Fast, Fresh and Filtered Approximate Nearest Neighbor Search

License

Notifications You must be signed in to change notification settings

epeshared/DiskANN

 
 

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

DiskANN

DiskANN Main PyPI version Downloads shield License: MIT

DiskANN Paper DiskANN Paper DiskANN Paper

DiskANN is a suite of scalable, accurate and cost-effective approximate nearest neighbor search algorithms for large-scale vector search that support real-time changes and simple filters. This code is based on ideas from the DiskANN, Fresh-DiskANN and the Filtered-DiskANN papers with further improvements. This code forked off from code for NSG algorithm.

This project has adopted the Microsoft Open Source Code of Conduct. For more information see the Code of Conduct FAQ or contact opencode@microsoft.com with any additional questions or comments.

See guidelines for contributing to this project.

Linux build:

Install the following packages through apt-get

sudo apt install make cmake g++ libaio-dev libgoogle-perftools-dev clang-format libboost-all-dev

Install Intel MKL

Ubuntu 20.04 or newer

sudo apt install libmkl-full-dev

Earlier versions of Ubuntu

Install Intel MKL either by downloading the oneAPI MKL installer or using apt (we tested with build 2019.4-070 and 2022.1.2.146).

# OneAPI MKL Installer
wget https://registrationcenter-download.intel.com/akdlm/irc_nas/18487/l_BaseKit_p_2022.1.2.146.sh
sudo sh l_BaseKit_p_2022.1.2.146.sh -a --components intel.oneapi.lin.mkl.devel --action install --eula accept -s

Build

mkdir build && cd build && cmake -DCMAKE_BUILD_TYPE=Release .. && make -j 

AVX-512 BF16 (optional acceleration)

DiskANN includes an optional AVX-512 BF16-accelerated kernel for bf16 distance computations.

  • Compile-time: the AVX-512 BF16 kernel is enabled only when the compiler supports the required flags; it is compiled for a single source file (src/bf16_simd_kernels.cpp) so the rest of the project is not forced to use AVX-512.
  • Runtime: bf16 distance code automatically dispatches to the AVX-512 BF16 kernel only when the running CPU/OS supports AVX-512 BF16; otherwise it falls back to the scalar implementation.

You can control this with the following CMake options (non-MSVC builds):

  • Default (try to enable when supported):
    cmake -S . -B build -DCMAKE_BUILD_TYPE=Release -DDISKANN_AVX512BF16=ON
    cmake --build build -j
  • Force disable:
    cmake -S . -B build -DCMAKE_BUILD_TYPE=Release -DDISKANN_AVX512BF16=OFF
    cmake --build build -j
  • Force enable (fail configure if compiler does not support AVX-512 BF16 flags):
    cmake -S . -B build -DCMAKE_BUILD_TYPE=Release -DDISKANN_FORCE_AVX512BF16=ON
    cmake --build build -j

AMX BF16 (optional acceleration)

DiskANN includes an optional AMX BF16-accelerated kernel for bf16 inner-product computations.

  • Compile-time: the AMX BF16 kernel is enabled only when the compiler supports the required flags; it is compiled for a single source file (src/bf16_amx_kernels.cpp) so the rest of the project is not forced to use AMX.
  • Runtime: bf16 distance code automatically dispatches to the AMX kernel only when the running CPU/OS supports AMX and the current thread is permitted to use AMX tile state (Linux arch_prctl request). If unavailable, it falls back to AVX-512 BF16 (if enabled) and then scalar.

You can control this with the following CMake options (non-MSVC builds):

  • Default (try to enable when supported):
    cmake -S . -B build-amx -DCMAKE_BUILD_TYPE=Release -DDISKANN_AMXBF16=ON
    cmake --build build-amx -j
  • Force disable:
    cmake -S . -B build -DCMAKE_BUILD_TYPE=Release -DDISKANN_AMXBF16=OFF
    cmake --build build -j
  • Force enable (fail configure if compiler does not support AMX flags):
    cmake -S . -B build-amx -DCMAKE_BUILD_TYPE=Release -DDISKANN_FORCE_AMXBF16=ON
    cmake --build build-amx -j

AVX-512 vs AMX (build one or the other)

If you want to do a strict A/B build where only one ISA path is compiled/used, configure two separate build directories.

  • AVX-512 BF16 only (no AMX codegen): bash cmake -S . -B build-avx512 -DCMAKE_BUILD_TYPE=Release \ -DDISKANN_FORCE_AVX512BF16=ON \ -DDISKANN_AMXBF16=OFF cmake --build build-avx512 -j

  • AMX BF16 only (no AVX-512 BF16 code path): bash cmake -S . -B build-amx -DCMAKE_BUILD_TYPE=Release \ -DDISKANN_FORCE_AMXBF16=ON \ -DDISKANN_AVX512BF16=OFF cmake --build build-amx -j

Note: some toolchains/build scripts add global -march=native. When AMX is disabled (-DDISKANN_AMXBF16=OFF), DiskANN explicitly compiles the AMX translation unit with -mno-amx-tile/-mno-amx-bf16 (when supported) to avoid accidentally emitting AMX instructions.

RaBitQ main-search approximate scoring (optional, runtime-gated)

DiskANN also supports using RaBitQ multi-bit codes as the main traversal approximate scorer in SSD search (inside PQFlashIndex::cached_beam_search).

  • Default behavior is unchanged: traversal uses the existing PQ distance lookup.
  • When enabled, traversal scoring uses RaBitQ approximate inner product (converted to a “distance” as -ip) while keeping the rest of the search logic intact.

Runtime enable

Set:

export DISKANN_USE_RABITQ_MAIN_APPROX=1

If the environment variable is set but RaBitQ main codes are missing or incompatible, DiskANN prints a one-time message and automatically falls back to PQ.

Main code file naming

Preferred sidecar file name:

<index_filepath>_rabitq_main.bin

For example, if your SSD index file is foo_disk.index, the RaBitQ main code file should be foo_disk.index_rabitq_main.bin.

Generating main codes during disk index build

You can generate the main-search sidecar automatically as part of disk index build:

./build/apps/build_disk_index \
    ... \
    --dist_fn mips \
    --build_rabitq_main_codes \
    --rabitq_nb_bits 4

This produces:

<index_path_prefix>_disk.index_rabitq_main.bin

Constraints

  • Currently supported only for dist_fn=mips / Metric::INNER_PRODUCT.
  • The RaBitQ code dim must match the index _data_dim (post any preprocessing/augmentation), otherwise main-search RaBitQ is disabled and the search falls back to PQ.
  • Ensure you run the updated search_disk_index/build_disk_index binaries from the same build directory that contains this feature.

Windows build:

The Windows version has been tested with Enterprise editions of Visual Studio 2022, 2019 and 2017. It should work with the Community and Professional editions as well without any changes.

Prerequisites:

  • CMake 3.15+ (available in VisualStudio 2019+ or from https://cmake.org)
  • NuGet.exe (install from https://www.nuget.org/downloads)
    • The build script will use NuGet to get MKL, OpenMP and Boost packages.
  • DiskANN git repository checked out together with submodules. To check out submodules after git clone:
git submodule init
git submodule update
  • Environment variables:
    • [optional] If you would like to override the Boost library listed in windows/packages.config.in, set BOOST_ROOT to your Boost folder.

Build steps:

  • Open the "x64 Native Tools Command Prompt for VS 2019" (or corresponding version) and change to DiskANN folder
  • Create a "build" directory inside it
  • Change to the "build" directory and run
cmake ..

OR for Visual Studio 2017 and earlier:

<full-path-to-installed-cmake>\cmake ..

This will create a diskann.sln solution. Now you can:

  • Open it from VisualStudio and build either Release or Debug configuration.
  • <full-path-to-installed-cmake>\cmake --build build
  • Use MSBuild:
msbuild.exe diskann.sln /m /nologo /t:Build /p:Configuration="Release" /property:Platform="x64"
  • This will also build gperftools submodule for libtcmalloc_minimal dependency.
  • Generated binaries are stored in the x64/Release or x64/Debug directories.

Usage:

Please see the following pages on using the compiled code:

Please cite this software in your work as:

@misc{diskann-github,
   author = {Simhadri, Harsha Vardhan and Krishnaswamy, Ravishankar and Srinivasa, Gopal and Subramanya, Suhas Jayaram and Antonijevic, Andrija and Pryce, Dax and Kaczynski, David and Williams, Shane and Gollapudi, Siddarth and Sivashankar, Varun and Karia, Neel and Singh, Aditi and Jaiswal, Shikhar and Mahapatro, Neelam and Adams, Philip and Tower, Bryan and Patel, Yash}},
   title = {{DiskANN: Graph-structured Indices for Scalable, Fast, Fresh and Filtered Approximate Nearest Neighbor Search}},
   url = {https://github.com/Microsoft/DiskANN},
   version = {0.6.1},
   year = {2023}
}

About

Graph-structured Indices for Scalable, Fast, Fresh and Filtered Approximate Nearest Neighbor Search

Resources

License

Code of conduct

Contributing

Security policy

Stars

Watchers

Forks

Packages

No packages published

Languages

  • C++ 65.6%
  • Rust 22.8%
  • Python 8.8%
  • CMake 1.6%
  • Shell 1.0%
  • Dockerfile 0.1%
  • C 0.1%