Function bind_for_value

python/tvm/tirx/script/parser/parser.py:134–167

Value binding methods when parsing for statement. e.g. binding i, j, k with T.grid(128, 128, 128), when parsing for i, j, k in T.grid(128, 128, 128). 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_for_value(self: Parser, node: doc.expr, var_name: str, value: Any) -> Any:
    """Value binding methods when parsing for statement.
    e.g. binding i, j, k with T.grid(128, 128, 128), when parsing
        for i, j, k in T.grid(128, 128, 128).

    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 | tvm.ir.Array):
        for i, v in enumerate(value):
            bind_for_value(self, node, f"{var_name}_{i}", v)
        return value
    elif isinstance(value, Var):
        IRBuilder.name(var_name, value)
        return value
    else:
        self.report_error(node, f"Do not know how to bind type: {type(value)} in for statement")
        raise NotImplementedError


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