wehub/apache--tvm
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
Lint / lint (push) Has been cancelled
Details
CI / MacOS (push) Has been cancelled
Details
CI / Windows (push) Has been cancelled
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:36:25 +08:00
commit 26446540fa
3151 changed files with 974126 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/tvm/relax/op/__init__.py
+240
View File
@@ -0,0 +1,240 @@
# isort: skip_file
# 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= redefined-builtin
"""Relax core operators."""
# Register operator gradient functions
from . import _op_gradient, builtin, ccl, distributed, grad, image, memory, nn, op_attrs
# Operators
from .base import (
assert_op,
call_builtin_with_ctx,
call_dps_packed,
call_inplace_packed,
call_pure_packed,
call_py_func,
call_tir,
call_tir_inplace,
call_tir_with_grad,
hint_on_device,
invoke_closure,
invoke_pure_closure,
make_closure,
null_value,
print,
register_gradient,
shape_of,
shape_to_tensor,
size,
tensor_to_shape,
to_vdevice,
)
from .binary import (
add,
atan2,
bitwise_and,
bitwise_or,
bitwise_xor,
divide,
equal,
floor_divide,
log_add_exp,
floor_mod,
greater,
greater_equal,
left_shift,
less,
less_equal,
logical_and,
logical_or,
logical_xor,
maximum,
minimum,
mod,
multiply,
not_equal,
power,
right_shift,
subtract,
)
from .create import (
arange,
full,
full_like,
hamming_window,
ones,
ones_like,
eye,
eye_like,
tril,
triu,
zeros,
zeros_like,
)
from .datatype import astype, wrap_param
from .index import dynamic_strided_slice, strided_slice, take
from .linear_algebra import einsum, linear, matmul, outer
from .manipulate import (
broadcast_to,
collapse_sum_like,
collapse_sum_to,
concat,
expand_dims,
flatten,
flip,
gather_elements,
gather_nd,
index_put,
index_tensor,
meshgrid,
layout_transform,
one_hot,
permute_dims,
repeat,
reshape,
reverse_sequence,
scatter_elements,
scatter_nd,
slice_scatter,
split,
squeeze,
stack,
tile,
)
from .mask import masked_fill
from .qdq import dequantize, quantize
from .sampling import multinomial_from_uniform
from .search import argmax, argmin, where, bucketize
from .set import nonzero, unique
from .sorting import argsort, sort, topk
from .statistical import cumprod, cumsum, max, mean, min, prod, std, sum, variance, median
from .ternary import ewise_fma
from .unary import (
abs,
acos,
acosh,
asin,
asinh,
atan,
atanh,
bitwise_not,
ceil,
clip,
cos,
cosh,
erf,
exp,
floor,
isfinite,
isinf,
isnan,
log,
logical_not,
negative,
round,
rsqrt,
sigmoid,
sign,
sin,
sinh,
sqrt,
square,
tan,
tanh,
trunc,
)
from .vision import (
all_class_non_max_suppression,
get_valid_counts,
multibox_transform_loc,
non_max_suppression,
roi_align,
roi_pool,
)
def _register_op_make():
# pylint: disable=import-outside-toplevel
from .. import expr
from tvm.ir import _tensor_expr_overload
from . import _ffi_api
expr._op_ffi_api = _ffi_api # type: ignore
def _add(lhs, rhs):
if isinstance(lhs.ty, expr.tvm.relax.TupleType) and isinstance(rhs, tuple):
return tuple([*lhs, *rhs])
return expr._binary_op_helper(lhs, rhs, _ffi_api.add)
def _rhs(_lhs, rhs):
return expr._binary_rhs_helper(rhs)
def _getitem(value, index):
try:
return expr.TupleGetItem(value, index)
except RuntimeError as err:
if "Index out of bounds" in err.args[0]:
raise IndexError from err
raise
_tensor_expr_overload.astype = lambda lhs, dtype, _span=None: _ffi_api.astype(lhs, dtype)
_tensor_expr_overload.__call__ = lambda func, *args, attrs=None: expr.tvm.ir.Call(
func, args, attrs=attrs
)
_tensor_expr_overload.__getitem__ = _getitem
_tensor_expr_overload.__neg__ = lambda lhs: _ffi_api.negative(lhs)
_tensor_expr_overload.__lt__ = lambda lhs, rhs: expr._binary_op_helper(lhs, rhs, _ffi_api.less)
_tensor_expr_overload.__le__ = lambda lhs, rhs: expr._binary_op_helper(
lhs, rhs, _ffi_api.less_equal
)
_tensor_expr_overload.__gt__ = lambda lhs, rhs: expr._binary_op_helper(
lhs, rhs, _ffi_api.greater
)
_tensor_expr_overload.__ge__ = lambda lhs, rhs: expr._binary_op_helper(
lhs, rhs, _ffi_api.greater_equal
)
_tensor_expr_overload.__add__ = _add
_tensor_expr_overload.__radd__ = _add
_tensor_expr_overload.__sub__ = lambda lhs, rhs: expr._binary_op_helper(
lhs, rhs, _ffi_api.subtract
)
_tensor_expr_overload.__rsub__ = _rhs
_tensor_expr_overload.__mul__ = lambda lhs, rhs: expr._binary_op_helper(
lhs, rhs, _ffi_api.multiply
)
_tensor_expr_overload.__rmul__ = _tensor_expr_overload.__mul__
_tensor_expr_overload.__div__ = lambda lhs, rhs: expr._binary_op_helper(
lhs, rhs, _ffi_api.divide
)
_tensor_expr_overload.__rdiv__ = _rhs
_tensor_expr_overload.__truediv__ = _tensor_expr_overload.__div__
_tensor_expr_overload.__rtruediv__ = _rhs
_tensor_expr_overload.__floordiv__ = lambda lhs, rhs: expr._binary_op_helper(
lhs, rhs, _ffi_api.floor_divide
)
_tensor_expr_overload.__rfloordiv__ = _rhs
_tensor_expr_overload.__mod__ = lambda lhs, rhs: expr._binary_op_helper(lhs, rhs, _ffi_api.mod)
_tensor_expr_overload.__rmod__ = _rhs
_tensor_expr_overload.__pow__ = lambda lhs, rhs: expr._binary_op_helper(
lhs, rhs, _ffi_api.power
)
_tensor_expr_overload.__rpow__ = _rhs
_register_op_make()
python/tvm/relax/op/_ffi_api.py
+20
View File
@@ -0,0 +1,20 @@
# 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
"""FFI APIs for tvm.relax.op"""
import tvm_ffi
tvm_ffi.init_ffi_api("relax.op", __name__)
python/tvm/relax/op/_op_gradient.py
+1317
View File
File diff suppressed because it is too large Load Diff
python/tvm/relax/op/base.py
+880
View File
@@ -0,0 +1,880 @@
# 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
# pylint: disable=redefined-builtin
# ruff: noqa: F821
"""The base Relax operators."""
from collections.abc import Callable
import tvm_ffi
import tvm
import tvm.runtime
from tvm.ir import Call
from tvm.runtime import Object, ObjectConvertible
from ..expr import Expr, ExternFunc, GlobalVar, ShapeExpr, StringImm, Var
from ..type import TensorType, Type
from ..utils import convert_to_expr
from . import _ffi_api
py_print = print # pylint: disable=invalid-name
def register_gradient(
op_name: str,
fgradient: Callable[[Var, Call, Var, "BlockBuilder"], list[Expr]] | None = None,
level: int = 10,
):
"""Register operator gradient function for a relax operator.
Parameters
----------
op_name: str
The name of the op.
fgradient: function (orig_var: Var, orig_call: Call, output_grad: Var, ctx: BlockBuilder)
-> partials: List[Expr]
The gradient function being used.
level: int
The priority level
"""
return tvm.ir.register_op_attr(op_name, "FPrimalGradient", fgradient, level)
def null_value() -> Call:
"""Create a call node that represents a null value object.
Returns
-------
ret: Call
The created call node.
"""
return _ffi_api.null_value() # type: ignore
def _wrap_inline_arg_tuple(args) -> Expr:
"""Helper function to wrap argument tuple
Normalize the arguments provided the functions that accept a tuple
of arguments, and require the tuple of arguments to be written
in-line. If the arguments provided are a single relax expression,
and are not a reference to a relax tuple, then wrap them into an
in-line relax Tuple.
"""
if isinstance(args, tuple | list):
return tvm.relax.Tuple([convert_to_expr(a) for a in args])
elif (
isinstance(args, Expr)
and not isinstance(args, tvm.relax.Tuple)
and (args.ty is None or not isinstance(args.ty, tvm.relax.TupleType))
):
return tvm.relax.Tuple([args])
else:
return args
def call_tir(
gvar: GlobalVar,
args: Expr,
out_ty: TensorType | list[TensorType],
tir_vars: ShapeExpr | tuple[Expr] | list[Expr] | None = None,
) -> Call:
"""
Call a tirx.prim_func and return the output.
Parameters
----------
gvar : GlobalVar
The GlobalVar referring to a tirx PrimFunc.
args : Expr
The input arguments.
out_ty : Union[TensorType, List[TensorType]]
The type information of the call_tir output.
It should be a single or a list of TensorType. Each one denotes the
type information of a returned tensor.
tir_vars : Optional[Union[ShapeExpr, Tuple[Expr], List[Expr]]]
ShapeExpr representing a tuple of integers to unpack when calling func. Is null if not used
Returns
-------
ret: Call
A call node for the call_tir operator.
"""
args = _wrap_inline_arg_tuple(args)
if not isinstance(out_ty, list):
out_ty = [out_ty]
if isinstance(tir_vars, list | tuple):
tir_vars = ShapeExpr(tir_vars)
return _ffi_api.call_tir(gvar, args, out_ty, tir_vars) # type: ignore
def call_tir_with_grad(
gvar: GlobalVar,
args: Expr,
out_ty: TensorType | list[TensorType],
te_grad_name: str,
te_grad_kwargs: dict[str, Object] | None = None,
tir_vars: ShapeExpr | tuple[Expr] | list[Expr] | None = None,
) -> Call:
"""
Call a tirx.prim_func and return the output. This intrinsic will bind a te gradient function
(refered by te_grad_name) to the call_tir_with_grad node. The te gradient function will be
called by the Gradient pass.
Parameters
----------
gvar : GlobalVar
The GlobalVar referring to a tirx PrimFunc.
args : Expr
The input arguments.
out_ty : Union[TensorType, List[TensorType]]
The type information of the call_tir_with_grad output.
It should be a single or a list of TensorType. Each one denotes the
type information of a returned tensor.
te_grad_name : str
The registered name of the te gradient function associated with the call_tir_with_grad
node. Must be provided as a keyword argument.
te_grad_kwargs : Dict[str, Object], optional
The keyword arguments passed to the te gradient function.
Optionally provided as a keyword argument. Default: {}.
tir_vars : Optional[Union[ShapeExpr, Tuple[Expr], List[Expr]]]
ShapeExpr representing a tuple of integers to unpack when calling func. Is null if not used
Returns
-------
ret: Call
A call node for the call_tir_with_grad operator.
"""
args = _wrap_inline_arg_tuple(args)
if not isinstance(out_ty, list):
out_ty = [out_ty]
if isinstance(tir_vars, list | tuple):
tir_vars = ShapeExpr(tir_vars)
if te_grad_kwargs is None:
te_grad_kwargs = {}
return _ffi_api.call_tir_with_grad( # type: ignore
gvar, args, out_ty, te_grad_name, te_grad_kwargs, tir_vars
)
def call_tir_inplace(
gvar: GlobalVar,
args: Expr,
inplace_indices: int | list[int],
out_ty: TensorType | list[TensorType],
tir_vars: ShapeExpr | tuple[Expr] | list[Expr] | None = None,
) -> Call:
"""
Call a TIR PrimFunc and return the result, doing the specified computations in-place
(based on the `inplace_indices` argument; outputs will alias the inputs
selected by in-place indices).
Warning: This operator is considered pure by the type system but actually mutates
the arguments specified by `inplace_indices`. This operator should not be used directly,
but rather should be inserted by passes that have checked whether it is safe to perform
operations in-place (i.e., none of the arguments specified as an output is aliased or is
live after calling call_tir_inplace).
Direct calls to this operator should be done for testing purposes only.
Parameters
----------
gvar : GlobalVar
The GlobalVar referring to a TIR PrimFunc.
args : Expr
The input arguments.
inplace_indices : Union[int, List[int]]
Specify which arguments should be used for in-place computations.
If `inplace_indices` is a single integer, it will be made into a singleton list.
Suppose `inplace_indices[i] = j`, where `j >= 0`. Then the `i`th output
will be an alias of `args[j]`.
If `inplace_indices[i] = -1`, then the `i`th output will be a freshly allocated tensor.
At least one member of `inplace_indices` must not be -1.
out_ty : Union[TensorType, List[TensorType]]
The type information of the call_tir_inplace output.
It should be a single `TensorType` or a list of `TensorType`.
Each one denotes the type information of a returned tensor.
If a list of `TensorType` is given, the result will be a tuple of `TensorType`.
tir_vars : Optional[Union[ShapeExpr, Tuple[Expr], List[Expr]]]
ShapeExpr representing a tuple of integers to unpack when calling func. Is null if not used
Returns
-------
ret: Call
A call node for the call_tir operator.
"""
args = _wrap_inline_arg_tuple(args)
if not isinstance(inplace_indices, list):
inplace_indices = [inplace_indices]
if not isinstance(out_ty, list):
out_ty = [out_ty]
if isinstance(tir_vars, list | tuple):
tir_vars = ShapeExpr(tir_vars)
return _ffi_api.call_tir_inplace( # type: ignore
gvar,
args,
inplace_indices,
out_ty,
tir_vars,
)
def call_dps_packed(
func: str | Expr,
args: Expr,
out_ty: TensorType | list[TensorType],
) -> Call:
"""
Call a destination-passing-style packed function and return the output.
Note: The called function is assumed to be _pure_ (other than modifying the designated
output arguments). If the function _does_ result in other side effects, then the compiler
may end up removing, reordering, or repeating those effects--no guarantees can be made.
Parameters
----------
func : Union[str, Expr]
The destination-passing-style function, can be ExternFunc.
args : Expr
The input arguments.
out_ty : Union[TensorType, List[TensorType]]
The type information of the call_dps_packed output.
It should be a single or a list of TensorType. Each one denotes the
type information of a returned tensor.
Returns
-------
ret: Call
A call node for the call_dps_packed operator.
"""
if isinstance(func, str):
func = ExternFunc(func)
args = _wrap_inline_arg_tuple(args)
if not isinstance(out_ty, list):
out_ty = [out_ty]
return _ffi_api.call_dps_packed(func, args, out_ty) # type: ignore
def call_py_func(
func_name: str,
args: Expr,
out_ty: TensorType | list[TensorType],
) -> Call:
"""
Call a Python function and return the output.
Parameters
----------
func_name : str
The name of the Python function to call. This should correspond to a function
in the IRModule's pyfuncs attribute.
args : Expr
The input arguments.
out_ty : Union[TensorType, List[TensorType]]
The type information of the call_py_func output.
It should be a single or a list of TensorType. Each one denotes the
type information of a returned tensor.
Returns
-------
ret: Call
A call node for the call_py_func operator.
"""
args = _wrap_inline_arg_tuple(args)
if not isinstance(out_ty, list):
out_ty = [out_ty]
return _ffi_api.call_py_func(func_name, args, out_ty) # type: ignore
def call_builtin_with_ctx(
func: str | Expr,
args: Expr,
*,
ty_args: Type | list[Type] | None = None,
) -> Call:
"""Call a builtin function func.
Parameters
----------
func : Expr
The builtin function to be called.
args : Expr
The input arguments.
ty_args: Optional[Union[Type, List[Type]]]
The type arguments to the call node.
Returns
-------
ret: Call
The created call node.
"""
if isinstance(func, str):
func = ExternFunc(func)
args = _wrap_inline_arg_tuple(args)
if ty_args is not None and not isinstance(ty_args, list | tuple):
ty_args = [ty_args]
return _ffi_api.call_builtin_with_ctx( # type: ignore
func,
args,
ty_args, # type: ignore
)
def make_closure(
func: Expr,
args: Expr,
) -> Object:
"""
Create a closure with free variables and return the closure.
Parameters
----------
func : Expr
The closure, can be ExternFunc or PrimFunc.
args : Expr
The input arguments.
Returns
-------
ret: Object
The VMClosure.
"""
args = _wrap_inline_arg_tuple(args)
return _ffi_api.make_closure(func, args) # type: ignore
def invoke_closure(
closure: Expr,
args: Expr,
ty_args: list[Type] | Type,
) -> Call:
"""
Invoke a closure.
Parameters
----------
closure : Expr
The VMClosure object.
args : Expr
The input arguments.
type_args: Union[List[Type], Type]
The type information arguments of the CallNode
Returns
-------
ret: Call
A call to `invoke_closure`.
"""
args = _wrap_inline_arg_tuple(args)
if not isinstance(ty_args, list | tuple):
ty_args = [ty_args]
return _ffi_api.invoke_closure(closure, args, ty_args) # type: ignore
def render_object(val: tvm.Object) -> str:
"""
Given a TVM Object, renders it in string form. Used for Relax printing and assertions.
Parameters
----------
val: tvm.Object
An object to render
Returns
-------
ret: str
A string representing the value, ideally human-readable
"""
if isinstance(val, tvm.runtime.Tensor):
return str(val)
if isinstance(val, tvm_ffi.Array):
fields = ", ".join([render_object(val[i]) for i in range(len(val))])
return f"({fields})"
return str(val)
@tvm.register_global_func("relax.run.shape_to_tensor")
def relax_shape_to_tensor(shape_tuple: tvm_ffi.Shape) -> tvm.runtime.Tensor:
"""
Takes a Shape and convert it to Tensor.
Parameters
----------
shape_tuple: tvm_ffi.Shape
Shape tuple that we want to convert to Tensor at runtime
"""
return tvm.runtime.tensor([int(v) for v in shape_tuple])
@tvm.register_global_func("relax.run.print")
def relax_print(format_str: str, *format_args: tvm.Object) -> None:
"""
Takes a list of values to print, formats with the given format string.
If the format string is empty, simply prints.
Call from TVM script like this:
`relax.print(value1, value2, ..., valueN, format=format_str)`
or
`relax.print(value1, value2, ..., valueN) # format_str defaults to ""`
Parameters
----------
format_str: str
The last argument is a Python-style format string for printing the value
format_args: List[Object]
The values to print.
"""
val_strs = map(render_object, format_args)
if format_str == "":
py_print(*val_strs)
else:
py_print(format_str.format(*val_strs))
def print(*values: list[Expr], format: str | Expr = "") -> Expr:
"""Print op to print the values
Parameters
----------
values : List[Expr]
The values to print.
format: Union[str, Expr]
The format string or StringImm.
Returns
-------
result : Expr
A relax Call, which will print the value during runtime.
"""
if isinstance(format, str):
format = StringImm(format)
return _ffi_api.print(values, format) # type: ignore # pylint: disable=no-member
@tvm.register_global_func("relax.run.assert_op")
def relax_assert_op(condition: tvm.Object, format_str: str, *format_args: tvm.Object) -> None:
"""
A variadic function. The first value serves as the assertion condition:
If the condition is true, then the operator does nothing.
If the condition is false, then the operator raises an assertion error.
Arguments after the first value serve as format arguments for the error message;
the last argument must be a format string for the error message (empty by default).
If the format string is the empty string, then the error message will simply include
a comma-separated list of the format arguments.
The condition argument is not included in the format string.
Parameters
----------
condition: tvm.Object
The assertion condition. Must be a boolean scalar.
format_str: str
The last argument is a Python-style format string for printing the value
format_args: List[tvm.Object]
Values used for formatting the string.
"""
if not isinstance(format_str, str):
raise ValueError(
f"The format string argument to assert must be a string, given {type(format_str)})"
)
if isinstance(condition, bool | int):
val = condition
elif isinstance(condition, tvm.runtime.Tensor):
# may happen if the original program had unknown shape or dtype for the tensor's type
dtype = condition.dtype
if dtype != "bool":
raise ValueError(f"The condition must be a bool scalar, but given a {dtype} tensor")
shape = condition.shape
if len(shape) != 0:
raise ValueError(f"The condition must be a scalar, but it has a shape of {shape}")
val = condition.numpy()
else:
# should be guaranteed by the type system
raise ValueError(
f"The condition for relax assert must be a bool, int, or Tensor, "
f"but received a {type(condition)}."
)
if not val:
error_message = "Assertion Failed"
if format_args or format_str != "":
rendered = map(render_object, format_args)
if format_str != "":
error_message = format_str.format(*rendered)
else:
error_message = ", ".join(rendered)
raise AssertionError(error_message)
def assert_op(
condition: Expr,
format_args: Expr | list[Expr] | None = None,
format: str | Expr = "",
) -> Expr:
"""
Create a call to Relax's assert_op operation (`assert` is reserved in Python,
so the name must be distinct).
Parameters
----------
condition: Expr
The assertion condition.
format_args: Optional[Union[Expr, List[Expr]]]
Format arguments for the error message if the condition fails.
format: Union[str, Expr]
The format string or StringImm for the error message.
Returns
-------
result : Expr
A Call to the Relax assert operation.
"""
if not isinstance(condition, Expr):
condition = tvm.relax.prim_value(condition)
if format_args is None:
format_args = []
elif isinstance(format_args, Expr):
format_args = [format_args]
if isinstance(format, str):
format = StringImm(format)
return _ffi_api.assert_op(condition, format_args, format) # type: ignore
def shape_of(expr: Expr) -> Expr:
"""Get shape of a tensor.
Parameters
----------
expr : Expr
The input Expr.
Returns
-------
result : Expr
A relax Call, which gets the shape of the input
"""
return _ffi_api.shape_of(expr) # type: ignore # pylint: disable=no-member
def size(expr: Expr) -> Expr:
"""Get the total number of elements in a tensor.
Parameters
----------
expr : Expr
The input tensor.
Returns
-------
result : Expr
A scalar tensor of dtype int64 containing the total number of elements.
"""
return _ffi_api.size(expr) # type: ignore # pylint: disable=no-member
def tensor_to_shape(expr: Expr) -> Expr:
"""Convert tensor to shape expr.
Parameters
----------
expr : Expr
The input Expr
Returns
-------
result : Expr
A relax Call, which transforms the tensor values to the shape
"""
return _ffi_api.tensor_to_shape(expr) # type: ignore # pylint: disable=no-member
def shape_to_tensor(expr: Expr) -> Expr:
"""Convert shape to tensor expr.
Parameters
----------
expr : Expr
The input Expr
Returns
-------
result : Expr
A relax Call, which transforms the shape values to the tensor
"""
return _ffi_api.shape_to_tensor(expr) # type: ignore # pylint: disable=no-member
def call_inplace_packed(
func: str | ExternFunc | GlobalVar,
*args: Expr,
inplace_indices: int | list[int],
ty_args: Type | list[Type],
) -> Expr:
"""
Construct a call to a packed function that consumes some of its arguments "in-place"
and returns the mutated arguments (aliased), but should be considered to be otherwise pure.
The `inplace_indices` argument indicates which of the outputs are mutated arguments.
The resulting call will have the same semantics as calling the packed function directly.
Note: This should be used for cases when the user knows that calling the packed function
with these arguments will **in reality** not cause any other side effects.
If it is used for a call that **does** result in other side effects, then the compiler
may end up removing, reordering, or repeating that call, with no guarantees
made about any side effects from the callee.
Warning: This operator as treated as pure by the type system even though it *is* performing
side effects (mutating some arguments). It is therefore incumbent upon the user to ensure
that it is being used safely (viz., that mutated arguments are not live after the mutation,
that they do not alias values live after the mutation).
Parameters
----------
func : Union[str, ExternFunc]
The name (global symbol) for a PackedFunc or an ExternFunc node.
args: Expr
The arguments for the PackedFunc.
inplace_indices : Union[int, List[int]]
Specify which arguments should be used for in-place computations.
If `inplace_indices` is a single integer, it will be made into a singleton list.
Suppose `inplace_indices[i] = j`, where `j >= 0`. Then the `i`th output
will be an alias of `args[j]`.
If `inplace_indices[i] = -1`, then the `i`th output will be a freshly allocated tensor.
At least one member of `inplace_indices` must not be -1.
ty_args: Union[Type, List[Type]]
The list of type information arguments (giving the type information for the returned value).
Returns
-------
result : Expr
A Relax call, corresponding to
`call_pure_packed(ExternFunc(func), args, DictAttrs(kwargs), ty_args)`
"""
if isinstance(func, ExternFunc):
func = func.global_symbol
op = ExternFunc(func)
args = tuple(convert_to_expr(a) for a in args)
if ty_args is None:
raise ValueError("R.call_pure_packed is required to have type_args")
if isinstance(ty_args, tuple): # type: ignore
ty_args = list(ty_args)
elif not isinstance(ty_args, list):
ty_args = [ty_args]
if not isinstance(inplace_indices, list):
inplace_indices = [inplace_indices]
return _ffi_api.call_inplace_packed(op, args, inplace_indices, ty_args) # type: ignore # pylint: disable=no-member
def call_pure_packed(
func: str | ExternFunc | GlobalVar,
*args: Expr,
ty_args: Type | list[Type],
) -> Expr:
"""
Construct a call to a packed function that should be treated as pure,
even though packed calls are normally not treated as pure.
The resulting call will have the same semantics as calling the packed function directly.
Note: This should be used for cases when the user knows that calling the packed function
with these arguments will **in reality** not cause any side effects.
If it is used for a call that **does** result in side effects, then the compiler
may end up removing, reordering, or repeating that call, with no guarantees
made about any side effects from the callee.
Parameters
----------
func : Union[str, ExternFunc]
The name (global symbol) for a PackedFunc or an ExternFunc node.
args: Expr
The arguments for the PackedFunc.
ty_args: Union[Type, List[Type]]
The list of type information arguments (giving the type information for the returned value).
Returns
-------
result : Expr
A Relax call, corresponding to
`call_pure_packed(ExternFunc(func), args, DictAttrs(kwargs), ty_args)`
"""
if isinstance(func, ExternFunc):
func = func.global_symbol
op = ExternFunc(func)
args = tuple(convert_to_expr(a) for a in args)
if ty_args is None:
raise ValueError("R.call_pure_packed is required to have type_args")
if isinstance(ty_args, tuple): # type: ignore
ty_args = list(ty_args)
elif not isinstance(ty_args, list):
ty_args = [ty_args]
ty_args = [
(ty() if callable(ty) else ty.asobject() if isinstance(ty, ObjectConvertible) else ty)
for ty in ty_args
]
# note: if we need attributes, we can also take them here
return _ffi_api.call_pure_packed(op, args, None, ty_args) # type: ignore # pylint: disable=no-member
def invoke_pure_closure(
closure: Expr,
args: Expr,
ty_args: list[Type] | Type,
) -> Call:
"""
Invoke a closure and indicate to the compiler that it is pure.
Note: This should be used for cases when the user knows that calling the closure
with these arguments will **in reality** not cause any side effects.
If it is used for a call that _does_ result in side effects, then the compiler
may end up removing, reordering, or repeating that call, with no guarantees
made about any side effects from the callee.
Parameters
----------
closure : Expr
The VMClosure object.
args : Expr
The input arguments.
type_args: Union[List[Type], Type]
The type information arguments of the CallNode
Returns
-------
ret: Call
A call to `invoke_pure_closure`.
"""
args = _wrap_inline_arg_tuple(args)
if not isinstance(ty_args, list | tuple):
ty_args = [ty_args]
return _ffi_api.invoke_pure_closure(closure, args, ty_args) # type: ignore
def to_vdevice(data, dst_vdevice) -> Expr:
"""Copy data to the destination device. This
operator helps data transferring between difference devices for
heterogeneous execution.
Parameters
----------
data : Expr
The tensor to be copied.
dst_device : VDevice
The destination device where the data is copied to.
Returns
-------
result : Expr
The copied result.
"""
return _ffi_api.to_vdevice(data, dst_vdevice) # type: ignore
def hint_on_device(data, dst_vdevice, memory_scope="global") -> Expr:
"""It provides a hint specifying the device on which the input data should be executed.
This hint is utilized by RealizeVDevice to propagate the virtual device."
Parameters
----------
data : Expr
The tensor to be copied.
dst_device : Device
The destination device where the data is supposed to be executed.
memory_scope: String
Memory scope of buffer on target device.
Returns
-------
result : Expr
The result.
"""
return _ffi_api.hint_on_device(data, dst_vdevice, memory_scope) # type: ignore
python/tvm/relax/op/binary.py
+484
View File
@@ -0,0 +1,484 @@
# 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=redefined-builtin, invalid-name
"""Relax binary arithmetic and comparison operators."""
from ..expr import Expr
from . import _ffi_api
###################### Arithmetic operators ######################
def add(x1: Expr, x2: Expr) -> Expr:
"""Addition with numpy-style broadcasting.
Parameters
----------
x1 : Expr
The first input tensor.
x2 : Expr
The second input tensor.
Returns
-------
result : Expr
The computed result.
Examples
--------
.. code:: python
bb = relax.BlockBuilder()
a = relax.Var("a", relax.TensorType(shape=(2, 3), dtype="float32"))
b = relax.Var("b", relax.TensorType(shape=(2, 1), dtype="float32"))
c = bb.normalize(relax.op.add(a, b)) # c has TensorType(shape=(2, 3), dtype="float32")
"""
return _ffi_api.add(x1, x2) # type: ignore
def divide(x1: Expr, x2: Expr) -> Expr:
"""Division with numpy-style broadcasting.
Parameters
----------
x1 : relax.Expr
The first input tensor.
x2 : relax.Expr
The second input tensor.
Returns
-------
result : relax.Expr
The computed result.
"""
return _ffi_api.divide(x1, x2) # type: ignore
def floor_divide(x1: Expr, x2: Expr) -> Expr:
"""Floor division with numpy-style broadcasting.
Parameters
----------
x1 : relax.Expr
The first input tensor.
x2 : relax.Expr
The second input tensor.
Returns
-------
result : relax.Expr
The computed result.
"""
return _ffi_api.floor_divide(x1, x2) # type: ignore
def log_add_exp(x1: Expr, x2: Expr) -> Expr:
"""
Compute the log of the sum of exponentials of the inputs, element-wise.
Parameters
----------
x1 : Expr
The first input tensor.
x2 : Expr
The second input tensor.
Returns
-------
Expr
The element-wise log-sum-exp of `x1` and `x2`.
"""
return _ffi_api.log_add_exp(x1, x2)
def multiply(x1: Expr, x2: Expr) -> Expr:
"""Multiplication with numpy-style broadcasting.
Parameters
----------
x1 : Expr
The first input tensor.
x2 : Expr
The second input tensor.
Returns
-------
result : Expr
The computed result.
"""
return _ffi_api.multiply(x1, x2) # type: ignore
def power(x1: Expr, x2: Expr):
"""Power with numpy-style broadcasting.
Parameters
----------
x1 : relax.Expr
The first input tensor.
x2 : relax.Expr
The second input tensor.
Returns
-------
result : relax.Expr
The computed result.
"""
return _ffi_api.power(x1, x2) # type: ignore
def atan2(x1: Expr, x2: Expr) -> Expr:
"""Atan2 with numpy-style broadcasting.
Parameters
----------
x1 : relax.Expr
The first input tensor (y-coordinates).
x2 : relax.Expr
The second input tensor (x-coordinates).
Returns
-------
result : relax.Expr
The computed result.
"""
return _ffi_api.atan2(x1, x2) # type: ignore
def subtract(x1: Expr, x2: Expr) -> Expr:
"""Subtraction with numpy-style broadcasting.
Parameters
----------
x1 : relax.Expr
The first input tensor.
x2 : relax.Expr
The second input tensor.
Returns
-------
result : relax.Expr
The computed result.
"""
return _ffi_api.subtract(x1, x2) # type: ignore
def mod(x1: Expr, x2: Expr) -> Expr:
"""Modulo with numpy-style broadcasting.
Parameters
----------
x1 : Expr
The first input tensor.
x2 : Expr
The second input tensor.
"""
return _ffi_api.mod(x1, x2) # type: ignore
def floor_mod(x1: Expr, x2: Expr) -> Expr:
"""Floor modulo with numpy-style broadcasting.
Parameters
----------
x1 : Expr
The first input tensor.
x2 : Expr
The second input tensor.
"""
return _ffi_api.floor_mod(x1, x2) # type: ignore
###################### Comparison operators ######################
def equal(x1: Expr, x2: Expr) -> Expr:
"""Broadcasted element-wise test for (lhs == rhs).
Parameters
----------
x1 : relax.Expr
The first input tensor.
x2 : relax.Expr
The second input tensor.
Returns
-------
result : relax.Expr
The computed result.
"""
return _ffi_api.equal(x1, x2) # type: ignore
def greater(x1: Expr, x2: Expr) -> Expr:
"""Broadcasted element-wise test for (lhs > rhs).
Parameters
----------
x1 : relax.Expr
The first input tensor.
x2 : relax.Expr
The second input tensor.
Returns
-------
result : relax.Expr
The computed result.
"""
return _ffi_api.greater(x1, x2) # type: ignore
def greater_equal(x1: Expr, x2: Expr) -> Expr:
"""Broadcasted element-wise test for (lhs >= rhs).
Parameters
----------
x1 : relax.Expr
The first input tensor.
x2 : relax.Expr
The second input tensor.
Returns
-------
result : relax.Expr
The computed result.
"""
return _ffi_api.greater_equal(x1, x2) # type: ignore
def less(x1: Expr, x2: Expr) -> Expr:
"""Broadcasted element-wise test for (lhs < rhs).
Parameters
----------
x1 : relax.Expr
The first input tensor.
x2 : relax.Expr
The second input tensor.
Returns
-------
result : relax.Expr
The computed result.
"""
return _ffi_api.less(x1, x2) # type: ignore
def less_equal(x1: Expr, x2: Expr) -> Expr:
"""Broadcasted element-wise test for (lhs <= rhs).
Parameters
----------
x1 : relax.Expr
The first input tensor.
x2 : relax.Expr
The second input tensor.
Returns
-------
result : relax.Expr
The computed result.
"""
return _ffi_api.less_equal(x1, x2) # type: ignore
def not_equal(x1: Expr, x2: Expr) -> Expr:
"""Broadcasted element-wise test for (lhs != rhs).
Parameters
----------
x1 : relax.Expr
The first input tensor.
x2 : relax.Expr
The second input tensor.
Returns
-------
result : relax.Expr
The computed result.
"""
return _ffi_api.not_equal(x1, x2) # type: ignore
def maximum(x1: Expr, x2: Expr) -> Expr:
"""Element-wise maximum
Parameters
----------
x1 : relax.Expr
The first input tensor.
x2 : relax.Expr
The second input tensor.
Returns
-------
result : relax.Expr
The computed result.
"""
return _ffi_api.maximum(x1, x2)
def minimum(x1: Expr, x2: Expr) -> Expr:
"""Element-wise minimum
Parameters
----------
x1 : relax.Expr
The first input tensor.
x2 : relax.Expr
The second input tensor.
Returns
-------
result : relax.Expr
The computed result.
"""
return _ffi_api.minimum(x1, x2)
###################### Logical operators ######################
def logical_and(x1: Expr, x2: Expr) -> Expr:
"""Logical AND
Parameters
----------
x1 : relax.Expr
The first input tensor.
x2 : relax.Expr
The second input tensor.
Returns
-------
result : relax.Expr
The computed result.
"""
return _ffi_api.logical_and(x1, x2)
def logical_or(x1: Expr, x2: Expr) -> Expr:
"""Logical OR
Parameters
----------
x1 : relax.Expr
The first input tensor.
x2 : relax.Expr
The second input tensor.
Returns
-------
result : relax.Expr
The computed result.
"""
return _ffi_api.logical_or(x1, x2)
def logical_xor(x1: Expr, x2: Expr) -> Expr:
"""Logical XOR
Parameters
----------
x1 : relax.Expr
The first input tensor.
x2 : relax.Expr
The second input tensor.
Returns
-------
result : relax.Expr
The computed result.
"""
return _ffi_api.logical_xor(x1, x2)
###################### Bitwise operators ######################
def bitwise_and(x1: Expr, x2: Expr) -> Expr:
"""Bitwise AND
Parameters
----------
x1 : relax.Expr
The first input tensor.
x2 : relax.Expr
The second input tensor.
Returns
-------
result : relax.Expr
The computed result.
"""
return _ffi_api.bitwise_and(x1, x2)
def bitwise_or(x1: Expr, x2: Expr) -> Expr:
"""Bitwise OR
Parameters
----------
x1 : relax.Expr
The first input tensor.
x2 : relax.Expr
The second input tensor.
Returns
-------
result : relax.Expr
The computed result.
"""
return _ffi_api.bitwise_or(x1, x2)
def bitwise_xor(x1: Expr, x2: Expr) -> Expr:
"""Bitwise XOR
Parameters
----------
x1 : relax.Expr
The first input tensor.
x2 : relax.Expr
The second input tensor.
Returns
-------
result : relax.Expr
The computed result.
"""
return _ffi_api.bitwise_xor(x1, x2)
def left_shift(x1: Expr, x2: Expr) -> Expr:
"""Bitwise Shift Left
Parameters
----------
x1 : relax.Expr
The input tensor to be shifted.
x2 : relax.Expr
The number of positions to shift.
Returns
-------
result : relax.Expr
The computed result.
"""
return _ffi_api.left_shift(x1, x2)
def right_shift(x1: Expr, x2: Expr) -> Expr:
"""Bitwise Shift Right
Parameters
----------
x1 : relax.Expr
The input tensor to be shifted.
x2 : relax.Expr
The number of positions to shift.
Returns
-------
result : relax.Expr
The computed result.
"""
return _ffi_api.right_shift(x1, x2)
python/tvm/relax/op/builtin/__init__.py
+19
View File
@@ -0,0 +1,19 @@
# 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.
"""Relax builtin operators."""
from .builtin import alloc_tensor, stop_lift_params
python/tvm/relax/op/builtin/_ffi_api.py
+20
View File
@@ -0,0 +1,20 @@
# 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
"""FFI APIs for tvm.relax.op.builtin"""
import tvm_ffi
tvm_ffi.init_ffi_api("relax.op.builtin", __name__)
python/tvm/relax/op/builtin/builtin.py
+85
View File
@@ -0,0 +1,85 @@
# 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
"""The builtin Relax operators."""
from tvm.ir import Call
from ...expr import DataTypeImm, Expr, StringImm, prim_value
from ...utils import convert_to_expr
from . import _ffi_api
def alloc_tensor(
shape: Expr,
dtype: str | Expr,
runtime_device_index: int | Expr,
storage_scope: str | Expr = "global",
) -> Call:
"""Construct a Call to allocate a tensor with specific shape, dtype, runtime_device_index.
Parameters
----------
shape : Expr
The shape of the tensor to be allocated.
dtype : Union[str, Expr]
The datatype of the tensor to be allocated.
runtime_device_index : Union[int, Expr]
The device index indicating on which device the tensor is to be allocated at runtime.
Index -1 is reserved for the host device.
storage_scope : Union[str, Expr]
The storage scope to allocate the storage to.
Returns
-------
result : Call
A relax Call, which gets the allocated tensor.
"""
if not isinstance(shape, Expr):
shape = convert_to_expr(shape)
if isinstance(dtype, str):
dtype = DataTypeImm(dtype)
if isinstance(runtime_device_index, int):
runtime_device_index = prim_value(runtime_device_index)
if isinstance(storage_scope, str):
storage_scope = StringImm(storage_scope)
if not isinstance(storage_scope, StringImm):
raise ValueError(
"relax.builtin.alloc_tensor expects string as the storage scope, "
f"but {storage_scope} is got."
)
return _ffi_api.alloc_tensor(shape, dtype, runtime_device_index, storage_scope) # type: ignore
def stop_lift_params(x: Expr) -> Expr:
"""
An indicator that the consumers of input tensor should not be
lifted to transform_params function
Parameters
----------
x: relax.Expr
The input data
Returns
-------
result : relax.Expr
The result tensor that is the same as input tensor
"""
return _ffi_api.stop_lift_params(x) # type: ignore
python/tvm/relax/op/ccl/__init__.py
+20
View File
@@ -0,0 +1,20 @@
# isort: skip_file
# 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.
"""CCL related operators."""
from .ccl import allgather, allreduce, broadcast_from_worker0, scatter_from_worker0
python/tvm/relax/op/ccl/_ffi_api.py
+21
View File
@@ -0,0 +1,21 @@
# 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.
"""Operators serving for Collective Communications Library (CCL) operators"""
import tvm_ffi
tvm_ffi.init_ffi_api("relax.op.ccl", __name__)
python/tvm/relax/op/ccl/ccl.py
+108
View File
@@ -0,0 +1,108 @@
# 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.
"""Relax Collective Communications Library (CCL) operators"""
from ...expr import Expr
from . import _ffi_api
def allreduce(x, op_type: str = "sum", in_group: bool = True): # pylint: disable=invalid-name
"""Allreduce operator
Parameters
----------
x : relax.Expr
The input tensor.
op_type : str
The type of reduction operation to be applied to the input data.
Now "sum", "prod", "min", "max" and "avg" are supported.
in_group : bool
Whether the reduction operation performs globally or in group as default.
Returns
-------
result : relax.Expr
The result of allreduce.
"""
supported_op_types = ["sum", "prod", "min", "max", "avg"]
assert op_type in supported_op_types, (
"Allreduce only supports limited reduction operations, "
f"including {supported_op_types}, but got {op_type}."
)
return _ffi_api.allreduce(x, op_type, in_group) # type: ignore # pylint: disable=no-member
def allgather(x, num_workers: int, in_group: bool = True): # pylint: disable=invalid-name
"""AllGather operator
Parameters
----------
x : relax.Expr
The input tensor.
num_worker : int
The number of workers to gather data from.
in_group : bool
Whether the gather operation performs globally or in group as default.
Returns
-------
result : relax.Expr
The result of allgather.
"""
return _ffi_api.allgather(x, num_workers, in_group) # type: ignore # pylint: disable=no-member
def broadcast_from_worker0(x: Expr) -> Expr:
"""Broadcast data from worker-0 to all other workers.
Parameters
----------
x : relax.Expr
The tensor to be broadcast.
Returns
-------
result : relax.Expr
The same tensor, which has been broadcast to all other workers.
"""
return _ffi_api.broadcast_from_worker0(x)
def scatter_from_worker0(x: Expr, num_workers: int, axis: int = 0) -> Expr:
"""Perform a scatter operation from worker-0, chunking the given buffer into equal parts.
Parameters
----------
x : relax.Expr
The buffer to be divided into equal parts and sent to each worker accordingly.
num_worker : int
The number of workers, i.e. the number of parts the given buffer should be chunked into.
axis : int
The dimension of the tensor to be scattered. Default is 0.
Returns
-------
result : relax.Expr
Chunked Tensor received by different workers.
"""
return _ffi_api.scatter_from_worker0(x, num_workers, axis)
python/tvm/relax/op/create.py
+378
View File
@@ -0,0 +1,378 @@
# 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.
"""Creation operators."""
from tvm import DataType, DataTypeCode
from tvm.ir import PrimType, is_prim_expr
from ..expr import Expr, ShapeExpr, prim_value
from . import _ffi_api
PrimExprLike = int | Expr
def _raw_dtype(dtype):
return dtype.dtype if isinstance(dtype, PrimType) else dtype
def _normalize_shape(shape):
if isinstance(shape, tuple | list):
return ShapeExpr(shape)
if not isinstance(shape, Expr) or is_prim_expr(shape):
raise TypeError("shape must be a tuple/list or a Relax shape expression")
return shape
def full(
shape: tuple[PrimExprLike] | Expr,
fill_value: Expr,
dtype: str | DataType | None = None,
) -> Expr:
"""Fill array with scalar value.
Parameters
----------
shape : Union[Tuple[PrimExprLike], Expr]
The shape of the created tensor.
fill_value : relax.Expr
The value to fill. Must be a scalar tensor.
dtype : Optional[str | DataType]
The data type of the created tensor.
If dtype is not given, it will by default use the dtype of fill_value.
Returns
-------
result : relax.Expr
The result tensor.
"""
shape = _normalize_shape(shape)
return _ffi_api.full(shape, fill_value, _raw_dtype(dtype)) # type: ignore
def full_like(x: Expr, fill_value: Expr, dtype: str | DataType | None = None) -> Expr:
"""Construct a tensor such that
- its shape is the same as the input data tensor's shape,
- its value is filled with the input scalar fill value.
Parameters
----------
x : relax.Expr
The input tensor, which provides the shape, and dtype
when the `dtype` field is not specified.
fill_value : relax.Expr
The value to fill. Must be a scalar tensor.
dtype : Optional[str | DataType]
The data type of the created tensor.
If dtype is not given, it will by default use the dtype of the input tensor.
Returns
-------
result : relax.Expr
The result tensor.
"""
return _ffi_api.full_like(x, fill_value, _raw_dtype(dtype)) # type: ignore
def ones(shape: tuple[PrimExprLike] | Expr, dtype: str | DataType) -> Expr:
"""Construct a tensor of all ones, with the input shape and dtype.
Parameters
----------
shape : Union[Tuple[PrimExprLike], Expr]
The shape of the created tensor.
dtype : str | DataType
The data type of the created tensor.
Returns
-------
result : relax.Expr
The result tensor.
"""
shape = _normalize_shape(shape)
return _ffi_api.ones(shape, _raw_dtype(dtype)) # type: ignore
def ones_like(x: Expr, dtype: str | DataType | None = None) -> Expr:
"""Construct a tensor with all ones, with shape of the input tensor shape.
Parameters
----------
x : relax.Expr
The input tensor, which provides the shape, and dtype
when the `dtype` field is not specified.
dtype : Optional[str | DataType]
The data type of the created tensor.
If dtype is not given, it will by default use the dtype of the input tensor.
Returns
-------
result : relax.Expr
The result tensor.
"""
return _ffi_api.ones_like(x, _raw_dtype(dtype)) # type: ignore
def zeros(shape: tuple[PrimExprLike] | Expr, dtype: str | DataType) -> Expr:
"""Construct a tensor of all zeros, with the input shape and dtype.
Parameters
----------
shape : Union[Tuple[PrimExprLike], Expr]
The shape of the created tensor.
dtype : str | DataType
The data type of the created tensor.
Returns
-------
result : relax.Expr
The result tensor.
"""
shape = _normalize_shape(shape)
return _ffi_api.zeros(shape, _raw_dtype(dtype)) # type: ignore
def zeros_like(x: Expr, dtype: str | DataType | None = None) -> Expr:
"""Construct a tensor with all zeros, with shape of the input tensor shape.
Parameters
----------
x : relax.Expr
The input tensor, which provides the shape, and dtype
when the `dtype` field is not specified.
dtype : Optional[str | DataType]
The data type of the created tensor.
If dtype is not given, it will by default use the dtype of the input tensor.
Returns
-------
result : relax.Expr
The result tensor.
"""
return _ffi_api.zeros_like(x, _raw_dtype(dtype)) # type: ignore
def eye(
n: PrimExprLike,
m: PrimExprLike | None = None,
k: PrimExprLike = 0,
dtype: str | DataType = "float32",
) -> Expr:
"""Construct a 2-D tensor with ones on the diagonal and zeros elsewhere.
Parameters
----------
n : PrimExprLike
Number of rows in the output.
m : Optional[PrimExprLike]
Number of columns in the output. If None, defaults to n.
k : PrimExprLike
Index of the diagonal: 0 (the default) refers to the main diagonal,
a positive value refers to an upper diagonal, and a negative value
to a lower diagonal.
dtype : str | DataType
The data type of the created tensor.
Returns
-------
result : relax.Expr
The result tensor.
"""
m = n if m is None else m
n = prim_value(n)
m = prim_value(m)
k = prim_value(k)
return _ffi_api.eye(n, m, k, _raw_dtype(dtype)) # type: ignore
def eye_like(
x: Expr,
k: PrimExprLike = 0,
dtype: str | DataType | None = None,
) -> Expr:
"""Return a 2-D tensor with ones on the diagonal and zeros elsewhere,
with the same shape as the input tensor.
Parameters
----------
x : relax.Expr
The input tensor, which provides the shape, and dtype
when the `dtype` field is not specified.
k : PrimExprLike
Index of the diagonal: 0 (the default) refers to the main diagonal,
a positive value refers to an upper diagonal, and a negative value
to a lower diagonal.
dtype : Optional[str | DataType]
The data type of the created tensor.
If dtype is not given, it will by default use the dtype of the input tensor.
Returns
-------
result : relax.Expr
The result tensor.
"""
k = prim_value(k)
return _ffi_api.eye_like(x, k, _raw_dtype(dtype)) # type: ignore
def arange(
start: PrimExprLike,
end: PrimExprLike | None = None,
step: PrimExprLike = 1,
dtype: str | DataType | None = None,
) -> Expr:
"""Construct a tensor with evenly spaced elements.
Parameters
----------
start : PrimExprLike
The start of the interval.
end : Optional[PrimExprLike]
The end of the interval. If not given, it will be set to start,
and start will be set to 0.
step : PrimExprLike
The step size.
dtype : Optional[str | DataType]
The data type of the created tensor.
Returns
-------
result : relax.Expr
The result tensor.
"""
if end is None:
end = start
start = 0
def is_int(expr):
if isinstance(expr, int):
return True
if is_prim_expr(expr):
return expr.ty.matches_code(DataTypeCode.INT)
return False
if dtype is None:
args = (start, end, step)
integer_args = all(is_int(arg) for arg in args)
dtype = "int64" if integer_args else "float32"
start = prim_value(start)
end = prim_value(end)
step = prim_value(step)
return _ffi_api.arange(start, end, step, dtype) # type: ignore
def hamming_window(window_size, periodic, alpha, beta, dtype):
"""Hamming window function.
Parameters
----------
window_size : Expr
The size of returned window.
periodic : Expr
If True, returns a window to be used as periodic function.
If False, return a symmetric window.
alpha : Expr
The co-efficient alpha.
beta : Expr
The co-efficient beta.
Returns
-------
ret : relax.Expr
The result tensor.
"""
if not is_prim_expr(window_size):
window_size = prim_value(window_size)
if not is_prim_expr(periodic):
periodic = prim_value(periodic)
if not is_prim_expr(alpha):
alpha = prim_value(alpha)
if not is_prim_expr(beta):
beta = prim_value(beta)
return _ffi_api.hamming_window(window_size, periodic, alpha, beta, dtype)
def tril(x: Expr, k: int | Expr = 0) -> Expr:
"""Return the lower triangular part of a matrix or a batch of matrices.
Parameters
----------
x : relax.Expr
The tensor that tril will be applied to.
It is required to have at least two dimensions.
k : int
The index indicating the diagonal above which to zero elements.
If k = 0, the diagonal is the main diagonal.
If k < 0, the diagonal is below the main diagonal.
If k > 0, the diagonal is above the main diagonal.
Returns
-------
ret : relax.Expr
The result tensor.
"""
if not is_prim_expr(k):
k = prim_value(k)
return _ffi_api.tril(x, k) # type: ignore
def triu(x: Expr, k: int | Expr = 0) -> Expr:
"""Return the upper triangular part of a matrix or a batch of matrices.
Parameters
----------
x : relax.Expr
The tensor that triu will be applied to.
It is required to have at least two dimensions.
k : int
The index indicating the diagonal below which to zero elements.
If k = 0, the diagonal is the main diagonal.
If k < 0, the diagonal is below the main diagonal.
If k > 0, the diagonal is above the main diagonal.
Returns
-------
ret : relax.Expr
The result tensor.
"""
if not is_prim_expr(k):
k = prim_value(k)
return _ffi_api.triu(x, k) # type: ignore
python/tvm/relax/op/datatype.py
+63
View File
@@ -0,0 +1,63 @@
# 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.
"""Datatype operators."""
from tvm import DataType
from tvm.ir import PrimType
from ..expr import Expr
from . import _ffi_api
def _raw_dtype(dtype):
return dtype.dtype if isinstance(dtype, PrimType) else dtype
def astype(x: Expr, dtype: str | DataType | PrimType) -> Expr:
"""Cast input tensor to the given data type.
Parameters
----------
x : relax.Expr
The input data to the operator.
dtype: Union[str, DataType]
The target data type
Returns
-------
result : relax.Expr
The casted result.
"""
return _ffi_api.astype(x, _raw_dtype(dtype)) # type: ignore
def wrap_param(data: Expr, dtype: str | DataType | PrimType = "float32") -> Expr:
"""Cast input tensor which is model param to data type if the dtype of the input data is not
the same as the given dtype.
Parameters
----------
data : relax.Expr
The input data to the operator.
dtype : Union[str, DataType]
The target data type
Returns
-------
result : relax.Expr
The casted result.
"""
return _ffi_api.wrap_param(data, _raw_dtype(dtype)) # type: ignore
python/tvm/relax/op/distributed/__init__.py
+25
View File
@@ -0,0 +1,25 @@
# isort: skip_file
# 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.
"""Operators serving for distributed Relax."""
from .distributed import (
annotate_sharding,
redistribute,
call_tir_local_view,
redistribute_replica_to_shard,
)
python/tvm/relax/op/distributed/_ffi_api.py
+21
View File
@@ -0,0 +1,21 @@
# 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.
"""FFI APIs for tvm.relax.op.distributed"""
import tvm_ffi
tvm_ffi.init_ffi_api("relax.op.dist", __name__)
python/tvm/relax/op/distributed/distributed.py
+137
View File
@@ -0,0 +1,137 @@
# 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=redefined-builtin
"""Operators for distributed Relax."""
from tvm.ir import Call
from tvm.relax.distributed import DeviceMesh, DTensorType, Placement
from ...expr import Expr, GlobalVar, ShapeExpr
from ...expr import Tuple as RxTuple
from ...utils import convert_to_expr
from . import _ffi_api
def annotate_sharding(input: Expr, device_mesh: DeviceMesh, placement: Placement) -> Expr:
"""Annotate sharding plan for tensor
Parameters
----------
input : relax.Expr
The input tensor.
device_mesh: DeviceMesh
The device mesh of the sharding plan
placement: Placement
The placement of the sharding plan
Returns
-------
result : relax.Expr
The tensor unmodified.
"""
return _ffi_api.annotate_sharding(input, device_mesh, placement) # type: ignore
def redistribute(input: Expr, device_mesh: DeviceMesh, placement: Placement) -> Expr:
"""Redistribute tensor
Parameters
----------
input : relax.Expr
The input tensor.
device_mesh: DeviceMesh
The device mesh after redistribution
placement: Placement
The placement after redistribution
Returns
-------
result : relax.Expr
The tensor after redistribution.
"""
return _ffi_api.redistribute(input, device_mesh, placement) # type: ignore
def call_tir_local_view(
gvar: GlobalVar,
args: Expr,
out_ty: DTensorType | list[DTensorType],
tir_vars: ShapeExpr | tuple[Expr] | list[Expr] | None = None,
) -> Call:
"""
Call a tirx.prim_func and return the output. The prim_func should be a worker-local function
that is actually executed on each worker, instead of the unpartitioned function.
The output of this operator is DTensor or a tuple of DTensors.
Parameters
----------
gvar : GlobalVar
The GlobalVar referring to a tirx PrimFunc.
args : Expr
The input arguments.
out_ty : Union[DTensorType, List[DTensorType]]
The type information of the call_tir output.
It should be a single or a list of DTensorType. Each one denotes the
type information of a returned tensor.
tir_vars : Optional[Union[ShapeExpr, Tuple[Expr], List[Expr]]]
ShapeExpr representing a tuple of integers to unpack when calling func. Is null if not used
Returns
-------
ret: Call
A call node for the call_tir_local_view operator.
"""
if isinstance(args, tuple | list):
args = RxTuple([convert_to_expr(a) for a in args])
elif isinstance(args, Expr) and not isinstance(args, RxTuple): # type: ignore
args = RxTuple((args,))
if not isinstance(out_ty, list):
out_ty = [out_ty]
if isinstance(tir_vars, list | tuple):
tir_vars = ShapeExpr(tir_vars)
return _ffi_api.call_tir_local_view(gvar, args, out_ty, tir_vars) # type: ignore
def redistribute_replica_to_shard(input: Expr, num_workers: int, axis: int) -> Expr:
"""Slice tensor into several parts along one axis,
and each worker takes one part.
input.ty.shape[axis] % num_workers == 0 is required.
Each worker must have an identical copy of the input.
This is a specialized version of redistribute op.
Parameters
----------
input : relax.Expr
The buffer to be sliced into equal parts.
num_worker : int
The number of workers, i.e. the number of parts the given buffer should be sliced into.
axis : int
The axis of the tensor to be sliced.
Returns
-------
result : relax.Expr
Sliced Tensor kept by each device.
"""
return _ffi_api.redistribute_replica_to_shard(input, num_workers, axis)
python/tvm/relax/op/grad/__init__.py
+19
View File
@@ -0,0 +1,19 @@
# 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.
"""Operators serving for finding gradient of relax operators."""
from .grad import *
python/tvm/relax/op/grad/_ffi_api.py
+21
View File
@@ -0,0 +1,21 @@
# 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.
"""FFI APIs for tvm.relax.op.grad"""
import tvm_ffi
tvm_ffi.init_ffi_api("relax.op.grad", __name__)
python/tvm/relax/op/grad/grad.py
+216
View File
@@ -0,0 +1,216 @@
# 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=redefined-builtin
"""Operators to implement operaor gradients. Used in `_op_gradient.py`.
We are trying to keep grad operators as simple as possible, and hope they are only used for finding
gradients for forward operators. The ty inference for grad operators just returns the
ty of the input.
"""
from ...expr import Expr
from . import _ffi_api
def no_grad(input: Expr) -> Expr:
"""No gradient dummy operator w.r.t. the input.
Parameters
----------
input : relax.Expr
The corresponding input tensor.
Returns
-------
result : relax.Expr
The no-gradient representation w.r.t. input.
"""
return _ffi_api.no_grad(input) # type: ignore
def start_checkpoint(input: Expr) -> Expr:
"""Mark the start of the checkpoint stage. The computation between start_checkpoint and
end_checkpoint will be marked as the checkpoint stage.
Rather than storing all intermediate activations of the entire computation graph for
computing backward, the checkpointed stage does not save intermediate activations, and instead
recomputes them in backward process.
For instance,
```
a = relax.Var("a", relax.TensorType((2, 2), "float32"))
b = relax.Var("b", relax.TensorType((2, 2), "float32"))
c = a * 2
d = b * 2
c_cp = start_checkpoint(c)
d_cp = start_checkpoint(d)
e = c_cp + d_cp
e_out = end_checkpoint(e)
```
Then `e` will be recomputed in the backward stage.
See tvm.relax.transform.Gradient, tvm.relax.testing.nn.checkpoint,
tvm.relax.op.grad.end_checkpoint for more information.
Parameters
----------
input : relax.Expr
The tensor marking the input of the checkpoint stage.
Returns
-------
result : relax.Expr
The same tensor as the input.
"""
return _ffi_api.start_checkpoint(input) # type: ignore
def end_checkpoint(input: Expr) -> Expr:
"""Mark the end of checkpoint stage. See tvm.relax.op.grad.start_checkpoint.
Parameters
----------
input : relax.Expr
The output of the checkpoint stage.
Returns
-------
result : relax.Expr
The same tensor as the input.
"""
return _ffi_api.end_checkpoint(input) # type: ignore
def nll_loss_backward(
output_grad: Expr,
predictions: Expr,
targets: Expr,
weights: Expr | None = None,
reduction: str = "mean",
ignore_index: int = -100,
) -> Expr:
"""Backward operator of relax.nn.nll_loss. All parameters except output_grad is the same as
relax.nn.nll_loss. Returns the gradient w.r.t. predictions.
Parameters
----------
output_grad : relax.Expr
The gradient w.r.t. the result of nll_loss.
Returns
-------
result : relax.Expr
The gradient w.r.t. predictions.
"""
return _ffi_api.nll_loss_backward( # type: ignore
output_grad, predictions, targets, weights, reduction, ignore_index
)
def max_pool2d_backward(
output_grad: Expr,
data: Expr,
pool_size: tuple[int, int] = (1, 1),
strides: tuple[int, int] = (1, 1),
padding: tuple[int, int, int, int] = (0, 0, 0, 0),
dilation: tuple[int, int] = (1, 1),
ceil_mode: bool = False,
count_include_pad: bool = False,
layout: str = "NCHW",
out_layout: str | None = None,
) -> Expr:
"""Backward operator of relax.nn.max_pool2d. All parameters except output_grad is the same as
relax.nn.max_pool2d. Returns the gradient w.r.t. data.
Parameters
----------
output_grad : relax.Expr
The gradient w.r.t. the result of max_pool2d.
Returns
-------
result : relax.Expr
The gradient w.r.t. data.
"""
return _ffi_api.max_pool2d_backward( # type: ignore
output_grad,
data,
pool_size,
strides,
padding,
dilation,
ceil_mode,
count_include_pad,
layout,
out_layout,
)
def avg_pool2d_backward(
output_grad: Expr,
data: Expr,
pool_size: tuple[int, int] = (1, 1),
strides: tuple[int, int] = (1, 1),
padding: tuple[int, int, int, int] = (0, 0, 0, 0),
dilation: tuple[int, int] = (1, 1),
ceil_mode: bool = False,
count_include_pad: bool = False,
layout: str = "NCHW",
out_layout: str | None = None,
) -> Expr:
"""Backward operator of relax.nn.avg_pool2d. All parameters except output_grad is the same as
relax.nn.avg_pool2d. Returns the gradient w.r.t. data.
Parameters
----------
output_grad : relax.Expr
The gradient w.r.t. the result of avg_pool2d.
Returns
-------
result : relax.Expr
The gradient w.r.t. data.
"""
return _ffi_api.avg_pool2d_backward( # type: ignore
output_grad,
data,
pool_size,
strides,
padding,
dilation,
ceil_mode,
count_include_pad,
layout,
out_layout,
)
def take_backward(output_grad: Expr, x: Expr, indices: Expr, axis: int | None = None) -> Expr:
"""Backward operator of relax.take. All parameters except output_grad is the same as
relax.take. Returns the gradient w.r.t. x.
Parameters
----------
output_grad : relax.Expr
The gradient w.r.t. the result of take.
Returns
-------
result : relax.Expr
The gradient w.r.t. x.
"""
return _ffi_api.take_backward(output_grad, x, indices, axis) # type: ignore
python/tvm/relax/op/image/__init__.py
+20
View File
@@ -0,0 +1,20 @@
# isort: skip_file
# 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.
"""Image operators."""
from .image import affine_grid, grid_sample, resize2d, resize3d
python/tvm/relax/op/image/_ffi_api.py
+21
View File
@@ -0,0 +1,21 @@
# 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.
"""Constructor APIs"""
import tvm_ffi
tvm_ffi.init_ffi_api("relax.op.image", __name__)
python/tvm/relax/op/image/image.py
+274
View File
@@ -0,0 +1,274 @@
# 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.
"""Image operators."""
from typing import cast
from tvm import DataType
from tvm.ir import is_prim_expr
from ...expr import Expr, ShapeExpr
from . import _ffi_api
PrimExprLike = int | Expr
SizeLike = PrimExprLike | tuple[PrimExprLike, ...]
def resize2d(
data: Expr,
size: SizeLike,
roi: float | tuple[float] | None = None,
layout: str = "NCHW",
method: str = "linear",
coordinate_transformation_mode: str = "half_pixel",
rounding_method: str = "round",
cubic_alpha: float = -0.75,
cubic_exclude: int = 0,
extrapolation_value: float = 0.0,
out_dtype: str | DataType | None = None,
) -> Expr:
"""Image resize2d operator.
This operator takes data as input and does 2D scaling to the given scale factor.
In the default case, where the data_layout is `NCHW`
with data of shape (n, c, h, w)
out will have a shape (n, c, size[0], size[1])
method indicates the algorithm to be used while calculating the out value
and method can be one of ("linear", "nearest_neighbor", "cubic")
Parameters
----------
data : relax.Expr
The input data to the operator.
size: SizeLike
The out size to which the image will be resized.
If specified as a list, it is required to have length either 1 or 2.
If specified as an Expr, it is required to have ndim 2.
roi: Optional[Union[float, Tuple[float]]]
The region of interest for cropping the input image. Expected to be of
size 4, and format [start_h, start_w, end_h, end_w].
Only used if coordinate_transformation_mode is tf_crop_and_resize.
layout : str
Layout of the input.
method : str
Scale method to used [nearest_neighbor, linear, cubic].
coordinate_transformation_mode : str
Describes how to transform the coordinate in the resized tensor
to the coordinate in the original tensor. Definitions can be found
in topi/image/resize.py.
[half_pixel, align_corners, asymmetric, pytorch_half_pixel,
tf_half_pixel_for_nn, and tf_crop_and_resize].
rounding_method: str
indicates how to find the "nearest" pixel in nearest_neighbor method
[round, floor, ceil]
cubic_alpha: float
Spline Coefficient for bicubic interpolation
cubic_exclude: int
Flag to exclude exterior of the image during bicubic interpolation
extrapolation_value: float
Fill value to use when roi is outside of the image
out_dtype : Optional[str | DataType]
The dtype of the output tensor.
It it is not specified, the output will have the same dtype as input if not specified.
Returns
-------
result: relax.Expr
The resized result.
"""
if roi is None:
roi = (0.0, 0.0, 0.0, 0.0) # type: ignore
elif isinstance(roi, float):
roi = (roi, roi, roi, roi) # type: ignore
elif isinstance(roi, tuple | list):
roi = tuple(val if isinstance(val, float) else float(val) for val in roi)
else:
raise NotImplementedError(f"Unsupported roi type {type(roi)}")
if isinstance(size, int) or is_prim_expr(size):
size = (size, size)
if isinstance(size, tuple | list):
if len(size) == 1:
size = ShapeExpr([size[0], size[0]])
else:
size = ShapeExpr(size)
return _ffi_api.resize2d( # type: ignore
data,
size,
roi,
layout,
method,
coordinate_transformation_mode,
rounding_method,
cubic_alpha,
cubic_exclude,
extrapolation_value,
out_dtype,
)
def resize3d(
data: Expr,
size: SizeLike,
roi: float | tuple[float] | None = None,
layout: str = "NCDHW",
method: str = "linear",
coordinate_transformation_mode: str = "half_pixel",
rounding_method: str = "",
cubic_alpha: float = -0.75,
cubic_exclude: int = 0,
extrapolation_value: float = 0.0,
out_dtype: str | DataType | None = None,
) -> Expr:
"""Image resize3d operator.
This operator takes data as input and does 3D scaling to the given output size.
In the default case, where data layout is `NCDHW`
with data of shape (n, c, d, h, w),
the output has shape (n, c, size[0], size[1], size[2]).
"""
if roi is None:
roi = (0.0, 0.0, 0.0, 0.0, 0.0, 0.0) # type: ignore
elif isinstance(roi, float):
roi = (roi, roi, roi, roi, roi, roi) # type: ignore
elif isinstance(roi, tuple | list):
roi = tuple(val if isinstance(val, float) else float(val) for val in roi)
else:
raise NotImplementedError(f"Unsupported roi type {type(roi)}")
if isinstance(size, int) or is_prim_expr(size):
size = (size, size, size)
if isinstance(size, tuple | list):
if len(size) == 1:
size = ShapeExpr([size[0], size[0], size[0]])
else:
size = ShapeExpr(size)
return _ffi_api.resize3d( # type: ignore
data,
size,
roi,
layout,
method,
coordinate_transformation_mode,
rounding_method,
cubic_alpha,
cubic_exclude,
extrapolation_value,
out_dtype,
)
def grid_sample(
data: Expr,
grid: Expr,
method: str = "bilinear",
layout: str = "NCHW",
padding_mode: str = "zeros",
align_corners: bool = False,
) -> Expr:
"""Applies grid sampling to input feature map.
Given data and grid, the output is computed by sampling from data using
the grid coordinates.
Parameters
----------
data : relax.Expr
The input data tensor with shape [N, C, H, W] for NCHW layout.
grid : relax.Expr
The grid tensor with shape [N, H_out, W_out, 2]. The values are normalized
to [-1, 1], where (-1, -1) is the top-left corner and (1, 1) is the bottom-right.
method : str
Interpolation method. Can be 'nearest', 'bilinear', or 'bicubic'.
layout : str
Layout of the input data. Default is 'NCHW'.
padding_mode : str
Padding mode for outside grid values. Can be 'zeros', 'border', or 'reflection'.
align_corners : bool
If True, the corner pixels of the input and output tensors are aligned.
Returns
-------
result : relax.Expr
The sampled output tensor with shape [N, C, H_out, W_out].
"""
return _ffi_api.grid_sample( # type: ignore
data,
grid,
method,
layout,
padding_mode,
align_corners,
)
def affine_grid(
data: Expr,
size: SizeLike,
align_corners: bool = True,
) -> Expr:
"""Generate a 2D or 3D sampling grid using an affine transformation matrix.
This operation is described in https://arxiv.org/pdf/1506.02025.pdf.
It generates a uniform sampling grid within the target shape, normalizes it
to [-1, 1], and applies the provided affine transformation.
Parameters
----------
data : relax.Expr
The input affine matrix tensor with shape [batch, 2, 3] for 2D or
[batch, 3, 4] for 3D.
size : SizeLike
The target output spatial shape, (H, W) for 2D or (D, H, W) for 3D. If a
single integer or PrimExpr is provided, it is interpreted as a square 2D
output shape (size, size).
align_corners : bool
If True, normalized grid coordinates map to corner pixels; if False, to
pixel centers (the PyTorch / ONNX default).
Returns
-------
result : relax.Expr
The output grid tensor with shape [batch, 2, H, W] for 2D or
[batch, 3, D, H, W] for 3D.
"""
if isinstance(size, int) or is_prim_expr(size):
size = (size, size)
if isinstance(size, tuple | list):
size = ShapeExpr(size)
return cast(Expr, _ffi_api.affine_grid(data, size, align_corners))
python/tvm/relax/op/index.py
+143
View File
@@ -0,0 +1,143 @@
# 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.
"""Indexing operators."""
from ..expr import Expr
from ..utils import convert_to_expr
from . import _ffi_api
PrimExprLike = int | Expr
def take(x: Expr, indices: Expr, axis: int | None = None, mode: str = "fast") -> Expr:
"""Take elements from a tensor along an axis.
Its semantic is mostly similar to `numpy.take`
(https://numpy.org/doc/stable/reference/generated/numpy.take.html),
which can cover `torch.take` (https://pytorch.org/docs/stable/generated/torch.take.html) and
`onnx.gather` (https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Gather-13).
Parameters
----------
x : relax.Expr
The source tensor.
indices : relax.Expr
The indices of the values to extract.
axis : Optional[int]
The axis over which to select values.
If it is none, the input tensor is required to be one-dimensional.
mode : str
Specifies how out-of-bounds indices will behave.
- fast (default): extra indices lead to seg fault (user must make sure indices are in-bound)
- nan: produce NaNs for out-of-bounds indices
- wrap: wrap around the indices
- clip: clip to the range
Returns
-------
ret : relax.Expr
The taken result.
"""
return _ffi_api.take(x, indices, axis, mode) # type: ignore
def strided_slice(
x: Expr,
axes: Expr,
begin: Expr,
end: Expr,
strides: Expr | None = None,
assume_inbound: bool = False,
) -> Expr:
"""Strided slice of a tensor.
Parameters
----------
x : relax.Expr
The source tensor to be sliced.
axes : List[int]
Axes along which slicing is applied.
begin : List[PrimExprLike]
The indices to begin with in the slicing, inclusive.
end : List[PrimExprLike]
The indices indicating end of the slice, exclusive.
strides : Optional[List[PrimExprLike]]
Specifies the stride values, it can be negative in that case,
the input tensor will be reversed in that particular axis.
If not specified, it by default is an list of ones of the same length as `axes`.
assume_inbound : bool
Whether to assume the indices are in bound. If it is set to false,
out of bound indices will be clipped to the bound.
Returns
-------
ret : relax.Expr
The sliced result.
Note
----
strided_slice require the input `begin`, `end` and `strides` to have the
same length as `axes`.
"""
axes = convert_to_expr(axes)
begin = convert_to_expr(begin)
end = convert_to_expr(end)
if strides is not None:
strides = convert_to_expr(strides)
return _ffi_api.strided_slice(x, axes, begin, end, strides, assume_inbound) # type: ignore
def dynamic_strided_slice(
x: Expr,
begin: Expr,
end: Expr,
strides: Expr,
) -> Expr:
"""Dynamic strided slice of a tensor. `begin`, `end`, `strides` can be computed at runtime.
Parameters
----------
x : Expr
The source tensor to be sliced.
begin : Expr
The indices to begin with in the slicing, inclusive.
end : Expr
The indices indicating end of the slice, exclusive.
strides : Expr
Specifies the stride values, it can be negative in that case,
the input tensor will be reversed in that particular axis.
If not specified, it by default is an list of ones of the same length as `axes`.
Returns
-------
ret : relax.Expr
The sliced result.
Note
----
dyn_strided_slice require the input `begin`, `end` and `strides` to have the
same length as rank of `data` tensor.
"""
return _ffi_api.dynamic_strided_slice(x, begin, end, strides) # type: ignore
python/tvm/relax/op/linear_algebra.py
+139
View File
@@ -0,0 +1,139 @@
# 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
"""Relax linear algebra operators"""
from tvm import DataType
from ..expr import Expr
from ..expr import Tuple as RxTuple
from . import _ffi_api
from .manipulate import permute_dims
def matmul(x1: Expr, x2: Expr, out_dtype: str | DataType | None = None) -> Expr:
"""General matrix multiplication of two tensors, with broadcasting on batched dimensions.
The semantics and output shape deduction rule is specified as
https://data-apis.org/array-api/latest/API_specification/generated/array_api.matmul.html.
Parameters
----------
x1 : relax.Expr
The first input tensor.
x2 : relax.Expr
The second input tensor.
out_dtype: Optional[Union[str, DataType]]
The data type of the matmul result.
When it is not specified, the output dtype will be the same as input dtype.
Returns
-------
result : relax.Expr
The computed result.
"""
return _ffi_api.matmul(x1, x2, out_dtype) # type: ignore
def linear(
data: Expr,
weight: Expr,
bias: Expr | None = None,
out_dtype: str | DataType | None = None,
) -> Expr:
"""Applies a linear transformation to the incoming data: y = xA^T + b
Parameters
----------
data : relax.Expr
The input data.
weight : relax.Expr
The weight tensor.
bias : Optional[Expr]
The bias tensor.
out_dtype: Optional[Union[str, DataType]]
The data type of the matmul result.
When it is not specified, the output dtype will be the same as input dtype.
Notes
-----
Relax does not regard the Linear Op as a primitive Op,
while combine the transpose, matmul and add op to implement it.
Returns
-------
result : relax.Expr
The computed result.
"""
# Since weight can be 1D or 2D, we use `axes=None` to support both cases.
x = matmul(data, permute_dims(weight, axes=None), out_dtype=out_dtype)
return x + bias if bias is not None else x
def einsum(operands, subscripts):
"""Evaluates the Einstein summation convention on data
Parameters
----------
operands : Union(List[relax.Expr], Tuple[relax.Expr])
A list of expression.
subscripts : str
The einsum expression string.
Returns
-------
result : relax.Expr
The output from the einsum op.
"""
if isinstance(operands, list | tuple):
operands = RxTuple(operands)
return _ffi_api.einsum(operands, subscripts) # type: ignore
def outer(x1: Expr, x2: Expr) -> Expr:
"""
Computes the outer product of two input expressions.
Parameters
----------
x1 : relax.Expr
The first input expression.
x2 : relax.Expr
The second input expression.
Notes
-----
This operation computes the outer product between two expressions,
resulting in a tensor where each element is the product of elements
from `x1` and `x2`. It is commonly used in tensor and matrix operations
to expand lower-dimensional inputs into higher-dimensional representations.
Returns
-------
result : relax.Expr
The resulting expression representing the outer product.
"""
return _ffi_api.outer(x1, x2)
python/tvm/relax/op/manipulate.py
+906
View File
@@ -0,0 +1,906 @@
# 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.
"""Manipulation operators."""
from collections.abc import Callable
from tvm.ir import is_prim_expr
from tvm.runtime import DataTypeCode
from tvm.tirx import FloatImm, IndexMap, IntImm
from ..expr import Expr, ShapeExpr, prim_value
from ..expr import Tuple as RxTuple
from . import _ffi_api
PrimExprLike = int | Expr
def broadcast_to(x: Expr, shape: tuple[PrimExprLike] | Expr) -> Expr:
"""Broadcasts a tensor to a specified shape.
Parameters
----------
x : relax.Expr
The input data to the operator.
shape : Union[Tuple[PrimExprLike], Expr]
The target shape.
Returns
-------
result : relax.Expr
The broadcasted tensor.
"""
if isinstance(shape, tuple | list):
shape = ShapeExpr(shape)
return _ffi_api.broadcast_to(x, shape) # type: ignore
def concat(tensors: Expr | list[Expr], axis: int | None = 0) -> Expr:
"""Concatenate the input tensors along the given axis.
Parameters
----------
tensors : Union[relax.Expr, List[relax.Expr]]
An Expr in Tuple type, containing the tensors to be concatenated,
or a list of Tensors.
axis : Optional[int]
The axis along which the tensors are concatenated.
If `axis` is `None`, the input tensor is required to be flattened before concatenation.
Returns
-------
result: relax.Expr
The concatenated tensor.
"""
if isinstance(tensors, list | tuple):
tensors = RxTuple(tensors)
return _ffi_api.concat(tensors, axis) # type: ignore
def expand_dims(x: Expr, axis: int | list[int]) -> Expr:
"""Insert new axes at the positions given by `axis`.
Parameters
----------
x : relax.Expr
The input data to the operator.
axis : Union[int, List[int]]
The axes at which the input array are expanded.
All values are required to lie in range `[-data.ndim - 1, data.ndim]`, with the convention
of negative indexing.
Returns
-------
result : relax.Expr
The transformed result.
"""
if isinstance(axis, int):
axis = [axis]
return _ffi_api.expand_dims(x, axis) # type: ignore
def flatten(x: Expr) -> Expr:
"""Flatten all the tensor dimensions into one.
Parameters
----------
x : relax.Expr
The input data to the operator.
Returns
-------
result : relax.Expr
The flattened result.
"""
return _ffi_api.flatten(x) # type: ignore
def layout_transform(
x: Expr,
index_map: Callable | IndexMap,
pad_value: int | float | Expr | None = None,
axis_separators: int | str | None = None, # str for IndexMap.AXIS_SEPARATOR
input_axis_separators: int | str | None = None, # str for IndexMap.AXIS_SEPARATOR
):
"""Modifies the layout of a tensor.
Parameters
----------
x : relax.Expr
The input tensor to the operator.
index_map : Callable | IndexMap
The transformation to apply.
pad_value : Optional[int | float | Expr]
The value used for padding if the transformation results in implicit padding.
If not specified, any value can be used.
axis_separators : Optional[int | IndexMap.AXIS_SEPARATOR]
The axis_separators for index_map to create non flat buffers.
Returns
-------
result : relax.Expr
The transformed tensor.
"""
default_index_dtype = "int64"
if callable(index_map):
index_map = IndexMap.from_func(index_map, index_dtype=default_index_dtype)
x_dtype = x.ty.dtype
# Explicitly convert python int/float pad_value to the x's type. If the default behavior
# is applied, it would be converted to int32/float32, which may not match the x's type.
if pad_value is None:
pass
elif not is_prim_expr(pad_value):
if x_dtype.matches_code(DataTypeCode.INT, DataTypeCode.UINT) and isinstance(pad_value, int):
pad_value = IntImm(x_dtype.dtype, pad_value)
elif x_dtype.matches_code(DataTypeCode.FLOAT, DataTypeCode.BFLOAT) and (
isinstance(pad_value, int | float)
):
pad_value = FloatImm(x_dtype.dtype, float(pad_value))
pad_value = prim_value(pad_value)
if axis_separators is None:
axis_separators = []
if input_axis_separators is None:
input_axis_separators = []
return _ffi_api.layout_transform(
x, index_map, pad_value, axis_separators, input_axis_separators
)
def permute_dims(x: Expr, axes: list[int] | None = None) -> Expr:
"""Permutes the dimensions of an array.
Parameters
----------
x : relax.Expr
The input data to the operator.
axes : Optional[List[int]]
The target axes order. If not specified, permute_dims will reverse the order of all axes.
Returns
-------
result : relax.Expr
The transposed result.
"""
return _ffi_api.permute_dims(x, axes) # type: ignore
def reshape(x: Expr, shape: tuple[PrimExprLike] | Expr) -> Expr:
"""Reshape the input array.
``-1`` infers the dimension of the output shape by using the remainder of
the input dimensions keeping the size of the new array same as that of the input array.
At most one dimension of shape can be -1.
.. code-block:: python
x.shape = (2, 3, 4), shape = (6, 1, -1), result.shape = (6, 1, 4)
x.shape = (2, 3, 4), shape = (3, -1, 8), result.shape = (3, 1, 8)
x.shape = (2, 3, 4), shape = (-1,), result.shape = (24,)
Parameters
----------
x : relax.Expr
The input data to the operator.
shape : Union[Tuple[PrimExprLike], Expr]
The new shape. Should be compatible with the original shape.
Returns
-------
result : relax.Expr
The reshaped result.
Note
----
The ``-1`` inference is only performed at compile-time.
That is to say, in any case the dimension length of ``-1`` cannot be inferred in
compile-time, an error will be thrown.
"""
if not isinstance(shape, tuple | list | Expr) or is_prim_expr(shape):
raise TypeError("shape must be a tuple/list or a Relax shape expression")
return _ffi_api.reshape(x, shape) # type: ignore
def split(
x: Expr,
indices_or_sections: int | list[PrimExprLike],
axis: int = 0,
) -> Expr:
"""Split input tensor along axis by sections or indices.
If indices_or_sections is an integer, the input will be divided equally
along given axis (if possible). Last section will be smaller if the tensor
size along the given dimension is not divisible by the integer.
If indices_or_sections is a tuple of mixture of int or Expr,
the entries indicate the indices where along axis the array is split.
Parameters
----------
x : relax.Expr
The tensor to be split.
indices_or_sections : Union[int, List[PrimExprLike]]
Indices or sections to split into. Accepts an int or a list.
axis : int
The axis over which to split.
Returns
-------
ret : relax.Expr
The computed result.
"""
if isinstance(indices_or_sections, int):
indices_or_sections = IntImm("int64", indices_or_sections)
return _ffi_api.split(x, indices_or_sections, axis) # type: ignore
def squeeze(x: Expr, axis: int | list[int] | None = None) -> Expr:
"""Squeeze axes in the array.
Parameters
----------
x : relax.Expr
The input data to the operator.
axis : Optional[Union[int, List[int]]
The set of axes to remove.
If axis = None, remove all axis of dimensions 1.
If any specified axis has dimension that does not equal 1, it is an error.
Returns
-------
result : relax.Expr
The squeezed result.
"""
if isinstance(axis, int):
axis = [axis]
return _ffi_api.squeeze(x, axis) # type: ignore
def stack(tensors: Expr | list[Expr], axis: int = 0) -> Expr:
"""Stack the input tensors along a new axis.
Parameters
----------
tensors : Union[relax.Expr, List[relax.Expr]]
An Expr in Tuple type, containing the tensors to be stacked,
or a list of Tensors. All input tensors must have the same shape.
axis : int
The axis in the resulting tensor along which the input tensors will be stacked.
Negative values wrap around. Default is 0.
Returns
-------
result: relax.Expr
The stacked tensor with an additional dimension compared to the input tensors.
"""
if isinstance(tensors, list | tuple):
tensors = RxTuple(tensors)
return _ffi_api.stack(tensors, axis) # type: ignore
def collapse_sum_like(data: Expr, collapse_target: Expr) -> Expr:
"""Return a summation of data to the shape of collapse_target.
For details, please see relax.op.collapse_sum_to.
Parameters
----------
data : relax.Expr
The input tensor.
collapse_target : relax.Expr
The tensor whose shape is the shape to collapse to.
Returns
-------
result : relax.Expr
The result tensor after summation.
"""
return _ffi_api.collapse_sum_like(data, collapse_target) # type: ignore
def collapse_sum_to(data: Expr, shape: tuple[PrimExprLike] | Expr) -> Expr:
"""Return a summation of data to the given shape.
collapse_sum_to is intended as the backward operator of tvm.relax.op.broadcast_to and
other broadcast operators in the automatic differentiation process.
We expect that data is the result of broadcasting some tensor of the given shape in some
broadcast operation. Thus the given `shape` and `data.shape` must follow broadcast rules.
During computation, all axes of `data.shape` and `shape` are checked from right to left.
For an axis, if it follows these rules, `data` will be summed over this axis:
- the axis exists in `data.shape` but not in `shape`, or
- the axis exists in `data.shape` and equals to 1 in `shape`.
Parameters
----------
data : relax.Expr
The input tensor.
shape : Union[Tuple[PrimExprLike], relax.Expr]
The shape to collapse to.
Returns
-------
result : relax.Expr
The result tensor of the given shape after summation.
"""
if isinstance(shape, tuple | list):
shape = ShapeExpr(shape)
return _ffi_api.collapse_sum_to(data, shape) # type: ignore
def repeat(data: Expr, repeats: int, axis: int | None = None) -> Expr:
"""Repeats elements of an array.
Parameters
----------
data : relax.Expr
The input tensor.
repeats : int
The number of repetitions.
axis: Optional[int]
The axis along which to repeat values. The negative numbers are interpreted
counting from the backward. By default, use the flattened input array, and
return a flat output array.
Returns
-------
ret : relax.Expr
The computed result.
Examples
--------
.. code-block:: python
x = R.const([[1, 2], [3, 4]])
lv1 = R.repeat(x, repeats=2) # lv1 == [1, 1, 2, 2, 3, 3, 4, 4]
lv2 = R.repeat(x, repeats=2, axis=1) # lv2 == [[1., 1., 2., 2.],
# [3., 3., 4., 4.]]
"""
return _ffi_api.repeat(data, repeats, axis) # type: ignore
def tile(data: Expr, repeats: int | tuple[int] | list[int]) -> Expr:
"""Construct an array by repeating data the number of times given by repeats.
If repeats has length l, and data has dimension d, the result will have dimension of max(l, d).
If d < l, data is promoted to be l-dimensional by prepending new axes. So a shape (3,) Tensor is
promoted to (1, 3) for 2-D replication, or shape (1, 1, 3) for 3-D replication. If this is not
the desired behavior, promote data to d-dimensions manually before calling this function.
If d > l, reps is promoted to length d by pre-pending 1's to it. Thus for a data of shape
(2, 3, 4, 5), a reps of (2, 2) is treated as (1, 1, 2, 2).
Parameters
----------
data : relax.Expr
The input data to the operator.
repeats : Union[int, Tuple[int], List[int]]
The number of repetitions of data along each axis.
Returns
-------
ret : relax.Expr
The computed result.
Examples
--------
.. code-block:: python
x = R.const([[1, 2], [3, 4]])
lv1 = R.tile(x, reps=(2, 3)) # lv1 = [[1., 2., 1., 2., 1., 2.],
# [3., 4., 3., 4., 3., 4.],
# [1., 2., 1., 2., 1., 2.],
# [3., 4., 3., 4., 3., 4.]]
lv2 = R.tile(x, reps=2) # lv2 = [[1., 2., 1., 2.],
# [3., 4., 3., 4.]]
"""
if isinstance(repeats, int):
repeats = [repeats]
return _ffi_api.tile(data, repeats) # type: ignore
def flip(data, axis):
"""Reverses the order of elements along given axis while preserving array shape.
Parameters
----------
data : relax.Expr
The input data to the operator.
axis: int
The axis along which to flip over.
Returns
-------
ret : relax.Expr
The computed result.
Examples
--------
.. code-block:: python
x = [[1., 2.], [3., 4.]]
relax.flip(x, axis=0) = [[3., 4.], [1., 2.]]
relax.flip(x, axis=1) = [[2., 1.], [4., 3.]]
"""
return _ffi_api.flip(data, axis) # type: ignore
def reverse_sequence(data: Expr, seq_lengths: Expr, seq_axis: int = 1, batch_axis: int = 0) -> Expr:
"""Reverses variable length slices.
Parameters
----------
data : relax.Expr
The input tensor.
seq_lengths : relax.Expr
A 1-D tensor containing sequence lengths for each batch.
seq_axis : int
The axis along which to reverse variable length slices.
batch_axis : int
The axis that indexes the batch.
Returns
-------
ret : relax.Expr
The computed result.
"""
return _ffi_api.reverse_sequence(data, seq_lengths, seq_axis, batch_axis) # type: ignore
def gather_elements(data: Expr, indices: Expr, axis: int = 0) -> Expr:
"""Gather elements from data according to indices along the specified axis.
Parameters
----------
data : relax.Expr
The input data to the operator.
indices : relax.Expr
The indices tensor, must have integer type.
axis : int
The axis along which to index. Default is 0.
Returns
-------
ret : relax.Expr
The computed result.
Examples
--------
.. code-block:: python
data = [[1, 2], [3, 4]]
indices = [[0, 0], [1, 0]]
axis = 1
output = [[1, 1], [4, 3]]
data = [[1, 2, 3], [4, 5, 6]]
indices = [[1, 1, 1]]
axis = 0
output = [[4, 5, 6]]
"""
return _ffi_api.gather_elements(data, indices, axis) # type: ignore
def gather_nd(data: Expr, indices: Expr, batch_dims: int = 0) -> Expr:
"""Update data at positions defined by indices with values in updates.
Parameters
----------
data : relax.Expr
The input data to the operator.
indices : relax.Expr
The indices tensor, must have integer type.
batch_dims : int
The number of batch dimensions. Default is 0.
Returns
-------
ret : relax.Expr
The computed result.
Examples
--------
.. code-block:: python
batch_dims = 0
data = [[0,1],[2,3]] # data_shape = [2, 2]
indices = [[0,0],[1,1]] # indices_shape = [2, 2]
output = [0,3] # output_shape = [2]
batch_dims = 1
data = [[[0,1],[2,3]],[[4,5],[6,7]]] # data_shape = [2, 2, 2]
indices = [[1],[0]] # indices_shape = [2, 1]
output = [[2,3],[4,5]] # output_shape = [2, 2]
"""
return _ffi_api.gather_nd(data, indices, batch_dims) # type: ignore
def index_tensor(data: Expr, indices: Expr | list[Expr]) -> Expr:
"""Advanced-tensor indexing (NumPy/PyTorch-style).
Given k index tensors ``indices = (I0, I1, …, Ik-1)`` this
operator selects elements from ``data`` as if one had written
``data[I0, I1, …, Ik-1]`` in NumPy/PyTorch:
All index tensors must have an integer dtype.
Their shapes are broadcast together to a common shape ``B`` in
the usual NumPy way.
The result shape is ``B + data.shape[k:]`` (i.e. the broadcast
shape followed by the remaining axes of ``data`` that are *not*
indexed).
At compile-time Relax checks that the number of index tensors
``k`` does not exceed ``data.ndim``, that the dtypes are integer,
and that the shapes are consitent (broadcast-compatible).
Parameters
----------
data : relax.Expr
The input tensor to be indexed.
indices : Union[relax.Expr, List[relax.Expr]]
A Tuple expression containing the index tensors,
or a Python ``list`` / ``tuple`` that will be promoted to a
tuple expression automatically. Each tensor must have an
integer dtype.
Returns
-------
result : relax.Expr
The tensor obtained after advanced indexing. Its dtype equals
``data.dtype``
Examples
--------
.. code-block:: python
import numpy as np
import tvm.relax as R
x = R.const(np.arange(9).reshape(3, 3).astype("float32"))
row = R.const(np.array([0, 2])) # shape (2,)
col = R.const(np.array([1, 0])) # shape (2,)
y = R.index_tensor(x, [row, col])
# y.shape == (2,) ; y == [1., 6.]
# Broadcasting: row : (2,1), col : (1,3) → B = (2,3)
row = R.const(np.array([[0],[1]]))
col = R.const(np.array([[0,1,2]]))
z = R.index_tensor(x, [row, col])
# z.shape == (2,3)
"""
if isinstance(indices, list | tuple):
indices = RxTuple(indices)
return _ffi_api.index_tensor(data, indices) # type: ignore
def index_put(
data: Expr,
indices: Expr | tuple[Expr],
values: Expr,
accumulate: bool = False,
) -> Expr:
"""This operation updates values in `data` at positions
specified by `indices` with corresponding values from `values`. The `indices` is a tuple
of tensors where each tensor corresponds to a dimension in `data`.
When `accumulate` is True, the operation performs accumulation (addition) rather than
replacement. The `reduction` parameter allows specifying different reduction operations.
Parameters
----------
data : relax.Expr
The input tensor to be modified
indices : Union[Expr, Tuple[Expr]]
Tuple of index tensors (one for each dimension) specifying positions to update
values : relax.Expr
Values to place at the specified indices
accumulate : bool
Whether to accumulate (add) values rather than replace (default: False)
Returns
-------
result : relax.Expr
A new tensor with the same shape as data but with specified positions updated
Examples
--------
.. code-block:: python
# inputs
data = torch.zeros(3, 3)
indices = (torch.tensor([0, 2]), torch.tensor([1, 1]))
values = torch.tensor([1.0, 2.0])
# output
output = [
[0.0, 1.0, 0.0],
[0.0, 0.0, 0.0],
[0.0, 2.0, 0.0],
]
# with accumulate=True
output = [
[0.0, 1.0, 0.0],
[0.0, 0.0, 0.0],
[0.0, 3.0, 0.0],
]
"""
if isinstance(indices, list | tuple):
indices = RxTuple(indices)
return _ffi_api.index_put(data, indices, values, accumulate) # type: ignore
def meshgrid(tensors: Expr | list[Expr], indexing: str | None = "ij") -> Expr:
"""Generate coordinate grids from input tensors.
Parameters
----------
tensors : Union[relax.Expr, List[relax.Expr]]
An Expr in Tuple type, containing 1D tensors (or scalars promoted to 1D)
to generate coordinate grids from, or a list of such tensors.
indexing : Optional[str]
The indexing mode, either "ij" (matrix indexing) or "xy" (Cartesian indexing).
Defaults to "ij".
Returns
-------
result : relax.Expr
A Tuple of tensors representing the coordinate grids.
"""
if isinstance(tensors, list | tuple):
tensors = RxTuple(tensors)
return _ffi_api.meshgrid(tensors, indexing)
def scatter_elements(
data: Expr, indices: Expr, updates: Expr, axis: int = 0, reduction: str = "update"
):
"""ONNX style scatter elements. This operation updates its value in `data` to values
specified by `updates` at specific index positions specified by `indices`.
For example, in 2D tensor, the update corresponding to the [i][j] entry is performed
as below:
.. code-block::
output[indices[i][j]][j] = updates[i][j] if axis = 0
output[i][indices[i][j]] = updates[i][j] if axis = 1
When the `reduction` is set to some reduction function `f`, the update corresponding to
[i][j] entry is performed as below:
.. code-block::
output[indices[i][j]][j] += f(output[indices[i][j]][j], updates[i][j]) if axis = 0
output[i][indices[i][j]] += f(output[i][indices[i][j]], updates[i][j]) if axis = 1
Where `f` is update, add, mul, mean, max, min.
Parameters
----------
data : relax.Expr
The input data to the operator.
indices: relax.Expr
The index positions to update in `data`.
updates: relax.Expr
Values to replace to.
axis: int
Axis to scatter on.
reduction: str
Type of reduction to apply: update, add, mul, mean, max, min.
It is "update" by default.
Returns
-------
result : relax.Expr
The result has the same size as data, and the same shape as data
Examples
--------
.. code-block:: python
# inputs
data = [
[0.0, 0.0, 0.0],
[0.0, 0.0, 0.0],
[0.0, 0.0, 0.0],
]
indices = [
[1, 0, 2],
[0, 2, 1],
]
updates = [
[1.0, 1.1, 1.2],
[2.0, 2.1, 2.2],
]
axis = 0
reduction = "update"
# output P
output = [
[2.0, 1.1, 0.0]
[1.0, 0.0, 2.2]
[0.0, 2.1, 1.2]
]
"""
return _ffi_api.scatter_elements(data, indices, updates, axis, reduction) # type: ignore
def scatter_nd(data: Expr, indices: Expr, updates: Expr, reduction: str = "update") -> Expr:
"""Scatter updates into an array according to indices.
Parameters
----------
data: relax.Expr
The input data to be updated.
indices: relax.Expr
The index positions to update in `data`.
updates: relax.Expr
Values to replace to.
reduction: str
Type of reduction to apply: update, add, mul, max, min.
It is "update" by default.
Returns
-------
result : relax.Expr
The result has the same shape as data.
Examples
--------
.. code-block:: python
# inputs
data = [1, 2, 3, 4, 5, 6, 7, 8]
indices = [[4], [3], [1], [7]]
updates = [9, 10, 11, 12]
# output
output = [1, 11, 3, 10, 9, 6, 7, 12]
"""
return _ffi_api.scatter_nd(data, indices, updates, reduction) # type: ignore
def slice_scatter(input_tensor: Expr, src: Expr, start, end, step, axis=0):
"""Embeds the values of the src tensor into input at the given dimension.
Parameters
----------
input_tensor: relax.Expr
The input tensor to be updated.
src: relax.Expr
The tensor to embed into input.
axis: int
The dimension to insert the slice into.
start:
The start index of where to insert the slice.
end:
The end index of where to insert the slice.
step:
The how many elements to skip in.
Returns
-------
result : relax.Expr
The computed result tensor with the same shape as `data`.
"""
if not is_prim_expr(start):
start = prim_value(start)
if not is_prim_expr(end):
end = prim_value(end)
if not is_prim_expr(step):
step = prim_value(step)
return _ffi_api.slice_scatter(input_tensor, src, axis, start, end, step)
def one_hot(
indices: Expr,
on_value: int | float | Expr,
off_value: int | float | Expr,
depth: int,
axis: int = -1,
) -> Expr:
"""Returns a one-hot tensor.
Parameters
----------
indices : relax.Expr
The indices to set to `on_value`.
on_value : int | float | Expr
The value to fill at `indices`.
off_value : int | float | Expr
The value to fill at other locations.
depth : int
The depth of the one-hot dimension.
axis : int, optional
The axis to fill. Default is -1 which adds a new dimension at the end.
Returns
-------
result : relax.Expr
The computed result.
Examples
--------
.. code-block:: python
indices = [0, 1, 2]
depth = 3
on_value = 1
off_value = 0
one_hot(indices, on_value, off_value, depth) =
[[1, 0, 0],
[0, 1, 0],
[0, 0, 1]]
"""
on_value = prim_value(on_value)
off_value = prim_value(off_value)
return _ffi_api.one_hot(indices, on_value, off_value, depth, axis) # type: ignore
python/tvm/relax/op/mask.py
+39
View File
@@ -0,0 +1,39 @@
# 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.
"""Operators with mask."""
from ..expr import Expr
from . import _ffi_api
def masked_fill(x: Expr, mask: Expr, value: Expr):
"""Fill a tensor by a specified value in places defined by a mask.
Parameters
----------
x : relax.Expr
The input data to the operator.
mask : relax.Expr
The mask.
value : relax.Expr
The value to set in the input tensor.
Returns
-------
result : relax.Expr
The filled tensor.
"""
values = _ffi_api.full_like(x, value, value.ty.dtype.dtype) # type: ignore
return _ffi_api.where(mask, values, x) # type: ignore
python/tvm/relax/op/memory/__init__.py
+21
View File
@@ -0,0 +1,21 @@
# isort: skip_file
# 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.
"""Relax memory primitives."""
from .memory import alloc_storage, alloc_tensor, kill_storage, kill_tensor
from .view import view, ensure_zero_offset
python/tvm/relax/op/memory/_ffi_api.py
+20
View File
@@ -0,0 +1,20 @@
# 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
"""FFI APIs for tvm.relax.op.memory"""
import tvm_ffi
tvm_ffi.init_ffi_api("relax.op.memory", __name__)
python/tvm/relax/op/memory/memory.py
+135
View File
@@ -0,0 +1,135 @@
# 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
"""Relax memory primitives."""
from tvm.ir import Call
from ...expr import DataTypeImm, Expr, StringImm, prim_value
from ...utils import convert_to_expr
from . import _ffi_api
def alloc_storage(
size: Expr,
virtual_device_index: int | Expr,
storage_scope: str | Expr,
dtype: str | Expr,
) -> Call:
"""Construct a Call to allocate a storage with specific size, virtual_device_index,
storage_scope and dtype.
Parameters
----------
size : Expr
The size of the storage to be allocated.
virtual_device_index : Union[int, Expr]
The virtual device index indicating on which device the storage is to be allocated.
Index -1 is reserved for the host device.
storage_scope : Union[str, Expr]
The storage scope to allocate the storage to.
dtype : Union[str, Expr]
The datatype of the storage to be allocated.
Returns
-------
result : Call
A relax Call, which gets the allocated storage.
"""
size = convert_to_expr(size)
if isinstance(dtype, str):
dtype = DataTypeImm(dtype)
if isinstance(storage_scope, str):
storage_scope = StringImm(storage_scope)
if isinstance(virtual_device_index, int):
virtual_device_index = prim_value(virtual_device_index)
return _ffi_api.alloc_storage(size, virtual_device_index, storage_scope, dtype) # type: ignore
def alloc_tensor(
storage: Expr,
offset: int | Expr,
shape: Expr,
dtype: str | Expr,
runtime_device_ind: int | Expr = prim_value(0),
) -> Call:
"""Construct a Call to allocate a tensor on a certain storage starting from the given offset.
Parameters
----------
storage : Expr
The storage to allocate the tensor to.
offset : Union[int, Expr]
The storage offset to allocate the tensor.
shape : Expr
The shape of the tensor to be allocated.
dtype : Union[str, Expr]
The datatype of the tensor to be allocated.
runtime_device_ind: Union[int, Expr]
The device index indicating on which device the tensor is to be
allocated at runtime. Index -1 is reserved for the host device.
Returns
-------
result : Call
A relax Call, which gets the allocated tensor.
"""
if isinstance(offset, int):
offset = prim_value(offset)
shape = convert_to_expr(shape)
if isinstance(dtype, str):
dtype = DataTypeImm(dtype)
if isinstance(runtime_device_ind, int):
runtime_device_ind = prim_value(runtime_device_ind)
return _ffi_api.alloc_tensor(storage, offset, shape, dtype, runtime_device_ind) # type: ignore
def kill_storage(storage: Expr) -> Call:
"""Construct a Call to kill a storage.
Parameters
----------
storage : Expr
The storage to be killed.
Returns
-------
result : Call
A relax Call to kill a storage.
"""
return _ffi_api.kill_storage(storage) # type: ignore
def kill_tensor(tensor: Expr) -> Call:
"""Construct a Call to kill a tensor.
Parameters
----------
tensor : Expr
The tensor to be killed.
Returns
-------
result : Call
A relax Call to kill a tensor.
"""
return _ffi_api.kill_tensor(tensor) # type: ignore
python/tvm/relax/op/memory/view.py
+116
View File
@@ -0,0 +1,116 @@
# 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.
"""Operations that act on the DLTensor container
While most operations require inspecting the values stored within the
allocated buffers, some operations only require updating the fields in
a `DLTensor`, without touching the values that are stored within it.
For example, given an array of shape `[16,16]`, the slice at
`[0:8,0:16]` can be generated by changing the `DLTensor::shape` field,
while keeping the same underlying data.
"""
from collections.abc import Sequence
from tvm.relax import DataTypeImm, Expr, ShapeExpr
from tvm.relax.expr import prim_value
from ..base import null_value
from . import _ffi_api
PrimExprLike = int | Expr
def view(
data: Expr,
shape: Sequence[PrimExprLike] | Expr | None = None,
dtype: Expr | None = None,
relative_byte_offset: Expr | None = None,
) -> Expr:
"""Provide a view into an existing tensor
The view may have a different shape, may be a different datatype,
and may start at an offset relative to the source array.
Regardless of which combination of these options are used, the
view may never access memory that was not accessible through the
input `data` array. This restriction applies even if the `data`
array is itself a view into a shared backing array.
Parameters
----------
data : relax.Expr
The input data to the operator.
shape : Optional[Union[Sequence[PrimExprLike], Expr]]
The target shape. Should be a `relax.ShapeExpr`, or a
collection that can be converted to a `relax.ShapeExpr`.
dtype : Optional[Expr]
The target datatype. Should be a `relax.ShapeExpr`, or a
collection that can be converted to a `relax.ShapeExpr`.
relative_byte_offset: Optional[Expr]
The offset of the output Tensor, relative to the byte offset
of `data`. If `None`, the offset of the view is the same as
the offset of `data`.
Returns
-------
result : relax.Expr
The tensor view
"""
def _normalize(expr, relax_cls):
if expr is None or isinstance(expr, Expr):
return expr
else:
return relax_cls(expr)
shape = _normalize(shape, ShapeExpr)
dtype = null_value() if dtype is None else _normalize(dtype, DataTypeImm)
relative_byte_offset = (
relative_byte_offset
if relative_byte_offset is None or isinstance(relative_byte_offset, Expr)
else prim_value(relative_byte_offset)
)
return _ffi_api.view(data, shape, dtype, relative_byte_offset) # type: ignore
def ensure_zero_offset(data: Expr) -> Expr:
"""
Ensure the tensor has elem_offset == 0. A copy will be made if necessary.
Parameters
----------
data : relax.Expr
The input tensor
Results
-------
result : relax.Expr
The tensor with elem_offset == 0
"""
return _ffi_api.ensure_zero_offset(data) # type: ignore
python/tvm/relax/op/nn/__init__.py
+61
View File
@@ -0,0 +1,61 @@
# isort: skip_file
# 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.
"""Neural network related operators."""
from .nn import (
adaptive_avg_pool1d,
adaptive_avg_pool2d,
adaptive_avg_pool3d,
attention,
attention_bias,
attention_var_len,
avg_pool1d,
avg_pool2d,
avg_pool3d,
batch_flatten,
batch_norm,
conv1d,
conv1d_transpose,
conv2d,
conv2d_transpose,
conv3d,
conv3d_transpose,
cross_entropy_with_logits,
dropout,
gelu,
gelu_tanh,
group_norm,
instance_norm,
layer_norm,
leakyrelu,
log_softmax,
max_pool1d,
max_pool2d,
max_pool3d,
nll_loss,
pad,
pixel_shuffle,
prelu,
relu,
relu6,
rms_norm,
selu,
silu,
softmax,
softplus,
)
python/tvm/relax/op/nn/_ffi_api.py
+21
View File
@@ -0,0 +1,21 @@
# 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.
"""Constructor APIs"""
import tvm_ffi
tvm_ffi.init_ffi_api("relax.op.nn", __name__)
python/tvm/relax/op/nn/nn.py
+2372
View File
File diff suppressed because it is too large Load Diff
python/tvm/relax/op/op_attrs.py
+401
View File
@@ -0,0 +1,401 @@
# 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.
"""The attributes node used for Relax operators"""
import tvm_ffi
from tvm.ir import Attrs
@tvm_ffi.register_object("relax.attrs.CallTIRWithGradAttrs")
class CallTIRWithGradAttrs(Attrs):
"""Attributes used in call_tir_with_grad operator"""
@tvm_ffi.register_object("relax.attrs.InitAttrs")
class InitAttrs(Attrs):
"""Attributes used in full/full_like, ones/ones_like, and zeros/zeros_like operator"""
@tvm_ffi.register_object("relax.attrs.TriluAttrs")
class TriluAttrs(Attrs):
"""Attributes used in tril and triu operator"""
@tvm_ffi.register_object("relax.attrs.AstypeAttrs")
class AstypeAttrs(Attrs):
"""Attributes used in astype operator"""
@tvm_ffi.register_object("relax.attrs.TakeAttrs")
class TakeAttrs(Attrs):
"""Attributes used in take operator"""
@tvm_ffi.register_object("relax.attrs.StridedSliceAttrs")
class StridedSliceAttrs(Attrs):
"""Attributes used in strided_slice operator"""
@tvm_ffi.register_object("relax.attrs.MatmulAttrs")
class MatmulAttrs(Attrs):
"""Attributes for matmul operator"""
@tvm_ffi.register_object("relax.attrs.Conv2DAttrs")
class Conv2DAttrs(Attrs):
"""Attributes for nn.conv2d"""
@tvm_ffi.register_object("relax.attrs.Conv3DAttrs")
class Conv3DAttrs(Attrs):
"""Attributes for nn.conv3d"""
@tvm_ffi.register_object("relax.attrs.Conv2DTransposeAttrs")
class Conv2DTransposeAttrs(Attrs):
"""Attributes for nn.conv2d_transpose"""
@tvm_ffi.register_object("relax.attrs.Conv3DTransposeAttrs")
class Conv3DTransposeAttrs(Attrs):
"""Attributes for nn.conv3d_transpose"""
@tvm_ffi.register_object("relax.attrs.Pool2DAttrs")
class Pool2DAttrs(Attrs):
"""Attributes for nn.max_pool2d"""
@tvm_ffi.register_object("relax.attrs.AdaptivePool2DAttrs")
class AdaptivePool2DAttrs(Attrs):
"""Attributes for 2d adaptive pool operator"""
@tvm_ffi.register_object("relax.attrs.SoftmaxAttrs")
class SoftmaxAttrs(Attrs):
"""Attributes for nn.softmax"""
@tvm_ffi.register_object("relax.attrs.BatchNormAttrs")
class BatchNormAttrs(Attrs):
"""Attributes used in batch_norm operator"""
@tvm_ffi.register_object("relax.attrs.LayerNormAttrs")
class LayerNormAttrs(Attrs):
"""Attributes used in layer_norm operator"""
@tvm_ffi.register_object("relax.attrs.InstanceNormAttrs")
class InstanceNormAttrs(Attrs):
"""Attributes used in instance_norm operator"""
@tvm_ffi.register_object("relax.attrs.DropoutAttrs")
class DropoutAttrs(Attrs):
"""Attributes for dropout operator"""
@tvm_ffi.register_object("relax.attrs.StatisticalAttrs")
class StatisticalAttrs(Attrs):
"""Attributes used in statistical operator"""
@tvm_ffi.register_object("relax.attrs.ConcatAttrs")
class ConcatAttrs(Attrs):
"""Attributes for concat operator"""
@tvm_ffi.register_object("relax.attrs.ExpandDimsAttrs")
class ExpandDimsAttrs(Attrs):
"""Attributes for expand_dims operator"""
@tvm_ffi.register_object("relax.attrs.PermuteDimsAttrs")
class PermuteDimsAttrs(Attrs):
"""Attributes for permute_dims operator"""
@tvm_ffi.register_object("relax.attrs.SortAttrs")
class SortAttrs(Attrs):
"""Attributes for sort operator"""
@tvm_ffi.register_object("relax.attrs.ArgsortAttrs")
class ArgsortAttrs(Attrs):
"""Attributes for argsort operator"""
@tvm_ffi.register_object("relax.attrs.SplitAttrs")
class SplitAttrs(Attrs):
"""Attributes used in split operator"""
@tvm_ffi.register_object("relax.attrs.SqueezeAttrs")
class SqueezeAttrs(Attrs):
"""Attributes for squeeze operator"""
@tvm_ffi.register_object("relax.attrs.StackAttrs")
class StackAttrs(Attrs):
"""Attributes for concat operator"""
@tvm_ffi.register_object("relax.attrs.IndexPutAttrs")
class IndexPutAttrs(Attrs):
"""Attributes for index_put operator"""
@tvm_ffi.register_object("relax.attrs.LayoutTransformAttrs")
class LayoutTransformAttrs(Attrs):
"""Attributes used in layout_transform operator"""
@tvm_ffi.register_object("relax.attrs.Resize2DAttrs")
class Resize2DAttrs(Attrs):
"""Attributes used in image resize2d operator"""
@tvm_ffi.register_object("relax.attrs.ArgmaxArgminAttrs")
class ArgmaxArgminAttrs(Attrs):
"""Attributes for argmax/argmin operator"""
@tvm_ffi.register_object("relax.attrs.RepeatAttrs")
class RepeatAttrs(Attrs):
"""Attributes for repeat operator"""
@tvm_ffi.register_object("relax.attrs.TileAttrs")
class TileAttrs(Attrs):
"""Attributes for tile operator"""
@tvm_ffi.register_object("relax.attrs.ScanopAttrs")
class ScanopAttrs(Attrs):
"""Attributes for scan operators"""
@tvm_ffi.register_object("relax.attrs.TopKAttrs")
class TopKAttrs(Attrs):
"""Attributes for topk operators"""
@tvm_ffi.register_object("relax.attrs.EinsumAttrs")
class EinsumAttrs(Attrs):
"""Attributes for einsum operator"""
@tvm_ffi.register_object("relax.attrs.FlipAttrs")
class FlipAttrs(Attrs):
"""Attributes for flip operator"""
@tvm_ffi.register_object("relax.attrs.ReverseSequenceAttrs")
class ReverseSequenceAttrs(Attrs):
"""Attributes for reverse_sequence operator"""
@tvm_ffi.register_object("relax.attrs.PadAttrs")
class PadAttrs(Attrs):
"""Attributes used in pad operator"""
@tvm_ffi.register_object("relax.attrs.MultinomialFromUniformAttrs")
class MultinomialFromUniformAttrs(Attrs):
"""Attributes for multinomial_from_uniform operator"""
@tvm_ffi.register_object("relax.attrs.CallInplacePackedAttrs")
class CallInplacePackedAttrs(Attrs):
"""Attributes used in call_inplace_packed operator"""
@tvm_ffi.register_object("relax.attrs.CallTIRInplaceAttrs")
class CallTIRInplaceAttrs(Attrs):
"""Attributes used in call_tir_inplace operator"""
@tvm_ffi.register_object("relax.attrs.ToVDeviceAttrs")
class ToVDeviceAttrs(Attrs):
"""Attributes used in to_vdevice operator"""
@tvm_ffi.register_object("relax.attrs.HintOnDeviceAttrs")
class HintOnDeviceAttrs(Attrs):
"""Attributes used in hint_on_device operator"""
@tvm_ffi.register_object("relax.attrs.ScatterCollectiveAttrs")
class ScatterCollectiveAttrs(Attrs):
"""Attributes used in scatter collective operators"""
@tvm_ffi.register_object("relax.attrs.AttentionAttrs")
class AttentionAttrs(Attrs):
"""Attributes used in attention operator"""
@tvm_ffi.register_object("relax.attrs.AllClassNonMaximumSuppressionAttrs")
class AllClassNonMaximumSuppressionAttrs(Attrs):
"""Attributes for vision.all_class_non_max_suppression"""
@tvm_ffi.register_object("relax.attrs.GetValidCountsAttrs")
class GetValidCountsAttrs(Attrs):
"""Attributes for vision.get_valid_counts"""
@tvm_ffi.register_object("relax.attrs.NonMaximumSuppressionAttrs")
class NonMaximumSuppressionAttrs(Attrs):
"""Attributes for vision.non_max_suppression"""
@tvm_ffi.register_object("relax.attrs.ROIAlignAttrs")
class ROIAlignAttrs(Attrs):
"""Attributes for vision.roi_align"""
@tvm_ffi.register_object("relax.attrs.ROIPoolAttrs")
class ROIPoolAttrs(Attrs):
"""Attributes for vision.roi_pool"""
@tvm_ffi.register_object("relax.attrs.MultiboxTransformLocAttrs")
class MultiboxTransformLocAttrs(Attrs):
"""Attributes for vision.multibox_transform_loc"""
@tvm_ffi.register_object("relax.attrs.Conv1DAttrs")
class Conv1DAttrs(Attrs):
"""Attributes for nn.conv1d"""
@tvm_ffi.register_object("relax.attrs.Conv1DTransposeAttrs")
class Conv1DTransposeAttrs(Attrs):
"""Attributes for nn.conv1d_transpose"""
@tvm_ffi.register_object("relax.attrs.Pool1DAttrs")
class Pool1DAttrs(Attrs):
"""Attributes for nn.max_pool1d and nn.avg_pool1d"""
@tvm_ffi.register_object("relax.attrs.Pool3DAttrs")
class Pool3DAttrs(Attrs):
"""Attributes for nn.max_pool3d and nn.avg_pool3d"""
@tvm_ffi.register_object("relax.attrs.AdaptivePool1DAttrs")
class AdaptivePool1DAttrs(Attrs):
"""Attributes for 1d adaptive pool operator"""
@tvm_ffi.register_object("relax.attrs.AdaptivePool3DAttrs")
class AdaptivePool3DAttrs(Attrs):
"""Attributes for 3d adaptive pool operator"""
@tvm_ffi.register_object("relax.attrs.LeakyReluAttrs")
class LeakyReluAttrs(Attrs):
"""Attributes used in leaky_relu operator"""
@tvm_ffi.register_object("relax.attrs.SoftplusAttrs")
class SoftplusAttrs(Attrs):
"""Attributes used in softplus operator"""
@tvm_ffi.register_object("relax.attrs.PReluAttrs")
class PReluAttrs(Attrs):
"""Attributes used in prelu operator"""
@tvm_ffi.register_object("relax.attrs.PixelShuffleAttrs")
class PixelShuffleAttrs(Attrs):
"""Attributes used in pixel_shuffle operator"""
@tvm_ffi.register_object("relax.attrs.GroupNormAttrs")
class GroupNormAttrs(Attrs):
"""Attributes used in group_norm operator"""
@tvm_ffi.register_object("relax.attrs.RMSNormAttrs")
class RMSNormAttrs(Attrs):
"""Attributes used in rms_norm operator"""
@tvm_ffi.register_object("relax.attrs.NLLLossAttrs")
class NLLLossAttrs(Attrs):
"""Attributes used in nll_loss operator"""
@tvm_ffi.register_object("relax.attrs.AllReduceAttrs")
class AllReduceAttrs(Attrs):
"""Attributes used in allreduce operator"""
@tvm_ffi.register_object("relax.attrs.AllGatherAttrs")
class AllGatherAttrs(Attrs):
"""Attributes used in allgather operator"""
@tvm_ffi.register_object("relax.attrs.WrapParamAttrs")
class WrapParamAttrs(Attrs):
"""Attributes used in wrap_param operator"""
@tvm_ffi.register_object("relax.attrs.QuantizeAttrs")
class QuantizeAttrs(Attrs):
"""Attributes used in quantize/dequantize operators"""
@tvm_ffi.register_object("relax.attrs.GatherElementsAttrs")
class GatherElementsAttrs(Attrs):
"""Attributes for gather_elements operator"""
@tvm_ffi.register_object("relax.attrs.GatherNDAttrs")
class GatherNDAttrs(Attrs):
"""Attributes for gather_nd operator"""
@tvm_ffi.register_object("relax.attrs.MeshgridAttrs")
class MeshgridAttrs(Attrs):
"""Attributes for meshgrid operator"""
@tvm_ffi.register_object("relax.attrs.ScatterElementsAttrs")
class ScatterElementsAttrs(Attrs):
"""Attributes for scatter_elements operator"""
@tvm_ffi.register_object("relax.attrs.ScatterNDAttrs")
class ScatterNDAttrs(Attrs):
"""Attributes for scatter_nd operator"""
@tvm_ffi.register_object("relax.attrs.SliceScatterAttrs")
class SliceScatterAttrs(Attrs):
"""Attributes for slice_scatter operator"""
@tvm_ffi.register_object("relax.attrs.OneHotAttrs")
class OneHotAttrs(Attrs):
"""Attributes for one_hot operator"""
python/tvm/relax/op/qdq.py
+88
View File
@@ -0,0 +1,88 @@
# 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.
"""Relax quantize/dequantize operators"""
from ..expr import Expr
from . import _ffi_api
def quantize(data: Expr, scale: Expr, zero_point: Expr, axis: int = -1, out_dtype: str = "int8"):
r"""Quantize op
This operator takes input and produces quantized output. The input tensor can be of any shape.
The output shape is the same as input shape.
Q_output = clamp((round(input_tensor/scale) + zero_point), out_dtype::min, out_dtype::max)
Parameters
----------
data : tvm.relax.Expr
The input tensor to be quantized.
scale : tvm.relax.Expr
The output scale.
zero_point : tvm.relax.Expr
The output zero_point.
axis : int
The channel axis for quantization. Default value is -1 which corresponds to the last axis.
out_dtype : str, optional
The data type of the output tensor.
Returns
-------
result : tvm.relax.Expr
The computed result.
"""
return _ffi_api.quantize(data, scale, zero_point, axis, out_dtype)
def dequantize(
data: Expr, scale: Expr, zero_point: Expr, axis: int = -1, out_dtype: str = "float32"
):
r"""Dequantize op
This operator takes input and produces dequantized output. The input tensor can be of any shape.
The output shape is the same as input shape.
output = clamp(scale * (input_tensor - zero_point), out_dtype::min, out_dtype::max)
Parameters
----------
data : tvm.relax.Expr
The input tensor to be dequantized.
scale : tvm.relax.Expr
The input scale.
zero_point : tvm.relax.Expr
The input zero_point.
axis : int
The channel axis for dequantization. Default value is -1 which corresponds to the last axis.
out_dtype : str, optional
The data type of the output tensor.
Returns
-------
result : tvm.relax.Expr
The computed result.
"""
return _ffi_api.dequantize(data, scale, zero_point, axis, out_dtype)
python/tvm/relax/op/sampling.py
+85
View File
@@ -0,0 +1,85 @@
# 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.
"""Sampling operators."""
from ..expr import Expr
from . import _ffi_api
def multinomial_from_uniform(
prob: Expr,
uniform_sample: Expr,
sample_indices: Expr,
dtype: str = "int64",
) -> Expr:
"""Returns a tensor where each row contains the index sampled from the multinomial
probability distribution located in the corresponding row of tensor prob.
Notes
-----
For better cpu performance, use 'vm.builtin.multinomial_from_uniform'.
For accurate results, ensure probabilities are between 0 and 1 and sum to 1.
Parameters
----------
prob : relax.Expr
A 2-D tensor of shape (batch, vocab_size) representing probability distributions.
Each row is a distribution across vocabulary for a batch, where:
Values range from [0, 1], indicating the probability of each vocabulary item.
The sum of values in each row is 1, forming a valid distribution.
uniform_sample : relax.Expr
The uniformly sampled 2-D tensor with the shape (n, 1).
Values range from 0 to 1, indicating probabilities sampled uniformly.
sample_indices : relax.Expr
The 2-D tensor with the shape [n, 1], which indicates the specific
probability distribution to sample from. The value of sample_indices[i]
determines that the ith token should be sampled from the sample_indices[i]th
probability distribution. For instance, if there are 3 distinct probability
distributions and the requirement is to sample 2, 3, and 4 tokens from each,
then sample_indices would be [0, 0, 1, 1, 1, 2, 2, 2, 2].
dtype : str
The data type of the output tensor.
Returns
-------
result : relax.Expr
The computed tensor with shape (n, 1).
Examples
--------
.. code-block:: python
prob = [[0.2, 0.3, 0.5], [0.3, 0.4, 0.3]]
usample = [[0.4], [0.9]]
sample_indices = [[0], [1]]
multinomial_from_uniform(prob, usample)
-> [[1], [2]]
multinomial_from_uniform(prob, usample, sample_indices)
-> [[1], [2]]
"""
return _ffi_api.multinomial_from_uniform( # type: ignore
prob,
uniform_sample,
sample_indices,
dtype,
)
python/tvm/relax/op/search.py
+128
View File
@@ -0,0 +1,128 @@
# 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
"""Search operators."""
from ..expr import Expr
from . import _ffi_api
def where(condition: Expr, x1: Expr, x2: Expr) -> Expr:
"""Selecting elements from either the input tensors depending on the value of the
condition.
For a given position, return the corresponding value in `x1` if `condition` is True,
and return the corresponding value in `x2` otherwise.
Parameters
----------
condition : relax.Expr
When True, yield `x1`; otherwise, yield `x2`.
Must be broadcasting compatible with `x1` and `x2`.
Must have boolean dtype.
x1 : relax.Expr
The first input tensor.
Must be broadcasting compatible with `condition` and `x2`.
x2 : relax.Expr
The second input tensor.
Must be broadcasting compatible with `condition` and `x1`.
Returns
-------
result : relax.Expr
The result tensor.
"""
return _ffi_api.where(condition, x1, x2) # type: ignore
def argmax(x: Expr, axis: int | None = None, keepdims: bool = False) -> Expr:
"""Computes the argmax of tensor elements over given axis.
Parameters
----------
x : relax.Expr
The input data tensor
axis : Optional[int]
Axis along which an argmax operation is performed.
The default, axis=None, will compute the argmax of all elements in the input tensor.
Negative indexing is supported.
keepdims : bool
If this is set to True, the axis being reduced is left in the result as dimensions
with size one.
With this option, the result will broadcast correctly against the input tensor.
Returns
-------
result : relax.Expr
The computed result.
"""
return _ffi_api.argmax(x, axis, keepdims) # type: ignore
def argmin(x: Expr, axis: int | None = None, keepdims: bool = False) -> Expr:
"""Computes the argmin of tensor elements over given axis.
Parameters
----------
x : relax.Expr
The input data tensor
axis : Optional[int]
Axis along which an argmin operation is performed.
The default, axis=None, will compute the argmin of all elements in the
input tensor. Negative indexing is supported.
keepdims : bool
If this is set to True, the axis being reduced is left in the result as
dimensions with size one.
With this option, the result will broadcast correctly against the input tensor.
Returns
-------
result : relax.Expr
The computed result.
"""
return _ffi_api.argmin(x, axis, keepdims) # type: ignore
def bucketize(input_tensor, boundaries, out_int32=False, right=False):
"""Returns the indices of the buckets to which each value in the input belongs.
Parameters
----------
input_tensor : relax.Expr
N-D tensor containing the search values.
boundaries : relax.Expr
1-D tensor, must contain a strictly increasing sequence, or the return value is undefined.
out_int32 : Optional[bool]
Indicate the output data type. int32 if True, int64 otherwise. Default=False
right : Optional[bool]
Determines the behavior for values in boundaries. Similar to torch.bucketize
Returns
-------
result : relax.Expr
The computed result with same shape as input_tensor.
"""
return _ffi_api.bucketize(input_tensor, boundaries, out_int32, right)
python/tvm/relax/op/set.py
+216
View File
@@ -0,0 +1,216 @@
# 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=import-outside-toplevel, redefined-builtin, unused-argument
"""Set operators."""
import numpy as np # type: ignore
import tvm
from ..expr import Expr, prim_value
from . import _ffi_api
def unique(
x: Expr,
sorted: bool | Expr = True,
return_index: bool | Expr = False,
return_inverse: bool | Expr = False,
return_counts: bool | Expr = False,
axis: int | Expr | None = None,
) -> Expr:
"""Find the unique elements in a given tensor.
In addition, it optionally returns
- the indices of the input tensor that give the unique values;
- the indices of the unique tensor that reconstruct the input tensor;
- the number of times each unique value comes up in the input tensor.
Parameters
----------
x : relax.Expr
The input tensor.
sorted : Union[bool, Expr]
Whether to sort the unique elements in ascending order before
returning as output.
return_index : Union[bool, Expr]
Whether to return an additional tensor with indices for where elements in
the unique tensor come from the original input.
return_inverse : Union[bool, Expr]
Whether to return an additional tensor with indices for where elements in
the original input ended up in the returned unique list.
return_counts : Union[bool, Expr]
Whether to return an additional tensor with counts of each unique elements.
axis : Optional
The dimension to apply unique.
If not specified, the unique values of the flattened input are returned.
Returns
-------
ret : relax.Expr
The created relax call with
"""
if isinstance(sorted, bool):
sorted = prim_value(sorted)
if isinstance(return_index, bool):
return_index = prim_value(return_index)
if isinstance(return_inverse, bool):
return_inverse = prim_value(return_inverse)
if isinstance(return_counts, bool):
return_counts = prim_value(return_counts)
if axis is not None and isinstance(axis, int):
axis = prim_value(axis)
return _ffi_api.unique( # type: ignore
x, sorted, return_index, return_inverse, return_counts, axis
)
@tvm.register_global_func("relax.run.unique")
def numpy_unique(
x: tvm.runtime.tensor,
sorted: int,
return_index: int,
return_inverse: int,
return_counts: int,
axis: int | None = None,
) -> tvm.runtime.tensor:
"""Returns the unique elements of the input tensor.
Uses numpy.unique to compute unique elements.
"""
import builtins
x_numpy = x.numpy()
# Call numpy.unique with all the requested return flags
result = np.unique(
x_numpy,
return_index=bool(return_index),
return_inverse=bool(return_inverse),
return_counts=bool(return_counts),
axis=axis,
)
# If no optional outputs requested, result is just the unique values
if not bool(return_index) and not bool(return_inverse) and not bool(return_counts):
unique_values = result
if not sorted:
indices = np.unique(x_numpy, return_index=True, axis=axis)[1]
unique_values = np.take(x_numpy, builtins.sorted(indices), axis=axis)
return tvm.runtime.tensor(unique_values)
# Otherwise, numpy returns a tuple
unique_values = result[0]
output_list = []
result_idx = 1
# Handle sorting for unique values
if not sorted and bool(return_index):
# Get the indices from numpy result
indices = result[result_idx]
result_idx += 1
# Sort indices to get original order
sort_order = np.argsort(indices)
unique_values = np.take(unique_values, sort_order, axis=axis)
indices = np.sort(indices)
output_list.append(tvm.runtime.tensor(unique_values))
output_list.append(tvm.runtime.tensor(indices))
elif not sorted:
# Need to get indices to reorder
_, indices = np.unique(x_numpy, return_index=True, axis=axis)
sort_order = np.argsort(indices)
unique_values = np.take(unique_values, sort_order, axis=axis)
output_list.append(tvm.runtime.tensor(unique_values))
if bool(return_index):
indices_from_result = result[result_idx]
result_idx += 1
output_list.append(tvm.runtime.tensor(np.sort(indices_from_result)))
else:
# Sorted case
output_list.append(tvm.runtime.tensor(unique_values))
if bool(return_index):
output_list.append(tvm.runtime.tensor(result[result_idx]))
result_idx += 1
if bool(return_inverse):
inverse_indices = result[result_idx]
if not sorted:
# Need to remap inverse indices to match reordered unique values
_, orig_indices = np.unique(x_numpy, return_index=True, axis=axis)
sort_order = np.argsort(orig_indices)
inverse_mapping = np.empty_like(sort_order)
inverse_mapping[sort_order] = np.arange(len(sort_order))
inverse_indices = inverse_mapping[inverse_indices]
# ONNX spec: inverse_indices is always 1D
# When axis is None, it has length X.size (flattened)
# When axis is specified, it has length X.shape[axis]
# numpy.unique already returns 1D inverse_indices, so no reshaping needed
output_list.append(tvm.runtime.tensor(inverse_indices))
result_idx += 1
if bool(return_counts):
counts = result[result_idx]
if not sorted:
# Reorder counts to match reordered unique values
_, orig_indices = np.unique(x_numpy, return_index=True, axis=axis)
sort_order = np.argsort(orig_indices)
counts = counts[sort_order]
output_list.append(tvm.runtime.tensor(counts))
return tuple(output_list)
def nonzero(x: Expr) -> Expr:
"""Find the indices of elements of a tensor that are non-zero.
Parameters
----------
x : relax.Expr
The input data tensor.
Returns
-------
result : relax.Expr
A 2-D tensor containing indices of non-zero elements.
Note
----
This function is equivalent to `onnx.nonzero`.
Examples
--------
.. code-block:: python
x = [[0, 1],
[2, 0]]
nonzero(x) = [[0, 1],
[1, 0]]
"""
return _ffi_api.nonzero(x) # type: ignore
@tvm.register_global_func("relax.run.nonzero")
def numpy_nonzero(x: tvm.runtime.tensor) -> tvm.runtime.tensor:
np_result = np.atleast_1d(x.numpy()).nonzero()
return tvm.runtime.tensor(np.stack(np_result, axis=0))
python/tvm/relax/op/sorting.py
+117
View File
@@ -0,0 +1,117 @@
# 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.
"""Sortings operators."""
from ..expr import Constant, Expr
from . import _ffi_api
def sort(x: Expr, axis: int = -1, descending: bool = False):
"""Performs sorting along the given axis and returns an array
in sorted order.
Parameters
----------
x : relax.Expr
The input tensor.
axis : int
Axis along which to sort the input tensor.
By default the last axis of the input is used.
descending : bool
Whether to sort in descending order, the default is False
Returns
-------
out : relax.Expr
Sorted tensor.
"""
return _ffi_api.sort(x, axis, descending) # type: ignore
def argsort(data: Expr, axis: int = -1, descending: bool = False, dtype: str = "int32"):
"""Performs sorting along the given axis and returns an array of indices
having same shape as an input array that index data in sorted order.
Parameters
----------
data : relax.Expr
The input data tensor.
axis : int
Axis long which to sort the input tensor.
descending : bool
Whether to sort in descending order, the default is False
dtype : str
The data type of the output indices.
Returns
-------
out : relax.Expr
Tensor with same shape as data.
"""
return _ffi_api.argsort(data, axis, descending, dtype) # type: ignore
def topk(
data: Expr,
k: int = 1,
axis: int = -1,
ret_type: str = "both",
largest: bool = True,
dtype: str = "int32",
):
"""Get the top k elements in an input tensor along the given axis.
ret_type specifies the return type, can be one of ("both", "values", "indices").
Parameters
----------
data : relax.Expr
The input data tensor.
k : int
Number of top elements to select. Return all elements if k < 1.
axis : int
Axis long which to sort the input tensor.
ret_type: str
The return type [both, values, indices].
"both": return both top k data and indices.
"values": return top k data only.
"indices": return top k indices only.
largest : bool
Whether to return largest or smallest elements.
The k smallest elements are returned if largest is False.
dtype : str
The data type of the indices output.
Returns
-------
out : relax.Expr or List[relax.Expr]
The computed result.
"""
if isinstance(k, Constant):
k = k.data.numpy().item()
return _ffi_api.topk(data, k, axis, ret_type, largest, dtype) # type: ignore
python/tvm/relax/op/statistical.py
+375
View File
@@ -0,0 +1,375 @@
# 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=redefined-builtin
"""Statistical operators."""
from tvm import DataType
from tvm.ir import PrimType
from ..expr import Expr
from . import _ffi_api
def _raw_dtype(dtype):
return dtype.dtype if isinstance(dtype, PrimType) else dtype
def max(x: Expr, axis: int | list[int] | None = None, keepdims: bool = False) -> Expr:
"""Computes the max of tensor elements over given axes.
Parameters
----------
x : relax.Expr
The input data tensor
axis : Optional[Union[int, List[int]]]
Axis or axes along which a max operation is performed.
The default, axis=None, will compute the max of all elements in the input tensor.
Negative indexing is supported.
keepdims : bool
If this is set to True, the axes which are reduced are left in the result as dimensions
with size one.
With this option, the result will broadcast correctly against the input tensor.
Returns
-------
result : relax.Expr
The computed result.
"""
if isinstance(axis, int):
axis = [axis]
return _ffi_api.max(x, axis, keepdims) # type: ignore
def mean(x: Expr, axis: int | list[int] | None = None, keepdims: bool = False) -> Expr:
"""Computes the mean of tensor elements over given axes.
Parameters
----------
x : relax.Expr
The input data tensor
axis : Optional[Union[int, List[int]]]
Axis or axes along which a mean operation is performed.
The default, axis=None, will compute the mean of all elements in the input tensor.
Negative indexing is supported.
keepdims : bool
If this is set to True, the axes which are reduced are left in the result as dimensions
with size one.
With this option, the result will broadcast correctly against the input tensor.
Returns
-------
result : relax.Expr
The computed result.
"""
if isinstance(axis, int):
axis = [axis]
return _ffi_api.mean(x, axis, keepdims) # type: ignore
def min(x: Expr, axis: int | list[int] | None = None, keepdims: bool = False) -> Expr:
"""Computes the min of tensor elements over given axes.
Parameters
----------
x : relax.Expr
The input data tensor
axis : Optional[Union[int, List[int]]]
Axis or axes along which a min operation is performed.
The default, axis=None, will compute the min of all elements in the input tensor.
Negative indexing is supported.
keepdims : bool
If this is set to True, the axes which are reduced are left in the result as dimensions
with size one.
With this option, the result will broadcast correctly against the input tensor.
Returns
-------
result : relax.Expr
The computed result.
"""
if isinstance(axis, int):
axis = [axis]
return _ffi_api.min(x, axis, keepdims) # type: ignore
def prod(x: Expr, axis: int | list[int] | None = None, keepdims: bool = False) -> Expr:
"""Computes the product of tensor elements over given axes.
Parameters
----------
x : relax.Expr
The input data tensor
axis : Optional[Union[int, List[int]]]
Axis or axes along which a product is performed.
The default, axis=None, will compute the product of all elements of the input tensor.
Negative indexing is supported.
keepdims : bool
If this is set to True, the axes which are reduced are left in the result as
dimensions with size one.
With this option, the result will broadcast correctly against the input tensor.
Returns
-------
result : relax.Expr
The computed result.
"""
if isinstance(axis, int):
axis = [axis]
return _ffi_api.prod(x, axis, keepdims) # type: ignore
def std(x: Expr, axis: int | list[int] | None = None, keepdims: bool = False) -> Expr:
"""Computes the standard deviation of tensor elements over given axes.
Parameters
----------
x : relax.Expr
The input data tensor
axis : Optional[Union[int, List[int]]]
Axis or axes along which a standard deviation is performed.
The default, axis=None, will compute the std of all elements of the input tensor.
Negative indexing is supported.
keepdims : bool
If this is set to True, the axes which are reduced are left in the result as
dimensions with size one.
With this option, the result will broadcast correctly against the input tensor.
Returns
-------
result : relax.Expr
The computed result.
"""
if isinstance(axis, int):
axis = [axis]
return _ffi_api.std(x, axis, keepdims) # type: ignore
def sum(x: Expr, axis: int | list[int] | None = None, keepdims: bool = False) -> Expr:
"""Computes the sum of tensor elements over given axes.
Parameters
----------
x : relax.Expr
The input data tensor
axis : Optional[Union[int, List[int]]]
Axis or axes along which a sum is performed.
The default, axis=None, will sum all of the elements of the input tensor.
Negative indexing is supported.
keepdims : bool
If this is set to True, the axes which are reduced are left in the result as
dimensions with size one.
With this option, the result will broadcast correctly against the input tensor.
Returns
-------
result : relax.Expr
The computed result.
"""
if isinstance(axis, int):
axis = [axis]
return _ffi_api.sum(x, axis, keepdims) # type: ignore
def cumprod(
data: Expr,
axis: int | None = None,
dtype: str | DataType | None = None,
exclusive: bool = False,
):
"""Numpy style cumprod op. Return the cumulative product of the elements along
a given axis.
Parameters
----------
data : relax.Expr
The input data to the operator.
axis : Optional[int]
Axis along which the cumulative product is computed. The default (None) is to compute
the cumprod over the flattened array.
dtype : Optional[Union[str, DataType]]
Type of the returned array and of the accumulator in which the elements are computed.
If dtype is not specified, it defaults to the dtype of data.
exclusive : bool
If false (default), all elements are included in the product. If
true, the first element is excluded from the product.
Returns
-------
result : relax.Expr
The result has the same size as data, and the same shape as data if axis is not None.
If axis is None, the result is a 1-d array.
Examples
--------
.. code-block:: python
a = [[1, 2, 3], [4, 5, 6]]
cumprod(a) # if axis is not provided, cumprod is done over the flattened input.
-> [ 1, 2, 6, 24, 120, 720]
cumprod(a, dtype="float32")
-> [ 1., 2., 6., 24., 120., 720.]
cumprod(a, axis=0) # multiply over rows for each of the 3 columns
-> [[1, 2, 3],
[4, 10, 18]]
cumprod(a, axis=1)
-> [[ 1, 2, 6],
[ 4, 20, 120]]
a = [1, 1, 1, 0, 1, 1, 0] # a is a boolean array
cumprod(a, dtype=int32) # dtype should be provided to get the expected results
-> [1, 1, 1, 0, 0, 0, 0]
"""
if exclusive is None:
exclusive = False
return _ffi_api.cumprod(data, axis, _raw_dtype(dtype), exclusive) # type: ignore
def cumsum(
data: Expr,
axis: int | None = None,
dtype: str | DataType | None = None,
exclusive: bool = False,
):
"""Numpy style cumsum op. Return the cumulative inclusive sum of the elements along
a given axis.
Parameters
----------
data : relax.Expr
The input data to the operator.
axis : Optional[int]
Axis along which the cumulative sum is computed. The default (None) is to compute
the cumsum over the flattened array.
dtype : Optional[Union[str, DataType]]
Type of the returned array and of the accumulator in which the elements are summed.
If dtype is not specified, it defaults to the dtype of data.
exclusive : bool
If false (default), all elements are included in the sum. If
true, the first element is excluded from the sum.
Returns
-------
result : relax.Expr
The result has the same size as data, and the same shape as data if axis is not None.
If axis is None, the result is a 1-d array.
Examples
--------
.. code-block:: python
a = [[1, 2, 3], [4, 5, 6]]
cumsum(a) # if axis is not provided, cumsum is done over the flattened input.
-> [ 1, 3, 6, 10, 15, 21]
cumsum(a, dtype="float32")
-> [ 1., 3., 6., 10., 15., 21.]
cumsum(a, axis=0) # sum over rows for each of the 3 columns
-> [[1, 2, 3],
[5, 7, 9]]
cumsum(a, axis=1)
-> [[ 1, 3, 6],
[ 4, 9, 15]]
a = [1, 0, 1, 0, 1, 1, 0] # a is a boolean array
cumsum(a, dtype=int32) # dtype should be provided to get the expected results
-> [1, 1, 2, 2, 3, 4, 4]
"""
if exclusive is None:
exclusive = False
return _ffi_api.cumsum(data, axis, _raw_dtype(dtype), exclusive) # type: ignore
def variance(x: Expr, axis: int | list[int] | None = None, keepdims: bool = False) -> Expr:
"""Computes the variance of tensor elements over given axes.
Parameters
----------
x : relax.Expr
The input data tensor
axis : Optional[Union[int, List[int]]]
Axis or axes along which a variance operation is performed.
The default, axis=None, will compute the variance of all elements in the input tensor.
Negative indexing is supported.
keepdims : bool
If this is set to True, the axes which are reduced are left in the result as dimensions
with size one.
With this option, the result will broadcast correctly against the input tensor.
Returns
-------
result : relax.Expr
The computed result.
"""
if isinstance(axis, int):
axis = [axis]
return _ffi_api.variance(x, axis, keepdims) # type: ignore
def median(x: Expr, axis: int | list[int] | None = None, keepdims: bool = False) -> Expr:
"""Computes the median of tensor elements over given axes.
Parameters
----------
x : relax.Expr
The input data tensor
axis : Optional[Union[int, List[int]]]
Axis along which the median is computed. The default (None) is to compute
the median of the entire flattened tensor.
keepdims : bool
If this is set to True, the axes which are reduced are left in the result as dimensions
with size one.
With this option, the result will broadcast correctly against the input tensor.
Returns
-------
result : relax.Expr
The computed result.
"""
if isinstance(axis, int):
axis = [axis]
return _ffi_api.median(x, axis, keepdims) # type: ignore
python/tvm/relax/op/ternary.py
+44
View File
@@ -0,0 +1,44 @@
# 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=redefined-builtin, invalid-name
"""Relax ternary arithmetic operators."""
from ..expr import Expr
from . import _ffi_api
def ewise_fma(x1: Expr, x2: Expr, x3: Expr) -> Expr:
"""Elementwise fused multiply-add operator
Returns elementwise result of :math:`x1 * x2 + x3`
Parameters
----------
x1 : relax.Expr
The left hand operand of the multiplication
x2 : relax.Expr
The right hand operand of the multiplication
x3 : relax.Expr
The operand of the addition
Returns
-------
result : relax.Expr
The computed result.
"""
return _ffi_api.ewise_fma(x1, x2, x3) # type: ignore
python/tvm/relax/op/unary.py
+617
View File
@@ -0,0 +1,617 @@
# 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=redefined-builtin, invalid-name
"""Relax unary arithmetic operators."""
from ..expr import Expr
from ..utils import convert_to_expr
from . import _ffi_api
###################### Arithmetic operators ######################
def abs(x: Expr) -> Expr:
"""Compute element-wise absolute value of the input data.
Parameters
----------
x : relax.Expr
The input data
Returns
-------
result : relax.Expr
The computed result.
"""
return _ffi_api.abs(x) # type: ignore
def acos(x: Expr) -> Expr:
"""Compute element-wise arc cos of the input data.
Parameters
----------
x : relax.Expr
The input data
Returns
-------
result : relax.Expr
The computed result.
Note
----
The input tensor is required to have float dtype
"""
return _ffi_api.acos(x) # type: ignore
def acosh(x: Expr) -> Expr:
"""Compute element-wise arc cosh of the input data.
Parameters
----------
x : relax.Expr
The input data
Returns
-------
result : relax.Expr
The computed result.
Note
----
The input tensor is required to have float dtype
"""
return _ffi_api.acosh(x) # type: ignore
def asin(x: Expr) -> Expr:
"""Compute element-wise arc sin of the input data.
Parameters
----------
x : relax.Expr
The input data
Returns
-------
result : relax.Expr
The computed result.
Note
----
The input tensor is required to have float dtype
"""
return _ffi_api.asin(x) # type: ignore
def asinh(x: Expr) -> Expr:
"""Compute element-wise arc sinh of the input data.
Parameters
----------
x : relax.Expr
The input data
Returns
-------
result : relax.Expr
The computed result.
Note
----
The input tensor is required to have float dtype
"""
return _ffi_api.asinh(x) # type: ignore
def atan(x: Expr) -> Expr:
"""Compute element-wise arc tan of the input data.
Parameters
----------
x : relax.Expr
The input data
Returns
-------
result : relax.Expr
The computed result.
Note
----
The input tensor is required to have float dtype
"""
return _ffi_api.atan(x) # type: ignore
def atanh(x: Expr) -> Expr:
"""Compute element-wise arc tanh of the input data.
Parameters
----------
x : relax.Expr
The input data
Returns
-------
result : relax.Expr
The computed result.
Note
----
The input tensor is required to have float dtype
"""
return _ffi_api.atanh(x) # type: ignore
def bitwise_not(x: Expr) -> Expr:
"""Compute bitwise NOT of the input data.
Parameters
----------
x : relax.Expr
The input data
Returns
-------
result : relax.Expr
The computed result.
"""
return _ffi_api.bitwise_not(x) # type: ignore
def ceil(x: Expr) -> Expr:
"""Take ceil of input data.
Parameters
----------
x : relax.Expr
The input data
Returns
-------
result : relax.Expr
The computed result.
"""
return _ffi_api.ceil(x) # type: ignore
def cos(x: Expr) -> Expr:
"""Compute element-wise cos of the input data.
Parameters
----------
x : relax.Expr
The input data
Returns
-------
result : relax.Expr
The computed result.
Note
----
The input tensor is required to have float dtype
"""
return _ffi_api.cos(x) # type: ignore
def cosh(x: Expr) -> Expr:
"""Compute element-wise cosh of the input data.
Parameters
----------
x : relax.Expr
The input data
Returns
-------
result : relax.Expr
The computed result.
Note
----
The input tensor is required to have float dtype
"""
return _ffi_api.cosh(x) # type: ignore
def exp(x: Expr) -> Expr:
"""Compute element-wise exp of data.
Parameters
----------
x : relax.Expr
The input data
Returns
-------
result : relax.Expr
The computed result.
Note
----
The input tensor is required to have float dtype
"""
return _ffi_api.exp(x) # type: ignore
def floor(x: Expr) -> Expr:
"""Take floor of input data.
Parameters
----------
x : relax.Expr
The input data
Returns
-------
result : relax.Expr
The computed result.
"""
return _ffi_api.floor(x) # type: ignore
def log(x: Expr) -> Expr:
"""Compute element-wise natural logarithm of the input data.
Parameters
----------
x : relax.Expr
The input data
Returns
-------
result : relax.Expr
The computed result.
Note
----
The input tensor is required to have float dtype
"""
return _ffi_api.log(x) # type: ignore
def logical_not(x: Expr) -> Expr:
"""Compute logical NOT of the input data.
Parameters
----------
x : relax.Expr
The input data
Returns
-------
result : relax.Expr
The computed result.
"""
return _ffi_api.logical_not(x) # type: ignore
def negative(x: Expr) -> Expr:
"""Compute element-wise negative of the input data.
Parameters
----------
x : relax.Expr
The input data
Returns
-------
result : relax.Expr
The computed result
"""
return _ffi_api.negative(x) # type: ignore
def round(x: Expr) -> Expr:
"""Rounds each element of the input data to nearest integer.
Parameters
----------
x : relax.Expr
The input data
Returns
-------
result : relax.Expr
The computed result.
"""
return _ffi_api.round(x) # type: ignore
def rsqrt(x: Expr) -> Expr:
"""Compute element-wise reciprocal square root of the input data.
.. math::
1/sqrt(x)
Parameters
----------
x : relax.Expr
The input data
Returns
-------
result : relax.Expr
The computed result.
Note
----
The input tensor is required to have float dtype
"""
return _ffi_api.rsqrt(x) # type: ignore
def sigmoid(x: Expr) -> Expr:
"""Compute element-wise sigmoid of the input data.
Parameters
----------
x : relax.Expr
The input data
Returns
-------
result : relax.Expr
The computed result.
Note
----
The input tensor is required to have float dtype
"""
return _ffi_api.sigmoid(x) # type: ignore
def sign(x: Expr) -> Expr:
"""Returns an indication of the sign of a number for each element of the input data.
Parameters
----------
x : relax.Expr
The input data
Returns
-------
result : relax.Expr
The computed result.
"""
return _ffi_api.sign(x) # type: ignore
def sin(x: Expr) -> Expr:
"""Compute element-wise sin of the input data.
Parameters
----------
x : relax.Expr
The input data
Returns
-------
result : relax.Expr
The computed result.
Note
----
The input tensor is required to have float dtype
"""
return _ffi_api.sin(x) # type: ignore
def sinh(x: Expr) -> Expr:
"""Compute element-wise sinh of the input data.
Parameters
----------
x : relax.Expr
The input data
Returns
-------
result : relax.Expr
The computed result.
Note
----
The input tensor is required to have float dtype
"""
return _ffi_api.sinh(x) # type: ignore
def square(x: Expr) -> Expr:
"""Squares each element of the input data.
Parameters
----------
x : relax.Expr
The input data
Returns
-------
result : relax.Expr
The computed result.
"""
return _ffi_api.square(x) # type: ignore
def sqrt(x: Expr) -> Expr:
"""Compute element-wise square root of the input data.
Parameters
----------
x : relax.Expr
The input data
Returns
-------
result : relax.Expr
The computed result.
Note
----
The input tensor is required to have float dtype
"""
return _ffi_api.sqrt(x) # type: ignore
def tan(x: Expr) -> Expr:
"""Compute element-wise tan of the input data.
Parameters
----------
x : relax.Expr
The input data
Returns
-------
result : relax.Expr
The computed result.
Note
----
The input tensor is required to have float dtype
"""
return _ffi_api.tan(x) # type: ignore
def tanh(x: Expr) -> Expr:
"""Compute element-wise tanh of the input data.
Parameters
----------
x : relax.Expr
The input data
Returns
-------
result : relax.Expr
The computed result.
Note
----
The input tensor is required to have float dtype
"""
return _ffi_api.tanh(x) # type: ignore
def trunc(x: Expr) -> Expr:
"""Take trunc of input data.
Parameters
----------
x : relax.Expr
The input data
Returns
-------
result : relax.Expr
The computed result.
"""
return _ffi_api.trunc(x) # type: ignore
def clip(x: Expr, min: Expr, max: Expr) -> Expr:
"""Clips tensor values to a specified min and max.
Parameters
----------
x : relax.Expr
The input data
min : relax.Expr
The minimum value
max : relax.Expr
The maximum value
Returns
-------
result : relax.Expr
The computed result.
"""
min = convert_to_expr(min)
max = convert_to_expr(max)
return _ffi_api.clip(x, min, max) # type: ignore
def erf(x: Expr) -> Expr:
"""Computes the error function of the input.
Parameters
----------
x : relax.Expr
The input data
Returns
-------
result : relax.Expr
Computed error function for each element.
"""
return _ffi_api.erf(x) # type: ignore
###################### Check operators ######################
def isfinite(x: Expr) -> Expr:
"""Check if input value is finite.
Parameters
----------
x : relax.Expr
The input data
Returns
-------
result : relax.Expr
The computed result.
"""
return _ffi_api.isfinite(x) # type: ignore
def isinf(x: Expr) -> Expr:
"""Check if input value is infinite.
Parameters
----------
x : relax.Expr
The input data
Returns
-------
result : relax.Expr
The computed result.
"""
return _ffi_api.isinf(x) # type: ignore
def isnan(x: Expr) -> Expr:
"""Check if input value is Nan.
Parameters
----------
x : relax.Expr
The input data
Returns
-------
result : relax.Expr
The computed result.
"""
return _ffi_api.isnan(x) # type: ignore
python/tvm/relax/op/vision/__init__.py
+23
View File
@@ -0,0 +1,23 @@
# isort: skip_file
# 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.
"""VISION operators."""
from .multibox_transform_loc import *
from .nms import *
from .roi_align import *
from .roi_pool import *
python/tvm/relax/op/vision/_ffi_api.py
+21
View File
@@ -0,0 +1,21 @@
# 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.
"""Constructor APIs"""
import tvm_ffi
tvm_ffi.init_ffi_api("relax.op.vision", __name__)
python/tvm/relax/op/vision/multibox_transform_loc.py
+85
View File
@@ -0,0 +1,85 @@
# 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.
"""Multibox location transform for object detection."""
from . import _ffi_api
def multibox_transform_loc(
cls_pred,
loc_pred,
anchor,
clip=False,
threshold=0.0,
variances=(1.0, 1.0, 1.0, 1.0),
keep_background=True,
):
"""SSD / TFLite-style decode: priors + offsets → boxes; logits → softmax scores.
Box decode follows TFLite ``DecodeCenterSizeBoxes``; expected tensor layout matches
``tflite_frontend.convert_detection_postprocess`` (loc reorder yxhw→xywh, anchor ltrb).
Parameters
----------
cls_pred : relax.Expr
``[B, C, N]`` class logits (pre-softmax).
loc_pred : relax.Expr
``[B, 4*N]`` per-anchor encodings as ``(x,y,w,h)`` after reorder (see above).
anchor : relax.Expr
``[1, N, 4]`` priors: ``(left, top, right, bottom)``.
clip : bool
If True, clip ``ymin,xmin,ymax,xmax`` to ``[0, 1]``.
threshold : float
After softmax, multiply scores by mask ``(score >= threshold)``.
variances : tuple of 4 floats
``(x,y,w,h)`` = TFLite ``1/x_scale, 1/y_scale, 1/w_scale, 1/h_scale``.
Use magnitudes consistent with the model: very large ``w``/``h`` entries scale the
encoded height/width terms inside ``exp(...)`` and can overflow in float32/float16.
keep_background : bool
If False, set output scores at class index 0 to zero.
Returns
-------
result : relax.Expr
Tuple ``(boxes, scores)``: ``boxes`` is ``[B, N, 4]`` as ``(ymin,xmin,ymax,xmax)``;
``scores`` is ``[B, C, N]`` softmax, post-processed like the implementation.
Notes
-----
**Shape/dtype (checked in ``FInferType`` when static):**
- ``cls_pred``: 3-D; ``loc_pred``: 2-D; ``anchor``: 3-D.
- ``cls_pred``, ``loc_pred``, ``anchor`` dtypes must match.
- ``N = cls_pred.shape[2]``; ``loc_pred.shape[1] == 4*N``; ``anchor.shape == [1,N,4]``.
- ``loc_pred.shape[1]`` must be divisible by 4.
- ``cls_pred.shape[0]`` must equal ``loc_pred.shape[0]`` (batch).
If ``cls_pred`` has **unknown** shape, inference only returns generic rank-3 tensor
type for the two outputs; it does **not** verify ``4*N`` vs ``loc_pred`` or
``anchor.shape[1]`` vs ``N``, because ``N`` is not available statically. Other checks
(ranks, dtypes, ``loc_pred.shape[1] % 4 == 0`` when known, batch match when both batch
axes are known, etc.) still run where applicable.
"""
return _ffi_api.multibox_transform_loc(
cls_pred,
loc_pred,
anchor,
clip,
threshold,
variances,
keep_background,
)
python/tvm/relax/op/vision/nms.py
+204
View File
@@ -0,0 +1,204 @@
# 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.
"""Non-maximum suppression operators."""
from . import _ffi_api
def all_class_non_max_suppression(
boxes,
scores,
max_output_boxes_per_class,
iou_threshold,
score_threshold,
output_format="onnx",
):
"""Non-maximum suppression operator for object detection, corresponding to ONNX
NonMaxSuppression and TensorFlow combined_non_max_suppression.
NMS is performed for each class separately.
Parameters
----------
boxes : relax.Expr
3-D tensor with shape (batch_size, num_boxes, 4)
scores: relax.Expr
3-D tensor with shape (batch_size, num_classes, num_boxes)
max_output_boxes_per_class : relax.Expr
The maxinum number of output selected boxes per class
iou_threshold : relax.Expr
IoU test threshold
score_threshold : relax.Expr
Score threshold to filter out low score boxes early
output_format : str, optional
"onnx" or "tensorflow", see below.
Returns
-------
out : relax.Expr
If `output_format` is "onnx", the output is two tensors. The first is `indices` of size
`(batch_size * num_class* num_boxes , 3)` and the second is a scalar tensor
`num_total_detection` of shape `(1,)` representing the total number of selected
boxes. The three values in `indices` encode batch, class, and box indices.
Rows of `indices` are ordered such that selected boxes from batch 0, class 0 come
first, in descending of scores, followed by boxes from batch 0, class 1 etc.
The output uses dynamic_strided_slice to trim to only valid detections,
so the first tensor has shape (num_total_detection, 3) containing only valid rows.
If `output_format` is "tensorflow", the output is three tensors, the first
is `indices` of size `(batch_size, num_class * num_boxes , 2)`, the second is `scores` of
size `(batch_size, num_class * num_boxes)`, and the third is `num_total_detection` of size
`(batch_size,)` representing the total number of selected boxes per batch. The two values
in `indices` encode class and box indices. Of num_class * num_boxes boxes in `indices` at
batch b, only the first `num_total_detection[b]` entries are valid. The second axis of
`indices` and `scores` are sorted within each class by box scores, but not across classes.
So the box indices and scores for the class 0 come first in a sorted order, followed by
the class 1 etc.
"""
return _ffi_api.all_class_non_max_suppression(
boxes, scores, max_output_boxes_per_class, iou_threshold, score_threshold, output_format
)
def get_valid_counts(data, score_threshold=0, id_index=0, score_index=1):
"""Get valid count of bounding boxes given a score threshold.
Also moves valid boxes to the top of input data.
Parameters
----------
data : relax.Expr
3-D tensor with shape [batch_size, num_anchors, elem_length].
score_threshold : float, optional
Lower limit of score for valid bounding boxes.
id_index : int, optional
Index of the class categories. Set to ``-1`` to disable the class-id check.
score_index : int, optional
Index of the scores/confidence of boxes.
Returns
-------
out : relax.Expr
A tuple ``(valid_count, out_tensor, out_indices)`` where ``valid_count``
has shape ``[batch_size]``, ``out_tensor`` has shape
``[batch_size, num_anchors, elem_length]``, and ``out_indices`` has shape
``[batch_size, num_anchors]``.
"""
return _ffi_api.get_valid_counts(data, score_threshold, id_index, score_index)
def non_max_suppression(
data,
valid_count,
indices,
max_output_size=-1,
iou_threshold=0.5,
force_suppress=False,
top_k=-1,
coord_start=2,
score_index=1,
id_index=0,
return_indices=True,
invalid_to_bottom=False,
soft_nms_sigma=0.0,
score_threshold=0.0,
):
"""Non-maximum suppression operator for object detection.
Parameters
----------
data : relax.Expr
3-D tensor with shape [batch_size, num_anchors, elem_length].
valid_count : relax.Expr
1-D tensor for valid number of boxes.
indices : relax.Expr
2-D tensor with shape [batch_size, num_anchors].
max_output_size : int, optional
Max number of output valid boxes, -1 for no limit.
iou_threshold : float, optional
Non-maximum suppression IoU threshold.
force_suppress : bool, optional
Whether to suppress all detections regardless of class_id. When
``id_index`` is ``-1``, all valid boxes are treated as belonging to the
same class, so this flag has the same effect as ``True``.
top_k : int, optional
Keep maximum top k detections before nms, -1 for no limit.
coord_start : int, optional
Start index of the consecutive 4 coordinates.
score_index : int, optional
Index of the scores/confidence of boxes.
id_index : int, optional
Index of the class categories. Set to ``-1`` to suppress boxes across
all classes.
return_indices : bool, optional
Whether to return box indices in input data.
invalid_to_bottom : bool, optional
Whether to move valid bounding boxes to the top of the returned tensor.
This option only affects the ``return_indices=False`` path.
soft_nms_sigma : float, optional
Sigma for soft-NMS Gaussian penalty. When ``0.0`` (default), standard
hard NMS is used. Positive values decay overlapping box scores instead
of suppressing them outright.
score_threshold : float, optional
Post-decay minimum score for a box to remain eligible during soft-NMS.
Only used when ``soft_nms_sigma > 0``. This is distinct from
``get_valid_counts.score_threshold``, which filters boxes before NMS.
Defaults to ``0.0``.
Returns
-------
out : relax.Expr
The return tuple shape depends on ``soft_nms_sigma``.
If ``return_indices`` is ``True`` and ``soft_nms_sigma`` is ``0.0``,
returns a 2-tuple ``(box_indices, valid_box_count)`` with shapes
``[batch_size, num_anchors]`` and ``[batch_size, 1]``.
If ``return_indices`` is ``True`` and ``soft_nms_sigma > 0``,
returns a 3-tuple ``(out_data, box_indices, valid_box_count)`` where
decayed ``out_data`` is prepended and has the same shape as the input
data.
Otherwise returns the modified data tensor.
"""
return _ffi_api.non_max_suppression(
data,
valid_count,
indices,
max_output_size,
iou_threshold,
force_suppress,
top_k,
coord_start,
score_index,
id_index,
return_indices,
invalid_to_bottom,
soft_nms_sigma,
score_threshold,
)
python/tvm/relax/op/vision/roi_align.py
+78
View File
@@ -0,0 +1,78 @@
# 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.
"""ROI Align operator"""
from ..base import Expr
from . import _ffi_api
def roi_align(
data: Expr,
rois: Expr,
pooled_size: int | tuple[int, int] | list[int],
spatial_scale: float,
sample_ratio: int = -1,
aligned: bool = False,
layout: str = "NCHW",
mode: str = "avg",
):
"""ROI Align operator.
Parameters
----------
data : relax.Expr
4-D input tensor.
rois : relax.Expr
2-D input tensor with shape `(num_roi, 5)` in
`[batch_idx, x1, y1, x2, y2]` format.
pooled_size : Union[int, Tuple[int, int], List[int]]
Output pooled size.
spatial_scale : float
Ratio of input feature map height (or width) to raw image height (or width).
sample_ratio : int, optional
Sampling ratio for ROI align. Non-positive values use adaptive sampling.
aligned : bool, optional
Whether to use aligned ROIAlign semantics without the legacy 1-pixel clamp.
layout : str, optional
Layout of the input data. Supported values are `NCHW` and `NHWC`.
mode : str, optional
Mode for ROI align. Supported values are `avg` and `max`.
Returns
-------
result : relax.Expr
The computed result.
"""
if isinstance(pooled_size, int):
pooled_size = (pooled_size, pooled_size)
return _ffi_api.roi_align(
data,
rois,
pooled_size,
spatial_scale,
sample_ratio,
aligned,
layout,
mode,
)
python/tvm/relax/op/vision/roi_pool.py
+57
View File
@@ -0,0 +1,57 @@
# 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.
"""ROI Pool operator"""
from ..base import Expr
from . import _ffi_api
def roi_pool(
data: Expr,
rois: Expr,
pooled_size: int | tuple[int, int] | list[int],
spatial_scale: float,
layout: str = "NCHW",
):
"""ROI Pool operator.
Parameters
----------
data : relax.Expr
4-D input tensor.
rois : relax.Expr
2-D input tensor with shape `(num_roi, 5)` in
`[batch_idx, x1, y1, x2, y2]` format.
pooled_size : Union[int, Tuple[int, int], List[int]]
Output pooled size.
spatial_scale : float
Ratio of input feature map height (or width) to raw image height (or width).
layout : str, optional
Layout of the input data. Currently only `NCHW` is supported.
Returns
-------
result : relax.Expr
The computed result.
"""
if isinstance(pooled_size, int):
pooled_size = (pooled_size, pooled_size)
return _ffi_api.roi_pool(data, rois, pooled_size, spatial_scale, layout)
python/tvm/relax/op/vm/__init__.py
+19
View File
@@ -0,0 +1,19 @@
# 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.
"""Relax vm primitives."""
from .vm import alloc_storage, alloc_tensor, call_tir_dyn, kill_object
python/tvm/relax/op/vm/_ffi_api.py
+20
View File
@@ -0,0 +1,20 @@
# 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
"""FFI APIs for tvm.relax.op.vm"""
import tvm_ffi
tvm_ffi.init_ffi_api("relax.op.vm", __name__)
python/tvm/relax/op/vm/vm.py
+144
View File
@@ -0,0 +1,144 @@
# 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
"""Relax vm primitives."""
from tvm.ir import Call
from ...expr import DataTypeImm, Expr, StringImm, Tuple, prim_value
from ...utils import convert_to_expr
from . import _ffi_api
def alloc_storage(
shape: Expr,
runtime_device_index: int | Expr,
dtype: str | Expr,
storage_scope: str | StringImm = "global",
) -> Call:
"""Construct a Call to allocate a storage with specific size,
runtime_device_index, and dtype.
Parameters
----------
shape : Expr
The shape of the storage to be allocated.
runtime_device_index : Union[int, Expr]
The device index indicating on which device the tensor is to
be allocated at runtime. Index -1 is reserved for the host device.
dtype : Union[str, Expr]
The datatype of the storage to be allocated.
storage_scope : Union[str, StringImm]
The storage scope of the storage to allocate. Default is global.
Returns
-------
result : Call
A relax Call, which gets the allocated storage.
"""
shape = convert_to_expr(shape)
if isinstance(dtype, str):
dtype = DataTypeImm(dtype)
if isinstance(storage_scope, str):
storage_scope = StringImm(storage_scope)
if isinstance(runtime_device_index, int):
runtime_device_index = prim_value(runtime_device_index)
return _ffi_api.alloc_storage(shape, runtime_device_index, dtype, storage_scope) # type: ignore
def alloc_tensor(
storage: Expr,
offset: int | Expr,
shape: Expr,
dtype: str | Expr,
runtime_device_ind: int | Expr = prim_value(0),
) -> Call:
"""Construct a Call to allocate a tensor on a certain storage starting from the given offset.
Parameters
----------
storage : Expr
The storage to allocate the tensor to.
offset : Union[int, Expr]
The storage offset to allocate the tensor.
shape : Expr
The shape of the tensor to be allocated.
dtype : Union[str, Expr]
The datatype of the tensor to be allocated.
runtime_device_ind: Union[int, Expr]
The device index indicating on which device the tensor is to be
allocated at runtime. Index -1 is reserved for the host device.
Returns
-------
result : Call
A relax Call, which gets the allocated tensor.
"""
if isinstance(offset, int):
offset = prim_value(offset)
shape = convert_to_expr(shape)
if isinstance(dtype, str):
dtype = DataTypeImm(dtype)
if isinstance(runtime_device_ind, int):
runtime_device_ind = prim_value(runtime_device_ind)
return _ffi_api.alloc_tensor(storage, offset, shape, dtype, runtime_device_ind) # type: ignore
def kill_object(obj: Expr) -> Call:
"""Construct a Call to set the register corresponding to the input object to
null at runtime, in order to kill the input object.
Parameters
----------
obj : Expr
The object to be killed.
Returns
-------
result : Call
CallNode that kills the input object.
"""
return _ffi_api.kill_object(obj) # type: ignore
def call_tir_dyn(func: Expr, args: Tuple) -> Call:
"""Construct a Call to call_tir_dyn (invoke the given TIR PrimFunc)
consisting of the input tensors and the shape of the result.
Parameters
----------
func : Expr
An expression evaluating to a TIR PrimFunc.
args : Tuple
The input args, includes a list of tensors, and a ShapeExpr.
Returns
-------
result : Call
A relax Call to call_tir_dyn.
"""
func = convert_to_expr(func)
if isinstance(args, list | tuple):
args = Tuple(args)
return _ffi_api.call_tir_dyn(func, args) # type: ignore