tilelang-rocm 0.1.4.post4__cp310-cp310-manylinux1_x86_64.whl → 0.1.4.post9__cp310-cp310-manylinux1_x86_64.whl

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,204 @@ 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)}")
+
+                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

tilelang/jit/adapter/wrapper.py

@@ -49,6 +49,29 @@ extern "C" int call({}) {{
 }}
 """
 
+L2_PERSISTENT_MAP_CREATE_HANDLE = """
+\tcudaStreamAttrValue stream_attribute;
+\tsize_t init_persisting_l2_cache_size;
+\tcudaDeviceGetLimit(&init_persisting_l2_cache_size, cudaLimitPersistingL2CacheSize);
+"""
+
+L2_PERSISTENT_MAP_INIT_FUNC = """
+\tstream_attribute.accessPolicyWindow.hitRatio = {1};
+\tstream_attribute.accessPolicyWindow.hitProp = cudaAccessPropertyPersisting;
+\tstream_attribute.accessPolicyWindow.missProp = cudaAccessPropertyStreaming;
+\tcudaDeviceSetLimit(cudaLimitPersistingL2CacheSize, {3});
+\tstream_attribute.accessPolicyWindow.base_ptr = (void*)({0});
+\tstream_attribute.accessPolicyWindow.num_bytes = {3};
+\tcudaStreamSetAttribute(stream, cudaStreamAttributeAccessPolicyWindow, &stream_attribute);
+"""
+
+L2_PERSISTENT_MAP_RESET_HANDLE = """
+\tstream_attribute.accessPolicyWindow.num_bytes = 0;
+\tcudaStreamSetAttribute(stream, cudaStreamAttributeAccessPolicyWindow, &stream_attribute);
+\tcudaCtxResetPersistingL2Cache();
+\tcudaDeviceSetLimit(cudaLimitPersistingL2CacheSize, init_persisting_l2_cache_size);
+"""
+
 TMA_DESC_INIT_FUNC = """
 \tCUtensorMap {0};
 \tCUtensorMapDataType {0}_type= (CUtensorMapDataType){1};
@@ -127,6 +150,7 @@ class TLCUDASourceWrapper(object):
         self.block_info: Union[List[int], Dict] = [1, 1, 1]
         self.grid_info: Union[List[int], Dict] = [1, 1, 1]
         self.tma_descriptor_args: Optional[Dict] = None
+        self.l2_persistent_map: Optional[Dict[str, Dict]] = {}
         self.parse_source_information()
         self.srcpath: Optional[str] = None
         self.libpath: Optional[str] = None
@@ -196,7 +220,15 @@ class TLCUDASourceWrapper(object):
                 p = int(p)
             return str(p).replace("//", "/")
 
+        has_l2_persistent_map = False
+        for function_name, _ in function_informations.items():
+            if function_name in self.l2_persistent_map:
+                has_l2_persistent_map = True
+                break
+
         kernel_launch_code = """"""
+        if has_l2_persistent_map:
+            kernel_launch_code += L2_PERSISTENT_MAP_CREATE_HANDLE
         desc_name_map: Dict[str, str] = {}
         for function_name, function_info in function_informations.items():
             block_info = function_info["block_info"]
@@ -221,16 +253,37 @@ class TLCUDASourceWrapper(object):
             grid_str = "dim3({}, {}, {})".format(
                 legalize_c(grid_info[0]), legalize_c(grid_info[1]), legalize_c(grid_info[2]))
             smem_str = 0 if dynamic_smem_buf is None else dynamic_smem_buf
+            init_l2_persistent_map = self.generate_l2_persistent_map(function_name)
+            kernel_launch_code += init_l2_persistent_map
             kernel_launch_code += "\t{}<<<{}, {}, {}, stream>>>({});\n".format(
                 function_name, grid_str, block_str, smem_str, call_args)
-            kernel_launch_code += "TILELANG_CHECK_LAST_ERROR(\"{}\");\n".format(function_name)
+            kernel_launch_code += "\tTILELANG_CHECK_LAST_ERROR(\"{}\");\n".format(function_name)
+            if has_l2_persistent_map:
+                kernel_launch_code += L2_PERSISTENT_MAP_RESET_HANDLE
 
