# 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.


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_before, mod_expected):
    mod_after = relax.transform.FuseTIR()(mod_before)
    tvm.ir.assert_structural_equal(mod_expected, mod_after)


def test_simple():
    def before():
        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"))
        p0 = relax.Var("p0", R.Tensor([], "float32"))
        with bb.function("main", [x, p0]):
            with bb.dataflow():
                gv = bb.emit_output(relax.Call(fused_add_exp_squeeze, [x, p0]))
            bb.emit_func_output(gv)

        return bb.get().with_attrs({"foo": "bar"})

    def expected():
        def fused_add_exp_squeeze(x, p0):
            add = topi.add(x, p0)
            exp = topi.exp(add)
            squeeze = topi.squeeze(exp)
            return squeeze

        bb = relax.BlockBuilder()
        x = relax.Var("x", R.Tensor([10, 20], "float32"))
        p0 = relax.Var("p0", R.Tensor([], "float32"))
        with bb.function("main", [x, p0]):
            with bb.dataflow():
                gv = bb.emit_output(bb.call_te(fused_add_exp_squeeze, x, p0))
            bb.emit_func_output(gv)
        return bb.get().with_attrs({"foo": "bar"})

    _check(before(), expected())


def test_conv2d_fuse():
    def before(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}):
            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}):
            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()

    def expected(dtype):
        def fused_conv2d_add1_add2(x, w, p):
            conv = topi.nn.conv2d(x, w, strides=1, padding=1, dilation=1)
            add = topi.add(p, conv)
            return topi.add(conv, add)

        def fused_conv2d1_add2(x, w, p):
            conv = topi.nn.conv2d(x, w, strides=1, padding=0, dilation=1)
            return topi.add(conv, p)

        bb = relax.BlockBuilder()

        # 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_te(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(bb.call_te(fused_conv2d1_add2, lv1, w2, lv2))
            bb.emit_func_output(gv)

        return bb.get()

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


def test_two_subfunction():
    def before():
        bb = relax.BlockBuilder()
        x1 = relax.Var("x1", R.Tensor([10, 20], "float32"))
        with bb.function("fused_exp_squeeze", [x1], attrs={"Primitive": True}):
            with bb.dataflow():
                lv1 = bb.emit_te(topi.exp, x1)
                gv = bb.emit_output(bb.call_te(topi.squeeze, lv1))
            bb.emit_func_output(gv)
        mod = bb.get()

        func_gv = mod.get_global_var("fused_exp_squeeze")
        x = relax.Var("x", R.Tensor([10, 20], "float32"))
        with bb.function("main", [x]):
            with bb.dataflow():
                lv = bb.emit(relax.Call(func_gv, [x]))
                lv2 = bb.emit(relax.Call(func_gv, [lv]))
                gv = bb.emit_output(lv2)
            bb.emit_func_output(gv)
        return bb.get()

    def expected():
        def fused_exp_squeeze(x):
            exp = topi.exp(x)
            squeeze = topi.squeeze(exp)
            return squeeze

        bb = relax.BlockBuilder()
        x = relax.Var("x", R.Tensor([10, 20], "float32"))
        with bb.function("main", [x]):
            with bb.dataflow():
                lv = bb.emit_te(fused_exp_squeeze, x)
                lv2 = bb.call_te(fused_exp_squeeze, lv)
                gv = bb.emit_output(lv2)
            bb.emit_func_output(gv)
        return bb.get()

    _check(before(), expected())


def test_fuse_same_primfunc():
    def before():
        bb = relax.BlockBuilder()
        x1 = relax.Var("x1", R.Tensor([10, 20], "float32"))
        with bb.function("fused_exp_exp_squeeze", [x1], attrs={"Primitive": True}):
            with bb.dataflow():
                lv1 = bb.emit_te(topi.exp, x1)
                lv2 = bb.emit_te(topi.exp, lv1)
                gv = bb.emit_output(bb.call_te(topi.squeeze, lv2))
            bb.emit_func_output(gv)
        mod = bb.get()

        func_gv = mod.get_global_var("fused_exp_exp_squeeze")
        x = relax.Var("x", R.Tensor([10, 20], "float32"))
        with bb.function("main", [x]):
            with bb.dataflow():
                lv = bb.emit(relax.Call(func_gv, [x]))
                gv = bb.emit_output(lv)
            bb.emit_func_output(gv)
        return bb.get()

    def expected():
        def fused_exp_exp_squeeze(x):
            exp = topi.exp(x)
            exp = topi.exp(exp)
            squeeze = topi.squeeze(exp)
            return squeeze

        bb = relax.BlockBuilder()
        x = relax.Var("x", R.Tensor([10, 20], "float32"))
        with bb.function("main", [x]):
            with bb.dataflow():
                lv = bb.call_te(fused_exp_exp_squeeze, x)
                gv = bb.emit_output(lv)
            bb.emit_func_output(gv)
        return bb.get()

    _check(before(), expected())


def test_fuse_with_tuple_as_param():
    def before():
        bb = relax.BlockBuilder()
        x = relax.Var("x", R.Tuple([R.Tensor([10], "float32"), R.Tensor([10], "float32")]))
        with bb.function("fused_exp_add", [x], attrs={"Primitive": True}, private=True):
            with bb.dataflow():
                lv0 = bb.emit(relax.TupleGetItem(x, 0))
                lv1 = bb.emit(relax.TupleGetItem(x, 1))
                lv2 = bb.emit_te(topi.exp, lv0)
                gv = bb.emit_output(bb.call_te(topi.add, lv2, lv1))
            bb.emit_func_output(gv)
        mod = bb.get()

        func_gv = mod.get_global_var("fused_exp_add")
        x = relax.Var("x", R.Tuple([R.Tensor([10], "float32"), R.Tensor([10], "float32")]))
        with bb.function("main", [x]):
            with bb.dataflow():
                gv = bb.emit_output(relax.Call(func_gv, [x]))
            bb.emit_func_output(gv)
        return bb.get()

    def expected():
        def fused_exp_add(x1, x2):
            exp = topi.exp(x1)
            return topi.add(exp, x2)

        bb = relax.BlockBuilder()
        x = relax.Var("x", R.Tuple([R.Tensor([10], "float32"), R.Tensor([10], "float32")]))
        with bb.function("main", [x]):
            with bb.dataflow():
                lv0 = bb.emit(relax.TupleGetItem(x, 0))
                lv1 = bb.emit(relax.TupleGetItem(x, 1))
                gv = bb.emit_output(bb.call_te(fused_exp_add, lv0, lv1))
            bb.emit_func_output(gv)
        return bb.get()

    _check(before(), expected())


def test_fuse_with_nested_tuple_as_param():
    tuple_ty = R.Tuple(
        [R.Tensor([10], "float32"), R.Tuple([R.Tensor([10], "float32"), R.Tensor([10], "float32")])]
    )

    def before():
        bb = relax.BlockBuilder()
        x = relax.Var("x", tuple_ty)
        with bb.function("fused_exp_add_add", [x], attrs={"Primitive": True}, private=True):
            with bb.dataflow():
                lv0 = bb.emit(relax.TupleGetItem(x, 0))
                lv0_exp = bb.emit_te(topi.exp, lv0)
                lv1 = bb.emit(relax.TupleGetItem(x, 1))
                lv1_0 = bb.emit(relax.TupleGetItem(lv1, 0))
                lv1_1 = bb.emit(relax.TupleGetItem(lv1, 1))
                lv2 = bb.emit_te(topi.add, lv1_0, lv1_1)
                gv = bb.emit_output(bb.call_te(topi.add, lv0_exp, lv2))
            bb.emit_func_output(gv)
        mod = bb.get()

        func_gv = mod.get_global_var("fused_exp_add_add")
        x = relax.Var("x", tuple_ty)
        with bb.function("main", [x]):
            with bb.dataflow():
                gv = bb.emit_output(relax.Call(func_gv, [x]))
            bb.emit_func_output(gv)
        return bb.get()

    def expected():
        def fused_exp_add_add(x1, x2, x3):
            exp = topi.exp(x1)
            add = topi.add(x2, x3)
            return topi.add(exp, add)

        bb = relax.BlockBuilder()
        x = relax.Var("x", tuple_ty)
        with bb.function("main", [x]):
            with bb.dataflow():
                lv0 = bb.emit(relax.TupleGetItem(x, 0))
                lv1 = bb.emit(relax.TupleGetItem(x, 1))
                lv2 = bb.emit(relax.TupleGetItem(lv1, 0))
                lv3 = bb.emit(relax.TupleGetItem(lv1, 1))
                gv = bb.emit_output(bb.call_te(fused_exp_add_add, lv0, lv2, lv3))
            bb.emit_func_output(gv)
        return bb.get()

    _check(before(), expected())


def test_fuse_with_call_tir_in_main():
    def before():
        bb = relax.BlockBuilder()
        x1 = relax.Var("x1", R.Tensor([10, 20], "float32"))
        with bb.function("fused_exp_squeeze", [x1], attrs={"Primitive": True}):
            with bb.dataflow():
                lv = bb.emit_te(topi.exp, x1)
                gv = bb.emit_output(bb.call_te(topi.squeeze, lv))
            bb.emit_func_output(gv)
        mod = bb.get()

        func_gv = mod.get_global_var("fused_exp_squeeze")
        x = relax.Var("x", R.Tensor([10, 20], "float32"))
        with bb.function("main", [x]):
            with bb.dataflow():
                lv0 = bb.emit(relax.Call(func_gv, [x]))
                lv1 = bb.emit_te(topi.add, lv0, relax.const(1, "float32"))
                gv = bb.emit_output(lv1)
            bb.emit_func_output(gv)
        return bb.get()

    def expected():
        def fused_exp_squeeze(x):
            exp = topi.exp(x)
            squeeze = topi.squeeze(exp)
            return squeeze

        bb = relax.BlockBuilder()
        x = relax.Var("x", R.Tensor([10, 20], "float32"))
        with bb.function("main", [x]):
            with bb.dataflow():
                lv = bb.emit_te(fused_exp_squeeze, x)
                lv2 = bb.call_te(topi.add, lv, relax.const(1, "float32"))
                gv = bb.emit_output(lv2)
            bb.emit_func_output(gv)
        return bb.get()

    _check(before(), expected())


