Function LegalizeOps

python/tvm/relax/transform/transform.py:1068–1179

Legalize high-level operator calls in Relax functions to call_tir with corresponding low-level TIR PrimFuncs. For each high-level operator, we register the way of legalizing it as a function, which takes a context BlockBuilder and the Call being legalized as input, and returns the l

(
    customize_legalize_map: dict[str, LegalizeFunc] | None = None,
    skip_ops: list[str] | None = None,
    enable_warning: bool = False,
)

def LegalizeOps(
    customize_legalize_map: dict[str, LegalizeFunc] | None = None,
    skip_ops: list[str] | None = None,
    enable_warning: bool = False,
):
    """Legalize high-level operator calls in Relax functions to call_tir
    with corresponding low-level TIR PrimFuncs.

    For each high-level operator, we register the way of legalizing it as a
    function, which takes a context BlockBuilder and the Call being legalized
    as input, and returns the legalized call. Here the input BlockBuilder is
    mainly used for adding the PrimFunc created by call_te into the context
    IRModule.

    The legalization function for each operator is registered as an attribute (with
    attribute key `FLegalize`) of the operator.

    This pass provides customizability for users to use their own legalization
    function for operators. The pass takes an optional customized map,
    with the key to be the operator name (`str`) and value to be the function
    (`LegalizeFunc`). The default legalization function will be overridden by the customized
    one.

    Parameters
    ----------
    customize_legalize_map : Optional[Dict[str, LegalizeFunc]]
        The customized operator legalization function map. The customized function will override
        the default one.

    skip_ops : Optional,List[str]]
        List of ops that need to be skipped from legalization

    enable_warning : bool
        A boolean value indicating if to print warnings for CallNode whose op's
        legalization function is not registered. By default we don't print
        warnings.

    Returns
    -------
    ret : tvm.transform.Pass
        The registered pass

    Examples
    --------
    The following code shows how to use this pass:

    .. code-block:: python

        # Define the pass input IRModule
        @tvm.script.ir_module
        class Module:
            @R.function
            def main(
                x: R.Tensor((2, 3), "float32"), y: R.Tensor((2, 3), "float32")
            ) -> R.Tensor((2, 3), "float32"):
                z: R.Tensor((2, 3), "float32") = R.add(x, y)
                r: R.Tensor((2, 3), "float32") = R.multiply(y, z)
                return r