tilelang.jit

This module provides an auto-tuning infrastructure for TileLang (tl) programs. It includes functionality to JIT-compile TileLang programs into a runnable kernel adapter using TVM.

Submodules

• tilelang.jit.adapter
• tilelang.jit.env
• tilelang.jit.kernel
• tilelang.jit.param

Attributes

logger

Functions

compile([func, out_idx, execution_backend, target, ...])	Compile the given TileLang PrimFunc with TVM and build a JITKernel.
jit([func, out_idx, target, target_host, ...])	Just-In-Time (JIT) compiler decorator for TileLang functions.

Package Contents

tilelang.jit.logger

tilelang.jit.compile(func=None, out_idx=None, execution_backend='cython', target='auto', target_host=None, verbose=False, pass_configs=None, compile_flags=None)

Compile the given TileLang PrimFunc with TVM and build a JITKernel. :param func: The TileLang TIR function to compile and wrap. :type func: tvm.tir.PrimFunc, optional :param out_idx: Index(es) of the output tensors to return (default: None). :type out_idx: Union[List[int], int], optional :param execution_backend: Execution backend to use for kernel execution (default: “cython”). :type execution_backend: Literal[“dlpack”, “ctypes”, “cython”, “nvrtc”], optional :param target: Compilation target, either as a string or a TVM Target object (default: “auto”). :type target: Union[str, Target], optional :param target_host: Target host for cross-compilation (default: None). :type target_host: Union[str, Target], optional :param verbose: Whether to enable verbose output (default: False). :type verbose: bool, optional :param pass_configs: Additional keyword arguments to pass to the Compiler PassContext.

Refer to tilelang.transform.PassConfigKey for supported options.

Parameters:

• func (tvm.tir.PrimFunc)

• out_idx (list[int] | int | None)

• execution_backend (Literal['dlpack', 'ctypes', 'cython', 'nvrtc'])

• target (str | tvm.target.Target)

• target_host (str | tvm.target.Target | None)

• verbose (bool)

• pass_configs (dict, optional)

• compile_flags (list[str] | str | None)

Return type:

kernel.JITKernel

tilelang.jit.jit(func=None, *, out_idx=None, target='auto', target_host=None, execution_backend='cython', verbose=False, pass_configs=None, debug_root_path=None, compile_flags=None)

Just-In-Time (JIT) compiler decorator for TileLang functions.

This decorator can be used without arguments (e.g., @tilelang.jit):

Applies JIT compilation with default settings.

Parameters:

• func_or_out_idx (Any, optional) – If using @tilelang.jit(…) to configure, this is the out_idx parameter. If using @tilelang.jit directly on a function, this argument is implicitly the function to be decorated (and out_idx will be None).

• target (Union[str, Target], optional) – Compilation target for TVM (e.g., “cuda”, “llvm”). Defaults to “auto”.

• target_host (Union[str, Target], optional) – Target host for cross-compilation. Defaults to None.

• execution_backend (Literal["dlpack", "ctypes", "cython", "nvrtc"], optional) – Backend for kernel execution and argument passing. Defaults to “cython”.

• verbose (bool, optional) – Enables verbose logging during compilation. Defaults to False.

• pass_configs (Optional[Dict[str, Any]], optional) – Configurations for TVM’s pass context. Defaults to None.

• debug_root_path (Optional[str], optional) – Directory to save compiled kernel source for debugging. Defaults to None.

• func (Callable[param._P, param._RProg] | tvm.tir.PrimFunc | None)

• out_idx (Any)

• compile_flags (list[str] | str | None)

Returns:

Either a JIT-compiled wrapper around the input function, or a configured decorator instance that can then be applied to a function.

Return type:

Callable