Skip to content

loop_order_bnb

loop_order_bfs(expr, stats_factory, stats_bindings, output_vars=None, *, k=None, cost_model=None)

Search prefixes with nonnegative incremental model costs as lower bounds.

Source code in src/finch/autoschedule/loop_orderer/loop_order_bnb.py 
37
38
39
40
41
42
43
44
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
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
def loop_order_bfs(
    expr: LogicExpression,
    stats_factory: StatsFactory[NS],
    stats_bindings: Mapping[Alias, NS],
    output_vars: tuple[Field, ...] | None = None,
    *,
    k: int | None = None,
    cost_model: CostModel[Prefix] | None = None,
) -> tuple[Field, ...]:
    """Search prefixes with nonnegative incremental model costs as lower bounds."""
    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))
    all_vars = tuple(dict.fromkeys(fields))
    if not all_vars:
        return ()

    best_order = greedy_loop_order(
        expr, stats_factory, stats_bindings, output_vars, cost_model=cost_model
    )
    best_cost = sum(
        cost_model.predict_component_cost(
            Prefix(best_order[:n], expr, output_vars), machine
        )
        for n in range(len(best_order) + 1)
    )
    initial_cost = cost_model.predict_component_cost(
        Prefix((), expr, output_vars), machine
    )
    prev_new_optimal_orders: list[tuple[tuple[Field, ...], float]] = [
        ((), initial_cost)
    ]

    for _ in all_vars:
        new_optimal_orders: list[tuple[tuple[Field, ...], float]] = []
        for prefix, prefix_cost in prev_new_optimal_orders:
            prefix_set = set(prefix)
            remaining = [field for field in all_vars if field not in prefix_set]
            # Preserves remaining / all_vars order (see connected_loop_candidates).
            candidates = connected_loop_candidates(prefix, remaining, output_vars)
            for field in candidates:
                new_prefix = prefix + (field,)
                new_cost = prefix_cost + cost_model.predict_component_cost(
                    Prefix(new_prefix, expr, output_vars), machine
                )
                if new_cost >= best_cost:
                    continue
                if len(new_prefix) == len(all_vars):
                    best_order, best_cost = new_prefix, new_cost
                else:
                    new_optimal_orders.append((new_prefix, new_cost))
        if k is not None:
            new_optimal_orders.sort(key=lambda state: state[1])
            new_optimal_orders = new_optimal_orders[:k]
        prev_new_optimal_orders = new_optimal_orders

    return best_order

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

Pick the loop order with the lowest model 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 
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
def loop_order_brute_force(
    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, ...]:
    """Pick the loop order with the lowest model 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.
    """
    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))
    all_vars = tuple(dict.fromkeys(fields))
    best_order = ()
    best_cost = float("inf")
    for order in itertools.permutations(all_vars):
        if (
            output_vars is not None
            and tuple(idx for idx in order if idx in output_vars) != output_vars
        ):
            continue
        cost = sum(
            cost_model.predict_component_cost(
                Prefix(order[:n], expr, output_vars), machine
            )
            for n in range(len(order) + 1)
        )
        if not best_order or cost < best_cost:
            best_order, best_cost = order, cost
    return best_order

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

Search prefixes with nonnegative incremental model costs as lower bounds.

Source code in src/finch/autoschedule/loop_orderer/loop_order_bnb.py 
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
def loop_order_dfs(
    expr: LogicExpression,
    stats_factory: StatsFactory[NS],
    stats_bindings: MutableMapping[Alias, NS],
    output_vars: tuple[Field, ...] | None = None,
    *,
    cost_model: CostModel[Prefix] | None = None,
) -> tuple[Field, ...]:
    """Search prefixes with nonnegative incremental model costs as lower bounds."""
    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))
    all_vars = tuple(dict.fromkeys(fields))
    if not all_vars:
        return ()

    best_order = greedy_loop_order(
        expr, stats_factory, stats_bindings, output_vars, cost_model=cost_model
    )
    best_cost = sum(
        cost_model.predict_component_cost(
            Prefix(best_order[:n], expr, output_vars), machine
        )
        for n in range(len(best_order) + 1)
    )
    initial_cost = cost_model.predict_component_cost(
        Prefix((), expr, output_vars), machine
    )
    stack: list[tuple[tuple[Field, ...], float]] = [((), initial_cost)]

    while stack:
        prefix, prefix_cost = stack.pop()
        prefix_set = set(prefix)
        remaining = [field for field in all_vars if field not in prefix_set]
        children: list[tuple[tuple[Field, ...], float]] = []

        # Preserves remaining / all_vars order (see connected_loop_candidates).
        candidates = connected_loop_candidates(prefix, remaining, output_vars)
        for field in candidates:
            new_prefix = prefix + (field,)
            new_cost = prefix_cost + cost_model.predict_component_cost(
                Prefix(new_prefix, expr, output_vars), machine
            )
            if new_cost >= best_cost:
                continue
            if len(new_prefix) == len(all_vars):
                best_order, best_cost = new_prefix, new_cost
            else:
                children.append((new_prefix, new_cost))

        children.sort(key=lambda child: child[1], reverse=True)
        stack.extend(children)

    return best_order