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,113 @@
+# 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.
+"""Testing utility functions in meta schedule"""
+
+from collections.abc import Callable
+
+import numpy as np  # type: ignore
+
+import tvm
+from tvm.runtime import Tensor
+
+
+def generate_input_data(
+    input_shape: list[int],
+    input_dtype: str,
+    *,
+    low: int | None = None,
+    high: int | None = None,
+) -> np.ndarray:
+    """Generate input date with given shape and data type.
+
+    Parameters
+    ----------
+    input_shape : List[int]
+        The shape of the input data.
+    input_dtype : str
+        The data type of the input date.
+
+    Returns
+    -------
+    input_data : np.ndarray
+        The generated input data with given shape and data type in numpy ndarray.
+    """
+    if input_dtype.startswith("float"):
+        return np.random.uniform(size=input_shape).astype(input_dtype)
+    range_map = {
+        "uint8": (0, 255),
+        "int8": (-128, 127),
+        "int32": (0, 10000),
+        "uint32": (0, 10000),
+        "int64": (0, 10000),
+        "uint64": (0, 10000),
+    }
+    if input_dtype in range_map:
+        _low, _high = range_map[input_dtype]
+        return np.random.randint(
+            low=_low if low is None else low,
+            high=_high if high is None else high,
+            size=input_shape,
+            dtype=input_dtype,
+        )
+    raise ValueError("Unsupported input datatype!")
+
+
+def create_calculator(backend: str) -> Callable:
+    """Create a function to fetch the computing result of running the given runtime module.
+
+    Parameters
+    ----------
+    backend : str
+        The backend to use, only tirx is supported for now.
+
+    Returns
+    -------
+    func : Callable
+        The function to fetch the computing result.
+    """
+
+    def f_calculator(
+        rt_mod: tvm.runtime.Module,
+        dev: tvm.runtime.Device,  # pylint: disable=unused-argument
+        input_data: dict[str, Tensor],
+    ) -> list[Tensor]:
+        """Fetch the result of running the given runtime module.
+
+        Parameters
+        ----------
+        rt_mod : tvm.runtime.Module
+            The runtime module.
+        dev : tvm.device
+            The device type to run workload.
+        input_data : Dict[str, np.ndarray]
+            The input data as a dictionary.
+        """
+        try:
+            if backend == "tirx":
+                data = [v for _, v in sorted(input_data.items(), key=lambda x: x[0])]
+                rt_mod(*data)
+                return data
+            else:
+                raise ValueError(f"Backend {backend} not supported in f_calculator!")
+
+        except Exception as exc:  # pylint: disable=broad-except
+            print(
+                f"Run module f_calculator via RPC failed, exception: {exc}",
+            )
+            return None
+
+    return f_calculator