loop_order_bnb

loop_order_brute_force(expr, stats_factory, stats_bindings, output_vars=None)

Pick the loop order with the lowest loop_order_cost over all permutations.

Ties keep the first permutation in expr.fields() order. Exponential in the number of indices; intended as an oracle for small nests and benchmarks.

Source code in src/finch/autoschedule/loop_orderer/loop_order_bnb.py

def loop_order_brute_force(
    expr: LogicExpression,
    stats_factory: StatsFactory[NS],
    stats_bindings: Mapping[Alias, NS],
    output_vars: tuple[Field, ...] | None = None,
) -> tuple[Field, ...]:
    """Pick the loop order with the lowest ``loop_order_cost`` over all permutations.

    Ties keep the first permutation in ``expr.fields()`` order. Exponential in the
    number of indices; intended as an oracle for small nests and benchmarks.
    """
    all_vars = tuple(expr.fields())
    if len(all_vars) == 0:
        return ()

    best_order = all_vars
    best_cost = loop_order_cost(
        expr, best_order, stats_factory, stats_bindings, output_vars
    )
    for order in itertools.permutations(all_vars):
        cost = loop_order_cost(expr, order, stats_factory, stats_bindings, output_vars)
        if cost < best_cost:
            best_order, best_cost = order, cost
    return best_order