wehub-resource-sync
298 lines
9.6 KiB
Python
298 lines
9.6 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.
|
|
"""An execution trace of a scheduling program"""
|
|
|
|
import os
|
|
from collections.abc import Callable
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
from tvm_ffi import Array, Map
|
|
from tvm_ffi import register_object as _register_object
|
|
|
|
from tvm.runtime import Object
|
|
from tvm.tirx.expr import FloatImm, IntImm
|
|
from tvm.tirx.function import IndexMap
|
|
|
|
from ...ir import save_json
|
|
from . import _ffi_api
|
|
from .instruction import ATTR_TYPE, INPUT_RV_TYPE, Instruction
|
|
|
|
if TYPE_CHECKING:
|
|
from .schedule import Schedule
|
|
|
|
|
|
DECISION_TYPE = Any
|
|
JSON_TYPE = Any
|
|
|
|
|
|
def _json_from_tvm(obj):
|
|
if obj is None:
|
|
return None
|
|
elif isinstance(obj, bool | int | float | str):
|
|
return obj
|
|
elif isinstance(obj, Array):
|
|
return [_json_from_tvm(i) for i in obj]
|
|
elif isinstance(obj, Map):
|
|
return {_json_from_tvm(k): _json_from_tvm(v) for k, v in obj.items()}
|
|
elif isinstance(obj, str):
|
|
return str(obj)
|
|
elif isinstance(obj, IntImm | FloatImm):
|
|
return obj
|
|
elif isinstance(obj, IndexMap):
|
|
return save_json(obj)
|
|
else:
|
|
raise TypeError("Not supported type: " + str(type(obj)))
|
|
|
|
|
|
@_register_object("s_tir.Trace")
|
|
class Trace(Object):
|
|
"""An execution trace of a scheduling program.
|
|
|
|
A trace has two parts:
|
|
1) The instructions invoked so far
|
|
2) The random decisions made upon those instructions, if any
|
|
|
|
A trace can be serialized to:
|
|
1) Roundtrippable JSON format: can be saved to file and loaded back
|
|
2) Python syntax: allows users to copy-paste the trace to reproduce the scheduling process
|
|
|
|
A trace can be applied to a TensorIR schedule by re-applying all its instructions possibly with
|
|
their decisions accordingly. Re-sampling is invoked if a sampling instruction doesn't have its
|
|
corresponding decision; Otherwise the existing decision will be reused accordingly.
|
|
|
|
Attributes
|
|
----------
|
|
insts : List[Instruction]
|
|
The instructions invoked so far in the program execution
|
|
decisions : Dict[Instruction, DECISION_TYPE]
|
|
The random decisions made upon those instructions
|
|
"""
|
|
|
|
insts: list[Instruction]
|
|
decisions: dict[Instruction, DECISION_TYPE]
|
|
|
|
def __init__(
|
|
self,
|
|
insts: list[Instruction],
|
|
decisions: dict[Instruction, DECISION_TYPE],
|
|
) -> None:
|
|
"""Constructor
|
|
|
|
Parameters
|
|
----------
|
|
insts : List[Instruction]
|
|
The instructions invoked so far in the program execution
|
|
decisions : Dict[Instruction, DECISION_TYPE]
|
|
The random decisions made upon those instructions
|
|
"""
|
|
self.__init_handle_by_constructor__(
|
|
_ffi_api.Trace, # type: ignore # pylint: disable=no-member
|
|
insts,
|
|
decisions,
|
|
)
|
|
|
|
def get_decision(self, inst: Instruction) -> DECISION_TYPE | None:
|
|
"""Retrieve the decision made on a specific instruction
|
|
|
|
Parameters
|
|
----------
|
|
insts : Instruction
|
|
The instruction whose decision is to be retrieved
|
|
|
|
Returns
|
|
-------
|
|
decision : Optional[DECISION_TYPE]
|
|
The corresponding decision; None if there is no decision made on the instruction
|
|
"""
|
|
return _ffi_api.TraceGetDecision(self, inst) # type: ignore # pylint: disable=no-member
|
|
|
|
def append(
|
|
self,
|
|
inst: Instruction,
|
|
decision: DECISION_TYPE | None = None,
|
|
) -> None:
|
|
"""Append a new instruction to the trace
|
|
|
|
Parameters
|
|
----------
|
|
insts : Instruction
|
|
The new instruction to be appended
|
|
decision : Optional[DECISION_TYPE] = None
|
|
The random decision made on this instruction
|
|
"""
|
|
_ffi_api.TraceAppend(self, inst, decision) # type: ignore # pylint: disable=no-member
|
|
|
|
def pop(self) -> Instruction | None:
|
|
"""Remove the last instruction, along with the decision made on that instruction, if any
|
|
|
|
Returns
|
|
-------
|
|
popped_inst : Instruction
|
|
Returns the instruction removed; std::nullopt if the trace is empty
|
|
"""
|
|
return _ffi_api.TracePop(self) # type: ignore # pylint: disable=no-member
|
|
|
|
def apply_to_schedule(
|
|
self,
|
|
sch: "Schedule",
|
|
remove_postproc: bool,
|
|
decision_provider: (
|
|
Callable[
|
|
[Instruction, list[INPUT_RV_TYPE], list[ATTR_TYPE], DECISION_TYPE], DECISION_TYPE
|
|
]
|
|
| None
|
|
) = None,
|
|
) -> None:
|
|
"""Apply the trace to a TensorIR schedule
|
|
|
|
Parameters
|
|
----------
|
|
sch : Schedule
|
|
The schedule to be applied onto
|
|
remove_postproc : bool
|
|
If postprocessing instructions are removed
|
|
decision_provider: Optional[Callable] = None
|
|
A callback that allows users to mutate decisions on the fly when applying instructions.
|
|
The signature of the callback is:
|
|
- The 1st argument: The instruction
|
|
- The 2nd argument: The input random variables
|
|
- The 3rd argument: The attributes
|
|
- The 4th argument: The decision
|
|
- Return: A new decision
|
|
"""
|
|
_ffi_api.TraceApplyToSchedule( # type: ignore # pylint: disable=no-member
|
|
self,
|
|
sch,
|
|
remove_postproc,
|
|
decision_provider,
|
|
)
|
|
|
|
def as_json(self, remove_postproc: bool = False) -> JSON_TYPE:
|
|
"""Serialize the trace as a JSON-style object
|
|
|
|
Parameters
|
|
----------
|
|
remove_postproc : bool = False
|
|
If postprocessing instructions are removed
|
|
|
|
Returns
|
|
-------
|
|
json: JSON_TYPE
|
|
The JSON-style object
|
|
"""
|
|
obj = _ffi_api.TraceAsJSON(self, remove_postproc) # type: ignore # pylint: disable=no-member
|
|
return _json_from_tvm(obj)
|
|
|
|
def as_python(self, remove_postproc: bool = False) -> list[str]:
|
|
"""Serialize the trace as a sequence of python statements
|
|
|
|
Parameters
|
|
----------
|
|
remove_postproc : bool = False
|
|
If postprocessing instructions are removed
|
|
|
|
Returns
|
|
-------
|
|
py_stmts: List[str]
|
|
A sequence of python statements
|
|
"""
|
|
return _ffi_api.TraceAsPython(self, remove_postproc) # type: ignore # pylint: disable=no-member
|
|
|
|
def with_decision(
|
|
self,
|
|
inst: Instruction,
|
|
decision: DECISION_TYPE,
|
|
remove_postproc: bool,
|
|
) -> "Trace":
|
|
"""Create a new trace with an instruction whose decision is changed,
|
|
assuming this instruction exists in the resulting trace
|
|
|
|
Parameters
|
|
----------
|
|
inst : Instruction
|
|
The instruction whose decision is to be changed
|
|
decision : DECISION_TYPE
|
|
The decision to be changed to
|
|
remove_postproc : bool
|
|
If postprocessing instructions are removed
|
|
|
|
Returns
|
|
-------
|
|
trace: Trace
|
|
The new trace with the decision changed
|
|
"""
|
|
return _ffi_api.TraceWithDecision( # type: ignore # pylint: disable=no-member
|
|
self,
|
|
inst,
|
|
decision,
|
|
remove_postproc,
|
|
)
|
|
|
|
def simplified(self, remove_postproc: bool) -> "Trace":
|
|
"""Simplify the trace with dead-code elimination
|
|
|
|
Parameters
|
|
----------
|
|
remove_postproc : bool
|
|
If postprocessing instructions are removed
|
|
|
|
Returns
|
|
-------
|
|
trace: Trace
|
|
A simplified trace
|
|
"""
|
|
return _ffi_api.TraceSimplified(self, remove_postproc) # type: ignore # pylint: disable=no-member
|
|
|
|
@staticmethod
|
|
def apply_json_to_schedule(json_obj: JSON_TYPE, sch: "Schedule") -> None:
|
|
"""Apply a JSON-serialized trace to a TensorIR schedule
|
|
|
|
Parameters
|
|
----------
|
|
json_obj : JSON_TYPE
|
|
The JSON-serialized trace
|
|
sch : Schedule
|
|
The TensorIR schedule
|
|
"""
|
|
_ffi_api.TraceApplyJSONToSchedule(json_obj, sch) # type: ignore # pylint: disable=no-member
|
|
|
|
def show(self, style: str | None = None, black_format: bool = False) -> None:
|
|
"""A sugar for print highlighted TVM script.
|
|
|
|
Parameters
|
|
----------
|
|
style : str, optional
|
|
|
|
Pygmentize printing style, auto-detected if None. See
|
|
`tvm.script.highlight.cprint` for more details.
|
|
|
|
black_format: bool
|
|
|
|
If true, use the formatter Black to format the TVMScript.
|
|
If None, determine based on the "TVM_BLACK_FORMAT" environment
|
|
variable.
|
|
"""
|
|
from tvm.script.highlight import ( # pylint: disable=import-outside-toplevel
|
|
cprint,
|
|
)
|
|
|
|
if black_format is None:
|
|
env = os.environ.get("TVM_BLACK_FORMAT")
|
|
black_format = bool(env and int(env))
|
|
|
|
cprint(str(self), style=style, black_format=black_format)
|