masahi commented on a change in pull request #7458:
URL: https://github.com/apache/tvm/pull/7458#discussion_r577182951
##########
File path: python/tvm/relay/op/strategy/rocm.py
##########
@@ -219,3 +221,93 @@ def batch_matmul_strategy_rocm(attrs, inputs, out_type,
target):
plevel=12,
)
return strategy
+
+
+def can_use_thrust(target, func_name):
+ return (
+ target.kind.name == "rocm"
+ and "thrust" in target.libs
+ and get_global_func(func_name, allow_missing=True)
+ )
Review comment:
Yes, currently thrust dispatch is done in a ad hoc way, unlike
cudnn/cublas etc where `-libs` option is required in addition to `USE_CUDNN=ON`
etc.
@tkonolige and I discussed this issue a bit and we agreed that we should
enforce `-libs=thrust` to enable thrust offload too. For rocm, this requirement
is already enforced in this PR, but for cuda it is not. That's what
`TODO(masahi)` will going to fix in the next PR.
We can discuss what the common helper should be / where it should be in the
next PR, would you agree?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
