mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-25 07:14:16 +08:00
Chaofan Lin
fef9f76092
* Refactor backend module registration * Define backends with explicit specs * Require explicit backend for codegen * Rename backend spec to backend module * Keep host codegen on selected backend * Introduce backend context for compilation * Remove legacy backend compatibility APIs * Restore backend API docstrings
138 lines
4.7 KiB
Python
138 lines
4.7 KiB
Python
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
import tilelang
|
|
from tilelang.backend import get_backend
|
|
from tvm.target import Target
|
|
import tilelang.backend.target as target_registry
|
|
from tilelang.backend.target import (
|
|
auto_detect_target,
|
|
determine_target,
|
|
list_target_detectors,
|
|
register_target_detector,
|
|
)
|
|
|
|
|
|
def test_tilelang_does_not_export_target_wrapper():
|
|
assert not hasattr(tilelang, "Target")
|
|
|
|
|
|
def test_default_target_env_accepts_json_string(monkeypatch):
|
|
monkeypatch.setenv("TILELANG_DEFAULT_TARGET", '{"kind": "cuda", "arch": "sm_100f", "code": ["sm_100a", "sm_103a"]}')
|
|
|
|
assert tilelang.env.get_default_target() == {
|
|
"kind": "cuda",
|
|
"arch": "sm_100f",
|
|
"code": ["sm_100a", "sm_103a"],
|
|
}
|
|
|
|
|
|
def test_default_target_env_rejects_unquoted_json_keys(monkeypatch):
|
|
monkeypatch.setenv("TILELANG_DEFAULT_TARGET", '{kind: "cuda", arch: "sm_100f"}')
|
|
|
|
with pytest.raises(ValueError, match="Use JSON syntax"):
|
|
tilelang.env.get_default_target()
|
|
|
|
|
|
def test_default_target_env_keeps_plain_string(monkeypatch):
|
|
monkeypatch.setenv("TILELANG_DEFAULT_TARGET", "cuda")
|
|
|
|
assert tilelang.env.get_default_target() == "cuda"
|
|
|
|
|
|
def test_bare_cuda_target_uses_detected_exact_arch(monkeypatch):
|
|
from tilelang.cuda import target as cuda_target
|
|
|
|
monkeypatch.setattr(cuda_target, "_detect_torch_cuda_arch", lambda: "sm_90a")
|
|
|
|
target = determine_target("cuda", return_object=True)
|
|
|
|
assert isinstance(target, Target)
|
|
assert target.kind.name == "cuda"
|
|
assert str(target.attrs["arch"]) == "sm_90a"
|
|
|
|
|
|
def test_cuda_target_code_attr_survives_target_normalization():
|
|
target = determine_target(
|
|
{"kind": "cuda", "arch": "sm_100f", "code": ["sm_100a", "sm_103a"]},
|
|
return_object=True,
|
|
)
|
|
|
|
assert isinstance(target, Target)
|
|
assert target.kind.name == "cuda"
|
|
assert str(target.attrs["arch"]) == "sm_100f"
|
|
assert list(target.attrs["code"]) == ["sm_100a", "sm_103a"]
|
|
assert list(target.export()["code"]) == ["sm_100a", "sm_103a"]
|
|
|
|
|
|
def test_cuda_target_code_attr_rejects_string_code():
|
|
with pytest.raises(AssertionError, match="valid target config dict"):
|
|
determine_target({"kind": "cuda", "arch": "sm_100f", "code": "sm_100a"}, return_object=True)
|
|
|
|
|
|
def test_cuda_target_rejects_compute_arch():
|
|
with pytest.raises(AssertionError, match="valid target config dict"):
|
|
determine_target({"kind": "cuda", "arch": "compute_90"}, return_object=True)
|
|
|
|
|
|
def test_auto_target_uses_registered_detectors():
|
|
name = "unit-auto-target"
|
|
old_detectors = dict(target_registry._TARGET_DETECTORS)
|
|
try:
|
|
target_registry._TARGET_DETECTORS.clear()
|
|
register_target_detector(name, lambda: Target({"kind": "llvm", "mcpu": "native"}), override=True)
|
|
|
|
target = auto_detect_target()
|
|
|
|
assert isinstance(target, Target)
|
|
assert target.kind.name == "llvm"
|
|
assert str(target.attrs["mcpu"]) == "native"
|
|
assert name in list_target_detectors()
|
|
finally:
|
|
target_registry._TARGET_DETECTORS.clear()
|
|
target_registry._TARGET_DETECTORS.update(old_detectors)
|
|
|
|
|
|
def test_auto_target_detector_falls_through_none_result():
|
|
first_name = "unit-auto-none"
|
|
second_name = "unit-auto-fallback"
|
|
old_detectors = dict(target_registry._TARGET_DETECTORS)
|
|
try:
|
|
target_registry._TARGET_DETECTORS.clear()
|
|
register_target_detector(first_name, lambda: None, override=True)
|
|
register_target_detector(second_name, lambda: "llvm", override=True)
|
|
|
|
assert auto_detect_target() == "llvm"
|
|
assert list_target_detectors() == (first_name, second_name)
|
|
finally:
|
|
target_registry._TARGET_DETECTORS.clear()
|
|
target_registry._TARGET_DETECTORS.update(old_detectors)
|
|
|
|
|
|
def test_backend_module_resolves_execution_policy():
|
|
backend = get_backend("cpu")
|
|
c_target = Target("c")
|
|
llvm_target = Target("llvm")
|
|
|
|
assert backend.allowed_execution_backends(c_target) == ("cython", "tvm_ffi")
|
|
assert backend.allowed_execution_backends(llvm_target) == ("tvm_ffi",)
|
|
assert backend.resolve_execution_backend("auto", c_target).name == "cython"
|
|
assert backend.resolve_execution_backend("auto", llvm_target).name == "tvm_ffi"
|
|
|
|
|
|
def test_execution_backend_registry_rejects_invalid_backend():
|
|
target = Target("llvm")
|
|
backend = get_backend("cpu")
|
|
|
|
with pytest.raises(ValueError, match="Invalid execution backend"):
|
|
backend.resolve_execution_backend("nvrtc", target)
|
|
|
|
|
|
def test_execution_backend_registry_rejects_removed_dlpack_backend():
|
|
target = Target("llvm")
|
|
backend = get_backend("cpu")
|
|
|
|
with pytest.raises(ValueError, match="Invalid execution backend 'dlpack'"):
|
|
backend.resolve_execution_backend("dlpack", target)
|