Skip to content

Commit 3a79fbe

Browse files
authored
Fix overflow in preprocess_data_kernel of NN Descent (rapidsai#1596)
`preprocess_data_kernel` in NN Descent had overflow issues. casting `blockIdx.x` to `size_t` to avoid overflow. Authors: - Jinsol Park (https://github.com/jinsolp) Approvers: - Divye Gala (https://github.com/divyegala) URL: rapidsai#1596
1 parent 55fca1e commit 3a79fbe

File tree

1 file changed

+5
-1
lines changed

1 file changed

+5
-1
lines changed

cpp/src/neighbors/detail/nn_descent.cuh

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -243,7 +243,11 @@ RAFT_KERNEL preprocess_data_kernel(
243243
Data_t* s_vec = (Data_t*)buffer;
244244
size_t list_id = list_offset + blockIdx.x;
245245

246-
load_vec(s_vec, input_data + blockIdx.x * dim, dim, dim, threadIdx.x % raft::warp_size());
246+
load_vec(s_vec,
247+
input_data + static_cast<size_t>(blockIdx.x) * dim,
248+
dim,
249+
dim,
250+
threadIdx.x % raft::warp_size());
247251
if (threadIdx.x == 0) { l2_norm = 0; }
248252
__syncthreads();
249253

0 commit comments

Comments
 (0)