This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 3edbc658e2 [TIRx][CUDA] Remove IKET dependency version restrictions
(#20477)
3edbc658e2 is described below
commit 3edbc658e2a74a47b90b47a6cc062ab5cfa4caec
Author: Shushi Hong <[email protected]>
AuthorDate: Sun Sep 27 18:21:09 2026 -0400
[TIRx][CUDA] Remove IKET dependency version restrictions (#20477)
IKET currently rejects otherwise working CUDA installations unless
NVIDIA wheel metadata and the NVRTC runtime match hard-coded versions.
For example, an installed `nvidia-cuda-nvdisasm==13.4.92` is rejected in
favor of `13.3.73`, and a
system NVRTC library is rejected when the corresponding Python wheel is
absent.
This pr removes the package version checks and NVRTC version equality
check. Continue checking the `run-iket` executable and entry point,
NVRTC availability, and the tracker/capture injection configuration.
Preserve the existing profile name for callers that select it
explicitly.
---
python/tvm/backend/cuda/iket.py | 110 +++----------------
src/backend/cuda/transforms/lower_iket.cc | 4 +-
tests/python/tirx/iket/iket_profile_workload.py | 4 +-
tests/python/tirx/iket/test_iket_orchestration.py | 51 ++++-----
tests/python/tirx/iket/test_iket_profiler.py | 122 ++++++++++++----------
5 files changed, 109 insertions(+), 182 deletions(-)
diff --git a/python/tvm/backend/cuda/iket.py b/python/tvm/backend/cuda/iket.py
index 62d65c4cce..5f492b1ed4 100644
--- a/python/tvm/backend/cuda/iket.py
+++ b/python/tvm/backend/cuda/iket.py
@@ -38,7 +38,6 @@ import tempfile
import threading
import time
import uuid
-import warnings
from collections import deque
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
@@ -50,30 +49,12 @@ from typing import Any
import tvm
from tvm.script import tirx as T
-_PROFILE_ENV = "TVM_IKET_OFFICIAL_PROFILE"
_INJECTED_CHILD_ENABLE_ENV = "TVM_IKET_INJECTED_CHILD_ENABLE"
-_DEFAULT_PROFILE = "cutlass-4.6.0"
_POSTPROCESS_CHOICES = frozenset(("perfetto", "json", "html", "none", "all"))
_INJECTION_ENV_VARS = ("CUDA_INJECTION64_PATH", "SMODEL_INJECTION_CONFIG")
_OUTPUT_TAIL_LINES = 100
_TERMINATION_GRACE_SECONDS = 5.0
-_OFFICIAL_PROFILES = {
- "cutlass-4.6.0": {
- "nvrtc_version": (13, 2),
- "minimum_versions": {
- "nvidia-cutlass-dsl": "4.6.0",
- "nvidia-cutlass-dsl-libs-base": "4.6.0",
- "nvidia-cutlass-dsl-libs-core": "4.6.0",
- "nvidia-cutlass-dsl-libs-cu13": "4.6.0",
- },
- "exact_versions": {
- "nvidia-cuda-nvdisasm": "13.3.73",
- "nvidia-cuda-nvrtc": "13.2.78",
- },
- }
-}
-
class IketProfileError(RuntimeError):
"""An error while validating or running an official IKET profile."""
@@ -171,24 +152,14 @@ def _profile_error(message: str) -> IketProfileError:
)
-def _is_newer_release(actual: str, expected: str) -> bool:
- """Compare the numeric release versions published by NVIDIA wheels."""
- try:
- actual_parts = tuple(int(part) for part in actual.split("."))
- expected_parts = tuple(int(part) for part in expected.split("."))
- except ValueError:
- return False
- width = max(len(actual_parts), len(expected_parts))
- return actual_parts + (0,) * (width - len(actual_parts)) > expected_parts
+ (0,) * (
- width - len(expected_parts)
- )
-
-
def _validate_run_iket_entrypoint() -> str:
executable = shutil.which("run-iket")
if executable is None or not os.access(executable, os.X_OK):
raise _profile_error("the run-iket executable is unavailable")
- entry_points =
metadata.distribution("nvidia-cutlass-dsl-libs-base").entry_points
+ try:
+ entry_points =
metadata.distribution("nvidia-cutlass-dsl-libs-base").entry_points
+ except metadata.PackageNotFoundError as err:
+ raise _profile_error("the run-iket entry point metadata is
unavailable") from err
if not any(
item.group == "console_scripts"
and item.name == "run-iket"
@@ -201,49 +172,21 @@ def _validate_run_iket_entrypoint() -> str:
return executable
-def _validate_nvrtc_version(expected_version: tuple[int, int]) -> None:
- expected_label = ".".join(str(part) for part in expected_version)
+def _validate_nvrtc_available() -> None:
try:
from cuda.bindings import nvrtc
- error, major, minor = nvrtc.nvrtcVersion()
+ error, _, _ = nvrtc.nvrtcVersion()
except (ImportError, OSError, RuntimeError) as err:
- raise _profile_error(f"CUDA NVRTC {expected_label} is unavailable")
from err
- actual_version = (int(major), int(minor))
- if int(error) != 0 or actual_version != expected_version:
- raise _profile_error(
- f"CUDA NVRTC {expected_label} is required, got
{actual_version[0]}.{actual_version[1]}"
- )
+ raise _profile_error("CUDA NVRTC is unavailable") from err
+ if int(error) != 0:
+ raise _profile_error(f"CUDA NVRTC is unavailable (error {int(error)})")
-def _validate_official_installation(profile_name: str) -> str:
- """Validate host-side package versions and return the official
executable."""
- if profile_name not in _OFFICIAL_PROFILES:
- raise _profile_error(f"unsupported profile {profile_name!r}; expected
{_DEFAULT_PROFILE!r}")
- profile_config = _OFFICIAL_PROFILES[profile_name]
- version_groups = (
- (profile_config["minimum_versions"], True),
- (profile_config["exact_versions"], False),
- )
- for versions, allow_newer in version_groups:
- for distribution_name, expected_version in versions.items():
- try:
- distribution = metadata.distribution(distribution_name)
- except metadata.PackageNotFoundError as err:
- operator = ">=" if allow_newer else "=="
- raise _profile_error(
- f"{distribution_name}{operator}{expected_version} is not
installed"
- ) from err
- if distribution.version != expected_version and not (
- allow_newer and _is_newer_release(distribution.version,
expected_version)
- ):
- requirement = f"{expected_version} or newer" if allow_newer
else expected_version
- raise _profile_error(
- f"{distribution_name} must be {requirement}, got
{distribution.version}"
- )
-
+def _validate_official_installation() -> str:
+ """Check the required tools are available and return the official
executable."""
executable = _validate_run_iket_entrypoint()
- _validate_nvrtc_version(profile_config["nvrtc_version"])
+ _validate_nvrtc_available()
return executable
@@ -278,12 +221,7 @@ def _validate_injection_environment() -> None:
def _validate_official_environment() -> str:
"""Validate the installation plus tracker/capture injection environment."""
- profile_name = os.environ.get(_PROFILE_ENV)
- if profile_name not in _OFFICIAL_PROFILES:
- raise _profile_error(
- f"{_PROFILE_ENV} must be set to {_DEFAULT_PROFILE}, got
{profile_name!r}"
- )
- executable = _validate_official_installation(profile_name)
+ executable = _validate_official_installation()
_validate_injection_environment()
return executable
@@ -409,11 +347,7 @@ def _validate_profile_options(
return timeout
-def _child_environment(
- profile_name: str, env: Mapping[str, str | os.PathLike[str]] | None
-) -> dict[str, str]:
- if not isinstance(profile_name, str):
- raise TypeError("profile_name must be a string")
+def _child_environment(env: Mapping[str, str | os.PathLike[str]] | None) ->
dict[str, str]:
child_env = os.environ.copy()
if env is not None:
if not isinstance(env, Mapping):
@@ -424,15 +358,6 @@ def _child_environment(
if not isinstance(normalized_key, str) or not
isinstance(normalized_value, str):
raise TypeError("env keys and values must resolve to strings")
child_env[normalized_key] = normalized_value
- inherited_profile = child_env.get(_PROFILE_ENV)
- if inherited_profile is not None and inherited_profile != profile_name:
- warnings.warn(
- f"Ignoring inherited {_PROFILE_ENV}={inherited_profile!r}; "
- f"profile_name={profile_name!r} takes precedence",
- RuntimeWarning,
- stacklevel=3,
- )
- child_env[_PROFILE_ENV] = profile_name
# LowerIket also requires the two run-iket injection variables before it
# honors this marker. This enables ordinary TIRx JIT compilation only in
# children started by this validated profiling entry point.
@@ -696,7 +621,6 @@ def profile(
command: Sequence[str | os.PathLike[str]],
*,
output_dir: str | os.PathLike[str],
- profile_name: str = _DEFAULT_PROFILE,
postprocess: str = "all",
clobber: bool = False,
cwd: str | os.PathLike[str] | None = None,
@@ -718,8 +642,8 @@ def profile(
if _path_exists(target) and not clobber:
raise FileExistsError(f"IKET output directory already exists:
{target}")
- child_env = _child_environment(profile_name, env)
- executable = _validate_official_installation(profile_name)
+ child_env = _child_environment(env)
+ executable = _validate_official_installation()
target.parent.mkdir(parents=True, exist_ok=True)
staging = Path(tempfile.mkdtemp(prefix=f".{target.name}.staging-",
dir=target.parent))
published = False
@@ -863,7 +787,6 @@ def run(
main,
*,
output_dir: str | os.PathLike[str],
- profile_name: str = _DEFAULT_PROFILE,
postprocess: str = "all",
clobber: bool = False,
cwd: str | os.PathLike[str] | None = None,
@@ -889,7 +812,6 @@ def run(
return profile(
command,
output_dir=output_dir,
- profile_name=profile_name,
postprocess=postprocess,
clobber=clobber,
cwd=cwd,
diff --git a/src/backend/cuda/transforms/lower_iket.cc
b/src/backend/cuda/transforms/lower_iket.cc
index 79d7c32b94..f619384dcc 100644
--- a/src/backend/cuda/transforms/lower_iket.cc
+++ b/src/backend/cuda/transforms/lower_iket.cc
@@ -1205,11 +1205,9 @@ class InstrumentOfficialKernel : public StmtExprMutator {
bool IketEnabled(const IRModule& module) {
if (module->HasNonzeroAttr("tirx.iket.enabled")) return true;
const char* child_enable = std::getenv("TVM_IKET_INJECTED_CHILD_ENABLE");
- const char* profile = std::getenv("TVM_IKET_OFFICIAL_PROFILE");
const char* injection = std::getenv("CUDA_INJECTION64_PATH");
const char* injection_config = std::getenv("SMODEL_INJECTION_CONFIG");
- return child_enable && std::string(child_enable) == "1" && profile &&
- std::string(profile) == "cutlass-4.6.0" && injection && injection[0]
!= '\0' &&
+ return child_enable && std::string(child_enable) == "1" && injection &&
injection[0] != '\0' &&
injection_config && injection_config[0] != '\0';
}
diff --git a/tests/python/tirx/iket/iket_profile_workload.py
b/tests/python/tirx/iket/iket_profile_workload.py
index 87c81dd24b..5cdc6a9a38 100644
--- a/tests/python/tirx/iket/iket_profile_workload.py
+++ b/tests/python/tirx/iket/iket_profile_workload.py
@@ -39,7 +39,7 @@ def canonical_iket_workload(out: T.Buffer((32,), "int32")):
T.device_entry()
profiler = iket.IketProfiler()
tx = T.thread_id([32])
- token = profiler.sentinel_token("token")
+ token: T.uint32 = profiler.sentinel_token("token")
profiler.range_end(token)
token = profiler.range_start("token")
profiler.mark("checkpoint")
@@ -64,7 +64,7 @@ def native_payload_workload(out: T.Buffer((32,), "int32")):
profiler.mark("bool_false_payload", tx != 0)
profiler.mark("float32_payload", T.float32(-3.25))
profiler.mark("float64_payload", T.float64(6.5))
- token = profiler.range_start("token_payload", tx + 200)
+ token: T.uint32 = profiler.range_start("token_payload", tx + 200)
profiler.range_end(token, tx + 300)
profiler.range_push("stack_payload", tx + 400)
profiler.range_pop()
diff --git a/tests/python/tirx/iket/test_iket_orchestration.py
b/tests/python/tirx/iket/test_iket_orchestration.py
index ff59170384..86d74af50b 100644
--- a/tests/python/tirx/iket/test_iket_orchestration.py
+++ b/tests/python/tirx/iket/test_iket_orchestration.py
@@ -73,7 +73,7 @@ def test_profile_rejects_nested_injection_before_validation(
):
validation_started = False
- def validate(_name):
+ def validate():
nonlocal validation_started
validation_started = True
@@ -122,25 +122,21 @@ def
test_profile_forwards_cwd_environment_timeout_and_publishes(tmp_path, monkey
cwd = tmp_path / "work"
cwd.mkdir()
captured = {}
- monkeypatch.setattr(
- iket, "_validate_official_installation", lambda profile_name:
"/opt/bin/run-iket"
- )
+ monkeypatch.setattr(iket, "_validate_official_installation", lambda:
"/opt/bin/run-iket")
monkeypatch.setattr(iket, "_run_process", _fake_profile_process(captured))
- monkeypatch.setenv("TVM_IKET_OFFICIAL_PROFILE", "inherited-profile")
+ monkeypatch.setenv("IKET_TEST_ENV", "inherited")
- with pytest.warns(RuntimeWarning, match="takes precedence"):
- result = iket.profile(
- (sys.executable, "workload.py"),
- output_dir=target,
- profile_name="cutlass-4.6.0",
- postprocess="all",
- clobber=False,
- cwd=cwd,
- env={"IKET_TEST_ENV": "present"},
- max_ts_cnt_per_warp=17,
- keep=True,
- timeout=12.5,
- )
+ result = iket.profile(
+ (sys.executable, "workload.py"),
+ output_dir=target,
+ postprocess="all",
+ clobber=False,
+ cwd=cwd,
+ env={"IKET_TEST_ENV": "present"},
+ max_ts_cnt_per_warp=17,
+ keep=True,
+ timeout=12.5,
+ )
staging = Path(captured["argv"][2])
assert staging.parent == target.parent
@@ -148,9 +144,8 @@ def
test_profile_forwards_cwd_environment_timeout_and_publishes(tmp_path, monkey
assert captured["cwd"] == cwd
assert captured["timeout"] == 12.5
assert captured["env"]["IKET_TEST_ENV"] == "present"
- assert captured["env"]["TVM_IKET_OFFICIAL_PROFILE"] == "cutlass-4.6.0"
assert captured["env"]["TVM_IKET_INJECTED_CHILD_ENABLE"] == "1"
- assert os.environ["TVM_IKET_OFFICIAL_PROFILE"] == "inherited-profile"
+ assert os.environ["IKET_TEST_ENV"] == "inherited"
assert result.output_dir == target
assert result.command == (sys.executable, "workload.py")
assert result.trace == {"launches": [{"index": 0}]}
@@ -164,7 +159,7 @@ def
test_profile_forwards_cwd_environment_timeout_and_publishes(tmp_path, monkey
def test_postprocess_artifact_contract(postprocess, tmp_path, monkeypatch):
target = tmp_path / postprocess
captured = {}
- monkeypatch.setattr(iket, "_validate_official_installation", lambda _name:
"run-iket")
+ monkeypatch.setattr(iket, "_validate_official_installation", lambda:
"run-iket")
monkeypatch.setattr(iket, "_run_process", _fake_profile_process(captured))
result = iket.profile(
@@ -190,7 +185,7 @@ def test_postprocess_artifact_contract(postprocess,
tmp_path, monkeypatch):
def test_multi_process_singular_properties_point_to_plural(tmp_path,
monkeypatch):
target = tmp_path / "multi"
captured = {}
- monkeypatch.setattr(iket, "_validate_official_installation", lambda _name:
"run-iket")
+ monkeypatch.setattr(iket, "_validate_official_installation", lambda:
"run-iket")
monkeypatch.setattr(iket, "_run_process", _fake_profile_process(captured,
artifact_count=2))
result = iket.profile(("python", "workload.py"), output_dir=target)
@@ -215,7 +210,7 @@ def
test_missing_artifact_rolls_back_existing_output(tmp_path, monkeypatch):
staging = Path(argv[argv.index("--output-dir") + 1])
(staging / "iket_pid_1.trace.json").write_text("{}", encoding="utf-8")
- monkeypatch.setattr(iket, "_validate_official_installation", lambda _name:
"run-iket")
+ monkeypatch.setattr(iket, "_validate_official_installation", lambda:
"run-iket")
monkeypatch.setattr(iket, "_run_process", incomplete)
with pytest.raises(iket.IketProfileError, match="Perfetto trace"):
iket.profile(("python", "workload.py"), output_dir=target,
clobber=True)
@@ -229,7 +224,7 @@ def
test_successful_clobber_replaces_only_after_profile_success(tmp_path, monkey
target.mkdir()
(target / "old.txt").write_text("old", encoding="utf-8")
captured = {}
- monkeypatch.setattr(iket, "_validate_official_installation", lambda _name:
"run-iket")
+ monkeypatch.setattr(iket, "_validate_official_installation", lambda:
"run-iket")
monkeypatch.setattr(iket, "_run_process", _fake_profile_process(captured))
result = iket.profile(("python", "workload.py"), output_dir=target,
clobber=True)
@@ -243,7 +238,7 @@ def
test_existing_output_fails_before_installation_validation(tmp_path, monkeypa
target.mkdir()
called = False
- def validate(_name):
+ def validate():
nonlocal called
called = True
@@ -257,7 +252,7 @@ def
test_existing_output_fails_before_installation_validation(tmp_path, monkeypa
def test_invalid_timeout_is_rejected_before_validation(timeout, tmp_path,
monkeypatch):
called = False
- def validate(_name):
+ def validate():
nonlocal called
called = True
@@ -292,7 +287,7 @@ def test_nonzero_exit_keeps_last_100_output_lines(tmp_path,
monkeypatch):
target.mkdir()
sentinel = target / "old.txt"
sentinel.write_text("old", encoding="utf-8")
- monkeypatch.setattr(iket, "_validate_official_installation", lambda _name:
str(executable))
+ monkeypatch.setattr(iket, "_validate_official_installation", lambda:
str(executable))
with pytest.raises(iket.IketProfileError) as error_info:
iket.profile(
@@ -371,7 +366,7 @@ def
test_timeout_kills_run_iket_workload_and_grandchild(tmp_path, monkeypatch):
time.sleep(3600)
""",
)
- monkeypatch.setattr(iket, "_validate_official_installation", lambda _name:
str(executable))
+ monkeypatch.setattr(iket, "_validate_official_installation", lambda:
str(executable))
monkeypatch.setattr(iket, "_TERMINATION_GRACE_SECONDS", 0.2)
with pytest.raises(iket.IketProfileError) as error_info:
diff --git a/tests/python/tirx/iket/test_iket_profiler.py
b/tests/python/tirx/iket/test_iket_profiler.py
index 83906bbe91..f7aa5e719f 100644
--- a/tests/python/tirx/iket/test_iket_profiler.py
+++ b/tests/python/tirx/iket/test_iket_profiler.py
@@ -14,7 +14,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-"""Tests for IKET lowering, metadata, installation versions, and trace
contracts."""
+"""Tests for IKET lowering, metadata, installation checks, and trace
contracts."""
import hashlib
import importlib
@@ -27,6 +27,7 @@ import subprocess
import sys
from importlib import metadata
from pathlib import Path
+from types import SimpleNamespace
import pytest
@@ -434,7 +435,7 @@ def test_regular_lowering_strips_annotations_and_tokens():
def test_verified_injected_child_automatically_enables_plain_jit(monkeypatch):
monkeypatch.setenv("TVM_IKET_INJECTED_CHILD_ENABLE", "1")
- monkeypatch.setenv("TVM_IKET_OFFICIAL_PROFILE", "cutlass-4.6.0")
+ monkeypatch.delenv("TVM_IKET_OFFICIAL_PROFILE", raising=False)
monkeypatch.setenv("CUDA_INJECTION64_PATH",
"/verified/libsmodel_injection.so")
monkeypatch.setenv("SMODEL_INJECTION_CONFIG", "/verified/config.json")
executable = tvm.compile(tvm.IRModule({"main": serial_a}), target=TARGET,
tir_pipeline="tirx")
@@ -449,7 +450,6 @@ def
test_verified_injected_child_automatically_enables_plain_jit(monkeypatch):
def test_injected_child_auto_enable_remains_fail_closed(monkeypatch,
missing_env):
values = {
"TVM_IKET_INJECTED_CHILD_ENABLE": "1",
- "TVM_IKET_OFFICIAL_PROFILE": "cutlass-4.6.0",
"CUDA_INJECTION64_PATH": "/verified/libsmodel_injection.so",
"SMODEL_INJECTION_CONFIG": "/verified/config.json",
}
@@ -779,77 +779,89 @@ def test_multi_kernel_module_has_no_tvm_control_plane():
def test_proxy_fails_closed_and_forbids_export(tmp_path, monkeypatch):
+ from tvm.tirx.cuda import iket as _iket_official
+
executable = _compile(serial_a)
with pytest.raises(RuntimeError, match="cannot be exported"):
executable.export_library(tmp_path / "official.so")
- monkeypatch.delenv("TVM_IKET_OFFICIAL_PROFILE", raising=False)
- with pytest.raises(RuntimeError, match="TVM_IKET_OFFICIAL_PROFILE must be
set"):
+ monkeypatch.setattr(_iket_official, "_validate_official_installation",
lambda: None)
+ monkeypatch.delenv("CUDA_INJECTION64_PATH", raising=False)
+ with pytest.raises(RuntimeError, match="CUDA_INJECTION64_PATH was not
supplied"):
executable.jit()
assert executable._executable._jitted_mod is None # pylint:
disable=protected-access
-def test_environment_validation_is_not_process_cached(monkeypatch):
+def test_environment_validation_is_not_process_cached(tmp_path, monkeypatch):
from tvm.tirx.cuda import iket as _iket_official
- profile = {
- "nvrtc_version": (13, 2),
- "minimum_versions": {"nvidia-cutlass-dsl-libs-base": "4.6.0"},
- "exact_versions": {},
- }
-
- class FakeDistribution:
- version = "4.6.0"
-
- monkeypatch.setitem(_iket_official._OFFICIAL_PROFILES, "cutlass-4.6.0",
profile)
- monkeypatch.setattr(_iket_official.metadata, "distribution", lambda _name:
FakeDistribution())
monkeypatch.setattr(_iket_official, "_validate_run_iket_entrypoint",
lambda: None)
- monkeypatch.setattr(_iket_official, "_validate_injection_environment",
lambda: None)
- monkeypatch.setattr(_iket_official, "_validate_nvrtc_version", lambda
_version: None)
- monkeypatch.setenv("TVM_IKET_OFFICIAL_PROFILE", "cutlass-4.6.0")
+ monkeypatch.setattr(_iket_official, "_validate_nvrtc_available", lambda:
None)
+ injection = tmp_path / "libsmodel_injection.so"
+ injection.write_bytes(b"injection")
+ config = tmp_path / "config.json"
+ config.write_text(json.dumps({"toolName": "tracker"}))
+ monkeypatch.setenv("CUDA_INJECTION64_PATH", str(injection))
+ monkeypatch.setenv("SMODEL_INJECTION_CONFIG", str(config))
+ monkeypatch.delenv("TVM_IKET_OFFICIAL_PROFILE", raising=False)
_iket_official.validate_official_environment()
- monkeypatch.delenv("TVM_IKET_OFFICIAL_PROFILE")
- with pytest.raises(RuntimeError, match="TVM_IKET_OFFICIAL_PROFILE must be
set"):
+ monkeypatch.delenv("CUDA_INJECTION64_PATH")
+ with pytest.raises(RuntimeError, match="CUDA_INJECTION64_PATH was not
supplied"):
_iket_official.validate_official_environment()
[email protected](
- ("version", "expected_error"),
- (
- ("4.5.0", "must be 4.6.0 or newer"),
- ("4.6.0", None),
- ("4.6.2", None),
- ),
-)
-def test_official_installation_accepts_newer_cutlass(monkeypatch, version,
expected_error):
[email protected]("version", ("4.5.0", "4.6.0", "4.6.2", "5.0.0.dev1"))
+def test_installation_ignores_package_versions(tmp_path, monkeypatch, version):
from tvm.tirx.cuda import iket as _iket_official
- profile = {
- "nvrtc_version": (13, 2),
- "minimum_versions": {"nvidia-cutlass-dsl-libs-base": "4.6.0"},
- "exact_versions": {},
- }
+ executable = tmp_path / "run-iket"
+ executable.write_text("#!/bin/sh\n")
+ executable.chmod(0o755)
+ entry_point = metadata.EntryPoint(
+ name="run-iket", value="iket.cli.main:entrypoint",
group="console_scripts"
+ )
- class FakeDistribution:
- pass
+ def distribution(name):
+ if name == "nvidia-cutlass-dsl-libs-base":
+ return SimpleNamespace(version=version,
entry_points=(entry_point,))
+ raise metadata.PackageNotFoundError(name)
- distribution = FakeDistribution()
- distribution.version = version
- monkeypatch.setitem(_iket_official._OFFICIAL_PROFILES, "cutlass-4.6.0",
profile)
- monkeypatch.setattr(_iket_official.metadata, "distribution", lambda _name:
distribution)
- monkeypatch.setattr(_iket_official, "_validate_run_iket_entrypoint",
lambda: None)
- monkeypatch.setattr(_iket_official, "_validate_nvrtc_version", lambda
_version: None)
-
- if expected_error:
- with pytest.raises(RuntimeError, match=expected_error):
- _iket_official._validate_official_installation( # pylint:
disable=protected-access
- "cutlass-4.6.0"
- )
- else:
- _iket_official._validate_official_installation( # pylint:
disable=protected-access
- "cutlass-4.6.0"
- )
+ monkeypatch.setattr(_iket_official.metadata, "distribution", distribution)
+ monkeypatch.setattr(_iket_official.shutil, "which", lambda _name:
str(executable))
+ monkeypatch.setattr(_iket_official, "_validate_nvrtc_available", lambda:
None)
+ assert _iket_official._validate_official_installation() == str(executable)
# pylint: disable=protected-access
+
+
[email protected]("version", ((12, 9), (13, 2), (13, 4), (14, 0)))
+def test_nvrtc_version_unrestricted(monkeypatch, version):
+ from tvm.tirx.cuda import iket as _iket_official
+
+ nvrtc = SimpleNamespace(nvrtcVersion=lambda: (0, *version))
+ monkeypatch.setitem(sys.modules, "cuda.bindings",
SimpleNamespace(nvrtc=nvrtc))
+ _iket_official._validate_nvrtc_available() # pylint:
disable=protected-access
+
+
+def test_nvrtc_error_status(monkeypatch):
+ from tvm.tirx.cuda import iket as _iket_official
+
+ nvrtc = SimpleNamespace(nvrtcVersion=lambda: (1, 0, 0))
+ monkeypatch.setitem(sys.modules, "cuda.bindings",
SimpleNamespace(nvrtc=nvrtc))
+ with pytest.raises(_iket_official.IketProfileError, match="CUDA NVRTC is
unavailable"):
+ _iket_official._validate_nvrtc_available() # pylint:
disable=protected-access
+
+
[email protected]("error_type", (ImportError, OSError, RuntimeError))
+def test_nvrtc_load_error(monkeypatch, error_type):
+ from tvm.tirx.cuda import iket as _iket_official
+
+ def unavailable_nvrtc():
+ raise error_type("NVRTC could not be loaded")
+
+ nvrtc = SimpleNamespace(nvrtcVersion=unavailable_nvrtc)
+ monkeypatch.setitem(sys.modules, "cuda.bindings",
SimpleNamespace(nvrtc=nvrtc))
+ with pytest.raises(_iket_official.IketProfileError, match="CUDA NVRTC is
unavailable"):
+ _iket_official._validate_nvrtc_available() # pylint:
disable=protected-access
def test_injection_environment_accepts_run_iket_two_passes(tmp_path,
monkeypatch):
@@ -958,7 +970,7 @@ def test_cutlass_4_6_0_oracle_manifest_integrity():
def test_external_trace_contract():
trace_path = os.environ.get("TVM_IKET_OFFICIAL_TRACE_JSON")
if trace_path is None:
- pytest.skip("set TVM_IKET_OFFICIAL_TRACE_JSON after the locked
run-iket workload")
+ pytest.skip("set TVM_IKET_OFFICIAL_TRACE_JSON after a run-iket
workload")
trace = json.loads(Path(trace_path).read_text(encoding="utf-8"))
assert len(trace["launches"]) == 3
launches = {launch["kernelName"]: launch for launch in trace["launches"]}
