# 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-function-docstring,missing-module-docstring
# ruff: noqa: E501

import pytest

import tvm
import tvm.testing
from tvm import te, tirx
from tvm.s_tir.schedule.testing import (
    assert_structural_equal_ignore_global_symbol,
    verify_trace_roundtrip,
)
from tvm.s_tir.tensor_intrin.arm_cpu import (
    DP4A_S8S8S32_INTRIN,
    DP4A_S8U8S32_INTRIN,
    DP4A_U8S8S32_INTRIN,
    DP4A_U8U8U32_INTRIN,
    ARM_DOT_4x4_i8_NEON_INTRIN,
    ARM_DOT_4x4_i8_SDOT_INTRIN,
)
from tvm.s_tir.tensor_intrin.hexagon import VDMPY_i16i16i32_INTRIN, VRMPY_u8u8i32_INTRIN
from tvm.s_tir.tensor_intrin.rocm import AMDGPU_SDOT4_INTRIN
from tvm.s_tir.tensor_intrin.x86 import AVX512_DOT_16x4_INTRIN, VNNI_DOT_16x4_INTRIN
from tvm.script import tirx as T

# fmt: off
# pylint: disable=no-member,invalid-name,unused-variable,line-too-long,redefined-outer-name,unexpected-keyword-arg,too-many-nested-blocks

@T.prim_func(s_tir=True)
def mma_desc(a: T.handle, b: T.handle, c: T.handle) -> None:
    A = T.match_buffer(a, (16, 16), align=64, offset_factor=1)
    B = T.match_buffer(b, (16, 16), align=64, offset_factor=1)
    C = T.match_buffer(c, (16, 16), align=64, offset_factor=1)

    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]


@T.prim_func(s_tir=True)
def mma_intrin(a: T.handle, b: T.handle, c: T.handle) -> None:
    A = T.match_buffer(a, (16, 16), align=64, offset_factor=1)
    B = T.match_buffer(b, (16, 16), align=64, offset_factor=1)
    C = T.match_buffer(c, (16, 16), align=64, offset_factor=1)

    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])
        T.evaluate(
            T.tvm_mma_sync(
                C.data,
                C.elem_offset // 256,
                A.data,
                A.elem_offset // 256,
                B.data,
                B.elem_offset // 256,
                C.data,
                C.elem_offset // 256,
                dtype="handle",
            )
        )


@T.prim_func(s_tir=True)
def dot_product_desc(a: T.handle, b: T.handle, c: T.handle) -> None:
    A = T.match_buffer(a, (4,))
    B = T.match_buffer(b, (4,))
    C = T.match_buffer(c, ())

    with T.sblock("root"):
        T.reads(C[()], A[0 : 4], B[0 : 4])
        T.writes(C[()])
        for i in range(0, 4):
            with T.sblock("update"):
                vi = T.axis.remap("R", [i])
                C[()] = C[()] + A[vi] * B[vi]


@T.prim_func(s_tir=True)
def dot_product_intrin(a: T.handle, b: T.handle, c: T.handle) -> None:
    A = T.match_buffer(a, (4,), offset_factor=1)
    B = T.match_buffer(b, (4,), offset_factor=1)
    C = T.match_buffer(c, (), offset_factor=1)

    with T.sblock("root"):
        T.reads(C[()], A[0 : 4], B[0 : 4])
        T.writes(C[()])
        T.evaluate(
            T.call_extern(
                "vec4add",
                C.data,
                C.elem_offset,
                A.data,
                A.elem_offset,
                B.data,
                B.elem_offset,
                dtype="int32",
            )
        )


@T.prim_func(s_tir=True)
def dot_product_intrin_annotated(a: T.handle, b: T.handle, c: T.handle) -> None:
    A = T.match_buffer(a, (4,), offset_factor=1)
    B = T.match_buffer(b, (4,), offset_factor=1)
    C = T.match_buffer(c, (), offset_factor=1)

    with T.sblock("root"):
        T.reads(C[()], A[0 : 4], B[0 : 4])
        T.writes(C[()])
        T.sblock_attr({"test_annotation": True})
        T.evaluate(
            T.call_extern(
                "vec4add",
                C.data,
                C.elem_offset,
                A.data,
                A.elem_offset,
                B.data,
                B.elem_offset,
                dtype="int32",
            )
        )


@T.prim_func(s_tir=True)
def outer_product_desc(a: T.handle, b: T.handle, c: T.handle) -> None:
    A = T.match_buffer(a, (16, 1), offset_factor=1)
    B = T.match_buffer(b, (16, 1), offset_factor=1)
    C = T.match_buffer(c, (16, 16), offset_factor=1)

    with T.sblock("root"):
        T.reads(
            C[0 : 16, 0 : 16],
            A[0 : 16, 0 : 1],
            B[0 : 16, 0 : 1],
        )
        T.writes(C[0 : 16, 0 : 16])
        for i, j in T.grid(16, 16):
            with T.sblock("update"):
                vii, vjj = T.axis.remap("SS", [i, j])
                C[vii, vjj] = C[vii, vjj] + A[vii, 0] * B[vjj, 0]


