diff --git a/crates/nargo_cli/tests/execution_success/references/src/main.nr b/crates/nargo_cli/tests/execution_success/references/src/main.nr index ec23f2f3a38..70de5cada3f 100644 --- a/crates/nargo_cli/tests/execution_success/references/src/main.nr +++ b/crates/nargo_cli/tests/execution_success/references/src/main.nr @@ -229,4 +229,4 @@ fn regression_2218_loop(x: Field, y: Field) { assert(*q1 == 20); } assert(*q1 == 2); -} \ No newline at end of file +} diff --git a/crates/nargo_cli/tests/execution_success/references_aliasing/Nargo.toml b/crates/nargo_cli/tests/execution_success/references_aliasing/Nargo.toml new file mode 100644 index 00000000000..b52fdcf77f0 --- /dev/null +++ b/crates/nargo_cli/tests/execution_success/references_aliasing/Nargo.toml @@ -0,0 +1,7 @@ +[package] +name = "references" +type = "bin" +authors = [""] +compiler_version = "0.5.1" + +[dependencies] diff --git a/crates/nargo_cli/tests/execution_success/references_aliasing/Prover.toml b/crates/nargo_cli/tests/execution_success/references_aliasing/Prover.toml new file mode 100644 index 00000000000..e69de29bb2d diff --git a/crates/nargo_cli/tests/execution_success/references_aliasing/src/main.nr b/crates/nargo_cli/tests/execution_success/references_aliasing/src/main.nr new file mode 100644 index 00000000000..4582444c8f7 --- /dev/null +++ b/crates/nargo_cli/tests/execution_success/references_aliasing/src/main.nr @@ -0,0 +1,10 @@ +fn increment(mut r: &mut Field) { + *r = *r + 1; +} + +fn main() { + let mut x = 100; + let mut xref = &mut x; + increment(xref); + assert(*xref == 101); +} diff --git a/crates/nargo_cli/tests/execution_success/references_aliasing/target/references.bytecode b/crates/nargo_cli/tests/execution_success/references_aliasing/target/references.bytecode new file mode 100644 index 00000000000..3e521923327 --- /dev/null +++ b/crates/nargo_cli/tests/execution_success/references_aliasing/target/references.bytecode @@ -0,0 +1 @@ +H4sIAAAAAAAA/+1c8ZONVRh+WCKykYqIiFRE37f37u5dkY2IEEJEtLv2bkJKQkpCSkJKIjLGNE3TmJqpmZqpmZrRH9L/Uu907szZzfjlPO+Z98yc88u+w8yz3/O853ue93zf3fsXgL/x/zXY/ex0P4uwVQ7ysCpFW7Vab2+pl5Wyu2jp6Km1FtXWnrZaWStba629LbVKpV6r1to7ejrai46yWqmXfa0dlT4HNpiAVe/7bzX9izHoJvwHkfkP5mEV/vU2efWQAb3z//82BU4Y8HsG6th8k3+j/nKNJjUp4A4Bb/Nr8R7C71HhQZrXtLEsG99QpGd8Q6FjfLd5dTa+QMyhSkINg23jE97D+D1SNT62po1l2fiGIz3jGw4d47vdq7PxBWIOd4KycUfAtvEJ7xH8HqkaH1vTxrJsfCORnvGNhI7x3eHV2fgCMUc6Qdm4o2Db+IT3KH6PVI2PrWljWTa+ZtCMrxbL+JqhY3x3enU2vkDMZicoG3c0bBuf8B7N71EyhjIGLEOp12MZyhjoGMpdXp0NJRBzjBOUjTsWtg1FeI/l9+iWhhKqA9Ok74bOTcHmzNxH9yTCmdnneyNxLsJWL7PP43jX1SIYbO+Rnoy7CW4o76vQ8bEmKv+yytzf44l9uQrWkBXvcRWRf78h6z6vzkNWIOZ4JygbdwJsD1nCewK/R6qPq9iaNhZ7SGCeLifS9CsrsYyPd839je9+r87GF4g50QnKxp0E28YnvCfxe6T5uKqVqelkIudYhsK8Zv96H/DqbCiBmJOdoGzcKbBtKMJ7Cr9HqpPUFMSZpIqwRX1OP5WAFfsIORU6xvegV2fjC8Sc6gRl406DbeMT3tP4PUrGUKaDZSjxXvxNh46hPOTV2VACMac7Qdm4M2DbUIT3DH6PVJ/1ME36YejcFGzOzH30SCKcmX1+NBLnImz1Mvs8k3ddKi/+pCczwX/xdw06PsZ+8cfc37OIfbkG2pDVG2vIIvLvN2Q95tV5yArEnOUEZePOhu0hS3jP5vdI9XEVW9PGsny6nAOW8cX7nPoc6Bjf416djS8Qc44TlI1bwLbxCe+C3yNFQylbmJqWRM6xDIV5zf71tnh1NpRAzNIJysatwLahCO8Kv0eqkxRb08ayPElVkd4kVYWO8bV6dTa+QMyqE5SN2wbbxie82/g9SsZQ2pGeobRDx1BqXp0NJRCz3QnKxu2AbUMR3h38HiVjKHORnqHMhY6hPOHV2VACMec6Qdm482DbUIT3PH6PkjGU+UjPUOZDx1Ce9OpsKIGY852gbNwFsG0ownsBv0fJGEon0jOUTugYylNenQ0lELPTCcrGXQjbhiI4C/k9SsZQFiE9Q1kEHUN52quzoQRiLnKCsnEXw7ahCO/F/B4lYyhLwDKUeH+NsQQ6hvKMV2dDCcRc4gRl4y6FbUMR3kv5PUrGUJYhPUNZBh1Dedars6EEYi5zgrJxl8O2oQjv5fweqf55F9OkV/A4t9yKcxG2qCa6EumZ6EromOhzXp1NNBBzpROUjbsKtk1UeK/i90jlWsXsV4D/N3XfQCc8mkjX2fiWEGYgryb2hahfMkG0BukF0RroBNHzXp2DKBBzjROUjbsWtoNIeK/l90jlWiUwV4MfRN8ijSBiDjXriH0h6pdMEK1HekG0HjpB9IJX5yAKxFzvBGXjboDtIBLeG/g9UrlWCcx14AfRd0gjiJhDzUZiX4j63TKIQjkzX0K/CNt+Jvf0RoV75Xukca8wfXcTsS9E/VS+yUn29SaFfXM90r4pglZZZXrEZmJfrhPvjViDL5F/v8H3Ja/Og28g5mYnKBt3C2wPvsJ7C79Hqt8/wNa0sSyf+LeCZnzRPuq6FTrG97JXZ+MLxNzqBGXjdsG28QnvLn6PkjGUbqRnKN3QMZQer86GEojZ7QRl426DbUMR3tv4PVL9ZFoXEasXOjcFmzNzH9UjcS7CFjU4+sAKjnjvnvqgExyveHUOjkDMPicoG3c7bAeH8N7O75FqcDBN9FWkERzMfbSDx7nfuyLLz4F3Evus9RzdcvDuAit4453YdkEneF/z6hy8gZi7nKBs3N2wHbzCeze/R8mc2F5HGsHL3EdvROJchC1qcOwBKzjindj2QCc43vTqHByBmHucoGzcvbAdHMJ7L79HqsHBNNG3kEZwMPfRPh5nlU8syYl8J/ifWPoB3P3N5i1PD3Yo8P4ROvd1E/k69xO1JPa61NLP8qB1AKxBK94J/QB0Bq23vToPWoGYB5ygbNyDsD1oCe+D/B4lc0J/Bzo3BZszcx+9G4lzEbaowXEIrOCId0I/BJ3geM+rc3AEYh5ygrJxD8N2cAjvw/weqQYH00TfRxrBwdxHR3icVU7o8gRmP/gn1Z/A3d9s3vK0aJ8C75+hc1+zT+hHiVoSe11q6Wd50DoG1qAV74R+DDqD1gdenQetQMxjTlA27nHYHrSE93F+j5I5oX8InZuCzZm5jz6KxLkIW9TgOAFWcMQ7oZ+ATnB87NU5OAIxTzhB2bgnYTs4hPdJfo9Ug4Npop8gjeBg7qNTPM4qJ3R5AnMU/JPqL+DubzZveVp0RIH3r9C5r9kn9NNELYm9LrX0szxonQFr0Ip3Qj8DnUHrU6/Og1Yg5hknKBv3LGwPWsL7LL9HyZzQP4POTcHmzNxHn0fiXIQtanCcAys44p3Qz0EnOL7w6hwcgZjnnKBs3POwHRzC+zy/R6rBwTTRL5FGcDD30QUeZ5UTujyBOQ3+SfU3cPc3m7c8LTqlwPt36NzX7BP6RaKWxF6XWvpZHqi/gu28lgy4qHCv/AHbHiF5dUGB959IwyMuEbUk9rpk68feN3I/X1LYNzcQZ98UQausdhE5Xyb25QbhumJ/1zWRf78D7ddenQ+0gZiXnaBs3CuwfaAV3lf4PVL9rmumpoMGXKO//gGokXOuaucAAA== diff --git a/crates/nargo_cli/tests/execution_success/references_aliasing/target/witness.tr b/crates/nargo_cli/tests/execution_success/references_aliasing/target/witness.tr new file mode 100644 index 00000000000..22a1c7fe109 Binary files /dev/null and b/crates/nargo_cli/tests/execution_success/references_aliasing/target/witness.tr differ diff --git a/crates/noirc_evaluator/src/ssa/opt/mem2reg.rs b/crates/noirc_evaluator/src/ssa/opt/mem2reg.rs index d83cda4a8b1..3e63aec63b4 100644 --- a/crates/noirc_evaluator/src/ssa/opt/mem2reg.rs +++ b/crates/noirc_evaluator/src/ssa/opt/mem2reg.rs @@ -1,69 +1,142 @@ -//! mem2reg implements a pass for promoting values stored in memory to values in registers where -//! possible. This is particularly important for converting our memory-based representation of -//! mutable variables into values that are easier to manipulate. -use std::collections::{BTreeMap, HashSet}; +//! The goal of the mem2reg SSA optimization pass is to replace any `Load` instructions to known +//! addresses with the value stored at that address, if it is also known. This pass will also remove +//! any `Store` instructions within a block that are no longer needed because no more loads occur in +//! between the Store in question and the next Store. +//! +//! The pass works as follows: +//! - Each block in each function is iterated in forward-order. +//! - The starting value of each reference in the block is the unification of the same references +//! at the end of each direct predecessor block to the current block. +//! - At each step, the value of each reference is either Known(ValueId) or Unknown. +//! - Two reference values unify to each other if they are exactly equal, or to Unknown otherwise. +//! - If a block has no predecessors, the starting value of each reference is Unknown. +//! - Throughout this pass, aliases of each reference are also tracked. +//! - References typically have 1 alias - themselves. +//! - A reference with multiple aliases means we will not be able to optimize out loads if the +//! reference is stored to. Note that this means we can still optimize out loads if these +//! aliased references are never stored to, or the store occurs after a load. +//! - A reference with 0 aliases means we were unable to find which reference this reference +//! refers to. If such a reference is stored to, we must conservatively invalidate every +//! reference in the current block. +//! +//! From there, to figure out the value of each reference at the end of block, iterate each instruction: +//! - On `Instruction::Allocate`: +//! - Register a new reference was made with itself as its only alias +//! - On `Instruction::Load { address }`: +//! - If `address` is known to only have a single alias (including itself) and if the value of +//! that alias is known, replace the value of the load with the known value. +//! - Furthermore, if the result of the load is a reference, mark the result as an alias +//! of the reference it dereferences to (if known). +//! - If which reference it dereferences to is not known, this load result has no aliases. +//! - On `Instruction::Store { address, value }`: +//! - If the address of the store is `Known(value)`: +//! - If the address has exactly 1 alias: +//! - Set the value of the address to `Known(value)`. +//! - If the address has more than 1 alias: +//! - Set the value of every possible alias to `Unknown`. +//! - If the address has 0 aliases: +//! - Conservatively mark every alias in the block to `Unknown`. +//! - If the address of the store is not known: +//! - Conservatively mark every alias in the block to `Unknown`. +//! - Additionally, if there were no Loads to any alias of the address between this Store and +//! the previous Store to the same address, the previous store can be removed. +//! - On `Instruction::Call { arguments }`: +//! - If any argument of the call is a reference, set the value of each alias of that +//! reference to `Unknown` +//! - Any builtin functions that may return aliases if their input also contains a +//! reference should be tracked. Examples: `slice_push_back`, `slice_insert`, `slice_remove`, etc. +//! +//! On a terminator instruction: +//! - If the terminator is a `Jmp`: +//! - For each reference argument of the jmp, mark the corresponding block parameter it is passed +//! to as an alias for the jmp argument. +//! +//! Finally, if this is the only block in the function, we can remove any Stores that were not +//! referenced by the terminator instruction. +//! +//! Repeating this algorithm for each block in the function in program order should result in +//! optimizing out most known loads. However, identifying all aliases correctly has been proven +//! undecidable in general (Landi, 1992). So this pass will not always optimize out all loads +//! that could theoretically be optimized out. This pass can be performed at any time in the +//! SSA optimization pipeline, although it will be more successful the simpler the program's CFG is. +//! This pass is currently performed several times to enable other passes - most notably being +//! performed before loop unrolling to try to allow for mutable variables used for loop indices. +use std::collections::{BTreeMap, BTreeSet}; use crate::ssa::{ ir::{ basic_block::BasicBlockId, cfg::ControlFlowGraph, - dfg::DataFlowGraph, dom::DominatorTree, function::Function, instruction::{Instruction, InstructionId, TerminatorInstruction}, post_order::PostOrder, - value::{Value, ValueId}, + types::Type, + value::ValueId, }, ssa_gen::Ssa, }; -use super::unrolling::{find_all_loops, Loops}; - impl Ssa { /// Attempts to remove any load instructions that recover values that are already available in /// scope, and attempts to remove stores that are subsequently redundant. - /// As long as they are not stores on memory used inside of loops pub(crate) fn mem2reg(mut self) -> Ssa { for function in self.functions.values_mut() { - let mut all_protected_allocations = HashSet::new(); - let mut context = PerFunctionContext::new(function); - - for block in function.reachable_blocks() { - // Maps Load instruction id -> value to replace the result of the load with - let mut loads_to_substitute_per_block = BTreeMap::new(); - - // Maps Load result id -> value to replace the result of the load with - let mut load_values_to_substitute = BTreeMap::new(); - - let allocations_protected_by_block = context - .analyze_allocations_and_eliminate_known_loads( - &mut function.dfg, - &mut loads_to_substitute_per_block, - &mut load_values_to_substitute, - block, - ); - all_protected_allocations.extend(allocations_protected_by_block.into_iter()); - } - - for block in function.reachable_blocks() { - context.remove_unused_stores(&mut function.dfg, &all_protected_allocations, block); - } + context.mem2reg(function); + context.remove_instructions(function); } self } } struct PerFunctionContext { - last_stores_with_block: BTreeMap<(AllocId, BasicBlockId), ValueId>, - // Maps Load result id -> (value, block_id) - // Used to replace the result of a load with the appropriate block - load_values_to_substitute_per_func: BTreeMap, - store_ids: Vec, cfg: ControlFlowGraph, post_order: PostOrder, dom_tree: DominatorTree, - loops: Loops, + + blocks: BTreeMap, + + /// Load and Store instructions that should be removed at the end of the pass. + /// + /// We avoid removing individual instructions as we go since removing elements + /// from the middle of Vecs many times will be slower than a single call to `retain`. + instructions_to_remove: BTreeSet, +} + +#[derive(Debug, Default, Clone)] +struct Block { + /// Maps a ValueId to the Expression it represents. + /// Multiple ValueIds can map to the same Expression, e.g. + /// dereferences to the same allocation. + expressions: BTreeMap, + + /// Each expression is tracked as to how many aliases it + /// may have. If there is only 1, we can attempt to optimize + /// out any known loads to that alias. Note that "alias" here + /// includes the original reference as well. + aliases: BTreeMap>, + + /// Each allocate instruction result (and some reference block parameters) + /// will map to a Reference value which tracks whether the last value stored + /// to the reference is known. + references: BTreeMap, +} + +/// An `Expression` here is used to represent a canonical key +/// into the aliases map since otherwise two dereferences of the +/// same address will be given different ValueIds. +#[derive(Debug, Clone, PartialOrd, Ord, PartialEq, Eq)] +enum Expression { + Dereference(Box), + Other(ValueId), +} + +/// Every reference's value is either Known and can be optimized away, or Unknown. +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +enum ReferenceValue { + Unknown, + Known(ValueId), } impl PerFunctionContext { @@ -71,308 +144,360 @@ impl PerFunctionContext { 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); + PerFunctionContext { - last_stores_with_block: BTreeMap::new(), - load_values_to_substitute_per_func: BTreeMap::new(), - store_ids: Vec::new(), cfg, post_order, dom_tree, - loops: find_all_loops(function), + blocks: BTreeMap::new(), + instructions_to_remove: BTreeSet::new(), } } -} -/// An AllocId is the ValueId returned from an allocate instruction. E.g. v0 in v0 = allocate. -/// This type alias is used to help signal where the only valid ValueIds are those that are from -/// an allocate instruction. -type AllocId = ValueId; + /// Apply the mem2reg pass to the given function. + /// + /// This function is expected to be the same one that the internal cfg, post_order, and + /// dom_tree were created from. + fn mem2reg(&mut self, function: &mut Function) { + // Iterate each block in reverse post order = forward order + let mut block_order = PostOrder::with_function(function).into_vec(); + block_order.reverse(); + + for block in block_order { + let references = self.find_starting_references(block); + self.analyze_block(function, block, references); + } + } -impl PerFunctionContext { - // 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 analyze_allocations_and_eliminate_known_loads( - &mut self, - dfg: &mut DataFlowGraph, - loads_to_substitute: &mut BTreeMap, - load_values_to_substitute_per_block: &mut BTreeMap, - block_id: BasicBlockId, - ) -> HashSet { - let mut protected_allocations = HashSet::new(); - let block = &dfg[block_id]; - - // Check whether the block has a common successor here to avoid analyzing - // the CFG for every block instruction. - let has_common_successor = self.has_common_successor(block_id); - - for instruction_id in block.instructions() { - match &dfg[*instruction_id] { - Instruction::Store { mut address, value } => { - if has_common_successor { - protected_allocations.insert(address); - } - address = self.fetch_load_value_to_substitute(block_id, address); + /// The value of each reference at the start of the given block is the unification + /// of the value of the same reference at the end of its predecessor blocks. + fn find_starting_references(&mut self, block: BasicBlockId) -> Block { + let mut predecessors = self.cfg.predecessors(block); + + if let Some(first_predecessor) = predecessors.next() { + let first = self.blocks.get(&first_predecessor).cloned().unwrap_or_default(); + + // Note that we have to start folding with the first block as the accumulator. + // If we started with an empty block, an empty block union'd with any other block + // is always also empty so we'd never be able to track any references across blocks. + predecessors.fold(first, |block, predecessor| { + let predecessor = self.blocks.entry(predecessor).or_default(); + block.unify(predecessor) + }) + } else { + Block::default() + } + } - self.last_stores_with_block.insert((address, block_id), *value); - self.store_ids.push(*instruction_id); + /// Analyze a block with the given starting reference values. + /// + /// This will remove any known loads in the block and track the value of references + /// as they are stored to. When this function is finished, the value of each reference + /// at the end of this block will be remembered in `self.blocks`. + fn analyze_block( + &mut self, + function: &mut Function, + block: BasicBlockId, + mut references: Block, + ) { + let instructions = function.dfg[block].instructions().to_vec(); + let mut last_stores = BTreeMap::new(); - if has_common_successor { - protected_allocations.insert(address); - } - } - Instruction::Load { mut address } => { - address = self.fetch_load_value_to_substitute(block_id, address); - - let found_last_value = self.find_load_to_substitute( - block_id, - address, - dfg, - instruction_id, - loads_to_substitute, - load_values_to_substitute_per_block, - ); - if !found_last_value { - // We want to protect allocations that do not have a load to substitute - protected_allocations.insert(address); - // We also want to check for allocations that share a value - // with the one we are protecting. - // This check prevents the pass from removing stores to a value that - // is used by reference aliases in different blocks - protected_allocations.insert(dfg.resolve(address)); - } - } - Instruction::Call { arguments, .. } => { - for arg in arguments { - if Self::value_is_from_allocation(*arg, dfg) { - protected_allocations.insert(*arg); - } - } - } - _ => { - // Nothing to do - } - } + for instruction in instructions { + self.analyze_instruction(function, &mut references, instruction, &mut last_stores); } - // Identify any arrays that are returned from this function - if let TerminatorInstruction::Return { return_values } = block.unwrap_terminator() { - for value in return_values { - if Self::value_is_from_allocation(*value, dfg) { - protected_allocations.insert(*value); - } - } - } + self.handle_terminator(function, block, &mut references, &mut last_stores); - // Substitute load result values - for (result_value, new_value) in load_values_to_substitute_per_block { - let result_value = dfg.resolve(*result_value); - dfg.set_value_from_id(result_value, *new_value); + // If there's only 1 block in the function total, we can remove any remaining last stores + // as well. We can't do this if there are multiple blocks since subsequent blocks may + // reference these stores. + if self.post_order.as_slice().len() == 1 { + self.remove_stores_that_do_not_alias_parameters(function, &references, last_stores); } - // Delete load instructions - // Even though we could let DIE handle this, doing it here makes the debug output easier - // to read. - dfg[block_id] - .instructions_mut() - .retain(|instruction| !loads_to_substitute.contains_key(instruction)); - - protected_allocations + self.blocks.insert(block, references); } - // This method will fetch already saved load values to substitute for a given address. - // The search starts at the block supplied as a parameter. If there is not a load to substitute - // the CFG is analyzed to determine whether a predecessor block has a load value to substitute. - // If there is no load value to substitute the original address is returned. - fn fetch_load_value_to_substitute( + /// Add all instructions in `last_stores` to `self.instructions_to_remove` which do not + /// possibly alias any parameters of the given function. + fn remove_stores_that_do_not_alias_parameters( &mut self, - block_id: BasicBlockId, - address: ValueId, - ) -> ValueId { - let mut stack = vec![block_id]; - let mut visited = HashSet::new(); - - while let Some(block) = stack.pop() { - visited.insert(block); - - // We do not want to substitute loads that take place within loops or yet to be simplified branches - // as this pass can occur before loop unrolling and flattening. - // The mem2reg pass should be ran again following all optimization passes as new values - // may be able to be promoted - for l in self.loops.yet_to_unroll.iter() { - // We do not want to substitute loads that take place within loops as this pass - // can occur before loop unrolling - // The pass should be ran again following loop unrolling as new values - if l.blocks.contains(&block) { - return address; - } - } + function: &Function, + references: &Block, + last_stores: BTreeMap, + ) { + let reference_parameters = function + .parameters() + .iter() + .filter(|param| function.dfg.type_of_value(**param) == Type::Reference) + .collect::>(); - // Check whether there is a load value to substitute in the current block. - // Return the value if found. - if let Some((value, load_block_id)) = - self.load_values_to_substitute_per_func.get(&address) - { - if *load_block_id == block { - return *value; - } - } + for (allocation, instruction) in last_stores { + if let Some(expression) = references.expressions.get(&allocation) { + if let Some(aliases) = references.aliases.get(expression) { + let allocation_aliases_parameter = + aliases.iter().any(|alias| reference_parameters.contains(alias)); - // If no load values to substitute have been found in the current block, check the block's predecessors. - if let Some(predecessor) = self.block_has_predecessor(block, &visited) { - stack.push(predecessor); + if !aliases.is_empty() && !allocation_aliases_parameter { + self.instructions_to_remove.insert(instruction); + } + } + } else { + self.instructions_to_remove.insert(instruction); } } - address } - // This method determines which loads should be substituted and saves them - // to be substituted further in the pass. - // Starting at the block supplied as a parameter, we check whether a store has occurred with the given address. - // If no store has occurred in the supplied block, the CFG is analyzed to determine whether - // a predecessor block has a store at the given address. - fn find_load_to_substitute( + fn analyze_instruction( &mut self, - block_id: BasicBlockId, - address: ValueId, - dfg: &DataFlowGraph, - instruction_id: &InstructionId, - loads_to_substitute: &mut BTreeMap, - load_values_to_substitute_per_block: &mut BTreeMap, - ) -> bool { - let mut stack = vec![block_id]; - let mut visited = HashSet::new(); - - while let Some(block) = stack.pop() { - visited.insert(block); - - // We do not want to substitute loads that take place within loops or yet to be simplified branches - // as this pass can occur before loop unrolling and flattening. - // The mem2reg pass should be ran again following all optimization passes as new values - // may be able to be promoted - for l in self.loops.yet_to_unroll.iter() { - // We do not want to substitute loads that take place within loops as this pass - // can occur before loop unrolling - // The pass should be ran again following loop unrolling as new values - if l.blocks.contains(&block) { - return false; + function: &mut Function, + references: &mut Block, + instruction: InstructionId, + last_stores: &mut BTreeMap, + ) { + match &function.dfg[instruction] { + Instruction::Load { address } => { + let address = function.dfg.resolve(*address); + + let result = function.dfg.instruction_results(instruction)[0]; + references.remember_dereference(function, address, result); + + // If the load is known, replace it with the known value and remove the load + if let Some(value) = references.get_known_value(address) { + let result = function.dfg.instruction_results(instruction)[0]; + function.dfg.set_value_from_id(result, value); + + self.instructions_to_remove.insert(instruction); + } else { + last_stores.remove(&address); } } + Instruction::Store { address, value } => { + let address = function.dfg.resolve(*address); + let value = function.dfg.resolve(*value); + + // If there was another store to this instruction without any (unremoved) loads or + // function calls in-between, we can remove the previous store. + if let Some(last_store) = last_stores.get(&address) { + self.instructions_to_remove.insert(*last_store); + } - // Check whether there has been a store instruction in the current block - // If there has been a store, add a load to be substituted. - if let Some(last_value) = self.last_stores_with_block.get(&(address, block)) { - let result_value = *dfg - .instruction_results(*instruction_id) - .first() - .expect("ICE: Load instructions should have single result"); - - loads_to_substitute.insert(*instruction_id, *last_value); - load_values_to_substitute_per_block.insert(result_value, *last_value); - self.load_values_to_substitute_per_func.insert(result_value, (*last_value, block)); - return true; + references.set_known_value(address, value, last_stores); + last_stores.insert(address, instruction); } - - // If no load values to substitute have been found in the current block, check the block's predecessors. - if let Some(predecessor) = self.block_has_predecessor(block, &visited) { - stack.push(predecessor); + Instruction::Call { arguments, .. } => { + self.mark_all_unknown(arguments, function, references, last_stores); } + Instruction::Allocate => { + // Register the new reference + let result = function.dfg.instruction_results(instruction)[0]; + references.expressions.insert(result, Expression::Other(result)); + + let mut aliases = BTreeSet::new(); + aliases.insert(result); + references.aliases.insert(Expression::Other(result), aliases); + } + _ => (), } - false } - // If no loads or stores have been found in the current block, check the block's predecessors. - // Only check blocks with one predecessor as otherwise a constant value could be propogated - // through successor blocks with multiple branches that rely on other simplification passes - // such as loop unrolling or flattening of the CFG. - fn block_has_predecessor( - &mut self, - block_id: BasicBlockId, - visited: &HashSet, - ) -> Option { - let mut predecessors = self.cfg.predecessors(block_id); - if predecessors.len() == 1 { - let predecessor = predecessors.next().unwrap(); - if self.dom_tree.is_reachable(predecessor) - && self.dom_tree.dominates(predecessor, block_id) - && !visited.contains(&predecessor) - { - return Some(predecessor); + fn mark_all_unknown( + &self, + values: &[ValueId], + function: &Function, + references: &mut Block, + last_stores: &mut BTreeMap, + ) { + for value in values { + if function.dfg.type_of_value(*value) == Type::Reference { + let value = function.dfg.resolve(*value); + references.set_unknown(value, last_stores); + last_stores.remove(&value); } } - None } - fn has_common_successor(&mut self, block_id: BasicBlockId) -> bool { - let mut predecessors = self.cfg.predecessors(block_id); - if let Some(predecessor) = predecessors.next() { - let pred_successors = self.find_all_successors(predecessor); - let current_successors: HashSet<_> = self.cfg.successors(block_id).collect(); - return pred_successors.into_iter().any(|b| current_successors.contains(&b)); + /// Remove any instructions in `self.instructions_to_remove` from the current function. + /// This is expected to contain any loads which were replaced and any stores which are + /// no longer needed. + fn remove_instructions(&self, function: &mut Function) { + // The order we iterate blocks in is not important + for block in self.post_order.as_slice() { + function.dfg[*block] + .instructions_mut() + .retain(|instruction| !self.instructions_to_remove.contains(instruction)); } - false } - fn find_all_successors(&self, block_id: BasicBlockId) -> HashSet { - let mut stack = vec![]; - let mut visited = HashSet::new(); - - // Fetch initial block successors - let successors = self.cfg.successors(block_id); - for successor in successors { - if !visited.contains(&successor) { - stack.push(successor); + fn handle_terminator( + &self, + function: &mut Function, + block: BasicBlockId, + references: &mut Block, + last_stores: &mut BTreeMap, + ) { + match function.dfg[block].unwrap_terminator() { + TerminatorInstruction::JmpIf { .. } => (), // Nothing to do + TerminatorInstruction::Jmp { destination, arguments, .. } => { + let destination_parameters = function.dfg[*destination].parameters(); + assert_eq!(destination_parameters.len(), arguments.len()); + + // Add an alias for each reference parameter + for (parameter, argument) in destination_parameters.iter().zip(arguments) { + if function.dfg.type_of_value(*parameter) == Type::Reference { + let argument = function.dfg.resolve(*argument); + + if let Some(expression) = references.expressions.get(&argument) { + if let Some(aliases) = references.aliases.get_mut(expression) { + // The argument reference is possibly aliased by this block parameter + aliases.insert(*parameter); + } + } + } + } + } + TerminatorInstruction::Return { return_values } => { + // Removing all `last_stores` for each returned reference is more important here + // than setting them all to ReferenceValue::Unknown since no other block should + // have a block with a Return terminator as a predecessor anyway. + self.mark_all_unknown(return_values, function, references, last_stores); } } + } +} - // Follow the CFG to fetch the remaining successors - while let Some(block) = stack.pop() { - visited.insert(block); - let successors = self.cfg.successors(block); - for successor in successors { - if !visited.contains(&successor) { - stack.push(successor); +impl Block { + /// If the given reference id points to a known value, return the value + fn get_known_value(&self, address: ValueId) -> Option { + if let Some(expression) = self.expressions.get(&address) { + if let Some(aliases) = self.aliases.get(expression) { + // We could allow multiple aliases if we check that the reference + // value in each is equal. + if aliases.len() == 1 { + let alias = aliases.first().expect("There should be exactly 1 alias"); + + if let Some(ReferenceValue::Known(value)) = self.references.get(alias) { + return Some(*value); + } } } } - visited + None + } + + /// If the given address is known, set its value to `ReferenceValue::Known(value)`. + fn set_known_value( + &mut self, + address: ValueId, + value: ValueId, + last_stores: &mut BTreeMap, + ) { + self.set_value(address, ReferenceValue::Known(value), last_stores); + } + + fn set_unknown( + &mut self, + address: ValueId, + last_stores: &mut BTreeMap, + ) { + self.set_value(address, ReferenceValue::Unknown, last_stores); } - /// Checks whether the given value id refers to an allocation. - fn value_is_from_allocation(value: ValueId, dfg: &DataFlowGraph) -> bool { - match &dfg[value] { - Value::Instruction { instruction, .. } => { - matches!(&dfg[*instruction], Instruction::Allocate) + fn set_value( + &mut self, + address: ValueId, + value: ReferenceValue, + last_stores: &mut BTreeMap, + ) { + if let Some(expression) = self.expressions.get(&address) { + if let Some(aliases) = self.aliases.get(expression) { + if aliases.is_empty() { + // uh-oh, we don't know at all what this reference refers to, could be anything. + // Now we have to invalidate every reference we know of + self.invalidate_all_references(last_stores); + } else if aliases.len() == 1 { + let alias = aliases.first().expect("There should be exactly 1 alias"); + self.references.insert(*alias, value); + } else { + // More than one alias. We're not sure which it refers to so we have to + // conservatively invalidate all references it may refer to. + for alias in aliases.iter() { + if let Some(reference_value) = self.references.get_mut(alias) { + *reference_value = ReferenceValue::Unknown; + } + } + } } - _ => false, } } - /// Removes all store instructions identified during analysis that aren't present in the - /// provided `protected_allocations` `HashSet`. - fn remove_unused_stores( - &self, - dfg: &mut DataFlowGraph, - protected_allocations: &HashSet, - block_id: BasicBlockId, - ) { - // Scan for unused stores - let mut stores_to_remove = HashSet::new(); + fn invalidate_all_references(&mut self, last_stores: &mut BTreeMap) { + for reference_value in self.references.values_mut() { + *reference_value = ReferenceValue::Unknown; + } - 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"), - }; + last_stores.clear(); + } + + fn unify(mut self, other: &Self) -> Self { + for (value_id, expression) in &other.expressions { + if let Some(existing) = self.expressions.get(value_id) { + assert_eq!(existing, expression, "Expected expressions for {value_id} to be equal"); + } else { + self.expressions.insert(*value_id, expression.clone()); + } + } + + for (expression, new_aliases) in &other.aliases { + let expression = expression.clone(); + + self.aliases + .entry(expression) + .and_modify(|aliases| { + for alias in new_aliases { + aliases.insert(*alias); + } + }) + .or_insert_with(|| new_aliases.clone()); + } - if !protected_allocations.contains(&address) { - stores_to_remove.insert(*instruction_id); + // Keep only the references present in both maps. + let mut intersection = BTreeMap::new(); + for (value_id, reference) in &other.references { + if let Some(existing) = self.references.get(value_id) { + intersection.insert(*value_id, existing.unify(*reference)); } } + self.references = intersection; - // Delete unused stores - dfg[block_id] - .instructions_mut() - .retain(|instruction| !stores_to_remove.contains(instruction)); + self + } + + /// Remember that `result` is the result of dereferencing `address`. This is important to + /// track aliasing when references are stored within other references. + fn remember_dereference(&mut self, function: &Function, address: ValueId, result: ValueId) { + if function.dfg.type_of_value(result) == Type::Reference { + if let Some(known_address) = self.get_known_value(address) { + self.expressions.insert(result, Expression::Other(known_address)); + } else { + let expression = Expression::Dereference(Box::new(Expression::Other(address))); + self.expressions.insert(result, expression); + // No known aliases to insert for this expression... can we find an alias + // even if we don't have a known address? If not we'll have to invalidate all + // known references if this reference is ever stored to. + } + } + } +} + +impl ReferenceValue { + fn unify(self, other: Self) -> Self { + if self == other { + self + } else { + ReferenceValue::Unknown + } } } @@ -423,8 +548,6 @@ mod tests { let ssa = builder.finish().mem2reg().fold_constants(); - println!("{ssa}"); - let func = ssa.main(); let block_id = func.entry_block(); @@ -495,7 +618,7 @@ mod tests { let func = ssa.main(); let block_id = func.entry_block(); - // Store affects outcome of returned array, and can't be removed + // Store is needed by the return value, and can't be removed assert_eq!(count_stores(block_id, &func.dfg), 1); let ret_val_id = match func.dfg[block_id].terminator().unwrap() { @@ -562,11 +685,13 @@ mod tests { assert_eq!(ssa.main().reachable_blocks().len(), 2); // Expected result: - // fn main { + // acir fn main f0 { // b0(): // v0 = allocate + // store Field 5 at v0 // jmp b1(Field 5) // b1(v3: Field): + // store Field 6 at v0 // return v3, Field 5, Field 6 // Optimized to constants 5 and 6 // } let ssa = ssa.mem2reg(); @@ -578,9 +703,9 @@ mod tests { assert_eq!(count_loads(main.entry_block(), &main.dfg), 0); assert_eq!(count_loads(b1, &main.dfg), 0); - // All stores should be removed - assert_eq!(count_stores(main.entry_block(), &main.dfg), 0); - assert_eq!(count_stores(b1, &main.dfg), 0); + // Neither store is removed since they are each the last in the block and there are multiple blocks + 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() { @@ -597,7 +722,7 @@ mod tests { // Test that a load in a predecessor block has been removed if the value // is later stored in a successor block #[test] - fn store_with_load_in_predecessor_block() { + fn load_aliases_in_predecessor_block() { // fn main { // b0(): // v0 = allocate @@ -647,14 +772,17 @@ mod tests { assert_eq!(ssa.main().reachable_blocks().len(), 2); // Expected result: - // fn main { - // b0(): - // v0 = allocate - // v2 = allocate - // jmp b1() - // b1(): - // v8 = eq Field 2, Field 2 - // return + // acir fn main f0 { + // b0(): + // v0 = allocate + // store Field 0 at v0 + // v2 = allocate + // store v0 at v2 + // jmp b1() + // b1(): + // store Field 2 at v0 + // v8 = eq Field 1, Field 2 + // return // } let ssa = ssa.mem2reg(); @@ -665,13 +793,15 @@ mod tests { assert_eq!(count_loads(main.entry_block(), &main.dfg), 0); assert_eq!(count_loads(b1, &main.dfg), 0); - // All stores should be removed - assert_eq!(count_stores(main.entry_block(), &main.dfg), 0); - assert_eq!(count_stores(b1, &main.dfg), 0); + // Only the first store in b1 is removed since there is another store to the same reference + // in the same block, and the store is not needed before the later store. + assert_eq!(count_stores(main.entry_block(), &main.dfg), 2); + assert_eq!(count_stores(b1, &main.dfg), 1); let b1_instructions = main.dfg[b1].instructions(); - // The first instruction should be a binary operation - match &main.dfg[b1_instructions[0]] { + + // The last instruction in b1 should be a binary operation + match &main.dfg[*b1_instructions.last().unwrap()] { Instruction::Binary(binary) => { let lhs = main.dfg.get_numeric_constant(binary.lhs).expect("Expected constant value"); diff --git a/crates/noirc_evaluator/src/ssa/opt/unrolling.rs b/crates/noirc_evaluator/src/ssa/opt/unrolling.rs index 9bf62bda1cb..f9c92f8335e 100644 --- a/crates/noirc_evaluator/src/ssa/opt/unrolling.rs +++ b/crates/noirc_evaluator/src/ssa/opt/unrolling.rs @@ -44,7 +44,7 @@ impl Ssa { } } -pub(crate) struct Loop { +struct Loop { /// The header block of a loop is the block which dominates all the /// other blocks in the loop. header: BasicBlockId, @@ -57,12 +57,12 @@ pub(crate) struct Loop { pub(crate) blocks: HashSet, } -pub(crate) struct Loops { +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, - pub(crate) yet_to_unroll: Vec, + yet_to_unroll: Vec, modified_blocks: HashSet, cfg: ControlFlowGraph, dom_tree: DominatorTree, @@ -70,7 +70,7 @@ pub(crate) struct Loops { /// 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. -pub(crate) fn find_all_loops(function: &Function) -> Loops { +fn find_all_loops(function: &Function) -> Loops { let cfg = ControlFlowGraph::with_function(function); let post_order = PostOrder::with_function(function); let mut dom_tree = DominatorTree::with_cfg_and_post_order(&cfg, &post_order); diff --git a/crates/noirc_frontend/src/monomorphization/mod.rs b/crates/noirc_frontend/src/monomorphization/mod.rs index 2ef980176d3..fb441a1bc17 100644 --- a/crates/noirc_frontend/src/monomorphization/mod.rs +++ b/crates/noirc_frontend/src/monomorphization/mod.rs @@ -882,8 +882,8 @@ impl<'interner> Monomorphizer<'interner> { } } let printable_type: PrintableType = typ.into(); - let abi_as_string = - serde_json::to_string(&printable_type).expect("ICE: expected PrintableType to serialize"); + let abi_as_string = serde_json::to_string(&printable_type) + .expect("ICE: expected PrintableType to serialize"); arguments.push(ast::Expression::Literal(ast::Literal::Str(abi_as_string))); }