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

import pytest

import tvm
import tvm.testing
from tvm.relax.transform import DeadCodeElimination
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 verify(input, expected):
    tvm.ir.assert_structural_equal(DeadCodeElimination()(input), expected)


def test_simple():
    @tvm.script.ir_module
    class Input:
        @R.function
        def main(
            x: R.Tensor((2, 3, 28, 28), dtype="float32"),
            w: R.Tensor((4, 3, 3, 3), dtype="float32"),
            bias: R.Tensor((26, 26), dtype="float32"),
        ):
            # block 0
            with R.dataflow():
                gv: R.Tensor((2, 28, 28, 3), dtype="float32") = R.permute_dims(x, axes=[0, 2, 3, 1])
                gv1: R.Tensor((4, 3, 3, 3), dtype="float32") = R.permute_dims(w, axes=[0, 2, 3, 1])
                gv2: R.Tensor((2, 26, 26, 4), dtype="float32") = R.nn.conv2d(
                    gv,
                    gv1,
                    strides=[1, 1],
                    padding=[0, 0, 0, 0],
                    dilation=[1, 1],
                    groups=1,
                    data_layout="NHWC",
                    kernel_layout="OHWI",
                    out_layout="NHWC",
                    out_dtype="float32",
                )
                gv21: R.Tensor((2, 4, 26, 26), dtype="float32") = R.permute_dims(
                    gv2, axes=[0, 3, 1, 2]
                )
                gv22: R.Tensor((2, 4, 26, 26), dtype="float32") = R.add(gv21, bias)
                R.output(gv2)
            return gv2

    @I.ir_module(s_tir=True)
    class Expected:
        @R.function
        def main(
            x: R.Tensor((2, 3, 28, 28), dtype="float32"),
            w: R.Tensor((4, 3, 3, 3), dtype="float32"),
            bias: R.Tensor((26, 26), dtype="float32"),
        ):
            with R.dataflow():
                gv: R.Tensor((2, 28, 28, 3), dtype="float32") = R.permute_dims(x, axes=[0, 2, 3, 1])
                gv1: R.Tensor((4, 3, 3, 3), dtype="float32") = R.permute_dims(w, axes=[0, 2, 3, 1])
                gv2: R.Tensor((2, 26, 26, 4), dtype="float32") = R.nn.conv2d(
                    gv,
                    gv1,
                    strides=[1, 1],
                    padding=[0, 0, 0, 0],
                    dilation=[1, 1],
                    groups=1,
                    data_layout="NHWC",
                    kernel_layout="OHWI",
                    out_layout="NHWC",
                    out_dtype="float32",
                )
                R.output(gv2)
            return gv2

    verify(Input, Expected)


def test_2block():
    @tvm.script.ir_module
    class Input:
        @R.function
        def main(
            x: R.Tensor((2, 3, 28, 28), dtype="float32"),
            w: R.Tensor((4, 3, 3, 3), dtype="float32"),
            bias: R.Tensor((26, 26), dtype="float32"),
        ):
            # block 0
            with R.dataflow():
                gv: R.Tensor((2, 28, 28, 3), dtype="float32") = R.permute_dims(x, axes=[0, 2, 3, 1])
                gv1: R.Tensor((4, 3, 3, 3), dtype="float32") = R.permute_dims(w, axes=[0, 2, 3, 1])
                gv2: R.Tensor((2, 26, 26, 4), dtype="float32") = R.nn.conv2d(
                    gv,
                    gv1,
                    strides=[1, 1],
                    padding=[0, 0, 0, 0],
                    dilation=[1, 1],
                    groups=1,
                    data_layout="NHWC",
                    kernel_layout="OHWI",
                    out_layout="NHWC",
                    out_dtype="float32",
                )
                gv21: R.Tensor((2, 4, 26, 26), dtype="float32") = R.permute_dims(
                    gv2, axes=[0, 3, 1, 2]
                )
                gv22: R.Tensor((2, 4, 26, 26), dtype="float32") = R.add(gv21, bias)
                R.output(gv2)
            gv3 = R.astype(gv2, dtype="float16")
            return gv3

    @I.ir_module(s_tir=True)
    class Expected:
        @R.function
        def main(
            x: R.Tensor((2, 3, 28, 28), dtype="float32"),
            w: R.Tensor((4, 3, 3, 3), dtype="float32"),
            bias: R.Tensor((26, 26), dtype="float32"),
        ):
            with R.dataflow():
                gv: R.Tensor((2, 28, 28, 3), dtype="float32") = R.permute_dims(x, axes=[0, 2, 3, 1])
                gv1: R.Tensor((4, 3, 3, 3), dtype="float32") = R.permute_dims(w, axes=[0, 2, 3, 1])
                gv2: R.Tensor((2, 26, 26, 4), dtype="float32") = R.nn.conv2d(
                    gv,
                    gv1,
                    strides=[1, 1],
                    padding=[0, 0, 0, 0],
                    dilation=[1, 1],
                    groups=1,
                    data_layout="NHWC",
                    kernel_layout="OHWI",
                    out_layout="NHWC",
                    out_dtype="float32",
                )
                R.output(gv2)
            gv3: R.Tensor((2, 26, 26, 4), dtype="float16") = R.astype(gv2, dtype="float16")
            return gv3

    verify(Input, Expected)