def test_fuse_with_const_in_argument():
    def before():
        bb = relax.BlockBuilder()
        x1 = relax.Var("x1", R.Tensor([10, 20], "float32"))
        x2 = relax.Var("x2", R.Tensor([], "float32"))
        with bb.function("fused_add_exp_squeeze", [x1, x2], attrs={"Primitive": True}):
            with bb.dataflow():
                lv0 = bb.emit_te(topi.add, x1, x2)
                lv1 = bb.emit_te(topi.exp, lv0)
                gv = bb.emit_output(bb.call_te(topi.squeeze, lv1))
            bb.emit_func_output(gv)
        mod = bb.get()

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

    def expected():
        def fused_add_exp_squeeze(x, y):
            add = topi.add(x, y)
            exp = topi.exp(add)
            squeeze = topi.squeeze(exp)
            return squeeze

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

    _check(before(), expected())


def test_fuse_tuple_output():
    def before():
        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", [x, p0], attrs={"Primitive": True}):
            with bb.dataflow():
                gv0 = bb.emit_output(bb.call_te(topi.add, x, p0))
                gv1 = bb.emit_output(bb.call_te(topi.exp, gv0))
            bb.emit_func_output(relax.Tuple([gv0, gv1]))
        fused_add_exp = bb.get().get_global_var("fused_add_exp")

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

        return bb.get()

    def expected():
        def fused_add_exp(x, p0):
            add = topi.add(x, p0)
            exp = topi.exp(add)
            return add, exp

        bb = relax.BlockBuilder()
        x = relax.Var("x", R.Tensor([10, 20], "float32"))
        p0 = relax.Var("p0", R.Tensor([], "float32"))
        with bb.function("main", [x, p0]):
            with bb.dataflow():
                gv = bb.emit_output(bb.call_te(fused_add_exp, x, p0))
            bb.emit_func_output(gv)
        return bb.get()

    _check(before(), expected())


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

        with bb.function("fused_add", [x, y], attrs={"Primitive": True}):
            with bb.dataflow():
                lv_tuple = bb.emit(relax.Tuple([x, relax.Tuple([x, y])]))
                lv_x = bb.emit(relax.TupleGetItem(lv_tuple, 0))
                lv0 = bb.emit(relax.TupleGetItem(lv_tuple, 1))
                lv_y = bb.emit(relax.TupleGetItem(lv0, 1))
                gv = bb.emit_output(bb.call_te(topi.add, lv_x, lv_y))
            bb.emit_func_output(gv)
        fused_add = bb.get().get_global_var("fused_add")

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

        return bb.get()

    def expected():
        bb = relax.BlockBuilder()
        x = relax.Var("x", R.Tensor([10, 20], "float32"))
        y = relax.Var("y", R.Tensor([10, 20], "float32"))
        with bb.function("main", [x, y]):
            with bb.dataflow():
                gv = bb.emit_output(bb.call_te(topi.add, x, y, primfunc_name_hint="fused_add"))
            bb.emit_func_output(gv)
        return bb.get()

    _check(before(), expected())


def test_fuse_return_partial_result():
    def te_argmax_idx_val(val):
        from tvm import te

        def f_combine(x, y):
            lhs = tvm.tirx.Select((x[1] >= y[1]), x[0], y[0])
            rhs = tvm.tirx.Select((x[1] >= y[1]), x[1], y[1])
            return lhs, rhs

        def f_identity(dtype0: tvm.DataType, dtype1: tvm.DataType):
            return tvm.tirx.const(-1, dtype0), tvm.te.min_value(dtype1)

        argmax = te.comm_reducer(f_combine, f_identity, name="argmax")
        m, n = val.shape
        k = te.reduce_axis((0, n), "k")
        max_idx, max_val = te.compute(
            (m,), lambda i: argmax((k.var, val[i, k]), axis=k), name="argmax"
        )
        return max_idx, max_val

    def before():
        bb = relax.BlockBuilder()
        x = relax.Var("x", R.Tensor([10, 20], "float32"))
        offset = relax.Var("offset", R.Tensor([10], "int32"))
        with bb.function("fused_argmax_add", [x, offset], attrs={"Primitive": True}):
            with bb.dataflow():
                lv = bb.emit_te(te_argmax_idx_val, x)
                idx = bb.emit(relax.TupleGetItem(lv, 0))
                gv = bb.emit_output(bb.call_te(topi.add, idx, offset))
            bb.emit_func_output(gv)
        mod = bb.get()

        func_gv = mod.get_global_var("fused_argmax_add")
        x = relax.Var("x", R.Tensor([10, 20], "float32"))
        offset = relax.Var("x", R.Tensor([10], "int32"))
        with bb.function("main", [x, offset]):
            with bb.dataflow():
                gv = bb.emit_output(relax.Call(func_gv, [x, offset]))
            bb.emit_func_output(gv)
        return bb.get()

    def expected():
        def fused_argmax_add(x, offset):
            idx, value = te_argmax_idx_val(x)
            idx = topi.add(idx, offset)
            return idx

        bb = relax.BlockBuilder()
        x = relax.Var("x", R.Tensor([10, 20], "float32"))
        offset = relax.Var("offset", R.Tensor([10], "int32"))
        with bb.function("main", [x, offset]):
            with bb.dataflow():
                gv = bb.emit_output(bb.call_te(fused_argmax_add, x, offset))
            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"))
        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()

    @I.ir_module(s_tir=True)
    class Expected:
        @R.function
        def func1(x: R.Tensor((10, 20), dtype="float32")) -> R.Tensor((10, 20), dtype="float32"):
            with R.dataflow():
                gv2 = R.call_tir(
                    Expected.fused_add_exp_squeeze,
                    (x, R.const(1, "float32")),
                    out_ty=R.Tensor((10, 20), dtype="float32"),
                )
                R.output(gv2)
            return gv2

        @R.function
        def func2(x: R.Tensor((20, 10), dtype="float32")) -> R.Tensor((20, 10), dtype="float32"):
            with R.dataflow():
                gv3 = R.call_tir(
                    Expected.fused_add1_exp1_squeeze1,
                    (x, R.const(1, "float32")),
                    out_ty=R.Tensor((20, 10), dtype="float32"),
                )
                R.output(gv3)
            return gv3

        @T.prim_func(private=True, s_tir=True)
        def fused_add1_exp1_squeeze1(
            x: T.Buffer((T.int64(20), T.int64(10)), "float32"),
            p0: T.Buffer((), "float32"),
            T_squeeze: T.Buffer((T.int64(20), T.int64(10)), "float32"),
        ):
            T.func_attr({"tirx.noalias": True})
            T_add = T.sblock_alloc_buffer((T.int64(20), T.int64(10)))
            compute = T.sblock_alloc_buffer((T.int64(20), T.int64(10)))
            for ax0, ax1 in T.grid(T.int64(20), T.int64(10)):
                with T.sblock("T_add"):
                    v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
                    T.reads(x[v_ax0, v_ax1], p0[()])
                    T.writes(T_add[v_ax0, v_ax1])
                    T_add[v_ax0, v_ax1] = x[v_ax0, v_ax1] + p0[()]
            for i0, i1 in T.grid(T.int64(20), T.int64(10)):
                with T.sblock("compute"):
                    v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
                    T.reads(T_add[v_i0, v_i1])
                    T.writes(compute[v_i0, v_i1])
                    compute[v_i0, v_i1] = T.exp(T_add[v_i0, v_i1])
            for ax0, ax1 in T.grid(T.int64(20), T.int64(10)):
                with T.sblock("T_squeeze"):
                    v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
                    T.reads(compute[v_ax0, v_ax1])
                    T.writes(T_squeeze[v_ax0, v_ax1])
                    T_squeeze[v_ax0, v_ax1] = compute[v_ax0, v_ax1]

        @T.prim_func(private=True, s_tir=True)
        def fused_add_exp_squeeze(
            x: T.Buffer((T.int64(10), T.int64(20)), "float32"),
            p0: T.Buffer((), "float32"),
            T_squeeze: T.Buffer((T.int64(10), T.int64(20)), "float32"),
        ):
            T.func_attr({"tirx.noalias": True})
            T_add = T.sblock_alloc_buffer((T.int64(10), T.int64(20)))
            compute = T.sblock_alloc_buffer((T.int64(10), T.int64(20)))
            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(x[v_ax0, v_ax1], p0[()])
                    T.writes(T_add[v_ax0, v_ax1])
                    T_add[v_ax0, v_ax1] = x[v_ax0, v_ax1] + p0[()]
            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(T_add[v_i0, v_i1])
                    T.writes(compute[v_i0, v_i1])
                    compute[v_i0, v_i1] = T.exp(T_add[v_i0, v_i1])
            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(compute[v_ax0, v_ax1])
                    T.writes(T_squeeze[v_ax0, v_ax1])
                    T_squeeze[v_ax0, v_ax1] = compute[v_ax0, v_ax1]

    _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

    # FuseTIR should do no change to it.
    _check(Module, Module)


def test_symbolic_shape_aware_fuse():
    @I.ir_module(s_tir=True)
    class Before:
        @R.function
        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 = Before
            with R.dataflow():
                gv = cls.fused_add_exp_squeeze(x, R.const(1, "float32"))
                R.output(gv)
            return gv

    def fused_add_exp_squeeze(x, p0):
        return topi.squeeze(topi.exp(topi.add(x, p0)))

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

    _check(Before, Expected)


