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

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,

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)