API for CPU side output indices and distances of all_neighbors#1905
API for CPU side output indices and distances of all_neighbors#1905jinsolp wants to merge 9 commits into
all_neighbors#1905Conversation
|
@aamijar yes, so eventually all-neighbors will be further optimized for host side distances/indices! |
📝 WalkthroughSummary by CodeRabbit
WalkthroughThis PR enables all-neighbors KNN to return indices and distances in host (CPU) memory instead of GPU-only memory. The core C++/CUDA implementation now accepts generic mdspan types for outputs, new public overloads expose host output options, and C/Python bindings route outputs to host or device based on buffer types or explicit flags. Tests validate all output placement modes. ChangesHost-side Output Support for All-Neighbors
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
cpp/src/neighbors/all_neighbors/all_neighbors.cuh (1)
197-211:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winAdd validation that
k(indices.extent(1)) is greater than zero.The code validates row counts but not column counts. If
indices.extent(1) == 0, subsequent operations likedistances.value()(i, k - 1)at line 261 would underflow the index, and shift operations would behave unexpectedly.🛡️ Proposed fix
auto build_algo = check_params_validity(params, core_distances.has_value()); +RAFT_EXPECTS(indices.extent(1) > 0, "k (number of neighbors) must be greater than 0"); + RAFT_EXPECTS(dataset.extent(0) == indices.extent(0), "number of rows in dataset should be the same as number of rows in indices matrix");The same validation should be added to the device-dataset overload at line 324.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/src/neighbors/all_neighbors/all_neighbors.cuh` around lines 197 - 211, Add a validation that the number of neighbors k (indices.extent(1)) is > 0 to prevent underflow when referencing distances.value()(i, k - 1); after computing build_algo = check_params_validity(...) use RAFT_EXPECTS(indices.extent(1) > 0, "k (number of neighbors) must be > 0") and include the same check in the device-dataset overload around the block starting at the device-dataset overload (the later overload near line 324) so both code paths validate k before any use of indices.extent(1) or distances.value()(i, k - 1).
🧹 Nitpick comments (5)
python/cuvs/cuvs/tests/test_all_neighbors.py (2)
202-202: ⚡ Quick winPrefix unused variable with underscore.
The
bf_distancesvariable is unpacked but never used. Use_or_bf_distancesto indicate this is intentional.🔧 Proposed fix
- bf_distances, bf_indices = brute_force.search(bf_index, X_device, k=k) + _, bf_indices = brute_force.search(bf_index, X_device, k=k)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cuvs/cuvs/tests/test_all_neighbors.py` at line 202, The variable bf_distances returned from brute_force.search is unpacked but never used; change the unpack to use a prefixed-underscore name (e.g., _bf_distances or _) when calling brute_force.search(bf_index, X_device, k=k) so the unused value is explicit and linter-friendly, leaving bf_indices as the used variable.
345-345: ⚡ Quick winPrefix unused variable with underscore.
The
bf_distancesvariable is unpacked but never used. Use_or_bf_distancesto indicate this is intentional.🔧 Proposed fix
- bf_distances, bf_indices = brute_force.search(bf_index, X_device, k=k) + _, bf_indices = brute_force.search(bf_index, X_device, k=k)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cuvs/cuvs/tests/test_all_neighbors.py` at line 345, The test unpacks brute_force.search into bf_distances and bf_indices but never uses bf_distances; change the unpack to use a throwaway name (e.g., _ or _bf_distances) so the unused variable is clearly intentional: update the call to brute_force.search(bf_index, X_device, k=k) to assign the first return to a prefixed underscore (e.g., _, bf_indices) or rename bf_distances to _bf_distances wherever it’s not used.python/cuvs/cuvs/neighbors/all_neighbors/all_neighbors.pyx (1)
333-338: 💤 Low valueConsider making the warning more specific.
The warning could be clearer about which parameter takes precedence. Consider specifying whether
indicesordistancesis driving the placement decision.💡 Suggested improvement
- warnings.warn( - "return_on_host is ignored when indices or distances buffers " - "are provided. Output placement is inferred from the provided " - "arrays instead.", - ) + which_provided = "indices" if indices is not None else "distances" + warnings.warn( + f"return_on_host is ignored because {which_provided} buffer " + f"was provided. Output placement is inferred from the provided " + f"array's memory location instead.", + )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cuvs/cuvs/neighbors/all_neighbors/all_neighbors.pyx` around lines 333 - 338, The current warning in all_neighbors.pyx is vague about precedence when return_on_host is true and indices/distances buffers are provided; update the warnings.warn call to state which buffer determines output placement by inspecting indices and distances and including that information in the message (e.g., "placement inferred from provided 'indices' buffer", "placement inferred from provided 'distances' buffer", or "placement inferred from provided 'indices' and 'distances' buffers"), so change the logic around the return_on_host check to compute which symbol(s) (indices, distances) are not None and interpolate them into the warning text.cpp/src/neighbors/all_neighbors/all_neighbors.cuh (2)
232-248: ⚖️ Poor tradeoffConsider extracting the shift logic into a reusable helper.
The shift-handling code (with and without core_distances) is duplicated across both
buildoverloads. This pattern appears 4 times with minor variations. Extracting a templated helper that branches onoutputs_are_hostwould reduce maintenance burden and risk of divergence.// Example signature template <typename T, typename IdxT, typename IndicesMdspan, typename DistancesMdspan, typename CoreDistMdspan> void apply_shift(const raft::resources& handle, IndicesMdspan indices, std::optional<DistancesMdspan> distances, std::optional<CoreDistMdspan> core_distances);This can be addressed in a follow-up.
Also applies to: 283-307, 353-370, 400-424
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/src/neighbors/all_neighbors/all_neighbors.cuh` around lines 232 - 248, The duplicated shift logic should be extracted into a single templated helper (e.g., apply_shift<T,IdxT,IndicesMdspan,DistancesMdspan,CoreDistMdspan>) that accepts the raft::resources& handle, the indices mdspan, optional distances mdspan, and optional core_distances mdspan and internally branches on the constexpr outputs_are_host: when host, call host_shift_columns(indices,1,...) and host_shift_columns for distances/core_distances (using std::make_optional<T>(0.0) for padding); when device, construct raft::make_device_matrix_view for indices/distances/core_distances and call raft::matrix::shift(handle, ..., 1[, std::make_optional<T>(0.0)]). Replace the four duplicated blocks in the build overloads with calls to this helper (preserving template parameters and types IdxT and T) so all variations (with/without distances and core_distances) are handled uniformly.
254-276: 💤 Low valueStream sync may be needed before the OpenMP loop accesses
distances.At line 259-262, the OpenMP loop reads from
distances.value()which was just populated bysingle_build. Whenoutputs_are_hostis true,single_builddoes sync the stream at line 180 before returning. However, if the dataflow changes in future refactors, this dependency on prior synchronization could be subtle.Consider adding a brief comment noting the sync dependency, or verifying that
single_buildalways syncs when returning host data.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/src/neighbors/all_neighbors/all_neighbors.cuh` around lines 254 - 276, The OpenMP loop in the outputs_are_host branch reads distances.value() (and writes core_distances.value()) immediately after single_build; to avoid a subtle race if single_build's sync behavior changes, ensure the device stream is synchronized before that host-access loop (or add a clear comment asserting single_build always synchronizes when returning host outputs). Concretely, in the outputs_are_host branch of all_neighbors.cuh (around the OpenMP loop that copies from distances.value() into core_distances.value()), either perform an explicit stream synchronization on the stream from raft::resource::get_cuda_stream(handle) before the loop or add a guarded comment referencing single_build and outputs_are_host to document the required sync contract.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@c/include/cuvs/neighbors/all_neighbors.h`:
- Line 100: Update the doc comment for the core_distances parameter in
all_neighbors.h to reflect that it may reside on host or device (matching the
memory location of indices/distances) rather than "on device" only; state it is
an optional 1D tensor [num_rows] of float32 that can be NULL and that its
location follows the indices/distances memory location so callers know
host/device placement is supported.
- Around line 31-32: The header comment incorrectly states that core_distances
can only be on device memory; update the documentation in all_neighbors.h to
state that core_distances follows the same memory placement rules as indices and
distances (i.e., it may be host- or device-memory depending on the arrays
passed), consistent with the validation logic in
c/src/neighbors/all_neighbors.cpp that checks core_distances along with
indices/distances (see the validation behavior around core_distances and the
indices/distances checks). Ensure the wording mirrors the behavior enforced in
the implementation so users know core_distances must match the memory location
semantics of indices/distances.
In `@cpp/src/neighbors/all_neighbors/all_neighbors_batched.cuh`:
- Around line 612-621: The async raft::copy calls that transfer to host-backed
mdspans (indices, distances) can complete after batch_build() returns, causing a
race when host_shift_columns(indices) runs; add a GPU stream synchronization
after those copies by calling raft::resource::sync_stream(handle) (matching the
single_build() pattern) immediately after the raft::copy blocks so the host
memory is safe to access before returning from
batch_build()/all_neighbors_batched.
---
Outside diff comments:
In `@cpp/src/neighbors/all_neighbors/all_neighbors.cuh`:
- Around line 197-211: Add a validation that the number of neighbors k
(indices.extent(1)) is > 0 to prevent underflow when referencing
distances.value()(i, k - 1); after computing build_algo =
check_params_validity(...) use RAFT_EXPECTS(indices.extent(1) > 0, "k (number of
neighbors) must be > 0") and include the same check in the device-dataset
overload around the block starting at the device-dataset overload (the later
overload near line 324) so both code paths validate k before any use of
indices.extent(1) or distances.value()(i, k - 1).
---
Nitpick comments:
In `@cpp/src/neighbors/all_neighbors/all_neighbors.cuh`:
- Around line 232-248: The duplicated shift logic should be extracted into a
single templated helper (e.g.,
apply_shift<T,IdxT,IndicesMdspan,DistancesMdspan,CoreDistMdspan>) that accepts
the raft::resources& handle, the indices mdspan, optional distances mdspan, and
optional core_distances mdspan and internally branches on the constexpr
outputs_are_host: when host, call host_shift_columns(indices,1,...) and
host_shift_columns for distances/core_distances (using
std::make_optional<T>(0.0) for padding); when device, construct
raft::make_device_matrix_view for indices/distances/core_distances and call
raft::matrix::shift(handle, ..., 1[, std::make_optional<T>(0.0)]). Replace the
four duplicated blocks in the build overloads with calls to this helper
(preserving template parameters and types IdxT and T) so all variations
(with/without distances and core_distances) are handled uniformly.
- Around line 254-276: The OpenMP loop in the outputs_are_host branch reads
distances.value() (and writes core_distances.value()) immediately after
single_build; to avoid a subtle race if single_build's sync behavior changes,
ensure the device stream is synchronized before that host-access loop (or add a
clear comment asserting single_build always synchronizes when returning host
outputs). Concretely, in the outputs_are_host branch of all_neighbors.cuh
(around the OpenMP loop that copies from distances.value() into
core_distances.value()), either perform an explicit stream synchronization on
the stream from raft::resource::get_cuda_stream(handle) before the loop or add a
guarded comment referencing single_build and outputs_are_host to document the
required sync contract.
In `@python/cuvs/cuvs/neighbors/all_neighbors/all_neighbors.pyx`:
- Around line 333-338: The current warning in all_neighbors.pyx is vague about
precedence when return_on_host is true and indices/distances buffers are
provided; update the warnings.warn call to state which buffer determines output
placement by inspecting indices and distances and including that information in
the message (e.g., "placement inferred from provided 'indices' buffer",
"placement inferred from provided 'distances' buffer", or "placement inferred
from provided 'indices' and 'distances' buffers"), so change the logic around
the return_on_host check to compute which symbol(s) (indices, distances) are not
None and interpolate them into the warning text.
In `@python/cuvs/cuvs/tests/test_all_neighbors.py`:
- Line 202: The variable bf_distances returned from brute_force.search is
unpacked but never used; change the unpack to use a prefixed-underscore name
(e.g., _bf_distances or _) when calling brute_force.search(bf_index, X_device,
k=k) so the unused value is explicit and linter-friendly, leaving bf_indices as
the used variable.
- Line 345: The test unpacks brute_force.search into bf_distances and bf_indices
but never uses bf_distances; change the unpack to use a throwaway name (e.g., _
or _bf_distances) so the unused variable is clearly intentional: update the call
to brute_force.search(bf_index, X_device, k=k) to assign the first return to a
prefixed underscore (e.g., _, bf_indices) or rename bf_distances to
_bf_distances wherever it’s not used.
🪄 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: Enterprise
Run ID: a683c13b-804f-432c-a08e-15d4bdd99cdb
📒 Files selected for processing (9)
c/include/cuvs/neighbors/all_neighbors.hc/src/neighbors/all_neighbors.cppcpp/include/cuvs/neighbors/all_neighbors.hppcpp/src/neighbors/all_neighbors/all_neighbors.cucpp/src/neighbors/all_neighbors/all_neighbors.cuhcpp/src/neighbors/all_neighbors/all_neighbors_batched.cuhcpp/tests/neighbors/all_neighbors.cuhpython/cuvs/cuvs/neighbors/all_neighbors/all_neighbors.pyxpython/cuvs/cuvs/tests/test_all_neighbors.py
| * - Outputs (indices, distances) can be on host memory (numpy arrays) | ||
| * or device memory (CUDA arrays). core_distances can only be on device memory. |
There was a problem hiding this comment.
Documentation inconsistency: core_distances memory placement.
The documentation states that core_distances can only be on device memory, but the implementation in c/src/neighbors/all_neighbors.cpp (lines 109-112) validates core_distances as either host-compatible or device-compatible, allowing both memory locations. This inconsistency will confuse API users.
Update the documentation to clarify that core_distances follows the same memory location as indices and distances (as enforced in the implementation at lines 154-158 and 243-247 of the C++ file).
📝 Proposed documentation fix
- * - Outputs (indices, distances) can be on host memory (numpy arrays)
- * or device memory (CUDA arrays). core_distances can only be on device memory.
+ * - Outputs (indices, distances, core_distances) can be on host memory (numpy arrays)
+ * or device memory (CUDA arrays). When core_distances is provided, it must be on the
+ * same memory location as indices and distances.📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| * - Outputs (indices, distances) can be on host memory (numpy arrays) | |
| * or device memory (CUDA arrays). core_distances can only be on device memory. | |
| * - Outputs (indices, distances, core_distances) can be on host memory (numpy arrays) | |
| * or device memory (CUDA arrays). When core_distances is provided, it must be on the | |
| * same memory location as indices and distances. |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@c/include/cuvs/neighbors/all_neighbors.h` around lines 31 - 32, The header
comment incorrectly states that core_distances can only be on device memory;
update the documentation in all_neighbors.h to state that core_distances follows
the same memory placement rules as indices and distances (i.e., it may be host-
or device-memory depending on the arrays passed), consistent with the validation
logic in c/src/neighbors/all_neighbors.cpp that checks core_distances along with
indices/distances (see the validation behavior around core_distances and the
indices/distances checks). Ensure the wording mirrors the behavior enforced in
the implementation so users know core_distances must match the memory location
semantics of indices/distances.
| * @param[out] distances Optional 2D tensor [num_rows x k] on device (float32); can be NULL | ||
| * @param[out] indices 2D tensor [num_rows x k] on host or device (int64) | ||
| * @param[out] distances Optional 2D tensor [num_rows x k] on host or device (float32); can be NULL | ||
| * @param[out] core_distances Optional 1D tensor [num_rows] on device (float32); can be NULL |
There was a problem hiding this comment.
Update core_distances parameter documentation.
The parameter description still states "on device" only, but the implementation allows both host and device memory (matching the location of indices/distances).
📝 Proposed fix
- * `@param`[out] core_distances Optional 1D tensor [num_rows] on device (float32); can be NULL
+ * `@param`[out] core_distances Optional 1D tensor [num_rows] on host or device (float32); can be NULL. Must be on the same memory location as indices and distances.📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| * @param[out] core_distances Optional 1D tensor [num_rows] on device (float32); can be NULL | |
| * `@param`[out] core_distances Optional 1D tensor [num_rows] on host or device (float32); can be NULL. Must be on the same memory location as indices and distances. |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@c/include/cuvs/neighbors/all_neighbors.h` at line 100, Update the doc comment
for the core_distances parameter in all_neighbors.h to reflect that it may
reside on host or device (matching the memory location of indices/distances)
rather than "on device" only; state it is an optional 1D tensor [num_rows] of
float32 that can be NULL and that its location follows the indices/distances
memory location so callers know host/device placement is supported.
| raft::copy(indices.data_handle(), | ||
| global_neighbors.data_handle(), | ||
| num_rows * k, | ||
| raft::resource::get_cuda_stream(handle)); | ||
| if (distances.has_value()) { | ||
| raft::copy( | ||
| handle, | ||
| raft::make_device_vector_view(distances.value().data_handle(), num_rows * k), | ||
| raft::make_device_vector_view<const T>(global_distances.data_handle(), num_rows * k)); | ||
| raft::copy(distances.value().data_handle(), | ||
| global_distances.data_handle(), | ||
| num_rows * k, | ||
| raft::resource::get_cuda_stream(handle)); | ||
| } |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Search for synchronization in the callers of batch_build
rg -n "batch_build" -A 10 -B 5 --type cpp --type cuh -g '!test*'Repository: rapidsai/cuvs
Length of output: 85
🏁 Script executed:
#!/bin/bash
# Check if there's a pattern of syncing after copy to host in similar code
rg -n "raft::copy.*host.*data_handle" -A 3 --type cpp --type cuh | rg -A 3 "sync_stream"Repository: rapidsai/cuvs
Length of output: 85
🏁 Script executed:
#!/bin/bash
# Search for batch_build calls with corrected file type
rg -n "batch_build" -A 10 -B 5 --type cpp --type c --type hRepository: rapidsai/cuvs
Length of output: 39
🏁 Script executed:
#!/bin/bash
# Also look at the actual file to understand context
head -c 30000 cpp/src/neighbors/all_neighbors/all_neighbors_batched.cuh | tail -c 15000Repository: rapidsai/cuvs
Length of output: 15057
🏁 Script executed:
#!/bin/bash
# Find all files that call batch_build
rg -l "batch_build" --type cpp --type c --type hRepository: rapidsai/cuvs
Length of output: 39
🏁 Script executed:
#!/bin/bash
# Search for explicit calls or instantiations of batch_build
rg "batch_build\s*\(" --type cpp --type c --type h -A 5 -B 2Repository: rapidsai/cuvs
Length of output: 39
🏁 Script executed:
#!/bin/bash
# Look for wrapper functions or public APIs that call batch_build
rg "batch_build" --type cpp --type c --type h -lRepository: rapidsai/cuvs
Length of output: 39
🏁 Script executed:
#!/bin/bash
# Check the function signature constraints and comments
sed -n '560,625p' cpp/src/neighbors/all_neighbors/all_neighbors_batched.cuhRepository: rapidsai/cuvs
Length of output: 2848
🏁 Script executed:
#!/bin/bash
# Search for public API headers that might wrap batch_build
fd -e hpp -e h . cpp/src/neighbors/all_neighbors/ | head -20Repository: rapidsai/cuvs
Length of output: 39
🏁 Script executed:
#!/bin/bash
# Look at the header file for this module
cat cpp/src/neighbors/all_neighbors.hppRepository: rapidsai/cuvs
Length of output: 125
🏁 Script executed:
#!/bin/bash
# Check if there's documentation about output memory location expectations
rg -i "indices.*host\|output.*host\|memory.*location" cpp/src/neighbors/all_neighbors/ -A 2 -B 2Repository: rapidsai/cuvs
Length of output: 39
🏁 Script executed:
#!/bin/bash
# Find all header files in the all_neighbors directory
find cpp/src/neighbors/all_neighbors -type f -name "*.hpp" -o -name "*.h"Repository: rapidsai/cuvs
Length of output: 39
🏁 Script executed:
#!/bin/bash
# Look for the main public API file for all_neighbors
find cpp/include -name "*all_neighbors*" -type fRepository: rapidsai/cuvs
Length of output: 102
🏁 Script executed:
#!/bin/bash
# Check what files are in the include directory structure
ls -la cpp/include/cuvs/neighbors/Repository: rapidsai/cuvs
Length of output: 1303
🏁 Script executed:
#!/bin/bash
# Read the all_neighbors public API
cat cpp/include/cuvs/neighbors/all_neighbors.hppRepository: rapidsai/cuvs
Length of output: 11892
🏁 Script executed:
#!/bin/bash
# Find the implementation of the public build() functions
find cpp/src/neighbors -name "*.cpp" -o -name "*.hpp" | xargs grep -l "void build" | grep -i all_neighborsRepository: rapidsai/cuvs
Length of output: 39
🏁 Script executed:
#!/bin/bash
# Check if there's a source file implementing these build functions
find cpp/src/neighbors/all_neighbors -type fRepository: rapidsai/cuvs
Length of output: 328
🏁 Script executed:
#!/bin/bash
# Look at the actual implementation that wraps batch_build
cat cpp/src/neighbors/all_neighbors/all_neighbors.cppRepository: rapidsai/cuvs
Length of output: 139
🏁 Script executed:
#!/bin/bash
# Look at the .cu file which should have the public API implementations
cat cpp/src/neighbors/all_neighbors/all_neighbors.cuRepository: rapidsai/cuvs
Length of output: 4653
🏁 Script executed:
#!/antml:bash
# Also check the .cuh file for inline/template implementations
head -200 cpp/src/neighbors/all_neighbors/all_neighbors.cuhRepository: rapidsai/cuvs
Length of output: 8895
🏁 Script executed:
#!/bin/bash
# Search for the batch_build template instantiations in all_neighbors.cuh to see if it's called with host mdspans
sed -n '1,100p' cpp/src/neighbors/all_neighbors/all_neighbors.cuh | tail -50Repository: rapidsai/cuvs
Length of output: 2945
🏁 Script executed:
#!/bin/bash
# Look for batch_build calls in the build() function to see what template arguments are passed
sed -n '100,300p' cpp/src/neighbors/all_neighbors/all_neighbors.cuhRepository: rapidsai/cuvs
Length of output: 8956
🏁 Script executed:
#!/bin/bash
# Check how batch_build is called from the build() function
sed -n '250,280p' cpp/src/neighbors/all_neighbors/all_neighbors.cuhRepository: rapidsai/cuvs
Length of output: 1578
🏁 Script executed:
#!/antml:bash
# Compare synchronization patterns in single_build vs batch_build
echo "=== single_build sync pattern ===" && sed -n '75,92p' cpp/src/neighbors/all_neighbors/all_neighbors.cuh && echo -e "\n=== batch_build end (no sync) ===" && sed -n '612,625p' cpp/src/neighbors/all_neighbors/all_neighbors_batched.cuhRepository: rapidsai/cuvs
Length of output: 1476
🏁 Script executed:
#!/an/bash
# Check if batch_build is called with host mdspans - look at the build() function calls
sed -n '240,290p' cpp/src/neighbors/all_neighbors/all_neighbors.cuhRepository: rapidsai/cuvs
Length of output: 2412
🏁 Script executed:
#!/bin/bash
# Find the exact location where batch_build is called in build() function
grep -n "batch_build" cpp/src/neighbors/all_neighbors/all_neighbors.cuhRepository: rapidsai/cuvs
Length of output: 314
🏁 Script executed:
#!/bin/bash
# Look at the build() template function to see if it does sync after batch_build
sed -n '170,230p' cpp/src/neighbors/all_neighbors/all_neighbors.cuhRepository: rapidsai/cuvs
Length of output: 2669
🏁 Script executed:
#!/bin/bash
# Check what happens after batch_build returns in the build() function - look for any host memory access patterns
sed -n '220,300p' cpp/src/neighbors/all_neighbors/all_neighbors.cuh | head -50Repository: rapidsai/cuvs
Length of output: 2367
🏁 Script executed:
#!/bin/bash
# Confirm that batch_build can receive host mdspans by checking the template signature
sed -n '480,520p' cpp/src/neighbors/all_neighbors/all_neighbors_batched.cuhRepository: rapidsai/cuvs
Length of output: 1857
Critical race condition: async copies to host memory without synchronization.
When indices or distances are host-memory mdspans, the async raft::copy operations (lines 612-620) complete asynchronously. However, the calling code in all_neighbors.cuh:231-234 immediately invokes host_shift_columns(indices) after batch_build() returns, which directly reads/writes the host buffer via OpenMP parallel loops without any GPU synchronization. This creates a data race where CPU threads access host memory while the GPU async copy is still in progress.
Compare with single_build() in all_neighbors.cuh:88, which correctly adds raft::resource::sync_stream(handle); after async copies to host. Apply the same pattern here:
raft::copy(indices.data_handle(),
global_neighbors.data_handle(),
num_rows * k,
raft::resource::get_cuda_stream(handle));
if (distances.has_value()) {
raft::copy(distances.value().data_handle(),
global_distances.data_handle(),
num_rows * k,
raft::resource::get_cuda_stream(handle));
}
raft::resource::sync_stream(handle); // Add this before function return🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@cpp/src/neighbors/all_neighbors/all_neighbors_batched.cuh` around lines 612 -
621, The async raft::copy calls that transfer to host-backed mdspans (indices,
distances) can complete after batch_build() returns, causing a race when
host_shift_columns(indices) runs; add a GPU stream synchronization after those
copies by calling raft::resource::sync_stream(handle) (matching the
single_build() pattern) immediately after the raft::copy blocks so the host
memory is safe to access before returning from
batch_build()/all_neighbors_batched.
dantegd
left a comment
dantegd
left a comment
There was a problem hiding this comment.
Just had a few questions and comments
| * | ||
| * Notes: | ||
| * - Outputs (indices, distances, core_distances) are expected to be on device memory. | ||
| * - Outputs (indices, distances) can be on host memory (numpy arrays) |
There was a problem hiding this comment.
It's kinda odd to see Python terminology in the C header, I guess one of the main consumers are the python apis, but still... not necessarily suggesting to change it, but was wondering ig it might be better to keep the docstring to C relevant concepts (like dlpack)
| * `n_clusters` clusters and assigns each row to `overlap_factor` nearest clusters. For device | ||
| * datasets, `n_clusters` must be 1 (no batching); `overlap_factor` is ignored. | ||
| * Outputs always reside in device memory. | ||
| * Outputs can be on host memory (numpy arrays) or device memory (CUDA arrays). |
There was a problem hiding this comment.
Same thing, it sticks out to me. To be honest I didn't check if we used it in other places, but also the other thing is that numpy arrays are just one potential origin for host memory, it could just be C arrays, or say other host python arrays from the python layer, and this seems to imply that there's something that would make these functions need numpy arrays, so I'm leaning towards removing numpy from C docs.
| #pragma omp parallel for | ||
| for (size_t i = 0; i < n_rows; i++) { | ||
| // Copy columns from right to left | ||
| for (size_t j = n_cols - 1; j >= k; j--) { |
There was a problem hiding this comment.
Minor thing, but j is unsigned, so j >= 0 is always true. The current RAFT_EXPECTS(n_cols > k, …) doesn't catch k == 0 which would make this be an infinite loop.
Today everyone calls this with k>=1, so it's not exploitable, but since it's a generic-looking helper in detail::, mind adding a RAFT_EXPECTS(k > 0, …) so a future caller can't trip on it?
| "output_location", | ||
| ["host_arrays", "device_arrays", "return_on_host", "return_on_device"], |
There was a problem hiding this comment.
Nice coverage of the placement combos. One gap is that there's no test that exercises core_distances with host outputs. Might be a good idea tp parametrize output_location on the existing mutual-reach test, or maybe add a small dedicated test
| cpp_res, cpp_params, dataset, indices, distances, core_distances, alpha); | ||
| auto dataset = cuvs::core::from_dlpack<dataset_mdspan_t>(dataset_tensor); | ||
|
|
||
| if (indices_is_host) { |
There was a problem hiding this comment.
The if (indices_is_host) { … } else { … } blocks here are essentially the same 30-ish lines duplicated four times across _build_host and _build_device (host branch and device branch in each). Would you be open to extracting a small templated helper that takes the dataset mdspan type + a host/device tag and does the optional unpacking + dispatch? Should cut this down a lot, or is this cod that'll significantly change in teh follow up and not worth the change?
| const_cast<IdxT*>(indices.data_handle()), indices.extent(0), indices.extent(1)); | ||
| raft::matrix::shift(handle, indices_d, 1); | ||
| if (distances.has_value()) { | ||
| auto distances_d = | ||
| raft::make_device_matrix_view<T, IdxT>(const_cast<T*>(distances.value().data_handle()), |
There was a problem hiding this comment.
Tiny nit (applies to all 8 of these in this file): I think the const_cast<IdxT*> and const_cast<T*> here are no-ops? We're in the else branch, so outputs_are_host is false, which means IndicesMdspan/DistancesMdspan are the device instantiations from all_neighbors.cu: device_matrix_view<IdxT, IdxT> and device_matrix_view<T, IdxT>, and data_handle() already gives us IdxT*/T*. single_build does the same thing right above without a cast. Mind dropping these? My worry is mostly that if someone later adds a device_matrix_view<const IdxT, …> instantiation, the cast would silently strip const on an output buffer.
That said, if you think the follow up will change it singificantly regardless, then happy to drop the comment here.
4 participants
Closes #1902
This PR allows indices/distances result of the
all_neighborsfunction to be on host memory. Focuses on exposing the API, and internally just does a copy.Internally, it still allocates GPU memory for the full indices/distances
n_clusters=1this is not a problem. If the user faces OOM issues they should usen_clusters>1.n_clusters>1we are still using managed memory regardless of the memory location of the user given indices/distances. Will open a follow-up PR with optimizations (related issue: [FEA] Improve batchedall-neighborsgiven CPU indices/distances #1903)