lib/segment/src/index/hnsw_index/gpu/gpu_insert_context.rs (1)

98-106: ⚠️ Potential issue

Potential buffer-overflow when more points than points_capacity are passed to init

insert_atomics_buffer and all visited-flags buffers are sized with:

let points_count = gpu_vector_storage.num_vectors();          // ≈ points_capacity

Later, init(remap) blindly uploads remap.len() elements.
If the caller adds extra virtual point IDs that are not yet present in the vector storage (e.g. during batched import) remap.len() can exceed points_capacity, causing out-of-bounds writes in:

GpuVisitedFlags::init → context.copy_gpu_buffer(...)

Recommendation:

 pub fn init(&mut self, remap: &[PointOffsetType]) -> OperationResult<()> {
-    self.gpu_visited_flags.init(remap)?;
+    if remap.len() > self.gpu_visited_flags.capacity() {
+        return Err(OperationError::service_error("Visited-flags capacity is too small for provided remap"));
+    }
+    self.gpu_visited_flags.init(remap)?;

…and add an assert/early-return inside GpuVisitedFlags::init as well.

Also applies to: 134-139

lib/segment/src/index/hnsw_index/gpu/gpu_level_builder.rs (1)

150-161: Avoid panicking in test helpers – propagate the error instead

gpu_search_context.init(...) is unwrapped here. If the allocation of the remap staging buffer inside GpuVisitedFlags::init fails (OOM on CI runners with small GPUs) the whole test binary will abort, making diagnostics harder.
Prefer bubbling the OperationResult up with ?, or at least add a helpful expect() message.

-gpu_search_context.init(batched_points.remap()).unwrap();
+gpu_search_context
+    .init(batched_points.remap())
+    .expect("initialising GPU insert context failed");

lib/segment/src/index/hnsw_index/gpu/gpu_insert_context.rs (1)

264-270: init re-allocates a fresh gpu::Context for every call inside GpuVisitedFlags::init

GpuVisitedFlags::init creates a brand-new gpu::Context, which is expensive (command-pool, fence, etc.).
Since GpuInsertContext already owns an operational Context, consider passing it down to reuse existing resources:

pub fn init(&mut self,
            gpu_ctx: &mut gpu::Context,
            points_remap: &[PointOffsetType]) -> OperationResult<()>

This will cut redundant allocations when init is invoked many times across large imports.

lib/segment/src/index/hnsw_index/gpu/gpu_graph_builder.rs (3)

15-15: Prefer const over static for immutable scalar values

GPU_MAX_VISITED_FLAGS_FACTOR is used as a compile-time constant. Declaring it as

pub const GPU_MAX_VISITED_FLAGS_FACTOR: usize = 32;

avoids putting an unnecessary symbol into the data section and prevents accidental mutability through unsafe { … }.

---

19-21: Annotate lifetime 'b more explicitly (optional)

The extra lifetime on GpuInsertContext<'b> is only needed if the context borrows from something with the same lifetime.
If the context now owns all internal buffers (post-refactor) you can drop the explicit 'b to simplify the signature:

pub fn build_hnsw_on_gpu<'a>(
    gpu_insert_context: &mut GpuInsertContext,
    …
)

This is optional but makes call-sites cleaner and eliminates an otherwise unused lifetime parameter.

---

149-171: Avoid cloning the full ids vector on every repeat

ids.clone() duplicates a potentially large buffer for each iteration.
Instead, pass a slice to build_hnsw_on_gpu and let the callee create its own copy only if needed:

-                build_hnsw_on_gpu(
+                build_hnsw_on_gpu(
                     &mut gpu_search_context,
                     &test.graph_layers_builder,
                     groups_count,
                     1,
                     cpu_linked_points_count,
-                    ids.clone(),
+                    &ids,

and change the parameter type in build_hnsw_on_gpu to &[PointOffsetType].
This cuts memory reallocations in tight benchmarks.

lib/segment/src/index/hnsw_index/hnsw.rs (2)

700-711: Potential duplication of expensive GPU context creation

GpuInsertContext::new is called inside build_main_graph_on_gpu, which is itself called from the builder.
Because the context can now be reused, consider propagating an existing GpuInsertContext instead of creating a fresh one here to avoid reallocating buffers (~hundreds of MB on large indices).

If life-time constraints prevent reuse, document why a fresh context is necessary.

---

764-775: Graceful fallback covers functional errors but not panics

build_hnsw_on_gpu errors are converted into a warning and the CPU path is taken, which is great.
However, if the GPU code panics (e.g., via unwrap() in lower levels) the whole process aborts.

Recommend wrapping the GPU build in std::panic::catch_unwind to keep the same graceful-degradation guarantee in release builds.

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

lib/segment/src/index/hnsw_index/hnsw.rs (1)

701-714: Avoid re-allocating a second GpuInsertContext for the main graph

A brand-new GpuInsertContext is constructed for the main graph here, while another one is created (and reused) for payload blocks later. That’s an extra set of large GPU buffers (~gigabytes on big indexes) that immediately gets dropped after the main build.

Reusing the same context throughout the whole build() run keeps allocations O(1). One possible approach:

// (1) create context once, before main-graph build)
-let gpu_vectors = Self::create_gpu_vectors(...)?;
-let mut gpu_insert_context = gpu_vectors.as_ref().map(|v| {
-    GpuInsertContext::new(
-        v,
-        get_gpu_groups_count(),
-        config.m,                     // main-graph params
-        config.m0,
-        config.ef_construct,
-        false,
-        1..=GPU_MAX_VISITED_FLAGS_FACTOR,
-    )
-}).transpose()?;
+let mut gpu_insert_context = gpu_vectors.as_ref().map(|v| {
+    GpuInsertContext::new(
+        v,
+        get_gpu_groups_count(),
+        config.m.max(payload_m),      // max of main/payload
+        config.m0.max(config.payload_m0.unwrap_or(config.m0)),
+        config.ef_construct,
+        false,
+        1..=GPU_MAX_VISITED_FLAGS_FACTOR,
+    )
+}).transpose()?;

// (2) pass `&mut gpu_insert_context` to BOTH main and payload builders
 Self::build_main_graph_on_gpu(..., gpu_insert_context.as_mut(), ...)?;
 ...
 Self::build_filtered_graph(..., &mut gpu_insert_context, ...)?;

This reduces memory pressure and honours the PR goal of not reallocating GPU resources between phases.

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro