wehub-resource-sync
546 lines
14 KiB
Python
546 lines
14 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=invalid-name
|
|
# ruff: noqa: RUF005
|
|
"""Common topi utilities"""
|
|
|
|
from numbers import Integral
|
|
|
|
import numpy as np
|
|
|
|
import tvm
|
|
from tvm import te
|
|
from tvm.s_tir import sbijective_layout, slayout
|
|
from tvm.tirx import Var
|
|
|
|
from . import cpp, tag
|
|
|
|
|
|
class InvalidShapeError(ValueError):
|
|
"""Invalid shape for a topi function. i.e. call winograd template for non-3x3 kernel)"""
|
|
|
|
|
|
def ncw_pack_layout(layout_info):
|
|
"""Check whether the layout type is NCWinic"""
|
|
return layout_info[:3] == "NCW" and "c" in layout_info and "n" in layout_info
|
|
|
|
|
|
def ncw_xc_layout(layout_info):
|
|
"""Check whether the layout type is NCWxc"""
|
|
return layout_info[:3] == "NCW" and "c" in layout_info and layout_info[3:-1].isnumeric()
|
|
|
|
|
|
def nchw_pack_layout(layout_info):
|
|
"""Check whether the layout type is NCHWinic"""
|
|
return layout_info[:4] == "NCHW" and "c" in layout_info and "n" in layout_info
|
|
|
|
|
|
def nchw_xc_layout(layout_info):
|
|
"""Check whether the layout type is NCHWxc"""
|
|
return layout_info[:4] == "NCHW" and "c" in layout_info and layout_info[4:-1].isnumeric()
|
|
|
|
|
|
def traverse_inline(s, final_op, callback):
|
|
"""Traverse computation graph and do auto inline
|
|
|
|
Parameters
|
|
----------
|
|
s: schedule
|
|
The schedule
|
|
final_op: Operation
|
|
The final output operator.
|
|
callback: callable
|
|
The callback function on each op
|
|
"""
|
|
visited = set()
|
|
|
|
def _traverse(op):
|
|
if op in visited:
|
|
return
|
|
visited.add(op)
|
|
if tag.is_injective(op.tag):
|
|
if op not in s.outputs:
|
|
s[op].compute_inline()
|
|
for tensor in op.input_tensors:
|
|
if isinstance(tensor.op, tvm.te.ComputeOp):
|
|
_traverse(tensor.op)
|
|
callback(op)
|
|
|
|
_traverse(final_op)
|
|
|
|
|
|
def prod(x):
|
|
"""Get the product of every items in the tuple.
|
|
|
|
Parameters
|
|
----------
|
|
x: tuple
|
|
Input tuple
|
|
|
|
Returns
|
|
-------
|
|
value : Expr
|
|
The result value
|
|
"""
|
|
if not x:
|
|
return tvm.tirx.const(1, "int32")
|
|
res = x[0]
|
|
for i in range(1, len(x)):
|
|
res = res * x[i]
|
|
return res
|
|
|
|
|
|
def get_const_int(expr):
|
|
"""Verifies expr is integer and get the constant value.
|
|
|
|
Parameters
|
|
----------
|
|
expr : tvm.Expr or int
|
|
The input expression.
|
|
|
|
Returns
|
|
-------
|
|
out_value : int
|
|
The output.
|
|
"""
|
|
if isinstance(expr, Integral):
|
|
return expr
|
|
if not isinstance(expr, tvm.tirx.IntImm):
|
|
ana = tvm.arith.Analyzer()
|
|
expr = ana.simplify(expr)
|
|
if not isinstance(expr, tvm.tirx.IntImm):
|
|
raise ValueError("Expect value to be constant int")
|
|
return int(expr.value)
|
|
|
|
|
|
def get_const_float(expr):
|
|
"""Verifies expr is a floating point and get the constant value.
|
|
|
|
Parameters
|
|
----------
|
|
expr : tvm.Expr or float
|
|
The input expression.
|
|
|
|
Returns
|
|
-------
|
|
out_value : float
|
|
The output.
|
|
"""
|
|
if isinstance(expr, float):
|
|
return float(expr)
|
|
if not isinstance(expr, tvm.tirx.FloatImm):
|
|
ana = tvm.arith.Analyzer()
|
|
expr = ana.simplify(expr)
|
|
if not isinstance(expr, tvm.tirx.FloatImm):
|
|
raise ValueError("Expect value to be constant float")
|
|
return float(expr.value)
|
|
|
|
|
|
def equal_const_int(expr, value):
|
|
"""Returns if expr equals value.
|
|
|
|
Parameters
|
|
----------
|
|
expr : tvm.Expr
|
|
The input expression.
|
|
|
|
Returns
|
|
-------
|
|
equal : bool
|
|
Whether they equals.
|
|
"""
|
|
if isinstance(expr, Integral):
|
|
return expr == value
|
|
if not isinstance(expr, tvm.tirx.IntImm):
|
|
ana = tvm.arith.Analyzer()
|
|
expr = ana.simplify(expr)
|
|
if not isinstance(expr, tvm.tirx.IntImm):
|
|
return False
|
|
return expr.value == value
|
|
|
|
|
|
def get_const_tuple(in_tuple):
|
|
"""Verifies input tuple is IntImm or Var, returns tuple of int or Var.
|
|
|
|
Parameters
|
|
----------
|
|
in_tuple : tuple of Expr
|
|
The input.
|
|
|
|
Returns
|
|
-------
|
|
out_tuple : tuple of int
|
|
The output.
|
|
"""
|
|
if isinstance(in_tuple, te.tensor.Tensor):
|
|
raise TypeError(
|
|
"get_const_tuple expects a tuple-like shape (e.g., tensor.shape), "
|
|
"but got a te.Tensor. Did you mean get_const_tuple(tensor.shape)?"
|
|
)
|
|
ret = []
|
|
ana = None
|
|
for elem in in_tuple:
|
|
if isinstance(elem, tvm.tirx.Var):
|
|
ret.append(elem)
|
|
elif not isinstance(elem, tvm.tirx.IntImm | int):
|
|
ana = tvm.arith.Analyzer() if ana is None else ana
|
|
elem = ana.simplify(elem)
|
|
if not isinstance(elem, tvm.tirx.IntImm):
|
|
ret.append(elem)
|
|
else:
|
|
ret.append(get_const_int(elem))
|
|
else:
|
|
ret.append(get_const_int(elem))
|
|
return tuple(ret)
|
|
|
|
|
|
def const_vector(vector, name="const_vector"):
|
|
"""convert a const numpy 1-dimensional vector to tvm tensor
|
|
|
|
Parameters
|
|
----------
|
|
vector: numpy.ndarray
|
|
Const input array
|
|
name: str, optional
|
|
The name of output op
|
|
|
|
Returns
|
|
-------
|
|
tensor: Tensor
|
|
The created tensor
|
|
"""
|
|
if not isinstance(vector, np.ndarray):
|
|
vector = np.array(vector)
|
|
row = vector.shape[0]
|
|
dtype = str(vector.dtype)
|
|
idxm = tvm.tirx.indexmod
|
|
|
|
def select_array(i):
|
|
now = tvm.tirx.const(0.0, dtype)
|
|
for ii in range(row):
|
|
now = tvm.tirx.Select(
|
|
tvm.tirx.all(idxm(i, row) == ii), tvm.tirx.const(vector[ii], dtype), now
|
|
)
|
|
return now
|
|
|
|
return te.compute(vector.shape, select_array, name=name)
|
|
|
|
|
|
def get_float_tuple(in_tuple):
|
|
"""Verifies input tuple is FloatImm, returns tuple of float.
|
|
|
|
Parameters
|
|
----------
|
|
in_tuple : tuple of Expr
|
|
The input.
|
|
|
|
Returns
|
|
-------
|
|
out_tuple : tuple of float
|
|
The output.
|
|
"""
|
|
return tuple(get_const_float(elem) for elem in in_tuple)
|
|
|
|
|
|
def simplify(expr):
|
|
"""Simplify the expression if it is Expr, directly return if it is int.
|
|
|
|
Parameters
|
|
----------
|
|
expr : Expr or int
|
|
The input.
|
|
|
|
Returns
|
|
-------
|
|
out : Expr or int
|
|
The simplified output
|
|
"""
|
|
if isinstance(expr, te.Tensor):
|
|
return te.compute(
|
|
expr.shape,
|
|
lambda *indices: tvm.arith.Analyzer().simplify(expr[indices]),
|
|
name="simplify_output",
|
|
tag="simplify",
|
|
)
|
|
elif tvm.ir.is_prim_expr(expr):
|
|
return tvm.arith.Analyzer().simplify(expr)
|
|
else:
|
|
return expr
|
|
|
|
|
|
def ravel_index(indices, shape):
|
|
"""Flatten the index tuple to 1D
|
|
|
|
Parameters
|
|
----------
|
|
indices : tuple of int or tvm.tirx.IntImm
|
|
The input coordinates
|
|
|
|
shape : tuple of int
|
|
Shape of the tensor.
|
|
|
|
Returns
|
|
-------
|
|
idx : int or Expr
|
|
The index after flattening
|
|
"""
|
|
idx = None
|
|
for i, (shape_val, ind) in enumerate(zip(shape, indices)):
|
|
if i != 0:
|
|
idx = idx * shape_val + ind
|
|
else:
|
|
idx = ind
|
|
return idx
|
|
|
|
|
|
def unravel_index(idx, shape):
|
|
"""Convert the flattened ind to the coordinate array
|
|
|
|
Parameters
|
|
----------
|
|
idx : int or tvm.tirx.IntImm
|
|
The 1D index
|
|
|
|
shape : tuple of int
|
|
Shape of the tensor
|
|
|
|
Returns
|
|
-------
|
|
indices : tuple of int or tvm.tirx.IntImm
|
|
Corresponding coordinate of the 1D index
|
|
"""
|
|
idxd = tvm.tirx.indexdiv
|
|
idxm = tvm.tirx.indexmod
|
|
indices = []
|
|
for i, dim in enumerate(reversed(shape)):
|
|
if dim == 0:
|
|
indices.append(0)
|
|
elif i == len(shape) - 1:
|
|
# Assuming the index is in-bounds, the last coordinate is
|
|
# already less than dim, and doesn't need the be remainder
|
|
# mod dim.
|
|
indices.append(idx)
|
|
else:
|
|
indices.append(idxm(idx, dim))
|
|
idx = idxd(idx, dim)
|
|
indices = indices[::-1]
|
|
return indices
|
|
|
|
|
|
def const_matrix(matrix, name="const_matrix", attrs=None):
|
|
"""convert a const numpy 2-dimensional matrix to tvm tensor
|
|
|
|
Parameters
|
|
----------
|
|
matrix: numpy.ndarray
|
|
Const input array
|
|
name: str, optional
|
|
The name of output op
|
|
|
|
Returns
|
|
-------
|
|
tensor: Tensor
|
|
The created tensor
|
|
"""
|
|
row, col = matrix.shape
|
|
dtype = str(matrix.dtype)
|
|
idxm = tvm.tirx.indexmod
|
|
|
|
def select_array(i, j):
|
|
now = tvm.tirx.const(0.0, dtype)
|
|
for ii in range(row):
|
|
for jj in range(col):
|
|
now = tvm.tirx.Select(
|
|
tvm.tirx.all(idxm(i, row) == ii, idxm(j, col) == jj),
|
|
tvm.tirx.const(matrix[ii][jj], dtype),
|
|
now,
|
|
)
|
|
return now
|
|
|
|
if attrs is None:
|
|
attrs = {"const_matrix": True, "schedule_rule": "None"}
|
|
|
|
return te.compute(matrix.shape, select_array, name=name, attrs=attrs)
|
|
|
|
|
|
def get_max_power2_factor(n, max_value=None):
|
|
"""Get max factor of n in power of 2. If max_value is specificed, max factor
|
|
value will be no more max_value,
|
|
|
|
Parameter
|
|
---------
|
|
n : int
|
|
The input value
|
|
|
|
max_value : int, optional
|
|
The max value for the factor
|
|
|
|
Returns
|
|
-------
|
|
factor : int
|
|
The max factor in power of 2.
|
|
"""
|
|
x = 1
|
|
while n % 2 == 0:
|
|
if max_value is not None and max_value < x * 2:
|
|
break
|
|
x *= 2
|
|
n /= 2
|
|
return x
|
|
|
|
|
|
def get_shape(src_shape, src_layout, dst_layout):
|
|
"""Given a source shape, a source layout and a destination layout, infer
|
|
the destination shape.
|
|
|
|
Parameter
|
|
---------
|
|
src_shape : tuple of int or IntImm
|
|
Source shape
|
|
|
|
src_layout : str or Layout
|
|
Source layout
|
|
|
|
dst_layout : str or Layout
|
|
Destination layout
|
|
|
|
Returns
|
|
-------
|
|
dst_shape : tuple of int
|
|
Destination shape
|
|
"""
|
|
if src_layout == dst_layout:
|
|
return get_const_tuple(src_shape)
|
|
|
|
if isinstance(src_layout, str):
|
|
src_layout = slayout(src_layout)
|
|
if isinstance(dst_layout, str):
|
|
dst_layout = slayout(dst_layout)
|
|
|
|
assert len(src_layout) == len(dst_layout), f"Incompatible layout {src_layout} vs {dst_layout}"
|
|
|
|
layout_mapping = sbijective_layout(src_layout, dst_layout)
|
|
dst_indices = layout_mapping.forward_index(tvm.runtime.convert(list(range(len(src_layout)))))
|
|
|
|
return get_const_tuple(tuple([src_shape[i.value] for i in dst_indices]))
|
|
|
|
|
|
def within_index(b, e, s, i):
|
|
"""Return a boolean value that indicates if i is within the given index.
|
|
|
|
Parameters
|
|
----------
|
|
b : Expr
|
|
beginning of the index
|
|
|
|
e : Expr
|
|
end of the index
|
|
|
|
s : Expr
|
|
strides of index
|
|
|
|
i : Expr
|
|
array position
|
|
|
|
Returns
|
|
-------
|
|
selected: Expr
|
|
bool expression that is True is the array position would be selected
|
|
by the index and False otherwise
|
|
"""
|
|
bc = tvm.tirx.Select(s < 0, i <= e, i < b)
|
|
ec = tvm.tirx.Select(s < 0, i > b, i >= e)
|
|
ss = te.if_then_else(s < 0, ((i - e) + (e % te.abs(s)) + 1) % te.abs(s), (i - b) % s)
|
|
return tvm.tirx.Select(tvm.tirx.Or(bc, ec), tvm.tirx.const(False), ss.equal(0))
|
|
|
|
|
|
def make_idx(b, e, s, z, i):
|
|
"""Return the array position in the selection that corresponds to an
|
|
array position in the full array.
|
|
|
|
The returned value is only meaningful if within_index() returns True
|
|
for the same set of parameters.
|
|
|
|
Parameters
|
|
----------
|
|
b : Expr
|
|
beginning of the index
|
|
|
|
e : Expr
|
|
end of the index
|
|
|
|
s : Expr
|
|
strides of index
|
|
|
|
z : Expr
|
|
size of the indexed dimension
|
|
|
|
i : Expr
|
|
array position
|
|
|
|
Returns
|
|
-------
|
|
position: Expr
|
|
int expression that corresponds to an array position in the selection.
|
|
"""
|
|
bc = tvm.tirx.Select(s < 0, i <= e, i < b)
|
|
ec = tvm.tirx.Select(s < 0, i > b, i >= e)
|
|
|
|
# Clamp to array size
|
|
b = tvm.tirx.Select(z < b, z - 1, b)
|
|
|
|
ss = tvm.tirx.if_then_else(s < 0, (b - i) // te.abs(s), (i - b) // s)
|
|
return tvm.tirx.if_then_else(tvm.tirx.Or(bc, ec), 88, ss)
|
|
|
|
|
|
def is_empty_shape(shape):
|
|
"""Check whether an input shape has dimesion with size 0.
|
|
|
|
Parameter
|
|
---------
|
|
shape : list of Expr
|
|
Input shape
|
|
|
|
Returns
|
|
-------
|
|
is_empty: bool
|
|
Whether input shape is empty or has dimesion with size 0.
|
|
"""
|
|
return cpp.utils.is_empty_shape(shape)
|
|
|
|
|
|
def ceil_div(a, b):
|
|
"""Return ceil division of a by b"""
|
|
return tvm.tirx.indexdiv(a + (b - 1), b)
|
|
|
|
|
|
def swap(arr, axis):
|
|
"""swap arr[axis] and arr[-1]"""
|
|
return arr[:axis] + [arr[-1]] + arr[axis + 1 : -1] + [arr[axis]]
|
|
|
|
|
|
def is_target(names):
|
|
"""Return True if the name of the current target is one of provided names"""
|
|
names = [names] if isinstance(names, str) else names
|
|
target = tvm.target.Target.current(allow_none=False)
|
|
return any(name in target.keys for name in names)
|
|
|
|
|
|
def is_dynamic_shape(shape):
|
|
"""Checks if any part of a shape is dynamic"""
|
|
return any(isinstance(x, Var) for x in shape)
|