# 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.
import tvm
from tvm import tir, te
from tvm.script import tir as T


def _check(original, transformed):
    func = original
    mod = tvm.IRModule.from_expr(func)
    mod = tvm.tir.transform.CompactBufferAllocation()(mod)
    mod = tvm.tir.transform.Simplify()(mod)
    tvm.ir.assert_structural_equal(mod["main"], transformed)


@T.prim_func
def elementwise_func(a: T.handle, c: T.handle) -> None:
    A = T.match_buffer(a, (16, 16), "float32")
    C = T.match_buffer(c, (16, 16), "float32")
    for i in range(0, 16):
        with T.block():
            T.reads(A[i, 0:16])
            T.writes(C[i, 0:16])
            B = T.alloc_buffer((16, 16), "float32")
            for j in range(0, 16):
                with T.block() as []:
                    T.reads(A[i, j])
                    T.writes(B[i, j])
                    B[i, j] = A[i, j] + 1.0
            for j in range(0, 16):
                with T.block() as []:
                    T.reads(B[i, j])
                    T.writes(C[i, j])
                    C[i, j] = B[i, j] * 2.0


@T.prim_func
def compacted_elementwise_func(a: T.handle, c: T.handle) -> None:
    A = T.match_buffer(a, (16, 16), "float32")
    C = T.match_buffer(c, (16, 16), "float32")
    for i in range(0, 16):
        with T.block():
            T.reads(A[i, 0:16])
            T.writes(C[i, 0:16])
            B = T.alloc_buffer((1, 16), "float32")
            for j in range(0, 16):
                with T.block() as []:
                    T.reads(A[i, j])
                    T.writes(B[0, j])
                    B[0, j] = A[i, j] + 1.0
            for j in range(0, 16):
                with T.block() as []:
                    T.reads(B[0, j])
                    T.writes(C[i, j])
                    C[i, j] = B[0, j] * 2.0


@T.prim_func
def unschedulable_func(a: T.handle, c: T.handle) -> None:
    A = T.match_buffer(a, (16, 16), "float32")
    C = T.match_buffer(c, (16, 16), "float32")
    for i in range(0, 16):
        with T.block():
            T.reads(A[i, 0:16])
            T.writes(C[i, 0:16])
            B = T.alloc_buffer((16, 16), "float32")
            for j in range(0, 16):
                T.store(B.data, i * 16 + j, A[i, j] + 1.0)
            for j in range(0, 16):
                C[i, j] = B[i, j] * 2.0


@T.prim_func
def param_buffer_access_func(a: T.handle, c: T.handle) -> None:
    A = T.match_buffer(a, (20, 20), "float32")
    B = T.match_buffer(c, (20, 20), "float32")
    for i in range(0, 16):
        with T.block():
            T.reads(A[i, 0:16])
            T.writes(B[i, 0:16])
            for j in range(0, 16):
                with T.block() as []:
                    T.reads(A[i, j])
                    T.writes(B[i, j])
                    B[i, j] = A[i, j] + 1.0


@T.prim_func
def shared_mem_func(a: T.handle, c: T.handle) -> None:
    A = T.match_buffer(a, (16, 16), "float32")
    C = T.match_buffer(c, (16, 16), "float32")
    for i0 in T.thread_binding(0, 2, thread="blockIdx.x"):
        for i1 in T.thread_binding(0, 2, thread="vthread"):
            for i2 in T.thread_binding(0, 4, thread="threadIdx.x"):
                with T.block():
                    T.reads(A[i0 * 8 + i1 * 4 + i2, 0:16])
                    T.writes(C[i0 * 8 + i1 * 4 + i2, 0:16])
                    B = T.alloc_buffer((16, 16), "float32", scope="shared")
                    for j in range(0, 16):
                        with T.block() as []:
                            T.reads(A[i0 * 8 + i1 * 4 + i2, j])
                            T.writes(B[i0 * 8 + i1 * 4 + i2, j])
                            B[i0 * 8 + i1 * 4 + i2, j] = A[i0 * 8 + i1 * 4 + i2, j] + 1.0
                    for j in range(0, 16):
                        with T.block() as []:
                            T.reads(B[i0 * 8 + i1 * 4 + i2, j])
                            T.writes(C[i0 * 8 + i1 * 4 + i2, j])
                            C[i0 * 8 + i1 * 4 + i2, j] = B[i0 * 8 + i1 * 4 + i2, j] * 2.0


