rapidsai / raft

RAFT contains fundamental widely-used algorithms and primitives for machine learning and information retrieval. The algorithms are CUDA-accelerated and form building blocks for more easily writing high performance applications.
https://docs.rapids.ai/api/raft/stable/
Apache License 2.0
683 stars 180 forks source link

Add support for nosync thrust exec policy #2293

Closed abc99lr closed 2 months ago

abc99lr commented 2 months ago

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 (https://github.com/rapidsai/raft/issues/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[bot] commented 2 months ago

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.

cjnolet commented 2 months ago

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 commented 2 months ago

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 commented 2 months ago

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 commented 2 months ago

Closing.

cjnolet commented 2 months ago

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 commented 2 months ago

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 commented 2 months ago

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