Add extra dimensions to arrays when they'd cause a race otherwise

hercules_opt/src/gcm.rs

@@ -1053,19 +1053,22 @@ fn add_extra_collection_dims(
     devices: &Vec<Device>,
     bbs: &BasicBlocks,
 ) -> bool {
-    if devices[editor.func_id().idx()] == Device::AsyncRust
-        && editor.func().name == "_1_laplacian_estimate"
-    {
+    if devices[editor.func_id().idx()] == Device::AsyncRust {
         // Look for collection constant nodes inside fork-joins that are mutated
         // inside the fork-join, aren't involved in any of the reduces of the
         // fork-join, and have a user that isn't a direct read based on all of
         // the thread IDs.
-        let nodes = &editor.func().nodes;
         let fco = &objects[&editor.func_id()];
-        for id in editor.node_ids().filter(|id| {
-            nodes[id.idx()].is_constant() && !editor.get_type(typing[id.idx()]).is_primitive()
-        }) {
+        let candidates: Vec<_> = editor
+            .node_ids()
+            .filter(|id| {
+                editor.func().nodes[id.idx()].is_constant()
+                    && !editor.get_type(typing[id.idx()]).is_primitive()
+            })
+            .collect();
+        for id in candidates {
             // Check all of the above conditions.
+            let nodes = &editor.func().nodes;
             if editor.get_users(id).len() != 1 {
                 continue;
             }
@@ -1115,7 +1118,42 @@ fn add_extra_collection_dims(
 
             // We know that this collection needs to be replicated across the
             // fork-join dimensions, so do that.
-            todo!()
+            let ty = typing[id.idx()];
+            let num_dims: Vec<_> = forks
+                .into_iter()
+                .rev()
+                .map(|id| nodes[id.idx()].try_fork().unwrap().1.len())
+                .collect();
+            let factors = forks
+                .into_iter()
+                .rev()
+                .flat_map(|id| nodes[id.idx()].try_fork().unwrap().1.into_iter())
+                .map(|dc| *dc)
+                .collect();
+            let array_ty = Type::Array(ty, factors);
+            let success = editor.edit(|mut edit| {
+                let new_ty = edit.add_type(array_ty);
+                let new_cons = edit.add_zero_constant(new_ty);
+                let new_cons = edit.add_node(Node::Constant { id: new_cons });
+                let mut tids = vec![];
+                for (fork, num_dims) in forks.into_iter().rev().zip(num_dims) {
+                    for dim in 0..num_dims {
+                        tids.push(edit.add_node(Node::ThreadID {
+                            control: *fork,
+                            dimension: dim,
+                        }));
+                    }
+                }
+                let read = edit.add_node(Node::Read {
+                    collect: new_cons,
+                    indices: Box::new([Index::Position(tids.into_boxed_slice())]),
+                });
+                edit = edit.replace_all_uses(id, read)?;
+                edit = edit.delete_node(id)?;
+                Ok(edit)
+            });
+            assert!(success);
+            return true;
         }
     }
     false