Add ebb macro to insert a new blocks in legalization output.

cranelift/codegen/meta/src/cdsl/ast.rs

@@ -4,9 +4,10 @@ use crate::cdsl::operands::{OperandKind, OperandKindFields};
 use crate::cdsl::types::ValueType;
 use crate::cdsl::typevar::{TypeSetBuilder, TypeVar};
 
-use cranelift_entity::{entity_impl, PrimaryMap};
+use cranelift_entity::{entity_impl, PrimaryMap, SparseMap, SparseMapValue};
 
 use std::fmt;
+use std::iter::IntoIterator;
 
 pub(crate) enum Expr {
     Var(VarIndex),
@@ -82,7 +83,7 @@ impl DefPool {
     pub fn next_index(&self) -> DefIndex {
         self.pool.next_key()
     }
-    pub fn create(&mut self, apply: Apply, defined_vars: Vec<VarIndex>) -> DefIndex {
+    pub fn create_inst(&mut self, apply: Apply, defined_vars: Vec<VarIndex>) -> DefIndex {
         self.pool.push(Def {
             apply,
             defined_vars,
@@ -94,6 +95,55 @@ impl DefPool {
 pub(crate) struct DefIndex(u32);
 entity_impl!(DefIndex);
 
+/// A definition which would lead to generate a block creation.
+#[derive(Clone)]
+pub(crate) struct Block {
+    /// Instruction index after which the block entry is set.
+    pub location: DefIndex,
+    /// Variable holding the new created block.
+    pub name: VarIndex,
+}
+
+pub(crate) struct BlockPool {
+    pool: SparseMap<DefIndex, Block>,
+}
+
+impl SparseMapValue<DefIndex> for Block {
+    fn key(&self) -> DefIndex {
+        self.location
+    }
+}
+
+impl BlockPool {
+    pub fn new() -> Self {
+        Self {
+            pool: SparseMap::new(),
+        }
+    }
+    pub fn get(&self, index: DefIndex) -> Option<&Block> {
+        self.pool.get(index)
+    }
+    pub fn create_block(&mut self, name: VarIndex, location: DefIndex) {
+        if self.pool.contains_key(location) {
+            panic!("Attempt to insert 2 blocks after the same instruction")
+        }
+        self.pool.insert(Block { location, name });
+    }
+    pub fn is_empty(&self) -> bool {
+        self.pool.is_empty()
+    }
+}
+
+// Implement IntoIterator such that we can iterate over blocks which are in the block pool.
+impl<'a> IntoIterator for &'a BlockPool {
+    type Item = <&'a SparseMap<DefIndex, Block> as IntoIterator>::Item;
+    type IntoIter = <&'a SparseMap<DefIndex, Block> as IntoIterator>::IntoIter;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.pool.into_iter()
+    }
+}
+
 #[derive(Clone, Debug)]
 pub(crate) enum Literal {
     /// A value of an enumerated immediate operand.
@@ -518,6 +568,7 @@ pub(crate) enum DummyExpr {
     Var(DummyVar),
     Literal(Literal),
     Apply(InstSpec, Vec<DummyExpr>),
+    Block(DummyVar),
 }
 
 #[derive(Clone)]
@@ -561,6 +612,11 @@ impl ExprBuilder {
             defined_vars,
         }
     }
+
+    pub fn block(name: DummyVar) -> Self {
+        let expr = DummyExpr::Block(name);
+        Self { expr }
+    }
 }
 
 macro_rules! def_rhs {
@@ -592,3 +648,11 @@ macro_rules! def {
         def_rhs!($($tt)*).assign_to(Vec::new())
     }
 }
+
+// Helper macro to define legalization recipes.
+macro_rules! ebb {
+    // An basic block definition, splitting the current block in 2.
+    ($block: ident) => {
+        ExprBuilder::block($block).assign_to(Vec::new())
+    };
+}

cranelift/codegen/meta/src/cdsl/xform.rs

@@ -1,5 +1,6 @@
 use crate::cdsl::ast::{
-    Apply, DefIndex, DefPool, DummyDef, DummyExpr, Expr, PatternPosition, VarIndex, VarPool,
+    Apply, BlockPool, DefIndex, DefPool, DummyDef, DummyExpr, Expr, PatternPosition, VarIndex,
+    VarPool,
 };
 use crate::cdsl::instructions::Instruction;
 use crate::cdsl::type_inference::{infer_transform, TypeEnvironment};
@@ -22,6 +23,7 @@ pub(crate) struct Transform {
     pub dst: Vec<DefIndex>,
     pub var_pool: VarPool,
     pub def_pool: DefPool,
+    pub block_pool: BlockPool,
     pub type_env: TypeEnvironment,
 }
 
@@ -31,6 +33,7 @@ impl Transform {
     fn new(src: DummyDef, dst: Vec<DummyDef>) -> Self {
         let mut var_pool = VarPool::new();
         let mut def_pool = DefPool::new();
+        let mut block_pool = BlockPool::new();
 
         let mut input_vars: Vec<VarIndex> = Vec::new();
         let mut defined_vars: Vec<VarIndex> = Vec::new();
@@ -47,6 +50,7 @@ impl Transform {
             &mut defined_vars,
             &mut var_pool,
             &mut def_pool,
+            &mut block_pool,
         )[0];
 
         let num_src_inputs = input_vars.len();
@@ -59,6 +63,7 @@ impl Transform {
             &mut defined_vars,
             &mut var_pool,
             &mut def_pool,
+            &mut block_pool,
         );
 
         // Sanity checks.
@@ -122,6 +127,7 @@ impl Transform {
             dst,
             var_pool,
             def_pool,
+            block_pool,
             type_env,
         }
     }
@@ -224,6 +230,9 @@ fn rewrite_expr(
             DummyExpr::Apply(..) => {
                 panic!("Recursive apply is not allowed.");
             }
+            DummyExpr::Block(_block) => {
+                panic!("Blocks are not valid arguments.");
+            }
         }
     }
 
@@ -238,8 +247,18 @@ fn rewrite_def_list(
     defined_vars: &mut Vec<VarIndex>,
     var_pool: &mut VarPool,
     def_pool: &mut DefPool,
+    block_pool: &mut BlockPool,
 ) -> Vec<DefIndex> {
     let mut new_defs = Vec::new();
+    // Register variable names of new blocks first as a block name can be used to jump forward. Thus
+    // the name has to be registered first to avoid misinterpreting it as an input-var.
+    for dummy_def in dummy_defs.iter() {
+        if let DummyExpr::Block(ref var) = dummy_def.expr {
+            var_index(var.name, symbol_table, defined_vars, var_pool);
+        }
+    }
+
+    // Iterate over the definitions and blocks, to map variables names to inputs or outputs.
     for dummy_def in dummy_defs {
         let def_index = def_pool.next_index();
 
@@ -251,7 +270,21 @@ fn rewrite_def_list(
             defined_vars,
             var_pool,
         );
-        let new_apply = rewrite_expr(position, dummy_def.expr, symbol_table, input_vars, var_pool);
+        if let DummyExpr::Block(var) = dummy_def.expr {
+            let var_index = *symbol_table
+                .get(var.name)
+                .or_else(|| {
+                    panic!(
+                        "Block {} was not registered during the first visit",
+                        var.name
+                    )
+                })
+                .unwrap();
+            var_pool.get_mut(var_index).set_def(position, def_index);
+            block_pool.create_block(var_index, def_index);
+        } else {
+            let new_apply =
+                rewrite_expr(position, dummy_def.expr, symbol_table, input_vars, var_pool);
 
             assert!(
                 def_pool.next_index() == def_index,
@@ -263,7 +296,8 @@ fn rewrite_def_list(
                 "number of Var results in instruction is incorrect"
             );
 
-        new_defs.push(def_pool.create(new_apply, new_defined_vars));
+            new_defs.push(def_pool.create_inst(new_apply, new_defined_vars));
+        }
     }
     new_defs
 }

cranelift/codegen/meta/src/gen_legalizer.rs

@@ -401,14 +401,30 @@ fn gen_transform<'a>(
     // Guard the actual expansion by `predicate`.
     fmt.line("if predicate {");
     fmt.indent(|fmt| {
+        // If we are adding some blocks, we need to recall the original block, such that we can
+        // recompute it.
+        if !transform.block_pool.is_empty() {
+            fmt.line("let orig_ebb = pos.current_ebb().unwrap();");
+        }
+
         // If we're going to delete `inst`, we need to detach its results first so they can be
         // reattached during pattern expansion.
         if !replace_inst {
             fmt.line("pos.func.dfg.clear_results(inst);");
         }
 
+        // Emit new block creation.
+        for block in &transform.block_pool {
+            let var = transform.var_pool.get(block.name);
+            fmtln!(fmt, "let {} = pos.func.dfg.make_ebb();", var.name);
+        }
+
         // Emit the destination pattern.
         for &def_index in &transform.dst {
+            if let Some(block) = transform.block_pool.get(def_index) {
+                let var = transform.var_pool.get(block.name);
+                fmtln!(fmt, "pos.insert_ebb({});", var.name);
+            }
             emit_dst_inst(
                 transform.def_pool.get(def_index),
                 &transform.def_pool,
@@ -417,17 +433,33 @@ fn gen_transform<'a>(
             );
         }
 
+        // Insert a new block after the last instruction, if needed.
+        let def_next_index = transform.def_pool.next_index();
+        if let Some(block) = transform.block_pool.get(def_next_index) {
+            let var = transform.var_pool.get(block.name);
+            fmtln!(fmt, "pos.insert_ebb({});", var.name);
+        }
+
         // Delete the original instruction if we didn't have an opportunity to replace it.
         if !replace_inst {
             fmt.line("let removed = pos.remove_inst();");
             fmt.line("debug_assert_eq!(removed, inst);");
         }
 
+        if transform.block_pool.is_empty() {
             if transform.def_pool.get(transform.src).apply.inst.is_branch {
                 // A branch might have been legalized into multiple branches, so we need to recompute
                 // the cfg.
                 fmt.line("cfg.recompute_ebb(pos.func, pos.current_ebb().unwrap());");
             }
+        } else {
+            // Update CFG for the new blocks.
+            fmt.line("cfg.recompute_ebb(pos.func, orig_ebb);");
+            for block in &transform.block_pool {
+                let var = transform.var_pool.get(block.name);
+                fmtln!(fmt, "cfg.recompute_ebb(pos.func, {});", var.name);
+            }
+        }
 
         fmt.line("return true;");
     });