[MetaSchedule][M3a] Instruction and Trace #8615

junrushao · 2021-08-01T00:38:00Z

This PR introduces the two core data structures in meta schedule:

Instruction
Trace

It is part of the meta schedule upstreaming effort: #8473.

comaniac

LGTM. Just a few nits

include/tvm/tir/schedule/instruction.h

include/tvm/tir/schedule/schedule.h

python/tvm/tir/schedule/instruction.py

python/tvm/tir/schedule/trace.py

tests/python/unittest/test_tir_schedule_trace.py

Co-authored-by: Siyuan Feng <Hzfengsy@sjtu.edu.cn> Co-authored-by: Bohan Hou <32121147+spectrometerHBH@users.noreply.github.com> Co-authored-by: Ruihang Lai <lairuihangdongdong@qq.com> Co-authored-by: Hongyi Jin <3231950289@qq.com> Co-authored-by: Wuwei Lin <wuwei@apache.org> Co-authored-by: Cody Yu <comaniac0422@gmail.com>

junrushao · 2021-08-01T20:57:09Z

Thanks @comaniac for the review! I fixed the RST format issue and turned all unittests related tir schedule to use pytest

comaniac · 2021-08-02T01:54:25Z

Thanks @junrushao1994

tkonolige

@comaniac this seemed to get merged really quickly over a weekend. In the future could we hold off merging important PRs like this for at least one work day so the community gets a chance to comment?

tkonolige · 2021-08-02T23:15:45Z

include/tvm/tir/schedule/instruction.h

+   */
+  bool is_pure{false};
+  /*! \brief A functor that applies the instruction to a TensorIR schedule */
+  FInstructionApply f_apply_to_schedule{nullptr};


Why are all these members typed packed functions and not member functions?

This is how registry works. Here are some examples: https://github.com/apache/tvm/blob/main/include/tvm/tir/op_attr_types.h#L50-L55

tkonolige · 2021-08-02T23:20:15Z

include/tvm/tir/schedule/instruction.h

+   * nothing, while `ComputeInline` is not because removing it leads to a different resulting
+   * schedule.
+   */
+  bool is_pure{false};


What is the purpose of marking a function as pure?

If the instruction is pure, which doesn't have any effect on the IR, e.g. GetBlock, then we can remove this instruction in dead-code elimination

tkonolige · 2021-08-02T23:22:13Z

include/tvm/tir/schedule/instruction.h

+class InstructionNode : public runtime::Object {
+ public:
+  /*! \brief The kind of the instruction */
+  InstructionKind kind;


What is the purpose of giving each instruction a kind. Why node use subclassing instead?

It is a common pattern in TVM codebase. We do the same thing for Target in TVM. @tqchen we should probably document it explicitly in the dev code so that new contributors could learn such pattern :-)

subclassing is OK if we have a relatively fix number of class. When it comes to things like operator or layers in ML frameworks, we will then need to switch to a registry pattern, where kind/op are being registered and matched.

Having an explicit kind that can be matched also allows specific pattern matching on the instruction kind itself(rather than rely on virtual functions).

Given that we do want to extend the set of scheduling primitives, the registry pattern fits better so that these instruction kinds could be defined distributedly in several files :-)

tkonolige · 2021-08-02T23:25:05Z

include/tvm/tir/schedule/instruction.h

