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

import tvm
import tvm.testing
from tvm.script import tirx as T


def nop():
    return tvm.tirx.Evaluate(1)


def test_remove_no_op():
    i = tvm.tirx.Var("i", "int32")
    j = tvm.tirx.Var("j", "int32")
    k = tvm.tirx.Var("k", "int32")
    m = tvm.tirx.Var("m", "int32")
    n = tvm.tirx.Var("n", "int32")
    dtype = "int64"
    Ab = tvm.tirx.decl_buffer((n,), dtype)
    stmt = tvm.tirx.For(
        i,
        0,
        4,
        tvm.tirx.ForKind.SERIAL,
        tvm.tirx.For(
            j,
            0,
            n,
            tvm.tirx.ForKind.SERIAL,
            tvm.tirx.For(
                k,
                0,
                m,
                tvm.tirx.ForKind.SERIAL,
                tvm.tirx.IfThenElse(
                    (i * m + j + k < n), tvm.tirx.Evaluate(m), tvm.tirx.Evaluate(n)
                ),
            ),
        ),
    )

    mod = tvm.IRModule.from_expr(tvm.tirx.PrimFunc([Ab], stmt))
    ret = tvm.tirx.transform.RemoveNoOp()(mod)["main"].body

    assert isinstance(ret, tvm.tirx.Evaluate)
    store = tvm.tirx.BufferStore(Ab, tvm.tirx.BufferLoad(Ab, [i]) + 1, [i + 1])
    stmt2 = tvm.tirx.SeqStmt([nop(), tvm.tirx.SeqStmt([store, nop()])])

    mod = tvm.IRModule.from_expr(tvm.tirx.PrimFunc([Ab], stmt2))
    ret = tvm.tirx.transform.RemoveNoOp()(mod)["main"].body
    assert ret == store

    # remove zero extent loop
    stmt3 = tvm.tirx.For(i, 0, 0, tvm.tirx.ForKind.SERIAL, store)
    mod = tvm.IRModule.from_expr(tvm.tirx.PrimFunc([Ab], stmt3))
    ret = tvm.tirx.transform.RemoveNoOp()(mod)["main"].body
    assert isinstance(ret, tvm.tirx.Evaluate)


def test_remove_no_op_with_invalid_extent():
    @T.prim_func(s_tir=True)
    def main(A: T.Buffer((16), "int32"), B: T.Buffer((16), "int32")) -> None:
        for i in T.serial(16):
            for j in T.serial(i - 20):
                B[i] = A[i] + j

    mod = tvm.ir.module.IRModule.from_expr(main)
    ret = tvm.tirx.transform.RemoveNoOp()(mod)["main"].body
    assert isinstance(ret, tvm.tirx.Evaluate)


def _apply_remove_no_op(mod, max_simplification_steps=0):
    """Helper function to apply RemoveNoOp transform with config."""
    config = {
        "tirx.RemoveNoOp": {
            "max_simplification_steps": max_simplification_steps,
        }
    }
    with tvm.transform.PassContext(config=config):
        mod = tvm.tirx.transform.RemoveNoOp()(mod)
    return mod


def test_remove_empty_for_loop():
    """A for-loop whose body is a no-op is itself a no-op."""

    @T.prim_func(private=True, s_tir=True)
    def before():
        for i in T.serial(16):
            T.evaluate(0)

    @T.prim_func(private=True, s_tir=True)
    def expected():
        T.evaluate(0)

    mod = tvm.IRModule.from_expr(before)
    mod = _apply_remove_no_op(mod)
    tvm.ir.assert_structural_equal(mod["main"], expected)


def test_remove_zero_extent_loop():
    """A for-loop with no extent is a no-op."""

    @T.prim_func(private=True, s_tir=True)
    def before(A: T.Buffer(16, "int32")):
        for i in T.serial(0):
            A[i] = 42

    @T.prim_func(private=True, s_tir=True)
    def expected(A: T.Buffer(16, "int32")):
        T.evaluate(0)

    mod = tvm.IRModule.from_expr(before)
    mod = _apply_remove_no_op(mod)
    tvm.ir.assert_structural_equal(mod["main"], expected)


