Ematch

ematch

EmatchFunctions dataclass

Bases: InterpreterFunctions

Interpreter functions for PDL patterns operating on e-graphs.

Source code in xdsl/interpreters/ematch.py

@register_impls
@dataclass
class EmatchFunctions(InterpreterFunctions):
    """Interpreter functions for PDL patterns operating on e-graphs."""

    known_ops: KnownOps = field(default_factory=KnownOps)
    """Used for hashconsing operations. When new operations are created, if they are identical to an existing operation,
    the existing operation is reused instead of creating a new one."""

    eclass_union_find: DisjointSet[equivalence.AnyClassOp] = field(
        default_factory=lambda: DisjointSet[equivalence.AnyClassOp]()
    )
    """Union-find structure tracking which e-classes are equivalent and should be merged."""

    worklist: list[equivalence.AnyClassOp] = field(
        default_factory=list[equivalence.AnyClassOp]
    )
    """Worklist of e-classes that need to be processed for matching."""

    analyses: list[SparseForwardDataFlowAnalysis[Lattice[Any]]] = field(
        default_factory=lambda: []
    )
    """The sparse forward analyses to be run during equality saturation.
    These must be registered with a NonPropagatingDataFlowSolver where `propagate` is False.
    This way, state propagation is handled purely by the equality saturation logic.
    """

    @impl(ematch.GetClassValsOp)
    def run_get_class_vals(
        self,
        interpreter: Interpreter,
        op: ematch.GetClassValsOp,
        args: tuple[Any, ...],
    ) -> tuple[Any, ...]:
        """
        Take a value and return all values in its equivalence class.

        If the value is an equivalence.class result, return the operands of the class,
        otherwise return a tuple containing just the value itself.
        """
        assert len(args) == 1
        val = args[0]

        if val is None:
            return ((val,),)

        assert isinstance(val, SSAValue)

        if isinstance(val, OpResult):
            defining_op = val.owner
            if isinstance(defining_op, equivalence.AnyClassOp):
                return (tuple(defining_op.operands),)

        # Value is not an eclass result, return it as a single-element tuple
        return ((val,),)

    @impl(ematch.GetClassRepresentativeOp)
    def run_get_class_representative(
        self,
        interpreter: Interpreter,
        op: ematch.GetClassRepresentativeOp,
        args: tuple[Any, ...],
    ) -> tuple[Any, ...]:
        """
        Get one of the values in the equivalence class of v.
        Returns the first operand of the equivalence class.
        """
        assert len(args) == 1
        val = args[0]

        if val is None:
            return (val,)

        assert isa(val, SSAValue)

        if isinstance(val, OpResult):
            defining_op = val.owner
            if isinstance(defining_op, equivalence.AnyClassOp):
                return (defining_op.operands[0],)

        # Value is not an eclass result, return it as-is
        return (val,)

    @impl(ematch.GetClassResultOp)
    def run_get_class_result(
        self,
        interpreter: Interpreter,
        op: ematch.GetClassResultOp,
        args: tuple[Any, ...],
    ) -> tuple[Any, ...]:
        """
        Get the equivalence.class result corresponding to the equivalence class of v.

        If v has exactly one use and that use is a ClassOp, return the ClassOp's result.
        Otherwise return v unchanged.
        """
        assert len(args) == 1
        val = args[0]

        if val is None:
            return (val,)

        assert isa(val, SSAValue)

        if val.has_one_use():
            user = val.get_user_of_unique_use()
            if isinstance(user, equivalence.AnyClassOp):
                return (user.result,)

        return (val,)

    @impl(ematch.GetClassResultsOp)
    def run_get_class_results(
        self,
        interpreter: Interpreter,
        op: ematch.GetClassResultsOp,
        args: tuple[Any, ...],
    ) -> tuple[Any, ...]:
        """
        Get the equivalence.class results corresponding to the equivalence classes
        of a range of values.
        """
        assert len(args) == 1
        vals = args[0]

        if vals is None:
            return ((),)

        results: list[SSAValue] = []
        for val in vals:
            if val is None:
                results.append(val)
            elif val.has_one_use():
                user = val.get_user_of_unique_use()
                if isinstance(user, equivalence.AnyClassOp):
                    results.append(user.result)
                else:
                    results.append(val)
            else:
                results.append(val)

        return (tuple(results),)

    def get_or_create_class(
        self, interpreter: Interpreter, val: SSAValue
    ) -> equivalence.AnyClassOp:
        """
        Get the equivalence class for a value, creating one if it doesn't exist.
        """
        if isinstance(val, OpResult):
            # If val is defined by a ClassOp, return it
            if isinstance(val.owner, equivalence.AnyClassOp):
                return val.owner
            insertpoint = InsertPoint.before(val.owner)
        else:
            assert isinstance(val.owner, Block)
            insertpoint = InsertPoint.at_start(val.owner)

        # If val has one use and it's a ClassOp, return it
        if (user := val.get_user_of_unique_use()) is not None:
            if isinstance(user, equivalence.AnyClassOp):
                return user

        # If the value is not part of an eclass yet, create one
        rewriter = PDLInterpFunctions.get_rewriter(interpreter)

        eclass_op = equivalence.ClassOp(val)
        rewriter.insert_op(eclass_op, insertpoint)
        self.eclass_union_find.add(eclass_op)

        # Replace uses of val with the eclass result (except in the eclass itself)
        rewriter.replace_uses_with_if(
            val, eclass_op.result, lambda use: use.operation is not eclass_op
        )

        return eclass_op

    def eclass_union(
        self,
        interpreter: Interpreter,
        a: equivalence.AnyClassOp,
        b: equivalence.AnyClassOp,
    ) -> bool:
        """Unions two eclasses, merging their operands and results.
        Returns True if the eclasses were merged, False if they were already the same."""
        a = self.eclass_union_find.find(a)
        b = self.eclass_union_find.find(b)

        if a == b:
            return False

        # Meet the analysis states of the two e-classes
        for analysis in self.analyses:
            a_lattice = analysis.get_lattice_element(a.result)
            b_lattice = analysis.get_lattice_element(b.result)
            a_lattice.meet(b_lattice)

        if isinstance(a, equivalence.ConstantClassOp):
            if isinstance(b, equivalence.ConstantClassOp):
                assert a.value == b.value, (
                    "Trying to union two different constant eclasses.",
                )
            to_keep, to_replace = a, b
            self.eclass_union_find.union_left(to_keep, to_replace)
        elif isinstance(b, equivalence.ConstantClassOp):
            to_keep, to_replace = b, a
            self.eclass_union_find.union_left(to_keep, to_replace)
        else:
            self.eclass_union_find.union(
                a,
                b,
            )
            to_keep = self.eclass_union_find.find(a)
            to_replace = b if to_keep is a else a
        # Operands need to be deduplicated because it can happen the same operand was
        # used by different parent eclasses after their children were merged:
        new_operands = OrderedSet(to_keep.operands)
        new_operands.update(to_replace.operands)
        to_keep.operands = new_operands

        for use in to_replace.result.uses:
            # uses are removed from the hashcons before the replacement is carried out.
            # (because the replacement changes the operations which means we cannot find them in the hashcons anymore)
            if use.operation in self.known_ops:
                self.known_ops.pop(use.operation)

        rewriter = PDLInterpFunctions.get_rewriter(interpreter)
        rewriter.replace_op(to_replace, new_ops=[], new_results=to_keep.results)
        return True

    def union_val(self, interpreter: Interpreter, a: SSAValue, b: SSAValue) -> None:
        """
        Union two values into the same equivalence class.
        """
        if a == b:
            return

        eclass_a = self.get_or_create_class(interpreter, a)
        eclass_b = self.get_or_create_class(interpreter, b)

        if self.eclass_union(interpreter, eclass_a, eclass_b):
            self.worklist.append(eclass_a)

    @impl(ematch.UnionOp)
    def run_union(
        self,
        interpreter: Interpreter,
        op: ematch.UnionOp,
        args: tuple[Any, ...],
    ) -> tuple[Any, ...]:
        """
        Merge two values, an operation and a value range, or two value ranges
        into equivalence class(es).

        Supported operand type combinations:
        - (value, value): merge two values
        - (operation, range<value>): merge operation results with values
        - (range<value>, range<value>): merge two value ranges
        """
        assert len(args) == 2
        lhs, rhs = args

        if isa(lhs, SSAValue) and isa(rhs, SSAValue):
            # (Value, Value) case
            self.union_val(interpreter, lhs, rhs)

        elif isinstance(lhs, Operation) and isa(rhs, Sequence[SSAValue]):
            # (Operation, ValueRange) case
            assert len(lhs.results) == len(rhs), (
                "Operation result count must match value range size"
            )
            for result, val in zip(lhs.results, rhs, strict=True):
                self.union_val(interpreter, result, val)

        elif isa(lhs, Sequence[SSAValue]) and isa(rhs, Sequence[SSAValue]):
            # (ValueRange, ValueRange) case
            assert len(lhs) == len(rhs), "Value ranges must have equal size"
            for val_lhs, val_rhs in zip(lhs, rhs, strict=True):
                self.union_val(interpreter, val_lhs, val_rhs)

        else:
            raise InterpretationError(
                f"union: unsupported argument types: {type(lhs)}, {type(rhs)}"
            )

        return ()

    @impl(ematch.DedupOp)
    def run_dedup(
        self,
        interpreter: Interpreter,
        op: ematch.DedupOp,
        args: tuple[Any, ...],
    ) -> tuple[Any, ...]:
        """
        Check if the operation already exists in the hashcons.

        If an equivalent operation exists, erase the input operation and return
        the existing one. Otherwise, insert the operation into the hashcons and
        return it.
        """
        assert len(args) == 1
        input_op = args[0]
        assert isinstance(input_op, Operation)

        # Check if an equivalent operation exists in hashcons
        existing = self.known_ops.get(input_op)

        if existing is not None and existing is not input_op:
            # Deduplicate: erase the new op and return existing
            rewriter = PDLInterpFunctions.get_rewriter(interpreter)
            rewriter.erase_op(input_op)
            return (existing,)

        # No duplicate found, insert into hashcons
        self.known_ops[input_op] = input_op
        return (input_op,)

    def repair(self, interpreter: Interpreter, eclass: equivalence.AnyClassOp):
        """
        Repair an e-class by finding and merging duplicate parent operations.

        This method:
        1. Finds all operations that use this e-class's result
        2. Identifies structurally equivalent operations among them
        3. Merges equivalent operations by unioning their result e-classes
        4. Updates dataflow analysis states
        """
        rewriter = PDLInterpFunctions.get_rewriter(interpreter)
        eclass = self.eclass_union_find.find(eclass)

        if eclass.parent is None:
            return

        unique_parents = KnownOps()

        # Collect e-class parents (operations that use this class's result)
        # Use OrderedSet to maintain deterministic ordering
        user_ops = OrderedSet(use.operation for use in eclass.result.uses)

        # Collect pairs of duplicate operations to merge AFTER the loop
        # This avoids modifying the hash map while iterating
        to_merge: list[tuple[Operation, Operation]] = []

        for op1 in user_ops:
            # Skip eclass operations themselves
            if isinstance(op1, equivalence.AnyClassOp):
                continue

            op2 = unique_parents.get(op1)

            if op2 is not None:
                # Found an equivalent operation - record for later merging
                to_merge.append((op1, op2))
            else:
                unique_parents[op1] = op1

        # Now perform all merges after we're done with the hash map
        for op1, op2 in to_merge:
            # Collect eclass pairs for ALL results before replacement
            eclass_pairs: list[
                tuple[equivalence.AnyClassOp, equivalence.AnyClassOp]
            ] = []
            for res1, res2 in zip(op1.results, op2.results, strict=True):
                eclass1 = self.get_or_create_class(interpreter, res1)
                eclass2 = self.get_or_create_class(interpreter, res2)
                eclass_pairs.append((eclass1, eclass2))

            # Replace op1 with op2's results
            rewriter.replace_op(op1, new_ops=(), new_results=op2.results)

            # Process each eclass pair
            for eclass1, eclass2 in eclass_pairs:
                if eclass1 == eclass2:
                    # Same eclass - just deduplicate operands
                    eclass1.operands = OrderedSet(eclass1.operands)
                else:
                    # Different eclasses - union them
                    if self.eclass_union(interpreter, eclass1, eclass2):
                        self.worklist.append(eclass1)

        # Update dataflow analysis for all parent operations
        eclass = self.eclass_union_find.find(eclass)
        for op in OrderedSet(use.operation for use in eclass.result.uses):
            if isinstance(op, equivalence.AnyClassOp):
                continue

            point = ProgramPoint.before(op)

            for analysis in self.analyses:
                operands = [
                    analysis.get_lattice_element_for(point, o) for o in op.operands
                ]
                results = [analysis.get_lattice_element(r) for r in op.results]

                if not results:
                    continue

                original_state: Any = None
                # For each result, reset to bottom and recompute
                for result in results:
                    original_state = result.value
                    result._value = result.value_cls()  # pyright: ignore[reportPrivateUsage]

                analysis.visit_operation_impl(op, operands, results)

                # Check if any result changed
                for result in results:
                    assert original_state is not None
                    changed = result.meet(type(result)(result.anchor, original_state))
                    if changed == ChangeResult.CHANGE:
                        # Find the eclass for this result and add to worklist
                        if (op_use := op.results[0].first_use) is not None:
                            if isinstance(
                                eclass_op := op_use.operation, equivalence.AnyClassOp
                            ):
                                self.worklist.append(eclass_op)
                        break  # Only need to add to worklist once per operation

    def rebuild(self, interpreter: Interpreter):
        while self.worklist:
            todo = OrderedSet(self.eclass_union_find.find(c) for c in self.worklist)
            self.worklist.clear()
            for c in todo:
                self.repair(interpreter, c)