def test_fuse_of_dynamic_kernel_with_var_params_and_static_args():
    @I.ir_module(s_tir=True)
    class Before:
        @T.prim_func(private=True, s_tir=True)
        def dynamic_tir_kernel(a: T.handle, b: T.handle):
            m = T.int64()
            n = T.int64()
            A = T.match_buffer(a, [m, n], "float32")
            B = T.match_buffer(b, [m, n], "float32")

            for iters in T.grid(m, n):
                with T.sblock("compute"):
                    i, j = T.axis.remap("SS", iters)
                    B[i, j] = A[i, j] * i + j

        @R.function(private=True)
        def fused_function(x: R.Tensor([16, 32], "float32")) -> R.Tensor([16, 32], dtype="float32"):
            R.func_attr({"Primitive": True})
            cls = Before
            with R.dataflow():
                y = R.call_tir(cls.dynamic_tir_kernel, [x], out_ty=R.Tensor([16, 32], "float32"))
                z = R.call_tir(cls.dynamic_tir_kernel, [y], out_ty=R.Tensor([16, 32], "float32"))
                R.output(z)
            return z

        @R.function
        def main(x: R.Tensor([16, 32], "float32")) -> R.Tensor([16, 32], dtype="float32"):
            cls = Before
            with R.dataflow():
                gv = cls.fused_function(x)
                R.output(gv)
            return gv

    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(private=True, s_tir=True)
        def fused_function(
            X: T.Buffer([T.int64(16), T.int64(32)], "float32"),
            Z: T.Buffer([T.int64(16), T.int64(32)], "float32"),
        ):
            T.func_attr({"tirx.noalias": True})
            Y = T.sblock_alloc_buffer(X.shape, "float32")
            for iters in T.grid(*X.shape):
                with T.sblock("compute_Y"):
                    i, j = T.axis.remap("SS", iters)
                    Y[i, j] = X[i, j] * i + j

            for iters in T.grid(*X.shape):
                with T.sblock("compute_Z"):
                    i, j = T.axis.remap("SS", iters)
                    Z[i, j] = Y[i, j] * i + j

        @R.function
        def main(x: R.Tensor([16, 32], "float32")) -> R.Tensor([16, 32], dtype="float32"):
            cls = Expected
            with R.dataflow():
                gv = R.call_tir(cls.fused_function, [x], out_ty=R.Tensor([16, 32], "float32"))
                R.output(gv)
            return gv

    _check(Before, Expected)


def test_fuse_of_dynamic_kernel_with_expression_params_and_static_args():
    """Parameters and arguments do not need to match structurally

    Here, the kernel requires arguments (m*n), and is provided
    """

    @I.ir_module(s_tir=True)
    class Before:
        @T.prim_func(private=True, s_tir=True)
        def dynamic_tir_kernel(a: T.handle, b: T.handle, c: T.handle, d: T.handle):
            m = T.int64()
            n = T.int64()
            A = T.match_buffer(a, [m * n], "float32")
            B = T.match_buffer(b, [m], "float32")
            C = T.match_buffer(c, [n], "float32")
            D = T.match_buffer(d, [m * n], "float32")

            for i, j in T.grid(m, n):
                with T.sblock("compute"):
                    vi, vj = T.axis.remap("SS", [i, j])
                    D[vi * 32 + vj] = A[vi * 32 + vj] * B[vi] + C[vj]

        @R.function(private=True)
        def fused_function(
            x: R.Tensor([16 * 32], "float32"),
            B: R.Tensor([16], "float32"),
            C: R.Tensor([32], "float32"),
        ):
            R.func_attr({"Primitive": True})
            cls = Before
            with R.dataflow():
                y = R.call_tir(
                    cls.dynamic_tir_kernel, [x, B, C], out_ty=R.Tensor([16 * 32], "float32")
                )
                z = R.call_tir(
                    cls.dynamic_tir_kernel, [y, B, C], out_ty=R.Tensor([16 * 32], "float32")
                )
                R.output(z)
            return z

        @R.function
        def main(
            x: R.Tensor([16 * 32], "float32"),
            B: R.Tensor([16], "float32"),
            C: R.Tensor([32], "float32"),
        ):
            cls = Before
            with R.dataflow():
                gv = cls.fused_function(x, B, C)
                R.output(gv)
            return gv

    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(private=True, s_tir=True)
        def fused_function(
            X: T.Buffer(T.int64(512), "float32"),
            B: T.Buffer(T.int64(16), "float32"),
            C: T.Buffer(T.int64(32), "float32"),
            Z: T.Buffer(T.int64(512), "float32"),
        ):
            T.func_attr({"tirx.noalias": True})
            Y = T.sblock_alloc_buffer((T.int64(512),))
            for i, j in T.grid(T.int64(16), T.int64(32)):
                with T.sblock("compute"):
                    vi, vj = T.axis.remap("SS", [i, j])
                    Y[vi * 32 + vj] = X[vi * 32 + vj] * B[vi] + C[vj]

            for i, j in T.grid(T.int64(16), T.int64(32)):
                with T.sblock("compute_1"):
                    vi, vj = T.axis.remap("SS", [i, j])
                    Z[vi * 32 + vj] = Y[vi * 32 + vj] * B[vi] + C[vj]

        @R.function
        def main(
            x: R.Tensor((512,), dtype="float32"),
            B: R.Tensor((16,), dtype="float32"),
            C: R.Tensor((32,), dtype="float32"),
        ) -> R.Tensor((512,), dtype="float32"):
            cls = Expected
            with R.dataflow():
                gv = R.call_tir(
                    cls.fused_function, (x, B, C), out_ty=R.Tensor((512,), dtype="float32")
                )
                R.output(gv)
            return gv

    _check(Before, Expected)


def test_symbolic_shape_aware_fuse_with_allocation():
    def te_mean(x, axis):
        return topi.divide(topi.sum(x, axis, keepdims=True), 4096)

    @I.ir_module(s_tir=True)
    class Before:
        @R.function
        def fused_mean_add_tir_sqrt_divide_multiply(
            x: R.Tensor((1, "n", 4096), dtype="float32"),
            y: R.Tensor((1, "n", 4096), dtype="float32"),
            rms_norm_weight: R.Tensor((4096,), dtype="float32"),
        ) -> R.Tensor((1, "n", 4096), dtype="float32"):
            R.func_attr({"Primitive": True})
            with R.dataflow():
                lv0 = R.emit_te(te_mean, x, axis=2)
                lv1 = R.emit_te(topi.add, lv0, lv0)
                lv2 = R.emit_te(topi.sqrt, lv1)
                lv3 = R.emit_te(topi.divide, y, lv2)
                gv = R.emit_te(topi.multiply, rms_norm_weight, lv3)
                R.output(gv)
            return gv

        @R.function
        def main(
            x: R.Tensor((1, "n", 4096), dtype="float32"),
            y: R.Tensor((1, "n", 4096), dtype="float32"),
            rms_norm_weight: R.Tensor((4096,), dtype="float32"),
        ) -> R.Tensor((1, "n", 4096), dtype="float32"):
            cls = Before
            with R.dataflow():
                gv = cls.fused_mean_add_tir_sqrt_divide_multiply(x, y, rms_norm_weight)
                R.output(gv)
            return gv

    def fused_mean_add_tir_sqrt_divide_multiply(x, y, rms_norm_weight):
        lv0 = te_mean(x, axis=2)
        lv1 = topi.add(lv0, lv0)
        lv2 = topi.sqrt(lv1)
        lv3 = topi.divide(y, lv2)
        return topi.multiply(rms_norm_weight, lv3)

    @I.ir_module(s_tir=True)
    class Expected:
        @R.function
        def main(
            x: R.Tensor((1, "n", 4096), dtype="float32"),
            y: R.Tensor((1, "n", 4096), dtype="float32"),
            rms_norm_weight: R.Tensor((4096,), dtype="float32"),
        ) -> R.Tensor((1, "n", 4096), dtype="float32"):
            with R.dataflow():
                gv = R.emit_te(fused_mean_add_tir_sqrt_divide_multiply, x, y, rms_norm_weight)
                R.output(gv)
            return gv

    _check(Before, Expected)


