wehub-resource-sync
578 lines
20 KiB
Python
578 lines
20 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
"""Abstraction for array data structures."""
|
|
|
|
import functools
|
|
from numbers import Integral
|
|
|
|
import tvm_ffi
|
|
|
|
import tvm
|
|
from tvm.ir import PointerType, PrimType, Range
|
|
from tvm.runtime import Object, Scriptable, convert
|
|
|
|
from . import _ffi_api
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.Buffer")
|
|
class Buffer(Object, Scriptable):
|
|
"""Symbolic data buffer in TVM.
|
|
|
|
Buffer provide a way to represent data layout
|
|
specialization of data structure in TVM.
|
|
|
|
Do not construct directly, use :py:func:`~decl_buffer` instead.
|
|
See the documentation of :py:func:`decl_buffer` for more details.
|
|
|
|
See Also
|
|
--------
|
|
decl_buffer : Declare a buffer
|
|
"""
|
|
|
|
READ = 1
|
|
WRITE = 2
|
|
|
|
def access_ptr(self, access_mask, ptr_type="handle", content_lanes=1, offset=0, extent=None):
|
|
"""Get an access pointer to the head of buffer.
|
|
|
|
This is the recommended method to get buffer data
|
|
ptress when interacting with external functions.
|
|
|
|
Parameters
|
|
----------
|
|
access_mask : int
|
|
The access pattern MASK. Indicate whether the
|
|
access will read or write to the data content.
|
|
|
|
ptr_type : str or tvm.ir.Type, optional
|
|
The data type of the result pointer. Do not specify
|
|
unless we want to cast pointer to specific type.
|
|
|
|
content_lanes: int, optional
|
|
The number of lanes for the data type. This value
|
|
is greater than one for vector types.
|
|
|
|
offset: Expr, optional
|
|
The offset of pointer. We can use it to offset by
|
|
the number of elements from the address of ptr.
|
|
|
|
extent: Expr, optional
|
|
The extent of pointer.
|
|
|
|
Examples
|
|
--------
|
|
.. code-block:: python
|
|
|
|
# Get access ptr for read
|
|
buffer.access_ptr("r")
|
|
# Get access ptr for read/write with bitmask
|
|
buffer.access_ptr(Buffer.READ | Buffer.WRITE)
|
|
# Get access ptr for read/write with str flag
|
|
buffer.access_ptr("rw")
|
|
# Get access ptr for read with offset
|
|
buffer.access_ptr("r", offset = 100)
|
|
# Get access ptr for read with extent
|
|
buffer.access_ptr("r", extent = 100)
|
|
"""
|
|
if isinstance(access_mask, str):
|
|
mask = 0
|
|
for value in access_mask:
|
|
if value == "r":
|
|
mask = mask | Buffer.READ
|
|
elif value == "w":
|
|
mask = mask | Buffer.WRITE
|
|
else:
|
|
raise ValueError(f"Unknown access_mask {access_mask}")
|
|
access_mask = mask
|
|
if isinstance(ptr_type, str):
|
|
ptr_type = (
|
|
PointerType(PrimType("void"))
|
|
if ptr_type == "handle"
|
|
else PointerType(PrimType(ptr_type))
|
|
)
|
|
elif isinstance(ptr_type, PrimType):
|
|
ptr_type = PointerType(ptr_type)
|
|
offset = convert(offset)
|
|
extent = convert(extent)
|
|
return _ffi_api.BufferAccessPtr(
|
|
self,
|
|
access_mask,
|
|
ptr_type,
|
|
content_lanes,
|
|
offset,
|
|
extent, # type: ignore
|
|
)
|
|
|
|
def vload(self, begin, dtype=None, predicate=None):
|
|
"""Generate an Expr that loads dtype from begin index.
|
|
|
|
Parameters
|
|
----------
|
|
begin : Array of Expr
|
|
The beginning index in unit of Buffer.dtype
|
|
|
|
dtype : str
|
|
The data type to be loaded,
|
|
can be vector type which have lanes that is multiple of Buffer.dtype
|
|
|
|
predicate : Optional[Expr]
|
|
A vector mask of boolean values indicating which lanes of a vector are to be
|
|
loaded. The number lanes of the mask must be equal to the number of lanes being loaded.
|
|
|
|
Returns
|
|
-------
|
|
load : Expr
|
|
The corresponding load expression.
|
|
"""
|
|
begin = (begin,) if isinstance(begin, int) or tvm.ir.is_prim_expr(begin) else begin
|
|
dtype = dtype if dtype else self.dtype
|
|
return _ffi_api.BufferVLoad(self, begin, dtype, predicate) # type: ignore
|
|
|
|
def vstore(self, begin, value, predicate=None):
|
|
"""Generate a Stmt that store value into begin index.
|
|
|
|
Parameters
|
|
----------
|
|
begin : Array of Expr
|
|
The beginning index in unit of Buffer.dtype
|
|
|
|
value : Expr
|
|
The value to be stored.
|
|
|
|
predicate : Optional[Expr]
|
|
A vector mask of boolean values indicating which lanes of a vector are to be
|
|
stored. The number lanes of the mask must be equal to the number of lanes in
|
|
value.
|
|
|
|
Returns
|
|
-------
|
|
store : Stmt
|
|
The corresponding store stmt.
|
|
"""
|
|
begin = (begin,) if isinstance(begin, int) or tvm.ir.is_prim_expr(begin) else begin
|
|
return _ffi_api.BufferVStore(self, begin, value, predicate) # type: ignore
|
|
|
|
def scope(self):
|
|
"""Return the storage scope associated with this buffer.
|
|
Returns
|
|
-------
|
|
scope : str
|
|
The storage scope associated with this buffer.
|
|
"""
|
|
return _ffi_api.BufferStorageScope(self) # type: ignore
|
|
|
|
def get_flattened_buffer(self):
|
|
"""Generate a Buffer that is a flattened version of this buffer.
|
|
|
|
Returns
|
|
-------
|
|
flattened : Buffer
|
|
The corresponding flat buffer.
|
|
"""
|
|
return _ffi_api.BufferGetFlattenedBuffer(self) # type: ignore
|
|
|
|
def with_allocated_addr(self, allocated_addr):
|
|
"""Return a new buffer with the allocated address."""
|
|
return _ffi_api.BufferWithAllocatedAddr(self, allocated_addr) # type: ignore
|
|
|
|
def with_dtype(self, dtype):
|
|
"""Return a new buffer with the dtype."""
|
|
return _ffi_api.BufferWithDtype(self, dtype) # type: ignore
|
|
|
|
def with_data(self, data):
|
|
"""Return a new buffer with the data."""
|
|
return _ffi_api.BufferWithData(self, data) # type: ignore
|
|
|
|
def offset_of(self, indices):
|
|
"""Determine the offset of the provided indices in the flattened buffer.
|
|
|
|
Parameters
|
|
----------
|
|
indices : Union[Expr, List[Expr]]
|
|
|
|
The indices of the element in the original buffer.
|
|
|
|
Returns
|
|
-------
|
|
flattened_indices: List[Expr]
|
|
|
|
The offset indices of the element in the flattened buffer.
|
|
"""
|
|
return _ffi_api.BufferOffsetOf(self, indices) # type: ignore
|
|
|
|
@property
|
|
def byte_offset(self):
|
|
"""Get the byte offset of the buffer."""
|
|
return self.elem_offset * tvm.DataType(self.dtype).bits // 8
|
|
|
|
def elem_offset_of(self, indices, inner=True):
|
|
"""Get the element offset of the buffer at the given indices.
|
|
Note that indices subject to buffer's layout mapping.
|
|
|
|
Parameters
|
|
----------
|
|
indices : Union[Expr, List[Expr]]
|
|
The indices of the element in the original buffer.
|
|
|
|
inner : bool, optional
|
|
If False, the offset is relative to the original buffer.
|
|
Default is True.
|
|
|
|
Returns
|
|
-------
|
|
offset: Expr
|
|
The element offset of the buffer at the given indices.
|
|
"""
|
|
if inner:
|
|
return _ffi_api.BufferOffsetOfp(self, indices)
|
|
return self.elem_offset + _ffi_api.BufferOffsetOfp(self, indices)
|
|
|
|
def byte_offset_of(self, indices, inner=True):
|
|
"""Get the byte offset of the buffer at the given indices.
|
|
Note that indices subject to buffer's layout mapping.
|
|
|
|
Parameters
|
|
----------
|
|
indices : Union[Expr, List[Expr]]
|
|
The indices of the element in the original buffer.
|
|
|
|
inner : bool, optional
|
|
If False, the offset is relative to the original buffer.
|
|
Default is True.
|
|
|
|
Returns
|
|
-------
|
|
offset: Expr
|
|
The byte offset of the buffer at the given indices.
|
|
"""
|
|
return self.elem_offset_of(indices, inner) * tvm.DataType(self.dtype).bits // 8
|
|
|
|
def is_scalar(self, alloc_or_decl=True):
|
|
"""Check if the buffer is a scalar.
|
|
|
|
Parameters
|
|
----------
|
|
alloc_or_decl : bool, optional
|
|
Whether to consider alloc_scalar and decl_scalar as scalar. True for alloc_scalar,
|
|
False for decl_scalar.
|
|
|
|
Returns
|
|
-------
|
|
bool: True if the buffer is a scalar, False otherwise.
|
|
"""
|
|
return _ffi_api.BufferIsScalar(self, alloc_or_decl)
|
|
|
|
def ptr_to(self, indices):
|
|
"""Get the pointer to the buffer at the given indices (logical indices).
|
|
|
|
Note that the bufferload inside requires LowerTIPp pass to apply the layout to get the physical indices.
|
|
""" # noqa: E501
|
|
assert len(indices) == len(self.shape), (
|
|
f"The number of indices {indices} does not match the shape of the buffer {self.shape}"
|
|
)
|
|
return tvm.tirx.address_of(self[tuple(indices)])
|
|
|
|
def view(self, *args, **kwargs) -> "Buffer":
|
|
"""Creates a new view of the buffer. (used by parser)
|
|
|
|
Supported signatures are ``view(*shape, layout=None)``, where shape can contain
|
|
``-1`` to indicate that the dimension size is auto-inferred, and
|
|
``view(dtype: Union[str, tvm.DataType])``.
|
|
|
|
Returns
|
|
-------
|
|
view : DeclBufferFrame
|
|
The corresponding view buffer.
|
|
"""
|
|
|
|
def _infer_shape(shape):
|
|
shape = list(shape)
|
|
if -1 in shape and shape.count(-1) == 1:
|
|
size = functools.reduce(lambda x, y: x * y, self.shape)
|
|
n_size = functools.reduce(lambda x, y: x * y, [s for s in shape if s != -1], 1)
|
|
shape[shape.index(-1)] = size // n_size
|
|
else:
|
|
# Only validate the shape product when both old and new shapes
|
|
# are fully concrete: a Expr `==` returns an `EQ` node, not
|
|
# a Python bool, and `assert <Expr>` raises (no __bool__).
|
|
if all(isinstance(s, int) for s in shape) and all(
|
|
isinstance(s, int) for s in self.shape
|
|
):
|
|
assert functools.reduce(lambda x, y: x * y, shape) == functools.reduce(
|
|
lambda x, y: x * y, self.shape
|
|
), (
|
|
"The shape of the buffer "
|
|
+ str(self.shape)
|
|
+ " and the new shape "
|
|
+ str(shape)
|
|
+ " are not compatible"
|
|
)
|
|
return shape
|
|
|
|
if len(args) == 1 and isinstance(args[0], str | tvm.DataType) and not kwargs:
|
|
cast_dtype = tvm.DataType(args[0])
|
|
cur_dtype = tvm.DataType(self.dtype)
|
|
if cast_dtype.bits > cur_dtype.bits:
|
|
# cast up
|
|
assert cast_dtype.bits % cur_dtype.bits == 0
|
|
ratio = cast_dtype.bits // cur_dtype.bits
|
|
layout = self.layout.pack(ratio)
|
|
shape = [s for s in self.shape[:-1]] + [self.shape[-1] // ratio]
|
|
new_elem_offset = self.elem_offset // ratio
|
|
else:
|
|
# cast down
|
|
assert cur_dtype.bits % cast_dtype.bits == 0
|
|
ratio = cur_dtype.bits // cast_dtype.bits
|
|
layout = self.layout.unpack(ratio)
|
|
shape = [s for s in self.shape[:-1]] + [self.shape[-1] * ratio]
|
|
new_elem_offset = self.elem_offset * ratio
|
|
return tvm.tirx.script.builder.decl_buffer(
|
|
shape,
|
|
cast_dtype,
|
|
self.data,
|
|
self.strides,
|
|
new_elem_offset,
|
|
None,
|
|
self.scope(),
|
|
self.data_alignment,
|
|
self.offset_factor,
|
|
"",
|
|
self.axis_separators,
|
|
layout,
|
|
)
|
|
else:
|
|
# --- Signature 1: view(*shape, **opts) ---
|
|
# Check if all positional args are integers/PrimExprs with dtype int32 or int64 (the shape) # noqa: E501
|
|
shape = args
|
|
assert all(
|
|
isinstance(arg, int)
|
|
or (tvm.ir.is_prim_expr(arg) and arg.ty.dtype in ["int32", "int64"])
|
|
for arg in shape
|
|
), "shape must be a list of integers or PrimExprs with dtype int32 or int64"
|
|
# Safely get optional keyword arguments
|
|
layout = kwargs.get("layout", None)
|
|
# Assert there are no other kwargs
|
|
assert set(kwargs.keys()).issubset({"layout"}), (
|
|
f"Unsupported kwargs for view: {set(kwargs.keys()) - {'layout'}}"
|
|
)
|
|
|
|
if layout is None:
|
|
shape = _infer_shape(shape)
|
|
|
|
return tvm.tirx.script.builder.decl_buffer(
|
|
shape,
|
|
self.dtype,
|
|
self.data,
|
|
self.strides,
|
|
self.elem_offset,
|
|
None,
|
|
self.scope(),
|
|
self.data_alignment,
|
|
self.offset_factor,
|
|
"",
|
|
self.axis_separators,
|
|
self.layout if layout is None else layout,
|
|
)
|
|
|
|
def local(self, *shape, layout=None) -> "Buffer":
|
|
"""Create a thread-local view of this buffer.
|
|
|
|
When called with no shape arguments, auto-infers a 1D shape from
|
|
the layout's non-thread component (i.e. ``layout.storage().shard``).
|
|
|
|
Parameters
|
|
----------
|
|
shape : tuple of Expr
|
|
The shape of the local view for indexing. If omitted, a 1D
|
|
shape is computed automatically.
|
|
|
|
layout : optional
|
|
Override layout. If None, uses the storage layout
|
|
(parent layout with thread axes removed).
|
|
|
|
Returns
|
|
-------
|
|
local : DeclBufferFrame
|
|
The corresponding local buffer.
|
|
"""
|
|
if not shape:
|
|
local_layout = self.layout.storage()
|
|
total = functools.reduce(
|
|
lambda x, y: x * y, [it.extent for it in local_layout.shard], 1
|
|
)
|
|
shape = (total,)
|
|
return tvm.tirx.script.builder.decl_buffer(
|
|
shape,
|
|
self.dtype,
|
|
self.data,
|
|
self.strides,
|
|
self.elem_offset,
|
|
None,
|
|
self.scope(),
|
|
self.data_alignment,
|
|
self.offset_factor,
|
|
"",
|
|
self.axis_separators,
|
|
self.layout.storage() if layout is None else layout,
|
|
)
|
|
|
|
def permute(self, *dims) -> "Buffer":
|
|
"""Permute the dimensions of the buffer.
|
|
|
|
Parameters
|
|
----------
|
|
dims : tuple of int
|
|
The permutation of dimensions.
|
|
|
|
Returns
|
|
-------
|
|
permuted : DeclBufferFrame
|
|
The buffer with permuted dimensions.
|
|
"""
|
|
new_shape = [self.shape[d] for d in dims]
|
|
# Permute *logical* dims, not the layout's fine-grained shard iters: a
|
|
# tcgen05/atom layout maps several shard iters to each logical axis, so
|
|
# group by the current shape first and permute whole groups. ``group``
|
|
# returns a regrouped layout (degenerate extent-1 iters folded away)
|
|
# plus seps over *that* layout — permute the regrouped one, not
|
|
# ``self.layout``. For a simple layout (one shard iter per axis) this
|
|
# reduces to ``permute_dims(dims)``.
|
|
grouped, seps = self.layout.group(list(self.shape))
|
|
new_layout = grouped.permute_by_groups(seps, list(dims))
|
|
return tvm.tirx.script.builder.decl_buffer(
|
|
new_shape,
|
|
self.dtype,
|
|
self.data,
|
|
self.strides,
|
|
self.elem_offset,
|
|
None,
|
|
self.scope(),
|
|
self.data_alignment,
|
|
self.offset_factor,
|
|
"",
|
|
self.axis_separators,
|
|
new_layout,
|
|
)
|
|
|
|
def __getitem__(self, indices):
|
|
from ..arith import Analyzer # pylint: disable=import-outside-toplevel
|
|
from .expr import BufferLoad, Ramp # pylint: disable=import-outside-toplevel
|
|
from .stmt import BufferRegion # pylint: disable=import-outside-toplevel
|
|
|
|
if not isinstance(indices, tuple | list):
|
|
indices = [indices]
|
|
has_slice = any(isinstance(i, slice) for i in indices)
|
|
has_step = any(
|
|
isinstance(i, slice) and (i.step is not None and i.step != 1) for i in indices
|
|
)
|
|
has_implicit_slice = len(indices) < len(self.shape)
|
|
analyzer = Analyzer()
|
|
if (has_slice and not has_step) or has_implicit_slice:
|
|
region = []
|
|
for i, index in enumerate(indices):
|
|
if isinstance(index, slice):
|
|
start = 0 if index.start is None else index.start
|
|
stop = self.shape[i] if index.stop is None else index.stop
|
|
region.append(Range.from_min_extent(start, analyzer.simplify(stop - start)))
|
|
else:
|
|
region.append(
|
|
Range.from_min_extent(
|
|
index,
|
|
tvm.tirx.expr.IntImm(index.ty, 1) if tvm.ir.is_prim_expr(index) else 1,
|
|
)
|
|
)
|
|
if has_implicit_slice:
|
|
for i in range(len(indices), len(self.shape)):
|
|
region.append(Range.from_min_extent(0, self.shape[i]))
|
|
return BufferRegion(self, region)
|
|
else:
|
|
expr_indices = []
|
|
for i, index in enumerate(indices):
|
|
if isinstance(index, slice):
|
|
start = 0 if index.start is None else index.start
|
|
stop = self.shape[i] if index.stop is None else index.stop
|
|
step = 1 if index.step is None else index.step
|
|
# We should ensure the dtype of start is the same with that of step.
|
|
if tvm.ir.is_prim_expr(start) and isinstance(step, int):
|
|
step = tvm.tirx.expr.IntImm(start.ty, step)
|
|
lanes = analyzer.simplify((stop - start + step - 1) // step)
|
|
if lanes == 1:
|
|
expr_indices.append(start)
|
|
else:
|
|
expr_indices.append(Ramp(start, step, int(lanes)))
|
|
else:
|
|
expr_indices.append(index)
|
|
return BufferLoad(self, expr_indices)
|
|
|
|
|
|
def decl_buffer(
|
|
shape,
|
|
dtype=None,
|
|
name="buffer",
|
|
data=None,
|
|
strides=None,
|
|
elem_offset=None,
|
|
scope="",
|
|
data_alignment=-1,
|
|
offset_factor=0,
|
|
buffer_type="",
|
|
axis_separators=None,
|
|
span=None,
|
|
layout="default",
|
|
):
|
|
# pylint: disable=import-outside-toplevel
|
|
from .expr import Var
|
|
from .layout import S, TileLayout
|
|
|
|
shape = (shape,) if tvm.ir.is_prim_expr(shape) or isinstance(shape, Integral) else shape
|
|
dtype = "float32" if dtype is None else dtype
|
|
strides = () if strides is None else strides
|
|
|
|
if axis_separators is None:
|
|
axis_separators = []
|
|
|
|
if layout == "default":
|
|
layout = TileLayout(S[tuple(shape)]) if shape else None
|
|
|
|
if offset_factor != 0 and elem_offset is None:
|
|
shape_ty = shape[0].ty if shape and tvm.ir.is_prim_expr(shape[0]) else "int32"
|
|
elem_offset = Var(f"{name}_elem_offset", shape_ty)
|
|
if data is None:
|
|
# Bool is represented as uint1 in the IR, but stored as int8
|
|
storage_type = dtype if isinstance(dtype, PrimType) else PrimType(dtype)
|
|
storage_type = PrimType("int8") if storage_type.dtype == "bool" else storage_type
|
|
data = Var(name, PointerType(storage_type, scope), span)
|
|
return _ffi_api.Buffer( # type: ignore
|
|
data,
|
|
dtype,
|
|
shape,
|
|
strides,
|
|
elem_offset,
|
|
name,
|
|
data_alignment,
|
|
offset_factor,
|
|
buffer_type,
|
|
axis_separators,
|
|
span,
|
|
layout,
|
|
)
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.DataProducer")
|
|
class DataProducer(Object):
|
|
pass
|