From 53f251c47200a9acaa3c11dc6e2b71bd9e8e8d8f Mon Sep 17 00:00:00 2001 From: Jake Fecher Date: Thu, 18 May 2023 12:02:11 -0500 Subject: [PATCH 1/5] Add simplify_cfg pass --- .../src/ssa_refactor/ir/basic_block.rs | 13 ++ .../src/ssa_refactor/ir/dfg.rs | 6 + .../src/ssa_refactor/ir/function.rs | 19 ++ .../src/ssa_refactor/opt/mod.rs | 1 + .../src/ssa_refactor/opt/simplify_cfg.rs | 172 ++++++++++++++++++ 5 files changed, 211 insertions(+) create mode 100644 crates/noirc_evaluator/src/ssa_refactor/opt/simplify_cfg.rs diff --git a/crates/noirc_evaluator/src/ssa_refactor/ir/basic_block.rs b/crates/noirc_evaluator/src/ssa_refactor/ir/basic_block.rs index 30526bc296e..a41563fe21c 100644 --- a/crates/noirc_evaluator/src/ssa_refactor/ir/basic_block.rs +++ b/crates/noirc_evaluator/src/ssa_refactor/ir/basic_block.rs @@ -40,6 +40,11 @@ impl BasicBlock { &self.parameters } + /// Removes all the parameters of this block + pub(crate) fn take_parameters(&mut self) -> Vec { + std::mem::take(&mut self.parameters) + } + /// Adds a parameter to this BasicBlock. /// Expects that the ValueId given should refer to a Value::Param /// instance with its position equal to self.parameters.len(). @@ -104,4 +109,12 @@ impl BasicBlock { }); self.instructions.remove(index); } + + /// Take ownership of this block's terminator, replacing it with an empty return terminator. + /// It is expected that this function is only called on blocks that are no longer reachable + /// as an optimization to copy over a terminator without cloning. + pub(crate) fn take_terminator(&mut self) -> TerminatorInstruction { + let terminator = self.terminator.as_mut().expect("Expected block to have a terminator"); + std::mem::replace(terminator, TerminatorInstruction::Return { return_values: Vec::new() }) + } } diff --git a/crates/noirc_evaluator/src/ssa_refactor/ir/dfg.rs b/crates/noirc_evaluator/src/ssa_refactor/ir/dfg.rs index fc15f3e2168..217472ad6d1 100644 --- a/crates/noirc_evaluator/src/ssa_refactor/ir/dfg.rs +++ b/crates/noirc_evaluator/src/ssa_refactor/ir/dfg.rs @@ -123,6 +123,12 @@ impl DataFlowGraph { self.values.insert(value) } + /// Set the value of value_to_replace to refer to the value referred to by new_value. + pub(crate) fn set_value_from_id(&mut self, value_to_replace: ValueId, new_value: ValueId) { + let new_value = self.values[new_value]; + self.values[value_to_replace] = new_value; + } + /// Creates a new constant value, or returns the Id to an existing one if /// one already exists. pub(crate) fn make_constant(&mut self, value: FieldElement, typ: Type) -> ValueId { diff --git a/crates/noirc_evaluator/src/ssa_refactor/ir/function.rs b/crates/noirc_evaluator/src/ssa_refactor/ir/function.rs index f37448462b7..06f2fd613b9 100644 --- a/crates/noirc_evaluator/src/ssa_refactor/ir/function.rs +++ b/crates/noirc_evaluator/src/ssa_refactor/ir/function.rs @@ -1,3 +1,5 @@ +use std::collections::HashSet; + use super::basic_block::BasicBlockId; use super::dfg::DataFlowGraph; use super::map::Id; @@ -61,6 +63,23 @@ impl Function { pub(crate) fn parameters(&self) -> &[ValueId] { self.dfg.block_parameters(self.entry_block) } + + /// Collects all the reachable blocks of this function. + /// + /// Note that self.dfg.basic_blocks_iter() iterates over all blocks, + /// whether reachable or not. This function should be used if you + /// want to iterate only reachable blocks. + pub(crate) fn reachable_blocks(&self) -> HashSet { + let mut blocks = HashSet::new(); + let mut stack = vec![self.entry_block]; + + while let Some(block) = stack.pop() { + if blocks.insert(block) { + stack.extend(self.dfg[block].successors()); + } + } + blocks + } } /// FunctionId is a reference for a function diff --git a/crates/noirc_evaluator/src/ssa_refactor/opt/mod.rs b/crates/noirc_evaluator/src/ssa_refactor/opt/mod.rs index 46ca7d443bc..2701b0bb73c 100644 --- a/crates/noirc_evaluator/src/ssa_refactor/opt/mod.rs +++ b/crates/noirc_evaluator/src/ssa_refactor/opt/mod.rs @@ -4,3 +4,4 @@ //! simpler form until the IR only has a single function remaining with 1 block within it. //! Generally, these passes are also expected to minimize the final amount of instructions. mod inlining; +mod simplify_cfg; diff --git a/crates/noirc_evaluator/src/ssa_refactor/opt/simplify_cfg.rs b/crates/noirc_evaluator/src/ssa_refactor/opt/simplify_cfg.rs new file mode 100644 index 00000000000..1d24edef146 --- /dev/null +++ b/crates/noirc_evaluator/src/ssa_refactor/opt/simplify_cfg.rs @@ -0,0 +1,172 @@ +//! This file contains the simplify cfg pass of the SSA IR. +//! +//! This is a rather simple pass that is expected to be cheap to perform. It: +//! 1. Removes blocks with no predecessors +//! 2. Inlines a block into its sole predecessor if that predecessor only has one successor. +//! 3. Removes any block arguments for blocks with only a single predecessor. +//! 4. Removes any blocks which have no instructions other than a single terminating jmp. +//! +//! Currently, only 2 is implemented. +use std::collections::HashSet; + +use crate::ssa_refactor::{ + ir::{ + basic_block::BasicBlockId, cfg::ControlFlowGraph, function::Function, + instruction::TerminatorInstruction, + }, + ssa_gen::Ssa, +}; + +impl Ssa { + /// Simplify each function's control flow graph by: + /// 1. Removing blocks with no predecessors + /// 2. Inlining a block into its sole predecessor if that predecessor only has one successor. + /// 3. Removing any block arguments for blocks with only a single predecessor. + /// 4. Removing any blocks which have no instructions other than a single terminating jmp. + /// + /// Currently, only 2 is implemented. + pub(crate) fn simplify_cfg(mut self) -> Self { + for function in self.functions.values_mut() { + simplify_function(function); + } + self + } +} + +/// Simplify a function's cfg by going through each block to check for any simple blocks that can +/// be inlined into their predecessor. +fn simplify_function(function: &mut Function) { + let mut cfg = ControlFlowGraph::with_function(function); + let mut stack = vec![function.entry_block()]; + let mut visited = HashSet::new(); + + while let Some(block) = stack.pop() { + if visited.insert(block) { + stack.extend(function.dfg[block].successors().filter(|block| !visited.contains(block))); + } + + let mut predecessors = cfg.predecessors(block); + + if predecessors.len() == 1 { + let predecessor = predecessors.next().expect("Already checked length of predecessors"); + drop(predecessors); + + if try_inline_into_predecessor(function, &mut cfg, block, predecessor) { + continue; + } + } + } +} + +/// Try to inline a block into its predecessor, returning true if successful. +/// +/// This will only occur if the predecessor's only successor is the given block. +/// It is also expected that the given block's only predecessor is the given one. +fn try_inline_into_predecessor( + function: &mut Function, + cfg: &mut ControlFlowGraph, + block_id: BasicBlockId, + predecessor_id: BasicBlockId, +) -> bool { + let mut successors = cfg.successors(predecessor_id); + if successors.len() == 1 && successors.next() == Some(block_id) { + drop(successors); + + // First remove all the instructions and terminator from the block we're removing + let block = &mut function.dfg[block_id]; + let mut instructions = std::mem::take(block.instructions_mut()); + let terminator = block.take_terminator(); + + // Then append each to the predecessor + let predecessor = &mut function.dfg[predecessor_id]; + predecessor.instructions_mut().append(&mut instructions); + + let jump_args = match predecessor.take_terminator() { + TerminatorInstruction::Jmp { arguments, .. } => arguments, + _ => unreachable!( + "Predecessor was already validated to have only a single jmp destination" + ), + }; + + predecessor.set_terminator(terminator); + + // Finally, replace any block parameters with their values, if there were any. + let block_params = function.dfg[block_id].take_parameters(); + assert_eq!(block_params.len(), jump_args.len()); + + for (param, arg) in block_params.iter().zip(jump_args) { + function.dfg.set_value_from_id(*param, arg); + } + + cfg.recompute_block(function, block_id); + cfg.recompute_block(function, predecessor_id); + + true + } else { + false + } +} + +#[cfg(test)] +mod test { + use crate::ssa_refactor::{ + ir::{instruction::TerminatorInstruction, map::Id, types::Type}, + ssa_builder::FunctionBuilder, + }; + + #[test] + fn inline_blocks() { + // fn main { + // b0(): + // jmp b1(Field 7) + // b1(v0: Field): + // jmp b2(v0) + // b2(v1: Field): + // return v1 + // } + let main_id = Id::test_new(0); + let mut builder = FunctionBuilder::new("main".into(), main_id); + + let b1 = builder.insert_block(); + let b2 = builder.insert_block(); + + let v0 = builder.add_block_parameter(b1, Type::field()); + let v1 = builder.add_block_parameter(b2, Type::field()); + + let expected_return = 7u128; + let seven = builder.field_constant(expected_return); + builder.terminate_with_jmp(b1, vec![seven]); + + builder.switch_to_block(b1); + builder.terminate_with_jmp(b2, vec![v0]); + + builder.switch_to_block(b2); + builder.terminate_with_return(vec![v1]); + + let ssa = builder.finish(); + assert_eq!(ssa.main().reachable_blocks().len(), 3); + + // Expected output: + // fn main { + // b0(): + // return Field 7 + // } + let ssa = ssa.simplify_cfg(); + let main = ssa.main(); + println!("{}", main); + assert_eq!(main.reachable_blocks().len(), 1); + + match main.dfg[main.entry_block()].terminator() { + Some(TerminatorInstruction::Return { return_values }) => { + assert_eq!(return_values.len(), 1); + let return_value = main + .dfg + .get_numeric_constant(return_values[0]) + .expect("Expected return value to be constant") + .to_u128(); + assert_eq!(return_value, expected_return); + } + other => panic!("Unexpected terminator {other:?}"), + } + } +} From 4d82337b1fab9847773344370fcbeda1405e7ec6 Mon Sep 17 00:00:00 2001 From: Jake Fecher Date: Thu, 18 May 2023 12:11:24 -0500 Subject: [PATCH 2/5] Amend comment --- .../noirc_evaluator/src/ssa_refactor/ir/basic_block.rs | 10 +++++++--- crates/noirc_evaluator/src/ssa_refactor/ir/dfg.rs | 4 +--- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/crates/noirc_evaluator/src/ssa_refactor/ir/basic_block.rs b/crates/noirc_evaluator/src/ssa_refactor/ir/basic_block.rs index a41563fe21c..3360122b2d2 100644 --- a/crates/noirc_evaluator/src/ssa_refactor/ir/basic_block.rs +++ b/crates/noirc_evaluator/src/ssa_refactor/ir/basic_block.rs @@ -110,9 +110,13 @@ impl BasicBlock { self.instructions.remove(index); } - /// Take ownership of this block's terminator, replacing it with an empty return terminator. - /// It is expected that this function is only called on blocks that are no longer reachable - /// as an optimization to copy over a terminator without cloning. + /// Take ownership of this block's terminator, replacing it with an empty return terminator + /// so that no clone is needed. + /// + /// It is expected that this function is used as an optimization on blocks that are no longer + /// reachable or will have their terminator overwritten afterwards. Using this on a reachable + /// block without setting the terminator afterward will result in the empty return terminator + /// being kept, which is likely unwanted. pub(crate) fn take_terminator(&mut self) -> TerminatorInstruction { let terminator = self.terminator.as_mut().expect("Expected block to have a terminator"); std::mem::replace(terminator, TerminatorInstruction::Return { return_values: Vec::new() }) diff --git a/crates/noirc_evaluator/src/ssa_refactor/ir/dfg.rs b/crates/noirc_evaluator/src/ssa_refactor/ir/dfg.rs index 217472ad6d1..b6b6496d861 100644 --- a/crates/noirc_evaluator/src/ssa_refactor/ir/dfg.rs +++ b/crates/noirc_evaluator/src/ssa_refactor/ir/dfg.rs @@ -167,9 +167,7 @@ impl DataFlowGraph { // Get all of the types that this instruction produces // and append them as results. - let typs = self.instruction_result_types(instruction_id, ctrl_typevars); - - for typ in typs { + for typ in self.instruction_result_types(instruction_id, ctrl_typevars) { self.append_result(instruction_id, typ); } } From 994716116b11ff5d22c014a9ec1241b3c5128cb6 Mon Sep 17 00:00:00 2001 From: Jake Fecher Date: Thu, 18 May 2023 13:33:52 -0500 Subject: [PATCH 3/5] Remove block arguments for blocks with 1 predecessor --- .../src/ssa_refactor/ir/basic_block.rs | 7 +++ .../src/ssa_refactor/opt/simplify_cfg.rs | 56 ++++++++++++------- 2 files changed, 44 insertions(+), 19 deletions(-) diff --git a/crates/noirc_evaluator/src/ssa_refactor/ir/basic_block.rs b/crates/noirc_evaluator/src/ssa_refactor/ir/basic_block.rs index 3360122b2d2..6184e2a21b0 100644 --- a/crates/noirc_evaluator/src/ssa_refactor/ir/basic_block.rs +++ b/crates/noirc_evaluator/src/ssa_refactor/ir/basic_block.rs @@ -121,4 +121,11 @@ impl BasicBlock { let terminator = self.terminator.as_mut().expect("Expected block to have a terminator"); std::mem::replace(terminator, TerminatorInstruction::Return { return_values: Vec::new() }) } + + /// Returns a mutable reference to the terminator of this block. + /// + /// Once this block has finished construction, this is expected to always be Some. + pub(crate) fn unwrap_terminator_mut(&mut self) -> &mut TerminatorInstruction { + self.terminator.as_mut().expect("Expected block to have terminator instruction") + } } diff --git a/crates/noirc_evaluator/src/ssa_refactor/opt/simplify_cfg.rs b/crates/noirc_evaluator/src/ssa_refactor/opt/simplify_cfg.rs index 1d24edef146..570715f795a 100644 --- a/crates/noirc_evaluator/src/ssa_refactor/opt/simplify_cfg.rs +++ b/crates/noirc_evaluator/src/ssa_refactor/opt/simplify_cfg.rs @@ -6,7 +6,7 @@ //! 3. Removes any block arguments for blocks with only a single predecessor. //! 4. Removes any blocks which have no instructions other than a single terminating jmp. //! -//! Currently, only 2 is implemented. +//! Currently, only 2 and 3 are implemented. use std::collections::HashSet; use crate::ssa_refactor::{ @@ -24,7 +24,7 @@ impl Ssa { /// 3. Removing any block arguments for blocks with only a single predecessor. /// 4. Removing any blocks which have no instructions other than a single terminating jmp. /// - /// Currently, only 2 is implemented. + /// Currently, only 2 and 3 are implemented. pub(crate) fn simplify_cfg(mut self) -> Self { for function in self.functions.values_mut() { simplify_function(function); @@ -51,6 +51,11 @@ fn simplify_function(function: &mut Function) { let predecessor = predecessors.next().expect("Already checked length of predecessors"); drop(predecessors); + // If the block has only 1 predecessor, we can safely remove its block parameters + remove_block_parameters(function, block, predecessor); + + // Note: this function relies on `remove_block_parameters` being called first. + // Otherwise the inlined block will refer to parameters that no longer exist. if try_inline_into_predecessor(function, &mut cfg, block, predecessor) { continue; } @@ -58,6 +63,36 @@ fn simplify_function(function: &mut Function) { } } +/// If the given block has block parameters, replace them with the jump arguments from the predecessor. +/// +/// Currently, if this function is needed, `try_inline_into_predecessor` will also always apply, +/// although in the future it is possible for only this function to apply if jmpif instructions +/// with block arguments are ever added. +fn remove_block_parameters( + function: &mut Function, + block: BasicBlockId, + predecessor: BasicBlockId, +) { + let block = &mut function.dfg[block]; + + if !block.parameters().is_empty() { + let block_params = block.take_parameters(); + + let jump_args = match function.dfg[predecessor].unwrap_terminator_mut() { + TerminatorInstruction::Jmp { arguments, .. } => std::mem::take(arguments), + TerminatorInstruction::JmpIf { .. } => unreachable!("If jmpif instructions are modified to support block arguments in the future, this match will need to be updated"), + _ => unreachable!( + "Predecessor was already validated to have only a single jmp destination" + ), + }; + + assert_eq!(block_params.len(), jump_args.len()); + for (param, arg) in block_params.iter().zip(jump_args) { + function.dfg.set_value_from_id(*param, arg); + } + } +} + /// Try to inline a block into its predecessor, returning true if successful. /// /// This will only occur if the predecessor's only successor is the given block. @@ -81,26 +116,9 @@ fn try_inline_into_predecessor( let predecessor = &mut function.dfg[predecessor_id]; predecessor.instructions_mut().append(&mut instructions); - let jump_args = match predecessor.take_terminator() { - TerminatorInstruction::Jmp { arguments, .. } => arguments, - _ => unreachable!( - "Predecessor was already validated to have only a single jmp destination" - ), - }; - predecessor.set_terminator(terminator); - - // Finally, replace any block parameters with their values, if there were any. - let block_params = function.dfg[block_id].take_parameters(); - assert_eq!(block_params.len(), jump_args.len()); - - for (param, arg) in block_params.iter().zip(jump_args) { - function.dfg.set_value_from_id(*param, arg); - } - cfg.recompute_block(function, block_id); cfg.recompute_block(function, predecessor_id); - true } else { false From eb2f65ecade448ca3e689b6e590f24b7e600ee77 Mon Sep 17 00:00:00 2001 From: Jake Fecher Date: Thu, 18 May 2023 13:53:43 -0500 Subject: [PATCH 4/5] Optimize out constant if conditions --- .../src/ssa_refactor/opt/simplify_cfg.rs | 87 ++++++++++++++++++- 1 file changed, 86 insertions(+), 1 deletion(-) diff --git a/crates/noirc_evaluator/src/ssa_refactor/opt/simplify_cfg.rs b/crates/noirc_evaluator/src/ssa_refactor/opt/simplify_cfg.rs index 570715f795a..10c542e1217 100644 --- a/crates/noirc_evaluator/src/ssa_refactor/opt/simplify_cfg.rs +++ b/crates/noirc_evaluator/src/ssa_refactor/opt/simplify_cfg.rs @@ -45,6 +45,10 @@ fn simplify_function(function: &mut Function) { stack.extend(function.dfg[block].successors().filter(|block| !visited.contains(block))); } + // This call is before try_inline_into_predecessor so that if it succeeds in changing a + // jmpif into a jmp, the block may then be inlined entirely into its predecessor in try_inline_into_predecessor. + check_for_constant_jmpif(function, block, &mut cfg); + let mut predecessors = cfg.predecessors(block); if predecessors.len() == 1 { @@ -63,6 +67,26 @@ fn simplify_function(function: &mut Function) { } } +/// Optimize a jmpif into a jmp if the condition is known +fn check_for_constant_jmpif( + function: &mut Function, + block: BasicBlockId, + cfg: &mut ControlFlowGraph, +) { + if let Some(TerminatorInstruction::JmpIf { condition, then_destination, else_destination }) = + function.dfg[block].terminator() + { + if let Some(constant) = function.dfg.get_numeric_constant(*condition) { + let destination = + if constant.is_zero() { *else_destination } else { *then_destination }; + + let jmp = TerminatorInstruction::Jmp { destination, arguments: Vec::new() }; + function.dfg[block].set_terminator(jmp); + cfg.recompute_block(function, block); + } + } +} + /// If the given block has block parameters, replace them with the jump arguments from the predecessor. /// /// Currently, if this function is needed, `try_inline_into_predecessor` will also always apply, @@ -128,7 +152,11 @@ fn try_inline_into_predecessor( #[cfg(test)] mod test { use crate::ssa_refactor::{ - ir::{instruction::TerminatorInstruction, map::Id, types::Type}, + ir::{ + instruction::{BinaryOp, TerminatorInstruction}, + map::Id, + types::Type, + }, ssa_builder::FunctionBuilder, }; @@ -187,4 +215,61 @@ mod test { other => panic!("Unexpected terminator {other:?}"), } } + + #[test] + fn remove_known_jmpif() { + // fn main { + // b0(v0: u1): + // v1 = eq v0, v0 + // jmpif v1, then: b1, else: b2 + // b1(): + // return Field 1 + // b2(): + // return Field 2 + // } + let main_id = Id::test_new(0); + let mut builder = FunctionBuilder::new("main".into(), main_id); + let v0 = builder.add_parameter(Type::bool()); + + let b1 = builder.insert_block(); + let b2 = builder.insert_block(); + + let one = builder.field_constant(1u128); + let two = builder.field_constant(2u128); + + let v1 = builder.insert_binary(v0, BinaryOp::Eq, v0); + builder.terminate_with_jmpif(v1, b1, b2); + + builder.switch_to_block(b1); + builder.terminate_with_return(vec![one]); + + builder.switch_to_block(b2); + builder.terminate_with_return(vec![two]); + + let ssa = builder.finish(); + assert_eq!(ssa.main().reachable_blocks().len(), 3); + + // Expected output: + // fn main { + // b0(): + // return Field 1 + // } + let ssa = ssa.simplify_cfg(); + let main = ssa.main(); + println!("{}", main); + assert_eq!(main.reachable_blocks().len(), 1); + + match main.dfg[main.entry_block()].terminator() { + Some(TerminatorInstruction::Return { return_values }) => { + assert_eq!(return_values.len(), 1); + let return_value = main + .dfg + .get_numeric_constant(return_values[0]) + .expect("Expected return value to be constant") + .to_u128(); + assert_eq!(return_value, 1u128); + } + other => panic!("Unexpected terminator {other:?}"), + } + } } From c39be5e87c83cee074696c2033d599eeac9b1e55 Mon Sep 17 00:00:00 2001 From: Jake Fecher Date: Thu, 18 May 2023 13:58:20 -0500 Subject: [PATCH 5/5] Edit comment --- .../noirc_evaluator/src/ssa_refactor/opt/simplify_cfg.rs | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/crates/noirc_evaluator/src/ssa_refactor/opt/simplify_cfg.rs b/crates/noirc_evaluator/src/ssa_refactor/opt/simplify_cfg.rs index 10c542e1217..b2bf406e0cd 100644 --- a/crates/noirc_evaluator/src/ssa_refactor/opt/simplify_cfg.rs +++ b/crates/noirc_evaluator/src/ssa_refactor/opt/simplify_cfg.rs @@ -5,8 +5,10 @@ //! 2. Inlines a block into its sole predecessor if that predecessor only has one successor. //! 3. Removes any block arguments for blocks with only a single predecessor. //! 4. Removes any blocks which have no instructions other than a single terminating jmp. +//! 5. Replaces any jmpifs with constant conditions with jmps. If this causes the block to have +//! only 1 successor then (2) also will be applied. //! -//! Currently, only 2 and 3 are implemented. +//! Currently, 1 and 4 are unimplemented. use std::collections::HashSet; use crate::ssa_refactor::{ @@ -23,8 +25,10 @@ impl Ssa { /// 2. Inlining a block into its sole predecessor if that predecessor only has one successor. /// 3. Removing any block arguments for blocks with only a single predecessor. /// 4. Removing any blocks which have no instructions other than a single terminating jmp. + /// 5. Replacing any jmpifs with constant conditions with jmps. If this causes the block to have + /// only 1 successor then (2) also will be applied. /// - /// Currently, only 2 and 3 are implemented. + /// Currently, 1 and 4 are unimplemented. pub(crate) fn simplify_cfg(mut self) -> Self { for function in self.functions.values_mut() { simplify_function(function);