known_ops: KnownOps = field(default_factory=KnownOps) class-attribute instance-attribute

Used for hashconsing operations. When new operations are created, if they are identical to an existing operation, the existing operation is reused instead of creating a new one.

eclass_union_find: DisjointSet[equivalence.AnyClassOp] = field(default_factory=(lambda: DisjointSet[equivalence.AnyClassOp]())) class-attribute instance-attribute

Union-find structure tracking which e-classes are equivalent and should be merged.

worklist: list[equivalence.AnyClassOp] = field(default_factory=(list[equivalence.AnyClassOp])) class-attribute instance-attribute

Worklist of e-classes that need to be processed for matching.

analyses: list[SparseForwardDataFlowAnalysis[Lattice[Any]]] = field(default_factory=(lambda: [])) class-attribute instance-attribute

The sparse forward analyses to be run during equality saturation. These must be registered with a NonPropagatingDataFlowSolver where propagate is False. This way, state propagation is handled purely by the equality saturation logic.

__init__(known_ops: KnownOps = KnownOps(), eclass_union_find: DisjointSet[equivalence.AnyClassOp] = (lambda: DisjointSet[equivalence.AnyClassOp]())(), worklist: list[equivalence.AnyClassOp] = list[equivalence.AnyClassOp](), analyses: list[SparseForwardDataFlowAnalysis[Lattice[Any]]] = (lambda: [])()) -> None

