[GitHub] [tvm] csullivan commented on a change in pull request #7458: [ROCM] Add Thrust support

[GitBox <gi...@apache.org>]

csullivan commented on a change in pull request #7458:
URL: https://github.com/apache/tvm/pull/7458#discussion_r577148555



##########
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:
       Thoughts on perhaps moving this helper to `python/tvm/contrib/rocm.py`? There are a few examples of helpers in python/tvm/contrib that help with op strategy dispatch based on conditions like these.

##########
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:
       Or given a few of your TODOs in this PR, like
   ```
   # TODO(masahi): Check -libs=thrust option
   ```
   maybe we want a common helper for use with both cuda and rocm?




----------------------------------------------------------------
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:
users@infra.apache.org