query_normalization

alpha_rename_expr(expr, mapping)

Apply a simple alpha-renaming to expr, replacing any Field that appears in mapping with its mapped value.

This traverses only the LogicExpression nodes that actually store fields and leaves all other nodes unchanged.

The goal is to rewrite the dimension variables of an alias definition to match the index variables at a particular call site.

Source code in src/finch/autoschedule/factorizer/galley_factorizer/query_normalization.py

def alpha_rename_expr(
    expr: LogicExpression, mapping: dict[Field, Field]
) -> LogicExpression:
    """
    Apply a simple alpha-renaming to expr, replacing any Field that appears
    in mapping with its mapped value.

    This traverses only the LogicExpression nodes that actually store fields
    and leaves all other nodes unchanged.

    The goal is to rewrite the dimension variables of an alias
    definition to match the index variables at a particular call site.
    """

    match expr:
        case Table(tns, idxs):
            new_idxs = tuple(_rename_field(idx, mapping) for idx in idxs)
            return Table(tns, new_idxs)
        case MapJoin(op, args):
            new_args = tuple(alpha_rename_expr(arg, mapping) for arg in args)
            return MapJoin(op, new_args)
        case Aggregate(op, init, arg, idxs):
            new_arg = alpha_rename_expr(arg, mapping)
            new_idxs = tuple(_rename_field(idx, mapping) for idx in idxs)
            return Aggregate(op, init, new_arg, new_idxs)
        case Reorder(arg, idxs):
            new_arg = alpha_rename_expr(arg, mapping)
            new_idxs = tuple(_rename_field(idx, mapping) for idx in idxs)
            return Reorder(new_arg, new_idxs)
        case Relabel(arg, idxs):
            new_arg = alpha_rename_expr(arg, mapping)
            new_idxs = tuple(_rename_field(idx, mapping) for idx in idxs)
            return Relabel(new_arg, new_idxs)
        case _:
            return expr

map_pointwise_aggregate(node)

Rewrite: Aggregate(op, init, arg, ()) -> MapJoin(op, (init, arg)) Galley starts each reduction from the init of the indices it reduces, so an aggregate which reduces no indices would drop an init which isn't an identity of its operator.

Source code in src/finch/autoschedule/factorizer/galley_factorizer/query_normalization.py

def map_pointwise_aggregate(node: LogicNode) -> LogicNode | None:
    """
    Rewrite: Aggregate(op, init, arg, ()) -> MapJoin(op, (init, arg))
    Galley starts each reduction from the init of the indices it reduces, so an
    aggregate which reduces no indices would drop an init which isn't an
    identity of its operator.
    """
    match node:
        case Aggregate(Literal(op), Literal(init), arg, ()) if (
            op not in (None, ffuncs.overwrite)
            and not is_dynamic(init)
            and not is_identity(op.ftype, init)
        ):
            return MapJoin(node.op, (node.init, arg))
        case _:
            return None

merge_queries(plan)

Merge per-alias queries into larger queries, one per produced tensor.

• Builds a mapping from Alias to its defining Query in the original plan.
• For each alias mentioned in the final Produces statement, constructs a new Query(Table(alias, idxs), merged_rhs) where merged_rhs inlines any Table(Alias, ...) references using _inline_tables_in_expr.
• Ensures that we never drop a Query whose lhs alias appears in a Produces statement

Source code in src/finch/autoschedule/factorizer/galley_factorizer/query_normalization.py

def merge_queries(plan: Plan) -> Plan:
    """
    Merge per-alias queries into larger queries, one per produced tensor.

    - Builds a mapping from Alias to its defining Query in the original plan.
    - For each alias mentioned in the final Produces statement, constructs a
      new Query(Table(alias, idxs), merged_rhs) where merged_rhs inlines any
      Table(Alias, ...) references using _inline_tables_in_expr.
    - Ensures that we never drop a Query whose lhs alias appears in a
      Produces statement
    """
    if not isinstance(plan, Plan):
        return plan

    bodies: tuple[LogicNode, ...] = plan.bodies  # type: ignore[assignment]
    if not bodies:
        return plan

    # Find the last Produces statement, if any.
    produces_stmt: Produces | None = None
    for body in reversed(bodies):
        if isinstance(body, Produces):
            produces_stmt = body
            break

    if produces_stmt is None:
        return plan

    # Map each alias to the last Query that defines it in the original plan.
    alias_to_query: dict[Alias, Query] = {}
    for body in bodies:
        match body:
            case Query(Table(Alias() as alias, _)):
                alias_to_query[alias] = body

    assert all(isinstance(arg, Alias) for arg in produces_stmt.args)
    produced_aliases: tuple[Alias, ...] = produces_stmt.args  # ty: ignore[invalid-assignment]
    produced_alias_set: set[Alias] = set(produced_aliases)

    # Emit the queries in the order they were defined, since a produced alias
    # may read another produced alias.
    ordered_aliases = sorted(
        (alias for alias in produced_alias_set if alias in alias_to_query),
        key=lambda alias: bodies.index(alias_to_query[alias]),
    )
    new_queries: list[Query] = []
    for alias in ordered_aliases:
        defining_query = alias_to_query[alias]
        # Don't inline other produced aliases, resue them
        other_produced = produced_alias_set - {alias}
        merged_rhs = _inline_tables_in_expr(
            defining_query.rhs, alias_to_query, other_produced
        )
        new_queries.append(Query(defining_query.lhs, merged_rhs))

    # Rebuild the plan as: [merged queries..., original Produces]
    if not new_queries:
        raise ValueError("Plan should not be empty after merging queries")

    new_bodies: list[LogicStatement] = []
    new_bodies.extend(new_queries)
    new_bodies.append(produces_stmt)
    return Plan(tuple(new_bodies))