run_get_class_vals(interpreter: Interpreter, op: ematch.GetClassValsOp, args: tuple[Any, ...]) -> tuple[Any, ...]

Take a value and return all values in its equivalence class.

If the value is an equivalence.class result, return the operands of the class, otherwise return a tuple containing just the value itself.

Source code in xdsl/interpreters/ematch.py

@impl(ematch.GetClassValsOp)
def run_get_class_vals(
    self,
    interpreter: Interpreter,
    op: ematch.GetClassValsOp,
    args: tuple[Any, ...],
) -> tuple[Any, ...]:
    """
    Take a value and return all values in its equivalence class.

    If the value is an equivalence.class result, return the operands of the class,
    otherwise return a tuple containing just the value itself.
    """
    assert len(args) == 1
    val = args[0]

    if val is None:
        return ((val,),)

    assert isinstance(val, SSAValue)

    if isinstance(val, OpResult):
        defining_op = val.owner
        if isinstance(defining_op, equivalence.AnyClassOp):
            return (tuple(defining_op.operands),)

    # Value is not an eclass result, return it as a single-element tuple
    return ((val,),)

run_get_class_representative(interpreter: Interpreter, op: ematch.GetClassRepresentativeOp, args: tuple[Any, ...]) -> tuple[Any, ...]

