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.export_library()
    • JITKernel.from_database()
    • JITKernel.from_tilelang_function()
    • JITKernel.get_host_source()
    • JITKernel.get_kernel_source()
    • JITKernel.get_profiler()
    • JITKernel.host_source
    • JITKernel.kernel_source
    • JITKernel.out_idx
    • JITKernel.params
    • JITKernel.run_once()
    • JITKernel.torch_function

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

tilelang.jit.jit(func: Optional[Callable] = None, *, out_idx: Optional[Union[List[int], int]] = None, execution_backend: Literal['dlpack', 'ctypes', 'cython'] = 'cython', target: Union[str, Target] = 'auto', verbose: bool = False, **pass_config_kwargs: Optional[Dict[str, Any]]) → BaseKernelAdapter

A decorator (or decorator factory) that JIT-compiles a given TileLang PrimFunc into a runnable kernel adapter using TVM. If called with arguments, it returns a decorator that can be applied to a function. If called without arguments, it directly compiles the given function.

Parameters:

• func (Callable, optional) – The TileLang PrimFunc to JIT-compile. If None, this function returns a decorator that expects a TileLang PrimFunc.

• out_idx (Union[List[int], int], optional) – The index (or list of indices) of the function outputs. This can be used to specify which outputs from the compiled function will be returned.

• execution_backend (Literal["dlpack", "ctypes"], optional) – The wrapper type to use for the kernel adapter. Currently, only “dlpack” and “ctypes” are supported.

• target (Union[str, Target], optional) – The compilation target for TVM. If set to “auto”, an appropriate target will be inferred automatically. Otherwise, must be one of the supported strings in AVALIABLE_TARGETS or a TVM Target instance.

Returns:

An adapter object that encapsulates the compiled function and can be used to execute it.

Return type:

BaseKernelAdapter

Raises:

AssertionError – If the provided target is an invalid string not present in AVALIABLE_TARGETS.