def test_symbolic_var_in_call_tir_args():
    @I.ir_module(s_tir=True)
    class Before:
        @T.prim_func(private=True, s_tir=True)
        def foo(
            X: T.Buffer((T.int64(1), T.int64(1), T.int64(32), T.int64(128)), "float32"),
            Y: T.Buffer((T.int64(2048), T.int64(128)), "float32"),
            m: T.int64,
            rotary: T.Buffer((T.int64(1), T.int64(1), T.int64(32), T.int64(128)), "float32"),
        ):
            for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(1), T.int64(32), T.int64(128)):
                with T.sblock("rotary"):
                    v0, v1, v2, v3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
                    rotary[v0, v1, v2, v3] = Y[m + v1 - 1, v3] * X[v0, v1, v2, v3]

        @R.function
        def fused(
            x: R.Tensor((1, 1, 32, 128), dtype="float32"),
            y: R.Tensor((2048, 128), dtype="float32"),
            len: R.Shape(["m"]),
        ) -> R.Tensor((1, 1, 32, 128), dtype="float32"):
            R.func_attr({"Primitive": True})
            m = T.int64()
            cls = Before
            with R.dataflow():
                lv1 = R.emit_te(topi.add, x, x)
                gv = R.call_tir(
                    cls.foo,
                    [lv1, y, m],
                    out_ty=R.Tensor((1, 1, 32, 128), dtype="float32"),
                )
                R.output(gv)
            return gv

        @R.function
        def main(
            x: R.Tensor((1, 1, 32, 128), dtype="float32"),
            y: R.Tensor((2048, 128), dtype="float32"),
            len: R.Shape(["m"]),
        ) -> R.Tensor((1, 1, 32, 128), dtype="float32"):
            cls = Before
            with R.dataflow():
                gv = cls.fused(x, y, len)
                R.output(gv)
            return gv

    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(private=True, s_tir=True)
        def fused(
            X: T.Buffer((T.int64(1), T.int64(1), T.int64(32), T.int64(128)), "float32"),
            Y: T.Buffer((T.int64(2048), T.int64(128)), "float32"),
            m: T.int64,
            rotary: T.Buffer((T.int64(1), T.int64(1), T.int64(32), T.int64(128)), "float32"),
        ):
            T.func_attr({"tirx.noalias": True})
            T_add = T.sblock_alloc_buffer((T.int64(1), T.int64(1), T.int64(32), T.int64(128)))
            for ax0, ax1, ax2, ax3 in T.grid(T.int64(1), T.int64(1), T.int64(32), T.int64(128)):
                with T.sblock("T_add"):
                    v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
                    T_add[v_ax0, v_ax1, v_ax2, v_ax3] = (
                        X[v_ax0, v_ax1, v_ax2, v_ax3] + X[v_ax0, v_ax1, v_ax2, v_ax3]
                    )
            for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(1), T.int64(32), T.int64(128)):
                with T.sblock("rotary"):
                    v0, v1, v2, v3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
                    rotary[v0, v1, v2, v3] = Y[m + v1 - T.int64(1), v3] * T_add[v0, v1, v2, v3]

        @R.function
        def main(
            x: R.Tensor((1, 1, 32, 128), dtype="float32"),
            y: R.Tensor((2048, 128), dtype="float32"),
            len: R.Shape(["m"]),
        ) -> R.Tensor((1, 1, 32, 128), dtype="float32"):
            m = T.int64()
            cls = Expected
            with R.dataflow():
                gv = R.call_tir(
                    cls.fused,
                    (x, y, m),
                    out_ty=R.Tensor([1, 1, 32, 128], "float32"),
                )
                R.output(gv)
            return gv

    _check(Before, Expected)


def test_same_buffer_multiple_read():
    @I.ir_module(s_tir=True)
    class Module:
        @T.prim_func(private=True, s_tir=True)
        def concatenate(
            rxplaceholder: T.Buffer((T.int64(1), T.int64(4), T.int64(64), T.int64(64)), "float32"),
            rxplaceholder_1: T.Buffer(
                (T.int64(1), T.int64(4), T.int64(64), T.int64(64)), "float32"
            ),
            T_concat: T.Buffer((T.int64(2), T.int64(4), T.int64(64), T.int64(64)), "float32"),
        ):
            T.func_attr({"op_pattern": 2, "tirx.noalias": True})
            for ax0, ax1, ax2, ax3 in T.grid(T.int64(2), T.int64(4), T.int64(64), T.int64(64)):
                with T.sblock("T_concat"):
                    v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
                    T.reads(
                        rxplaceholder_1[v_ax0 - T.int64(1), v_ax1, v_ax2, v_ax3],
                        rxplaceholder[v_ax0, v_ax1, v_ax2, v_ax3],
                    )
                    T.writes(T_concat[v_ax0, v_ax1, v_ax2, v_ax3])
                    T_concat[v_ax0, v_ax1, v_ax2, v_ax3] = T.if_then_else(
                        T.int64(1) <= v_ax0,
                        rxplaceholder_1[v_ax0 - T.int64(1), v_ax1, v_ax2, v_ax3],
                        rxplaceholder[v_ax0, v_ax1, v_ax2, v_ax3],
                    )

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

        @R.function
        def fused_concatenate_transpose2(
            inp_0: R.Tensor((1, 4, 64, 64), dtype="float32"),
        ) -> R.Tensor((2, 64, 64, 4), dtype="float32"):
            R.func_attr({"Primitive": True})
            cls = Module
            with R.dataflow():
                lv = R.call_tir(
                    cls.concatenate,
                    (inp_0, inp_0),
                    out_ty=R.Tensor((2, 4, 64, 64), dtype="float32"),
                )
                gv = R.call_tir(
                    cls.transpose2, (lv,), out_ty=R.Tensor((2, 64, 64, 4), dtype="float32")
                )
                R.output(gv)
            return gv

        @R.function
        def main(inp_0: R.Tensor((1, 4, 64, 64), dtype="float32")) -> R.Tensor(
            (2, 64, 64, 4), dtype="float32"
        ):
            R.func_attr({"num_input": 3})
            cls = Module
            with R.dataflow():
                lv = cls.fused_concatenate_transpose2(inp_0)
                R.output(lv)
            return lv

    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(private=True, s_tir=True)
        def fused_concatenate_transpose2(
            inp_0: T.Buffer((T.int64(1), T.int64(4), T.int64(64), T.int64(64)), "float32"),
            T_transpose_handle_intermediate: T.Buffer(
                (T.int64(2), T.int64(64), T.int64(64), T.int64(4)), "float32"
            ),
        ):
            T.func_attr({"tirx.noalias": True})
            T_concat_handle_intermediate = T.sblock_alloc_buffer(
                (T.int64(2), T.int64(4), T.int64(64), T.int64(64))
            )
            for ax0, ax1, ax2, ax3 in T.grid(T.int64(2), T.int64(4), T.int64(64), T.int64(64)):
                with T.sblock("T_concat"):
                    v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
                    T.reads(inp_0[v_ax0 - T.int64(1), v_ax1, v_ax2, v_ax3])
                    T.writes(T_concat_handle_intermediate[v_ax0, v_ax1, v_ax2, v_ax3])
                    T_concat_handle_intermediate[v_ax0, v_ax1, v_ax2, v_ax3] = T.if_then_else(
                        T.int64(1) <= v_ax0,
                        inp_0[v_ax0 - T.int64(1), v_ax1, v_ax2, v_ax3],
                        inp_0[v_ax0, v_ax1, v_ax2, v_ax3],
                    )
            for ax0, ax1, ax2, ax3 in T.grid(T.int64(2), T.int64(64), T.int64(64), T.int64(4)):
                with T.sblock("T_transpose"):
                    v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
                    T.reads(T_concat_handle_intermediate[v_ax0, v_ax3, v_ax1, v_ax2])
                    T.writes(T_transpose_handle_intermediate[v_ax0, v_ax1, v_ax2, v_ax3])
                    T_transpose_handle_intermediate[v_ax0, v_ax1, v_ax2, v_ax3] = (
                        T_concat_handle_intermediate[v_ax0, v_ax3, v_ax1, v_ax2]
                    )

        @R.function
        def main(inp_0: R.Tensor((1, 4, 64, 64), dtype="float32")) -> R.Tensor(
            (2, 64, 64, 4), dtype="float32"
        ):
            R.func_attr({"num_input": 3})
            cls = Expected
            with R.dataflow():
                lv = R.call_tir(
                    cls.fused_concatenate_transpose2,
                    (inp_0,),
                    out_ty=R.Tensor((2, 64, 64, 4), dtype="float32"),
                )
                R.output(lv)
            return lv

    _check(Module, Expected)


def test_tir_expression_in_shape():
    @I.ir_module(s_tir=True)
    class Module:
        @R.function
        def fused_transpose_matmul(
            x: R.Tensor((3, 4), dtype="float32"),
            y: R.Tensor(("n - 1", 4), dtype="float32"),
            tir_vars: R.Shape(["n"]),
        ) -> R.Tensor(("n - 1", 3), dtype="float32"):
            R.func_attr({"Primitive": True})
            with R.dataflow():
                lv = R.emit_te(topi.transpose, x)
                gv = R.emit_te(topi.matmul, y, lv)
                R.output(gv)
            return gv

        @R.function
        def main(
            x: R.Tensor((3, 4), dtype="float32"),
            y: R.Tensor(("n - 1", 4), dtype="float32"),
            tir_vars: R.Shape(["n"]),
        ) -> R.Tensor(("n - 1", 3), dtype="float32"):
            cls = Module
            with R.dataflow():
                lv = cls.fused_transpose_matmul(x, y, tir_vars)
                R.output(lv)
            return lv

    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(private=True, s_tir=True)
        def fused_transpose_matmul(
            x: T.Buffer((T.int64(3), T.int64(4)), "float32"),
            p_y: T.handle,
            n: T.int64,
            p_output0: T.handle,
        ):
            T.func_attr({"tirx.noalias": True})
            y = T.match_buffer(p_y, (n - T.int64(1), T.int64(4)))
            var_T_matmul_intermediate = T.match_buffer(p_output0, (n - T.int64(1), T.int64(3)))
            var_T_transpose_intermediate = T.sblock_alloc_buffer((T.int64(4), T.int64(3)))
            for ax0, ax1 in T.grid(T.int64(4), T.int64(3)):
                with T.sblock("T_transpose"):
                    v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
                    var_T_transpose_intermediate[v_ax0, v_ax1] = x[v_ax1, v_ax0]
            for ax0, ax1, k in T.grid(n - T.int64(1), T.int64(3), T.int64(4)):
                with T.sblock("T_matmul"):
                    v_ax0, v_ax1, v_k = T.axis.remap("SSR", [ax0, ax1, k])
                    with T.init():
                        var_T_matmul_intermediate[v_ax0, v_ax1] = T.float32(0)
                    var_T_matmul_intermediate[v_ax0, v_ax1] = (
                        var_T_matmul_intermediate[v_ax0, v_ax1]
                        + y[v_ax0, v_k] * var_T_transpose_intermediate[v_k, v_ax1]
                    )

        @R.function
        def main(
            x: R.Tensor((3, 4), dtype="float32"),
            y: R.Tensor(("n - 1", 4), dtype="float32"),
            tir_vars: R.Shape(["n"]),
        ) -> R.Tensor(("n - 1", 3), dtype="float32"):
            n = T.int64()
            cls = Expected
            with R.dataflow():
                lv = R.call_tir(
                    cls.fused_transpose_matmul,
                    (x, y, n),
                    out_ty=R.Tensor((n - 1, 3), dtype="float32"),
                )
                R.output(lv)
            return lv

    _check(Module, Expected)


