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

import pytest

import tvm
import tvm.script
import tvm.testing
from tvm import IRModule, relax, tirx, topi
from tvm.ir import DummyGlobalInfo, VDevice
from tvm.script.parser import ir as I
from tvm.script.parser import relax as R
from tvm.script.parser import tirx as T


def _check(
    parsed: relax.Function | IRModule,
    expect: relax.Function | IRModule | None = None,
):
    test = parsed.script(show_meta=True)
    roundtrip_mod = tvm.script.from_source(test)
    tvm.ir.assert_structural_equal(parsed, roundtrip_mod)
    if isinstance(parsed, IRModule) and isinstance(roundtrip_mod, IRModule):
        relax.analysis.well_formed(parsed)
        relax.analysis.well_formed(roundtrip_mod)
    if expect:
        tvm.ir.assert_structural_equal(parsed, expect)


def test_simple_func():
    @R.function
    def foo(x: R.Tensor((128, 128), "float32")) -> R.Tensor((128, 128), "float32"):
        R.func_attr({"Primitive": True})
        gv0 = R.call_dps_packed("extern_func", x, R.Tensor((128, 128), dtype="float32"))
        gv1 = R.call_dps_packed("extern_dps_func", gv0, R.Tensor((128, 128), dtype="float32"))
        return gv1

    x = relax.Var("x", R.Tensor((128, 128), "float32"))
    bb = relax.BlockBuilder()
    with bb.function("foo", (x,), attrs={"Primitive": True}):
        y = bb.emit(relax.call_dps_packed("extern_func", x, R.Tensor((128, 128), dtype="float32")))
        out = bb.emit(
            relax.call_dps_packed("extern_dps_func", y, R.Tensor((128, 128), dtype="float32"))
        )
        bb.emit_func_output(out)

    _check(foo, bb.get()["foo"])


def test_error_report():
    with pytest.raises(tvm.error.DiagnosticError):

        @R.function
        def foo(x: R.Tensor((128, 128), "float32")) -> R.Tensor(None, "float32", ndim=2):
            # error: a = b = c is not allowed.
            gv0 = gv1 = R.call_tir("extern_func", x, R.Tensor((128, 128), dtype="float32"))
            return gv0


def test_mismatch_cast_dims_and_ndim():
    with pytest.raises(Exception):

        @R.function
        def f(
            x: R.Tensor((2, 3), "float32", ndim=3),
        ):  # error: ndim and the shape dims are mismatch
            return x


def test_unexpected_num_kw_args():
    with pytest.raises(Exception):

        @R.function
        def f(x: R.Tensor(dtype="float32", ndim=1, foo=2)):  # error: unexpected kw args foo
            return x


def test_unexpected_ndim():
    with pytest.raises(Exception):

        @R.function
        # error: dim is expected to be non-negative int or -1 for unknown
        def f(x: R.Tensor(dtype="float32", ndim=-2)):
            return x


def test_unexpected_ndim_type():
    with pytest.raises(Exception):

        @R.function
        def f(x: R.Tensor(dtype="float32", ndim="1")):  # error: dim is expected to be int
            return x


def test_unexpected_tir_cast_args():
    with pytest.raises(tvm.error.DiagnosticError):

        @R.function
        def f(x: R.Tensor(("m",), "float32")):
            m = T.int64()
            # tirx.cast expects 2 arguments, but got 3
            return R.call_tir("foo", (x,), R.Tensor((T.cast("int32", m, 1),), dtype="float32"))


def test_unexpected_tir_args():
    with pytest.raises(tvm.error.DiagnosticError):

        @tvm.script.ir_module
        class TestWellCallTIR:
            @T.prim_func(s_tir=True)
            def tir_addone(A: T.Buffer((16, 16), "int32"), B: T.Buffer((16, 16), "int32")) -> None:
                T.func_attr({"global_symbol": "tir_addone"})
                for i, j in T.grid(16, 16):
                    with T.sblock("tir_addone"):
                        vi, vj = T.axis.remap("SS", [i, j])
                        B[vi, vj] = A[vi, vj] + T.int32(1)

            @R.function
            def foo(x: R.Tensor(("m", "m"), "float32")):
                m = T.int64()
                # tirx.max expects 2 arguments, but got 1
                gv = R.call_tir(tir_addone, (x,), R.Tensor((T.max(16),), dtype="float32"))
                return gv

    with pytest.raises(tvm.error.DiagnosticError):

        @R.function
        def f(x: R.Tensor(("m", "n"), "float32")):
            m = T.int64()
            # call_tir expected a tirx prim_func
            return relax.call_tir("extern_func", (x,), R.Tensor((T.max(m),), dtype="float32"))


def test_func_type_annotation_fail():
    with pytest.raises(tvm.error.DiagnosticError):

        @R.function
        def f(x, y):  # error: the parameter type annotation is missing
            z = R.add(x, y)
            y = z
            return y


def test_if_mismatch_var_fail():
    with pytest.raises(tvm.error.DiagnosticError):

        @R.function
        def f(cond: R.Tensor((), "bool"), x: R.Tensor((1,), "float32")):
            if cond:
                w = R.add(x, x)
                y = R.multiply(w, w)
            else:
                w = R.multiply(x, x)
                z = R.add(w, w)  # error: The binding var is expected to `y`
            return z


def test_unassigned_call_fail():
    with pytest.raises(tvm.error.DiagnosticError):

        @R.function
        def f(x: R.Tensor):
            R.add(x, x)
            return x


def test_incorrect_tensor_shape():
    with pytest.raises(tvm.error.DiagnosticError):

        @R.function
        def f(x: R.Tensor([16])):
            y: R.Tensor(16) = R.add(x, x)
            return y


def test_simple_module():
    @I.ir_module(s_tir=True)
    class TestModule:
        @T.prim_func(private=True, s_tir=True)
        def tir_func(
            x: T.Buffer((T.int64(128), T.int64(128)), "float32"),
            y: T.Buffer((T.int64(128), T.int64(128)), "float32"),
        ):
            T.func_attr({"tirx.noalias": True})
            for i, j in T.grid(T.int64(128), T.int64(128)):
                with T.sblock():
                    vi, vj = T.axis.remap("SS", [i, j])
                    y[vi, vj] = x[vi, vj] + 1.0

        @R.function
        def foo(x: R.Tensor((128, 128), "float32")) -> R.Tensor((128, 128), "float32"):
            cls = TestModule
            gv0 = R.call_tir(cls.tir_func, x, R.Tensor((128, 128), dtype="float32"))
            return gv0

    x = relax.Var("x", R.Tensor((128, 128), "float32"))
    bb = relax.BlockBuilder()
    with bb.function("foo", (x,), {"global_symbol": "foo"}):
        out = bb.emit_te(lambda x: x + 1, x, primfunc_name_hint="tir_func")
        bb.emit_func_output(out)

    _check(TestModule, bb.get())


def test_emit_te_primfunc_attrs():
    @I.ir_module(s_tir=True)
    class TestModule:
        @T.prim_func(private=True, s_tir=True)
        def plus_one(
            x: T.Buffer((T.int64(128), T.int64(128)), "float32"),
            y: T.Buffer((T.int64(128), T.int64(128)), "float32"),
        ):
            T.func_attr({"some_attr": "foo", "another_attr": True, "tirx.noalias": True})
            for i, j in T.grid(T.int64(128), T.int64(128)):
                with T.sblock():
                    vi, vj = T.axis.remap("SS", [i, j])
                    y[vi, vj] = x[vi, vj] + 1.0

        @R.function
        def foo(x: R.Tensor((128, 128), "float32")) -> R.Tensor((128, 128), "float32"):
            cls = TestModule
            gv0 = R.call_tir(cls.plus_one, x, R.Tensor((128, 128), dtype="float32"))
            return gv0

    x = relax.Var("x", R.Tensor((128, 128), "float32"))
    bb = relax.BlockBuilder()
    with bb.function("foo", (x,), {"global_symbol": "foo"}):
        out = bb.emit_te(
            lambda x: x + 1,
            x,
            primfunc_name_hint="plus_one",
            primfunc_attrs={"some_attr": "foo", "another_attr": True},
        )
        bb.emit_func_output(out)
    _check(TestModule, bb.get())


def test_emit_te():
    @I.ir_module(s_tir=True)
    class EmitTE:
        @R.function
        def main(x: R.Tensor((10, 20), "float32")) -> R.Tensor((10, 20), dtype="float32"):
            lv1 = R.emit_te(topi.add, x, x)
            out = R.emit_te(topi.multiply, lv1, lv1)
            return out

    bb = relax.BlockBuilder()
    x = relax.Var("x", relax.TensorType([10, 20], "float32"))
    with bb.function("main", [x], {"global_symbol": "main"}):
        lv1 = bb.emit_te(topi.add, x, x)
        out = bb.emit_te(topi.multiply, lv1, lv1)
        bb.emit_func_output(out)

    _check(EmitTE, bb.get())


def test_module_with_attr_and_global_info():
    @I.ir_module(s_tir=True)
    class TestModule:
        I.module_attrs({"attr": 10})
        I.module_global_infos(
            {
                "dummy": [
                    I.dummy_global_info(),  # dummy[0]
                    I.dummy_global_info(),  # dummy[1]
                ]
            }
        )

        @T.prim_func(private=True, s_tir=True)
        def tir_func(
            x: T.Buffer((T.int64(128), T.int64(128)), "float32"),
            y: T.Buffer((T.int64(128), T.int64(128)), "float32"),
        ):
            T.func_attr({"tirx.noalias": True})
            for i, j in T.grid(T.int64(128), T.int64(128)):
                with T.sblock():
                    vi, vj = T.axis.remap("SS", [i, j])
                    y[vi, vj] = x[vi, vj] + 1.0

        @R.function
        def foo(x: R.Tensor((128, 128), "float32")) -> R.Tensor((128, 128), "float32"):
            cls = TestModule
            gv0 = R.call_tir(cls.tir_func, x, R.Tensor((128, 128), dtype="float32"))
            return gv0

    x = relax.Var("x", R.Tensor((128, 128), "float32"))
    bb = relax.BlockBuilder()
    with bb.function("foo", (x,), {"global_symbol": "foo"}):
        out = bb.emit_te(lambda x: x + 1, x, primfunc_name_hint="tir_func")
        bb.emit_func_output(out)
    mod = bb.get()
    mod.update_global_info("dummy", [DummyGlobalInfo(), DummyGlobalInfo()])
    mod = mod.with_attr("attr", 10)
    _check(TestModule, mod)


