diff --git a/cmake/libs/libdiskann.cmake b/cmake/libs/libdiskann.cmake index 9c6ce5eb3..ca3b64f46 100644 --- a/cmake/libs/libdiskann.cmake +++ b/cmake/libs/libdiskann.cmake @@ -1,6 +1,6 @@ add_definitions(-DKNOWHERE_WITH_DISKANN) -find_package(Boost REQUIRED COMPONENTS program_options) -include_directories(${Boost_INCLUDE_DIR}) +find_package(Boost REQUIRED COMPONENTS program_options CONFIG) +include_directories(${Boost_INCLUDE_DIRS}) find_package(aio REQUIRED) include_directories(${AIO_INCLUDE}) find_package(fmt REQUIRED) diff --git a/cmake/libs/libfaiss.cmake b/cmake/libs/libfaiss.cmake index 056b530e6..b20dbe261 100644 --- a/cmake/libs/libfaiss.cmake +++ b/cmake/libs/libfaiss.cmake @@ -593,3 +593,41 @@ if(__PPC64) knowhere_utils) target_compile_definitions(faiss PRIVATE FINTEGER=int) endif() + +# GPU HNSW CUDA sources — compiled when WITH_CUVS is enabled. +# +# This `faiss_gpu_hnsw` OBJECT library is the ONLY place the faiss GPU sources +# are compiled in the knowhere build: +# * The main `faiss` STATIC target above is built from FAISS_SRCS, which globs +# only thirdparty/faiss/faiss/{*.cpp,impl,utils,...} — it deliberately does +# NOT include thirdparty/faiss/faiss/gpu/*, so none of the files below are +# also compiled into `faiss`. +# * Upstream faiss ships its own GPU build file +# (thirdparty/faiss/faiss/gpu/CMakeLists.txt → the `faiss_gpu_objs` target +# over FAISS_GPU_SRC). Knowhere assembles faiss manually via this .cmake and +# never `add_subdirectory()`s the vendored faiss tree, so that file (and +# `faiss_gpu_objs`) is never part of the knowhere build graph. +# Therefore the two GPU build paths are mutually exclusive here — there is no +# duplicate compilation or ODR hazard. Keep it that way: add GPU sources to the +# list below, NOT by pulling in the vendored gpu/CMakeLists.txt. +if(WITH_CUVS) + set(FAISS_GPU_HNSW_SRCS + thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu + thirdparty/faiss/faiss/gpu/GpuIndex.cu + thirdparty/faiss/faiss/gpu/GpuResources.cpp + thirdparty/faiss/faiss/gpu/StandardGpuResources.cpp + thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu + thirdparty/faiss/faiss/gpu/impl/IndexUtils.cu + thirdparty/faiss/faiss/gpu/utils/DeviceUtils.cu + thirdparty/faiss/faiss/gpu/utils/StackDeviceMemory.cpp + thirdparty/faiss/faiss/gpu/utils/Timer.cpp + ) + add_library(faiss_gpu_hnsw OBJECT ${FAISS_GPU_HNSW_SRCS}) + target_include_directories(faiss_gpu_hnsw PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/thirdparty/faiss + ${Boost_INCLUDE_DIRS} + ) + target_compile_definitions(faiss_gpu_hnsw PRIVATE FINTEGER=int) + target_link_libraries(faiss_gpu_hnsw PRIVATE CUDA::cudart) + target_link_libraries(faiss PUBLIC faiss_gpu_hnsw) +endif() diff --git a/include/knowhere/bitsetview.h b/include/knowhere/bitsetview.h index ee6db3a47..84eddcb08 100644 --- a/include/knowhere/bitsetview.h +++ b/include/knowhere/bitsetview.h @@ -69,6 +69,11 @@ class BitsetView { return out_ids_ != nullptr; } + size_t + id_offset() const { + return id_offset_; + } + void set_out_ids(const uint32_t* out_ids, size_t num_internal_ids, std::optional num_filtered_out_ids = std::nullopt) { diff --git a/include/knowhere/comp/index_param.h b/include/knowhere/comp/index_param.h index 9066b3641..93cb96d01 100644 --- a/include/knowhere/comp/index_param.h +++ b/include/knowhere/comp/index_param.h @@ -53,6 +53,8 @@ constexpr const char* INDEX_GPU_BRUTEFORCE = "GPU_BRUTE_FORCE"; constexpr const char* INDEX_GPU_IVFFLAT = "GPU_IVF_FLAT"; constexpr const char* INDEX_GPU_IVFPQ = "GPU_IVF_PQ"; constexpr const char* INDEX_GPU_CAGRA = "GPU_CAGRA"; +constexpr const char* INDEX_GPU_HNSW = "GPU_HNSW"; +constexpr const char* INDEX_GPU_HNSW_SQ = "GPU_HNSW_SQ"; constexpr const char* INDEX_HNSW = "HNSW"; constexpr const char* INDEX_HNSW_SQ = "HNSW_SQ"; diff --git a/include/knowhere/index/index_static.h b/include/knowhere/index/index_static.h index f6279f613..001e7796f 100644 --- a/include/knowhere/index/index_static.h +++ b/include/knowhere/index/index_static.h @@ -19,8 +19,16 @@ namespace knowhere { struct Resource { - uint64_t memoryCost; // in bytes + uint64_t memoryCost; // retained host memory after load completes, in bytes uint64_t diskCost; // in bytes + // Peak transient host memory required *during* load, in bytes. Defaults to 0 + // so existing indexes are unaffected: consumers should fall back to their + // memoryCost-based heuristic when this is 0. Indexes whose peak load + // footprint differs from the retained footprint (e.g. GPU_HNSW frees its CPU + // copy after uploading to VRAM, so memoryCost≈0 but the peak covers the + // download buffer + deserialized CPU index + decode/graph staging) set this + // to the peak so the loader reserves enough host RAM and does not OOM. + uint64_t maxMemoryCost = 0; }; #define DEFINE_HAS_STATIC_FUNC(func_name) \ diff --git a/include/knowhere/index/index_table.h b/include/knowhere/index/index_table.h index d807305e9..0b300ef66 100644 --- a/include/knowhere/index/index_table.h +++ b/include/knowhere/index/index_table.h @@ -82,6 +82,14 @@ static std::set> legal_knowhere_index = { {IndexEnum::INDEX_GPU_CAGRA, VecType::VECTOR_FLOAT16}, {IndexEnum::INDEX_GPU_CAGRA, VecType::VECTOR_INT8}, {IndexEnum::INDEX_GPU_CAGRA, VecType::VECTOR_BINARY}, + {IndexEnum::INDEX_GPU_HNSW, VecType::VECTOR_FLOAT}, + {IndexEnum::INDEX_GPU_HNSW, VecType::VECTOR_FLOAT16}, + {IndexEnum::INDEX_GPU_HNSW, VecType::VECTOR_BFLOAT16}, + {IndexEnum::INDEX_GPU_HNSW, VecType::VECTOR_INT8}, + {IndexEnum::INDEX_GPU_HNSW_SQ, VecType::VECTOR_FLOAT}, + {IndexEnum::INDEX_GPU_HNSW_SQ, VecType::VECTOR_FLOAT16}, + {IndexEnum::INDEX_GPU_HNSW_SQ, VecType::VECTOR_BFLOAT16}, + {IndexEnum::INDEX_GPU_HNSW_SQ, VecType::VECTOR_INT8}, // hnsw {IndexEnum::INDEX_HNSW, VecType::VECTOR_FLOAT}, diff --git a/src/index/hnsw/faiss_hnsw.cc b/src/index/hnsw/faiss_hnsw.cc index 63101aa21..7617c5817 100644 --- a/src/index/hnsw/faiss_hnsw.cc +++ b/src/index/hnsw/faiss_hnsw.cc @@ -21,12 +21,14 @@ #include #include +#include #include #include #include #include #include #include +#include #include #include #include @@ -3286,4 +3288,598 @@ KNOWHERE_SIMPLE_REGISTER_DENSE_FLOAT_ALL_GLOBAL(HNSW_PRQ, BaseFaissRegularIndexH KNOWHERE_SIMPLE_REGISTER_DENSE_INT_GLOBAL(HNSW_PRQ, BaseFaissRegularIndexHNSWPRQNodeTemplate, knowhere::feature::MMAP | knowhere::feature::MV | knowhere::feature::EMB_LIST) +#ifdef KNOWHERE_WITH_CUVS +} // namespace knowhere — temporarily close to include GPU headers at file scope +// ── GPU HNSW ───────────────────────────────────────────────────────────────── +#include +#include +namespace knowhere { // reopen namespace knowhere + +// Process-global StandardGpuResources shared by all GpuHnswIndexNode instances. +// Avoids per-segment 256 MiB pinned memory, cuBLAS handle, and CUDA stream +// allocations that accumulate to tens of GiB with many segments. +// +// Device model: this assumes a single visible GPU per process, which is how +// Milvus GPU querynodes are deployed (one device pinned via CUDA_VISIBLE_DEVICES +// / MIG). GpuIndexHNSW below is constructed with the default device (0), so all +// segments in a process land on that one GPU. StandardGpuResources can manage +// multiple devices, but true multi-GPU-per-process would additionally require +// explicit per-segment device selection in the GpuIndexHNSW config; that is a +// follow-up and is intentionally not wired here (cannot be validated without +// multi-GPU capacity). +static std::shared_ptr& +GetSharedGpuResources() { + static std::once_flag flag; + static std::shared_ptr instance; + std::call_once(flag, [] { + instance = std::make_shared(); + instance->setTempMemory(0); + instance->setPinnedMemory(0); + }); + return instance; +} + +// Serialize GpuIndexHNSW construction across segments. +// StandardGpuResourcesImpl::initializeForDevice is not thread-safe; +// concurrent constructors race on the allocs_ map assertion. +static std::mutex& +GetGpuConstructionMutex() { + static std::mutex mtx; + return mtx; +} + +// GPU HNSW index node. Registered for F32, FP16, BF16 and INT8 (see +// registrations at the bottom of this file); it loads CPU-serialized HNSW (F32 +// Flat) or HNSW_SQ (SQ8/INT8, fp16, bf16) binaries at load time and uploads them +// to the GPU via faiss::gpu::GpuIndexHNSW. FP16/BF16/INT8 stay in their native +// low-precision layout on the device (2 bytes for fp16/bf16, 1 byte for int8) +// and are up-converted to fp32 per element inside the search kernel. The CPU +// copy is freed after upload, so this node exposes no CPU raw data. +class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { + public: + GpuHnswIndexNode(const int32_t& version, const Object& object) + : BaseFaissRegularIndexHNSWNode(version, object, DataFormatEnum::fp32) { + } + + // Constructor for native data-type variants (e.g. int8 DP4A path). + GpuHnswIndexNode(const int32_t& version, const Object& object, DataFormatEnum query_format) + : BaseFaissRegularIndexHNSWNode(version, object, query_format) { + } + + static std::unique_ptr + StaticCreateConfig() { + return std::make_unique(); + } + + static bool + StaticHasRawData(const knowhere::BaseConfig& config, const IndexVersion& version) { + // The CPU copy is freed after the GPU upload, so no CPU-reconstructable + // raw data is retained. Keep this consistent with the instance + // HasRawData() override below. + return false; + } + + static expected + StaticEstimateLoadResource(const uint64_t file_size_in_bytes, const int64_t num_rows, const int64_t dim, + const knowhere::BaseConfig& config, const IndexVersion& version) { + // GPU HNSW stores the vectors and graph in VRAM; the CPU copy is freed + // after upload in Deserialize()/DeserializeFromFile(). Phase-separated + // accounting: + // memoryCost (retained after load): ~0 — the CPU index is released + // once the graph/vectors are on the GPU, so a loaded segment holds + // no steady-state host RAM (its data lives in VRAM). + // maxMemoryCost (transient peak during load): the host RAM that is + // briefly live before the GPU upload frees the CPU copy. + // + // The on-disk index for this collection is the compact FAISS int8-SQ + // form (IndexHNSWSQCosine: signed-int8 codes + graph), the SAME form the + // CPU HNSW cluster loads resident at ~file_size. Deserialize reads it + // into host RAM at ~file_size. For native int8/fp16/bf16 upload paths + // the codes are uploaded to the device directly with no fp32 + // reconstruct/decode staging. (Note: unsigned SQ8 via QT_8bit does call + // sa_decode() into a float buffer, so that path would incur fp32 + // staging — but it is not used in production GPU HNSW loads here.) + // The transient peak is therefore dominated by the serialized read + // buffer plus the deserialized compact index, ~2x the on-disk file + // size — not the fp32 rows*dim*4 blowup that a Flat/hnswlib path would + // incur. Basing the estimate on file_size also avoids relying on + // num_rows/dim, which can arrive as 0 at estimate time. + // + // KNOWN LIMITATION: this 2x is correct for the native int8/fp16/bf16 + // and int8-SQ upload paths (no fp32 staging). It UNDER-counts the two + // paths that materialize a full fp32 buffer during load — VECTOR_FLOAT + // IndexHNSWFlat (from_faiss_hnsw_flat does a reconstruct_n into a + // host float[]) and unsigned QT_8bit SQ8 (sa_decode into float) — whose + // true transient peak is closer to file_size + rows*dim*4. It is not + // corrected here because this static entry point receives no data-type + // discriminator (the same non-templated StaticEstimateLoadResource is + // registered for fp32/fp16/bf16/int8), so adding a blanket rows*dim*4 + // term would badly over-reserve host RAM on the compact int8 path that + // production actually uses. Fixing it properly needs a per-DataType + // estimator; tracked as a follow-up. + (void)num_rows; + (void)dim; + (void)config; + (void)version; + + const uint64_t peak = file_size_in_bytes * 2; + + LOG_KNOWHERE_INFO_ << "GPU_HNSW load estimate (compact int8): file_size=" << file_size_in_bytes + << " transient_peak=" << peak << " retained=0"; + + return Resource{.memoryCost = 0, .diskCost = 0, .maxMemoryCost = peak}; + } + + std::unique_ptr + CreateConfig() const override { + return StaticCreateConfig(); + } + + std::string + Type() const override { + return IndexEnum::INDEX_GPU_HNSW; + } + + // After the eager GPU upload the CPU index (indexes[0]) is freed, so the + // inherited Count()/Dim()/HasRawData()/GetVectorByIds() would report an + // empty index. Serve vector count and dimension from the metadata captured + // before the CPU copy was released. + int64_t + Count() const override { + if (gpu_ready_.load(std::memory_order_acquire)) { + return gpu_ntotal_; + } + return BaseFaissRegularIndexHNSWNode::Count(); + } + + int64_t + Dim() const override { + if (gpu_ready_.load(std::memory_order_acquire)) { + return gpu_dim_; + } + return BaseFaissRegularIndexHNSWNode::Dim(); + } + + bool + HasRawData(const std::string& metric_type) const override { + // The CPU copy is freed after the GPU upload; GPU_HNSW does not expose + // CPU-reconstructable raw data (kept consistent with StaticHasRawData). + if (gpu_ready_.load(std::memory_order_acquire)) { + return false; + } + return BaseFaissRegularIndexHNSWNode::HasRawData(metric_type); + } + + expected + GetVectorByIds(const DataSetPtr dataset, milvus::OpContext* op_context) const override { + if (gpu_ready_.load(std::memory_order_acquire)) { + return expected::Err(Status::not_implemented, + "GPU_HNSW keeps vectors on GPU; raw data retrieval is not supported"); + } + return BaseFaissRegularIndexHNSWNode::GetVectorByIds(dataset, op_context); + } + + protected: + Status + TrainInternal(const DataSetPtr /*dataset*/, const Config& /*cfg*/) override { + return Status::not_implemented; + } + + public: + Status + Deserialize(const BinarySet& binset, std::shared_ptr cfg) override { + std::unique_lock lock(gpu_mutex_); + UnpublishGpuIndexLocked(); + + // Accept CPU-built HNSW (F32) or HNSW_SQ (quantized) binaries. + // The base loader reads the payload via binset.GetByName(Type()), so + // the incoming binary must be filed under *this* node's Type(): + // INDEX_GPU_HNSW for the base node, INDEX_GPU_HNSW_SQ for the SQ + // subclass. Aliasing under a hardcoded INDEX_GPU_HNSW would make + // GPU_HNSW_SQ loads fail — the SQ loader looks up INDEX_GPU_HNSW_SQ and + // would find nothing. + Status status; + const std::string self_key = Type(); + if (!binset.Contains(self_key)) { + BinarySet aliased = binset; + for (const char* key : {IndexEnum::INDEX_GPU_HNSW, IndexEnum::INDEX_GPU_HNSW_SQ, + IndexEnum::INDEX_HNSW_SQ, IndexEnum::INDEX_HNSW}) { + if (binset.Contains(key)) { + aliased.Append(self_key, binset.GetByName(key)); + break; + } + } + status = BaseFaissRegularIndexHNSWNode::Deserialize(aliased, cfg); + } else { + status = BaseFaissRegularIndexHNSWNode::Deserialize(binset, cfg); + } + if (status != Status::success) { + return status; + } + + return UploadCpuIndexToGpuLocked(); + } + + // Milvus loads sealed segments through the mmap/file path + // (VectorMemIndex::LoadFromFile -> DeserializeFromFile), NOT Deserialize(). + // The base only materializes the CPU index, so without this override the + // eager GPU upload would never run on that path: the CPU-built HNSW would + // stay resident in host RAM (fp32-expanded for quantized inputs) and the GPU + // would sit unused. Load the CPU index via the base, then upload + free it. + Status + DeserializeFromFile(const std::string& filename, std::shared_ptr config) override { + std::unique_lock lock(gpu_mutex_); + UnpublishGpuIndexLocked(); + + auto status = BaseFaissRegularIndexHNSWNode::DeserializeFromFile(filename, config); + if (status != Status::success) { + return status; + } + + return UploadCpuIndexToGpuLocked(); + } + + expected + Search(const DataSetPtr dataset, std::unique_ptr cfg, const BitsetView& bitset, + milvus::OpContext* op_context) const override { + // Filtered search (deletes / TTL / partition / visibility bitset) is + // applied inside the GPU kernel: filtered rows stay graph waypoints but + // never enter the top-k, and a brute-force fallback guarantees k live + // results under heavy deletes (see faiss GpuHnswSearch.cuh / + // GpuHnswSearchKernel.cuh). The kernel tests the bitset by raw FAISS + // node id, which equals the segment row offset == BitsetView index + // (proven by the gpu_hnsw_id_mapping test). That direct mapping only + // holds when the view has no id remapping and no id offset, so those + // rarer shapes (never produced for a sealed GPU segment's delete + // bitset) are still rejected rather than silently mis-filtered. + const bool has_filter = !bitset.empty() && bitset.data() != nullptr; + if (has_filter && (bitset.has_out_ids() || bitset.id_offset() != 0)) { + return expected::Err( + Status::invalid_args, + "GPU_HNSW filtered search does not support id-remapped or " + "id-offset bitsets"); + } + + // Take a shared-ownership snapshot of the GPU index under gpu_mutex_. The + // snapshot keeps the index alive for the whole search even if a concurrent + // Deserialize() reload resets/rebuilds gpu_index_; the lock is held only + // for the pointer copy, not the search itself. (std::atomic + // is C++20/libstdc++-12 only, so guard the plain shared_ptr with the + // existing mutex instead.) + std::shared_ptr gpu_snapshot; + { + std::lock_guard lock(gpu_mutex_); + gpu_snapshot = gpu_index_; + } + if (!gpu_snapshot) { + std::unique_lock lock(gpu_mutex_); + gpu_snapshot = gpu_index_; + if (!gpu_snapshot) { + if (GetFaissHnswIndex() == nullptr) { + return expected::Err(Status::empty_index, "index not loaded"); + } + // Reuse the eager upload path (it holds no lock itself and we + // already hold gpu_mutex_). This keeps metric detection + // identical to Deserialize()/DeserializeFromFile(): the metric + // is derived from the FAISS index type (dynamic_cast for cosine, + // metric_type for IP), NOT from the query config. Detecting from + // config here was a latent bug — a search whose config omits + // metric_type would default to L2 and silently run a cosine/IP + // index with the wrong metric. + Status upload_status = UploadCpuIndexToGpuLocked(); + if (upload_status != Status::success) { + return expected::Err(Status::cuvs_inner_error, + "failed to build GPU HNSW index"); + } + gpu_snapshot = gpu_index_; + if (!gpu_snapshot) { + return expected::Err(Status::empty_index, "index not loaded"); + } + } + } + + const auto& hnsw_cfg = static_cast(*cfg); + auto k = hnsw_cfg.k.value(); + auto nq = dataset->GetRows(); + auto dim = dataset->GetDim(); + auto ef = hnsw_cfg.ef.value_or(200); + + // Derive the metric from the built GPU index (captured at upload time + // from the FAISS index type), NOT from the search config. A config that + // omits metric_type defaults to L2, which would skip cosine query + // normalization and silently return wrong scores for a cosine index. + const bool is_cosine = gpu_snapshot->isCosine(); + + // The kernel indexes the bitset by node id in [0, ntotal); a bitset + // shorter than ntotal would read out of bounds. Milvus sizes the delete + // bitset to the segment row count (== ntotal), so this only guards + // against a malformed caller. + if (has_filter && + static_cast(bitset.size()) < gpu_snapshot->ntotal) { + return expected::Err( + Status::invalid_args, + "GPU_HNSW filtered search: bitset shorter than index size"); + } + + // Fill the filtered-search fields of a search-params struct from the + // incoming bitset. A null/empty bitset leaves them defaulted, taking + // the unfiltered fast path. + auto set_filter_params = [&](faiss::gpu::GpuHnswSearchParams& gsp) { + if (!has_filter) { + return; + } + gsp.bitset_data = bitset.data(); + gsp.bitset_nbits = static_cast(bitset.size()); + gsp.bitset_filtered_count = static_cast(bitset.count()); + gsp.disable_fallback_brute_force = + hnsw_cfg.disable_fallback_brute_force.value_or(false); + }; + + auto h_ids = std::make_unique(nq * k); + auto h_dist = std::make_unique(nq * k); + + // Native int8 DP4A path: queries arrive as raw int8 (no MockWrapper upcast). + if (data_format == DataFormatEnum::int8) { + const auto* h_queries_i8 = + reinterpret_cast(dataset->GetTensor()); + try { + faiss::gpu::GpuHnswSearchParams gsp; + gsp.ef = ef; + set_filter_params(gsp); + gpu_snapshot->searchHostInt8(nq, h_queries_i8, k, h_dist.get(), h_ids.get(), gsp); + } catch (const std::exception& e) { + LOG_KNOWHERE_ERROR_ << "GPU_HNSW int8 search failed: " << e.what(); + return expected::Err(Status::cuvs_inner_error, + std::string("GPU HNSW int8 search failed: ") + e.what()); + } + + // For COSINE, post-scale distances by 1/||q|| so scores are in + // [-1, 1]. The kernel returns dot(shifted_q, shifted_db) * (1/||db||) + // (the true cosine numerator scaled by the db inverse norm, already + // sign-corrected to "larger == closer"); dividing by the query norm + // finalizes the cosine similarity. + if (is_cosine) { + for (int64_t i = 0; i < static_cast(nq); i++) { + const int8_t* q = h_queries_i8 + i * dim; + int32_t sq_norm = 0; + for (int64_t d = 0; d < static_cast(dim); d++) { + sq_norm += static_cast(q[d]) * static_cast(q[d]); + } + float inv_qnorm = (sq_norm > 0) ? (1.0f / std::sqrt(static_cast(sq_norm))) : 1.0f; + for (int64_t j = 0; j < static_cast(k); j++) { + h_dist[i * k + j] *= inv_qnorm; + } + } + } + + // No sign flip: faiss now returns the true similarity for IP/COSINE + // (larger == closer), negating its internal min-first score on + // copy-out. Previously the kernel returned the negated score and + // this path flipped it back. + return GenResultDataSet(nq, k, h_ids.release(), h_dist.release()); + } + + // fp32 path (fp16/bf16 have already been upcast by MockWrapper). + const auto* h_queries_raw = reinterpret_cast(dataset->GetTensor()); + + // For COSINE metric, normalize queries to unit length. + const float* h_queries = h_queries_raw; + std::unique_ptr normalized_queries; + if (is_cosine) { + normalized_queries = std::make_unique(nq * dim); + for (int64_t i = 0; i < nq; i++) { + const float* src = h_queries_raw + i * dim; + float* dst = normalized_queries.get() + i * dim; + float sq_norm = 0.0f; + for (int64_t d = 0; d < dim; d++) sq_norm += src[d] * src[d]; + float inv = (sq_norm > 0.0f) ? (1.0f / std::sqrt(sq_norm)) : 1.0f; + for (int64_t d = 0; d < dim; d++) dst[d] = src[d] * inv; + } + h_queries = normalized_queries.get(); + } + + try { + faiss::gpu::GpuHnswSearchParams gsp; + gsp.ef = ef; + set_filter_params(gsp); + gpu_snapshot->searchHost(nq, h_queries, k, h_dist.get(), h_ids.get(), gsp); + } catch (const std::exception& e) { + LOG_KNOWHERE_ERROR_ << "GPU_HNSW search failed: " << e.what(); + return expected::Err(Status::cuvs_inner_error, + std::string("GPU HNSW search failed: ") + e.what()); + } + + // No sign flip: faiss now returns the true similarity for IP/COSINE + // (larger == closer), negating its internal min-first score on copy-out. + return GenResultDataSet(nq, k, h_ids.release(), h_dist.release()); + } + + ~GpuHnswIndexNode() override = default; + + private: + // Requires gpu_mutex_ held. Unpublish before tearing down: on a reload a + // stale gpu_ready_==true combined with a reset (or failed re-upload of) + // gpu_index_ would let the lock-free reader in Search() dereference a null + // index. Clearing it here leaves the node honestly "not ready" if the + // subsequent (re-)upload fails. + void + UnpublishGpuIndexLocked() { + gpu_ready_.store(false, std::memory_order_release); + gpu_index_ = nullptr; + } + + // Requires gpu_mutex_ held. Uploads the just-deserialized CPU HNSW to the GPU + // via faiss::gpu::GpuIndexHNSW, releases the host copy, and publishes the + // device index. Shared by Deserialize() (BinarySet) and DeserializeFromFile() + // (Milvus mmap/file load path) so both entrypoints upload identically. A + // no-op success when there is no CPU index to upload (empty segment). + Status + UploadCpuIndexToGpuLocked() const { + const auto* faiss_idx = GetFaissHnswIndex(); + if (faiss_idx == nullptr) { + return Status::success; + } + try { + // Detect metric from the FAISS index type rather than config, because + // deserialize may be called without metric_type in the config (e.g. + // an empty json defaults metric_type to L2). + bool is_cosine = + dynamic_cast(faiss_idx) != nullptr; + bool use_ip = is_cosine || (faiss_idx->metric_type == ::faiss::METRIC_INNER_PRODUCT); + + std::shared_ptr local; + { + std::lock_guard gpu_ctor_lock(GetGpuConstructionMutex()); + gpu_resources_ = GetSharedGpuResources(); + local = std::make_shared(gpu_resources_.get(), faiss_idx->d, + faiss_idx->metric_type); + } + local->copyFromWithMetric(faiss_idx, use_ip, is_cosine); + // Capture count/dim before releasing the CPU copy so Count()/Dim() + // keep working once indexes[0] is freed. + gpu_ntotal_ = faiss_idx->ntotal; + gpu_dim_ = faiss_idx->d; + // Release CPU copy — vectors and graph are now on GPU. This is what + // drops the transient host-RAM peak back to ~0. const_cast because + // this helper is const (reused from the const Search() lazy path); + // it only frees the now-uploaded host copy. + const_cast&>(indexes[0]).reset(); + // Publish last: Search() takes a locked snapshot of gpu_index_, and + // gpu_ready_ is released for the lock-free Count()/Dim() readers. + gpu_index_ = local; + gpu_ready_.store(true, std::memory_order_release); + } catch (const std::exception& e) { + // Fail the load rather than reporting success with no GPU index: a + // silently "loaded" segment would only surface the error on the first + // search. Leave the member null so gpu_ready_ stays false and return + // an error so the caller treats the load as failed. + LOG_KNOWHERE_ERROR_ << "eager GPU upload failed: " << e.what(); + gpu_index_ = nullptr; + return Status::cuda_runtime_error; + } + return Status::success; + } + + const ::faiss::cppcontrib::knowhere::IndexHNSW* + GetFaissHnswIndex() const { + if (indexes.empty() || !indexes[0]) + return nullptr; + return dynamic_cast(indexes[0].get()); + } + + mutable std::mutex gpu_mutex_; + mutable std::shared_ptr gpu_resources_; + // Guarded by gpu_mutex_. Search() copies this into a local shared_ptr under + // the lock so the GPU index stays alive for the entire (unlocked) search even + // if a concurrent Deserialize() reload resets/rebuilds the member. A bare + // pointer read on the search fast path would be a use-after-free, and + // std::atomic is C++20/libstdc++-12 only (the build + // toolchain is GCC 11), so the mutex-guarded snapshot is used instead. + mutable std::shared_ptr gpu_index_; + // Published (release) after gpu_index_ is fully built; read lock-free + // (acquire) by Count()/Dim()/HasRawData()/GetVectorByIds(). + mutable std::atomic gpu_ready_{false}; + // Vector count/dim captured before the CPU copy is freed (written from the + // const lazy-upload path in Search(), hence mutable). + mutable int64_t gpu_ntotal_ = 0; + mutable int64_t gpu_dim_ = 0; +}; + +// GPU_HNSW_SQ is an alias of GPU_HNSW: the GPU search path is identical (a +// CPU-built HNSW/HNSW_SQ index uploaded to GpuIndexHNSW), so it reuses +// GpuHnswIndexNode and only reports a distinct Type() so Milvus's configured +// index_type matches the loaded node. +class GpuHnswSQIndexNode : public GpuHnswIndexNode { + public: + GpuHnswSQIndexNode(const int32_t& version, const Object& object) : GpuHnswIndexNode(version, object) { + } + + GpuHnswSQIndexNode(const int32_t& version, const Object& object, DataFormatEnum query_format) + : GpuHnswIndexNode(version, object, query_format) { + } + + std::string + Type() const override { + return IndexEnum::INDEX_GPU_HNSW_SQ; + } +}; + +// Register GPU_HNSW / GPU_HNSW_SQ in the static config map at process startup. +__attribute__((constructor)) static void +register_gpu_hnsw_static_config() { + IndexStaticFaced::Instance().RegisterStaticFunc(IndexEnum::INDEX_GPU_HNSW); + IndexStaticFaced::Instance().RegisterStaticFunc(IndexEnum::INDEX_GPU_HNSW); + IndexStaticFaced::Instance().RegisterStaticFunc(IndexEnum::INDEX_GPU_HNSW); + IndexStaticFaced::Instance().RegisterStaticFunc(IndexEnum::INDEX_GPU_HNSW); + IndexStaticFaced::Instance().RegisterStaticFunc(IndexEnum::INDEX_GPU_HNSW_SQ); + IndexStaticFaced::Instance().RegisterStaticFunc(IndexEnum::INDEX_GPU_HNSW_SQ); + IndexStaticFaced::Instance().RegisterStaticFunc(IndexEnum::INDEX_GPU_HNSW_SQ); + IndexStaticFaced::Instance().RegisterStaticFunc(IndexEnum::INDEX_GPU_HNSW_SQ); +} + +// GPU_HNSW is not registered as MMAP-capable: the CPU index is deserialized +// into host RAM, uploaded to VRAM, then the CPU copy is freed. For a native +// int8 index the transient CPU copy is compact (~1 byte/dim), matching how the +// CPU HNSW int8 node loads, so no memory-mapped-file path is required. +KNOWHERE_REGISTER_GLOBAL( + GPU_HNSW, + [](const int32_t& version, const Object& object) { return Index::Create(version, object); }, fp32, + true, (feature::GPU_ANN_FLOAT_INDEX)); + +KNOWHERE_REGISTER_GLOBAL( + GPU_HNSW, + [](const int32_t& version, const Object& object) { + return Index>::Create(std::make_unique(version, object)); + }, + fp16, true, (feature::FP16 | feature::GPU)); + +KNOWHERE_REGISTER_GLOBAL( + GPU_HNSW, + [](const int32_t& version, const Object& object) { + return Index>::Create(std::make_unique(version, object)); + }, + bf16, true, (feature::BF16 | feature::GPU)); + +KNOWHERE_REGISTER_GLOBAL( + GPU_HNSW, + [](const int32_t& version, const Object& object) { + // Native int8 path: bypass MockWrapper upcast, pass int8 queries directly + // to searchHostInt8(), which uploads the signed int8 queries verbatim (no + // +128 shift; the dataset upload already reverses FAISS's +128 SQ bias) + // and uses the DP4A kernel. + return Index::Create(version, object, DataFormatEnum::int8); + }, + int8, true, (feature::INT8 | feature::GPU)); + +KNOWHERE_REGISTER_GLOBAL( + GPU_HNSW_SQ, + [](const int32_t& version, const Object& object) { return Index::Create(version, object); }, + fp32, true, (feature::GPU_ANN_FLOAT_INDEX)); + +KNOWHERE_REGISTER_GLOBAL( + GPU_HNSW_SQ, + [](const int32_t& version, const Object& object) { + return Index>::Create(std::make_unique(version, object)); + }, + fp16, true, (feature::FP16 | feature::GPU)); + +KNOWHERE_REGISTER_GLOBAL( + GPU_HNSW_SQ, + [](const int32_t& version, const Object& object) { + return Index>::Create(std::make_unique(version, object)); + }, + bf16, true, (feature::BF16 | feature::GPU)); + +KNOWHERE_REGISTER_GLOBAL( + GPU_HNSW_SQ, + [](const int32_t& version, const Object& object) { + // Native int8 path: bypass MockWrapper upcast, pass int8 queries directly + // to searchHostInt8(), which uploads the signed int8 queries verbatim (no + // +128 shift; the dataset upload already reverses FAISS's +128 SQ bias) + // and uses the DP4A kernel. + return Index::Create(version, object, DataFormatEnum::int8); + }, + int8, true, (feature::INT8 | feature::GPU)); +#endif // KNOWHERE_WITH_CUVS + } // namespace knowhere diff --git a/tests/ut/test_gpu_search.cc b/tests/ut/test_gpu_search.cc index 4f7f21ed7..adaf63ff7 100644 --- a/tests/ut/test_gpu_search.cc +++ b/tests/ut/test_gpu_search.cc @@ -9,7 +9,16 @@ // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express // or implied. See the License for the specific language governing permissions and limitations under the License. +#include +#include +#include +#include +#include +#include #include +#include +#include +#include #include "catch2/catch_approx.hpp" #include "catch2/catch_test_macros.hpp" @@ -23,6 +32,10 @@ #ifdef KNOWHERE_WITH_CUVS +// Host-safe header (no device kernels) exposing GpuHnswUploadFaultInjection, the +// test-only hook used by the CUDA-upload fault-injection section below. +#include + template void check_search(const int64_t nb, const int64_t nq, const int64_t dim, const int64_t seed, std::string name, int version, @@ -443,5 +456,1697 @@ TEST_CASE("Test All GPU Index", "[search]") { CHECK(GetRelativeLoss(gt_dist[i], dist[i]) < 0.1f); } } + + // GPU_HNSW tests: build on CPU as HNSW, serialize, deserialize as GPU_HNSW, search on GPU. + SECTION("Test GPU HNSW Search (CPU build -> GPU search)") { + // Build a CPU HNSW index + knowhere::Json hnsw_json; + hnsw_json[knowhere::meta::DIM] = dim; + hnsw_json[knowhere::meta::METRIC_TYPE] = knowhere::metric::L2; + hnsw_json[knowhere::meta::TOPK] = 1; + hnsw_json[knowhere::indexparam::HNSW_M] = 16; + hnsw_json[knowhere::indexparam::EFCONSTRUCTION] = 200; + hnsw_json[knowhere::indexparam::EF] = 200; + + auto train_ds = GenDataSet(nb, dim, seed); + auto query_ds = GenDataSet(nq, dim, seed + 2); + + auto cpu_idx = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_HNSW, version).value(); + auto res = cpu_idx.Build(train_ds, hnsw_json); + REQUIRE(res == knowhere::Status::success); + + // Serialize the CPU index + knowhere::BinarySet bs; + cpu_idx.Serialize(bs); + + // Deserialize as GPU_HNSW + auto gpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version) + .value(); + auto deser_res = gpu_idx.Deserialize(bs); + REQUIRE(deser_res == knowhere::Status::success); + + // Count()/Dim() must survive the CPU-copy release after GPU upload; + // regression guard for returning -1 (segment treated as empty). + REQUIRE(gpu_idx.Count() == nb); + REQUIRE(gpu_idx.Dim() == dim); + + // Search on GPU + auto results = gpu_idx.Search(query_ds, hnsw_json, nullptr); + REQUIRE(results.has_value()); + + // Compare against brute force + auto gt = knowhere::BruteForce::Search(train_ds, query_ds, hnsw_json, nullptr); + REQUIRE(gt.has_value()); + float recall = GetKNNRecall(*gt.value(), *results.value()); + REQUIRE(recall >= 0.9f); + } + + SECTION("Test GPU HNSW Search Cosine Metric") { + knowhere::Json hnsw_json; + hnsw_json[knowhere::meta::DIM] = dim; + hnsw_json[knowhere::meta::METRIC_TYPE] = knowhere::metric::COSINE; + hnsw_json[knowhere::meta::TOPK] = 1; + hnsw_json[knowhere::indexparam::HNSW_M] = 16; + hnsw_json[knowhere::indexparam::EFCONSTRUCTION] = 200; + hnsw_json[knowhere::indexparam::EF] = 200; + + auto train_ds = GenDataSet(nb, dim, seed); + auto query_ds = GenDataSet(nq, dim, seed + 1); + + auto cpu_idx = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_HNSW, version).value(); + auto res = cpu_idx.Build(train_ds, hnsw_json); + REQUIRE(res == knowhere::Status::success); + + knowhere::BinarySet bs; + cpu_idx.Serialize(bs); + + auto gpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version) + .value(); + auto deser_res = gpu_idx.Deserialize(bs); + REQUIRE(deser_res == knowhere::Status::success); + + auto results = gpu_idx.Search(query_ds, hnsw_json, nullptr); + REQUIRE(results.has_value()); + + auto gt = knowhere::BruteForce::Search(train_ds, query_ds, hnsw_json, nullptr); + REQUIRE(gt.has_value()); + float recall = GetKNNRecall(*gt.value(), *results.value()); + // GPU HNSW cosine recall tracks L2 (measured ~0.98 on L40S); keep the + // floor in line with the L2/IP sections so real regressions are caught. + REQUIRE(recall >= 0.80f); + } + + SECTION("Test GPU HNSW Cosine Metric Omitted From Search Config") { + // Regression for I1: the GPU node must derive the metric from the built + // GPU index (via GpuIndexHNSW::isCosine()), NOT from the search config. + // A search config that omits metric_type defaults to L2, which would + // skip cosine query normalization and silently return wrong scores. + // Here the index is COSINE but the search json carries no metric_type; + // results must match a search whose json explicitly sets COSINE. + knowhere::Json build_json; + build_json[knowhere::meta::DIM] = dim; + build_json[knowhere::meta::METRIC_TYPE] = knowhere::metric::COSINE; + build_json[knowhere::meta::TOPK] = 10; + build_json[knowhere::indexparam::HNSW_M] = 16; + build_json[knowhere::indexparam::EFCONSTRUCTION] = 200; + build_json[knowhere::indexparam::EF] = 200; + + auto train_ds = GenDataSet(nb, dim, seed); + auto query_ds = GenDataSet(nq, dim, seed + 1); + + auto cpu_idx = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_HNSW, version).value(); + REQUIRE(cpu_idx.Build(train_ds, build_json) == knowhere::Status::success); + + knowhere::BinarySet bs; + cpu_idx.Serialize(bs); + + auto gpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version) + .value(); + REQUIRE(gpu_idx.Deserialize(bs) == knowhere::Status::success); + + // Baseline: search with metric_type explicitly set to COSINE. + auto with_metric = gpu_idx.Search(query_ds, build_json, nullptr); + REQUIRE(with_metric.has_value()); + + // Same search, but metric_type omitted from the config. + knowhere::Json no_metric_json = build_json; + no_metric_json.erase(knowhere::meta::METRIC_TYPE); + auto without_metric = gpu_idx.Search(query_ds, no_metric_json, nullptr); + REQUIRE(without_metric.has_value()); + + // The metric is index-derived, so both searches must be identical. + const int64_t topk = build_json[knowhere::meta::TOPK].get(); + const auto* ids_a = with_metric.value()->GetIds(); + const auto* ids_b = without_metric.value()->GetIds(); + const auto* dist_a = with_metric.value()->GetDistance(); + const auto* dist_b = without_metric.value()->GetDistance(); + for (int64_t i = 0; i < nq * topk; ++i) { + REQUIRE(ids_a[i] == ids_b[i]); + REQUIRE(dist_a[i] == Approx(dist_b[i]).epsilon(1e-5)); + } + } + + SECTION("Test GPU HNSW Search TopK") { + knowhere::Json hnsw_json; + hnsw_json[knowhere::meta::DIM] = dim; + hnsw_json[knowhere::meta::METRIC_TYPE] = knowhere::metric::L2; + hnsw_json[knowhere::meta::TOPK] = 1; + hnsw_json[knowhere::indexparam::HNSW_M] = 16; + hnsw_json[knowhere::indexparam::EFCONSTRUCTION] = 200; + hnsw_json[knowhere::indexparam::EF] = 200; + + auto train_ds = GenDataSet(nb, dim, seed); + auto query_ds = GenDataSet(nq, dim, seed + 2); + + auto cpu_idx = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_HNSW, version).value(); + auto res = cpu_idx.Build(train_ds, hnsw_json); + REQUIRE(res == knowhere::Status::success); + + knowhere::BinarySet bs; + cpu_idx.Serialize(bs); + + auto gpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version) + .value(); + auto deser_res = gpu_idx.Deserialize(bs); + REQUIRE(deser_res == knowhere::Status::success); + + const auto topk_values = { + std::make_tuple(5, 0.85f), + std::make_tuple(25, 0.85f), + std::make_tuple(100, 0.85f), + }; + + for (const auto& [topk, threshold] : topk_values) { + hnsw_json[knowhere::meta::TOPK] = topk; + auto results = gpu_idx.Search(query_ds, hnsw_json, nullptr); + REQUIRE(results.has_value()); + auto gt = knowhere::BruteForce::Search(train_ds, query_ds, hnsw_json, nullptr); + float recall = GetKNNRecall(*gt.value(), *results.value()); + REQUIRE(recall >= threshold); + } + } + + SECTION("Test GPU HNSW Serialize/Deserialize Round Trip") { + knowhere::Json hnsw_json; + hnsw_json[knowhere::meta::DIM] = dim; + hnsw_json[knowhere::meta::METRIC_TYPE] = knowhere::metric::L2; + hnsw_json[knowhere::meta::TOPK] = 1; + hnsw_json[knowhere::indexparam::HNSW_M] = 16; + hnsw_json[knowhere::indexparam::EFCONSTRUCTION] = 200; + hnsw_json[knowhere::indexparam::EF] = 200; + + auto train_ds = GenDataSet(nb, dim, seed); + + auto cpu_idx = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_HNSW, version).value(); + auto res = cpu_idx.Build(train_ds, hnsw_json); + REQUIRE(res == knowhere::Status::success); + + knowhere::BinarySet bs; + cpu_idx.Serialize(bs); + + auto gpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version) + .value(); + auto deser_res = gpu_idx.Deserialize(bs); + REQUIRE(deser_res == knowhere::Status::success); + + // Self-search: each vector should find itself as nearest neighbor. + // The query set is train_ds (nb rows), so check all nb results — using + // nq only exercised the first nq of nb vectors. + auto results = gpu_idx.Search(train_ds, hnsw_json, nullptr); + REQUIRE(results.has_value()); + auto ids = results.value()->GetIds(); + int correct = 0; + for (int i = 0; i < nb; ++i) { + if (ids[i] == i) + correct++; + } + float self_recall = static_cast(correct) / nb; + REQUIRE(self_recall >= 0.95f); + } + + SECTION("Test GPU HNSW SQ8 Deserialization") { + // Build a CPU HNSW_SQ index, then load as GPU_HNSW + knowhere::Json hnsw_json; + hnsw_json[knowhere::meta::DIM] = dim; + hnsw_json[knowhere::meta::METRIC_TYPE] = knowhere::metric::L2; + hnsw_json[knowhere::meta::TOPK] = 1; + hnsw_json[knowhere::indexparam::HNSW_M] = 16; + hnsw_json[knowhere::indexparam::EFCONSTRUCTION] = 200; + hnsw_json[knowhere::indexparam::EF] = 200; + + auto train_ds = GenDataSet(nb, dim, seed); + auto query_ds = GenDataSet(nq, dim, seed + 2); + + auto cpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_HNSW_SQ, version) + .value(); + auto res = cpu_idx.Build(train_ds, hnsw_json); + REQUIRE(res == knowhere::Status::success); + + knowhere::BinarySet bs; + cpu_idx.Serialize(bs); + + // GPU_HNSW should accept HNSW_SQ binaries + auto gpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version) + .value(); + auto deser_res = gpu_idx.Deserialize(bs); + REQUIRE(deser_res == knowhere::Status::success); + + auto results = gpu_idx.Search(query_ds, hnsw_json, nullptr); + REQUIRE(results.has_value()); + + auto gt = knowhere::BruteForce::Search(train_ds, query_ds, hnsw_json, nullptr); + REQUIRE(gt.has_value()); + float recall = GetKNNRecall(*gt.value(), *results.value()); + // SQ8 has lower recall than flat due to quantization + REQUIRE(recall >= 0.7f); + } + + SECTION("Test GPU HNSW INT8 Deserialization") { + // Build a CPU int8 HNSW (HNSW_SQ QT_8bit_direct_signed storage), then + // load and search on GPU. int8 vectors stay in their native 1-byte + // layout on the device (upload_int8_dataset) and are up-converted to + // fp32 per element in the search kernel. This exercises the native + // signed-int8 device path, distinct from the unsigned SQ8 test above + // which routes through the decode-to-fp32 upload branch. + knowhere::Json hnsw_json; + hnsw_json[knowhere::meta::DIM] = dim; + hnsw_json[knowhere::meta::METRIC_TYPE] = knowhere::metric::L2; + hnsw_json[knowhere::meta::TOPK] = 1; + hnsw_json[knowhere::indexparam::HNSW_M] = 16; + hnsw_json[knowhere::indexparam::EFCONSTRUCTION] = 200; + hnsw_json[knowhere::indexparam::EF] = 200; + + auto train_ds = knowhere::ConvertToDataTypeIfNeeded(GenDataSet(nb, dim, seed)); + auto query_ds = knowhere::ConvertToDataTypeIfNeeded(GenDataSet(nq, dim, seed + 2)); + + auto cpu_idx = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_HNSW, version).value(); + auto res = cpu_idx.Build(train_ds, hnsw_json); + REQUIRE(res == knowhere::Status::success); + + knowhere::BinarySet bs; + cpu_idx.Serialize(bs); + + auto gpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version) + .value(); + auto deser_res = gpu_idx.Deserialize(bs); + REQUIRE(deser_res == knowhere::Status::success); + + auto results = gpu_idx.Search(query_ds, hnsw_json, nullptr); + REQUIRE(results.has_value()); + + auto gt = knowhere::BruteForce::Search(train_ds, query_ds, hnsw_json, nullptr); + REQUIRE(gt.has_value()); + float recall = GetKNNRecall(*gt.value(), *results.value()); + // int8 native path should match the CPU int8 HNSW recall closely. + REQUIRE(recall >= 0.9f); + } + + SECTION("Test GPU HNSW FP16 Deserialization") { + // Build a CPU HNSW (fp16 -> HNSW_SQ QT_fp16 storage), then load and + // search on GPU. fp16 vectors stay in their native 2-byte layout on the + // device and are up-converted to fp32 per element in the search kernel. + knowhere::Json hnsw_json; + hnsw_json[knowhere::meta::DIM] = dim; + hnsw_json[knowhere::meta::METRIC_TYPE] = knowhere::metric::L2; + hnsw_json[knowhere::meta::TOPK] = 1; + hnsw_json[knowhere::indexparam::HNSW_M] = 16; + hnsw_json[knowhere::indexparam::EFCONSTRUCTION] = 200; + hnsw_json[knowhere::indexparam::EF] = 200; + + auto train_ds = knowhere::ConvertToDataTypeIfNeeded(GenDataSet(nb, dim, seed)); + auto query_ds = knowhere::ConvertToDataTypeIfNeeded(GenDataSet(nq, dim, seed + 2)); + + auto cpu_idx = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_HNSW, version).value(); + auto res = cpu_idx.Build(train_ds, hnsw_json); + REQUIRE(res == knowhere::Status::success); + + knowhere::BinarySet bs; + cpu_idx.Serialize(bs); + + auto gpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version) + .value(); + auto deser_res = gpu_idx.Deserialize(bs); + REQUIRE(deser_res == knowhere::Status::success); + + auto results = gpu_idx.Search(query_ds, hnsw_json, nullptr); + REQUIRE(results.has_value()); + + auto gt = knowhere::BruteForce::Search(train_ds, query_ds, hnsw_json, nullptr); + REQUIRE(gt.has_value()); + float recall = GetKNNRecall(*gt.value(), *results.value()); + // fp16 is near-lossless (10-bit mantissa); recall should stay close to + // the fp32 flat path. + REQUIRE(recall >= 0.9f); + } + + SECTION("Test GPU HNSW BF16 Deserialization") { + // Build a CPU HNSW (bf16 -> HNSW_SQ QT_bf16 storage), then load and + // search on GPU. bf16 vectors stay in their native 2-byte layout on the + // device and are up-converted to fp32 per element in the search kernel. + knowhere::Json hnsw_json; + hnsw_json[knowhere::meta::DIM] = dim; + hnsw_json[knowhere::meta::METRIC_TYPE] = knowhere::metric::L2; + hnsw_json[knowhere::meta::TOPK] = 1; + hnsw_json[knowhere::indexparam::HNSW_M] = 16; + hnsw_json[knowhere::indexparam::EFCONSTRUCTION] = 200; + hnsw_json[knowhere::indexparam::EF] = 200; + + auto train_ds = knowhere::ConvertToDataTypeIfNeeded(GenDataSet(nb, dim, seed)); + auto query_ds = knowhere::ConvertToDataTypeIfNeeded(GenDataSet(nq, dim, seed + 2)); + + auto cpu_idx = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_HNSW, version).value(); + auto res = cpu_idx.Build(train_ds, hnsw_json); + REQUIRE(res == knowhere::Status::success); + + knowhere::BinarySet bs; + cpu_idx.Serialize(bs); + + auto gpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version) + .value(); + auto deser_res = gpu_idx.Deserialize(bs); + REQUIRE(deser_res == knowhere::Status::success); + + auto results = gpu_idx.Search(query_ds, hnsw_json, nullptr); + REQUIRE(results.has_value()); + + auto gt = knowhere::BruteForce::Search(train_ds, query_ds, hnsw_json, nullptr); + REQUIRE(gt.has_value()); + float recall = GetKNNRecall(*gt.value(), *results.value()); + // bf16 has a 7-bit mantissa (coarser than fp16) but a wide exponent; + // recall stays high but with a slightly looser floor. + REQUIRE(recall >= 0.85f); + } + + SECTION("Test GPU HNSW non-power-of-two 2*M staging (fp32 + int8 DP4A)") { + // M=24 => max_degree0 = 2*M = 48, which is NOT a power of two. The + // parallel bitonic-merge kernel pads the layer-0 staging capacity up to + // the next power of two (64) and grows the block accordingly. Before the + // padding fix this configuration hard-failed every GPU search. Cover + // both the generic fp32-query kernel and the native int8 DP4A kernel + // (dim=128 is divisible by 4, so int8 COSINE takes the DP4A path). + const int64_t m = 24; + + SECTION("fp32 cosine, non-pow2 2*M") { + knowhere::Json hnsw_json; + hnsw_json[knowhere::meta::DIM] = dim; + hnsw_json[knowhere::meta::METRIC_TYPE] = knowhere::metric::COSINE; + hnsw_json[knowhere::meta::TOPK] = 10; + hnsw_json[knowhere::indexparam::HNSW_M] = m; + hnsw_json[knowhere::indexparam::EFCONSTRUCTION] = 200; + hnsw_json[knowhere::indexparam::EF] = 200; + + auto train_ds = GenDataSet(nb, dim, seed); + auto query_ds = GenDataSet(nq, dim, seed + 2); + + auto cpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_HNSW, version) + .value(); + REQUIRE(cpu_idx.Build(train_ds, hnsw_json) == knowhere::Status::success); + + knowhere::BinarySet bs; + cpu_idx.Serialize(bs); + + auto gpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version) + .value(); + REQUIRE(gpu_idx.Deserialize(bs) == knowhere::Status::success); + + auto results = gpu_idx.Search(query_ds, hnsw_json, nullptr); + REQUIRE(results.has_value()); + + auto gt = knowhere::BruteForce::Search(train_ds, query_ds, hnsw_json, nullptr); + REQUIRE(gt.has_value()); + float recall = GetKNNRecall(*gt.value(), *results.value()); + REQUIRE(recall >= 0.80f); + } + + SECTION("int8 cosine DP4A, non-pow2 2*M") { + knowhere::Json hnsw_json; + hnsw_json[knowhere::meta::DIM] = dim; + hnsw_json[knowhere::meta::METRIC_TYPE] = knowhere::metric::COSINE; + hnsw_json[knowhere::meta::TOPK] = 10; + hnsw_json[knowhere::indexparam::HNSW_M] = m; + hnsw_json[knowhere::indexparam::EFCONSTRUCTION] = 200; + hnsw_json[knowhere::indexparam::EF] = 200; + + auto train_ds = knowhere::ConvertToDataTypeIfNeeded(GenDataSet(nb, dim, seed)); + auto query_ds = knowhere::ConvertToDataTypeIfNeeded(GenDataSet(nq, dim, seed + 2)); + + auto cpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_HNSW, version) + .value(); + REQUIRE(cpu_idx.Build(train_ds, hnsw_json) == knowhere::Status::success); + + knowhere::BinarySet bs; + cpu_idx.Serialize(bs); + + auto gpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version) + .value(); + REQUIRE(gpu_idx.Deserialize(bs) == knowhere::Status::success); + + auto results = gpu_idx.Search(query_ds, hnsw_json, nullptr); + REQUIRE(results.has_value()); + + auto gt = knowhere::BruteForce::Search(train_ds, query_ds, hnsw_json, nullptr); + REQUIRE(gt.has_value()); + float recall = GetKNNRecall(*gt.value(), *results.value()); + REQUIRE(recall >= 0.80f); + } + } +} + +TEST_CASE("Test CPU vs GPU HNSW Comparison", "[gpu_hnsw_compare]") { + using Catch::Approx; + + int64_t nb = 10000, nq = 100; + int64_t dim = 128; + int64_t seed = 42; + + auto version = GenTestVersionList(); + + auto run_comparison = [&](const std::string& metric, float min_recall, float max_dist_drift) { + CAPTURE(metric, nb, nq, dim); + + knowhere::Json hnsw_json; + hnsw_json[knowhere::meta::DIM] = dim; + hnsw_json[knowhere::meta::METRIC_TYPE] = metric; + hnsw_json[knowhere::meta::TOPK] = 10; + hnsw_json[knowhere::indexparam::HNSW_M] = 16; + hnsw_json[knowhere::indexparam::EFCONSTRUCTION] = 200; + hnsw_json[knowhere::indexparam::EF] = 200; + + auto train_ds = GenDataSet(nb, dim, seed); + auto query_ds = GenDataSet(nq, dim, seed + 2); + + // --- Build CPU HNSW --- + auto cpu_idx = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_HNSW, version).value(); + auto res = cpu_idx.Build(train_ds, hnsw_json); + REQUIRE(res == knowhere::Status::success); + + // Serialize for GPU deserialization + knowhere::BinarySet bs; + cpu_idx.Serialize(bs); + + // --- CPU Search with timing --- + auto cpu_start = std::chrono::high_resolution_clock::now(); + auto cpu_results = cpu_idx.Search(query_ds, hnsw_json, nullptr); + auto cpu_end = std::chrono::high_resolution_clock::now(); + REQUIRE(cpu_results.has_value()); + double cpu_ms = std::chrono::duration(cpu_end - cpu_start).count(); + + // --- Build GPU HNSW from serialized CPU index --- + auto gpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version) + .value(); + auto deser_res = gpu_idx.Deserialize(bs); + REQUIRE(deser_res == knowhere::Status::success); + + // Warm-up search (first GPU call has kernel launch overhead) + auto warmup = gpu_idx.Search(query_ds, hnsw_json, nullptr); + REQUIRE(warmup.has_value()); + + // --- GPU Search with timing --- + auto gpu_start = std::chrono::high_resolution_clock::now(); + auto gpu_results = gpu_idx.Search(query_ds, hnsw_json, nullptr); + auto gpu_end = std::chrono::high_resolution_clock::now(); + REQUIRE(gpu_results.has_value()); + double gpu_ms = std::chrono::duration(gpu_end - gpu_start).count(); + + // --- Brute force ground truth --- + auto gt = knowhere::BruteForce::Search(train_ds, query_ds, hnsw_json, nullptr); + REQUIRE(gt.has_value()); + + // --- Recall comparison --- + float cpu_recall = GetKNNRecall(*gt.value(), *cpu_results.value()); + float gpu_recall = GetKNNRecall(*gt.value(), *gpu_results.value()); + + // --- Distance accuracy: compare GPU distances against CPU distances --- + auto cpu_dist = cpu_results.value()->GetDistance(); + auto gpu_dist = gpu_results.value()->GetDistance(); + auto k = hnsw_json[knowhere::meta::TOPK].get(); + + double total_rel_error = 0.0; + int valid_pairs = 0; + for (int64_t i = 0; i < nq * k; ++i) { + if (std::abs(cpu_dist[i]) > 1e-6f) { + total_rel_error += std::abs((gpu_dist[i] - cpu_dist[i]) / cpu_dist[i]); + valid_pairs++; + } + } + double avg_rel_dist_error = (valid_pairs > 0) ? (total_rel_error / valid_pairs) : 0.0; + + // --- ID overlap: how many of the same results are returned --- + auto cpu_ids = cpu_results.value()->GetIds(); + auto gpu_ids = gpu_results.value()->GetIds(); + int id_overlap = 0; + for (int64_t q = 0; q < nq; ++q) { + std::set cpu_set(cpu_ids + q * k, cpu_ids + (q + 1) * k); + for (int64_t j = 0; j < k; ++j) { + if (cpu_set.count(gpu_ids[q * k + j])) { + id_overlap++; + } + } + } + float id_overlap_ratio = static_cast(id_overlap) / (nq * k); + + // --- Print comparison report --- + fprintf(stderr, "\n=== CPU vs GPU HNSW Comparison (%s, nb=%ld, nq=%ld, dim=%ld, k=%ld) ===\n", metric.c_str(), + (long)nb, (long)nq, (long)dim, (long)k); + fprintf(stderr, " CPU recall@%ld: %.4f\n", (long)k, cpu_recall); + fprintf(stderr, " GPU recall@%ld: %.4f\n", (long)k, gpu_recall); + fprintf(stderr, " Recall delta (GPU - CPU): %+.4f\n", gpu_recall - cpu_recall); + fprintf(stderr, " CPU search time: %.2f ms\n", cpu_ms); + fprintf(stderr, " GPU search time: %.2f ms (includes H2D/D2H transfers)\n", gpu_ms); + fprintf(stderr, " Speedup: %.2fx\n", cpu_ms / gpu_ms); + fprintf(stderr, " Avg relative distance error: %.6f\n", avg_rel_dist_error); + fprintf(stderr, " ID overlap (CPU vs GPU): %.4f (%d/%ld)\n", id_overlap_ratio, id_overlap, (long)(nq * k)); + fprintf(stderr, "===\n\n"); + + // --- Assertions --- + // GPU recall should be close to CPU recall (within a small tolerance) + REQUIRE(gpu_recall >= min_recall); + // GPU should return at least 80% of the same IDs as CPU + REQUIRE(id_overlap_ratio >= 0.8f); + // Average relative distance error should be small + REQUIRE(avg_rel_dist_error <= max_dist_drift); + }; + + SECTION("L2 metric comparison") { + run_comparison(knowhere::metric::L2, 0.85f, 0.05); + } + + SECTION("IP metric comparison") { + run_comparison(knowhere::metric::IP, 0.85f, 0.05); + } + + SECTION("COSINE metric comparison") { + run_comparison(knowhere::metric::COSINE, 0.60f, 0.10); + } + + SECTION("Varying ef comparison") { + knowhere::Json hnsw_json; + hnsw_json[knowhere::meta::DIM] = dim; + hnsw_json[knowhere::meta::METRIC_TYPE] = knowhere::metric::L2; + hnsw_json[knowhere::meta::TOPK] = 10; + hnsw_json[knowhere::indexparam::HNSW_M] = 16; + hnsw_json[knowhere::indexparam::EFCONSTRUCTION] = 200; + + auto train_ds = GenDataSet(nb, dim, seed); + auto query_ds = GenDataSet(nq, dim, seed + 2); + + auto cpu_idx = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_HNSW, version).value(); + auto res = cpu_idx.Build(train_ds, hnsw_json); + REQUIRE(res == knowhere::Status::success); + + knowhere::BinarySet bs; + cpu_idx.Serialize(bs); + + auto gpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version) + .value(); + auto deser_res = gpu_idx.Deserialize(bs); + REQUIRE(deser_res == knowhere::Status::success); + + auto gt = knowhere::BruteForce::Search(train_ds, query_ds, hnsw_json, nullptr); + REQUIRE(gt.has_value()); + + fprintf(stderr, "\n=== ef sweep (L2, nb=%ld, nq=%ld, k=10) ===\n", (long)nb, (long)nq); + fprintf(stderr, " %6s %10s %10s %10s %10s\n", "ef", "cpu_recall", "gpu_recall", "cpu_ms", "gpu_ms"); + + for (int ef : {16, 32, 64, 128, 256, 512}) { + hnsw_json[knowhere::indexparam::EF] = ef; + + auto t0 = std::chrono::high_resolution_clock::now(); + auto cpu_res = cpu_idx.Search(query_ds, hnsw_json, nullptr); + auto t1 = std::chrono::high_resolution_clock::now(); + auto gpu_res = gpu_idx.Search(query_ds, hnsw_json, nullptr); + auto t2 = std::chrono::high_resolution_clock::now(); + + REQUIRE(cpu_res.has_value()); + REQUIRE(gpu_res.has_value()); + + float cpu_recall = GetKNNRecall(*gt.value(), *cpu_res.value()); + float gpu_recall = GetKNNRecall(*gt.value(), *gpu_res.value()); + double cpu_ms = std::chrono::duration(t1 - t0).count(); + double gpu_ms = std::chrono::duration(t2 - t1).count(); + + fprintf(stderr, " %6d %10.4f %10.4f %10.2f %10.2f\n", ef, cpu_recall, gpu_recall, cpu_ms, gpu_ms); + + // This is a GPU-vs-CPU comparison: GPU HNSW must track the CPU HNSW oracle at + // every ef, not hit an arbitrary absolute floor. At small ef (e.g. ef=16, ~= k) + // HNSW recall is legitimately low on both engines (CPU itself is < 0.5 here), so + // an absolute floor is meaningless; parity with CPU is the correct invariant. + REQUIRE(gpu_recall >= cpu_recall - 0.05f); + } + fprintf(stderr, "===\n\n"); + } +} + +// Regression tests for the three Codex P1 findings on the gpu-hnsw-faiss branch: +// #1 quantized-cosine inverse-norm semantics (native low-precision kernels + +// SQ8/FP16/BF16 cosine parity), +// #2 search-vs-reload lifetime safety, +// #3 phase-separated load-resource accounting. +TEST_CASE("Test GPU HNSW Codex P1 Regressions", "[gpu_hnsw_p1]") { + using Catch::Approx; + + const int64_t nb = 4000, nq = 200; + const int64_t dim = 64; + const int64_t seed = 42; + const auto version = GenTestVersionList(); + + auto sq_json = [&](const std::string& metric, const std::string& sq_type, int64_t topk) { + knowhere::Json json; + json[knowhere::meta::DIM] = dim; + json[knowhere::meta::METRIC_TYPE] = metric; + json[knowhere::meta::TOPK] = topk; + json[knowhere::indexparam::HNSW_M] = 16; + json[knowhere::indexparam::EFCONSTRUCTION] = 200; + json[knowhere::indexparam::EF] = 200; + json[knowhere::indexparam::SQ_TYPE] = sq_type; + return json; + }; + + // Build an HNSW_SQ index on the fp32 (real, non-mock) node so the serialized + // storage carries the requested SQ qtype. Deserializing that as GPU_HNSW on + // the fp32 node selects the native device representation from the storage + // qtype (FP16 -> half kernel, BF16 -> __nv_bfloat16 kernel, SQ8 -> fp32 + // decode). This bypasses the int8/fp16/bf16 IndexNodeDataMockWrapper, which + // otherwise converts data to fp32 and routes the fp32 kernel. + + // #1a: native low-precision (FP16/BF16) and SQ8 device paths — L2. + SECTION("Native low-precision device path (L2): FP16 / BF16 / SQ8") { + const auto sq_type = GENERATE(std::string("fp16"), std::string("bf16"), std::string("sq8")); + CAPTURE(sq_type); + auto json = sq_json(knowhere::metric::L2, sq_type, 10); + + auto train_ds = GenDataSet(nb, dim, seed); + auto query_ds = GenDataSet(nq, dim, seed + 2); + + auto cpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_HNSW_SQ, version) + .value(); + REQUIRE(cpu_idx.Build(train_ds, json) == knowhere::Status::success); + knowhere::BinarySet bs; + cpu_idx.Serialize(bs); + + auto gpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version) + .value(); + REQUIRE(gpu_idx.Deserialize(bs) == knowhere::Status::success); + + auto results = gpu_idx.Search(query_ds, json, nullptr); + REQUIRE(results.has_value()); + auto gt = knowhere::BruteForce::Search(train_ds, query_ds, json, nullptr); + REQUIRE(gt.has_value()); + float recall = GetKNNRecall(*gt.value(), *results.value()); + // Lossy SQ paths; a conservative floor. fp16/bf16 are near-lossless, sq8 + // coarser. + REQUIRE(recall >= (sq_type == "sq8" ? 0.65f : 0.85f)); + } + + // #4: Milvus loads sealed segments through the mmap/file path + // (VectorMemIndex::LoadFromFile -> DeserializeFromFile), NOT + // Deserialize(BinarySet). The eager GPU upload must fire on that path too; + // otherwise the CPU-built HNSW stays resident in host RAM and the GPU is + // never used. This was the root cause of the mpd_v2 querynode OOM: each + // segment's (fp32-expanded) CPU index — ~5.9 GiB — was retained in host + // anon memory instead of being uploaded to VRAM and freed, so ~14 + // segments/pod overran the memory limit while VRAM stayed near-idle. + // + // The discriminating assertion is Size()==0 *before any search*: the eager + // upload releases indexes[0], so the base Size() (computed from indexes[0]) + // is 0. Pre-fix, DeserializeFromFile skipped the upload and left the CPU + // copy resident, so Size() reported the full retained host footprint. A + // recall-only check would NOT catch the regression, because the first + // Search() lazily uploads and frees the copy — masking the load-time leak + // that OOMs before any search runs. + SECTION("DeserializeFromFile uploads to GPU and frees the host copy") { + const bool enable_mmap = GENERATE(false, true); + CAPTURE(enable_mmap); + auto json = sq_json(knowhere::metric::COSINE, "sq8", 10); + + auto train_ds = GenDataSet(nb, dim, seed); + auto query_ds = GenDataSet(nq, dim, seed + 2); + + auto cpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_HNSW_SQ, version) + .value(); + REQUIRE(cpu_idx.Build(train_ds, json) == knowhere::Status::success); + knowhere::BinarySet bs; + REQUIRE(cpu_idx.Serialize(bs) == knowhere::Status::success); + auto binary = bs.GetByName(cpu_idx.Type()); + REQUIRE(binary != nullptr); + + const auto path = std::filesystem::temp_directory_path() / + ("knowhere_gpu_hnsw_deserialize_from_file_" + std::to_string(enable_mmap) + ".index"); + std::filesystem::remove(path); + { + std::ofstream out(path, std::ios::binary); + REQUIRE(out.is_open()); + out.write(reinterpret_cast(binary->data.get()), binary->size); + } + + auto gpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version) + .value(); + auto load_json = json; + load_json["enable_mmap"] = enable_mmap; + REQUIRE(gpu_idx.DeserializeFromFile(path.string(), load_json) == knowhere::Status::success); + + // Eager upload ran during load (no search yet): host copy released. + REQUIRE(gpu_idx.Size() == 0); + REQUIRE(gpu_idx.Count() == nb); + REQUIRE(gpu_idx.HasRawData(knowhere::metric::COSINE) == false); + + // And the GPU index actually serves searches. + auto results = gpu_idx.Search(query_ds, json, nullptr); + REQUIRE(results.has_value()); + auto gt = knowhere::BruteForce::Search(train_ds, query_ds, json, nullptr); + REQUIRE(gt.has_value()); + float recall = GetKNNRecall(*gt.value(), *results.value()); + REQUIRE(recall >= 0.65f); + + std::filesystem::remove(path); + } + + // #1b: quantized-cosine parity. The discriminating input is vectors that + // share directions but have widely different magnitudes: cosine must ignore + // magnitude. The CPU cosine index records inverse norms from the ORIGINAL + // input vectors; the GPU path must upload those same norms rather than + // recomputing them from lossily-decoded codes. With the pre-fix behavior + // (recompute-from-decoded / normalize-decoded) the SQ8 scores and top-1 IDs + // diverge from the CPU index. Assert per-query top-1 ID and per-result score, + // not just aggregate recall. + SECTION("Quantized cosine parity (CPU vs GPU): SQ8 / FP16 / BF16") { + const auto sq_type = GENERATE(std::string("sq8"), std::string("fp16"), std::string("bf16")); + CAPTURE(sq_type); + const int64_t topk = 10; + auto json = sq_json(knowhere::metric::COSINE, sq_type, topk); + + // Directions from a fixed RNG, then scale each row by a magnitude that + // spans three orders of magnitude so magnitude cannot be ignored by luck. + auto make_scaled = [&](int64_t rows, int64_t s) { + auto ds = GenDataSet(rows, dim, s); + auto* data = const_cast(static_cast(ds->GetTensor())); + std::mt19937 rng(static_cast(s + 7)); + std::uniform_real_distribution mag(0.5f, 500.0f); + for (int64_t r = 0; r < rows; ++r) { + float scale = mag(rng); + for (int64_t d = 0; d < dim; ++d) { + data[r * dim + d] *= scale; + } + } + return ds; + }; + auto train_ds = make_scaled(nb, seed); + auto query_ds = make_scaled(nq, seed + 2); + + auto cpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_HNSW_SQ, version) + .value(); + REQUIRE(cpu_idx.Build(train_ds, json) == knowhere::Status::success); + auto cpu_results = cpu_idx.Search(query_ds, json, nullptr); + REQUIRE(cpu_results.has_value()); + + knowhere::BinarySet bs; + cpu_idx.Serialize(bs); + auto gpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version) + .value(); + REQUIRE(gpu_idx.Deserialize(bs) == knowhere::Status::success); + auto gpu_results = gpu_idx.Search(query_ds, json, nullptr); + REQUIRE(gpu_results.has_value()); + + const auto* cpu_ids = cpu_results.value()->GetIds(); + const auto* gpu_ids = gpu_results.value()->GetIds(); + const auto* cpu_dist = cpu_results.value()->GetDistance(); + const auto* gpu_dist = gpu_results.value()->GetDistance(); + + // Top-1 must match the CPU index for the overwhelming majority of + // queries: identical cosine semantics (same norms) means identical + // nearest neighbor modulo rare quantization ties. + int top1_match = 0; + for (int64_t q = 0; q < nq; ++q) { + if (cpu_ids[q * topk] == gpu_ids[q * topk]) { + top1_match++; + } + } + float top1_ratio = static_cast(top1_match) / nq; + CAPTURE(top1_ratio); + REQUIRE(top1_ratio >= 0.95f); + + // Per-result cosine scores must agree closely (both decode identical SQ + // codes and apply the SAME original inverse norms). A loose floor would + // hide the pre-fix norm mismatch, so keep the tolerance tight. + int compared = 0, close = 0; + for (int64_t q = 0; q < nq; ++q) { + // Only compare positions where CPU and GPU agree on the id, so score + // comparison is like-for-like. + for (int64_t j = 0; j < topk; ++j) { + if (cpu_ids[q * topk + j] == gpu_ids[q * topk + j]) { + compared++; + if (GetRelativeLoss(cpu_dist[q * topk + j], gpu_dist[q * topk + j]) < 0.02f) { + close++; + } + } + } + } + REQUIRE(compared > 0); + float close_ratio = static_cast(close) / compared; + CAPTURE(close_ratio); + REQUIRE(close_ratio >= 0.98f); + } + + // #2: search must not use-after-free a GPU index that a concurrent reload + // resets/rebuilds. Search continuously from several threads while another + // thread repeatedly Deserialize()s (reloads) the same index. Run under + // TSAN/ASAN to catch the lifetime race the atomic snapshot fixes. + SECTION("Search vs concurrent reload (lifetime safety)") { + auto json = sq_json(knowhere::metric::L2, "fp16", 10); + json.erase(knowhere::indexparam::SQ_TYPE); // plain fp32 HNSW for the reload driver + auto train_ds = GenDataSet(nb, dim, seed); + auto query_ds = GenDataSet(nq, dim, seed + 2); + + auto cpu_idx = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_HNSW, version).value(); + REQUIRE(cpu_idx.Build(train_ds, json) == knowhere::Status::success); + knowhere::BinarySet bs; + cpu_idx.Serialize(bs); + + auto gpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version) + .value(); + REQUIRE(gpu_idx.Deserialize(bs) == knowhere::Status::success); + + std::atomic stop{false}; + std::atomic failed{false}; + std::vector searchers; + for (int t = 0; t < 6; ++t) { + searchers.emplace_back([&]() { + while (!stop.load(std::memory_order_relaxed)) { + auto r = gpu_idx.Search(query_ds, json, nullptr); + // A search may race a reload window; success or a benign + // "not ready" error are both acceptable, a crash is not. + if (!r.has_value() && r.error() != knowhere::Status::empty_index && + r.error() != knowhere::Status::cuda_runtime_error) { + failed.store(true, std::memory_order_relaxed); + } + } + }); + } + std::thread reloader([&]() { + for (int i = 0; i < 30 && !stop.load(std::memory_order_relaxed); ++i) { + if (gpu_idx.Deserialize(bs) != knowhere::Status::success) { + failed.store(true, std::memory_order_relaxed); + } + } + stop.store(true, std::memory_order_relaxed); + }); + reloader.join(); + stop.store(true, std::memory_order_relaxed); + for (auto& th : searchers) { + th.join(); + } + REQUIRE(!failed.load()); + // Index is still usable after the reload storm. + auto after = gpu_idx.Search(query_ds, json, nullptr); + REQUIRE(after.has_value()); + } + + // Codex coverage #3: more than four simultaneous searches (the device scratch + // pool has four slots), with varying nq/k/ef, each validated against its own + // brute-force ground truth. Exercises scratch-pool contention/serialization. + SECTION("More than four concurrent searches with varying nq/k/ef") { + auto build_json = sq_json(knowhere::metric::L2, "fp16", 10); + build_json.erase(knowhere::indexparam::SQ_TYPE); + auto train_ds = GenDataSet(nb, dim, seed); + + auto cpu_idx = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_HNSW, version).value(); + REQUIRE(cpu_idx.Build(train_ds, build_json) == knowhere::Status::success); + knowhere::BinarySet bs; + cpu_idx.Serialize(bs); + auto gpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version) + .value(); + REQUIRE(gpu_idx.Deserialize(bs) == knowhere::Status::success); + + struct Case { + int64_t nq; + int64_t k; + int ef; + }; + const std::vector cases = {{50, 5, 32}, {120, 10, 64}, {200, 20, 128}, {80, 50, 256}, + {30, 10, 512}, {150, 1, 48}, {64, 100, 300}, {90, 25, 96}}; + std::atomic ok{true}; + std::vector threads; + for (const auto& c : cases) { + threads.emplace_back([&, c]() { + knowhere::Json j; + j[knowhere::meta::DIM] = dim; + j[knowhere::meta::METRIC_TYPE] = knowhere::metric::L2; + j[knowhere::meta::TOPK] = c.k; + j[knowhere::indexparam::EF] = c.ef; + auto q = GenDataSet(c.nq, dim, seed + 100 + c.ef); + auto r = gpu_idx.Search(q, j, nullptr); + if (!r.has_value()) { + ok.store(false, std::memory_order_relaxed); + return; + } + auto gt = knowhere::BruteForce::Search(train_ds, q, j, nullptr); + if (!gt.has_value() || GetKNNRecall(*gt.value(), *r.value()) < 0.7f) { + ok.store(false, std::memory_order_relaxed); + } + }); + } + for (auto& th : threads) { + th.join(); + } + REQUIRE(ok.load()); + } + + // #3: phase-separated load-resource accounting. GPU_HNSW frees its CPU copy + // after uploading to VRAM, so retained memoryCost is 0 while the transient + // peak (maxMemoryCost) must be non-zero and cover the serialized read buffer + // + deserialized compact index (~2x file_size). This estimate is a static + // computation and needs no live GPU. + SECTION("Load-resource estimate: memoryCost=0, maxMemoryCost covers peak") { + knowhere::Json params; + params[knowhere::meta::METRIC_TYPE] = knowhere::metric::L2; + params[knowhere::meta::DIM] = dim; + + const uint64_t file_size = static_cast(nb) * dim * sizeof(float); + auto gpu_res = knowhere::IndexStaticFaced::EstimateLoadResource( + knowhere::IndexEnum::INDEX_GPU_HNSW, version, file_size, nb, dim, params); + REQUIRE(gpu_res.has_value()); + // Retained host memory is ~0 after the CPU copy is freed. + REQUIRE(gpu_res.value().memoryCost == 0); + REQUIRE(gpu_res.value().diskCost == 0); + // Transient peak is non-zero and at least 2x the on-disk file size + // (serialized read buffer + deserialized compact index). + REQUIRE(gpu_res.value().maxMemoryCost > 0); + REQUIRE(gpu_res.value().maxMemoryCost >= 2 * file_size); + + // The estimate is driven by file_size * 2 (compact int8 upload path: + // no fp32 expansion staging). For a compressed (e.g. int8) file the + // peak is 2 * int8_file — NOT the fp32-expanded footprint that the + // old hnswlib-based estimate used. + const uint64_t int8_file = static_cast(nb) * dim * sizeof(int8_t); + auto compressed_res = knowhere::IndexStaticFaced::EstimateLoadResource( + knowhere::IndexEnum::INDEX_GPU_HNSW, version, int8_file, nb, dim, params); + REQUIRE(compressed_res.has_value()); + REQUIRE(compressed_res.value().maxMemoryCost >= 2 * int8_file); + + // When row/dim metadata is missing (num_rows arrives as 0 from the load + // info), the estimate still works because it depends only on file_size. + auto norows_res = knowhere::IndexStaticFaced::EstimateLoadResource( + knowhere::IndexEnum::INDEX_GPU_HNSW, version, int8_file, 0, dim, params); + REQUIRE(norows_res.has_value()); + REQUIRE(norows_res.value().maxMemoryCost >= 2 * int8_file); + + // A plain CPU HNSW keeps its data resident: memoryCost is non-zero and it + // does not opt into the peak field (maxMemoryCost stays 0 -> Milvus falls + // back to its 2*memoryCost heuristic). + auto cpu_res = knowhere::IndexStaticFaced::EstimateLoadResource( + knowhere::IndexEnum::INDEX_HNSW, version, file_size, nb, dim, params); + REQUIRE(cpu_res.has_value()); + REQUIRE(cpu_res.value().memoryCost > 0); + REQUIRE(cpu_res.value().maxMemoryCost == 0); + } + + // Codex coverage #5: CUDA-upload fault injection during Deserialize(). Using + // the test-only GpuHnswUploadFaultInjection hook, force the eager GPU upload + // to fail both at the very first wrapped CUDA call (immediate) and mid-upload + // (after several allocations have already succeeded, exercising the device + // index destructor's partial-allocation cleanup). Assert: (1) the load + // reports failure, (2) nothing is half-published (a search while re-armed + // errors cleanly rather than using a partial index or crashing), (3) VRAM + // returns to baseline after the failed loads + node teardown (no leak), and + // (4) a subsequent fault-free load succeeds and searches correctly. + SECTION("CUDA upload fault injection: load fails, nothing half-published, retry succeeds") { + auto json = sq_json(knowhere::metric::L2, "fp16", 10); + json.erase(knowhere::indexparam::SQ_TYPE); // plain fp32 HNSW + auto train_ds = GenDataSet(nb, dim, seed); + auto query_ds = GenDataSet(nq, dim, seed + 3); + + auto cpu_idx = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_HNSW, version).value(); + REQUIRE(cpu_idx.Build(train_ds, json) == knowhere::Status::success); + knowhere::BinarySet bs; + cpu_idx.Serialize(bs); + + size_t free_before = 0, total = 0; + REQUIRE(cudaMemGetInfo(&free_before, &total) == cudaSuccess); + + // fail_at = 1 fails the first cudaMalloc (immediate); fail_at = 4 fails + // after several buffers are already on the device (mid-upload cleanup). + for (int fail_at : {1, 4}) { + CAPTURE(fail_at); + auto gpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version) + .value(); + + faiss::gpu::GpuHnswUploadFaultInjection::arm(fail_at); + auto st = gpu_idx.Deserialize(bs); + faiss::gpu::GpuHnswUploadFaultInjection::disarm(); + // (1) load fails rather than reporting a silently-broken segment. + REQUIRE(st != knowhere::Status::success); + + // (2) no half-published index: re-arm so the lazy-upload retry inside + // Search() also faults, proving the failed load left no usable index + // behind and the search path errors cleanly (no crash / no results). + faiss::gpu::GpuHnswUploadFaultInjection::arm(fail_at); + auto r = gpu_idx.Search(query_ds, json, nullptr); + faiss::gpu::GpuHnswUploadFaultInjection::disarm(); + REQUIRE(!r.has_value()); + } + + // (3) no leak: after the failed loads and node teardown, free VRAM must + // return to ~baseline (allow slack for allocator/context fragmentation). + REQUIRE(cudaDeviceSynchronize() == cudaSuccess); + size_t free_after = 0; + REQUIRE(cudaMemGetInfo(&free_after, &total) == cudaSuccess); + const size_t slack = static_cast(64) * 1024 * 1024; // 64 MiB + REQUIRE(free_after + slack >= free_before); + + // (4) retry with the fault disarmed succeeds and returns valid results. + faiss::gpu::GpuHnswUploadFaultInjection::disarm(); + auto gpu_ok = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version) + .value(); + REQUIRE(gpu_ok.Deserialize(bs) == knowhere::Status::success); + auto r_ok = gpu_ok.Search(query_ds, json, nullptr); + REQUIRE(r_ok.has_value()); + auto gt = knowhere::BruteForce::Search(train_ds, query_ds, json, nullptr); + REQUIRE(gt.has_value()); + REQUIRE(GetKNNRecall(*gt.value(), *r_ok.value()) >= 0.8f); + } +} + +// ============================================================================ +// Full dtype x metric coverage for the GPU HNSW layer-0 search kernel. +// +// One TEST_CASE per dtype so each runs in its OWN process invocation (a CUDA +// illegal-access in one kernel instantiation corrupts the context and aborts +// the process; separate tags let the harness run every dtype independently and +// still collect per-dtype results instead of losing everything after the first +// crash). Every dtype is exercised for L2, IP and COSINE. +// +// All cases use the SAME config as the fp16 P1 regression that v82 OOBs on +// (nb=4000, dim=64, M=16, ef=200) so the fp32 case below is the control: if +// fp32 (the kernel) is clean here while fp16 +// (<__half,float,false>) faults, the defect is specific to the low-precision +// storage kernel; if fp32 also faults, it is config- rather than dtype-driven. +// dim=64 is divisible by 4, so the int8 IP/COSINE cases take the DP4A path. +namespace { +constexpr int64_t kDtypeNb = 4000; +constexpr int64_t kDtypeNq = 200; +constexpr int64_t kDtypeDim = 64; +constexpr int64_t kDtypeSeed = 42; +constexpr int64_t kDtypeTopk = 10; + +knowhere::Json +dtype_json(const std::string& metric, const std::string& sq_type) { + knowhere::Json json; + json[knowhere::meta::DIM] = kDtypeDim; + json[knowhere::meta::METRIC_TYPE] = metric; + json[knowhere::meta::TOPK] = kDtypeTopk; + json[knowhere::indexparam::HNSW_M] = 16; + json[knowhere::indexparam::EFCONSTRUCTION] = 200; + json[knowhere::indexparam::EF] = 200; + if (!sq_type.empty()) { + json[knowhere::indexparam::SQ_TYPE] = sq_type; + } + return json; +} + +// Build a CPU index of `index_type` on data type T, serialize, deserialize as +// GPU_HNSW, search, and return recall vs a brute-force reference on T. Shared by +// every per-dtype case so the only variables are the data type, the storage +// index type and the metric. +template +float +gpu_dtype_recall(const std::string& index_type, const knowhere::Json& json, int version) { + auto train_ds = knowhere::ConvertToDataTypeIfNeeded(GenDataSet(kDtypeNb, kDtypeDim, kDtypeSeed)); + auto query_ds = knowhere::ConvertToDataTypeIfNeeded(GenDataSet(kDtypeNq, kDtypeDim, kDtypeSeed + 2)); + + auto cpu_idx = knowhere::IndexFactory::Instance().Create(index_type, version).value(); + REQUIRE(cpu_idx.Build(train_ds, json) == knowhere::Status::success); + knowhere::BinarySet bs; + REQUIRE(cpu_idx.Serialize(bs) == knowhere::Status::success); + + auto gpu_idx = knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version).value(); + REQUIRE(gpu_idx.Deserialize(bs) == knowhere::Status::success); + + auto results = gpu_idx.Search(query_ds, json, nullptr); + REQUIRE(results.has_value()); + auto gt = knowhere::BruteForce::Search(train_ds, query_ds, json, nullptr); + REQUIRE(gt.has_value()); + return GetKNNRecall(*gt.value(), *results.value()); +} + +// Filtered-search recall helper. Builds a CPU index of `index_type` on T, +// uploads to GPU_HNSW, then for a given filter `ratio` searches with a random +// delete bitset. Returns {gpu_recall, cpu_recall} where both are measured vs a +// brute-force reference that uses the SAME bitset, so CPU HNSW is the parity +// baseline (not brute force alone). Also asserts no returned GPU id is a +// filtered bit. ratio == 0 exercises the HAS_FILTER=false fast path (empty +// view => data()==nullptr => unfiltered). +template +std::pair +gpu_filtered_recall(const std::string& index_type, const knowhere::Json& json, int version, float ratio, + uint32_t bitset_seed) { + auto train_ds = knowhere::ConvertToDataTypeIfNeeded(GenDataSet(kDtypeNb, kDtypeDim, kDtypeSeed)); + auto query_ds = knowhere::ConvertToDataTypeIfNeeded(GenDataSet(kDtypeNq, kDtypeDim, kDtypeSeed + 2)); + + auto cpu_idx = knowhere::IndexFactory::Instance().Create(index_type, version).value(); + REQUIRE(cpu_idx.Build(train_ds, json) == knowhere::Status::success); + knowhere::BinarySet bs; + REQUIRE(cpu_idx.Serialize(bs) == knowhere::Status::success); + + auto gpu_idx = knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version).value(); + REQUIRE(gpu_idx.Deserialize(bs) == knowhere::Status::success); + + std::vector bitset_data((kDtypeNb + 7) / 8, 0); + int64_t filtered_count = 0; + std::mt19937 rng(bitset_seed); + std::uniform_real_distribution u(0.0f, 1.0f); + for (int64_t r = 0; r < kDtypeNb; r++) { + if (u(rng) < ratio) { + bitset_data[r / 8] |= static_cast(1 << (r % 8)); + filtered_count++; + } + } + const bool has_filter = filtered_count > 0; + knowhere::BitsetView bitset = + has_filter ? knowhere::BitsetView(bitset_data.data(), kDtypeNb, filtered_count) : knowhere::BitsetView(); + + auto gpu_res = gpu_idx.Search(query_ds, json, bitset); + REQUIRE(gpu_res.has_value()); + const int64_t* gids = gpu_res.value()->GetIds(); + for (int64_t i = 0; i < kDtypeNq * kDtypeTopk; i++) { + const int64_t id = gids[i]; + if (id < 0) { + continue; + } + REQUIRE(id < kDtypeNb); + if (has_filter) { + REQUIRE_FALSE((bitset_data[id / 8] & (1 << (id % 8))) != 0); + } + } + + auto cpu_res = cpu_idx.Search(query_ds, json, bitset); + REQUIRE(cpu_res.has_value()); + auto gt = knowhere::BruteForce::Search(train_ds, query_ds, json, bitset); + REQUIRE(gt.has_value()); + return {GetKNNRecall(*gt.value(), *gpu_res.value()), GetKNNRecall(*gt.value(), *cpu_res.value())}; +} +} // namespace + +// FP32 flat storage -> generic kernel. The config control. +TEST_CASE("GPU HNSW dtype: FP32 (L2/IP/COSINE)", "[gpu_hnsw_dtype_fp32]") { + const auto version = GenTestVersionList(); + const auto metric = GENERATE(knowhere::metric::L2, knowhere::metric::IP, knowhere::metric::COSINE); + CAPTURE(metric); + float recall = gpu_dtype_recall(knowhere::IndexEnum::INDEX_HNSW, dtype_json(metric, ""), version); + REQUIRE(recall >= 0.85f); +} + +// FP16 native 2-byte storage -> <__half,float,false> kernel. Built via HNSW_SQ +// on the fp32 node with SQ_TYPE=fp16 so the device selects the native half +// representation (the path the mock-wrapper fp16 tests do NOT reach). +TEST_CASE("GPU HNSW dtype: FP16 native (L2/IP/COSINE)", "[gpu_hnsw_dtype_fp16]") { + const auto version = GenTestVersionList(); + const auto metric = GENERATE(knowhere::metric::L2, knowhere::metric::IP, knowhere::metric::COSINE); + CAPTURE(metric); + float recall = + gpu_dtype_recall(knowhere::IndexEnum::INDEX_HNSW_SQ, dtype_json(metric, "fp16"), version); + REQUIRE(recall >= 0.85f); +} + +// BF16 native 2-byte storage -> <__nv_bfloat16,float,false> kernel. +TEST_CASE("GPU HNSW dtype: BF16 native (L2/IP/COSINE)", "[gpu_hnsw_dtype_bf16]") { + const auto version = GenTestVersionList(); + const auto metric = GENERATE(knowhere::metric::L2, knowhere::metric::IP, knowhere::metric::COSINE); + CAPTURE(metric); + float recall = + gpu_dtype_recall(knowhere::IndexEnum::INDEX_HNSW_SQ, dtype_json(metric, "bf16"), version); + // bf16 has a 7-bit mantissa (coarser than fp16); slightly looser floor. + REQUIRE(recall >= 0.80f); +} + +// SQ8 storage -> decoded to fp32 on upload -> generic +// kernel (distinct from the native half/bf16 paths above). +TEST_CASE("GPU HNSW dtype: SQ8 decoded (L2/IP/COSINE)", "[gpu_hnsw_dtype_sq8]") { + const auto version = GenTestVersionList(); + const auto metric = GENERATE(knowhere::metric::L2, knowhere::metric::IP, knowhere::metric::COSINE); + CAPTURE(metric); + float recall = + gpu_dtype_recall(knowhere::IndexEnum::INDEX_HNSW_SQ, dtype_json(metric, "sq8"), version); + // SQ8 is coarse (8-bit per component); conservative floor. + REQUIRE(recall >= 0.65f); +} + +// INT8 native 1-byte storage. IP/COSINE with dim%4==0 take the DP4A kernel +// (); L2 takes the generic decoded path (). +TEST_CASE("GPU HNSW dtype: INT8 native (L2/IP/COSINE)", "[gpu_hnsw_dtype_int8]") { + const auto version = GenTestVersionList(); + const auto metric = GENERATE(knowhere::metric::L2, knowhere::metric::IP, knowhere::metric::COSINE); + CAPTURE(metric); + float recall = gpu_dtype_recall(knowhere::IndexEnum::INDEX_HNSW, dtype_json(metric, ""), version); + REQUIRE(recall >= 0.85f); +} + +// v84 regression: a CPU HNSW_SQ index must load into the *GPU_HNSW_SQ* node. +// GpuHnswSQIndexNode inherits Deserialize from GpuHnswIndexNode, which re-files +// the incoming payload under this node's own Type(). The SQ node's Type() is +// INDEX_GPU_HNSW_SQ and the base loader reads the payload via +// GetByName(Type()); the pre-v84 code aliased under a hardcoded +// INDEX_GPU_HNSW, so the SQ lookup missed and the load failed. Loading into the +// SQ node (not the base GPU_HNSW node) is what exercises the fix. +TEST_CASE("GPU HNSW_SQ node deserialize alias", "[gpu_hnsw_sq_alias]") { + const auto version = GenTestVersionList(); + const auto metric = GENERATE(knowhere::metric::L2, knowhere::metric::COSINE); + CAPTURE(metric); + + auto json = dtype_json(metric, "sq8"); + auto train_ds = GenDataSet(kDtypeNb, kDtypeDim, kDtypeSeed); + auto query_ds = GenDataSet(kDtypeNq, kDtypeDim, kDtypeSeed + 2); + + auto cpu_idx = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_HNSW_SQ, version).value(); + REQUIRE(cpu_idx.Build(train_ds, json) == knowhere::Status::success); + knowhere::BinarySet bs; + REQUIRE(cpu_idx.Serialize(bs) == knowhere::Status::success); + + // Load into the GPU_HNSW_SQ node specifically (not the base GPU_HNSW node). + auto gpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW_SQ, version) + .value(); + REQUIRE(gpu_idx.Deserialize(bs) == knowhere::Status::success); + + auto results = gpu_idx.Search(query_ds, json, nullptr); + REQUIRE(results.has_value()); + auto gt = knowhere::BruteForce::Search(train_ds, query_ds, json, nullptr); + REQUIRE(gt.has_value()); + REQUIRE(GetKNNRecall(*gt.value(), *results.value()) >= 0.65f); +} + +// v84: IP/COSINE distances must be returned as the true similarity (larger == +// closer), matching faiss's METRIC_INNER_PRODUCT contract. The GPU kernel keeps +// a negated, min-first score internally and negates it back on copy-out; +// Knowhere no longer re-negates. This checks the numeric distance value (sign + +// magnitude) against the similarity recomputed from the returned neighbor, which +// recall-only checks cannot catch. A re-introduced double negation would flip +// every score negative and fail here. +TEST_CASE("GPU HNSW IP/COSINE distance sign", "[gpu_hnsw_ip_sign]") { + const auto version = GenTestVersionList(); + const auto metric = GENERATE(knowhere::metric::IP, knowhere::metric::COSINE); + CAPTURE(metric); + const bool is_cosine = (metric == knowhere::metric::COSINE); + + auto json = dtype_json(metric, ""); + auto train_ds = GenDataSet(kDtypeNb, kDtypeDim, kDtypeSeed); + auto query_ds = GenDataSet(kDtypeNq, kDtypeDim, kDtypeSeed + 2); + const auto* train = reinterpret_cast(train_ds->GetTensor()); + const auto* query = reinterpret_cast(query_ds->GetTensor()); + + auto cpu_idx = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_HNSW, version).value(); + REQUIRE(cpu_idx.Build(train_ds, json) == knowhere::Status::success); + knowhere::BinarySet bs; + REQUIRE(cpu_idx.Serialize(bs) == knowhere::Status::success); + + auto gpu_idx = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version).value(); + REQUIRE(gpu_idx.Deserialize(bs) == knowhere::Status::success); + + auto results = gpu_idx.Search(query_ds, json, nullptr); + REQUIRE(results.has_value()); + const int64_t* ids = results.value()->GetIds(); + const float* dist = results.value()->GetDistance(); + + int positive_top1 = 0; + for (int64_t q = 0; q < kDtypeNq; q++) { + // Distances must be sorted best-first, i.e. descending for a + // larger-is-closer similarity metric. + for (int64_t j = 1; j < kDtypeTopk; j++) { + REQUIRE(dist[q * kDtypeTopk + j] <= dist[q * kDtypeTopk + j - 1] + 1e-4f); + } + // Recompute the true similarity between the query and the returned + // top-1 neighbor and compare to the reported distance. + int64_t nn = ids[q * kDtypeTopk]; + REQUIRE(nn >= 0); + REQUIRE(nn < kDtypeNb); + double dot = 0.0, qn = 0.0, dn = 0.0; + for (int64_t d = 0; d < kDtypeDim; d++) { + double qv = query[q * kDtypeDim + d]; + double dv = train[nn * kDtypeDim + d]; + dot += qv * dv; + qn += qv * qv; + dn += dv * dv; + } + double expected = is_cosine ? dot / (std::sqrt(qn) * std::sqrt(dn) + 1e-12) : dot; + double got = dist[q * kDtypeTopk]; + REQUIRE(std::fabs(got - expected) <= 1e-2 * (std::fabs(expected) + 1.0)); + if (got > 0.0) + positive_top1++; + } + // For this data the nearest neighbor's similarity is positive for the vast + // majority of queries; a sign inversion would make all of them negative. + REQUIRE(positive_top1 > kDtypeNq * 0.8); +} + +// v84: every requested result slot must be a real neighbor (no UINT64_MAX/-1 id +// or FLT_MAX sentinel). ef is guaranteed >= k: Knowhere's config bumps an unset +// ef to max(k, 16) and the GPU kernel additionally auto-clamps ef=max(ef,k) as +// defense-in-depth. Sentinel leakage into the last (k-ef) slots was the k>ef +// symptom. ef is intentionally left unset here so the requested topk drives it. +TEST_CASE("GPU HNSW returns k valid results (ef>=k)", "[gpu_hnsw_k_valid]") { + const auto version = GenTestVersionList(); + const int64_t topk = 100; + knowhere::Json json; + json[knowhere::meta::DIM] = kDtypeDim; + json[knowhere::meta::METRIC_TYPE] = knowhere::metric::L2; + json[knowhere::meta::TOPK] = topk; + json[knowhere::indexparam::HNSW_M] = 16; + json[knowhere::indexparam::EFCONSTRUCTION] = 200; + // EF intentionally unset: Knowhere sets ef = max(topk, 16) = topk here. + + auto train_ds = GenDataSet(kDtypeNb, kDtypeDim, kDtypeSeed); + auto query_ds = GenDataSet(kDtypeNq, kDtypeDim, kDtypeSeed + 2); + + auto cpu_idx = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_HNSW, version).value(); + REQUIRE(cpu_idx.Build(train_ds, json) == knowhere::Status::success); + knowhere::BinarySet bs; + REQUIRE(cpu_idx.Serialize(bs) == knowhere::Status::success); + + auto gpu_idx = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version).value(); + REQUIRE(gpu_idx.Deserialize(bs) == knowhere::Status::success); + + auto results = gpu_idx.Search(query_ds, json, nullptr); + REQUIRE(results.has_value()); + const int64_t* ids = results.value()->GetIds(); + const float* dist = results.value()->GetDistance(); + for (int64_t i = 0; i < kDtypeNq * topk; i++) { + REQUIRE(ids[i] >= 0); + REQUIRE(ids[i] < kDtypeNb); + REQUIRE(dist[i] < std::numeric_limits::max()); + } +} + +// Blocking prerequisite for GPU filtered search: the kernel tests the delete +// bitset by raw FAISS node id, so correctness hinges on +// GPU node id == storage row offset == BitsetView bit index. +// This proves that mapping empirically with known ordered data and known +// filters, independent of recall. If this ever fails, filtered search silently +// excludes the wrong rows and every downstream filtered test is meaningless. +TEST_CASE("GPU HNSW G-ID mapping: node id == row offset == bitset index", "[gpu_hnsw_id_mapping]") { + const auto version = GenTestVersionList(); + const int64_t nb = 2000; + const int64_t dim = 32; + + knowhere::Json json; + json[knowhere::meta::DIM] = dim; + json[knowhere::meta::METRIC_TYPE] = knowhere::metric::L2; + json[knowhere::meta::TOPK] = 1; + json[knowhere::indexparam::HNSW_M] = 16; + json[knowhere::indexparam::EFCONSTRUCTION] = 200; + json[knowhere::indexparam::EF] = 200; + + // Randomly generated training data; regardless of order, every row is + // trivially its own exact nearest neighbor (L2 self-distance 0), so an + // unfiltered self-query for row i must return id i. + auto train_ds = GenDataSet(nb, dim, 42); + + auto cpu_idx = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_HNSW, version).value(); + REQUIRE(cpu_idx.Build(train_ds, json) == knowhere::Status::success); + knowhere::BinarySet bs; + REQUIRE(cpu_idx.Serialize(bs) == knowhere::Status::success); + + auto gpu_idx = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version).value(); + REQUIRE(gpu_idx.Deserialize(bs) == knowhere::Status::success); + REQUIRE(gpu_idx.Count() == nb); + + // 1) Unfiltered self-query: node id == storage row offset. + { + auto res = gpu_idx.Search(train_ds, json, nullptr); + REQUIRE(res.has_value()); + const int64_t* ids = res.value()->GetIds(); + for (int64_t i = 0; i < nb; i++) { + REQUIRE(ids[i] == i); + } + } + + // 2) Filtered self-query: bit index == storage row offset. Filter every + // third row; a returned id must never be a filtered bit, and a filtered + // row's own id must be excluded from its self-query. + { + std::vector bitset_data((nb + 7) / 8, 0); + int64_t filtered_count = 0; + for (int64_t r = 0; r < nb; r += 3) { + bitset_data[r / 8] |= static_cast(1 << (r % 8)); + filtered_count++; + } + knowhere::BitsetView bitset(bitset_data.data(), nb, filtered_count); + + auto res = gpu_idx.Search(train_ds, json, bitset); + REQUIRE(res.has_value()); + const int64_t* ids = res.value()->GetIds(); + for (int64_t i = 0; i < nb; i++) { + const int64_t id = ids[i]; + if (id == -1) { + continue; + } + const bool id_filtered = (bitset_data[id / 8] & (1 << (id % 8))) != 0; + REQUIRE_FALSE(id_filtered); + const bool query_filtered = (bitset_data[i / 8] & (1 << (i % 8))) != 0; + if (query_filtered) { + REQUIRE(id != i); + } + } + } +} + +// Filtered-search recall/exclusion parity across delete ratios and metrics. +// Exercises the two-tier beam + alpha gate + brute-force fallback. For every +// (metric, ratio) the helper asserts no returned id is a filtered bit; here we +// assert GPU recall tracks the CPU HNSW baseline (both vs the SAME-bitset brute +// force), which is the design's parity requirement. Parity is asserted by +// recall, not bit-identical traversal (the alpha gate is deterministic but not +// CPU encounter-order identical). ratio 0.0 exercises the HAS_FILTER=false fast +// path; 0.3/0.7 cover the mid-range where the alpha-gate order divergence is +// most pronounced; 0.95/0.99 trip the up-front brute-force fallback. +namespace { +void +check_filtered_ratios(const std::string& index_type, const knowhere::Json& json) { + const auto version = GenTestVersionList(); + uint32_t seed = 123; + for (const float ratio : {0.0f, 0.1f, 0.3f, 0.5f, 0.7f, 0.9f, 0.95f, 0.99f}) { + CAPTURE(ratio); + auto [gpu_recall, cpu_recall] = gpu_filtered_recall(index_type, json, version, ratio, seed++); + CAPTURE(gpu_recall); + CAPTURE(cpu_recall); + // Primary parity gate: GPU must not lag CPU HNSW by more than 0.05. + REQUIRE(gpu_recall >= cpu_recall - 0.05f); + if (ratio >= 0.93f) { + // Up-front brute force is exact -> recall should be ~1.0. + REQUIRE(gpu_recall >= 0.90f); + } + } +} +} // namespace + +TEST_CASE("GPU HNSW filtered search recall (FP32 L2/IP/COSINE)", "[gpu_hnsw_filtered_recall]") { + const auto metric = GENERATE(knowhere::metric::L2, knowhere::metric::IP, knowhere::metric::COSINE); + CAPTURE(metric); + check_filtered_ratios(knowhere::IndexEnum::INDEX_HNSW, dtype_json(metric, "")); +} + +// INT8 native storage: IP/COSINE with dim%4==0 take the DP4A brute-force + +// graph path; L2 takes the generic decoded path. Confirms the filter/BF kernels +// are threaded through searchHostInt8, not just the fp32 searchHost. +TEST_CASE("GPU HNSW filtered search recall (INT8 native L2/IP/COSINE)", "[gpu_hnsw_filtered_recall_int8]") { + const auto version = GenTestVersionList(); + const auto metric = GENERATE(knowhere::metric::L2, knowhere::metric::IP, knowhere::metric::COSINE); + CAPTURE(metric); + uint32_t seed = 777; + for (const float ratio : {0.0f, 0.5f, 0.9f, 0.99f}) { + CAPTURE(ratio); + auto [gpu_recall, cpu_recall] = gpu_filtered_recall( + knowhere::IndexEnum::INDEX_HNSW, dtype_json(metric, ""), version, ratio, seed++); + CAPTURE(gpu_recall); + CAPTURE(cpu_recall); + REQUIRE(gpu_recall >= cpu_recall - 0.05f); + if (ratio >= 0.93f) { + REQUIRE(gpu_recall >= 0.90f); + } + } +} + +// Edge cases the ratio sweep does not cover: everything filtered (no live rows) +// and k larger than the live-row count. Both must return sentinels for the +// missing slots rather than a filtered id or an error. +TEST_CASE("GPU HNSW filtered search edge cases", "[gpu_hnsw_filtered_edge]") { + const auto version = GenTestVersionList(); + const int64_t nb = 4000; + const int64_t dim = 64; + const int64_t nq = 32; + + auto json = dtype_json(knowhere::metric::L2, ""); + json[knowhere::meta::TOPK] = 10; + auto train_ds = GenDataSet(nb, dim, 42); + auto query_ds = GenDataSet(nq, dim, 44); + + auto cpu_idx = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_HNSW, version).value(); + REQUIRE(cpu_idx.Build(train_ds, json) == knowhere::Status::success); + knowhere::BinarySet bs; + REQUIRE(cpu_idx.Serialize(bs) == knowhere::Status::success); + auto gpu_idx = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version).value(); + REQUIRE(gpu_idx.Deserialize(bs) == knowhere::Status::success); + + // 1) All rows filtered: every result slot must be a sentinel (-1), never a + // filtered id, and the search must not error. + { + std::vector bitset_data((nb + 7) / 8, 0xFF); + knowhere::BitsetView bitset(bitset_data.data(), nb, nb); + auto res = gpu_idx.Search(query_ds, json, bitset); + REQUIRE(res.has_value()); + const int64_t* ids = res.value()->GetIds(); + for (int64_t i = 0; i < nq * 10; i++) { + REQUIRE(ids[i] == -1); + } + } + + // 2) k > live_count: only a handful of rows survive; the first `live` slots + // per query must be live ids, the remainder sentinels. + { + const int64_t live = 3; + std::vector bitset_data((nb + 7) / 8, 0xFF); + for (int64_t r = 0; r < live; r++) { + bitset_data[r / 8] &= static_cast(~(1 << (r % 8))); + } + knowhere::BitsetView bitset(bitset_data.data(), nb, nb - live); + auto res = gpu_idx.Search(query_ds, json, bitset); + REQUIRE(res.has_value()); + const int64_t* ids = res.value()->GetIds(); + for (int64_t q = 0; q < nq; q++) { + for (int64_t j = 0; j < 10; j++) { + const int64_t id = ids[q * 10 + j]; + if (id == -1) { + continue; + } + REQUIRE(id < live); + } + } + } +} + +// Visited-bitmap OOM fix: the layer-0 search processes queries in nq-chunks so +// the grow-only visited bitmap (nq * ceil(N/32) * 4 bytes per scratch slot) is +// bounded regardless of batch size / search concurrency. Chunking must be +// result-invariant: each query is independent, and every per-query buffer +// (queries, entry points, neighbors, distances, visited bitmap, BF worklist) is +// indexed by the chunk-local block index against a caller-offset base pointer. +// This runs the SAME search twice -- once with the default cap (single pass) and +// once with GPU_HNSW_BITMAP_BYTES forced small enough to split the batch into +// many chunks -- and asserts the ids and distances match element-for-element. +namespace { +template +void +check_chunk_parity(const std::string& index_type, const knowhere::Json& json, const knowhere::BitsetView& bitset) { + using Catch::Approx; + const auto version = GenTestVersionList(); + auto train_ds = knowhere::ConvertToDataTypeIfNeeded(GenDataSet(kDtypeNb, kDtypeDim, kDtypeSeed)); + auto query_ds = knowhere::ConvertToDataTypeIfNeeded(GenDataSet(kDtypeNq, kDtypeDim, kDtypeSeed + 2)); + + auto cpu_idx = knowhere::IndexFactory::Instance().Create(index_type, version).value(); + REQUIRE(cpu_idx.Build(train_ds, json) == knowhere::Status::success); + knowhere::BinarySet bs; + REQUIRE(cpu_idx.Serialize(bs) == knowhere::Status::success); + auto gpu_idx = knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version).value(); + REQUIRE(gpu_idx.Deserialize(bs) == knowhere::Status::success); + + // Baseline: default cap => chunk == nq (single pass). + ::unsetenv("GPU_HNSW_BITMAP_BYTES"); + ::unsetenv("GPU_HNSW_BITMAP_MB"); + auto base = gpu_idx.Search(query_ds, json, bitset); + REQUIRE(base.has_value()); + const int64_t n_out = static_cast(kDtypeNq) * kDtypeTopk; + std::vector base_ids(base.value()->GetIds(), base.value()->GetIds() + n_out); + std::vector base_dist(base.value()->GetDistance(), base.value()->GetDistance() + n_out); + + // Force chunking: cap the bitmap to ~2 queries' worth so the batch is split + // into ~nq/2 chunks, exercising the multi-launch offset math and the + // per-chunk bitmap/worklist resets. + const long per_query = static_cast((kDtypeNb + 31) / 32) * 4; + ::setenv("GPU_HNSW_BITMAP_BYTES", std::to_string(2 * per_query).c_str(), 1); + auto chunked = gpu_idx.Search(query_ds, json, bitset); + ::unsetenv("GPU_HNSW_BITMAP_BYTES"); + REQUIRE(chunked.has_value()); + const int64_t* cids = chunked.value()->GetIds(); + const float* cdist = chunked.value()->GetDistance(); + + for (int64_t i = 0; i < n_out; i++) { + REQUIRE(cids[i] == base_ids[i]); + REQUIRE(cdist[i] == Approx(base_dist[i]).epsilon(1e-5)); + } +} + +std::vector +make_random_bitset(float ratio, uint32_t seed, int64_t& filtered_count) { + std::vector bd((kDtypeNb + 7) / 8, 0); + filtered_count = 0; + std::mt19937 rng(seed); + std::uniform_real_distribution u(0.0f, 1.0f); + for (int64_t r = 0; r < kDtypeNb; r++) { + if (u(rng) < ratio) { + bd[r / 8] |= static_cast(1 << (r % 8)); + filtered_count++; + } + } + return bd; +} +} // namespace + +TEST_CASE("GPU HNSW nq-chunking is result-invariant (visited-bitmap OOM fix)", "[gpu_hnsw_chunk_parity]") { + const auto metric = GENERATE(knowhere::metric::L2, knowhere::metric::IP, knowhere::metric::COSINE); + CAPTURE(metric); + + // Unfiltered (HAS_FILTER=false fast path). + SECTION("fp32 unfiltered") { + check_chunk_parity(knowhere::IndexEnum::INDEX_HNSW, dtype_json(metric, ""), + knowhere::BitsetView()); + } + // int8: IP/COSINE with dim%4==0 take the DP4A path, where the layer-0 (int8) + // and upper-layer (fp32) query buffers are distinct, so BOTH chunk offsets + // must be correct; L2 takes the generic decoded path. + SECTION("int8 unfiltered") { + check_chunk_parity(knowhere::IndexEnum::INDEX_HNSW, dtype_json(metric, ""), + knowhere::BitsetView()); + } + // Mid-ratio filter: HAS_FILTER=true two-tier beam + per-chunk worklist reset. + SECTION("fp32 filtered (~30%)") { + int64_t fc = 0; + auto bd = make_random_bitset(0.3f, 7, fc); + knowhere::BitsetView bitset(bd.data(), kDtypeNb, fc); + check_chunk_parity(knowhere::IndexEnum::INDEX_HNSW, dtype_json(metric, ""), bitset); + } + // High-ratio filter: trips the up-front brute-force fallback (also chunked). + SECTION("fp32 filtered (~99% -> up-front BF)") { + int64_t fc = 0; + auto bd = make_random_bitset(0.99f, 9, fc); + knowhere::BitsetView bitset(bd.data(), kDtypeNb, fc); + check_chunk_parity(knowhere::IndexEnum::INDEX_HNSW, dtype_json(metric, ""), bitset); + } } #endif diff --git a/thirdparty/faiss/faiss/gpu/CMakeLists.txt b/thirdparty/faiss/faiss/gpu/CMakeLists.txt index 50740450a..7ab9ffd39 100644 --- a/thirdparty/faiss/faiss/gpu/CMakeLists.txt +++ b/thirdparty/faiss/faiss/gpu/CMakeLists.txt @@ -29,9 +29,11 @@ set(FAISS_GPU_SRC GpuIndexIVF.cu GpuIndexIVFFlat.cu GpuIndexIVFPQ.cu + GpuIndexHNSW.cu GpuIndexIVFScalarQuantizer.cu GpuResources.cpp StandardGpuResources.cpp + impl/GpuHnswTypes.cu impl/BinaryDistance.cu impl/BinaryFlatIndex.cu impl/BroadcastSum.cu @@ -96,10 +98,16 @@ set(FAISS_GPU_HEADERS GpuIndexIVF.h GpuIndexIVFFlat.h GpuIndexIVFPQ.h + GpuIndexHNSW.h GpuIndexIVFScalarQuantizer.h GpuIndicesOptions.h GpuResources.h StandardGpuResources.h + impl/GpuHnswBruteForce.cuh + impl/GpuHnswBuild.cuh + impl/GpuHnswSearch.cuh + impl/GpuHnswSearchKernel.cuh + impl/GpuHnswTypes.h impl/BinaryDistance.cuh impl/BinaryFlatIndex.cuh impl/BroadcastSum.cuh diff --git a/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu b/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu new file mode 100644 index 000000000..0f7b5a0c2 --- /dev/null +++ b/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu @@ -0,0 +1,421 @@ +// @lint-ignore-every LICENSELINT +/** + * Copyright (c) Facebook, Inc. and its affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ +/* + * Copyright (c) 2026, 6sense Insights Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace faiss { +namespace gpu { + +namespace { + +// Upload the filtered-search bitset (deletes / TTL / partition) into the +// scratch slot's device buffer on the search stream. gpu_hnsw_search reads +// sp.bitset_data only to decide the filtered path; the bytes themselves come +// from sc.d_bitset. The host BitsetView bytes referenced by sp.bitset_data +// must stay alive until the stream is synchronized (the caller synchronizes +// before returning, which happens in every search path here). +inline void upload_bitset_if_needed( + GpuHnswSearchScratch& sc, + const GpuHnswSearchParams& sp, + int nq, + cudaStream_t stream) { + if (sp.bitset_data == nullptr || sp.bitset_nbits <= 0) { + return; + } + size_t bytes = static_cast((sp.bitset_nbits + 7) / 8); + sc.ensure_filter(nq, bytes); + GPU_HNSW_CUDA_CHECK(cudaMemcpyAsync( + sc.d_bitset, + sp.bitset_data, + bytes, + cudaMemcpyHostToDevice, + stream)); +} + +} // namespace + +GpuIndexHNSW::GpuIndexHNSW( + GpuResourcesProvider* provider, + int dims, + faiss::MetricType metric, + GpuIndexHNSWConfig config) + : GpuIndex(provider->getResources(), dims, metric, 0.0f, config), + hnswConfig_(config) { + this->is_trained = false; +} + +GpuIndexHNSW::~GpuIndexHNSW() = default; + +void GpuIndexHNSW::copyFrom( + const faiss::cppcontrib::knowhere::IndexHNSW* index) { + FAISS_THROW_IF_NOT_MSG(index, "index must not be null"); + FAISS_THROW_IF_NOT_MSG(index->ntotal > 0, "index must not be empty"); + + DeviceScope scope(config_.device); + + this->d = index->d; + this->metric_type = index->metric_type; + this->ntotal = index->ntotal; + + // Detect cosine from index type (IndexHNSWFlatCosine / IndexHNSWSQCosine + // implement HasInverseL2Norms). + bool is_cosine = + dynamic_cast( + index) != nullptr; + bool use_ip = + is_cosine || (index->metric_type == faiss::METRIC_INNER_PRODUCT); + + is_cosine_ = is_cosine; + use_ip_ = use_ip; + + if (dynamic_cast(index->storage)) { + deviceIndex_ = + from_faiss_hnsw_sq(*index, use_ip, is_cosine, config_.device); + } else { + deviceIndex_ = + from_faiss_hnsw_flat(*index, use_ip, is_cosine, config_.device); + } + + this->is_trained = true; +} + +void GpuIndexHNSW::copyFromWithMetric( + const faiss::cppcontrib::knowhere::IndexHNSW* index, + bool use_ip, + bool is_cosine) { + FAISS_THROW_IF_NOT_MSG(index, "index must not be null"); + FAISS_THROW_IF_NOT_MSG(index->ntotal > 0, "index must not be empty"); + + DeviceScope scope(config_.device); + + this->d = index->d; + this->metric_type = index->metric_type; + this->ntotal = index->ntotal; + + is_cosine_ = is_cosine; + use_ip_ = use_ip; + + if (dynamic_cast(index->storage)) { + deviceIndex_ = + from_faiss_hnsw_sq(*index, use_ip, is_cosine, config_.device); + } else { + deviceIndex_ = + from_faiss_hnsw_flat(*index, use_ip, is_cosine, config_.device); + } + + this->is_trained = true; +} + +void GpuIndexHNSW::reset() { + deviceIndex_.reset(); + this->ntotal = 0; + this->is_trained = false; + is_cosine_ = false; + use_ip_ = false; +} + +void GpuIndexHNSW::setSearchParams(const GpuHnswSearchParams& params) const { + std::lock_guard lock(searchParamsMutex_); + directSearchParams_ = params; + hasDirectSearchParams_ = true; +} + +bool GpuIndexHNSW::addImplRequiresIDs_() const { + return false; +} + +void GpuIndexHNSW::addImpl_(idx_t, const float*, const idx_t*) { + FAISS_THROW_MSG( + "GpuIndexHNSW does not support add(). " + "Build on CPU with IndexHNSW, then call copyFrom()."); +} + +void GpuIndexHNSW::searchImpl_( + idx_t n, + const float* x, + int k, + float* distances, + idx_t* labels, + const SearchParameters* search_params) const { + FAISS_THROW_IF_NOT_MSG( + this->is_trained && deviceIndex_, + "Index not loaded. Call copyFrom() first."); + FAISS_THROW_IF_NOT_MSG(n > 0, "n must be > 0"); + + auto& idx = *deviceIndex_; + + GpuHnswSearchParams sp; + bool got_params = false; + + // Prefer direct params set via setSearchParams() — avoids dynamic_cast. + // These are sticky: they stay in effect for every subsequent search until + // overwritten, so a later search never silently falls back to defaults. + { + std::lock_guard lock(searchParamsMutex_); + if (hasDirectSearchParams_) { + sp = directSearchParams_; + got_params = true; + } + } + + // Fallback: try dynamic_cast from SearchParameters. + if (!got_params && search_params) { + auto* params = + dynamic_cast(search_params); + if (params) { + sp.ef = params->ef; + sp.search_width = params->search_width; + sp.max_iterations = params->max_iterations; + sp.thread_block_size = params->thread_block_size; + got_params = true; + } + } + + ScratchPoolGuard guard(*idx.scratch_pool); + auto* slot = guard.get(); + auto& sc = slot->scratch; + cudaStream_t stream = slot->stream; + + int nq = static_cast(n); + int dim = static_cast(idx.dim); + sc.ensure(nq, k, dim, static_cast(idx.n_rows)); + + // The parent GpuIndex::search copied host queries into `x` (and will later + // copy outputs back) on the resources' default stream. Our scratch-pool + // slot has its own private stream, so make it wait for that H2D/D2D copy to + // complete before we read `x` — otherwise the search can race ahead and + // operate on partially-written query vectors. + cudaStream_t defaultStream = resources_->getDefaultStream(config_.device); + cudaEvent_t inputReady; + GPU_HNSW_CUDA_CHECK( + cudaEventCreateWithFlags(&inputReady, cudaEventDisableTiming)); + GPU_HNSW_CUDA_CHECK(cudaEventRecord(inputReady, defaultStream)); + GPU_HNSW_CUDA_CHECK(cudaStreamWaitEvent(stream, inputReady, 0)); + GPU_HNSW_CUDA_CHECK(cudaEventDestroy(inputReady)); + + // D2D: query vectors (GpuIndex::search passes device pointers) + GPU_HNSW_CUDA_CHECK(cudaMemcpyAsync( + sc.d_queries, + x, + static_cast(nq) * dim * sizeof(float), + cudaMemcpyDeviceToDevice, + stream)); + + gpu_hnsw_search(stream, sp, idx, sc, nq, k); + + // D2D: distances (output is a device pointer from GpuIndex::search) + GPU_HNSW_CUDA_CHECK(cudaMemcpyAsync( + distances, + sc.d_distances, + static_cast(nq) * k * sizeof(float), + cudaMemcpyDeviceToDevice, + stream)); + + // Labels: D2H stage (uint64_t→idx_t conversion), then H2D back + auto tmp = std::make_unique(nq * k); + GPU_HNSW_CUDA_CHECK(cudaMemcpyAsync( + tmp.get(), + sc.d_neighbors, + static_cast(nq) * k * sizeof(uint64_t), + cudaMemcpyDeviceToHost, + stream)); + GPU_HNSW_CUDA_CHECK(cudaStreamSynchronize(stream)); + + auto h_labels = std::make_unique(nq * k); + for (int i = 0; i < nq * k; i++) { + h_labels[i] = (tmp[i] == UINT64_MAX) ? -1 : static_cast(tmp[i]); + } + + GPU_HNSW_CUDA_CHECK(cudaMemcpyAsync( + labels, + h_labels.get(), + static_cast(nq) * k * sizeof(idx_t), + cudaMemcpyHostToDevice, + stream)); + GPU_HNSW_CUDA_CHECK(cudaStreamSynchronize(stream)); +} + +void GpuIndexHNSW::searchHost( + idx_t n, + const float* x_host, + int k, + float* distances_host, + idx_t* labels_host, + const GpuHnswSearchParams& sp) const { + FAISS_THROW_IF_NOT_MSG( + this->is_trained && deviceIndex_, + "Index not loaded. Call copyFrom() first."); + FAISS_THROW_IF_NOT_MSG(n > 0, "n must be > 0"); + + GPU_HNSW_CUDA_CHECK(cudaSetDevice(config_.device)); + DeviceScope scope(config_.device); + auto& idx = *deviceIndex_; + + ScratchPoolGuard guard(*idx.scratch_pool); + auto* slot = guard.get(); + auto& sc = slot->scratch; + cudaStream_t stream = slot->stream; + + int nq = static_cast(n); + int dim = static_cast(idx.dim); + sc.ensure(nq, k, dim, static_cast(idx.n_rows)); + + GPU_HNSW_CUDA_CHECK(cudaMemcpyAsync( + sc.d_queries, + x_host, + static_cast(nq) * dim * sizeof(float), + cudaMemcpyDefault, + stream)); + + upload_bitset_if_needed(sc, sp, nq, stream); + + gpu_hnsw_search(stream, sp, idx, sc, nq, k); + + GPU_HNSW_CUDA_CHECK(cudaMemcpyAsync( + distances_host, + sc.d_distances, + static_cast(nq) * k * sizeof(float), + cudaMemcpyDeviceToHost, + stream)); + + auto tmp = std::make_unique(nq * k); + GPU_HNSW_CUDA_CHECK(cudaMemcpyAsync( + tmp.get(), + sc.d_neighbors, + static_cast(nq) * k * sizeof(uint64_t), + cudaMemcpyDeviceToHost, + stream)); + + GPU_HNSW_CUDA_CHECK(cudaStreamSynchronize(stream)); + + for (int i = 0; i < nq * k; i++) { + labels_host[i] = + (tmp[i] == UINT64_MAX) ? -1 : static_cast(tmp[i]); + } +} + +void GpuIndexHNSW::searchHostInt8( + idx_t n, + const int8_t* x_host, + int k, + float* distances_host, + idx_t* labels_host, + const GpuHnswSearchParams& sp) const { + FAISS_THROW_IF_NOT_MSG( + this->is_trained && deviceIndex_, + "Index not loaded. Call copyFrom() first."); + FAISS_THROW_IF_NOT_MSG(n > 0, "n must be > 0"); + + auto& idx = *deviceIndex_; + // If dim is not divisible by 4, DP4A cannot be used; fall back to fp32 path. + if (idx.dim % 4 != 0) { + auto fp32_fallback = std::make_unique( + static_cast(n) * idx.dim); + for (int64_t i = 0; i < static_cast(n) * idx.dim; i++) { + fp32_fallback[i] = static_cast(x_host[i]); + } + searchHost(n, fp32_fallback.get(), k, distances_host, labels_host, sp); + return; + } + + GPU_HNSW_CUDA_CHECK(cudaSetDevice(config_.device)); + DeviceScope scope(config_.device); + + ScratchPoolGuard guard(*idx.scratch_pool); + auto* slot = guard.get(); + auto& sc = slot->scratch; + cudaStream_t stream = slot->stream; + + int nq = static_cast(n); + int dim = static_cast(idx.dim); + int64_t nelem = static_cast(nq) * dim; + + sc.ensure(nq, k, dim, static_cast(idx.n_rows), /*use_i8_queries=*/true); + + // Upload int8 queries directly — dataset on GPU is already in signed int8 + // (upload_int8_dataset applies codes[i]-128, reversing FAISS's +128 bias, + // yielding the original signed user values). Queries arrive as the same + // signed int8 user values; no shift is needed. + GPU_HNSW_CUDA_CHECK(cudaMemcpyAsync( + sc.d_queries_i8, + x_host, + nelem * sizeof(int8_t), + cudaMemcpyHostToDevice, + stream)); + + // Also upload fp32 queries to d_queries for upper-layer greedy search. + auto fp32_q = std::make_unique(nelem); + for (int64_t i = 0; i < nelem; i++) { + fp32_q[i] = static_cast(x_host[i]); + } + GPU_HNSW_CUDA_CHECK(cudaMemcpyAsync( + sc.d_queries, + fp32_q.get(), + nelem * sizeof(float), + cudaMemcpyHostToDevice, + stream)); + + // Synchronize to ensure host buffers (fp32_q) are not freed + // while the async copies are in flight. + GPU_HNSW_CUDA_CHECK(cudaStreamSynchronize(stream)); + + upload_bitset_if_needed(sc, sp, nq, stream); + + gpu_hnsw_search(stream, sp, idx, sc, nq, k); + + GPU_HNSW_CUDA_CHECK(cudaMemcpyAsync( + distances_host, + sc.d_distances, + static_cast(nq) * k * sizeof(float), + cudaMemcpyDeviceToHost, + stream)); + + auto tmp = std::make_unique(nq * k); + GPU_HNSW_CUDA_CHECK(cudaMemcpyAsync( + tmp.get(), + sc.d_neighbors, + static_cast(nq) * k * sizeof(uint64_t), + cudaMemcpyDeviceToHost, + stream)); + + GPU_HNSW_CUDA_CHECK(cudaStreamSynchronize(stream)); + + for (int i = 0; i < nq * k; i++) { + labels_host[i] = + (tmp[i] == UINT64_MAX) ? -1 : static_cast(tmp[i]); + } +} + +} // namespace gpu +} // namespace faiss diff --git a/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h b/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h new file mode 100644 index 000000000..f374e3fe0 --- /dev/null +++ b/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h @@ -0,0 +1,193 @@ +// @lint-ignore-every LICENSELINT +/** + * Copyright (c) Facebook, Inc. and its affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ +/* + * Copyright (c) 2026, 6sense Insights Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include + +#include +#include + +namespace faiss { + +namespace cppcontrib { +namespace knowhere { +struct IndexHNSW; +} // namespace knowhere +} // namespace cppcontrib + +namespace gpu { + +struct GpuHnswDeviceIndex; + +struct GpuIndexHNSWConfig : public GpuIndexConfig {}; + +struct SearchParametersGpuHNSW : SearchParameters { + /// Search ef — number of candidates maintained during search. + /// Higher values improve recall at the cost of speed. + int ef = 200; + + /// Number of candidates to expand per iteration. + int search_width = 4; + + /// Maximum search iterations (0 = auto). + int max_iterations = 0; + + /// Thread block size (0 = auto, default 128). + int thread_block_size = 0; +}; + +/// GPU implementation of HNSW search. +/// +/// This index type does NOT build an HNSW graph on GPU — it takes a +/// CPU-built faiss::IndexHNSW (Flat or SQ storage), converts the graph +/// to a GPU-friendly dense format, and runs the search on GPU using a +/// parallel beam search kernel. +/// +/// Supports L2, inner product, and cosine metrics. +/// Supports float32, fp16, bf16, and int8 (QT_8bit_direct_signed) data. +/// Low-precision formats (int8/fp16/bf16) are kept in their native on-device +/// byte layout and up-converted to fp32 per element inside the search kernel. +/// +/// Two search entry points exist: +/// - searchHost(): the path Knowhere uses. Host in/out pointers, a single +/// device sync at the end. Preferred for production. +/// - searchImpl_(): the faiss-standard GpuIndex::search() override. It does a +/// D2H copy of labels, a CPU uint64->idx_t conversion, then an H2D copy +/// back, with a stream sync on each side. Correct but not latency-optimal; +/// a GPU-side label-conversion kernel to avoid the round-trip is a +/// documented follow-up. Knowhere does not use this path. +struct GpuIndexHNSW : public GpuIndex { + public: + GpuIndexHNSW( + GpuResourcesProvider* provider, + int dims, + faiss::MetricType metric = faiss::METRIC_L2, + GpuIndexHNSWConfig config = GpuIndexHNSWConfig()); + + ~GpuIndexHNSW() override; + + /// Load an HNSW index from CPU to GPU. + /// The CPU index must have been built and trained already. + /// Supports IndexHNSWFlat (float32) and IndexHNSWSQ + /// (QT_8bit_direct_signed for INT8, QT_fp16/QT_bf16 kept native, or + /// dequantized for other SQ types). + void copyFrom(const faiss::cppcontrib::knowhere::IndexHNSW* index); + + /// Like copyFrom(), but with the metric interpretation supplied by the + /// caller instead of being detected from the index type. + /// \param use_ip treat the metric as inner product + /// \param is_cosine the storage carries cosine (inverse L2) norms + void copyFromWithMetric( + const faiss::cppcontrib::knowhere::IndexHNSW* index, + bool use_ip, + bool is_cosine); + + void reset() override; + + /// Set search parameters directly, bypassing SearchParameters. + /// Mutex-guarded. The params are sticky: once set they apply to every + /// subsequent search() until overwritten by another setSearchParams() + /// call. Prefer passing SearchParametersGpuHNSW per-search (or using + /// searchHost) when different concurrent searches need different params. + void setSearchParams(const GpuHnswSearchParams& params) const; + + /// Search with host pointers directly, bypassing GpuIndex::search. + /// All input/output pointers must be host memory. + /// This avoids the GpuIndex::search_ex temp allocation chain + /// which can cause SIGSEGV from pointer lifetime issues. + /// This is the preferred entry point (single device sync); prefer it + /// over the searchImpl_/GpuIndex::search() path, which round-trips the + /// labels D2H then H2D. + /// + /// Distance convention: inner-product and cosine metrics return the true + /// similarity score (larger == more similar), matching faiss's + /// METRIC_INNER_PRODUCT contract. Internally the kernel keeps a negated + /// (smaller == more similar) score for a uniform min-first ordering and + /// negates it back on copy-out. L2 distances are returned as-is. + void searchHost( + idx_t n, + const float* x_host, + int k, + float* distances_host, + idx_t* labels_host, + const GpuHnswSearchParams& params) const; + + /// Search with int8 host query vectors using the native DP4A path. + /// The queries are the caller's signed int8 values and are uploaded + /// verbatim — no bias/shift is applied, matching the dataset upload which + /// already reverses FAISS's +128 SQ bias. When dim % 4 != 0 (DP4A requires + /// groups of four int8 lanes) this transparently falls back to the fp32 + /// searchHost() path. Same IP/cosine distance convention as searchHost() + /// (true similarity returned). All input/output pointers must be host + /// memory. + void searchHostInt8( + idx_t n, + const int8_t* x_host, + int k, + float* distances_host, + idx_t* labels_host, + const GpuHnswSearchParams& params) const; + + /// Metric interpretation captured at copyFrom()/copyFromWithMetric() time + /// from the FAISS index type, so callers can post-process results (cosine + /// query normalization, IP handling) without re-deriving the metric from a + /// per-search config — the config may omit metric_type and default to L2, + /// which would silently mishandle a cosine/IP index. + bool isCosine() const { + return is_cosine_; + } + bool useInnerProduct() const { + return use_ip_; + } + + protected: + bool addImplRequiresIDs_() const override; + + void addImpl_(idx_t n, const float* x, const idx_t* ids) override; + + void searchImpl_( + idx_t n, + const float* x, + int k, + float* distances, + idx_t* labels, + const SearchParameters* search_params) const override; + + private: + GpuIndexHNSWConfig hnswConfig_; + + std::unique_ptr deviceIndex_; + + mutable std::mutex searchParamsMutex_; + mutable GpuHnswSearchParams directSearchParams_; + mutable bool hasDirectSearchParams_ = false; + + // Metric interpretation, set at copy time (see isCosine()/useInnerProduct()). + bool is_cosine_ = false; + bool use_ip_ = false; +}; + +} // namespace gpu +} // namespace faiss diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswBruteForce.cuh b/thirdparty/faiss/faiss/gpu/impl/GpuHnswBruteForce.cuh new file mode 100644 index 000000000..0f0045db3 --- /dev/null +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswBruteForce.cuh @@ -0,0 +1,162 @@ +// @lint-ignore-every LICENSELINT +/** + * Copyright (c) Facebook, Inc. and its affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ +/* + * Copyright (c) 2026, 6sense Insights Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +#include +#include + +#include + +namespace faiss { +namespace gpu { +namespace hnsw_kernel { + +// Fixed upper bound on the launch block size for the brute-force kernel; the +// tree reduction requires a power-of-two block, so callers must launch with a +// power-of-two blockDim <= kBruteForceMaxBlock. +constexpr int kBruteForceMaxBlock = 256; + +// Brute-force top-k over the live (non-filtered) rows for a set of queries, +// operating entirely from GPU-resident vectors (Knowhere frees the CPU HNSW +// index after upload, so the fallback cannot call it). Distances use the same +// "smaller == closer" internal convention as the graph kernel (IP negated), +// and are negated back on copy-out so callers receive the true similarity. +// +// One block per work item. The work item maps to a query index via d_worklist, +// or identity when d_worklist == nullptr (the up-front all-queries path). Rows +// filtered out by the bitset are skipped. Selection is iterative in +// (distance, id) lexicographic order: k block-wide reductions, each finding the +// smallest live candidate strictly after the previous winner. This needs no +// per-query top-k buffer and is duplicate-free because (dist, id) is a total +// order. Cost is O(k * N / blockDim) per query. When fewer than k live rows +// exist the remaining slots are filled with sentinels (-1 id, worst score). +// +// Requires: blockDim.x is a power of two, <= kBruteForceMaxBlock. For the DP4A +// int8 path (USE_DP4A) dim must be a multiple of 4, matching the graph kernel. +template +__global__ void brute_force_topk_kernel( + const QueryT* __restrict__ d_queries, + const DataT* __restrict__ d_dataset, + const float* __restrict__ d_inv_norms, + const uint8_t* __restrict__ d_bitset, + const uint32_t* __restrict__ d_worklist, // null => identity mapping + int num_items, + const int* __restrict__ d_num_items, // null => use num_items directly + int N, + int dim, + int k, + bool use_inner_product, + uint64_t* __restrict__ d_neighbors, + float* __restrict__ d_distances) { + // The per-query fallback launches a full num_queries grid and reads the + // device-side worklist length here, so no host sync is needed between the + // graph kernel and this pass. The up-front path passes num_items directly. + int count = (d_num_items != nullptr) ? *d_num_items : num_items; + int item = blockIdx.x; + if (item >= count) + return; + int query_idx = (d_worklist != nullptr) + ? static_cast(d_worklist[item]) + : item; + + const QueryT* query = d_queries + static_cast(query_idx) * dim; + int dim4 = dim / 4; + int tid = threadIdx.x; + + __shared__ float s_dist[kBruteForceMaxBlock]; + __shared__ uint32_t s_id[kBruteForceMaxBlock]; + __shared__ float win_dist_s; + __shared__ uint32_t win_id_s; + + // (-inf, 0): the first round selects the global minimum since any finite + // distance is strictly greater than -FLT_MAX. + float prev_dist = -FLT_MAX; + uint32_t prev_id = 0; + + for (int t = 0; t < k; t++) { + float local_best = FLT_MAX; + uint32_t local_id = UINT32_MAX; + for (int row = tid; row < N; row += blockDim.x) { + uint32_t r = static_cast(row); + if (is_bitset_filtered(d_bitset, r)) + continue; + float d = layer0_distance( + query, d_dataset, d_inv_norms, r, dim, dim4, + use_inner_product); + bool after = (d > prev_dist) || (d == prev_dist && r > prev_id); + if (!after) + continue; + if (d < local_best || (d == local_best && r < local_id)) { + local_best = d; + local_id = r; + } + } + s_dist[tid] = local_best; + s_id[tid] = local_id; + __syncthreads(); + for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) { + if (tid < stride) { + float od = s_dist[tid + stride]; + uint32_t oi = s_id[tid + stride]; + if (od < s_dist[tid] || + (od == s_dist[tid] && oi < s_id[tid])) { + s_dist[tid] = od; + s_id[tid] = oi; + } + } + __syncthreads(); + } + if (tid == 0) { + win_dist_s = s_dist[0]; + win_id_s = s_id[0]; + } + __syncthreads(); + float win_dist = win_dist_s; + uint32_t win_id = win_id_s; + + int64_t out = static_cast(query_idx) * k + t; + if (win_id == UINT32_MAX) { + // No live candidate remains: sentinel-fill this and the rest. + for (int tt = t + tid; tt < k; tt += blockDim.x) { + int64_t o2 = static_cast(query_idx) * k + tt; + d_neighbors[o2] = UINT64_MAX; + d_distances[o2] = use_inner_product ? -FLT_MAX : FLT_MAX; + } + break; + } + if (tid == 0) { + d_neighbors[out] = static_cast(win_id); + d_distances[out] = use_inner_product ? -win_dist : win_dist; + } + prev_dist = win_dist; + prev_id = win_id; + __syncthreads(); + } +} + +} // namespace hnsw_kernel +} // namespace gpu +} // namespace faiss diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh b/thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh new file mode 100644 index 000000000..7d2305d72 --- /dev/null +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh @@ -0,0 +1,448 @@ +// @lint-ignore-every LICENSELINT +/** + * Copyright (c) Facebook, Inc. and its affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ +/* + * Copyright (c) 2026, 6sense Insights Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Converts a FAISS HNSW index (CSR graph) + vectors into a +// GpuHnswDeviceIndex on device memory. +// +// This version is adapted for Knowhere's FAISS fork which uses +// faiss::cppcontrib::knowhere::IndexHNSW / HNSW types. + +#pragma once + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include +#include + +// GpuHnswUploadFaultInjection (used below) lives in GpuHnswTypes.h so it is +// reachable from host-compiled unit tests without pulling in device kernels. +#define GPU_HNSW_BUILD_CUDA_CHECK(expr) \ + do { \ + if (faiss::gpu::GpuHnswUploadFaultInjection::should_fail()) { \ + throw std::runtime_error( \ + std::string("CUDA error (injected): simulated ") + \ + "upload failure at " + __FILE__ + ":" + \ + std::to_string(__LINE__)); \ + } \ + cudaError_t _e = (expr); \ + if (_e != cudaSuccess) { \ + throw std::runtime_error( \ + std::string("CUDA error: ") + \ + cudaGetErrorString(_e) + " at " + __FILE__ + ":" + \ + std::to_string(__LINE__)); \ + } \ + } while (0) + +namespace faiss { +namespace gpu { + +/// Extract HNSW graph layers from a Knowhere HNSW struct. +/// Template parameter HnswT can be faiss::cppcontrib::knowhere::HNSW +/// or faiss::HNSW — any type exposing this interface: +/// - neighbor_range(node, layer, &begin, &end) +/// - nb_neighbors(layer) -> int +/// - neighbors : flat neighbor array indexed by neighbor_range +/// - levels : per-node array; levels[i] is the 1-based layer count, so +/// node i lives on layers 0..levels[i]-1 and "i is on layer L" +/// is the test levels[i] > L +/// - entry_point, max_level +/// The levels[] convention above matches faiss::HNSW (HNSW.cpp uses +/// pt_level = levels[i] - 1); a variant with different semantics would need a +/// different membership test at the levels[i] > layer check below. +template +inline void extract_hnsw_layers( + const HnswT& hnsw, + int64_t n_rows, + std::vector& h_upper_layers, + std::vector& h_layer0_flat, + uint32_t& entry_point, + int& M, + int& max_degree0, + int& num_layers) { + const int maxM0 = hnsw.nb_neighbors(0); + const int maxM = hnsw.nb_neighbors(1); + const int max_lv = hnsw.max_level; + + // The kernels dereference d_dataset + entry_point * dim directly, so a + // malformed CPU index carrying an invalid/sentinel entry_point (e.g. -1 + // cast to UINT32_MAX) would be an out-of-bounds device read. Validate here + // rather than faulting on the GPU. + if (hnsw.entry_point < 0 || hnsw.entry_point >= n_rows) { + throw std::runtime_error( + std::string("gpu_hnsw: invalid HNSW entry_point ") + + std::to_string(static_cast(hnsw.entry_point)) + + " (n_rows=" + std::to_string(n_rows) + ")"); + } + entry_point = static_cast(hnsw.entry_point); + M = maxM; + max_degree0 = maxM0; + num_layers = max_lv + 1; + + // Layer 0: dense [n_rows x maxM0] + h_layer0_flat.assign(n_rows * maxM0, UINT32_MAX); + for (int64_t i = 0; i < n_rows; i++) { + size_t begin, end; + hnsw.neighbor_range(i, 0, &begin, &end); + uint32_t count = static_cast(end - begin); + for (uint32_t j = 0; j < count; j++) { + auto nb = hnsw.neighbors[begin + j]; + if (nb >= 0) + h_layer0_flat[i * maxM0 + j] = static_cast(nb); + } + } + + // Upper layers (1..max_level): sparse [num_nodes_at_L x maxM] + h_upper_layers.resize(max_lv); + for (int layer = 1; layer <= max_lv; layer++) { + auto& ul = h_upper_layers[layer - 1]; + ul.max_degree = static_cast(maxM); + + std::vector node_ids; + for (int64_t i = 0; i < n_rows; i++) { + if (hnsw.levels[i] > layer) + node_ids.push_back(static_cast(i)); + } + ul.num_nodes = static_cast(node_ids.size()); + + std::vector h_neighbors(ul.num_nodes * maxM, UINT32_MAX); + + for (uint32_t idx = 0; idx < ul.num_nodes; idx++) { + int64_t i = node_ids[idx]; + size_t begin, end; + hnsw.neighbor_range(i, layer, &begin, &end); + uint32_t count = static_cast(end - begin); + for (uint32_t j = 0; j < count; j++) { + auto nb = hnsw.neighbors[begin + j]; + if (nb >= 0) + h_neighbors[idx * maxM + j] = static_cast(nb); + } + } + + GPU_HNSW_BUILD_CUDA_CHECK( + cudaMalloc(&ul.d_node_ids, ul.num_nodes * sizeof(uint32_t))); + GPU_HNSW_BUILD_CUDA_CHECK(cudaMemcpy( + ul.d_node_ids, + node_ids.data(), + ul.num_nodes * sizeof(uint32_t), + cudaMemcpyHostToDevice)); + + GPU_HNSW_BUILD_CUDA_CHECK(cudaMalloc( + &ul.d_neighbors, + ul.num_nodes * maxM * sizeof(uint32_t))); + GPU_HNSW_BUILD_CUDA_CHECK(cudaMemcpy( + ul.d_neighbors, + h_neighbors.data(), + ul.num_nodes * maxM * sizeof(uint32_t), + cudaMemcpyHostToDevice)); + } +} + +inline void normalize_vectors( + std::vector& h_vectors, + int64_t n_rows, + int64_t dim) { + for (int64_t i = 0; i < n_rows; i++) { + float* v = h_vectors.data() + i * dim; + float sq_norm = 0.0f; + for (int64_t d = 0; d < dim; d++) + sq_norm += v[d] * v[d]; + if (sq_norm > 0.0f) { + float inv = 1.0f / std::sqrt(sq_norm); + for (int64_t d = 0; d < dim; d++) + v[d] *= inv; + } + } +} + +template +inline void upload_graph_to_gpu( + GpuHnswDeviceIndex& idx, + const HnswT& hnsw, + int64_t n_rows) { + std::vector h_layer0_flat; + extract_hnsw_layers( + hnsw, + n_rows, + idx.upper_layers, + h_layer0_flat, + idx.entry_point, + idx.M, + idx.max_degree0, + idx.num_layers); + + size_t graph0_bytes = + static_cast(n_rows) * idx.max_degree0 * sizeof(uint32_t); + GPU_HNSW_BUILD_CUDA_CHECK(cudaMalloc(&idx.d_layer0_graph, graph0_bytes)); + GPU_HNSW_BUILD_CUDA_CHECK(cudaMemcpy( + idx.d_layer0_graph, + h_layer0_flat.data(), + graph0_bytes, + cudaMemcpyHostToDevice)); + + int num_upper = static_cast(idx.upper_layers.size()); + idx.num_upper_layers_built = num_upper; + if (num_upper > 0) { + using kernel_ptrs = hnsw_kernel::upper_layer_ptrs; + std::vector h_ptrs(num_upper); + for (int i = 0; i < num_upper; i++) { + const auto& ul = idx.upper_layers[i]; + h_ptrs[i] = { + ul.d_node_ids, ul.d_neighbors, ul.num_nodes, ul.max_degree}; + } + size_t ptrs_bytes = num_upper * sizeof(kernel_ptrs); + GPU_HNSW_BUILD_CUDA_CHECK( + cudaMalloc(&idx.d_upper_layer_ptrs, ptrs_bytes)); + GPU_HNSW_BUILD_CUDA_CHECK(cudaMemcpy( + idx.d_upper_layer_ptrs, + h_ptrs.data(), + ptrs_bytes, + cudaMemcpyHostToDevice)); + } + +} + +// Upload precomputed per-row inverse L2 norms to the device. These must be the +// norms of the *original* input vectors as recorded by the CPU cosine index +// (HasInverseL2Norms::get_inverse_l2_norms()), NOT norms recomputed from +// lossily-decoded codes — otherwise cosine scores and graph traversal diverge +// from the CPU index for lossy SQ (SQ8 / fp16 / bf16). The search kernel +// computes score = IP(query, decoded_db) * inv_norm[row], mirroring the CPU +// WithCosineNormDistanceComputer. +inline void upload_inv_norms( + GpuHnswDeviceIndex& idx, + const float* inv_norms, + int64_t n_rows) { + size_t norms_bytes = static_cast(n_rows) * sizeof(float); + GPU_HNSW_BUILD_CUDA_CHECK(cudaMalloc(&idx.d_inv_norms, norms_bytes)); + GPU_HNSW_BUILD_CUDA_CHECK(cudaMemcpy( + idx.d_inv_norms, inv_norms, norms_bytes, cudaMemcpyHostToDevice)); +} + +inline void upload_fp32_dataset( + GpuHnswDeviceIndex& idx, + std::vector& h_vectors, + int64_t n_rows, + bool is_cosine, + const float* stored_inv_norms = nullptr) { + int64_t dim = idx.dim; + // Flat cosine (stored_inv_norms == nullptr): the stored vectors are the + // exact originals, so normalizing them in place yields cosine via plain + // inner product. Lossy-SQ cosine decoded to fp32 (stored_inv_norms != null): + // keep the decoded vectors un-normalized and apply the CPU index's original + // inverse norms at search time, matching CPU semantics for lossy codes. + if (is_cosine && stored_inv_norms == nullptr) + normalize_vectors(h_vectors, n_rows, dim); + + size_t dataset_bytes = static_cast(n_rows) * dim * sizeof(float); + GPU_HNSW_BUILD_CUDA_CHECK(cudaMalloc(&idx.d_dataset, dataset_bytes)); + GPU_HNSW_BUILD_CUDA_CHECK(cudaMemcpy( + idx.d_dataset, + h_vectors.data(), + dataset_bytes, + cudaMemcpyHostToDevice)); + idx.dataset_type = GpuHnswDatasetType::FP32; + + if (is_cosine && stored_inv_norms != nullptr) + upload_inv_norms(idx, stored_inv_norms, n_rows); +} + +// Upload fp16 (QT_fp16) or bf16 (QT_bf16) ScalarQuantizer codes to the GPU in +// their native 2-byte layout. faiss stores these codes row-major as raw IEEE +// half / bfloat16, which are bit-compatible with CUDA half / __nv_bfloat16, so +// the bytes are copied verbatim (no up-conversion to fp32). For cosine, the CPU +// index's original inverse L2 norms are applied at search time — mirroring the +// int8 path, since the stored codes are not normalized. +inline void upload_halfwidth_dataset( + GpuHnswDeviceIndex& idx, + const uint8_t* codes, + int64_t n_rows, + bool is_cosine, + GpuHnswDatasetType dtype, + const float* stored_inv_norms) { + int64_t dim = idx.dim; + size_t dataset_bytes = static_cast(n_rows) * dim * 2; + GPU_HNSW_BUILD_CUDA_CHECK(cudaMalloc(&idx.d_dataset, dataset_bytes)); + GPU_HNSW_BUILD_CUDA_CHECK(cudaMemcpy( + idx.d_dataset, codes, dataset_bytes, cudaMemcpyHostToDevice)); + idx.dataset_type = dtype; + + // The stored fp16/bf16 codes are not normalized, so cosine needs the CPU + // index's original inverse norms (not norms recomputed from the lossy + // decoded values). Mirrors the int8 path. + if (is_cosine) + upload_inv_norms(idx, stored_inv_norms, n_rows); +} + +inline void upload_int8_dataset( + GpuHnswDeviceIndex& idx, + const uint8_t* codes, + int64_t n_rows, + bool is_cosine, + const float* stored_inv_norms) { + int64_t dim = idx.dim; + size_t dataset_bytes = static_cast(n_rows) * dim; + + std::vector signed_codes(dataset_bytes); + for (size_t i = 0; i < dataset_bytes; i++) { + signed_codes[i] = + static_cast(static_cast(codes[i]) - 128); + } + + GPU_HNSW_BUILD_CUDA_CHECK(cudaMalloc(&idx.d_dataset, dataset_bytes)); + GPU_HNSW_BUILD_CUDA_CHECK(cudaMemcpy( + idx.d_dataset, + signed_codes.data(), + dataset_bytes, + cudaMemcpyHostToDevice)); + idx.dataset_type = GpuHnswDatasetType::INT8; + + // Apply the CPU index's original inverse norms for cosine. For + // QT_8bit_direct_signed the codes are the original data (lossless), so these + // match a recompute; using the stored norms keeps all cosine paths uniform + // and bit-exact with the CPU index. + if (is_cosine) + upload_inv_norms(idx, stored_inv_norms, n_rows); +} + +/// Build from Knowhere's HNSW index with SQ storage. +inline std::unique_ptr from_faiss_hnsw_sq( + const faiss::cppcontrib::knowhere::IndexHNSW& hnsw_index, + bool use_ip, + bool is_cosine = false, + int device = 0) { + const auto* sq_storage = + dynamic_cast( + hnsw_index.storage); + if (!sq_storage) + throw std::runtime_error( + "gpu_hnsw: storage is not IndexScalarQuantizer"); + + int64_t n_rows = hnsw_index.ntotal; + int64_t dim = hnsw_index.d; + + auto idx = std::make_unique(); + idx->n_rows = n_rows; + idx->dim = dim; + idx->use_ip = use_ip; + idx->device = device; + idx->scratch_pool = std::make_unique(4, device); + + auto qtype = sq_storage->sq.qtype; + + // For cosine, use the CPU index's inverse L2 norms computed from the + // *original* input vectors (HasInverseL2Norms::get_inverse_l2_norms()). + // Recomputing norms from lossily-decoded/quantized codes would diverge from + // the CPU index's scores and graph traversal for lossy SQ. + const float* stored_inv_norms = nullptr; + if (is_cosine) { + // The cosine norms live on the SQ storage (IndexScalarQuantizerCosine), + // which implements HasInverseL2Norms — not on the outer IndexHNSW. + const auto* cos = dynamic_cast< + const faiss::cppcontrib::knowhere::HasInverseL2Norms*>( + sq_storage); + if (!cos || !cos->get_inverse_l2_norms()) + throw std::runtime_error( + "gpu_hnsw: cosine SQ index missing inverse L2 norms"); + stored_inv_norms = cos->get_inverse_l2_norms(); + } + + if (qtype == faiss::ScalarQuantizer::QT_8bit_direct_signed) { + upload_int8_dataset( + *idx, sq_storage->codes.data(), n_rows, is_cosine, + stored_inv_norms); + } else if ( + qtype == faiss::ScalarQuantizer::QT_fp16 || + qtype == faiss::ScalarQuantizer::QT_bf16) { + // Keep fp16/bf16 in their native 2-byte layout on the GPU. + GpuHnswDatasetType dtype = + (qtype == faiss::ScalarQuantizer::QT_fp16) + ? GpuHnswDatasetType::FP16 + : GpuHnswDatasetType::BF16; + upload_halfwidth_dataset( + *idx, + sq_storage->codes.data(), + n_rows, + is_cosine, + dtype, + stored_inv_norms); + } else { + // Other SQ types (e.g. unsigned QT_8bit / SQ8) are decoded to fp32 and + // uploaded un-normalized; the stored original inverse norms are applied + // in the kernel, matching the CPU cosine computer for lossy codes. + std::vector h_vectors(n_rows * dim); + sq_storage->sa_decode( + n_rows, sq_storage->codes.data(), h_vectors.data()); + upload_fp32_dataset( + *idx, h_vectors, n_rows, is_cosine, stored_inv_norms); + } + + upload_graph_to_gpu(*idx, hnsw_index.hnsw, n_rows); + return idx; +} + +/// Build from Knowhere's HNSW index with Flat storage. +inline std::unique_ptr from_faiss_hnsw_flat( + const faiss::cppcontrib::knowhere::IndexHNSW& hnsw_index, + bool use_ip, + bool is_cosine = false, + int device = 0) { + const auto* flat_storage = + dynamic_cast(hnsw_index.storage); + if (!flat_storage) + throw std::runtime_error("gpu_hnsw: storage is not IndexFlat"); + + int64_t n_rows = hnsw_index.ntotal; + int64_t dim = hnsw_index.d; + + // Use reconstruct_n instead of get_xb() — get_xb() can return a + // device pointer in GPU querynode context, causing SIGSEGV when + // accessed from CPU. + std::vector h_vectors(n_rows * dim); + flat_storage->reconstruct_n(0, n_rows, h_vectors.data()); + + auto idx = std::make_unique(); + idx->n_rows = n_rows; + idx->dim = dim; + idx->use_ip = use_ip; + idx->device = device; + idx->scratch_pool = std::make_unique(4, device); + + upload_fp32_dataset(*idx, h_vectors, n_rows, is_cosine); + upload_graph_to_gpu(*idx, hnsw_index.hnsw, n_rows); + return idx; +} + +} // namespace gpu +} // namespace faiss diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh new file mode 100644 index 000000000..df128d4c8 --- /dev/null +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh @@ -0,0 +1,485 @@ +// @lint-ignore-every LICENSELINT +/** + * Copyright (c) Facebook, Inc. and its affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ +/* + * Copyright (c) 2026, 6sense Insights Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +#define GPU_HNSW_CUDA_CHECK(expr) \ + do { \ + cudaError_t _e = (expr); \ + if (_e != cudaSuccess) { \ + throw std::runtime_error( \ + std::string("CUDA error: ") + \ + cudaGetErrorString(_e) + " at " + __FILE__ + ":" + \ + std::to_string(__LINE__)); \ + } \ + } while (0) + +namespace faiss { +namespace gpu { + +inline void gpu_hnsw_search( + cudaStream_t stream, + const GpuHnswSearchParams& params, + const GpuHnswDeviceIndex& idx, + GpuHnswSearchScratch& sc, + int num_queries, + int k) { + + int ef = params.ef; + int sw = params.search_width; + // Reject degenerate params up front: ef and search_width feed the + // shared-memory sizing and the auto-iteration bound (2*ef/sw below), so a + // value of 0 would divide by zero and a negative value would under-size + // the buffers. + if (ef <= 0 || sw <= 0) { + throw std::runtime_error( + std::string("gpu_hnsw: ef and search_width must be > 0 (ef=") + + std::to_string(ef) + ", search_width=" + std::to_string(sw) + + ")"); + } + // Auto-clamp ef up to k. The kernel tracks exactly ef candidates in the + // result beam, so ef < k could only ever return ef real neighbors with the + // remaining k-ef slots padded with sentinels. Raising ef to k matches the + // CPU / Knowhere HNSW contract (always return k valid results). The + // shared-memory-fit check below still throws if even k slots cannot fit. + if (k > ef) { + ef = k; + } + int max_iter = params.max_iterations > 0 + ? params.max_iterations + : 2 * ef / sw + 10; + int dim = static_cast(idx.dim); + int num_upper_layers = idx.num_upper_layers_built; + + // --- Filtered-search setup (deletes / TTL / partition bitset) --- + // The bitset bytes are uploaded into sc.d_bitset by the caller (searchHost/ + // searchHostInt8) on this same stream before gpu_hnsw_search runs. When + // bitset_data is null we take the unfiltered fast path. + bool has_filter = (params.bitset_data != nullptr); + int64_t nbits = params.bitset_nbits; + int64_t filtered_count = params.bitset_filtered_count; + int64_t live_count = has_filter ? (nbits - filtered_count) + : static_cast(idx.n_rows); + if (live_count < 0) { + live_count = 0; + } + // kAlpha rate-limits how often filtered nodes become waypoints; it scales + // with the delete ratio, matching CPU HNSW (kAlpha = filter_ratio * 0.7). + float filter_ratio = (has_filter && nbits > 0) + ? static_cast(filtered_count) / static_cast(nbits) + : 0.0f; + float kAlpha = filter_ratio * 0.7f; + bool disable_bf = params.disable_fallback_brute_force; + + // Up-front brute force: when almost everything is filtered, or k is a large + // fraction of the live set, the graph walk cannot beat a full scan and + // risks returning fewer than k live results. Mirrors CPU HNSW's + // kHnswSearchKnnBFFilterThreshold (0.93) and kHnswSearchBFTopkThreshold + // (0.5, applied to the live count). Disabled when the caller disables the + // brute-force fallback. + bool up_front_bf = has_filter && !disable_bf && nbits > 0 && + (static_cast(filtered_count) >= + 0.93 * static_cast(nbits) || + static_cast(k) >= 0.5 * static_cast(live_count)); + + // Layer-0 is templated on the dataset type (DataT), the layer-0 query type + // (QueryT: float generic, int8_t for the native DP4A path) and USE_DP4A. + // The upper-layer greedy descent always uses the fp32 queries (sc.d_queries). + auto launch_kernels = [&]( + const DataT* d_data, + const float* d_inv_norms, + const QueryT* d_layer0_queries) { + int N_int = static_cast(idx.n_rows); + + // Brute-force launcher (shares the current dtype specialization). It + // operates on a chunk of `num_items` queries whose scratch pointers + // (q0/nb0/ds0) are already offset by the caller; when non-null the + // worklist holds chunk-local query indices, matching the offset outputs. + // Grid is num_items: the up-front path uses the identity mapping, the + // per-query fallback reads its worklist length from d_num on the device + // (no host sync between graph and BF). Block must be a power of two <= + // kBruteForceMaxBlock for the tree reduction. + int bf_block = 128; + auto launch_bf = [&](const QueryT* q0, + uint64_t* nb0, + float* ds0, + const uint32_t* worklist, + const int* d_num, + int num_items) { + hnsw_kernel::brute_force_topk_kernel + <<>>( + q0, + d_data, + d_inv_norms, + sc.d_bitset, + worklist, + num_items, + d_num, + N_int, + dim, + k, + idx.use_ip, + nb0, + ds0); + GPU_HNSW_CUDA_CHECK(cudaGetLastError()); + }; + + // --- Shared-memory sizing / ef finalization --- + // Query-count independent (depends only on ef/search_width/M and the + // device smem limit), so compute it once before the chunk loop. + int block_size = + params.thread_block_size > 0 ? params.thread_block_size : 128; + + // Per-block dynamic shared-memory budget for this device. The default + // limit is 48 KiB, but Volta+ GPUs can opt into more via + // cudaFuncSetAttribute; query the real limit instead of assuming 48 KiB. + int smem_max = 49152; + { + int device = 0; + int optin = 0; + if (cudaGetDevice(&device) == cudaSuccess && + cudaDeviceGetAttribute( + &optin, + cudaDevAttrMaxSharedMemoryPerBlockOptin, + device) == cudaSuccess && + optin > smem_max) { + smem_max = optin; + } + } + + // The bitonic sort in the parallel merge needs a power-of-two staging + // capacity, one thread per slot. Pad up to the next power of two + // (handles non-power-of-two 2*M) and grow the block so every staging + // slot is owned by a thread. + int max_staging = + hnsw_kernel::padded_staging_capacity(sw, idx.max_degree0); + { + if (block_size < max_staging) { + block_size = max_staging; + } + // The layer-0 expansion is warp-cooperative (one warp per candidate + // edge, 32 lanes striding the vector for coalesced loads), so the + // block must be a whole number of warps for the full-mask __shfl + // reductions to be well-defined. The default (128) and the + // power-of-two staging bump are already multiples of 32; this only + // rounds up a caller-supplied non-multiple thread_block_size. + block_size = ((block_size + 31) / 32) * 32; + if (block_size > 1024) { + throw std::runtime_error( + std::string("gpu_hnsw: padded staging capacity ") + + std::to_string(max_staging) + + " exceeds the max CUDA block size (1024); reduce " + "search_width or M (max_degree0=" + + std::to_string(idx.max_degree0) + ")"); + } + // Fixed overhead: staging (max_staging*8) + parent_ids (sw*4) + // + meta (12, or 24 on the filtered path with 6 slots). + int meta_bytes = has_filter ? 24 : 12; + int smem_overhead = max_staging * 8 + sw * 4 + meta_bytes; + // Per-ef cost: 3 result arrays + 3 merge arrays = 6 × 4 = 24 + // bytes/slot. The filtered path adds the invalid frontier (3 more + // arrays); with ef_inv capped at ef (the default) that is a further + // 12 bytes/slot, so divide by 36 to bound ef conservatively. + int per_ef = has_filter ? 36 : 24; + int max_ef = (smem_max - smem_overhead) / per_ef; + if (max_ef < 1) { + throw std::runtime_error( + std::string("gpu_hnsw: search_width=") + + std::to_string(sw) + + " too large for device shared memory (" + + std::to_string(smem_max) + + " bytes); reduce search_width"); + } + // ef was auto-clamped up to at least k above. If even k candidate + // slots cannot fit in shared memory, fail loudly rather than + // silently returning fewer than k valid results. + if (k > max_ef) { + throw std::runtime_error( + std::string("gpu_hnsw: k=") + std::to_string(k) + + " needs k candidate slots (ef>=k), only " + + std::to_string(max_ef) + + " fit in device shared memory (" + + std::to_string(smem_max) + + " bytes); reduce k or raise thread_block_size"); + } + if (ef > max_ef) { + fprintf(stderr, + "[gpu_hnsw] warning: ef=%d exceeds the per-block " + "shared-memory budget (%d bytes); clamping ef to %d. " + "Recall may be reduced; raise thread_block_size or " + "lower search_width to restore the requested ef.\n", + ef, smem_max, max_ef); + ef = max_ef; + } + } + + // Invalid-frontier capacity: default to the (now finalized) ef, but + // allow the caller to cap it smaller to save shared memory. Capped at + // ef so the /36 clamp above stays conservative. Only used on the + // filtered path. + int ef_inv = 0; + if (has_filter) { + ef_inv = (params.ef_inv > 0) ? std::min(params.ef_inv, ef) : ef; + } + + size_t smem_size = hnsw_kernel::calc_layer0_smem_size( + ef, sw, idx.max_degree0, has_filter ? ef_inv : 0); + + // Graph search + per-query brute-force fallback for one chunk of `cnq` + // queries. The scratch pointers (q0/ep0/nb0/ds0) are offset by the + // caller so the kernels' chunk-local blockIdx.x maps to the right + // global query; the visited bitmap is the base buffer, indexed by the + // same chunk-local index. The graph kernel is templated on HAS_FILTER; + // the filtered and unfiltered instantiations are distinct __global__ + // functions, so each needs its own cudaFuncSetAttribute high-water + // tracking (the statics below are per-<...,HAS_FILTER>). + auto run_graph = [&]( + const QueryT* q0, + uint32_t* ep0, + uint64_t* nb0, + float* ds0, + int cnq) { + // Opt into >48 KiB dynamic shared memory when the device supports + // it; without this the kernel launch would fail for large ef. This + // is cached per kernel instantiation: cudaFuncSetAttribute + // configures the kernel's max dynamic shared memory globally, so it + // only needs to run once per (kernel, high-water size). + // + // The mutex makes the check-set-store atomic and monotonic: without + // it, two concurrent searches with different smem sizes can + // interleave so a smaller-ef search's cudaFuncSetAttribute runs + // *after* a larger one, downgrading the kernel's global attribute + // below what a recorded high-water mark implies, and a later + // intermediate search then skips the set (thinks it's configured) + // and fails to launch. Under the lock the attribute only ever grows, + // and is always >= the current launch's requirement before we + // proceed. + if (smem_size > 49152) { + static std::mutex configured_smem_mutex; + static size_t configured_smem = 0; + std::lock_guard lock(configured_smem_mutex); + if (smem_size > configured_smem) { + GPU_HNSW_CUDA_CHECK(cudaFuncSetAttribute( + hnsw_kernel::layer0_beam_search_kernel< + DataT, + QueryT, + USE_DP4A, + HAS_FILTER>, + cudaFuncAttributeMaxDynamicSharedMemorySize, + static_cast(smem_size))); + configured_smem = smem_size; + } + } + + // Zero the per-query BF worklist counter before the graph launch + // (same stream, so ordered ahead of the kernel that appends to it). + if constexpr (HAS_FILTER) { + if (!disable_bf) { + GPU_HNSW_CUDA_CHECK(cudaMemsetAsync( + sc.d_needs_bf_count, 0, sizeof(int), stream)); + } + } + + const uint8_t* d_bitset = HAS_FILTER ? sc.d_bitset : nullptr; + uint32_t* d_needs_bf = HAS_FILTER ? sc.d_needs_bf : nullptr; + int* d_needs_bf_count = HAS_FILTER ? sc.d_needs_bf_count : nullptr; + + hnsw_kernel::layer0_beam_search_kernel< + DataT, + QueryT, + USE_DP4A, + HAS_FILTER><<>>( + q0, + d_data, + d_inv_norms, + idx.d_layer0_graph, + ep0, + sc.d_visited_bitmaps, + nb0, + ds0, + cnq, + N_int, + dim, + idx.max_degree0, + k, + ef, + sw, + max_iter, + idx.use_ip, + d_bitset, + ef_inv, + kAlpha, + static_cast(live_count), + disable_bf, + d_needs_bf, + d_needs_bf_count); + GPU_HNSW_CUDA_CHECK(cudaGetLastError()); + + // Per-query brute-force fallback: queries whose graph search fell + // short (fewer than k live results) were appended to sc.d_needs_bf + // (chunk-local indices) by the graph kernel. Launch a full-width + // grid; each block reads the device worklist length and early-exits + // when out of range, so no host sync is needed between the two + // kernels. + if constexpr (HAS_FILTER) { + if (!disable_bf) { + launch_bf(q0, nb0, ds0, sc.d_needs_bf, + sc.d_needs_bf_count, cnq); + } + } + }; + + // --- Query chunking to bound the visited-bitmap VRAM --- + // Process the batch in chunks of at most chunk_nq queries so the visited + // bitmap (sized in GpuHnswSearchScratch::ensure with the same chunk) + // stays within the VRAM cap regardless of nq / search concurrency. For + // small segments chunk_nq == num_queries -> a single pass, identical to + // the pre-chunk behavior. + int chunk_nq = gpu_hnsw_bitmap_chunk(num_queries, N_int); + for (int c0 = 0; c0 < num_queries; c0 += chunk_nq) { + int cnq = std::min(chunk_nq, num_queries - c0); + // Chunk-offset scratch pointers. q0 is the layer-0 query buffer + // (int8 on the DP4A path, else fp32); fq0 is always the fp32 buffer + // used by the upper-layer greedy descent. + const QueryT* q0 = + d_layer0_queries + static_cast(c0) * dim; + const float* fq0 = sc.d_queries + static_cast(c0) * dim; + uint32_t* ep0 = sc.d_entry_points + c0; + uint64_t* nb0 = sc.d_neighbors + static_cast(c0) * k; + float* ds0 = sc.d_distances + static_cast(c0) * k; + + // Up-front brute force skips the graph search entirely (does not + // touch the visited bitmap, but still runs per chunk so its outputs + // land at the right offsets). + if (up_front_bf) { + launch_bf(q0, nb0, ds0, nullptr, nullptr, cnq); + continue; + } + + if (num_upper_layers > 0) { + auto* d_layer_ptrs = + static_cast( + idx.d_upper_layer_ptrs); + + int warps_per_block = 4; + int threads_per_block = warps_per_block * 32; + int num_blocks = + (cnq + warps_per_block - 1) / warps_per_block; + + hnsw_kernel::upper_layer_search_kernel + <<>>( + fq0, + d_data, + d_inv_norms, + d_layer_ptrs, + ep0, + idx.entry_point, + cnq, + dim, + num_upper_layers, + idx.use_ip); + GPU_HNSW_CUDA_CHECK(cudaGetLastError()); + } else { + std::vector h_eps(cnq, idx.entry_point); + GPU_HNSW_CUDA_CHECK(cudaMemcpyAsync( + ep0, + h_eps.data(), + static_cast(cnq) * sizeof(uint32_t), + cudaMemcpyHostToDevice, + stream)); + // h_eps is stack-local; synchronize before it is destroyed so + // the copy never reads freed memory (safe even if the source + // buffer is ever switched to pinned host memory). + GPU_HNSW_CUDA_CHECK(cudaStreamSynchronize(stream)); + } + + // Zero this chunk's visited bitmap (reused across chunks). + size_t bitmap_bytes = + hnsw_kernel::calc_visited_bitmap_size(cnq, N_int); + GPU_HNSW_CUDA_CHECK(cudaMemsetAsync( + sc.d_visited_bitmaps, 0, bitmap_bytes, stream)); + + if (has_filter) { + run_graph.template operator()(q0, ep0, nb0, ds0, cnq); + } else { + run_graph.template operator()(q0, ep0, nb0, ds0, cnq); + } + } + }; + + switch (idx.dataset_type) { + case GpuHnswDatasetType::INT8: + // Native DP4A path: requires int8 queries staged on device, an + // inner-product/cosine metric (DP4A only computes dot products) and + // dim % 4 == 0. Otherwise fall back to the generic fp32-query path. + if (sc.d_queries_i8 != nullptr && idx.use_ip && (dim % 4 == 0)) { + launch_kernels.template operator()( + static_cast(idx.d_dataset), + idx.d_inv_norms, + sc.d_queries_i8); + } else { + launch_kernels.template operator()( + static_cast(idx.d_dataset), + idx.d_inv_norms, + sc.d_queries); + } + break; + case GpuHnswDatasetType::FP16: + launch_kernels.template operator()( + static_cast(idx.d_dataset), + idx.d_inv_norms, + sc.d_queries); + break; + case GpuHnswDatasetType::BF16: + launch_kernels.template operator()<__nv_bfloat16, float, false>( + static_cast(idx.d_dataset), + idx.d_inv_norms, + sc.d_queries); + break; + case GpuHnswDatasetType::FP32: + default: + launch_kernels.template operator()( + static_cast(idx.d_dataset), + idx.d_inv_norms, + sc.d_queries); + break; + } +} + +} // namespace gpu +} // namespace faiss diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh new file mode 100644 index 000000000..66ab00709 --- /dev/null +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh @@ -0,0 +1,1134 @@ +// @lint-ignore-every LICENSELINT +/** + * Copyright (c) Facebook, Inc. and its affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ +/* + * Copyright (c) 2026, 6sense Insights Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include + +#include +#include + +namespace faiss { +namespace gpu { +namespace hnsw_kernel { + +// ============================================================================ +// Distance computation helpers (templated on dataset element type) +// ============================================================================ + +__device__ __forceinline__ float load_elem(const float* ptr, int idx) { + return __ldg(&ptr[idx]); +} +__device__ __forceinline__ float load_elem(const half* ptr, int idx) { + return __half2float(__ldg(&ptr[idx])); +} +__device__ __forceinline__ float load_elem(const __nv_bfloat16* ptr, int idx) { + return __bfloat162float(__ldg(&ptr[idx])); +} +__device__ __forceinline__ float load_elem(const int8_t* ptr, int idx) { + return static_cast(__ldg(&ptr[idx])); +} + +template +__device__ __forceinline__ float thread_l2_distance( + const float* __restrict__ query, + const DataT* __restrict__ vec, + int dim) { + float sum = 0.0f; +#pragma unroll 8 + for (int d = 0; d < dim; d++) { + float diff = query[d] - load_elem(vec, d); + sum += diff * diff; + } + return sum; +} + +template +__device__ __forceinline__ float thread_ip_distance( + const float* __restrict__ query, + const DataT* __restrict__ vec, + int dim) { + float sum = 0.0f; +#pragma unroll 8 + for (int d = 0; d < dim; d++) { + sum += query[d] * load_elem(vec, d); + } + return -sum; +} + +// Native int8 inner-product distance using DP4A (4 int8 MADs/cycle, SM_61+). +// Both query and vec must point at int8 data reinterpreted as packed int32, +// i.e. dim4 == dim / 4 (dim must be a multiple of 4). Returns the negated dot +// product to match thread_ip_distance()'s convention (smaller == closer). +__device__ __forceinline__ float thread_ip_distance_dp4a( + const int32_t* __restrict__ query_packed, + const int32_t* __restrict__ vec_packed, + int dim4) { + int sum = 0; +#pragma unroll 8 + for (int d = 0; d < dim4; d++) { + sum = __dp4a(query_packed[d], vec_packed[d], sum); + } + return static_cast(-sum); +} + +// Smallest power of two >= x (x >= 1). Used to pad the staging capacity so the +// bitonic sort (which requires a power-of-two capacity) never rejects graphs +// whose search_width * max_degree0 is not already a power of two. +__host__ __device__ __forceinline__ int next_pow2_int(int x) { + int p = 1; + while (p < x) + p <<= 1; + return p; +} + +// Padded staging capacity: power-of-two >= search_width * max_degree0. +__host__ __device__ __forceinline__ int padded_staging_capacity( + int search_width, + int max_degree0) { + return next_pow2_int(search_width * max_degree0); +} + +// Unified layer-0 distance: DP4A for native int8+IP, generic fp32 path +// otherwise. The branch is resolved at compile time so only the relevant +// distance function is instantiated for each specialization. +template +__device__ __forceinline__ float layer0_distance( + const QueryT* __restrict__ query, + const DataT* __restrict__ d_dataset, + const float* __restrict__ d_inv_norms, + uint32_t id, + int dim, + int dim4, + bool use_inner_product) { + if constexpr (USE_DP4A) { + const int32_t* vec_packed = reinterpret_cast( + d_dataset + static_cast(id) * dim); + float dist = thread_ip_distance_dp4a( + reinterpret_cast(query), vec_packed, dim4); + if (d_inv_norms) + dist *= __ldg(&d_inv_norms[id]); + return dist; + } else { + const DataT* vec = d_dataset + static_cast(id) * dim; + float dist; + if (use_inner_product) { + dist = thread_ip_distance(query, vec, dim); + if (d_inv_norms) + dist *= __ldg(&d_inv_norms[id]); + } else { + dist = thread_l2_distance(query, vec, dim); + } + return dist; + } +} + +__device__ __forceinline__ float warp_reduce_sum(float v) { +#pragma unroll + for (int off = 16; off > 0; off >>= 1) + v += __shfl_down_sync(0xffffffff, v, off); + return v; +} + +__device__ __forceinline__ int warp_reduce_sum(int v) { +#pragma unroll + for (int off = 16; off > 0; off >>= 1) + v += __shfl_down_sync(0xffffffff, v, off); + return v; +} + +// Warp-cooperative layer-0 distance. The 32 lanes stride the vector (lane d +// reads elements d, d+32, ...), so the dataset loads coalesce into 1 +// sector/request instead of the ~12 the thread-per-candidate path produced +// (see the perf analysis: uncoalesced loads were 49% of warp stalls / DRAM at +// 38% of peak). The per-lane partial sums are warp-reduced; the returned +// distance is valid on lane 0 only (the reduction accumulates there), so +// callers that need it on other lanes must broadcast. Same convention as +// layer0_distance() (smaller == closer, IP negated, COSINE scaled by inv_norm). +template +__device__ __forceinline__ float layer0_distance_warp( + const QueryT* __restrict__ query, + const DataT* __restrict__ d_dataset, + const float* __restrict__ d_inv_norms, + uint32_t id, + int dim, + int dim4, + bool use_inner_product, + int lane) { + if constexpr (USE_DP4A) { + const int32_t* vec_packed = reinterpret_cast( + d_dataset + static_cast(id) * dim); + const int32_t* query_packed = + reinterpret_cast(query); + int sum = 0; + for (int d = lane; d < dim4; d += 32) { + sum = __dp4a(query_packed[d], vec_packed[d], sum); + } + sum = warp_reduce_sum(sum); + float dist = static_cast(-sum); + if (d_inv_norms) + dist *= __ldg(&d_inv_norms[id]); + return dist; + } else { + const DataT* vec = d_dataset + static_cast(id) * dim; + if (use_inner_product) { + float sum = 0.0f; + for (int d = lane; d < dim; d += 32) { + sum += query[d] * load_elem(vec, d); + } + sum = warp_reduce_sum(sum); + float dist = -sum; + if (d_inv_norms) + dist *= __ldg(&d_inv_norms[id]); + return dist; + } else { + float sum = 0.0f; + for (int d = lane; d < dim; d += 32) { + float diff = query[d] - load_elem(vec, d); + sum += diff * diff; + } + return warp_reduce_sum(sum); + } + } +} + +// ============================================================================ +// Phase 1: Upper-layer greedy search +// ============================================================================ + +struct upper_layer_ptrs { + const uint32_t* d_node_ids; + const uint32_t* d_neighbors; + uint32_t num_nodes; + uint32_t max_degree; +}; + +__device__ __forceinline__ uint32_t +binary_search_node(const uint32_t* d_node_ids, uint32_t n, uint32_t global_id) { + uint32_t lo = 0, hi = n; + while (lo < hi) { + uint32_t mid = (lo + hi) / 2; + if (__ldg(&d_node_ids[mid]) < global_id) { + lo = mid + 1; + } else { + hi = mid; + } + } + if (lo < n && __ldg(&d_node_ids[lo]) == global_id) + return lo; + return UINT32_MAX; +} + +template +__global__ void upper_layer_search_kernel( + const float* __restrict__ d_queries, + const DataT* __restrict__ d_dataset, + const float* __restrict__ d_inv_norms, + const upper_layer_ptrs* __restrict__ d_layer_ptrs, + uint32_t* __restrict__ d_entry_points, + uint32_t global_entry_point, + int num_queries, + int dim, + int num_upper_layers, + bool use_inner_product) { + int warp_id = (blockIdx.x * blockDim.x + threadIdx.x) / 32; + int lane = threadIdx.x % 32; + if (warp_id >= num_queries) + return; + + const float* query = d_queries + static_cast(warp_id) * dim; + uint32_t current = global_entry_point; + + float best_dist; + if (lane == 0) { + if (use_inner_product) { + best_dist = thread_ip_distance( + query, + d_dataset + static_cast(current) * dim, + dim); + if (d_inv_norms) + best_dist *= __ldg(&d_inv_norms[current]); + } else { + best_dist = thread_l2_distance( + query, + d_dataset + static_cast(current) * dim, + dim); + } + } + best_dist = __shfl_sync(0xffffffff, best_dist, 0); + + for (int li = num_upper_layers - 1; li >= 0; li--) { + const upper_layer_ptrs& lp = d_layer_ptrs[li]; + bool improved = true; + while (improved) { + improved = false; + uint32_t local_idx = + binary_search_node(lp.d_node_ids, lp.num_nodes, current); + if (local_idx == UINT32_MAX) + break; + + uint32_t best_nbr = UINT32_MAX; + float best_nbr_dist = best_dist; + + for (uint32_t j = lane; j < lp.max_degree; j += 32) { + uint32_t nbr = lp.d_neighbors + [static_cast(local_idx) * + lp.max_degree + + j]; + float dist = FLT_MAX; + if (nbr != UINT32_MAX) { + const DataT* nbr_vec = + d_dataset + static_cast(nbr) * dim; + if (use_inner_product) { + dist = thread_ip_distance(query, nbr_vec, dim); + if (d_inv_norms) + dist *= d_inv_norms[nbr]; + } else { + dist = thread_l2_distance(query, nbr_vec, dim); + } + } + if (dist < best_nbr_dist) { + best_nbr_dist = dist; + best_nbr = nbr; + } + } + + for (int offset = 16; offset > 0; offset >>= 1) { + float other_dist = + __shfl_down_sync(0xffffffff, best_nbr_dist, offset); + uint32_t other_id = + __shfl_down_sync(0xffffffff, best_nbr, offset); + if (other_dist < best_nbr_dist) { + best_nbr_dist = other_dist; + best_nbr = other_id; + } + } + best_nbr_dist = __shfl_sync(0xffffffff, best_nbr_dist, 0); + best_nbr = __shfl_sync(0xffffffff, best_nbr, 0); + + if (best_nbr != UINT32_MAX && best_nbr_dist < best_dist) { + best_dist = best_nbr_dist; + current = best_nbr; + improved = true; + } + } + } + + if (lane == 0) { + d_entry_points[warp_id] = current; + } +} + +// ============================================================================ +// Phase 1: Parallel merge helpers (bitonic sort + parallel merge) +// ============================================================================ + +__device__ __forceinline__ void bitonic_sort_staging( + uint32_t* ids, float* dists, int active_count, int capacity) { + int i = threadIdx.x; + if (i < capacity && i >= active_count) { + ids[i] = UINT32_MAX; + dists[i] = FLT_MAX; + } + __syncthreads(); + + for (int k = 2; k <= capacity; k <<= 1) { + for (int j = k >> 1; j > 0; j >>= 1) { + int partner = i ^ j; + // Compare-exchange in two barrier-separated phases. Each thread + // owns writes to slot i only, but it READS slot `partner` — which + // another thread owns and may overwrite in the SAME phase. When the + // partner is in a different warp (j >= warpSize) the read and that + // write are unsynchronized -> a shared-memory hazard (racecheck + // flagged ~1.6M in the half kernel). So: (1) every thread reads its + // pair into registers and decides whether to swap, (2) __syncthreads + // so all reads complete before any write, (3) threads that must swap + // write their own slot, (4) __syncthreads before the next phase's + // reads. Loop bounds (k, j) are uniform across the block, so every + // thread — including i >= capacity — reaches every barrier. + float dp = 0.0f; + uint32_t ip = 0u; + bool do_swap = false; + if (i < capacity && partner < capacity) { + float di = dists[i]; + dp = dists[partner]; + uint32_t ii = ids[i]; + ip = ids[partner]; + // Ascending sort when the direction bit (i & k) is 0. + // In ascending order: lower index should hold the smaller value. + bool ascending = ((i & k) == 0); + bool i_is_lower = (i < partner); + // i should hold the smaller value iff: ascending and ipartner. + bool want_smaller = ascending ? i_is_lower : !i_is_lower; + bool i_is_smaller = (di < dp) || (di == dp && ii < ip); + do_swap = (want_smaller != i_is_smaller); + } + __syncthreads(); + if (do_swap) { + dists[i] = dp; + ids[i] = ip; + } + __syncthreads(); + } + } +} + +__device__ __forceinline__ void parallel_merge_into_result( + uint32_t* result_ids, + float* result_dists, + uint32_t* is_expanded, + uint32_t* staging_ids, + float* staging_dists, + uint32_t* merged_ids, + float* merged_dists, + uint32_t* merged_expanded, + int* meta, + int ef, + int max_staging) { + int sc = min(meta[1], max_staging); + int rc = meta[0]; + + bitonic_sort_staging(staging_ids, staging_dists, sc, max_staging); + + int T = rc + sc; + int merge_limit = min(T, ef); + + for (int pos = threadIdx.x; pos < merge_limit; pos += blockDim.x) { + // Binary search for si: number of staging elements ranked before pos. + // Valid range: si in [max(0,pos-rc), min(pos,sc)]. + // Half-open: lo inclusive, hi exclusive = min(pos,sc)+1. + // ri = pos-si is always >= 0 within this range, but after exit si + // may equal min(pos,sc)+1 (all staging ranked first); guard ri reads. + int lo = max(0, pos - rc); + int hi = min(pos, sc) + 1; + + while (lo < hi) { + int mid = (lo + hi) / 2; + // Merge-path diagonal test: staging[mid] is taken ahead of the + // result element it would displace, which is result[pos-mid-1] + // (NOT result[pos-mid]). Getting this index wrong mis-partitions + // every merge (empirically ~4% recall) and can promote a padded + // UINT32_MAX sentinel into the result list, which is later used as + // a graph/dataset index -> illegal memory access. + int ri_mid = pos - mid - 1; + bool ri_mid_valid = (ri_mid >= 0 && ri_mid < rc); + int ri_mid_safe = max(0, ri_mid); + float sv = (mid < sc) ? staging_dists[mid] : FLT_MAX; + uint32_t si_id = (mid < sc) ? staging_ids[mid] : UINT32_MAX; + float rv; + uint32_t ri_id; + if (ri_mid < 0) { + // No result element precedes this split: staging[mid] must not + // be taken ahead of it -> compare against -inf (predicate false). + rv = -FLT_MAX; + ri_id = 0u; + } else if (ri_mid_valid) { + rv = result_dists[ri_mid_safe]; + ri_id = result_ids[ri_mid_safe]; + } else { + // ri_mid >= rc: no result element there -> +inf (take staging). + rv = FLT_MAX; + ri_id = UINT32_MAX; + } + bool sv_le = (sv < rv) || (sv == rv && si_id < ri_id); + if (sv_le) + lo = mid + 1; + else + hi = mid; + } + + int si = lo; + // Clamp ri to 0 to prevent negative-index speculative loads; + // the ri_valid predicate gates whether the value is actually used. + int ri = pos - si; + bool ri_valid = (ri >= 0 && ri < rc); + int ri_safe = max(0, ri); + + float sv = (si < sc) ? staging_dists[si] : FLT_MAX; + float rv = ri_valid ? result_dists[ri_safe] : FLT_MAX; + uint32_t si_id = (si < sc) ? staging_ids[si] : UINT32_MAX; + uint32_t ri_id = ri_valid ? result_ids[ri_safe] : UINT32_MAX; + bool take_staging = (si < sc) && + (!ri_valid || (sv < rv) || (sv == rv && si_id < ri_id)); + + if (take_staging) { + merged_ids[pos] = si_id; + merged_dists[pos] = sv; + merged_expanded[pos] = 0; + } else { + merged_ids[pos] = ri_id; + merged_dists[pos] = rv; + merged_expanded[pos] = ri_valid ? is_expanded[ri_safe] : 0; + } + } + __syncthreads(); + + int new_rc = min(T, ef); + for (int i = threadIdx.x; i < new_rc; i += blockDim.x) { + result_ids[i] = merged_ids[i]; + result_dists[i] = merged_dists[i]; + is_expanded[i] = merged_expanded[i]; + } + if (threadIdx.x == 0) { + meta[0] = new_rc; + } + __syncthreads(); +} + +// ============================================================================ +// Phase 2: Layer-0 parallel beam search kernel +// ============================================================================ + +__device__ __forceinline__ bool bitmap_visit( + uint32_t* bitmap, + uint32_t node_id) { + uint32_t word = node_id >> 5; + uint32_t bit = 1u << (node_id & 31); + uint32_t old = atomicOr(&bitmap[word], bit); + return (old & bit) == 0; +} + +// True if row `id` is filtered OUT (deleted / not visible). Matches Knowhere's +// BitsetView byte/bit order: byte id/8, bit id%8, LSB-first. A null bitset (the +// unfiltered fast path) is never consulted — callers gate on HAS_FILTER. +__device__ __forceinline__ bool is_bitset_filtered( + const uint8_t* __restrict__ bitset, + uint32_t id) { + return (bitset[id >> 3] >> (id & 7)) & 1u; +} + +// Serial insert of (id, dist) into a distance-ascending beam of current size +// *count and capacity cap. Evicts the worst (last) element when full and the +// new distance is better; sets expanded[pos]=0 for the freshly inserted node +// and shifts existing expanded flags with their elements so an already-expanded +// node keeps its flag. MUST be called from a single thread (thread 0) — it is +// the serialized post-sort section that mirrors CPU NeighborSetPopList::insert +// (the sequential encounter order the alpha gate depends on). Returns true if +// the node was inserted. +__device__ __forceinline__ bool beam_insert_serial( + uint32_t* ids, + float* dists, + uint32_t* expanded, + int* count, + int cap, + uint32_t id, + float dist) { + int c = *count; + if (c >= cap && !(dist < dists[cap - 1])) { + return false; // full and not better than the current worst + } + // Number of elements retained ahead of the insert: drop the worst if full. + int n = (c < cap) ? c : cap - 1; + int p = n; + while (p > 0 && dists[p - 1] > dist) { + p--; + } + for (int i = n; i > p; i--) { + ids[i] = ids[i - 1]; + dists[i] = dists[i - 1]; + expanded[i] = expanded[i - 1]; + } + ids[p] = id; + dists[p] = dist; + expanded[p] = 0; + *count = (c < cap) ? (c + 1) : cap; + return true; +} + +// Two-tier merge for the filtered path (thread 0 only). Walks the already +// bitonic-sorted staging candidates in ascending-distance order and routes +// each one: +// - not filtered -> result beam (valid results, ef capacity) +// - filtered -> alpha gate; if it fires, admit to the invalid frontier +// only while closer than the valid beam's worst +// (at_search_back_dist: FLT_MAX until the valid beam is +// full, else result_dists[ef-1]). +// The alpha gate (accumulated_alpha, carried per query in meta[4]) rate-limits +// how often filtered nodes become waypoints, matching CPU +// NeighborSetDoublePopList. Deterministic and rate-limiting, but not +// bit-identical to CPU graph-neighbor encounter order — parity is asserted by +// the recall gate, not identical traversal. +__device__ __forceinline__ void filtered_merge_serial( + uint32_t* result_ids, + float* result_dists, + uint32_t* is_expanded, + uint32_t* invalid_ids, + float* invalid_dists, + uint32_t* invalid_exp, + uint32_t* staging_ids, + float* staging_dists, + int* meta, + const uint8_t* __restrict__ bitset, + int ef, + int ef_inv, + int max_staging, + float kAlpha) { + int sc = min(meta[1], max_staging); + int rc = meta[0]; + int ic = meta[3]; + float* alpha_ptr = reinterpret_cast(&meta[4]); + float accumulated_alpha = *alpha_ptr; + + for (int i = 0; i < sc; i++) { + uint32_t id = staging_ids[i]; + if (id == UINT32_MAX) + continue; + float d = staging_dists[i]; + if (!is_bitset_filtered(bitset, id)) { + beam_insert_serial( + result_ids, result_dists, is_expanded, &rc, ef, id, d); + } else { + accumulated_alpha += kAlpha; + if (accumulated_alpha < 1.0f) + continue; + accumulated_alpha -= 1.0f; + float valid_worst = (rc >= ef) ? result_dists[ef - 1] : FLT_MAX; + if (d < valid_worst) { + beam_insert_serial( + invalid_ids, + invalid_dists, + invalid_exp, + &ic, + ef_inv, + id, + d); + } + } + } + meta[0] = rc; + meta[3] = ic; + *alpha_ptr = accumulated_alpha; +} + +// Cross-beam parent selection for the filtered path (thread 0 only). Repeatedly +// takes the globally-closest UNEXPANDED node across the valid result beam and +// the invalid frontier (both distance-ascending), up to search_width parents. +// This is the GPU analog of CPU NeighborSetDoublePopList::pop_based_on_distance +// + has_next: an invalid node is only eligible while it is closer than the +// valid beam's worst (FLT_MAX until the valid beam fills). Returns the number +// of parents chosen; 0 means the search has converged. Filtered parents are +// expanded as waypoints but never enter the top-k (they live only in the +// invalid frontier). +__device__ __forceinline__ int select_parents_cross_beam( + uint32_t* result_ids, + float* result_dists, + uint32_t* is_expanded, + int rc, + uint32_t* invalid_ids, + float* invalid_dists, + uint32_t* invalid_exp, + int ic, + uint32_t* parent_ids, + int ef, + int search_width) { + int num_parents = 0; + float valid_worst = (rc >= ef) ? result_dists[ef - 1] : FLT_MAX; + int ri = 0, ii = 0; + while (num_parents < search_width) { + while (ri < rc && is_expanded[ri]) + ri++; + while (ii < ic && invalid_exp[ii]) + ii++; + bool hasR = (ri < rc); + bool hasI = (ii < ic) && (invalid_dists[ii] < valid_worst); + if (!hasR && !hasI) + break; + if (hasR && (!hasI || result_dists[ri] <= invalid_dists[ii])) { + parent_ids[num_parents++] = result_ids[ri]; + is_expanded[ri] = 1; + ri++; + } else { + parent_ids[num_parents++] = invalid_ids[ii]; + invalid_exp[ii] = 1; + ii++; + } + } + return num_parents; +} + +// Templated on the dataset element type (DataT), the query element type +// (QueryT: float for the generic path, int8_t for the native DP4A path), and +// USE_DP4A which selects the DP4A int8 inner-product distance at compile time. +template +__global__ void layer0_beam_search_kernel( + const QueryT* __restrict__ d_queries, + const DataT* __restrict__ d_dataset, + const float* __restrict__ d_inv_norms, + const uint32_t* __restrict__ d_layer0_graph, + const uint32_t* __restrict__ d_entry_points, + uint32_t* __restrict__ d_visited_bitmaps, + uint64_t* __restrict__ d_neighbors, + float* __restrict__ d_distances, + int num_queries, + int N, + int dim, + int max_degree0, + int k, + int ef, + int search_width, + int max_iterations, + bool use_inner_product, + // --- Filtered-search args (ignored when HAS_FILTER is false; the + // branches are compiled out via `if constexpr` so the unfiltered + // instantiation is codegen-identical to the append-only kernel). --- + const uint8_t* __restrict__ d_bitset, + int ef_inv, + float kAlpha, + int live_count, + bool disable_bf, + uint32_t* __restrict__ d_needs_bf, + int* __restrict__ d_needs_bf_count) { + int query_idx = blockIdx.x; + if (query_idx >= num_queries) + return; + + const QueryT* query = d_queries + static_cast(query_idx) * dim; + int dim4 = dim / 4; + + // Warp-cooperative neighbor expansion: one warp owns one candidate edge and + // its 32 lanes compute the distance together (coalesced loads). The host + // rounds block_size up to a whole number of warps so every lane below is + // part of a full 32-lane warp (required for the full-mask __shfl reductions + // in layer0_distance_warp). + const int lane = threadIdx.x & 31; + const int warp_in_block = threadIdx.x >> 5; + const int num_warps = blockDim.x >> 5; + + extern __shared__ char smem[]; + + // Pad to a power of two so the bitonic sort in parallel_merge_into_result + // has a valid capacity for any max_degree0 (e.g. non-power-of-two 2*M). + int max_staging = padded_staging_capacity(search_width, max_degree0); + + // The filtered path carries 3 extra meta slots (invalid count + alpha + + // spare) and appends an ef_inv-sized invalid frontier after the merge + // scratch. The unfiltered path keeps exactly today's layout. + constexpr int kMetaSlots = HAS_FILTER ? 6 : 3; + + uint32_t* result_ids = reinterpret_cast(smem); + float* result_dists = reinterpret_cast(result_ids + ef); + uint32_t* is_expanded = reinterpret_cast(result_dists + ef); + uint32_t* staging_ids = is_expanded + ef; + float* staging_dists = reinterpret_cast(staging_ids + max_staging); + uint32_t* parent_ids = + reinterpret_cast(staging_dists + max_staging); + int* meta = reinterpret_cast(parent_ids + search_width); + uint32_t* merged_ids = reinterpret_cast(meta + kMetaSlots); + float* merged_dists = reinterpret_cast(merged_ids + ef); + uint32_t* merged_expanded = reinterpret_cast(merged_dists + ef); + // Invalid frontier (filtered path only): filtered nodes retained as + // waypoints. Never emitted; expanded when they are the globally-closest + // unexpanded node. + uint32_t* invalid_ids = merged_expanded + ef; + float* invalid_dists = reinterpret_cast(invalid_ids + ef_inv); + uint32_t* invalid_exp = reinterpret_cast(invalid_dists + ef_inv); + + int bitmap_words = (N + 31) / 32; + uint32_t* visited_bmap = + d_visited_bitmaps + static_cast(query_idx) * bitmap_words; + + for (int i = threadIdx.x; i < ef; i += blockDim.x) { + result_ids[i] = UINT32_MAX; + result_dists[i] = FLT_MAX; + is_expanded[i] = 0; + } + if (threadIdx.x == 0) { + meta[0] = 0; + meta[1] = 0; + meta[2] = 0; + if constexpr (HAS_FILTER) { + meta[3] = 0; // invalid frontier count + // accumulated_alpha, carried per query in meta[4], init 1.0f to + // match CPU NeighborSetDoublePopList (first filtered node always + // admitted once the gate accumulates). + *reinterpret_cast(&meta[4]) = 1.0f; + } + } + __syncthreads(); + + // --- Seed with entry point --- + uint32_t ep = d_entry_points[query_idx]; + // Defensive: the entry point must be a real node (< N) before it indexes + // the dataset (ep_dist) or the layer-0 graph (neighbor seeding). A bad ep + // would otherwise fault far out of bounds. On the healthy path ep is always + // valid; if it is not, seed nothing and let the search return sentinels. + bool ep_valid = (ep < static_cast(N)); + if (threadIdx.x == 0) { + if (ep_valid) { + float ep_dist = layer0_distance( + query, d_dataset, d_inv_norms, ep, dim, dim4, + use_inner_product); + if constexpr (HAS_FILTER) { + // A filtered entry point is a waypoint only: seed it into the + // invalid frontier so it is never emitted but can still be + // expanded. Otherwise seed the valid result beam as usual. + if (is_bitset_filtered(d_bitset, ep)) { + invalid_ids[0] = ep; + invalid_dists[0] = ep_dist; + invalid_exp[0] = 0; + meta[3] = 1; + meta[0] = 0; + } else { + result_ids[0] = ep; + result_dists[0] = ep_dist; + is_expanded[0] = 0; + meta[0] = 1; + } + } else { + result_ids[0] = ep; + result_dists[0] = ep_dist; + is_expanded[0] = 0; + meta[0] = 1; + } + bitmap_visit(visited_bmap, ep); + } else { + meta[0] = 0; + } + } + __syncthreads(); + + // --- Seed with entry point's neighbors --- + if (threadIdx.x == 0) + meta[1] = 0; + __syncthreads(); + + for (int j = warp_in_block; ep_valid && j < max_degree0; j += num_warps) { + // Lane 0 resolves the neighbor and claims the first-visit; the go/no-go + // and id are broadcast so the whole warp stays convergent for the + // cooperative distance below. `proceed` is warp-uniform, so the + // `continue` never splits the warp (no __shfl hazard). + uint32_t nbr = UINT32_MAX; + bool proceed = false; + if (lane == 0) { + uint32_t cand = + d_layer0_graph[static_cast(ep) * max_degree0 + j]; + if (cand != UINT32_MAX && cand < static_cast(N) && + bitmap_visit(visited_bmap, cand)) { + nbr = cand; + proceed = true; + } + } + proceed = __shfl_sync(0xffffffff, proceed, 0); + if (!proceed) + continue; + nbr = __shfl_sync(0xffffffff, nbr, 0); + + float dist = layer0_distance_warp( + query, d_dataset, d_inv_norms, nbr, dim, dim4, + use_inner_product, lane); + + if (lane == 0) { + int slot = atomicAdd(&meta[1], 1); + if (slot < max_staging) { + staging_ids[slot] = nbr; + staging_dists[slot] = dist; + } + } + } + __syncthreads(); + + if constexpr (HAS_FILTER) { + bitonic_sort_staging( + staging_ids, + staging_dists, + min(meta[1], max_staging), + max_staging); + if (threadIdx.x == 0) { + filtered_merge_serial( + result_ids, + result_dists, + is_expanded, + invalid_ids, + invalid_dists, + invalid_exp, + staging_ids, + staging_dists, + meta, + d_bitset, + ef, + ef_inv, + max_staging, + kAlpha); + } + } else { + parallel_merge_into_result( + result_ids, + result_dists, + is_expanded, + staging_ids, + staging_dists, + merged_ids, + merged_dists, + merged_expanded, + meta, + ef, + max_staging); + } + __syncthreads(); + // Mark the entry point expanded (its neighbors were already seeded above). + // It lives in the valid beam, or in the invalid frontier if it was + // filtered. + if (threadIdx.x == 0) { + int rc = meta[0]; + bool found = false; + for (int i = 0; i < rc; i++) { + if (result_ids[i] == ep) { + is_expanded[i] = 1; + found = true; + break; + } + } + if constexpr (HAS_FILTER) { + if (!found) { + int ic = meta[3]; + for (int i = 0; i < ic; i++) { + if (invalid_ids[i] == ep) { + invalid_exp[i] = 1; + break; + } + } + } + } + } + __syncthreads(); + + // --- Unified main loop --- + for (int iter = 0; iter < max_iterations; iter++) { + if (threadIdx.x == 0) { + int num_parents = 0; + int rc = meta[0]; + + if constexpr (HAS_FILTER) { + num_parents = select_parents_cross_beam( + result_ids, + result_dists, + is_expanded, + rc, + invalid_ids, + invalid_dists, + invalid_exp, + meta[3], + parent_ids, + ef, + search_width); + } else { + for (int i = 0; i < rc && num_parents < search_width; i++) { + if (!is_expanded[i]) { + parent_ids[num_parents++] = result_ids[i]; + is_expanded[i] = 1; + } + } + } + + meta[2] = num_parents; + } + __syncthreads(); + + int num_parents = meta[2]; + if (num_parents == 0) + break; + + if (threadIdx.x == 0) + meta[1] = 0; + __syncthreads(); + + int total_work = num_parents * max_degree0; + for (int wi = warp_in_block; wi < total_work; wi += num_warps) { + int parent_idx = wi / max_degree0; + int nbr_slot = wi % max_degree0; + + uint32_t parent = parent_ids[parent_idx]; + + // Lane 0 resolves the neighbor id, bounds-checks it, and claims the + // first-visit; the go/no-go and id are broadcast so the warp stays + // convergent for the cooperative distance. All predicates below are + // warp-uniform (`parent` is identical across lanes), so the + // `continue` never splits the warp and the __shfl calls in + // layer0_distance_warp always see a full 32-lane mask. + uint32_t nbr = UINT32_MAX; + bool proceed = false; + if (lane == 0) { + // Defensive: a parent id must be a real node (< N). result_ids + // only ever holds the entry point or graph neighbors (both < N) + // once the merge is correct, so this never fires on the healthy + // path; it caps a stray/garbage id at O(N) instead of letting it + // index the graph at parent*max_degree0 and fault far out of + // bounds. Mirrors the neighbor guard below. + if (parent < static_cast(N)) { + uint32_t cand = d_layer0_graph + [static_cast(parent) * max_degree0 + + nbr_slot]; + if (cand != UINT32_MAX && cand < static_cast(N) && + bitmap_visit(visited_bmap, cand)) { + nbr = cand; + proceed = true; + } + } + } + proceed = __shfl_sync(0xffffffff, proceed, 0); + if (!proceed) + continue; + nbr = __shfl_sync(0xffffffff, nbr, 0); + + float dist = layer0_distance_warp( + query, d_dataset, d_inv_norms, nbr, dim, dim4, + use_inner_product, lane); + + if (lane == 0) { + int slot = atomicAdd(&meta[1], 1); + if (slot < max_staging) { + staging_ids[slot] = nbr; + staging_dists[slot] = dist; + } + } + } + __syncthreads(); + + if constexpr (HAS_FILTER) { + bitonic_sort_staging( + staging_ids, + staging_dists, + min(meta[1], max_staging), + max_staging); + if (threadIdx.x == 0) { + filtered_merge_serial( + result_ids, + result_dists, + is_expanded, + invalid_ids, + invalid_dists, + invalid_exp, + staging_ids, + staging_dists, + meta, + d_bitset, + ef, + ef_inv, + max_staging, + kAlpha); + } + __syncthreads(); + } else { + parallel_merge_into_result( + result_ids, + result_dists, + is_expanded, + staging_ids, + staging_dists, + merged_ids, + merged_dists, + merged_expanded, + meta, + ef, + max_staging); + } + + // Stagnation detected when num_parents==0 (all expanded) → break above + } + + // --- Copy top-k results to global memory --- + // The beam search works internally with a "smaller == closer" convention: + // thread_ip_distance() returns the negated dot product so inner-product and + // cosine can share the same min-first ordering as L2. Negate back on + // copy-out for those metrics so callers receive the true similarity score + // (faiss's METRIC_INNER_PRODUCT contract is "larger == closer"). L2 is + // already a genuine distance and is copied verbatim. Empty slots use the + // worst score for the metric (-FLT_MAX for IP/cosine, +FLT_MAX for L2) so + // they always sort last regardless of the caller's ordering. + int rc = meta[0]; + for (int i = threadIdx.x; i < k; i += blockDim.x) { + if (i < rc && result_ids[i] < static_cast(N)) { + d_neighbors[static_cast(query_idx) * k + i] = + static_cast(result_ids[i]); + d_distances[static_cast(query_idx) * k + i] = + use_inner_product ? -result_dists[i] : result_dists[i]; + } else { + d_neighbors[static_cast(query_idx) * k + i] = UINT64_MAX; + d_distances[static_cast(query_idx) * k + i] = + use_inner_product ? -FLT_MAX : FLT_MAX; + } + } + + // --- Per-query brute-force fallback (filtered path) --- + // If the filtered graph search returned fewer than k valid results while + // more than that many live rows exist, this query short-fell (its near + // neighbors were filtered out). Append it to the device worklist for a + // batched BF pass, matching CPU HNSW's per-query fallback. Skipped when + // disable_bf mirrors hnsw_cfg.disable_fallback_brute_force (short queries + // then keep their padded sentinels). + if constexpr (HAS_FILTER) { + if (threadIdx.x == 0 && !disable_bf) { + int rc = meta[0]; + int real_topk = min(rc, k); + if (real_topk < k && real_topk < live_count) { + int pos = atomicAdd(d_needs_bf_count, 1); + d_needs_bf[pos] = static_cast(query_idx); + } + } + } +} + +inline size_t calc_layer0_smem_size( + int ef, + int search_width, + int max_degree0, + int ef_inv = 0) { + // Must match the kernel's padded staging capacity exactly. + int max_staging = padded_staging_capacity(search_width, max_degree0); + + // ef_inv > 0 selects the filtered layout: 3 extra meta slots and an + // ef_inv-sized invalid frontier (12 B / ef_inv). ef_inv == 0 is the + // unfiltered layout, byte-identical to the append-only kernel. + bool has_filter = (ef_inv > 0); + int meta_slots = has_filter ? 6 : 3; + + size_t size = 0; + size += ef * sizeof(uint32_t); // result_ids + size += ef * sizeof(float); // result_dists + size += ef * sizeof(uint32_t); // is_expanded + size += max_staging * sizeof(uint32_t); // staging_ids + size += max_staging * sizeof(float); // staging_dists + size += search_width * sizeof(uint32_t); // parent_ids + size += meta_slots * sizeof(int); // meta + size += ef * sizeof(uint32_t); // merged_ids + size += ef * sizeof(float); // merged_dists + size += ef * sizeof(uint32_t); // merged_expanded + if (has_filter) { + size += ef_inv * sizeof(uint32_t); // invalid_ids + size += ef_inv * sizeof(float); // invalid_dists + size += ef_inv * sizeof(uint32_t); // invalid_exp + } + return size; +} + +inline size_t calc_visited_bitmap_size(int num_queries, int N) { + int bitmap_words = (N + 31) / 32; + return static_cast(num_queries) * bitmap_words * sizeof(uint32_t); +} + +} // namespace hnsw_kernel +} // namespace gpu +} // namespace faiss diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu new file mode 100644 index 000000000..1d6ca4d2b --- /dev/null +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu @@ -0,0 +1,237 @@ +// @lint-ignore-every LICENSELINT +/** + * Copyright (c) Facebook, Inc. and its affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ +/* + * Copyright (c) 2026, 6sense Insights Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include + +#include +#include + +namespace faiss { +namespace gpu { + +namespace { + +inline void check_cuda(cudaError_t err, const char* file, int line) { + if (err != cudaSuccess) { + throw std::runtime_error( + std::string("CUDA error in scratch alloc: ") + + cudaGetErrorString(err) + " at " + file + ":" + + std::to_string(line)); + } +} + +#define SCRATCH_CUDA_CHECK(expr) check_cuda((expr), __FILE__, __LINE__) + +} // namespace + +void GpuHnswSearchScratch::ensure( + int nq, + int k, + int dim, + int N, + bool use_i8_queries) { + size_t need_q = static_cast(nq) * dim * sizeof(float); + if (need_q > queries_bytes) { + if (d_queries) + cudaFree(d_queries); + SCRATCH_CUDA_CHECK(cudaMalloc(&d_queries, need_q)); + queries_bytes = need_q; + } + size_t need_n = static_cast(nq) * k * sizeof(uint64_t); + if (need_n > neighbors_bytes) { + if (d_neighbors) + cudaFree(d_neighbors); + SCRATCH_CUDA_CHECK(cudaMalloc(&d_neighbors, need_n)); + neighbors_bytes = need_n; + } + size_t need_d = static_cast(nq) * k * sizeof(float); + if (need_d > distances_bytes) { + if (d_distances) + cudaFree(d_distances); + SCRATCH_CUDA_CHECK(cudaMalloc(&d_distances, need_d)); + distances_bytes = need_d; + } + if (nq > entry_cap) { + if (d_entry_points) + cudaFree(d_entry_points); + SCRATCH_CUDA_CHECK(cudaMalloc( + &d_entry_points, + static_cast(nq) * sizeof(uint32_t))); + entry_cap = nq; + } + // The visited bitmap is only ever indexed by the chunk-local query index + // (see the launch loop in gpu_hnsw_search), so size it for one chunk rather + // than the full batch. This is the fix for the grow-only OOM: without the + // chunk cap, a large nq (or high-concurrency high-water mark) sized this at + // nq * ceil(N/32) * 4 and never released it. Must match the chunk used at + // launch time. + int bm_nq = gpu_hnsw_bitmap_chunk(nq, N); + int bitmap_words = (N + 31) / 32; + size_t need_bm = + static_cast(bm_nq) * bitmap_words * sizeof(uint32_t); + if (need_bm > bitmap_bytes) { + if (d_visited_bitmaps) + cudaFree(d_visited_bitmaps); + SCRATCH_CUDA_CHECK(cudaMalloc(&d_visited_bitmaps, need_bm)); + bitmap_bytes = need_bm; + } + if (use_i8_queries) { + size_t need_i8 = static_cast(nq) * dim * sizeof(int8_t); + if (need_i8 > queries_i8_bytes) { + if (d_queries_i8) + cudaFree(d_queries_i8); + SCRATCH_CUDA_CHECK(cudaMalloc(&d_queries_i8, need_i8)); + queries_i8_bytes = need_i8; + } + } +} + +void GpuHnswSearchScratch::ensure_filter(int nq, size_t bitset_bytes_needed) { + // Bind this scratch slot's owning device before allocating so cudaMalloc + // lands on the right GPU even if the active device was changed elsewhere + // (matches the destructor). searchHost already sets it, but keep the + // allocation self-contained on multi-GPU systems. + cudaSetDevice(device); + // Device bitset: reallocate when the segment row count (hence the byte + // count) grows. ensure() is called per search with the current n_rows, so + // a segment that gains rows via reload re-sizes the bitset here. + if (bitset_bytes_needed > bitset_bytes) { + if (d_bitset) + cudaFree(d_bitset); + SCRATCH_CUDA_CHECK(cudaMalloc(&d_bitset, bitset_bytes_needed)); + bitset_bytes = bitset_bytes_needed; + } + // needs_bf worklist: one uint32 query index per query, plus a single + // atomic counter. Sized to nq (grow-only). + if (nq > needs_bf_cap) { + if (d_needs_bf) + cudaFree(d_needs_bf); + SCRATCH_CUDA_CHECK(cudaMalloc( + &d_needs_bf, static_cast(nq) * sizeof(uint32_t))); + needs_bf_cap = nq; + } + if (d_needs_bf_count == nullptr) { + SCRATCH_CUDA_CHECK(cudaMalloc(&d_needs_bf_count, sizeof(int))); + } +} + +GpuHnswSearchScratch::~GpuHnswSearchScratch() { + // Set the owning device before freeing so cudaFree runs in the correct + // context on multi-GPU systems. + cudaSetDevice(device); + if (d_queries) + cudaFree(d_queries); + if (d_neighbors) + cudaFree(d_neighbors); + if (d_distances) + cudaFree(d_distances); + if (d_entry_points) + cudaFree(d_entry_points); + if (d_visited_bitmaps) + cudaFree(d_visited_bitmaps); + if (d_queries_i8) + cudaFree(d_queries_i8); + if (d_bitset) + cudaFree(d_bitset); + if (d_needs_bf) + cudaFree(d_needs_bf); + if (d_needs_bf_count) + cudaFree(d_needs_bf_count); +} + +GpuHnswScratchSlot::~GpuHnswScratchSlot() { + // Set the owning device before destroying the stream so it runs in the + // correct context on multi-GPU systems (scratch.device is assigned when the + // slot is created in GpuHnswScratchPool::init_once()). + cudaSetDevice(scratch.device); + if (stream) + cudaStreamDestroy(stream); +} + +GpuHnswScratchPool::GpuHnswScratchPool(int pool_size, int device) + : pool_size_(pool_size), device_(device) {} + +void GpuHnswScratchPool::init_once() { + if (initialized_) + return; + slots_.reserve(pool_size_); + available_.reserve(pool_size_); + for (int i = 0; i < pool_size_; i++) { + auto slot = std::make_unique(); + slot->scratch.device = device_; + SCRATCH_CUDA_CHECK(cudaSetDevice(device_)); + SCRATCH_CUDA_CHECK( + cudaStreamCreateWithFlags(&slot->stream, cudaStreamNonBlocking)); + available_.push_back(slot.get()); + slots_.push_back(std::move(slot)); + } + initialized_ = true; +} + +GpuHnswScratchSlot* GpuHnswScratchPool::acquire() { + std::unique_lock lock(mutex_); + init_once(); + cv_.wait(lock, [this] { return !available_.empty(); }); + auto* slot = available_.back(); + available_.pop_back(); + return slot; +} + +void GpuHnswScratchPool::release(GpuHnswScratchSlot* slot) { + // Safety net: ensure all GPU work on this slot's stream has completed + // before returning it to the pool. Callers (searchHost / searchImpl_) + // already synchronize before release, so this is normally a no-op; it + // guards a future caller that skips its own sync, whose in-flight buffers + // could otherwise be freed+realloced by the next acquirer's ensure() + // (use-after-free). + cudaSetDevice(device_); + cudaStreamSynchronize(slot->stream); + { + std::lock_guard lock(mutex_); + available_.push_back(slot); + } + cv_.notify_one(); +} + +GpuHnswDeviceIndex::~GpuHnswDeviceIndex() { + // Set the owning device before freeing so cudaFree runs in the correct + // context on multi-GPU systems. + cudaSetDevice(device); + if (d_dataset) + cudaFree(d_dataset); + if (d_inv_norms) + cudaFree(d_inv_norms); + if (d_layer0_graph) + cudaFree(d_layer0_graph); + for (auto& ul : upper_layers) { + if (ul.d_node_ids) + cudaFree(ul.d_node_ids); + if (ul.d_neighbors) + cudaFree(ul.d_neighbors); + } + if (d_upper_layer_ptrs) + cudaFree(d_upper_layer_ptrs); +} + +} // namespace gpu +} // namespace faiss diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h new file mode 100644 index 000000000..3778da7ba --- /dev/null +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h @@ -0,0 +1,309 @@ +// @lint-ignore-every LICENSELINT +/** + * Copyright (c) Facebook, Inc. and its affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ +/* + * Copyright (c) 2026, 6sense Insights Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace faiss { +namespace gpu { + +// Test-only fault injection for the device-upload path (consulted by +// GPU_HNSW_BUILD_CUDA_CHECK in GpuHnswBuild.cuh). Production code never arms it +// (the countdown stays 0), so the check is a single relaxed atomic load per +// wrapped CUDA call with no runtime effect. Unit tests call arm(n) so the n-th +// subsequent wrapped CUDA call reports a simulated failure, exercising +// Deserialize()'s upload error / partial-allocation cleanup / retry path +// without a real OOM. Defined here (host-safe header) rather than in the .cuh so +// host-compiled tests can arm it without pulling in device kernels. The static +// lives in an inline function, so all translation units share one instance. +struct GpuHnswUploadFaultInjection { + static std::atomic& countdown() { + static std::atomic c{0}; + return c; + } + // Arm so the n-th following wrapped CUDA call fails (n >= 1). 0 disarms. + static void arm(int n) { + countdown().store(n, std::memory_order_relaxed); + } + static void disarm() { + countdown().store(0, std::memory_order_relaxed); + } + // Returns true exactly once, on the n-th call after arm(n); self-disarms. + static bool should_fail() { + int cur = countdown().load(std::memory_order_relaxed); + if (cur <= 0) + return false; + return countdown().fetch_sub(1, std::memory_order_relaxed) == 1; + } +}; + +// Element type of the device-resident dataset. The graph walk kernel is +// templated on this; each value selects a load_elem specialization so the +// vectors stay in their native precision on the GPU (no up-conversion to +// fp32 at upload time). +enum class GpuHnswDatasetType { + FP32 = 0, + INT8 = 1, + FP16 = 2, + BF16 = 3, +}; + +// Cap (bytes) on the per-search visited bitmap for a single scratch slot. The +// bitmap is nq * ceil(N/32) * 4 bytes and is grow-only per slot, so a large +// query batch or high search concurrency (one bitmap per pool slot) can grow it +// until it exhausts device memory. Observed regression: 16 concurrent batch=512 +// searches on a 538M-row segment grew the pool to ~97 of ~98 GB and OOM'd every +// subsequent allocation. Bounding the bitmap and processing queries in +// nq-chunks caps it regardless of batch size / concurrency. +// +// Tunable via environment (re-read on each call so tests can toggle it; getenv +// is negligible next to a GPU search): GPU_HNSW_BITMAP_BYTES sets the cap in +// bytes (used by tests to force the multi-chunk path on tiny inputs) and takes +// precedence over GPU_HNSW_BITMAP_MB (megabytes). Default 256 MiB. The value is +// only a chunking bound, so any positive value is safe (the chunk is clamped to +// >= 1 query regardless). +inline size_t gpu_hnsw_bitmap_cap_bytes() { + if (const char* eb = std::getenv("GPU_HNSW_BITMAP_BYTES")) { + long long v = std::atoll(eb); + if (v > 0) { + return static_cast(v); + } + } + size_t mb = 256; + if (const char* e = std::getenv("GPU_HNSW_BITMAP_MB")) { + long v = std::atol(e); + if (v > 0) { + mb = static_cast(v); + } + } + return mb * (static_cast(1) << 20); +} + +// Number of queries to process per launch so the visited bitmap stays within +// gpu_hnsw_bitmap_cap_bytes(). Always in [1, nq]; returns nq when the whole +// batch already fits, so small segments are chunked into a single pass and see +// no behavior change. Queries are independent in HNSW search, so chunking does +// not change per-query results. Must stay in lockstep with the bitmap sizing in +// GpuHnswSearchScratch::ensure() and the launch loop in gpu_hnsw_search(). +inline int gpu_hnsw_bitmap_chunk(int nq, int N) { + if (nq <= 0) { + return nq; + } + size_t per_query = + static_cast((N + 31) / 32) * sizeof(uint32_t); + if (per_query == 0) { + return nq; + } + size_t cap_q = gpu_hnsw_bitmap_cap_bytes() / per_query; + if (cap_q < 1) { + cap_q = 1; + } + return (cap_q >= static_cast(nq)) ? nq + : static_cast(cap_q); +} + +struct GpuHnswSearchParams { + int ef = 200; + int search_width = 4; + int max_iterations = 0; + int thread_block_size = 0; + + // --- Filtered search (deletes / TTL / partition / visibility bitset) --- + // + // When bitset_data == nullptr the search runs the unfiltered fast path, + // which is codegen-identical to the append-only kernel (the filter + // branches are compiled out via `if constexpr`). When non-null it enables + // CPU-HNSW-parity filtered search: filtered nodes stay graph waypoints but + // are never emitted, an alpha gate rate-limits their expansion, and a + // brute-force fallback guarantees k live results under heavy deletes. + // + // Semantics match Knowhere's BitsetView: a set bit means the row at that + // index is filtered OUT (deleted / not visible). The bit order is + // LSB-first per byte: row r is byte r/8, bit r%8. The index space is the + // storage add-order row id, which the GPU returns directly (see the design + // doc's ID-mapping section and the gpu_hnsw_id_mapping test). + const uint8_t* bitset_data = nullptr; // host ptr, uploaded per search batch + int64_t bitset_nbits = 0; // rows the bitset covers (== n_rows) + int64_t bitset_filtered_count = 0; // popcount: rows filtered out + // Invalid-frontier capacity. 0 => default to ef. Cappable to bound the + // extra shared memory (12 B / ef_inv on top of the 24 B / ef base). + int ef_inv = 0; + // Mirrors Knowhere's hnsw_cfg.disable_fallback_brute_force. When true the + // short-result BF fallback is skipped and short queries keep padded + // sentinels, matching CPU HNSW with fallback disabled. + bool disable_fallback_brute_force = false; +}; + +struct GpuHnswDeviceUpperLayer { + uint32_t* d_node_ids = nullptr; + uint32_t* d_neighbors = nullptr; + uint32_t num_nodes = 0; + uint32_t max_degree = 0; +}; + +struct GpuHnswSearchScratch { + float* d_queries = nullptr; + uint64_t* d_neighbors = nullptr; + float* d_distances = nullptr; + uint32_t* d_entry_points = nullptr; + uint32_t* d_visited_bitmaps = nullptr; + int8_t* d_queries_i8 = nullptr; // int8 queries for the native DP4A path + + // Filtered-search scratch. d_bitset holds the uploaded BitsetView bytes + // (ceil(nbits/8)); d_needs_bf is the device-side worklist of query indices + // whose graph search returned fewer than k live results (per-query BF + // fallback), and d_needs_bf_count is its atomic length. All allocated + // lazily and only when a filtered search runs. + uint8_t* d_bitset = nullptr; + uint32_t* d_needs_bf = nullptr; + int* d_needs_bf_count = nullptr; + + size_t queries_bytes = 0; + size_t neighbors_bytes = 0; + size_t distances_bytes = 0; + int entry_cap = 0; + size_t bitmap_bytes = 0; + size_t queries_i8_bytes = 0; + size_t bitset_bytes = 0; + int needs_bf_cap = 0; + + // Device this scratch's allocations live on; used to set the CUDA device + // context before freeing in the destructor (multi-GPU correctness). + int device = 0; + + void ensure(int nq, int k, int dim, int N, bool use_i8_queries = false); + + // Ensure the per-search filter scratch is large enough: a device bitset of + // `bitset_bytes_needed` bytes (reallocated when n_rows grows) and a + // needs_bf worklist sized for `nq` queries plus its counter. Separate from + // ensure() because filtering is optional and the bitset size tracks the + // segment row count, not nq/k/dim. + void ensure_filter(int nq, size_t bitset_bytes_needed); + + ~GpuHnswSearchScratch(); + + GpuHnswSearchScratch() = default; + GpuHnswSearchScratch(const GpuHnswSearchScratch&) = delete; + GpuHnswSearchScratch& operator=(const GpuHnswSearchScratch&) = delete; +}; + +/// One slot in the scratch pool: a scratch buffer + its own CUDA stream. +struct GpuHnswScratchSlot { + GpuHnswSearchScratch scratch; + cudaStream_t stream = nullptr; + + ~GpuHnswScratchSlot(); + GpuHnswScratchSlot() = default; + GpuHnswScratchSlot(const GpuHnswScratchSlot&) = delete; + GpuHnswScratchSlot& operator=(const GpuHnswScratchSlot&) = delete; +}; + +/// Pool of scratch buffers allowing concurrent GPU searches on the same segment. +/// Each acquire() returns a slot with its own scratch + CUDA stream. +/// Callers block if all slots are in use. +class GpuHnswScratchPool { + public: + /// Create a pool. CUDA streams are allocated lazily on first acquire(). + explicit GpuHnswScratchPool(int pool_size = 4, int device = 0); + ~GpuHnswScratchPool() = default; + + GpuHnswScratchPool(const GpuHnswScratchPool&) = delete; + GpuHnswScratchPool& operator=(const GpuHnswScratchPool&) = delete; + + /// Acquire a scratch slot (blocks until one is available). + GpuHnswScratchSlot* acquire(); + /// Release a previously acquired scratch slot back to the pool. + void release(GpuHnswScratchSlot* slot); + + int pool_size() const { return pool_size_; } + + private: + void init_once(); + + std::mutex mutex_; + std::condition_variable cv_; + int pool_size_; + int device_; + bool initialized_ = false; + std::vector> slots_; + std::vector available_; +}; + +/// RAII guard: acquires a scratch slot on construction, releases on destruction. +class ScratchPoolGuard { + public: + ScratchPoolGuard(GpuHnswScratchPool& pool) + : pool_(pool), slot_(pool.acquire()) {} + ~ScratchPoolGuard() { pool_.release(slot_); } + GpuHnswScratchSlot* get() const { return slot_; } + + ScratchPoolGuard(const ScratchPoolGuard&) = delete; + ScratchPoolGuard& operator=(const ScratchPoolGuard&) = delete; + + private: + GpuHnswScratchPool& pool_; + GpuHnswScratchSlot* slot_; +}; + +struct GpuHnswDeviceIndex { + void* d_dataset = nullptr; + GpuHnswDatasetType dataset_type = GpuHnswDatasetType::FP32; + float* d_inv_norms = nullptr; + uint32_t* d_layer0_graph = nullptr; + std::vector upper_layers; + + int64_t n_rows = 0; + int64_t dim = 0; + uint32_t entry_point = 0; + // Total layer count (max_level + 1), informational only; the search path + // uses num_upper_layers_built (below) for the number of uploaded upper + // layers. Kept for diagnostics/logging. + int num_layers = 0; + int M = 0; + int max_degree0 = 0; + bool use_ip = false; + + void* d_upper_layer_ptrs = nullptr; + int num_upper_layers_built = 0; + + // Device this index's allocations live on; used to set the CUDA device + // context before freeing in the destructor (multi-GPU correctness). + int device = 0; + + mutable std::unique_ptr scratch_pool; + + ~GpuHnswDeviceIndex(); +}; + +} // namespace gpu +} // namespace faiss