def test_tuple_input_unused_field():
    @I.ir_module(s_tir=True)
    class Module:
        @T.prim_func(private=True, s_tir=True)
        def reshape(
            A: T.Buffer((T.int64(4), T.int64(8), T.int64(2048)), "float32"),
            T_reshape: T.Buffer((T.int64(4), T.int64(8), T.int64(32), T.int64(64)), "float32"),
        ):
            T.func_attr({"op_pattern": 2, "tirx.noalias": True})
            # with T.sblock("root"):
            for ax0, ax1, ax2, ax3 in T.grid(T.int64(4), T.int64(8), T.int64(32), T.int64(64)):
                with T.sblock("T_reshape"):
                    v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
                    T.reads(
                        A[
                            (
                                ((v_ax2 * T.int64(64) + v_ax3) // T.int64(2048) + v_ax1)
                                // T.int64(8)
                                + v_ax0
                            )
                            % T.int64(4),
                            ((v_ax2 * T.int64(64) + v_ax3) // T.int64(2048) + v_ax1) % T.int64(8),
                            (v_ax2 * T.int64(64) + v_ax3) % T.int64(2048),
                        ]
                    )
                    T.writes(T_reshape[v_ax0, v_ax1, v_ax2, v_ax3])
                    T_reshape[v_ax0, v_ax1, v_ax2, v_ax3] = A[
                        (
                            ((v_ax2 * T.int64(64) + v_ax3) // T.int64(2048) + v_ax1) // T.int64(8)
                            + v_ax0
                        )
                        % T.int64(4),
                        ((v_ax2 * T.int64(64) + v_ax3) // T.int64(2048) + v_ax1) % T.int64(8),
                        (v_ax2 * T.int64(64) + v_ax3) % T.int64(2048),
                    ]

        @R.function(private=True)
        def fused_reshape(
            lv: R.Tuple(
                R.Tensor((4, 8, 2048), dtype="float32"), R.Tensor((4, 8, 2048), dtype="float32")
            ),
        ) -> R.Tensor((4, 8, 32, 64), dtype="float32"):
            R.func_attr({"Primitive": True})
            cls = Module
            with R.dataflow():
                lv1: R.Tensor((4, 8, 2048), dtype="float32") = lv[0]
                gv = R.call_tir(
                    cls.reshape, (lv1,), out_ty=R.Tensor((4, 8, 32, 64), dtype="float32")
                )
                R.output(gv)
            return gv

        @R.function
        def main(
            tup: R.Tuple(
                R.Tensor((4, 8, 2048), dtype="float32"), R.Tensor((4, 8, 2048), dtype="float32")
            ),
        ) -> R.Tensor((4, 8, 32, 64), dtype="float32"):
            cls = Module
            with R.dataflow():
                lv_1: R.Tensor((4, 8, 32, 64), dtype="float32") = cls.fused_reshape(tup)
                R.output(lv_1)
            return lv_1

    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(private=True, s_tir=True)
        def fused_reshape(
            lv_0: T.Buffer((T.int64(4), T.int64(8), T.int64(2048)), "float32"),
            T_reshape_handle_intermediate: T.Buffer(
                (T.int64(4), T.int64(8), T.int64(32), T.int64(64)), "float32"
            ),
        ):
            T.func_attr({"tirx.noalias": True})
            # with T.sblock("root"):
            for ax0, ax1, ax2, ax3 in T.grid(T.int64(4), T.int64(8), T.int64(32), T.int64(64)):
                with T.sblock("T_reshape"):
                    v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
                    T.reads(
                        lv_0[
                            (
                                ((v_ax2 * T.int64(64) + v_ax3) // T.int64(2048) + v_ax1)
                                // T.int64(8)
                                + v_ax0
                            )
                            % T.int64(4),
                            ((v_ax2 * T.int64(64) + v_ax3) // T.int64(2048) + v_ax1) % T.int64(8),
                            (v_ax2 * T.int64(64) + v_ax3) % T.int64(2048),
                        ]
                    )
                    T.writes(T_reshape_handle_intermediate[v_ax0, v_ax1, v_ax2, v_ax3])
                    T_reshape_handle_intermediate[v_ax0, v_ax1, v_ax2, v_ax3] = lv_0[
                        (
                            ((v_ax2 * T.int64(64) + v_ax3) // T.int64(2048) + v_ax1) // T.int64(8)
                            + v_ax0
                        )
                        % T.int64(4),
                        ((v_ax2 * T.int64(64) + v_ax3) // T.int64(2048) + v_ax1) % T.int64(8),
                        (v_ax2 * T.int64(64) + v_ax3) % T.int64(2048),
                    ]

        @R.function
        def main(
            tup: R.Tuple(
                R.Tensor((4, 8, 2048), dtype="float32"), R.Tensor((4, 8, 2048), dtype="float32")
            ),
        ) -> R.Tensor((4, 8, 32, 64), dtype="float32"):
            cls = Expected
            with R.dataflow():
                lv: R.Tensor((4, 8, 2048), dtype="float32") = tup[0]
                lv_1 = R.call_tir(
                    cls.fused_reshape, (lv,), out_ty=R.Tensor((4, 8, 32, 64), dtype="float32")
                )
                R.output(lv_1)
            return lv_1

    _check(Module, Expected)


def test_unique_duplicated_buffer_allocation():
    @I.ir_module(s_tir=True)
    class Module:
        @T.prim_func(private=True, s_tir=True)
        def add(
            A: T.Buffer((T.int64(4096), T.int64(4096)), "float16"),
            Out: T.Buffer((T.int64(4096), T.int64(4096)), "float16"),
        ):
            for i, j in T.grid(T.int64(4096), T.int64(4096)):
                with T.sblock("add"):
                    vi, vj = T.axis.remap("SS", [i, j])
                    Out[vi, vj] = A[vi, vj] + T.float16(1.0)

        @T.prim_func(private=True, s_tir=True)
        def add1(
            A: T.Buffer((T.int64(4096), T.int64(4096)), "float16"),
            Out: T.Buffer((T.int64(4096), T.int64(4096)), "float16"),
        ):
            for i, j in T.grid(T.int64(4096), T.int64(4096)):
                with T.sblock("add"):
                    vi, vj = T.axis.remap("SS", [i, j])
                    Out[vi, vj] = A[vi, vj] + T.float16(2.0)

        @R.function
        def main(
            input_embeds: R.Tensor((4096, 4096), dtype="float16"),
        ) -> R.Tensor((4096, 4096), dtype="float16"):
            cls = Module
            with R.dataflow():
                gv: R.Tensor((4096, 4096), dtype="float16") = cls.fused_func(input_embeds)
                R.output(gv)
            return gv

        @R.function(private=True)
        def fused_func(
            input_embeds: R.Tensor((4096, 4096), dtype="float16"),
        ) -> R.Tensor((4096, 4096), dtype="float16"):
            R.func_attr({"Primitive": True})
            cls = Module
            with R.dataflow():
                lv = R.call_tir(
                    cls.add, (input_embeds,), out_ty=R.Tensor((4096, 4096), dtype="float16")
                )
                gv = R.call_tir(cls.add1, (lv,), out_ty=R.Tensor((4096, 4096), dtype="float16"))
                R.output(gv)
            return gv

    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(private=True, s_tir=True)
        def fused_func(
            input_embeds: T.Buffer((T.int64(4096), T.int64(4096)), "float16"),
            Out_intermediate_1: T.Buffer((T.int64(4096), T.int64(4096)), "float16"),
        ):
            T.func_attr({"tirx.noalias": True})
            Out_intermediate = T.sblock_alloc_buffer((T.int64(4096), T.int64(4096)), "float16")
            for i, j in T.grid(T.int64(4096), T.int64(4096)):
                with T.sblock("add"):
                    vi, vj = T.axis.remap("SS", [i, j])
                    T.reads(input_embeds[vi, vj])
                    T.writes(Out_intermediate[vi, vj])
                    Out_intermediate[vi, vj] = input_embeds[vi, vj] + T.float16(1)
            for i, j in T.grid(T.int64(4096), T.int64(4096)):
                with T.sblock("add_1"):
                    vi, vj = T.axis.remap("SS", [i, j])
                    T.reads(Out_intermediate[vi, vj])
                    T.writes(Out_intermediate_1[vi, vj])
                    Out_intermediate_1[vi, vj] = Out_intermediate[vi, vj] + T.float16(2)

        @R.function
        def main(input_embeds: R.Tensor((4096, 4096), dtype="float16")) -> R.Tensor(
            (4096, 4096), dtype="float16"
        ):
            cls = Expected
            with R.dataflow():
                gv = R.call_tir(
                    cls.fused_func,
                    (input_embeds,),
                    out_ty=R.Tensor((4096, 4096), dtype="float16"),
                )
                R.output(gv)
            return gv

    _check(Module, Expected)


def test_extern_func():
    bb = relax.BlockBuilder()
    bb.add_func(relax.extern("extern_func"), "extern_func")
    mod = bb.get()
    # FuseTIR should keep the ExternFunc in the IRModule.
    _check(mod, mod)


def test_symbolic_var_in_buffer_shape():
    """A PrimFunc may have dynamic buffer shapes

    Symbolic variables in a PrimFunc may be present in the buffer
    shape without a corresponding parameter.  These symbolic variables
    are inferred from the buffer's shape.  (Or, at runtime, they are
    typically determined from the DLTensor's known shape.)
    """

    @I.ir_module(s_tir=True)
    class Before:
        @T.prim_func(private=True, s_tir=True)
        def foo(
            X_handle: T.handle,
            Y: T.Buffer((T.int64(2048), T.int64(128)), "float32"),
            m: T.int64,
            rotary_handle: T.handle,
        ):
            sequence_length = T.int64()

            X = T.match_buffer(
                X_handle, [T.int64(1), sequence_length, T.int64(32), T.int64(128)], "float32"
            )
            rotary = T.match_buffer(
                rotary_handle, [T.int64(1), sequence_length, T.int64(32), T.int64(128)], "float32"
            )

            for i0, i1, i2, i3 in T.grid(T.int64(1), sequence_length, T.int64(32), T.int64(128)):
                with T.sblock("rotary"):
                    v0, v1, v2, v3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
                    rotary[v0, v1, v2, v3] = Y[m + v1 - 1, v3] * X[v0, v1, v2, v3]

        @R.function
        def fused(
            x: R.Tensor((1, "sequence_length", 32, 128), dtype="float32"),
            y: R.Tensor((2048, 128), dtype="float32"),
            len: R.Shape(["m"]),
        ) -> R.Tensor((1, "sequence_length", 32, 128), dtype="float32"):
            R.func_attr({"Primitive": True})
            sequence_length = T.int64()
            m = T.int64()
            cls = Before
            with R.dataflow():
                lv1 = R.emit_te(topi.add, x, x)
                gv = R.call_tir(
                    cls.foo,
                    [lv1, y, m],
                    out_ty=R.Tensor((1, sequence_length, 32, 128), dtype="float32"),
                )
                R.output(gv)
            return gv

        @R.function
        def main(
            x: R.Tensor((1, "sequence_length", 32, 128), dtype="float32"),
            y: R.Tensor((2048, 128), dtype="float32"),
            len: R.Shape(["m"]),
        ) -> R.Tensor((1, "sequence_length", 32, 128), dtype="float32"):
            cls = Before
            with R.dataflow():
                gv = cls.fused(x, y, len)
                R.output(gv)
            return gv

    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(private=True, s_tir=True)
        def fused(
            X_handle: T.handle,
            Y: T.Buffer((T.int64(2048), T.int64(128)), "float32"),
            m: T.int64,
            rotary_handle: T.handle,
        ):
            T.func_attr({"tirx.noalias": True})

            sequence_length = T.int64()

            X = T.match_buffer(
                X_handle, [T.int64(1), sequence_length, T.int64(32), T.int64(128)], "float32"
            )
            rotary = T.match_buffer(
                rotary_handle, [T.int64(1), sequence_length, T.int64(32), T.int64(128)], "float32"
            )

            T_add = T.sblock_alloc_buffer((T.int64(1), sequence_length, T.int64(32), T.int64(128)))
            for ax0, ax1, ax2, ax3 in T.grid(
                T.int64(1), sequence_length, T.int64(32), T.int64(128)
            ):
                with T.sblock("T_add"):
                    v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
                    T_add[v_ax0, v_ax1, v_ax2, v_ax3] = (
                        X[v_ax0, v_ax1, v_ax2, v_ax3] + X[v_ax0, v_ax1, v_ax2, v_ax3]
                    )
            for i0, i1, i2, i3 in T.grid(T.int64(1), sequence_length, T.int64(32), T.int64(128)):
                with T.sblock("rotary"):
                    v0, v1, v2, v3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
                    rotary[v0, v1, v2, v3] = Y[m + v1 - T.int64(1), v3] * T_add[v0, v1, v2, v3]

        @R.function
        def main(
            x: R.Tensor((1, "sequence_length", 32, 128), dtype="float32"),
            y: R.Tensor((2048, 128), dtype="float32"),
            len: R.Shape(["m"]),
        ) -> R.Tensor((1, "sequence_length", 32, 128), dtype="float32"):
            sequence_length = T.int64()
            m = T.int64()
            cls = Expected
            with R.dataflow():
                gv = R.call_tir(
                    cls.fused,
                    (x, y, m),
                    out_ty=R.Tensor([1, sequence_length, 32, 128], "float32"),
                )
                R.output(gv)
            return gv

    _check(Before, Expected)


def test_symbolic_var_called_with_static_shape():
    """A dynamic PrimFunc may be called with a static shape"""

    @I.ir_module(s_tir=True)
    class Before:
        @T.prim_func(private=True, s_tir=True)
        def sum_1d(
            X_handle: T.handle,
            Y: T.Buffer([T.int64(1)], "float32"),
        ):
            num_elements = T.int64()

            X = T.match_buffer(X_handle, [num_elements], "float32")

            for i in range(num_elements):
                with T.sblock("sum"):
                    vi = T.axis.remap("R", [i])
                    with T.init():
                        Y[0] = 0.0
                    Y[0] = Y[0] + X[vi]

        @R.function(private=True)
        def fused(
            x: R.Tensor([64], dtype="float32"),
        ) -> R.Tensor([1], dtype="float32"):
            R.func_attr({"Primitive": True})
            cls = Before
            with R.dataflow():
                gv = R.call_tir(
                    cls.sum_1d,
                    [x],
                    out_ty=R.Tensor([1], dtype="float32"),
                )
                R.output(gv)
            return gv

        @R.function
        def main(
            x: R.Tensor([64], dtype="float32"),
        ) -> R.Tensor([1], dtype="float32"):
            cls = Before
            with R.dataflow():
                gv = cls.fused(x)
                R.output(gv)
            return gv

    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(private=True, s_tir=True)
        def fused(
            X: T.Buffer([T.int64(64)], "float32"),
            Y: T.Buffer([T.int64(1)], "float32"),
        ):
            T.func_attr({"tirx.noalias": True})

            for i in range(T.int64(64)):
                with T.sblock("sum"):
                    vi = T.axis.remap("R", [i])
                    with T.init():
                        Y[0] = 0.0
                    Y[0] = Y[0] + X[vi]

        @R.function
        def main(
            x: R.Tensor([64], dtype="float32"),
        ) -> R.Tensor([1], dtype="float32"):
            cls = Expected
            with R.dataflow():
                gv = R.call_tir(cls.fused, (x,), out_ty=R.Tensor((1,), dtype="float32"))
                R.output(gv)
            return gv

    _check(Before, Expected)


def test_symbolic_var_called_with_multiple_static_shapes():
    """A dynamic PrimFunc may be called with different shapes each time"""

    @I.ir_module(s_tir=True)
    class Before:
        @T.prim_func(private=True, s_tir=True)
        def sum_1d(
            X_handle: T.handle,
            Sum: T.Buffer([T.int64(1)], "float32"),
        ):
            num_elements = T.int64()

            X = T.match_buffer(X_handle, [num_elements], "float32")

            for i in range(num_elements):
                with T.sblock("sum"):
                    vi = T.axis.remap("R", [i])
                    with T.init():
                        Sum[0] = 0.0
                    Sum[0] = Sum[0] + X[vi]

        @T.prim_func(private=True, s_tir=True)
        def sum_scalar(
            X: T.Buffer([T.int64(1)], "float32"),
            Y: T.Buffer([T.int64(1)], "float32"),
            Sum: T.Buffer([T.int64(1)], "float32"),
        ):
            for i in range(T.int64(1)):
                with T.sblock("Out"):
                    vi = T.axis.remap("S", [i])
                    Sum[vi] = X[vi] + Y[vi]

        @R.function(private=True)
        def fused(
            x: R.Tensor([64], dtype="float32"),
            y: R.Tensor([16], dtype="float32"),
        ) -> R.Tensor([1], dtype="float32"):
            R.func_attr({"Primitive": True})
            cls = Before
            with R.dataflow():
                x_sum = R.call_tir(
                    cls.sum_1d,
                    [x],
                    out_ty=R.Tensor([1], dtype="float32"),
                )
                y_sum = R.call_tir(
                    cls.sum_1d,
                    [y],
                    out_ty=R.Tensor([1], dtype="float32"),
                )
                gv = R.call_tir(
                    cls.sum_scalar,
                    [x_sum, y_sum],
                    out_ty=R.Tensor([1], dtype="float32"),
                )
                R.output(gv)
            return gv

        @R.function
        def main(
            x: R.Tensor([64], dtype="float32"),
            y: R.Tensor([16], dtype="float32"),
        ) -> R.Tensor([1], dtype="float32"):
            cls = Before
            with R.dataflow():
                gv = cls.fused(x, y)
                R.output(gv)
            return gv

    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(private=True, s_tir=True)
        def fused(
            X: T.Buffer([T.int64(64)], "float32"),
            Y: T.Buffer([T.int64(16)], "float32"),
            Out: T.Buffer([T.int64(1)], "float32"),
        ):
            T.func_attr({"tirx.noalias": True})

            XSum = T.sblock_alloc_buffer([T.int64(1)], "float32")
            YSum = T.sblock_alloc_buffer([T.int64(1)], "float32")

            for i in range(T.int64(64)):
                with T.sblock("XSum"):
                    vi = T.axis.remap("R", [i])
                    with T.init():
                        XSum[0] = 0.0
                    XSum[0] = XSum[0] + X[vi]

            for i in range(T.int64(16)):
                with T.sblock("YSum"):
                    vi = T.axis.remap("R", [i])
                    with T.init():
                        YSum[0] = 0.0
                    YSum[0] = YSum[0] + Y[vi]

            for i in range(T.int64(1)):
                with T.sblock("Out"):
                    vi = T.axis.remap("S", [i])
                    Out[vi] = XSum[vi] + YSum[vi]

        @R.function
        def main(
            x: R.Tensor([64], dtype="float32"),
            y: R.Tensor([16], dtype="float32"),
        ) -> R.Tensor([1], dtype="float32"):
            cls = Expected
            with R.dataflow():
                gv = R.call_tir(cls.fused, (x, y), out_ty=R.Tensor((1,), dtype="float32"))
                R.output(gv)
            return gv

    _check(Before, Expected)


def test_symbolic_var_called_with_static_argument():
    """A dynamic PrimFunc may accept a static argument

    Primitive arguments in `R.call_tir` contain definitions for all TIR
    variables explicitly listed in the function signature, and contain
    the TIR expression to be passed for each parameter.

    This test is identical to the earlier test named
    "test_symbolic_var_called_with_static_shape", except for the
    explicit parameter in `sum_1d`.
    """

    @I.ir_module(s_tir=True)
    class Before:
        @T.prim_func(private=True, s_tir=True)
        def sum_1d(
            X_handle: T.handle,
            num_elements: T.int64,
            Y: T.Buffer([T.int64(1)], "float32"),
        ):
            X = T.match_buffer(X_handle, [num_elements], "float32")

            for i in range(num_elements):
                with T.sblock("sum"):
                    vi = T.axis.remap("R", [i])
                    with T.init():
                        Y[0] = 0.0
                    Y[0] = Y[0] + X[vi]

        @R.function(private=True)
        def fused(
            x: R.Tensor([64], dtype="float32"),
        ) -> R.Tensor([1], dtype="float32"):
            R.func_attr({"Primitive": True})
            cls = Before
            with R.dataflow():
                gv = R.call_tir(
                    cls.sum_1d,
                    [x, 64],
                    out_ty=R.Tensor([1], dtype="float32"),
                )
                R.output(gv)
            return gv

        @R.function
        def main(
            x: R.Tensor([64], dtype="float32"),
        ) -> R.Tensor([1], dtype="float32"):
            cls = Before
            with R.dataflow():
                gv = cls.fused(x)
                R.output(gv)
            return gv

    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(private=True, s_tir=True)
        def fused(
            X: T.Buffer([T.int64(64)], "float32"),
            Y: T.Buffer([T.int64(1)], "float32"),
        ):
            T.func_attr({"tirx.noalias": True})

            for i in range(T.int64(64)):
                with T.sblock("sum"):
                    vi = T.axis.remap("R", [i])
                    with T.init():
                        Y[0] = 0.0
                    Y[0] = Y[0] + X[vi]

        @R.function
        def main(
            x: R.Tensor([64], dtype="float32"),
        ) -> R.Tensor([1], dtype="float32"):
            cls = Expected
            with R.dataflow():
                gv = R.call_tir(cls.fused, (x,), out_ty=R.Tensor((1,), dtype="float32"))
                R.output(gv)
            return gv

    _check(Before, Expected)


def test_gather():
    @I.ir_module(s_tir=True)
    class Before:
        @T.prim_func(private=True, s_tir=True)
        def add(
            A: T.Buffer((T.int64(4096), T.int64(4096)), "float16"),
            Out: T.Buffer((T.int64(4096), T.int64(4096)), "float16"),
        ):
            for i, j in T.grid(T.int64(4096), T.int64(4096)):
                with T.sblock("add"):
                    vi, vj = T.axis.remap("SS", [i, j])
                    Out[vi, vj] = A[vi, vj] + T.float16(1.0)

        @T.prim_func(private=True, s_tir=True)
        def take(
            A: T.Buffer((T.int64(4096), T.int64(4096)), "float16"),
            B: T.Buffer((T.int64(1),), "int32"),
            T_take: T.Buffer((T.int64(1), T.int64(4096)), "float16"),
        ):
            for ax0, ax1 in T.grid(T.int64(1), T.int64(4096)):
                with T.sblock("T_take"):
                    v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
                    T_take[v_ax0, v_ax1] = A[B[v_ax0], v_ax1]

        @R.function
        def main(
            input_ids: R.Tensor((1,), dtype="int32"),
            input_embeds: R.Tensor((4096, 4096), dtype="float16"),
        ) -> R.Tensor((1, 4096), dtype="float16"):
            cls = Before
            with R.dataflow():
                gv: R.Tensor((1, 4096), dtype="float16") = cls.fused_func(input_ids, input_embeds)
                R.output(gv)
            return gv

        @R.function(private=True)
        def fused_func(
            input_ids: R.Tensor((1,), dtype="int32"),
            input_embeds: R.Tensor((4096, 4096), dtype="float16"),
        ) -> R.Tensor((1, 4096), dtype="float16"):
            R.func_attr({"Primitive": True})
            cls = Before
            with R.dataflow():
                lv = R.call_tir(
                    cls.add, (input_embeds,), out_ty=R.Tensor((4096, 4096), dtype="float16")
                )
                gv = R.call_tir(
                    cls.take, (lv, input_ids), out_ty=R.Tensor((1, 4096), dtype="float16")
                )
                R.output(gv)
            return gv

    @I.ir_module(s_tir=True)
    class After:
        @T.prim_func(private=True, s_tir=True)
        def fused_func(
            input_ids: T.Buffer((T.int64(1),), "int32"),
            input_embeds: T.Buffer((T.int64(4096), T.int64(4096)), "float16"),
            T_take: T.Buffer((T.int64(1), T.int64(4096)), "float16"),
        ):
            T.func_attr({"tirx.noalias": True})
            Out_handle_intermediate = T.sblock_alloc_buffer(
                (T.int64(4096), T.int64(4096)), "float16"
            )
            for i, j in T.grid(T.int64(4096), T.int64(4096)):
                with T.sblock("add"):
                    vi, vj = T.axis.remap("SS", [i, j])
                    Out_handle_intermediate[vi, vj] = input_embeds[vi, vj] + T.float16(1)
            for ax0, ax1 in T.grid(T.int64(1), T.int64(4096)):
                with T.sblock("T_take"):
                    v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
                    T_take[v_ax0, v_ax1] = Out_handle_intermediate[input_ids[v_ax0], v_ax1]

        @R.function
        def main(
            input_ids: R.Tensor((1,), dtype="int32"),
            input_embeds: R.Tensor((4096, 4096), dtype="float16"),
        ) -> R.Tensor((1, 4096), dtype="float16"):
            cls = After
            with R.dataflow():
                gv = R.call_tir(
                    cls.fused_func,
                    (input_ids, input_embeds),
                    out_ty=R.Tensor((1, 4096), dtype="float16"),
                )
                R.output(gv)
            return gv

    _check(Before, After)


def test_inplace_simple():
    @I.ir_module(s_tir=True)
    class Module:
        I.module_attrs({"foo": "bar"})

        @T.prim_func(private=True, s_tir=True)
        def add_inplace(
            A: T.Buffer((T.int64(10), T.int64(20)), "float32"), B: T.Buffer((), "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(A[v_ax0, v_ax1])
                    A[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(private=True)
        def fused_add_exp_squeeze(
            x: R.Tensor((10, 20), dtype="float32"), p0: R.Tensor((), dtype="float32")
        ) -> R.Tensor((10, 20), dtype="float32"):
            R.func_attr({"Primitive": True})
            cls = Module
            with R.dataflow():
                # This overwrites x and is actually evil because the function is marked as pure
                # but we are doing it just to test the pass. The automatic DataflowUseInplaceCalls
                # transformation will not produce code like this, but it may make sense to do it
                # if ownership of x is fully and truly transferred.
                # Users should apply with caution!
                lv = R.call_tir_inplace(
                    cls.add_inplace,
                    (x, p0),
                    inplace_indices=[0],
                    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 = Module
            with R.dataflow():
                gv1: R.Tensor((10, 20), dtype="float32") = cls.fused_add_exp_squeeze(x, p0)
                R.output(gv1)
            return gv1

    @I.ir_module(s_tir=True)
    class Expected:
        I.module_attrs({"foo": "bar"})

        @T.prim_func(private=True, s_tir=True)
        def fused_add_exp_squeeze(
            x: T.Buffer((T.int64(10), T.int64(20)), "float32"), p0: T.Buffer((), "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])
                    x[v_ax0, v_ax1] = x[v_ax0, v_ax1] + p0[()]
            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])
                    x[v_i0, v_i1] = T.exp(x[v_i0, v_i1])
            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])
                    x[v_ax0, v_ax1] = x[v_ax0, v_ax1]

        # note that this will clobber x! Use with caution
        @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") = R.call_tir_inplace(
                    cls.fused_add_exp_squeeze,
                    (x, p0),
                    out_ty=R.Tensor((10, 20), dtype="float32"),
                    inplace_indices=[0],
                )
                R.output(gv1)
            return gv1

    _check(Module, Expected)


def test_fuse_inplace_and_non_inplace():
    @I.ir_module(s_tir=True)
    class Module:
        I.module_attrs({"foo": "bar"})

        @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])
                    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])
                    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])
                    A[v_ax0, v_ax1] = A[v_ax0, v_ax1]

        @R.function(private=True)
        def fused_add_exp_squeeze(
            x: R.Tensor((10, 20), dtype="float32"), p0: R.Tensor((), dtype="float32")
        ) -> R.Tensor((10, 20), dtype="float32"):
            R.func_attr({"Primitive": True})
            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

        @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():
                gv1: R.Tensor((10, 20), dtype="float32") = cls.fused_add_exp_squeeze(x, p0)
                R.output(gv1)
            return gv1

    @I.ir_module(s_tir=True)
    class Expected:
        I.module_attrs({"foo": "bar"})

        @T.prim_func(private=True, s_tir=True)
        def fused_add_exp_squeeze(
            x: T.Buffer((T.int64(10), T.int64(20)), "float32"),
            p0: T.Buffer((), "float32"),
            p_output0: 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])
                    p_output0[v_ax0, v_ax1] = x[v_ax0, v_ax1] + p0[()]
            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])
                    p_output0[v_i0, v_i1] = T.exp(p_output0[v_i0, v_i1])
            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])
                    p_output0[v_ax0, v_ax1] = p_output0[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 = Expected
            with R.dataflow():
                gv1: R.Tensor((10, 20), dtype="float32") = R.call_tir(
                    cls.fused_add_exp_squeeze,
                    (x, p0),
                    out_ty=R.Tensor((10, 20), dtype="float32"),
                )
                R.output(gv1)
            return gv1

    _check(Module, Expected)


def test_use_as_inplace_and_dps():
    @I.ir_module(s_tir=True)
    class Module:
        # we will use it both in-place and normally (DPS)
        @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])
                    Out[v_ax0, v_ax1] = A[v_ax0, v_ax1] + B[()]

        @R.function(private=True)
        def fused_sums(
            x: R.Tensor((10, 20), dtype="float32"), p0: R.Tensor((), dtype="float32")
        ) -> R.Tensor((10, 20), dtype="float32"):
            R.func_attr({"Primitive": True})
            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.add,
                    (x, p0, lv),
                    inplace_indices=[2],
                    out_ty=R.Tensor((10, 20), dtype="float32"),
                )
                lv2 = R.call_tir_inplace(
                    cls.add,
                    (x, p0, lv1),
                    inplace_indices=[2],
                    out_ty=R.Tensor((10, 20), dtype="float32"),
                )
                R.output(lv2)
            return lv2

        @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():
                gv1: R.Tensor((10, 20), dtype="float32") = cls.fused_sums(x, p0)
                R.output(gv1)
            return gv1

    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(private=True, s_tir=True)
        def fused_sums(
            x: T.Buffer((T.int64(10), T.int64(20)), "float32"),
            p0: T.Buffer((), "float32"),
            p_output0: 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])
                    p_output0[v_ax0, v_ax1] = x[v_ax0, v_ax1] + p0[()]
            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])
                    p_output0[v_ax0, v_ax1] = x[v_ax0, v_ax1] + p0[()]
            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])
                    p_output0[v_ax0, v_ax1] = x[v_ax0, v_ax1] + p0[()]

        @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") = R.call_tir(
                    cls.fused_sums,
                    (x, p0),
                    out_ty=R.Tensor((10, 20), dtype="float32"),
                )
                R.output(gv1)
            return gv1

    _check(Module, Expected)