def check_if_func_exists(mod, func_name):
    gvs = [gv.name_hint for gv in mod.get_global_vars()]
    return func_name in gvs


def test_unused_relax_func():
    @tvm.script.ir_module
    class InputModule:
        @T.prim_func(s_tir=True)
        def tir_add(
            x: T.Buffer((16, 16), "float32"),
            y: T.Buffer((16, 16), "float32"),
            z: T.Buffer((16, 16), "float32"),
        ) -> None:
            for i, j in T.grid(16, 16):
                with T.sblock("add"):
                    vi, vj = T.axis.remap("SS", [i, j])
                    z[vi, vj] = x[vi, vj] + y[vi, vj]

        @R.function(private=True)
        def unused_func(x: R.Tensor((16, 16), "float32"), w: R.Tensor((16, 16), "float32")):
            gv0 = R.add(x, w)
            return gv0

        @R.function
        def main(x: R.Tensor((16, 16), "float32"), w: R.Tensor((16, 16), "float32")) -> R.Tensor(
            (16, 16), "float32"
        ):
            gv0 = R.call_tir(InputModule.tir_add, (x, w), R.Tensor((16, 16), dtype="float32"))
            return gv0

    mod = InputModule
    assert mod
    new_mod = DeadCodeElimination()(mod)
    assert check_if_func_exists(new_mod, "main")
    assert check_if_func_exists(new_mod, "tir_add")
    assert not check_if_func_exists(new_mod, "unused_func")


provide_entry_func_name = tvm.testing.parameter(True, False)


def test_unused_relax_func_custom_entry_func(provide_entry_func_name):
    @tvm.script.ir_module
    class InputModule:
        @T.prim_func(private=True, s_tir=True)
        def tir_add(
            x: T.Buffer((16, 16), "float32"),
            y: T.Buffer((16, 16), "float32"),
            z: T.Buffer((16, 16), "float32"),
        ) -> None:
            for i, j in T.grid(16, 16):
                with T.sblock("add"):
                    vi, vj = T.axis.remap("SS", [i, j])
                    z[vi, vj] = x[vi, vj] + y[vi, vj]

        @R.function(private=True)
        def unused_func(x: R.Tensor((16, 16), "float32"), w: R.Tensor((16, 16), "float32")):
            gv0 = R.add(x, w)
            return gv0

        @R.function
        def foo(x: R.Tensor((16, 16), "float32"), w: R.Tensor((16, 16), "float32")) -> R.Tensor(
            (16, 16), "float32"
        ):
            gv0 = R.call_tir(InputModule.tir_add, (x, w), R.Tensor((16, 16), dtype="float32"))
            return gv0

    mod = InputModule
    assert mod

    if provide_entry_func_name:
        entry_functions = ["foo"]
    else:
        entry_functions = None

    # Test entry function other than "main".
    new_mod = DeadCodeElimination(entry_functions=entry_functions)(mod)
    assert check_if_func_exists(new_mod, "foo")
    assert check_if_func_exists(new_mod, "tir_add")
    assert not check_if_func_exists(new_mod, "unused_func")


