Pragmas in TensorIR

Questions
1

Hello, is there a way to add pragmas in TIR schedules? As far as I saw annotate* primitive does not appear in lowered schedules.

*: sch.annotate(global_A_block, “ann_key”, “ann_value”)

2

In the lowering phase currently, all keys with “pragma_” prefix are converted to attribute nodes, just the same with TE-schedule, otherwise:

  • If you annotate on a block, the keys are dropped.
  • If you annotate on a loop, the keys are preserved in loop annotations.
3

It worked, thanks.

so basic implementation would be as in following

global_A_block = sch.get_block(“A_global”)

lp =sch.get_loops(global_A_block)

sch.annotate(lp[0], “pragma_key”, “pragma_value”) # annotated on only first iteration variable of the nest loops

1 Like