Multi partition cagra search - #2035
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
📝 WalkthroughWalkthroughThis pull request introduces multi-segment CAGRA search functionality, GPU-based selection, and RMM memory pool management. It adds C and C++ APIs for multi-segment approximate nearest neighbor search, async memory allocation, and select-K operations, alongside corresponding Java bindings and internal implementations supporting concurrent per-segment GPU processing and stream-based result aggregation. Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~75 minutes 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 17
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraSearchParams.java (1)
293-322:⚠️ Potential issue | 🟡 MinorThe
toString()method does not include the new persistent fields.The three new fields (
persistent,persistentLifetime,persistentDeviceUsage) are not included in thetoString()output. While not critical, this creates an inconsistency where debugging output won't show the full configuration.🔧 Proposed fix to include persistent fields in toString()
+ ", randXORMask=" + randXORMask + + ", persistent=" + + persistent + + ", persistentLifetime=" + + persistentLifetime + + ", persistentDeviceUsage=" + + persistentDeviceUsage + "]";🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraSearchParams.java` around lines 293 - 322, The toString() in class CagraSearchParams currently omits the new fields persistent, persistentLifetime, and persistentDeviceUsage; update the CagraSearchParams.toString() method to append these three fields (with their names and values) into the returned string (same formatting style as the other fields) so debugging output shows the full configuration.
🧹 Nitpick comments (3)
java/cuvs-java/src/test/java/com/nvidia/cuvs/CheckedCuVSResources.java (1)
59-62: Add destroyed-state guard in the new delegate method.
setWorkspacePool(...)should callcheckNotDestroyed()before delegating, consistent with this wrapper’s defensive contract.Proposed change
`@Override` public void setWorkspacePool(long sizeBytes) { + checkNotDestroyed(); inner.setWorkspacePool(sizeBytes); }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@java/cuvs-java/src/test/java/com/nvidia/cuvs/CheckedCuVSResources.java` around lines 59 - 62, The new delegate method setWorkspacePool(long sizeBytes) is missing the wrapper's destroyed-state guard; modify CheckedCuVSResources.setWorkspacePool to call checkNotDestroyed() at the start (before delegating to inner.setWorkspacePool(sizeBytes)) so it matches the class's defensive contract and uses the same destroyed-state check as other methods.java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraSearchParams.java (1)
527-537: Consider adding validation forpersistentDeviceUsagebounds.The Javadoc states the value "must be greater than 0.0 and not greater than 1.0", but the setter doesn't validate this constraint. Invalid values would only fail at the native layer.
🛡️ Optional: Add validation in the builder
public Builder withPersistentDeviceUsage(float persistentDeviceUsage) { + if (persistentDeviceUsage <= 0.0f || persistentDeviceUsage > 1.0f) { + throw new IllegalArgumentException( + "persistentDeviceUsage must be > 0.0 and <= 1.0, got: " + persistentDeviceUsage); + } this.persistentDeviceUsage = persistentDeviceUsage; return this; }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraSearchParams.java` around lines 527 - 537, The builder setter withPersistentDeviceUsage currently assigns persistentDeviceUsage without enforcing the Javadoc constraint; update Builder.withPersistentDeviceUsage to validate that persistentDeviceUsage > 0.0f and <= 1.0f and throw an IllegalArgumentException (including the invalid value in the message) when the check fails so invalid values are rejected early before reaching native code.cpp/src/neighbors/detail/cagra/cagra_search.cuh (1)
388-417: Consider hoisting query_norms allocation outside the loop for CosineExpanded metric.When multiple segments use
CosineExpandedmetric,query_normsis allocated and computed redundantly for each segment. Since queries are the same across segments (repeated query vector), the norms could be computed once.This is a minor optimization opportunity. The current implementation is correct; the redundant computation is bounded by the number of segments.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cpp/src/neighbors/detail/cagra/cagra_search.cuh` around lines 388 - 417, Hoist the allocation and computation of query_norms out of the segment loop and reuse it for every segment whose indices[i]->metric() == cuvs::distance::DistanceType::CosineExpanded: allocate query_norms once (using raft::make_device_vector) before the for-loop, run raft::linalg::reduce on the shared queries data to compute the norms, then inside the loop call raft::linalg::matrix_vector_op (as currently done) using the precomputed query_norms for each CosineExpanded segment; after the loop free or let query_norms go out of scope. Ensure you still call cuvs::neighbors::ivf::detail::postprocess_distances for non-CosineExpanded branches and keep all existing ops (raft::compose_op, raft::sq_op, raft::div_const_op, raft::cast_op, raft::add_const_op, raft::div_checkzero_op) unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@c/include/cuvs/core/c_api.h`:
- Around line 229-240: The current implementation of
cuvsRMMAsyncMemoryResourceEnable stores cuda_async_memory_resource in a
thread_local async_mr and passes it to rmm::mr::set_current_device_resource(),
creating a lifetime mismatch vs the API doc that says the change is global;
change the implementation so the async memory resource is process-scoped (not
thread_local) by replacing thread_local async_mr with a static/process-global
instance (e.g., a static unique_ptr or static object) so it outlives the thread
and remains valid for rmm::mr::set_current_device_resource(), and ensure
cuvsRMMAsyncMemoryResourceEnable and any cleanup use that same process-global
symbol (cuda_async_memory_resource / async_mr) when setting or resetting the
current device resource; alternatively, if thread-local semantics are intended,
update the cuvsRMMAsyncMemoryResourceEnable documentation to state it only
affects the calling thread and keep async_mr thread_local.
In `@c/src/neighbors/cagra.cpp`:
- Around line 710-733: The loop currently casts every indices[i]->addr and
builds device views without validating each segment's types/devices; update the
loop that fills idx_vec, q_vec, n_vec, d_vec to perform the same per-segment
checks used in cuvsCagraSearch: assert indices[i] != nullptr and addr != 0
(already present), then verify indices[i]->dtype.code == kDLFloat &&
indices[i]->dtype.bits == 32 for every i, and validate that queries[i],
neighbors[i], and distances[i] are device-backed DLPack tensors with the
expected element types (float for queries/distances, uint32 for neighbors)
before calling reinterpret_cast<const IndexT*>(indices[i]->addr) and
cuvs::core::from_dlpack to populate q_vec[i], n_vec[i], d_vec[i]; replace the
blind casts with RAFT_EXPECTS that include i in the error messages so a bad
segment fails fast and clearly.
- Around line 726-736: The loop that builds idx_vec/q_vec/n_vec/d_vec must also
validate that all segment indices use the same distance metric as the first
segment to prevent mixing incompatible metrics; inside the for-loop in cagra.cpp
(after the existing RAFT_EXPECTS that checks indices[i] non-null) compare
indices[i]->metric (or the actual metric field name on your IndexT struct/class)
to indices[0]->metric and fail fast with RAFT_EXPECTS (or equivalent) and a
clear message like "Mixed distance metrics across segments: expected %s but got
%s at segment %u"; keep this check before pushing idx_vec[i] and before calling
cuvs::neighbors::cagra::search_multi_segment so the function only runs when all
segments share the same metric.
In `@c/src/selection/select_k.cpp`:
- Around line 14-40: cuvsSelectK dereferences shape[1] and casts buffers without
validating the DLPack tensors; add explicit validation at the top of cuvsSelectK
for in_val, out_val, out_idx (non-null), then check each
DLManagedTensor->dl_tensor for expected ndim (==2), shapes (rows match expected
1 or compatible), dtype (in_val/out_val float32, out_idx int64), device type
(CUDA) and device id, byte_offset == 0, and contiguous row-major
strides/compatibility before creating device views with
raft::make_device_matrix_view; if any check fails return an appropriate
cuvsError_t (or throw inside translate_exceptions) instead of proceeding to
casts and calling cuvs::selection::select_k so malformed callers cannot crash or
corrupt memory.
In `@cpp/include/cuvs/neighbors/cagra.hpp`:
- Around line 1752-1806: Add Doxygen documentation for each overloaded
search_multi_segment declaration so they appear in generated API docs; either
add brief doxygen blocks above each overload or use `@copydoc` to reference the
primary search_multi_segment doc block (e.g., use `@copydoc`
search_multi_segment(raft::resources const&,
cuvs::neighbors::cagra::search_params const&, const std::vector<const
cuvs::neighbors::cagra::index<float, uint32_t>*>&, const
std::vector<raft::device_matrix_view<const float, int64_t, raft::row_major>>&,
const std::vector<raft::device_matrix_view<int64_t, int64_t, raft::row_major>>&,
const std::vector<raft::device_matrix_view<float, int64_t, raft::row_major>>&))
for the overloads with half/int8_t/uint8_t and uint32_t/int64_t neighbor types
so each signature (the overloads of search_multi_segment) is documented.
- Around line 1744-1806: The header declares the template parameter order as <T,
OutputIdxT, IdxT> but the instantiation macro (used with <T, IdxT, OutputIdxT>
e.g. (data_t, uint32_t, int64_t)) expects <T, IdxT, OutputIdxT>, causing the
IdxT/OutputIdxT swap that trips the static_assert in cagra_search.cuh:276 and
breaks link-time overloads for int64_t; fix by changing the template parameter
order in the search_multi_segment declarations to <T, IdxT, OutputIdxT> (or
alternatively update the instantiation macro to match the declared order) so
types map correctly, and add missing Doxygen for overloads 2–8 by inserting a
`@copydoc` search_multi_segment (or equivalent documentation block) above each of
those overloaded search_multi_segment declarations to satisfy the public API
docs requirement.
In `@cpp/src/neighbors/cagra.cuh`:
- Around line 409-420: The wrapper search_multi_segment currently forwards
indices, queries, neighbors, and distances without validation; add the same
upfront shape checks used by search() before calling
cagra::detail::search_multi_segment: verify indices.size() == queries.size() ==
neighbors.size() == distances.size(), then for each segment i ensure
queries[i].n_rows == neighbors[i].n_rows, neighbors[i].n_cols ==
distances[i].n_cols (k matches), queries[i].n_cols equals the index dimension
for indices[i] (or indices[i]->dim() / appropriate accessor), and that k > 0 and
dims are consistent; if any check fails, return/throw a clear error (or use
RAFT/CUASSERT used elsewhere) rather than forwarding to the detail
implementation.
In `@java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSResources.java`:
- Around line 60-76: The Java API should validate workspace pool sizes before
calling native code: update the Javadoc for setWorkspacePool to state the valid
range is > 0, and in the implementation of setWorkspacePool (the method that
currently invokes the native cuvsResourcesSetWorkspacePool) add a check that
rejects non-positive values (<= 0) by throwing an appropriate Java exception
(e.g., IllegalArgumentException) with a clear message; only call
cuvsResourcesSetWorkspacePool when the value is positive to avoid
signed->unsigned wraparound in native size_t.
In `@java/cuvs-java/src/main/java/com/nvidia/cuvs/MultiSegmentSearchResults.java`:
- Around line 24-27: MultiSegmentSearchResults currently stores native uint32
ordinals in an int[] (field ordinals) which corrupts values > Integer.MAX_VALUE;
change ordinals from int[] to long[] (and any constructor/getter signatures) and
ensure code that decodes native uint32_t values writes unsigned values into the
long (e.g., value & 0xFFFFFFFFL) so ordinals remain non-negative; update any
consumers (notably MultiSegmentCagraSearch) to compare against a long sentinel
(e.g., -1L) or otherwise handle long ordinals instead of treating negative ints
as sentinels.
In `@java/cuvs-java/src/main/java/com/nvidia/cuvs/SynchronizedCuVSResources.java`:
- Around line 43-46: The setWorkspacePool method in SynchronizedCuVSResources is
not using the shared lock and must be serialized like access(); modify
setWorkspacePool to acquire the same lock used by access() (e.g., wrap the call
to inner.setWorkspacePool(sizeBytes) in the synchronized block or lock guard
used by access()) so that mutations to workspace pool are protected by the same
synchronization as access().
In `@java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java`:
- Around line 391-396: The offset calculation uses segmentIdx * topK
(neighborByteOffset/distanceByteOffset) but the code builds tensors with shape
{numQueries, topK} (numQueries from queryVectors.size()), so when numQueries > 1
the buffer offsets are wrong; either enforce single-query by adding a guard (if
(numQueries != 1) throw new IllegalArgumentException(...)) near where numQueries
is computed (queryVectors) or update the offsets to multiply by numQueries
(neighborByteOffset = segmentIdx * numQueries * topK * C_INT_BYTE_SIZE and
distanceByteOffset = segmentIdx * numQueries * topK * Float.BYTES) before
creating neighborSlice/distanceSlice (globalNeighborsDP/globalDistancesDP).
In `@java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CudaStreamPool.java`:
- Around line 115-119: Make nextSlot atomic and wrap the returned slot index by
changing slotCounter to an AtomicInteger and using an atomic get-and-add plus
modulo; specifically, replace the non-atomic increment in nextSlot with
something like int start = slotCounter.getAndAdd(count); then return
Math.floorMod(start, poolSize) (or equivalent using your pool size field) so the
operation is thread-safe and the returned slot is bounded by the pool size.
- Around line 125-131: In CudaStreamPool.close(), wrap the cuvsResourcesDestroy
call with the same error validation used elsewhere by replacing the raw
cuvsResourcesDestroy(resources[i]) invocation with a call to
checkCuVSError(cuvsResourcesDestroy(resources[i]), "cuvsResourcesDestroy") so
cleanup failures are logged/handled; keep the surrounding loop and existing
calls to checkCudaError(cudaEventDestroy(events[i]), "cudaEventDestroy") and
checkCudaError(cudaStreamDestroy(streams[i]), "cudaStreamDestroy") unchanged and
reference the close(), cuvsResourcesDestroy, checkCuVSError, events, resources,
streams, and size symbols to locate the change.
- Around line 59-85: The constructor CudaStreamPool currently leaks native
handles if a create call fails mid-loop; modify the CudaStreamPool(int size)
constructor to perform rollback cleanup on failure by tracking the current index
and, if any checkCudaError/checkCuVSError throws, iterating over
already-initialized entries in resources[], streams[], and events[] to call the
corresponding destroy functions (cudaStreamDestroy for streams[],
cudaEventDestroy for events[], and the cuvs resources destroy routine for
resources[]) before rethrowing the exception; implement this by wrapping the
allocation loop in try/catch (or try/finally with a success flag) and invoking
the same cleanup logic as close() for indices < currentIndex so no native
handles are leaked if construction aborts.
In
`@java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSResourcesImpl.java`:
- Around line 31-33: The CudaStreamPool is currently allocated in the field
initializer (streamPool) which can leak native resources if
CuVSResourcesImpl(Path) throws before construction completes; move creation of
the CudaStreamPool from the field initializer into the CuVSResourcesImpl(Path)
constructor (use Integer.getInteger(CudaStreamPool.SIZE_PROPERTY,
CudaStreamPool.DEFAULT_SIZE) to determine size), assign it to the streamPool
field there, and in the constructor failure path ensure you call
streamPool.close() (or otherwise tear it down) before rethrowing so native
resources are not leaked; keep the streamPool field declaration but initialize
it only in the constructor and ensure close() is reachable on exceptions.
In `@java/cuvs-java/src/main/java22/com/nvidia/cuvs/MultiSegmentCagraSearch.java`:
- Around line 119-145: The code must enforce the single-query contract or size
outputs from the actual query row count: in MultiSegmentCagraSearch, after
obtaining var queryVectors = (CuVSMatrixInternal)
queries.get(i).getQueryVectors(), read its row count (e.g.,
queryVectors.getRowCount()/numRows()/size(0) — use the actual accessor on
CuVSMatrixInternal) into int nq; if nq > 1 throw an IllegalArgumentException
rejecting multi-row queries, or alternatively set segShape = new long[] {nq, k}
and compute neighbor/distance byte offsets and tensor sizes using nq*k (update
nByteOffset/dByteOffset and prepareTensor calls for
neighborsArray/distancesArray accordingly) so prepareTensor and the
globalNeighborsDP/globalDistancesDP slices match the query row count.
- Around line 96-111: The code in MultiSegmentCagraSearch currently uses
CagraSearchParams built from queries.get(0) (via
CuVSParamsHelper.buildCagraSearchParams) and thus drops per-segment settings
from subsequent CagraQuery entries; fix by either validating that all CagraQuery
instances in queries have identical search params and no per-segment filters
(throw IllegalArgumentException from MultiSegmentCagraSearch if any CagraQuery
differs from queries.get(0)), or plumb per-segment parameters through the native
call: extend the native wrapper (the cuvsCagraSearchMultiSegment binding) and
CuVSParamsHelper to accept/allocate an array of CagraSearchParams (build one
MemorySegment per CagraQuery) and pass that array/handle when invoking the
multi-segment search so each segment’s filters/params are applied.
---
Outside diff comments:
In `@java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraSearchParams.java`:
- Around line 293-322: The toString() in class CagraSearchParams currently omits
the new fields persistent, persistentLifetime, and persistentDeviceUsage; update
the CagraSearchParams.toString() method to append these three fields (with their
names and values) into the returned string (same formatting style as the other
fields) so debugging output shows the full configuration.
---
Nitpick comments:
In `@cpp/src/neighbors/detail/cagra/cagra_search.cuh`:
- Around line 388-417: Hoist the allocation and computation of query_norms out
of the segment loop and reuse it for every segment whose indices[i]->metric() ==
cuvs::distance::DistanceType::CosineExpanded: allocate query_norms once (using
raft::make_device_vector) before the for-loop, run raft::linalg::reduce on the
shared queries data to compute the norms, then inside the loop call
raft::linalg::matrix_vector_op (as currently done) using the precomputed
query_norms for each CosineExpanded segment; after the loop free or let
query_norms go out of scope. Ensure you still call
cuvs::neighbors::ivf::detail::postprocess_distances for non-CosineExpanded
branches and keep all existing ops (raft::compose_op, raft::sq_op,
raft::div_const_op, raft::cast_op, raft::add_const_op, raft::div_checkzero_op)
unchanged.
In `@java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraSearchParams.java`:
- Around line 527-537: The builder setter withPersistentDeviceUsage currently
assigns persistentDeviceUsage without enforcing the Javadoc constraint; update
Builder.withPersistentDeviceUsage to validate that persistentDeviceUsage > 0.0f
and <= 1.0f and throw an IllegalArgumentException (including the invalid value
in the message) when the check fails so invalid values are rejected early before
reaching native code.
In `@java/cuvs-java/src/test/java/com/nvidia/cuvs/CheckedCuVSResources.java`:
- Around line 59-62: The new delegate method setWorkspacePool(long sizeBytes) is
missing the wrapper's destroyed-state guard; modify
CheckedCuVSResources.setWorkspacePool to call checkNotDestroyed() at the start
(before delegating to inner.setWorkspacePool(sizeBytes)) so it matches the
class's defensive contract and uses the same destroyed-state check as other
methods.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 626e7fa1-95fc-4f40-b30d-d7fbcb6521a6
📒 Files selected for processing (40)
c/CMakeLists.txtc/include/cuvs/core/c_api.hc/include/cuvs/neighbors/cagra.hc/include/cuvs/selection/select_k.hc/src/core/c_api.cppc/src/neighbors/cagra.cppc/src/selection/select_k.cppcpp/include/cuvs/neighbors/cagra.hppcpp/src/neighbors/cagra.cuhcpp/src/neighbors/cagra_search_inst.cu.incpp/src/neighbors/detail/cagra/cagra_search.cuhcpp/src/neighbors/detail/cagra/search_single_cta.cuhcpp/src/neighbors/detail/cagra/search_single_cta_inst.cuhcpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuhcpp/src/neighbors/detail/cagra/search_single_cta_kernel.cuhjava/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.javajava/cuvs-java/src/main/java/com/nvidia/cuvs/CagraSearchParams.javajava/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.javajava/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.javajava/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSResources.javajava/cuvs-java/src/main/java/com/nvidia/cuvs/HnswAceParams.javajava/cuvs-java/src/main/java/com/nvidia/cuvs/HnswIndex.javajava/cuvs-java/src/main/java/com/nvidia/cuvs/HnswIndexParams.javajava/cuvs-java/src/main/java/com/nvidia/cuvs/MultiSegmentSearchResults.javajava/cuvs-java/src/main/java/com/nvidia/cuvs/SynchronizedCuVSResources.javajava/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.javajava/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.javajava/cuvs-java/src/main/java22/com/nvidia/cuvs/MultiSegmentCagraSearch.javajava/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BufferedCagraSearch.javajava/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.javajava/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixBaseImpl.javajava/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixInternal.javajava/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSParamsHelper.javajava/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSResourcesImpl.javajava/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CudaStreamPool.javajava/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/HnswIndexImpl.javajava/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/SelectKHelper.javajava/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/LinkerHelper.javajava/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.javajava/cuvs-java/src/test/java/com/nvidia/cuvs/CheckedCuVSResources.java
|
Pareto frontiers from initial benchmarking: Hardware: 32/64-core/thread CPU, L40S 48GB
CAGRA_SEARCH (CAGRA index used directly for search on GPU) parameter sweep ranges:
|
aamijar
added
non-breaking
| * @param[out] neighbors array of num_segments DLManagedTensor* (device, uint32, [nq, topk]) | ||
| * @param[out] distances array of num_segments DLManagedTensor* (device, float32, [nq, topk]) | ||
| */ | ||
| cuvsError_t cuvsCagraSearchMultiSegment(cuvsResources_t res, |
There was a problem hiding this comment.
Just to align on nomenclature a bit, I wonder if we can think of a more general name. Maybe "Partition"? Segment is pretty closely coupled to databases, and more specifically to LSM-based databases, but cuVS the library is more general that that. cuVS is at the level of "hash partitioning" or "blind sharding" (those are the terms we tend to use in this context). I think "MultiPartition" would be a more fitting name.
There was a problem hiding this comment.
Aligning to "partition" for now. FYI, also considered: "MultiShard", "MultiIndex", "Federated", but these might come with unintended connotations.
| * @param[out] out_idx DLManagedTensor* shape [1, k], int64, device memory | ||
| * @return cuvsError_t | ||
| */ | ||
| cuvsError_t cuvsSelectK(cuvsResources_t res, |
There was a problem hiding this comment.
Oh this is great. I was just working on code examples for the new docs and realized we only have a C++ API for select_k. It'll be great to get the C APis, and later on the Python and other language wrappers for select-k.
There was a problem hiding this comment.
With the refactoring prompted by your other comment, select-k is no longer needed for this work. Leaving the C API intact in case it might be useful to others.
| cuvs::neighbors::cagra::search_params const& params, | ||
| const std::vector<const cuvs::neighbors::cagra::index<float, uint32_t>*>& indices, | ||
| const std::vector<raft::device_matrix_view<const float, int64_t, raft::row_major>>& queries, | ||
| const std::vector<raft::device_matrix_view<int64_t, int64_t, raft::row_major>>& neighbors, |
There was a problem hiding this comment.
It seems a bit odd to return a vector of outputs instead of a single output. Can't we perform the reduction for the final single neighborhood outputs?
There was a problem hiding this comment.
Refactored this API to return the global top-k results (i.e. after select_k) instead of returning per-partition top-k results which then have to be reduced via select_k.
|
Pareto frontiers from initial benchmarking with prefiltering: 
Hardware: 32/64-core/thread CPU, L40S 48GB CAGRA_HNSW (CAGRA index converted to HNSW, search on CPU) parameter sweep ranges: "efSearch": [512, 1024] CAGRA_SEARCH (CAGRA index used directly for search on GPU) parameter sweep ranges: "efSearch": [512] |
|
Pareto frontiers unfiltered search with top-k = 1500: 
Hardware: 32/64-core/thread CPU, L40S 48GB CAGRA_HNSW (CAGRA index converted to HNSW, search on CPU) parameter sweep ranges: "efSearch": [1500, 3000] CAGRA_SEARCH (CAGRA index used directly for search on GPU) parameter sweep ranges: "efSearch": 1500 |
| void search_multi_partition( | ||
| raft::resources const& res, | ||
| cuvs::neighbors::cagra::search_params const& params, | ||
| const std::vector<const cuvs::neighbors::cagra::index<float, uint32_t>*>& indices, | ||
| raft::device_matrix_view<const float, int64_t, raft::row_major> queries, | ||
| raft::device_matrix_view<uint32_t, int64_t, raft::row_major> partition_ids, | ||
| raft::device_matrix_view<uint32_t, int64_t, raft::row_major> neighbors, | ||
| raft::device_matrix_view<float, int64_t, raft::row_major> distances, | ||
| const cuvs::neighbors::filtering::base_filter& sample_filter = | ||
| cuvs::neighbors::filtering::none_sample_filter{}); | ||
|
|
||
| void search_multi_partition( | ||
| raft::resources const& res, | ||
| cuvs::neighbors::cagra::search_params const& params, | ||
| const std::vector<const cuvs::neighbors::cagra::index<float, uint32_t>*>& indices, | ||
| raft::device_matrix_view<const float, int64_t, raft::row_major> queries, | ||
| raft::device_matrix_view<uint32_t, int64_t, raft::row_major> partition_ids, | ||
| raft::device_matrix_view<int64_t, int64_t, raft::row_major> neighbors, | ||
| raft::device_matrix_view<float, int64_t, raft::row_major> distances, | ||
| const cuvs::neighbors::filtering::base_filter& sample_filter = | ||
| cuvs::neighbors::filtering::none_sample_filter{}); | ||
|
|
||
| void search_multi_partition( | ||
| raft::resources const& res, | ||
| cuvs::neighbors::cagra::search_params const& params, | ||
| const std::vector<const cuvs::neighbors::cagra::index<half, uint32_t>*>& indices, | ||
| raft::device_matrix_view<const half, int64_t, raft::row_major> queries, | ||
| raft::device_matrix_view<uint32_t, int64_t, raft::row_major> partition_ids, | ||
| raft::device_matrix_view<uint32_t, int64_t, raft::row_major> neighbors, | ||
| raft::device_matrix_view<float, int64_t, raft::row_major> distances, | ||
| const cuvs::neighbors::filtering::base_filter& sample_filter = | ||
| cuvs::neighbors::filtering::none_sample_filter{}); | ||
|
|
||
| void search_multi_partition( | ||
| raft::resources const& res, | ||
| cuvs::neighbors::cagra::search_params const& params, | ||
| const std::vector<const cuvs::neighbors::cagra::index<half, uint32_t>*>& indices, | ||
| raft::device_matrix_view<const half, int64_t, raft::row_major> queries, | ||
| raft::device_matrix_view<uint32_t, int64_t, raft::row_major> partition_ids, | ||
| raft::device_matrix_view<int64_t, int64_t, raft::row_major> neighbors, | ||
| raft::device_matrix_view<float, int64_t, raft::row_major> distances, | ||
| const cuvs::neighbors::filtering::base_filter& sample_filter = | ||
| cuvs::neighbors::filtering::none_sample_filter{}); | ||
|
|
||
| void search_multi_partition( | ||
| raft::resources const& res, | ||
| cuvs::neighbors::cagra::search_params const& params, | ||
| const std::vector<const cuvs::neighbors::cagra::index<int8_t, uint32_t>*>& indices, | ||
| raft::device_matrix_view<const int8_t, int64_t, raft::row_major> queries, | ||
| raft::device_matrix_view<uint32_t, int64_t, raft::row_major> partition_ids, | ||
| raft::device_matrix_view<uint32_t, int64_t, raft::row_major> neighbors, | ||
| raft::device_matrix_view<float, int64_t, raft::row_major> distances, | ||
| const cuvs::neighbors::filtering::base_filter& sample_filter = | ||
| cuvs::neighbors::filtering::none_sample_filter{}); | ||
|
|
||
| void search_multi_partition( | ||
| raft::resources const& res, | ||
| cuvs::neighbors::cagra::search_params const& params, | ||
| const std::vector<const cuvs::neighbors::cagra::index<int8_t, uint32_t>*>& indices, | ||
| raft::device_matrix_view<const int8_t, int64_t, raft::row_major> queries, | ||
| raft::device_matrix_view<uint32_t, int64_t, raft::row_major> partition_ids, | ||
| raft::device_matrix_view<int64_t, int64_t, raft::row_major> neighbors, | ||
| raft::device_matrix_view<float, int64_t, raft::row_major> distances, | ||
| const cuvs::neighbors::filtering::base_filter& sample_filter = | ||
| cuvs::neighbors::filtering::none_sample_filter{}); | ||
|
|
||
| void search_multi_partition( | ||
| raft::resources const& res, | ||
| cuvs::neighbors::cagra::search_params const& params, | ||
| const std::vector<const cuvs::neighbors::cagra::index<uint8_t, uint32_t>*>& indices, | ||
| raft::device_matrix_view<const uint8_t, int64_t, raft::row_major> queries, | ||
| raft::device_matrix_view<uint32_t, int64_t, raft::row_major> partition_ids, | ||
| raft::device_matrix_view<uint32_t, int64_t, raft::row_major> neighbors, | ||
| raft::device_matrix_view<float, int64_t, raft::row_major> distances, | ||
| const cuvs::neighbors::filtering::base_filter& sample_filter = | ||
| cuvs::neighbors::filtering::none_sample_filter{}); | ||
|
|
||
| void search_multi_partition( | ||
| raft::resources const& res, | ||
| cuvs::neighbors::cagra::search_params const& params, | ||
| const std::vector<const cuvs::neighbors::cagra::index<uint8_t, uint32_t>*>& indices, | ||
| raft::device_matrix_view<const uint8_t, int64_t, raft::row_major> queries, | ||
| raft::device_matrix_view<uint32_t, int64_t, raft::row_major> partition_ids, | ||
| raft::device_matrix_view<int64_t, int64_t, raft::row_major> neighbors, | ||
| raft::device_matrix_view<float, int64_t, raft::row_major> distances, | ||
| const cuvs::neighbors::filtering::base_filter& sample_filter = | ||
| cuvs::neighbors::filtering::none_sample_filter{}); | ||
|
|
There was a problem hiding this comment.
Rather than creating a new set of overloads specific for CAGRA, maybe we'd better introduce a new index type, e.g. multi_partition_cagra? This whole thing looks like it can be represented as a "virtual merge" index we've already introduced once.
This has couple benefits:
- Public API stays the same homogeneous - no new search functions for a user to discover - just a new index type and its constructor.
- the new index search would be composable with dynamic batching (which takes some index and calls search on it).
There was a problem hiding this comment.
Thanks, @achirkin, for the insightful comment. Assuming that composite_index is the "virtual merge" index you mentioned, it does look very attractive as an existing infrastructure for multi-partition search. There are two wrinkles that prevent immediate adoption:
composite_indeximplements the per-index search in a for loop, which does not expose partition-level parallelism needed to maximize device utilization when the batch size is 1 as is always the case in cuVS-Lucene. For the current work, we actually had to create multi-partition versions of the entire call stack down to the kernels and device functions (e.g. search_multi_cta_mp_jit). Now, the multi-partition code can in theory entirely subsume the original single-partition code by essentially settingnum_partition=1as needed but I don't think we should make such a sweeping change in this PR.- For cuVS-Lucene, we need to produce the partition-local ordinals of the neighbors instead of the global indices. However,
composite_indexinternally consumes the ordinals to produce only the global indices.
Both issues above can possibly be resolved with additional configuration parameters and corresponding code paths in composite_index.
On the other hand, making multi_partition_cagra its own header probably involves easier changes. Do you think we should make a push for extending/customizing composite_index at this time?
There was a problem hiding this comment.
Rather than creating a new set of overloads specific for CAGRA, maybe we'd better introduce a new index type, e.g. multi_partition_cagra? This whole thing looks like it can be represented as a "virtual merge" index we've already introduced once.
I don't know that I agree with this. This needs to be a phyiscal merge, not a logical merge. This is critical for LSM-based databases and it'll be widely used outside of just Lucene. I'm not saying we can't (and shouldn't) improve the API experience wherever possible, but I'm not convinced the "logical merge" (or composite index) is not the way to do that.
There was a problem hiding this comment.
Hi Corey, I totally agree with you comment. I meant to say structurally the code is very similar and I think it is a good pattern to use for this case too.
There was a problem hiding this comment.
I like your thinking here for sure- if there's a good reusable abstraction for this that can consolidate APIs (ideally without introducing a complex class hierarchy on the cuvs impl side) then I'm all for it. I say "cuvs impl side" because abstractions like device_resources, while they have established a hierarchy on the user side, are just meant to be simple user convenience and we use raft::resources container on the implementation side).
There was a problem hiding this comment.
Sorry. I would not suggest this design @jamxia155. Let's keep this simple. Index shoukd also not have search methods (we use free functions in cuVS not methods on objects).
Again to reiterate- composite index is going away. We should not be using that as a guide for designs.
There was a problem hiding this comment.
After further discussions with @cjnolet, the plan is to keep things simple and essentially rename search_multi_partition as search for consistency, while allowing for necessary parameters in the signature even if they are not common to other search functions.
There was a problem hiding this comment.
I'd still suggest to go one step further and make this a new index as I proposed before. If the composite_index goes away, this could be something new like this:
template <typename BaseIndex>
struct multi_segment_index {
std::vector<BaseIndex> segments;
}The new necessary search parameters can be passed via a new parameter struct that envelopes member index parameters.
This is essentially the same as your latest version that searches the vector of indices, but with the added benefit that it is compatible with the dynamic batching index (by wrapping the new index type in the dynamic batching index).
There was a problem hiding this comment.
We are going to stick with what we have for right now so that we don't rush in a new abstraction. This feature is needed today direclty by a customer, so we want to enable them.
Let's continue the abstraction conversation in the meantime. This is a fairly low-risk change because at thie moment it's almost exclusively focused on a specific Lucene optimization. My hope is that this design will translate well to other databases, but I'd like us to really think through the abstraction layer.
There was a problem hiding this comment.
Created cuVS issue #2281 and added a TODO comment referencing it.
|
Pareto frontiers 84% filtered search with top-k = 1500: 
Hardware: 32/64-core/thread CPU, L40S 48GB CAGRA_HNSW (CAGRA index converted to HNSW, search on CPU) parameter sweep ranges: "efSearch": [1500, 2000, 3000] CAGRA_SEARCH (CAGRA index used directly for search on GPU) parameter sweep ranges: "efSearch": 1500 |
f150cc9 to
22e337c
Compare
|
/ok to test 7a7c3e6 |
|
/ok to test 1770971 |
|
/ok to test fa29064 |
|
Thanks @jamxia155. |
|
/ok to test d4c5ff7 |
imotov
left a comment
imotov
left a comment
There was a problem hiding this comment.
I reviewed the java portion and left a few comments.
| DeviceData getOrUpload(long cuvsRes) { | ||
| if (closed) throw new IllegalStateException("FilterBitsetHandle has been closed"); | ||
| DeviceData data = sharedDeviceData; | ||
| if (data != null) return data; |
There was a problem hiding this comment.
This is a slightly different variation of the cache issue in cuvs-lucene I mentioned yesterday, but we may still encounter a race condition where an acquired resource can be closed while it is still in use. I think both cases require a different approach. Perhaps we should switch to a reference-counting pattern with explicit release for managing these resources.
There was a problem hiding this comment.
Implemented reference counting. Thanks.
| var stream = getStream(cuvsRes); | ||
| // Host arena must outlive the stream sync that confirms the H2D copy. | ||
| try (var arena = Arena.ofConfined()) { | ||
| MemorySegment hostBitset = arena.allocate(combinedBitsetBytes, Long.BYTES); | ||
| MemorySegment.copy( | ||
| combinedLongs, 0, hostBitset, ValueLayout.JAVA_LONG, 0, combinedLongs.length); | ||
| cudaMemcpyAsync( | ||
| combinedBitsetDP.handle(), hostBitset, combinedBitsetBytes, HOST_TO_DEVICE, stream); | ||
|
|
||
| checkCuVSError(cuvsStreamSync(cuvsRes), "cuvsStreamSync in FilterBitsetHandle.upload"); | ||
| } |
There was a problem hiding this comment.
If any of these lines throw an exception, it appears that we may leak an unclosed CloseableRMMAllocation.
There was a problem hiding this comment.
Strengthened exception handling logic. Thanks.
| MemorySegment filterSeg = cuvsFilter.allocate(arena); | ||
| if (filter != null) { | ||
| FilterBitsetHandleImpl.DeviceData dev = | ||
| ((FilterBitsetHandleImpl) filter).getOrUpload(cuvsRes); |
There was a problem hiding this comment.
The API suggests that any implementation of FilterBitsetHandle can be used as a filter parameter. However, the current implementation assumes the object is a FilterBitsetHandleImpl and will throw a ClassCastException otherwise.
There was a problem hiding this comment.
Added a runtime check with a clearer error message, and documentation stating the requirement that the handle must come from FilterBitsetHandle.create().
| cuvsCagraSearchParams.max_queries(seg, params.getMaxQueries()); | ||
| cuvsCagraSearchParams.itopk_size(seg, params.getITopKSize()); | ||
| cuvsCagraSearchParams.max_iterations(seg, params.getMaxIterations()); | ||
| cuvsCagraSearchParams.algo(seg, params.getCagraSearchAlgo().value); |
There was a problem hiding this comment.
Since CagraSearchParams.Builder doesn't enforce non-null values, we shouldn't assume that params.getCagraSearchAlgo() is never null. Same with hashMapMode below.
There was a problem hiding this comment.
Mapped null values to AUTO/AUTO_HASH.
| CagraIndexImpl[] buffered = new CagraIndexImpl[numPartitions]; | ||
| for (int i = 0; i < numPartitions; i++) { | ||
| CagraIndex idx = indices.get(i); | ||
| if (!(idx instanceof CagraIndexImpl)) { |
There was a problem hiding this comment.
This seems somewhat fragile, though it's not unique to this change since we already rely on a similar pattern in other places. We should probably take another look at the CagraIndex abstraction at some point.
There was a problem hiding this comment.
Agree about this- ideally we'd allow the polymorphism to work it's magic here instead of having to do explicit type checking. I agree this is fairly brittle. Hopefully we can work towards improving this.
|
/ok to test 518dfcc |
|
/ok to test f775d81 |
|
/ok to test 483348d |
|
/ok to test d85a21f |
| auto scaled_sq_op = raft::compose_op(raft::sq_op{}, | ||
| raft::div_const_op<DistanceT>{DistanceT(kScale)}, | ||
| raft::cast_op<DistanceT>()); |
There was a problem hiding this comment.
Perhaps it would be cleaner to construct the mapping object directly rather than manually dividing the scales:
| auto scaled_sq_op = raft::compose_op(raft::sq_op{}, | |
| raft::div_const_op<DistanceT>{DistanceT(kScale)}, | |
| raft::cast_op<DistanceT>()); | |
| auto scaled_sq_op = raft::compose_op(raft::sq_op{}, | |
| cuvs::spatial::knn::detail::utils::mapping<DistanceT>{}); |
There was a problem hiding this comment.
Updated as suggested. Thanks!
| lightweight_uvector<DistanceT> intermediate_distances(res); | ||
| lightweight_uvector<DistanceT> transposed_distances(res); | ||
| lightweight_uvector<uint32_t> positions_buf(res); | ||
| auto query_norms = raft::make_device_vector<DistanceT, int64_t>(res, max_queries); |
There was a problem hiding this comment.
Performance danger zone: you're allocating during search. At the very least, use the workspace resource for allocating the vector, so we can rely on the user initializing it with a pool.
Then, I'd suggest to consider whether to use lightweight_uvector (if you know it's only ever used in one thread and one stream), and whether it makes sense to clump all these five allocations in one "workspace" allocation and manually map the pointers.
There was a problem hiding this comment.
Consolidated the 5 internal allocations into one lightweight_uvector. Thanks!
| raft::linalg::map( | ||
| res, | ||
| partition_ids_slice, | ||
| [per_partition_topk_u32] __device__(uint32_t pos) { return pos / per_partition_topk_u32; }, |
There was a problem hiding this comment.
You've used so nicely the raft op structs above, maybe you can use here too to avoid a lambda? :)
There was a problem hiding this comment.
Updated. Thanks for the spot.
| std::vector<part_desc_t> host_part_descs(num_partitions); | ||
|
|
||
| std::vector<dataset_descriptor_host<T, graph_idx_type, DistanceT>> part_dataset_descs; | ||
| part_dataset_descs.reserve(num_partitions); |
There was a problem hiding this comment.
Would a raft::make_host_vector<...>(res, num_partitions) work here? It would help to account for the allocation with memory tracking resources and the dry run mode in future.
There was a problem hiding this comment.
Updated host_part_descs to use RAFT host vector. However, dataset_descriptor_host is a non-trivial type and cannot be used in host_vector. Having said that, part_dataset_descs is only 48 bytes * num_partitions so hopefully it's a rounding error for memory tracking?
|
/ok to test 0ba38f0 |
|
/ok to test de32dde |
|
/ok to test d99349e |
|
/ok to test 4a60735 |
4a60735 to
33da211
Compare
|
/ok to test 33da211 |
|
/merge |
github-project-automation
Bot
moved this from In Progress
to Done
in Unstructured Data Processing
Companion PR to [cuvs!2035](NVIDIA/cuvs#2035), addresses #124. Existing CAGRA search code path for each search query: - Call CAGRA search API on one index segment - Copy results back to host - Add results into host-side global top-k priority queue - Repeat for all index segments Proposed change: - Leverage new multi-segment CAGRA search API to launch all per-segment searches in one API call - Leave results on device and run GPU-accelerated select-k API to compute global top-k - Copy final top-k results to host Authors: - James Xia (https://github.com/jamxia155) Approvers: - Kyle Edwards (https://github.com/KyleFromNVIDIA) - Corey J. Nolet (https://github.com/cjnolet) - Bradley Dice (https://github.com/bdice) URL: #133
9 participants
This PR adds a CAGRA search API and associated implementation for batching workloads involving a single query and multiple indexes (e.g. index segments built from one dataset by cuvs-lucene). This promotes GPU utilization for non-batched queries by parallelizing work over multiple CTAs, while minimizing the number of concurrent kernels for potential host-side parallelism.
Addresses cuvs-lucene issue 124, issue 158, issue 159, and cuvs issue 2166.
Associated cuvs-lucene PR: NVIDIA/cuvs-lucene#133.