tilelang-rocm 0.1.4.post5__cp310-cp310-manylinux1_x86_64.whl → 0.1.4.post10__cp310-cp310-manylinux1_x86_64.whl

tilelang/intrinsics/mfma_macro_generator.py

@@ -30,6 +30,7 @@ class MatrixCoreIntrinEmitter(object):
         "int32": "int32",
         "e4m3_float8": "e4m3",
         "e5m2_float8": "e5m2",
+        "float8_e4m3fnuz": "e4m3fnuz",
     }
 
     # k_pack represents the number of elements in a vectorized instruction
@@ -83,10 +84,14 @@ class MatrixCoreIntrinEmitter(object):
 
     def _initialize_k_dim(self, a_dtype="float16"):
         if isinstance(a_dtype, str):
+            if a_dtype in ["float8_e4m3fnuz"]:
+                self.k_dim = 32
+                return
             a_dtype = DataType(a_dtype)
+
         if a_dtype.bits == 32:
             self.k_dim = 4
-        elif a_dtype.bits in [16, 8]:
+        elif a_dtype.bits in {16, 8}:
             self.k_dim = 16
         else:
             raise ValueError(f"Unsupported a_dtype = {a_dtype}")
@@ -115,10 +120,14 @@ class MatrixCoreIntrinEmitter(object):
             "float16": "f16",
             "float32": "f32",
             "int8": "i8",
-            "int32": "i32"
+            "int32": "i32",
+            "float8_e4m3fnuz": "fp8",
         }[in_dtype]
 
-        self.mfma_suffix = f"{out_dtype_abbrv}_{M_DIM}x{N_DIM}x{k_dim}{in_dtype_abbrv}"
+        if in_dtype_abbrv == "fp8":
+            self.mfma_suffix = f"{out_dtype_abbrv}_{M_DIM}x{N_DIM}x{k_dim}_fp8_fp8"
+        else:
+            self.mfma_suffix = f"{out_dtype_abbrv}_{M_DIM}x{N_DIM}x{k_dim}{in_dtype_abbrv}"
 
     def _initialize_micro_size(self, m_dim=16, n_dim=16, k_dim=16):
         self.micro_size_x = m_dim
@@ -141,12 +150,16 @@ class MatrixCoreIntrinEmitter(object):
             shared_16x16_to_local_64x4_layout_B,
             shared_16x32_to_local_64x8_layout_A,
             shared_16x32_to_local_64x8_layout_B,
+            shared_16x64_to_local_64x16_layout_A,
+            shared_16x64_to_local_64x16_layout_B,
             thread_id_shared_access_64x1_to_16x4_layout_A,
             thread_id_shared_access_64x1_to_4x16_layout_B,
             thread_id_shared_access_64x4_to_16x16_layout_A,
             thread_id_shared_access_64x4_to_16x16_layout_B,
             thread_id_shared_access_64x8_to_16x32_layout_A,
             thread_id_shared_access_64x8_to_16x32_layout_B,
+            thread_id_shared_access_64x16_to_16x64_layout_A,
+            thread_id_shared_access_64x16_to_16x64_layout_B,
         )
 
         k_dim = self.k_dim * self.k_pack
@@ -171,8 +184,15 @@ class MatrixCoreIntrinEmitter(object):
             if is_b:
                 index_map = shared_16x32_to_local_64x8_layout_A if transposed else shared_16x32_to_local_64x8_layout_B
                 reverse_index_map = thread_id_shared_access_64x8_to_16x32_layout_A if transposed else thread_id_shared_access_64x8_to_16x32_layout_B
