wehub-resource-sync
884 lines
18 KiB
Python
884 lines
18 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.
|
|
"""Elementwise operators"""
|
|
|
|
# pylint: disable=redefined-builtin,unused-argument
|
|
import tvm
|
|
from tvm import DataTypeCode, te
|
|
|
|
from . import cpp, tag
|
|
from .utils import get_const_tuple
|
|
|
|
|
|
def _require_float_tensor(op_name, x):
|
|
if not x.dtype.matches_code(DataTypeCode.FLOAT, DataTypeCode.BFLOAT):
|
|
raise TypeError(f"topi.{op_name} only supports floating-point inputs, but got {x.dtype}")
|
|
return x
|
|
|
|
|
|
def _is_integer_tensor(x):
|
|
return x.dtype.matches_code(DataTypeCode.INT, DataTypeCode.UINT)
|
|
|
|
|
|
@tvm.te.tag_scope(tag=tag.ELEMWISE)
|
|
def identity(x):
|
|
"""Take identity of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
# pylint: disable=unnecessary-lambda
|
|
return te.compute(x.shape, lambda *i: x(*i))
|
|
|
|
|
|
@tvm.te.tag_scope(tag=tag.ELEMWISE)
|
|
def negative(x):
|
|
"""Take negation of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
# pylint: disable=unnecessary-lambda
|
|
return te.compute(x.shape, lambda *i: -x(*i))
|
|
|
|
|
|
@tvm.te.tag_scope(tag=tag.ELEMWISE)
|
|
def exp(x):
|
|
"""Take exponential of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
return te.compute(x.shape, lambda *i: te.exp(x(*i)))
|
|
|
|
|
|
@tvm.te.tag_scope(tag=tag.ELEMWISE)
|
|
def erf(x):
|
|
"""Take gauss error function of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
return te.compute(x.shape, lambda *i: te.erf(x(*i)))
|
|
|
|
|
|
@tvm.te.tag_scope(tag=tag.ELEMWISE)
|
|
def tanh(x):
|
|
"""Take hyperbolic tanh of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
return te.compute(x.shape, lambda *i: te.tanh(x(*i)))
|
|
|
|
|
|
@tvm.te.tag_scope(tag=tag.ELEMWISE)
|
|
def tan(x):
|
|
"""Take tan of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
return te.compute(x.shape, lambda *i: te.tan(x(*i)))
|
|
|
|
|
|
@tvm.te.tag_scope(tag=tag.ELEMWISE)
|
|
def cos(x):
|
|
"""Take cos of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
return te.compute(x.shape, lambda *i: te.cos(x(*i)))
|
|
|
|
|
|
@tvm.te.tag_scope(tag=tag.ELEMWISE)
|
|
def cosh(x):
|
|
"""Take cosh of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
return te.compute(x.shape, lambda *i: te.cosh(x(*i)))
|
|
|
|
|
|
@tvm.te.tag_scope(tag=tag.ELEMWISE)
|
|
def sin(x):
|
|
"""Take sin of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
return te.compute(x.shape, lambda *i: te.sin(x(*i)))
|
|
|
|
|
|
@tvm.te.tag_scope(tag=tag.ELEMWISE)
|
|
def sinh(x):
|
|
"""Take sinh of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
return te.compute(x.shape, lambda *i: te.sinh(x(*i)))
|
|
|
|
|
|
@tvm.te.tag_scope(tag=tag.ELEMWISE)
|
|
def acos(x):
|
|
"""Take arc cos of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
x = _require_float_tensor("acos", x)
|
|
return te.compute(x.shape, lambda *i: te.acos(x(*i)))
|
|
|
|
|
|
@tvm.te.tag_scope(tag=tag.ELEMWISE)
|
|
def acosh(x):
|
|
"""Take arc cosh of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
x = _require_float_tensor("acosh", x)
|
|
return te.compute(x.shape, lambda *i: te.acosh(x(*i)))
|
|
|
|
|
|
@tvm.te.tag_scope(tag=tag.ELEMWISE)
|
|
def asin(x):
|
|
"""Take arc sin of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
x = _require_float_tensor("asin", x)
|
|
return te.compute(x.shape, lambda *i: te.asin(x(*i)))
|
|
|
|
|
|
@tvm.te.tag_scope(tag=tag.ELEMWISE)
|
|
def asinh(x):
|
|
"""Take arc sinh of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
x = _require_float_tensor("asinh", x)
|
|
return te.compute(x.shape, lambda *i: te.asinh(x(*i)))
|
|
|
|
|
|
@tvm.te.tag_scope(tag=tag.ELEMWISE)
|
|
def atan(x):
|
|
"""Take atan of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
return te.compute(x.shape, lambda *i: te.atan(x(*i)))
|
|
|
|
|
|
@tvm.te.tag_scope(tag=tag.ELEMWISE)
|
|
def atanh(x):
|
|
"""Take atanh of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
x = _require_float_tensor("atanh", x)
|
|
return te.compute(x.shape, lambda *i: te.atanh(x(*i)))
|
|
|
|
|
|
@tvm.te.tag_scope(tag=tag.ELEMWISE)
|
|
def floor(x):
|
|
"""Take floor of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
return te.compute(x.shape, lambda *i: te.floor(x(*i)))
|
|
|
|
|
|
@tvm.te.tag_scope(tag=tag.ELEMWISE)
|
|
def ceil(x):
|
|
"""Take ceil of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
return te.compute(x.shape, lambda *i: te.ceil(x(*i)))
|
|
|
|
|
|
def sign(x):
|
|
"""Returns -1, 0, 1 based on sign of x.
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
return cpp.sign(x)
|
|
|
|
|
|
@tvm.te.tag_scope(tag=tag.ELEMWISE)
|
|
def trunc(x):
|
|
"""Take truncated value of the input of x, element-wise.
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
return te.compute(x.shape, lambda *i: te.trunc(x(*i)))
|
|
|
|
|
|
@tvm.te.tag_scope(tag=tag.ELEMWISE)
|
|
def abs(x):
|
|
"""Take absolute value of the input of x, element-wise.
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
return te.compute(x.shape, lambda *i: te.abs(x(*i)))
|
|
|
|
|
|
@tvm.te.tag_scope(tag=tag.ELEMWISE)
|
|
def isnan(x):
|
|
"""Check if value of x is NaN, element-wise.
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
return te.compute(x.shape, lambda *i: te.isnan(x(*i)))
|
|
|
|
|
|
@tvm.te.tag_scope(tag=tag.ELEMWISE)
|
|
def isfinite(x):
|
|
"""Check if value of x is finite, element-wise.
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
return te.compute(x.shape, lambda *i: te.isfinite(x(*i)))
|
|
|
|
|
|
@tvm.te.tag_scope(tag=tag.ELEMWISE)
|
|
def isinf(x):
|
|
"""Check if value of x is infinite, element-wise.
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
return te.compute(x.shape, lambda *i: te.isinf(x(*i)))
|
|
|
|
|
|
@tvm.te.tag_scope(tag=tag.ELEMWISE)
|
|
def round(x):
|
|
"""Round elements of x to nearest integer using ties-to-even (banker's rounding).
|
|
|
|
Ties are broken by rounding to the nearest even integer, matching the ONNX Round
|
|
specification and IEEE 754 default rounding mode.
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
return te.compute(x.shape, lambda *i: te.nearbyint(x(*i)))
|
|
|
|
|
|
def log(x):
|
|
"""Take logarithm of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
if x.dtype.matches_code(DataTypeCode.INT):
|
|
x = te.compute(x.shape, lambda *i: x(*i).astype("float32"))
|
|
return te.compute(x.shape, lambda *i: te.log(x(*i)), tag=tag.ELEMWISE)
|
|
|
|
|
|
def log2(x):
|
|
"""Take logarithm to the base 2 of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
if x.dtype.matches_code(DataTypeCode.INT):
|
|
x = te.compute(x.shape, lambda *i: x(*i).astype("float32"))
|
|
return te.compute(x.shape, lambda *i: te.log2(x(*i)), tag=tag.ELEMWISE)
|
|
|
|
|
|
def log10(x):
|
|
"""Take logarithm to the base 10 of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
if x.dtype.matches_code(DataTypeCode.INT):
|
|
x = te.compute(x.shape, lambda *i: x(*i).astype("float32"))
|
|
return te.compute(x.shape, lambda *i: te.log10(x(*i)), tag=tag.ELEMWISE)
|
|
|
|
|
|
@tvm.te.tag_scope(tag=tag.ELEMWISE)
|
|
def sqrt(x):
|
|
"""Take square root of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
if x.dtype.matches_code(DataTypeCode.INT):
|
|
x = te.compute(x.shape, lambda *i: x(*i).astype("float32"))
|
|
return te.compute(x.shape, lambda *i: te.sqrt(x(*i)))
|
|
|
|
|
|
@tvm.te.tag_scope(tag=tag.ELEMWISE)
|
|
def rsqrt(x):
|
|
"""Take inverse square root of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
if x.dtype.matches_code(DataTypeCode.INT):
|
|
x = te.compute(x.shape, lambda *i: x(*i).astype("float32"))
|
|
return te.compute(x.shape, lambda *i: te.rsqrt(x(*i)))
|
|
|
|
|
|
@tvm.te.tag_scope(tag=tag.ELEMWISE)
|
|
def sigmoid(x):
|
|
"""Take sigmoid tanh of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
return te.compute(x.shape, lambda *i: te.sigmoid(x(*i)))
|
|
|
|
|
|
@tvm.te.tag_scope(tag=tag.ELEMWISE)
|
|
def left_shift(x, n):
|
|
"""Take n bits left shift of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor
|
|
Input argument.
|
|
n : int
|
|
Number of bits.
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
return te.compute(x.shape, lambda *i: x(*i) << n)
|
|
|
|
|
|
@tvm.te.tag_scope(tag=tag.ELEMWISE)
|
|
def right_shift(x, n):
|
|
"""Take n bits right shift of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor
|
|
Input argument.
|
|
n : int
|
|
Number of bits.
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
return te.compute(x.shape, lambda *i: x(*i) >> n)
|
|
|
|
|
|
@tvm.te.tag_scope(tag=tag.ELEMWISE)
|
|
def clip(x, a_min, a_max):
|
|
"""Clip (limit) the values in an array. Given an interval, values
|
|
outside the interval are clipped to the interval edges.
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor
|
|
Input argument.
|
|
a_min : tvm.tirx.Expr
|
|
Minimum value.
|
|
a_max : tvm.tirx.Expr
|
|
Maximum value.
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
|
|
def _compute(*indices):
|
|
value = x(*indices)
|
|
const_min = (
|
|
tvm.tirx.Cast(value.ty, a_min)
|
|
if tvm.ir.is_prim_expr(a_min)
|
|
else tvm.tirx.const(a_min, value.ty)
|
|
)
|
|
const_max = (
|
|
tvm.tirx.Cast(value.ty, a_max)
|
|
if tvm.ir.is_prim_expr(a_max)
|
|
else tvm.tirx.const(a_max, value.ty)
|
|
)
|
|
return tvm.te.max(tvm.te.min(value, const_max), const_min)
|
|
|
|
return te.compute(x.shape, _compute)
|
|
|
|
|
|
@tvm.te.tag_scope(tag=tag.ELEMWISE)
|
|
def fixed_point_multiply(x, multiplier, shift):
|
|
"""Fixed point multiplication between data and a fixed point
|
|
constant expressed as multiplier * 2^(-shift), where multiplier
|
|
is a Q-number with 31 fractional bits
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor or Expr
|
|
Input argument.
|
|
multiplier : int
|
|
Multiplier of a fixed floating point number described as multiplier*2^(-shift).
|
|
shift : int
|
|
Shift of a fixed floating point number described as multiplier*2^(-shift).
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
|
|
def _compute(*indices):
|
|
value = x(*indices)
|
|
return tvm.tirx.q_multiply_shift(
|
|
value,
|
|
tvm.tirx.const(multiplier, "int32"),
|
|
tvm.tirx.const(31, "int32"),
|
|
tvm.tirx.const(shift, "int32"),
|
|
)
|
|
|
|
return te.compute(x.shape, _compute)
|
|
|
|
|
|
@tvm.te.tag_scope(tag=tag.BROADCAST)
|
|
def fixed_point_multiply_per_axis(
|
|
x: te.Tensor,
|
|
y: te.Tensor,
|
|
lshift: te.Tensor,
|
|
rshift: te.Tensor,
|
|
is_lshift_required: int,
|
|
is_rshift_required: int,
|
|
axes,
|
|
):
|
|
"""Fixed point multiplication between data and a fixed point constant expressed as
|
|
multiplier * 2^(-shift), where multiplier is a Q-number with 31 fractional bits
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor
|
|
Input argument.
|
|
y : tvm.te.Tensor
|
|
Multiplier of a fixed floating point number described as multiplier*2^(-shift).
|
|
lshift : tvm.te.Tensor
|
|
Left shifts of a fixed floating point number described as multiplier*2^(-shift).
|
|
rshift : tvm.te.Tensor
|
|
Right shifts of a fixed floating point number described as multiplier*2^(-shift).
|
|
is_lshift_required : int
|
|
Whether we need to do left shift or not.
|
|
is_rshift_required : int
|
|
Whether we need to do right shift or not.
|
|
|
|
Returns
|
|
-------
|
|
z : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
|
|
def _compute(*indices):
|
|
elements = []
|
|
for element in get_const_tuple(axes):
|
|
elements += [indices[element]]
|
|
param_indices = tuple(elements)
|
|
|
|
value = x(*indices)
|
|
m = y(*param_indices)
|
|
l_shift = lshift(*param_indices)
|
|
r_shift = rshift(*param_indices)
|
|
return tvm.tirx.q_multiply_shift_per_axis(
|
|
value,
|
|
m,
|
|
l_shift,
|
|
r_shift,
|
|
tvm.tirx.const(31, "int32"),
|
|
tvm.tirx.const(is_lshift_required, "bool"),
|
|
tvm.tirx.const(is_rshift_required, "bool"),
|
|
)
|
|
|
|
return te.compute(x.shape, _compute)
|
|
|
|
|
|
def cast(x, dtype, span=None):
|
|
"""Cast input to specified data type.
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor or Expr
|
|
Input argument.
|
|
|
|
dtype : str
|
|
Data type.
|
|
|
|
span : Optional[Span]
|
|
The location of the cast in the source.
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
if isinstance(x, te.tensor.Tensor):
|
|
return te.compute(x.shape, lambda *i: x(*i).astype(dtype), tag=tag.ELEMWISE)
|
|
# pylint: disable=import-outside-toplevel
|
|
from tvm.tirx import _ffi_api
|
|
|
|
return _ffi_api._cast(dtype, x, span)
|
|
|
|
|
|
def reinterpret(x, dtype):
|
|
"""Reinterpret input to specified data type.
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor
|
|
Input argument.
|
|
|
|
dtype : str
|
|
Data type.
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
return cpp.reinterpret(x, dtype)
|
|
|
|
|
|
def fast_exp(x):
|
|
"""Take exponential of input x using fast_exp implementation
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
if _is_integer_tensor(x):
|
|
x = cast(x, "float32")
|
|
return cpp.fast_exp(x, x.dtype, tag.ELEMWISE)
|
|
|
|
|
|
def fast_tanh(x):
|
|
"""Take hyperbolic tangent of input x using fast_tanh implementation
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
if _is_integer_tensor(x):
|
|
x = cast(x, "float32")
|
|
return cpp.fast_tanh(x, x.dtype, tag.ELEMWISE)
|
|
|
|
|
|
def fast_erf(x):
|
|
"""Take gauss error function of input x using fast_erf implementation.
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
return cpp.fast_erf(x, x.dtype, tag.ELEMWISE)
|
|
|
|
|
|
def ceil_log2(x):
|
|
"""Compute integer ceil log2 with a special code path for vulkan
|
|
SPIR-V does not support log2 on fp64. Instead, we compute integer ceil_log2 via clz
|
|
intrinsic when the target is vulkan.
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
if not tvm.ir.is_prim_expr(x):
|
|
x = tvm.tirx.const(x)
|
|
|
|
if x.ty.matches_code(DataTypeCode.FLOAT, DataTypeCode.BFLOAT):
|
|
return tvm.tirx.ceil(tvm.tirx.log2(x))
|
|
|
|
target = tvm.target.Target.current()
|
|
|
|
if target is not None:
|
|
target_name = target.kind.name
|
|
if "vulkan" in target_name:
|
|
clz = tvm.tirx.clz(x)
|
|
bits = x.ty.dtype.bits
|
|
res = tvm.tirx.if_then_else(x & (x - 1) == 0, bits - clz - 1, bits - clz)
|
|
if res.ty != x.ty:
|
|
return cast(res, x.ty)
|
|
return res
|
|
|
|
if "adreno" in str(target.attrs.get("device", "")) or target_name in [
|
|
"metal",
|
|
"rocm",
|
|
"webgpu",
|
|
]:
|
|
return cast(tvm.tirx.ceil(tvm.tirx.log2(cast(x, "float32"))), x.ty)
|
|
|
|
return cast(tvm.tirx.ceil(tvm.tirx.log2(cast(x, "float64"))), x.ty)
|