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,133 @@
+# 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=unused-argument
+"""Tools to compare libraries."""
+
+from collections.abc import Iterable
+
+import tvm
+import tvm.testing
+
+
+class LibCompareVMInstrument:
+    """Instrument class to compare libs.
+
+    This class build an instrument function that
+    pair tests an existing compiled relax vm implementation
+    and an extra module, which can sits in another backend
+    but offers a same subset of compiled TIR functions.
+
+    The instrumentation enables us to automatically
+    check and compare each ops being called in the pipeline
+    by looking up the same name in the provided mod and run testing.
+
+    Parameters
+    ----------
+    mod: runtime.Module
+        The module of interest to be validated.
+
+    device: runtime.Device
+        The device to run the target module on.
+
+    verbose: bool
+        Whether print out messages.
+
+    rtol: float
+        rtol used in validation
+
+    atol: float
+        atol used in validation
+    """
+
+    def __init__(self, mod, device, verbose=True, rtol=1e-5, atol=1e-5):
+        self.mod = mod
+        self.device = device
+        self.verbose = verbose
+        self.counter = 0
+        self.rtol = rtol
+        self.atol = atol
+
+    def compare(
+        self,
+        name: str,
+        ref_args: list[tvm.runtime.Tensor] | tuple[tvm.runtime.Tensor, ...],
+        new_args: list[tvm.runtime.Tensor] | tuple[tvm.runtime.Tensor, ...],
+        ret_indices: Iterable[int],
+    ):
+        """Comparison function, can be overloaded.
+
+        Parameters
+        ----------
+        name: str
+            Name of the function.
+
+        ref_args:
+            The reference arguments.
+
+        new_args:
+            The args to be passed to the comparison function.
+
+        ret_indices:
+            List of indices to validate return values.
+        """
+        my_func = self.mod.get_function(name, query_imports=True)
+        if self.verbose:
+            print(f"[{self.counter}] Validating {name} ...")
+        my_func(*new_args)
+        for rindex in ret_indices:
+            tvm.testing.assert_allclose(
+                new_args[rindex].numpy(), ref_args[rindex].numpy(), atol=self.atol, rtol=self.rtol
+            )
+        if self.verbose:
+            print(f"[{self.counter}] Validating {name}, passed.")
+        self.counter += 1
+
+    def skip_instrument(self, func, name, before_run, ret_val, *args):
+        return False
+
+    def __call__(self, func, name, before_run, ret_val, *args):
+        if before_run:
+            return
+        if name.startswith("vm.builtin."):
+            return
+        if any(not isinstance(x, tvm.runtime.Tensor) for x in args):
+            return
+        try:
+            self.mod.get_function(name, query_imports=True)
+        except AttributeError:
+            if self.verbose:
+                print(f"Cannot find {name}, skip...")
+            return
+
+        if self.skip_instrument(func, name, before_run, ret_val, *args):
+            return
+
+        new_args = []
+        # not always true, true for most ops.
+        ret_indices = (len(args) - 1,)
+        temp_args = []
+        for i, arg in enumerate(args):
+            arr = tvm.runtime.empty(arg.shape, arg.dtype, device=self.device)
+            # copy from cpu since we look at different device
+            if i not in ret_indices:
+                temp_cpu = arg.copyto(tvm.cpu())
+                temp_args.append(temp_cpu)
+                arr.copyfrom(temp_cpu)
+            new_args.append(arr)
+        # wait until all copy complete before we release temp_cpu
+        self.device.sync()
+        self.compare(name, args, new_args, ret_indices)