mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-21 06:40:23 +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
124 lines
4.4 KiB
Python
124 lines
4.4 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import replace
|
|
|
|
import pytest
|
|
|
|
from tilelang import tvm
|
|
from tilelang.backend import BackendContext, create_backend_context, get_backend, list_backends, register_backend
|
|
|
|
|
|
def test_builtin_backend_modules_are_explicit():
|
|
expected = {
|
|
"cuda": (("cuda",), ["tvm_ffi", "nvrtc", "cython"]),
|
|
"cutedsl": (("cuda",), ["cutedsl"]),
|
|
"rocm": (("hip",), ["tvm_ffi", "cython"]),
|
|
"cpu": (("c", "llvm"), ["cython", "tvm_ffi"]),
|
|
"metal": (("metal",), ["torch", "tvm_ffi"]),
|
|
"webgpu": (("webgpu",), ["tvm_ffi"]),
|
|
}
|
|
|
|
assert {
|
|
name: (backend.target_kinds, [spec.name for spec in backend.execution_backends]) for name, backend in list_backends().items()
|
|
} == expected
|
|
|
|
|
|
def test_list_backends_returns_copy():
|
|
backends = list_backends()
|
|
backends.clear()
|
|
|
|
assert list_backends()
|
|
|
|
|
|
def test_backend_registration_is_idempotent():
|
|
backend = get_backend("cuda")
|
|
|
|
assert register_backend(backend) is backend
|
|
|
|
|
|
def test_backend_rejects_different_redeclaration():
|
|
backend = get_backend("cuda")
|
|
replacement = replace(backend, callbacks={**backend.callbacks, "unit.callback": lambda: None})
|
|
|
|
with pytest.raises(ValueError, match="different declaration"):
|
|
register_backend(replacement)
|
|
|
|
|
|
def test_backend_rejects_duplicate_execution_names():
|
|
execution_backend = get_backend("cuda").execution_backends[0]
|
|
|
|
with pytest.raises(ValueError, match="must be unique"):
|
|
replace(get_backend("cuda"), execution_backends=(execution_backend, execution_backend))
|
|
|
|
|
|
def test_backend_variants_require_target_predicates():
|
|
backend = replace(get_backend("cuda"), name="unit-cuda", supports_target=None)
|
|
|
|
with pytest.raises(ValueError, match="must define supports_target"):
|
|
register_backend(backend)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("target_kind", "backend_name", "component_name"),
|
|
[
|
|
("c", "cpu", "c"),
|
|
("llvm", "cpu", "llvm"),
|
|
("cuda", "cuda", "cuda"),
|
|
("hip", "rocm", "hip"),
|
|
("metal", "metal", "metal"),
|
|
("webgpu", "webgpu", "webgpu"),
|
|
],
|
|
)
|
|
def test_backend_methods_are_the_primary_component_interface(target_kind: str, backend_name: str, component_name: str):
|
|
context = create_backend_context(target_kind, "c", "auto")
|
|
target = context.target
|
|
backend = context.module
|
|
|
|
assert backend.name == backend_name
|
|
assert backend.matches(target)
|
|
assert backend.get_device_codegen(target).name == component_name
|
|
assert context.execution_backend.name in backend.allowed_execution_backends(target)
|
|
if any(spec.enable_host_codegen for spec in backend.execution_backends):
|
|
assert backend.get_host_codegen(tvm.target.Target("c")).name == "c"
|
|
|
|
|
|
def test_backend_owns_compile_callbacks():
|
|
assert "tilelang_callback_cuda_validate" in get_backend("cuda").callbacks
|
|
assert "tilelang_callback_cuda_compile" in get_backend("cuda").callbacks
|
|
assert "tilelang_callback_hip_compile" in get_backend("rocm").callbacks
|
|
assert tvm.ffi.get_global_func("tilelang_callback_cuda_compile")
|
|
assert tvm.ffi.get_global_func("tilelang_callback_hip_compile")
|
|
|
|
|
|
def test_cutedsl_backend_reuses_cuda_pipeline():
|
|
cuda_target = tvm.target.Target("cuda")
|
|
cutedsl_target = tvm.target.Target({"kind": "cuda", "keys": ["cuda", "gpu", "cutedsl"]})
|
|
cuda_backend = get_backend("cuda")
|
|
cutedsl_backend = get_backend("cutedsl")
|
|
|
|
assert cuda_backend.name == "cuda"
|
|
assert cutedsl_backend.name == "cutedsl"
|
|
assert cutedsl_backend.get_pipeline(cutedsl_target) is cuda_backend.get_pipeline(cuda_target)
|
|
assert cutedsl_backend.get_device_codegen(cutedsl_target).name == "cutedsl"
|
|
assert cutedsl_backend.allowed_execution_backends(cutedsl_target) == ("cutedsl",)
|
|
|
|
|
|
def test_webgpu_only_exposes_tvm_ffi_execution():
|
|
target = tvm.target.Target("webgpu")
|
|
backend = get_backend("webgpu")
|
|
|
|
assert backend.allowed_execution_backends(target) == ("tvm_ffi",)
|
|
|
|
|
|
def test_create_backend_context_binds_compile_state():
|
|
context = create_backend_context("cuda", "c", "tvm_ffi")
|
|
|
|
assert isinstance(context, BackendContext)
|
|
assert context.module is get_backend("cuda")
|
|
assert context.target.kind.name == "cuda"
|
|
assert context.target_host.kind.name == "c"
|
|
assert context.execution_backend.name == "tvm_ffi"
|
|
|
|
with pytest.raises(AttributeError):
|
|
context.target = tvm.target.Target("llvm")
|