normalize_reorders_in_plan(plan)

Apply normalize_reorders_in_query to each Query body in the plan.

Source code in src/finch/autoschedule/factorizer/galley_factorizer/query_normalization.py

def normalize_reorders_in_plan(plan: Plan) -> Plan:
    """Apply `normalize_reorders_in_query` to each Query body in the plan."""
    if not isinstance(plan, Plan):
        return plan

    new_bodies: list[LogicStatement] = []
    for body in plan.bodies:
        if isinstance(body, Query):
            new_bodies.append(normalize_reorders_in_query(body))
        else:
            new_bodies.append(body)

    return Plan(tuple(new_bodies))

normalize_reorders_in_query(query)

Normalize Reorder usage in a single query.

For a given Query(lhs, rhs): 1. Strip all Reorder nodes from rhs. 2. Remove any aggregates that references idxs created by reorders

Source code in src/finch/autoschedule/factorizer/galley_factorizer/query_normalization.py

def normalize_reorders_in_query(query: Query) -> Query:
    """
    Normalize Reorder usage in a single query.

    For a given Query(lhs, rhs):
      1. Strip all Reorder nodes from rhs.
      2. Remove any aggregates that references idxs created by reorders
    """

    def strip_reorder(node):
        match node:
            case Reorder(arg, _):
                return arg
            case _:
                return node

    def remove_extra_aggregates(node):
        match node:
            case Aggregate(op, init, arg, idxs):
                valid_idxs = [idx for idx in idxs if idx in arg.fields()]
                if len(valid_idxs) > 0:
                    return Aggregate(op, init, arg, tuple(valid_idxs))
                return arg
            case _:
                return node

    inner_rhs = Rewrite(PostWalk(Chain([strip_reorder, remove_extra_aggregates])))(
        query.rhs
    )
    return Query(query.lhs, inner_rhs)

postprocess_plan_after_galley(plan)

Postprocessing used after running Galley greedy optimization.

• Rebuilds the plan so that, for each alias mentioned in the final Produces, there is only one Query computing it. This is necessary because the greedy optimizer may produce multiple queries with the same lhs alias as it reduces different indices, but downstream components expect at most one query per alias.

Source code in src/finch/autoschedule/factorizer/galley_factorizer/query_normalization.py

def postprocess_plan_after_galley(plan: Plan) -> Plan:
    """
    Postprocessing used after running Galley greedy optimization.

    - Rebuilds the plan so that, for each alias mentioned in the final
      Produces, there is only one Query computing it. This is necessary
      because the greedy optimizer may produce multiple queries with the same
      lhs alias as it reduces different indices, but downstream components
      expect at most one query per alias.
    """
    new_bodies: list[LogicStatement] = []
    for body in plan.bodies:
        if isinstance(body, Query):
            new_rhs = Rewrite(PreWalk(split_mapjoin))(body.rhs)
            new_bodies.append(Query(body.lhs, new_rhs))
        else:
            new_bodies.append(body)
    return Plan(tuple(new_bodies))

preprocess_plan_for_galley(plan)

End-to-end preprocessing used before running Galley greedy optimization.

• First writes aggregates which reduce no indices as maps, so their inits are kept.
• Merges alias-based queries to produce roughly one query per produced alias.
• Pushes aggregates up and merges adjacent same-op aggregates so reduction indices are not parent-related (enables cost-optimal reduction order).
• Then normalizes Reorder usage so that no query RHS has Reorder nodes.

Source code in src/finch/autoschedule/factorizer/galley_factorizer/query_normalization.py

def preprocess_plan_for_galley(plan: Plan) -> Plan:
    """
    End-to-end preprocessing used before running Galley greedy optimization.

    - First writes aggregates which reduce no indices as maps, so their inits
      are kept.
    - Merges alias-based queries to produce roughly one query per produced
      alias.
    - Pushes aggregates up and merges adjacent same-op aggregates so reduction
      indices are not parent-related (enables cost-optimal reduction order).
    - Then normalizes Reorder usage so that no query RHS has Reorder nodes.
    """
    plan = Rewrite(PostWalk(map_pointwise_aggregate))(plan)
    merged = merge_queries(plan)
    merged = normalize_reorders_in_plan(merged)
    new_bodies: list[LogicStatement] = []
    for body in merged.bodies:
        if isinstance(body, Query):
            pushed_rhs = push_aggregates_up(body.rhs)
            merged_rhs = Rewrite(PostWalk(merge_mapjoin_rule))(pushed_rhs)
            new_bodies.append(Query(body.lhs, merged_rhs))
        else:
            new_bodies.append(body)
    merged = Plan(tuple(new_bodies))
    return normalize_reorders_in_plan(merged)

push_aggregates_up(expr)

Apply push-aggregate and merge-aggregate rewrites to fixpoint. Produces an expression where reduction indices are not parent-related

Source code in src/finch/autoschedule/factorizer/galley_factorizer/query_normalization.py

def push_aggregates_up(expr: LogicExpression) -> LogicExpression:
    """
    Apply push-aggregate and merge-aggregate rewrites to fixpoint.
    Produces an expression where reduction indices are not parent-related
    """

    def rule(node: LogicNode) -> LogicNode | None:
        if isinstance(node, LogicExpression):
            r = _merge_adjacent_aggregates(node)
            if r is not None:
                return r
            r = _push_aggregate_up(node)
            if r is not None:
                return r
        return None

    result = Rewrite(Fixpoint(PostWalk(rule)))(expr)
    return result if result is not None else expr