Function visit_expr_stmt

python/tvm/tirx/script/parser/parser.py:715–763

The expr statement visiting method for tirx. Parameters ---------- self : Parser The visiting parser. node : doc.Expr The doc AST Expr node.

(self: Parser, node: doc.Expr)

@dispatch.register(token="tirx", type_name="Expr")
def visit_expr_stmt(self: Parser, node: doc.Expr) -> None:
    """The expr statement visiting method for tirx.

    Parameters
    ----------
    self : Parser
        The visiting parser.

    node : doc.Expr
        The doc AST Expr node.
    """

    res = self.eval_expr(node.value)
    if res is None:
        pass
    elif isinstance(res, Frame):
        res.add_callback(partial(res.__exit__, None, None, None))
        res.__enter__()
    elif hasattr(res, "frames") and hasattr(res, "__enter__"):
        # _FrameScope from T.attr({...}) — enter each inner frame for concise scoping
        for f in res.frames:
            f.add_callback(partial(f.__exit__, None, None, None))
            f.__enter__()
    elif isinstance(res, Var):
        # Standalone Var expression (e.g. from T.bind(value, var=v)) --
        # the Bind statement was already emitted to the parent frame by the FFI call,
        # so just discard the returned Var.
        pass
    elif isinstance(res, PrimExpr):
        T.evaluate(res)
    elif isinstance(res, int | bool):
        T.evaluate(tvm.tirx.const(res))
    elif isinstance(res, tvm.relax.Call) and not res.args:
        # Using GlobalVar.__call__ with no arguments is ambiguous, as
        # each IR has a different function Call representation.  If
        # this occurs, convert to the TIR representation.
        T.evaluate(tvm.tirx.call_tir(res.op))
    elif isinstance(res, str):
        # Ignore docstrings
        pass
    elif isinstance(res, tvm.tirx.stmt.BufferStore):
        T.buffer_store(res.buffer, res.value, res.indices, res.predicate)
    elif isinstance(res, tvm.tirx.Buffer):
        # ``T.match_buffer(...)`` used as a bare statement (no LHS) — the
        # buffer object is discarded; the underlying side effect (the
        # match_buffer node) has already been emitted into the frame.
        pass
    else:
        self.report_error(node, f"Parsing resulted in unexpected type {type(res)}")


@dispatch.register(token="tirx", type_name="If")