Schedule annotations of For and Block are all Map<String, ObjectRef>. But certain pragma annotations can not get lowerer to T.attr，only those of expression typed values are allowed.
Would you like to elaborate? Currently the type of AttrStmtNode::value is PrimExpr, but which type cannot be supported by TVMScript printer and why?
Hi~ I think this is not the issue of tvmscript. For example, though List[Integer] is supported by script, it would fail in lowering with Illegal attribute of key pragma_key, value type Array not supported, since the annotation can not convert to an attr stmt.
import tvm
from tvm.script import tir as T
@T.prim_func
def fun(x: T.Buffer((16,), "int32")) -> None:
for k in T.serial(0, 16, annotations={"pragma_key": [1, 2, 3]}):
x[k] = 1
print(tvm.lower(fun).script())
This is definitely interesting usecases which unifies “AttrStmt” with definitions of Attr otherwhere in the codebase. Given AttrStmt is something we wanted to move away from, I would love to confirm with @tqchen that the change is acceptable
AttrStmt was our historical way to create extensions. While it offered quite some convenience, it also created a dilemma that it may introduce new language extensions that may become incompatible with the existing semantics. e.g. an old pass may no longer be compatible with a new extension semantics.
Moving forward, we ideally would like to think about these extension points in a more structured way:
-
annotationsin blocks and for loops are “hints” to the program transformations about what can be done, which means they can be safely ignored (e.g. a parallel for annotation) by most of the transformations. This enables most passes to safely ignore annotations when possible. Importantly, the object in annotations are all constant(so we do not need to worry about mutating its value when we change the IR). - When there are possible operations that changes the semantics of the execution block, we would like to be explicit and consider IR change that would require examination of all pass compatibility(e.g. While loop, bound of a variable which depends on other variables.)
Given AttrStmt was the legacy extension point we might want to leave it was it is and deprecate(or restrict to annotations) as we continue to improve the lowering. In the meantime, we can try to leverage annotation fields for new code generation behaviors.
2 Likes