def test_remove_unused_let():
    """With flat Bind, unused bindings are preserved by remove_no_op.

    Unused Bind elimination requires a separate two-pass approach
    and is not handled by the current remove_no_op pass.
    """

    @T.prim_func(private=True, s_tir=True)
    def before(A: T.Buffer(16, "int32")):
        x = 5
        for i in T.serial(16):
            A[i] = 0

    @T.prim_func(private=True, s_tir=True)
    def expected(A: T.Buffer(16, "int32")):
        x = 5
        for i in T.serial(16):
            A[i] = 0

    mod = tvm.IRModule.from_expr(before)
    mod = _apply_remove_no_op(mod)
    tvm.ir.assert_structural_equal(mod["main"], expected)


def test_remove_let_used_only_in_no_op():
    """With flat Bind, unused bindings are preserved by remove_no_op.

    The zero-extent for loop is removed, but the Bind for x remains
    since unused Bind elimination is not handled by remove_no_op.
    """

    @T.prim_func(private=True, s_tir=True)
    def before(A: T.Buffer(16, "int32")):
        x = 5
        for i in T.serial(0):
            A[i] = x

    @T.prim_func(private=True, s_tir=True)
    def expected(A: T.Buffer(16, "int32")):
        x = 5
        T.evaluate(0)

    mod = tvm.IRModule.from_expr(before)
    mod = _apply_remove_no_op(mod)
    tvm.ir.assert_structural_equal(mod["main"], expected)


def test_keep_side_effects_of_let():
    """Side-effect Bind is preserved as-is by remove_no_op."""

    @T.prim_func(private=True, s_tir=True)
    def before():
        x = T.call_extern("extern_func", dtype="int32")
        T.evaluate(0)

    @T.prim_func(private=True, s_tir=True)
    def expected():
        x = T.call_extern("extern_func", dtype="int32")
        T.evaluate(0)

    mod = tvm.IRModule.from_expr(before)
    mod = _apply_remove_no_op(mod)
    tvm.ir.assert_structural_equal(mod["main"], expected)


def test_remove_empty_then_case():
    """A no-op then_case can be removed."""

    @T.prim_func(private=True, s_tir=True)
    def before(A: T.Buffer(16, "int32")):
        for i in T.serial(16):
            if i < 8:
                T.evaluate(0)
            else:
                A[i] = 42

    @T.prim_func(private=True, s_tir=True)
    def expected(A: T.Buffer(16, "int32")):
        for i in T.serial(16):
            if not (i < 8):
                A[i] = 42

    mod = tvm.IRModule.from_expr(before)
    mod = _apply_remove_no_op(mod)
    tvm.ir.assert_structural_equal(mod["main"], expected)


def test_remove_empty_else_case():
    """A no-op else_case can be removed."""

    @T.prim_func(private=True, s_tir=True)
    def before(A: T.Buffer(16, "int32")):
        for i in T.serial(16):
            if i < 8:
                A[i] = 42
            else:
                T.evaluate(0)

    @T.prim_func(private=True, s_tir=True)
    def expected(A: T.Buffer(16, "int32")):
        for i in T.serial(16):
            if i < 8:
                A[i] = 42

    mod = tvm.IRModule.from_expr(before)
    mod = _apply_remove_no_op(mod)
    tvm.ir.assert_structural_equal(mod["main"], expected)


def test_suppress_removal_of_unused_write():
    """Sequential writes to the same location are not removed.

    Dataflow analysis is no longer supported.
    """

    @T.prim_func(private=True, s_tir=True)
    def before(A: T.Buffer(16, "int32")):
        for i in T.serial(16):
            A[i] = 100
            A[i] = 42

    mod = tvm.IRModule.from_expr(before)
    mod = _apply_remove_no_op(mod)
    tvm.ir.assert_structural_equal(mod["main"], before)


def test_keep_first_write_when_used():
    """For two sequential writes, keep the first if it is used"""

    @T.prim_func(private=True, s_tir=True)
    def before(A: T.Buffer(16, "int32")):
        for i in T.serial(16):
            A[i] = 100
            A[i] = A[i] + 1

    mod = tvm.IRModule.from_expr(before)
    mod = _apply_remove_no_op(mod)
    tvm.ir.assert_structural_equal(mod["main"], before)


