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

tilelang/jit/kernel.py

@@ -31,11 +31,16 @@ class JITKernel(object):
     torch_function : Callable
         The compiled function that can be invoked as a PyTorch-compatible function.
     """
-
+    prim_func: PrimFunc = None
     artifact: CompiledArtifact = None
     adapter: BaseKernelAdapter = None
     torch_function: Callable = None
 
+    # tuner result
+    latency: float = None
+    config: Dict[str, Any] = None
+    ref_latency: float = None
+
     def __init__(
         self,
         func: PrimFunc = None,
@@ -74,6 +79,7 @@ class JITKernel(object):
         from_database : bool, optional
             Whether to create a TorchFunction from a database.
         """
+        self.prim_func = func
         self.execution_backend = execution_backend
         self.target = target
         self.target_host = target_host
@@ -344,6 +350,51 @@ class JITKernel(object):
     def run_once(self, func: Optional[Callable] = None) -> None:
         return self.get_profiler().run_once(func)
 
+    def update_tuner_result(self, latency: float, config: Dict[str, Any],
+                            ref_latency: float) -> "JITKernel":
+        """
+        Updates the tuning results for this kernel.
+
+        Parameters
+        ----------
+        latency : float
+            The measured latency of this kernel configuration.
+        config : Dict[str, Any]
+            The configuration parameters used for this kernel.
+        ref_latency : float
+            The reference latency to compare against.
+
+        Returns
+        -------
+        None
+        """
+        self.latency = latency
+        self.config = config
+        self.ref_latency = ref_latency
+
+        return self
+
+    def get_tuner_result(self) -> Dict[str, Any]:
+        """
+        Gets the tuning results for this kernel.
+
+        Returns
+        -------
+        Dict[str, Any]
+            A dictionary containing:
+            - latency: The measured latency of this kernel
+            - config: The configuration parameters used
+            - ref_latency: The reference latency for comparison
+        """
+        if self.latency is None:
+            raise ValueError("Tuning results are not available. Please tune the kernel first.")
+
+        return {
+            "latency": self.latency,
+            "config": self.config,
+            "ref_latency": self.ref_latency,
+        }
+
     @property
     def out_idx(self) -> List[int]:
         return self.adapter.result_idx

tilelang/jit/param.py

@@ -0,0 +1,45 @@
+# Copyright (c) Tile-AI Corporation.
+# Licensed under the MIT License.
+
+from typing import (
+    Any,
+    TypeVar,
+)
+from typing_extensions import ParamSpec
+
+
+# --- Mocking dependencies for the example to run ---
+# In your actual code, these would be your real types.
+class Program:
+    """Placeholder for the type returned by the original decorated function."""
+
+    def __init__(self, data: str):
+        self.data = data
+
+    def __repr__(self):
+        return f"Program('{self.data}')"
+
+
+class Kernel:
+    """Placeholder for the type of the compiled kernel."""
+
+    def __init__(self, source: str, out_idx: Any):
+        self.source_code = source
+        self.out_idx = out_idx
+
+    def get_kernel_source(self) -> str:
+        return self.source_code
+
+    def __repr__(self):
+        return f"Kernel('{self.source_code[:20]}...')"
+
+
+# --- End Mocking ---
+
+# P (Parameters) captures the argument types of the decorated function.
+_P = ParamSpec("_P")
+# R_prog (Return type of Program) captures the return type of the original decorated function.
+# We assume the original function returns something compatible with 'Program'.
+_RProg = TypeVar("_RProg", bound=Program)
+
+__all__ = ["Program", "Kernel", "_P", "_RProg"]

tilelang/language/__init__.py

@@ -57,6 +57,7 @@ from .print import print  # noqa: F401
 from .customize import (
     atomic_add,  # noqa: F401
     atomic_addx2,  # noqa: F401
+    atomic_addx4,  # noqa: F401
     dp4a,  # noqa: F401
     clamp,  # noqa: F401
     reshape,  # noqa: F401
@@ -81,12 +82,93 @@ def use_swizzle(panel_size: int, order: str = "row", enable: bool = True):
                 f"tl::{device_func}<{panel_size}>") if enable else None
 
 
