From 1baceec4f306919aa04485384a54205116a11ecf Mon Sep 17 00:00:00 2001 From: SiriusNEO Date: Tue, 2 Jun 2026 13:48:05 +0800 Subject: [PATCH 1/2] Refactor backend resolution policy --- .../backend/test_tilelang_backend_registry.py | 97 +++++++++ tilelang/backend/README.md | 118 ++++++++-- tilelang/backend/__init__.py | 17 +- tilelang/backend/backend.py | 160 ++++++++++++++ tilelang/backend/codegen.py | 21 ++ tilelang/backend/common.py | 40 +++- tilelang/backend/registry.py | 77 +++++++ tilelang/cpu/__init__.py | 1 + tilelang/cpu/backend.py | 37 ++++ tilelang/cuda/__init__.py | 1 + tilelang/cuda/backend.py | 204 ++++++++++++++++++ tilelang/engine/lower.py | 160 +------------- tilelang/jit/adapter/utils.py | 4 +- tilelang/jit/execution_backend.py | 44 ++++ tilelang/jit/kernel.py | 42 ++-- tilelang/metal/__init__.py | 1 + tilelang/metal/backend.py | 40 ++++ tilelang/rocm/__init__.py | 1 + tilelang/rocm/backend.py | 72 +++++++ 19 files changed, 944 insertions(+), 193 deletions(-) create mode 100644 testing/python/backend/test_tilelang_backend_registry.py create mode 100644 tilelang/backend/backend.py create mode 100644 tilelang/backend/codegen.py create mode 100644 tilelang/backend/registry.py create mode 100644 tilelang/cpu/backend.py create mode 100644 tilelang/cuda/backend.py create mode 100644 tilelang/jit/execution_backend.py create mode 100644 tilelang/metal/backend.py create mode 100644 tilelang/rocm/backend.py diff --git a/testing/python/backend/test_tilelang_backend_registry.py b/testing/python/backend/test_tilelang_backend_registry.py new file mode 100644 index 0000000000..2cbb859f8b --- /dev/null +++ b/testing/python/backend/test_tilelang_backend_registry.py @@ -0,0 +1,97 @@ +from __future__ import annotations + +import pytest + +from tilelang import tvm +from tilelang.backend import Backend, ExecutionBackendSpec, list_backends, resolve_backend +from tilelang.jit.execution_backend import allowed_backends_for_target, resolve_execution_backend + + +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_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")) diff --git a/tilelang/backend/README.md b/tilelang/backend/README.md index 547e1bb613..0d5cc053d9 100644 --- a/tilelang/backend/README.md +++ b/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-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//`, 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/ diff --git a/tilelang/backend/__init__.py b/tilelang/backend/__init__.py index cfd99caee8..8d2fe52117 100644 --- a/tilelang/backend/__init__.py +++ b/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 diff --git a/tilelang/backend/backend.py b/tilelang/backend/backend.py new file mode 100644 index 0000000000..a509ac58b4 --- /dev/null +++ b/tilelang/backend/backend.py @@ -0,0 +1,160 @@ +from __future__ import annotations + +from collections.abc import Callable, Mapping +from dataclasses import dataclass, field + +from tvm import IRModule +from tvm.target import Target + +TargetLike = str | dict[str, object] | Target +LowerFunc = Callable[[IRModule, Target], IRModule] +CodegenFunc = Callable[[IRModule, Target], IRModule] +TargetPredicate = Callable[[Target], bool] +TargetNormalizer = Callable[[TargetLike], Target | None] +FeatureQuery = Callable[[Target], object] +AvailabilityCheck = Callable[[], bool] +CallbackRegistrar = Callable[[], None] +ExecutionBackendSelector = Callable[[Target], str] + + +def _always_available() -> bool: + return True + + +@dataclass(frozen=True, slots=True) +class ExecutionBackendSpec: + """JIT execution policy owned by a compiler backend.""" + + name: str + adapter: str | None = None + is_available: AvailabilityCheck = _always_available + supports_target: TargetPredicate | None = None + enable_host_codegen: bool = False + enable_device_compile: bool = False + priority: int = 0 + + def __post_init__(self) -> None: + if not self.name: + raise ValueError("ExecutionBackendSpec.name must not be empty") + if self.adapter is None: + object.__setattr__(self, "adapter", self.name) + + def matches(self, target: Target) -> bool: + if self.supports_target is not None: + return self.supports_target(target) + return True + + +@dataclass(frozen=True, slots=True) +class Backend: + """Compiler backend descriptor. + + A backend owns the target-specific lowering pipeline, codegen hooks, + callback registration, and feature queries for one or more TVM target kinds. + """ + + name: str + target_kinds: tuple[str, ...] + import_path: str | None = None + + pipeline: object | LowerFunc | None = None + supports_target: TargetPredicate | None = None + normalize_target: TargetNormalizer | None = None + is_available: AvailabilityCheck = _always_available + features: Mapping[str, FeatureQuery] = field(default_factory=dict) + + device_codegen: CodegenFunc | None = None + device_codegen_without_compile: CodegenFunc | None = None + host_pre_codegen: CodegenFunc | None = None + register_callbacks: CallbackRegistrar | None = None + + execution_backends: Mapping[str, ExecutionBackendSpec] = field(default_factory=dict) + default_execution_backend: str | ExecutionBackendSelector | None = None + jit_execution_backends: tuple[str, ...] = () + cmake_name: str | None = None + language_modules: tuple[str, ...] = () + priority: int = 0 + + def __post_init__(self) -> None: + if not self.name: + raise ValueError("Backend.name must not be empty") + if not self.target_kinds: + raise ValueError(f"Backend {self.name!r} must register at least one target kind") + object.__setattr__(self, "target_kinds", tuple(self.target_kinds)) + execution_backends = dict(self.execution_backends) + if not execution_backends and self.jit_execution_backends: + execution_backends = {name: ExecutionBackendSpec(name=name, adapter=name) for name in tuple(self.jit_execution_backends)} + object.__setattr__(self, "execution_backends", execution_backends) + object.__setattr__(self, "jit_execution_backends", tuple(execution_backends)) + object.__setattr__(self, "language_modules", tuple(self.language_modules)) + + def matches(self, target: Target) -> bool: + if self.supports_target is not None: + return self.supports_target(target) + return target.kind.name in self.target_kinds + + def lower(self, mod: IRModule, target: Target) -> IRModule: + if self.pipeline is None: + raise ValueError(f"Backend {self.name!r} has no pass pipeline") + if hasattr(self.pipeline, "lower"): + return self.pipeline.lower(mod, target) # type: ignore[no-any-return, attr-defined] + return self.pipeline(mod, target) # type: ignore[misc, no-any-return] + + def codegen(self, mod: IRModule, target: Target, *, compile: bool) -> IRModule: + hook = self.device_codegen if compile else self.device_codegen_without_compile + if hook is None: + mode = "compiled" if compile else "source" + raise ValueError(f"Backend {self.name!r} has no {mode} device codegen hook") + return hook(mod, target) + + def allowed_execution_backends(self, target: Target, *, include_unavailable: bool = True) -> tuple[str, ...]: + """Return execution backend names that this backend allows for *target*.""" + + specs = [spec for spec in self.execution_backends.values() if spec.matches(target)] + if not include_unavailable: + specs = [spec for spec in specs if spec.is_available()] + specs.sort(key=lambda spec: spec.priority, reverse=True) + return tuple(spec.name for spec in specs) + + def resolve_execution_backend(self, requested: str | None, target: Target) -> ExecutionBackendSpec: + """Resolve a requested JIT execution backend under this compiler backend.""" + + requested_name = None if requested in (None, "auto") else requested + if requested_name is None: + if callable(self.default_execution_backend): + requested_name = self.default_execution_backend(target) + else: + requested_name = self.default_execution_backend + + if requested_name is None: + available = self.allowed_execution_backends(target, include_unavailable=False) + if not available: + raise ValueError(f"Backend {self.name!r} has no available execution backend for target {target}") + requested_name = available[0] + + spec = self.execution_backends.get(requested_name) + if spec is None or not spec.matches(target): + allowed = ", ".join(self.allowed_execution_backends(target)) or "" + raise ValueError( + f"Invalid execution backend {requested_name!r} for Backend {self.name!r} " + f"and target {target.kind.name!r}. Allowed: {allowed}. Tip: use execution_backend='auto'." + ) + + if not spec.is_available(): + available = ", ".join(self.allowed_execution_backends(target, include_unavailable=False)) or "" + raise ValueError(f"Execution backend {requested_name!r} for Backend {self.name!r} is unavailable. Available: {available}.") + + return spec + + def preprocess_host_codegen(self, mod: IRModule, target: Target) -> IRModule: + if self.host_pre_codegen is None: + return mod + return self.host_pre_codegen(mod, target) + + def feature(self, name: str, target: Target, default: object = None) -> object: + query = self.features.get(name) + return default if query is None else query(target) + + def ensure_callbacks_registered(self) -> None: + if self.register_callbacks is not None: + self.register_callbacks() diff --git a/tilelang/backend/codegen.py b/tilelang/backend/codegen.py new file mode 100644 index 0000000000..c9f8c5b679 --- /dev/null +++ b/tilelang/backend/codegen.py @@ -0,0 +1,21 @@ +from __future__ import annotations + +from tilelang import tvm +from tvm import tirx +from tvm.target import Target + +import tilelang.transform + + +def prepare_device_codegen(device_mod: tvm.IRModule) -> tvm.IRModule: + """Run common cleanup before backend device codegen.""" + + device_mod = tilelang.transform.LowerIntrin()(device_mod) + device_mod = tirx.transform.Simplify()(device_mod) + device_mod = tilelang.transform.HoistBroadcastValues()(device_mod) + return device_mod + + +def build_device_with_global_func(device_mod: tvm.IRModule, target: Target, global_func_name: str) -> tvm.IRModule: + device_mod = prepare_device_codegen(device_mod) + return tvm.ffi.get_global_func(global_func_name)(device_mod, target) diff --git a/tilelang/backend/common.py b/tilelang/backend/common.py index 16d5bda13f..94530d1cb7 100644 --- a/tilelang/backend/common.py +++ b/tilelang/backend/common.py @@ -1,14 +1,46 @@ from __future__ import annotations -from tilelang.backend.execution_backend import ExecutionBackendSpec, register_execution_backend +from tvm.target import Target + +from tilelang import tvm +from tilelang.backend import Backend, ExecutionBackendSpec, register_backend +from tilelang.backend.codegen import build_device_with_global_func +from tilelang.backend.execution_backend import ( + ExecutionBackendSpec as LegacyExecutionBackendSpec, + register_execution_backend, +) +from tilelang.webgpu.pipeline import WebGPUPassPipelineBody -# Importing the package registers its pass pipeline. import tilelang.webgpu # noqa: F401 -register_execution_backend("webgpu", ExecutionBackendSpec("cython"), override=True) +def webgpu_device_codegen_without_compile(device_mod: tvm.IRModule, target: Target) -> tvm.IRModule: + return build_device_with_global_func(device_mod, target, "target.build.webgpu") + + +webgpu_backend = Backend( + name="webgpu", + target_kinds=("webgpu",), + import_path="tilelang.backend.common", + pipeline=WebGPUPassPipelineBody, + device_codegen_without_compile=webgpu_device_codegen_without_compile, + execution_backends={ + "tvm_ffi": ExecutionBackendSpec( + name="tvm_ffi", + enable_host_codegen=True, + enable_device_compile=True, + ), + "cython": ExecutionBackendSpec(name="cython"), + }, + default_execution_backend="tvm_ffi", + cmake_name="WEBGPU", +) + +register_backend(webgpu_backend, override=True) + +register_execution_backend("webgpu", LegacyExecutionBackendSpec("cython"), override=True) register_execution_backend( "webgpu", - ExecutionBackendSpec("tvm_ffi", enable_host_codegen=True, enable_device_compile=True), + LegacyExecutionBackendSpec("tvm_ffi", enable_host_codegen=True, enable_device_compile=True), override=True, ) diff --git a/tilelang/backend/registry.py b/tilelang/backend/registry.py new file mode 100644 index 0000000000..88adbf8bb6 --- /dev/null +++ b/tilelang/backend/registry.py @@ -0,0 +1,77 @@ +from __future__ import annotations + +from importlib import import_module + +from tvm.target import Target + +from .backend import Backend + +_BACKENDS: dict[str, Backend] = {} +_TARGET_INDEX: dict[str, list[str]] = {} +_LAZY_IMPORTS: dict[str, str] = {} + + +def register_backend(backend: Backend, *, override: bool = False) -> Backend: + """Register a TileLang backend descriptor.""" + + if backend.name in _BACKENDS and not override: + raise ValueError(f"Backend {backend.name!r} is already registered") + + if backend.name in _BACKENDS: + old = _BACKENDS[backend.name] + for kind in old.target_kinds: + names = _TARGET_INDEX.get(kind, []) + if old.name in names: + names.remove(old.name) + + _BACKENDS[backend.name] = backend + for kind in backend.target_kinds: + _TARGET_INDEX.setdefault(kind, []).append(backend.name) + return backend + + +def register_lazy_backend(target_kind: str, import_path: str) -> None: + """Register an import path that can populate a target kind on demand.""" + + _LAZY_IMPORTS[target_kind] = import_path + + +def get_backend(name: str) -> Backend: + if name not in _BACKENDS: + raise ValueError(f"No backend registered with name {name!r}") + return _BACKENDS[name] + + +def list_backends() -> dict[str, Backend]: + return dict(_BACKENDS) + + +def _ensure_loaded_for_kind(target_kind: str) -> None: + if target_kind in _TARGET_INDEX: + return + import_path = _LAZY_IMPORTS.get(target_kind) + if import_path is not None: + import_module(import_path) + + +def resolve_backend(target: Target) -> Backend: + """Resolve exactly one backend for a TVM target.""" + + kind = target.kind.name + _ensure_loaded_for_kind(kind) + + candidates = [_BACKENDS[name] for name in _TARGET_INDEX.get(kind, []) if _BACKENDS[name].matches(target)] + if not candidates: + available = ", ".join(sorted(_BACKENDS)) or "" + raise ValueError(f"No TileLang backend registered for target {target}. Registered backends: {available}") + + candidates.sort(key=lambda backend: backend.priority, reverse=True) + if len(candidates) > 1 and candidates[0].priority == candidates[1].priority: + names = ", ".join(backend.name for backend in candidates) + raise ValueError(f"Multiple TileLang backends match target {target}: {names}") + + backend = candidates[0] + if not backend.is_available(): + raise ValueError(f"TileLang backend {backend.name!r} is registered but unavailable") + backend.ensure_callbacks_registered() + return backend diff --git a/tilelang/cpu/__init__.py b/tilelang/cpu/__init__.py index f3b82b90a5..b660b082f1 100644 --- a/tilelang/cpu/__init__.py +++ b/tilelang/cpu/__init__.py @@ -2,3 +2,4 @@ from . import op # noqa: F401 from . import pipeline # noqa: F401 from . import execution_backend # noqa: F401 +from . import backend # noqa: F401 diff --git a/tilelang/cpu/backend.py b/tilelang/cpu/backend.py new file mode 100644 index 0000000000..2fe2dce328 --- /dev/null +++ b/tilelang/cpu/backend.py @@ -0,0 +1,37 @@ +from __future__ import annotations + +from tvm.target import Target + +from tilelang import tvm +from tilelang.backend import Backend, ExecutionBackendSpec, register_backend +from tilelang.backend.codegen import build_device_with_global_func +from tilelang.cpu.pipeline import CPUPassPipelineBody + + +def cpu_device_codegen_without_compile(device_mod: tvm.IRModule, target: Target) -> tvm.IRModule: + if target.kind.name == "c": + return build_device_with_global_func(device_mod, target, "target.build.tilelang_c") + if target.kind.name == "llvm": + return build_device_with_global_func(device_mod, target, "target.build.llvm") + raise ValueError(f"Target {target.kind.name} is not supported by the CPU backend") + + +cpu_backend = Backend( + name="cpu", + target_kinds=("c", "llvm"), + import_path="tilelang.cpu", + pipeline=CPUPassPipelineBody, + device_codegen_without_compile=cpu_device_codegen_without_compile, + execution_backends={ + "cython": ExecutionBackendSpec(name="cython"), + "tvm_ffi": ExecutionBackendSpec( + name="tvm_ffi", + enable_host_codegen=True, + enable_device_compile=False, + ), + }, + default_execution_backend="cython", + cmake_name="CPU", +) + +register_backend(cpu_backend, override=True) diff --git a/tilelang/cuda/__init__.py b/tilelang/cuda/__init__.py index c1d8b142ca..ae30509f8c 100644 --- a/tilelang/cuda/__init__.py +++ b/tilelang/cuda/__init__.py @@ -5,3 +5,4 @@ from . import target # noqa: F401 from . import execution_backend # noqa: F401 from . import transform # noqa: F401 +from . import backend # noqa: F401 diff --git a/tilelang/cuda/backend.py b/tilelang/cuda/backend.py new file mode 100644 index 0000000000..1634292555 --- /dev/null +++ b/tilelang/cuda/backend.py @@ -0,0 +1,204 @@ +from __future__ import annotations + +import re + +import tvm_ffi +from tvm import tirx +from tvm.target import Target + +from tilelang import tvm +from tilelang.backend import Backend, ExecutionBackendSpec, register_backend +from tilelang.backend.codegen import build_device_with_global_func +from tilelang.contrib import nvcc +from tilelang.cuda.pipeline import cuda_pipeline +from tilelang.env import CUTLASS_INCLUDE_DIR, TILELANG_TEMPLATE_PATH +from tilelang.transform import PassConfigKey + +_CUDA_GLOBAL_KERNEL_PATTERN = re.compile(r'(?:extern\s+"C"\s+)?__global__\s+void\s+(?:__launch_bounds__\([^\)]*\)\s+)?(\w+)') +_CALLBACKS_REGISTERED = False + + +def _is_cutedsl_target(target: Target) -> bool: + return target.kind.name == "cuda" and "cutedsl" in target.keys + + +def _is_plain_cuda_target(target: Target) -> bool: + return target.kind.name == "cuda" and "cutedsl" not in target.keys + + +def _default_cuda_execution_backend(target: Target) -> str: + return "cutedsl" if _is_cutedsl_target(target) else "tvm_ffi" + + +def target_get_warp_size(target: Target) -> int: + return int(target.attrs.get("thread_warp_size", 32)) + + +def _is_nvrtc_available() -> bool: + from tilelang.jit.adapter.nvrtc import is_nvrtc_available + + return bool(is_nvrtc_available) + + +def _is_cutedsl_available() -> bool: + try: + from tilelang.jit.adapter.cutedsl.checks import check_cutedsl_available + + check_cutedsl_available() + except ImportError: + return False + return True + + +def _collect_external_cuda_kernel_names(source: str) -> list[str]: + kernel_names: list[str] = [] + seen_names: set[str] = set() + for match in _CUDA_GLOBAL_KERNEL_PATTERN.finditer(source): + kernel_name = match.group(1) + if kernel_name not in seen_names: + kernel_names.append(kernel_name) + seen_names.add(kernel_name) + return kernel_names + + +def _tilelang_callback_cuda_validate(device_mod): + for _, base_func in device_mod.functions.items(): + if not isinstance(base_func, tirx.PrimFunc) or not base_func.attrs: + continue + + code_block_source = base_func.attrs.get("code_block_source") + if code_block_source is None: + continue + + global_symbol = base_func.attrs.get("global_symbol") + if global_symbol is None: + raise ValueError("CodeGenTileLangCUDA expects source-kernel PrimFunc to have the global_symbol attribute") + + expected_name = str(global_symbol) + code_block_entry_name = base_func.attrs.get("code_block_entry_name") + if code_block_entry_name is not None and str(code_block_entry_name) != expected_name: + raise ValueError("T.CUDASourceCodeKernel expects the lowered device global_symbol to match entry_name") + + kernel_names = _collect_external_cuda_kernel_names(str(code_block_source)) + if not kernel_names: + raise ValueError("T.CUDASourceCodeKernel expects external CUDA source to declare at least one __global__ kernel") + if expected_name not in kernel_names: + raise ValueError( + "T.CUDASourceCodeKernel expected device global_symbol " + f"`{expected_name}` to match a __global__ kernel in the provided CUDA source. " + f"Available entries: {', '.join(kernel_names)}" + ) + + +def _tilelang_callback_cuda_compile(code, target, pass_config=None): + target_arch = nvcc.get_target_arch(nvcc.get_target_compute_version(target)) + + arch = [f"-arch=sm_{target_arch}"] + compile_format = "cubin" + + cfg = pass_config or {} + enable_fast_math = bool(cfg.get(PassConfigKey.TL_ENABLE_FAST_MATH, False)) + + ptxas_usage_level = cfg.get(PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL, None) + if ptxas_usage_level is not None: + ptxas_usage_level = int(ptxas_usage_level) + verbose_ptxas_output = bool(cfg.get(PassConfigKey.TL_ENABLE_PTXAS_VERBOSE_OUTPUT, False)) + + options = [ + "-std=c++20", + "-I" + TILELANG_TEMPLATE_PATH, + "-I" + CUTLASS_INCLUDE_DIR, + ] + extra_flags = cfg.get(PassConfigKey.TL_DEVICE_COMPILE_FLAGS, None) + if extra_flags: + import shlex + + if isinstance(extra_flags, str): + tokens = shlex.split(extra_flags) + else: + tokens = [] + for flag in extra_flags: + if isinstance(flag, str): + tokens.extend(shlex.split(flag)) + else: + tokens.append(str(flag)) + options += tokens + + verbose = False + if enable_fast_math: + options.append("--use_fast_math") + if ptxas_usage_level is not None: + options.append(f"--ptxas-options=--register-usage-level={ptxas_usage_level}") + if verbose_ptxas_output: + options.append("--ptxas-options=--verbose") + options.append("-w") + verbose = True + + return nvcc.compile_cuda( + code, + compile_format, + arch, + options=options, + verbose=verbose, + ) + + +def register_cuda_callbacks() -> None: + global _CALLBACKS_REGISTERED + if _CALLBACKS_REGISTERED: + return + tvm_ffi.register_global_func("tilelang_callback_cuda_validate", f=_tilelang_callback_cuda_validate, override=True) + tvm_ffi.register_global_func("tilelang_callback_cuda_compile", f=_tilelang_callback_cuda_compile, override=True) + _CALLBACKS_REGISTERED = True + + +def cuda_device_codegen(device_mod: tvm.IRModule, target: Target) -> tvm.IRModule: + global_func = "target.build.tilelang_" + ("cutedsl" if "cutedsl" in target.keys else "cuda") + return build_device_with_global_func(device_mod, target, global_func) + + +def cuda_device_codegen_without_compile(device_mod: tvm.IRModule, target: Target) -> tvm.IRModule: + global_func = "target.build.tilelang_" + ("cutedsl" if "cutedsl" in target.keys else "cuda") + "_without_compile" + return build_device_with_global_func(device_mod, target, global_func) + + +cuda_backend = Backend( + name="cuda", + target_kinds=("cuda",), + import_path="tilelang.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, + "supports_tma": nvcc.have_tma, + "supports_pdl": nvcc.have_pdl, + }, + execution_backends={ + "tvm_ffi": ExecutionBackendSpec( + name="tvm_ffi", + supports_target=_is_plain_cuda_target, + enable_host_codegen=True, + enable_device_compile=True, + ), + "cython": ExecutionBackendSpec( + name="cython", + supports_target=_is_plain_cuda_target, + ), + "nvrtc": ExecutionBackendSpec( + name="nvrtc", + is_available=_is_nvrtc_available, + supports_target=_is_plain_cuda_target, + ), + "cutedsl": ExecutionBackendSpec( + name="cutedsl", + is_available=_is_cutedsl_available, + supports_target=_is_cutedsl_target, + ), + }, + default_execution_backend=_default_cuda_execution_backend, + cmake_name="CUDA", +) + +register_backend(cuda_backend, override=True) diff --git a/tilelang/engine/lower.py b/tilelang/engine/lower.py index d37dc6e4af..e9ea45c336 100644 --- a/tilelang/engine/lower.py +++ b/tilelang/engine/lower.py @@ -2,27 +2,21 @@ from __future__ import annotations -import re from collections.abc import Callable import tilelang.transform from tilelang import tvm as tvm from tvm import tirx -import tvm_ffi from tvm.ir import CallingConv from tvm.target import Target -from tilelang.contrib import hipcc, nvcc -from tilelang.env import COMPOSABLE_KERNEL_INCLUDE_DIR, CUTLASS_INCLUDE_DIR, TILELANG_TEMPLATE_PATH -from tilelang.transform import PassConfigKey from tilelang.engine.param import KernelParam, CompiledArtifact from tilelang.engine.semantic_check import PreLowerSemanticCheck -from tilelang.backend.device_codegen import resolve_device_codegen -from tilelang.backend.host_codegen import apply_host_codegen_hooks, resolve_host_codegen +from tilelang.backend import resolve_backend +from tilelang.backend.host_codegen import resolve_host_codegen from tilelang.backend.target import determine_target -from tilelang.backend.pass_pipeline import resolve_pipeline def is_cpu_device_backend(target: Target): - return target.kind.name == "c" + return target.kind.name in {"c", "llvm"} def has_device_kernel_launch(attrs) -> bool: @@ -35,8 +29,8 @@ def is_device_call_c_device(func: tirx.PrimFunc): calling_conv = attrs.get("calling_conv", CallingConv.DEFAULT) is_cpacked = calling_conv == CallingConv.C_PACKED_FUNC - # Check if it's a C target - if "target" in attrs and attrs["target"].kind.name == "c" and not is_cpacked: + # CPU targets keep callable PrimFuncs in the device module for source/codegen. + if "target" in attrs and attrs["target"].kind.name in {"c", "llvm"} and not is_cpacked: return True return has_device_kernel_launch(attrs) @@ -54,130 +48,6 @@ def get_host_call(is_device_c: bool = False) -> Callable[[tirx.PrimFunc], bool]: return lambda func: not get_device_call(is_device_c)(func) -_CUDA_GLOBAL_KERNEL_PATTERN = re.compile(r'(?:extern\s+"C"\s+)?__global__\s+void\s+(?:__launch_bounds__\([^\)]*\)\s+)?(\w+)') - - -def _collect_external_cuda_kernel_names(source: str) -> list[str]: - kernel_names: list[str] = [] - seen_names: set[str] = set() - for match in _CUDA_GLOBAL_KERNEL_PATTERN.finditer(source): - kernel_name = match.group(1) - if kernel_name not in seen_names: - kernel_names.append(kernel_name) - seen_names.add(kernel_name) - return kernel_names - - -@tvm_ffi.register_global_func("tilelang_callback_cuda_validate", override=True) -def tilelang_callback_cuda_validate(device_mod): - for _, base_func in device_mod.functions.items(): - if not isinstance(base_func, tirx.PrimFunc) or not base_func.attrs: - continue - - code_block_source = base_func.attrs.get("code_block_source") - if code_block_source is None: - continue - - global_symbol = base_func.attrs.get("global_symbol") - if global_symbol is None: - raise ValueError("CodeGenTileLangCUDA expects source-kernel PrimFunc to have the global_symbol attribute") - - expected_name = str(global_symbol) - code_block_entry_name = base_func.attrs.get("code_block_entry_name") - if code_block_entry_name is not None and str(code_block_entry_name) != expected_name: - raise ValueError("T.CUDASourceCodeKernel expects the lowered device global_symbol to match entry_name") - - kernel_names = _collect_external_cuda_kernel_names(str(code_block_source)) - if not kernel_names: - raise ValueError("T.CUDASourceCodeKernel expects external CUDA source to declare at least one __global__ kernel") - if expected_name not in kernel_names: - raise ValueError( - "T.CUDASourceCodeKernel expected device global_symbol " - f"`{expected_name}` to match a __global__ kernel in the provided CUDA source. " - f"Available entries: {', '.join(kernel_names)}" - ) - - -@tvm_ffi.register_global_func("tilelang_callback_cuda_compile", override=True) -def tilelang_callback_cuda_compile(code, target, pass_config=None): - target_arch = nvcc.get_target_arch(nvcc.get_target_compute_version(target)) - - arch = [f"-arch=sm_{target_arch}"] - compile_format = "cubin" - - # Read pass-config keys (string-valued) like in jit.adapter.libgen.compile_lib - cfg = pass_config or {} - enable_fast_math = bool(cfg.get(PassConfigKey.TL_ENABLE_FAST_MATH, False)) - - ptxas_usage_level = cfg.get(PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL, None) - if ptxas_usage_level is not None: - ptxas_usage_level = int(ptxas_usage_level) - verbose_ptxas_output = bool(cfg.get(PassConfigKey.TL_ENABLE_PTXAS_VERBOSE_OUTPUT, False)) - - options = [ - # tl_templates/cuda/reduce.h uses explicit lambda template parameters - # (`[&](T) { ... }`) which require C++20. - "-std=c++20", - "-I" + TILELANG_TEMPLATE_PATH, - "-I" + CUTLASS_INCLUDE_DIR, - ] - # Merge extra device compiler flags from pass config, if provided - extra_flags = cfg.get(PassConfigKey.TL_DEVICE_COMPILE_FLAGS, None) - if extra_flags: - import shlex - - if isinstance(extra_flags, str): - tokens = shlex.split(extra_flags) - else: - tokens = [] - for flag in extra_flags: - if isinstance(flag, str): - tokens.extend(shlex.split(flag)) - else: - tokens.append(str(flag)) - options += tokens - - verbose = False - if enable_fast_math: - options.append("--use_fast_math") - if ptxas_usage_level is not None: - options.append(f"--ptxas-options=--register-usage-level={ptxas_usage_level}") - if verbose_ptxas_output: - options.append("--ptxas-options=--verbose") - options.append("-w") # Suppress warnings to make ptxas output more readable - verbose = True - - ptx = nvcc.compile_cuda( - code, - compile_format, - arch, - options=options, - verbose=verbose, - ) - - return ptx - - -@tvm_ffi.register_global_func("tilelang_callback_hip_compile", override=True) -def tilelang_callback_hip_compile(code, target): - from tilelang.rocm.target import target_get_mcpu - - arch = target_get_mcpu(target) - hsaco = hipcc.compile_hip( - code, - target_format="hsaco", - arch=arch, - options=[ - "-std=c++17", - "-I" + TILELANG_TEMPLATE_PATH, - "-I" + COMPOSABLE_KERNEL_INCLUDE_DIR, - ], - verbose=False, - ) - - return hsaco - - def extrac_params(func: tirx.PrimFunc) -> list[KernelParam]: tensor_types = [] for var in func.params: @@ -218,25 +88,17 @@ def host_codegen(host_mod: tvm.IRModule, target_host: Target, target: Target | N combine_context_call = getattr(tirx.transform, "CombineContextCall", None) if combine_context_call is not None: host_mod = combine_context_call()(host_mod) - host_mod = apply_host_codegen_hooks(host_mod, target_host, target) + if target is not None: + host_mod = resolve_backend(target).preprocess_host_codegen(host_mod, target) return resolve_host_codegen(target_host).lower(host_mod, target_host) -def _prepare_device_codegen_mod(device_mod: tvm.IRModule) -> tvm.IRModule: - device_mod = tilelang.transform.LowerIntrin()(device_mod) - device_mod = tirx.transform.Simplify()(device_mod) - device_mod = tilelang.transform.HoistBroadcastValues()(device_mod) - return device_mod - - def device_codegen(device_mod: tvm.IRModule, target: Target) -> tvm.IRModule: - device_mod = _prepare_device_codegen_mod(device_mod) - return resolve_device_codegen(target).lower(device_mod, target, compile_device=True) + return resolve_backend(target).codegen(device_mod, target, compile=True) def device_codegen_without_compile(device_mod: tvm.IRModule, target: Target) -> tvm.IRModule: - device_mod = _prepare_device_codegen_mod(device_mod) - return resolve_device_codegen(target).lower(device_mod, target, compile_device=False) + return resolve_backend(target).codegen(device_mod, target, compile=False) def lower_to_host_device_ir( @@ -268,8 +130,8 @@ def lower_to_host_device_ir( # Run backend-independent semantic checks before target-specific lowering. PreLowerSemanticCheck(mod) - pipeline = resolve_pipeline(target) - mod = pipeline.lower(mod, target) + backend = resolve_backend(target) + mod = backend.lower(mod, target) host_mod = tirx.transform.Filter(_is_host_call)(mod) device_mod = tirx.transform.Filter(_is_device_call)(mod) diff --git a/tilelang/jit/adapter/utils.py b/tilelang/jit/adapter/utils.py index df8dff461b..23a0e51f30 100644 --- a/tilelang/jit/adapter/utils.py +++ b/tilelang/jit/adapter/utils.py @@ -13,7 +13,7 @@ canon_target_host, is_cpu_device_backend, ) -from tilelang.backend.pass_pipeline.pipeline import resolve_pipeline +from tilelang.backend import resolve_backend from tilelang.engine.semantic_check import PreLowerSemanticCheck @@ -140,7 +140,7 @@ def get_annotated_mod( # Apply transformations PreLowerSemanticCheck(mod) - mod = resolve_pipeline(target).lower(mod, target) + mod = resolve_backend(target).lower(mod, target) # Define dispatch dictionary for different model types dispatch = { diff --git a/tilelang/jit/execution_backend.py b/tilelang/jit/execution_backend.py new file mode 100644 index 0000000000..c923fea240 --- /dev/null +++ b/tilelang/jit/execution_backend.py @@ -0,0 +1,44 @@ +from __future__ import annotations + +from tvm.target import Target + +from tilelang.backend import resolve_backend +from tilelang.backend.backend import ExecutionBackendSpec + +# Canonical names for execution backends used internally +_CANONICAL_MAP = { + "dlpack": "tvm_ffi", # historical alias +} + + +def _canon_backend(name: str | None) -> str | None: + if name is None: + return None + key = str(name).lower() + return _CANONICAL_MAP.get(key, key) + + +def allowed_backends_for_target(target: Target, *, include_unavailable: bool = True) -> list[str]: + """Return execution backends allowed by the resolved compiler backend. + + include_unavailable: if False, this will filter out backends that are known + to be unavailable at runtime (e.g., NVRTC without cuda-python installed). + """ + return list(resolve_backend(target).allowed_execution_backends(target, include_unavailable=include_unavailable)) + + +def resolve_execution_backend_spec(requested: str | None, target: Target) -> ExecutionBackendSpec: + """Resolve a user request to a concrete execution-backend spec.""" + + return resolve_backend(target).resolve_execution_backend(_canon_backend(requested), target) + + +def resolve_execution_backend(requested: str | None, target: Target) -> str: + """Resolve an execution backend string to a concrete backend. + + - Supports the alias "dlpack" -> "tvm_ffi". + - Supports the sentinel "auto" which selects a sensible default per target. + - Validates the combination (target, backend) and raises with helpful + alternatives when invalid. + """ + return resolve_execution_backend_spec(requested, target).name diff --git a/tilelang/jit/kernel.py b/tilelang/jit/kernel.py index 1035d623ec..6583d486d2 100644 --- a/tilelang/jit/kernel.py +++ b/tilelang/jit/kernel.py @@ -8,7 +8,6 @@ import tilelang from tilelang import tvm from tilelang import env -from tilelang.backend.execution_backend import resolve_execution_backend_spec from tvm.target import Target from tilelang.engine.param import CompiledArtifact, KernelParam from tilelang.jit.adapter import ( @@ -19,6 +18,7 @@ TVMFFIKernelAdapter, MetalKernelAdapter, ) +from tilelang.jit.execution_backend import resolve_execution_backend_spec from tilelang.profiler import Profiler, TensorSupplyType from tilelang.backend.target import determine_target from tilelang.contrib import nvcc as tl_nvcc @@ -106,11 +106,10 @@ def __init__( # The wrapper is normalized here because lower/codegen still consume TVM Target. self.target = determine_target(target, return_object=True) - self.execution_backend_spec = resolve_execution_backend_spec(execution_backend, self.target) self.execution_backend = self.execution_backend_spec.name - if self.execution_backend == "cython": + if self.execution_backend_spec.adapter == "cython": from tilelang.contrib.cc import get_cplus_compiler assert get_cplus_compiler() is not None, "Cython backend requires a C++ compiler, please install or use other backends." @@ -219,6 +218,7 @@ def _compile_and_create_adapter(self, tilelang_func: PrimFunc, out_idx: list[int target_host = self.target_host execution_backend = self.execution_backend + execution_backend_spec = self.execution_backend_spec pass_configs = dict(self.pass_configs) if self.pass_configs else {} compile_flags = self.compile_flags @@ -228,9 +228,9 @@ def _compile_and_create_adapter(self, tilelang_func: PrimFunc, out_idx: list[int compile_flags_cfg + compile_flags if compile_flags_cfg is not None else compile_flags ) - # Compile the function with TVM, optimizing with shared memory lowering. - enable_host_codegen = self.execution_backend_spec.enable_host_codegen - enable_device_compile = self.execution_backend_spec.enable_device_compile + # The compiler backend owns how each execution backend should request codegen. + enable_host_codegen = execution_backend_spec.enable_host_codegen + enable_device_compile = execution_backend_spec.enable_device_compile # Additional pass instruments pass_instruments = [] @@ -268,8 +268,9 @@ def create_adapter(adapter_cls: Callable[..., BaseKernelAdapter], **kwargs: Any) with jit_phase("adapter", verbose=verbose, **phase_context): return adapter_cls(**kwargs) - # Create an adapter based on the specified execution backend. - if execution_backend == "tvm_ffi": + # Create an adapter based on the backend-resolved adapter key. + adapter_name = execution_backend_spec.adapter + if adapter_name == "tvm_ffi": # Use TVMFFIKernelAdapter for interoperability with PyTorch via DLPack. # But we need to ensure that the runtime is enabled and the runtime module is not None. assert artifact.rt_mod is not None, "tvm_ffi backend requires a runtime module." @@ -287,7 +288,7 @@ def create_adapter(adapter_cls: Callable[..., BaseKernelAdapter], **kwargs: Any) pass_configs=pass_configs, compile_flags=compile_flags, ) - elif execution_backend == "cython": + elif adapter_name == "cython": adapter = create_adapter( CythonKernelAdapter, params=artifact.params, @@ -301,7 +302,7 @@ def create_adapter(adapter_cls: Callable[..., BaseKernelAdapter], **kwargs: Any) pass_configs=pass_configs, compile_flags=compile_flags, ) - elif execution_backend == "nvrtc": + elif adapter_name == "nvrtc": from tilelang.jit.adapter import NVRTCKernelAdapter adapter = create_adapter( @@ -317,7 +318,7 @@ def create_adapter(adapter_cls: Callable[..., BaseKernelAdapter], **kwargs: Any) pass_configs=pass_configs, compile_flags=compile_flags, ) - elif execution_backend == "torch": + elif adapter_name == "torch": assert is_metal_target(target) adapter = create_adapter( MetalKernelAdapter, @@ -332,7 +333,7 @@ def create_adapter(adapter_cls: Callable[..., BaseKernelAdapter], **kwargs: Any) # pass_configs=pass_configs, # compile_flags=compile_flags, ) - elif execution_backend == "cutedsl": + elif adapter_name == "cutedsl": assert is_cutedsl_target(target) adapter = create_adapter( CuTeDSLKernelAdapter, @@ -348,8 +349,7 @@ def create_adapter(adapter_cls: Callable[..., BaseKernelAdapter], **kwargs: Any) compile_flags=compile_flags, ) else: - # Handle invalid backend. - raise ValueError(f"Invalid execution backend: {execution_backend}") + raise ValueError(f"Execution backend {execution_backend!r} resolved unknown JIT adapter {adapter_name!r}") if capture_resources: self._resource_usage = pop_recorded() @@ -370,9 +370,10 @@ def _create_adapter_from_database( ) -> BaseKernelAdapter: target = self.target execution_backend = self.execution_backend + adapter_name = self.execution_backend_spec.adapter - # Create an adapter based on the specified execution backend. - if execution_backend == "tvm_ffi": + # Create an adapter based on the backend-resolved adapter key. + if adapter_name == "tvm_ffi": adapter = TVMFFIKernelAdapter.from_database( params=params, result_idx=result_idx, @@ -384,7 +385,7 @@ def _create_adapter_from_database( pass_configs=pass_configs, compile_flags=compile_flags, ) - elif execution_backend == "cython": + elif adapter_name == "cython": adapter = CythonKernelAdapter.from_database( params=params, result_idx=result_idx, @@ -395,7 +396,7 @@ def _create_adapter_from_database( kernel_lib_path=kernel_lib_path, pass_configs=pass_configs, ) - elif execution_backend == "nvrtc": + elif adapter_name == "nvrtc": from tilelang.jit.adapter import NVRTCKernelAdapter adapter = NVRTCKernelAdapter.from_database( @@ -409,7 +410,7 @@ def _create_adapter_from_database( pass_configs=pass_configs, compile_flags=compile_flags, ) - elif execution_backend == "cutedsl": + elif adapter_name == "cutedsl": adapter = CuTeDSLKernelAdapter.from_database( params=params, result_idx=result_idx, @@ -422,8 +423,7 @@ def _create_adapter_from_database( compile_flags=compile_flags, ) else: - # Handle invalid backend. - raise ValueError(f"Invalid execution backend: {execution_backend}") + raise ValueError(f"Execution backend {execution_backend!r} resolved unknown or uncached JIT adapter {adapter_name!r}") return adapter diff --git a/tilelang/metal/__init__.py b/tilelang/metal/__init__.py index c1d8b142ca..ae30509f8c 100644 --- a/tilelang/metal/__init__.py +++ b/tilelang/metal/__init__.py @@ -5,3 +5,4 @@ from . import target # noqa: F401 from . import execution_backend # noqa: F401 from . import transform # noqa: F401 +from . import backend # noqa: F401 diff --git a/tilelang/metal/backend.py b/tilelang/metal/backend.py new file mode 100644 index 0000000000..9b5a5ed875 --- /dev/null +++ b/tilelang/metal/backend.py @@ -0,0 +1,40 @@ +from __future__ import annotations + +from tvm.target import Target + +from tilelang import tvm +from tilelang.backend import Backend, ExecutionBackendSpec, register_backend +from tilelang.backend.codegen import build_device_with_global_func +from tilelang.metal.pipeline import metal_pipeline +from tilelang.metal.transform import MarkHostMetalContext + + +def metal_device_codegen(device_mod: tvm.IRModule, target: Target) -> tvm.IRModule: + return build_device_with_global_func(device_mod, target, "target.build.tilelang_metal") + + +def metal_host_pre_codegen(host_mod: tvm.IRModule, target: Target) -> tvm.IRModule: + return MarkHostMetalContext()(host_mod) + + +metal_backend = Backend( + name="metal", + target_kinds=("metal",), + import_path="tilelang.metal", + pipeline=metal_pipeline, + device_codegen=metal_device_codegen, + device_codegen_without_compile=metal_device_codegen, + host_pre_codegen=metal_host_pre_codegen, + execution_backends={ + "tvm_ffi": ExecutionBackendSpec( + name="tvm_ffi", + enable_host_codegen=True, + enable_device_compile=True, + ), + "torch": ExecutionBackendSpec(name="torch"), + }, + default_execution_backend="tvm_ffi", + cmake_name="METAL", +) + +register_backend(metal_backend, override=True) diff --git a/tilelang/rocm/__init__.py b/tilelang/rocm/__init__.py index c8c00e0480..ebfd889160 100644 --- a/tilelang/rocm/__init__.py +++ b/tilelang/rocm/__init__.py @@ -4,3 +4,4 @@ from . import pipeline # noqa: F401 from . import target # noqa: F401 from . import execution_backend # noqa: F401 +from . import backend # noqa: F401 diff --git a/tilelang/rocm/backend.py b/tilelang/rocm/backend.py new file mode 100644 index 0000000000..d12bba9c0f --- /dev/null +++ b/tilelang/rocm/backend.py @@ -0,0 +1,72 @@ +from __future__ import annotations + +import tvm_ffi +from tvm.target import Target + +from tilelang import tvm +from tilelang.backend import Backend, ExecutionBackendSpec, register_backend +from tilelang.backend.codegen import build_device_with_global_func +from tilelang.contrib import hipcc +from tilelang.env import COMPOSABLE_KERNEL_INCLUDE_DIR, TILELANG_TEMPLATE_PATH +from tilelang.rocm.pipeline import rocm_pipeline +from tilelang.rocm.target import target_get_mcpu, target_get_warp_size + +_CALLBACKS_REGISTERED = False + + +def _tilelang_callback_hip_compile(code, target): + arch = target_get_mcpu(target) + return hipcc.compile_hip( + code, + target_format="hsaco", + arch=arch, + options=[ + "-std=c++17", + "-I" + TILELANG_TEMPLATE_PATH, + "-I" + COMPOSABLE_KERNEL_INCLUDE_DIR, + ], + verbose=False, + ) + + +def register_rocm_callbacks() -> None: + global _CALLBACKS_REGISTERED + if _CALLBACKS_REGISTERED: + return + tvm_ffi.register_global_func("tilelang_callback_hip_compile", f=_tilelang_callback_hip_compile, override=True) + _CALLBACKS_REGISTERED = True + + +def rocm_device_codegen(device_mod: tvm.IRModule, target: Target) -> tvm.IRModule: + return build_device_with_global_func(device_mod, target, "target.build.tilelang_hip") + + +def rocm_device_codegen_without_compile(device_mod: tvm.IRModule, target: Target) -> tvm.IRModule: + return build_device_with_global_func(device_mod, target, "target.build.tilelang_hip_without_compile") + + +rocm_backend = Backend( + name="rocm", + target_kinds=("hip",), + import_path="tilelang.rocm", + pipeline=rocm_pipeline, + device_codegen=rocm_device_codegen, + device_codegen_without_compile=rocm_device_codegen_without_compile, + register_callbacks=register_rocm_callbacks, + features={ + "mcpu": target_get_mcpu, + "warp_size": target_get_warp_size, + }, + execution_backends={ + "tvm_ffi": ExecutionBackendSpec( + name="tvm_ffi", + enable_host_codegen=True, + enable_device_compile=True, + ), + "cython": ExecutionBackendSpec(name="cython"), + }, + default_execution_backend="tvm_ffi", + cmake_name="ROCM", +) + +register_backend(rocm_backend, override=True) From 784b6cbcbd4216cd2c91c1f22168a37647bd1b2a Mon Sep 17 00:00:00 2001 From: SiriusNEO Date: Tue, 2 Jun 2026 17:18:08 +0800 Subject: [PATCH 2/2] Address backend registry review feedback --- .../backend/test_tilelang_backend_registry.py | 64 +++++++++++++++++++ tilelang/backend/backend.py | 6 +- tilelang/backend/registry.py | 10 ++- tilelang/cuda/backend.py | 11 ++-- tilelang/jit/kernel.py | 6 +- tilelang/rocm/backend.py | 3 +- 6 files changed, 89 insertions(+), 11 deletions(-) diff --git a/testing/python/backend/test_tilelang_backend_registry.py b/testing/python/backend/test_tilelang_backend_registry.py index 2cbb859f8b..8b7a79aef5 100644 --- a/testing/python/backend/test_tilelang_backend_registry.py +++ b/testing/python/backend/test_tilelang_backend_registry.py @@ -4,7 +4,10 @@ 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(): @@ -44,6 +47,44 @@ def test_list_backends_returns_copy(): 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( @@ -95,3 +136,26 @@ def test_cuda_only_execution_backend_policy_is_backend_owned(): 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" diff --git a/tilelang/backend/backend.py b/tilelang/backend/backend.py index a509ac58b4..77f0b1cb08 100644 --- a/tilelang/backend/backend.py +++ b/tilelang/backend/backend.py @@ -2,6 +2,7 @@ from collections.abc import Callable, Mapping from dataclasses import dataclass, field +from types import MappingProxyType from tvm import IRModule from tvm.target import Target @@ -61,7 +62,7 @@ class Backend: supports_target: TargetPredicate | None = None normalize_target: TargetNormalizer | None = None is_available: AvailabilityCheck = _always_available - features: Mapping[str, FeatureQuery] = field(default_factory=dict) + features: Mapping[str, FeatureQuery] = field(default_factory=lambda: MappingProxyType({})) device_codegen: CodegenFunc | None = None device_codegen_without_compile: CodegenFunc | None = None @@ -81,10 +82,11 @@ def __post_init__(self) -> None: if not self.target_kinds: raise ValueError(f"Backend {self.name!r} must register at least one target kind") object.__setattr__(self, "target_kinds", tuple(self.target_kinds)) + object.__setattr__(self, "features", MappingProxyType(dict(self.features))) execution_backends = dict(self.execution_backends) if not execution_backends and self.jit_execution_backends: execution_backends = {name: ExecutionBackendSpec(name=name, adapter=name) for name in tuple(self.jit_execution_backends)} - object.__setattr__(self, "execution_backends", execution_backends) + object.__setattr__(self, "execution_backends", MappingProxyType(execution_backends)) object.__setattr__(self, "jit_execution_backends", tuple(execution_backends)) object.__setattr__(self, "language_modules", tuple(self.language_modules)) diff --git a/tilelang/backend/registry.py b/tilelang/backend/registry.py index 88adbf8bb6..f46441a9a1 100644 --- a/tilelang/backend/registry.py +++ b/tilelang/backend/registry.py @@ -9,6 +9,7 @@ _BACKENDS: dict[str, Backend] = {} _TARGET_INDEX: dict[str, list[str]] = {} _LAZY_IMPORTS: dict[str, str] = {} +_CALLBACKS_REGISTERED: set[str] = set() def register_backend(backend: Backend, *, override: bool = False) -> Backend: @@ -23,6 +24,9 @@ def register_backend(backend: Backend, *, override: bool = False) -> Backend: names = _TARGET_INDEX.get(kind, []) if old.name in names: names.remove(old.name) + if not names: + _TARGET_INDEX.pop(kind, None) + _CALLBACKS_REGISTERED.discard(backend.name) _BACKENDS[backend.name] = backend for kind in backend.target_kinds: @@ -47,7 +51,7 @@ def list_backends() -> dict[str, Backend]: def _ensure_loaded_for_kind(target_kind: str) -> None: - if target_kind in _TARGET_INDEX: + if _TARGET_INDEX.get(target_kind): return import_path = _LAZY_IMPORTS.get(target_kind) if import_path is not None: @@ -73,5 +77,7 @@ def resolve_backend(target: Target) -> Backend: backend = candidates[0] if not backend.is_available(): raise ValueError(f"TileLang backend {backend.name!r} is registered but unavailable") - backend.ensure_callbacks_registered() + if backend.name not in _CALLBACKS_REGISTERED: + backend.ensure_callbacks_registered() + _CALLBACKS_REGISTERED.add(backend.name) return backend diff --git a/tilelang/cuda/backend.py b/tilelang/cuda/backend.py index 1634292555..eb9cae627a 100644 --- a/tilelang/cuda/backend.py +++ b/tilelang/cuda/backend.py @@ -35,9 +35,12 @@ def target_get_warp_size(target: Target) -> int: def _is_nvrtc_available() -> bool: - from tilelang.jit.adapter.nvrtc import is_nvrtc_available + try: + from tilelang.jit.adapter.nvrtc import is_nvrtc_available + except ImportError: + return False - return bool(is_nvrtc_available) + return bool(is_nvrtc_available() if callable(is_nvrtc_available) else is_nvrtc_available) def _is_cutedsl_available() -> bool: @@ -153,12 +156,12 @@ def register_cuda_callbacks() -> None: def cuda_device_codegen(device_mod: tvm.IRModule, target: Target) -> tvm.IRModule: - global_func = "target.build.tilelang_" + ("cutedsl" if "cutedsl" in target.keys else "cuda") + global_func = "target.build.tilelang_" + ("cutedsl" if _is_cutedsl_target(target) else "cuda") return build_device_with_global_func(device_mod, target, global_func) def cuda_device_codegen_without_compile(device_mod: tvm.IRModule, target: Target) -> tvm.IRModule: - global_func = "target.build.tilelang_" + ("cutedsl" if "cutedsl" in target.keys else "cuda") + "_without_compile" + global_func = "target.build.tilelang_" + ("cutedsl" if _is_cutedsl_target(target) else "cuda") + "_without_compile" return build_device_with_global_func(device_mod, target, global_func) diff --git a/tilelang/jit/kernel.py b/tilelang/jit/kernel.py index 6583d486d2..3b32a0cb4f 100644 --- a/tilelang/jit/kernel.py +++ b/tilelang/jit/kernel.py @@ -471,7 +471,8 @@ def get_kernel_source(self, kernel_only: bool = True) -> str: str The source code of the compiled kernel function. """ - if self.execution_backend in {"cython", "nvrtc", "tvm_ffi", "cutedsl"}: + adapter_name = self.execution_backend_spec.adapter + if adapter_name in {"cython", "nvrtc", "tvm_ffi", "cutedsl"}: return self.adapter.get_kernel_source(kernel_only=kernel_only) return self.artifact.kernel_source @@ -479,7 +480,8 @@ def get_host_source(self) -> str: """ Returns the source code of the host function. """ - if self.execution_backend in {"cython", "nvrtc", "tvm_ffi", "cutedsl"}: + adapter_name = self.execution_backend_spec.adapter + if adapter_name in {"cython", "nvrtc", "tvm_ffi", "cutedsl"}: return self.adapter.get_host_source() assert self.artifact.host_mod is not None, "host_mod is not available" return str(self.artifact.host_mod) diff --git a/tilelang/rocm/backend.py b/tilelang/rocm/backend.py index d12bba9c0f..8f2513e76d 100644 --- a/tilelang/rocm/backend.py +++ b/tilelang/rocm/backend.py @@ -6,7 +6,6 @@ from tilelang import tvm from tilelang.backend import Backend, ExecutionBackendSpec, register_backend from tilelang.backend.codegen import build_device_with_global_func -from tilelang.contrib import hipcc from tilelang.env import COMPOSABLE_KERNEL_INCLUDE_DIR, TILELANG_TEMPLATE_PATH from tilelang.rocm.pipeline import rocm_pipeline from tilelang.rocm.target import target_get_mcpu, target_get_warp_size @@ -15,6 +14,8 @@ def _tilelang_callback_hip_compile(code, target): + from tilelang.contrib import hipcc + arch = target_get_mcpu(target) return hipcc.compile_hip( code,