tilelang.jit package

Subpackages

• tilelang.jit.adapter package
  • Subpackages
    • tilelang.jit.adapter.ctypes package
      • Submodules
      • Module contents
    • tilelang.jit.adapter.cython package
      • Submodules
      • Module contents
  • Submodules
    • tilelang.jit.adapter.base module
      • BaseKernelAdapter
    • tilelang.jit.adapter.dlpack module
      • TorchDLPackKernelAdapter
    • tilelang.jit.adapter.libgen module
      • LibraryGenerator
    • tilelang.jit.adapter.utils module
      • get_annotated_mod()
      • is_cpu_target()
      • is_cuda_target()
      • is_hip_target()
      • match_declare_kernel()
      • match_declare_kernel_cpu()
      • match_global_kernel()
    • tilelang.jit.adapter.wrapper module
      • BaseWrapper
      • TLCPUSourceWrapper
      • TLCUDASourceWrapper
      • TLHIPSourceWrapper
      • TLWrapper
  • Module contents

Submodules

• tilelang.jit.env module
• tilelang.jit.kernel module
  • JITKernel
    • JITKernel.artifact
    • JITKernel.adapter
    • JITKernel.torch_function
    • JITKernel.adapter
    • JITKernel.artifact
    • JITKernel.config
    • JITKernel.export_library()
    • JITKernel.from_database()
    • JITKernel.from_tilelang_function()
    • JITKernel.get_host_source()
    • JITKernel.get_kernel_source()
    • JITKernel.get_profiler()
    • JITKernel.get_tuner_result()
    • JITKernel.host_source
    • JITKernel.kernel_source
    • JITKernel.latency
    • JITKernel.out_idx
    • JITKernel.params
    • JITKernel.prim_func
    • JITKernel.ref_latency
    • JITKernel.run_once()
    • JITKernel.torch_function
    • JITKernel.update_tuner_result()
• tilelang.jit.param module
  • Kernel
    • Kernel.get_kernel_source()
  • Program

Module contents

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.

tilelang.jit.compile(func: Optional[PrimFunc] = None, out_idx: Optional[Union[List[int], int]] = None, execution_backend: Literal['dlpack', 'ctypes', 'cython'] = 'cython', target: Union[str, Target] = 'auto', target_host: Optional[Union[str, Target]] = None, verbose: bool = False, pass_configs: Optional[Dict[str, Any]] = None) → JITKernel

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: “dlpack”). :type execution_backend: Literal[“dlpack”, “ctypes”], 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.

Available options:

“tir.disable_vectorize”: bool, default: False “tl.disable_tma_lower”: bool, default: False “tl.disable_warp_specialized”: bool, default: False “tl.config_index_bitwidth”: int, default: None “tl.disable_dynamic_tail_split”: bool, default: False “tl.dynamic_vectorize_size_bits”: int, default: 128 “tl.disable_safe_memory_legalize”: bool, default: False

tilelang.jit.jit(func: Optional[Union[Callable[[_P], _RProg], PrimFunc]] = None, *, out_idx: Optional[Any] = None, target: Union[str, Target] = 'auto', target_host: Optional[Union[str, Target]] = None, execution_backend: Literal['dlpack', 'ctypes', 'cython'] = 'cython', verbose: bool = False, pass_configs: Optional[Dict[str, Any]] = None, debug_root_path: Optional[str] = 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"], 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.

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