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..828a47817 100644 --- a/cmake/libs/libfaiss.cmake +++ b/cmake/libs/libfaiss.cmake @@ -593,3 +593,26 @@ if(__PPC64) knowhere_utils) target_compile_definitions(faiss PRIVATE FINTEGER=int) endif() + +# GPU HNSW CUDA sources — compiled when WITH_CUVS is enabled +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/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..64ef7d5f4 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,543 @@ 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. + (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. + Status status; + if (!binset.Contains(IndexEnum::INDEX_GPU_HNSW)) { + BinarySet aliased = binset; + for (const char* key : {IndexEnum::INDEX_GPU_HNSW_SQ, IndexEnum::INDEX_HNSW_SQ, IndexEnum::INDEX_HNSW}) { + if (binset.Contains(key)) { + aliased.Append(IndexEnum::INDEX_GPU_HNSW, 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 { + // GPU_HNSW does not apply the bitset, so reject any search that carries + // one. Guard on data()!=nullptr rather than count(), because count() + // defaults to 0 when the caller omits num_filtered_out_bits and would + // let a real filter slip through. + if (!bitset.empty() && bitset.data() != nullptr) { + return expected::Err(Status::invalid_args, "GPU_HNSW does not support filtered search"); + } + + // 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) { + const auto* faiss_idx = GetFaissHnswIndex(); + if (!faiss_idx) { + return expected::Err(Status::empty_index, "index not loaded"); + } + try { + const auto& hnsw_cfg = static_cast(*cfg); + bool is_cosine = IsMetricType(hnsw_cfg.metric_type.value(), metric::COSINE); + bool use_ip = IsMetricType(hnsw_cfg.metric_type.value(), metric::IP) || is_cosine; + + 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); + gpu_ntotal_ = faiss_idx->ntotal; + gpu_dim_ = faiss_idx->d; + const_cast&>(indexes[0]).reset(); + gpu_index_ = local; + gpu_ready_.store(true, std::memory_order_release); + gpu_snapshot = local; + } catch (const std::exception& e) { + return expected::Err(Status::cuvs_inner_error, + std::string("failed to build GPU HNSW index: ") + e.what()); + } + } + } + + 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); + + 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; + 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 computes -dot(shifted_q, shifted_db) which equals the + // unscaled negative cosine numerator; divide by query norm to finalize. + if (IsMetricType(hnsw_cfg.metric_type.value(), metric::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; + } + } + } + + // Negate back to positive for IP and COSINE. + if (IsMetricType(hnsw_cfg.metric_type.value(), metric::IP) || + IsMetricType(hnsw_cfg.metric_type.value(), metric::COSINE)) { + for (int64_t i = 0; i < static_cast(nq * k); i++) { + h_dist[i] = -h_dist[i]; + } + } + + 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 (IsMetricType(hnsw_cfg.metric_type.value(), metric::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; + 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()); + } + + // Negate back to positive for IP and COSINE. + if (IsMetricType(hnsw_cfg.metric_type.value(), metric::IP) || + IsMetricType(hnsw_cfg.metric_type.value(), metric::COSINE)) { + for (int64_t i = 0; i < static_cast(nq * k); i++) { + h_dist[i] = -h_dist[i]; + } + } + + 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 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. + 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 applies the +128 shift and uses 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 applies the +128 shift and uses 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..4d3a7274f 100644 --- a/tests/ut/test_gpu_search.cc +++ b/tests/ut/test_gpu_search.cc @@ -9,7 +9,13 @@ // 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 "catch2/catch_approx.hpp" #include "catch2/catch_test_macros.hpp" @@ -23,6 +29,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 +453,959 @@ 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 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); + } +} + +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); + + // GPU recall should be reasonable at all ef values + REQUIRE(gpu_recall >= 0.5f); + } + 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); + } } #endif diff --git a/thirdparty/faiss/faiss/gpu/CMakeLists.txt b/thirdparty/faiss/faiss/gpu/CMakeLists.txt index 50740450a..59c7e7eab 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,15 @@ set(FAISS_GPU_HEADERS GpuIndexIVF.h GpuIndexIVFFlat.h GpuIndexIVFPQ.h + GpuIndexHNSW.h GpuIndexIVFScalarQuantizer.h GpuIndicesOptions.h GpuResources.h StandardGpuResources.h + 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..63eb18e91 --- /dev/null +++ b/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu @@ -0,0 +1,381 @@ +// @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 { + +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); + + 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; + + 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; +} + +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)); + + 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)); + + 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..05be2b9c4 --- /dev/null +++ b/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h @@ -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 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); + + /// Load with explicit metric specification. + 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. + 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. + /// Applies +128 shift to queries to match upload_int8_dataset encoding + /// (which stores as value - 128). dim must be divisible by 4. + /// 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; + + 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; +}; + +} // 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..f051bcd73 --- /dev/null +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh @@ -0,0 +1,438 @@ +// @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; + + 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..4422c20f0 --- /dev/null +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh @@ -0,0 +1,320 @@ +// @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 + +#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; + 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; + + auto launch_kernels = [&]( + const DataT* d_data, + const float* d_inv_norms) { + 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 = + (num_queries + warps_per_block - 1) / warps_per_block; + + hnsw_kernel::upper_layer_search_kernel + <<>>( + sc.d_queries, + d_data, + d_inv_norms, + d_layer_ptrs, + sc.d_entry_points, + idx.entry_point, + num_queries, + dim, + num_upper_layers, + idx.use_ip); + GPU_HNSW_CUDA_CHECK(cudaGetLastError()); + } else { + std::vector h_eps(num_queries, idx.entry_point); + GPU_HNSW_CUDA_CHECK(cudaMemcpyAsync( + sc.d_entry_points, + h_eps.data(), + num_queries * 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)); + } + + 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; + } + } + + { + int smem_overhead = sw * idx.max_degree0 * 8 + sw * 4 + 12; + int max_ef = (smem_max - smem_overhead) / 12; + // A search_width so large that even ef=1 does not fit in shared + // memory leaves no valid beam. Fail clearly rather than clamping ef + // to a negative/zero value (which would launch with a bad geometry + // or read out of bounds). + 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"); + } + if (ef > max_ef) { + ef = max_ef; + } + } + + size_t smem_size = hnsw_kernel::calc_layer0_smem_size( + ef, sw, idx.max_degree0); + + // Opt into >48 KiB dynamic shared memory when the device supports it; + // without this the kernel launch would fail for large ef. + if (smem_size > 49152) { + GPU_HNSW_CUDA_CHECK(cudaFuncSetAttribute( + hnsw_kernel::layer0_beam_search_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, + static_cast(smem_size))); + } + + int N_int = static_cast(idx.n_rows); + size_t bitmap_bytes = hnsw_kernel::calc_visited_bitmap_size( + num_queries, N_int); + + GPU_HNSW_CUDA_CHECK( + cudaMemsetAsync(sc.d_visited_bitmaps, 0, bitmap_bytes, stream)); + + hnsw_kernel::layer0_beam_search_kernel + <<>>( + sc.d_queries, + d_data, + d_inv_norms, + idx.d_layer0_graph, + sc.d_entry_points, + sc.d_visited_bitmaps, + sc.d_neighbors, + sc.d_distances, + num_queries, + N_int, + dim, + idx.max_degree0, + k, + ef, + sw, + max_iter, + idx.use_ip); + GPU_HNSW_CUDA_CHECK(cudaGetLastError()); + }; + + switch (idx.dataset_type) { + case GpuHnswDatasetType::INT8: + // Use the DP4A native int8 kernel when int8 queries are available + // and the metric is IP/COSINE. L2 falls through to the fp32 path + // because the DP4A kernel only computes dot products, not L2 distances. + if (sc.d_queries_i8 != nullptr && idx.use_ip) { + // Upper-layer greedy search still uses fp32 queries (d_queries). + 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 = + (num_queries + warps_per_block - 1) / warps_per_block; + hnsw_kernel::upper_layer_search_kernel + <<>>( + sc.d_queries, + static_cast(idx.d_dataset), + idx.d_inv_norms, + d_layer_ptrs, + sc.d_entry_points, + idx.entry_point, + num_queries, + dim, + num_upper_layers, + idx.use_ip); + GPU_HNSW_CUDA_CHECK(cudaGetLastError()); + } else { + std::vector h_eps(num_queries, idx.entry_point); + GPU_HNSW_CUDA_CHECK(cudaMemcpyAsync( + sc.d_entry_points, + h_eps.data(), + num_queries * sizeof(uint32_t), + cudaMemcpyHostToDevice, + stream)); + GPU_HNSW_CUDA_CHECK(cudaStreamSynchronize(stream)); + } + + int block_size = + params.thread_block_size > 0 ? params.thread_block_size : 128; + + int smem_max_i8 = 49152; + { + int device = 0; + int optin = 0; + if (cudaGetDevice(&device) == cudaSuccess && + cudaDeviceGetAttribute( + &optin, + cudaDevAttrMaxSharedMemoryPerBlockOptin, + device) == cudaSuccess && + optin > smem_max_i8) { + smem_max_i8 = optin; + } + } + + { + int smem_overhead = sw * idx.max_degree0 * 8 + sw * 4 + 12; + int max_ef = (smem_max_i8 - smem_overhead) / 12; + if (max_ef < 1) { + throw std::runtime_error( + std::string("gpu_hnsw_int8: search_width=") + + std::to_string(sw) + + " too large for device shared memory (" + + std::to_string(smem_max_i8) + + " bytes); reduce search_width"); + } + if (ef > max_ef) { + ef = max_ef; + } + } + + size_t smem_size_i8 = hnsw_kernel::calc_layer0_smem_size( + ef, sw, idx.max_degree0); + + if (smem_size_i8 > 49152) { + GPU_HNSW_CUDA_CHECK(cudaFuncSetAttribute( + hnsw_kernel::layer0_beam_search_kernel_int8, + cudaFuncAttributeMaxDynamicSharedMemorySize, + static_cast(smem_size_i8))); + } + + int N_int = static_cast(idx.n_rows); + size_t bitmap_bytes = hnsw_kernel::calc_visited_bitmap_size( + num_queries, N_int); + GPU_HNSW_CUDA_CHECK(cudaMemsetAsync( + sc.d_visited_bitmaps, 0, bitmap_bytes, stream)); + + hnsw_kernel::layer0_beam_search_kernel_int8 + <<>>( + reinterpret_cast(sc.d_queries_i8), + static_cast(idx.d_dataset), + idx.d_inv_norms, + idx.d_layer0_graph, + sc.d_entry_points, + sc.d_visited_bitmaps, + sc.d_neighbors, + sc.d_distances, + num_queries, + N_int, + dim, + idx.max_degree0, + k, + ef, + sw, + max_iter, + idx.use_ip); + GPU_HNSW_CUDA_CHECK(cudaGetLastError()); + break; + } + // Fall through to generic int8 fp32-query path if no i8 queries. + launch_kernels( + static_cast(idx.d_dataset), idx.d_inv_norms); + break; + case GpuHnswDatasetType::FP16: + launch_kernels( + static_cast(idx.d_dataset), idx.d_inv_norms); + break; + case GpuHnswDatasetType::BF16: + launch_kernels( + static_cast(idx.d_dataset), + idx.d_inv_norms); + break; + case GpuHnswDatasetType::FP32: + default: + launch_kernels( + static_cast(idx.d_dataset), idx.d_inv_norms); + 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..c0a3ae471 --- /dev/null +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh @@ -0,0 +1,772 @@ +// @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; +} + +// DP4A inner-product distance for int8 queries vs int8 dataset. +// Both query and dataset must be in the same encoding: signed int8 user values. +// Dataset is stored as-is (upload_int8_dataset reverses FAISS's +128 bias). +// Queries are passed without any shift (user signed int8 values directly). +// Computes -dot(query, vec) as a float (max-IP / negative dot convention). +// dim4 = dim / 4 (caller must ensure dim % 4 == 0). +__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); // negate: max-IP convention +} + +// ============================================================================ +// 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 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; +} + +template +__global__ void layer0_beam_search_kernel( + const float* __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) { + int query_idx = blockIdx.x; + if (query_idx >= num_queries) + return; + + const float* query = d_queries + static_cast(query_idx) * dim; + + extern __shared__ char smem[]; + + int max_staging = search_width * max_degree0; + + 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); + + 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; + } + __syncthreads(); + + // --- Seed with entry point --- + uint32_t ep = d_entry_points[query_idx]; + if (threadIdx.x == 0) { + float ep_dist; + if (use_inner_product) { + ep_dist = thread_ip_distance( + query, d_dataset + static_cast(ep) * dim, dim); + if (d_inv_norms) + ep_dist *= __ldg(&d_inv_norms[ep]); + } else { + ep_dist = thread_l2_distance( + query, d_dataset + static_cast(ep) * dim, dim); + } + result_ids[0] = ep; + result_dists[0] = ep_dist; + is_expanded[0] = 0; + meta[0] = 1; + bitmap_visit(visited_bmap, ep); + } + __syncthreads(); + + // --- Seed with entry point's neighbors --- + if (threadIdx.x == 0) + meta[1] = 0; + __syncthreads(); + + for (int j = threadIdx.x; j < max_degree0; j += blockDim.x) { + uint32_t nbr = + d_layer0_graph[static_cast(ep) * max_degree0 + j]; + if (nbr == UINT32_MAX || nbr >= static_cast(N)) + continue; + if (!bitmap_visit(visited_bmap, nbr)) + continue; + + const DataT* nbr_vec = d_dataset + static_cast(nbr) * dim; + float dist; + 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); + } + + int slot = atomicAdd(&meta[1], 1); + if (slot < max_staging) { + staging_ids[slot] = nbr; + staging_dists[slot] = dist; + } + } + __syncthreads(); + + if (threadIdx.x == 0) { + int staging_count = min(meta[1], max_staging); + int rc = meta[0]; + + for (int s = 0; s < staging_count; s++) { + uint32_t sid = staging_ids[s]; + float sdist = staging_dists[s]; + if (rc >= ef && sdist >= result_dists[rc - 1]) + continue; + + int lo = 0, hi = rc; + while (lo < hi) { + int mid = (lo + hi) / 2; + if (result_dists[mid] < sdist) + lo = mid + 1; + else + hi = mid; + } + int insert_end = rc < ef ? rc : ef - 1; + for (int i = insert_end; i > lo; i--) { + result_ids[i] = result_ids[i - 1]; + result_dists[i] = result_dists[i - 1]; + is_expanded[i] = is_expanded[i - 1]; + } + result_ids[lo] = sid; + result_dists[lo] = sdist; + is_expanded[lo] = 0; + if (rc < ef) + rc++; + } + + for (int i = 0; i < rc; i++) { + if (result_ids[i] == ep) { + is_expanded[i] = 1; + break; + } + } + meta[0] = rc; + } + __syncthreads(); + + // --- Unified main loop --- + for (int iter = 0; iter < max_iterations; iter++) { + if (threadIdx.x == 0) { + int num_parents = 0; + int rc = meta[0]; + + 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 = threadIdx.x; wi < total_work; wi += blockDim.x) { + int parent_idx = wi / max_degree0; + int nbr_slot = wi % max_degree0; + + uint32_t parent = parent_ids[parent_idx]; + uint32_t nbr = + d_layer0_graph + [static_cast(parent) * max_degree0 + + nbr_slot]; + if (nbr == UINT32_MAX || nbr >= static_cast(N)) + continue; + if (!bitmap_visit(visited_bmap, nbr)) + continue; + + const DataT* nbr_vec = d_dataset + static_cast(nbr) * dim; + float dist; + 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); + } + + int slot = atomicAdd(&meta[1], 1); + if (slot < max_staging) { + staging_ids[slot] = nbr; + staging_dists[slot] = dist; + } + } + __syncthreads(); + + if (threadIdx.x == 0) { + int staging_count = min(meta[1], max_staging); + int rc = meta[0]; + + for (int s = 0; s < staging_count; s++) { + uint32_t sid = staging_ids[s]; + float sdist = staging_dists[s]; + if (rc >= ef && sdist >= result_dists[rc - 1]) + continue; + + int lo = 0, hi = rc; + while (lo < hi) { + int mid = (lo + hi) / 2; + if (result_dists[mid] < sdist) + lo = mid + 1; + else + hi = mid; + } + int insert_end = rc < ef ? rc : ef - 1; + for (int i = insert_end; i > lo; i--) { + result_ids[i] = result_ids[i - 1]; + result_dists[i] = result_dists[i - 1]; + is_expanded[i] = is_expanded[i - 1]; + } + result_ids[lo] = sid; + result_dists[lo] = sdist; + is_expanded[lo] = 0; + if (rc < ef) + rc++; + } + + meta[0] = rc; + } + __syncthreads(); + + // Stagnation detected when num_parents==0 (all expanded) → break above + } + + // --- Copy top-k results to global memory --- + int rc = meta[0]; + for (int i = threadIdx.x; i < k; i += blockDim.x) { + if (i < rc) { + d_neighbors[static_cast(query_idx) * k + i] = + static_cast(result_ids[i]); + d_distances[static_cast(query_idx) * k + i] = + result_dists[i]; + } else { + d_neighbors[static_cast(query_idx) * k + i] = UINT64_MAX; + d_distances[static_cast(query_idx) * k + i] = FLT_MAX; + } + } +} + +// ============================================================================ +// Phase 2 (int8/DP4A variant): Layer-0 parallel beam search kernel +// Uses __dp4a for 4x int8 MADs per cycle instead of fp32 FMA. +// d_queries_packed: int8 queries reinterpreted as int32 (signed user values, +// no shift — matches dataset encoding after upload_int8_dataset -128 reversal); +// dim must be divisible by 4. +// d_dataset: raw int8 dataset on device (already shifted -128 at upload time, +// now consistent with query shift so DP4A computes the correct dot product). +// ============================================================================ +__global__ void layer0_beam_search_kernel_int8( + const int32_t* __restrict__ d_queries_packed, + const int8_t* __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) { + int query_idx = blockIdx.x; + if (query_idx >= num_queries) + return; + + int dim4 = dim / 4; + const int32_t* query = d_queries_packed + static_cast(query_idx) * dim4; + + extern __shared__ char smem[]; + + int max_staging = search_width * max_degree0; + + 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); + + 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; + } + __syncthreads(); + + // --- Seed with entry point --- + uint32_t ep = d_entry_points[query_idx]; + if (threadIdx.x == 0) { + const int32_t* ep_packed = + reinterpret_cast(d_dataset + static_cast(ep) * dim); + // DP4A kernel is only dispatched for IP/COSINE (use_ip=true); L2 + // falls back to the fp32 kernel before this point. + float ep_dist = thread_ip_distance_dp4a(query, ep_packed, dim4); + if (d_inv_norms) + ep_dist *= __ldg(&d_inv_norms[ep]); + result_ids[0] = ep; + result_dists[0] = ep_dist; + is_expanded[0] = 0; + meta[0] = 1; + bitmap_visit(visited_bmap, ep); + } + __syncthreads(); + + // --- Seed with entry point's neighbors --- + if (threadIdx.x == 0) + meta[1] = 0; + __syncthreads(); + + for (int j = threadIdx.x; j < max_degree0; j += blockDim.x) { + uint32_t nbr = + d_layer0_graph[static_cast(ep) * max_degree0 + j]; + if (nbr == UINT32_MAX || nbr >= static_cast(N)) + continue; + if (!bitmap_visit(visited_bmap, nbr)) + continue; + + const int32_t* nbr_packed = + reinterpret_cast(d_dataset + static_cast(nbr) * dim); + float dist = thread_ip_distance_dp4a(query, nbr_packed, dim4); + if (d_inv_norms) + dist *= d_inv_norms[nbr]; + + int slot = atomicAdd(&meta[1], 1); + if (slot < max_staging) { + staging_ids[slot] = nbr; + staging_dists[slot] = dist; + } + } + __syncthreads(); + + if (threadIdx.x == 0) { + int staging_count = min(meta[1], max_staging); + int rc = meta[0]; + + for (int s = 0; s < staging_count; s++) { + uint32_t sid = staging_ids[s]; + float sdist = staging_dists[s]; + if (rc >= ef && sdist >= result_dists[rc - 1]) + continue; + + int lo = 0, hi = rc; + while (lo < hi) { + int mid = (lo + hi) / 2; + if (result_dists[mid] < sdist) + lo = mid + 1; + else + hi = mid; + } + int insert_end = rc < ef ? rc : ef - 1; + for (int i = insert_end; i > lo; i--) { + result_ids[i] = result_ids[i - 1]; + result_dists[i] = result_dists[i - 1]; + is_expanded[i] = is_expanded[i - 1]; + } + result_ids[lo] = sid; + result_dists[lo] = sdist; + is_expanded[lo] = 0; + if (rc < ef) + rc++; + } + + for (int i = 0; i < rc; i++) { + if (result_ids[i] == ep) { + is_expanded[i] = 1; + break; + } + } + meta[0] = rc; + } + __syncthreads(); + + // --- Unified main loop --- + for (int iter = 0; iter < max_iterations; iter++) { + if (threadIdx.x == 0) { + int num_parents = 0; + int rc = meta[0]; + + 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 = threadIdx.x; wi < total_work; wi += blockDim.x) { + int parent_idx = wi / max_degree0; + int nbr_slot = wi % max_degree0; + + uint32_t parent = parent_ids[parent_idx]; + uint32_t nbr = + d_layer0_graph + [static_cast(parent) * max_degree0 + + nbr_slot]; + if (nbr == UINT32_MAX || nbr >= static_cast(N)) + continue; + if (!bitmap_visit(visited_bmap, nbr)) + continue; + + const int32_t* nbr_packed = + reinterpret_cast(d_dataset + static_cast(nbr) * dim); + float dist = thread_ip_distance_dp4a(query, nbr_packed, dim4); + if (d_inv_norms) + dist *= d_inv_norms[nbr]; + + int slot = atomicAdd(&meta[1], 1); + if (slot < max_staging) { + staging_ids[slot] = nbr; + staging_dists[slot] = dist; + } + } + __syncthreads(); + + if (threadIdx.x == 0) { + int staging_count = min(meta[1], max_staging); + int rc = meta[0]; + + for (int s = 0; s < staging_count; s++) { + uint32_t sid = staging_ids[s]; + float sdist = staging_dists[s]; + if (rc >= ef && sdist >= result_dists[rc - 1]) + continue; + + int lo = 0, hi = rc; + while (lo < hi) { + int mid = (lo + hi) / 2; + if (result_dists[mid] < sdist) + lo = mid + 1; + else + hi = mid; + } + int insert_end = rc < ef ? rc : ef - 1; + for (int i = insert_end; i > lo; i--) { + result_ids[i] = result_ids[i - 1]; + result_dists[i] = result_dists[i - 1]; + is_expanded[i] = is_expanded[i - 1]; + } + result_ids[lo] = sid; + result_dists[lo] = sdist; + is_expanded[lo] = 0; + if (rc < ef) + rc++; + } + + meta[0] = rc; + } + __syncthreads(); + } + + // --- Copy top-k results to global memory --- + int rc = meta[0]; + for (int i = threadIdx.x; i < k; i += blockDim.x) { + if (i < rc) { + d_neighbors[static_cast(query_idx) * k + i] = + static_cast(result_ids[i]); + d_distances[static_cast(query_idx) * k + i] = + result_dists[i]; + } else { + d_neighbors[static_cast(query_idx) * k + i] = UINT64_MAX; + d_distances[static_cast(query_idx) * k + i] = FLT_MAX; + } + } +} + +inline size_t calc_layer0_smem_size(int ef, int search_width, int max_degree0) { + int max_staging = search_width * max_degree0; + + size_t size = 0; + size += ef * sizeof(uint32_t); + size += ef * sizeof(float); + size += ef * sizeof(uint32_t); + size += max_staging * sizeof(uint32_t); + size += max_staging * sizeof(float); + size += search_width * sizeof(uint32_t); + size += 3 * sizeof(int); + 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..67d00365a --- /dev/null +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu @@ -0,0 +1,191 @@ +// @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; + } + int bitmap_words = (N + 31) / 32; + size_t need_bm = + static_cast(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; + } + } +} + +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); +} + +GpuHnswScratchSlot::~GpuHnswScratchSlot() { + 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..1f15a3cca --- /dev/null +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h @@ -0,0 +1,210 @@ +// @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 + +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, +}; + +struct GpuHnswSearchParams { + int ef = 200; + int search_width = 4; + int max_iterations = 0; + int thread_block_size = 0; +}; + +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 DP4A path (shifted +128) + + 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; + + // 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); + + ~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