Skip to content

loop_order_greedy

connected_loop_candidates(prefix, remaining, output_vars=None)

Preserve candidate order while respecting the output's index order.

Source code in src/finch/autoschedule/loop_orderer/loop_order_greedy.py 
33
34
35
36
37
38
39
40
41
42
def connected_loop_candidates(
    prefix: tuple[Field, ...],
    remaining: Sequence[Field],
    output_vars: tuple[Field, ...] | None = None,
) -> list[Field]:
    """Preserve candidate order while respecting the output's index order."""
    # TODO this could eventually be improved to consider output transpose.
    output_vars = output_vars or ()
    next_output = next((idx for idx in output_vars if idx not in prefix), None)
    return [idx for idx in remaining if idx not in output_vars or idx == next_output]

greedy_loop_order(expr, stats_factory, stats_bindings, output_vars=None, *, cost_model=None)

Build a loop order by appending the cheapest candidate prefix.

Source code in src/finch/autoschedule/loop_orderer/loop_order_greedy.py 
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
def greedy_loop_order(
    expr: LogicExpression,
    stats_factory: StatsFactory[NS],
    stats_bindings: Mapping[Alias, NS],
    output_vars: tuple[Field, ...] | None = None,
    *,
    cost_model: CostModel[Prefix] | None = None,
) -> tuple[Field, ...]:
    """Build a loop order by appending the cheapest candidate prefix."""
    if cost_model is None:
        cost_model = Sum(PigeonCostModel(), TransposeCostModel())
    machine = StatsMachine(stats_factory, dict(stats_bindings))
    match expr:
        case Aggregate(arg=arg):
            pass
        case _:
            arg = expr
    fields = (*arg.fields(), *(expr.fields() if output_vars is None else output_vars))
    remaining = list(dict.fromkeys(fields))
    prefix: list[Field] = []

    def candidate_cost(field: Field) -> float:
        return cost_model.predict_component_cost(
            Prefix((*prefix, field), expr, output_vars), machine
        )

    while remaining:
        candidates = connected_loop_candidates(tuple(prefix), remaining, output_vars)
        best = min(candidates, key=candidate_cost)
        prefix.append(best)
        remaining.remove(best)

    return tuple(prefix)