Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 14 additions & 0 deletions kernels/src/kernels/variants.py
Original file line number Diff line number Diff line change
Expand Up @@ -474,6 +474,20 @@ def _check_variants(
)
)
continue
# Same for ROCm: given a system ROCm version of x.y, only ROCm
# versions x.z, where z <= y qualify.
elif isinstance(selected_backend, ROCm) and isinstance(v.arch.backend, ROCm):
if (
v.arch.backend.version.major != selected_backend.version.major
or v.arch.backend.version.minor > selected_backend.version.minor
):
result.append(
VariantRejected(
variant=v,
reason=f"ROCm version ({v.arch.backend.version}) is not compatible with system ROCm version ({selected_backend.version})",
)
)
continue
elif v.arch.backend.variant_str != selected_backend.variant_str:
result.append(
VariantRejected(
Expand Down
36 changes: 36 additions & 0 deletions kernels/tests/test_variants.py
Original file line number Diff line number Diff line change
Expand Up @@ -265,6 +265,42 @@ def test_resolve_rocm(linux_abi):
assert {vs.variant for vs in trace} == set(variants)


ROCM_VARIANTS = [
parse_variant(f"torch215-cxx11-{rocm}-x86_64-linux") for rocm in ["rocm72", "rocm100", "rocm101", "rocm102"]
]


def _resolve_rocm(rocm_version: str):
return _resolve_variant_for_system(
variants=ROCM_VARIANTS,
selected_backend=ROCm(Version(rocm_version)),
cpu="x86_64",
os="linux",
torch_version=Version("2.15"),
torch_cxx11_abi=True,
tvm_ffi_version=None,
)


def test_resolve_rocm_best_older_minor():
# Newest minor that is not newer than the system, same major only.
result, _ = _resolve_rocm("10.1")
assert [v.variant_str for v in result] == [
"torch215-cxx11-rocm101-x86_64-linux",
"torch215-cxx11-rocm100-x86_64-linux",
]


def test_resolve_rocm_no_different_major():
result, _ = _resolve_rocm("7.2")
assert [v.variant_str for v in result] == ["torch215-cxx11-rocm72-x86_64-linux"]


def test_resolve_rocm_no_newer_minor():
result, _ = _resolve_rocm("10.0")
assert [v.variant_str for v in result] == ["torch215-cxx11-rocm100-x86_64-linux"]


def test_resolve_cpu_linux(linux_abi):
variants = _resolve_variants(linux_abi)
result, trace = _resolve_variant_for_system(
Expand Down
Loading