# 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, topi
from tvm.script import ir as I
from tvm.script import relax as R
from tvm.script import tirx as T


def _check(mod_actual, mod_expected):
    mod_actual = relax.transform.AnnotateTIROpPattern()(mod_actual)
    mod_actual = relax.transform.FuseOps()(mod_actual)
    mod_expected = relax.transform.AnnotateTIROpPattern()(mod_expected)
    tvm.ir.assert_structural_equal(mod_actual, mod_expected)


def test_fuse_simple():
    """Simple testcase."""

    def before():
        bb = relax.BlockBuilder()
        x = relax.Var("x", R.Tensor([10, 20], "float32"))
        with bb.function("main", [x]):
            with bb.dataflow():
                lv0 = bb.emit_te(topi.add, x, relax.const(1, "float32"))
                lv1 = bb.emit_te(topi.exp, lv0)
                gv = bb.emit_output(bb.call_te(topi.squeeze, lv1))
            bb.emit_func_output(gv)

        return bb.get()

    def expected():
        bb = relax.BlockBuilder()
        x = relax.Var("x", R.Tensor([10, 20], "float32"))
        p0 = relax.Var("p0", R.Tensor((), "float32"))

        with bb.function("fused_add_exp_squeeze", [x, p0], attrs={"Primitive": True}, private=True):
            with bb.dataflow():
                lv0 = bb.emit_te(topi.add, x, p0)
                lv1 = bb.emit_te(topi.exp, lv0)
                gv = bb.emit_output(bb.call_te(topi.squeeze, lv1))
            bb.emit_func_output(gv)
        fused_add_exp_squeeze = bb.get().get_global_var("fused_add_exp_squeeze")

        x = relax.Var("x", R.Tensor([10, 20], "float32"))
        with bb.function("main", [x]):
            with bb.dataflow():
                gv = bb.emit_output(
                    relax.Call(fused_add_exp_squeeze, [x, relax.const(1, "float32")])
                )
            bb.emit_func_output(gv)

        return bb.get()

    _check(before(), expected())


def test_conv2d_fuse():
    """Test fusion case of conv2d"""

    def before(dtype):
        bb = relax.BlockBuilder()
        x = relax.Var("x", R.Tensor((1, 16, 64, 64), dtype))
        w1 = relax.Var("w1", R.Tensor((16, 16, 3, 3), dtype))
        w2 = relax.Var("w2", R.Tensor((16, 16, 1, 1), dtype))
        w3 = relax.Var("w3", R.Tensor((16, 16, 3, 3), dtype))
        with bb.function("main", [x, w1, w2, w3]):
            with bb.dataflow():
                lv0 = bb.emit_te(topi.add, x, relax.const(1, dtype))
                lv1 = bb.emit_te(topi.nn.conv2d, lv0, w1, strides=1, padding=1, dilation=1)
                # this is the next dominator.
                lv2 = bb.emit_te(topi.add, relax.const(1, dtype), lv1)
                lv3 = bb.emit_te(topi.add, lv1, lv2)
                # second path
                lv4 = bb.emit_te(topi.nn.conv2d, lv3, w2, strides=1, padding=0, dilation=1)
                lv5 = bb.emit_te(topi.nn.conv2d, lv3, w3, strides=1, padding=1, dilation=1)
                gv = bb.emit_output(bb.call_te(topi.add, lv4, lv5))
            bb.emit_func_output(gv)

        return bb.get()

    def expected(dtype):
        bb = relax.BlockBuilder()

        # Grouped function 1
        x = relax.Var("x", R.Tensor((1, 16, 64, 64), dtype))
        w = relax.Var("w", R.Tensor((16, 16, 3, 3), dtype))
        p0 = relax.Var("p0", R.Tensor((), dtype))
        with bb.function(
            "fused_conv2d_add1_add2", [x, w, p0], attrs={"Primitive": True}, private=True
        ):
            with bb.dataflow():
                lv0 = bb.emit_te(
                    topi.nn.conv2d,
                    x,
                    w,
                    strides=1,
                    padding=1,
                    dilation=1,
                    primfunc_name_hint="conv2d",
                )
                lv1 = bb.emit_te(topi.add, p0, lv0, primfunc_name_hint="add1")
                gv = bb.emit_output(bb.call_te(topi.add, lv0, lv1, primfunc_name_hint="add2"))
            bb.emit_func_output(gv)

        # Grouped function 2
        x = relax.Var("x", R.Tensor((1, 16, 64, 64), dtype))
        w = relax.Var("w", R.Tensor((16, 16, 1, 1), dtype))
        y = relax.Var("y", R.Tensor((1, 16, 64, 64), dtype))
        with bb.function("fused_conv2d1_add2", [x, w, y], attrs={"Primitive": True}, private=True):
            with bb.dataflow():
                lv0 = bb.emit_te(
                    topi.nn.conv2d,
                    x,
                    w,
                    strides=1,
                    padding=0,
                    dilation=1,
                    primfunc_name_hint="conv2d1",
                )
                gv = bb.emit_output(bb.call_te(topi.add, lv0, y, primfunc_name_hint="add2"))
            bb.emit_func_output(gv)

        # Get the global variables of the grouped functions
        mod = bb.get()
        fused_conv2d_add1_add2 = mod.get_global_var("fused_conv2d_add1_add2")
        fused_conv2d1_add2 = mod.get_global_var("fused_conv2d1_add2")

        # Main function
        x = relax.Var("x", R.Tensor((1, 16, 64, 64), dtype))
        w1 = relax.Var("w1", R.Tensor((16, 16, 3, 3), dtype))
        w2 = relax.Var("w2", R.Tensor((16, 16, 1, 1), dtype))
        w3 = relax.Var("w3", R.Tensor((16, 16, 3, 3), dtype))
        with bb.function("main", [x, w1, w2, w3]):
            with bb.dataflow():
                lv0 = bb.emit_te(topi.add, x, relax.const(1, dtype))
                lv1 = bb.emit(relax.Call(fused_conv2d_add1_add2, [lv0, w1, relax.const(1, dtype)]))
                lv2 = bb.emit_te(
                    topi.nn.conv2d,
                    lv1,
                    w3,
                    strides=1,
                    padding=1,
                    dilation=1,
                )
                gv = bb.emit_output(relax.Call(fused_conv2d1_add2, [lv1, w2, lv2]))
            bb.emit_func_output(gv)

        return bb.get()

    _check(before("float32"), expected("float32"))
    _check(before("float16"), expected("float16"))
    _check(before("int8"), expected("int8"))


def test_concatenate():
    """Test fusion case involving concat op and Tuple node"""

    def before():
        bb = relax.BlockBuilder()
        x = relax.Var("x", R.Tensor((1, 16, 64, 64), "float32"))
        with bb.function("main", [x]):
            with bb.dataflow():
                lv0 = bb.emit_te(
                    topi.nn.pool2d,
                    x,
                    kernel=(2, 2),
                    stride=(2, 2),
                    dilation=(1, 1),
                    padding=(0, 0, 0, 0),
                    pool_type="max",
                )
                lv1 = bb.emit_te(topi.nn.upsampling, lv0, scale_h=2.0, scale_w=2.0)
                lv2 = bb.emit_te(topi.concatenate, (lv1, x), axis=1)
                gv = bb.emit_output(bb.call_te(topi.add, lv2, relax.const(1, "float32")))
            bb.emit_func_output(gv)

        return bb.get()

    def expected():
        bb = relax.BlockBuilder()

        # Grouped function
        x = relax.Var("x", R.Tensor((1, 16, 64, 64), "float32"))
        w = relax.Var("w", R.Tensor((1, 16, 32, 32), "float32"))
        p0 = relax.Var("p0", R.Tensor((), "float32"))
        with bb.function(
            "fused_upsampling_concatenate_add", [w, x, p0], attrs={"Primitive": True}, private=True
        ):
            with bb.dataflow():
                lv0 = bb.emit_te(topi.nn.upsampling, w, scale_h=2.0, scale_w=2.0)
                lv1 = bb.emit_te(topi.concatenate, (lv0, x), axis=1)
                gv = bb.emit_output(bb.call_te(topi.add, lv1, p0))
            bb.emit_func_output(gv)

        # Get the global variables of the grouped functions
        fused_upsampling_concatenate_add = bb.get().get_global_var(
            "fused_upsampling_concatenate_add"
        )

        # Main function
        x = relax.Var("x", R.Tensor((1, 16, 64, 64), "float32"))
        with bb.function("main", [x]):
            with bb.dataflow():
                lv0 = bb.emit_te(
                    topi.nn.pool2d,
                    x,
                    kernel=(2, 2),
                    stride=(2, 2),
                    dilation=(1, 1),
                    padding=(0, 0, 0, 0),
                    pool_type="max",
                )
                gv = bb.emit_output(
                    relax.Call(
                        fused_upsampling_concatenate_add, (lv0, x, relax.const(1, "float32"))
                    )
                )
            bb.emit_func_output(gv)

        return bb.get()

    _check(before(), expected())


def test_tuple_root():
    """Test fusion case where Tuple node is the root in its group"""

    def before():
        bb = relax.BlockBuilder()
        x = relax.Var("x", R.Tensor((1, 16, 64, 64), "float32"))
        with bb.function("main", [x]):
            with bb.dataflow():
                lv0 = bb.emit_te(
                    topi.nn.pool2d,
                    x,
                    kernel=(2, 2),
                    stride=(2, 2),
                    dilation=(1, 1),
                    padding=(0, 0, 0, 0),
                    pool_type="max",
                )
                lv1 = bb.emit_te(topi.nn.upsampling, lv0, scale_h=2.0, scale_w=2.0)
                gv = bb.emit_output((lv1, x))
            bb.emit_func_output(gv)

        return bb.get()

    # The fusion is supposed to make no change.
    _check(before(), before())