Get one of the values in the equivalence class of v. Returns the first operand of the equivalence class.

Source code in xdsl/interpreters/ematch.py

@impl(ematch.GetClassRepresentativeOp)
def run_get_class_representative(
    self,
    interpreter: Interpreter,
    op: ematch.GetClassRepresentativeOp,
    args: tuple[Any, ...],
) -> tuple[Any, ...]:
    """
    Get one of the values in the equivalence class of v.
    Returns the first operand of the equivalence class.
    """
    assert len(args) == 1
    val = args[0]

    if val is None:
        return (val,)

    assert isa(val, SSAValue)

    if isinstance(val, OpResult):
        defining_op = val.owner
        if isinstance(defining_op, equivalence.AnyClassOp):
            return (defining_op.operands[0],)

    # Value is not an eclass result, return it as-is
    return (val,)

run_get_class_result(interpreter: Interpreter, op: ematch.GetClassResultOp, args: tuple[Any, ...]) -> tuple[Any, ...]

Get the equivalence.class result corresponding to the equivalence class of v.

If v has exactly one use and that use is a ClassOp, return the ClassOp's result. Otherwise return v unchanged.

Source code in xdsl/interpreters/ematch.py

@impl(ematch.GetClassResultOp)
def run_get_class_result(
    self,
    interpreter: Interpreter,
    op: ematch.GetClassResultOp,
    args: tuple[Any, ...],
) -> tuple[Any, ...]:
    """
    Get the equivalence.class result corresponding to the equivalence class of v.

    If v has exactly one use and that use is a ClassOp, return the ClassOp's result.
    Otherwise return v unchanged.
    """
    assert len(args) == 1
    val = args[0]

    if val is None:
        return (val,)

    assert isa(val, SSAValue)

    if val.has_one_use():
        user = val.get_user_of_unique_use()
        if isinstance(user, equivalence.AnyClassOp):
            return (user.result,)

    return (val,)

