wehub-resource-sync
6128 lines
224 KiB
Python
6128 lines
224 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
# ruff: noqa: E501, E731, E741, RUF005
|
|
"""ONNX: Open Neural Network Exchange importer for Relax.
|
|
|
|
This module implements the required functionality to read ONNX models
|
|
and convert them into equivalent Relax functions. The entry point that encapsulates
|
|
this functionality is the function from_onnx.
|
|
|
|
In order to extend the functionality of the importer, you can add new
|
|
operators to the operator registry. The operator registry is a dictionary
|
|
that maps operator names to operator converters. The registry is defined
|
|
in the _get_converter_map function. To add a new operator, you can define
|
|
a new class that inherits from the OnnxOpConverter class and implement
|
|
the _impl method.
|
|
|
|
By default, ONNX defines models in terms of dynamic shapes. The ONNX importer
|
|
retains dynamic shapes upon import, and when possible, the compiler attempts to
|
|
convert the model to use static shapes at compile time.
|
|
If this fails, there may still be dynamic operations in the model.
|
|
Not all TVM kernels currently support dynamic shapes, please file an issue on
|
|
github.com/apache/tvm/issues if you hit an error with dynamic kernels.
|
|
"""
|
|
|
|
import contextlib
|
|
import functools
|
|
import math
|
|
import operator
|
|
import re
|
|
import warnings
|
|
from collections.abc import Callable
|
|
from typing import Any
|
|
|
|
import numpy as _np
|
|
|
|
try:
|
|
import onnx.onnx_ml_pb2
|
|
except ImportError as err:
|
|
raise ImportError(
|
|
"onnx is required by the ONNX frontend. Install it with: pip install onnx"
|
|
) from err
|
|
|
|
import tvm_ffi
|
|
|
|
import tvm
|
|
from tvm import relax, tirx, topi
|
|
from tvm.ir import IRModule
|
|
from tvm.ir.supply import UniqueNameSupply
|
|
from tvm.runtime import DataType, DataTypeCode
|
|
from tvm.topi.utils import get_const_tuple
|
|
|
|
from ..common import autopad
|
|
|
|
|
|
def _relax_dtype_is_floating_point(dtype: str) -> bool:
|
|
"""Whether a Relax dtype string is a floating point type."""
|
|
try:
|
|
code = DataType(dtype).type_code
|
|
except (ValueError, TypeError, RuntimeError):
|
|
return False
|
|
return (
|
|
code == DataTypeCode.FLOAT
|
|
or code == DataTypeCode.BFLOAT
|
|
or (code >= DataTypeCode.Float8E3M4 and code <= DataTypeCode.Float4E2M1FN)
|
|
)
|
|
|
|
|
|
def get_type(elem_type: str | int) -> str:
|
|
"""Converts onnx integer datatype to numpy datatype"""
|
|
# If a string was passed instead of a tensor type, it does not need
|
|
# conversion and can be returned.
|
|
if isinstance(elem_type, str):
|
|
return elem_type
|
|
|
|
try:
|
|
from onnx.helper import ( # pylint: disable=import-outside-toplevel
|
|
tensor_dtype_to_np_dtype,
|
|
)
|
|
except ImportError as exception:
|
|
raise ImportError(f"Unable to import onnx which is required {exception}")
|
|
|
|
return str(tensor_dtype_to_np_dtype(elem_type))
|
|
|
|
|
|
def get_constant(
|
|
var: relax.Constant | relax.Var,
|
|
params: list[dict[str, relax.Var]],
|
|
) -> relax.Constant | relax.Var:
|
|
"""Attempt to convert a variable to a constant if possible.
|
|
This is the primary function meant to interact with params.
|
|
|
|
Parameters
|
|
----------
|
|
var: Union[relax.Constant, relax.Var]
|
|
The input value to try to convert to a constant.
|
|
params: List[Dict[str, relax.Var]]
|
|
The parameters for the graph. Contains both the global registry of nodes
|
|
for the graph and the parameter dictionary. The global registry is updated
|
|
with a constant value if possible.
|
|
|
|
Returns
|
|
-------
|
|
var : Union[relax.Constant, relax.Var]
|
|
The input value converted to a constant if possible. If the value
|
|
isn't found in params, the input variable is returned unmodified.
|
|
"""
|
|
# Params is actually both the graph nodes and param dictionary, unpack them.
|
|
graph_nodes, params = params
|
|
# Convert if possible
|
|
if isinstance(var, relax.Var) and var.name_hint in params:
|
|
# When converting a parameter to a constant, update references to it as well.
|
|
_, value = params[var.name_hint]
|
|
const_value = relax.const(value)
|
|
graph_nodes[var.name_hint] = const_value
|
|
return const_value
|
|
# Otherwise return variable.
|
|
else:
|
|
return var
|
|
|
|
|
|
def get_value(token, value_dict: dict[str, tvm.tirx.Var]) -> int | tvm.tirx.Var:
|
|
"""Convert a token to an integer constant or symbolic variable.
|
|
|
|
Parameters
|
|
----------
|
|
token: str
|
|
current token to decode.
|
|
|
|
value_dict: Dict
|
|
The dictionary mapping ValueInfoProto names to symbolic variables.
|
|
|
|
Returns
|
|
-------
|
|
Union[int, tvm.tirx.Var]
|
|
The decoded token
|
|
"""
|
|
|
|
try:
|
|
return int(token)
|
|
except ValueError:
|
|
if token not in value_dict or token == "?":
|
|
value_dict[token] = tvm.tirx.Var(token, "int64")
|
|
value = value_dict[token]
|
|
return value
|
|
|
|
|
|
def parse_shape_name(name: str, value_dict: dict[str, tvm.tirx.Var]) -> tirx.Expr | tvm.tirx.Var:
|
|
"""Converts expressions in the shape dimension name to prim expressions.
|
|
|
|
Parameters
|
|
----------
|
|
name: str
|
|
name of shape dimension.
|
|
|
|
value_dict: Dict
|
|
The dictionary mapping ValueInfoProto names to symbolic variables.
|
|
|
|
Returns
|
|
-------
|
|
Union[tirx.Expr, tvm.tirx.Var]
|
|
The expression of the shape dimension.
|
|
"""
|
|
|
|
tokens = re.split(r"(\+|\-|\*|\/\/|\/)", name.replace(" ", ""))
|
|
|
|
operators = {
|
|
"+": operator.add,
|
|
"-": operator.sub,
|
|
"*": operator.mul,
|
|
"/": operator.floordiv, # is floordiv since the operands are always int
|
|
"//": operator.floordiv,
|
|
}
|
|
|
|
value_stack = []
|
|
operator_stack = []
|
|
|
|
for token in tokens:
|
|
if token in operators:
|
|
operator_stack.append(token)
|
|
else:
|
|
value = get_value(token, value_dict)
|
|
if value_stack and operator_stack:
|
|
prev_value = value_stack.pop()
|
|
op = operator_stack.pop()
|
|
result = operators[op](prev_value, value)
|
|
value_stack.append(result)
|
|
else:
|
|
value_stack.append(value)
|
|
|
|
if value_stack:
|
|
return value_stack[0]
|
|
else:
|
|
raise Exception("Shape dimension could not be inferred")
|
|
|
|
|
|
def get_info(
|
|
info_proto: onnx.onnx_ml_pb2.ValueInfoProto, value_dict: dict[str, tvm.tirx.Var]
|
|
) -> tuple[str, list, str, list, dict]:
|
|
"""Extract the shape from a ValueInfoProto.
|
|
|
|
Parameters
|
|
----------
|
|
info_proto: onnx.onnx_ml_pb2.ValueInfoProto
|
|
The ValueInfoProto to extract the info from.
|
|
|
|
value_dict: Dict
|
|
The dictionary mapping ValueInfoProto names to symbolic variables.
|
|
|
|
Returns
|
|
-------
|
|
Tuple[str, List, str, List, Dict]
|
|
The name, shape, type, and shape name of the ValueInfoProto, and the
|
|
value_dict.
|
|
"""
|
|
shape = []
|
|
shape_name = []
|
|
for dim in info_proto.type.tensor_type.shape.dim:
|
|
name = dim.dim_param
|
|
value = dim.dim_value
|
|
if value is None or value == 0:
|
|
value = parse_shape_name(name, value_dict)
|
|
shape_name.append(name)
|
|
else:
|
|
shape_name.append(value)
|
|
shape.append(value)
|
|
|
|
name = info_proto.name
|
|
if info_proto.type.tensor_type.elem_type:
|
|
dtype = get_type(info_proto.type.tensor_type.elem_type)
|
|
else:
|
|
dtype = None
|
|
return name, shape, dtype, shape_name, value_dict
|
|
|
|
|
|
def get_numpy(tensor_proto: onnx.onnx_ml_pb2.TensorProto) -> _np.ndarray:
|
|
"""Grab data in TensorProto and convert to numpy array."""
|
|
try:
|
|
from onnx.numpy_helper import to_array # pylint: disable=import-outside-toplevel
|
|
except ImportError as exception:
|
|
raise ImportError(f"Unable to import onnx which is required {exception}")
|
|
return to_array(tensor_proto)
|
|
|
|
|
|
def get_prim_expr_list(
|
|
inputs: relax.Constant | relax.ShapeExpr,
|
|
) -> list[int | tirx.Expr]:
|
|
"""Attempt to convert a variable to list of Expr if possible.
|
|
|
|
Parameters
|
|
----------
|
|
inputs : Union[relax.Constant, relax.ShapeExpr, tvm.tirx.Expr]
|
|
The input value to try to convert to a list of Expr.
|
|
|
|
Returns
|
|
-------
|
|
ret : List[Union[int, tirx.Expr]]
|
|
The input value converted to a list of Expr if possible.
|
|
"""
|
|
if isinstance(inputs, relax.Constant):
|
|
np_value = inputs.data.numpy()
|
|
if np_value.ndim != 1:
|
|
raise ValueError(f"Cannot cast {type(inputs)} to list of Expr")
|
|
return np_value.tolist()
|
|
elif isinstance(inputs, relax.ShapeExpr):
|
|
return inputs.values
|
|
elif tvm.ir.is_prim_expr(inputs):
|
|
return [inputs]
|
|
else:
|
|
raise ValueError(f"Cannot cast {type(inputs)} to list of Expr")
|
|
|
|
|
|
class onnx_input(list): # pylint: disable=invalid-name
|
|
"""A list that returns None when out-of-bounds indices are accessed."""
|
|
|
|
def __getitem__(self, item):
|
|
if isinstance(item, slice):
|
|
if item.stop is None:
|
|
stop = len(self)
|
|
else:
|
|
stop = item.stop
|
|
indices = list(range(stop)[item])
|
|
return [self[i] for i in indices]
|
|
if isinstance(item, int):
|
|
return list(self)[item] if item < len(self) else None
|
|
raise TypeError(f"list indices must be integers or slices, not {type(item).__name__}")
|
|
|
|
|
|
# pylint: disable=invalid-name, len-as-condition, unused-argument, too-many-lines, redefined-builtin
|
|
class OnnxOpConverter:
|
|
"""A helper class for holding the common logic for ONNX op converters.
|
|
Each converter maps to a single ONNX op and defines the equivalent
|
|
functionality using Relax expressions. The converter can define multiple versions
|
|
of the op and the version is selected based on the opset version of the model.
|
|
"""
|
|
|
|
@classmethod
|
|
def get_converter(cls, opset):
|
|
"""Get converter matches given opset.
|
|
|
|
Parameters
|
|
----------
|
|
opset: int
|
|
opset from model.
|
|
|
|
Returns
|
|
-------
|
|
converter, which should be `_impl_vx`. Number x is the biggest
|
|
number smaller than or equal to opset belongs to all support versions.
|
|
"""
|
|
impl_versions = sorted(int(d.replace("_impl_v", "")) for d in dir(cls) if "_impl_v" in d)
|
|
# Select the largest implemented version that is <= opset.
|
|
# If opset is below all implementations, fall back to the smallest.
|
|
candidates = [v for v in impl_versions if v <= opset]
|
|
version = max(candidates) if candidates else impl_versions[0]
|
|
if hasattr(cls, f"_impl_v{version}"):
|
|
return getattr(cls, f"_impl_v{version}")
|
|
raise NotImplementedError(f"opset version {version} of {cls.__name__} not implemented")
|
|
|
|
|
|
class QuantizeLinear(OnnxOpConverter):
|
|
@classmethod
|
|
def _impl_v10(cls, bb, inputs, attr, params):
|
|
x, scale = inputs[0], inputs[1]
|
|
zp = inputs[2] if len(inputs) > 2 and inputs[2] is not None else None
|
|
axis = attr.get("axis", 1)
|
|
if hasattr(x.ty, "ndim") and x.ty.ndim <= 1 and axis == 1:
|
|
axis = 0
|
|
out_dtype = "uint8" if zp is None else zp.ty.dtype.dtype
|
|
if zp is None:
|
|
zp = relax.const(0, out_dtype)
|
|
return relax.op.quantize(x, scale, zp, axis=axis, out_dtype=out_dtype)
|
|
|
|
@classmethod
|
|
def _impl_v13(cls, bb, inputs, attr, params):
|
|
x, scale = inputs[0], inputs[1]
|
|
zp = inputs[2] if len(inputs) > 2 and inputs[2] is not None else None
|
|
axis = attr.get("axis", 1)
|
|
if hasattr(x.ty, "ndim") and x.ty.ndim <= 1 and axis == 1:
|
|
axis = 0
|
|
out_dtype = "uint8" if zp is None else zp.ty.dtype.dtype
|
|
if zp is None:
|
|
zp = relax.const(0, out_dtype)
|
|
return relax.op.quantize(x, scale, zp, axis=axis, out_dtype=out_dtype)
|
|
|
|
|
|
class DequantizeLinear(OnnxOpConverter):
|
|
@classmethod
|
|
def _impl_v10(cls, bb, inputs, attr, params):
|
|
x, scale = inputs[0], inputs[1]
|
|
zp = inputs[2] if len(inputs) > 2 and inputs[2] is not None else None
|
|
axis = attr.get("axis", 1)
|
|
if hasattr(x.ty, "ndim") and x.ty.ndim <= 1 and axis == 1:
|
|
axis = 0
|
|
if zp is None:
|
|
zp = relax.const(0, x.ty.dtype.dtype)
|
|
return relax.op.dequantize(x, scale, zp, axis=axis, out_dtype="float32")
|
|
|
|
@classmethod
|
|
def _impl_v13(cls, bb, inputs, attr, params):
|
|
x, scale = inputs[0], inputs[1]
|
|
zp = inputs[2] if len(inputs) > 2 and inputs[2] is not None else None
|
|
axis = attr.get("axis", 1)
|
|
if hasattr(x.ty, "ndim") and x.ty.ndim <= 1 and axis == 1:
|
|
axis = 0
|
|
if zp is None:
|
|
zp = relax.const(0, x.ty.dtype.dtype)
|
|
return relax.op.dequantize(x, scale, zp, axis=axis, out_dtype="float32")
|
|
|
|
|
|
class DynamicQuantizeLinear(OnnxOpConverter):
|
|
@classmethod
|
|
def _impl_v11(cls, bb, inputs, attr, params):
|
|
x = inputs[0]
|
|
x_dtype = x.ty.dtype.dtype
|
|
qmin = relax.const(0, x_dtype)
|
|
qmax = relax.const(255, x_dtype)
|
|
|
|
x_max = relax.op.maximum(qmin, relax.op.max(x))
|
|
x_min = relax.op.minimum(qmin, relax.op.min(x))
|
|
y_scale = relax.op.divide(relax.op.subtract(x_max, x_min), qmax)
|
|
|
|
zp_fp = relax.op.subtract(qmin, relax.op.divide(x_min, y_scale))
|
|
y_zero_point = relax.op.astype(relax.op.round(relax.op.clip(zp_fp, 0, 255)), "uint8")
|
|
|
|
y = relax.op.quantize(x, y_scale, y_zero_point, axis=0, out_dtype="uint8")
|
|
return relax.Tuple([y, y_scale, y_zero_point])
|
|
|
|
|
|
class MatMul(OnnxOpConverter):
|
|
"""Converts an onnx MatMul node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v13(cls, bb, inputs, attr, params):
|
|
return relax.op.matmul(inputs[0], inputs[1])
|
|
|
|
|
|
class MatMulInteger16(OnnxOpConverter):
|
|
"""Converts an ONNX MatMulInteger16 node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
if len(inputs) != 2:
|
|
raise ValueError(f"MatMulInteger16 expects two inputs, but got {len(inputs)}")
|
|
a, b = inputs
|
|
valid_types = ["int16", "uint16"]
|
|
a_dtype = a.ty.dtype.dtype
|
|
b_dtype = b.ty.dtype.dtype
|
|
if a_dtype not in valid_types:
|
|
raise ValueError(
|
|
"MatMulInteger16 expects input A to have int16 or uint16 dtype, "
|
|
f"but got {a.ty.dtype}"
|
|
)
|
|
if b_dtype not in valid_types:
|
|
raise ValueError(
|
|
"MatMulInteger16 expects input B to have int16 or uint16 dtype, "
|
|
f"but got {b.ty.dtype}"
|
|
)
|
|
|
|
out_dtype = "uint32" if a_dtype == "uint16" and b_dtype == "uint16" else "int32"
|
|
return relax.op.matmul(
|
|
relax.op.astype(a, out_dtype),
|
|
relax.op.astype(b, out_dtype),
|
|
)
|
|
|
|
|
|
def _to_numpy(x):
|
|
if tvm.ir.is_prim_expr(x):
|
|
if isinstance(x, tirx.IntImm | tirx.FloatImm):
|
|
return _np.array(x.value)
|
|
return x
|
|
else:
|
|
return x.data.numpy()
|
|
|
|
|
|
class _EmptyOptional:
|
|
"""Sentinel object that preserves an empty ONNX Optional during import."""
|
|
|
|
def __init__(self, type_proto: onnx.onnx_ml_pb2.TypeProto):
|
|
self.type_proto = type_proto
|
|
|
|
|
|
def _is_empty_optional(value: Any) -> bool:
|
|
"""Returns whether the given value represents an empty ONNX Optional."""
|
|
|
|
return isinstance(value, _EmptyOptional)
|
|
|
|
|
|
class BinaryBase(OnnxOpConverter):
|
|
"""Converts an onnx BinaryBase node into an equivalent Relax expression."""
|
|
|
|
numpy_op: Callable = None
|
|
relax_op: Callable = None
|
|
|
|
@classmethod
|
|
def base_impl(cls, bb, inputs, attr, params):
|
|
"""Base implementation for binary operations."""
|
|
if cls.numpy_op is None or cls.relax_op is None:
|
|
raise ValueError("Numpy and Relax operators must be defined for BinaryBase.")
|
|
if all([not isinstance(inp, tvm.ir.Call | relax.Var) for inp in inputs]):
|
|
has_prim_expr = any([tvm.ir.is_prim_expr(inp) for inp in inputs])
|
|
x = _to_numpy(inputs[0])
|
|
y = _to_numpy(inputs[1])
|
|
output = cls.numpy_op(x, y) # pylint: disable=not-callable
|
|
if has_prim_expr:
|
|
if hasattr(output, "item"):
|
|
output = output.item()
|
|
return relax.prim_value(output)
|
|
if x.dtype == y.dtype:
|
|
# no numpy precision widening
|
|
output = output.astype(x.dtype)
|
|
if all([isinstance(inp, relax.Constant) for inp in inputs]):
|
|
return relax.const(output, output.dtype) # pylint: disable=not-callable
|
|
|
|
return cls.relax_op(inputs[0], inputs[1]) # pylint: disable=not-callable
|
|
|
|
|
|
class Add(BinaryBase):
|
|
"""Converts an onnx Add node into an equivalent Relax expression."""
|
|
|
|
numpy_op = _np.add
|
|
relax_op = relax.op.add
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
return cls.base_impl(bb, inputs, attr, params)
|
|
|
|
|
|
class Sub(BinaryBase):
|
|
"""Converts an onnx Sub node into an equivalent Relax expression."""
|
|
|
|
numpy_op = _np.subtract
|
|
relax_op = relax.op.subtract
|
|
|
|
@classmethod
|
|
def _impl_v7(cls, bb, inputs, attr, params):
|
|
return cls.base_impl(bb, inputs, attr, params)
|
|
|
|
|
|
class Mul(BinaryBase):
|
|
"""Converts an onnx Mul node into an equivalent Relax expression."""
|
|
|
|
numpy_op = _np.multiply
|
|
relax_op = relax.op.multiply
|
|
|
|
@classmethod
|
|
def _impl_v7(cls, bb, inputs, attr, params):
|
|
return cls.base_impl(bb, inputs, attr, params)
|
|
|
|
|
|
class Div(BinaryBase):
|
|
"""Converts an onnx Div node into an equivalent Relax expression."""
|
|
|
|
numpy_op = _np.divide
|
|
relax_op = relax.op.divide
|
|
|
|
@classmethod
|
|
def _impl_v7(cls, bb, inputs, attr, params):
|
|
try:
|
|
lhs_code = DataType(inputs[0].ty.dtype.dtype).type_code
|
|
rhs_code = DataType(inputs[1].ty.dtype.dtype).type_code
|
|
except (AttributeError, ValueError, TypeError, RuntimeError):
|
|
return cls.base_impl(bb, inputs, attr, params)
|
|
|
|
lhs_is_integer = lhs_code == DataTypeCode.INT or lhs_code == DataTypeCode.UINT
|
|
rhs_is_integer = rhs_code == DataTypeCode.INT or rhs_code == DataTypeCode.UINT
|
|
if not (lhs_is_integer and rhs_is_integer):
|
|
return cls.base_impl(bb, inputs, attr, params)
|
|
|
|
if isinstance(inputs[1], relax.Constant) and bool(_np.any(inputs[1].data.numpy() == 0)):
|
|
raise ValueError("ONNX Div with integer inputs encountered divisor value 0.")
|
|
|
|
return cls.base_impl(bb, inputs, attr, params)
|
|
|
|
|
|
class Pow(BinaryBase):
|
|
"""Converts an onnx Pow node into an equivalent Relax expression."""
|
|
|
|
numpy_op = _np.power
|
|
relax_op = relax.op.power
|
|
|
|
@classmethod
|
|
def _impl_v7(cls, bb, inputs, attr, params):
|
|
return cls.base_impl(bb, inputs, attr, params)
|
|
|
|
|
|
class Mod(BinaryBase):
|
|
"""Converts an onnx Mod node into an equivalent Relax expression."""
|
|
|
|
numpy_op = _np.mod
|
|
relax_op = relax.op.mod
|
|
|
|
@classmethod
|
|
def _impl_v10(cls, bb, inputs, attr, params):
|
|
if attr.get("fmod", 0) == 0:
|
|
cls.numpy_op = _np.fmod
|
|
cls.relax_op = relax.op.floor_mod
|
|
else:
|
|
cls.numpy_op = _np.mod
|
|
cls.relax_op = relax.op.mod
|
|
return cls.base_impl(bb, inputs, attr, params)
|
|
|
|
|
|
class And(BinaryBase):
|
|
"""Converts an onnx And node into an equivalent Relax expression."""
|
|
|
|
numpy_op = _np.logical_and
|
|
relax_op = relax.op.logical_and
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
return cls.base_impl(bb, inputs, attr, params)
|
|
|
|
|
|
class Or(BinaryBase):
|
|
"""Converts an onnx Or node into an equivalent Relax expression."""
|
|
|
|
numpy_op = _np.logical_or
|
|
relax_op = relax.op.logical_or
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
return cls.base_impl(bb, inputs, attr, params)
|
|
|
|
|
|
class Xor(BinaryBase):
|
|
"""Converts an onnx Xor node into an equivalent Relax expression."""
|
|
|
|
numpy_op = _np.logical_xor
|
|
relax_op = relax.op.logical_xor
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
return cls.base_impl(bb, inputs, attr, params)
|
|
|
|
|
|
class Less(BinaryBase):
|
|
"""Converts an onnx Less node into an equivalent Relax expression."""
|
|
|
|
numpy_op = _np.less
|
|
relax_op = relax.op.less
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
return cls.base_impl(bb, inputs, attr, params)
|
|
|
|
|
|
class LessOrEqual(BinaryBase):
|
|
"""Converts an onnx LessEqual node into an equivalent Relax expression."""
|
|
|
|
numpy_op = _np.less_equal
|
|
relax_op = relax.op.less_equal
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
return cls.base_impl(bb, inputs, attr, params)
|
|
|
|
|
|
class Greater(BinaryBase):
|
|
"""Converts an onnx Greater node into an equivalent Relax expression."""
|
|
|
|
numpy_op = _np.greater
|
|
relax_op = relax.op.greater
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
return cls.base_impl(bb, inputs, attr, params)
|
|
|
|
|
|
class GreaterOrEqual(BinaryBase):
|
|
"""Converts an onnx GreaterEqual node into an equivalent Relax expression."""
|
|
|
|
numpy_op = _np.greater_equal
|
|
relax_op = relax.op.greater_equal
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
return cls.base_impl(bb, inputs, attr, params)
|
|
|
|
|
|
class Equal(OnnxOpConverter):
|
|
"""Converts an onnx Equal node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v13(cls, bb, inputs, attr, params):
|
|
if all([isinstance(inp, relax.Constant) for inp in inputs]):
|
|
output = inputs[0].data.numpy() == inputs[1].data.numpy()
|
|
return relax.const(output, output.dtype)
|
|
elif all([isinstance(inp, relax.Constant | relax.ShapeExpr) for inp in inputs]):
|
|
lhs = get_prim_expr_list(inputs[0])
|
|
rhs = get_prim_expr_list(inputs[1])
|
|
if len(lhs) != len(rhs):
|
|
raise ValueError("Cannot compare two tensors with different shapes")
|
|
output = [tvm_ffi.structural_equal(l, r) for l, r in zip(lhs, rhs)]
|
|
return relax.const(output, "bool")
|
|
return relax.op.equal(inputs[0], inputs[1])
|
|
|
|
|
|
class BitwiseBase(BinaryBase):
|
|
"""Converts an onnx BitwiseBase node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def base_impl(cls, bb, inputs, attr, params):
|
|
"""Base implementation for bitwise operations."""
|
|
valid_types = ["int8", "int16", "int32", "int64", "uint8", "uint16", "uint32", "uint64"]
|
|
for num, inp in enumerate(inputs):
|
|
if inp.ty.dtype.dtype not in valid_types:
|
|
raise ValueError(
|
|
f"Bitwise operations expect all inputs to have integer types, "
|
|
f"got {inp.ty.dtype} for input {num}"
|
|
)
|
|
return super().base_impl(bb, inputs, attr, params)
|
|
|
|
|
|
class BitwiseAnd(BitwiseBase):
|
|
"""Converts an onnx BitwiseAnd node into an equivalent Relax expression."""
|
|
|
|
numpy_op = _np.bitwise_and
|
|
relax_op = relax.op.bitwise_and
|
|
|
|
@classmethod
|
|
def _impl_v18(cls, bb, inputs, attr, params):
|
|
return cls.base_impl(bb, inputs, attr, params)
|
|
|
|
|
|
class BitwiseOr(BitwiseBase):
|
|
"""Converts an onnx BitwiseOr node into an equivalent Relax expression."""
|
|
|
|
numpy_op = _np.bitwise_or
|
|
relax_op = relax.op.bitwise_or
|
|
|
|
@classmethod
|
|
def _impl_v18(cls, bb, inputs, attr, params):
|
|
return cls.base_impl(bb, inputs, attr, params)
|
|
|
|
|
|
class BitwiseXor(BitwiseBase):
|
|
"""Converts an onnx BitwiseXor node into an equivalent Relax expression."""
|
|
|
|
numpy_op = _np.bitwise_xor
|
|
relax_op = relax.op.bitwise_xor
|
|
|
|
@classmethod
|
|
def _impl_v18(cls, bb, inputs, attr, params):
|
|
return cls.base_impl(bb, inputs, attr, params)
|
|
|
|
|
|
class BitwiseNot(OnnxOpConverter):
|
|
"""Converts an onnx BitwiseNot node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v18(cls, bb, inputs, attr, params):
|
|
if isinstance(inputs[0], relax.Constant):
|
|
return relax.const(_np.bitwise_not(inputs[0].data.numpy()), inputs[0].ty.dtype)
|
|
return relax.op.bitwise_not(inputs[0])
|
|
|
|
|
|
class BitShift(BitwiseBase):
|
|
"""Converts an onnx BitShift node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v11(cls, bb, inputs, attr, params):
|
|
direction = attr.get("direction", "LEFT").decode("ascii")
|
|
if direction == "LEFT":
|
|
cls.numpy_op = _np.left_shift
|
|
cls.relax_op = relax.op.left_shift
|
|
elif direction == "RIGHT":
|
|
cls.numpy_op = _np.right_shift
|
|
cls.relax_op = relax.op.right_shift
|
|
else:
|
|
raise ValueError("Unsupported Shift Direction: " + direction)
|
|
|
|
return cls.base_impl(bb, inputs, attr, params)
|
|
|
|
|
|
class Sigmoid(OnnxOpConverter):
|
|
"""Converts an onnx Sigmoid node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v13(cls, bb, inputs, attr, params):
|
|
return relax.op.sigmoid(inputs[0])
|
|
|
|
|
|
def _normalize_legacy_softmax_axis(axis: int, rank: int, op_name: str) -> int:
|
|
"""Normalize axis for ONNX Softmax/LogSoftmax/Hardmax opset <= 12 semantics.
|
|
|
|
Legacy semantics allow axis in [-rank, rank], where axis == rank means the
|
|
last dimension after flattening has extent 1.
|
|
"""
|
|
|
|
if axis < -rank or axis > rank:
|
|
raise ValueError(f"{op_name} axis {axis} is out of range for rank {rank}.")
|
|
if axis < 0:
|
|
axis += rank
|
|
return axis
|
|
|
|
|
|
def _shape_product(dims: list[int | tirx.Expr]) -> int | tirx.Expr:
|
|
"""Compute product of a list of shape dims (supports symbolic dims)."""
|
|
|
|
prod = 1
|
|
for dim in dims:
|
|
if isinstance(dim, tirx.IntImm):
|
|
dim = int(dim.value)
|
|
prod = prod * dim
|
|
return prod
|
|
|
|
|
|
def _legacy_softmax_prepare(
|
|
data: relax.Expr, axis: int, op_name: str
|
|
) -> tuple[relax.Expr, tuple[int | tirx.Expr, ...]] | None:
|
|
"""Build legacy 2D view for Softmax-family opset <= 12 semantics.
|
|
|
|
Returns (reshaped_data, original_shape). If rank/shape isn't statically
|
|
available, returns None so caller can choose a permissive fallback.
|
|
"""
|
|
|
|
rank = _get_known_tensor_rank(data)
|
|
if rank is None:
|
|
return None
|
|
|
|
axis = _normalize_legacy_softmax_axis(axis, rank, op_name)
|
|
ty = data.ty
|
|
if not isinstance(ty, relax.TensorType):
|
|
return None
|
|
if not isinstance(ty.shape, relax.ShapeExpr):
|
|
return None
|
|
|
|
original_shape = list(ty.shape.values)
|
|
if len(original_shape) != rank:
|
|
return None
|
|
|
|
dim0 = _shape_product(original_shape[:axis])
|
|
dim1 = _shape_product(original_shape[axis:])
|
|
flattened = relax.op.reshape(data, (dim0, dim1))
|
|
return flattened, tuple(original_shape)
|
|
|
|
|
|
def _get_axis_extent(data: relax.Expr, axis: int, op_name: str) -> tuple[int, int | tirx.Expr]:
|
|
"""Return normalized axis and axis extent when rank/shape are known."""
|
|
|
|
rank = _get_known_tensor_rank(data)
|
|
if rank is None:
|
|
raise ValueError(f"{op_name} requires a statically known input rank.")
|
|
|
|
normalized_axis = _normalize_constant_axes([axis], rank, op_name)[0]
|
|
ty = data.ty
|
|
if isinstance(ty, relax.TensorType) and isinstance(ty.shape, relax.ShapeExpr):
|
|
axis_extent = ty.shape.values[normalized_axis]
|
|
if isinstance(axis_extent, tirx.IntImm):
|
|
axis_extent = int(axis_extent.value)
|
|
return normalized_axis, axis_extent
|
|
|
|
raise ValueError(f"{op_name} requires a statically known axis extent.")
|
|
|
|
|
|
class Softmax(OnnxOpConverter):
|
|
"""Converts an onnx Softmax node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
axis = attr.get("axis", 1)
|
|
prepared = _legacy_softmax_prepare(inputs[0], axis, "Softmax")
|
|
if prepared is None:
|
|
warnings.warn(
|
|
"Softmax opset<=12 fallback: static rank/shape is unavailable, "
|
|
"falling back to axis-based softmax semantics."
|
|
)
|
|
return relax.op.nn.softmax(inputs[0], axis=axis)
|
|
|
|
flattened, original_shape = prepared
|
|
out = relax.op.nn.softmax(flattened, axis=-1)
|
|
return relax.op.reshape(out, original_shape)
|
|
|
|
_impl_v11 = _impl_v1
|
|
|
|
@classmethod
|
|
def _impl_v13(cls, bb, inputs, attr, params):
|
|
axis = attr.get("axis", -1)
|
|
return relax.op.nn.softmax(inputs[0], axis=axis)
|
|
|
|
|
|
class LogSoftmax(OnnxOpConverter):
|
|
"""Converts an onnx LogSoftmax node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
axis = attr.get("axis", 1)
|
|
prepared = _legacy_softmax_prepare(inputs[0], axis, "LogSoftmax")
|
|
if prepared is None:
|
|
warnings.warn(
|
|
"LogSoftmax opset<=12 fallback: static rank/shape is unavailable, "
|
|
"falling back to axis-based log_softmax semantics."
|
|
)
|
|
return relax.op.nn.log_softmax(inputs[0], axis=axis)
|
|
|
|
flattened, original_shape = prepared
|
|
out = relax.op.nn.log_softmax(flattened, axis=-1)
|
|
return relax.op.reshape(out, original_shape)
|
|
|
|
_impl_v11 = _impl_v1
|
|
|
|
@classmethod
|
|
def _impl_v13(cls, bb, inputs, attr, params):
|
|
axis = attr.get("axis", -1)
|
|
return relax.op.nn.log_softmax(inputs[0], axis=axis)
|
|
|
|
|
|
class Hardmax(OnnxOpConverter):
|
|
"""Converts an onnx Hardmax node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _hardmax_impl(cls, *args):
|
|
"""Hardmax core implementation.
|
|
|
|
Compatibility note:
|
|
- New signature: _hardmax_impl(bb, data, axis)
|
|
- Legacy signature: _hardmax_impl(data, axis)
|
|
"""
|
|
if len(args) == 3:
|
|
bb, data, axis = args
|
|
elif len(args) == 2:
|
|
bb = None
|
|
data, axis = args
|
|
else:
|
|
raise TypeError("Hardmax._hardmax_impl expects (bb, data, axis) or (data, axis).")
|
|
|
|
if bb is not None:
|
|
data = bb.normalize(data)
|
|
normalized_axis, axis_extent = _get_axis_extent(data, axis, "Hardmax")
|
|
dtype = data.ty.dtype
|
|
argmax = relax.op.argmax(data, axis=normalized_axis)
|
|
on_value = relax.prim_value(tvm.tirx.const(1.0, dtype))
|
|
off_value = relax.prim_value(tvm.tirx.const(0.0, dtype))
|
|
return relax.op.one_hot(argmax, on_value, off_value, axis_extent, normalized_axis)
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
axis = attr.get("axis", 1)
|
|
prepared = _legacy_softmax_prepare(inputs[0], axis, "Hardmax")
|
|
if prepared is None:
|
|
warnings.warn(
|
|
"Hardmax opset<=12 fallback: static rank/shape is unavailable, "
|
|
"falling back to axis-based hardmax semantics."
|
|
)
|
|
hardmax_input = inputs[0]
|
|
hardmax_axis = axis
|
|
original_shape = None
|
|
else:
|
|
hardmax_input, original_shape = prepared
|
|
hardmax_axis = -1
|
|
|
|
out = cls._hardmax_impl(bb, hardmax_input, hardmax_axis)
|
|
return out if original_shape is None else relax.op.reshape(out, original_shape)
|
|
|
|
_impl_v11 = _impl_v1
|
|
|
|
@classmethod
|
|
def _impl_v13(cls, bb, inputs, attr, params):
|
|
axis = attr.get("axis", -1)
|
|
return cls._hardmax_impl(bb, inputs[0], axis)
|
|
|
|
|
|
class Transpose(OnnxOpConverter):
|
|
"""Converts an onnx Transpose node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v13(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
axes = attr.get("perm", None)
|
|
|
|
if hasattr(data.ty, "ndim"):
|
|
input_ndim = data.ty.ndim
|
|
elif hasattr(data.ty, "shape") and data.ty.shape:
|
|
input_ndim = len(data.ty.shape)
|
|
else:
|
|
if isinstance(data, relax.Constant):
|
|
input_ndim = data.data.numpy().ndim
|
|
else:
|
|
input_ndim = None
|
|
|
|
if input_ndim == 0:
|
|
return data
|
|
|
|
if input_ndim is not None and axes is not None:
|
|
if len(axes) != input_ndim:
|
|
raise ValueError(
|
|
f"Transpose: number of axes in perm attribute ({len(axes)}) "
|
|
f"must equal the number of input tensor dimensions ({input_ndim})"
|
|
)
|
|
|
|
if isinstance(data, relax.Constant):
|
|
output = _np.transpose(data.data.numpy(), axes)
|
|
return relax.const(output, output.dtype)
|
|
|
|
return relax.op.permute_dims(data, axes)
|
|
|
|
|
|
class Unsqueeze(OnnxOpConverter):
|
|
"""Converts an onnx Unsqueeze node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
axes = list(attr.get("axes"))
|
|
inputs = inputs + [relax.const(axes, "int64")]
|
|
return cls._impl_v13(bb, inputs, attr, params)
|
|
|
|
@classmethod
|
|
def _impl_v13(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
axes = get_constant(inputs[1], params)
|
|
data_ndim = _get_known_tensor_rank(data)
|
|
|
|
if tvm.ir.is_prim_expr(data) and isinstance(axes, relax.Constant):
|
|
constant_axes = _normalize_constant_axes(
|
|
list(map(int, axes.data.numpy().tolist())), 1, "Unsqueeze"
|
|
)
|
|
if constant_axes == [0]:
|
|
return relax.ShapeExpr([data])
|
|
raise NotImplementedError("Unsqueeze with symbolic scalar inputs only supports axis 0.")
|
|
if isinstance(data, relax.Constant) and isinstance(axes, relax.Constant):
|
|
constant_axes = _normalize_constant_axes(
|
|
list(map(int, axes.data.numpy().tolist())),
|
|
data.data.numpy().ndim + axes.data.numpy().size,
|
|
"Unsqueeze",
|
|
)
|
|
constant_axes = sorted(constant_axes)
|
|
expanded = data.data.numpy()
|
|
output_rank = expanded.ndim + len(constant_axes)
|
|
new_shape = []
|
|
input_dims_iter = iter(expanded.shape)
|
|
for i in range(output_rank):
|
|
if i in constant_axes:
|
|
new_shape.append(1)
|
|
else:
|
|
new_shape.append(next(input_dims_iter))
|
|
expanded = expanded.reshape(new_shape)
|
|
return relax.const(expanded, data.ty.dtype)
|
|
|
|
if isinstance(axes, relax.Constant):
|
|
if data_ndim is None:
|
|
raise ValueError("Unsqueeze requires a statically known input rank.")
|
|
constant_axes = _normalize_constant_axes(
|
|
list(map(int, axes.data.numpy().tolist())),
|
|
data_ndim + axes.data.numpy().size,
|
|
"Unsqueeze",
|
|
)
|
|
constant_axes = sorted(constant_axes)
|
|
for axis in constant_axes:
|
|
data = relax.op.expand_dims(data, axis=axis)
|
|
return data
|
|
|
|
if data_ndim is None:
|
|
raise ValueError("Unsqueeze with dynamic axes requires a statically known input rank.")
|
|
axes_len = _get_known_tensor_length(axes)
|
|
if axes_len is None:
|
|
raise ValueError("Unsqueeze requires a statically known axes length.")
|
|
data_shape = bb.normalize(relax.op.shape_of(data))
|
|
data_shape_tensor = bb.normalize(relax.op.shape_to_tensor(data_shape))
|
|
output_shape_tensor = _build_unsqueezed_shape_tensor(bb, data_shape_tensor, axes, data_ndim)
|
|
output_shape = _tensor_to_shape_expr(
|
|
bb, output_shape_tensor, data_ndim + axes_len, "unsqueeze_dim"
|
|
)
|
|
return relax.op.reshape(data, output_shape)
|
|
|
|
|
|
class Concat(OnnxOpConverter):
|
|
"""Convert an onnx Concat node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v13(cls, bb, inputs, attr, params):
|
|
axis = attr.get("axis", 0)
|
|
_, param_dict = params
|
|
|
|
def is_shape_like(x: Any) -> bool:
|
|
if isinstance(x, relax.ShapeExpr):
|
|
return True
|
|
elif isinstance(x, relax.Constant):
|
|
return x.ty.ndim == 1 and x.ty.dtype == "int64"
|
|
else:
|
|
return False
|
|
|
|
# Resolve 1D-int64 param Vars to constants only for the shape-like
|
|
# fast path; tensor fallback keeps the original Vars so runtime
|
|
# weights aren't folded under keep_params_in_input=True.
|
|
def resolve(x):
|
|
if isinstance(x, relax.Var) and x.name_hint in param_dict:
|
|
arr = param_dict[x.name_hint][1].numpy()
|
|
if arr.ndim == 1 and arr.dtype == _np.int64:
|
|
return relax.const(arr, "int64")
|
|
return x
|
|
|
|
resolved = [resolve(inp) for inp in inputs]
|
|
|
|
# If all inputs are shape expr, perform computation directly.
|
|
if all([is_shape_like(inp) for inp in resolved]):
|
|
const_inputs = []
|
|
for inp in resolved:
|
|
if isinstance(inp, relax.ShapeExpr):
|
|
const_inputs.extend(inp.values)
|
|
elif isinstance(inp, relax.Constant):
|
|
const_inputs.extend(inp.data.numpy().tolist())
|
|
else:
|
|
raise NotImplementedError(f"Unsupported input type: {type(inp)}")
|
|
return relax.ShapeExpr(const_inputs)
|
|
|
|
# If all inputs are constant, perform computation directly.
|
|
if all([isinstance(inp, relax.Constant) for inp in inputs]):
|
|
const_inputs = []
|
|
for inp in inputs:
|
|
const_inputs.append(inp.data.numpy())
|
|
out = _np.concatenate(const_inputs, axis=axis)
|
|
dtype = inputs[0].ty.dtype
|
|
return relax.const(out, dtype)
|
|
|
|
return relax.op.concat(inputs, axis=axis)
|
|
|
|
|
|
class Cast(OnnxOpConverter):
|
|
"""Convert an onnx Cast node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v13(cls, bb, inputs, attr, params):
|
|
to_type = get_type(attr["to"])
|
|
if isinstance(inputs[0], relax.ShapeExpr):
|
|
shape = inputs[0]
|
|
if all([isinstance(x, tirx.IntImm) for x in shape]):
|
|
shape = [int(x) for x in shape]
|
|
return relax.const(shape, to_type)
|
|
if isinstance(inputs[0], relax.Constant):
|
|
output = inputs[0].data.numpy().astype(to_type)
|
|
return relax.const(output, to_type)
|
|
if tvm.ir.is_prim_expr(inputs[0]):
|
|
if isinstance(inputs[0], tirx.IntImm | tirx.FloatImm):
|
|
return tvm.tirx.const(inputs[0].value, to_type)
|
|
return inputs[0].astype(to_type)
|
|
|
|
try:
|
|
np_dst = _np.dtype(str(to_type))
|
|
except Exception:
|
|
return relax.op.astype(inputs[0], to_type)
|
|
|
|
if np_dst.kind in ("i", "u"):
|
|
src = inputs[0]
|
|
src_dtype = getattr(getattr(src, "ty", None), "dtype", None) or getattr(
|
|
src, "dtype", None
|
|
)
|
|
if src_dtype is not None and _relax_dtype_is_floating_point(src_dtype):
|
|
x_sanitized = bb.emit(
|
|
relax.op.where(
|
|
relax.op.logical_not(relax.op.isfinite(src)),
|
|
relax.const(0.0, src_dtype),
|
|
src,
|
|
)
|
|
)
|
|
dst_str = str(to_type)
|
|
if dst_str.startswith("uint"):
|
|
signed = False
|
|
bits = int(dst_str[4:])
|
|
elif dst_str.startswith("int"):
|
|
signed = True
|
|
bits = int(dst_str[3:])
|
|
else:
|
|
return relax.op.astype(x_sanitized, to_type)
|
|
|
|
if bits == 64:
|
|
return relax.op.astype(x_sanitized, to_type)
|
|
|
|
temp_dtype = "int64" if bits >= 32 else "int32"
|
|
t = relax.op.astype(x_sanitized, temp_dtype)
|
|
if bits == 32:
|
|
two_pow = relax.const(1 << bits, temp_dtype)
|
|
uw = relax.op.floor_mod(t, two_pow)
|
|
else:
|
|
mask_val = (1 << bits) - 1
|
|
mask = relax.const(mask_val, temp_dtype)
|
|
uw = relax.op.bitwise_and(t, mask)
|
|
if signed:
|
|
half = 1 << (bits - 1)
|
|
half_c = relax.const(half, temp_dtype)
|
|
if bits == 32:
|
|
two_pow = relax.const(1 << bits, temp_dtype)
|
|
else:
|
|
two_pow = relax.op.add(mask, relax.const(1, temp_dtype))
|
|
wrapped = relax.op.where(
|
|
relax.op.greater_equal(uw, half_c),
|
|
relax.op.subtract(uw, two_pow),
|
|
uw,
|
|
)
|
|
else:
|
|
wrapped = uw
|
|
return relax.op.astype(wrapped, to_type)
|
|
|
|
return relax.op.astype(inputs[0], to_type)
|
|
|
|
|
|
class Gather(OnnxOpConverter):
|
|
"""Convert an onnx Gather node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v13(cls, bb, inputs, attr, params):
|
|
# Unpack inputs
|
|
data = inputs[0]
|
|
indices = inputs[1]
|
|
axis = attr.get("axis", 0)
|
|
|
|
# If all inputs are constant, we can compute directly.
|
|
if all([isinstance(inp, relax.Constant) for inp in [data, indices]]):
|
|
output = _np.take(data.data.numpy(), indices.data.numpy(), axis=axis)
|
|
return relax.const(output, output.dtype)
|
|
|
|
# If input is a shape expression, take a value from that shape and return it as a constant.
|
|
if isinstance(data, relax.ShapeExpr):
|
|
assert isinstance(indices, relax.Constant), (
|
|
"Only constant indices supported for shape gather."
|
|
)
|
|
np_index = indices.data.numpy()
|
|
if len(np_index.shape) == 1:
|
|
np_index = np_index[0]
|
|
np_index = int(np_index)
|
|
shape_val = data[np_index]
|
|
return relax.prim_value(shape_val)
|
|
|
|
indices_dtype = indices.ty.dtype.dtype
|
|
if not indices_dtype.startswith("uint"):
|
|
data_shape = bb.normalize(relax.op.shape_of(data))
|
|
data_shape_tensor = bb.normalize(relax.op.shape_to_tensor(data_shape))
|
|
axis_extent = bb.normalize(
|
|
relax.op.take(data_shape_tensor, relax.const(axis, "int64"), axis=0, mode="wrap")
|
|
)
|
|
|
|
if indices_dtype != "int64":
|
|
axis_extent = bb.normalize(relax.op.astype(axis_extent, indices_dtype))
|
|
|
|
indices = bb.normalize(
|
|
relax.op.where(
|
|
relax.op.less(indices, relax.const(0, indices_dtype)),
|
|
relax.op.add(indices, axis_extent),
|
|
indices,
|
|
)
|
|
)
|
|
|
|
return relax.op.take(data, indices, axis)
|
|
|
|
|
|
class GatherElements(OnnxOpConverter):
|
|
"""Convert an onnx GatherElements node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v13(cls, bb, inputs, attr, params):
|
|
axis = attr.get("axis", 0)
|
|
return relax.op.gather_elements(inputs[0], inputs[1], axis)
|
|
|
|
|
|
class GatherND(OnnxOpConverter):
|
|
"""Convert an onnx GatherND node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v13(cls, bb, inputs, attr, params):
|
|
batch_dims = attr.get("batch_dims", 0)
|
|
return relax.op.gather_nd(inputs[0], inputs[1], batch_dims)
|
|
|
|
|
|
class Scatter(OnnxOpConverter):
|
|
"""Convert an onnx Scatter node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v9(cls, bb, inputs, attr, params):
|
|
axis = attr.get("axis", 0)
|
|
return relax.op.scatter_elements(inputs[0], inputs[1], inputs[2], axis=axis)
|
|
|
|
@classmethod
|
|
def _impl_v11(cls, bb, inputs, attr, params):
|
|
raise ValueError("Scatter is deprecated in ONNX 11")
|
|
|
|
|
|
def _get_onnx_reduction(attr, valid_reductions: list[str]):
|
|
reduction = attr.get("reduction", None)
|
|
reduction = reduction or b"update"
|
|
if isinstance(reduction, bytes):
|
|
reduction = reduction.decode("utf-8")
|
|
reduction = "update" if reduction == "none" else reduction
|
|
if reduction not in valid_reductions:
|
|
raise ValueError(f"Only {valid_reductions} reductions are supported, but got {reduction}")
|
|
|
|
return reduction
|
|
|
|
|
|
class ScatterElements(OnnxOpConverter):
|
|
"""Convert an onnx ScatterElements node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v11(cls, bb, inputs, attr, params):
|
|
axis = attr.get("axis", 0)
|
|
return relax.op.scatter_elements(inputs[0], inputs[1], inputs[2], axis=axis)
|
|
|
|
@classmethod
|
|
def _impl_v16(cls, bb, inputs, attr, params):
|
|
axis = attr.get("axis", 0)
|
|
reduction = _get_onnx_reduction(attr, ["update", "add", "mul"])
|
|
return relax.op.scatter_elements(
|
|
inputs[0], inputs[1], inputs[2], axis=axis, reduction=reduction
|
|
)
|
|
|
|
@classmethod
|
|
def _impl_v18(cls, bb, inputs, attr, params):
|
|
axis = attr.get("axis", 0)
|
|
reduction = _get_onnx_reduction(attr, ["update", "add", "mul", "min", "max"])
|
|
return relax.op.scatter_elements(
|
|
inputs[0], inputs[1], inputs[2], axis=axis, reduction=reduction
|
|
)
|
|
|
|
|
|
class ScatterND(OnnxOpConverter):
|
|
"""Convert an onnx ScatterND node into an equivalent Relax expression."""
|
|
|
|
@staticmethod
|
|
def _reduction_check(attr, valid_reductions: list[str]):
|
|
return _get_onnx_reduction(attr, valid_reductions)
|
|
|
|
@classmethod
|
|
def _impl_v11(cls, bb, inputs, attr, params):
|
|
return relax.op.scatter_nd(inputs[0], inputs[1], inputs[2])
|
|
|
|
@classmethod
|
|
def _impl_v16(cls, bb, inputs, attr, params):
|
|
reduction = cls._reduction_check(attr, ["update", "add", "mul"])
|
|
return relax.op.scatter_nd(inputs[0], inputs[1], inputs[2], reduction)
|
|
|
|
@classmethod
|
|
def _impl_v18(cls, bb, inputs, attr, params):
|
|
reduction = cls._reduction_check(attr, ["update", "add", "mul", "min", "max"])
|
|
return relax.op.scatter_nd(inputs[0], inputs[1], inputs[2], reduction)
|
|
|
|
|
|
class Compress(OnnxOpConverter):
|
|
"""Convert an onnx Compress node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v11(cls, bb, inputs, attr, params):
|
|
tensor, condition = inputs
|
|
axis = attr.get("axis", None)
|
|
|
|
# Change one hot tensor to indices e.g. [0, 1, 1, 0, 1] -> [1, 2, 4]
|
|
if condition.ty.dtype != "bool":
|
|
raise ValueError("Condition tensor is expected to be a boolean tensor")
|
|
if condition.ty.ndim != 1:
|
|
raise ValueError("Condition tensor is expected to be a 1D boolean tensor")
|
|
indices = relax.op.nonzero(condition)
|
|
num_nonzero = tirx.Var("num_nonzero", "int64")
|
|
indices = bb.match_cast(indices, relax.TensorType([1, num_nonzero], "int64"))
|
|
indices = relax.op.reshape(indices, [-1])
|
|
|
|
if axis is not None:
|
|
return relax.op.take(tensor, indices, axis=axis)
|
|
|
|
# if axis is None, flatten input tensor before selection
|
|
tensor = relax.op.reshape(tensor, (-1,))
|
|
return relax.op.take(tensor, indices, axis=0)
|
|
|
|
|
|
class Size(OnnxOpConverter):
|
|
"""Convert an onnx Size node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
return relax.op.size(inputs[0])
|
|
|
|
|
|
class EyeLike(OnnxOpConverter):
|
|
"""Convert an onnx EyeLike node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v9(cls, bb, inputs, attr, params):
|
|
k = attr.get("k", 0)
|
|
input_dtype = inputs[0].ty.dtype.dtype
|
|
if "dtype" in attr and get_type(attr["dtype"]) != input_dtype:
|
|
raise ValueError(
|
|
f"dtype mismatch between input ({input_dtype}) and attribute ({attr['dtype']})"
|
|
)
|
|
return relax.op.eye_like(inputs[0], k, input_dtype)
|
|
|
|
|
|
class Gemm(OnnxOpConverter):
|
|
"""Convert an onnx Gemm node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v13(cls, bb, inputs, attr, params):
|
|
alpha = attr.get("alpha", None)
|
|
beta = attr.get("beta", None)
|
|
transA = attr.get("transA", False)
|
|
transB = attr.get("transB", False)
|
|
A = inputs[0]
|
|
B = inputs[1]
|
|
C = inputs[2]
|
|
dtype = A.ty.dtype
|
|
|
|
# Compute Y = alpha * A X B + beta * C
|
|
|
|
if alpha is not None and alpha != 1.0:
|
|
A = relax.op.multiply(A, relax.const(alpha, dtype=dtype))
|
|
|
|
if transA:
|
|
A = relax.op.permute_dims(A, [1, 0])
|
|
if transB:
|
|
B = relax.op.permute_dims(B, [1, 0])
|
|
Y = relax.op.matmul(A, B)
|
|
|
|
if C is not None:
|
|
if beta is not None and beta != 1.0:
|
|
C = relax.op.multiply(C, relax.const(beta, dtype=dtype))
|
|
Y = relax.op.add(Y, C)
|
|
|
|
return Y
|
|
|
|
|
|
class Reshape(OnnxOpConverter):
|
|
"""Convert an onnx Reshape node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v13(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
new_shape = get_constant(inputs[1], params)
|
|
|
|
if isinstance(data, relax.ShapeExpr):
|
|
# Preserve identity flatten for shape values to keep shape-specialized
|
|
# handling in downstream shape-construction patterns.
|
|
if isinstance(new_shape, relax.Constant):
|
|
new_shape_values = new_shape.data.numpy().tolist()
|
|
if new_shape_values == [-1]:
|
|
return data
|
|
|
|
# Other reshape targets follow regular int64 tensor reshape semantics.
|
|
data = bb.normalize(relax.op.shape_to_tensor(data))
|
|
|
|
if isinstance(data, relax.Constant) and isinstance(new_shape, relax.Constant):
|
|
out = _np.reshape(data.data.numpy(), new_shape.data.numpy().tolist())
|
|
return relax.const(out, out.dtype)
|
|
if isinstance(new_shape, relax.Constant):
|
|
new_shape = new_shape.data.numpy().tolist()
|
|
out = relax.op.reshape(data, new_shape)
|
|
return out
|
|
|
|
|
|
class Where(OnnxOpConverter):
|
|
"""Convert an onnx Where node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v16(cls, bb, inputs, attr, params):
|
|
if all([isinstance(inp, relax.Constant) for inp in inputs]):
|
|
np_inputs = [inp.data.numpy() for inp in inputs]
|
|
output = _np.where(*np_inputs)
|
|
return relax.const(output, output.dtype)
|
|
if all([isinstance(inp, relax.Constant | relax.ShapeExpr) for inp in inputs]):
|
|
condition, x, y = [get_prim_expr_list(inp) for inp in inputs]
|
|
if len(condition) != len(x) or len(condition) != len(y):
|
|
raise ValueError("Cannot broadcast condition to x and y")
|
|
output = [x if c else y for c, x, y in zip(condition, x, y)]
|
|
return relax.ShapeExpr(output)
|
|
return relax.op.where(inputs[0], inputs[1], inputs[2])
|
|
|
|
|
|
class Clip(OnnxOpConverter):
|
|
"""Converts an onnx Clip node into an equivalent Relax expression."""
|
|
|
|
@staticmethod
|
|
def _sanitize_nan_clip_bound(bb, bound: relax.Expr, *, for_min: bool) -> relax.Expr:
|
|
"""ONNX/ORT treat NaN clip bounds as unbounded; plain max/min with NaN poisons output."""
|
|
dtype = bound.ty.dtype
|
|
if not _relax_dtype_is_floating_point(dtype):
|
|
return bound
|
|
repl = -_np.inf if for_min else _np.inf
|
|
return bb.emit(relax.op.where(relax.op.isnan(bound), relax.const(repl, dtype), bound))
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
min = float(attr.get("min", -_np.inf))
|
|
max = float(attr.get("max", _np.inf))
|
|
results = inputs[0]
|
|
results = bb.emit_te(topi.maximum, results, min)
|
|
results = bb.emit_te(topi.minimum, results, max)
|
|
return results
|
|
|
|
@classmethod
|
|
def _impl_v11(cls, bb, inputs, attr, params):
|
|
# Opset 11 changed Clip from attribute-based min/max to input-based.
|
|
return cls._impl_v13(bb, inputs, attr, params)
|
|
|
|
@classmethod
|
|
def _impl_v13(cls, bb, inputs, attr, params):
|
|
x: Any = inputs[0]
|
|
results = x
|
|
if inputs[1] is not None:
|
|
lo = cls._sanitize_nan_clip_bound(bb, inputs[1], for_min=True)
|
|
results = bb.emit_te(topi.maximum, results, lo)
|
|
if inputs[2] is not None:
|
|
hi = cls._sanitize_nan_clip_bound(bb, inputs[2], for_min=False)
|
|
results = bb.emit_te(topi.minimum, results, hi)
|
|
return results
|
|
|
|
|
|
class Shape(OnnxOpConverter):
|
|
"""Converts an onnx Equal node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v13(cls, bb, inputs, attr, params):
|
|
data_info = inputs[0].ty
|
|
|
|
if isinstance(data_info, relax.ShapeType):
|
|
if data_info.ndim == -1:
|
|
raise ValueError("The ndim of ShapeExpr is expected to a real number, but got -1.")
|
|
return relax.ShapeExpr([data_info.ndim])
|
|
|
|
# If no shape is defined in the type, it must be computed at runtime.
|
|
if not data_info.shape:
|
|
data_shape = bb.normalize(relax.op.shape_of(inputs[0]))
|
|
return data_shape
|
|
|
|
return data_info.shape
|
|
|
|
|
|
class Trilu(OnnxOpConverter):
|
|
"""Given a 2-D matrix or batches of 2-D matrices, returns the upper or
|
|
lower triangular part of the tensor(s)
|
|
"""
|
|
|
|
@classmethod
|
|
def _impl_v14(cls, bb, inputs, attr, params):
|
|
upper = attr.get("upper", True)
|
|
x = inputs[0]
|
|
k = inputs[1] if len(inputs) > 1 else 0
|
|
|
|
if len(inputs) > 1:
|
|
k = get_constant(inputs[1], params)
|
|
if isinstance(k, relax.Constant):
|
|
k = int(k.data.numpy().item())
|
|
else:
|
|
raise ValueError("Currently only support constant k for Trilu op.")
|
|
else:
|
|
k = 0
|
|
|
|
if upper:
|
|
return relax.op.triu(x, k)
|
|
else:
|
|
return relax.op.tril(x, k)
|
|
|
|
|
|
class Relu(OnnxOpConverter):
|
|
"""Converts an onnx Relu node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v13(cls, bb, inputs, attr, params):
|
|
return relax.op.nn.relu(inputs[0])
|
|
|
|
|
|
class Elu(OnnxOpConverter):
|
|
"""Converts an onnx Elu node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
alpha = float(attr.get("alpha", 1.0))
|
|
return relax.expr.const(-alpha) * relax.op.nn.relu(
|
|
relax.expr.const(1.0) - relax.op.exp(inputs[0])
|
|
) + relax.op.nn.relu(inputs[0])
|
|
|
|
|
|
class Selu(OnnxOpConverter):
|
|
"""Converts an onnx Selu node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
alpha = attr.get("alpha", 1.67326319217681884765625)
|
|
gamma = attr.get("gamma", 1.05070102214813232421875)
|
|
return relax.const(gamma) * (
|
|
relax.const(-alpha) * relax.op.nn.relu(relax.const(1.0) - relax.op.exp(inputs[0]))
|
|
+ relax.op.nn.relu(inputs[0])
|
|
)
|
|
|
|
|
|
class Mish(OnnxOpConverter):
|
|
"""Converts an onnx Mish node into an equivalent Relax expression.
|
|
|
|
mish(x) = x * tanh(softplus(x)) = x * tanh(ln(1 + e^{x}))
|
|
"""
|
|
|
|
@classmethod
|
|
def _impl_v18(cls, bb, inputs, attr, params):
|
|
dtype = inputs[0].ty.dtype
|
|
return inputs[0] * relax.op.tanh(
|
|
relax.op.log(relax.const(1.0, dtype) + relax.op.exp(inputs[0]))
|
|
)
|
|
|
|
|
|
class PRelu(OnnxOpConverter):
|
|
"""Converts an onnx PRelu node into an equivalent Relax expression.
|
|
|
|
f(x) = slope * x for x < 0, x for x >= 0
|
|
"""
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
x = inputs[0]
|
|
slope = inputs[1]
|
|
|
|
x_shape = x.ty.shape
|
|
slope_shape = slope.ty.shape
|
|
|
|
ndim = len(x_shape)
|
|
s_ndim = len(slope_shape)
|
|
|
|
if all(ss == 1 for ss in slope_shape) or s_ndim == 1:
|
|
slope = relax.op.reshape(slope, (slope_shape[0],))
|
|
return relax.op.nn.prelu(x, slope, ndim - 1)
|
|
|
|
if s_ndim == ndim:
|
|
non_one_axes = [i for i, ss in enumerate(slope_shape) if ss != 1]
|
|
|
|
# Must have only ONE non-broadcast axis
|
|
if len(non_one_axes) != 1:
|
|
raise ValueError(
|
|
f"Invalid PRelu slope shape (multiple non-broadcast dims): {slope_shape}"
|
|
)
|
|
axis = non_one_axes[0]
|
|
|
|
slope = relax.op.reshape(slope, (slope_shape[axis],))
|
|
return relax.op.nn.prelu(x, slope, axis)
|
|
|
|
raise ValueError(f"Unsupported PRelu slope shape: {slope_shape}")
|
|
|
|
|
|
class ThresholdedRelu(OnnxOpConverter):
|
|
"""Converts an onnx ThresholdedRelu node into an equivalent Relax expression.
|
|
|
|
f(x) = x for x > alpha, 0 otherwise
|
|
"""
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
x = inputs[0]
|
|
alpha = attr.get("alpha", 1.0)
|
|
return relax.op.greater(x, relax.const(alpha)).astype("float32") * x
|
|
|
|
|
|
class LeakyRelu(OnnxOpConverter):
|
|
"""Converts an onnx LeakyRelu node into an equivalent Relax expression.
|
|
|
|
f(x) = x for x > 0, alpha * x otherwise
|
|
"""
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
x = inputs[0]
|
|
alpha = attr.get("alpha", 0.01)
|
|
return relax.op.nn.leakyrelu(x, alpha)
|
|
|
|
|
|
class Gelu(OnnxOpConverter):
|
|
"""Operator converter for Gelu.
|
|
|
|
Supports both Microsoft onnxruntime contrib opset and ONNX Opset 20+.
|
|
|
|
gelu(x) = 0.5x(1 + erf(x/sqrt(2)))
|
|
|
|
When approximate="tanh" (ONNX Opset 20):
|
|
gelu(x) = 0.5x(1 + tanh(sqrt(2/pi)(x + 0.044715x^3)))
|
|
"""
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
return relax.op.nn.gelu(inputs[0])
|
|
|
|
@classmethod
|
|
def _impl_v20(cls, bb, inputs, attr, params):
|
|
approximate = attr.get("approximate", b"none").decode("utf-8")
|
|
if approximate == "tanh":
|
|
return relax.op.nn.gelu_tanh(inputs[0])
|
|
if approximate == "none":
|
|
return relax.op.nn.gelu(inputs[0])
|
|
raise ValueError(f"Unsupported approximate mode for Gelu: {approximate}")
|
|
|
|
|
|
class FastGelu(OnnxOpConverter):
|
|
"""Operator converter for FastGelu from Microsoft onnxruntime contrib opset.
|
|
|
|
fast_gelu(x) = 0.5x(1 + tanh(sqrt(2/pi)(x + 0.044715x^3)))
|
|
= 0.5x(1 + tanh((sqrt(2/pi)x + 0.044715(sqrt(2/pi)x^3)))
|
|
= 0.5x(1 + tanh(c1 * x + c2 * x^3)))
|
|
, where
|
|
c1 = sqrt(2/pi)
|
|
c2 = 0.044715 * sqrt(2/pi)
|
|
"""
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
x = inputs[0]
|
|
if len(inputs) > 1 and inputs[1] is not None:
|
|
bias = inputs[1]
|
|
bias_shape = bias.ty.shape
|
|
assert len(bias_shape) == 1, "bias term must be a 1D tensor"
|
|
x = bb.emit(relax.op.add(x, bias))
|
|
|
|
# Declare consts
|
|
const_dtype = x.ty.dtype
|
|
half = relax.const(0.5, dtype=const_dtype)
|
|
one = relax.const(1.0, dtype=const_dtype)
|
|
const1 = relax.const(math.sqrt(2 / math.pi), dtype=const_dtype)
|
|
const2 = relax.const(0.044715 * math.sqrt(2 / math.pi), dtype=const_dtype)
|
|
|
|
# Compute FastGelu
|
|
term1 = bb.emit(relax.op.multiply(half, x))
|
|
term2 = bb.emit(relax.op.multiply(const1, x))
|
|
# use x^3 = x * x * x instead of pow(x, 3) for better performance
|
|
x_cubed = bb.emit(relax.op.multiply(relax.op.multiply(x, x), x))
|
|
term3 = bb.emit(relax.op.multiply(const2, x_cubed))
|
|
tanh = bb.emit(relax.op.tanh(relax.op.add(term2, term3)))
|
|
return bb.emit(relax.op.multiply(term1, relax.op.add(one, tanh)))
|
|
|
|
|
|
class BiasGelu(OnnxOpConverter):
|
|
"""Operator converter for BiasGelu from Microsoft onnxruntime contrib opset.
|
|
|
|
bias_gelu(x, b) = 0.5(x + b)(1 + erf((x + b)/sqrt(2)))
|
|
"""
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
inp = relax.op.add(inputs[0], inputs[1])
|
|
return relax.op.nn.gelu(inp)
|
|
|
|
|
|
class Shrink(OnnxOpConverter):
|
|
"""Converts an onnx Shrink node into an equivalent Relax expression.
|
|
|
|
f(x) = x + bias if x > lambd, x - bias if x < -lambd, 0 otherwise
|
|
"""
|
|
|
|
@classmethod
|
|
def _impl_v9(cls, bb, inputs, attr, params):
|
|
x = inputs[0]
|
|
dtype = x.ty.dtype
|
|
lambd = relax.const(attr.get("lambd", 0.5), dtype)
|
|
bias = relax.const(attr.get("bias", 0.0), dtype)
|
|
zeros = relax.op.zeros_like(x)
|
|
return relax.op.where(x > lambd, x - bias, zeros) + relax.op.where(
|
|
x < -lambd, x + bias, zeros
|
|
)
|
|
|
|
|
|
class Conv(OnnxOpConverter):
|
|
"""Convert an onnx Conv node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v11(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
if hasattr(inputs[0].ty, "ndim"):
|
|
ndim = inputs[0].ty.ndim
|
|
else:
|
|
ndim = len(inputs[0].ty.shape)
|
|
|
|
if "kernel_shape" not in attr:
|
|
attr["kernel_shape"] = inputs[1].ty.shape.values[2:]
|
|
|
|
if ndim == 3:
|
|
op = relax.op.nn.conv1d
|
|
data_layout = "NCW"
|
|
kernel_layout = "OIW"
|
|
elif ndim == 4:
|
|
op = relax.op.nn.conv2d
|
|
data_layout = "NCHW"
|
|
kernel_layout = "OIHW"
|
|
elif ndim == 5:
|
|
op = relax.op.nn.conv3d
|
|
data_layout = "NCDHW"
|
|
kernel_layout = "OIDHW"
|
|
else:
|
|
raise NotImplementedError("Ndim > 5 not supported for convolution.")
|
|
|
|
if "auto_pad" in attr:
|
|
attr["auto_pad"] = attr["auto_pad"].decode("utf-8")
|
|
if attr["auto_pad"] in ("SAME_UPPER", "SAME_LOWER"):
|
|
data = autopad(
|
|
bb,
|
|
inputs[0],
|
|
attr.get("strides", [1] * (ndim - 2)),
|
|
attr["kernel_shape"],
|
|
attr.get("dilations", [1] * (ndim - 2)),
|
|
mode=attr["auto_pad"],
|
|
deconv=False,
|
|
)
|
|
elif attr["auto_pad"] == "VALID":
|
|
attr["pads"] = [0 for _ in range(ndim - 2)]
|
|
elif attr["auto_pad"] == "NOTSET":
|
|
pass
|
|
else:
|
|
msg = (
|
|
f'Value {attr["auto_pad"]} in attribute "auto_pad" of operator Conv is invalid.'
|
|
)
|
|
raise tvm.error.OpAttributeInvalid(msg)
|
|
attr.pop("auto_pad")
|
|
|
|
conv_out = bb.normalize(
|
|
op(
|
|
data=data,
|
|
weight=inputs[1],
|
|
strides=attr.get("strides", 1),
|
|
padding=attr.get("pads", 0),
|
|
dilation=attr.get("dilations", 1),
|
|
groups=attr.get("group", 1),
|
|
data_layout=data_layout,
|
|
kernel_layout=kernel_layout,
|
|
)
|
|
)
|
|
if inputs[2] is not None:
|
|
bias = relax.op.reshape(inputs[2], [1, -1] + [1] * (ndim - 2))
|
|
conv_out = relax.op.add(conv_out, bias)
|
|
|
|
return conv_out
|
|
|
|
|
|
class ConvTranspose(OnnxOpConverter):
|
|
"""Converts an onnx ConvTranspose node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
if hasattr(inputs[0].ty, "ndim"):
|
|
ndim = inputs[0].ty.ndim
|
|
else:
|
|
ndim = len(inputs[0].ty.shape)
|
|
|
|
if ndim == 3:
|
|
op = relax.op.nn.conv1d_transpose
|
|
data_layout = "NCW"
|
|
kernel_layout = "IOW"
|
|
elif ndim == 4:
|
|
op = relax.op.nn.conv2d_transpose
|
|
data_layout = "NCHW"
|
|
kernel_layout = "IOHW"
|
|
elif ndim == 5:
|
|
op = relax.op.nn.conv3d_transpose
|
|
data_layout = "NCDHW"
|
|
kernel_layout = "IODHW"
|
|
else:
|
|
raise NotImplementedError("Ndim > 5 not supported for convolution.")
|
|
|
|
spatial_dims = ndim - 2
|
|
strides = attr.get("strides", [1] * spatial_dims)
|
|
dilations = attr.get("dilations", [1] * spatial_dims)
|
|
output_padding = attr.get("output_padding", [0] * spatial_dims)
|
|
if "kernel_shape" in attr:
|
|
kernel_shape = list(attr["kernel_shape"])
|
|
else:
|
|
kernel_shape = [int(s) for s in inputs[1].ty.shape.values[2:]]
|
|
|
|
# Resolve `auto_pad` per ONNX ConvTranspose spec. Unlike Conv, the spec
|
|
# derives `pads` from `output_shape`/`strides` when auto_pad is SAME_*,
|
|
# so we cannot reuse `autopad()` (which pads the input data instead).
|
|
if "auto_pad" in attr:
|
|
auto_pad = attr["auto_pad"]
|
|
if isinstance(auto_pad, bytes):
|
|
auto_pad = auto_pad.decode("utf-8")
|
|
if auto_pad in ("SAME_UPPER", "SAME_LOWER"):
|
|
# Per ONNX ConvTranspose spec, when output_shape is unspecified
|
|
# the target output size is `input_size * stride`. Substituting
|
|
# this into the spec's total_padding formula cancels the
|
|
# input-size term, leaving a value that depends only on the
|
|
# kernel/dilation/stride/output_padding attributes. Avoiding the
|
|
# input shape keeps the converter usable when spatial dims are
|
|
# symbolic (`tir.Var`).
|
|
pads_begin: list[int] = []
|
|
pads_end: list[int] = []
|
|
for i in range(spatial_dims):
|
|
total_pad = (
|
|
(kernel_shape[i] - 1) * dilations[i] + 1 + output_padding[i] - strides[i]
|
|
)
|
|
total_pad = max(total_pad, 0)
|
|
if auto_pad == "SAME_UPPER":
|
|
pad_begin = total_pad // 2
|
|
else:
|
|
pad_begin = total_pad - total_pad // 2
|
|
pads_begin.append(pad_begin)
|
|
pads_end.append(total_pad - pad_begin)
|
|
attr["pads"] = pads_begin + pads_end
|
|
elif auto_pad == "VALID":
|
|
attr["pads"] = [0] * (2 * spatial_dims)
|
|
elif auto_pad == "NOTSET":
|
|
pass
|
|
else:
|
|
raise tvm.error.OpAttributeInvalid(
|
|
f'Value {auto_pad} in attribute "auto_pad" of operator '
|
|
"ConvTranspose is invalid."
|
|
)
|
|
attr.pop("auto_pad")
|
|
|
|
conv_out = op(
|
|
data=inputs[0],
|
|
weight=inputs[1],
|
|
strides=strides,
|
|
padding=attr.get("pads", 0),
|
|
output_padding=output_padding,
|
|
dilation=dilations,
|
|
groups=attr.get("group", 1),
|
|
data_layout=data_layout,
|
|
kernel_layout=kernel_layout,
|
|
)
|
|
|
|
if inputs[2] is not None:
|
|
bias = relax.op.reshape(inputs[2], [1, -1] + [1] * (ndim - 2))
|
|
conv_out = relax.op.add(conv_out, bias)
|
|
|
|
return conv_out
|
|
|
|
|
|
class Erf(OnnxOpConverter):
|
|
"""Converts an onnx Erf node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v13(cls, bb, inputs, attr, params):
|
|
return relax.op.erf(inputs[0])
|
|
|
|
|
|
class CumSum(OnnxOpConverter):
|
|
"""Converts an onnx CumSum node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v14(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
axis_input = get_constant(inputs[1], params)
|
|
exclusive = attr.get("exclusive", 0) != 0
|
|
|
|
if isinstance(axis_input, relax.Constant):
|
|
axis_data = axis_input.data.numpy()
|
|
if axis_data.ndim == 0:
|
|
axis = int(axis_data.item())
|
|
elif axis_data.ndim == 1 and axis_data.shape[0] == 1:
|
|
axis = int(axis_data.item())
|
|
else:
|
|
raise ValueError(
|
|
"CumSum axis input must be a scalar (0-D) or a single-element 1-D tensor, "
|
|
f"got shape {axis_data.shape}"
|
|
)
|
|
elif isinstance(axis_input, relax.Var):
|
|
axis_shape = axis_input.ty.shape if hasattr(axis_input.ty, "shape") else None
|
|
raise ValueError(
|
|
"CumSum with non-constant axis input is not supported yet. "
|
|
"ONNX permits runtime axis tensors, but Relax/TE currently requires a compile-time "
|
|
f"constant axis for cumsum/flip. Got axis shape {axis_shape}"
|
|
)
|
|
else:
|
|
raise TypeError("CumSum axis input must be a Constant or Var")
|
|
|
|
if attr.get("reverse", 0) != 0:
|
|
data = bb.emit_te(topi.flip, data, axis=axis)
|
|
|
|
data = relax.op.cumsum(data, axis, exclusive=exclusive)
|
|
data = bb.normalize(data)
|
|
|
|
if attr.get("reverse", 0) != 0:
|
|
data = bb.emit_te(topi.flip, data, axis=axis)
|
|
|
|
return data
|
|
|
|
|
|
class Squeeze(OnnxOpConverter):
|
|
"""Converts an onnx Squeeze node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v13(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
axis = get_constant(inputs[1], params)
|
|
if isinstance(axis, relax.Constant):
|
|
axis = tuple([int(x) for x in axis.data.numpy()])
|
|
|
|
# If data is constant, perform computation directly.
|
|
if isinstance(data, relax.Constant):
|
|
if isinstance(axis, tuple | type(None)):
|
|
out_data = _np.squeeze(data.data.numpy(), axis)
|
|
else:
|
|
raise NotImplementedError("Squeeze with symbolic axes not supported")
|
|
|
|
return relax.const(out_data, data.ty.dtype)
|
|
|
|
if isinstance(data, relax.ShapeExpr):
|
|
shape_tensor_ndim = 1
|
|
if axis is None:
|
|
if len(data) == 1:
|
|
return relax.prim_value(data[0])
|
|
return data
|
|
normalized_axes = _normalize_constant_axes(list(axis), shape_tensor_ndim, "Squeeze")
|
|
if normalized_axes == [0] and len(data) == 1:
|
|
return relax.prim_value(data[0])
|
|
raise NotImplementedError(
|
|
"Squeeze on symbolic shape tensors only supports removing the sole axis."
|
|
)
|
|
|
|
if axis is None:
|
|
return relax.op.squeeze(data)
|
|
|
|
if isinstance(axis, tuple):
|
|
return relax.op.squeeze(data, list(axis))
|
|
|
|
data_ndim = _get_known_tensor_rank(data)
|
|
if data_ndim is None:
|
|
raise ValueError("Squeeze with dynamic axes requires a statically known input rank.")
|
|
axes_len = _get_known_tensor_length(axis)
|
|
if axes_len is None:
|
|
raise ValueError("Squeeze requires a statically known axes length.")
|
|
data_shape = bb.normalize(relax.op.shape_of(data))
|
|
data_shape_tensor = bb.normalize(relax.op.shape_to_tensor(data_shape))
|
|
output_shape_tensor = _build_squeezed_shape_tensor(bb, data_shape_tensor, axis, data_ndim)
|
|
output_shape = _tensor_to_shape_expr(
|
|
bb, output_shape_tensor, data_ndim - axes_len, "squeeze_dim"
|
|
)
|
|
return relax.op.reshape(data, output_shape)
|
|
|
|
|
|
class Constant(OnnxOpConverter):
|
|
"""Converts an onnx Constant node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v13(cls, bb, inputs, attr, params):
|
|
if "value" not in attr:
|
|
raise ValueError("no value in Constant")
|
|
value = attr.pop("value")
|
|
# Constants may rarely have string types. These are likely exported
|
|
# from other frameworks and not actually used in TVM. We'll just use
|
|
# a zero valued constant for compatibility.
|
|
if isinstance(value, bytes):
|
|
np_value = _np.asarray([0]).astype("int64")
|
|
else:
|
|
np_value = get_numpy(value)
|
|
dtype = np_value.dtype.name
|
|
value = relax.const(np_value, dtype)
|
|
return value
|
|
|
|
|
|
class ConstantOfShape(OnnxOpConverter):
|
|
"""Converts an onnx ConstantOfShape node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v9(cls, bb, inputs, attr, params):
|
|
shape = inputs[0]
|
|
# ONNX spec: `value` is optional and defaults to a zero float32 scalar.
|
|
# `get_numpy` requires a TensorProto, so dispatch on presence first.
|
|
attr_value = attr.get("value")
|
|
value = get_numpy(attr_value) if attr_value is not None else 0
|
|
if isinstance(value, _np.ndarray):
|
|
dtype = str(value.dtype)
|
|
else:
|
|
dtype = "float32"
|
|
# If shape is a constant, treat it as a ShapeExpr.
|
|
if isinstance(shape, relax.Constant):
|
|
shape = relax.ShapeExpr(list(shape.data.numpy()))
|
|
|
|
# Special case where requested shape are constant
|
|
if len(shape) == 1 and all([isinstance(x, tirx.IntImm) for x in shape]):
|
|
shape = [int(x) for x in shape]
|
|
return relax.const(_np.full(shape, value, dtype), dtype)
|
|
|
|
# Convert to shape expression from tensor if needed.
|
|
if not isinstance(shape, relax.ShapeExpr):
|
|
shape = relax.op.tensor_to_shape(shape)
|
|
|
|
return relax.op.broadcast_to(relax.const(value, dtype), shape)
|
|
|
|
|
|
class Sin(OnnxOpConverter):
|
|
"""Converts an onnx Sin node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v7(cls, bb, inputs, attr, params):
|
|
return relax.op.sin(inputs[0])
|
|
|
|
|
|
class Sinh(OnnxOpConverter):
|
|
"""Converts an onnx Sinh node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v9(cls, bb, inputs, attr, params):
|
|
return relax.op.sinh(inputs[0])
|
|
|
|
|
|
class Cos(OnnxOpConverter):
|
|
"""Converts an onnx Cos node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v7(cls, bb, inputs, attr, params):
|
|
return relax.op.cos(inputs[0])
|
|
|
|
|
|
class Cosh(OnnxOpConverter):
|
|
"""Converts an onnx Cosh node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v9(cls, bb, inputs, attr, params):
|
|
return relax.op.cosh(inputs[0])
|
|
|
|
|
|
class Tan(OnnxOpConverter):
|
|
"""Converts an onnx Tan node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v7(cls, bb, inputs, attr, params):
|
|
return relax.op.tan(inputs[0])
|
|
|
|
|
|
class Tanh(OnnxOpConverter):
|
|
"""Converts an onnx Tanh node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v7(cls, bb, inputs, attr, params):
|
|
return relax.op.tanh(inputs[0])
|
|
|
|
|
|
class Acos(OnnxOpConverter):
|
|
"""Converts an onnx Acos node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v7(cls, bb, inputs, attr, params):
|
|
return relax.op.acos(inputs[0])
|
|
|
|
|
|
class Acosh(OnnxOpConverter):
|
|
"""Converts an onnx Acosh node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v9(cls, bb, inputs, attr, params):
|
|
return relax.op.acosh(inputs[0])
|
|
|
|
|
|
class Asin(OnnxOpConverter):
|
|
"""Converts an onnx Asin node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v7(cls, bb, inputs, attr, params):
|
|
return relax.op.asin(inputs[0])
|
|
|
|
|
|
class Asinh(OnnxOpConverter):
|
|
"""Converts an onnx Asinh node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v9(cls, bb, inputs, attr, params):
|
|
return relax.op.asinh(inputs[0])
|
|
|
|
|
|
class Atan(OnnxOpConverter):
|
|
"""Converts an onnx Atan node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v7(cls, bb, inputs, attr, params):
|
|
return relax.op.atan(inputs[0])
|
|
|
|
|
|
class Atanh(OnnxOpConverter):
|
|
"""Converts an onnx Atanh node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v9(cls, bb, inputs, attr, params):
|
|
return relax.op.atanh(inputs[0])
|
|
|
|
|
|
class Neg(OnnxOpConverter):
|
|
"""Converts an onnx Neg node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v13(cls, bb, inputs, attr, params):
|
|
if isinstance(inputs[0], relax.Constant):
|
|
data_np = inputs[0].data.numpy()
|
|
return relax.const(_np.negative(data_np), inputs[0].ty.dtype)
|
|
if tvm.ir.is_prim_expr(inputs[0]):
|
|
return -inputs[0]
|
|
return relax.op.negative(inputs[0])
|
|
|
|
|
|
class Abs(OnnxOpConverter):
|
|
"""Converts an onnx Abs node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v13(cls, bb, inputs, attr, params):
|
|
if isinstance(inputs[0], relax.Constant):
|
|
output = _np.abs(inputs[0].data.numpy())
|
|
return relax.const(output, output.dtype)
|
|
return relax.op.abs(inputs[0])
|
|
|
|
|
|
class Reciprocal(OnnxOpConverter):
|
|
"""Converts an onnx Reciprocal node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v13(cls, bb, inputs, attr, params):
|
|
input_dtype = inputs[0].ty.dtype
|
|
return relax.op.divide(relax.const(1, dtype=input_dtype), inputs[0])
|
|
|
|
|
|
class Floor(OnnxOpConverter):
|
|
"""Converts an onnx Floor node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
return relax.op.floor(inputs[0])
|
|
|
|
|
|
class Ceil(OnnxOpConverter):
|
|
"""Converts an onnx Ceil node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
return relax.op.ceil(inputs[0])
|
|
|
|
|
|
class Round(OnnxOpConverter):
|
|
"""Converts an onnx Round node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
return relax.op.round(inputs[0])
|
|
|
|
|
|
class IsInf(OnnxOpConverter):
|
|
"""Converts an onnx IsInf node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v10(cls, bb, inputs, attr, params):
|
|
return relax.op.isinf(inputs[0])
|
|
|
|
|
|
class IsNaN(OnnxOpConverter):
|
|
"""Converts an onnx IsNaN node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v9(cls, bb, inputs, attr, params):
|
|
return relax.op.isnan(inputs[0])
|
|
|
|
|
|
class Sqrt(OnnxOpConverter):
|
|
"""Converts an onnx Sqrt node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
return relax.op.sqrt(inputs[0])
|
|
|
|
|
|
def compute_broadcast_shape(shape_a, shape_b):
|
|
"""Compute target shape for Multidirectional Broadcasting"""
|
|
rank = max(len(shape_a), len(shape_b))
|
|
|
|
a = (1,) * (rank - len(shape_a)) + tuple(shape_a)
|
|
b = (1,) * (rank - len(shape_b)) + tuple(shape_b)
|
|
|
|
target = []
|
|
for ai, bi in zip(a, b):
|
|
if ai == bi or ai == 1 or bi == 1:
|
|
target.append(max(ai, bi))
|
|
else:
|
|
raise ValueError(f"Cannot broadcast {ai} and {bi}")
|
|
return tuple(target)
|
|
|
|
|
|
class MultiInputBase(OnnxOpConverter):
|
|
"""Converts an onnx MultiInputBase node into an equivalent Relax expression."""
|
|
|
|
numpy_op: Callable = None
|
|
relax_op: Callable = None
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
if cls.numpy_op is None or cls.relax_op is None:
|
|
raise NotImplementedError("numpy_op and relax_op must be defined for MultiInputBase")
|
|
if all([isinstance(inp, relax.Constant) for inp in inputs]):
|
|
np_inputs = [inp.data.numpy() for inp in inputs]
|
|
output = cls.numpy_op(*np_inputs) # pylint: disable=not-callable
|
|
return relax.const(output, output.dtype)
|
|
|
|
input_shapes = [inp.ty.shape for inp in inputs]
|
|
target_shape = functools.reduce(compute_broadcast_shape, input_shapes)
|
|
|
|
# broadcast_to, stack them, then perform minimum over the new axis.
|
|
inputs = [bb.normalize(relax.op.broadcast_to(i, target_shape)) for i in inputs]
|
|
stacked_tensor = bb.normalize(relax.op.stack(inputs, axis=0))
|
|
return cls.relax_op(stacked_tensor, axis=0) # pylint: disable=not-callable
|
|
|
|
|
|
class Min(MultiInputBase):
|
|
"""Converts an onnx Min node into an equivalent Relax expression."""
|
|
|
|
numpy_op = _np.min
|
|
relax_op = relax.op.min
|
|
|
|
|
|
class Max(MultiInputBase):
|
|
"""Converts an onnx Max node into an equivalent Relax expression."""
|
|
|
|
numpy_op = _np.max
|
|
relax_op = relax.op.max
|
|
|
|
|
|
class Mean(MultiInputBase):
|
|
"""Converts an onnx Mean node into an equivalent Relax expression."""
|
|
|
|
numpy_op = _np.mean
|
|
relax_op = relax.op.mean
|
|
|
|
|
|
class Sum(MultiInputBase):
|
|
"""Converts an onnx Sum node into an equivalent Relax expression."""
|
|
|
|
numpy_op = _np.sum
|
|
relax_op = relax.op.sum
|
|
|
|
|
|
class Log(OnnxOpConverter):
|
|
"""Converts an onnx Log node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v13(cls, bb, inputs, attr, params):
|
|
if isinstance(inputs[0], relax.Constant):
|
|
return relax.const(_np.log(inputs[0].data.numpy()), inputs[0].ty.dtype)
|
|
return relax.op.log(inputs[0])
|
|
|
|
|
|
class Exp(OnnxOpConverter):
|
|
"""Converts an onnx Exp node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _check_type(cls, dtype, valid_types):
|
|
assert dtype in valid_types, f"Types {valid_types} are supported only, but {dtype} is given"
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
valid_types = ["float", "float32", "double", "float64", "float16"]
|
|
cls._check_type(data.ty.dtype, valid_types)
|
|
|
|
return relax.op.exp(data)
|
|
|
|
@classmethod
|
|
def _impl_v13(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
valid_types = ["float", "float32", "double", "float64", "float16", "bfloat16"]
|
|
cls._check_type(data.ty.dtype, valid_types)
|
|
|
|
return relax.op.exp(data)
|
|
|
|
|
|
class Softplus(OnnxOpConverter):
|
|
"""Converts an onnx Softplus node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
dtype = inputs[0].ty.dtype
|
|
threshold = 10.0 if dtype == "float16" else 20.0
|
|
return relax.op.nn.softplus(inputs[0], threshold=threshold)
|
|
|
|
|
|
class Softsign(OnnxOpConverter):
|
|
"""Converts an onnx Softsign node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
dtype = inputs[0].ty.dtype
|
|
return inputs[0] / (relax.op.abs(inputs[0]) + relax.const(1, dtype=dtype))
|
|
|
|
|
|
class Split(OnnxOpConverter):
|
|
"""Converts an onnx Split node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
splits = attr.get("split", None)
|
|
if splits is not None and len(splits) > 1:
|
|
indices = []
|
|
index = 0
|
|
for i in splits[:-1]:
|
|
index += i
|
|
indices.append(index)
|
|
# When splits isnt specified divide evenly over axis.
|
|
else:
|
|
indices = attr["tvm_custom"]["num_outputs"]
|
|
return relax.op.split(inputs[0], indices, attr.get("axis", 0))
|
|
|
|
@classmethod
|
|
def _impl_v13(cls, bb, inputs, attr, params):
|
|
splits = inputs[1]
|
|
splits_rank = None
|
|
if splits is not None:
|
|
splits_rank = splits.ty.ndim
|
|
if splits is not None and splits_rank > 0:
|
|
if isinstance(splits, relax.Constant):
|
|
splits = splits.data.numpy()
|
|
indices = []
|
|
index = 0
|
|
for i in splits[:-1]:
|
|
index += i
|
|
indices.append(index.item())
|
|
else:
|
|
raise ValueError("Dynamic Split not yet supported")
|
|
# When splits isnt specified divide evenly over axis.
|
|
else:
|
|
indices = attr["tvm_custom"]["num_outputs"]
|
|
return relax.op.split(inputs[0], indices, attr.get("axis", 0))
|
|
|
|
|
|
def get_prim_value_list(values):
|
|
new_values = []
|
|
for v in list(values):
|
|
if tvm.ir.is_prim_expr(v):
|
|
new_values.append(relax.prim_value(v))
|
|
else:
|
|
new_values.append(v)
|
|
return new_values
|
|
|
|
|
|
def _get_known_tensor_rank(expr: relax.Expr) -> int | None:
|
|
"""Return the statically known rank of an expression when available."""
|
|
|
|
if isinstance(expr, relax.Constant):
|
|
return len(expr.data.numpy().shape)
|
|
if isinstance(expr, relax.ShapeExpr):
|
|
return 1
|
|
if tvm.ir.is_prim_expr(expr):
|
|
return 0
|
|
ty = expr.ty
|
|
if isinstance(ty, relax.TensorType):
|
|
return None if ty.ndim == -1 else ty.ndim
|
|
return None
|
|
|
|
|
|
def _get_known_tensor_length(expr: relax.Expr | None) -> int | None:
|
|
"""Return the statically known length of a 1-D tensor-like expression."""
|
|
|
|
if expr is None:
|
|
return None
|
|
if isinstance(expr, relax.Constant):
|
|
np_value = expr.data.numpy()
|
|
if np_value.ndim != 1:
|
|
raise ValueError(f"Expected a 1-D tensor, but got ndim={np_value.ndim}.")
|
|
return int(np_value.shape[0])
|
|
if isinstance(expr, relax.ShapeExpr):
|
|
return len(expr.values)
|
|
if tvm.ir.is_prim_expr(expr):
|
|
return 1
|
|
ty = expr.ty
|
|
if not isinstance(ty, relax.TensorType):
|
|
return None
|
|
if ty.ndim == -1:
|
|
return None
|
|
if ty.ndim != 1:
|
|
raise ValueError(f"Expected a 1-D tensor, but got ndim={ty.ndim}.")
|
|
if isinstance(ty.shape, relax.ShapeExpr):
|
|
dim = ty.shape.values[0]
|
|
if isinstance(dim, tirx.IntImm):
|
|
return int(dim.value)
|
|
if isinstance(dim, int):
|
|
return dim
|
|
return None
|
|
|
|
|
|
def _normalize_constant_axes(axes: list[int], rank: int, op_name: str) -> list[int]:
|
|
"""Normalize a list of constant axes and validate their uniqueness."""
|
|
|
|
normalized_axes = []
|
|
for axis in axes:
|
|
original_axis = axis
|
|
if axis < 0:
|
|
axis += rank
|
|
if axis < 0 or axis >= rank:
|
|
raise ValueError(f"{op_name} axis {original_axis} is out of range for rank {rank}.")
|
|
normalized_axes.append(axis)
|
|
if len(normalized_axes) != len(set(normalized_axes)):
|
|
raise ValueError(f"{op_name} axes must be unique.")
|
|
return normalized_axes
|
|
|
|
|
|
def _as_int64_tensor(bb: relax.BlockBuilder, expr: relax.Expr) -> relax.Expr:
|
|
"""Convert a tensor-like expression to an int64 tensor expression."""
|
|
|
|
if isinstance(expr, relax.ShapeExpr):
|
|
return bb.normalize(relax.op.shape_to_tensor(expr))
|
|
if tvm.ir.is_prim_expr(expr):
|
|
return bb.normalize(relax.op.full((1,), expr, dtype="int64"))
|
|
if isinstance(expr, relax.Constant):
|
|
if expr.ty.dtype == "int64":
|
|
return expr
|
|
return bb.normalize(relax.op.astype(expr, "int64"))
|
|
if isinstance(expr.ty, relax.TensorType) and expr.ty.dtype != "int64":
|
|
return bb.normalize(relax.op.astype(expr, "int64"))
|
|
return expr
|
|
|
|
|
|
def _tensor_to_shape_expr(
|
|
bb: relax.BlockBuilder, shape_tensor: relax.Expr, shape_ndim: int, prefix: str
|
|
) -> relax.ShapeExpr:
|
|
"""Convert a statically sized int64 tensor into a ShapeExpr."""
|
|
|
|
shape_tensor = bb.match_cast(shape_tensor, relax.TensorType([shape_ndim], "int64"))
|
|
shape_dataflow_var = bb.emit(relax.op.tensor_to_shape(shape_tensor))
|
|
shape_vars = [tirx.Var(f"{prefix}_{i}", "int64") for i in range(shape_ndim)]
|
|
bb.match_cast(shape_dataflow_var, relax.ShapeType(shape_vars))
|
|
return relax.ShapeExpr(shape_vars)
|
|
|
|
|
|
def _build_unsqueezed_shape_tensor(
|
|
bb: relax.BlockBuilder, data_shape_tensor: relax.Expr, axes: relax.Expr, data_ndim: int
|
|
) -> relax.Expr:
|
|
"""Build the output shape tensor for Unsqueeze with runtime axes."""
|
|
|
|
axes = _as_int64_tensor(bb, axes)
|
|
axes_len = _get_known_tensor_length(axes)
|
|
if axes_len is None:
|
|
raise ValueError("Unsqueeze requires a statically known axes length.")
|
|
|
|
output_ndim = data_ndim + axes_len
|
|
axes = bb.normalize(
|
|
relax.op.where(
|
|
relax.op.less(axes, relax.const(0, "int64")),
|
|
relax.op.add(axes, relax.const(output_ndim, "int64")),
|
|
axes,
|
|
)
|
|
)
|
|
positions = relax.op.arange(output_ndim, dtype="int64")
|
|
positions = bb.normalize(relax.op.expand_dims(positions, axis=1))
|
|
axes = bb.normalize(relax.op.expand_dims(axes, axis=0))
|
|
insert_mask = bb.normalize(
|
|
relax.op.sum(relax.op.astype(relax.op.equal(positions, axes), "int64"), axis=1)
|
|
)
|
|
keep_mask = bb.normalize(relax.op.subtract(relax.const(1, "int64"), insert_mask))
|
|
input_indices = bb.normalize(
|
|
relax.op.subtract(relax.op.cumsum(keep_mask, axis=0), relax.const(1, "int64"))
|
|
)
|
|
safe_indices = bb.normalize(
|
|
relax.op.where(
|
|
relax.op.less(input_indices, relax.const(0, "int64")),
|
|
relax.const(0, "int64"),
|
|
input_indices,
|
|
)
|
|
)
|
|
kept_dims = bb.normalize(relax.op.take(data_shape_tensor, safe_indices, axis=0))
|
|
return bb.normalize(
|
|
relax.op.where(
|
|
relax.op.greater(insert_mask, relax.const(0, "int64")),
|
|
relax.const(1, "int64"),
|
|
kept_dims,
|
|
)
|
|
)
|
|
|
|
|
|
def _build_squeezed_shape_tensor(
|
|
bb: relax.BlockBuilder, data_shape_tensor: relax.Expr, axes: relax.Expr, data_ndim: int
|
|
) -> relax.Expr:
|
|
"""Build the output shape tensor for Squeeze with runtime axes."""
|
|
|
|
axes = _as_int64_tensor(bb, axes)
|
|
axes = bb.normalize(
|
|
relax.op.where(
|
|
relax.op.less(axes, relax.const(0, "int64")),
|
|
relax.op.add(axes, relax.const(data_ndim, "int64")),
|
|
axes,
|
|
)
|
|
)
|
|
positions = relax.op.arange(data_ndim, dtype="int64")
|
|
positions = bb.normalize(relax.op.expand_dims(positions, axis=1))
|
|
axes = bb.normalize(relax.op.expand_dims(axes, axis=0))
|
|
remove_mask = bb.normalize(
|
|
relax.op.sum(relax.op.astype(relax.op.equal(positions, axes), "int64"), axis=1)
|
|
)
|
|
keep_mask = bb.normalize(relax.op.equal(remove_mask, relax.const(0, "int64")))
|
|
keep_indices = bb.normalize(relax.op.nonzero(keep_mask))
|
|
num_keep_dims = tirx.Var("squeeze_num_keep_dims", "int64")
|
|
keep_indices = bb.match_cast(keep_indices, relax.TensorType([1, num_keep_dims], "int64"))
|
|
keep_indices = bb.normalize(relax.op.reshape(keep_indices, [-1]))
|
|
return bb.normalize(relax.op.take(data_shape_tensor, keep_indices, axis=0))
|
|
|
|
|
|
class Slice(OnnxOpConverter):
|
|
"""Converts an onnx Slice node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v13(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
starts = get_constant(inputs[1], params)
|
|
ends = get_constant(inputs[2], params)
|
|
axes = get_constant(inputs[3], params)
|
|
steps = get_constant(inputs[4], params)
|
|
all_constant_params = all(
|
|
isinstance(param, relax.Constant | relax.ShapeExpr)
|
|
or tvm.ir.is_prim_expr(param)
|
|
or param is None
|
|
for param in [starts, ends, axes, steps]
|
|
)
|
|
if all_constant_params:
|
|
starts = get_prim_expr_list(starts)
|
|
ends = get_prim_expr_list(ends)
|
|
if len(starts) != len(ends):
|
|
raise ValueError(
|
|
f"Slice expects starts and ends to have the same length, but got "
|
|
f"{len(starts)} and {len(ends)}."
|
|
)
|
|
if axes is not None:
|
|
axes = get_prim_expr_list(axes)
|
|
if len(axes) != len(starts):
|
|
raise ValueError(
|
|
f"Slice expects axes and starts to have the same length, but got "
|
|
f"{len(axes)} and {len(starts)}."
|
|
)
|
|
else:
|
|
axes = list(range(len(starts)))
|
|
|
|
data_ndim = _get_known_tensor_rank(data)
|
|
if data_ndim is None:
|
|
raise ValueError("Slice requires a statically known input rank.")
|
|
axes = _normalize_constant_axes(list(axes), data_ndim, "Slice")
|
|
if steps is not None:
|
|
steps = get_prim_expr_list(steps)
|
|
if len(steps) != len(starts):
|
|
raise ValueError(
|
|
f"Slice expects steps and starts to have the same length, but got "
|
|
f"{len(steps)} and {len(starts)}."
|
|
)
|
|
else:
|
|
steps = [1] * len(axes)
|
|
if any(
|
|
(isinstance(step, int) and step == 0)
|
|
or (isinstance(step, tirx.IntImm) and int(step) == 0)
|
|
for step in steps
|
|
):
|
|
raise ValueError("Slice step values must be non-zero.")
|
|
if isinstance(data, relax.ShapeExpr):
|
|
shape_data = list(data)
|
|
assert all(len(i) == 1 for i in [starts, ends, steps])
|
|
sliced_values = shape_data[starts[0] : ends[0] : steps[0]]
|
|
|
|
if all([isinstance(val, tirx.IntImm | int) for val in sliced_values]):
|
|
return relax.const([x.value for x in sliced_values], "int64")
|
|
return relax.ShapeExpr(sliced_values)
|
|
|
|
assume_inbound = not all(
|
|
[isinstance(param, tirx.IntImm | int) for param in [*starts, *ends, *steps]]
|
|
)
|
|
starts = get_prim_value_list(starts)
|
|
ends = get_prim_value_list(ends)
|
|
steps = get_prim_value_list(steps)
|
|
|
|
return relax.op.strided_slice(
|
|
data, axes, starts, ends, steps, assume_inbound=assume_inbound
|
|
)
|
|
|
|
data_ndim = _get_known_tensor_rank(data)
|
|
if data_ndim is None:
|
|
raise ValueError(
|
|
"Slice with dynamic parameters requires a statically known input rank."
|
|
)
|
|
|
|
if isinstance(data, relax.ShapeExpr):
|
|
raise ValueError("Slice with dynamic parameters does not support ShapeExpr input.")
|
|
data_expr = data
|
|
|
|
starts_tensor = _as_int64_tensor(bb, starts)
|
|
ends_tensor = _as_int64_tensor(bb, ends)
|
|
axes_len = _get_known_tensor_length(starts_tensor)
|
|
if axes_len is None:
|
|
raise ValueError("Slice requires a statically known starts length.")
|
|
ends_len = _get_known_tensor_length(ends_tensor)
|
|
if ends_len is None:
|
|
raise ValueError("Slice requires a statically known ends length.")
|
|
if ends_len != axes_len:
|
|
raise ValueError(
|
|
f"Slice expects starts and ends to have the same length, but got "
|
|
f"{axes_len} and {ends_len}."
|
|
)
|
|
|
|
if axes is None:
|
|
axes_tensor = relax.op.arange(axes_len, dtype="int64")
|
|
else:
|
|
axes_tensor = _as_int64_tensor(bb, axes)
|
|
axes_tensor_len = _get_known_tensor_length(axes_tensor)
|
|
if axes_tensor_len is None:
|
|
raise ValueError("Slice requires a statically known axes length.")
|
|
if axes_tensor_len != axes_len:
|
|
raise ValueError(
|
|
f"Slice expects axes and starts to have the same length, but got "
|
|
f"{axes_tensor_len} and {axes_len}."
|
|
)
|
|
if steps is None:
|
|
steps_tensor = relax.const(_np.ones((axes_len,), dtype="int64"), "int64")
|
|
else:
|
|
steps_tensor = _as_int64_tensor(bb, steps)
|
|
steps_len = _get_known_tensor_length(steps_tensor)
|
|
if steps_len is None:
|
|
raise ValueError("Slice requires a statically known steps length.")
|
|
if steps_len != axes_len:
|
|
raise ValueError(
|
|
f"Slice expects steps and starts to have the same length, but got "
|
|
f"{steps_len} and {axes_len}."
|
|
)
|
|
if isinstance(steps_tensor, relax.Constant) and _np.any(steps_tensor.data.numpy() == 0):
|
|
raise ValueError("Slice step values must be non-zero.")
|
|
|
|
axes_tensor = bb.normalize(
|
|
relax.op.where(
|
|
relax.op.less(axes_tensor, relax.const(0, "int64")),
|
|
relax.op.add(axes_tensor, relax.const(data_ndim, "int64")),
|
|
axes_tensor,
|
|
)
|
|
)
|
|
|
|
data_shape = bb.normalize(relax.op.shape_of(data_expr))
|
|
data_shape_tensor = bb.normalize(relax.op.shape_to_tensor(data_shape))
|
|
full_starts = relax.const(_np.zeros((data_ndim,), dtype="int64"), "int64")
|
|
full_steps = relax.const(_np.ones((data_ndim,), dtype="int64"), "int64")
|
|
full_starts = bb.normalize(
|
|
relax.op.scatter_elements(full_starts, axes_tensor, starts_tensor)
|
|
)
|
|
full_ends = bb.normalize(
|
|
relax.op.scatter_elements(data_shape_tensor, axes_tensor, ends_tensor)
|
|
)
|
|
full_steps = bb.normalize(relax.op.scatter_elements(full_steps, axes_tensor, steps_tensor))
|
|
return relax.op.dynamic_strided_slice(data_expr, full_starts, full_ends, full_steps)
|
|
|
|
|
|
class Pad(OnnxOpConverter):
|
|
"""Converts an onnx Pad node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v2(cls, bb, inputs, attr, params):
|
|
pads = attr.get("pads")
|
|
pads = relax.const(_np.array(pads), inputs[0].ty.shape[0].ty)
|
|
constant_value = attr.get("value")
|
|
if constant_value is None:
|
|
constant_value = 0.0
|
|
|
|
if isinstance(pads, relax.Constant):
|
|
pad_before, pad_after = _np.split(pads.data.numpy(), 2)
|
|
pad_before = _np.ndarray.tolist(pad_before)
|
|
pad_after = _np.ndarray.tolist(pad_after)
|
|
else:
|
|
raise ValueError("Dynamic pads are not supported yet.")
|
|
|
|
pad_mode = attr.get("mode", b"constant").decode("utf-8")
|
|
if pad_mode not in ["constant", "edge", "reflect"]:
|
|
raise tvm.error.OpAttributeInvalid(
|
|
"Value " + pad_mode + ' in attribute "mode" is invalid for operator Pad.'
|
|
)
|
|
|
|
if pad_mode == "constant":
|
|
return bb.emit_te(topi.nn.pad, inputs[0], pad_before, pad_after, constant_value)
|
|
elif pad_mode == "reflect":
|
|
return bb.emit_te(topi.nn.mirror_pad, inputs[0], pad_before, pad_after, "REFLECT")
|
|
else:
|
|
# edge mode - replicate border values
|
|
return bb.emit_te(topi.nn.replicate_pad, inputs[0], pad_before, pad_after)
|
|
|
|
@classmethod
|
|
def _impl_v11(cls, bb, inputs, attr, params):
|
|
pads = get_constant(inputs[1], params)
|
|
constant_value = get_constant(inputs[2], params)
|
|
if constant_value is not None:
|
|
constant_value = constant_value.data.numpy().item()
|
|
else:
|
|
constant_value = 0.0
|
|
|
|
if isinstance(pads, relax.Constant):
|
|
pad_before, pad_after = _np.split(pads.data.numpy(), 2)
|
|
pad_before = _np.ndarray.tolist(pad_before)
|
|
pad_after = _np.ndarray.tolist(pad_after)
|
|
else:
|
|
raise ValueError("Dynamic pads are not supported yet.")
|
|
|
|
pad_mode = attr.get("mode", b"constant").decode("utf-8")
|
|
if pad_mode not in ["constant", "edge", "reflect"]:
|
|
raise tvm.error.OpAttributeInvalid(
|
|
"Value " + pad_mode + ' in attribute "mode" is invalid for operator Pad.'
|
|
)
|
|
|
|
if pad_mode == "constant":
|
|
return bb.emit_te(topi.nn.pad, inputs[0], pad_before, pad_after, constant_value)
|
|
elif pad_mode == "reflect":
|
|
return bb.emit_te(topi.nn.mirror_pad, inputs[0], pad_before, pad_after, "REFLECT")
|
|
else:
|
|
# edge mode - replicate border values
|
|
return bb.emit_te(topi.nn.replicate_pad, inputs[0], pad_before, pad_after)
|
|
|
|
@classmethod
|
|
def _impl_v19(cls, bb, inputs, attr, params):
|
|
pads = get_constant(inputs[1], params)
|
|
constant_value = get_constant(inputs[2], params)
|
|
if constant_value is not None:
|
|
constant_value = constant_value.data.numpy().item()
|
|
else:
|
|
constant_value = 0.0
|
|
|
|
if isinstance(pads, relax.Constant):
|
|
pad_before, pad_after = _np.split(pads.data.numpy(), 2)
|
|
pad_before = _np.ndarray.tolist(pad_before)
|
|
pad_after = _np.ndarray.tolist(pad_after)
|
|
else:
|
|
raise ValueError("Dynamic pads are not supported yet.")
|
|
|
|
axes_input = inputs[3] if len(inputs) > 3 else None
|
|
if axes_input is not None:
|
|
axes_const = get_constant(axes_input, params)
|
|
if not isinstance(axes_const, relax.Constant):
|
|
raise ValueError("Dynamic axes are not supported for Pad yet.")
|
|
|
|
axes = axes_const.data.numpy().tolist()
|
|
if len(pad_before) != len(axes):
|
|
raise ValueError(
|
|
f"Pad expects pads length 2 * len(axes), got "
|
|
f"{len(pad_before) + len(pad_after)} pads and {len(axes)} axes."
|
|
)
|
|
|
|
rank = _get_known_tensor_rank(inputs[0])
|
|
if rank is None:
|
|
raise ValueError("Pad with axes requires a statically known input rank.")
|
|
|
|
axes = _normalize_constant_axes([int(a) for a in axes], rank, "Pad")
|
|
full_before = [0] * rank
|
|
full_after = [0] * rank
|
|
for i, ax in enumerate(axes):
|
|
full_before[ax] = pad_before[i]
|
|
full_after[ax] = pad_after[i]
|
|
pad_before, pad_after = full_before, full_after
|
|
|
|
pad_mode = attr.get("mode", b"constant").decode("utf-8")
|
|
if pad_mode not in ["constant", "edge", "reflect", "wrap"]:
|
|
raise tvm.error.OpAttributeInvalid(
|
|
"Value " + pad_mode + ' in attribute "mode" is invalid for operator Pad.'
|
|
)
|
|
|
|
if pad_mode == "constant":
|
|
return bb.emit_te(topi.nn.pad, inputs[0], pad_before, pad_after, constant_value)
|
|
elif pad_mode == "reflect":
|
|
return bb.emit_te(topi.nn.mirror_pad, inputs[0], pad_before, pad_after, "REFLECT")
|
|
elif pad_mode == "wrap":
|
|
return bb.emit_te(topi.nn.circular_pad, inputs[0], pad_before, pad_after)
|
|
else:
|
|
# edge mode - replicate border values
|
|
return bb.emit_te(topi.nn.replicate_pad, inputs[0], pad_before, pad_after)
|
|
|
|
|
|
class Tile(OnnxOpConverter):
|
|
"""Converts an onnx Tile node into an equivalent Relax expression."""
|
|
|
|
@staticmethod
|
|
def _tensor_length(expr):
|
|
shape = expr.ty.shape
|
|
if not isinstance(shape, relax.ShapeExpr):
|
|
return None
|
|
|
|
length = shape.values[0]
|
|
if not isinstance(length, tirx.IntImm):
|
|
return None
|
|
return length.value
|
|
|
|
@classmethod
|
|
def _impl_v13(cls, bb, inputs, attr, params):
|
|
reps = get_constant(inputs[1], params)
|
|
if isinstance(reps, relax.Constant):
|
|
reps = reps.data.numpy().tolist()
|
|
return bb.emit_te(topi.tile, inputs[0], reps)
|
|
|
|
data = inputs[0]
|
|
data_ndim = data.ty.ndim
|
|
reps_len = cls._tensor_length(reps)
|
|
if data_ndim == -1 or reps_len is None:
|
|
raise ValueError("Dynamic Tile requires known input rank and repeats length.")
|
|
|
|
if reps.ty.dtype != "int64":
|
|
reps = bb.normalize(relax.op.astype(reps, "int64"))
|
|
|
|
data_shape = bb.normalize(relax.op.shape_of(data))
|
|
data_shape_tensor = bb.normalize(relax.op.shape_to_tensor(data_shape))
|
|
output_shape_tensor = reps
|
|
|
|
if data_ndim > reps_len:
|
|
reps_prefix = relax.const(_np.ones((data_ndim - reps_len,), dtype="int64"), "int64")
|
|
output_shape_tensor = bb.normalize(
|
|
relax.op.concat([reps_prefix, output_shape_tensor], axis=0)
|
|
)
|
|
elif reps_len > data_ndim:
|
|
data_prefix = relax.const(_np.ones((reps_len - data_ndim,), dtype="int64"), "int64")
|
|
data_shape_tensor = bb.normalize(
|
|
relax.op.concat([data_prefix, data_shape_tensor], axis=0)
|
|
)
|
|
|
|
output_shape_tensor = bb.normalize(
|
|
relax.op.multiply(output_shape_tensor, data_shape_tensor)
|
|
)
|
|
output_shape = bb.normalize(relax.op.tensor_to_shape(output_shape_tensor))
|
|
output_shape_vars = [
|
|
tirx.Var(f"tile_dim_{i}", "int64") for i in range(max(data_ndim, reps_len))
|
|
]
|
|
bb.match_cast(output_shape, relax.ShapeType(output_shape_vars))
|
|
return bb.emit_te(topi.dyn_tile, data, output_shape_vars, reps_len)
|
|
|
|
|
|
class Expand(OnnxOpConverter):
|
|
"""Converts an onnx Expand node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v13(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
shape = inputs[1]
|
|
if isinstance(shape, relax.ShapeExpr):
|
|
data_shape = list(data.ty.shape)
|
|
target_shape = list(shape.values)
|
|
original_data_shape = [
|
|
dim.value if hasattr(dim, "value") else str(dim) for dim in data_shape
|
|
]
|
|
original_target_shape = [
|
|
dim.value if hasattr(dim, "value") else str(dim) for dim in target_shape
|
|
]
|
|
data_shape = [1] * (len(target_shape) - len(data_shape)) + data_shape
|
|
assert len(data_shape) == len(target_shape)
|
|
# Apply ONNX v13 Expand broadcasting rules
|
|
for i, s in enumerate(target_shape):
|
|
if isinstance(s, tvm.tirx.IntImm):
|
|
if s.value == -1:
|
|
# -1 means preserve the input dimension
|
|
target_shape[i] = data_shape[i]
|
|
elif isinstance(data_shape[i], tvm.tirx.IntImm) and data_shape[i].value == 1:
|
|
# Input dimension is 1, can broadcast to any target dimension >= 1
|
|
if s.value < 1:
|
|
raise ValueError(
|
|
f"ONNX Expand: Invalid target dimension {s.value} "
|
|
f"at possition {i}. Target dimensions must be >= 1."
|
|
)
|
|
elif (
|
|
isinstance(data_shape[i], tvm.tirx.IntImm)
|
|
and s.value == data_shape[i].value
|
|
):
|
|
# Dimensions match, no change needed
|
|
pass
|
|
elif s.value == 1:
|
|
# Target dimension is 1 but input dimension is not 1
|
|
# This would "squeeze" the dimension - preserve input for safety
|
|
target_shape[i] = data_shape[i]
|
|
else:
|
|
if isinstance(data_shape[i], tvm.tirx.IntImm):
|
|
raise ValueError(
|
|
f"ONNX Expand: Cannot broadcast input shape {original_data_shape} "
|
|
f"to target shape {original_target_shape}. "
|
|
f"At dimension {i}: input size {data_shape[i].value} is "
|
|
f"incompatible with target size {s.value}. "
|
|
f"ONNX broadcasting requires corresponding dimensions to have "
|
|
f"the same value or one of them to be 1."
|
|
)
|
|
# For dynamic shapes, let broadcast_to handle it
|
|
if target_shape == data_shape:
|
|
return data
|
|
return relax.op.broadcast_to(data, relax.ShapeExpr(target_shape))
|
|
|
|
# If possible, directly expand to constant shape.
|
|
if isinstance(shape, relax.Constant):
|
|
new_shape = shape.data.numpy().tolist()
|
|
# ONNX Expand operator requires preserving target rank and broadcasting
|
|
# according to standard rules. Dimensions are right-aligned.
|
|
data_shape = [dim.value for dim in data.ty.shape]
|
|
original_data_shape = data_shape.copy()
|
|
original_new_shape = new_shape.copy()
|
|
|
|
# Right-align the shapes
|
|
if len(new_shape) > len(data_shape):
|
|
data_shape = [1] * (len(new_shape) - len(data_shape)) + data_shape
|
|
else:
|
|
new_shape = [1] * (len(data_shape) - len(new_shape)) + new_shape
|
|
# Fix small target shapes - if target dim is smaller than input dim
|
|
# use the input dim (ONNX-specific behavior).
|
|
for i in range(len(new_shape)):
|
|
if new_shape[i] == -1:
|
|
# -1 means preserve the input dimension
|
|
new_shape[i] = data_shape[i]
|
|
elif data_shape[i] == 1:
|
|
# Input dimension is 1, can broadcast to any target dimension >= 1
|
|
if new_shape[i] < 1:
|
|
raise ValueError(
|
|
f"ONNX Expand: Invalid target dimension {new_shape[i]} "
|
|
f"at possition {i}. Target dimensions must be >= 1."
|
|
)
|
|
elif new_shape[i] == data_shape[i]:
|
|
# Dimensions match, no change needed
|
|
pass
|
|
elif new_shape[i] == 1:
|
|
# Target dimension is 1 but input dimension is not 1
|
|
# This would "squeeze" the dimension - preserve input for safety
|
|
new_shape[i] = data_shape[i]
|
|
else:
|
|
raise ValueError(
|
|
f"ONNX Expand: Cannot broadcast input shape {original_data_shape} "
|
|
f"to target shape {original_new_shape}. "
|
|
f"At dimension {i}: input size {data_shape[i]} is incompatible "
|
|
f"with target size {new_shape[i]}. "
|
|
f"ONNX broadcasting requires corresponding dimensions to have the same "
|
|
f"value or one of them to be 1."
|
|
)
|
|
return relax.op.broadcast_to(data, relax.ShapeExpr(new_shape))
|
|
|
|
# Otherwise handle dynamic shapes.
|
|
shape_ndim = next(dim.value for dim in shape.ty.shape.values)
|
|
shape_dataflow_var = bb.emit(
|
|
relax.Call(
|
|
relax.ExternFunc("vm.builtin.tensor_to_shape"),
|
|
[shape],
|
|
ty_args=[relax.ShapeType(ndim=shape_ndim)],
|
|
)
|
|
)
|
|
|
|
shape_vars = []
|
|
for i in range(shape_ndim):
|
|
shape_vars.append(tvm.tirx.Var(f"x_{i}", "int64"))
|
|
bb.match_cast(shape_dataflow_var, relax.ShapeType(shape_vars))
|
|
|
|
# Applying broadcasting rules for dynamic shapes
|
|
data_shape = list(data.ty.shape)
|
|
data_ndim = len(data_shape)
|
|
target_ndim = shape_ndim
|
|
padded_data = data
|
|
|
|
if target_ndim > data_ndim:
|
|
padded_data_shape = [tirx.IntImm("int64", 1)] * (target_ndim - data_ndim) + data_shape
|
|
padded_data = bb.normalize(relax.op.reshape(data, relax.ShapeExpr(padded_data_shape)))
|
|
|
|
return bb.normalize(relax.op.broadcast_to(padded_data, relax.ShapeExpr(shape_vars)))
|
|
|
|
|
|
class Attention(OnnxOpConverter):
|
|
"""Converts an onnx.microsoft Attention node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
num_heads = attr["num_heads"]
|
|
|
|
assert "do_rotary" not in attr, "rotary position embedding is not currently supported"
|
|
assert "past_present_share_buffer" not in attr, (
|
|
"past state for key and value is not currently supported"
|
|
)
|
|
assert "scale" not in attr, "custom scale is not currently supported"
|
|
assert "unidirectional" not in attr, "unidirectional attention is not currently supported"
|
|
|
|
if "mask_filter_value" in attr:
|
|
mask_filter_value = attr["mask_filter_value"]
|
|
else:
|
|
mask_filter_value = -10000.0
|
|
|
|
# (batch_size, sequence_length, input_hidden_size)
|
|
input_emb = bb.normalize(inputs[0])
|
|
|
|
# (input_hidden_size, hidden_size + hidden_size + v_hidden_size)
|
|
weight = bb.normalize(inputs[1])
|
|
|
|
def optional_input(k: int):
|
|
if inputs[k] is not None:
|
|
return bb.normalize(inputs[k])
|
|
else:
|
|
return None
|
|
|
|
# (hidden_size + hidden_size + v_hidden_size)
|
|
bias = optional_input(2)
|
|
|
|
# 1. ( batch_size, 1, max_seq_len, max_seq_len,)
|
|
# 2. ( batch_size, total_seq_len,)
|
|
# 3. ( batch_size, seq_len, total_seq_len,)
|
|
# 4. ( batch_size,)
|
|
# 5. (2 * batch_size,)
|
|
# For now, we only support case 2 & 3.
|
|
mask_index = optional_input(3)
|
|
|
|
# (2, batch_size, num_heads, past_sequence_length, head_size)
|
|
assert inputs[4] is None, "past state for key and value is not currently supported"
|
|
|
|
# (batch_size, num_heads, sequence_length, total_sequence_length)
|
|
qk_bias = optional_input(5)
|
|
|
|
assert inputs[6] is None, "past_sequence_length is not currently supported"
|
|
|
|
(batch_size, seq_len, input_hidden_size) = [val.value for val in input_emb.ty.shape.values]
|
|
weight_shape = [val.value for val in weight.ty.shape.values]
|
|
|
|
assert weight_shape[0] == input_hidden_size, (
|
|
"input and weight should share the same input hiden size"
|
|
)
|
|
|
|
if "qkv_hidden_sizes" in attr:
|
|
assert attr["qkv_hidden_sizes"][0] == attr["qkv_hidden_sizes"][1], (
|
|
"Q and K should share the same hidden sizes"
|
|
)
|
|
hidden_size, _, hidden_size_v = attr["qkv_hidden_sizes"]
|
|
else:
|
|
hidden_size = hidden_size_v = weight_shape[1] // 3
|
|
|
|
assert hidden_size % num_heads == 0, (
|
|
"hidden size should be divisible by number of attention heads"
|
|
)
|
|
head_size = hidden_size // num_heads
|
|
head_size_v = hidden_size_v // num_heads
|
|
|
|
if mask_index is not None:
|
|
mask_index_shape = [val.value for val in mask_index.ty.shape.values]
|
|
assert mask_index_shape in (
|
|
[batch_size, seq_len],
|
|
[
|
|
batch_size,
|
|
seq_len,
|
|
seq_len,
|
|
],
|
|
), """mask index should be in shape of (batch_size, seq_len),
|
|
or (batch_size, seq_len, seq_len)"""
|
|
mask_bias = relax.op.subtract(relax.const(1, dtype=mask_index.ty.dtype), mask_index)
|
|
mask_bias = relax.op.astype(mask_bias, dtype=input_emb.ty.dtype.dtype)
|
|
mask_bias = bb.normalize(
|
|
relax.op.multiply(
|
|
mask_bias,
|
|
relax.const(mask_filter_value, dtype=input_emb.ty.dtype),
|
|
)
|
|
)
|
|
if qk_bias is None:
|
|
qk_bias = mask_bias
|
|
else:
|
|
if len(mask_index_shape) == 2:
|
|
mask_bias = bb.normalize(
|
|
relax.op.reshape(mask_bias, [batch_size, 1, 1, seq_len])
|
|
)
|
|
elif len(mask_index_shape) == 3:
|
|
mask_bias = bb.normalize(
|
|
relax.op.reshape(mask_bias, [batch_size, 1, seq_len, seq_len])
|
|
)
|
|
qk_bias = bb.normalize(relax.op.add(qk_bias, mask_bias))
|
|
|
|
QKV = relax.op.matmul(input_emb, weight)
|
|
|
|
if bias:
|
|
bias_shape = [val.value for val in bias.ty.shape.values]
|
|
assert bias_shape[0] == weight_shape[1], (
|
|
"bias and weight should share the same hidden size sum"
|
|
)
|
|
QKV = relax.op.add(QKV, bias)
|
|
|
|
QKV = relax.op.split(QKV, [hidden_size, hidden_size * 2], 2)
|
|
Q, K, V = QKV[0], QKV[1], QKV[2]
|
|
|
|
Q = bb.normalize(relax.op.reshape(Q, (batch_size, seq_len, num_heads, head_size)))
|
|
K = bb.normalize(relax.op.reshape(K, (batch_size, seq_len, num_heads, head_size)))
|
|
V = bb.normalize(relax.op.reshape(V, (batch_size, seq_len, num_heads, head_size_v)))
|
|
output = relax.op.nn.attention(Q, K, V, qk_bias)
|
|
output = bb.normalize(
|
|
relax.op.reshape(output, (batch_size, seq_len, num_heads * head_size_v))
|
|
)
|
|
# add placeholder for optional present state supported in the future
|
|
placeholder = relax.const(0, dtype="float32")
|
|
return relax.Tuple([output, placeholder])
|
|
|
|
|
|
class Identity(OnnxOpConverter):
|
|
"""Converts an onnx Identity node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
return inputs[0]
|
|
|
|
|
|
class Dropout(OnnxOpConverter):
|
|
"""Converts an onnx Dropout node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
ratio = float(attr.get("ratio", 0.5))
|
|
return relax.op.nn.dropout(inputs[0], ratio)
|
|
|
|
@classmethod
|
|
def _impl_v12(cls, bb, inputs, attr, params):
|
|
# Since opset 12 ratio is the optional second input rather than an attribute.
|
|
ratio = 0.5
|
|
if len(inputs) >= 2 and inputs[1] is not None:
|
|
const = get_constant(inputs[1], params)
|
|
if isinstance(const, relax.Constant):
|
|
ratio = float(const.data.numpy())
|
|
return relax.op.nn.dropout(inputs[0], ratio)
|
|
|
|
|
|
def _onnx_resize_spatial_roi_vector(roi_full: relax.Expr, rank: int) -> relax.Expr:
|
|
"""Map ONNX ROI [starts..., ends...] to TOPI spatial ROI (drop N/C axes)."""
|
|
return relax.op.concat(
|
|
[
|
|
relax.op.strided_slice(roi_full, axes=[0], begin=[2], end=[rank]),
|
|
relax.op.strided_slice(roi_full, axes=[0], begin=[rank + 2], end=[2 * rank]),
|
|
],
|
|
axis=0,
|
|
)
|
|
|
|
|
|
def _topi_resize3d_roi_from_onnx_ncdhw_spatial(roi_spatial: list[float]) -> list[float]:
|
|
"""Reorder spatial ROI for NCDHW ONNX layout to TOPI resize3d convention.
|
|
|
|
ONNX spatial slice after dropping N/C is ordered (D, H, W) for starts then ends.
|
|
TOPI ``resize3d`` with layout NCDHW expects
|
|
``(start_w, start_h, start_d, end_w, end_h, end_d)`` (see topi/image/resize.py).
|
|
"""
|
|
if len(roi_spatial) != 6:
|
|
return roi_spatial
|
|
d0, h0, w0, d1, h1, w1 = roi_spatial
|
|
return [w0, h0, d0, w1, h1, d1]
|
|
|
|
|
|
def _emit_resize_topi_dynamic_roi(
|
|
bb: relax.BlockBuilder,
|
|
data: relax.Expr,
|
|
roi_spatial_vec: relax.Expr,
|
|
sizes_spatial: list,
|
|
rank: int,
|
|
topi_mode: str,
|
|
coord_mode: str,
|
|
rounding_method: str,
|
|
cubic_coeff_a: float,
|
|
exclude_outside: int,
|
|
extrapolation_value: float,
|
|
) -> relax.Expr:
|
|
"""Lower Resize with runtime ROI via TOPI, which supports Expr ROI."""
|
|
if rank == 3:
|
|
|
|
def resize1d_dyn(d, r, s0):
|
|
return topi.image.resize1d(
|
|
d,
|
|
(r[0], r[1]),
|
|
[s0],
|
|
"NCW",
|
|
topi_mode,
|
|
coord_mode,
|
|
rounding_method,
|
|
cubic_coeff_a,
|
|
exclude_outside,
|
|
extrapolation_value,
|
|
)
|
|
|
|
return bb.emit_te(resize1d_dyn, data, roi_spatial_vec, sizes_spatial[0])
|
|
|
|
if rank == 4:
|
|
|
|
def resize2d_dyn(d, r, s0, s1):
|
|
return topi.image.resize2d(
|
|
d,
|
|
(r[0], r[1], r[2], r[3]),
|
|
(s0, s1),
|
|
layout="NCHW",
|
|
method=topi_mode,
|
|
coordinate_transformation_mode=coord_mode,
|
|
rounding_method=rounding_method,
|
|
bicubic_alpha=cubic_coeff_a,
|
|
bicubic_exclude=exclude_outside,
|
|
extrapolation_value=extrapolation_value,
|
|
)
|
|
|
|
return bb.emit_te(resize2d_dyn, data, roi_spatial_vec, sizes_spatial[0], sizes_spatial[1])
|
|
|
|
def resize3d_dyn(d, r, s0, s1, s2):
|
|
# r is ONNX order (D,H,W) x2; TOPI expects (W,H,D) x2.
|
|
return topi.image.resize3d(
|
|
d,
|
|
(r[2], r[1], r[0], r[5], r[4], r[3]),
|
|
(s0, s1, s2),
|
|
layout="NCDHW",
|
|
method=topi_mode,
|
|
coordinate_transformation_mode=coord_mode,
|
|
rounding_method=rounding_method,
|
|
bicubic_alpha=cubic_coeff_a,
|
|
bicubic_exclude=exclude_outside,
|
|
extrapolation_value=extrapolation_value,
|
|
)
|
|
|
|
return bb.emit_te(
|
|
resize3d_dyn,
|
|
data,
|
|
roi_spatial_vec,
|
|
sizes_spatial[0],
|
|
sizes_spatial[1],
|
|
sizes_spatial[2],
|
|
)
|
|
|
|
|
|
class Resize(OnnxOpConverter):
|
|
"""Converts an onnx Resize node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v18(cls, bb, inputs, attr, params):
|
|
# Extract the many attributes of resize.
|
|
coord_mode = attr.get("coordinate_transformation_mode", b"half_pixel").decode("ascii")
|
|
cubic_coeff_a = attr.get("cubic_coeff_a", -0.75)
|
|
exclude_outside = attr.get("exclude_outside", 0)
|
|
extrapolation_value = attr.get("extrapolation_value", 0.0)
|
|
mode = attr.get("mode", b"nearest").decode("ascii")
|
|
rounding_method = attr.get("nearest_mode", b"round_prefer_floor").decode("ascii")
|
|
|
|
# Adapt attributes to fit TVM definition.
|
|
if mode == "nearest":
|
|
relax_mode = "nearest_neighbor"
|
|
else:
|
|
relax_mode = mode
|
|
topi_mode = relax_mode
|
|
|
|
# Unpack inputs.
|
|
x = inputs[0]
|
|
roi = get_constant(inputs[1], params) if len(inputs) > 1 and inputs[1] is not None else None
|
|
scales = get_constant(inputs[2], params) if len(inputs) > 2 else None
|
|
sizes = get_constant(inputs[3], params) if len(inputs) > 3 else None
|
|
ndims = len(x.ty.shape)
|
|
assert ndims in (3, 4, 5), "Only resize1d/resize2d/resize3d are supported."
|
|
|
|
assert scales is None or sizes is None, (
|
|
"Only one of scales and sizes can be provided in Resize."
|
|
)
|
|
|
|
# ROI can be a static list (for relax.image.resize*) or dynamic tensor (TOPI path).
|
|
roi_static: list[float] | None = None
|
|
roi_dynamic_vec: relax.Expr | None = None
|
|
if roi is not None:
|
|
if isinstance(roi, relax.Constant):
|
|
roi_np = roi.data.numpy().tolist()
|
|
if len(roi_np) == 2 * ndims:
|
|
roi_static = roi_np[2:ndims] + roi_np[ndims + 2 : 2 * ndims]
|
|
elif len(roi_np) == 0:
|
|
roi_static = [0.0] * (2 * (ndims - 2))
|
|
elif len(roi_np) == 2 * (ndims - 2):
|
|
# Some exporters already provide spatial-only ROI.
|
|
roi_static = roi_np
|
|
else:
|
|
roi_static = roi_np
|
|
else:
|
|
roi_dynamic_vec = bb.normalize(_onnx_resize_spatial_roi_vector(roi, ndims))
|
|
else:
|
|
roi_static = [0.0] * (2 * (ndims - 2))
|
|
|
|
use_dynamic_roi = roi_dynamic_vec is not None
|
|
|
|
# Convert scales to sizes if needed.
|
|
if scales is not None:
|
|
if isinstance(scales, relax.Constant):
|
|
scales = scales.data.numpy()
|
|
elif isinstance(scales, relax.expr.ShapeExpr):
|
|
scales = [int(val.value) for val in scales.values]
|
|
else:
|
|
raise ValueError(f"Type {type(scales)} for scale is currently unsupported.")
|
|
sizes = []
|
|
|
|
for i, dim in enumerate(x.ty.shape):
|
|
sizes.append((scales[i] * dim).astype("int64"))
|
|
sizes = sizes[2:]
|
|
else:
|
|
if isinstance(sizes, relax.Constant):
|
|
sizes = sizes.data.numpy().astype("int64").tolist()[2:]
|
|
elif isinstance(sizes, relax.expr.ShapeExpr):
|
|
sizes = [int(val.value) for val in sizes.values][2:]
|
|
else:
|
|
raise ValueError(f"Type {type(sizes)} for size is currently unsupported.")
|
|
|
|
if use_dynamic_roi:
|
|
return _emit_resize_topi_dynamic_roi(
|
|
bb,
|
|
x,
|
|
roi_dynamic_vec,
|
|
sizes,
|
|
ndims,
|
|
topi_mode,
|
|
coord_mode,
|
|
rounding_method,
|
|
cubic_coeff_a,
|
|
exclude_outside,
|
|
extrapolation_value,
|
|
)
|
|
|
|
if ndims == 3:
|
|
return bb.emit_te(
|
|
topi.image.resize1d,
|
|
x,
|
|
roi_static,
|
|
sizes,
|
|
"NCW",
|
|
topi_mode,
|
|
coord_mode,
|
|
rounding_method,
|
|
cubic_coeff_a,
|
|
exclude_outside,
|
|
extrapolation_value,
|
|
)
|
|
elif ndims == 4:
|
|
return relax.op.image.resize2d(
|
|
x,
|
|
size=relax.ShapeExpr(sizes),
|
|
roi=roi_static,
|
|
layout="NCHW",
|
|
method=relax_mode,
|
|
coordinate_transformation_mode=coord_mode,
|
|
rounding_method=rounding_method,
|
|
cubic_alpha=cubic_coeff_a,
|
|
cubic_exclude=exclude_outside,
|
|
extrapolation_value=extrapolation_value,
|
|
)
|
|
else: # ndims == 5
|
|
roi3d = _topi_resize3d_roi_from_onnx_ncdhw_spatial(roi_static)
|
|
return relax.op.image.resize3d(
|
|
x,
|
|
size=relax.ShapeExpr(sizes),
|
|
roi=roi3d,
|
|
layout="NCDHW",
|
|
method=relax_mode,
|
|
coordinate_transformation_mode=coord_mode,
|
|
rounding_method=rounding_method,
|
|
cubic_alpha=cubic_coeff_a,
|
|
cubic_exclude=exclude_outside,
|
|
extrapolation_value=extrapolation_value,
|
|
)
|
|
|
|
|
|
class AffineGrid(OnnxOpConverter):
|
|
"""Converts an onnx AffineGrid node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v20(cls, bb, inputs, attr, params):
|
|
theta = inputs[0] # [N, 2, 3] for 2D
|
|
size = get_constant(inputs[1], params) # [N, C, H, W] for 2D
|
|
align_corners = bool(attr.get("align_corners", 0))
|
|
|
|
# Extract size values
|
|
if isinstance(size, relax.Constant):
|
|
size_vals = size.data.numpy().astype("int64").tolist()
|
|
elif isinstance(size, relax.expr.ShapeExpr):
|
|
size_vals = [int(v.value) for v in size.values]
|
|
else:
|
|
raise NotImplementedError(f"Dynamic size of type {type(size)} is not supported")
|
|
|
|
if len(size_vals) not in (4, 5):
|
|
raise ValueError("AffineGrid expects size to be [N,C,H,W] (2D) or [N,C,D,H,W] (3D)")
|
|
|
|
# relax affine_grid outputs [N, spatial, *spatial_dims]; move the coord axis
|
|
# last to match the ONNX convention [N, *spatial_dims, spatial].
|
|
grid = bb.emit(relax.op.image.affine_grid(theta, tuple(size_vals[2:]), align_corners))
|
|
axes = [0, *range(2, len(size_vals)), 1]
|
|
return bb.emit(relax.op.permute_dims(grid, axes=axes))
|
|
|
|
|
|
class Einsum(OnnxOpConverter):
|
|
"""Converts an onnx Einsum node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v12(cls, bb, inputs, attr, params):
|
|
equation = attr["equation"].decode("utf-8")
|
|
return bb.emit_te(topi.einsum, equation, *inputs)
|
|
|
|
|
|
class RoiAlign(OnnxOpConverter):
|
|
"""Converts an onnx RoiAlign node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl(cls, bb, inputs, attr, params, default_coordinate_transformation_mode):
|
|
if len(inputs) != 3:
|
|
raise ValueError("RoiAlign expects exactly 3 inputs")
|
|
|
|
data = inputs[0]
|
|
rois = inputs[1]
|
|
batch_indices = inputs[2]
|
|
rois_dtype = rois.ty.dtype
|
|
|
|
mode = attr.get("mode", b"avg")
|
|
if isinstance(mode, bytes):
|
|
mode = mode.decode("ascii")
|
|
if mode not in ("avg", "max"):
|
|
raise NotImplementedError("RoiAlign in Relax only supports avg and max modes")
|
|
|
|
output_height = attr.get("output_height", 1)
|
|
output_width = attr.get("output_width", 1)
|
|
sampling_ratio = attr.get("sampling_ratio", 0)
|
|
spatial_scale = attr.get("spatial_scale", 1.0)
|
|
coordinate_transformation_mode = attr.get(
|
|
"coordinate_transformation_mode", default_coordinate_transformation_mode
|
|
)
|
|
if isinstance(coordinate_transformation_mode, bytes):
|
|
coordinate_transformation_mode = coordinate_transformation_mode.decode("ascii")
|
|
|
|
if coordinate_transformation_mode == "half_pixel":
|
|
offset = relax.const([-0.5, -0.5, -0.5, -0.5], rois_dtype)
|
|
rois = relax.op.add(rois, offset)
|
|
aligned = True
|
|
elif coordinate_transformation_mode != "output_half_pixel":
|
|
raise NotImplementedError(
|
|
"RoiAlign only supports coordinate_transformation_mode "
|
|
"'half_pixel' and 'output_half_pixel'"
|
|
)
|
|
else:
|
|
aligned = False
|
|
|
|
batch_indices = relax.op.expand_dims(batch_indices, axis=1)
|
|
batch_indices = relax.op.astype(batch_indices, rois_dtype)
|
|
rois = relax.op.concat([batch_indices, rois], axis=1)
|
|
|
|
return relax.op.vision.roi_align(
|
|
data,
|
|
rois,
|
|
pooled_size=(output_height, output_width),
|
|
spatial_scale=spatial_scale,
|
|
sample_ratio=sampling_ratio,
|
|
aligned=aligned,
|
|
layout="NCHW",
|
|
mode=mode,
|
|
)
|
|
|
|
@classmethod
|
|
def _impl_v10(cls, bb, inputs, attr, params):
|
|
return cls._impl(bb, inputs, attr, params, b"output_half_pixel")
|
|
|
|
@classmethod
|
|
def _impl_v16(cls, bb, inputs, attr, params):
|
|
return cls._impl(bb, inputs, attr, params, b"half_pixel")
|
|
|
|
|
|
class MaxRoiPool(OnnxOpConverter):
|
|
"""Converts an onnx MaxRoiPool node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
if len(inputs) != 2:
|
|
raise ValueError("MaxRoiPool expects exactly 2 inputs")
|
|
|
|
pooled_shape = attr.get("pooled_shape")
|
|
if pooled_shape is None:
|
|
raise ValueError("MaxRoiPool requires pooled_shape attribute")
|
|
|
|
spatial_scale = attr.get("spatial_scale", 1.0)
|
|
return relax.op.vision.roi_pool(
|
|
inputs[0],
|
|
inputs[1],
|
|
pooled_size=tuple(pooled_shape),
|
|
spatial_scale=spatial_scale,
|
|
layout="NCHW",
|
|
)
|
|
|
|
|
|
class Range(OnnxOpConverter):
|
|
"""Converts an onnx Range node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v12(cls, bb, inputs, attr, params):
|
|
start = get_constant(inputs[0], params)
|
|
limit = get_constant(inputs[1], params)
|
|
delta = get_constant(inputs[2], params)
|
|
out_dtype = start.ty.dtype
|
|
|
|
if isinstance(start, relax.Constant):
|
|
start = start.data.numpy().tolist()
|
|
|
|
if isinstance(limit, relax.Constant):
|
|
limit = limit.data.numpy().tolist()
|
|
|
|
assert isinstance(delta, relax.Constant), "Constant delta required for Range."
|
|
step = delta.data.numpy().tolist()
|
|
|
|
# If all inputs are constant, compute directly.
|
|
if isinstance(start, int) and isinstance(limit, int):
|
|
out_range = _np.arange(start=start, stop=limit, step=step)
|
|
return relax.const(out_range, out_dtype)
|
|
|
|
# Otherwise compute in graph.
|
|
return relax.op.arange(start, limit, step, out_dtype)
|
|
|
|
|
|
class InstanceNormalization(OnnxOpConverter):
|
|
"""Converts an onnx InstanceNormalization node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v6(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
scale = inputs[1]
|
|
B = inputs[2]
|
|
epsilon = attr.get("epsilon", 1e-05)
|
|
epsilon = relax.const(epsilon, dtype=data.ty.dtype)
|
|
|
|
ndim = len(data.ty.shape)
|
|
redux_axes = list(range(2, ndim))
|
|
|
|
mean = relax.op.mean(data, axis=redux_axes, keepdims=True)
|
|
var = relax.op.variance(data, axis=redux_axes, keepdims=True)
|
|
sqrt = relax.op.sqrt(relax.op.add(var, epsilon))
|
|
out = relax.op.divide(relax.op.subtract(data, mean), sqrt)
|
|
broadcast_shape = [-1] + [
|
|
1,
|
|
] * (ndim - 2)
|
|
if scale is not None:
|
|
scale = relax.op.reshape(scale, broadcast_shape)
|
|
out = relax.op.multiply(out, scale)
|
|
if B is not None:
|
|
B = relax.op.reshape(B, broadcast_shape)
|
|
out = relax.op.add(out, B)
|
|
return out
|
|
|
|
|
|
class BatchNormalization(OnnxOpConverter):
|
|
"""Converts an onnx BatchNormalization node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v15(cls, bb, inputs, attr, params):
|
|
# Unpack inputs
|
|
data = inputs[0]
|
|
scale = inputs[1]
|
|
bias = inputs[2]
|
|
mean = inputs[3]
|
|
var = inputs[4]
|
|
epsilon = attr.get("epsilon", 1e-05)
|
|
momentum = attr.get("momentum", 0.9)
|
|
training_mode = attr.get("training_mode", 0)
|
|
|
|
data_dtype = data.ty.dtype
|
|
scale_dtype = scale.ty.dtype
|
|
bias_dtype = bias.ty.dtype
|
|
mean_dtype = mean.ty.dtype
|
|
var_dtype = var.ty.dtype
|
|
|
|
if scale_dtype != bias_dtype:
|
|
raise ValueError(
|
|
"ONNX BatchNormalization requires scale and bias to have the same "
|
|
f"dtype, but received {scale_dtype} and {bias_dtype}."
|
|
)
|
|
|
|
if mean_dtype != var_dtype:
|
|
raise ValueError(
|
|
"ONNX BatchNormalization requires mean and var to have the same "
|
|
f"dtype, but received {mean_dtype} and {var_dtype}."
|
|
)
|
|
|
|
if data_dtype == scale_dtype == mean_dtype:
|
|
compute_dtype = data_dtype
|
|
elif (
|
|
data_dtype == "float16"
|
|
and scale_dtype in ("float16", "float32")
|
|
and mean_dtype in ("float16", "float32")
|
|
):
|
|
compute_dtype = "float32"
|
|
else:
|
|
raise NotImplementedError(
|
|
"ONNX BatchNormalization with mixed input dtypes is currently "
|
|
"supported only for float16 data with float16/float32 parameters "
|
|
"and statistics, but received "
|
|
f"data={data_dtype}, scale/bias={scale_dtype}, mean/var={mean_dtype}."
|
|
)
|
|
|
|
# ONNX requires float computation for float16 training statistics to avoid overflow.
|
|
if training_mode and data_dtype == "float16":
|
|
compute_dtype = "float32"
|
|
|
|
def cast_for_compute(expr, source_dtype):
|
|
if source_dtype == compute_dtype:
|
|
return expr
|
|
return relax.op.astype(expr, compute_dtype)
|
|
|
|
output = relax.op.nn.batch_norm(
|
|
cast_for_compute(data, data_dtype),
|
|
gamma=cast_for_compute(scale, scale_dtype),
|
|
beta=cast_for_compute(bias, bias_dtype),
|
|
moving_mean=cast_for_compute(mean, mean_dtype),
|
|
moving_var=cast_for_compute(var, var_dtype),
|
|
axis=1,
|
|
epsilon=epsilon,
|
|
momentum=momentum,
|
|
training=bool(training_mode),
|
|
)
|
|
|
|
y = relax.TupleGetItem(output, 0)
|
|
running_mean = relax.TupleGetItem(output, 1)
|
|
running_var = relax.TupleGetItem(output, 2)
|
|
|
|
if compute_dtype != data_dtype:
|
|
y = relax.op.astype(y, data_dtype)
|
|
if compute_dtype != mean_dtype:
|
|
running_mean = relax.op.astype(running_mean, mean_dtype)
|
|
if compute_dtype != var_dtype:
|
|
running_var = relax.op.astype(running_var, var_dtype)
|
|
|
|
return relax.Tuple([y, running_mean, running_var])
|
|
|
|
|
|
class MeanVarianceNormalization(OnnxOpConverter):
|
|
"""Converts an onnx MeanVarianceNormalization node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v9(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
axis = attr.get("axes", (0, 2, 3))
|
|
data_mean = relax.op.mean(data, axis=axis, keepdims=True)
|
|
data_mean_squared = relax.op.power(data_mean, relax.const(2, dtype="float32"))
|
|
data_squared = relax.op.power(data, relax.const(2, dtype="float32"))
|
|
data_squared_mean = relax.op.mean(data_squared, axis=axis, keepdims=True)
|
|
return (data - data_mean) / relax.op.sqrt(data_squared_mean - data_mean_squared)
|
|
|
|
|
|
class LocalResponseNormalization(OnnxOpConverter):
|
|
"""Converts an onnx LocalResponseNormalization node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v13(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
size = attr["size"]
|
|
alpha = attr.get("alpha", 0.0001)
|
|
beta = attr.get("beta", 0.75)
|
|
bias = attr.get("bias", 1.0)
|
|
|
|
if hasattr(data.ty, "ndim"):
|
|
ndim = data.ty.ndim
|
|
else:
|
|
ndim = len(data.ty.shape)
|
|
|
|
if ndim not in [3, 4]:
|
|
raise ValueError(f"LRN only supports 3D or 4D input, got {ndim}D.")
|
|
|
|
data_squared = relax.op.multiply(data, data)
|
|
data_expanded = relax.op.expand_dims(data_squared, axis=1)
|
|
pad_len = size // 2
|
|
if ndim == 3:
|
|
pool_padding = [pad_len, 0, pad_len, 0]
|
|
pool_op = relax.op.nn.avg_pool2d
|
|
pool_size = (size, 1)
|
|
layout = "NCHW"
|
|
strides = (1, 1)
|
|
else:
|
|
pool_padding = [pad_len, 0, 0, pad_len, 0, 0]
|
|
pool_op = relax.op.nn.avg_pool3d
|
|
pool_size = (size, 1, 1)
|
|
layout = "NCDHW"
|
|
strides = (1, 1, 1)
|
|
|
|
data_avgpool = pool_op(
|
|
data_expanded,
|
|
pool_size=pool_size,
|
|
strides=strides,
|
|
padding=pool_padding,
|
|
layout=layout,
|
|
ceil_mode=False,
|
|
count_include_pad=True,
|
|
)
|
|
data_squeezed = relax.op.squeeze(data_avgpool, axis=1)
|
|
|
|
const_alpha = relax.const(alpha, dtype="float32")
|
|
const_bias = relax.const(bias, dtype="float32")
|
|
const_beta = relax.const(beta, dtype="float32")
|
|
|
|
scale = relax.op.multiply(data_squeezed, const_alpha)
|
|
scale = relax.op.add(scale, const_bias)
|
|
denominator = relax.op.power(scale, const_beta)
|
|
|
|
return relax.op.divide(data, denominator)
|
|
|
|
|
|
class Pool(OnnxOpConverter):
|
|
"""A helper class for pool op converters."""
|
|
|
|
name = ""
|
|
|
|
@classmethod
|
|
def get_pad_pair(cls, input1d, kernel1d, stride1d, mode):
|
|
"""infer pad size"""
|
|
if input1d % stride1d == 0:
|
|
pad = max(kernel1d - stride1d, 0)
|
|
else:
|
|
pad = max(kernel1d - (input1d % stride1d), 0)
|
|
pad_before = pad // 2
|
|
pad_after = pad - pad_before
|
|
if "LOWER" in mode:
|
|
return [pad_after, pad_before]
|
|
return [pad_before, pad_after]
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
# Unpack inputs and attributes.
|
|
data = inputs[0]
|
|
input_shape = data.ty.shape
|
|
ndim = len(input_shape)
|
|
|
|
auto_pad = attr.get("auto_pad", b"NOTSET").decode("utf-8")
|
|
ceil_mode = attr.get("ceil_mode", 0)
|
|
dilations = attr.get("dilations", [1] * (ndim - 2))
|
|
kernel_shape = attr.get("kernel_shape")
|
|
pads = attr.get("pads", 0)
|
|
strides = attr.get("strides", [1] * (ndim - 2))
|
|
count_include_pad = attr.get("count_include_pad", False)
|
|
|
|
assert len(kernel_shape) in [1, 2, 3], "Currently only 1D/2D/3D/ pooling is supported."
|
|
|
|
assert auto_pad in [
|
|
"NOTSET",
|
|
"SAME_UPPER",
|
|
"SAME_LOWER",
|
|
"VALID",
|
|
], f"Value {auto_pad} in attribute auto_pad is invalid."
|
|
|
|
if auto_pad in ("SAME_UPPER", "SAME_LOWER"):
|
|
pads = []
|
|
if cls.name == "avg_pool":
|
|
for axis in range(len(input_shape) - 2):
|
|
axis_shape = int(input_shape[2 + axis])
|
|
stride = strides[axis]
|
|
kernel = kernel_shape[axis]
|
|
pad = cls.get_pad_pair(axis_shape, kernel, stride, auto_pad)
|
|
pads.append(pad)
|
|
else:
|
|
input_spatial_shape = cls._get_input_spatial_shape(data)
|
|
output_spatial_shape = [0 for _ in input_spatial_shape]
|
|
|
|
for i, _ in enumerate(input_spatial_shape):
|
|
if auto_pad == "SAME_UPPER":
|
|
output_spatial_shape[i] = int(_np.ceil(input_spatial_shape[i] / strides[i]))
|
|
else:
|
|
output_spatial_shape[i] = int(
|
|
_np.floor(input_spatial_shape[i] / strides[i])
|
|
)
|
|
pad_i = (
|
|
(output_spatial_shape[i] - 1) * strides[i]
|
|
+ ((kernel_shape[i] - 1) * dilations[i] + 1)
|
|
- input_spatial_shape[i]
|
|
)
|
|
|
|
if auto_pad == "SAME_UPPER":
|
|
pads.append([pad_i // 2, pad_i - pad_i // 2])
|
|
else:
|
|
pads.append([pad_i - pad_i // 2, pad_i // 2])
|
|
|
|
pads = tuple([val for pair in zip(*pads) for val in pair])
|
|
|
|
op = getattr(relax.op.nn, cls.name + str(len(kernel_shape)) + "d")
|
|
return op(data, kernel_shape, strides, pads, dilations, ceil_mode, count_include_pad)
|
|
|
|
@classmethod
|
|
def _get_input_spatial_shape(cls, tensor):
|
|
# shape is (N x C x D1 x D2 ... Dn)
|
|
return _np.array([int(d) for d in tensor.ty.shape], dtype="int64")[2:]
|
|
|
|
|
|
class MaxPool(Pool):
|
|
"""Converts an onnx MaxPool node into an equivalent Relax expression."""
|
|
|
|
name = "max_pool"
|
|
|
|
|
|
class AveragePool(Pool):
|
|
"""Converts an onnx MaxPool node into an equivalent Relax expression."""
|
|
|
|
name = "avg_pool"
|
|
|
|
|
|
class LpPool(OnnxOpConverter):
|
|
"""Converts an onnx LpPool node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
dtype = inputs[0].ty.dtype
|
|
p = attr.get("p", 2.0)
|
|
reci_p = relax.const(1.0 / p, dtype=dtype)
|
|
# emit for get ty
|
|
data = bb.emit(relax.op.power(inputs[0], relax.const(p, dtype=dtype)))
|
|
attr.update({"count_include_pad": True})
|
|
avg_pool = AveragePool._impl_v1(bb, [data], attr, params)
|
|
kernels = attr["kernel_shape"]
|
|
out = avg_pool * relax.const(_np.prod(kernels).astype(dtype))
|
|
return relax.op.power(out, reci_p)
|
|
|
|
|
|
class GlobalAveragePool(OnnxOpConverter):
|
|
"""Converts an onnx GlobalAveragePool node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
rank = len(inputs[0].ty.shape)
|
|
axes = list(range(2, rank))
|
|
return relax.op.mean(inputs[0], axis=axes, keepdims=True)
|
|
|
|
|
|
class GlobalMaxPool(OnnxOpConverter):
|
|
"""Converts an onnx GlobalMaxPool node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
rank = len(inputs[0].ty.shape)
|
|
axes = list(range(2, rank))
|
|
return relax.op.max(inputs[0], axis=axes, keepdims=True)
|
|
|
|
|
|
class GlobalLpPool(OnnxOpConverter):
|
|
"""Converts an onnx GlobalLpPool node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v2(cls, bb, inputs, attr, params):
|
|
p = attr.get("p", 2.0)
|
|
dtype = inputs[0].ty.dtype
|
|
rank = len(inputs[0].ty.shape)
|
|
axes = list(range(2, rank))
|
|
x_abs = relax.op.abs(inputs[0])
|
|
x_p = relax.op.power(x_abs, relax.const(p, dtype=dtype))
|
|
x_sum = relax.op.sum(x_p, axes, keepdims=True)
|
|
return relax.op.power(x_sum, relax.const(1.0 / p, dtype=dtype))
|
|
|
|
|
|
class MaxUnpool(OnnxOpConverter):
|
|
"""Converts an onnx MaxUnpool node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v9(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
indices = inputs[1]
|
|
output_shape = inputs[2]
|
|
kernel_shape = attr.get("kernel_shape")
|
|
pads = attr.get("pads", [0] * len(kernel_shape) * 2)
|
|
strides = attr.get("strides", [1] * len(kernel_shape))
|
|
|
|
multiplier = _np.concatenate([[1, 1], list(strides)])
|
|
shape = [v.value for v in data.ty.shape]
|
|
total_output_shape = multiplier * shape
|
|
# Add extra dimensions from kernel size and stride mismatch
|
|
total_output_shape += _np.concatenate([[0, 0], list(kernel_shape)], axis=0)
|
|
total_output_shape -= _np.concatenate([[0, 0], list(strides)], axis=0)
|
|
|
|
if output_shape is not None:
|
|
total_output_shape = output_shape
|
|
|
|
elif pads is not None:
|
|
# Get pads in the proper format
|
|
pads = _np.concatenate([[0, 0, 0, 0], list(pads)], axis=0)
|
|
pads = _np.reshape(pads, [-1, 2])
|
|
# Compute the total padding per axis.
|
|
total_pad = _np.sum(pads, axis=-1)
|
|
# Reversing maxpool means that padding actually makes our output smaller.
|
|
total_output_shape = total_output_shape - total_pad
|
|
|
|
# Create a tensor of zeros then scatter our data through it.
|
|
relax_shape = relax.ShapeExpr(total_output_shape.tolist())
|
|
zeros_tensor = bb.emit(relax.op.zeros(relax_shape, data.ty.dtype.dtype))
|
|
# We need to flatten all our tensors before scattering.
|
|
flat_tensor = relax.op.scatter_elements(
|
|
relax.op.reshape(zeros_tensor, [-1]),
|
|
relax.op.reshape(indices, [-1]),
|
|
relax.op.reshape(data, [-1]),
|
|
axis=0,
|
|
)
|
|
# Reshape our flattened data back to normal.
|
|
output = relax.op.reshape(flat_tensor, relax_shape)
|
|
return output
|
|
|
|
|
|
class Flatten(OnnxOpConverter):
|
|
"""Converts an onnx Flatten node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v13(cls, bb, inputs, attr, params):
|
|
axis = attr.get("axis", 1)
|
|
data_shape = list(inputs[0].ty.shape)
|
|
|
|
if axis == 0:
|
|
new_shape = (1, -1)
|
|
else:
|
|
shape_flags = [isinstance(x, tvm.script.tirx.IntImm) for x in data_shape[0:axis]]
|
|
|
|
if all(shape_flags):
|
|
data_shape = [x.value for x in data_shape[0:axis]]
|
|
new_shape = (_np.prod(data_shape).astype("int64"), -1)
|
|
else:
|
|
batch_size = 1
|
|
|
|
for el in data_shape[0:axis]:
|
|
batch_size = batch_size * el
|
|
|
|
new_shape = (batch_size, -1)
|
|
|
|
return relax.op.reshape(inputs[0], new_shape)
|
|
|
|
|
|
class LayerNormalization(OnnxOpConverter):
|
|
"""Converts an onnx LayerNormalization node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v17(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
scale = inputs[1]
|
|
bias = inputs[2]
|
|
axis = attr.get("axis", -1)
|
|
epsilon = attr.get("epsilon", 1e-05)
|
|
|
|
gamma_shape = get_const_tuple(scale.ty.shape)
|
|
|
|
if bias is None:
|
|
bias = relax.const(_np.zeros(gamma_shape), dtype=scale.ty.dtype)
|
|
else:
|
|
beta_shape = get_const_tuple(bias.ty.shape)
|
|
if gamma_shape != beta_shape:
|
|
raise ValueError("gamma and beta shapes do not match")
|
|
|
|
axis = list(axis) if isinstance(axis, list | tuple) else [axis]
|
|
if len(axis) < len(gamma_shape):
|
|
axis.extend(range(axis[-1] + 1, axis[-1] + 1 + len(gamma_shape) - len(axis)))
|
|
|
|
output = relax.op.nn.layer_norm(data, scale, bias, axis, epsilon)
|
|
# Onnx layernorm has 3 outputs but only the first is used.
|
|
# We construct two empty constants for this.
|
|
placeholder = relax.const(0, dtype="float32")
|
|
return relax.Tuple([output, placeholder, placeholder])
|
|
|
|
|
|
class RMSNormalization(OnnxOpConverter):
|
|
"""Converts an onnx RMSNormalization node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v23(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
scale = inputs[1]
|
|
axis = attr.get("axis", -1)
|
|
epsilon = attr.get("epsilon", 1e-05)
|
|
stash_type = attr.get("stash_type", 1)
|
|
|
|
# Determine normalization axes: from `axis` to the last dimension
|
|
ndim = _get_known_tensor_rank(data)
|
|
if ndim is None:
|
|
raise ValueError("RMSNormalization requires a statically known input rank.")
|
|
axis = _normalize_constant_axes([axis], ndim, "RMSNormalization")[0]
|
|
axes = list(range(axis, ndim))
|
|
|
|
# If stash_type requires float32 computation and input is not float32, cast
|
|
input_dtype = data.ty.dtype.dtype
|
|
if stash_type == 1 and input_dtype != "float32":
|
|
data_compute = relax.op.astype(data, "float32")
|
|
scale_compute = relax.op.astype(scale, "float32")
|
|
else:
|
|
data_compute = data
|
|
scale_compute = scale
|
|
|
|
output = relax.op.nn.rms_norm(data_compute, scale_compute, axes, epsilon)
|
|
|
|
# Cast back to original dtype if needed
|
|
if stash_type == 1 and input_dtype != "float32":
|
|
output = relax.op.astype(output, input_dtype)
|
|
|
|
return output
|
|
|
|
|
|
class GroupNormalization(OnnxOpConverter):
|
|
"""Converts an onnx GroupNormalization node into an equivalent Relax expression"""
|
|
|
|
@classmethod
|
|
def _impl_v18(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
scale = inputs[1]
|
|
bias = inputs[2]
|
|
num_groups = attr["num_groups"]
|
|
epsilon = attr.get("epsilon", 1e-05)
|
|
|
|
ndim = _get_known_tensor_rank(data)
|
|
if ndim is None:
|
|
raise ValueError("GroupNormalization requires a statically known input rank.")
|
|
|
|
ty = data.ty
|
|
if not isinstance(ty, relax.TensorType) or len(ty.shape) < 2:
|
|
raise ValueError(
|
|
"GroupNormalization-18 requires a statically typed input with rank >= 2."
|
|
)
|
|
|
|
input_dtype = ty.dtype
|
|
if input_dtype != "float32":
|
|
raise ValueError("GroupNormalization-18 currently only supports float32 inputs.")
|
|
|
|
if num_groups <= 0:
|
|
raise ValueError(
|
|
f"GroupNormalization requires num_groups to be positive, got {num_groups}."
|
|
)
|
|
|
|
channel_dim = ty.shape[1]
|
|
if not isinstance(channel_dim, tirx.IntImm):
|
|
raise ValueError(
|
|
"GroupNormalization-18 requires a statically known channel count "
|
|
"to expand per-group scale/bias to per-channel."
|
|
)
|
|
|
|
channels = int(channel_dim)
|
|
if channels % num_groups != 0:
|
|
raise ValueError(
|
|
f"GroupNormalization requires num_groups to divide channel count, "
|
|
f"but got C={channels} and num_groups={num_groups}."
|
|
)
|
|
|
|
channels_per_group = channels // num_groups
|
|
|
|
scale = relax.op.reshape(scale, [num_groups, 1])
|
|
scale = relax.op.broadcast_to(scale, [num_groups, channels_per_group])
|
|
scale = relax.op.reshape(scale, [channels])
|
|
|
|
bias = relax.op.reshape(bias, [num_groups, 1])
|
|
bias = relax.op.broadcast_to(bias, [num_groups, channels_per_group])
|
|
bias = relax.op.reshape(bias, [channels])
|
|
|
|
axes = list(range(2, ndim))
|
|
return relax.op.nn.group_norm(
|
|
data, scale, bias, num_groups, channel_axis=1, axes=axes, epsilon=epsilon
|
|
)
|
|
|
|
@classmethod
|
|
def _impl_v21(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
scale = inputs[1]
|
|
bias = inputs[2]
|
|
num_groups = attr["num_groups"]
|
|
epsilon = attr.get("epsilon", 1e-05)
|
|
stash_type = attr.get("stash_type", 1)
|
|
|
|
if stash_type != 1:
|
|
raise ValueError(
|
|
f"GroupNormalization currently only supports stash_type=1 (FLOAT), "
|
|
f"but got stash_type={stash_type}."
|
|
)
|
|
|
|
ndim = _get_known_tensor_rank(data)
|
|
if ndim is None:
|
|
raise ValueError("GroupNormalization requires a statically known input rank.")
|
|
|
|
ty = data.ty
|
|
if not isinstance(ty, relax.TensorType) or len(ty.shape) < 2:
|
|
raise ValueError("GroupNormalization requires a statically typed input with rank >= 2.")
|
|
|
|
if num_groups <= 0:
|
|
raise ValueError(
|
|
f"GroupNormalization requires num_groups to be positive, got {num_groups}."
|
|
)
|
|
|
|
channel_dim = ty.shape[1]
|
|
if isinstance(channel_dim, tirx.IntImm):
|
|
channels = int(channel_dim)
|
|
if channels % num_groups != 0:
|
|
raise ValueError(
|
|
f"GroupNormalization requires num_groups to divide channel count, "
|
|
f"but got C={channels} and num_groups={num_groups}."
|
|
)
|
|
|
|
axes = list(range(2, ndim))
|
|
input_dtype = ty.dtype
|
|
|
|
orig_scale = scale
|
|
orig_bias = bias
|
|
|
|
if input_dtype != "float32":
|
|
data = relax.op.astype(data, "float32")
|
|
scale = relax.op.astype(scale, "float32")
|
|
bias = relax.op.astype(bias, "float32")
|
|
|
|
norm_scale = relax.op.ones_like(scale)
|
|
norm_bias = relax.op.zeros_like(bias)
|
|
|
|
output = relax.op.nn.group_norm(
|
|
data,
|
|
norm_scale,
|
|
norm_bias,
|
|
num_groups,
|
|
channel_axis=1,
|
|
axes=axes,
|
|
epsilon=epsilon,
|
|
center=False,
|
|
scale=False,
|
|
)
|
|
|
|
if input_dtype != "float32":
|
|
output = relax.op.astype(output, input_dtype)
|
|
|
|
affine_shape = [channel_dim] + [1] * (ndim - 2)
|
|
orig_scale = relax.op.reshape(orig_scale, affine_shape)
|
|
orig_bias = relax.op.reshape(orig_bias, affine_shape)
|
|
output = relax.op.multiply(output, orig_scale)
|
|
output = relax.op.add(output, orig_bias)
|
|
return output
|
|
|
|
|
|
class ReduceMax(OnnxOpConverter):
|
|
"""Converts an onnx ReduceMax node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v11(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
axes = attr.get("axes", None)
|
|
keepdims = attr.get("keepdims", 1)
|
|
return relax.op.max(data, axes, keepdims)
|
|
|
|
@classmethod
|
|
def _impl_v18(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
keepdims = attr.get("keepdims", 1)
|
|
noop_with_empty_axes = attr.get("noop_with_empty_axes", 0)
|
|
|
|
# Optional axes input
|
|
axes = None
|
|
if len(inputs) > 1 and inputs[1] is not None:
|
|
axes_const = get_constant(inputs[1], params)
|
|
assert isinstance(axes_const, relax.Constant), "Only constant axes currently supported"
|
|
axes = axes_const.data.numpy().tolist()
|
|
|
|
# If axes is empty and noop_with_empty_axes is False, reduce all dims
|
|
if not axes and not noop_with_empty_axes:
|
|
return relax.op.max(data, None, keepdims)
|
|
# If axes is empty and noop_with_empty_axes is True, return input unchanged
|
|
elif not axes and noop_with_empty_axes:
|
|
return data
|
|
# Otherwise reduce over specified axes
|
|
else:
|
|
return relax.op.max(data, axes, keepdims)
|
|
|
|
|
|
class ReduceMin(OnnxOpConverter):
|
|
"""Converts an onnx ReduceMin node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v11(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
axes = attr.get("axes", None)
|
|
keepdims = attr.get("keepdims", 1)
|
|
return relax.op.min(data, axes, keepdims)
|
|
|
|
@classmethod
|
|
def _impl_v18(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
keepdims = attr.get("keepdims", 1)
|
|
noop_with_empty_axes = attr.get("noop_with_empty_axes", 0)
|
|
|
|
# Optional axes input
|
|
axes = None
|
|
if len(inputs) > 1 and inputs[1] is not None:
|
|
axes_const = get_constant(inputs[1], params)
|
|
assert isinstance(axes_const, relax.Constant), "Only constant axes currently supported"
|
|
axes = axes_const.data.numpy().tolist()
|
|
|
|
# If axes is empty and noop_with_empty_axes is False, reduce all dims
|
|
if not axes and not noop_with_empty_axes:
|
|
return relax.op.min(data, None, keepdims)
|
|
# If axes is empty and noop_with_empty_axes is True, return input unchanged
|
|
elif not axes and noop_with_empty_axes:
|
|
return data
|
|
# Otherwise reduce over specified axes
|
|
else:
|
|
return relax.op.min(data, axes, keepdims)
|
|
|
|
|
|
class ReduceSum(OnnxOpConverter):
|
|
"""Converts an onnx ReduceSum node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v11(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
axes = attr.get("axes", None)
|
|
keepdims = attr.get("keepdims", 1)
|
|
return relax.op.sum(data, axes, keepdims)
|
|
|
|
@classmethod
|
|
def _impl_v13(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
keepdims = attr.get("keepdims", 1)
|
|
noop_with_empty_axes = attr.get("noop_with_empty_axes", 0)
|
|
|
|
# Optional axes input
|
|
axes = None
|
|
if len(inputs) > 1 and inputs[1] is not None:
|
|
axes_const = get_constant(inputs[1], params)
|
|
assert isinstance(axes_const, relax.Constant), "Only constant axes currently supported"
|
|
axes = axes_const.data.numpy().tolist()
|
|
|
|
# If axes is empty and noop_with_empty_axes is 0, reduce all dimensions
|
|
if not axes and not noop_with_empty_axes:
|
|
return relax.op.sum(data, None, keepdims)
|
|
# If axes is empty and noop_with_empty_axes is 1, return the input data unchanged.
|
|
elif not axes and noop_with_empty_axes:
|
|
return data
|
|
# If axes is provided, reduce over the specified axes
|
|
else:
|
|
return relax.op.sum(data, axes, keepdims)
|
|
|
|
|
|
class ReduceMean(OnnxOpConverter):
|
|
"""Converts an onnx ReduceMean node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v13(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
axes = attr.get("axes", None)
|
|
keepdims = attr.get("keepdims", 1)
|
|
return relax.op.mean(data, axes, keepdims)
|
|
|
|
@classmethod
|
|
def _impl_v18(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
keepdims = attr.get("keepdims", 1)
|
|
noop_with_empty_axes = attr.get("noop_with_empty_axes", 0)
|
|
|
|
# Optional axes input
|
|
axes = None
|
|
if len(inputs) > 1 and inputs[1] is not None:
|
|
axes_const = get_constant(inputs[1], params)
|
|
assert isinstance(axes_const, relax.Constant), "Only constant axes currently supported"
|
|
axes = axes_const.data.numpy().tolist()
|
|
|
|
# If axes is empty and noop_with_empty_axes is 0, reduce all dimensions
|
|
if not axes and not noop_with_empty_axes:
|
|
return relax.op.mean(data, None, keepdims)
|
|
# If axes is empty and noop_with_empty_axes is 1, return the input data unchanged.
|
|
elif not axes and noop_with_empty_axes:
|
|
return data
|
|
# If axes is provided, reduce over the specified axes
|
|
else:
|
|
return relax.op.mean(data, axes, keepdims)
|
|
|
|
|
|
class ReduceProd(OnnxOpConverter):
|
|
"""Converts an onnx ReduceProd node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v13(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
axes = attr.get("axes", None)
|
|
keepdims = attr.get("keepdims", 1)
|
|
return relax.op.prod(data, axes, keepdims)
|
|
|
|
@classmethod
|
|
def _impl_v18(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
keepdims = attr.get("keepdims", 1)
|
|
noop_with_empty_axes = attr.get("noop_with_empty_axes", 0)
|
|
|
|
# Optional axes input
|
|
axes = None
|
|
if len(inputs) > 1 and inputs[1] is not None:
|
|
axes_const = get_constant(inputs[1], params)
|
|
assert isinstance(axes_const, relax.Constant), "Only constant axes currently supported"
|
|
axes = axes_const.data.numpy().tolist()
|
|
|
|
# If axes is empty and noop_with_empty_axes is 0, reduce all dimensions
|
|
if not axes and not noop_with_empty_axes:
|
|
return relax.op.prod(data, None, keepdims)
|
|
# If axes is empty and noop_with_empty_axes is 1, return the input data unchanged.
|
|
elif not axes and noop_with_empty_axes:
|
|
return data
|
|
# If axes is provided, reduce over the specified axes
|
|
else:
|
|
return relax.op.prod(data, axes, keepdims)
|
|
|
|
|
|
class ReduceLogSumExp(OnnxOpConverter):
|
|
"""Converts an onnx ReduceLogSumExp node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v13(cls, bb, inputs, attr, params):
|
|
x = inputs[0]
|
|
axes = attr.get("axes", None)
|
|
keepdims = attr.get("keepdims", 1)
|
|
max_x = relax.op.max(x, axes, True)
|
|
exp_x = relax.op.exp(relax.op.subtract(x, max_x))
|
|
sum_x = relax.op.sum(exp_x, axes, True)
|
|
out_x = relax.op.add(relax.op.log(sum_x), max_x)
|
|
if not keepdims:
|
|
out_x = relax.op.squeeze(out_x, axes)
|
|
return out_x
|
|
|
|
@classmethod
|
|
def _impl_v18(cls, bb, inputs, attr, params):
|
|
x = inputs[0]
|
|
keepdims = attr.get("keepdims", 1)
|
|
noop_with_empty_axes = attr.get("noop_with_empty_axes", 0)
|
|
|
|
# Optional axes input (second input)
|
|
axes = None
|
|
if len(inputs) > 1 and inputs[1] is not None:
|
|
axes_const = get_constant(inputs[1], params)
|
|
assert isinstance(axes_const, relax.Constant), "Only constant axes currently supported"
|
|
axes = axes_const.data.numpy().tolist()
|
|
|
|
# Calculate LogSumExp
|
|
log_sum_exp = lambda axes: (
|
|
max_x := relax.op.max(x, axes, True),
|
|
exp_x := relax.op.exp(relax.op.subtract(x, max_x)),
|
|
sum_x := relax.op.sum(exp_x, axes, True),
|
|
out_x := relax.op.add(relax.op.log(sum_x), max_x),
|
|
relax.op.squeeze(out_x, axes) if not keepdims else out_x,
|
|
)[-1]
|
|
|
|
# If axes is empty and noop_with_empty_axes is 0, reduce all dimensions
|
|
if not axes and not noop_with_empty_axes:
|
|
return log_sum_exp(None)
|
|
# If axes is empty and noop_with_empty_axes is 1, return the input data unchanged.
|
|
elif not axes and noop_with_empty_axes:
|
|
return x
|
|
# If axes is provided, reduce over the specified axes
|
|
else:
|
|
return log_sum_exp(axes)
|
|
|
|
|
|
class ReduceLogSum(OnnxOpConverter):
|
|
"""Converts an onnx ReduceLogSum node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v13(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
axes = attr.get("axes", None)
|
|
keepdims = attr.get("keepdims", 1)
|
|
return relax.op.log(relax.op.sum(data, axes, keepdims))
|
|
|
|
@classmethod
|
|
def _impl_v18(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
keepdims = attr.get("keepdims", 1)
|
|
noop_with_empty_axes = attr.get("noop_with_empty_axes", 0)
|
|
|
|
# Optional axes input
|
|
axes = None
|
|
if len(inputs) > 1 and inputs[1] is not None:
|
|
axes_const = get_constant(inputs[1], params)
|
|
assert isinstance(axes_const, relax.Constant), "Only constant axes currently supported"
|
|
axes = axes_const.data.numpy().tolist()
|
|
|
|
# If axes is empty and noop_with_empty_axes is 0, reduce all dimensions
|
|
if not axes and not noop_with_empty_axes:
|
|
return relax.op.log(relax.op.sum(data, None, keepdims))
|
|
# If axes is empty and noop_with_empty_axes is 1, return the input data unchanged.
|
|
elif not axes and noop_with_empty_axes:
|
|
return data
|
|
# If axes is provided, reduce over the specified axes
|
|
else:
|
|
return relax.op.log(relax.op.sum(data, axes, keepdims))
|
|
|
|
|
|
class ReduceSumSquare(OnnxOpConverter):
|
|
"""Converts an onnx ReduceSumSquare node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v13(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
axes = attr.get("axes", None)
|
|
keepdims = attr.get("keepdims", 1)
|
|
return relax.op.sum(relax.op.multiply(data, data), axes, keepdims)
|
|
|
|
@classmethod
|
|
def _impl_v18(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
keepdims = attr.get("keepdims", 1)
|
|
noop_with_empty_axes = attr.get("noop_with_empty_axes", 0)
|
|
|
|
# Optional axes input
|
|
axes = None
|
|
if len(inputs) > 1 and inputs[1] is not None:
|
|
axes_const = get_constant(inputs[1], params)
|
|
assert isinstance(axes_const, relax.Constant), "Only constant axes currently supported"
|
|
axes = axes_const.data.numpy().tolist()
|
|
|
|
# If axes is empty and noop_with_empty_axes is 0, reduce all dimensions
|
|
if not axes and not noop_with_empty_axes:
|
|
return relax.op.sum(relax.op.multiply(data, data), None, keepdims)
|
|
# If axes is empty and noop_with_empty_axes is 1, return the input data unchanged.
|
|
elif not axes and noop_with_empty_axes:
|
|
return data
|
|
# If axes is provided, reduce over the specified axes
|
|
else:
|
|
return relax.op.sum(relax.op.multiply(data, data), axes, keepdims)
|
|
|
|
|
|
class ReduceL1(OnnxOpConverter):
|
|
"""Converts an onnx ReduceL1 node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v13(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
axes = attr.get("axes", None)
|
|
keepdims = attr.get("keepdims", 1)
|
|
return relax.op.sum(relax.op.abs(data), axes, keepdims)
|
|
|
|
@classmethod
|
|
def _impl_v18(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
keepdims = attr.get("keepdims", 1)
|
|
noop_with_empty_axes = attr.get("noop_with_empty_axes", 0)
|
|
|
|
# Optional axes input
|
|
axes = None
|
|
if len(inputs) > 1 and inputs[1] is not None:
|
|
axes_const = get_constant(inputs[1], params)
|
|
assert isinstance(axes_const, relax.Constant), "Only constant axes currently supported"
|
|
axes = axes_const.data.numpy().tolist()
|
|
|
|
# If axes is empty and noop_with_empty_axes is 0, reduce all dimensions
|
|
if not axes and not noop_with_empty_axes:
|
|
return relax.op.sum(relax.op.abs(data), None, keepdims)
|
|
# If axes is empty and noop_with_empty_axes is 1, return the input data unchanged.
|
|
elif not axes and noop_with_empty_axes:
|
|
return data
|
|
# If axes is provided, reduce over the specified axes
|
|
else:
|
|
return relax.op.sum(relax.op.abs(data), axes, keepdims)
|
|
|
|
|
|
class ReduceL2(OnnxOpConverter):
|
|
"""Converts an onnx ReduceL2 node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v13(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
axes = attr.get("axes", None)
|
|
keepdims = attr.get("keepdims", 1)
|
|
return relax.op.sqrt(relax.op.sum(relax.op.multiply(data, data), axes, keepdims))
|
|
|
|
@classmethod
|
|
def _impl_v18(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
keepdims = attr.get("keepdims", 1)
|
|
noop_with_empty_axes = attr.get("noop_with_empty_axes", 0)
|
|
|
|
# Optional axes input
|
|
axes = None
|
|
if len(inputs) > 1 and inputs[1] is not None:
|
|
axes_const = get_constant(inputs[1], params)
|
|
assert isinstance(axes_const, relax.Constant), "Only constant axes currently supported"
|
|
axes = axes_const.data.numpy().tolist()
|
|
|
|
# If axes is empty and noop_with_empty_axes is 0, reduce all dimensions
|
|
if not axes and not noop_with_empty_axes:
|
|
return relax.op.sqrt(relax.op.sum(relax.op.multiply(data, data), None, keepdims))
|
|
# If axes is empty and noop_with_empty_axes is 1, return the input data unchanged.
|
|
elif not axes and noop_with_empty_axes:
|
|
return data
|
|
# If axes is provided, reduce over the specified axes
|
|
else:
|
|
return relax.op.sqrt(relax.op.sum(relax.op.multiply(data, data), axes, keepdims))
|
|
|
|
|
|
def _argreduce_select_last_index(bb, data, axis, keepdims, op):
|
|
"""Helper for ArgMax/ArgMin with select_last_index=1.
|
|
|
|
Reverses the tensor along the reduction axis, runs the reduction op,
|
|
then remaps the index back: last_idx = (axis_size - 1) - flipped_idx.
|
|
Handles both static and dynamic axis sizes.
|
|
"""
|
|
data_flipped = relax.op.flip(data, axis=axis)
|
|
flipped_idx = bb.normalize(op(data_flipped, axis, keepdims))
|
|
axis_size = data.ty.shape[axis]
|
|
if isinstance(axis_size, tirx.IntImm):
|
|
offset = relax.const(int(axis_size) - 1, "int64")
|
|
else:
|
|
# dynamic: get axis size at runtime and subtract 1
|
|
shape_tensor = bb.normalize(relax.op.shape_to_tensor(bb.normalize(relax.op.shape_of(data))))
|
|
offset = bb.normalize(
|
|
relax.op.subtract(
|
|
bb.normalize(relax.op.take(shape_tensor, relax.const(axis, "int64"), axis=0)),
|
|
relax.const(1, "int64"),
|
|
)
|
|
)
|
|
return relax.op.subtract(offset, flipped_idx)
|
|
|
|
|
|
class ArgMax(OnnxOpConverter):
|
|
"""Converts an onnx ArgMax node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _check_attrs(cls, data, attr, shift_axis=True):
|
|
dims_num = len(data.ty.shape)
|
|
axis = attr.get("axis", 0)
|
|
if shift_axis and axis < 0:
|
|
axis += dims_num
|
|
assert 0 <= axis < dims_num, "Axis is out of bounds"
|
|
keepdims = attr.get("keepdims", True)
|
|
return axis, keepdims
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
axis, keepdims = cls._check_attrs(data, attr, False)
|
|
return relax.op.argmax(data, axis, keepdims)
|
|
|
|
@classmethod
|
|
def _impl_v11(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
axis, keepdims = cls._check_attrs(data, attr)
|
|
return relax.op.argmax(data, axis, keepdims)
|
|
|
|
@classmethod
|
|
def _impl_v12(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
axis, keepdims = cls._check_attrs(data, attr)
|
|
select_last_index = attr.get("select_last_index", False)
|
|
if select_last_index:
|
|
return _argreduce_select_last_index(bb, data, axis, keepdims, relax.op.argmax)
|
|
return relax.op.argmax(data, axis, keepdims)
|
|
|
|
|
|
class ArgMin(OnnxOpConverter):
|
|
"""Converts an onnx ArgMin node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _check_attrs(cls, data, attr, shift_axis=True):
|
|
dims_num = len(data.ty.shape)
|
|
axis = attr.get("axis", 0)
|
|
if shift_axis and axis < 0:
|
|
axis += dims_num
|
|
assert 0 <= axis < dims_num, "Axis is out of bounds"
|
|
keepdims = attr.get("keepdims", True)
|
|
return axis, keepdims
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
axis, keepdims = cls._check_attrs(data, attr, False)
|
|
return relax.op.argmin(data, axis, keepdims)
|
|
|
|
@classmethod
|
|
def _impl_v11(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
axis, keepdims = cls._check_attrs(data, attr)
|
|
return relax.op.argmin(data, axis, keepdims)
|
|
|
|
@classmethod
|
|
def _impl_v12(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
axis, keepdims = cls._check_attrs(data, attr)
|
|
select_last_index = attr.get("select_last_index", False)
|
|
if select_last_index:
|
|
return _argreduce_select_last_index(bb, data, axis, keepdims, relax.op.argmin)
|
|
return relax.op.argmin(data, axis, keepdims)
|
|
|
|
|
|
class TopK(OnnxOpConverter):
|
|
"""Converts an onnx TopK node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v11(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
k = get_constant(inputs[1], params)
|
|
if not isinstance(k, relax.Constant):
|
|
raise ValueError("TopK k must be a constant")
|
|
k = int(k.data.numpy().item())
|
|
axis = attr.get("axis", -1)
|
|
largest = attr.get("largest", 1)
|
|
sorted = attr.get("sorted", 1)
|
|
if sorted != 1:
|
|
raise ValueError("TopK sorted must be 1 for Relax frontend")
|
|
|
|
return relax.op.topk(data, k, axis, ret_type="both", largest=largest, dtype="int64")
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
k = attr.get("k", 1)
|
|
axis = attr.get("axis", -1)
|
|
return relax.op.topk(data, k, axis, ret_type="both", dtype="int64")
|
|
|
|
|
|
class SkipLayerNormalization(OnnxOpConverter):
|
|
"""Converts a microsoft contrib SkipLayerNormalization node into a Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
skip = inputs[1]
|
|
gamma = inputs[2]
|
|
beta = inputs[3]
|
|
bias = inputs[4]
|
|
|
|
assert beta is not None and bias is not None, (
|
|
"SkipLayerNormalization import currently only supports required beta and bias"
|
|
)
|
|
|
|
epsilon = attr.get("epsilon", 1e-12)
|
|
|
|
data = relax.op.add(data, skip)
|
|
if bias is not None:
|
|
data = relax.op.add(data, bias)
|
|
|
|
output = relax.op.nn.layer_norm(data, gamma, beta, axes=-1, epsilon=epsilon)
|
|
|
|
# Expects three outputs though only the first is used. Construct a placeholder for others.
|
|
placeholder = relax.const(0, dtype="float32")
|
|
return relax.Tuple([output, placeholder, placeholder])
|
|
|
|
|
|
class EmbedLayerNormalization(OnnxOpConverter):
|
|
"""Converts a microsoft contrib EmbedLayerNormalization node into a Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
input_ids = inputs[0]
|
|
segment_ids = inputs[1]
|
|
word_emb = inputs[2]
|
|
pos_emb = inputs[3]
|
|
segment_emb = inputs[4]
|
|
gamma = inputs[5]
|
|
beta = inputs[6]
|
|
mask = inputs[7]
|
|
pos_ids = inputs[8]
|
|
|
|
epsilon = attr.get("epsilon", 1e-12)
|
|
|
|
(batch_size, seq_len) = [dim.value for dim in input_ids.ty.shape]
|
|
|
|
if segment_ids:
|
|
assert segment_emb
|
|
|
|
if pos_ids is None:
|
|
pos_ids = relax.const([list(range(seq_len))] * batch_size, dtype="int64")
|
|
word_vec = relax.op.take(word_emb, input_ids, axis=0)
|
|
if segment_ids:
|
|
segment_vec = relax.op.take(segment_emb, segment_ids, axis=0)
|
|
pos_vec = relax.op.take(pos_emb, pos_ids, axis=0)
|
|
|
|
vec_sum = relax.op.add(word_vec, pos_vec)
|
|
if segment_ids:
|
|
vec_sum = relax.op.add(vec_sum, segment_vec)
|
|
|
|
ln = relax.op.nn.layer_norm(vec_sum, gamma, beta, axes=-1, epsilon=epsilon)
|
|
|
|
mask_index = relax.const(_np.zeros((batch_size,), dtype="int64"))
|
|
if mask:
|
|
# Caculate number of words per sentence.
|
|
mask_index = relax.op.sum(mask, axis=1)
|
|
|
|
return relax.Tuple([ln, mask_index])
|
|
|
|
|
|
class OneHot(OnnxOpConverter):
|
|
"""Converts an onnx OneHot node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v11(cls, bb, inputs, attr, params):
|
|
indices = inputs[0]
|
|
depth = get_constant(inputs[1], params)
|
|
values = get_constant(inputs[2], params)
|
|
axis = attr.get("axis", -1)
|
|
assert isinstance(depth, relax.Constant), "Only constant depth currently supported."
|
|
depth = depth.data.numpy().tolist()
|
|
assert isinstance(values, relax.Constant), "Only constant values currently supported."
|
|
values = values.data.numpy().tolist()
|
|
off_value, on_value = values
|
|
off_value, on_value = (
|
|
relax.prim_value(off_value),
|
|
relax.prim_value(on_value),
|
|
)
|
|
return relax.op.one_hot(indices, on_value, off_value, depth, axis)
|
|
|
|
|
|
class Unique(OnnxOpConverter):
|
|
"""Converts an onnx Unique node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v11(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
axis = attr.get("axis", None)
|
|
sorted_flag = bool(attr.get("sorted", 1))
|
|
num_outputs = attr["tvm_custom"]["num_outputs"]
|
|
|
|
return_index = num_outputs > 1
|
|
return_inverse = num_outputs > 2
|
|
return_counts = num_outputs > 3
|
|
|
|
unique = relax.op.unique(
|
|
data,
|
|
sorted=sorted_flag,
|
|
return_index=return_index,
|
|
return_inverse=return_inverse,
|
|
return_counts=return_counts,
|
|
axis=axis,
|
|
)
|
|
|
|
unique_numbers = tirx.Var("unique_numbers", "int64")
|
|
input_shape = data.ty.shape
|
|
dtype = data.ty.dtype
|
|
|
|
if axis is None:
|
|
output_shape = (unique_numbers,)
|
|
else:
|
|
axis = axis if axis >= 0 else len(input_shape) + axis
|
|
if axis < 0 or axis >= len(input_shape):
|
|
raise ValueError(f"Axis {axis} is out of bounds")
|
|
output_shape = [
|
|
input_shape[i] if i != axis else unique_numbers for i in range(len(input_shape))
|
|
]
|
|
|
|
if num_outputs == 1:
|
|
return bb.match_cast(unique, relax.TensorType(output_shape, dtype))
|
|
|
|
outputs = [bb.match_cast(unique[0], relax.TensorType(output_shape, dtype))]
|
|
tuple_idx = 1 # Track which index in the tuple we're at
|
|
|
|
if return_index:
|
|
index_shape = (unique_numbers,)
|
|
index_ty = relax.TensorType(index_shape, "int64")
|
|
outputs.append(bb.match_cast(unique[tuple_idx], index_ty))
|
|
tuple_idx += 1
|
|
|
|
if return_inverse:
|
|
# ONNX spec: inverse_indices is always 1D
|
|
# When axis is None: shape is [X.size]
|
|
# When axis is specified: shape is [X.shape[axis]]
|
|
inverse_shape = (tirx.Var("inverse_numbers", "int64"),)
|
|
inverse_ty = relax.TensorType(inverse_shape, "int64")
|
|
outputs.append(bb.match_cast(unique[tuple_idx], inverse_ty))
|
|
tuple_idx += 1
|
|
|
|
if return_counts:
|
|
count_shape = (unique_numbers,)
|
|
count_ty = relax.TensorType(count_shape, "int64")
|
|
outputs.append(bb.match_cast(unique[tuple_idx], count_ty))
|
|
|
|
return relax.Tuple(outputs)
|
|
|
|
|
|
class NonZero(OnnxOpConverter):
|
|
"""Converts an onnx NonZero node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v9(cls, bb, inputs, attr, params):
|
|
ndim = inputs[0].ty.ndim
|
|
ndim = 1 if ndim == 0 else ndim
|
|
nonzero_numbers = tirx.Var("nonzero_numbers", "int64")
|
|
return bb.match_cast(
|
|
relax.op.nonzero(inputs[0]), relax.TensorType((ndim, nonzero_numbers), "int64")
|
|
)
|
|
|
|
|
|
class Upsample(OnnxOpConverter):
|
|
"""Operator converter for Upsample (nearest mode)."""
|
|
|
|
@classmethod
|
|
def _impl_v9(cls, bb, inputs, attr, params):
|
|
scales = attr.get("scales")
|
|
assert len(scales) == 4
|
|
assert scales[0] == scales[1] == 1
|
|
|
|
inp_shape = [int(x) for x in inputs[0].ty.shape]
|
|
assert len(inp_shape) == 4
|
|
out_shape2d = [int(dim * scale) for dim, scale in zip(inp_shape[2:], scales[2:])]
|
|
|
|
mode = attr.get("mode", b"nearest").decode("ascii")
|
|
if mode == "nearest":
|
|
mode = "nearest_neighbor"
|
|
msg = f'Value {mode} in attribute "mode" of operator Upsample is not valid.'
|
|
assert mode in ("linear", "nearest_neighbor", "cubic"), msg
|
|
|
|
return relax.op.image.resize2d(
|
|
data=inputs[0],
|
|
roi=None,
|
|
size=relax.ShapeExpr(out_shape2d), # (H, W)
|
|
layout="NCHW",
|
|
method=mode,
|
|
coordinate_transformation_mode="asymmetric", # Align with Upsample
|
|
)
|
|
|
|
|
|
class HardSigmoid(OnnxOpConverter):
|
|
"""Converts an onnx HardSigmoid node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
x = inputs[0]
|
|
dtype = x.ty.dtype
|
|
alpha = float(attr.get("alpha", 0.2))
|
|
alpha = relax.const(alpha, dtype=dtype)
|
|
beta = float(attr.get("beta", 0.5))
|
|
beta = relax.const(beta, dtype=dtype)
|
|
return relax.op.clip(relax.op.add(relax.op.multiply(alpha, x), beta), 0, 1)
|
|
|
|
|
|
class HardSwish(OnnxOpConverter):
|
|
"""Converts an onnx HardSwish node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v14(cls, bb, inputs, attr, params):
|
|
x = inputs[0]
|
|
dtype = x.ty.dtype
|
|
return relax.op.multiply(
|
|
x,
|
|
relax.op.divide(
|
|
relax.op.clip(relax.op.add(x, relax.const(3, dtype)), 0, 6),
|
|
relax.expr.const(6, dtype),
|
|
),
|
|
)
|
|
|
|
|
|
class Sign(OnnxOpConverter):
|
|
"""Converts an onnx Sign node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v9(cls, bb, inputs, attr, params):
|
|
return relax.op.sign(inputs[0])
|
|
|
|
|
|
class Not(OnnxOpConverter):
|
|
"""Converts an onnx Not node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
return relax.op.logical_not(inputs[0])
|
|
|
|
|
|
class DepthToSpace(OnnxOpConverter):
|
|
"""Converts an onnx DepthToSpace node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v11(cls, bb, inputs, attr, params):
|
|
block_size = int(attr["blocksize"])
|
|
mode = attr.get("mode", b"DCR").decode("utf-8")
|
|
b, c, h, w = inputs[0].ty.shape
|
|
if mode == "DCR":
|
|
x = relax.op.reshape(inputs[0], (b, block_size, block_size, c // (block_size**2), h, w))
|
|
x = relax.op.permute_dims(x, [0, 3, 4, 1, 5, 2])
|
|
return relax.op.reshape(x, (b, c // (block_size**2), h * block_size, w * block_size))
|
|
elif mode == "CRD":
|
|
x = relax.op.reshape(inputs[0], (b, c // (block_size**2), block_size, block_size, h, w))
|
|
x = relax.op.permute_dims(x, [0, 1, 4, 2, 5, 3])
|
|
return relax.op.reshape(x, (b, c // (block_size**2), h * block_size, w * block_size))
|
|
else:
|
|
raise ValueError(f"Unsupported mode: {mode}, expected DCR or CRD")
|
|
|
|
|
|
class SpaceToDepth(OnnxOpConverter):
|
|
"""Converts an onnx SpaceToDepth node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
block_size = int(attr["blocksize"])
|
|
b, c, h, w = inputs[0].ty.shape
|
|
x = relax.op.reshape(
|
|
inputs[0], (b, c, h // block_size, block_size, w // block_size, block_size)
|
|
)
|
|
x = relax.op.permute_dims(x, [0, 3, 5, 1, 2, 4])
|
|
return relax.op.reshape(
|
|
x, (b, c * block_size * block_size, h // block_size, w // block_size)
|
|
)
|
|
|
|
|
|
class Optional_(OnnxOpConverter):
|
|
"""Converts an ONNX Optional node into an erased or empty Optional representation."""
|
|
|
|
@classmethod
|
|
def _impl_v15(cls, bb, inputs, attr, params):
|
|
if len(inputs) > 1:
|
|
raise ValueError(f"Optional accepts at most one input, but got {len(inputs)}")
|
|
if len(inputs) == 0 or inputs[0] is None:
|
|
if "type" not in attr:
|
|
raise ValueError("Optional without an input must specify the type attribute.")
|
|
return _EmptyOptional(attr["type"])
|
|
return inputs[0]
|
|
|
|
_impl_v18 = _impl_v15
|
|
|
|
|
|
class OptionalHasElement(OnnxOpConverter):
|
|
"""Converts an ONNX OptionalHasElement node into a boolean constant."""
|
|
|
|
@classmethod
|
|
def _impl_v15(cls, bb, inputs, attr, params):
|
|
if len(inputs) != 1:
|
|
raise ValueError(f"OptionalHasElement expects one input, but got {len(inputs)}")
|
|
if inputs[0] is None or _is_empty_optional(inputs[0]):
|
|
return relax.const(False, dtype="bool")
|
|
return relax.const(True, dtype="bool")
|
|
|
|
_impl_v18 = _impl_v15
|
|
|
|
|
|
class OptionalGetElement(OnnxOpConverter):
|
|
"""Converts an ONNX OptionalGetElement node by unwrapping a non-empty Optional."""
|
|
|
|
@classmethod
|
|
def _impl_v15(cls, bb, inputs, attr, params):
|
|
if len(inputs) != 1:
|
|
raise ValueError(f"OptionalGetElement expects one input, but got {len(inputs)}")
|
|
if inputs[0] is None or _is_empty_optional(inputs[0]):
|
|
raise ValueError("OptionalGetElement cannot access an empty optional.")
|
|
return inputs[0]
|
|
|
|
_impl_v18 = _impl_v15
|
|
|
|
|
|
class SequenceConstruct(OnnxOpConverter):
|
|
"""Operator converter for sequence construction op."""
|
|
|
|
@classmethod
|
|
def _impl_v11(cls, bb, inputs, attr, params):
|
|
# Construct a tuple from input tensors.
|
|
return relax.Tuple(inputs)
|
|
|
|
|
|
class SequenceEmpty(OnnxOpConverter):
|
|
"""Operator converter for sequence empty op."""
|
|
|
|
@classmethod
|
|
def _impl_v11(cls, bb, inputs, attr, params):
|
|
# Construct an empty tuple.
|
|
return relax.Tuple([])
|
|
|
|
|
|
class SequenceErase(OnnxOpConverter):
|
|
"""Operator converter for sequence erase op."""
|
|
|
|
@classmethod
|
|
def _impl_v11(cls, bb, inputs, attr, params):
|
|
# Erase tensor from sequence on specified position
|
|
input_sequence = inputs[0]
|
|
|
|
if len(inputs) == 2:
|
|
position = inputs[1]
|
|
# Non constant position is not supported.
|
|
if isinstance(position, relax.Constant):
|
|
position = int(position.data.numpy())
|
|
else:
|
|
raise NotImplementedError("Position must be a constant.")
|
|
else:
|
|
position = -1
|
|
|
|
seq_len = len(input_sequence)
|
|
if not -seq_len <= position < seq_len:
|
|
raise ValueError(
|
|
f"Position is out of bounds, expected [-{seq_len}, {seq_len}), got {position}"
|
|
)
|
|
|
|
if position < 0:
|
|
position = seq_len + position
|
|
seq_list = list(input_sequence)
|
|
items = [t for i, t in enumerate(seq_list) if i != position]
|
|
return relax.Tuple(items)
|
|
|
|
|
|
class SequenceInsert(OnnxOpConverter):
|
|
"""Operator converter for sequence insert op."""
|
|
|
|
@classmethod
|
|
def _impl_v11(cls, bb, inputs, attr, params):
|
|
# Insert a new tensor into a tuple of tensors.
|
|
input_sequence = inputs[0]
|
|
tensor_to_insert = inputs[1]
|
|
|
|
if len(inputs) == 3:
|
|
position = inputs[2]
|
|
# Non constant position is not supported.
|
|
if isinstance(position, relax.Constant):
|
|
position = int(position.data.numpy())
|
|
else:
|
|
raise NotImplementedError("Position must be a constant.")
|
|
else:
|
|
position = -1
|
|
|
|
seq_len = len(input_sequence)
|
|
if position < 0:
|
|
position = seq_len + position + 1
|
|
# Upper bound is inclusive: position == seq_len appends at the end.
|
|
if not 0 <= position <= seq_len:
|
|
raise ValueError(
|
|
f"SequenceInsert position out of bounds for length {seq_len}, got {position}"
|
|
)
|
|
tensor_list = list(input_sequence)
|
|
tensor_list.insert(position, tensor_to_insert)
|
|
return relax.Tuple(tensor_list)
|
|
|
|
|
|
class SequenceLength(OnnxOpConverter):
|
|
"""Operator converter for sequence length op."""
|
|
|
|
@classmethod
|
|
def _impl_v11(cls, bb, inputs, attr, params):
|
|
# Get length of input sequence
|
|
return relax.const(len(inputs[0]), dtype="int64")
|
|
|
|
|
|
class ConcatFromSequence(OnnxOpConverter):
|
|
"""Operator converter for sequence concatenation op."""
|
|
|
|
@classmethod
|
|
def _impl_v11(cls, bb, inputs, attr, params):
|
|
axis = attr.get("axis", 0)
|
|
new_axis = attr.get("new_axis", 0)
|
|
|
|
if new_axis not in (0, 1):
|
|
raise ValueError(f"ConcatFromSequence only supports new_axis in (0, 1), got {new_axis}")
|
|
|
|
tensors = list(inputs[0])
|
|
if new_axis == 1:
|
|
tensors = [relax.op.expand_dims(t, axis=axis) for t in tensors]
|
|
|
|
return relax.op.concat(tensors, axis=axis)
|
|
|
|
|
|
class SplitToSequence(OnnxOpConverter):
|
|
"""Operator converter for split to sequence op."""
|
|
|
|
@classmethod
|
|
def _impl_v11(cls, bb, inputs, attr, params):
|
|
axis = attr.get("axis", 0)
|
|
keepdims = attr.get("keepdims", 1)
|
|
|
|
input_tensor = inputs[0]
|
|
input_shape = input_tensor.ty.shape
|
|
|
|
if len(inputs) == 1:
|
|
split = _np.array(1)
|
|
else:
|
|
split = inputs[1]
|
|
if not isinstance(split, relax.Constant):
|
|
raise ValueError("Only constant split supported for SplitToSequence")
|
|
split = split.data.numpy()
|
|
|
|
if len(split.shape) == 1 and split.shape[0] > 1:
|
|
split = _np.cumsum(split)
|
|
split = list(split[:-1])
|
|
else:
|
|
chunk_size = int(split)
|
|
dim_size = input_shape[axis]
|
|
|
|
if isinstance(dim_size, int | tirx.IntImm):
|
|
dim_size_int = int(dim_size)
|
|
split = math.ceil(dim_size_int / chunk_size)
|
|
else:
|
|
raise NotImplementedError(
|
|
"SplitToSequence with dynamic dim size and scalar split is not supported."
|
|
)
|
|
|
|
output = relax.op.split(input_tensor, split, axis=axis)
|
|
|
|
# keepdims=0 applies when split is a scalar (whether provided or defaulted to 1)
|
|
# Per ONNX spec: "If input 'split' is specified, this attribute is ignored."
|
|
if not keepdims and len(inputs) == 1:
|
|
output = bb.emit(output)
|
|
n = len(output.ty.fields)
|
|
squeezed = [
|
|
relax.op.squeeze(bb.emit(relax.TupleGetItem(output, i)), axis=[axis])
|
|
for i in range(n)
|
|
]
|
|
return relax.Tuple(squeezed)
|
|
|
|
return output
|
|
|
|
|
|
class SequenceAt(OnnxOpConverter):
|
|
"""Operator converter for sequence at op."""
|
|
|
|
@classmethod
|
|
def _impl_v11(cls, bb, inputs, attr, params):
|
|
input_sequence = inputs[0]
|
|
position = inputs[1]
|
|
assert isinstance(position, relax.Constant), (
|
|
"Only constant position supported for SequenceAt"
|
|
)
|
|
position = int(position.data.numpy())
|
|
return input_sequence[position]
|
|
|
|
|
|
class NonMaxSuppression(OnnxOpConverter):
|
|
"""Converts an onnx NonMaxSuppression node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v10(cls, bb, inputs, attr, params):
|
|
"""
|
|
NonMaxSuppression performs non-maximum suppression (NMS) on all classes.
|
|
|
|
Inputs:
|
|
- boxes: (N, 4) tensor of bounding boxes in format [x1, y1, x2, y2]
|
|
- scores: (N, C) tensor of scores for each box and class
|
|
- max_output_boxes_per_class: maximum number of boxes to keep per class
|
|
- iou_threshold: IoU threshold for NMS
|
|
- score_threshold: score threshold for filtering
|
|
|
|
Outputs:
|
|
- selected_indices: (M, 3) tensor with [batch_idx, class_idx, box_idx]
|
|
"""
|
|
boxes = inputs[0]
|
|
scores = inputs[1]
|
|
max_output_boxes_per_class = inputs[2] if len(inputs) > 2 else None
|
|
iou_threshold = inputs[3] if len(inputs) > 3 else None
|
|
score_threshold = inputs[4] if len(inputs) > 4 else None
|
|
|
|
center_point_box = attr.get("center_point_box", 0)
|
|
|
|
if max_output_boxes_per_class is not None and isinstance(
|
|
max_output_boxes_per_class, relax.Constant
|
|
):
|
|
max_output_boxes_per_class = int(max_output_boxes_per_class.data.numpy().item())
|
|
elif max_output_boxes_per_class is not None and isinstance(
|
|
max_output_boxes_per_class, relax.Var
|
|
):
|
|
var_name = max_output_boxes_per_class.name_hint
|
|
if var_name in params[1]:
|
|
_, param_value = params[1][var_name]
|
|
max_output_boxes_per_class = int(param_value.numpy().item())
|
|
else:
|
|
max_output_boxes_per_class = 0 # Default value
|
|
else:
|
|
max_output_boxes_per_class = 0 # Default value
|
|
|
|
if iou_threshold is not None and isinstance(iou_threshold, relax.Constant):
|
|
iou_threshold = float(iou_threshold.data.numpy().item())
|
|
elif iou_threshold is not None and isinstance(iou_threshold, relax.Var):
|
|
var_name = iou_threshold.name_hint
|
|
if var_name in params[1]:
|
|
_, param_value = params[1][var_name]
|
|
iou_threshold = float(param_value.numpy().item())
|
|
else:
|
|
iou_threshold = 0.5 # Default value
|
|
else:
|
|
iou_threshold = 0.5 # Default value
|
|
|
|
if score_threshold is not None and isinstance(score_threshold, relax.Constant):
|
|
score_threshold = float(score_threshold.data.numpy().item())
|
|
elif score_threshold is not None and isinstance(score_threshold, relax.Var):
|
|
var_name = score_threshold.name_hint
|
|
if var_name in params[1]:
|
|
_, param_value = params[1][var_name]
|
|
score_threshold = float(param_value.numpy().item())
|
|
else:
|
|
score_threshold = 0.0 # Default value
|
|
else:
|
|
score_threshold = 0.0 # Default value
|
|
|
|
if center_point_box != 0:
|
|
split_result = relax.op.split(boxes, 4, axis=2)
|
|
xc = split_result[0]
|
|
yc = split_result[1]
|
|
w = split_result[2]
|
|
h = split_result[3]
|
|
half_w = w / relax.const(2.0, boxes.ty.dtype)
|
|
half_h = h / relax.const(2.0, boxes.ty.dtype)
|
|
x1 = xc - half_w
|
|
x2 = xc + half_w
|
|
y1 = yc - half_h
|
|
y2 = yc + half_h
|
|
boxes = relax.op.concat([y1, x1, y2, x2], axis=2)
|
|
|
|
nms_out = bb.normalize(
|
|
relax.op.vision.all_class_non_max_suppression(
|
|
boxes,
|
|
scores,
|
|
relax.const(max_output_boxes_per_class, dtype="int64"),
|
|
relax.const(iou_threshold, dtype="float32"),
|
|
relax.const(score_threshold, dtype="float32"),
|
|
output_format="onnx",
|
|
)
|
|
)
|
|
|
|
selected_indices = bb.emit(relax.TupleGetItem(nms_out, 0))
|
|
|
|
return selected_indices
|
|
|
|
|
|
class AllClassNMS(OnnxOpConverter):
|
|
"""Converts an onnx AllClassNMS node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v1(cls, bb, inputs, attr, params):
|
|
"""
|
|
AllClassNMS performs non-maximum suppression (NMS) on all classes.
|
|
|
|
Inputs:
|
|
- boxes: (N, 4) tensor of bounding boxes in format [x1, y1, x2, y2]
|
|
- scores: (N, C) tensor of scores for each box and class
|
|
- max_output_boxes_per_class: maximum number of boxes to keep per class
|
|
- iou_threshold: IoU threshold for NMS
|
|
- score_threshold: score threshold for filtering
|
|
|
|
Outputs:
|
|
- selected_indices: (M, 3) tensor with [batch_idx, class_idx, box_idx]
|
|
"""
|
|
boxes = inputs[0]
|
|
scores = inputs[1]
|
|
max_output_boxes_per_class = inputs[2] if len(inputs) > 2 else None
|
|
iou_threshold = inputs[3] if len(inputs) > 3 else None
|
|
score_threshold = inputs[4] if len(inputs) > 4 else None
|
|
|
|
center_point_box = attr.get("center_point_box", 0)
|
|
|
|
if max_output_boxes_per_class is not None and isinstance(
|
|
max_output_boxes_per_class, relax.Constant
|
|
):
|
|
max_output_boxes_per_class = int(max_output_boxes_per_class.data.numpy().item())
|
|
elif max_output_boxes_per_class is not None and isinstance(
|
|
max_output_boxes_per_class, relax.Var
|
|
):
|
|
var_name = max_output_boxes_per_class.name_hint
|
|
if var_name in params[1]:
|
|
_, param_value = params[1][var_name]
|
|
max_output_boxes_per_class = int(param_value.numpy().item())
|
|
else:
|
|
max_output_boxes_per_class = 0 # Default value
|
|
else:
|
|
max_output_boxes_per_class = 0 # Default value
|
|
|
|
if iou_threshold is not None and isinstance(iou_threshold, relax.Constant):
|
|
iou_threshold = float(iou_threshold.data.numpy().item())
|
|
elif iou_threshold is not None and isinstance(iou_threshold, relax.Var):
|
|
var_name = iou_threshold.name_hint
|
|
if var_name in params[1]:
|
|
_, param_value = params[1][var_name]
|
|
iou_threshold = float(param_value.numpy().item())
|
|
else:
|
|
iou_threshold = 0.5 # Default value
|
|
else:
|
|
iou_threshold = 0.5 # Default value
|
|
|
|
if score_threshold is not None and isinstance(score_threshold, relax.Constant):
|
|
score_threshold = float(score_threshold.data.numpy().item())
|
|
elif score_threshold is not None and isinstance(score_threshold, relax.Var):
|
|
var_name = score_threshold.name_hint
|
|
if var_name in params[1]:
|
|
_, param_value = params[1][var_name]
|
|
score_threshold = float(param_value.numpy().item())
|
|
else:
|
|
score_threshold = 0.0 # Default value
|
|
else:
|
|
score_threshold = 0.0 # Default value
|
|
|
|
if center_point_box != 0:
|
|
split_result = relax.op.split(boxes, 4, axis=2)
|
|
xc = split_result[0]
|
|
yc = split_result[1]
|
|
w = split_result[2]
|
|
h = split_result[3]
|
|
half_w = w / relax.const(2.0, boxes.ty.dtype)
|
|
half_h = h / relax.const(2.0, boxes.ty.dtype)
|
|
x1 = xc - half_w
|
|
x2 = xc + half_w
|
|
y1 = yc - half_h
|
|
y2 = yc + half_h
|
|
boxes = relax.op.concat([y1, x1, y2, x2], axis=2)
|
|
|
|
nms_out = bb.normalize(
|
|
relax.op.vision.all_class_non_max_suppression(
|
|
boxes,
|
|
scores,
|
|
relax.const(max_output_boxes_per_class, dtype="int64"),
|
|
relax.const(iou_threshold, dtype="float32"),
|
|
relax.const(score_threshold, dtype="float32"),
|
|
output_format="onnx",
|
|
)
|
|
)
|
|
|
|
return nms_out
|
|
|
|
|
|
class GridSample(OnnxOpConverter):
|
|
"""Converts an onnx GridSample node into an equivalent Relax expression."""
|
|
|
|
@classmethod
|
|
def _impl_v16(cls, bb, inputs, attr, params):
|
|
data = inputs[0]
|
|
grid = inputs[1]
|
|
|
|
method = attr.get("mode", b"bilinear")
|
|
if isinstance(method, bytes):
|
|
method = method.decode("ascii")
|
|
|
|
# Translate ONNX mode names to TVM method names
|
|
if method == "linear":
|
|
method = "bilinear"
|
|
elif method == "cubic":
|
|
method = "bicubic"
|
|
|
|
padding_mode = attr.get("padding_mode", b"zeros")
|
|
if isinstance(padding_mode, bytes):
|
|
padding_mode = padding_mode.decode("ascii")
|
|
|
|
align_corners = bool(attr.get("align_corners", 0))
|
|
|
|
if hasattr(data.ty, "ndim"):
|
|
ndim = data.ty.ndim
|
|
else:
|
|
ndim = len(data.ty.shape)
|
|
|
|
if ndim == 5 and method == "bicubic":
|
|
raise NotImplementedError(
|
|
"5D (volumetric) GridSample with mode='cubic' is not supported "
|
|
"(TOPI 3D grid_sample supports only bilinear and nearest)."
|
|
)
|
|
|
|
if ndim == 4:
|
|
# ONNX grid shape: [N, H_out, W_out, 2]
|
|
# TVM grid shape: [N, 2, H_out, W_out]
|
|
grid = relax.op.permute_dims(grid, [0, 3, 1, 2])
|
|
layout = "NCHW"
|
|
elif ndim == 5:
|
|
# ONNX grid shape: [N, D_out, H_out, W_out, 3]
|
|
# TVM grid shape: [N, 3, D_out, H_out, W_out]
|
|
grid = relax.op.permute_dims(grid, [0, 4, 1, 2, 3])
|
|
layout = "NCDHW"
|
|
else:
|
|
raise NotImplementedError(f"GridSample only supports 4D or 5D input, got {ndim}D.")
|
|
|
|
return relax.op.image.grid_sample(
|
|
data,
|
|
grid,
|
|
method=method,
|
|
layout=layout,
|
|
padding_mode=padding_mode,
|
|
align_corners=align_corners,
|
|
)
|
|
|
|
|
|
class MatMulInteger(OnnxOpConverter):
|
|
"""
|
|
Converts ONNX MatMulInteger (INT8/UINT8 quantized matrix multiply).
|
|
|
|
Computes: output = (A - a_zero_point) * (B - b_zero_point)
|
|
in int32 accumulation, per ONNX spec v10.
|
|
|
|
Zero-point shapes per spec:
|
|
a_zero_point: scalar | [M] (per-row) | [D1, D2, M, 1] (N-D per-row)
|
|
b_zero_point: scalar | [N] (per-col) | [D1, D2, 1, N] (N-D per-col)
|
|
"""
|
|
|
|
@classmethod
|
|
def _impl_v10(cls, bb, inputs, attr, params):
|
|
a = inputs[0]
|
|
b = inputs[1]
|
|
|
|
# Optional zero points with default of None (treated as 0)
|
|
a_zero_point = inputs[2] if len(inputs) > 2 and inputs[2] is not None else None
|
|
b_zero_point = inputs[3] if len(inputs) > 3 and inputs[3] is not None else None
|
|
|
|
# Widen to int32 before any arithmetic to prevent overflow
|
|
a = relax.op.astype(a, "int32")
|
|
b = relax.op.astype(b, "int32")
|
|
|
|
if a_zero_point is not None:
|
|
a_zp = relax.op.astype(
|
|
a_zero_point, "int32"
|
|
) # Ensure zero point is int32 for subtraction
|
|
a_zp = bb.normalize(a_zp) # Normalize the expr so ty gets populated
|
|
a_zp_ndim = len(a_zp.ty.shape)
|
|
|
|
# Per-row case: [M] -> [M, 1] so it broadcasts over [M, K] row-wise
|
|
# N-D case: spec says shape is [D1, D2, M, 1], which already broadcasts correctly (no need to reshape)
|
|
if a_zp_ndim == 1:
|
|
a_zp = relax.op.expand_dims(a_zp, axis=-1)
|
|
|
|
a = relax.op.subtract(a, a_zp)
|
|
|
|
if b_zero_point is not None:
|
|
b_zp = relax.op.astype(b_zero_point, "int32")
|
|
b_zp = bb.normalize(b_zp)
|
|
b_zp_ndim = len(b_zp.ty.shape)
|
|
|
|
# Per-col case: [N] -> [1, N] so it broadcasts over [K, N] column-wise
|
|
# N-D case: [D1, D2, 1, N] already broadcasts correctly
|
|
if b_zp_ndim == 1:
|
|
b_zp = relax.op.expand_dims(b_zp, axis=0)
|
|
|
|
b = relax.op.subtract(b, b_zp)
|
|
|
|
return relax.op.matmul(a, b, out_dtype="int32") # Output is int32 per ONNX spec
|
|
|
|
|
|
def _get_convert_map():
|
|
return {
|
|
# defs/experimental
|
|
"Optional": Optional_,
|
|
"OptionalHasElement": OptionalHasElement,
|
|
"OptionalGetElement": OptionalGetElement,
|
|
# Binary operators
|
|
"Add": Add,
|
|
"Sub": Sub,
|
|
"Mul": Mul,
|
|
"Div": Div,
|
|
"Mod": Mod,
|
|
"Less": Less,
|
|
"LessOrEqual": LessOrEqual,
|
|
"Greater": Greater,
|
|
"GreaterOrEqual": GreaterOrEqual,
|
|
"Equal": Equal,
|
|
"BitwiseAnd": BitwiseAnd,
|
|
"BitwiseOr": BitwiseOr,
|
|
"BitwiseXor": BitwiseXor,
|
|
"BitShift": BitShift,
|
|
"And": And,
|
|
"Or": Or,
|
|
"Xor": Xor,
|
|
"Not": Not,
|
|
# Unary operators
|
|
"BitwiseNot": BitwiseNot,
|
|
"Log": Log,
|
|
"Exp": Exp,
|
|
"Acos": Acos,
|
|
"Acosh": Acosh,
|
|
"Asin": Asin,
|
|
"Asinh": Asinh,
|
|
"Atan": Atan,
|
|
"Atanh": Atanh,
|
|
"Cos": Cos,
|
|
"Cosh": Cosh,
|
|
"Sin": Sin,
|
|
"Sinh": Sinh,
|
|
"Tan": Tan,
|
|
"Tanh": Tanh,
|
|
"Neg": Neg,
|
|
"Abs": Abs,
|
|
"Reciprocal": Reciprocal,
|
|
"Floor": Floor,
|
|
"Ceil": Ceil,
|
|
"Round": Round,
|
|
"IsInf": IsInf,
|
|
"IsNaN": IsNaN,
|
|
"Sqrt": Sqrt,
|
|
"Relu": Relu,
|
|
"Selu": Selu,
|
|
"Mish": Mish,
|
|
"Trilu": Trilu,
|
|
"PRelu": PRelu,
|
|
"LeakyRelu": LeakyRelu,
|
|
"ThresholdedRelu": ThresholdedRelu,
|
|
"Elu": Elu,
|
|
"Gelu": Gelu,
|
|
"FastGelu": FastGelu,
|
|
"BiasGelu": BiasGelu,
|
|
"HardSigmoid": HardSigmoid,
|
|
"HardSwish": HardSwish,
|
|
"Sign": Sign,
|
|
"Softplus": Softplus,
|
|
"Softsign": Softsign,
|
|
"Shrink": Shrink,
|
|
"Erf": Erf,
|
|
"Sum": Sum,
|
|
"Min": Min,
|
|
"Max": Max,
|
|
"Mean": Mean,
|
|
"Cast": Cast,
|
|
"Gemm": Gemm,
|
|
"MatMul": MatMul,
|
|
"MatMulInteger": MatMulInteger,
|
|
"MatMulInteger16": MatMulInteger16,
|
|
"Reshape": Reshape,
|
|
"Sigmoid": Sigmoid,
|
|
"Softmax": Softmax,
|
|
"LogSoftmax": LogSoftmax,
|
|
"Hardmax": Hardmax,
|
|
"Transpose": Transpose,
|
|
"Unsqueeze": Unsqueeze,
|
|
"Where": Where,
|
|
"Concat": Concat,
|
|
"Clip": Clip,
|
|
"Shape": Shape,
|
|
"Pow": Pow,
|
|
"CumSum": CumSum,
|
|
"Squeeze": Squeeze,
|
|
"Constant": Constant,
|
|
"Gather": Gather,
|
|
"GatherElements": GatherElements,
|
|
"GatherND": GatherND,
|
|
"Scatter": Scatter,
|
|
"ScatterElements": ScatterElements,
|
|
"ScatterND": ScatterND,
|
|
"Compress": Compress,
|
|
"Size": Size,
|
|
"EyeLike": EyeLike,
|
|
# Normalization
|
|
"BatchNormalization": BatchNormalization,
|
|
"LayerNormalization": LayerNormalization,
|
|
"RMSNormalization": RMSNormalization,
|
|
"GroupNormalization": GroupNormalization,
|
|
"SkipLayerNormalization": SkipLayerNormalization,
|
|
"EmbedLayerNormalization": EmbedLayerNormalization,
|
|
"InstanceNormalization": InstanceNormalization,
|
|
"MeanVarianceNormalization": MeanVarianceNormalization,
|
|
"LRN": LocalResponseNormalization,
|
|
# defs/reduction
|
|
"ReduceMax": ReduceMax,
|
|
"ReduceMin": ReduceMin,
|
|
"ReduceSum": ReduceSum,
|
|
"ReduceMean": ReduceMean,
|
|
"ReduceProd": ReduceProd,
|
|
"ReduceLogSumExp": ReduceLogSumExp,
|
|
"ReduceLogSum": ReduceLogSum,
|
|
"ReduceSumSquare": ReduceSumSquare,
|
|
"ReduceL1": ReduceL1,
|
|
"ReduceL2": ReduceL2,
|
|
"ArgMax": ArgMax,
|
|
"ArgMin": ArgMin,
|
|
"TopK": TopK,
|
|
"Expand": Expand,
|
|
"ConstantOfShape": ConstantOfShape,
|
|
"Slice": Slice,
|
|
"Attention": Attention,
|
|
"Pad": Pad,
|
|
"Split": Split,
|
|
"Tile": Tile,
|
|
"AveragePool": AveragePool,
|
|
"MaxPool": MaxPool,
|
|
"LpPool": LpPool,
|
|
"GlobalAveragePool": GlobalAveragePool,
|
|
"GlobalMaxPool": GlobalMaxPool,
|
|
"GlobalLpPool": GlobalLpPool,
|
|
"MaxUnpool": MaxUnpool,
|
|
"Conv": Conv,
|
|
"ConvTranspose": ConvTranspose,
|
|
"Flatten": Flatten,
|
|
"Identity": Identity,
|
|
"Dropout": Dropout,
|
|
"Resize": Resize,
|
|
"Einsum": Einsum,
|
|
"Range": Range,
|
|
"OneHot": OneHot,
|
|
"Unique": Unique,
|
|
"NonZero": NonZero,
|
|
"MaxRoiPool": MaxRoiPool,
|
|
"RoiAlign": RoiAlign,
|
|
"NonMaxSuppression": NonMaxSuppression,
|
|
"AllClassNMS": AllClassNMS,
|
|
"GridSample": GridSample,
|
|
"AffineGrid": AffineGrid,
|
|
"Upsample": Upsample,
|
|
# others
|
|
"DepthToSpace": DepthToSpace,
|
|
"SpaceToDepth": SpaceToDepth,
|
|
# Sequence operators
|
|
"SequenceConstruct": SequenceConstruct,
|
|
"SequenceEmpty": SequenceEmpty,
|
|
"SequenceErase": SequenceErase,
|
|
"SequenceInsert": SequenceInsert,
|
|
"SequenceLength": SequenceLength,
|
|
"ConcatFromSequence": ConcatFromSequence,
|
|
"SplitToSequence": SplitToSequence,
|
|
"SequenceAt": SequenceAt,
|
|
# Quantization
|
|
"QuantizeLinear": QuantizeLinear,
|
|
"DequantizeLinear": DequantizeLinear,
|
|
"DynamicQuantizeLinear": DynamicQuantizeLinear,
|
|
}
|
|
|
|
|
|
class ONNXGraphImporter:
|
|
"""A helper class for handling Relax expression copying from pb2.GraphProto.
|
|
Definition: https://github.com/onnx/onnx/blob/main/onnx/onnx.proto
|
|
|
|
Parameters
|
|
----------
|
|
shape_dict : dict of str to tuple, optional
|
|
The input shape to the graph
|
|
dtype_dict : str or dict of str to str
|
|
The input types to the graph
|
|
keep_params_in_input : bool
|
|
If True, parameters will be treated as input variables. If false,
|
|
parameters are treated as constant and folded directly into the graph.
|
|
sanitize : bool
|
|
Whether to sanitize the input names to be valid Relax identifiers.
|
|
"""
|
|
|
|
current = None
|
|
|
|
def __init__(
|
|
self,
|
|
shape_dict: dict[str, list],
|
|
dtype_dict: str | dict[str, str],
|
|
keep_params_in_input: bool = False,
|
|
sanitize: bool = True,
|
|
):
|
|
self._nodes: dict[str, relax.Expr] = {}
|
|
self._inputs: dict[str, relax.Var] = {}
|
|
self._num_input: int = 0
|
|
self._shape = shape_dict.copy() if shape_dict else {}
|
|
self._input_names: list[str] = []
|
|
self._dtype = dtype_dict
|
|
self.opset: int = None
|
|
self._name_supply = UniqueNameSupply()
|
|
self._keep_params_in_input = keep_params_in_input
|
|
self._sanitize: bool = sanitize
|
|
self.bb: relax.BlockBuilder = relax.BlockBuilder() # pylint: disable=invalid-name
|
|
self._params = {}
|
|
|
|
def from_onnx(self, graph: onnx.onnx_ml_pb2.ModelProto, opset: int) -> IRModule:
|
|
"""Construct Relax expressions from the ONNX graph.
|
|
Onnx graph is a python protobuf object.
|
|
|
|
Parameters
|
|
----------
|
|
graph : onnx protobuf object
|
|
The loaded onnx graph
|
|
opset : opset version
|
|
Returns
|
|
-------
|
|
mod : tvm.IRModule
|
|
The returned relax module
|
|
"""
|
|
has_if = any(node.op_type == "If" for node in graph.node)
|
|
|
|
self.opset = opset
|
|
self._parse_graph_initializers(graph)
|
|
self._parse_graph_input(graph)
|
|
self._check_for_unsupported_ops(graph)
|
|
|
|
func_attrs = {"num_input": self._num_input}
|
|
input_list = [value for value in self._inputs.values() if isinstance(value, relax.Var)]
|
|
if self._keep_params_in_input and self._params:
|
|
param_var_list, param_value_list = map(list, zip(*self._params.values()))
|
|
input_list = input_list + param_var_list
|
|
func_attrs["params"] = param_value_list
|
|
|
|
# Enter the function with its parameters already known. This lets
|
|
# BlockBuilder derive non-negative constraints from shape positions
|
|
# before constructing and simplifying the body.
|
|
with self.bb.function("main", params=input_list):
|
|
with contextlib.ExitStack() as stack:
|
|
if not has_if:
|
|
stack.enter_context(self.bb.dataflow())
|
|
|
|
self._construct_nodes(graph)
|
|
|
|
# now return the outputs
|
|
output_names = [self._parse_value_proto(output) for output in graph.output]
|
|
outputs = []
|
|
for output_name in output_names:
|
|
output_value = self._nodes[output_name]
|
|
if _is_empty_optional(output_value):
|
|
raise ValueError(
|
|
"ONNX graph output "
|
|
f"{output_name} is an empty optional. Empty optional graph outputs "
|
|
"are not supported by the Relax ONNX frontend."
|
|
)
|
|
outputs.append(output_value)
|
|
outputs = outputs[0] if len(outputs) == 1 else relax.Tuple(outputs)
|
|
|
|
if has_if:
|
|
output_var = outputs
|
|
else:
|
|
output_var = self.bb.emit_output(outputs)
|
|
|
|
# ExitStack closes here — dataflow block is now closed
|
|
self.bb.emit_func_output(output_var)
|
|
|
|
relax_mod = self.bb.get()
|
|
relax_mod["main"] = relax_mod["main"].with_attrs(func_attrs)
|
|
return relax_mod
|
|
|
|
def _parse_graph_initializers(self, graph: onnx.onnx_ml_pb2.GraphProto):
|
|
"""Parse network inputs to relax, aka parameters."""
|
|
for init_tensor in graph.initializer:
|
|
# There are two cases for handling parameters, they are either
|
|
# treated as variables or constants.
|
|
if not init_tensor.name.strip():
|
|
raise ValueError("Tensor's name is required.")
|
|
array = self._parse_array(init_tensor)
|
|
# Create variables for constants.
|
|
if self._keep_params_in_input:
|
|
# Pytorch sometimes inserts silly weight prefix. Remove it.
|
|
var_name = init_tensor.name.strip("onnx::")
|
|
init_var = self._new_var(var_name, shape=array.shape, dtype=array.dtype)
|
|
self._nodes[init_tensor.name] = init_var
|
|
# We need to keep track of both the real value and variable for this variable.
|
|
self._params[var_name] = (init_var, array)
|
|
# Otherwise we can use the weight as a constant.
|
|
else:
|
|
self._nodes[init_tensor.name] = relax.const(array)
|
|
|
|
def _sanitize_name(self, name: str) -> str:
|
|
"""Sanitize a name to make it a valid identifier.
|
|
If the name is None, returns a string input_0, input_1, etc.
|
|
If the input is an empty string, returns empty_0, empty_1, etc.
|
|
If the input is a string that does not start with a letter or underscore,
|
|
returns input_<name>. Otherwise, returns an unique input name.
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
The name to sanitize
|
|
Returns
|
|
-------
|
|
new_name : str
|
|
"""
|
|
|
|
if name == "":
|
|
return self._name_supply.fresh_name("empty_")
|
|
|
|
new_name = name.replace(".", "_")
|
|
if not new_name[0].isalpha() and new_name[0] != "_":
|
|
new_name = str(self._name_supply.fresh_name("input_" + new_name))
|
|
else:
|
|
new_name = str(self._name_supply.fresh_name(new_name))
|
|
|
|
if new_name != name:
|
|
warnings.warn(f"Renaming name {name} to {new_name}")
|
|
return new_name
|
|
|
|
def _new_var(self, var_name: str, shape: list, dtype: str = "float32"):
|
|
"""Creates a new Relax variable."""
|
|
return relax.Var(name_hint=var_name, ty=relax.TensorType(shape=shape, dtype=dtype))
|
|
|
|
def _parse_graph_input(self, graph: onnx.onnx_ml_pb2.GraphProto):
|
|
"""Parse model inputs to Relax parameters."""
|
|
value_dict = {}
|
|
for i in graph.input:
|
|
# from onnx v0.2, GraphProto.input has type ValueInfoProto,
|
|
# and the name is 'i.name'
|
|
i_name, i_shape, d_type, i_shape_name, value_dict = get_info(i, value_dict)
|
|
if i_name not in self._nodes:
|
|
self._num_input += 1
|
|
self._input_names.append(i_name)
|
|
if i_name in self._shape:
|
|
i_shape = self._shape[i_name]
|
|
else:
|
|
if "?" in str(i_shape):
|
|
warning_msg = (
|
|
f"Input {i_name} has unknown dimension shapes: {i_shape_name!s}. "
|
|
"Specifying static values may improve performance"
|
|
)
|
|
warnings.warn(warning_msg)
|
|
if isinstance(self._dtype, dict):
|
|
dtype = self._dtype[i_name] if i_name in self._dtype else d_type
|
|
else:
|
|
dtype = d_type
|
|
var_name = self._sanitize_name(i_name) if self._sanitize else i_name
|
|
self._nodes[i_name] = self._new_var(var_name, shape=i_shape, dtype=dtype)
|
|
self._inputs[i_name] = self._nodes[i_name]
|
|
|
|
def _check_for_unsupported_ops(self, graph: onnx.onnx_ml_pb2.GraphProto):
|
|
convert_map = _get_convert_map()
|
|
# Ops handled directly in _construct_nodes rather than via the converter map.
|
|
directly_handled_ops = {"If"}
|
|
unsupported_ops = set()
|
|
for node in graph.node:
|
|
op_name = node.op_type
|
|
if (
|
|
op_name not in convert_map
|
|
and op_name not in directly_handled_ops
|
|
and op_name != "Constant"
|
|
):
|
|
unsupported_ops.add(op_name)
|
|
if unsupported_ops:
|
|
msg = "The following operators are not supported for frontend ONNX: "
|
|
msg += ", ".join(unsupported_ops)
|
|
raise tvm.error.OpNotImplemented(msg)
|
|
|
|
def _construct_nodes(self, graph: onnx.onnx_ml_pb2.GraphProto):
|
|
"""Nodes are stored as directed acyclic graph."""
|
|
for node in graph.node:
|
|
op_name = node.op_type
|
|
attr = self._parse_attr(node.attribute)
|
|
# Create and populate input list.
|
|
inputs = onnx_input()
|
|
for i in node.input:
|
|
if i != "":
|
|
inputs.append(self._nodes[i])
|
|
else:
|
|
inputs.append(None)
|
|
i_name = self._parse_value_proto(node)
|
|
outputs = node.output
|
|
attr["tvm_custom"] = {}
|
|
attr["tvm_custom"]["name"] = i_name
|
|
attr["tvm_custom"]["num_outputs"] = len(outputs)
|
|
|
|
if op_name == "If":
|
|
cond = inputs[0]
|
|
then_expr = self._convert_subgraph(self.bb, attr["then_branch"])
|
|
else_expr = self._convert_subgraph(self.bb, attr["else_branch"])
|
|
then_seq = relax.SeqExpr(blocks=[], body=then_expr)
|
|
else_seq = relax.SeqExpr(blocks=[], body=else_expr)
|
|
if_result = self.bb.emit(relax.If(cond, then_seq, else_seq))
|
|
if len(outputs) == 1:
|
|
self._nodes[outputs[0]] = if_result
|
|
else:
|
|
for i, k in enumerate(outputs):
|
|
self._nodes[k] = self.bb.emit(relax.TupleGetItem(if_result, i))
|
|
continue
|
|
|
|
# Perform special handling for shape expressions. If an input is a
|
|
# shape expr, make sure the current op can handle it, otherwise
|
|
# convert it to a tensor.
|
|
shape_compatible_ops = [
|
|
"Reshape",
|
|
"Resize",
|
|
"ConstantOfShape",
|
|
"Gather",
|
|
"Slice",
|
|
"Shape",
|
|
"Expand",
|
|
"Concat",
|
|
"Equal",
|
|
"Where",
|
|
"Cast",
|
|
"Squeeze",
|
|
]
|
|
return_tuple_ops = [
|
|
"Optional",
|
|
"OptionalGetElement",
|
|
"SequenceConstruct",
|
|
"SequenceEmpty",
|
|
"SequenceErase",
|
|
"SequenceInsert",
|
|
"ConcatFromSequence",
|
|
"SplitToSequence",
|
|
]
|
|
for i, inp in enumerate(inputs):
|
|
if (
|
|
inp is not None
|
|
and isinstance(inp, relax.Expr)
|
|
and isinstance(inp.ty, relax.ShapeType)
|
|
and op_name not in shape_compatible_ops
|
|
):
|
|
raise ValueError(f"Node {node.name} cannot handle ShapeExpr inputs.")
|
|
try:
|
|
op = self._convert_operator(op_name, inputs, attr, self.opset)
|
|
# Create type information for the new operator.
|
|
if isinstance(op, relax.Expr):
|
|
op = self.bb.normalize(op)
|
|
except Exception as err: # pylint: disable=broad-exception-caught
|
|
print(f"Error converting operator {op_name}, with inputs: {inputs}")
|
|
raise err
|
|
|
|
if op_name in return_tuple_ops:
|
|
outputs_num = 1
|
|
elif not isinstance(op, relax.Tuple):
|
|
if isinstance(op.ty, relax.TupleType):
|
|
# This is a var bound to a tuple. We need to unpack it and create
|
|
# a new tuple.
|
|
tuple_items = []
|
|
for i in range(len(op.ty.fields)):
|
|
tuple_items.append(self.bb.emit(relax.TupleGetItem(op, i)))
|
|
op = relax.Tuple(tuple_items)
|
|
outputs_num = len(tuple_items)
|
|
else:
|
|
outputs_num = 1
|
|
else:
|
|
outputs_num = len(op)
|
|
assert len(outputs) <= outputs_num, (
|
|
f"Missing outputs during conversion. Expected {len(outputs)} but Got {outputs_num} in {op_name}."
|
|
)
|
|
if outputs_num == 1:
|
|
self._nodes[outputs[0]] = op
|
|
else:
|
|
for i, k in enumerate(outputs):
|
|
self._nodes[k] = op[i]
|
|
|
|
def _parse_value_proto(self, value_proto: onnx.onnx_ml_pb2.GraphProto):
|
|
"""Parse ValueProto or raw str."""
|
|
try:
|
|
name = value_proto.name
|
|
except AttributeError:
|
|
name = value_proto
|
|
return name
|
|
|
|
def _parse_array(self, tensor_proto: onnx.onnx_ml_pb2.TensorProto) -> tvm.runtime.tensor:
|
|
np_array = get_numpy(tensor_proto).reshape(tuple(tensor_proto.dims))
|
|
return tvm.runtime.tensor(np_array)
|
|
|
|
def _parse_attr(self, attr_proto: onnx.onnx_ml_pb2.AttributeProto) -> dict[str, Any]:
|
|
"""Convert a list of AttributeProto to a dict, with names as keys."""
|
|
attrs = {}
|
|
for a in attr_proto:
|
|
for f in ["f", "i", "s", "g"]:
|
|
if a.HasField(f):
|
|
attrs[a.name] = getattr(a, f)
|
|
for f in ["floats", "ints", "strings"]:
|
|
if list(getattr(a, f)):
|
|
assert a.name not in attrs, "Only one type of attr is allowed"
|
|
attrs[a.name] = tuple(getattr(a, f))
|
|
for f in ["t", "tp"]:
|
|
if hasattr(a, f) and a.HasField(f):
|
|
attrs[a.name] = getattr(a, f)
|
|
for f in ["tensors", "type_protos"]:
|
|
if hasattr(a, f) and list(getattr(a, f)):
|
|
assert a.name not in attrs, "Only one type of attr is allowed"
|
|
attrs[a.name] = tuple(getattr(a, f))
|
|
for f in ["graphs"]:
|
|
if list(getattr(a, f)):
|
|
assert a.name not in attrs, "Only one type of attr is allowed"
|
|
attrs[a.name] = tuple(getattr(a, f))
|
|
if a.name not in attrs:
|
|
raise ValueError(f"Cannot parse attribute: \n{a}\n.")
|
|
return attrs
|
|
|
|
def _convert_operator(
|
|
self,
|
|
op_name: str,
|
|
inputs: list[relax.Expr],
|
|
attrs: dict,
|
|
opset: int,
|
|
) -> relax.Expr:
|
|
"""Convert ONNX operator into a Relax operator.
|
|
The converter must specify conversions explicitly for incompatible name, and
|
|
apply handlers to operator attributes.
|
|
|
|
Parameters
|
|
----------
|
|
op_name : str
|
|
Operator name, such as Convolution, FullyConnected
|
|
inputs : list of tvm.relax.function.Function
|
|
List of inputs.
|
|
attrs : dict
|
|
Dict of operator attributes
|
|
opset : int
|
|
Opset version
|
|
Returns
|
|
-------
|
|
sym : tvm.relax.function.Function
|
|
Converted relax function
|
|
"""
|
|
convert_map = _get_convert_map()
|
|
if op_name in convert_map:
|
|
convert_class = convert_map[op_name]
|
|
op_function = convert_class.get_converter(opset)
|
|
sym = op_function(self.bb, inputs, attrs, [self._nodes, self._params])
|
|
else:
|
|
raise NotImplementedError(f"Operator {op_name} not implemented.")
|
|
return sym
|
|
|
|
def _convert_subgraph(self, bb, graph):
|
|
"""
|
|
Walk an ONNX GraphProto (a branch body) and return a Relax SeqExpr.
|
|
Outer-scope nodes are visible because we copy self._nodes into the
|
|
local lookup table before processing.
|
|
"""
|
|
outer_nodes = dict(self._nodes)
|
|
|
|
try:
|
|
for init_tensor in graph.initializer:
|
|
array = self._parse_array(init_tensor)
|
|
self._nodes[init_tensor.name] = relax.const(array)
|
|
|
|
for node in graph.node:
|
|
op_name = node.op_type
|
|
attr = self._parse_attr(node.attribute)
|
|
|
|
inputs = onnx_input()
|
|
for i in node.input:
|
|
if i != "":
|
|
inputs.append(self._nodes.get(i, outer_nodes.get(i)))
|
|
else:
|
|
inputs.append(None)
|
|
|
|
attr["tvm_custom"] = {}
|
|
attr["tvm_custom"]["name"] = node.name
|
|
attr["tvm_custom"]["num_outputs"] = len(node.output)
|
|
|
|
# Handle nested If recursively.
|
|
if op_name == "If":
|
|
cond = inputs[0]
|
|
then_expr = self._convert_subgraph(bb, attr["then_branch"])
|
|
else_expr = self._convert_subgraph(bb, attr["else_branch"])
|
|
then_seq = relax.SeqExpr(blocks=[], body=then_expr)
|
|
else_seq = relax.SeqExpr(blocks=[], body=else_expr)
|
|
op = bb.emit(relax.If(cond, then_seq, else_seq))
|
|
outputs = node.output
|
|
if len(outputs) == 1:
|
|
self._nodes[outputs[0]] = op
|
|
else:
|
|
for i, k in enumerate(outputs):
|
|
self._nodes[k] = bb.emit(relax.TupleGetItem(op, i))
|
|
continue
|
|
|
|
op = self._convert_operator(op_name, inputs, attr, self.opset)
|
|
try:
|
|
_ = op.ty
|
|
has_ty = True
|
|
except tvm.error.InternalError:
|
|
has_ty = False
|
|
|
|
if not has_ty:
|
|
op = bb.normalize(op)
|
|
|
|
if not isinstance(op, relax.Tuple):
|
|
if isinstance(op.ty, relax.TupleType):
|
|
tuple_items = [relax.TupleGetItem(op, i) for i in range(len(op.ty.fields))]
|
|
op = relax.Tuple(tuple_items)
|
|
|
|
outputs = node.output
|
|
if len(outputs) == 1:
|
|
self._nodes[outputs[0]] = op
|
|
else:
|
|
for i, k in enumerate(outputs):
|
|
self._nodes[k] = op[i]
|
|
|
|
branch_outputs = [self._nodes[o.name] for o in graph.output]
|
|
result = branch_outputs[0] if len(branch_outputs) == 1 else relax.Tuple(branch_outputs)
|
|
|
|
self._nodes = outer_nodes
|
|
return result
|
|
finally:
|
|
self._nodes = outer_nodes
|
|
|
|
|
|
def from_onnx(
|
|
model: onnx.onnx_ml_pb2.GraphProto,
|
|
shape_dict: dict[str, list] | None = None,
|
|
dtype_dict: str | dict[str, str] | None = "float32",
|
|
opset: int | None = None,
|
|
keep_params_in_input: bool = False,
|
|
sanitize_input_names: bool = True,
|
|
) -> IRModule:
|
|
"""Convert a ONNX model into an equivalent Relax Function.
|
|
ONNX graphs are represented as Python Protobuf objects.
|
|
|
|
The current implementation assumes that the input model is after ONNX v1.1.0.
|
|
|
|
Parameters
|
|
----------
|
|
model : protobuf object
|
|
ONNX ModelProto after ONNX v1.1.0
|
|
shape_dict : dict of str to tuple, optional
|
|
The input shape to the graph
|
|
dtype_dict : str or dict of str to str, optional
|
|
The input types to the graph
|
|
opset : int, optional
|
|
Override to autodetected opset.
|
|
This can be helpful for some testing.
|
|
keep_params_in_input : bool
|
|
If True, parameters will be treated as input variables. If false,
|
|
parameters are treated as constant and folded directly into the graph.
|
|
sanitize_input_names : bool, optional
|
|
Whether to sanitize the input names to ensure they are valid Relax identifiers.
|
|
|
|
Returns
|
|
-------
|
|
mod : tvm.IRModule
|
|
The relax module for compilation
|
|
"""
|
|
# Error if the model version is below 1.1.0
|
|
if model.ir_version < 3:
|
|
raise ValueError(
|
|
f"Model IR version {model.ir_version} not supported. Must be at least after 1.1.0."
|
|
)
|
|
|
|
try:
|
|
import onnx # pylint: disable=import-outside-toplevel, redefined-outer-name
|
|
|
|
if hasattr(onnx.checker, "check_model"):
|
|
# try use onnx's own model checker before converting any model
|
|
try:
|
|
onnx.checker.check_model(model)
|
|
except Exception as exception: # pylint: disable=c-extension-no-member, broad-except
|
|
# the checker is a bit violent about errors, so simply print warnings here
|
|
warnings.warn(str(exception))
|
|
except ImportError as error:
|
|
raise ImportError(f"Unable to import onnx which is required {error}")
|
|
|
|
g = ONNXGraphImporter(
|
|
shape_dict,
|
|
dtype_dict,
|
|
keep_params_in_input=keep_params_in_input,
|
|
sanitize=sanitize_input_names,
|
|
)
|
|
graph = model.graph
|
|
|
|
try:
|
|
opset_in_model = 1
|
|
if model.opset_import:
|
|
# TODO: for now we only really support ai.onnx op set
|
|
# TODO: handle other namespaces well see https://github.com/apache/tvm/issues/10950
|
|
for opset_identifier in model.opset_import:
|
|
# As per https://github.com/onnx/onnx/blob/main/docs/IR.md
|
|
# All operator sets except the default one must specify the operator version
|
|
if str(opset_identifier.domain) in ["ai.onnx", ""]:
|
|
opset_in_model = opset_identifier.version
|
|
break
|
|
except AttributeError:
|
|
opset_in_model = 1
|
|
|
|
if opset is None:
|
|
opset = opset_in_model
|
|
elif opset < opset_in_model:
|
|
warnings.warn(
|
|
""
|
|
f"You are overwritting original opset ver = {opset_in_model} by lower ver = {opset}. "
|
|
f"That might cause model conversion errors."
|
|
)
|
|
|
|
# Use the graph proto as a scope so that ops can access other nodes if needed.
|
|
return g.from_onnx(graph, opset)
|