Refinement E-Graphs: Incorporando Relaciones de Orden Parcial en la Optimización de Programas

Los e-graphs (equality graphs) han demostrado ser una herramienta poderosa para la optimización de programas y la reescritura de términos, al representar conjuntos de términos equivalentes en una estructura compacta y permitir la aplicación eficiente de reglas de reescritura. Sin embargo, muchas optimizaciones en compiladores no son meras equivalencias, sino 'refinamientos': transformaciones que mueven un programa de una especificación más abstracta o ambigua a una más concreta y ejecutable, sin perder corrección. Por ejemplo, la resolución de 'don't care conditions' en el diseño de circuitos digitales o la especificación de comportamientos ante desbordamientos de enteros.

El problema fundamental que los refinement e-graphs buscan resolver es cómo integrar esta noción de orden parcial (<=) directamente en la estructura de datos y los algoritmos de los e-graphs, de manera que sea tan 'baked-in' como la igualdad. Esto permite que las optimizaciones que son unidireccionales por naturaleza sean modeladas y aplicadas de forma más natural y eficiente, en lugar de ser codificadas de forma ad-hoc o como equivalencias forzadas que podrían limitar el espacio de optimización o introducir errores semánticos.

Arquitectura del Sistema

La arquitectura de un refinement e-graph se construye sobre la base de un e-graph tradicional, que combina un Union-Find para gestionar las clases de equivalencia (e-classes) y un hash consing para la deduplicación de nodos (enodes). La innovación principal es la sustitución del Union-Find estándar por un 'Inequality Union-Find'. Este componente no solo gestiona la igualdad, sino que también mantiene conjuntos de 'uppers' y 'lowers' para cada e-class, representando las relaciones 'mayor o igual que' (>=) y 'menor o igual que' (<=) respectivamente. Estas relaciones se propagan transitivamente y colapsan ciclos a igualdad.

Los algoritmos clave de los e-graphs —congruence closure, e-matching y extracción— se adaptan para manejar la relación de orden. La 'refinement closure' reemplaza la congruence closure, aplicando reglas de inferencia que respetan la monotonicidad de las funciones. Por ejemplo, si a <= b y c >= d, y una función diff es monótona en su primer argumento y anti-monótona en el segundo, entonces diff(a,c) <= diff(b,d). Para esto, cada símbolo de función se anota con una 'variance signature' (MONO, ANTI, UNKNOWN) para cada uno de sus argumentos. El e-matching y la extracción se parametrizan con un 'Mode' (LE, EQ, GE) que indica la dirección de la búsqueda, y este modo se ajusta según la varianza de los argumentos al descender por el árbol de términos. La extracción, por ejemplo, puede buscar el término 'más refinado' (más concreto) o el 'menos refinado' (más abstracto) que satisfaga una relación dada.

class LE_EGraph():
    parents : list[Id] = field(default_factory=list)
    memo : dict[Node, Id] = field(default_factory=dict)
    uppers : list[set[Id]] = field(default_factory=list)
    lowers : list[set[Id]] = field(default_factory=list)
    variance : dict[str, tuple[Variance, ...]] = field(default_factory=dict)

    def find(self, a : Id) -> Id:
        while self.parents[a] != a:
            a = self.parents[a]
        return a

    def union(self, a : Id, b : Id) -> None:
        a,b = self.find(a), self.find(b)
        if a == b:
            return
        self.parents[b] = a
        self.uppers[a] |= self.uppers[b]
        self.lowers[a] |= self.lowers[b]

    def add_le(self, a : Id, b : Id) -> None:
        a,b = self.find(a), self.find(b)
        if a == b:
            return
        self.uppers[a].add(b)
        self.lowers[b].add(a)

    def all_le(self, a : Id) -> set[Id]:
        a = self.find(a)
        u = {self.find(x) for x in self.uppers[a]}
        u.add(a)
        todo = list(u)
        while todo:
            x = self.find(todo.pop())
            for y in self.uppers[x]:
                y = self.find(y)
                if y not in u:
                    u.add(y)
                    todo.append(y)
        return u

class Mode(Enum):
    LE = -1
    EQ = 0
    GE = 1

class Variance(Enum):
    MONO = 1
    ANTI = -1
    UNKNOWN = 0

    def act(self, mode: Mode) -> Mode:
        match self:
            case Variance.MONO:
                return mode
            case Variance.ANTI:
                return Mode(-mode.value)
            case Variance.UNKNOWN:
                return Mode.EQ

Fundamentos Teóricos

La noción de e-graphs tiene sus raíces en el trabajo de Robert Tarjan sobre el algoritmo de Union-Find para la detección de ciclos y la gestión de conjuntos disjuntos, y en el concepto de 'congruence closure' desarrollado por Nelson y Oppen en la década de 1970 para la decisión de la satisfacibilidad de teorías combinadas. La extensión a 'refinement' se conecta con la teoría de la reescritura de términos y los sistemas de tipos con subtipos, donde las relaciones de orden parcial son fundamentales.

Conceptos como la monotonicidad y la contravarianza de funciones en relación con un orden parcial son bien conocidos en la lógica y la teoría de tipos, especialmente en el contexto de la inferencia de tipos y la programación funcional. La idea de 'don't care conditions' se remonta a los fundamentos del diseño de circuitos digitales y la minimización de funciones booleanas, donde se utilizan para simplificar expresiones lógicas y optimizar la implementación de hardware. La integración de estas ideas en una estructura de e-graph proporciona un marco unificado para aplicar principios de optimización que históricamente han sido tratados de forma separada.