From a69e2c09a93dafb790a39691154e9f67aba4fd35 Mon Sep 17 00:00:00 2001 From: Savio Sou <72797635+Savio-Sou@users.noreply.github.com> Date: Tue, 23 May 2023 00:23:33 +0800 Subject: [PATCH 1/9] chore(github): Update GitHub Feature Request Template (#1369) --- .github/ISSUE_TEMPLATE/feature_request.yml | 44 ++++++++++++++-------- 1 file changed, 29 insertions(+), 15 deletions(-) diff --git a/.github/ISSUE_TEMPLATE/feature_request.yml b/.github/ISSUE_TEMPLATE/feature_request.yml index 3eeef2f4a43..ae48534c009 100644 --- a/.github/ISSUE_TEMPLATE/feature_request.yml +++ b/.github/ISSUE_TEMPLATE/feature_request.yml @@ -1,40 +1,54 @@ -name: Feature request -description: Suggest an idea for this project. +name: Feature Request +description: Suggest an idea. labels: [enhancement] body: - type: markdown attributes: value: | ## Description - Thanks for taking the time to fill out this feature request to help us improve Noir! + Thanks for taking the time to create the Issue, and welcome to the Noirot family! - type: textarea id: problem attributes: label: Problem - description: Describe the problem your suggestion sets out to solve. + description: Describe what you feel lacking. Supply code / step-by-step examples if applicable. validations: required: true - type: textarea id: solution attributes: - label: Proposed solution - description: Describe your proposed solution. + label: Happy Case + description: Describe how you think it should work. Supply pseudocode / step-by-step examples if applicable. validations: required: true - type: textarea id: alternatives attributes: - label: Alternatives considered - description: Describe any alternative solutions you have considered. + label: Alternatives Considered + description: Describe less-happy cases you have considered, if any. - type: textarea id: additional attributes: - label: Additional context - description: Please provide any additional context that may be applicable. - - type: checkboxes - id: checklist + label: Additional Context + description: Supplement further information if applicable. + - type: markdown + attributes: + value: | + ## Pull Request + - type: dropdown + id: pr-preference attributes: - label: Submission Checklist - description: For core contributors. + label: Would you like to submit a PR for this Issue? + description: Fellow contributors are happy to provide support where applicable. + multiple: false options: - - label: Once I hit submit, I will assign this issue to the Project Board with the appropriate tags. + - "No" + - "Maybe" + - "Yes" + validations: + required: true + - type: textarea + id: pr-support + attributes: + label: Support Needs + description: Support from other contributors you are looking for to create a PR for this Issue. From c4d8858b933e7cede3cc5efd3e16d038d822fcfe Mon Sep 17 00:00:00 2001 From: jfecher Date: Mon, 22 May 2023 12:58:53 -0400 Subject: [PATCH 2/9] chore(ssa refactor): Add pass to simplify the control flow graph (#1373) * Add simplify_cfg pass * Amend comment * Remove block arguments for blocks with 1 predecessor * Add comment --- .../src/ssa_refactor/ir/basic_block.rs | 24 +++ .../src/ssa_refactor/ir/dfg.rs | 28 +-- .../src/ssa_refactor/ir/function.rs | 19 ++ .../src/ssa_refactor/opt/mod.rs | 1 + .../src/ssa_refactor/opt/simplify_cfg.rs | 192 ++++++++++++++++++ 5 files changed, 252 insertions(+), 12 deletions(-) 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 c68efa44127..72954f2d0bd 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(). @@ -102,4 +107,23 @@ impl BasicBlock { }); self.instructions.remove(index); } + + /// 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() }) + } + + /// 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/ir/dfg.rs b/crates/noirc_evaluator/src/ssa_refactor/ir/dfg.rs index f4f6004d41c..eaec4920790 100644 --- a/crates/noirc_evaluator/src/ssa_refactor/ir/dfg.rs +++ b/crates/noirc_evaluator/src/ssa_refactor/ir/dfg.rs @@ -123,6 +123,21 @@ impl DataFlowGraph { self.values.insert(value) } + /// Replaces the value specified by the given ValueId with a new Value. + /// + /// This is the preferred method to call for optimizations simplifying + /// values since other instructions referring to the same ValueId need + /// not be modified to refer to a new ValueId. + pub(crate) fn set_value(&mut self, value_id: ValueId, new_value: Value) { + self.values[value_id] = new_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 { @@ -161,9 +176,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); } } @@ -269,15 +282,6 @@ impl DataFlowGraph { ) { self.blocks[block].set_terminator(terminator); } - - /// Replaces the value specified by the given ValueId with a new Value. - /// - /// This is the preferred method to call for optimizations simplifying - /// values since other instructions referring to the same ValueId need - /// not be modified to refer to a new ValueId. - pub(crate) fn set_value(&mut self, value_id: ValueId, new_value: Value) { - self.values[value_id] = new_value; - } } impl std::ops::Index for DataFlowGraph { 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 f4b2a6ca683..d9d31f8f356 100644 --- a/crates/noirc_evaluator/src/ssa_refactor/opt/mod.rs +++ b/crates/noirc_evaluator/src/ssa_refactor/opt/mod.rs @@ -5,3 +5,4 @@ //! Generally, these passes are also expected to minimize the final amount of instructions. mod inlining; mod mem2reg; +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..7c91b5f0fe5 --- /dev/null +++ b/crates/noirc_evaluator/src/ssa_refactor/opt/simplify_cfg.rs @@ -0,0 +1,192 @@ +//! 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 and 3 are 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 and 3 are 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 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 successful, `block` will be empty and unreachable after this call, so any + // optimizations performed after this point on the same block should check if + // the inlining here was successful before continuing. + try_inline_into_predecessor(function, &mut cfg, block, predecessor); + } + } +} + +/// 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. +/// 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); + + predecessor.set_terminator(terminator); + 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 a57c5749cefd5a2999e3b4bba9110097f6d9e915 Mon Sep 17 00:00:00 2001 From: jfecher Date: Mon, 22 May 2023 13:59:32 -0400 Subject: [PATCH 3/9] chore(ssa refactor): Update mem2reg pass to work with multiple functions and blocks (#1375) * Change mem2reg pass to work with multiple functions and blocks * PR feedback: update method name * Fix method name * Fix comment --- .../src/ssa_refactor/opt/mem2reg.rs | 294 +++++++++++------- 1 file changed, 181 insertions(+), 113 deletions(-) diff --git a/crates/noirc_evaluator/src/ssa_refactor/opt/mem2reg.rs b/crates/noirc_evaluator/src/ssa_refactor/opt/mem2reg.rs index a020230fdab..15bfdc1bce9 100644 --- a/crates/noirc_evaluator/src/ssa_refactor/opt/mem2reg.rs +++ b/crates/noirc_evaluator/src/ssa_refactor/opt/mem2reg.rs @@ -19,20 +19,28 @@ impl Ssa { /// scope, and attempts to remove store that are subsequently redundant, as long as they are /// not stores on memory that will be passed into a function call. /// - /// This pass assumes that the whole program has been inlined into a single block, such that - /// we can be sure that store instructions cannot have side effects outside of this block - /// (apart from intrinsic function calls). - /// /// This pass also assumes that constant folding has been run, such that all addresses given /// as input to store/load instructions are represented as one of: /// - a value that directly resolves to an allocate instruction /// - a value that directly resolves to a binary add instruction which has a allocate /// instruction and a numeric constant as its operands - pub(crate) fn mem2reg_final(mut self) -> Ssa { - let func = self.main_mut(); - assert_eq!(func.dfg.basic_blocks_iter().count(), 1); - let block_id = func.entry_block(); - PerBlockContext::new(&mut func.dfg, block_id).eliminate_store_load(); + pub(crate) fn mem2reg(mut self) -> Ssa { + let mut first_context = None; + + for function in self.functions.values_mut() { + for block in function.reachable_blocks() { + let mut context = PerBlockContext::new(block); + context.eliminate_known_loads(&mut function.dfg); + first_context = Some(context); + } + } + + // If there is only one block in total, remove any unused stores as well since we + // know there is no other block they can impact. + if self.functions.len() == 1 && self.main().dfg.basic_blocks_iter().len() == 1 { + first_context.unwrap().remove_unused_stores(&mut self.main_mut().dfg); + } + self } } @@ -52,59 +60,58 @@ impl Address { } } -struct PerBlockContext<'dfg> { - dfg: &'dfg mut DataFlowGraph, +struct PerBlockContext { block_id: BasicBlockId, + last_stores: BTreeMap, + store_ids: Vec, + failed_substitutes: BTreeSet
, + alloc_ids_in_calls: BTreeSet, } -impl<'dfg> PerBlockContext<'dfg> { - fn new(dfg: &'dfg mut DataFlowGraph, block_id: BasicBlockId) -> Self { - PerBlockContext { dfg, block_id } +impl PerBlockContext { + fn new(block_id: BasicBlockId) -> Self { + PerBlockContext { + block_id, + last_stores: BTreeMap::new(), + store_ids: Vec::new(), + failed_substitutes: BTreeSet::new(), + alloc_ids_in_calls: BTreeSet::new(), + } } - // Attempts to remove redundant load & store instructions for constant addresses. Returns the - // count of remaining store instructions. - // - // This method assumes the entire program is now represented in a single block (minus any - // intrinsic function calls). Therefore we needn't be concerned with store instructions having - // an effect beyond the scope of this block. - fn eliminate_store_load(&mut self) -> u32 { - let mut store_count: u32 = 0; - let mut last_stores: BTreeMap = BTreeMap::new(); - let mut loads_to_substitute: Vec<(InstructionId, Value)> = Vec::new(); - let mut store_ids: Vec = Vec::new(); - let mut failed_substitutes: BTreeSet
= BTreeSet::new(); - let mut alloc_ids_in_calls: BTreeSet = BTreeSet::new(); - - let block = &self.dfg[self.block_id]; + // Attempts to remove load instructions for which the result is already known from previous + // store instructions to the same address in the same block. + fn eliminate_known_loads(&mut self, dfg: &mut DataFlowGraph) { + let mut loads_to_substitute = Vec::new(); + let block = &dfg[self.block_id]; + for instruction_id in block.instructions() { - match &self.dfg[*instruction_id] { + match &dfg[*instruction_id] { Instruction::Store { address, value } => { - store_count += 1; - if let Some(address) = self.try_const_address(*address) { + if let Some(address) = self.try_const_address(*address, dfg) { // We can only track the address if it is a constant offset from an // allocation. A previous constant folding pass should make such addresses // possible to identify. - last_stores.insert(address, *value); + self.last_stores.insert(address, *value); } // TODO: Consider if it's worth falling back to storing addresses by their // value id such we can shallowly check for dynamic address reuse. - store_ids.push(*instruction_id); + self.store_ids.push(*instruction_id); } Instruction::Load { address } => { - if let Some(address) = self.try_const_address(*address) { - if let Some(last_value) = last_stores.get(&address) { - let last_value = self.dfg[*last_value]; + if let Some(address) = self.try_const_address(*address, dfg) { + if let Some(last_value) = self.last_stores.get(&address) { + let last_value = dfg[*last_value]; loads_to_substitute.push((*instruction_id, last_value)); } else { - failed_substitutes.insert(address); + self.failed_substitutes.insert(address); } } } Instruction::Call { arguments, .. } => { for arg in arguments { - if let Some(address) = self.try_const_address(*arg) { - alloc_ids_in_calls.insert(address.alloc_id()); + if let Some(address) = self.try_const_address(*arg, dfg) { + self.alloc_ids_in_calls.insert(address.alloc_id()); } } } @@ -116,74 +123,78 @@ impl<'dfg> PerBlockContext<'dfg> { // Substitute load result values for (instruction_id, new_value) in &loads_to_substitute { - let result_value = *self - .dfg + let result_value = *dfg .instruction_results(*instruction_id) .first() .expect("ICE: Load instructions should have single result"); - self.dfg.set_value(result_value, *new_value); + dfg.set_value(result_value, *new_value); } // Delete load instructions // TODO: should we let DCE do this instead? - let block = &mut self.dfg[self.block_id]; + let block = &mut dfg[self.block_id]; for (instruction_id, _) in loads_to_substitute { block.remove_instruction(instruction_id); } + } + fn remove_unused_stores(self, dfg: &mut DataFlowGraph) { // Scan for unused stores let mut stores_to_remove: Vec = Vec::new(); - for instruction_id in store_ids { - let address = match &self.dfg[instruction_id] { + for instruction_id in &self.store_ids { + let address = match &dfg[*instruction_id] { Instruction::Store { address, .. } => *address, _ => unreachable!("store_ids should contain only store instructions"), }; - if let Some(address) = self.try_const_address(address) { - if !failed_substitutes.contains(&address) - && !alloc_ids_in_calls.contains(&address.alloc_id()) + + if let Some(address) = self.try_const_address(address, dfg) { + if !self.failed_substitutes.contains(&address) + && !self.alloc_ids_in_calls.contains(&address.alloc_id()) { - stores_to_remove.push(instruction_id); + stores_to_remove.push(*instruction_id); } } } // Delete unused stores - let block = &mut self.dfg[self.block_id]; + let block = &mut dfg[self.block_id]; for instruction_id in stores_to_remove { - store_count -= 1; block.remove_instruction(instruction_id); } - - store_count } // Attempts to normalize the given value into a const address - fn try_const_address(&self, value_id: ValueId) -> Option
{ - let value = &self.dfg[value_id]; + fn try_const_address(&self, value_id: ValueId, dfg: &DataFlowGraph) -> Option
{ + let value = &dfg[value_id]; let instruction_id = match value { Value::Instruction { instruction, .. } => *instruction, _ => return None, }; - let instruction = &self.dfg[instruction_id]; + let instruction = &dfg[instruction_id]; match instruction { Instruction::Allocate { .. } => Some(Address::Zeroth(instruction_id)), Instruction::Binary(binary) => { if binary.operator != BinaryOp::Add { return None; } - let lhs = &self.dfg[binary.lhs]; - let rhs = &self.dfg[binary.rhs]; - self.try_const_address_offset(lhs, rhs) - .or_else(|| self.try_const_address_offset(rhs, lhs)) + let lhs = &dfg[binary.lhs]; + let rhs = &dfg[binary.rhs]; + self.try_const_address_offset(lhs, rhs, dfg) + .or_else(|| self.try_const_address_offset(rhs, lhs, dfg)) } _ => None, } } // Tries val1 as an allocation instruction id and val2 as a constant offset - fn try_const_address_offset(&self, val1: &Value, val2: &Value) -> Option
{ + fn try_const_address_offset( + &self, + val1: &Value, + val2: &Value, + dfg: &DataFlowGraph, + ) -> Option
{ let alloc_id = match val1 { - Value::Instruction { instruction, .. } => match &self.dfg[*instruction] { + Value::Instruction { instruction, .. } => match &dfg[*instruction] { Instruction::Allocate { .. } => *instruction, _ => return None, }, @@ -203,6 +214,8 @@ mod tests { use crate::ssa_refactor::{ ir::{ + basic_block::BasicBlockId, + dfg::DataFlowGraph, instruction::{BinaryOp, Instruction, Intrinsic, TerminatorInstruction}, map::Id, types::Type, @@ -210,18 +223,17 @@ mod tests { ssa_builder::FunctionBuilder, }; - use super::PerBlockContext; - #[test] fn test_simple() { - // func() { - // block0(): + // fn func() { + // b0(): // v0 = alloc 2 // v1 = add v0, Field 1 // store v1, Field 1 // v2 = add v0, Field 1 // v3 = load v1 // return v3 + // } let func_id = Id::test_new(0); let mut builder = FunctionBuilder::new("func".into(), func_id); @@ -233,32 +245,15 @@ mod tests { let v3 = builder.insert_load(v0, const_one, Type::field()); builder.terminate_with_return(vec![v3]); - let mut ssa = builder.finish(); + let ssa = builder.finish().mem2reg(); - let mut func = ssa.functions.remove(&func_id).unwrap(); + let func = ssa.main(); let block_id = func.entry_block(); - let mut mem2reg_context = PerBlockContext::new(&mut func.dfg, block_id); - let remaining_stores = mem2reg_context.eliminate_store_load(); - - assert_eq!(remaining_stores, 0); + assert_eq!(count_loads(block_id, &func.dfg), 0); + assert_eq!(count_stores(block_id, &func.dfg), 0); - let block = &func.dfg[block_id]; - let load_count = block - .instructions() - .iter() - .filter(|instruction_id| matches!(func.dfg[**instruction_id], Instruction::Load { .. })) - .count(); - assert_eq!(load_count, 0); - let store_count = block - .instructions() - .iter() - .filter(|instruction_id| { - matches!(func.dfg[**instruction_id], Instruction::Store { .. }) - }) - .count(); - assert_eq!(store_count, 0); - let ret_val_id = match block.terminator().unwrap() { + let ret_val_id = match func.dfg[block_id].terminator().unwrap() { TerminatorInstruction::Return { return_values } => return_values.first().unwrap(), _ => unreachable!(), }; @@ -267,8 +262,8 @@ mod tests { #[test] fn test_simple_with_call() { - // func() { - // block0(): + // fn func { + // b0(): // v0 = alloc 2 // v1 = add v0, Field 1 // store v1, Field 1 @@ -276,6 +271,7 @@ mod tests { // v3 = load v1 // v4 = call f0, v0 // return v3 + // } let func_id = Id::test_new(0); let mut builder = FunctionBuilder::new("func".into(), func_id); @@ -289,38 +285,110 @@ mod tests { builder.insert_call(f0, vec![v0], vec![Type::Unit]); builder.terminate_with_return(vec![v3]); - let mut ssa = builder.finish(); + let ssa = builder.finish().mem2reg(); - let mut func = ssa.functions.remove(&func_id).unwrap(); + let func = ssa.main(); let block_id = func.entry_block(); - let mut mem2reg_context = PerBlockContext::new(&mut func.dfg, block_id); - let remaining_stores = mem2reg_context.eliminate_store_load(); + assert_eq!(count_loads(block_id, &func.dfg), 0); + assert_eq!(count_stores(block_id, &func.dfg), 1); - assert_eq!( - remaining_stores, 1, - "Store cannot be removed as it affects intrinsic function call" - ); + let ret_val_id = match func.dfg[block_id].terminator().unwrap() { + TerminatorInstruction::Return { return_values } => return_values.first().unwrap(), + _ => unreachable!(), + }; + assert_eq!(func.dfg[*ret_val_id], func.dfg[const_one]); + } - let block = &func.dfg[block_id]; - let load_count = block + fn count_stores(block: BasicBlockId, dfg: &DataFlowGraph) -> usize { + dfg[block] .instructions() .iter() - .filter(|instruction_id| matches!(func.dfg[**instruction_id], Instruction::Load { .. })) - .count(); - assert_eq!(load_count, 0); - let store_count = block + .filter(|instruction_id| matches!(dfg[**instruction_id], Instruction::Store { .. })) + .count() + } + + fn count_loads(block: BasicBlockId, dfg: &DataFlowGraph) -> usize { + dfg[block] .instructions() .iter() - .filter(|instruction_id| { - matches!(func.dfg[**instruction_id], Instruction::Store { .. }) - }) - .count(); - assert_eq!(store_count, 1); - let ret_val_id = match block.terminator().unwrap() { - TerminatorInstruction::Return { return_values } => return_values.first().unwrap(), + .filter(|instruction_id| matches!(dfg[**instruction_id], Instruction::Load { .. })) + .count() + } + + // Test that loads across multiple blocks are not removed + #[test] + fn multiple_blocks() { + // fn main { + // b0(): + // v0 = alloc 1 + // store v0, Field 5 + // v1 = load v0 + // jmp b1(v1): + // b1(v2: Field): + // v3 = load v0 + // store v0, Field 6 + // v4 = load v0 + // return v2, v3, v4 + // } + let main_id = Id::test_new(0); + let mut builder = FunctionBuilder::new("main".into(), main_id); + + let v0 = builder.insert_allocate(1); + + let five = builder.field_constant(5u128); + builder.insert_store(v0, five); + + let zero = builder.field_constant(0u128); + let v1 = builder.insert_load(v0, zero, Type::field()); + let b1 = builder.insert_block(); + builder.terminate_with_jmp(b1, vec![v1]); + + builder.switch_to_block(b1); + let v2 = builder.add_block_parameter(b1, Type::field()); + let v3 = builder.insert_load(v0, zero, Type::field()); + + let six = builder.field_constant(6u128); + builder.insert_store(v0, six); + let v4 = builder.insert_load(v0, zero, Type::field()); + + builder.terminate_with_return(vec![v2, v3, v4]); + + let ssa = builder.finish(); + assert_eq!(ssa.main().reachable_blocks().len(), 2); + + // Expected result: + // fn main { + // b0(): + // v0 = alloc 1 + // store v0, Field 5 + // jmp b1(Field 5): // Optimized to constant 5 + // b1(v2: Field): + // v3 = load v0 // kept in program + // store v0, Field 6 + // return v2, v3, Field 6 // Optimized to constant 6 + // } + let ssa = ssa.mem2reg(); + let main = ssa.main(); + assert_eq!(main.reachable_blocks().len(), 2); + + // Only the load from the entry block should be removed + assert_eq!(count_loads(main.entry_block(), &main.dfg), 0); + assert_eq!(count_loads(b1, &main.dfg), 1); + + // All stores should be kept + assert_eq!(count_stores(main.entry_block(), &main.dfg), 1); + assert_eq!(count_stores(b1, &main.dfg), 1); + + // The jmp to b1 should also be a constant 5 now + match main.dfg[main.entry_block()].terminator() { + Some(TerminatorInstruction::Jmp { arguments, .. }) => { + assert_eq!(arguments.len(), 1); + let argument = + main.dfg.get_numeric_constant(arguments[0]).expect("Expected constant value"); + assert_eq!(argument.to_u128(), 5); + } _ => unreachable!(), }; - assert_eq!(func.dfg[*ret_val_id], func.dfg[const_one]); } } From 72a2adbf9eff6145c440ec5cd6a25a0b56431901 Mon Sep 17 00:00:00 2001 From: Savio Sou <72797635+Savio-Sou@users.noreply.github.com> Date: Tue, 23 May 2023 15:03:39 +0800 Subject: [PATCH 4/9] chore(github): Update GitHub Bug Report Template (#1368) * chore(github): Update GitHub Bug Report Template * Update TypeScript packages section with placeholder text --------- Co-authored-by: Globallager <72797635+Globallager@users.noreply.github.com> --- .github/ISSUE_TEMPLATE/bug_report.yml | 71 ++++++++++++++------------- 1 file changed, 36 insertions(+), 35 deletions(-) diff --git a/.github/ISSUE_TEMPLATE/bug_report.yml b/.github/ISSUE_TEMPLATE/bug_report.yml index 873d209ee2a..112da342e10 100644 --- a/.github/ISSUE_TEMPLATE/bug_report.yml +++ b/.github/ISSUE_TEMPLATE/bug_report.yml @@ -1,12 +1,12 @@ name: Bug Report -description: File a bug report to help us improve. +description: Report an unexpected behavior. labels: ["bug"] body: - type: markdown attributes: value: | - ## Description - Thanks for taking the time to fill out this bug report to help us improve Noir! + # Description + Thanks for taking the time to create the Issue, and welcome to the Noirot family! - type: textarea id: aim attributes: @@ -17,7 +17,7 @@ body: - type: textarea id: expected attributes: - label: Expected behavior + label: Expected Behavior description: Describe what you expected to happen. validations: required: true @@ -31,7 +31,7 @@ body: - type: textarea id: reproduction attributes: - label: To reproduce + label: To Reproduce description: Describe the steps to reproduce the behavior. value: | 1. @@ -41,16 +41,16 @@ body: - type: markdown attributes: value: | - ## Environment + # Environment Specify your versions of Noir releases used. - type: markdown attributes: value: | - ### Using Nargo? + ## Using Nargo? - type: dropdown id: nargo-install attributes: - label: Installation method + label: Installation Method description: How did you install Nargo? multiple: false options: @@ -59,40 +59,41 @@ body: - type: input id: nargo-version attributes: - label: Nargo version + label: Nargo Version description: What is the output of the `nargo --version` command? - placeholder: "nargo 0.2.0 (git version hash: e927a39dc3d6517f233509b8349dfd9c7f79471d, is dirty: false)" + placeholder: "nargo 0.6.0 (git version hash: 0181813203a9e3e46c6d8c3169ad5d25971d4282, is dirty: false)" - type: markdown attributes: value: | - ### Using TypeScript? - - type: input - id: noir_wasm-version - attributes: - label: "@noir-lang/noir_wasm version" - description: What is the version number? e.g. version in yarn.lock - placeholder: "0.2.0-ca986a4" - - type: input - id: barretenberg-version - attributes: - label: "@noir-lang/barretenberg version" - description: What is the version number? e.g. version in yarn.lock - placeholder: "2.19.0" - - type: input - id: aztec_backend-version + ## Using TypeScript? + Please await for our new set of packages. + You can find our target release timeframe on the [Noir Roadmap](https://github.com/orgs/noir-lang/projects/1/views/16). + - type: markdown attributes: - label: "@noir-lang/aztec_backend version" - description: What is the version number? e.g. version in yarn.lock - placeholder: "0.12.0" + value: | + # Misc - type: textarea id: additional attributes: - label: Additional context - description: Please provide any additional context that may be applicable. - - type: checkboxes - id: checklist + label: Additional Context + description: Supplement further information if applicable. + - type: markdown attributes: - label: Submission Checklist - description: For core contributors. + value: | + # Pull Request + - type: dropdown + id: pr_preference + attributes: + label: Would you like to submit a PR for this Issue? + description: Fellow contributors are happy to provide support where applicable. options: - - label: Once I hit submit, I will assign this issue to the Project Board with the appropriate tags. + - "No" + - "Maybe" + - "Yes" + validations: + required: true + - type: textarea + id: pr_support + attributes: + label: Support Needs + description: Support from other contributors you are looking for to create a PR for this Issue. From 66744f3ea0a5891001225e57b1e2ee960c4d1392 Mon Sep 17 00:00:00 2001 From: Savio Sou <72797635+Savio-Sou@users.noreply.github.com> Date: Tue, 23 May 2023 15:43:45 +0800 Subject: [PATCH 5/9] chore(github): Update GitHub Pull Request Template (#1370) * chore(github): Update GitHub Pull Request Template * Make Example section optional * Revise doc comment wording --------- Co-authored-by: Globallager <72797635+Globallager@users.noreply.github.com> --- .github/pull_request_template.md | 60 +++++++++++++++++++++----------- 1 file changed, 40 insertions(+), 20 deletions(-) diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index 2973a8bfd03..26c48012457 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -1,36 +1,56 @@ -# Related issue(s) +# Description - + + -Resolves # +## Problem\* -# Description + -## Summary of changes +Resolves - +## Summary\* -## Dependency additions / changes + - +This PR sets out to -## Test additions / changes +### Example - + -# Checklist +Before: -- [ ] I have tested the changes locally. -- [ ] I have formatted the changes with [Prettier](https://prettier.io/) and/or `cargo fmt` with default settings. -- [ ] I have [linked](https://docs.github.com/en/issues/tracking-your-work-with-issues/linking-a-pull-request-to-an-issue) this PR to the issue(s) that it resolves. -- [ ] I have reviewed the changes on GitHub, line by line. -- [ ] I have ensured all changes are covered in the description. +``` + +``` + +After: + +``` + +``` + +## Documentation -## Documentation needs - [ ] This PR requires documentation updates when merged. - + -# Additional context + - [ ] I will submit a noir-lang/docs PR. - + + + - [ ] I will request for and support Dev Rel's help in documenting this PR. + + + + +## Additional Context + + + +# PR Checklist\* + +- [ ] I have tested the changes locally. +- [ ] I have formatted the changes with [Prettier](https://prettier.io/) and/or `cargo fmt` on default settings. From e7a0d5c7b3b86587861401533d4e6784d0353404 Mon Sep 17 00:00:00 2001 From: jfecher Date: Tue, 23 May 2023 04:47:43 -0400 Subject: [PATCH 6/9] feat: Allow warnings by default (#1383) Change allow_warnings to deny_warnings flag --- crates/nargo_cli/src/cli/check_cmd.rs | 2 +- crates/nargo_cli/src/cli/mod.rs | 2 +- crates/noirc_driver/src/lib.rs | 10 +++++----- crates/noirc_errors/src/reporter.rs | 16 ++++++++-------- 4 files changed, 15 insertions(+), 15 deletions(-) diff --git a/crates/nargo_cli/src/cli/check_cmd.rs b/crates/nargo_cli/src/cli/check_cmd.rs index c57bb2a1fd9..4453616f32f 100644 --- a/crates/nargo_cli/src/cli/check_cmd.rs +++ b/crates/nargo_cli/src/cli/check_cmd.rs @@ -197,7 +197,7 @@ d2 = ["", "", ""] .join(format!("{TEST_DATA_DIR}/pass_dev_mode")); let backend = crate::backends::ConcreteBackend::default(); - let config = CompileOptions { allow_warnings: true, ..Default::default() }; + let config = CompileOptions { deny_warnings: false, ..Default::default() }; let paths = std::fs::read_dir(pass_dir).unwrap(); for path in paths.flatten() { diff --git a/crates/nargo_cli/src/cli/mod.rs b/crates/nargo_cli/src/cli/mod.rs index c61a970905b..738b6a8511f 100644 --- a/crates/nargo_cli/src/cli/mod.rs +++ b/crates/nargo_cli/src/cli/mod.rs @@ -89,7 +89,7 @@ pub fn prove_and_verify(proof_name: &str, program_dir: &Path) -> bool { let compile_options = CompileOptions { show_ssa: false, print_acir: false, - allow_warnings: false, + deny_warnings: false, show_output: false, experimental_ssa: false, }; diff --git a/crates/noirc_driver/src/lib.rs b/crates/noirc_driver/src/lib.rs index 0b2c6a6ffcf..6adf8d64f0c 100644 --- a/crates/noirc_driver/src/lib.rs +++ b/crates/noirc_driver/src/lib.rs @@ -44,9 +44,9 @@ pub struct CompileOptions { #[arg(short, long)] pub print_acir: bool, - /// Issue a warning for each unused variable instead of an error + /// Treat all warnings as errors #[arg(short, long)] - pub allow_warnings: bool, + pub deny_warnings: bool, /// Display output of `println` statements #[arg(long)] @@ -62,7 +62,7 @@ impl Default for CompileOptions { Self { show_ssa: false, print_acir: false, - allow_warnings: false, + deny_warnings: false, show_output: true, experimental_ssa: false, } @@ -174,7 +174,7 @@ impl Driver { let mut errs = vec![]; CrateDefMap::collect_defs(LOCAL_CRATE, &mut self.context, &mut errs); let error_count = - reporter::report_all(&self.context.file_manager, &errs, options.allow_warnings); + reporter::report_all(&self.context.file_manager, &errs, options.deny_warnings); reporter::finish_report(error_count) } @@ -315,7 +315,7 @@ impl Driver { // Errors will be shown at the call site without a stacktrace let file = err.location.map(|loc| loc.file); let files = &self.context.file_manager; - let error = reporter::report(files, &err.into(), file, options.allow_warnings); + let error = reporter::report(files, &err.into(), file, options.deny_warnings); reporter::finish_report(error as u32)?; Err(ReportedError) } diff --git a/crates/noirc_errors/src/reporter.rs b/crates/noirc_errors/src/reporter.rs index 1ec5fa3dbb5..2c8be90a280 100644 --- a/crates/noirc_errors/src/reporter.rs +++ b/crates/noirc_errors/src/reporter.rs @@ -106,11 +106,11 @@ impl CustomLabel { pub fn report_all( files: &fm::FileManager, diagnostics: &[FileDiagnostic], - allow_warnings: bool, + deny_warnings: bool, ) -> u32 { diagnostics .iter() - .map(|error| report(files, &error.diagnostic, Some(error.file_id), allow_warnings) as u32) + .map(|error| report(files, &error.diagnostic, Some(error.file_id), deny_warnings) as u32) .sum() } @@ -119,24 +119,24 @@ pub fn report( files: &fm::FileManager, custom_diagnostic: &CustomDiagnostic, file: Option, - allow_warnings: bool, + deny_warnings: bool, ) -> bool { let writer = StandardStream::stderr(ColorChoice::Always); let config = codespan_reporting::term::Config::default(); - let diagnostic = convert_diagnostic(custom_diagnostic, file, allow_warnings); + let diagnostic = convert_diagnostic(custom_diagnostic, file, deny_warnings); term::emit(&mut writer.lock(), &config, files.as_simple_files(), &diagnostic).unwrap(); - !allow_warnings || custom_diagnostic.is_error() + deny_warnings || custom_diagnostic.is_error() } fn convert_diagnostic( cd: &CustomDiagnostic, file: Option, - allow_warnings: bool, + deny_warnings: bool, ) -> Diagnostic { - let diagnostic = match (cd.kind, allow_warnings) { - (DiagnosticKind::Warning, true) => Diagnostic::warning(), + let diagnostic = match (cd.kind, deny_warnings) { + (DiagnosticKind::Warning, false) => Diagnostic::warning(), _ => Diagnostic::error(), }; From 6e0facbc8deca2633a0922f350a51e0195ba4f58 Mon Sep 17 00:00:00 2001 From: jfecher Date: Tue, 23 May 2023 05:05:09 -0400 Subject: [PATCH 7/9] chore(ssa refactor): Optimize constant `jmpif`s into `jmp`s (#1374) * Add simplify_cfg pass * Amend comment * Remove block arguments for blocks with 1 predecessor * Optimize out constant if conditions * Edit comment --- .../src/ssa_refactor/opt/simplify_cfg.rs | 95 ++++++++++++++++++- 1 file changed, 92 insertions(+), 3 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 7c91b5f0fe5..419120fc4e8 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); @@ -45,6 +49,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 { @@ -65,6 +73,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, @@ -130,7 +158,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, }; @@ -189,4 +221,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 dfb24ba02e45702bb6bffecb79593fd9acf187f2 Mon Sep 17 00:00:00 2001 From: joss-aztec <94053499+joss-aztec@users.noreply.github.com> Date: Tue, 23 May 2023 10:46:37 +0100 Subject: [PATCH 8/9] chore: Add terms added in ssa refactor to `cspell` (#1385) chore: cspell additions --- cspell.json | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/cspell.json b/cspell.json index f0031d724b2..60b38f6534c 100644 --- a/cspell.json +++ b/cspell.json @@ -13,6 +13,7 @@ "coeff", "combinators", "comptime", + "cranelift", "desugared", "endianness", "forall", @@ -25,6 +26,10 @@ "injective", "interner", "intrinsics", + "jmp", + "jmpif", + "jmpifs", + "jmps", "keccak", "krate", "lvalue", From a21b2519dddbe415e1747a8d4721faea5b71d262 Mon Sep 17 00:00:00 2001 From: jfecher Date: Wed, 24 May 2023 02:54:23 -0700 Subject: [PATCH 9/9] chore(ssa refactor): Add loop unrolling pass (#1364) * Start inlining pass * Get most of pass working * Finish function inlining pass * Add basic test * Address PR comments * Start block inlining * Add basic instruction simplification * Cargo fmt * Add comments * Add context object * Add push_instruction * Fix bug in inlining pass * Reorder loop unrolling pass * Get it working for most loops. Still missing loops with if inside * Rework entire pass from scratch * Finish loop unrolling * Add doc comment * Fix bad merge * Add test * Remove outdated parts of PR * Correctly handle loops with non-const indices * Address PR comments * Fix inlining bug and add a test for loops which fail to unroll * Update simplify_cfg to use new inline_block method * Remove now-unneeded test helper function --- crates/noirc_evaluator/src/ssa_refactor.rs | 25 +- .../src/ssa_refactor/acir_gen/mod.rs | 2 +- .../src/ssa_refactor/ir/basic_block.rs | 57 +- .../src/ssa_refactor/ir/dfg.rs | 59 +- .../src/ssa_refactor/ir/instruction.rs | 116 ++- .../src/ssa_refactor/opt/inlining.rs | 13 +- .../src/ssa_refactor/opt/mod.rs | 1 + .../src/ssa_refactor/opt/simplify_cfg.rs | 23 +- .../src/ssa_refactor/opt/unrolling.rs | 659 ++++++++++++++++++ .../src/ssa_refactor/ssa_builder/mod.rs | 6 +- 10 files changed, 868 insertions(+), 93 deletions(-) create mode 100644 crates/noirc_evaluator/src/ssa_refactor/opt/unrolling.rs diff --git a/crates/noirc_evaluator/src/ssa_refactor.rs b/crates/noirc_evaluator/src/ssa_refactor.rs index 34061227336..8330931a4b4 100644 --- a/crates/noirc_evaluator/src/ssa_refactor.rs +++ b/crates/noirc_evaluator/src/ssa_refactor.rs @@ -7,13 +7,13 @@ //! This module heavily borrows from Cranelift #![allow(dead_code)] -use crate::errors::RuntimeError; +use crate::errors::{RuntimeError, RuntimeErrorKind}; use acvm::{acir::circuit::Circuit, compiler::transformers::IsOpcodeSupported, Language}; use noirc_abi::Abi; use noirc_frontend::monomorphization::ast::Program; -use self::acir_gen::Acir; +use self::ssa_gen::Ssa; mod acir_gen; mod ir; @@ -24,9 +24,15 @@ pub mod ssa_gen; /// Optimize the given program by converting it into SSA /// form and performing optimizations there. When finished, /// convert the final SSA into ACIR and return it. -pub fn optimize_into_acir(program: Program) -> Acir { - ssa_gen::generate_ssa(program).inline_functions().into_acir() +pub fn optimize_into_acir(program: Program) { + ssa_gen::generate_ssa(program) + .print("Initial SSA:") + .inline_functions() + .print("After Inlining:") + .unroll_loops() + .print("After Unrolling:"); } + /// Compiles the Program into ACIR and applies optimizations to the arithmetic gates /// This is analogous to `ssa:create_circuit` and this method is called when one wants /// to use the new ssa module to process Noir code. @@ -37,5 +43,14 @@ pub fn experimental_create_circuit( _enable_logging: bool, _show_output: bool, ) -> Result<(Circuit, Abi), RuntimeError> { - todo!("this is a stub function for the new SSA refactor module") + optimize_into_acir(_program); + let error_kind = RuntimeErrorKind::Spanless("Acir-gen is unimplemented".into()); + Err(RuntimeError::new(error_kind, None)) +} + +impl Ssa { + fn print(self, msg: &str) -> Ssa { + println!("{msg}\n{self}"); + self + } } diff --git a/crates/noirc_evaluator/src/ssa_refactor/acir_gen/mod.rs b/crates/noirc_evaluator/src/ssa_refactor/acir_gen/mod.rs index a0959db5db8..ddda689a0da 100644 --- a/crates/noirc_evaluator/src/ssa_refactor/acir_gen/mod.rs +++ b/crates/noirc_evaluator/src/ssa_refactor/acir_gen/mod.rs @@ -6,7 +6,7 @@ use super::ssa_gen::Ssa; struct Context {} /// The output of the Acir-gen pass -pub struct Acir {} +pub(crate) struct Acir {} impl Ssa { pub(crate) fn into_acir(self) -> Acir { 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 72954f2d0bd..590b780d4df 100644 --- a/crates/noirc_evaluator/src/ssa_refactor/ir/basic_block.rs +++ b/crates/noirc_evaluator/src/ssa_refactor/ir/basic_block.rs @@ -29,10 +29,10 @@ pub(crate) struct BasicBlock { pub(crate) type BasicBlockId = Id; impl BasicBlock { - /// Create a new BasicBlock with the given instructions. + /// Create a new BasicBlock with the given parameters. /// Parameters can also be added later via BasicBlock::add_parameter - pub(crate) fn new(instructions: Vec) -> Self { - Self { parameters: Vec::new(), instructions, terminator: None } + pub(crate) fn new() -> Self { + Self { parameters: Vec::new(), instructions: Vec::new(), terminator: None } } /// Returns the parameters of this block @@ -52,6 +52,12 @@ impl BasicBlock { self.parameters.push(parameter); } + /// Replace this block's current parameters with that of the given Vec. + /// This does not perform any checks that any previous parameters were unused. + pub(crate) fn set_parameters(&mut self, parameters: Vec) { + self.parameters = parameters; + } + /// Insert an instruction at the end of this block pub(crate) fn insert_instruction(&mut self, instruction: InstructionId) { self.instructions.push(instruction); @@ -83,6 +89,32 @@ impl BasicBlock { self.terminator.as_ref() } + /// Returns the terminator of this block, panics if there is None. + /// + /// Once this block has finished construction, this is expected to always be Some. + pub(crate) fn unwrap_terminator(&self) -> &TerminatorInstruction { + self.terminator().expect("Expected block to have terminator instruction") + } + + /// 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") + } + + /// 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() }) + } + /// Iterate over all the successors of the currently block, as determined by /// the blocks jumped to in the terminator instruction. If there is no terminator /// instruction yet, this will iterate 0 times. @@ -107,23 +139,4 @@ impl BasicBlock { }); self.instructions.remove(index); } - - /// 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() }) - } - - /// 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/ir/dfg.rs b/crates/noirc_evaluator/src/ssa_refactor/ir/dfg.rs index eaec4920790..0048761ff6f 100644 --- a/crates/noirc_evaluator/src/ssa_refactor/ir/dfg.rs +++ b/crates/noirc_evaluator/src/ssa_refactor/ir/dfg.rs @@ -1,5 +1,7 @@ use std::collections::HashMap; +use crate::ssa_refactor::ir::instruction::SimplifyResult; + use super::{ basic_block::{BasicBlock, BasicBlockId}, constant::{NumericConstant, NumericConstantId}, @@ -13,6 +15,7 @@ use super::{ }; use acvm::FieldElement; +use iter_extended::vecmap; /// The DataFlowGraph contains most of the actual data in a function including /// its blocks, instructions, and values. This struct is largely responsible for @@ -65,7 +68,27 @@ impl DataFlowGraph { /// After being created, the block is unreachable in the current function /// until another block is made to jump to it. pub(crate) fn make_block(&mut self) -> BasicBlockId { - self.blocks.insert(BasicBlock::new(Vec::new())) + self.blocks.insert(BasicBlock::new()) + } + + /// Create a new block with the same parameter count and parameter + /// types from the given block. + /// This is a somewhat niche operation used in loop unrolling but is included + /// here as doing it outside the DataFlowGraph would require cloning the parameters. + pub(crate) fn make_block_with_parameters_from_block( + &mut self, + block: BasicBlockId, + ) -> BasicBlockId { + let new_block = self.make_block(); + let parameters = self.blocks[block].parameters(); + + let parameters = vecmap(parameters.iter().enumerate(), |(position, param)| { + let typ = self.values[*param].get_type(); + self.values.insert(Value::Param { block: new_block, position, typ }) + }); + + self.blocks[new_block].set_parameters(parameters); + new_block } /// Get an iterator over references to each basic block within the dfg, paired with the basic @@ -101,17 +124,19 @@ impl DataFlowGraph { } /// Inserts a new instruction at the end of the given block and returns its results - pub(crate) fn insert_instruction( + pub(crate) fn insert_instruction_and_results( &mut self, instruction: Instruction, block: BasicBlockId, ctrl_typevars: Option>, ) -> InsertInstructionResult { + use InsertInstructionResult::*; match instruction.simplify(self) { - Some(simplification) => InsertInstructionResult::SimplifiedTo(simplification), - None => { + SimplifyResult::SimplifiedTo(simplification) => SimplifiedTo(simplification), + SimplifyResult::Remove => InstructionRemoved, + SimplifyResult::None => { let id = self.make_instruction(instruction, ctrl_typevars); - self.insert_instruction_in_block(block, id); + self.blocks[block].insert_instruction(id); InsertInstructionResult::Results(self.instruction_results(id)) } } @@ -246,16 +271,6 @@ impl DataFlowGraph { parameter } - /// Insert an instruction at the end of a given block. - /// If the block already has a terminator, the instruction is inserted before the terminator. - pub(crate) fn insert_instruction_in_block( - &mut self, - block: BasicBlockId, - instruction: InstructionId, - ) { - self.blocks[block].insert_instruction(instruction); - } - /// Returns the field element represented by this value if it is a numeric constant. /// Returns None if the given value is not a numeric constant. pub(crate) fn get_numeric_constant(&self, value: Id) -> Option { @@ -282,6 +297,20 @@ impl DataFlowGraph { ) { self.blocks[block].set_terminator(terminator); } + + /// Moves the entirety of the given block's contents into the destination block. + /// The source block afterward will be left in a valid but emptied state. The + /// destination block will also have its terminator overwritten with that of the + /// source block. + pub(crate) fn inline_block(&mut self, source: BasicBlockId, destination: BasicBlockId) { + let source = &mut self.blocks[source]; + let mut instructions = std::mem::take(source.instructions_mut()); + let terminator = source.take_terminator(); + + let destination = &mut self.blocks[destination]; + destination.instructions_mut().append(&mut instructions); + destination.set_terminator(terminator); + } } impl std::ops::Index for DataFlowGraph { diff --git a/crates/noirc_evaluator/src/ssa_refactor/ir/instruction.rs b/crates/noirc_evaluator/src/ssa_refactor/ir/instruction.rs index 42968568dee..7ca23c6f8a9 100644 --- a/crates/noirc_evaluator/src/ssa_refactor/ir/instruction.rs +++ b/crates/noirc_evaluator/src/ssa_refactor/ir/instruction.rs @@ -160,10 +160,16 @@ impl Instruction { /// Try to simplify this instruction. If the instruction can be simplified to a known value, /// that value is returned. Otherwise None is returned. - pub(crate) fn simplify(&self, dfg: &mut DataFlowGraph) -> Option { + pub(crate) fn simplify(&self, dfg: &mut DataFlowGraph) -> SimplifyResult { + use SimplifyResult::*; match self { Instruction::Binary(binary) => binary.simplify(dfg), - Instruction::Cast(value, typ) => (*typ == dfg.type_of_value(*value)).then_some(*value), + Instruction::Cast(value, typ) => { + match (*typ == dfg.type_of_value(*value)).then_some(*value) { + Some(value) => SimplifiedTo(value), + _ => None, + } + } Instruction::Not(value) => { match &dfg[*value] { // Limit optimizing ! on constants to only booleans. If we tried it on fields, @@ -171,12 +177,12 @@ impl Instruction { // would be incorrect however since the extra bits on the field would not be flipped. Value::NumericConstant { constant, typ } if *typ == Type::bool() => { let value = dfg[*constant].value().is_zero() as u128; - Some(dfg.make_constant(value.into(), Type::bool())) + SimplifiedTo(dfg.make_constant(value.into(), Type::bool())) } Value::Instruction { instruction, .. } => { // !!v => v match &dfg[*instruction] { - Instruction::Not(value) => Some(*value), + Instruction::Not(value) => SimplifiedTo(*value), _ => None, } } @@ -186,8 +192,7 @@ impl Instruction { Instruction::Constrain(value) => { if let Some(constant) = dfg.get_numeric_constant(*value) { if constant.is_one() { - // "simplify" to a unit literal that will just be thrown away anyway - return Some(dfg.make_constant(0u128.into(), Type::Unit)); + return Remove; } } None @@ -248,6 +253,44 @@ pub(crate) enum TerminatorInstruction { Return { return_values: Vec }, } +impl TerminatorInstruction { + /// Map each ValueId in this terminator to a new value. + pub(crate) fn map_values( + &self, + mut f: impl FnMut(ValueId) -> ValueId, + ) -> TerminatorInstruction { + use TerminatorInstruction::*; + match self { + JmpIf { condition, then_destination, else_destination } => JmpIf { + condition: f(*condition), + then_destination: *then_destination, + else_destination: *else_destination, + }, + Jmp { destination, arguments } => { + Jmp { destination: *destination, arguments: vecmap(arguments, |value| f(*value)) } + } + Return { return_values } => { + Return { return_values: vecmap(return_values, |value| f(*value)) } + } + } + } + + /// Mutate each BlockId to a new BlockId specified by the given mapping function. + pub(crate) fn mutate_blocks(&mut self, mut f: impl FnMut(BasicBlockId) -> BasicBlockId) { + use TerminatorInstruction::*; + match self { + JmpIf { then_destination, else_destination, .. } => { + *then_destination = f(*then_destination); + *else_destination = f(*else_destination); + } + Jmp { destination, .. } => { + *destination = f(*destination); + } + Return { .. } => (), + } + } +} + /// A binary instruction in the IR. #[derive(Debug, PartialEq, Eq, Hash, Clone)] pub(crate) struct Binary { @@ -269,13 +312,16 @@ impl Binary { } /// Try to simplify this binary instruction, returning the new value if possible. - fn simplify(&self, dfg: &mut DataFlowGraph) -> Option { + fn simplify(&self, dfg: &mut DataFlowGraph) -> SimplifyResult { let lhs = dfg.get_numeric_constant(self.lhs); let rhs = dfg.get_numeric_constant(self.rhs); let operand_type = dfg.type_of_value(self.lhs); if let (Some(lhs), Some(rhs)) = (lhs, rhs) { - return self.eval_constants(dfg, lhs, rhs, operand_type); + return match self.eval_constants(dfg, lhs, rhs, operand_type) { + Some(value) => SimplifyResult::SimplifiedTo(value), + None => SimplifyResult::None, + }; } let lhs_is_zero = lhs.map_or(false, |lhs| lhs.is_zero()); @@ -287,71 +333,80 @@ impl Binary { match self.operator { BinaryOp::Add => { if lhs_is_zero { - return Some(self.rhs); + return SimplifyResult::SimplifiedTo(self.rhs); } if rhs_is_zero { - return Some(self.lhs); + return SimplifyResult::SimplifiedTo(self.lhs); } } BinaryOp::Sub => { if rhs_is_zero { - return Some(self.lhs); + return SimplifyResult::SimplifiedTo(self.lhs); } } BinaryOp::Mul => { if lhs_is_one { - return Some(self.rhs); + return SimplifyResult::SimplifiedTo(self.rhs); } if rhs_is_one { - return Some(self.lhs); + return SimplifyResult::SimplifiedTo(self.lhs); } } BinaryOp::Div => { if rhs_is_one { - return Some(self.lhs); + return SimplifyResult::SimplifiedTo(self.lhs); } } BinaryOp::Mod => { if rhs_is_one { - return Some(self.lhs); + let zero = dfg.make_constant(FieldElement::zero(), operand_type); + return SimplifyResult::SimplifiedTo(zero); } } BinaryOp::Eq => { if self.lhs == self.rhs { - return Some(dfg.make_constant(FieldElement::one(), Type::bool())); + let one = dfg.make_constant(FieldElement::one(), Type::bool()); + return SimplifyResult::SimplifiedTo(one); } } BinaryOp::Lt => { if self.lhs == self.rhs { - return Some(dfg.make_constant(FieldElement::zero(), Type::bool())); + let zero = dfg.make_constant(FieldElement::zero(), Type::bool()); + return SimplifyResult::SimplifiedTo(zero); } } BinaryOp::And => { if lhs_is_zero || rhs_is_zero { - return Some(dfg.make_constant(FieldElement::zero(), operand_type)); + let zero = dfg.make_constant(FieldElement::zero(), operand_type); + return SimplifyResult::SimplifiedTo(zero); } } BinaryOp::Or => { if lhs_is_zero { - return Some(self.rhs); + return SimplifyResult::SimplifiedTo(self.rhs); } if rhs_is_zero { - return Some(self.lhs); + return SimplifyResult::SimplifiedTo(self.lhs); + } + } + BinaryOp::Xor => { + if self.lhs == self.rhs { + let zero = dfg.make_constant(FieldElement::zero(), Type::bool()); + return SimplifyResult::SimplifiedTo(zero); } } - BinaryOp::Xor => (), BinaryOp::Shl => { if rhs_is_zero { - return Some(self.lhs); + return SimplifyResult::SimplifiedTo(self.lhs); } } BinaryOp::Shr => { if rhs_is_zero { - return Some(self.lhs); + return SimplifyResult::SimplifiedTo(self.lhs); } } } - None + SimplifyResult::None } /// Evaluate the two constants with the operation specified by self.operator. @@ -477,3 +532,16 @@ impl std::fmt::Display for BinaryOp { } } } + +/// Contains the result to Instruction::simplify, specifying how the instruction +/// should be simplified. +pub(crate) enum SimplifyResult { + /// Replace this function's result with the given value + SimplifiedTo(ValueId), + + /// Remove the instruction, it is unnecessary + Remove, + + /// Instruction could not be simplified + None, +} diff --git a/crates/noirc_evaluator/src/ssa_refactor/opt/inlining.rs b/crates/noirc_evaluator/src/ssa_refactor/opt/inlining.rs index c63cac520bf..53c3f566a9a 100644 --- a/crates/noirc_evaluator/src/ssa_refactor/opt/inlining.rs +++ b/crates/noirc_evaluator/src/ssa_refactor/opt/inlining.rs @@ -395,25 +395,21 @@ impl<'function> PerFunctionContext<'function> { block_id: BasicBlockId, block_queue: &mut Vec, ) -> Option<(BasicBlockId, Vec)> { - match self.source_function.dfg[block_id].terminator() { - Some(TerminatorInstruction::Jmp { destination, arguments }) => { + match self.source_function.dfg[block_id].unwrap_terminator() { + TerminatorInstruction::Jmp { destination, arguments } => { let destination = self.translate_block(*destination, block_queue); let arguments = vecmap(arguments, |arg| self.translate_value(*arg)); self.context.builder.terminate_with_jmp(destination, arguments); None } - Some(TerminatorInstruction::JmpIf { - condition, - then_destination, - else_destination, - }) => { + TerminatorInstruction::JmpIf { condition, then_destination, else_destination } => { let condition = self.translate_value(*condition); let then_block = self.translate_block(*then_destination, block_queue); let else_block = self.translate_block(*else_destination, block_queue); self.context.builder.terminate_with_jmpif(condition, then_block, else_block); None } - Some(TerminatorInstruction::Return { return_values }) => { + TerminatorInstruction::Return { return_values } => { let return_values = vecmap(return_values, |value| self.translate_value(*value)); if self.inlining_main { self.context.builder.terminate_with_return(return_values.clone()); @@ -421,7 +417,6 @@ impl<'function> PerFunctionContext<'function> { let block_id = self.translate_block(block_id, block_queue); Some((block_id, return_values)) } - None => unreachable!("Block has no terminator instruction"), } } } diff --git a/crates/noirc_evaluator/src/ssa_refactor/opt/mod.rs b/crates/noirc_evaluator/src/ssa_refactor/opt/mod.rs index d9d31f8f356..d4d75221685 100644 --- a/crates/noirc_evaluator/src/ssa_refactor/opt/mod.rs +++ b/crates/noirc_evaluator/src/ssa_refactor/opt/mod.rs @@ -6,3 +6,4 @@ mod inlining; mod mem2reg; mod simplify_cfg; +mod unrolling; 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 419120fc4e8..35d3ebc8733 100644 --- a/crates/noirc_evaluator/src/ssa_refactor/opt/simplify_cfg.rs +++ b/crates/noirc_evaluator/src/ssa_refactor/opt/simplify_cfg.rs @@ -130,25 +130,16 @@ fn remove_block_parameters( fn try_inline_into_predecessor( function: &mut Function, cfg: &mut ControlFlowGraph, - block_id: BasicBlockId, - predecessor_id: BasicBlockId, + block: BasicBlockId, + predecessor: BasicBlockId, ) -> bool { - let mut successors = cfg.successors(predecessor_id); - if successors.len() == 1 && successors.next() == Some(block_id) { + let mut successors = cfg.successors(predecessor); + if successors.len() == 1 && successors.next() == Some(block) { drop(successors); + function.dfg.inline_block(block, predecessor); - // 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); - - predecessor.set_terminator(terminator); - cfg.recompute_block(function, block_id); - cfg.recompute_block(function, predecessor_id); + cfg.recompute_block(function, block); + cfg.recompute_block(function, predecessor); true } else { false diff --git a/crates/noirc_evaluator/src/ssa_refactor/opt/unrolling.rs b/crates/noirc_evaluator/src/ssa_refactor/opt/unrolling.rs new file mode 100644 index 00000000000..dba64dde6b4 --- /dev/null +++ b/crates/noirc_evaluator/src/ssa_refactor/opt/unrolling.rs @@ -0,0 +1,659 @@ +//! This file contains the loop unrolling pass for the new SSA IR. +//! +//! This pass is divided into a few steps: +//! 1. Find all loops in the program (`find_all_loops`) +//! 2. For each loop: +//! a. If the loop is in our list of loops that previously failed to unroll, skip it. +//! b. If we have previously modified any of the blocks in the loop, +//! restart from step 1 to refresh the context. +//! c. If not, try to unroll the loop. If successful, remember the modified +//! blocks. If not, remember that the loop failed to unroll and leave it +//! unmodified. +//! +//! Note that this pass also often creates superfluous jmp instructions in the +//! program that will need to be removed by a later simplify cfg pass. +use std::collections::{HashMap, HashSet}; + +use iter_extended::vecmap; + +use crate::ssa_refactor::{ + ir::{ + basic_block::BasicBlockId, + cfg::ControlFlowGraph, + dfg::InsertInstructionResult, + dom::DominatorTree, + function::Function, + instruction::{InstructionId, TerminatorInstruction}, + post_order::PostOrder, + value::ValueId, + }, + ssa_gen::Ssa, +}; + +impl Ssa { + /// Unroll all loops in each SSA function. + /// If any loop cannot be unrolled, it is left as-is or in a partially unrolled state. + pub(crate) fn unroll_loops(mut self) -> Ssa { + for function in self.functions.values_mut() { + find_all_loops(function).unroll_each_loop(function); + } + self + } +} + +struct Loop { + /// The header block of a loop is the block which dominates all the + /// other blocks in the loop. + header: BasicBlockId, + + /// The start of the back_edge n -> d is the block n at the end of + /// the loop that jumps back to the header block d which restarts the loop. + back_edge_start: BasicBlockId, + + /// All the blocks contained within the loop, including `header` and `back_edge_start`. + blocks: HashSet, +} + +struct Loops { + /// The loops that failed to be unrolled so that we do not try to unroll them again. + /// Each loop is identified by its header block id. + failed_to_unroll: HashSet, + + yet_to_unroll: Vec, + modified_blocks: HashSet, + cfg: ControlFlowGraph, + dom_tree: DominatorTree, +} + +/// Find a loop in the program by finding a node that dominates any predecessor node. +/// The edge where this happens will be the back-edge of the loop. +fn find_all_loops(function: &Function) -> Loops { + let cfg = ControlFlowGraph::with_function(function); + let post_order = PostOrder::with_function(function); + let dom_tree = DominatorTree::with_cfg_and_post_order(&cfg, &post_order); + + let mut loops = vec![]; + + for (block, _) in function.dfg.basic_blocks_iter() { + // These reachable checks wouldn't be needed if we only iterated over reachable blocks + if dom_tree.is_reachable(block) { + for predecessor in cfg.predecessors(block) { + if dom_tree.is_reachable(predecessor) && dom_tree.dominates(block, predecessor) { + // predecessor -> block is the back-edge of a loop + loops.push(find_blocks_in_loop(block, predecessor, &cfg)); + } + } + } + } + + // Sort loops by block size so that we unroll the smaller, nested loops first as an optimization. + loops.sort_by(|loop_a, loop_b| loop_b.blocks.len().cmp(&loop_a.blocks.len())); + + Loops { + failed_to_unroll: HashSet::new(), + yet_to_unroll: loops, + modified_blocks: HashSet::new(), + cfg, + dom_tree, + } +} + +impl Loops { + /// Unroll all loops within a given function. + /// Any loops which fail to be unrolled (due to using non-constant indices) will be unmodified. + fn unroll_each_loop(mut self, function: &mut Function) { + while let Some(next_loop) = self.yet_to_unroll.pop() { + // If we've previously modified a block in this loop we need to refresh the context. + // This happens any time we have nested loops. + if next_loop.blocks.iter().any(|block| self.modified_blocks.contains(block)) { + let mut new_context = find_all_loops(function); + new_context.failed_to_unroll = self.failed_to_unroll; + return new_context.unroll_each_loop(function); + } + + // Don't try to unroll the loop again if it is known to fail + if !self.failed_to_unroll.contains(&next_loop.header) { + if unroll_loop(function, &self.cfg, &next_loop).is_ok() { + self.modified_blocks.extend(next_loop.blocks); + } else { + self.failed_to_unroll.insert(next_loop.header); + } + } + } + } +} + +/// Return each block that is in a loop starting in the given header block. +/// Expects back_edge_start -> header to be the back edge of the loop. +fn find_blocks_in_loop( + header: BasicBlockId, + back_edge_start: BasicBlockId, + cfg: &ControlFlowGraph, +) -> Loop { + let mut blocks = HashSet::new(); + blocks.insert(header); + + let mut insert = |block, stack: &mut Vec| { + if !blocks.contains(&block) { + blocks.insert(block); + stack.push(block); + } + }; + + // Starting from the back edge of the loop, each predecessor of this block until + // the header is within the loop. + let mut stack = vec![]; + insert(back_edge_start, &mut stack); + + while let Some(block) = stack.pop() { + for predecessor in cfg.predecessors(block) { + insert(predecessor, &mut stack); + } + } + + Loop { header, back_edge_start, blocks } +} + +/// Unroll a single loop in the function. +/// Returns Err(()) if it failed to unroll and Ok(()) otherwise. +fn unroll_loop(function: &mut Function, cfg: &ControlFlowGraph, loop_: &Loop) -> Result<(), ()> { + let mut unroll_into = get_pre_header(cfg, loop_); + let mut jump_value = get_induction_variable(function, unroll_into)?; + + while let Some(context) = unroll_loop_header(function, loop_, unroll_into, jump_value) { + let (last_block, last_value) = context.unroll_loop_iteration(); + unroll_into = last_block; + jump_value = last_value; + } + + Ok(()) +} + +/// The loop pre-header is the block that comes before the loop begins. Generally a header block +/// is expected to have 2 predecessors: the pre-header and the final block of the loop which jumps +/// back to the beginning. +fn get_pre_header(cfg: &ControlFlowGraph, loop_: &Loop) -> BasicBlockId { + let mut pre_header = cfg + .predecessors(loop_.header) + .filter(|predecessor| *predecessor != loop_.back_edge_start) + .collect::>(); + + assert_eq!(pre_header.len(), 1); + pre_header.remove(0) +} + +/// Return the induction value of the current iteration of the loop, from the given block's jmp arguments. +/// +/// Expects the current block to terminate in `jmp h(N)` where h is the loop header and N is +/// a Field value. +fn get_induction_variable(function: &Function, block: BasicBlockId) -> Result { + match function.dfg[block].terminator() { + Some(TerminatorInstruction::Jmp { arguments, .. }) => { + // This assumption will no longer be valid if e.g. mutable variables are represented as + // block parameters. If that becomes the case we'll need to figure out which variable + // is generally constant and increasing to guess which parameter is the induction + // variable. + assert_eq!(arguments.len(), 1, "It is expected that a loop's induction variable is the only block parameter of the loop header"); + let value = arguments[0]; + if function.dfg.get_numeric_constant(value).is_some() { + Ok(value) + } else { + Err(()) + } + } + _ => Err(()), + } +} + +/// Unrolls the header block of the loop. This is the block that dominates all other blocks in the +/// loop and contains the jmpif instruction that lets us know if we should continue looping. +/// Returns Some(iteration context) if we should perform another iteration. +fn unroll_loop_header<'a>( + function: &'a mut Function, + loop_: &'a Loop, + unroll_into: BasicBlockId, + induction_value: ValueId, +) -> Option> { + // We insert into a fresh block first and move instructions into the unroll_into block later + // only once we verify the jmpif instruction has a constant condition. If it does not, we can + // just discard this fresh block and leave the loop unmodified. + let fresh_block = function.dfg.make_block(); + + let mut context = LoopIteration::new(function, loop_, fresh_block, loop_.header); + let source_block = &context.function.dfg[context.source_block]; + assert_eq!(source_block.parameters().len(), 1, "Expected only 1 argument in loop header"); + + // Insert the current value of the loop induction variable into our context. + let first_param = source_block.parameters()[0]; + context.values.insert(first_param, induction_value); + context.inline_instructions_from_block(); + + match context.function.dfg[fresh_block].unwrap_terminator() { + TerminatorInstruction::JmpIf { condition, then_destination, else_destination } => { + let next_blocks = context.handle_jmpif(*condition, *then_destination, *else_destination); + + // If there is only 1 next block the jmpif evaluated to a single known block. + // This is the expected case and lets us know if we should loop again or not. + if next_blocks.len() == 1 { + context.function.dfg.inline_block(fresh_block, unroll_into); + + // The fresh block is gone now so we're committing to insert into the original + // unroll_into block from now on. + context.insert_block = unroll_into; + + loop_.blocks.contains(&context.source_block).then_some(context) + } else { + // If this case is reached the loop either uses non-constant indices or we need + // another pass, such as mem2reg to resolve them to constants. + None + } + } + other => unreachable!("Expected loop header to terminate in a JmpIf to the loop body, but found {other:?} instead"), + } +} + +/// The context object for each loop iteration. +/// Notably each loop iteration maps each loop block to a fresh, unrolled block. +struct LoopIteration<'f> { + function: &'f mut Function, + loop_: &'f Loop, + + /// Maps pre-unrolled ValueIds to unrolled ValueIds. + /// These will often be the exact same as before, unless the ValueId was + /// dependent on the loop induction variable which is changing on each iteration. + values: HashMap, + + /// Maps pre-unrolled block ids from within the loop to new block ids of each loop + /// block for each loop iteration. + blocks: HashMap, + + /// Maps unrolled block ids back to the original source block ids + original_blocks: HashMap, + visited_blocks: HashSet, + + insert_block: BasicBlockId, + source_block: BasicBlockId, + + /// The induction value (and the block it was found in) is the new value for + /// the variable traditionally called `i` on each iteration of the loop. + /// This is None until we visit the block which jumps back to the start of the + /// loop, at which point we record its value and the block it was found in. + induction_value: Option<(BasicBlockId, ValueId)>, +} + +impl<'f> LoopIteration<'f> { + fn new( + function: &'f mut Function, + loop_: &'f Loop, + insert_block: BasicBlockId, + source_block: BasicBlockId, + ) -> Self { + Self { + function, + loop_, + insert_block, + source_block, + values: HashMap::new(), + blocks: HashMap::new(), + original_blocks: HashMap::new(), + visited_blocks: HashSet::new(), + induction_value: None, + } + } + + /// Unroll a single iteration of the loop. + /// + /// Note that after unrolling a single iteration, the loop is _not_ in a valid state. + /// It is expected the terminator instructions are set up to branch into an empty block + /// for further unrolling. When the loop is finished this will need to be mutated to + /// jump to the end of the loop instead. + fn unroll_loop_iteration(mut self) -> (BasicBlockId, ValueId) { + let mut next_blocks = self.unroll_loop_block(); + + while let Some(block) = next_blocks.pop() { + self.insert_block = block; + self.source_block = self.get_original_block(block); + + if !self.visited_blocks.contains(&self.source_block) { + let mut blocks = self.unroll_loop_block(); + next_blocks.append(&mut blocks); + } + } + + self.induction_value + .expect("Expected to find the induction variable by end of loop iteration") + } + + /// Unroll a single block in the current iteration of the loop + fn unroll_loop_block(&mut self) -> Vec { + let mut next_blocks = self.unroll_loop_block_helper(); + next_blocks.retain(|block| { + let b = self.get_original_block(*block); + self.loop_.blocks.contains(&b) + }); + next_blocks + } + + /// Unroll a single block in the current iteration of the loop + fn unroll_loop_block_helper(&mut self) -> Vec { + self.inline_instructions_from_block(); + self.visited_blocks.insert(self.source_block); + + match self.function.dfg[self.insert_block].unwrap_terminator() { + TerminatorInstruction::JmpIf { condition, then_destination, else_destination } => { + self.handle_jmpif(*condition, *then_destination, *else_destination) + } + TerminatorInstruction::Jmp { destination, arguments } => { + if self.get_original_block(*destination) == self.loop_.header { + assert_eq!(arguments.len(), 1); + self.induction_value = Some((self.insert_block, arguments[0])); + } + vec![*destination] + } + TerminatorInstruction::Return { .. } => vec![], + } + } + + /// Find the next branch(es) to take from a jmpif terminator and return them. + /// If only one block is returned, it means the jmpif condition evaluated to a known + /// constant and we can safely take only the given branch. + fn handle_jmpif( + &mut self, + condition: ValueId, + then_destination: BasicBlockId, + else_destination: BasicBlockId, + ) -> Vec { + let condition = self.get_value(condition); + + match self.function.dfg.get_numeric_constant(condition) { + Some(constant) => { + let destination = + if constant.is_zero() { else_destination } else { then_destination }; + + self.source_block = self.get_original_block(destination); + + let jmp = TerminatorInstruction::Jmp { destination, arguments: Vec::new() }; + self.function.dfg.set_block_terminator(self.insert_block, jmp); + vec![destination] + } + None => vec![then_destination, else_destination], + } + } + + /// Map a ValueId in the original pre-unrolled ssa to its new id in the unrolled SSA. + /// This is often the same ValueId as most values don't change while unrolling. The main + /// exception is instructions referencing the induction variable (or the variable itself) + /// which may have been simplified to another form. Block parameters or values outside the + /// loop shouldn't change at all and won't be present inside self.values. + fn get_value(&self, value: ValueId) -> ValueId { + self.values.get(&value).copied().unwrap_or(value) + } + + /// Translate a block id to a block id in the unrolled loop. If the given + /// block id is not within the loop, it is returned as-is. + fn get_or_insert_block(&mut self, block: BasicBlockId) -> BasicBlockId { + if let Some(new_block) = self.blocks.get(&block) { + return *new_block; + } + + // If the block is in the loop we create a fresh block for each iteration + if self.loop_.blocks.contains(&block) { + let new_block = self.function.dfg.make_block_with_parameters_from_block(block); + + let old_parameters = self.function.dfg.block_parameters(block); + let new_parameters = self.function.dfg.block_parameters(new_block); + + for (param, new_param) in old_parameters.iter().zip(new_parameters) { + // Don't overwrite any existing entries to avoid overwriting the induction variable + self.values.entry(*param).or_insert(*new_param); + } + + self.blocks.insert(block, new_block); + self.original_blocks.insert(new_block, block); + new_block + } else { + block + } + } + + fn get_original_block(&self, block: BasicBlockId) -> BasicBlockId { + self.original_blocks.get(&block).copied().unwrap_or(block) + } + + fn inline_instructions_from_block(&mut self) { + let source_block = &self.function.dfg[self.source_block]; + let instructions = source_block.instructions().to_vec(); + + // We cannot directly append each instruction since we need to substitute any + // instances of the induction variable or any values that were changed as a result + // of the new induction variable value. + for instruction in instructions { + self.push_instruction(instruction); + } + + let mut terminator = self.function.dfg[self.source_block] + .unwrap_terminator() + .map_values(|value| self.get_value(value)); + + terminator.mutate_blocks(|block| self.get_or_insert_block(block)); + self.function.dfg.set_block_terminator(self.insert_block, terminator); + } + + fn push_instruction(&mut self, id: InstructionId) { + let instruction = self.function.dfg[id].map_values(|id| self.get_value(id)); + let results = self.function.dfg.instruction_results(id).to_vec(); + + let ctrl_typevars = instruction + .requires_ctrl_typevars() + .then(|| vecmap(&results, |result| self.function.dfg.type_of_value(*result))); + + let new_results = self.function.dfg.insert_instruction_and_results( + instruction, + self.insert_block, + ctrl_typevars, + ); + + Self::insert_new_instruction_results(&mut self.values, &results, new_results); + } + + /// Modify the values HashMap to remember the mapping between an instruction result's previous + /// ValueId (from the source_function) and its new ValueId in the destination function. + fn insert_new_instruction_results( + values: &mut HashMap, + old_results: &[ValueId], + new_results: InsertInstructionResult, + ) { + assert_eq!(old_results.len(), new_results.len()); + + match new_results { + InsertInstructionResult::SimplifiedTo(new_result) => { + values.insert(old_results[0], new_result); + } + InsertInstructionResult::Results(new_results) => { + for (old_result, new_result) in old_results.iter().zip(new_results) { + values.insert(*old_result, *new_result); + } + } + InsertInstructionResult::InstructionRemoved => (), + } + } +} + +#[cfg(test)] +mod tests { + use crate::ssa_refactor::{ + ir::{instruction::BinaryOp, map::Id, types::Type}, + ssa_builder::FunctionBuilder, + }; + + #[test] + fn unroll_nested_loops() { + // fn main() { + // for i in 0..3 { + // for j in 0..4 { + // assert(i + j > 10); + // } + // } + // } + // + // fn main f0 { + // b0(): + // jmp b1(Field 0) + // b1(v0: Field): // header of outer loop + // v1 = lt v0, Field 3 + // jmpif v1, then: b2, else: b3 + // b2(): + // jmp b4(Field 0) + // b4(v2: Field): // header of inner loop + // v3 = lt v2, Field 4 + // jmpif v3, then: b5, else: b6 + // b5(): + // v4 = add v0, v2 + // v5 = lt Field 10, v4 + // constrain v5 + // v6 = add v2, Field 1 + // jmp b4(v6) + // b6(): // end of inner loop + // v7 = add v0, Field 1 + // jmp b1(v7) + // b3(): // end of outer loop + // return Field 0 + // } + let main_id = Id::test_new(0); + + // Compiling main + let mut builder = FunctionBuilder::new("main".into(), main_id); + + let b1 = builder.insert_block(); + let b2 = builder.insert_block(); + let b3 = builder.insert_block(); + let b4 = builder.insert_block(); + let b5 = builder.insert_block(); + let b6 = builder.insert_block(); + + let v0 = builder.add_block_parameter(b1, Type::field()); + let v2 = builder.add_block_parameter(b4, Type::field()); + + let zero = builder.field_constant(0u128); + let one = builder.field_constant(1u128); + let three = builder.field_constant(3u128); + let four = builder.field_constant(4u128); + let ten = builder.field_constant(10u128); + + builder.terminate_with_jmp(b1, vec![zero]); + + // b1 + builder.switch_to_block(b1); + let v1 = builder.insert_binary(v0, BinaryOp::Lt, three); + builder.terminate_with_jmpif(v1, b2, b3); + + // b2 + builder.switch_to_block(b2); + builder.terminate_with_jmp(b4, vec![zero]); + + // b3 + builder.switch_to_block(b3); + builder.terminate_with_return(vec![zero]); + + // b4 + builder.switch_to_block(b4); + let v3 = builder.insert_binary(v2, BinaryOp::Lt, four); + builder.terminate_with_jmpif(v3, b5, b6); + + // b5 + builder.switch_to_block(b5); + let v4 = builder.insert_binary(v0, BinaryOp::Add, v2); + let v5 = builder.insert_binary(ten, BinaryOp::Lt, v4); + builder.insert_constrain(v5); + let v6 = builder.insert_binary(v2, BinaryOp::Add, one); + builder.terminate_with_jmp(b4, vec![v6]); + + // b6 + builder.switch_to_block(b6); + let v7 = builder.insert_binary(v0, BinaryOp::Add, one); + builder.terminate_with_jmp(b1, vec![v7]); + + let ssa = builder.finish(); + assert_eq!(ssa.main().reachable_blocks().len(), 7); + + // Expected output: + // + // fn main f0 { + // b0(): + // constrain Field 0 + // constrain Field 0 + // constrain Field 0 + // constrain Field 0 + // jmp b23() + // b23(): + // constrain Field 0 + // constrain Field 0 + // constrain Field 0 + // constrain Field 0 + // jmp b27() + // b27(): + // constrain Field 0 + // constrain Field 0 + // constrain Field 0 + // constrain Field 0 + // jmp b31() + // b31(): + // jmp b3() + // b3(): + // return Field 0 + // } + // The final block count is not 1 because unrolling creates some unnecessary jmps. + // If a simplify cfg pass is ran afterward, the expected block count will be 1. + let ssa = ssa.unroll_loops(); + assert_eq!(ssa.main().reachable_blocks().len(), 5); + } + + // Test that the pass can still be run on loops which fail to unroll properly + #[test] + fn fail_to_unroll_loop() { + // fn main f0 { + // b0(v0: Field): + // jmp b1(v0) + // b1(v1: Field): + // v2 = lt v1, 5 + // jmpif v2, then: b2, else: b3 + // b2(): + // v3 = add v1, Field 1 + // jmp b1(v3) + // b3(): + // return Field 0 + // } + 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 b3 = builder.insert_block(); + + let v0 = builder.add_parameter(Type::field()); + let v1 = builder.add_block_parameter(b1, Type::field()); + + builder.terminate_with_jmp(b1, vec![v0]); + + builder.switch_to_block(b1); + let five = builder.field_constant(5u128); + let v2 = builder.insert_binary(v1, BinaryOp::Lt, five); + builder.terminate_with_jmpif(v2, b2, b3); + + builder.switch_to_block(b2); + let one = builder.field_constant(1u128); + let v3 = builder.insert_binary(v1, BinaryOp::Add, one); + builder.terminate_with_jmp(b1, vec![v3]); + + builder.switch_to_block(b3); + let zero = builder.field_constant(0u128); + builder.terminate_with_return(vec![zero]); + + let ssa = builder.finish(); + assert_eq!(ssa.main().reachable_blocks().len(), 4); + + // Expected ssa is unchanged + let ssa = ssa.unroll_loops(); + assert_eq!(ssa.main().reachable_blocks().len(), 4); + } +} diff --git a/crates/noirc_evaluator/src/ssa_refactor/ssa_builder/mod.rs b/crates/noirc_evaluator/src/ssa_refactor/ssa_builder/mod.rs index 60379097523..3fe0b885bde 100644 --- a/crates/noirc_evaluator/src/ssa_refactor/ssa_builder/mod.rs +++ b/crates/noirc_evaluator/src/ssa_refactor/ssa_builder/mod.rs @@ -110,7 +110,11 @@ impl FunctionBuilder { instruction: Instruction, ctrl_typevars: Option>, ) -> InsertInstructionResult { - self.current_function.dfg.insert_instruction(instruction, self.current_block, ctrl_typevars) + self.current_function.dfg.insert_instruction_and_results( + instruction, + self.current_block, + ctrl_typevars, + ) } /// Switch to inserting instructions in the given block.