Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
szerdick wants to merge 73 commits into
base: jumerckx/interp_ematch
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
73 commits
Select commit Hold shift + click to select a range
51c1ceb
Merge branch 'main' into jumerckx/pdl-to-pdl-interp-pass
superlopuh Jan 16, 2026
fa1688c
Merge branch 'main' into jumerckx/pdl-to-pdl-interp-pass
jumerckx Jan 20, 2026
6c45648
Add choosenode
jumerckx Jan 19, 2026
2578775
handle choosenode in _optimize_tree
jumerckx Jan 20, 2026
bc7f2da
insert get_eq_vals + foreach operations before get_defining_op
jumerckx Jan 19, 2026
a1c60f7
fix tests
jumerckx Jan 26, 2026
2c9d1cd
add option to print predicatetree for debugging
jumerckx Jan 20, 2026
3c76e40
zip strictly
jumerckx Jan 20, 2026
419d2de
handle inserting new predicates after a choosenode by rotating the tree
jumerckx Jan 20, 2026
3ae9192
fix generate_choose_node
jumerckx Jan 22, 2026
ccf089e
eqsat instrumentation for get_result in matching
jumerckx Jan 22, 2026
eb82b9a
eqsat instrumentation for create_operation
jumerckx Jan 23, 2026
1274de6
eqsat instrumentation for replace
jumerckx Jan 23, 2026
56c9fd7
Add more get_class_vals
jumerckx Jan 25, 2026
684fe7a
pop foreach failure block after a choice is visited
jumerckx Jan 25, 2026
3a39220
Do not pass ClassOp results to rewriters
jumerckx Jan 26, 2026
c8f2a47
insert foreach operations conditionally
jumerckx Jan 26, 2026
18ab0ad
fix docs issue
jumerckx Jan 26, 2026
75a2deb
make worklist a deque
jumerckx Jan 29, 2026
69ace3f
cleanup init
jumerckx Jan 29, 2026
956df7a
outline functions
jumerckx Jan 29, 2026
99a4565
add lit test for debug tree printing
jumerckx Jan 29, 2026
adbc6a1
anotha one
jumerckx Jan 29, 2026
508390d
Merge branch 'main' into jumerckx/pdl-to-pdl-interp-eqsat
jumerckx Jan 29, 2026
e9061e6
sasha's review
jumerckx Feb 3, 2026
f15112d
list comprehension
jumerckx Feb 3, 2026
921dcb5
Merge branch 'main' into jumerckx/pdl-to-pdl-interp-eqsat
jumerckx Feb 4, 2026
795a3ab
revert defaultdict change
jumerckx Feb 4, 2026
f021fff
pdl_interp: defer rewrite application
jumerckx Feb 4, 2026
b8544af
update apply_eqsat_pdl_interp
jumerckx Feb 4, 2026
c3d50e9
pattern_rewriter: add replace_uses_with_if
jumerckx Feb 4, 2026
bc78c8d
test replace_uses_with_if
jumerckx Feb 4, 2026
bc9674e
pdl-to-pdl-interp: generate ematch ops instead of rewrites
jumerckx Feb 4, 2026
22f6f10
ematch interpreter implementations
jumerckx Feb 4, 2026
2666de2
add ematch-saturate pass
jumerckx Feb 4, 2026
14af4f5
fixup! add ematch-saturate pass
jumerckx Feb 4, 2026
db16657
[TEMP] foreach interpreter function
jumerckx Feb 4, 2026
7e1160e
pdl_interp.create_range interpreter method
jumerckx Feb 4, 2026
1748779
equivalence.graph add operand
jumerckx Feb 5, 2026
78ccb7e
add binom_prod example
jumerckx Feb 5, 2026
4d6b861
fix issue where known ops entries are corrupted due to collision
jumerckx Feb 5, 2026
8ba74c2
Merge branch 'main' into jumerckx/pdl-to-pdl-interp-eqsat
jumerckx Feb 4, 2026
2b9401f
revert defaultdict change
jumerckx Feb 4, 2026
fea9a48
pdl_interp: defer rewrite application
jumerckx Feb 4, 2026
6f86659
update apply_eqsat_pdl_interp
jumerckx Feb 4, 2026
e30ef47
pdl-to-pdl-interp: generate ematch ops instead of rewrites
jumerckx Feb 4, 2026
ce9be60
ematch interpreter implementations
jumerckx Feb 4, 2026
ce34c0e
add ematch-saturate pass
jumerckx Feb 4, 2026
3cd4b95
fixup! add ematch-saturate pass
jumerckx Feb 4, 2026
cdb2a4a
pdl_interp.create_range interpreter method
jumerckx Feb 4, 2026
5728274
equivalence.graph add operand
jumerckx Feb 5, 2026
d532bad
clean up equivalence assembly format
jumerckx Feb 5, 2026
cac2feb
add binom_prod example
jumerckx Feb 5, 2026
42006ed
fix issue where known ops entries are corrupted due to collision
jumerckx Feb 5, 2026
b5f5693
call _stable_topological_sort
jumerckx Feb 11, 2026
498cd20
eqsat: add bookkeeper and equivalence pattern rewriter
Feb 11, 2026
57886d2
eqsat: add bookkeeper and equivalence pattern rewriter
Feb 11, 2026
7d70eee
Merge branch 'interp_ematch' of https://github.com/szerdick/xdsl into…
Feb 13, 2026
ce8baf3
eqsat: update pattern rewriter and equivalence pattern rewriter
Feb 13, 2026
cb7b773
add test and small modifications
Feb 13, 2026
5fafa7b
pdl_interp: defer rewrite application
jumerckx Feb 4, 2026
2bfcda5
update apply_eqsat_pdl_interp
jumerckx Feb 4, 2026
0d6270c
pdl-to-pdl-interp: generate ematch ops instead of rewrites
jumerckx Feb 4, 2026
a873e1e
ematch interpreter implementations
jumerckx Feb 4, 2026
bac24a9
add ematch-saturate pass
jumerckx Feb 4, 2026
2f0177a
fixup! add ematch-saturate pass
jumerckx Feb 4, 2026
d38a6f9
pdl_interp.create_range interpreter method
jumerckx Feb 4, 2026
6639a7e
equivalence.graph add operand
jumerckx Feb 5, 2026
df74bf3
clean up equivalence assembly format
jumerckx Feb 5, 2026
f7e62ac
add binom_prod example
jumerckx Feb 5, 2026
e64571a
fix issue where known ops entries are corrupted due to collision
jumerckx Feb 5, 2026
074d8b8
Merge remote-tracking branch 'upstream/jumerckx/interp_ematch' into i…
Feb 13, 2026
3d2f953
Merge branch 'jumerckx/interp_ematch' into interp_ematch
szerdick Feb 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions tests/filecheck/mlir-conversion/with-mlir/apply-eqsat-pdl/apply_eqsat_pdl_extra_file.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
// XFAIL: *