def test_private_nonprimitive_func():
    """Input IRModule may contain calls to non-primitive functions

    This is a regression test.  Prior implementations did not preserve
    relax-to-relax function calls.
    """

    @I.ir_module(s_tir=True)
    class Before:
        @R.function
        def main(
            input_ids: R.Tensor((1,), dtype="int32"),
            input_embeds: R.Tensor((4096, 4096), dtype="float16"),
        ) -> R.Tensor((1, 4096), dtype="float16"):
            cls = Before
            with R.dataflow():
                gv = cls.fused_func(input_ids, input_embeds)
                R.output(gv)
            return gv

        @R.function(private=True)
        def fused_func(
            input_ids: R.Tensor((1,), dtype="int32"),
            input_embeds: R.Tensor((4096, 4096), dtype="float16"),
        ) -> R.Tensor((1, 4096), dtype="float16"):
            cls = Before
            with R.dataflow():
                lv = R.call_tir(
                    cls.add, (input_embeds,), out_ty=R.Tensor((4096, 4096), dtype="float16")
                )
                gv = R.call_tir(
                    cls.take, (lv, input_ids), out_ty=R.Tensor((1, 4096), dtype="float16")
                )
                R.output(gv)
            return gv

        @T.prim_func(private=True, s_tir=True)
        def add(
            A: T.Buffer((T.int64(4096), T.int64(4096)), "float16"),
            Out: T.Buffer((T.int64(4096), T.int64(4096)), "float16"),
        ):
            for i, j in T.grid(T.int64(4096), T.int64(4096)):
                with T.sblock("add"):
                    vi, vj = T.axis.remap("SS", [i, j])
                    Out[vi, vj] = A[vi, vj] + T.float16(1.0)

        @T.prim_func(private=True, s_tir=True)
        def take(
            A: T.Buffer((T.int64(4096), T.int64(4096)), "float16"),
            B: T.Buffer((T.int64(1),), "int32"),
            T_take: T.Buffer((T.int64(1), T.int64(4096)), "float16"),
        ):
            for ax0, ax1 in T.grid(T.int64(1), T.int64(4096)):
                with T.sblock("T_take"):
                    v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
                    T_take[v_ax0, v_ax1] = A[B[v_ax0], v_ax1]

    _check(Before, Before)


