Add set ops to scheduler

Cargo.lock

@@ -1459,6 +1459,7 @@ dependencies = [
 name = "juno_scheduler"
 version = "0.0.1"
 dependencies = [
+ "bitvec",
  "cfgrammar",
  "hercules_cg",
  "hercules_ir",

hercules_opt/src/editor.rs

@@ -108,9 +108,8 @@ impl<'a: 'b, 'b> FunctionEditor<'a> {
         }
     }
 
-    // Constructs an editor but only makes the nodes with at least one of the
-    // set of labels as mutable.
-    pub fn new_labeled(
+    // Constructs an editor with a specified mask determining which nodes are mutable
+    pub fn new_mask(
         function: &'a mut Function,
         function_id: FunctionID,
         constants: &'a RefCell<Vec<Constant>>,
@@ -118,7 +117,7 @@ impl<'a: 'b, 'b> FunctionEditor<'a> {
         types: &'a RefCell<Vec<Type>>,
         labels: &'a RefCell<Vec<String>>,
         def_use: &ImmutableDefUseMap,
-        with_labels: &HashSet<LabelID>,
+        mask: BitVec<u8, Lsb0>,
     ) -> Self {
         let mut_def_use = (0..function.nodes.len())
             .map(|idx| {
@@ -130,14 +129,6 @@ impl<'a: 'b, 'b> FunctionEditor<'a> {
             })
             .collect();
 
-        let mut mutable_nodes = bitvec![u8, Lsb0; 0; function.nodes.len()];
-        // Add all nodes which have some label which is in the with_labels set
-        for (idx, labels) in function.labels.iter().enumerate() {
-            if !labels.is_disjoint(with_labels) {
-                mutable_nodes.set(idx, true);
-            }
-        }
-
         FunctionEditor {
             function,
             function_id,
@@ -146,7 +137,7 @@ impl<'a: 'b, 'b> FunctionEditor<'a> {
             types,
             labels,
             mut_def_use,
-            mutable_nodes,
+            mutable_nodes: mask,
             modified: false,
         }
     }

juno_scheduler/Cargo.toml

@@ -24,4 +24,5 @@ hercules_ir = { path = "../hercules_ir" }
 hercules_opt = { path = "../hercules_opt" }
 juno_utils = { path = "../juno_utils" }
 postcard = { version = "*", features = ["alloc"] }
-serde = { version = "*", features = ["derive"] }
+serde = { version = "*", features = ["derive"] }
\ No newline at end of file
+bitvec = "*"

juno_scheduler/src/compile.rs

@@ -473,6 +473,20 @@ fn compile_expr(
             }
             Ok(ExprResult::Expr(ir::ScheduleExp::Record { fields: result }))
         }
+        parser::Expr::SetOp {
+            span: _,
+            op,
+            lhs,
+            rhs,
+        } => {
+            let lhs = compile_exp_as_expr(*lhs, lexer, macrostab, macros)?;
+            let rhs = compile_exp_as_expr(*rhs, lexer, macrostab, macros)?;
+            Ok(ExprResult::Expr(ir::ScheduleExp::SetOp {
+                op,
+                lhs: Box::new(lhs),
+                rhs: Box::new(rhs),
+            }))
+        }
     }
 }
 

juno_scheduler/src/ir.rs

 use hercules_ir::ir::{Device, Schedule};
 
+use crate::parser;
+
 #[derive(Debug, Copy, Clone)]
 pub enum Pass {
     ArraySLF,
@@ -117,6 +119,11 @@ pub enum ScheduleExp {
         body: Vec<ScheduleStmt>,
         res: Box<ScheduleExp>,
     },
+    SetOp {
+        op: parser::SetOp,
+        lhs: Box<ScheduleExp>,
+        rhs: Box<ScheduleExp>,
+    },
     // This is used to "box" a selection by evaluating it at one point and then
     // allowing it to be used as a selector later on
     Selection {

juno_scheduler/src/lang.l

@@ -39,6 +39,10 @@ false           "false"
 \{              "{"
 \}              "}"
 
+\\              "\\"
+\|              "|"
+&               "&"
+
 panic[\t \n\r]+after "panic_after"
 print[\t \n\r]+iter  "print_iter"
 stop[\t \n\r]+after  "stop_after"

juno_scheduler/src/lang.y

@@ -3,6 +3,11 @@
 %avoid_insert "ID" "INT" "STRING"
 %expect-unused Unmatched 'UNMATCHED'
 
+%left '\\'
+%left '|'
+%left '&'
+%left '.' '@'
+
 %%
 
 Schedule -> OperationList
@@ -59,6 +64,12 @@ Expr -> Expr
       { Expr::BlockExpr { span: $span, body: Box::new($2) } }
   | '<' Fields '>'
       { Expr::Record { span: $span, fields: rev($2) } }
+  | Expr '\\' Expr
+      { Expr::SetOp { span: $span, op: SetOp::Difference, lhs: Box::new($1), rhs: Box::new($3) } }
+  | Expr '|' Expr
+      { Expr::SetOp { span: $span, op: SetOp::Union, lhs: Box::new($1), rhs: Box::new($3) } }
+  | Expr '&' Expr
+      { Expr::SetOp { span: $span, op: SetOp::Intersection, lhs: Box::new($1), rhs: Box::new($3) } }
   ;
 
 Args -> Vec<Expr>
@@ -151,6 +162,13 @@ pub enum FixpointLimit {
   PrintIter  { span: Span },
 }
 
+#[derive(Copy, Clone, Debug)]
+pub enum SetOp {
+  Difference,
+  Union,
+  Intersection,
+}
+
 pub enum Expr {
   Function    { span: Span, name: Span, args: Vec<Expr>, selection: Selector },
   Macro       { span: Span, name: Span, args: Vec<Expr>, selection: Selector },
@@ -161,6 +179,7 @@ pub enum Expr {
   Field       { span: Span, lhs: Box<Expr>, field: Span },
   BlockExpr   { span: Span, body: Box<OperationList> },
   Record      { span: Span, fields: Vec<(Span, Expr)> },
+  SetOp       { span: Span, op: SetOp, lhs: Box<Expr>, rhs: Box<Expr> },
 }
 
 pub enum Selector {