def test_tracking_through_externally_exposed_func(provide_entry_func_name):
    @tvm.script.ir_module
    class InputModule:
        @T.prim_func(private=True, s_tir=True)
        def tir_add(
            x: T.Buffer((16, 16), "float32"),
            y: T.Buffer((16, 16), "float32"),
            z: T.Buffer((16, 16), "float32"),
        ) -> None:
            for i, j in T.grid(16, 16):
                with T.sblock("add"):
                    vi, vj = T.axis.remap("SS", [i, j])
                    z[vi, vj] = x[vi, vj] + y[vi, vj]

        @R.function(private=True)
        def unused_func(x: R.Tensor((16, 16), "float32"), w: R.Tensor((16, 16), "float32")):
            gv0 = R.add(x, w)
            return gv0

        @R.function
        def foo(x: R.Tensor((16, 16), "float32"), w: R.Tensor((16, 16), "float32")) -> R.Tensor(
            (16, 16), "float32"
        ):
            gv0 = R.call_tir(InputModule.tir_add, (x, w), R.Tensor((16, 16), dtype="float32"))
            return gv0

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

    mod = InputModule
    assert mod

    # Test tracking of usage through externally-exposed function
    new_mod = DeadCodeElimination(entry_functions=["main"])(mod)
    assert check_if_func_exists(new_mod, "main")
    assert check_if_func_exists(new_mod, "foo")
    assert check_if_func_exists(new_mod, "tir_add")
    assert not check_if_func_exists(new_mod, "unused_func")


def test_unused_relax_func_symbolic_shape():
    # Test with relax function w/ symbolic shape.
    @tvm.script.ir_module(check_well_formed=False)
    class InputModule:
        @T.prim_func(s_tir=True)
        def tir_matmul(
            x_handle: T.handle,
            y_handle: T.handle,
            z_handle: T.handle,
        ) -> None:
            m = T.int64()
            n = T.int64()
            k = T.int64()
            x = T.match_buffer(x_handle, (m, n), "float32")
            y = T.match_buffer(y_handle, (n, k), "float32")
            z = T.match_buffer(z_handle, (m, k), "float32")
            for i, j, k in T.grid(m, k, n):
                with T.sblock("matmul"):
                    vi, vj, vk = T.axis.remap("SSR", [i, j, k])
                    with T.init():
                        z[vi, vj] = 0.0
                    z[vi, vj] = z[vi, vj] + x[vi, vk] * y[vk, vj]

        @R.function(private=True)
        def unused_func(x: R.Tensor(("m", "n"), "float32"), w: R.Tensor(("n", "k"), "float32")):
            gv0 = R.add(x, w)
            return gv0

        @R.function
        def main(x: R.Tensor(("m", "n"), "float32"), w: R.Tensor(("n", "k"), "float32")):
            m, k = T.int64(), T.int64()
            gv0 = R.call_tir(InputModule.tir_matmul, (x, w), R.Tensor((m, k), dtype="float32"))
            return gv0

    mod = InputModule
    assert mod

    new_mod = DeadCodeElimination()(mod)
    assert check_if_func_exists(new_mod, "main")
    assert check_if_func_exists(new_mod, "tir_matmul")
    assert not check_if_func_exists(new_mod, "unused_func")


def test_unused_prim_func():
    @tvm.script.ir_module
    class InputModule:
        @T.prim_func(s_tir=True)
        def unused_func(
            x: T.Buffer((16, 16), "float32"),
            y: T.Buffer((16, 16), "float32"),
            z: T.Buffer((16, 16), "float32"),
        ) -> None:
            T.func_attr({"global_symbol": "tir_unused"})
            for i, j in T.grid(16, 16):
                with T.sblock("add"):
                    vi, vj = T.axis.remap("SS", [i, j])
                    z[vi, vj] = x[vi, vj] + y[vi, vj]

        @R.function
        def relax_add(x: R.Tensor((16, 16), "float32"), w: R.Tensor((16, 16), "float32")):
            gv0 = R.add(x, w)
            return gv0

        @R.function
        def main(x: R.Tensor((16, 16), "float32"), w: R.Tensor((16, 16), "float32")) -> R.Tensor(
            (16, 16), "float32"
        ):
            gv0 = InputModule.relax_add(x, w)
            return gv0

    mod = InputModule
    assert mod
    new_mod = DeadCodeElimination()(mod)
    assert check_if_func_exists(new_mod, "main")
    assert check_if_func_exists(new_mod, "relax_add")
    # RemoveUnusedFunction pass won't remove the function with global symbol for the external linkage.
    assert check_if_func_exists(new_mod, "unused_func")


def test_preserve_indirectly_used_prim_func():
    @tvm.script.ir_module
    class InputModule:
        @R.function
        def main(x: R.Tensor((16, 16), "float32"), w: R.Tensor((16, 16), "float32")) -> R.Tensor(
            (16, 16), "float32"
        ):
            gv0 = R.call_tir(
                InputModule.tir_add_tensors,
                [x, w],
                out_ty=R.Tensor((16, 16), "float32"),
            )
            return gv0

        @T.prim_func(private=True, s_tir=True)
        def tir_add_tensors(
            x: T.Buffer((16, 16), "float32"),
            y: T.Buffer((16, 16), "float32"),
            z: T.Buffer((16, 16), "float32"),
        ):
            for i, j in T.grid(16, 16):
                with T.sblock("add"):
                    vi, vj = T.axis.remap("SS", [i, j])
                    z[vi, vj] = InputModule.tir_add_float32(x[vi, vj], y[vi, vj])

        @T.prim_func(private=True, s_tir=True)
        def tir_add_float32(x: T.float32, y: T.float32) -> T.float32:
            return x + y

    mod = InputModule
    assert mod
    new_mod = DeadCodeElimination()(mod)

    tvm.ir.assert_structural_equal(mod, new_mod)


