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

tilelang/cache/tuner_cache.py

@@ -0,0 +1,356 @@
+# Copyright (c) Tile-AI Corporation.
+# Licensed under the MIT License.
+"""The cache utils with class and database persistence - KernelCache Class"""
+
+import os
+import json
+import shutil
+from pathlib import Path
+from hashlib import sha256
+from typing import Callable, List, Literal, Union, Optional
+from tvm.target import Target
+from tvm.tir import PrimFunc
+from tilelang.jit import JITKernel
+from tilelang.engine.param import KernelParam
+import threading
+import cloudpickle
+import logging
+
+from tilelang.env import TILELANG_CACHE_DIR, is_cache_enabled
+from tilelang.version import __version__
+
+KERNEL_PATH = "kernel.cu"
+WRAPPED_KERNEL_PATH = "wrapped_kernel.cu"
+KERNEL_LIB_PATH = "kernel_lib.so"
+PARAMS_PATH = "params.pkl"
+
+
+class AutoTunerCache:
+    """
+    Caches compiled kernels using a class and database persistence to avoid redundant compilation.
+    Cache files:
+        kernel.cu: The compiled kernel source code
+        wrapped_kernel.cu: The compiled wrapped kernel source code
+        kernel_lib.so: The compiled kernel library
+        params.pkl: The compiled kernel parameters
+    """
+
+    _instance = None  # For implementing singleton pattern
+    _lock = threading.Lock()  # For thread safety
+    _memory_cache = {}  # In-memory cache dictionary
+
+    cache_dir: Path = Path(TILELANG_CACHE_DIR)
+
+    def __new__(cls, cache_dir=TILELANG_CACHE_DIR):
+        """
+        Implements singleton pattern for KernelCache class.
+
+        Args:
+            cache_dir (str): Directory path for storing kernel cache. Defaults to TILELANG_CACHE_DIR.
+
+        Returns:
+            KernelCache: The singleton instance of KernelCache.
+        """
+        if cls._instance is None:
+            with cls._lock:
+                if cls._instance is None:  # Double-checked locking
+                    instance = super().__new__(cls)
+                    instance.cache_dir = Path(cache_dir)
+                    os.makedirs(instance.cache_dir, exist_ok=True)
+
+                    instance.logger = logging.getLogger(__name__)
+                    instance.logger.setLevel(logging.ERROR)
+                    instance._memory_cache = {}  # Initialize memory cache
+                    cls._instance = instance
+        return cls._instance
+
+    def _generate_key(
+        self,
+        func: Callable,
+        out_idx: List[int],
+        execution_backend: Literal["dlpack", "ctypes", "cython"] = "cython",
+        args=None,
+        target: Union[str, Target] = "auto",
+        target_host: Union[str, Target] = None,
+        pass_configs: dict = None,
+    ) -> str:
+        """
+        Generates a unique hash key for caching compiled kernels.
+
+        Args:
+            func (Callable): The function to be compiled.
+            out_idx (List[int]): Indices specifying which outputs to return.
+            execution_backend (Literal): Backend type for execution. Defaults to "cython".
+            args: Arguments passed to the function.
+            target (Union[str, Target]): Compilation target platform. Defaults to "auto".
+            target_host (Union[str, Target], optional): Host target platform.
+
+        Returns:
+            str: SHA256 hash key for the kernel configuration.
+        """
+        func_binary = cloudpickle.dumps(func.script())
+        key_data = {
+            "version": __version__,
+            "func": sha256(func_binary).hexdigest(),  # Use SHA256 to generate hash key
+            "out_idx": (tuple(out_idx) if isinstance(out_idx, (list, tuple)) else [out_idx]),
+            "args_repr": tuple(
+                repr(arg) for arg in args
+            ),  # Use repr to serialize arguments, may need more robust serialization
+            "target": str(target),
+            "target_host": str(target_host) if target_host else None,
+            "execution_backend": execution_backend,
+            "pass_configs": pass_configs,
+        }
+        key_string = json.dumps(key_data, sort_keys=True)  # Sort keys to ensure consistency
+        return sha256(key_string.encode()).hexdigest()  # Use SHA256 to generate hash key
+
+    def cached(
+        self,
+        func: PrimFunc = None,
+        out_idx: List[int] = None,
+        *args,
+        target: Union[str, Target] = "auto",
+        target_host: Union[str, Target] = None,
+        execution_backend: Literal["dlpack", "ctypes", "cython"] = "cython",
+        verbose: bool = False,
+        pass_configs: dict = None,
+    ) -> JITKernel:
+        """
+        Caches and reuses compiled kernels to avoid redundant compilation.
+
+        Args:
+            func: Function to be compiled or a prepared PrimFunc
+            out_idx: Indices specifying which outputs to return
+            target: Compilation target platform
+            target_host: Host target platform
+            *args: Arguments passed to func
+
+        Returns:
+            JITKernel: The compiled kernel, either freshly compiled or from cache
+        """
+        if not is_cache_enabled():
+            return JITKernel(
+                func,
+                out_idx=out_idx,
+                execution_backend=execution_backend,
+                target=target,
+                target_host=target_host,
+                verbose=verbose,
+                pass_configs=pass_configs,
+            )
+
+        key = self._generate_key(
+            func=func,
+            out_idx=out_idx,
+            execution_backend=execution_backend,
+            args=args,
+            target=target,
+            target_host=target_host,
+            pass_configs=pass_configs,
+        )
+        with self._lock:
+            # First check in-memory cache
+            if key in self._memory_cache:
+                self.logger.warning("Found kernel in memory cache. For better performance," \
+                                    " consider using `@tilelang.jit` instead of direct kernel caching.")
+                return self._memory_cache[key]
+
+            # Then check disk cache
+            kernel = self._load_kernel_from_disk(key, target, target_host, out_idx,
+                                                 execution_backend, pass_configs, func)
+            if kernel is not None:
+                # Populate memory cache with disk result
+                self._memory_cache[key] = kernel
+                return kernel
+
+        # Compile kernel if cache miss; leave critical section
+        kernel = JITKernel(
+            func,
+            out_idx=out_idx,
+            execution_backend=execution_backend,
+            target=target,
+            target_host=target_host,
+            verbose=verbose,
+            pass_configs=pass_configs,
+        )
+        if execution_backend == "dlpack":
+            self.logger.warning("DLPack backend does not support cache saving to disk.")
+        else:
+            with self._lock:  # enter critical section again to check and update disk cache
+                disk_kernel = self._load_kernel_from_disk(
+                    key,
+                    target,
+                    target_host,
+                    out_idx,
+                    execution_backend,
+                    pass_configs,
+                    func,
+                )
+                if disk_kernel is None:
+                    self._save_kernel_to_disk(key, kernel, func)
+
+        # Store in memory cache after compilation
+        self._memory_cache[key] = kernel
+        return kernel
+
+    def set_cache_dir(self, cache_dir: str):
+        """
+        Sets the cache directory for the kernel cache.
+        """
+        self.cache_dir = Path(cache_dir)
+
+    def get_cache_dir(self) -> Path:
+        """
+        Gets the cache directory for the kernel cache.
+        """
+        return self.cache_dir
+
+    def clear_cache(self):
+        """
+        Clears the entire kernel cache, including both in-memory and disk cache.
+        """
+        with self._lock:
+            self._memory_cache.clear()  # Clear in-memory cache
+            self._clear_disk_cache()  # Clear disk cache
+
+    def _get_cache_path(self, key: str) -> str:
+        """
+        Gets the filesystem path for a cached kernel.
+
+        Args:
+            key (str): The hash key identifying the kernel.
+
+        Returns:
+            str: Absolute path to the cache directory for this kernel.
+        """
+        return os.path.join(self.cache_dir, key)
+
+    def _save_kernel_to_disk(self, key: str, kernel: JITKernel, func: Callable = None):
+        """
+        Persists a compiled kernel to disk cache.
+
+        Args:
+            key (str): The hash key identifying the kernel.
+            kernel (JITKernel): The compiled kernel to be saved.
+            func (Callable, optional): The original function.
+
+        Note:
+            Saves the following files:
+            - kernel.cu: The compiled kernel source code
+            - wrapped_kernel.cu: The wrapped kernel source code
+            - kernel_lib.so: The compiled kernel library
+            - params.pkl: The serialized kernel parameters
+        """
+        cache_path = self._get_cache_path(key)
+        os.makedirs(cache_path, exist_ok=True)  # Ensure directory exists
+
+        # Save kernel source code
+        try:
+            kernel_path = os.path.join(cache_path, KERNEL_PATH)
+            with open(kernel_path, "w") as f:
+                f.write(kernel.artifact.kernel_source)
+        except Exception as e:
+            self.logger.error(f"Error saving kernel source code to disk: {e}")
+
+        # Save wrapped kernel source code
+        try:
+            wrapped_kernel_path = os.path.join(cache_path, WRAPPED_KERNEL_PATH)
+            with open(wrapped_kernel_path, "w") as f:
+                f.write(kernel.adapter.get_kernel_source())
+        except Exception as e:
+            self.logger.error(f"Error saving wrapped kernel source code to disk: {e}")
+
+        # Save kernel library
+        try:
+            kernel_lib_path = os.path.join(cache_path, KERNEL_LIB_PATH)
+            src_lib_path = kernel.adapter.libpath
+            shutil.copy(src_lib_path, kernel_lib_path)
+        except Exception as e:
+            self.logger.error(f"Error saving kernel library to disk: {e}")
+
+        # Save kernel parameters
+        try:
+            params_path = os.path.join(cache_path, PARAMS_PATH)
+            with open(params_path, "wb") as f:
+                cloudpickle.dump(kernel.params, f)
+        except Exception as e:
+            self.logger.error(f"Error saving kernel parameters to disk: {e}")
+
+    def _load_kernel_from_disk(
+        self,
+        key: str,
+        target: Union[str, Target] = "auto",
+        target_host: Union[str, Target] = None,
+        out_idx: List[int] = None,
+        execution_backend: Literal["dlpack", "ctypes", "cython"] = "cython",
+        pass_configs: dict = None,
+        func: Callable = None,
+    ) -> JITKernel:
+        """
+        Loads a previously compiled kernel from disk cache.
+
+        Args:
+            key (str): The hash key identifying the kernel.
+            target (Union[str, Target]): Compilation target platform. Defaults to "auto".
+            target_host (Union[str, Target], optional): Host target platform.
+            out_idx (List[int], optional): Indices specifying which outputs to return.
+            execution_backend (Literal): Backend type for execution. Defaults to "cython".
+            pass_configs (dict, optional): Configuration for compiler passes.
+            func (Callable, optional): The original function.
+
+        Returns:
+            JITKernel: The loaded kernel if found, None otherwise.
+        """
+        cache_path = self._get_cache_path(key)
+        if not os.path.exists(cache_path):
+            return None
+
+        kernel_global_source: Optional[str] = None
+        kernel_params: Optional[List[KernelParam]] = None
+
+        try:
+            wrapped_kernel_path = os.path.join(cache_path, WRAPPED_KERNEL_PATH)
+            with open(wrapped_kernel_path, "r") as f:
+                kernel_global_source = f.read()
+        except Exception as e:
+            self.logger.error(f"Error loading wrapped kernel source code from disk: {e}")
+
+        kernel_lib_path = os.path.join(cache_path, KERNEL_LIB_PATH)
+
+        # Load kernel parameters
+        try:
+            params_path = os.path.join(cache_path, PARAMS_PATH)
+            with open(params_path, "rb") as f:
+                kernel_params = cloudpickle.load(f)
+        except Exception as e:
+            self.logger.error(f"Error loading kernel parameters from disk: {e}")
+
+        if kernel_global_source and kernel_params:
+            return JITKernel.from_database(
+                func=func,
+                kernel_global_source=kernel_global_source,
+                kernel_lib_path=kernel_lib_path,
+                params=kernel_params,
+                target=target,
+                target_host=target_host,
+                out_idx=out_idx,
+                execution_backend=execution_backend,
+                pass_configs=pass_configs,
+            )
+        else:
+            return None
+
+    def _clear_disk_cache(self):
+        """
+        Removes all cached kernels from disk.
+        
+        Note:
+            This operation will delete the entire cache directory and recreate it empty.
+            Use with caution as this operation cannot be undone.
+        """
+        try:
+            if os.path.exists(self.cache_dir):
+                shutil.rmtree(self.cache_dir)  # Delete entire cache directory
+            os.makedirs(self.cache_dir, exist_ok=True)  # Re-create cache directory
+        except Exception as e:
+            self.logger.error(f"Error clearing disk cache: {e}")