def test_fuse_tuple_get_elemwise():
    def before(dim: int):
        bb = relax.BlockBuilder()
        x = relax.Var("x", R.Tensor((1, dim), "float32"))
        w = relax.Var("w", R.Tensor((3 * dim, dim), "float32"))
        with bb.function("main", [x, w]):
            with bb.dataflow():
                lv0 = bb.emit_te(topi.nn.dense, x, w)
                lv1 = bb.emit_te(topi.split, lv0, indices_or_sections=3, axis=1)
                lv2 = bb.emit(relax.TupleGetItem(lv1, 0))
                lv3 = bb.emit_te(topi.sigmoid, lv2)
                lv4 = bb.emit(relax.TupleGetItem(lv1, 1))
                lv5 = bb.emit_te(topi.tanh, lv4)
                lv6 = bb.emit(relax.TupleGetItem(lv1, 2))
                lv7 = bb.emit_te(topi.exp, lv6)
                lv8 = bb.emit_te(topi.multiply, lv5, lv7)
                gv = bb.emit_output(bb.call_te(topi.add, lv3, lv8))
            bb.emit_func_output(gv)

        return bb.get()

    def expected(dim: int):
        bb = relax.BlockBuilder()

        # Grouped function
        dense = relax.Var("dense", R.Tensor((1, 3 * dim), "float32"))
        with bb.function(
            "fused_split_sigmoid_tanh_exp_multiply_add",
            [dense],
            attrs={"Primitive": True},
            private=True,
        ):
            with bb.dataflow():
                lv0 = bb.emit_te(topi.split, dense, indices_or_sections=3, axis=1)
                lv1 = bb.emit(relax.TupleGetItem(lv0, 0))
                lv2 = bb.emit_te(topi.sigmoid, lv1)
                lv3 = bb.emit(relax.TupleGetItem(lv0, 1))
                lv4 = bb.emit_te(topi.tanh, lv3)
                lv5 = bb.emit(relax.TupleGetItem(lv0, 2))
                lv6 = bb.emit_te(topi.exp, lv5)
                lv7 = bb.emit_te(topi.multiply, lv4, lv6)
                gv = bb.emit_output(bb.call_te(topi.add, lv2, lv7))
            bb.emit_func_output(gv)

        # Get the global variables of the grouped functions
        fused_split_sigmoid_tanh_exp_multiply_add = bb.get().get_global_var(
            "fused_split_sigmoid_tanh_exp_multiply_add"
        )

        # Main function
        x = relax.Var("x", R.Tensor((1, dim), "float32"))
        w = relax.Var("w", R.Tensor((3 * dim, dim), "float32"))
        with bb.function("main", [x, w]):
            with bb.dataflow():
                lv0 = bb.emit_te(topi.nn.dense, x, w)
                gv = bb.emit_output(relax.Call(fused_split_sigmoid_tanh_exp_multiply_add, (lv0,)))
            bb.emit_func_output(gv)

        return bb.get()

    dim = 10
    _check(before(dim), expected(dim))


def test_tuple_get_root():
    def before(dim: int):
        bb = relax.BlockBuilder()
        x = relax.Var("x", R.Tensor((1, 3 * dim), "float32"))
        w = relax.Var("w", R.Tensor((dim, dim), "float32"))
        with bb.function("main", [x, w]):
            with bb.dataflow():
                lv0 = bb.emit_te(topi.split, x, indices_or_sections=3, axis=1)
                lv1 = bb.emit(relax.TupleGetItem(lv0, 0))
                gv = bb.emit_output(bb.call_te(topi.nn.dense, lv1, w))
            bb.emit_func_output(gv)

        return bb.get()

    def expected(dim: int):
        bb = relax.BlockBuilder()

        # Grouped function
        x = relax.Var("x", R.Tensor((1, 3 * dim), "float32"))
        with bb.function("fused_split", [x], attrs={"Primitive": True}, private=True):
            with bb.dataflow():
                lv0 = bb.emit_te(topi.split, x, indices_or_sections=3, axis=1)
                gv = bb.emit_output(relax.TupleGetItem(lv0, 0))
            bb.emit_func_output(gv)

        # Get the global variables of the grouped functions
        fused_split = bb.get().get_global_var("fused_split")

        # Main function
        x = relax.Var("x", R.Tensor((1, 3 * dim), "float32"))
        w = relax.Var("w", R.Tensor((dim, dim), "float32"))
        with bb.function("main", [x, w]):
            with bb.dataflow():
                lv0 = bb.emit(relax.Call(fused_split, (x,)))
                gv = bb.emit_output(bb.call_te(topi.nn.dense, lv0, w))
            bb.emit_func_output(gv)

        return bb.get()

    dim = 10
    _check(before(dim), expected(dim))


def test_tuple_intermediate():
    def before():
        bb = relax.BlockBuilder()

        x = relax.Var("x", R.Tensor((1, 16, 64, 64), "float32"))
        with bb.function("main", [x]):
            with bb.dataflow():
                lv0 = bb.emit_te(topi.squeeze, x)
                lv1 = bb.emit_te(topi.add, lv0, relax.const(1, "float32"))
                lv2 = bb.emit_te(topi.squeeze, lv0)
                lv3 = bb.emit_te(topi.add, lv2, relax.const(1, "float32"))
                lv4 = bb.emit_te(topi.add, lv3, relax.const(1, "float32"))
                lv5 = bb.emit_te(topi.add, lv0, relax.const(1, "float32"))
                lv6 = bb.emit_te(topi.concatenate, (lv1, lv4, lv5), axis=1)
                lv7 = bb.emit_te(topi.squeeze, lv6)
                gv = bb.emit_output(bb.call_te(topi.add, lv7, relax.const(1, "float32")))
            bb.emit_func_output(gv)

        return bb.get()

    def expected():
        bb = relax.BlockBuilder()

        # Grouped function
        x = relax.Var("x", R.Tensor((1, 16, 64, 64), "float32"))
        p0 = relax.Var("p0", R.Tensor((), "float32"))
        p1 = relax.Var("p1", R.Tensor((), "float32"))
        p2 = relax.Var("p2", R.Tensor((), "float32"))
        p3 = relax.Var("p3", R.Tensor((), "float32"))
        p4 = relax.Var("p4", R.Tensor((), "float32"))
        with bb.function(
            "fused_squeeze_add_squeeze1_add_add_add_concatenate_squeeze2_add1",
            [x, p0, p1, p2, p3, p4],
            attrs={"Primitive": True},
            private=True,
        ):
            with bb.dataflow():
                lv0 = bb.emit_te(topi.squeeze, x)
                lv1 = bb.emit_te(topi.add, lv0, p0)
                lv2 = bb.emit_te(topi.squeeze, lv0)
                lv3 = bb.emit_te(topi.add, lv2, p1)
                lv4 = bb.emit_te(topi.add, lv3, p2)
                lv5 = bb.emit_te(topi.add, lv0, p3)
                lv6 = bb.emit_te(topi.concatenate, (lv1, lv4, lv5), axis=1)
                lv7 = bb.emit_te(topi.squeeze, lv6)
                gv = bb.emit_output(bb.call_te(topi.add, lv7, p4))
            bb.emit_func_output(gv)

        # Get the global variables of the grouped functions
        fused_func = bb.get().get_global_var(
            "fused_squeeze_add_squeeze1_add_add_add_concatenate_squeeze2_add1"
        )

        # Main func
        x = relax.Var("x", R.Tensor((1, 16, 64, 64), "float32"))
        with bb.function("main", [x]):
            with bb.dataflow():
                gv = bb.emit_output(
                    relax.Call(
                        fused_func,
                        (
                            x,
                            relax.const(1, "float32"),
                            relax.const(1, "float32"),
                            relax.const(1, "float32"),
                            relax.const(1, "float32"),
                            relax.const(1, "float32"),
                        ),
                    )
                )
            bb.emit_func_output(gv)

        return bb.get()

    _check(before(), expected())


def test_tuple_consecutive():
    def before():
        bb = relax.BlockBuilder()

        x = relax.Var("x", R.Tensor((1, 16, 64, 64), "float32"))
        with bb.function("main", [x]):
            with bb.dataflow():
                lv0 = bb.emit_te(topi.add, x, relax.const(1, "float32"))
                lv1 = bb.emit_te(topi.add, x, relax.const(1, "float32"))
                lv2 = bb.emit_te(topi.add, x, relax.const(1, "float32"))
                lv3 = bb.emit_te(topi.concatenate, (lv0, lv1, lv2), axis=1)
                lv4 = bb.emit_te(topi.add, lv3, relax.const(1, "float32"))
                lv5 = bb.emit_te(topi.add, x, relax.const(1, "float32"))
                lv6 = bb.emit_te(topi.add, x, relax.const(1, "float32"))
                lv7 = bb.emit_te(topi.add, x, relax.const(1, "float32"))
                lv8 = bb.emit_te(topi.concatenate, (lv5, lv6, lv7), axis=1)
                lv9 = bb.emit_te(topi.add, lv8, relax.const(1, "float32"))
                lv10 = bb.emit_te(topi.add, x, relax.const(1, "float32"))
                lv11 = bb.emit_te(topi.add, x, relax.const(1, "float32"))
                lv12 = bb.emit_te(topi.add, x, relax.const(1, "float32"))
                lv13 = bb.emit_te(topi.concatenate, (lv10, lv11, lv12), axis=1)
                lv14 = bb.emit_te(topi.add, lv13, relax.const(1, "float32"))
                lv15 = bb.emit_te(topi.concatenate, (lv4, lv9, lv14), axis=1)
                lv16 = bb.emit_te(
                    topi.nn.pool2d,
                    lv15,
                    kernel=(2, 2),
                    stride=(2, 2),
                    dilation=(1, 1),
                    padding=(0, 0, 0, 0),
                    pool_type="max",
                )
                lv17 = bb.emit_te(topi.add, lv16, relax.const(1, "float32"))
                lv18 = bb.emit_te(topi.add, lv17, relax.const(1, "float32"))
                gv = bb.emit_output((lv17, lv18))
            bb.emit_func_output(gv)

        return bb.get()

    def expected():
        bb = relax.BlockBuilder()

        # Grouped function 1
        x = relax.Var("x", R.Tensor((1, 16, 64, 64), "float32"))
        p0 = relax.Var("p0", R.Tensor((), "float32"))
        p1 = relax.Var("p1", R.Tensor((), "float32"))
        p2 = relax.Var("p2", R.Tensor((), "float32"))
        p3 = relax.Var("p3", R.Tensor((), "float32"))
        p4 = relax.Var("p4", R.Tensor((), "float32"))
        p5 = relax.Var("p5", R.Tensor((), "float32"))
        p6 = relax.Var("p6", R.Tensor((), "float32"))
        p7 = relax.Var("p7", R.Tensor((), "float32"))
        p8 = relax.Var("p8", R.Tensor((), "float32"))
        p9 = relax.Var("p9", R.Tensor((), "float32"))
        p10 = relax.Var("p10", R.Tensor((), "float32"))
        p11 = relax.Var("p11", R.Tensor((), "float32"))
        with bb.function(
            "fused_add_add_add_concatenate_add1_add_add_add_concatenate_add1_add_add_add_concatenate_add1_concatenate1",
            [x, p0, p1, p2, p3, p4, p5, p6, p7, p8, p9, p10, p11],
            attrs={"Primitive": True},
            private=True,
        ):
            with bb.dataflow():
                lv0 = bb.emit_te(topi.add, x, p0)
                lv1 = bb.emit_te(topi.add, x, p1)
                lv2 = bb.emit_te(topi.add, x, p2)
                lv3 = bb.emit_te(topi.concatenate, (lv0, lv1, lv2), axis=1)
                lv4 = bb.emit_te(topi.add, lv3, p3)
                lv5 = bb.emit_te(topi.add, x, p4)
                lv6 = bb.emit_te(topi.add, x, p5)
                lv7 = bb.emit_te(topi.add, x, p6)
                lv8 = bb.emit_te(topi.concatenate, (lv5, lv6, lv7), axis=1)
                lv9 = bb.emit_te(topi.add, lv8, p7)
                lv10 = bb.emit_te(topi.add, x, p8)
                lv11 = bb.emit_te(topi.add, x, p9)
                lv12 = bb.emit_te(topi.add, x, p10)
                lv13 = bb.emit_te(topi.concatenate, (lv10, lv11, lv12), axis=1)
                lv14 = bb.emit_te(topi.add, lv13, p11)
                gv = bb.emit_output(bb.call_te(topi.concatenate, (lv4, lv9, lv14), axis=1))
            bb.emit_func_output(gv)

        # Grouped function 2
        concat = relax.Var("concat", R.Tensor((1, 144, 64, 64), "float32"))
        p0 = relax.Var("p0", R.Tensor((), "float32"))
        with bb.function(
            "fused_pool2d_add2", [concat, p0], attrs={"Primitive": True}, private=True
        ):
            with bb.dataflow():
                lv0 = bb.emit_te(
                    topi.nn.pool2d,
                    concat,
                    kernel=(2, 2),
                    stride=(2, 2),
                    dilation=(1, 1),
                    padding=(0, 0, 0, 0),
                    pool_type="max",
                )
                gv = bb.emit_output(bb.call_te(topi.add, lv0, p0))
            bb.emit_func_output(gv)

        # Get the global variables of the grouped functions
        mod = bb.get()
        fused_func1 = mod.get_global_var(
            "fused_add_add_add_concatenate_add1_add_add_add_concatenate_add1_add_add_add_concatenate_add1_concatenate1"
        )
        fused_func2 = mod.get_global_var("fused_pool2d_add2")

        # Main function
        x = relax.Var("x", R.Tensor((1, 16, 64, 64), "float32"))
        with bb.function("main", [x]):
            with bb.dataflow():
                lv0 = bb.emit(
                    relax.Call(
                        fused_func1,
                        (
                            x,
                            relax.const(1, "float32"),
                            relax.const(1, "float32"),
                            relax.const(1, "float32"),
                            relax.const(1, "float32"),
                            relax.const(1, "float32"),
                            relax.const(1, "float32"),
                            relax.const(1, "float32"),
                            relax.const(1, "float32"),
                            relax.const(1, "float32"),
                            relax.const(1, "float32"),
                            relax.const(1, "float32"),
                            relax.const(1, "float32"),
                        ),
                    )
                )
                lv1 = bb.emit(relax.Call(fused_func2, (lv0, relax.const(1, "float32"))))
                lv2 = bb.emit_te(topi.add, lv1, relax.const(1, "float32"))
                gv = bb.emit_output((lv1, lv2))
            bb.emit_func_output(gv)

        return bb.get()

    _check(before(), expected())


def test_inception_like():
    def before():
        bb = relax.BlockBuilder()

        x = relax.Var("x", R.Tensor((1, 16, 64, 64), "float32"))
        w0 = relax.Var("w0", R.Tensor((16, 16, 3, 3), "float32"))
        w1 = relax.Var("w1", R.Tensor((16, 16, 3, 3), "float32"))
        w2 = relax.Var("w2", R.Tensor((16, 32, 3, 3), "float32"))
        w3 = relax.Var("w3", R.Tensor((16, 32, 3, 3), "float32"))
        with bb.function("main", [x, w0, w1, w2, w3]):
            with bb.dataflow():
                lv0 = bb.emit_te(topi.nn.conv2d, x, w0, strides=1, padding=1, dilation=1)
                lv1 = bb.emit_te(topi.nn.relu, lv0)
                lv2 = bb.emit_te(topi.nn.conv2d, x, w1, strides=1, padding=1, dilation=1)
                lv3 = bb.emit_te(topi.nn.relu, lv2)
                lv4 = bb.emit_te(topi.concatenate, (lv1, lv3), axis=1)
                lv5 = bb.emit_te(topi.nn.conv2d, lv4, w2, strides=1, padding=1, dilation=1)
                lv6 = bb.emit_te(topi.nn.relu, lv5)
                lv7 = bb.emit_te(topi.nn.conv2d, lv4, w3, strides=1, padding=1, dilation=1)
                lv8 = bb.emit_te(topi.nn.relu, lv7)
                gv = bb.emit_output(bb.call_te(topi.concatenate, (lv6, lv8), axis=1))
            bb.emit_func_output(gv)

        return bb.get()

    def expected():
        bb = relax.BlockBuilder()

        # Grouped function 1
        x = relax.Var("x", R.Tensor((1, 16, 64, 64), "float32"))
        w = relax.Var("w", R.Tensor((16, 16, 3, 3), "float32"))
        with bb.function("fused_conv2d_relu", [x, w], attrs={"Primitive": True}, private=True):
            with bb.dataflow():
                lv0 = bb.emit_te(
                    topi.nn.conv2d,
                    x,
                    w,
                    strides=1,
                    padding=1,
                    dilation=1,
                    primfunc_name_hint="conv2d",
                )
                gv = bb.emit_output(bb.call_te(topi.nn.relu, lv0))
            bb.emit_func_output(gv)

        # Grouped function 2
        x = relax.Var("x", R.Tensor((1, 32, 64, 64), "float32"))
        w = relax.Var("w", R.Tensor((16, 32, 3, 3), "float32"))
        with bb.function("fused_conv2d1_relu", [x, w], attrs={"Primitive": True}, private=True):
            with bb.dataflow():
                lv0 = bb.emit_te(
                    topi.nn.conv2d,
                    x,
                    w,
                    strides=1,
                    padding=1,
                    dilation=1,
                    primfunc_name_hint="conv2d1",
                )
                gv = bb.emit_output(bb.call_te(topi.nn.relu, lv0))
            bb.emit_func_output(gv)

        # Get the global variables of the grouped functions
        mod = bb.get()
        fused_conv2d_relu1 = mod.get_global_var("fused_conv2d_relu")
        fused_conv2d_relu2 = mod.get_global_var("fused_conv2d1_relu")

        # Main function
        x = relax.Var("x", R.Tensor((1, 16, 64, 64), "float32"))
        w0 = relax.Var("w0", R.Tensor((16, 16, 3, 3), "float32"))
        w1 = relax.Var("w1", R.Tensor((16, 16, 3, 3), "float32"))
        w2 = relax.Var("w2", R.Tensor((16, 32, 3, 3), "float32"))
        w3 = relax.Var("w3", R.Tensor((16, 32, 3, 3), "float32"))
        with bb.function("main", [x, w0, w1, w2, w3]):
            with bb.dataflow():
                lv0 = bb.emit(relax.Call(fused_conv2d_relu1, (x, w0)))
                lv1 = bb.emit(relax.Call(fused_conv2d_relu1, (x, w1)))
                lv2 = bb.emit_te(topi.concatenate, (lv0, lv1), axis=1)
                lv3 = bb.emit(relax.Call(fused_conv2d_relu2, (lv2, w2)))
                lv4 = bb.emit(relax.Call(fused_conv2d_relu2, (lv2, w3)))
                gv = bb.emit_output(bb.call_te(topi.concatenate, (lv3, lv4), axis=1))
            bb.emit_func_output(gv)

        return bb.get()

    _check(before(), expected())


def test_fuse_parallel_injective():
    def before():
        bb = relax.BlockBuilder()

        x = relax.Var("x", R.Tensor((10, 20), "int32"))
        with bb.function("main", [x]):
            with bb.dataflow():
                lv0 = bb.emit_te(topi.add, x, relax.const(1, "int32"))
                lv1 = bb.emit_te(topi.squeeze, lv0)
                lv2 = bb.emit_te(topi.transpose, lv0, axes=[1, 0])
                lv3 = bb.emit_te(topi.transpose, lv2, axes=[1, 0])
                gv = bb.emit_output(bb.call_te(topi.left_shift, lv1, lv3))
            bb.emit_func_output(gv)

        return bb.get()

    def expected():
        bb = relax.BlockBuilder()

        # Grouped function
        x = relax.Var("x", R.Tensor((10, 20), "int32"))
        p0 = relax.Var("p0", R.Tensor((), "int32"))
        with bb.function(
            "fused_add_squeeze_transpose_transpose1_left_shift",
            [x, p0],
            attrs={"Primitive": True},
            private=True,
        ):
            with bb.dataflow():
                lv0 = bb.emit_te(topi.add, x, p0)
                lv1 = bb.emit_te(topi.squeeze, lv0)
                lv2 = bb.emit_te(topi.transpose, lv0, axes=[1, 0])
                lv3 = bb.emit_te(topi.transpose, lv2, axes=[1, 0], primfunc_name_hint="transpose1")
                gv = bb.emit_output(bb.call_te(topi.left_shift, lv1, lv3))
            bb.emit_func_output(gv)

        # Get the global variables of the grouped functions
        fused_func = bb.get().get_global_var("fused_add_squeeze_transpose_transpose1_left_shift")

        # Main function
        x = relax.Var("x", R.Tensor((10, 20), "int32"))
        with bb.function("main", [x]):
            with bb.dataflow():
                gv = bb.emit_output(relax.Call(fused_func, (x, relax.const(1, "int32"))))
            bb.emit_func_output(gv)

        return bb.get()

    _check(before(), expected())


def test_softmax():
    """Test if softmax can be fused with following ops."""

    def before():
        bb = relax.BlockBuilder()

        x = relax.Var("x", R.Tensor((16, 16), "float32"))
        with bb.function("main", [x]):
            with bb.dataflow():
                lv0 = bb.emit_te(topi.nn.softmax, x)
                gv = bb.emit_output(bb.call_te(topi.cast, lv0, dtype="float16"))
            bb.emit_func_output(gv)

        return bb.get()

    def expected():
        bb = relax.BlockBuilder()

        # Grouped function
        x = relax.Var("x", R.Tensor((16, 16), "float32"))
        with bb.function("fused_softmax_cast", [x], attrs={"Primitive": True}, private=True):
            with bb.dataflow():
                lv0 = bb.emit_te(topi.nn.softmax, x)
                gv = bb.emit_output(bb.call_te(topi.cast, lv0, dtype="float16"))
            bb.emit_func_output(gv)

        # Get the global variables of the grouped functions
        fused_func = bb.get().get_global_var("fused_softmax_cast")

        # Main function
        x = relax.Var("x", R.Tensor((16, 16), "float32"))
        with bb.function("main", [x]):
            with bb.dataflow():
                gv = bb.emit_output(relax.Call(fused_func, (x,)))
            bb.emit_func_output(gv)

        return bb.get()

    _check(before(), expected())


def test_multiple_relax_functions():
    def before():
        bb = relax.BlockBuilder()
        x = relax.Var("x", R.Tensor([10, 20], "float32"))
        with bb.function("func1", [x]):
            with bb.dataflow():
                lv0 = bb.emit_te(topi.add, x, relax.const(1, "float32"))
                lv1 = bb.emit_te(topi.exp, lv0)
                gv = bb.emit_output(bb.call_te(topi.squeeze, lv1))
            bb.emit_func_output(gv)

        x = relax.Var("x", R.Tensor([20, 10], "float32"))
        with bb.function("func2", [x]):
            with bb.dataflow():
                lv0 = bb.emit_te(topi.add, x, relax.const(1, "float32"))
                lv1 = bb.emit_te(topi.exp, lv0)
                gv = bb.emit_output(bb.call_te(topi.squeeze, lv1))
            bb.emit_func_output(gv)

        return bb.get()

    def expected():
        bb = relax.BlockBuilder()

        x = relax.Var("x", R.Tensor([10, 20], "float32"))
        p0 = relax.Var("p0", R.Tensor((), "float32"))
        with bb.function("fused_add_exp_squeeze", [x, p0], attrs={"Primitive": True}, private=True):
            with bb.dataflow():
                lv0 = bb.emit_te(topi.add, x, p0)
                lv1 = bb.emit_te(topi.exp, lv0)
                gv = bb.emit_output(bb.call_te(topi.squeeze, lv1))
            bb.emit_func_output(gv)
        fused_add_exp_squeeze = bb.get().get_global_var("fused_add_exp_squeeze")

        x = relax.Var("x", R.Tensor([20, 10], "float32"))
        p0 = relax.Var("p0", R.Tensor((), "float32"))
        with bb.function(
            "fused_add1_exp1_squeeze1", [x, p0], attrs={"Primitive": True}, private=True
        ):
            with bb.dataflow():
                lv0 = bb.emit_te(topi.add, x, p0)
                lv1 = bb.emit_te(topi.exp, lv0)
                gv = bb.emit_output(bb.call_te(topi.squeeze, lv1))
            bb.emit_func_output(gv)
        fused_add1_exp1_squeeze1 = bb.get().get_global_var("fused_add1_exp1_squeeze1")

        x = relax.Var("x", R.Tensor([10, 20], "float32"))
        with bb.function("func1", [x]):
            with bb.dataflow():
                gv = bb.emit_output(
                    relax.Call(fused_add_exp_squeeze, [x, relax.const(1, "float32")])
                )
            bb.emit_func_output(gv)

        x = relax.Var("x", R.Tensor([20, 10], "float32"))
        with bb.function("func2", [x]):
            with bb.dataflow():
                gv = bb.emit_output(
                    relax.Call(fused_add1_exp1_squeeze1, [x, relax.const(1, "float32")])
                )
            bb.emit_func_output(gv)

        return bb.get()

    _check(before(), expected())


def test_skip_call_dps_packed():
    @I.ir_module(s_tir=True)
    class Module:
        @R.function
        def main(x: R.Tensor((2, 3), "float32")):
            with R.dataflow():
                y = R.call_dps_packed("func_packed_dps", x, R.Tensor((2, 3), "float32"))
                R.output(y)
            return y

    # FuseOps should does no change to it.
    _check(Module, Module)


def test_edge_with_call_dps_packed():
    @I.ir_module(s_tir=True)
    class Module:
        @R.function
        def main(x: R.Tensor((2, 3), "float32")):
            cls = Module
            with R.dataflow():
                a = R.call_tir(cls.exp, (x,), out_ty=R.Tensor((2, 3), "float32"))
                b = R.call_tir(cls.exp, (a,), out_ty=R.Tensor((2, 3), "float32"))
                c = R.call_dps_packed("packed_dps", (a,), out_ty=R.Tensor((2, 3), "float32"))
                R.output(b, c)
            return R.tuple(b, c)

        @T.prim_func(private=True, s_tir=True)
        def exp(A: T.Buffer((2, 3), "float32"), B: T.Buffer((2, 3), "float32")):
            T.evaluate(0)

    # FuseOps should does no change to it.
    _check(Module, Module)


def test_layer_norm_silu():
    # fmt: off
    @I.ir_module(s_tir=True)
    class Module:
        @R.function
        def main(x: R.Tensor((1, 512, 64, 64), "float32"), mean: R.Tensor((64, 64), "float32"), var: R.Tensor((64, 64), "float32")):
            cls = Module
            with R.dataflow():
                gv0 = R.call_tir(cls.layer_norm, (x, mean, var), out_ty=R.Tensor((1, 512, 64, 64), 'float32'))
                gv1 = R.call_tir(cls.relu, gv0, out_ty=R.Tensor((1, 512, 64, 64), "float32"))
                R.output(gv1)
            return gv1

        @T.prim_func(private=True, s_tir=True)
        def layer_norm(A: T.Buffer((T.int64(1), T.int64(512), T.int64(64), T.int64(64)), "float32"), gamma: T.Buffer((T.int64(64), T.int64(64)), "float32"), beta: T.Buffer((T.int64(64), T.int64(64)), "float32"), T_layer_norm: T.Buffer((T.int64(1), T.int64(512), T.int64(64), T.int64(64)), "float32")):
            rxplaceholder_red_temp_v0 = T.sblock_alloc_buffer([T.int64(64), T.int64(64)], dtype="float32")
            rxplaceholder_red_temp_v1 = T.sblock_alloc_buffer([T.int64(64), T.int64(64)], dtype="float32")
            for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(512), T.int64(64), T.int64(64)):
                with T.sblock("rxplaceholder_red_temp"):
                    ax0, ax1, k2, k3 = T.axis.remap("SSRR", [i0, i1, i2, i3])
                    T.reads(A[ax0, ax1, k2, k3])
                    T.writes(rxplaceholder_red_temp_v0[ax0, ax1], rxplaceholder_red_temp_v1[ax0, ax1])
                    with T.init():
                        rxplaceholder_red_temp_v0[ax0, ax1] = T.float32(0)
                        rxplaceholder_red_temp_v1[ax0, ax1] = T.float32(0)
                    v_rxplaceholder_red_temp_v0: T.let[T.float32] = rxplaceholder_red_temp_v0[ax0, ax1] + A[ax0, ax1, k2, k3]
                    v_rxplaceholder_red_temp_v1: T.let[T.float32] = rxplaceholder_red_temp_v1[ax0, ax1] + A[ax0, ax1, k2, k3] * A[ax0, ax1, k2, k3]
                    rxplaceholder_red_temp_v0[ax0, ax1] = v_rxplaceholder_red_temp_v0
                    rxplaceholder_red_temp_v1[ax0, ax1] = v_rxplaceholder_red_temp_v1
            for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(512), T.int64(64), T.int64(64)):
                with T.sblock("T_layer_norm"):
                    ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
                    T.reads(A[ax0, ax1, ax2, ax3], rxplaceholder_red_temp_v0[ax0, ax1], rxplaceholder_red_temp_v1[ax0, ax1], gamma[ax2, ax3], beta[ax2, ax3])
                    T.writes(T_layer_norm[ax0, ax1, ax2, ax3])
                    T_layer_norm[ax0, ax1, ax2, ax3] = (A[ax0, ax1, ax2, ax3] - rxplaceholder_red_temp_v0[ax0, ax1] * T.float32(0.05)) * T.rsqrt(rxplaceholder_red_temp_v1[ax0, ax1] * T.float32(0.05) - rxplaceholder_red_temp_v0[ax0, ax1] * T.float32(0.05) * (rxplaceholder_red_temp_v0[ax0, ax1] * T.float32(0.05)) + T.float32(1e-05), dtype="float32") * gamma[ax2, ax3] + beta[ax2, ax3]

        @T.prim_func(private=True, s_tir=True)
        def relu(A: T.Buffer((T.int64(1), T.int64(512), T.int64(64), T.int64(64)), "float32"), B: T.Buffer((T.int64(1), T.int64(512), T.int64(64), T.int64(64)), "float32")):
            for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(512), T.int64(64), T.int64(64)):
                with T.sblock("relu"):
                    v_i0, v_i1, v_i2, v_i3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
                    T.reads(A[v_i0, v_i1, v_i2, v_i3])
                    T.writes(B[v_i0, v_i1, v_i2, v_i3])
                    B[v_i0, v_i1, v_i2, v_i3] = T.max(A[v_i0, v_i1, v_i2, v_i3], T.float32(0))

    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(private=True, s_tir=True)
        def layer_norm(A: T.Buffer((T.int64(1), T.int64(512), T.int64(64), T.int64(64)), "float32"), gamma: T.Buffer((T.int64(64), T.int64(64)), "float32"), beta: T.Buffer((T.int64(64), T.int64(64)), "float32"), T_layer_norm: T.Buffer((T.int64(1), T.int64(512), T.int64(64), T.int64(64)), "float32")):
            T.func_attr({"op_pattern": 4})
            # with T.sblock("root"):
            rxplaceholder_red_temp_v0 = T.sblock_alloc_buffer((T.int64(64), T.int64(64)))
            rxplaceholder_red_temp_v1 = T.sblock_alloc_buffer((T.int64(64), T.int64(64)))
            for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(512), T.int64(64), T.int64(64)):
                with T.sblock("rxplaceholder_red_temp"):
                    ax0, ax1, k2, k3 = T.axis.remap("SSRR", [i0, i1, i2, i3])
                    T.reads(A[ax0, ax1, k2, k3])
                    T.writes(rxplaceholder_red_temp_v0[ax0, ax1], rxplaceholder_red_temp_v1[ax0, ax1])
                    with T.init():
                        rxplaceholder_red_temp_v0[ax0, ax1] = T.float32(0)
                        rxplaceholder_red_temp_v1[ax0, ax1] = T.float32(0)
                    v_rxplaceholder_red_temp_v0: T.let[T.float32] = rxplaceholder_red_temp_v0[ax0, ax1] + A[ax0, ax1, k2, k3]
                    v_rxplaceholder_red_temp_v1: T.let[T.float32] = rxplaceholder_red_temp_v1[ax0, ax1] + A[ax0, ax1, k2, k3] * A[ax0, ax1, k2, k3]
                    rxplaceholder_red_temp_v0[ax0, ax1] = v_rxplaceholder_red_temp_v0
                    rxplaceholder_red_temp_v1[ax0, ax1] = v_rxplaceholder_red_temp_v1
            for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(512), T.int64(64), T.int64(64)):
                with T.sblock("T_layer_norm"):
                    ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
                    T.reads(A[ax0, ax1, ax2, ax3], rxplaceholder_red_temp_v0[ax0, ax1], rxplaceholder_red_temp_v1[ax0, ax1], gamma[ax2, ax3], beta[ax2, ax3])
                    T.writes(T_layer_norm[ax0, ax1, ax2, ax3])
                    T_layer_norm[ax0, ax1, ax2, ax3] = (A[ax0, ax1, ax2, ax3] - rxplaceholder_red_temp_v0[ax0, ax1] * T.float32(0.050000000000000003)) * T.rsqrt(rxplaceholder_red_temp_v1[ax0, ax1] * T.float32(0.050000000000000003) - rxplaceholder_red_temp_v0[ax0, ax1] * T.float32(0.050000000000000003) * (rxplaceholder_red_temp_v0[ax0, ax1] * T.float32(0.050000000000000003)) + T.float32(1.0000000000000001e-05)) * gamma[ax2, ax3] + beta[ax2, ax3]

        @T.prim_func(private=True, s_tir=True)
        def relu(A: T.Buffer((T.int64(1), T.int64(512), T.int64(64), T.int64(64)), "float32"), B: T.Buffer((T.int64(1), T.int64(512), T.int64(64), T.int64(64)), "float32")):
            T.func_attr({"op_pattern": 0})
            # with T.sblock("root"):
            for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(512), T.int64(64), T.int64(64)):
                with T.sblock("relu"):
                    v_i0, v_i1, v_i2, v_i3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
                    T.reads(A[v_i0, v_i1, v_i2, v_i3])
                    T.writes(B[v_i0, v_i1, v_i2, v_i3])
                    B[v_i0, v_i1, v_i2, v_i3] = T.max(A[v_i0, v_i1, v_i2, v_i3], T.float32(0))

        @R.function(private=True)
        def fused_layer_norm_relu(x: R.Tensor((1, 512, 64, 64), dtype="float32"), mean: R.Tensor((64, 64), dtype="float32"), var: R.Tensor((64, 64), dtype="float32")) -> R.Tensor((1, 512, 64, 64), dtype="float32"):
            R.func_attr({"Primitive": True})
            cls = Expected
            with R.dataflow():
                gv0 = R.call_tir(cls.layer_norm, (x, mean, var), out_ty=R.Tensor((1, 512, 64, 64), 'float32'))
                gv = R.call_tir(cls.relu, (gv0,), out_ty=R.Tensor((1, 512, 64, 64), dtype="float32"))
                R.output(gv)
            return gv

        @R.function
        def main(x: R.Tensor((1, 512, 64, 64), dtype="float32"), mean: R.Tensor((64, 64), dtype="float32"), var: R.Tensor((64, 64), dtype="float32")) -> R.Tensor((1, 512, 64, 64), dtype="float32"):
            cls = Expected
            with R.dataflow():
                gv: R.Tensor((1, 512, 64, 64), dtype="float32") = cls.fused_layer_norm_relu(x, mean, var)
                R.output(gv)
            return gv
    # fmt: on

    _check(Module, Expected)


