UMAP cuda_illegal_memory_access when nnz close to `INT32_MAX`

[aamijar]

While working on the spectral embedding scale up with UMAP integration I came across cuda_illegal_memory_access that is also present in init=random. This is a separate issue that the one that is specific to init=spectral.

It seems that this is happening when the dataset creates a fuzzy_simpl_set graph that has a nnz size that is close to the maximum value supported by int. When nnz requires the usage of int64_t the memory access problem does not happen.

Will update this issue with a reproducer and more details.

Observations

n_neighbors=1-97 works n_neighbors=98-107 fails n_neighbors=108+ works

MRE

import cupy as cp
import numpy as np
from cuml.manifold import UMAP

INT32_MAX = 2147483647

n_neighbors = 60
n_components = 2

n_rows_needed = (INT32_MAX // (2 * n_neighbors)) + 1

n_rows = min(n_rows_needed, 25000000)
n_features = 100

n_rows = 9990000
n_features = 96
n_neighbors = 101

print(f"INT32_MAX: {INT32_MAX}")
print(f"n_rows: {n_rows}")
print(f"n_neighbors: {n_neighbors}")
print(f"2 * n_rows * n_neighbors: {2 * n_rows * n_neighbors}")
print(f"Will use uint64_t: {2 * n_rows * n_neighbors > INT32_MAX}")

X = cp.random.randn(n_rows, n_features, dtype=cp.float32) * 0.1
X = X / cp.linalg.norm(X, axis=1, keepdims=True)

umap = UMAP(
    n_neighbors=n_neighbors,
    n_components=n_components,
    init='random',
    # min_dist=0.1,
    # metric='euclidean',
    # random_state=42,
    verbose=True
)

embedding = umap.fit_transform(X)
print(f"\nDataset shape: {X.shape}")
print(f"Embedding shape: {embedding.shape}")