@T.prim_func
def compacted_shared_mem_func(a: T.handle, c: T.handle) -> None:
    A = T.match_buffer(a, (16, 16), "float32")
    C = T.match_buffer(c, (16, 16), "float32")
    for i0 in T.thread_binding(0, 2, thread="blockIdx.x"):
        for i1 in T.thread_binding(0, 2, thread="vthread"):
            for i2 in T.thread_binding(0, 4, thread="threadIdx.x"):
                with T.block():
                    T.reads(A[i0 * 8 + i1 * 4 + i2, 0:16])
                    T.writes(C[i0 * 8 + i1 * 4 + i2, 0:16])
                    B = T.alloc_buffer((8, 16), "float32", scope="shared")
                    for j in range(0, 16):
                        with T.block() as []:
                            T.reads(A[i0 * 8 + i1 * 4 + i2, j])
                            T.writes(B[i1 * 4 + i2, j])
                            B[i1 * 4 + i2, j] = A[i0 * 8 + i1 * 4 + i2, j] + 1.0
                    for j in range(0, 16):
                        with T.block() as []:
                            T.reads(B[i1 * 4 + i2, j])
                            T.writes(C[i0 * 8 + i1 * 4 + i2, j])
                            C[i0 * 8 + i1 * 4 + i2, j] = B[i1 * 4 + i2, j] * 2.0


@T.prim_func
def warp_mem_func(a: T.handle, c: T.handle) -> None:
    A = T.match_buffer(a, (16, 16), "float32")
    C = T.match_buffer(c, (16, 16), "float32")
    for i0 in T.thread_binding(0, 2, thread="blockIdx.x"):
        for i1 in T.thread_binding(0, 2, thread="vthread"):
            for i2 in T.thread_binding(0, 4, thread="threadIdx.x"):
                with T.block():
                    T.reads(A[i0 * 8 + i1 * 4 + i2, 0:16])
                    T.writes(C[i0 * 8 + i1 * 4 + i2, 0:16])
                    B = T.alloc_buffer((16, 16), "float32", scope="warp")
                    for j in range(0, 16):
                        with T.block() as []:
                            T.reads(A[i0 * 8 + i1 * 4 + i2, j])
                            T.writes(B[i0 * 8 + i1 * 4 + i2, j])
                            B[i0 * 8 + i1 * 4 + i2, j] = A[i0 * 8 + i1 * 4 + i2, j] + 1.0
                    for j in range(0, 16):
                        with T.block() as []:
                            T.reads(B[i0 * 8 + i1 * 4 + i2, j])
                            T.writes(C[i0 * 8 + i1 * 4 + i2, j])
                            C[i0 * 8 + i1 * 4 + i2, j] = B[i0 * 8 + i1 * 4 + i2, j] * 2.0


@T.prim_func
def compacted_warp_mem_func(a: T.handle, c: T.handle) -> None:
    A = T.match_buffer(a, (16, 16), "float32")
    C = T.match_buffer(c, (16, 16), "float32")
    for i0 in T.thread_binding(0, 2, thread="blockIdx.x"):
        for i1 in T.thread_binding(0, 2, thread="vthread"):
            for i2 in T.thread_binding(0, 4, thread="threadIdx.x"):
                with T.block():
                    T.reads(A[i0 * 8 + i1 * 4 + i2, 0:16])
                    T.writes(C[i0 * 8 + i1 * 4 + i2, 0:16])
                    B = T.alloc_buffer((4, 16), "float32", scope="warp")
                    for j in range(0, 16):
                        with T.block() as []:
                            T.reads(A[i0 * 8 + i1 * 4 + i2, j])
                            T.writes(B[i2, j])
                            B[i2, j] = A[i0 * 8 + i1 * 4 + i2, j] + 1.0
                    for j in range(0, 16):
                        with T.block() as []:
                            T.reads(B[i2, j])
                            T.writes(C[i0 * 8 + i1 * 4 + i2, j])
                            C[i0 * 8 + i1 * 4 + i2, j] = B[i2, j] * 2.0


@T.prim_func
def symbolic_func(a: T.handle, c: T.handle, n: T.int32) -> None:
    A = T.match_buffer(a, (n * 8,), "float32")
    C = T.match_buffer(c, (n * 8,), "float32")
    for i in range(0, n):
        with T.block():
            T.reads(A[i * 8 : i * 8 + 8])
            T.writes(C[i * 8 : i * 8 + 8])
            B = T.alloc_buffer((n * 8,), "float32")
            for j in range(0, 8):
                with T.block() as []:
                    T.reads(A[i * 8 + j])
                    T.writes(B[i * 8 + j])
                    B[i * 8 + j] = A[i * 8 + j] + 1.0
            for j in range(0, 8):
                with T.block() as []:
                    T.reads(B[i * 8 + j])
                    T.writes(C[i * 8 + j])
                    C[i * 8 + j] = B[i * 8 + j] * 2.0


