# 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 numpy as np

import tvm
from tvm import te
from tvm.driver.build_module import schedule_to_module


def test_const():
    x = tvm.te.const(1, "int32")
    assert x.dtype == "int32"
    assert isinstance(x, tvm.tir.IntImm)


def test_schedule0():
    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")
    s = te.create_schedule(A1.op)

    mod = schedule_to_module(s, [A, A1])
    assert isinstance(mod["main"], tvm.tir.PrimFunc)


def test_schedule1():
    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")

    s = te.create_schedule(A1.op)
    xo, xi = s[A1].split(A1.op.axis[0], 8)
    s[A1].pragma(xo, "auto_unroll_max_step", 10)

    mod = schedule_to_module(s, [A, A1])
    assert isinstance(mod["main"], tvm.tir.PrimFunc)


def test_schedule2():
    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(A2.op.axis[0], 8)
    s[A1].compute_at(s[A2], xo)

    mod = schedule_to_module(s, [A, A2])
    assert isinstance(mod["main"], tvm.tir.PrimFunc)


def test_schedule_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])
    res = tvm.te.scan(s_init, s_update, s_state)

    assert tuple(res.shape) == (m, n)
    s = te.create_schedule(res.op)
    s = s.normalize()
    ir = tvm.lower(s, [s_state], simple_mode=True)
    bounds = tvm.te.schedule.InferBound(s)
    assert bounds[res.op.scan_axis].min.value == 1
    stmt = tvm.te.schedule.ScheduleOps(s, bounds)


def test_inline_multi_reduce():
    def argmax_comp(x, y):
        idx = tvm.tir.Select((x[1] >= y[1]), x[0], y[0])
        val = tvm.tir.Select((x[1] >= y[1]), x[1], y[1])
        return idx, val

    def argmax_init(idx_typ, val_typ):
        return tvm.tir.const(-1, idx_typ), tvm.te.min_value(val_typ)

    argmax = te.comm_reducer(argmax_comp, argmax_init, name="argmax")
    m = te.var("m")
    n = te.var("n")
    val = te.placeholder((m, n), name="val", dtype="float32")
    val1 = te.compute((m, n), lambda i, j: val[i, j] + 1, name="val1")
    val2 = te.compute((m, n), lambda i, j: te.exp(val1[i, j]), name="val2")
    k = te.reduce_axis((0, n), "k")
    T_idx, T_val = te.compute((m,), lambda i: argmax((k.var, val2[i, k]), axis=k), name="T")
    s = te.create_schedule(T_idx.op)
    s[val1].compute_inline()
    s = s.normalize()
    bounds = tvm.te.schedule.InferBound(s)
    stmt = tvm.te.schedule.ScheduleOps(s, bounds)


def test_auto_inline():
    def elemwise():
        m = te.var("m")
        n = te.var("n")
        A = te.placeholder((m, n), name="A")
        B = te.placeholder((m, n), name="B")
        C = te.placeholder((m, n), name="C")
        T1 = te.compute((m, n), lambda i, j: A(i, j) * B(i, j), name="T1")
        T2 = te.compute((m, n), lambda i, j: T1(i, j) + C(i, j), name="T2")

        return te.create_schedule(T2.op), T1

    def broadcast():
        m = te.var("m")
        n = te.var("n")
        A = te.placeholder((1,), name="A")
        B = te.placeholder((m, n), name="B")
        C = te.placeholder((m, n), name="C")
        T1 = te.compute((m, n), lambda i, j: A(0) * B(i, j), name="T1", tag="broadcast")
        T2 = te.compute((m, n), lambda i, j: T1(i, j) + C(i, j), name="T2")

        return te.create_schedule(T2.op), T1

    def injective():
        m = te.var("m")
        n = te.var("n")
        A = te.placeholder((m,), name="A")
        B = te.placeholder((m, n), name="B")
        C = te.placeholder((m, n), name="C")
        T1 = te.compute((m, n), lambda i, j: A(i) * B(i, j), name="T1")
        T2 = te.compute((m, n), lambda i, j: T1(i, j) + C(i, j), name="T2")

        return te.create_schedule(T2.op), T1

    def check_auto_inline(schedule_func, auto_inline_func):
        s, T1 = schedule_func()
        # before auto inline the attach type is AttachType.kGroupRoot
        assert s[T1].attach_type == 1
        auto_inline_func(s)
        # after auto inline the attach type is AttachType.kInline
        assert s[T1].attach_type == 2
        s = s.normalize()
        bounds = tvm.te.schedule.InferBound(s)
        stmt = tvm.te.schedule.ScheduleOps(s, bounds)

    check_auto_inline(elemwise, tvm.te.schedule.AutoInlineElemWise)
    check_auto_inline(broadcast, tvm.te.schedule.AutoInlineBroadcast)
    check_auto_inline(injective, tvm.te.schedule.AutoInlineInjective)


