Ejemplo n.º 1
0
    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))
Ejemplo n.º 2
0
 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