@T.prim_func(s_tir=True)
def outer_product_intrin(a: T.handle, b: T.handle, c: T.handle) -> None:
    A = T.match_buffer(a, (16, 1), offset_factor=1)
    B = T.match_buffer(b, (16, 1), offset_factor=1)
    C = T.match_buffer(c, (16, 16), offset_factor=1)

    with T.sblock("root"):
        T.reads(
            C[0 : 16, 0 : 16],
            A[0 : 16, 0 : 1],
            B[0 : 16, 0 : 1],
        )
        T.writes(C[0 : 16, 0 : 16])
        T.evaluate(
            T.call_extern(
                "outer_product",
                C.data,
                C.elem_offset,
                A.data,
                A.elem_offset,
                B.data,
                B.elem_offset,
                dtype="int32",
            )
        )


@T.prim_func(s_tir=True)
def matmul(
    A: T.Buffer((128, 128), "float32"),
    B: T.Buffer((128, 128), "float32"),
    C: T.Buffer((128, 128), "float32"),
) -> None:
    for i, j, k in T.grid(128, 128, 128):
        with T.sblock("update"):
            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
            with T.init():
                C[vi, vj] = T.float32(0)
            C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]


@T.prim_func(s_tir=True)
def tensorized_matmul(a: T.handle, b: T.handle, c: T.handle) -> None:
    C = T.match_buffer(c, [128, 128], elem_offset=0, align=64, offset_factor=1)
    B = T.match_buffer(b, [128, 128], elem_offset=0, align=64, offset_factor=1)
    A = T.match_buffer(a, [128, 128], elem_offset=0, align=64, offset_factor=1)

    for i_outer, j_outer in T.grid(8, 8):
        for i_inner_init, j_inner_init in T.grid(16, 16):
            with T.sblock("init"):
                vi_init = T.axis.S(128, ((i_outer * 16) + i_inner_init))
                vj_init = T.axis.S(128, ((j_outer * 16) + j_inner_init))
                C[vi_init, vj_init] = T.float32(0)
        for k_outer in T.grid(8):
            with T.sblock("update"):
                vi, vj, vk = T.axis.remap("SSR", [i_outer, j_outer, k_outer])
                T.reads(
                    [
                        C[vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16],
                        A[vi * 16 : vi * 16 + 16, vk * 16 : vk * 16 + 16],
                        B[vj * 16 : vj * 16 + 16, vk * 16 : vk * 16 + 16],
                    ]
                )
                T.writes(C[vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16])
                A_elem_offset = T.int32()
                B_elem_offset = T.int32()
                C_elem_offset = T.int32()
                A_sub = T.match_buffer(
                    A[vi * 16 : vi * 16 + 16, vk * 16 : vk * 16 + 16],
                    [16, 16],
                    elem_offset=A_elem_offset,
                )
                B_sub = T.match_buffer(
                    B[vj * 16 : vj * 16 + 16, vk * 16 : vk * 16 + 16],
                    [16, 16],
                    elem_offset=B_elem_offset,
                )
                C_sub = T.match_buffer(
                    C[vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16],
                    [16, 16],
                    elem_offset=C_elem_offset,
                )
                T.evaluate(
                    T.tvm_mma_sync(
                        C_sub.data,
                        T.floordiv(C_sub.elem_offset, 256),
                        A_sub.data,
                        T.floordiv(A_sub.elem_offset, 256),
                        B_sub.data,
                        T.floordiv(B_sub.elem_offset, 256),
                        C_sub.data,
                        T.floordiv(C_sub.elem_offset, 256),
                        dtype="handle",
                    )
                )


@T.prim_func(s_tir=True)
def batch_matmul(
    A: T.Buffer((16, 128, 128), "float32"),
    B: T.Buffer((16, 128, 128), "float32"),
    C: T.Buffer((16, 128, 128), "float32"),
) -> None:
    for n, i, j in T.grid(16, 128, 128):
        with T.sblock("init"):
            vn, vi, vj = T.axis.remap("SSS", [n, i, j])
            C[vn, vi, vj] = T.float32(0)

    for n, i, j, k in T.grid(16, 128, 128, 128):
        with T.sblock("update"):
            vn, vi, vj, vk = T.axis.remap("SSSR", [n, i, j, k])
            C[vn, vi, vj] = C[vn, vi, vj] + A[vn, vi, vk] * B[vn, vj, vk]


