Fix vectorizable inference

hercules_cg/src/gpu.rs

@@ -3,8 +3,6 @@ extern crate hercules_ir;
 
 use std::collections::{BTreeMap, HashMap, HashSet};
 use std::fmt::{Error, Write};
-use std::fs::{File, OpenOptions};
-use std::io::Write as _;
 
 use self::hercules_ir::*;
 

hercules_opt/src/schedule.rs

@@ -122,26 +122,46 @@ pub fn infer_parallel_reduce(
 
 /*
  * Infer vectorizable fork-joins. Just check that there are no control nodes
- * between a fork and its join and the factor is a constant.
+ * between a fork and its join that depend on thread IDs and the factor is a constant.
  */
 pub fn infer_vectorizable(editor: &mut FunctionEditor, fork_join_map: &HashMap<NodeID, NodeID>) {
-    for id in editor.node_ids() {
-        let func = editor.func();
-        if !func.nodes[id.idx()].is_join() {
+    for (fork, join) in fork_join_map {
+        let nodes = &editor.func().nodes;
+
+        let factors = nodes[fork.idx()].try_fork().unwrap().1;
+        if factors.len() != 1
+            || evaluate_dynamic_constant(factors[0], &editor.get_dynamic_constants()).is_none()
+        {
             continue;
         }
 
-        let u = get_uses(&func.nodes[id.idx()]).as_ref()[0];
-        if let Some(join) = fork_join_map.get(&u)
-            && *join == id
-        {
-            let factors = func.nodes[u.idx()].try_fork().unwrap().1;
-            if factors.len() == 1
-                && evaluate_dynamic_constant(factors[0], &editor.get_dynamic_constants()).is_some()
-            {
-                editor.edit(|edit| edit.add_schedule(u, Schedule::Vectorizable));
+        let mut worklist: Vec<_> = editor
+            .get_users(*fork)
+            .filter(|id| nodes[id.idx()].is_thread_id())
+            .collect();
+        let mut set: HashSet<_> = worklist.iter().map(|id| *id).collect();
+        let mut found_control = false;
+        'worklist: while let Some(item) = worklist.pop() {
+            for u in editor.get_users(item) {
+                if nodes[u.idx()].is_control() {
+                    found_control = true;
+                    break 'worklist;
+                }
+
+                let terminate = nodes[u.idx()]
+                    .try_reduce()
+                    .map(|(control, _, _)| control == *join)
+                    .unwrap_or(false);
+                if !set.contains(&u) && !terminate {
+                    worklist.push(u);
+                }
+                set.insert(u);
             }
         }
+
+        if !found_control {
+            editor.edit(|edit| edit.add_schedule(*fork, Schedule::Vectorizable));
+        }
     }
 }
 
@@ -155,13 +175,10 @@ pub fn infer_monoid_reduce(
     reduce_cycles: &HashMap<NodeID, HashSet<NodeID>>,
 ) {
     let is_binop_monoid = |op| {
-        matches!(
-            op,
-            BinaryOperator::Add
-                | BinaryOperator::Mul
-                | BinaryOperator::Or
-                | BinaryOperator::And
-        )
+        op == BinaryOperator::Add
+            || op == BinaryOperator::Mul
+            || op == BinaryOperator::Or
+            || op == BinaryOperator::And
     };
     let is_intrinsic_monoid = |intrinsic| matches!(intrinsic, Intrinsic::Max | Intrinsic::Min);