// RUN: xdsl-opt %s -p 'apply-eqsat-pdl{pdl_file="%p/extra_file.mlir"}' | filecheck %s

// CHECK: %x_c = equivalence.class %x : i32
Expand Down
2 changes: 2 additions & 0 deletions tests/filecheck/mlir-conversion/with-mlir/apply-eqsat-pdl/egg_example.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
// XFAIL: *

// RUN: xdsl-opt %s -p apply-eqsat-pdl | filecheck %s
// RUN: xdsl-opt %s -p apply-eqsat-pdl{individual_patterns=true} | filecheck %s --check-prefix=INDIVIDUAL

Expand Down
2 changes: 2 additions & 0 deletions tests/filecheck/transforms/apply-eqsat-pdl-interp/egg_example.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
// XFAIL: *

// RUN: xdsl-opt %s -p apply-eqsat-pdl-interp | filecheck %s

func.func @impl() -> i32 {
Expand Down
111 changes: 111 additions & 0 deletions ...s/filecheck/transforms/convert-pdl-to-pdl-interp/pdl-to-pdl-interp-debug-tree-larger.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,111 @@
// RUN: xdsl-opt -p convert-pdl-to-pdl-interp{print-debug-info=true} %s | filecheck %s

// CHECK: Bool[root.result[0]] IsNotNullQuestion -> TrueAnswer()
// CHECK-NEXT: ├── success:
// CHECK-NEXT: │ └── Switch[root] OperationNameQuestion
// CHECK-NEXT: │ ├── case StringAnswer(value='arith.divui'):
// CHECK-NEXT: │ │ └── Bool[root] OperandCountQuestion -> UnsignedAnswer(value=2)
// CHECK-NEXT: │ │ └── success:
// CHECK-NEXT: │ │ └── Bool[root] ResultCountQuestion -> UnsignedAnswer(value=1)
// CHECK-NEXT: │ │ └── success:
// CHECK-NEXT: │ │ └── Bool[root.operand[0]] IsNotNullQuestion -> TrueAnswer()
// CHECK-NEXT: │ │ └── success:
// CHECK-NEXT: │ │ └── Bool[root.operand[1]] IsNotNullQuestion -> TrueAnswer()
// CHECK-NEXT: │ │ ├── success:
// CHECK-NEXT: │ │ │ └── Bool[root.operand[0].defining_op] IsNotNullQuestion -> TrueAnswer()
// CHECK-NEXT: │ │ │ └── success:
// CHECK-NEXT: │ │ │ └── Bool[root.operand[0].defining_op] OperationNameQuestion -> StringAnswer(value='arith.muli')
// CHECK-NEXT: │ │ │ └── success:
// CHECK-NEXT: │ │ │ └── Bool[root.operand[0].defining_op] OperandCountQuestion -> UnsignedAnswer(value=2)
// CHECK-NEXT: │ │ │ └── success:
// CHECK-NEXT: │ │ │ └── Bool[root.operand[0].defining_op] ResultCountQuestion -> UnsignedAnswer(value=1)
// CHECK-NEXT: │ │ │ └── success:
// CHECK-NEXT: │ │ │ └── Bool[root.operand[0].defining_op.operand[0]] IsNotNullQuestion -> TrueAnswer()
// CHECK-NEXT: │ │ │ └── success:
// CHECK-NEXT: │ │ │ └── Bool[root.operand[0].defining_op.operand[1]] IsNotNullQuestion -> TrueAnswer()
// CHECK-NEXT: │ │ │ └── success:
// CHECK-NEXT: │ │ │ └── Bool[root.operand[0].defining_op.result[0]] IsNotNullQuestion -> TrueAnswer()
// CHECK-NEXT: │ │ │ └── success:
// CHECK-NEXT: │ │ │ └── Bool[root.operand[0].defining_op.result[0]] EqualToQuestion -> TrueAnswer()
// CHECK-NEXT: │ │ │ └── success:
// CHECK-NEXT: │ │ │ └── Bool[root.operand[0].defining_op.result[0].type] EqualToQuestion -> TrueAnswer()
// CHECK-NEXT: │ │ │ └── success:
// CHECK-NEXT: │ │ │ └── SUCCESS(anonymous)
// CHECK-NEXT: │ │ └── failure:
// CHECK-NEXT: │ │ └── Bool[root.operand[0]] EqualToQuestion -> TrueAnswer()
// CHECK-NEXT: │ │ └── success:
// CHECK-NEXT: │ │ └── SUCCESS(anonymous)
// CHECK-NEXT: │ └── case StringAnswer(value='arith.muli'):
// CHECK-NEXT: │ └── Bool[root] OperandCountQuestion -> UnsignedAnswer(value=2)
// CHECK-NEXT: │ └── success:
// CHECK-NEXT: │ └── Bool[root] ResultCountQuestion -> UnsignedAnswer(value=1)
// CHECK-NEXT: │ └── success:
// CHECK-NEXT: │ └── Bool[root.operand[0]] IsNotNullQuestion -> TrueAnswer()
// CHECK-NEXT: │ └── success:
// CHECK-NEXT: │ └── Bool[root.operand[1]] IsNotNullQuestion -> TrueAnswer()
// CHECK-NEXT: │ └── success:
// CHECK-NEXT: │ └── Bool[root.operand[1].defining_op] IsNotNullQuestion -> TrueAnswer()
// CHECK-NEXT: │ └── success:
// CHECK-NEXT: │ └── Bool[root.operand[1].defining_op] OperationNameQuestion -> StringAnswer(value='arith.constant')
// CHECK-NEXT: │ └── success:
// CHECK-NEXT: │ └── Bool[root.operand[1].defining_op] OperandCountQuestion -> UnsignedAnswer(value=0)
// CHECK-NEXT: │ └── success:
// CHECK-NEXT: │ └── Bool[root.operand[1].defining_op] ResultCountQuestion -> UnsignedAnswer(value=1)
// CHECK-NEXT: │ └── success:
// CHECK-NEXT: │ └── Bool[root.operand[1].defining_op.attribute[value]] IsNotNullQuestion -> TrueAnswer()
// CHECK-NEXT: │ └── success:
// CHECK-NEXT: │ └── Bool[root.operand[1].defining_op.attribute[value]] AttributeConstraintQuestion -> AttributeAnswer(value=IntegerAttr(value=IntAttr(data=1), type=IntegerType(32)))
// CHECK-NEXT: │ └── success:
// CHECK-NEXT: │ └── Bool[root.operand[1].defining_op.result[0]] IsNotNullQuestion -> TrueAnswer()
// CHECK-NEXT: │ └── success:
// CHECK-NEXT: │ └── Bool[root.operand[1].defining_op.result[0]] EqualToQuestion -> TrueAnswer()
// CHECK-NEXT: │ └── success:
// CHECK-NEXT: │ └── Bool[root.operand[1].defining_op.result[0].type] EqualToQuestion -> TrueAnswer()
// CHECK-NEXT: │ └── success:
// CHECK-NEXT: │ └── SUCCESS(anonymous)
// CHECK-NEXT: └── failure:
// CHECK-NEXT: └── EXIT

// (x * y) / z -> x * (y/z)
pdl.pattern : benefit(1) {
%x = pdl.operand
%y = pdl.operand
%z = pdl.operand
%type = pdl.type
%mulop = pdl.operation "arith.muli" (%x, %y : !pdl.value, !pdl.value) -> (%type : !pdl.type)
%mul = pdl.result 0 of %mulop
%resultop = pdl.operation "arith.divui" (%mul, %z : !pdl.value, !pdl.value) -> (%type : !pdl.type)
%result = pdl.result 0 of %resultop
pdl.rewrite %resultop {
%newdivop = pdl.operation "arith.divui" (%y, %z : !pdl.value, !pdl.value) -> (%type : !pdl.type)
%newdiv = pdl.result 0 of %newdivop
%newresultop = pdl.operation "arith.muli" (%x, %newdiv : !pdl.value, !pdl.value) -> (%type : !pdl.type)
%newresult = pdl.result 0 of %newresultop
pdl.replace %resultop with %newresultop
}
}

// x / x -> 1
pdl.pattern : benefit(1) {
%x = pdl.operand
%type = pdl.type
%resultop = pdl.operation "arith.divui" (%x, %x : !pdl.value, !pdl.value) -> (%type : !pdl.type)
pdl.rewrite %resultop {
%2 = pdl.attribute = 1 : i32
%3 = pdl.operation "arith.constant" {"value" = %2} -> (%type : !pdl.type)
pdl.replace %resultop with %3
}
}

// x * 1 -> x
pdl.pattern : benefit(1) {
%x = pdl.operand
%type = pdl.type
%one = pdl.attribute = 1 : i32
%constop = pdl.operation "arith.constant" {"value" = %one} -> (%type : !pdl.type)
%const = pdl.result 0 of %constop
%mulop = pdl.operation "arith.muli" (%x, %const : !pdl.value, !pdl.value) -> (%type : !pdl.type)
pdl.rewrite %mulop {
pdl.replace %mulop with (%x : !pdl.value)
}
}
89 changes: 89 additions & 0 deletions tests/filecheck/transforms/convert-pdl-to-pdl-interp/pdl-to-pdl-interp-debug-tree.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,89 @@
// RUN: xdsl-opt -p convert-pdl-to-pdl-interp{print-debug-info=true} %s | filecheck %s

// CHECK: Bool[root.result[0]] IsNotNullQuestion -> TrueAnswer()
// CHECK-NEXT: ├── success:
// CHECK-NEXT: │ └── Bool[root] OperationNameQuestion -> StringAnswer(value='arith.addf')
// CHECK-NEXT: │ └── success:
// CHECK-NEXT: │ └── Bool[root] OperandCountQuestion -> UnsignedAnswer(value=2)
// CHECK-NEXT: │ └── success:
// CHECK-NEXT: │ └── Bool[root] ResultCountQuestion -> UnsignedAnswer(value=1)
// CHECK-NEXT: │ └── success:
// CHECK-NEXT: │ └── Bool[root.operand[0]] IsNotNullQuestion -> TrueAnswer()
// CHECK-NEXT: │ └── success:
// CHECK-NEXT: │ └── Bool[root.operand[1]] IsNotNullQuestion -> TrueAnswer()
// CHECK-NEXT: │ └── success:
// CHECK-NEXT: │ └── Bool[root.operand[0].defining_op] IsNotNullQuestion -> TrueAnswer()
// CHECK-NEXT: │ ├── success:
// CHECK-NEXT: │ │ └── Bool[root.operand[0].defining_op] OperationNameQuestion -> StringAnswer(value='arith.absf')
// CHECK-NEXT: │ │ └── success:
// CHECK-NEXT: │ │ └── Bool[root.operand[0].defining_op] OperandCountQuestion -> UnsignedAnswer(value=1)
// CHECK-NEXT: │ │ └── success:
// CHECK-NEXT: │ │ └── Bool[root.operand[0].defining_op] ResultCountQuestion -> UnsignedAnswer(value=1)
// CHECK-NEXT: │ │ └── success:
// CHECK-NEXT: │ │ └── Bool[root.operand[0].defining_op.operand[0]] IsNotNullQuestion -> TrueAnswer()
// CHECK-NEXT: │ │ └── success:
// CHECK-NEXT: │ │ └── Bool[root.operand[0].defining_op.result[0]] IsNotNullQuestion -> TrueAnswer()
// CHECK-NEXT: │ │ └── success:
// CHECK-NEXT: │ │ └── Bool[root.operand[0].defining_op.result[0]] EqualToQuestion -> TrueAnswer()
// CHECK-NEXT: │ │ └── success:
// CHECK-NEXT: │ │ └── Bool[root.operand[0].defining_op.operand[0].type] EqualToQuestion -> TrueAnswer()
// CHECK-NEXT: │ │ └── success:
// CHECK-NEXT: │ │ └── Bool[root.operand[0].defining_op.operand[0].type] EqualToQuestion -> TrueAnswer()
// CHECK-NEXT: │ │ └── success:
// CHECK-NEXT: │ │ └── Bool[root.operand[0].defining_op.operand[0].type] EqualToQuestion -> TrueAnswer()
// CHECK-NEXT: │ │ └── success:
// CHECK-NEXT: │ │ └── Bool[root.operand[0].defining_op.operand[0].type] TypeConstraintQuestion -> TypeAnswer(value=Float32Type())
// CHECK-NEXT: │ │ └── success:
// CHECK-NEXT: │ │ └── SUCCESS(add_absf_left)
// CHECK-NEXT: │ └── failure:
// CHECK-NEXT: │ └── Bool[root.operand[1].defining_op] IsNotNullQuestion -> TrueAnswer()
// CHECK-NEXT: │ └── success:
// CHECK-NEXT: │ └── Bool[root.operand[0].type] EqualToQuestion -> TrueAnswer()
// CHECK-NEXT: │ └── success:
// CHECK-NEXT: │ └── Bool[root.operand[0].type] TypeConstraintQuestion -> TypeAnswer(value=Float32Type())
// CHECK-NEXT: │ └── success:
// CHECK-NEXT: │ └── Bool[root.operand[1].defining_op] OperationNameQuestion -> StringAnswer(value='arith.absf')
// CHECK-NEXT: │ └── success:
// CHECK-NEXT: │ └── Bool[root.operand[1].defining_op] OperandCountQuestion -> UnsignedAnswer(value=1)
// CHECK-NEXT: │ └── success:
// CHECK-NEXT: │ └── Bool[root.operand[1].defining_op] ResultCountQuestion -> UnsignedAnswer(value=1)
// CHECK-NEXT: │ └── success:
// CHECK-NEXT: │ └── Bool[root.operand[1].defining_op.operand[0]] IsNotNullQuestion -> TrueAnswer()
// CHECK-NEXT: │ └── success:
// CHECK-NEXT: │ └── Bool[root.operand[1].defining_op.result[0]] IsNotNullQuestion -> TrueAnswer()
// CHECK-NEXT: │ └── success:
// CHECK-NEXT: │ └── Bool[root.operand[1].defining_op.result[0]] EqualToQuestion -> TrueAnswer()
// CHECK-NEXT: │ └── success:
// CHECK-NEXT: │ └── Bool[root.operand[1].defining_op.operand[0].type] EqualToQuestion -> TrueAnswer()
// CHECK-NEXT: │ └── success:
// CHECK-NEXT: │ └── Bool[root.operand[1].defining_op.result[0].type] EqualToQuestion -> TrueAnswer()
// CHECK-NEXT: │ └── success:
// CHECK-NEXT: │ └── SUCCESS(add_absf_right)
// CHECK-NEXT: └── failure:
// CHECK-NEXT: └── EXIT

pdl.pattern @add_absf_left : benefit(1) {
%0 = pdl.type : f32
%x = pdl.operand : %0
%y = pdl.operand : %0
%1 = pdl.operation "arith.absf" (%x : !pdl.value) -> (%0 : !pdl.type)
%2 = pdl.result 0 of %1
%3 = pdl.operation "arith.addf" (%2, %y : !pdl.value, !pdl.value) -> (%0 : !pdl.type)
%4 = pdl.result 0 of %3
pdl.rewrite %3 {
pdl.attribute = "hello"
}
}

pdl.pattern @add_absf_right : benefit(1) {
%0 = pdl.type : f32
%x = pdl.operand : %0
%y = pdl.operand : %0
%1 = pdl.operation "arith.absf" (%y : !pdl.value) -> (%0 : !pdl.type)
%2 = pdl.result 0 of %1
%3 = pdl.operation "arith.addf" (%x, %2 : !pdl.value, !pdl.value) -> (%0 : !pdl.type)
%4 = pdl.result 0 of %3
pdl.rewrite %3 {
pdl.attribute = "hello"
}
}
Loading
Loading