Function FuseOpsByPattern

python/tvm/relax/transform/transform.py:905–962

Apply pattern matching to each function in the given module, and group matched expressions into a new function. The end result is similar to FuseOps, but fusion is driven completely by the provided patterns. Note: Only operates within dataflow blocks. ConvertToDataflow may need to be c

(
    patterns: list[FusionPattern | tuple],
    bind_constants: bool = True,
    annotate_codegen: bool = False,
    entry_functions: list[str] | None = None,
)

def FuseOpsByPattern(
    patterns: list[FusionPattern | tuple],
    bind_constants: bool = True,
    annotate_codegen: bool = False,
    entry_functions: list[str] | None = None,
) -> tvm.ir.transform.Pass:
    """Apply pattern matching to each function in the given module, and group matched expressions
    into a new function.

    The end result is similar to FuseOps, but fusion is driven completely by the provided patterns.

    Note: Only operates within dataflow blocks. ConvertToDataflow may need to be called first.

    Parameters
    ----------
    patterns : List[Union[FusionPattern, Tuple]]
        A list of patterns to be matched. The order of the patterns determines the order of priority
        in which they are matched. Higher-priority patterns should come earlier in the list.

        In addition to FusionPattern, a tuple can be passed as item of this list. The pattern
        will be constructed through :code:`FusionPattern(*item)`

    bind_constants : bool
        Whether or not to keep bound constants in the grouped function.

    annotate_codegen : bool
        If True, wrap each created composite function with another function, whose body consists
        only of a call to the composite function, and annotate the outer function with "Codegen"
        and "global_symbol" attributes. The "Codegen" attribute is set as the prefix of the
        corresponding pattern name. For example, "dnnl" if the pattern name is "dnnl.conv2d_relu".

        This must be True if the created composite functions are intended to be offloaded to
        an external backend without using the MergeCompositeFunctions pass.

    entry_functions : Optional[List[str]]
        The set of entry functions to start from.

    Returns
    -------
    ret : tvm.transform.Pass
        The registered pass for pattern-based fusion.

    """
    converted_patterns = []
    for pattern in patterns:
        if isinstance(pattern, tuple):
            converted_patterns.append(FusionPattern(*pattern))
        elif isinstance(pattern, FusionPattern):
            converted_patterns.append(pattern)
        else:
            raise ValueError(f"Invalid pattern: {pattern}")

    return _ffi_api.FuseOpsByPattern(
        converted_patterns,
        bind_constants,
        annotate_codegen,
        entry_functions or [],
    )  # type: ignore