gemini-code-assist[bot] commented on code in PR #19777: URL: https://github.com/apache/tvm/pull/19777#discussion_r3411425583
########## python/tvm/testing/env.py: ########## @@ -0,0 +1,466 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Thin capability probes for test gating. + +This module exposes small ``has_*`` predicates that report whether the +current environment can run a given feature. They are meant to be used +with plain pytest markers and ``skipif``:: + + import pytest + import tvm.testing + + @pytest.mark.cuda + @pytest.mark.skipif(not tvm.testing.env.has_cuda(), reason="need cuda") + def test_my_cuda_kernel(): + ... + +Every probe is memoized with :func:`functools.lru_cache`, so the +underlying device query / ``nvcc`` subprocess runs at most once per +process even though ``skipif`` evaluates the predicate at import time for +every decorated test. Probes never raise: when support is absent they +return ``False`` (or a zero version tuple) rather than propagating an +error out of collection. + +Three kinds of probe live here: + +* **runtime device** probes (``has_cuda``, ``has_gpu`` …) ask whether a + usable device of a given kind is present; +* **build-support** probes (``has_cutlass``, ``has_cudnn`` …) ask whether + an optional library was compiled into the runtime; +* **version / capability** probes (``has_cuda_compute``, + ``has_tensorcore`` …) ask about a finer capability of a present device + or toolchain. +""" + +import functools +import os +import platform + +import tvm + +__all__ = [ + "has_aarch64_sme", + "has_aarch64_sve", + "has_adreno_opencl", + "has_aprofile_aem_fvp", + "has_arm_dot", + "has_arm_fp16", + # cpu features + "has_cpu_feature", + "has_cublas", + # runtime device + "has_cuda", + # version / capability + "has_cuda_compute", + "has_cudagraph", + # build support + "has_cudnn", + "has_cutlass", + "has_gpu", + # toolchain / environment + "has_hexagon", + "has_hexagon_toolchain", + "has_hipblas", + "has_llvm", + "has_llvm_min_version", + "has_matrixcore", + "has_metal", + "has_mrvl", + "has_multi_gpu", + "has_nccl", + "has_nnapi", + "has_nvcc_version", + "has_nvptx", + "has_nvshmem", + "has_opencl", + "has_openclml", + "has_rocm", + "has_rpc", + "has_tensorcore", + "has_vulkan", + "has_x86_amx", + "has_x86_avx512", + "has_x86_vnni", + "is_aarch64", + # host architecture + "is_x86", +] + + [email protected] +def _device_exists(kind: str, index: int = 0) -> bool: + """Return whether ``tvm.device(kind, index)`` is present and usable.""" + try: + return bool(tvm.device(kind, index).exist) + except Exception: # pylint: disable=broad-except + # A missing backend / driver must skip the test, not crash collection. + return False + + [email protected] +def _build_flag_enabled(flag: str) -> bool: + """Return whether an optional build flag (e.g. ``USE_CUTLASS``) is on. + + Mirrors the historical ``Feature`` check: a flag counts as enabled + unless it is explicitly disabled, so library flags carrying a path + still register as present. + """ + try: + value = tvm.support.libinfo().get(flag, "OFF") + return value.lower() not in ("off", "false", "0") + except Exception: # pylint: disable=broad-except + return False + + +# --- runtime device probes ------------------------------------------------- + + +def has_cuda() -> bool: + """True if a CUDA device is present and the runtime supports it.""" + return _device_exists("cuda") + + +def has_rocm() -> bool: + """True if a ROCm device is present and the runtime supports it.""" + return _device_exists("rocm") + + +def has_vulkan() -> bool: + """True if a Vulkan device is present and the runtime supports it.""" + return _device_exists("vulkan") + + +def has_metal() -> bool: + """True if a Metal device is present and the runtime supports it.""" + return _device_exists("metal") + + +def has_opencl() -> bool: + """True if an OpenCL device is present and the runtime supports it.""" + return _device_exists("opencl") + + +def has_nvptx() -> bool: + """True if an NVPTX device is present and the runtime supports it.""" + return _device_exists("nvptx") + + +def has_llvm() -> bool: + """True if the LLVM backend is available.""" + return _device_exists("llvm") + + +def has_gpu() -> bool: + """True if any GPU backend (cuda/rocm/opencl/metal/vulkan) is present.""" + return ( + _device_exists("cuda") + or _device_exists("rocm") + or _device_exists("opencl") + or _device_exists("metal") + or _device_exists("vulkan") + ) + + [email protected] +def has_multi_gpu(count: int = 2) -> bool: + """True if at least ``count`` devices of a single GPU backend exist.""" + for kind in ("cuda", "rocm", "opencl", "metal", "vulkan"): + if all(_device_exists(kind, index) for index in range(count)): + return True + return False + + +# --- build-support probes -------------------------------------------------- +# +# These wrap the optional-library build flags. Features that extend CUDA / +# ROCm additionally require the parent device to be present. + + +def has_cudnn() -> bool: + """True if cuDNN was built in and a CUDA device is present.""" + return has_cuda() and _build_flag_enabled("USE_CUDNN") + + +def has_cublas() -> bool: + """True if cuBLAS was built in and a CUDA device is present.""" + return has_cuda() and _build_flag_enabled("USE_CUBLAS") + + +def has_nccl() -> bool: + """True if NCCL was built in and a CUDA device is present.""" + return has_cuda() and _build_flag_enabled("USE_NCCL") + + +def has_hipblas() -> bool: + """True if hipBLAS was built in and a ROCm device is present.""" + return has_rocm() and _build_flag_enabled("USE_HIPBLAS") + + +def has_cutlass() -> bool: + """True if CUTLASS support was built into the runtime.""" + return _build_flag_enabled("USE_CUTLASS") + + +def has_rpc() -> bool: + """True if RPC support was built into the runtime.""" + return _build_flag_enabled("USE_RPC") + + +def has_nnapi() -> bool: + """True if NNAPI codegen support was built into the runtime.""" + return _build_flag_enabled("USE_NNAPI_CODEGEN") + + +def has_openclml() -> bool: + """True if OpenCLML (CLML) support was built into the runtime.""" + return _build_flag_enabled("USE_CLML") + + +def has_mrvl() -> bool: + """True if the Marvell (MRVL) backend was built into the runtime.""" + return _build_flag_enabled("USE_MRVL") + + [email protected] +def has_nvshmem() -> bool: + """True if the disco NVSHMEM runtime is available (requires CUDA). + + Probes the runtime global function rather than the ``USE_NVSHMEM`` build + flag, since the flag can be set in builds that do not ship the runtime. + """ + try: + return has_cuda() and ( + tvm.get_global_func("runtime.disco.nvshmem.init_nvshmem_uid", allow_missing=True) + is not None + ) + except Exception: # pylint: disable=broad-except + return False + + +# --- version / capability probes ------------------------------------------- + + [email protected] +def _cuda_compute_version() -> tuple: + """Return the (major, minor) CUDA compute version, or (0, 0) if unknown.""" + try: + from tvm.support import nvcc # pylint: disable=import-outside-toplevel + + arch = nvcc.get_target_compute_version() + return nvcc.parse_compute_version(arch) + except Exception: # pylint: disable=broad-except + return (0, 0) + + +def has_cuda_compute(major: int, minor: int = 0, exact: bool = False) -> bool: + """True if the CUDA compute capability satisfies ``(major, minor)``. + + When ``exact`` is False (default) the check is ``compute >= (major, + minor)``; when True it requires an exact match. Returns False when no + CUDA device is present, so it implies :func:`has_cuda`. + """ + compute = _cuda_compute_version() + want = (major, minor) + if exact: + return compute == want + return compute >= want Review Comment:  The docstring states that `has_cuda_compute` should return `False` when no CUDA device is present, implying `has_cuda()`. However, if no CUDA device is present, `_cuda_compute_version()` returns `(0, 0)`. If `has_cuda_compute` is called with `(0, 0)` (or if a check evaluates against `(0, 0)`), it would incorrectly return `True` because `(0, 0) >= (0, 0)` is `True`. To align with the docstring and ensure robustness, we should explicitly check `has_cuda()` at the beginning of the function, similar to how `has_nvcc_version` and other capability probes do. ```suggestion if not has_cuda(): return False compute = _cuda_compute_version() want = (major, minor) if exact: return compute == want return compute >= want ``` ########## python/tvm/testing/env.py: ########## @@ -0,0 +1,466 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Thin capability probes for test gating. + +This module exposes small ``has_*`` predicates that report whether the +current environment can run a given feature. They are meant to be used +with plain pytest markers and ``skipif``:: + + import pytest + import tvm.testing + + @pytest.mark.cuda + @pytest.mark.skipif(not tvm.testing.env.has_cuda(), reason="need cuda") + def test_my_cuda_kernel(): + ... + +Every probe is memoized with :func:`functools.lru_cache`, so the +underlying device query / ``nvcc`` subprocess runs at most once per +process even though ``skipif`` evaluates the predicate at import time for +every decorated test. Probes never raise: when support is absent they +return ``False`` (or a zero version tuple) rather than propagating an +error out of collection. + +Three kinds of probe live here: + +* **runtime device** probes (``has_cuda``, ``has_gpu`` …) ask whether a + usable device of a given kind is present; +* **build-support** probes (``has_cutlass``, ``has_cudnn`` …) ask whether + an optional library was compiled into the runtime; +* **version / capability** probes (``has_cuda_compute``, + ``has_tensorcore`` …) ask about a finer capability of a present device + or toolchain. +""" + +import functools +import os +import platform + +import tvm + +__all__ = [ + "has_aarch64_sme", + "has_aarch64_sve", + "has_adreno_opencl", + "has_aprofile_aem_fvp", + "has_arm_dot", + "has_arm_fp16", + # cpu features + "has_cpu_feature", + "has_cublas", + # runtime device + "has_cuda", + # version / capability + "has_cuda_compute", + "has_cudagraph", + # build support + "has_cudnn", + "has_cutlass", + "has_gpu", + # toolchain / environment + "has_hexagon", + "has_hexagon_toolchain", + "has_hipblas", + "has_llvm", + "has_llvm_min_version", + "has_matrixcore", + "has_metal", + "has_mrvl", + "has_multi_gpu", + "has_nccl", + "has_nnapi", + "has_nvcc_version", + "has_nvptx", + "has_nvshmem", + "has_opencl", + "has_openclml", + "has_rocm", + "has_rpc", + "has_tensorcore", + "has_vulkan", + "has_x86_amx", + "has_x86_avx512", + "has_x86_vnni", + "is_aarch64", + # host architecture + "is_x86", +] + + [email protected] +def _device_exists(kind: str, index: int = 0) -> bool: + """Return whether ``tvm.device(kind, index)`` is present and usable.""" + try: + return bool(tvm.device(kind, index).exist) + except Exception: # pylint: disable=broad-except + # A missing backend / driver must skip the test, not crash collection. + return False + + [email protected] +def _build_flag_enabled(flag: str) -> bool: + """Return whether an optional build flag (e.g. ``USE_CUTLASS``) is on. + + Mirrors the historical ``Feature`` check: a flag counts as enabled + unless it is explicitly disabled, so library flags carrying a path + still register as present. + """ + try: + value = tvm.support.libinfo().get(flag, "OFF") + return value.lower() not in ("off", "false", "0") Review Comment:  To prevent a potential `AttributeError` if `tvm.support.libinfo().get(flag)` returns a non-string value (such as a boolean or integer), it is safer to cast `value` to a string before calling `.lower()`. ```suggestion value = tvm.support.libinfo().get(flag, "OFF") return str(value).lower() not in ("off", "false", "0") ``` -- 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]
