# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
# pylint: disable=missing-function-docstring,missing-module-docstring
# ruff: noqa: E501, F401
import pytest

import tvm
import tvm.testing
from tvm import te, tirx
from tvm.s_tir.schedule.testing import (
    assert_structural_equal_ignore_global_symbol,
    verify_trace_roundtrip,
)
from tvm.script import tirx as T

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

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


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


@T.prim_func(s_tir=True)
def blockized_1(a: T.handle, c: T.handle) -> None:
    A = T.match_buffer(a, [128, 128], "float32")
    B = T.sblock_alloc_buffer([128, 128], "float32")
    C = T.match_buffer(c, [128, 128], "float32")
    for i, j in T.grid(128, 128):
        with T.sblock("B"):
            vi, vj = T.axis.remap("SS", [i, j])
            B[vi, vj] = A[vi, vj] * 2.0
    for i, j in T.grid(8, 8):
        with T.sblock("C_outer"):
            vi_o, vj_o = T.axis.remap("SS", [i, j])
            T.reads([B[
                vi_o * 16 : vi_o * 16 + 16,
                vj_o * 16 : vj_o * 16 + 16,
            ]])
            T.writes([C[
                vi_o * 16 : vi_o * 16 + 16,
                vj_o * 16 : vj_o * 16 + 16
            ]])
            for i_i, j_i in T.grid(16, 16):
                with T.sblock("C_inner"):
                    vi = T.axis.S(128, vi_o * 16 + i_i)
                    vj = T.axis.S(128, vj_o * 16 + j_i)
                    C[vi, vj] = B[vi, vj] + 1.0


@T.prim_func(s_tir=True)
def blockized_after_compute_at(a: T.handle, c: T.handle) -> None:
    A = T.match_buffer(a, [128, 128], "float32")
    B = T.sblock_alloc_buffer([128, 128], "float32")
    C = T.match_buffer(c, [128, 128], "float32")
    for i0_0, i1_0 in T.grid(8, 8):
        for ax0, ax1 in T.grid(16, 16):
            with T.sblock("B"):
                vi = T.axis.S(128, i0_0 * 16 + ax0)
                vj = T.axis.S(128, i1_0 * 16 + ax1)
                B[vi, vj] = A[vi, vj] * 2.0
        with T.sblock("C_outer"):
            vi_o, vj_o = T.axis.remap("SS", [i0_0, i1_0])
            T.reads([B[
                vi_o * 16 : vi_o * 16 + 16,
                vj_o * 16 : vj_o * 16 + 16,
            ]])
            T.writes([C[
                vi_o * 16 : vi_o * 16 + 16,
                vj_o * 16 : vj_o * 16 + 16
            ]])
            for i0_1, i1_1 in T.grid(16, 16):
                with T.sblock("C_inner"):
                    vi = T.axis.S(128, vi_o * 16 + i0_1)
                    vj = T.axis.S(128, vj_o * 16 + i1_1)
                    C[vi, vj] = B[vi, vj] + 1.0


@T.prim_func(s_tir=True)
def blockized_2(a: T.handle, c: T.handle) -> None:
    A = T.match_buffer(a, [128, 128], "float32")
    B = T.sblock_alloc_buffer([128, 128], "float32")
    C = T.match_buffer(c, [128, 128], "float32")
    for i_o, j_o in T.grid(8, 8):
        with T.sblock("B_outer"):
            vio, vjo = T.axis.remap("SS", [i_o, j_o])
            T.reads([A[
                vio * 16 : vio * 16 + 16,
                vjo * 16 : vjo * 16 + 16,
            ]])
            T.writes([B[
                vio * 16 : vio * 16 + 16,
                vjo * 16 : vjo * 16 + 16
            ]])
            for i_i, j_i in T.grid(16, 16):
                with T.sblock("B_inner"):
                    vi = T.axis.S(128, vio * 16 + i_i)
                    vj = T.axis.S(128, vjo * 16 + j_i)
                    B[vi, vj] = A[vi, vj] * 2.0
    for i_o, j_o, i_i, j_i in T.grid(4, 4, 32, 32):
        with T.sblock("C"):
            vi = T.axis.S(128, i_o * 32 + i_i)
            vj = T.axis.S(128, j_o * 32 + j_i)
            C[vi, vj] = B[vi, vj] + 1.0


@T.prim_func(s_tir=True)
def blockized_2_after_reverse_compute_at(a: T.handle, c: T.handle) -> None:
    A = T.match_buffer(a, [128, 128], "float32")
    B = T.sblock_alloc_buffer([128, 128], "float32")
    C = T.match_buffer(c, [128, 128], "float32")
    for i_o, j_o in T.grid(8, 8):
        with T.sblock("B_outer"):
            vio, vjo = T.axis.remap("SS", [i_o, j_o])
            T.reads([A[
                vio * 16 : vio * 16 + 16,
                vjo * 16 : vjo * 16 + 16,
            ]])
            T.writes([B[
                vio * 16 : vio * 16 + 16,
                vjo * 16 : vjo * 16 + 16
            ]])
            for i_i, j_i in T.grid(16, 16):
                with T.sblock("B_inner"):
                    vi = T.axis.S(128, vio * 16 + i_i)
                    vj = T.axis.S(128, vjo * 16 + j_i)
                    B[vi, vj] = A[vi, vj] * 2.0
        for ax0, ax1 in T.grid(16, 16):
            with T.sblock("C"):
                vi = T.axis.S(128, i_o * 16 + ax0)
                vj = T.axis.S(128, j_o * 16 + ax1)
                T.reads([B[vi, vj]])
                T.writes([C[vi, vj]])
                C[vi, vj] = B[vi, vj] + 1.0


@T.prim_func(s_tir=True)
def blockized_2_after_compute_at(a: T.handle, c: T.handle) -> None:
    A = T.match_buffer(a, [128, 128], "float32")
    B = T.sblock_alloc_buffer([128, 128], "float32")
    C = T.match_buffer(c, [128, 128], "float32")
    for i_o, j_o in T.grid(4, 4):
        for ax0, ax1 in T.grid(2, 2):
            with T.sblock("blockized_B"):
                vio = T.axis.S(8, i_o * 2 + ax0)
                vjo = T.axis.S(8, j_o * 2 + ax1)
                T.reads([A[
                    vio * 16 : vio * 16 + 16,
                    vjo * 16 : vjo * 16 + 16,
                ]])
                T.writes([B[
                    vio * 16 : vio * 16 + 16,
                    vjo * 16 : vjo * 16 + 16,
                ]])
                for i_i, j_i in T.grid(16, 16):
                    with T.sblock("B"):
                        vi = T.axis.S(128, vio * 16 + i_i)
                        vj = T.axis.S(128, vjo * 16 + j_i)
                        B[vi, vj] = A[vi, vj] * 2.0
        for i_i, j_i in T.grid(32, 32):
            with T.sblock("C"):
                vi = T.axis.S(128, i_o * 32 + i_i)
                vj = T.axis.S(128, j_o * 32 + j_i)
                C[vi, vj] = B[vi, vj] + 1.0

