# 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()

[((20, 1, 20, 26), 'tvm.tag_scope', 'tvm.tag_scope', (), '', False, 'import tvm\n'), ((27, 9, 27, 44), 'tvm.reduce_axis', 'tvm.reduce_axis', (), '', False, 'import tvm\n'), ((28, 9, 28, 44), 'tvm.reduce_axis', 'tvm.reduce_axis', (), '', False, 'import tvm\n'), ((29, 9, 29, 44), 'tvm.reduce_axis', 'tvm.reduce_axis', (), '', False, 'import tvm\n'), ((36, 8, 36, 25), 'tvm.size_var', 'tvm.size_var', ({(36, 21, 36, 24): '"""n"""'}, {}), "('n')", False, 'import tvm\n'), ((37, 8, 37, 25), 'tvm.size_var', 'tvm.size_var', ({(37, 21, 37, 24): '"""m"""'}, {}), "('m')", False, 'import tvm\n'), ((38, 8, 38, 25), 'tvm.size_var', 'tvm.size_var', ({(38, 21, 38, 24): '"""l"""'}, {}), "('l')", False, 'import tvm\n'), ((40, 8, 40, 41), 'tvm.placeholder', 'tvm.placeholder', (), '', False, 'import tvm\n'), ((41, 8, 41, 41), 'tvm.placeholder', 'tvm.placeholder', (), '', False, 'import tvm\n'), ((59, 8, 59, 25), 'tvm.size_var', 'tvm.size_var', ({(59, 21, 59, 24): '"""n"""'}, {}), "('n')", False, 'import tvm\n'), ((60, 8, 60, 25), 'tvm.size_var', 'tvm.size_var', ({(60, 21, 60, 24): '"""c"""'}, {}), "('c')", False, 'import tvm\n'), ((61, 8, 61, 25), 'tvm.size_var', 'tvm.size_var', ({(61, 21, 61, 24): '"""h"""'}, {}), "('h')", False, 'import tvm\n'), ((62, 8, 62, 25), 'tvm.size_var', 'tvm.size_var', ({(62, 21, 62, 24): '"""w"""'}, {}), "('w')", False, 'import tvm\n'), ((63, 9, 63, 27), 'tvm.size_var', 'tvm.size_var', ({(63, 22, 63, 26): '"""kh"""'}, {}), "('kh')", False, 'import tvm\n'), ((64, 9, 64, 27), 'tvm.size_var', 'tvm.size_var', ({(64, 22, 64, 26): '"""kw"""'}, {}), "('kw')", False, 'import tvm\n'), ((66, 8, 66, 47), 'tvm.placeholder', 'tvm.placeholder', (), '', False, 'import tvm\n'), ((67, 8, 67, 49), 'tvm.placeholder', 'tvm.placeholder', (), '', False, 'import tvm\n'), ((73, 8, 73, 25), 'tvm.size_var', 'tvm.size_var', ({(73, 21, 73, 24): '"""n"""'}, {}), "('n')", False, 'import tvm\n'), ((74, 8, 74, 25), 'tvm.size_var', 'tvm.size_var', ({(74, 21, 74, 24): '"""c"""'}, {}), "('c')", False, 'import tvm\n'), ((75, 8, 75, 25), 'tvm.size_var', 'tvm.size_var', ({(75, 21, 75, 24): '"""h"""'}, {}), "('h')", False, 'import tvm\n'), ((76, 8, 76, 25), 'tvm.size_var', 'tvm.size_var', ({(76, 21, 76, 24): '"""w"""'}, {}), "('w')", False, 'import tvm\n'), ((77, 9, 77, 27), 'tvm.size_var', 'tvm.size_var', ({(77, 22, 77, 26): '"""kh"""'}, {}), "('kh')", False, 'import tvm\n'), ((78, 9, 78, 27), 'tvm.size_var', 'tvm.size_var', ({(78, 22, 78, 26): '"""kw"""'}, {}), "('kw')", False, 'import tvm\n'), ((80, 8, 80, 47), 'tvm.placeholder', 'tvm.placeholder', (), '', False, 'import tvm\n'), ((81, 8, 81, 49), 'tvm.placeholder', 'tvm.placeholder', (), '', False, 'import tvm\n'), ((42, 9, 42, 34), 'tvm.tag_scope', 'tvm.tag_scope', (), '', False, 'import tvm\n'), ((43, 12, 43, 45), 'tvm.reduce_axis', 'tvm.reduce_axis', (), '', False, 'import tvm\n'), ((51, 23, 51, 39), 'tvm.save_json', 'tvm.save_json', ({(51, 37, 51, 38): 'C'}, {}), '(C)', False, 'import tvm\n'), ((32, 8, 33, 34), 'tvm.sum', 'tvm.sum', (), '', False, 'import tvm\n'), ((83, 13, 83, 38), 'tvm.tag_scope', 'tvm.tag_scope', (), '', False, 'import tvm\n'), ((44, 45, 44, 79), 'tvm.sum', 'tvm.sum', (), '', False, 'import tvm\n')]