Optimizando E-Graphs con Estructuras de Datos Personalizadas y Funciones de Orden Superior

Los E-graphs son una potente estructura de datos para el razonamiento sobre equivalencia de programas y optimización de compiladores. Sin embargo, su rendimiento se degrada significativamente cuando se aplican identidades algebraicas como la asociatividad, conmutatividad y distributividad (A/C/D), lo que provoca una explosión exponencial en el número de nodos del grafo. Este problema fundamental de la computación, la explosión del espacio de búsqueda en la optimización simbólica, se aborda tradicionalmente limitando la aplicación de reglas o el tamaño del E-graph, sacrificando la exhaustividad de la optimización.

La tesis central de este artículo es que, en lugar de modificar el núcleo de los sistemas de E-graphs o limitar las reglas, se pueden utilizar estructuras de datos personalizadas (contenedores) y funciones de orden superior para codificar estas identidades algebraicas directamente en la representación de los datos. Esto permite que el E-graph opere sobre formas normalizadas de expresiones, reduciendo drásticamente la redundancia y el tamaño del grafo, sin perder la capacidad de aplicar reglas de reescritura complejas como la factorización o la plegado de constantes. Este enfoque permite una optimización más eficiente y escalable, especialmente relevante para expresiones simbólicas grandes en dominios como la simulación física o el cálculo de gradientes.

Arquitectura del Sistema

El sistema Egglog, implementado en Rust con bindings en Python, es el framework de E-graphs utilizado. Su arquitectura permite la definición de tipos personalizados, funciones no interpretadas y reglas de reescritura. La clave de la solución reside en la capacidad de Egglog para manejar 'primitivas' y 'contenedores'. Las primitivas son valores opacos (ej. i64, String) sobre los cuales Egglog no razona internamente, pero permite definir funciones. Los contenedores son tipos personalizados que 'contienen' referencias a otras e-classes y deben implementar una operación de 'rebuilding' para mantener la congruencia cuando las e-classes internas se fusionan.

Para mitigar la explosión de nodos, se introduce un tipo de contenedor MultiSet[Num] para representar sumas o productos, donde el orden y la asociatividad son intrínsecos a la estructura de datos, eliminando la necesidad de reglas explícitas de A/C. Por ejemplo, sum_(MultiSet(a, b)) es indistinguible de sum_(MultiSet(b, a)). Para realizar operaciones como el plegado de constantes sobre estos contenedores, se utilizan funciones de orden superior (map, fold) que operan sobre los elementos del MultiSet. Esto evita la necesidad de 'matching' directo sobre el contenido interno del contenedor, que Egglog no soporta de forma nativa. En el caso de la factorización de polinomios, se utiliza un MultiSet[MultiSet[Value]] para representar polinomios como sumas de monomios (productos de términos), permitiendo un algoritmo de factorización de Horner multivariado codificado como reglas de reescritura que operan sobre estas estructuras normalizadas.

from __future__ import annotations
from egglog import *

class Num(Expr):
    def __init__(self, value: i64Like) -> None: ...
    @classmethod
    def var(cls, name: StringLike) -> Num: ...
    def __add__(self, other: Num) -> Num: ...
    def __mul__(self, other: Num) -> Num: ...

@ruleset
def comm_dist_fold(a: Num, b: Num, c: Num, i: i64, j: i64):
    yield rewrite(a + b).to(b + a)
    yield rewrite(a * (b + c)).to((a * b) + (a * c))
    yield rewrite(Num(i) + Num(j)).to(Num(i + j))
    yield rewrite(Num(i) * Num(j)).to(Num(i * j))

egraph = EGraph()
expr1 = egraph.let("expr1", Num(2) * (Num.var("x") + Num(3)))
expr2 = egraph.let("expr2", Num(6) + Num(2) * Num.var("x"))
egraph.saturate(comm_dist_fold)
egraph.check(expr1 == expr2)

@function
def sum_(xs: MultiSet[Num]) -> Num: ...

@function
def get_i64(x: Num) -> i64: ...

@ruleset
def set_get_i64(i: i64):
    yield rule(Num(i)).then(set_(get_i64(Num(i))).to(i))

@ruleset
def constant_fold_sum(initial: Num, xs: MultiSet[Num]):
    constants = xs.map(get_i64)
    remaining = xs - constants.map(Num)
    folded = multiset_fold(i64.__add__, i64(0), constants)
    yield rewrite(sum_(xs)).to(
        sum_(remaining.insert(Num(folded))),
        constants.length() > 1,
    )

egraph = EGraph()
a, b = Num.var("a"), Num.var("b")
new_expr = egraph.let(
    "new_expr", sum_(MultiSet(Num(2), a, b, b, Num(3))))

egraph.run(set_get_i64.saturate() + constant_fold_sum.saturate())
egraph.extract(new_expr)

Fundamentos Teóricos

El problema de la explosión del espacio de búsqueda en la optimización simbólica y la reescritura de términos es un desafío clásico en la informática teórica y la ingeniería de compiladores. Los E-graphs, popularizados por el paper "egg: Fast and Extensible E-Graph Rewriting" (Willard et al., 2021), son una evolución de los sistemas de reescritura de términos y los grafos de compartición de expresiones (expression DAGs). La idea de codificar propiedades algebraicas directamente en estructuras de datos subyacentes tiene raíces en la teoría de los sistemas de reescritura y la lógica de la igualdad, donde la congruencia es un concepto fundamental. La implementación de MultiSet para manejar asociatividad y conmutatividad se alinea con el concepto de formas canónicas o normalizadas, un principio clave en la simplificación algebraica y la demostración automática de teoremas, buscando una representación única para cada clase de equivalencia. La factorización de Horner, aunque originalmente para polinomios univariados, es un ejemplo de algoritmo greedy que busca una representación más compacta, un problema de optimización combinatoria que se beneficia de representaciones de datos que reducen el espacio de búsqueda.