@T.prim_func(s_tir=True)
def cuda_matmul_0(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")
    A_shared = T.sblock_alloc_buffer([2048, 2048], "float32", scope="shared")
    B_shared = T.sblock_alloc_buffer([2048, 2048], "float32", scope="shared")
    A_shared_local = T.sblock_alloc_buffer([2048, 2048], "float32", scope="local")
    B_shared_local = T.sblock_alloc_buffer([2048, 2048], "float32", scope="local")
    C_local = T.sblock_alloc_buffer([2048, 2048], "float32", scope="local")
    for i, j in T.grid(2048, 2048):
        with T.sblock("A_shared"):
            v0, v1 = T.axis.remap("SS", [i, j])
            A_shared[v0, v1] = A[v0, v1]
    for i, j in T.grid(2048, 2048):
        with T.sblock("B_shared"):
            v0, v1 = T.axis.remap("SS", [i, j])
            B_shared[v0, v1] = B[v0, v1]
    for i, j in T.grid(2048, 2048):
        with T.sblock("A_shared_local"):
            v0, v1 = T.axis.remap("SS", [i, j])
            A_shared_local[v0, v1] = A_shared[v0, v1]
    for i, j in T.grid(2048, 2048):
        with T.sblock("B_shared_local"):
            v0, v1 = T.axis.remap("SS", [i, j])
            B_shared_local[v0, v1] = B_shared[v0, v1]
    for i, j, k in T.grid(2048, 2048, 2048):
        with T.sblock("C"):
            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
            with T.init():
                C_local[vi, vj] = 0.0
            C_local[vi, vj] = C_local[vi, vj] + A_shared_local[vk, vi] * B_shared_local[vk, vj]
    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 i, j in T.grid(4, 4):
                                with T.sblock("C_local"):
                                    v0_4 = T.axis.S(2048, by * 64 + vy * 32 + ty * 4 + i)
                                    v1_4 = T.axis.S(2048, bx * 64 + vx * 32 + tx * 4 + j)
                                    C[v0_4, v1_4] = C_local[v0_4, v1_4]


@T.prim_func(s_tir=True)
def cuda_matmul_0_after_compute_at(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")
    A_shared = T.sblock_alloc_buffer([2048, 2048], "float32", scope="shared")
    B_shared = T.sblock_alloc_buffer([2048, 2048], "float32", scope="shared")
    A_shared_local = T.sblock_alloc_buffer([2048, 2048], "float32", scope="local")
    B_shared_local = T.sblock_alloc_buffer([2048, 2048], "float32", scope="local")
    C_local = T.sblock_alloc_buffer([2048, 2048], "float32", scope="local")
    for i, j in T.grid(2048, 2048):
        with T.sblock("A_shared"):
            v0, v1 = T.axis.remap("SS", [i, j])
            A_shared[v0, v1] = A[v0, v1]
    for i, j in T.grid(2048, 2048):
        with T.sblock("B_shared"):
            v0, v1 = T.axis.remap("SS", [i, j])
            B_shared[v0, v1] = B[v0, v1]
    for i, j in T.grid(2048, 2048):
        with T.sblock("A_shared_local"):
            v0, v1 = T.axis.remap("SS", [i, j])
            A_shared_local[v0, v1] = A_shared[v0, v1]
    for i, j in T.grid(2048, 2048):
        with T.sblock("B_shared_local"):
            v0, v1 = T.axis.remap("SS", [i, j])
            B_shared_local[v0, v1] = B_shared[v0, v1]
    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 i, j, k in T.grid(4, 4, 2048):
                                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, k)
                                    with T.init():
                                        C_local[vi, vj] = 0.0
                                    C_local[vi, vj] = C_local[vi, vj] + A_shared_local[vk, vi] * B_shared_local[vk, vj]
                            for i, j in T.grid(4, 4):
                                with T.sblock("C_local"):
                                    vi = T.axis.S(2048, by * 64 + vy * 32 + ty * 4 + i)
                                    vj = T.axis.S(2048, bx * 64 + vx * 32 + tx * 4 + j)
                                    C[vi, vj] = C_local[vi, vj]


@T.prim_func(s_tir=True)
def cuda_matmul_1(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")
    A_shared = T.sblock_alloc_buffer([2048, 2048], "float32", scope="shared")
    B_shared = T.sblock_alloc_buffer([2048, 2048], "float32", scope="shared")
    A_shared_local = T.sblock_alloc_buffer([2048, 2048], "float32", scope="local")
    B_shared_local = T.sblock_alloc_buffer([2048, 2048], "float32", scope="local")
    C_local = T.sblock_alloc_buffer([2048, 2048], "float32", scope="local")
    for i, j in T.grid(2048, 2048):
        with T.sblock("A_shared"):
            v0, v1 = T.axis.remap("SS", [i, j])
            A_shared[v0, v1] = A[v0, v1]
    for i, j in T.grid(2048, 2048):
        with T.sblock("B_shared"):
            v0, v1 = T.axis.remap("SS", [i, j])
            B_shared[v0, v1] = B[v0, v1]
    for i, j in T.grid(2048, 2048):
        with T.sblock("A_shared_local"):
            v0, v1 = T.axis.remap("SS", [i, j])
            A_shared_local[v0, v1] = A_shared[v0, v1]
    for i, j in T.grid(2048, 2048):
        with T.sblock("B_shared_local"):
            v0, v1 = T.axis.remap("SS", [i, j])
            B_shared_local[v0, v1] = B_shared[v0, v1]
    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 k_0 in T.serial(0, 256):
                                for k_1 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, k_0 * 8 + k_1)
                                            with T.init():
                                                C_local[vi, vj] = 0.0
                                            C_local[vi, vj] = C_local[vi, vj] + A_shared_local[vk, vi] * B_shared_local[vk, vj]
                            for i, j in T.grid(4, 4):
                                with T.sblock("C_local"):
                                    vi = T.axis.S(2048, by * 64 + vy * 32 + ty * 4 + i)
                                    vj = T.axis.S(2048, bx * 64 + vx * 32 + tx * 4 + j)
                                    C[vi, vj] = C_local[vi, vj]


@T.prim_func(s_tir=True)
def cuda_matmul_2(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")
    A_shared = T.sblock_alloc_buffer([2048, 2048], "float32", scope="shared")
    B_shared = T.sblock_alloc_buffer([2048, 2048], "float32", scope="shared")
    A_shared_local = T.sblock_alloc_buffer([2048, 2048], "float32", scope="local")
    B_shared_local = T.sblock_alloc_buffer([2048, 2048], "float32", scope="local")
    C_local = T.sblock_alloc_buffer([2048, 2048], "float32", scope="local")
    for i, j in T.grid(2048, 2048):
        with T.sblock("A_shared"):
            v0, v1 = T.axis.remap("SS", [i, j])
            A_shared[v0, v1] = A[v0, v1]
    for i, j in T.grid(2048, 2048):
        with T.sblock("B_shared"):
            v0, v1 = T.axis.remap("SS", [i, j])
            B_shared[v0, v1] = B[v0, v1]
    for i, j in T.grid(2048, 2048):
        with T.sblock("B_shared_local"):
            v0, v1 = T.axis.remap("SS", [i, j])
            B_shared_local[v0, v1] = B_shared[v0, v1]
    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 k_0 in T.serial(0, 256):
                                for k_1 in T.unroll(0, 8):
                                    for i, j in T.grid(1, 4):
                                        with T.sblock("A_shared_local"):
                                            v0 = T.axis.S(2048, k_0 * 8 + k_1 + i)
                                            v1 = T.axis.S(2048, by * 64 + vy * 32 + ty * 4 + j)
                                            A_shared_local[v0, v1] = A_shared[v0, v1]
                                    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, k_0 * 8 + k_1)
                                            with T.init():
                                                C_local[vi, vj] = T.float32(0)
                                            C_local[vi, vj] = C_local[vi, vj] + A_shared_local[vk, vi] * B_shared_local[vk, vj]
                            for i, j in T.grid(4, 4):
                                with T.sblock("C_local"):
                                    v0 = T.axis.S(2048, by * 64 + vy * 32 + ty * 4 + i)
                                    v1 = T.axis.S(2048, bx * 64 + vx * 32 + tx * 4 + j)
                                    C[v0, v1] = C_local[v0, v1]


@T.prim_func(s_tir=True)
def cuda_matmul_3(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")
    A_shared = T.sblock_alloc_buffer([2048, 2048], "float32", scope="shared")
    B_shared = T.sblock_alloc_buffer([2048, 2048], "float32", scope="shared")
    A_shared_local = T.sblock_alloc_buffer([2048, 2048], "float32", scope="local")
    B_shared_local = T.sblock_alloc_buffer([2048, 2048], "float32", scope="local")
    C_local = T.sblock_alloc_buffer([2048, 2048], "float32", scope="local")
    for i, j in T.grid(2048, 2048):
        with T.sblock("A_shared"):
            v0, v1 = T.axis.remap("SS", [i, j])
            A_shared[v0, v1] = A[v0, v1]
    for i, j in T.grid(2048, 2048):
        with T.sblock("B_shared"):
            v0, v1 = T.axis.remap("SS", [i, j])
            B_shared[v0, v1] = B[v0, v1]
    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):
                                        with T.sblock("A_shared_local"):
                                            v0 = T.axis.S(2048, k0 * 8 + k1 + i)
                                            v1 = T.axis.S(2048, by * 64 + vy * 32 + ty * 4 + j)
                                            A_shared_local[v0, v1] = A_shared[v0, v1]
                                    for i, j in T.grid(1, 4):
                                        with T.sblock("B_shared_local"):
                                            v0 = T.axis.S(2048, k0 * 8 + k1 + i)
                                            v1 = T.axis.S(2048, bx * 64 + vx * 32 + tx * 4 + j)
                                            B_shared_local[v0, v1] = B_shared[v0, v1]
                                    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)
                                            with T.init():
                                                C_local[vi, vj] = T.float32(0)
                                            C_local[vi, vj] = C_local[vi, vj] + A_shared_local[vk, vi] * B_shared_local[vk, vj]
                            for i, j in T.grid(4, 4):
                                with T.sblock("C_local"):
                                    v0 = T.axis.S(2048, by * 64 + vy * 32 + ty * 4 + i)
                                    v1 = T.axis.S(2048, bx * 64 + vx * 32 + tx * 4 + j)
                                    C[v0, v1] = C_local[v0, v1]


@T.prim_func(s_tir=True)
def cuda_matmul_4(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")
    A_shared = T.sblock_alloc_buffer([2048, 2048], "float32", scope="shared")
    B_shared = T.sblock_alloc_buffer([2048, 2048], "float32", scope="shared")
    A_shared_local = T.sblock_alloc_buffer([2048, 2048], "float32", scope="local")
    B_shared_local = T.sblock_alloc_buffer([2048, 2048], "float32", scope="local")
    C_local = T.sblock_alloc_buffer([2048, 2048], "float32", scope="local")
    for i, j in T.grid(2048, 2048):
        with T.sblock("B_shared"):
            v0, v1 = T.axis.remap("SS", [i, j])
            B_shared[v0, v1] = B[v0, v1]
    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 i, j in T.grid(8, 64):
                                    with T.sblock("A_shared"):
                                        v0 = T.axis.S(2048, k0 * 8 + i)
                                        v1 = T.axis.S(2048, by * 64 + j)
                                        A_shared[v0, v1] = A[v0, v1]
                                for k1 in T.unroll(0, 8):
                                    for i, j in T.grid(1, 4):
                                        with T.sblock("A_shared_local"):
                                            v0 = T.axis.S(2048, k0 * 8 + k1 + i)
                                            v1 = T.axis.S(2048, by * 64 + vy * 32 + ty * 4 + j)
                                            A_shared_local[v0, v1] = A_shared[v0, v1]
                                    for i, j in T.grid(1, 4):
                                        with T.sblock("B_shared_local"):
                                            v0 = T.axis.S(2048, k0 * 8 + k1 + i)
                                            v1 = T.axis.S(2048, bx * 64 + vx * 32 + tx * 4 + j)
                                            B_shared_local[v0, v1] = B_shared[v0, v1]
                                    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)
                                            with T.init():
                                                C_local[vi, vj] = 0.0
                                            C_local[vi, vj] = C_local[vi, vj] + A_shared_local[vk, vi] * B_shared_local[vk, vj]
                            for i, j in T.grid(4, 4):
                                with T.sblock("C_local"):
                                    v0 = T.axis.S(2048, by * 64 + vy * 32 + ty * 4 + i)
                                    v1 = T.axis.S(2048, bx * 64 + vx * 32 + tx * 4 + j)
                                    C[v0, v1] = C_local[v0, v1]


@T.prim_func(s_tir=True)
def cuda_matmul_5(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")
    A_shared = T.sblock_alloc_buffer([2048, 2048], "float32", scope="shared")
    B_shared = T.sblock_alloc_buffer([2048, 2048], "float32", scope="shared")
    A_shared_local = T.sblock_alloc_buffer([2048, 2048], "float32", scope="local")
    B_shared_local = T.sblock_alloc_buffer([2048, 2048], "float32", scope="local")
    C_local = T.sblock_alloc_buffer([2048, 2048], "float32", scope="local")
    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 i, j in T.grid(8, 64):
                                    with T.sblock("A_shared"):
                                        v0 = T.axis.S(2048, k0 * 8 + i)
                                        v1 = T.axis.S(2048, by * 64 + j)
                                        A_shared[v0, v1] = A[v0, v1]
                                for i, j in T.grid(8, 64):
                                    with T.sblock("B_shared"):
                                        v0 = T.axis.S(2048, k0 * 8 + i)
                                        v1 = T.axis.S(2048, bx * 64 + j)
                                        B_shared[v0, v1] = B[v0, v1]
                                for k1 in T.unroll(0, 8):
                                    for i, j in T.grid(1, 4):
                                        with T.sblock("A_shared_local"):
                                            v0 = T.axis.S(2048, k0 * 8 + k1 + i)
                                            v1 = T.axis.S(2048, by * 64 + vy * 32 + ty * 4 + j)
                                            A_shared_local[v0, v1] = A_shared[v0, v1]
                                    for i, j in T.grid(1, 4):
                                        with T.sblock("B_shared_local"):
                                            v0 = T.axis.S(2048, k0 * 8 + k1 + i)
                                            v1 = T.axis.S(2048, bx * 64 + vx * 32 + tx * 4 + j)
                                            B_shared_local[v0, v1] = B_shared[v0, v1]
                                    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)
                                            with T.init():
                                                C_local[vi, vj] = 0.0
                                            C_local[vi, vj] = C_local[vi, vj] + A_shared_local[vk, vi] * B_shared_local[vk, vj]
                            for i, j in T.grid(4, 4):
                                with T.sblock("C_local"):
                                    v0 = T.axis.S(2048, by * 64 + vy * 32 + ty * 4 + i)
                                    v1 = T.axis.S(2048, bx * 64 + vx * 32 + tx * 4 + j)
                                    C[v0, v1] = C_local[v0, v1]


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


@T.prim_func(s_tir=True)
def tiled_after_reverse_compute_at(a: T.handle, c: T.handle) -> None:
    A = T.match_buffer(a, [128, 128], "float32")
    B = T.sblock_alloc_buffer([128, 128], "float32")
    C = T.match_buffer(c, [128, 128], "float32")
    for i_0, j_0, i_1 in T.grid(8, 8, 16):
        for j_1 in T.serial(0, 16):
            with T.sblock("B"):
                vi = T.axis.S(128, i_0 * 16 + i_1)
                vj = T.axis.S(128, j_0 * 16 + j_1)
                B[vi, vj] = A[vi, vj] * 2.0
        for j_1 in T.serial(0, 16):
            with T.sblock("C"):
                vi = T.axis.S(128, i_0 * 16 + i_1)
                vj = T.axis.S(128, j_0 * 16 + j_1)
                C[vi, vj] = B[vi, vj] + 1.0


