chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,578 @@
|
||||
# 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)
|
||||
Reference in New Issue
Block a user