[Backend] Add unified backend resolution policy

testing/python/backend/test_tilelang_backend_registry.py

@@ -0,0 +1,161 @@
+from __future__ import annotations
+
+import pytest
+
+from tilelang import tvm
+from tilelang.backend import Backend, ExecutionBackendSpec, list_backends, resolve_backend
+import tilelang.backend.registry as backend_registry
+from tilelang.engine.lower import is_cpu_device_backend
+from tilelang.jit.execution_backend import allowed_backends_for_target, resolve_execution_backend
+from tilelang.jit.kernel import JITKernel
+
+
+def test_backend_descriptor_accepts_callable_pipeline():
+    target = tvm.target.Target("c")
+    mod = tvm.IRModule({})
+    backend = Backend("unit-test", ("c",), pipeline=lambda input_mod, input_target: input_mod)
+
+    assert backend.matches(target)
+    assert backend.lower(mod, target) is mod
+
+
+def test_builtin_backend_resolution():
+    expected = {
+        "c": "cpu",
+        "llvm": "cpu",
+        "cuda": "cuda",
+        "hip": "rocm",
+        "metal": "metal",
+        "webgpu": "webgpu",
+    }
+
+    for target_kind, backend_name in expected.items():
+        assert resolve_backend(tvm.target.Target(target_kind)).name == backend_name
+
+
+def test_missing_codegen_hook_reports_backend_name():
+    backend = Backend("unit-test", ("c",))
+
+    with pytest.raises(ValueError, match="unit-test"):
+        backend.codegen(tvm.IRModule({}), tvm.target.Target("c"), compile=False)
+
+
+def test_list_backends_returns_copy():
+    registered = list_backends()
+    registered.clear()
+
+    assert list_backends()
+
+
+def test_backend_descriptor_freezes_nested_mappings():
+    backend = Backend(
+        "unit-test",
+        ("c",),
+        features={"feature": lambda target: True},
+        execution_backends={"fast": ExecutionBackendSpec("fast")},
+    )
+
+    with pytest.raises(TypeError):
+        backend.features["other"] = lambda target: False
+    with pytest.raises(TypeError):
+        backend.execution_backends["slow"] = ExecutionBackendSpec("slow")
+
+
+def test_backend_override_prunes_old_target_bucket():
+    backend_name = "unit-prune"
+    old_kind = "unit-old"
+    new_kind = "unit-new"
+
+    try:
+        backend_registry.register_backend(Backend(backend_name, (old_kind,)), override=True)
+        assert backend_name in backend_registry._TARGET_INDEX[old_kind]
+
+        backend_registry.register_backend(Backend(backend_name, (new_kind,)), override=True)
+
+        assert old_kind not in backend_registry._TARGET_INDEX
+        assert backend_name in backend_registry._TARGET_INDEX[new_kind]
+    finally:
+        backend_registry._BACKENDS.pop(backend_name, None)
+        backend_registry._CALLBACKS_REGISTERED.discard(backend_name)
+        for kind in (old_kind, new_kind):
+            names = backend_registry._TARGET_INDEX.get(kind, [])
+            if backend_name in names:
+                names.remove(backend_name)
+            if not names:
+                backend_registry._TARGET_INDEX.pop(kind, None)
+
+
+def test_backend_resolves_execution_backend_policy():
+    target = tvm.target.Target("c")
+    backend = Backend(
+        "unit-test",
+        ("c",),
+        execution_backends={
+            "slow": ExecutionBackendSpec("slow", adapter="cython"),
+            "fast": ExecutionBackendSpec(
+                "fast",
+                adapter="tvm_ffi",
+                enable_host_codegen=True,
+                enable_device_compile=True,
+            ),
+        },
+        default_execution_backend="fast",
+    )
+
+    spec = backend.resolve_execution_backend("auto", target)
+
+    assert spec.name == "fast"
+    assert spec.adapter == "tvm_ffi"
+    assert spec.enable_host_codegen
+    assert spec.enable_device_compile
+
+
+def test_execution_backend_resolves_through_backend_descriptor():
+    expected = {
+        "c": "cython",
+        "llvm": "cython",
+        "cuda": "tvm_ffi",
+        "hip": "tvm_ffi",
+        "metal": "tvm_ffi",
+        "webgpu": "tvm_ffi",
+    }
+
+    for target_kind, execution_backend in expected.items():
+        assert resolve_execution_backend("auto", tvm.target.Target(target_kind)) == execution_backend
+
+
+def test_invalid_execution_backend_reports_resolved_backend():
+    with pytest.raises(ValueError, match="Backend 'rocm'"):
+        resolve_execution_backend("nvrtc", tvm.target.Target("hip"))
+
+
+def test_cuda_only_execution_backend_policy_is_backend_owned():
+    target = tvm.target.Target("cuda")
+    backend = resolve_backend(target)
+
+    assert "nvrtc" in allowed_backends_for_target(target)
+    assert "nvrtc" in backend.allowed_execution_backends(target)
+    assert "nvrtc" not in resolve_backend(tvm.target.Target("hip")).allowed_execution_backends(tvm.target.Target("hip"))
+
+
+def test_llvm_is_cpu_device_backend():
+    assert is_cpu_device_backend(tvm.target.Target("c"))
+    assert is_cpu_device_backend(tvm.target.Target("llvm"))
+
+
+def test_jit_source_helpers_use_resolved_adapter_key():
+    class FakeAdapter:
+        def get_kernel_source(self, kernel_only=True):
+            return f"kernel:{kernel_only}"
+
+        def get_host_source(self):
+            return "host"
+
+    kernel = JITKernel.__new__(JITKernel)
+    kernel.execution_backend = "fast"
+    kernel.execution_backend_spec = ExecutionBackendSpec("fast", adapter="tvm_ffi")
+    kernel.adapter = FakeAdapter()
+    kernel.artifact = None
+
+    assert kernel.get_kernel_source(kernel_only=True) == "kernel:True"
+    assert kernel.get_host_source() == "host"

