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/kernel/index
   /autoapi/tilelang/jit/param/index


Attributes
----------

.. autoapisummary::

   tilelang.jit.logger


Functions
---------

.. autoapisummary::

   tilelang.jit.compile
   tilelang.jit.jit


Package Contents
----------------

.. py:data:: logger

.. py:function:: 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.
                        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
   :type pass_configs: dict, optional


.. py:function:: 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.

   :param func_or_out_idx: 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`).
   :type func_or_out_idx: Any, optional
   :param target: Compilation target for TVM (e.g., "cuda", "llvm"). Defaults to "auto".
   :type target: Union[str, Target], optional
   :param target_host: Target host for cross-compilation. Defaults to None.
   :type target_host: Union[str, Target], optional
   :param execution_backend: Backend for kernel execution and argument passing. Defaults to "cython".
   :type execution_backend: Literal["dlpack", "ctypes", "cython"], optional
   :param verbose: Enables verbose logging during compilation. Defaults to False.
   :type verbose: bool, optional
   :param pass_configs: Configurations for TVM's pass context. Defaults to None.
   :type pass_configs: Optional[Dict[str, Any]], optional
   :param debug_root_path: Directory to save compiled kernel source for debugging. Defaults to None.
   :type debug_root_path: Optional[str], optional

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