Merge branch 'gpu_scheduled_tests' into 'main'

Misc. improvements for GPU testing

See merge request !126

hercules_ir/src/dot.rs

@@ -219,7 +219,13 @@ fn write_subgraph_header<W: Write>(
     } else {
         write!(w, "label=\"{}\"\n", function.name)?;
     }
-    write!(w, "bgcolor=ivory4\n")?;
+    let color = match function.device {
+        Some(Device::LLVM) => "paleturquoise1",
+        Some(Device::CUDA) => "darkseagreen1",
+        Some(Device::AsyncRust) => "peachpuff1",
+        None => "ivory2",
+    };
+    write!(w, "bgcolor={}\n", color)?;
     write!(w, "cluster=true\n")?;
     Ok(())
 }

hercules_ir/src/loops.rs

@@ -152,16 +152,7 @@ pub fn loops(
         })
         .collect();
 
-    // Step 6: compute the inverse loop map - this maps control nodes to which
-    // loop they are in (keyed by header), if they are in one.
-    let mut inverse_loops = HashMap::new();
-    for (header, (contents, _)) in loops.iter() {
-        for idx in contents.iter_ones() {
-            inverse_loops.insert(NodeID::new(idx), *header);
-        }
-    }
-
-    // Step 7: compute loop tree nesting.
+    // Step 6: compute loop tree nesting.
     let mut nesting = HashMap::new();
     let mut worklist: VecDeque<NodeID> = loops.keys().map(|id| *id).collect();
     while let Some(header) = worklist.pop_front() {
@@ -175,6 +166,24 @@ pub fn loops(
         }
     }
 
+    // Step 7: compute the inverse loop map - this maps control nodes to which
+    // loop they are in (identified by header), if they are in one. Pick the
+    // most nested loop as the loop they are in.
+    let mut inverse_loops = HashMap::new();
+    for (header, (contents, _)) in loops.iter() {
+        for idx in contents.iter_ones() {
+            let id = NodeID::new(idx);
+            if let Some(old_header) = inverse_loops.get(&id)
+                && nesting[old_header] > nesting[header]
+            {
+                // If the inserted header is more deeply nested, don't do anything.
+                assert!(nesting[old_header] != nesting[header] || old_header == header);
+            } else {
+                inverse_loops.insert(id, *header);
+            }
+        }
+    }
+
     LoopTree {
         root,
         loops,

hercules_opt/src/gcm.rs

@@ -5,7 +5,6 @@ use bitvec::prelude::*;
 use either::Either;
 use union_find::{QuickFindUf, UnionBySize, UnionFind};
 
-use hercules_cg::*;
 use hercules_ir::*;
 
 use crate::*;
@@ -837,19 +836,16 @@ fn liveness_dataflow(
         liveness.insert(NodeID::new(bb_idx), vec![BTreeSet::new(); insts.len() + 1]);
     }
     let mut num_phis_reduces = vec![0; function.nodes.len()];
-    let mut reducing = vec![false; function.nodes.len()];
+    let mut has_phi = vec![false; function.nodes.len()];
+    let mut has_reduce = vec![false; function.nodes.len()];
     for (node_idx, bb) in bbs.0.iter().enumerate() {
         let node = &function.nodes[node_idx];
         if node.is_phi() || node.is_reduce() {
             num_phis_reduces[bb.idx()] += 1;
-            // Phis and reduces can't be in the same basic block.
-            if node.is_reduce() {
-                assert!(num_phis_reduces[bb.idx()] == 0 || reducing[bb.idx()]);
-                reducing[bb.idx()] = true;
-            } else {
-                assert!(!reducing[bb.idx()]);
-            }
         }
+        has_phi[bb.idx()] = node.is_phi();
+        has_reduce[bb.idx()] = node.is_reduce();
+        assert!(!has_phi[bb.idx()] || !has_reduce[bb.idx()]);
     }
     let is_obj = |id: NodeID| !objects[&func_id].objects(id).is_empty();
 
@@ -861,7 +857,7 @@ fn liveness_dataflow(
             let last_pt = bbs.1[bb.idx()].len();
             let old_value = &liveness[&bb][last_pt];
             let mut new_value = BTreeSet::new();
-            for succ in control_subgraph.succs(*bb).chain(if reducing[bb.idx()] {
+            for succ in control_subgraph.succs(*bb).chain(if has_reduce[bb.idx()] {
                 Either::Left(once(*bb))
             } else {
                 Either::Right(empty())

hercules_samples/dot/Cargo.toml

@@ -4,6 +4,9 @@ version = "0.1.0"
 authors = ["Russel Arbore <rarbore2@illinois.edu>"]
 edition = "2021"
 
+[features]
+cuda = ["hercules_rt/cuda"]
+
 [build-dependencies]
 juno_build = { path = "../../juno_build" }
 

hercules_samples/dot/build.rs

@@ -4,6 +4,9 @@ fn main() {
     JunoCompiler::new()
         .ir_in_src("dot.hir")
         .unwrap()
+        //.schedule_in_src(if cfg!(feature = "cuda") { "gpu.sch" } else { "cpu.sch" })
+        .schedule_in_src("cpu.sch")
+        .unwrap()
         .build()
         .unwrap();
 }

hercules_samples/dot/src/cpu.sch

+gvn(*);
+phi-elim(*);
+dce(*);
+
+auto-outline(*);
+
+ip-sroa(*);
+sroa(*);
+unforkify(*);
+dce(*);
+
+gcm(*);

hercules_samples/dot/src/gpu.sch

+gvn(*);
+phi-elim(*);
+dce(*);
+
+auto-outline(*);
+gpu(*);
+host(dot);
+
+ip-sroa(*);
+sroa(*);
+dce(*);
+
+gcm(*);

hercules_samples/matmul/Cargo.toml

@@ -4,6 +4,9 @@ version = "0.1.0"
 authors = ["Russel Arbore <rarbore2@illinois.edu>"]
 edition = "2021"
 
+[features]
+cuda = ["hercules_rt/cuda"]
+
 [build-dependencies]
 juno_build = { path = "../../juno_build" }
 

hercules_samples/matmul/build.rs

@@ -4,6 +4,9 @@ fn main() {
     JunoCompiler::new()
         .ir_in_src("matmul.hir")
         .unwrap()
+        //.schedule_in_src(if cfg!(feature = "cuda") { "gpu.sch" } else { "cpu.sch" })
+        .schedule_in_src("cpu.sch")
+        .unwrap()
         .build()
         .unwrap();
 }

hercules_samples/matmul/src/cpu.sch

+gvn(*);
+phi-elim(*);
+dce(*);
+
+auto-outline(*);
+
+ip-sroa(*);
+sroa(*);
+fork-split(*);
+unforkify(*);
+dce(*);
+float-collections(*);
+
+gcm(*);

hercules_samples/matmul/src/gpu.sch

+gvn(*);
+phi-elim(*);
+dce(*);
+
+auto-outline(*);
+gpu(*);
+host(matmul);
+
+ip-sroa(*);
+sroa(*);
+dce(*);
+float-collections(*);
+
+gcm(*);
+xdot[true](*);