@T.prim_func(s_tir=True)
def tensorized_batch_matmul_mma(
    A: T.Buffer((16, 128, 128), "float32"),
    B: T.Buffer((16, 128, 128), "float32"),
    C: T.Buffer((16, 128, 128), "float32"),
) -> None:
    for n, i, j in T.grid(16, 128, 128):
        with T.sblock("init"):
            vn, vi, vj = T.axis.remap("SSS", [n, i, j])
            T.reads()
            T.writes(C[vn, vi, vj])
            C[vn, vi, vj] = T.float32(0)
    for n in range(0, 16):
        for i, j, k in T.grid(8, 8, 8):
            with T.sblock("update"):
                vn, vi, vj, vk = T.axis.remap("SSSR", [n, i, j, k])
                T.reads(
                    C[vn : vn + 1, vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16],
                    A[vn : vn + 1, vi * 16 : vi * 16 + 16, vk * 16 : vk * 16 + 16],
                    B[vn : vn + 1, vj * 16 : vj * 16 + 16, vk * 16 : vk * 16 + 16],
                )
                T.writes(C[vn : vn + 1, vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16])
                A_elem_offset = T.int32()
                B_elem_offset = T.int32()
                C_elem_offset = T.int32()
                A_sub = T.match_buffer(
                    A[vn : vn + 1, vi * 16 : vi * 16 + 16, vk * 16 : vk * 16 + 16],
                    (16, 16),
                    elem_offset=A_elem_offset,
                )
                B_sub = T.match_buffer(
                    B[vn : vn + 1, vj * 16 : vj * 16 + 16, vk * 16 : vk * 16 + 16],
                    (16, 16),
                    elem_offset=B_elem_offset,
                )
                C_sub = T.match_buffer(
                    C[vn : vn + 1, vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16],
                    (16, 16),
                    elem_offset=C_elem_offset,
                )
                T.evaluate(
                    T.tvm_mma_sync(
                        C_sub.data,
                        T.floordiv(C_sub.elem_offset, 256),
                        A_sub.data,
                        T.floordiv(A_sub.elem_offset, 256),
                        B_sub.data,
                        T.floordiv(B_sub.elem_offset, 256),
                        C_sub.data,
                        T.floordiv(C_sub.elem_offset, 256),
                        dtype="handle",
                    )
                )


@T.prim_func(s_tir=True)
def tensorized_batch_matmul_dot_product(
    A: T.Buffer((16, 128, 128), "float32"),
    B: T.Buffer((16, 128, 128), "float32"),
    C: T.Buffer((16, 128, 128), "float32"),
) -> None:
    for n, i, j in T.grid(16, 128, 128):
        with T.sblock("init"):
            vn, vi, vj = T.axis.remap("SSS", [n, i, j])
            T.reads()
            T.writes(C[vn, vi, vj])
            C[vn, vi, vj] = T.float32(0)
    for n, i, j, k_0 in T.grid(16, 128, 128, 32):
        with T.sblock("blockized_update"):
            vn, vi, vj, vko = T.axis.remap("SSSR", [n, i, j, k_0])
            T.reads(
                C[vn, vi, vj], A[vn, vi, vko * 4 : vko * 4 + 4], B[vn, vj, vko * 4 : vko * 4 + 4]
            )
            T.writes(C[vn, vi, vj])
            A_1 = T.match_buffer(
                A[vn, vi, vko * 4 : vko * 4 + 4], [4], dtype="float32", offset_factor=1
            )
            B_1 = T.match_buffer(
                B[vn, vj, vko * 4 : vko * 4 + 4], [4], dtype="float32", offset_factor=1
            )
            C_1 = T.match_buffer(C[vn, vi, vj], [], dtype="float32", offset_factor=1)
            T.evaluate(
                T.call_extern(
                    "vec4add",
                    C_1.data,
                    C_1.elem_offset,
                    A_1.data,
                    A_1.elem_offset,
                    B_1.data,
                    B_1.elem_offset,
                    dtype="int32",
                )
            )


@T.prim_func(s_tir=True)
def tensorized_batch_matmul_outer_product(
    A: T.Buffer((16, 128, 128), "float32"),
    B: T.Buffer((16, 128, 128), "float32"),
    C: T.Buffer((16, 128, 128), "float32"),
) -> None:
    for n, i, j in T.grid(16, 128, 128):
        with T.sblock("init"):
            vn, vi, vj = T.axis.remap("SSS", [n, i, j])
            T.reads()
            T.writes(C[vn, vi, vj])
            C[vn, vi, vj] = T.float32(0)
    for n, i_0, j_0, k in T.grid(16, 8, 8, 128):
        with T.sblock("blockized_update"):
            vn, vio, vjo, vk = T.axis.remap("SSSR", [n, i_0, j_0, k])
            T.reads(
                C[vn, vio * 16 : vio * 16 + 16, vjo * 16 : vjo * 16 + 16],
                A[vn, vio * 16 : vio * 16 + 16, vk],
                B[vn, vjo * 16 : vjo * 16 + 16, vk],
            )
            T.writes(C[vn, vio * 16 : vio * 16 + 16, vjo * 16 : vjo * 16 + 16])
            A_1 = T.match_buffer(A[vn, vio * 16 : vio * 16 + 16, vk], [16, 1], dtype="float32", offset_factor=1)
            B_1 = T.match_buffer(B[vn, vjo * 16 : vjo * 16 + 16, vk], [16, 1], dtype="float32", offset_factor=1
            )
            C_1 = T.match_buffer(
                C[vn, vio * 16 : vio * 16 + 16, vjo * 16 : vjo * 16 + 16], [16, 16], dtype="float32", offset_factor=1
            )
            T.evaluate(
                T.call_extern("outer_product", C_1.data, C_1.elem_offset, A_1.data, A_1.elem_offset,
                              B_1.data, B_1.elem_offset, dtype="int32"
                )
            )