tilelang/carver/arch/driver/__init__.py

@@ -7,4 +7,6 @@ from .cuda_driver import (
     get_shared_memory_per_block,  # noqa: F401
     get_device_attribute,  # noqa: F401
     get_max_dynamic_shared_size_bytes,  # noqa: F401
+    get_persisting_l2_cache_max_size,  # noqa: F401
+    get_num_sms,  # noqa: F401
 )

tilelang/carver/arch/driver/cuda_driver.py

@@ -165,3 +165,31 @@ def get_max_dynamic_shared_size_bytes(device_id: int = 0, format: str = "bytes")
             raise RuntimeError("Invalid format. Must be one of: bytes, kb, mb")
     else:
         raise RuntimeError("Failed to get device properties.")
+
+
+def get_persisting_l2_cache_max_size(device_id: int = 0) -> int:
+    prop = get_cuda_device_properties(device_id)
+    if prop:
+        return prop.persistingL2CacheMaxSize
+    else:
+        raise RuntimeError("Failed to get device properties for persisting L2 cache max size.")
+
+
+def get_num_sms(device_id: int = 0) -> int:
+    """
+    Get the number of streaming multiprocessors (SMs) on the CUDA device.
+
+    Args:
+        device_id (int, optional): The CUDA device ID. Defaults to 0.
+
+    Returns:
+        int: The number of SMs on the device.
+
+    Raises:
+        RuntimeError: If unable to get the device properties.
+    """
+    prop = get_cuda_device_properties(device_id)
+    if prop:
+        return prop.multiProcessorCount
+    else:
+        raise RuntimeError("Failed to get device properties.")

