# 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.
# pylint: disable=missing-docstring
# ruff: noqa: F401, F841

import pytest

import tvm
import tvm.testing
from tvm.s_tir import Schedule
from tvm.s_tir.meta_schedule.testing import te_workload
from tvm.s_tir.schedule.analysis import (
    TensorizeInfo,
    get_auto_tensorize_mapping_info,
    get_tensorize_loop_mapping,
    is_output_block,
    suggest_index_map,
)
from tvm.s_tir.tensor_intrin.cuda import (
    WMMA_SYNC_16x16x16_f16f16f16_INTRIN,
    WMMA_SYNC_16x16x16_f16f16f32_INTRIN,
)
from tvm.s_tir.tensor_intrin.x86 import dot_product_16x4_u8i8i32_desc
from tvm.script import tirx as T
from tvm.te import create_prim_func
from tvm.tirx import (
    Evaluate,
    For,
    ForKind,
    IndexMap,
    Var,
    decl_buffer,
    floordiv,
    floormod,
)
from tvm.tirx.analysis import expr_deep_equal
from tvm.tirx.function import TensorIntrin
from tvm.tirx.stmt_functor import pre_order_visit


def _make_vars(*args: str) -> list[Var]:
    return [Var(arg, dtype="int32") for arg in args]


def _make_loops(loop_vars: list[Var], extents: list[int]) -> list[For]:
    assert len(loop_vars) == len(extents)
    return [
        For(
            loop_var=loop_var,
            min=0,
            extent=extent,
            kind=ForKind.SERIAL,
            body=Evaluate(0),
        )
        for loop_var, extent in zip(loop_vars, extents)
    ]


def test_suggest_index_map_simple():
    i, j = _make_vars("i", "j")
    index_map = suggest_index_map(
        buffer=decl_buffer(shape=[8, 256]),
        indices=[
            floordiv(i, 16) * 4 + floordiv(j, 16),
            floormod(i, 16) * 16 + floormod(j, 16),
        ],
        loops=_make_loops(
            loop_vars=[i, j],
            extents=[32, 64],
        ),
        predicate=True,
    )
    expected_index_map = IndexMap.from_func(
        lambda x, y: [
            floordiv(x, 4),
            floordiv(y, 16),
            floormod(x, 4),
            floormod(y, 16),
        ],
    )
    assert index_map.is_equivalent_to(expected_index_map)


def test_suggest_index_map_bijective():
    i, j = _make_vars("i", "j")
    index_map = suggest_index_map(
        buffer=decl_buffer(shape=[8]),
        indices=[floormod(j, 4) * 2 + i],
        loops=_make_loops(
            loop_vars=[i, j],
            extents=[2, 32],
        ),
        predicate=True,
    )
    expected_index_map = IndexMap.from_func(
        lambda x: [
            floormod(x, 2),
            floordiv(x, 2),
        ],
    )
    assert index_map.is_equivalent_to(expected_index_map)


def test_suggest_index_map_winograd():
    """use case in winograd conv where the indices are complicated"""
    fused_outer, i3_3_fused, i4_0, i4_1 = _make_vars("fused_outer", "i3_3_fused", "i4_0", "i4_1")
    eps = floordiv(fused_outer, 336) * 2 + floordiv(floormod(fused_outer, 16), 8)
    nu = floordiv(floormod(fused_outer, 336), 112) * 2 + floordiv(floormod(fused_outer, 8), 4)
    co = floormod(fused_outer, 4) * 32 + i3_3_fused
    ci = (i4_0 * 32) + i4_1
    buffer = decl_buffer(shape=[6, 6, 128, 128])
    index_map = suggest_index_map(
        buffer=buffer,
        indices=[eps, nu, co, ci],
        loops=_make_loops(
            loop_vars=[fused_outer, i3_3_fused, i4_0, i4_1],
            extents=[1008, 32, 4, 32],
        ),
        predicate=True,
    )
    expected_index_map = IndexMap.from_func(
        lambda i0, i1, i2, i3: (
            floordiv(i0, 2),
            floordiv(i1, 2),
            floormod(i0, 2),
            floormod(i1, 2) * 4 + floordiv(i2, 32),
            floormod(i2, 32),
            floordiv(i3, 32),
            floormod(i3, 32),
        )
    )
    assert index_map.is_equivalent_to(expected_index_map)
    inverse_index_map = index_map.inverse(buffer.shape)
    expected_inverse_index_map = IndexMap.from_func(
        lambda i0, i1, i2, i3, i4, i5, i6: (
            ((i0 * 2) + i2),
            i1 * 2 + floordiv(i3, 4),
            floormod(i3, 4) * 32 + i4,
            ((i5 * 32) + i6),
        )
    )
    assert inverse_index_map.is_equivalent_to(expected_inverse_index_map)