@T.prim_func
def compacted_symbolic_func(a: T.handle, c: T.handle, n: T.int32) -> None:
    A = T.match_buffer(a, (n * 8,), "float32")
    C = T.match_buffer(c, (n * 8,), "float32")
    for i in range(0, n):
        with T.block():
            T.reads(A[i * 8 : i * 8 + 8])
            T.writes(C[i * 8 : i * 8 + 8])
            B = T.alloc_buffer((8,), "float32")
            for j in range(0, 8):
                with T.block() as []:
                    T.reads(A[i * 8 + j])
                    T.writes(B[j])
                    B[j] = A[i * 8 + j] + 1.0
            for j in range(0, 8):
                with T.block() as []:
                    T.reads(B[j])
                    T.writes(C[i * 8 + j])
                    C[i * 8 + j] = B[j] * 2.0


@T.prim_func
def complex_func(a: T.handle, c: T.handle, n: T.int32) -> None:
    A = T.match_buffer(a, (8, 8), "float32")
    C = T.match_buffer(c, (8, 8), "float32")
    for i in range(0, 8):
        with T.block():
            T.reads(A[0, 8])
            T.writes(C[0, 8])
            B = T.alloc_buffer((8, 8), "float32")
            for j in range(0, 4):
                with T.block() as []:
                    D = T.alloc_buffer((8, 8), "float32")
                    T.reads(A[i, j])
                    T.writes(B[i, j])
                    for k in range(4, 8):
                        D[k, j] = 1.0
                    for k in range(2, 4):
                        T.store(B.data, j, A[i, j] + D[k, j])
            for j in range(3, 5):
                with T.block() as []:
                    T.reads(B[i, j])
                    T.writes(C[i, j])
                    C[i, j] = B[i, j]
            for j in range(6, 8):
                with T.block() as []:
                    T.reads(B[i, j])
                    T.writes(C[i, j])
                    C[i, j] = B[i, j]


@T.prim_func
def compacted_complex_func(a: T.handle, c: T.handle, n: T.int32) -> None:
    A = T.match_buffer(a, (8, 8), "float32")
    C = T.match_buffer(c, (8, 8), "float32")
    for i in range(0, 8):
        with T.block():
            T.reads(A[0, 8])
            T.writes(C[0, 8])
            B = T.alloc_buffer((1, 8), "float32")
            for j in range(0, 4):
                with T.block() as []:
                    D = T.alloc_buffer((6, 1), "float32")
                    T.reads(A[i, j])
                    T.writes(B[0, j])
                    for k in range(4, 8):
                        D[k - 2, 0] = 1.0
                    for k in range(2, 4):
                        T.store(B.data, j, A[i, j] + D[k - 2, 0])
            for j in range(3, 5):
                with T.block() as []:
                    T.reads(B[0, j])
                    T.writes(C[i, j])
                    C[i, j] = B[0, j]
            for j in range(6, 8):
                with T.block() as []:
                    T.reads(B[0, j])
                    T.writes(C[i, j])
                    C[i, j] = B[0, j]


@T.prim_func
def match_buffer_func(a: T.handle, c: T.handle) -> None:
    A = T.match_buffer(a, (16, 16))
    C = T.match_buffer(c, (16, 16))
    for i in range(0, 16):
        with T.block():
            A0 = T.match_buffer(A[i, 0:16], (16))
            C0 = T.match_buffer(C[i, 0:16], (16))
            B = T.alloc_buffer((16, 16))
            with T.block():
                B0 = T.match_buffer(B[i, 0:16], (16))
                for j in range(0, 16):
                    with T.block() as []:
                        A1 = T.match_buffer(A0[j], ())
                        B1 = T.match_buffer(B0[j], ())
                        B1[()] = A1[()] + 1.0
            for j in range(0, 16):
                with T.block() as []:
                    C1 = T.match_buffer(C0[j], ())
                    B2 = T.match_buffer(B[i, j], ())
                    C1[()] = B2[()] * 2.0