def test_multiple_paths():
    # fmt: off
    @I.ir_module(s_tir=True)
    class Module:
        @R.function
        def main(
            inp_0: R.Tensor((2, 320, 64, 64), dtype="float32"),
            inp_1: R.Tensor((2, 1280), dtype="float32"),
            w1: R.Tensor((320, 320, 3, 3), dtype="float32"),
            b1: R.Tensor((320,), "float32"),
            w2: R.Tensor((320, 1280), "float32"),
            b2: R.Tensor((320,), "float32"),
        ):
            R.func_attr({"num_input": 2})
            with R.dataflow():
                lv27: R.Tensor((2, 320, 64, 64), dtype="float32") = R.nn.conv2d(inp_0, w1, strides=[1, 1], padding=[1, 1, 1, 1], dilation=[1, 1], groups=1, data_layout="NCHW", kernel_layout="OIHW", out_layout="NCHW", out_dtype="float32")
                lv28: R.Tensor((1, 320, 1, 1), dtype="float32") = R.reshape(b1, R.shape([1, 320, 1, 1]))  ##
                lv29: R.Tensor((2, 320, 64, 64), dtype="float32") = R.add(lv27, lv28)
                lv31: R.Tensor((1280, 320), dtype="float32") = R.permute_dims(w2, axes=None)  ##
                lv32: R.Tensor((2, 320), dtype="float32") = R.matmul(inp_1, lv31, out_dtype="float32")
                lv33: R.Tensor((2, 320), dtype="float32") = R.add(lv32, b2)
                lv35: R.Tensor((2, 320, 1, 1), dtype="float32") = R.reshape(lv33, R.shape([2, 320, 1, 1]))
                lv36: R.Tensor((2, 320, 64, 64), dtype="float32") = R.add(lv29, lv35)
                gv = lv36
                R.output(gv)
            return gv

    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(private=True, s_tir=True)
        def add(rxplaceholder: T.Buffer((T.int64(2), T.int64(320), T.int64(64), T.int64(64)), "float32"), rxplaceholder_1: T.Buffer((T.int64(1), T.int64(320), T.int64(1), T.int64(1)), "float32"), T_add: T.Buffer((T.int64(2), T.int64(320), T.int64(64), T.int64(64)), "float32")):
            T.func_attr({"op_pattern": 0, "tirx.noalias": True})
            for ax0, ax1, ax2, ax3 in T.grid(T.int64(2), T.int64(320), T.int64(64), T.int64(64)):
                with T.sblock("T_add"):
                    v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
                    T.reads(rxplaceholder[v_ax0, v_ax1, v_ax2, v_ax3], rxplaceholder_1[T.int64(0), v_ax1, T.int64(0), T.int64(0)])
                    T.writes(T_add[v_ax0, v_ax1, v_ax2, v_ax3])
                    T_add[v_ax0, v_ax1, v_ax2, v_ax3] = rxplaceholder[v_ax0, v_ax1, v_ax2, v_ax3] + rxplaceholder_1[T.int64(0), v_ax1, T.int64(0), T.int64(0)]

        @T.prim_func(private=True, s_tir=True)
        def add1(rxplaceholder: T.Buffer((T.int64(2), T.int64(320)), "float32"), rxplaceholder_1: T.Buffer((T.int64(320),), "float32"), T_add: T.Buffer((T.int64(2), T.int64(320)), "float32")):
            T.func_attr({"op_pattern": 0, "tirx.noalias": True})
            for ax0, ax1 in T.grid(T.int64(2), T.int64(320)):
                with T.sblock("T_add"):
                    v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
                    T.reads(rxplaceholder[v_ax0, v_ax1], rxplaceholder_1[v_ax1])
                    T.writes(T_add[v_ax0, v_ax1])
                    T_add[v_ax0, v_ax1] = rxplaceholder[v_ax0, v_ax1] + rxplaceholder_1[v_ax1]

        @T.prim_func(private=True, s_tir=True)
        def add2(rxplaceholder: T.Buffer((T.int64(2), T.int64(320), T.int64(64), T.int64(64)), "float32"), rxplaceholder_1: T.Buffer((T.int64(2), T.int64(320), T.int64(1), T.int64(1)), "float32"), T_add: T.Buffer((T.int64(2), T.int64(320), T.int64(64), T.int64(64)), "float32")):
            T.func_attr({"op_pattern": 0, "tirx.noalias": True})
            for ax0, ax1, ax2, ax3 in T.grid(T.int64(2), T.int64(320), T.int64(64), T.int64(64)):
                with T.sblock("T_add"):
                    v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
                    T.reads(rxplaceholder[v_ax0, v_ax1, v_ax2, v_ax3], rxplaceholder_1[v_ax0, v_ax1, T.int64(0), T.int64(0)])
                    T.writes(T_add[v_ax0, v_ax1, v_ax2, v_ax3])
                    T_add[v_ax0, v_ax1, v_ax2, v_ax3] = rxplaceholder[v_ax0, v_ax1, v_ax2, v_ax3] + rxplaceholder_1[v_ax0, v_ax1, T.int64(0), T.int64(0)]

        @T.prim_func(private=True, s_tir=True)
        def conv2d(rxplaceholder: T.Buffer((T.int64(2), T.int64(320), T.int64(64), T.int64(64)), "float32"), rxplaceholder_1: T.Buffer((T.int64(320), T.int64(320), T.int64(3), T.int64(3)), "float32"), conv2d_nchw: T.Buffer((T.int64(2), T.int64(320), T.int64(64), T.int64(64)), "float32")):
            T.func_attr({"op_pattern": 4, "tirx.noalias": True})
            pad_temp = T.sblock_alloc_buffer((T.int64(2), T.int64(320), T.int64(66), T.int64(66)))
            for i0, i1, i2, i3 in T.grid(T.int64(2), T.int64(320), T.int64(66), T.int64(66)):
                with T.sblock("pad_temp"):
                    v_i0, v_i1, v_i2, v_i3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
                    T.reads(rxplaceholder[v_i0, v_i1, v_i2 - T.int64(1), v_i3 - T.int64(1)])
                    T.writes(pad_temp[v_i0, v_i1, v_i2, v_i3])
                    pad_temp[v_i0, v_i1, v_i2, v_i3] = T.if_then_else(T.int64(1) <= v_i2 and v_i2 < T.int64(65) and T.int64(1) <= v_i3 and v_i3 < T.int64(65), rxplaceholder[v_i0, v_i1, v_i2 - T.int64(1), v_i3 - T.int64(1)], T.float32(0))
            for nn, ff, yy, xx, rc, ry, rx in T.grid(T.int64(2), T.int64(320), T.int64(64), T.int64(64), T.int64(320), T.int64(3), T.int64(3)):
                with T.sblock("conv2d_nchw"):
                    v_nn, v_ff, v_yy, v_xx, v_rc, v_ry, v_rx = T.axis.remap("SSSSRRR", [nn, ff, yy, xx, rc, ry, rx])
                    T.reads(pad_temp[v_nn, v_rc, v_yy + v_ry, v_xx + v_rx], rxplaceholder_1[v_ff, v_rc, v_ry, v_rx])
                    T.writes(conv2d_nchw[v_nn, v_ff, v_yy, v_xx])
                    with T.init():
                        conv2d_nchw[v_nn, v_ff, v_yy, v_xx] = T.float32(0)
                    conv2d_nchw[v_nn, v_ff, v_yy, v_xx] = conv2d_nchw[v_nn, v_ff, v_yy, v_xx] + pad_temp[v_nn, v_rc, v_yy + v_ry, v_xx + v_rx] * rxplaceholder_1[v_ff, v_rc, v_ry, v_rx]

        @T.prim_func(private=True, s_tir=True)
        def matmul(rxplaceholder: T.Buffer((T.int64(2), T.int64(1280)), "float32"), rxplaceholder_1: T.Buffer((T.int64(1280), T.int64(320)), "float32"), matmul: T.Buffer((T.int64(2), T.int64(320)), "float32")):
            T.func_attr({"op_pattern": 4, "tirx.noalias": True})
            for i0, i1, k in T.grid(T.int64(2), T.int64(320), T.int64(1280)):
                with T.sblock("matmul"):
                    v_i0, v_i1, v_k = T.axis.remap("SSR", [i0, i1, k])
                    T.reads(rxplaceholder[v_i0, v_k], rxplaceholder_1[v_k, v_i1])
                    T.writes(matmul[v_i0, v_i1])
                    with T.init():
                        matmul[v_i0, v_i1] = T.float32(0)
                    matmul[v_i0, v_i1] = matmul[v_i0, v_i1] + rxplaceholder[v_i0, v_k] * rxplaceholder_1[v_k, v_i1]

        @T.prim_func(private=True, s_tir=True)
        def reshape(rxplaceholder: T.Buffer((T.int64(320),), "float32"), T_reshape: T.Buffer((T.int64(1), T.int64(320), T.int64(1), T.int64(1)), "float32")):
            T.func_attr({"op_pattern": 2, "tirx.noalias": True})
            for ax0, ax1, ax2, ax3 in T.grid(T.int64(1), T.int64(320), T.int64(1), T.int64(1)):
                with T.sblock("T_reshape"):
                    v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
                    T.reads(rxplaceholder[(v_ax1 + v_ax2 + 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_ax1 + v_ax2 + v_ax3) % T.int64(320)]

        @T.prim_func(private=True, s_tir=True)
        def reshape1(rxplaceholder: T.Buffer((T.int64(2), T.int64(320)), "float32"), T_reshape: T.Buffer((T.int64(2), T.int64(320), T.int64(1), T.int64(1)), "float32")):
            T.func_attr({"op_pattern": 2, "tirx.noalias": True})
            for ax0, ax1, ax2, ax3 in T.grid(T.int64(2), T.int64(320), T.int64(1), T.int64(1)):
                with T.sblock("T_reshape"):
                    v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
                    T.reads(rxplaceholder[((v_ax1 + v_ax2 + v_ax3) // T.int64(320) + v_ax0) % T.int64(2), (v_ax1 + v_ax2 + 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_ax1 + v_ax2 + v_ax3) // T.int64(320) + v_ax0) % T.int64(2), (v_ax1 + v_ax2 + v_ax3) % T.int64(320)]

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

        @R.function(private=True)
        def fused_conv2d_add_add2(inp_0: R.Tensor((2, 320, 64, 64), dtype="float32"), w1: R.Tensor((320, 320, 3, 3), dtype="float32"), lv28: R.Tensor((1, 320, 1, 1), dtype="float32"), lv35: R.Tensor((2, 320, 1, 1), dtype="float32")) -> R.Tensor((2, 320, 64, 64), dtype="float32"):
            R.func_attr({"Primitive": True})
            cls = Expected
            with R.dataflow():
                lv27 = R.call_tir(cls.conv2d, (inp_0, w1), out_ty=R.Tensor((2, 320, 64, 64), dtype="float32"))
                lv29 = R.call_tir(cls.add, (lv27, lv28), out_ty=R.Tensor((2, 320, 64, 64), dtype="float32"))
                gv = R.call_tir(cls.add2, (lv29, lv35), out_ty=R.Tensor((2, 320, 64, 64), dtype="float32"))
                R.output(gv)
            return gv

        @R.function(private=True)
        def fused_matmul_add1(inp_1: R.Tensor((2, 1280), dtype="float32"), lv31: R.Tensor((1280, 320), dtype="float32"), b2: R.Tensor((320,), dtype="float32")) -> R.Tensor((2, 320), dtype="float32"):
            cls = Expected
            R.func_attr({"Primitive": True})
            with R.dataflow():
                lv32 = R.call_tir(cls.matmul, (inp_1, lv31), out_ty=R.Tensor((2, 320), dtype="float32"))
                gv = R.call_tir(cls.add1, (lv32, b2), out_ty=R.Tensor((2, 320), dtype="float32"))
                R.output(gv)
            return gv

        @R.function
        def main(inp_0: R.Tensor((2, 320, 64, 64), dtype="float32"), inp_1: R.Tensor((2, 1280), dtype="float32"), w1: R.Tensor((320, 320, 3, 3), dtype="float32"), b1: R.Tensor((320,), dtype="float32"), w2: R.Tensor((320, 1280), dtype="float32"), b2: R.Tensor((320,), dtype="float32")) -> R.Tensor((2, 320, 64, 64), dtype="float32"):
            R.func_attr({"num_input": 2})
            cls = Expected
            with R.dataflow():
                lv28 = R.call_tir(cls.reshape, (b1,), out_ty=R.Tensor((1, 320, 1, 1), dtype="float32"))
                lv31 = R.call_tir(cls.transpose, (w2,), out_ty=R.Tensor((1280, 320), dtype="float32"))
                lv: R.Tensor((2, 320), dtype="float32") = cls.fused_matmul_add1(inp_1, lv31, b2)
                lv35 = R.call_tir(cls.reshape1, (lv,), out_ty=R.Tensor((2, 320, 1, 1), dtype="float32"))
                lv1: R.Tensor((2, 320, 64, 64), dtype="float32") = cls.fused_conv2d_add_add2(inp_0, w1, lv28, lv35)
                gv: R.Tensor((2, 320, 64, 64), dtype="float32") = lv1
                R.output(gv)
            return gv
    # fmt: on

    mod = relax.transform.LegalizeOps()(Module)
    mod = relax.transform.AnnotateTIROpPattern()(mod)
    mod = relax.transform.FuseOps()(mod)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_dead_group():
    # fmt: off

    @I.ir_module(s_tir=True)
    class Module:
        @R.function
        def main(inp_0: R.Tensor((1, 784), dtype="float32"), inp_1: R.Tensor((1, 128), dtype="float32"), linear1_bias: R.Tensor((128,), dtype="float32"), linear1_weight: R.Tensor((128, 784), dtype="float32"), linear2_bias: R.Tensor((10,), dtype="float32"), linear2_weight: R.Tensor((10, 128), dtype="float32")) -> R.Tensor((1, 10), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((784, 128), dtype="float32") = R.permute_dims(linear1_weight, axes=None)
                lv1: R.Tensor((1, 128), dtype="float32") = R.matmul(inp_0, lv, out_dtype="float32")
                lv2: R.Tensor((1, 128), dtype="float32") = R.add(lv1, linear1_bias)
                lv3: R.Tensor((1, 128), dtype="float32") = R.nn.relu(lv2)
                lv4: R.Tensor((128, 10), dtype="float32") = R.permute_dims(linear2_weight, axes=None)
                lv5: R.Tensor((1, 10), dtype="float32") = R.matmul(inp_1, lv4, out_dtype="float32")
                lv6: R.Tensor((1, 10), dtype="float32") = R.add(lv5, linear2_bias)
                gv: R.Tensor((1, 10), dtype="float32") = lv6
                R.output(gv)
            return gv

    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(private=True, s_tir=True)
        def add(rxplaceholder: T.Buffer((T.int64(1), T.int64(128)), "float32"), rxplaceholder_1: T.Buffer((T.int64(128),), "float32"), T_add: T.Buffer((T.int64(1), T.int64(128)), "float32")):
            T.func_attr({"op_pattern": 0, "tirx.noalias": True})
            # with T.sblock("root"):
            for ax0, ax1 in T.grid(T.int64(1), T.int64(128)):
                with T.sblock("T_add"):
                    v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
                    T.reads(rxplaceholder[v_ax0, v_ax1], rxplaceholder_1[v_ax1])
                    T.writes(T_add[v_ax0, v_ax1])
                    T_add[v_ax0, v_ax1] = rxplaceholder[v_ax0, v_ax1] + rxplaceholder_1[v_ax1]

        @T.prim_func(private=True, s_tir=True)
        def add1(rxplaceholder: T.Buffer((T.int64(1), T.int64(10)), "float32"), rxplaceholder_1: T.Buffer((T.int64(10),), "float32"), T_add: T.Buffer((T.int64(1), T.int64(10)), "float32")):
            T.func_attr({"op_pattern": 0, "tirx.noalias": True})
            # with T.sblock("root"):
            for ax0, ax1 in T.grid(T.int64(1), T.int64(10)):
                with T.sblock("T_add"):
                    v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
                    T.reads(rxplaceholder[v_ax0, v_ax1], rxplaceholder_1[v_ax1])
                    T.writes(T_add[v_ax0, v_ax1])
                    T_add[v_ax0, v_ax1] = rxplaceholder[v_ax0, v_ax1] + rxplaceholder_1[v_ax1]

        @T.prim_func(private=True, s_tir=True)
        def matmul(rxplaceholder: T.Buffer((T.int64(1), T.int64(784)), "float32"), rxplaceholder_1: T.Buffer((T.int64(784), T.int64(128)), "float32"), matmul_1: T.Buffer((T.int64(1), T.int64(128)), "float32")):
            T.func_attr({"op_pattern": 4, "tirx.noalias": True})
            # with T.sblock("root"):
            for i0, i1, k in T.grid(T.int64(1), T.int64(128), T.int64(784)):
                with T.sblock("matmul"):
                    v_i0, v_i1, v_k = T.axis.remap("SSR", [i0, i1, k])
                    T.reads(rxplaceholder[v_i0, v_k], rxplaceholder_1[v_k, v_i1])
                    T.writes(matmul_1[v_i0, v_i1])
                    with T.init():
                        matmul_1[v_i0, v_i1] = T.float32(0)
                    matmul_1[v_i0, v_i1] = matmul_1[v_i0, v_i1] + rxplaceholder[v_i0, v_k] * rxplaceholder_1[v_k, v_i1]

        @T.prim_func(private=True, s_tir=True)
        def matmul1(rxplaceholder: T.Buffer((T.int64(1), T.int64(128)), "float32"), rxplaceholder_1: T.Buffer((T.int64(128), T.int64(10)), "float32"), matmul: T.Buffer((T.int64(1), T.int64(10)), "float32")):
            T.func_attr({"op_pattern": 4, "tirx.noalias": True})
            # with T.sblock("root"):
            for i0, i1, k in T.grid(T.int64(1), T.int64(10), T.int64(128)):
                with T.sblock("matmul"):
                    v_i0, v_i1, v_k = T.axis.remap("SSR", [i0, i1, k])
                    T.reads(rxplaceholder[v_i0, v_k], rxplaceholder_1[v_k, v_i1])
                    T.writes(matmul[v_i0, v_i1])
                    with T.init():
                        matmul[v_i0, v_i1] = T.float32(0)
                    matmul[v_i0, v_i1] = matmul[v_i0, v_i1] + rxplaceholder[v_i0, v_k] * rxplaceholder_1[v_k, v_i1]

        @T.prim_func(private=True, s_tir=True)
        def relu(rxplaceholder: T.Buffer((T.int64(1), T.int64(128)), "float32"), compute: T.Buffer((T.int64(1), T.int64(128)), "float32")):
            T.func_attr({"op_pattern": 0, "tirx.noalias": True})
            # with T.sblock("root"):
            for i0, i1 in T.grid(T.int64(1), T.int64(128)):
                with T.sblock("compute"):
                    v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
                    T.reads(rxplaceholder[v_i0, v_i1])
                    T.writes(compute[v_i0, v_i1])
                    compute[v_i0, v_i1] = T.max(rxplaceholder[v_i0, v_i1], T.float32(0))

        @T.prim_func(private=True, s_tir=True)
        def transpose(rxplaceholder: T.Buffer((T.int64(128), T.int64(784)), "float32"), T_transpose: T.Buffer((T.int64(784), T.int64(128)), "float32")):
            T.func_attr({"op_pattern": 2, "tirx.noalias": True})
            # with T.sblock("root"):
            for ax0, ax1 in T.grid(T.int64(784), T.int64(128)):
                with T.sblock("T_transpose"):
                    v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
                    T.reads(rxplaceholder[v_ax1, v_ax0])
                    T.writes(T_transpose[v_ax0, v_ax1])
                    T_transpose[v_ax0, v_ax1] = rxplaceholder[v_ax1, v_ax0]

        @T.prim_func(private=True, s_tir=True)
        def transpose1(rxplaceholder: T.Buffer((T.int64(10), T.int64(128)), "float32"), T_transpose: T.Buffer((T.int64(128), T.int64(10)), "float32")):
            T.func_attr({"op_pattern": 2, "tirx.noalias": True})
            # with T.sblock("root"):
            for ax0, ax1 in T.grid(T.int64(128), T.int64(10)):
                with T.sblock("T_transpose"):
                    v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
                    T.reads(rxplaceholder[v_ax1, v_ax0])
                    T.writes(T_transpose[v_ax0, v_ax1])
                    T_transpose[v_ax0, v_ax1] = rxplaceholder[v_ax1, v_ax0]

        @R.function(private=True)
        def fused_matmul1_add1(inp_1: R.Tensor((1, 128), dtype="float32"), lv4: R.Tensor((128, 10), dtype="float32"), linear2_bias: R.Tensor((10,), dtype="float32")) -> R.Tensor((1, 10), dtype="float32"):
            R.func_attr({"Primitive": True})
            cls = Expected
            with R.dataflow():
                lv5 = R.call_tir(cls.matmul1, (inp_1, lv4), out_ty=R.Tensor((1, 10), dtype="float32"))
                gv = R.call_tir(cls.add1, (lv5, linear2_bias), out_ty=R.Tensor((1, 10), dtype="float32"))
                R.output(gv)
            return gv

        @R.function
        def main(inp_0: R.Tensor((1, 784), dtype="float32"), inp_1: R.Tensor((1, 128), dtype="float32"), linear1_bias: R.Tensor((128,), dtype="float32"), linear1_weight: R.Tensor((128, 784), dtype="float32"), linear2_bias: R.Tensor((10,), dtype="float32"), linear2_weight: R.Tensor((10, 128), dtype="float32")) -> R.Tensor((1, 10), dtype="float32"):
            R.func_attr({"num_input": 1})
            cls = Expected
            with R.dataflow():
                lv = R.call_tir(cls.transpose, (linear1_weight,), out_ty=R.Tensor((784, 128), dtype="float32"))
                lv4 = R.call_tir(cls.transpose1, (linear2_weight,), out_ty=R.Tensor((128, 10), dtype="float32"))
                lv_1: R.Tensor((1, 10), dtype="float32") = cls.fused_matmul1_add1(inp_1, lv4, linear2_bias)
                gv: R.Tensor((1, 10), dtype="float32") = lv_1
                R.output(gv)
            return gv

    # fmt: on

    mod = relax.transform.LegalizeOps()(Module)
    _check(mod, Expected)


def test_symbolic_shape_aware_fuse():
    @I.ir_module(s_tir=True)
    class Before:
        @R.function
        def main(x: R.Tensor(["n", "m"], "float32")):
            with R.dataflow():
                lv0 = R.emit_te(topi.add, x, R.const(1, "float32"))
                lv1 = R.emit_te(topi.exp, lv0)
                gv = R.emit_te(topi.squeeze, lv1)
                R.output(gv)
            return gv

    @I.ir_module(s_tir=True)
    class Expected:
        @R.function(private=True)
        def fused_add_exp_squeeze(
            x: R.Tensor(["n", "m"], "float32"), p0: R.Tensor([], "float32")
        ) -> R.Tensor(["n", "m"], dtype="float32"):
            R.func_attr({"Primitive": True})
            with R.dataflow():
                lv0 = R.emit_te(topi.add, x, p0)
                lv1 = R.emit_te(topi.exp, lv0)
                gv = R.emit_te(topi.squeeze, lv1)
                R.output(gv)
            return gv

        @R.function
        def main(x: R.Tensor(["n", "m"], "float32")) -> R.Tensor(["n", "m"], dtype="float32"):
            cls = Expected
            with R.dataflow():
                gv = cls.fused_add_exp_squeeze(x, R.const(1, "float32"))
                R.output(gv)
            return gv

    _check(Before, Expected)


def test_symbolic_shape_aware_fuse_2():
    @I.ir_module(s_tir=True)
    class Before:
        @R.function
        def main(s: R.Shape(["n"])):
            n = T.int64()
            with R.dataflow():
                lv0 = R.emit_te(topi.full, [n, n], "float32", 0)
                lv1 = R.emit_te(topi.trilu, lv0, tvm.tirx.const(1, "int32"), upper=True)
                gv = R.emit_te(topi.broadcast_to, lv1, [1, 1, n, n])
                R.output(gv)
            return gv

    @I.ir_module(s_tir=True)
    class Expected:
        @R.function(private=True)
        def fused_full_trilu_broadcast_to(
            s: R.Shape(["n"]),
        ) -> R.Tensor([1, 1, "n", "n"], "float32"):
            R.func_attr({"Primitive": True})
            n = T.int64()
            with R.dataflow():
                lv0 = R.emit_te(topi.full, [n, n], "float32", 0)
                lv1 = R.emit_te(topi.trilu, lv0, tvm.tirx.const(1, "int32"), upper=True)
                gv = R.emit_te(topi.broadcast_to, lv1, [1, 1, n, n])
                R.output(gv)
            return gv

        @R.function
        def main(s: R.Shape(["n"])) -> R.Tensor((1, 1, "n", "n"), dtype="float32"):
            cls = Expected
            n = T.int64()
            with R.dataflow():
                gv: R.Tensor([1, 1, n, n], "float32") = cls.fused_full_trilu_broadcast_to(
                    R.shape([n])
                )
                R.output(gv)
            return gv

    _check(Before, Expected)


def test_shape_expr_arg():
    @I.ir_module(s_tir=True)
    class Before:
        @R.function
        def main(s: R.Shape(["n"]), kv_cache: R.Any):
            n = T.int64()
            with R.dataflow():
                lv0 = R.emit_te(topi.full, [n, n], "float32", 0)
                lv1 = R.emit_te(topi.trilu, lv0, tvm.tirx.const(1, "int32"), upper=True)
                lv2 = R.emit_te(topi.broadcast_to, lv1, [1, 1, n, n])
                gv = R.call_pure_packed(
                    "vm.builtin.attention_kv_cache_view",
                    kv_cache,
                    R.shape([1 + n, 32, 128]),
                    ty_args=(R.Tensor((1 + n, 32, 128), dtype="float32"),),
                )
                R.output(gv, lv2)
            return gv, lv2

    @I.ir_module(s_tir=True)
    class Expected:
        @R.function(private=True)
        def fused_full_trilu_broadcast_to(
            s: R.Shape(["n"]),
        ) -> R.Tensor([1, 1, "n", "n"], "float32"):
            R.func_attr({"Primitive": True})
            n = T.int64()
            with R.dataflow():
                lv0 = R.emit_te(topi.full, [n, n], "float32", 0)
                lv1 = R.emit_te(topi.trilu, lv0, tvm.tirx.const(1, "int32"), upper=True)
                gv = R.emit_te(topi.broadcast_to, lv1, [1, 1, n, n])
                R.output(gv)
            return gv

        @R.function
        def main(s: R.Shape(["n"]), kv_cache: R.Any):
            cls = Expected
            n = T.int64()
            with R.dataflow():
                lv: R.Tensor([1, 1, n, n], "float32") = cls.fused_full_trilu_broadcast_to(
                    R.shape([n])
                )
                gv = R.call_pure_packed(
                    "vm.builtin.attention_kv_cache_view",
                    kv_cache,
                    R.shape([1 + n, 32, 128]),
                    ty_args=(R.Tensor((1 + n, 32, 128), dtype="float32"),),
                )
                R.output(gv, lv)
            return gv, lv

    _check(Before, Expected)


def test_skipping_match_cast():
    @I.ir_module(s_tir=True)
    class Module:
        @R.function
        def main(A: R.Tensor((10, 20), dtype="float32")) -> R.Tensor(dtype="float32", ndim=2):
            m = T.int64()
            n = T.int64()
            with R.dataflow():
                lv: R.Tensor((m, n), dtype="float32") = R.match_cast(
                    A, R.Tensor((m, n), dtype="float32")
                )
                gv: R.Tensor((m, n), dtype="float32") = lv
                R.output(gv)
            return gv

    _check(Module, Module)


def test_skipping_primvalue():
    @I.ir_module(s_tir=True)
    class Module:
        @R.function
        def main(inp: R.Tensor((2, 2), dtype="float32")) -> R.Tensor((2, 2), dtype="float32"):
            with R.dataflow():
                lv = R.call_pure_packed(
                    "my_func1", inp, R.prim_value(0), ty_args=[R.Tensor((2, 2), dtype="float32")]
                )
                lv1 = R.call_pure_packed(
                    "my_func2", lv, R.str("str"), ty_args=[R.Tensor((2, 2), dtype="float32")]
                )
                gv = R.call_pure_packed(
                    "my_func3",
                    lv1,
                    R.dtype("float32"),
                    ty_args=[R.Tensor((2, 2), dtype="float32")],
                )
                R.output(gv)
            return gv

    _check(Module, Module)


def test_partially_used_tuple_param():
    @I.ir_module(s_tir=True)
    class Module:
        @R.function
        def main(
            x: R.Tuple(
                R.Tensor((2,), "float32"),
                R.Tensor((2,), "float32"),
                R.Tensor((2,), "float32"),
                R.Tensor((2,), "float32"),
                R.Tensor((2,), "float32"),
                R.Tensor((2,), "float32"),
            ),
        ):
            with R.dataflow():
                x0 = x[0]
                y0 = R.emit_te(topi.add, x0, R.const(1, "float32"))
                y1 = R.emit_te(topi.divide, y0, R.const(1, "float32"))
                gv = y1
                R.output(gv)
            return gv

    @I.ir_module(s_tir=True)
    class Expected:
        @R.function(private=True)
        def fused_add_divide(
            x_0: R.Tensor((2,), dtype="float32"),
            param_0: R.Tensor((), dtype="float32"),
            param_1: R.Tensor((), dtype="float32"),
        ) -> R.Tensor((2,), dtype="float32"):
            R.func_attr({"Primitive": True})
            with R.dataflow():
                y0 = R.emit_te(topi.add, x_0, param_0)
                gv = R.emit_te(topi.divide, y0, param_1)
                R.output(gv)
            return gv

        @R.function
        def main(
            x: R.Tuple(
                R.Tensor((2,), dtype="float32"),
                R.Tensor((2,), dtype="float32"),
                R.Tensor((2,), dtype="float32"),
                R.Tensor((2,), dtype="float32"),
                R.Tensor((2,), dtype="float32"),
                R.Tensor((2,), dtype="float32"),
            ),
        ) -> R.Tensor((2,), dtype="float32"):
            cls = Expected
            with R.dataflow():
                lv: R.Tensor((2,), dtype="float32") = x[0]
                lv1: R.Tensor((2,), dtype="float32") = cls.fused_add_divide(
                    lv, R.const(1, "float32"), R.const(1, "float32")
                )
                gv: R.Tensor((2,), dtype="float32") = lv1
                R.output(gv)
            return gv

    _check(Module, Expected)


def test_call_tir_inplace():
    @I.ir_module(s_tir=True)
    class Module:
        @T.prim_func(private=True, s_tir=True)
        def add(
            A: T.Buffer((T.int64(10), T.int64(20)), "float32"),
            B: T.Buffer((), "float32"),
            Out: T.Buffer((T.int64(10), T.int64(20)), "float32"),
        ):
            T.func_attr({"tirx.noalias": True})
            for ax0, ax1 in T.grid(T.int64(10), T.int64(20)):
                with T.sblock("T_add"):
                    v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
                    T.reads(A[v_ax0, v_ax1], B[()])
                    T.writes(Out[v_ax0, v_ax1])
                    Out[v_ax0, v_ax1] = A[v_ax0, v_ax1] + B[()]

        @T.prim_func(private=True, s_tir=True)
        def exp_inplace(A: T.Buffer((T.int64(10), T.int64(20)), "float32")):
            T.func_attr({"tirx.noalias": True})
            for i0, i1 in T.grid(T.int64(10), T.int64(20)):
                with T.sblock("compute"):
                    v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
                    T.reads(A[v_i0, v_i1])
                    T.writes(A[v_i0, v_i1])
                    A[v_i0, v_i1] = T.exp(A[v_i0, v_i1])

        @T.prim_func(private=True, s_tir=True)
        def squeeze_inplace(A: T.Buffer((T.int64(10), T.int64(20)), "float32")):
            T.func_attr({"tirx.noalias": True})
            for ax0, ax1 in T.grid(T.int64(10), T.int64(20)):
                with T.sblock("T_squeeze"):
                    v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
                    T.reads(A[v_ax0, v_ax1])
                    T.writes(A[v_ax0, v_ax1])
                    A[v_ax0, v_ax1] = A[v_ax0, v_ax1]

        @R.function
        def main(
            x: R.Tensor((10, 20), dtype="float32"), p0: R.Tensor((), dtype="float32")
        ) -> R.Tensor((10, 20), dtype="float32"):
            cls = Module
            with R.dataflow():
                lv = R.call_tir(
                    cls.add,
                    (x, p0),
                    out_ty=R.Tensor((10, 20), dtype="float32"),
                )
                lv1 = R.call_tir_inplace(
                    cls.exp_inplace,
                    (lv,),
                    inplace_indices=[0],
                    out_ty=R.Tensor((10, 20), dtype="float32"),
                )
                gv = R.call_tir_inplace(
                    cls.squeeze_inplace,
                    (lv1,),
                    inplace_indices=[0],
                    out_ty=R.Tensor((10, 20), dtype="float32"),
                )
                R.output(gv)
            return gv

    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(private=True, s_tir=True)
        def add(
            A: T.Buffer((T.int64(10), T.int64(20)), "float32"),
            B: T.Buffer((), "float32"),
            Out: T.Buffer((T.int64(10), T.int64(20)), "float32"),
        ):
            T.func_attr({"tirx.noalias": True, "op_pattern": 0})
            for ax0, ax1 in T.grid(T.int64(10), T.int64(20)):
                with T.sblock("T_add"):
                    v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
                    T.reads(A[v_ax0, v_ax1], B[()])
                    T.writes(Out[v_ax0, v_ax1])
                    Out[v_ax0, v_ax1] = A[v_ax0, v_ax1] + B[()]

        @T.prim_func(private=True, s_tir=True)
        def exp_inplace(A: T.Buffer((T.int64(10), T.int64(20)), "float32")):
            T.func_attr({"tirx.noalias": True, "op_pattern": 0})
            for i0, i1 in T.grid(T.int64(10), T.int64(20)):
                with T.sblock("compute"):
                    v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
                    T.reads(A[v_i0, v_i1])
                    T.writes(A[v_i0, v_i1])
                    A[v_i0, v_i1] = T.exp(A[v_i0, v_i1])

        @T.prim_func(private=True, s_tir=True)
        def squeeze_inplace(A: T.Buffer((T.int64(10), T.int64(20)), "float32")):
            T.func_attr({"tirx.noalias": True, "op_pattern": 0})
            for ax0, ax1 in T.grid(T.int64(10), T.int64(20)):
                with T.sblock("T_squeeze"):
                    v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
                    T.reads(A[v_ax0, v_ax1])
                    T.writes(A[v_ax0, v_ax1])
                    A[v_ax0, v_ax1] = A[v_ax0, v_ax1]

        @R.function(private=True)
        def fused_add_exp_inplace_squeeze_inplace(
            x: R.Tensor((10, 20), dtype="float32"), p0: R.Tensor((), dtype="float32")
        ) -> R.Tensor((10, 20), dtype="float32"):
            R.func_attr({"Primitive": True})
            cls = Expected
            with R.dataflow():
                lv = R.call_tir(
                    cls.add,
                    (x, p0),
                    out_ty=R.Tensor((10, 20), dtype="float32"),
                )
                lv1 = R.call_tir_inplace(
                    cls.exp_inplace,
                    (lv,),
                    inplace_indices=[0],
                    out_ty=R.Tensor((10, 20), dtype="float32"),
                )
                gv = R.call_tir_inplace(
                    cls.squeeze_inplace,
                    (lv1,),
                    inplace_indices=[0],
                    out_ty=R.Tensor((10, 20), dtype="float32"),
                )
                R.output(gv)
            return gv

        @R.function
        def main(
            x: R.Tensor((10, 20), dtype="float32"), p0: R.Tensor((), dtype="float32")
        ) -> R.Tensor((10, 20), dtype="float32"):
            cls = Expected
            with R.dataflow():
                gv1: R.Tensor((10, 20), dtype="float32") = (
                    cls.fused_add_exp_inplace_squeeze_inplace(x, p0)
                )
                R.output(gv1)
            return gv1

    _check(Module, Expected)


def test_packed_params():
    # fmt: off
    @I.ir_module(s_tir=True)
    class Before:
        @T.prim_func(private=True, s_tir=True)
        def cast(lv: T.Buffer((T.int64(16), T.int64(16)), "float16"), compute: T.Buffer((T.int64(16), T.int64(16)), "float32")):
            T.func_attr({"tirx.noalias": True})
            # with T.sblock("root"):
            for i0, i1 in T.grid(T.int64(16), T.int64(16)):
                with T.sblock("compute"):
                    v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
                    T.reads(lv[v_i0, v_i1])
                    T.writes(compute[v_i0, v_i1])
                    compute[v_i0, v_i1] = T.Cast("float32", lv[v_i0, v_i1])

        @T.prim_func(private=True, s_tir=True)
        def matmul(x: T.Buffer((T.int64(16), T.int64(16)), "float32"), lv2: T.Buffer((T.int64(16), T.int64(16)), "float32"), T_matmul: T.Buffer((T.int64(16), T.int64(16)), "float32")):
            T.func_attr({"tirx.noalias": True})
            # with T.sblock("root"):
            for ax0, ax1, k in T.grid(T.int64(16), T.int64(16), T.int64(16)):
                with T.sblock("T_matmul"):
                    v_ax0, v_ax1, v_k = T.axis.remap("SSR", [ax0, ax1, k])
                    T.reads(x[v_ax0, v_k], lv2[v_k, v_ax1])
                    T.writes(T_matmul[v_ax0, v_ax1])
                    with T.init():
                        T_matmul[v_ax0, v_ax1] = T.float32(0)
                    T_matmul[v_ax0, v_ax1] = T_matmul[v_ax0, v_ax1] + x[v_ax0, v_k] * lv2[v_k, v_ax1]

        @R.function
        def main(x: R.Tensor((16, 16), dtype="float32"), packed_params: R.Tuple(R.Tensor((16, 16), dtype="float16"), R.Tensor((16, 16), dtype="float16"))) -> R.Tensor((16, 16), dtype="float32"):
            R.func_attr({"num_input": 1})
            cls = Before
            with R.dataflow():
                lv: R.Tensor((16, 16), dtype="float16") = packed_params[0]
                lv1: R.Tensor((16, 16), dtype="float16") = packed_params[1]
                lv2 = R.call_tir(cls.cast, (lv,), out_ty=R.Tensor((16, 16), dtype="float32"))
                lv3 = R.call_tir(cls.matmul, (x, lv2), out_ty=R.Tensor((16, 16), dtype="float32"))
                lv4 = R.call_tir(cls.cast, (lv1,), out_ty=R.Tensor((16, 16), dtype="float32"))
                lv5 = R.call_tir(cls.matmul, (lv3, lv4), out_ty=R.Tensor((16, 16), dtype="float32"))
                gv: R.Tensor((16, 16), dtype="float32") = lv5
                R.output(gv)
            return gv
    # fmt: on

    Expected = Before
    _check(Before, Expected)


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