# 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
import tvm.testing
from tvm import te


def test_bound1():
    m = te.var("m")
    l = te.var("l")
    A = te.placeholder((m, l), name="A")
    A1 = te.compute((m, l), lambda i, j: A[i, j], name="A1")
    A2 = te.compute((m, l), lambda i, j: A1[i, j] + 3, name="A2")

    s = te.create_schedule([A2.op])
    xo, xi = s[A2].split(s[A2].op.axis[0], 8)
    s[A1].compute_at(s[A2], xo)
    bounds = tvm.te.schedule.InferBound(s)
    assert isinstance(bounds, tvm.container.Map)
    assert bounds[A1.op.axis[0]].extent.value == 8


def test_bound2():
    m = te.var("m")
    l = te.var("l")
    A = te.placeholder((m, l), name="A")
    A1 = te.compute((m, l), lambda i, j: A[i, j], name="A1")
    A2 = te.compute((m, l), lambda i, j: A1[i, j] + 3, name="A2")
    s = te.create_schedule(A2.op)
    xo, yo, xi, yi = s[A2].tile(A2.op.axis[0], A2.op.axis[1], 8, 8)
    # test normalize not affecting schedule
    _ = s.normalize()
    s[A1].compute_at(s[A2], yo)
    bounds = tvm.te.schedule.InferBound(s)
    assert isinstance(bounds, tvm.container.Map)
    assert bounds[A1.op.axis[0]].extent.value == 8
    assert bounds[A1.op.axis[1]].extent.value == 8


def test_bound3():
    m = te.var("m")
    l = te.var("l")
    A = te.placeholder((m, l), name="A")
    A1 = te.compute((m, l), lambda i, j: A[i, j], name="A1")
    A2 = te.compute((m, l), lambda i, j: A1[i, j] + 3, name="A2")

    s = te.create_schedule(A2.op)
    s[A1].set_scope("shared")
    xo, xi = s[A2].split(A2.op.axis[0], 32)
    xi0, xi1 = s[A2].split(xi, nparts=16)
    s[A2].bind(xi0, te.thread_axis("threadIdx.x"))
    yo, yi = s[A2].split(A2.op.axis[1], 16)
    # test normalize not affecting schedule
    _ = s.normalize()
    s[A2].reorder(xo, xi0, yo, xi1, yi)
    s[A1].compute_at(s[A2], yo)

    bounds = tvm.te.schedule.InferBound(s)
    assert isinstance(bounds, tvm.container.Map)
    assert bounds[A1.op.axis[0]].extent.value == 32
    assert bounds[A1.op.axis[1]].extent.value == 16


def test_bound_split_ext_less_than_factor():
    m = 8
    I = te.placeholder((m,), name="I")
    EF = te.compute((m,), lambda i: I[i] * 2, name="EF")
    E = te.compute((m,), lambda i: EF[i] * 2, name="E")
    s = te.create_schedule([E.op])
    xo, xi = s[E].split(s[E].op.axis[0], factor=32)
    s[EF].compute_at(s[E], xo)

    bounds = tvm.te.schedule.InferBound(s)
    assert isinstance(bounds, tvm.container.Map)
    assert bounds[xi].extent.value == m


def test_bound_split_ext_less_than_naprts():
    m = 8
    I = te.placeholder((m,), name="I")
    EF = te.compute((m,), lambda i: I[i] * 2, name="EF")
    E = te.compute((m,), lambda i: EF[i] * 2, name="E")
    s = te.create_schedule([E.op])
    xo, xi = s[E].split(s[E].op.axis[0], nparts=32)
    s[EF].compute_at(s[E], xo)

    bounds = tvm.te.schedule.InferBound(s)
    assert isinstance(bounds, tvm.container.Map)
    assert bounds[xo].extent.value == m


def test_bound_split_divisible():
    m = te.var("m")
    l = te.var("l")
    A = te.placeholder((8 * m, l), name="A")
    B = te.compute((8 * m, l), lambda i, j: A[i, j], name="B")
    s = te.create_schedule(B.op)
    xo, xi = s[B].split(B.op.axis[0], 8)
    bounds = tvm.te.schedule.InferBound(s)
    assert isinstance(bounds, tvm.container.Map)
    assert bounds[xo].extent == m
    assert bounds[xi].extent.value == 8


