-
Notifications
You must be signed in to change notification settings - Fork 3.5k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[CUDA][THRUST] Enforce -libs=thrust to allow thrust offload #7468
Conversation
if target and ( | ||
can_use_thrust(target, "tvm.contrib.thrust.sort") | ||
or can_use_rocthrust(target, "tvm.contrib.thrust.sort") | ||
): | ||
sort_tensor = argsort_thrust(score_tensor, axis=1, is_ascend=False, dtype="int32") | ||
else: | ||
sort_tensor = argsort(score_tensor, axis=1, is_ascend=False, dtype="int32") |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Just a thought and feel free to ignore: Since THRUST won't be used without -libs=thrust
after this PR, would that be better to add a warning message here if tvm.contrib.thrust.sort
is available (i.e., users built TVM with THRUST enabled) and target doesn't have -libs=thrust
? We can then remove the warning in the next release or so.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Sure I'll do that after I get more feedback
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Warning added in f688f64
Please have a look
If there is no comment, I assume everyone is cool with this change, and I'll ask someone to merge this. |
LGTM, just a small comment: could you also add a comment in the deploy ssd tutorial to let people know this change if they want to use cuda + thrust? |
I have no complaints with this but also very little experience with this part of the code, so I'm happy with it if no one objects. |
607f2d5
to
d0eaca4
Compare
This reverts commit c1629b3.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM. Although the warning is per op, it should be fine as most models have only one or a few thrust supported ops.
thanks @comaniac @Laurawly @mbrookhart |
) * add contrib/thrust.py * update cuda strategy * remove is_thrust_available, update nms, scan, sort and tests * remove unused import * trigger CI * update * add note on how to enable thrust in ssd tutorial * add warning * Revert "update" This reverts commit c1629b3. Co-authored-by: masa <masa@pop-os.localdomain>
) * add contrib/thrust.py * update cuda strategy * remove is_thrust_available, update nms, scan, sort and tests * remove unused import * trigger CI * update * add note on how to enable thrust in ssd tutorial * add warning * Revert "update" This reverts commit c1629b3. Co-authored-by: masa <masa@pop-os.localdomain>
Currently, thrust offload is done in a rather ad hoc way, unlike other libs such as cudnn which requires using
-libs=cudnn
in addition to-DUSE_CUDNN=ON
.This PR enforces a requirement
-libs=thrust
to enable thrust offload, to make it consistent with other libs. It also makes it easier to compare performance of thrust and TIR sort. @mbrookhart is working on a new TIR sort implementation, based on mergepath, that is competitive or even faster than thrust, so hopefully we don't have to depend on thrust in the near future.NOTE: If you want to use thrust, you must update your code to add
-libs=thrust
after this PR.@tkonolige @mbrookhart @csullivan @jwfromm @zhiics @kevinthesun @anijain2305 @trevor-m @Laurawly