tilelang/backend/README.md

@@ -8,12 +8,12 @@ TileLang language surface backend-neutral.
 
 The Python backend layer is split into two parts:
 
-- `tilelang/backend/`: common backend infrastructure, especially pass-pipeline
-  registration, host/device-codegen registration, and shared pipeline
-  utilities.
+- `tilelang/backend/`: common backend infrastructure, especially the `Backend`
+  descriptor, backend registry, pass-pipeline registration, host/device-codegen
+  registration, target registration, and shared pipeline utilities.
 - `tilelang/<backend>/`: backend-owned Python implementation files, such as
-  pass pipelines, host/device-codegen entry registration, tile-op
-  implementation registration, and backend intrinsics.
+  backend descriptors, pass pipelines, tile-op implementation registration,
+  callbacks, codegen hooks, and backend intrinsics.
 
 The native side mirrors this split under `src/<backend>/`, where C++ op
 lowering, codegen, runtime modules, stubs, and backend-local CMake files live.
@@ -22,18 +22,23 @@ lowering, codegen, runtime modules, stubs, and backend-local CMake files live.
 ## Lowering Entry
 
 `tilelang/engine/lower.py` owns the high-level lowering entry. It runs
-backend-independent semantic checks first, then resolves a pass pipeline from
-the TVM target kind:
+backend-independent semantic checks first, then resolves a `Backend` from the
+TVM target kind:
 
 ```text
 PreLowerSemanticCheck(mod)
-pipeline = resolve_pipeline(target)
-mod = pipeline.lower(mod, target)
+backend = resolve_backend(target)
+mod = backend.lower(mod, target)
+codegen_mod = backend.codegen(device_mod, target, compile=enable_device_compile)
 ```
 
-The resolver is implemented in `tilelang/backend/pass_pipeline/pipeline.py`.
-Backends register a `PassPipeline(name, lower)` at import time. The pipeline
-name should match `target.kind.name`.
+The resolver is implemented in `tilelang/backend/registry.py`. Backends register
+a `Backend` descriptor at import time. The descriptor owns target matching,
+pipeline lowering, device codegen hooks, optional host pre-codegen hooks,
+callback registration, feature queries, JIT metadata, and CMake metadata.
+
+The existing `PassPipeline` API remains available and is wrapped by
+`Backend.pipeline` during migration.
 
 Device codegen follows the same ownership model after host/device splitting:
 
