gemini-code-assist[bot] commented on code in PR #19819:
URL: https://github.com/apache/tvm/pull/19819#discussion_r3430963972
##########
python/tvm/testing/plugin.py:
##########
@@ -279,44 +279,39 @@ def sort_key(item):
items.sort(key=sort_key)
-def _gpu_mark_and_skip(has_fn, reason):
- """A GPU-family target: the ``gpu`` selection marker plus an env skip."""
- return [pytest.mark.gpu, pytest.mark.skipif(not has_fn(), reason=reason)]
-
-
-def _skip_only(has_fn, reason):
- """A non-GPU target: an env skip with no selection marker."""
- return [pytest.mark.skipif(not has_fn(), reason=reason)]
+# GPU-family target kinds carry the ``gpu`` selection marker; CPU-family kinds
+# (llvm, hexagon) only skip. The skip condition is the matching tvm.testing.env
+# probe, resolved by name, so there is no per-kind ladder of has_* calls.
+_GPU_TARGET_KINDS = frozenset(
+ {"cuda", "cudnn", "cublas", "rocm", "vulkan", "nvptx", "metal", "opencl"}
+)
+_CPU_TARGET_KINDS = frozenset({"llvm", "hexagon"})
def _target_to_requirement(target):
if isinstance(target, str | dict):
target = tvm.target.Target(target)
- # GPU-family kinds get the `gpu` selection marker; CPU-family kinds only
skip.
+ # A cuda target carrying an accelerator library gates on that library's
probe
+ # (cudnn before cublas) instead of plain cuda.
kind = target.kind.name
- if kind == "cuda" and "cudnn" in target.attrs.get("libs", []):
- return _gpu_mark_and_skip(env.has_cudnn, "need cudnn")
- if kind == "cuda" and "cublas" in target.attrs.get("libs", []):
- return _gpu_mark_and_skip(env.has_cublas, "need cublas")
if kind == "cuda":
- return _gpu_mark_and_skip(env.has_cuda, "need cuda")
- if kind == "rocm":
- return _gpu_mark_and_skip(env.has_rocm, "need rocm")
- if kind == "vulkan":
- return _gpu_mark_and_skip(env.has_vulkan, "need vulkan")
- if kind == "nvptx":
- return _gpu_mark_and_skip(env.has_nvptx, "need nvptx")
- if kind == "metal":
- return _gpu_mark_and_skip(env.has_metal, "need metal")
- if kind == "opencl":
- return _gpu_mark_and_skip(env.has_opencl, "need opencl")
- if kind == "llvm":
- return _skip_only(env.has_llvm, "need llvm")
- if kind == "hexagon":
- return _skip_only(env.has_hexagon, "need hexagon")
-
- return []
+ libs = target.attrs.get("libs", [])
+ if "cudnn" in libs:
+ kind = "cudnn"
+ elif "cublas" in libs:
+ kind = "cublas"
+
+ if kind in _GPU_TARGET_KINDS:
+ is_gpu = True
+ elif kind in _CPU_TARGET_KINDS:
+ is_gpu = False
+ else:
+ return []
+
+ marks = [pytest.mark.gpu] if is_gpu else []
+ marks.append(pytest.mark.skipif(not getattr(env, f"has_{kind}")(),
reason=f"need {kind}"))
Review Comment:
Using `getattr` dynamically on `env` without a fallback or check can lead to
a cryptic `AttributeError` if a developer adds a new target kind to
`_GPU_TARGET_KINDS` or `_CPU_TARGET_KINDS` but forgets to define the
corresponding `has_{kind}` helper in `env.py`. Adding a defensive check with a
clear, actionable error message improves maintainability and developer
experience.
```suggestion
has_fn = getattr(env, f"has_{kind}", None)
if has_fn is None:
raise AttributeError(
f"Target kind '{kind}' is registered in plugin.py but "
f"tvm.testing.env is missing the corresponding 'has_{kind}'
probe function."
)
marks.append(pytest.mark.skipif(not has_fn(), reason=f"need {kind}"))
```
##########
tests/python/testing/test_env.py:
##########
@@ -123,22 +113,6 @@ def test_has_multi_gpu_is_bool():
assert env.has_multi_gpu(1) or not env.has_multi_gpu(2)
[email protected](
- "probe,flag",
- [
- (env.has_cutlass, "USE_CUTLASS"),
- (env.has_rpc, "USE_RPC"),
- (env.has_nnapi, "USE_NNAPI_CODEGEN"),
- (env.has_openclml, "USE_CLML"),
- (env.has_mrvl, "USE_MRVL"),
- ],
- ids=lambda v: getattr(v, "__name__", v),
-)
-def test_build_flag_probe_matches_libinfo(probe, flag):
- """Pure build-flag probes agree with the build-info flag they wrap."""
- assert probe() == env._build_flag_enabled(flag) # pylint:
disable=protected-access
-
-
def test_llvm_min_version_is_monotone():
Review Comment:
The `build_flag_enabled` function has been promoted to a public API in
`tvm.testing.env` (added to `__all__`), but the existing tests that covered it
indirectly via the inlined helpers were removed. We should add a direct unit
test for `build_flag_enabled` to ensure test coverage and prevent regressions.
```suggestion
def test_build_flag_enabled():
assert isinstance(env.build_flag_enabled("USE_LLVM"), bool)
assert not env.build_flag_enabled("NON_EXISTENT_FLAG")
def test_llvm_min_version_is_monotone():
```
--
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.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