def test_block_name_numeric_suffix_deduplication():
    @I.ir_module(s_tir=True)
    class Before:
        @T.prim_func(private=True, s_tir=True)
        def add1(x: T.Buffer((10,), "float32"), y: T.Buffer((10,), "float32")):
            T.func_attr({"tirx.noalias": True})
            for i in range(10):
                with T.sblock("compute1"):
                    vi = T.axis.spatial(10, i)
                    y[vi] = x[vi] + T.float32(1.0)

        @T.prim_func(private=True, s_tir=True)
        def mul1(x: T.Buffer((10,), "float32"), y: T.Buffer((10,), "float32")):
            T.func_attr({"tirx.noalias": True})
            for i in range(10):
                with T.sblock("compute1"):
                    vi = T.axis.spatial(10, i)
                    y[vi] = x[vi] * T.float32(2.0)

        @R.function(private=True)
        def fused_add_mul(x: R.Tensor((10,), "float32")) -> R.Tensor((10,), dtype="float32"):
            R.func_attr({"Primitive": True})
            cls = Before
            with R.dataflow():
                lv1 = R.call_tir(cls.add1, (x,), out_ty=R.Tensor((10,), dtype="float32"))
                lv2 = R.call_tir(cls.mul1, (lv1,), out_ty=R.Tensor((10,), dtype="float32"))
                R.output(lv2)
            return lv2

        @R.function
        def main(x: R.Tensor((10,), dtype="float32")) -> R.Tensor((10,), dtype="float32"):
            cls = Before
            with R.dataflow():
                gv = cls.fused_add_mul(x)
                R.output(gv)
            return gv

    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(private=True, s_tir=True)
        def fused_add_mul(p_x: T.handle, p_output0: T.handle):
            T.func_attr({"tirx.noalias": True})
            x = T.match_buffer(p_x, (T.int64(10),))
            y_intermediate_1 = T.match_buffer(p_output0, (T.int64(10),), elem_offset=T.int32(0))
            with T.sblock("root"):
                T.reads()
                T.writes()
                y_intermediate = T.sblock_alloc_buffer((T.int64(10),), elem_offset=T.int32(0))
                for i in range(10):
                    with T.sblock("compute1"):
                        vi = T.axis.spatial(10, i)
                        T.reads(x[vi])
                        T.writes(y_intermediate[vi])
                        y_intermediate[vi] = x[vi] + T.float32(1.0)
                for i in range(10):
                    with T.sblock("compute2"):
                        vi = T.axis.spatial(10, i)
                        T.reads(y_intermediate[vi])
                        T.writes(y_intermediate_1[vi])
                        y_intermediate_1[vi] = y_intermediate[vi] * T.float32(2.0)

        @R.function
        def main(x: R.Tensor((10,), dtype="float32")) -> R.Tensor((10,), dtype="float32"):
            cls = Expected
            with R.dataflow():
                gv = R.call_tir(cls.fused_add_mul, (x,), out_ty=R.Tensor((10,), dtype="float32"))
                R.output(gv)
            return gv

    _check(Before, Expected)


