# 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 json
import tvm

@tvm.tag_scope(tag="conv")
def compute_conv(data, weight):
    N, IC, H, W = data.shape
    OC, IC, KH, KW = weight.shape
    OH = H - KH + 1
    OW = W - KW + 1

    ic = tvm.reduce_axis((0, IC), name='ic')
    dh = tvm.reduce_axis((0, KH), name='dh')
    dw = tvm.reduce_axis((0, KW), name='dw')

    return tvm.compute((N, OC, OH, OW), lambda i, oc, h, w: \
        tvm.sum(data[i, ic, h+dh, w+dw] * weight[oc, ic, dh, dw],
                axis=[ic, dh, dw]))

def test_with():
    n = tvm.size_var('n')
    m = tvm.size_var('m')
    l = tvm.size_var('l')

    A = tvm.placeholder((n, l), name='A')
    B = tvm.placeholder((m, l), name='B')
    with tvm.tag_scope(tag="gemm"):
        k = tvm.reduce_axis((0, l), name='k')
        C = tvm.compute((n, m), lambda i, j: tvm.sum(A[i, k] * B[j, k], axis=k),
                        attrs={"hello" : 1, "arr": [10, 12]})

    assert C.op.tag == 'gemm'
    assert "hello" in C.op.attrs
    assert "xx" not in C.op.attrs
    assert C.op.attrs["hello"].value == 1
    CC = tvm.load_json(tvm.save_json(C))
    assert CC.op.attrs["hello"].value == 1
    assert CC.op.attrs["arr"][0].value == 10
    # str format happened to be json compatible
    assert json.loads(str(CC.op.attrs))["arr"][1] == 12


def test_decorator():
    n = tvm.size_var('n')
    c = tvm.size_var('c')
    h = tvm.size_var('h')
    w = tvm.size_var('w')
    kh = tvm.size_var('kh')
    kw = tvm.size_var('kw')

    A = tvm.placeholder((n, c, h, w), name='A')
    B = tvm.placeholder((c, c, kh, kw), name='B')
    C = compute_conv(A, B)
    assert C.op.tag == 'conv'
    assert len(C.op.attrs) == 0

def test_nested():
    n = tvm.size_var('n')
    c = tvm.size_var('c')
    h = tvm.size_var('h')
    w = tvm.size_var('w')
    kh = tvm.size_var('kh')
    kw = tvm.size_var('kw')

    A = tvm.placeholder((n, c, h, w), name='A')
    B = tvm.placeholder((c, c, kh, kw), name='B')
    try:
        with tvm.tag_scope(tag='conv'):
            C = compute_conv(A, B)
        assert False
    except ValueError:
        pass


if __name__ == "__main__":
    test_with()
    test_decorator()
    test_nested()

[((810, 835), 'tvm.tag_scope', 'tvm.tag_scope', ([], {'tag': '"""conv"""'}), "(tag='conv')\n", (823, 835), False, 'import tvm\n'), ((981, 1016), 'tvm.reduce_axis', 'tvm.reduce_axis', (['(0, IC)'], {'name': '"""ic"""'}), "((0, IC), name='ic')\n", (996, 1016), False, 'import tvm\n'), ((1026, 1061), 'tvm.reduce_axis', 'tvm.reduce_axis', (['(0, KH)'], {'name': '"""dh"""'}), "((0, KH), name='dh')\n", (1041, 1061), False, 'import tvm\n'), ((1071, 1106), 'tvm.reduce_axis', 'tvm.reduce_axis', (['(0, KW)'], {'name': '"""dw"""'}), "((0, KW), name='dw')\n", (1086, 1106), False, 'import tvm\n'), ((1298, 1315), 'tvm.size_var', 'tvm.size_var', (['"""n"""'], {}), "('n')\n", (1310, 1315), False, 'import tvm\n'), ((1324, 1341), 'tvm.size_var', 'tvm.size_var', (['"""m"""'], {}), "('m')\n", (1336, 1341), False, 'import tvm\n'), ((1350, 1367), 'tvm.size_var', 'tvm.size_var', (['"""l"""'], {}), "('l')\n", (1362, 1367), False, 'import tvm\n'), ((1377, 1410), 'tvm.placeholder', 'tvm.placeholder', (['(n, l)'], {'name': '"""A"""'}), "((n, l), name='A')\n", (1392, 1410), False, 'import tvm\n'), ((1419, 1452), 'tvm.placeholder', 'tvm.placeholder', (['(m, l)'], {'name': '"""B"""'}), "((m, l), name='B')\n", (1434, 1452), False, 'import tvm\n'), ((2083, 2100), 'tvm.size_var', 'tvm.size_var', (['"""n"""'], {}), "('n')\n", (2095, 2100), False, 'import tvm\n'), ((2109, 2126), 'tvm.size_var', 'tvm.size_var', (['"""c"""'], {}), "('c')\n", (2121, 2126), False, 'import tvm\n'), ((2135, 2152), 'tvm.size_var', 'tvm.size_var', (['"""h"""'], {}), "('h')\n", (2147, 2152), False, 'import tvm\n'), ((2161, 2178), 'tvm.size_var', 'tvm.size_var', (['"""w"""'], {}), "('w')\n", (2173, 2178), False, 'import tvm\n'), ((2188, 2206), 'tvm.size_var', 'tvm.size_var', (['"""kh"""'], {}), "('kh')\n", (2200, 2206), False, 'import tvm\n'), ((2216, 2234), 'tvm.size_var', 'tvm.size_var', (['"""kw"""'], {}), "('kw')\n", (2228, 2234), False, 'import tvm\n'), ((2244, 2283), 'tvm.placeholder', 'tvm.placeholder', (['(n, c, h, w)'], {'name': '"""A"""'}), "((n, c, h, w), name='A')\n", (2259, 2283), False, 'import tvm\n'), ((2292, 2333), 'tvm.placeholder', 'tvm.placeholder', (['(c, c, kh, kw)'], {'name': '"""B"""'}), "((c, c, kh, kw), name='B')\n", (2307, 2333), False, 'import tvm\n'), ((2451, 2468), 'tvm.size_var', 'tvm.size_var', (['"""n"""'], {}), "('n')\n", (2463, 2468), False, 'import tvm\n'), ((2477, 2494), 'tvm.size_var', 'tvm.size_var', (['"""c"""'], {}), "('c')\n", (2489, 2494), False, 'import tvm\n'), ((2503, 2520), 'tvm.size_var', 'tvm.size_var', (['"""h"""'], {}), "('h')\n", (2515, 2520), False, 'import tvm\n'), ((2529, 2546), 'tvm.size_var', 'tvm.size_var', (['"""w"""'], {}), "('w')\n", (2541, 2546), False, 'import tvm\n'), ((2556, 2574), 'tvm.size_var', 'tvm.size_var', (['"""kh"""'], {}), "('kh')\n", (2568, 2574), False, 'import tvm\n'), ((2584, 2602), 'tvm.size_var', 'tvm.size_var', (['"""kw"""'], {}), "('kw')\n", (2596, 2602), False, 'import tvm\n'), ((2612, 2651), 'tvm.placeholder', 'tvm.placeholder', (['(n, c, h, w)'], {'name': '"""A"""'}), "((n, c, h, w), name='A')\n", (2627, 2651), False, 'import tvm\n'), ((2660, 2701), 'tvm.placeholder', 'tvm.placeholder', (['(c, c, kh, kw)'], {'name': '"""B"""'}), "((c, c, kh, kw), name='B')\n", (2675, 2701), False, 'import tvm\n'), ((1462, 1487), 'tvm.tag_scope', 'tvm.tag_scope', ([], {'tag': '"""gemm"""'}), "(tag='gemm')\n", (1475, 1487), False, 'import tvm\n'), ((1501, 1534), 'tvm.reduce_axis', 'tvm.reduce_axis', (['(0, l)'], {'name': '"""k"""'}), "((0, l), name='k')\n", (1516, 1534), False, 'import tvm\n'), ((1841, 1857), 'tvm.save_json', 'tvm.save_json', (['C'], {}), '(C)\n', (1854, 1857), False, 'import tvm\n'), ((1178, 1263), 'tvm.sum', 'tvm.sum', (['(data[i, ic, h + dh, w + dw] * weight[oc, ic, dh, dw])'], {'axis': '[ic, dh, dw]'}), '(data[i, ic, h + dh, w + dw] * weight[oc, ic, dh, dw], axis=[ic, dh, dw]\n    )\n', (1185, 1263), False, 'import tvm\n'), ((2724, 2749), 'tvm.tag_scope', 'tvm.tag_scope', ([], {'tag': '"""conv"""'}), "(tag='conv')\n", (2737, 2749), False, 'import tvm\n'), ((1580, 1614), 'tvm.sum', 'tvm.sum', (['(A[i, k] * B[j, k])'], {'axis': 'k'}), '(A[i, k] * B[j, k], axis=k)\n', (1587, 1614), False, 'import tvm\n')]