util

desugar_query_into(root)

Replace each in-place QueryInto with the equivalent Query, which reads the previous value of its left-hand side explicitly.

Source code in src/finch/autoschedule/util.py

def desugar_query_into(root: LogicStatement) -> LogicStatement:
    """
    Replace each in-place `QueryInto` with the equivalent `Query`, which reads
    the previous value of its left-hand side explicitly.
    """

    def rule(stmt):
        match stmt:
            case QueryInto() as q:
                return q.as_query()

    return Rewrite(PostWalk(rule))(root)

drop_query_reorders(root)

drop_query_reorders(root: LogicStatement) -> LogicStatement

drop_query_reorders(root: LogicNode) -> LogicNode

Remove a Reorder at the root of each query's right-hand side. The table on the left-hand side of a query already fixes the order of its result.

Source code in src/finch/autoschedule/util.py

def drop_query_reorders(root):
    """
    Remove a Reorder at the root of each query's right-hand side. The table on
    the left-hand side of a query already fixes the order of its result.
    """

    def rule(stmt):
        match stmt:
            case Query(lhs, Reorder(arg, _)):
                return Query(lhs, arg)
            case QueryInto(lhs, op, Reorder(Aggregate() as arg, _)):
                return QueryInto(lhs, op, arg)

    return Rewrite(PostWalk(rule))(root)

reorder_to(ex, idxs)

ex with its fields in the order idxs.

Source code in src/finch/autoschedule/util.py

def reorder_to(ex: LogicExpression, idxs: tuple[Field, ...]) -> LogicExpression:
    """`ex` with its fields in the order `idxs`."""
    return ex if ex.fields() == idxs else Reorder(ex, idxs)

resugar_query_into(root)

Replace each query of an aggregate which starts from a tensor with an in-place QueryInto. The aggregate folds its argument into a copy of its init, so the copy is skipped when the init is the table being written. Inits which are broadcast over the output can't be copied into it, so they are left as they are.

Source code in src/finch/autoschedule/util.py

def resugar_query_into(root: LogicStatement) -> LogicStatement:
    """
    Replace each query of an aggregate which starts from a tensor with an
    in-place `QueryInto`. The aggregate folds its argument into a copy of its
    init, so the copy is skipped when the init is the table being written.
    Inits which are broadcast over the output can't be copied into it, so they
    are left as they are.
    """

    def rule(stmt):
        match stmt:
            case Query(
                Table(Alias() as lhs, idxs) as tbl,
                Aggregate(op, init, arg, agg_idxs),
            ) if (
                not isinstance(init, Literal)
                and lhs not in PostOrderDFS(arg)
                and set(setdiff(arg.fields(), agg_idxs)) == set(idxs)
                and set(init.fields()) == set(idxs)
                and (init == tbl or lhs not in PostOrderDFS(init))
            ):
                update = QueryInto(tbl, op, arg)
                return update if init == tbl else Plan((Query(tbl, init), update))

    return Rewrite(PostWalk(rule))(root)

split_aggregate_inits(root, bindings, broadcast_only=False)

Replace each aggregate which starts from a tensor with a map of that tensor and a reduction which starts from the identity of the aggregate's operator. With broadcast_only, only the inits which are broadcast over some fields of the aggregate are split.

Source code in src/finch/autoschedule/util.py

def split_aggregate_inits(
    root: LogicStatement,
    bindings: dict[Alias, TensorFType],
    broadcast_only: bool = False,
) -> LogicStatement:
    """
    Replace each aggregate which starts from a tensor with a map of that tensor
    and a reduction which starts from the identity of the aggregate's operator.
    With `broadcast_only`, only the inits which are broadcast over some fields
    of the aggregate are split.
    """

    # Element types are only needed to find identities, so they are inferred
    # lazily.
    @cache
    def element_types():
        return root.infer_element_type(
            {var: tns.element_type for var, tns in bindings.items()}
        )

    def rule(node):
        match node:
            case Aggregate(Literal(op) as op_lit, init, arg, idxs) if not (
                isinstance(init, Literal)
                or (broadcast_only and set(init.fields()) == set(node.fields()))
            ):
                if idxs:
                    z = init_value(op.ftype, init.element_type(element_types()))
                    arg = Aggregate(op_lit, Literal(z), arg, idxs)
                return MapJoin(op_lit, (init, arg))

    return Rewrite(PostWalk(rule))(root)