Using loop-inversion to optimizes translation overhead for while loops

dag_in_context/src/lib.rs

@@ -63,6 +63,7 @@ pub fn prologue() -> String {
         include_str!("optimizations/loop_unroll.egg"),
         include_str!("optimizations/passthrough.egg"),
         include_str!("optimizations/loop_strength_reduction.egg"),
+        include_str!("optimizations/ivt.egg"),
         include_str!("utility/debug-helper.egg"),
     ]
     .join("\n")

dag_in_context/src/optimizations/ivt.egg

@@ -0,0 +1,133 @@
+
+;; A Perm is a reverse list of integers
+(datatype Perm (PermPush i64 Perm) (PNil))
+;; expr1 is a list of expressions of the form (Get expr2 i),
+;; where all the i's form a permutation
+(relation IVTPermutationAnalysisDemand (Expr))
+;;                                    expr1 curr  expr2
+(relation IVTPermutationAnalysisImpl (Expr  Expr  Expr Perm))
+;;                                expr1 expr2
+(relation IVTPermutationAnalysis (Expr Expr Perm))
+
+(rule (
+    (DoWhile inpW outW)
+) (
+    (IVTPermutationAnalysisDemand outW)
+) :ruleset always-run)
+
+(rule (
+    (IVTPermutationAnalysisDemand loop-body)
+    (= loop-body (Concat (Single (Get expr ith)) rest))
+) (
+    (let perm (PermPush ith (PNil)))
+    (IVTPermutationAnalysisImpl loop-body rest expr perm)
+) :ruleset always-run)
+
+(rule (
+    (IVTPermutationAnalysisImpl loop-body curr expr perm)
+    (= curr (Concat (Single (Get expr ith)) rest))
+) (
+    (let new-perm (PermPush ith perm))
+    (IVTPermutationAnalysisImpl loop-body rest expr new-perm)
+) :ruleset always-run)
+
+(rule (
+    (IVTPermutationAnalysisImpl loop-body (Single last) expr perm)
+    (= last (Get expr ith))
+) (
+    (let new-perm (PermPush ith perm))
+    (IVTPermutationAnalysis loop-body expr new-perm)
+) :ruleset always-run)
+
+(function ApplyPerm (Perm Expr) Expr)
+
+(rewrite (ApplyPerm (PermPush ith rest) expr) 
+         (Concat (ApplyPerm rest expr) (Single (Get expr ith)))
+    :when ((= rest (PermPush _jth _rest)))
+    :ruleset always-run)
+(rewrite (ApplyPerm (PermPush ith (PNil)) expr) (Single (Get expr ith))
+    :ruleset always-run)
+
+(ruleset loop-inversion)
+
+;; This is for unified handling of thn/els branches
+(relation ith-arg-is-bool (Expr Expr Expr i64 Expr Expr))
+
+(rule (
+    (= loop (DoWhile inpW outW))
+    (IVTPermutationAnalysis outW conditional perm)
+    (= conditional (If condI inpI thn els))
+
+    (= (Get thn ith) (Const (Bool true) _unused1 _unused2))
+    (= (Get els ith) (Const (Bool false) _unused3 _unused4))
+) (
+    (ith-arg-is-bool conditional condI inpI ith thn els)
+) :ruleset always-run)
+
+(rule (
+    (= loop (DoWhile inpW outW))
+    (IVTPermutationAnalysis outW conditional perm)
+    (= conditional (If condI inpI thn els))
+
+    (= (Get thn ith) (Const (Bool false) _unused1 _unused2))
+    (= (Get els ith) (Const (Bool true) _unused3 _unused4))
+) (
+    ;; TODO: this may introduce overhead, but is the only way to
+    ;; not have two rules
+    (ith-arg-is-bool conditional (Uop (Not) condI) inpI ith els thn)
+) :ruleset always-run)
+
+(rule (
+    (= loop (DoWhile inpW outW))
+    (IVTPermutationAnalysis outW conditional perm)
+    ;; This generalizes the following conditions:
+    ;;   (= conditional (If condI inpI thn els))
+    ;;   (= (Get thn ith) (Const (Bool true) _unused1 _unused2))
+    ;;   (= (Get els ith) (Const (Bool false) _unused3 _unused4))
+    (ith-arg-is-bool conditional condI inpI ith thn els)
+
+    (ContextOf inpW outer-ctx)
+    (ContextOf inpI if-ctx)
+    (HasType inpW argW)
+) (
+    ;; first peeled condition
+    (let new-if-cond (Subst outer-ctx inpW condI))
+    ;; if contexts
+    (let new-if-true-ctx (InIf true new-if-cond inpW))
+    (let new-if-false-ctx (InIf false new-if-cond inpW))
+
+    (let new-loop-context (TmpCtx))
+
+    ;; body
+    (let new-loop-outputs_ 
+        (TupleRemoveAt
+            (ApplyPerm perm (Subst new-loop-context inpI thn))
+            0))
+    ;; the first element of body was true but should be the condition.
+    (let new-loop-outputs
+        (Concat
+            ; (Single (Subst new-loop-context (Arg argW new-loop-context) condI))
+            (Single (Subst new-loop-context new-loop-outputs_ condI))
+            new-loop-outputs_))
+
+    (let new-loop (DoWhile (Arg argW new-if-true-ctx) new-loop-outputs))
+    (let new-if
+        (If new-if-cond inpW
+            new-loop
+            (Arg argW new-if-false-ctx)))
+
+    ;; Apply the body of the false branch as an afterprocessing wrapper
+    (let new-expr_
+        (Subst outer-ctx new-if 
+            (Subst if-ctx inpI els)))
+    (let new-expr 
+        (TupleRemoveAt 
+            (ApplyPerm perm new-expr_) 
+            0))
+
+    (union new-expr loop)
+    (union new-loop-context (InLoop (Arg argW new-if-true-ctx) new-loop-outputs))
+
+    (delete (DoWhile inpW outW))
+    (delete (TmpCtx))
+) :ruleset loop-inversion)

dag_in_context/src/schedule.rs

@@ -62,6 +62,7 @@ pub fn mk_schedule() -> String {
     switch_rewrite
     loop-inv-motion
     loop-strength-reduction
+    loop-inversion
     loop-peel
 )