def test_primitive_scalar_parameter_preserves_identity():
    @I.ir_module(s_tir=True)
    class Before:
        @T.prim_func(private=True, s_tir=True)
        def add_scalar(x: T.Buffer((4,), "int64"), p: T.int64, y: T.Buffer((1,), "int64")):
            for i in range(1):
                with T.sblock("add"):
                    vi = T.axis.spatial(1, i)
                    y[vi] = x[0] + p

        @R.function(private=True)
        def fused(x: R.Tensor((4,), "int64"), p: R.Prim("int64")) -> R.Tensor((1,), "int64"):
            R.func_attr({"Primitive": True})
            cls = Before
            out = R.call_tir(
                cls.add_scalar,
                (x, p),
                out_ty=R.Tensor((1,), "int64"),
            )
            return out

        @R.function
        def main(x: R.Tensor((4,), "int64"), p: R.Prim("int64")) -> R.Tensor((1,), "int64"):
            cls = Before
            out = cls.fused(x, p)
            return out

    after = relax.transform.FuseTIR()(Before)
    assert relax.analysis.check_well_formed(after)
    assert tvm.tirx.analysis.verify_well_formed(after["fused"])


def test_inplace_argument_after_primitive_scalar():
    @I.ir_module(s_tir=True)
    class Before:
        @T.prim_func(private=True, s_tir=True)
        def add_scalar_inplace(p: T.int64, x: T.Buffer((4,), "int64")):
            for i in range(4):
                with T.sblock("add"):
                    vi = T.axis.spatial(4, i)
                    x[vi] = x[vi] + p

        @R.function(private=True)
        def fused(p: R.Prim("int64"), x: R.Tensor((4,), "int64")) -> R.Tensor((4,), "int64"):
            R.func_attr({"Primitive": True})
            cls = Before
            out = R.call_tir_inplace(
                cls.add_scalar_inplace,
                (p, x),
                inplace_indices=[1],
                out_ty=R.Tensor((4,), "int64"),
            )
            return out

        @R.function
        def main(p: R.Prim("int64"), x: R.Tensor((4,), "int64")) -> R.Tensor((4,), "int64"):
            cls = Before
            out = cls.fused(p, x)
            return out

    after = relax.transform.FuseTIR()(Before)
    assert relax.analysis.check_well_formed(after)
    assert tvm.tirx.analysis.verify_well_formed(after["fused"])


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