# 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,line-too-long,invalid-name,too-few-public-methods,too-many-locals
# ruff: noqa: E501, F841

import tvm.testing
from tvm.ir import assert_structural_equal
from tvm.s_tir import dlight as dl
from tvm.script import ir as I
from tvm.script import tirx as T
from tvm.target import Target


def test_decode_gemv_1():
    # NK layout + K as decode dim
    # fmt: off
    @I.ir_module(s_tir=True)
    class Before:
        @T.prim_func(s_tir=True)
        def func(W: T.Buffer((4096, 512), "uint32"), S: T.Buffer((4096, 128), "float16"), V: T.Buffer((1, 1, 4096), "float16"), C: T.Buffer((1, 1, 4096), "float16")):
            T.func_attr({"global_symbol": "main", "tirx.noalias": True})
            # with T.sblock("root"):
            B = T.sblock_alloc_buffer((4096, 4096), "float16")
            for i, j in T.grid(4096, 4096):
                with T.sblock("decode"):
                    v_i, v_j = T.axis.remap("SS", [i, j])
                    T.reads(W[v_i, v_j // 8], S[v_i, v_j // 32])
                    T.writes(B[v_i, v_j])
                    B[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(W[v_i, v_j // 8], T.Cast("uint32", v_j % 8) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * S[v_i, v_j // 32]
            for i0, i1, i2, k in T.grid(1, 1, 4096, 4096):
                with T.sblock("matmul"):
                    v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
                    T.reads(V[v_i0, v_i1, v_k], B[v_i2, v_k])
                    T.writes(C[v_i0, v_i1, v_i2])
                    with T.init():
                        C[v_i0, v_i1, v_i2] = T.float16(0)
                    C[v_i0, v_i1, v_i2] = C[v_i0, v_i1, v_i2] + V[v_i0, v_i1, v_k] * B[v_i2, v_k]


    @I.ir_module(s_tir=True)
    class After:
        @T.prim_func(s_tir=True)
        def func(W_handle: T.handle, S_handle: T.handle, V_handle: T.handle, C_handle: T.handle):
            T.func_attr({"global_symbol": "main", "tirx.is_scheduled": True, "tirx.noalias": True})
            W = T.match_buffer(W_handle, (4096, 512), "uint32")
            S = T.match_buffer(S_handle, (4096, 128), "float16")
            V = T.match_buffer(V_handle, (1, 1, 4096), "float16")
            C = T.match_buffer(C_handle, (1, 1, 4096), "float16")
            with T.sblock("root"):
                T.reads()
                T.writes()
                C_rf_local = T.sblock_alloc_buffer((512, 1, 1, 4096), "float16", scope="local")
                for ax0_fused in T.thread_binding(4096, thread="blockIdx.x"):
                    for ax1_0_fused_1 in T.thread_binding(512, thread="threadIdx.x", annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
                        with T.sblock("matmul_rf_init"):
                            vax1_0_fused_1 = T.axis.spatial(512, ax1_0_fused_1)
                            v0 = T.axis.spatial(4096, ax0_fused)
                            T.reads()
                            T.writes(C_rf_local[vax1_0_fused_1, 0, 0, v0])
                            C_rf_local[vax1_0_fused_1, 0, 0, v0] = T.float16(0)
                        for ax1_0_fused_0 in range(1):
                            for ax1_1 in range(8):
                                with T.sblock("matmul_rf_update"):
                                    vax1_0_fused_1 = T.axis.spatial(512, ax1_0_fused_1)
                                    v0 = T.axis.spatial(4096, ax0_fused)
                                    vax1_0_fused_0 = T.axis.reduce(1, ax1_0_fused_0)
                                    vax1_1 = T.axis.reduce(8, ax1_1)
                                    T.reads(C_rf_local[vax1_0_fused_1, 0, 0, v0], V[0, 0, vax1_0_fused_0 * 4096 + vax1_0_fused_1 * 8 + vax1_1], W[v0, (vax1_0_fused_0 * 4096 + vax1_0_fused_1 * 8 + vax1_1) // 8], S[v0, (vax1_0_fused_0 * 4096 + vax1_0_fused_1 * 8 + vax1_1) // 32])
                                    T.writes(C_rf_local[vax1_0_fused_1, 0, 0, v0])
                                    C_rf_local[vax1_0_fused_1, 0, 0, v0] = C_rf_local[vax1_0_fused_1, 0, 0, v0] + V[0, 0, vax1_0_fused_0 * 4096 + vax1_0_fused_1 * 8 + vax1_1] * ((T.Cast("float16", T.bitwise_and(T.shift_right(W[v0, (vax1_0_fused_0 * 4096 + vax1_0_fused_1 * 8 + vax1_1) // 8], T.Cast("uint32", (vax1_0_fused_0 * 4096 + vax1_0_fused_1 * 8 + vax1_1) % 8) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * S[v0, (vax1_0_fused_0 * 4096 + vax1_0_fused_1 * 8 + vax1_1) // 32])
                    for ax1_fused in range(1):
                        for ax0 in T.thread_binding(512, thread="threadIdx.x"):
                            with T.sblock("matmul"):
                                vax1_0_fused_1 = T.axis.reduce(512, ax0)
                                v0 = T.axis.spatial(4096, ax0_fused)
                                T.reads(C_rf_local[vax1_0_fused_1, 0, 0, v0])
                                T.writes(C[0, 0, v0])
                                with T.init():
                                    C[0, 0, v0] = T.float16(0)
                                C[0, 0, v0] = C[0, 0, v0] + C_rf_local[vax1_0_fused_1, 0, 0, v0]
    # fmt: on

    target = Target("nvidia/geforce-rtx-3090-ti")
    with target:
        mod = dl.ApplyDefaultSchedule(dl.gpu.Reduction())(Before)  # pylint: disable=not-callable
    assert_structural_equal(mod, After)


def test_decode_gemv_2():
    # KN layout + K as decode dim
    # fmt: off
    @I.ir_module(s_tir=True)
    class Before:
        @T.prim_func(s_tir=True)
        def func(W: T.Buffer((512, 4096), "uint32"), S: T.Buffer((128, 4096), "float16"), V: T.Buffer((1, 1, 4096), "float16"), C: T.Buffer((1, 1, 4096), "float16")):
            T.func_attr({"global_symbol": "main", "tirx.noalias": True})
            # with T.sblock("root"):
            B = T.sblock_alloc_buffer((4096, 4096), "float16")
            for i, j in T.grid(4096, 4096):
                with T.sblock("decode"):
                    v_i, v_j = T.axis.remap("SS", [i, j])
                    T.reads(W[v_i // 8, v_j], S[v_i // 32, v_j])
                    T.writes(B[v_i, v_j])
                    B[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(W[v_i // 8, v_j], T.Cast("uint32", v_i % 8) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * S[v_i // 32, v_j]
            for i0, i1, i2, k in T.grid(1, 1, 4096, 4096):
                with T.sblock("matmul"):
                    v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
                    T.reads(V[v_i0, v_i1, v_k], B[v_k, v_i2])
                    T.writes(C[v_i0, v_i1, v_i2])
                    with T.init():
                        C[v_i0, v_i1, v_i2] = T.float16(0)
                    C[v_i0, v_i1, v_i2] = C[v_i0, v_i1, v_i2] + V[v_i0, v_i1, v_k] * B[v_k, v_i2]


    @I.ir_module(s_tir=True)
    class After:
        @T.prim_func(s_tir=True)
        def func(W: T.Buffer((512, 4096), "uint32"), S: T.Buffer((128, 4096), "float16"), V: T.Buffer((1, 1, 4096), "float16"), C: T.Buffer((1, 1, 4096), "float16")):
            T.func_attr({"global_symbol": "main", "tirx.is_scheduled": True, "tirx.noalias": True})
            # with T.sblock("root"):
            C_rf_local = T.sblock_alloc_buffer((16, 1, 1, 4096), "float16", scope="local")
            for i2_i0_i1_fused_0 in T.thread_binding(256, thread="blockIdx.x"):
                for i2_i0_i1_fused_1 in T.thread_binding(16, thread="threadIdx.x"):
                    for k_0_fused_1 in T.thread_binding(16, thread="threadIdx.y"):
                        with T.sblock("matmul_rf_init"):
                            vk_0_fused_1 = T.axis.spatial(16, k_0_fused_1)
                            v_i2 = T.axis.spatial(4096, i2_i0_i1_fused_0 * 16 + i2_i0_i1_fused_1)
                            C_rf_local[vk_0_fused_1, 0, 0, v_i2] = T.float16(0)
                        for k_0_fused_0, k_1 in T.grid(32, 8):
                            with T.sblock("matmul_rf_update"):
                                vk_0_fused_1 = T.axis.spatial(16, k_0_fused_1)
                                v_i2 = T.axis.spatial(4096, i2_i0_i1_fused_0 * 16 + i2_i0_i1_fused_1)
                                vk_0_fused_0, vk_1 = T.axis.remap("RR", [k_0_fused_0, k_1])
                                C_rf_local[vk_0_fused_1, 0, 0, v_i2] = C_rf_local[vk_0_fused_1, 0, 0, v_i2] + V[0, 0, vk_0_fused_0 * 128 + vk_0_fused_1 * 8 + vk_1] * ((T.Cast("float16", T.bitwise_and(T.shift_right(W[(vk_0_fused_0 * 128 + vk_0_fused_1 * 8 + vk_1) // 8, v_i2], T.Cast("uint32", (vk_0_fused_0 * 128 + vk_0_fused_1 * 8 + vk_1) % 8) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * S[(vk_0_fused_0 * 128 + vk_0_fused_1 * 8 + vk_1) // 32, v_i2])
                for ax1_ax2_ax3_fused in T.thread_binding(16, thread="threadIdx.x"):
                    for ax0_fused in T.thread_binding(16, thread="threadIdx.y"):
                        with T.sblock("matmul"):
                            vk_0_fused_1 = T.axis.reduce(16, ax0_fused)
                            v_i2 = T.axis.spatial(4096, i2_i0_i1_fused_0 * 16 + ax1_ax2_ax3_fused)
                            with T.init():
                                C[0, 0, v_i2] = T.float16(0)
                            C[0, 0, v_i2] = C[0, 0, v_i2] + C_rf_local[vk_0_fused_1, 0, 0, v_i2]

    # fmt: on

    target = Target("nvidia/geforce-rtx-3090-ti")
    with target:
        mod = dl.ApplyDefaultSchedule(dl.gpu.Reduction())(Before)  # pylint: disable=not-callable
    assert_structural_equal(mod, After)


def test_decode_gemv_3():
    # NK layout + N as decode dim
    # fmt: off
    @I.ir_module(s_tir=True)
    class Before:
        @T.prim_func(s_tir=True)
        def func(W: T.Buffer((512, 4096), "uint32"), S: T.Buffer((128, 4096), "float16"), V: T.Buffer((1, 1, 4096), "float16"), C: T.Buffer((1, 1, 4096), "float16")):
            T.func_attr({"global_symbol": "main", "tirx.noalias": True})
            # with T.sblock("root"):
            B = T.sblock_alloc_buffer((4096, 4096), "float16")
            for i, j in T.grid(4096, 4096):
                with T.sblock("decode"):
                    v_i, v_j = T.axis.remap("SS", [i, j])
                    T.reads(W[v_i // 8, v_j], S[v_i // 32, v_j])
                    T.writes(B[v_i, v_j])
                    B[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(W[v_i // 8, v_j], T.Cast("uint32", v_i % 8) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * S[v_i // 32, v_j]
            for i0, i1, i2, k in T.grid(1, 1, 4096, 4096):
                with T.sblock("matmul"):
                    v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
                    T.reads(V[v_i0, v_i1, v_k], B[v_i2, v_k])
                    T.writes(C[v_i0, v_i1, v_i2])
                    with T.init():
                        C[v_i0, v_i1, v_i2] = T.float16(0)
                    C[v_i0, v_i1, v_i2] = C[v_i0, v_i1, v_i2] + V[v_i0, v_i1, v_k] * B[v_i2, v_k]

    @I.ir_module(s_tir=True)
    class After:
        @T.prim_func(s_tir=True)
        def func(W_handle: T.handle, S_handle: T.handle, V_handle: T.handle, C_handle: T.handle):
            T.func_attr({"global_symbol": "main", "tirx.is_scheduled": True, "tirx.noalias": True})
            W = T.match_buffer(W_handle, (512, 4096), "uint32")
            S = T.match_buffer(S_handle, (128, 4096), "float16")
            V = T.match_buffer(V_handle, (1, 1, 4096), "float16")
            C = T.match_buffer(C_handle, (1, 1, 4096), "float16")
            with T.sblock("root"):
                T.reads()
                T.writes()
                C_rf_local = T.sblock_alloc_buffer((1024, 1, 1, 4096), "float16", scope="local")
                for ax0_0_fused in T.thread_binding(512, thread="blockIdx.x"):
                    for ax1_fused_1 in T.thread_binding(1024, thread="threadIdx.x", annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
                        for ax0_1_init in range(8):
                            with T.sblock("matmul_rf_init"):
                                vax1_fused_1 = T.axis.spatial(1024, ax1_fused_1)
                                v0 = T.axis.spatial(4096, ax0_0_fused * 8 + ax0_1_init)
                                T.reads()
                                T.writes(C_rf_local[vax1_fused_1, 0, 0, v0])
                                C_rf_local[vax1_fused_1, 0, 0, v0] = T.float16(0)
                        for ax1_fused_0 in range(4):
                            for ax0_1 in range(8):
                                with T.sblock("matmul_rf_update"):
                                    vax1_fused_1 = T.axis.spatial(1024, ax1_fused_1)
                                    v0 = T.axis.spatial(4096, ax0_0_fused * 8 + ax0_1)
                                    vax1_fused_0 = T.axis.reduce(4, ax1_fused_0)
                                    T.reads(C_rf_local[vax1_fused_1, 0, 0, v0], V[0, 0, vax1_fused_0 * 1024 + vax1_fused_1], W[v0 // 8, vax1_fused_0 * 1024 + vax1_fused_1], S[v0 // 32, vax1_fused_0 * 1024 + vax1_fused_1])
                                    T.writes(C_rf_local[vax1_fused_1, 0, 0, v0])
                                    C_rf_local[vax1_fused_1, 0, 0, v0] = C_rf_local[vax1_fused_1, 0, 0, v0] + V[0, 0, vax1_fused_0 * 1024 + vax1_fused_1] * ((T.Cast("float16", T.bitwise_and(T.shift_right(W[v0 // 8, vax1_fused_0 * 1024 + vax1_fused_1], T.Cast("uint32", v0 % 8) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * S[v0 // 32, vax1_fused_0 * 1024 + vax1_fused_1])
                    for ax1_fused_0 in range(1):
                        for ax0 in T.thread_binding(1024, thread="threadIdx.x"):
                            for ax1_fused_1 in range(8):
                                with T.sblock("matmul"):
                                    vax1_fused_1 = T.axis.reduce(1024, ax0)
                                    v0 = T.axis.spatial(4096, ax0_0_fused * 8 + ax1_fused_1)
                                    T.reads(C_rf_local[vax1_fused_1, 0, 0, v0])
                                    T.writes(C[0, 0, v0])
                                    with T.init():
                                        C[0, 0, v0] = T.float16(0)
                                    C[0, 0, v0] = C[0, 0, v0] + C_rf_local[vax1_fused_1, 0, 0, v0]

    # fmt: on

    target = Target("nvidia/geforce-rtx-3090-ti")
    with target:
        mod = dl.ApplyDefaultSchedule(dl.gpu.Reduction())(Before)  # pylint: disable=not-callable
    assert_structural_equal(mod, After)


def test_decode_gemv_4():
    # KN layout + N as decode dim
    # fmt: off
    @I.ir_module(s_tir=True)
    class Before:
        @T.prim_func(s_tir=True)
        def func(W: T.Buffer((4096, 512), "uint32"), S: T.Buffer((4096, 128), "float16"), V: T.Buffer((1, 1, 4096), "float16"), C: T.Buffer((1, 1, 4096), "float16")):
            T.func_attr({"global_symbol": "main", "tirx.noalias": True})
            # with T.sblock("root"):
            B = T.sblock_alloc_buffer((4096, 4096), "float16")
            for i, j in T.grid(4096, 4096):
                with T.sblock("decode"):
                    v_i, v_j = T.axis.remap("SS", [i, j])
                    T.reads(W[v_i, v_j // 8], S[v_i, v_j // 32])
                    T.writes(B[v_i, v_j])
                    B[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(W[v_i, v_j // 8], T.Cast("uint32", v_j % 8) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * S[v_i, v_j // 32]
            for i0, i1, i2, k in T.grid(1, 1, 4096, 4096):
                with T.sblock("matmul"):
                    v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
                    T.reads(V[v_i0, v_i1, v_k], B[v_k, v_i2])
                    T.writes(C[v_i0, v_i1, v_i2])
                    with T.init():
                        C[v_i0, v_i1, v_i2] = T.float16(0)
                    C[v_i0, v_i1, v_i2] = C[v_i0, v_i1, v_i2] + V[v_i0, v_i1, v_k] * B[v_k, v_i2]


    @I.ir_module(s_tir=True)
    class After:
        @T.prim_func(s_tir=True)
        def func(W: T.Buffer((4096, 512), "uint32"), S: T.Buffer((4096, 128), "float16"), V: T.Buffer((1, 1, 4096), "float16"), C: T.Buffer((1, 1, 4096), "float16")):
            T.func_attr({"global_symbol": "main", "tirx.is_scheduled": True, "tirx.noalias": True})
            # with T.sblock("root"):
            C_rf_local = T.sblock_alloc_buffer((16, 1, 1, 4096), "float16", scope="local")
            for i2_0_i0_i1_fused_0 in T.thread_binding(32, thread="blockIdx.x"):
                for i2_0_i0_i1_fused_1 in T.thread_binding(16, thread="threadIdx.x"):
                    for k_fused_1 in T.thread_binding(16, thread="threadIdx.y"):
                        for i2_1_init in range(8):
                            with T.sblock("matmul_rf_init"):
                                vk_fused_1 = T.axis.spatial(16, k_fused_1)
                                v_i2 = T.axis.spatial(4096, i2_0_i0_i1_fused_0 * 128 + i2_0_i0_i1_fused_1 * 8 + i2_1_init)
                                C_rf_local[vk_fused_1, 0, 0, v_i2] = T.float16(0)
                        for k_fused_0, i2_1 in T.grid(256, 8):
                            with T.sblock("matmul_rf_update"):
                                vk_fused_1 = T.axis.spatial(16, k_fused_1)
                                v_i2 = T.axis.spatial(4096, i2_0_i0_i1_fused_0 * 128 + i2_0_i0_i1_fused_1 * 8 + i2_1)
                                vk_fused_0 = T.axis.reduce(256, k_fused_0)
                                C_rf_local[vk_fused_1, 0, 0, v_i2] = C_rf_local[vk_fused_1, 0, 0, v_i2] + V[0, 0, vk_fused_0 * 16 + vk_fused_1] * ((T.Cast("float16", T.bitwise_and(T.shift_right(W[vk_fused_0 * 16 + vk_fused_1, v_i2 // 8], T.Cast("uint32", v_i2 % 8) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * S[vk_fused_0 * 16 + vk_fused_1, v_i2 // 32])
                for ax1_ax2_ax3_fused_0 in T.thread_binding(16, thread="threadIdx.x"):
                    for ax1_ax2_ax3_fused_1 in range(8):
                        for ax0_fused in T.thread_binding(16, thread="threadIdx.y"):
                            with T.sblock("matmul"):
                                vk_fused_1 = T.axis.reduce(16, ax0_fused)
                                v_i2 = T.axis.spatial(4096, i2_0_i0_i1_fused_0 * 128 + ax1_ax2_ax3_fused_0 * 8 + ax1_ax2_ax3_fused_1)
                                with T.init():
                                    C[0, 0, v_i2] = T.float16(0)
                                C[0, 0, v_i2] = C[0, 0, v_i2] + C_rf_local[vk_fused_1, 0, 0, v_i2]

    # fmt: on

    target = Target("nvidia/geforce-rtx-3090-ti")
    with target:
        mod = dl.ApplyDefaultSchedule(dl.gpu.Reduction())(Before)  # pylint: disable=not-callable
    assert_structural_equal(mod, After)


def test_decode_gemv_sigmoid():
    # NK layout + K as decode dim
    # fmt: off
    @I.ir_module(s_tir=True)
    class Before:
        @T.prim_func(s_tir=True)
        def func(W: T.Buffer((4096, 512), "uint32"), S: T.Buffer((4096, 128), "float16"), V: T.Buffer((1, 1, 4096), "float16"), D: T.Buffer((1, 1, 4096), "float16")):
            T.func_attr({"global_symbol": "main", "tirx.noalias": True})
            # with T.sblock("root"):
            B = T.sblock_alloc_buffer((4096, 4096), "float16")
            C = T.sblock_alloc_buffer((1, 1, 4096), "float16")
            for i, j in T.grid(4096, 4096):
                with T.sblock("decode"):
                    v_i, v_j = T.axis.remap("SS", [i, j])
                    T.reads(W[v_i, v_j // 8], S[v_i, v_j // 32])
                    T.writes(B[v_i, v_j])
                    B[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(W[v_i, v_j // 8], T.Cast("uint32", v_j % 8) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * S[v_i, v_j // 32]
            for i0, i1, i2, k in T.grid(1, 1, 4096, 4096):
                with T.sblock("matmul"):
                    v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
                    T.reads(V[v_i0, v_i1, v_k], B[v_i2, v_k])
                    T.writes(C[v_i0, v_i1, v_i2])
                    with T.init():
                        C[v_i0, v_i1, v_i2] = T.float16(0)
                    C[v_i0, v_i1, v_i2] = C[v_i0, v_i1, v_i2] + V[v_i0, v_i1, v_k] * B[v_i2, v_k]
            for i0, i1, i2 in T.grid(1, 1, 4096):
                with T.sblock("sigmoid"):
                    v_i0, v_i1, v_i2 = T.axis.remap("SSS", [i0, i1, i2])
                    T.reads(C[v_i0, v_i1, v_i2])
                    T.writes(D[v_i0, v_i1, v_i2])
                    D[v_i0, v_i1, v_i2] = T.sigmoid(C[v_i0, v_i1, v_i2])

    @I.ir_module(s_tir=True)
    class After:
        @T.prim_func(s_tir=True)
        def func(W_handle: T.handle, S_handle: T.handle, V_handle: T.handle, D_handle: T.handle):
            T.func_attr({"global_symbol": "main", "tirx.is_scheduled": True, "tirx.noalias": True})
            W = T.match_buffer(W_handle, (4096, 512), "uint32")
            S = T.match_buffer(S_handle, (4096, 128), "float16")
            V = T.match_buffer(V_handle, (1, 1, 4096), "float16")
            D = T.match_buffer(D_handle, (1, 1, 4096), "float16")
            with T.sblock("root"):
                T.reads()
                T.writes()
                C_local = T.sblock_alloc_buffer((1, 1, 4096), "float16", scope="local")
                C_rf_local = T.sblock_alloc_buffer((512, 1, 1, 4096), "float16", scope="local")
                for ax0_fused in T.thread_binding(4096, thread="blockIdx.x"):
                    for ax1_0_fused_1 in T.thread_binding(512, thread="threadIdx.x", annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
                        with T.sblock("matmul_rf_init"):
                            vax1_0_fused_1 = T.axis.spatial(512, ax1_0_fused_1)
                            v0 = T.axis.spatial(4096, ax0_fused)
                            T.reads()
                            T.writes(C_rf_local[vax1_0_fused_1, 0, 0, v0])
                            C_rf_local[vax1_0_fused_1, 0, 0, v0] = T.float16(0)
                        for ax1_0_fused_0 in range(1):
                            for ax1_1 in range(8):
                                with T.sblock("matmul_rf_update"):
                                    vax1_0_fused_1 = T.axis.spatial(512, ax1_0_fused_1)
                                    v0 = T.axis.spatial(4096, ax0_fused)
                                    vax1_0_fused_0 = T.axis.reduce(1, ax1_0_fused_0)
                                    vax1_1 = T.axis.reduce(8, ax1_1)
                                    T.reads(C_rf_local[vax1_0_fused_1, 0, 0, v0], V[0, 0, vax1_0_fused_0 * 4096 + vax1_0_fused_1 * 8 + vax1_1], W[v0, (vax1_0_fused_0 * 4096 + vax1_0_fused_1 * 8 + vax1_1) // 8], S[v0, (vax1_0_fused_0 * 4096 + vax1_0_fused_1 * 8 + vax1_1) // 32])
                                    T.writes(C_rf_local[vax1_0_fused_1, 0, 0, v0])
                                    C_rf_local[vax1_0_fused_1, 0, 0, v0] = C_rf_local[vax1_0_fused_1, 0, 0, v0] + V[0, 0, vax1_0_fused_0 * 4096 + vax1_0_fused_1 * 8 + vax1_1] * ((T.Cast("float16", T.bitwise_and(T.shift_right(W[v0, (vax1_0_fused_0 * 4096 + vax1_0_fused_1 * 8 + vax1_1) // 8], T.Cast("uint32", (vax1_0_fused_0 * 4096 + vax1_0_fused_1 * 8 + vax1_1) % 8) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * S[v0, (vax1_0_fused_0 * 4096 + vax1_0_fused_1 * 8 + vax1_1) // 32])
                    for ax1_fused in range(1):
                        for ax0 in T.thread_binding(512, thread="threadIdx.x"):
                            with T.sblock("matmul"):
                                vax1_0_fused_1 = T.axis.reduce(512, ax0)
                                v0 = T.axis.spatial(4096, ax0_fused)
                                T.reads(C_rf_local[vax1_0_fused_1, 0, 0, v0])
                                T.writes(C_local[0, 0, v0])
                                with T.init():
                                    C_local[0, 0, v0] = T.float16(0)
                                C_local[0, 0, v0] = C_local[0, 0, v0] + C_rf_local[vax1_0_fused_1, 0, 0, v0]
                    for ax0 in range(1):
                        with T.sblock("sigmoid"):
                            v0 = T.axis.spatial(4096, ax0_fused + ax0)
                            T.reads(C_local[0, 0, v0])
                            T.writes(D[0, 0, v0])
                            D[0, 0, v0] = T.sigmoid(C_local[0, 0, v0])

    # fmt: on

    target = Target("nvidia/geforce-rtx-3090-ti")
    with target:
        mod = dl.ApplyDefaultSchedule(dl.gpu.Reduction())(Before)  # pylint: disable=not-callable
    assert_structural_equal(mod, After)


def test_decode_gemv_1_fp32():
    # NK layout + K as decode dim
    # fmt: off
    @I.ir_module(s_tir=True)
    class Before:
        @T.prim_func(s_tir=True)
        def func(W: T.Buffer((4096, 512), "uint32"), S: T.Buffer((4096, 128), "float16"), V: T.Buffer((1, 1, 4096), "float16"), C: T.Buffer((1, 1, 4096), "float16")):
            T.func_attr({"global_symbol": "main", "tirx.noalias": True})
            # with T.sblock("root"):
            B = T.sblock_alloc_buffer((4096, 4096), "float16")
            C_fp32 = T.sblock_alloc_buffer((1, 1, 4096), "float32")
            for i, j in T.grid(4096, 4096):
                with T.sblock("decode"):
                    v_i, v_j = T.axis.remap("SS", [i, j])
                    T.reads(W[v_i, v_j // 8], S[v_i, v_j // 32])
                    T.writes(B[v_i, v_j])
                    B[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(W[v_i, v_j // 8], T.Cast("uint32", v_j % 8) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * S[v_i, v_j // 32]
            for i0, i1, i2, k in T.grid(1, 1, 4096, 4096):
                with T.sblock("matmul"):
                    v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
                    T.reads(V[v_i0, v_i1, v_k], B[v_i2, v_k])
                    T.writes(C_fp32[v_i0, v_i1, v_i2])
                    with T.init():
                        C_fp32[v_i0, v_i1, v_i2] = T.float16(0)
                    C_fp32[v_i0, v_i1, v_i2] = C_fp32[v_i0, v_i1, v_i2] + T.Cast("float32", V[v_i0, v_i1, v_k]) * T.Cast("float32", B[v_i2, v_k])
            for i0, i1, i2 in T.grid(1, 1, 4096):
                with T.sblock("cast"):
                    v_i0, v_i1, v_i2 = T.axis.remap("SSS", [i0, i1, i2])
                    T.reads(C_fp32[v_i0, v_i1, v_i2])
                    T.writes(C[v_i0, v_i1, v_i2])
                    C[v_i0, v_i1, v_i2] = T.Cast("float16", C_fp32[v_i0, v_i1, v_i2])

    @I.ir_module(s_tir=True)
    class After:
        @T.prim_func(s_tir=True)
        def func(W_handle: T.handle, S_handle: T.handle, V_handle: T.handle, C_handle: T.handle):
            T.func_attr({"global_symbol": "main", "tirx.is_scheduled": True, "tirx.noalias": True})
            W = T.match_buffer(W_handle, (4096, 512), "uint32")
            S = T.match_buffer(S_handle, (4096, 128), "float16")
            V = T.match_buffer(V_handle, (1, 1, 4096), "float16")
            C = T.match_buffer(C_handle, (1, 1, 4096), "float16")
            with T.sblock("root"):
                T.reads()
                T.writes()
                C_fp32_local = T.sblock_alloc_buffer((1, 1, 4096), scope="local")
                C_fp32_rf_local = T.sblock_alloc_buffer((512, 1, 1, 4096), scope="local")
                for ax0_fused in T.thread_binding(4096, thread="blockIdx.x"):
                    for ax1_0_fused_1 in T.thread_binding(512, thread="threadIdx.x", annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
                        with T.sblock("matmul_rf_init"):
                            vax1_0_fused_1 = T.axis.spatial(512, ax1_0_fused_1)
                            v0 = T.axis.spatial(4096, ax0_fused)
                            T.reads()
                            T.writes(C_fp32_rf_local[vax1_0_fused_1, 0, 0, v0])
                            C_fp32_rf_local[vax1_0_fused_1, 0, 0, v0] = T.float32(0)
                        for ax1_0_fused_0 in range(1):
                            for ax1_1 in range(8):
                                with T.sblock("matmul_rf_update"):
                                    vax1_0_fused_1 = T.axis.spatial(512, ax1_0_fused_1)
                                    v0 = T.axis.spatial(4096, ax0_fused)
                                    vax1_0_fused_0 = T.axis.reduce(1, ax1_0_fused_0)
                                    vax1_1 = T.axis.reduce(8, ax1_1)
                                    T.reads(C_fp32_rf_local[vax1_0_fused_1, 0, 0, v0], V[0, 0, vax1_0_fused_0 * 4096 + vax1_0_fused_1 * 8 + vax1_1], W[v0, (vax1_0_fused_0 * 4096 + vax1_0_fused_1 * 8 + vax1_1) // 8], S[v0, (vax1_0_fused_0 * 4096 + vax1_0_fused_1 * 8 + vax1_1) // 32])
                                    T.writes(C_fp32_rf_local[vax1_0_fused_1, 0, 0, v0])
                                    C_fp32_rf_local[vax1_0_fused_1, 0, 0, v0] = C_fp32_rf_local[vax1_0_fused_1, 0, 0, v0] + T.Cast("float32", V[0, 0, vax1_0_fused_0 * 4096 + vax1_0_fused_1 * 8 + vax1_1]) * T.Cast("float32", (T.Cast("float16", T.bitwise_and(T.shift_right(W[v0, (vax1_0_fused_0 * 4096 + vax1_0_fused_1 * 8 + vax1_1) // 8], T.Cast("uint32", (vax1_0_fused_0 * 4096 + vax1_0_fused_1 * 8 + vax1_1) % 8) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * S[v0, (vax1_0_fused_0 * 4096 + vax1_0_fused_1 * 8 + vax1_1) // 32])
                    for ax1_fused in range(1):
                        for ax0 in T.thread_binding(512, thread="threadIdx.x"):
                            with T.sblock("matmul"):
                                vax1_0_fused_1 = T.axis.reduce(512, ax0)
                                v0 = T.axis.spatial(4096, ax0_fused)
                                T.reads(C_fp32_rf_local[vax1_0_fused_1, 0, 0, v0])
                                T.writes(C_fp32_local[0, 0, v0])
                                with T.init():
                                    C_fp32_local[0, 0, v0] = T.float32(0)
                                C_fp32_local[0, 0, v0] = C_fp32_local[0, 0, v0] + C_fp32_rf_local[vax1_0_fused_1, 0, 0, v0]
                    for ax0 in range(1):
                        with T.sblock("cast"):
                            v0 = T.axis.spatial(4096, ax0_fused + ax0)
                            T.reads(C_fp32_local[0, 0, v0])
                            T.writes(C[0, 0, v0])
                            C[0, 0, v0] = T.Cast("float16", C_fp32_local[0, 0, v0])

    # fmt: on

    target = Target("nvidia/geforce-rtx-3090-ti")
    with target:
        mod = dl.ApplyDefaultSchedule(dl.gpu.Reduction())(Before)  # pylint: disable=not-callable
    assert_structural_equal(mod, After)


def test_reduction_no_spatial():
    # fmt: off
    @I.ir_module(s_tir=True)
    class Before:
        @T.prim_func(s_tir=True)
        def main(A: T.Buffer((1, 1, 4096), "float16"), B: T.Buffer((4096,), "float16"), rms_norm: T.Buffer((1, 4096), "float16")):
            T.func_attr({"global_symbol": "main", "tirx.noalias": True})
            Ared_temp = T.sblock_alloc_buffer((1, 1))
            for ax0 in range(4096):
                with T.sblock("Ared_temp"):
                    v0 = T.axis.reduce(4096, ax0)
                    with T.init():
                        Ared_temp[0, 0] = T.float32(0)
                    Ared_temp[0, 0] = Ared_temp[0, 0] + T.Cast("float32", A[0, 0, v0]) * T.Cast("float32", A[0, 0, v0])
            for ax0 in range(4096):
                with T.sblock("rms_norm"):
                    v0 = T.axis.spatial(4096, ax0)
                    rms_norm[0, v0] = T.Cast("float16", T.Cast("float32", B[v0]) * (T.Cast("float32", A[0, 0, v0]) / T.sqrt(Ared_temp[0, 0] * T.float32(0.000244140625) + T.float32(9.9999999999999995e-07))))

    @I.ir_module(s_tir=True)
    class After:
        @T.prim_func(s_tir=True)
        def main(A_handle: T.handle, B_handle: T.handle, rms_norm_handle: T.handle):
            T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
            A = T.match_buffer(A_handle, (1, 1, 4096), "float16")
            B = T.match_buffer(B_handle, (4096,), "float16")
            rms_norm = T.match_buffer(rms_norm_handle, (1, 4096), "float16")
            with T.sblock("root"):
                T.reads()
                T.writes()
                Ared_temp_shared = T.sblock_alloc_buffer((1, 1), scope="shared")
                Ared_temp_rf_local = T.sblock_alloc_buffer((1024, 1, 1), scope="local")
                for ax0_fused in T.thread_binding(T.int64(1), thread="blockIdx.x"):
                    for ax1_fused_1 in T.thread_binding(1024, thread="threadIdx.x", annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
                        with T.sblock("Ared_temp_rf_init"):
                            vax1_fused_1 = T.axis.spatial(1024, ax1_fused_1)
                            v0 = T.axis.spatial(T.int64(1), T.int64(0))
                            T.reads()
                            T.writes(Ared_temp_rf_local[vax1_fused_1, 0, 0])
                            Ared_temp_rf_local[vax1_fused_1, 0, 0] = T.float32(0)
                        for ax1_fused_0 in range(4):
                            for u in range(1):
                                with T.sblock("Ared_temp_rf_update"):
                                    vax1_fused_1 = T.axis.spatial(1024, ax1_fused_1)
                                    v0 = T.axis.spatial(T.int64(1), T.int64(0))
                                    vax1_fused_0 = T.axis.reduce(4, ax1_fused_0)
                                    T.reads(Ared_temp_rf_local[vax1_fused_1, 0, 0], A[0, 0, vax1_fused_0 * 1024 + vax1_fused_1])
                                    T.writes(Ared_temp_rf_local[vax1_fused_1, 0, 0])
                                    Ared_temp_rf_local[vax1_fused_1, 0, 0] = Ared_temp_rf_local[vax1_fused_1, 0, 0] + T.Cast("float32", A[0, 0, vax1_fused_0 * 1024 + vax1_fused_1]) * T.Cast("float32", A[0, 0, vax1_fused_0 * 1024 + vax1_fused_1])
                    for ax1_fused in range(T.int64(1)):
                        for ax0 in T.thread_binding(1024, thread="threadIdx.x"):
                            with T.sblock("Ared_temp"):
                                vax1_fused_1 = T.axis.reduce(1024, ax0)
                                v0 = T.axis.spatial(T.int64(1), T.int64(0))
                                T.reads(Ared_temp_rf_local[vax1_fused_1, 0, 0])
                                T.writes(Ared_temp_shared[0, 0])
                                with T.init():
                                    Ared_temp_shared[0, 0] = T.float32(0)
                                Ared_temp_shared[0, 0] = Ared_temp_shared[0, 0] + Ared_temp_rf_local[vax1_fused_1, 0, 0]
                    for ax0_fused_0 in range(4):
                        for ax0_fused_1 in T.thread_binding(1024, thread="threadIdx.x"):
                            with T.sblock("rms_norm"):
                                v0 = T.axis.spatial(4096, ax0_fused_0 * 1024 + ax0_fused_1)
                                T.reads(B[v0], A[0, 0, v0], Ared_temp_shared[0, 0])
                                T.writes(rms_norm[0, v0])
                                rms_norm[0, v0] = T.Cast("float16", T.Cast("float32", B[v0]) * (T.Cast("float32", A[0, 0, v0]) / T.sqrt(Ared_temp_shared[0, 0] * T.float32(0.000244140625) + T.float32(9.9999999999999995e-07))))
    # fmt: on
    target = Target("nvidia/geforce-rtx-3090-ti")
    with target:
        mod = dl.ApplyDefaultSchedule(dl.gpu.Reduction())(Before)  # pylint: disable=not-callable
    assert_structural_equal(mod, After)


def test_spatial_inner_no_broadcasting():
    # fmt: off
    @I.ir_module(s_tir=True)
    class Module:
        @T.prim_func(s_tir=True)
        def main(lv575: T.Buffer((1376, 4096), "uint32"), lv576: T.Buffer((344, 4096), "float16"), lv574: T.Buffer((1, 1, 11008), "float16"), lv570: T.Buffer((1, 1, 4096), "float16"), p_output0_intermediate: T.Buffer((1, 1, 4096), "float16")):
            T.func_attr({"tirx.noalias": True})
            p_output0_intermediate_1 = T.sblock_alloc_buffer((11008, 4096), "float16")
            var_matmul_intermediate = T.sblock_alloc_buffer((1, 1, 4096), "float16")
            for i, j in T.grid(11008, 4096):
                with T.sblock("decode"):
                    v_i, v_j = T.axis.remap("SS", [i, j])
                    T.reads(lv575[v_i // 8, v_j], lv576[v_i // 32, v_j])
                    T.writes(p_output0_intermediate_1[v_i, v_j])
                    p_output0_intermediate_1[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(lv575[v_i // 8, v_j], T.Cast("uint32", v_i % 8) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * lv576[v_i // 32, v_j]
            for i0, i1, i2, k in T.grid(1, 1, 4096, 11008):
                with T.sblock("matmul"):
                    v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
                    T.reads(lv574[v_i0, v_i1, v_k], p_output0_intermediate_1[v_k, v_i2])
                    T.writes(var_matmul_intermediate[v_i0, v_i1, v_i2])
                    with T.init():
                        var_matmul_intermediate[v_i0, v_i1, v_i2] = T.float16(0)
                    var_matmul_intermediate[v_i0, v_i1, v_i2] = var_matmul_intermediate[v_i0, v_i1, v_i2] + lv574[v_i0, v_i1, v_k] * p_output0_intermediate_1[v_k, v_i2]
            for ax0, ax1, ax2 in T.grid(1, 1, 4096):
                with T.sblock("T_add"):
                    v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
                    T.reads(lv570[v_ax0, v_ax1, v_ax2], var_matmul_intermediate[v_ax0, v_ax1, v_ax2])
                    T.writes(p_output0_intermediate[v_ax0, v_ax1, v_ax2])
                    p_output0_intermediate[v_ax0, v_ax1, v_ax2] = lv570[v_ax0, v_ax1, v_ax2] + var_matmul_intermediate[v_ax0, v_ax1, v_ax2]

    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(s_tir=True)
        def main(lv575: T.Buffer((1376, 4096), "uint32"), lv576: T.Buffer((344, 4096), "float16"), lv574: T.Buffer((1, 1, 11008), "float16"), lv570: T.Buffer((1, 1, 4096), "float16"), p_output0_intermediate: T.Buffer((1, 1, 4096), "float16")):
            T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
            var_matmul_intermediate_local = T.sblock_alloc_buffer((1, 1, 4096), "float16", scope="local")
            var_matmul_intermediate_rf_local = T.sblock_alloc_buffer((16, 1, 1, 4096), "float16", scope="local")
            for ax0_fused_0 in T.thread_binding(256, thread="blockIdx.x"):
                for ax0_fused_1 in T.thread_binding(16, thread="threadIdx.x"):
                    for ax1_0_fused_1 in T.thread_binding(16, thread="threadIdx.y"):
                        with T.sblock("matmul_rf_init"):
                            vax1_0_fused_1 = T.axis.spatial(16, ax1_0_fused_1)
                            v0 = T.axis.spatial(4096, ax0_fused_0 * 16 + ax0_fused_1)
                            T.reads()
                            T.writes(var_matmul_intermediate_rf_local[vax1_0_fused_1, 0, 0, v0])
                            var_matmul_intermediate_rf_local[vax1_0_fused_1, 0, 0, v0] = T.float16(0)
                        for ax1_0_fused_0, ax1_1 in T.grid(86, 8):
                            with T.sblock("matmul_rf_update"):
                                vax1_0_fused_1 = T.axis.spatial(16, ax1_0_fused_1)
                                v0 = T.axis.spatial(4096, ax0_fused_0 * 16 + ax0_fused_1)
                                vax1_0_fused_0, vax1_1 = T.axis.remap("RR", [ax1_0_fused_0, ax1_1])
                                T.reads(var_matmul_intermediate_rf_local[vax1_0_fused_1, 0, 0, v0], lv574[0, 0, vax1_0_fused_0 * 128 + vax1_0_fused_1 * 8 + vax1_1], lv575[(vax1_0_fused_0 * 128 + vax1_0_fused_1 * 8 + vax1_1) // 8, v0], lv576[(vax1_0_fused_0 * 128 + vax1_0_fused_1 * 8 + vax1_1) // 32, v0])
                                T.writes(var_matmul_intermediate_rf_local[vax1_0_fused_1, 0, 0, v0])
                                var_matmul_intermediate_rf_local[vax1_0_fused_1, 0, 0, v0] = var_matmul_intermediate_rf_local[vax1_0_fused_1, 0, 0, v0] + lv574[0, 0, vax1_0_fused_0 * 128 + vax1_0_fused_1 * 8 + vax1_1] * ((T.Cast("float16", T.bitwise_and(T.shift_right(lv575[(vax1_0_fused_0 * 128 + vax1_0_fused_1 * 8 + vax1_1) // 8, v0], T.Cast("uint32", (vax1_0_fused_0 * 128 + vax1_0_fused_1 * 8 + vax1_1) % 8) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * lv576[(vax1_0_fused_0 * 128 + vax1_0_fused_1 * 8 + vax1_1) // 32, v0])
                for ax1_fused in T.thread_binding(16, thread="threadIdx.x"):
                    for ax0 in T.thread_binding(16, thread="threadIdx.y"):
                        with T.sblock("matmul"):
                            vax1_0_fused_1 = T.axis.reduce(16, ax0)
                            v0 = T.axis.spatial(4096, ax0_fused_0 * 16 + ax1_fused)
                            T.reads(var_matmul_intermediate_rf_local[vax1_0_fused_1, 0, 0, v0])
                            T.writes(var_matmul_intermediate_local[0, 0, v0])
                            with T.init():
                                var_matmul_intermediate_local[0, 0, v0] = T.float16(0)
                            var_matmul_intermediate_local[0, 0, v0] = var_matmul_intermediate_local[0, 0, v0] + var_matmul_intermediate_rf_local[vax1_0_fused_1, 0, 0, v0]
                for ax0_fused_0_1 in T.thread_binding(16, thread="threadIdx.x"):
                    for ax0_fused_1 in range(1):
                        with T.sblock("T_add"):
                            v0 = T.axis.spatial(4096, ax0_fused_0 * 16 + ax0_fused_0_1 + ax0_fused_1)
                            T.reads(lv570[0, 0, v0], var_matmul_intermediate_local[0, 0, v0])
                            T.writes(p_output0_intermediate[0, 0, v0])
                            p_output0_intermediate[0, 0, v0] = lv570[0, 0, v0] + var_matmul_intermediate_local[0, 0, v0]
    # fmt: on

    target = Target("nvidia/geforce-rtx-3090-ti")
    with target:
        mod = dl.ApplyDefaultSchedule(dl.gpu.Reduction())(Module)  # pylint: disable=not-callable
    assert_structural_equal(mod, Expected)


def test_spatial_inner_broadcasting():
    # fmt: off
    @I.ir_module(s_tir=True)
    class Module:
        @T.prim_func(s_tir=True)
        def main(A: T.Buffer((256, 256), "float32"), B: T.Buffer((256, 256), "float32")):
            T.func_attr({"tirx.noalias": True})
            temp_local = T.sblock_alloc_buffer((256,))
            for j in T.serial(256):
                for k in T.serial(256):
                    with T.sblock("sum"):
                        vj, vk = T.axis.remap("SR", [j, k])
                        T.reads(A[vk, vj])
                        T.writes(temp_local[vj])
                        with T.init():
                            temp_local[vj] = T.float32(0)
                        temp_local[vj] = temp_local[vj] + A[vk, vj]
            for i, j in T.grid(256, 256):
                with T.sblock("add"):
                    vi, vj = T.axis.remap("SS", [i, j])
                    T.reads(temp_local[vj])
                    T.writes(B[vi, vj])
                    B[vi, vj] = A[vi, vj] + temp_local[vj]

    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(s_tir=True)
        def main(A: T.Buffer((256, 256), "float32"), B: T.Buffer((256, 256), "float32")):
            T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
            temp_local_shared = T.sblock_alloc_buffer((256,), scope="shared")
            temp_local_rf_local = T.sblock_alloc_buffer((16, 256), scope="local")
            for ax0_fused_0 in T.thread_binding(16, thread="blockIdx.x"):
                for ax0_fused_1 in T.thread_binding(16, thread="threadIdx.x"):
                    for ax1_fused_1 in T.thread_binding(16, thread="threadIdx.y"):
                        with T.sblock("sum_rf_init"):
                            vax1_fused_1 = T.axis.spatial(16, ax1_fused_1)
                            v0 = T.axis.spatial(256, ax0_fused_0 * 16 + ax0_fused_1)
                            T.reads()
                            T.writes(temp_local_rf_local[vax1_fused_1, v0])
                            temp_local_rf_local[vax1_fused_1, v0] = T.float32(0)
                        for ax1_fused_0, u in T.grid(16, 1):
                            with T.sblock("sum_rf_update"):
                                vax1_fused_1 = T.axis.spatial(16, ax1_fused_1)
                                v0 = T.axis.spatial(256, ax0_fused_0 * 16 + ax0_fused_1)
                                vax1_fused_0 = T.axis.reduce(16, ax1_fused_0)
                                T.reads(temp_local_rf_local[vax1_fused_1, v0], A[vax1_fused_0 * 16 + vax1_fused_1, v0])
                                T.writes(temp_local_rf_local[vax1_fused_1, v0])
                                temp_local_rf_local[vax1_fused_1, v0] = temp_local_rf_local[vax1_fused_1, v0] + A[vax1_fused_0 * 16 + vax1_fused_1, v0]
                for ax1_fused in T.thread_binding(16, thread="threadIdx.x"):
                    for ax0 in T.thread_binding(16, thread="threadIdx.y"):
                        with T.sblock("sum"):
                            vax1_fused_1 = T.axis.reduce(16, ax0)
                            v0 = T.axis.spatial(256, ax0_fused_0 * 16 + ax1_fused)
                            T.reads(temp_local_rf_local[vax1_fused_1, v0])
                            T.writes(temp_local_shared[v0])
                            with T.init():
                                temp_local_shared[v0] = T.float32(0)
                            temp_local_shared[v0] = temp_local_shared[v0] + temp_local_rf_local[vax1_fused_1, v0]
                for ax0_ax1_fused_0 in range(16):
                    for ax0_ax1_fused_1 in T.thread_binding(16, thread="threadIdx.x"):
                        for ax0_ax1_fused_2 in T.thread_binding(16, thread="threadIdx.y"):
                            with T.sblock("add"):
                                v0 = T.axis.spatial(256, (ax0_ax1_fused_0 * 256 + ax0_ax1_fused_1 * 16 + ax0_ax1_fused_2) // 16)
                                v1 = T.axis.spatial(256, ax0_fused_0 * 16 + (ax0_ax1_fused_0 * 256 + ax0_ax1_fused_1 * 16 + ax0_ax1_fused_2) % 16)
                                T.reads(temp_local_shared[v1])
                                T.writes(B[v0, v1])
                                B[v0, v1] = A[v0, v1] + temp_local_shared[v1]
    # fmt: on

    target = Target("nvidia/geforce-rtx-3090-ti")
    with target:
        mod = dl.ApplyDefaultSchedule(dl.gpu.Reduction())(Module)  # pylint: disable=not-callable
    assert_structural_equal(mod, Expected)


def test_reduction_inner_no_broadcasting():
    # fmt: off
    @I.ir_module(s_tir=True)
    class Module:
        @T.prim_func(s_tir=True)
        def main(A: T.Buffer((256, 256), "float32"), B: T.Buffer((256,), "float32")):
            T.func_attr({"tirx.noalias": True})
            temp_local = T.sblock_alloc_buffer((256,))
            for i in T.serial(256):
                for k in T.serial(256):
                    with T.sblock("sum"):
                        vi, vk = T.axis.remap("SR", [i, k])
                        T.reads(A[vi, vk])
                        T.writes(temp_local[vi])
                        with T.init():
                            temp_local[vi] = T.float32(0)
                        temp_local[vi] = temp_local[vi] + A[vi, vk]
            for i in T.grid(256):
                with T.sblock("add"):
                    vi = T.axis.remap("S", [i])
                    T.reads(temp_local[vi])
                    T.writes(B[vi,])
                    B[vi] = temp_local[vi] + T.float32(1)

    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(s_tir=True)
        def main(A: T.Buffer((256, 256), "float32"), B: T.Buffer((256,), "float32")):
            T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
            # with T.sblock("root"):
            temp_local_local = T.sblock_alloc_buffer((256,), scope="local")
            temp_local_rf_local = T.sblock_alloc_buffer((256, 256), scope="local")
            for ax0_fused in T.thread_binding(256, thread="blockIdx.x"):
                for ax1_fused_1 in T.thread_binding(256, thread="threadIdx.x", annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
                    with T.sblock("sum_rf_init"):
                        vax1_fused_1, v0 = T.axis.remap("SS", [ax1_fused_1, ax0_fused])
                        T.reads()
                        T.writes(temp_local_rf_local[vax1_fused_1, v0])
                        temp_local_rf_local[vax1_fused_1, v0] = T.float32(0)
                    for ax1_fused_0, u in T.grid(1, 1):
                        with T.sblock("sum_rf_update"):
                            vax1_fused_1, v0, vax1_fused_0 = T.axis.remap("SSR", [ax1_fused_1, ax0_fused, ax1_fused_0])
                            T.reads(temp_local_rf_local[vax1_fused_1, v0], A[v0, vax1_fused_0 * 256 + vax1_fused_1])
                            T.writes(temp_local_rf_local[vax1_fused_1, v0])
                            temp_local_rf_local[vax1_fused_1, v0] = temp_local_rf_local[vax1_fused_1, v0] + A[v0, vax1_fused_0 * 256 + vax1_fused_1]
                for ax1_fused in range(1):
                    for ax0 in T.thread_binding(256, thread="threadIdx.x"):
                        with T.sblock("sum"):
                            vax1_fused_1, v0 = T.axis.remap("RS", [ax0, ax0_fused])
                            T.reads(temp_local_rf_local[vax1_fused_1, v0])
                            T.writes(temp_local_local[v0])
                            with T.init():
                                temp_local_local[v0] = T.float32(0)
                            temp_local_local[v0] = temp_local_local[v0] + temp_local_rf_local[vax1_fused_1, v0]
                for ax0 in range(1):
                    with T.sblock("add"):
                        v0 = T.axis.spatial(256, ax0_fused + ax0)
                        T.reads(temp_local_local[v0])
                        T.writes(B[v0])
                        B[v0] = temp_local_local[v0] + T.float32(1)
    # fmt: on

    target = Target("nvidia/geforce-rtx-3090-ti")
    with target:
        mod = dl.ApplyDefaultSchedule(dl.gpu.Reduction())(Module)  # pylint: disable=not-callable
    assert_structural_equal(mod, Expected)


def test_reduction_inner_no_broadcasting2():
    # fmt: off
    @I.ir_module(s_tir=True)
    class Module:
        @T.prim_func(s_tir=True)
        def main(lv9: T.Buffer((2560, 320), "uint32"), lv10: T.Buffer((2560, 80), "float16"), lv1: T.Buffer((1, 2560), "float16"), p_output0_intermediate: T.Buffer((1, 2560), "float32")):
            T.func_attr({"tirx.noalias": True})
            # with T.sblock("root"):
            p_output0_intermediate_1 = T.sblock_alloc_buffer((2560, 2560), "float16")
            var_matmul_intermediate = T.sblock_alloc_buffer((1, 2560), "float16")
            for i, j in T.grid(2560, 2560):
                with T.sblock("decode"):
                    v_i, v_j = T.axis.remap("SS", [i, j])
                    T.reads(lv9[v_i, v_j // 8], lv10[v_i, v_j // 32])
                    T.writes(p_output0_intermediate_1[v_i, v_j])
                    p_output0_intermediate_1[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(lv9[v_i, v_j // 8], T.Cast("uint32", v_j % 8) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * lv10[v_i, v_j // 32]
            for i0, i1, k in T.grid(1, 2560, 2560):
                with T.sblock("matmul"):
                    v_i0, v_i1, v_k = T.axis.remap("SSR", [i0, i1, k])
                    T.reads(lv1[v_i0, v_k], p_output0_intermediate_1[v_k, v_i1])
                    T.writes(var_matmul_intermediate[v_i0, v_i1])
                    with T.init():
                        var_matmul_intermediate[v_i0, v_i1] = T.float16(0)
                    var_matmul_intermediate[v_i0, v_i1] = var_matmul_intermediate[v_i0, v_i1] + lv1[v_i0, v_k] * p_output0_intermediate_1[v_k, v_i1]
            for i0, i1 in T.grid(1, 2560):
                with T.sblock("compute"):
                    v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
                    T.reads(var_matmul_intermediate[v_i0, v_i1])
                    T.writes(p_output0_intermediate[v_i0, v_i1])
                    p_output0_intermediate[v_i0, v_i1] = T.Cast("float32", var_matmul_intermediate[v_i0, v_i1])

    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(s_tir=True)
        def main(lv9: T.Buffer((2560, 320), "uint32"), lv10: T.Buffer((2560, 80), "float16"), lv1: T.Buffer((1, 2560), "float16"), p_output0_intermediate: T.Buffer((1, 2560), "float32")):
            T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
            # with T.sblock("root"):
            var_matmul_intermediate_local = T.sblock_alloc_buffer((1, 2560), "float16", scope="local")
            var_matmul_intermediate_rf_local = T.sblock_alloc_buffer((16, 1, 2560), "float16", scope="local")
            for ax0_0_fused_0 in T.thread_binding(20, thread="blockIdx.x"):
                for ax0_0_fused_1 in T.thread_binding(16, thread="threadIdx.x"):
                    for ax1_fused_1 in T.thread_binding(16, thread="threadIdx.y"):
                        for ax0_1_init in range(8):
                            with T.sblock("matmul_rf_init"):
                                vax1_fused_1 = T.axis.spatial(16, ax1_fused_1)
                                v0 = T.axis.spatial(2560, ax0_0_fused_0 * 128 + ax0_0_fused_1 * 8 + ax0_1_init)
                                T.reads()
                                T.writes(var_matmul_intermediate_rf_local[vax1_fused_1, 0, v0])
                                var_matmul_intermediate_rf_local[vax1_fused_1, 0, v0] = T.float16(0)
                        for ax1_fused_0, ax0_1 in T.grid(160, 8):
                            with T.sblock("matmul_rf_update"):
                                vax1_fused_1 = T.axis.spatial(16, ax1_fused_1)
                                v0 = T.axis.spatial(2560, ax0_0_fused_0 * 128 + ax0_0_fused_1 * 8 + ax0_1)
                                vax1_fused_0 = T.axis.reduce(160, ax1_fused_0)
                                T.reads(var_matmul_intermediate_rf_local[vax1_fused_1, 0, v0], lv1[0, vax1_fused_0 * 16 + vax1_fused_1], lv9[vax1_fused_0 * 16 + vax1_fused_1, v0 // 8], lv10[vax1_fused_0 * 16 + vax1_fused_1, v0 // 32])
                                T.writes(var_matmul_intermediate_rf_local[vax1_fused_1, 0, v0])
                                var_matmul_intermediate_rf_local[vax1_fused_1, 0, v0] = var_matmul_intermediate_rf_local[vax1_fused_1, 0, v0] + lv1[0, vax1_fused_0 * 16 + vax1_fused_1] * ((T.Cast("float16", T.bitwise_and(T.shift_right(lv9[vax1_fused_0 * 16 + vax1_fused_1, v0 // 8], T.Cast("uint32", v0 % 8) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * lv10[vax1_fused_0 * 16 + vax1_fused_1, v0 // 32])
                for ax1_fused_0 in T.thread_binding(16, thread="threadIdx.x"):
                    for ax1_fused_1 in range(8):
                        for ax0 in T.thread_binding(16, thread="threadIdx.y"):
                            with T.sblock("matmul"):
                                vax1_fused_1 = T.axis.reduce(16, ax0)
                                v0 = T.axis.spatial(2560, ax0_0_fused_0 * 128 + ax1_fused_0 * 8 + ax1_fused_1)
                                T.reads(var_matmul_intermediate_rf_local[vax1_fused_1, 0, v0])
                                T.writes(var_matmul_intermediate_local[0, v0])
                                with T.init():
                                    var_matmul_intermediate_local[0, v0] = T.float16(0)
                                var_matmul_intermediate_local[0, v0] = var_matmul_intermediate_local[0, v0] + var_matmul_intermediate_rf_local[vax1_fused_1, 0, v0]
                for ax0_fused_0 in T.thread_binding(16, thread="threadIdx.x"):
                    for ax0_fused_1 in range(8):
                        with T.sblock("compute"):
                            v0 = T.axis.spatial(2560, ax0_0_fused_0 * 128 + ax0_fused_0 * 8 + ax0_fused_1)
                            T.reads(var_matmul_intermediate_local[0, v0])
                            T.writes(p_output0_intermediate[0, v0])
                            p_output0_intermediate[0, v0] = T.Cast("float32", var_matmul_intermediate_local[0, v0])
    # fmt: on

    with Target("nvidia/geforce-rtx-3090-ti"):
        mod = dl.ApplyDefaultSchedule(dl.gpu.Reduction())(Module)  # pylint: disable=not-callable
    assert_structural_equal(mod, Expected)


def test_reduction_inner_spatial_choose_perfect_factor():
    # fmt: off
    @I.ir_module(s_tir=True)
    class Module:
        @T.prim_func(s_tir=True)
        def main(var_A: T.handle, var_B: T.handle, matmul: T.Buffer((T.int64(1), T.int64(32), T.int64(1), T.int64(100)), "float16")):
            T.func_attr({"tirx.noalias": True})
            n = T.int64()
            A = T.match_buffer(var_A, (T.int64(1), T.int64(32), T.int64(1), n), "float16")
            B = T.match_buffer(var_B, (T.int64(1), T.int64(32), n, T.int64(100)), "float16")
            # with T.sblock("root"):
            for i0, i1, i2, i3, k in T.grid(T.int64(1), T.int64(32), T.int64(1), T.int64(100), n):
                with T.sblock("matmul"):
                    v_i0, v_i1, v_i2, v_i3, v_k = T.axis.remap("SSSSR", [i0, i1, i2, i3, k])
                    T.reads(A[v_i0, v_i1, v_i2, v_k], B[v_i0, v_i1, v_k, v_i3])
                    T.writes(matmul[v_i0, v_i1, v_i2, v_i3])
                    with T.init():
                        matmul[v_i0, v_i1, v_i2, v_i3] = T.float16(0)
                    matmul[v_i0, v_i1, v_i2, v_i3] = matmul[v_i0, v_i1, v_i2, v_i3] + A[v_i0, v_i1, v_i2, v_k] * B[v_i0, v_i1, v_k, v_i3]
    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(s_tir=True)
        def main(var_A: T.handle, var_B: T.handle, matmul: T.Buffer((T.int64(1), T.int64(32), T.int64(1), T.int64(100)), "float16")):
            T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
            n = T.int64()
            A = T.match_buffer(var_A, (T.int64(1), T.int64(32), T.int64(1), n), "float16")
            B = T.match_buffer(var_B, (T.int64(1), T.int64(32), n, T.int64(100)), "float16")
            # with T.sblock("root"):
            matmul_rf_local = T.sblock_alloc_buffer((T.int64(16), T.int64(1), T.int64(32), T.int64(1), T.int64(100)), "float16", scope="local")
            for ax0_ax1_fused_0 in T.thread_binding(T.int64(320), thread="blockIdx.x"):
                for ax0_ax1_fused_1 in T.thread_binding(T.int64(10), thread="threadIdx.x"):
                    for ax2_fused_1 in T.thread_binding(T.int64(16), thread="threadIdx.y"):
                        with T.sblock("matmul_rf_init"):
                            vax2_fused_1 = T.axis.spatial(T.int64(16), ax2_fused_1)
                            v0 = T.axis.spatial(T.int64(32), (ax0_ax1_fused_0 * T.int64(10) + ax0_ax1_fused_1) // T.int64(100))
                            v1 = T.axis.spatial(T.int64(100), (ax0_ax1_fused_0 * T.int64(10) + ax0_ax1_fused_1) % T.int64(100))
                            T.reads()
                            T.writes(matmul_rf_local[vax2_fused_1, T.int64(0), v0, T.int64(0), v1])
                            matmul_rf_local[vax2_fused_1, T.int64(0), v0, T.int64(0), v1] = T.float16(0)
                        for ax2_fused_0, u in T.grid((n + T.int64(15)) // T.int64(16), 1):
                            with T.sblock("matmul_rf_update"):
                                vax2_fused_1 = T.axis.spatial(T.int64(16), ax2_fused_1)
                                v0 = T.axis.spatial(T.int64(32), (ax0_ax1_fused_0 * T.int64(10) + ax0_ax1_fused_1) // T.int64(100))
                                v1 = T.axis.spatial(T.int64(100), (ax0_ax1_fused_0 * T.int64(10) + ax0_ax1_fused_1) % T.int64(100))
                                vax2_fused_0 = T.axis.reduce((n + T.int64(15)) // T.int64(16), ax2_fused_0)
                                T.where(ax2_fused_0 * T.int64(16) + ax2_fused_1 < n)
                                T.reads(matmul_rf_local[vax2_fused_1, T.int64(0), v0, T.int64(0), v1], A[T.int64(0), v0, T.int64(0), vax2_fused_0 * T.int64(16) + vax2_fused_1], B[T.int64(0), v0, vax2_fused_0 * T.int64(16) + vax2_fused_1, v1])
                                T.writes(matmul_rf_local[vax2_fused_1, T.int64(0), v0, T.int64(0), v1])
                                matmul_rf_local[vax2_fused_1, T.int64(0), v0, T.int64(0), v1] = matmul_rf_local[vax2_fused_1, T.int64(0), v0, T.int64(0), v1] + A[T.int64(0), v0, T.int64(0), vax2_fused_0 * T.int64(16) + vax2_fused_1] * B[T.int64(0), v0, vax2_fused_0 * T.int64(16) + vax2_fused_1, v1]
                for ax1_ax2_fused in T.thread_binding(T.int64(10), thread="threadIdx.x"):
                    for ax0 in T.thread_binding(T.int64(16), thread="threadIdx.y"):
                        with T.sblock("matmul"):
                            vax2_fused_1 = T.axis.reduce(T.int64(16), ax0)
                            v0 = T.axis.spatial(T.int64(32), ax0_ax1_fused_0 // T.int64(10))
                            v1 = T.axis.spatial(T.int64(100), ax0_ax1_fused_0 % T.int64(10) * T.int64(10) + ax1_ax2_fused)
                            T.reads(matmul_rf_local[vax2_fused_1, T.int64(0), v0, T.int64(0), v1])
                            T.writes(matmul[T.int64(0), v0, T.int64(0), v1])
                            with T.init():
                                matmul[T.int64(0), v0, T.int64(0), v1] = T.float16(0)
                            matmul[T.int64(0), v0, T.int64(0), v1] = matmul[T.int64(0), v0, T.int64(0), v1] + matmul_rf_local[vax2_fused_1, T.int64(0), v0, T.int64(0), v1]
    # fmt: on

    with Target("nvidia/geforce-rtx-3090-ti"):
        mod = dl.ApplyDefaultSchedule(dl.gpu.Reduction())(Module)  # pylint: disable=not-callable
    assert_structural_equal(mod, Expected)


def test_repeat_transpose_gemv():
    # fmt: off

    @I.ir_module(s_tir=True)
    class Before:
        @T.prim_func(private=True, s_tir=True)
        def fused_relax_repeat_relax_permute_dims_relax_matmul1(p_lv716: T.handle, p_astype66: T.handle, var_matmul_intermediate: T.Buffer((T.int64(1), T.int64(32), T.int64(1), T.int64(128)), "float16")):
            T.func_attr({"tirx.noalias": True})
            kv_seq_len = T.int64()
            lv716 = T.match_buffer(p_lv716, (T.int64(1), kv_seq_len, T.int64(8), T.int64(128)), "float16")
            astype66 = T.match_buffer(p_astype66, (T.int64(1), T.int64(32), T.int64(1), kv_seq_len), "float16")
            # with T.sblock("root"):
            var_T_repeat_intermediate = T.sblock_alloc_buffer((T.int64(1), kv_seq_len, T.int64(32), T.int64(128)), "float16")
            var_T_transpose_intermediate = T.sblock_alloc_buffer((T.int64(1), T.int64(32), kv_seq_len, T.int64(128)), "float16")
            for ax0, ax1, ax2, ax3 in T.grid(T.int64(1), kv_seq_len, T.int64(32), T.int64(128)):
                with T.sblock("T_repeat"):
                    v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
                    T.reads(lv716[v_ax0, v_ax1, v_ax2 // T.int64(4), v_ax3])
                    T.writes(var_T_repeat_intermediate[v_ax0, v_ax1, v_ax2, v_ax3])
                    var_T_repeat_intermediate[v_ax0, v_ax1, v_ax2, v_ax3] = lv716[v_ax0, v_ax1, v_ax2 // T.int64(4), v_ax3]
            for ax0, ax1, ax2, ax3 in T.grid(T.int64(1), T.int64(32), kv_seq_len, T.int64(128)):
                with T.sblock("T_transpose"):
                    v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
                    T.reads(var_T_repeat_intermediate[v_ax0, v_ax2, v_ax1, v_ax3])
                    T.writes(var_T_transpose_intermediate[v_ax0, v_ax1, v_ax2, v_ax3])
                    var_T_transpose_intermediate[v_ax0, v_ax1, v_ax2, v_ax3] = var_T_repeat_intermediate[v_ax0, v_ax2, v_ax1, v_ax3]
            for i0, i1, i2, i3, k in T.grid(T.int64(1), T.int64(32), T.int64(1), T.int64(128), kv_seq_len):
                with T.sblock("matmul"):
                    v_i0, v_i1, v_i2, v_i3, v_k = T.axis.remap("SSSSR", [i0, i1, i2, i3, k])
                    T.reads(astype66[v_i0, v_i1, v_i2, v_k], var_T_transpose_intermediate[v_i0, v_i1, v_k, v_i3])
                    T.writes(var_matmul_intermediate[v_i0, v_i1, v_i2, v_i3])
                    with T.init():
                        var_matmul_intermediate[v_i0, v_i1, v_i2, v_i3] = T.float16(0)
                    var_matmul_intermediate[v_i0, v_i1, v_i2, v_i3] = var_matmul_intermediate[v_i0, v_i1, v_i2, v_i3] + astype66[v_i0, v_i1, v_i2, v_k] * var_T_transpose_intermediate[v_i0, v_i1, v_k, v_i3]
    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(private=True, s_tir=True)
        def fused_relax_repeat_relax_permute_dims_relax_matmul1(p_lv716: T.handle, p_astype66: T.handle, var_matmul_intermediate: T.Buffer((T.int64(1), T.int64(32), T.int64(1), T.int64(128)), "float16")):
            T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
            kv_seq_len = T.int64()
            lv716 = T.match_buffer(p_lv716, (T.int64(1), kv_seq_len, T.int64(8), T.int64(128)), "float16")
            astype66 = T.match_buffer(p_astype66, (T.int64(1), T.int64(32), T.int64(1), kv_seq_len), "float16")
            # with T.sblock("root"):
            var_matmul_intermediate_rf_local = T.sblock_alloc_buffer((T.int64(16), T.int64(1), T.int64(32), T.int64(1), T.int64(128)), "float16", scope="local")
            for ax0_0_ax1_fused_0 in T.thread_binding(T.int64(64), thread="blockIdx.x"):
                for ax0_0_ax1_fused_1 in T.thread_binding(T.int64(16), thread="threadIdx.x"):
                    for ax2_fused_1 in T.thread_binding(T.int64(16), thread="threadIdx.y"):
                        for ax0_1_init in range(T.int64(4)):
                            with T.sblock("matmul_rf_init"):
                                vax2_fused_1 = T.axis.spatial(T.int64(16), ax2_fused_1)
                                v0 = T.axis.spatial(T.int64(32), (ax0_0_ax1_fused_0 * T.int64(16) + ax0_0_ax1_fused_1) // T.int64(128) * T.int64(4) + ax0_1_init)
                                v1 = T.axis.spatial(T.int64(128), (ax0_0_ax1_fused_0 * T.int64(16) + ax0_0_ax1_fused_1) % T.int64(128))
                                T.reads()
                                T.writes(var_matmul_intermediate_rf_local[vax2_fused_1, T.int64(0), v0, T.int64(0), v1])
                                var_matmul_intermediate_rf_local[vax2_fused_1, T.int64(0), v0, T.int64(0), v1] = T.float16(0)
                        for ax2_fused_0, ax0_1 in T.grid((kv_seq_len + T.int64(15)) // T.int64(16), T.int64(4)):
                            with T.sblock("matmul_rf_update"):
                                vax2_fused_1 = T.axis.spatial(T.int64(16), ax2_fused_1)
                                v0 = T.axis.spatial(T.int64(32), (ax0_0_ax1_fused_0 * T.int64(16) + ax0_0_ax1_fused_1) // T.int64(128) * T.int64(4) + ax0_1)
                                v1 = T.axis.spatial(T.int64(128), (ax0_0_ax1_fused_0 * T.int64(16) + ax0_0_ax1_fused_1) % T.int64(128))
                                vax2_fused_0 = T.axis.reduce((kv_seq_len + T.int64(15)) // T.int64(16), ax2_fused_0)
                                T.where(ax2_fused_0 * T.int64(16) + ax2_fused_1 < kv_seq_len)
                                T.reads(var_matmul_intermediate_rf_local[vax2_fused_1, T.int64(0), v0, T.int64(0), v1], astype66[T.int64(0), v0, T.int64(0), vax2_fused_0 * T.int64(16) + vax2_fused_1], lv716[T.int64(0), vax2_fused_0 * T.int64(16) + vax2_fused_1, v0 // T.int64(4), v1])
                                T.writes(var_matmul_intermediate_rf_local[vax2_fused_1, T.int64(0), v0, T.int64(0), v1])
                                var_matmul_intermediate_rf_local[vax2_fused_1, T.int64(0), v0, T.int64(0), v1] = var_matmul_intermediate_rf_local[vax2_fused_1, T.int64(0), v0, T.int64(0), v1] + astype66[T.int64(0), v0, T.int64(0), vax2_fused_0 * T.int64(16) + vax2_fused_1] * lv716[T.int64(0), vax2_fused_0 * T.int64(16) + vax2_fused_1, v0 // T.int64(4), v1]
                for ax1_0_ax2_fused in T.thread_binding(T.int64(16), thread="threadIdx.x"):
                    for ax1_1 in range(T.int64(4)):
                        for ax0 in T.thread_binding(T.int64(16), thread="threadIdx.y"):
                            with T.sblock("matmul"):
                                vax2_fused_1 = T.axis.reduce(T.int64(16), ax0)
                                v0 = T.axis.spatial(T.int64(32), ax0_0_ax1_fused_0 // T.int64(8) * T.int64(4) + ax1_1)
                                v1 = T.axis.spatial(T.int64(128), ax0_0_ax1_fused_0 % T.int64(8) * T.int64(16) + ax1_0_ax2_fused)
                                T.reads(var_matmul_intermediate_rf_local[vax2_fused_1, T.int64(0), v0, T.int64(0), v1])
                                T.writes(var_matmul_intermediate[T.int64(0), v0, T.int64(0), v1])
                                with T.init():
                                    var_matmul_intermediate[T.int64(0), v0, T.int64(0), v1] = T.float16(0)
                                var_matmul_intermediate[T.int64(0), v0, T.int64(0), v1] = var_matmul_intermediate[T.int64(0), v0, T.int64(0), v1] + var_matmul_intermediate_rf_local[vax2_fused_1, T.int64(0), v0, T.int64(0), v1]
    # fmt: on

    with Target("nvidia/geforce-rtx-3090-ti"):
        mod = dl.ApplyDefaultSchedule(dl.gpu.Reduction())(Before)  # pylint: disable=not-callable
    assert_structural_equal(mod, Expected)


def test_gemv_dyn_shape_epilogue():
    @I.ir_module(s_tir=True)
    class Module:
        @T.prim_func(private=True, s_tir=True)
        def main(
            var_A: T.handle,
            B: T.Buffer((T.int64(1), T.int64(1), T.int64(4096)), "float16"),
            var_C: T.handle,
        ):
            T.func_attr({"tirx.noalias": True})
            vocab_size = T.int64()
            A = T.match_buffer(var_A, (T.int64(4096), vocab_size), "float16")
            C = T.match_buffer(var_C, (T.int64(1), T.int64(1), vocab_size))
            C_temp = T.sblock_alloc_buffer((T.int64(1), T.int64(1), vocab_size), "float16")
            for i0, i1, i2, k in T.grid(T.int64(1), T.int64(1), vocab_size, T.int64(4096)):
                with T.sblock("matmul"):
                    v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
                    T.reads(B[v_i0, v_i1, v_k], A[v_k, v_i2])
                    T.writes(C_temp[v_i0, v_i1, v_i2])
                    with T.init():
                        C_temp[v_i0, v_i1, v_i2] = T.float16(0)
                    C_temp[v_i0, v_i1, v_i2] = (
                        C_temp[v_i0, v_i1, v_i2] + B[v_i0, v_i1, v_k] * A[v_k, v_i2]
                    )
            for i0, i1, i2 in T.grid(T.int64(1), T.int64(1), vocab_size):
                with T.sblock("epilogue"):
                    v_i0, v_i1, v_i2 = T.axis.remap("SSS", [i0, i1, i2])
                    T.reads(C_temp[v_i0, v_i1, v_i2])
                    T.writes(C[v_i0, v_i1, v_i2])
                    C[v_i0, v_i1, v_i2] = T.Cast("float32", C_temp[v_i0, v_i1, v_i2])

    # fmt: off
    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(private=True, s_tir=True)
        def main(var_A: T.handle, B: T.Buffer((T.int64(1), T.int64(1), T.int64(4096)), "float16"), var_C: T.handle):
            T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
            vocab_size = T.int64()
            A = T.match_buffer(var_A, (T.int64(4096), vocab_size), "float16")
            C = T.match_buffer(var_C, (T.int64(1), T.int64(1), vocab_size))
            # with T.sblock("root"):
            C_temp_local = T.sblock_alloc_buffer((T.int64(1), T.int64(1), vocab_size), "float16", scope="local")
            C_temp_rf_local = T.sblock_alloc_buffer((T.int64(16), T.int64(1), T.int64(1), vocab_size), "float16", scope="local")
            for ax0_fused_0 in T.thread_binding(vocab_size, thread="blockIdx.x"):
                for ax0_fused_1 in T.thread_binding(T.int64(1), thread="threadIdx.x"):
                    for ax1_fused_1 in T.thread_binding(T.int64(16), thread="threadIdx.y"):
                        with T.sblock("matmul_rf_init"):
                            vax1_fused_1 = T.axis.spatial(T.int64(16), ax1_fused_1)
                            v0 = T.axis.spatial(vocab_size, ax0_fused_0 + ax0_fused_1)
                            T.reads()
                            T.writes(C_temp_rf_local[vax1_fused_1, T.int64(0), T.int64(0), v0])
                            C_temp_rf_local[vax1_fused_1, T.int64(0), T.int64(0), v0] = T.float16(0)
                        for ax1_fused_0, u in T.grid(T.int64(256), 1):
                            with T.sblock("matmul_rf_update"):
                                vax1_fused_1 = T.axis.spatial(T.int64(16), ax1_fused_1)
                                v0 = T.axis.spatial(vocab_size, ax0_fused_0 + ax0_fused_1)
                                vax1_fused_0 = T.axis.reduce(T.int64(256), ax1_fused_0)
                                T.reads(C_temp_rf_local[vax1_fused_1, T.int64(0), T.int64(0), v0], B[T.int64(0), T.int64(0), vax1_fused_0 * T.int64(16) + vax1_fused_1], A[vax1_fused_0 * T.int64(16) + vax1_fused_1, v0])
                                T.writes(C_temp_rf_local[vax1_fused_1, T.int64(0), T.int64(0), v0])
                                C_temp_rf_local[vax1_fused_1, T.int64(0), T.int64(0), v0] = C_temp_rf_local[vax1_fused_1, T.int64(0), T.int64(0), v0] + B[T.int64(0), T.int64(0), vax1_fused_0 * T.int64(16) + vax1_fused_1] * A[vax1_fused_0 * T.int64(16) + vax1_fused_1, v0]
                for ax1_fused in T.thread_binding(T.int64(1), thread="threadIdx.x"):
                    for ax0 in T.thread_binding(T.int64(16), thread="threadIdx.y"):
                        with T.sblock("matmul"):
                            vax1_fused_1, v0 = T.axis.remap("RS", [ax0, ax0_fused_0])
                            T.reads(C_temp_rf_local[vax1_fused_1, T.int64(0), T.int64(0), v0])
                            T.writes(C_temp_local[T.int64(0), T.int64(0), v0])
                            with T.init():
                                C_temp_local[T.int64(0), T.int64(0), v0] = T.float16(0)
                            C_temp_local[T.int64(0), T.int64(0), v0] = C_temp_local[T.int64(0), T.int64(0), v0] + C_temp_rf_local[vax1_fused_1, T.int64(0), T.int64(0), v0]
                for ax0_fused_0_1 in T.thread_binding(T.int64(1), thread="threadIdx.x"):
                    for ax0_fused_1 in range(T.int64(1)):
                        with T.sblock("epilogue"):
                            v0 = T.axis.spatial(vocab_size, ax0_fused_0)
                            T.reads(C_temp_local[T.int64(0), T.int64(0), v0])
                            T.writes(C[T.int64(0), T.int64(0), v0])
                            C[T.int64(0), T.int64(0), v0] = T.Cast("float32", C_temp_local[T.int64(0), T.int64(0), v0])
    # fmt: on

    with Target("nvidia/geforce-rtx-3090-ti"):
        mod = dl.ApplyDefaultSchedule(dl.gpu.Reduction())(Module)  # pylint: disable=not-callable
    assert_structural_equal(mod, Expected)


def test_gemv_output_one_element():
    # fmt: off
    @I.ir_module(s_tir=True)
    class Before:
        @T.prim_func(private=True, s_tir=True)
        def main(A: T.Buffer((T.int64(1), T.int64(2048)), "float16"), weight: T.Buffer((T.int64(1), T.int64(2048)), "float16"), out: T.Buffer((T.int64(1), T.int64(1)), "float16")):
            T.func_attr({"tirx.noalias": True})
            NT_matmul_intermediate = T.sblock_alloc_buffer((T.int64(1), T.int64(1)), "float16")
            for i0, i1, k in T.grid(T.int64(1), T.int64(1), T.int64(2048)):
                with T.sblock("NT_matmul"):
                    v_i0, v_i1, v_k = T.axis.remap("SSR", [i0, i1, k])
                    with T.init():
                        NT_matmul_intermediate[v_i0, v_i1] = T.float16(0)
                    NT_matmul_intermediate[v_i0, v_i1] = NT_matmul_intermediate[v_i0, v_i1] + A[v_i0, v_k] * weight[v_i1, v_k]
            for i0, i1 in T.grid(T.int64(1), T.int64(1)):
                with T.sblock("compute"):
                    v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
                    out[v_i0, v_i1] = T.sigmoid(NT_matmul_intermediate[v_i0, v_i1])


    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(private=True, s_tir=True)
        def main(A: T.Buffer((T.int64(1), T.int64(2048)), "float16"), weight: T.Buffer((T.int64(1), T.int64(2048)), "float16"), out: T.Buffer((T.int64(1), T.int64(1)), "float16")):
            T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
            NT_matmul_intermediate_shared = T.sblock_alloc_buffer((T.int64(1), T.int64(1)), "float16", scope="shared")
            NT_matmul_intermediate_rf_local = T.sblock_alloc_buffer((T.int64(1024), T.int64(1), T.int64(1)), "float16", scope="local")
            for ax0_fused in T.thread_binding(T.int64(1), thread="blockIdx.x"):
                for ax1_fused_1 in T.thread_binding(T.int64(1024), thread="threadIdx.x", annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
                    with T.sblock("NT_matmul_rf_init"):
                        vax1_fused_1 = T.axis.spatial(T.int64(1024), ax1_fused_1)
                        v0 = T.axis.spatial(T.int64(1), T.int64(0))
                        NT_matmul_intermediate_rf_local[vax1_fused_1, T.int64(0), T.int64(0)] = T.float16(0)
                    for ax1_fused_0, u in T.grid(T.int64(2), 1):
                        with T.sblock("NT_matmul_rf_update"):
                            vax1_fused_1 = T.axis.spatial(T.int64(1024), ax1_fused_1)
                            v0 = T.axis.spatial(T.int64(1), T.int64(0))
                            vax1_fused_0 = T.axis.reduce(T.int64(2), ax1_fused_0)
                            NT_matmul_intermediate_rf_local[vax1_fused_1, T.int64(0), T.int64(0)] = NT_matmul_intermediate_rf_local[vax1_fused_1, T.int64(0), T.int64(0)] + A[T.int64(0), vax1_fused_0 * T.int64(1024) + vax1_fused_1] * weight[T.int64(0), vax1_fused_0 * T.int64(1024) + vax1_fused_1]
                for ax1_fused in range(T.int64(1)):
                    for ax0 in T.thread_binding(T.int64(1024), thread="threadIdx.x"):
                        with T.sblock("NT_matmul"):
                            vax1_fused_1 = T.axis.reduce(T.int64(1024), ax0)
                            v0 = T.axis.spatial(T.int64(1), T.int64(0))
                            with T.init():
                                NT_matmul_intermediate_shared[T.int64(0), T.int64(0)] = T.float16(0)
                            NT_matmul_intermediate_shared[T.int64(0), T.int64(0)] = NT_matmul_intermediate_shared[T.int64(0), T.int64(0)] + NT_matmul_intermediate_rf_local[vax1_fused_1, T.int64(0), T.int64(0)]
                for ax0_fused_0 in range(T.int64(1)):
                    for ax0_fused_1 in T.thread_binding(T.int64(1024), thread="threadIdx.x"):
                        with T.sblock("compute"):
                            v0 = T.axis.spatial(T.int64(1), T.int64(0))
                            T.where(ax0_fused_0 * T.int64(1024) + ax0_fused_1 < T.int64(1))
                            out[T.int64(0), T.int64(0)] = T.sigmoid(NT_matmul_intermediate_shared[T.int64(0), T.int64(0)])
    # fmt: on

    with Target("nvidia/geforce-rtx-3090-ti"):
        mod = dl.ApplyDefaultSchedule(dl.gpu.Reduction())(Before)  # pylint: disable=not-callable
    assert_structural_equal(mod, Expected)


def test_no_reduction_loop_check():
    # The normalized prime func will not contain a reduction loop since its extent is one.
    # This checks that the Reduction schedule is correctly not applied in this case
    # fmt: off
    @I.ir_module(s_tir=True)
    class Before:
        @T.prim_func(private=True, s_tir=True)
        def matmul(lv43: T.Buffer((T.int64(1), T.int64(32), T.int64(1)), "float16"), lv44: T.Buffer((T.int64(1), T.int64(1), T.int64(1)), "float16"), matmul: T.Buffer((T.int64(1), T.int64(32), T.int64(1)), "float16")):
            T.func_attr({"op_pattern": 4, "tirx.noalias": True})
            # with T.sblock("root"):
            for i0, i1, i2, k in T.grid(T.int64(1), T.int64(32), T.int64(1), T.int64(1)):
                with T.sblock("matmul"):
                    v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
                    T.reads(lv43[v_i0, v_i1, v_k], lv44[v_i0, v_k, v_i2])
                    T.writes(matmul[v_i0, v_i1, v_i2])
                    with T.init():
                        matmul[v_i0, v_i1, v_i2] = T.float16(0.0)
                    matmul[v_i0, v_i1, v_i2] = matmul[v_i0, v_i1, v_i2] + lv43[v_i0, v_i1, v_k] * lv44[v_i0, v_k, v_i2]
    # fmt: on

    target = Target("nvidia/geforce-rtx-3090-ti")
    with target:
        mod = dl.ApplyDefaultSchedule(dl.gpu.Reduction())(Before)  # pylint: disable=not-callable
    assert_structural_equal(mod, Before)


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