Make loop passthrough work; performance issues though

dag_in_context/src/loop_iteration_analysis.egg

@@ -1,5 +1,5 @@
 ;; Analysis to get the number of iterations of a loop
-(ruleset loop-iter-analysis)
+(ruleset loop-iters-analysis)
 
 ;;                      inputs, outputs -> number of iterations
 ;; The minimum possible guess is 1 because of do-while loops
@@ -11,15 +11,15 @@
 ;; by default, guess that all loops run 1000 times
 (rule ((DoWhile inputs outputs))
       ((set (LoopNumItersGuess inputs outputs) 1000))
-      :ruleset loop-iter-analysis)
+      :ruleset loop-iters-analysis)
 
 ;; For a loop that is false, its num iters is 1
 (rule
   ((= loop (DoWhile inputs outputs))
    (= (Const (Bool false) ty ctx) (Get outputs 0)))
   ((set (LoopNumItersGuess inputs outputs) 1)
    (TerminatingLoop inputs outputs))
-:ruleset loop-iter-analysis)
+:ruleset loop-iters-analysis)
 
 ;; Figure out number of iterations for a loop with constant bounds and initial value
 ;; and i is updated before checking pred
@@ -33,6 +33,7 @@
    ;; updated counter at counter_i
    (= next_counter (Get outputs (+ counter_i 1)))
    ;; increments by some constant each loop
+   ;; TODO: how to handle the invariant case?
    (= next_counter (Bop (Add) (Get (Arg _ty _ctx) counter_i)
                               (Const (Int increment) _ty2 _ctx2)))
    (> increment 0)
@@ -46,7 +47,7 @@
     (set (LoopNumItersGuess inputs outputs) (/ (- end_constant start_const) increment))
     (TerminatingLoop inputs outputs)
   )
-  :ruleset loop-iter-analysis)
+  :ruleset loop-iters-analysis)
 
 ;; Figure out number of iterations for a loop with constant bounds and initial value
 ;; and i is updated after checking pred
@@ -72,5 +73,5 @@
     (set (LoopNumItersGuess inputs outputs) (+ (/ (- end_constant start_const) increment) 1))
     (TerminatingLoop inputs outputs)
   )
-  :ruleset loop-iter-analysis)
+  :ruleset loop-iters-analysis)
 

dag_in_context/src/optimizations/loop_unroll.egg

@@ -2,7 +2,6 @@
 ;; Depends on loop iteration analysis
 (ruleset loop-unroll)
 (ruleset loop-peel)
-(ruleset loop-iters-analysis)
 
 ;; loop peeling rule
 ;; Only peel loops that we know iterate < 3 times

dag_in_context/src/optimizations/passthrough.egg

@@ -14,14 +14,31 @@
        ((union lhs (Get inputs i)))
        :ruleset passthrough)
 
-;; Pass through thetas: state edge case
-(rule ((= lhs (Get loop i))
-        (= loop (DoWhile inputs pred-outputs))
+; ;; Pass through thetas: state edge case
+(rule ((= loop (DoWhile inputs pred-outputs))
         (= (Get pred-outputs (+ i 1)) (Get (Arg _ty _ctx) i))
         ;; It is OK to pass through state edges as long as the loop terminates
-        (TerminatingLoop inputs pred-outputs)
-        )
-       ((union lhs (Get inputs i)))
+        (TerminatingLoop inputs pred-outputs))
+      (
+       ;; To maintain the linearity invariant, we must remove the state edge
+       ;; from the loop.
+       (let new-inputs (TupleRemoveAt inputs i))
+       (let removed-outputs (TupleRemoveAt pred-outputs (+ i 1)))
+       (let new-outputs (DropAt (TmpCtx) i removed-outputs))
+
+       (let projected-old-loop (TupleRemoveAt loop i))
+       (let new-loop (DoWhile new-inputs new-outputs))
+       (union new-loop projected-old-loop)
+
+       ;; Resolve the temporary context
+       (union (TmpCtx) (InLoop new-inputs new-outputs))
+       (delete (TmpCtx))
+
+       ;; State edge can be gotten without the loop now
+       (union (Get loop i) (Get inputs i))
+
+       ;; Subsume the loop later
+       (ToSubsumeLoop inputs pred-outputs))
        :ruleset passthrough)
 
 ;; Pass through switch arguments

dag_in_context/src/utility/util.egg

@@ -75,4 +75,8 @@
       ((subsume (If a b c d)))
       :ruleset subsume-after-helpers)
 
+(relation ToSubsumeLoop (Expr Expr))
+(rule ((ToSubsumeLoop in p-out))
+      ((subsume (DoWhile in p-out)))
+      :ruleset subsume-after-helpers)
 

tests/passing/small/dead_loop_deletion.bril

@@ -0,0 +1,15 @@
+@main: int {
+    i: int = const 1;
+    forty: int = const 40;
+    one: int = const 1;
+
+.loop_body:
+    i: int = add i one;
+    cond: bool = lt i forty;
+    br cond .loop_body .loop_end;
+
+.loop_end:
+    j: int = const 2;
+
+    ret j;
+}

tests/passing/small/loop_state_pass_through.bril

@@ -0,0 +1,14 @@
+# ARGS: 5
+@main(input: int) {
+    one: int = const 1;
+    i: int = const 1;
+    jmp .loop;
+.loop:
+    max: int = const 10;
+    cond: bool = lt i max;
+    i: int = add i one;
+    br cond .loop .exit;
+.exit:
+    res: int = add i input;
+    print res;
+}