def test_bound_tile_divisible():
    m = te.var("m")
    l = te.var("l")
    shape = (8 * m, 32 * l)
    A = te.placeholder(shape, name="A")
    B = te.compute(shape, lambda i, j: A[i, j], name="B")
    s = te.create_schedule(B.op)
    xo, yo, xi, yi = s[B].tile(B.op.axis[0], B.op.axis[1], 8, 32)
    bounds = tvm.te.schedule.InferBound(s)
    assert isinstance(bounds, tvm.container.Map)
    assert bounds[xo].extent == m
    assert bounds[xi].extent.value == 8
    assert bounds[yo].extent == l
    assert bounds[yi].extent.value == 32


def test_bound_fusesplit1():
    m = te.var("m")
    l = te.var("l")
    split1 = te.var("s")
    A = te.placeholder((m, l), name="A")
    A1 = te.compute((m, l), lambda i, j: A[i, j], name="A1")
    A2 = te.compute((m, l), lambda i, j: A1[i, j] + 3, name="A2")

    s = te.create_schedule(A2.op)
    fused_axes = s[A2].fuse(A2.op.axis[0], A2.op.axis[1])
    xo, xi = s[A2].split(fused_axes, split1)
    s[A1].compute_at(s[A2], xo)

    bounds = tvm.te.schedule.InferBound(s)
    assert isinstance(bounds, tvm.container.Map)
    idxdiv = tvm.tir.indexdiv
    tvm.testing.assert_prim_expr_equal(bounds[A1.op.axis[0]].min, idxdiv(xo * split1, l))

    expected_extent = idxdiv((xo + 1) * split1 - 1, l) - idxdiv(xo * split1, l) + 1
    for i in range(1, 6):
        for j in range(1, 6):
            for k in range(1, 6):
                vars = tvm.runtime.convert(
                    {
                        split1: tvm.tir.const(i, "int32"),
                        l: tvm.tir.const(j, "int32"),
                        xo.var: tvm.tir.const(k, "int32"),
                    }
                )
                tvm.testing.assert_prim_expr_equal(
                    tvm.tir.stmt_functor.substitute(bounds[A1.op.axis[0]].extent, vars),
                    tvm.tir.stmt_functor.substitute(expected_extent, vars),
                )

    tvm.testing.assert_prim_expr_equal(bounds[A1.op.axis[1]].extent, l)


def test_bound_fusesplit2():
    m = te.var("m")
    l = tvm.runtime.convert(6)
    split = tvm.runtime.convert(3)
    A = te.placeholder((m, l), name="A")
    A1 = te.compute((m, l), lambda i, j: A[i, j], name="A1")
    A2 = te.compute((m, l), lambda i, j: A1[i, j] + 3, name="A2")

    s = te.create_schedule(A2.op)
    fused_axes = s[A2].fuse(A2.op.axis[0], A2.op.axis[1])
    xo, xi = s[A2].split(fused_axes, split)
    s[A1].compute_at(s[A2], xo)

    bounds = tvm.te.schedule.InferBound(s)
    assert isinstance(bounds, tvm.container.Map)
    vars = tvm.runtime.convert({xo.var: tvm.tir.const(5, "int32")})
    tvm.testing.assert_prim_expr_equal(
        tvm.tir.stmt_functor.substitute(bounds[A1.op.axis[0]].min, vars), 2
    )
    tvm.testing.assert_prim_expr_equal(
        tvm.tir.stmt_functor.substitute(bounds[A1.op.axis[1]].min, vars), 3
    )
    tvm.testing.assert_prim_expr_equal(
        tvm.tir.stmt_functor.substitute(bounds[A1.op.axis[0]].extent, vars), 1
    )
    tvm.testing.assert_prim_expr_equal(
        tvm.tir.stmt_functor.substitute(bounds[A1.op.axis[1]].extent, vars), 3
    )


def test_bound_warp():
    m = te.var("m")
    l = te.var("l")
    A = te.placeholder((m, l), name="A")
    A1 = te.compute((m, l), lambda i, j: A[i, j], name="A1")
    A2 = te.compute((m, l), lambda i, j: A1[i, j] + 3, name="A2")

    s = te.create_schedule(A2.op)
    s[A1].set_scope("warp")
    xo, xi = s[A2].split(A2.op.axis[0], 32)
    xi0, xi1 = s[A2].split(xi, factor=16)
    tx = te.thread_axis("threadIdx.x")
    s[A2].bind(xi1, tx)
    s[A2].bind(xi0, te.thread_axis("threadIdx.y"))
    y = s[A2].op.axis[1]
    s[A1].compute_at(s[A2], y)
    xo, xi = s[A1].split(s[A1].op.axis[0], factor=16)
    s[A1].bind(xi, tx)
    bounds = tvm.te.schedule.InferBound(s)
    assert isinstance(bounds, tvm.container.Map)
    assert bounds[A1.op.axis[0]].extent.value == 16


def test_bound_scan():
    m = te.var("m")
    n = te.var("n")
    X = te.compute((m, n), lambda i, j: tvm.tir.const(1, "float32"), name="x")
    s_state = te.placeholder((m, n))
    s_init = te.compute((1, n), lambda _, i: X[0, i])
    s_update = te.compute((m, n), lambda t, i: s_state[t - 1, i] + X[t, i])
    s_scan = tvm.te.scan(s_init, s_update, s_state)

    assert tuple(s_scan.shape) == (m, n)
    s = te.create_schedule(s_scan.op)
    XX = s.cache_read(X, "local", s_update)
    xo, xi = s[s_update].split(s_update.op.axis[1], factor=4)
    s[XX].compute_at(s[s_update], xo)
    s = s.normalize()
    bounds = tvm.te.schedule.InferBound(s)
    stmt = tvm.te.schedule.ScheduleOps(s, bounds)
    assert bounds[XX.op.axis[1]].extent.value == 4


def test_bound_conv1d():
    n = te.var("n")
    A = te.compute((n + 2), lambda i: 1, name="A")

    def computeB(ii):
        i = ii + 1
        return A[i - 1] + A[i] + A[i + 1]

    B = te.compute(n, computeB, name="B")
    s = te.create_schedule(B.op)
    s[A].compute_at(s[B], B.op.axis[0])
    s = s.normalize()
    bounds = tvm.te.schedule.InferBound(s)
    assert bounds[A.op.axis[0]].extent.value == 3


def test_bound_blur():
    n = tvm.runtime.convert(12)
    A = te.compute((n, n), lambda i, j: 1, name="A")

    def computeB(ii, jj):
        # set the correct center
        i = ii + 1
        j = jj + 1
        return A[i][j] + A[i - 1][j] + A[i + 1][j] + A[i][j + 1] + A[i][j - 1]

    B = te.compute((n - 2, n - 2), computeB, name="B")
    s = te.create_schedule(B.op)
    s[A].compute_at(s[B], B.op.axis[1])
    s = s.normalize()
    bounds = tvm.te.schedule.InferBound(s)
    assert bounds[A.op.axis[0]].extent.value == 3
    assert bounds[A.op.axis[1]].extent.value == 3


def test_bound_rfactor():
    n = te.var("n")
    A = te.placeholder((n,), name="A")
    k = te.reduce_axis((0, n))
    B = te.compute((1,), lambda i: te.sum(A[k], axis=k, where=(i > 1)), name="B")
    # schedule
    s = te.create_schedule(B.op)
    kf, ki = s[B].split(k, nparts=4)
    BF = s.rfactor(B, kf)
    s = s.normalize()
    bounds = tvm.te.schedule.InferBound(s)

    assert bounds[BF.op.axis[0]].extent.value == 4
    assert bounds[BF.op.axis[1]].extent.value == 1


def test_bound_group_schedule():
    m = te.var("m")
    n = te.var("n")
    x = te.compute((m, n), lambda i, j: tvm.tir.const(1, "float32"), name="x")
    x1 = te.compute(x.shape, lambda *i: x(*i) + 1, name="x1")
    x2 = te.compute(x.shape, lambda *i: x1(*i) + 2, name="x2")
    s = te.create_schedule(x2.op)
    g = s.create_group(outputs=x1, inputs=x, include_inputs=True)
    g.compute_at(s[x2], x2.op.axis[0])
    assert s[x1].group == g
    assert s[x].group == g
    s = s.normalize()
    bounds = tvm.te.schedule.InferBound(s)
    assert bounds[x.op.axis[0]].extent.value == 1
    assert bounds[x.op.axis[1]].extent == n


