Primitives to set tir block attr?

Questions
2

There are some testcases to describe how annotations are lowered: https://github.com/apache/tvm/blob/main/tests/python/unittest/test_tir_transform_lower_opaque_block.py#L324. As a summary:

  • only annotation keys start with pragma_ are converted to T.attr
  • loop annotations are preserved as it is if not start with pragma_

The community generally has a trend to discard T.attr in TensorIR lowering flow (refer to Can we lift tir.AttrStmt value type to ObjectRef? - #5 by tqchen). I am not sure why inject_ptx_async_copy require current fashion, maybe we could adjust it to make consistent with TensorIR annotation system? cc authors @junrushao @masahi for help:)

2 Likes