Skip to content

Equivalence Pattern Rewriter to integrate expansion of Softmax into equality saturation#5655

Open
szerdick wants to merge 73 commits intoxdslproject:jumerckx/interp_ematchfrom
szerdick:interp_ematch
Open

Equivalence Pattern Rewriter to integrate expansion of Softmax into equality saturation#5655
szerdick wants to merge 73 commits intoxdslproject:jumerckx/interp_ematchfrom
szerdick:interp_ematch

Conversation

@szerdick
Copy link
Collaborator

No description provided.

superlopuh and others added 30 commits January 22, 2026 15:19
With the foreach approach, I believe part of this logic can be
simplified/omitted, but this is a good first step.
We need to be more careful with the bindings selected by matchers.
Right now, values are sometimes passed as their class result to a
rewriter:
```
x = …
y = …

c = class x, y

rewrite “myrewrite”(c)
```

This is problematic in MLIR as other rewrites can invalidate c. While we
are guaranteed to not erase any operation in the “program domain”, this
is not true for the operations of the “e-graph domain”.
The solution is to ensure in the pdl lowering that only values from the
program domain are passed to rewriters and the conversion to e-graph
domain happens at the latest possible point, in the rewriter itself.
(This is not an issue in xDSL because erasing an e-class does not
discard of the object)
non-eqsat lowering should stay the same
(not entirely done)

Co-authored-by: Sasha Lopoukhine <superlopuh@gmail.com>
to show how it should be done, the change would also need to be made to
apply_eqsat_pdl.
repair detects when two parent operations have become equal due to their
children having been merged.
At this point, there are two identical operations, but the hashcons
(`known_ops`) only tracks one: there is a collision.
One of the two operations is replaced by the other. If the hashcons
happened to store the operation that was replaced, instead of the
(identical) replacement, the hashcons is corrupt.

This is fixed by explicitly updating the hashcons to point to the
operation that is not replaced.
often it doesn't seem to do anything on top of the sorted list, but
better safe than sorry?
@szerdick szerdick requested a review from jumerckx February 11, 2026 19:12
@codecov
Copy link

codecov bot commented Feb 11, 2026

Codecov Report

❌ Patch coverage is 47.33096% with 148 lines in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (jumerckx/interp_ematch@e64571a). Learn more about missing BASE report.

Files with missing lines Patch % Lines
xdsl/eqsat_bookkeeper.py 29.20% 132 Missing and 11 partials ⚠️
xdsl/pattern_rewriter_eq.py 91.89% 0 Missing and 3 partials ⚠️
xdsl/transforms/ematch_exp.py 94.44% 1 Missing and 1 partial ⚠️
Additional details and impacted files
@@                    Coverage Diff                    @@
##             jumerckx/interp_ematch    #5655   +/-   ##
=========================================================
  Coverage                          ?   85.55%           
=========================================================
  Files                             ?      408           
  Lines                             ?    57585           
  Branches                          ?     6721           
=========================================================
  Hits                              ?    49264           
  Misses                            ?     6743           
  Partials                          ?     1578           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

_worklist: Worklist = field(default_factory=Worklist, init=False)
"""The worklist of operations to walk over."""

rewriter_factory: type[PatternRewriter] = PatternRewriter
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
rewriter_factory: type[PatternRewriter] = PatternRewriter
rewriter_factory: type[PatternRewriter] = PatternRewriter
"""Factory method that takes an operation and returns a PatternRewriter"""

@jumerckx
Copy link
Collaborator

This looks mostly correct to me. Would be good to have some tests to see if things actually work.
What wouldn't work yet:

  • insert for an op that already exists in IR and is used by an eclass (if that op result is used in the rewrite, it would not point to the eclass result)
  • rebuild/repair (i.e. when merging two classes makes two other values become equal)
  • insert with multiple ops.

I think the first two can be tackled separately. But you could already write tests that verify that things are broken and we can mark these as expectedly failing.
The bookkeeping now isn't used in ematch.py yet, I'm also not sure if everything that in there now needs to be there (i.e. the run_* methods should perhaps need to be renamed), but this refactor should also be a separate pr.

For now adding tests seems most useful, we can talk about cleaning up afterwards.

Mia Sophie Zerdick and others added 14 commits February 13, 2026 11:13
… interp_ematch

# Conflicts:
#	tests/filecheck/transforms/ematch-saturate/binom_prod.mlir
#	xdsl/dialects/equivalence.py
#	xdsl/interpreters/ematch.py
#	xdsl/pattern_rewriter.py
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
to show how it should be done, the change would also need to be made to
apply_eqsat_pdl.
repair detects when two parent operations have become equal due to their
children having been merged.
At this point, there are two identical operations, but the hashcons
(`known_ops`) only tracks one: there is a collision.
One of the two operations is replaced by the other. If the hashcons
happened to store the operation that was replaced, instead of the
(identical) replacement, the hashcons is corrupt.

This is fixed by explicitly updating the hashcons to point to the
operation that is not replaced.
@jumerckx jumerckx force-pushed the jumerckx/interp_ematch branch from 42006ed to e64571a Compare February 13, 2026 13:38
@jumerckx jumerckx force-pushed the jumerckx/interp_ematch branch from e64571a to ba97362 Compare February 13, 2026 14:04
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants