def create_loop_info(
self,
begin: PrimExpr,
end: PrimExpr,
kind: ForKind,
thread_binding: Optional[str] = None,
annotations: Optional[Mapping[str, Object]] = None,
) -> None:
"""
Helper function for creating For in TVM Script parser.
Parameters
----------
begin : PrimExpr
The beginning value.
end : PrimExpr
The endding value.
kind : ForKind
The type of the for.
thread_binding: Optional[str]
The thread this loop binds to.
annotations : Optional[Mapping[str, Object]]
Additional annotation hints.
span : Optional[Span]
The location of this for in the source code.
Returns
-------
for : For
The constructed For.
"""
begin, end = [convert(_) for _ in [begin, end]]
assert self.context and self.node, "call 'exit_scope' before 'enter_scope'"
extent = end if begin == 0 else self.context.analyzer.simplify(end -
begin)
if begin == 0 and isinstance(extent, PrimExpr):
begin = IntImm(extent.dtype, 0, begin.span)
self.annotations: Mapping[str, Object] = {}
if annotations is not None:
self.annotations = {
key: String(val) if isinstance(val, str) else val
for key, val in annotations.items()
}
self.loop_info.append(
LoopInfo(begin, extent, kind, thread_binding, annotations))
def block_attr(attrs: Mapping[str, Object], span: Span = None):
assert self.context, "call 'exit_scope' before 'enter_scope'"
block_scope = self.context.current_block_scope()
if block_scope is None:
self.context.report_error(
"Expected to declare block annotations inside a block.",
span,
)
if block_scope.annotations is not None:
self.context.report_error(
"Duplicate block annotations declaration, "
+ "previous one is "
+ str(block_scope.annotations),
span,
)
attrs = {
key: String(val) if isinstance(val, str) else val for key, val in attrs.items()
}
block_scope.annotations = attrs