From ec4d42e7f69a76f2dfd2e1eab6b9cf42d9e33a93 Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Tue, 15 Aug 2023 16:55:38 +0200 Subject: [PATCH] gpu hnsw add gpu indexing flag add gpu vector storage time measurement add chmod flag build hnsw using gpu scoring speedup move gl_GlobalInvocationID to mains glslc with o flag fix tests use gpu atomics instead of multiple runs enable atomics add timers profile search let shader work while points collection use generations for visited flags add gpu init timer avoid entries download remove gpu call timers for a while rename to combined graph builder move gpu graph builder restore levels in link access restore links_map add gpu start condition use graph builder links memory remove mut's list point ids before processing download-upload links using methods only cpu graph builder as separate struct add cpu threads count parallel cpu build single threaded cpu preprocess try to fix glove accuracy unsafe Send and Sync for gpu builder copy graph layers builder instead of moving remove obsolete mut fix build move builders to arc fix bfb deadlock move uneccessary clear run cpu and gpu in separate threads many vectors storage buffers debug gpu vector storage fix vector storage upload bug add cpu ful graph build don't use non-ready links for point try to resolve conflicts while upsert add indexing unit test for measurements add test utilize cpu more efficient dont clear bhaep debug gpu runs try to use vec4 instead of mat4 dump scores count estimate gpu usage print usage per run start using working groups build shaders for vulkan v1.3 start shaders subgroup vk instance version 1.3 subgroup vector storage fix test nearest heap nearest heap test are you happy fmt add assert to nearest heap test are you happy fmt gpu visited flags start candidates heap candidates heap push debug nearest heap remove subgroupExclusiveMin fix build are you happy clippy remove links dependency from nearest heap provide input count to nearest test are you happy fmt provide ef to shader are you happy fmt get subgroup size move compiled shaders are you happy fmt fix workgroups count for vector storage compute test return true subgroup size do nearest heap on gpu with partial sorting remove tmp shader unite scores and indices are you happy fmt gpu nearest sorting are you happy fmt fix nearest heap test links test with subgroups gpu candidates start fix candidates sorting are you happy fmt fix candidates test on m1 chip move obsolete and test shaders start new gpu search context search context test fix links uploading debug hnsw searh on level test fix gpu hnsw search on level test move test initialization to separate function gpu greedy search shader test test heuristic shader gpu greedy search test gpu test heuristic move greedy search into searcher start insertion shader find hnsw patch on gpu fix gpu patches test start gpu graph builder equivalency test fix layers check add timers are you happy fmt quality test and fix multithreaded bug debug tests are you happy fmt fix levels count and clear links fix equivalency test remove obsolete solution trivial gpu indexing integration gpu config and skip conflicts for large hnsw glsl generic vector storage element type gpu generic vector storage enable vulkan features for f16 and u8 force_half_precision option add force f16 test fix u8 and f16 unit tests apply generic shaders in hnsw construction remove conflicts check; start cpu utilization cpu prebuild upload first point links to gpu mark all points in cpu graph as ready multithreaded cpu and gpu fix equivalency test remove cpu links copy parallel patches clearing fix multithreaded env external layers source more logs to debug cpu+gpu syncronization clear gpu processed points count remove some temp logs fix cpu stop condition bug new insert vector shader greedy search returns only ids download whole layer links apply new links on gpu directly update win shader script parallel links loading provide memory amount as setting and estimate groups count start bq gpu bq test fix tests shaders invocations refactor revert max groups settings gix build after rebase fix bug after rebase upload links to gpu in cpu thread fix greedy search iteration atomic checker for conflicts dump graph changes for debugging move iterator to separate file start combined builder refactor start cpu builder as separate struct separate gpu and cpu construction fix tests fix atomics count bug upload links to gpu as separate fn glue cpu and gpu gpu thread link points on cpu while gpu is busy fix gpu bq scoring test use xor instead of calc it reallocate with less groups if not enough gpu memory update compiled shaders don't use llvm gpu emulator debug panic if no gpu remove tmp debug dumping dockerfile for gpu customize candidates count fix gpu storage size runtime shader compilation dim as define nearest heap params as macro candidates and links capacities as macro parallel bitonic sort use bubble sort too nearest heap as shared limited shared candidates heap combine buble and bitonic sorts one-argumented similarity cache vector for scoring dont check candidates overflow do greedy before search visited hashtable Revert "visited hashtable" This reverts commit 7540d04d3859d02f17908c874927d18076f3eb8e. use less flushes bulk sum calculation Revert "bulk sum calculation" This reverts commit 6ef78ceb8d6ba7d9bf2077804548df162052c749. load vector in one read Revert "load vector in one read" This reverts commit ea980d0399b90f44638590d40fdbfb1d6ffa9504. fix build after rebase fix build after rebase sq support simplify bq uploading using constructed quantization metrics alignment depends on subgroup size are you happy fmt exact flag use lazy static for device access remove prints test reallocation factor visited flags capacity as a compile time constant build gpu hnsw for payload blocks add dynamically sized subgroup size support fix build update rust version for nvidia docker dynamically sized subgroups start device filter test device filter are you happy clippy devices manager provide device instead of flag, better locking partial cpu permit release free cpus when gpu enabled use cpu is waiting is disabled create device only when gpu is enabled cpu permit release all cases parallel indices count provide queue index small refactor of shader builder are you happy fmt quantization params as separate struct pq refactor fmt remove upload_mapped_ptr from gpu buffer refactor gpu buffer check buffer range buffer refactor refactor pipeline move shader compiler to instance refactor context more refactor no unwraps in gpu crate clippy apply pr changes are you happy fmt dont use shaderc in segment more gpu tests pq bq tests provide bq option fix bq for small dim bind pq buffers pq more complex test cover all storage types by tests test multi and vectors as iterator multivectors upload data multivectors shader fix multivector tests refactor shader builder move storage tests into separate file move quantization to separate file refactor gpu vector storage creation refactor gpu quantization api dont convert to float when unnecessary more comments small upload vectors refactor fix build after rebase fix tests after rebase different vk queue priorities capacity vectors upload stopper amd dockerfile fix build after rebase --- Cargo.lock | 2 + Cargo.toml | 3 +- DockerfileGpuAmd | 192 +++ DockerfileGpuNvidia | 195 +++ config/config.yaml | 3 + lib/common/common/src/cpu.rs | 13 + lib/segment/Cargo.toml | 3 + lib/segment/benches/multi_vector_search.rs | 1 + .../index/hnsw_index/gpu/batched_points.rs | 159 +++ .../index/hnsw_index/gpu/cpu_level_builder.rs | 260 ++++ .../index/hnsw_index/gpu/devices_manager.rs | 84 ++ .../hnsw_index/gpu/gpu_candidates_heap.rs | 217 +++ .../index/hnsw_index/gpu/gpu_graph_builder.rs | 304 ++++ .../index/hnsw_index/gpu/gpu_level_builder.rs | 211 +++ .../src/index/hnsw_index/gpu/gpu_links.rs | 298 ++++ .../index/hnsw_index/gpu/gpu_nearest_heap.rs | 259 ++++ .../hnsw_index/gpu/gpu_search_context.rs | 1234 +++++++++++++++++ .../index/hnsw_index/gpu/gpu_visited_flags.rs | 124 ++ lib/segment/src/index/hnsw_index/gpu/mod.rs | 345 +++++ .../index/hnsw_index/gpu/shaders/bheap.comp | 123 ++ .../gpu/shaders/candidates_heap.comp | 75 + .../hnsw_index/gpu/shaders/iterators.comp | 27 + .../index/hnsw_index/gpu/shaders/links.comp | 46 + .../hnsw_index/gpu/shaders/nearest_heap.comp | 205 +++ .../hnsw_index/gpu/shaders/run_get_patch.comp | 139 ++ .../gpu/shaders/run_greedy_search.comp | 44 + .../gpu/shaders/run_insert_vector.comp | 129 ++ .../gpu/shaders/search_context.comp | 140 ++ .../shaders/tests/test_candidates_heap.comp | 51 + .../gpu/shaders/tests/test_heuristic.comp | 50 + .../gpu/shaders/tests/test_hnsw_search.comp | 53 + .../gpu/shaders/tests/test_nearest_heap.comp | 58 + .../hnsw_index/gpu/shaders/visited_flags.comp | 37 + .../src/index/hnsw_index/graph_layers.rs | 33 + .../index/hnsw_index/graph_layers_builder.rs | 153 +- lib/segment/src/index/hnsw_index/hnsw.rs | 208 ++- .../tests/test_graph_connectivity.rs | 1 + .../segment_constructor/segment_builder.rs | 20 + .../segment_constructor_base.rs | 3 + .../quantized/quantized_vectors.rs | 10 +- .../tests/integration/batch_search_test.rs | 1 + .../integration/byte_storage_hnsw_test.rs | 1 + .../byte_storage_quantization_test.rs | 1 + .../tests/integration/exact_search_test.rs | 1 + .../tests/integration/filtrable_hnsw_test.rs | 1 + .../tests/integration/hnsw_discover_test.rs | 2 + .../integration/hnsw_quantized_search_test.rs | 1 + .../multivector_filtrable_hnsw_test.rs | 1 + .../integration/multivector_hnsw_test.rs | 2 + .../multivector_quantization_test.rs | 1 + nvidia_icd.json | 7 + run.bat | 7 + run.sh | 9 + src/main.rs | 18 + src/settings.rs | 27 + 55 files changed, 5533 insertions(+), 59 deletions(-) create mode 100644 DockerfileGpuAmd create mode 100644 DockerfileGpuNvidia create mode 100644 lib/segment/src/index/hnsw_index/gpu/batched_points.rs create mode 100644 lib/segment/src/index/hnsw_index/gpu/cpu_level_builder.rs create mode 100644 lib/segment/src/index/hnsw_index/gpu/devices_manager.rs create mode 100644 lib/segment/src/index/hnsw_index/gpu/gpu_candidates_heap.rs create mode 100644 lib/segment/src/index/hnsw_index/gpu/gpu_graph_builder.rs create mode 100644 lib/segment/src/index/hnsw_index/gpu/gpu_level_builder.rs create mode 100644 lib/segment/src/index/hnsw_index/gpu/gpu_links.rs create mode 100644 lib/segment/src/index/hnsw_index/gpu/gpu_nearest_heap.rs create mode 100644 lib/segment/src/index/hnsw_index/gpu/gpu_search_context.rs create mode 100644 lib/segment/src/index/hnsw_index/gpu/gpu_visited_flags.rs create mode 100644 lib/segment/src/index/hnsw_index/gpu/shaders/bheap.comp create mode 100644 lib/segment/src/index/hnsw_index/gpu/shaders/candidates_heap.comp create mode 100644 lib/segment/src/index/hnsw_index/gpu/shaders/iterators.comp create mode 100644 lib/segment/src/index/hnsw_index/gpu/shaders/links.comp create mode 100644 lib/segment/src/index/hnsw_index/gpu/shaders/nearest_heap.comp create mode 100644 lib/segment/src/index/hnsw_index/gpu/shaders/run_get_patch.comp create mode 100644 lib/segment/src/index/hnsw_index/gpu/shaders/run_greedy_search.comp create mode 100644 lib/segment/src/index/hnsw_index/gpu/shaders/run_insert_vector.comp create mode 100644 lib/segment/src/index/hnsw_index/gpu/shaders/search_context.comp create mode 100644 lib/segment/src/index/hnsw_index/gpu/shaders/tests/test_candidates_heap.comp create mode 100644 lib/segment/src/index/hnsw_index/gpu/shaders/tests/test_heuristic.comp create mode 100644 lib/segment/src/index/hnsw_index/gpu/shaders/tests/test_hnsw_search.comp create mode 100644 lib/segment/src/index/hnsw_index/gpu/shaders/tests/test_nearest_heap.comp create mode 100644 lib/segment/src/index/hnsw_index/gpu/shaders/visited_flags.comp create mode 100644 nvidia_icd.json create mode 100644 run.bat create mode 100755 run.sh diff --git a/Cargo.lock b/Cargo.lock index 7a9f7320ba..f13a378efd 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4945,6 +4945,7 @@ dependencies = [ "prost 0.11.9", "pyroscope", "pyroscope_pprofrs", + "quantization", "raft", "raft-proto", "rand 0.8.5", @@ -5973,6 +5974,7 @@ dependencies = [ "is_sorted", "issues", "itertools 0.13.0", + "lazy_static", "log", "macro_rules_attribute", "memmap2 0.9.5", diff --git a/Cargo.toml b/Cargo.toml index 13d022c165..23867125c6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -17,7 +17,7 @@ default-run = "qdrant" workspace = true [features] -default = ["web", "parking_lot"] +default = ["web", "parking_lot", "gpu"] web = ["actix-web"] multiling-chinese = ["segment/multiling-chinese"] multiling-japanese = ["segment/multiling-japanese"] @@ -100,6 +100,7 @@ common = { path = "lib/common/common" } cancel = { path = "lib/common/cancel" } memory = { path = "lib/common/memory" } issues = { path = "lib/common/issues" } +quantization = { path = "lib/quantization" } segment = { path = "lib/segment" } collection = { path = "lib/collection" } storage = { path = "lib/storage" } diff --git a/DockerfileGpuAmd b/DockerfileGpuAmd new file mode 100644 index 0000000000..a0b19a531e --- /dev/null +++ b/DockerfileGpuAmd @@ -0,0 +1,192 @@ +# Cross-compiling using Docker multi-platform builds/images and `xx`. +# +# https://docs.docker.com/build/building/multi-platform/ +# https://github.com/tonistiigi/xx +FROM --platform=${BUILDPLATFORM:-linux/amd64} tonistiigi/xx AS xx + +# Utilizing Docker layer caching with `cargo-chef`. +# +# https://www.lpalmieri.com/posts/fast-rust-docker-builds/ +FROM --platform=${BUILDPLATFORM:-linux/amd64} lukemathwalker/cargo-chef:latest-rust-1.82.0 AS chef + + +FROM chef AS planner +WORKDIR /qdrant +COPY . . +RUN cargo chef prepare --recipe-path recipe.json + + +FROM chef AS builder +WORKDIR /qdrant + +COPY --from=xx / / + +# Relative order of `ARG` and `RUN` commands in the Dockerfile matters. +# +# If you pass a different `ARG` to `docker build`, it would invalidate Docker layer cache +# for the next steps. (E.g., the following steps may depend on a new `ARG` value, so Docker would +# have to re-execute them instead of using a cached layer from a previous run.) +# +# Steps in this stage are ordered in a way that should maximize Docker layer cache utilization, +# so, please, don't reorder them without prior consideration. 🥲 + +RUN apt-get update \ + && apt-get install -y clang lld cmake protobuf-compiler jq \ + && rustup component add rustfmt + +# `ARG`/`ENV` pair is a workaround for `docker build` backward-compatibility. +# +# https://github.com/docker/buildx/issues/510 +ARG BUILDPLATFORM +ENV BUILDPLATFORM=${BUILDPLATFORM:-linux/amd64} + +ARG MOLD_VERSION=2.31.0 + +RUN case "$BUILDPLATFORM" in \ + */amd64 ) PLATFORM=x86_64 ;; \ + */arm64 | */arm64/* ) PLATFORM=aarch64 ;; \ + * ) echo "Unexpected BUILDPLATFORM '$BUILDPLATFORM'" >&2; exit 1 ;; \ + esac; \ + \ + mkdir -p /opt/mold; \ + cd /opt/mold; \ + \ + TARBALL="mold-$MOLD_VERSION-$PLATFORM-linux.tar.gz"; \ + curl -sSLO "https://github.com/rui314/mold/releases/download/v$MOLD_VERSION/$TARBALL"; \ + tar -xf "$TARBALL" --strip-components 1; \ + rm "$TARBALL" + +# `ARG`/`ENV` pair is a workaround for `docker build` backward-compatibility. +# +# https://github.com/docker/buildx/issues/510 +ARG TARGETPLATFORM +ENV TARGETPLATFORM=${TARGETPLATFORM:-linux/amd64} + +RUN xx-apt-get install -y pkg-config gcc g++ libc6-dev libunwind-dev + +# Select Cargo profile (e.g., `release`, `dev` or `ci`) +ARG PROFILE=release + +# Enable crate features +ARG FEATURES + +# Pass custom `RUSTFLAGS` (e.g., `--cfg tokio_unstable` to enable Tokio tracing/`tokio-console`) +ARG RUSTFLAGS + +# Select linker (e.g., `mold`, `lld` or an empty string for the default linker) +ARG LINKER=mold + +COPY --from=planner /qdrant/recipe.json recipe.json +# `PKG_CONFIG=...` is a workaround for `xx-cargo` bug for crates using `pkg-config`! +# +# https://github.com/tonistiigi/xx/issues/107 +# https://github.com/tonistiigi/xx/pull/108 +RUN PKG_CONFIG="/usr/bin/$(xx-info)-pkg-config" \ + PATH="$PATH:/opt/mold/bin" \ + RUSTFLAGS="${LINKER:+-C link-arg=-fuse-ld=}$LINKER $RUSTFLAGS" \ + xx-cargo chef cook --profile $PROFILE ${FEATURES:+--features} $FEATURES --features=stacktrace --recipe-path recipe.json + +COPY . . +# Include git commit into Qdrant binary during build +ARG GIT_COMMIT_ID +# `PKG_CONFIG=...` is a workaround for `xx-cargo` bug for crates using `pkg-config`! +# +# https://github.com/tonistiigi/xx/issues/107 +# https://github.com/tonistiigi/xx/pull/108 +RUN PKG_CONFIG="/usr/bin/$(xx-info)-pkg-config" \ + PATH="$PATH:/opt/mold/bin" \ + RUSTFLAGS="${LINKER:+-C link-arg=-fuse-ld=}$LINKER $RUSTFLAGS" \ + xx-cargo build --profile $PROFILE ${FEATURES:+--features} $FEATURES --features=stacktrace --bin qdrant \ + && PROFILE_DIR=$(if [ "$PROFILE" = dev ]; then echo debug; else echo $PROFILE; fi) \ + && mv target/$(xx-cargo --print-target-triple)/$PROFILE_DIR/qdrant /qdrant/qdrant + +# Download and extract web UI +RUN mkdir /static && STATIC_DIR=/static ./tools/sync-web-ui.sh + + +FROM rocm/dev-ubuntu-22.04 + +ARG DEBIAN_FRONTEND=noninteractive + +RUN apt-get update && apt-get install -y --no-install-recommends \ + libglvnd0 \ + libgl1 \ + libglx0 \ + libegl1 \ + libgles2 \ + libxcb1-dev \ + wget \ + && rm -rf /var/lib/apt/lists/* + +RUN apt-get update \ + && apt-get install -y \ + libxext6 \ + libvulkan1 \ + libvulkan-dev \ + vulkan-tools + +# Install additional packages into the container. +# E.g., the debugger of choice: gdb/gdbserver/lldb. +ARG PACKAGES + +RUN apt-get install -y --no-install-recommends ca-certificates tzdata libunwind8 $PACKAGES \ + && rm -rf /var/lib/apt/lists/* + +# Copy Qdrant source files into the container. Useful for debugging. +# +# To enable, set `SOURCES` to *any* non-empty string. E.g., 1/true/enable/whatever. +# (Note, that *any* non-empty string would work, so 0/false/disable would enable the option as well.) +ARG SOURCES + +# Dockerfile does not support conditional `COPY` instructions (e.g., it's impossible to do something +# like `if [ -n "$SOURCES" ]; then COPY ...; fi`), so we *hack* conditional `COPY` by abusing +# parameter expansion and `COPY` wildcards support. 😎 + +ENV DIR=${SOURCES:+/qdrant/src} +COPY --from=builder ${DIR:-/null?} $DIR/ + +ENV DIR=${SOURCES:+/qdrant/lib} +COPY --from=builder ${DIR:-/null?} $DIR/ + +ENV DIR=${SOURCES:+/usr/local/cargo/registry/src} +COPY --from=builder ${DIR:-/null?} $DIR/ + +ENV DIR=${SOURCES:+/usr/local/cargo/git/checkouts} +COPY --from=builder ${DIR:-/null?} $DIR/ + +ENV DIR= + +ARG APP=/qdrant + +COPY --from=builder /qdrant/qdrant "$APP"/qdrant +COPY --from=builder /qdrant/config "$APP"/config +COPY --from=builder /qdrant/tools/entrypoint.sh "$APP"/entrypoint.sh +COPY --from=builder /static "$APP"/static + +WORKDIR "$APP" + +ARG USER_ID=0 + +RUN if [ "$USER_ID" != 0 ]; then \ + groupadd --gid "$USER_ID" qdrant; \ + useradd --uid "$USER_ID" --gid "$USER_ID" -m qdrant; \ + mkdir -p "$APP"/storage "$APP"/snapshots; \ + chown -R "$USER_ID:$USER_ID" "$APP"; \ + fi + +USER "$USER_ID:$USER_ID" + +ENV TZ=Etc/UTC \ + RUN_MODE=production + +EXPOSE 6333 +EXPOSE 6334 + +LABEL org.opencontainers.image.title="Qdrant" +LABEL org.opencontainers.image.description="Official Qdrant image" +LABEL org.opencontainers.image.url="https://qdrant.com/" +LABEL org.opencontainers.image.documentation="https://qdrant.com/docs" +LABEL org.opencontainers.image.source="https://github.com/qdrant/qdrant" +LABEL org.opencontainers.image.vendor="Qdrant" + +CMD ["./entrypoint.sh"] diff --git a/DockerfileGpuNvidia b/DockerfileGpuNvidia new file mode 100644 index 0000000000..0a3f91591f --- /dev/null +++ b/DockerfileGpuNvidia @@ -0,0 +1,195 @@ +# Cross-compiling using Docker multi-platform builds/images and `xx`. +# +# https://docs.docker.com/build/building/multi-platform/ +# https://github.com/tonistiigi/xx +FROM --platform=${BUILDPLATFORM:-linux/amd64} tonistiigi/xx AS xx + +# Utilizing Docker layer caching with `cargo-chef`. +# +# https://www.lpalmieri.com/posts/fast-rust-docker-builds/ +FROM --platform=${BUILDPLATFORM:-linux/amd64} lukemathwalker/cargo-chef:latest-rust-1.82.0 AS chef + + +FROM chef AS planner +WORKDIR /qdrant +COPY . . +RUN cargo chef prepare --recipe-path recipe.json + + +FROM chef AS builder +WORKDIR /qdrant + +COPY --from=xx / / + +# Relative order of `ARG` and `RUN` commands in the Dockerfile matters. +# +# If you pass a different `ARG` to `docker build`, it would invalidate Docker layer cache +# for the next steps. (E.g., the following steps may depend on a new `ARG` value, so Docker would +# have to re-execute them instead of using a cached layer from a previous run.) +# +# Steps in this stage are ordered in a way that should maximize Docker layer cache utilization, +# so, please, don't reorder them without prior consideration. 🥲 + +RUN apt-get update \ + && apt-get install -y clang lld cmake protobuf-compiler jq \ + && rustup component add rustfmt + +# `ARG`/`ENV` pair is a workaround for `docker build` backward-compatibility. +# +# https://github.com/docker/buildx/issues/510 +ARG BUILDPLATFORM +ENV BUILDPLATFORM=${BUILDPLATFORM:-linux/amd64} + +ARG MOLD_VERSION=2.31.0 + +RUN case "$BUILDPLATFORM" in \ + */amd64 ) PLATFORM=x86_64 ;; \ + */arm64 | */arm64/* ) PLATFORM=aarch64 ;; \ + * ) echo "Unexpected BUILDPLATFORM '$BUILDPLATFORM'" >&2; exit 1 ;; \ + esac; \ + \ + mkdir -p /opt/mold; \ + cd /opt/mold; \ + \ + TARBALL="mold-$MOLD_VERSION-$PLATFORM-linux.tar.gz"; \ + curl -sSLO "https://github.com/rui314/mold/releases/download/v$MOLD_VERSION/$TARBALL"; \ + tar -xf "$TARBALL" --strip-components 1; \ + rm "$TARBALL" + +# `ARG`/`ENV` pair is a workaround for `docker build` backward-compatibility. +# +# https://github.com/docker/buildx/issues/510 +ARG TARGETPLATFORM +ENV TARGETPLATFORM=${TARGETPLATFORM:-linux/amd64} + +RUN xx-apt-get install -y pkg-config gcc g++ libc6-dev libunwind-dev + +# Select Cargo profile (e.g., `release`, `dev` or `ci`) +ARG PROFILE=release + +# Enable crate features +ARG FEATURES + +# Pass custom `RUSTFLAGS` (e.g., `--cfg tokio_unstable` to enable Tokio tracing/`tokio-console`) +ARG RUSTFLAGS + +# Select linker (e.g., `mold`, `lld` or an empty string for the default linker) +ARG LINKER=mold + +COPY --from=planner /qdrant/recipe.json recipe.json +# `PKG_CONFIG=...` is a workaround for `xx-cargo` bug for crates using `pkg-config`! +# +# https://github.com/tonistiigi/xx/issues/107 +# https://github.com/tonistiigi/xx/pull/108 +RUN PKG_CONFIG="/usr/bin/$(xx-info)-pkg-config" \ + PATH="$PATH:/opt/mold/bin" \ + RUSTFLAGS="${LINKER:+-C link-arg=-fuse-ld=}$LINKER $RUSTFLAGS" \ + xx-cargo chef cook --profile $PROFILE ${FEATURES:+--features} $FEATURES --features=stacktrace --recipe-path recipe.json + +COPY . . +# Include git commit into Qdrant binary during build +ARG GIT_COMMIT_ID +# `PKG_CONFIG=...` is a workaround for `xx-cargo` bug for crates using `pkg-config`! +# +# https://github.com/tonistiigi/xx/issues/107 +# https://github.com/tonistiigi/xx/pull/108 +RUN PKG_CONFIG="/usr/bin/$(xx-info)-pkg-config" \ + PATH="$PATH:/opt/mold/bin" \ + RUSTFLAGS="${LINKER:+-C link-arg=-fuse-ld=}$LINKER $RUSTFLAGS" \ + xx-cargo build --profile $PROFILE ${FEATURES:+--features} $FEATURES --features=stacktrace --bin qdrant \ + && PROFILE_DIR=$(if [ "$PROFILE" = dev ]; then echo debug; else echo $PROFILE; fi) \ + && mv target/$(xx-cargo --print-target-triple)/$PROFILE_DIR/qdrant /qdrant/qdrant + +# Download and extract web UI +RUN mkdir /static && STATIC_DIR=/static ./tools/sync-web-ui.sh + + +FROM nvidia/opengl:1.0-glvnd-devel-ubuntu22.04 + +ARG DEBIAN_FRONTEND=noninteractive + +RUN apt-get update && apt-get install -y --no-install-recommends \ + libglvnd0 \ + libgl1 \ + libglx0 \ + libegl1 \ + libgles2 \ + libxcb1-dev \ + wget \ + && rm -rf /var/lib/apt/lists/* + +RUN apt-get update \ + && apt-get install -y \ + libxext6 \ + libvulkan1 \ + libvulkan-dev \ + vulkan-tools + +COPY nvidia_icd.json /etc/vulkan/icd.d/ +ENV NVIDIA_DRIVER_CAPABILITIES compute,graphics,utility + +# Install additional packages into the container. +# E.g., the debugger of choice: gdb/gdbserver/lldb. +ARG PACKAGES + +RUN apt-get install -y --no-install-recommends ca-certificates tzdata libunwind8 $PACKAGES \ + && rm -rf /var/lib/apt/lists/* + +# Copy Qdrant source files into the container. Useful for debugging. +# +# To enable, set `SOURCES` to *any* non-empty string. E.g., 1/true/enable/whatever. +# (Note, that *any* non-empty string would work, so 0/false/disable would enable the option as well.) +ARG SOURCES + +# Dockerfile does not support conditional `COPY` instructions (e.g., it's impossible to do something +# like `if [ -n "$SOURCES" ]; then COPY ...; fi`), so we *hack* conditional `COPY` by abusing +# parameter expansion and `COPY` wildcards support. 😎 + +ENV DIR=${SOURCES:+/qdrant/src} +COPY --from=builder ${DIR:-/null?} $DIR/ + +ENV DIR=${SOURCES:+/qdrant/lib} +COPY --from=builder ${DIR:-/null?} $DIR/ + +ENV DIR=${SOURCES:+/usr/local/cargo/registry/src} +COPY --from=builder ${DIR:-/null?} $DIR/ + +ENV DIR=${SOURCES:+/usr/local/cargo/git/checkouts} +COPY --from=builder ${DIR:-/null?} $DIR/ + +ENV DIR= + +ARG APP=/qdrant + +COPY --from=builder /qdrant/qdrant "$APP"/qdrant +COPY --from=builder /qdrant/config "$APP"/config +COPY --from=builder /qdrant/tools/entrypoint.sh "$APP"/entrypoint.sh +COPY --from=builder /static "$APP"/static + +WORKDIR "$APP" + +ARG USER_ID=0 + +RUN if [ "$USER_ID" != 0 ]; then \ + groupadd --gid "$USER_ID" qdrant; \ + useradd --uid "$USER_ID" --gid "$USER_ID" -m qdrant; \ + mkdir -p "$APP"/storage "$APP"/snapshots; \ + chown -R "$USER_ID:$USER_ID" "$APP"; \ + fi + +USER "$USER_ID:$USER_ID" + +ENV TZ=Etc/UTC \ + RUN_MODE=production + +EXPOSE 6333 +EXPOSE 6334 + +LABEL org.opencontainers.image.title="Qdrant" +LABEL org.opencontainers.image.description="Official Qdrant image" +LABEL org.opencontainers.image.url="https://qdrant.com/" +LABEL org.opencontainers.image.documentation="https://qdrant.com/docs" +LABEL org.opencontainers.image.source="https://github.com/qdrant/qdrant" +LABEL org.opencontainers.image.vendor="Qdrant" + +CMD ["./entrypoint.sh"] diff --git a/config/config.yaml b/config/config.yaml index 4eef7e78b4..bea2b7b184 100644 --- a/config/config.yaml +++ b/config/config.yaml @@ -347,3 +347,6 @@ tls: # Only works for HTTPS endpoints. Does not support gRPC (and intra-cluster communication). # If `null` - TTL is disabled. cert_ttl: 3600 + +gpu: + indexing: false diff --git a/lib/common/common/src/cpu.rs b/lib/common/common/src/cpu.rs index 4d2cf9bb06..bcda34179d 100644 --- a/lib/common/common/src/cpu.rs +++ b/lib/common/common/src/cpu.rs @@ -194,6 +194,19 @@ impl CpuPermit { pub fn release(&mut self) { self.permit.take(); } + + /// Partial release CPU permit, giving them back to the semaphore. + pub fn release_count(&mut self, release_count: u32) { + if self.num_cpus > release_count { + self.num_cpus -= release_count; + let permit = self.permit.take(); + self.permit = permit + .map(|mut permit| permit.split(self.num_cpus as usize)) + .flatten(); + } else { + self.release(); + } + } } impl Drop for CpuPermit { diff --git a/lib/segment/Cargo.toml b/lib/segment/Cargo.toml index 1f1d68e1ba..66a569613d 100644 --- a/lib/segment/Cargo.toml +++ b/lib/segment/Cargo.toml @@ -12,6 +12,7 @@ edition = "2021" workspace = true [features] +default = ["gpu"] multiling-chinese = ["charabia/chinese-segmentation", "charabia/chinese-normalization", "charabia/chinese-normalization-pinyin"] multiling-japanese = ["charabia/japanese"] multiling-korean = ["charabia/korean"] @@ -85,6 +86,8 @@ is_sorted = "0.1.1" strum = { workspace = true } byteorder = { workspace = true } zerocopy = { workspace = true } +lazy_static = "1.5.0" +env_logger = "0.11" sysinfo = "0.32" charabia = { version = "0.9.1", default-features = false, features = ["greek", "hebrew", "thai"] } diff --git a/lib/segment/benches/multi_vector_search.rs b/lib/segment/benches/multi_vector_search.rs index 5a5a5af1fa..5217ccde18 100644 --- a/lib/segment/benches/multi_vector_search.rs +++ b/lib/segment/benches/multi_vector_search.rs @@ -85,6 +85,7 @@ fn multi_vector_search_benchmark(c: &mut Criterion) { payload_index: segment.payload_index.clone(), hnsw_config, permit: Some(permit), + gpu_device: None, stopped: &stopped, }) .unwrap(); diff --git a/lib/segment/src/index/hnsw_index/gpu/batched_points.rs b/lib/segment/src/index/hnsw_index/gpu/batched_points.rs new file mode 100644 index 0000000000..d26cf22ecd --- /dev/null +++ b/lib/segment/src/index/hnsw_index/gpu/batched_points.rs @@ -0,0 +1,159 @@ +use std::ops::Range; +use std::sync::atomic::AtomicU32; + +use ahash::HashSet; +use common::types::PointOffsetType; + +use crate::common::operation_error::OperationResult; + +#[derive(Debug)] +pub struct PointLinkingData { + pub point_id: PointOffsetType, + pub level: usize, + pub batch_index: usize, + pub entry: AtomicU32, +} + +#[derive(Debug)] +pub struct Batch<'a> { + pub points: &'a [PointLinkingData], + pub level: usize, + pub end_index: usize, +} + +pub struct BatchedPoints { + pub points: Vec, + pub batches: Vec>, + pub ids_by_batches: Vec>, + pub first_point_id: PointOffsetType, + pub levels_count: usize, +} + +impl BatchedPoints { + pub fn new( + level_fn: impl Fn(PointOffsetType) -> usize, + mut ids: Vec, + groups_count: usize, + ) -> OperationResult { + Self::sort_points_by_level(&level_fn, &mut ids); + let first_point_id = ids.remove(0); + + let batches = Self::build_initial_batches(&level_fn, &ids, groups_count); + + let mut points = Vec::with_capacity(ids.len()); + for (batch_index, batch) in batches.iter().enumerate() { + for i in batch.clone() { + let point_id = ids[i]; + let level = level_fn(point_id); + points.push(PointLinkingData { + point_id, + level, + batch_index, + entry: first_point_id.into(), + }); + } + } + + let ids_by_batches = batches + .iter() + .map(|batch| { + batch + .clone() + .map(|i| ids[i]) + .collect::>() + }) + .collect(); + + Ok(Self { + points, + batches, + ids_by_batches, + first_point_id, + levels_count: level_fn(first_point_id) + 1, + }) + } + + pub fn is_same_batch( + &self, + linking_point: &PointLinkingData, + other_point: PointOffsetType, + ) -> bool { + self.ids_by_batches[linking_point.batch_index].contains(&other_point) + } + + pub fn iter_batches(&self, skip_count: usize) -> impl Iterator { + self.batches + .iter() + .filter(move |batch| batch.end > skip_count) + .map(move |batch| { + let intersected_batch = std::cmp::max(batch.start, skip_count)..batch.end; + let level = self.points[intersected_batch.start].level; + Batch { + end_index: intersected_batch.end, + points: &self.points[intersected_batch], + level, + } + }) + } + + pub fn sort_points_by_level( + level_fn: impl Fn(PointOffsetType) -> usize, + ids: &mut [PointOffsetType], + ) { + ids.sort_by(|&a, &b| { + let a_level = level_fn(a); + let b_level = level_fn(b); + match b_level.cmp(&a_level) { + std::cmp::Ordering::Less => std::cmp::Ordering::Less, + std::cmp::Ordering::Greater => std::cmp::Ordering::Greater, + std::cmp::Ordering::Equal => a.cmp(&b), + } + }); + } + + fn build_initial_batches( + level_fn: impl Fn(PointOffsetType) -> usize, + ids: &[PointOffsetType], + groups_count: usize, + ) -> Vec> { + let timer = std::time::Instant::now(); + + let num_vectors = ids.len(); + let mut batches: Vec<_> = (0..num_vectors.div_ceil(groups_count)) + .map(|start| { + groups_count * start..std::cmp::min(groups_count * (start + 1), num_vectors) + }) + .collect(); + + let mut batch_index = 0usize; + while batch_index < batches.len() { + let batch = batches[batch_index].clone(); + let point_id = ids[batch.start]; + let batch_level = level_fn(point_id); + for i in 1..batch.len() { + let point_id = ids[batch.start + i]; + let level = level_fn(point_id); + // divide batch by level. all batches must be on the same level + if level != batch_level { + let batch1 = batch.start..batch.start + i; + let batch2 = batch.start + i..batch.end; + batches[batch_index] = batch1; + batches.insert(batch_index + 1, batch2); + break; + } + } + + batch_index += 1; + } + + for batch_pair in batches.windows(2) { + if batch_pair.len() == 2 { + assert_eq!(batch_pair[0].end, batch_pair[1].start); + } + } + + log::trace!("Initial batchs time: {:?}", timer.elapsed()); + + batches + } +} diff --git a/lib/segment/src/index/hnsw_index/gpu/cpu_level_builder.rs b/lib/segment/src/index/hnsw_index/gpu/cpu_level_builder.rs new file mode 100644 index 0000000000..61ef257c0b --- /dev/null +++ b/lib/segment/src/index/hnsw_index/gpu/cpu_level_builder.rs @@ -0,0 +1,260 @@ +use std::sync::atomic::Ordering; +use std::sync::Mutex; + +use common::types::PointOffsetType; +use rayon::iter::{IntoParallelIterator, ParallelIterator}; +use rayon::ThreadPool; + +use super::batched_points::{BatchedPoints, PointLinkingData}; +use super::gpu_search_context::GpuRequest; +use crate::common::operation_error::OperationResult; +use crate::index::hnsw_index::graph_layers::GraphLayersBase; +use crate::index::hnsw_index::graph_layers_builder::GraphLayersBuilder; +use crate::index::hnsw_index::point_scorer::FilteredScorer; +use crate::payload_storage::FilterContext; +use crate::vector_storage::RawScorer; + +// Build level on CPU. Returns amount of processed points. +pub fn build_level_on_cpu<'a>( + pool: &ThreadPool, + graph_layers_builder: &GraphLayersBuilder, + batched_points: &BatchedPoints, + level: usize, + stop_condition: impl Fn(usize) -> bool + Send + Sync, + points_scorer_builder: impl Fn( + PointOffsetType, + ) + -> OperationResult<(Box, Option>)> + + Send + + Sync, +) -> OperationResult { + let retry_mutex: Mutex<()> = Default::default(); + let index = Mutex::new(0usize); + + pool.install(|| { + (0..graph_layers_builder.links_layers.len()) + .into_par_iter() + .try_for_each(|_| -> OperationResult<()> { + let index = { + let mut locked_index = index.lock().unwrap(); + let index = *locked_index; + if index >= batched_points.points.len() { + return Ok(()); + } + + if stop_condition(index) { + return Ok(()); + } + *locked_index += 1; + index + }; + + let linking_point = &batched_points.points[index]; + if graph_layers_builder.get_point_level(linking_point.point_id) < level { + update_entry_on_cpu( + graph_layers_builder, + linking_point, + level, + &points_scorer_builder, + ) + } else { + link_point_on_cpu( + &retry_mutex, + graph_layers_builder, + batched_points, + linking_point, + level, + &points_scorer_builder, + ) + } + }) + })?; + + let processed_count = index.lock().unwrap().min(batched_points.points.len()); + Ok(processed_count) +} + +fn update_entry_on_cpu<'a>( + graph_layers_builder: &GraphLayersBuilder, + linking_point: &PointLinkingData, + level: usize, + points_scorer_builder: impl Fn( + PointOffsetType, + ) + -> OperationResult<(Box, Option>)> + + Send + + Sync, +) -> OperationResult<()> { + let (raw_scorer, filter_context) = points_scorer_builder(linking_point.point_id)?; + let mut points_scorer = FilteredScorer::new(raw_scorer.as_ref(), filter_context.as_deref()); + let new_entry = graph_layers_builder + .search_entry_on_level( + linking_point.entry.load(Ordering::Relaxed), + level, + &mut points_scorer, + ) + .idx; + linking_point.entry.store(new_entry, Ordering::Relaxed); + raw_scorer.take_hardware_counter().discard_results(); + Ok(()) +} + +fn link_point_on_cpu<'a>( + retry_mutex: &Mutex<()>, + graph_layers_builder: &GraphLayersBuilder, + batched_points: &BatchedPoints, + linking_point: &PointLinkingData, + level: usize, + points_scorer_builder: impl Fn( + PointOffsetType, + ) + -> OperationResult<(Box, Option>)> + + Send + + Sync, +) -> OperationResult<()> { + let (raw_scorer, filter_context) = points_scorer_builder(linking_point.point_id)?; + let mut points_scorer = FilteredScorer::new(raw_scorer.as_ref(), filter_context.as_deref()); + let (patches, new_entries) = graph_layers_builder.get_patch( + GpuRequest { + id: linking_point.point_id, + entry: linking_point.entry.load(Ordering::Relaxed), + }, + level, + &mut points_scorer, + ); + + let new_entries = if !graph_layers_builder.try_apply_patch(level, patches) { + let _retry_guard = retry_mutex.lock().unwrap(); + + // try again in single-thread mode + let (patches, new_entries) = graph_layers_builder.get_patch( + GpuRequest { + id: linking_point.point_id, + entry: linking_point.entry.load(Ordering::Relaxed), + }, + level, + &mut points_scorer, + ); + graph_layers_builder.apply_patch(level, patches); + new_entries + } else { + new_entries + }; + + raw_scorer.take_hardware_counter().discard_results(); + + new_entries + .into_iter() + .filter(|&new_entry| !batched_points.is_same_batch(linking_point, new_entry)) + .next() + .map(|new_entry| { + linking_point.entry.store(new_entry, Ordering::Relaxed); + }); + + Ok(()) +} + +#[cfg(test)] +mod tests { + use common::types::PointOffsetType; + + use super::*; + use crate::fixtures::index_fixtures::FakeFilterContext; + use crate::index::hnsw_index::gpu::batched_points::BatchedPoints; + use crate::index::hnsw_index::gpu::create_graph_layers_builder; + use crate::index::hnsw_index::gpu::tests::{ + check_graph_layers_builders_quality, compare_graph_layers_builders, + create_gpu_graph_test_data, GpuGraphTestData, + }; + use crate::vector_storage::chunked_vector_storage::VectorOffsetType; + + fn build_cpu_graph(test: &GpuGraphTestData, threads: usize) -> GraphLayersBuilder { + let num_vectors = test.graph_layers_builder.links_layers.len(); + let m = test.graph_layers_builder.m; + let m0 = test.graph_layers_builder.m0; + let ef = test.graph_layers_builder.ef_construct; + + let batched_points = BatchedPoints::new( + |point_id| test.graph_layers_builder.get_point_level(point_id), + (0..num_vectors as PointOffsetType).collect(), + threads, + ) + .unwrap(); + + let graph_layers_builder = + create_graph_layers_builder(&batched_points, num_vectors, m, m0, ef, 1).unwrap(); + + let pool = rayon::ThreadPoolBuilder::new() + .thread_name(|idx| format!("hnsw-build-{idx}")) + .num_threads(threads) + .build() + .unwrap(); + + for level in (0..batched_points.levels_count).rev() { + build_level_on_cpu( + &pool, + &graph_layers_builder, + &batched_points, + level, + |_| false, + |point_id| { + let fake_filter_context = FakeFilterContext {}; + let added_vector = test + .vector_holder + .vectors + .get(point_id as VectorOffsetType) + .to_vec(); + let raw_scorer = test + .vector_holder + .get_raw_scorer(added_vector.clone()) + .unwrap(); + Ok((raw_scorer, Some(Box::new(fake_filter_context)))) + }, + ) + .unwrap(); + } + + graph_layers_builder + } + + #[test] + fn test_cpu_hnsw_level_equivalency() { + let _ = env_logger::builder() + .is_test(true) + .filter_level(log::LevelFilter::Trace) + .try_init(); + + let num_vectors = 1024; + let dim = 64; + let m = 8; + let m0 = 16; + let ef = 32; + + let test = create_gpu_graph_test_data(num_vectors, dim, m, m0, ef, 0); + let graph_layers_builder = build_cpu_graph(&test, 1); + + compare_graph_layers_builders(&test.graph_layers_builder, &graph_layers_builder); + } + + #[test] + fn test_cpu_hnsw_level_quality() { + let _ = env_logger::builder() + .is_test(true) + .filter_level(log::LevelFilter::Trace) + .try_init(); + + let num_vectors = 1024; + let dim = 64; + let m = 8; + let m0 = 16; + let ef = 32; + let threads = 4; + let searches_count = 10; + let top = 10; + + let test = create_gpu_graph_test_data(num_vectors, dim, m, m0, ef, searches_count); + let graph_layers_builder = build_cpu_graph(&test, threads); + + check_graph_layers_builders_quality(graph_layers_builder, test, top, ef, 0.9) + } +} diff --git a/lib/segment/src/index/hnsw_index/gpu/devices_manager.rs b/lib/segment/src/index/hnsw_index/gpu/devices_manager.rs new file mode 100644 index 0000000000..6393639596 --- /dev/null +++ b/lib/segment/src/index/hnsw_index/gpu/devices_manager.rs @@ -0,0 +1,84 @@ +use std::sync::Arc; + +use parking_lot::{Mutex, MutexGuard}; + +use crate::common::operation_error::OperationResult; + +pub struct DevicesMaganer { + pub devices: Vec>>, + pub wait_free: bool, +} + +pub struct LockedDevice<'a> { + pub locked_device: MutexGuard<'a, Arc>, +} + +impl DevicesMaganer { + pub fn new( + instance: Arc, + filter: &str, + start_index: usize, + count: usize, + wait_free: bool, + parallel_indexes: usize, + ) -> OperationResult { + let filter = filter.to_lowercase(); + let mut devices = Vec::new(); + for queue_index in 0..parallel_indexes { + devices.extend( + instance + .physical_devices() + .iter() + .filter(|device| { + let device_name = device.name.to_lowercase(); + device_name.contains(&filter) + }) + .skip(start_index) + .take(count) + .filter_map(|physical_device| { + match gpu::Device::new_with_queue_index( + instance.clone(), + physical_device, + queue_index, + ) { + Ok(device) => { + log::info!("Initialized GPU device: {:?}", &physical_device.name); + Some(Mutex::new(device)) + } + Err(err) => { + log::error!( + "Failed to create GPU device: {:?}, error: {:?}", + &physical_device.name, + err + ); + None + } + } + }), + ); + } + Ok(Self { devices, wait_free }) + } + + pub fn lock_device(&self) -> Option { + if self.devices.is_empty() { + return None; + } + loop { + // TODO(gpu): Add timeout + for device in &self.devices { + if let Some(guard) = device.try_lock() { + return Some(LockedDevice { + locked_device: guard, + }); + } + } + + if !self.wait_free { + return None; + } + + std::thread::sleep(std::time::Duration::from_millis(500)); + } + } +} diff --git a/lib/segment/src/index/hnsw_index/gpu/gpu_candidates_heap.rs b/lib/segment/src/index/hnsw_index/gpu/gpu_candidates_heap.rs new file mode 100644 index 0000000000..4825391343 --- /dev/null +++ b/lib/segment/src/index/hnsw_index/gpu/gpu_candidates_heap.rs @@ -0,0 +1,217 @@ +use std::collections::HashMap; +use std::sync::Arc; + +use super::shader_builder::ShaderBuilderParameters; + +pub struct GpuCandidatesHeap { + pub capacity: usize, + pub device: Arc, +} + +impl GpuCandidatesHeap { + pub fn new(device: Arc, capacity: usize) -> gpu::GpuResult { + let ceiled_capacity = capacity.div_ceil(device.subgroup_size()) * device.subgroup_size(); + Ok(Self { + capacity: ceiled_capacity, + device, + }) + } +} + +impl ShaderBuilderParameters for GpuCandidatesHeap { + fn shader_includes(&self) -> HashMap { + HashMap::from([( + "candidates_heap.comp".to_string(), + include_str!("shaders/candidates_heap.comp").to_string(), + )]) + } + + fn shader_defines(&self) -> HashMap> { + let mut defines = HashMap::new(); + defines.insert( + "CANDIDATES_HEAP_CAPACITY".to_owned(), + Some(self.capacity.to_string()), + ); + defines + } +} + +#[cfg(test)] +mod tests { + use std::collections::BinaryHeap; + + use common::types::{PointOffsetType, ScoredPointOffset}; + use rand::rngs::StdRng; + use rand::{Rng, SeedableRng}; + + use super::*; + use crate::index::hnsw_index::gpu::shader_builder::ShaderBuilder; + use crate::index::hnsw_index::gpu::GPU_TIMEOUT; + + #[repr(C)] + struct TestParams { + input_counts: u32, + } + + #[test] + fn test_gpu_candidates_heap() { + let _ = env_logger::builder() + .is_test(true) + .filter_level(log::LevelFilter::Trace) + .try_init(); + + let capacity = 128; + let points_count = 128; + let groups_count = 8; + let inputs_count = points_count; + + let mut rng = StdRng::seed_from_u64(42); + let inputs_data: Vec = (0..inputs_count * groups_count) + .map(|i| ScoredPointOffset { + idx: i as PointOffsetType, + score: rng.gen_range(-1.0..1.0), + }) + .collect(); + + let debug_messenger = gpu::PanicIfErrorMessenger {}; + let instance = gpu::Instance::new(Some(&debug_messenger), None, false).unwrap(); + let device = gpu::Device::new(instance.clone(), &instance.physical_devices()[0]).unwrap(); + + let gpu_candidates_heap = GpuCandidatesHeap::new(device.clone(), capacity).unwrap(); + + let shader = ShaderBuilder::new(device.clone()) + .with_shader_code(include_str!("shaders/tests/test_candidates_heap.comp")) + .with_parameters(&gpu_candidates_heap) + .build() + .unwrap(); + + let input_points_buffer = gpu::Buffer::new( + device.clone(), + "Input points buffer", + gpu::BufferType::Storage, + inputs_count * groups_count * std::mem::size_of::(), + ) + .unwrap(); + + let upload_staging_buffer = gpu::Buffer::new( + device.clone(), + "Candidates heap upload staging buffer", + gpu::BufferType::CpuToGpu, + inputs_count * groups_count * std::mem::size_of::(), + ) + .unwrap(); + upload_staging_buffer.upload_slice(&inputs_data, 0).unwrap(); + + let mut context = gpu::Context::new(device.clone()).unwrap(); + context + .copy_gpu_buffer( + upload_staging_buffer.clone(), + input_points_buffer.clone(), + 0, + 0, + input_points_buffer.size(), + ) + .unwrap(); + context.run().unwrap(); + context.wait_finish(GPU_TIMEOUT).unwrap(); + + let test_params_buffer = gpu::Buffer::new( + device.clone(), + "Test params buffer", + gpu::BufferType::Uniform, + std::mem::size_of::(), + ) + .unwrap(); + upload_staging_buffer + .upload( + &TestParams { + input_counts: inputs_count as u32, + }, + 0, + ) + .unwrap(); + context + .copy_gpu_buffer( + upload_staging_buffer, + test_params_buffer.clone(), + 0, + 0, + test_params_buffer.size(), + ) + .unwrap(); + context.run().unwrap(); + context.wait_finish(GPU_TIMEOUT).unwrap(); + + let scores_output_buffer = gpu::Buffer::new( + device.clone(), + "Scores output buffer", + gpu::BufferType::Storage, + inputs_count * groups_count * std::mem::size_of::(), + ) + .unwrap(); + + let descriptor_set_layout = gpu::DescriptorSetLayout::builder() + .add_uniform_buffer(0) + .add_storage_buffer(1) + .add_storage_buffer(2) + .build(device.clone()) + .unwrap(); + + let descriptor_set = gpu::DescriptorSet::builder(descriptor_set_layout.clone()) + .add_uniform_buffer(0, test_params_buffer.clone()) + .add_storage_buffer(1, input_points_buffer.clone()) + .add_storage_buffer(2, scores_output_buffer.clone()) + .build() + .unwrap(); + + let pipeline = gpu::Pipeline::builder() + .add_descriptor_set_layout(0, descriptor_set_layout.clone()) + .add_shader(shader.clone()) + .build(device.clone()) + .unwrap(); + + context + .bind_pipeline(pipeline, &[descriptor_set.clone()]) + .unwrap(); + context.dispatch(groups_count, 1, 1).unwrap(); + context.run().unwrap(); + context.wait_finish(GPU_TIMEOUT).unwrap(); + + let mut scores_cpu = vec![]; + for group in 0..groups_count { + let mut heap = BinaryHeap::::new(); + for i in 0..inputs_count { + let scored_point = inputs_data[group * inputs_count + i]; + heap.push(scored_point); + } + while !heap.is_empty() { + scores_cpu.push(heap.pop().unwrap()); + } + } + + let download_staging_buffer = gpu::Buffer::new( + device.clone(), + "Candidates heap download staging buffer", + gpu::BufferType::GpuToCpu, + scores_output_buffer.size(), + ) + .unwrap(); + context + .copy_gpu_buffer( + scores_output_buffer.clone(), + download_staging_buffer.clone(), + 0, + 0, + scores_output_buffer.size(), + ) + .unwrap(); + context.run().unwrap(); + context.wait_finish(GPU_TIMEOUT).unwrap(); + let mut scores_gpu = vec![ScoredPointOffset::default(); inputs_count * groups_count]; + download_staging_buffer + .download_slice(&mut scores_gpu, 0) + .unwrap(); + + assert_eq!(scores_gpu, scores_cpu); + } +} diff --git a/lib/segment/src/index/hnsw_index/gpu/gpu_graph_builder.rs b/lib/segment/src/index/hnsw_index/gpu/gpu_graph_builder.rs new file mode 100644 index 0000000000..7d70a5f9bf --- /dev/null +++ b/lib/segment/src/index/hnsw_index/gpu/gpu_graph_builder.rs @@ -0,0 +1,304 @@ +use std::sync::atomic::{AtomicBool, AtomicUsize}; +use std::sync::Arc; +use std::thread::JoinHandle; + +use common::types::PointOffsetType; +use parking_lot::Mutex; +use rayon::ThreadPool; + +use crate::common::operation_error::OperationResult; +use crate::index::hnsw_index::gpu::batched_points::BatchedPoints; +use crate::index::hnsw_index::gpu::cpu_level_builder::build_level_on_cpu; +use crate::index::hnsw_index::gpu::create_graph_layers_builder; +use crate::index::hnsw_index::gpu::gpu_level_builder::build_level_on_gpu; +use crate::index::hnsw_index::gpu::gpu_search_context::GpuSearchContext; +use crate::index::hnsw_index::graph_layers_builder::GraphLayersBuilder; +use crate::payload_storage::FilterContext; +use crate::vector_storage::quantized::quantized_vectors::QuantizedVectors; +use crate::vector_storage::{RawScorer, VectorStorageEnum}; + +#[allow(clippy::too_many_arguments)] +pub fn build_hnsw_on_gpu<'a>( + device: Arc, + pool: &ThreadPool, + reference_graph: &GraphLayersBuilder, + max_groups_count: usize, + vector_storage: &VectorStorageEnum, + quantized_storage: Option<&QuantizedVectors>, + entry_points_num: usize, + force_half_precision: bool, + min_cpu_linked_points_count: usize, + exact: bool, + ids: Vec, + points_scorer_builder: impl Fn( + PointOffsetType, + ) + -> OperationResult<(Box, Option>)> + + Send + + Sync, + stopped: &AtomicBool, +) -> OperationResult { + let num_vectors = reference_graph.links_layers.len(); + let m = reference_graph.m; + let m0 = reference_graph.m0; + let ef = reference_graph.ef_construct; + + let gpu_search_context = Arc::new(Mutex::new(GpuSearchContext::new( + device, + max_groups_count, + vector_storage, + quantized_storage, + m, + m0, + ef, + num_vectors, + force_half_precision, + exact, + stopped, + )?)); + + let batched_points = Arc::new(BatchedPoints::new( + |point_id| reference_graph.get_point_level(point_id), + ids, + gpu_search_context.lock().groups_count, + )?); + + let graph_layers_builder = Arc::new(create_graph_layers_builder( + &batched_points, + num_vectors, + m, + m0, + ef, + entry_points_num, + )?); + + let mut gpu_thread_handle: Option>> = None; + let gpu_built_points = Arc::new(AtomicUsize::new(batched_points.points.len())); + + let cpu_stop_condition = |cpu_built_points| { + assert_ne!(cpu_built_points, batched_points.points.len()); + + // if CPU is faster than GPU, wait for GPU to catch up + while cpu_built_points >= gpu_built_points.load(std::sync::atomic::Ordering::Relaxed) { + log::trace!( + "Waiting for GPU to catch up (CPU ready {}, GPU ready {})", + cpu_built_points, + gpu_built_points.load(std::sync::atomic::Ordering::Relaxed) + ); + std::thread::sleep(std::time::Duration::from_millis(100)); + } + + let is_gpu_ready = gpu_built_points.load(std::sync::atomic::Ordering::Relaxed) + == batched_points.points.len(); + let min_required_points_reached = cpu_built_points >= min_cpu_linked_points_count; + + // stop id GPU is ready and min required points are linked + min_required_points_reached && is_gpu_ready + }; + + for level in (0..batched_points.levels_count).rev() { + // Start building level on CPU + let gpu_start_index = build_level_on_cpu( + &pool, + &graph_layers_builder, + &batched_points, + level, + &cpu_stop_condition, + &points_scorer_builder, + )?; + + // CPU is ready, finish prev level GPU thread to continue with this level + if let Some(handle) = gpu_thread_handle.take() { + handle.join().unwrap()?; + } + + // Clear GPU built points counter + gpu_built_points.store(0, std::sync::atomic::Ordering::Relaxed); + + log::debug!( + "Starting GPU level {}, skipped points {}", + level, + gpu_start_index + ); + + let gpu_search_context = gpu_search_context.clone(); + let graph_layers_builder = graph_layers_builder.clone(); + let batched_points = batched_points.clone(); + let gpu_built_points = gpu_built_points.clone(); + // Continue level building on GPU in separate thread + gpu_thread_handle = Some(std::thread::spawn(move || -> OperationResult<()> { + let mut gpu_search_context = gpu_search_context.lock(); + gpu_search_context.upload_links(level, &graph_layers_builder)?; + build_level_on_gpu( + &mut gpu_search_context, + &batched_points, + gpu_start_index, + level, + |count| { + gpu_built_points.store(count, std::sync::atomic::Ordering::Relaxed); + }, + )?; + gpu_search_context.download_links(level, &graph_layers_builder)?; + Ok(()) + })); + } + + // Wait for the last GPU level + if let Some(handle) = gpu_thread_handle.take() { + handle.join().unwrap()?; + } + + { + let gpu_search_context = gpu_search_context.lock(); + log::debug!( + "Gpu graph patches time: {:?}, count {:?}, avg {:?}", + &gpu_search_context.patches_timer, + gpu_search_context.patches_count, + gpu_search_context + .patches_timer + .checked_div(gpu_search_context.patches_count as u32) + .unwrap_or_default(), + ); + log::debug!( + "Gpu graph update entries time: {:?}, count {:?}, avg {:?}", + &gpu_search_context.updates_timer, + gpu_search_context.updates_count, + gpu_search_context + .updates_timer + .checked_div(gpu_search_context.updates_count as u32) + .unwrap_or_default(), + ); + } + + Ok(Arc::into_inner(graph_layers_builder).unwrap()) +} + +#[cfg(test)] +mod tests { + use std::borrow::Borrow; + + use super::*; + use crate::fixtures::index_fixtures::FakeFilterContext; + use crate::index::hnsw_index::gpu::tests::{ + check_graph_layers_builders_quality, compare_graph_layers_builders, + create_gpu_graph_test_data, GpuGraphTestData, + }; + use crate::vector_storage::chunked_vector_storage::VectorOffsetType; + + fn build_gpu_graph( + test: &GpuGraphTestData, + groups_count: usize, + min_cpu_linked_points_count: usize, + exact: bool, + ) -> GraphLayersBuilder { + let num_vectors = test.graph_layers_builder.links_layers.len(); + + let pool = rayon::ThreadPoolBuilder::new() + .thread_name(|idx| format!("hnsw-build-{idx}")) + .num_threads(groups_count) + .build() + .unwrap(); + + let debug_messenger = gpu::PanicIfErrorMessenger {}; + let instance = gpu::Instance::new(Some(&debug_messenger), None, false).unwrap(); + let device = gpu::Device::new(instance.clone(), &instance.physical_devices()[0]).unwrap(); + + let ids = (0..num_vectors as PointOffsetType).collect(); + build_hnsw_on_gpu( + device, + &pool, + &test.graph_layers_builder, + groups_count, + &test.vector_storage.borrow(), + None, + 1, + false, + min_cpu_linked_points_count, + exact, + ids, + |point_id| { + let fake_filter_context = FakeFilterContext {}; + let added_vector = test + .vector_holder + .vectors + .get(point_id as VectorOffsetType) + .to_vec(); + let raw_scorer = test + .vector_holder + .get_raw_scorer(added_vector.clone()) + .unwrap(); + Ok((raw_scorer, Some(Box::new(fake_filter_context)))) + }, + &false.into(), + ) + .unwrap() + } + + #[test] + fn test_gpu_hnsw_equivalency() { + let _ = env_logger::builder() + .is_test(true) + .filter_level(log::LevelFilter::Trace) + .try_init(); + + let num_vectors = 1024; + let dim = 64; + let m = 8; + let m0 = 16; + let ef = 32; + let min_cpu_linked_points_count = 64; + + let test = create_gpu_graph_test_data(num_vectors, dim, m, m0, ef, 0); + let graph_layers_builder = build_gpu_graph(&test, 1, min_cpu_linked_points_count, true); + + compare_graph_layers_builders(&test.graph_layers_builder, &graph_layers_builder); + } + + #[test] + fn test_gpu_hnsw_quality_exact() { + let _ = env_logger::builder() + .is_test(true) + .filter_level(log::LevelFilter::Trace) + .try_init(); + + let num_vectors = 1024; + let dim = 64; + let m = 8; + let m0 = 16; + let ef = 32; + let groups_count = 4; + let searches_count = 20; + let top = 10; + let min_cpu_linked_points_count = 64; + + let test = create_gpu_graph_test_data(num_vectors, dim, m, m0, ef, searches_count); + let graph_layers_builder = + build_gpu_graph(&test, groups_count, min_cpu_linked_points_count, true); + + check_graph_layers_builders_quality(graph_layers_builder, test, top, ef, 0.8) + } + + #[test] + fn test_gpu_hnsw_quality() { + let _ = env_logger::builder() + .is_test(true) + .filter_level(log::LevelFilter::Trace) + .try_init(); + + let num_vectors = 1024; + let dim = 64; + let m = 8; + let m0 = 16; + let ef = 32; + let groups_count = 4; + let searches_count = 20; + let top = 10; + let min_cpu_linked_points_count = 64; + + let test = create_gpu_graph_test_data(num_vectors, dim, m, m0, ef, searches_count); + let graph_layers_builder = + build_gpu_graph(&test, groups_count, min_cpu_linked_points_count, false); + + check_graph_layers_builders_quality(graph_layers_builder, test, top, ef, 0.8) + } +} diff --git a/lib/segment/src/index/hnsw_index/gpu/gpu_level_builder.rs b/lib/segment/src/index/hnsw_index/gpu/gpu_level_builder.rs new file mode 100644 index 0000000000..7a8fabb7ea --- /dev/null +++ b/lib/segment/src/index/hnsw_index/gpu/gpu_level_builder.rs @@ -0,0 +1,211 @@ +use std::sync::atomic::Ordering; + +use common::types::PointOffsetType; + +use super::batched_points::{Batch, BatchedPoints}; +use super::gpu_search_context::GpuSearchContext; +use crate::common::operation_error::OperationResult; +use crate::index::hnsw_index::gpu::gpu_search_context::GpuRequest; + +// Build level on GPU +pub fn build_level_on_gpu( + gpu_search_context: &mut GpuSearchContext, + batched_points: &BatchedPoints, + skip_count: usize, + level: usize, + processed_points_callback: impl Fn(usize), +) -> OperationResult<()> { + let mut prev_batch = None; + + // skip_count points are already processed + processed_points_callback(skip_count); + + for batch in batched_points.iter_batches(skip_count) { + if level > batch.level { + gpu_batched_update_entries(gpu_search_context, &batch, prev_batch.as_ref())?; + } else { + gpu_batched_insert(gpu_search_context, &batch, prev_batch.as_ref())?; + } + + // Prev batch entries are updated, mark them as processed + if let Some(prev_batch) = prev_batch { + processed_points_callback(prev_batch.end_index); + } + + prev_batch = Some(batch); + } + + if let Some(prev_batch) = prev_batch { + let new_entries = gpu_search_context.download_responses(prev_batch.points.len())?; + gpu_batched_apply_entries(&prev_batch, new_entries); + processed_points_callback(prev_batch.end_index); + } else { + processed_points_callback(0); + } + + Ok(()) +} + +fn gpu_batched_update_entries( + gpu_search_context: &mut GpuSearchContext, + batch: &Batch, + prev_batch: Option<&Batch>, +) -> OperationResult<()> { + let mut requests = Vec::with_capacity(batch.points.len()); + for linking_point in batch.points { + requests.push(GpuRequest { + id: linking_point.point_id, + entry: linking_point.entry.load(Ordering::Relaxed), + }) + } + + let prev_batch_len = prev_batch + .map(|prev_batch| prev_batch.points.len()) + .unwrap_or(0); + + let new_entries = gpu_search_context.greedy_search(&requests, prev_batch_len)?; + + if let Some(prev_batch) = prev_batch { + gpu_batched_apply_entries(prev_batch, new_entries); + } + Ok(()) +} + +fn gpu_batched_insert( + gpu_search_context: &mut GpuSearchContext, + batch: &Batch, + prev_batch: Option<&Batch>, +) -> OperationResult<()> { + let mut requests = Vec::with_capacity(batch.points.len()); + for linking_point in batch.points { + requests.push(GpuRequest { + id: linking_point.point_id, + entry: linking_point.entry.load(Ordering::Relaxed), + }) + } + + let prev_batch_len = prev_batch + .map(|prev_batch| prev_batch.points.len()) + .unwrap_or(0); + + let new_entries = gpu_search_context.run_insert_vector(&requests, prev_batch_len)?; + + if let Some(prev_batch) = prev_batch { + gpu_batched_apply_entries(prev_batch, new_entries); + } + Ok(()) +} + +fn gpu_batched_apply_entries(batch: &Batch, new_entries: Vec) { + assert_eq!(batch.points.len(), new_entries.len()); + for (linking_point, new_entry) in batch.points.iter().zip(new_entries) { + linking_point.entry.store(new_entry, Ordering::Relaxed); + } +} + +#[cfg(test)] +mod tests { + use std::borrow::Borrow; + + use common::types::PointOffsetType; + + use super::*; + use crate::index::hnsw_index::gpu::batched_points::BatchedPoints; + use crate::index::hnsw_index::gpu::create_graph_layers_builder; + use crate::index::hnsw_index::gpu::tests::{ + check_graph_layers_builders_quality, compare_graph_layers_builders, + create_gpu_graph_test_data, GpuGraphTestData, + }; + use crate::index::hnsw_index::graph_layers::GraphLayersBase; + use crate::index::hnsw_index::graph_layers_builder::GraphLayersBuilder; + + fn build_gpu_graph(test: &GpuGraphTestData, groups_count: usize) -> GraphLayersBuilder { + let num_vectors = test.graph_layers_builder.links_layers.len(); + let m = test.graph_layers_builder.m; + let m0 = test.graph_layers_builder.m0; + let ef = test.graph_layers_builder.ef_construct; + + let batched_points = BatchedPoints::new( + |point_id| test.graph_layers_builder.get_point_level(point_id), + (0..num_vectors as PointOffsetType).collect(), + groups_count, + ) + .unwrap(); + + let graph_layers_builder = + create_graph_layers_builder(&batched_points, num_vectors, m, m0, ef, 1).unwrap(); + + let debug_messenger = gpu::PanicIfErrorMessenger {}; + let instance = gpu::Instance::new(Some(&debug_messenger), None, false).unwrap(); + let device = gpu::Device::new(instance.clone(), &instance.physical_devices()[0]).unwrap(); + + let mut gpu_search_context = GpuSearchContext::new( + device, + groups_count, + &test.vector_storage.borrow(), + None, + m, + m0, + ef, + num_vectors, + false, + true, + &false.into(), + ) + .unwrap(); + + for level in (0..batched_points.levels_count).rev() { + let level_m = graph_layers_builder.get_m(level); + gpu_search_context.clear(level_m).unwrap(); + + build_level_on_gpu(&mut gpu_search_context, &batched_points, 0, level, |_| {}).unwrap(); + + gpu_search_context + .download_links(level, &graph_layers_builder) + .unwrap(); + } + + graph_layers_builder + } + + #[test] + fn test_gpu_hnsw_level_equivalency() { + let _ = env_logger::builder() + .is_test(true) + .filter_level(log::LevelFilter::Trace) + .try_init(); + + let num_vectors = 1024; + let dim = 64; + let m = 8; + let m0 = 16; + let ef = 32; + + let test = create_gpu_graph_test_data(num_vectors, dim, m, m0, ef, 0); + let graph_layers_builder = build_gpu_graph(&test, 1); + + compare_graph_layers_builders(&test.graph_layers_builder, &graph_layers_builder); + } + + #[test] + fn test_gpu_hnsw_level_quality() { + let _ = env_logger::builder() + .is_test(true) + .filter_level(log::LevelFilter::Trace) + .try_init(); + + let num_vectors = 1024; + let dim = 64; + let m = 8; + let m0 = 16; + let ef = 32; + let groups_count = 4; + let searches_count = 20; + let top = 10; + + let test = create_gpu_graph_test_data(num_vectors, dim, m, m0, ef, searches_count); + let graph_layers_builder = build_gpu_graph(&test, groups_count); + + check_graph_layers_builders_quality(graph_layers_builder, test, top, ef, 0.8) + } +} diff --git a/lib/segment/src/index/hnsw_index/gpu/gpu_links.rs b/lib/segment/src/index/hnsw_index/gpu/gpu_links.rs new file mode 100644 index 0000000000..693e96248d --- /dev/null +++ b/lib/segment/src/index/hnsw_index/gpu/gpu_links.rs @@ -0,0 +1,298 @@ +use std::collections::HashMap; +use std::sync::Arc; + +use common::types::PointOffsetType; + +use super::shader_builder::ShaderBuilderParameters; +use crate::common::operation_error::{OperationError, OperationResult}; +use crate::index::hnsw_index::gpu::GPU_TIMEOUT; +use crate::index::hnsw_index::graph_layers::GraphLayersBase; +use crate::index::hnsw_index::graph_layers_builder::GraphLayersBuilder; + +#[repr(C)] +struct GpuLinksParamsBuffer { + m: u32, + links_capacity: u32, +} + +pub struct GpuLinks { + pub m: usize, + pub links_capacity: usize, + pub points_count: usize, + pub max_patched_points: usize, + pub device: Arc, + pub links_buffer: Arc, + pub params_buffer: Arc, + pub patch_buffer: Arc, + pub patched_points: Vec<(PointOffsetType, usize)>, + pub descriptor_set_layout: Arc, + pub descriptor_set: Arc, +} + +impl ShaderBuilderParameters for GpuLinks { + fn shader_includes(&self) -> HashMap { + HashMap::from([ + ( + "iterators.comp".to_string(), + include_str!("shaders/iterators.comp").to_string(), + ), + ( + "links.comp".to_string(), + include_str!("shaders/links.comp").to_string(), + ), + ]) + } + + fn shader_defines(&self) -> HashMap> { + let mut defines = HashMap::new(); + defines.insert( + "LINKS_CAPACITY".to_owned(), + Some(self.links_capacity.to_string()), + ); + defines + } +} + +impl GpuLinks { + pub fn new( + device: Arc, + m: usize, + links_capacity: usize, + points_count: usize, + max_patched_points: usize, + ) -> gpu::GpuResult { + let links_buffer = gpu::Buffer::new( + device.clone(), + "Links buffer", + gpu::BufferType::Storage, + points_count * (links_capacity + 1) * std::mem::size_of::(), + )?; + let params_buffer = gpu::Buffer::new( + device.clone(), + "Links params buffer", + gpu::BufferType::Uniform, + std::mem::size_of::(), + )?; + let links_patch_capacity = + max_patched_points * (links_capacity + 1) * std::mem::size_of::(); + let patch_buffer = gpu::Buffer::new( + device.clone(), + "Links patch buffer", + gpu::BufferType::CpuToGpu, + links_patch_capacity + std::mem::size_of::(), + )?; + + let params = GpuLinksParamsBuffer { + m: m as u32, + links_capacity: links_capacity as u32, + }; + patch_buffer.upload(¶ms, 0)?; + + let mut upload_context = gpu::Context::new(device.clone())?; + upload_context.copy_gpu_buffer( + patch_buffer.clone(), + params_buffer.clone(), + 0, + 0, + std::mem::size_of::(), + )?; + upload_context.clear_buffer(links_buffer.clone())?; + upload_context.run()?; + upload_context.wait_finish(GPU_TIMEOUT)?; + + let descriptor_set_layout = gpu::DescriptorSetLayout::builder() + .add_uniform_buffer(0) + .add_storage_buffer(1) + .build(device.clone())?; + + let descriptor_set = gpu::DescriptorSet::builder(descriptor_set_layout.clone()) + .add_uniform_buffer(0, params_buffer.clone()) + .add_storage_buffer(1, links_buffer.clone()) + .build()?; + + Ok(Self { + m, + links_capacity, + points_count, + max_patched_points, + device, + links_buffer, + params_buffer, + patch_buffer, + patched_points: vec![], + descriptor_set_layout, + descriptor_set, + }) + } + + pub fn update_params(&mut self, context: &mut gpu::Context, m: usize) -> OperationResult<()> { + self.m = m; + + let params = GpuLinksParamsBuffer { + m: m as u32, + links_capacity: self.links_capacity as u32, + }; + let links_patch_capacity = self.max_patched_points + * (self.links_capacity + 1) + * std::mem::size_of::(); + self.patch_buffer.upload(¶ms, links_patch_capacity)?; + + context.copy_gpu_buffer( + self.patch_buffer.clone(), + self.params_buffer.clone(), + links_patch_capacity, + 0, + std::mem::size_of::(), + )?; + Ok(()) + } + + pub fn clear(&mut self, gpu_context: &mut gpu::Context) -> OperationResult<()> { + if !self.patched_points.is_empty() { + self.patched_points.clear(); + } + gpu_context.clear_buffer(self.links_buffer.clone())?; + Ok(()) + } + + pub fn apply_gpu_patches(&mut self, gpu_context: &mut gpu::Context) -> OperationResult<()> { + for (i, &(patched_point_id, patched_links_count)) in self.patched_points.iter().enumerate() + { + let patch_start_index = + i * (self.links_capacity + 1) * std::mem::size_of::(); + let patch_size = (patched_links_count + 1) * std::mem::size_of::(); + let links_start_index = patched_point_id as usize + * (self.links_capacity + 1) + * std::mem::size_of::(); + gpu_context.copy_gpu_buffer( + self.patch_buffer.clone(), + self.links_buffer.clone(), + patch_start_index, + links_start_index, + patch_size, + )?; + } + self.patched_points.clear(); + Ok(()) + } + + pub fn set_links( + &mut self, + point_id: PointOffsetType, + links: &[PointOffsetType], + ) -> OperationResult<()> { + if self.patched_points.len() >= self.max_patched_points { + return Err(OperationError::service_error("Gpu links patches are full")); + } + + let mut patch_start_index = self.patched_points.len() + * (self.links_capacity + 1) + * std::mem::size_of::(); + self.patch_buffer + .upload(&(links.len() as u32), patch_start_index)?; + patch_start_index += std::mem::size_of::(); + self.patch_buffer.upload_slice(links, patch_start_index)?; + self.patched_points.push((point_id, links.len())); + + Ok(()) + } + + pub fn upload_links( + &mut self, + level: usize, + graph_layers_builder: &GraphLayersBuilder, + context: &mut gpu::Context, + ) -> OperationResult<()> { + self.update_params(context, graph_layers_builder.get_m(level))?; + self.clear(context)?; + context.run()?; + context.wait_finish(GPU_TIMEOUT)?; + + let timer = std::time::Instant::now(); + let points: Vec<_> = (0..graph_layers_builder.links_layers.len()) + .filter(|&point_id| { + graph_layers_builder.get_point_level(point_id as PointOffsetType) >= level + }) + .filter(|&point_id| { + !graph_layers_builder.links_layers[point_id][level] + .read() + .is_empty() + }) + .collect(); + + for points_slice in points.chunks(self.max_patched_points) { + for &point_id in points_slice { + let links = graph_layers_builder.links_layers[point_id][level].read(); + self.set_links(point_id as PointOffsetType, &links)?; + } + self.apply_gpu_patches(context)?; + context.run()?; + context.wait_finish(GPU_TIMEOUT)?; + } + + log::trace!("Upload links on level {level} time: {:?}", timer.elapsed()); + Ok(()) + } + + pub fn download_links( + &mut self, + level: usize, + graph_layers_builder: &GraphLayersBuilder, + context: &mut gpu::Context, + ) -> OperationResult<()> { + let timer = std::time::Instant::now(); + + let links_patch_capacity = self.max_patched_points + * (self.links_capacity + 1) + * std::mem::size_of::(); + let download_buffer = gpu::Buffer::new( + self.device.clone(), + "Download links staging buffer", + gpu::BufferType::GpuToCpu, + links_patch_capacity, + ) + .unwrap(); + + let points = (0..graph_layers_builder.links_layers.len() as PointOffsetType) + .filter(|&point_id| graph_layers_builder.get_point_level(point_id) >= level) + .collect::>(); + + for chunk_index in 0..points.len().div_ceil(self.max_patched_points) { + let start = chunk_index * self.max_patched_points; + let end = (start + self.max_patched_points).min(points.len()); + let chunk_size = end - start; + for (i, &point_id) in points[start..end].iter().enumerate() { + let links_size = (self.links_capacity + 1) * std::mem::size_of::(); + context.copy_gpu_buffer( + self.links_buffer.clone(), + download_buffer.clone(), + point_id as usize * links_size, + i * links_size, + links_size, + )?; + } + context.run()?; + context.wait_finish(GPU_TIMEOUT)?; + + let mut links = + vec![PointOffsetType::default(); chunk_size * (self.links_capacity + 1)]; + download_buffer.download_slice(&mut links, 0)?; + + for (index, chunk) in links.chunks(self.links_capacity + 1).enumerate() { + let point_id = points[start + index] as usize; + let links_count = chunk[0] as usize; + let links = &chunk[1..=links_count]; + let mut dst = graph_layers_builder.links_layers[point_id][level].write(); + dst.clear(); + dst.extend_from_slice(links); + } + } + + log::trace!( + "Download links for level {} in time {:?}", + level, + timer.elapsed() + ); + Ok(()) + } +} diff --git a/lib/segment/src/index/hnsw_index/gpu/gpu_nearest_heap.rs b/lib/segment/src/index/hnsw_index/gpu/gpu_nearest_heap.rs new file mode 100644 index 0000000000..8961f777f8 --- /dev/null +++ b/lib/segment/src/index/hnsw_index/gpu/gpu_nearest_heap.rs @@ -0,0 +1,259 @@ +use std::collections::HashMap; +use std::sync::Arc; + +use super::shader_builder::ShaderBuilderParameters; + +pub struct GpuNearestHeap { + pub ef: usize, + pub capacity: usize, + pub device: Arc, +} + +impl ShaderBuilderParameters for GpuNearestHeap { + fn shader_includes(&self) -> HashMap { + HashMap::from([( + "nearest_heap.comp".to_string(), + include_str!("shaders/nearest_heap.comp").to_string(), + )]) + } + + fn shader_defines(&self) -> HashMap> { + let mut defines = HashMap::new(); + defines.insert("NEAREST_HEAP_EF".to_owned(), Some(self.ef.to_string())); + defines.insert( + "NEAREST_HEAP_CAPACITY".to_owned(), + Some(self.capacity.to_string()), + ); + defines + } +} + +impl GpuNearestHeap { + pub fn new(device: Arc, ef: usize, capacity: usize) -> gpu::GpuResult { + assert!(capacity >= ef); + let capacity = capacity.div_ceil(device.subgroup_size()) * device.subgroup_size(); + Ok(Self { + ef, + capacity, + device, + }) + } +} + +#[cfg(test)] +mod tests { + use common::fixed_length_priority_queue::FixedLengthPriorityQueue; + use common::types::{PointOffsetType, ScoredPointOffset}; + use rand::rngs::StdRng; + use rand::{Rng, SeedableRng}; + + use super::*; + use crate::index::hnsw_index::gpu::shader_builder::ShaderBuilder; + use crate::index::hnsw_index::gpu::GPU_TIMEOUT; + + #[repr(C)] + struct TestParams { + input_counts: u32, + } + + #[test] + fn test_gpu_nearest_heap() { + let _ = env_logger::builder() + .is_test(true) + .filter_level(log::LevelFilter::Trace) + .try_init(); + + let ef = 100; + let points_count = 1024; + let groups_count = 8; + let inputs_count = points_count; + + let mut rng = StdRng::seed_from_u64(42); + let inputs_data: Vec = (0..inputs_count * groups_count) + .map(|i| ScoredPointOffset { + idx: i as PointOffsetType, + score: rng.gen_range(-1.0..1.0), + }) + .collect(); + + let debug_messenger = gpu::PanicIfErrorMessenger {}; + let instance = gpu::Instance::new(Some(&debug_messenger), None, false).unwrap(); + let device = gpu::Device::new(instance.clone(), &instance.physical_devices()[0]).unwrap(); + + let gpu_nearest_heap = GpuNearestHeap::new(device.clone(), ef, ef).unwrap(); + + let shader = ShaderBuilder::new(device.clone()) + .with_shader_code(include_str!("shaders/tests/test_nearest_heap.comp")) + .with_parameters(&gpu_nearest_heap) + .build() + .unwrap(); + + let input_points_buffer = gpu::Buffer::new( + device.clone(), + "Nearest heap input points buffer", + gpu::BufferType::Storage, + inputs_count * groups_count * std::mem::size_of::(), + ) + .unwrap(); + + let upload_staging_buffer = gpu::Buffer::new( + device.clone(), + "Nearest heap upload staging buffer", + gpu::BufferType::CpuToGpu, + inputs_count * groups_count * std::mem::size_of::(), + ) + .unwrap(); + upload_staging_buffer.upload_slice(&inputs_data, 0).unwrap(); + + let mut context = gpu::Context::new(device.clone()).unwrap(); + context + .copy_gpu_buffer( + upload_staging_buffer.clone(), + input_points_buffer.clone(), + 0, + 0, + input_points_buffer.size(), + ) + .unwrap(); + context.run().unwrap(); + context.wait_finish(GPU_TIMEOUT).unwrap(); + + let test_params_buffer = gpu::Buffer::new( + device.clone(), + "Nearest heap test params buffer", + gpu::BufferType::Uniform, + std::mem::size_of::(), + ) + .unwrap(); + upload_staging_buffer + .upload( + &TestParams { + input_counts: inputs_count as u32, + }, + 0, + ) + .unwrap(); + context + .copy_gpu_buffer( + upload_staging_buffer, + test_params_buffer.clone(), + 0, + 0, + test_params_buffer.size(), + ) + .unwrap(); + context.run().unwrap(); + context.wait_finish(GPU_TIMEOUT).unwrap(); + + let scores_output_buffer = gpu::Buffer::new( + device.clone(), + "Nearest heap scores output buffer", + gpu::BufferType::Storage, + inputs_count * groups_count * std::mem::size_of::(), + ) + .unwrap(); + let sorted_output_buffer = gpu::Buffer::new( + device.clone(), + "Nearest heap sorted output buffer", + gpu::BufferType::Storage, + ef * groups_count * std::mem::size_of::(), + ) + .unwrap(); + + let descriptor_set_layout = gpu::DescriptorSetLayout::builder() + .add_uniform_buffer(0) + .add_storage_buffer(1) + .add_storage_buffer(2) + .add_storage_buffer(3) + .build(device.clone()) + .unwrap(); + + let descriptor_set = gpu::DescriptorSet::builder(descriptor_set_layout.clone()) + .add_uniform_buffer(0, test_params_buffer.clone()) + .add_storage_buffer(1, input_points_buffer.clone()) + .add_storage_buffer(2, scores_output_buffer.clone()) + .add_storage_buffer(3, sorted_output_buffer.clone()) + .build() + .unwrap(); + + let pipeline = gpu::Pipeline::builder() + .add_descriptor_set_layout(0, descriptor_set_layout.clone()) + .add_shader(shader.clone()) + .build(device.clone()) + .unwrap(); + + context + .bind_pipeline(pipeline, &[descriptor_set.clone()]) + .unwrap(); + context.dispatch(groups_count, 1, 1).unwrap(); + context.run().unwrap(); + context.wait_finish(GPU_TIMEOUT).unwrap(); + + let download_staging_buffer = gpu::Buffer::new( + device.clone(), + "Nearest heap download staging buffer", + gpu::BufferType::GpuToCpu, + std::cmp::max(scores_output_buffer.size(), sorted_output_buffer.size()), + ) + .unwrap(); + context + .copy_gpu_buffer( + scores_output_buffer.clone(), + download_staging_buffer.clone(), + 0, + 0, + scores_output_buffer.size(), + ) + .unwrap(); + context.run().unwrap(); + context.wait_finish(GPU_TIMEOUT).unwrap(); + let mut scores_output = vec![0.0; inputs_count * groups_count]; + download_staging_buffer + .download_slice(&mut scores_output, 0) + .unwrap(); + + let mut scores_output_cpu = vec![0.0; inputs_count * groups_count]; + let mut sorted_output_cpu = vec![PointOffsetType::default(); ef * groups_count]; + for group in 0..groups_count { + let mut queue = FixedLengthPriorityQueue::::new(ef); + for i in 0..inputs_count { + let scored_point = inputs_data[group * inputs_count + i]; + queue.push(scored_point); + scores_output_cpu[group * inputs_count + i] = queue.top().unwrap().score; + } + let sorted = queue.into_vec(); + for i in 0..ef { + sorted_output_cpu[group * ef + i] = sorted[i].idx; + } + } + + let mut nearest_gpu: Vec = + vec![Default::default(); gpu_nearest_heap.ef * groups_count]; + context + .copy_gpu_buffer( + sorted_output_buffer.clone(), + download_staging_buffer.clone(), + 0, + 0, + nearest_gpu.len() * std::mem::size_of::(), + ) + .unwrap(); + context.run().unwrap(); + context.wait_finish(GPU_TIMEOUT).unwrap(); + download_staging_buffer + .download_slice(nearest_gpu.as_mut_slice(), 0) + .unwrap(); + + let mut sorted_output_gpu = Vec::new(); + for group in 0..groups_count { + let mut nearest_group = Vec::new(); + for i in 0..ef { + nearest_group.push(nearest_gpu[group * gpu_nearest_heap.ef + i]); + } + sorted_output_gpu.extend(nearest_group); + } + + assert_eq!(scores_output, scores_output_cpu); + assert_eq!(sorted_output_gpu, sorted_output_cpu); + } +} diff --git a/lib/segment/src/index/hnsw_index/gpu/gpu_search_context.rs b/lib/segment/src/index/hnsw_index/gpu/gpu_search_context.rs new file mode 100644 index 0000000000..907f9d10e0 --- /dev/null +++ b/lib/segment/src/index/hnsw_index/gpu/gpu_search_context.rs @@ -0,0 +1,1234 @@ +use std::collections::HashMap; +use std::sync::atomic::AtomicBool; +use std::sync::Arc; + +use common::types::{PointOffsetType, ScoredPointOffset}; +use itertools::Itertools; + +use super::gpu_candidates_heap::GpuCandidatesHeap; +use super::gpu_links::GpuLinks; +use super::gpu_nearest_heap::GpuNearestHeap; +use super::gpu_vector_storage::GpuVectorStorage; +use super::gpu_visited_flags::GpuVisitedFlags; +use super::shader_builder::ShaderBuilderParameters; +use super::GPU_TIMEOUT; +use crate::common::operation_error::{OperationError, OperationResult}; +use crate::index::hnsw_index::gpu::shader_builder::ShaderBuilder; +use crate::index::hnsw_index::graph_layers_builder::GraphLayersBuilder; +use crate::vector_storage::quantized::quantized_vectors::QuantizedVectors; +use crate::vector_storage::{VectorStorage, VectorStorageEnum}; + +#[derive(Clone, Copy, Debug)] +#[repr(C)] +pub struct GpuRequest { + pub id: PointOffsetType, + pub entry: PointOffsetType, +} + +#[derive(Clone, Debug)] +pub struct GpuGraphLinksPatch { + pub id: PointOffsetType, + pub links: Vec, +} + +pub struct GpuSearchContext { + pub device: Arc, + pub context: gpu::Context, + pub groups_count: usize, + pub gpu_vector_storage: GpuVectorStorage, + pub gpu_links: GpuLinks, + pub gpu_nearest_heap: GpuNearestHeap, + pub gpu_candidates_heap: GpuCandidatesHeap, + pub gpu_visited_flags: GpuVisitedFlags, + pub is_dirty_links: bool, + + pub upload_staging_buffer: Arc, + pub download_staging_buffer: Arc, + + pub requests_buffer: Arc, + pub responses_buffer: Arc, + + pub greedy_descriptor_set: Arc, + pub greedy_pipeline: Arc, + + pub search_responses_buffer: Arc, + pub search_descriptor_set: Arc, + pub search_pipeline: Arc, + + pub patches_responses_buffer: Arc, + pub patches_descriptor_set: Arc, + pub patches_pipeline: Arc, + + pub insert_atomics_buffer: Arc, + pub insert_descriptor_set: Arc, + pub insert_pipeline: Arc, + + pub updates_timer: std::time::Duration, + pub updates_count: usize, + pub patches_timer: std::time::Duration, + pub patches_count: usize, + pub exact: bool, +} + +struct GpuSearchContextParams { + exact: bool, +} + +impl ShaderBuilderParameters for GpuSearchContextParams { + fn shader_includes(&self) -> HashMap { + HashMap::from([( + "search_context.comp".to_string(), + include_str!("shaders/search_context.comp").to_string(), + )]) + } + + fn shader_defines(&self) -> HashMap> { + let mut defines = HashMap::new(); + if self.exact { + defines.insert("EXACT".to_owned(), None); + } + defines + } +} + +struct GpuSearchContextGroupAllocation { + groups_count: usize, + gpu_nearest_heap: GpuNearestHeap, + gpu_candidates_heap: GpuCandidatesHeap, + gpu_visited_flags: GpuVisitedFlags, + upload_staging_buffer: Arc, + download_staging_buffer: Arc, + requests_buffer: Arc, + responses_buffer: Arc, + search_responses_buffer: Arc, + patches_responses_buffer: Arc, + insert_atomics_buffer: Arc, +} + +impl GpuSearchContext { + #[allow(clippy::too_many_arguments)] + pub fn new( + device: Arc, + max_groups_count: usize, + vector_storage: &VectorStorageEnum, + quantized_storage: Option<&QuantizedVectors>, + m: usize, + m0: usize, + ef: usize, + max_patched_points: usize, + force_half_precision: bool, + exact: bool, + stopped: &AtomicBool, + ) -> OperationResult { + let points_count = vector_storage.total_vector_count(); + let search_context_params = GpuSearchContextParams { exact }; + + let gpu_vector_storage = GpuVectorStorage::new( + device.clone(), + vector_storage, + quantized_storage, + force_half_precision, + stopped, + )?; + let gpu_links = GpuLinks::new(device.clone(), m, m0, points_count, max_patched_points)?; + + let allocation_timer = std::time::Instant::now(); + let GpuSearchContextGroupAllocation { + groups_count, + gpu_nearest_heap, + gpu_candidates_heap, + gpu_visited_flags, + upload_staging_buffer, + download_staging_buffer, + requests_buffer, + responses_buffer, + search_responses_buffer, + patches_responses_buffer, + insert_atomics_buffer, + } = Self::allocate_grouped_data(device.clone(), max_groups_count, points_count, ef, m0) + .map_err(|_| OperationError::service_error("Failed to allocate gpu data"))?; + log::debug!( + "GPU groups count = {groups_count} (max = {max_groups_count}), allocation time: {:?}", + allocation_timer.elapsed() + ); + + let greedy_search_shader = ShaderBuilder::new(device.clone()) + .with_shader_code(include_str!("shaders/run_greedy_search.comp")) + .with_parameters(&gpu_vector_storage) + .with_parameters(&gpu_nearest_heap) + .with_parameters(&gpu_candidates_heap) + .with_parameters(&gpu_links) + .with_parameters(&gpu_visited_flags) + .with_parameters(&search_context_params) + .build()?; + + let insert_shader = ShaderBuilder::new(device.clone()) + .with_shader_code(include_str!("shaders/run_insert_vector.comp")) + .with_parameters(&gpu_vector_storage) + .with_parameters(&gpu_nearest_heap) + .with_parameters(&gpu_candidates_heap) + .with_parameters(&gpu_links) + .with_parameters(&gpu_visited_flags) + .with_parameters(&search_context_params) + .build()?; + + let search_shader = ShaderBuilder::new(device.clone()) + .with_shader_code(include_str!("shaders/tests/test_hnsw_search.comp")) + .with_parameters(&gpu_vector_storage) + .with_parameters(&gpu_nearest_heap) + .with_parameters(&gpu_candidates_heap) + .with_parameters(&gpu_links) + .with_parameters(&gpu_visited_flags) + .with_parameters(&search_context_params) + .build()?; + + let patches_shader = ShaderBuilder::new(device.clone()) + .with_shader_code(include_str!("shaders/run_get_patch.comp")) + .with_parameters(&gpu_vector_storage) + .with_parameters(&gpu_nearest_heap) + .with_parameters(&gpu_candidates_heap) + .with_parameters(&gpu_links) + .with_parameters(&gpu_visited_flags) + .with_parameters(&search_context_params) + .build()?; + + let greedy_descriptor_set_layout = gpu::DescriptorSetLayout::builder() + .add_storage_buffer(0) + .add_storage_buffer(1) + .build(device.clone())?; + + let greedy_descriptor_set = + gpu::DescriptorSet::builder(greedy_descriptor_set_layout.clone()) + .add_storage_buffer(0, requests_buffer.clone()) + .add_storage_buffer(1, responses_buffer.clone()) + .build()?; + + let greedy_pipeline = gpu::Pipeline::builder() + .add_descriptor_set_layout(0, greedy_descriptor_set_layout.clone()) + .add_descriptor_set_layout(1, gpu_vector_storage.descriptor_set_layout()) + .add_descriptor_set_layout(2, gpu_links.descriptor_set_layout.clone()) + .add_descriptor_set_layout(3, gpu_visited_flags.descriptor_set_layout.clone()) + .add_shader(greedy_search_shader.clone()) + .build(device.clone())?; + + let search_descriptor_set_layout = gpu::DescriptorSetLayout::builder() + .add_storage_buffer(0) + .add_storage_buffer(1) + .build(device.clone())?; + + let search_descriptor_set = + gpu::DescriptorSet::builder(search_descriptor_set_layout.clone()) + .add_storage_buffer(0, requests_buffer.clone()) + .add_storage_buffer(1, search_responses_buffer.clone()) + .build()?; + + let search_pipeline = gpu::Pipeline::builder() + .add_descriptor_set_layout(0, search_descriptor_set_layout.clone()) + .add_descriptor_set_layout(1, gpu_vector_storage.descriptor_set_layout()) + .add_descriptor_set_layout(2, gpu_links.descriptor_set_layout.clone()) + .add_descriptor_set_layout(3, gpu_visited_flags.descriptor_set_layout.clone()) + .add_shader(search_shader.clone()) + .build(device.clone())?; + + let patches_descriptor_set_layout = gpu::DescriptorSetLayout::builder() + .add_storage_buffer(0) + .add_storage_buffer(1) + .add_storage_buffer(2) + .build(device.clone())?; + + let patches_descriptor_set = + gpu::DescriptorSet::builder(patches_descriptor_set_layout.clone()) + .add_storage_buffer(0, requests_buffer.clone()) + .add_storage_buffer(1, patches_responses_buffer.clone()) + .add_storage_buffer(2, responses_buffer.clone()) + .build()?; + + let patches_pipeline = gpu::Pipeline::builder() + .add_descriptor_set_layout(0, patches_descriptor_set_layout.clone()) + .add_descriptor_set_layout(1, gpu_vector_storage.descriptor_set_layout()) + .add_descriptor_set_layout(2, gpu_links.descriptor_set_layout.clone()) + .add_descriptor_set_layout(3, gpu_visited_flags.descriptor_set_layout.clone()) + .add_shader(patches_shader.clone()) + .build(device.clone())?; + + let insert_descriptor_set_layout = gpu::DescriptorSetLayout::builder() + .add_storage_buffer(0) + .add_storage_buffer(1) + .add_storage_buffer(2) + .build(device.clone())?; + + let insert_descriptor_set = + gpu::DescriptorSet::builder(insert_descriptor_set_layout.clone()) + .add_storage_buffer(0, requests_buffer.clone()) + .add_storage_buffer(1, responses_buffer.clone()) + .add_storage_buffer(2, insert_atomics_buffer.clone()) + .build()?; + + let insert_pipeline = gpu::Pipeline::builder() + .add_descriptor_set_layout(0, insert_descriptor_set_layout.clone()) + .add_descriptor_set_layout(1, gpu_vector_storage.descriptor_set_layout()) + .add_descriptor_set_layout(2, gpu_links.descriptor_set_layout.clone()) + .add_descriptor_set_layout(3, gpu_visited_flags.descriptor_set_layout.clone()) + .add_shader(insert_shader.clone()) + .build(device.clone())?; + + let context = gpu::Context::new(device.clone())?; + Ok(Self { + gpu_vector_storage, + gpu_links, + gpu_nearest_heap, + gpu_candidates_heap, + gpu_visited_flags, + device, + context, + groups_count, + is_dirty_links: false, + upload_staging_buffer, + download_staging_buffer, + requests_buffer, + responses_buffer, + greedy_descriptor_set, + greedy_pipeline, + search_responses_buffer, + search_descriptor_set, + search_pipeline, + patches_responses_buffer, + patches_descriptor_set, + patches_pipeline, + insert_atomics_buffer, + insert_descriptor_set, + insert_pipeline, + updates_timer: Default::default(), + updates_count: 0, + patches_timer: Default::default(), + patches_count: 0, + exact, + }) + } + + fn allocate_grouped_data( + device: Arc, + max_groups_count: usize, + points_count: usize, + ef: usize, + m0: usize, + ) -> gpu::GpuResult { + let mut visited_flags_factor = 1.0; + // TODO(gpu): move 32 to config + while visited_flags_factor < 32.0 { + // TODO(gpu): decrease only gpu visited flags + // TODO(gpu): restart only if OOM, for other errors return an error + if let Ok(alloc) = Self::try_allocate_grouped_data( + device.clone(), + max_groups_count, + points_count, + visited_flags_factor, + ef, + m0, + ) { + return Ok(alloc); + } + visited_flags_factor *= 1.5; + } + panic!("Failed to allocate gpu data") + } + + fn try_allocate_grouped_data( + device: Arc, + groups_count: usize, + points_count: usize, + visited_flags_factor: f32, + ef: usize, + m0: usize, + ) -> OperationResult { + let requests_buffer = gpu::Buffer::new( + device.clone(), + "Search requests buffer", + gpu::BufferType::Storage, + groups_count * std::mem::size_of::(), + )?; + let responses_buffer = gpu::Buffer::new( + device.clone(), + "Search responses buffer", + gpu::BufferType::Storage, + groups_count * std::mem::size_of::(), + )?; + + let gpu_nearest_heap = GpuNearestHeap::new(device.clone(), ef, std::cmp::max(ef, m0 + 1))?; + let gpu_candidates_heap = + GpuCandidatesHeap::new(device.clone(), std::cmp::max(ef, m0 + 1))?; + let gpu_visited_flags = GpuVisitedFlags::new( + device.clone(), + groups_count, + (points_count as f32 / visited_flags_factor) as usize, + )?; + + // todo(gpu): remove this buffer + let patches_responses_buffer = gpu::Buffer::new( + device.clone(), + "Patch responses buffer", + gpu::BufferType::Storage, + groups_count * ((m0 + 1) * (m0 + 2)) * std::mem::size_of::(), + )?; + + let upload_staging_buffer = gpu::Buffer::new( + device.clone(), + "Search context upload staging buffer", + gpu::BufferType::CpuToGpu, + requests_buffer.size(), + )?; + let download_staging_buffer = gpu::Buffer::new( + device.clone(), + "Search context download staging buffer", + gpu::BufferType::GpuToCpu, + patches_responses_buffer.size() + responses_buffer.size(), + )?; + + let search_responses_buffer = gpu::Buffer::new( + device.clone(), + "Search responses buffer", + gpu::BufferType::Storage, + groups_count * ef * std::mem::size_of::(), + )?; + let insert_atomics_buffer = gpu::Buffer::new( + device.clone(), + "Insert atomics buffer", + gpu::BufferType::Storage, + points_count * std::mem::size_of::(), + )?; + + Ok(GpuSearchContextGroupAllocation { + groups_count, + gpu_nearest_heap, + gpu_candidates_heap, + gpu_visited_flags, + upload_staging_buffer, + download_staging_buffer, + requests_buffer, + responses_buffer, + search_responses_buffer, + patches_responses_buffer, + insert_atomics_buffer, + }) + } + + pub fn download_responses(&mut self, count: usize) -> OperationResult> { + self.context.copy_gpu_buffer( + self.responses_buffer.clone(), + self.download_staging_buffer.clone(), + 0, + 0, + count * std::mem::size_of::(), + )?; + self.run_context()?; + let mut gpu_responses = vec![PointOffsetType::default(); count]; + self.download_staging_buffer + .download_slice(&mut gpu_responses, 0)?; + + // TODO(gpu) if response if uint::MAX, we need to reallocate + for i in 0..count { + if gpu_responses[i] == PointOffsetType::MAX { + return Err(OperationError::service_error("Candidates heap overflow")); + } + } + + Ok(gpu_responses) + } + + pub fn search( + &mut self, + requests: &[GpuRequest], + ) -> OperationResult>> { + if requests.len() > self.groups_count { + return Err(OperationError::service_error( + "Too many gpu greedy search requests", + )); + } + + if self.is_dirty() { + self.apply_links_patch().unwrap(); + } + self.gpu_visited_flags.clear(&mut self.context)?; + + self.upload_staging_buffer.upload_slice(requests, 0)?; + self.context.copy_gpu_buffer( + self.upload_staging_buffer.clone(), + self.requests_buffer.clone(), + 0, + 0, + std::mem::size_of_val(requests), + )?; + self.run_context()?; + + self.context.bind_pipeline( + self.search_pipeline.clone(), + &[ + self.search_descriptor_set.clone(), + self.gpu_vector_storage.descriptor_set(), + self.gpu_links.descriptor_set.clone(), + self.gpu_visited_flags.descriptor_set.clone(), + ], + )?; + self.context.dispatch(requests.len(), 1, 1)?; + self.run_context()?; + + // Download response + self.context.copy_gpu_buffer( + self.search_responses_buffer.clone(), + self.download_staging_buffer.clone(), + 0, + 0, + requests.len() * self.gpu_nearest_heap.ef * std::mem::size_of::(), + )?; + self.run_context()?; + let mut gpu_responses = + vec![ScoredPointOffset::default(); requests.len() * self.gpu_nearest_heap.ef]; + self.download_staging_buffer + .download_slice(&mut gpu_responses, 0)?; + Ok(gpu_responses + .chunks(self.gpu_nearest_heap.ef) + .map(|r| { + r.iter() + .take_while(|s| s.idx != PointOffsetType::MAX) + .cloned() + .collect_vec() + }) + .collect_vec()) + } + + pub fn greedy_search( + &mut self, + requests: &[GpuRequest], + prev_results_count: usize, + ) -> OperationResult> { + if requests.len() > self.groups_count { + return Err(OperationError::service_error( + "Too many gpu greedy search requests", + )); + } + + let timer = std::time::Instant::now(); + + if self.is_dirty() { + self.apply_links_patch().unwrap(); + } + + // upload requests + self.upload_staging_buffer.upload_slice(requests, 0)?; + self.context.copy_gpu_buffer( + self.upload_staging_buffer.clone(), + self.requests_buffer.clone(), + 0, + 0, + std::mem::size_of_val(requests), + )?; + + // download previous results + if prev_results_count > 0 { + self.context.copy_gpu_buffer( + self.responses_buffer.clone(), + self.download_staging_buffer.clone(), + 0, + 0, + prev_results_count * std::mem::size_of::(), + )?; + } + self.run_context()?; + + self.context.bind_pipeline( + self.greedy_pipeline.clone(), + &[ + self.greedy_descriptor_set.clone(), + self.gpu_vector_storage.descriptor_set(), + self.gpu_links.descriptor_set.clone(), + self.gpu_visited_flags.descriptor_set.clone(), + ], + )?; + self.context.dispatch(requests.len(), 1, 1)?; + self.run_context()?; + + self.updates_timer += timer.elapsed(); + self.updates_count += 1; + + if prev_results_count > 0 { + let mut gpu_responses = vec![PointOffsetType::default(); prev_results_count]; + self.download_staging_buffer + .download_slice(&mut gpu_responses, 0)?; + Ok(gpu_responses) + } else { + Ok(vec![]) + } + } + + pub fn run_insert_vector( + &mut self, + requests: &[GpuRequest], + prev_results_count: usize, + ) -> OperationResult> { + if requests.len() > self.groups_count { + return Err(OperationError::service_error("Too many gpu patch requests")); + } + + let timer = std::time::Instant::now(); + + if self.is_dirty() { + self.apply_links_patch().unwrap(); + } + self.gpu_visited_flags.clear(&mut self.context)?; + + // clear atomics + if self.gpu_visited_flags.params.generation == 1 { + self.context + .clear_buffer(self.insert_atomics_buffer.clone())?; + } + + // upload requests + self.upload_staging_buffer.upload_slice(requests, 0)?; + self.context.copy_gpu_buffer( + self.upload_staging_buffer.clone(), + self.requests_buffer.clone(), + 0, + 0, + std::mem::size_of_val(requests), + )?; + + // download previous results + if prev_results_count > 0 { + self.context.copy_gpu_buffer( + self.responses_buffer.clone(), + self.download_staging_buffer.clone(), + 0, + 0, + prev_results_count * std::mem::size_of::(), + )?; + } + self.run_context()?; + + self.context.bind_pipeline( + self.insert_pipeline.clone(), + &[ + self.insert_descriptor_set.clone(), + self.gpu_vector_storage.descriptor_set(), + self.gpu_links.descriptor_set.clone(), + self.gpu_visited_flags.descriptor_set.clone(), + ], + )?; + self.context.dispatch(requests.len(), 1, 1)?; + self.run_context()?; + + self.patches_timer += timer.elapsed(); + self.patches_count += 1; + + if prev_results_count > 0 { + let mut gpu_responses = vec![PointOffsetType::default(); prev_results_count]; + self.download_staging_buffer + .download_slice(&mut gpu_responses, 0)?; + Ok(gpu_responses) + } else { + Ok(vec![]) + } + } + + pub fn run_get_patch( + &mut self, + requests: &[GpuRequest], + ) -> OperationResult<(Vec>, Vec)> { + if requests.len() > self.groups_count { + return Err(OperationError::service_error("Too many gpu patch requests")); + } + + if self.is_dirty() { + self.apply_links_patch().unwrap(); + } + self.gpu_visited_flags.clear(&mut self.context)?; + + self.upload_staging_buffer.upload_slice(requests, 0)?; + self.context.copy_gpu_buffer( + self.upload_staging_buffer.clone(), + self.requests_buffer.clone(), + 0, + 0, + std::mem::size_of_val(requests), + )?; + self.run_context()?; + + self.context.bind_pipeline( + self.patches_pipeline.clone(), + &[ + self.patches_descriptor_set.clone(), + self.gpu_vector_storage.descriptor_set(), + self.gpu_links.descriptor_set.clone(), + self.gpu_visited_flags.descriptor_set.clone(), + ], + )?; + self.context.dispatch(requests.len(), 1, 1)?; + + let timer = std::time::Instant::now(); + self.run_context()?; + self.patches_timer += timer.elapsed(); + + // Download response + self.context.copy_gpu_buffer( + self.responses_buffer.clone(), + self.download_staging_buffer.clone(), + 0, + 0, + requests.len() * std::mem::size_of::(), + )?; + self.context.copy_gpu_buffer( + self.patches_responses_buffer.clone(), + self.download_staging_buffer.clone(), + 0, + self.responses_buffer.size(), + self.patches_responses_buffer.size(), + )?; + self.run_context()?; + let mut new_entries = vec![PointOffsetType::default(); requests.len()]; + self.download_staging_buffer + .download_slice(&mut new_entries, 0)?; + + let mut patches_data = vec![ + PointOffsetType::default(); + self.patches_responses_buffer.size() + / std::mem::size_of::() + ]; + self.download_staging_buffer + .download_slice(&mut patches_data, self.responses_buffer.size())?; + + let m = self.gpu_links.m; + let mut all_patches = vec![]; + for i in 0..requests.len() { + let patch_size = m + 2; + let all_patches_size = (m + 1) * patch_size; + let mut patches_offset = i * all_patches_size; + + let mut patches = vec![]; + for _ in 0..m + 1 { + let point_id = patches_data[patches_offset]; + if point_id == PointOffsetType::MAX { + break; + } + let links_count = patches_data[patches_offset + 1] as usize; + let links = &patches_data[patches_offset + 2..patches_offset + 2 + links_count]; + patches.push(GpuGraphLinksPatch { + id: point_id, + links: links.to_vec(), + }); + patches_offset += patch_size; + } + all_patches.push(patches); + } + + Ok((all_patches, new_entries)) + } + + pub fn set_links( + &mut self, + point_id: PointOffsetType, + links: &[PointOffsetType], + ) -> OperationResult<()> { + self.is_dirty_links = true; + self.gpu_links.set_links(point_id, links) + } + + pub fn upload_links( + &mut self, + level: usize, + graph_layers_builder: &GraphLayersBuilder, + ) -> OperationResult<()> { + self.gpu_links + .upload_links(level, graph_layers_builder, &mut self.context) + } + + pub fn download_links( + &mut self, + level: usize, + graph_layers_builder: &GraphLayersBuilder, + ) -> OperationResult<()> { + self.gpu_links + .download_links(level, graph_layers_builder, &mut self.context) + } + + pub fn clear(&mut self, new_m: usize) -> OperationResult<()> { + self.gpu_links.update_params(&mut self.context, new_m)?; + self.gpu_links.clear(&mut self.context)?; + self.run_context()?; + Ok(()) + } + + pub fn apply_links_patch(&mut self) -> OperationResult<()> { + self.gpu_links.apply_gpu_patches(&mut self.context)?; + self.is_dirty_links = false; + Ok(()) + } + + pub fn run_context(&mut self) -> OperationResult<()> { + self.context.run()?; + self.context.wait_finish(GPU_TIMEOUT)?; + Ok(()) + } + + fn is_dirty(&self) -> bool { + self.is_dirty_links + } +} + +#[cfg(test)] +mod tests { + use common::types::ScoredPointOffset; + use itertools::Itertools; + use rand::rngs::StdRng; + use rand::SeedableRng; + + use super::*; + use crate::common::rocksdb_wrapper::{open_db, DB_VECTOR_CF}; + use crate::fixtures::index_fixtures::{FakeFilterContext, TestRawScorerProducer}; + use crate::index::hnsw_index::graph_layers::GraphLayersBase; + use crate::index::hnsw_index::graph_layers_builder::GraphLayersBuilder; + use crate::index::hnsw_index::point_scorer::FilteredScorer; + use crate::spaces::simple::DotProductMetric; + use crate::types::Distance; + use crate::vector_storage::chunked_vector_storage::VectorOffsetType; + use crate::vector_storage::dense::simple_dense_vector_storage::open_simple_dense_vector_storage; + + #[repr(C)] + struct TestSearchRequest { + id: PointOffsetType, + entry: PointOffsetType, + } + + struct TestData { + gpu_search_context: GpuSearchContext, + gpu_search_context_params: GpuSearchContextParams, + vector_holder: TestRawScorerProducer, + graph_layers_builder: GraphLayersBuilder, + } + + fn create_test_data( + num_vectors: usize, + groups_count: usize, + dim: usize, + m: usize, + ef: usize, + ) -> TestData { + // Generate random vectors + let mut rng = StdRng::seed_from_u64(42); + let vector_holder = TestRawScorerProducer::::new( + dim, + num_vectors + groups_count, + &mut rng, + ); + + // upload vectors to storage + let dir = tempfile::Builder::new().prefix("db_dir").tempdir().unwrap(); + let db = open_db(dir.path(), &[DB_VECTOR_CF]).unwrap(); + let mut storage = + open_simple_dense_vector_storage(db, DB_VECTOR_CF, dim, Distance::Dot, &false.into()) + .unwrap(); + for idx in 0..(num_vectors + groups_count) { + let v = vector_holder.get_vector(idx as PointOffsetType); + storage + .insert_vector(idx as PointOffsetType, v.as_vec_ref()) + .unwrap(); + } + + // Build HNSW index + let mut graph_layers_builder = GraphLayersBuilder::new(num_vectors, m, m, ef, 1, true); + for idx in 0..(num_vectors as PointOffsetType) { + let level = graph_layers_builder.get_random_layer(&mut rng); + graph_layers_builder.set_levels(idx, level); + } + for idx in 0..(num_vectors as PointOffsetType) { + let fake_filter_context = FakeFilterContext {}; + let added_vector = vector_holder.vectors.get(idx as VectorOffsetType).to_vec(); + let raw_scorer = vector_holder.get_raw_scorer(added_vector.clone()).unwrap(); + let scorer = FilteredScorer::new(raw_scorer.as_ref(), Some(&fake_filter_context)); + graph_layers_builder.link_new_point(idx, scorer); + raw_scorer.take_hardware_counter().discard_results(); + } + + // Create GPU search context + let debug_messenger = gpu::PanicIfErrorMessenger {}; + let instance = gpu::Instance::new(Some(&debug_messenger), None, false).unwrap(); + let device = gpu::Device::new(instance.clone(), &instance.physical_devices()[0]).unwrap(); + + let mut gpu_search_context = GpuSearchContext::new( + device, + groups_count, + &storage, + None, + m, + m, + ef, + num_vectors, + false, + true, + &false.into(), + ) + .unwrap(); + + // Upload HNSW links to GPU + for idx in 0..(num_vectors as PointOffsetType) { + let mut links = vec![]; + graph_layers_builder.links_map(idx, 0, |link| links.push(link)); + gpu_search_context.set_links(idx, &links).unwrap(); + } + gpu_search_context.apply_links_patch().unwrap(); + gpu_search_context.run_context().unwrap(); + + TestData { + gpu_search_context, + vector_holder, + graph_layers_builder, + gpu_search_context_params: GpuSearchContextParams { exact: true }, + } + } + + #[test] + fn test_gpu_hnsw_search_on_level() { + let _ = env_logger::builder() + .is_test(true) + .filter_level(log::LevelFilter::Trace) + .try_init(); + + let num_vectors = 1024; + let groups_count = 8; + let dim = 64; + let m = 16; + let ef = 32; + + let mut test = create_test_data(num_vectors, groups_count, dim, m, ef); + + // create request data + let mut search_requests = vec![]; + for i in 0..groups_count { + search_requests.push(GpuRequest { + id: (num_vectors + i) as PointOffsetType, + entry: 0, + }); + } + + let gpu_responses_1 = test.gpu_search_context.search(&search_requests).unwrap(); + + // restart search to check reset + let gpu_responses_2 = test.gpu_search_context.search(&search_requests).unwrap(); + + // Check response + for i in 0..groups_count { + let fake_filter_context = FakeFilterContext {}; + let added_vector = test.vector_holder.vectors.get(num_vectors + i).to_vec(); + let raw_scorer = test + .vector_holder + .get_raw_scorer(added_vector.clone()) + .unwrap(); + let mut scorer = FilteredScorer::new(raw_scorer.as_ref(), Some(&fake_filter_context)); + let entry = ScoredPointOffset { + idx: 0, + score: scorer.score_point(0), + }; + let search_result = test + .graph_layers_builder + .search_on_level(entry, 0, ef, &mut scorer) + .into_vec(); + for (cpu, (gpu_1, gpu_2)) in search_result + .iter() + .zip(gpu_responses_1[i].iter().zip(gpu_responses_2[i].iter())) + { + assert_eq!(cpu.idx, gpu_1.idx); + assert_eq!(cpu.idx, gpu_2.idx); + assert!((cpu.score - gpu_1.score).abs() < 1e-5); + assert!((cpu.score - gpu_2.score).abs() < 1e-5); + } + raw_scorer.take_hardware_counter().discard_results(); + } + } + + #[test] + fn test_gpu_greedy_search() { + let _ = env_logger::builder() + .is_test(true) + .filter_level(log::LevelFilter::Trace) + .try_init(); + + let num_vectors = 1024; + let groups_count = 8; + let dim = 64; + let m = 16; + let ef = 32; + + let mut test = create_test_data(num_vectors, groups_count, dim, m, ef); + + // create request data + let mut search_requests = vec![]; + for i in 0..groups_count { + search_requests.push(GpuRequest { + id: (num_vectors + i) as PointOffsetType, + entry: 0, + }); + } + + test.gpu_search_context + .greedy_search(&search_requests, 0) + .unwrap(); + let gpu_responses = test + .gpu_search_context + .download_responses(groups_count) + .unwrap(); + + // Check response + for i in 0..groups_count { + let fake_filter_context = FakeFilterContext {}; + let added_vector = test.vector_holder.vectors.get(num_vectors + i).to_vec(); + let raw_scorer = test + .vector_holder + .get_raw_scorer(added_vector.clone()) + .unwrap(); + let mut scorer = FilteredScorer::new(raw_scorer.as_ref(), Some(&fake_filter_context)); + let search_result = test + .graph_layers_builder + .search_entry_on_level(0, 0, &mut scorer); + assert_eq!(search_result.idx, gpu_responses[i]); + raw_scorer.take_hardware_counter().discard_results(); + } + } + + #[test] + fn test_gpu_hnsw_patch() { + let _ = env_logger::builder() + .is_test(true) + .filter_level(log::LevelFilter::Trace) + .try_init(); + + let num_vectors = 1024; + let groups_count = 8; + let dim = 64; + let m = 16; + let ef = 32; + + let mut test = create_test_data(num_vectors, groups_count, dim, m, ef); + + // create request data + let mut requests = vec![]; + for i in 0..groups_count { + requests.push(GpuRequest { + id: (num_vectors + i) as PointOffsetType, + entry: 0, + }); + } + + let (patches, new_entries) = test.gpu_search_context.run_get_patch(&requests).unwrap(); + + for (i, gpu_patches) in patches.iter().enumerate() { + let fake_filter_context = FakeFilterContext {}; + let added_vector = test.vector_holder.vectors.get(num_vectors + i).to_vec(); + let raw_scorer = test + .vector_holder + .get_raw_scorer(added_vector.clone()) + .unwrap(); + let mut scorer = FilteredScorer::new(raw_scorer.as_ref(), Some(&fake_filter_context)); + + let (cpu_patches, cpu_new_entry) = + test.graph_layers_builder + .get_patch(requests[i], 0, &mut scorer); + raw_scorer.take_hardware_counter().discard_results(); + + assert_eq!(new_entries[i], *cpu_new_entry.first().unwrap()); + assert_eq!(gpu_patches.len(), cpu_patches.len()); + for (gpu_patch, cpu_patch) in gpu_patches.iter().zip(cpu_patches.iter()) { + assert_eq!(gpu_patch.id, cpu_patch.id); + assert_eq!(gpu_patch.links, cpu_patch.links); + } + } + } + + #[test] + fn test_gpu_heuristic() { + let _ = env_logger::builder() + .is_test(true) + .filter_level(log::LevelFilter::Trace) + .try_init(); + + let num_vectors = 1024; + let groups_count = 8; + let dim = 64; + let m = 16; + let ef = 32; + + let mut test = create_test_data(num_vectors, groups_count, dim, m, ef); + + // create request data + let mut search_requests = vec![]; + for i in 0..groups_count { + search_requests.push(TestSearchRequest { + id: (num_vectors + i) as PointOffsetType, + entry: 0, + }); + } + + // upload search requests to GPU + let search_requests_buffer = gpu::Buffer::new( + test.gpu_search_context.device.clone(), + "Search requests buffer", + gpu::BufferType::Storage, + search_requests.len() * std::mem::size_of::(), + ) + .unwrap(); + let upload_staging_buffer = gpu::Buffer::new( + test.gpu_search_context.device.clone(), + "Search context upload staging buffer", + gpu::BufferType::CpuToGpu, + search_requests.len() * std::mem::size_of::(), + ) + .unwrap(); + upload_staging_buffer + .upload_slice(&search_requests, 0) + .unwrap(); + test.gpu_search_context + .context + .copy_gpu_buffer( + upload_staging_buffer.clone(), + search_requests_buffer.clone(), + 0, + 0, + search_requests_buffer.size(), + ) + .unwrap(); + test.gpu_search_context.run_context().unwrap(); + + // create response and response staging buffers + let responses_buffer = gpu::Buffer::new( + test.gpu_search_context.device.clone(), + "Search responses buffer", + gpu::BufferType::Storage, + groups_count * ef * std::mem::size_of::(), + ) + .unwrap(); + let responses_staging_buffer = gpu::Buffer::new( + test.gpu_search_context.device.clone(), + "Search responses staging buffer", + gpu::BufferType::GpuToCpu, + responses_buffer.size(), + ) + .unwrap(); + + // Create test pipeline + let shader = ShaderBuilder::new(test.gpu_search_context.device.clone()) + .with_shader_code(include_str!("shaders/tests/test_heuristic.comp")) + .with_parameters(&test.gpu_search_context.gpu_vector_storage) + .with_parameters(&test.gpu_search_context.gpu_nearest_heap) + .with_parameters(&test.gpu_search_context.gpu_candidates_heap) + .with_parameters(&test.gpu_search_context.gpu_links) + .with_parameters(&test.gpu_search_context.gpu_visited_flags) + .with_parameters(&test.gpu_search_context_params) + .build() + .unwrap(); + + let descriptor_set_layout = gpu::DescriptorSetLayout::builder() + .add_storage_buffer(0) + .add_storage_buffer(1) + .build(test.gpu_search_context.device.clone()) + .unwrap(); + + let descriptor_set = gpu::DescriptorSet::builder(descriptor_set_layout.clone()) + .add_storage_buffer(0, search_requests_buffer.clone()) + .add_storage_buffer(1, responses_buffer.clone()) + .build() + .unwrap(); + + let pipeline = gpu::Pipeline::builder() + .add_descriptor_set_layout(0, descriptor_set_layout.clone()) + .add_descriptor_set_layout( + 1, + test.gpu_search_context + .gpu_vector_storage + .descriptor_set_layout(), + ) + .add_descriptor_set_layout( + 2, + test.gpu_search_context + .gpu_links + .descriptor_set_layout + .clone(), + ) + .add_descriptor_set_layout( + 3, + test.gpu_search_context + .gpu_visited_flags + .descriptor_set_layout + .clone(), + ) + .add_shader(shader.clone()) + .build(test.gpu_search_context.device.clone()) + .unwrap(); + + test.gpu_search_context + .context + .bind_pipeline( + pipeline.clone(), + &[ + descriptor_set.clone(), + test.gpu_search_context.gpu_vector_storage.descriptor_set(), + test.gpu_search_context.gpu_links.descriptor_set.clone(), + test.gpu_search_context + .gpu_visited_flags + .descriptor_set + .clone(), + ], + ) + .unwrap(); + test.gpu_search_context + .context + .dispatch(groups_count, 1, 1) + .unwrap(); + test.gpu_search_context.run_context().unwrap(); + + // Download response + test.gpu_search_context + .context + .copy_gpu_buffer( + responses_buffer.clone(), + responses_staging_buffer.clone(), + 0, + 0, + responses_buffer.size(), + ) + .unwrap(); + test.gpu_search_context.run_context().unwrap(); + let mut gpu_responses = vec![ScoredPointOffset::default(); groups_count * ef]; + responses_staging_buffer + .download_slice(&mut gpu_responses, 0) + .unwrap(); + let gpu_responses = gpu_responses + .chunks_exact(ef) + .map(|r| r.to_owned()) + .collect_vec(); + + // Check response + for i in 0..groups_count { + let fake_filter_context = FakeFilterContext {}; + let added_vector = test.vector_holder.vectors.get(num_vectors + i).to_vec(); + let raw_scorer = test + .vector_holder + .get_raw_scorer(added_vector.clone()) + .unwrap(); + let mut scorer = FilteredScorer::new(raw_scorer.as_ref(), Some(&fake_filter_context)); + let entry = ScoredPointOffset { + idx: 0, + score: scorer.score_point(0), + }; + let search_result = + test.graph_layers_builder + .search_on_level(entry, 0, ef, &mut scorer); + + let scorer_fn = |a, b| scorer.score_internal(a, b); + + let heuristic = + GraphLayersBuilder::select_candidates_with_heuristic(search_result, m, scorer_fn); + + for (&cpu, gpu) in heuristic.iter().zip(gpu_responses[i].iter()) { + assert_eq!(cpu, gpu.idx); + } + raw_scorer.take_hardware_counter().discard_results(); + } + } +} diff --git a/lib/segment/src/index/hnsw_index/gpu/gpu_visited_flags.rs b/lib/segment/src/index/hnsw_index/gpu/gpu_visited_flags.rs new file mode 100644 index 0000000000..6e270dd403 --- /dev/null +++ b/lib/segment/src/index/hnsw_index/gpu/gpu_visited_flags.rs @@ -0,0 +1,124 @@ +use std::collections::HashMap; +use std::sync::Arc; + +use super::shader_builder::ShaderBuilderParameters; +use super::GPU_TIMEOUT; +use crate::common::operation_error::OperationResult; + +#[repr(C)] +pub struct GpuVisitedFlagsParamsBuffer { + pub generation: u32, +} + +pub struct GpuVisitedFlags { + pub device: Arc, + pub params: GpuVisitedFlagsParamsBuffer, + pub params_buffer: Arc, + pub params_staging_buffer: Arc, + pub visited_flags_buffer: Arc, + pub descriptor_set_layout: Arc, + pub descriptor_set: Arc, + pub capacity: usize, +} + +impl ShaderBuilderParameters for GpuVisitedFlags { + fn shader_includes(&self) -> HashMap { + HashMap::from([( + "visited_flags.comp".to_string(), + include_str!("shaders/visited_flags.comp").to_string(), + )]) + } + + fn shader_defines(&self) -> HashMap> { + let mut defines = HashMap::new(); + defines.insert( + "VISITED_FLAGS_CAPACITY".to_owned(), + Some(self.capacity.to_string()), + ); + defines + } +} + +impl GpuVisitedFlags { + pub fn new( + device: Arc, + groups_count: usize, + points_count: usize, + ) -> OperationResult { + let alignment = std::mem::size_of::(); + let points_count = points_count.div_ceil(alignment) * alignment; + + let params_buffer = gpu::Buffer::new( + device.clone(), + "Visited flags params buffer", + gpu::BufferType::Uniform, + std::mem::size_of::(), + )?; + let params_staging_buffer = gpu::Buffer::new( + device.clone(), + "Visited flags params staging buffer", + gpu::BufferType::CpuToGpu, + std::mem::size_of::(), + )?; + let visited_flags_buffer = gpu::Buffer::new( + device.clone(), + "Visited flags buffer", + gpu::BufferType::Storage, + groups_count * points_count * std::mem::size_of::(), + )?; + + let params = GpuVisitedFlagsParamsBuffer { generation: 1 }; + params_staging_buffer.upload(¶ms, 0)?; + + let mut upload_context = gpu::Context::new(device.clone())?; + upload_context.copy_gpu_buffer( + params_staging_buffer.clone(), + params_buffer.clone(), + 0, + 0, + std::mem::size_of::(), + )?; + upload_context.run()?; + upload_context.wait_finish(GPU_TIMEOUT)?; + + let descriptor_set_layout = gpu::DescriptorSetLayout::builder() + .add_uniform_buffer(0) + .add_storage_buffer(1) + .build(device.clone())?; + + let descriptor_set = gpu::DescriptorSet::builder(descriptor_set_layout.clone()) + .add_uniform_buffer(0, params_buffer.clone()) + .add_storage_buffer(1, visited_flags_buffer.clone()) + .build()?; + + Ok(Self { + device, + params, + params_buffer, + params_staging_buffer, + visited_flags_buffer, + descriptor_set_layout, + descriptor_set, + capacity: points_count, + }) + } + + pub fn clear(&mut self, gpu_context: &mut gpu::Context) -> OperationResult<()> { + if self.params.generation == 255 { + self.params.generation = 1; + gpu_context.clear_buffer(self.visited_flags_buffer.clone())?; + } else { + self.params.generation += 1; + } + + self.params_staging_buffer.upload(&self.params, 0)?; + gpu_context.copy_gpu_buffer( + self.params_staging_buffer.clone(), + self.params_buffer.clone(), + 0, + 0, + self.params_buffer.size(), + )?; + Ok(()) + } +} diff --git a/lib/segment/src/index/hnsw_index/gpu/mod.rs b/lib/segment/src/index/hnsw_index/gpu/mod.rs index b6af2b5f1f..ca8d46b2ee 100644 --- a/lib/segment/src/index/hnsw_index/gpu/mod.rs +++ b/lib/segment/src/index/hnsw_index/gpu/mod.rs @@ -1,8 +1,353 @@ +pub mod batched_points; +pub mod cpu_level_builder; +pub mod devices_manager; +pub mod gpu_candidates_heap; +pub mod gpu_graph_builder; +pub mod gpu_level_builder; +pub mod gpu_links; +pub mod gpu_nearest_heap; +pub mod gpu_search_context; pub mod gpu_vector_storage; +pub mod gpu_visited_flags; pub mod shader_builder; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::sync::Arc; + +use batched_points::BatchedPoints; +use bitvec::vec::BitVec; +use devices_manager::DevicesMaganer; +use lazy_static::lazy_static; +use parking_lot::{Mutex, RwLock}; + +use super::graph_layers_builder::GraphLayersBuilder; +use crate::common::operation_error::{OperationError, OperationResult}; + +lazy_static! { + static ref GPU_INSTANCE: OperationResult> = create_gpu_instance(); + static ref GPU_DEVICE: OperationResult> = create_gpu_device(); + static ref GPU_DEVICE_FILER: Mutex = Mutex::new("".to_string()); + pub static ref GPU_DEVICES_MANAGER: OperationResult = init_devices_manager(); +} + /// Each GPU operation has a timeout by Vulkan API specification. /// Choose large enough timeout. /// We cannot use too small timeout and check stopper in the loop because /// GPU resources should be alive while GPU operation is in progress. static GPU_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(60); + +static GPU_INDEXING: AtomicBool = AtomicBool::new(false); +static GPU_WAIT_FREE: AtomicBool = AtomicBool::new(true); +static GPU_FORCE_HALF_PRECISION: AtomicBool = AtomicBool::new(false); +static GPU_MAX_GROUPS: AtomicUsize = AtomicUsize::new(GPU_MAX_GROUPS_COUNT_DEFAULT); +pub const GPU_MAX_GROUPS_COUNT_DEFAULT: usize = 512; +static GPU_MIN_POINTS_COUNT: AtomicUsize = AtomicUsize::new(10_000); +static GPU_DEVICE_START_INDEX: AtomicUsize = AtomicUsize::new(0); +static GPU_DEVICES_COUNT: AtomicUsize = AtomicUsize::new(usize::MAX); +static GPU_PARALLEL_INDEXES: AtomicUsize = AtomicUsize::new(0); + +fn create_gpu_instance() -> OperationResult> { + Ok(gpu::Instance::new(None, None, false)?) +} + +fn init_devices_manager() -> OperationResult { + let instance = GPU_INSTANCE.clone()?; + let filter = GPU_DEVICE_FILER.lock().clone(); + // TODO(gpu): add start index and count + let devices_manager = DevicesMaganer::new( + instance, + &filter, + GPU_DEVICE_START_INDEX.load(Ordering::Relaxed), + GPU_DEVICES_COUNT.load(Ordering::Relaxed), + GPU_WAIT_FREE.load(Ordering::Relaxed), + GPU_PARALLEL_INDEXES.load(Ordering::Relaxed), + )?; + Ok(devices_manager) +} + +fn create_gpu_device() -> OperationResult> { + let instance = GPU_INSTANCE.clone()?; + let filter = GPU_DEVICE_FILER.lock().clone(); + + for physical_device in instance.physical_devices() { + let device_name = physical_device.name.to_owned(); + if filter.is_empty() || device_name.contains(&filter) { + log::info!("Found GPU device: {}", device_name); + return Ok(gpu::Device::new(instance.clone(), physical_device)?); + } + } + + Err(OperationError::service_error("No GPU device found")) +} + +pub fn set_gpu_indexing(gpu_indexing: bool) { + GPU_INDEXING.store(gpu_indexing, Ordering::Relaxed); +} + +pub fn get_gpu_indexing() -> bool { + GPU_INDEXING.load(Ordering::Relaxed) +} + +pub fn set_wait_free(wait_free: bool) { + GPU_WAIT_FREE.store(wait_free, Ordering::Relaxed); +} + +pub fn set_gpu_force_half_precision(force_half_precision: bool) { + GPU_FORCE_HALF_PRECISION.store(force_half_precision, Ordering::Relaxed); +} + +pub fn get_gpu_force_half_precision() -> bool { + GPU_FORCE_HALF_PRECISION.load(Ordering::Relaxed) +} + +pub fn set_gpu_max_groups(max_groups: Option) { + if let Some(max_groups) = max_groups { + GPU_MAX_GROUPS.store(max_groups, Ordering::Relaxed); + } +} + +pub fn get_gpu_max_groups() -> usize { + GPU_MAX_GROUPS.load(Ordering::Relaxed) +} + +pub fn set_device_index(device_index: Option) { + if let Some(device_index) = device_index { + GPU_DEVICE_START_INDEX.store(device_index, Ordering::Relaxed); + } +} + +pub fn set_devices_count(devices_count: Option) { + if let Some(devices_count) = devices_count { + GPU_DEVICES_COUNT.store(devices_count, Ordering::Relaxed); + } +} + +pub fn set_device_filter(device_filter: &str) { + let mut filter = GPU_DEVICE_FILER.lock(); + *filter = device_filter.to_string(); +} + +pub fn set_gpu_parallel_indexes(parallel_indexes: usize) { + GPU_PARALLEL_INDEXES.store(parallel_indexes, Ordering::Relaxed); +} + +pub fn get_gpu_min_points_count() -> usize { + GPU_MIN_POINTS_COUNT.load(Ordering::Relaxed) +} + +pub fn create_graph_layers_builder( + batched_points: &BatchedPoints, + num_vectors: usize, + m: usize, + m0: usize, + ef: usize, + entry_points_num: usize, +) -> OperationResult { + // create graph layers builder + let mut graph_layers_builder = + GraphLayersBuilder::new(num_vectors, m, m0, ef, entry_points_num, true); + + // mark all vectors as ready + graph_layers_builder.ready_list = RwLock::new(BitVec::repeat(true, num_vectors)); + + // set first entry point + graph_layers_builder.set_levels( + batched_points.first_point_id, + batched_points.levels_count - 1, + ); + graph_layers_builder.entry_points.lock().new_point( + batched_points.first_point_id, + batched_points.levels_count - 1, + |_| true, + ); + + // set levels + for linking_point in &batched_points.points { + graph_layers_builder.set_levels(linking_point.point_id, linking_point.level); + let mut entry_points = graph_layers_builder.entry_points.lock(); + entry_points.new_point(linking_point.point_id, linking_point.level, |_| true); + } + + Ok(graph_layers_builder) +} + +#[cfg(test)] +mod tests { + use ahash::HashSet; + use common::types::PointOffsetType; + use rand::rngs::StdRng; + use rand::SeedableRng; + use tempfile::TempDir; + + use super::batched_points::BatchedPoints; + use crate::common::rocksdb_wrapper::{open_db, DB_VECTOR_CF}; + use crate::data_types::vectors::DenseVector; + use crate::fixtures::index_fixtures::{FakeFilterContext, TestRawScorerProducer}; + use crate::fixtures::payload_fixtures::random_vector; + use crate::index::hnsw_index::graph_layers::GraphLayers; + use crate::index::hnsw_index::graph_layers_builder::GraphLayersBuilder; + use crate::index::hnsw_index::graph_links::GraphLinksRam; + use crate::index::hnsw_index::point_scorer::FilteredScorer; + use crate::spaces::simple::CosineMetric; + use crate::types::Distance; + use crate::vector_storage::chunked_vector_storage::VectorOffsetType; + use crate::vector_storage::dense::simple_dense_vector_storage::open_simple_dense_vector_storage; + use crate::vector_storage::{VectorStorage, VectorStorageEnum}; + + #[allow(dead_code)] + pub struct GpuGraphTestData { + pub dir: TempDir, + pub vector_storage: VectorStorageEnum, + pub vector_holder: TestRawScorerProducer, + pub graph_layers_builder: GraphLayersBuilder, + pub search_vectors: Vec, + } + + pub fn create_gpu_graph_test_data( + num_vectors: usize, + dim: usize, + m: usize, + m0: usize, + ef: usize, + search_counts: usize, + ) -> GpuGraphTestData { + // Generate random vectors + let mut rng = StdRng::seed_from_u64(42); + let vector_holder = TestRawScorerProducer::::new(dim, num_vectors, &mut rng); + + // upload vectors to storage + let dir = tempfile::Builder::new().prefix("db_dir").tempdir().unwrap(); + let db = open_db(dir.path(), &[DB_VECTOR_CF]).unwrap(); + let mut storage = open_simple_dense_vector_storage( + db, + DB_VECTOR_CF, + dim, + Distance::Cosine, + &false.into(), + ) + .unwrap(); + for idx in 0..num_vectors { + let v = vector_holder.get_vector(idx as PointOffsetType); + storage + .insert_vector(idx as PointOffsetType, v.as_vec_ref()) + .unwrap(); + } + + // Build HNSW index + let mut graph_layers_builder = GraphLayersBuilder::new(num_vectors, m, m0, ef, 1, true); + for idx in 0..(num_vectors as PointOffsetType) { + let level = graph_layers_builder.get_random_layer(&mut rng); + graph_layers_builder.set_levels(idx, level); + } + + let mut ids: Vec<_> = (0..num_vectors as PointOffsetType).collect(); + BatchedPoints::sort_points_by_level( + &|point_id| graph_layers_builder.get_point_level(point_id), + &mut ids, + ); + + for &idx in &ids { + let fake_filter_context = FakeFilterContext {}; + let added_vector = vector_holder.vectors.get(idx as VectorOffsetType).to_vec(); + let raw_scorer = vector_holder.get_raw_scorer(added_vector.clone()).unwrap(); + let scorer = FilteredScorer::new(raw_scorer.as_ref(), Some(&fake_filter_context)); + graph_layers_builder.link_new_point(idx, scorer); + raw_scorer.take_hardware_counter().discard_results(); + } + + let search_vectors = (0..search_counts) + .map(|_| random_vector(&mut rng, dim)) + .collect(); + + GpuGraphTestData { + dir, + vector_storage: storage, + vector_holder, + graph_layers_builder, + search_vectors, + } + } + + pub fn compare_graph_layers_builders( + graph_a: &GraphLayersBuilder, + graph_b: &GraphLayersBuilder, + ) { + assert_eq!(graph_a.links_layers.len(), graph_b.links_layers.len()); + let num_vectors = graph_a.links_layers.len(); + for point_id in 0..num_vectors as PointOffsetType { + let levels_a = graph_a.get_point_level(point_id); + let levels_b = graph_b.get_point_level(point_id); + assert_eq!(levels_a, levels_b); + + for level in (0..levels_a + 1).rev() { + let links_a = graph_a.links_layers[point_id as usize][level] + .read() + .clone(); + let links_b = graph_b.links_layers[point_id as usize][level] + .read() + .clone(); + if links_a != links_b { + log::error!("Wrong links point_id={} at level {}", point_id, level); + } + assert_eq!(links_a, links_b); + } + } + } + + pub fn check_graph_layers_builders_quality( + graph: GraphLayersBuilder, + test: GpuGraphTestData, + top: usize, + ef: usize, + accuracy: f32, + ) { + let graph: GraphLayers = graph.into_graph_layers(None).unwrap(); + let ref_graph: GraphLayers = + test.graph_layers_builder.into_graph_layers(None).unwrap(); + + let mut total_sames = 0; + let total_top = top * test.search_vectors.len(); + for search_vector in &test.search_vectors { + let fake_filter_context = FakeFilterContext {}; + let raw_scorer = test + .vector_holder + .get_raw_scorer(search_vector.clone()) + .unwrap(); + let scorer = FilteredScorer::new(raw_scorer.as_ref(), Some(&fake_filter_context)); + + let search_result_gpu = graph.search(top, ef, scorer, None); + raw_scorer.take_hardware_counter().discard_results(); + + let fake_filter_context = FakeFilterContext {}; + let raw_scorer = test + .vector_holder + .get_raw_scorer(search_vector.clone()) + .unwrap(); + let scorer = FilteredScorer::new(raw_scorer.as_ref(), Some(&fake_filter_context)); + + let search_result_cpu = ref_graph.search(top, ef, scorer, None); + raw_scorer.take_hardware_counter().discard_results(); + + let mut gpu_set = HashSet::default(); + let mut cpu_set = HashSet::default(); + for (gpu_id, cpu_id) in search_result_gpu.iter().zip(search_result_cpu.iter()) { + gpu_set.insert(gpu_id.idx); + cpu_set.insert(cpu_id.idx); + } + + total_sames += gpu_set.intersection(&cpu_set).count(); + } + log::debug!( + "total_sames: {}, total_top: {}, div {}", + total_sames, + total_top, + total_sames as f32 / total_top as f32 + ); + assert!( + total_sames as f32 >= total_top as f32 * accuracy, + "sames: {}, total_top: {}", + total_sames, + total_top + ); + } +} diff --git a/lib/segment/src/index/hnsw_index/gpu/shaders/bheap.comp b/lib/segment/src/index/hnsw_index/gpu/shaders/bheap.comp new file mode 100644 index 0000000000..d5d8ae9a75 --- /dev/null +++ b/lib/segment/src/index/hnsw_index/gpu/shaders/bheap.comp @@ -0,0 +1,123 @@ +/****** +Shader binary heap. + +Usage example. +Step 1: defines + +#define BHEAP_NAME bheap // provide binary heap name. Name is required for cases when there are many binary heaps in one shader +float heap_data[8]; // reserve storage for binary heap +#define BHEAP_DATA heap_data // provide storage name +#define BHEAP_TYPE float // binary tree element type +#define BHEAP_CMP(a, b) a > b // cmp function +#include "bheap.comp" // define all binary tree functions + +Step 2: init +BHEAP(BHEAP_NAME, init)(0.0, 8, 0); // default value, capacity and offset in storage + +Step 3: usage +BHEAP(BHEAP_NAME, capacity) // get capacity +BHEAP(BHEAP_NAME, offset) // get offset +BHEAP(BHEAP_NAME, size); // get size +BHEAP(BHEAP_NAME, top)(); // get top value +BHEAP(BHEAP_NAME, pop)(); // pop value +BHEAP(BHEAP_NAME, change_top)(1.0); // pop and push +BHEAP(BHEAP_NAME, push)(1.0); // push value +BHEAP(BHEAP_NAME, into_array)(); // convert to sorted array and return array size + +******/ + +#include "common.comp" + +#define BHEAP(NAME, a) CONCAT(NAME, a) + +uint BHEAP(BHEAP_NAME, offset) = 0; +uint BHEAP(BHEAP_NAME, capacity) = 0; +uint BHEAP(BHEAP_NAME, size) = 0; + +// swap +void BHEAP(BHEAP_NAME, swap)(uint i, uint j) { + BHEAP_TYPE tmp = BHEAP_DATA[BHEAP(BHEAP_NAME, offset) + i]; + BHEAP_DATA[BHEAP(BHEAP_NAME, offset) + i] = BHEAP_DATA[BHEAP(BHEAP_NAME, offset) + j]; + BHEAP_DATA[BHEAP(BHEAP_NAME, offset) + j] = tmp; +} + +// compare, call BHEAP_CMP function +bool BHEAP(BHEAP_NAME, cmp)(uint i, uint j) { + return BHEAP_CMP( + (BHEAP_DATA[BHEAP(BHEAP_NAME, offset) + i]), + (BHEAP_DATA[BHEAP(BHEAP_NAME, offset) + j]) + ); +} + +// sift down +void BHEAP(BHEAP_NAME, sift_down)(uint i) { + while (2 * i + 1 < BHEAP(BHEAP_NAME, size)) { + uint left = 2 * i + 1; + uint right = 2 * i + 2; + uint j = left; + if (right < BHEAP(BHEAP_NAME, size) && BHEAP(BHEAP_NAME, cmp)(right, left)) { + j = right; + } + if (BHEAP(BHEAP_NAME, cmp)(i, j)) { + break; + } + BHEAP(BHEAP_NAME, swap)(i, j); + i = j; + } +} + +// push +void BHEAP(BHEAP_NAME, push)(BHEAP_TYPE score) { + uint i = BHEAP(BHEAP_NAME, size); + BHEAP(BHEAP_NAME, size)++; + BHEAP_DATA[BHEAP(BHEAP_NAME, offset) + i] = score; + while (i > 0 && BHEAP(BHEAP_NAME, cmp)(i, (i - 1) / 2)) { + BHEAP(BHEAP_NAME, swap)(i, (i - 1) / 2); + i = (i - 1) / 2; + } +} + +// get top value +BHEAP_TYPE BHEAP(BHEAP_NAME, top)() { + return BHEAP_DATA[BHEAP(BHEAP_NAME, offset)]; +} + +// pop value +BHEAP_TYPE BHEAP(BHEAP_NAME, pop)() { + BHEAP_TYPE top_value = BHEAP(BHEAP_NAME, top)(); + BHEAP(BHEAP_NAME, size)--; + BHEAP_DATA[BHEAP(BHEAP_NAME, offset)] = BHEAP_DATA[BHEAP(BHEAP_NAME, offset) + BHEAP(BHEAP_NAME, size)]; + BHEAP(BHEAP_NAME, sift_down)(0); + return top_value; +} + +// replace top value +void BHEAP(BHEAP_NAME, change_top)(BHEAP_TYPE new_value) { + BHEAP_DATA[BHEAP(BHEAP_NAME, offset)] = new_value; + BHEAP(BHEAP_NAME, sift_down)(0); +} + +// into_array +uint BHEAP(BHEAP_NAME, into_array)() { + uint array_size = BHEAP(BHEAP_NAME, size); + while (BHEAP(BHEAP_NAME, size) > 0) { + uint i = BHEAP(BHEAP_NAME, size) - 1; + BHEAP_TYPE value = BHEAP(BHEAP_NAME, pop)(); + BHEAP_DATA[BHEAP(BHEAP_NAME, offset) + i] = value; + } + return array_size; +} + +// init +void BHEAP(BHEAP_NAME, init)(BHEAP_TYPE none_value, uint capacity, uint offset) { + BHEAP(BHEAP_NAME, offset) = offset; + BHEAP(BHEAP_NAME, capacity) = capacity; + //for (uint i = 0; i < capacity; i++) { + // BHEAP_DATA[BHEAP(BHEAP_NAME, offset) + i] = none_value; + //} +} + +#undef BHEAP_NAME +#undef BHEAP_DATA +#undef BHEAP_TYPE +#undef BHEAP_CMP diff --git a/lib/segment/src/index/hnsw_index/gpu/shaders/candidates_heap.comp b/lib/segment/src/index/hnsw_index/gpu/shaders/candidates_heap.comp new file mode 100644 index 0000000000..5549b480b7 --- /dev/null +++ b/lib/segment/src/index/hnsw_index/gpu/shaders/candidates_heap.comp @@ -0,0 +1,75 @@ +#ifndef CANDIDATES_HEAP_HEADER +#define CANDIDATES_HEAP_HEADER + +#include "common.comp" + +shared ScoredPoint candidates[CANDIDATES_HEAP_CAPACITY]; +uint candidates_count; + +void candidates_heap_init() { + candidates_count = 0; + for (uint i = gl_SubgroupInvocationID; i < CANDIDATES_HEAP_CAPACITY; i += SUBGROUP_SIZE) { + candidates[i] = ScoredPoint(0, positive_infinity); + } +} + +ScoredPoint pop_candidate() { + if (candidates_count == 0) { + return ScoredPoint(0, 0.0); + } + + ScoredPoint max_scored_point = ScoredPoint(0, negative_infinity); + uint max_index = 0; + for (uint i = gl_SubgroupInvocationID; i < candidates_count; i += SUBGROUP_SIZE) { + if (candidates[i].score > max_scored_point.score) { + max_scored_point = candidates[i]; + max_index = i; + } + } + + candidates_count--; + + float subgroup_largest = subgroupMax(max_scored_point.score); + bool is_max_thread = max_scored_point.score == subgroup_largest; + uvec4 subgroup_max_threads = subgroupBallot(is_max_thread); + uint subgroup_max_thread = subgroupBallotFindLSB(subgroup_max_threads); + max_index = subgroupShuffle(max_index, subgroup_max_thread); + max_scored_point.score = subgroupShuffle(max_scored_point.score, subgroup_max_thread); + max_scored_point.id = subgroupShuffle(max_scored_point.id, subgroup_max_thread); + + if (subgroupElect()) { + candidates[max_index] = candidates[candidates_count]; + } + return max_scored_point; +} + +bool push_candidate(ScoredPoint scored_point) { + if (candidates_count == CANDIDATES_HEAP_CAPACITY) { + ScoredPoint min_scored_point = ScoredPoint(0, positive_infinity); + uint min_index = 0; + for (uint i = gl_SubgroupInvocationID; i < candidates_count; i += SUBGROUP_SIZE) { + if (candidates[i].score < min_scored_point.score) { + min_scored_point = candidates[i]; + min_index = i; + } + } + + float subgroup_smallest = subgroupMin(min_scored_point.score); + bool is_min_thread = min_scored_point.score == subgroup_smallest; + uvec4 subgroup_min_threads = subgroupBallot(is_min_thread); + uint subgroup_min_thread = subgroupBallotFindLSB(subgroup_min_threads); + min_index = subgroupShuffle(min_index, subgroup_min_thread); + + if (subgroupElect()) { + candidates[min_index] = scored_point; + } + } else { + if (subgroupElect()) { + candidates[candidates_count] = scored_point; + } + candidates_count++; + } + return true; +} + +#endif diff --git a/lib/segment/src/index/hnsw_index/gpu/shaders/iterators.comp b/lib/segment/src/index/hnsw_index/gpu/shaders/iterators.comp new file mode 100644 index 0000000000..837858c091 --- /dev/null +++ b/lib/segment/src/index/hnsw_index/gpu/shaders/iterators.comp @@ -0,0 +1,27 @@ +#ifndef ITERATORS_HEADER +#define ITERATORS_HEADER + +struct Iterator { + uint begin; + uint end; +}; + +#define NEW_ITERATOR(BEGIN, END) \ + { \ + (BEGIN), \ + (END) \ + } + +#define FOREACH_ITERATOR_VARIABLE_NAME _foreach_iterator + +#define FOREACH(ITERATOR) \ + for ( \ + Iterator FOREACH_ITERATOR_VARIABLE_NAME = ITERATOR; \ + FOREACH_ITERATOR_VARIABLE_NAME.begin < FOREACH_ITERATOR_VARIABLE_NAME.end; \ + FOREACH_ITERATOR_VARIABLE_NAME.begin += 1 \ + ) + +#define FOREACH_VALUE(CONTAINER) \ + CONTAINER[FOREACH_ITERATOR_VARIABLE_NAME.begin] + +#endif diff --git a/lib/segment/src/index/hnsw_index/gpu/shaders/links.comp b/lib/segment/src/index/hnsw_index/gpu/shaders/links.comp new file mode 100644 index 0000000000..a0d136c88e --- /dev/null +++ b/lib/segment/src/index/hnsw_index/gpu/shaders/links.comp @@ -0,0 +1,46 @@ +#ifndef LINKS_HEADER +#define LINKS_HEADER + +#ifndef LINKS_LAYOUT_SET +#error "LINKS_LAYOUT_SET not defined" +#endif + +#include "common.comp" +#include "iterators.comp" + +layout(set = LINKS_LAYOUT_SET, binding = 0) +readonly uniform LinksParams { + uint m; +} links_params; + +layout(set = LINKS_LAYOUT_SET, binding = 1) +buffer Links { + POINT_ID data[]; +} links; + +#define LINKS_CONTAINER links.data + +#define LEVEL_M links_params.m + +#define LINKS_CONTAINER_OFFSET(ID) \ + ID * (LINKS_CAPACITY + 1) + +#define LINKS_NEW_ITERATOR(ID) \ + NEW_ITERATOR((LINKS_CONTAINER_OFFSET(ID) + 1), (LINKS_CONTAINER_OFFSET(ID) + 1 + LINKS_COUNT(ID))) + +#define LINKS_COUNT(ID) \ + LINKS_CONTAINER[LINKS_CONTAINER_OFFSET(ID)] + +#define GET_LINK(ID, I) \ + LINKS_CONTAINER[LINKS_CONTAINER_OFFSET(ID) + 1 + I] + +#define LINKS_FOREACH(ID) \ + FOREACH(LINKS_NEW_ITERATOR(ID)) + +#define LINKS_FOREACH_VALUE \ + FOREACH_VALUE(LINKS_CONTAINER) + +#define LINKS_SET_SIZE(ID, SIZE) \ + LINKS_CONTAINER[LINKS_CONTAINER_OFFSET(ID)] = SIZE + +#endif diff --git a/lib/segment/src/index/hnsw_index/gpu/shaders/nearest_heap.comp b/lib/segment/src/index/hnsw_index/gpu/shaders/nearest_heap.comp new file mode 100644 index 0000000000..bbc2955d0f --- /dev/null +++ b/lib/segment/src/index/hnsw_index/gpu/shaders/nearest_heap.comp @@ -0,0 +1,205 @@ +#ifndef NEAREST_HEAP_HEADER +#define NEAREST_HEAP_HEADER + +#include "common.comp" + +shared ScoredPoint nearest[NEAREST_HEAP_CAPACITY]; + +ScoredPoint nearest_smaller; +uint nearest_smaller_index; +uint nearest_count; + +void nearest_heap_init() { + nearest_count = 0; + nearest_smaller_index = 0; + nearest_smaller = ScoredPoint(UINT_MAX, positive_infinity); + + uint capacity = NEAREST_HEAP_CAPACITY; + for (uint i = gl_SubgroupInvocationID; i < capacity; i += SUBGROUP_SIZE) { + uint buffer_index = i; + nearest[buffer_index] = ScoredPoint(UINT_MAX, negative_infinity); + } +} + +// Push the nearest `ScoredPoint`. Returns true if the `ScoredPoint` was added. +bool push_nearest(ScoredPoint scored_point) { + uint ef = NEAREST_HEAP_EF; + if (nearest_count == ef) { + if (scored_point.score > nearest_smaller.score) { + if (subgroupElect()) { + uint buffer_index = nearest_smaller_index; + nearest[buffer_index] = scored_point; + } + nearest_smaller_index = 0; + nearest_smaller.score = positive_infinity; + for (uint i = gl_SubgroupInvocationID; i < ef; i += SUBGROUP_SIZE) { + uint buffer_index = i; + float score = nearest[buffer_index].score; + if (score < nearest_smaller.score) { + nearest_smaller.score = score; + nearest_smaller_index = i; + } + } + + float subgroup_smaller = subgroupMin(nearest_smaller.score); + bool is_min_thread = nearest_smaller.score == subgroup_smaller; + uvec4 subgroup_min_threads = subgroupBallot(is_min_thread); + uint subgroup_min_thread = subgroupBallotFindLSB(subgroup_min_threads); + nearest_smaller_index = subgroupShuffle(nearest_smaller_index, subgroup_min_thread); + nearest_smaller.score = subgroupShuffle(nearest_smaller.score, subgroup_min_thread); + nearest_smaller.id = nearest[nearest_smaller_index].id; + + return true; + } else { + return false; + } + } else { + if (subgroupElect()) { + uint buffer_index = nearest_count; + nearest[buffer_index] = scored_point; + } + if (scored_point.score < nearest_smaller.score) { + nearest_smaller = scored_point; + nearest_smaller_index = nearest_count; + } + nearest_count++; + return true; + } +} + +uint subgroup_bubble_sort(float score) { + uint invocation = gl_SubgroupInvocationID; + bool is_even = invocation % 2 == 0; + uint invocation_forward = is_even ? invocation + 1 : invocation - 1; + uint invocation_backward = invocation; + if (invocation != SUBGROUP_SIZE - 1 && invocation != 0) { + invocation_backward = is_even ? invocation - 1 : invocation + 1; + } + + uint index = invocation; + for (uint j = 0; j < SUBGROUP_SIZE / 2; j++) { + bool is_even = gl_SubgroupInvocationID % 2 == 0; + bool was_changed = false; + + // step one: compare with the next element + { + float next_score = subgroupShuffle(score, invocation_forward); + uint next_index = subgroupShuffle(index, invocation_forward); + bool need_exchange = is_even ? score < next_score : score > next_score; + if (need_exchange) { + was_changed = true; + score = next_score; + index = next_index; + } + } + + // step two: compare with previous element + { + float next_score = subgroupShuffle(score, invocation_backward); + uint next_index = subgroupShuffle(index, invocation_backward); + bool need_exchange = is_even ? score > next_score : score < next_score; + if (need_exchange) { + was_changed = true; + score = next_score; + index = next_index; + } + } + + if (!subgroupAny(was_changed)) { + break; + } + } + return index; +} + +void sort_nearest_bubble() { + uint iterations_count = nearest_count / (SUBGROUP_SIZE / 2) + uint(sign(float(nearest_count % (SUBGROUP_SIZE / 2)))); + for (uint _cnt = 0; _cnt < iterations_count; _cnt++) { + // do bubble sort in each subgroup to reduce the number of memory accesses + for (uint i = gl_SubgroupInvocationID; i < NEAREST_HEAP_CAPACITY; i += SUBGROUP_SIZE) { + ScoredPoint scored_point = nearest[i]; + + if (i >= nearest_count) { + scored_point.score = negative_infinity; + } + uint permut = subgroup_bubble_sort(scored_point.score); + float shuffled_scored_point_score = subgroupShuffle(scored_point.score, permut); + uint shuffled_scored_point_id = subgroupShuffle(scored_point.id, permut); + nearest[i] = ScoredPoint(shuffled_scored_point_id, shuffled_scored_point_score); + subgroupMemoryBarrierShared(); + } + + if (NEAREST_HEAP_CAPACITY > SUBGROUP_SIZE) { + // do bubble sort in each subgroup to reduce the number of memory accesses + for (uint i = gl_SubgroupInvocationID + (SUBGROUP_SIZE / 2); i < NEAREST_HEAP_CAPACITY - (SUBGROUP_SIZE / 2); i += SUBGROUP_SIZE) { + ScoredPoint scored_point = nearest[i]; + + if (i >= nearest_count) { + scored_point.score = negative_infinity; + } + uint permut = subgroup_bubble_sort(scored_point.score); + float shuffled_scored_point_score = subgroupShuffle(scored_point.score, permut); + uint shuffled_scored_point_id = subgroupShuffle(scored_point.id, permut); + nearest[i] = ScoredPoint(shuffled_scored_point_id, shuffled_scored_point_score); + subgroupMemoryBarrierShared(); + } + } + } +} + +void bitonic_cmp(uint i, uint j) { + uint b_i = i; + uint b_j = j; + if (i < j && j < nearest_count && nearest[b_j].score > nearest[b_i].score) { + ScoredPoint temp = nearest[b_i]; + nearest[b_i] = nearest[b_j]; + nearest[b_j] = temp; + } +} + +void sort_nearest() { + for (uint group_offset = 0; group_offset < nearest_count; group_offset += SUBGROUP_SIZE) { + ScoredPoint scored_point = ScoredPoint(0, negative_infinity); + uint i = group_offset + gl_SubgroupInvocationID; + if (i < nearest_count) { + scored_point = nearest[i]; + } + + uint permut = subgroup_bubble_sort(scored_point.score); + scored_point = ScoredPoint( + subgroupShuffle(scored_point.id, permut), + subgroupShuffle(scored_point.score, permut) + ); + + if (i < nearest_count) { + nearest[i] = scored_point; + } + } + + for (uint k = 64; (k >> 1) < nearest_count; k <<= 1) { + for (uint i = gl_SubgroupInvocationID; i < nearest_count; i += SUBGROUP_SIZE) { + bitonic_cmp(i, i ^ (k - 1)); + } + subgroupMemoryBarrierShared(); + + for (uint j = k >> 1; 0 < j; j >>= 1) { + for (uint i = gl_SubgroupInvocationID; i < nearest_count; i += SUBGROUP_SIZE) { + bitonic_cmp(i, i ^ j); + } + subgroupMemoryBarrierShared(); + } + } +} + +ScoredPoint get_nearest(uint i) { + return nearest[i]; +} + +void set_nearest(uint i, ScoredPoint scored_point) { + if (subgroupElect()) { + nearest[i] = scored_point; + } + subgroupMemoryBarrierShared(); +} + +#endif diff --git a/lib/segment/src/index/hnsw_index/gpu/shaders/run_get_patch.comp b/lib/segment/src/index/hnsw_index/gpu/shaders/run_get_patch.comp new file mode 100644 index 0000000000..488b08ee9a --- /dev/null +++ b/lib/segment/src/index/hnsw_index/gpu/shaders/run_get_patch.comp @@ -0,0 +1,139 @@ +#version 450 + +#include "extensions.comp" + +#include "common.comp" +layout(local_size_x = SUBGROUP_SIZE, local_size_y = 1, local_size_z = 1) in; + +#define VECTOR_STORAGE_LAYOUT_SET 1 +#include "vector_storage.comp" + +#define LINKS_LAYOUT_SET 2 +#include "links.comp" + +#define VISITED_FLAGS_LAYOUT_SET 3 +#include "visited_flags.comp" + +#include "search_context.comp" + +struct Request { + uint id; + uint entry; +}; + +layout(set = 0, binding = 0) buffer Requests { + Request data[]; +} requests; + +layout(set = 0, binding = 1) buffer Patches { + uint data[]; +} patches; + +layout(set = 0, binding = 2) buffer NewEntries { + uint data[]; +} new_entries; + +void clear_patches() { + uint subgroup_index = SUBGROUP_ID; + uint patch_size = LEVEL_M + 2; // id + links_count + links + uint all_patches_size = (LEVEL_M + 1) * patch_size; + uint patches_offset = subgroup_index * all_patches_size; + for (uint i = gl_SubgroupInvocationID; i < all_patches_size; i += SUBGROUP_SIZE) { + patches.data[patches_offset + i] = UINT_MAX; + } + subgroupMemoryBarrierShared(); +} + +void main() { + clear_patches(); + search_context_init(); + uint subgroup_index = SUBGROUP_ID; + Request request = requests.data[subgroup_index]; + + set_target(request.id); + float entry_score = similarity(request.entry); + uint search_count = search(request.id, ScoredPoint(request.entry, entry_score)); + if (search_count == UINT_MAX) { + // error while searching. propably candidates heap overflow + return; + } + + if (subgroupElect()) { + if (search_count > 0) { + new_entries.data[subgroup_index] = get_nearest(0).id; + } else { + new_entries.data[subgroup_index] = request.entry; + } + } + + uint count = run_heuristic(); + + uint patch_size = LEVEL_M + 2; // id + links_count + links + uint all_patches_size = (LEVEL_M + 1) * patch_size; + uint patches_offset = subgroup_index * all_patches_size; + + // write request.id patch + if (subgroupElect()) { + patches.data[patches_offset] = request.id; + patches.data[patches_offset + 1] = count; + for (uint i = 0; i < count; i++) { + patches.data[patches_offset + i + 2] = get_nearest(i).id; + } + } + subgroupMemoryBarrierShared(); + + for (uint i = 0; i < count; i++) { + uint other_id = patches.data[patches_offset + i + 2]; + uint other_links_count = LINKS_COUNT(other_id); + uint other_patches_offset = patches_offset + (i + 1) * patch_size; + if (other_links_count < LEVEL_M) { + if (subgroupElect()) { + patches.data[other_patches_offset] = other_id; + other_patches_offset++; + patches.data[other_patches_offset] = other_links_count + 1; + other_patches_offset++; + LINKS_FOREACH(other_id) { + patches.data[other_patches_offset] = LINKS_FOREACH_VALUE; + other_patches_offset++; + } + patches.data[other_patches_offset] = request.id; + } + } else { + set_target(other_id); + nearest_count = 0; + LINKS_FOREACH(other_id) { + set_nearest( + nearest_count, + ScoredPoint( + LINKS_FOREACH_VALUE, + similarity(LINKS_FOREACH_VALUE)) + ); + nearest_count++; + } + set_nearest( + nearest_count, + ScoredPoint( + request.id, + similarity(request.id)) + ); + nearest_count++; + + sort_nearest(); + uint other_new_links_count = run_heuristic(); + + if (subgroupElect()) { + patches.data[other_patches_offset] = other_id; + other_patches_offset++; + patches.data[other_patches_offset] = other_new_links_count; + other_patches_offset++; + } + for (uint j = 0; j < other_new_links_count; j++) { + if (subgroupElect()) { + uint l = get_nearest(j).id; + patches.data[other_patches_offset] = l; + other_patches_offset++; + } + } + } + } +} diff --git a/lib/segment/src/index/hnsw_index/gpu/shaders/run_greedy_search.comp b/lib/segment/src/index/hnsw_index/gpu/shaders/run_greedy_search.comp new file mode 100644 index 0000000000..d675d3fd0b --- /dev/null +++ b/lib/segment/src/index/hnsw_index/gpu/shaders/run_greedy_search.comp @@ -0,0 +1,44 @@ +#version 450 + +#include "extensions.comp" + +#include "common.comp" +layout(local_size_x = SUBGROUP_SIZE, local_size_y = 1, local_size_z = 1) in; + +#define VECTOR_STORAGE_LAYOUT_SET 1 +#include "vector_storage.comp" + +#define LINKS_LAYOUT_SET 2 +#include "links.comp" + +#define VISITED_FLAGS_LAYOUT_SET 3 +#include "visited_flags.comp" + +#include "search_context.comp" + +struct SearchRequest { + uint id; + uint entry; +}; + +layout(set = 0, binding = 0) buffer SearchRequests { + SearchRequest data[]; +} search_requests; + +layout(set = 0, binding = 1) buffer SearchResults { + uint data[]; +} search_results; + +void main() { + search_context_init(); + uint subgroup_index = SUBGROUP_ID; + SearchRequest search_request = search_requests.data[subgroup_index]; + + set_target(search_request.id); + float entry_score = similarity(search_request.entry); + ScoredPoint result = greedy_search(search_request.id, ScoredPoint(search_request.entry, entry_score)); + + if (subgroupElect()) { + search_results.data[subgroup_index] = result.id; + } +} diff --git a/lib/segment/src/index/hnsw_index/gpu/shaders/run_insert_vector.comp b/lib/segment/src/index/hnsw_index/gpu/shaders/run_insert_vector.comp new file mode 100644 index 0000000000..d8c04fdab3 --- /dev/null +++ b/lib/segment/src/index/hnsw_index/gpu/shaders/run_insert_vector.comp @@ -0,0 +1,129 @@ +#version 450 + +#include "extensions.comp" + +#include "common.comp" +layout(local_size_x = SUBGROUP_SIZE, local_size_y = 1, local_size_z = 1) in; + +#define VECTOR_STORAGE_LAYOUT_SET 1 +#include "vector_storage.comp" + +#define LINKS_LAYOUT_SET 2 +#include "links.comp" + +#define VISITED_FLAGS_LAYOUT_SET 3 +#include "visited_flags.comp" + +#include "search_context.comp" + +struct Request { + uint id; + uint entry; +}; + +layout(set = 0, binding = 0) buffer Requests { + Request data[]; +} requests; + +layout(set = 0, binding = 1) buffer NewEntries { + uint data[]; +} new_entries; + +layout(set = 0, binding = 2) buffer Atomics { + uint data[]; +} atomics; + +void main() { + search_context_init(); + uint subgroup_index = SUBGROUP_ID; + Request request = requests.data[subgroup_index]; + + set_target(request.id); + float entry_score = similarity(request.entry); +#ifdef NO_EQUIVALENCY + ScoredPoint new_entry = greedy_search(request.id, ScoredPoint(request.entry, entry_score)); + uint search_count = search(request.id, new_entry); +#else + uint search_count = search(request.id, ScoredPoint(request.entry, entry_score)); +#endif + if (search_count == UINT_MAX) { + // error while searching. propably candidates heap overflow + return; + } + + if (subgroupElect()) { + if (search_count > 0) { + new_entries.data[subgroup_index] = get_nearest(0).id; + } else { + new_entries.data[subgroup_index] = request.entry; + } + } + + uint count = run_heuristic(); + + // sync all threads to make sure that we don't do search while updating links + barrier(); + + // update request.id links + for (uint i = gl_SubgroupInvocationID; i < count; i += SUBGROUP_SIZE) { + GET_LINK(request.id, i) = get_nearest(i).id; + } + if (subgroupElect()) { + LINKS_SET_SIZE(request.id, count); + } + + subgroupMemoryBarrier(); + + uint generation = VISITED_GENERATION; + + for (uint i = 0; i < count; i++) { + uint other_id = GET_LINK(request.id, i); + + bool is_changed = false; + if (subgroupElect()) { + uint other_atomic = atomicExchange(atomics.data[other_id], generation); + is_changed = other_atomic == generation; + } + if (subgroupAny(is_changed)) { + // other_id is already being processed by another subgroup + continue; + } + + uint other_links_count = LINKS_COUNT(other_id); + if (other_links_count < LEVEL_M) { + if (subgroupElect()) { + GET_LINK(other_id, other_links_count) = request.id; + LINKS_SET_SIZE(other_id, other_links_count + 1); + } + } else { + set_target(other_id); + nearest_count = 0; + LINKS_FOREACH(other_id) { + set_nearest( + nearest_count, + ScoredPoint( + LINKS_FOREACH_VALUE, + similarity(LINKS_FOREACH_VALUE)) + ); + nearest_count++; + } + set_nearest( + nearest_count, + ScoredPoint( + request.id, + similarity(request.id)) + ); + nearest_count++; + + sort_nearest(); + uint other_new_links_count = run_heuristic(); + + for (uint j = gl_SubgroupInvocationID; j < other_new_links_count; j += SUBGROUP_SIZE) { + GET_LINK(other_id, j) = get_nearest(j).id; + } + if (subgroupElect()) { + LINKS_SET_SIZE(other_id, other_new_links_count); + } + } + } +} diff --git a/lib/segment/src/index/hnsw_index/gpu/shaders/search_context.comp b/lib/segment/src/index/hnsw_index/gpu/shaders/search_context.comp new file mode 100644 index 0000000000..0f14770646 --- /dev/null +++ b/lib/segment/src/index/hnsw_index/gpu/shaders/search_context.comp @@ -0,0 +1,140 @@ +#ifndef SEARCH_CONTEXT_HEADER +#define SEARCH_CONTEXT_HEADER + +#include "common.comp" +#include "vector_storage.comp" +#include "links.comp" +#include "nearest_heap.comp" +#include "candidates_heap.comp" +#include "visited_flags.comp" + +void process_candidate(ScoredPoint scored_point) { + bool not_presented = nearest_smaller.id != scored_point.id; + bool was_added = push_nearest(scored_point); + if (was_added && not_presented) { + push_candidate(scored_point); + subgroupMemoryBarrierShared(); + } +} + +void process_candidate(POINT_ID candidate_id) { + uint links_count = LINKS_COUNT(candidate_id); + uint subgroup_link = 0; + bool subgroup_is_visited = true; + uint invocation = gl_SubgroupInvocationID; + + for (uint i = 0; i < links_count; i++) { + if (i % SUBGROUP_SIZE == 0 && i + invocation < links_count) { + subgroup_link = GET_LINK(candidate_id, i + invocation); + subgroup_is_visited = is_visited(subgroup_link); + if (!subgroup_is_visited) { + mark_visited(subgroup_link); + } + } + + uint link = subgroupShuffle(subgroup_link, i % SUBGROUP_SIZE); + bool is_visited = subgroupShuffle(subgroup_is_visited, i % SUBGROUP_SIZE); + + if (!is_visited) { + float score = similarity(link); + ScoredPoint scored_point = ScoredPoint(link, score); + process_candidate(scored_point); + } + } + + // flush all changed visited flags + subgroupMemoryBarrierBuffer(); +} + +ScoredPoint greedy_search(POINT_ID point_id, ScoredPoint entry) { + bool changed = true; + ScoredPoint result = entry; + + uint subgroup_link = 0; + uint invocation = gl_SubgroupInvocationID; + + while (changed) { + changed = false; + + uint current = result.id; + uint links_count = LINKS_COUNT(current); + for (uint i = 0; i < links_count; i++) { + if (i % SUBGROUP_SIZE == 0 && i + invocation < links_count) { + subgroup_link = GET_LINK(current, i + invocation); + } + uint link = subgroupShuffle(subgroup_link, i % SUBGROUP_SIZE); + + float score = similarity(link); + if (score > result.score) { + changed = true; + result = ScoredPoint(link, score); + } + } + } + return result; +} + +uint search(POINT_ID point_id, ScoredPoint entry) { +#ifndef EXACT + entry = greedy_search(point_id, entry); +#endif + + if (subgroupElect()) { + mark_visited(point_id); + mark_visited(entry.id); + } + push_nearest(entry); + push_candidate(entry); + groupMemoryBarrier(); + + while (candidates_count > 0) { + ScoredPoint candidate = pop_candidate(); + if (candidate.score < nearest_smaller.score) { + continue; + } + process_candidate(candidate.id); + } + + process_candidate(point_id); + + sort_nearest(); + return nearest_count; +} + +uint run_heuristic() { + uint result_count = 0; + + for (uint i = 0; i < nearest_count; i++) { + if (result_count >= LEVEL_M) { + break; + } + + ScoredPoint current_closest = get_nearest(i); + set_target(current_closest.id); + bool is_good = true; + + for (uint j = 0; j < result_count; j++) { + ScoredPoint selected_point = get_nearest(j); + float dist_to_already_selected = similarity(selected_point.id); + if (dist_to_already_selected > current_closest.score) { + is_good = false; + break; + } + } + + if (is_good) { + set_nearest(result_count, current_closest); + result_count += 1; + } + } + + nearest_count = result_count; + return result_count; +} + +void search_context_init() { + nearest_heap_init(); + candidates_heap_init(); +} + +#endif diff --git a/lib/segment/src/index/hnsw_index/gpu/shaders/tests/test_candidates_heap.comp b/lib/segment/src/index/hnsw_index/gpu/shaders/tests/test_candidates_heap.comp new file mode 100644 index 0000000000..40b8a6b99a --- /dev/null +++ b/lib/segment/src/index/hnsw_index/gpu/shaders/tests/test_candidates_heap.comp @@ -0,0 +1,51 @@ +#version 450 + +#include "extensions.comp" + +#include "common.comp" +layout(local_size_x = SUBGROUP_SIZE, local_size_y = 1, local_size_z = 1) in; + +#include "candidates_heap.comp" + +layout(set = 0, binding = 0) +readonly uniform TestParams { + uint input_counts; +} test_params; + +layout(set = 0, binding = 1) buffer ScoredPointsInput { + ScoredPoint data[]; +} input_points; + +layout(set = 0, binding = 2) buffer ScoresOuput { + ScoredPoint data[]; +} scores_output; + +void main() { + candidates_heap_init(); + uint subgroup_index = SUBGROUP_ID; + + for (uint i = 0; i < test_params.input_counts; i++) { + uint point_index = subgroup_index * test_params.input_counts + i; + ScoredPoint point = input_points.data[point_index]; + push_candidate(point); + } + + while (candidates_count > 0) { + ScoredPoint candidate = pop_candidate(); + } + + for (uint i = 0; i < test_params.input_counts; i++) { + uint point_index = subgroup_index * test_params.input_counts + i; + ScoredPoint point = input_points.data[point_index]; + push_candidate(point); + } + + uint output_index = subgroup_index * test_params.input_counts; + while (candidates_count > 0) { + ScoredPoint candidate = pop_candidate(); + if (subgroupElect()) { + scores_output.data[output_index] = candidate; + } + output_index++; + } +} diff --git a/lib/segment/src/index/hnsw_index/gpu/shaders/tests/test_heuristic.comp b/lib/segment/src/index/hnsw_index/gpu/shaders/tests/test_heuristic.comp new file mode 100644 index 0000000000..0ca5974643 --- /dev/null +++ b/lib/segment/src/index/hnsw_index/gpu/shaders/tests/test_heuristic.comp @@ -0,0 +1,50 @@ +#version 450 + +#include "extensions.comp" + +#include "common.comp" +layout(local_size_x = SUBGROUP_SIZE, local_size_y = 1, local_size_z = 1) in; + +#define VECTOR_STORAGE_LAYOUT_SET 1 +#include "vector_storage.comp" + +#define LINKS_LAYOUT_SET 2 +#include "links.comp" + +#define VISITED_FLAGS_LAYOUT_SET 3 +#include "visited_flags.comp" + +#include "search_context.comp" + +struct SearchRequest { + uint id; + uint entry; +}; + +layout(set = 0, binding = 0) buffer SearchRequests { + SearchRequest data[]; +} search_requests; + +layout(set = 0, binding = 1) buffer SearchResults { + ScoredPoint data[]; +} search_results; + +void main() { + search_context_init(); + uint subgroup_index = SUBGROUP_ID; + SearchRequest search_request = search_requests.data[subgroup_index]; + + set_target(search_request.id); + float entry_score = similarity(search_request.entry); + search(search_request.id, ScoredPoint(search_request.entry, entry_score)); + uint count = run_heuristic(); + + barrier(); + + if (subgroupElect()) { + uint result_index_offset = NEAREST_HEAP_EF * subgroup_index; + for (uint i = 0; i < count; i++) { + search_results.data[result_index_offset + i] = get_nearest(i); + } + } +} diff --git a/lib/segment/src/index/hnsw_index/gpu/shaders/tests/test_hnsw_search.comp b/lib/segment/src/index/hnsw_index/gpu/shaders/tests/test_hnsw_search.comp new file mode 100644 index 0000000000..e67178ec18 --- /dev/null +++ b/lib/segment/src/index/hnsw_index/gpu/shaders/tests/test_hnsw_search.comp @@ -0,0 +1,53 @@ +#version 450 + +#include "extensions.comp" + +#include "common.comp" +layout(local_size_x = SUBGROUP_SIZE, local_size_y = 1, local_size_z = 1) in; + +#define VECTOR_STORAGE_LAYOUT_SET 1 +#include "vector_storage.comp" + +#define LINKS_LAYOUT_SET 2 +#include "links.comp" + +#define VISITED_FLAGS_LAYOUT_SET 3 +#include "visited_flags.comp" + +#include "search_context.comp" + +struct SearchRequest { + uint id; + uint entry; +}; + +layout(set = 0, binding = 0) buffer SearchRequests { + SearchRequest data[]; +} search_requests; + +layout(set = 0, binding = 1) buffer SearchResults { + ScoredPoint data[]; +} search_results; + +void main() { + search_context_init(); + uint subgroup_index = SUBGROUP_ID; + SearchRequest search_request = search_requests.data[subgroup_index]; + + set_target(search_request.id); + float entry_score = similarity(search_request.entry); + uint count = search(search_request.id, ScoredPoint(search_request.entry, entry_score)); + + barrier(); + + if (subgroupElect()) { + uint result_index_offset = NEAREST_HEAP_EF * subgroup_index; + uint i = 0; + for (; i < count; i++) { + search_results.data[result_index_offset + i] = get_nearest(i); + } + for (; i < NEAREST_HEAP_EF; i++) { + search_results.data[result_index_offset + i] = ScoredPoint(UINT_MAX, 0.0); + } + } +} diff --git a/lib/segment/src/index/hnsw_index/gpu/shaders/tests/test_nearest_heap.comp b/lib/segment/src/index/hnsw_index/gpu/shaders/tests/test_nearest_heap.comp new file mode 100644 index 0000000000..c96da1d1c8 --- /dev/null +++ b/lib/segment/src/index/hnsw_index/gpu/shaders/tests/test_nearest_heap.comp @@ -0,0 +1,58 @@ +#version 450 + +#include "extensions.comp" + +#include "common.comp" +layout(local_size_x = SUBGROUP_SIZE, local_size_y = 1, local_size_z = 1) in; + +#include "nearest_heap.comp" + +layout(set = 0, binding = 0) +readonly uniform TestParams { + uint input_counts; +} test_params; + +layout(set = 0, binding = 1) buffer ScoredPointsInput { + ScoredPoint data[]; +} input_points; + +layout(set = 0, binding = 2) buffer ScoresOuput { + float data[]; +} scores_output; + +layout(set = 0, binding = 3) buffer SortedOuput { + uint data[]; +} sorted_output; + +void main() { + nearest_heap_init(); + uint subgroup_index = SUBGROUP_ID; + + for (uint i = 0; i < test_params.input_counts; i++) { + uint point_index = subgroup_index * test_params.input_counts + i; + ScoredPoint point = input_points.data[point_index]; + push_nearest(point); + + groupMemoryBarrier(); + memoryBarrier(); + + float nearest_smaller_score_check = subgroupMin(nearest_smaller.score); + bool is_different = nearest_smaller_score_check != nearest_smaller.score; + if (subgroupAny(is_different)) { + if (subgroupElect()) { + scores_output.data[point_index] = positive_infinity; + } + } else if (subgroupElect()) { + scores_output.data[point_index] = nearest_smaller.score; + } + } + + sort_nearest(); + if (subgroupElect()) { + for (uint i = 0; i < NEAREST_HEAP_EF; i++) { + uint point_index = subgroup_index * NEAREST_HEAP_EF + i; + ScoredPoint point = get_nearest(i); + sorted_output.data[point_index] = point.id; + } + } +} diff --git a/lib/segment/src/index/hnsw_index/gpu/shaders/visited_flags.comp b/lib/segment/src/index/hnsw_index/gpu/shaders/visited_flags.comp new file mode 100644 index 0000000000..68cddaa544 --- /dev/null +++ b/lib/segment/src/index/hnsw_index/gpu/shaders/visited_flags.comp @@ -0,0 +1,37 @@ +#ifndef VISITED_FLAGS_HEADER +#define VISITED_FLAGS_HEADER + +#ifndef VISITED_FLAGS_LAYOUT_SET +#error "VISITED_FLAGS_LAYOUT_SET not defined" +#endif + +#include "common.comp" +#include "vector_storage.comp" + +layout(set = VISITED_FLAGS_LAYOUT_SET, binding = 0) +readonly uniform VisitedFlagsParams { + uint generation; +} visited_flags_params; + +layout(set = VISITED_FLAGS_LAYOUT_SET, binding = 1) +buffer VisitedFlagsBuffer { + uint8_t data[]; +} visited_flags; + +#define BYTE_MASK 0xFF +#define VISITED_GENERATION visited_flags_params.generation + +bool is_visited(POINT_ID point_id) { + uint subgroup_index = SUBGROUP_ID; + uint index = subgroup_index * VISITED_FLAGS_CAPACITY + point_id % VISITED_FLAGS_CAPACITY; + uint prev_generation = uint(visited_flags.data[index]); + return prev_generation == VISITED_GENERATION; +} + +void mark_visited(POINT_ID point_id) { + uint subgroup_index = SUBGROUP_ID; + uint index = subgroup_index * VISITED_FLAGS_CAPACITY + point_id % VISITED_FLAGS_CAPACITY; + visited_flags.data[index] = uint8_t(VISITED_GENERATION); +} + +#endif diff --git a/lib/segment/src/index/hnsw_index/graph_layers.rs b/lib/segment/src/index/hnsw_index/graph_layers.rs index 4d654379b3..21f8d8b4a6 100644 --- a/lib/segment/src/index/hnsw_index/graph_layers.rs +++ b/lib/segment/src/index/hnsw_index/graph_layers.rs @@ -137,6 +137,39 @@ pub trait GraphLayersBase { } current_point } + + fn search_entry_on_level( + &self, + entry_point: PointOffsetType, + level: usize, + points_scorer: &mut FilteredScorer, + ) -> ScoredPointOffset { + let limit = self.get_m(level); + let mut links: Vec = Vec::with_capacity(2 * self.get_m(0)); + let mut current_point = ScoredPointOffset { + idx: entry_point, + score: points_scorer.score_point(entry_point), + }; + + let mut changed = true; + while changed { + changed = false; + + links.clear(); + self.links_map(current_point.idx, level, |link| { + links.push(link); + }); + + let scores = points_scorer.score_points(&mut links, limit); + scores.iter().copied().for_each(|score_point| { + if score_point.score > current_point.score { + changed = true; + current_point = score_point; + } + }); + } + current_point + } } impl GraphLayersBase for GraphLayers { diff --git a/lib/segment/src/index/hnsw_index/graph_layers_builder.rs b/lib/segment/src/index/hnsw_index/graph_layers_builder.rs index 5e77595d18..8bc69430aa 100644 --- a/lib/segment/src/index/hnsw_index/graph_layers_builder.rs +++ b/lib/segment/src/index/hnsw_index/graph_layers_builder.rs @@ -6,6 +6,7 @@ use std::sync::atomic::AtomicUsize; use bitvec::prelude::BitVec; use common::fixed_length_priority_queue::FixedLengthPriorityQueue; use common::types::{PointOffsetType, ScoreType, ScoredPointOffset}; +use itertools::Itertools; use parking_lot::{Mutex, MutexGuard, RwLock}; use rand::distributions::Uniform; use rand::Rng; @@ -25,22 +26,22 @@ pub type LockedLayersContainer = Vec; /// Same as `GraphLayers`, but allows to build in parallel /// Convertible to `GraphLayers` pub struct GraphLayersBuilder { - max_level: AtomicUsize, - m: usize, - m0: usize, - ef_construct: usize, + pub max_level: AtomicUsize, + pub m: usize, + pub m0: usize, + pub ef_construct: usize, // Factor of level probability - level_factor: f64, + pub level_factor: f64, // Exclude points according to "not closer than base" heuristic? - use_heuristic: bool, - links_layers: Vec, - entry_points: Mutex, + pub use_heuristic: bool, + pub links_layers: Vec, + pub entry_points: Mutex, // Fields used on construction phase only - visited_pool: VisitedPool, + pub visited_pool: VisitedPool, // List of bool flags, which defines if the point is already indexed or not - ready_list: RwLock, + pub ready_list: RwLock, } impl GraphLayersBase for GraphLayersBuilder { @@ -206,7 +207,7 @@ impl GraphLayersBuilder { picked_level.round() as usize } - fn get_point_level(&self, point_id: PointOffsetType) -> usize { + pub fn get_point_level(&self, point_id: PointOffsetType) -> usize { self.links_layers[point_id as usize].len() - 1 } @@ -286,7 +287,7 @@ impl GraphLayersBuilder { } /// - fn select_candidates_with_heuristic( + pub(crate) fn select_candidates_with_heuristic( candidates: FixedLengthPriorityQueue, m: usize, score_internal: F, @@ -321,12 +322,14 @@ impl GraphLayersBuilder { // Let's find closest one on same level // greedy search for a single closest point - self.search_entry( + let new_entry = self.search_entry( entry_point.point_id, entry_point.level, level, &mut points_scorer, - ) + ); + + new_entry } else { ScoredPointOffset { idx: entry_point.point_id, @@ -477,6 +480,128 @@ impl GraphLayersBuilder { sum as f32 / count as f32 } } + + pub fn get_patch( + &self, + request: super::gpu::gpu_search_context::GpuRequest, + level: usize, + points_scorer: &mut FilteredScorer, + ) -> ( + Vec, + Vec, + ) { + let entry = request.entry; + let point_id = request.id; + let mut patches = vec![]; + + let level_m = self.get_m(level); + let mut visited_list = self.get_visited_list_from_pool(); + + visited_list.check_and_update_visited(entry); + + let level_entry = ScoredPointOffset { + idx: entry, + score: points_scorer.score_point(entry), + }; + let mut search_context = SearchContext::new(level_entry, self.ef_construct); + + self._search_on_level(&mut search_context, level, &mut visited_list, points_scorer); + + let new_entries = search_context + .nearest + .clone() + .into_vec() + .iter() + .map(|p| p.idx) + .collect_vec(); + + let scorer = |a, b| points_scorer.score_internal(a, b); + + let selected_nearest = { + let selected_nearest = + Self::select_candidates_with_heuristic(search_context.nearest, level_m, scorer); + patches.push(super::gpu::gpu_search_context::GpuGraphLinksPatch { + id: point_id, + links: selected_nearest.clone(), + }); + selected_nearest + }; + + for &other_point in &selected_nearest { + let mut other_point_links = self.links_layers[other_point as usize][level] + .read() + .clone(); + if other_point_links.len() < level_m { + // If linked point is lack of neighbours + other_point_links.push(point_id); + } else { + let mut candidates = BinaryHeap::with_capacity(level_m + 1); + candidates.push(ScoredPointOffset { + idx: point_id, + score: scorer(point_id, other_point), + }); + for other_point_link in other_point_links.iter().take(level_m).copied() { + candidates.push(ScoredPointOffset { + idx: other_point_link, + score: scorer(other_point_link, other_point), + }); + } + let selected_candidates = Self::select_candidate_with_heuristic_from_sorted( + candidates.into_sorted_vec().into_iter().rev(), + level_m, + scorer, + ); + other_point_links.clear(); // this do not free memory, which is good + for selected in selected_candidates.iter().copied() { + other_point_links.push(selected); + } + } + patches.push(super::gpu::gpu_search_context::GpuGraphLinksPatch { + id: other_point, + links: other_point_links, + }); + } + + (patches, new_entries) + } + + pub fn try_apply_patch( + &self, + level: usize, + patches: Vec, + ) -> bool { + let mut locks = Vec::with_capacity(patches.len()); + let mut links_to_apply = Vec::with_capacity(patches.len()); + for patch in patches { + if let Some(lock) = self.links_layers[patch.id as usize][level].try_write() { + //if lock.as_slice() != patch.old_links.as_slice() { + // return false; + //} + + locks.push(lock); + links_to_apply.push(patch.links); + } else { + return false; + } + } + + for (links, lock) in links_to_apply.into_iter().zip(locks.iter_mut()) { + **lock = links; + } + + true + } + + pub fn apply_patch( + &self, + level: usize, + patches: Vec, + ) { + for patch in patches { + let mut lock = self.links_layers[patch.id as usize][level].write(); + *lock = patch.links; + } + } } #[cfg(test)] diff --git a/lib/segment/src/index/hnsw_index/hnsw.rs b/lib/segment/src/index/hnsw_index/hnsw.rs index 227edda96a..6a8e841d1e 100644 --- a/lib/segment/src/index/hnsw_index/hnsw.rs +++ b/lib/segment/src/index/hnsw_index/hnsw.rs @@ -20,6 +20,7 @@ use rand::thread_rng; use rayon::prelude::*; use rayon::ThreadPool; +use super::gpu::get_gpu_min_points_count; use super::graph_links::{GraphLinks, GraphLinksMmap}; use crate::common::operation_error::{check_process_stopped, OperationError, OperationResult}; use crate::common::operation_time_statistics::{ @@ -31,6 +32,10 @@ use crate::data_types::vectors::{QueryVector, VectorInternal, VectorRef}; use crate::id_tracker::IdTrackerSS; use crate::index::hnsw_index::build_condition_checker::BuildConditionChecker; use crate::index::hnsw_index::config::HnswGraphConfig; +use crate::index::hnsw_index::gpu::gpu_graph_builder::build_hnsw_on_gpu; +use crate::index::hnsw_index::gpu::{ + get_gpu_force_half_precision, get_gpu_indexing, get_gpu_max_groups, +}; use crate::index::hnsw_index::graph_layers::GraphLayers; use crate::index::hnsw_index::graph_layers_builder::GraphLayersBuilder; use crate::index::hnsw_index::point_scorer::FilteredScorer; @@ -39,6 +44,7 @@ use crate::index::sample_estimation::sample_check_cardinality; use crate::index::struct_payload_index::StructPayloadIndex; use crate::index::visited_pool::{VisitedListHandle, VisitedPool}; use crate::index::{PayloadIndex, VectorIndex}; +use crate::payload_storage::FilterContext; use crate::telemetry::VectorIndexSearchesTelemetry; use crate::types::Condition::Field; use crate::types::{ @@ -91,6 +97,7 @@ pub struct HnswIndexOpenArgs<'a> { pub payload_index: Arc>, pub hnsw_config: HnswConfig, pub permit: Option>, + pub gpu_device: Option>, pub stopped: &'a AtomicBool, } @@ -104,6 +111,7 @@ impl HNSWIndex { payload_index, hnsw_config, permit, + gpu_device, stopped, } = args; @@ -162,6 +170,7 @@ impl HNSWIndex { &payload_index.borrow(), &hnsw_config, num_cpus, + gpu_device, stopped, )?; @@ -209,6 +218,7 @@ impl HNSWIndex { payload_index: &StructPayloadIndex, hnsw_config: &HnswConfig, num_cpus: usize, + gpu_device: Option>, stopped: &AtomicBool, ) -> OperationResult<(HnswGraphConfig, GraphLayers)> { let total_vector_count = vector_storage.total_vector_count(); @@ -237,20 +247,22 @@ impl HNSWIndex { let mut rng = thread_rng(); let deleted_bitslice = vector_storage.deleted_vector_bitslice(); + // TODO(gpu): log device name debug!("building HNSW for {total_vector_count} vectors with {num_cpus} CPUs"); + let num_entries = std::cmp::max( + 1, + total_vector_count + .checked_div(full_scan_threshold) + .unwrap_or(0) + * 10, + ); let mut graph_layers_builder = GraphLayersBuilder::new( total_vector_count, config.m, config.m0, config.ef_construct, - std::cmp::max( - 1, - total_vector_count - .checked_div(full_scan_threshold) - .unwrap_or(0) - * 10, - ), + num_entries, HNSW_USE_HEURISTIC, ); @@ -289,49 +301,108 @@ impl HNSWIndex { let mut indexed_vectors = 0; if config.m > 0 { - let mut ids_iterator = id_tracker.iter_ids_excluding(deleted_bitslice); + let timer = std::time::Instant::now(); - let first_few_ids: Vec<_> = ids_iterator - .by_ref() - .take(SINGLE_THREADED_HNSW_BUILD_THRESHOLD) - .collect(); - let ids: Vec<_> = ids_iterator.collect(); + let use_gpu = total_vector_count > get_gpu_min_points_count() && get_gpu_indexing(); + let gpu_constructed_graph = if let Some(gpu_device) = &gpu_device { + if use_gpu { + let ids: Vec<_> = id_tracker.iter_ids_excluding(deleted_bitslice).collect(); + indexed_vectors = ids.len(); - indexed_vectors = ids.len() + first_few_ids.len(); + let points_scorer_builder = |vector_id| { + let vector = vector_storage.get_vector(vector_id); + let vector = vector.as_vec_ref().into(); + let raw_scorer = if let Some(quantized_storage) = quantized_vectors.as_ref() + { + quantized_storage.raw_scorer( + vector, + id_tracker.deleted_point_bitslice(), + vector_storage.deleted_vector_bitslice(), + stopped, + ) + } else { + new_raw_scorer( + vector, + &vector_storage, + id_tracker.deleted_point_bitslice(), + ) + }?; + Ok((raw_scorer, None)) + }; - let insert_point = |vector_id| { - check_process_stopped(stopped)?; - let vector = vector_storage.get_vector(vector_id); - let vector = vector.as_vec_ref().into(); - let raw_scorer = if let Some(quantized_storage) = quantized_vectors.as_ref() { - quantized_storage.raw_scorer( - vector, - id_tracker.deleted_point_bitslice(), - vector_storage.deleted_vector_bitslice(), + Some(build_hnsw_on_gpu( + gpu_device.clone(), + &pool, + &graph_layers_builder, + get_gpu_max_groups(), + &vector_storage, + quantized_vectors.as_ref(), + num_entries, + get_gpu_force_half_precision(), + SINGLE_THREADED_HNSW_BUILD_THRESHOLD, + false, + ids, + points_scorer_builder, stopped, - ) + )) } else { - new_raw_scorer(vector, vector_storage, id_tracker.deleted_point_bitslice()) - }?; - let points_scorer = FilteredScorer::new(raw_scorer.as_ref(), None); - - graph_layers_builder.link_new_point(vector_id, points_scorer); - - // Ignore hardware counter, for internal operations - raw_scorer.take_hardware_counter().discard_results(); - - Ok::<_, OperationError>(()) + None + } + } else { + None }; - for vector_id in first_few_ids { - insert_point(vector_id)?; + if let Some(Err(gpu_error)) = &gpu_constructed_graph { + log::warn!("Failed to build HNSW on GPU: {gpu_error}. Falling back to CPU."); } - if !ids.is_empty() { - pool.install(|| ids.into_par_iter().try_for_each(insert_point))?; + if let Some(Ok(gpu_graph)) = gpu_constructed_graph { + graph_layers_builder = gpu_graph; + } else { + let mut ids_iterator = id_tracker.iter_ids_excluding(deleted_bitslice); + + let first_few_ids: Vec<_> = ids_iterator + .by_ref() + .take(SINGLE_THREADED_HNSW_BUILD_THRESHOLD) + .collect(); + let ids: Vec<_> = ids_iterator.collect(); + + indexed_vectors = ids.len() + first_few_ids.len(); + + let insert_point = |vector_id| { + check_process_stopped(stopped)?; + let vector = vector_storage.get_vector(vector_id); + let vector = vector.as_vec_ref().into(); + let raw_scorer = if let Some(quantized_storage) = quantized_vectors.as_ref() { + quantized_storage.raw_scorer( + vector, + id_tracker.deleted_point_bitslice(), + vector_storage.deleted_vector_bitslice(), + stopped, + ) + } else { + new_raw_scorer(vector, &vector_storage, id_tracker.deleted_point_bitslice()) + }?; + let points_scorer = FilteredScorer::new(raw_scorer.as_ref(), None); + + graph_layers_builder.link_new_point(vector_id, points_scorer); + + // Ignore hardware counter, for internal operations + raw_scorer.take_hardware_counter().discard_results(); + + Ok::<_, OperationError>(()) + }; + + for vector_id in first_few_ids { + insert_point(vector_id)?; + } + + if !ids.is_empty() { + pool.install(|| ids.into_par_iter().try_for_each(insert_point))?; + } } - debug!("finish main graph"); + debug!("finish main graph in time {:?}", timer.elapsed()); } else { debug!("skip building main HNSW graph"); } @@ -390,6 +461,7 @@ impl HNSWIndex { quantized_vectors, payload_index, &pool, + gpu_device.clone(), stopped, &mut additional_graph, payload_block.condition, @@ -437,6 +509,7 @@ impl HNSWIndex { quantized_vectors: &Option, payload_index: &StructPayloadIndex, pool: &ThreadPool, + gpu_device: Option>, stopped: &AtomicBool, graph_layers_builder: &mut GraphLayersBuilder, condition: FieldCondition, @@ -468,6 +541,63 @@ impl HNSWIndex { } } + let use_gpu = points_to_index.len() > get_gpu_min_points_count() && get_gpu_indexing(); + let gpu_constructed_graph = if let Some(gpu_device) = gpu_device { + if use_gpu { + let points_scorer_builder = |block_point_id| -> OperationResult<_> { + let vector = vector_storage.get_vector(block_point_id); + let vector = vector.as_vec_ref().into(); + let raw_scorer = match quantized_vectors.as_ref() { + Some(quantized_storage) => quantized_storage.raw_scorer( + vector, + id_tracker.deleted_point_bitslice(), + deleted_bitslice, + stopped, + ), + None => new_raw_scorer( + vector, + vector_storage, + id_tracker.deleted_point_bitslice(), + ), + }?; + let block_condition_checker: Box = + Box::new(BuildConditionChecker { + filter_list: block_filter_list, + current_point: block_point_id, + }); + Ok((raw_scorer, Some(block_condition_checker))) + }; + Some(build_hnsw_on_gpu( + gpu_device.clone(), + &pool, + &graph_layers_builder, + get_gpu_max_groups(), + &vector_storage, + quantized_vectors.as_ref(), + 1, + get_gpu_force_half_precision(), + SINGLE_THREADED_HNSW_BUILD_THRESHOLD, + false, + points_to_index.clone(), + points_scorer_builder, + stopped, + )) + } else { + None + } + } else { + None + }; + + if let Some(Err(gpu_error)) = &gpu_constructed_graph { + log::warn!("Failed to build HNSW on GPU: {gpu_error}. Falling back to CPU."); + } + + if let Some(Ok(gpu_graph)) = gpu_constructed_graph { + *graph_layers_builder = gpu_graph; + return Ok(()); + } + let insert_points = |block_point_id| { check_process_stopped(stopped)?; diff --git a/lib/segment/src/index/hnsw_index/tests/test_graph_connectivity.rs b/lib/segment/src/index/hnsw_index/tests/test_graph_connectivity.rs index 479da784e3..49c064ace4 100644 --- a/lib/segment/src/index/hnsw_index/tests/test_graph_connectivity.rs +++ b/lib/segment/src/index/hnsw_index/tests/test_graph_connectivity.rs @@ -86,6 +86,7 @@ fn test_graph_connectivity() { payload_index: payload_index_ptr, hnsw_config, permit: Some(permit), + gpu_device: None, stopped: &stopped, }) .unwrap(); diff --git a/lib/segment/src/segment_constructor/segment_builder.rs b/lib/segment/src/segment_constructor/segment_builder.rs index 7233c667e8..f8b29a15e9 100644 --- a/lib/segment/src/segment_constructor/segment_builder.rs +++ b/lib/segment/src/segment_constructor/segment_builder.rs @@ -427,6 +427,23 @@ impl SegmentBuilder { id_tracker.versions_flusher()()?; let id_tracker_arc = Arc::new(AtomicRefCell::new(id_tracker)); + let gpu_device = if crate::index::hnsw_index::gpu::get_gpu_indexing() { + crate::index::hnsw_index::gpu::GPU_DEVICES_MANAGER + .as_ref() + .map(|devices_manager| devices_manager.lock_device()) + .ok() + .flatten() + } else { + None + }; + /* + if let Some(_gpu_device) = &gpu_device { + if permit.num_cpus > 1 { + permit.release_count(permit.num_cpus - 1); + } + } + */ + // Arc permit to share it with each vector store let permit = Arc::new(permit); @@ -500,6 +517,9 @@ impl SegmentBuilder { payload_index_arc.clone(), quantized_vectors_arc, Some(permit.clone()), + gpu_device + .as_ref() + .map(|gpu_device| gpu_device.locked_device.clone()), stopped, )?; } diff --git a/lib/segment/src/segment_constructor/segment_constructor_base.rs b/lib/segment/src/segment_constructor/segment_constructor_base.rs index 35cc13cdca..8e854fbe09 100644 --- a/lib/segment/src/segment_constructor/segment_constructor_base.rs +++ b/lib/segment/src/segment_constructor/segment_constructor_base.rs @@ -364,6 +364,7 @@ pub(crate) fn create_vector_index( payload_index: Arc>, quantized_vectors: Arc>>, permit: Option>, + gpu_device: Option>, stopped: &AtomicBool, ) -> OperationResult { let vector_index = match &vector_config.index { @@ -379,6 +380,7 @@ pub(crate) fn create_vector_index( payload_index, hnsw_config: vector_hnsw_config.clone(), permit, + gpu_device, stopped, }; if vector_hnsw_config.on_disk == Some(true) { @@ -559,6 +561,7 @@ fn create_segment( payload_index.clone(), quantized_vectors.clone(), None, + None, stopped, )?); diff --git a/lib/segment/src/vector_storage/quantized/quantized_vectors.rs b/lib/segment/src/vector_storage/quantized/quantized_vectors.rs index 4c47f42f54..452c0ab57c 100644 --- a/lib/segment/src/vector_storage/quantized/quantized_vectors.rs +++ b/lib/segment/src/vector_storage/quantized/quantized_vectors.rs @@ -111,11 +111,11 @@ impl fmt::Debug for QuantizedVectorStorage { #[derive(Debug)] pub struct QuantizedVectors { - storage_impl: QuantizedVectorStorage, - config: QuantizedVectorsConfig, - path: PathBuf, - distance: Distance, - datatype: VectorStorageDatatype, + pub storage_impl: QuantizedVectorStorage, + pub config: QuantizedVectorsConfig, + pub path: PathBuf, + pub distance: Distance, + pub datatype: VectorStorageDatatype, } impl QuantizedVectors { diff --git a/lib/segment/tests/integration/batch_search_test.rs b/lib/segment/tests/integration/batch_search_test.rs index f82ed58521..c627c9191b 100644 --- a/lib/segment/tests/integration/batch_search_test.rs +++ b/lib/segment/tests/integration/batch_search_test.rs @@ -170,6 +170,7 @@ fn test_batch_and_single_request_equivalency() { payload_index: payload_index_ptr, hnsw_config, permit: Some(permit), + gpu_device: None, stopped: &stopped, }) .unwrap(); diff --git a/lib/segment/tests/integration/byte_storage_hnsw_test.rs b/lib/segment/tests/integration/byte_storage_hnsw_test.rs index 81ebcbb58c..79296bac09 100644 --- a/lib/segment/tests/integration/byte_storage_hnsw_test.rs +++ b/lib/segment/tests/integration/byte_storage_hnsw_test.rs @@ -223,6 +223,7 @@ fn test_byte_storage_hnsw( payload_index: segment_byte.payload_index.clone(), hnsw_config, permit: Some(permit), + gpu_device: None, stopped: &stopped, }) .unwrap(); diff --git a/lib/segment/tests/integration/byte_storage_quantization_test.rs b/lib/segment/tests/integration/byte_storage_quantization_test.rs index 2fc1ef2c66..5ccb04fdf2 100644 --- a/lib/segment/tests/integration/byte_storage_quantization_test.rs +++ b/lib/segment/tests/integration/byte_storage_quantization_test.rs @@ -350,6 +350,7 @@ fn test_byte_storage_binary_quantization_hnsw( payload_index: segment_byte.payload_index.clone(), hnsw_config, permit: Some(permit), + gpu_device: None, stopped: &stopped, }) .unwrap(); diff --git a/lib/segment/tests/integration/exact_search_test.rs b/lib/segment/tests/integration/exact_search_test.rs index 1bd987d7b3..4f8c181e8d 100644 --- a/lib/segment/tests/integration/exact_search_test.rs +++ b/lib/segment/tests/integration/exact_search_test.rs @@ -140,6 +140,7 @@ fn exact_search_test() { payload_index: payload_index_ptr.clone(), hnsw_config, permit: Some(permit), + gpu_device: None, stopped: &stopped, }) .unwrap(); diff --git a/lib/segment/tests/integration/filtrable_hnsw_test.rs b/lib/segment/tests/integration/filtrable_hnsw_test.rs index 60695df6ed..c3a693d1f0 100644 --- a/lib/segment/tests/integration/filtrable_hnsw_test.rs +++ b/lib/segment/tests/integration/filtrable_hnsw_test.rs @@ -200,6 +200,7 @@ fn _test_filterable_hnsw( payload_index: payload_index_ptr.clone(), hnsw_config, permit: Some(permit), + gpu_device: None, stopped: &stopped, }) .unwrap(); diff --git a/lib/segment/tests/integration/hnsw_discover_test.rs b/lib/segment/tests/integration/hnsw_discover_test.rs index 91eb7cf545..a261d1c127 100644 --- a/lib/segment/tests/integration/hnsw_discover_test.rs +++ b/lib/segment/tests/integration/hnsw_discover_test.rs @@ -118,6 +118,7 @@ fn hnsw_discover_precision() { payload_index: payload_index_ptr, hnsw_config, permit: Some(permit), + gpu_device: None, stopped: &stopped, }) .unwrap(); @@ -241,6 +242,7 @@ fn filtered_hnsw_discover_precision() { payload_index: payload_index_ptr, hnsw_config, permit: Some(permit), + gpu_device: None, stopped: &stopped, }) .unwrap(); diff --git a/lib/segment/tests/integration/hnsw_quantized_search_test.rs b/lib/segment/tests/integration/hnsw_quantized_search_test.rs index 7ec28bc65d..690864251b 100644 --- a/lib/segment/tests/integration/hnsw_quantized_search_test.rs +++ b/lib/segment/tests/integration/hnsw_quantized_search_test.rs @@ -141,6 +141,7 @@ fn hnsw_quantized_search_test( payload_index: segment.payload_index.clone(), hnsw_config, permit: Some(permit), + gpu_device: None, stopped: &stopped, }) .unwrap(); diff --git a/lib/segment/tests/integration/multivector_filtrable_hnsw_test.rs b/lib/segment/tests/integration/multivector_filtrable_hnsw_test.rs index affa192d7e..729dcd5494 100644 --- a/lib/segment/tests/integration/multivector_filtrable_hnsw_test.rs +++ b/lib/segment/tests/integration/multivector_filtrable_hnsw_test.rs @@ -192,6 +192,7 @@ fn test_multi_filterable_hnsw( payload_index: payload_index_ptr, hnsw_config, permit: Some(permit), + gpu_device: None, stopped: &stopped, }) .unwrap(); diff --git a/lib/segment/tests/integration/multivector_hnsw_test.rs b/lib/segment/tests/integration/multivector_hnsw_test.rs index da45a13b17..1980ade427 100644 --- a/lib/segment/tests/integration/multivector_hnsw_test.rs +++ b/lib/segment/tests/integration/multivector_hnsw_test.rs @@ -150,6 +150,7 @@ fn test_single_multi_and_dense_hnsw_equivalency() { payload_index: segment.payload_index.clone(), hnsw_config: hnsw_config.clone(), permit: Some(permit.clone()), + gpu_device: None, stopped: &stopped, }) .unwrap(); @@ -164,6 +165,7 @@ fn test_single_multi_and_dense_hnsw_equivalency() { payload_index: segment.payload_index.clone(), hnsw_config, permit: Some(permit), + gpu_device: None, stopped: &stopped, }) .unwrap(); diff --git a/lib/segment/tests/integration/multivector_quantization_test.rs b/lib/segment/tests/integration/multivector_quantization_test.rs index e6f0c7a786..5740565992 100644 --- a/lib/segment/tests/integration/multivector_quantization_test.rs +++ b/lib/segment/tests/integration/multivector_quantization_test.rs @@ -320,6 +320,7 @@ fn test_multivector_quantization_hnsw( payload_index: segment.payload_index.clone(), hnsw_config, permit: Some(permit), + gpu_device: None, stopped: &stopped, }) .unwrap(); diff --git a/nvidia_icd.json b/nvidia_icd.json new file mode 100644 index 0000000000..d421809814 --- /dev/null +++ b/nvidia_icd.json @@ -0,0 +1,7 @@ +{ + "file_format_version" : "1.0.0", + "ICD": { + "library_path": "libGLX_nvidia.so.0", + "api_version" : "1.3" + } +} diff --git a/run.bat b/run.bat new file mode 100644 index 0000000000..315d1cd17a --- /dev/null +++ b/run.bat @@ -0,0 +1,7 @@ +SET QDRANT__GPU__indexing=1 +SET QDRANT__GPU__max_groups=512 +SET QDRANT__GPU__force_half_precision=0 + +SET QDRANT__storage__optimizers__max_optimization_threads=1 + +cargo run --release diff --git a/run.sh b/run.sh new file mode 100755 index 0000000000..4c039f9373 --- /dev/null +++ b/run.sh @@ -0,0 +1,9 @@ +export QDRANT__GPU__indexing=1 +export QDRANT__GPU__max_groups=512 +export QDRANT__GPU__force_half_precision=0 + +export QDRANT__GPU__device_filter=LLVM + +export QDRANT__storage__optimizers__max_optimization_threads=1 + +cargo run --release diff --git a/src/main.rs b/src/main.rs index fcc44a31b3..91308ee8d2 100644 --- a/src/main.rs +++ b/src/main.rs @@ -168,6 +168,24 @@ fn main() -> anyhow::Result<()> { .async_scorer .unwrap_or_default(), ); + if let Some(settings_gpu) = &settings.gpu { + segment::index::hnsw_index::gpu::set_gpu_indexing(settings_gpu.indexing); + segment::index::hnsw_index::gpu::set_gpu_force_half_precision( + settings_gpu.force_half_precision, + ); + segment::index::hnsw_index::gpu::set_gpu_max_groups(settings_gpu.max_groups); + segment::index::hnsw_index::gpu::set_device_index(settings_gpu.device_index); + segment::index::hnsw_index::gpu::set_devices_count(settings_gpu.devices_count); + segment::index::hnsw_index::gpu::set_device_filter(&settings_gpu.device_filter); + segment::index::hnsw_index::gpu::set_wait_free(settings_gpu.wait_free.unwrap_or(true)); + segment::index::hnsw_index::gpu::set_gpu_parallel_indexes( + settings_gpu.parallel_indexes.unwrap_or(1), + ); + // initialize GPU devices + let _ = segment::index::hnsw_index::gpu::GPU_DEVICES_MANAGER + .as_ref() + .map(|devices_manager| devices_manager.lock_device()); + } welcome(&settings); diff --git a/src/settings.rs b/src/settings.rs index 9d9b930f0e..41f8be59b2 100644 --- a/src/settings.rs +++ b/src/settings.rs @@ -141,6 +141,30 @@ pub struct TlsConfig { pub cert_ttl: Option, } +#[derive(Clone, Debug, Deserialize, Validate)] +#[allow(dead_code)] +pub struct GpuConfig { + #[serde(default)] + pub indexing: bool, + #[serde(default)] + pub force_half_precision: bool, + #[serde(default)] + #[validate(range(min = 1))] + pub max_groups: Option, + #[serde(default)] + pub device_index: Option, + #[serde(default)] + pub devices_count: Option, + #[serde(default)] + pub device_filter: String, + #[serde(default)] + pub min_points: Option, + #[serde(default)] + pub wait_free: Option, + #[serde(default)] + pub parallel_indexes: Option, +} + #[derive(Debug, Deserialize, Clone, Validate)] #[allow(dead_code)] // necessary because some field are only used in main.rs pub struct Settings { @@ -168,6 +192,9 @@ pub struct Settings { pub load_errors: Vec, #[serde(default)] pub inference: Option, + #[serde(default)] + #[validate(nested)] + pub gpu: Option, } impl Settings {