[FEA] Possible improvements for sparse KNN

[viclafargue]

Possible improvements for sparse KNN, some of which might be necessary to implement MNMG sparse KNN :

• Provide the possibility to specify different types for indices of sparse matrix inputs, and indices of dense nearest neighbors output
• Allow multiple index inputs to be provided and the KNN results to be merged in similar fashion to what is possible in the dense case
• Fix issue for large batch sizes

[cjnolet]

Hey @viclafargue these improvements look useful. Are you also planning to implement them as part of the sparse knn update work?

[viclafargue]

Hey @cjnolet, I've already been working on making sparse KNN compile with int64_t, and also looked into ways to have different types for inputs and outputs indices. However, sparse KNN is quite complex and I feel like RAFT/KNN experts would be better suited to work at least on the two last points mentioned above. That said, if no one is available I can work on these.

[cjnolet]

@viclafargue yes it would be very helpful if you have cycles to take these on. As you probably know, we are are short in resources across the board, but when issues like this arise in cuml it's somewhat expected that the corresponding raft pieces will be updated as well.

[cjnolet]

@viclafargue, I'm reivewing your PR and I'm trying to think through why this is necessary:

Provide the possibility to specify different types for indices of sparse matrix inputs, and indices of dense nearest neighbors output

Can you explain the need for this?

[viclafargue]

Sorry @cjnolet, I missed your question. This is actually something optional, I was thinking that we could systematically return indices as int64 as with the dense algorithm. But, when I think of it, compiling more templates is probably not the best option. In addition, input indices could be int32 or int64, and it's getting complicated to create comms for MNMG KNN to send both datatypes. We should probably just systematically convert input indices to int64 in the MNGM KNN Python layer and only allow int64 inputs from the MNGM KNN C++ API.

[viclafargue]

Closing this issue as it does not seem to be relevant anymore.