# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
# ruff: noqa: E501, F841

import tvm
import tvm.testing
from tvm import relax
from tvm.script import ir as I
from tvm.script import relax as R
from tvm.script import tirx as T


def test_reshape_expand_dims():
    @tvm.script.ir_module
    class Module:
        @T.prim_func(s_tir=True)
        def reshape(
            rxplaceholder: T.Buffer((T.int64(8), T.int64(3)), "float32"),
            T_reshape: T.Buffer((T.int64(2), T.int64(4), T.int64(3)), "float32"),
        ):
            for ax0, ax1, ax2 in T.grid(T.int64(2), T.int64(4), T.int64(3)):
                with T.sblock("T_reshape"):
                    v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
                    T.reads(
                        rxplaceholder[
                            (v_ax0 * 12 + v_ax1 * 3 + v_ax2) // T.int64(3),
                            (v_ax0 * 12 + v_ax1 * 3 + v_ax2) % T.int64(3),
                        ]
                    )
                    T.writes(T_reshape[v_ax0, v_ax1, v_ax2])
                    T_reshape[v_ax0, v_ax1, v_ax2] = rxplaceholder[
                        (v_ax0 * 12 + v_ax1 * 3 + v_ax2) // T.int64(3),
                        (v_ax0 * 12 + v_ax1 * 3 + v_ax2) % T.int64(3),
                    ]

        @T.prim_func(s_tir=True)
        def expand_dims(
            rxplaceholder: T.Buffer((T.int64(2), T.int64(4), T.int64(3)), "float32"),
            expand_dims: T.Buffer(
                (T.int64(2), T.int64(1), T.int64(4), T.int64(1), T.int64(3)), "float32"
            ),
        ):
            for i0, i1, i2, i3, i4 in T.grid(
                T.int64(2), T.int64(1), T.int64(4), T.int64(1), T.int64(3)
            ):
                with T.sblock("expand_dims"):
                    i0_1, i1_1, i2_1, i3_1, i4_1 = T.axis.remap("SSSSS", [i0, i1, i2, i3, i4])
                    T.reads(rxplaceholder[i0_1, i2_1, i4_1])
                    T.writes(expand_dims[i0_1, i1_1, i2_1, i3_1, i4_1])
                    expand_dims[i0_1, i1_1, i2_1, i3_1, i4_1] = rxplaceholder[i0_1, i2_1, i4_1]

        @R.function
        def main(x: R.Tensor((8, 3), dtype="float32")) -> R.Tensor(
            (2, 1, 4, 1, 3), dtype="float32"
        ):
            cls = Module
            with R.dataflow():
                y = R.call_tir(cls.reshape, (x,), out_ty=R.Tensor((2, 4, 3), dtype="float32"))
                z = R.call_tir(cls.expand_dims, (y,), out_ty=R.Tensor((2, 1, 4, 1, 3), "float32"))
                R.output(z)
            return z

    @tvm.script.ir_module
    class Expected:
        @T.prim_func(s_tir=True)
        def reshape(
            rxplaceholder: T.Buffer((T.int64(8), T.int64(3)), "float32"),
            T_reshape: T.Buffer((T.int64(2), T.int64(4), T.int64(3)), "float32"),
        ):
            for ax0, ax1, ax2 in T.grid(T.int64(2), T.int64(4), T.int64(3)):
                with T.sblock("T_reshape"):
                    v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
                    T.reads(
                        rxplaceholder[
                            (v_ax0 * T.int64(12) + v_ax1 * T.int64(3) + v_ax2) // T.int64(3),
                            (v_ax0 * T.int64(12) + v_ax1 * T.int64(3) + v_ax2) % T.int64(3),
                        ]
                    )
                    T.writes(T_reshape[v_ax0, v_ax1, v_ax2])
                    T_reshape[v_ax0, v_ax1, v_ax2] = rxplaceholder[
                        (v_ax0 * T.int64(12) + v_ax1 * T.int64(3) + v_ax2) // T.int64(3),
                        (v_ax0 * T.int64(12) + v_ax1 * T.int64(3) + v_ax2) % T.int64(3),
                    ]

        @T.prim_func(s_tir=True)
        def expand_dims(
            rxplaceholder: T.Buffer((T.int64(2), T.int64(4), T.int64(3)), "float32"),
            expand_dims: T.Buffer(
                (T.int64(2), T.int64(1), T.int64(4), T.int64(1), T.int64(3)), "float32"
            ),
        ):
            for i0, i1, i2, i3, i4 in T.grid(
                T.int64(2), T.int64(1), T.int64(4), T.int64(1), T.int64(3)
            ):
                with T.sblock("expand_dims"):
                    i0_1, i1_1, i2_1, i3_1, i4_1 = T.axis.remap("SSSSS", [i0, i1, i2, i3, i4])
                    T.reads(rxplaceholder[i0_1, i2_1, i4_1])
                    T.writes(expand_dims[i0_1, i1_1, i2_1, i3_1, i4_1])
                    expand_dims[i0_1, i1_1, i2_1, i3_1, i4_1] = rxplaceholder[i0_1, i2_1, i4_1]

        @R.function
        def main(x: R.Tensor((8, 3), dtype="float32")) -> R.Tensor(
            (2, 1, 4, 1, 3), dtype="float32"
        ):
            with R.dataflow():
                cls = Expected
                y: R.Tensor((2, 4, 3), "float32") = R.reshape(x, (2, 4, 3))
                # Note: `z` is the output var of the dataflow block, and is thus
                # not expected to be rewritten.
                z = R.call_tir(
                    cls.expand_dims, (y,), out_ty=R.Tensor((2, 1, 4, 1, 3), dtype="float32")
                )
                R.output(z)
            return z

    assert relax.analysis.has_reshape_pattern(Module["expand_dims"])
    mod = relax.transform.RewriteDataflowReshape()(Module)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_reshape_pattern_detect():
    # fmt: off
    @tvm.script.ir_module
    class Module:
        @T.prim_func(s_tir=True)
        def reshape(rxplaceholder: T.Buffer((T.int64(2), T.int64(4096), T.int64(320)), "float32"), T_reshape: T.Buffer((T.int64(2), T.int64(4096), T.int64(5), T.int64(64)), "float32")):
            for ax0_ax1_ax2_ax3_fused_1 in T.thread_binding(T.int64(256), thread="blockIdx.x"):
                for ax0_ax1_ax2_ax3_fused_2 in T.thread_binding(T.int64(1024), thread="threadIdx.x"):
                    for ax0_ax1_ax2_ax3_fused_0 in range(T.int64(10)):
                        with T.sblock("T_reshape"):
                            v_ax0 = T.axis.spatial(T.int64(2), (ax0_ax1_ax2_ax3_fused_0 * T.int64(262144) + ax0_ax1_ax2_ax3_fused_1 * T.int64(1024) + ax0_ax1_ax2_ax3_fused_2) // T.int64(1310720))
                            v_ax1 = T.axis.spatial(T.int64(4096), (ax0_ax1_ax2_ax3_fused_0 * T.int64(262144) + ax0_ax1_ax2_ax3_fused_1 * T.int64(1024) + ax0_ax1_ax2_ax3_fused_2) % T.int64(1310720) // T.int64(320))
                            v_ax2 = T.axis.spatial(T.int64(5), (ax0_ax1_ax2_ax3_fused_0 * T.int64(262144) + ax0_ax1_ax2_ax3_fused_1 * T.int64(1024) + ax0_ax1_ax2_ax3_fused_2) % T.int64(320) // T.int64(64))
                            v_ax3 = T.axis.spatial(T.int64(64), (ax0_ax1_ax2_ax3_fused_0 * T.int64(262144) + ax0_ax1_ax2_ax3_fused_1 * T.int64(1024) + ax0_ax1_ax2_ax3_fused_2) % T.int64(64))
                            T.reads(rxplaceholder[(((v_ax2 * T.int64(64) + v_ax3) // T.int64(320) + v_ax1) // T.int64(4096) + v_ax0) % T.int64(2), ((v_ax2 * T.int64(64) + v_ax3) // T.int64(320) + v_ax1) % T.int64(4096), (v_ax2 * T.int64(64) + v_ax3) % T.int64(320)])
                            T.writes(T_reshape[v_ax0, v_ax1, v_ax2, v_ax3])
                            T_reshape[v_ax0, v_ax1, v_ax2, v_ax3] = rxplaceholder[(((v_ax2 * T.int64(64) + v_ax3) // T.int64(320) + v_ax1) // T.int64(4096) + v_ax0) % T.int64(2), ((v_ax2 * T.int64(64) + v_ax3) // T.int64(320) + v_ax1) % T.int64(4096), (v_ax2 * T.int64(64) + v_ax3) % T.int64(320)]

        @T.prim_func(s_tir=True)
        def expand_dims(
            rxplaceholder: T.Buffer((T.int64(2), T.int64(4096), T.int64(5), T.int64(64)), "float32"),
            expand_dims: T.Buffer(
                (T.int64(2), T.int64(1), T.int64(4096), T.int64(1), T.int64(5), T.int64(64)),
                "float32",
            ),
        ):
            for i0, i1, i2, i3, i4, i5 in T.grid(
                T.int64(2), T.int64(1), T.int64(4096), T.int64(1), T.int64(5), T.int64(64)
            ):
                with T.sblock("expand_dims"):
                    i0_1, i1_1, i2_1, i3_1, i4_1, i5_1 = T.axis.remap("SSSSSS", [i0, i1, i2, i3, i4, i5])
                    T.reads(rxplaceholder[i0_1, i2_1, i4_1, i5_1])
                    T.writes(expand_dims[i0_1, i1_1, i2_1, i3_1, i4_1, i5_1])
                    expand_dims[i0_1, i1_1, i2_1, i3_1, i4_1, i5_1] = rxplaceholder[i0_1, i2_1, i4_1, i5_1]

        @R.function
        def main(
            x: R.Tensor((2, 4096, 320), dtype="float32")
        ) -> R.Tensor((2, 1, 4096, 1, 5, 64), dtype="float32"):
            cls = Module
            with R.dataflow():
                y = R.call_tir(cls.reshape, (x,), out_ty=R.Tensor((2, 4096, 5, 64), dtype="float32"))
                z = R.call_tir(
                    cls.expand_dims, (y,), out_ty=R.Tensor((2, 1, 4096, 1, 5, 64), "float32")
                )
                R.output(z)
            return z

    @tvm.script.ir_module
    class Expected:
        @T.prim_func(s_tir=True)
        def expand_dims(rxplaceholder: T.Buffer((T.int64(2), T.int64(4096), T.int64(5), T.int64(64)), "float32"), expand_dims_1: T.Buffer((T.int64(2), T.int64(1), T.int64(4096), T.int64(1), T.int64(5), T.int64(64)), "float32")):
            # with T.sblock("root"):
            for i0, i1, i2, i3, i4, i5 in T.grid(T.int64(2), T.int64(1), T.int64(4096), T.int64(1), T.int64(5), T.int64(64)):
                with T.sblock("expand_dims"):
                    i0_1, i1_1, i2_1, i3_1, i4_1, i5_1 = T.axis.remap("SSSSSS", [i0, i1, i2, i3, i4, i5])
                    T.reads(rxplaceholder[i0_1, i2_1, i4_1, i5_1])
                    T.writes(expand_dims_1[i0_1, i1_1, i2_1, i3_1, i4_1, i5_1])
                    expand_dims_1[i0_1, i1_1, i2_1, i3_1, i4_1, i5_1] = rxplaceholder[i0_1, i2_1, i4_1, i5_1]

        @T.prim_func(s_tir=True)
        def reshape(rxplaceholder: T.Buffer((T.int64(2), T.int64(4096), T.int64(320)), "float32"), T_reshape: T.Buffer((T.int64(2), T.int64(4096), T.int64(5), T.int64(64)), "float32")):
            # with T.sblock("root"):
            for ax0_ax1_ax2_ax3_fused_1 in T.thread_binding(T.int64(256), thread="blockIdx.x"):
                for ax0_ax1_ax2_ax3_fused_2 in T.thread_binding(T.int64(1024), thread="threadIdx.x"):
                    for ax0_ax1_ax2_ax3_fused_0 in range(T.int64(10)):
                        with T.sblock("T_reshape"):
                            v_ax0 = T.axis.spatial(T.int64(2), (ax0_ax1_ax2_ax3_fused_0 * T.int64(262144) + ax0_ax1_ax2_ax3_fused_1 * T.int64(1024) + ax0_ax1_ax2_ax3_fused_2) // T.int64(1310720))
                            v_ax1 = T.axis.spatial(T.int64(4096), (ax0_ax1_ax2_ax3_fused_0 * T.int64(262144) + ax0_ax1_ax2_ax3_fused_1 * T.int64(1024) + ax0_ax1_ax2_ax3_fused_2) % T.int64(1310720) // T.int64(320))
                            v_ax2 = T.axis.spatial(T.int64(5), (ax0_ax1_ax2_ax3_fused_0 * T.int64(262144) + ax0_ax1_ax2_ax3_fused_1 * T.int64(1024) + ax0_ax1_ax2_ax3_fused_2) % T.int64(320) // T.int64(64))
                            v_ax3 = T.axis.spatial(T.int64(64), (ax0_ax1_ax2_ax3_fused_0 * T.int64(262144) + ax0_ax1_ax2_ax3_fused_1 * T.int64(1024) + ax0_ax1_ax2_ax3_fused_2) % T.int64(64))
                            T.reads(rxplaceholder[(((v_ax2 * T.int64(64) + v_ax3) // T.int64(320) + v_ax1) // T.int64(4096) + v_ax0) % T.int64(2), ((v_ax2 * T.int64(64) + v_ax3) // T.int64(320) + v_ax1) % T.int64(4096), (v_ax2 * T.int64(64) + v_ax3) % T.int64(320)])
                            T.writes(T_reshape[v_ax0, v_ax1, v_ax2, v_ax3])
                            T_reshape[v_ax0, v_ax1, v_ax2, v_ax3] = rxplaceholder[(((v_ax2 * T.int64(64) + v_ax3) // T.int64(320) + v_ax1) // T.int64(4096) + v_ax0) % T.int64(2), ((v_ax2 * T.int64(64) + v_ax3) // T.int64(320) + v_ax1) % T.int64(4096), (v_ax2 * T.int64(64) + v_ax3) % T.int64(320)]

        @R.function
        def main(x: R.Tensor((2, 4096, 320), dtype="float32")) -> R.Tensor((2, 1, 4096, 1, 5, 64), dtype="float32"):
            cls = Expected
            with R.dataflow():
                y: R.Tensor((2, 4096, 5, 64), dtype="float32") = R.reshape(x, R.shape([2, 4096, 5, 64]))
                z = R.call_tir(cls.expand_dims, (y,), out_ty=R.Tensor((2, 1, 4096, 1, 5, 64), dtype="float32"))
                R.output(z)
            return z
    # fmt: on

    assert relax.analysis.has_reshape_pattern(Module["reshape"])
    mod = relax.transform.RewriteDataflowReshape()(Module)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_reshape_dynamic_shape():
    @tvm.script.ir_module
    class Module:
        @T.prim_func(private=True, s_tir=True)
        def reshape(var_A: T.handle, var_T_reshape: T.handle):
            T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
            n = T.int32()
            A = T.match_buffer(var_A, (n, 16, 128), "float16")
            T_reshape = T.match_buffer(var_T_reshape, (1, n, 16, 128), "float16")
            # with T.sblock("root"):
            for ax0_ax1_ax2_fused_0 in T.thread_binding(n * 2, thread="blockIdx.x"):
                for ax0_ax1_ax2_fused_1 in T.thread_binding(1024, thread="threadIdx.x"):
                    with T.sblock("T_reshape"):
                        v0 = T.axis.spatial(
                            n, (ax0_ax1_ax2_fused_0 * 1024 + ax0_ax1_ax2_fused_1) // 2048
                        )
                        v1 = T.axis.spatial(
                            16, (ax0_ax1_ax2_fused_0 * 1024 + ax0_ax1_ax2_fused_1) % 2048 // 128
                        )
                        v2 = T.axis.spatial(
                            128, (ax0_ax1_ax2_fused_0 * 1024 + ax0_ax1_ax2_fused_1) % 128
                        )
                        T.reads(
                            A[((v2 // 128 + v1) // 32 + v0) % n, (v2 // 128 + v1) % 32, v2 % 128]
                        )
                        T.writes(T_reshape[0, v0, v1, v2])
                        T_reshape[0, v0, v1, v2] = A[
                            ((v2 // 128 + v1) // 32 + v0) % n, (v2 // 128 + v1) % 32, v2 % 128
                        ]

        @R.function
        def main(x: R.Tensor((8, 16, 128), dtype="float16")) -> R.Tensor(
            (1, 8, 16, 128), dtype="float16"
        ):
            cls = Module
            with R.dataflow():
                y = R.call_tir(cls.reshape, (x,), out_ty=R.Tensor((1, 8, 16, 128), dtype="float16"))
                z = R.add(y, R.const(1, "float16"))
                R.output(z)
            return z

    @tvm.script.ir_module
    class Expected:
        @T.prim_func(private=True, s_tir=True)
        def reshape(var_A: T.handle, var_T_reshape: T.handle):
            T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
            n = T.int32()
            A = T.match_buffer(var_A, (n, 16, 128), "float16")
            T_reshape = T.match_buffer(var_T_reshape, (1, n, 16, 128), "float16")
            # with T.sblock("root"):
            for ax0_ax1_ax2_fused_0 in T.thread_binding(n * 2, thread="blockIdx.x"):
                for ax0_ax1_ax2_fused_1 in T.thread_binding(1024, thread="threadIdx.x"):
                    with T.sblock("T_reshape"):
                        v0 = T.axis.spatial(
                            n, (ax0_ax1_ax2_fused_0 * 1024 + ax0_ax1_ax2_fused_1) // 2048
                        )
                        v1 = T.axis.spatial(
                            16, (ax0_ax1_ax2_fused_0 * 1024 + ax0_ax1_ax2_fused_1) % 2048 // 128
                        )
                        v2 = T.axis.spatial(
                            128, (ax0_ax1_ax2_fused_0 * 1024 + ax0_ax1_ax2_fused_1) % 128
                        )
                        T.reads(
                            A[((v2 // 128 + v1) // 32 + v0) % n, (v2 // 128 + v1) % 32, v2 % 128]
                        )
                        T.writes(T_reshape[0, v0, v1, v2])
                        T_reshape[0, v0, v1, v2] = A[
                            ((v2 // 128 + v1) // 32 + v0) % n, (v2 // 128 + v1) % 32, v2 % 128
                        ]

        @R.function
        def main(x: R.Tensor((8, 16, 128), dtype="float16")) -> R.Tensor(
            (1, 8, 16, 128), dtype="float16"
        ):
            with R.dataflow():
                y: R.Tensor((1, 8, 16, 128), dtype="float16") = R.reshape(
                    x, R.shape([1, 8, 16, 128])
                )
                z: R.Tensor((1, 8, 16, 128), dtype="float16") = R.add(y, R.const(1, "float16"))
                R.output(z)
            return z

    assert relax.analysis.has_reshape_pattern(Module["reshape"])
    mod = relax.transform.RewriteDataflowReshape()(Module)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_reshape_non_dataflow():
    @tvm.script.ir_module
    class Module:
        @T.prim_func(s_tir=True)
        def reshape(
            rxplaceholder: T.Buffer((T.int64(8), T.int64(3)), "float32"),
            T_reshape: T.Buffer((T.int64(2), T.int64(4), T.int64(3)), "float32"),
        ):
            for ax0, ax1, ax2 in T.grid(T.int64(2), T.int64(4), T.int64(3)):
                with T.sblock("T_reshape"):
                    v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
                    T.reads(
                        rxplaceholder[
                            (v_ax0 * 12 + v_ax1 * 3 + v_ax2) // T.int64(3),
                            (v_ax0 * 12 + v_ax1 * 3 + v_ax2) % T.int64(3),
                        ]
                    )
                    T.writes(T_reshape[v_ax0, v_ax1, v_ax2])
                    T_reshape[v_ax0, v_ax1, v_ax2] = rxplaceholder[
                        (v_ax0 * 12 + v_ax1 * 3 + v_ax2) // T.int64(3),
                        (v_ax0 * 12 + v_ax1 * 3 + v_ax2) % T.int64(3),
                    ]

        @R.function
        def main(x: R.Tensor((8, 3), dtype="float32")) -> R.Tensor((2, 4, 3), dtype="float32"):
            cls = Module
            y = R.call_tir(cls.reshape, (x,), out_ty=R.Tensor((2, 4, 3), dtype="float32"))
            return y

    assert relax.analysis.has_reshape_pattern(Module["reshape"])
    # The binding var of the call_tir is not a DataflowVar. So the pass does no change.
    mod = relax.transform.RewriteDataflowReshape()(Module)
    tvm.ir.assert_structural_equal(mod, Module)


def test_tuple_get_reshape():
    @tvm.script.ir_module
    class Module:
        @T.prim_func(s_tir=True)
        def fused_reshape5(
            lv2_0: T.Buffer((T.int64(2), T.int64(4096), T.int64(320)), "float16"),
            lv2_1: T.Buffer((T.int64(2), T.int64(4096), T.int64(320)), "float16"),
            lv2_2: T.Buffer((T.int64(2), T.int64(4096), T.int64(320)), "float16"),
            T_reshape_handle_intermediate: T.Buffer(
                (T.int64(2), T.int64(4096), T.int64(8), T.int64(40)), "float16"
            ),
        ):
            T.func_attr({"tirx.noalias": True})
            # with T.sblock("root"):
            for ax0, ax1, ax2, ax3 in T.grid(T.int64(2), T.int64(4096), T.int64(8), T.int64(40)):
                with T.sblock("T_reshape"):
                    v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
                    T.reads(
                        lv2_0[
                            (
                                ((v_ax2 * T.int64(40) + v_ax3) // T.int64(320) + v_ax1)
                                // T.int64(4096)
                                + v_ax0
                            )
                            % T.int64(2),
                            ((v_ax2 * T.int64(40) + v_ax3) // T.int64(320) + v_ax1) % T.int64(4096),
                            (v_ax2 * T.int64(40) + v_ax3) % T.int64(320),
                        ]
                    )
                    T.writes(T_reshape_handle_intermediate[v_ax0, v_ax1, v_ax2, v_ax3])
                    T_reshape_handle_intermediate[v_ax0, v_ax1, v_ax2, v_ax3] = lv2_0[
                        (
                            ((v_ax2 * T.int64(40) + v_ax3) // T.int64(320) + v_ax1) // T.int64(4096)
                            + v_ax0
                        )
                        % T.int64(2),
                        ((v_ax2 * T.int64(40) + v_ax3) // T.int64(320) + v_ax1) % T.int64(4096),
                        (v_ax2 * T.int64(40) + v_ax3) % T.int64(320),
                    ]

        @R.function
        def main(
            lv41_1: R.Tuple(
                R.Tensor((2, 4096, 320), dtype="float16"),
                R.Tensor((2, 4096, 320), dtype="float16"),
                R.Tensor((2, 4096, 320), dtype="float16"),
            ),
        ) -> R.Tensor((2, 4096, 8, 40), dtype="float16"):
            cls = Module
            with R.dataflow():
                lv: R.Tensor((2, 4096, 320), dtype="float16") = lv41_1[0]
                lv1: R.Tensor((2, 4096, 320), dtype="float16") = lv41_1[1]
                lv2: R.Tensor((2, 4096, 320), dtype="float16") = lv41_1[2]
                lv645 = R.call_tir(
                    cls.fused_reshape5,
                    (lv, lv1, lv2),
                    out_ty=R.Tensor((2, 4096, 8, 40), dtype="float16"),
                )
                out: R.Tensor((2, 4096, 8, 40), dtype="float16") = R.add(lv645, lv645)
                R.output(out)
            return out

    @tvm.script.ir_module
    class Expected:
        @T.prim_func(s_tir=True)
        def fused_reshape5(
            lv2_0: T.Buffer((T.int64(2), T.int64(4096), T.int64(320)), "float16"),
            lv2_1: T.Buffer((T.int64(2), T.int64(4096), T.int64(320)), "float16"),
            lv2_2: T.Buffer((T.int64(2), T.int64(4096), T.int64(320)), "float16"),
            T_reshape_handle_intermediate: T.Buffer(
                (T.int64(2), T.int64(4096), T.int64(8), T.int64(40)), "float16"
            ),
        ):
            T.func_attr({"tirx.noalias": True})
            # with T.sblock("root"):
            for ax0, ax1, ax2, ax3 in T.grid(T.int64(2), T.int64(4096), T.int64(8), T.int64(40)):
                with T.sblock("T_reshape"):
                    v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
                    T.reads(
                        lv2_0[
                            (
                                ((v_ax2 * T.int64(40) + v_ax3) // T.int64(320) + v_ax1)
                                // T.int64(4096)
                                + v_ax0
                            )
                            % T.int64(2),
                            ((v_ax2 * T.int64(40) + v_ax3) // T.int64(320) + v_ax1) % T.int64(4096),
                            (v_ax2 * T.int64(40) + v_ax3) % T.int64(320),
                        ]
                    )
                    T.writes(T_reshape_handle_intermediate[v_ax0, v_ax1, v_ax2, v_ax3])
                    T_reshape_handle_intermediate[v_ax0, v_ax1, v_ax2, v_ax3] = lv2_0[
                        (
                            ((v_ax2 * T.int64(40) + v_ax3) // T.int64(320) + v_ax1) // T.int64(4096)
                            + v_ax0
                        )
                        % T.int64(2),
                        ((v_ax2 * T.int64(40) + v_ax3) // T.int64(320) + v_ax1) % T.int64(4096),
                        (v_ax2 * T.int64(40) + v_ax3) % T.int64(320),
                    ]

        @R.function
        def main(
            lv41_1: R.Tuple(
                R.Tensor((2, 4096, 320), dtype="float16"),
                R.Tensor((2, 4096, 320), dtype="float16"),
                R.Tensor((2, 4096, 320), dtype="float16"),
            ),
        ) -> R.Tensor((2, 4096, 8, 40), dtype="float16"):
            with R.dataflow():
                lv: R.Tensor((2, 4096, 320), dtype="float16") = lv41_1[0]
                lv1: R.Tensor((2, 4096, 320), dtype="float16") = lv41_1[1]
                lv2: R.Tensor((2, 4096, 320), dtype="float16") = lv41_1[2]
                lv645: R.Tensor((2, 4096, 8, 40), dtype="float16") = R.reshape(
                    lv, R.shape([2, 4096, 8, 40])
                )
                out: R.Tensor((2, 4096, 8, 40), dtype="float16") = R.add(lv645, lv645)
                R.output(out)
            return out

    rewritten = relax.transform.RewriteDataflowReshape()(Module)
    tvm.ir.assert_structural_equal(rewritten, Expected)


def test_invalid_reshape():
    @tvm.script.ir_module
    class Module:
        # The strided_slice op has the reshape pattern, but it can take only a part of the input.
        # It can't be replaced with the reshape op because reshape expects to preserve the "volume"
        # of the input.
        @T.prim_func(s_tir=True)
        def strided_slice(
            A: T.Buffer((T.int64(1), T.int64(1024)), "int32"),
            T_strided_slice: T.Buffer((T.int64(1), T.int64(1000)), "int32"),
        ):
            T.func_attr({"tirx.noalias": True})
            for ax0, ax1 in T.grid(T.int64(1), T.int64(1000)):
                with T.sblock("T_strided_slice"):
                    v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
                    T.reads(A[v_ax0, v_ax1])
                    T.writes(T_strided_slice[v_ax0, v_ax1])
                    T_strided_slice[v_ax0, v_ax1] = A[v_ax0, v_ax1]

        @T.prim_func(s_tir=True)
        def add_one(
            A: T.Buffer((T.int64(1), T.int64(1000)), "int32"),
            T_add_one: T.buffer((T.int64(1), T.int64(1000)), "int32"),
        ):
            for ax0, ax1 in T.grid(T.int64(1), T.int64(1000)):
                with T.sblock("T_add_one"):
                    v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
                    T.reads(A[v_ax0, v_ax1])
                    T.writes(T_add_one[v_ax0, v_ax1])
                    T_add_one[v_ax0, v_ax1] = A[v_ax0, v_ax1] + 1

        @R.function
        def main(A: R.Tensor((1, 1024), dtype="int32")) -> R.Tensor((1, 1000), dtype="int32"):
            with R.dataflow():
                cls = Module
                S = R.call_tir(cls.strided_slice, (A,), out_ty=R.Tensor((1, 1000), dtype="int32"))
                A = R.call_tir(cls.add_one, (S,), out_ty=R.Tensor((1, 1000), dtype="int32"))
                R.output(A)
            return A

    assert relax.analysis.has_reshape_pattern(Module["strided_slice"])
    rewritten = relax.transform.RewriteDataflowReshape()(Module)
    tvm.ir.assert_structural_equal(rewritten, Module)


def test_reshape_detect_nop():
    @tvm.script.ir_module
    class Module:
        @R.function
        def main(x: R.Tensor((8, 8), dtype="float16")) -> R.Tensor((8, 8), dtype="float16"):
            with R.dataflow():
                gv = R.call_pure_packed("foo", x, x, ty_args=(R.Tensor((8, 8), dtype="float16"),))
                out = R.call_pure_packed(
                    "foo", gv, gv, ty_args=(R.Tensor((8, 8), dtype="float16"),)
                )
                R.output(out)
            return out

    rewritten = relax.transform.RewriteDataflowReshape()(Module)
    tvm.ir.assert_structural_equal(rewritten, Module)


def test_reshape_scalar():
    @tvm.script.ir_module
    class Module:
        @R.function
        def main(x: R.Tensor((), dtype="float32")) -> R.Tensor((1,), dtype="float32"):
            with R.dataflow():
                lv1: R.Tensor((1,), dtype="float32") = R.reshape(x, [1])
                lv2: R.Tensor((1,), dtype="float32") = R.add(lv1, lv1)
                R.output(lv2)
            return lv2

    @tvm.script.ir_module
    class Expected:
        @T.prim_func(private=True, s_tir=True)
        def add(
            A: T.Buffer((T.int64(1),), "float32"),
            B: T.Buffer((T.int64(1),), "float32"),
            T_add: T.Buffer((T.int64(1),), "float32"),
        ):
            T.func_attr({"tirx.noalias": True})
            # with T.sblock("root"):
            for ax0 in range(T.int64(1)):
                with T.sblock("T_add"):
                    v_ax0 = T.axis.spatial(T.int64(1), ax0)
                    T.reads(A[v_ax0], B[v_ax0])
                    T.writes(T_add[v_ax0])
                    T_add[v_ax0] = A[v_ax0] + B[v_ax0]

        @T.prim_func(private=True, s_tir=True)
        def reshape(A: T.Buffer((), "float32"), T_reshape: T.Buffer((T.int64(1),), "float32")):
            T.func_attr({"tirx.noalias": True})
            # with T.sblock("root"):
            for ax0 in range(T.int64(1)):
                with T.sblock("T_reshape"):
                    v_ax0 = T.axis.spatial(T.int64(1), ax0)
                    T.reads(A[()])
                    T.writes(T_reshape[v_ax0])
                    T_reshape[v_ax0] = A[()]

        @R.function
        def main(x: R.Tensor((), dtype="float32")) -> R.Tensor((1,), dtype="float32"):
            cls = Expected
            with R.dataflow():
                lv1: R.Tensor((1,), dtype="float32") = R.reshape(x, R.shape([1]))
                lv2 = R.call_tir(cls.add, (lv1, lv1), out_ty=R.Tensor((1,), dtype="float32"))
                R.output(lv2)
            return lv2

    mod = Module
    mod = relax.transform.LegalizeOps()(mod)
    rewritten = relax.transform.RewriteDataflowReshape()(mod)
    tvm.ir.assert_structural_equal(rewritten, Expected)


def test_rewrite_static_reshape():
    @I.ir_module(s_tir=True)
    class Before:
        @R.function
        def main(x: R.Tensor([256], dtype="float32")):
            with R.dataflow():
                y = R.reshape(x, [64, 4])
                z = R.add(y, y)
                R.output(z)
            return z

    @I.ir_module(s_tir=True)
    class Expected:
        @R.function
        def main(x: R.Tensor((256,), dtype="float32")):
            cls = Expected

            with R.dataflow():
                y = R.reshape(x, R.shape([64, 4]))
                z = R.call_tir(cls.add, (y, y), out_ty=R.Tensor((64, 4), dtype="float32"))
                R.output(z)
            return z

        @T.prim_func(private=True, s_tir=True)
        def add(
            y1: T.Buffer((T.int64(64), T.int64(4)), "float32"),
            y2: T.Buffer((T.int64(64), T.int64(4)), "float32"),
            z: T.Buffer((T.int64(64), T.int64(4)), "float32"),
        ):
            T.func_attr({"tirx.noalias": True})

            for iters in T.grid(T.int64(64), T.int64(4)):
                with T.sblock("T_add"):
                    i, j = T.axis.remap("SS", iters)
                    z[i, j] = y1[i, j] + y2[i, j]

    After = tvm.ir.transform.Sequential(
        [
            # Lower both R.reshape and R.add from Relax to TIR
            relax.transform.LegalizeOps(),
            # Identify reshapes, raise calls to cls.reshape from TIR
            # to Relax
            relax.transform.RewriteDataflowReshape(),
            # Clean up afterwards, removing the no-longer-required
            # PrimFunc "reshape"
            relax.transform.DeadCodeElimination(),
        ]
    )(Before)

    tvm.ir.assert_structural_equal(Expected, After)


# def test_rewrite_dynamic_reshape():
#     @I.ir_module
#     class Before:
#         @R.function
#         def main(x: R.Tensor(["N"], dtype="float32")):
#             N = T.int64()
#             with R.dataflow():
#                 y = R.reshape(x, [N // 4, 4])
#                 z = R.add(y, y)
#                 R.output(z)
#             return z

#     @I.ir_module
#     class Expected:
#         @R.function
#         def main(x: R.Tensor(["N"], dtype="float32")):
#             N = T.int64()
#             cls = Expected

#             with R.dataflow():
#                 y = R.reshape(x, R.shape([N // 4, 4]))
#                 z = R.call_tir(
#                     cls.add,
#                     (y, y),
#                     tir_vars=[N],
#                     out_ty=R.Tensor((N // 4, 4), dtype="float32"),
#                 )
#                 R.output(z)
#             return z

#         @T.prim_func(private=True)
#         def add(
#             y1_handle: T.handle,
#             y2_handle: T.handle,
#             z_handle: T.handle,
#             N: T.int64,
#         ):

#             y1 = T.match_buffer(y1_handle, [N // 4, 4], "float32")
#             y2 = T.match_buffer(y2_handle, [N // 4, 4], "float32")
#             z = T.match_buffer(z_handle, [N // 4, 4], "float32")

#             T.func_attr({"tirx.noalias": True})

#             for iters in T.grid(T.int64(64), T.int64(4)):
#                 with T.sblock("T_add"):
#                     i, j = T.axis.remap("SS", iters)
#                     z[i, j] = y1[i, j] + y2[i, j]

#     After = tvm.ir.transform.Sequential(
#         [
#             # Lower both R.reshape and R.add from Relax to TIR
#             relax.transform.LegalizeOps(),
#             # Identify reshapes, raise calls to cls.reshape from TIR
#             # to Relax
#             relax.transform.RewriteDataflowReshape(),
#             # Clean up afterwards, removing the no-longer-required
#             # PrimFunc "reshape"
#             relax.transform.DeadCodeElimination(),
#         ]
#     )(Before)
#     After.show()
#     tvm.ir.assert_structural_equal(Expected, After)


def test_rewrite_dynamic_reshape():
    @I.ir_module(s_tir=True)
    class Before:
        @R.function
        def main(x: R.Tensor(["N", 16], dtype="float32")):
            N = T.int64()
            with R.dataflow():
                y = R.reshape(x, [N * 4, T.int64(4)])
                z = R.add(y, y)
                R.output(z)
            return z

    @I.ir_module(s_tir=True)
    class Expected:
        @R.function
        def main(x: R.Tensor(["N", 16], dtype="float32")):
            N = T.int64()
            cls = Expected

            with R.dataflow():
                y = R.reshape(x, R.shape([N * 4, T.int64(4)]))
                z = R.call_tir(
                    cls.add,
                    (y, y),
                    tir_vars=[N],
                    out_ty=R.Tensor((N * 4, 4), dtype="float32"),
                )
                R.output(z)
            return z

        @T.prim_func(private=True, s_tir=True)
        def add(
            y1_handle: T.handle,
            y2_handle: T.handle,
            z_handle: T.handle,
            N: T.int64,
        ):
            y1 = T.match_buffer(y1_handle, [N * 4, T.int64(4)], "float32")
            y2 = T.match_buffer(y2_handle, [N * 4, T.int64(4)], "float32")
            z = T.match_buffer(z_handle, [N * 4, T.int64(4)], "float32")

            T.func_attr({"tirx.noalias": True})

            for iters in T.grid(N * 4, T.int64(4)):
                with T.sblock("T_add"):
                    i, j = T.axis.remap("SS", iters)
                    z[i, j] = y1[i, j] + y2[i, j]

    After = tvm.ir.transform.Sequential(
        [
            # Lower both R.reshape and R.add from Relax to TIR
            relax.transform.LegalizeOps(),
            # Identify reshapes, raise calls to cls.reshape from TIR
            # to Relax
            relax.transform.RewriteDataflowReshape(),
            # Clean up afterwards, removing the no-longer-required
            # PrimFunc "reshape"
            relax.transform.DeadCodeElimination(),
        ]
    )(Before)
    tvm.ir.assert_structural_equal(Expected, After)


if __name__ == "__main__":
    tvm.testing.main()