# 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
import pytest
import tvm
from tvm import tir
from tvm.script import ty

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


@tvm.script.tir
def elementwise(a: ty.handle, c: ty.handle) -> None:
    A = tir.match_buffer(a, (128, 128))
    B = tir.alloc_buffer((128, 128))
    C = tir.match_buffer(c, (128, 128))
    with tir.block([128, 128], "B") as [vi, vj]:
        B[vi, vj] = A[vi, vj] * 2.0
    with tir.block([128, 128], "C") as [vi, vj]:
        C[vi, vj] = B[vi, vj] + 1.0


@tvm.script.tir
def elementwise_multi_producer_consumer(a: ty.handle, c: ty.handle, d: ty.handle) -> None:
    A = tir.match_buffer(a, (128, 128))
    B = tir.alloc_buffer((128, 128))
    C = tir.match_buffer(c, (128, 128))
    D = tir.match_buffer(d, (128, 128))
    with tir.block([128, 128], "B") as [vi, vj]:
        B[vi, vj] = A[vi, vj] * 2.0  # B has two consumers
    with tir.block([128, 128], "C") as [vi, vj]:
        C[vi, vj] = B[vi, vj] + 1.0
    with tir.block([128, 128], "D") as [vi, vj]:
        D[vi, vj] = B[vi, vj] + 2.0 + C[vi, vj]  # D has two producers


@tvm.script.tir
def elementwise_multi_consumer_inlined(a: ty.handle, c: ty.handle, d: ty.handle) -> None:
    A = tir.match_buffer(a, (128, 128))
    C = tir.match_buffer(c, (128, 128))
    D = tir.match_buffer(d, (128, 128))
    with tir.block([128, 128], "C") as [vi, vj]:
        C[vi, vj] = A[vi, vj] * 2.0 + 1.0
    with tir.block([128, 128], "D") as [vi, vj]:
        D[vi, vj] = A[vi, vj] * 2.0 + 2.0 + C[vi, vj]


@tvm.script.tir
def elementwise_standalone(a: ty.handle, c: ty.handle) -> None:
    A = tir.match_buffer(a, (128, 128))
    B = tir.alloc_buffer((128, 128))
    C = tir.match_buffer(c, (128, 128))
    with tir.block([128, 128], "B") as [vi, vj]:
        B[vi, vj] = A[vi, vj] * 2.0
    with tir.block([128, 128], "C") as [vi, vj]:
        C[vi, vj] = A[vi, vj] + 1.0


@tvm.script.tir
def elementwise_standalone_dce(a: ty.handle, c: ty.handle) -> None:
    A = tir.match_buffer(a, (128, 128))
    C = tir.match_buffer(c, (128, 128))
    with tir.block([128, 128], "C") as [vi, vj]:
        C[vi, vj] = A[vi, vj] + 1.0


@tvm.script.tir
def elementwise_under_loop(a: ty.handle, c: ty.handle) -> None:
    A = tir.match_buffer(a, (128, 128))
    C = tir.match_buffer(c, (128, 128))
    B = tir.alloc_buffer((128, 128))
    for i in tir.serial(0, 128):
        for j in tir.serial(0, 128):
            with tir.block([128, 128], "B") as [vi, vj]:
                tir.bind(vi, i)
                tir.bind(vj, j)
                B[vi, vj] = A[vi, vj] * 2.0
        for j in tir.serial(0, 128):
            with tir.block([128, 128], "C") as [vi, vj]:
                tir.bind(vi, i)
                tir.bind(vj, j)
                C[vi, vj] = B[vi, vj] + 1.0


@tvm.script.tir
def elementwise_inlined(a: ty.handle, c: ty.handle) -> None:
    A = tir.match_buffer(a, (128, 128))
    C = tir.match_buffer(c, (128, 128))
    with tir.block([128, 128], "C") as [vi, vj]:
        C[vi, vj] = A[vi, vj] * 2.0 + 1.0


@tvm.script.tir
def fail_multi_reader_writer(a: ty.handle, d: ty.handle) -> None:
    A = tir.match_buffer(a, (128, 128))
    B = tir.alloc_buffer((128, 128))
    C = tir.alloc_buffer((128, 128))
    D = tir.match_buffer(d, (128, 128))
    with tir.block([128, 128], "B") as [vi, vj]:
        B[vi, vj] = A[vi, vj] * 2.0
        C[vi, vj] = A[vi, vj] + 2.0
    with tir.block([128, 128], "C") as [vi, vj]:
        D[vi, vj] = B[vi, vj] + C[vi, vj]


@tvm.script.tir
def elementwise_multi_reverse_loads(a: ty.handle, c: ty.handle) -> None:
    A = tir.match_buffer(a, (128, 128))
    B = tir.alloc_buffer((128, 128))
    C = tir.match_buffer(c, (128, 128))
    with tir.block([128, 128], "B") as [vi, vj]:
        B[vi, vj] = A[vi, vj] * 2.0
    with tir.block([128, 128], "C") as [vi, vj]:
        C[vi, vj] = (B[vi, vj] + 1.0) * (B[vi, vj] * 2.0) + 3.0


@tvm.script.tir
def elementwise_multi_reverse_loads_inlined(a: ty.handle, c: ty.handle) -> None:
    A = tir.match_buffer(a, (128, 128))
    C = tir.match_buffer(c, (128, 128))
    with tir.block([128, 128], "B") as [vi, vj]:
        C[vi, vj] = (A[vi, vj] * 2.0 + 1.0) * (A[vi, vj] * 2.0 * 2.0) + 3.0


@tvm.script.tir
def opaque_access_load(a: ty.handle, c: ty.handle) -> None:
    A = tir.match_buffer(a, (128, 128))
    B = tir.alloc_buffer((128, 128))
    C = tir.match_buffer(c, (128, 128))
    with tir.block([128, 128], "B") as [vi, vj]:
        B[vi, vj] = A[vi, vj] * 2.0
    with tir.block([128, 128], "C") as [vi, vj]:
        tir.reads(B[0:128, 0:128])
        tir.writes(C[0:128, 0:128])
        C[vi, vj] = tir.load("float32", B.data, vi * 128 + vj) + 1.0


@tvm.script.tir
def opaque_access_store(a: ty.handle, c: ty.handle) -> None:
    A = tir.match_buffer(a, (128, 128))
    B = tir.alloc_buffer((128, 128))
    C = tir.match_buffer(c, (128, 128))
    with tir.block([128, 128], "B") as [vi, vj]:
        B[vi, vj] = A[vi, vj] * 2.0
    with tir.block([128, 128], "C") as [vi, vj]:
        tir.reads(B[0:128, 0:128])
        tir.writes(C[0:128, 0:128])
        tir.store(C.data, vi * 128 + vj, B[vi, vj] + 1.0)
        C[vi, vj] = tir.load("float32", B.data, vi * 16 + vj) + 1.0


@tvm.script.tir
def buffer_matched(a: ty.handle, c: ty.handle) -> None:
    A = tir.match_buffer(a, (128, 128))
    B = tir.alloc_buffer((128, 128))
    C = tir.match_buffer(c, (128, 128))
    with tir.block([128, 128], "B") as [vi, vj]:
        B[vi, vj] = A[vi, vj] * 2.0
    with tir.block([128, 128], "C") as [vi, vj]:
        Bb = tir.match_buffer(B[vi : vi + 1, vj], (1, 1))
        C[vi, vj] = Bb[0, 0] + 1.0


@tvm.script.tir
def elementwise_predicate(a: ty.handle, c: ty.handle) -> None:
    A = tir.match_buffer(a, (128, 128))
    B = tir.alloc_buffer((128, 128))
    C = tir.match_buffer(c, (128, 128))
    with tir.block([128, 128], "B") as [vi, vj]:
        B[vi, vj] = A[vi, vj] * 2.0
    for i, j in tir.grid(128, 128):
        with tir.block([128, 128], "C") as [vi, vj]:
            tir.where(B[i, j] < 10.0)
            C[vi, vj] = B[vi, vj] + 1.0


@tvm.script.tir
def elementwise_predicate_inlined(a: ty.handle, c: ty.handle) -> None:
    A = tir.match_buffer(a, (128, 128))
    C = tir.match_buffer(c, (128, 128))
    for i, j in tir.grid(128, 128):
        with tir.block([128, 128], "C") as [vi, vj]:
            tir.where(A[i, j] * 2.0 < 10.0)
            C[vi, vj] = A[vi, vj] * 2.0 + 1.0


@tvm.script.tir
def elementwise_multi_loads(a: ty.handle, c: ty.handle) -> None:
    A = tir.match_buffer(a, (128, 128))
    B = tir.alloc_buffer((128, 128))
    C = tir.match_buffer(c, (128, 128))
    with tir.block([128, 128], "B") as [vi, vj]:
        B[vi, vj] = A[vi, vj] * 2.0
    with tir.block([128, 126], "C") as [vi, vj]:
        C[vi, vj] = B[vi, vj] + B[vi, vj + 1] + B[vi, vj + 2]


@tvm.script.tir
def elementwise_multi_loads_inlined(a: ty.handle, c: ty.handle) -> None:
    A = tir.match_buffer(a, (128, 128))
    C = tir.match_buffer(c, (128, 128))
    with tir.block([128, 126], "C") as [vi, vj]:
        C[vi, vj] = A[vi, vj] * 2.0 + A[vi, vj + 1] * 2.0 + A[vi, vj + 2] * 2.0


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


def test_compute_inline_elementwise():
    sch = tir.Schedule(elementwise, debug_mode=True)
    block_b = sch.get_block("B")
    block_c = sch.get_block("C")
    sch.compute_inline(block_b)
    tvm.ir.assert_structural_equal(elementwise_inlined, sch.mod["main"])
    assert sch.get(block_c).name_hint == "C"


def test_compute_inline_under_loop():
    sch = tir.Schedule(elementwise_under_loop, debug_mode=True)
    block_b = sch.get_block("B")
    block_c = sch.get_block("C")
    sch.compute_inline(block_b)
    tvm.ir.assert_structural_equal(elementwise_inlined, sch.mod["main"])
    assert sch.get(block_c).name_hint == "C"


def test_compute_inline_as_dce():
    sch = tir.Schedule(elementwise_standalone, debug_mode=True)
    block_b = sch.get_block("B")
    block_c = sch.get_block("C")
    sch.compute_inline(block_b)
    tvm.ir.assert_structural_equal(elementwise_standalone_dce, sch.mod["main"])
    assert sch.get(block_c).name_hint == "C"


def test_compute_inline_multi_consumer():
    sch = tir.Schedule(elementwise_multi_producer_consumer, debug_mode=True)
    block_b = sch.get_block("B")
    block_c = sch.get_block("C")
    block_d = sch.get_block("D")
    sch.compute_inline(block_b)
    tvm.ir.assert_structural_equal(elementwise_multi_consumer_inlined, sch.mod["main"])
    assert sch.get(block_c).name_hint == "C"
    assert sch.get(block_d).name_hint == "D"


def test_compute_inline_fail_multi_writer():
    sch = tir.Schedule(fail_multi_reader_writer, debug_mode=True, error_render_level="detail")
    block_b = sch.get_block("B")
    with pytest.raises(tvm.tir.ScheduleError):
        sch.compute_inline(block_b)


def test_reverse_compute_inline_elementwise():
    sch = tir.Schedule(elementwise, debug_mode=True)
    block_b = sch.get_block("B")
    block_c = sch.get_block("C")
    sch.reverse_compute_inline(block_c)
    tvm.ir.assert_structural_equal(elementwise_inlined, sch.mod["main"])
    assert sch.get(block_b).name_hint == "B"


def test_reverse_compute_inline_under_loop():
    sch = tir.Schedule(elementwise_under_loop, debug_mode=True)
    block_b = sch.get_block("B")
    block_c = sch.get_block("C")
    sch.reverse_compute_inline(block_c)
    tvm.ir.assert_structural_equal(elementwise_inlined, sch.mod["main"])
    assert sch.get(block_b).name_hint == "B"


def test_reverse_compute_inline_fail_as_dce():
    sch = tir.Schedule(elementwise_standalone, debug_mode=True)
    block_b = sch.get_block("B")
    with pytest.raises(tvm.tir.ScheduleError):
        sch.reverse_compute_inline(block_b)


def test_reverse_compute_inline_fail_multi_producer():
    sch = tir.Schedule(elementwise_multi_producer_consumer, debug_mode=True)
    block_d = sch.get_block("D")
    with pytest.raises(tvm.tir.ScheduleError):
        sch.reverse_compute_inline(block_d)


def test_reverse_compute_inline_fail_multi_reader():
    sch = tir.Schedule(fail_multi_reader_writer, debug_mode=True)
    block_c = sch.get_block("C")
    with pytest.raises(tvm.tir.ScheduleError):
        sch.reverse_compute_inline(block_c)


def test_reverse_compute_multi_reverse_loads():
    sch = tir.Schedule(elementwise_multi_reverse_loads, debug_mode=True)
    block_c = sch.get_block("C")
    sch.reverse_compute_inline(block_c)
    tvm.ir.assert_structural_equal(elementwise_multi_reverse_loads_inlined, sch.mod["main"])


def test_reverse_compute_fail_multi_reverse_loads():
    sch = tir.Schedule(elementwise_multi_loads, debug_mode=True)
    block_c = sch.get_block("C")
    with pytest.raises(tvm.tir.ScheduleError):
        sch.reverse_compute_inline(block_c)


def test_opaque_access_load():
    sch = tir.Schedule(opaque_access_load, debug_mode=True)
    block_b = sch.get_block("B")
    with pytest.raises(tvm.tir.ScheduleError):
        sch.compute_inline(block_b)


def test_opaque_access_store():
    sch = tir.Schedule(opaque_access_store, debug_mode=True)
    block_b = sch.get_block("B")
    with pytest.raises(tvm.tir.ScheduleError):
        sch.compute_inline(block_b)


def test_buffer_matched():
    sch = tir.Schedule(buffer_matched, debug_mode=True)
    block_b = sch.get_block("B")
    with pytest.raises(tvm.tir.ScheduleError):
        sch.compute_inline(block_b)


def test_compute_inline_predicate():
    sch = tir.Schedule(elementwise_predicate, debug_mode=True)
    block_b = sch.get_block("B")
    sch.compute_inline(block_b)
    tvm.ir.assert_structural_equal(elementwise_predicate_inlined, sch.mod["main"])


def test_compute_inline_multi_loads():
    sch = tir.Schedule(elementwise_multi_loads, debug_mode=True)
    block_b = sch.get_block("B")
    sch.compute_inline(block_b)
    tvm.ir.assert_structural_equal(elementwise_multi_loads_inlined, sch.mod["main"])


if __name__ == "__main__":
    test_compute_inline_elementwise()
    test_compute_inline_under_loop()
    test_compute_inline_as_dce()
    test_compute_inline_multi_consumer()
    test_compute_inline_fail_multi_writer()
    test_reverse_compute_inline_elementwise()
    test_reverse_compute_inline_under_loop()
    test_reverse_compute_inline_fail_as_dce()
    test_reverse_compute_inline_fail_multi_producer()
    test_reverse_compute_inline_fail_multi_reader()
    test_reverse_compute_multi_reverse_loads()
    test_reverse_compute_fail_multi_reverse_loads()
    test_opaque_access_load()
    test_opaque_access_store()
    test_buffer_matched()
    test_compute_inline_predicate()
    test_compute_inline_multi_loads()

