Class DataflowBlockRewrite

python/tvm/relax/binding_rewrite.py:29–158

A binding/statement-level dataflow block rewriter. Notes ----- Due to the immutable and copy-on-write nature of TVM AST nodes, the rewriting is not done in place. Instead, a new DataflowBlock is created and returned with mutated_dfb. Similarly, its new root Function is crea

@tvm_ffi.register_object("relax.DataflowBlockRewrite")
class DataflowBlockRewrite(Object):
    """
    A binding/statement-level dataflow block rewriter.

    Notes
    -----
    Due to the immutable and copy-on-write nature of TVM AST nodes, the rewriting is not done in
    place. Instead, a new DataflowBlock is created and returned with mutated_dfb. Similarly, its new
    root Function is created and returned by mutated_root_fn. To apply this change for an IRModule,
    use mutate_irmodule which rewrites the old function that registered in the constructor.
    """

    __slots__ = ("__dict__",)

    def __init__(self, dfb: DataflowBlock, root_fn: Function):
        """
        Construct a rewriter with the DataflowBlock to rewrite and its root function.

        Parameters
        ----------
        dfb : DataflowBlock
            The DataflowBlock to rewrite.
        root_fn : Function
            The root function of the DataflowBlock.
        """
        self.func_name = root_fn.__name__ if hasattr(root_fn, "__name__") else None
        self.__init_handle_by_constructor__(
            _ffi_api.DataflowBlockRewrite,
            dfb,
            root_fn,  # type: ignore
        )

    def replace_all_uses(self, old_var: Var, new_var: Var) -> None:
        """
        Replace all uses of old_var with new_var.

        Parameters
        ----------
        old_var : Var
            The old variable to replace.
        new_var : Var
            The new variable to replace with.
        """
        _ffi_api.dfb_rewrite_replace_all_uses(self, old_var, new_var)  # type: ignore

    def add_binding(self, binding: Binding) -> None:
        return _ffi_api.dfb_rewrite_add_binding(self, binding)  # type: ignore

    def add(self, expr: Expr, name: str | None = None, is_dfvar: bool = False) -> None:
        """
        Add a new statement to the DataflowBlock with an automatically generated variable name.

        Parameters
        ----------
        expr : Expr
            The expression to add.
        name : Optional[str], optional
            Variable name, by default None