run_get_class_results(interpreter: Interpreter, op: ematch.GetClassResultsOp, args: tuple[Any, ...]) -> tuple[Any, ...]

Get the equivalence.class results corresponding to the equivalence classes of a range of values.

Source code in xdsl/interpreters/ematch.py

@impl(ematch.GetClassResultsOp)
def run_get_class_results(
    self,
    interpreter: Interpreter,
    op: ematch.GetClassResultsOp,
    args: tuple[Any, ...],
) -> tuple[Any, ...]:
    """
    Get the equivalence.class results corresponding to the equivalence classes
    of a range of values.
    """
    assert len(args) == 1
    vals = args[0]

    if vals is None:
        return ((),)

    results: list[SSAValue] = []
    for val in vals:
        if val is None:
            results.append(val)
        elif val.has_one_use():
            user = val.get_user_of_unique_use()
            if isinstance(user, equivalence.AnyClassOp):
                results.append(user.result)
            else:
                results.append(val)
        else:
            results.append(val)

    return (tuple(results),)

get_or_create_class(interpreter: Interpreter, val: SSAValue) -> equivalence.AnyClassOp

Get the equivalence class for a value, creating one if it doesn't exist.

Source code in xdsl/interpreters/ematch.py

def get_or_create_class(
    self, interpreter: Interpreter, val: SSAValue
) -> equivalence.AnyClassOp:
    """
    Get the equivalence class for a value, creating one if it doesn't exist.
    """
    if isinstance(val, OpResult):
        # If val is defined by a ClassOp, return it
        if isinstance(val.owner, equivalence.AnyClassOp):
            return val.owner
        insertpoint = InsertPoint.before(val.owner)
    else:
        assert isinstance(val.owner, Block)
        insertpoint = InsertPoint.at_start(val.owner)

    # If val has one use and it's a ClassOp, return it
    if (user := val.get_user_of_unique_use()) is not None:
        if isinstance(user, equivalence.AnyClassOp):
            return user

    # If the value is not part of an eclass yet, create one
    rewriter = PDLInterpFunctions.get_rewriter(interpreter)

    eclass_op = equivalence.ClassOp(val)
    rewriter.insert_op(eclass_op, insertpoint)
    self.eclass_union_find.add(eclass_op)

    # Replace uses of val with the eclass result (except in the eclass itself)
    rewriter.replace_uses_with_if(
        val, eclass_op.result, lambda use: use.operation is not eclass_op
    )

    return eclass_op

