# 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, E731, F841
import pytest

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

##################### Manipulation #####################


def test_broadcast_to():
    # fmt: off
    @tvm.script.ir_module
    class BroadcastTo:
        @R.function
        def main(x: R.Tensor((2, 1, 3), "float32")) -> R.Tensor((4, 2, 5, 3), "float32"):
            gv: R.Tensor((4, 2, 5, 3), "float32") = R.broadcast_to(x, (4, 2, 5, 3))
            return gv

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor((2, 1, 3), "float32")) -> R.Tensor((4, 2, 5, 3), "float32"):
            gv = R.call_tir(Expected.broadcast_to, (x,), R.Tensor((4, 2, 5, 3), dtype="float32"))
            return gv

        @T.prim_func(private=True, s_tir=True)
        def broadcast_to(rxplaceholder: T.Buffer((T.int64(2), T.int64(1), T.int64(3)), "float32"), T_broadcast_to: T.Buffer((T.int64(4), T.int64(2), T.int64(5), T.int64(3)), "float32")):
            T.func_attr({"tirx.noalias": True})
            for i0, i1, i2, i3 in T.grid(T.int64(4), T.int64(2), T.int64(5), T.int64(3)):
                with T.sblock("T_broadcast_to"):
                    ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
                    T.reads(rxplaceholder[ax1, T.int64(0), ax3])
                    T.writes(T_broadcast_to[ax0, ax1, ax2, ax3])
                    T_broadcast_to[ax0, ax1, ax2, ax3] = rxplaceholder[ax1, T.int64(0), ax3]
    # fmt: on

    mod = LegalizeOps()(BroadcastTo)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_broadcast_to_symbolic():
    # fmt: off
    @tvm.script.ir_module
    class BroadcastTo:
        @R.function
        def main(dumb_param: R.Tensor(("a", "c")), x: R.Tensor(("b", 1, "d"), "float32")) -> R.Tensor(("a", "b", "c", "d"), "float32"):
            a = T.int64()
            b = T.int64()
            c = T.int64()
            d = T.int64()
            gv: R.Tensor((a, b, c, d), "float32") = R.broadcast_to(x, (a, b, c, d))
            return gv

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(dumb_param: R.Tensor(("a", "c")), x: R.Tensor(("b", 1, "d"), "float32")) -> R.Tensor(("a", "b", "c", "d"), "float32"):
            a = T.int64()
            b = T.int64()
            c = T.int64()
            d = T.int64()
            gv = R.call_tir(Expected.broadcast_to, (x,), R.Tensor((a, b, c, d), dtype="float32"))
            return gv

        @T.prim_func(private=True, s_tir=True)
        def broadcast_to(var_rxplaceholder: T.handle, var_T_broadcast_to: T.handle):
            T.func_attr({"tirx.noalias": True})
            a = T.int64()
            b = T.int64()
            c = T.int64()
            d = T.int64()
            rxplaceholder = T.match_buffer(var_rxplaceholder, [b, T.int64(1), d], dtype="float32")
            T_broadcast_to = T.match_buffer(var_T_broadcast_to, [a, b, c, d], dtype="float32")
            for i0, i1, i2, i3 in T.grid(a, b, c, d):
                with T.sblock("T_broadcast_to"):
                    ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
                    T.reads(rxplaceholder[ax1, T.int64(0), ax3])
                    T.writes(T_broadcast_to[ax0, ax1, ax2, ax3])
                    T_broadcast_to[ax0, ax1, ax2, ax3] = rxplaceholder[ax1, T.int64(0), ax3]
    # fmt: on

    mod = LegalizeOps()(BroadcastTo)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_concat():
    # fmt: off
    @tvm.script.ir_module
    class Concat:
        @R.function
        def main(x1: R.Tensor((1, 2, 3), "float32"), x2: R.Tensor((1, 3, 3), "float32"), x3: R.Tensor((1, 4, 3), "float32")) -> R.Tensor((1, 9, 3), "float32"):
            gv: R.Tensor((1, 9, 3), "float32") = R.concat((x1, x2, x3), axis=1)
            return gv

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(x1: R.Tensor((1, 2, 3), "float32"), x2: R.Tensor((1, 3, 3), "float32"), x3: R.Tensor((1, 4, 3), "float32")) -> R.Tensor((1, 9, 3), "float32"):
            gv = R.call_tir(Expected.concatenate, (x1, x2, x3), R.Tensor((1, 9, 3), dtype="float32"))
            return gv

        @T.prim_func(private=True, s_tir=True)
        def concatenate(rxplaceholder: T.Buffer((T.int64(1), T.int64(2), T.int64(3)), "float32"), rxplaceholder_1: T.Buffer((T.int64(1), T.int64(3), T.int64(3)), "float32"), rxplaceholder_2: T.Buffer((T.int64(1), T.int64(4), T.int64(3)), "float32"), T_concat: T.Buffer((T.int64(1), T.int64(9), T.int64(3)), "float32")):
            T.func_attr({"tirx.noalias": True})
            for i0, i1, i2 in T.grid(T.int64(1), T.int64(9), T.int64(3)):
                with T.sblock("T_concat"):
                    ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
                    T.reads(rxplaceholder_2[ax0, ax1 - T.int64(5), ax2], rxplaceholder_1[ax0, ax1 - T.int64(2), ax2], rxplaceholder[ax0, ax1, ax2])
                    T.writes(T_concat[ax0, ax1, ax2])
                    T_concat[ax0, ax1, ax2] = T.if_then_else(T.int64(5) <= ax1, rxplaceholder_2[ax0, ax1 - T.int64(5), ax2], T.if_then_else(T.int64(2) <= ax1, rxplaceholder_1[ax0, ax1 - T.int64(2), ax2], rxplaceholder[ax0, ax1, ax2]))
    # fmt: on

    mod = LegalizeOps()(Concat)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_concat_input_tuple_var():
    # fmt: off
    @tvm.script.ir_module
    class Concat:
        @R.function
        def main(t: R.Tuple(R.Tensor((3, 4), "float32"), R.Tensor((3, 5), "float32"))) -> R.Tensor((3, 9), "float32"):
            gv: R.Tensor((3, 9), "float32") = R.concat(t, axis=1)
            return gv

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(t: R.Tuple(R.Tensor((3, 4), "float32"), R.Tensor((3, 5), "float32"))) -> R.Tensor((3, 9), "float32"):
            gv: R.Tensor((3, 4), dtype="float32") = t[0]
            gv1: R.Tensor((3, 5), dtype="float32") = t[1]
            gv2 = R.call_tir(Expected.concatenate, (gv, gv1), R.Tensor((3, 9), dtype="float32"))
            return gv2

        @T.prim_func(private=True, s_tir=True)
        def concatenate(rxplaceholder: T.Buffer((T.int64(3), T.int64(4)), "float32"), rxplaceholder_1: T.Buffer((T.int64(3), T.int64(5)), "float32"), T_concat: T.Buffer((T.int64(3), T.int64(9)), "float32")):
            T.func_attr({"tirx.noalias": True})
            for i0, i1 in T.grid(T.int64(3), T.int64(9)):
                with T.sblock("T_concat"):
                    ax0, ax1 = T.axis.remap("SS", [i0, i1])
                    T.reads(rxplaceholder_1[ax0, ax1 - T.int64(4)], rxplaceholder[ax0, ax1])
                    T.writes(T_concat[ax0, ax1])
                    T_concat[ax0, ax1] = T.if_then_else(T.int64(4) <= ax1, rxplaceholder_1[ax0, ax1 - T.int64(4)], rxplaceholder[ax0, ax1])
    # fmt: on

    mod = LegalizeOps()(Concat)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_concat_input_tuple_var_symbolic():
    # fmt: off
    @tvm.script.ir_module
    class Concat:
        @R.function
        def main(t: R.Tuple(R.Tensor(("a", "b0"), "float32"), R.Tensor(("a", "b1"), "float32"), R.Tensor(("a", "b2"), "float32"))) -> R.Tensor(("a", "b0 + b1 + b2"), "float32"):
            a = T.int64()
            b0 = T.int64()
            b1 = T.int64()
            b2 = T.int64()
            gv: R.Tensor((a, b0 + b1 + b2), "float32") = R.concat(t, axis=1)
            return gv

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(t: R.Tuple(R.Tensor(("a", "b0"), "float32"), R.Tensor(("a", "b1"), "float32"), R.Tensor(("a", "b2"), "float32"))) -> R.Tensor(("a", "b0 + b1 + b2"), "float32"):
            a = T.int64()
            b0 = T.int64()
            b1 = T.int64()
            b2 = T.int64()
            gv: R.Tensor((a, b0), dtype="float32") = t[0]
            gv1: R.Tensor((a, b1), dtype="float32") = t[1]
            gv2: R.Tensor((a, b2), dtype="float32") = t[2]
            gv3 = R.call_tir(Expected.concatenate, (gv, gv1, gv2), R.Tensor((a, ((b0 + b1) + b2)), dtype="float32"))
            return gv3

        @T.prim_func(private=True, s_tir=True)
        def concatenate(var_rxplaceholder: T.handle, var_rxplaceholder_1: T.handle, var_rxplaceholder_2: T.handle, var_T_concat: T.handle):
            T.func_attr({"tirx.noalias": True})
            a = T.int64()
            b0 = T.int64()
            b1 = T.int64()
            b2 = T.int64()
            rxplaceholder = T.match_buffer(var_rxplaceholder, [a, b0], dtype="float32")
            rxplaceholder_1 = T.match_buffer(var_rxplaceholder_1, [a, b1], dtype="float32")
            rxplaceholder_2 = T.match_buffer(var_rxplaceholder_2, [a, b2], dtype="float32")
            T_concat = T.match_buffer(var_T_concat, [a, b0 + b1 + b2], dtype="float32")
            for i0, i1 in T.grid(a, b0 + b1 + b2):
                with T.sblock("T_concat"):
                    ax0, ax1 = T.axis.remap("SS", [i0, i1])
                    T.reads(rxplaceholder_2[ax0, ax1 - b0 - b1], rxplaceholder_1[ax0, ax1 - b0], rxplaceholder[ax0, ax1])
                    T.writes(T_concat[ax0, ax1])
                    T_concat[ax0, ax1] = T.if_then_else(T.int64(0) <= ax1 - b0 - b1, rxplaceholder_2[ax0, ax1 - b0 - b1], T.if_then_else(T.int64(0) <= ax1 - b0, rxplaceholder_1[ax0, ax1 - b0], rxplaceholder[ax0, ax1]))
    # fmt: on

    mod = LegalizeOps()(Concat)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_expand_dims():
    # fmt: off
    @tvm.script.ir_module
    class ExpandDims:
        @R.function
        def main(x: R.Tensor((2, 3, 4), "float32")) -> R.Tensor((2, 1, 1, 1, 3, 1, 4, 1), "float32"):
            gv: R.Tensor((2, 1, 1, 1, 3, 1, 4, 1), "float32") = R.expand_dims(x, axis=[-1, 1, -6, 3, 5])
            return gv

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor((2, 3, 4), "float32")) -> R.Tensor((2, 1, 1, 1, 3, 1, 4, 1), "float32"):
            gv = R.call_tir(Expected.expand_dims, (x,), R.Tensor((2, 1, 1, 1, 3, 1, 4, 1), dtype="float32"))
            return gv

        @T.prim_func(private=True, s_tir=True)
        def expand_dims(rxplaceholder: T.Buffer((T.int64(2), T.int64(3), T.int64(4)), "float32"), expand_dims: T.Buffer((T.int64(2), T.int64(1), T.int64(1), T.int64(1), T.int64(3), T.int64(1), T.int64(4), T.int64(1)), "float32")):
            T.func_attr({"tirx.noalias": True})
            for i0, i1, i2, i3, i4, i5, i6, i7 in T.grid(T.int64(2), T.int64(1), T.int64(1), T.int64(1), T.int64(3), T.int64(1), T.int64(4), T.int64(1)):
                with T.sblock("expand_dims"):
                    i0_1, i1_1, i2_1, i3_1, i4_1, i5_1, i6_1, i7_1 = T.axis.remap("SSSSSSSS", [i0, i1, i2, i3, i4, i5, i6, i7])
                    T.reads(rxplaceholder[i0_1, i4_1, i6_1])
                    T.writes(expand_dims[i0_1, i1_1, i2_1, i3_1, i4_1, i5_1, i6_1, i7_1])
                    expand_dims[i0_1, i1_1, i2_1, i3_1, i4_1, i5_1, i6_1, i7_1] = rxplaceholder[i0_1, i4_1, i6_1]
    # fmt: on

    mod = LegalizeOps()(ExpandDims)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_expand_dims_symbolic():
    # fmt: off
    @tvm.script.ir_module
    class ExpandDims:
        @R.function
        def main(x: R.Tensor(("a", "b", "c"), "float32")) -> R.Tensor(("a", 1, "b", 1, "c", 1), "float32"):
            a = T.int64()
            b = T.int64()
            c = T.int64()
            gv: R.Tensor((a, 1, b, 1, c, 1), "float32") = R.expand_dims(x, axis=[1, 3, 5])
            return gv

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor(("a", "b", "c"), "float32")) -> R.Tensor(("a", 1, "b", 1, "c", 1), "float32"):
            a = T.int64()
            b = T.int64()
            c = T.int64()
            gv = R.call_tir(Expected.expand_dims, (x,), R.Tensor((a, 1, b, 1, c, 1), dtype="float32"))
            return gv

        @T.prim_func(private=True, s_tir=True)
        def expand_dims(var_rxplaceholder: T.handle, var_expand_dims: T.handle):
            T.func_attr({"tirx.noalias": True})
            a = T.int64()
            b = T.int64()
            c = T.int64()
            rxplaceholder = T.match_buffer(var_rxplaceholder, [a, b, c], dtype="float32")
            expand_dims = T.match_buffer(var_expand_dims, [a, T.int64(1), b, T.int64(1), c, T.int64(1)], dtype="float32")
            for i0, i1, i2, i3, i4, i5 in T.grid(a, T.int64(1), b, T.int64(1), c, T.int64(1)):
                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])
                    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]
    # fmt: on

    mod = LegalizeOps()(ExpandDims)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_flatten():
    # fmt: off
    @tvm.script.ir_module
    class Flatten:
        @R.function
        def main(x: R.Tensor((2, 3, 4), "float32")) -> R.Tensor((24,), "float32"):
            gv: R.Tensor((24,), "float32") = R.flatten(x)
            return gv

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor((2, 3, 4), "float32")) -> R.Tensor((24,), "float32"):
            gv = R.call_tir(Expected.reshape, (x,), R.Tensor((24,), dtype="float32"))
            return gv

        @T.prim_func(private=True, s_tir=True)
        def reshape(rxplaceholder: T.Buffer((T.int64(2), T.int64(3), T.int64(4)), "float32"), T_reshape: T.Buffer(T.int64(24), "float32")):
            T.func_attr({"tirx.noalias": True})
            for i0 in T.serial(T.int64(24)):
                with T.sblock("T_reshape"):
                    ax0 = T.axis.spatial(T.int64(24), i0)
                    T.reads(rxplaceholder[ax0 % T.int64(24) // T.int64(12), ax0 % T.int64(12) // T.int64(4), ax0 % T.int64(4)])
                    T.writes(T_reshape[ax0])
                    T_reshape[ax0] = rxplaceholder[ax0 % T.int64(24) // T.int64(12), ax0 % T.int64(12) // T.int64(4), ax0 % T.int64(4)]
    # fmt: on

    mod = LegalizeOps()(Flatten)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_flatten_zero_rank():
    # fmt: off
    @tvm.script.ir_module
    class Flatten:
        @R.function
        def main(x: R.Tensor((), "float32")) -> R.Tensor((1,), "float32"):
            gv: R.Tensor((1,), "float32") = R.flatten(x)
            return gv

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor((), "float32")) -> R.Tensor((1,), "float32"):
            gv = R.call_tir(Expected.reshape, (x,), R.Tensor((1,), dtype="float32"))
            return gv

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

    mod = LegalizeOps()(Flatten)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_flatten_symbolic():
    # fmt: off
    @tvm.script.ir_module
    class Flatten:
        @R.function
        def main(x: R.Tensor(("a", "b", "c"), "float32")) -> R.Tensor(("a * b * c",), "float32"):
            a = T.int64()
            b = T.int64()
            c = T.int64()
            gv: R.Tensor((a * b * c,), "float32") = R.flatten(x)
            return gv

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor(("a", "b", "c"), "float32")) -> R.Tensor(("a * b * c",), "float32"):
            a = T.int64()
            b = T.int64()
            c = T.int64()
            gv = R.call_tir(Expected.reshape, (x,), R.Tensor((((a * b) * c),), dtype="float32"))
            return gv

        @T.prim_func(private=True, s_tir=True)
        def reshape(var_rxplaceholder: T.handle, var_T_reshape: T.handle):
            T.func_attr({"tirx.noalias": True})
            a = T.int64()
            b = T.int64()
            c = T.int64()
            rxplaceholder = T.match_buffer(var_rxplaceholder, [a, b, c], dtype="float32")
            T_reshape = T.match_buffer(var_T_reshape, [a * b * c], dtype="float32")
            for i0 in T.serial(a * b * c):
                with T.sblock("T_reshape"):
                    ax0 = T.axis.spatial(a * b * c, i0)
                    T.reads(rxplaceholder[ax0 // c // b % a, ax0 // c % b, ax0 % c])
                    T.writes(T_reshape[ax0])
                    T_reshape[ax0] = rxplaceholder[ax0 // c // b % a, ax0 // c % b, ax0 % c]
    # fmt: on

    mod = LegalizeOps()(Flatten)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_permute_dims():
    # fmt: off
    @tvm.script.ir_module
    class PermuteDims:
        @R.function
        def main(x: R.Tensor((1, 2, 3, 4), "float32")) -> R.Tensor((2, 4, 3, 1), "float32"):
            gv: R.Tensor((2, 4, 3, 1), "float32") = R.permute_dims(x, axes=[1, -1, 2, -4])
            return gv

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor((1, 2, 3, 4), "float32")) -> R.Tensor((2, 4, 3, 1), "float32"):
            gv = R.call_tir(Expected.transpose, (x,), R.Tensor((2, 4, 3, 1), dtype="float32"))
            return gv

        @T.prim_func(private=True, s_tir=True)
        def transpose(rxplaceholder: T.Buffer((T.int64(1), T.int64(2), T.int64(3), T.int64(4)), "float32"), T_transpose: T.Buffer((T.int64(2), T.int64(4), T.int64(3), T.int64(1)), "float32")):
            T.func_attr({"tirx.noalias": True})
            for i0, i1, i2, i3 in T.grid(T.int64(2), T.int64(4), T.int64(3), T.int64(1)):
                with T.sblock("T_transpose"):
                    ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
                    T.reads(rxplaceholder[ax3, ax0, ax2, ax1])
                    T.writes(T_transpose[ax0, ax1, ax2, ax3])
                    T_transpose[ax0, ax1, ax2, ax3] = rxplaceholder[ax3, ax0, ax2, ax1]
    # fmt: on

    mod = LegalizeOps()(PermuteDims)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_permute_dims_symbolic():
    # fmt: off
    @tvm.script.ir_module
    class PermuteDims:
        @R.function
        def main(x: R.Tensor(("a", "b", "c", "d"), "float32")) -> R.Tensor(("b", "d", "c", "a"), "float32"):
            a = T.int64()
            b = T.int64()
            c = T.int64()
            d = T.int64()
            gv: R.Tensor((b, d, c, a), "float32") = R.permute_dims(x, axes=[1, -1, 2, -4])
            return gv

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor(("a", "b", "c", "d"), dtype="float32")) -> R.Tensor(("b", "d", "c", "a"), dtype="float32"):
            b = T.int64()
            d = T.int64()
            c = T.int64()
            a = T.int64()
            gv = R.call_tir(Expected.transpose, (x,), R.Tensor((b, d, c, a), dtype="float32"))
            return gv

        @T.prim_func(private=True, s_tir=True)
        def transpose(var_rxplaceholder: T.handle, var_T_transpose: T.handle):
            T.func_attr({"tirx.noalias": True})
            a = T.int64()
            b = T.int64()
            c = T.int64()
            d = T.int64()
            rxplaceholder = T.match_buffer(var_rxplaceholder, [a, b, c, d], dtype="float32")
            T_transpose = T.match_buffer(var_T_transpose, [b, d, c, a], dtype="float32")
            for i0, i1, i2, i3 in T.grid(b, d, c, a):
                with T.sblock("T_transpose"):
                    ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
                    T.reads(rxplaceholder[ax3, ax0, ax2, ax1])
                    T.writes(T_transpose[ax0, ax1, ax2, ax3])
                    T_transpose[ax0, ax1, ax2, ax3] = rxplaceholder[ax3, ax0, ax2, ax1]
    # fmt: on

    mod = LegalizeOps()(PermuteDims)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_reshape():
    # fmt: off
    @tvm.script.ir_module
    class Reshape:
        @R.function
        def main(x: R.Tensor((1, 2, 3, 4), "float32")) -> R.Tensor((8, 3), "float32"):
            gv: R.Tensor((8, 3), "float32") = R.reshape(x, (8, 3))
            return gv

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor((1, 2, 3, 4), "float32")) -> R.Tensor((8, 3), "float32"):
            gv = R.call_tir(Expected.reshape, (x,), R.Tensor((8, 3), dtype="float32"))
            return gv

        @T.prim_func(private=True, s_tir=True)
        def reshape(rxplaceholder: T.Buffer((T.int64(1), T.int64(2), T.int64(3), T.int64(4)), "float32"), T_reshape: T.Buffer((T.int64(8), T.int64(3)), "float32")):
            T.func_attr({"tirx.noalias": True})
            for i0, i1 in T.grid(T.int64(8), T.int64(3)):
                with T.sblock("T_reshape"):
                    ax0, ax1 = T.axis.remap("SS", [i0, i1])
                    T.reads(rxplaceholder[T.int64(0), (ax0 * T.int64(3) + ax1) % T.int64(24) // T.int64(12), (ax0 * T.int64(3) + ax1) % T.int64(12) // T.int64(4), (ax0 * T.int64(3) + ax1) % T.int64(4)])
                    T.writes(T_reshape[ax0, ax1])
                    T_reshape[ax0, ax1] = rxplaceholder[T.int64(0), (ax0 * T.int64(3) + ax1) % T.int64(24) // T.int64(12), (ax0 * T.int64(3) + ax1) % T.int64(12) // T.int64(4), (ax0 * T.int64(3) + ax1) % T.int64(4)]
    # fmt: on

    mod = LegalizeOps()(Reshape)
    tvm.ir.assert_structural_equal(mod, Expected)

    # fmt: off
    # ShapeExpr might be produced by shape computation
    @tvm.script.ir_module
    class Reshape2:
        @R.function
        def main(x: R.Tensor((1, 2, 3, 4), "float32")) -> R.Tensor((8, 3), "float32"):
            lv: R.Shape((8, 3)) = R.shape((8, 3))
            gv: R.Tensor((8, 3), "float32") = R.reshape(x, lv)
            return gv

    # After lowering, redundant var might be removed by later dead code elimination
    @tvm.script.ir_module
    class Expected2:
        @T.prim_func(private=True, s_tir=True)
        def reshape(
            rxplaceholder: T.Buffer((T.int64(1), T.int64(2), T.int64(3), T.int64(4)), "float32"),
            T_reshape: T.Buffer((T.int64(8), T.int64(3)), "float32"),
        ):
            T.func_attr({"tirx.noalias": True})
            # with T.sblock("root"):
            for ax0, ax1 in T.grid(T.int64(8), T.int64(3)):
                with T.sblock("T_reshape"):
                    v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
                    T.reads(
                        rxplaceholder[
                            T.int64(0),
                            (v_ax0 * T.int64(3) + v_ax1) % T.int64(24) // T.int64(12),
                            (v_ax0 * T.int64(3) + v_ax1) % T.int64(12) // T.int64(4),
                            (v_ax0 * T.int64(3) + v_ax1) % T.int64(4),
                        ]
                    )
                    T.writes(T_reshape[v_ax0, v_ax1])
                    T_reshape[v_ax0, v_ax1] = rxplaceholder[
                        T.int64(0),
                        (v_ax0 * T.int64(3) + v_ax1) % T.int64(24) // T.int64(12),
                        (v_ax0 * T.int64(3) + v_ax1) % T.int64(12) // T.int64(4),
                        (v_ax0 * T.int64(3) + v_ax1) % T.int64(4),
                    ]

        @R.function
        def main(x: R.Tensor((1, 2, 3, 4), dtype="float32")) -> R.Tensor((8, 3), dtype="float32"):
            lv: R.Shape((8, 3)) = R.shape((8, 3))
            gv = R.call_tir(Expected2.reshape, (x,), out_ty=R.Tensor((8, 3), dtype="float32"))
            return gv
    # fmt: on

    mod2 = LegalizeOps()(Reshape2)
    tvm.ir.assert_structural_equal(mod2, Expected2)


def test_reshape_symbolic():
    # fmt: off
    @tvm.script.ir_module
    class Reshape:
        @R.function
        def main(x: R.Tensor(("a", "b"), "float32")) -> R.Tensor(("a // 2", "b * 2"), "float32"):
            a = T.int64()
            b = T.int64()
            gv: R.Tensor((a // 2, b * 2), "float32") = R.reshape(x, (a // 2, b * 2))
            return gv

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor(("a", "b"), "float32")) -> R.Tensor(("a // 2", "b * 2"), "float32"):
            a = T.int64()
            b = T.int64()
            gv = R.call_tir(Expected.reshape, (x,), R.Tensor(((a // 2), (b * 2)), dtype="float32"))
            return gv

        @T.prim_func(private=True, s_tir=True)
        def reshape(var_rxplaceholder: T.handle, var_T_reshape: T.handle):
            T.func_attr({"tirx.noalias": True})
            a = T.int64()
            b = T.int64()
            rxplaceholder = T.match_buffer(var_rxplaceholder, [a, b], dtype="float32")
            T_reshape = T.match_buffer(var_T_reshape, [a // T.int64(2), b * T.int64(2)], dtype="float32")
            for i0, i1 in T.grid(a // T.int64(2), b * T.int64(2)):
                with T.sblock("T_reshape"):
                    ax0, ax1 = T.axis.remap("SS", [i0, i1])
                    T.reads(rxplaceholder[(ax0 * b * T.int64(2) + ax1) // b % a, (ax0 * b * T.int64(2) + ax1) % b])
                    T.writes(T_reshape[ax0, ax1])
                    T_reshape[ax0, ax1] = rxplaceholder[(ax0 * b * T.int64(2) + ax1) // b % a, (ax0 * b * T.int64(2) + ax1) % b]
    # fmt: on

    mod = LegalizeOps()(Reshape)
    tvm.ir.assert_structural_equal(mod, Expected)

    # ShapeExpr might be produced by shape computation
    @tvm.script.ir_module
    class Reshape2:
        @R.function
        def main(x: R.Tensor(("a", "b"), "float32")) -> R.Tensor(("a // 2", "b * 2"), "float32"):
            a = T.int64()
            b = T.int64()
            lv: R.Shape((a // 2, b * 2)) = R.shape((a // 2, b * 2))
            gv: R.Tensor((a // 2, b * 2), "float32") = R.reshape(x, lv)
            return gv

    # After lowering, redundant var might be removed by later dead code elimination
    @tvm.script.ir_module
    class Expected2:
        @R.function
        def main(x: R.Tensor(("a", "b"), "float32")) -> R.Tensor(("a // 2", "b * 2"), "float32"):
            a = T.int64()
            b = T.int64()
            lv: R.Shape((a // 2, b * 2)) = R.shape((a // 2, b * 2))
            gv = R.call_tir(Expected2.reshape, (x,), R.Tensor(((a // 2), (b * 2)), dtype="float32"))
            return gv

        @T.prim_func(private=True, s_tir=True)
        def reshape(var_rxplaceholder: T.handle, var_T_reshape: T.handle):
            T.func_attr({"tirx.noalias": True})
            a = T.int64()
            b = T.int64()
            rxplaceholder = T.match_buffer(var_rxplaceholder, [a, b], dtype="float32")
            T_reshape = T.match_buffer(
                var_T_reshape, [a // T.int64(2), b * T.int64(2)], dtype="float32"
            )
            for i0, i1 in T.grid(a // T.int64(2), b * T.int64(2)):
                with T.sblock("T_reshape"):
                    ax0, ax1 = T.axis.remap("SS", [i0, i1])
                    T.reads(
                        rxplaceholder[
                            (ax0 * b * T.int64(2) + ax1) // b % a,
                            (ax0 * b * T.int64(2) + ax1) % b,
                        ]
                    )
                    T.writes(T_reshape[ax0, ax1])
                    T_reshape[ax0, ax1] = rxplaceholder[
                        (ax0 * b * T.int64(2) + ax1) // b % a, (ax0 * b * T.int64(2) + ax1) % b
                    ]

    mod2 = LegalizeOps()(Reshape2)
    tvm.ir.assert_structural_equal(mod2, Expected2)

    # ShapeExpr might be produced by shape computation
    @I.ir_module(s_tir=True)
    class Reshape3:
        @R.function
        def main(x: R.Tensor((10, "b"), "float32")) -> R.Tensor((5, "b * 2"), "float32"):
            a = T.int64()
            b = T.int64()
            lv: R.Shape((5, b * 2)) = R.shape((5, b * 2))
            gv: R.Tensor((5, b * 2), "float32") = R.reshape(x, lv)
            return gv

    # After lowering, redundant var might be removed by later dead code elimination
    @I.ir_module(s_tir=True)
    class Expected3:
        @T.prim_func(private=True, s_tir=True)
        def reshape(var_rxplaceholder: T.handle, var_T_reshape: T.handle):
            T.func_attr({"tirx.noalias": True})
            b = T.int64()
            rxplaceholder = T.match_buffer(var_rxplaceholder, (T.int64(10), b))
            T_reshape = T.match_buffer(var_T_reshape, (T.int64(5), b * T.int64(2)))
            # with T.sblock("root"):
            for ax0, ax1 in T.grid(T.int64(5), b * T.int64(2)):
                with T.sblock("T_reshape"):
                    v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
                    T.reads(
                        rxplaceholder[
                            (v_ax0 * b * T.int64(2) + v_ax1) // b % T.int64(10),
                            (v_ax0 * b * T.int64(2) + v_ax1) % b,
                        ]
                    )
                    T.writes(T_reshape[v_ax0, v_ax1])
                    T_reshape[v_ax0, v_ax1] = rxplaceholder[
                        (v_ax0 * b * T.int64(2) + v_ax1) // b % T.int64(10),
                        (v_ax0 * b * T.int64(2) + v_ax1) % b,
                    ]

        @R.function
        def main(
            x: R.Tensor((10, "b"), dtype="float32"),
        ) -> R.Tensor((5, "b * 2"), dtype="float32"):
            b = T.int64()
            lv: R.Shape([5, b * 2]) = R.shape([5, b * 2])
            gv = R.call_tir(Expected3.reshape, (x,), out_ty=R.Tensor((5, b * 2), dtype="float32"))
            return gv

    mod3 = LegalizeOps()(Reshape3)
    tvm.ir.assert_structural_equal(mod3, Expected3)


def test_data_dependent_reshape():
    # fmt: off
    @tvm.script.ir_module
    class DDReshape:
        @R.function
        def main(
            x: R.Tensor([2], dtype="int64"),
            y: R.Tensor([16],dtype='float32'),
        ):
            lv: R.Shape(ndim=2) = R.tensor_to_shape(x)
            gv = R.reshape(y, lv)
            return gv
    # fmt: on

    relax.analysis.well_formed(DDReshape)
    mod = relax.transform.DecomposeOpsForInference()(DDReshape)
    out_mod = relax.transform.LegalizeOps()(mod)

    # fmt: off
    @I.ir_module(s_tir=True)
    class Expected:
        @R.function
        def main(
                x: R.Tensor([2], dtype="int64"),
                y: R.Tensor([16],dtype="float32"),
        ) -> R.Tensor(ndim=2, dtype="float32"):
            M = T.int64()
            N = T.int64()
            gv = R.call_pure_packed("vm.builtin.tensor_to_shape", x, ty_args=(R.Shape(ndim=2),))
            _ = R.match_cast(gv, R.Shape([M,N]))
            _ = R.shape([M,N])
            gv_1 = R.call_tir(Expected.reshape, (y,), out_ty=R.Tensor([M,N], dtype="float32"))
            return gv_1

        @T.prim_func(private=True, s_tir=True)
        def reshape(
            rxplaceholder: T.Buffer(T.int64(16), "float32"),
            var_T_reshape: T.handle,
        ):
            T.func_attr({"tirx.noalias": True})
            M = T.int64()
            N = T.int64()
            T_reshape = T.match_buffer(var_T_reshape, [M,N], "float32")
            for i,j in T.grid(M,N):
                with T.sblock("T_reshape"):
                    vi,vj = T.axis.remap('SS',[i,j])
                    T.reads(rxplaceholder[(vi*N + vj) % 16])
                    T.writes(T_reshape[vi,vj])
                    T_reshape[vi,vj] = rxplaceholder[(vi*N + vj) % 16]

    # fmt: on
    tvm.ir.assert_structural_equal(out_mod, Expected)


def test_split_by_indices():
    # fmt: off
    @tvm.script.ir_module
    class Split:
        @R.function
        def main(x: R.Tensor((2, 10, 4), "float32")) -> R.Tuple([R.Tensor((2, 3, 4), "float32"), R.Tensor((2, 4, 4), "float32"), R.Tensor((2, 3, 4), "float32")]):
            gv: R.Tuple([R.Tensor((2, 3, 4), "float32"), R.Tensor((2, 4, 4), "float32"), R.Tensor((2, 3, 4), "float32")]) = R.split(x, [3, 7], axis=1)
            return gv

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor((2, 10, 4), "float32")) -> R.Tuple([R.Tensor((2, 3, 4), "float32"), R.Tensor((2, 4, 4), "float32"), R.Tensor((2, 3, 4), "float32")]):
            gv = R.call_tir(Expected.split, (x,), [R.Tensor((2, 3, 4), "float32"), R.Tensor((2, 4, 4), "float32"), R.Tensor((2, 3, 4), "float32")])
            return gv

        @T.prim_func(private=True, s_tir=True)
        def split(rxplaceholder: T.Buffer((T.int64(2), T.int64(10), T.int64(4)), "float32"), T_split: T.Buffer((T.int64(2), T.int64(3), T.int64(4)), "float32"), T_split_1: T.Buffer((T.int64(2), T.int64(4), T.int64(4)), "float32"), T_split_2: T.Buffer((T.int64(2), T.int64(3), T.int64(4)), "float32")):
            T.func_attr({"tirx.noalias": True})
            for i0, i1, i2 in T.grid(T.int64(2), T.int64(3), T.int64(4)):
                with T.sblock("T_split"):
                    ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
                    T.reads(rxplaceholder[ax0, ax1, ax2])
                    T.writes(T_split[ax0, ax1, ax2])
                    T_split[ax0, ax1, ax2] = rxplaceholder[ax0, ax1, ax2]
            for i0, i1, i2 in T.grid(T.int64(2), T.int64(4), T.int64(4)):
                with T.sblock("T_split_1"):
                    ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
                    T.reads(rxplaceholder[ax0, ax1 + T.int64(3), ax2])
                    T.writes(T_split_1[ax0, ax1, ax2])
                    T_split_1[ax0, ax1, ax2] = rxplaceholder[ax0, ax1 + T.int64(3), ax2]
            for i0, i1, i2 in T.grid(T.int64(2), T.int64(3), T.int64(4)):
                with T.sblock("T_split_2"):
                    ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
                    T.reads(rxplaceholder[ax0, ax1 + T.int64(7), ax2])
                    T.writes(T_split_2[ax0, ax1, ax2])
                    T_split_2[ax0, ax1, ax2] = rxplaceholder[ax0, ax1 + T.int64(7), ax2]
    # fmt: on

    mod = LegalizeOps()(Split)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_split_by_indices_n_section_indivisible():
    # fmt: off
    @tvm.script.ir_module
    class Split:
        @R.function
        def main(x: R.Tensor((2, 10, 4), "float32")) -> R.Tuple([R.Tensor((2, 4, 4), "float32"), R.Tensor((2, 4, 4), "float32"), R.Tensor((2, 2, 4), "float32")]):
            gv: R.Tuple([R.Tensor((2, 4, 4), "float32"), R.Tensor((2, 4, 4), "float32"), R.Tensor((2, 2, 4), "float32")]) = R.split(x, indices_or_sections=3, axis=1)
            return gv

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor((2, 10, 4), "float32")) -> R.Tuple([R.Tensor((2, 4, 4), "float32"), R.Tensor((2, 4, 4), "float32"), R.Tensor((2, 2, 4), "float32")]):
            gv = R.call_tir(Expected.split, (x,), [R.Tensor((2, 4, 4), "float32"), R.Tensor((2, 4, 4), "float32"), R.Tensor((2, 2, 4), "float32")])
            return gv

        @T.prim_func(private=True, s_tir=True)
        def split(rxplaceholder: T.Buffer((T.int64(2), T.int64(10), T.int64(4)), "float32"), T_split_sections: T.Buffer((T.int64(2), T.int64(4), T.int64(4)), "float32"), T_split_sections_1: T.Buffer((T.int64(2), T.int64(4), T.int64(4)), "float32"), T_split_sections_2: T.Buffer((T.int64(2), T.int64(2), T.int64(4)), "float32")):
            T.func_attr({"tirx.noalias": True})
            for i0, i1, i2 in T.grid(T.int64(2), T.int64(4), T.int64(4)):
                with T.sblock("T_split_sections"):
                    ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
                    T.reads(rxplaceholder[ax0, ax1, ax2])
                    T.writes(T_split_sections[ax0, ax1, ax2])
                    T_split_sections[ax0, ax1, ax2] = rxplaceholder[ax0, ax1, ax2]
            for i0, i1, i2 in T.grid(T.int64(2), T.int64(4), T.int64(4)):
                with T.sblock("T_split_sections_1"):
                    ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
                    T.reads(rxplaceholder[ax0, ax1 + T.int64(4), ax2])
                    T.writes(T_split_sections_1[ax0, ax1, ax2])
                    T_split_sections_1[ax0, ax1, ax2] = rxplaceholder[ax0, ax1 + T.int64(4), ax2]
            for i0, i1, i2 in T.grid(T.int64(2), T.int64(2), T.int64(4)):
                with T.sblock("T_split_sections_2"):
                    ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
                    T.reads(rxplaceholder[ax0, ax1 + T.int64(8), ax2])
                    T.writes(T_split_sections_2[ax0, ax1, ax2])
                    T_split_sections_2[ax0, ax1, ax2] = rxplaceholder[ax0, ax1 + T.int64(8), ax2]

    # fmt: on

    mod = LegalizeOps()(Split)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_split_by_indices_n_section_divisible():
    # fmt: off
    @tvm.script.ir_module
    class Split:
        @R.function
        def main(x: R.Tensor((2, 10, 4), "float32")) -> R.Tuple([R.Tensor((2, 5, 4), "float32"), R.Tensor((2, 5, 4), "float32")]):
            gv: R.Tuple([R.Tensor((2, 5, 4), "float32"), R.Tensor((2, 5, 4), "float32")]) = R.split(x, 2, axis=1)
            return gv

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor((2, 10, 4), "float32")) -> R.Tuple([R.Tensor((2, 5, 4), "float32"), R.Tensor((2, 5, 4), "float32")]):
            gv = R.call_tir(Expected.split, (x,), [R.Tensor((2, 5, 4), "float32"), R.Tensor((2, 5, 4), "float32")])
            return gv

        @T.prim_func(private=True, s_tir=True)
        def split(rxplaceholder: T.Buffer((T.int64(2), T.int64(10), T.int64(4)), "float32"), T_split_sections: T.Buffer((T.int64(2), T.int64(5), T.int64(4)), "float32"), T_split_sections_1: T.Buffer((T.int64(2), T.int64(5), T.int64(4)), "float32")):
            T.func_attr({"tirx.noalias": True})
            for i0, i1, i2 in T.grid(T.int64(2), T.int64(5), T.int64(4)):
                with T.sblock("T_split_sections"):
                    ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
                    T.reads(rxplaceholder[ax0, ax1, ax2])
                    T.writes(T_split_sections[ax0, ax1, ax2])
                    T_split_sections[ax0, ax1, ax2] = rxplaceholder[ax0, ax1, ax2]
            for i0, i1, i2 in T.grid(T.int64(2), T.int64(5), T.int64(4)):
                with T.sblock("T_split_sections_1"):
                    ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
                    T.reads(rxplaceholder[ax0, ax1 + T.int64(5), ax2])
                    T.writes(T_split_sections_1[ax0, ax1, ax2])
                    T_split_sections_1[ax0, ax1, ax2] = rxplaceholder[ax0, ax1 + T.int64(5), ax2]
    # fmt: on

    mod = LegalizeOps()(Split)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_split_by_indices_n_section_divisible_symbolic():
    # fmt: off
    @tvm.script.ir_module
    class Split:
        @R.function
        def main(dumb_param: R.Tensor(("n",)), x: R.Tensor(("m", "n * 3"), "float32")) -> R.Tuple([R.Tensor(("m", "n"), "float32"), R.Tensor(("m", "n"), "float32"), R.Tensor(("m", "n"), "float32")]):
            m = T.int64()
            n = T.int64()
            gv: R.Tuple([R.Tensor((m, n), "float32"), R.Tensor((m, n), "float32"), R.Tensor((m, n), "float32")]) = R.split(x, 3, axis=1)
            return gv

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(dumb_param: R.Tensor(("n",)), x: R.Tensor(("m", "(n * 3)"), "float32")) -> R.Tuple(R.Tensor(("m", "((n * 3) // 3)"), "float32"), R.Tensor(("m", "((((n * 3) // 3) * 2) - ((n * 3) // 3))"), "float32"), R.Tensor(("m", "((n * 3) - (((n * 3) // 3) * 2))"), "float32")):
            m = T.int64()
            n = T.int64()
            gv = R.call_tir(Expected.split, (x, n), [R.Tensor((m, ((n * 3 + 3 - 1) // 3)), "float32"), R.Tensor((m, ((((n * 3 + 3 - 1) // 3) * 2) - ((n * 3 + 3 - 1) // 3))), "float32"), R.Tensor((m, ((n * 3) - (((n * 3 + 3 - 1) // 3) * 2))), "float32")])
            return gv

        @T.prim_func(private=True, s_tir=True)
        def split(var_rxplaceholder: T.handle, n: T.int64, var_T_split_sections: T.handle, var_T_split_sections_1: T.handle, var_T_split_sections_2: T.handle):
            T.func_attr({"tirx.noalias": True})
            m = T.int64()
            rxplaceholder = T.match_buffer(var_rxplaceholder, [m, n * T.int64(3)], dtype="float32")
            T_split_sections = T.match_buffer(var_T_split_sections, [m, (n * T.int64(3) + T.int64(3) - T.int64(1)) // T.int64(3)], dtype="float32")
            T_split_sections_1 = T.match_buffer(var_T_split_sections_1, [m, (n * T.int64(3) + T.int64(3) - T.int64(1)) // T.int64(3) * T.int64(2) - (n * T.int64(3) + T.int64(3) - T.int64(1)) // T.int64(3)], dtype="float32")
            T_split_sections_2 = T.match_buffer(var_T_split_sections_2, [m, n * T.int64(3) - (n * T.int64(3) + T.int64(3) - T.int64(1)) // T.int64(3) * T.int64(2)], dtype="float32")
            for i0, i1 in T.grid(m, n):
                with T.sblock("T_split_sections"):
                    ax0, ax1 = T.axis.remap("SS", [i0, i1])
                    T.reads(rxplaceholder[ax0, ax1])
                    T.writes(T_split_sections[ax0, ax1])
                    T_split_sections[ax0, ax1] = rxplaceholder[ax0, ax1]
            for i0, i1 in T.grid(m, n):
                with T.sblock("T_split_sections_1"):
                    ax0, ax1 = T.axis.remap("SS", [i0, i1])
                    T.reads(rxplaceholder[ax0, ax1 + n])
                    T.writes(T_split_sections_1[ax0, ax1])
                    T_split_sections_1[ax0, ax1] = rxplaceholder[ax0, ax1 + n]
            for i0, i1 in T.grid(m, n):
                with T.sblock("T_split_sections_2"):
                    ax0, ax1 = T.axis.remap("SS", [i0, i1])
                    T.reads(rxplaceholder[ax0, n * T.int64(2) + ax1])
                    T.writes(T_split_sections_2[ax0, ax1])
                    T_split_sections_2[ax0, ax1] = rxplaceholder[ax0, n * T.int64(2) + ax1]
    # fmt: on

    mod = LegalizeOps()(Split)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_squeeze():
    # fmt: off
    @tvm.script.ir_module
    class Squeeze:
        @R.function
        def main(x: R.Tensor((2, 1, 3, 1, 1, 4), "float32")) -> R.Tensor((2, 3, 1, 4), "float32"):
            gv: R.Tensor((2, 3, 1, 4), "float32") = R.squeeze(x, [1, 4])
            return gv

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor((2, 1, 3, 1, 1, 4), "float32")) -> R.Tensor((2, 3, 1, 4), "float32"):
            gv = R.call_tir(Expected.squeeze, (x,), R.Tensor((2, 3, 1, 4), dtype="float32"))
            return gv

        @T.prim_func(private=True, s_tir=True)
        def squeeze(rxplaceholder: T.Buffer((T.int64(2), T.int64(1), T.int64(3), T.int64(1), T.int64(1), T.int64(4)), "float32"), T_squeeze: T.Buffer((T.int64(2), T.int64(3), T.int64(1), T.int64(4)), "float32")):
            T.func_attr({"tirx.noalias": True})
            for i0, i1, i2, i3 in T.grid(T.int64(2), T.int64(3), T.int64(1), T.int64(4)):
                with T.sblock("T_squeeze"):
                    ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
                    T.reads(rxplaceholder[ax0, T.int64(0), ax1, ax2, T.int64(0), ax3])
                    T.writes(T_squeeze[ax0, ax1, ax2, ax3])
                    T_squeeze[ax0, ax1, ax2, ax3] = rxplaceholder[ax0, T.int64(0), ax1, ax2, T.int64(0), ax3]
    # fmt: on

    mod = LegalizeOps()(Squeeze)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_squeeze_no_axis():
    # fmt: off
    @tvm.script.ir_module
    class Squeeze:
        @R.function
        def main(x: R.Tensor((2, 1, 3, 1, 1, 4), "float32")) :
            gv: R.Tensor((2, 3, 4), "float32") = R.squeeze(x)
            return gv

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor((2, 1, 3, 1, 1, 4), "float32")) :
            gv = R.call_tir(Expected.squeeze, (x,), R.Tensor((2, 3, 4), dtype="float32"))
            return gv

        @T.prim_func(private=True, s_tir=True)
        def squeeze(rxplaceholder: T.Buffer((T.int64(2), T.int64(1), T.int64(3), T.int64(1), T.int64(1), T.int64(4)), "float32"), T_squeeze: T.Buffer((T.int64(2), T.int64(3), T.int64(4)), "float32")):
            T.func_attr({"tirx.noalias": True})
            for i0, i1, i2 in T.grid(T.int64(2), T.int64(3), T.int64(4)):
                with T.sblock("T_squeeze"):
                    ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
                    T.reads(rxplaceholder[ax0, T.int64(0), ax1, T.int64(0), T.int64(0), ax2])
                    T.writes(T_squeeze[ax0, ax1, ax2])
                    T_squeeze[ax0, ax1, ax2] = rxplaceholder[ax0, T.int64(0), ax1, T.int64(0), T.int64(0), ax2]
    # fmt: on

    mod = LegalizeOps()(Squeeze)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_squeeze_symbolic():
    # fmt: off
    @tvm.script.ir_module
    class Squeeze:
        @R.function
        def main(x: R.Tensor(("a", 1, "b", 1), "float32")) -> R.Tensor(("a", "b", 1), "float32"):
            a = T.int64()
            b = T.int64()
            gv: R.Tensor((a, b, 1), "float32") = R.squeeze(x, [1])
            return gv

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor(("a", 1, "b", 1), "float32")) -> R.Tensor(("a", "b", 1), "float32"):
            a = T.int64()
            b = T.int64()
            gv = R.call_tir(Expected.squeeze, (x,), R.Tensor((a, b, 1), dtype="float32"))
            return gv

        @T.prim_func(private=True, s_tir=True)
        def squeeze(var_rxplaceholder: T.handle, var_T_squeeze: T.handle):
            T.func_attr({"tirx.noalias": True})
            a = T.int64()
            b = T.int64()
            rxplaceholder = T.match_buffer(var_rxplaceholder, [a, T.int64(1), b, T.int64(1)], dtype="float32")
            T_squeeze = T.match_buffer(var_T_squeeze, [a, b, T.int64(1)], dtype="float32")
            for i0, i1, i2 in T.grid(a, b, T.int64(1)):
                with T.sblock("T_squeeze"):
                    ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
                    T.reads(rxplaceholder[ax0, T.int64(0), ax1, ax2])
                    T.writes(T_squeeze[ax0, ax1, ax2])
                    T_squeeze[ax0, ax1, ax2] = rxplaceholder[ax0, T.int64(0), ax1, ax2]
    # fmt: on

    mod = LegalizeOps()(Squeeze)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_collapse_sum_like():
    # fmt: off
    @tvm.script.ir_module
    class CollapseSumLike:
        @R.function
        def main(x: R.Tensor((2, 3), "float32"), y: R.Tensor((1, 3), "float32")) -> R.Tensor((1, 3), "float32"):
            gv: R.Tensor((1, 3), "float32") = R.collapse_sum_like(x, y)
            return gv

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor((2, 3), "float32"), y: R.Tensor((1, 3), "float32")) -> R.Tensor((1, 3), "float32"):
            gv = R.call_tir(Expected.collapse_sum, (x,), R.Tensor((1, 3), dtype="float32"))
            return gv

        @T.prim_func(private=True, s_tir=True)
        def collapse_sum(rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"), rxplaceholder_red: T.Buffer((T.int64(1), T.int64(3)), "float32")):
            T.func_attr({"tirx.noalias": True})
            for i0, i1, i2 in T.grid(T.int64(1), T.int64(3), T.int64(2)):
                with T.sblock("rxplaceholder_red"):
                    ax0, ax1, k0 = T.axis.remap("SSR", [i0, i1, i2])
                    T.reads(rxplaceholder[k0, ax1])
                    T.writes(rxplaceholder_red[ax0, ax1])
                    with T.init():
                        rxplaceholder_red[ax0, ax1] = T.float32(0)
                    rxplaceholder_red[ax0, ax1] = rxplaceholder_red[ax0, ax1] + rxplaceholder[k0, ax1]
    # fmt: on

    mod = LegalizeOps()(CollapseSumLike)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_collapse_sum_to():
    # fmt: off
    @tvm.script.ir_module
    class CollapseSumTo:
        @R.function
        def main(x: R.Tensor((3, 2, 3), "float32")) -> R.Tensor((2, 1), "float32"):
            gv: R.Tensor((2, 1), "float32") = R.collapse_sum_to(x, (2, 1))
            return gv

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(
            x: R.Tensor((3, 2, 3), dtype="float32")
        ) -> R.Tensor((2, 1), dtype="float32"):
            # block 0
            gv = R.call_tir(Expected.collapse_sum, (x,), R.Tensor((2, 1), dtype="float32"))
            return gv

        @T.prim_func(private=True, s_tir=True)
        def collapse_sum(rxplaceholder: T.Buffer((T.int64(3), T.int64(2), T.int64(3)), "float32"), rxplaceholder_red: T.Buffer((T.int64(2), T.int64(1)), "float32")):
            T.func_attr({"tirx.noalias": True})
            for ax0, ax1, k0, k2 in T.grid(T.int64(2), T.int64(1), T.int64(3), T.int64(3)):
                with T.sblock("rxplaceholder_red"):
                    v_ax0, v_ax1, v_k0, v_k2 = T.axis.remap("SSRR", [ax0, ax1, k0, k2])
                    T.reads(rxplaceholder[v_k0, v_ax0, v_k2])
                    T.writes(rxplaceholder_red[v_ax0, v_ax1])
                    with T.init():
                        rxplaceholder_red[v_ax0, v_ax1] = T.float32(0)
                    rxplaceholder_red[v_ax0, v_ax1] = (rxplaceholder_red[v_ax0, v_ax1] + rxplaceholder[v_k0, v_ax0, v_k2])
    # fmt: on

    mod = LegalizeOps()(CollapseSumTo)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_repeat():
    # fmt: off
    @I.ir_module(s_tir=True)
    class Repeat:
        @R.function
        def main(x: R.Tensor((3, 2, 3), "float32")):
            gv = R.repeat(x, 2, 0)
            return gv

    @I.ir_module(s_tir=True)
    class Expected:
        @R.function
        def main(x: R.Tensor((3, 2, 3), dtype="float32")) -> R.Tensor((6, 2, 3), dtype="float32"):
            gv = R.call_tir(Expected.repeat, (x,), out_ty=R.Tensor((6, 2, 3), dtype="float32"))
            return gv

        @T.prim_func(private=True, s_tir=True)
        def repeat(rxplaceholder: T.Buffer((T.int64(3), T.int64(2), T.int64(3)), "float32"), T_repeat: T.Buffer((T.int64(6), T.int64(2), T.int64(3)), "float32")):
            T.func_attr({"tirx.noalias": True})
            # with T.sblock("root"):
            for ax0, ax1, ax2 in T.grid(T.int64(6), T.int64(2), T.int64(3)):
                with T.sblock("T_repeat"):
                    v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
                    T.reads(rxplaceholder[v_ax0 // T.int64(2), v_ax1, v_ax2])
                    T.writes(T_repeat[v_ax0, v_ax1, v_ax2])
                    T_repeat[v_ax0, v_ax1, v_ax2] = rxplaceholder[v_ax0 // T.int64(2), v_ax1, v_ax2]
    # fmt: on

    mod = LegalizeOps()(Repeat)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_repeat_no_axis():
    # fmt: off
    @I.ir_module(s_tir=True)
    class Repeat:
        @R.function
        def main(x: R.Tensor((3, 2, 3), "float32")):
            gv = R.repeat(x, 2)
            return gv

    @I.ir_module(s_tir=True)
    class Expected:
        @R.function
        def main(
            x: R.Tensor((3, 2, 3), dtype="float32")
        ) -> R.Tensor((36,), dtype="float32"):
            gv = R.call_tir(Expected.repeat, (x,), out_ty=R.Tensor((36,), dtype="float32"))
            return gv

        @T.prim_func(private=True, s_tir=True)
        def repeat(
            rxplaceholder: T.Buffer((T.int64(3), T.int64(2), T.int64(3)), "float32"),
            T_repeat: T.Buffer((T.int64(36),), "float32"),
        ):
            T.func_attr({"tirx.noalias": True})
            # with T.sblock("root"):
            T_reshape = T.sblock_alloc_buffer((T.int64(18),))
            for ax0 in range(T.int64(18)):
                with T.sblock("T_reshape"):
                    v_ax0 = T.axis.spatial(T.int64(18), ax0)
                    T.reads(
                        rxplaceholder[
                            v_ax0 % T.int64(18) // T.int64(6),
                            v_ax0 % T.int64(6) // T.int64(3),
                            v_ax0 % T.int64(3),
                        ]
                    )
                    T.writes(T_reshape[v_ax0])
                    T_reshape[v_ax0] = rxplaceholder[
                        v_ax0 % T.int64(18) // T.int64(6),
                        v_ax0 % T.int64(6) // T.int64(3),
                        v_ax0 % T.int64(3),
                    ]
            for ax0 in range(T.int64(36)):
                with T.sblock("T_repeat"):
                    v_ax0 = T.axis.spatial(T.int64(36), ax0)
                    T.reads(T_reshape[v_ax0 // T.int64(2)])
                    T.writes(T_repeat[v_ax0])
                    T_repeat[v_ax0] = T_reshape[v_ax0 // T.int64(2)]
    # fmt: on

    mod = LegalizeOps()(Repeat)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_repeat_symbolic():
    # fmt: off
    @I.ir_module(s_tir=True)
    class Repeat:
        @R.function
        def main(x: R.Tensor(("a", "b", "c"), "float32")):
            gv = R.repeat(x, 2, 0)
            return gv

    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(private=True, s_tir=True)
        def repeat(var_rxplaceholder: T.handle, var_T_repeat: T.handle):
            T.func_attr({"tirx.noalias": True})
            a = T.int64()
            b = T.int64()
            c = T.int64()
            rxplaceholder = T.match_buffer(var_rxplaceholder, (a, b, c))
            T_repeat = T.match_buffer(var_T_repeat, (T.int64(2) * a, b, c))
            # with T.sblock("root"):
            for ax0, ax1, ax2 in T.grid(a * T.int64(2), b, c):
                with T.sblock("T_repeat"):
                    v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
                    T.reads(rxplaceholder[v_ax0 // T.int64(2), v_ax1, v_ax2])
                    T.writes(T_repeat[v_ax0, v_ax1, v_ax2])
                    T_repeat[v_ax0, v_ax1, v_ax2] = rxplaceholder[v_ax0 // T.int64(2), v_ax1, v_ax2]

        @R.function
        def main(x: R.Tensor(("a", "b", "c"), dtype="float32")) -> R.Tensor(("2 * a", "b", "c"), dtype="float32"):
            a = T.int64()
            b = T.int64()
            c = T.int64()
            gv = R.call_tir(Expected.repeat, (x,), out_ty=R.Tensor((2 * a, b, c), dtype="float32"))
            return gv
    # fmt: on

    mod = LegalizeOps()(Repeat)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_tile():
    # fmt: off
    @I.ir_module(s_tir=True)
    class Tile:
        @R.function
        def main(x: R.Tensor((3, 2, 3), "float32")):
            gv = R.tile(x, (2, 1, 2, 3))
            return gv

    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(private=True, s_tir=True)
        def tile(rxplaceholder: T.Buffer((T.int64(3), T.int64(2), T.int64(3)), "float32"), T_tile: T.Buffer((T.int64(2), T.int64(3), T.int64(4), T.int64(9)), "float32")):
            T.func_attr({"tirx.noalias": True})
            # with T.sblock("root"):
            for ax0, ax1, ax2, ax3 in T.grid(T.int64(2), T.int64(3), T.int64(4), T.int64(9)):
                with T.sblock("T_tile"):
                    v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
                    T.reads(rxplaceholder[v_ax1 % T.int64(3), v_ax2 % T.int64(2), v_ax3 % T.int64(3)])
                    T.writes(T_tile[v_ax0, v_ax1, v_ax2, v_ax3])
                    T_tile[v_ax0, v_ax1, v_ax2, v_ax3] = rxplaceholder[v_ax1 % T.int64(3), v_ax2 % T.int64(2), v_ax3 % T.int64(3)]

        @R.function
        def main(x: R.Tensor((3, 2, 3), dtype="float32")) -> R.Tensor((2, 3, 4, 9), dtype="float32"):
            gv = R.call_tir(Expected.tile, (x,), out_ty=R.Tensor((2, 3, 4, 9), dtype="float32"))
            return gv
    # fmt: on

    mod = LegalizeOps()(Tile)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_tile_symbolic():
    # fmt: off
    @I.ir_module(s_tir=True)
    class Tile:
        @R.function
        def main(x: R.Tensor(("a", "b", "c"), "float32")):
            gv = R.tile(x, (2, 1, 2, 3))
            return gv

    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(private=True, s_tir=True)
        def tile(var_rxplaceholder: T.handle, var_T_tile: T.handle):
            T.func_attr({"tirx.noalias": True})
            a = T.int64()
            b = T.int64()
            c = T.int64()
            rxplaceholder = T.match_buffer(var_rxplaceholder, (a, b, c))
            T_tile = T.match_buffer(var_T_tile, (T.int64(2), a, b * T.int64(2), c * T.int64(3)))
            # with T.sblock("root"):
            for ax0, ax1, ax2, ax3 in T.grid(T.int64(2), a, b * T.int64(2), c * T.int64(3)):
                with T.sblock("T_tile"):
                    v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
                    T.reads(rxplaceholder[v_ax1 % a, v_ax2 % b, v_ax3 % c])
                    T.writes(T_tile[v_ax0, v_ax1, v_ax2, v_ax3])
                    T_tile[v_ax0, v_ax1, v_ax2, v_ax3] = rxplaceholder[v_ax1 % a, v_ax2 % b, v_ax3 % c]

        @R.function
        def main(x: R.Tensor(("a", "b", "c"), dtype="float32")) -> R.Tensor((2, "a", "b * 2", "c * 3"), dtype="float32"):
            a = T.int64()
            b = T.int64()
            c = T.int64()
            gv = R.call_tir(Expected.tile, (x,), out_ty=R.Tensor((2, a, b * 2, c * 3), dtype="float32"))
            return gv
    # fmt: on
    mod = LegalizeOps()(Tile)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_flip():
    # fmt: off
    @I.ir_module(s_tir=True)
    class Flip:
        @R.function
        def main(x: R.Tensor((2, 3), "float32")):
            gv = R.flip(x, axis=0)
            return gv

    @I.ir_module(s_tir=True)
    class Expected:
        @R.function
        def main(x: R.Tensor((2, 3), dtype="float32")) -> R.Tensor((2, 3), dtype="float32"):
            cls = Expected
            gv = R.call_tir(cls.flip, (x,), out_ty=R.Tensor((2, 3), dtype="float32"))
            return gv

        @T.prim_func(private=True, s_tir=True)
        def flip(
            rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"),
            T_reverse_sequence: T.Buffer((T.int64(2), T.int64(3)), "float32"),
        ):
            T.func_attr({"tirx.noalias": True})
            for ax0, ax1 in T.grid(T.int64(2), T.int64(3)):
                with T.sblock("T_reverse_sequence"):
                    v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
                    T.reads(rxplaceholder[T.int64(1) - v_ax0, v_ax1])
                    T.writes(T_reverse_sequence[v_ax0, v_ax1])
                    T_reverse_sequence[v_ax0, v_ax1] = rxplaceholder[
                        T.int64(1) - v_ax0, v_ax1
                    ]

    # fmt: on

    mod = LegalizeOps()(Flip)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_flip_symbolic():
    # fmt: off
    @I.ir_module(s_tir=True)
    class Flip:
        @R.function
        def main(x: R.Tensor(("a", "b"), "float32")):
            gv = R.flip(x, axis=1)
            return gv

    @I.ir_module(s_tir=True)
    class Expected:
        @R.function
        def main(
            x: R.Tensor(("a", "b"), dtype="float32")
        ) -> R.Tensor(("a", "b"), dtype="float32"):
            a = T.int64()
            b = T.int64()
            cls = Expected
            gv = R.call_tir(cls.flip, (x,), out_ty=R.Tensor((a, b), dtype="float32"))
            return gv

        @T.prim_func(private=True, s_tir=True)
        def flip(var_rxplaceholder: T.handle, var_T_reverse_sequence: T.handle):
            T.func_attr({"tirx.noalias": True})
            a, b = T.int64(), T.int64()
            rxplaceholder = T.match_buffer(var_rxplaceholder, (a, b))
            T_reverse_sequence = T.match_buffer(var_T_reverse_sequence, (a, b))
            for ax0, ax1 in T.grid(a, b):
                with T.sblock("T_reverse_sequence"):
                    v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
                    T.reads(rxplaceholder[v_ax0, b - v_ax1 - T.int64(1)])
                    T.writes(T_reverse_sequence[v_ax0, v_ax1])
                    T_reverse_sequence[v_ax0, v_ax1] = rxplaceholder[
                        v_ax0, b - v_ax1 - T.int64(1)
                    ]

    # fmt: on

    mod = LegalizeOps()(Flip)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_reverse_sequence():
    # fmt: off
    @I.ir_module(s_tir=True)
    class ReverseSequence:
        @R.function
        def main(x: R.Tensor((4, 2, 3), "float32"), seq_lengths: R.Tensor((2,), "int64")):
            gv = R.reverse_sequence(x, seq_lengths, seq_axis=0, batch_axis=1)
            return gv

    @I.ir_module(s_tir=True)
    class Expected:
        @R.function
        def main(
            x: R.Tensor((4, 2, 3), dtype="float32"),
            seq_lengths: R.Tensor((2,), dtype="int64"),
        ) -> R.Tensor((4, 2, 3), dtype="float32"):
            cls = Expected
            gv = R.call_tir(
                cls.reverse_sequence,
                (x, seq_lengths),
                out_ty=R.Tensor((4, 2, 3), dtype="float32"),
            )
            return gv

        @T.prim_func(private=True, s_tir=True)
        def reverse_sequence(
            rxplaceholder: T.Buffer((T.int64(4), T.int64(2), T.int64(3)), "float32"),
            seq_lengths: T.Buffer((T.int64(2),), "int64"),
            T_reverse_sequence: T.Buffer((T.int64(4), T.int64(2), T.int64(3)), "float32"),
        ):
            T.func_attr({"tirx.noalias": True})
            for ax0, ax1, ax2 in T.grid(T.int64(4), T.int64(2), T.int64(3)):
                with T.sblock("T_reverse_sequence"):
                    v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
                    T.reads(rxplaceholder[T.int64(0):T.int64(4), v_ax1, v_ax2], seq_lengths[v_ax1])
                    T.writes(T_reverse_sequence[v_ax0, v_ax1, v_ax2])
                    T_reverse_sequence[v_ax0, v_ax1, v_ax2] = rxplaceholder[
                        T.if_then_else(
                            seq_lengths[v_ax1] <= T.int64(1) or seq_lengths[v_ax1] <= v_ax0,
                            v_ax0,
                            T.if_then_else(
                                T.int64(4) < seq_lengths[v_ax1],
                                T.int64(3) - v_ax0,
                                seq_lengths[v_ax1] - v_ax0 - T.int64(1),
                            ),
                        ),
                        v_ax1,
                        v_ax2,
                    ]

    # fmt: on

    mod = LegalizeOps()(ReverseSequence)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_scatter_elements():
    # fmt: off
    @I.ir_module(s_tir=True)
    class ScatterElements:
        @R.function
        def main(x: R.Tensor((4,4), "float32"), indices: R.Tensor((2,2), "int64"), updates: R.Tensor((2,2), "float32")):
            gv = R.scatter_elements(x, indices, updates, axis=1)
            return gv
    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(private=True, s_tir=True)
        def scatter_elements(
            var_rxplaceholder: T.handle,
            var_rxplaceholder_1: T.handle,
            var_rxplaceholder_2: T.handle,
            out_buf: T.Buffer((T.int64(4), T.int64(4)), "float32"),
        ):
            T.func_attr({"tirx.noalias": True})
            rxplaceholder = T.match_buffer(
                var_rxplaceholder, (T.int64(4), T.int64(4)), offset_factor=1
            )
            rxplaceholder_1 = T.match_buffer(
                var_rxplaceholder_1, (T.int64(2), T.int64(2)), "int64", offset_factor=1
            )
            rxplaceholder_2 = T.match_buffer(
                var_rxplaceholder_2, (T.int64(2), T.int64(2)), offset_factor=1
            )
            with T.sblock("scatter_elements_generic"):
                T.attr(0, "pragma_scope", "seq")
                for i in T.parallel(T.int64(16)):
                    out_buf[i // T.int64(4), i % T.int64(4)] = rxplaceholder[
                        i // T.int64(4), i % T.int64(4)
                    ]
                for fused in T.parallel(T.int64(2)):
                    for k in range(T.int64(2)):
                        out_buf[
                            (
                                fused * T.int64(4)
                                + (
                                    rxplaceholder_1[
                                        (fused * T.int64(2) + k) // T.int64(2),
                                        (fused * T.int64(2) + k) % T.int64(2),
                                    ]
                                    + T.Cast(
                                        "int64",
                                        rxplaceholder_1[
                                            (fused * T.int64(2) + k) // T.int64(2),
                                            (fused * T.int64(2) + k) % T.int64(2),
                                        ]
                                        < T.int64(0),
                                    )
                                    * T.int64(4)
                                )
                            )
                            // T.int64(4),
                            (
                                fused * T.int64(4)
                                + (
                                    rxplaceholder_1[
                                        (fused * T.int64(2) + k) // T.int64(2),
                                        (fused * T.int64(2) + k) % T.int64(2),
                                    ]
                                    + T.Cast(
                                        "int64",
                                        rxplaceholder_1[
                                            (fused * T.int64(2) + k) // T.int64(2),
                                            (fused * T.int64(2) + k) % T.int64(2),
                                        ]
                                        < T.int64(0),
                                    )
                                    * T.int64(4)
                                )
                            )
                            % T.int64(4),
                        ] = rxplaceholder_2[
                            (fused * T.int64(2) + k) // T.int64(2),
                            (fused * T.int64(2) + k) % T.int64(2),
                        ]

        @R.function
        def main(
            x: R.Tensor((4, 4), dtype="float32"),
            indices: R.Tensor((2, 2), dtype="int64"),
            updates: R.Tensor((2, 2), dtype="float32"),
        ) -> R.Tensor((4, 4), dtype="float32"):
            gv = R.call_tir(
                Expected.scatter_elements,
                (x, indices, updates),
                out_ty=R.Tensor((4, 4), dtype="float32"),
            )
            return gv

    # fmt: on
    mod = LegalizeOps()(ScatterElements)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_scatter_elements_symbolic():
    # fmt: off
    @I.ir_module(s_tir=True)
    class ScatterElements:
        @R.function
        def main(x: R.Tensor(("a", "b"), "float32"), indices:R.Tensor(("m", "n"), "int64"), updates:R.Tensor(("m","n"), "float32")):
            gv = R.scatter_elements(x, indices, updates, axis=1)
            return gv
    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(private=True, s_tir=True)
        def scatter_elements(
            var_rxplaceholder: T.handle,
            var_rxplaceholder_1: T.handle,
            var_rxplaceholder_2: T.handle,
            var_scatter_elements_generic: T.handle,
        ):
            T.func_attr({"tirx.noalias": True})
            a, b = T.int64(), T.int64()
            rxplaceholder = T.match_buffer(var_rxplaceholder, (a, b), offset_factor=1)
            m, n = T.int64(), T.int64()
            rxplaceholder_1 = T.match_buffer(
                var_rxplaceholder_1, (m, n), "int64", offset_factor=1
            )
            rxplaceholder_2 = T.match_buffer(var_rxplaceholder_2, (m, n), offset_factor=1)
            out_buf = T.match_buffer(var_scatter_elements_generic, (a, b))
            with T.sblock("scatter_elements_generic"):
                T.attr(0, "pragma_scope", "seq")
                for i in T.parallel(a * b):
                    out_buf[i // b, i % b] = rxplaceholder[i // b, i % b]
                for fused in T.parallel(m):
                    for k in range(n):
                        out_buf[
                            (
                                fused * b
                                + (
                                    rxplaceholder_1[
                                        (fused * n + k) // n, (fused * n + k) % n
                                    ]
                                    + T.Cast(
                                        "int64",
                                        rxplaceholder_1[
                                            (fused * n + k) // n, (fused * n + k) % n
                                        ]
                                        < T.int64(0),
                                    )
                                    * b
                                )
                            )
                            // b,
                            (
                                fused * b
                                + (
                                    rxplaceholder_1[
                                        (fused * n + k) // n, (fused * n + k) % n
                                    ]
                                    + T.Cast(
                                        "int64",
                                        rxplaceholder_1[
                                            (fused * n + k) // n, (fused * n + k) % n
                                        ]
                                        < T.int64(0),
                                    )
                                    * b
                                )
                            )
                            % b,
                        ] = rxplaceholder_2[(fused * n + k) // n, (fused * n + k) % n]

        @R.function
        def main(
            x: R.Tensor(("a", "b"), dtype="float32"),
            indices: R.Tensor(("m", "n"), dtype="int64"),
            updates: R.Tensor(("m", "n"), dtype="float32"),
        ) -> R.Tensor(("a", "b"), dtype="float32"):
            a = T.int64()
            b = T.int64()
            m = T.int64()
            n = T.int64()
            gv = R.call_tir(
                Expected.scatter_elements,
                (x, indices, updates),
                out_ty=R.Tensor((a, b), dtype="float32"),
            )
            return gv
    # fmt: on

    mod = LegalizeOps()(ScatterElements)
    tvm.ir.assert_structural_equal(mod, Expected)


@pytest.mark.gpu
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not enabled")
def test_scatter_elements_gpu():
    """scatter_elements lowered for GPU must build"""
    target = "cuda"

    @I.ir_module(s_tir=True)
    class Mod:
        @R.function
        def main(
            x: R.Tensor((4, 8), "float32"),
            indices: R.Tensor((2, 8), "int64"),
            updates: R.Tensor((2, 8), "float32"),
        ):
            with R.dataflow():
                lv = R.scatter_elements(x, indices, updates, axis=0)
                gv = lv
                R.output(gv)
            return gv

    with tvm.target.Target(target):
        mod = LegalizeOps()(Mod)
    relax.build(mod, target=target)


def test_layout_transform():
    transformation = lambda a, b, c: (a, c, b // 3, b % 3)
    pad_value = 2

    # fmt: off
    @I.ir_module(s_tir=True)
    class LayoutTransform:
        @R.function
        def main(x: R.Tensor((10, 21, 30), "float32")):
            gv = R.layout_transform(
                x, index_map=transformation, pad_value=pad_value
            )
            return gv

    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(private=True, s_tir=True)
        def te_layout_transform(A: T.Buffer((T.int64(10), T.int64(21), T.int64(30)), "float32"), te_layout_transform_1: T.Buffer((T.int64(10), T.int64(30), T.int64(7), T.int64(3)), "float32")):
            T.func_attr({"tirx.noalias": True})
            # with T.sblock("root"):
            for i0, i1, i2 in T.grid(T.int64(10), T.int64(21), T.int64(30)):
                with T.sblock("te_layout_transform"):
                    v_i0, v_i1, v_i2 = T.axis.remap("SSS", [i0, i1, i2])
                    T.reads(A[v_i0, v_i1, v_i2])
                    T.writes(te_layout_transform_1[v_i0, v_i2, v_i1 // T.int64(3), v_i1 % T.int64(3)])
                    te_layout_transform_1[v_i0, v_i2, v_i1 // T.int64(3), v_i1 % T.int64(3)] = A[v_i0, v_i1, v_i2]

        @R.function
        def main(x: R.Tensor((10, 21, 30), dtype="float32")) -> R.Tensor((10, 30, 7, 3), dtype="float32"):
            cls = Expected
            gv = R.call_tir(cls.te_layout_transform, (x,), out_ty=R.Tensor((10, 30, 7, 3), dtype="float32"))
            return gv
    # fmt: on

    mod = LegalizeOps()(LayoutTransform)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_layout_transform_with_pad():
    transformation = lambda a, b, c: (a, c, b // 3, b % 3)
    pad_value = 2

    # fmt: off
    @I.ir_module(s_tir=True)
    class LayoutTransform:
        @R.function
        def main(x: R.Tensor((10, 20, 30), "float32")):
            gv = R.layout_transform(
                x, index_map=transformation, pad_value=pad_value
            )
            return gv

    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(private=True, s_tir=True)
        def te_layout_transform_with_pad(A: T.Buffer((T.int64(10), T.int64(20), T.int64(30)), "float32"), te_layout_transform_with_pad_1: T.Buffer((T.int64(10), T.int64(30), T.int64(7), T.int64(3)), "float32")):
            T.func_attr({"tirx.noalias": True})
            # with T.sblock("root"):
            for axis0, axis1, axis2, axis3 in T.grid(T.int64(10), T.int64(30), T.int64(7), T.int64(3)):
                with T.sblock("te_layout_transform_with_pad"):
                    v_axis0, v_axis1, v_axis2, v_axis3 = T.axis.remap("SSSS", [axis0, axis1, axis2, axis3])
                    T.reads(A[v_axis0, v_axis2 * T.int64(3) + v_axis3, v_axis1])
                    T.writes(te_layout_transform_with_pad_1[v_axis0, v_axis1, v_axis2, v_axis3])
                    te_layout_transform_with_pad_1[v_axis0, v_axis1, v_axis2, v_axis3] = T.if_then_else(v_axis2 == T.int64(6) and v_axis3 == T.int64(2), T.float32(2), A[v_axis0, v_axis2 * T.int64(3) + v_axis3, v_axis1])

        @R.function
        def main(x: R.Tensor((10, 20, 30), dtype="float32")) -> R.Tensor((10, 30, 7, 3), dtype="float32"):
            cls = Expected
            gv = R.call_tir(cls.te_layout_transform_with_pad, (x,), out_ty=R.Tensor((10, 30, 7, 3), dtype="float32"))
            return gv
    # fmt: on

    mod = LegalizeOps()(LayoutTransform)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_layout_transform_symbolic():
    transformation = lambda a, b, c: (a, c, b // 3, b % 3)
    pad_value = 2

    # fmt: off
    @I.ir_module(s_tir=True)
    class LayoutTransform:
        @R.function
        def main(x: R.Tensor(("a", "b", "c"), "float32")):
            gv = R.layout_transform(
                x, index_map=transformation, pad_value=pad_value
            )
            return gv

    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(private=True, s_tir=True)
        def te_layout_transform_with_pad(var_A: T.handle, var_te_layout_transform_with_pad: T.handle):
            T.func_attr({"tirx.noalias": True})
            a, b, c = T.int64(), T.int64(), T.int64()
            A = T.match_buffer(var_A, (a, b, c))
            te_layout_transform_with_pad_1 = T.match_buffer(var_te_layout_transform_with_pad, (a, c, (b - b % T.int64(-3)) // T.int64(3), T.int64(3)))
            # with T.sblock("root"):
            for axis0, axis1, axis2, axis3 in T.grid(a, c, (b - b % T.int64(-3)) // T.int64(3), T.int64(3)):
                with T.sblock("te_layout_transform_with_pad_with_pad"):
                    v_axis0, v_axis1, v_axis2, v_axis3 = T.axis.remap("SSSS", [axis0, axis1, axis2, axis3])
                    T.reads(A[v_axis0, v_axis2 * T.int64(3) + v_axis3, v_axis1])
                    T.writes(te_layout_transform_with_pad_1[v_axis0, v_axis1, v_axis2, v_axis3])
                    te_layout_transform_with_pad_1[v_axis0, v_axis1, v_axis2, v_axis3] = T.if_then_else(b % T.int64(-3) < T.int64(0) and v_axis2 == b // T.int64(3) and b % T.int64(3) <= v_axis3, T.float32(2), A[v_axis0, v_axis2 * T.int64(3) + v_axis3, v_axis1])

        @R.function
        def main(x: R.Tensor(("a", "b", "c"), dtype="float32")) -> R.Tensor(("a", "c", "(b - b % -3) // 3", 3), dtype="float32"):
            a = T.int64()
            c = T.int64()
            b = T.int64()
            cls = Expected
            gv = R.call_tir(cls.te_layout_transform_with_pad, (x,), out_ty=R.Tensor((a, c, (b - b % -3) // 3, 3), dtype="float32"))
            return gv
    # fmt: on

    mod = LegalizeOps()(LayoutTransform)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_func_ty_of_legalized_layout_transform():
    """PrimFunc shape information must be correct

    This is a regression test.  Previously, the legalization of
    `R.layout_transform` produced a PrimFunc with `FuncType`
    different than its actual signature.  This resulted in errors
    when later passes attempted to infer the Type.
    """

    @I.ir_module(s_tir=True)
    class Before:
        @R.function
        def main(
            x: R.Tensor((16,), dtype="float32"), y: R.Tensor((16,), dtype="float32")
        ) -> R.Tensor((16,), dtype="float32"):
            R.func_attr({"relax.force_pure": True})
            with R.dataflow():
                lv: R.Tensor((4, 4), dtype="float32") = R.layout_transform(
                    x, index_map=lambda i: (i // 4, i % 4), pad_value=None
                )
                gv: R.Tensor((4, 4), dtype="float32") = lv
                R.output(gv)
            return gv

    After = tvm.ir.transform.Sequential(
        [
            relax.transform.LegalizeOps(),
            relax.transform.ToNonDataflow(),
            relax.transform.RemovePurityChecking(),
            relax.transform.CallTIRRewrite(),
        ]
    )(Before)

    @I.ir_module(s_tir=True)
    class Expected:
        @R.function
        def main(
            x: R.Tensor((16,), dtype="float32"),
            y: R.Tensor((16,), dtype="float32"),
        ):
            R.func_attr({"relax.force_pure": True})
            cls = Expected
            alloc: R.Tensor((4, 4), dtype="float32") = R.emit_with_ty(
                "relax.builtin.alloc_tensor",
                (R.shape([4, 4]), R.dtype("float32"), R.prim_value(0), R.str("global")),
                (R.Tensor((4, 4), dtype="float32"),),
            )
            cls.te_layout_transform(x, alloc)
            lv = alloc
            gv = lv
            return gv

        @T.prim_func(private=True, s_tir=True)
        def te_layout_transform(
            A: T.Buffer((T.int64(16),), "float32"),
            te_layout_transform: T.Buffer((T.int64(4), T.int64(4)), "float32"),
        ):
            T.func_attr({"tirx.noalias": True})
            for i in range(T.int64(16)):
                with T.sblock("te_layout_transform"):
                    vi = T.axis.spatial(T.int64(16), i)
                    te_layout_transform[vi // T.int64(4), vi % T.int64(4)] = A[vi]

    tvm.ir.assert_structural_equal(Expected, After)


def test_scatter_nd():
    # fmt: off
    @I.ir_module(s_tir=True)
    class Before:
        @R.function
        def main(
            data: R.Tensor((8,), "float32"),
            indices: R.Tensor((4, 1), "int64"),
            updates: R.Tensor((4,), "float32"),
        ) -> R.Tensor((8,), "float32"):
            gv: R.Tensor((8,), "float32") = R.scatter_nd(data, indices, updates, reduction="update")
            return gv

    After = relax.transform.LegalizeOps()(Before)

    @I.ir_module(s_tir=True)
    class Expected:
        @R.function
        def main(
            data: R.Tensor((8,), "float32"),
            indices: R.Tensor((4, 1), "int64"),
            updates: R.Tensor((4,), "float32"),
        ) -> R.Tensor((8,), "float32"):
            gv = R.call_tir(
                Expected.scatter_nd, (data, indices, updates), R.Tensor((8,), dtype="float32")
            )
            return gv

        @T.prim_func(private=True, s_tir=True)
        def scatter_nd(var_data: T.handle, var_indices: T.handle, var_updates: T.handle, var_scatter_nd_generic: T.handle):
            T.func_attr({"tirx.noalias": True})
            data = T.match_buffer(var_data, (T.int64(8),), offset_factor=1)
            indices = T.match_buffer(var_indices, (T.int64(4), T.int64(1)), "int64")
            updates = T.match_buffer(var_updates, (T.int64(4),), offset_factor=1)
            out_buf = T.match_buffer(var_scatter_nd_generic, (T.int64(8),))
            with T.sblock("root"):
                T.reads()
                T.writes()
                T_transpose = T.sblock_alloc_buffer((T.int64(1), T.int64(4)), "int64")
                for ax0 in range(T.int64(1)):
                    for ax1 in range(T.int64(4)):
                        with T.sblock("T_transpose"):
                            v_ax0 = T.axis.spatial(T.int64(1), ax0)
                            v_ax1 = T.axis.spatial(T.int64(4), ax1)
                            T.reads(indices[v_ax1, v_ax0])
                            T.writes(T_transpose[v_ax0, v_ax1])
                            T_transpose[v_ax0, v_ax1] = indices[v_ax1, v_ax0]
                with T.sblock("scatter_nd_generic"):
                    T.reads()
                    T.writes()
                    T.attr(0, "pragma_scope", "seq")
                    for i in range(T.int64(8)):
                        out_buf[i] = data[i]
                    for j in range(T.int64(4)):
                        for k in T.parallel(T.int64(1)):
                            out_buf[k + T_transpose[j // T.int64(4), j % T.int64(4)]] = updates[j + k]

    # fmt: on
    tvm.ir.assert_structural_equal(After, Expected)


@pytest.mark.gpu
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not enabled")
def test_scatter_nd_gpu():
    """scatter_nd lowered for GPU must build"""
    target = "cuda"

    @I.ir_module(s_tir=True)
    class Mod:
        @R.function
        def main(
            data: R.Tensor((4, 8), "float32"),
            indices: R.Tensor((3, 2), "int64"),
            updates: R.Tensor((3,), "float32"),
        ):
            with R.dataflow():
                lv = R.scatter_nd(data, indices, updates)
                gv = lv
                R.output(gv)
            return gv

    with tvm.target.Target(target):
        mod = LegalizeOps()(Mod)
    relax.build(mod, target=target)


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