def test_global_info_vdevice():
    vdevices = [
        VDevice("llvm"),
        VDevice("cuda", 0),
        VDevice({"kind": "cuda", "arch": "sm_80"}, 0),
        VDevice("metal", 0, "global"),
    ]

    @I.ir_module(s_tir=True)
    class TestModule:
        I.module_attrs({"attr": 10})
        I.module_global_infos(
            {
                "vdevice": [
                    I.vdevice("llvm"),
                    I.vdevice("cuda", 0),
                    I.vdevice({"kind": "cuda", "arch": "sm_80"}, 0),
                    I.vdevice("metal", 0, "global"),
                ]
            }
        )

        @T.prim_func(private=True, s_tir=True)
        def tir_func(
            x: T.Buffer((T.int64(128), T.int64(128)), "float32"),
            y: T.Buffer((T.int64(128), T.int64(128)), "float32"),
        ):
            T.func_attr({"tirx.noalias": True})
            for i, j in T.grid(T.int64(128), T.int64(128)):
                with T.sblock():
                    vi, vj = T.axis.remap("SS", [i, j])
                    y[vi, vj] = x[vi, vj] + 1.0

        @R.function
        def foo(x: R.Tensor((128, 128), "float32")) -> R.Tensor((128, 128), "float32"):
            cls = TestModule
            gv0 = R.call_tir(cls.tir_func, x, R.Tensor((128, 128), dtype="float32"))
            return gv0

    x = relax.Var("x", R.Tensor((128, 128), "float32"))
    bb = relax.BlockBuilder()
    with bb.function("foo", (x,)):
        out = bb.emit_te(lambda x: x + 1, x, primfunc_name_hint="tir_func")
        bb.emit_func_output(out)
    mod = bb.get()
    mod.update_global_info("vdevice", vdevices)
    mod = mod.with_attr("attr", 10)
    _check(TestModule, mod)


def test_relax_tensor_op():
    @R.function
    def foo(x: R.Tensor((4, 4), "float32")) -> R.Tensor((4, 4), "float32"):
        y = R.add(x, x)
        z = R.multiply(x, y)
        return z

    x = relax.Var("x", R.Tensor((4, 4), "float32"))
    bb = relax.BlockBuilder()
    with bb.function("foo", (x,)):
        y = bb.emit(relax.op.add(x, x))
        z = bb.emit(relax.op.multiply(x, y))
        bb.emit_func_output(z)

    _check(foo, bb.get()["foo"])


def test_relax_base_op():
    @R.function
    def foo(x: R.Tensor((4, 4), "float32")):
        alloc = R.builtin.alloc_tensor(R.shape([4, 4]), runtime_device_index=0, dtype="float32")
        shape = R.shape_of(alloc)
        return shape

    x = relax.Var("x", R.Tensor((4, 4), "float32"))
    bb = relax.BlockBuilder()
    with bb.function("foo", (x,)):
        alloc = bb.emit(relax.op.builtin.alloc_tensor(relax.ShapeExpr((4, 4)), "float32", 0))
        shape = bb.emit(relax.op.shape_of(alloc))
        bb.emit_func_output(shape)

    _check(foo, bb.get()["foo"])


def test_relax_shape_to_tensor():
    @R.function
    def foo(x: R.Shape((4, 4))):
        tensor = R.shape_to_tensor(x)
        return tensor

    x = relax.Var("x", R.Shape((4, 4)))
    bb = relax.BlockBuilder()
    with bb.function("foo", (x,)):
        tensor = bb.emit(relax.op.shape_to_tensor(x))
        bb.emit_func_output(tensor)

    _check(foo, bb.get()["foo"])


