tilelang.language.utils¶
Utils in TileLang operators.
Functions¶
|
Create a tl.region call for a BufferLoad and extents. |
|
Clamp extents and return a tl.region call. |
|
Convert a flat (linear) index into multi-dimensional coordinates for a given shape. |
|
Compute a flat (linear) index from multi-dimensional coordinates and strides. |
|
Get the buffer region from a buffer load. |
|
Return the inferred extent (shape) of a buffer-like object. |
Module Contents¶
- tilelang.language.utils.region(buffer, access_type, *args)¶
Create a tl.region call for a BufferLoad and extents.
- Parameters:
buffer (tvm.tir.BufferLoad)
access_type (str)
args (tvm.tir.PrimExpr)
- Return type:
tvm.tir.PrimExpr
- tilelang.language.utils.buffer_region_to_tile_region(buffer_region, access_type, extents)¶
Clamp extents and return a tl.region call.
- Parameters:
buffer_region (tvm.tir.BufferRegion)
access_type (str)
extents (list[tvm.tir.PrimExpr])
- Return type:
tvm.tir.PrimExpr
- tilelang.language.utils.index_to_coordinates(index, shape)¶
Convert a flat (linear) index into multi-dimensional coordinates for a given shape.
Given a linear index and a shape (sequence of dimension extents), returns a list of coordinates (one per dimension) such that converting those coordinates back to a linear index using the usual row-major / C-order formula yields the original index. The computation iterates from the last dimension to the first using modulo and integer division, then reverses the collected coordinates.
- Parameters:
index (int or PrimExpr) – The flat index to convert.
shape (Sequence[int]) – The extents of each dimension (length >= 1).
- Returns:
Coordinates for each dimension in the same order as shape.
- Return type:
List[PrimExpr]
- tilelang.language.utils.linear_index(*args)¶
Compute a flat (linear) index from multi-dimensional coordinates and strides.
The function accepts a sequence of PrimExpr arguments where the first portion are coordinates and the trailing portion are the corresponding strides. The number of strides must equal (number of coordinates - 1). The linear index is computed as:
linear = coords[0] for each (coord, stride) in zip(coords[1:], strides):
linear = linear * stride + coord
Examples
linear_index(i) -> i
linear_index(i, j) -> i * j_stride + j (requires j_stride provided as stride when needed)
linear_index(i, j, stride_j) -> i * stride_j + j
linear_index(i, j, k, stride_j, stride_k) -> i*stride_j*stride_k + j*stride_k + k
linear_index(i, tx, v, threads, local_size) -> i*threads*local_size + tx*local_size + v
- Raises:
ValueError – If called with no arguments, or if the number of strides is not one less than the number of coordinates.
- Returns:
The computed linear index expression.
- Return type:
PrimExpr
- Parameters:
args (tvm.tir.PrimExpr)
- tilelang.language.utils.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.
If the buffer load has ramp indices, we will use the ramp’s base and lanes to create the region. Otherwise, return None since the load cannot be converted to a region.
- Parameters:
buffer_load (tvm.tir.BufferLoad)
extents (list[tvm.tir.PrimExpr] | None)
- Return type:
tvm.tir.BufferRegion | None
- tilelang.language.utils.get_extent(data)¶
Return the inferred extent (shape) of a buffer-like object.
If data is a Var bound to a let value, the let value is resolved before inspection.
- Parameters:
data (tilelang._typing.BufferLikeType) – A Var, Buffer, BufferLoad or BufferRegion to inspect.
- Returns:
The shape/extents as a list-like of PrimExpr (Buffer.shape or list of region item extents), or None if the extent cannot be determined.
- Return type:
tilelang._typing.ShapeType | None