[SiriusNEO]

Summary

• add a backend host codegen registry with lazy host-target registration
• move host build dispatch for c/llvm into tilelang.cpu.codegen
• move Metal host-context preparation into a backend-owned host codegen hook
• update backend layout docs for host/device codegen ownership

Overview

This PR refactors the backend architecture by introducing a centralized host codegen registry system that decentralizes host code generation dispatch logic. The changes move backend-specific host build logic out of the engine's lowering layer and into their respective backend packages, improving separation of concerns and maintainability.

Key Changes

New Infrastructure: Host Codegen Registry

tilelang/backend/host_codegen.py introduces a new registry system with:

• Type aliases for host codegen functions and hooks (HostCodegenFunc, HostCodegenHookFunc, TargetPredicate)
• HostCodegen dataclass to register host code generators with optional target predicates
• HostCodegenHook dataclass to register backend-owned host preparation hooks
• Registry functions with lazy-loading support:
  • register_host_codegen() and register_lazy_host_codegen() for host code generators
  • register_host_codegen_hook() and register_lazy_host_codegen_hooks() for host hooks
• Resolution functions:
  • allowed_host_codegens_for_target() to query available host codegens
  • resolve_host_codegen() to select the matching host codegen
  • apply_host_codegen_hooks() to apply all matching hooks sequentially
• Helper global_func_host_codegen() to wrap TVM global functions

Backend Integration

tilelang/backend/__init__.py exposes the new host codegen infrastructure and registers lazy entries for "c" and "llvm" host codegens and "metal" host codegen hooks.

tilelang/cpu/codegen.py now registers host code generators for both "c" and "llvm" target kinds, wiring them to their respective TVM build functions (target.build.tilelang_c_host and target.build.llvm).

tilelang/metal/codegen.py introduces _mark_host_metal_context() hook function and registers it via the new host codegen hook system for Metal-specific host context preparation, keeping host-side Metal logic self-contained within the Metal backend.

Simplified Engine Lowering

tilelang/engine/lower.py refactors the host_codegen function to use the new registry system instead of inline dispatch logic. It now:

• Applies registered host codegen hooks using apply_host_codegen_hooks()
• Delegates final host compilation to resolve_host_codegen(target_host).lower()

This eliminates hardcoded conditionals and centralizes host build dispatch logic in backend packages.

Documentation

tilelang/backend/README.md updated to clarify backend ownership boundaries, distinguishing tilelang/backend/ as shared infrastructure from tilelang/<backend>/ as backend-owned implementations. Added documentation for the host-codegen resolution flow, including the new apply_host_codegen_hooks step and how backends register host build entries and hooks.

[github-actions]

👋 Hi! Thank you for contributing to the TileLang project.

Please remember to run pre-commit run --all-files in the root directory of the project to ensure your changes are properly linted and formatted. This will help ensure your contribution passes the format check.

We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀

[coderabbitai]

📝 Walkthrough

Walkthrough

A new host_codegen.py module is added to tilelang/backend providing a registry for HostCodegen entries and HostCodegenHook entries, with lazy-import loading, override semantics, and target-predicate filtering. The engine's host_codegen function in lower.py is refactored to use apply_host_codegen_hooks and resolve_host_codegen instead of inline Metal/LLVM/C dispatch. CPU and Metal backend packages register their entries into this registry, and documentation is updated accordingly.

Changes

Host Codegen Registry System

Layer / File(s)	Summary
Host codegen registry: types, dataclasses, and public API tilelang/backend/host_codegen.py	Introduces HostCodegenFunc, HostCodegenHookFunc, TargetPredicate type aliases; global_func_host_codegen factory; HostCodegen and HostCodegenHook dataclasses with matches/lower methods; module-level registries with lazy-import bookkeeping; register_host_codegen, register_lazy_host_codegen, register_host_codegen_hook, register_lazy_host_codegen_hooks with override semantics; lazy-loading helpers; and public allowed_host_codegens_for_target, apply_host_codegen_hooks, resolve_host_codegen.
Engine lower.py: replace inline dispatch with registry tilelang/engine/lower.py	Imports apply_host_codegen_hooks and resolve_host_codegen; removes inline per-target_host conditional compilation branches (Metal MarkHostMetalContext, LLVM/C TVM global func calls, unsupported error) and replaces with a unified apply_host_codegen_hooks then resolve_host_codegen().lower() call.
CPU and Metal backend registrations and package re-exports tilelang/cpu/codegen.py, tilelang/metal/codegen.py, tilelang/backend/__init__.py	cpu/codegen.py registers HostCodegen entries for "c" and "llvm" targets; metal/codegen.py defines _mark_host_metal_context and registers a "metal_context" HostCodegenHook; backend/__init__.py re-exports all host codegen symbols and adds lazy registration calls for "c", "llvm", and "metal".
README: document host codegen registry architecture tilelang/backend/README.md	Updates Overview, Lowering Entry, file inventory, Backend Packages guidance, and Guidelines to reflect the new host codegen registry design and ownership rules.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Poem

🐇 A registry for hosts, how neat and clean,
No more inline if metal scattered between,
Each backend now owns its dispatch and hook,
Just register and match — come take a look!
The engine stays pure, the backends stay free,
Lazy-loaded at runtime, as tidy as can be. ✨

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 42.86% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title 'Dispatch host CodeGen through backend registry' directly and clearly summarizes the main change: introducing a host codegen registry system for backend packages to handle host code generation dispatch.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

🧹 Nitpick comments (1)

tilelang/backend/README.md (1)

185-188: 🧹 Nitpick | 🔵 Trivial | 💤 Low value

Vary "Keep" sentence starters to improve readability.

Lines 185–188 have three consecutive bullet points opening with "Keep," which creates repetitive phrasing. Reword one or two to improve variety while maintaining clarity.

 - Keep `tilelang/language` and `tilelang/tileop` backend-neutral.
 - Keep backend-specific pass ordering in the backend package.
-- Keep backend-specific host-codegen dispatch and host preparation hooks in the
-  backend package.
+- Backend-specific host-codegen dispatch and host preparation hooks should stay
+  in the backend package.
 - Keep backend-specific device-codegen dispatch in the backend package.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tilelang/backend/README.md` around lines 185 - 188, The three consecutive
bullet points in the backend package architecture section (covering pass
ordering, host-codegen dispatch and hooks, and device-codegen dispatch) all
begin with "Keep," creating repetitive phrasing that reduces readability. Reword
one or two of these bullet points to use alternative sentence starters or
phrasings that convey the same information while maintaining clarity and
variety, such as restructuring around the component names or using
action-oriented language instead of "Keep."

Source: Linters/SAST tools

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@tilelang/backend/README.md`:
- Around line 185-188: The three consecutive bullet points in the backend
package architecture section (covering pass ordering, host-codegen dispatch and
hooks, and device-codegen dispatch) all begin with "Keep," creating repetitive
phrasing that reduces readability. Reword one or two of these bullet points to
use alternative sentence starters or phrasings that convey the same information
while maintaining clarity and variety, such as restructuring around the
component names or using action-oriented language instead of "Keep."

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: eabaf237-b57d-44a8-8ef6-dd155f93f839

📥 Commits

Reviewing files that changed from the base of the PR and between 607a914 and 55de4b6.

📒 Files selected for processing (6)

• tilelang/backend/README.md
• tilelang/backend/__init__.py
• tilelang/backend/host_codegen.py
• tilelang/cpu/codegen.py
• tilelang/engine/lower.py
• tilelang/metal/codegen.py