wehub-resource-sync
126 lines
4.5 KiB
Python
126 lines
4.5 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.
|
|
"""Common functionality for legalization."""
|
|
|
|
from collections.abc import Callable
|
|
|
|
import tvm
|
|
from tvm import te
|
|
from tvm.ir import Call
|
|
from tvm.runtime import DataTypeCode
|
|
from tvm.tirx import FloatImm, IntImm
|
|
|
|
from ...block_builder import BlockBuilder
|
|
from ...expr import Constant, Expr
|
|
|
|
##################### Types #####################
|
|
|
|
|
|
# The function type of a TE function, which accepts TE Tensors and
|
|
# other attributes, and returns the output TE Tensor.
|
|
TEFunc = Callable[..., te.Tensor]
|
|
|
|
# The function type of a legalization function, which takes a
|
|
# BlockBuilder and the Call to be legalized, and outputs the legalization
|
|
# result Expr.
|
|
LegalizeFunc = Callable[[BlockBuilder, Call], Expr]
|
|
|
|
|
|
def _is_relax_expr(expr: object) -> bool:
|
|
return isinstance(expr, Expr) and not tvm.ir.is_prim_expr(expr)
|
|
|
|
|
|
def _try_convert_to_scalar_const(
|
|
expr: Expr, python_native: bool = False
|
|
) -> Expr | FloatImm | IntImm | bool | float | int:
|
|
"""Check if the input Expr is a scalar constant.
|
|
If it is, return its plain value with the same data type or in native python type.
|
|
If it is not, return the input expr.
|
|
|
|
Note that if the python_native flag is True, the returned value will be in native python type,
|
|
this might cause loss of data type for example, a float16 constant will be converted to float32
|
|
and a int64 constant will be converted to int32.
|
|
|
|
Parameters
|
|
----------
|
|
expr : Expr
|
|
The expr to be checked and converted.
|
|
|
|
Returns
|
|
-------
|
|
ret : Union[Expr, FloatImm, IntImm, bool, float, int]
|
|
Return a FloatImm or IntImm if the given expr is a scalar integer or float constant, and the
|
|
python native flag is False. Or return the plain value of the constant in native python type
|
|
if the python native flag is True.
|
|
Or return the input itself if it is not a scalar constant.
|
|
"""
|
|
if isinstance(expr, Constant) and expr.ty.ndim == 0:
|
|
# get the value of the scalar constant
|
|
value = expr.data.numpy()[()].item()
|
|
dtype = expr.ty.dtype
|
|
dtype_str = str(dtype.dtype)
|
|
if python_native:
|
|
return value
|
|
# preserve the data type of the constant
|
|
if dtype.matches_code(DataTypeCode.FLOAT, DataTypeCode.BFLOAT):
|
|
return tvm.tirx.FloatImm(dtype_str, value)
|
|
elif dtype.matches_code(DataTypeCode.INT, DataTypeCode.UINT, DataTypeCode.BOOL):
|
|
return tvm.tirx.IntImm(dtype_str, value)
|
|
return expr
|
|
|
|
|
|
def _call_topi_without_attr(te_func: TEFunc, primfunc_name: str | None = None) -> LegalizeFunc:
|
|
"""A common wrapper util for the ops who has no attributes and whose
|
|
legalization is simply passing its arguments to some TE function.
|
|
|
|
Parameters
|
|
----------
|
|
te_func : TEFunc
|
|
The input TE function which is to be converted to PrimFunc.
|
|
|
|
primfunc_name : Optional[str]
|
|
The name of the generated PrimFunc.
|
|
If it is not specified, the name of `te_func` will be used by default.
|
|
|
|
Returns
|
|
-------
|
|
func : LegalizeFunc
|
|
The legalization wrapper function, which wraps the input TE function.
|
|
"""
|
|
if primfunc_name is None:
|
|
primfunc_name = te_func.__name__
|
|
return lambda bb, call: bb.call_te(te_func, *call.args, primfunc_name_hint=primfunc_name)
|
|
|
|
|
|
##################### Decorators #####################
|
|
|
|
_LEGALIZE_ATTR_NAME = "FLegalize"
|
|
|
|
|
|
def register_legalize(op_name: str, legal_func: LegalizeFunc = None):
|
|
"""Register legal transformation function for a Relax op.
|
|
|
|
Parameters
|
|
----------
|
|
op_name : str
|
|
The name of the operator
|
|
|
|
legal_func: function (bb: BlockBuilder, call: Call) -> new_expr: Expr
|
|
The function for transforming an expr to another expr.
|
|
"""
|
|
return tvm.ir.register_op_attr(op_name, _LEGALIZE_ATTR_NAME, legal_func)
|