From 6ad5af8c0d76b8340dca33b7f1adeb40f69c8aa8 Mon Sep 17 00:00:00 2001 From: Ary Borenszweig Date: Fri, 11 Apr 2025 13:55:10 -0300 Subject: [PATCH 1/5] chore: don't use `set_value_from_id` in `remove_truncate_after_range_checks` --- compiler/noirc_evaluator/src/ssa/ir/dfg.rs | 42 +++++++++++++++++++ .../opt/remove_truncate_after_range_check.rs | 6 +-- 2 files changed, 44 insertions(+), 4 deletions(-) diff --git a/compiler/noirc_evaluator/src/ssa/ir/dfg.rs b/compiler/noirc_evaluator/src/ssa/ir/dfg.rs index 690bb121634..9b5a9f37b53 100644 --- a/compiler/noirc_evaluator/src/ssa/ir/dfg.rs +++ b/compiler/noirc_evaluator/src/ssa/ir/dfg.rs @@ -383,6 +383,48 @@ impl DataFlowGraph { self.instructions[id] = instruction; } + /// Replaces all values in the given block with the values in the given map. + /// + /// This method should be preferred over `set_value_from_id` which might eventually be removed. + pub(crate) fn replace_values_in_block( + &mut self, + block: BasicBlockId, + values_to_replace: &std::collections::HashMap, + ) { + if values_to_replace.is_empty() { + return; + } + + let replacement_fn = |value_id| { + if let Some(replacement_id) = values_to_replace.get(&value_id) { + *replacement_id + } else { + value_id + } + }; + + // Replace in all the block's instructions + for instruction_id in self.blocks[block].instructions() { + let instruction = &mut self.instructions[*instruction_id]; + instruction.map_values_mut(replacement_fn); + + // Make sure we also replace the instruction results + let results = self.results.get_mut(instruction_id); + if let Some(results) = results { + for result in results { + if let Some(replacement_id) = values_to_replace.get(result) { + *result = *replacement_id; + } + } + } + } + + // Finally, the value might show up in a terminator + if self[block].terminator().is_some() { + self[block].unwrap_terminator_mut().map_values_mut(replacement_fn); + } + } + /// Set the value of value_to_replace to refer to the value referred to by new_value. /// /// This is the preferred method to call for optimizations simplifying diff --git a/compiler/noirc_evaluator/src/ssa/opt/remove_truncate_after_range_check.rs b/compiler/noirc_evaluator/src/ssa/opt/remove_truncate_after_range_check.rs index 8b8d6a05dc8..241fc9646ed 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/remove_truncate_after_range_check.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/remove_truncate_after_range_check.rs @@ -24,7 +24,7 @@ impl Function { // Keeps the minimum bit size a value was range-checked against let mut range_checks: HashMap = HashMap::new(); let mut instructions_to_remove = HashSet::new(); - let mut values_to_replace = HashMap::::new(); + let mut values_to_replace = std::collections::HashMap::::new(); for instruction_id in self.dfg[block].instructions() { let instruction = &self.dfg[*instruction_id]; @@ -70,9 +70,7 @@ impl Function { .instructions_mut() .retain(|instruction| !instructions_to_remove.contains(instruction)); - for (old_value, new_value) in values_to_replace { - self.dfg.set_value_from_id(old_value, new_value); - } + self.dfg.replace_values_in_block(block, &values_to_replace); } } } From 3a136fd7d234e870f0b5f8583b51021a60b14fc2 Mon Sep 17 00:00:00 2001 From: Ary Borenszweig Date: Fri, 11 Apr 2025 14:10:32 -0300 Subject: [PATCH 2/5] Use FxHashMap --- compiler/noirc_evaluator/src/ssa/ir/dfg.rs | 2 +- .../src/ssa/opt/remove_truncate_after_range_check.rs | 9 +++++---- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/compiler/noirc_evaluator/src/ssa/ir/dfg.rs b/compiler/noirc_evaluator/src/ssa/ir/dfg.rs index 9b5a9f37b53..180ab8c525d 100644 --- a/compiler/noirc_evaluator/src/ssa/ir/dfg.rs +++ b/compiler/noirc_evaluator/src/ssa/ir/dfg.rs @@ -389,7 +389,7 @@ impl DataFlowGraph { pub(crate) fn replace_values_in_block( &mut self, block: BasicBlockId, - values_to_replace: &std::collections::HashMap, + values_to_replace: &HashMap, ) { if values_to_replace.is_empty() { return; diff --git a/compiler/noirc_evaluator/src/ssa/opt/remove_truncate_after_range_check.rs b/compiler/noirc_evaluator/src/ssa/opt/remove_truncate_after_range_check.rs index 241fc9646ed..6e70e44a4a5 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/remove_truncate_after_range_check.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/remove_truncate_after_range_check.rs @@ -1,4 +1,5 @@ -use im::{HashMap, HashSet}; +use fxhash::FxHashMap as HashMap; +use fxhash::FxHashSet as HashSet; use crate::ssa::{ ir::{function::Function, instruction::Instruction, value::ValueId}, @@ -22,9 +23,9 @@ impl Function { pub(crate) fn remove_truncate_after_range_check(&mut self) { for block in self.reachable_blocks() { // Keeps the minimum bit size a value was range-checked against - let mut range_checks: HashMap = HashMap::new(); - let mut instructions_to_remove = HashSet::new(); - let mut values_to_replace = std::collections::HashMap::::new(); + let mut range_checks: HashMap = HashMap::default(); + let mut instructions_to_remove = HashSet::default(); + let mut values_to_replace = HashMap::::default(); for instruction_id in self.dfg[block].instructions() { let instruction = &self.dfg[*instruction_id]; From 0e4b91207df734cbc850e0c0dd023a9ca6f3ae2a Mon Sep 17 00:00:00 2001 From: Ary Borenszweig Date: Fri, 11 Apr 2025 14:41:11 -0300 Subject: [PATCH 3/5] Replace in all blocks --- compiler/noirc_evaluator/src/ssa/ir/dfg.rs | 38 ++++++++++--------- .../opt/remove_truncate_after_range_check.rs | 12 ++++-- 2 files changed, 28 insertions(+), 22 deletions(-) diff --git a/compiler/noirc_evaluator/src/ssa/ir/dfg.rs b/compiler/noirc_evaluator/src/ssa/ir/dfg.rs index 180ab8c525d..0c7427dc228 100644 --- a/compiler/noirc_evaluator/src/ssa/ir/dfg.rs +++ b/compiler/noirc_evaluator/src/ssa/ir/dfg.rs @@ -383,12 +383,12 @@ impl DataFlowGraph { self.instructions[id] = instruction; } - /// Replaces all values in the given block with the values in the given map. + /// Replaces all values in the given blocks with the values in the given map. /// /// This method should be preferred over `set_value_from_id` which might eventually be removed. - pub(crate) fn replace_values_in_block( + pub(crate) fn replace_values_in_blocks( &mut self, - block: BasicBlockId, + blocks: impl Iterator, values_to_replace: &HashMap, ) { if values_to_replace.is_empty() { @@ -403,25 +403,27 @@ impl DataFlowGraph { } }; - // Replace in all the block's instructions - for instruction_id in self.blocks[block].instructions() { - let instruction = &mut self.instructions[*instruction_id]; - instruction.map_values_mut(replacement_fn); - - // Make sure we also replace the instruction results - let results = self.results.get_mut(instruction_id); - if let Some(results) = results { - for result in results { - if let Some(replacement_id) = values_to_replace.get(result) { - *result = *replacement_id; + for block in blocks { + // Replace in all the block's instructions + for instruction_id in self.blocks[block].instructions() { + let instruction = &mut self.instructions[*instruction_id]; + instruction.map_values_mut(replacement_fn); + + // Make sure we also replace the instruction results + let results = self.results.get_mut(instruction_id); + if let Some(results) = results { + for result in results { + if let Some(replacement_id) = values_to_replace.get(result) { + *result = *replacement_id; + } } } } - } - // Finally, the value might show up in a terminator - if self[block].terminator().is_some() { - self[block].unwrap_terminator_mut().map_values_mut(replacement_fn); + // Finally, the value might show up in a terminator + if self[block].terminator().is_some() { + self[block].unwrap_terminator_mut().map_values_mut(replacement_fn); + } } } diff --git a/compiler/noirc_evaluator/src/ssa/opt/remove_truncate_after_range_check.rs b/compiler/noirc_evaluator/src/ssa/opt/remove_truncate_after_range_check.rs index 6e70e44a4a5..9074585a9d8 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/remove_truncate_after_range_check.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/remove_truncate_after_range_check.rs @@ -21,11 +21,15 @@ impl Ssa { impl Function { pub(crate) fn remove_truncate_after_range_check(&mut self) { - for block in self.reachable_blocks() { + let mut values_to_replace = HashMap::::default(); + + let blocks = self.reachable_blocks(); + for block in &blocks { + let block = *block; + // Keeps the minimum bit size a value was range-checked against let mut range_checks: HashMap = HashMap::default(); let mut instructions_to_remove = HashSet::default(); - let mut values_to_replace = HashMap::::default(); for instruction_id in self.dfg[block].instructions() { let instruction = &self.dfg[*instruction_id]; @@ -70,9 +74,9 @@ impl Function { self.dfg[block] .instructions_mut() .retain(|instruction| !instructions_to_remove.contains(instruction)); - - self.dfg.replace_values_in_block(block, &values_to_replace); } + + self.dfg.replace_values_in_blocks(blocks.into_iter(), &values_to_replace); } } From ad73aa4856e3ca63620bb6b0db14422dda169fd4 Mon Sep 17 00:00:00 2001 From: Ary Borenszweig Date: Fri, 11 Apr 2025 14:51:11 -0300 Subject: [PATCH 4/5] Make sure values are replaced across blocks, and fix optimization --- .../src/ssa/opt/remove_truncate_after_range_check.rs | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/compiler/noirc_evaluator/src/ssa/opt/remove_truncate_after_range_check.rs b/compiler/noirc_evaluator/src/ssa/opt/remove_truncate_after_range_check.rs index 9074585a9d8..aeee55525f2 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/remove_truncate_after_range_check.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/remove_truncate_after_range_check.rs @@ -22,13 +22,12 @@ impl Ssa { impl Function { pub(crate) fn remove_truncate_after_range_check(&mut self) { let mut values_to_replace = HashMap::::default(); + // Keeps the minimum bit size a value was range-checked against + let mut range_checks: HashMap = HashMap::default(); let blocks = self.reachable_blocks(); for block in &blocks { let block = *block; - - // Keeps the minimum bit size a value was range-checked against - let mut range_checks: HashMap = HashMap::default(); let mut instructions_to_remove = HashSet::default(); for instruction_id in self.dfg[block].instructions() { @@ -94,6 +93,8 @@ mod tests { b0(v0: Field): range_check v0 to 64 bits // This is to make sure we keep the smallest one range_check v0 to 32 bits + jmp b1() + b1(): v1 = truncate v0 to 32 bits, max_bit_size: 254 return v1 } @@ -106,6 +107,8 @@ mod tests { b0(v0: Field): range_check v0 to 64 bits range_check v0 to 32 bits + jmp b1() + b1(): return v0 } "); From b49a9bbdc339c76d16204094aace2102c68a33cc Mon Sep 17 00:00:00 2001 From: Ary Borenszweig Date: Fri, 11 Apr 2025 14:54:01 -0300 Subject: [PATCH 5/5] Add a comment --- .../src/ssa/opt/remove_truncate_after_range_check.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/compiler/noirc_evaluator/src/ssa/opt/remove_truncate_after_range_check.rs b/compiler/noirc_evaluator/src/ssa/opt/remove_truncate_after_range_check.rs index aeee55525f2..9273453d6cb 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/remove_truncate_after_range_check.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/remove_truncate_after_range_check.rs @@ -93,7 +93,7 @@ mod tests { b0(v0: Field): range_check v0 to 64 bits // This is to make sure we keep the smallest one range_check v0 to 32 bits - jmp b1() + jmp b1() // Make sure the optimization is applied across blocks b1(): v1 = truncate v0 to 32 bits, max_bit_size: 254 return v1