def test_bound_nest_group():
    m = te.var("m")
    n = te.var("n")
    x = te.compute((m, n), lambda i, j: tvm.tir.const(1, "float32"), name="x")
    x1 = te.compute(x.shape, lambda *i: x(*i) + 1, name="x1")
    x2 = te.compute(x.shape, lambda *i: x1(*i) + 2, name="x2")
    s = te.create_schedule(x2.op)
    g1 = s.create_group(outputs=x, inputs=x, include_inputs=True)
    g2 = s.create_group(outputs=x1, inputs=x, include_inputs=True)
    assert s[x].group == g1
    assert s[x1].group == g2
    g2.compute_at(s[x2], x2.op.axis[0])
    g1.compute_at(s[x1], s[x1].op.axis[1])
    s = s.normalize()
    bounds = tvm.te.schedule.InferBound(s)
    assert bounds[x.op.axis[0]].extent.value == 1
    assert bounds[x.op.axis[1]].extent.value == 1
    assert bounds[x1.op.axis[0]].extent.value == 1
    assert bounds[x1.op.axis[1]].extent == n


def test_bound_nest_thread():
    m = te.var("m")
    A = te.placeholder((m), name="A")
    A1 = te.compute((m,), lambda i: A[i], name="A1")
    A2 = te.compute((m,), lambda i: A1[i] + 2, name="A2")
    A3 = te.compute((m,), lambda i: A2[i] + 3, name="A3")

    s = te.create_schedule(A3.op)
    s[A2].set_scope("shared")
    s[A1].set_scope("local")

    block_x = te.thread_axis("blockIdx.x")
    thread_x = te.thread_axis("threadIdx.x")
    bx, tx = s[A3].split(A3.op.axis[0], factor=32)
    s[A3].bind(bx, block_x)
    s[A3].bind(tx, thread_x)
    s[A2].compute_at(s[A3], tx)
    _, xi = s[A2].split(A2.op.axis[0], nparts=1)
    s[A2].bind(xi, thread_x)
    s[A1].compute_at(s[A3], tx)
    s = s.normalize()
    bounds = tvm.te.schedule.InferBound(s)
    assert bounds[A1.op.axis[0]].extent.value == 1
    assert bounds[A2.op.axis[0]].extent.value == 32
    assert bounds[A3.op.axis[0]].extent == m


def test_gemm_bound():
    nn = 1024
    n = tvm.runtime.convert(nn)
    A = te.placeholder((n, n), name="A")
    B = te.placeholder((n, n), name="B")
    k = te.reduce_axis((0, n), name="k")
    C = te.compute((n, n), lambda ii, jj: te.sum(A[ii, k] * B[jj, k], axis=k), name="CC")
    # schedule
    s = te.create_schedule(C.op)
    xtile, ytile = 32, 32
    scale = 8
    num_thread = 8
    block_factor = scale * num_thread
    block_x = te.thread_axis("blockIdx.x")
    thread_x = te.thread_axis("threadIdx.x")
    block_y = te.thread_axis("blockIdx.y")
    thread_y = te.thread_axis("threadIdx.y")

    CC = s.cache_write(C, "local")
    AA = s.cache_read(A, "shared", [CC])
    BB = s.cache_read(B, "shared", [CC])
    by, yi = s[C].split(C.op.axis[0], factor=block_factor)
    bx, xi = s[C].split(C.op.axis[1], factor=block_factor)
    s[C].reorder(by, bx, yi, xi)
    s[C].bind(by, block_y)
    s[C].bind(bx, block_x)
    ty, yi = s[C].split(yi, nparts=num_thread)
    tx, xi = s[C].split(xi, nparts=num_thread)
    s[C].reorder(ty, tx, yi, xi)
    s[C].bind(ty, thread_y)
    s[C].bind(tx, thread_x)
    yo, xo = CC.op.axis
    s[CC].reorder(k, yo, xo)

    s[CC].compute_at(s[C], tx)
    s[AA].compute_at(s[CC], k)
    s[BB].compute_at(s[CC], k)

    ty, xi = s[AA].split(s[AA].op.axis[0], nparts=num_thread)
    tx, xi = s[AA].split(xi, nparts=num_thread)
    s[AA].bind(ty, thread_y)
    s[AA].bind(tx, thread_x)

    ty, xi = s[BB].split(s[BB].op.axis[0], nparts=num_thread)
    tx, xi = s[BB].split(xi, nparts=num_thread)
    s[BB].bind(ty, thread_y)
    s[BB].bind(tx, thread_x)
    s = s.normalize()
    bounds = tvm.te.schedule.InferBound(s)
    assert bounds[BB.op.axis[0]].extent.value == 64
    assert bounds[AA.op.axis[0]].extent.value == 64
    assert bounds[CC.op.axis[0]].extent.value == 8
    assert bounds[CC.op.axis[1]].extent.value == 8


