Fixing overflow in expanded distances (#365)

Closes rapidsai/cuml#4289 Authors: - Corey J. Nolet (https://github.com/cjnolet) Approvers: - Dante Gama Dessavre (https://github.com/dantegd) URL: #365
rapidsai · Oct 27, 2021 · 091cfdf · 091cfdf
1 parent ccba252
commit 091cfdf
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/cpp/include/raft/sparse/distance/detail/bin_distance.cuh b/cpp/include/raft/sparse/distance/detail/bin_distance.cuh
@@ -54,7 +54,7 @@ __global__ void compute_binary_warp_kernel(value_t *__restrict__ C,
                                            const value_t *__restrict__ R_norms,
                                            value_idx n_rows, value_idx n_cols,
                                            expansion_f expansion_func) {
-  value_idx tid = blockDim.x * blockIdx.x + threadIdx.x;
+  std::size_t tid = blockDim.x * blockIdx.x + threadIdx.x;
   value_idx i = tid / n_cols;
   value_idx j = tid % n_cols;
 

diff --git a/cpp/include/raft/sparse/distance/detail/l2_distance.cuh b/cpp/include/raft/sparse/distance/detail/l2_distance.cuh
@@ -64,7 +64,7 @@ __global__ void compute_euclidean_warp_kernel(
   value_t *__restrict__ C, const value_t *__restrict__ Q_sq_norms,
   const value_t *__restrict__ R_sq_norms, value_idx n_rows, value_idx n_cols,
   expansion_f expansion_func) {
-  value_idx tid = blockDim.x * blockIdx.x + threadIdx.x;
+  std::size_t tid = blockDim.x * blockIdx.x + threadIdx.x;
   value_idx i = tid / n_cols;
   value_idx j = tid % n_cols;
 
@@ -85,7 +85,7 @@ __global__ void compute_correlation_warp_kernel(
   const value_t *__restrict__ R_sq_norms, const value_t *__restrict__ Q_norms,
   const value_t *__restrict__ R_norms, value_idx n_rows, value_idx n_cols,
   value_idx n) {
-  value_idx tid = blockDim.x * blockIdx.x + threadIdx.x;
+  std::size_t tid = blockDim.x * blockIdx.x + threadIdx.x;
   value_idx i = tid / n_cols;
   value_idx j = tid % n_cols;