[TIR] Preserve annotations after lower opaque block

[wrongtest-intellif]

Hi, there. The change aims to improve customization experience of TIR lowering using annotations.

Currently, we have LowerOpaqueBlock pass to convert s-tir(schedulable tir with blocks) to non-stir. All annotations are dropped then except those used by legacy TE pragma attributes, they are converted to legacy AttrStmt. For example:

# before
for i in range(10, annotations={"pragma_myattr0": 1, "my_attr0": 1}): ....

# after
with T.attr(i, "pragma_myattr0", 1):
    for i in range(10): ....

This introduce difficulties when we want to leverage schedule phase annotations in lowering passes.

1. pragma_ annotations are converted to legacy AttrStmt, which only accepts PrimExpr typed value. But the loop/block annotations before are originally much flexible, one could use list/dict to encode more rich compile time hints.
2. Using AttrStmt loses the roundtrip property of TIR, thus we fail to write testcases of lowering pass with friendly before/after script comparations.
3. It is more cohesive to encode annotations just with related IR node (like ForNode). Or else the developers always have to write separate logics handling attr node and for node.

The change just add a pass configuration "tir.LowerOpaqueBlock": {"preserved_annotations": List[str]}}. For each annotation key value of loop or block:

1. if key starts with pragma_, keep the original behaviour.
2. if key is registered in preserved_annotations list, the annotation is preserved.
3. or else it is dropped.

The alternative may be preserve all annotations without any extra configuration. But I think current way take the least impact on overall lowering stream. If we do not add customized pass config into pass context, all things should be same.

cc @Hzfengsy @junrushao1994

[Hzfengsy]

It‘s a bit tricky to add a unit loop here. Can we keep the current behavior or add an AttrStmt?

[junrushao]

I'm a bit uncertain if haivng an AttrStmt might make sense in our workflow, given it's something in the low-level TIR that we might want to minimize its usage (CC: @tqchen). I'm aware that there could be any implications if we directly remove a unit loop without preserving its annotation too. Not sure what the best approach is.

[tqchen]

On AttrStmt, we should move away from it when possible as an effort to limit the scope and clarify the semantics. Of course, we should also document the set of possible annotatons we use in For so they become manageable

[Hzfengsy]

Is it possible to preserve all annotations by default rather than write a pass config?

[wrongtest-intellif]

A tracing of usages in current test_tir_ and test_meta_schedule_ show that for loop annotations, only software pipeline specified keys are used, they are consumed before lower opaque block. The other keys are mostly meta_schedule.* in block annotations.

What about that for non-pragma keys?

• Always preserve loop annotations.

• Always drop block annotations, since they are mostly designed for meta-schedule thus make no sense to subsequent passes, and it seems not clear now how to lower the block's attrs after blocks get eliminated.

For my circumstance the loop annotations are quite useful while block annotations are not used for lowering passes. Could the decision about block annotation get delayed to when it is indeed needed :)?

[Hzfengsy]

cc @junrushao