diff --git a/python/tvm/backend/cuda/iket.py b/python/tvm/backend/cuda/iket.py index c1edbed2effb..62d65c4cce72 100644 --- a/python/tvm/backend/cuda/iket.py +++ b/python/tvm/backend/cuda/iket.py @@ -24,7 +24,6 @@ from __future__ import annotations import functools -import hashlib import importlib.util import inspect import json @@ -59,48 +58,19 @@ _OUTPUT_TAIL_LINES = 100 _TERMINATION_GRACE_SECONDS = 5.0 -# Hashes are SHA-256 digests of files in the public 4.6.0 wheels. Only -# ABI-independent runtime/compiler binaries are pinned so this profile works -# with every Python version supported by that CUTLASS DSL release. _OFFICIAL_PROFILES = { "cutlass-4.6.0": { "nvrtc_version": (13, 2), - "versions": { + "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", }, - "files": { - "nvidia-cutlass-dsl-libs-base": { - "nvidia_cutlass_dsl/dsl_packages/iket/libiket_cubin_info.so": ( - "7ee839130c6bd129b04908a807c066118a459ebea644a59ecb6e41fbb323c103" - ), - "nvidia_cutlass_dsl/dsl_packages/iket/profiler/libsmodel_injection.so": ( - "83be54bd06e2cd82b2f6c17bbee6c925d049acae8d880242d4a5d5509a29e122" - ), - }, - "nvidia-cutlass-dsl-libs-cu13": { - "nvidia_cutlass_dsl/cu13/lib/libcute_dsl_runtime.so": ( - "2fa9809047485ae420ca99cab0678846de692e9608a179b0020834994311dd2f" - ), - }, - "nvidia-cuda-nvdisasm": { - "nvidia/cu13/bin/nvdisasm": ( - "5842e6adf9e232c9503a804915f158a576473e542577c070da3be49390474140" - ), - }, - "nvidia-cuda-nvrtc": { - "nvidia/cu13/lib/libnvrtc.so.13": ( - "c673cf3b5099d83b98a388a2bb21e5d6f481be3c4bb956e2d74c39cb714d8c63" - ), - "nvidia/cu13/lib/libnvrtc-builtins.so.13.2": ( - "6b1c571cc730d5fcfd57f322e1fa7e0e65de7454b2239ff6d552a09b82d47dbe" - ), - }, - }, } } @@ -201,12 +171,17 @@ def _profile_error(message: str) -> IketProfileError: ) -def _sha256(path: Path) -> str: - digest = hashlib.sha256() - with path.open("rb") as file_obj: - for chunk in iter(lambda: file_obj.read(1 << 20), b""): - digest.update(chunk) - return digest.hexdigest() +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: @@ -220,7 +195,9 @@ def _validate_run_iket_entrypoint() -> str: and item.value == "iket.cli.main:entrypoint" for item in entry_points ): - raise _profile_error("the run-iket entry point does not match the locked CUTLASS profile") + raise _profile_error( + "the run-iket entry point does not match the supported CUTLASS profile" + ) return executable @@ -240,34 +217,29 @@ def _validate_nvrtc_version(expected_version: tuple[int, int]) -> None: def _validate_official_installation(profile_name: str) -> str: - """Validate pinned host-side packages and return the official executable.""" + """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] - distributions = {} - for distribution_name, expected_version in profile_config["versions"].items(): - try: - distribution = metadata.distribution(distribution_name) - except metadata.PackageNotFoundError as err: - raise _profile_error( - f"{distribution_name}=={expected_version} is not installed" - ) from err - if distribution.version != expected_version: - raise _profile_error( - f"{distribution_name} must be {expected_version}, got {distribution.version}" - ) - distributions[distribution_name] = distribution - - for distribution_name, expected_files in profile_config["files"].items(): - distribution = distributions[distribution_name] - for relative_path, expected_digest in expected_files.items(): - path = Path(distribution.locate_file(relative_path)) - if not path.is_file(): - raise _profile_error(f"profile binary is missing: {relative_path}") - actual_digest = _sha256(path) - if actual_digest != expected_digest: + 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"profile binary hash mismatch for {relative_path}: {actual_digest}" + 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}" ) executable = _validate_run_iket_entrypoint() @@ -275,13 +247,11 @@ def _validate_official_installation(profile_name: str) -> str: return executable -def _validate_injection_environment(expected_injection_digest: str) -> None: +def _validate_injection_environment() -> None: injection_value = os.environ.get("CUDA_INJECTION64_PATH") injection_path = Path(injection_value) if injection_value else None if injection_path is None or not injection_path.is_file(): raise _profile_error("CUDA_INJECTION64_PATH was not supplied by run-iket") - if _sha256(injection_path) != expected_injection_digest: - raise _profile_error("CUDA_INJECTION64_PATH does not match the locked run-iket binary") config_value = os.environ.get("SMODEL_INJECTION_CONFIG") config_path = Path(config_value) if config_value else None @@ -314,16 +284,12 @@ def _validate_official_environment() -> str: f"{_PROFILE_ENV} must be set to {_DEFAULT_PROFILE}, got {profile_name!r}" ) executable = _validate_official_installation(profile_name) - profile_config = _OFFICIAL_PROFILES[profile_name] - injection_digest = profile_config["files"]["nvidia-cutlass-dsl-libs-base"][ - "nvidia_cutlass_dsl/dsl_packages/iket/profiler/libsmodel_injection.so" - ] - _validate_injection_environment(injection_digest) + _validate_injection_environment() return executable def validate_official_environment() -> None: - """Validate the exact official runtime before an instrumented CUBIN is loaded.""" + """Validate the supported official runtime before an instrumented CUBIN is loaded.""" _validate_official_environment() @@ -469,7 +435,7 @@ def _child_environment( 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 locked profiling entry point. + # children started by this validated profiling entry point. child_env[_INJECTED_CHILD_ENABLE_ENV] = "1" return child_env diff --git a/tests/python/tirx/iket/test_iket_profiler.py b/tests/python/tirx/iket/test_iket_profiler.py index 4ad200e912b0..ba6f30f3d7f5 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 locks, and trace contracts.""" +"""Tests for IKET lowering, metadata, installation versions, and trace contracts.""" import hashlib import importlib @@ -769,35 +769,22 @@ def test_proxy_fails_closed_and_forbids_export(tmp_path, monkeypatch): assert executable._executable._jitted_mod is None # pylint: disable=protected-access -def test_environment_validation_is_not_process_cached(tmp_path, monkeypatch): +def test_environment_validation_is_not_process_cached(monkeypatch): from tvm.tirx.cuda import iket as _iket_official - injection_path = tmp_path / "libsmodel_injection.so" - injection_path.write_bytes(b"locked") - injection_relative = "nvidia_cutlass_dsl/dsl_packages/iket/profiler/libsmodel_injection.so" profile = { "nvrtc_version": (13, 2), - "versions": {"nvidia-cutlass-dsl-libs-base": "4.6.0"}, - "files": { - "nvidia-cutlass-dsl-libs-base": { - injection_relative: hashlib.sha256(b"locked").hexdigest() - } - }, + "minimum_versions": {"nvidia-cutlass-dsl-libs-base": "4.6.0"}, + "exact_versions": {}, } class FakeDistribution: version = "4.6.0" - @staticmethod - def locate_file(_relative_path): - return injection_path - 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 _expected_digest: 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") @@ -807,18 +794,55 @@ def locate_file(_relative_path): _iket_official.validate_official_environment() +@pytest.mark.parametrize( + ("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): + 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: + pass + + 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" + ) + + def test_injection_environment_accepts_run_iket_two_passes(tmp_path, monkeypatch): from tvm.tirx.cuda import iket as _iket_official injection = tmp_path / "libsmodel_injection.so" - injection.write_bytes(b"locked-injection") - expected_digest = hashlib.sha256(injection.read_bytes()).hexdigest() + injection.write_bytes(b"injection") config_path = tmp_path / "config.json" monkeypatch.setenv("CUDA_INJECTION64_PATH", str(injection)) monkeypatch.setenv("SMODEL_INJECTION_CONFIG", str(config_path)) config_path.write_text(json.dumps({"toolName": "tracker", "toolConfig": {}})) - _iket_official._validate_injection_environment(expected_digest) # pylint: disable=protected-access + _iket_official._validate_injection_environment() # pylint: disable=protected-access instrument = tmp_path / "instrument.config.json" instrument.write_text("{}", encoding="utf-8") @@ -831,15 +855,10 @@ def test_injection_environment_accepts_run_iket_two_passes(tmp_path, monkeypatch ), encoding="utf-8", ) - _iket_official._validate_injection_environment(expected_digest) # pylint: disable=protected-access - - with pytest.raises(RuntimeError, match="locked run-iket binary"): - _iket_official._validate_injection_environment("0" * 64) # pylint: disable=protected-access + _iket_official._validate_injection_environment() # pylint: disable=protected-access config_path.write_text(json.dumps({"toolName": "other"}), encoding="utf-8") with pytest.raises(RuntimeError, match="not generated by run-iket profile"): - _iket_official._validate_injection_environment( # pylint: disable=protected-access - expected_digest - ) + _iket_official._validate_injection_environment() # pylint: disable=protected-access def test_cutlass_4_6_0_oracle_manifest_integrity():