# 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

# 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
import sys

import pytest

import tvm
from tvm import tirx
from tvm.s_tir.schedule.testing import (
    assert_structural_equal_ignore_global_symbol,
    verify_trace_roundtrip,
)
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,not-callable

@T.prim_func(s_tir=True)
def cuda_matmul(a: T.handle, b: T.handle, c: T.handle) -> None:  # pylint: disable=undefined-loop-variable
    A = T.match_buffer(a, [2048, 2048], "float32")
    B = T.match_buffer(b, [2048, 2048], "float32")
    C = T.match_buffer(c, [2048, 2048], "float32")
    for by in T.thread_binding(0, 32, thread = "blockIdx.y"):
        for bx in T.thread_binding(0, 32, thread = "blockIdx.x"):
            for vy in T.thread_binding(0, 2, thread = "vthread.y"):
                for vx in T.thread_binding(0, 2, thread = "vthread.x"):
                    for ty in T.thread_binding(0, 8, thread = "threadIdx.y"):
                        for tx in T.thread_binding(0, 8, thread = "threadIdx.x"):
                            for k0 in T.serial(0, 256):
                                for k1 in T.unroll(0, 8):
                                    for _, i, j in T.grid(1, 4, 4):
                                        with T.sblock("C"):
                                            vi = T.axis.S(2048, by * 64 + vy * 32 + ty * 4 + i)
                                            vj = T.axis.S(2048, bx * 64 + vx * 32 + tx * 4 + j)
                                            vk = T.axis.R(2048, k0 * 8 + k1)
                                            T.reads([C[vi, vj], A[vi, vk], B[vk, vj]])
                                            T.writes([C[vi, vj]])
                                            with T.init():
                                                C[vi, vj] = 0.0
                                            C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]


@T.prim_func(s_tir=True)
def cuda_matmul_read_at_a(a: T.handle, b: T.handle, c: T.handle) -> None:
    A = T.match_buffer(a, [2048, 2048], dtype="float32")
    B = T.match_buffer(b, [2048, 2048], dtype="float32")
    C = T.match_buffer(c, [2048, 2048], dtype="float32")
    A_shared = T.sblock_alloc_buffer([2048, 2048], dtype="float32", scope="shared")
    for by in T.thread_binding(0, 32, thread="blockIdx.y"):
        for bx in T.thread_binding(0, 32, thread="blockIdx.x"):
            for vy in T.thread_binding(0, 2, thread="vthread.y"):
                for vx in T.thread_binding(0, 2, thread="vthread.x"):
                    for ty in T.thread_binding(0, 8, thread="threadIdx.y"):
                        for tx in T.thread_binding(0, 8, thread="threadIdx.x"):
                            for k0 in T.serial(0, 256):
                                with T.sblock("A_shared"):
                                    v0 = T.axis.S(32, by)
                                    v1 = T.axis.S(256, k0)
                                    T.reads([A[v0 * 64 : v0 * 64 + 64, v1 * 8 : v1 * 8 + 8]])
                                    T.writes([A_shared[v0 * 64 : v0 * 64 + 64, v1 * 8 : v1 * 8 + 8]])
                                    T.sblock_attr({"auto_copy": True})
                                    for ax0, ax1 in T.grid(64, 8):
                                        A_shared[v0 * 64 + ax0, v1 * 8 + ax1] = A[v0 * 64 + ax0, v1 * 8 + ax1]
                                for k1 in T.unroll(0, 8):
                                    for v_, i, j in T.grid(1, 4, 4):
                                        with T.sblock("C"):
                                            vi = T.axis.S(2048, by * 64 + vy * 32 + ty * 4 + i)
                                            vj = T.axis.S(2048, bx * 64 + vx * 32 + tx * 4 + j)
                                            vk = T.axis.R(2048, k0 * 8 + k1)
                                            T.reads([C[vi, vj], A_shared[vi, vk], B[vk, vj]])
                                            T.writes([C[vi, vj]])
                                            with T.init():
                                                C[vi, vj] = T.float32(0)
                                            C[vi, vj] = C[vi, vj] + A_shared[vi, vk] * B[vk, vj]


