
TileLang 后端架构解析以垂直切片模型解耦语言方言、Pass 流水线与执行后端【免费下载链接】tilelangDomain-specific language designed to streamline the development of high-performance GPU/CPU/Accelerators kernels项目地址: https://gitcode.com/GitHub_Trending/ti/tilelang导读TileLang 将一个目标后端target backend视为一条从语言方言language dialect到目标代码生成target code generation的完整编译器垂直切片而把 Build、加载与 Launch 交给独立的、可复用的执行后端execution backend处理。本文以 tilelang/backend/README.md 为骨架结合仓库中 模块注册与上下文解析、执行后端策略、代码生成组件、CUDA 后端清单 与 编译器入口 等源码完整讲解 TileLang 后端的所有权边界、注册机制、上下文解析流程以及新增一个后端所需遵循的目录结构与规则。读完本文你将理解目标后端与执行后端为何刻意分离以及如何为 CUDA、ROCm、CPU、Metal、WebGPU 等目标新增或扩展一个后端实现。设计模型后端的四个实现区域TileLang 把目标后端定义为拥有以下四个实现区域的垂直切片部分职责语言方言Language dialect在公共 TileLang 语言之上扩展后端特有的操作op与内联函数intrinsic。上下文贡献Context contribution定义目标检测/归一化target detection/normalization、目标归属target ownership以及共享上下文解析时使用的兼容执行后端。Pass 流水线Pass pipeline通过一条显式、由后端拥有的 Pass 序列降低公共 IR 与后端特有 IR。宿主/设备代码生成Host/device codegen将降低后的宿主host与设备deviceIR 转换为源码或运行时模块并提供目标特有的工具链钩子toolchain hooks。目标后端还会声明哪些共享执行后端可以消费它的输出。这只是一个兼容性声明而非第五个实现区域。新目标后端通常会复用tvm_ffi、nvrtc、cython、torch或其它既有执行实现而不是从零实现一套 JIT 适配器或运行时。后端整体逻辑流程如下target-backend language dialect | v frontend IR | v shared BackendContext resolution - selects target backend - selects compatible ExecutionBackend | v target-backend PassPipeline | v host/device split | ----------- HostCodegen ---- | | ----------- DeviceCodegen -- | v shared ExecutionBackend Build - Load - LaunchBackendContext在整个编译过程中只创建一次并贯穿流水线与代码生成阶段它同时记录了选中的共享执行后端。后续任何阶段都不得再次推断或解析这两个后端这正是所有权边界显式化的核心体现。关键术语目标后端target backend如 CUDA、ROCm、CPU、Metal、WebGPU拥有方言、降低、代码生成以及目标特有工具链原语。执行后端execution backend如tvm_ffi、nvrtc、cython、torch、cutedsl是可复用的 Build/JIT/Runtime 实现可为多个兼容的目标后端服务。执行兼容性声明目标后端声明兼容性并满足所选执行实现的源码、产物与启动元数据契约它通常不亲自实现另一个 JIT 适配器或运行时。BackendModule面向编译器的注册清单manifest描述当前实现所用的组件但它不是后端实现的全部。BackendContext一次编译中已解析、不可变的状态。目标后端与执行后端刻意分离例如 CUDA 可以通过tvm_ffi、nvrtc或cython执行选择nvrtc不会改变目标架构或 Pass 流水线。源码布局公共基础设施与后端自有包Python 侧的实现被划分为公共基础设施与后端自有包tilelang/backend/后端注册表、组件接口、上下文解析以及少量共享辅助工具。tilelang/backend/某个目标后端的方言、目标辅助、流水线、代码生成、工具链钩子、操作实现、内联函数与执行兼容性声明。tilelang/jit/共享 JIT 基础设施与当前执行适配器adapter适配器具体实现在 tilelang/jit/adapter/ 下包含tvm_ffi、cython、nvrtc、torch、cutedsl等适配器。原生侧镜像目标后端的所有权位于src/backend/下C 的 op 降低、代码生成、运行时模块、工具链 stubs 与后端本地 CMake 文件都在此处src/backend/则保留给共享的原生后端辅助实现。后端清单BackendModule注册即宣言每个后端的backend.py发布一个BackendModule作为编译器注册表使用的类型化清单BACKEND register_backend( BackendModule( namecuda, target_kinds(cuda,), supports_targetis_plain_cuda_target, pipelines{...}, device_codegens{...}, host_codegens{...}, host_codegen_hooks{...}, execution_backends( ExecutionBackendSpec(tvm_ffi), ExecutionBackendSpec(nvrtc), ), callbacks{...}, ) )清单当前声明自身名称与所拥有的 TVM target kind一个可选的 target 谓词predicate用于区分同 kind 下的变体为每个所拥有的 target kind 提供恰好一条Pass 流水线与一个设备代码生成入口可选的宿主代码生成入口与代码生成前钩子pre-codegen hooks按auto偏好顺序排列的兼容共享执行后端供后端校验或目标工具链集成使用的 FFI 回调。对应实现位于 tilelang/backend/module.pyBackendModule的__post_init__会对声明做严格校验名称非空、target kinds 唯一、每个 target kind 必须恰好对应一条流水线且流水线名与 target kind 一致、必须为每个 target kind 提供设备代码生成入口、至少声明一个执行后端、若某个执行后端启用宿主代码生成则必须提供宿主代码生成目标、回调名非空等。任何一项不满足都会在导入期抛出ValueError。register_backend()tilelang/backend/module.py会完成完整校验并发布一次它还维护_TARGET_KIND_INDEX索引并强制约束多个后端共享同一 target kind 时各自必须定义supports_target谓词否则视为歧义拒绝注册。此外注册时会把清单中的 FFI 回调通过tvm_ffi.register_global_func注册到全局函数表中。后端包在 TileLang 导入期间即完成初始化因此注册表不维护导入路径、加载状态或同步。多个清单可以共享一个 TVM target kind——只要谓词能让变体无歧义。CUDA 与 CuTeDSL 是两个独立清单都匹配cudatarget kind它们显式复用 CUDA 流水线但声明了不同的设备代码生成与执行兼容性。流水线复用必须是显式的绝不能来自引擎中基于 target 的分支。后端上下文准备一次解析全程复用公共编译器或 JIT 入口创建一个不可变的BackendContextcontext create_backend_context(target, target_host, execution_backend)上下文准备执行三步操作见 tilelang/backend/module.py 的create_backend_context与 tilelang/backend/target.py 的determine_target归一化设备目标与宿主目标targetauto时优先取Target.current()否则依次调用已注册的 target detector 自动探测宿主目标未指定时默认llvmLLVM 可用时否则c。精确选择一个BackendModule依据target.kind.name必要时再用supports_target谓词过滤多个候选匹配时直接报错。解析一个可用的ExecutionBackendSpec用户显式指定或按后端声明的auto顺序选择第一个可用项。得到的上下文绑定BackendContext module selected target-backend manifest target normalized device target target_host normalized host target execution_backend selected shared Build/JIT/Runtime implementation缓存、降低、代码生成与 JIT 代码全部传递同一个上下文实例。后端特有的 target 解析与规范化属于后端包后端选择本身留在共享上下文工厂中。BackendContext还提供了lower、codegen_device、preprocess_host_codegen、codegen_host等便捷方法并在构造时校验执行后端确实由该后端模块声明且执行后端匹配该 target见 tilelang/backend/module.py。codegen_device未显式传compile_device时会回落到execution_backend.enable_device_compile策略值。语言方言按导入路径静态选择tilelang/language定义后端无关的语言表面language surface。后端可以在tilelang/backend/language下自由组合扩展from tilelang import language as T # common CUDA compatibility facade from tilelang.cuda import language as T # common CUDA extensions from tilelang.rocm import language as T # common ROCm extensions方言选择是静态的、跟随导入路径的不依赖进程级可变方言注册表。后端特有的语言 API 会发出对应后端流水线知道如何降低的 IR 操作或注解语言层只负责描述语义、构造 IR而不执行目标代码生成。例如 CUDA 拥有 WGMMA、TCGEN05 语言辅助与内联函数发射器ROCm 拥有 MFMA 与 WMMA 辅助。使用后端特有符号的库代码、测试与示例应优先使用显式后端语言导入以便静态分析与自动补全识别方言。注意语言方言是逻辑后端的一部分即使它并不是当前BackendModule清单的一个字段。Pass 流水线后端显式拥有完整降低序列每个目标后端必须在共享语义检查之后显式拥有完整的降低序列PreLowerSemanticCheck(mod) # shared frontend boundary mod context.lower(mod) # selected backend pipeline在 tilelang/engine/lower.py 的lower_to_host_device_ir中可以看到这条边界的实际落点先执行后端无关的PreLowerSemanticCheck(mod)再调用context.lower(mod)最后用Filter按宿主/设备调用约定把模块拆分为host_mod与device_mod。完整的编译器入口负责 Pass 插桩会话instrumentation session。BackendContext.lower与PassPipeline.lower从不会隐式创建会话只有当会话已激活时流水线才贡献其后端特有的阶段作用域。PassPipelinetilelang/backend/pass_pipeline/pipeline.py在无会话时也会正常跑完降低流水线只是不产生开发者工具插桩它还统一在异常处调用enrich_error把 IR 节点携带的-- file:line:col标记注入诊断。有序 Pass 列表位于tilelang/backend/pipeline.py。后端专用 Pass 必须在那里被调用而不是从tilelang/engine/lower.py分发。后端流水线可以使用 tilelang/backend/pass_pipeline 下的小型共享辅助如布局可视化、向量化开关、共享内存复用等见 pipeline_utils.py但排序与目标特有选择必须可见于后端包中。目标变体可以在 IR 降低完全一致时显式引用另一个后端的流水线例如 CuTeDSL 当前就引用pipeline.CUDA_PIPELINE见 tilelang/cuda/cutedsl_backend.py。宿主与设备代码生成流水线之后编译器把模块拆分为宿主 IR 与设备 IR两条代码生成路径都通过BackendContext选择device_mod context.codegen_device(device_mod, compile_device...) host_mod context.preprocess_host_codegen(host_mod) host_mod context.codegen_host(host_mod)DeviceCodegentilelang/backend/device_codegen.py在后端同时支持时提供已编译与仅源码两个入口build与build_without_compile由compile_device参数选择未提供某个模式时抛出明确报错。global_func_device_codegen可以把 TVM 全局函数包装成设备代码生成回调并自动套上run_codegen_with_instrumentation插桩。HostCodegen依据具体宿主 target 选择通常是c或llvm仓库在 tilelang/backend/host_codegen.py 提供了STANDARD_HOST_CODEGENS分别绑定target.build.tilelang_c_host与target.build.llvm两个全局函数。HostCodegenHook允许设备后端在宿主代码生成前准备宿主 IR而无需在共享引擎中添加 target 检查——Metal 正是用此类钩子标记需要 Metal 运行时上下文的函数。引擎内部不得出现基于target.kind.name的分发表来分发后端代码生成CUDA 与 CuTeDSL 等变体在自己的清单中声明各自代码生成入口。实际编译路径中device_codegen会先做LowerIntrin、Simplify、HoistBroadcastValues等预备变换再调用context.codegen_devicehost_codegen则依次执行BindTarget、FP8StorageLegalize、BF16StorageLegalize、LowerTVMBuiltin、LowerCustomDatatypes、LowerIntrin、CombineContextCall然后才进入context.preprocess_host_codegen与context.codegen_host见 tilelang/engine/lower.py。执行后端兼容性声明契约而非重新实现目标后端通常不实现 JIT 或运行时行为只声明哪些共享执行后端可以消费其生成的输出execution_backends: tvm_ffi - compatible with TVM FFI runtime modules nvrtc - compatible with CUDA source driver launch cython - compatible with generated Cython wrappers torch - compatible with framework-provided compile/launchExecutionBackendSpectilelang/backend/execution_backend.py当前记录四个字段执行后端名称、可用性检查is_available默认恒可用、target 谓词supports_target、是否需要宿主代码生成enable_host_codegen、是否需要即时设备编译enable_device_compile。兼容的目标后端必须提供所选实现期望的代码生成模式与输出契约按需提供已编译或仅源码的设备代码生成按需提供宿主代码生成提供源码、产物、全局符号、参数与启动元数据提供目标特有校验、编译器回调与工具链辅助。共享执行后端负责缓存编排、wrapper 创建、产物加载、参数/流绑定、启动与运行时生命周期。它可以在 Build 期间调用目标拥有的编译器回调但这些回调不会让 JIT/Runtime 变成目标后端的实现区域。CUDA 的实际声明见 tilelang/cuda/execution_backend.pyCUDA_EXECUTION_BACKENDS依次为tvm_ffienable_host_codegenTrue、enable_device_compileTrue、nvrtc带可用性检查与cythonCuTeDSL 变体则单独声明cutedsl带check_cutedsl_available可用性检查。这也印证了原文档选择nvrtc不会改变目标架构或 Pass 流水线的论断——执行后端只是消费产物。只有当没有任何既有实现能消费目标输出或驱动其运行时 API 时才应新增一个执行后端——这是一个独立的集成任务。具体适配器当前位于 tilelang/jit/adapter适配器分发仍发生在共享 JIT 层见 tilelang/jit/kernel.py其中按tvm_ffi/cython/nvrtc/torch/cutedsl分别选用不同适配器新的执行行为应放到该共享接口之后而不是塞进目标 Pass 或代码生成分发中。部分 CodeGen 与 Build 仍组合在原生target.build.*函数以及历史遗留的DeviceCodegen.build名称之后这应被视为边界处的实现细节目标代码生成提供操作选中的执行后端决定是否以及何时调用它。已注册目标后端一览Python 包Target kind说明tilelang/cuda/backend.pycuda普通 CUDA 代码生成、编译器回调与执行兼容性。tilelang/cuda/cutedsl_backend.pycudaCuTeDSL 变体显式复用 CUDA 流水线。tilelang/rocmhipROCm/HIP 流水线、代码生成、编译器回调与 MFMA/WMMA 扩展。tilelang/cpuc、llvmCPU 流水线、代码生成与标量 CPU tile-op 实现。tilelang/metalmetalMetal 流水线、代码生成、宿主钩子与 Metal 语言扩展。tilelang/webgpu/backend.pywebgpuWebGPU 编译器组件注册。以 CUDA 清单为实例tilelang/cuda/backend.py它声明namecuda、target_kinds(cuda,)、supports_targetcodegen.is_plain_cuda_target流水线与设备代码生成都绑定到cuda并注册了两个 FFI 回调——tilelang_callback_cuda_validate校验T.CUDASourceCodeKernel等外部 CUDA 源码内核与global_symbol一致性与tilelang_callback_cuda_compile通过nvcc编译、处理fast_math、ptxas-options、-stdc20、模板与 CUTLASS 头文件路径并经由CUDABinaryCache做二进制缓存。公共后端基础设施小而聚焦tilelang/backend应保持精简只包含共享接口与管道不包含目标特有实现tilelang/backend/ __init__.py module.py target.py device_codegen.py host_codegen.py execution_backend.py pass_pipeline/ __init__.py pipeline.py pipeline_utils.pymodule.py定义BackendModule、BackendContext、注册、校验与上下文解析。target.py提供公共 target 检测与归一化管道register_target_detector、register_target_normalizer、auto_detect_target、determine_target。pass_pipeline/pipeline.py定义PassPipeline。device_codegen.py 与 host_codegen.py定义代码生成组件类型与共享全局函数辅助。execution_backend.py定义当前执行后端选择与能力描述符。pass_pipeline/pipeline_utils.pyPass 配置、可视化、向量化开关与共享内存复用等小工具。后端包布局与原生布局一个典型目标后端包的形状如下tilelang/backend/ __init__.py backend.py manifest and backend toolchain callbacks target.py target parsing and normalization helpers execution_backend.py compatible execution paths and auto preference language/ backend language dialect pipeline.py complete lowering sequence codegen.py host/device codegen entries transform/ backend-only Python passes op/ tile-op implementations and registration intrinsics/ backend intrinsic emitters and helpers并非每个后端都需要全部文件。组件文件定义实现backend.py是唯一组装并注册清单的地方。导入期注册必须保持确定且轻量。原生侧对应关系src/ cpu/ cuda/ metal/ rocm/ webgpu/ src/backend/ common/典型的后端本地子目录op/原生 tile-op 降低辅助transform/原生后端专用 Passcodegen/目标代码生成、工具链入口与运行时模块集成stubs/可选的懒加载驱动/运行时 stubsCMakeLists.txt后端本地源码选择与工具链配置。不依赖任何目标运行时特性的共享原生辅助应放在src/backend/common。所有权规则新增后端的行动清单保持公共语言表面与共享语义检查的后端无关性。后端语言扩展放在tilelang/backend/language并在该后端流水线中降低。在创建BackendContext时一次性解析 target 与执行状态。后端完整的 Pass 顺序保留在tilelang/backend/pipeline.py。宿主/设备代码生成分发、宿主准备钩子、编译器回调与目标特有工具链辅助归目标后端所有。目标后端声明执行兼容性而不是重写 JIT/Runtime只有出现新的产物或运行时契约时才新增执行后端。不要把执行模式检查散布到编译器 Pass 中。后端注册保持显式不要从目录名推断 target 归属——例如tilelang/rocm拥有的是 TVMhiptarget kind。清单在常规包初始化期间注册不维护第二套懒加载注册表。实现新后端的编码代理必须遵循tilelang-backend技能.agents/skills/tilelang-backend/SKILL.md。结合上述规则一个后端从方言到执行兼容性声明的全链路——注册清单、上下文解析、流水线拥有、双路代码生成、执行后端契约——即构成一条完整、自洽且边界清晰的编译器垂直切片。【免费下载链接】tilelangDomain-specific language designed to streamline the development of high-performance GPU/CPU/Accelerators kernels项目地址: https://gitcode.com/GitHub_Trending/ti/tilelang创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考