@@ -67,12 +72,73 @@ host functions that need Metal runtime context.
 
 | Python package | Target kind | Notes |
 | --- | --- | --- |
-| `tilelang/cuda` | `cuda` | CUDA-specific pass sequence, CUDA tile ops, MMA/WGMMA/TCGEN05 intrinsics, CUDA transform wrappers. |
-| `tilelang/rocm` | `hip` | ROCm/HIP pass sequence and MFMA/WMMA tile-op implementations. |
-| `tilelang/cpu` | `c`, `llvm` | CPU pass sequence and scalar CPU tile-op implementations. |
-| `tilelang/metal` | `metal` | Metal pass sequence and Metal GEMM registration. |
+| `tilelang/cuda/backend.py` | `cuda` | CUDA-specific pass sequence, CUDA tile ops, MMA/WGMMA/TCGEN05 intrinsics, CUDA transform wrappers, CUDA compile callbacks. |
+| `tilelang/rocm/backend.py` | `hip` | ROCm/HIP pass sequence, MFMA/WMMA tile-op implementations, HIP compile callback. |
+| `tilelang/cpu/backend.py` | `c`, `llvm` | CPU pass sequence and scalar CPU tile-op implementations. |
+| `tilelang/metal/backend.py` | `metal` | Metal pass sequence, Metal GEMM registration, Metal host pre-codegen hook. |
 | `tilelang/backend/common.py` | `webgpu` | Temporary/common registration for targets that do not yet own a dedicated Python backend package. |
 
+The backend package name does not have to match `target.kind.name`. ROCm is the
+main example: the package is `tilelang/rocm`, but it registers target kind
+`hip`.
+
+## Backend Descriptor
+
+`Backend` is a Python-side descriptor:
+
+```python
+Backend(
+    name="cuda",
+    target_kinds=("cuda",),
+    pipeline=cuda_pipeline,
+    device_codegen=cuda_device_codegen,
+    device_codegen_without_compile=cuda_device_codegen_without_compile,
+    register_callbacks=register_cuda_callbacks,
+    features={"warp_size": target_get_warp_size},
+    execution_backends={
+        "tvm_ffi": ExecutionBackendSpec(
+            "tvm_ffi",
+            enable_host_codegen=True,
+            enable_device_compile=True,
+        ),
+        "cython": ExecutionBackendSpec("cython"),
+        "nvrtc": ExecutionBackendSpec("nvrtc", is_available=is_nvrtc_available),
+    },
+    default_execution_backend="tvm_ffi",
+    cmake_name="CUDA",
+)
+```
+
+The registry resolves exactly one backend for a concrete TVM target. If multiple
+backends match the same target kind, `priority` must break the tie; otherwise
+resolution fails with an explicit ambiguity error.
+
+The backend descriptor also owns target-dependent JIT execution policy. For
+example, CUDA declares `nvrtc`, ROCm does not; CPU can choose `cython` by
+default; and a backend can choose whether an execution mode requests host
+codegen and device compilation.
+
+Lazy import is supported through `register_lazy_backend(target_kind,
+import_path)`, so optional backends can delay importing toolchain-specific Python
+modules until a matching target is requested.
+
+## TVM/tvm-ffi Boundary
+
+The `Backend` registry is intentionally Python-side because it stores Python
+callables, lazy import policy, feature query functions, and diagnostics.
+
+TVM/tvm-ffi remains the right boundary for cross-language registration:
+
+- Python compile callbacks use `tvm_ffi.register_global_func`, for example
+  `tilelang_callback_cuda_compile` and `tilelang_callback_hip_compile`.
+- Native transforms and codegen entry points stay registered from C++ through
+  TVM FFI global functions, for example `target.build.tilelang_cuda`.
+- Python `_ffi_api.py` modules should continue using `tvm_ffi.init_ffi_api`.
+
+The first implementation should not make `Backend` a TVM FFI `ObjectRef`;
+there is no native enumeration requirement yet, and doing so would make the
+Python orchestration path more complex.
+
 ## `tilelang/backend`
 
 `tilelang/backend` should stay small. It contains shared backend plumbing, not
