wehub-resource-sync
579 lines
19 KiB
Python
579 lines
19 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.
|
|
# pylint: disable=unrecognized-inline-option
|
|
"""Function data types."""
|
|
|
|
import collections
|
|
import inspect
|
|
from collections.abc import Callable, Mapping
|
|
from typing import Optional
|
|
|
|
import tvm_ffi
|
|
|
|
import tvm
|
|
import tvm.runtime
|
|
from tvm.ir import BaseFunc, Range
|
|
from tvm.runtime import Object, Scriptable
|
|
|
|
from ..runtime._tensor import Tensor
|
|
from . import _ffi_api
|
|
from .buffer import Buffer
|
|
from .expr import Expr, Var
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.PrimFunc")
|
|
class PrimFunc(BaseFunc, Scriptable):
|
|
"""A function declaration expression.
|
|
|
|
Parameters
|
|
----------
|
|
params: List[Union[tvm.tirx.Var, tvm.tirx.Buffer]]
|
|
List of input parameters to the function.
|
|
|
|
body: tvm.tirx.Stmt
|
|
The body of the function.
|
|
|
|
ret_type: tvm.ir.Type
|
|
The return type annotation of the function.
|
|
|
|
buffer_map : Map[tvm.tirx.Var, tvm.tirx.Buffer]
|
|
The buffer binding map.
|
|
|
|
attrs: Optional[tvm.Attrs]
|
|
Attributes of the function, can be None
|
|
|
|
span : Optional[Span]
|
|
The location of this itervar in the source code.
|
|
"""
|
|
|
|
def __init__(self, params, body, ret_type=None, buffer_map=None, attrs=None, span=None):
|
|
# Legacy compatibility: expand body-carrying leaf stmt wrappers
|
|
# (e.g. DeclBuffer/AllocBuffer forms) into SeqStmt form.
|
|
from .stmt import _normalize_legacy_stmt
|
|
|
|
body = _normalize_legacy_stmt(body)
|
|
if ret_type is None:
|
|
ret_type = tvm.ir.Type.missing()
|
|
param_list = []
|
|
buffer_map = {} if buffer_map is None else buffer_map
|
|
for x in params:
|
|
x = tvm.runtime.convert(x) if not isinstance(x, Object) else x
|
|
if isinstance(x, Buffer):
|
|
var = Var(x.name, dtype="handle")
|
|
param_list.append(var)
|
|
buffer_map[var] = x
|
|
elif isinstance(x, Var):
|
|
param_list.append(x)
|
|
else:
|
|
raise TypeError("params can only contain Var or Buffer")
|
|
|
|
if attrs is None:
|
|
attrs = tvm.ir.make_node("ir.DictAttrs")
|
|
|
|
self.__init_handle_by_constructor__(
|
|
_ffi_api.PrimFunc,
|
|
param_list,
|
|
body,
|
|
ret_type,
|
|
buffer_map,
|
|
attrs,
|
|
span,
|
|
) # type: ignore
|
|
|
|
def with_body(self, new_body, span=None):
|
|
"""Create a new PrimFunc with the same set signatures but a new body.
|
|
|
|
Parameters
|
|
----------
|
|
new_body : Stmt
|
|
The new body.
|
|
|
|
span : Optional[Span]
|
|
The location of this itervar in the source code.
|
|
|
|
Returns
|
|
-------
|
|
new_func : PrimFunc
|
|
The created new function.
|
|
"""
|
|
return PrimFunc(
|
|
self.params,
|
|
new_body,
|
|
self.ret_type,
|
|
self.buffer_map,
|
|
self.attrs,
|
|
span,
|
|
)
|
|
|
|
def specialize(self, param_map: Mapping[Var, Expr | Buffer]):
|
|
"""Specialize parameters of PrimFunc
|
|
|
|
Parameters
|
|
----------
|
|
|
|
param_map : Mapping[Var, Union[Expr, Buffer]]
|
|
The mapping from function params to the instance
|
|
|
|
Examples
|
|
--------
|
|
We can define a Meta TIR function with symbolic shape:
|
|
|
|
.. code-block:: python
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def mem_copy(a: T.handle, b: T.handle, m: T.int32, n: T.int32) -> None:
|
|
A = T.match_buffer(a, (m, n), "float32")
|
|
B = T.match_buffer(b, (m, n), "float32")
|
|
|
|
for i, j in T.grid(m, n):
|
|
with T.sblock():
|
|
vi, vj = T.axis.remap("SS", [i, j])
|
|
B[vi, vj] = A[vi, vj]
|
|
|
|
Then we can make it specialized with given shapes or buffers.
|
|
|
|
.. code-block:: python
|
|
|
|
a, _, m, n = mem_copy.params
|
|
func = mem_copy.specialize({a: tirx.decl_buffer((16, 16))})
|
|
# or
|
|
func = mem_copy.specialize({n: 16, m: 16})
|
|
|
|
The specialized function:
|
|
|
|
.. code-block:: python
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def mem_copy_16_16(a: T.handle, b: T.handle) -> None:
|
|
A = T.match_buffer(a, (16, 16), "float32")
|
|
B = T.match_buffer(b, (16, 16), "float32")
|
|
|
|
for i, j in T.grid(16, 16):
|
|
with T.sblock():
|
|
vi, vj = T.axis.remap("SS", [i, j])
|
|
B[vi, vj] = A[vi, vj]
|
|
|
|
Returns
|
|
-------
|
|
func : PrimFunc
|
|
The new function with parameter specialized
|
|
"""
|
|
return _ffi_api.Specialize(self, param_map) # type: ignore
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.TensorIntrin")
|
|
class TensorIntrin(Object):
|
|
"""A tensor intrinsic.
|
|
|
|
Parameters
|
|
----------
|
|
desc : PrimFunc
|
|
The function to describe the computation.
|
|
|
|
impl : PrimFunc
|
|
The function of the implementation for the execution.
|
|
"""
|
|
|
|
def __init__(self, desc, impl):
|
|
self.__init_handle_by_constructor__(_ffi_api.TensorIntrin, desc, impl)
|
|
|
|
@staticmethod
|
|
def register(name: str, desc: PrimFunc, impl: PrimFunc, override: bool = False):
|
|
"""Register a tensor intrinsic with its name.
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
The name of the TensorIntrin to register.
|
|
desc : PrimFunc
|
|
The function to describe the computation.
|
|
impl : PrimFunc
|
|
The function of the implementation for the execution.
|
|
override: bool
|
|
Whether override existing intrinsic.
|
|
"""
|
|
return _ffi_api.TensorIntrinRegister(name, TensorIntrin(desc, impl), override) # type: ignore
|
|
|
|
@staticmethod
|
|
def get(name: str, allow_missing: bool = False) -> Optional["TensorIntrin"]:
|
|
"""Look up a tensor intrinsic by its name.
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
The name of the TensorIntrin to look up.
|
|
|
|
allow_missing : bool
|
|
Whether to allow missing tensor intrin. If False, raise an error if the tensor intrin
|
|
doesn't exist.
|
|
|
|
Returns
|
|
-------
|
|
result : Optional[TensorIntrin]
|
|
The TensorIntrin with the specified name, or None if not found.
|
|
"""
|
|
return _ffi_api.TensorIntrinGet(name, allow_missing) # pylint: type: ignore
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.IndexMap")
|
|
class IndexMap(Object):
|
|
"""A mapping from multi-dimensional indices to another set of multi-dimensional indices
|
|
|
|
Parameters
|
|
----------
|
|
initial_indices : List[Var]
|
|
Variables representing the indices prior to remapping.
|
|
final_indices : List[Expr]
|
|
Expressions defining the indices after remapping.
|
|
inverse_index_map : Union[Callable, Optional[IndexMap]]
|
|
The optional pre-defined inverse index map.
|
|
When this is defined, IndexMap::Inverse will return the pre-defined inverse index map.
|
|
Otherwise, the inverse index map will be computed on the fly.
|
|
It is the user's responsibility to ensure the correctness of the pre-defined inverse
|
|
index map.
|
|
"""
|
|
|
|
initial_indices: list[Var]
|
|
final_indices: list[Expr]
|
|
|
|
# Sentinel value used to indicate which groups of pre-flattening axes
|
|
# should be used to post-flattening axes axes. See
|
|
# Stage.transform_layout for more details.
|
|
AXIS_SEPARATOR = "axis_separator"
|
|
|
|
def __init__(self, initial_indices, final_indices, inverse_index_map):
|
|
if isinstance(inverse_index_map, Callable):
|
|
inverse_index_map = IndexMap.from_func(inverse_index_map)
|
|
self.__init_handle_by_constructor__(
|
|
_ffi_api.IndexMap, initial_indices, final_indices, inverse_index_map
|
|
)
|
|
|
|
@staticmethod
|
|
def from_func(
|
|
mapping_function: Callable,
|
|
ndim: int | None = None,
|
|
inverse_index_map: Callable | Optional["IndexMap"] = None,
|
|
*,
|
|
index_dtype: str = "int64",
|
|
):
|
|
"""Create an index map from a function
|
|
|
|
Parameters
|
|
----------
|
|
mapping_function : Callable
|
|
|
|
The function to map from source indices to target indices.
|
|
The function should accept `tirx.Var` parameters and return
|
|
a either a `tirx.Expr`, or a list of `tirx.Expr`.
|
|
Returning a `tirx.Expr` is equivalent to returning a
|
|
list of length 1 containing that `tirx.Expr`.
|
|
|
|
ndim: Optional[int]
|
|
|
|
The dimensionality of the buffer to which this
|
|
transformation should be applied. If mapping_function uses
|
|
variadic argument `*args`, `ndim` must be specified. If
|
|
mapping_function does not use variadic arguments, ndim is
|
|
optional.
|
|
|
|
inverse_index_map : Union[Callable, Optional[IndexMap]]
|
|
The optional pre-defined inverse index map.
|
|
When this is defined, IndexMap::Inverse will return the pre-defined inverse index map.
|
|
Otherwise, the inverse index map will be computed on the fly.
|
|
It is the user's responsibility to ensure the correctness of the pre-defined inverse
|
|
index map.
|
|
|
|
Returns
|
|
-------
|
|
index_map: IndexMap
|
|
|
|
Returns an IndexMap representing the `mapping_function`.
|
|
|
|
"""
|
|
index_map, axis_separators = IndexMap.from_func_with_separators(
|
|
mapping_function,
|
|
ndim,
|
|
inverse_index_map,
|
|
index_dtype=index_dtype,
|
|
)
|
|
assert not axis_separators, (
|
|
"The mapping_function provided to IndexMap.from_func "
|
|
"may not return IndexMap.AXIS_SEPARATOR. "
|
|
"If required, please use IndexMap.from_func_with_separators instead."
|
|
)
|
|
return index_map
|
|
|
|
@staticmethod
|
|
def from_func_with_separators(
|
|
mapping_function: Callable,
|
|
ndim: int | None = None,
|
|
inverse_index_map: Callable | Optional["IndexMap"] = None,
|
|
*,
|
|
index_dtype: str = "int64",
|
|
):
|
|
"""Create an index map from a function
|
|
|
|
Parameters
|
|
----------
|
|
mapping_function : Callable
|
|
|
|
The function to map from source indices to target indices.
|
|
The function should accept tirx.Var parameters and return
|
|
either a `tirx.Expr` or a list. Each element of the
|
|
returned list should be either a `tirx.Expr` or the
|
|
object `IndexMap.AXIS_SEPARATOR`. Returning a
|
|
`tirx.Expr` is equivalent to returning a list of length
|
|
1 containing that `tirx.Expr`.
|
|
|
|
ndim: Optional[int]
|
|
|
|
The dimensionality of the buffer to which this
|
|
transformation should be applied. If mapping_function uses
|
|
variadic argument `*args`, ndim must be specified. If
|
|
mapping_function does not use variadic arguments, ndim is
|
|
optional.
|
|
|
|
inverse_index_map : Union[Callable, Optional[IndexMap]]
|
|
The optional pre-defined inverse index map.
|
|
When this is defined, IndexMap::Inverse will return the pre-defined inverse index map.
|
|
Otherwise, the inverse index map will be computed on the fly.
|
|
It is the user's responsibility to ensure the correctness of the pre-defined inverse
|
|
index map.
|
|
|
|
index_dtype : str
|
|
The default index dtype to use for input iters in the mapping function.
|
|
|
|
Returns
|
|
-------
|
|
ret: Tuple[IndexMap, List[int]]
|
|
|
|
Returns a tuple whose first element is an IndexMap
|
|
representing the `mapping_function`, and whose second index
|
|
is a list of indices at which `IndexMap.AXIS_SEPARATOR`
|
|
occurred.
|
|
|
|
"""
|
|
params = inspect.signature(mapping_function).parameters
|
|
|
|
args = []
|
|
var_arg_name = None
|
|
kwargs = collections.OrderedDict()
|
|
|
|
for name, param in params.items():
|
|
if param.kind in [
|
|
inspect.Parameter.POSITIONAL_ONLY,
|
|
inspect.Parameter.POSITIONAL_OR_KEYWORD,
|
|
]:
|
|
args.append(tvm.tirx.Var(name, index_dtype))
|
|
|
|
elif param.kind == inspect.Parameter.VAR_POSITIONAL:
|
|
var_arg_name = name
|
|
|
|
elif param.kind == inspect.Parameter.KEYWORD_ONLY:
|
|
kwargs[name] = tvm.tirx.Var(name, index_dtype)
|
|
|
|
else:
|
|
raise ValueError("transform_layout mapping may not have *args")
|
|
|
|
# Now that all the named arguments have been collected,
|
|
# everything that remains should go to the *args, if
|
|
# specified.
|
|
if var_arg_name is not None:
|
|
assert ndim is not None, "ndim must be specified when *args is used"
|
|
num_var_args = ndim - len(args) - len(kwargs)
|
|
for i in range(num_var_args):
|
|
args.append(tvm.tirx.Var(f"{var_arg_name}_{i}", index_dtype))
|
|
|
|
mapping = mapping_function(*args, **kwargs)
|
|
|
|
initial_indices = args + list(kwargs.values())
|
|
|
|
final_indices = []
|
|
axis_separators = []
|
|
|
|
try:
|
|
iter(mapping)
|
|
is_iterable = True
|
|
except TypeError:
|
|
is_iterable = False
|
|
|
|
if is_iterable:
|
|
for val in mapping:
|
|
if tvm.ir.is_prim_expr(val):
|
|
final_indices.append(val)
|
|
elif val is IndexMap.AXIS_SEPARATOR:
|
|
axis_separators.append(len(final_indices))
|
|
else:
|
|
raise TypeError(
|
|
"Expected mapping function to return list of "
|
|
"either tvm.ir.Expr or IndexMap.AXIS_SEPARATOR. "
|
|
f"Instead received {val} of type {type(val)}."
|
|
)
|
|
else:
|
|
final_indices.append(mapping)
|
|
|
|
return IndexMap(initial_indices, final_indices, inverse_index_map), axis_separators
|
|
|
|
def is_equivalent_to(self, other_map: "IndexMap", analyzer=None) -> bool:
|
|
"""Return if the index maps are equivalent.
|
|
|
|
Parameters
|
|
----------
|
|
other_map: IndexMap
|
|
|
|
The IndexMap to which the comparison should be made.
|
|
|
|
analyzer : Optional[tvm.arith.Analyzer]
|
|
|
|
The analyzer to use while comparing the mapped indices. When
|
|
provided, its accumulated bindings and constraints are reused so
|
|
that maps that are only equivalent under those bindings can be
|
|
proven equal.
|
|
|
|
Returns
|
|
-------
|
|
is_equivalent: bool
|
|
|
|
True if the two mappings represent the same
|
|
transformation, otherwise False
|
|
"""
|
|
if len(self.initial_indices) != len(other_map.initial_indices):
|
|
return False
|
|
if len(self.final_indices) != len(other_map.final_indices):
|
|
return False
|
|
|
|
if analyzer is None:
|
|
analyzer = tvm.arith.Analyzer()
|
|
|
|
mapped_other_final_indices = other_map.map_indices(self.initial_indices, analyzer=analyzer)
|
|
for self_index, other_index in zip(self.final_indices, mapped_other_final_indices):
|
|
if not analyzer.can_prove_equal(self_index, other_index):
|
|
return False
|
|
|
|
return True
|
|
|
|
def map_indices(self, indices: list[Expr], analyzer=None) -> list[Expr]:
|
|
"""Apply the index map to a set of indices
|
|
|
|
Parameters
|
|
----------
|
|
indices : List[Expr]
|
|
The indices to be mapped
|
|
analyzer : Optional[tvm.arith.Analyzer]
|
|
The analyzer to use while simplifying mapped indices.
|
|
|
|
Returns
|
|
-------
|
|
result : List[Expr]
|
|
The mapped indices
|
|
"""
|
|
return _ffi_api.IndexMapMapIndices(self, indices, analyzer)
|
|
|
|
def map_shape(self, shape: list[Expr], analyzer=None) -> list[Expr]:
|
|
"""Apply the index map to a buffer shape
|
|
|
|
Parameters
|
|
----------
|
|
shape : List[Expr]
|
|
The buffer shape to be mapped
|
|
analyzer : Optional[tvm.arith.Analyzer]
|
|
The analyzer to use while simplifying mapped shape expressions.
|
|
|
|
Returns
|
|
-------
|
|
result : List[Expr]
|
|
The mapped shape
|
|
"""
|
|
return _ffi_api.IndexMapMapShape(self, shape, analyzer)
|
|
|
|
def map_tensor(self, arr_src: Tensor) -> Tensor:
|
|
"""Apply thie index map to transform the layout of the input Tensor
|
|
|
|
Parameters
|
|
----------
|
|
arr_src : runtime.Tensor
|
|
The Tensor to be transformed
|
|
|
|
Returns
|
|
-------
|
|
arr_dst : runtime.Tensor
|
|
The transformed Tensor
|
|
"""
|
|
return _ffi_api.IndexMapMapTensor(self, arr_src)
|
|
|
|
def inverse(self, shape: list[Range | Expr], analyzer=None) -> "IndexMap":
|
|
"""Return the inverse of the map
|
|
|
|
Throws an error if the function is not bijective.
|
|
|
|
Parameters
|
|
----------
|
|
shape: List[Union[Range,Expr]]
|
|
|
|
The region over which the inverse should be determined.
|
|
Used for validating that the mapping is bijective over
|
|
this range.
|
|
analyzer : Optional[tvm.arith.Analyzer]
|
|
The analyzer to use while deriving and validating the inverse.
|
|
|
|
Returns
|
|
-------
|
|
inverse : IndexMap
|
|
|
|
The inverse
|
|
"""
|
|
|
|
shape = [dim if isinstance(dim, Range) else Range(0, dim) for dim in shape]
|
|
return _ffi_api.IndexMapInverse(self, shape, analyzer)
|
|
|
|
def non_surjective_inverse(
|
|
self, shape: list[Range | Expr], analyzer=None
|
|
) -> tuple["IndexMap", Expr]:
|
|
"""Return the inverse of the map
|
|
|
|
Can be applied to transformations that introduce padding.
|
|
|
|
Parameters
|
|
----------
|
|
shape: List[Union[Range,Expr]]
|
|
|
|
The region over which the inverse should be determined.
|
|
Used for determining the predicate.
|
|
analyzer : Optional[tvm.arith.Analyzer]
|
|
The analyzer to use while deriving the inverse and padding predicate.
|
|
|
|
Returns
|
|
-------
|
|
result : Tuple[IndexMap, Expr]
|
|
|
|
The inverse, and a predicate for which the inverse maps to
|
|
a valid index in the input range.
|
|
|
|
Examples
|
|
--------
|
|
|
|
.. code-block:: python
|
|
|
|
index_map = IndexMap.from_func(lambda i: [i//4, i%4])
|
|
inverse_map, predicate = index_map.non_surjective_inverse([14])
|
|
assert inverse_map.is_equivalent_to(IndexMap.from_func(lambda j,k: [4*j + k])
|
|
print(predicate) # Prints "(axis0==3) && (axis2 >= 2)"
|
|
"""
|
|
|
|
shape = [dim if isinstance(dim, Range) else Range(0, dim) for dim in shape]
|
|
return _ffi_api.IndexMapNonSurjectiveInverse(self, shape, analyzer)
|