Fork improvements

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,27 @@ pub fn forkify_loop(
             };
             let reduce_id = edit.add_node(reduce);
 
+            if (!edit.get_node(init).is_reduce()
+                && edit.get_schedule(init).contains(&Schedule::ParallelReduce))
+                || (!edit.get_node(continue_latch).is_reduce()
+                    && edit
+                        .get_schedule(continue_latch)
+                        .contains(&Schedule::ParallelReduce))
+            {
+                edit = edit.add_schedule(reduce_id, Schedule::ParallelReduce)?;
+            }
+            if (!edit.get_node(init).is_reduce()
+                && edit
+                    .get_schedule(init)
+                    .contains(&Schedule::TightAssociative))
+                || (!edit.get_node(continue_latch).is_reduce()
+                    && 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 +435,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 +485,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 +533,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()),
         }

juno_scheduler/src/pm.rs

@@ -780,8 +780,15 @@ impl PassManager {
             .stdin(Stdio::piped())
             .stdout(Stdio::piped())
             .spawn()
-            .expect("Error running clang. Is it installed?");
-        assert!(clang_process.wait().unwrap().success());
+            .expect("PANIC: Error running clang. Is it installed?");
+        if clang_process
+            .wait()
+            .map(|status| !status.success())
+            .unwrap_or(false)
+        {
+            let path = tmp_dir.into_path();
+            panic!("PANIC: Clang failed to compile the LLVM IR module. Persisting temporary directory ({}).", path.display());
+        }
 
         let mut ar_args = vec!["crus", &output_archive, &llvm_object];
 
@@ -806,8 +813,15 @@ impl PassManager {
                 .arg(&cuda_object)
                 .arg(&cuda_path)
                 .spawn()
-                .expect("Error running nvcc. Is it installed?");
-            assert!(nvcc_process.wait().unwrap().success());
+                .expect("PANIC: Error running NVCC. Is it installed?");
+            if nvcc_process
+                .wait()
+                .map(|status| !status.success())
+                .unwrap_or(false)
+            {
+                let path = tmp_dir.into_path();
+                panic!("PANIC: NVCC failed to compile the CUDA module. Persisting temporary directory ({}).", path.display());
+            }
 
             ar_args.push(&cuda_object);
         }
@@ -816,7 +830,17 @@ impl PassManager {
             .args(&ar_args)
             .spawn()
             .expect("Error running ar. Is it installed?");
-        assert!(ar_process.wait().unwrap().success());
+        if ar_process
+            .wait()
+            .map(|status| !status.success())
+            .unwrap_or(false)
+        {
+            let path = tmp_dir.into_path();
+            panic!(
+                "PANIC: Ar failed to create a static library. Persisting temporary directory ({}).",
+                path.display()
+            );
+        }
 
         // Write the Rust runtime into a file.
         let output_rt = format!("{}/rt_{}.hrt", output_dir, module_name);