Method split

tests/python/relax/test_transform_fold_constant.py:452–464

(
            A: T.Buffer((4, 4), "float32"),
            B: T.Buffer((2, 4), "float32"),
            C: T.Buffer((2, 4), "float32"),
        )

    class Module:
        @T.prim_func(s_tir=True)
        def split(
            A: T.Buffer((4, 4), "float32"),
            B: T.Buffer((2, 4), "float32"),
            C: T.Buffer((2, 4), "float32"),
        ) -> None:
            for i, j in T.grid(2, 4):
                with T.sblock("upper"):
                    vi, vj = T.axis.remap("SS", [i, j])
                    B[vi, vj] = A[vi, vj]
            for i, j in T.grid(2, 4):
                with T.sblock("lower"):
                    vi, vj = T.axis.remap("SS", [i, j])
                    C[vi, vj] = A[vi + 2, vj]

        @R.function
        def before(c0: R.Tensor((4, 4), "float32")):