tilelang.utils.language ======================= .. py:module:: tilelang.utils.language Functions --------- .. autoapisummary:: tilelang.utils.language.is_global tilelang.utils.language.is_shared tilelang.utils.language.is_shared_dynamic tilelang.utils.language.is_tensor_memory tilelang.utils.language.is_local tilelang.utils.language.is_fragment tilelang.utils.language.get_buffer_elems tilelang.utils.language.array_reduce tilelang.utils.language.retrieve_func_from_module tilelang.utils.language.get_buffer_region_from_load tilelang.utils.language.to_buffer_region tilelang.utils.language.retrieve_shape tilelang.utils.language.retrieve_stride tilelang.utils.language.retrive_ptr_from_buffer_region tilelang.utils.language.retrieve_ptr tilelang.utils.language.retrieve_offset tilelang.utils.language.bits_product tilelang.utils.language.prim_expr_equal tilelang.utils.language.legalize_pairwise_extents tilelang.utils.language.is_full_region Module Contents --------------- .. py:function:: is_global(buffer) Check if the buffer is in the global memory scope. :param buffer: The TVM buffer, BufferLoad, or BufferRegion to check. :returns: True if the buffer is in global memory, False otherwise. :rtype: bool .. py:function:: is_shared(buffer, allow_dynamic = True) Check if the buffer is in the shared memory scope. :param buffer: The TVM buffer, BufferLoad, or BufferRegion to check. :returns: True if the buffer is in shared memory, False otherwise. :rtype: bool .. py:function:: is_shared_dynamic(buffer) Check if the buffer is in the dynamic shared memory scope. :param buffer: The TVM buffer, BufferLoad, or BufferRegion to check. :returns: True if the buffer is in dynamic shared memory, False otherwise. :rtype: bool .. py:function:: is_tensor_memory(buffer) Check if the buffer is in tensor memory scope (e.g., shared.tmem). :param buffer: The TVM buffer, BufferLoad, or BufferRegion to check. :returns: True if the buffer is in tensor memory, False otherwise. :rtype: bool .. py:function:: is_local(buffer) Check if the buffer is in the local memory scope. :param buffer: The TVM buffer, BufferLoad, or BufferRegion to check. :returns: True if the buffer is in local memory, False otherwise. :rtype: bool .. py:function:: is_fragment(buffer) Check if the buffer is a fragment (e.g., for matrix multiplication operations). :param buffer: The TVM buffer, BufferLoad, or BufferRegion to check. :returns: True if the buffer is a fragment, False otherwise. :rtype: bool .. py:function:: get_buffer_elems(buffer) Get the number of elements in the buffer. .. py:function:: array_reduce(array) Reduce an array of integers to a single integer. :param array: The array of integers to reduce. :type array: List[int] :returns: The reduced integer. :rtype: int .. py:function:: retrieve_func_from_module(ir_module) Retrieve the single PrimFunc from an IRModule. :param ir_module: The TVM IRModule to extract the function from. The module should contain exactly one global function. :type ir_module: IRModule :returns: The single function contained in the module. :rtype: PrimFunc :raises ValueError: If ir_module is not an IRModule. :raises AssertionError: If the module contains more than one global function. .. py:function:: get_buffer_region_from_load(buffer_load, extents = None) Get the buffer region from a buffer load. May encounter buffer load like C[0:128, 0:32], ref to pull request for buffer wise op: https://github.com/apache/tvm/pull/14693 convert load to region .. py:function:: to_buffer_region(obj, access_type = 'rw', extents = None) Convert to/from the tl.region representation. - Buffer/BufferLoad/BufferRegion -> returns a tl.region call (PrimExpr) - tl.region Call -> returns the decoded BufferRegion for analysis .. py:function:: retrieve_shape(obj) Retrieve shape-like extents for a buffer-like object. - Buffer -> its `shape` - BufferRegion -> list of each range's `extent` - BufferLoad -> extents from `get_buffer_region_from_load(obj)` .. py:function:: retrieve_stride(obj) Retrieve row-major strides for a buffer-like object based on its buffer.shape. For BufferRegion and BufferLoad, uses the underlying buffer's `shape`. .. py:function:: retrive_ptr_from_buffer_region(buffer_or_load_or_region, access_type = 'r') .. py:function:: retrieve_ptr(obj, access_type = 'r', ignore_last_ndim = 0) Retrieve a pointer to the start of a (possibly sliced) buffer region. - Buffer -> base pointer - BufferRegion -> pointer with byte offset computed from region minima - BufferLoad -> pointer offset computed from indices or derived region :param obj: Buffer-like object :param access_type: TVM Buffer access mask, e.g. "r", "w", "rw" :param ignore_last_ndim: do not offset the last N dimensions .. py:function:: retrieve_offset(obj) Retrieve per-dimension minima offsets. - Buffer -> [0, 0, ...] - BufferRegion -> [r.min for r in region] - BufferLoad -> indices (or derived region minima) .. py:function:: bits_product(shape, dtype) Compute the number of bits in a Buffer (shape with dtype). .. py:function:: prim_expr_equal(lhs, rhs) Robust equality for PrimExpr shapes/extents. Tries structural_equal first, then falls back to expr_deep_equal. Python ints are converted to IntImm for comparison. .. py:function:: legalize_pairwise_extents(src_extents, dst_extents) Right-align and broadcast two extent lists to be mutually compatible. Early-exit rule: - If the number of non-1 dimensions in `src_extents` equals that in `dst_extents`, no adjustment is made; the original extents are returned unchanged. This preserves the per-dimension iteration mapping (one loop var per non-1 dim) and avoids creating extra varying axes on either side. Otherwise, for each pair of tail-aligned dimensions (x, y): - if x == y: keep both - elif x == 1: set x = y - elif y == 1: set y = x - else: promote both to tir.max(x, y) to handle dynamic-vs-static safely Leading unmatched dimensions are kept as-is. Returns a tuple of new lists (src_new, dst_new). .. py:function:: is_full_region(buffer_region) Check whether a BufferRegion covers the full buffer region. A full region means each dimension has start 0 and extent equal to the corresponding dimension in the buffer's shape. :param buffer_region: The TVM BufferRegion to check. :returns: True if the region is full; otherwise False. :rtype: bool