Add support for nosync thrust exec policy

[abc99lr]

Currently, all the thrust calls used in RAFT are sync calls. There is another RMM execution policy that support async (or nosync) thrust calls. Supporting async calls is important. For example, we need the kmeans predict to be async in order to achieve kernel/copy overlapping in IVF-Flat index build (#2106).

This PR

1. Add support for nosync thrust policy in raft::resource
2. Change the thrust calls in kmeans predict to nosync versions. This would enable us achieve memcpy and kernel overlapping in IVF-Flat index building. Based on discussions, we should use raft::linalg::map instead. Will open another PR for this change.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[tfeher]

Thanks Rui for fixing this, it is indeed important to switch to async policy to improve vector search build time. Overall it looks good!

@achirkin could you also have a look at the changes in resource handling?

[achirkin]

Thanks for the PR! I welcome introduction of the thrust_nosync_policy, it could be useful to us in theory. However, both places you changed in kmean_balanced.cuh could better be done with raft's own linalg::map.

[cjnolet]

However, both places you changed in kmean_balanced.cuh could better be done with raft's own linalg::map

@abc99lr, I just want to give an additional +1 to @achirkin's comment here. In general, we prefer to reuse raft functions where at all possible across the codebase, even when the raft function itself might be a trivial wrapper around thrust, cub, or one of the math libs. The reason for this is that it centralizes these calls so that we can properly instrument, improve, or even fix bugs as they arise in a single place, rather than having to scrape through the codebase and fix them in many places.

The other reason for this is API consistency- over time, raft has improved to become quite a pleasant API experience by being able to compose larger algorithms out of a series of raft functions, which improves code readability.

[abc99lr]

Thanks for the comments. I'll change the thrust calls in kmeans_balanced::predict to raft::linalg::map in a separate PR. For this one, I'll only add the nosync policy and remove the changes in kmeans_balanced::predict. Sound good?

[cjnolet]

For this one, I'll only add the nosync policy and remove the changes in kmeans_balanced::predict. Sound good?

If we aren't going to be using this nosync policy, I'd like to avoid merging changes just for the sake of merging changes.

[abc99lr]

Closing.

[cjnolet]

Most functions in RAFT are assumed to be async, so I suspect we could probably scrape through all of the places we use the thrust_policy and replace them with nosync_policy. @abc99lr I'm not against merging this just for that reason alone, but would you mind creating an issue for the above so that we don't lose sight of it? Just want to avoid this getting stale and remaining unused for the non-foreseeable future.

[abc99lr]

If we aren't going to be using this nosync policy, I'd like to avoid merging changes just for the sake of merging changes.

Understand. But I do think RAFT should provide this functionality to developers and tell people to use nosync policy when could. Otherwise, there would be many unnecessary syncs introduced.

[abc99lr]

Talked with @cjnolet, we think it's worth to try using nosync_policy by default. Going to create another PR for that