[TIR] Add schedule primitive SetAxisSeparator

[vinx13]

This PR added a schedule primitive set_axis_separator, it modifies axis_separators attribute of the target buffer, which affects the physical dimension after flattening.

cc @Lunderberg @csullivan @junrushao1994

[Lunderberg]

This looks fantastic! I really like the ReplaceBufferMutator utility, and look forward to using it elsewhere as well!

[tmoreau89]

cc @quic-sanirudh

[quic-sanirudh]

@Lunderberg @vinx13 This is great, thanks a lot, I would love to start using this.

I do have a small doubt however. Is there a reason why we're introducing a new schedule primitive instead of allowing the user to pass something similar to te.AXIS_SEPARATOR to layout_transform and internally call both ScheduleTransformLayout and SetAxisSeparator as done for te?

The reason for this question is that the list to be passed to axis_separators argument seems like an internal detail and might confuse users. Also it might make it difficult for us to modify how axis_separators are handled internally later if needed once we expose this detail into user facing API.

[csullivan]

Thank you for adding axis separator support to TIR scheduling! I also really love ReplaceBufferMutator, not only very useful but edifying as a concise description of all the places buffer rewrites may need to occur. Just a few comments, only one really needs addressing.

[vinx13]

@quic-sanirudh Thanks for the comments. The motivation to have a separate schedule primitive is to decouple the logical-physical mapping from generic layout transformation which doesn't have to deal with physical layout. @Lunderberg suggested that we can have a user convenience API backed by these two schedule primitives

[quic-sanirudh]

we can have a user convenience API backed by these two schedule primitives

Ah okay, thanks for the reply @vinx13

[csullivan]

Many thanks @vinx13 @quic-sanirudh @junrushao1994 @Lunderberg, this is merged!