From 2da30f5ebbfc651dc558a8e42c40a1207b478376 Mon Sep 17 00:00:00 2001 From: Junru Shao Date: Sun, 1 Aug 2021 20:54:04 +0000 Subject: [PATCH] [Meta Schedule][M3a] Instruction and Trace Co-authored-by: Siyuan Feng Co-authored-by: Bohan Hou <32121147+spectrometerHBH@users.noreply.github.com> Co-authored-by: Ruihang Lai Co-authored-by: Hongyi Jin <3231950289@qq.com> Co-authored-by: Wuwei Lin Co-authored-by: Cody Yu --- include/tvm/tir/schedule/instruction.h | 288 ++++++++++ include/tvm/tir/schedule/schedule.h | 19 +- include/tvm/tir/schedule/state.h | 8 - include/tvm/tir/schedule/trace.h | 164 ++++++ python/tvm/tir/schedule/__init__.py | 4 +- .../{_ffi_api_schedule.py => _ffi_api.py} | 0 python/tvm/tir/schedule/block_scope.py | 14 +- python/tvm/tir/schedule/instruction.py | 166 ++++++ python/tvm/tir/schedule/schedule.py | 130 +++-- python/tvm/tir/schedule/state.py | 26 +- python/tvm/tir/schedule/trace.py | 260 +++++++++ src/tir/schedule/analysis.h | 16 - src/tir/schedule/analysis/analysis.cc | 34 -- src/tir/schedule/concrete_schedule.h | 22 +- src/tir/schedule/instruction.cc | 102 ++++ src/tir/schedule/instruction_traits.h | 536 ++++++++++++++++++ src/tir/schedule/primitive.h | 36 +- src/tir/schedule/primitive/compute_inline.cc | 51 ++ src/tir/schedule/primitive/get_block_loop.cc | 113 ++++ .../schedule/primitive/loop_transformation.cc | 74 +++ src/tir/schedule/primitive/reduction.cc | 29 + src/tir/schedule/schedule.cc | 38 +- src/tir/schedule/state.cc | 17 +- src/tir/schedule/trace.cc | 533 +++++++++++++++++ src/tir/schedule/utils.h | 4 + .../unittest/test_tir_schedule_block_scope.py | 7 +- .../test_tir_schedule_compute_inline.py | 20 +- .../unittest/test_tir_schedule_error.py | 7 +- .../unittest/test_tir_schedule_instruction.py | 68 +++ .../unittest/test_tir_schedule_reduction.py | 5 +- .../unittest/test_tir_schedule_split_fuse.py | 4 +- .../unittest/test_tir_schedule_state.py | 17 +- .../test_tir_schedule_state_cached_flags.py | 19 +- .../unittest/test_tir_schedule_trace.py | 241 ++++++++ .../unittest/test_tir_schedule_utilities.py | 9 +- 35 files changed, 2836 insertions(+), 245 deletions(-) create mode 100644 include/tvm/tir/schedule/instruction.h create mode 100644 include/tvm/tir/schedule/trace.h rename python/tvm/tir/schedule/{_ffi_api_schedule.py => _ffi_api.py} (100%) create mode 100644 python/tvm/tir/schedule/instruction.py create mode 100644 python/tvm/tir/schedule/trace.py create mode 100644 src/tir/schedule/instruction.cc create mode 100644 src/tir/schedule/instruction_traits.h create mode 100644 src/tir/schedule/primitive/get_block_loop.cc create mode 100644 src/tir/schedule/trace.cc create mode 100644 tests/python/unittest/test_tir_schedule_instruction.py create mode 100644 tests/python/unittest/test_tir_schedule_trace.py diff --git a/include/tvm/tir/schedule/instruction.h b/include/tvm/tir/schedule/instruction.h new file mode 100644 index 000000000000..5a9e687dc8c7 --- /dev/null +++ b/include/tvm/tir/schedule/instruction.h @@ -0,0 +1,288 @@ +/* + * 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. + */ +#ifndef TVM_TIR_SCHEDULE_INSTRUCTION_H_ +#define TVM_TIR_SCHEDULE_INSTRUCTION_H_ + +#include + +#include + +namespace tvm { + +// Forward declaration +template +class AttrRegistry; + +namespace tir { + +// Forward declaration +class Schedule; + +/*! + * \brief Type of the functor that applies the instruction to a TensorIR schedule + * \param sch The schedule to be applied on + * \param inputs The input random variables + * \param attrs Instruction attributes + * \param decision Decisions made on the instruction + * \return The functor returns an array of output random variables + */ +using FInstructionApply = runtime::TypedPackedFunc( + Schedule sch, const Array& inputs, const Array& attrs, + const Optional& decision)>; + +/*! + * \brief Type of the functor that converts the instruction to a statement in python syntax + * \param inputs Names of the input random variables + * \param attrs Instruction attributes + * \param decisions Decisions made on the instruction + * \param outputs Names of the output random variables + * \return A string representing the python api call + */ +using FInstructionAsPython = runtime::TypedPackedFunc& inputs, const Array& attrs, + const Optional& decision, const Array& outputs)>; + +/*! + * \brief Type of the functor that serialize its attributes to JSON + * \param attrs The attributes to be serialized + * \return An array, serialized attributes + * \note This functor is nullable + */ +using FInstructionAttrsAsJSON = runtime::TypedPackedFunc attrs)>; + +/*! + * \brief Type of the functor that deserialize its attributes from JSON + * \param json_attrs The attributes to be serialized + * \return An array, deserialized attributes + * \note This functor is nullable + */ +using FInstructionAttrsFromJSON = runtime::TypedPackedFunc(ObjectRef json_attrs)>; + +/*! + * \brief Kind of an instruction, e.g. Split, Reorder, etc. + * Besides the name, every kind of instruction has its own properties, including: + * 1) A boolean indicating if the instruction is pure, i.e. change nothing in the schedule state + * 2) A functor that applies the instruction to a TensorIR schedule + * 3) A functor that converts the instruction to a statement in python syntax + * 4) A functor that serialize its attributes to JSON + * 5) A functor that deserialize its attributes from JSON + * + * Unlike `tvm::OpNode`, `InstructionKindNode` doesn't support unstructured properties, + * mainly because there is no such usecase yet to add any other property. + */ +class InstructionKindNode : public runtime::Object { + public: + /*! \brief The name of a kind of instructions */ + String name; + /*! + * \brief Indicates if the instruction is pure, i.e. removing it alone doesn't mutate the schedule + * state. For example, the instruction `GetBlock` is pure because it changes + * nothing, while `ComputeInline` is not because removing it leads to a different resulting + * schedule. + */ + bool is_pure{false}; + /*! \brief A functor that applies the instruction to a TensorIR schedule */ + FInstructionApply f_apply_to_schedule{nullptr}; + /*! \brief A functor that converts the instruction to a statement in python syntax */ + FInstructionAsPython f_as_python{nullptr}; + /*! + * \brief A functor that serialize its attributes to JSON + * \note If the functor is null, it means no conversion is needed + */ + FInstructionAttrsAsJSON f_attrs_as_json{nullptr}; + /*! + * \brief A functor that deserialize its attributes from JSON + * \note If the functor is null, it means no conversion is needed + */ + FInstructionAttrsFromJSON f_attrs_from_json{nullptr}; + + void VisitAttrs(tvm::AttrVisitor* v) { + v->Visit("name", &name); + v->Visit("_is_pure", &is_pure); + // not visited: f_apply_to_schedule + // not visited: f_as_python + // not visited: f_attrs_as_json + // not visited: f_attrs_from_json + } + + static constexpr const char* _type_key = "tir.InstructionKind"; + TVM_DECLARE_FINAL_OBJECT_INFO(InstructionKindNode, runtime::Object); +}; + +/*! + * \brief Managed reference to InstructionKindNode + * \sa InstructionKindNode + */ +class InstructionKind : public runtime::ObjectRef { + public: + /*! + * \brief Retrieve an InstructionKind using its name + * \param name The registered name of the InstructionKind + * \return The InstructionKind retrieved + */ + static InstructionKind Get(const String& name); + TVM_DEFINE_OBJECT_REF_METHODS(InstructionKind, runtime::ObjectRef, InstructionKindNode); +}; + +/*! \brief Schedule instructions each corresponds to a schedule primitive */ +class InstructionNode : public runtime::Object { + public: + /*! \brief The kind of the instruction */ + InstructionKind kind; + /*! + * \brief The input random variables of the instruction, and the type of each element can be one + * of the following: + * - BlockRV + * - LoopRV + * - ExprRV + * - FloatImm + * - IntImm + * - String + * - null pointer + */ + Array inputs; + /*! + * \brief The attributes of the instruction. Similar to attributes of an operator, + * attributes of an instruction are arbitrary constant metadata required by the instructions. + * For example, the name of the block to be retrieved in `GetBlock`. + */ + Array attrs; + /*! \brief The output random variables of the instruction, and the type of each element can be one + * of the following: + * - BlockRV + * - LoopRV + * - ExprRV, atomic variables only, won't be constants or composite PrimExpr + */ + Array outputs; + + void VisitAttrs(tvm::AttrVisitor* v) { + v->Visit("kind", &kind); + v->Visit("inputs", &inputs); + v->Visit("attrs", &attrs); + v->Visit("outputs", &outputs); + } + + static constexpr const char* _type_key = "tir.Instruction"; + TVM_DECLARE_FINAL_OBJECT_INFO(InstructionNode, runtime::Object); +}; + +/*! + * \brief Managed reference to InstructionNode + * \sa InstructionNode + */ +class Instruction : public runtime::ObjectRef { + public: + /*! + * \brief Constructor + * \param kind The kind of the instruction + * \param inputs The input random variables of the instruction + * \param attrs The attributes of the instruction + * \param outputs The output random variables of the instruction + */ + explicit Instruction(InstructionKind kind, Array inputs, Array attrs, + Array outputs); + + TVM_DEFINE_OBJECT_REF_METHODS(Instruction, runtime::ObjectRef, InstructionNode); +}; + +/*! + * \brief A helper macro to register InstructionKind, only used in `TVM_REGISTER_INST_KIND` + * \note This macro is not user-facing. + * \sa TVM_REGISTER_INST_KIND + */ +#define TVM_INST_KIND_REGISTER_VAR_DEF \ + static DMLC_ATTRIBUTE_UNUSED ::tvm::tir::InstructionKindRegEntry& __make_##InstructionKind + +/*! + * \brief Register an InstructionKind + * \param InstructionKindName The name of the InstructionKind + * + * Example: + * + * \code + * + * TVM_REGISTER_INST_KIND("ComputeInline") + * .set_is_pure(false) + * .set_apply_to_schedule(ApplyToSchedule) + * .set_attrs_as_json(AttrsAsJSON) + * .set_attrs_from_json(AttrsFromJSON) + * .set_as_python(AsPython); + * + * \endcode + */ +#define TVM_REGISTER_INST_KIND(InstructionKindName) \ + TVM_STR_CONCAT(TVM_INST_KIND_REGISTER_VAR_DEF, __COUNTER__) = \ + ::tvm::tir::InstructionKindRegEntry::RegisterOrGet(InstructionKindName).set_name() + +/*! \brief An entry in the registry of InstructionKind */ +class InstructionKindRegEntry { + public: + static InstructionKindRegEntry& RegisterOrGet(const String& name); + + InstructionKindRegEntry& set_name() { + get_mutable()->name = this->name; + return *this; + } + + InstructionKindRegEntry& set_is_pure(bool is_pure) { + get_mutable()->is_pure = is_pure; + return *this; + } + + InstructionKindRegEntry& set_apply_to_schedule(FInstructionApply f_apply_to_schedule) { + get_mutable()->f_apply_to_schedule = std::move(f_apply_to_schedule); + return *this; + } + + InstructionKindRegEntry& set_as_python(FInstructionAsPython f_as_python) { + get_mutable()->f_as_python = std::move(f_as_python); + return *this; + } + + InstructionKindRegEntry& set_attrs_as_json(FInstructionAttrsAsJSON f_attrs_as_json) { + get_mutable()->f_attrs_as_json = std::move(f_attrs_as_json); + return *this; + } + + InstructionKindRegEntry& set_attrs_from_json(FInstructionAttrsFromJSON f_attrs_from_json) { + get_mutable()->f_attrs_from_json = std::move(f_attrs_from_json); + return *this; + } + + private: + /*! \brief Private constructor, used only by AttrRegistry */ + explicit InstructionKindRegEntry(uint32_t reg_index); + /*! \brief Get the mutable reference to the internal InstructionKind */ + InstructionKindNode* get_mutable() const { + return const_cast(inst_kind_.get()); + } + + /*! \brief The name of the registry entry */ + String name; + /*! \brief The instruction kind */ + InstructionKind inst_kind_; + template + friend class ::tvm::AttrRegistry; + friend class InstructionKind; +}; + +} // namespace tir +} // namespace tvm + +#endif // TVM_TIR_SCHEDULE_INSTRUCTION_H_ diff --git a/include/tvm/tir/schedule/schedule.h b/include/tvm/tir/schedule/schedule.h index 245a904b91ee..bd2377397626 100644 --- a/include/tvm/tir/schedule/schedule.h +++ b/include/tvm/tir/schedule/schedule.h @@ -180,7 +180,8 @@ class ScheduleNode : public runtime::Object { virtual void RemoveRV(const ExprRV& expr_rv) = 0; public: - /******** Block/Loop relation ********/ + /******** Schedule: Sampling ********/ + /******** Schedule: Get blocks & loops ********/ /*! * \brief Retrieve a block in a specific function with its name * \param name The name of the block to be retrieved @@ -195,7 +196,7 @@ class ScheduleNode : public runtime::Object { * \return A list of loops above the given block in its scope, from outer to inner */ virtual Array GetLoops(const BlockRV& block_rv) = 0; - /******** Schedule: loops manipulation ********/ + /******** Schedule: Transform loops ********/ /*! * \brief Fuse a list of consecutive loops into one. It requires: * 1) The loops can't have annotations or thread bindings. @@ -215,7 +216,9 @@ class ScheduleNode : public runtime::Object { * \return The new loops after split */ virtual Array Split(const LoopRV& loop_rv, const Array>& factors) = 0; - /******** Schedule: compute location ********/ + /******** Schedule: Manipulate ForKind ********/ + /******** Schedule: Insert cache stages ********/ + /******** Schedule: Compute location ********/ /*! * \brief Inline a block into its consumer(s). It requires: * 1) The block is a complete non-root block, which only produces one buffer @@ -239,9 +242,7 @@ class ScheduleNode : public runtime::Object { * \param block The block to be inlined to its producer */ virtual void ReverseComputeInline(const BlockRV& block) = 0; - /******** Schedule: loop binding/annotation ********/ - /******** Schedule: cache read/write ********/ - /******** Schedule: reduction ********/ + /******** Schedule: Reduction ********/ /*! * \brief Factorize an associative reduction block by the specified loop. * \details An associative reduction cannot be parallelized directly, @@ -260,7 +261,11 @@ class ScheduleNode : public runtime::Object { * \return The rfactor block */ virtual BlockRV RFactor(const LoopRV& loop_rv, int factor_axis) = 0; - /******** Schedule: blockize & tensorize ********/ + /******** Schedule: Blockize & Tensorize ********/ + /******** Schedule: Annotation ********/ + /******** Schedule: Misc ********/ + /*! \brief A no-op that marks the start of postprocessing phase of scheduling */ + virtual void EnterPostproc() = 0; }; /*! diff --git a/include/tvm/tir/schedule/state.h b/include/tvm/tir/schedule/state.h index 83ac7150543f..077bf938f48a 100644 --- a/include/tvm/tir/schedule/state.h +++ b/include/tvm/tir/schedule/state.h @@ -190,14 +190,6 @@ class ScheduleState : public ObjectRef { * and each time after calling the Replace method. */ TVM_DLL explicit ScheduleState(IRModule mod, int debug_mode = 0); - /*! - * \brief Construct a schedule state from a PrimFunc - * \param func The PrimFunc to be scheduled. A new IRModule will be created with - * this specific PrimFunc as "main" function in the module to be scheduled - * \param debug_mode Do extra correctness checking after the class creation - * and each time after calling the Replace method. - */ - TVM_DLL explicit ScheduleState(PrimFunc func, int debug_mode = 0); /*! \return The mutable pointer to the ScheduleStateNode */ ScheduleStateNode* get() const { return static_cast(data_.get()); } diff --git a/include/tvm/tir/schedule/trace.h b/include/tvm/tir/schedule/trace.h new file mode 100644 index 000000000000..b6b3b57226c8 --- /dev/null +++ b/include/tvm/tir/schedule/trace.h @@ -0,0 +1,164 @@ +/* + * 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. + */ +#ifndef TVM_TIR_SCHEDULE_TRACE_H_ +#define TVM_TIR_SCHEDULE_TRACE_H_ + +#include + +namespace tvm { +namespace tir { + +// Forward declaration +class Trace; + +/*! + * \brief A callback that allows users to mutate decisions on the fly + * when applying instructions. The signature of the callback is: + * \param inst The instruction + * \param inputs The input random variables + * \param attrs The attributes + * \param decision The original decision + * \return A new decision + */ +using FTraceDecisionProvider = runtime::TypedPackedFunc& inputs, const Array& attrs, + const Optional& decision)>; + +/*! + * \brief An execution trace of a scheduling program + * + * A trace has two parts: + * 1) The instructions invoked so far in the program execution + * 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. + */ +class TraceNode : public runtime::Object { + public: + /*! \brief The instructions invoked so far in the program execution */ + Array insts; + /*! \brief The random decisions made upon those instructions */ + Map decisions; + + void VisitAttrs(tvm::AttrVisitor* v) { + v->Visit("insts", &insts); + v->Visit("decisions", &decisions); + } + + static constexpr const char* _type_key = "tir.Trace"; + TVM_DECLARE_FINAL_OBJECT_INFO(TraceNode, runtime::Object); + + public: + /*! + * \brief Retrieve the decision made on a specific instruction + * \param inst The instruction whose decision is to be retrieved + * \return The corresponding decision; NullOpt if there is no decision made on the instruction + */ + Optional GetDecision(const Instruction& inst) const; + /*! + * \brief Append a new instruction to the trace + * \param inst The new instruction to be appended + */ + void Append(Instruction inst); + /*! + * \brief Append a new instruction with a random decision to the trace + * \param inst The new instruction to be appended + * \param decision The random decision made on this instruction + * The type of `decision` depends on the instruction, e.g. + * the decision of `SamplePerfectTile` has type `Array` + */ + void Append(Instruction inst, ObjectRef decision); + /*! + * \brief Remove the last instruction, along with the decision made on that instruction, if any + * \return The instruction removed; NullOpt if the trace is empty + */ + Optional Pop(); + /*! + * \brief Apply the trace to a TensorIR schedule + * \param sch The schedule to be applied onto + * \param remove_postproc If postprocessing instructions are removed + * \param decision_provider A callback that allows users to mutate decisions on the fly + * when applying instructions. + * \sa FTraceDecisionProvider + */ + void ApplyToSchedule(Schedule sch, bool remove_postproc, + FTraceDecisionProvider decision_provider = nullptr) const; + /*! + * \brief Serialize the trace as a JSON-style object + * \param remove_postproc If postprocessing instructions are removed + * \return The JSON-style object + */ + ObjectRef AsJSON(bool remove_postproc) const; + /*! + * \brief Serialize the trace as a sequence of python statements + * \param remove_postproc If postprocessing instructions are removed + * \return A sequence of python statements + */ + Array AsPython(bool remove_postproc) const; + /*! + * \brief Create a new trace with an instruction whose decision is changed, + * assuming this instruction exists in the resulting trace + * \param inst The instruction whose decision is to be changed + * \param decision The decision to be changed to + * \param remove_postproc If postprocessing instructions are removed + * \return The new trace with the decision changed + */ + Trace WithDecision(Instruction inst, ObjectRef decision, bool remove_postproc) const; + /*! + * \brief Simplify the trace with dead-code elimination + * \param remove_postproc If postprocessing instructions are removed + * \return A simplified trace + */ + Trace Simplified(bool remove_postproc) const; +}; + +/*! + * \brief Managed reference to TraceNode + * \sa TraceNode + */ +class Trace : public runtime::ObjectRef { + public: + /*! \brief Default constructor. Creating an empty trace. */ + Trace(); + /*! + * \brief Constructor. Creating a trace from existing instructions and their decisions + * \param insts The instructions used + * \param decisions The decisions made in sampling + */ + explicit Trace(Array insts, Map decisions); + /*! + * \brief Apply a JSON-serialized trace to a TensorIR schedule + * \param json The JSON-serialized trace + * \param sch The TensorIR schedule + */ + static void ApplyJSONToSchedule(ObjectRef json, Schedule sch); + + TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(Trace, runtime::ObjectRef, TraceNode); +}; + +} // namespace tir +} // namespace tvm + +#endif // TVM_TIR_SCHEDULE_TRACE_H_ diff --git a/python/tvm/tir/schedule/__init__.py b/python/tvm/tir/schedule/__init__.py index ef1cab1fb663..5f0e169c43e3 100644 --- a/python/tvm/tir/schedule/__init__.py +++ b/python/tvm/tir/schedule/__init__.py @@ -18,5 +18,7 @@ """Namespace for the TensorIR schedule API.""" from .block_scope import BlockScope, Dependency, DepKind, StmtSRef +from .instruction import Instruction, InstructionKind +from .schedule import BlockRV, ExprRV, LoopRV, Schedule, ScheduleError from .state import ScheduleDebugMask, ScheduleState -from .schedule import LoopRV, BlockRV, ExprRV, RAND_VAR_TYPE, Schedule, ScheduleError +from .trace import Trace diff --git a/python/tvm/tir/schedule/_ffi_api_schedule.py b/python/tvm/tir/schedule/_ffi_api.py similarity index 100% rename from python/tvm/tir/schedule/_ffi_api_schedule.py rename to python/tvm/tir/schedule/_ffi_api.py diff --git a/python/tvm/tir/schedule/block_scope.py b/python/tvm/tir/schedule/block_scope.py index 061a472ad9a9..30e047b4f78a 100644 --- a/python/tvm/tir/schedule/block_scope.py +++ b/python/tvm/tir/schedule/block_scope.py @@ -22,7 +22,7 @@ from tvm.runtime import Object from tvm.tir import Block, For -from . import _ffi_api_schedule +from . import _ffi_api @register_object("tir.StmtSRef") @@ -45,24 +45,24 @@ class StmtSRef(Object): @property def stmt(self) -> Optional[Union[Block, For]]: """The block/for stmt the object refers to""" - return _ffi_api_schedule.StmtSRefStmt(self) # type: ignore # pylint: disable=no-member + return _ffi_api.StmtSRefStmt(self) # type: ignore # pylint: disable=no-member @property def parent(self) -> Optional["StmtSRef"]: """The parent sref""" - return _ffi_api_schedule.StmtSRefParent(self) # type: ignore # pylint: disable=no-member + return _ffi_api.StmtSRefParent(self) # type: ignore # pylint: disable=no-member @staticmethod def inline_mark() -> "StmtSRef": """A special StmtSRef, which doesn't point to any stmt in the AST, only serving as a "mark" to hint compute-at to do the work of compute-inline""" - return _ffi_api_schedule.StmtSRefInlineMark() # type: ignore # pylint: disable=no-member + return _ffi_api.StmtSRefInlineMark() # type: ignore # pylint: disable=no-member @staticmethod def root_mark() -> "StmtSRef": """A special StmtSRef, which doesn't point to any stmt in the AST, only serving as a "mark" to hint compute-at to do nothing""" - return _ffi_api_schedule.StmtSRefRootMark() # type: ignore # pylint: disable=no-member + return _ffi_api.StmtSRefRootMark() # type: ignore # pylint: disable=no-member class DepKind(IntEnum): @@ -137,7 +137,7 @@ def get_deps_by_src(self, block: StmtSRef) -> List[Dependency]: blocks: List[Dependency] The dependencies """ - return _ffi_api_schedule.BlockScopeGetDepsBySrc(self, block) # type: ignore # pylint: disable=no-member + return _ffi_api.BlockScopeGetDepsBySrc(self, block) # type: ignore # pylint: disable=no-member def get_deps_by_dst(self, block: StmtSRef) -> List[Dependency]: """Get all dependencies whose `dst` is the target `block`. @@ -152,4 +152,4 @@ def get_deps_by_dst(self, block: StmtSRef) -> List[Dependency]: blocks: List[Dependency] The dependencies """ - return _ffi_api_schedule.BlockScopeGetDepsByDst(self, block) # type: ignore # pylint: disable=no-member + return _ffi_api.BlockScopeGetDepsByDst(self, block) # type: ignore # pylint: disable=no-member diff --git a/python/tvm/tir/schedule/instruction.py b/python/tvm/tir/schedule/instruction.py new file mode 100644 index 000000000000..09b2d70dc321 --- /dev/null +++ b/python/tvm/tir/schedule/instruction.py @@ -0,0 +1,166 @@ +# 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. +"""Schedule instructions each corresponds to a schedule primitive""" +from typing import TYPE_CHECKING, Any, List, Union + +from tvm._ffi import register_object as _register_object +from tvm.runtime import Object + +from . import _ffi_api + +if TYPE_CHECKING: + from .schedule import RAND_VAR_TYPE + + INPUT_RV_TYPE = Union[RAND_VAR_TYPE, float, int, str, None] # pylint: disable=invalid-name + OUTPUT_RV_TYPE = Union[RAND_VAR_TYPE] # pylint: disable=invalid-name + ATTR_TYPE = Any +else: + INPUT_RV_TYPE = OUTPUT_RV_TYPE = ATTR_TYPE = Any + + +@_register_object("tir.InstructionKind") +class InstructionKind(Object): + """Kind of an instruction, e.g. Split, Reorder, etc. + Besides the name, every kind of instruction has its own properties, including: + 1) A boolean indicating if the instruction is pure, i.e. change nothing in the schedule state + 2) A functor that applies the instruction to a TensorIR schedule + 3) A functor that converts the instruction to a statement in python syntax + 4) A functor that serialize its attributes to JSON + 5) A functor that deserialize its attributes from JSON + + Unlike `tvm.ir.op`, `InstructionKind` doesn't support unstructured properties, + mainly because there is no such usecase yet to add any other property. + + Attributes + ---------- + name : str + The name of a kind of instructions + + Note + ---- + The functor properties are not exposed on python side at the moment + """ + + name: str + + @property + def is_pure(self) -> bool: + """Indicates if the instruction is pure, i.e. removing it alone doesn't mutate the schedule + state. For example, the instruction `GetBlock` is pure because it changes + nothing, while `ComputeInline` is not because removing it leads to a different resulting + schedule. + + Returns + ------- + pure : bool + The boolean flag indicating if the instruction is pure + """ + return bool(self._is_pure) + + @staticmethod + def get(name: str) -> "InstructionKind": + """Retrieve an InstructionKind using its name + + Parameters + ---------- + name : str + The registered name of the InstructionKind + + Returns + ------- + kind : InstructionKind + The InstructionKind retrieved + """ + return _ffi_api.InstructionKindGet(name) # type: ignore # pylint: disable=no-member + + +@_register_object("tir.Instruction") +class Instruction(Object): + """Schedule instructions each corresponds to a schedule primitive + + Attributes + ---------- + kind : InstructionKind + The kind of the instruction + inputs : List[INPUT_RV_TYPE] + The input random variables of the instruction, + and the type of each element can be one of the following: + - BlockRV + - LoopRV + - ExprRV + - float + - int + - str + - None + attrs : List[ATTR_TYPE] + The attributes of the instruction. Similar to attributes of an operator, + attributes of an instruction are arbitrary constant metadata required by the instructions. + For example, the name of the block to be retrieved in `GetBlock`. + outputs : List[OUTPUT_RV_TYPE] + The output random variables of the instruction, + and the type of each element can be one of the following: + - BlockRV + - LoopRV + - ExprRV, atomic variables only, won't be constants or composite PrimExpr + """ + + kind: InstructionKind + inputs: List[INPUT_RV_TYPE] + attrs: List[ATTR_TYPE] + outputs: List[OUTPUT_RV_TYPE] + + def __init__( + self, + kind: InstructionKind, + inputs: List[INPUT_RV_TYPE], + attrs: List[ATTR_TYPE], + outputs: List[OUTPUT_RV_TYPE], + ) -> None: + """Constructor + + Parameters + ---------- + kind : InstructionKind + The kind of the instruction + inputs : List[INPUT_RV_TYPE] + The input random variables of the instruction, + and the type of each element can be one of the following: + - BlockRV + - LoopRV + - ExprRV + - float + - int + - str + - None + attrs : List[ATTR_TYPE] + The attributes of the instruction. Similar to attributes of an operator, + attributes of an instruction are arbitrary constant metadata required by the + instructions. For example, the name of the block to be retrieved in `GetBlock`. + outputs : List[OUTPUT_RV_TYPE] + The output random variables of the instruction, + and the type of each element can be one of the following: + - BlockRV + - LoopRV + - ExprRV, atomic variables only, won't be constants or composite PrimExpr + """ + self.__init_handle_by_constructor__( + _ffi_api.Instruction, # type: ignore # pylint: disable=no-member + kind, + inputs, + attrs, + outputs, + ) diff --git a/python/tvm/tir/schedule/schedule.py b/python/tvm/tir/schedule/schedule.py index e3af8d3191a6..22c08398df33 100644 --- a/python/tvm/tir/schedule/schedule.py +++ b/python/tvm/tir/schedule/schedule.py @@ -16,15 +16,15 @@ # under the License. # pylint: disable=unused-import """The TensorIR schedule class""" -from typing import List, Optional, Union, Tuple +from typing import List, Optional, Union from tvm._ffi import register_object as _register_object from tvm.error import TVMError, register_error from tvm.ir import IRModule, PrimExpr from tvm.runtime import Object -from tvm.tir import Block, For, IntImm, PrimFunc, Var +from tvm.tir import Block, For, IntImm, PrimFunc -from . import _ffi_api_schedule +from . import _ffi_api from .state import ScheduleState, StmtSRef @@ -37,18 +37,33 @@ class ScheduleError(TVMError): class LoopRV(Object): """A random variable that refers to a loop""" + def __init__(self) -> None: + """Construct a new LoopRV.""" + self.__init_handle_by_constructor__( + _ffi_api.LoopRV # type: ignore # pylint: disable=no-member + ) + @_register_object("tir.BlockRV") class BlockRV(Object): """A random variable that refers to a block""" + def __init__(self) -> None: + """Construct a new BlockRV.""" + self.__init_handle_by_constructor__( + _ffi_api.BlockRV # type: ignore # pylint: disable=no-member + ) + # It is a workaround for mypy: https://github.com/python/mypy/issues/7866#issuecomment-549454370 # This feature is not supported until python 3.10: # https://docs.python.org/3.10/whatsnew/3.10.html#pep-613-typealias ExprRV = Union[PrimExpr] # A random variable that evaluates to an integer -RAND_VAR_TYPE = Union[ExprRV, BlockRV, LoopRV] # type: ignore # pylint: disable=invalid-name +RAND_VAR_TYPE = Union[ExprRV, BlockRV, LoopRV] # pylint: disable=invalid-name + +# Update to `Literal["detail", "fast", "none"]` once upgraded to python3.8 +ERROR_RENDER_LEVEL_CANDIDATES = Union[str] # pylint: disable=invalid-name @_register_object("tir.Schedule") @@ -66,20 +81,24 @@ class Schedule(Object): Link to tutorial: https://tvm.apache.org/docs/tutorials/language/schedule_primitives.html """ - ERROR_RENDER_LEVEL = {"detail": 0, "fast": 1, "none": 2} + ERROR_RENDER_LEVEL = { + "detail": 0, + "fast": 1, + "none": 2, + } def __init__( self, - func_or_mod: Union[PrimFunc, IRModule], + mod: Union[PrimFunc, IRModule], *, debug_mode: Union[bool, int] = False, - error_render_level: str = "detail", - ): + error_render_level: ERROR_RENDER_LEVEL_CANDIDATES = "detail", + ) -> None: """Construct a concrete TensorIR schedule from an IRModule or a PrimFunc Parameters ---------- - func_or_mod : Union[PrimFunc, IRModule] + mod : Union[PrimFunc, IRModule] The IRModule or PrimFunc to be scheduled debug_mode : Union[bool, int] Do extra correctness checking after the class creation and each time @@ -91,11 +110,13 @@ def __init__( "none": Do not show any error message. Note - ---------- + ---- The checks performed includes: 1) VerifySRefTree 2) VerifyCachedFlags """ + if isinstance(mod, PrimFunc): + mod = IRModule({"main": mod}) if isinstance(debug_mode, bool): if debug_mode: debug_mode = -1 @@ -108,12 +129,11 @@ def __init__( 'error_render_level can be "detail", "fast", or "none", but got: ' + f"{error_render_level}" ) - error_render_level = Schedule.ERROR_RENDER_LEVEL.get(error_render_level) # type: ignore self.__init_handle_by_constructor__( - _ffi_api_schedule.ConcreteSchedule, # type: ignore # pylint: disable=no-member - func_or_mod, + _ffi_api.ConcreteSchedule, # type: ignore # pylint: disable=no-member + mod, debug_mode, - error_render_level, + Schedule.ERROR_RENDER_LEVEL.get(error_render_level), ) ########## Utilities ########## @@ -121,12 +141,12 @@ def __init__( @property def mod(self) -> IRModule: """Returns the AST of the module being scheduled""" - return _ffi_api_schedule.ScheduleModule(self) # type: ignore # pylint: disable=no-member + return _ffi_api.ScheduleModule(self) # type: ignore # pylint: disable=no-member @property def state(self) -> ScheduleState: """Returns the ScheduleState in the current schedule class""" - return _ffi_api_schedule.ScheduleGetState(self) # type: ignore # pylint: disable=no-member + return _ffi_api.ScheduleGetState(self) # type: ignore # pylint: disable=no-member def copy(self) -> "Schedule": """Returns a copy of the schedule, including both the state and the symbol table, @@ -135,30 +155,34 @@ def copy(self) -> "Schedule": * 2) The IRModule being scheduled is untouched; * 3) All the random variables are valid in the copy, pointing to the correpsonding sref * reconstructed + Returns ------- copy : Schedule A new copy of the schedule """ - return _ffi_api_schedule.ScheduleCopy(self) # type: ignore # pylint: disable=no-member + return _ffi_api.ScheduleCopy(self) # type: ignore # pylint: disable=no-member def seed(self, seed: int) -> None: """Seed the randomness + Parameters ---------- seed : int The new random seed, -1 if use device random, otherwise non-negative """ - return _ffi_api_schedule.ScheduleSeed(self, seed) # type: ignore # pylint: disable=no-member + return _ffi_api.ScheduleSeed(self, seed) # type: ignore # pylint: disable=no-member def show(self, rand_var: RAND_VAR_TYPE) -> str: """Returns a string representation of the value that the random variable evaluates to + Parameters ---------- rand_var : Union[ExprRV, BlockRV, LoopRV] The random variable to be evaluated + Returns - ---------- + ------- str_repr : str The string representation """ @@ -176,18 +200,20 @@ def get( - the corresponding integer that a ExprRV evaluates to; - the corresponding Block that a block sref points to; - the corresponding For that a loop sref points to; + Parameters ---------- rand_var_or_sref : Union[ExprRV, BlockRV, LoopRV, StmtSRef] The random variable / sref to be evaluated + Returns - ---------- + ------- result : Optional[Union[int, Block, For]] The correpsonding result """ if isinstance(rand_var_or_sref, StmtSRef): return rand_var_or_sref.stmt - result = _ffi_api_schedule.ScheduleGet(self, rand_var_or_sref) # type: ignore # pylint: disable=no-member + result = _ffi_api.ScheduleGet(self, rand_var_or_sref) # type: ignore # pylint: disable=no-member if isinstance(result, IntImm): result = result.value return result @@ -198,49 +224,55 @@ def get_sref(self, rand_var_or_stmt: Union[BlockRV, LoopRV, Block, For]) -> Opti 2) BlockRV 3) Block 4) For + Parameters ---------- rand_var_or_stmt : Union[BlockRV, LoopRV, Block, For] The random variable / sref to be evaluated + Returns - ---------- + ------- result : Optional[StmtSRef] The correpsonding result """ - return _ffi_api_schedule.ScheduleGetSRef( # type: ignore # pylint: disable=no-member + return _ffi_api.ScheduleGetSRef( # type: ignore # pylint: disable=no-member self, rand_var_or_stmt ) def remove_rv(self, rand_var: RAND_VAR_TYPE) -> None: """Remove a random variable from the symbol table + Parameters ---------- rand_var : Union[BlockRV, LoopRV, ExprRV] The random variable to be removed """ - return _ffi_api_schedule.ScheduleRemoveRV(self, rand_var) # type: ignore # pylint: disable=no-member + return _ffi_api.ScheduleRemoveRV(self, rand_var) # type: ignore # pylint: disable=no-member - ########## Block/Loop relation ########## + ########## Schedule: Sampling ########## + ########## Schedule: Get blocks & loops ########## def get_block( self, name: str, func_name: str = "main", ) -> BlockRV: """Retrieve a block in a specific function with its name + Parameters ---------- name : str The name of the block func_name : str = "main" The name of the function + Returns - ---------- + ------- block : BlockRV The block retrieved IndexError is raised if 0 or multiple blocks exist with the specific name. """ - return _ffi_api_schedule.ScheduleGetBlock( # type: ignore # pylint: disable=no-member + return _ffi_api.ScheduleGetBlock( # type: ignore # pylint: disable=no-member self, name, func_name, @@ -248,18 +280,20 @@ def get_block( def get_loops(self, block: BlockRV) -> List[LoopRV]: """Get the parent loops of the block in its scope, from outer to inner + Parameters ---------- block : BlockRV The query block + Returns - ---------- + ------- loops : List[LoopRV] A list of loops above the given block in its scope, from outer to inner """ - return _ffi_api_schedule.ScheduleGetLoops(self, block) # type: ignore # pylint: disable=no-member + return _ffi_api.ScheduleGetLoops(self, block) # type: ignore # pylint: disable=no-member - ########## Schedule: loops manipulation ########## + ########## Schedule: Transform loops ########## def fuse(self, *loops: List[LoopRV]) -> LoopRV: """Fuse a list of consecutive loops into one. It requires: 1) The loops can't have annotations or thread bindings. @@ -272,7 +306,7 @@ def fuse(self, *loops: List[LoopRV]) -> LoopRV: The loops to be fused Returns - ---------- + ------- fused_loop : LoopRV The new loop after fusion @@ -316,7 +350,7 @@ def after_fuse(a: ty.handle, b: ty.handle) -> None: B[vi, vj] = A[vi, vj] * 2.0 """ - return _ffi_api_schedule.ScheduleFuse(self, loops) # type: ignore # pylint: disable=no-member + return _ffi_api.ScheduleFuse(self, loops) # type: ignore # pylint: disable=no-member def split( self, @@ -343,7 +377,7 @@ def split( - Nonnegative constant integers Returns - ---------- + ------- split_loops : List[LoopRV] The new loops after split @@ -389,9 +423,14 @@ def after_split(a: ty.handle, b: ty.handle) -> None: """ # it will be checked later in C++ implementation # that there is at most one None in `factors` - return _ffi_api_schedule.ScheduleSplit(self, loop, factors) # type: ignore # pylint: disable=no-member + return _ffi_api.ScheduleSplit(self, loop, factors) # type: ignore # pylint: disable=no-member + + ########## Schedule: Manipulate ForKind ########## + + ########## Schedule: Insert cache stages ########## + + ########## Schedule: Compute location ########## - ########## Schedule: compute location ########## def compute_inline(self, block: BlockRV) -> None: """Inline a block into its consumer(s). It requires: @@ -447,7 +486,7 @@ def after_inline(a: ty.handle, c: ty.handle) -> None: C[vi, vj] = A[vi, vj] * 2.0 + 1.0 """ - _ffi_api_schedule.ScheduleComputeInline(self, block) # type: ignore # pylint: disable=no-member + _ffi_api.ScheduleComputeInline(self, block) # type: ignore # pylint: disable=no-member def reverse_compute_inline(self, block: BlockRV) -> None: """Inline a block into its only producer. It requires: @@ -507,11 +546,10 @@ def after_inline(a: ty.handle, c: ty.handle) -> None: C[vi, vj] = A[vi, vj] * 2.0 + 1.0 """ - _ffi_api_schedule.ScheduleReverseComputeInline(self, block) # type: ignore # pylint: disable=no-member + _ffi_api.ScheduleReverseComputeInline(self, block) # type: ignore # pylint: disable=no-member + + ########## Schedule: Reduction ########## - ########## Schedule: loop binding/annotation ########## - ########## Schedule: cache read/write ########## - ########## Schedule: reduction ########## def rfactor(self, loop: LoopRV, factor_axis: int) -> LoopRV: """Factorize an associative reduction block by the specified loop. @@ -653,9 +691,17 @@ def after_rfactor(a: ty.handle, b: ty.handle) -> None: where `B` is the buffer that the reduction block writes to. Negative indexing is normalized according to numpy convention. """ - return _ffi_api_schedule.ScheduleRFactor(self, loop, factor_axis) # type: ignore # pylint: disable=no-member + return _ffi_api.ScheduleRFactor(self, loop, factor_axis) # type: ignore # pylint: disable=no-member + + ########## Schedule: Blockize & Tensorize ########## + + ########## Schedule: Annotation ########## + + ########## Schedule: Misc ########## - ########## Schedule: blockize & tensorize ########## + def enter_postproc(self) -> None: + """A no-op that marks the start of postprocessing phase of scheduling""" + _ffi_api.ScheduleEnterPostproc(self) # type: ignore # pylint: disable=no-member @_register_object("tir.ConcreteSchedule") diff --git a/python/tvm/tir/schedule/state.py b/python/tvm/tir/schedule/state.py index 845e1db5cb83..cc2415f150c9 100644 --- a/python/tvm/tir/schedule/state.py +++ b/python/tvm/tir/schedule/state.py @@ -24,7 +24,7 @@ from tvm.runtime import Object from tvm.tir import Block, BlockRealize, For, PrimFunc -from . import _ffi_api_schedule +from . import _ffi_api from .block_scope import BlockScope, StmtSRef CachedFlags = namedtuple("CachedFlags", ["affine_binding", "region_cover", "stage_pipeline"]) @@ -75,14 +75,14 @@ class ScheduleState(Object): def __init__( self, - func_or_mod: Union[PrimFunc, IRModule], + mod: Union[PrimFunc, IRModule], debug_mode: Union[bool, int] = False, - ): + ) -> None: """Construct a schedule state from an IRModule or a PrimFunc Parameters ---------- - func_or_mod : Union[PrimFunc, IRModule] + mod : Union[PrimFunc, IRModule] The IRModule or PrimFunc to be scheduled debug_mode : Union[bool, int] Do extra correctness checking after the class creation and each time @@ -92,6 +92,8 @@ def __init__( 2) False - Turn off all the checks 3) An integer - Turn on checks according to the bitmasks provided in ScheduleDebugMask """ + if isinstance(mod, PrimFunc): + mod = IRModule({"main": mod}) if isinstance(debug_mode, bool): if debug_mode: debug_mode = -1 @@ -100,8 +102,8 @@ def __init__( if not isinstance(debug_mode, int): raise TypeError(f"`debug_mode` should be integer or boolean, but gets: {debug_mode}") self.__init_handle_by_constructor__( - _ffi_api_schedule.ScheduleState, # type: ignore # pylint: disable=no-member - func_or_mod, + _ffi_api.ScheduleState, # type: ignore # pylint: disable=no-member + mod, debug_mode, ) @@ -118,7 +120,7 @@ def get_sref(self, stmt: Union[Block, For]) -> Optional[StmtSRef]: sref : StmtSRef The corresponding sref """ - return _ffi_api_schedule.ScheduleStateGetSRef(self, stmt) # type: ignore # pylint: disable=no-member + return _ffi_api.ScheduleStateGetSRef(self, stmt) # type: ignore # pylint: disable=no-member def get_block_scope(self, block_sref: StmtSRef) -> BlockScope: """Get the BlockScope correpsonding to the block sref @@ -133,7 +135,7 @@ def get_block_scope(self, block_sref: StmtSRef) -> BlockScope: sref : StmtSRef The corresponding sref """ - return _ffi_api_schedule.ScheduleStateGetBlockScope( # type: ignore # pylint: disable=no-member + return _ffi_api.ScheduleStateGetBlockScope( # type: ignore # pylint: disable=no-member self, block_sref ) @@ -151,14 +153,14 @@ def _get_cached_flags(self, block_sref: StmtSRef) -> CachedFlags: Three flags: affine_binding, region_cover, stage_pipeline Note - ------- + ---- It is an API intended for internal testing use. """ ( affine_binding, region_cover, stage_pipeline, - ) = _ffi_api_schedule.ScheduleStateGetCachedFlags( # type: ignore # pylint: disable=no-member + ) = _ffi_api.ScheduleStateGetCachedFlags( # type: ignore # pylint: disable=no-member self, block_sref ) return CachedFlags( @@ -199,12 +201,12 @@ def replace( the sref that points to the old block will point to the new one Note - ---------- + ---- The reuse of loop srefs are detected automatically according to the reuse of loop vars. """ if block_sref_reuse is None: block_sref_reuse = {} - _ffi_api_schedule.ScheduleStateReplace( # type: ignore # pylint: disable=no-member + _ffi_api.ScheduleStateReplace( # type: ignore # pylint: disable=no-member self, src_sref, tgt_stmt, diff --git a/python/tvm/tir/schedule/trace.py b/python/tvm/tir/schedule/trace.py new file mode 100644 index 000000000000..18bcca373dbb --- /dev/null +++ b/python/tvm/tir/schedule/trace.py @@ -0,0 +1,260 @@ +# 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""" +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional + +from tvm._ffi import register_object as _register_object +from tvm.runtime import Object + +from ...ir import Array, Map +from ...runtime import String +from ..expr import FloatImm, IntImm +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 + if isinstance(obj, Array): + return [_json_from_tvm(i) for i in obj] + if isinstance(obj, Map): + return {_json_from_tvm(k): _json_from_tvm(v) for k, v in obj.items()} + if isinstance(obj, String): + return str(obj) + if isinstance(obj, (IntImm, FloatImm)): + return obj.value + raise TypeError("Not supported type: " + str(type(obj))) + + +@_register_object("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) -> Optional[DECISION_TYPE]: + """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: Optional[DECISION_TYPE] = 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) -> Optional[Instruction]: + """Remove the last instruction, along with the decision made on that instruction, if any + + Returns + ------- + popped_inst : Instruction + Returns the instruction removed; 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: Optional[ + Callable[ + [Instruction, List[INPUT_RV_TYPE], List[ATTR_TYPE], DECISION_TYPE], DECISION_TYPE + ] + ] = 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 diff --git a/src/tir/schedule/analysis.h b/src/tir/schedule/analysis.h index 440c41246193..9baf4b5245ea 100644 --- a/src/tir/schedule/analysis.h +++ b/src/tir/schedule/analysis.h @@ -168,22 +168,6 @@ bool GetVarsTouchedByBlockIters(const BlockRealize& block_realize, std::unordered_set* reduce_vars); /******** Block-loop relation ********/ -/*! - * \brief Retrieves blocks in a specific function with its name - * \param self The schedule state - * \param name The name of the blocks to be retrieved - * \param func_name The name of the function - * \return A list of blocks with the specific name - */ -Array GetBlocks(const ScheduleState& self, const String& name, const String& func_name); - -/*! - * \brief Gets the parent loops of the block in its scope, from outer to inner - * \param self The schedule state - * \param block_sref The query block - * \return A list of loops above the given block in its scope, from outer to inner - */ -Array GetLoops(const StmtSRef& block_sref); /*! * \brief Gets StmtSRefs of leaf blocks of a scope where a specific block/loop is in diff --git a/src/tir/schedule/analysis/analysis.cc b/src/tir/schedule/analysis/analysis.cc index 31607a0b27b0..3ee98ec5b7d2 100644 --- a/src/tir/schedule/analysis/analysis.cc +++ b/src/tir/schedule/analysis/analysis.cc @@ -417,40 +417,6 @@ bool GetVarsTouchedByBlockIters(const BlockRealize& block_realize, /******** Block-loop relation ********/ -Array GetBlocks(const ScheduleState& self, const String& name, const String& func_name) { - struct Finder : public StmtVisitor { - explicit Finder(const ScheduleState& self, const String& name) : self_(self), name_(name) {} - - void VisitStmt_(const BlockNode* block) override { - if (block->name_hint == name_) { - auto it = self_->stmt2ref.find(block); - ICHECK(it != self_->stmt2ref.end()); - results_.push_back(it->second); - } - StmtVisitor::VisitStmt_(block); - } - - const ScheduleState& self_; - const String& name_; - Array results_; - }; - - BaseFunc func = self->mod->Lookup(func_name); - const auto* prim_func = TVM_TYPE_AS(prim_func, func, PrimFuncNode); - Finder finder(self, name); - finder(prim_func->body); - return std::move(finder.results_); -} - -Array GetLoops(const StmtSRef& block_sref) { - std::vector result; - for (StmtSRefNode* parent = block_sref->parent; parent && parent->stmt->IsInstance(); - parent = parent->parent) { - result.push_back(GetRef(parent)); - } - return {result.rbegin(), result.rend()}; -} - Array GetChildBlockSRefOnSRefTree(const ScheduleState& self, const StmtSRef& parent_sref) { Array child_block_realize = GetChildBlockRealizeOnSRefTree(parent_sref); diff --git a/src/tir/schedule/concrete_schedule.h b/src/tir/schedule/concrete_schedule.h index 5925cc59ded7..c44ec05d660b 100644 --- a/src/tir/schedule/concrete_schedule.h +++ b/src/tir/schedule/concrete_schedule.h @@ -76,25 +76,27 @@ class ConcreteScheduleNode : public ScheduleNode { using ScheduleNode::GetSRef; public: - /******** Block/Loop relation ********/ + /******** Schedule: Sampling ********/ + /******** Schedule: Get blocks & loops ********/ BlockRV GetBlock(const String& name, const String& func_name = "main") override; Array GetLoops(const BlockRV& block_rv) override; - /******** Schedule: loops manipulation ********/ + /******** Schedule: Transform loops ********/ LoopRV Fuse(const Array& loop_rvs) override; Array Split(const LoopRV& loop_rv, const Array>& factors) override; - /******** Schedule: compute location ********/ + /******** Schedule: Manipulate ForKind ********/ + /******** Schedule: Insert cache stages ********/ + /******** Schedule: Compute location ********/ void ComputeInline(const BlockRV& block) override; void ReverseComputeInline(const BlockRV& block) override; - /******** Schedule: loop binding/annotation ********/ - /******** Schedule: cache read/write ********/ - /******** Schedule: reduction ********/ - /******** Schedule: blockize & tensorize ********/ - - /******** Schedule: reduction ********/ + /******** Schedule: Reduction ********/ BlockRV RFactor(const LoopRV& loop_rv, int factor_axis) override; + /******** Schedule: Blockize & Tensorize ********/ + /******** Schedule: Annotation ********/ + /******** Schedule: Misc ********/ + void EnterPostproc() override {} - /******** Utility functions ********/ protected: + /******** Utility functions ********/ /*! * \brief Copy the schedule state, as well as the symbol table * \param new_state The ScheduleState copied diff --git a/src/tir/schedule/instruction.cc b/src/tir/schedule/instruction.cc new file mode 100644 index 000000000000..af721767c32f --- /dev/null +++ b/src/tir/schedule/instruction.cc @@ -0,0 +1,102 @@ +/* + * 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. + */ +#include "./utils.h" + +namespace tvm { +namespace tir { + +Instruction::Instruction(InstructionKind kind, Array inputs, Array attrs, + Array outputs) { + ObjectPtr n = make_object(); + n->kind = std::move(kind); + n->inputs = std::move(inputs); + n->attrs = std::move(attrs); + n->outputs = std::move(outputs); + this->data_ = std::move(n); +} + +using InstructionKindRegistry = AttrRegistry; + +InstructionKind InstructionKind::Get(const String& name) { + const InstructionKindRegEntry* reg = InstructionKindRegistry::Global()->Get(name); + ICHECK(reg != nullptr) << "AttributeError: Instruction kind " << name << " is not registered"; + return reg->inst_kind_; +} + +InstructionKindRegEntry::InstructionKindRegEntry(uint32_t reg_index) { + this->inst_kind_ = InstructionKind(make_object()); +} + +InstructionKindRegEntry& InstructionKindRegEntry::RegisterOrGet(const String& name) { + return InstructionKindRegistry::Global()->RegisterOrGet(name); +} + +/**************** Repr ****************/ + +TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable) + .set_dispatch([](const ObjectRef& obj, ReprPrinter* p) { + const auto* self = obj.as(); + ICHECK_NOTNULL(self); + Array inputs; + inputs.reserve(self->inputs.size()); + for (const ObjectRef& obj : self->inputs) { + if (!obj.defined()) { + inputs.push_back(String("None")); + } else if (obj->IsInstance() || obj->IsInstance()) { + inputs.push_back(String("_")); + } else if (const auto* str_obj = obj.as()) { + inputs.push_back(String('"' + std::string(str_obj->data) + '"')); + } else if (obj->IsInstance() || obj->IsInstance()) { + inputs.push_back(obj); + } else if (const auto* expr = obj.as()) { + PrimExpr new_expr = + Substitute(GetRef(expr), [](const Var& var) -> Optional { + ObjectPtr new_var = make_object(*var.get()); + new_var->name_hint = "_"; + return Var(new_var); + }); + std::ostringstream os; + os << new_expr; + inputs.push_back(String(os.str())); + } else { + LOG(FATAL) << "TypeError: Stringifying is not supported for type: " << obj->GetTypeKey(); + throw; + } + } + p->stream << self->kind->f_as_python( + /*inputs=*/inputs, + /*attrs=*/self->attrs, + /*decision=*/NullOpt, + /*outputs=*/Array(self->outputs.size(), String("_"))); + }); + +/**************** FFI ****************/ + +TVM_REGISTER_NODE_TYPE(InstructionNode); +TVM_REGISTER_NODE_TYPE(InstructionKindNode); + +TVM_REGISTER_GLOBAL("tir.schedule.InstructionKindGet").set_body_typed(InstructionKind::Get); +TVM_REGISTER_GLOBAL("tir.schedule.Instruction") + .set_body_typed([](InstructionKind kind, Array inputs, Array attrs, + Array outputs) -> Instruction { + return Instruction(kind, inputs, attrs, outputs); + }); + +} // namespace tir +} // namespace tvm diff --git a/src/tir/schedule/instruction_traits.h b/src/tir/schedule/instruction_traits.h new file mode 100644 index 000000000000..95d636467aa0 --- /dev/null +++ b/src/tir/schedule/instruction_traits.h @@ -0,0 +1,536 @@ +/* + * 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. + */ +#ifndef TVM_TIR_SCHEDULE_INSTRUCTION_TRAITS_H_ +#define TVM_TIR_SCHEDULE_INSTRUCTION_TRAITS_H_ + +#include +#include + +#include +#include +#include + +namespace tvm { +namespace tir { + +/*! + * \brief Register an InstructionKind using a trait class + * \param InstructionKindTraits A traits class of an InstructionKind + * + * Example: + * + * \code + * + * struct SomeInstructionKindTraits { + * static constexpr const char* kName = "name-of-the-instruction"; + * static constexpr bool kIsPure = false; + * + * // Convertible to `InstructionKindNode::FInstructionApply` + * static Array ApplyToSchedule( + * const tir::Schedule& sch, + * const Array& inputs, + * const Array& attrs, + * const Optional& decision); + * + * // Convertible to `InstructionKindNode::FInstructionAsPython` + * static String AsPython( + * const Array& inputs, + * const Array& attrs, + * const Optional& decision, + * const Array& outputs); + * + * // Convertible to `InstructionKindNode::FInstructionAttrsAsJSON` + * static Array AttrsAsJSON( + * const Array& attrs); + * + * // Convertible to `InstructionKindNode::FInstructionAttrsFromJSON` + * static Array AttrsFromJSON( + * const Array& attrs_record); + * }; + * + * TVM_REGISTER_INST_KIND_TRAITS(SomeInstructionKindTraits); + * + * \endcode + */ +#define TVM_REGISTER_INST_KIND_TRAITS(InstructionKindTraits) \ + TVM_REGISTER_INST_KIND(InstructionKindTraits::kName) \ + .set_is_pure(InstructionKindTraits::kIsPure) \ + .set_apply_to_schedule(InstructionKindTraits::ApplyToSchedule) \ + .set_attrs_as_json(InstructionKindTraits::AttrsAsJSON) \ + .set_attrs_from_json(InstructionKindTraits::AttrsFromJSON) \ + .set_as_python(InstructionKindTraits::AsPython) + +/*! + * \brief A helper to conveniently register an InstructionKind. When inherited in curiously + * recursive template pattern, the derived class `TTraits` only needs to define two functions on the + * unpacked inputs, and the helper handles unpacking and downcasting. See the example for more + * details. + * + * \tparam TTraits The derived class + * + * Example: + * + * \code + * + * struct SamplePerfectTileTraits : public UnpackedInstTraits { + * // The name of this kind of instruction + * static constexpr const char* kName = "SamplePerfectTile"; + * // A boolean indicating if the instruction is pure, i.e. change nothing in the schedule state + * static constexpr bool kIsPure = true; + * // The number of inputs in this kind of instruction + * static constexpr size_t kNumInputs = 1; + * // The number of attributes in this kind of instruction + * static constexpr size_t kNumAttrs = 2; + * // The number of decisions in this kind of instruction (only 0 or 1 is allowed) + * static constexpr size_t kNumDecisions = 1; + * + * // Calling convention: + * // - All the arguments must be ObjectRef + * // - The 1st argument is Schedule + * // - The next `kNumInputs` arguments are input random variables + * // - The next `kNumAttrs` arguments are attributes + * // - The next argument is decision, if `kNumDecisions == 1` + * static Array UnpackedApplyToSchedule( + * Schedule sch, + * LoopRV loop_rv, + * Integer n, + * Integer max_innermost_factor, + * Optional> decision) { + * return sch->SamplePerfectTile(loop_rv, n->value, max_innermost_factor->value, decision); + * } + * + * // Calling convention: + * // - All the arguments must be ObjectRef + * // - The 1st argument is an array containing names of output random variables + * // - The next `kNumInputs` arguments are names of input random variables + * // - The next `kNumAttrs` arguments are attributes + * // - The next argument is decision, if `kNumDecisions == 1` + * static String UnpackedAsPython( + * Array outputs, + * String loop_rv, + * Integer n, + * Integer max_innermost_factor, + * Optional> decision) { + * PythonAPICall py("sample_perfect_tile"); + * py.Input("loop", loop_rv); + * py.Input("n", n->value); + * py.Input("max_innermost_factor", max_innermost_factor->value); + * py.Decision(decision); + * py.OutputList(outputs); + * return py.Str(); + * } + * + * template + * friend struct UnpackedInstTraits; + * }; + * + * TVM_REGISTER_INST_KIND(SamplePerfectTileTraits); + * \endcode + */ +template +struct UnpackedInstTraits { + /*! + * \brief Unpack the arguments in the calling convention, and feed them into + * `TTraits::UnpackedApplyToSchedule` + * \sa InstructionKindNode::f_apply_to_schedule + */ + static Array ApplyToSchedule(const Schedule& sch, const Array& inputs, + const Array& attrs, + const Optional& decision); + + /*! + * \brief Unpack the arguments in the calling convention, and feed them into + * `TTraits::UnpackedAsPython` + * \sa InstructionKindNode::f_as_python + */ + static String AsPython(const Array& inputs, const Array& attrs, + const Optional& decision, const Array& outputs); + + /*! \brief No customized serializer by default */ + static constexpr std::nullptr_t AttrsAsJSON = nullptr; + + /*! \brief No customized deserializer by default */ + static constexpr std::nullptr_t AttrsFromJSON = nullptr; + + protected: + template + static TVM_ALWAYS_INLINE void _SetInputs(const runtime::TVMArgsSetter& setter, + const Array& inputs); + template + static TVM_ALWAYS_INLINE void _SetAttrs(const runtime::TVMArgsSetter& setter, + const Array& attrs); + template + static TVM_ALWAYS_INLINE void _SetDecision(const runtime::TVMArgsSetter& setter, + const Optional& decision); + static TVM_ALWAYS_INLINE Array _ConvertOutputs(const TVMRetValue& rv); +}; + +/*! + * \brief A helper class that constructs schedule API call in python syntax, + * which helps convert an Inst to a python statement. + * \sa InstructionKindNode::f_as_python + */ +class PythonAPICall { + public: + /*! + * \brief Constructor + * \param method_name The name of the schedule API to be called + */ + explicit PythonAPICall(String method_name) : method_name_(method_name), output_(NullOpt) {} + /*! \brief Add an intger input */ + inline void Input(String arg_name, int arg); + /*! \brief Add an intger input */ + inline void Input(String arg_name, int64_t arg); + /*! \brief Add a double input */ + inline void Input(String arg_name, double arg); + /*! \brief Add an input random variable */ + inline void Input(String arg_name, String arg); + /*! \brief Add an input, dispatched to different implementations according to the object's type */ + inline void Input(String arg_name, ObjectRef arg); + /*! \brief Add the decision */ + inline void Decision(ObjectRef decision); + /*! + * \brief Add a single output random variable + * \param unit_array An array containing only one element + */ + inline void SingleOutput(Array unit_array); + /*! \brief Add a list of output random variables */ + inline void OutputList(Array outputs); + /*! \returns The schedule API call in python syntax */ + inline String Str() const; + + private: + /*! \brief Converts a TVM object to python string and print to the output stream */ + inline void AsPythonString(const ObjectRef& obj, std::ostream& os); + + private: + /*! \brief The name of the API to call */ + String method_name_; + /*! \brief The output of the instruction */ + Optional output_; + /*! \brief The names of input arguments */ + std::vector arg_names_; + /*! \brief The values of input arguments */ + std::vector args_; +}; + +/********** implementation details **********/ + +// forward declaration +namespace details { + +template +struct _ArgsPacker; + +template <> +struct _ArgsPacker<> { + static constexpr bool checked = true; +}; + +template +struct _ArgsPacker { + static constexpr bool checked = + std::is_base_of::value && _ArgsPacker::checked; +}; + +template +struct _MethodType {}; + +template +struct _MethodType { + using return_type = TReturn; + using argument_type = _ArgsPacker; +}; + +template +struct _NumArgs {}; + +template +struct _NumArgs { + static constexpr size_t value = sizeof...(Args); +}; + +template +struct _IsTVMArray : std::false_type {}; + +template +struct _IsTVMArray> : std::true_type {}; + +template +struct _IsSingleObject + : std::integral_constant::value && !_IsTVMArray::value> { +}; + +template +using ReturnType = typename _MethodType>::return_type; + +template +static constexpr bool ArgumentAreAllObjects = + _MethodType>::argument_type::checked; + +template +static constexpr size_t NumArgs = _NumArgs>::value; + +template +static constexpr int IsTVMArray = _IsTVMArray>::value; + +template +static constexpr int IsSingleObject = _IsSingleObject>::value; + +}; // namespace details + +template +Array UnpackedInstTraits::ApplyToSchedule(const Schedule& sch, + const Array& inputs, + const Array& attrs, + const Optional& decision) { + using method_type = decltype(TTraits::UnpackedApplyToSchedule); + using return_type = details::ReturnType; + static_assert(details::ArgumentAreAllObjects, + "All arguments to `UnpackedApplyToSchedule` must be subclasses of ObjectRef"); + constexpr size_t kNumArgs = details::NumArgs; + constexpr size_t kNumInputs = TTraits::kNumInputs; + constexpr size_t kNumAttrs = TTraits::kNumAttrs; + constexpr size_t kNumDecisions = TTraits::kNumDecisions; + static_assert(kNumArgs == 1 + kNumInputs + kNumAttrs + kNumDecisions, + "length of argument list mismatch"); + TVMValue tvm_values[kNumArgs]; + int tvm_type_codes[kNumArgs]; + runtime::TVMArgsSetter setter(tvm_values, tvm_type_codes); + setter(0, sch); + TTraits::template _SetInputs<1>(setter, inputs); + TTraits::template _SetAttrs<1 + kNumInputs>(setter, attrs); + TTraits::template _SetDecision<1 + kNumInputs + kNumAttrs>(setter, decision); + PackedFunc pf([](const TVMArgs& args, TVMRetValue* rv) -> void { + using runtime::detail::unpack_call; + constexpr size_t kNumArgs = details::NumArgs; + ICHECK_EQ(args.size(), kNumArgs); + unpack_call(nullptr, TTraits::UnpackedApplyToSchedule, args, rv); + }); + TVMRetValue rv; + pf.CallPacked(TVMArgs(tvm_values, tvm_type_codes, kNumArgs), &rv); + return TTraits::_ConvertOutputs(rv); +} + +template +String UnpackedInstTraits::AsPython(const Array& inputs, + const Array& attrs, + const Optional& decision, + const Array& outputs) { + using method_type = decltype(TTraits::UnpackedAsPython); + using return_type = details::ReturnType; + static_assert(details::ArgumentAreAllObjects, + "All arguments to `UnpackedAsPython` must be subclasses of ObjectRef"); + constexpr size_t kNumArgs = details::NumArgs; + constexpr size_t kNumInputs = TTraits::kNumInputs; + constexpr size_t kNumAttrs = TTraits::kNumAttrs; + constexpr size_t kNumDecisions = TTraits::kNumDecisions; + static_assert(kNumArgs == 1 + kNumInputs + kNumAttrs + kNumDecisions, + "length of argument list mismatch"); + TVMValue tvm_values[kNumArgs]; + int tvm_type_codes[kNumArgs]; + runtime::TVMArgsSetter setter(tvm_values, tvm_type_codes); + setter(0, outputs); + TTraits::template _SetInputs<1>(setter, inputs); + TTraits::template _SetAttrs<1 + kNumInputs>(setter, attrs); + TTraits::template _SetDecision<1 + kNumInputs + kNumAttrs>(setter, decision); + PackedFunc pf([](const TVMArgs& args, TVMRetValue* rv) -> void { + using runtime::detail::unpack_call; + constexpr size_t kNumArgs = details::NumArgs; + ICHECK_EQ(args.size(), kNumArgs); + unpack_call(nullptr, TTraits::UnpackedAsPython, args, rv); + }); + TVMRetValue rv; + pf.CallPacked(TVMArgs(tvm_values, tvm_type_codes, kNumArgs), &rv); + String result = rv; + return result; +} + +template +template +TVM_ALWAYS_INLINE void UnpackedInstTraits::_SetInputs(const runtime::TVMArgsSetter& setter, + const Array& inputs) { + constexpr size_t kNumInputs = TTraits::kNumInputs; + ICHECK_EQ(kNumInputs, inputs.size()) + << "ValueError: Incorrect kNumInputs for instruction: " << TTraits::kName; + const ObjectRef* ptr = inputs.template as()->begin(); + for (size_t i = 0; i < kNumInputs; ++i) { + setter(i + index_offset, *(ptr + i)); + } +} + +template +template +TVM_ALWAYS_INLINE void UnpackedInstTraits::_SetAttrs(const runtime::TVMArgsSetter& setter, + const Array& attrs) { + constexpr size_t kNumAttrs = TTraits::kNumAttrs; + ICHECK_EQ(kNumAttrs, attrs.size()) + << "ValueError: Incorrect kNumAttrs for instruction: " << TTraits::kName; + const ObjectRef* ptr = attrs.as()->begin(); + for (size_t i = 0; i < kNumAttrs; ++i) { + setter(i + index_offset, *(ptr + i)); + } +} + +template +template +TVM_ALWAYS_INLINE void UnpackedInstTraits::_SetDecision( + const runtime::TVMArgsSetter& setter, const Optional& decision) { + constexpr size_t kNumDecisions = TTraits::kNumDecisions; + static_assert(kNumDecisions <= 1, "an instruction is supposed to have at most 1 decision"); + if (kNumDecisions == 1) { + setter(index_offset, decision); + } else { + ICHECK(!decision.defined()); + } +} + +template +TVM_ALWAYS_INLINE Array UnpackedInstTraits::_ConvertOutputs( + const TVMRetValue& rv) { + using method_type = decltype(TTraits::UnpackedApplyToSchedule); + using return_type = details::ReturnType; + constexpr int is_array = details::IsTVMArray; + constexpr int is_single_obj = details::IsSingleObject; + constexpr int is_void = std::is_void::value; + static_assert(is_array || is_single_obj || is_void, "return type not supported"); + static_assert(is_array + is_single_obj + is_void == 1, "internal template error"); + if (is_void) { + return {}; + } else if (is_single_obj) { + ObjectRef obj = rv; + return {obj}; + } else if (is_array) { + ObjectRef obj = rv; + const ArrayNode* array = obj.as(); + return GetRef>(array); + } +} + +/********** PythonAPICall **********/ + +inline void PythonAPICall::AsPythonString(const ObjectRef& obj, std::ostream& os) { + if (const auto* str = obj.as()) { + os << str->data; + } else if (const auto* int_imm = obj.as()) { + os << int_imm->value; + } else if (const auto* float_imm = obj.as()) { + os.precision(17); + os << float_imm->value; + } else if (const auto* array = obj.as()) { + os << '['; + bool is_first = true; + for (const ObjectRef& e : *array) { + if (is_first) { + is_first = false; + } else { + os << ", "; + } + AsPythonString(e, os); + } + os << ']'; + } else { + LOG(FATAL) << "ValueError: Cannot translate type '" << obj->GetTypeKey() + << "' to python. Its value is: " << obj; + throw; + } +} + +void PythonAPICall::Input(String arg_name, int arg) { + arg_names_.emplace_back(std::move(arg_name)); + args_.push_back(std::to_string(arg)); +} + +void PythonAPICall::Input(String arg_name, int64_t arg) { + arg_names_.emplace_back(std::move(arg_name)); + args_.push_back(std::to_string(arg)); +} + +void PythonAPICall::Input(String arg_name, double arg) { + arg_names_.emplace_back(std::move(arg_name)); + std::ostringstream os; + os.precision(17); + os << arg; + args_.push_back(os.str()); +} + +void PythonAPICall::Input(String arg_name, String arg) { + arg_names_.emplace_back(std::move(arg_name)); + args_.emplace_back(std::move(arg)); +} + +void PythonAPICall::Input(String arg_name, ObjectRef arg) { + arg_names_.emplace_back(std::move(arg_name)); + std::ostringstream os; + AsPythonString(arg, os); + args_.push_back(os.str()); +} + +void PythonAPICall::Decision(ObjectRef decision) { + if (decision.defined()) { + this->Input("decision", decision); + } +} + +void PythonAPICall::SingleOutput(Array unit_array) { + ICHECK_EQ(unit_array.size(), 1); + this->output_ = unit_array[0]; +} + +void PythonAPICall::OutputList(Array outputs) { + if (outputs.empty()) { + return; + } + if (outputs.size() == 1) { + this->output_ = outputs[0] + ","; + return; + } + std::ostringstream os; + os << outputs[0]; + for (int i = 1, n = outputs.size(); i < n; ++i) { + os << ", " << outputs[i]; + } + this->output_ = os.str(); +} + +String PythonAPICall::Str() const { + std::ostringstream os; + if (output_.defined()) { + os << output_.value() << " = "; + } + os << "sch." << method_name_ << '('; + int n = args_.size(); + for (int i = 0; i < n; ++i) { + if (i > 0) { + os << ", "; + } + if (arg_names_[i].empty()) { + os << args_[i]; + } else { + os << arg_names_[i] << '=' << args_[i]; + } + } + os << ')'; + return os.str(); +} + +} // namespace tir +} // namespace tvm + +#endif // TVM_TIR_SCHEDULE_INSTRUCTION_TRAITS_H_ diff --git a/src/tir/schedule/primitive.h b/src/tir/schedule/primitive.h index cf96c4362422..22e25f1c54a7 100644 --- a/src/tir/schedule/primitive.h +++ b/src/tir/schedule/primitive.h @@ -24,7 +24,25 @@ namespace tvm { namespace tir { -/******** Schedule: loops manipulation ********/ +/******** Schedule: Sampling ********/ +/******** Schedule: Get blocks & loops ********/ +/*! + * \brief Retrieves blocks in a specific function with its name + * \param self The schedule state + * \param name The name of the blocks to be retrieved + * \param func_name The name of the function + * \return A list of blocks with the specific name + */ +Array GetBlocks(const ScheduleState& self, const String& name, const String& func_name); +/*! + * \brief Gets the parent loops of the block in its scope, from outer to inner + * \param self The schedule state + * \param block_sref The query block + * \return A list of loops above the given block in its scope, from outer to inner + */ +Array GetLoops(const StmtSRef& block_sref); +/******** Schedule: Transform loops ********/ + /*! * Split a loop into a list of consecutive loops. It requires: * 1) The loop can't have annotation or thread binding. @@ -46,7 +64,9 @@ TVM_DLL Array Split(ScheduleState self, const StmtSRef& loop_sref, * \return The sref to the fused loop */ TVM_DLL StmtSRef Fuse(ScheduleState self, const Array& loop_srefs); -/******** Schedule: compute location ********/ +/******** Schedule: Manipulate ForKind ********/ +/******** Schedule: Insert cache stages ********/ +/******** Schedule: Compute location ********/ /*! * \brief Inline a block into its consumer(s). It requires: * 1) The block is a complete non-root block, which only produces one buffer @@ -72,12 +92,7 @@ TVM_DLL void ComputeInline(ScheduleState self, const StmtSRef& block_sref); * \param block_sref The sref to the block to be inlined to its producer */ TVM_DLL void ReverseComputeInline(ScheduleState self, const StmtSRef& block_sref); - -/******** Schedule: loop binding/annotation ********/ - -/******** Schedule: cache read/write ********/ - -/******** Schedule: reduction ********/ +/******** Schedule: Reduction ********/ /*! * \brief Factor a reduction block by the specified loop * \details See python/tvm/tir/schedule/schedule.py @@ -89,8 +104,9 @@ TVM_DLL void ReverseComputeInline(ScheduleState self, const StmtSRef& block_sref * \return The sref of the rfactor block */ TVM_DLL StmtSRef RFactor(ScheduleState self, const StmtSRef& loop_sref, int factor_axis); - -/******** Schedule: blockize & tensorize ********/ +/******** Schedule: Blockize & Tensorize ********/ +/******** Schedule: Annotation ********/ +/******** Schedule: Misc ********/ } // namespace tir } // namespace tvm diff --git a/src/tir/schedule/primitive/compute_inline.cc b/src/tir/schedule/primitive/compute_inline.cc index 3892f358e0ec..2583b21227e4 100644 --- a/src/tir/schedule/primitive/compute_inline.cc +++ b/src/tir/schedule/primitive/compute_inline.cc @@ -675,5 +675,56 @@ void ReverseComputeInline(ScheduleState self, const StmtSRef& consumer_block_sre self->Replace(scope_root_sref, tgt_stmt, inliner.block_reuse); } +/******** Instruction Registration ********/ + +struct ComputeInlineTraits : public UnpackedInstTraits { + static constexpr const char* kName = "ComputeInline"; + static constexpr bool kIsPure = false; + + private: + static constexpr size_t kNumInputs = 1; + static constexpr size_t kNumAttrs = 0; + static constexpr size_t kNumDecisions = 0; + + static void UnpackedApplyToSchedule(Schedule sch, BlockRV block_rv) { + return sch->ComputeInline(block_rv); + } + + static String UnpackedAsPython(Array outputs, String block_rv) { + PythonAPICall py("compute_inline"); + py.Input("block", block_rv); + return py.Str(); + } + + template + friend struct ::tvm::tir::UnpackedInstTraits; +}; + +struct ReverseComputeInlineTraits : public UnpackedInstTraits { + static constexpr const char* kName = "ReverseComputeInline"; + static constexpr bool kIsPure = false; + + private: + static constexpr size_t kNumInputs = 1; + static constexpr size_t kNumAttrs = 0; + static constexpr size_t kNumDecisions = 0; + + static void UnpackedApplyToSchedule(Schedule sch, BlockRV block_rv) { + return sch->ReverseComputeInline(block_rv); + } + + static String UnpackedAsPython(Array outputs, String block_rv) { + PythonAPICall py("reverse_compute_inline"); + py.Input("block", block_rv); + return py.Str(); + } + + template + friend struct ::tvm::tir::UnpackedInstTraits; +}; + +TVM_REGISTER_INST_KIND_TRAITS(ComputeInlineTraits); +TVM_REGISTER_INST_KIND_TRAITS(ReverseComputeInlineTraits); + } // namespace tir } // namespace tvm diff --git a/src/tir/schedule/primitive/get_block_loop.cc b/src/tir/schedule/primitive/get_block_loop.cc new file mode 100644 index 000000000000..a8d9c5a69dc9 --- /dev/null +++ b/src/tir/schedule/primitive/get_block_loop.cc @@ -0,0 +1,113 @@ +/* + * 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. + */ +#include "../utils.h" + +namespace tvm { +namespace tir { + +Array GetBlocks(const ScheduleState& self, const String& name, const String& func_name) { + struct Finder : public StmtVisitor { + explicit Finder(const ScheduleState& self, const String& name) : self_(self), name_(name) {} + + void VisitStmt_(const BlockNode* block) override { + if (block->name_hint == name_) { + auto it = self_->stmt2ref.find(block); + ICHECK(it != self_->stmt2ref.end()); + results_.push_back(it->second); + } + StmtVisitor::VisitStmt_(block); + } + + const ScheduleState& self_; + const String& name_; + Array results_; + }; + + BaseFunc func = self->mod->Lookup(func_name); + const auto* prim_func = TVM_TYPE_AS(prim_func, func, PrimFuncNode); + Finder finder(self, name); + finder(prim_func->body); + return std::move(finder.results_); +} + +Array GetLoops(const StmtSRef& block_sref) { + std::vector result; + for (StmtSRefNode* parent = block_sref->parent; parent && parent->stmt->IsInstance(); + parent = parent->parent) { + result.push_back(GetRef(parent)); + } + return {result.rbegin(), result.rend()}; +} + +/******** Instruction Registration ********/ + +struct GetBlockTraits : public UnpackedInstTraits { + static constexpr const char* kName = "GetBlock"; + static constexpr bool kIsPure = true; + + private: + static constexpr size_t kNumInputs = 0; + static constexpr size_t kNumAttrs = 2; + static constexpr size_t kNumDecisions = 0; + + static BlockRV UnpackedApplyToSchedule(Schedule sch, String name, String func_name) { + return sch->GetBlock(name, func_name); + } + + static String UnpackedAsPython(Array outputs, String name, String func_name) { + PythonAPICall py("get_block"); + py.Input("name", name); + py.Input("func_name", func_name); + py.SingleOutput(outputs); + return py.Str(); + } + + template + friend struct ::tvm::tir::UnpackedInstTraits; +}; + +struct GetLoopsTraits : public UnpackedInstTraits { + static constexpr const char* kName = "GetLoops"; + static constexpr bool kIsPure = true; + + private: + static constexpr size_t kNumInputs = 1; + static constexpr size_t kNumAttrs = 0; + static constexpr size_t kNumDecisions = 0; + + static Array UnpackedApplyToSchedule(Schedule sch, BlockRV block_rv) { + return sch->GetLoops(block_rv); + } + + static String UnpackedAsPython(Array outputs, String block_rv) { + PythonAPICall py("get_loops"); + py.Input("block", block_rv); + py.OutputList(outputs); + return py.Str(); + } + + template + friend struct ::tvm::tir::UnpackedInstTraits; +}; + +TVM_REGISTER_INST_KIND_TRAITS(GetBlockTraits); +TVM_REGISTER_INST_KIND_TRAITS(GetLoopsTraits); + +} // namespace tir +} // namespace tvm diff --git a/src/tir/schedule/primitive/loop_transformation.cc b/src/tir/schedule/primitive/loop_transformation.cc index 2a2d9ed2a888..d1875df61ac7 100644 --- a/src/tir/schedule/primitive/loop_transformation.cc +++ b/src/tir/schedule/primitive/loop_transformation.cc @@ -385,5 +385,79 @@ StmtSRef Fuse(ScheduleState self, const Array& loop_srefs) { return self->stmt2ref.at(new_stmt.get()); } +/******** Instruction Registration ********/ + +struct SplitTraits : public UnpackedInstTraits { + static constexpr const char* kName = "Split"; + static constexpr bool kIsPure = false; + + private: + static constexpr size_t kNumInputs = 2; + static constexpr size_t kNumAttrs = 0; + static constexpr size_t kNumDecisions = 0; + + template + static TVM_ALWAYS_INLINE void _SetInputs(const runtime::TVMArgsSetter& setter, + const Array& inputs) { + thread_local ObjectRef loop_rv{nullptr}; + thread_local Array factors{nullptr}; + loop_rv = inputs[0]; + factors = Array{inputs.begin() + 1, inputs.end()}; + setter(delta, loop_rv); + setter(delta + 1, factors); + } + + static Array UnpackedApplyToSchedule(Schedule sch, LoopRV loop_rv, + Array> factors) { + return sch->Split(loop_rv, factors); + } + + static String UnpackedAsPython(Array outputs, String loop_rv, Array factors) { + PythonAPICall py("split"); + py.Input("loop", loop_rv); + py.Input("factors", factors); + py.OutputList(outputs); + return py.Str(); + } + + template + friend struct ::tvm::tir::UnpackedInstTraits; +}; + +struct FuseTraits : public UnpackedInstTraits { + static constexpr const char* kName = "Fuse"; + static constexpr bool kIsPure = false; + + private: + static constexpr size_t kNumInputs = 1; + static constexpr size_t kNumAttrs = 0; + static constexpr size_t kNumDecisions = 0; + + template + static TVM_ALWAYS_INLINE void _SetInputs(const runtime::TVMArgsSetter& setter, + const Array& inputs) { + setter(delta, inputs); + } + + static LoopRV UnpackedApplyToSchedule(Schedule sch, Array loop_rvs) { + return sch->Fuse(loop_rvs); + } + + static String UnpackedAsPython(Array outputs, Array loop_rvs) { + PythonAPICall py("fuse"); + for (const String& loop_rv : loop_rvs) { + py.Input("", loop_rv); + } + py.SingleOutput(outputs); + return py.Str(); + } + + template + friend struct ::tvm::tir::UnpackedInstTraits; +}; + +TVM_REGISTER_INST_KIND_TRAITS(SplitTraits); +TVM_REGISTER_INST_KIND_TRAITS(FuseTraits); + } // namespace tir } // namespace tvm diff --git a/src/tir/schedule/primitive/reduction.cc b/src/tir/schedule/primitive/reduction.cc index a4b07964d42c..bf29ceb1ef9f 100644 --- a/src/tir/schedule/primitive/reduction.cc +++ b/src/tir/schedule/primitive/reduction.cc @@ -952,6 +952,35 @@ StmtSRef RFactor(ScheduleState self, const StmtSRef& rf_loop_sref, int factor_ax return new_block_srefs[0]; } +/******** Instruction Registration ********/ + +struct RFactorTraits : public UnpackedInstTraits { + static constexpr const char* kName = "RFactor"; + static constexpr bool kIsPure = false; + + private: + static constexpr size_t kNumInputs = 1; + static constexpr size_t kNumAttrs = 1; + static constexpr size_t kNumDecisions = 0; + + static BlockRV UnpackedApplyToSchedule(Schedule sch, LoopRV loop_rv, Integer factor_axis) { + return sch->RFactor(loop_rv, factor_axis->value); + } + + static String UnpackedAsPython(Array outputs, String loop_rv, Integer factor_axis) { + PythonAPICall py("rfactor"); + py.Input("loop", loop_rv); + py.Input("factor_axis", factor_axis->value); + py.SingleOutput(outputs); + return py.Str(); + } + + template + friend struct ::tvm::tir::UnpackedInstTraits; +}; + +TVM_REGISTER_INST_KIND_TRAITS(RFactorTraits); + /******** FFI ********/ TVM_REGISTER_GLOBAL("tir.schedule.RegisterReducer") diff --git a/src/tir/schedule/schedule.cc b/src/tir/schedule/schedule.cc index eae04bc76d55..eda6ac27d283 100644 --- a/src/tir/schedule/schedule.cc +++ b/src/tir/schedule/schedule.cc @@ -16,7 +16,7 @@ * specific language governing permissions and limitations * under the License. */ -#include +#include "./utils.h" namespace tvm { namespace tir { @@ -55,17 +55,10 @@ TVM_REGISTER_GLOBAL("tir.schedule.ScheduleCopy") // /**************** (FFI) Constructor ****************/ +TVM_REGISTER_GLOBAL("tir.schedule.BlockRV").set_body_typed([]() { return BlockRV(); }); +TVM_REGISTER_GLOBAL("tir.schedule.LoopRV").set_body_typed([]() { return LoopRV(); }); TVM_REGISTER_GLOBAL("tir.schedule.ConcreteSchedule") - .set_body_typed([](ObjectRef obj, int debug_mode, int error_render_level) -> Schedule { - IRModule mod{nullptr}; - if (const auto* func = obj.as()) { - mod = IRModule({{GlobalVar("main"), GetRef(func)}}); - } else if (const auto* p_mod = obj.as()) { - mod = GetRef(p_mod); - } else { - LOG(FATAL) << "TypeError: Expects `IRModule` or `PrimFunc`, but gets: " - << obj->GetTypeKey(); - } + .set_body_typed([](IRModule mod, int debug_mode, int error_render_level) -> Schedule { return Schedule::Concrete(mod, debug_mode, static_cast(error_render_level)); }); @@ -116,29 +109,30 @@ TVM_REGISTER_GLOBAL("tir.schedule.ScheduleRemoveRV") throw; }); -/***** (FFI) Block/Loop relation *****/ - +/******** (FFI) Sampling ********/ +/******** (FFI) Get blocks & loops ********/ TVM_REGISTER_GLOBAL("tir.schedule.ScheduleGetBlock") .set_body_method(&ScheduleNode::GetBlock); TVM_REGISTER_GLOBAL("tir.schedule.ScheduleGetLoops") .set_body_method(&ScheduleNode::GetLoops); -/******** (FFI) loops manipulation ********/ +/******** (FFI) Transform loops ********/ TVM_REGISTER_GLOBAL("tir.schedule.ScheduleFuse").set_body_method(&ScheduleNode::Fuse); TVM_REGISTER_GLOBAL("tir.schedule.ScheduleSplit").set_body_method(&ScheduleNode::Split); -/******** (FFI) compute location ********/ +/******** (FFI) Manipulate ForKind ********/ +/******** (FFI) Insert cache stages ********/ +/******** (FFI) Compute location ********/ TVM_REGISTER_GLOBAL("tir.schedule.ScheduleComputeInline") .set_body_method(&ScheduleNode::ComputeInline); TVM_REGISTER_GLOBAL("tir.schedule.ScheduleReverseComputeInline") .set_body_method(&ScheduleNode::ReverseComputeInline); -/******** (FFI) loop binding/annotation ********/ -/******** (FFI) cache read/write ********/ -/******** (FFI) reduction ********/ -/******** (FFI) blockize & tensorize ********/ - -/******** (FFI) reduction ********/ - +/******** (FFI) Reduction ********/ TVM_REGISTER_GLOBAL("tir.schedule.ScheduleRFactor") .set_body_method(&ScheduleNode::RFactor); +/******** (FFI) Blockize & Tensorize ********/ +/******** (FFI) Annotation ********/ +/******** (FFI) Misc ********/ +TVM_REGISTER_GLOBAL("tir.schedule.ScheduleEnterPostproc") + .set_body_method(&ScheduleNode::EnterPostproc); } // namespace tir } // namespace tvm diff --git a/src/tir/schedule/state.cc b/src/tir/schedule/state.cc index 6865e41ffb3c..8f0284f2901e 100644 --- a/src/tir/schedule/state.cc +++ b/src/tir/schedule/state.cc @@ -416,9 +416,6 @@ ScheduleState::ScheduleState(IRModule mod, int debug_mode) { data_ = StateCreator::Create(mod, debug_mode); } -ScheduleState::ScheduleState(PrimFunc func, int debug_mode) - : ScheduleState(IRModule({{GlobalVar("main"), func}}), debug_mode) {} - /**************** Replace ****************/ /* @@ -1035,16 +1032,10 @@ TVM_DLL Array GetCachedFlags(const ScheduleState& self, const StmtSRef& bl /**************** FFI ****************/ TVM_REGISTER_NODE_TYPE(ScheduleStateNode); -TVM_REGISTER_GLOBAL("tir.schedule.ScheduleState").set_body_typed([](ObjectRef obj, int debug_mode) { - if (const auto* func = obj.as()) { - return ScheduleState(GetRef(func), debug_mode); - } - if (const auto* mod = obj.as()) { - return ScheduleState(GetRef(mod), debug_mode); - } - LOG(FATAL) << "TypeError: Expects `IRModule` or `PrimFunc`, but gets: " << obj->GetTypeKey(); - throw; -}); +TVM_REGISTER_GLOBAL("tir.schedule.ScheduleState") + .set_body_typed([](IRModule mod, int debug_mode) -> ScheduleState { + return ScheduleState(mod, debug_mode); + }); TVM_REGISTER_GLOBAL("tir.schedule.ScheduleStateGetBlockScope") .set_body_method(&ScheduleStateNode::GetBlockScope); TVM_REGISTER_GLOBAL("tir.schedule.ScheduleStateReplace") diff --git a/src/tir/schedule/trace.cc b/src/tir/schedule/trace.cc new file mode 100644 index 000000000000..d8c18f0de0d6 --- /dev/null +++ b/src/tir/schedule/trace.cc @@ -0,0 +1,533 @@ +/* + * 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. + */ +#include "./utils.h" + +namespace tvm { +namespace tir { + +/**************** Constructors ****************/ + +Trace::Trace() { data_ = make_object(); } + +Trace::Trace(Array insts, Map decisions) { + ObjectPtr n = make_object(); + n->insts = std::move(insts); + n->decisions = std::move(decisions); + data_ = std::move(n); +} + +/**************** Utilities ****************/ + +bool IsPostproc(const InstructionKind& inst_kind) { + static InstructionKind inst_enter_postproc = InstructionKind::Get("EnterPostproc"); + return inst_kind.same_as(inst_enter_postproc); +} + +int GetNumValidInstructions(const Array& insts, bool remove_postproc) { + if (!remove_postproc) { + return insts.size(); + } + int n_insts = 0; + for (const Instruction& inst : insts) { + if (!IsPostproc(inst->kind)) { + ++n_insts; + } else { + break; + } + } + return n_insts; +} + +/**************** TranslateInputRVs ****************/ + +Array TranslateInputRVs(const Array& inputs, + const std::unordered_map& rv_map) { + Array result; + result.reserve(inputs.size()); + for (const ObjectRef& input : inputs) { + if (!input.defined() || // constant: nullptr + input->IsInstance() || // constant: string + input->IsInstance() || // constant: integer + input->IsInstance()) { // constant: float + result.push_back(input); + } else if (input->IsInstance() || // RV: block + input->IsInstance() || // RV: loop + input->IsInstance()) { // RV: var + auto it = rv_map.find(input.get()); + ICHECK(it != rv_map.end()) << "IndexError: Random variable doesn't exist: " << input; + result.push_back(GetRef(it->second)); + } else if (const auto* expr = input.as()) { // RV: Expr + result.push_back( + Substitute(GetRef(expr), [&rv_map](const Var& var) -> Optional { + auto it = rv_map.find(var.get()); + if (it == rv_map.end()) { + return NullOpt; + } + const Object* dst = it->second; + ICHECK(dst->IsInstance()) + << "TypeError: Expect 'tir.Var', but gets: " << dst->GetTypeKey(); + return GetRef(static_cast(dst)); + })); + } else { + ICHECK(false) << "TypeError: Cannot recognize the type of an input random variable: " + << input->GetTypeKey(); + throw; + } + } + return result; +} + +Array TranslateInputRVs( + const Array& inputs, + const std::unordered_map& rv_names) { + Array results; + results.reserve(inputs.size()); + for (const ObjectRef& input : inputs) { + if (!input.defined()) { + // Case 0. nullptr => None + results.push_back(String("None")); + continue; + } + auto it = rv_names.find(input); + if (it != rv_names.end()) { + // Case 1. BlockRV, LoopRV, VarRV + results.push_back(it->second); + } else if (const auto* str_obj = input.as()) { + // Case 2. string => "content" + results.push_back(String('"' + std::string(str_obj->data) + '"')); + } else if (input->IsInstance() || input->IsInstance()) { + // Case 3. integer or floating-point number + results.push_back(input); + } else if (input->IsInstance() || inputs->IsInstance() || + inputs->IsInstance()) { + LOG(FATAL) << "IndexError: Random variable is not defined " << input; + throw; + } else { + LOG(FATAL) << "TypeError: Stringifying is not supported for type: " << input->GetTypeKey(); + throw; + } + } + return results; +} + +Array TranslateInputRVs(const Array& inputs, + const std::unordered_map& named_rvs) { + Array results; + results.reserve(inputs.size()); + for (const ObjectRef& input : inputs) { + // Case 3. integer or floating-point number + if (input->IsInstance() || input->IsInstance()) { + results.push_back(input); + continue; + } + const auto* str = input.as(); + CHECK(str) << "TypeError: Expect String, but gets: " << input->GetTypeKey(); + CHECK_GT(str->size, 0) << "ValueError: Empty string is not allowed in input names"; + const char* name = str->data; + int64_t size = str->size; + // Case 2. string + if (size > 2 && name[0] == '"' && name[size - 1] == '"') { + results.push_back(String(std::string(name + 1, size - 2))); + continue; + } + // Case 0 & 1. None, BlockRV, LoopRV, VarRV + auto it = named_rvs.find(name); + CHECK(it != named_rvs.end()) << "ValueError: The random variable is not defined: " << name; + results.push_back(it->second); + } + return results; +} + +/**************** TranslateAddOutputRVs ****************/ + +void TranslateAddOutputRVs(const Array& old_outputs, const Array& new_outputs, + std::unordered_map* rv_map) { + ICHECK_EQ(old_outputs.size(), new_outputs.size()); + int n = old_outputs.size(); + const ObjectRef* p_old = old_outputs.GetArrayNode()->begin(); + const ObjectRef* p_new = new_outputs.GetArrayNode()->begin(); + for (int i = 0; i < n; ++i) { + (*rv_map)[p_old[i].get()] = p_new[i].get(); + } +} + +Array TranslateAddOutputRVs( + const Array& outputs, + std::unordered_map* rv_names) { + Array results; + results.reserve(outputs.size()); + for (const ObjectRef& output : outputs) { + int i = rv_names->size(); + ICHECK(!rv_names->count(output)) + << "ValueError: The random variable has been produced once: " << rv_names->at(output); + String result{ObjectPtr{nullptr}}; + if (output->IsInstance()) { + result = "b" + std::to_string(i); + } else if (output->IsInstance()) { + result = "l" + std::to_string(i); + } else if (output->IsInstance()) { + result = "v" + std::to_string(i); + } else { + LOG(FATAL) << "TypeError: Cannot recognize the type of the random variable: " + << output->GetTypeKey(); + throw; + } + results.push_back(result); + rv_names->emplace(output, std::move(result)); + } + return results; +} + +void TranslateAddOutputRVs(const Array& old_outputs, const Array& new_outputs, + std::unordered_map* named_rvs) { + ICHECK_EQ(old_outputs.size(), new_outputs.size()); + int n = old_outputs.size(); + const ObjectRef* p_old = old_outputs.GetArrayNode()->begin(); + const ObjectRef* p_new = new_outputs.GetArrayNode()->begin(); + for (int i = 0; i < n; ++i) { + const auto* name = static_cast(p_old[i].get()); + named_rvs->emplace(std::string(name->data, name->size), p_new[i]); + } +} + +/**************** Add/Remove/Get ****************/ + +Optional TraceNode::GetDecision(const Instruction& inst) const { + auto it = this->decisions.find(inst); + return it == this->decisions.end() ? Optional(NullOpt) : (*it).second; +} + +void TraceNode::Append(Instruction inst) { insts.push_back(std::move(inst)); } + +void TraceNode::Append(Instruction inst, ObjectRef decision) { + decisions.Set(inst, std::move(decision)); + insts.push_back(std::move(inst)); +} + +Optional TraceNode::Pop() { + if (insts.empty()) { + return NullOpt; + } + Instruction inst = insts.back(); + insts.pop_back(); + if (decisions.count(inst)) { + decisions.erase(inst); + } + return inst; +} + +/**************** Interfacing with InstructionKind ****************/ + +void TraceNode::ApplyToSchedule( + Schedule sch, bool remove_postproc, + runtime::TypedPackedFunc& inputs, // + const Array& attrs, // + const Optional& decision)> + decision_provider) const { + std::unordered_map rv_map; + for (const Instruction& inst : this->insts) { + if (remove_postproc && IsPostproc(inst->kind)) { + break; + } + Array inputs = TranslateInputRVs(inst->inputs, rv_map); + Array attrs = inst->attrs; + Optional decision = this->GetDecision(inst); + if (decision_provider != nullptr) { + decision = decision_provider(inst, inputs, attrs, decision); + } + Array outputs = inst->kind->f_apply_to_schedule(sch, inputs, attrs, decision); + TranslateAddOutputRVs(inst->outputs, outputs, &rv_map); + } +} + +ObjectRef TraceNode::AsJSON(bool remove_postproc) const { + std::unordered_map rv_names; + Array json_insts; + Array json_decisions; + json_insts.reserve(this->insts.size()); + json_decisions.reserve(this->insts.size()); + + int i = 0; + for (const Instruction& inst : this->insts) { + const InstructionKind& kind = inst->kind; + if (remove_postproc && IsPostproc(kind)) { + break; + } + json_insts.push_back(Array{ + /* 0: inst name */ kind->name, + /* 1: inputs */ TranslateInputRVs(inst->inputs, rv_names), + /* 2: attrs */ kind->f_attrs_as_json != nullptr ? kind->f_attrs_as_json(inst->attrs) + : ObjectRef(inst->attrs), + /* 3: outputs */ TranslateAddOutputRVs(inst->outputs, &rv_names), + }); + if (Optional decision = this->GetDecision(inst)) { + json_decisions.push_back(Array{ + /* 0: index */ Integer(i), + /* 1: decision */ decision.value(), + }); + } + ++i; + } + return Array{ + /* 0: trace */ std::move(json_insts), + /* 1: decision */ std::move(json_decisions), + }; +} + +Array TraceNode::AsPython(bool remove_postproc) const { + std::unordered_map rv_names; + Array py_trace; + py_trace.reserve(this->insts.size()); + for (const Instruction& inst : this->insts) { + if (remove_postproc && IsPostproc(inst->kind)) { + break; + } + Array attrs; + attrs.reserve(inst->attrs.size()); + for (const ObjectRef& obj : inst->attrs) { + if (const auto* str = obj.as()) { + attrs.push_back(String('"' + std::string(str->data) + '"')); + } else { + attrs.push_back(obj); + } + } + py_trace.push_back( + inst->kind->f_as_python(/*inputs=*/TranslateInputRVs(inst->inputs, rv_names), + /*attrs=*/attrs, + /*decision=*/this->GetDecision(inst), + /*outputs=*/TranslateAddOutputRVs(inst->outputs, &rv_names))); + } + return py_trace; +} + +void Trace::ApplyJSONToSchedule(ObjectRef json, Schedule sch) { + Array json_insts{nullptr}; + Array json_decisions{nullptr}; + // Parse `json` into `json_insts` and `json_decisions` + try { + const ArrayNode* arr = json.as(); + ICHECK(arr && arr->size() == 2); + const auto* arr0 = arr->at(0).as(); + const auto* arr1 = arr->at(1).as(); + ICHECK(arr0 && arr1); + json_insts = GetRef>(arr0); + json_decisions = GetRef>(arr1); + } catch (const tvm::Error& e) { + LOG(FATAL) << "ValueError: The json entry of a trace should contain two arrays, an array of " + "instructions and an array of decisions, but gets: " + << json; + throw; + } + // Parse `json_decisions` + std::vector> decisions(json_insts.size(), NullOpt); + for (const ObjectRef& decision_entry : json_decisions) { + int index = -1; + ObjectRef decision{nullptr}; + try { + const ArrayNode* arr = decision_entry.as(); + ICHECK(arr && arr->size() == 2); + const IntImmNode* arr0 = arr->at(0).as(); + ICHECK(arr0); + index = arr0->value; + decision = arr->at(1); + } catch (const tvm::Error& e) { + LOG(FATAL) << "ValueError: Each entry of a json decision should be a tuple [index, " + "decision], but gets: " + << decision_entry; + throw; + } + decisions[index] = std::move(decision); + } + // Parse `json_insts` + std::unordered_map named_rvs{{"None", ObjectRef{nullptr}}}; + int i = 0; + for (const ObjectRef& inst_entry : json_insts) { + InstructionKind kind{nullptr}; + Array inputs{nullptr}; + Array attrs{nullptr}; + Array outputs{ObjectPtr{nullptr}}; + // Parse the entry + try { + const auto* arr = inst_entry.as(); + ICHECK(arr && arr->size() == 4); + const auto* arr0 = arr->at(0).as(); + const auto* arr1 = arr->at(1).as(); + const auto* arr2 = arr->at(2).as(); + const auto* arr3 = arr->at(3).as(); + ICHECK(arr0 && arr1 && arr2 && arr3); + for (const ObjectRef& str : *arr3) { + ICHECK(str->IsInstance()); + } + kind = InstructionKind::Get(arr0->data); + inputs = GetRef>(arr1); + attrs = GetRef>(arr2); + outputs = GetRef>(arr3); + } catch (const tvm::Error& e) { + LOG(FATAL) << "ValueError: Each entry of a json instruction should be a tuple [inst_name, " + "inputs, attrs, outputs], but gets: " + << inst_entry; + throw; + } + // Parse inputs + inputs = TranslateInputRVs(inputs, named_rvs); + // Parse attrs + if (kind->f_attrs_from_json != nullptr) { + attrs = kind->f_attrs_from_json(attrs); + } + // Apply to the schedule + Array new_outputs = kind->f_apply_to_schedule(sch, inputs, attrs, decisions[i]); + // Parse outputs + TranslateAddOutputRVs(outputs, new_outputs, &named_rvs); + ++i; + } +} + +/**************** Creation ****************/ + +Trace TraceNode::WithDecision(Instruction inst, ObjectRef decision, bool remove_postproc) const { + int n_insts = GetNumValidInstructions(this->insts, remove_postproc); + Array new_insts = + Array{this->insts.begin(), this->insts.begin() + n_insts}; + Map new_decisions{this->decisions.begin(), this->decisions.end()}; + new_decisions.Set(std::move(inst), std::move(decision)); + return Trace(new_insts, new_decisions); +} + +Trace TraceNode::Simplified(bool remove_postproc) const { + int n_insts = GetNumValidInstructions(this->insts, remove_postproc); + std::unordered_set used_rvs; + std::vector new_insts; + std::unordered_map new_decisions; + new_insts.reserve(n_insts); + new_decisions.reserve(this->decisions.size()); + for (int inst_idx = n_insts - 1; inst_idx >= 0; --inst_idx) { + const Instruction& inst = this->insts[inst_idx]; + // Check if all the variables the instruction defined are dead + // If so, and the instruction is pure, we can safely remove this instruction + bool all_defs_dead = inst->kind->is_pure; + if (all_defs_dead) { + for (const ObjectRef& obj : inst->outputs) { + if (used_rvs.count(obj.get())) { + all_defs_dead = false; + break; + } + } + } + // Remove this instruction + if (all_defs_dead) { + continue; + } + // Otherwise this instruction is not dead + new_insts.push_back(inst); + if (Optional decision = this->GetDecision(inst)) { + new_decisions.emplace(inst, std::move(decision)); + } + // Add its inputs as "used" ones + for (const ObjectRef& obj : inst->inputs) { + if (obj->IsInstance() || obj->IsInstance() || + obj->IsInstance()) { + used_rvs.insert(obj.get()); + continue; + } else if (obj->IsInstance()) { + PostOrderVisit(obj, [&used_rvs](const ObjectRef& obj) -> void { + if (obj->IsInstance()) { + used_rvs.insert(obj.get()); + } + }); + } + } + } + return Trace(Array(new_insts.rbegin(), new_insts.rend()), + Map(new_decisions)); +} + +/**************** Repr ****************/ + +TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable) + .set_dispatch([](const ObjectRef& obj, ReprPrinter* p) { + const auto* self = obj.as(); + ICHECK_NOTNULL(self); + Array repr = self->AsPython(/*remove_postproc=*/false); + bool is_first = true; + for (const String& line : repr) { + if (is_first) { + is_first = false; + } else { + p->stream << std::endl; + } + p->stream << line; + } + }); + +/**************** Instruction Registration ****************/ + +struct EnterPostprocTraits : public UnpackedInstTraits { + static constexpr const char* kName = "EnterPostproc"; + static constexpr bool kIsPure = false; + + private: + static constexpr size_t kNumInputs = 0; + static constexpr size_t kNumAttrs = 0; + static constexpr size_t kNumDecisions = 0; + + static void UnpackedApplyToSchedule(Schedule sch) { return sch->EnterPostproc(); } + + static String UnpackedAsPython(Array outputs) { + PythonAPICall py("enter_postproc"); + return py.Str(); + } + + template + friend struct ::tvm::tir::UnpackedInstTraits; +}; + +TVM_REGISTER_INST_KIND_TRAITS(EnterPostprocTraits); + +/**************** FFI ****************/ + +TVM_REGISTER_NODE_TYPE(TraceNode); +TVM_REGISTER_GLOBAL("tir.schedule.Trace") + .set_body_typed([](Optional> insts, + Optional> decisions) { + return Trace(insts.value_or(Array()), + decisions.value_or(Map())); + }); +TVM_REGISTER_GLOBAL("tir.schedule.TraceGetDecision") + .set_body_method(&TraceNode::GetDecision); +TVM_REGISTER_GLOBAL("tir.schedule.TraceAppend") + .set_body_typed([](Trace self, Instruction inst, Optional decision) { + if (decision.defined()) { + return self->Append(inst, decision.value()); + } else { + return self->Append(inst); + } + }); +TVM_REGISTER_GLOBAL("tir.schedule.TracePop").set_body_method(&TraceNode::Pop); +TVM_REGISTER_GLOBAL("tir.schedule.TraceApplyToSchedule") + .set_body_method(&TraceNode::ApplyToSchedule); +TVM_REGISTER_GLOBAL("tir.schedule.TraceAsJSON").set_body_method(&TraceNode::AsJSON); +TVM_REGISTER_GLOBAL("tir.schedule.TraceAsPython").set_body_method(&TraceNode::AsPython); +TVM_REGISTER_GLOBAL("tir.schedule.TraceWithDecision") + .set_body_method(&TraceNode::WithDecision); +TVM_REGISTER_GLOBAL("tir.schedule.TraceSimplified").set_body_method(&TraceNode::Simplified); +TVM_REGISTER_GLOBAL("tir.schedule.TraceApplyJSONToSchedule") + .set_body_typed(Trace::ApplyJSONToSchedule); + +} // namespace tir +} // namespace tvm diff --git a/src/tir/schedule/utils.h b/src/tir/schedule/utils.h index d31cea578139..8ccf8da731b5 100644 --- a/src/tir/schedule/utils.h +++ b/src/tir/schedule/utils.h @@ -25,18 +25,22 @@ #include #include #include +#include #include #include +#include #include #include #include +#include "../../node/attr_registry.h" #include "../../printer/text_printer.h" #include "../../runtime/thread_storage_scope.h" #include "../../support/array.h" #include "./analysis.h" #include "./error.h" +#include "./instruction_traits.h" #include "./primitive.h" namespace tvm { diff --git a/tests/python/unittest/test_tir_schedule_block_scope.py b/tests/python/unittest/test_tir_schedule_block_scope.py index 4a914f5063f8..ced8d78ff11a 100644 --- a/tests/python/unittest/test_tir_schedule_block_scope.py +++ b/tests/python/unittest/test_tir_schedule_block_scope.py @@ -15,6 +15,9 @@ # specific language governing permissions and limitations # under the License. # pylint: disable=missing-function-docstring,missing-module-docstring +import sys + +import pytest import tvm from tvm import tir from tvm.script import ty @@ -140,6 +143,4 @@ def test_war_dependency(): if __name__ == "__main__": - test_elementwise_dependency() - test_matmul_dependency() - test_war_dependency() + sys.exit(pytest.main([__file__] + sys.argv[1:])) diff --git a/tests/python/unittest/test_tir_schedule_compute_inline.py b/tests/python/unittest/test_tir_schedule_compute_inline.py index 0a33db09aef1..d6934c6f407f 100644 --- a/tests/python/unittest/test_tir_schedule_compute_inline.py +++ b/tests/python/unittest/test_tir_schedule_compute_inline.py @@ -15,6 +15,8 @@ # specific language governing permissions and limitations # under the License. # pylint: disable=missing-function-docstring,missing-module-docstring +import sys + import pytest import tvm from tvm import tir @@ -354,20 +356,4 @@ def test_compute_inline_multi_loads(): if __name__ == "__main__": - test_compute_inline_elementwise() - test_compute_inline_under_loop() - test_compute_inline_as_dce() - test_compute_inline_multi_consumer() - test_compute_inline_fail_multi_writer() - test_reverse_compute_inline_elementwise() - test_reverse_compute_inline_under_loop() - test_reverse_compute_inline_fail_as_dce() - test_reverse_compute_inline_fail_multi_producer() - test_reverse_compute_inline_fail_multi_reader() - test_reverse_compute_multi_reverse_loads() - test_reverse_compute_fail_multi_reverse_loads() - test_opaque_access_load() - test_opaque_access_store() - test_buffer_matched() - test_compute_inline_predicate() - test_compute_inline_multi_loads() + sys.exit(pytest.main([__file__] + sys.argv[1:])) diff --git a/tests/python/unittest/test_tir_schedule_error.py b/tests/python/unittest/test_tir_schedule_error.py index 1fa658feabe3..6f56eb598894 100644 --- a/tests/python/unittest/test_tir_schedule_error.py +++ b/tests/python/unittest/test_tir_schedule_error.py @@ -15,12 +15,13 @@ # specific language governing permissions and limitations # under the License. # pylint: disable=missing-function-docstring,missing-module-docstring +import sys + import pytest import tvm from tvm import tir from tvm.script import ty - # pylint: disable=no-member,invalid-name,unused-variable @@ -65,6 +66,4 @@ def test_tir_schedule_error_none(): if __name__ == "__main__": - test_tir_schedule_error_detail() - test_tir_schedule_error_fast() - test_tir_schedule_error_none() + sys.exit(pytest.main([__file__] + sys.argv[1:])) diff --git a/tests/python/unittest/test_tir_schedule_instruction.py b/tests/python/unittest/test_tir_schedule_instruction.py new file mode 100644 index 000000000000..9e6f447dd3e6 --- /dev/null +++ b/tests/python/unittest/test_tir_schedule_instruction.py @@ -0,0 +1,68 @@ +# 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-function-docstring,missing-module-docstring +# mypy: ignore-errors +import sys + +import pytest +from tvm.tir.schedule import BlockRV, Instruction, InstructionKind, LoopRV + + +def test_inst_kind_get(): + kind = InstructionKind.get("EnterPostproc") + assert not kind.is_pure + assert kind.name == "EnterPostproc" + + +def test_inst_construct_1(): + block = BlockRV() + loop0 = LoopRV() + loop1 = LoopRV() + inst = Instruction( + kind=InstructionKind.get("GetLoops"), + inputs=[block], + attrs=[], + outputs=[loop0, loop1], + ) + assert str(inst) == "_, _ = sch.get_loops(block=_)" + assert len(inst.inputs) == 1 + assert len(inst.attrs) == 0 + assert len(inst.outputs) == 2 + assert inst.kind.same_as(InstructionKind.get("GetLoops")) + assert inst.inputs[0].same_as(block) + assert inst.outputs[0].same_as(loop0) + assert inst.outputs[1].same_as(loop1) + + +def test_inst_construct_2(): + block = BlockRV() + inst = Instruction( + kind=InstructionKind.get("ComputeInline"), + inputs=[block], + attrs=[], + outputs=[], + ) + assert str(inst) == "sch.compute_inline(block=_)" + assert len(inst.inputs) == 1 + assert len(inst.attrs) == 0 + assert len(inst.outputs) == 0 + assert inst.kind.same_as(InstructionKind.get("ComputeInline")) + assert inst.inputs[0].same_as(block) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__] + sys.argv[1:])) diff --git a/tests/python/unittest/test_tir_schedule_reduction.py b/tests/python/unittest/test_tir_schedule_reduction.py index b22183bf2958..b285f72ca59f 100644 --- a/tests/python/unittest/test_tir_schedule_reduction.py +++ b/tests/python/unittest/test_tir_schedule_reduction.py @@ -14,9 +14,10 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -import pytest +import sys import numpy as np +import pytest import tvm import tvm.testing from tvm import tir @@ -671,4 +672,4 @@ def test_reduction_rfactor_outermost_loop_multiple_children(): if __name__ == "__main__": - pytest.main([__file__]) + sys.exit(pytest.main([__file__] + sys.argv[1:])) diff --git a/tests/python/unittest/test_tir_schedule_split_fuse.py b/tests/python/unittest/test_tir_schedule_split_fuse.py index 4c5c49a1a039..9ac15b8c1986 100644 --- a/tests/python/unittest/test_tir_schedule_split_fuse.py +++ b/tests/python/unittest/test_tir_schedule_split_fuse.py @@ -15,6 +15,8 @@ # specific language governing permissions and limitations # under the License. # pylint: disable=missing-function-docstring,missing-module-docstring +import sys + import pytest import tvm from tvm import tir @@ -450,4 +452,4 @@ def test_split_symbolic(): if __name__ == "__main__": - pytest.main([__file__]) + sys.exit(pytest.main([__file__] + sys.argv[1:])) diff --git a/tests/python/unittest/test_tir_schedule_state.py b/tests/python/unittest/test_tir_schedule_state.py index 34041120f252..ca2ee796a2ba 100644 --- a/tests/python/unittest/test_tir_schedule_state.py +++ b/tests/python/unittest/test_tir_schedule_state.py @@ -15,9 +15,10 @@ # specific language governing permissions and limitations # under the License. # pylint: disable=missing-function-docstring,missing-module-docstring - import gc +import sys +import pytest import tvm from tvm import tir from tvm.ir import IRModule @@ -338,16 +339,4 @@ def test_replace_ir_module(): if __name__ == "__main__": - test_replace_direct_write0() - test_replace_direct_write1() - test_replace_copy() - test_replace_partial_copy0() - test_replace_partial_copy1() - test_replace_root_write() - test_replace_root_copy0() - test_replace_root_copy1() - test_replace_root_copy2() - test_replace_root_copy3() - test_replace_block_remap() - test_replace_block_in_opaque_block() - test_replace_ir_module() + sys.exit(pytest.main([__file__] + sys.argv[1:])) diff --git a/tests/python/unittest/test_tir_schedule_state_cached_flags.py b/tests/python/unittest/test_tir_schedule_state_cached_flags.py index a320812b339f..f77ec0318eea 100644 --- a/tests/python/unittest/test_tir_schedule_state_cached_flags.py +++ b/tests/python/unittest/test_tir_schedule_state_cached_flags.py @@ -15,7 +15,9 @@ # specific language governing permissions and limitations # under the License. # pylint: disable=missing-function-docstring,missing-module-docstring +import sys +import pytest import tvm from tvm import tir from tvm.script import ty @@ -651,19 +653,4 @@ def test_warp_memory_negative(): if __name__ == "__main__": - test_elementwise() - test_matmul() - test_block_in_opaque_block() - test_write_after_read() - test_loop_carried_dependency() - test_concatenate_multi_producer_covered() - test_concatenate_multi_producer_uncovered() - test_lca_at_loop() - test_multi_producer_consumer() - test_elementwise_affine_producer() - test_subblock() - test_subblock_uncovered() - test_thread_binding() - test_equal_ranked_threads() - test_warp_memory() - test_warp_memory_negative() + sys.exit(pytest.main([__file__] + sys.argv[1:])) diff --git a/tests/python/unittest/test_tir_schedule_trace.py b/tests/python/unittest/test_tir_schedule_trace.py new file mode 100644 index 000000000000..cafc6fe1d292 --- /dev/null +++ b/tests/python/unittest/test_tir_schedule_trace.py @@ -0,0 +1,241 @@ +# 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-function-docstring,missing-module-docstring +# mypy: ignore-errors +import sys + +import pytest +import tvm +from tvm import tir +from tvm.script import ty +from tvm.tir.schedule import BlockRV, Instruction, InstructionKind, LoopRV, Trace + +# pylint: disable=no-member,invalid-name,unused-variable + + +@tvm.script.tir +def elementwise(a: ty.handle, c: ty.handle) -> None: + A = tir.match_buffer(a, (128, 128)) + B = tir.alloc_buffer((128, 128)) + C = tir.match_buffer(c, (128, 128)) + with tir.block([128, 128], "B") as [vi, vj]: + B[vi, vj] = A[vi, vj] * 2.0 + with tir.block([128, 128], "C") as [vi, vj]: + C[vi, vj] = B[vi, vj] + 1.0 + + +@tvm.script.tir +def elementwise_inlined(a: ty.handle, c: ty.handle) -> None: + A = tir.match_buffer(a, (128, 128)) + C = tir.match_buffer(c, (128, 128)) + with tir.block([128, 128], "C") as [vi, vj]: + C[vi, vj] = A[vi, vj] * 2.0 + 1.0 + + +# pylint: enable=no-member,invalid-name,unused-variable + + +def _make_get_block(name, output): + return Instruction( + kind=InstructionKind.get("GetBlock"), + inputs=[], + attrs=[name, "main"], + outputs=[output], + ) + + +def _make_get_loops(input, outputs): # pylint: disable=redefined-builtin + return Instruction( + kind=InstructionKind.get("GetLoops"), + inputs=[input], + attrs=[], + outputs=outputs, + ) + + +def _make_compute_inline(input): # pylint: disable=redefined-builtin + return Instruction( + kind=InstructionKind.get("ComputeInline"), + inputs=[input], + attrs=[], + outputs=[], + ) + + +def _make_enter_postproc(): + return Instruction( + kind=InstructionKind.get("EnterPostproc"), + inputs=[], + attrs=[], + outputs=[], + ) + + +def _make_trace_1(b0, l1, l2): # pylint: disable=invalid-name + return Trace( + insts=[ + _make_get_block(name="block", output=b0), + _make_get_loops(input=b0, outputs=[l1, l2]), + ], + decisions={}, + ) + + +def _make_trace_2(b0): # pylint: disable=invalid-name + return Trace( + insts=[ + _make_get_block(name="B", output=b0), + _make_compute_inline(input=b0), + ], + decisions={}, + ) + + +def _make_trace_3(b0, b1, add_postproc): # pylint: disable=invalid-name + if add_postproc: + insts = [ + _make_get_block(name="B", output=b0), + _make_compute_inline(input=b0), + _make_get_block(name="C", output=b1), + _make_enter_postproc(), + _make_compute_inline(input=b1), + ] + else: + insts = [ + _make_get_block(name="B", output=b0), + _make_compute_inline(input=b0), + _make_get_block(name="C", output=b1), + ] + return Trace(insts=insts, decisions={}) + + +def test_trace_construct_1(): + trace = _make_trace_1(BlockRV(), LoopRV(), LoopRV()) + assert str(trace) == "\n".join( + ( + 'b0 = sch.get_block(name="block", func_name="main")', + "l1, l2 = sch.get_loops(block=b0)", + ) + ) + assert len(trace.insts) == 2 + assert len(trace.decisions) == 0 + + +def test_trace_construct_get_decision_1(): + trace = _make_trace_1(BlockRV(), LoopRV(), LoopRV()) + assert trace.get_decision(trace.insts[0]) is None + assert trace.get_decision(trace.insts[1]) is None + + +def test_trace_construct_append_1(): + trace = _make_trace_1(BlockRV(), LoopRV(), LoopRV()) + trace.append(inst=_make_get_block("block2", BlockRV())) + assert str(trace) == "\n".join( + ( + 'b0 = sch.get_block(name="block", func_name="main")', + "l1, l2 = sch.get_loops(block=b0)", + 'b3 = sch.get_block(name="block2", func_name="main")', + ) + ) + + +def test_trace_construct_pop_1(): + trace = _make_trace_1(BlockRV(), LoopRV(), LoopRV()) + last_inst = trace.insts[-1] + assert trace.pop().same_as(last_inst) + assert str(trace) == 'b0 = sch.get_block(name="block", func_name="main")' + + +def test_trace_construct_pop_2(): + trace = Trace([], {}) + assert str(trace) == "" + assert trace.pop() is None + assert str(trace) == "" + + +def test_trace_apply_to_schedule(): + trace = _make_trace_2(BlockRV()) + sch = tir.Schedule(elementwise, debug_mode=True) + trace.apply_to_schedule(sch, remove_postproc=False, decision_provider=None) + tvm.ir.assert_structural_equal(elementwise_inlined, sch.mod["main"]) + + +def test_trace_as_json_1(): + trace = _make_trace_1(BlockRV(), LoopRV(), LoopRV()) + obj = trace.as_json() + assert obj == [ + [ + ["GetBlock", [], ["block", "main"], ["b0"]], + ["GetLoops", ["b0"], [], ["l1", "l2"]], + ], + [], + ] + + +def test_trace_simplified_1(): + trace = _make_trace_3(BlockRV(), BlockRV(), add_postproc=True) + assert str(trace) == "\n".join( + ( + 'b0 = sch.get_block(name="B", func_name="main")', + "sch.compute_inline(block=b0)", + 'b1 = sch.get_block(name="C", func_name="main")', + "sch.enter_postproc()", + "sch.compute_inline(block=b1)", + ) + ) + trace = trace.simplified(remove_postproc=True) + assert str(trace) == "\n".join( + ( + 'b0 = sch.get_block(name="B", func_name="main")', + "sch.compute_inline(block=b0)", + ) + ) + + +def test_trace_simplified_2(): + trace = _make_trace_3(BlockRV(), BlockRV(), add_postproc=True) + assert str(trace) == "\n".join( + ( + 'b0 = sch.get_block(name="B", func_name="main")', + "sch.compute_inline(block=b0)", + 'b1 = sch.get_block(name="C", func_name="main")', + "sch.enter_postproc()", + "sch.compute_inline(block=b1)", + ) + ) + trace = trace.simplified(remove_postproc=False) + assert str(trace) == "\n".join( + ( + 'b0 = sch.get_block(name="B", func_name="main")', + "sch.compute_inline(block=b0)", + 'b1 = sch.get_block(name="C", func_name="main")', + "sch.enter_postproc()", + "sch.compute_inline(block=b1)", + ) + ) + + +def test_apply_json_to_schedule_1(): + trace = _make_trace_2(BlockRV()) + json_obj = trace.as_json() + sch = tir.Schedule(elementwise, debug_mode=True) + Trace.apply_json_to_schedule(json_obj, sch) + tvm.ir.assert_structural_equal(elementwise_inlined, sch.mod["main"]) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__] + sys.argv[1:])) diff --git a/tests/python/unittest/test_tir_schedule_utilities.py b/tests/python/unittest/test_tir_schedule_utilities.py index af89ca252738..07658978db52 100644 --- a/tests/python/unittest/test_tir_schedule_utilities.py +++ b/tests/python/unittest/test_tir_schedule_utilities.py @@ -15,13 +15,14 @@ # specific language governing permissions and limitations # under the License. # pylint: disable=missing-function-docstring,missing-module-docstring +import sys + import pytest import tvm from tvm import tir from tvm.ir import IRModule from tvm.script import ty - # pylint: disable=no-member,invalid-name,unused-variable @@ -108,8 +109,4 @@ def test_tir_schedule_remove_rv(): if __name__ == "__main__": - test_tir_schedule_creation() - test_tir_schedule_get_block() - test_tir_schedule_get_loops() - test_tir_schedule_copy() - test_tir_schedule_remove_rv() + sys.exit(pytest.main([__file__] + sys.argv[1:]))