# Scheduling

In this lesson, we will learn how to use run-schedule to improve the performance of egglog. We start by using the same language as the previous lesson.

(datatype Expr
    (Num BigRat)
    (Var String)
    (Add Expr Expr)
    (Mul Expr Expr)
    (Div Expr Expr))

(let zero (bigrat (bigint 0) (bigint 1)))
(let one (bigrat (bigint 1) (bigint 1)))
(let two (bigrat (bigint 2) (bigint 1)))

Rulesets

Different from lesson 3, we organize our rules into "rulesets" A ruleset is exactly what it sounds like; a set of rules. We can declare rulesets using the ruleset command.

(ruleset optimizations)
(ruleset analysis)

We can add rules to rulesets using the :ruleset annotation on rules, rewrites, and birewrites. Leaving off :ruleset causes the rule to be added to the default ruleset, which is what we've shown so far. We can run rulesets using (run ruleset iters), or (run iters) for running the default ruleset.

Here, we add leq rules to the analysis ruleset, because they don't add new Expr nodes to the e-graph.

(relation leq (Expr Expr))
(rule ((leq e1 e2) (leq e2 e3)) ((leq e1 e3)) :ruleset analysis)
(rule ((= e1 (Num n1)) (= e2 (Num n2)) (<= n1 n2)) ((leq e1 e2)) :ruleset analysis)
(rule ((= v (Var x))) ((leq v v)) :ruleset analysis)
(rule ((= e1 (Add e1a e1b)) (= e2 (Add e2a e2b)) (leq e1a e2a) (leq e1b e2b))
      ((leq e1 e2))
      :ruleset analysis)

In contrast, the following axiomatic rules are doing optimizations, so we add them to the optimizations ruleset.

(birewrite (Add x (Add y z)) (Add (Add x y) z) :ruleset optimizations)
(birewrite (Mul x (Mul y z)) (Mul (Mul x y) z) :ruleset optimizations)
(rewrite (Add x y) (Add y x) :ruleset optimizations)
(rewrite (Mul x y) (Mul y x) :ruleset optimizations)
(rewrite (Mul x (Add y z)) (Add (Mul x y) (Mul x z)) :ruleset optimizations)
(rewrite (Add x (Num zero)) x :ruleset optimizations)
(rewrite (Mul x (Num one)) x :ruleset optimizations)
(rewrite (Add (Num a) (Num b)) (Num (+ a b)) :ruleset optimizations)
(rewrite (Mul (Num a) (Num b)) (Num (* a b)) :ruleset optimizations)

It is tedious and error-prone to manually annotate each rule with a ruleset, so egglog-experimental provides a convenience syntax with-ruleset:

 (with-ruleset optimizations
   (birewrite (Add x (Add y z)) (Add (Add x y) z))
   ...
 )

From now on, we will use this syntax.

Here we add the rest of the rules from the last section, but tagged with the appropriate rulesets.

(function upper-bound (Expr) BigRat :merge (min old new))
(function lower-bound (Expr) BigRat :merge (max old new))
(with-ruleset analysis
    (rule ((leq e (Num n))) ((set (upper-bound e) n)))
    (rule ((leq (Num n) e)) ((set (lower-bound e) n)))
    (rule ((= e (Add e1 e2)) (= (upper-bound e1) u1) (= (upper-bound e2) u2))
          ((set (upper-bound e) (+ u1 u2))))
    (rule ((= e (Add e1 e2)) (= (lower-bound e1) l1) (= (lower-bound e2) l2))
        ((set (lower-bound e) (+ l1 l2))))

    (rule ((= e (Mul e1 e2))
           (= le1 (lower-bound e1)) (= le2 (lower-bound e2))
           (= ue1 (upper-bound e1)) (= ue2 (upper-bound e2)))
         ((set (lower-bound e)
              (min (* le1 le2) (min (* le1 ue2) (min (* ue1 le2) (* ue1 ue2)))))
          (set (upper-bound e)
              (max (* le1 le2) (max (* le1 ue2) (max (* ue1 le2) (* ue1 ue2))))))
    )
)

(relation non-zero (Expr))
(with-ruleset analysis
    (rule ((< (upper-bound e) zero)) ((non-zero e)))
    (rule ((> (lower-bound e) zero)) ((non-zero e)))
)

Finally, we have optimization rules that depend on the analysis rules we defined above.

(with-ruleset optimizations
    (rewrite (Div x x)         (Num one) :when ((non-zero x)))
    (rewrite (Mul x (Div y x)) y         :when ((non-zero x)))
)

Now consider the following program, which consists of a long sequence of additions inside a cancelling division.

(push)
; a + (b + (c + (d + (e + f))))
(let addition-chain 
    (Add (Var "a") (Add (Var "b") 
        (Add (Var "c") (Add (Var "d") 
            (Add (Var "e") (Var "f")))))))
; 1 + 1 + 1 + 1 + 2
(let nonzero-expr 
    (Add (Num one) (Add (Num one) 
        (Add (Num one) (Add (Num one) (Num two))))))
(let expr (Mul nonzero-expr (Div addition-chain nonzero-expr)))

We want the following check to pass after running the rules.

(fail (check (= expr addition-chain)))

To make this check pass, we have to first discover that nonzero-expr is indeed non-zero, which allows the rule from (Mul x (Div y x)) to y to fire. On the other hand, if we apply the optimization rules, we risk the exponential blowup from the associative and commutative permutations of the addition-chain.

Therefore, if we try to run both rulesets directly, egglog will spend lots of effort reassociating and commuting the terms in the addition-chain, even though the optimization that we actually want to run only takes one iteration. However, that optimization requires knowing a fact that takes multiple iterations to compute (propagating lower- and upper-bounds through nonzero-expr). We can build a more efficient schedule.

(push)

Schedules

Our schedule starts by saturating the analysis rules, fully propagating the non-zero information without adding any e-nodes to the e-graph.

(run-schedule (saturate (run analysis)))

Then, just run one iteration of the optimizations ruleset.

(run optimizations 1)

Or equivalently,

 (run-schedule 
     (saturate (run analysis))
     (run optimizations))

This makes our check pass

(check (= expr addition-chain))
(pop)

While the above program is effective at optimizing that specific program, it would fail if we had a slightly more complex program where we had to interleave the optimizations and analyses to derive the optimal program. For expressing more complex schedules like these, egglog supports a scheduling sub-language, with primitives repeat, seq, saturate, and run.

The idea behind the following schedule is to always saturate analyses before running optimizations. This combination is wrapped in a repeat block to give us control over how long to run egglog. With repeat 1 it is the same schedule as before, but now we can increase the iteration count if we want to optimize harder with more time and space budget.

(run-schedule
    (repeat 2
        (saturate (run analysis))
        (run optimizations)))
(pop)

Running more iterations does not help our above example per se, but if we had started with a slightly more complex program to optimize...

; a + (b + (c + (d + (e + f))))
(let addition-chain 
    (Add (Var "a") (Add (Var "b") 
        (Add (Var "c") (Add (Var "d") 
            (Add (Var "e") (Var "f")))))))
; x * 0
(let x-times-zero (Mul (Var "x") (Num zero)))
(let nonzero-expr 
    (Add (Num one) (Add (Num one) 
        (Add (Num one) (Add (Num one) 
            x-times-zero)))))
(let expr (Mul nonzero-expr (Div addition-chain nonzero-expr)))

For the purpose of this example, we add this rule

(rewrite (Mul x (Num zero)) (Num zero) :ruleset optimizations)

To prove expr is equivalent to addition-chian by applying the cancellation law, we need to prove nonzero-expr is nonzero, which requires proving x-times-zero's bound. To show x-time-zero's bound, we need to apply an optimization rule to rewrite it to 0. In other words, this requires running analyses in between two runs of optimization rules (the cancellation law and Mul's identity law)

Therefore, only running our schedule with one iteration (repeat 1) does not give us the optimal program.

(push)
(run-schedule
    (repeat 1
        (saturate (run analysis))
        (run optimizations)))
(extract expr)
(pop)

Instead, we need to increase the iteration number.

(push)
(run-schedule
    (repeat 2
        (saturate (run analysis))
        (run optimizations)))
(extract expr)
(pop)

Using custom schedulers

However, sometimes just having an iteration number does not give you enough control. For example, for many rules, such as associativity and commutativity (AC), the size of the e-graph grows hyper-exponentially with respect to the number of iterations.

Let's go back to this example, and run for five iterations.

(push)
(run-schedule
    (repeat 5
        (saturate (run analysis))
        (run optimizations)))
(print-size Mul)
(pop)

At iteration 5, the Mul function has size 582. However, if we bump that to 6, the size of the Mul function will increase to 13285! Therefore, the iteration number is too coarse of a granularity for defining the search space.

To this end, egglog provides a scheduler mechanism. A scheduler can decide which matches are important and need to be applied, while others can be delayed or skipped. To use scheduler, there are two operations: let-scheduler and run-with. (let-scheduler sched ..) instantiates a scheduler sched, and (run-with sched ruleset) rules a ruleset with that scheduler.

Currently, egglog-experimental implements one scheduler, back-off. (We will implement our own scheduler in Section 6.) The idea of back-off is that it will ban a rule from applying if that rule grows the e-graph too fast. The decision to ban is based on a threshold, which is initially small and increases as rules are banned. This scheduler works well when the ruleset contains explosive rules like AC.

In this example, the back-off scheduler can prevent the associativity rule from dominating the equality saturation: when the the associativity rule (or any other rule) is fired too much, the scheduler will automatically ban this rule for a few iterations, so that other rules can catch up.

(run-schedule
    (let-scheduler bo (back-off))
    (repeat 10 (run-with bo optimizations)))
(print-size Mul)

It is important that the scheduler bo is instantiated outside the repeat loop, since each scheduler carries some state that is updated when run. For example, the following schedule has a very semantics than the schedule above.

 (run-schedule
     (repeat 10
         (let-scheduler bo (back-off))
         (run-with bo optimizations)))

This schedule instantiates a (fresh) back-off scheduler for each run-with, so the ten iterations of rulesets are all run with the initial configuration of the back-off scheduler, which has a very low threshold for banning rules.