chore: import upstream snapshot with attribution

python/tvm/relax/transform/legalize_ops/common.py

@@ -0,0 +1,125 @@
+# 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)