@T.prim_func(s_tir=True)
def tiled_trivial_binding(a: T.handle, c: T.handle) -> None:
    A = T.match_buffer(a, [1, 128, 128], "float32")
    B = T.sblock_alloc_buffer([1, 128, 128], "float32")
    C = T.match_buffer(c, [1, 128, 128], "float32")
    for i_0, j_0, i_1, j_1 in T.grid(8, 8, 16, 16):
        with T.sblock("B"):
            vi = T.axis.S(128, i_0 * 16 + i_1)
            vj = T.axis.S(128, j_0 * 16 + j_1)
            B[0, vi, vj] = A[0, vi, vj] * 2.0
    for i, j in T.grid(128, 128):
        with T.sblock("C"):
            vi, vj = T.axis.remap("SS", [i, j])
            C[0, vi, vj] = B[0, vi, vj] + 1.0


@T.prim_func(s_tir=True)
def tiled_trivial_binding_after_reverse_compute_at(a: T.handle, c: T.handle) -> None:
    A = T.match_buffer(a, [1, 128, 128], "float32")
    B = T.sblock_alloc_buffer([1, 128, 128], "float32")
    C = T.match_buffer(c, [1, 128, 128], "float32")
    for i_0, j_0, i_1 in T.grid(8, 8, 16):
        for j_1 in T.serial(0, 16):
            with T.sblock("B"):
                vi = T.axis.S(128, i_0 * 16 + i_1)
                vj = T.axis.S(128, j_0 * 16 + j_1)
                B[0, vi, vj] = A[0, vi, vj] * 2.0
        for j_1 in T.serial(0, 16):
            with T.sblock("C"):
                vi = T.axis.S(128, i_0 * 16 + i_1)
                vj = T.axis.S(128, j_0 * 16 + j_1)
                C[0, vi, vj] = B[0, vi, vj] + 1.0


@T.prim_func(s_tir=True)
def factorized(a: T.handle, b: T.handle) -> None:
    A = T.match_buffer(a, [16, 16, 16], "float32")
    B = T.match_buffer(b, [16], "float32")
    B_rf_local = T.sblock_alloc_buffer([16, 16], "float32", scope="local")
    for j in T.thread_binding(0, 16, thread = "blockIdx.x"):
        for i_o in T.thread_binding(0, 4, thread = "threadIdx.x"):
            for i_i, k in T.grid(4, 16):
                with T.sblock("B_rf"):
                    vi = T.axis.S(16, i_o * 4 + i_i)
                    vj, vk = T.axis.remap("SR", [j, k])
                    with T.init():
                        B_rf_local[vi, vj] = 0.0
                    B_rf_local[vi, vj] = B_rf_local[vi, vj] + A[vj, vi, vk]
    for i, k in T.grid(16, 16):
        with T.sblock("B"):
            vi, vk = T.axis.remap("SR", [i, k])
            with T.init():
                B[vi] = 0.0
            B[vi] = B[vi] + B_rf_local[vk, vi]


@T.prim_func(s_tir=True)
def factorized_after_reverse_compute_at(a: T.handle, b: T.handle) -> None:
    A = T.match_buffer(a, [16, 16, 16], "float32")
    B = T.match_buffer(b, [16], "float32")
    B_rf_local = T.sblock_alloc_buffer([16, 16], "float32", scope="local")
    for j in T.thread_binding(0, 16, thread = "blockIdx.x"):
        for i_o in T.thread_binding(0, 4, thread = "threadIdx.x"):
            for i_i, k in T.grid(4, 16):
                with T.sblock("B_rf"):
                    vi = T.axis.S(16, i_o * 4 + i_i)
                    vj = T.axis.S(16, j)
                    vk = T.axis.R(16, k)
                    with T.init():
                        B_rf_local[vi, vj] = 0.0
                    B_rf_local[vi, vj] = B_rf_local[vi, vj] + A[vj, vi, vk]
            for k in T.serial(0, 4):
                with T.sblock("B"):
                    vi = T.axis.S(16, j)
                    vk = T.axis.R(16, i_o * 4 + k)
                    with T.init():
                        B[vi] = 0.0
                    B[vi] = B[vi] + B_rf_local[vk, vi]


@T.prim_func(s_tir=True)
def not_all_compact_data_flow(a: T.handle, c: T.handle):
    A = T.match_buffer(a, (128, 128), "float32")
    B = T.sblock_alloc_buffer((128, 128), "float32")
    C = T.match_buffer(c, (128, 128), "float32")
    for i, j in T.grid(128, 128):
        with T.sblock("B"):
            vi, vj = T.axis.remap("SS", [i, j])
            B[vi, vj] = A[vi, vj]
    for i, j in T.grid(128, 64):
        with T.sblock("C_1"):
            vi, vj = T.axis.remap("SS", [i, j])
            C[vi, vj * 2] = B[vi, vj * 2] + 1.0
        with T.sblock("C_2"):
            vi, vj = T.axis.remap("SS", [i, j])
            C[vi, vj * 2 + 1] = B[vi, vj * 2 + 1] * 2.0


@T.prim_func(s_tir=True)
def not_all_compact_data_flow_after_compute_at(a: T.handle, c: T.handle):
    A = T.match_buffer(a, (128, 128), "float32")
    B = T.sblock_alloc_buffer((128, 128), "float32")
    C = T.match_buffer(c, (128, 128), "float32")
    for i, j in T.grid(128, 64):
        for t in range(2):
            with T.sblock("B"):
                vi = T.axis.S(128, i)
                vj = T.axis.S(128, j * 2 + t)
                B[vi, vj] = A[vi, vj]
        with T.sblock("C_1"):
            vi, vj = T.axis.remap("SS", [i, j])
            C[vi, vj * 2] = B[vi, vj * 2] + 1.0
        with T.sblock("C_2"):
            vi, vj = T.axis.remap("SS", [i, j])
            C[vi, vj * 2 + 1] = B[vi, vj * 2 + 1] * 2.0


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


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


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


@T.prim_func(s_tir=True)
def read_out_of_bound(a: T.handle, c:T.handle) -> None:
    A = T.match_buffer(a, [16], "float32")
    B = T.sblock_alloc_buffer([16], "float32")
    C = T.match_buffer(c, [16], "float32")
    for i in T.serial(0, 16):
        with T.sblock("B"):
            v = T.axis.S(16, i)
            B[v] = A[v]
    for j in T.serial(0, 16):
        with T.sblock("C"):
            v = T.axis.S(16, j)
            T.reads(B[v : v + 2])
            C[v] = T.if_then_else(v < 15, T.max(B[v], B[v + 1]), B[v], dtype="float32")


@T.prim_func(s_tir=True)
def read_out_of_bound_after_compute_at(a: T.handle, c: T.handle) -> None:
    A = T.match_buffer(a, [16], "float32")
    B = T.sblock_alloc_buffer([16], "float32")
    C = T.match_buffer(c, [16], "float32")
    for j in T.serial(0, 16):
        for i in T.serial(0, 2):
            with T.sblock("B"):
                v = T.axis.S(16, j + i)
                T.where(j + i < 16)
                B[v] = A[v]
        with T.sblock("C"):
            v = T.axis.S(16, j)
            T.reads([B[v : v + 2]])
            C[v] = T.if_then_else(v < 15, T.max(B[v], B[v + 1]), B[v], dtype="float32")


@T.prim_func(s_tir=True)
def multi_reduction(A: T.Buffer((16, 16), "float32"), C: T.Buffer((), "float32")):
    B = T.sblock_alloc_buffer((16, ), dtype="float32")
    for i, k in T.grid(16, 16):
        with T.sblock("B"):
            vi, vk = T.axis.remap("SR", [i, k])
            with T.init():
                B[vi] = 0.0
            B[vi] += A[vi, vk]
    for k in T.grid(16):
        with T.sblock("C"):
            vk = T.axis.remap("R", [k])
            with T.init():
                C[()] = 0.0
            C[()] += B[vk]


@T.prim_func(s_tir=True)
def multi_reduction_after_compute_at(
    A: T.Buffer((16, 16), "float32"),
    C:T.Buffer((), "float32"),
):
    B = T.sblock_alloc_buffer((16, ), dtype="float32")
    for k in T.grid(16):
        for kk in T.grid(16):
            with T.sblock("B"):
                vi, vk = T.axis.remap("SR", [k, kk])
                with T.init():
                    B[vi] = 0.0
                B[vi] += A[vi, vk]
        with T.sblock("C"):
            vk = T.axis.remap("R", [k])
            with T.init():
                C[()] = 0.0
            C[()] += B[vk]


@T.prim_func(s_tir=True)
def tiled_pooling_read_cache(a: T.handle, b: T.handle) -> None:
    X = T.match_buffer(a, [224, 224], dtype="float32")
    Y = T.match_buffer(b, [224, 224], dtype="float32")
    cache = T.sblock_alloc_buffer([224, 224], dtype="float32")
    for hh, ww in T.grid(224, 224):
        with T.sblock("cache"):
            h, w = T.axis.remap("SS", [hh, ww])
            cache[h, w] = X[h, w]
    for hh_0, ww_0, hh_1, ww_1, khh, kww in T.grid(28, 28, 8, 8, 3, 3):
        with T.sblock("compute"):
            h = T.axis.spatial(224, hh_0 * 8 + hh_1)
            w = T.axis.spatial(224, ww_0 * 8 + ww_1)
            kh, kw = T.axis.remap("RR", [khh, kww])
            with T.init():
                Y[h, w] = 0.0
            Y[h, w] = T.max(Y[h, w], T.if_then_else(
                T.likely(1 <= h + kh, dtype="bool") and \
                T.likely(h + kh < 225, dtype="bool") and \
                T.likely(1 <= w + kw, dtype="bool") and \
                T.likely(w + kw < 225, dtype="bool"),
                cache[h + kh - 1, w + kw - 1], 0.0, dtype="float32"))

@T.prim_func(s_tir=True)
def tiled_pooling_read_cache_after_compute_at(a: T.handle, b: T.handle) -> None:
    X = T.match_buffer(a, [224, 224], dtype="float32")
    Y = T.match_buffer(b, [224, 224], dtype="float32")
    cache = T.sblock_alloc_buffer([224, 224], dtype="float32")
    for hh_0, ww_0 in T.grid(28, 28):
        for ax0, ax1 in T.grid(10, 10):
            with T.sblock("cache"):
                h = T.axis.spatial(224, hh_0 * 8 - 1 + ax0)
                w = T.axis.spatial(224, ww_0 * 8 - 1 + ax1)
                T.where(1 <= hh_0 * 8 + ax0 and hh_0 * 8 + ax0 < 225 and 1 <= ww_0 * 8 + ax1 and ww_0 * 8 + ax1 < 225)
                cache[h, w] = X[h, w]
        for hh_1, ww_1, khh, kww in T.grid(8, 8, 3, 3):
            with T.sblock("compute"):
                h = T.axis.spatial(224, hh_0 * 8 + hh_1)
                w = T.axis.spatial(224, ww_0 * 8 + ww_1)
                kh, kw = T.axis.remap("RR", [khh, kww])
                with T.init():
                    Y[h, w] = 0.0
                Y[h, w] = T.max(Y[h, w], T.if_then_else(
                    T.likely(1 <= h + kh, dtype="bool") and \
                    T.likely(h + kh < 225, dtype="bool") and \
                    T.likely(1 <= w + kw, dtype="bool") and \
                    T.likely(w + kw < 225, dtype="bool"),
                    cache[h + kh - 1, w + kw - 1], 0.0, dtype="float32"))