def test_multiple_unused_funcs():
    @tvm.script.ir_module
    class InputModule:
        @T.prim_func(s_tir=True)
        def unused_func1(
            x: T.Buffer((16, 16), "float32"),
            y: T.Buffer((16, 16), "float32"),
            z: T.Buffer((16, 16), "float32"),
        ) -> None:
            T.func_attr({"global_symbol": "tir_unused"})
            for i, j in T.grid(16, 16):
                with T.sblock("add"):
                    vi, vj = T.axis.remap("SS", [i, j])
                    z[vi, vj] = x[vi, vj] + y[vi, vj]

        @R.function(private=True)
        def unused_func2(x: R.Tensor((16, 16), "float32"), w: R.Tensor((16, 16), "float32")):
            gv0 = R.add(x, w)
            return gv0

        @R.function
        def main(x: R.Tensor((16, 16), "float32"), w: R.Tensor((16, 16), "float32")) -> R.Tensor(
            (16, 16), "float32"
        ):
            gv0 = R.add(x, w)
            return gv0

    mod = InputModule
    assert mod

    new_mod = DeadCodeElimination()(mod)
    assert check_if_func_exists(new_mod, "main")
    # RemoveUnusedFunction pass won't remove the function with global symbol for the external linkage.
    assert check_if_func_exists(new_mod, "unused_func1")
    assert not check_if_func_exists(new_mod, "unused_func2")


def test_unused_dfb():
    # test if an unused dataflow block can be removed.
    @tvm.script.ir_module
    class Input:
        @R.function
        def main(
            x: R.Tensor((2, 3, 28, 28), dtype="float32"),
            w: R.Tensor((4, 3, 3, 3), dtype="float32"),
        ):
            # block 0
            with R.dataflow():
                lv0: R.Tensor((2, 28, 28, 3), dtype="float32") = R.permute_dims(
                    x, axes=[0, 2, 3, 1]
                )
                lv1: R.Tensor((4, 3, 3, 3), dtype="float32") = R.permute_dims(w, axes=[0, 2, 3, 1])
                lv2: R.Tensor((2, 26, 26, 4), dtype="float32") = R.nn.conv2d(
                    lv0,
                    lv1,
                    data_layout="NHWC",
                    kernel_layout="OHWI",
                    out_layout="NHWC",
                )
                lv3: R.Tensor((2, 4, 26, 26), dtype="float32") = R.permute_dims(
                    lv2, axes=[0, 3, 1, 2]
                )
                R.output(lv2)
            gv3 = R.astype(lv2, dtype="float16")
            # dead block
            with R.dataflow():
                lv4: R.Tensor((2, 4, 26, 26), dtype="float16") = R.permute_dims(
                    gv3, axes=[0, 3, 1, 2]
                )
                R.output(lv4)
            return gv3

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(
            x: R.Tensor((2, 3, 28, 28), dtype="float32"),
            w: R.Tensor((4, 3, 3, 3), dtype="float32"),
        ):
            # block 0
            with R.dataflow():
                lv0: R.Tensor((2, 28, 28, 3), dtype="float32") = R.permute_dims(
                    x, axes=[0, 2, 3, 1]
                )
                lv1: R.Tensor((4, 3, 3, 3), dtype="float32") = R.permute_dims(w, axes=[0, 2, 3, 1])
                lv2: R.Tensor((2, 26, 26, 4), dtype="float32") = R.nn.conv2d(
                    lv0,
                    lv1,
                    data_layout="NHWC",
                    kernel_layout="OHWI",
                    out_layout="NHWC",
                )
                R.output(lv2)
            gv3 = R.astype(lv2, dtype="float16")
            return gv3

    verify(Input, Expected)