eclass_union(interpreter: Interpreter, a: equivalence.AnyClassOp, b: equivalence.AnyClassOp) -> bool

Unions two eclasses, merging their operands and results. Returns True if the eclasses were merged, False if they were already the same.

Source code in xdsl/interpreters/ematch.py

def eclass_union(
    self,
    interpreter: Interpreter,
    a: equivalence.AnyClassOp,
    b: equivalence.AnyClassOp,
) -> bool:
    """Unions two eclasses, merging their operands and results.
    Returns True if the eclasses were merged, False if they were already the same."""
    a = self.eclass_union_find.find(a)
    b = self.eclass_union_find.find(b)

    if a == b:
        return False

    # Meet the analysis states of the two e-classes
    for analysis in self.analyses:
        a_lattice = analysis.get_lattice_element(a.result)
        b_lattice = analysis.get_lattice_element(b.result)
        a_lattice.meet(b_lattice)

    if isinstance(a, equivalence.ConstantClassOp):
        if isinstance(b, equivalence.ConstantClassOp):
            assert a.value == b.value, (
                "Trying to union two different constant eclasses.",
            )
        to_keep, to_replace = a, b
        self.eclass_union_find.union_left(to_keep, to_replace)
    elif isinstance(b, equivalence.ConstantClassOp):
        to_keep, to_replace = b, a
        self.eclass_union_find.union_left(to_keep, to_replace)
    else:
        self.eclass_union_find.union(
            a,
            b,
        )
        to_keep = self.eclass_union_find.find(a)
        to_replace = b if to_keep is a else a
    # Operands need to be deduplicated because it can happen the same operand was
    # used by different parent eclasses after their children were merged:
    new_operands = OrderedSet(to_keep.operands)
    new_operands.update(to_replace.operands)
    to_keep.operands = new_operands

    for use in to_replace.result.uses:
        # uses are removed from the hashcons before the replacement is carried out.
        # (because the replacement changes the operations which means we cannot find them in the hashcons anymore)
        if use.operation in self.known_ops:
            self.known_ops.pop(use.operation)

    rewriter = PDLInterpFunctions.get_rewriter(interpreter)
    rewriter.replace_op(to_replace, new_ops=[], new_results=to_keep.results)
    return True

union_val(interpreter: Interpreter, a: SSAValue, b: SSAValue) -> None