+ public:
+  static InstructionKindRegEntry& RegisterOrGet(const String& name);
+
+  InstructionKindRegEntry& set_name() {


Why do we want registered instructions to be mutable. It seems like this would cause issues where instructions could be mutated mid tuning.

Hmm yet, it is a common pattern in TVM, it is used to make sure the following syntax works:

TVM_REGISTER_xxx() .set_xxx() .set_yyy();

Some examples:

Our global registry: https://github.com/apache/tvm/blob/main/include/tvm/runtime/registry.h#L111

Operator registry: https://github.com/apache/tvm/blob/main/include/tvm/ir/op.h#L252

TargetKind registry: https://github.com/apache/tvm/blob/main/include/tvm/target/target_kind.h#L163

tkonolige · 2021-08-02T23:32:20Z

include/tvm/tir/schedule/trace.h

+ * 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


What does "possibly with their decisions accordingly" mean? Is it optional to use the original decisions? Does re-sampling mean you are choosing a new decision?

tkonolige · 2021-08-02T23:35:52Z

include/tvm/tir/schedule/trace.h

+  /*! \brief The instructions invoked so far in the program execution */
+  Array<Instruction> insts;
+  /*! \brief The random decisions made upon those instructions */
+  Map<Instruction, ObjectRef> decisions;


If decisions can be arbitrary objects, how do we guarantee that a Trace is serializable?

Decisions are not arbitrary objects. The schedule primitive itself should guarantee it is serializable

tkonolige · 2021-08-02T23:38:44Z

include/tvm/tir/schedule/trace.h

+   * \brief Append a new instruction to the trace
+   * \param inst The new instruction to be appended
+   */
+  void Append(Instruction inst);


Given most internal datastructures are COW, should this be too?

To clarify, It is a pattern in TVM that the IR is immutable, where COW is applied, but other data structures are allowed to be mutable.

More context: TVM provides a macro TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS in its core object system for users to define mutable objects.

For example, there are a lot of data structures in TVM codebase that are not COW for efficiency:

te::Schedule

AutoScheduler's cost model

Relay operator strategy

...

tkonolige · 2021-08-02T23:58:06Z

src/tir/schedule/trace.cc

+            return GetRef<Var>(static_cast<const VarNode*>(dst));
+          }));
+    } else {
+      ICHECK(false) << "TypeError: Cannot recognize the type of an input random variable: "


Use LOG(FATAL)

We were thinking about ILOG(FATAL) but it doesnt exist :-(

Any suggestions?

ICHECK uses LOG(FATAL) but just adds a banner to the error message. We could add a version of LOG(FATAL) that includes the banner.

tkonolige · 2021-08-03T00:05:19Z

include/tvm/tir/schedule/trace.h

+   * \sa FTraceDecisionProvider
+   */
+  void ApplyToSchedule(Schedule sch, bool remove_postproc,
+                       FTraceDecisionProvider decision_provider = nullptr) const;


Callback can often lead to spaghetti code. Instead of providing a callback, can the flow to do this just be that you modify the trace to contain your decisions and then you apply it to the schedule?

This callback is nullptr In most of the cases - so no spaghetti in most of the cases.

In some circumstances, the caller has to analyze the intermediate schedule result TIR to decide which decision to make.

Therefore, we leave this callback to satisfy such needs :-)

junrushao requested review from areusch, comaniac, jroesch, kparzysz-quic, masahi, merrymercy, tqchen, vinx13, yzhliu, zhiics and ZihengJiang as code owners August 1, 2021 00:38

junrushao changed the title ~~[M3a] Instruction and Trace~~ [Meta Schedule][M3a] Instruction and Trace Aug 1, 2021

junrushao mentioned this pull request Aug 1, 2021

[RFC][Tracking Issue] Meta Schedule (AutoTIR) #8473

Closed

62 tasks

junrushao assigned comaniac Aug 1, 2021

comaniac approved these changes Aug 1, 2021

View reviewed changes

junrushao requested a review from comaniac August 1, 2021 20:55

comaniac added the status: accepted label Aug 2, 2021

comaniac merged commit 7653972 into apache:main Aug 2, 2021

tkonolige reviewed Aug 3, 2021

View reviewed changes

MasterJH5574 mentioned this pull request Aug 6, 2021

[M2a] Storage align tlc-pack/tvm-tensorir#422

Merged

ylc pushed a commit to ylc/tvm that referenced this pull request Sep 29, 2021

[Meta Schedule][M3a] Instruction and Trace (apache#8615)

97ebf8f

junrushao mentioned this pull request Nov 1, 2021

Apache TVM v0.8 Release Note Candidate #9416

Closed

ylc pushed a commit to ylc/tvm that referenced this pull request Jan 13, 2022

[Meta Schedule][M3a] Instruction and Trace (apache#8615)

81bca08

junrushao changed the title ~~[Meta Schedule][M3a] Instruction and Trace~~ [MetaSchedule][M3a] Instruction and Trace Jan 26, 2022

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[MetaSchedule][M3a] Instruction and Trace #8615

[MetaSchedule][M3a] Instruction and Trace #8615

junrushao commented Aug 1, 2021 •

edited

Loading

comaniac left a comment

junrushao commented Aug 1, 2021

comaniac commented Aug 2, 2021

tkonolige left a comment

tkonolige Aug 2, 2021

junrushao Aug 3, 2021

tkonolige Aug 2, 2021

junrushao Aug 3, 2021

tkonolige Aug 2, 2021

junrushao Aug 3, 2021 •

edited

Loading

tqchen Aug 3, 2021

junrushao Aug 3, 2021

tkonolige Aug 2, 2021

junrushao Aug 3, 2021 •

edited

Loading

tkonolige Aug 2, 2021

tkonolige Aug 2, 2021

junrushao Aug 3, 2021

tkonolige Aug 2, 2021

junrushao Aug 3, 2021

tkonolige Aug 2, 2021

junrushao Aug 3, 2021

tkonolige Aug 6, 2021

tkonolige Aug 3, 2021

junrushao Aug 3, 2021

		* 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

[MetaSchedule][M3a] Instruction and Trace #8615

[MetaSchedule][M3a] Instruction and Trace #8615

Conversation

junrushao commented Aug 1, 2021 • edited Loading

comaniac left a comment

Choose a reason for hiding this comment

junrushao commented Aug 1, 2021

comaniac commented Aug 2, 2021

tkonolige left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

junrushao Aug 3, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

junrushao Aug 3, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

junrushao commented Aug 1, 2021 •

edited

Loading

junrushao Aug 3, 2021 •

edited

Loading

junrushao Aug 3, 2021 •

edited

Loading