[((1063, 1094), 'tvm.tir.match_buffer', 'tir.match_buffer', (['a', '(128, 128)'], {}), '(a, (128, 128))\n', (1079, 1094), False, 'from tvm import tir\n'), ((1103, 1131), 'tvm.tir.alloc_buffer', 'tir.alloc_buffer', (['(128, 128)'], {}), '((128, 128))\n', (1119, 1131), False, 'from tvm import tir\n'), ((1140, 1171), 'tvm.tir.match_buffer', 'tir.match_buffer', (['c', '(128, 128)'], {}), '(c, (128, 128))\n', (1156, 1171), False, 'from tvm import tir\n'), ((1459, 1490), 'tvm.tir.match_buffer', 'tir.match_buffer', (['a', '(128, 128)'], {}), '(a, (128, 128))\n', (1475, 1490), False, 'from tvm import tir\n'), ((1499, 1527), 'tvm.tir.alloc_buffer', 'tir.alloc_buffer', (['(128, 128)'], {}), '((128, 128))\n', (1515, 1527), False, 'from tvm import tir\n'), ((1536, 1567), 'tvm.tir.match_buffer', 'tir.match_buffer', (['c', '(128, 128)'], {}), '(c, (128, 128))\n', (1552, 1567), False, 'from tvm import tir\n'), ((1576, 1607), 'tvm.tir.match_buffer', 'tir.match_buffer', (['d', '(128, 128)'], {}), '(d, (128, 128))\n', (1592, 1607), False, 'from tvm import tir\n'), ((2037, 2068), 'tvm.tir.match_buffer', 'tir.match_buffer', (['a', '(128, 128)'], {}), '(a, (128, 128))\n', (2053, 2068), False, 'from tvm import tir\n'), ((2077, 2108), 'tvm.tir.match_buffer', 'tir.match_buffer', (['c', '(128, 128)'], {}), '(c, (128, 128))\n', (2093, 2108), False, 'from tvm import tir\n'), ((2117, 2148), 'tvm.tir.match_buffer', 'tir.match_buffer', (['d', '(128, 128)'], {}), '(d, (128, 128))\n', (2133, 2148), False, 'from tvm import tir\n'), ((2433, 2464), 'tvm.tir.match_buffer', 'tir.match_buffer', (['a', '(128, 128)'], {}), '(a, (128, 128))\n', (2449, 2464), False, 'from tvm import tir\n'), ((2473, 2501), 'tvm.tir.alloc_buffer', 'tir.alloc_buffer', (['(128, 128)'], {}), '((128, 128))\n', (2489, 2501), False, 'from tvm import tir\n'), ((2510, 2541), 'tvm.tir.match_buffer', 'tir.match_buffer', (['c', '(128, 128)'], {}), '(c, (128, 128))\n', (2526, 2541), False, 'from tvm import tir\n'), ((2806, 2837), 'tvm.tir.match_buffer', 'tir.match_buffer', (['a', '(128, 128)'], {}), '(a, (128, 128))\n', (2822, 2837), False, 'from tvm import tir\n'), ((2846, 2877), 'tvm.tir.match_buffer', 'tir.match_buffer', (['c', '(128, 128)'], {}), '(c, (128, 128))\n', (2862, 2877), False, 'from tvm import tir\n'), ((3053, 3084), 'tvm.tir.match_buffer', 'tir.match_buffer', (['a', '(128, 128)'], {}), '(a, (128, 128))\n', (3069, 3084), False, 'from tvm import tir\n'), ((3093, 3124), 'tvm.tir.match_buffer', 'tir.match_buffer', (['c', '(128, 128)'], {}), '(c, (128, 128))\n', (3109, 3124), False, 'from tvm import tir\n'), ((3133, 3161), 'tvm.tir.alloc_buffer', 'tir.alloc_buffer', (['(128, 128)'], {}), '((128, 128))\n', (3149, 3161), False, 'from tvm import tir\n'), ((3175, 3193), 'tvm.tir.serial', 'tir.serial', (['(0)', '(128)'], {}), '(0, 128)\n', (3185, 3193), False, 'from tvm import tir\n'), ((3686, 3717), 'tvm.tir.match_buffer', 'tir.match_buffer', (['a', '(128, 128)'], {}), '(a, (128, 128))\n', (3702, 3717), False, 'from tvm import tir\n'), ((3726, 3757), 'tvm.tir.match_buffer', 'tir.match_buffer', (['c', '(128, 128)'], {}), '(c, (128, 128))\n', (3742, 3757), False, 'from tvm import tir\n'), ((3941, 3972), 'tvm.tir.match_buffer', 'tir.match_buffer', (['a', '(128, 128)'], {}), '(a, (128, 128))\n', (3957, 3972), False, 'from tvm import tir\n'), ((3981, 4009), 'tvm.tir.alloc_buffer', 'tir.alloc_buffer', (['(128, 128)'], {}), '((128, 128))\n', (3997, 4009), False, 'from tvm import tir\n'), ((4018, 4046), 'tvm.tir.alloc_buffer', 'tir.alloc_buffer', (['(128, 128)'], {}), '((128, 128))\n', (4034, 4046), False, 'from tvm import tir\n'), ((4055, 4086), 'tvm.tir.match_buffer', 'tir.match_buffer', (['d', '(128, 128)'], {}), '(d, (128, 128))\n', (4071, 4086), False, 'from tvm import tir\n'), ((4398, 4429), 'tvm.tir.match_buffer', 'tir.match_buffer', (['a', '(128, 128)'], {}), '(a, (128, 128))\n', (4414, 4429), False, 'from tvm import tir\n'), ((4438, 4466), 'tvm.tir.alloc_buffer', 'tir.alloc_buffer', (['(128, 128)'], {}), '((128, 128))\n', (4454, 4466), False, 'from tvm import tir\n'), ((4475, 4506), 'tvm.tir.match_buffer', 'tir.match_buffer', (['c', '(128, 128)'], {}), '(c, (128, 128))\n', (4491, 4506), False, 'from tvm import tir\n'), ((4812, 4843), 'tvm.tir.match_buffer', 'tir.match_buffer', (['a', '(128, 128)'], {}), '(a, (128, 128))\n', (4828, 4843), False, 'from tvm import tir\n'), ((4852, 4883), 'tvm.tir.match_buffer', 'tir.match_buffer', (['c', '(128, 128)'], {}), '(c, (128, 128))\n', (4868, 4883), False, 'from tvm import tir\n'), ((5095, 5126), 'tvm.tir.match_buffer', 'tir.match_buffer', (['a', '(128, 128)'], {}), '(a, (128, 128))\n', (5111, 5126), False, 'from tvm import tir\n'), ((5135, 5163), 'tvm.tir.alloc_buffer', 'tir.alloc_buffer', (['(128, 128)'], {}), '((128, 128))\n', (5151, 5163), False, 'from tvm import tir\n'), ((5172, 5203), 'tvm.tir.match_buffer', 'tir.match_buffer', (['c', '(128, 128)'], {}), '(c, (128, 128))\n', (5188, 5203), False, 'from tvm import tir\n'), ((5565, 5596), 'tvm.tir.match_buffer', 'tir.match_buffer', (['a', '(128, 128)'], {}), '(a, (128, 128))\n', (5581, 5596), False, 'from tvm import tir\n'), ((5605, 5633), 'tvm.tir.alloc_buffer', 'tir.alloc_buffer', (['(128, 128)'], {}), '((128, 128))\n', (5621, 5633), False, 'from tvm import tir\n'), ((5642, 5673), 'tvm.tir.match_buffer', 'tir.match_buffer', (['c', '(128, 128)'], {}), '(c, (128, 128))\n', (5658, 5673), False, 'from tvm import tir\n'), ((6087, 6118), 'tvm.tir.match_buffer', 'tir.match_buffer', (['a', '(128, 128)'], {}), '(a, (128, 128))\n', (6103, 6118), False, 'from tvm import tir\n'), ((6127, 6155), 'tvm.tir.alloc_buffer', 'tir.alloc_buffer', (['(128, 128)'], {}), '((128, 128))\n', (6143, 6155), False, 'from tvm import tir\n'), ((6164, 6195), 'tvm.tir.match_buffer', 'tir.match_buffer', (['c', '(128, 128)'], {}), '(c, (128, 128))\n', (6180, 6195), False, 'from tvm import tir\n'), ((6512, 6543), 'tvm.tir.match_buffer', 'tir.match_buffer', (['a', '(128, 128)'], {}), '(a, (128, 128))\n', (6528, 6543), False, 'from tvm import tir\n'), ((6552, 6580), 'tvm.tir.alloc_buffer', 'tir.alloc_buffer', (['(128, 128)'], {}), '((128, 128))\n', (6568, 6580), False, 'from tvm import tir\n'), ((6589, 6620), 'tvm.tir.match_buffer', 'tir.match_buffer', (['c', '(128, 128)'], {}), '(c, (128, 128))\n', (6605, 6620), False, 'from tvm import tir\n'), ((6722, 6740), 'tvm.tir.grid', 'tir.grid', (['(128)', '(128)'], {}), '(128, 128)\n', (6730, 6740), False, 'from tvm import tir\n'), ((6970, 7001), 'tvm.tir.match_buffer', 'tir.match_buffer', (['a', '(128, 128)'], {}), '(a, (128, 128))\n', (6986, 7001), False, 'from tvm import tir\n'), ((7010, 7041), 'tvm.tir.match_buffer', 'tir.match_buffer', (['c', '(128, 128)'], {}), '(c, (128, 128))\n', (7026, 7041), False, 'from tvm import tir\n'), ((7058, 7076), 'tvm.tir.grid', 'tir.grid', (['(128)', '(128)'], {}), '(128, 128)\n', (7066, 7076), False, 'from tvm import tir\n'), ((7312, 7343), 'tvm.tir.match_buffer', 'tir.match_buffer', (['a', '(128, 128)'], {}), '(a, (128, 128))\n', (7328, 7343), False, 'from tvm import tir\n'), ((7352, 7380), 'tvm.tir.alloc_buffer', 'tir.alloc_buffer', (['(128, 128)'], {}), '((128, 128))\n', (7368, 7380), False, 'from tvm import tir\n'), ((7389, 7420), 'tvm.tir.match_buffer', 'tir.match_buffer', (['c', '(128, 128)'], {}), '(c, (128, 128))\n', (7405, 7420), False, 'from tvm import tir\n'), ((7716, 7747), 'tvm.tir.match_buffer', 'tir.match_buffer', (['a', '(128, 128)'], {}), '(a, (128, 128))\n', (7732, 7747), False, 'from tvm import tir\n'), ((7756, 7787), 'tvm.tir.match_buffer', 'tir.match_buffer', (['c', '(128, 128)'], {}), '(c, (128, 128))\n', (7772, 7787), False, 'from tvm import tir\n'), ((8026, 8068), 'tvm.tir.Schedule', 'tir.Schedule', (['elementwise'], {'debug_mode': '(True)'}), '(elementwise, debug_mode=True)\n', (8038, 8068), False, 'from tvm import tir\n'), ((8171, 8239), 'tvm.ir.assert_structural_equal', 'tvm.ir.assert_structural_equal', (['elementwise_inlined', "sch.mod['main']"], {}), "(elementwise_inlined, sch.mod['main'])\n", (8201, 8239), False, 'import tvm\n'), ((8335, 8388), 'tvm.tir.Schedule', 'tir.Schedule', (['elementwise_under_loop'], {'debug_mode': '(True)'}), '(elementwise_under_loop, debug_mode=True)\n', (8347, 8388), False, 'from tvm import tir\n'), ((8491, 8559), 'tvm.ir.assert_structural_equal', 'tvm.ir.assert_structural_equal', (['elementwise_inlined', "sch.mod['main']"], {}), "(elementwise_inlined, sch.mod['main'])\n", (8521, 8559), False, 'import tvm\n'), ((8651, 8704), 'tvm.tir.Schedule', 'tir.Schedule', (['elementwise_standalone'], {'debug_mode': '(True)'}), '(elementwise_standalone, debug_mode=True)\n', (8663, 8704), False, 'from tvm import tir\n'), ((8807, 8882), 'tvm.ir.assert_structural_equal', 'tvm.ir.assert_structural_equal', (['elementwise_standalone_dce', "sch.mod['main']"], {}), "(elementwise_standalone_dce, sch.mod['main'])\n", (8837, 8882), False, 'import tvm\n'), ((8982, 9048), 'tvm.tir.Schedule', 'tir.Schedule', (['elementwise_multi_producer_consumer'], {'debug_mode': '(True)'}), '(elementwise_multi_producer_consumer, debug_mode=True)\n', (8994, 9048), False, 'from tvm import tir\n'), ((9184, 9272), 'tvm.ir.assert_structural_equal', 'tvm.ir.assert_structural_equal', (['elementwise_multi_consumer_inlined', "sch.mod['main']"], {}), "(elementwise_multi_consumer_inlined, sch.mod[\n    'main'])\n", (9214, 9272), False, 'import tvm\n'), ((9415, 9504), 'tvm.tir.Schedule', 'tir.Schedule', (['fail_multi_reader_writer'], {'debug_mode': '(True)', 'error_render_level': '"""detail"""'}), "(fail_multi_reader_writer, debug_mode=True, error_render_level=\n    'detail')\n", (9427, 9504), False, 'from tvm import tir\n'), ((9675, 9717), 'tvm.tir.Schedule', 'tir.Schedule', (['elementwise'], {'debug_mode': '(True)'}), '(elementwise, debug_mode=True)\n', (9687, 9717), False, 'from tvm import tir\n'), ((9828, 9896), 'tvm.ir.assert_structural_equal', 'tvm.ir.assert_structural_equal', (['elementwise_inlined', "sch.mod['main']"], {}), "(elementwise_inlined, sch.mod['main'])\n", (9858, 9896), False, 'import tvm\n'), ((10000, 10053), 'tvm.tir.Schedule', 'tir.Schedule', (['elementwise_under_loop'], {'debug_mode': '(True)'}), '(elementwise_under_loop, debug_mode=True)\n', (10012, 10053), False, 'from tvm import tir\n'), ((10164, 10232), 'tvm.ir.assert_structural_equal', 'tvm.ir.assert_structural_equal', (['elementwise_inlined', "sch.mod['main']"], {}), "(elementwise_inlined, sch.mod['main'])\n", (10194, 10232), False, 'import tvm\n'), ((10337, 10390), 'tvm.tir.Schedule', 'tir.Schedule', (['elementwise_standalone'], {'debug_mode': '(True)'}), '(elementwise_standalone, debug_mode=True)\n', (10349, 10390), False, 'from tvm import tir\n'), ((10582, 10648), 'tvm.tir.Schedule', 'tir.Schedule', (['elementwise_multi_producer_consumer'], {'debug_mode': '(True)'}), '(elementwise_multi_producer_consumer, debug_mode=True)\n', (10594, 10648), False, 'from tvm import tir\n'), ((10838, 10893), 'tvm.tir.Schedule', 'tir.Schedule', (['fail_multi_reader_writer'], {'debug_mode': '(True)'}), '(fail_multi_reader_writer, debug_mode=True)\n', (10850, 10893), False, 'from tvm import tir\n'), ((11078, 11140), 'tvm.tir.Schedule', 'tir.Schedule', (['elementwise_multi_reverse_loads'], {'debug_mode': '(True)'}), '(elementwise_multi_reverse_loads, debug_mode=True)\n', (11090, 11140), False, 'from tvm import tir\n'), ((11218, 11311), 'tvm.ir.assert_structural_equal', 'tvm.ir.assert_structural_equal', (['elementwise_multi_reverse_loads_inlined', "sch.mod['main']"], {}), "(elementwise_multi_reverse_loads_inlined, sch\n    .mod['main'])\n", (11248, 11311), False, 'import tvm\n'), ((11372, 11426), 'tvm.tir.Schedule', 'tir.Schedule', (['elementwise_multi_loads'], {'debug_mode': '(True)'}), '(elementwise_multi_loads, debug_mode=True)\n', (11384, 11426), False, 'from tvm import tir\n'), ((11594, 11643), 'tvm.tir.Schedule', 'tir.Schedule', (['opaque_access_load'], {'debug_mode': '(True)'}), '(opaque_access_load, debug_mode=True)\n', (11606, 11643), False, 'from tvm import tir\n'), ((11804, 11854), 'tvm.tir.Schedule', 'tir.Schedule', (['opaque_access_store'], {'debug_mode': '(True)'}), '(opaque_access_store, debug_mode=True)\n', (11816, 11854), False, 'from tvm import tir\n'), ((12010, 12055), 'tvm.tir.Schedule', 'tir.Schedule', (['buffer_matched'], {'debug_mode': '(True)'}), '(buffer_matched, debug_mode=True)\n', (12022, 12055), False, 'from tvm import tir\n'), ((12221, 12273), 'tvm.tir.Schedule', 'tir.Schedule', (['elementwise_predicate'], {'debug_mode': '(True)'}), '(elementwise_predicate, debug_mode=True)\n', (12233, 12273), False, 'from tvm import tir\n'), ((12343, 12421), 'tvm.ir.assert_structural_equal', 'tvm.ir.assert_structural_equal', (['elementwise_predicate_inlined', "sch.mod['main']"], {}), "(elementwise_predicate_inlined, sch.mod['main'])\n", (12373, 12421), False, 'import tvm\n'), ((12473, 12527), 'tvm.tir.Schedule', 'tir.Schedule', (['elementwise_multi_loads'], {'debug_mode': '(True)'}), '(elementwise_multi_loads, debug_mode=True)\n', (12485, 12527), False, 'from tvm import tir\n'), ((12597, 12682), 'tvm.ir.assert_structural_equal', 'tvm.ir.assert_structural_equal', (['elementwise_multi_loads_inlined', "sch.mod['main']"], {}), "(elementwise_multi_loads_inlined, sch.mod['main']\n    )\n", (12627, 12682), False, 'import tvm\n'), ((1181, 1207), 'tvm.tir.block', 'tir.block', (['[128, 128]', '"""B"""'], {}), "([128, 128], 'B')\n", (1190, 1207), False, 'from tvm import tir\n'), ((1266, 1292), 'tvm.tir.block', 'tir.block', (['[128, 128]', '"""C"""'], {}), "([128, 128], 'C')\n", (1275, 1292), False, 'from tvm import tir\n'), ((1617, 1643), 'tvm.tir.block', 'tir.block', (['[128, 128]', '"""B"""'], {}), "([128, 128], 'B')\n", (1626, 1643), False, 'from tvm import tir\n'), ((1725, 1751), 'tvm.tir.block', 'tir.block', (['[128, 128]', '"""C"""'], {}), "([128, 128], 'C')\n", (1734, 1751), False, 'from tvm import tir\n'), ((1810, 1836), 'tvm.tir.block', 'tir.block', (['[128, 128]', '"""D"""'], {}), "([128, 128], 'D')\n", (1819, 1836), False, 'from tvm import tir\n'), ((2158, 2184), 'tvm.tir.block', 'tir.block', (['[128, 128]', '"""C"""'], {}), "([128, 128], 'C')\n", (2167, 2184), False, 'from tvm import tir\n'), ((2249, 2275), 'tvm.tir.block', 'tir.block', (['[128, 128]', '"""D"""'], {}), "([128, 128], 'D')\n", (2258, 2275), False, 'from tvm import tir\n'), ((2551, 2577), 'tvm.tir.block', 'tir.block', (['[128, 128]', '"""B"""'], {}), "([128, 128], 'B')\n", (2560, 2577), False, 'from tvm import tir\n'), ((2636, 2662), 'tvm.tir.block', 'tir.block', (['[128, 128]', '"""C"""'], {}), "([128, 128], 'C')\n", (2645, 2662), False, 'from tvm import tir\n'), ((2887, 2913), 'tvm.tir.block', 'tir.block', (['[128, 128]', '"""C"""'], {}), "([128, 128], 'C')\n", (2896, 2913), False, 'from tvm import tir\n'), ((3212, 3230), 'tvm.tir.serial', 'tir.serial', (['(0)', '(128)'], {}), '(0, 128)\n', (3222, 3230), False, 'from tvm import tir\n'), ((3414, 3432), 'tvm.tir.serial', 'tir.serial', (['(0)', '(128)'], {}), '(0, 128)\n', (3424, 3432), False, 'from tvm import tir\n'), ((3767, 3793), 'tvm.tir.block', 'tir.block', (['[128, 128]', '"""C"""'], {}), "([128, 128], 'C')\n", (3776, 3793), False, 'from tvm import tir\n'), ((4096, 4122), 'tvm.tir.block', 'tir.block', (['[128, 128]', '"""B"""'], {}), "([128, 128], 'B')\n", (4105, 4122), False, 'from tvm import tir\n'), ((4217, 4243), 'tvm.tir.block', 'tir.block', (['[128, 128]', '"""C"""'], {}), "([128, 128], 'C')\n", (4226, 4243), False, 'from tvm import tir\n'), ((4516, 4542), 'tvm.tir.block', 'tir.block', (['[128, 128]', '"""B"""'], {}), "([128, 128], 'B')\n", (4525, 4542), False, 'from tvm import tir\n'), ((4601, 4627), 'tvm.tir.block', 'tir.block', (['[128, 128]', '"""C"""'], {}), "([128, 128], 'C')\n", (4610, 4627), False, 'from tvm import tir\n'), ((4893, 4919), 'tvm.tir.block', 'tir.block', (['[128, 128]', '"""B"""'], {}), "([128, 128], 'B')\n", (4902, 4919), False, 'from tvm import tir\n'), ((5213, 5239), 'tvm.tir.block', 'tir.block', (['[128, 128]', '"""B"""'], {}), "([128, 128], 'B')\n", (5222, 5239), False, 'from tvm import tir\n'), ((5298, 5324), 'tvm.tir.block', 'tir.block', (['[128, 128]', '"""C"""'], {}), "([128, 128], 'C')\n", (5307, 5324), False, 'from tvm import tir\n'), ((5346, 5372), 'tvm.tir.reads', 'tir.reads', (['B[0:128, 0:128]'], {}), '(B[0:128, 0:128])\n', (5355, 5372), False, 'from tvm import tir\n'), ((5381, 5408), 'tvm.tir.writes', 'tir.writes', (['C[0:128, 0:128]'], {}), '(C[0:128, 0:128])\n', (5391, 5408), False, 'from tvm import tir\n'), ((5683, 5709), 'tvm.tir.block', 'tir.block', (['[128, 128]', '"""B"""'], {}), "([128, 128], 'B')\n", (5692, 5709), False, 'from tvm import tir\n'), ((5768, 5794), 'tvm.tir.block', 'tir.block', (['[128, 128]', '"""C"""'], {}), "([128, 128], 'C')\n", (5777, 5794), False, 'from tvm import tir\n'), ((5816, 5842), 'tvm.tir.reads', 'tir.reads', (['B[0:128, 0:128]'], {}), '(B[0:128, 0:128])\n', (5825, 5842), False, 'from tvm import tir\n'), ((5851, 5878), 'tvm.tir.writes', 'tir.writes', (['C[0:128, 0:128]'], {}), '(C[0:128, 0:128])\n', (5861, 5878), False, 'from tvm import tir\n'), ((5887, 5936), 'tvm.tir.store', 'tir.store', (['C.data', '(vi * 128 + vj)', '(B[vi, vj] + 1.0)'], {}), '(C.data, vi * 128 + vj, B[vi, vj] + 1.0)\n', (5896, 5936), False, 'from tvm import tir\n'), ((6205, 6231), 'tvm.tir.block', 'tir.block', (['[128, 128]', '"""B"""'], {}), "([128, 128], 'B')\n", (6214, 6231), False, 'from tvm import tir\n'), ((6290, 6316), 'tvm.tir.block', 'tir.block', (['[128, 128]', '"""C"""'], {}), "([128, 128], 'C')\n", (6299, 6316), False, 'from tvm import tir\n'), ((6343, 6387), 'tvm.tir.match_buffer', 'tir.match_buffer', (['B[vi:vi + 1, (vj)]', '(1, 1)'], {}), '(B[vi:vi + 1, (vj)], (1, 1))\n', (6359, 6387), False, 'from tvm import tir\n'), ((6630, 6656), 'tvm.tir.block', 'tir.block', (['[128, 128]', '"""B"""'], {}), "([128, 128], 'B')\n", (6639, 6656), False, 'from tvm import tir\n'), ((7430, 7456), 'tvm.tir.block', 'tir.block', (['[128, 128]', '"""B"""'], {}), "([128, 128], 'B')\n", (7439, 7456), False, 'from tvm import tir\n'), ((7515, 7541), 'tvm.tir.block', 'tir.block', (['[128, 126]', '"""C"""'], {}), "([128, 126], 'C')\n", (7524, 7541), False, 'from tvm import tir\n'), ((7797, 7823), 'tvm.tir.block', 'tir.block', (['[128, 126]', '"""C"""'], {}), "([128, 126], 'C')\n", (7806, 7823), False, 'from tvm import tir\n'), ((9542, 9578), 'pytest.raises', 'pytest.raises', (['tvm.tir.ScheduleError'], {}), '(tvm.tir.ScheduleError)\n', (9555, 9578), False, 'import pytest\n'), ((10433, 10469), 'pytest.raises', 'pytest.raises', (['tvm.tir.ScheduleError'], {}), '(tvm.tir.ScheduleError)\n', (10446, 10469), False, 'import pytest\n'), ((10691, 10727), 'pytest.raises', 'pytest.raises', (['tvm.tir.ScheduleError'], {}), '(tvm.tir.ScheduleError)\n', (10704, 10727), False, 'import pytest\n'), ((10936, 10972), 'pytest.raises', 'pytest.raises', (['tvm.tir.ScheduleError'], {}), '(tvm.tir.ScheduleError)\n', (10949, 10972), False, 'import pytest\n'), ((11469, 11505), 'pytest.raises', 'pytest.raises', (['tvm.tir.ScheduleError'], {}), '(tvm.tir.ScheduleError)\n', (11482, 11505), False, 'import pytest\n'), ((11686, 11722), 'pytest.raises', 'pytest.raises', (['tvm.tir.ScheduleError'], {}), '(tvm.tir.ScheduleError)\n', (11699, 11722), False, 'import pytest\n'), ((11897, 11933), 'pytest.raises', 'pytest.raises', (['tvm.tir.ScheduleError'], {}), '(tvm.tir.ScheduleError)\n', (11910, 11933), False, 'import pytest\n'), ((12098, 12134), 'pytest.raises', 'pytest.raises', (['tvm.tir.ScheduleError'], {}), '(tvm.tir.ScheduleError)\n', (12111, 12134), False, 'import pytest\n'), ((5429, 5471), 'tvm.tir.load', 'tir.load', (['"""float32"""', 'B.data', '(vi * 128 + vj)'], {}), "('float32', B.data, vi * 128 + vj)\n", (5437, 5471), False, 'from tvm import tir\n'), ((5957, 5998), 'tvm.tir.load', 'tir.load', (['"""float32"""', 'B.data', '(vi * 16 + vj)'], {}), "('float32', B.data, vi * 16 + vj)\n", (5965, 5998), False, 'from tvm import tir\n'), ((6755, 6781), 'tvm.tir.block', 'tir.block', (['[128, 128]', '"""C"""'], {}), "([128, 128], 'C')\n", (6764, 6781), False, 'from tvm import tir\n'), ((6807, 6832), 'tvm.tir.where', 'tir.where', (['(B[i, j] < 10.0)'], {}), '(B[i, j] < 10.0)\n', (6816, 6832), False, 'from tvm import tir\n'), ((7091, 7117), 'tvm.tir.block', 'tir.block', (['[128, 128]', '"""C"""'], {}), "([128, 128], 'C')\n", (7100, 7117), False, 'from tvm import tir\n'), ((7143, 7174), 'tvm.tir.where', 'tir.where', (['(A[i, j] * 2.0 < 10.0)'], {}), '(A[i, j] * 2.0 < 10.0)\n', (7152, 7174), False, 'from tvm import tir\n'), ((3249, 3275), 'tvm.tir.block', 'tir.block', (['[128, 128]', '"""B"""'], {}), "([128, 128], 'B')\n", (3258, 3275), False, 'from tvm import tir\n'), ((3305, 3320), 'tvm.tir.bind', 'tir.bind', (['vi', 'i'], {}), '(vi, i)\n', (3313, 3320), False, 'from tvm import tir\n'), ((3337, 3352), 'tvm.tir.bind', 'tir.bind', (['vj', 'j'], {}), '(vj, j)\n', (3345, 3352), False, 'from tvm import tir\n'), ((3451, 3477), 'tvm.tir.block', 'tir.block', (['[128, 128]', '"""C"""'], {}), "([128, 128], 'C')\n", (3460, 3477), False, 'from tvm import tir\n'), ((3507, 3522), 'tvm.tir.bind', 'tir.bind', (['vi', 'i'], {}), '(vi, i)\n', (3515, 3522), False, 'from tvm import tir\n'), ((3539, 3554), 'tvm.tir.bind', 'tir.bind', (['vj', 'j'], {}), '(vj, j)\n', (3547, 3554), False, 'from tvm import tir\n')]