From 0b328b856350a8ad08e5f353d7fe05b68913ddb8 Mon Sep 17 00:00:00 2001
From: oflatt <oflatt@gmail.com>
Date: Mon, 11 Nov 2024 16:41:40 -0800
Subject: [PATCH] add inlining after optimization

---
 dag_in_context/src/lib.rs                     |  7 ++++--
 .../src/optimizations/rec_to_loop.egg         |  3 +--
 dag_in_context/src/schedule.rs                | 22 ++++++++++++++-----
 3 files changed, 22 insertions(+), 10 deletions(-)

diff --git a/dag_in_context/src/lib.rs b/dag_in_context/src/lib.rs
index ae954586..38c68e7d 100644
--- a/dag_in_context/src/lib.rs
+++ b/dag_in_context/src/lib.rs
@@ -293,8 +293,11 @@ pub fn optimize(
 
         // if we are inlining, save the program
         // TODO we inline on the first pass, but this should be configurable from the schedule
-        // TODO inline
-        let inline_program = if i == 1000 { Some(res.clone()) } else { None };
+        let inline_program = if schedule.contains("INLINE") {
+            Some(res.clone())
+        } else {
+            None
+        };
 
         for func in fns {
             log::info!("Running pass {} on {}", i, func);
diff --git a/dag_in_context/src/optimizations/rec_to_loop.egg b/dag_in_context/src/optimizations/rec_to_loop.egg
index 6f9fa2bc..2a774d38 100644
--- a/dag_in_context/src/optimizations/rec_to_loop.egg
+++ b/dag_in_context/src/optimizations/rec_to_loop.egg
@@ -86,8 +86,7 @@
    (= always-runs-len (tuple-length always-runs))
    (= start-ty (TupleT start-ty-list))
    (HasType body func-ty))
-  ((panic "good")
-    (let loop-ty
+  ((let loop-ty
      (TupleT (TLConcat start-ty-list (TCons (IntT) (TNil)))))
    ;; recursive case in the loop
    (let new-rec-case
diff --git a/dag_in_context/src/schedule.rs b/dag_in_context/src/schedule.rs
index 60df9e9c..190d237a 100644
--- a/dag_in_context/src/schedule.rs
+++ b/dag_in_context/src/schedule.rs
@@ -50,7 +50,6 @@ fn optimizations() -> Vec<String> {
         "loop-inv-motion",
         "loop-strength-reduction",
         "loop-peel",
-        "rec-to-loop",
     ]
     .iter()
     .map(|opt| opt.to_string())
@@ -118,12 +117,22 @@ pub fn mk_sequential_schedule() -> Vec<String> {
 pub fn parallel_schedule() -> Vec<String> {
     let helpers = helpers();
 
-    vec![format!(
-        "
+    vec![
+        format!(
+            "
 (run-schedule
+    (repeat 2
+      {helpers}
+      swap-if)
     {helpers}
-    (repeat 2 swap-if)
-
+    rec-to-loop
+    {helpers})"
+        ),
+        format!(
+            "
+;; HACK: when INLINE appears in this string
+;; we perform inlining in this pass
+(run-schedule
     (repeat 2
         {helpers}
         all-optimizations
@@ -137,5 +146,6 @@ pub fn parallel_schedule() -> Vec<String> {
     {helpers}
 )
 "
-    )]
+        ),
+    ]
 }