Method replace

python/tvm/s_tir/schedule/state.py:199–238

Replace the part of the AST, as being pointed to by `src_sref`, with a specific statement `tgt_stmt`, and maintain the sref tree accordingly. Replace will try to perform copy on write as much as possible when the ScheduleState holds the only copy to the IRModule and

(
        self,
        src_sref: StmtSRef,
        tgt_stmt: SBlock | For | SBlockRealize,
        block_sref_reuse: dict[SBlock, SBlock] | None = None,
    )

        )

    def replace(
        self,
        src_sref: StmtSRef,
        tgt_stmt: SBlock | For | SBlockRealize,
        block_sref_reuse: dict[SBlock, SBlock] | None = None,
    ) -> None:
        """
        Replace the part of the AST, as being pointed to by `src_sref`,
        with a specific statement `tgt_stmt`, and maintain the sref tree accordingly.
        Replace will try to perform copy on write as much as possible when the ScheduleState holds
        the only copy to the IRModule and IR nodes.

        Only 3 types of replacements are allowed: from `src_sref->stmt` to `tgt_stmt`.
        1) SBlock -> SBlock
        2) Loop -> Loop
        3) Loop -> BlockRealize

        Parameters
        ----------
        src_sref : StmtSRef
            The sref to the statement to be replaced in the TensorIR AST

        tgt_stmt : Union[Block, For, BlockRealize]
            The statement to be replaced to

        block_sref_reuse : Optional[Dict[Block, Block]] = None
            Maps an old block (to be replaced in the subtree under `src_sref->stmt`)
            to a new block (replaced to, in the subtree under `tgt_stmt`), and enforces
            reuse of srefs between them (rather than create new srefs) i.e. after being replaced,
            the sref that points to the old block will point to the new one

        Note
        ----
        The reuse of loop srefs are detected automatically according to the reuse of loop vars.
        """
        if block_sref_reuse is None:
            block_sref_reuse = {}
        _ffi_api.ScheduleStateReplace(  # type: ignore # pylint: disable=no-member
            self, src_sref, tgt_stmt, block_sref_reuse
        )