Support custom extraction with cost functions

[saulshanabrook]

Currently, the extraction in egglog is rather limited. It does a tree-based extraction (meaning that if a node shows up twice, it will be counted twice) and requires static costs per function.

The first issue, the type of extractor, could be alleviated by using some extractors from extraction gym. The second, having some custom costs per item could be addressed upstream in egglog (egraphs-good/egglog#294) but is not on the immediate roadmap.

Either way, it would also be nice to have fully custom extraction. Being able to iterate through the e-graph and do what you will...

Currently, it's "possible" by serializing the e-graph to JSON. But this is not ideal because then you have to look at JSON with random keys and they might not map to your Python function names and it's not type safe and... Yeah it's just a real pain!

So I think it would make sense to add an interface that allows:

1. Using custom extractors from extraction-gym without leaving the Python bindings by building that with egglog.
2. Being able to set custom costs per record before extracting.
3. Being able to query costs and write your own extractor in Python
4. Do all of this while keeping static-type safety.
5. Reduce overhead as much as possible in terms of serialization and wrapping/unwrapping.

Possible Design

Here is a possible API design for the extractors

"""
Examples using egglog.
"""

from __future__ import annotations

from typing import Literal, Protocol, TypeVar

from egglog import Expr

EXPR = TypeVar("EXPR", bound=Expr)


class EGraph:
    def serialize(self) -> SerializedEGraph:
        """
        Serializes the e-graph into a format that can be passed to an extractor.
        """
        raise NotImplementedError

    def extract(
        self, x: EXPR, /, extractor: Extractor, include_cost: Literal["tree", "dag"] | None
    ) -> EXPR | tuple[EXPR, int]:
        """
        Extracts the given expression using the given extractor, optionally including the cost of the extraction.
        """
        extract_result = extractor.extract(self.serialize(), [x])
        res = extract_result.chosen(x)
        if include_cost is None:
            return res
        cost = extract_result.tree_cost([x]) if include_cost == "tree" else extract_result.dag_cost([x])
        return res, cost


class SerializedEGraph:
    def equivalent(self, x: EXPR, /) -> list[EXPR]:
        """
        Returns all equivalent expressions. i.e. all expressions with the same e-class.
        """
        raise NotImplementedError

    def self_cost(self, x: Expr, /) -> int:
        """
        Returns the cost of just that function, not including its children.
        """
        raise NotImplementedError

    def set_self_cost(self, x: Expr, cost: int, /) -> None:
        """
        Sets the cost of just that function, not including its children.
        """
        raise NotImplementedError


class Extractor(Protocol):
    def extract(self, egraph: SerializedEGraph, roots: list[Expr]) -> ExtractionResult: ...


class ExtractionResult:
    """
    An extraction result is a mapping from an e-class to chosen nodes, paired with the original extracted e-graph.

    Based off of https://github.com/egraphs-good/extraction-gym/blob/main/src/extract/mod.rs but removed e-classes
    since they are not present in Python bindings and instead let you pass in any member of that e-class and get out
    representative nodes.
    """

    egraph: SerializedEGraph

    def __init__(self, egraph: SerializedEGraph) -> None: ...

    def choose(self, class_: EXPR, chosen_node: EXPR, /) -> None:
        """
        Choose an expression in the e-graph.
        """

    def chosen(self, x: EXPR, /) -> EXPR:
        """
        Given an expr that is in the e-graph, it recursively returns the chosen expressions in each e-class.
        """
        raise NotImplementedError

    def check(self) -> None:
        """
        Check the extraction result for consistency.
        """
        raise NotImplementedError

    def find_cycles(self, roots: list[Expr]) -> list[Expr]:
        """
        Returns all classes that are in a cycle, which is reachable from the roots.
        """
        raise NotImplementedError

    def tree_cost(self, roots: list[Expr]) -> int:
        """
        Returns the "tree cost" (counting duplicate nodes twice) of the trees rooted at the given roots.
        """
        raise NotImplementedError

    def dag_cost(self, roots: list[Expr]) -> int:
        """
        Returns the "dag cost" (counting duplicate nodes once) of the dag rooted at the given roots.
        """
        raise NotImplementedError

Using this interface, you could use the default costs form egglog and use a custom extractor, as shown in the helper extract method.

However, you could also set custom costs before serializing, overriding any from egglog:

serialized = egraph.serialize()
serialized.set_self_cost(x, 10)
serialized.set_self_cost(y, 100)
extractor.extract(serialized, [x]).chosen(x)

How would you be able to traverse an expression at runtime and see its children? I think with three small additions, we could be able to do this with our current API:

1. Primitives: Allow any primitive to be converted to a Python object with, i.e. int(i64(0))
2. User Defined Constants: Allow bool(eq(x).to(y)) which will resolve to whether the two sides are exactly syntactically equal.
3. User Defined Functions: Support a new way to get the args in a type safe manner based on a function, i.e. fn_matches(x, f) would return a boolean to say whether the function matches, and then fn_args(x, f) would return a list of the args. They could be typed like this:

class _FnMatchesBuilder(Generic[EXPR]):
    def fn(self, y: Callable[[Unpack[EXPRS]], EXPR], /) -> tuple[Unpack[EXPRS]] | None:
        """
        Returns the list of args or None
        """
        raise NotImplementedError


EXPRS = TypeVarTuple("EXPRS")


def matches(x: EXPR) -> _FnMatchesBuilder[EXPR]:
    raise NotImplementedError


if args := matches(x).fn(y):
    x, y, z = args

Alternatively, how would you create a custom extractor? We would want to add one more way to traverse expressions... This time not caring about what particular expression they are, just their args and a way to re-ccreate them with different args. Using that, we could write a simple tree based extractor:

def decompose(x: EXPR, /) -> tuple[ReturnsExpr[EXPR], list[Expr]]:
    """
    Decomposes an expression into a callable that will reconstruct it based on its args.

    For all expressions, constants or functions, this should hold:

    >>> fn, args = decompose(x)
    >>> assert fn(*args) == x

    This can be used to change the args of a function and then reconstruct it.

    If you are looking for a type safe way to deal with a particular constructor, you can use either `eq(x).to(y)` for
    constants or `match(x).fn(y)` for functions to get their args in a type safe manner.
    """
    raise NotImplementedError


def tree_based_extractor(serialized: SerializedEGraph, expr: EXPR, /) -> tuple[EXPR, int]:
    """
    Returns the lowest cost equivalent expression and the cost of that expression, based on a tree based extraction.
    """
    min_expr, min_cost = expr_cost(serialized, expr)
    for eq in serialized.equivalent(expr):
        new_expr, new_cost = expr_cost(serialized, eq)
        if new_cost < min_cost:
            min_cost = new_cost
            min_expr = new_expr
    return min_expr, min_cost


def expr_cost(serialized: SerializedEGraph, expr: EXPR, /) -> tuple[EXPR, int]:
    """
    Returns the cost of the given expression.
    """
    cost = serialized.self_cost(expr)
    constructor, children = decompose(expr)
    best_children = []
    for child in children:
        best_child, child_cost = tree_based_extractor(serialized, child)
        cost += child_cost
        best_children.append(best_child)
    return constructor(*best_children), cost

[saulshanabrook]

I think we would also need a way to get all parent nodes from a node in the serialized format for doing custom cost traversal... For example you might set some kind of length of a vec, and then want to look that up when computing costs.

[sklam]

For your reference, I wrote a PoC custom extraction with custom cost model at sklam/prototypes_2025@983b81d
It is heavily inspired by Tensat (https://github.com/uwplse/tensat/blob/master/src/optimize.rs#L57C12-L57C25).

The example expands x ** 4 to x * x * x * x, such that AST size cost will not work. The custom cost-model will penalize the Pow a lot so it will select the Mul variant.

How would you be able to traverse an expression at runtime and see its children? I think with three small additions, we could be able to do this with our current API:

This is currently hard to do and therefore omitted in my PoC. I would want to know that it is a Pow(x, 4) and compute cost knowing the 4.

For short term, is there a way to associate node in the serialized json back to the egglog-python Expr object?

Alternatively, how would you create a custom extractor?

I'm very interested in the decompose() and constructor(). Our workflow is compiler IR -> egglog -> compiler IR. We need the extracted result to be translated back to the IR nodes.

[orenht]

Is using an extractor like an ILP extractor currently supported in upstream egglog? Or support would be needed both there and in the python bindings?

[orenht]

If it is of any interest, I built a small and (hopefully) extensible python library that supports custom extraction and custom cost models. You can implement a cost model or extractor independently, and I plan to port some of extraction-gym extractors for my needs. Then using some ugly deserialization, I convert it back to an egglog object for further processing or code generation.

It is based on @sklam POC, i.e. parsing the serialized JSON and using networkx.

It does not support the matching logic you described in the issue without serializing. I believe the proposed matching might be too "specific", in the sense that I don't always want to match against a full expression (instead of a partial one - i.e. I raise to a power and there is an addition in the exponent). In addition, how does it support cycles in the graph?

It is not fully in the spirit of the library, where everything is strongly typed, but it was necessary to me to support at least a custom cost model that can "peek" to the children of nodes.

[saulshanabrook]

Oh cool! Yeah, definitely of interest, feel free to post a link. This is what I would like to work on next, so in particular, if there are test cases or use cases that use the logic you wrote, those would be super helpful to look at to see if I can get something upstreamed here that covers them...

[sklam]

I have been making progress with my POC. Here's a long notebook that goes through compiling a Python function, encoding it into egraph, custom extraction, and MLIR codegen: https://numba.github.io/sealir/demo_geglu_approx_mlir.html#extracting-an-optimized-representation

I have changed the extraction logic since the POC. The previous code mishandles cycles. The new algorithm approximate the solution using an iterative relaxation and it naturally deals with cycles without much complexity: https://github.com/sklam/sealir/blob/05e85a1ed2c7ff4c6f5ec5f7dd901833aa3516ee/sealir/eqsat/rvsdg_extract.py#L145-L202 . For the cost-modeling, my next step is to look into Pareto curve for time-cost vs power-cost (or vs precision). The new algorithm is easier to reason about (for me) in that kind of scenario. So, I'll want the cost to be a ND vector soon.