def test_unused_dfb2():
    # test if an unused dataflow block can be removed.
    @tvm.script.ir_module
    class Input:
        @R.function
        def main(
            x: R.Tensor((2, 3, 28, 28), dtype="float32"),
            w: R.Tensor((4, 3, 3, 3), dtype="float32"),
        ):
            # dead block
            with R.dataflow():
                lv0: R.Tensor((2, 28, 28, 3), dtype="float32") = R.permute_dims(
                    x, axes=[0, 2, 3, 1]
                )
                R.output(lv0)

            gv_x = R.astype(x, dtype="float16")
            gv_w = R.astype(w, dtype="float16")

            with R.dataflow():
                lv1: R.Tensor((2, 28, 28, 3), dtype="float16") = R.permute_dims(
                    gv_x, axes=[0, 2, 3, 1]
                )
                lv2: R.Tensor((4, 3, 3, 3), dtype="float16") = R.permute_dims(
                    gv_w, axes=[0, 2, 3, 1]
                )
                lv3: R.Tensor((2, 26, 26, 4), dtype="float16") = R.nn.conv2d(
                    lv1,
                    lv2,
                    data_layout="NHWC",
                    kernel_layout="OHWI",
                    out_layout="NHWC",
                )
                # dead instruction -> usee lv1 also dead.
                lv4: R.Tensor((2, 3, 28, 28), dtype="float32") = R.permute_dims(
                    lv0, axes=[0, 3, 1, 2]
                )
                R.output(lv3)
            return lv3

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(
            x: R.Tensor((2, 3, 28, 28), dtype="float32"),
            w: R.Tensor((4, 3, 3, 3), dtype="float32"),
        ):
            gv_x = R.astype(x, dtype="float16")
            gv_w = R.astype(w, dtype="float16")

            with R.dataflow():
                lv1: R.Tensor((2, 28, 28, 3), dtype="float16") = R.permute_dims(
                    gv_x, axes=[0, 2, 3, 1]
                )
                lv2: R.Tensor((4, 3, 3, 3), dtype="float16") = R.permute_dims(
                    gv_w, axes=[0, 2, 3, 1]
                )
                lv3: R.Tensor((2, 26, 26, 4), dtype="float16") = R.nn.conv2d(
                    lv1,
                    lv2,
                    data_layout="NHWC",
                    kernel_layout="OHWI",
                    out_layout="NHWC",
                )
                R.output(lv3)
            return lv3

    verify(Input, Expected)


def test_extern_func():
    """DeadCodeElimination should retain the ExternFunc in the IRModule."""

    builder = tvm.relax.BlockBuilder()
    builder.add_func(tvm.relax.extern("extern_func"), "extern_func")
    before = builder.get()

    verify(before, before)


def test_recursively_defined_lambda():
    """DCE may be applied to recursively-defined functions

    While most expressions may only contain references to
    previously-defined variables, local Relax function definitions may
    contain references to themselves.

    This is a regression test.  In previous implementations, the
    recursive use of `while_loop` resulted in an error, as
    `while_loop` was not considered in-scope by the `CollectVarUsage`
    utility until after the body of `while_loop` had been visited.

    """

    @I.ir_module(s_tir=True)
    class Before:
        @R.function
        def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor:
            @R.function
            def while_loop(i: R.Tensor((), "int32"), s: R.Tensor((2, 3), "float32")) -> R.Tensor(
                (2, 3), "float32"
            ):
                cond = R.call_pure_packed(
                    "test.vm.less", i, R.const(10), ty_args=R.Tensor((), dtype="bool")
                )
                c = R.const(1, dtype="int32")
                if cond:
                    new_i = R.add(i, c)
                    new_s = R.add(s, x)
                    r = while_loop(new_i, new_s)
                else:
                    r = s
                return r

            gv = while_loop(R.const(0), x)
            return gv

    Expected = Before

    verify(Before, Expected)


def test_recursively_defined_closure():
    """DCE may be applied to recursively-defined closures

    This test is identical to `test_recursively_defined_lambda`,
    except that the threshold for recursion is defined in an enclosed
    variable outside of the recursive function.

    """

    @I.ir_module(s_tir=True)
    class Before:
        @R.function
        def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor:
            threshold = R.const(10)

            @R.function
            def while_loop(i: R.Tensor((), "int32"), s: R.Tensor((2, 3), "float32")) -> R.Tensor(
                (2, 3), "float32"
            ):
                cond = R.call_pure_packed(
                    "test.vm.less", i, threshold, ty_args=R.Tensor((), dtype="bool")
                )
                c = R.const(1, dtype="int32")
                if cond:
                    new_i = R.add(i, c)
                    new_s = R.add(s, x)
                    r = while_loop(new_i, new_s)
                else:
                    r = s
                return r

            gv = while_loop(R.const(0), x)
            return gv

    Expected = Before

    verify(Before, Expected)


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