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,107 @@
+# 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=dangerous-default-value
+"""Testing utilities for the TensorIR schedule API"""
+
+from collections.abc import Sequence
+from typing import Any
+
+import tvm
+from tvm.ir import IRModule, assert_structural_equal
+from tvm.s_tir.schedule import Schedule, Trace
+from tvm.tirx import PrimFunc
+
+
+def assert_structural_equal_ignore_global_symbol(
+    func1: PrimFunc,
+    func2: PrimFunc,
+    *args: Any,
+    **kwargs: Any,
+) -> None:
+    """
+    Asserts that PrimFuncs func1 and func2 are structurally equal, setting both
+    their global symbol attributes to main so that the global symbol
+    will not be a point of comparison.
+    """
+    assert_structural_equal(
+        func1.with_attr("global_symbol", "main"),
+        func2.with_attr("global_symbol", "main"),
+        *args,
+        **kwargs,
+    )
+
+
+def verify_trace_roundtrip(
+    sch: Schedule,
+    mod: PrimFunc | IRModule,
+    *,
+    debug_mask: str | int = "all",
+    text_format: str | Sequence[str] = ["python", "json"],
+) -> Schedule:
+    """Serialize a traced schedule to JSON, then replay the JSON trace by applying to
+    a fresh new schedule, verifying the reproducibility of scheduling.
+
+    Parameters
+    ----------
+    sch : s_tir.Schedule
+        The traced TensorIR schedule to be verified
+    mod : Union[PrimFunc, IRModule]
+        The IRModule or PrimFunc to construct the fresh new schedule
+    debug_mask : Union[str, int]
+        Do extra correctness checking after the class creation and each time
+        after calling the Replace method.
+        Possible choices of `debug_mask`:
+        1) "all" - Turn on all the checks
+        2) "none" - Turn off all the checks
+        3) An integer - Turn on checks according to the bitmasks provided in ScheduleDebugMask
+    text_format: Union[str, Sequence[str]]
+        The text format or formats whose round-trip behavior should be
+        validated.  If a single string, validate round-trips through
+    """
+    from tvm.script import tirx as T  # pylint: disable=import-outside-toplevel
+
+    if not isinstance(text_format, str):
+        for opt in text_format:
+            new_sch = verify_trace_roundtrip(sch, mod, debug_mask=debug_mask, text_format=opt)
+        return new_sch
+
+    trace = sch.trace
+    assert trace is not None
+
+    # Step 1. Perform a round-trip through the text-format
+    new_sch = Schedule(mod=mod, debug_mask=debug_mask)
+    if text_format == "json":
+        json_obj = trace.as_json()
+        Trace.apply_json_to_schedule(json_obj=json_obj, sch=new_sch)
+    elif text_format == "python":
+        py_trace = "\n".join(trace.as_python())
+        vars_dict = {"T": T}
+        vars_dict.update(tvm.tirx.__dict__)
+        exec(py_trace, vars_dict, {"sch": new_sch})  # pylint: disable=exec-used
+    else:
+        assert text_format in ("json", "python"), f"Unknown text format: {text_format}"
+
+    # Step 2. Verify that the round-trip produced the same scheduling
+    assert_structural_equal(new_sch.mod, sch.mod)
+
+    # Step 3. Check the consistency of the text format between the old and new traces
+    py_repr = "\n".join(trace.as_python())
+    new_py_repr = "\n".join(new_sch.trace.as_python())
+    assert py_repr == new_py_repr
+
+    # Step 4. Return the new schedule in case it could be useful
+    return new_sch