[Backend] Add unified backend resolution policy (#2318)

* 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
This commit is contained in:
Chaofan Lin
2026-08-06 02:41:57 +08:00
committed by GitHub
parent 752043ec1a
commit fef9f76092
51 changed files with 1142 additions and 1001 deletions
+9 -10
View File
@@ -1,11 +1,11 @@
from __future__ import annotations
import cloudpickle
import importlib
import os
import tilelang
import tilelang.cache.kernel_cache as kernel_cache_mod
from tilelang.backend import create_backend_context
from tilelang.cache.cuda_binary_cache import CUDABinaryCache
from tilelang.cache.kernel_cache import KernelCache
from tilelang.env import env
@@ -37,7 +37,8 @@ def test_kernel_cache_namespace_includes_host_platform(monkeypatch):
def test_cuda_binary_cache_hit_skips_nvcc_compile(monkeypatch, tmp_path):
_set_cache_dirs(monkeypatch, tmp_path)
lower = importlib.import_module("tilelang.engine.lower")
from tilelang.cuda import backend as cuda_backend
monkeypatch.setattr(env, "TILELANG_KERNEL_CACHE_USE_LIB_STAMP", "0")
compile_calls = []
@@ -46,7 +47,7 @@ def test_cuda_binary_cache_hit_skips_nvcc_compile(monkeypatch, tmp_path):
compile_calls.append((code, target_format, tuple(arch), tuple(options or ())))
return bytearray(b"fake-cubin")
monkeypatch.setattr(lower.nvcc, "compile_cuda", fake_compile_cuda)
monkeypatch.setattr(cuda_backend.nvcc, "compile_cuda", fake_compile_cuda)
target = Target({"kind": "cuda", "arch": "sm_90a"})
source = 'extern "C" __global__ void kernel() {}'
@@ -56,12 +57,12 @@ def test_cuda_binary_cache_hit_skips_nvcc_compile(monkeypatch, tmp_path):
tilelang.PassConfigKey.TL_DEVICE_COMPILE_FLAGS: ["--extra-device-vectorization"],
}
first = lower.tilelang_callback_cuda_compile(source, target)
second = lower.tilelang_callback_cuda_compile(source, target)
first = cuda_backend.tilelang_callback_cuda_compile(source, target)
second = cuda_backend.tilelang_callback_cuda_compile(source, target)
# Different compiler options (e.g. --use_fast_math) change the generated
# SASS without changing the source, so they must NOT share a cache entry.
third = lower.tilelang_callback_cuda_compile(source, target, fast_math_pass_configs)
fourth = lower.tilelang_callback_cuda_compile(source, target, fast_math_pass_configs)
third = cuda_backend.tilelang_callback_cuda_compile(source, target, fast_math_pass_configs)
fourth = cuda_backend.tilelang_callback_cuda_compile(source, target, fast_math_pass_configs)
assert bytes(first) == b"fake-cubin"
assert bytes(second) == b"fake-cubin"
@@ -93,10 +94,8 @@ def test_disk_cache_load_failure_is_cache_miss(monkeypatch, tmp_path):
loaded = cache._load_kernel_from_disk(
key,
target="cuda",
target_host=None,
backend_context=create_backend_context("cuda", execution_backend="tvm_ffi"),
out_idx=[0],
execution_backend="tvm_ffi",
pass_configs=None,
compile_flags=None,
func=None,