Function bind_with_value

python/tvm/tirx/script/parser/parser.py:98–131

Value binding methods when parsing with statement. e.g. binding i, j, k with T.grid(128, 128, 128), when parsing with T.grid(128, 128, 18) as i, j, k. Parameters ---------- self : Parser The current parser. node : doc.expr The doc AST expression node for

(self: Parser, node: doc.expr, var_name: str, value: Any)

def bind_with_value(self: Parser, node: doc.expr, var_name: str, value: Any) -> Any:
    """Value binding methods when parsing with statement.
    e.g. binding i, j, k with T.grid(128, 128, 128), when parsing
        with T.grid(128, 128, 18) as i, j, k.

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

    node : doc.expr
        The doc AST expression node for error reporting.

    var_name : str
        The variable name.

    value : Any
        The value to be bound with.

    Returns
    -------
    res : Any
        The bound value.
    """
    if isinstance(value, list | tuple):
        for i, v in enumerate(value):
            bind_with_value(self, node, f"{var_name}_{i}", v)
        return value
    elif isinstance(value, Buffer | Var):
        IRBuilder.name(var_name, value)
        return value
    else:
        self.report_error(node, f"Do not know how to bind type: {type(value)} in with statement")
        raise NotImplementedError


def bind_for_value(self: Parser, node: doc.expr, var_name: str, value: Any) -> Any: