test_pass_remove_unused_functions.py 3.95 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
# 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.
17
import pytest
18 19 20 21 22
import tvm
from tvm import relay
from tvm.relay import transform
from tvm.relay.prelude import Prelude

23

24 25 26 27 28 29 30 31 32
def test_remove_all_prelude_functions():
    mod = relay.Module()
    p = Prelude(mod)
    x = relay.var("x", shape=(1, 16))
    mod["main"] = relay.Function([x], x)
    mod = relay.transform.RemoveUnusedFunctions()(mod)
    l = set([x[0].name_hint for x in mod.functions.items()])
    assert l == set(['main'])

33

34 35 36 37 38 39 40 41 42 43 44 45 46
def test_remove_all_prelude_functions_but_referenced_functions():
    mod = relay.Module()
    p = Prelude(mod)
    x = relay.var("x", shape=(1, 16))
    id_func = relay.Function([x], x)
    id_name = relay.GlobalVar('id_func')
    mod[id_name] = id_func

    mod["main"] = relay.Function([x], id_name(x))
    mod = relay.transform.RemoveUnusedFunctions()(mod)
    l = set([x[0].name_hint for x in mod.functions.items()])
    assert l == set(['id_func', 'main'])

47

48 49 50 51 52 53 54 55 56 57 58 59
def test_keep_only_referenced_prelude_functions():
    mod = relay.Module()
    p = Prelude(mod)
    l = p.nil()
    for i in [4, 3, 2, 1, 0]:
        l = p.cons(relay.const(i), l)
    body = p.hd(p.tl(p.tl(l)))
    mod["main"] = relay.Function([], body)
    mod = relay.transform.RemoveUnusedFunctions()(mod)
    l = set([x[0].name_hint for x in mod.functions.items()])
    assert l == set(['tl', 'hd', 'main'])

60

61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78
def test_multiple_entry_functions():
    mod = relay.Module()
    p = Prelude(mod)
    l = p.nil()
    for i in [4, 3, 2, 1, 0]:
        l = p.cons(relay.const(i), l)
    body = p.hd(p.tl(p.tl(l)))
    mod["main1"] = relay.Function([], body)

    x = relay.var("x", shape=(1, 16))
    id_func = relay.Function([x], x)
    id_name = relay.GlobalVar('id_func')
    mod[id_name] = id_func
    mod["main2"] = relay.Function([x], id_name(x))
    mod = relay.transform.RemoveUnusedFunctions(['main1', 'main2'])(mod)
    l = set([x[0].name_hint for x in mod.functions.items()])
    assert l == set(['tl', 'hd', 'main2', 'id_func', 'main1'])

79

80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95
def test_globalvar_as_call_arg():
    mod = relay.Module()
    p = Prelude(mod)
    tensor_array = p.get_var('tensor_array', 'int32')
    tensor1 = p.get_var('tensor1', 'int32')
    write = p.get_var('tensor_array_write', 'int32')
    stack = p.get_var('tensor_array_stack', 'int32')
    v = relay.var('v')
    init_tensor_array = tensor_array(relay.const(3))
    tensor_array1 = write(init_tensor_array, relay.const(0), tensor1(v))
    tensor_array2 = stack(tensor_array1)
    mod["main"] = relay.Function([v], tensor_array2)
    mod = relay.transform.RemoveUnusedFunctions()(mod)
    l = set([x[0].name_hint for x in mod.functions.items()])
    assert 'tensor_array_int32' in l

96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114

def test_call_globalvar_without_args():
    def get_mod():
        mod = relay.Module({})
        fn1 = relay.Function([], relay.const(1))
        fn2 = relay.Function([], relay.const(2))
        g1 = relay.GlobalVar('g1')
        g2 = relay.GlobalVar('g2')
        mod[g1] = fn1
        mod[g2] = fn2
        p = relay.var('p', 'bool')
        mod['main'] = relay.Function([p], relay.Call(relay.If(p, g1, g2), []))
        return mod
    mod = get_mod()
    ref_mod = get_mod()
    mod = relay.transform.RemoveUnusedFunctions()(mod)
    assert relay.alpha_equal(mod, ref_mod)


115 116
if __name__ == '__main__':
    pytest.main()