def test_keep_partially_overwritten_loop():
    """Keep partially overwritten regions

    If the second loop doesn't entirely overwrite the first, the first
    may not be removed be kept.
    """

    @T.prim_func(private=True, s_tir=True)
    def before(A: T.Buffer(16, "int32")):
        for i in T.serial(16):
            A[i] = 100

        for i in T.serial(16):
            if i < 12:
                A[i] = 42

    mod = tvm.IRModule.from_expr(before)
    mod = _apply_remove_no_op(mod)
    tvm.ir.assert_structural_equal(mod["main"], before)


def test_remove_read_write():
    """Writing a value to the same location as was just read is a no-op."""

    @T.prim_func(private=True, s_tir=True)
    def before(A: T.Buffer(1, "int32")):
        A[0] = A[0]

    @T.prim_func(private=True, s_tir=True)
    def expected(A: T.Buffer(1, "int32")):
        T.evaluate(0)

    mod = tvm.IRModule.from_expr(before)
    mod = _apply_remove_no_op(mod)
    tvm.ir.assert_structural_equal(mod["main"], expected)


def test_keep_read_write_to_different_indices():
    """Writing a value to a different index should not be removed"""

    @T.prim_func(private=True, s_tir=True)
    def before(A: T.Buffer(16, "int32")):
        for i in T.serial(15):
            A[i] = A[i + 1]

    mod = tvm.IRModule.from_expr(before)
    mod = _apply_remove_no_op(mod)
    tvm.ir.assert_structural_equal(mod["main"], before)


def test_remove_read_write_same_index_different_expression():
    """Writing a value to the same location as the read is a no-op.

    If the value of the index can be proven to be the same, then the
    store can be removed. With flat Bind, the Bind for i and the
    enclosing loops remain since unused Bind elimination is not
    handled by remove_no_op.
    """

    @T.prim_func(private=True, s_tir=True)
    def before(A: T.Buffer(16, "int32")):
        for io, ii in T.grid(4, 4):
            i: T.let[T.int32] = 4 * io + ii
            A[4 * io + ii] = A[i]

    @T.prim_func(private=True, s_tir=True)
    def expected(A: T.Buffer(16, "int32")):
        for io in range(4):
            for ii in range(4):
                i: T.let[T.int32] = 4 * io + ii

    mod = tvm.IRModule.from_expr(before)
    mod = _apply_remove_no_op(mod)
    tvm.ir.assert_structural_equal(mod["main"], expected)


def test_remove_read_write_same_index_using_constraint():
    """Writing a value to the same location as the read is a no-op.

    If the value of the index can be proven to be the same, then the
    no-op can be removed.  This may require using the a constraint
    that is known from a conditional containing the read/write.
    """

    @T.prim_func(private=True, s_tir=True)
    def before(A: T.Buffer(16, "int32")):
        for i in T.serial(16):
            if i != 0:
                A[i] = A[i - 1]
            else:
                A[i] = A[0]

    @T.prim_func(private=True, s_tir=True)
    def expected(A: T.Buffer(16, "int32")):
        for i in T.serial(16):
            if i != 0:
                A[i] = A[i - 1]

    mod = tvm.IRModule.from_expr(before)
    mod = _apply_remove_no_op(mod)
    tvm.ir.assert_structural_equal(mod["main"], expected)


@pytest.mark.xfail(reason="Dead alloc removal not yet implemented for flat AllocBuffer")
def test_remove_empty_temporary():
    """An allocation with a no-op body is a no-op."""

    @T.prim_func(private=True, s_tir=True)
    def before():
        A = T.alloc_buffer((16,), "int32", scope="local")
        T.evaluate(0)

    @T.prim_func(private=True, s_tir=True)
    def expected():
        T.evaluate(0)

    mod = tvm.IRModule.from_expr(before)
    mod = _apply_remove_no_op(mod)
    tvm.ir.assert_structural_equal(mod["main"], expected)