+        elif k_dim == 64:
+            index_map = shared_16x64_to_local_64x16_layout_B if transposed else shared_16x64_to_local_64x16_layout_A
+            reverse_index_map = thread_id_shared_access_64x16_to_16x64_layout_B if transposed else thread_id_shared_access_64x16_to_16x64_layout_A
+
+            if is_b:
+                index_map = shared_16x64_to_local_64x16_layout_A if transposed else shared_16x64_to_local_64x16_layout_B
+                reverse_index_map = thread_id_shared_access_64x16_to_16x64_layout_A if transposed else thread_id_shared_access_64x16_to_16x64_layout_B
         else:
-            raise ValueError("k_dim must be 4 or 16 currently")
+            raise ValueError("k_dim must be 4 or 16 or 32 or 64 currently")
 
         return index_map, reverse_index_map
 
@@ -231,7 +251,7 @@ class MatrixCoreIntrinEmitter(object):
                 for i in T.serial(warp_rows):
                     for local_id in T.vectorized(k_pack * local_size_a):
                         row, col = T.meta_var(reverse_index_map(tx, local_id))
-                        l, r = (rk * chunk + ki * micro_size_k,
+                        l, r = (rk * chunk + ki * (k_pack * micro_size_k),
                                 warp_m * warp_row_tiles + i * micro_size_x)
                         A_local_buf[i * k_pack * local_size_a + local_id] = A_shared_buf[l + row,
                                                                                          r + col]
@@ -240,7 +260,7 @@ class MatrixCoreIntrinEmitter(object):
                     for local_id in T.vectorized(k_pack * local_size_a):
                         row, col = T.meta_var(reverse_index_map(tx, local_id))
                         l, r = (warp_m * warp_row_tiles + i * micro_size_x,
-                                rk * chunk + ki * micro_size_k)
+                                rk * chunk + ki * (k_pack * micro_size_k))
                         A_local_buf[i * k_pack * local_size_a + local_id] = A_shared_buf[l + row,
                                                                                          r + col]
 
@@ -275,7 +295,7 @@ class MatrixCoreIntrinEmitter(object):
                         row, col = T.meta_var(reverse_index_map(tx, local_id))
                         l, r = (
                             warp_n * warp_col_tiles + j * micro_size_y,
-                            rk * chunk + ki * micro_size_k,
+                            rk * chunk + ki * (k_pack * micro_size_k),
                         )
                         B_local_buf[j * k_pack * local_size_b + local_id] = B_shared_buf[l + row,
                                                                                          r + col]
