tilelang.jit ============ .. py:module:: tilelang.jit .. autoapi-nested-parse:: 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 ---------- .. toctree:: :maxdepth: 1 /autoapi/tilelang/jit/adapter/index /autoapi/tilelang/jit/env/index /autoapi/tilelang/jit/exceptions/index /autoapi/tilelang/jit/execution_backend/index /autoapi/tilelang/jit/kernel/index /autoapi/tilelang/jit/param/index Attributes ---------- .. autoapisummary:: tilelang.jit.logger tilelang.jit.ExecutionBackend Classes ------- .. autoapisummary:: tilelang.jit.JITImpl Functions --------- .. autoapisummary:: tilelang.jit.compile tilelang.jit.par_compile tilelang.jit.jit Package Contents ---------------- .. py:data:: logger .. py:function:: compile(func = None, out_idx = None, execution_backend = None, target = None, target_host = None, verbose = None, 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. If None, reads from TILELANG_EXECUTION_BACKEND environment variable (defaults to "auto"). :type execution_backend: Literal["auto", "dlpack", "tvm_ffi", "cython", "nvrtc", "torch", "cutedsl"], optional :param target: Compilation target, either as a string or a TVM Target object. If None, reads from TILELANG_TARGET environment variable (defaults to "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. If None, reads from TILELANG_VERBOSE environment variable (defaults to 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. :type pass_configs: dict, optional :param Environment Variables: :param ---------------------: :param TILELANG_TARGET: Default compilation target (e.g., "cuda", "llvm"). Defaults to "auto". :type TILELANG_TARGET: str :param TILELANG_EXECUTION_BACKEND: Default execution backend. Defaults to "auto". :type TILELANG_EXECUTION_BACKEND: str :param TILELANG_VERBOSE: Set to "1", "true", "yes", or "on" to enable verbose compilation by default. :type TILELANG_VERBOSE: str .. py:function:: par_compile(funcs, out_idx = None, execution_backend = None, target = None, target_host = None, verbose = None, 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. If None, reads from TILELANG_EXECUTION_BACKEND environment variable (defaults to "auto"). :type execution_backend: Literal["auto", "dlpack", "tvm_ffi", "cython", "nvrtc", "torch", "cutedsl"], optional :param target: Compilation target, either as a string or a TVM Target object. If None, reads from TILELANG_TARGET environment variable (defaults to "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. If None, reads from TILELANG_VERBOSE environment variable (defaults to 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. :type pass_configs: dict, optional :param Environment Variables: :param ---------------------: :param TILELANG_TARGET: Default compilation target (e.g., "cuda", "llvm"). Defaults to "auto". :type TILELANG_TARGET: str :param TILELANG_EXECUTION_BACKEND: Default execution backend. Defaults to "auto". :type TILELANG_EXECUTION_BACKEND: str :param TILELANG_VERBOSE: Set to "1", "true", "yes", or "on" to enable verbose compilation by default. :type TILELANG_VERBOSE: str .. py:class:: JITImpl Bases: :py:obj:`Generic`\ [\ :py:obj:`_P`\ , :py:obj:`_KP`\ , :py:obj:`_T`\ , :py:obj:`_Ret`\ ] Just-In-Time compilation wrapper for TileLang programs. This class provides a unified interface for compiling and executing TileLang kernels. It supports two execution modes that are automatically inferred: Execution Modes --------------- - **lazy**: The decorated function returns a PrimFunc explicitly. Calling the JIT wrapper returns a compiled kernel object, which can be invoked separately. This mode is useful when you want to inspect or reuse the kernel object. Example (lazy mode):: @tilelang.jit(out_idx=[-1]) def matmul(M, N, K, block_M, block_N, block_K): @T.prim_func def kernel(A: T.Tensor((M, K), dtype), ...): ... return kernel # explicitly return PrimFunc kernel = matmul(1024, 1024, 1024, 128, 128, 32) # returns kernel result = kernel(a, b) # execute separately - **eager**: The decorated function uses the DSL builder pattern with tensor type annotations. Calling the JIT wrapper compiles and immediately executes the kernel, returning the result directly. Example (eager mode):: @tilelang.jit def gemm(A, B, C, block_M: int = 64): M, N, K = T.const("M N K") A: T.Tensor[[M, K], dtype] # tensor shape via annotation B: T.Tensor[[K, N], dtype] C: T.Tensor[[M, N], dtype] with T.Kernel(...): ... gemm(A, B, C) # compiles and executes immediately The mode is automatically inferred based on whether the function returns a PrimFunc (lazy) or uses the builder pattern (eager). .. attribute:: out_idx Index(es) of output tensor(s) to return (lazy mode only). :type: list[int] | int | None .. attribute:: execution_backend Backend for kernel execution ("auto", "dlpack", "tvm_ffi", etc.). :type: str | None .. attribute:: target TVM compilation target (e.g., "cuda", "llvm", "auto"). :type: str | Target | None .. attribute:: target_host Host target for cross-compilation. :type: str | Target | None .. attribute:: verbose Enable verbose compilation output. :type: bool | None .. attribute:: pass_configs TVM pass configuration options. :type: dict[str, Any] | None .. attribute:: debug_root_path Directory to save compiled kernel source for debugging. :type: str | None .. attribute:: compile_flags Additional compiler flags. :type: list[str] | str | None .. attribute:: func_source Original Python source code of the decorated function. :type: str .. attribute:: signature Function signature of the original function. :type: inspect.Signature .. attribute:: mode Execution mode. "auto" infers from function behavior. :type: Literal["auto", "lazy", "eager"] .. attribute:: func The wrapped function object. :type: JITFunc .. py:attribute:: out_idx :type: list[int] | int | None .. py:attribute:: execution_backend :type: Literal['auto', 'dlpack', 'tvm_ffi', 'cython', 'nvrtc', 'torch', 'cutedsl'] | None .. py:attribute:: target :type: str | tvm.target.Target | None .. py:attribute:: target_host :type: str | tvm.target.Target | None .. py:attribute:: verbose :type: bool | None .. py:attribute:: pass_configs :type: dict[str, Any] | None .. py:attribute:: debug_root_path :type: str | None .. py:attribute:: compile_flags :type: list[str] | str | None .. py:attribute:: func_source :type: str .. py:attribute:: signature :type: inspect.Signature .. py:attribute:: mode :type: Literal['auto', 'lazy', 'eager'] .. py:attribute:: func :type: tilelang.language.eager.JITFunc[_KP, _T] .. py:method:: __post_init__() .. py:method:: get_tir(*args, **kwargs) Retrieve a TIR (Tensor Intermediate Representation) PrimFunc from the stored callable or object. .. py:method:: initialize_jit_mode(*args, **kwargs) .. py:method:: par_compile(configs, num_workers = None, ignore_error = False) Parallel compile multiple TileLang PrimFunc with TVM and build JITKernels. :param configs: The configurations to elaborate and compile. Each config can be either a dictionary mapping keyword arguments to values, or a tuple of positional arguments. :type configs: Iterable[Union[dict[str, Any], tuple[Any, ...]]] :param num_workers: Number of parallel workers to use for compilation. Defaults to None, which lets the system decide. :type num_workers: int, optional :param ignore_error: If True, compilation errors for individual configs will be logged as warnings and the corresponding result will be None. If False, any compilation error will raise an exception. Defaults to False. :type ignore_error: bool, optional :returns: A list of compiled JITKernel objects corresponding to the provided configs. :rtype: List[JITKernel] .. py:method:: compile(*args, **kwargs) .. py:method:: parse_cache_key(*args, **kwargs) .. py:method:: get_kernel_source(*args, **kwargs) .. py:method:: __call__(*args, **kwargs) .. py:data:: ExecutionBackend .. py:function:: jit(func: Callable[_KP, _T]) -> JITImpl[_KP, _KP, _T, _T] jit(*, out_idx: Any = None, target: str | tvm.target.Target | None = None, target_host: str | tvm.target.Target | None = None, execution_backend: ExecutionBackend | None = None, verbose: bool | None = None, pass_configs: dict[str, Any] | None = None, debug_root_path: str | None = None, compile_flags: list[str] | str | None = None) -> Callable[[Callable[_KP, _T]], JITImpl[_KP, _KP, _T, _T]] JIT compiler decorator for TileLang functions. Supports two execution modes (automatically inferred): - **lazy**: Function returns PrimFunc explicitly. Returns compiled kernel object. - **eager**: Function uses DSL builder pattern. Executes kernel immediately. :param out_idx: Output tensor index(es). Only supported in lazy mode. :type out_idx: list[int] | int | None :param target: TVM compilation target (e.g., "cuda", "llvm", "auto"). :type target: str | Target | None :param target_host: Host target for cross-compilation. :type target_host: str | Target | None :param execution_backend: Backend for kernel execution. :type execution_backend: ExecutionBackend | None :param verbose: Enable verbose compilation output. :type verbose: bool | None :param pass_configs: TVM pass configuration options. :type pass_configs: dict[str, Any] | None :param debug_root_path: Directory to save compiled kernel source for debugging. :type debug_root_path: str | None :param compile_flags: Additional compiler flags. :type compile_flags: list[str] | str | None