# 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)