def test_schedule_const_bound():
    n = 128
    A = te.placeholder((n,), name="A")
    A1 = te.compute((n,), lambda i: A[i] + 1, name="A1")
    s = te.create_schedule(A1.op)
    xo, xi = s[A1].split(A1.op.axis[0], 8)
    bounds = tvm.te.schedule.InferBound(s)
    assert isinstance(bounds, tvm.container.Map)
    stmt = tvm.te.schedule.ScheduleOps(s, bounds)


def test_inline_mixed():
    n = te.var("n")
    A = te.placeholder((n,), name="A")
    A1 = te.compute(A.shape, lambda *i: A(*i) + 1, name="A1")
    A2 = te.compute(A.shape, lambda *i: A1(*i) + 2, name="A2")
    C = te.compute((n,), lambda i: A2[i] + A1[i], name="C")

    s = te.create_schedule(C.op)
    xo, xi = s[C].split(C.op.axis[0], factor=8)
    s[A1].compute_at(s[C], xo)
    s[A2].compute_inline()
    s = s.normalize()
    bounds = tvm.te.schedule.InferBound(s)
    stmt = tvm.te.schedule.ScheduleOps(s, bounds)

    def check(x):
        if isinstance(x, tvm.tir.Call):
            assert x.func != A2

    tvm.tir.stmt_functor.post_order_visit(s[C].op.body[0], check)


def test_scan_inline1():
    m = te.var("m")
    n = te.var("n")
    x = te.compute((m, n), lambda i, j: tvm.tir.const(1, "float32"), name="x")
    s_state1 = te.placeholder((m, n))
    s_state2 = te.placeholder((m, n))
    s_init1 = te.compute((1, n), lambda _, i: x[0, i])
    s_init2 = te.compute((1, n), lambda _, i: x[0, i])
    s_x1 = te.compute((m, n), lambda t, i: s_state1[t - 1, i] + x[t, i], name="x1")
    s_x2 = te.compute((m, n), lambda t, i: s_state2[t - 1, i] + 1, name="x2")
    s_update1 = te.compute((m, n), lambda t, i: s_x1[t, i], "u1")
    s_update2 = te.compute((m, n), lambda t, i: s_x2[t, i], "u2")
    res1, res2 = tvm.te.scan([s_init1, s_init2], [s_update1, s_update2], [s_state1, s_state2])
    s = te.create_schedule(res1.op)
    s[s_x1].compute_inline()
    stmt = tvm.lower(s, [x, res1, res2])


def test_scan_inline2():
    m = te.var("m")
    n = te.var("n")
    x = te.compute((m, n), lambda i, j: tvm.tir.const(1, "float32"), name="x")
    s_state1 = te.placeholder((m, n))
    s_state2 = te.placeholder((m, n))
    s_init1 = te.compute((1, n), lambda _, i: x[0, i])
    s_init2 = te.compute((1, n), lambda _, i: x[0, i])
    s_xx = te.compute((m, n), lambda t, i: s_state1[t - 1, i] + x[t, i], name="xx")
    s_x1 = te.compute((m, n), lambda t, i: s_xx[t, i] + 1, name="x1")
    s_x2 = te.compute((m, n), lambda t, i: s_xx[t, i] + s_state2[t - 1, 2], name="x2")
    s_update1 = te.compute((m, n), lambda t, i: s_x1[t, i], "u1")
    s_update2 = te.compute((m, n), lambda t, i: s_x2[t, i], "u2")
    res1, res2 = tvm.te.scan([s_init1, s_init2], [s_update1, s_update2], [s_state1, s_state2])
    s = te.create_schedule(res1.op)
    s[s_xx].compute_inline()
    s[s_x1].compute_inline()
    s[s_x2].compute_inline()
    stmt = tvm.lower(s, [x, res1, res2])


def test_schedule_cache():
    m = te.var("m")
    n = te.var("n")
    A = te.placeholder((m, n), name="A")
    B = te.placeholder((m, n), name="B")
    C = te.compute((m, n), lambda i, j: A(i, j) * B(i, j), name="C")

    s = te.create_schedule(C.op)
    AA = s.cache_read(A, "shared", readers=[C])
    CC = s.cache_write(C, "shared")
    s[AA].compute_at(s[CC], CC.op.axis[0])
    bounds = tvm.te.schedule.InferBound(s)
    stmt = tvm.te.schedule.ScheduleOps(s, bounds)


def test_schedule_middle_cache():
    m = te.var("m")
    n = te.var("n")
    A = te.placeholder((m, n), name="A")
    B = te.placeholder((m, n), name="B")

    C = te.compute((m, n), lambda i, j: A(i, j) * B(i, j), name="C")
    D = te.compute((m, n), lambda i, j: C(i, j), name="D")

    s = te.create_schedule(D.op)
    AA = s.cache_read(A, "local", readers=[C])
    BB = s.cache_read(B, "local", readers=[C])
    CC = s.cache_read(C, "local", readers=[D])
    DD = s.cache_write(D, "local")
    # s[AA].compute_at(s[CC], CC.op.axis[0])
    bounds = tvm.te.schedule.InferBound(s)
    stmt = tvm.te.schedule.ScheduleOps(s, bounds)


def test_schedule_cache_relayout1():
    m = te.var("m")
    n = te.var("n")
    A = te.placeholder((m, n), name="A")
    B = te.placeholder((m, n), name="B")
    C = te.compute((m, n), lambda i, j: A(i, j) * B(i, j), name="C")

    s = te.create_schedule(C.op)
    s[C].reorder(C.op.axis[1], C.op.axis[0])
    CC = s.cache_write(C, "global")
    bounds = tvm.te.schedule.InferBound(s)
    stmt = tvm.te.schedule.ScheduleOps(s, bounds)


def test_schedule_cache_relayout2():
    m = te.var("m")
    n = te.var("n")
    A = te.placeholder((m * 4, n), name="A")
    B = te.placeholder((m * 4, n), name="B")
    C = te.compute(A.shape, lambda i, j: A(i, j) * B(i, j), name="C")
    s = te.create_schedule(C.op)
    x, y = C.op.axis
    xo, xi = s[C].split(x, factor=4)
    s[C].reorder(xo, y, xi)
    CC = s.cache_write(C, "global")
    s = s.normalize()
    bounds = tvm.te.schedule.InferBound(s)
    stmt = tvm.te.schedule.ScheduleOps(s, bounds)


def test_schedule_cache_relayout3():
    m = te.var("m")
    n = te.var("n")
    A = te.placeholder((m * 4, n), name="A")
    B = te.placeholder((m * 4, n), name="B")
    k = te.reduce_axis((0, n), "k")
    C = te.compute((A.shape[0],), lambda i: te.sum(A(i, k) * B(i, k), axis=k), name="C")
    s = te.create_schedule(C.op)
    x = C.op.axis[0]
    xo, xi = s[C].split(x, factor=4)
    CC = s.cache_write(C, "global")
    s = s.normalize()
    bounds = tvm.te.schedule.InferBound(s)
    stmt = tvm.te.schedule.ScheduleOps(s, bounds)


def test_schedule_cache_relayout4():
    def _compute(*indice):
        return A(*indice) + 1, B(*indice) / 2

    m = te.var("m")
    n = te.var("n")
    A = te.placeholder((m * 4, n), name="A")
    B = te.placeholder((m * 4, n), name="B")
    C1, C2 = te.compute(A.shape, _compute, name="C")
    s = te.create_schedule([C1.op, C2.op])
    C1_cache, C2_cache = s.cache_write([C1, C2], "local")
    s = s.normalize()
    bounds = tvm.te.schedule.InferBound(s)
    stmt = tvm.te.schedule.ScheduleOps(s, bounds)


def intrin_gemv(m, n):
    w = te.placeholder((m, n), name="w")
    x = te.placeholder((n,), name="x")
    k = te.reduce_axis((0, n), name="k")
    z = te.compute((m,), lambda i: te.sum(w[i, k] * x[k], axis=k), name="z")
    Wb = tvm.tir.decl_buffer(
        w.shape, w.dtype, name="W", offset_factor=16, strides=[te.var("ldw"), 1]
    )

    def intrin_func(ins, outs):
        ww, xx = ins
        zz = outs[0]
        ww_ptr = ww.access_ptr("r")
        xx_ptr = xx.access_ptr("r")
        zz_ptr = zz.access_ptr("w")
        body = tvm.tir.call_packed("gemm", ww_ptr, xx_ptr, zz_ptr, n, ww.strides[0])
        reset = tvm.tir.call_packed("fill_zero", zz_ptr, n)
        update = tvm.tir.call_packed("gemv_add", ww_ptr, xx_ptr, zz_ptr, n, ww.strides[0])
        return body, reset, update

    buffer_params = {"data_alignment": 16, "offset_factor": 16}
    return te.decl_tensor_intrin(
        z.op, intrin_func, binds={w: Wb}, default_buffer_params=buffer_params
    )


