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 {