@T.prim_func(s_tir=True)
def cuda_matmul_read_at_ab(a: T.handle, b: T.handle, c: T.handle) -> None:
    A = T.match_buffer(a, [2048, 2048], dtype="float32")
    B = T.match_buffer(b, [2048, 2048], dtype="float32")
    C = T.match_buffer(c, [2048, 2048], dtype="float32")
    A_shared = T.sblock_alloc_buffer([2048, 2048], dtype="float32", scope="shared")
    B_shared = T.sblock_alloc_buffer([2048, 2048], dtype="float32", scope="shared")
    for by in T.thread_binding(0, 32, thread="blockIdx.y"):
        for bx in T.thread_binding(0, 32, thread="blockIdx.x"):
            for vy in T.thread_binding(0, 2, thread="vthread.y"):
                for vx in T.thread_binding(0, 2, thread="vthread.x"):
                    for ty in T.thread_binding(0, 8, thread="threadIdx.y"):
                        for tx in T.thread_binding(0, 8, thread="threadIdx.x"):
                            for k0 in T.serial(0, 256):
                                with T.sblock("A_shared"):
                                    v0 = T.axis.S(32, by)
                                    v1 = T.axis.S(256, k0)
                                    T.reads([A[v0 * 64 : v0 * 64 + 64, v1 * 8 : v1 * 8 + 8]])
                                    T.writes([A_shared[v0 * 64 : v0 * 64 + 64, v1 * 8 : v1 * 8 + 8]])
                                    T.sblock_attr({"auto_copy": True})
                                    for ax0, ax1 in T.grid(64, 8):
                                        A_shared[v0 * 64 + ax0, v1 * 8 + ax1] = A[v0 * 64 + ax0, v1 * 8 + ax1]
                                with T.sblock("B_shared"):
                                    v0 = T.axis.S(256, k0)
                                    v1 = T.axis.S(32, bx)
                                    T.reads([B[v0 * 8 : v0 * 8 + 8, v1 * 64 : v1 * 64 + 64]])
                                    T.writes([B_shared[v0 * 8 : v0 * 8 + 8, v1 * 64 : v1 * 64 + 64]])
                                    T.sblock_attr({"auto_copy": True})
                                    for ax0, ax1 in T.grid(8, 64):
                                        B_shared[v0 * 8 + ax0, v1 * 64 + ax1] = B[v0 * 8 + ax0, v1 * 64 + ax1]
                                for k1 in T.unroll(0, 8):
                                    for v_, i, j in T.grid(1, 4, 4):
                                        with T.sblock("C"):
                                            vi = T.axis.S(2048, by * 64 + vy * 32 + ty * 4 + i)
                                            vj = T.axis.S(2048, bx * 64 + vx * 32 + tx * 4 + j)
                                            vk = T.axis.R(2048, k0 * 8 + k1)
                                            T.reads([C[vi, vj], A_shared[vi, vk], B_shared[vk, vj]])
                                            T.writes([C[vi, vj]])
                                            with T.init():
                                                C[vi, vj] = T.float32(0)
                                            C[vi, vj] = C[vi, vj] + A_shared[vi, vk] * B_shared[vk, vj]

