Set the loop order of a query. An aggregate query holds its loop order in
a Reorder of the aggregate's argument, and an in-place update holds it in a
Reorder of its right-hand side. An in-place update can't change the layout
of the table it updates, so its loop order is permuted to visit the fields
of that table in order. Aggregate queries also include output-only unit
dimensions and visit the output fields in order.
Source code in src/finch/autoschedule/loop_orderer/loop_ordering.py
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197 | def with_loop_order(
stmt: LogicStatement, loop_order: tuple[Field, ...]
) -> LogicStatement:
"""
Set the loop order of a query. An aggregate query holds its loop order in
a Reorder of the aggregate's argument, and an in-place update holds it in a
Reorder of its right-hand side. An in-place update can't change the layout
of the table it updates, so its loop order is permuted to visit the fields
of that table in order. Aggregate queries also include output-only unit
dimensions and visit the output fields in order.
"""
match stmt:
case Query(Table(_, lhs_idxs) as lhs, Aggregate(op, init, arg, idxs)):
loop_order = with_subsequence(intersect(lhs_idxs, loop_order), loop_order)
loop_order = toposort([list(loop_order), list(lhs_idxs)])
return Query(lhs, Aggregate(op, init, Reorder(arg, loop_order), idxs))
case QueryInto(Table(_, lhs_idxs) as lhs, op, arg):
loop_order = with_subsequence(intersect(lhs_idxs, loop_order), loop_order)
return QueryInto(lhs, op, Reorder(arg, loop_order))
case _:
raise ValueError(f"Expected an aggregate or in-place query, got {stmt}")
|