def test_schedule_tensor_compute1():
    # basic: split, reorder, tile
    M, N, L = 2048, 1024, 512
    factor, rfactor = 16, 16
    A = te.placeholder((N // factor, L // rfactor, factor, rfactor), name="A")
    B = te.placeholder((M, L // rfactor, rfactor), name="B")
    k = te.reduce_axis((0, L // rfactor), name="k")

    gemv = intrin_gemv(factor, rfactor)
    C = te.compute(
        (N, M // factor, factor),
        lambda i, j: gemv(A[i, k, 0:factor, 0:factor], B[j, k, 0:rfactor], reduce_axis=k),
        name="C",
    )

    s = te.create_schedule(C.op)
    ai, aj, ax = s[C].op.axis
    aio, aii = s[C].split(ai, 16)
    s[C].reorder(aio, aj, aii)
    aioo, ajo, aioi, aji = s[C].tile(aio, aj, 16, 4)

    s = s.normalize()
    bounds = tvm.te.schedule.InferBound(s)
    stmt = tvm.te.schedule.ScheduleOps(s, bounds)


def intrin_vadd(n, cache_read=False, cache_write=False):
    scope_ubuf = "local"
    dtype = "float32"
    x = te.placeholder((n,), dtype=dtype, name="vx")
    y = te.placeholder((n,), dtype=dtype, name="vy")
    z = te.compute(x.shape, lambda i: x[i] + y[i], name="z")
    s = te.create_schedule(z.op)

    def create_buffer(t):
        return tvm.tir.decl_buffer(
            t.shape, t.dtype, name="W" + t.name, scope=scope_ubuf, offset_factor=16
        )

    binds = {}
    if cache_read:
        binds[x] = create_buffer(x)
        binds[y] = create_buffer(y)
    if cache_write:
        binds[z] = create_buffer(z)

    def intrin_func(ins, outs):
        ib = tvm.tir.ir_builder.create()
        ib.emit(
            tvm.tir.call_extern(
                outs[0].dtype,
                "vadd",
                ins[0].access_ptr("r"),
                ins[1].access_ptr("r"),
                outs[0].access_ptr("wr"),
            )
        )
        return ib.get()

    return te.decl_tensor_intrin(
        z.op, intrin_func, binds=binds, default_buffer_params={"offset_factor": 16}
    )


def test_schedule_tensor_compute2():
    # cache_read, cache_write
    M = 1024
    factor = 16
    dtype = "float32"
    scope_ubuf = "local"

    A = te.placeholder((M // factor, factor), name="A", dtype=dtype)
    B = te.placeholder((M // factor, factor), name="B", dtype=dtype)

    vadd = intrin_vadd(factor, True, True)
    C = te.compute((M // factor, factor), lambda i: vadd(A[i, 0:factor], B[i, 0:factor]), name="C")

    s = te.create_schedule(C.op)
    AL = s.cache_read(A, scope_ubuf, C)
    BL = s.cache_read(B, scope_ubuf, C)
    CL = s.cache_write(C, scope_ubuf)
    s = s.normalize()
    bounds = tvm.te.schedule.InferBound(s)
    stmt = tvm.te.schedule.ScheduleOps(s, bounds)


def test_schedule_tensor_compute3():
    # compute_at
    M = 1024
    factor = 16
    dtype = "float32"
    A = te.placeholder((M // factor, factor), name="A", dtype=dtype)
    B = te.placeholder((M // factor, factor), name="B", dtype=dtype)
    Bi = te.compute((M // factor, factor), lambda i, j: B[i, j] + 5, name="Bi")

    vadd = intrin_vadd(factor)
    C = te.compute((M // factor, factor), lambda i: vadd(A[i, 0:factor], Bi[i, 0:factor]), name="C")
    s = te.create_schedule(C.op)
    s[Bi].compute_at(s[C], C.op.axis[0])
    s = s.normalize()
    bounds = tvm.te.schedule.InferBound(s)
    stmt = tvm.te.schedule.ScheduleOps(s, bounds)


def test_loop_dep_reduce():
    X = te.placeholder(shape=(10,), name="x")

    def f(n):
        rv = te.reduce_axis((0, n))
        return te.sum(X[rv], axis=rv)

    Y = te.compute(X.shape, f, name="y")
    s = te.create_schedule([Y.op])
    f = tvm.build(s, [X, Y])


def test_loop_dep_reduce_cache_write():
    X = te.placeholder(shape=(10,), name="x")

    def f(n):
        rv = te.reduce_axis((0, n))
        init = lambda dtype: tvm.tir.Select(n > 1, tvm.tir.const(0, dtype), n.astype(dtype))
        sum = te.comm_reducer(lambda x, y: tvm.te.max(x + y, n.astype("float32")), init, name="sum")
        return sum(X[rv], axis=rv)

    Y = te.compute(X.shape, f, name="y")
    s = te.create_schedule([Y.op])
    s.cache_write(Y, "local")
    f = tvm.build(s, [X, Y])


def test_reduction_and_dummy_fuse_split():
    n = 10
    X = te.placeholder(shape=(n,), dtype="int32", name="X")
    k = te.reduce_axis((0, n))
    Y = te.compute((), lambda: te.sum(X[k], k), name="Y")
    s = te.create_schedule([Y.op])
    ax = s[Y.op].fuse(*Y.op.axis)
    axo, axi = s[Y.op].split(ax, nparts=20)
    f = tvm.build(s, [Y, X])

    args = [tvm.nd.empty((), "int32")] + [tvm.nd.array(np.ones((n,), dtype="int32"))]
    f(*args)
    assert args[0].numpy() == n

    n = 10
    X = te.placeholder(shape=(n,), dtype="int32", name="X")
    k = te.reduce_axis((0, n))
    Y = te.compute((n,), lambda i: te.sum(X[k], k), name="Y")
    s = te.create_schedule([Y.op])
    ax = s[Y.op].fuse(*(list(Y.op.axis) + list(Y.op.reduce_axis)))
    f = tvm.build(s, [Y, X])

    args = [tvm.nd.array(np.ones((n,), dtype="int32"))] + [
        tvm.nd.array(np.ones((n,), dtype="int32"))
    ]
    f(*args)
    assert np.all(args[0].numpy() == n)


def test_schedule_compute_inline():
    shape = [10, 1024]
    A = te.placeholder(shape, name="A")
    B = te.placeholder(shape, name="B")
    C = te.compute(shape, lambda *index: A(*index) + B(*index), name="C")

    def _compute(*index):
        return C(*index), C(*index) * B(*index)

    F, E = te.compute(shape, _compute, name="F")

    s = te.create_schedule([F.op, E.op])
    AL = s.cache_read(A, "local", [C])
    BL = s.cache_read(B, "local", [C, E])
    CL = s.cache_write(C, "local")
    FL, EL = s.cache_write([F, E], "local")
    s[C].compute_inline()

    s = s.normalize()
    bounds = tvm.te.schedule.InferBound(s)
    stmt = tvm.te.schedule.ScheduleOps(s, bounds)


def test_local_stage_predicate():
    m = 1
    n = 3
    p = 2
    A = tvm.te.placeholder((m, n, p), name="A")
    B = tvm.te.compute((m, n, p), lambda bi, bj, bk: A[bi, bj, bk], name="B")
    C = tvm.te.compute((m, n, p), lambda ci, cj, ck: B[ci, cj, ck], name="C")
    by = tvm.te.thread_axis("blockIdx.y")
    tx = tvm.te.thread_axis("threadIdx.x")
    vx = tvm.te.thread_axis("vthread")

    def schedule(thread_tag, mem_scope):
        s = tvm.te.create_schedule(C.op)
        s[B].compute_at(s[C], s[C].op.axis[0])
        s[B].set_scope(mem_scope)
        bno, bni = s[B].split(s[B].op.axis[1], n)
        bx = tvm.te.thread_axis("blockIdx.x")
        s[C].bind(s[C].op.axis[0], bx)
        s[C].bind(s[C].op.axis[1], thread_tag)
        s[B].bind(bni, thread_tag)
        return s

    def collect_visit(stmt, f):
        ret = []
        tvm.tir.stmt_functor.post_order_visit(stmt, lambda x: ret.append(f(x)))
        return ret

    # local vs. threadIdx
    s = schedule(tx, "local")
    lowered_body = tvm.lower(s, [A, C])["main"].body
    assert not any(collect_visit(lowered_body, lambda x: isinstance(x, tvm.tir.IfThenElse)))
    # local vs. vthread
    s = schedule(vx, "local")
    lowered_body = tvm.lower(s, [A, C])["main"].body
    assert not any(collect_visit(lowered_body, lambda x: isinstance(x, tvm.tir.IfThenElse)))
    # shared vs. blockIdx
    s = schedule(by, "shared")
    lowered_body = tvm.lower(s, [A, C])["main"].body
    assert not any(collect_visit(lowered_body, lambda x: isinstance(x, tvm.tir.IfThenElse)))


def test_local_stage_predicate2():
    A = tvm.te.placeholder((128,), name="A")
    B = tvm.te.compute((128,), lambda bi: A[bi] + 1, name="B")
    C = tvm.te.compute((128,), lambda ci: B[ci] + 2, name="C")
    s = tvm.te.create_schedule(C.op)
    AA = s.cache_read(A, "local", [B])
    s[B].set_scope("shared")
    block_x = tvm.te.thread_axis("blockIdx.x")
    thread_x = tvm.te.thread_axis((0, 32), "threadIdx.x")
    oc, ic = s[C].split(s[C].op.axis[0], factor=64)
    ooc, ioc = s[C].split(oc, factor=2)
    oic, iic = s[C].split(ic, factor=32)
    s[C].bind(ooc, block_x)
    s[C].bind(iic, thread_x)
    s[B].compute_at(s[C], ioc)
    ob, ib = s[B].split(s[B].op.axis[0], factor=32)
    s[B].bind(ib, thread_x)
    s[AA].compute_root()
    s[AA].compute_at(s[C], ooc)
    oaa, iaa = s[AA].split(s[AA].op.axis[0], factor=32)
    s[AA].bind(iaa, thread_x)
    lowered_body = tvm.lower(s, [A, C])["main"].body

    def collect_visit(stmt, f):
        ret = []
        tvm.tir.stmt_functor.post_order_visit(stmt, lambda x: ret.append(f(x)))
        return ret

    def visit_stmt(op):
        if isinstance(op, tvm.tir.Allocate):
            return op.extents[0].value == 97
        return False

    assert not any(collect_visit(lowered_body, lambda x: isinstance(x, tvm.tir.IfThenElse)))
    assert any(collect_visit(lowered_body, visit_stmt))


def test_schedule_record_gemm():
    with tvm.transform.PassContext(config={"te.keep_schedule_record": True}):
        M, K, N = 1024, 1024, 1024
        k = te.reduce_axis((0, K), "k")
        A = te.placeholder((M, K), name="A")
        B = te.placeholder((K, N), name="B")
        C = te.compute((M, N), lambda m, n: te.sum(A[m, k] * B[k, n], axis=k), name="C")
        s = te.create_schedule(C.op)
        # currently there are no other applied primitives
        # size of schedule record is expected to be 1 (vanilla schedule)
        assert len(s.schedule_record) == 1
        # apply sequential optimizatoin primitives
        block_size, factor = 32, 8
        # tile -> split + split + reorder
        mo, no, mi, ni = s[C].tile(C.op.axis[0], C.op.axis[1], block_size, block_size)
        ko, ki = s[C].split(k, factor=factor)
        s[C].reorder(mo, ko, no, mi, ki, ni)
        s[C].vectorize(ni)
        s[C].parallel(mo)
        assert len(s.schedule_record) == 8
        # compare primitive names
        expected_names = [
            "vanilla",
            "split",
            "split",
            "reorder",
            "split",
            "reorder",
            "vectorize",
            "parallel",
        ]
        for i in range(len(s.schedule_record)):
            assert s.primitive_record[i] == expected_names[i]


def test_schedule_record_misc():
    s = te.create_schedule([])
    # size of schedule record is expected to be 0 (no storing behavior)
    assert len(s.schedule_record) == 0

    with tvm.transform.PassContext(config={"te.keep_schedule_record": True}):
        s = te.create_schedule([])
        # size of schedule record is expected to be 1 (vanilla schedule)
        assert len(s.schedule_record) == 1

        stg = te.compute((), lambda *args: 0, name="empty_op")
        s = te.create_schedule(stg.op)
        # size of schedule record is expected to be 1 (vanilla schedule)
        assert len(s.schedule_record) == 1


if __name__ == "__main__":
    test_loop_dep_reduce()
    test_loop_dep_reduce_cache_write()
    test_schedule_middle_cache()
    test_inline_multi_reduce()
    test_schedule_cache_relayout4()
    test_schedule_cache_relayout3()
    test_schedule_cache_relayout2()
    test_schedule_cache_relayout1()
    test_schedule_const_bound()
    test_scan_inline1()
    test_scan_inline2()
    test_inline_mixed()
    test_auto_inline()
    test_schedule_scan()
    test_schedule0()
    test_schedule1()
    test_schedule2()
    test_schedule_cache()
    test_schedule_tensor_compute1()
    test_schedule_tensor_compute2()
    test_schedule_tensor_compute3()
    test_reduction_and_dummy_fuse_split()
    test_schedule_compute_inline()
    test_local_stage_predicate()
    test_local_stage_predicate2()
    test_schedule_record_gemm()
    test_schedule_record_misc()
