Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[MetaSchedule][M3a] Instruction and Trace #8615

Merged
merged 1 commit into from
Aug 2, 2021
Merged

[MetaSchedule][M3a] Instruction and Trace #8615

merged 1 commit into from
Aug 2, 2021

Conversation

junrushao
Copy link
Member

@junrushao junrushao commented Aug 1, 2021

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

  • Instruction
  • Trace

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

@junrushao junrushao changed the title [M3a] Instruction and Trace [Meta Schedule][M3a] Instruction and Trace Aug 1, 2021
Copy link
Contributor

@comaniac comaniac left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM. Just a few nits

include/tvm/tir/schedule/instruction.h Show resolved Hide resolved
include/tvm/tir/schedule/schedule.h Show resolved Hide resolved
python/tvm/tir/schedule/instruction.py Outdated Show resolved Hide resolved
python/tvm/tir/schedule/trace.py Outdated Show resolved Hide resolved
python/tvm/tir/schedule/trace.py Outdated Show resolved Hide resolved
python/tvm/tir/schedule/trace.py Outdated Show resolved Hide resolved
python/tvm/tir/schedule/trace.py Outdated Show resolved Hide resolved
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 junrushao requested a review from comaniac August 1, 2021 20:55
@junrushao
Copy link
Member Author

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

@comaniac comaniac merged commit 7653972 into apache:main Aug 2, 2021
@comaniac
Copy link
Contributor

comaniac commented Aug 2, 2021

Thanks @junrushao1994

Copy link
Contributor

@tkonolige tkonolige left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@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?

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

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

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

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What is the purpose of marking a function as pure?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

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

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Member Author

@junrushao junrushao Aug 3, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

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

InstructionKindRegEntry& set_name() {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Member Author

@junrushao junrushao Aug 3, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

Comment on lines +54 to +55
* 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
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

/*! \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;
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

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

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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
  • ...

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

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Use LOG(FATAL)

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Any suggestions?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

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

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

ylc pushed a commit to ylc/tvm that referenced this pull request Sep 29, 2021
ylc pushed a commit to ylc/tvm that referenced this pull request Jan 13, 2022
@junrushao junrushao changed the title [Meta Schedule][M3a] Instruction and Trace [MetaSchedule][M3a] Instruction and Trace Jan 26, 2022
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants