Efficient E-Graph Extraction via Linear Programming

Haichen Dong

Deyuan (Mike) He

Background

2

Term rewriting is extremely common in Compilers:

a * 2

a << 1

a / pot

a >> log2(pot)

a / const

a * (1 / const)

Background

3

However, determining the order of applying rewrite rules is HARD!

M₁

[N, M]

M₂

[M, K]

M₃

[K, P]

M₄

[K, Q]

Background

4

However, determining the order of applying rewrite rules is HARD!

M₁

[N, M]

M₂

[M, K]

M₃

[K, P]

M₄

[K, Q]

Background

5

M₁

M₂

M₃

*

*

M₁

M₂

M₄

*

*

1. Common Subexpression Elimination (CSE)
2. Associativity of Matrix Multiplication

2NMK + NKP + NKQ multiplications

Background

6

M₁

M₂

M₃

*

*

M₁

M₂

M₄

*

*

M₃

*

M₄

*

M₁

M₂

*

M₁₂

=

M₁₂

M₁₂

2NMK + NKP + NKQ multiplications

NMK + NKP + NKQ multiplications

CSE Rewrite

Background

7

M₁

M₂

M₃

*

*

M₁

M₂

M₄

*

*

2NMK + NKP + NKQ multiplications

NKP + NKQ + MKP + MKQ multiplications

Associativity Rewrite

M₁

M₂

M₃

*

*

M₁

M₂

M₄

*

*

Background

8

NKP + NKQ + MKP + MKQ

Associativity Rewrite

M₁

M₂

M₃

*

*

M₁

M₂

M₄

*

*

M₃

*

M₄

*

M₁

M₂

*

M₁₂

=

M₁₂

M₁₂

NMK + NKP + NKQ

CSE Rewrite

Background

9

NKP + NKQ + MKP + MKQ

Associativity Rewrite

M₁

M₂

M₃

*

*

M₁

M₂

M₄

*

*

M₃

*

M₄

*

M₁

M₂

*

M₁₂

=

M₁₂

M₁₂

NMK + NKP + NKQ

CSE Rewrite

Background

10

Case 1 (Assoc better than CSE):

NMK + NKP + NKQ > NKP + NKQ + MKP + MKQ

⇒ NMK > (P+Q)MK

⇒ N > P+Q

​

Case 2 (CSE better than Assoc):

NMK + NKP + NKQ < NKP + NKQ + MKP + MKQ

⇒ NMK < (P+Q)MK

⇒ N < P+Q

​

Background

3

Compilers may have hundreds of passes.

​

How to determine the order to ensure the product program is Optimal ?

Background

4

Compilers may have hundreds of passes.

​

How to determine the order to ensure the product program is Optimal ?

​

Phase Ordering Problem

E-Graph

5

A compact, tree data structure that maintains equivalences over terms with respect to a set of axioms (rewrite rules)

1. E-Classes (dashed boxes): A set of equivalent terms
2. E-Nodes (solid boxes): Operators, variables or literals

x

a

2

+

1

Represents the term “2”

Represents the term “1+1”

E-Graph

6

Syntactic Rewrites: an initial pattern and a target pattern

x

a

2

+

1

Apply

with ?x ⇒ a

E-Graph

7

Syntactic Rewrites: initiate and add new expressions based on the matched bindings

x

a

2

+

1

x

a

2

+

1

<<

Apply

with ?x ⇒ a

instantiate

?x << 1 to a << 1

E-Graph

8

Syntactic Rewrites: initiate and add new expressions based on the matched bindings

x

a

2

+

1

x

a

2

+

1

<<

Apply

with ?x ⇒ a

instantiate

?x << 1 to a << 1

Non-destructive Rewrite

E-Graph

9

x

a

2

+

1

<<

10

1

0

5

0

1

Numbers in E-Nodes are example costs

Extraction: Given a root E-Class, pick the “best” term

(minimizing the sum of costs of E-Nodes given by a cost model)

E-Graph

10

Extraction: Given a root E-Class, pick the “best” term

(minimizing the sum of costs of E-Nodes given by a cost model)

x

a

2

+

1

<<

10

1

0

5

0

1

Numbers in E-Nodes are example costs

E-Graph

11

x

a

2

+

1

<<

10

1

0

5

0

1

Numbers in E-Nodes are example costs

Extraction: Given a root E-Class, pick the “best” term

(minimizing the sum of costs of E-Nodes given by a cost model)

The Problem

12

Greedy does not always work!

Yang, Yichen et al. "Equality Saturation for Tensor Graph Superoptimization." (2021).

The Problem

13

Greedy does not always work!

Yang, Yichen et al. "Equality Saturation for Tensor Graph Superoptimization." (2021).

Suppose

b + d > c + e

a + b + d < a + c + e + d

The Problem

14

Greedy does not always work!

Suppose

b + d > c + e

a + b + d < a + c + e + d

Yang, Yichen et al. "Equality Saturation for Tensor Graph Superoptimization." (2021).

Extraction as ILP

15

For each E-Class k, create a binary variable w_k

w_R

w_A

w_B

w_C

Extraction as ILP

16

For each E-Class k, create a binary variable w_k

w_R

w_A

w_B

w_C

For each E-Node n, create a binary variable w_n

w_a

w_b

w_c

w_e

w_d

Extraction as ILP

17

w_R

w_A

w_B

w_C

w_a

w_b

w_c

w_e

w_d

Constraints:

1. We must pick at least one E-Node from the root E-Class (in this case, R)

​

• If an E-Class K is picked, then at least one E-Node in K must also be picked

​

• If an E-Node n is picked, then all of its children E-Classes must be picked

For any E-Class/E-Node, w_x = 1 denotes the E-Class/E-Node is active (i.e. they are picked in the extraction)

Extraction as ILP

18

w_R

w_A

w_B

w_C

w_a

w_b

w_c

w_e

w_d

Constraints:

• Root constraints
1. w_a = 1

​

• Node constraints
• -w_R + w_a > 0
• -w_A + w_b + w_c > 0
• -w_B + w_d > 0
• -w_C + w_e > 0

​

• Children constraints
• w_B ≥ w_a w_A ≥ w_a
• w_B ≥ w_b
• w_C ≥ w_c

For any E-Class/E-Node, w_x = 1 denotes the E-Class/E-Node is active (i.e. they are picked in the extraction)

Extraction as LP

19

LP Relaxation

20

Analysis

21

Counter-example:

All operations have uniform costs

​

Observation: OPT_ILP = k

Furthermore, the cost of any valid solution to this extraction problem is at least k!

​

Possible solution for LP:

Halve weights in every layer.

The total cost is only 3!

Evaluation

22

Setup:

• Implemented LP encoding in egg library (~200 LoC), invoking CPLEX solver
• Adapted Tensat (Yang et al.) evaluation pipeline to test our rounding scheme
• Evaluated graph rewrites on
• BERT, VGG-16, ResNexT-50, NasRNN, InceptionV3

Optimization Comparison

23

Costs of extracted program (normalized by the cost of inputs)

Optimization Comparison

24

Cases LP extraction doesn’t work well

(due to the drawback we discussed)

Solver Time Comparison

25

Solver time comparison on solving LP/ILP extractions

Solver Time Comparison

26

Solver time comparison on solving LP/ILP extractions

Solver Time Comparison

27

For cases LP extraction works, it has up to 200x~ speed-up

Discussion

28

1. If it is not guaranteed to work, why do we do this?
1. It is natural to try LP relaxation for ILP formulations
2. The proof of “It may not work” is also a result
3. It may worth to try LP first before ILP if the approximation works well since most of the cases, E-Graph does not saturates
2. Can it be improved? If so, how?
• LP itself probably does not have a good space for improvement.
• However, this finding led us to explore other strategies: Non-convex optimization and Weighted partial MaxSAT. (See report for more details)
3. What’s the future plan?
• Try implementing other extraction algorithms and compare with current ones.
• We are considering a workshop paper next year.