# 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)))

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). The leq rules are purely analysis rules, because they don't add any 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, these axiomatic rules are doing optimizations.

(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. 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 the rules that depend on analyses but actually do optimization.

(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)))

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 (cancelling division) 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). Thus, we can build a more efficient schedule for this specific program.

(push)

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

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

Then, just run one iteration of the optimizations ruleset.

(run optimizations 1)

This makes our check pass

(check (= expr addition-chain))

Print the size of the e-graph and the best expression reached for later.

(print-size)
(extract expr)

While the above program is effective at optimizing that specific program, it would fail if we had a program where we had to interleave the optimizations and analyses. There is a common pattern of schedules which we use to fix this, which use all four of the variants in our scheduling language: repeat, seq, saturate, and run.

(pop)
(push)

The idea behind this schedule is to always saturate analyses before running optimizations. We then wrap that schedule in a repeat block to give us control over how long to run egglog. Using repeat 1 gives us the same schedule as before, but now we can increase the iteration count if we want better results and have the time/space budget for it.

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

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

(pop)
(pop)
(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")))))))
; 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)

Then, only running with repeat 1...

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

does not give us the optimal program. Instead, we need to increase the iteration number.

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

Because it takes one iteration of optimizations to simplify x-times-zero to zero, and this is required to determine that nonzero-expr is actually nonzero.

However, sometimes just having an iteration number does not give you enough control. As we can see from the (print-size)s above, adding one iteration can increase the size of the e-graph hyper-exponentially.

To this end, egglog has an 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 add 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. This scheduler works well when the ruleset contains explosive rules like associativity.

(pop)
(push)

Let's go back to this example, and increase the number of iterations slightly.

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

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! This is because the associativity is exponential and can easily blow up the e-graph in one iteration.

This is where the back-off scheduler comes in handy. It will prevent the associativity rule from dominating the equality saturation, and will ban associativity for some iterations so that other rules can catch up.

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