add new test

hercules_opt/src/forkify.rs

@@ -10,8 +10,8 @@ use hercules_ir::*;
 
 use crate::*;
 
-/* 
- * TODO: Forkify currently makes a bunch of small edits - this needs to be 
+/*
+ * TODO: Forkify currently makes a bunch of small edits - this needs to be
  * changed so that every loop that gets forkified corresponds to a single edit
  * + sub-edits. This would allow us to run forkify on a subset of a function.
  */
@@ -349,6 +349,23 @@ pub fn forkify_loop(
             };
             let reduce_id = edit.add_node(reduce);
 
+            if edit.get_schedule(init).contains(&Schedule::ParallelReduce)
+                || edit
+                    .get_schedule(continue_latch)
+                    .contains(&Schedule::ParallelReduce)
+            {
+                edit = edit.add_schedule(reduce_id, Schedule::ParallelReduce)?;
+            }
+            if edit
+                .get_schedule(init)
+                .contains(&Schedule::TightAssociative)
+                || edit
+                    .get_schedule(continue_latch)
+                    .contains(&Schedule::TightAssociative)
+            {
+                edit = edit.add_schedule(reduce_id, Schedule::TightAssociative)?;
+            }
+
             edit = edit.replace_all_uses_where(phi, reduce_id, |usee| *usee != reduce_id)?;
             edit = edit.replace_all_uses_where(continue_latch, reduce_id, |usee| {
                 !loop_nodes.contains(usee) && *usee != reduce_id
@@ -414,9 +431,8 @@ pub fn analyze_phis<'a>(
     phis: &'a [NodeID],
     loop_nodes: &'a HashSet<NodeID>,
 ) -> impl Iterator<Item = LoopPHI> + 'a {
-
-    // Find data cycles within the loop of this phi, 
-    // Start from the phis loop_continue_latch, and walk its uses until we find the original phi. 
+    // Find data cycles within the loop of this phi,
+    // Start from the phis loop_continue_latch, and walk its uses until we find the original phi.
 
     phis.into_iter().map(move |phi| {
         let stop_on: HashSet<NodeID> = editor
@@ -465,7 +481,7 @@ pub fn analyze_phis<'a>(
             .unwrap();
 
         let loop_continue_latch = editor.node(phi).try_phi().unwrap().1[continue_idx];
-        
+
         let uses = walk_all_uses_stop_on(loop_continue_latch, editor, stop_on.clone());
         let users = walk_all_users_stop_on(*phi, editor, stop_on.clone());
 
@@ -513,7 +529,7 @@ pub fn analyze_phis<'a>(
             // If some other node in the cycle is used, there is not a valid node to assign it after making the cycle a reduce.
             if intersection
                 .iter()
-                .filter(|node| **node != loop_continue_latch )
+                .filter(|node| **node != loop_continue_latch)
                 .any(|data_node| {
                     editor
                         .get_users(*data_node)

juno_samples/fork_join_tests/src/fork_join_tests.jn

@@ -59,3 +59,16 @@ fn test4(input : i32) -> i32[4, 4] {
   }
   return arr;
 }
+
+#[entry]
+fn test5(input : i32) -> i32[4] {
+  let arr1 : i32[4];
+  for i = 0 to 4 {
+    let red = arr1[i];
+    for k = 0 to 3 {
+      red += k as i32;
+    }
+    @reduce arr1[i] = red + input;
+  }
+  return arr1;
+}

juno_samples/fork_join_tests/src/gpu.sch

+parallel-reduce(test5@reduce);
+
 gvn(*);
 phi-elim(*);
 dce(*);
@@ -7,6 +9,7 @@ gpu(out.test1);
 gpu(out.test2);
 gpu(out.test3);
 gpu(out.test4);
+gpu(out.test5);
 
 ip-sroa(*);
 sroa(*);

juno_samples/fork_join_tests/src/main.rs

@@ -37,6 +37,11 @@ fn main() {
         let output = r.run(9).await;
         let correct = vec![63i32; 16];
         assert(correct, output);
+
+        let mut r = runner!(test5);
+        let output = r.run(4).await;
+        let correct = vec![7i32; 4];
+        assert(correct, output);
     });
 }
 

juno_scheduler/src/compile.rs

@@ -130,6 +130,7 @@ impl FromStr for Appliable {
             "parallel-fork" => Ok(Appliable::Schedule(Schedule::ParallelFork)),
             "parallel-reduce" => Ok(Appliable::Schedule(Schedule::ParallelReduce)),
             "vectorize" => Ok(Appliable::Schedule(Schedule::Vectorizable)),
+            "no-memset" | "no-reset" => Ok(Appliable::Schedule(Schedule::NoResetConstant)),
 
             _ => Err(s.to_string()),
         }