# 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: E741, F401
import sys

import pytest

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

# pylint: disable=no-member,invalid-name,unused-variable

########## Function before schedule ##########


@T.prim_func(s_tir=True)
def elementwise(a: T.handle, c: T.handle) -> None:
    A = T.match_buffer(a, (128, 128))
    B = T.sblock_alloc_buffer((128, 128))
    C = T.match_buffer(c, (128, 128))
    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 elementwise_shape_int64(a: T.handle, c: T.handle) -> None:
    A = T.match_buffer(a, (T.int64(128), T.int64(128)))
    B = T.sblock_alloc_buffer((T.int64(128), T.int64(128)))
    C = T.match_buffer(c, (T.int64(128), T.int64(128)))
    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 elementwise_reindex_cache_read(
    A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")
):
    B = T.sblock_alloc_buffer((128, 128))
    B_shared = T.sblock_alloc_buffer((128, 64, 2), scope="shared")
    for i, j in T.grid(128, 128):
        with T.sblock("B"):
            vi, vj = T.axis.remap("SS", [i, j])
            T.reads(A[vi, vj])
            T.writes(B[vi, vj])
            B[vi, vj] = A[vi, vj] * T.float32(2)
    for i, j in T.grid(128, 128):
        with T.sblock("B_shared"):
            vi, vj = T.axis.remap("SS", [i, j])
            T.reads(B[vi, vj])
            T.writes(B_shared[vj, vi // 2, vi % 2])
            B_shared[vj, vi // 2, vi % 2] = B[vi, vj]
    for i, j in T.grid(128, 128):
        with T.sblock("C"):
            vi, vj = T.axis.remap("SS", [i, j])
            T.reads(B_shared[vj, vi // 2, vi % 2])
            T.writes(C[vi, vj])
            C[vi, vj] = B_shared[vj, vi // 2, vi % 2] + T.float32(1)


@T.prim_func(s_tir=True)
def elementwise_reindex_cache_write(
    A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")
):
    B = T.sblock_alloc_buffer((128, 128))
    B_shared = T.sblock_alloc_buffer((128, 128), scope="shared")
    for i, j in T.grid(128, 128):
        with T.sblock("B"):
            vi, vj = T.axis.remap("SS", [i, j])
            T.reads(A[vi, vj])
            T.writes(B_shared[vj, vi])
            B_shared[vj, vi] = A[vi, vj] * T.float32(2)
    for i, j in T.grid(128, 128):
        with T.sblock("B_shared"):
            vi, vj = T.axis.remap("SS", [i, j])
            T.reads(B_shared[vj, vi])
            T.writes(B[vi, vj])
            B[vi, vj] = B_shared[vj, vi]
    for i, j in T.grid(128, 128):
        with T.sblock("C"):
            vi, vj = T.axis.remap("SS", [i, j])
            T.reads(B[vi, vj])
            T.writes(C[vi, vj])
            C[vi, vj] = B[vi, vj] + T.float32(1)


@T.prim_func(s_tir=True)
def reduce(A: T.Buffer((128, 128, 128, 128), "float32"), C: T.Buffer((128, 128), "float32")):
    B = T.sblock_alloc_buffer((128, 128, 128), dtype="float32")
    for i, j, k in T.grid(128, 128, 128):
        for l in range(128):
            with T.sblock("B"):
                vi, vj, vk, vl = T.axis.remap("SSSR", [i, j, k, l])
                with T.init():
                    B[vi, vj, vk] = T.float32(0)
                B[vi, vj, vk] = B[vi, vj, vk] + A[vi, vj, vk, vl]
        with T.sblock("C"):
            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
            with T.init():
                C[vi, vj] = T.float32(0)
            C[vi, vj] = C[vi, vj] + B[vi, vj, vk]


@T.prim_func(s_tir=True)
def reduce_reindex_cache_write_0(
    A: T.Buffer((128, 128, 128, 128), "float32"), C: T.Buffer((128, 128), "float32")
):
    B = T.sblock_alloc_buffer((128, 128, 128))
    B_shared = T.sblock_alloc_buffer((128, 128, 128), scope="shared")
    for i, j, k in T.grid(128, 128, 128):
        for l in range(128):
            with T.sblock("B"):
                vi, vj, vk, vl = T.axis.remap("SSSR", [i, j, k, l])
                T.reads(A[vi, vj, vk, vl])
                T.writes(B_shared[vj, vi, vk])
                with T.init():
                    B_shared[vj, vi, vk] = T.float32(0)
                B_shared[vj, vi, vk] = B_shared[vj, vi, vk] + A[vi, vj, vk, vl]
        with T.sblock("B_shared"):
            vi, vj, vk = T.axis.remap("SSS", [i, j, k])
            T.reads(B_shared[vj, vi, vk])
            T.writes(B[vi, vj, vk])
            B[vi, vj, vk] = B_shared[vj, vi, vk]
        with T.sblock("C"):
            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
            T.reads(B[vi, vj, vk])
            T.writes(C[vi, vj])
            with T.init():
                C[vi, vj] = T.float32(0)
            C[vi, vj] = C[vi, vj] + B[vi, vj, vk]


@T.prim_func(s_tir=True)
def reduce_reindex_cache_write_1(
    A: T.Buffer((128, 128, 128, 128), "float32"), C: T.Buffer((128, 128), "float32")
):
    B = T.sblock_alloc_buffer((128, 128, 128))
    B_shared = T.sblock_alloc_buffer((128, 128, 128), scope="shared")
    C_shared = T.sblock_alloc_buffer((128, 128), scope="shared")
    for i, j, k in T.grid(128, 128, 128):
        for l in range(128):
            with T.sblock("B"):
                vi, vj, vk, vl = T.axis.remap("SSSR", [i, j, k, l])
                T.reads(A[vi, vj, vk, vl])
                T.writes(B_shared[vj, vi, vk])
                with T.init():
                    B_shared[vj, vi, vk] = T.float32(0)
                B_shared[vj, vi, vk] = B_shared[vj, vi, vk] + A[vi, vj, vk, vl]
        with T.sblock("B_shared"):
            vi, vj, vk = T.axis.remap("SSS", [i, j, k])
            T.reads(B_shared[vj, vi, vk])
            T.writes(B[vi, vj, vk])
            B[vi, vj, vk] = B_shared[vj, vi, vk]
        with T.sblock("C"):
            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
            T.reads(B[vi, vj, vk])
            T.writes(C_shared[vj, vi])
            with T.init():
                C_shared[vj, vi] = T.float32(0)
            C_shared[vj, vi] = C_shared[vj, vi] + B[vi, vj, vk]
    for i, j in T.grid(128, 128):
        with T.sblock("C_shared"):
            vi, vj = T.axis.remap("SS", [i, j])
            T.reads(C_shared[vj, vi])
            T.writes(C[vi, vj])
            C[vi, vj] = C_shared[vj, vi]


@T.prim_func(s_tir=True)
def func_nested_seq(b: T.handle, c: T.handle) -> None:
    A = T.sblock_alloc_buffer((128, 128))
    B = T.match_buffer(b, (128, 128))
    C = T.match_buffer(c, (128, 128))

    for i, j in T.grid(128, 128):
        with T.sblock("A"):
            vi, vj = T.axis.remap("SS", [i, j])
            A[vi, vj] = 2.0
    for i, j in T.grid(8, 8):
        for x, y in T.grid(16, 16):
            with T.sblock("B0"):
                vi = T.axis.S(128, i * 16 + x)
                vj = T.axis.S(128, j * 16 + y)
                B[vi, vj] = 1.0
        for x, y in T.grid(16, 16):
            with T.sblock("B1"):
                vi = T.axis.S(128, i * 16 + x)
                vj = T.axis.S(128, j * 16 + y)
                B[vi, vj] = A[vi, vj] + B[vi, vj]
    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] * 2.0


@T.prim_func(s_tir=True)
def access_under_scope(b: T.handle, c: T.handle) -> None:
    A = T.sblock_alloc_buffer((128, 128))
    B = T.match_buffer(b, (128, 128))
    C = T.match_buffer(c, (128, 128))

    for i0, j0 in T.grid(8, 8):
        with T.sblock("scope"):
            i, j = T.axis.remap("SS", [i0, j0])
            for x, y in T.grid(16, 16):
                with T.sblock("A"):
                    vi = T.axis.S(128, i * 16 + x)
                    vj = T.axis.S(128, j * 16 + y)
                    A[vi, vj] = 1.0
            for x, y in T.grid(16, 16):
                with T.sblock("B"):
                    vi = T.axis.S(128, i * 16 + x)
                    vj = T.axis.S(128, j * 16 + y)
                    B[vi, vj] = A[vi, vj] + 1.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] * 2.0


@T.prim_func(s_tir=True)
def opaque_access(a: T.handle, b: T.handle, c: T.handle, d: T.handle) -> None:
    A = T.match_buffer(a, (128, 128), dtype="float16")
    B = T.match_buffer(b, (128, 128), dtype="float16")
    C = T.match_buffer(c, (128, 128), dtype="float16")
    D = T.match_buffer(d, (128, 128), dtype="float16")

    for i, j in T.grid(128, 128):
        with T.sblock("load_store"):
            vi, vj = T.axis.remap("SS", [i, j])
            T.reads(A[vi, vj])
            T.writes(D[vi, vj])
            D[vi, vj] = A[vi, vj]
    for i, j in T.grid(8, 8):
        with T.sblock("opaque"):
            vi, vj = T.axis.remap("SS", [i, j])
            T.reads(A[vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16])
            T.writes(B[vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16])
            T.evaluate(
                T.tvm_load_matrix_sync(
                    B.data,
                    16,
                    16,
                    16,
                    vi * 8 + vj,
                    T.tvm_access_ptr(
                        T.type_annotation(dtype="float16"),
                        A.data,
                        vi * 2048 + vj * 16,
                        128,
                        1,
                        dtype="handle",
                    ),
                    128,
                    "row_major",
                    dtype="handle",
                )
            )
    for i, j in T.grid(8, 8):
        with T.sblock("match_buffer"):
            vi, vj = T.axis.remap("SS", [i, j])
            T.reads(A[vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16])
            T.writes(C[vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16])
            A0 = T.match_buffer(
                A[
                    vi * 16 : vi * 16 + 16,
                    vj * 16 : vj * 16 + 16,
                ],
                (16, 16),
                "float16",
                strides=[128, 1],
                offset_factor=1,
            )
            C0 = T.match_buffer(
                C[
                    vi * 16 : vi * 16 + 16,
                    vj * 16 : vj * 16 + 16,
                ],
                (16, 16),
                "float16",
                strides=[128, 1],
                offset_factor=1,
            )
            T.evaluate(
                T.tvm_load_matrix_sync(
                    C0.data,
                    16,
                    16,
                    16,
                    vi * 8 + vj,
                    T.tvm_access_ptr(
                        T.type_annotation(dtype="float16"),
                        A0.data,
                        A0.elem_offset,
                        A0.strides[0],
                        1,
                        dtype="handle",
                    ),
                    128,
                    "row_major",
                    dtype="handle",
                )
            )


@T.prim_func(s_tir=True)
def func_multi_consumer() -> None:
    A = T.sblock_alloc_buffer(128)
    B = T.sblock_alloc_buffer(128)
    C = T.sblock_alloc_buffer(128)
    for i in T.grid(8):
        for j in T.grid(16):
            with T.sblock("A"):
                vi = T.axis.S(128, i * 16 + j)
                A[vi] = 1.0
        for j in T.grid(16):
            with T.sblock("B"):
                vi = T.axis.S(128, i * 16 + j)
                B[vi] = A[vi] + 1.0
    for i in T.grid(128):
        with T.sblock("C"):
            vi = T.axis.S(128, i)
            C[vi] = A[vi]


@T.prim_func(s_tir=True)
def reindex_cache_read_multi_consumer() -> None:
    A = T.sblock_alloc_buffer((128,))
    B = T.sblock_alloc_buffer((128,))
    C = T.sblock_alloc_buffer((128,))
    A_shared = T.sblock_alloc_buffer((4, 32), scope="shared")
    for i in range(8):
        for j in range(16):
            with T.sblock("A"):
                vi = T.axis.spatial(128, i * 16 + j)
                T.reads()
                T.writes(A[vi])
                A[vi] = T.float32(1)
        for j in range(16):
            with T.sblock("A_shared"):
                vi = T.axis.spatial(128, i * 16 + j)
                T.reads(A[vi])
                T.writes(A_shared[vi // 32, vi % 32])
                A_shared[vi // 32, vi % 32] = A[vi]
        for j in range(16):
            with T.sblock("B"):
                vi = T.axis.spatial(128, i * 16 + j)
                T.reads(A_shared[vi // 32, vi % 32])
                T.writes(B[vi])
                B[vi] = A_shared[vi // 32, vi % 32] + T.float32(1)
    for i in range(128):
        with T.sblock("C"):
            vi = T.axis.spatial(128, i)
            T.reads(A[vi])
            T.writes(C[vi])
            C[vi] = A[vi]


@T.prim_func(s_tir=True)
def func_multi_producer() -> None:
    A = T.sblock_alloc_buffer(128)
    B = T.sblock_alloc_buffer(128)
    for i in range(128):
        with T.sblock("A0"):
            vi = T.axis.S(128, i)
            A[vi] = 1.0
    for i in range(128):
        with T.sblock("A1"):
            vi = T.axis.S(128, i)
            A[vi] = 2.0
    for i in range(128):
        with T.sblock("B"):
            vi = T.axis.S(128, i)
            B[vi] = A[vi]


@T.prim_func(s_tir=True)
def func_with_block_predicate() -> None:
    A = T.sblock_alloc_buffer(120)
    B = T.sblock_alloc_buffer(120)
    for i, j in T.grid(16, 8):
        with T.sblock("producer"):
            T.where(i * 8 + j < 120)
            ax = T.axis.S(120, i * 8 + j)
            A[ax] = 0.0
    for i, j in T.grid(16, 8):
        with T.sblock("consumer"):
            T.where(i * 8 + j < 120)
            ax = T.axis.S(120, i * 8 + j)
            B[ax] = A[ax] + 1.0


@T.prim_func(s_tir=True)
def inplace_func(data_io: T.Buffer((64), "int32")):
    data_1d = T.sblock_alloc_buffer([64], dtype="int32")
    for i0 in T.serial(64):
        with T.sblock("copy_in"):
            v0 = T.axis.remap("S", [i0])
            data_1d[v0] = data_io[v0]
    for i0 in T.serial(1):
        with T.sblock("ext_call"):
            T.reads(data_1d[:64])
            T.writes(data_1d[:64])
            T.evaluate(T.call_extern("call_impl", data_1d.data, dtype=""))
    for i0 in T.serial(64):
        with T.sblock("copy_out"):
            v0 = T.axis.remap("S", [i0])
            data_io[v0] = data_1d[v0]


@T.prim_func(s_tir=True)
def inplace_call(data_io: T.Buffer((64), "int32")):
    for i0 in T.serial(1):
        with T.sblock("ext_call"):
            T.reads(data_io[:64])
            T.writes(data_io[:64])
            T.evaluate(T.call_extern("call_impl", data_io.data, dtype=""))


@T.prim_func(s_tir=True)
def cache_read_nested_seq_target(
    B: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")
) -> None:
    A = T.sblock_alloc_buffer([128, 128], dtype="float32")
    A_global = T.sblock_alloc_buffer([128, 128], dtype="float32")
    for i, j in T.grid(128, 128):
        with T.sblock("A"):
            vi, vj = T.axis.remap("SS", [i, j])
            T.reads()
            T.writes(A[vi, vj])
            A[vi, vj] = T.float32(2)
    for i, j in T.grid(8, 8):
        for x, y in T.grid(16, 16):
            with T.sblock("B0"):
                vi = T.axis.spatial(128, i * 16 + x)
                vj = T.axis.spatial(128, j * 16 + y)
                T.reads()
                T.writes(B[vi, vj])
                B[vi, vj] = T.float32(1)
        for x, y in T.grid(16, 16):
            with T.sblock("B1"):
                vi = T.axis.spatial(128, i * 16 + x)
                vj = T.axis.spatial(128, j * 16 + y)
                T.reads(A[vi, vj], B[vi, vj])
                T.writes(B[vi, vj])
                B[vi, vj] = A[vi, vj] + B[vi, vj]
    for ax0, ax1 in T.grid(128, 128):
        with T.sblock("A_global"):
            v0, v1 = T.axis.remap("SS", [ax0, ax1])
            T.reads(A[v0, v1])
            T.writes(A_global[v0, v1])
            A_global[v0, v1] = A[v0, v1]
    for i, j in T.grid(128, 128):
        with T.sblock("C"):
            vi, vj = T.axis.remap("SS", [i, j])
            T.reads(A_global[vi, vj])
            T.writes(C[vi, vj])
            C[vi, vj] = A_global[vi, vj] * T.float32(2)


@T.prim_func(s_tir=True)
def nested_buffer_access(var_A: T.handle, var_B: T.handle, var_C: T.handle):
    A = T.match_buffer(var_A, (T.int64(7), T.int64(512)), dtype="float32")
    B = T.match_buffer(var_B, T.int64(1), dtype="int32")
    C = T.match_buffer(var_C, (T.int64(1), T.int64(512)), dtype="float32")
    for ax0, ax1 in T.grid(T.int64(1), T.int64(512)):
        with T.sblock("C"):
            v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
            T.reads(A[B[v_ax0], v_ax1], B[v_ax0])
            T.writes(C[v_ax0, v_ax1])
            C[v_ax0, v_ax1] = A[B[v_ax0], v_ax1]


########## Expected function after cache_read ##########


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


@T.prim_func(s_tir=True)
def cache_read_under_scope(b: T.handle, c: T.handle) -> None:
    A = T.sblock_alloc_buffer((128, 128))
    B = T.match_buffer(b, (128, 128))
    C = T.match_buffer(c, (128, 128))
    A_global = T.sblock_alloc_buffer((128, 128))

    for i0, j0 in T.grid(8, 8):
        with T.sblock("scope"):
            i, j = T.axis.remap("SS", [i0, j0])
            A_local = T.sblock_alloc_buffer((16, 16), scope="local")
            for x, y in T.grid(16, 16):
                with T.sblock("A"):
                    vi = T.axis.S(128, i * 16 + x)
                    vj = T.axis.S(128, j * 16 + y)
                    A[vi, vj] = 1.0
            for x, y in T.grid(16, 16):
                with T.sblock("A_local"):
                    vi = T.axis.S(16, x)
                    vj = T.axis.S(16, y)
                    A_local[vi, vj] = A[i * 16 + vi, j * 16 + vj]
            for x, y in T.grid(16, 16):
                with T.sblock("B"):
                    vi = T.axis.S(128, i * 16 + x)
                    vj = T.axis.S(128, j * 16 + y)
                    B[vi, vj] = A_local[vi - i * 16, vj - j * 16] + 1.0
    for i, j in T.grid(128, 128):
        with T.sblock("A_global"):
            vi, vj = T.axis.remap("SS", [i, j])
            A_global[vi, vj] = A[vi, vj]
    for i, j in T.grid(128, 128):
        with T.sblock("C"):
            vi, vj = T.axis.remap("SS", [i, j])
            C[vi, vj] = A_global[vi, vj] * 2.0


@T.prim_func(s_tir=True)
def cache_read_opaque_access(a: T.handle, b: T.handle, c: T.handle, d: T.handle) -> None:
    A = T.match_buffer(a, (128, 128), dtype="float16")
    B = T.match_buffer(b, (128, 128), dtype="float16")
    C = T.match_buffer(c, (128, 128), dtype="float16")
    D = T.match_buffer(d, (128, 128), dtype="float16")
    A_global = T.sblock_alloc_buffer((128, 128), dtype="float16")

    for i, j in T.grid(128, 128):
        with T.sblock("A_global"):
            vi, vj = T.axis.remap("SS", [i, j])
            A_global[vi, vj] = A[vi, vj]
    for i, j in T.grid(128, 128):
        with T.sblock("load_store"):
            vi, vj = T.axis.remap("SS", [i, j])
            T.reads(A_global[vi, vj])
            T.writes(D[vi, vj])
            D[vi, vj] = A_global[vi, vj]
    for i, j in T.grid(8, 8):
        with T.sblock("opaque"):
            vi, vj = T.axis.remap("SS", [i, j])
            T.reads(A_global[vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16])
            T.writes(B[vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16])
            T.evaluate(
                T.tvm_load_matrix_sync(
                    B.data,
                    16,
                    16,
                    16,
                    vi * 8 + vj,
                    T.tvm_access_ptr(
                        T.type_annotation(dtype="float16"),
                        A_global.data,
                        vi * 2048 + vj * 16,
                        128,
                        1,
                        dtype="handle",
                    ),
                    128,
                    "row_major",
                    dtype="handle",
                )
            )
    for i, j in T.grid(8, 8):
        with T.sblock("match_buffer"):
            vi, vj = T.axis.remap("SS", [i, j])
            T.reads(A_global[vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16])
            T.writes(C[vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16])
            A0 = T.match_buffer(
                A_global[
                    vi * 16 : vi * 16 + 16,
                    vj * 16 : vj * 16 + 16,
                ],
                (16, 16),
                "float16",
                strides=[128, 1],
                offset_factor=1,
            )
            C0 = T.match_buffer(
                C[
                    vi * 16 : vi * 16 + 16,
                    vj * 16 : vj * 16 + 16,
                ],
                (16, 16),
                "float16",
                strides=[128, 1],
                offset_factor=1,
            )
            T.evaluate(
                T.tvm_load_matrix_sync(
                    C0.data,
                    16,
                    16,
                    16,
                    vi * 8 + vj,
                    T.tvm_access_ptr(
                        T.type_annotation(dtype="float16"),
                        A0.data,
                        A0.elem_offset,
                        A0.strides[0],
                        1,
                        dtype="handle",
                    ),
                    128,
                    "row_major",
                    dtype="handle",
                )
            )


@T.prim_func(s_tir=True)
def cache_read_multi_consumer() -> None:
    A = T.sblock_alloc_buffer(128)
    B = T.sblock_alloc_buffer(128)
    C = T.sblock_alloc_buffer(128)
    A_global = T.sblock_alloc_buffer(128)
    for i in T.grid(8):
        for j in T.grid(16):
            with T.sblock("A"):
                vi = T.axis.S(128, i * 16 + j)
                A[vi] = 1.0
        for j in T.grid(16):
            with T.sblock("A"):
                vi = T.axis.S(128, i * 16 + j)
                A_global[vi] = A[vi]
        for j in T.grid(16):
            with T.sblock("B"):
                vi = T.axis.S(128, i * 16 + j)
                B[vi] = A_global[vi] + 1.0

    for i in T.grid(128):
        with T.sblock("C"):
            vi = T.axis.S(128, i)
            C[vi] = A_global[vi]


@T.prim_func(s_tir=True)
def cache_read_multi_consumer_target() -> None:
    A = T.sblock_alloc_buffer(128)
    B = T.sblock_alloc_buffer(128)
    C = T.sblock_alloc_buffer(128)
    A_global = T.sblock_alloc_buffer(128)
    for i in T.grid(8):
        for j in T.grid(16):
            with T.sblock("A"):
                vi = T.axis.S(128, i * 16 + j)
                A[vi] = 1.0
        for j in T.grid(16):
            with T.sblock("B"):
                vi = T.axis.S(128, i * 16 + j)
                B[vi] = A[vi] + 1.0

    for i in T.grid(128):
        with T.sblock("A"):
            vi = T.axis.S(128, i)
            A_global[vi] = A[vi]
    for i in T.grid(128):
        with T.sblock("C"):
            vi = T.axis.S(128, i)
            C[vi] = A_global[vi]


@T.prim_func(s_tir=True)
def continuous_cache_read(a: T.handle, c: T.handle) -> None:
    A = T.match_buffer(a, (128, 128))
    C = T.match_buffer(c, (128, 128))
    B = T.sblock_alloc_buffer((128, 128))
    B_shared = T.sblock_alloc_buffer((128, 128), scope="shared")
    B_local = T.sblock_alloc_buffer((128, 128), scope="local")
    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("B_shared"):
            vi, vj = T.axis.remap("SS", [i, j])
            B_shared[vi, vj] = B[vi, vj]
    for i, j in T.grid(128, 128):
        with T.sblock("B_local"):
            vi, vj = T.axis.remap("SS", [i, j])
            B_local[vi, vj] = B_shared[vi, vj]
    for i, j in T.grid(128, 128):
        with T.sblock("C"):
            vi, vj = T.axis.remap("SS", [i, j])
            C[vi, vj] = B_local[vi, vj] + 1.0


@T.prim_func(s_tir=True)
def block_predicate_cache_read() -> None:
    A = T.sblock_alloc_buffer([120], dtype="float32")
    B = T.sblock_alloc_buffer([120], dtype="float32")
    A_shared = T.sblock_alloc_buffer([120], dtype="float32", scope="shared")
    for i, j in T.grid(16, 8):
        with T.sblock("producer"):
            ax = T.axis.spatial(120, i * 8 + j)
            T.where(i * 8 + j < 120)
            A[ax] = T.float32(0)
    for ax0 in T.serial(120):
        with T.sblock("A_shared"):
            v0 = T.axis.spatial(120, ax0)
            A_shared[v0] = A[v0]
    for i, j in T.grid(16, 8):
        with T.sblock("consumer"):
            ax = T.axis.spatial(120, i * 8 + j)
            T.where(i * 8 + j < 120)
            B[ax] = A_shared[ax] + T.float32(1)


@T.prim_func(s_tir=True)
def cache_read_shape_int64(var_A: T.handle, var_C: T.handle) -> None:
    A = T.match_buffer(var_A, (T.int64(128), T.int64(128)), dtype="float32")
    C = T.match_buffer(var_C, (T.int64(128), T.int64(128)), dtype="float32")
    B = T.sblock_alloc_buffer([T.int64(128), T.int64(128)], dtype="float32")
    A_global = T.sblock_alloc_buffer([T.int64(128), T.int64(128)], dtype="float32")
    for ax0, ax1 in T.grid(T.int64(128), T.int64(128)):
        with T.sblock("A_global"):
            v0, v1 = T.axis.remap("SS", [ax0, ax1])
            T.reads(A[v0, v1])
            T.writes(A_global[v0, v1])
            A_global[v0, v1] = A[v0, v1]
    for i, j in T.grid(128, 128):
        with T.sblock("B"):
            vi, vj = T.axis.remap("SS", [i, j])
            T.reads(A_global[vi, vj])
            T.writes(B[vi, vj])
            B[vi, vj] = A_global[vi, vj] * T.float32(2)
    for i, j in T.grid(128, 128):
        with T.sblock("C"):
            vi, vj = T.axis.remap("SS", [i, j])
            T.reads(B[vi, vj])
            T.writes(C[vi, vj])
            C[vi, vj] = B[vi, vj] + T.float32(1)


@T.prim_func(s_tir=True)
def cache_read_inplace(data_io: T.Buffer(64, "int32")) -> None:
    data_1d = T.sblock_alloc_buffer([64], dtype="int32")
    data_io_local = T.sblock_alloc_buffer([64], dtype="int32", scope="local")
    for ax0 in T.serial(64):
        with T.sblock("data_io_local"):
            v0 = T.axis.spatial(64, ax0)
            T.reads(data_io[v0])
            T.writes(data_io_local[v0])
            data_io_local[v0] = data_io[v0]
    for i0 in T.serial(64):
        with T.sblock("copy_in"):
            v0 = T.axis.spatial(64, i0)
            T.reads(data_io_local[v0])
            T.writes(data_1d[v0])
            data_1d[v0] = data_io_local[v0]
    for i0 in T.serial(1):
        with T.sblock("ext_call"):
            T.reads(data_1d[0:64])
            T.writes(data_1d[0:64])
            T.evaluate(T.call_extern("call_impl", data_1d.data, dtype=""))
    for i0 in T.serial(64):
        with T.sblock("copy_out"):
            v0 = T.axis.spatial(64, i0)
            T.reads(data_1d[v0])
            T.writes(data_io[v0])
            data_io[v0] = data_1d[v0]


@T.prim_func(s_tir=True)
def cache_inplace_buffer(data_io: T.Buffer(64, "int32")) -> None:
    data_io_local = T.sblock_alloc_buffer([64], dtype="int32", scope="local")
    data_io_global = T.sblock_alloc_buffer([64], dtype="int32")
    data_io_global_1 = T.sblock_alloc_buffer([64], dtype="int32")
    for ax0 in T.serial(64):
        with T.sblock("data_io_global"):
            v0 = T.axis.spatial(64, ax0)
            T.reads(data_io[v0])
            T.writes(data_io_global[v0])
            data_io_global[v0] = data_io[v0]
    for i0 in T.serial(1):
        for ax0 in T.serial(64):
            with T.sblock("data_io_local"):
                v0 = T.axis.spatial(64, ax0)
                T.reads(data_io_global[v0])
                T.writes(data_io_local[v0])
                data_io_local[v0] = data_io_global[v0]
        with T.sblock("ext_call"):
            T.reads(data_io_local[0:64])
            T.writes(data_io_local[0:64])
            T.evaluate(T.call_extern("call_impl", data_io_local.data, dtype=""))
        for ax0 in T.serial(64):
            with T.sblock("data_io_local"):
                v0 = T.axis.spatial(64, ax0)
                T.reads(data_io_local[v0])
                T.writes(data_io_global_1[v0])
                data_io_global_1[v0] = data_io_local[v0]
    for ax0 in T.serial(64):
        with T.sblock("data_io_global"):
            v0 = T.axis.spatial(64, ax0)
            T.reads(data_io_global_1[v0])
            T.writes(data_io[v0])
            data_io[v0] = data_io_global_1[v0]


@T.prim_func(s_tir=True)
def cache_read_nested_buffer_access(var_A: T.handle, var_B: T.handle, var_C: T.handle):
    A = T.match_buffer(var_A, (T.int64(7), T.int64(512)), dtype="float32")
    B = T.match_buffer(var_B, T.int64(1), dtype="int32")
    C = T.match_buffer(var_C, (T.int64(1), T.int64(512)), dtype="float32")
    B_global = T.sblock_alloc_buffer((T.int64(1),), "int32")
    for ax0 in range(T.int64(1)):
        with T.sblock("B_global"):
            v0 = T.axis.spatial(T.int64(1), ax0)
            T.reads(B[v0])
            T.writes(B_global[v0])
            B_global[v0] = B[v0]
    for ax0, ax1 in T.grid(T.int64(1), T.int64(512)):
        with T.sblock("C"):
            v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
            T.reads(A[B_global[v_ax0], v_ax1], B_global[v_ax0])
            T.writes(C[v_ax0, v_ax1])
            C[v_ax0, v_ax1] = A[B_global[v_ax0], v_ax1]


########## Expected function after cache_write ##########


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


@T.prim_func(s_tir=True)
def cache_write_under_scope(b: T.handle, c: T.handle) -> None:
    A = T.sblock_alloc_buffer((128, 128))
    B = T.match_buffer(b, (128, 128))
    C = T.match_buffer(c, (128, 128))
    A_global = T.sblock_alloc_buffer((128, 128))

    for i0, j0 in T.grid(8, 8):
        with T.sblock("scope"):
            i, j = T.axis.remap("SS", [i0, j0])
            A_local = T.sblock_alloc_buffer((16, 16), scope="local")
            B_global = T.sblock_alloc_buffer((16, 16))
            for x, y in T.grid(16, 16):
                with T.sblock("A_local"):
                    vi = T.axis.S(128, i * 16 + x)
                    vj = T.axis.S(128, j * 16 + y)
                    A_local[vi - i * 16, vj - j * 16] = 1.0
            for x, y in T.grid(16, 16):
                with T.sblock("A"):
                    vi = T.axis.S(16, x)
                    vj = T.axis.S(16, y)
                    A_global[i * 16 + vi, j * 16 + vj] = A_local[vi, vj]
            for x, y in T.grid(16, 16):
                with T.sblock("B"):
                    vi = T.axis.S(128, i * 16 + x)
                    vj = T.axis.S(128, j * 16 + y)
                    B_global[vi - i * 16, vj - j * 16] = A_global[vi, vj] + 1.0
            for x, y in T.grid(16, 16):
                with T.sblock("B_global"):
                    vi = T.axis.S(16, x)
                    vj = T.axis.S(16, y)
                    B[i * 16 + vi, j * 16 + vj] = B_global[vi, vj]
    for i, j in T.grid(128, 128):
        with T.sblock("A_global"):
            vi, vj = T.axis.remap("SS", [i, j])
            A[vi, vj] = A_global[vi, vj]
    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] * 2.0


@T.prim_func(s_tir=True)
def cache_write_opaque_access(a: T.handle, b: T.handle, c: T.handle, d: T.handle) -> None:
    A = T.match_buffer(a, (128, 128), dtype="float16")
    B = T.match_buffer(b, (128, 128), dtype="float16")
    C = T.match_buffer(c, (128, 128), dtype="float16")
    D = T.match_buffer(d, (128, 128), dtype="float16")
    D_global = T.sblock_alloc_buffer((128, 128), dtype="float16")
    B_global = T.sblock_alloc_buffer((128, 128), dtype="float16")
    C_global = T.sblock_alloc_buffer((128, 128), dtype="float16")

    for i, j in T.grid(128, 128):
        with T.sblock("load_store"):
            vi, vj = T.axis.remap("SS", [i, j])
            T.reads(A[vi, vj])
            T.writes(D_global[vi, vj])
            D_global[vi, vj] = A[vi, vj]
    for i, j in T.grid(8, 8):
        with T.sblock("opaque"):
            vi, vj = T.axis.remap("SS", [i, j])
            T.reads(A[vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16])
            T.writes(B_global[vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16])
            T.evaluate(
                T.tvm_load_matrix_sync(
                    B_global.data,
                    16,
                    16,
                    16,
                    vi * 8 + vj,
                    T.tvm_access_ptr(
                        T.type_annotation(dtype="float16"),
                        A.data,
                        vi * 2048 + vj * 16,
                        128,
                        1,
                        dtype="handle",
                    ),
                    128,
                    "row_major",
                    dtype="handle",
                )
            )
    for i, j in T.grid(8, 8):
        with T.sblock("match_buffer"):
            vi, vj = T.axis.remap("SS", [i, j])
            T.reads(A[vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16])
            T.writes(C_global[vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16])
            A0 = T.match_buffer(
                A[
                    vi * 16 : vi * 16 + 16,
                    vj * 16 : vj * 16 + 16,
                ],
                (16, 16),
                "float16",
                strides=[128, 1],
                offset_factor=1,
            )
            C0 = T.match_buffer(
                C_global[
                    vi * 16 : vi * 16 + 16,
                    vj * 16 : vj * 16 + 16,
                ],
                (16, 16),
                "float16",
                strides=[128, 1],
                offset_factor=1,
            )
            T.evaluate(
                T.tvm_load_matrix_sync(
                    C0.data,
                    16,
                    16,
                    16,
                    vi * 8 + vj,
                    T.tvm_access_ptr(
                        T.type_annotation(dtype="float16"),
                        A0.data,
                        A0.elem_offset,
                        A0.strides[0],
                        1,
                        dtype="handle",
                    ),
                    128,
                    "row_major",
                    dtype="handle",
                )
            )

    for i, j in T.grid(128, 128):
        with T.sblock("D"):
            vi, vj = T.axis.remap("SS", [i, j])
            D[vi, vj] = D_global[vi, vj]
    for i, j in T.grid(128, 128):
        with T.sblock("B"):
            vi, vj = T.axis.remap("SS", [i, j])
            B[vi, vj] = B_global[vi, vj]
    for i, j in T.grid(128, 128):
        with T.sblock("C"):
            vi, vj = T.axis.remap("SS", [i, j])
            C[vi, vj] = C_global[vi, vj]


@T.prim_func(s_tir=True)
def cache_write_multi_consumer() -> None:
    A = T.sblock_alloc_buffer(128)
    B = T.sblock_alloc_buffer(128)
    C = T.sblock_alloc_buffer(128)
    A_global = T.sblock_alloc_buffer(128)
    for i in T.grid(8):
        for j in T.grid(16):
            with T.sblock("A_global"):
                vi = T.axis.S(128, i * 16 + j)
                A_global[vi] = 1.0
        for j in T.grid(16):
            with T.sblock("A"):
                vi = T.axis.S(128, i * 16 + j)
                A[vi] = A_global[vi]
        for j in T.grid(16):
            with T.sblock("B"):
                vi = T.axis.S(128, i * 16 + j)
                B[vi] = A[vi] + 1.0

    for i in T.grid(128):
        with T.sblock("C"):
            vi = T.axis.S(128, i)
            C[vi] = A[vi]


@T.prim_func(s_tir=True)
def cache_write_multi_consumer_B_consume_cache():
    A = T.sblock_alloc_buffer([128], dtype="float32")
    B = T.sblock_alloc_buffer([128], dtype="float32")
    C = T.sblock_alloc_buffer([128], dtype="float32")
    A_global = T.sblock_alloc_buffer([128], dtype="float32")
    for i in T.serial(8):
        for j in T.serial(16):
            with T.sblock("A"):
                vi = T.axis.spatial(128, i * 16 + j)
                A_global[vi] = 1.0
        for j in T.serial(16):
            with T.sblock("B"):
                vi = T.axis.spatial(128, i * 16 + j)
                B[vi] = A_global[vi] + 1.0
    for ax0 in T.serial(128):
        with T.sblock("A_global"):
            v0 = T.axis.spatial(128, ax0)
            A[v0] = A_global[v0]
    for i in T.serial(128):
        with T.sblock("C"):
            vi = T.axis.spatial(128, i)
            C[vi] = A[vi]


@T.prim_func(s_tir=True)
def cache_write_multi_consumer_C_consume_cache():
    A = T.sblock_alloc_buffer([128], dtype="float32")
    B = T.sblock_alloc_buffer([128], dtype="float32")
    C = T.sblock_alloc_buffer([128], dtype="float32")
    A_global = T.sblock_alloc_buffer([128], dtype="float32")
    for i in T.serial(8):
        for j in T.serial(16):
            with T.sblock("A"):
                vi = T.axis.spatial(128, i * 16 + j)
                A_global[vi] = T.float32(1)
        for ax0 in T.serial(16):
            with T.sblock("A_global"):
                v0 = T.axis.spatial(128, i * 16 + ax0)
                A[v0] = A_global[v0]
        for j in T.serial(16):
            with T.sblock("B"):
                vi = T.axis.spatial(128, i * 16 + j)
                B[vi] = A[vi] + T.float32(1)
    for i in T.serial(128):
        with T.sblock("C"):
            vi = T.axis.spatial(128, i)
            C[vi] = A_global[vi]


@T.prim_func(s_tir=True)
def cache_write_multi_consumer_all_consume_cache():
    A = T.sblock_alloc_buffer([128], dtype="float32")
    B = T.sblock_alloc_buffer([128], dtype="float32")
    C = T.sblock_alloc_buffer([128], dtype="float32")
    A_global = T.sblock_alloc_buffer([128], dtype="float32")
    for i in T.serial(8):
        for j in T.serial(16):
            with T.sblock("A"):
                vi = T.axis.spatial(128, i * 16 + j)
                A_global[vi] = T.float32(1)
        for j in T.serial(16):
            with T.sblock("B"):
                vi = T.axis.spatial(128, i * 16 + j)
                B[vi] = A_global[vi] + T.float32(1)
    for i in T.serial(128):
        with T.sblock("C"):
            vi = T.axis.spatial(128, i)
            C[vi] = A_global[vi]
    for ax0 in T.serial(128):
        with T.sblock("A_global"):
            v0 = T.axis.spatial(128, ax0)
            A[v0] = A_global[v0]


@T.prim_func(s_tir=True)
def continuous_cache_write(a: T.handle, c: T.handle) -> None:
    A = T.match_buffer(a, (128, 128))
    B = T.sblock_alloc_buffer((128, 128))
    C = T.match_buffer(c, (128, 128))
    B_shared = T.sblock_alloc_buffer((128, 128), scope="shared")
    B_local = T.sblock_alloc_buffer((128, 128), scope="local")
    for i, j in T.grid(128, 128):
        with T.sblock("B"):
            vi, vj = T.axis.remap("SS", [i, j])
            B_local[vi, vj] = A[vi, vj] * 2.0
    for i, j in T.grid(128, 128):
        with T.sblock("B"):
            vi, vj = T.axis.remap("SS", [i, j])
            B_shared[vi, vj] = B_local[vi, vj]
    for i, j in T.grid(128, 128):
        with T.sblock("B"):
            vi, vj = T.axis.remap("SS", [i, j])
            B[vi, vj] = B_shared[vi, vj]
    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 block_predicate_cache_write_intermediate_buf() -> None:
    A = T.sblock_alloc_buffer([120], dtype="float32")
    B = T.sblock_alloc_buffer([120], dtype="float32")
    A_shared = T.sblock_alloc_buffer([120], dtype="float32", scope="shared")
    for i, j in T.grid(16, 8):
        with T.sblock("producer"):
            ax = T.axis.spatial(120, i * 8 + j)
            T.where(i * 8 + j < 120)
            A_shared[ax] = T.float32(0)
    for ax0 in T.serial(120):
        with T.sblock("A_shared"):
            v0 = T.axis.spatial(120, ax0)
            A[v0] = A_shared[v0]
    for i, j in T.grid(16, 8):
        with T.sblock("consumer"):
            ax = T.axis.spatial(120, i * 8 + j)
            T.where(i * 8 + j < 120)
            B[ax] = A[ax] + 1.0


@T.prim_func(s_tir=True)
def block_predicate_cache_write_output_buf() -> None:
    A = T.sblock_alloc_buffer([120], dtype="float32")
    B = T.sblock_alloc_buffer([120], dtype="float32")
    B_shared = T.sblock_alloc_buffer([120], dtype="float32", scope="shared")
    for i, j in T.grid(16, 8):
        with T.sblock("producer"):
            ax = T.axis.spatial(120, i * 8 + j)
            T.where(i * 8 + j < 120)
            A[ax] = T.float32(0)
    for i, j in T.grid(16, 8):
        with T.sblock("consumer"):
            ax = T.axis.spatial(120, i * 8 + j)
            T.where(i * 8 + j < 120)
            B_shared[ax] = A[ax] + T.float32(1)
    for ax0 in T.serial(120):
        with T.sblock("B_shared"):
            v0 = T.axis.spatial(120, ax0)
            B[v0] = B_shared[v0]


@T.prim_func(s_tir=True)
def symbolic_matmul_blocked(var_A: T.handle, var_B: T.handle, var_C: T.handle, n: T.int32):
    A = T.match_buffer(var_A, ((n + 31) // 32 * 32, 4))
    B = T.match_buffer(var_B, (4, (n + 31) // 32 * 32))
    C = T.match_buffer(var_C, ((n + 31) // 32 * 32, (n + 31) // 32 * 32))
    for i0_0, i1_0 in T.grid((n + 31) // 32, (n + 31) // 32):
        with T.sblock("matmul_o"):
            v_i0_o, v_i1_o = T.axis.remap("SS", [i0_0, i1_0])
            T.reads(
                A[v_i0_o * 32 : v_i0_o * 32 + 32, 0:4],
                B[0:4, v_i1_o * 32 : v_i1_o * 32 + 32],
            )
            T.writes(C[v_i0_o * 32 : v_i0_o * 32 + 32, v_i1_o * 32 : v_i1_o * 32 + 32])
            for i0_1, i1_1, k in T.grid(32, 32, 4):
                with T.sblock("matmul"):
                    v_i0_i, v_i1_i, v_k_i = T.axis.remap("SSR", [i0_1, i1_1, k])
                    T.reads(A[v_i0_o * 32 + v_i0_i, v_k_i], B[v_k_i, v_i1_o * 32 + v_i1_i])
                    T.writes(C[v_i0_o * 32 + v_i0_i, v_i1_o * 32 + v_i1_i])
                    with T.init():
                        C[v_i0_o * 32 + v_i0_i, v_i1_o * 32 + v_i1_i] = T.float32(0)
                    C[v_i0_o * 32 + v_i0_i, v_i1_o * 32 + v_i1_i] = (
                        C[v_i0_o * 32 + v_i0_i, v_i1_o * 32 + v_i1_i]
                        + A[v_i0_o * 32 + v_i0_i, v_k_i] * B[v_k_i, v_i1_o * 32 + v_i1_i]
                    )


@T.prim_func(s_tir=True)
def symbolic_matmul_blocked_cache_read(
    var_A: T.handle, var_B: T.handle, var_C: T.handle, n: T.int32
):
    A = T.match_buffer(var_A, ((n + 31) // 32 * 32, 4))
    B = T.match_buffer(var_B, (4, (n + 31) // 32 * 32))
    C = T.match_buffer(var_C, ((n + 31) // 32 * 32, (n + 31) // 32 * 32))
    for i0_0, i1_0 in T.grid((n + 31) // 32, (n + 31) // 32):
        with T.sblock("matmul_o"):
            v_i0_o, v_i1_o = T.axis.remap("SS", [i0_0, i1_0])
            T.reads(
                A[v_i0_o * 32 : v_i0_o * 32 + 32, 0:4],
                B[0:4, v_i1_o * 32 : v_i1_o * 32 + 32],
            )
            T.writes(C[v_i0_o * 32 : v_i0_o * 32 + 32, v_i1_o * 32 : v_i1_o * 32 + 32])
            A_shared = T.sblock_alloc_buffer((32, 4), scope="shared")
            for ax0, ax1 in T.grid(32, 4):
                with T.sblock("A_shared"):
                    v0 = T.axis.spatial(32, ax0)
                    v1 = T.axis.spatial(4, ax1)
                    T.reads(A[v_i0_o * 32 + v0, v1])
                    T.writes(A_shared[v0, v1])
                    A_shared[v0, v1] = A[v_i0_o * 32 + v0, v1]
            for i0_1, i1_1, k in T.grid(32, 32, 4):
                with T.sblock("matmul"):
                    v_i0_i, v_i1_i, v_k_i = T.axis.remap("SSR", [i0_1, i1_1, k])
                    T.reads(A_shared[v_i0_i, v_k_i], B[v_k_i, v_i1_o * 32 + v_i1_i])
                    T.writes(C[v_i0_o * 32 + v_i0_i, v_i1_o * 32 + v_i1_i])
                    with T.init():
                        C[v_i0_o * 32 + v_i0_i, v_i1_o * 32 + v_i1_i] = T.float32(0)
                    C[v_i0_o * 32 + v_i0_i, v_i1_o * 32 + v_i1_i] = (
                        C[v_i0_o * 32 + v_i0_i, v_i1_o * 32 + v_i1_i]
                        + A_shared[v_i0_i, v_k_i] * B[v_k_i, v_i1_o * 32 + v_i1_i]
                    )


@T.prim_func(s_tir=True)
def symbolic_matmul_blocked_cache_write(
    var_A: T.handle, var_B: T.handle, var_C: T.handle, n: T.int32
):
    A = T.match_buffer(var_A, ((n + 31) // 32 * 32, 4))
    B = T.match_buffer(var_B, (4, (n + 31) // 32 * 32))
    C = T.match_buffer(var_C, ((n + 31) // 32 * 32, (n + 31) // 32 * 32))
    for i0_0, i1_0 in T.grid((n + 31) // 32, (n + 31) // 32):
        with T.sblock("matmul_o"):
            v_i0_o, v_i1_o = T.axis.remap("SS", [i0_0, i1_0])
            T.reads(
                A[v_i0_o * 32 : v_i0_o * 32 + 32, 0:4],
                B[0:4, v_i1_o * 32 : v_i1_o * 32 + 32],
            )
            T.writes(C[v_i0_o * 32 : v_i0_o * 32 + 32, v_i1_o * 32 : v_i1_o * 32 + 32])
            C_pad_local = T.sblock_alloc_buffer((32, 32), scope="local")
            for i0_1, i1_1, k in T.grid(32, 32, 4):
                with T.sblock("matmul"):
                    v_i0_i, v_i1_i, v_k_i = T.axis.remap("SSR", [i0_1, i1_1, k])
                    T.reads(A[v_i0_o * 32 + v_i0_i, v_k_i], B[v_k_i, v_i1_o * 32 + v_i1_i])
                    T.writes(C_pad_local[v_i0_i, v_i1_i])
                    with T.init():
                        C_pad_local[v_i0_i, v_i1_i] = T.float32(0)
                    C_pad_local[v_i0_i, v_i1_i] = (
                        C_pad_local[v_i0_i, v_i1_i]
                        + A[v_i0_o * 32 + v_i0_i, v_k_i] * B[v_k_i, v_i1_o * 32 + v_i1_i]
                    )
            for ax0, ax1 in T.grid(32, 32):
                with T.sblock("C_pad_local"):
                    v0 = T.axis.spatial(32, ax0)
                    v1 = T.axis.spatial(32, ax1)
                    T.reads(C_pad_local[v0, v1])
                    T.writes(C[v_i0_o * 32 + v0, v_i1_o * 32 + v1])
                    C[v_i0_o * 32 + v0, v_i1_o * 32 + v1] = C_pad_local[v0, v1]


########## Testcases for cache_read ##########

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


def test_cache_read_elementwise(use_block_name):
    sch = tvm.s_tir.Schedule(elementwise, debug_mask="all")
    block_b = sch.get_sblock("B")
    block_c = sch.get_sblock("C")
    if use_block_name:
        cached_a = sch.cache_read("B", "A", "global")
        cached_b = sch.cache_read("C", "B", "local")
    else:
        cached_a = sch.cache_read(block_b, 0, "global")
        cached_b = sch.cache_read(block_c, 0, "local")
    assert sch.get(cached_a) == sch.get(sch.get_sblock("A_global"))
    assert sch.get(cached_b) == sch.get(sch.get_sblock("B_local"))
    assert sch.get(block_b) == sch.get(sch.get_sblock("B"))
    assert sch.get(block_c) == sch.get(sch.get_sblock("C"))
    assert_structural_equal_ignore_global_symbol(cache_read_elementwise, sch.mod["main"])
    verify_trace_roundtrip(sch=sch, mod=elementwise)


def test_cache_read_under_scope(use_block_name):
    sch = tvm.s_tir.Schedule(access_under_scope, debug_mask="all")
    block_b = "B" if use_block_name else sch.get_sblock("B")
    block_c = "C" if use_block_name else sch.get_sblock("C")
    sch.cache_read(block_b, 0, "local")
    sch.cache_read(block_c, 0, "global")
    assert_structural_equal_ignore_global_symbol(cache_read_under_scope, sch.mod["main"])
    verify_trace_roundtrip(sch=sch, mod=access_under_scope)


def test_cache_read_opaque_access(use_block_name):
    sch = tvm.s_tir.Schedule(opaque_access, debug_mask="all")
    block = "load_store" if use_block_name else sch.get_sblock("load_store")
    sch.cache_read(block, 0, "global")
    assert_structural_equal_ignore_global_symbol(cache_read_opaque_access, sch.mod["main"])
    verify_trace_roundtrip(sch=sch, mod=opaque_access)


def test_cache_read_location(use_block_name):
    sch = tvm.s_tir.Schedule(func_multi_consumer, debug_mask="all")
    block_b = "B" if use_block_name else sch.get_sblock("B")
    sch.cache_read(block_b, 0, "global")
    assert_structural_equal_ignore_global_symbol(cache_read_multi_consumer, sch.mod["main"])
    verify_trace_roundtrip(sch=sch, mod=func_multi_consumer)

    # Test that specific consumer block targeting works.
    sch = tvm.s_tir.Schedule(func_multi_consumer, debug_mask="all")
    block_b = "B" if use_block_name else sch.get_sblock("B")
    block_c = "C" if use_block_name else sch.get_sblock("C")
    sch.cache_read(block_b, 0, "global", consumer_blocks=[block_c])
    assert_structural_equal_ignore_global_symbol(cache_read_multi_consumer_target, sch.mod["main"])
    verify_trace_roundtrip(sch=sch, mod=func_multi_consumer)

    # Also test setting multiple consumers yields same result as unspecified.
    sch = tvm.s_tir.Schedule(func_multi_consumer, debug_mask="all")
    block_b = "B" if use_block_name else sch.get_sblock("B")
    block_c = "C" if use_block_name else sch.get_sblock("C")
    sch.cache_read(block_b, 0, "global", consumer_blocks=[block_b, block_c])
    assert_structural_equal_ignore_global_symbol(cache_read_multi_consumer, sch.mod["main"])
    verify_trace_roundtrip(sch=sch, mod=func_multi_consumer)


def test_continuous_cache_read(use_block_name):
    sch = tvm.s_tir.Schedule(elementwise, debug_mask="all")
    block_c = "C" if use_block_name else sch.get_sblock("C")
    sch.cache_read(block_c, 0, "shared")
    sch.cache_read(block_c, 0, "local")
    assert_structural_equal_ignore_global_symbol(continuous_cache_read, sch.mod["main"])
    verify_trace_roundtrip(sch=sch, mod=elementwise)


def test_cache_read_with_block_predicate(use_block_name):
    sch = tvm.s_tir.Schedule(func_with_block_predicate, debug_mask="all")
    block = "consumer" if use_block_name else sch.get_sblock("consumer")
    sch.cache_read(block, 0, "shared")
    assert_structural_equal_ignore_global_symbol(block_predicate_cache_read, sch.mod["main"])
    verify_trace_roundtrip(sch=sch, mod=func_with_block_predicate)


def test_cache_read_non_int32_shape(use_block_name):
    sch = tvm.s_tir.Schedule(elementwise_shape_int64, debug_mask="all")
    block_b = "B" if use_block_name else sch.get_sblock("B")
    sch.cache_read(block_b, 0, "global")
    assert_structural_equal_ignore_global_symbol(cache_read_shape_int64, sch.mod["main"])
    verify_trace_roundtrip(sch=sch, mod=elementwise_shape_int64)


def test_cache_read_nested_buffer_access(use_block_name):
    sch = tvm.s_tir.Schedule(nested_buffer_access, debug_mask="all")
    block_c = "C" if use_block_name else sch.get_sblock("C")
    sch.cache_read(block_c, 1, "global")
    assert_structural_equal_ignore_global_symbol(cache_read_nested_buffer_access, sch.mod["main"])
    verify_trace_roundtrip(sch=sch, mod=nested_buffer_access)


def test_cache_read_fail_multi_producer(use_block_name):
    sch = tvm.s_tir.Schedule(func_multi_producer, debug_mask="all")
    block_b = "B" if use_block_name else sch.get_sblock("B")
    with pytest.raises(tvm.s_tir.ScheduleError):
        sch.cache_read(block_b, 0, "global")


def test_cache_read_fail_index_out_of_bound(use_block_name):
    sch = tvm.s_tir.Schedule(elementwise, debug_mask="all")
    block_b = "B" if use_block_name else sch.get_sblock("B")
    with pytest.raises(tvm.s_tir.ScheduleError):
        sch.cache_read(block_b, 1, "global")


def test_cache_read_fail_invalid_storage_scope(use_block_name):
    sch = tvm.s_tir.Schedule(elementwise, debug_mask="all")
    block_b = "B" if use_block_name else sch.get_sblock("B")
    with pytest.raises(tvm.s_tir.ScheduleError):
        sch.cache_read(block_b, 0, "test_scope")


def test_inplace_cache_read():
    sch = tvm.s_tir.Schedule(inplace_func, debug_mask="all")
    block = sch.get_sblock("copy_in")
    sch.cache_read(block, 0, "local", [block])
    assert_structural_equal_ignore_global_symbol(cache_read_inplace, sch.mod["main"])
    verify_trace_roundtrip(sch=sch, mod=inplace_func)


def test_cache_inplace():
    # cache_inplace could introduce WAR, which is expected but stage pipeline property changes
    debug_mask = tvm.s_tir.schedule.state.ScheduleDebugMask.VERIFY_SREF_TREE
    sch = tvm.s_tir.Schedule(inplace_call, debug_mask=debug_mask)
    block = sch.get_sblock("ext_call")
    blocks = sch.cache_inplace(block, 0, "local")
    block = sch.cache_read(blocks[0], 0, "global", [blocks[0]])
    block = sch.cache_write(blocks[1], 0, "global")

    assert_structural_equal_ignore_global_symbol(cache_inplace_buffer, sch.mod["main"])
    verify_trace_roundtrip(sch=sch, mod=inplace_call, debug_mask=debug_mask)


def test_cache_read_nested_seq(use_block_name):
    sch = tvm.s_tir.Schedule(func_nested_seq, debug_mask="all")
    block_c = "C" if use_block_name else sch.get_sblock("C")
    sch.cache_read(block_c, 0, "global", consumer_blocks=[block_c])
    assert_structural_equal_ignore_global_symbol(cache_read_nested_seq_target, sch.mod["main"])
    verify_trace_roundtrip(sch=sch, mod=func_nested_seq)


########## Testcases for cache_write ##########


def test_cache_write_elementwise(use_block_name):
    sch = tvm.s_tir.Schedule(elementwise, debug_mask="all")
    block_b = sch.get_sblock("B")
    block_c = sch.get_sblock("C")
    cached_b = sch.cache_write("B" if use_block_name else block_b, 0, "local")
    cached_c = sch.cache_write("C" if use_block_name else block_c, 0, "global")
    assert sch.get(cached_b) == sch.get(sch.get_sblock("B_local"))
    assert sch.get(cached_c) == sch.get(sch.get_sblock("C_global"))
    assert sch.get(block_b) == sch.get(sch.get_sblock("B"))
    assert sch.get(block_c) == sch.get(sch.get_sblock("C"))
    assert_structural_equal_ignore_global_symbol(cache_write_elementwise, sch.mod["main"])
    verify_trace_roundtrip(sch=sch, mod=elementwise)


def test_cache_write_under_scope(use_block_name):
    sch = tvm.s_tir.Schedule(access_under_scope, debug_mask="all")
    block_a = "A" if use_block_name else sch.get_sblock("A")
    block_b = "B" if use_block_name else sch.get_sblock("B")
    block_scope = sch.get_sblock("scope")
    sch.cache_write(block_a, 0, "local")
    sch.cache_write(block_b, 0, "global")
    sch.cache_write(block_scope, 0, "global")
    assert_structural_equal_ignore_global_symbol(cache_write_under_scope, sch.mod["main"])
    verify_trace_roundtrip(sch=sch, mod=access_under_scope)


def test_cache_write_opaque_access(use_block_name):
    sch = tvm.s_tir.Schedule(opaque_access, debug_mask="all")
    block_store = "load_store" if use_block_name else sch.get_sblock("load_store")
    block_opaque = "opaque" if use_block_name else sch.get_sblock("opaque")
    block_match_buffer = "match_buffer" if use_block_name else sch.get_sblock("match_buffer")
    sch.cache_write(block_store, 0, "global")
    sch.cache_write(block_opaque, 0, "global")
    sch.cache_write(block_match_buffer, 0, "global")
    assert_structural_equal_ignore_global_symbol(cache_write_opaque_access, sch.mod["main"])
    verify_trace_roundtrip(sch=sch, mod=opaque_access)


def test_cache_write_location(use_block_name):
    sch = tvm.s_tir.Schedule(func_multi_consumer, debug_mask="all")
    block_a = "A" if use_block_name else sch.get_sblock("A")
    sch.cache_write(block_a, 0, "global")
    assert_structural_equal_ignore_global_symbol(cache_write_multi_consumer, sch.mod["main"])
    verify_trace_roundtrip(sch=sch, mod=func_multi_consumer)

    # Test that specific consumer block targeting works.
    # B read cache buffer and C read original output buffer
    sch = tvm.s_tir.Schedule(func_multi_consumer, debug_mask="all")
    block_a = "A" if use_block_name else sch.get_sblock("A")
    block_b = "B" if use_block_name else sch.get_sblock("B")
    sch.cache_write(block_a, 0, "global", consumer_blocks=[block_b])
    assert_structural_equal_ignore_global_symbol(
        cache_write_multi_consumer_B_consume_cache, sch.mod["main"]
    )
    verify_trace_roundtrip(sch=sch, mod=func_multi_consumer)

    # Test that specific consumer block targeting works.
    # B read original output buffer and C read cache buffer
    sch = tvm.s_tir.Schedule(func_multi_consumer, debug_mask="all")
    block_a = "A" if use_block_name else sch.get_sblock("A")
    block_c = "C" if use_block_name else sch.get_sblock("C")
    sch.cache_write(block_a, 0, "global", consumer_blocks=[block_c])
    assert_structural_equal_ignore_global_symbol(
        cache_write_multi_consumer_C_consume_cache, sch.mod["main"]
    )
    verify_trace_roundtrip(sch=sch, mod=func_multi_consumer)

    # Test that specific consumer block targeting works.
    # B and C read cache buffer
    sch = tvm.s_tir.Schedule(func_multi_consumer, debug_mask="all")
    block_a = "A" if use_block_name else sch.get_sblock("A")
    block_b = "B" if use_block_name else sch.get_sblock("B")
    block_c = "C" if use_block_name else sch.get_sblock("C")
    sch.cache_write(block_a, 0, "global", consumer_blocks=[block_b, block_c])
    assert_structural_equal_ignore_global_symbol(
        cache_write_multi_consumer_all_consume_cache, sch.mod["main"]
    )
    verify_trace_roundtrip(sch=sch, mod=func_multi_consumer)


def test_continuous_cache_write(use_block_name):
    sch = tvm.s_tir.Schedule(elementwise, debug_mask="all")
    block_b = "B" if use_block_name else sch.get_sblock("B")
    sch.cache_write(block_b, 0, "shared")
    sch.cache_write(block_b, 0, "local")
    assert_structural_equal_ignore_global_symbol(continuous_cache_write, sch.mod["main"])
    verify_trace_roundtrip(sch=sch, mod=elementwise)


def test_cache_write_with_block_predicate(use_block_name):
    # cache write for intermediate buffer
    sch = tvm.s_tir.Schedule(func_with_block_predicate, debug_mask="all")
    block = "producer" if use_block_name else sch.get_sblock("producer")
    sch.cache_write(block, 0, "shared")
    assert_structural_equal_ignore_global_symbol(
        block_predicate_cache_write_intermediate_buf, sch.mod["main"]
    )
    verify_trace_roundtrip(sch=sch, mod=func_with_block_predicate)
    # cache write for external buffer
    sch = tvm.s_tir.Schedule(func_with_block_predicate, debug_mask="all")
    block = "consumer" if use_block_name else sch.get_sblock("consumer")
    sch.cache_write(block, 0, "shared")
    assert_structural_equal_ignore_global_symbol(
        block_predicate_cache_write_output_buf, sch.mod["main"]
    )
    verify_trace_roundtrip(sch=sch, mod=func_with_block_predicate)


def test_cache_write_fail_multi_producer(use_block_name):
    sch = tvm.s_tir.Schedule(func_multi_producer, debug_mask="all")
    block_a0 = "A0" if use_block_name else sch.get_sblock("A0")
    block_a1 = "A1" if use_block_name else sch.get_sblock("A1")
    with pytest.raises(tvm.s_tir.ScheduleError):
        sch.cache_write(block_a0, 0, "global")
    with pytest.raises(tvm.s_tir.ScheduleError):
        sch.cache_write(block_a1, 0, "global")


def test_cache_write_fail_index_out_of_bound(use_block_name):
    sch = tvm.s_tir.Schedule(elementwise, debug_mask="all")
    block_b = "B" if use_block_name else sch.get_sblock("B")
    with pytest.raises(tvm.s_tir.ScheduleError):
        sch.cache_write(block_b, 1, "global")


def test_cache_write_fail_invalid_storage_scope(use_block_name):
    sch = tvm.s_tir.Schedule(elementwise, debug_mask="all")
    block_b = "B" if use_block_name else sch.get_sblock("B")
    with pytest.raises(tvm.s_tir.ScheduleError):
        sch.cache_write(block_b, 0, "test_scope")


def test_reindex_cache_read():
    sch = tvm.s_tir.Schedule(elementwise, debug_mask="all")
    sch.reindex_cache_read("C", 0, "shared", lambda i, j: (j, i // 2, i % 2))
    assert_structural_equal_ignore_global_symbol(elementwise_reindex_cache_read, sch.mod["main"])
    verify_trace_roundtrip(sch=sch, mod=elementwise)


def test_reindex_cache_read_multi_consumer():
    sch = tvm.s_tir.Schedule(func_multi_consumer)
    sch.reindex_cache_read("B", 0, "shared", lambda i: (i // 32, i % 32))
    assert_structural_equal_ignore_global_symbol(reindex_cache_read_multi_consumer, sch.mod["main"])
    # NOTE(zihao): we do not verify trace roundtrip because of in set analysis issues.


def test_reindex_cache_read_fail_not_match():
    sch = tvm.s_tir.Schedule(elementwise, debug_mask="all")
    with pytest.raises(tvm.s_tir.ScheduleError):
        sch.reindex_cache_read(
            "C",
            0,
            "shared",
            lambda i, j: j * 2,
        )


def test_reindex_cache_read_failed_not_single_point():
    sch = tvm.s_tir.Schedule(access_under_scope, debug_mask="all")
    with pytest.raises(tvm.s_tir.ScheduleError):
        sch.reindex_cache_read("scope", 0, "shared", lambda i, j: (i, j))


def test_reindex_cache_write():
    sch = tvm.s_tir.Schedule(elementwise, debug_mask="all")
    sch.reindex_cache_write("B", 0, "shared", lambda i, j: (j, i))
    assert_structural_equal_ignore_global_symbol(elementwise_reindex_cache_write, sch.mod["main"])
    verify_trace_roundtrip(sch=sch, mod=elementwise)


def test_reindex_cache_write_reduce():
    sch = tvm.s_tir.Schedule(reduce, debug_mask="all")
    sch.reindex_cache_write("B", 0, "shared", lambda i, j, k, l: (j, i, k))
    assert_structural_equal_ignore_global_symbol(reduce_reindex_cache_write_0, sch.mod["main"])
    sch.reindex_cache_write("C", 0, "shared", lambda i, j, k: [j, i])
    assert_structural_equal_ignore_global_symbol(reduce_reindex_cache_write_1, sch.mod["main"])
    verify_trace_roundtrip(sch=sch, mod=reduce)


def test_reindex_cache_write_fail_not_match():
    sch = tvm.s_tir.Schedule(elementwise, debug_mask="all")
    with pytest.raises(tvm.s_tir.ScheduleError):
        sch.reindex_cache_write(
            "B",
            0,
            "shared",
            lambda i, j: i,
        )


def test_reindex_cache_write_fail_not_single_point():
    sch = tvm.s_tir.Schedule(access_under_scope, debug_mask="all")
    with pytest.raises(tvm.s_tir.ScheduleError):
        sch.reindex_cache_write("scope", 0, "shared", lambda i, j: (i, j))


def test_symbolic_matmul_blocked_cache_read(use_block_name):
    sch = tvm.s_tir.Schedule(symbolic_matmul_blocked, debug_mask="all")
    block = "matmul" if use_block_name else sch.get_sblock("matmul")
    sch.cache_read(block=block, read_buffer_index=0, storage_scope="shared")
    assert_structural_equal_ignore_global_symbol(
        sch.mod["main"], symbolic_matmul_blocked_cache_read
    )
    verify_trace_roundtrip(sch=sch, mod=symbolic_matmul_blocked)


def test_symbolic_matmul_blocked_cache_write(use_block_name):
    sch = tvm.s_tir.Schedule(symbolic_matmul_blocked, debug_mask="all")
    block = "matmul" if use_block_name else sch.get_sblock("matmul")
    sch.cache_write(block=block, write_buffer_index=0, storage_scope="local")
    assert_structural_equal_ignore_global_symbol(
        sch.mod["main"], symbolic_matmul_blocked_cache_write
    )
    verify_trace_roundtrip(sch=sch, mod=symbolic_matmul_blocked)


def test_cache_write_with_nested_block_predicate():
    @T.prim_func(s_tir=True)
    def main(A: T.handle, C: T.handle) -> None:
        A_buf = T.match_buffer(A, (12, 24), "float32")
        C_buf = T.match_buffer(C, (10, 20), "float32")

        for i, j in T.grid(12, 24):
            with T.sblock("compute"):
                vi, vj = T.axis.remap("SS", [i, j])

                with T.sblock("inner"):
                    T.where(vi < 10 and vj < 20)
                    C_buf[vi, vj] = A_buf[vi, vj] * 2.0

    @T.prim_func(s_tir=True)
    def expected(A_buf: T.Buffer((12, 24), "float32"), C_buf: T.Buffer((10, 20), "float32")):
        with T.sblock("root"):
            C_buf_local = T.sblock_alloc_buffer((10, 20), scope="local")
            for i, j in T.grid(12, 24):
                with T.sblock("compute"):
                    vi, vj = T.axis.remap("SS", [i, j])
                    T.reads(A_buf[vi, vj])
                    T.writes(C_buf_local[vi, vj])
                    with T.sblock("inner"):
                        T.where(vi < 10 and vj < 20)
                        T.reads(A_buf[vi, vj])
                        T.writes(C_buf_local[vi, vj])
                        C_buf_local[vi, vj] = A_buf[vi, vj] * T.float32(2)
            for ax0, ax1 in T.grid(10, 20):
                with T.sblock("C_buf_local"):
                    v0, v1 = T.axis.remap("SS", [ax0, ax1])
                    T.reads(C_buf_local[v0, v1])
                    T.writes(C_buf[v0, v1])
                    C_buf[v0, v1] = C_buf_local[v0, v1]

    sch = tvm.s_tir.Schedule(main)
    block = sch.get_sblock("compute")
    sch.cache_write(block, 0, "local")
    assert_structural_equal_ignore_global_symbol(expected, sch.mod["main"])


def test_cache_read_with_nested_block_predicate():
    @T.prim_func(s_tir=True)
    def main(A: T.handle, C: T.handle) -> None:
        A_buf = T.match_buffer(A, (12, 24), "float32")
        C_buf = T.match_buffer(C, (10, 20), "float32")

        for i, j in T.grid(12, 24):
            with T.sblock("compute"):
                vi, vj = T.axis.remap("SS", [i, j])

                with T.sblock("inner"):
                    T.where(vi < 10 and vj < 20)
                    C_buf[vi, vj] = A_buf[vi, vj] * 2.0

    @T.prim_func(s_tir=True)
    def expected(A_buf: T.Buffer((12, 24), "float32"), C_buf: T.Buffer((10, 20), "float32")):
        with T.sblock("root"):
            A_buf_local = T.sblock_alloc_buffer((10, 20), scope="local")
            for ax0, ax1 in T.grid(10, 20):
                with T.sblock("A_buf_local"):
                    v0, v1 = T.axis.remap("SS", [ax0, ax1])
                    T.reads(A_buf[v0, v1])
                    T.writes(A_buf_local[v0, v1])
                    A_buf_local[v0, v1] = A_buf[v0, v1]
            for i, j in T.grid(12, 24):
                with T.sblock("compute"):
                    vi, vj = T.axis.remap("SS", [i, j])
                    T.reads(A_buf_local[vi, vj])
                    T.writes(C_buf[vi, vj])
                    with T.sblock("inner"):
                        T.where(vi < 10 and vj < 20)
                        T.reads(A_buf_local[vi, vj])
                        T.writes(C_buf[vi, vj])
                        C_buf[vi, vj] = A_buf_local[vi, vj] * T.float32(2)

    sch = tvm.s_tir.Schedule(main)
    block = sch.get_sblock("compute")
    sch.cache_read(block, 0, "local")
    assert_structural_equal_ignore_global_symbol(expected, sch.mod["main"])


def test_cache_write_sibling_nested_block_predicates_use_union():
    """Regression: cache_write with sibling nested blocks must union their predicates.

    Two sibling nested sblocks access the same buffer under *different* predicates:
      left  block: T.where(vi < 8)   — writes rows 0-7, all columns
      top   block: T.where(vj < 16)  — writes all rows, columns 0-15

    The cache must cover the UNION of both access sets.  The bounding box of that
    union is (12, 24) — the full buffer shape.

    Bug: CollectNestedBlockPredicates ANDs the predicates of all found nested blocks,
    giving (vi < 8) AND (vj < 16).  RelaxBufferRegion under that intersection predicate
    yields the bounding box of the *intersection* instead: (8, 16), which is too small.
    The "left" block then writes C_buf_local[vi, vj] for vi in [8,12) — indices that
    were never loaded into C_buf_local — resulting in incorrect output.
    """

    @T.prim_func(s_tir=True)
    def main(A: T.handle, C: T.handle) -> None:
        A_buf = T.match_buffer(A, (12, 24), "float32")
        C_buf = T.match_buffer(C, (12, 24), "float32")
        for i, j in T.grid(12, 24):
            with T.sblock("compute"):
                vi, vj = T.axis.remap("SS", [i, j])
                with T.sblock("left"):
                    T.where(vi < 8)
                    C_buf[vi, vj] = A_buf[vi, vj] * 2.0
                with T.sblock("top"):
                    T.where(vj < 16)
                    C_buf[vi, vj] = A_buf[vi, vj] * 3.0

    sch = tvm.s_tir.Schedule(main)
    block = sch.get_sblock("compute")
    sch.cache_write(block, 0, "local")

    # Extract the alloc buffer shape from the resulting IR.
    result_script = sch.mod["main"].script()
    # The cache must be large enough to hold the union of both write regions.
    # Union bounding box = full (12, 24).  The buggy AND gives (8, 16).
    assert "sblock_alloc_buffer((12, 24)" in result_script, (
        f"Expected cache shape (12, 24) covering the union of both write regions, "
        f"but got a smaller shape. Full IR:\n{result_script}"
    )


def test_cache_read_sibling_nested_block_predicates_use_union():
    """Regression: cache_read with sibling nested blocks must union their predicates.

    Two sibling nested sblocks read the same input buffer under different predicates:
      left  block: T.where(vi < 8)   — reads rows 0-7, all columns
      top   block: T.where(vj < 16)  — reads all rows, columns 0-15

    The cache must cover the UNION of both read sets.  The bounding box of that
    union is (12, 24) — the full buffer shape.

    Bug: CollectNestedBlockPredicates ANDs the two predicates, giving (vi < 8) AND
    (vj < 16).  Case 2 of CacheRead calls RelaxBufferRegion under that intersection
    predicate, producing a cache of shape (8, 16).  The "left" block then tries to
    read A_buf_local[vi, vj] for vi in [8,12) — indices outside the cache — which
    is incorrect.
    """

    @T.prim_func(s_tir=True)
    def main(A: T.handle, C: T.handle) -> None:
        A_buf = T.match_buffer(A, (12, 24), "float32")
        C_buf = T.match_buffer(C, (12, 24), "float32")
        for i, j in T.grid(12, 24):
            with T.sblock("compute"):
                vi, vj = T.axis.remap("SS", [i, j])
                with T.sblock("left"):
                    T.where(vi < 8)
                    C_buf[vi, vj] = A_buf[vi, vj] * 2.0
                with T.sblock("top"):
                    T.where(vj < 16)
                    C_buf[vi, vj] = A_buf[vi, vj] * 3.0

    sch = tvm.s_tir.Schedule(main)
    block = sch.get_sblock("compute")
    sch.cache_read(block, 0, "local")

    result_script = sch.mod["main"].script()
    # Cache must cover the union bounding box (12, 24).  Buggy AND gives (8, 16).
    assert "sblock_alloc_buffer((12, 24)" in result_script, (
        f"Expected cache shape (12, 24) covering the union of both read regions, "
        f"but got a smaller shape. Full IR:\n{result_script}"
    )


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