def test_symbolic_shape():
    @R.function
    def foo(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(("m", "n"), "float32"):
        m = T.int64()
        n = T.int64()
        gv0 = R.call_dps_packed("extern_func", x, R.Tensor((m, n), dtype="float32"))
        return gv0

    @R.function
    def bar(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(("m", "n"), "float32"):
        m = T.int64()
        n = T.int64()
        gv0 = R.call_dps_packed("extern_func", x, R.Tensor((m, n), dtype="float32"))
        return gv0

    with pytest.raises(tvm.error.DiagnosticError):

        @R.function
        def mismatch_dtype(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(None, "float32", ndim=2):
            m = T.int64()
            n = T.int32()  # The shape dtype should be int64
            gv0 = R.call_dps_packed("extern_func", x, R.Tensor((m, n), dtype="float32"))
            return gv0

    def _expected(name: str):
        n, m = tirx.Var("n", "int64"), tirx.Var("m", "int64")
        x = relax.Var("x", R.Tensor([m, n], "float32"))
        bb = relax.BlockBuilder()
        with bb.function(name, (x,)):
            out = bb.emit(
                relax.call_dps_packed("extern_func", x, R.Tensor((m, n), dtype="float32"))
            )
            bb.emit_func_output(out)
        return bb.get()[name]

    _check(foo, _expected("foo"))
    _check(bar, _expected("bar"))


def test_shadowing():
    @R.function
    def foo(x: R.Tensor((4, 4), "float32")):
        y = R.add(x, x)
        z = R.multiply(x, y)
        y = R.add(x, y)
        y = z
        y = R.multiply(y, x)
        z = y
        return z

    x = relax.Var("x", R.Tensor((4, 4), "float32"))
    bb = relax.BlockBuilder()
    with bb.function("foo", (x,)):
        y = bb.emit(relax.op.add(x, x))
        z = bb.emit(relax.op.multiply(x, y))
        y = bb.emit(relax.op.add(x, y))
        y = bb.emit(z)
        y = bb.emit(relax.op.multiply(y, x))
        z = bb.emit(y)
        bb.emit_func_output(z)

    _check(foo, bb.get()["foo"])


def test_match_cast():
    @R.function
    def foo(x: R.Tensor("float32"), y: R.Tensor("float32")):
        m = T.int64()
        n = T.int64()
        x0 = R.match_cast(x, R.Tensor([m], "float32"))
        with R.dataflow():
            y0 = R.match_cast(y, R.Tensor([n], "float32"))
            gv = y0
            R.output(gv)
        return (x0, R.shape([m, n * 2]))

    x = relax.Var("x", R.Tensor("float32"))
    y = relax.Var("y", R.Tensor("float32"))
    m = tirx.Var("m", ty="int64")
    n = tirx.Var("n", ty="int64")
    y2 = relax.Var("y", R.Tensor([n], "float32"))
    bb = relax.BlockBuilder()
    with bb.function("foo", (x, y)):
        x0 = bb.match_cast(x, R.Tensor([m], "float32"))
        with bb.dataflow():
            y0 = bb.match_cast(y, R.Tensor([n], "float32"))
            bb.emit_output(y0)
        bb.emit_func_output(relax.Tuple([x0, relax.ShapeExpr([m, n * 2])]))

    _check(foo, bb.get()["foo"])


def test_tuple_return():
    @R.function
    def foo(x: R.Tensor((4, 4), "float32")):
        gv0 = R.call_dps_packed("extern_func_0", x, R.Tensor((4, 4), dtype="float32"))
        gv1 = R.call_dps_packed("extern_func_1", x, R.Tensor((4, 4), dtype="float32"))
        return (gv0, gv1)

    x = relax.Var("x", R.Tensor((4, 4), "float32"))
    bb = relax.BlockBuilder()
    with bb.function("foo", (x,)):
        gv0 = bb.emit(relax.call_dps_packed("extern_func_0", x, R.Tensor((4, 4), dtype="float32")))
        gv1 = bb.emit(relax.call_dps_packed("extern_func_1", x, R.Tensor((4, 4), dtype="float32")))
        bb.emit_func_output(relax.Tuple((gv0, gv1)))

    _check(foo, bb.get()["foo"])


def test_tuple_return_2():
    @R.function
    def foo(x: R.Tensor("float32", ndim=2)):
        n, m = T.int64(), T.int64()
        x0 = R.match_cast(x, R.Tensor((n, m), "float32"))
        return (x0, R.shape([n + 1, m, 1]))

    x = relax.Var("x", R.Tensor("float32", ndim=2))
    n, m = tirx.Var("n", "int64"), tirx.Var("m", "int64")
    bb = relax.BlockBuilder()
    with bb.function("foo", (x,)):
        x0 = bb.match_cast(x, R.Tensor((n, m), "float32"))
        bb.emit_func_output(relax.Tuple([x0, relax.ShapeExpr([n + 1, m, 1])]))

    _check(foo, bb.get()["foo"])


def test_tuple_binding():
    @R.function
    def foo(x: R.Tensor("float32", ndim=2)):
        n, m = T.int64(), T.int64()
        x0 = R.match_cast(x, R.Tensor((n, m), "float32"))
        t0 = (x, x0)
        t1 = (x, R.shape([n, m]), t0)
        return t1

    x = relax.Var("x", R.Tensor("float32", ndim=2))
    n, m = tirx.Var("n", "int64"), tirx.Var("m", "int64")
    bb = relax.BlockBuilder()
    with bb.function("foo", (x,)):
        x0 = bb.match_cast(x, R.Tensor((n, m), "float32"))
        t0 = bb.emit(relax.Tuple([x, x0]))
        t1 = bb.emit(relax.Tuple([x, relax.ShapeExpr([n, m]), t0]))
        bb.emit_func_output(t1)

    _check(foo, bb.get()["foo"])


def test_tuple_get_item():
    @R.function
    def foo(x: R.Tensor, y: R.Tensor):
        t1 = R.tuple(x, y)
        t2 = (x, y)
        a = t1[0]
        b = R.TupleGetItem(t2, 1)
        c = R.add(a, b)
        return c

    x = relax.Var("x", R.Tensor())
    y = relax.Var("y", R.Tensor())
    bb = relax.BlockBuilder()
    with bb.function("foo", (x, y)):
        t1 = bb.emit(relax.Tuple([x, y]))
        t2 = bb.emit(relax.Tuple([x, y]))
        a = bb.emit(relax.TupleGetItem(t1, 0))
        b = bb.emit(relax.TupleGetItem(t2, 1))
        c = bb.emit(relax.op.add(a, b))
        bb.emit_func_output(c)

    _check(foo, bb.get()["foo"])


def test_dataflow_block():
    @R.function
    def foo(x: R.Tensor((128, 128), "float32")) -> R.Tensor(None, "float32", ndim=2):
        with R.dataflow():
            lv0 = R.call_dps_packed("extern_func", x, R.Tensor((128, 128), dtype="float32"))
            lv1 = R.call_dps_packed("extern_func", lv0, R.Tensor((128, 128), dtype="float32"))
            gv = lv1
            R.output(gv)
        return gv

    x = relax.Var("x", R.Tensor((128, 128), "float32"))
    bb = relax.BlockBuilder()
    with bb.function("foo", (x,)):
        with bb.dataflow():
            lv0 = bb.emit(
                relax.call_dps_packed("extern_func", x, R.Tensor((128, 128), dtype="float32"))
            )
            lv1 = bb.emit(
                relax.call_dps_packed("extern_func", lv0, R.Tensor((128, 128), dtype="float32"))
            )
            gv = bb.emit_output(lv1)
        bb.emit_func_output(gv)

    _check(foo, bb.get()["foo"])


def test_dataflow_block_advanced():
    @R.function
    def foo(x: R.Tensor((128, 128), "float32")) -> R.Tensor(None, "float32", ndim=2):
        gv0 = R.call_dps_packed("extern_func", x, R.Tensor((128, 128), dtype="float32"))
        gv1 = R.call_dps_packed("extern_func", gv0, R.Tensor((128, 128), dtype="float32"))
        with R.dataflow():
            m = T.int64()
            n = T.int64()
            lv0 = R.call_dps_packed("extern_func", gv1, R.Tensor((128, 128), dtype="float32"))
            lv1 = R.match_cast(lv0, R.Tensor((m, n), "float32"))
            gv2 = R.call_dps_packed("extern_func", lv0, R.Tensor((128, 128), dtype="float32"))
            gv2 = R.call_dps_packed("extern_func", gv2, R.Tensor((128, 128), dtype="float32"))
            gv3 = R.match_cast(gv2, R.Tensor((m, n), "float32"))
            gv3 = R.match_cast(lv0, R.Tensor((m, n), "float32"))
            gv4 = gv3
            gv5 = gv2
            R.output(gv5, gv4)
        gv6 = R.call_dps_packed("extern_func", gv5, R.Tensor((128, 128), dtype="float32"))
        gv7 = R.call_dps_packed("extern_func", gv6, R.Tensor((128, 128), dtype="float32"))
        return gv7

    x = relax.Var("x", R.Tensor((128, 128), "float32"))
    bb = relax.BlockBuilder()
    m = tirx.Var("m", ty="int64")
    n = tirx.Var("n", ty="int64")
    with bb.function("foo", (x,)):
        gv0 = bb.emit(
            relax.call_dps_packed("extern_func", x, R.Tensor((128, 128), dtype="float32"))
        )
        gv1 = bb.emit(
            relax.call_dps_packed("extern_func", gv0, R.Tensor((128, 128), dtype="float32"))
        )
        with bb.dataflow():
            lv0 = bb.emit(
                relax.call_dps_packed("extern_func", gv1, R.Tensor((128, 128), dtype="float32"))
            )
            lv1 = bb.match_cast(lv0, R.Tensor((m, n), "float32"))
            gv2 = bb.emit(
                relax.call_dps_packed("extern_func", lv0, R.Tensor((128, 128), dtype="float32"))
            )
            gv21 = bb.emit(
                relax.call_dps_packed("extern_func", gv2, R.Tensor((128, 128), dtype="float32"))
            )
            gv3 = bb.match_cast(gv21, R.Tensor((m, n), "float32"))
            gv31 = bb.match_cast(lv0, R.Tensor((m, n), "float32"))
            gv32 = bb.emit_output(gv31)
            gv22 = bb.emit_output(gv21)
        gv4 = bb.emit(
            relax.call_dps_packed("extern_func", gv22, R.Tensor((128, 128), dtype="float32"))
        )
        gv5 = bb.emit(
            relax.call_dps_packed("extern_func", gv4, R.Tensor((128, 128), dtype="float32"))
        )
        bb.emit_func_output(gv5)

    _check(foo, bb.get()["foo"])


def test_dataflow_binding_after_output():
    with pytest.raises(tvm.error.DiagnosticError):

        @R.function
        def foo(x: R.Tensor((128, 128), "float32")) -> R.Tensor(None, "float32", ndim=2):
            with R.dataflow():
                gv = R.call_tir("extern_func", x, R.Tensor((128, 128), dtype="float32"))
                R.output(gv)
                lv = R.call_tir("extern_func", gv, R.Tensor((128, 128), dtype="float32"))
            return gv


def test_dataflow_output_global_var():
    with pytest.raises(tvm.error.DiagnosticError):

        @R.function
        def foo(x: R.Tensor((128, 128), "float32")) -> R.Tensor(None, "float32", ndim=2):
            gv0 = R.call_tir("extern_func", x, R.Tensor((128, 128), dtype="float32"))
            with R.dataflow():
                gv1 = R.call_tir("extern_func", gv0, R.Tensor((128, 128), dtype="float32"))
                R.output(gv0, gv1)
            return gv1


def test_dataflow_multiple_output():
    with pytest.raises(tvm.error.DiagnosticError):

        @R.function
        def foo(x: R.Tensor((128, 128), "float32")) -> R.Tensor(None, "float32", ndim=2):
            with R.dataflow():
                gv = R.call_tir("extern_func", x, R.Tensor((128, 128), dtype="float32"))
                R.output(gv)
                R.output(gv)
            return gv


def test_dataflow_output_outside_dataflow_block():
    with pytest.raises(tvm.error.DiagnosticError):

        @R.function
        def foo(x: R.Tensor((128, 128), "float32")) -> R.Tensor(None, "float32", ndim=2):
            gv = R.call_tir("extern_func", x, R.Tensor((128, 128), dtype="float32"))
            R.output(gv)
            return gv


def test_dataflow_scope_fail():
    with pytest.raises(tvm.error.DiagnosticError):

        @R.function
        def f(x: R.Tensor(ndim=2)):
            with R.dataflow():
                y = R.add(x, x)
                z = R.multiply(y, x)
                w = R.add(z, x)
                R.output(y, w)
            t = R.multiply(y, z)  # z is not in the outer scope
            return t


def test_return_without_binding():
    @R.function
    def foo(x: R.Tensor((128, 128), "float32")):
        return x

    x = relax.Var("x", R.Tensor((128, 128), "float32"))
    bb = relax.BlockBuilder()
    with bb.function("foo", (x,)):
        bb.emit_func_output(x)

    _check(foo, bb.get()["foo"])


def test_multiple_return():
    with pytest.raises(tvm.error.DiagnosticError):

        @R.function
        def foo(x: R.Tensor((128, 128), "float32")):
            return x
            return x


def test_function_without_return():
    with pytest.raises(tvm.error.DiagnosticError):

        @R.function
        def foo(x: R.Tensor((128, 128), "float32")):
            gv0 = R.call_tir("extern_func", x, R.Tensor((128, 128), dtype="float32"))


def test_tensor_type_without_args():
    @R.function
    def foo(x: R.Tensor((32, 32), "float32")) -> R.Tensor:
        v = R.call_dps_packed("extern_relu", x, R.Tensor((32, 32), dtype="float32"))
        return v

    x = relax.Var("x", R.Tensor((32, 32), "float32"))
    bb = relax.BlockBuilder()
    with bb.function("foo", (x)):
        v = bb.emit(relax.call_dps_packed("extern_relu", x, R.Tensor((32, 32), dtype="float32")))
        bb.emit_func_output(v)

    _check(foo, bb.get()["foo"])


def test_tensor_with_vdevice():
    vdevices = [
        VDevice("llvm"),
        VDevice("cuda", 0),
        VDevice("metal", 0, "global"),
        VDevice({"kind": "cuda", "arch": "sm_80"}, 0),
    ]

    @I.ir_module(s_tir=True)
    class TestModule:
        I.module_attrs({"attr": 10})
        I.module_global_infos(
            {
                "vdevice": [
                    I.vdevice("llvm"),
                    I.vdevice("cuda", 0),
                    I.vdevice("metal", 0, "global"),
                    I.vdevice({"kind": "cuda", "arch": "sm_80"}, 0),
                ]
            }
        )

        @R.function
        def foo(
            a: R.Tensor((128, 128), "float32", "cuda:1"),
            b: R.Tensor((128, 128), "float32", "llvm"),
            c: R.Tensor((128, 128), "float32", "vdevice:3"),
        ) -> R.Tensor((128, 128), "float32", "cuda:1"):
            s = R.add(a, c)
            return s

    a = relax.Var("a", R.Tensor((128, 128), "float32", vdevices[3]))
    b = relax.Var("b", R.Tensor((128, 128), "float32", vdevices[0]))
    c = relax.Var("c", R.Tensor((128, 128), "float32", vdevices[3]))
    bb = relax.BlockBuilder()
    with bb.function("foo", (a, b, c)):
        out = bb.emit(relax.op.add(a, c))
        bb.emit_func_output(out)
    mod = bb.get()
    mod = mod.with_attr("attr", 10)
    mod.update_global_info("vdevice", vdevices)

    _check(TestModule, mod)


def test_direct_return():
    @R.function
    def foo(x: R.Tensor((32, 32), "float32")) -> R.Tensor((32, 32), "float32"):
        return x

    x = relax.Var("x", R.Tensor((32, 32), "float32"))
    bb = relax.BlockBuilder()
    with bb.function("foo", (x)):
        bb.emit_func_output(x)

    _check(foo, bb.get()["foo"])


def test_call_packed():
    @R.function(pure=False)
    def foo(x: R.Tensor((32, 32), "float32")) -> R.Tensor:
        z = R.call_packed("vm.builtin.copy", x, ty_args=R.Tensor((32, 32), "float32"))
        return z

    x = relax.Var("x", R.Tensor((32, 32), "float32"))
    bb = relax.BlockBuilder()
    with bb.function("foo", (x), pure=False):
        z = bb.emit(
            relax.Call(
                relax.ExternFunc("vm.builtin.copy"),
                (x,),
                None,
                ty_args=[R.Tensor((32, 32), "float32")],
            )
        )
        bb.emit_func_output(z)

    _check(foo, bb.get()["foo"])


def test_call_packed_without_ty_args():
    @R.function(pure=False)
    def foo(x: R.Any) -> R.Any:
        z = R.call_packed("test", x)
        return z

    x = relax.Var("x", R.Any())
    bb = relax.BlockBuilder()
    with bb.function("foo", (x), pure=False):
        z = bb.emit(
            relax.Call(
                relax.ExternFunc("test"),
                (x,),
                None,
                ty_args=[],
            )
        )
        bb.emit_func_output(z)

    _check(foo, bb.get()["foo"])


def test_object_proxy_compat_alias():
    @R.function
    def foo(x: R.Object) -> R.Object:
        return x

    assert isinstance(foo.ret_ty, relax.AnyType)


def test_annotation():
    @R.function(pure=False)
    def foo(
        x: R.Tensor((32, "m"), "float32"),
        y: R.Tensor(("m",), "float32"),
        r: R.Tensor(dtype="int64"),
    ) -> R.Any:
        m = T.int64()
        z: R.Tensor((32, m), "float32") = R.multiply(x, y)
        w: R.Tensor(ndim=2) = R.multiply(z, z)
        q: R.Tensor = R.add(w, w)
        t = R.add(w, z)
        sh: R.Shape = R.call_packed("shape_of", x, ty_args=R.Shape)
        lv: R.Tensor(sh, dtype="float32") = R.reshape(x, sh)
        o: R.Any = R.call_packed("contrib.tensor_array_stack", x, y, ty_args=R.Any)
        return o

    def _check_ty(binding, expected_ty):
        tvm.ir.assert_structural_equal(binding.var.ty, expected_ty)
        tvm.ir.assert_structural_equal(binding.value.ty, expected_ty)

    # Cannot use block builder here because we need to check the annotated type,
    # which may be inconsistent with deduced type.
    assert isinstance(foo.ret_ty, relax.AnyType)
    m = relax.get_shape_of(foo.params[0])[1]
    bindings = foo.body.blocks[0].bindings
    sh = bindings[4].var

    _check_ty(bindings[0], relax.TensorType([32, m], "float32"))
    _check_ty(bindings[1], relax.TensorType(dtype=None, ndim=2))
    _check_ty(bindings[2], relax.TensorType(dtype=None, ndim=-1))
    _check_ty(bindings[3], relax.TensorType(dtype=None, ndim=2))
    _check_ty(bindings[4], relax.ShapeType(ndim=-1))
    _check_ty(bindings[5], relax.TensorType(sh))
    _check_ty(bindings[6], relax.AnyType())


def test_annotate_override():
    @R.function
    def foo(x: R.Tensor):
        y = x
        # z will be treated as Any even though it's a tensor
        z: R.Any = R.add(x, y)
        return z

    assert isinstance(foo.ret_ty, relax.AnyType)
    y_bind, z_bind = foo.body.blocks[0].bindings
    assert isinstance(y_bind.var.ty, relax.TensorType)
    assert isinstance(z_bind.var.ty, relax.AnyType)

    with pytest.raises(tvm.error.DiagnosticError):

        @R.function
        def test(x: R.Tensor):
            # Error: x is of Tensor Type, which can not annotate to R.Shape.
            z: R.Shape = x
            return z

    @R.function
    def bar(x: R.Tensor):
        # x is of Tensor Type, the annotation of `z` is ignored.
        z: R.Any = x
        return z

    assert isinstance(bar.ret_ty, relax.TensorType)
    (z_bind,) = bar.body.blocks[0].bindings
    assert isinstance(z_bind.var.ty, relax.TensorType)


def test_call_dps_packed_empty_shape():
    @R.function
    def foo(x: R.Tensor((), "float32")):
        z = R.call_dps_packed("scalar_add", x, R.Tensor((), dtype="float32"))
        return z

    (z_bind,) = foo.body.blocks[0].bindings
    shape_expr = z_bind.value.ty_args[0].shape

    assert isinstance(shape_expr, relax.ShapeExpr)
    assert len(shape_expr.values) == 0


def test_call_tir_empty_tuple_arg():
    bb = relax.BlockBuilder()
    dummy_param = relax.Var("dummy_param", R.Tensor(()))
    with bb.function("foo", [dummy_param], {"global_symbol": "foo"}):
        output = bb.emit_te(topi.full, shape=(16, 32), dtype="float32", fill_value=1.0)
        bb.emit_func_output(output)

    _check(bb.get())


def test_call_tir_with_tir_var():
    @I.ir_module(s_tir=True)
    class Module:
        @R.function
        def main(
            dumb_param: R.Tensor(("n",), "float32"), x: R.Tensor(("n * 2",), "float32")
        ) -> R.Tensor(("n * 2",), "float32"):
            n = T.int64()
            cls = Module
            y = R.call_tir(cls.copy, (x, n), R.Tensor((n * 2,), dtype="float32"))
            return y

        @T.prim_func(s_tir=True)
        def copy(var_x: T.handle, n: T.int64, var_y: T.handle):
            X = T.match_buffer(var_x, (n * 2,), dtype="float32")
            Y = T.match_buffer(var_y, (n * 2,), dtype="float32")
            for i in T.grid(n * 2):
                with T.sblock("block"):
                    vi = T.axis.remap("S", [i])
                    Y[vi] = X[vi]

    _check(Module)


def test_call_tir_with_grad():
    @I.ir_module(s_tir=True)
    class Module:
        @T.prim_func(s_tir=True)
        def identity_tir(a: T.handle, b: T.handle) -> None:
            A = T.match_buffer(a, [54, 96])
            B = T.match_buffer(b, [54, 96])

            for i, j in T.grid(54, 96):
                with T.sblock("compute"):
                    vi, vj = T.axis.remap("SS", [i, j])
                    B[vi, vj] = A[vi, vj]

        @R.function
        def main(v0: R.Tensor([54, 96], "float32")):
            cls = Module
            out = R.call_tir_with_grad(
                cls.identity_tir,
                (v0,),
                R.Tensor((54, 96), "float32"),
                te_grad_name="identity_k_grad",
                te_grad_kwargs={"k": 1.0},
            )
            return out

    _check(Module)


def test_call_tir_inplace():
    @tvm.script.ir_module
    class Module:
        @T.prim_func(s_tir=True)
        def copy(
            A: T.Buffer((2, 3), "int32"),
            B: T.Buffer((2, 3), "int32"),
            out1: T.Buffer((2, 3), "int32"),
        ):
            # copies the contents of B into A and out1
            T.func_attr({"tirx.noalias": True})
            for i0, i1 in T.grid(T.int64(2), T.int64(3)):
                with T.sblock("T_zeros"):
                    ax0, ax1 = T.axis.remap("SS", [i0, i1])
                    T.reads(B[ax0, ax1])
                    T.writes(A[ax0, ax1], out1[ax0, ax1])
                    A[ax0, ax1] = B[ax0, ax1]
                    out1[ax0, ax1] = B[ax0, ax1]

        @R.function
        def main(x: R.Tensor((2, 3), "int32"), y: R.Tensor((2, 3), "int32")) -> R.Tuple(
            R.Tensor((2, 3), "int32"), R.Tensor((2, 3), "int32"), R.Tensor((2, 3), "int32")
        ):
            res = R.call_tir_inplace(
                Module.copy,
                (x, y),
                [0, -1],
                [R.Tensor((2, 3), "int32"), R.Tensor((2, 3), "int32")],
            )
            return res

    _check(Module)


def test_call_tir_inplace_with_tuple_var_raises_error():
    with pytest.raises(tvm.error.DiagnosticError):

        @tvm.script.ir_module
        class Module:
            @R.function
            def main(x: R.Tensor((2, 3), "int32"), y: R.Tensor((2, 3), "int32")):
                cls = Module
                args = (x, y)
                res = R.call_tir_inplace(
                    cls.copy,
                    # The `args` tuple must be an in-line tuple, not a
                    # reference to a tuple.  This error should be
                    # caught and raised during parsing.
                    args,
                    inplace_indices=[0, -1],
                    out_ty=[R.Tensor((2, 3), "int32"), R.Tensor((2, 3), "int32")],
                )
                return res

            @T.prim_func(s_tir=True)
            def copy(
                A: T.Buffer((2, 3), "int32"),
                B: T.Buffer((2, 3), "int32"),
                out1: T.Buffer((2, 3), "int32"),
            ):
                # copies the contents of B into A and out1
                T.func_attr({"tirx.noalias": True})
                for iters in T.grid(T.int64(2), T.int64(3)):
                    with T.sblock("T_zeros"):
                        i, j = T.axis.remap("SS", iters)
                        A[i, j] = B[i, j]
                        out1[i, j] = B[i, j]


def test_local_function():
    @R.function
    def main(x: R.Tensor((2, 3), "float32"), y: R.Tensor((2, 3), "float32")) -> R.Tensor(
        (2, 3), "float32"
    ):
        @R.function
        def outer_func(
            c1: R.Tensor((2, 3), "float32"),
        ) -> R.Callable((R.Tensor(None, "float32", ndim=2),), R.Tensor(None, "float32", ndim=2)):
            @R.function
            def inner_func(x1: R.Tensor((2, 3), "float32")):
                s: R.Tensor((2, 3), "float32") = R.add(x1, c1)
                return s

            return inner_func

        in_call = outer_func(x)
        res = in_call(y)
        return res

    main_bindings = main.body.blocks[0].bindings
    assert len(main_bindings) == 3
    outer_func = main_bindings[0].value
    assert isinstance(outer_func, relax.Function)

    outer_func_bindings = outer_func.body.blocks[0].bindings
    assert len(outer_func_bindings) == 1
    inner_func = outer_func_bindings[0].value
    assert isinstance(inner_func, relax.Function)


def test_inline_prim_func():
    with pytest.raises(tvm.error.DiagnosticError):

        @I.ir_module(s_tir=True)
        class TestModule:
            @R.function
            def f(x: R.Tensor((128, 128), "float32"), y: R.Tensor((128, 128), "float32")):
                @T.prim_func(s_tir=True)
                def my_matmul(a: T.handle, b: T.handle, c: T.handle) -> None:
                    A = T.match_buffer(a, (128, 128))
                    B = T.match_buffer(b, (128, 128))
                    C = T.match_buffer(c, (128, 128))

                    for i, j, k in T.grid(128, 128, 128):
                        with T.sblock():
                            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
                            with T.init():
                                C[vi, vj] = 0.0
                            C[vi, vj] += A[vi, vk] * B[vj, vk]

                z = relax.call_tir(my_matmul, (x, y), R.Tensor((128, 128), dtype="float32"))
                return z


def test_cross_function_call():
    @I.ir_module(s_tir=True)
    class Mod0:
        @R.function
        def foo(x: R.Tensor((10, 5), "float32")):
            s = R.add(x, x)
            return s

        @R.function
        def main(x: R.Tensor((10, 5), "float32")):
            cls = Mod0
            inner = cls.foo
            gv1 = inner(x)
            gv2 = Mod0.foo(x)
            return (inner, gv1, gv2)

    @I.ir_module(s_tir=True)
    class Mod1:
        @R.function
        def main(x: R.Tensor((10, 5), "float32")):
            cls = Mod1
            inner = cls.foo
            gv1 = inner(x)
            gv2 = Mod1.foo(x)
            return (inner, gv1, gv2)

        @R.function
        def foo(x: R.Tensor((10, 5), "float32")) -> R.Tensor((10, 5), "float32"):
            s = R.add(x, x)
            return s


def test_if_branch():
    @R.function
    def foo(cond: R.Tensor((), "bool"), x: R.Tensor((1,), "float32")) -> R.Tensor((1,), "float32"):
        if cond:
            w = R.add(x, x)
            y = R.multiply(w, w)
        else:
            w = R.multiply(x, x)
            y = R.add(w, w)
        return y

    cond, x = foo.params
    y_bind = foo.body.blocks[0].bindings[0]
    y, ite = y_bind.var, y_bind.value

    assert isinstance(y, relax.Var)
    assert y.name == "y"

    assert isinstance(ite, relax.If)
    assert isinstance(ite.true_branch, relax.SeqExpr)
    assert isinstance(ite.false_branch, relax.SeqExpr)

    def check_call(call, op, args):
        assert isinstance(call, relax.Call)
        if isinstance(op, str):
            assert call.op.name == op
        else:
            assert call.op == op
        tvm.ir.assert_structural_equal(call.args, args)

    w_bind = ite.true_branch.blocks[0].bindings[0]
    # the seq exprts in the branches are normalized to bind any call
    # in the seq expr "body" to a var
    y_bind = ite.true_branch.blocks[-1].bindings[-1]
    assert w_bind.var.name == "w"
    check_call(w_bind.value, "relax.add", [x, x])
    check_call(y_bind.value, "relax.multiply", [w_bind.var, w_bind.var])

    w_bind = ite.false_branch.blocks[0].bindings[0]
    y_bind = ite.false_branch.blocks[-1].bindings[-1]
    assert w_bind.var.name == "w"
    check_call(w_bind.value, "relax.multiply", [x, x])
    check_call(y_bind.value, "relax.add", [w_bind.var, w_bind.var])


def test_if_branch_with_match_cast():
    """The last branch of a relax::If node may be a MatchCast

    This is a regression test.  In previous implementations, using
    R.match_cast as the last binding would cause a segfault while
    parsing.
    """

    @R.function
    def func(A: R.Tensor([16, 16]), is_bfloat16: R.Prim("bool")):
        if is_bfloat16:
            A = R.match_cast(A, R.Tensor([16, 16], "bfloat16"))
            B = A.astype("float16")
        else:
            B = R.match_cast(A, R.Tensor([16, 16], "float16"))
        return B

    A, is_bfloat16 = func.params
    (block,) = func.body.blocks
    (B_binding,) = block.bindings

    B_var = B_binding.var
    assert isinstance(B_var, relax.Var)
    assert B_var.name == "B"

    if_then_else = B_binding.value
    assert isinstance(if_then_else, relax.If)
    assert isinstance(if_then_else.true_branch, relax.SeqExpr)
    assert isinstance(if_then_else.false_branch, relax.SeqExpr)

    else_branch = if_then_else.false_branch
    (else_block,) = else_branch.blocks

    assert isinstance(else_block.bindings[-1], relax.MatchCast)

    # If the `R.match_cast` were removed, the function would infer the
    # return value as `R.Tensor([16,16])`, with an unknown dtype.
    # With the `R.match_cast` retained, the output dtype is known.
    tvm.ir.assert_structural_equal(func.ret_ty, R.Tensor([16, 16], "float16"))


def test_if_inside_dataflow():
    with pytest.raises(tvm.error.DiagnosticError):

        @R.function
        def foo(cond: R.Tensor((), "bool"), x: R.Tensor((1,), "float32")):
            with R.dataflow():
                if cond:
                    w = R.add(x, x)
                    y = R.multiply(w, w)
                else:
                    w = R.multiply(x, x)
                    y = R.add(w, w)
                R.output(y)
            return y


def test_var_if_scoping_fail():
    with pytest.raises(tvm.error.DiagnosticError):

        @R.function
        def f(cond: R.Tensor((), "bool"), x: R.Tensor((1,), "float32")):
            if cond:
                w = R.add(x, x)
                y = R.multiply(w, w)
            else:
                w = R.multiply(x, x)
                y = R.add(w, w)
            return w  # error: The w is not defined in the outer scope


def test_scalar_tensor_as_branch_condition():
    """Branch condition can be 0-d tensor"""

    @R.function
    def func(cond: R.Tensor([], "bool"), x: R.Tensor((1,), "float32")):
        if cond:
            out = R.add(x, x)
        else:
            out = R.multiply(x, x)
        return out

    if_else = func.body.blocks[0].bindings[0].value
    assert isinstance(if_else.cond, relax.Var)
    tvm.ir.assert_structural_equal(if_else.cond.ty, R.Tensor([], "bool"))


def test_prim_annotation_requires_dtype():
    with pytest.raises(TypeError, match="missing 1 required positional argument: 'dtype'"):
        R.Prim()

    with pytest.raises(TypeError, match="unexpected keyword argument 'value'"):
        R.Prim(value="n")


def test_prim_value_as_branch_condition():
    """In addition to scalar tensor, can use R.Prim condition"""

    @R.function
    def func(cond: R.Prim("bool"), x: R.Tensor((1,), "float32")):
        if cond:
            out = R.add(x, x)
        else:
            out = R.multiply(x, x)
        return out

    if_else = func.body.blocks[0].bindings[0].value
    assert isinstance(if_else.cond, relax.Var)
    tvm.ir.assert_structural_equal(if_else.cond.ty, R.Prim("bool"))


def test_computed_prim_value_as_branch_condition():
    """The R.Prim condition may be computed within the function"""

    @R.function
    def func(x: R.Tensor(["N"], "float32")):
        N = T.int64()
        if R.prim_value(N % 16 == 0):
            out = R.call_pure_packed("fast_vectorized_impl", x, ty_args=[x.ty])
        else:
            out = R.call_pure_packed("slow_non_vectorized_impl", x, ty_args=[x.ty])
        return out

    N = func.params[0].ty.shape[0]
    if_else = func.body.blocks[0].bindings[0].value
    assert tvm.ir.is_prim_expr(if_else.cond)
    tvm.ir.assert_structural_equal(N % 16 == 0, if_else.cond)
    tvm.ir.assert_structural_equal(if_else.cond.ty, R.Prim("bool"))


def test_tir_expr_as_branch_condition():
    """Syntactic sugar, use Expr directly"""

    @R.function(private=True)
    def sugared(x: R.Tensor(["N"], "float32")):
        N = T.int64()
        if N % 16 == 0:
            out = R.call_pure_packed("fast_vectorized_impl", x, ty_args=[x.ty])
        else:
            out = R.call_pure_packed("slow_non_vectorized_impl", x, ty_args=[x.ty])
        return out

    @R.function(private=True)
    def unsugared(x: R.Tensor(["N"], "float32")):
        N = T.int64()
        if R.prim_value(N % 16 == 0):
            out = R.call_pure_packed("fast_vectorized_impl", x, ty_args=[x.ty])
        else:
            out = R.call_pure_packed("slow_non_vectorized_impl", x, ty_args=[x.ty])
        return out

    tvm.ir.assert_structural_equal(unsugared, sugared)


def test_scalar_tensor_as_assert_condition():
    """Branch condition can be 0-d tensor"""

    @R.function(pure=False)
    def func(cond: R.Tensor([], "bool"), x: R.Tensor((1,), "float32")):
        _ = R.assert_op(cond)
        out = R.add(x, x)
        return out

    assert_op = func.body.blocks[0].bindings[0].value
    condition = assert_op.args[0]
    assert isinstance(condition, relax.Var)
    tvm.ir.assert_structural_equal(condition.ty, R.Tensor([], "bool"))


def test_prim_value_as_assert_condition():
    """In addition to scalar tensor, can use R.Prim condition"""

    @R.function(pure=False)
    def func(cond: R.Prim("bool"), x: R.Tensor((1,), "float32")):
        _ = R.assert_op(cond)
        out = R.add(x, x)
        return out

    assert_op = func.body.blocks[0].bindings[0].value
    condition = assert_op.args[0]
    assert isinstance(condition, relax.Var)
    tvm.ir.assert_structural_equal(condition.ty, R.Prim("bool"))


def test_computed_prim_value_as_assert_condition():
    """The R.Prim condition may be computed within the function"""

    @R.function(pure=False)
    def func(x: R.Tensor(["N"], "float32")):
        N = T.int64()
        _ = R.assert_op(R.prim_value(N % 16 == 0))
        out = R.call_packed("fast_vectorized_impl", x, ty_args=[x.ty])
        return out

    N = func.params[0].ty.shape[0]
    assert_op = func.body.blocks[0].bindings[0].value
    condition = assert_op.args[0]
    assert tvm.ir.is_prim_expr(condition)
    tvm.ir.assert_structural_equal(N % 16 == 0, condition)
    tvm.ir.assert_structural_equal(condition.ty, R.Prim("bool"))


def test_tir_expr_as_assert_condition():
    """Syntactic sugar, use Expr directly"""

    @R.function(pure=False, private=True)
    def sugared(x: R.Tensor(["N"], "float32")):
        N = T.int64()
        _ = R.assert_op(N % 16 == 0)
        out = R.call_packed("fast_vectorized_impl", x, ty_args=[x.ty])
        return out

    @R.function(pure=False, private=True)
    def unsugared(x: R.Tensor(["N"], "float32")):
        N = T.int64()
        _ = R.assert_op(R.prim_value(N % 16 == 0))
        out = R.call_packed("fast_vectorized_impl", x, ty_args=[x.ty])
        return out

    tvm.ir.assert_structural_equal(unsugared, sugared)


def test_erase_to_well_defined_removes_internal_vars():
    @R.function
    def foo(x: R.Tensor):
        q = x
        m, n = T.int64(), T.int64()
        z = R.match_cast(q, R.Tensor((m, n)))
        w = z
        return w

    tvm.ir.assert_structural_equal(foo.ret_ty, R.Tensor(ndim=2))
    assert foo.ret_ty.shape is None
    _check(foo)


def test_erase_to_well_defined_keeps_variables_exposed_by_tensor_shape():
    @R.function
    def foo(x: R.Tensor(["m", "n"])):
        q = x
        m, n = T.int64(), T.int64()
        z = R.match_cast(q, R.Tensor((m, n)))
        w = z
        return w

    assert foo.ret_ty.shape is not None
    _check(foo)


def test_erase_to_well_defined_keeps_variants_exposed_by_shape_expr():
    @R.function
    def foo(x: R.Tensor, _: R.Shape(["m", "n"])):
        q = x
        m, n = T.int64(), T.int64()
        z = R.match_cast(q, R.Tensor((m, n)))
        w = z
        return w

    assert foo.ret_ty.shape is not None
    _check(foo)


def test_erase_to_well_defined_infers_from_shape_expr():
    @I.ir_module(s_tir=True)
    class Module:
        # The subroutine's symbolic variables are only in-scope for the subroutine.
        @R.function
        def subroutine(x: R.Tensor, _: R.Shape(["m", "n"])) -> R.Tensor(["m", "n"]):
            q = x
            m, n = T.int64(), T.int64()
            z = R.match_cast(q, R.Tensor((m, n)))
            w = z
            return w

        # However, struct inference can make the symbolic variables in
        # the main function to the symbolic variables in the
        # subroutine.  Therefore, the shape of the tensor returned
        # from main can have a well-defined shape.
        @R.function
        def main(x: R.Tensor, shape: R.Shape(["m", "n"])):
            output = Module.subroutine(x, shape)
            return output

    assert Module["main"].ret_ty.shape is not None
    _check(Module)


def test_empty_tuple():
    @R.function
    def foo(x: R.Tuple()):
        y: R.Tuple() = R.tuple()
        return y

    x = relax.Var("x", relax.TupleType([]))
    bb = relax.BlockBuilder()
    with bb.function("foo", (x,)):
        y = bb.emit(relax.Tuple([]))
        bb.emit_func_output(y)

    _check(foo, bb.get()["foo"])


def test_symbolic_vars_in_tensor_shape_with_usage_first():
    """First param may use symbolic variable defined in second param"""

    @R.function
    def foo(x: R.Tensor(("m + 1",), "float32"), y: R.Tensor(("m", 1), "float32")):
        z = R.add(x, y)
        return z

    m = tirx.Var("m", "int64")
    x = relax.Var("x", relax.TensorType([m + 1], "float32"))
    y = relax.Var("y", relax.TensorType([m, 1], "float32"))
    bb = relax.BlockBuilder()
    with bb.function("foo", (x, y)):
        z = bb.emit(relax.op.add(x, y))
        bb.emit_func_output(z)

    _check(foo, bb.get()["foo"])


def test_symbolic_vars_in_tensor_shape_with_definition_first():
    """Second param may use symbolic variable defined in first param"""

    @R.function
    def bar(x: R.Tensor(("m",), "float32"), y: R.Tensor(("T.max(m, 20)",), "float32")) -> R.Tensor(
        ("T.max(m, 20) + 1",), "float32"
    ):
        m = T.int64()
        z = R.call_dps_packed("test_intrin", (x, y), R.Tensor((T.max(m, 20) + 1,), dtype="float32"))
        return z

    m = tirx.Var("m", "int64")
    x = relax.Var("x", relax.TensorType([m], "float32"))
    y = relax.Var("y", relax.TensorType([tirx.max(m, 20)], "float32"))
    bb = relax.BlockBuilder()
    with bb.function("bar", (x, y)):
        z = bb.emit(
            relax.call_dps_packed(
                "test_intrin", (x, y), R.Tensor((tirx.max(m, 20) + 1,), dtype="float32")
            )
        )
        bb.emit_func_output(z)

    _check(bar, bb.get()["bar"])


def test_bound_prim_param_reused_in_dependent_annotations():
    func = tvm.script.from_source(
        """
@R.function
def main(
    n: R.Prim("int64"),
    direct: R.Tensor([n], "float32"),
    string_direct: R.Tensor(["n"], "float32"),
    shape: R.Shape(["n"]),
    compound: R.Tensor(["n + 1"], "float32"),
) -> R.Tensor(["n + 1"], "float32"):
    return compound
"""
    )

    n, direct, string_direct, shape, compound = func.params
    assert direct.ty.shape[0].same_as(n)
    assert string_direct.ty.shape[0].same_as(n)
    assert shape.ty.values[0].same_as(n)
    assert compound.ty.shape[0].a.same_as(n)
    assert func.ret_ty.shape[0].a.same_as(n)
    _check(func)


def test_bound_prim_param_reused_in_declared_function_signature():
    mod = tvm.script.from_source(
        """
@I.ir_module
class Module:
    @R.function
    def main(n: R.Prim("int64"), x: R.Tensor(["n + 1"], "float32")) -> R.Tensor(
        ["n + 1"], "float32"
    ):
        return x
"""
    )

    func = mod["main"]
    n, x = func.params
    assert x.ty.shape[0].a.same_as(n)
    assert func.ret_ty.shape[0].a.same_as(n)
    _check(mod)


def test_later_prim_param_not_adopted_by_usage_first_symbol():
    with pytest.raises(tvm.error.DiagnosticError):
        tvm.script.from_source(
            """
@R.function
def main(x: R.Tensor(["n"], "float32"), n: R.Prim("int64")):
    return x
"""
        )


def test_non_int64_prim_param_rejected_in_shape_annotation():
    with pytest.raises(tvm.error.DiagnosticError):
        tvm.script.from_source(
            """
@R.function
def main(n: R.Prim("int32"), x: R.Tensor(["n"], "float32")):
    return x
"""
        )


def test_recursive_local_function_reuses_earlier_prim_param_in_signature():
    func = tvm.script.from_source(
        """
@R.function
def main(n: R.Prim("int64"), x: R.Tensor([n], "float32")):
    @R.function
    def recurse(current: R.Prim("int64"), value: R.Tensor([current], "float32")) -> R.Tensor(
        [current], "float32"
    ):
        return recurse(current, value)

    return recurse(n, x)
"""
    )

    recursive_binding = next(
        binding
        for block in func.body.blocks
        for binding in block.bindings
        if isinstance(binding.value, relax.Function)
    )
    recursive_func = recursive_binding.value
    current, value = recursive_func.params
    assert value.ty.shape[0].same_as(current)
    assert recursive_func.ret_ty.shape[0].same_as(current)
    relax.analysis.well_formed(func)


def test_symbolic_vars_in_shape():
    """Symbolic variable may be defined in R.Shape"""

    @R.function
    def baz(x: R.Shape(("m",)), y: R.Tensor(("m * 2",), "float32")):
        m = T.int64()
        z = R.call_dps_packed("test_intrin", y, R.Tensor((m * 2,), dtype="float32"))
        return z

    m = tirx.Var("m", "int64")
    x = relax.Var("x", relax.ShapeType([m]))
    y = relax.Var("y", relax.TensorType([m * 2], "float32"))
    bb = relax.BlockBuilder()
    with bb.function("baz", (x, y)):
        z = bb.emit(relax.call_dps_packed("test_intrin", (y), R.Tensor((m * 2,), dtype="float32")))
        bb.emit_func_output(z)

    _check(baz, bb.get()["baz"])


def test_undefined_symbolic_var_raises_error():
    """An undefined symbolic variable in an error

    A symbolic variables is defined at the first site where it appears
    as a shape parameter without any modification.  TVMScript does not
    support solving for a symbolic variable in terms of the argument
    shape.  That is, this test case raises an error, and will not
    attempt to define `m` as either `x.shape[0]-1` or `x.shape[1]//2`.
    """
    with pytest.raises(tvm.error.DiagnosticError):

        @R.function
        def foo(x: R.Tensor(("m + 1", "m * 2"), "float32")):  # name 'm' is not defined
            z = R.add(x, x)
            return z


def test_arith_operators():
    @R.function
    def foo(x: R.Tensor(("m", "n"), "float32"), y: R.Tensor(("m", "n"), "float32")):
        a0 = -x
        a1 = x + y
        a2 = x - y
        a3 = x * y
        a4 = x / y
        a5 = x // y
        a6 = x**y

        c0 = x > y
        c1 = x < y
        c2 = x >= y
        c3 = x <= y

        tuple_expr = ((x, x), y)
        t0 = tuple_expr[0]
        t1 = tuple_expr[1]
        t2 = tuple_expr[0][0]  # <= Will normalize to two bindings
        return (a0, a1, a2, a3, a4, a5, a6, c0, c1, c2, c3, t0, t1, t2)

    m = tirx.Var("m", "int64")
    n = tirx.Var("n", "int64")
    x = relax.Var("x", relax.TensorType([m, n], "float32"))
    y = relax.Var("y", relax.TensorType([m, n], "float32"))
    bb = relax.BlockBuilder()
    with bb.function("foo", (x, y)):
        a0 = bb.emit(relax.op.negative(x))
        a1 = bb.emit(relax.op.add(x, y))
        a2 = bb.emit(relax.op.subtract(x, y))
        a3 = bb.emit(relax.op.multiply(x, y))
        a4 = bb.emit(relax.op.divide(x, y))
        a5 = bb.emit(relax.op.floor_divide(x, y))
        a6 = bb.emit(relax.op.power(x, y))

        c0 = bb.emit(relax.op.greater(x, y))
        c1 = bb.emit(relax.op.less(x, y))
        c2 = bb.emit(relax.op.greater_equal(x, y))
        c3 = bb.emit(relax.op.less_equal(x, y))

        tuple_expr = bb.emit(relax.Tuple((relax.Tuple((x, x)), y)))
        t0 = bb.emit(relax.TupleGetItem(tuple_expr, 0))
        t1 = bb.emit(relax.TupleGetItem(tuple_expr, 1))
        tmp = bb.emit(relax.TupleGetItem(tuple_expr, 0))
        t2 = bb.emit(relax.TupleGetItem(tmp, 0))
        bb.emit_func_output(relax.Tuple((a0, a1, a2, a3, a4, a5, a6, c0, c1, c2, c3, t0, t1, t2)))

    _check(foo, bb.get()["foo"])


def test_memory_ops():
    @R.function
    def foo(x: R.Tensor(("m", "n"), dtype="float32")):
        m = T.int64()
        n = T.int64()
        storage = R.memory.alloc_storage(
            R.shape([4 * m * n]), virtual_device_index=0, storage_scope="global", dtype="float32"
        )
        alloc = R.memory.alloc_tensor(storage, offset=0, shape=R.shape([m, n]), dtype="float32")
        tensor = R.builtin.alloc_tensor(R.shape([m, n]), dtype="float32", runtime_device_index=0)
        gv = tensor
        return alloc, gv

    _check(foo)


def test_vm_ops():
    @R.function(pure=False)
    def foo(x: R.Tensor(("m", "n"), dtype="float32")):
        m = T.int64()
        n = T.int64()
        storage = R.vm.alloc_storage(R.shape([4 * m * n]), runtime_device_index=0, dtype="uint8")
        alloc = R.vm.alloc_tensor(storage, offset=0, shape=R.shape([m, n]), dtype="float32")
        tensor = R.builtin.alloc_tensor(R.shape([m, n]), dtype="float32", runtime_device_index=0)
        tir_dym = R.vm.call_tir_dyn("te_func", (x, tensor, R.ShapeExpr((m, n))))
        return alloc, tir_dym

    _check(foo)


def test_builtin_ops():
    @R.function
    def foo(x: R.Tensor(("m", "n"), dtype="float32")):
        tensor = R.builtin.stop_lift_params(x)
        gv = tensor
        return gv

    _check(foo)


def test_prim_value():
    @R.function(pure=False)
    def foo():
        gv = R.call_packed("test", 1, ty_args=R.Tensor((32, 32), "float32"))
        return gv

    _check(foo)


def test_string_imm():
    @R.function(pure=False)
    def foo():
        gv = R.call_packed("test", "hello", ty_args=R.Tensor((32, 32), "float32"))
        return gv

    _check(foo)


def test_datatype_imm():
    @R.function(pure=False)
    def foo():
        gv = R.call_packed("test", R.dtype("float32"), ty_args=R.Tensor((32, 32), "float32"))
        return gv

    _check(foo)


def test_function_void_return_type():
    @tvm.script.ir_module
    class Foo:
        @R.function
        def main(x: R.Tensor((3, 3), dtype="float32")):
            res = Foo.mul(x)
            return res

        @R.function
        def mul(x: R.Tensor((3, 3), dtype="float32")):
            res = R.multiply(x, x)
            return res

    _check(Foo)
    # Since the return type of function `mul` is not annotated,
    # the function `main` regards it as a generic return type.
    assert isinstance(Foo["main"].ret_ty, relax.AnyType)
    assert isinstance(Foo["mul"].ret_ty, relax.TensorType)

    @tvm.script.ir_module
    class Bar:
        @R.function
        def main(x1: R.Tensor((3, 3), dtype="float32")):
            res1 = Bar.mul(x1)
            return res1

        @R.function
        def mul(x: R.Tensor((3, 3), dtype="float32")) -> None:
            res = R.multiply(x, x)
            return res

    # Since the return type of function `mul` is not annotated,
    # the function `main` regards it as a generic return type.
    _check(Bar)
    tvm.ir.assert_structural_equal(Bar["main"].ret_ty, relax.TupleType([]))
    tvm.ir.assert_structural_equal(Bar["mul"].ret_ty, relax.TupleType([]))


def test_class_normalize():
    @tvm.script.ir_module
    class InputModule:
        @R.function
        def mul_add(x: R.Tensor) -> R.Tensor:
            return R.multiply(R.add(x, x), R.add(x, x))

    # The parser automatically normalizes the input AST to the following ANF form
    @tvm.script.ir_module
    class OutputModule:
        @R.function
        def mul_add(x: R.Tensor) -> R.Tensor:
            gv = R.add(x, x)
            gv1 = R.add(x, x)
            return R.multiply(gv, gv1)

    _check(InputModule, OutputModule)


def test_context_aware_parsing(monkeypatch):
    @tvm.script.ir_module
    class Module:
        @T.prim_func(s_tir=True)
        def add(
            X: T.Buffer([T.int64(2), T.int64(4)], "float32"),
            Y: T.Buffer((), "float32"),
            Z: T.Buffer([T.int64(2), T.int64(4)], "float32"),
        ):
            T.evaluate(0)

        @R.function
        def main(x: R.Tensor((2, 4), dtype="float32")) -> R.Tensor((10,), dtype="float32"):
            R.func_attr({"relax.force_pure": True})
            cls = Module
            alloc = R.builtin.alloc_tensor(R.shape([2, 4]), dtype="float32", runtime_device_index=0)
            _: R.Tuple() = cls.add(x, R.const(1, "float32"), alloc)
            return alloc

    _check(Module)

    # Break the env settings, but context-aware parsing can still handle it
    def _break_env(self, *args):
        raise RuntimeError("Fail to pass context-aware parsing")

    monkeypatch.setattr(tvm.ir.GlobalVar, "__call__", _break_env)

    _check(Module)


def test_unit_tuple_on_rhs_of_assign():
    @I.ir_module(s_tir=True)
    class Module:
        @R.function
        def main(input: R.Tensor((5, 5))) -> R.Tuple(R.Tensor((5, 5))):
            gv = (input,)
            return gv

    _check(Module)


def test_empty_tuple_on_rhs_of_assign():
    @I.ir_module(s_tir=True)
    class Module:
        @R.function
        def main(input: R.Tensor((5, 5))) -> R.Tuple():
            gv = ()
            return gv

    _check(Module)


def test_global_var_ty():
    @I.ir_module(s_tir=True)
    class Module:
        @R.function
        def foo(x: R.Tensor((128, 128), "float32")):
            gv0 = R.emit_te(topi.add, x, x)
            return gv0

    target_ty = R.Callable(
        (R.Tensor((128, 128), dtype="float32"),), R.Tensor((128, 128), dtype="float32")
    )
    gv = Module.get_global_var("foo")
    tvm.ir.assert_structural_equal(gv.ty, target_ty)
    tvm.ir.assert_structural_equal(Module["foo"].ty, target_ty)
    _check(Module)


def test_assert_op():
    @I.ir_module(s_tir=True)
    class AssertOp:
        @R.function(pure=False)
        def main(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
            y = R.assert_op(R.const(False, dtype="bool"), x, format="x: {}")
            return x

    _check(AssertOp)


def test_assert_outside_of_class():
    @R.function(pure=False)
    def func(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
        y = R.assert_op(R.const(False, dtype="bool"), x, format="x: {}")
        return x

    # this just makes sure that the machinery regarding the pure attribute parses
    # in the case where the function is outside of a class too
    _check(func)


def test_impure_inner_function():
    @R.function
    def f(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
        # we will not actually call it
        @R.function(pure=False)
        def g(y: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
            z = R.assert_op(R.const(False, dtype="bool"), y, format="y: {}")
            return y

        return x

    assert f.is_pure
    # definition of g
    assert not f.body.blocks[0].bindings[0].value.is_pure

    # make sure we are not incorrectly passing state for inner functions
    _check(f)


def test_impure_inner_function_in_class():
    @I.ir_module(s_tir=True)
    class ImpureInner:
        @R.function
        def main(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
            # we will not actually call it
            @R.function(pure=False)
            def g(y: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
                z = R.assert_op(R.const(False, dtype="bool"), y, format="y: {}")
                return y

            return x

    assert ImpureInner["main"].is_pure
    # definition of g
    assert not ImpureInner["main"].body.blocks[0].bindings[0].value.is_pure

    # make sure we are not incorrectly passing state for inner functions
    _check(ImpureInner)


def test_print():
    @I.ir_module(s_tir=True)
    class Print:
        @R.function(pure=False)
        def main(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
            y = R.print(x, format="x: {}")
            return x

    _check(Print)


def test_parse_multiple_pure_and_impure_funcs():
    @I.ir_module(s_tir=True)
    class Mixture:
        @R.function(pure=False)
        def print(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
            y = R.print(x, format="x: {}")
            return x

        @R.function(pure=False)
        def assert_func(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
            y = R.assert_op(R.const(False, dtype="bool"), x, format="x: {}")
            return x

        @R.function
        def main(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
            return x

    assert not Mixture["print"].is_pure
    assert not Mixture["assert_func"].is_pure
    assert Mixture["main"].is_pure
    _check(Mixture)


def test_function_with_void_return_type_may_be_used_as_statements():
    """Void return of calls do not need to be assigned"""

    @I.ir_module(s_tir=True)
    class Unsugared:
        @R.function(pure=False)
        def print(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
            y = R.print(x, format="x: {}")
            return x

        @R.function(pure=False)
        def assert_func(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
            y = R.assert_op(R.const(False, dtype="bool"), x, format="x: {}")
            return x

    @I.ir_module(s_tir=True)
    class Sugared:
        @R.function(pure=False)
        def print(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
            R.print(x, format="x: {}")
            return x

        @R.function(pure=False)
        def assert_func(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
            R.assert_op(R.const(False, dtype="bool"), x, format="x: {}")
            return x

    tvm.ir.assert_structural_equal(Unsugared, Sugared)


def test_function_with_non_void_return_type_must_be_assigned():
    """Non-void results must be assigned to a variable"""

    with pytest.raises(tvm.error.DiagnosticError):

        @R.function(pure=False)
        def func(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
            R.add(x, x)
            return x


def test_function_with_void_return_type_in_if_else():
    """Last statement in if/else may be a void return"""

    @I.ir_module(s_tir=True)
    class Unsugared:
        @R.function(pure=False)
        def conditional(x: R.Tensor((), "int32"), condition: R.Tensor((), "bool")) -> R.Tensor(
            (), "int32"
        ):
            if condition:
                y = R.print(x, format="True condition: {}")
            else:
                y = R.print(x, format="False condition: {}")
            return x

    @I.ir_module(s_tir=True)
    class Sugared:
        @R.function(pure=False)
        def conditional(x: R.Tensor((), "int32"), condition: R.Tensor((), "bool")) -> R.Tensor(
            (), "int32"
        ):
            if condition:
                R.print(x, format="True condition: {}")
            else:
                R.print(x, format="False condition: {}")
            return x

    _check(Sugared, Unsugared)


def test_call_pure_packed():
    @R.function
    def foo(x: R.Tensor((32, 32), "float32")) -> R.Tensor:
        z = R.call_pure_packed("vm.builtin.copy", x, ty_args=R.Tensor((32, 32), "float32"))
        return z

    x = relax.Var("x", R.Tensor((32, 32), "float32"))
    bb = relax.BlockBuilder()
    with bb.function("foo", (x)):
        z = bb.emit(
            R.call_pure_packed("vm.builtin.copy", x, ty_args=[R.Tensor((32, 32), "float32")])
        )
        bb.emit_func_output(z)

    _check(foo, bb.get()["foo"])


def test_call_pure_packed_returning_object():
    @R.function
    def foo() -> R.Any:
        z = R.call_pure_packed("dummy_func", ty_args=R.Any)
        return z

    bb = relax.BlockBuilder()
    with bb.function("foo", params=[]):
        z = bb.emit(R.call_pure_packed("dummy_func", ty_args=[relax.AnyType()]))
        bb.emit_func_output(z)

    _check(foo, bb.get()["foo"])


def test_private_function():
    @I.ir_module(s_tir=True)
    class Addition:
        @R.function(private=True)
        def main(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
            y = R.add(x, x)
            return y

    x = relax.Var("x", R.Tensor((), "int32"))
    bb = relax.BlockBuilder()
    with bb.function("main", (x), private=True):
        y = bb.emit(R.add(x, x))
        bb.emit_func_output(y)

    _check(Addition, bb.get())


def test_private_function_with_global_symbol_fail():
    with pytest.raises(tvm.error.DiagnosticError):

        @I.ir_module(s_tir=True)
        class Addition:
            @R.function(private=True)
            def main(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
                # it is an error to simultaneously mark a function private
                # and give it a global symbol manually
                R.func_attr({"global_symbol": "main"})
                y = R.add(x, x)
                return y

        # should not execute
        _check(Addition)


def test_private_function_with_global_symbol_no_module_fail():
    with pytest.raises(tvm.error.DiagnosticError):

        @R.function(private=True)
        def func(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
            R.func_attr({"global_symbol": "main"})
            y = R.add(x, x)
            return y

        # should not execute
        _check(func)


def test_macro_hygienic():
    x = R.prim_value(2)

    @R.macro(hygienic=True)
    def alloc_and_shape(dtype: str):
        alloc = R.builtin.alloc_tensor(R.shape([4, 4]), runtime_device_index=x, dtype=dtype)
        shape = R.shape_of(alloc)
        return shape

    x = R.prim_value(1)

    @R.function(private=True)
    def func(z: R.Tensor((4, 4), "float32")):
        shape = alloc_and_shape(dtype="float32")
        return shape

    @R.function(private=True)
    def expect(z: R.Tensor((4, 4), dtype="float32")) -> R.Shape([4, 4]):
        alloc: R.Tensor((4, 4), dtype="float32") = R.builtin.alloc_tensor(
            R.shape([4, 4]),
            R.dtype("float32"),
            R.prim_value(2),  # Make sure prim_value is 2
        )
        shape: R.Shape([4, 4]) = R.shape_of(alloc)
        shape_1: R.Shape([4, 4]) = shape
        return shape_1

    _check(func, expect)


def test_macro_non_hygienic():
    global global_x_var  # Lookup doesn't find this variable if it's not global

    global_x_var = R.prim_value(2)

    @R.macro(hygienic=False)
    def alloc_and_shape(dtype: str):
        alloc = R.builtin.alloc_tensor(
            R.shape([4, 4]), runtime_device_index=global_x_var, dtype=dtype
        )
        shape = R.shape_of(alloc)
        return shape

    global_x_var = R.prim_value(1)

    @R.function(private=True)
    def func(z: R.Tensor((4, 4), "float32")):
        shape = alloc_and_shape(dtype="float32")
        return shape

    @R.function(private=True)
    def expect(z: R.Tensor((4, 4), dtype="float32")) -> R.Shape([4, 4]):
        alloc: R.Tensor((4, 4), dtype="float32") = R.builtin.alloc_tensor(
            R.shape([4, 4]),
            R.dtype("float32"),
            R.prim_value(1),  # Make sure prim_value is 1
        )
        shape: R.Shape([4, 4]) = R.shape_of(alloc)
        shape_1: R.Shape([4, 4]) = shape
        return shape_1

    _check(func, expect)


def test_macro_no_variable_leak():
    with pytest.raises(tvm.error.DiagnosticError):

        @R.macro(hygienic=True)
        def add_two(value):
            x = value + R.const(1)  # `x` defined in macro
            y = x + R.const(1)
            return y

        @R.function(private=True)
        def func(t: R.Tensor((), "int32")):
            u = add_two(t)
            return x  # Should be undefined here


def test_reused_extern_func():
    """ExternFunc lookups can become bindings in EliminateCommonSubexpr"""

    @R.function(private=True)
    def parsed(x: R.Tensor((128, 128), "float32")) -> R.Tensor((128, 128), "float32"):
        func = R.ExternFunc("extern_func")
        gv0 = R.call_dps_packed(func, x, R.Tensor((128, 128), dtype="float32"))
        gv1 = R.call_dps_packed(func, gv0, R.Tensor((128, 128), dtype="float32"))
        return gv1

    x = relax.Var("x", R.Tensor((128, 128), "float32"))
    bb = relax.BlockBuilder()
    with bb.function("main", [x], private=True):
        func = bb.emit(relax.ExternFunc("extern_func"))
        y = bb.emit(relax.call_dps_packed(func, x, out_ty=R.Tensor((128, 128), "float32")))
        z = bb.emit(relax.call_dps_packed(func, y, out_ty=R.Tensor((128, 128), "float32")))
        bb.emit_func_output(z)

    expected = bb.get()["main"]

    _check(parsed, expected)


def test_extern_func_in_module():
    """Module-level parsing may produce function bindings"""

    @I.ir_module(s_tir=True)
    class parsed_module:
        my_ext = R.ExternFunc("my_ext")

        @R.function
        def func(a: R.Tensor((10, 10))) -> R.Tensor((10, 10)):
            return a

    @R.function
    def func(a: R.Tensor((10, 10))) -> R.Tensor((10, 10)):
        return a

    expected = tvm.IRModule({"my_ext": relax.ExternFunc("my_ext"), "func": func})

    _check(parsed_module, expected)


def test_define_relax_function_using_global_var():
    """A @R.function may call a GlobalVar

    When parsing a @R.function, the function's body may reference
    GlobalVar instances available in the calling python scope.  The
    resulting function should pass TVMScript's well-formed check, as
    the GlobalVar may be available in the IRModule for which the
    function is being defined.
    """

    @I.ir_module(s_tir=True)
    class DefinedAllAtOnce:
        @R.function
        def main(A: R.Tensor, B: R.Tensor):
            return DefinedAllAtOnce.subroutine(A, B)

        @R.function(private=True)
        def subroutine(A: R.Tensor, B: R.Tensor) -> R.Tensor:
            return R.matmul(A, B)

    @I.ir_module(s_tir=True)
    class MainDefinedLater:
        @R.function(private=True)
        def subroutine(A: R.Tensor, B: R.Tensor) -> R.Tensor:
            return R.matmul(A, B)

    subroutine_gvar = MainDefinedLater.get_global_var("subroutine")

    @R.function
    def main(A: R.Tensor, B: R.Tensor):
        return subroutine_gvar(A, B)

    MainDefinedLater["main"] = main

    tvm.ir.assert_structural_equal(DefinedAllAtOnce, MainDefinedLater)


def test_function_attributes_are_defined():
    """func.attrs defaults to an empty DictAttrs"""

    @I.ir_module(s_tir=True)
    class Module:
        @R.function
        def main(x: R.Tensor, shape: R.Shape(["m", "n"])):
            output = Module.subroutine(x, shape)
            return output

        @R.function
        def subroutine(x: R.Tensor, _: R.Shape(["m", "n"])) -> R.Tensor(["m", "n"]):
            q = x
            m, n = T.int64(), T.int64()
            z = R.match_cast(q, R.Tensor((m, n)))
            w = z
            return w

    for gvar, func in Module.functions.items():
        assert func.attrs is not None


@pytest.mark.xfail(reason="Bug: Implicit bounds not provided when parsing")
def test_function_symbolic_variables_are_annotated():
    """Symbolic variables must be exposed for struct inference

    Because Relax struct inference is performed while the function is
    being built, all constraints on symbolic variables that are used
    for simplifications must be provided to the analyzer.
    """

    @R.function(private=True)
    def inferred_ty(A: R.Tensor(["extent"])):
        extent = T.int64()
        output = R.strided_slice(A, [0], [0], [extent - 1])
        return output

    @R.function(private=True)
    def expected(A: R.Tensor(["extent"])) -> R.Tensor(["extent-1"]):
        extent = T.int64()
        output: R.Tensor([extent - 1]) = R.strided_slice(A, [0], [0], [extent - 1])
        return output

    tvm.ir.assert_structural_equal(inferred_ty, expected)


def test_non_declaration_prim_expr_emits_binding():
    """Only zero-argument dtype calls declare symbolic variables."""

    @R.function(private=True)
    def func(A: R.Tensor(["extent"], "float32")):
        extent = T.int64(4)
        output = A
        return output

    symbolic_extent = func.params[0].ty.shape[0]
    extent_binding = func.body.blocks[0].bindings[0]
    assert isinstance(extent_binding, relax.VarBinding)
    assert tvm.ir.is_prim_var(symbolic_extent)
    assert tvm.ir.is_prim_var(extent_binding.var)
    assert not symbolic_extent.same_as(extent_binding.var)
    tvm.ir.assert_structural_equal(extent_binding.value, T.int64(4))
    _check(func)


def test_primitive_assignments_emit_fresh_bindings():
    """Primitive expressions follow ordinary Relax assignment semantics."""

    @R.function(private=True)
    def func(scalar: R.Prim("int64"), tensor: R.Tensor([1], "float32")):
        alias = scalar
        arithmetic = alias + 1
        integer: R.Prim("int32") = 2
        floating: R.Prim("float32") = 2.5
        boolean: R.Prim("bool") = True
        compatibility = R.prim_value(arithmetic)
        tensor_alias = tensor
        return tensor_alias

    bindings = func.body.blocks[0].bindings
    assert len(bindings) == 7
    for binding in bindings:
        assert isinstance(binding, relax.VarBinding)
        assert not binding.var.same_as(binding.value)

    alias, arithmetic, integer, floating, boolean, compatibility, tensor_alias = bindings
    assert alias.value.same_as(func.params[0])
    assert compatibility.value.same_as(arithmetic.var)
    assert tensor_alias.value.same_as(func.params[1])
    for binding, dtype in [
        (integer, "int32"),
        (floating, "float32"),
        (boolean, "bool"),
    ]:
        assert isinstance(binding.var.ty, tvm.ir.PrimType)
        assert binding.var.ty.dtype == dtype
        assert tvm.ir.is_prim_expr(binding.value)
        assert not isinstance(binding.value, relax.Constant)

    source = func.script(show_all_ty=False)
    assert "R.prim_value" not in source
    for annotation in [
        "alias: T.int64",
        "integer: T.int32",
        "floating: T.float32",
        "boolean: T.bool",
    ]:
        assert annotation in source
    _check(func)


def test_shared_meta_var_skips_relax_bindings():
    """I.meta_var and T.meta_var are one explicit parser-time escape."""

    assert I.meta_var is T.meta_var

    @R.function(private=True)
    def func(A: R.Tensor(["N"], "float32")):
        N: R.Prim("int64") = T.int64()
        via_i = I.meta_var(N)
        via_t = T.meta_var(via_i)
        output = R.reshape(A, R.shape([via_t]))
        return output

    assert len(func.body.blocks[0].bindings) == 1
    source = func.script(show_all_ty=False)
    assert "meta_var" not in source
    _check(func)

    with pytest.raises(tvm.error.DiagnosticError):

        @R.function(private=True)
        def mismatched_declaration():
            value: R.Prim("float32") = T.int64()
            return value


def test_primitive_if_emits_fresh_result():
    """A primitive If has a fresh Relax result and typed branch terminators."""

    @R.function(private=True)
    def func(cond: R.Prim("bool"), lhs: R.Prim("int64"), rhs: R.Prim("int64")):
        if cond:
            output = lhs
        else:
            output = rhs
        return output

    outer_binding = func.body.blocks[0].bindings[0]
    assert isinstance(outer_binding, relax.VarBinding)
    assert isinstance(outer_binding.value, relax.If)
    assert isinstance(outer_binding.var.ty, tvm.ir.PrimType)
    assert not outer_binding.var.same_as(func.params[1])
    assert not outer_binding.var.same_as(func.params[2])
    for branch, param in [
        (outer_binding.value.true_branch, func.params[1]),
        (outer_binding.value.false_branch, func.params[2]),
    ]:
        assert not branch.blocks
        assert branch.body.same_as(param)

    source = func.script(show_all_ty=False)
    assert "R.prim_value" not in source
    assert source.count("output: T.int64") == 2
    _check(func)


def test_conditional_may_use_symbolic_variables_from_function_scope():
    """Symbolic variables from function scope may be used in branch

    This is a regression test.  In earlier implementations, the
    branches of `relax::If` were normalized with
    `EraseToWellDefinedInScope`, using a fresh variable scope.  While
    this had the intended behavior of preventing variables defined in
    a single branch from being usable outside of the conditional, it
    also caused the conditional's branches to treat function-scope
    symbolic variables as if they were undefined.

    """

    @R.function(private=True)
    def explicit_ty(
        A: R.Tensor(["N"], "float32"),
        B: R.Tensor(["N"], "float32"),
        cond: R.Prim("bool"),
    ) -> R.Tensor(["N"], "float32"):
        N = T.int64()

        if cond:
            out: R.Tensor([N], "float32") = A + B
        else:
            out: R.Tensor([N], "float32") = A * B

        return out

    @R.function(private=True)
    def inferred_ty(
        A: R.Tensor(["N"], "float32"),
        B: R.Tensor(["N"], "float32"),
        cond: R.Prim("bool"),
    ):
        N = T.int64()
        if cond:
            out = A + B
        else:
            out = A * B

        return out

    tvm.ir.assert_structural_equal(explicit_ty, inferred_ty)


def test_return_from_dataflow_block():
    """Return statements imply

    The `R.output` statement in a `R.dataflow()` block marks a
    variable that should be a `relax.Var` instead of a
    `relax.DataflowVar`, allowing it to be used outside of the
    `DataflowBlock` that defined it.  A relax function's output is not
    part of any binding, and must not contain any `DataflowVar`, so
    these are exposed implicitly.

    """

    @R.function(private=True)
    def output_then_return(A: R.Tensor([16], "float16")):
        with R.dataflow():
            B = R.add(A, A)
            C = R.multiply(B, B)
            R.output(C)

        return C

    @R.function(private=True)
    def return_inside_dataflow(A: R.Tensor([16], "float16")):
        with R.dataflow():
            B = R.add(A, A)
            C = R.multiply(B, B)
            return C

    tvm.ir.assert_structural_equal(output_then_return, return_inside_dataflow)


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