# 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 tvm from tvm import te import numpy as np from tvm import relay from tvm.relay.transform import memory_alloc def check_vm_alloc(func, check_fn): mod = tvm.IRModule() mod['main'] = func ex = relay.create_executor('vm', mod) args = [] for param in func.params: param = param.type_annotation sh = [int(sh) for sh in param.shape] data = np.random.rand(*sh).astype(param.dtype) args.append(tvm.nd.array(data)) result = ex.evaluate(mod['main'])(*args) py_res = check_fn(*[arg.asnumpy() for arg in args]) np.testing.assert_allclose(result.asnumpy(), py_res) def storage_type(mod): return relay.TypeCall(mod.get_global_type_var("Storage"), []) def test_tyck_alloc_storage(): mod = tvm.IRModule() mod.import_from_std("core.rly") def test_tyck_alloc_tensor(): mod = tvm.IRModule() mod.import_from_std("core.rly") sto = relay.Var("x", storage_type(mod)) sh = relay.const(np.array([1, 2]), dtype="int64") at = relay.op.memory.alloc_tensor(sto, sh) mod['main'] = relay.Function([sto], at) relay.transform.InferType()(mod) def check_add(x): return x + x def test_add(): x = relay.var('x', shape=(2,)) z = x + x func = relay.Function([x,], z) check_vm_alloc(func, check_add) def check_add_sub(x, y): z = x + x return z - y def test_add_sub(): x = relay.var('x', shape=(10,)) y = relay.var('y', shape=(10,)) z = x + x z = z - y func = relay.Function([x, y], z) check_vm_alloc(func, check_add_sub) if __name__ == "__main__": test_tyck_alloc_tensor() test_add() test_add_sub()