@T.prim_func
def compacted_match_buffer_func(a: T.handle, c: T.handle) -> None:
    A = T.match_buffer(a, (16, 16))
    C = T.match_buffer(c, (16, 16))
    for i in range(0, 16):
        with T.block():
            A0 = T.match_buffer(A[i, 0:16], (16))
            C0 = T.match_buffer(C[i, 0:16], (16))
            B = T.alloc_buffer((1, 16))
            with T.block():
                B0 = T.match_buffer(B[0, 0:16], (16))
                for j in range(0, 16):
                    with T.block() as []:
                        A1 = T.match_buffer(A0[j], ())
                        B1 = T.match_buffer(B0[j], ())
                        B1[()] = A1[()] + 1.0
            for j in range(0, 16):
                with T.block() as []:
                    C1 = T.match_buffer(C0[j], ())
                    B2 = T.match_buffer(B[0, j], ())
                    C1[()] = B2[()] * 2.0


@T.prim_func
def storage_align_func(a: T.handle, c: T.handle) -> None:
    A = T.match_buffer(a, (16, 16), "float32")
    C = T.match_buffer(c, (16, 16), "float32")
    for i in range(0, 16):
        with T.block():
            T.reads(A[i, 0:16])
            T.writes(C[i, 0:16])
            B = T.alloc_buffer((16, 16), "float32")
            for j in range(0, 16):
                with T.block() as []:
                    T.reads(A[i, j])
                    T.writes(B[i, j])
                    T.block_attr({"buffer_dim_align": [[0, 0, 16, 15]]})
                    B[i, j] = A[i, j] + 1.0
            for j in range(0, 16):
                with T.block() as []:
                    T.reads(B[i, j])
                    T.writes(C[i, j])
                    C[i, j] = B[i, j] * 2.0


@T.prim_func
def compacted_storage_align_func(a: T.handle, c: T.handle) -> None:
    A = T.match_buffer(a, (16, 16), "float32")
    C = T.match_buffer(c, (16, 16), "float32")
    for i in range(0, 16):
        with T.block():
            T.reads(A[i, 0:16])
            T.writes(C[i, 0:16])
            B = T.alloc_buffer((1, 16), strides=(31, 1), dtypes="float32")
            for j in range(0, 16):
                with T.block() as []:
                    T.reads(A[i, j])
                    T.writes(B[0, j])
                    T.block_attr({"buffer_dim_align": [[0, 0, 16, 15]]})
                    B[0, j] = A[i, j] + 1.0
            for j in range(0, 16):
                with T.block() as []:
                    T.reads(B[0, j])
                    T.writes(C[i, j])
                    C[i, j] = B[0, j] * 2.0


@T.prim_func
def padding_pattern_func(a: T.handle, c: T.handle) -> None:
    A = T.match_buffer(a, (16, 16), "float32")
    C = T.match_buffer(c, (20, 20), "float32")
    with T.block():
        B = T.alloc_buffer((20, 20), dtypes="float32")
        for i, j in T.grid(16, 16):
            with T.block():
                B[i, j] = A[i, j]
        for i, j in T.grid(20, 20):
            with T.block():
                C[i, j] = T.if_then_else(
                    2 <= i and i < 18 and 2 <= j and j < 18,
                    B[i - 2, j - 2],
                    0.0,
                    dtype="float32",
                )


@T.prim_func
def compacted_padding_pattern_func(a: T.handle, c: T.handle) -> None:
    A = T.match_buffer(a, [16, 16], dtype="float32")
    C = T.match_buffer(c, [20, 20], dtype="float32")
    with T.block():
        B = T.alloc_buffer([16, 16], dtype="float32")
        for i, j in T.grid(16, 16):
            with T.block():
                B[i, j] = A[i, j]
        for i, j in T.grid(20, 20):
            with T.block():
                C[i, j] = T.if_then_else(
                    2 <= i and i < 18 and 2 <= j and j < 18, B[i - 2, j - 2], 0.0, dtype="float32"
                )


@T.prim_func
def mem_access_in_branch_func(a: T.handle) -> None:
    A = T.match_buffer(a, (224, 224), "float32")
    with T.block():
        B1 = T.alloc_buffer((224, 224), dtypes="float32")
        B2 = T.alloc_buffer((224, 224), dtypes="float32")
        B3 = T.alloc_buffer((224, 224), dtypes="float32")
        B4 = T.alloc_buffer((224, 224), dtypes="float32")
        for i in range(0, 224):
            for j in range(0, 224):
                with T.block():
                    if i < 112 and j < 112:
                        B1[i, j] = A[i, j] * 2.0
                    else:
                        B2[i, j] = A[i, j] + 3.0
        for i in range(0, 224):
            for j in range(0, 224):
                with T.block():
                    if i < 112 or j < 112:
                        B3[i, j] = A[i, j] * 2.0
                    else:
                        B4[i, j] = A[i, j] + 3.0


@T.prim_func
def compacted_mem_access_in_branch_func(a: T.handle) -> None:
    A = T.match_buffer(a, [224, 224], dtype="float32")
    with T.block():
        B1 = T.alloc_buffer([112, 112], dtype="float32")
        B2 = T.alloc_buffer([224, 224], dtype="float32")
        B3 = T.alloc_buffer([224, 224], dtype="float32")
        B4 = T.alloc_buffer([112, 112], dtype="float32")
        for i, j in T.grid(224, 224):
            with T.block():
                if i < 112 and j < 112:
                    B1[i, j] = A[i, j] * 2.0
                else:
                    B2[i, j] = A[i, j] + 3.0
        for i, j in T.grid(224, 224):
            with T.block():
                if i < 112 or j < 112:
                    B3[i, j] = A[i, j] * 2.0
                else:
                    B4[i - 112, j - 112] = A[i, j] + 3.0


@T.prim_func
def opaque_access_annotated_func(a: T.handle) -> None:
    A = T.match_buffer(a, (1024,), "float32")
    with T.block():
        B = T.alloc_buffer((1024,), dtypes="float32")
        C = T.alloc_buffer((1024,), dtypes="float32")
        for i in range(0, 512):
            with T.block():
                # no annotation, opaque access will cover full region
                T.reads([])
                T.writes([])
                T.store(B.data, i, "float32", A[i])
            with T.block():
                # treat opaque access only access annotated regions, even if
                # they are not compatible with actual buffer accesses.
                T.reads([B[i]])
                T.writes([C[i : i + 9]])
                T.store(C.data, i, T.load("float32", B.data, i))


@T.prim_func
def compacted_opaque_access_annotated_func(a: T.handle) -> None:
    A = T.match_buffer(a, (1024,), "float32")
    with T.block():
        B = T.alloc_buffer((1024,), dtypes="float32")
        C = T.alloc_buffer((520,), dtypes="float32")
        for i in range(0, 512):
            with T.block():
                # no annotation, opaque access will cover full region
                T.reads([])
                T.writes([])
                T.store(B.data, i, "float32", A[i])
            with T.block():
                # treat opaque access only access annotated regions, even if
                # they are not compatible with actual buffer accesses.
                T.reads([B[i]])
                T.writes([C[i : i + 9]])
                T.store(C.data, i, T.load("float32", B.data, i))


def test_elementwise():
    _check(elementwise_func, compacted_elementwise_func)


def test_unschedulable_block():
    _check(unschedulable_func, unschedulable_func)  # changes nothing


def test_param_access():
    _check(param_buffer_access_func, param_buffer_access_func)  # changes nothing


def test_shared_mem():
    _check(shared_mem_func, compacted_shared_mem_func)


def test_warp_mem():
    _check(warp_mem_func, compacted_warp_mem_func)


def test_symbolic():
    _check(symbolic_func, compacted_symbolic_func)


def test_complex():
    _check(complex_func, compacted_complex_func)


def test_match_buffer():
    _check(match_buffer_func, compacted_match_buffer_func)


def test_lower_te():
    x = te.placeholder((1,))
    y = te.compute((1,), lambda i: x[i] + 2)
    s = te.create_schedule(y.op)
    orig_mod = tvm.driver.build_module.schedule_to_module(s, [x, y])
    mod = tvm.tir.transform.CompactBufferAllocation()(orig_mod)
    tvm.ir.assert_structural_equal(mod, orig_mod)  # CompactBufferAllocation should do nothing on TE


def test_storage_align():
    _check(storage_align_func, compacted_storage_align_func)


def test_padding_pattern():
    _check(padding_pattern_func, compacted_padding_pattern_func)


def test_mem_access_in_branch_func():
    _check(mem_access_in_branch_func, compacted_mem_access_in_branch_func)


def test_opaque_access_annotated_func():
    _check(opaque_access_annotated_func, compacted_opaque_access_annotated_func)


if __name__ == "__main__":
    test_elementwise()
    test_unschedulable_block()
    test_param_access()
    test_shared_mem()
    test_warp_mem()
    test_symbolic()
    test_complex()
    test_match_buffer()
    test_storage_align()
    test_lower_te()
    test_padding_pattern()
    test_mem_access_in_branch_func()
    test_opaque_access_annotated_func()