@T.prim_func(s_tir=True)
def annotated_mma_desc(a: T.handle, b: T.handle, c: T.handle) -> None:
    A = T.match_buffer(a, (16, 16), align=64, offset_factor=1)
    B = T.match_buffer(b, (16, 16), align=64, offset_factor=1)
    C = T.match_buffer(c, (16, 16), align=64, offset_factor=1)

    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"):
                T.sblock_attr({"test_annotation": True})
                vii, vjj, vkk = T.axis.remap("SSR", [i, j, k])
                C[vii, vjj] = C[vii, vjj] + A[vii, vkk] * B[vjj, vkk]


@T.prim_func(s_tir=True)
def annotated_matmul(
    A: T.Buffer((128, 128), "float32"),
    B: T.Buffer((128, 128), "float32"),
    C: T.Buffer((128, 128), "float32"),
) -> None:
    for i, j, k in T.grid(128, 128, 128):
        with T.sblock("update"):
            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
            T.sblock_attr({"test_annotation": True})
            with T.init():
                C[vi, vj] = T.float32(0)
            C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]


@T.prim_func(s_tir=True)
def annotated_tensorized_matmul(a: T.handle, b: T.handle, c: T.handle) -> None:
    C = T.match_buffer(c, [128, 128], elem_offset=0, align=64, offset_factor=1)
    B = T.match_buffer(b, [128, 128], elem_offset=0, align=64, offset_factor=1)
    A = T.match_buffer(a, [128, 128], elem_offset=0, align=64, offset_factor=1)

    for i_outer, j_outer in T.grid(8, 8):
        for i_inner_init, j_inner_init in T.grid(16, 16):
            with T.sblock("init"):
                vi_init = T.axis.S(128, ((i_outer * 16) + i_inner_init))
                vj_init = T.axis.S(128, ((j_outer * 16) + j_inner_init))
                T.sblock_attr({"test_annotation": True})
                C[vi_init, vj_init] = T.float32(0)
        for k_outer in T.grid(8):
            with T.sblock("update"):
                vi, vj, vk = T.axis.remap("SSR", [i_outer, j_outer, k_outer])
                T.reads(
                    [
                        C[vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16],
                        A[vi * 16 : vi * 16 + 16, vk * 16 : vk * 16 + 16],
                        B[vj * 16 : vj * 16 + 16, vk * 16 : vk * 16 + 16],
                    ]
                )
                T.writes(C[vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16])
                A_elem_offset = T.int32()
                B_elem_offset = T.int32()
                C_elem_offset = T.int32()
                A_sub = T.match_buffer(
                    A[vi * 16 : vi * 16 + 16, vk * 16 : vk * 16 + 16],
                    [16, 16],
                    elem_offset=A_elem_offset,
                )
                B_sub = T.match_buffer(
                    B[vj * 16 : vj * 16 + 16, vk * 16 : vk * 16 + 16],
                    [16, 16],
                    elem_offset=B_elem_offset,
                )
                C_sub = T.match_buffer(
                    C[vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16],
                    [16, 16],
                    elem_offset=C_elem_offset,
                )
                T.evaluate(
                    T.tvm_mma_sync(
                        C_sub.data,
                        T.floordiv(C_sub.elem_offset, 256),
                        A_sub.data,
                        T.floordiv(A_sub.elem_offset, 256),
                        B_sub.data,
                        T.floordiv(B_sub.elem_offset, 256),
                        C_sub.data,
                        T.floordiv(C_sub.elem_offset, 256),
                        dtype="handle",
                    )
                )


# fmt: off
# pylint: disable=no-member,invalid-name,unused-variable,line-too-long,redefined-outer-name,unexpected-keyword-arg,too-many-nested-blocks

tirx.TensorIntrin.register("test_mma_intrin", mma_desc, mma_intrin)
tirx.TensorIntrin.register("test_annotated_mma_intrin", annotated_mma_desc, mma_intrin)
tirx.TensorIntrin.register("test_dot_product_intrin", dot_product_desc, dot_product_intrin)
tirx.TensorIntrin.register("test_outer_product_intrin", outer_product_desc, outer_product_intrin)
tirx.TensorIntrin.register("test_dot_product_intrin_annotated", dot_product_desc, dot_product_intrin_annotated)


def test_tensorize_matmul():
    func = matmul
    # schedule
    s = tvm.s_tir.Schedule(func, debug_mask="all")
    update = s.get_sblock("update")
    i, j, k = s.get_loops(update)
    io, ii = s.split(i, factors=[None, 16])
    jo, ji = s.split(j, factors=[None, 16])
    ko, ki = s.split(k, factors=[None, 16])
    s.reorder(io, jo, ko, ii, ji, ki)
    s.decompose_reduction(update, ko)
    s.tensorize(ii, "test_mma_intrin")
    assert_structural_equal_ignore_global_symbol(tensorized_matmul, s.mod["main"])
    verify_trace_roundtrip(sch=s, mod=func)


