Fix parallel code gen in RT backend

• Properly use Arc<Mutex<_Option<JoinHandle<_, _>>_>> for task parallel calls.
  • Include Arc clones right before opening async closure (for async calls and for fork-joins).
• Make data parallel fork-joins not task parallel by default (can still be outlined for task parallelism).
• Fix edge case in function inlining.

hercules_cg/src/rt.rs

@@ -138,6 +138,7 @@ struct RTContext<'a> {
 #[derive(Debug, Clone, Default)]
 struct RustBlock {
     prologue: String,
+    arc_clones: String,
     data: String,
     phi_tmp_assignments: String,
     phi_assignments: String,
@@ -268,8 +269,9 @@ impl<'a> RTContext<'a> {
             } else {
                 write!(
                     w,
-                    "{}{}{}{}{}",
+                    "{}{}{}{}{}{}",
                     block.prologue,
+                    block.arc_clones,
                     block.data,
                     block.phi_tmp_assignments,
                     block.phi_assignments,
@@ -416,6 +418,14 @@ impl<'a> RTContext<'a> {
                 let succ = self.control_subgraph.succs(id).next().unwrap();
                 write!(epilogue, "{} => {{", id.idx())?;
 
+                // Await the empty futures for the fork-joins, waiting for them
+                // to complete.
+                write!(
+                    epilogue,
+                    "for fut in fork_{}.drain(..) {{ fut.await; }}; ::std::sync::atomic::fence(::std::sync::atomic::Ordering::Acquire);",
+                    fork.idx(),
+                )?;
+
                 // Emit the assignments to the reduce variables in the
                 // surrounding context. It's very unfortunate that we have to do
                 // it while lowering the join node (rather than the reduce nodes
@@ -558,12 +568,14 @@ impl<'a> RTContext<'a> {
                 assert_eq!(control, bb);
                 // The device backends ensure that device functions have the
                 // same interface as AsyncRust functions.
-                let block = &mut blocks.get_mut(&bb).unwrap().data;
+                let block = &mut blocks.get_mut(&bb).unwrap();
+                let arc_clones = &mut block.arc_clones;
+                let block = &mut block.data;
                 let is_async = func.schedules[id.idx()].contains(&Schedule::AsyncCall);
                 if is_async {
                     for arg in args {
-                        if let Some(arc) = self.clone_arc(*arg, false) {
-                            write!(block, "{}", arc)?;
+                        if let Some(arc) = self.clone_arc(*arg) {
+                            write!(arc_clones, "{}", arc)?;
                         }
                     }
                 }
@@ -572,17 +584,19 @@ impl<'a> RTContext<'a> {
                     (Device::AsyncRust, false) | (_, false) => {
                         format!("{} = ", self.get_value(id, bb, true))
                     }
-                    (_, true) => format!(
-                        "{}::async_std::task::spawn(async move {{ async_call_sender_{}.send(",
-                        self.clone_arc(id, true).unwrap(),
-                        id.idx()
-                    ),
+                    (_, true) => {
+                        write!(arc_clones, "{}", self.clone_arc(id).unwrap())?;
+                        format!(
+                            "*async_call_{}.lock().await = Some(::async_std::task::spawn(async move {{ ",
+                            id.idx(),
+                        )
+                    }
                 };
                 let postfix = match (device, is_async) {
                     (Device::AsyncRust, false) => ".await",
                     (_, false) => "",
-                    (Device::AsyncRust, true) => ".await).await})",
-                    (_, true) => ").await})",
+                    (Device::AsyncRust, true) => ".await}))",
+                    (_, true) => "}))",
                 };
                 write!(
                     block,
@@ -1079,14 +1093,14 @@ impl<'a> RTContext<'a> {
             }
 
             // If the node is a call with an AsyncCall schedule, it should be
-            // lowered to a channel.
+            // lowered to a Arc<Mutex<>> over the future.
             let is_async_call =
                 func.nodes[idx].is_call() && func.schedules[idx].contains(&Schedule::AsyncCall);
             if is_async_call {
                 write!(
                     w,
-                    "let mut async_call_channel_{} = ::async_std::channel::bounded(1);let async_call_sender_{} = ::std::sync::Arc::new(async_call_channel_{}.0);let async_call_receiver_{} = ::std::sync::Arc::new(async_call_channel_{}.1);",
-                    idx, idx, idx, idx, idx
+                    "let mut async_call_{} = ::std::sync::Arc::new(::async_std::sync::Mutex::new(None));",
+                    idx,
                 )?;
             } else {
                 write!(
@@ -1353,40 +1367,30 @@ impl<'a> RTContext<'a> {
     fn get_value(&self, id: NodeID, bb: NodeID, lhs: bool) -> String {
         let func = self.get_func();
         if let Some((control, _, _)) = func.nodes[id.idx()].try_reduce()
-            && control == bb
+            && (control == bb
+                || !self.nodes_in_fork_joins[&self.join_fork_map[&control]].contains(&bb))
         {
             format!("reduce_{}", id.idx())
-        } else if let Some((control, _, _)) = func.nodes[id.idx()].try_reduce()
-            && let fork = self.join_fork_map[&control]
-            && !self.nodes_in_fork_joins[&fork].contains(&bb)
-        {
-            // Before using the value of a reduction outside the fork-join,
-            // await the futures.
-            format!(
-                "{{for fut in fork_{}.drain(..) {{ fut.await; }}; ::std::sync::atomic::fence(::std::sync::atomic::Ordering::Acquire); reduce_{}}}",
-                fork.idx(),
-                id.idx()
-            )
         } else if func.nodes[id.idx()].is_call()
             && func.schedules[id.idx()].contains(&Schedule::AsyncCall)
         {
             assert!(!lhs);
-            format!("async_call_receiver_{}.recv().await.unwrap()", id.idx(),)
+            format!(
+                "async_call_{}.lock().await.as_mut().unwrap().await",
+                id.idx(),
+            )
         } else {
             format!("node_{}", id.idx())
         }
     }
 
-    fn clone_arc(&self, id: NodeID, lhs: bool) -> Option<String> {
+    fn clone_arc(&self, id: NodeID) -> Option<String> {
         let func = self.get_func();
         if func.nodes[id.idx()].is_call() && func.schedules[id.idx()].contains(&Schedule::AsyncCall)
         {
-            let kind = if lhs { "sender" } else { "receiver" };
             Some(format!(
-                "let async_call_{}_{} = async_call_{}_{}.clone();",
-                kind,
+                "let async_call_{} = async_call_{}.clone();",
                 id.idx(),
-                kind,
                 id.idx()
             ))
         } else {