chore: import upstream snapshot with attribution

python/tvm/relax/testing/vm.py

@@ -0,0 +1,92 @@
+# 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=invalid-name
+"""Testing utilities for relax VM"""
+
+from typing import Any
+
+import numpy as np  # type: ignore
+
+import tvm
+from tvm import relax
+from tvm.runtime import Object
+
+
+@tvm.register_global_func("test.vm.move")
+def move(src):
+    return src
+
+
+@tvm.register_global_func("test.vm.add")
+def add(a, b):
+    ret = a.numpy() + b.numpy()
+    return tvm.runtime.tensor(ret)
+
+
+@tvm.register_global_func("test.vm.mul")
+def mul(a, b):
+    ret = a.numpy() * b.numpy()
+    return tvm.runtime.tensor(ret)
+
+
+@tvm.register_global_func("test.vm.equal_zero")
+def equal_zero(a):
+    ret = np.all(a.numpy() == 0)
+    return tvm.runtime.tensor(ret)
+
+
+@tvm.register_global_func("test.vm.subtract_one")
+def subtract_one(a):
+    ret = np.subtract(a.numpy(), 1)
+    return tvm.runtime.tensor(ret)
+
+
+@tvm.register_global_func("test.vm.identity")
+def identity_packed(a, b):
+    b[:] = tvm.runtime.tensor(a.numpy())
+
+
+@tvm.register_global_func("test.vm.tile")
+def tile_packed(a, b):
+    b[:] = tvm.runtime.tensor(np.tile(a.numpy(), (1, 2)))
+
+
+@tvm.register_global_func("test.vm.add_scalar")
+def add_scalar(a, b):
+    return a + b
+
+
+@tvm.register_global_func("test.vm.get_device_id")
+def get_device_id(device):
+    return device.index
+
+
+def check_saved_func(vm: relax.VirtualMachine, func_name: str, *inputs: list[Any]) -> Object:
+    # uses save_function to create a closure with the given inputs
+    # and ensure the result is the same
+    # (assumes the functions return tensors and that they're idempotent)
+    saved_name = f"{func_name}_saved"
+    vm.save_function(func_name, saved_name, *inputs)
+    res1 = vm[func_name](*inputs)
+    res2 = vm[saved_name]()
+    tvm.testing.assert_allclose(res1.numpy(), res2.numpy(), rtol=1e-7, atol=1e-7)
+    return res1
+
+
+@tvm.register_global_func("test.vm.check_if_defined")
+def check_if_defined(obj: tvm.Object) -> tvm.tirx.IntImm:
+    return tvm.runtime.convert(obj is not None)