-def annotate_layout(layout_map):
+def annotate_layout(layout_map: Dict):
+    """Annotate the layout of the buffer
+
+    Args:
+        layout_map (Dict): a dictionary of buffer to layout
+
+    Returns:
+        block_attr: a block attribute
+    
+    Example:
+        @T.prim_func
+        def main(
+                A: T.Tensor((M, N), dtype),
+                B: T.Tensor((M, N), dtype),
+        ):
+            # Initialize Kernel Context
+            with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):
+                A_shared = T.alloc_shared((block_M, block_N), dtype)
+
+                T.annotate_layout({A_shared: layout})
+                for i, j in T.Parallel(block_M, block_N):
+                    A_shared[i, j] = A[by * block_M + i, bx * block_N + j]
+
+                for i, j in T.Parallel(block_M, block_N):
+                    B[by * block_M + i, bx * block_N + j] = A_shared[i, j]
+
+        return main
+    """
     # layout_map is a dictionary of buffer to layout
     layout_map = {buffer.data: layout for buffer, layout in layout_map.items()}
     return block_attr({"layout_map": layout_map})
 
 
+def annotate_padding(padding_map: Dict):
+    """Annotate the padding of the buffer
+
+    Args:
+        padding_map (dict): a dictionary of buffer to padding value
+
+    Returns:
+        block_attr: a block attribute
+    
+    Example:
+        @T.prim_func
+        def main(
+                A: T.Tensor((M, N), dtype),
+                B: T.Tensor((M, N), dtype),
+        ):
+            # Initialize Kernel Context
+            with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):
+                A_shared = T.alloc_shared((block_M, block_N), dtype)
+
+                T.annotate_padding({A_shared: pad_value})
+                for i, j in T.Parallel(block_M, block_N):
+                    A_shared[i, j] = A[by * block_M + i - 10, bx * block_N + j]
+
+                for i, j in T.Parallel(block_M, block_N):
+                    B[by * block_M + i, bx * block_N + j] = A_shared[i, j]
+
+        return main
+    """
+    # padding_map is a dictionary of buffer to padding value
+    _padding_map = {}
+    for buffer, padding_value in padding_map.items():
+        # assert not global
+        assert buffer.scope() != "global", "padding can not be applied to global buffers"
+        _padding_map[buffer.data] = padding_value
+    return block_attr({"padding_map": _padding_map})
+
+
+def annotate_l2_hit_ratio(l2_hit_ratio_map: Dict):
+    """Annotate the L2 hit ratio of the buffer, detailed explanation please refer to:
+    https://docs.nvidia.com/cuda/cuda-c-programming-guide/#l2-policy-for-persisting-accesses
+
+    Args:
+        l2_hit_ratio_map (dict): a dictionary of buffer to L2 hit ratio value
+    Example:
+        # 0.5 is the hit ratio
+        T.annotate_l2_hit_ratio({A: 0.5})
+    """
+    _l2_hit_ratio_map = {}
+    for buffer, hit_ratio in l2_hit_ratio_map.items():
+        assert buffer.scope() == "global", "persistent L2 can only be applied to global buffers"
+        _l2_hit_ratio_map[buffer.data] = hit_ratio
+    return block_attr({"l2_hit_ratio_map": _l2_hit_ratio_map})
+
+
 def import_source(source: Optional[str] = None):
     # source is the source code to be imported
     return block_attr({"pragma_import_c": source}) if source is not None else None

tilelang/language/builtin.py

@@ -3,7 +3,8 @@
 """The language interface for tl programs."""
 
 from tilelang import tvm as tvm
-from tilelang.language import ptx_arrive_barrier
+from tilelang.language import ptx_arrive_barrier, evaluate
+from tilelang.language.kernel import get_thread_bindings, get_block_extents
 from tvm import tir
 from typing import Union, Any
 from tvm.tir import PrimExpr, Var, Call