def test_tensorize_batch_matmul():
    func = batch_matmul
    s = tvm.s_tir.Schedule(func, debug_mask="all")
    update = s.get_sblock("update")
    _, i, j, k = s.get_loops(update)
    io, ii = s.split(i, factors=[None, 16])
    jo, ji = s.split(j, factors=[None, 16])
    ko, ki = s.split(k, factors=[None, 16])
    s.reorder(io, jo, ko, ii, ji, ki)
    s.tensorize(ii, "test_mma_intrin")
    assert_structural_equal_ignore_global_symbol(tensorized_batch_matmul_mma, s.mod["main"])
    verify_trace_roundtrip(sch=s, mod=batch_matmul)


def test_tensorize_dot_product():
    func = batch_matmul
    s = tvm.s_tir.Schedule(func, debug_mask="all")
    C = s.get_sblock("update")
    _, _, _, k = s.get_loops(C)
    _, ki = s.split(k, factors=[None, 4])
    s.tensorize(ki, "test_dot_product_intrin")
    assert_structural_equal_ignore_global_symbol(tensorized_batch_matmul_dot_product, s.mod["main"])
    verify_trace_roundtrip(sch=s, mod=func)


def test_tensorize_outer_product():
    func = batch_matmul
    s = tvm.s_tir.Schedule(func, debug_mask="all")
    C = s.get_sblock("update")
    _, i, j, k = s.get_loops(C)
    io, ii = s.split(i, factors=[None, 16])
    jo, ji = s.split(j, factors=[None, 16])
    s.reorder(io, jo, k, ii, ji)
    s.tensorize(ii, "test_outer_product_intrin")
    assert_structural_equal_ignore_global_symbol(tensorized_batch_matmul_outer_product, s.mod["main"])
    verify_trace_roundtrip(sch=s, mod=func)


def test_tensorize_with_annotation():
    func = annotated_matmul
    s = tvm.s_tir.Schedule(func, debug_mask="all")
    update = s.get_sblock("update")
    i, j, k = s.get_loops(update)
    io, ii = s.split(i, factors=[None, 16])
    jo, ji = s.split(j, factors=[None, 16])
    ko, ki = s.split(k, factors=[None, 16])
    s.reorder(io, jo, ko, ii, ji, ki)
    s.decompose_reduction(update, ko)
    s.tensorize(ii, "test_annotated_mma_intrin")
    assert_structural_equal_ignore_global_symbol(annotated_tensorized_matmul, s.mod["main"])
    verify_trace_roundtrip(sch=s, mod=func)


def test_tensorize_intrinsic_with_annotation():
    func = matmul
    s = tvm.s_tir.Schedule(func, debug_mask="all")
    update = s.get_sblock("update")
    _, _, k = s.get_loops(update)
    ko, ki = s.split(k, factors=[None, 4])
    s.decompose_reduction(update, ko)
    s.tensorize(ki, "test_dot_product_intrin_annotated")

    b = s.get(s.get_sblock("update_update_o"))
    assert b.annotations["test_annotation"] == T.bool(True)
    verify_trace_roundtrip(sch=s, mod=func)


def get_matmul_packed(m, n, k, lhs_type, rhs_dtype="int8"):
    X = te.placeholder((m, k), name="X", dtype=lhs_type)
    W = te.placeholder((n, k), name="W", dtype=rhs_dtype)

    ak = te.reduce_axis((0, k), name="k")
    matmul = te.compute(
        (m, n),
        lambda i, j: te.sum(
            X[i, ak].astype("int32") * W[j, ak].astype("int32"),
            axis=ak,
        ),
        name="compute",
    )

    return te.create_prim_func([X, W, matmul])


