Skip to content

Commit 59bc6bd

Browse files
authored
Use 64 bit types for dataset size calculation in CAGRA graph optimizer (#380)
This PR fixes #375. Authors: - Tamas Bela Feher (https://github.com/tfeher) Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: #380
1 parent d093986 commit 59bc6bd

File tree

1 file changed

+4
-4
lines changed

1 file changed

+4
-4
lines changed

cpp/src/neighbors/detail/cagra/graph_core.cuh

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -475,12 +475,12 @@ void sort_knn_graph(
475475
{
476476
RAFT_EXPECTS(dataset.extent(0) == knn_graph.extent(0),
477477
"dataset size is expected to have the same number of graph index size");
478-
const uint32_t dataset_size = dataset.extent(0);
479-
const uint32_t dataset_dim = dataset.extent(1);
478+
const uint64_t dataset_size = dataset.extent(0);
479+
const uint64_t dataset_dim = dataset.extent(1);
480480
const DataT* dataset_ptr = dataset.data_handle();
481481

482482
const IdxT graph_size = dataset_size;
483-
const uint32_t input_graph_degree = knn_graph.extent(1);
483+
const uint64_t input_graph_degree = knn_graph.extent(1);
484484
IdxT* const input_graph_ptr = knn_graph.data_handle();
485485

486486
auto large_tmp_mr = raft::resource::get_large_workspace_resource(res);
@@ -528,7 +528,7 @@ void sort_knn_graph(
528528
kernel_sort = kern_sort<DataT, IdxT, numElementsPerThread>;
529529
} else {
530530
RAFT_FAIL(
531-
"The degree of input knn graph is too large (%u). "
531+
"The degree of input knn graph is too large (%lu). "
532532
"It must be equal to or smaller than %d.",
533533
input_graph_degree,
534534
1024);

0 commit comments

Comments
 (0)