UMAP fix int32 overflow causing illegal mem access

[aamijar]

Resolves #7517

The optimize_batch_kernel and optimize_batch_kernel_reg use the following condition to check for out of bounds

while (row < nnz)

Inside the loop row += skip_size is used for iteration. However, when row is close to INT32_MAX it can cause an overflow, which leads to the value of row wrapping around to become a negative number. The loop will then be run again since row < nnz still, which will lead to a cuda illegal memory access since row is negative.

To solve this we can check for the overflow case and break from the while loop
if (row > nnz - skip_size) break;

Update: We can use size_t for row and skip_size instead.

See #7517 for a concrete reproducer.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[jinsolp]

can we instead use size_t as the type for row and skip_size (rather than using nnz_t which can in int or size_t)?
Or change our nnz_t dispatch algorithm?

[aamijar]

can we instead use size_t as the type for row and skip_size (rather than using nnz_t which can in int or size_t)? Or change our nnz_t dispatch algorithm?

Yes, I've updated it to use size_t 👍

[jinsolp]

Thanks! LGTM

[aamijar]

/merge