Union two values into the same equivalence class.

Source code in xdsl/interpreters/ematch.py

def union_val(self, interpreter: Interpreter, a: SSAValue, b: SSAValue) -> None:
    """
    Union two values into the same equivalence class.
    """
    if a == b:
        return

    eclass_a = self.get_or_create_class(interpreter, a)
    eclass_b = self.get_or_create_class(interpreter, b)

    if self.eclass_union(interpreter, eclass_a, eclass_b):
        self.worklist.append(eclass_a)

run_union(interpreter: Interpreter, op: ematch.UnionOp, args: tuple[Any, ...]) -> tuple[Any, ...]

Merge two values, an operation and a value range, or two value ranges into equivalence class(es).

Supported operand type combinations: - (value, value): merge two values - (operation, range): merge operation results with values - (range, range): merge two value ranges

Source code in xdsl/interpreters/ematch.py

@impl(ematch.UnionOp)
def run_union(
    self,
    interpreter: Interpreter,
    op: ematch.UnionOp,
    args: tuple[Any, ...],
) -> tuple[Any, ...]:
    """
    Merge two values, an operation and a value range, or two value ranges
    into equivalence class(es).

    Supported operand type combinations:
    - (value, value): merge two values
    - (operation, range<value>): merge operation results with values
    - (range<value>, range<value>): merge two value ranges
    """
    assert len(args) == 2
    lhs, rhs = args

    if isa(lhs, SSAValue) and isa(rhs, SSAValue):
        # (Value, Value) case
        self.union_val(interpreter, lhs, rhs)

    elif isinstance(lhs, Operation) and isa(rhs, Sequence[SSAValue]):
        # (Operation, ValueRange) case
        assert len(lhs.results) == len(rhs), (
            "Operation result count must match value range size"
        )
        for result, val in zip(lhs.results, rhs, strict=True):
            self.union_val(interpreter, result, val)

    elif isa(lhs, Sequence[SSAValue]) and isa(rhs, Sequence[SSAValue]):
        # (ValueRange, ValueRange) case
        assert len(lhs) == len(rhs), "Value ranges must have equal size"
        for val_lhs, val_rhs in zip(lhs, rhs, strict=True):
            self.union_val(interpreter, val_lhs, val_rhs)

    else:
        raise InterpretationError(
            f"union: unsupported argument types: {type(lhs)}, {type(rhs)}"
        )

    return ()

run_dedup(interpreter: Interpreter, op: ematch.DedupOp, args: tuple[Any, ...]) -> tuple[Any, ...]

Check if the operation already exists in the hashcons.

If an equivalent operation exists, erase the input operation and return the existing one. Otherwise, insert the operation into the hashcons and return it.

Source code in xdsl/interpreters/ematch.py

@impl(ematch.DedupOp)
def run_dedup(
    self,
    interpreter: Interpreter,
    op: ematch.DedupOp,
    args: tuple[Any, ...],
) -> tuple[Any, ...]:
    """
    Check if the operation already exists in the hashcons.

    If an equivalent operation exists, erase the input operation and return
    the existing one. Otherwise, insert the operation into the hashcons and
    return it.
    """
    assert len(args) == 1
    input_op = args[0]
    assert isinstance(input_op, Operation)

    # Check if an equivalent operation exists in hashcons
    existing = self.known_ops.get(input_op)

    if existing is not None and existing is not input_op:
        # Deduplicate: erase the new op and return existing
        rewriter = PDLInterpFunctions.get_rewriter(interpreter)
        rewriter.erase_op(input_op)
        return (existing,)

    # No duplicate found, insert into hashcons
    self.known_ops[input_op] = input_op
    return (input_op,)

repair(interpreter: Interpreter, eclass: equivalence.AnyClassOp)

Repair an e-class by finding and merging duplicate parent operations.

This method: 1. Finds all operations that use this e-class's result 2. Identifies structurally equivalent operations among them 3. Merges equivalent operations by unioning their result e-classes 4. Updates dataflow analysis states

