# 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: F821

import pytest

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

replace_by_tir_var = tvm.testing.parameter(
    by_dict={"replace-by-string": False, "replace-by-tirx-var": True}
)


def test_bind_static_value(replace_by_tir_var):
    """Symbolic vars may be replaced

    The replaced variables may be given either as strings, or as TIR variables
    """

    @R.function(private=True)
    def before(A: R.Tensor(("M", "K")), B: R.Tensor(("K", "N"))) -> R.Tensor(("M", "N")):
        return R.matmul(A, B)

    @R.function(private=True)
    def expected(A: R.Tensor((128, 64)), B: R.Tensor((64, 32))) -> R.Tensor((128, 32)):
        return R.matmul(A, B)

    if replace_by_tir_var:
        M, K = before.params[0].ty.shape
        _, N = before.params[1].ty.shape
        symbolic_var_map = {M: 128, K: 64, N: 32}
    else:
        symbolic_var_map = {"M": 128, "K": 64, "N": 32}

    after = before.bind_symbolic_vars(symbolic_var_map)
    tvm.ir.assert_structural_equal(expected, after)


def test_error_with_duplicate_var_names():
    """Duplicate variable names may not be replaced by string

    Two TIR variables may have the same name.  If two symbolic
    variables share the same name, the replacement map may not refer
    to that variable by string.
    """
    N1 = tvm.tirx.Var("N", "int64")
    N2 = tvm.tirx.Var("N", "int64")

    @R.function(private=True)
    def func(A: R.Tensor((N1, N1)), B: R.Tensor((N1, N2))) -> R.Tensor((N1, N2)):
        out: R.Tensor((N1, N2)) = R.matmul(A, B)
        return out

    with pytest.raises(RuntimeError):
        func.bind_symbolic_vars({"N": 64})


def test_string_var_when_other_var_has_duplicate_var_names():
    """Like test_error_with_duplicate_var_names, but replacing a different variable

    If two TIR variables share the same name, the restriction against
    replacing variables by name only applies to those duplicate names.
    Other variables may still be replaced by name.
    """
    N1 = tvm.tirx.Var("N", "int64")
    N2 = tvm.tirx.Var("N", "int64")
    BatchSize = tvm.tirx.Var("BatchSize", "int64")

    @R.function(private=True)
    def before(A: R.Tensor((BatchSize, N1, N1)), B: R.Tensor((N1, N2))) -> R.Tensor(
        (BatchSize, N1, N2)
    ):
        out: R.Tensor((BatchSize, N1, N2)) = R.matmul(A, B)
        return out

    @R.function(private=True)
    def expected(A: R.Tensor((16, N1, N1)), B: R.Tensor((N1, N2))) -> R.Tensor((16, N1, N2)):
        out: R.Tensor((16, N1, N2)) = R.matmul(A, B)
        return out

    after = before.bind_symbolic_vars({"BatchSize": 16})
    tvm.ir.assert_structural_equal(expected, after)


def test_error_with_nonexisting_var_name():
    """A string name of a symbolic var must be used by the function"""

    @R.function(private=True)
    def func(A: R.Tensor(("M", "N"))):
        return A

    with pytest.raises(RuntimeError):
        func.bind_symbolic_vars({"non_existing_symbolic_var": 64})


def test_error_with_nonexisting_tir_var():
    """A TIR symbolic var must be a symbolic var of the function"""

    @R.function(private=True)
    def func(A: R.Tensor(["M", "N"])):
        return A

    with pytest.raises(RuntimeError):
        func.bind_symbolic_vars({tvm.tirx.Var("M", "int64"): 64})


def test_error_with_multiple_definitions():
    """The string/TIR var syntaxes may not define the same variable"""

    @R.function(private=True)
    def func(A: R.Tensor(["M", "N"])):
        return A

    tir_var = func.params[0].ty.shape[0]
    symbolic_var_map = {tir_var: 0, "M": 0}

    with pytest.raises(RuntimeError):
        func.bind_symbolic_vars(symbolic_var_map)