tilelang/contrib/dlpack.py

@@ -38,7 +38,7 @@ def convert_func(tvm_func, tensor_type, to_dlpack_func):
 
     float8_dtype_map = {
         torch.float8_e4m3fn: "e4m3_float8",
-        torch.float8_e4m3fnuz: "e4m3_float8",
+        torch.float8_e4m3fnuz: "float8_e4m3fnuz",
         torch.float8_e5m2: "e5m2_float8",
         torch.float8_e5m2fnuz: "e5m2_float8",
     }

tilelang/contrib/nvcc.py

@@ -270,7 +270,7 @@ def get_target_compute_version(target=None):
     Returns
     -------
     compute_version : str
-        compute capability of a GPU (e.g. "8.6")
+        compute capability of a GPU (e.g. "8.6" or "9.0")
     """
     # 1. input target object
     # 2. Target.current()
@@ -279,10 +279,17 @@ def get_target_compute_version(target=None):
         arch = target.arch.split("_")[1]
         if len(arch) == 2:
             major, minor = arch
+            # Handle old format like sm_89
             return major + "." + minor
         elif len(arch) == 3:
-            # This is for arch like "sm_90a"
-            major, minor, suffix = arch
+            major = int(arch[0])
+            if major < 2:
+                major = arch[0:2]
+                minor = arch[2]
+                return major + "." + minor
+            else:
+                # This is for arch like "sm_90a"
+                major, minor, suffix = arch
             return major + "." + minor + "." + suffix
 
     # 3. GPU compute version
@@ -416,6 +423,23 @@ def have_fp8(compute_version):
     return any(conditions)
 
 
+@tvm._ffi.register_func("tvm.contrib.nvcc.supports_tma", override=True)
+def have_tma(target):
+    """Whether TMA support is provided in the specified compute capability or not
+
+    Parameters
+    ----------
+    target : tvm.target.Target
+        The compilation target
+    """
+    compute_version = get_target_compute_version(target)
+    major, minor = parse_compute_version(compute_version)
+    # TMA is supported in Ada Lovelace (9.0) or later architectures.
+    conditions = [False]
+    conditions.append(major >= 9)
+    return any(conditions)
+
+
 def get_nvcc_compiler() -> str:
     """Get the path to the nvcc compiler"""
     return os.path.join(find_cuda_path(), "bin", "nvcc")

tilelang/engine/phase.py

@@ -4,23 +4,41 @@ from tvm import tir, IRModule
 from tvm.target import Target
 import tilelang
 from tilelang.transform import PassContext
+from tilelang.contrib.nvcc import have_tma
 from typing import Optional
 
 
+def allow_warp_specialized(pass_ctx: Optional[PassContext] = None,
+                           target: Optional[Target] = None) -> bool:
+    # avoid circular import
+    from tilelang.jit.adapter.utils import is_cuda_target
+
+    if pass_ctx is None:
+        pass_ctx = tilelang.transform.get_pass_context()
+    if not is_cuda_target(target):
+        return False
+    disable_warp_specialized = pass_ctx.config.get("tl.disable_warp_specialized", False)
+    return not disable_warp_specialized
+
+
 def allow_tma_and_warp_specialized(pass_ctx: Optional[PassContext] = None,
                                    target: Optional[Target] = None) -> bool:
+    # avoid circular import
+    from tilelang.jit.adapter.utils import is_cuda_target
+
     if pass_ctx is None:
         pass_ctx = tilelang.transform.get_pass_context()
-    if target.arch not in {"sm_90", "sm_90a"}:
+    if not is_cuda_target(target) or not have_tma(target):
         return False
     disable_tma_lower = pass_ctx.config.get("tl.disable_tma_lower", False)
-    disable_tma_lower = pass_ctx.config.get("tl.disable_tma_lower", False)
-    disable_warp_specialized = pass_ctx.config.get("tl.disable_warp_specialized", False)
-    return not (disable_tma_lower and disable_warp_specialized)
+    return not disable_tma_lower and allow_warp_specialized(pass_ctx=pass_ctx, target=target)
 
 
 def allow_fence_proxy(target: Optional[Target] = None) -> bool:
-    return target.arch in {"sm_90", "sm_90a"}
+    # avoid circular import
+    from tilelang.jit.adapter.utils import is_cuda_target
+
+    return is_cuda_target(target) and have_tma(target)
 
 
 def allow_vectorize(pass_ctx: Optional[PassContext] = None) -> bool:
@@ -42,6 +60,8 @@ def LowerAndLegalize(mod: IRModule, target: Target) -> IRModule:
     mod = tilelang.transform.LayoutInference()(mod)
     # Lower high-level tile operations to low-level operations
     mod = tilelang.transform.LowerTileOp()(mod)
+    # Lower l2 persistent map
+    mod = tilelang.transform.LowerL2Persistent()(mod)
     # Legalize vectorized loops to ensure they are valid
     mod = tilelang.transform.LegalizeVectorizedLoop()(mod)
     # Add safety checks for memory accesses
@@ -86,12 +106,12 @@ def OptimizeForTarget(mod: IRModule, target: Target) -> IRModule:
             mod = tilelang.transform.InjectFenceProxy()(mod)
 
     mod = tir.transform.LowerOpaqueBlock()(mod)
-    mod = tilelang.transform.FlattenBuffer()(mod)
     mod = tir.transform.NarrowDataType(32)(mod)
+    mod = tilelang.transform.ConfigIndexBitwidth()(mod)
+    mod = tilelang.transform.FlattenBuffer()(mod)
     mod = tir.transform.Simplify()(mod)
 
     mod = tilelang.transform.VectorizeLoop(enable_vectorize=allow_vectorize(pass_ctx=pass_ctx))(mod)
-
     mod = tir.transform.StorageRewrite()(mod)
     mod = tir.transform.UnrollLoop()(mod)
     mod = tir.transform.RenormalizeSplitPattern()(mod)
@@ -115,15 +135,28 @@ def OptimizeForTarget(mod: IRModule, target: Target) -> IRModule:
     mod = tir.transform.InferFragment()(mod)
     mod = tir.transform.LowerThreadAllreduce()(mod)
     mod = tilelang.transform.LowerHopperIntrin()(mod)
-    mod = tilelang.transform.ConfigIndexBitwidth()(mod)
+
+    # Global Barrier Synchronization must be applied before
+    # SplitHostDevice pass, as the global barrier
+    mod = tilelang.transform.ThreadSync("global")(mod)
+    mod = tilelang.transform.AnnotateDeviceRegions()(mod)
+    mod = tir.transform.SplitHostDevice()(mod)
+
+    if allow_warp_specialized(pass_ctx=pass_ctx, target=target):
+        # This is a workaround to avoid the bug in the MergeSharedMemoryAllocations pass
+        # when warp specialization is enabled, as different warp threads may access different
+        # buffers, but the liveness analysis is hard because we need to do pipeline.
+        mod = tir.transform.MergeSharedMemoryAllocations()(mod)
+    else:
+        mod = tilelang.transform.MergeSharedMemoryAllocations()(mod)
+
     mod = tilelang.transform.ThreadSync("shared")(mod)
     mod = tilelang.transform.ThreadSync("shared.dyn")(mod)
     mod = tilelang.transform.EliminateStorageSyncForMBarrier()(mod)
+    # Inject PTX async copy must behind the thread sync pass
+    # as ptx async copy won't be recognized as a valid buffer load
     mod = tilelang.transform.InjectPTXAsyncCopy()(mod)
 
-    mod = tilelang.transform.AnnotateDeviceRegions()(mod)
-    mod = tir.transform.SplitHostDevice()(mod)
-    mod = tir.transform.MergeSharedMemoryAllocations()(mod)
     mod = tilelang.transform.MakePackedAPI()(mod)
     mod = tir.transform.LowerDeviceKernelLaunch()(mod)
 

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