Method split

python/tvm/s_tir/schedule/schedule.py:738–834

Split a loop into a list of consecutive loops. It requires: - The loop can't have annotation or thread binding. - The loop must start with 0. Predicates may be added to ensure the total loop numbers keeps unchanged. In `factors`, at most one of the factors can be No

(
        self,
        loop: LoopRV,
        factors: list[int | ExprRV | None],
        preserve_unit_iters: bool = True,
        disable_predication: bool = False,
    )

    @type_checked
    def split(
        self,
        loop: LoopRV,
        factors: list[int | ExprRV | None],
        preserve_unit_iters: bool = True,
        disable_predication: bool = False,
    ) -> list[LoopRV]:
        """Split a loop into a list of consecutive loops. It requires:

        - The loop can't have annotation or thread binding.
        - The loop must start with 0.

        Predicates may be added to ensure the total loop numbers keeps unchanged.
        In `factors`, at most one of the factors can be None,
        which will be automatically inferred.

        Parameters
        ----------
        loop : LoopRV
            The loop to be split

        factors: List[int | ExprRV | None]
            The splitting factors
            Potential inputs are:

            - None
            - ExprRV
            - Positive constant integers

        preserve_unit_iters : bool
            Whether or not to preserve unit iterators in block bindings

        disable_predication : bool
            If enabled, don't create a predicate for guarding the loop. This can
            be useful when splitting with scalable factors that the schedule writer
            knows are divisible by the loop bound.

            Warning: enabling this feature may result in incorrect code generation
            if not used carefully.

        Returns
        -------
        split_loops : List[LoopRV]
            The new loops after split

        Examples
        --------

        Before split, in TensorIR, the IR is:

        .. code-block:: python

            @T.prim_func(s_tir=True)
            def before_split(a: T.handle, b: T.handle) -> None:
                A = T.match_buffer(a, (128, 128))
                B = T.match_buffer(b, (128, 128))
                for i, j in T.grid(128, 128):
                    with T.sblock("B"):