def tensorize_16x4_test(intrin=VNNI_DOT_16x4_INTRIN):
    m, n, k = 128, 128, 128

    func = get_matmul_packed(m, n, k, "uint8")

    sch = tvm.s_tir.Schedule(func, debug_mask="all")
    block = sch.get_sblock("compute")
    sch.transform_layout(block, "W", lambda i, j: [i//16, j//4, i%16, j%4])
    _, j, k = sch.get_loops(block)

    _, ji = sch.split(j, factors=[None, 16])
    ko, ki = sch.split(k, factors=[None, 4])
    sch.reorder(ko, ji, ki)

    sch.decompose_reduction(block, ko)
    sch.tensorize(ji, intrin)

    verify_trace_roundtrip(sch=sch, mod=func)


def test_tensorize_vnni():
    tensorize_16x4_test()


def test_tensorize_avx512():
    tensorize_16x4_test(AVX512_DOT_16x4_INTRIN)


def test_tensorize_arm_dot():
    m, n, k = 128, 128, 128

    func = get_matmul_packed(m, n, k, "int8")

    for intrin in [ARM_DOT_4x4_i8_SDOT_INTRIN, ARM_DOT_4x4_i8_NEON_INTRIN]:
        sch = tvm.s_tir.Schedule(func, debug_mask="all")
        block = sch.get_sblock("compute")
        sch.transform_layout(block, "W", lambda i, j: [i//4, j//4, i%4, j%4])
        _, j, k = sch.get_loops(block)

        _, ji = sch.split(j, factors=[None, 4])
        ko, ki = sch.split(k, factors=[None, 4])
        sch.reorder(ko, ji, ki)

        sch.decompose_reduction(block, ko)
        sch.tensorize(ji, intrin)

        verify_trace_roundtrip(sch=sch, mod=func)


def test_tensorize_vrmpy():
    m, n, k = 128, 128, 128

    func = get_matmul_packed(m, n, k, "uint8", "uint8")

    sch = tvm.s_tir.Schedule(func, debug_mask="all")
    block = sch.get_sblock("compute")
    sch.transform_layout(block, "W", lambda i, j: [i//32, j//4, i%32, j%4])
    _, j, k = sch.get_loops(block)

    _, ji = sch.split(j, factors=[None, 32])
    ko, ki = sch.split(k, factors=[None, 4])
    sch.reorder(ko, ji, ki)

    sch.decompose_reduction(block, ko)
    sch.tensorize(ji, VRMPY_u8u8i32_INTRIN)

    verify_trace_roundtrip(sch=sch, mod=func)


def test_tensorize_vdmpy():
    m, n, k = 128, 128, 128

    func = get_matmul_packed(m, n, k, "int16", "int16")

    sch = tvm.s_tir.Schedule(func, debug_mask="all")
    block = sch.get_sblock("compute")
    sch.transform_layout(block, "W", lambda i, j: [i//32, j//2, i%32, j%2])
    _, j, k = sch.get_loops(block)

    _, ji = sch.split(j, factors=[None, 32])
    ko, ki = sch.split(k, factors=[None, 2])
    sch.reorder(ko, ji, ki)

    sch.decompose_reduction(block, ko)
    sch.tensorize(ji, VDMPY_i16i16i32_INTRIN)

    verify_trace_roundtrip(sch=sch, mod=func)


def test_tensorize_dp4a():
    # pylint: disable=too-many-locals
    def _test_intrin(dtype_a, dtype_b, dtype_c, intrin):
        m, n, k = 128, 128, 128
        X = te.placeholder((m, k), name="X", dtype=dtype_a)
        W = te.placeholder((n, k), name="W", dtype=dtype_b)
        ak = te.reduce_axis((0, k), name="k")

        matmul = te.compute(
            (m, n),
            lambda i, j: te.sum(
                X[i, ak].astype(dtype_c) * W[j, ak].astype(dtype_c),
                axis=ak,
            ),
            name="compute",
        )

        func = te.create_prim_func([X, W, matmul])

        sch = tvm.s_tir.Schedule(func, debug_mask="all")
        block = sch.get_sblock("compute")
        i, j, k = sch.get_loops(block)

        by, ty, yi = sch.split(i, factors=sch.sample_perfect_tile(i, n=3))
        bx, tx, xi = sch.split(j, factors=sch.sample_perfect_tile(j, n=3))
        ko, ki = sch.split(k, [None, 4])
        ko, kt = sch.split(ko, factors=sch.sample_perfect_tile(ko, n=2))

        sch.reorder(by, bx, ty, tx, yi, xi)
        CC = sch.cache_write(block, 0, "local")
        sch.reverse_compute_at(CC, tx)

        def fetch_to_shared(block, idx):
            block_read = sch.cache_read(block, idx, "shared")
            sch.compute_at(block_read, ko, True)
            return block_read

        fetch_to_shared(block, 0)
        fetch_to_shared(block, 1)

        sch.decompose_reduction(block, ko)
        sch.tensorize(ki, intrin)

        verify_trace_roundtrip(sch=sch, mod=func)

    for args in [
        ("int8", "int8", "int32", AMDGPU_SDOT4_INTRIN),
        ("int8", "int8", "int32", DP4A_S8S8S32_INTRIN),
        ("int8", "uint8", "int32", DP4A_S8U8S32_INTRIN),
        ("uint8", "int8", "int32", DP4A_U8S8S32_INTRIN),
        ("uint8", "uint8", "uint32", DP4A_U8U8U32_INTRIN),
    ]:
        _test_intrin(*args)


def test_tensor_intrin_look_up():
    intrin_name = 'non_existent_intrin'
    assert tirx.TensorIntrin.get(intrin_name, allow_missing=True) is None
    with pytest.raises(ValueError):
        tirx.TensorIntrin.get(intrin_name)


def test_tensorize_matmul_mixed_dtype():
    # fmt: off
    @T.prim_func(s_tir=True)
    def matmul_int64_shape(
        A: T.Buffer((T.int64(128), T.int64(128)), "float32"),
        B: T.Buffer((T.int64(128), T.int64(128)), "float32"),
        C: T.Buffer((T.int64(128), T.int64(128)), "float32")
    ) -> None:
        for i_0, j_0 in T.grid(T.int64(8), T.int64(8)):
            for i_1_init, j_1_init in T.grid(T.int64(16), T.int64(16)):
                with T.sblock("init"):
                    vi = T.axis.spatial(T.int64(128), i_0 * T.int64(16) + i_1_init)
                    vj = T.axis.spatial(T.int64(128), j_0 * T.int64(16) + j_1_init)
                    C[vi, vj] = T.float32(0)
            for k_0, i_1, j_1, k_1 in T.grid(T.int64(8), T.int64(16), T.int64(16), T.int64(16)):
                with T.sblock("update"):
                    vi = T.axis.spatial(T.int64(128), i_0 * T.int64(16) + i_1)
                    vj = T.axis.spatial(T.int64(128), j_0 * T.int64(16) + j_1)
                    vk = T.axis.reduce(T.int64(128), k_0 * T.int64(16) + k_1)
                    C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]

    @T.prim_func(s_tir=True)
    def tensorized_matmul_int64_shape(
        A: T.Buffer((T.int64(128), T.int64(128)), "float32"),
        B: T.Buffer((T.int64(128), T.int64(128)), "float32"),
        C: T.Buffer((T.int64(128), T.int64(128)), "float32")
    ) -> None:
        for i_outer, j_outer in T.grid(T.int64(8), T.int64(8)):
            for i_inner_init, j_inner_init in T.grid(T.int64(16), T.int64(16)):
                with T.sblock("init"):
                    vi = T.axis.spatial(T.int64(128), i_outer * T.int64(16) + i_inner_init)
                    vj = T.axis.spatial(T.int64(128), j_outer * T.int64(16) + j_inner_init)
                    C[vi, vj] = T.float32(0)
            for k_outer in T.grid(T.int64(8)):
                with T.sblock("update"):
                    vi, vj, vk = T.axis.remap("SSR", [i_outer, j_outer, k_outer])
                    T.reads(
                        [
                            C[vi * T.int64(16) : vi * T.int64(16) + T.int64(16), vj * T.int64(16) : vj * T.int64(16) + T.int64(16)],
                            A[vi * T.int64(16) : vi * T.int64(16) + T.int64(16), vk * T.int64(16) : vk * T.int64(16) + T.int64(16)],
                            B[vj * T.int64(16) : vj * T.int64(16) + T.int64(16), vk * T.int64(16) : vk * T.int64(16) + T.int64(16)],
                        ]
                    )
                    T.writes(C[vi * T.int64(16) : vi * T.int64(16) + T.int64(16), vj * T.int64(16) : vj * T.int64(16) + T.int64(16)])
                    A_elem_offset = T.int64()
                    B_elem_offset = T.int64()
                    C_elem_offset = T.int64()
                    A_sub = T.match_buffer(
                        A[vi * T.int64(16) : vi * T.int64(16) + T.int64(16), vk * T.int64(16) : vk * T.int64(16) + T.int64(16)],
                        [T.int64(16), T.int64(16)],
                        elem_offset=A_elem_offset,
                    )
                    B_sub = T.match_buffer(
                        B[vj * T.int64(16) : vj * T.int64(16) + T.int64(16), vk * T.int64(16) : vk * T.int64(16) + T.int64(16)],
                        [T.int64(16), T.int64(16)],
                        elem_offset=B_elem_offset,
                    )
                    C_sub = T.match_buffer(
                        C[vi * T.int64(16) : vi * T.int64(16) + T.int64(16), vj * T.int64(16) : vj * T.int64(16) + T.int64(16)],
                        [T.int64(16), T.int64(16)],
                        elem_offset=C_elem_offset,
                    )
                    T.evaluate(
                        T.tvm_mma_sync(
                            C_sub.data,
                            T.floordiv(C_sub.elem_offset, T.int64(256)),
                            A_sub.data,
                            T.floordiv(A_sub.elem_offset, T.int64(256)),
                            B_sub.data,
                            T.floordiv(B_sub.elem_offset, T.int64(256)),
                            C_sub.data,
                            T.floordiv(C_sub.elem_offset, T.int64(256)),
                            dtype="handle",
                        )
                    )
    # fmt: on

    s = tvm.s_tir.Schedule(matmul_int64_shape, debug_mask="all")
    update = s.get_sblock("update")
    ii = s.get_loops(update)[-3]
    s.tensorize(ii, "test_mma_intrin")
    assert_structural_equal_ignore_global_symbol(s.mod["main"], tensorized_matmul_int64_shape)
    verify_trace_roundtrip(sch=s, mod=matmul_int64_shape)

def _tir_packed_int_to_int_to_float(storage_nbit: int):
    storage_dtype = "int" + str(storage_nbit)

    def f_convert(nbit: int, val: tirx.Expr, pos: tirx.Expr, dtype: str):
        assert val.ty.dtype == storage_dtype
        mask = tirx.const((1 << nbit) - 1, "int32")
        unextended = (val >> (pos.astype("int32") * tirx.const(nbit, "int32"))) & mask
        return tirx.Cast(dtype, (unextended << tirx.const(32 - nbit, "int32")) >> tirx.const(32 - nbit, "int32"))

    return f_convert

@T.prim_func(s_tir=True)
def decode_i4s_to_f16_desc(compressed: T.handle, decompressed: T.handle) -> None:
    Compressed = T.match_buffer(
        compressed,
        [
            1,
        ],
        dtype="int32",
        scope="local",
    )
    Decompressed = T.match_buffer(
        decompressed,
        [
            8,
        ],
        dtype="float16",
        scope="local",
    )

    with T.sblock("root"):
        T.reads(Compressed[0:1])
        T.writes(Decompressed[0:8])
        for i in T.grid(8):
            with T.sblock("decode"):
                vi = T.axis.remap("S", [i])
                Decompressed[vi] = _tir_packed_int_to_int_to_float(32)(
                    4,
                    Compressed[vi // 8],
                    vi % 8,
                    dtype="float16",
                )

@T.prim_func(s_tir=True)
def decode_i4s_to_f16_impl(compressed: T.handle, decompressed: T.handle) -> None:
    Compressed = T.match_buffer(
        compressed,
        [
            1,
        ],
        dtype="int32",
        scope="local",
    )
    Decompressed = T.match_buffer(
        decompressed,
        [
            8,
        ],
        dtype="float16",
        scope="local",
    )

    with T.sblock("root"):
        T.reads(Compressed[0:1])
        T.writes(Decompressed[0:8])
        T.call_extern(
            "handle",
            "test_decode_i4s_to_f16",
            Compressed.data,
            Decompressed.data,
            8,
        )

tirx.TensorIntrin.register("test_decode_i4s_to_f16_intrin", decode_i4s_to_f16_desc, decode_i4s_to_f16_impl)

def test_tensorize_arith_simplification():
    # fmt: off
    @T.prim_func(s_tir=True)
    def decode_i4s_to_int32_to_f16():
        B_decode_local = T.sblock_alloc_buffer((16384, 16384), "float16", scope="local")
        B_local = T.sblock_alloc_buffer((16384, 2048), "int32", scope="local")
        for ax0_0 in T.thread_binding(8192, thread="blockIdx.x"):
            for ax0_1 in T.thread_binding(2, thread="threadIdx.y"):
                for ax1_0 in range(32):
                    for ax1_1 in T.thread_binding(64, thread="threadIdx.x"):
                        for ax0, ax1 in T.grid(1, 8):
                            with T.sblock("B_decode_local"):
                                v0 = T.axis.spatial(16384, ax0_0 * 2 + ax0_1 + ax0)
                                v1 = T.axis.spatial(16384, ax1_0 * 512 + ax1_1 * 8 + ax1)
                                T.reads(B_local[v0, v1 // 8])
                                T.writes(B_decode_local[v0, v1])
                                B_decode_local[v0, v1] = T.Cast("float16", T.shift_right(T.shift_left(T.bitwise_and(T.shift_right(B_local[v0, v1 // 8], v1 % 8 * 4), 15), 28), 28))

    @T.prim_func(s_tir=True)
    def tensorized_decode_i4s_to_int32_to_f16():
        B_decode_local = T.sblock_alloc_buffer((16384, 16384), "float16", scope="local")
        B_local = T.sblock_alloc_buffer((16384, 2048), "int32", scope="local")
        for ax0_0 in T.thread_binding(8192, thread="blockIdx.x"):
            for ax0_1 in T.thread_binding(2, thread="threadIdx.y"):
                for ax1_0 in range(32):
                    for ax1_1 in T.thread_binding(64, thread="threadIdx.x"):
                        for ax0 in range(1):
                            with T.sblock("B_decode_local_o"):
                                v0_o = T.axis.spatial(16384, ax0_0 * 2 + ax0_1 + ax0)
                                v1_o = T.axis.spatial(2048, ax1_0 * 64 + ax1_1)
                                T.reads(B_local[v0_o, v1_o])
                                T.writes(B_decode_local[v0_o, v1_o * 8:v1_o * 8 + 8])
                                Compressed = T.match_buffer(B_local[v0_o, v1_o], (1,), "int32", scope="local")
                                Decompressed = T.match_buffer(B_decode_local[v0_o, v1_o * 8:v1_o * 8 + 8], (8,), "float16", scope="local")
                                T.call_extern("handle", "test_decode_i4s_to_f16", Compressed.data, Decompressed.data, 8)

    s = tvm.s_tir.Schedule(decode_i4s_to_int32_to_f16, debug_mask="all")
    update = s.get_sblock("B_decode_local")
    ii = s.get_loops(update)[-1]
    s.tensorize(ii, "test_decode_i4s_to_f16_intrin")
    assert_structural_equal_ignore_global_symbol(s.mod["main"], tensorized_decode_i4s_to_int32_to_f16)
    verify_trace_roundtrip(sch=s, mod=decode_i4s_to_int32_to_f16)


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