Function has_nvshmem

python/tvm/testing/env.py:286–298

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.

()

@functools.cache
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 -------------------------------------------