@pytest.mark.xfail(reason="Dead alloc removal not yet implemented for flat AllocBuffer")
def test_remove_empty_temporary_with_decl_buffer():
    """Remove DeclBuffer alongside Allocate

    If an unused allocation is removed, any DeclBuffer instances that
    refer to it should also be removed.
    """

    @T.prim_func(private=True, s_tir=True)
    def before():
        A = T.decl_buffer([4, 4], "int32", scope="local")
        A_flat = T.decl_buffer(16, "int32", scope="local", data=A.data)
        T.evaluate(0)

    @T.prim_func(private=True, s_tir=True)
    def expected():
        T.evaluate(0)

    mod = tvm.IRModule.from_expr(before)
    mod = _apply_remove_no_op(mod)
    tvm.ir.assert_structural_equal(mod["main"], expected)


@pytest.mark.xfail(reason="Not implemented yet")
def test_remove_unused_temporary():
    """An unused allocation is a no-op."""

    @T.prim_func(private=True, s_tir=True)
    def before(A: T.Buffer(16, "int32")):
        B = T.alloc_buffer((16,), "int32", scope="local")
        for i in T.serial(16):
            A[i] = 1

    @T.prim_func(private=True, s_tir=True)
    def expected(A: T.Buffer(16, "int32")):
        for i in T.serial(16):
            A[i] = 1

    mod = tvm.IRModule.from_expr(before)
    mod = _apply_remove_no_op(mod)
    tvm.ir.assert_structural_equal(mod["main"], expected)


@pytest.mark.xfail(reason="Not implemented yet")
def test_remove_unused_write_into_temporary():
    """A write that only impacts a temporary allocation is a no-op."""

    @T.prim_func(private=True, s_tir=True)
    def before():
        A = T.decl_buffer([16], "int32", scope="local")
        for i in T.serial(16):
            A[i] = 0

    @T.prim_func(private=True, s_tir=True)
    def expected():
        T.evaluate(0)

    mod = tvm.IRModule.from_expr(before)
    mod = _apply_remove_no_op(mod)
    tvm.ir.assert_structural_equal(mod["main"], expected)


def test_keep_used_write_into_temporary():
    """A write into a temporary that is used later must be kept."""

    @T.prim_func(private=True, s_tir=True)
    def before(B: T.Buffer(16, "int32")):
        A = T.decl_buffer([16], "int32", scope="local")
        for i in T.serial(16):
            A[i] = 0

        for i in T.serial(16):
            B[i] = A[i]

    mod = tvm.IRModule.from_expr(before)
    mod = _apply_remove_no_op(mod)
    tvm.ir.assert_structural_equal(mod["main"], before)


@pytest.mark.xfail(reason="Not implemented yet")
def test_remove_write_into_temporary():
    """A write that only impacts a temporary allocation is a no-op."""

    @T.prim_func(private=True, s_tir=True)
    def before(A: T.Buffer(16, "int32"), C: T.Buffer(1, "int32")):
        B = T.decl_buffer([16], "int32", scope="local")
        for i in T.serial(16):
            B[i] = A[i]

        C[0] = 0
        for i in T.serial(16):
            C[0] = C[0] + B[i]

        for i in T.serial(16):
            B[i] = 0

    @T.prim_func(private=True, s_tir=True)
    def expected(A: T.Buffer(16, "int32"), C: T.Buffer(1, "int32")):
        B = T.decl_buffer([16], "int32", scope="local")
        for i in T.serial(16):
            B[i] = A[i]

        C[0] = 0
        for i in T.serial(16):
            C[0] = C[0] + B[i]

    mod = tvm.IRModule.from_expr(before)
    mod = _apply_remove_no_op(mod)
    tvm.ir.assert_structural_equal(mod["main"], expected)


def test_certain_condition():
    """The conditon of the If-Else node is certain.
    This would cause `Segmentation fault` error before."""

    @T.prim_func(private=True, s_tir=True)
    def before():
        if True:
            T.evaluate(0)
        else:
            T.evaluate(0)

    @T.prim_func(private=True, s_tir=True)
    def expected():
        T.evaluate(0)

    mod = tvm.IRModule.from_expr(before)
    mod = _apply_remove_no_op(mod)
    tvm.ir.assert_structural_equal(mod["main"], expected)


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