@@ -235,3 +236,91 @@ def wait_wgmma(*args):
         tir.Call: A handle to the WGMMA wait operation
     """
     return tir.call_intrin("handle", tir.op.Op.get("tl.wait_wgmma"), *args)
+
+
+def barrier_wait(barrier_id: Union[int, PrimExpr, tir.Call], parity: Union[int, Var, None] = None):
+    """Wait for a memory barrier to complete.
+
+    Args:
+        barrier_id: Optional[int, PrimExpr]
+            The memory barrier to wait on
+        parity: Optional[int, Var]
+            The parity value to wait for
+    Returns:
+        tir.Call: A handle to the barrier wait operation
+    Current implementation is a sugar syntax for mbarrier_wait_parity, as we only support parity 0 and 1.
+    """
+    return mbarrier_wait_parity(barrier_id, parity)
+
+
+def barrier_arrive(barrier_id: Union[int, PrimExpr, tir.Call]):
+    """Arrive at a memory barrier.
+
+    Args:
+        barrier_id: Optional[int, PrimExpr]
+            The memory barrier to arrive at
+    """
+    return mbarrier_arrive(barrier_id)
+
+
+def shfl_xor(value: Union[int, PrimExpr, tir.Call], offset: Union[int, PrimExpr, tir.Call]):
+    """Perform a shuffle operation with XOR offset.
+
+    Args:
+        value: Optional[int, PrimExpr]
+            The value to shuffle
+        offset: Optional[int, PrimExpr]
+            The offset for the shuffle operation
+    Returns:
+        tir.Call: A handle to the shuffle operation
+    """
+    return tir.call_extern(value.dtype, "__shfl_xor_sync", 0xffffffff, value, offset)
+
+
+def shfl_down(value: Union[int, PrimExpr, tir.Call], offset: Union[int, PrimExpr, tir.Call]):
+    """Perform a shuffle operation with down offset.
+
+    Args:
+        value: Optional[int, PrimExpr]
+            The value to shuffle
+    """
+    return tir.call_extern(value.dtype, "__shfl_down_sync", 0xffffffff, value, offset)
+
+
+def shfl_up(value: Union[int, PrimExpr, tir.Call], offset: Union[int, PrimExpr, tir.Call]):
+    """Perform a shuffle operation with up offset.
+
+    Args:
+        value: Optional[int, PrimExpr]
+            The value to shuffle
+    """
+    return tir.call_extern(value.dtype, "__shfl_up_sync", 0xffffffff, value, offset)
+
+
+def sync_threads():
+    """Synchronize all threads in a warp.
+    """
+    return tir.op.tvm_storage_sync("shared")
+
+
+def sync_thread_partial(barrier_id: Union[int, PrimExpr, tir.Call]):
+    """Synchronize threads within a warp.
+
+    Args:
+        barrier_id: Optional[int, PrimExpr]
+            The memory barrier to synchronize
+
+    Returns:
+        tir.Call: A handle to the synchronization operation
+    """
+    return tir.call_intrin("handle", tir.op.Op.get("tl.sync_thread_partial"), barrier_id)
+
+
+def sync_global():
+    """Synchronize all threads in a block.
+    """
+    tx, ty, tz = get_thread_bindings()
+    ex, ey, ez = get_block_extents()
+    print(tx, ty, tz, ex, ey, ez)
+    args = ["global", tx == 0 and ty == 0 and tz == 0, ex * ey * ez]
+    return evaluate(tir.Call("handle", "tir.tvm_storage_sync", args))

tilelang/language/copy.py

@@ -62,7 +62,8 @@ def buffer_load_to_tile_region(load: tir.BufferLoad, access_type: str, extents:
     return region(load, access_type, *extents)
 
 
-def buffer_region_to_tile_region(buffer_region: tir.BufferRegion, access_type: str):
+def buffer_region_to_tile_region(buffer_region: tir.BufferRegion, access_type: str,
+                                 extents: List[tir.PrimExpr]):
     """Convert a buffer region to a tile region descriptor.
 
     Args:
@@ -73,8 +74,12 @@ def buffer_region_to_tile_region(buffer_region: tir.BufferRegion, access_type: s
         tir.Call: A region descriptor for the specified buffer region
     """
     mins = [x.min for x in buffer_region.region]
-    extents = [x.extent for x in buffer_region.region]
-    return region(T.BufferLoad(buffer_region.buffer, mins), access_type, *extents)
+    region_extents = [x.extent for x in buffer_region.region]
+    assert len(region_extents) >= len(
+        extents
+    ), f"region_extents must be >= extents, region_extents = {region_extents}, extents = {extents}"
+
+    return region(T.BufferLoad(buffer_region.buffer, mins), access_type, *region_extents)
 
 
 def copy(
@@ -110,13 +115,10 @@ def copy(
 
     src_extent = get_extent(src)
     dst_extent = get_extent(dst)
-
-    if src_extent:
-        extent = src_extent
-    elif dst_extent:
-        extent = dst_extent
-    else:
-        raise TypeError("Can't deduce copy extents from args")
+    assert src_extent or dst_extent, "Can't deduce copy extents from args"
+    src_extent = list(src_extent) if src_extent else [1] * len(dst_extent)
+    dst_extent = list(dst_extent) if dst_extent else [1] * len(src_extent)
+    extent = max(src_extent, dst_extent)
 
     def _to_region(data, access_type):
         if isinstance(data, tir.Var) and T.has_let_value(data):
@@ -124,7 +126,7 @@ def copy(
         if isinstance(data, tir.Buffer):
             return buffer_to_tile_region(data, access_type)
         elif isinstance(data, tir.BufferRegion):
-            return buffer_region_to_tile_region(data, access_type)
+            return buffer_region_to_tile_region(data, access_type, extent)
         else:
             return buffer_load_to_tile_region(data, access_type, extent)
 

tilelang/language/customize.py

@@ -33,6 +33,19 @@ def atomic_addx2(dst: Buffer, value: PrimExpr) -> PrimExpr:
     return T.call_extern("handle", "AtomicAddx2", T.address_of(dst), T.address_of(value))
 
 
+def atomic_addx4(dst: Buffer, value: PrimExpr) -> PrimExpr:
+    """Perform an atomic addition operation with double-width operands.
+
+    Args:
+        dst (Buffer): Destination buffer where the atomic addition will be performed
+        value (PrimExpr): Value to be atomically added (double-width)
+
+    Returns:
+        PrimExpr: Handle to the double-width atomic addition operation
+    """
+    return T.call_extern("handle", "AtomicAddx4", T.address_of(dst), T.address_of(value))
+
+
 def dp4a(A: Buffer, B: Buffer, C: Buffer) -> PrimExpr:
     """Perform a 4-element dot product with accumulation (DP4A).
 

tilelang/language/print.py

@@ -44,6 +44,24 @@ def print_var_with_condition(condition: tir.PrimExpr,
         tir.call_extern("handle", "debug_print_var", msg, var)
 
 
+@macro
+def print_global_buffer_with_condition(condition: tir.PrimExpr,
+                                       buffer: tir.Buffer,
+                                       elems: int,
+                                       msg: str = "") -> tir.PrimExpr:
+    """
+    Conditionally prints the values of a flattened TIR buffer if the condition is True.
+    """
+    if condition:
+        # Iterate through the buffer elements and print each one.
+        for i in serial(elems):
+            coords = index_to_coordinates(i, buffer.shape)
+            tir.call_extern("handle", "debug_print_buffer_value", msg, buffer.name, i,
+                            buffer[coords])
+    else:
+        tir.call_extern("handle", "debug_print_buffer_value", msg, buffer.name, i, buffer[coords])
+
+
 @macro
 def print_shared_buffer_with_condition(condition: tir.PrimExpr,
                                        buffer: tir.Buffer,
@@ -172,6 +190,15 @@ def print(obj: Any, msg: str = "") -> tir.PrimExpr:
             if not msg:
                 msg = f"buffer<{buffer.name}, {buffer.dtype}>"
             return print_shared_buffer_with_condition(condition, buffer, elems, msg)
+        elif buffer.scope() == "global":
+            # Get the number of elements in the buffer.
+            elems = 1
+            for dim in buffer.shape:
+                elems *= dim
+            condition = True
+            return print_global_buffer_with_condition(condition, buffer, elems, msg)
+        else:
+            raise ValueError(f"Unsupported buffer scope: {buffer.scope()}")
 
     elif isinstance(obj, tir.PrimExpr):
         if not msg:

tilelang/language/reduce.py

@@ -7,6 +7,12 @@ from typing import Optional
 from tilelang.language import copy, macro, alloc_shared
 
 
+def _legalize_dim(buffer: tir.Buffer, dim: int):
+    if dim < 0:
+        dim = len(buffer.shape) + dim
+    return dim
+
+
 def reduce(buffer: tir.Buffer, out: tir.Buffer, reduce_type: str, dim: int, clear: bool):
     """Perform a reduction operation on a buffer along a specified dimension.
 
@@ -33,7 +39,7 @@ def reduce(buffer: tir.Buffer, out: tir.Buffer, reduce_type: str, dim: int, clea
     )
 
 
-def reduce_max(buffer: tir.Buffer, out: tir.Buffer, dim: int, clear: bool = True):
+def reduce_max(buffer: tir.Buffer, out: tir.Buffer, dim: int = -1, clear: bool = True):
     """Perform reduce max on input buffer, store the result to output buffer
 
     Parameters
@@ -50,10 +56,11 @@ def reduce_max(buffer: tir.Buffer, out: tir.Buffer, dim: int, clear: bool = True
     -------
     handle : PrimExpr
     """
+    dim = _legalize_dim(buffer, dim)
     return reduce(buffer, out, "max", dim, clear)
 
 
-def reduce_min(buffer: tir.Buffer, out: tir.Buffer, dim: int, clear: bool = True):
+def reduce_min(buffer: tir.Buffer, out: tir.Buffer, dim: int = -1, clear: bool = True):
     """Perform reduce min on input buffer, store the result to output buffer.
 
     Args:
@@ -65,10 +72,11 @@ def reduce_min(buffer: tir.Buffer, out: tir.Buffer, dim: int, clear: bool = True
     Returns:
         tir.Call: Handle to the reduction operation
     """
+    dim = _legalize_dim(buffer, dim)
     return reduce(buffer, out, "min", dim, clear)
 
 
-def reduce_sum(buffer: tir.Buffer, out: tir.Buffer, dim: int, clear: bool = True):
+def reduce_sum(buffer: tir.Buffer, out: tir.Buffer, dim: int = -1, clear: bool = True):
     """Perform reduce sum on input buffer, store the result to output buffer.
 
     Args:
@@ -89,10 +97,11 @@ def reduce_sum(buffer: tir.Buffer, out: tir.Buffer, dim: int, clear: bool = True
     Returns:
         tir.Call: Handle to the reduction operation
     """
+    dim = _legalize_dim(buffer, dim)
     return reduce(buffer, out, "sum", dim, clear)
 
 
-def reduce_abssum(buffer: tir.Buffer, out: tir.Buffer, dim: int):
+def reduce_abssum(buffer: tir.Buffer, out: tir.Buffer, dim: int = -1):
     """Perform reduce absolute sum on input buffer, store the result to output buffer.
 
     Args:
@@ -103,10 +112,11 @@ def reduce_abssum(buffer: tir.Buffer, out: tir.Buffer, dim: int):
     Returns:
         tir.Call: Handle to the reduction operation
     """
+    dim = _legalize_dim(buffer, dim)
     return reduce(buffer, out, "abssum", dim, True)
 
 
-def reduce_absmax(buffer: tir.Buffer, out: tir.Buffer, dim: int, clear: bool = True):
+def reduce_absmax(buffer: tir.Buffer, out: tir.Buffer, dim: int = -1, clear: bool = True):
     """Perform reduce absolute max on input buffer, store the result to output buffer.
 
     Args:
@@ -117,6 +127,7 @@ def reduce_absmax(buffer: tir.Buffer, out: tir.Buffer, dim: int, clear: bool = T
     Returns:
         tir.Call: Handle to the reduction operation
     """
+    dim = _legalize_dim(buffer, dim)
     return reduce(buffer, out, "absmax", dim, clear)
 
 

tilelang/language/tir/op.py

@@ -2603,6 +2603,21 @@ def isinf(x, span=None):
     return _tvm_op.isinf(x, span)
 
 
+def pow_of_int(x: PrimExpr, y: int) -> PrimExpr:
+    """Fast power operation than pow(float, float).
+    
+    Args:
+        x (PrimExpr): Base value
+        y (int): Exponent value
+    """
+    return call_intrin(
+        x.dtype,
+        tvm.tir.op.Op.get("tl.pow_of_int"),
+        x,
+        y,
+    )
+
+
 def power(x, y, span=None):
     """x power y
 
@@ -2622,6 +2637,8 @@ def power(x, y, span=None):
     z : PrimExpr
         The result.
     """
+    if isinstance(y, (int, IntImm)):
+        return pow_of_int(x, y)
     return _tvm_op.power(x, y, span)
 
 
@@ -2644,6 +2661,8 @@ def pow(x, y, span=None):
     z : PrimExpr
         The result.
     """
+    if isinstance(y, (int, IntImm)):
+        return pow_of_int(x, y)
     return _tvm_op.pow(x, y, span)
 
 

tilelang/language/warpgroup.py

@@ -37,8 +37,13 @@ def WarpSpecialize(*warp_group_idx):
         >>> T.ws(0, 1) -> if tx < 128 or (tx >= 128 and tx < 256)
     """
     id_x, id_y, id_z = get_thread_bindings()
-    ex_x, ex_y, _ = get_thread_extents()
-    tid = id_z * (ex_y * ex_x) + id_y * ex_x + id_x
+    ex_x, ex_y, ex_z = get_thread_extents()
+    tid = id_x
+    if ex_y > 1:
+        tid = id_y * ex_x + tid
+    if ex_z > 1:
+        tid = id_z * (ex_y * ex_x) + tid
+
     # only available for nvidia gpus.
     warp_group_size = 128