wehub-resource-sync
108 lines
4.0 KiB
Python
108 lines
4.0 KiB
Python
# 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
|