Skip to content

Commit

Permalink
[UnitTests] Correct parametrization of cudnn target.
Browse files Browse the repository at this point in the history
Previous checks for enabled runtimes were based only on the target
kind.  CuDNN is the same target kind as "cuda", and therefore needs
special handling.
  • Loading branch information
Lunderberg committed Aug 10, 2021
1 parent f744b1f commit a9cc505
Showing 1 changed file with 8 additions and 2 deletions.
10 changes: 8 additions & 2 deletions python/tvm/testing.py
Original file line number Diff line number Diff line change
Expand Up @@ -388,8 +388,14 @@ def _get_targets(target_str=None):
targets = []
for target in target_names:
target_kind = target.split()[0]
is_enabled = tvm.runtime.enabled(target_kind)
is_runnable = is_enabled and tvm.device(target_kind).exist

if target_kind == "cuda" and "cudnn" in tvm.target.Target(target).attrs.get("libs", []):
is_enabled = tvm.support.libinfo()["USE_CUDNN"].lower() in ["on", "true", "1"]
is_runnable = is_enabled and cudnn.exists()
else:
is_enabled = tvm.runtime.enabled(target_kind)
is_runnable = is_enabled and tvm.device(target_kind).exist

targets.append(
{
"target": target,
Expand Down

0 comments on commit a9cc505

Please sign in to comment.