Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Feature] Add check for NNZ in COOToCSR #7459

Merged
merged 13 commits into from
Jun 17, 2024
1 change: 1 addition & 0 deletions src/array/cpu/spmat_op_impl_coo.cc
Original file line number Diff line number Diff line change
Expand Up @@ -746,6 +746,7 @@ inline COOToCSRAlg WhichCOOToCSR(const COOMatrix &coo) {

template <DGLDeviceType XPU, typename IdType>
CSRMatrix COOToCSR(COOMatrix coo) {
CHECK_NO_OVERFLOW(coo.row->dtype, coo.row->shape[0]);
switch (WhichCOOToCSR<IdType>(coo)) {
case COOToCSRAlg::sorted:
return SortedCOOToCSR<IdType>(coo);
Expand Down
1 change: 1 addition & 0 deletions src/array/cuda/coo2csr.cu
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@ CSRMatrix COOToCSR<kDGLCUDA, int32_t>(COOMatrix coo) {
}

const int64_t nnz = coo.row->shape[0];
CHECK_NO_OVERFLOW(coo.row->dtype, nnz);
// TODO(minjie): Many of our current implementation assumes that CSR must have
// a data array. This is a temporary workaround. Remove this after:
// - The old immutable graph implementation is deprecated.
Expand Down
Loading