@tvm.script.ir_module
class DenseTIRModule:
    @T.prim_func(s_tir=True)
    def main(
        placeholder: T.Buffer((1024, 1024), "uint8"),
        placeholder_1: T.Buffer((64, 256, 16, 4), "int8"),
        compute: T.Buffer((1024, 1024), "int32"),
    ) -> None:
        T.func_attr({"global_symbol": "main", "tirx.noalias": True})
        with T.sblock("root"):
            T.reads()
            T.writes()
            for i0, i1, i2 in T.grid(1024, 1024, 1024):
                with T.sblock("compute"):
                    i, j, k = T.axis.remap("SSR", [i0, i1, i2])
                    T.reads(placeholder[i, k], placeholder_1[j // 16, k // 4, j % 16, k % 4])
                    T.writes(compute[i, j])
                    with T.init():
                        compute[i, j] = 0
                    compute[i, j] = compute[i, j] + T.cast(placeholder[i, k], "int32") * T.cast(
                        placeholder_1[j // 16, k // 4, j % 16, k % 4], "int32"
                    )


@tvm.script.ir_module
class Conv2dNCHWcTIRModule:
    @T.prim_func(s_tir=True)
    def main(
        placeholder: T.Buffer((1, 4, 56, 56, 16), "uint8"),
        placeholder_1: T.Buffer((16, 4, 1, 1, 4, 16, 4), "int8"),
        conv2d_NCHWc_int8: T.Buffer((1, 16, 56, 56, 16), "int32"),
    ) -> None:
        T.func_attr({"global_symbol": "main", "tirx.noalias": True})
        for i0, i1, i2, i3, i4, i5, i6, i7, i8, i9 in T.grid(1, 16, 56, 56, 16, 1, 1, 4, 4, 4):
            with T.sblock("conv2d_NCHWc_int8"):
                (
                    n,
                    oc_chunk,
                    oh,
                    ow,
                    oc_block,
                    kh,
                    kw,
                    ic_outer,
                    ic_f_inner,
                    ic_s_inner,
                ) = T.axis.remap("SSSSSRRRRR", [i0, i1, i2, i3, i4, i5, i6, i7, i8, i9])
                T.reads(
                    placeholder[n, ic_outer, oh + kh, ow + kw, ic_f_inner * 4 + ic_s_inner],
                    placeholder_1[oc_chunk, ic_outer, kh, kw, ic_f_inner, oc_block, ic_s_inner],
                )
                T.writes(conv2d_NCHWc_int8[n, oc_chunk, oh, ow, oc_block])
                with T.init():
                    conv2d_NCHWc_int8[n, oc_chunk, oh, ow, oc_block] = 0
                conv2d_NCHWc_int8[n, oc_chunk, oh, ow, oc_block] = conv2d_NCHWc_int8[
                    n, oc_chunk, oh, ow, oc_block
                ] + T.cast(
                    placeholder[n, ic_outer, oh + kh, ow + kw, ic_f_inner * 4 + ic_s_inner],
                    "int32",
                ) * T.cast(
                    placeholder_1[oc_chunk, ic_outer, kh, kw, ic_f_inner, oc_block, ic_s_inner],
                    "int32",
                )


def collect_loops(prim_func):
    loops = []

    def callback(node):
        if isinstance(node, tvm.tirx.For):
            loops.append(node)
        return True

    pre_order_visit(prim_func.body, callback)

    return loops


def test_get_tensorize_loop_mapping_dense_16x4():
    s = Schedule(DenseTIRModule)
    block = s.get_sblock("compute")

    info = get_tensorize_loop_mapping(s, block, dot_product_16x4_u8i8i32_desc)

    assert isinstance(info, TensorizeInfo)

    desc_loop_to_sref = dict((v, k) for k, v in info.loop_map.items())

    desc_loops = collect_loops(dot_product_16x4_u8i8i32_desc)
    _, loop_j, loop_k = s.get_loops(block)

    assert desc_loops[0] in desc_loop_to_sref and desc_loops[1] in desc_loop_to_sref
    assert s.get(desc_loop_to_sref[desc_loops[0]]) == s.get(loop_j)
    assert s.get(desc_loop_to_sref[desc_loops[1]]) == s.get(loop_k)


def test_get_tensorize_loop_mapping_conv2d_nchwc_16x4():
    s = Schedule(Conv2dNCHWcTIRModule)
    block = s.get_sblock("conv2d_NCHWc_int8")

    info = get_tensorize_loop_mapping(s, block, dot_product_16x4_u8i8i32_desc)

    desc_loop_to_sref = dict((v, k) for k, v in info.loop_map.items())

    desc_loops = collect_loops(dot_product_16x4_u8i8i32_desc)

    # i4 corresonds to the inner output channel axis of the NCHWc output tensor
    # for i0, i1, i2, i3, i4, i5, i6, i7, i8, i9 in T.grid(1, 16, 56, 56, 16, 1, 1, 4, 4, 4):
    _, _, _, _, i4, _, _, _, _, i9 = s.get_loops(block)

    assert desc_loops[0] in desc_loop_to_sref and desc_loops[1] in desc_loop_to_sref
    assert s.get(desc_loop_to_sref[desc_loops[0]]) == s.get(i4)
    assert s.get(desc_loop_to_sref[desc_loops[1]]) == s.get(i9)


def test_get_tensorize_loop_mapping_matmul_mma():
    @T.prim_func(s_tir=True)
    def matmul_16x16x16xf16f16f16_desc(
        A: T.Buffer((16, 16), "float16", align=64, offset_factor=1),
        B: T.Buffer((16, 16), "float16", align=64, offset_factor=1),
        C: T.Buffer((16, 16), "float16", align=64, offset_factor=1),
    ) -> None:
        with T.sblock("root"):
            T.reads(C[0:16, 0:16], A[0:16, 0:16], B[0:16, 0:16])
            T.writes(C[0:16, 0:16])
            for i, j, k in T.grid(16, 16, 16):
                with T.sblock("update"):
                    vii, vjj, vkk = T.axis.remap("SSR", [i, j, k])
                    C[vii, vjj] = C[vii, vjj] + A[vii, vkk] * B[vjj, vkk]

    matmul = create_prim_func(
        te_workload.matmul_relu(
            n=512,
            m=512,
            k=512,
        )
    )

    s = Schedule(matmul)
    block = s.get_sblock("C")
    i0, i1, i2 = s.get_loops(block)
    desc_loops = collect_loops(matmul_16x16x16xf16f16f16_desc)

    for do_reorder in [False, True]:
        # Mapping should be invariant to the loop permutation
        if do_reorder:
            s.reorder(i2, i0, i1)

        info = get_tensorize_loop_mapping(s, block, matmul_16x16x16xf16f16f16_desc)
        assert info is not None
        desc_loop_to_sref = dict((v, k) for k, v in info.loop_map.items())

        for i in range(3):
            assert desc_loops[i] in desc_loop_to_sref

        assert s.get(desc_loop_to_sref[desc_loops[0]]) == s.get(i0)
        assert s.get(desc_loop_to_sref[desc_loops[1]]) == s.get(i1)
        assert s.get(desc_loop_to_sref[desc_loops[2]]) == s.get(i2)


def test_get_tensorize_loop_mapping_padding_matmul():
    matmul = create_prim_func(
        te_workload.matmul_relu(
            n=127,
            m=256,
            k=65,
            in_dtype="float16",
            out_dtype="float16",
        )
    )
    s = Schedule(matmul)
    block = s.get_sblock("C")

    desc = TensorIntrin.get(WMMA_SYNC_16x16x16_f16f16f16_INTRIN).desc
    info = get_tensorize_loop_mapping(s, block, desc, allow_padding=True)
    assert info is not None
    expected_padding = [16, 1, 16]
    actual_padding = info.block_iter_paddings
    assert actual_padding is not None
    assert len(actual_padding) == len(expected_padding)
    for actual, expected in zip(actual_padding, expected_padding):
        assert actual == expected


def check_index_map(workload, block_name, intrin_name, expected_index_map):
    s = Schedule(workload)
    block = s.get_sblock(block_name)
    desc_func = TensorIntrin.get(intrin_name).desc
    info = get_auto_tensorize_mapping_info(s, block, desc_func)
    if expected_index_map is None:
        assert info is None
        return
    assert len(info.mappings) == 1
    assert IndexMap.from_func(expected_index_map).is_equivalent_to(info.mappings[0])


def test_get_auto_tensorize_mapping_info_conv2d():
    conv2d = create_prim_func(
        te_workload.conv2d_nhwc(4, 16, 16, 64, 64, 3, 1, 1, in_dtype="float16", out_dtype="float32")
    )
    check_index_map(
        conv2d,
        "conv2d_nhwc",
        WMMA_SYNC_16x16x16_f16f16f32_INTRIN,
        lambda n, h, w, c, rh, rw, rc: (n * 256 + h * 16 + w, c, rh * 192 + rw * 64 + rc),
    )


def test_get_auto_tensorize_mapping_info_conv2d_unit_batch():
    conv2d = create_prim_func(
        te_workload.conv2d_nhwc(1, 16, 16, 64, 64, 3, 1, 1, in_dtype="float16", out_dtype="float32")
    )
    check_index_map(
        conv2d,
        "conv2d_nhwc",
        WMMA_SYNC_16x16x16_f16f16f32_INTRIN,
        lambda n, h, w, c, rh, rw, rc: (n * 256 + h * 16 + w, c, rh * 192 + rw * 64 + rc),
    )


@pytest.mark.parametrize("b,m,n,k", [(1, 512, 512, 512), (16, 32, 32, 32)])
def test_get_auto_tensorize_mapping_info_batch_matmul(b, m, n, k):
    matmul = create_prim_func(
        te_workload.batch_matmul_nkkm(b, m, n, k, in_dtype="float16", out_dtype="float32")
    )
    check_index_map(
        matmul, "Z", WMMA_SYNC_16x16x16_f16f16f32_INTRIN, lambda b, m, n, k: (b, m, n, k)
    )


@pytest.mark.parametrize(
    "n,m,k,expected",
    [
        (
            512,
            512,
            512,
            lambda n, m, k: (
                n,
                m,
                k,
            ),
        ),
        (1, 32, 32, lambda n, m, k: (n, m, k)),
    ],
)
def test_get_auto_tensorize_mapping_info_matmul(n, m, k, expected):
    matmul = create_prim_func(te_workload.matmul(n, m, k, in_dtype="float16", out_dtype="float32"))
    check_index_map(matmul, "C", WMMA_SYNC_16x16x16_f16f16f32_INTRIN, expected)


def test_is_output_block():
    @T.prim_func(s_tir=True)
    def two_elementwise(a: T.handle, c: T.handle) -> None:
        A = T.match_buffer(a, (128, 128), "float32")
        B = T.sblock_alloc_buffer((128, 128), "float32")
        C = T.match_buffer(c, (128, 128), "float32")
        for i, j in T.grid(128, 128):
            with T.sblock("B"):
                vi, vj = T.axis.remap("SS", [i, j])
                B[vi, vj] = A[vi, vj] * 2.0
        for i, j in T.grid(128, 128):
            with T.sblock("C"):
                vi, vj = T.axis.remap("SS", [i, j])
                C[vi, vj] = B[vi, vj] + 1.0

    sch = tvm.s_tir.Schedule(two_elementwise)
    block_rv = sch.get_sblock("C")
    assert is_output_block(sch, block_rv)


def test_empty_grid():
    @T.prim_func(s_tir=True)
    def foo(out: T.Buffer((T.int64(1), T.int64(8), T.int64(8)), "int32")):
        act = T.sblock_alloc_buffer((1, 8, 8), "int32")
        for z2, y2, x2 in T.grid(1, 8, 8):
            with T.sblock("b0"):
                az, ay, ax = T.axis.remap("SSS", [z2, y2, x2])
                T.writes(act[az, ay, ax])
                act[az, ay, az] = T.int32(0)
        # Empty grid:
        for z1, y1, x1 in T.grid(0, 8, 8):
            with T.sblock("b1"):
                az, ay, ax = T.axis.remap("SSS", [z1, y1, x1])
                T.reads(act[az + 1, ay, ax])
                T.writes(out[az, ay, ax])
                out[az, ay, ax] = act[az + 1, ay, ax]
        # The block below is not needed to show the bug, but the 'out'
        # buffer would be undefined without it.
        for z2, y2, x2 in T.grid(1, 8, 8):
            with T.sblock("b2"):
                az, ay, ax = T.axis.remap("SSS", [z2, y2, x2])
                T.writes(out[az, ay, ax])
                out[az, ay, az] = T.int32(0)

    # This caused a crash before.
    sch = tvm.s_tir.Schedule(foo)


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