wehub-resource-sync
151 lines
4.9 KiB
Python
151 lines
4.9 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=missing-module-docstring,missing-function-docstring,missing-class-docstring
|
|
|
|
# isort: off
|
|
from typing import Literal
|
|
|
|
# isort: on
|
|
|
|
import tvm_ffi
|
|
|
|
from tvm.ir import IRModule
|
|
from tvm.s_tir import Schedule
|
|
from tvm.s_tir import meta_schedule as ms
|
|
from tvm.s_tir.schedule import Trace
|
|
from tvm.s_tir.schedule.testing import verify_trace_roundtrip
|
|
from tvm.target import Target
|
|
|
|
|
|
def get_rules(
|
|
kind: Literal["llvm", "cuda", "cuda-tensorcore", "hexagon"],
|
|
types: type | tuple[type, ...],
|
|
) -> list[ms.ScheduleRule]:
|
|
"""Get default schedule rules"""
|
|
rules = ms.ScheduleRule.create(kind)
|
|
return [rule for rule in rules if isinstance(rule, types)]
|
|
|
|
|
|
def structural_equal_no_gs(mod1: IRModule, mod2: IRModule) -> bool:
|
|
"""
|
|
Checks structural equality but ignores global symbols
|
|
"""
|
|
|
|
# for every function in the modules, remove global symbols from the attrs and then compare
|
|
def remove_global_symbols(mod: IRModule) -> IRModule:
|
|
stripped_mod = IRModule()
|
|
for global_var in mod.get_global_vars():
|
|
func = mod[global_var]
|
|
stripped_mod[global_var] = func.without_attr("global_symbol")
|
|
return stripped_mod
|
|
|
|
return tvm_ffi.structural_equal(remove_global_symbols(mod1), remove_global_symbols(mod2))
|
|
|
|
|
|
def generate_design_space(
|
|
kind: Literal["llvm", "cuda", "cuda-tensorcore", "hexagon"],
|
|
mod: IRModule,
|
|
target: Target,
|
|
types: type | tuple[type, ...],
|
|
sch_rules: list[ms.ScheduleRule] | None = None,
|
|
) -> list[Schedule]:
|
|
if sch_rules is None:
|
|
sch_rules = get_rules(kind, types)
|
|
else:
|
|
assert types is None
|
|
return ms.TuneContext(
|
|
mod=mod,
|
|
target=target,
|
|
space_generator=ms.space_generator.PostOrderApply(
|
|
sch_rules=sch_rules,
|
|
postprocs=[],
|
|
mutator_probs={},
|
|
),
|
|
task_name="test",
|
|
).generate_design_space()
|
|
|
|
|
|
def _find_match_sketch_id(
|
|
mod: IRModule,
|
|
sketches: list[Schedule],
|
|
expected_mod: IRModule,
|
|
expected_decision: list[tuple[str, list[int]]],
|
|
*,
|
|
debug_mask="all",
|
|
) -> int | None:
|
|
for sketch_id, sketch in enumerate(sketches):
|
|
i = 0
|
|
new_decisions = {}
|
|
for inst in sketch.trace.insts:
|
|
if not inst.kind.name.startswith("Sample"):
|
|
continue
|
|
assert i < len(expected_decision)
|
|
if inst.kind.name == expected_decision[i][0]:
|
|
new_decisions[inst] = expected_decision[i][1]
|
|
i += 1
|
|
if len(new_decisions) != len(expected_decision):
|
|
continue
|
|
sch = Schedule(mod, debug_mask=debug_mask)
|
|
Trace(
|
|
insts=sketch.trace.insts,
|
|
decisions=new_decisions,
|
|
).apply_to_schedule(sch, remove_postproc=True)
|
|
if structural_equal_no_gs(sch.mod, expected_mod):
|
|
verify_trace_roundtrip(sch=sch, mod=mod, debug_mask=debug_mask, text_format="json")
|
|
return sketch_id
|
|
return None
|
|
|
|
|
|
def check_sketches(
|
|
mod: IRModule,
|
|
sketches: list[Schedule],
|
|
expected_mods: list[IRModule],
|
|
expected_decisions: list[list[tuple[str, list[int]]]],
|
|
*,
|
|
debug_mask="all",
|
|
):
|
|
assert len(expected_mods) == len(expected_decisions)
|
|
assert len(sketches) == len(expected_mods)
|
|
expected_mods = [
|
|
IRModule({"main": m}) if not isinstance(m, IRModule) else m for m in expected_mods
|
|
]
|
|
sketches = list(sketches)
|
|
for expected_id, (expected_mod, expected_decision) in enumerate(
|
|
zip(expected_mods, expected_decisions)
|
|
):
|
|
sketch_id = _find_match_sketch_id(
|
|
mod,
|
|
sketches,
|
|
expected_mod,
|
|
expected_decision,
|
|
debug_mask=debug_mask,
|
|
)
|
|
if sketch_id is None:
|
|
raise AssertionError(
|
|
f"Expected sketch #{expected_id} doesn't exist in the generated sketches."
|
|
)
|
|
sketches.pop(sketch_id)
|
|
|
|
|
|
def print_sketches(sketches: list[Schedule]):
|
|
for i, sch in enumerate(sketches):
|
|
print(f"###### {i}")
|
|
sch.mod.show(black_format=False)
|
|
for inst in sch.trace.insts:
|
|
if inst in sch.trace.decisions:
|
|
print(f'("{inst.kind.name}", {sch.trace.decisions[inst]}),')
|