@@ -284,7 +304,7 @@ class MatrixCoreIntrinEmitter(object):
                     for local_id in T.vectorized(k_pack * local_size_b):
                         row, col = T.meta_var(reverse_index_map(tx, local_id))
                         l, r = (
-                            rk * chunk + ki * micro_size_k,
+                            rk * chunk + ki * (k_pack * micro_size_k),
                             warp_n * warp_col_tiles + j * micro_size_y,
                         )
                         B_local_buf[j * k_pack * local_size_b + local_id] = B_shared_buf[l + row,
@@ -347,7 +367,7 @@ class MatrixCoreIntrinEmitter(object):
         def _warp_stmatrix_shared(C_local_buf, C_buf, thread_binding):
             tx, warp_n, warp_m = self.extract_thread_binding(thread_binding)
             for i, j in T.grid(warp_rows, warp_cols):
-                for local_id in T.serial(local_size_out):
+                for local_id in T.vectorized(local_size_out):
                     row, col = T.meta_var(mfma_store_index_map(tx, local_id))
                     C_buf[warp_m * warp_rows + i, warp_n * warp_cols + j, row,
                           col] = C_local_buf[i * warp_cols * local_size_out + j * local_size_out +
@@ -357,7 +377,7 @@ class MatrixCoreIntrinEmitter(object):
         def _warp_stmatrix_global(C_local_buf, C_buf, thread_binding):
             tx, warp_n, warp_m = self.extract_thread_binding(thread_binding)
             for i, j in T.grid(warp_rows, warp_cols):
-                for local_id in T.serial(local_size_out):
+                for local_id in T.vectorized(local_size_out):
                     row, col = T.meta_var(mfma_store_index_map(tx, local_id))
                     C_buf[(pid_m * BLOCK_M + warp_m * warp_rows + i) * M_DIM + row,
                           (pid_n * BLOCK_N + warp_n * warp_cols + j) * N_DIM +

tilelang/intrinsics/mma_layout.py

@@ -92,60 +92,42 @@ def shared_32x16_to_mma_32x16_smoothlayout(i, j):
     return (i * 2 + j // 16, j % 16)
 
 
-def get_swizzle_layout(row_idx, col_idx, row_size, dtype: Union[DataType, str]):
+def get_swizzle_layout(row_idx, col_idx, row_size, dtype: Union[DataType, str], swizzle_bytes=None):
     ana = arith.Analyzer()
-    BANK_SIZE_BYTES = 128
     if isinstance(dtype, str):
         dtype = DataType(dtype)
-    col_idx_outer, col_idx_inner = col_idx // (BANK_SIZE_BYTES // dtype.bits), col_idx % (
-        BANK_SIZE_BYTES // dtype.bits)
-    #  use transaction bits to support diverse dtype.
-    #  for fp16, 64 elems * 16 bits = 1024 bits, 32 elems * 32 bits = 512 bits
-    #  for int8, 128 elems * 8 bits = 1024 bits, 64 elems * 8 bits = 512 bits
-    coalescent_bits = dtype.bits * row_size
-    # permutation on 4 banks, each bank has 32 bits
-    bank_elems = BANK_SIZE_BYTES // dtype.bits
-    new_col_idx_outer = None
-
-    if coalescent_bits % 1024 == 0:
-        #   Use 8 * 8 permuted layout
-        #   Every number below corresponds to 8 consecutive fp16 number in shared mem, i.e. one read
-        #   Every row below corresponds to 32 banks
-        #   0  1  2  3  4  5  6  7    ==>    0  1  2  3  4  5  6  7
-        #   0  1  2  3  4  5  6  7    ==>    1  0  3  2  5  4  7  6
-        #   0  1  2  3  4  5  6  7    ==>    2  3  0  1  6  7  4  5
-        #   0  1  2  3  4  5  6  7    ==>    3  2  1  0  7  6  5  4
-        #   0  1  2  3  4  5  6  7    ==>    4  5  6  7  0  1  2  3
-        #   0  1  2  3  4  5  6  7    ==>    5  4  7  6  1  0  3  2
-        #   0  1  2  3  4  5  6  7    ==>    6  7  4  5  2  3  0  1
-        #   0  1  2  3  4  5  6  7    ==>    7  6  5  4  3  2  1  0
-        row_idx_sub = row_idx % bank_elems
-        new_col_idx_outer = col_idx_outer ^ row_idx_sub
-    else:
-        assert coalescent_bits % 512 == 0
-        #  Use 8 * 4 permuted layout
-        #  Every number below corresponds to 8 consecutive fp16 number in shared mem, i.e. one read
-        #  Every row below corresponds to 16 banks
-        #  0  1  2  3    ==>    0  1  2  3
-        #  0  1  2  3    ==>    0  1  2  3
-        #  0  1  2  3    ==>    1  0  3  2
-        #  0  1  2  3    ==>    1  0  3  2
-        #  0  1  2  3    ==>    2  3  0  1
-        #  0  1  2  3    ==>    2  3  0  1
-        #  0  1  2  3    ==>    3  2  1  0
-        #  0  1  2  3    ==>    3  2  1  0
-        #  View with 8 elements per row:
-        #  0  1  2  3  4  0  1  2  3    ==>    0  1  2  3  0  1  2  3
-        #  0  1  2  3  4  0  1  2  3    ==>    1  0  3  2  1  0  3  2
-        #  0  1  2  3  4  0  1  2  3    ==>    2  3  0  1  2  3  0  1
-        #  0  1  2  3  4  0  1  2  3    ==>    3  2  1  0  3  2  1  0
-        row_idx_sub = row_idx % bank_elems
-        #  Interleave elems per byte
-        interleave_elems = 32 // dtype.bits
-        new_col_idx_outer = col_idx_outer ^ (row_idx_sub // interleave_elems)
-
-    assert (new_col_idx_outer is not None), f"Unsupported dtype {dtype} with {coalescent_bits} bits"
-    return row_idx, ana.simplify(new_col_idx_outer * bank_elems + col_idx_inner)
+    row_bytes = dtype.bits * row_size // 8
+    assert row_bytes % 32 == 0, "Row size must be multiple of 32B."
+    if swizzle_bytes is None:
+        swizzle_bytes = min(128, row_bytes)
+    # 128B swizzle
+    #   Use 8 * 8 permuted layout
+    #   Every number below corresponds to 16B
+    #   0  1  2  3  4  5  6  7    ==>    0  1  2  3  4  5  6  7
+    #   0  1  2  3  4  5  6  7    ==>    1  0  3  2  5  4  7  6
+    #   0  1  2  3  4  5  6  7    ==>    2  3  0  1  6  7  4  5
+    #   0  1  2  3  4  5  6  7    ==>    3  2  1  0  7  6  5  4
+    #   0  1  2  3  4  5  6  7    ==>    4  5  6  7  0  1  2  3
+    #   0  1  2  3  4  5  6  7    ==>    5  4  7  6  1  0  3  2
+    #   0  1  2  3  4  5  6  7    ==>    6  7  4  5  2  3  0  1
+    #   0  1  2  3  4  5  6  7    ==>    7  6  5  4  3  2  1  0
+    # 64B swizzle
+    #  Use 8 * 4 permuted layout
+    #  Every number below corresponds to 16B
+    #  0  1  2  3  4  0  1  2  3    ==>    0  1  2  3  0  1  2  3
+    #  0  1  2  3  4  0  1  2  3    ==>    1  0  3  2  1  0  3  2
+    #  0  1  2  3  4  0  1  2  3    ==>    2  3  0  1  2  3  0  1
+    #  0  1  2  3  4  0  1  2  3    ==>    3  2  1  0  3  2  1  0
+    # 32B swizzle
+    #  Use 8 * 2 permuted layout
+    #  Every number below corresponds to 16B
+    #  0  1  2  3  4  5  6  7    ==>    0  1  2  3  4  5  6  7
+    #  0  1  2  3  4  5  6  7    ==>    1  0  3  2  5  4  7  6
+    elem_per_16B = 128 // dtype.bits
+    col_idx_16B = col_idx // elem_per_16B
+    col_idx_in_16B = col_idx % elem_per_16B
+    new_col_idx_16B = col_idx_16B ^ (row_idx % (swizzle_bytes // 16))
+    return row_idx, ana.simplify(new_col_idx_16B * elem_per_16B + col_idx_in_16B)
 
 
 def make_mma_swizzle_layout(shared_buf, is_smooth: bool = False):

tilelang/jit/__init__.py

@@ -6,109 +6,31 @@ It includes functionality to JIT-compile TileLang programs into a runnable
 kernel adapter using TVM.
 """
 
-from typing import Callable, List, Literal, Union, Any, Optional, Dict
-
+from typing import (
+    Any,
+    List,
+    Union,
+    Callable,
+    Tuple,
+    overload,
+    Literal,
+    Dict,  # For type hinting dicts
+    Optional,
+)
 from tilelang import tvm as tvm
 from tvm.tir import PrimFunc
 from tvm.target import Target
 
-from tilelang.jit.adapter import BaseKernelAdapter
 from tilelang.jit.kernel import JITKernel
-from tilelang.utils.target import determine_target, AVALIABLE_TARGETS
 from tilelang.cache import cached
+from os import path, makedirs
 from logging import getLogger
+import functools
+from tilelang.jit.param import Kernel, _P, _RProg
 
 logger = getLogger(__name__)
 
 
-def jit(
-    func: Callable = None,
-    *,  # Enforce keyword-only arguments from here on
-    out_idx: 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
-    -------
-    BaseKernelAdapter
-        An adapter object that encapsulates the compiled function and can be
-        used to execute it.
-
-    Raises
-    ------
-    AssertionError
-        If the provided target is an invalid string not present in AVALIABLE_TARGETS.
-    """
-
-    # If the target is specified as a string, ensure it is valid and convert to a TVM Target.
-    if isinstance(target, str):
-        assert target in AVALIABLE_TARGETS, f"Invalid target: {target}"
-        target = determine_target(target)
-
-    target = Target(target)
-
-    assert execution_backend in ["dlpack", "ctypes", "cython"], "Invalid execution backend."
-
-    def _compile_and_create_adapter(tilelang_func: PrimFunc) -> BaseKernelAdapter:
-        """
-        Compile the given TileLang PrimFunc with TVM and build a kernel adapter.
-
-        Parameters
-        ----------
-        tilelang_func : tvm.tir.PrimFunc
-            The TileLang (TVM TIR) function to compile.
-
-        Returns
-        -------
-        BaseKernelAdapter
-            The compiled and ready-to-run kernel adapter.
-        """
-        if verbose:
-            logger.info(f"Compiling TileLang function:\n{tilelang_func}")
-        return compile(
-            tilelang_func,
-            target=target,
-            verbose=verbose,
-            execution_backend=execution_backend,
-            out_idx=out_idx,
-            **pass_config_kwargs,
-        ).adapter
-
-    # If `func` was given, compile it immediately and return the adapter.
-    if func is not None:
-        return _compile_and_create_adapter(func)
-
-    # Otherwise, return a decorator that expects a function to compile.
-    def real_decorator(tilelang_func: PrimFunc) -> BaseKernelAdapter:
-        return _compile_and_create_adapter(tilelang_func)
-
-    return real_decorator
-
-
 def compile(
     func: PrimFunc = None,
     out_idx: Union[List[int], int, None] = None,
@@ -153,4 +75,207 @@ def compile(
         target_host=target_host,
         verbose=verbose,
         pass_configs=pass_configs,
-    )
+    )
+
+
+class _JitImplementation:
+
+    out_idx: Any
+    target: Union[str, Target]
+    target_host: Union[str, Target]
+    execution_backend: Literal["dlpack", "ctypes", "cython"]
+    verbose: bool
+    pass_configs: Optional[Dict[str, Any]]
+    debug_root_path: Optional[str]
+
+    def __init__(self,
+                 out_idx: Any = None,
+                 target: Union[str, Target] = "auto",
+                 target_host: 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):
+        """
+        Initializes the JIT compiler decorator.
+
+        Parameters
+        ----------
+        out_idx : Any, optional
+            Index(es) of the output tensors to return from the compiled kernel
+            (default: None, meaning all outputs are returned or determined by the kernel itself).
+        target : Union[str, Target], optional
+            Compilation target for TVM. Can be a string (e.g., "cuda", "llvm")
+            or a TVM Target object. If "auto", the target is determined automatically
+            (default: "auto").
+        target_host : Union[str, Target], optional
+            Target host for cross-compilation, similar to `target` (default: None).
+        execution_backend : Literal["dlpack", "ctypes", "cython"], optional
+            The backend used for kernel execution and argument passing.
+            "dlpack" is generally preferred for zero-copy tensor passing with compatible frameworks.
+            "ctypes" uses standard C types. "cython" uses Cython for potentially faster execution.
+            (default: "cython").
+        verbose : bool, optional
+            If True, enables verbose logging during compilation (default: False).
+        pass_configs : Optional[Dict[str, Any]], optional
+            A dictionary of configurations for TVM's pass context. These can fine-tune
+            the compilation process. Examples include "tir.disable_vectorize"
+            (default: None).
+        debug_root_path : Optional[str], optional
+            If provided, the compiled kernel's source code will be saved to a file
+            in this directory. This is useful for debugging the generated code.
+            If None, no debug information is saved (default: None).
+            If a relative path is given, it's made absolute relative to the project root
+            or current working directory.
+        """
+        self.out_idx = out_idx
+        self.execution_backend = execution_backend
+        self.target = target
+        self.target_host = target_host
+        self.verbose = verbose
+        self.pass_configs = pass_configs
+
+        # Corrected debug_root_path handling
+        self.debug_root_path = debug_root_path
+        if self.debug_root_path is not None and not path.isabs(self.debug_root_path):
+            try:
+                base_path = path.dirname(path.dirname(path.dirname(__file__)))
+                self.debug_root_path = path.join(base_path, self.debug_root_path)
+            except NameError:
+                self.debug_root_path = path.abspath(self.debug_root_path)
+
+        self._kernel_cache: Dict[tuple, Kernel] = {}
+
+    # This tells the type checker what the *wrapper* function will return.
+    # this is for linting, please do not remove it.
+    @overload
+    def __call__(self, func: Callable[_P, _RProg]) -> Callable[_P, Tuple[_RProg, Kernel]]:
+        ...
+
+    @overload
+    def __call__(self, func: Callable[_P, _RProg]) -> Callable[_P, Kernel]:
+        ...
+
+    # Actual implementation of __call__
+    def __call__(
+        self,
+        func: Callable[_P, _RProg]  # func is Union[Callable[_P, _RProg], PrimFunc] in original
+    ) -> Callable[_P, Any]:
+
+        @functools.wraps(func)
+        def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> Any:
+            # Separate out the tuning parameters from the user's kwargs
+            tune_params = kwargs.pop('__tune_params', {})
+
+            key_args_tuple = args
+            key_kwargs_tuple = tuple(sorted(kwargs.items()))
+            key = (key_args_tuple, key_kwargs_tuple)
+
+            if key not in self._kernel_cache:
+                # Ensure 'func' (the original user function) is used correctly
+                program_result_source = func
+                if isinstance(program_result_source, PrimFunc):
+                    program_result = program_result_source
+                elif callable(program_result_source):
+                    program_result = program_result_source(*args, **kwargs, **tune_params)
+                else:
+                    raise ValueError(f"Invalid function type: {type(program_result_source)}")
+
+                if self.verbose:
+                    logger.info(f"Verbose: Compiling for program \n {program_result.script()}")
+
+                kernel_result = compile(
+                    program_result,
+                    out_idx=self.out_idx,
+                    execution_backend=self.execution_backend,
+                    target=self.target,
+                    target_host=self.target_host,
+                    verbose=self.verbose,
+                    pass_configs=self.pass_configs,
+                )
+
+                if self.debug_root_path:
+                    func_name = getattr(func, '__name__', 'jit_kernel')  # Use func for name
+                    kernel_file = f'tilelang_jit_kernel_{func_name}.c'
+                    program_file = f'tilelang_jit_program_{func_name}.py'
+                    makedirs(self.debug_root_path, exist_ok=True)
+                    with open(path.join(self.debug_root_path, kernel_file), 'w') as f:
+                        print(kernel_result.get_kernel_source(), file=f)
+                    with open(path.join(self.debug_root_path, program_file), 'w') as f:
+                        print(program_result.script(), file=f)
+
+                self._kernel_cache[key] = kernel_result
+
+            return self._kernel_cache[key]
+
+        return wrapper
+
+
+def jit(  # This is the new public interface
+        func: Union[Callable[_P, _RProg], PrimFunc, None] = None,
+        *,  # Indicates subsequent arguments are keyword-only
+        out_idx: Any = None,
+        target: Union[str, Target] = "auto",
+        target_host: 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
+    -------
+    Callable
+        Either a JIT-compiled wrapper around the input function, or a configured decorator
+        instance that can then be applied to a function.
+    """
+    if callable(func):
+        # Case 1: Used as @jit (func_or_out_idx is the function, others are defaults)
+        # Create a default _JitImplementation instance and apply it to the function.
+        default_decorator = _JitImplementation(
+            out_idx=out_idx,  # Explicitly None for the default case
+            target=target,
+            target_host=target_host,
+            execution_backend=execution_backend,
+            verbose=verbose,
+            pass_configs=pass_configs,
+            debug_root_path=debug_root_path)
+        return default_decorator(func)
+    elif isinstance(func, PrimFunc):
+        raise ValueError("Use tilelang.jit to decorate prim_func is not supported yet.")
+    else:
+        # Case 2: Used as @jit(...) to configure, or func_or_out_idx is meant as out_idx.
+        # Create a _JitImplementation instance with the provided/defaulted arguments.
+        # This instance is a decorator that will be applied to the function later.
+        configured_decorator = _JitImplementation(
+            out_idx=out_idx,  # Pass along; could be an actual out_idx or None
+            target=target,
+            target_host=target_host,
+            execution_backend=execution_backend,
+            verbose=verbose,
+            pass_configs=pass_configs,
+            debug_root_path=debug_root_path)
+        return configured_decorator

tilelang/jit/adapter/cython/cython_wrapper.pyx

@@ -116,20 +116,23 @@ cdef class CythonKernelWrapper:
         # Convert tensor pointers to C void pointers for kernel call
         call_args = []
         for i in range(len(tensor_list)):
-            if isinstance(tensor_list[i], torch.Tensor):
-                call_args.append(ctypes.c_void_p(tensor_list[i].data_ptr()))
-            elif isinstance(tensor_list[i], int):
+            tensor = tensor_list[i]
+            if isinstance(tensor, torch.Tensor):
+                if not tensor.is_contiguous():
+                    raise ValueError(f"Input tensor at index {i} must be contiguous")
+                call_args.append(ctypes.c_void_p(tensor.data_ptr()))
+            elif isinstance(tensor, int):
                 # Dynamic symbolics which are passed as integer arguments
                 if i in self.ptr_map:
-                    call_args.append(ctypes.c_void_p(tensor_list[i]))
+                    call_args.append(ctypes.c_void_p(tensor))
                 else:
-                    call_args.append(tensor_list[i])
-            elif isinstance(tensor_list[i], float):
-                call_args.append(ctypes.c_float(tensor_list[i]))
-            elif isinstance(tensor_list[i], bool):
-                call_args.append(ctypes.c_bool(tensor_list[i]))
+                    call_args.append(tensor)
+            elif isinstance(tensor, float):
+                call_args.append(ctypes.c_float(tensor))
+            elif isinstance(tensor, bool):
+                call_args.append(ctypes.c_bool(tensor))
             else:
-                raise ValueError(f"Unsupported tensor type: {type(tensor_list[i])}")
+                raise ValueError(f"Unsupported tensor type: {type(tensor)}")
 
         # Check buffer device
         # cdef str tensor_list_device_type = tensor_list[0].device.type

tilelang/jit/adapter/libgen.py

@@ -98,13 +98,15 @@ class LibraryGenerator(object):
 
         src.write(self.lib_code)
         src.flush()
+
         try:
             ret = subprocess.run(command, timeout=timeout)
         except Exception as e:
             raise RuntimeError(f"Compile kernel failed because of {e}") from e
 
         if ret.returncode != 0:
-            raise RuntimeError(f"Compilation Failed! {command}")
+            raise RuntimeError(f"Compilation Failed! {command}"
+                               f"\n {self.lib_code}")
 
         self.srcpath = src.name
         self.libpath = libpath