@T.prim_func(s_tir=True)
def non_uniform_tiled_conv(x: T.Buffer((1, 3, 100, 100), "float32"),
                           w: T.Buffer((16, 3, 3, 3), "float32"),
                           y: T.Buffer((1, 16, 98, 98), "float32")) -> None:
    x_global = T.sblock_alloc_buffer([1, 3, 100, 100], dtype="float32")
    for ax0, ax1, ax2, ax3 in T.grid(1, 3, 100, 100):
        with T.sblock("cache"):
            v0, v1, v2, v3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
            x_global[v0, v1, v2, v3] = x[v0, v1, v2, v3]
    for h_o, w_o, n, c_o, h_i, w_i, c_i, kh, kw in T.grid(7, 7, 1, 16, 15, 15, 3, 3, 3):
        with T.sblock("compute"):
            nn = T.axis.spatial(1, 0)
            cc = T.axis.spatial(16, c_o)
            hh = T.axis.spatial(98, h_o * 15 + h_i)
            ww = T.axis.spatial(98, w_o * 15 + w_i)
            rc, rh, rw = T.axis.remap("RRR", [c_i, kh, kw])
            T.where(h_o * 15 + h_i < 98 and w_o * 15 + w_i < 98)
            with T.init():
                y[nn, cc, hh, ww] = T.float32(0)
            y[nn, cc, hh, ww] = y[nn, cc, hh, ww] + \
                x_global[nn, cc // 16 * 3 + rc, hh + rh, ww + rw] * w[cc, rc, rh, rw]

@T.prim_func(s_tir=True)
def non_uniform_tiled_conv_after_compute_at(x: T.Buffer((1, 3, 100, 100), "float32"),
                                            w: T.Buffer((16, 3, 3, 3), "float32"),
                                            y: T.Buffer((1, 16, 98, 98), "float32")) -> None:
    x_global = T.sblock_alloc_buffer([1, 3, 100, 100], dtype="float32")
    for h_o, w_o in T.grid(7, 7):
        for ax0, ax1, ax2 in T.grid(3, 17, 17):
            with T.sblock("cache"):
                v0 = T.axis.spatial(1, 0)
                v1 = T.axis.spatial(3, ax0)
                v2 = T.axis.spatial(100, h_o * 15 + ax1)
                v3 = T.axis.spatial(100, w_o * 15 + ax2)
                T.where(h_o * 15 + ax1 < 100 and w_o * 15 + ax2 < 100)
                x_global[v0, v1, v2, v3] = x[v0, v1, v2, v3]
        for n, c_o, h_i, w_i, c_i, kh, kw in T.grid(1, 16, 15, 15, 3, 3, 3):
            with T.sblock("compute"):
                nn = T.axis.spatial(1, 0)
                cc = T.axis.spatial(16, c_o)
                hh = T.axis.spatial(98, h_o * 15 + h_i)
                ww = T.axis.spatial(98, w_o * 15 + w_i)
                rc, rh, rw = T.axis.remap("RRR", [c_i, kh, kw])
                T.where(h_o * 15 + h_i < 98 and w_o * 15 + w_i < 98)
                with T.init():
                    y[nn, cc, hh, ww] = T.float32(0)
                y[nn, cc, hh, ww] = y[nn, cc, hh, ww] + \
                    x_global[nn, cc // 16 * 3 + rc, hh + rh, ww + rw] * w[cc, rc, rh, rw]

@T.prim_func(s_tir=True)
def concat_two_elemwise(x: T.Buffer((16,), "float32"),
                        y: T.Buffer((8,), "float32"),
                        T_concat: T.Buffer((24,), "float32")) -> None:
    T_add_1 = T.sblock_alloc_buffer([16], dtype="float32")
    T_add_2 = T.sblock_alloc_buffer([8], dtype="float32")
    for i in T.serial(16):
        with T.sblock("T_add_1"):
            ax = T.axis.spatial(16, i)
            T_add_1[ax] = x[ax] + T.float32(1)
    for i in T.serial(8):
        with T.sblock("T_add_2"):
            ax = T.axis.spatial(8, i)
            T_add_2[ax] = y[ax] + T.float32(2)
    for i in T.serial(24):
        with T.sblock("T_concat"):
            ax = T.axis.spatial(24, i)
            T_concat[ax] = T.if_then_else(16 <= ax, T_add_2[ax - 16], T_add_1[ax], dtype="float32")

@T.prim_func(s_tir=True)
def concat_two_elemwise_after_compute_at(x: T.Buffer((16,), "float32"),
                                         y: T.Buffer((8,), "float32"),
                                         T_concat: T.Buffer((24,), "float32")) -> None:
    T_add_1 = T.sblock_alloc_buffer([16], dtype="float32")
    T_add_2 = T.sblock_alloc_buffer([8], dtype="float32")
    for i in T.serial(24):
        with T.sblock("T_add_1"):
            ax = T.axis.spatial(16, i)
            T.where(i < 16)
            T_add_1[ax] = x[ax] + T.float32(1)
        with T.sblock("T_add_2"):
            ax = T.axis.spatial(8, i - 16)
            T.where(16 <= i)
            T_add_2[ax] = y[ax] + T.float32(2)
        with T.sblock("T_concat"):
            ax = T.axis.spatial(24, i)
            T_concat[ax] = T.if_then_else(16 <= ax, T_add_2[ax - 16], T_add_1[ax], dtype="float32")

@T.prim_func(s_tir=True)
def floordiv_and_floormod_indices(a: T.handle, b: T.handle) -> None:
    X = T.match_buffer(a, [16, 16])
    Y = T.match_buffer(b, [256])
    temp = T.sblock_alloc_buffer([16, 16])
    for i, j in T.grid(16, 16):
        with T.sblock("A"):
            v_i, v_j = T.axis.remap("SS", [i, j])
            temp[v_i, v_j] = X[v_j, v_i] + 1.0
    for i in T.serial(0, 256):
        with T.sblock("B"):
            v_i = T.axis.remap("S", [i])
            Y[v_i] = temp[v_i // 16, v_i % 16]

@T.prim_func(s_tir=True)
def floordiv_and_floormod_indices_after_reverse_compute_at(a: T.handle, b: T.handle) -> None:
    X = T.match_buffer(a, [16, 16], dtype="float32")
    Y = T.match_buffer(b, [256], dtype="float32")
    temp = T.sblock_alloc_buffer([16, 16], dtype="float32")
    for i in T.serial(0, 16):
        for j in T.serial(0, 16):
            with T.sblock("A"):
                v_i, v_j = T.axis.remap("SS", [i, j])
                temp[v_i, v_j] = X[v_j, v_i] + T.float32(1)
        for ax0 in T.serial(0, 16):
            with T.sblock("B"):
                v_i = T.axis.spatial(256, i * 16 + ax0)
                Y[v_i] = temp[v_i // 16, v_i % 16]


@T.prim_func(s_tir=True)
def recursive_floordiv_floormod(A: T.Buffer((16, 64, 1, 8, 8, 32), "float32"),
                                C: T.Buffer((3, 512, 512), "float32")) -> None:
    T.func_attr({"tirx.noalias": True})
    # with T.sblock("root"):
    B = T.sblock_alloc_buffer((1, 128, 16, 8, 2, 32, 2), "float32")
    for axis1, axis2, axis3, axis4, axis5, axis6, axis7 in T.grid(1, 128, 16, 8, 2, 32, 2):
        with T.sblock("In"):
            v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6, v_axis7 = T.axis.remap("SSSSSSS", [axis1, axis2, axis3, axis4, axis5, axis6, axis7])
            T.reads(A[(v_axis2 * 4 + v_axis5 * 2 + v_axis7) // 32, (v_axis3 * 32 + v_axis6) // 8, (v_axis1 * 8 + v_axis4) // 8, (v_axis3 * 32 + v_axis6) % 8, v_axis1 * 8 + v_axis4, (v_axis2 * 4 + v_axis5 * 2 + v_axis7) % 32])
            T.writes(B[v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6, v_axis7])
            B[v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6, v_axis7] = A[(v_axis2 * 4 + v_axis5 * 2 + v_axis7) // 32, (v_axis3 * 32 + v_axis6) // 8, (v_axis1 * 8 + v_axis4) // 8, (v_axis3 * 32 + v_axis6) % 8, v_axis1 * 8 + v_axis4, (v_axis2 * 4 + v_axis5 * 2 + v_axis7) % 32] + 3
    for ax1, ax2, ax3 in T.grid(3, 512, 512):
        with T.sblock("Out"):
            v1, v2, v3 = T.axis.remap("SSS", [ax1, ax2, ax3])
            T.reads(B[v1 // 8, v2 // 4, v3 // 32, v1, v2 % 4 // 2, v3 % 32, v2 % 2])
            T.writes(C[v1, v2, v3])
            C[v1, v2, v3] = B[v1 // 8, v2 // 4, v3 // 32, v1, v2 % 4 // 2, v3 % 32, v2 % 2] * 2


@T.prim_func(s_tir=True)
def recursive_floordiv_floormod_after_reverse_compute_at(A: T.Buffer((16, 64, 1, 8, 8, 32), "float32"), C: T.Buffer((3, 512, 512), "float32")) -> None:
    T.func_attr({"tirx.noalias": True})
    # with T.sblock("root"):
    B = T.sblock_alloc_buffer((1, 128, 16, 8, 2, 32, 2))
    for axis1, axis2, axis3 in T.grid(1, 128, 16):
        for axis4, axis5, axis6, axis7 in T.grid(8, 2, 32, 2):
            with T.sblock("In"):
                v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6, v_axis7 = T.axis.remap("SSSSSSS", [axis1, axis2, axis3, axis4, axis5, axis6, axis7])
                T.reads(A[(v_axis2 * 4 + v_axis5 * 2 + v_axis7) // 32, (v_axis3 * 32 + v_axis6) // 8, (v_axis1 * 8 + v_axis4) // 8, (v_axis3 * 32 + v_axis6) % 8, v_axis1 * 8 + v_axis4, (v_axis2 * 4 + v_axis5 * 2 + v_axis7) % 32])
                T.writes(B[v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6, v_axis7])
                B[v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6, v_axis7] = A[(v_axis2 * 4 + v_axis5 * 2 + v_axis7) // 32, (v_axis3 * 32 + v_axis6) // 8, (v_axis1 * 8 + v_axis4) // 8, (v_axis3 * 32 + v_axis6) % 8, v_axis1 * 8 + v_axis4, (v_axis2 * 4 + v_axis5 * 2 + v_axis7) % 32] + T.float32(3)
        for ax0, ax1, ax2 in T.grid(3, 4, 32):
            with T.sblock("Out"):
                v1 = T.axis.spatial(3, ax0)
                v2 = T.axis.spatial(512, axis2 * 4 + ax1)
                v3 = T.axis.spatial(512, axis3 * 32 + ax2)
                T.reads(B[v1 // 8, v2 // 4, v3 // 32, v1, v2 % 4 // 2, v3 % 32, v2 % 2])
                T.writes(C[v1, v2, v3])
                C[v1, v2, v3] = B[v1 // 8, v2 // 4, v3 // 32, v1, v2 % 4 // 2, v3 % 32, v2 % 2] * T.float32(2)


@T.prim_func(s_tir=True)
def tiled_repeat_op(x: T.Buffer((4,), "float32"), T_repeat: T.Buffer((64,), "float32")) -> None:
    T_add = T.sblock_alloc_buffer([4], dtype="float32")
    for i0 in T.serial(4):
        with T.sblock("T_add"):
            ax0 = T.axis.spatial(4, i0)
            T_add[ax0] = x[ax0] + 1.0
    for i0_0, i0_1 in T.grid(8, 8):
        with T.sblock("T_repeat"):
            ax0 = T.axis.spatial(64, i0_0 * 8 + i0_1)
            T_repeat[ax0] = T_add[ax0 // 16]

@T.prim_func(s_tir=True)
def tiled_repeat_op_after_compute_at(x: T.Buffer((4,), "float32"), T_repeat: T.Buffer((64,), "float32")) -> None:
    T_add = T.sblock_alloc_buffer([4], dtype="float32")
    for i0_0 in T.serial(8):
        with T.sblock("T_add"):
            ax0 = T.axis.spatial(4, i0_0 // 2)
            T_add[ax0] = x[ax0] + T.float32(1)
        for i0_1 in T.serial(8):
            with T.sblock("T_repeat"):
                ax0 = T.axis.spatial(64, i0_0 * 8 + i0_1)
                T_repeat[ax0] = T_add[ax0 // 16]

@T.prim_func(s_tir=True)
def static_bound(A: T.Buffer((32, 1), "float32"), C: T.Buffer((32, 1), "float32")) -> None:
    B = T.sblock_alloc_buffer((32, 1), "float32")
    for i, j in T.grid(32, 1):
        with T.sblock("B"):
            vi = T.axis.spatial(32, i)
            vj = T.axis.spatial(1, j)
            B[vi, vj] = A[vi, vj] * 2.0
    for i, j in T.grid(32, 32):
        with T.sblock("C"):
            vi = T.axis.spatial(32, i)
            vj = T.axis.spatial(1, j)
            T.where(j < 1)
            C[vi, vj] = B[vi, vj] + 1.0

@T.prim_func(s_tir=True)
def static_bound_after_compute_at(A: T.Buffer((32, 1), "float32"), C: T.Buffer((32, 1), "float32")) -> None:
    B = T.sblock_alloc_buffer((32, 1), "float32")
    for i in range(32):
        for ax0, ax1 in T.grid(1, 1):
            with T.sblock("B"):
                vi = T.axis.spatial(32, i + ax0)
                vj = T.axis.spatial(1, ax1)
                B[vi, vj] = A[vi, vj] * 2.0
        for j in range(32):
            with T.sblock("C"):
                vi = T.axis.spatial(32, i)
                vj = T.axis.spatial(1, j)
                T.where(j < 1)
                C[vi, vj] = B[vi, vj] + 1.0
# pylint: enable=no-member,invalid-name,unused-variable,line-too-long,redefined-outer-name,unexpected-keyword-arg,too-many-nested-blocks
# fmt: on

use_block_name = tvm.testing.parameter(by_dict={"block_obj": False, "block_name": True})


def test_compute_at_two_elementwise(use_block_name):
    sch = tvm.s_tir.Schedule(two_elementwise, debug_mask="all")
    block = "B" if use_block_name else sch.get_sblock("B")
    loop, _ = sch.get_loops("C" if use_block_name else sch.get_sblock("C"))
    sch.compute_at(block, loop, preserve_unit_loops=True)
    assert_structural_equal_ignore_global_symbol(two_elementwise_after_compute_at, sch.mod["main"])
    verify_trace_roundtrip(sch=sch, mod=two_elementwise)


def test_compute_at_blockized_1(use_block_name):
    sch = tvm.s_tir.Schedule(blockized_1, debug_mask="all")
    block = sch.get_sblock("B")
    _, loop = sch.get_loops(sch.get_sblock("C_outer"))
    sch.compute_at(block, loop, preserve_unit_loops=True)
    assert_structural_equal_ignore_global_symbol(blockized_after_compute_at, sch.mod["main"])
    verify_trace_roundtrip(sch=sch, mod=blockized_1)


def test_compute_at_blockized_2(use_block_name):
    sch = tvm.s_tir.Schedule(blockized_2, debug_mask="all")
    block = sch.get_sblock("B_outer")
    _, loop, _, _ = sch.get_loops(sch.get_sblock("C"))
    sch.compute_at(block, loop, preserve_unit_loops=True)
    assert_structural_equal_ignore_global_symbol(blockized_2_after_compute_at, sch.mod["main"])
    verify_trace_roundtrip(sch=sch, mod=blockized_2)


def test_compute_at_cuda_matmul_0(use_block_name):
    sch = tvm.s_tir.Schedule(cuda_matmul_0, debug_mask="all")
    block = sch.get_sblock("C")
    _, _, _, _, _, loop, _, _ = sch.get_loops(sch.get_sblock("C_local"))
    sch.compute_at(block, loop, preserve_unit_loops=True)
    assert_structural_equal_ignore_global_symbol(cuda_matmul_0_after_compute_at, sch.mod["main"])
    verify_trace_roundtrip(sch=sch, mod=cuda_matmul_0)


def test_compute_at_cuda_matmul_1(use_block_name):
    sch = tvm.s_tir.Schedule(cuda_matmul_1, debug_mask="all")
    block = sch.get_sblock("A_shared_local")
    _, _, _, _, _, _, _, loop, _, _, _ = sch.get_loops(sch.get_sblock("C"))
    sch.compute_at(block, loop, preserve_unit_loops=True)
    assert_structural_equal_ignore_global_symbol(cuda_matmul_2, sch.mod["main"])
    verify_trace_roundtrip(sch=sch, mod=cuda_matmul_1)


def test_compute_at_cuda_matmul_2(use_block_name):
    sch = tvm.s_tir.Schedule(cuda_matmul_2, debug_mask="all")
    block = sch.get_sblock("B_shared_local")
    _, _, _, _, _, _, _, loop, _, _, _ = sch.get_loops(sch.get_sblock("C"))
    sch.compute_at(block, loop, preserve_unit_loops=True)
    assert_structural_equal_ignore_global_symbol(cuda_matmul_3, sch.mod["main"])
    verify_trace_roundtrip(sch=sch, mod=cuda_matmul_2)


def test_compute_at_cuda_matmul_3(use_block_name):
    sch = tvm.s_tir.Schedule(cuda_matmul_3, debug_mask="all")
    block = sch.get_sblock("A_shared")
    _, _, _, _, _, _, loop, _, _, _, _ = sch.get_loops(sch.get_sblock("C"))
    sch.compute_at(block, loop, preserve_unit_loops=True)
    assert_structural_equal_ignore_global_symbol(cuda_matmul_4, sch.mod["main"])
    verify_trace_roundtrip(sch=sch, mod=cuda_matmul_3)


def test_compute_at_cuda_matmul_4(use_block_name):
    sch = tvm.s_tir.Schedule(cuda_matmul_4, debug_mask="all")
    block = sch.get_sblock("B_shared")
    _, _, _, _, _, _, loop, _, _, _, _ = sch.get_loops(sch.get_sblock("C"))
    sch.compute_at(block, loop, preserve_unit_loops=True)
    assert_structural_equal_ignore_global_symbol(cuda_matmul_5, sch.mod["main"])
    verify_trace_roundtrip(sch=sch, mod=cuda_matmul_4)


def test_compute_at_reduction_block(use_block_name):
    sch = tvm.s_tir.Schedule(multi_reduction, debug_mask="all")
    block = sch.get_sblock("B")
    (loop,) = sch.get_loops(sch.get_sblock("C"))
    sch.compute_at(block, loop, preserve_unit_loops=False)
    assert_structural_equal_ignore_global_symbol(multi_reduction_after_compute_at, sch.mod["main"])
    verify_trace_roundtrip(sch=sch, mod=multi_reduction)


def test_compute_at_tiled_pooling_read_cache(use_block_name):
    sch = tvm.s_tir.Schedule(tiled_pooling_read_cache, debug_mask="all")
    compute = sch.get_sblock("compute")
    _, w_o, _, _, _, _ = sch.get_loops(compute)
    cache = sch.get_sblock("cache")
    sch.compute_at(cache, w_o)
    assert_structural_equal_ignore_global_symbol(
        tiled_pooling_read_cache_after_compute_at, sch.mod["main"]
    )
    verify_trace_roundtrip(sch=sch, mod=tiled_pooling_read_cache)


def test_compute_at_non_uniform_tiled_conv(use_block_name):
    sch = tvm.s_tir.Schedule(non_uniform_tiled_conv, debug_mask="all")
    compute = sch.get_sblock("compute")
    sch.compute_at(sch.get_sblock("cache"), sch.get_loops(compute)[1])
    assert_structural_equal_ignore_global_symbol(
        non_uniform_tiled_conv_after_compute_at, sch.mod["main"]
    )
    verify_trace_roundtrip(sch=sch, mod=non_uniform_tiled_conv)


def test_compute_at_concat(use_block_name):
    sch = tvm.s_tir.Schedule(concat_two_elemwise, debug_mask="all")
    concat = sch.get_sblock("T_concat")
    add1 = sch.get_sblock("T_add_1")
    add2 = sch.get_sblock("T_add_2")
    axis = sch.get_loops(concat)[0]
    sch.compute_at(add1, axis)
    sch.compute_at(add2, axis)
    assert_structural_equal_ignore_global_symbol(
        concat_two_elemwise_after_compute_at, sch.mod["main"]
    )
    verify_trace_roundtrip(sch=sch, mod=concat_two_elemwise)


def test_compute_at_tiled_repeat_op(use_block_name):
    sch = tvm.s_tir.Schedule(tiled_repeat_op, debug_mask="all")
    outer_ax, _ = sch.get_loops(sch.get_sblock("T_repeat"))
    sch.compute_at(sch.get_sblock("T_add"), outer_ax)
    assert_structural_equal_ignore_global_symbol(tiled_repeat_op_after_compute_at, sch.mod["main"])
    verify_trace_roundtrip(sch=sch, mod=tiled_repeat_op)


def test_compute_at_rev_iter():
    @T.prim_func(s_tir=True)
    def before(X: T.Buffer((10, 10), "float32"), Z: T.Buffer((10, 10), "float32")):
        Y = T.sblock_alloc_buffer([10, 10], "float32")
        for i, j in T.grid(10, 10):
            with T.sblock("b0"):
                vi, vj = T.axis.remap("SS", [i, j])
                Y[9 - vi, 9 - vj] = X[vi, vj] + 1.0
        for i, j in T.grid(10, 10):
            with T.sblock("b1"):
                vi, vj = T.axis.remap("SS", [i, j])
                Z[vi, vj] = Y[vj, vi] + 2.0

    @T.prim_func(s_tir=True)
    def after(X: T.Buffer((10, 10), "float32"), Z: T.Buffer((10, 10), "float32")):
        Y = T.sblock_alloc_buffer([10, 10], "float32")
        for i in range(10):
            for j in range(10):
                with T.sblock("b0"):
                    vi = T.axis.spatial(10, j)
                    vj = T.axis.spatial(10, 9 - i)
                    Y[9 - vi, 9 - vj] = X[vi, vj] + 1.0
            for j in range(10):
                with T.sblock("b1"):
                    vi, vj = T.axis.remap("SS", [i, j])
                    Z[vi, vj] = Y[vj, vi] + 2.0

    sch = tvm.s_tir.Schedule(before, debug_mask="all")
    axis = sch.get_loops(sch.get_sblock("b1"))[0]
    sch.compute_at(sch.get_sblock("b0"), axis)
    assert_structural_equal_ignore_global_symbol(after, sch.mod["main"])
    verify_trace_roundtrip(sch=sch, mod=before)


def test_reverse_compute_at_tiled(use_block_name):
    sch = tvm.s_tir.Schedule(tiled, debug_mask="all")
    block = sch.get_sblock("C")
    _, _, loop, _ = sch.get_loops(sch.get_sblock("B"))
    sch.reverse_compute_at(block, loop, preserve_unit_loops=False)
    assert_structural_equal_ignore_global_symbol(tiled_after_reverse_compute_at, sch.mod["main"])
    verify_trace_roundtrip(sch=sch, mod=tiled)


def test_reverse_compute_at_tiled_trivial_binding(use_block_name):
    sch = tvm.s_tir.Schedule(tiled_trivial_binding, debug_mask="all")
    block = sch.get_sblock("C")
    _, _, loop, _ = sch.get_loops(sch.get_sblock("B"))
    sch.reverse_compute_at(block, loop, preserve_unit_loops=False)
    assert_structural_equal_ignore_global_symbol(
        tiled_trivial_binding_after_reverse_compute_at, sch.mod["main"]
    )
    verify_trace_roundtrip(sch=sch, mod=tiled_trivial_binding)


def test_reverse_compute_at_blockized_2(use_block_name):
    sch = tvm.s_tir.Schedule(blockized_2, debug_mask="all")
    block = sch.get_sblock("C")
    _, loop = sch.get_loops(sch.get_sblock("B_outer"))
    sch.reverse_compute_at(block, loop, preserve_unit_loops=True)
    assert_structural_equal_ignore_global_symbol(
        blockized_2_after_reverse_compute_at, sch.mod["main"]
    )
    verify_trace_roundtrip(sch=sch, mod=blockized_2)


def test_reverse_compute_at_factorized(use_block_name):
    sch = tvm.s_tir.Schedule(factorized, debug_mask="all")
    block = sch.get_sblock("B")
    _, loop, _, _ = sch.get_loops(sch.get_sblock("B_rf"))
    sch.reverse_compute_at(block, loop, preserve_unit_loops=False)
    assert_structural_equal_ignore_global_symbol(
        factorized_after_reverse_compute_at, sch.mod["main"]
    )
    verify_trace_roundtrip(sch=sch, mod=factorized)


def test_reverse_compute_at_floordiv_and_floormod_indices(use_block_name):
    sch = tvm.s_tir.Schedule(floordiv_and_floormod_indices, debug_mask="all")
    A = sch.get_sblock("A")
    B = sch.get_sblock("B")
    sch.reverse_compute_at(B, sch.get_loops(A)[0])
    assert_structural_equal_ignore_global_symbol(
        floordiv_and_floormod_indices_after_reverse_compute_at, sch.mod["main"]
    )
    verify_trace_roundtrip(sch=sch, mod=floordiv_and_floormod_indices)


def test_reverse_compute_at_floordiv_and_floormod_recursive(use_block_name):
    sch = tvm.s_tir.Schedule(recursive_floordiv_floormod, debug_mask="all")
    write_block = sch.get_sblock("Out")
    sch.reverse_compute_at(write_block, sch.get_loops("In")[2])
    assert_structural_equal_ignore_global_symbol(
        recursive_floordiv_floormod_after_reverse_compute_at, sch.mod["main"]
    )
    verify_trace_roundtrip(sch=sch, mod=recursive_floordiv_floormod)


def test_read_out_of_bound(use_block_name):
    sch = tvm.s_tir.Schedule(read_out_of_bound, debug_mask="all")
    block = sch.get_sblock("B")
    (loop,) = sch.get_loops(sch.get_sblock("C"))
    sch.compute_at(block, loop)
    assert_structural_equal_ignore_global_symbol(
        read_out_of_bound_after_compute_at, sch.mod["main"]
    )
    verify_trace_roundtrip(sch=sch, mod=read_out_of_bound)


def test_compact_dataflow(use_block_name):
    sch = tvm.s_tir.Schedule(not_all_compact_data_flow, debug_mask="all")
    block = sch.get_sblock("B")
    _, loop = sch.get_loops(sch.get_sblock("C_1"))
    sch.compute_at(block, loop)
    assert_structural_equal_ignore_global_symbol(
        not_all_compact_data_flow_after_compute_at, sch.mod["main"]
    )
    verify_trace_roundtrip(sch=sch, mod=not_all_compact_data_flow)


def test_compute_at_simplify_static_bound(use_block_name):
    sch = tvm.s_tir.Schedule(static_bound, debug_mask="all")
    block = sch.get_sblock("B")
    loop, _ = sch.get_loops(sch.get_sblock("C"))
    sch.compute_at(block, loop, preserve_unit_loops=True)
    assert_structural_equal_ignore_global_symbol(static_bound_after_compute_at, sch.mod["main"])
    verify_trace_roundtrip(sch=sch, mod=static_bound)


def test_compute_at_simplify_symbolic_predicate():
    @tvm.script.ir_module
    class Before:
        @T.prim_func(s_tir=True)
        def main(x: T.handle, y: T.handle, n: T.int64):
            X = T.match_buffer(x, (T.int64(8), n * 32), "float32")
            Y = T.match_buffer(y, (T.int64(8), n * 32), "float32")
            for i, k in T.grid(T.int64(8), n * 32):
                with T.sblock("Y"):
                    vi, vk = T.axis.remap("SS", [i, k])
                    Y[vi, vk] = X[vi, vk]

    @tvm.script.ir_module
    class After:
        @T.prim_func(s_tir=True)
        def main(x: T.handle, y: T.handle, n: T.int64):
            X = T.match_buffer(x, (T.int64(8), n * T.int64(32)))
            Y = T.match_buffer(y, (T.int64(8), n * T.int64(32)))
            X_global = T.sblock_alloc_buffer((T.int64(8), n * T.int64(32)))

            for i, k_0 in T.grid(T.int64(8), n):
                for ax0 in range(T.int64(32)):
                    with T.sblock("X_global"):
                        v0 = T.axis.spatial(T.int64(8), i)
                        v1 = T.axis.spatial(n * T.int64(32), k_0 * T.int64(32) + ax0)
                        X_global[v0, v1] = X[v0, v1]
                for k_1 in range(T.int64(32)):
                    with T.sblock("Y"):
                        vi = T.axis.spatial(T.int64(8), i)
                        vk = T.axis.spatial(n * T.int64(32), k_0 * T.int64(32) + k_1)
                        Y[vi, vk] = X_global[vi, vk]

    sch = tvm.s_tir.Schedule(Before, debug_mask="all")
    block = sch.get_sblock("Y")
    i, k = sch.get_loops(sch.get_sblock("Y"))
    ko, ki = sch.split(k, [None, 32])
    XX = sch.cache_read(block, 0, "global")
    sch.compute_at(XX, ko)
    tvm.ir.assert_structural_equal(sch.mod, After)


def test_compute_at_non_perfect_channel_group(use_block_name):
    @T.prim_func(s_tir=True)
    def grouped_channel_bias(
        X: T.Buffer((720, 8, 8), "float32"), Y: T.Buffer((720, 8, 8), "float32")
    ):
        B = T.sblock_alloc_buffer([45], dtype="float32", scope="")
        for i in T.grid(45):
            with T.sblock("init"):
                vi = T.axis.remap("S", [i])
                B[vi] = vi
        for c_o, h, w, c_i in T.grid(2, 8, 8, 360):
            with T.sblock("compute"):
                hh, ww = T.axis.remap("SS", [h, w])
                cc = T.axis.spatial(720, c_o * 360 + c_i)
                Y[cc, hh, ww] = X[cc, hh, ww] + B[cc // 16]

    @T.prim_func(s_tir=True)
    def grouped_channel_bias_non_perfect_tiled(
        X: T.Buffer((720, 8, 8), "float32"), Y: T.Buffer((720, 8, 8), "float32")
    ):
        B = T.sblock_alloc_buffer([45], dtype="float32")
        for c_o in range(2):
            for ax0 in range(23):
                with T.sblock("init"):
                    vi = T.axis.spatial(45, c_o * 22 + ax0)
                    B[vi] = vi
            for h, w, c_i in T.grid(8, 8, 360):
                with T.sblock("compute"):
                    hh, ww = T.axis.remap("SS", [h, w])
                    cc = T.axis.spatial(720, c_o * 360 + c_i)
                    Y[cc, hh, ww] = X[cc, hh, ww] + B[cc // 16]

    sch = tvm.s_tir.Schedule(grouped_channel_bias, debug_mask="all")
    loop = sch.get_loops(sch.get_sblock("compute"))[0]
    sch.compute_at(sch.get_sblock("init"), loop)
    assert_structural_equal_ignore_global_symbol(
        sch.mod["main"], grouped_channel_bias_non_perfect_tiled
    )


def test_fail_subtree_complete_block(use_block_name):
    sch = tvm.s_tir.Schedule(fail_subtree_compact_dataflow, debug_mask="all")
    block = sch.get_sblock("B_0")
    loop, _ = sch.get_loops(sch.get_sblock("C"))
    with pytest.raises(tvm.s_tir.ScheduleError, match="complete block"):
        sch.compute_at(block, loop)


def test_fail_not_in_same_scope(use_block_name):
    sch = tvm.s_tir.Schedule(blockized_1, debug_mask="all")
    block = "B" if use_block_name else sch.get_sblock("B")
    loop, _ = sch.get_loops(sch.get_sblock("C_inner"))
    with pytest.raises(tvm.s_tir.ScheduleError, match="same block scope"):
        sch.compute_at(block, loop)


def test_fail_loop_is_ancestor_of_block(use_block_name):
    sch = tvm.s_tir.Schedule(two_elementwise, debug_mask="all")
    block = "B" if use_block_name else sch.get_sblock("B")
    loop, _ = sch.get_loops(sch.get_sblock("B"))
    with pytest.raises(tvm.s_tir.ScheduleError, match="ancestor of block"):
        sch.compute_at(block, loop)


def test_fail_output_block(use_block_name):
    sch = tvm.s_tir.Schedule(tiled, debug_mask="all")
    block = "C" if use_block_name else sch.get_sblock("C")
    loop, _, _, _ = sch.get_loops(sch.get_sblock("B"))
    with pytest.raises(tvm.s_tir.ScheduleError, match="output block"):
        sch.compute_at(block, loop)


def test_fail_all_consumers_under_loop(use_block_name):
    sch = tvm.s_tir.Schedule(fail_all_consumers_under_loop, debug_mask="all")
    block = "B" if use_block_name else sch.get_sblock("B")
    loop, _ = sch.get_loops(sch.get_sblock("C"))
    with pytest.raises(tvm.s_tir.ScheduleError, match="requires all the consumer"):
        sch.compute_at(block, loop)


def test_fail_all_producers_under_loop(use_block_name):
    sch = tvm.s_tir.Schedule(fail_all_producers_under_loop, debug_mask="all")
    block = "D" if use_block_name else sch.get_sblock("D")
    loop, _ = sch.get_loops(sch.get_sblock("C"))
    with pytest.raises(tvm.s_tir.ScheduleError, match="requires all the producer"):
        sch.reverse_compute_at(block, loop)


def test_compute_at_int64_loop(use_block_name):
    def _create_prim_func():
        n = te.var("n", dtype="int64")
        m = te.var("m", dtype="int64")
        A = te.placeholder((n, m), name="A", dtype="float32")
        B = te.placeholder((n, m), name="B", dtype="float32")
        C = te.compute((n, m), lambda i, j: A[i, j] + B[i, j], name="C")
        D = te.compute((n, m), lambda i, j: C[i, j] + 1.0, name="D")
        return te.create_prim_func([A, B, D])

    mod = _create_prim_func()
    sch = tvm.s_tir.Schedule(mod, debug_mask="all")
    block_c = "C" if use_block_name else sch.get_sblock("C")
    block_d = "D" if use_block_name else sch.get_sblock("D")
    i, _ = sch.get_loops(block_d)
    sch.compute_at(block_c, i)
    verify_trace_roundtrip(sch=sch, mod=mod)


def test_compute_at_to_index():
    @T.prim_func(s_tir=True)
    def multi_producers_conv(
        data: T.Buffer((1, 3, 224, 224), "int8"),
        w: T.Buffer((16, 3, 7, 7), "int8"),
        conv: T.Buffer((1, 16, 112, 112), "int32"),
    ) -> None:
        pad = T.sblock_alloc_buffer([1, 3, 230, 230], dtype="int8")
        wbuf = T.sblock_alloc_buffer([16, 3, 7, 7], dtype="int8")
        for i0, i1, i2, i3 in T.grid(1, 3, 230, 230):
            with T.sblock("pad"):
                i0_1, i1_1, i2_1, i3_1 = T.axis.remap("SSSS", [i0, i1, i2, i3])
                T.reads(data[i0_1, i1_1, i2_1 - 3, i3_1 - 3])
                T.writes(pad[i0_1, i1_1, i2_1, i3_1])
                pad[i0_1, i1_1, i2_1, i3_1] = T.if_then_else(
                    3 <= i2_1 and i2_1 < 227 and 3 <= i3_1 and i3_1 < 227,
                    data[i0_1, i1_1, i2_1 - 3, i3_1 - 3],
                    T.int8(0),
                    dtype="int8",
                )
        for i0 in T.serial(1):
            for ax0, ax1, ax2, ax3 in T.grid(16, 3, 7, 7):
                with T.sblock("wbuf"):
                    v0, v1, v2, v3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
                    T.reads(w[v0, v1, v2, v3])
                    T.writes(wbuf[v0, v1, v2, v3])
                    wbuf[v0, v1, v2, v3] = w[v0, v1, v2, v3]
            for i1, i2, i3, i4, i5, i6 in T.grid(16, 112, 112, 3, 7, 7):
                with T.sblock("conv"):
                    nn, ff, yy, xx, rc, ry, rx = T.axis.remap(
                        "SSSSRRR", [i0, i1, i2, i3, i4, i5, i6]
                    )
                    T.reads(pad[nn, rc, yy * 2 + ry, xx * 2 + rx], wbuf[ff, rc, ry, rx])
                    T.writes(conv[nn, ff, yy, xx])
                    with T.init():
                        conv[nn, ff, yy, xx] = 0
                    conv[nn, ff, yy, xx] = conv[nn, ff, yy, xx] + T.cast(
                        pad[nn, rc, yy * 2 + ry, xx * 2 + rx], "int32"
                    ) * T.cast(wbuf[ff, rc, ry, rx], "int32")

    @T.prim_func(s_tir=True)
    def multi_producers_after_compute_at(
        data: T.Buffer((1, 3, 224, 224), "int8"),
        w: T.Buffer((16, 3, 7, 7), "int8"),
        conv: T.Buffer((1, 16, 112, 112), "int32"),
    ) -> None:
        pad = T.sblock_alloc_buffer([1, 3, 230, 230], dtype="int8")
        wbuf = T.sblock_alloc_buffer([16, 3, 7, 7], dtype="int8")
        for i0 in T.serial(1):
            for ax0, ax1, ax2 in T.grid(3, 229, 229):
                with T.sblock("pad"):
                    i0_1 = T.axis.spatial(1, 0)
                    i1_1 = T.axis.spatial(3, ax0)
                    i2_1 = T.axis.spatial(230, ax1)
                    i3_1 = T.axis.spatial(230, ax2)
                    T.reads(data[i0_1, i1_1, i2_1 - 3, i3_1 - 3])
                    T.writes(pad[i0_1, i1_1, i2_1, i3_1])
                    pad[i0_1, i1_1, i2_1, i3_1] = T.if_then_else(
                        3 <= i2_1 and i2_1 < 227 and 3 <= i3_1 and i3_1 < 227,
                        data[i0_1, i1_1, i2_1 - 3, i3_1 - 3],
                        T.int8(0),
                        dtype="int8",
                    )
            for ax0, ax1, ax2, ax3 in T.grid(16, 3, 7, 7):
                with T.sblock("wbuf"):
                    v0, v1, v2, v3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
                    T.reads(w[v0, v1, v2, v3])
                    T.writes(wbuf[v0, v1, v2, v3])
                    wbuf[v0, v1, v2, v3] = w[v0, v1, v2, v3]
            for i1, i2, i3, i4, i5, i6 in T.grid(16, 112, 112, 3, 7, 7):
                with T.sblock("conv"):
                    nn, ff, yy, xx, rc, ry, rx = T.axis.remap(
                        "SSSSRRR", [i0, i1, i2, i3, i4, i5, i6]
                    )
                    T.reads(pad[nn, rc, yy * 2 + ry, xx * 2 + rx], wbuf[ff, rc, ry, rx])
                    T.writes(conv[nn, ff, yy, xx])
                    with T.init():
                        conv[nn, ff, yy, xx] = 0
                    conv[nn, ff, yy, xx] = conv[nn, ff, yy, xx] + T.cast(
                        pad[nn, rc, yy * 2 + ry, xx * 2 + rx], "int32"
                    ) * T.cast(wbuf[ff, rc, ry, rx], "int32")

    sch = tvm.s_tir.Schedule(multi_producers_conv, debug_mask="all")
    block_c = sch.get_sblock("pad")
    axis = sch.get_loops("conv")[0]
    sch.compute_at(block_c, axis, index=-2)
    assert_structural_equal_ignore_global_symbol(multi_producers_after_compute_at, sch.mod["main"])


def test_reverse_compute_at_to_index():
    @T.prim_func(s_tir=True)
    def main(A: T.Buffer((128, 128), "float32"), D: T.Buffer((128, 128), "float32")) -> None:
        B = T.sblock_alloc_buffer([128, 128], dtype="float32")
        C = T.sblock_alloc_buffer([128, 128], dtype="float32")
        for i_0, j_0, i_1 in T.grid(8, 8, 16):
            for j_1 in T.serial(16):
                with T.sblock("B"):
                    vi = T.axis.spatial(128, i_0 * 16 + i_1)
                    vj = T.axis.spatial(128, j_0 * 16 + j_1)
                    T.reads(A[vi, vj])
                    T.writes(B[vi, vj])
                    B[vi, vj] = A[vi, vj] * T.float32(2)
            for ax0 in T.serial(16):
                with T.sblock("C"):
                    vi = T.axis.spatial(128, i_0 * 16 + i_1)
                    vj = T.axis.spatial(128, j_0 * 16 + ax0)
                    T.reads(B[vi, vj])
                    T.writes(C[vi, vj])
                    C[vi, vj] = B[vi, vj] + T.float32(1)
        for i, j in T.grid(128, 128):
            with T.sblock("D"):
                vi, vj = T.axis.remap("SS", [i, j])
                T.reads(B[vi, vj])
                T.writes(D[vi, vj])
                D[vi, vj] = B[vi, vj] + T.float32(1)

    @T.prim_func(s_tir=True)
    def main_reverse_compute_at(
        A: T.Buffer((128, 128), "float32"), D: T.Buffer((128, 128), "float32")
    ) -> None:
        B = T.sblock_alloc_buffer([128, 128], dtype="float32")
        C = T.sblock_alloc_buffer([128, 128], dtype="float32")
        for i_0, j_0, i_1 in T.grid(8, 8, 16):
            for j_1 in T.serial(16):
                with T.sblock("B"):
                    vi = T.axis.spatial(128, i_0 * 16 + i_1)
                    vj = T.axis.spatial(128, j_0 * 16 + j_1)
                    T.reads(A[vi, vj])
                    T.writes(B[vi, vj])
                    B[vi, vj] = A[vi, vj] * T.float32(2)
            for ax0 in T.serial(16):
                with T.sblock("D"):
                    vi = T.axis.spatial(128, i_0 * 16 + i_1)
                    vj = T.axis.spatial(128, j_0 * 16 + ax0)
                    T.reads(B[vi, vj])
                    T.writes(D[vi, vj])
                    D[vi, vj] = B[vi, vj] + T.float32(1)
            for ax0 in T.serial(16):
                with T.sblock("C"):
                    vi = T.axis.spatial(128, i_0 * 16 + i_1)
                    vj = T.axis.spatial(128, j_0 * 16 + ax0)
                    T.reads(B[vi, vj])
                    T.writes(C[vi, vj])
                    C[vi, vj] = B[vi, vj] + T.float32(1)

    sch = tvm.s_tir.Schedule(main, debug_mask="all")
    block_c = sch.get_sblock("D")
    axis = sch.get_loops("B")[2]
    sch.reverse_compute_at(block_c, axis, index=1)
    assert_structural_equal_ignore_global_symbol(main_reverse_compute_at, sch.mod["main"])


def test_reverse_compute_at_with_unit_loop():
    @T.prim_func(s_tir=True)
    def main(A: T.Buffer((128, 128), "float32"), D: T.Buffer((1, 2, 1), "float32")) -> None:
        B = T.sblock_alloc_buffer([128, 128], dtype="float32")
        for i_0, j_0, i_1 in T.grid(T.int64(8), T.int64(8), T.int64(16)):
            for j_1 in T.serial(T.int64(16)):
                with T.sblock("B"):
                    vi = T.axis.spatial(T.int64(128), i_0 * T.int64(16) + i_1)
                    vj = T.axis.spatial(T.int64(128), j_0 * T.int64(16) + j_1)
                    T.reads(A[vi, vj])
                    T.writes(B[vi, vj])
                    B[vi, vj] = A[vi, vj] * T.float32(2)
        for ax0, ax1, ax2 in T.grid(T.int64(1), T.int64(2), T.int64(1)):
            with T.sblock("D"):
                v0, v1, v2 = T.axis.remap("SSS", [ax0, ax1, ax2])
                T.reads(B[v0, v1])
                T.writes(D[v0, v1, v2])
                D[v0, v1, v2] = B[v0, v1] + T.float32(1)

    @T.prim_func(s_tir=True)
    def main_reverse_compute_at(
        A: T.Buffer((128, 128), "float32"), D: T.Buffer((1, 2, 1), "float32")
    ):
        B = T.sblock_alloc_buffer([128, 128], dtype="float32")
        for i_0, j_0, i_1 in T.grid(T.int64(8), T.int64(8), T.int64(16)):
            for j_1 in T.serial(T.int64(16)):
                with T.sblock("B"):
                    vi = T.axis.spatial(T.int64(128), i_0 * T.int64(16) + i_1)
                    vj = T.axis.spatial(T.int64(128), j_0 * T.int64(16) + j_1)
                    T.reads(A[vi, vj])
                    T.writes(B[vi, vj])
                    B[vi, vj] = A[vi, vj] * T.float32(2)
            for ax0, ax1, ax2 in T.grid(T.int64(1), T.int64(16), T.int64(1)):
                with T.sblock("D"):
                    T.where(
                        i_0 * T.int64(16) + i_1 < T.int64(1)
                        and j_0 * T.int64(16) + ax1 < T.int64(2)
                    )
                    v0 = T.axis.spatial(T.int64(1), i_0 * T.int64(16) + i_1 + ax0)
                    v1 = T.axis.spatial(T.int64(2), j_0 * T.int64(16) + ax1)
                    v2 = T.axis.spatial(T.int64(1), ax2)
                    T.reads(B[v0, v1])
                    T.writes(D[v0, v1, v2])
                    D[v0, v1, v2] = B[v0, v1] + T.float32(1)

    sch = tvm.s_tir.Schedule(main, debug_mask="all")
    block_d = sch.get_sblock("D")
    axis = sch.get_loops("B")[2]
    sch.reverse_compute_at(block_d, axis, preserve_unit_loops=True, index=1)
    assert_structural_equal_ignore_global_symbol(main_reverse_compute_at, sch.mod["main"])


def test_reverse_compute_at_layout_trans():
    @T.prim_func(s_tir=True)
    def before(A: T.Buffer((1, 3, 5, 5, 16), "float32"), C: T.Buffer((1, 6, 5, 5, 8), "float32")):
        B = T.sblock_alloc_buffer((1, 3, 5, 5, 16))
        for i0, i1, i2, i3, i4 in T.grid(1, 3, 5, 5, 16):
            with T.sblock("compute"):
                v_i0, v_i1, v_i2, v_i3, v_i4 = T.axis.remap("SSSSS", [i0, i1, i2, i3, i4])
                B[v_i0, v_i1, v_i2, v_i3, v_i4] = A[v_i0, v_i1, v_i2, v_i3, v_i4] + T.float32(1)
        for ax0, ax1, ax2, ax3, ax4 in T.grid(1, 6, 5, 5, 8):
            with T.sblock("T_layout_trans"):
                v_ax0, v_ax1, v_ax2, v_ax3, v_ax4 = T.axis.remap("SSSSS", [ax0, ax1, ax2, ax3, ax4])
                C[v_ax0, v_ax1, v_ax2, v_ax3, v_ax4] = B[
                    v_ax0, (v_ax1 * 8 + v_ax4) // 16, v_ax2, v_ax3, (v_ax1 * 8 + v_ax4) % 16
                ]

    @T.prim_func(s_tir=True)
    def after(A: T.Buffer((1, 3, 5, 5, 16), "float32"), C: T.Buffer((1, 6, 5, 5, 8), "float32")):
        B = T.sblock_alloc_buffer((1, 3, 5, 5, 16))
        for i0, i1 in T.grid(1, 3):
            for i2, i3, i4 in T.grid(5, 5, 16):
                with T.sblock("compute"):
                    v_i0, v_i1, v_i2, v_i3, v_i4 = T.axis.remap("SSSSS", [i0, i1, i2, i3, i4])
                    B[v_i0, v_i1, v_i2, v_i3, v_i4] = A[v_i0, v_i1, v_i2, v_i3, v_i4] + T.float32(1)
            for ax0, ax1, ax2, ax3 in T.grid(2, 5, 5, 8):
                with T.sblock("T_layout_trans"):
                    v_ax0 = T.axis.spatial(1, 0)
                    v_ax1 = T.axis.spatial(6, i1 * 2 + ax0)
                    v_ax2, v_ax3, v_ax4 = T.axis.remap("SSS", [ax1, ax2, ax3])
                    C[v_ax0, v_ax1, v_ax2, v_ax3, v_ax4] = B[
                        v_ax0, (v_ax1 * 8 + v_ax4) // 16, v_ax2, v_ax3, (v_ax1 * 8 + v_ax4) % 16
                    ]

    sch = tvm.s_tir.Schedule(before, debug_mask="all")
    trans = sch.get_sblock("T_layout_trans")
    axis = sch.get_loops("compute")[1]
    sch.reverse_compute_at(trans, axis)
    assert_structural_equal_ignore_global_symbol(after, sch.mod["main"])
    verify_trace_roundtrip(sch=sch, mod=before)


def test_shape_var_as_bound():
    # fmt: off
    @T.prim_func(s_tir=True)
    def before(a: T.handle, b: T.handle, c: T.handle):
        n = T.int32()
        A = T.match_buffer(a, (32, 1, 128))
        B = T.match_buffer(b, (32, n, 128))
        C = T.match_buffer(c, (32, 1, n))
        # with T.sblock("root"):
        C_rf = T.sblock_alloc_buffer((128, 32, 1, n))
        for ax0_ax1_fused, ax2_fused_1, ax2_fused_0 in T.grid(n * 32, 128, 1):
            with T.sblock("NT_matmul_rf"):
                vax2_fused_1 = T.axis.spatial(128, ax2_fused_1)
                v0 = T.axis.spatial(32, ax0_ax1_fused // n)
                v1 = T.axis.spatial(n, ax0_ax1_fused % n)
                vax2_fused_0 = T.axis.reduce(1, ax2_fused_0)
                T.reads(A[v0, 0, vax2_fused_0 * 128 + vax2_fused_1], B[v0, v1, vax2_fused_0 * 128 + vax2_fused_1])
                T.writes(C_rf[vax2_fused_1, v0, 0, v1])
                with T.init():
                    C_rf[vax2_fused_1, v0, 0, v1] = T.float32(0)
                C_rf[vax2_fused_1, v0, 0, v1] = C_rf[vax2_fused_1, v0, 0, v1] + A[v0, 0, vax2_fused_0 * 128 + vax2_fused_1] * B[v0, v1, vax2_fused_0 * 128 + vax2_fused_1]
        for ax0_ax1_fused, ax2_fused_1 in T.grid(n * 32, 128):
            with T.sblock("NT_matmul"):
                vax2_fused_1 = T.axis.reduce(128, ax2_fused_1)
                v0 = T.axis.spatial(32, ax0_ax1_fused // n)
                v1 = T.axis.spatial(n, ax0_ax1_fused % n)
                T.reads(C_rf[vax2_fused_1, v0, 0, v1])
                T.writes(C[v0, 0, v1])
                with T.init():
                    C[v0, 0, v1] = T.float32(0)
                C[v0, 0, v1] = C[v0, 0, v1] + C_rf[vax2_fused_1, v0, 0, v1]

    @T.prim_func(s_tir=True)
    def expected(A: T.Buffer((32, 1, 128), "float32"), b: T.handle, c: T.handle):
        n = T.int32()
        B = T.match_buffer(b, (32, n, 128))
        C = T.match_buffer(c, (32, 1, n))
        # with T.sblock("root"):
        C_rf = T.sblock_alloc_buffer((128, 32, 1, n))
        for ax0_ax1_fused in range(n * 32):
            for ax2_fused_1, ax2_fused_0 in T.grid(128, 1):
                with T.sblock("NT_matmul_rf"):
                    vax2_fused_1 = T.axis.spatial(128, ax2_fused_1)
                    v0 = T.axis.spatial(32, ax0_ax1_fused // n)
                    v1 = T.axis.spatial(n, ax0_ax1_fused % n)
                    vax2_fused_0 = T.axis.reduce(1, ax2_fused_0)
                    T.reads(A[v0, 0, vax2_fused_0 * 128 + vax2_fused_1], B[v0, v1, vax2_fused_0 * 128 + vax2_fused_1])
                    T.writes(C_rf[vax2_fused_1, v0, 0, v1])
                    with T.init():
                        C_rf[vax2_fused_1, v0, 0, v1] = T.float32(0)
                    C_rf[vax2_fused_1, v0, 0, v1] = C_rf[vax2_fused_1, v0, 0, v1] + A[v0, 0, vax2_fused_0 * 128 + vax2_fused_1] * B[v0, v1, vax2_fused_0 * 128 + vax2_fused_1]
            for ax0, ax1, ax2 in T.grid(128, 1, 1):
                with T.sblock("NT_matmul"):
                    vax2_fused_1 = T.axis.reduce(128, ax0)
                    v0 = T.axis.spatial(32, ax0_ax1_fused // n + ax1)
                    v1 = T.axis.spatial(n, ax0_ax1_fused % n + ax2)
                    T.reads(C_rf[vax2_fused_1, v0, 0, v1])
                    T.writes(C[v0, 0, v1])
                    with T.init():
                        C[v0, 0, v1] = T.float32(0)
                    C[v0, 0, v1] = C[v0, 0, v1] + C_rf[vax2_fused_1, v0, 0, v1]
    # fmt: on
    sch = tvm.s_tir.Schedule(before.with_attr("global_symbol", "main"), debug_mask="all")
    block = sch.get_sblock("NT_matmul")
    loop, _, _ = sch.get_loops(sch.get_sblock("NT_matmul_rf"))
    sch.reverse_compute_at(block, loop, preserve_unit_loops=True)
    tvm.ir.assert_structural_equal(
        sch.mod["main"], expected.with_attr("global_symbol", "main"), True
    )


def test_compute_at_sliced_concatenate():
    @T.prim_func(s_tir=True)
    def before():
        X = T.sblock_alloc_buffer((1, 16, 28, 64), "float32")
        Y = T.sblock_alloc_buffer((1, 32, 28, 64), "float32")
        Z = T.sblock_alloc_buffer((1, 53, 28, 64), "float32")
        Concat = T.sblock_alloc_buffer((1, 101, 28, 64), "float32")
        Slice = T.sblock_alloc_buffer((1, 87, 28, 64), "float32")
        for ax0, ax1, ax2, ax3 in T.grid(1, 16, 28, 64):
            with T.sblock("compute"):
                v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
                X[v_ax0, v_ax1, v_ax2, v_ax3] = 1.0
        for ax0, ax1, ax2, ax3 in T.grid(1, 101, 28, 64):
            with T.sblock("T_concat"):
                v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
                Concat[v_ax0, v_ax1, v_ax2, v_ax3] = T.if_then_else(
                    85 <= v_ax1,
                    X[v_ax0, v_ax1 - 85, v_ax2, v_ax3],
                    T.if_then_else(
                        53 <= v_ax1,
                        Y[v_ax0, v_ax1 - 53, v_ax2, v_ax3],
                        Z[v_ax0, v_ax1, v_ax2, v_ax3],
                    ),
                )
        for ax0, ax1, ax2, ax3 in T.grid(1, 87, 28, 64):
            with T.sblock("T_strided_slice"):
                v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
                Slice[v_ax0, v_ax1, v_ax2, v_ax3] = Concat[v_ax0, v_ax1, v_ax2, v_ax3]

    @T.prim_func(s_tir=True)
    def expect():
        X = T.sblock_alloc_buffer((1, 16, 28, 64))
        Y = T.sblock_alloc_buffer((1, 32, 28, 64))
        Z = T.sblock_alloc_buffer((1, 53, 28, 64))
        Concat = T.sblock_alloc_buffer((1, 101, 28, 64))
        Slice = T.sblock_alloc_buffer((1, 87, 28, 64))
        for ax0 in range(1):
            for ax0_1, ax1, ax2 in T.grid(2, 28, 64):
                with T.sblock("compute"):
                    v_ax0 = T.axis.spatial(1, 0)
                    v_ax1 = T.axis.spatial(16, ax0_1)
                    v_ax2, v_ax3 = T.axis.remap("SS", [ax1, ax2])
                    X[v_ax0, v_ax1, v_ax2, v_ax3] = T.float32(1)
            for ax0_1, ax1, ax2 in T.grid(87, 28, 64):
                with T.sblock("T_concat"):
                    v_ax0 = T.axis.spatial(1, 0)
                    v_ax1 = T.axis.spatial(101, ax0_1)
                    v_ax2, v_ax3 = T.axis.remap("SS", [ax1, ax2])
                    Concat[v_ax0, v_ax1, v_ax2, v_ax3] = T.if_then_else(
                        85 <= v_ax1,
                        X[v_ax0, v_ax1 - 85, v_ax2, v_ax3],
                        T.if_then_else(
                            53 <= v_ax1,
                            Y[v_ax0, v_ax1 - 53, v_ax2, v_ax3],
                            Z[v_ax0, v_ax1, v_ax2, v_ax3],
                        ),
                    )
            for ax1, ax2, ax3 in T.grid(87, 28, 64):
                with T.sblock("T_strided_slice"):
                    v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
                    Slice[v_ax0, v_ax1, v_ax2, v_ax3] = Concat[v_ax0, v_ax1, v_ax2, v_ax3]

    sch = tvm.s_tir.Schedule(before, debug_mask="all")
    blk1 = sch.get_sblock("compute")
    blk2 = sch.get_sblock("T_concat")
    blk3 = sch.get_sblock("T_strided_slice")
    loop = sch.get_loops(blk3)[0]
    sch.compute_at(blk2, loop)
    sch.compute_at(blk1, loop)
    after = sch.mod["main"]
    assert_structural_equal_ignore_global_symbol(expect, after)
    verify_trace_roundtrip(sch=sch, mod=before)


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