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

Attributes

logger

Classes

JITImpl

Abstract base class for generic types.

Functions

compile([func, out_idx, execution_backend, target, ...])

Compile the given TileLang PrimFunc with TVM and build a JITKernel.

par_compile(funcs[, out_idx, execution_backend, ...])

Parallel compile multiple TileLang PrimFunc with TVM and build JITKernels.

jit(…)

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

Package Contents

tilelang.jit.logger
tilelang.jit.compile(func=None, out_idx=None, execution_backend='auto', 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. Use “auto” to pick a sensible

default per target (cuda->tvm_ffi, metal->torch, others->cython).

Parameters:

  • target (Union[str, Target], optional) – Compilation target, either as a string or a TVM Target object (default: “auto”).

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

  • verbose (bool, optional) – Whether to enable verbose output (default: False).

  • pass_configs (dict, optional) – Additional keyword arguments to pass to the Compiler PassContext. Refer to tilelang.transform.PassConfigKey for supported options.

  • func (tilelang.language.v2.PrimFunc[_KP, _T])

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

  • execution_backend (Literal["auto", "dlpack", "tvm_ffi", "ctypes", "cython", "nvrtc", "torch"], optional)

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

Return type:

kernel.JITKernel[_KP, _T]

tilelang.jit.par_compile(funcs, out_idx=None, execution_backend='auto', target='auto', target_host=None, verbose=False, pass_configs=None, compile_flags=None, num_workers=None, ignore_error=False)

Parallel compile multiple TileLang PrimFunc with TVM and build JITKernels. :param funcs: The TileLang TIR functions to compile and wrap. :type funcs: Iterable[tvm.tir.PrimFunc] :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. Use “auto” to pick a sensible

default per target (cuda->tvm_ffi, metal->torch, others->cython).

Parameters:

  • target (Union[str, Target], optional) – Compilation target, either as a string or a TVM Target object (default: “auto”).

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

  • verbose (bool, optional) – Whether to enable verbose output (default: False).

  • pass_configs (dict, optional) – Additional keyword arguments to pass to the Compiler PassContext. Refer to tilelang.transform.PassConfigKey for supported options.

  • funcs (collections.abc.Iterable[tilelang.language.v2.PrimFunc[_KP, _T]])

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

  • execution_backend (Literal["auto", "dlpack", "tvm_ffi", "ctypes", "cython", "nvrtc", "torch"], optional)

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

  • num_workers (int)

  • ignore_error (bool)

Return type:

list[kernel.JITKernel[_KP, _T]]

class tilelang.jit.JITImpl

Bases: Generic[_P, _KP, _T]

Abstract base class for generic types.

A generic type is typically declared by inheriting from this class parameterized with one or more type variables. For example, a generic mapping type might be defined as:

class Mapping(Generic[KT, VT]):
    def __getitem__(self, key: KT) -> VT:
        ...
    # Etc.

This class can then be used as follows:

def lookup_name(mapping: Mapping[KT, VT], key: KT, default: VT) -> VT:
    try:
        return mapping[key]
    except KeyError:
        return default
func: Callable[_P, _T] | tilelang.language.v2.PrimFunc[_KP, _T]
out_idx: list[int] | int | None
execution_backend: Literal['auto', 'dlpack', 'tvm_ffi', 'ctypes', 'cython', 'nvrtc', 'torch']
target: str | tvm.target.Target
target_host: str | tvm.target.Target
verbose: bool
pass_configs: dict[str, Any] | None
debug_root_path: str | None
compile_flags: list[str] | str | None
func_source: str
signature: inspect.Signature
__post_init__()
get_tir(*args, **kwargs)
Parameters:

  • args (_P)

  • kwargs (_P)

Return type:

tilelang.language.v2.PrimFunc[_KP, _T]

par_compile(configs, num_workers=None, ignore_error=False)
Parameters:

  • configs (collections.abc.Iterable[dict[str, Any] | tuple[str, Any]])

  • num_workers (int)

  • ignore_error (bool)

Return type:

list[kernel.JITKernel[_KP, _T]]

compile(*args, **kwargs)
Parameters:

  • args (_P)

  • kwargs (_P)

Return type:

kernel.JITKernel[_KP, _T]

__call__(*args, **kwargs)
Parameters:

  • args (_P)

  • kwargs (_P)

Return type:

kernel.JITKernel[_KP, _T]

tilelang.jit.jit(func: Callable[_P, tilelang.language.v2.PrimFunc[_KP, _T]]) → JITImpl[_P, _KP, _T]
tilelang.jit.jit(*, out_idx: Any = None, target: str | tvm.target.Target = 'auto', target_host: str | tvm.target.Target = None, execution_backend: Literal['auto', 'dlpack', 'tvm_ffi', 'ctypes', 'cython', 'nvrtc', 'torch'] = 'auto', verbose: bool = False, pass_configs: dict[str, Any] | None = None, debug_root_path: str | None = None, compile_flags: list[str] | str | None = None) → Callable[[Callable[_P, tilelang.language.v2.PrimFunc[_KP, _T]]], JITImpl[_P, _KP, _T]]

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["auto", "dlpack", "tvm_ffi", "ctypes", "cython", "nvrtc", "torch"], optional) – Backend for kernel execution and argument passing. Use “auto” to pick a sensible default per target (cuda->tvm_ffi, metal->torch, others->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.

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