-        kernel_launch_code = self.generate_tma_descriptor_args(desc_name_map) + kernel_launch_code
+        init_tma_descriptor_args = self.generate_tma_descriptor_args(desc_name_map)
+        kernel_launch_code = init_tma_descriptor_args + kernel_launch_code
 
         # Wrap the kernel dispatch logic in an external C function
         host_func = PREDEF_HOST_FUNC.format(def_args, kernel_launch_code)
         return host_func
 
+    def generate_l2_persistent_map(self, function_name: str) -> str:
+        if function_name not in self.l2_persistent_map:
+            return ""
+        init_l2_persistent_map = ""
+        for buffer_name, (hit_ratio,
+                          size_in_bytes) in self.l2_persistent_map[function_name].items():
+            # get persisting_l2_cache_max_size
+            from tilelang.carver.arch.driver import get_persisting_l2_cache_max_size
+            persisting_l2_cache_max_size = get_persisting_l2_cache_max_size()
+            num_bytes = min(size_in_bytes, persisting_l2_cache_max_size)
+
+            init_l2_persistent_map += L2_PERSISTENT_MAP_INIT_FUNC.format(
+                buffer_name, float(hit_ratio), size_in_bytes, num_bytes)
+
+        return init_l2_persistent_map
+
     def generate_tma_descriptor_args(self, desc_name_map: Dict[str, str]) -> str:
         tma_descripter_init = ""
         if self.tma_descriptor_args is None:
@@ -263,10 +316,19 @@ class TLCUDASourceWrapper(object):
             box_dim = remaining_args[2 * tensor_rank:3 * tensor_rank]
             element_strides = remaining_args[3 * tensor_rank:4 * tensor_rank]
 
-            global_dim = [str(i) for i in global_dim]
-            global_stride = [str(i) for i in global_stride]
-            box_dim = [str(i) for i in box_dim]
-            element_strides = [str(i) for i in element_strides]
+            def legalize_c2s(p):
+                # Convert TIR expressions to legal C expressions
+                # Directly convert to string since the special case handling
+                # does not alter the string representation for `tvm.tir.Var` and `IntImm`.
+                # Replace Python's floor division operator with C's division operator
+                if isinstance(p, tvm.tir.IntImm):
+                    p = int(p)
+                return str(p)
+
+            global_dim = [legalize_c2s(i) for i in global_dim]
+            global_stride = [legalize_c2s(i) for i in global_stride]
+            box_dim = [legalize_c2s(i) for i in box_dim]
+            element_strides = [legalize_c2s(i) for i in element_strides]
 
             # Extract remaining parameters
             try:
@@ -331,6 +393,9 @@ class TLCUDASourceWrapper(object):
         for _, func in self.host_mod.functions.items():
             if "tma_descriptor_args" in func.attrs:
                 self.tma_descriptor_args = func.attrs["tma_descriptor_args"]
+            if "l2_persistent_map" in func.attrs:
+                self.l2_persistent_map[function_name] = func.attrs["l2_persistent_map"]
+
             host_code = str(func)
             for function_name in function_names:
                 index = host_code.index(f'T.call_packed("{function_name}"')
@@ -413,6 +478,26 @@ class TLHIPSourceWrapper(TLCUDASourceWrapper):
     A wrapper class for the TileLang HIP backend.
     """
 
+    _TYPE_MAP = {
+        "float32": "float",
+        "float16": "half_t",
+        "bfloat16": "bfloat16_t",
+        "e4m3_float8": "fp8_e4_t",
+        "e5m2_float8": "fp8_e5_t",
+        "float8_e4m3fnuz": "fp8_e4_t",
+        "e4m3fnuz_float8": "fp8_e4_t",
+        "float64": "double",
+        "int64": "int64_t",
+        "int32": "int",
+        "uint32": "unsigned int",
+        "bool": "int8_t",
+        "int8": "int8_t",
+        "uint8": "uint8_t",
+        "int16": "int16_t",
+        "uint16": "uint16_t",
+        "uchar": "uint8_t",
+    }
+
     def __init__(self,
                  scheduled_ir_module: IRModule,
                  source: str,