chore: import upstream snapshot with attribution

2026-07-13 13:36:25 +08:00
commit 26446540fa
3151 changed files with 974126 additions and 0 deletions
@@ -0,0 +1,103 @@
+# 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.
+# ruff: noqa: RUF005
+import numpy as np
+
+import tvm
+import tvm.testing
+from tvm import relax
+from tvm.relax.testing import nn
+from tvm.relax.testing.lib_comparator import LibCompareVMInstrument
+
+
+def get_exec(data_shape):
+    builder = relax.BlockBuilder()
+    weight1_np = np.random.randn(64, 64).astype("float32")
+    weight2_np = np.random.randn(64, 64).astype("float32")
+
+    with builder.function("main"):
+        model = nn.Sequential(
+            nn.Linear(data_shape[1], weight1_np.shape[0], bias=False),
+            nn.ReLU(),
+            nn.Linear(weight2_np.shape[0], weight2_np.shape[1], bias=False),
+            nn.ReLU(),
+        )
+        data = nn.Placeholder(data_shape, name="data")
+        output = model(data)
+        params = [data] + model.parameters()
+        builder.emit_func_output(output, params=params)
+
+    mod = builder.get()
+
+    params = {"linear_weight": weight1_np, "linear_weight1": weight2_np}
+    mod = relax.transform.BindParams("main", params)(mod)
+
+    target = "llvm"
+    return tvm.compile(mod, target)
+
+
+def get_exec_int32(data_shape):
+    builder = relax.BlockBuilder()
+
+    with builder.function("main"):
+        model = nn.ReLU()
+        data = nn.Placeholder(data_shape, dtype="int32", name="data")
+        output = model(data)
+        params = [data] + model.parameters()
+        builder.emit_func_output(output, params=params)
+
+    mod = builder.get()
+    target = "llvm"
+    return tvm.compile(mod, target)
+
+
+def test_conv2d_cpu():
+    data_np = np.random.randn(1, 64).astype("float32")
+    ex = get_exec(data_np.shape)
+    vm = relax.VirtualMachine(ex, tvm.cpu())
+    hit_count = {}
+
+    def instrument(func, name, before_run, ret_val, *args):
+        if (name, before_run) not in hit_count:
+            hit_count[(name, before_run)] = 0
+        hit_count[(name, before_run)] += 1
+        assert callable(func)
+        if before_run:
+            assert ret_val is None
+        if name == "matmul":
+            return relax.VMInstrumentReturnKind.SKIP_RUN
+
+    vm.set_instrument(instrument)
+    vm["main"](tvm.runtime.tensor(data_np))
+    assert hit_count[("matmul", True)] == 2
+    assert ("matmul", False) not in hit_count
+    assert hit_count[("relu", True)] == 2
+    assert hit_count[("relu", False)] == 2
+
+
+def test_lib_comparator():
+    data_np = np.random.randn(1, 64).astype("int32")
+    ex = get_exec_int32(data_np.shape)
+    vm = relax.VirtualMachine(ex, tvm.cpu())
+    # compare against library module
+    cmp = LibCompareVMInstrument(vm.module.imports[0], tvm.cpu(), verbose=False)
+    vm.set_instrument(cmp)
+    vm["main"](tvm.runtime.tensor(data_np))
+
+
+if __name__ == "__main__":
+    tvm.testing.main()