UMAP spectral initialization with `int64_t` nnz types.

[aamijar]

UMAP with large datasets dispatches the COO graph with uint64_t nnz types. This needs to be properly handled in spectral initialization. We are aiming to resolve #6910 with this.

There are several changes and considerations to be made when working on this.

1. Exposing int64_t types in cuvs spectral embedding algorithm. Currently there is only support for int.
2. uint64_t in spectral embedding is not possible because there are downstream primitives that do not work with it (I forget the exact ones, but they were blockers when I was working on this.)
3. There are several bugs in the raft kernel primitives for int64_t that need to be resolved.
4. UMAP is handing off row/col indices in int always. This could be a problem since when converting to CSR the row offsets needs to be int64.
5. However, allocating space for int64 and copying again, is problematic because it causes OOM.
6. What is the best way to efficiently handle the int64 row/col support?
7. It would be great is row/col were already in int64 during _get_graph in UMAP. Is it possible to do this? I tried doing this however, it seems like an extensive change.

Update:
Now that I think about it, allocating int64 on the fly for the spectral embedding coo_to_csr step will have the same memory overhead as using row/col int64 in the COO creation upfront.

Scenario 1. COO in umap has row and col with int.
Then during coo_to_csr the output row_offsets must be int64. The col_indices could be int, but that will not work with cusparse spmv primitives in the lanczos solver. Therefore we need to allocate space for a temporary col_indices int64 array.
Scenario 2. COO in umap has row and col with int64
Again during coo_to_csr the output row_offsets must be int64. Now we don't need a temp col_indices since its already in int64.

In both scenarios we create the same amount of extra memory overhead. Either we have the temp col_indices of size nnz * int64. Or we need to increase the size of row/col to int64 upfront. This increase would be from int so it would be a total increase of (nnz * int * 2) + (nnz * int * 2). This is the same increase as nnz * int64.

Therefore, scenario 1 is preferrable, since it doesn't require extensive changes to the COO umap types and needs temporary memory increase.

[cjnolet]

@aamijar this sounds like a workaround, and a dangerous one because it appears to be putting a bandaid on several issues with the primitives instead of fixing those issues. Is that correct? Regardless of the memory footprint changes, if the primitives aren’t properly handing a type, we should be addressing that directly instead of working around it.

[aamijar]

As discussed offline with @cjnolet, here are updates on some of the concerns:

uint64_t in spectral embedding is not possible because there are downstream primitives that do not work with it (I forget the exact ones, but they were blockers when I was working on this.)

It should be reasonable to reinterpret nnz uint64_t to int64_t. The range of int64_t is more than enough for the datasets we are working with (up to 2^63 nnz; much more than anything we need as we will likely be in the range of billions of nnz as a upper limit).

There are several bugs in the raft kernel primitives for int64_t that need to be resolved.

This was already addressed and merged in rapidsai/raft#2793

It would be great is row/col were already in int64 during _get_graph in UMAP. Is it possible to do this? I tried doing this however, it seems like an extensive change.

As in scenario 1, keeping COO with row/col as int is fine. During csr conversion we should try and keep the col indices in int and row offsets in int64. In order to do this we need to reinvestigate this limitation of cusparse primitives used in the lanczos solver. Need to revisit to see what the exact issue was, it seems that cusparse should support mixed types.

[aamijar]

Here is the exact issue that I ran into when trying to cusparse csr with mixed types after adding the appropriate raft cusparse wrappers:

** On entry to cusparseCreateCsr(): different index sizes for 'csrRowOffsetsType' (64I) and 'csrColIndType' (32I) is not currently supported

A possible path forward is to see whether we can drop CSR altogether. Let's see if it is possible to use COO for everything in spectral embedding path.

[aamijar]

I've implemented COO support throughout the spectral embedding algorithm and confirmed that it works. Need to do a chunk of refactoring now and connect it with UMAP.