@T.prim_func(s_tir=True)
def cuda_matmul_write_at_c(a: T.handle, b: T.handle, c: T.handle) -> None:
    A = T.match_buffer(a, [2048, 2048], dtype="float32")
    B = T.match_buffer(b, [2048, 2048], dtype="float32")
    C = T.match_buffer(c, [2048, 2048], dtype="float32")
    A_shared = T.sblock_alloc_buffer([2048, 2048], dtype="float32", scope="shared")
    B_shared = T.sblock_alloc_buffer([2048, 2048], dtype="float32", scope="shared")
    C_shared = T.sblock_alloc_buffer([2048, 2048], dtype="float32", scope="shared")
    for by in T.thread_binding(0, 32, thread="blockIdx.y"):
        for bx in T.thread_binding(0, 32, thread="blockIdx.x"):
            for vy in T.thread_binding(0, 2, thread="vthread.y"):
                for vx in T.thread_binding(0, 2, thread="vthread.x"):
                    for ty in T.thread_binding(0, 8, thread="threadIdx.y"):
                        for tx in T.thread_binding(0, 8, thread="threadIdx.x"):
                            for k0 in T.serial(0, 256):
                                with T.sblock("A_shared"):
                                    v0 = T.axis.S(32, by)
                                    v1 = T.axis.S(256, k0)
                                    T.reads([A[v0 * 64 : v0 * 64 + 64, v1 * 8 : v1 * 8 + 8]])
                                    T.writes([A_shared[v0 * 64 : v0 * 64 + 64, v1 * 8 : v1 * 8 + 8]])
                                    T.sblock_attr({"auto_copy": True})
                                    for ax0, ax1 in T.grid(64, 8):
                                        A_shared[v0 * 64 + ax0, v1 * 8 + ax1] = A[v0 * 64 + ax0, v1 * 8 + ax1]
                                with T.sblock("B_shared"):
                                    v0 = T.axis.S(256, k0)
                                    v1 = T.axis.S(32, bx)
                                    T.reads([B[v0 * 8 : v0 * 8 + 8, v1 * 64 : v1 * 64 + 64]])
                                    T.writes([B_shared[v0 * 8 : v0 * 8 + 8, v1 * 64 : v1 * 64 + 64]])
                                    T.sblock_attr({"auto_copy": True})
                                    for ax0, ax1 in T.grid(8, 64):
                                        B_shared[v0 * 8 + ax0, v1 * 64 + ax1] = B[v0 * 8 + ax0, v1 * 64 + ax1]
                                for k1 in T.unroll(0, 8):
                                    for v_, i, j in T.grid(1, 4, 4):
                                        with T.sblock("C"):
                                            vi = T.axis.S(2048, by * 64 + vy * 32 + ty * 4 + i)
                                            vj = T.axis.S(2048, bx * 64 + vx * 32 + tx * 4 + j)
                                            vk = T.axis.R(2048, k0 * 8 + k1)
                                            T.reads([C_shared[vi, vj], A_shared[vi, vk], B_shared[vk, vj]])
                                            T.writes([C_shared[vi, vj]])
                                            with T.init():
                                                C_shared[vi, vj] = T.float32(0)
                                            C_shared[vi, vj] = C_shared[vi, vj] + A_shared[vi, vk] * B_shared[vk, vj]
                            with T.sblock("C_shared"):
                                v0 = T.axis.S(32, by)
                                v1 = T.axis.S(32, bx)
                                T.reads([C_shared[v0 * 64 : v0 * 64 + 64, v1 * 64 : v1 * 64 + 64]])
                                T.writes([C[v0 * 64 : v0 * 64 + 64, v1 * 64 : v1 * 64 + 64]])
                                T.sblock_attr({"auto_copy": True})
                                for ax0, ax1 in T.grid(64, 64):
                                    C[v0 * 64 + ax0, v1 * 64 + ax1] = C_shared[v0 * 64 + ax0, v1 * 64 + ax1]


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


def test_read_at_global_to_shared_a():
    sch = tvm.s_tir.Schedule(cuda_matmul, debug_mask="all")
    block = sch.get_sblock("C")
    # pylint: disable=invalid-name
    _by, _bx, _vy, _vx, _ty, _tx, k0, _k1, _, _i, _j = sch.get_loops(block)
    # pylint: enable=invalid-name
    sch.read_at(k0, block, 1, "shared")
    assert_structural_equal_ignore_global_symbol(sch.mod["main"], cuda_matmul_read_at_a)
    verify_trace_roundtrip(sch, cuda_matmul)


def test_read_at_global_to_shared_ab():
    sch = tvm.s_tir.Schedule(cuda_matmul_read_at_a, debug_mask="all")
    block = sch.get_sblock("C")
    # pylint: disable=invalid-name
    _by, _bx, _vy, _vx, _ty, _tx, k0, _k1, _, _i, _j = sch.get_loops(block)
    # pylint: enable=invalid-name
    sch.read_at(k0, block, 2, "shared")
    assert_structural_equal_ignore_global_symbol(sch.mod["main"], cuda_matmul_read_at_ab)
    verify_trace_roundtrip(sch, cuda_matmul_read_at_a)


def test_read_at_local_to_shared_c():
    sch = tvm.s_tir.Schedule(cuda_matmul_read_at_ab, debug_mask="all")
    block = sch.get_sblock("C")
    # pylint: disable=invalid-name
    _by, _bx, _vy, _vx, _ty, tx, _k0, _k1, _, _i, _j = sch.get_loops(block)
    # pylint: enable=invalid-name
    sch.write_at(tx, block, 0, "shared")
    assert_structural_equal_ignore_global_symbol(sch.mod["main"], cuda_matmul_write_at_c)
    verify_trace_roundtrip(sch, cuda_matmul_read_at_ab)


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