TileLang Kernel 编译与执行流程
1. 概述
本文档详细介绍了 TileLang 中从 Kernel 定义到编译为目标设备二进制代码，再到最终执行的完整流程。该流程主要包括以下几个关键阶段：
- Kernel 定义与装饰
- JIT 编译触发
- 中间表示（IR）生成
- IR 优化与代码生成
- 二进制代码编译
- 执行与结果返回
2. 详细流程分析
2.1 Kernel 定义与装饰
2.1.1 用户定义 Kernel
用户通过 TileLang 提供的 API 定义计算 Kernel，通常使用 @tilelang.jit 装饰器来标记需要 JIT 编译的函数：
import tilelang as tl
@tl.jit
@tl.prim_func
def matmul(A: tl.tensor, B: tl.tensor, C: tl.tensor):
# Kernel 实现代码
...
2.1.2 JIT 装饰器实现
@tilelang.jit 装饰器的实现位于 /tilelang/tilelang/jit/__init__.py 中，它是整个编译流程的入口点。
装饰器主要完成以下工作：
- 解析用户提供的编译配置参数（如目标设备、执行后端等）
- 创建
JITImpl实例，封装用户函数和编译配置 - 返回一个可调用对象，替代原始函数
关键代码：
def jit(func: Callable[_P, _T] | PrimFunc | None = None, *, ...):
def decorator(func: Callable[_P, _T]) -> JITImpl[_P, _T]:
if isinstance(func, (PrimFunc, PrimFuncCreater)):
orig_func = func.orig_func
else:
orig_func = func
return JITImpl(
func=func,
out_idx=out_idx,
execution_backend=execution_backend,
target=target,
target_host=target_host,
verbose=verbose,
pass_configs=pass_configs,
debug_root_path=debug_root_path,
compile_flags=compile_flags,
func_source=inspect.getsource(orig_func),
signature=inspect.signature(orig_func),
lazy_jit=False,
)
if func is not None:
return decorator(func)
else:
return decorator
2.2 JIT 编译触发
当用户调用被装饰的函数时，JITImpl.__call__ 方法会被触发，开始 JIT 编译流程：
def __call__(self, *args: _P.args, **kwargs: _P.kwargs) -> _Ret:
# 解析缓存键
key = self.parse_cache_key(*args, **kwargs)
# 检查缓存中是否已有编译结果
kernel = self._kernel_cache.get(key, None)
if kernel is None:
# 编译函数
kernel = self.compile(*args, **kwargs)
# 缓存编译结果
self._kernel_cache[key] = kernel
# 执行并返回结果
return kernel
2.3 中间表示（IR）生成
在 compile 方法中，首先调用 get_tir 方法获取 TIR（Tensor Intermediate Representation）格式的 PrimFunc：
def compile(self, *args: _P.args, **kwargs: _P.kwargs) -> _Ret:
# 获取 TIR PrimFunc
func = self.get_tir(*args, **kwargs)
# 编译 PrimFunc
kernel_result = compile(
func,
out_idx=self.out_idx,
execution_backend=self.execution_backend,
target=self.target,
target_host=self.target_host,
verbose=self.verbose,
pass_configs=self.pass_configs,
compile_flags=self.compile_flags,
)
# 返回编译结果
return kernel_result
get_tir 方法根据函数类型生成或返回 TIR PrimFunc：
def get_tir(self, *args: _P.args, **kwargs: _P.kwargs) -> PrimFunc[_KP, _T]:
if isinstance(self.func, PrimFuncCreater):
# 如果是 PrimFunc 创建器，则调用生成 PrimFunc
tir = self.func(*args, **kwargs)
elif isinstance(self.func, PrimFunc):
# 如果已经是 PrimFunc，则直接返回
tir = self.func
elif callable(self.func):
# 如果是可调用对象，则调用生成 PrimFunc
tir = self.func(*args, **kwargs)
else:
raise ValueError(f"Invalid function type: {type(self.func)}")
assert isinstance(tir, PrimFunc), f"target function must be a PrimFunc but got {type(tir)}"
return tir
2.4 IR 优化与代码生成
2.4.1 compile 函数调用
compile 函数（位于 tilelang/tilelang/jit/__init__.py）是编译的核心入口，它调用 tilelang.lower 来生成编译产物：
def compile(func: PrimFunc[_KP, _T] = None, ...) -> JITKernel[_KP, _T]:
# 解析执行后端
execution_backend = resolve_execution_backend(requested_backend, target)
# 缓存编译结果
return cached(
func=func,
out_idx=out_idx,
execution_backend=execution_backend,
target=target,
target_host=target_host,
verbose=verbose,
pass_configs=pass_configs,
compile_flags=compile_flags,
)
2.4.2 tilelang.lower 调用
tilelang.lower 函数（位于 tilelang/tilelang/engine/lower.py）是 IR 优化和代码生成的核心函数：
def lower(func_or_mod: tir.PrimFunc | tvm.IRModule, ...) -> CompiledArtifact:
# 将函数转换为 IRModule
if isinstance(func_or_mod, tir.PrimFunc):
func = func_or_mod
params = extrac_params(func) if not runtime_only else None
mod = tvm.IRModule({func.attrs["global_symbol"]: func})
# 确定编译目标
if isinstance(target, str):
target = determine_target(target)
# 规范化目标主机
target_host = canon_target_host(target, target_host)
# 语义检查
PreLowerSemanticCheck(mod)
# 阶段 1: IR 降低与合法化
mod = LowerAndLegalize(mod, target)
# 阶段 2: 针对目标设备优化 IR
mod = OptimizeForTarget(mod, target)
# 分离主机和设备代码
host_mod = tir.transform.Filter(_is_host_call)(mod)
device_mod = tir.transform.Filter(_is_device_call)(mod)
# 设备代码生成
codegen_mod = device_codegen(device_mod, target) if enable_device_compile else device_codegen_without_compile(device_mod, target)
# 主机代码生成
if enable_host_codegen:
host_mod = host_codegen(host_mod, target_host)
host_mod.import_module(codegen_mod)
return CompiledArtifact(host_mod, device_mod, params, codegen_mod.inspect_source(), rt_mod=host_mod)
return CompiledArtifact(host_mod, device_mod, params, codegen_mod.inspect_source())
2.4.3 host_codegen 函数
host_codegen 函数（位于 tilelang/tilelang/engine/lower.py）负责主机端代码的优化和生成：
def host_codegen(host_mod: tvm.IRModule, target_host: Target) -> tvm.IRModule:
# 将目标主机信息绑定到 IRModule 上
host_mod = tir.transform.BindTarget(target_host)(host_mod)
# 对 FP8 存储格式进行合法化转换
host_mod = tir.transform.FP8StorageLegalize()(host_mod)
# 对 BF16 存储格式进行合法化转换
host_mod = tir.transform.BF16StorageLegalize()(host_mod)
# 降低 TVM 内建函数调用为更底层的表示
host_mod = tir.transform.LowerTVMBuiltin()(host_mod)
# 降低自定义数据类型
host_mod = tir.transform.LowerCustomDatatypes()(host_mod)
# 应用 tilelang 特有的内联函数降低
host_mod = tilelang.transform.LowerIntrin()(host_mod)
# 降低设备存储访问信息
host_mod = tilelang.transform.LowerDeviceStorageAccessInfo()(host_mod)
# 合并上下文调用
host_mod = tir.transform.CombineContextCall()(host_mod)
# 根据目标主机类型选择对应的构建函数
if target_host.kind.name == "llvm":
host_mod = tvm.ffi.get_global_func("target.build.llvm")(host_mod, target_host)
elif target_host.kind.name == "c":
host_mod = tvm.ffi.get_global_func("target.build.tilelang_c")(host_mod, target_host)
else:
raise ValueError(f"Target host {target_host.kind.name} is not supported")
return host_mod
2.5 二进制代码编译
在 _compile_and_create_adapter 方法（位于 tilelang/tilelang/jit/kernel.py）中，根据执行后端创建相应的 Kernel 适配器：
def _compile_and_create_adapter(self, tilelang_func: PrimFunc, out_idx: list[int]) -> BaseKernelAdapter:
# 编译函数，生成 CompiledArtifact
enable_host_codegen = execution_backend == "tvm_ffi"
enable_device_compile = execution_backend == "tvm_ffi"
with tvm.transform.PassContext(opt_level=3, config=pass_configs), self.target:
artifact = tilelang.lower(
tilelang_func,
target=target,
target_host=target_host,
enable_host_codegen=enable_host_codegen,
enable_device_compile=enable_device_compile,
)
# 创建适配器
if execution_backend == "tvm_ffi":
# 使用 TVMFFI 后端
adapter = TVMFFIKernelAdapter(
params=artifact.params,
result_idx=out_idx,
target=target,
func_or_mod=tilelang_func,
host_mod=artifact.host_mod,
device_mod=artifact.device_mod,
rt_mod=artifact.rt_mod,
device_kernel_source=artifact.kernel_source,
verbose=verbose,
pass_configs=pass_configs,
compile_flags=compile_flags,
)
elif execution_backend == "ctypes":
# 使用 CTypes 后端
...
elif execution_backend == "cython":
# 使用 Cython 后端
...
return adapter
2.6 执行与结果返回
编译完成后，生成的 Kernel 适配器会被缓存，当用户再次调用相同的 Kernel 时，可以直接使用缓存的适配器执行，无需重新编译。
执行过程由 Kernel 适配器的 func 方法处理，根据不同的执行后端，将参数传递给编译好的二进制代码，并返回执行结果。
3. 关键组件与接口
3.1 JITImpl 类
JITImpl 类是 JIT 编译的核心组件，它封装了用户函数、编译配置和编译结果，提供了以下关键方法：
__call__: 触发 JIT 编译和执行get_tir: 生成或获取 TIR PrimFunccompile: 编译 PrimFunc 为可执行 Kernelpar_compile: 并行编译多个配置
3.2 CompiledArtifact 类
CompiledArtifact 类封装了编译结果，包括：
host_mod: 主机端 IRModuledevice_mod: 设备端 IRModuleparams: Kernel 参数信息kernel_source: 生成的设备端源代码rt_mod: 运行时模块（如果启用了主机代码生成）
3.3 Kernel 适配器
根据执行后端的不同，TileLang 提供了多种 Kernel 适配器：
TVMFFIKernelAdapter: 使用 TVM FFI 接口执行CtypesKernelAdapter: 使用 CTypes 接口执行CythonKernelAdapter: 使用 Cython 接口执行NVRTCKernelAdapter: 使用 NVRTC 接口执行TorchKernelAdapter: 使用 PyTorch 接口执行
4. 执行流程图
用户定义 Kernel
│
▼
@tilelang.jit 装饰器
│
▼
创建 JITImpl 实例
│
▼
用户调用装饰后的函数
│
▼
JITImpl.__call__()
│
▼
JITImpl.get_tir() → 获取 TIR PrimFunc
│
▼
JITImpl.compile()
│
▼
compile() 函数
│
▼
tilelang.lower()
│
├───────────────────┐
▼ ▼
阶段 1: IR 降低与合法化 阶段 2: 针对目标设备优化
│ │
▼ ▼
分离主机和设备代码 分离主机和设备代码
│ │
▼ ▼
host_mod device_mod
│ │
▼ ▼
如果 enable_host_codegen → host_codegen()
│
▼
device_codegen() 或 device_codegen_without_compile()
│
▼
codegen_mod
│ │
└───────────────────┘
│
▼
创建 CompiledArtifact
│
▼
创建 Kernel 适配器
│
▼
缓存并返回 Kernel 适配器
│
▼
执行 Kernel 并返回结果
5. 优化与加速技术
5.1 编译缓存
TileLang 使用缓存机制避免重复编译相同的 Kernel，提高执行效率。缓存键通常基于以下因素：
- Kernel 函数定义
- 输入参数类型和形状
- 编译目标设备
- 编译配置参数
5.2 多后端支持
TileLang 支持多种执行后端，根据目标设备和使用场景选择最合适的后端：
tvm_ffi: 适用于 CUDA 等 GPU 设备ctypes: 适用于通用 CPU 设备cython: 适用于需要高性能的 CPU 设备nvrtc: 适用于 NVIDIA GPU 设备torch: 适用于与 PyTorch 集成的场景
5.3 并行编译
对于需要编译多个 Kernel 或同一 Kernel 的多个配置的场景，TileLang 提供了并行编译功能，可以显著提高编译效率：
def par_compile(funcs, ...) -> list[JITKernel]:
with concurrent.futures.ThreadPoolExecutor(num_workers, "tl-par-comp") as executor:
futures = []
for func in funcs:
future = executor.submit(compile, func, ...)
futures.append(future)
# 收集编译结果
...
return results
6. 总结
TileLang 的 Kernel 编译与执行流程是一个复杂但高效的系统，它将高级的 TileLang 代码逐步转换为目标设备的二进制代码，并提供了多种优化和加速技术。该流程的设计充分考虑了灵活性和性能，支持多种目标设备和执行后端，可以满足不同场景的需求。
通过理解这个完整流程，用户可以更好地利用 TileLang 的功能，编写高效的计算 Kernel，并根据实际需求选择合适的编译配置和执行后端。
浙公网安备 33010602011771号