Skip to content

loop_order_greedy

connected_loop_candidates(prefix, remaining, conjunct_stats, disjunct_stats)

All fields still available for the next loop index.

Source code in src/finch/autoschedule/loop_orderer/loop_order_greedy.py 
32
33
34
35
36
37
38
39
def connected_loop_candidates(
    prefix: tuple[Field, ...],
    remaining: Sequence[Field],
    conjunct_stats: Sequence[TensorStats],
    disjunct_stats: Sequence[TensorStats],
) -> list[Field]:
    """All fields still available for the next loop index."""
    return list(remaining)

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

Build a loop order one index at a time, appending the cheapest candidate.

Each step is scored with the same terms loop_order_cost charges for that prefix -- get_prefix_cost for the reads and writes, plus the reformats the new index forces -- so the summed score of the returned order equals its loop_order_cost.

Source code in src/finch/autoschedule/loop_orderer/loop_order_greedy.py 
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
def greedy_loop_order(
    expr: LogicExpression,
    stats_factory: StatsFactory,
    stats_bindings: Mapping[Alias, NS],
    output_vars: tuple[Field, ...] | None = None,
) -> tuple[Field, ...]:
    """Build a loop order one index at a time, appending the cheapest candidate.

    Each step is scored with the same terms ``loop_order_cost`` charges for that
    prefix -- ``get_prefix_cost`` for the reads and writes, plus the reformats
    the new index forces -- so the summed score of the returned order equals its
    ``loop_order_cost``.
    """
    all_vars = tuple(dict.fromkeys(expr.fields()))
    if not all_vars:
        return ()

    stats_bindings_2 = dict(stats_bindings)
    cache: dict[object, NS] = {}
    conjunct_stats, disjunct_stats = get_conjunctive_and_disjunctive_inputs(
        expr, stats_factory, stats_bindings_2, cache
    )

    # Deduplicated by identity, matching how loop_order_cost charges reformats.
    input_stats: list[NS] = []
    for stat in itertools.chain(conjunct_stats, disjunct_stats):
        if not any(stat is seen for seen in input_stats):
            input_stats.append(stat)

    prefix: list[Field] = []
    remaining = list(all_vars)
    charged: frozenset[int] = frozenset()

    def candidate_cost(field: Field) -> float:
        new_prefix = tuple(prefix) + (field,)
        prefix_cost = get_prefix_cost(
            new_prefix, conjunct_stats, disjunct_stats, stats_factory, output_vars
        )
        return prefix_cost + transpose_penalty(input_stats, new_prefix, charged)

    while remaining:
        candidates = connected_loop_candidates(
            tuple(prefix), remaining, conjunct_stats, disjunct_stats
        )
        best = min(candidates, key=candidate_cost)
        prefix.append(best)
        remaining.remove(best)
        charged = get_reformat_set(input_stats, tuple(prefix))

    return tuple(prefix)

transpose_penalty(input_stats, prefix, charged)

Cost of the reformats that prefix newly forces.

needs_reformat is monotone in the prefix: an index placed before one that precedes it in a tensor's storage order already forces that tensor to be reformatted, and extending the prefix cannot undo it. So each tensor is charged exactly once, at the step where its reformat becomes unavoidable, and the total over a full order matches loop_order_cost.

Source code in src/finch/autoschedule/loop_orderer/loop_order_greedy.py 
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
def transpose_penalty(
    input_stats: Sequence[TensorStats],
    prefix: tuple[Field, ...],
    charged: frozenset[int],
) -> float:
    """Cost of the reformats that ``prefix`` newly forces.

    ``needs_reformat`` is monotone in the prefix: an index placed before one
    that precedes it in a tensor's storage order already forces that tensor to
    be reformatted, and extending the prefix cannot undo it. So each tensor is
    charged exactly once, at the step where its reformat becomes unavoidable,
    and the total over a full order matches ``loop_order_cost``.
    """
    return sum(
        cost_of_reformat(input_stats[i])
        for i in get_reformat_set(input_stats, prefix) - charged
    )