def test_bound_tensor_compute_op():
    def intrin_test():
        m1 = te.var("m1")
        n1 = te.var("n1")
        a = te.placeholder((m1, n1), name="a")
        c = te.compute((1, n1), lambda i, j: a[0, j] + a[1, j] + a[2, j], name="c")

        Ab = tvm.tir.decl_buffer(a.shape, name="Abuf", offset_factor=1)
        Cb = tvm.tir.decl_buffer(c.shape, name="Cbuf", offset_factor=1)

        def intrin_func(ins, outs):
            aa = ins[0]
            cc = outs[0]

            def _body():
                ib = tvm.tir.ir_builder.create()
                ib.emit(
                    tvm.tir.call_extern("int32", "test", cc.access_ptr("w"), aa.access_ptr("r"))
                )
                return ib.get()

            return _body()

        return te.decl_tensor_intrin(c.op, intrin_func, binds={a: Ab, c: Cb})

    test_func = intrin_test()
    A = te.placeholder((20, 20), name="A")
    B = te.compute(A.shape, lambda i, j: A[i, j], name="B")
    C = te.compute((10, 20), lambda i: test_func(B[i:10, 0:20]), name="C")
    s = te.create_schedule(C.op)
    bounds = tvm.te.schedule.InferBound(s)
    assert isinstance(bounds, tvm.container.Map)
    assert bounds[B.op.axis[0]].extent.value == 10


def test_bound_simplification_failure():
    # Check that the bounds are not expanded
    A = te.compute((2,), lambda j: j, "A")

    def _check(B, A=A):
        s = te.create_schedule(B.op)
        s = s.normalize()
        bounds = tvm.te.schedule.InferBound(s)
        stmt = tvm.lower(s, [B, A], simple_mode=True)
        if not bounds[A.op.axis[0]].extent.value <= 2:
            print(stmt)
            assert bounds[A.op.axis[0]].extent.value <= 2

    tdiv = tvm.tir.truncdiv
    # These are hard to simplify, moreover we don't simplify them
    _check(te.compute((10,), lambda i: A[tvm.te.min(3 * i, 4 * i) + tvm.te.min(-3 * i, -2 * i)]))
    _check(te.compute((10,), lambda i: A[tvm.te.min(3 * i, 4 * i) + tvm.te.max(-3 * i, -4 * i)]))
    _check(te.compute((10,), lambda i: A[-2 * tdiv(i, 2) - tvm.te.min(i, 0 - i)]))
    _check(te.compute((10,), lambda i: A[i + (0 - i)]))
    # This would cause out of bounds, but we nevertheless include it
    _check(te.compute((10,), lambda i: A[i]))


def test_bound_block():
    def _check(shape, expected, block_size=4):
        N, C, H, W = shape
        tail = C % block_size
        chunks = C // block_size
        if tail != 0:
            chunks += 1
        A = te.placeholder((N, C, H, W), name="A")
        pad_value = tvm.tir.const(0, A.dtype)

        def _reorder_data_nchw(*indices):
            condition = []
            condition.append(indices[1] == chunks - 1)
            condition.append(indices[4] >= tail)
            condition = tvm.tir.all(*condition)
            return tvm.tir.if_then_else(
                condition,
                pad_value,
                A[indices[0], indices[1] * block_size + indices[4], indices[2], indices[3]],
            )

        repack = te.compute((N, chunks, H, W, block_size), _reorder_data_nchw, name="repack")
        B = te.compute(
            (N, C, H, W),
            lambda n, c, h, w: repack[n, c // block_size, h, w, c % block_size],
            name="back_repack",
        )
        s = te.create_schedule([B.op])
        bounds = tvm.te.schedule.InferBound(s)
        # Block for intermediate compute function should be equal to 4 for all cases except than number of channels is less than 4
        assert bounds[repack.op.axis[4]].extent.value == expected

    _check((1, 4, 6, 6), 4)
    _check((1, 7, 6, 6), 4)
    _check((1, 3, 6, 6), 3)


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