def test_error_if_output_has_undefined():
    """The replacements may not introduce undefined symbolic vars"""

    @R.function(private=True)
    def func(A: R.Tensor(["M", "N"])):
        return A

    outside_var = tvm.tirx.Var("outside_var", "int64")

    with pytest.raises(RuntimeError):
        func.bind_symbolic_vars({"M": outside_var * 2})


def test_replacements_may_produce_new_symbolic_vars():
    """The output may introduce symbolic vars, but they must be bound"""

    @R.function(private=True)
    def before(A: R.Tensor(["M", "N"])):
        return A

    @R.function(private=True)
    def expected(A: R.Tensor(["outside_var * 2", "outside_var"])):
        return A

    outside_var = tvm.tirx.Var("outside_var", "int64")

    after = before.bind_symbolic_vars({"M": outside_var * 2, "N": outside_var})
    tvm.ir.assert_structural_equal(expected, after)


def test_bind_symbolic_vars_in_tensor_shape():
    """The bound variable should be replaced when appearing in type"""

    @R.function(private=True)
    def before(A: R.Tensor(["M", "N"])):
        M = T.int64()
        N = T.int64()
        B = R.call_dps_packed("dummy_func", [A], out_ty=R.Tensor([2 * M * N]))
        return B

    @R.function(private=True)
    def expected(A: R.Tensor(["M", 16])):
        M = T.int64()
        B = R.call_dps_packed("dummy_func", [A], out_ty=R.Tensor([M * 32]))
        return B

    after = before.bind_symbolic_vars({"N": 16})
    tvm.ir.assert_structural_equal(expected, after)


def test_bind_symbolic_vars_in_shape_expr():
    """The bound variable should be replaced when appearing in R.Shape"""

    @R.function(private=True)
    def before(A: R.Tensor(["M * N"]), x: R.Shape(["M", "N"])):
        M = T.int64()
        N = T.int64()
        B = R.call_dps_packed("dummy_func", [A], out_ty=R.Tensor([2 * M * N]))
        return B

    @R.function(private=True)
    def expected(A: R.Tensor(["M * 16"]), x: R.Shape(["M", 16])):
        B = R.call_dps_packed("dummy_func", [A], out_ty=R.Tensor([M * 32]))
        return B

    after = before.bind_symbolic_vars({"N": 16})
    tvm.ir.assert_structural_equal(expected, after)


def test_bind_strided_slice():
    """relax.op.strided_slice stores Expr attributes"""

    @R.function(private=True)
    def before(A: R.Tensor(["M", "N"])):
        N = T.int64()
        B = R.strided_slice(A, [1], [0], [N // 4])
        return B

    @R.function(private=True)
    def expected(A: R.Tensor(["M", 32])):
        # Binding substitutes runtime primitive arguments without applying
        # shape-only analyzer simplification to them.
        B = R.strided_slice(A, [1], [0], [T.FloorDiv(T.int64(32), T.int64(4))])
        return B

    after = before.bind_symbolic_vars({"N": 32})
    tvm.ir.assert_structural_equal(expected, after)


def test_bind_inside_match_cast():
    """Symbolic variables may occur within R.match_cast"""

    @R.function(private=True)
    def before(A: R.Tensor(["M", "N"]), B: R.Tensor(ndim=2)):
        M = T.int64()
        N = T.int64()
        C = R.match_cast(B, R.Tensor([M, N]))
        D = R.add(A, C)
        return D

    @R.function(private=True)
    def expected(A: R.Tensor(["M", 32]), B: R.Tensor(ndim=2)):
        M = T.int64()
        C = R.match_cast(B, R.Tensor([M, 32]))
        D = R.add(A, C)
        return D

    after = before.bind_symbolic_vars({"N": 32})
    tvm.ir.assert_structural_equal(expected, after)


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