@@ -81,15 +147,24 @@ backend-specific implementation details.
 ```text
 tilelang/backend/
   __init__.py
+  backend.py
   common.py
+  codegen.py
   device_codegen.py
+  execution_backend.py
   host_codegen.py
+  registry.py
+  target.py
   pass_pipeline/
     __init__.py
     pipeline.py
     pipeline_utils.py
 ```
 
+- `backend.py` defines the `Backend` descriptor and hook types.
+- `registry.py` defines `register_backend`, `register_lazy_backend`,
+  `resolve_backend`, `get_backend`, and `list_backends`.
+- `codegen.py` contains shared device-codegen cleanup helpers.
 - `pass_pipeline/pipeline.py` defines `PassPipeline`, `register_pipeline`, and
   `resolve_pipeline`.
 - `device_codegen.py` defines `DeviceCodegen`, `register_device_codegen`, and
@@ -112,26 +187,37 @@ for that backend.
 
 ```text
 tilelang/cuda/
+  backend.py
   codegen.py
+  execution_backend.py
   pipeline.py
+  target.py
   transform/
   op/
   intrinsics/
 
 tilelang/rocm/
+  backend.py
   codegen.py
+  execution_backend.py
   pipeline.py
+  target.py
   op/
   intrinsics/
 
 tilelang/cpu/
+  backend.py
   codegen.py
+  execution_backend.py
   pipeline.py
   op/
 
 tilelang/metal/
+  backend.py
   codegen.py
+  execution_backend.py
   pipeline.py
+  target.py
   transform/
   op/
   intrinsics/

tilelang/backend/__init__.py

@@ -1,3 +1,4 @@
+from .backend import Backend, ExecutionBackendSpec  # noqa: F401
 from .pass_pipeline import PassPipeline, register_pipeline, resolve_pipeline  # noqa: F401
 from .device_codegen import (  # noqa: F401
     DeviceCodegen,
@@ -18,7 +19,6 @@
     resolve_host_codegen,
 )
 from .execution_backend import (  # noqa: F401
-    ExecutionBackendSpec,
     allowed_backends_for_target,
     canonicalize_execution_backend,
     register_execution_backend,
@@ -32,12 +32,20 @@
     register_target_detector,
     register_target_normalizer,
 )
+from .registry import (  # noqa: F401
+    get_backend,
+    list_backends,
+    register_backend,
+    register_lazy_backend,
+    resolve_backend,
+)
 
 register_lazy_execution_backends("cuda", "tilelang.cuda.execution_backend")
 register_lazy_execution_backends("hip", "tilelang.rocm.execution_backend")
 register_lazy_execution_backends("c", "tilelang.cpu.execution_backend")
 register_lazy_execution_backends("llvm", "tilelang.cpu.execution_backend")
 register_lazy_execution_backends("metal", "tilelang.metal.execution_backend")
+register_lazy_execution_backends("webgpu", "tilelang.backend.common")
 
 register_lazy_device_codegen("cuda", "tilelang.cuda.codegen")
 register_lazy_device_codegen("hip", "tilelang.rocm.codegen")
@@ -50,4 +58,11 @@
 register_lazy_host_codegen("llvm", "tilelang.cpu.codegen")
 register_lazy_host_codegen_hooks("metal", "tilelang.metal.codegen")
 
+register_lazy_backend("cuda", "tilelang.cuda.backend")
+register_lazy_backend("hip", "tilelang.rocm.backend")
+register_lazy_backend("c", "tilelang.cpu.backend")
+register_lazy_backend("llvm", "tilelang.cpu.backend")
+register_lazy_backend("metal", "tilelang.metal.backend")
+register_lazy_backend("webgpu", "tilelang.backend.common")
+
 from . import common as common  # noqa: F401,E402