Source code in xdsl/interpreters/ematch.py

def repair(self, interpreter: Interpreter, eclass: equivalence.AnyClassOp):
    """
    Repair an e-class by finding and merging duplicate parent operations.

    This method:
    1. Finds all operations that use this e-class's result
    2. Identifies structurally equivalent operations among them
    3. Merges equivalent operations by unioning their result e-classes
    4. Updates dataflow analysis states
    """
    rewriter = PDLInterpFunctions.get_rewriter(interpreter)
    eclass = self.eclass_union_find.find(eclass)

    if eclass.parent is None:
        return

    unique_parents = KnownOps()

    # Collect e-class parents (operations that use this class's result)
    # Use OrderedSet to maintain deterministic ordering
    user_ops = OrderedSet(use.operation for use in eclass.result.uses)

    # Collect pairs of duplicate operations to merge AFTER the loop
    # This avoids modifying the hash map while iterating
    to_merge: list[tuple[Operation, Operation]] = []

    for op1 in user_ops:
        # Skip eclass operations themselves
        if isinstance(op1, equivalence.AnyClassOp):
            continue

        op2 = unique_parents.get(op1)

        if op2 is not None:
            # Found an equivalent operation - record for later merging
            to_merge.append((op1, op2))
        else:
            unique_parents[op1] = op1

    # Now perform all merges after we're done with the hash map
    for op1, op2 in to_merge:
        # Collect eclass pairs for ALL results before replacement
        eclass_pairs: list[
            tuple[equivalence.AnyClassOp, equivalence.AnyClassOp]
        ] = []
        for res1, res2 in zip(op1.results, op2.results, strict=True):
            eclass1 = self.get_or_create_class(interpreter, res1)
            eclass2 = self.get_or_create_class(interpreter, res2)
            eclass_pairs.append((eclass1, eclass2))

        # Replace op1 with op2's results
        rewriter.replace_op(op1, new_ops=(), new_results=op2.results)

        # Process each eclass pair
        for eclass1, eclass2 in eclass_pairs:
            if eclass1 == eclass2:
                # Same eclass - just deduplicate operands
                eclass1.operands = OrderedSet(eclass1.operands)
            else:
                # Different eclasses - union them
                if self.eclass_union(interpreter, eclass1, eclass2):
                    self.worklist.append(eclass1)

    # Update dataflow analysis for all parent operations
    eclass = self.eclass_union_find.find(eclass)
    for op in OrderedSet(use.operation for use in eclass.result.uses):
        if isinstance(op, equivalence.AnyClassOp):
            continue

        point = ProgramPoint.before(op)

        for analysis in self.analyses:
            operands = [
                analysis.get_lattice_element_for(point, o) for o in op.operands
            ]
            results = [analysis.get_lattice_element(r) for r in op.results]

            if not results:
                continue

            original_state: Any = None
            # For each result, reset to bottom and recompute
            for result in results:
                original_state = result.value
                result._value = result.value_cls()  # pyright: ignore[reportPrivateUsage]

            analysis.visit_operation_impl(op, operands, results)

            # Check if any result changed
            for result in results:
                assert original_state is not None
                changed = result.meet(type(result)(result.anchor, original_state))
                if changed == ChangeResult.CHANGE:
                    # Find the eclass for this result and add to worklist
                    if (op_use := op.results[0].first_use) is not None:
                        if isinstance(
                            eclass_op := op_use.operation, equivalence.AnyClassOp
                        ):
                            self.worklist.append(eclass_op)
                    break  # Only need to add to worklist once per operation

rebuild(interpreter: Interpreter)

Source code in xdsl/interpreters/ematch.py

def rebuild(self, interpreter: Interpreter):
    while self.worklist:
        todo = OrderedSet(self.eclass_union_find.find(c) for c in self.worklist)
        self.worklist.clear()
        for c in todo:
            self.repair(interpreter, c)