From 63d84a30fcbc117443cd3b404e872cb3c2f26111 Mon Sep 17 00:00:00 2001 From: jfecher Date: Tue, 16 May 2023 07:12:14 -0400 Subject: [PATCH] chore(ssa refactor): Add basic instruction simplification (#1329) * Add basic instruction simplification * Cargo fmt * Add comments --- .../src/ssa_refactor/ir/dfg.rs | 70 +++++- .../src/ssa_refactor/ir/instruction.rs | 202 +++++++++++++++++- .../src/ssa_refactor/opt/inlining.rs | 19 +- .../src/ssa_refactor/ssa_builder/mod.rs | 19 +- 4 files changed, 285 insertions(+), 25 deletions(-) diff --git a/crates/noirc_evaluator/src/ssa_refactor/ir/dfg.rs b/crates/noirc_evaluator/src/ssa_refactor/ir/dfg.rs index 3ab345f06b..fc15f3e216 100644 --- a/crates/noirc_evaluator/src/ssa_refactor/ir/dfg.rs +++ b/crates/noirc_evaluator/src/ssa_refactor/ir/dfg.rs @@ -100,15 +100,21 @@ impl DataFlowGraph { id } - /// Replace an instruction id with another. - /// - /// This function should generally be avoided if possible in favor of inserting new - /// instructions since it does not check whether the instruction results of the removed - /// instruction are still in use. Users of this function thus need to ensure the old - /// instruction's results are no longer in use or are otherwise compatible with the - /// new instruction's result count and types. - pub(crate) fn replace_instruction(&mut self, id: Id, instruction: Instruction) { - self.instructions[id] = instruction; + /// Inserts a new instruction at the end of the given block and returns its results + pub(crate) fn insert_instruction( + &mut self, + instruction: Instruction, + block: BasicBlockId, + ctrl_typevars: Option>, + ) -> InsertInstructionResult { + match instruction.simplify(self) { + Some(simplification) => InsertInstructionResult::SimplifiedTo(simplification), + None => { + let id = self.make_instruction(instruction, ctrl_typevars); + self.insert_instruction_in_block(block, id); + InsertInstructionResult::Results(self.instruction_results(id)) + } + } } /// Insert a value into the dfg's storage and return an id to reference it. @@ -300,6 +306,52 @@ impl std::ops::IndexMut for DataFlowGraph { } } +// The result of calling DataFlowGraph::insert_instruction can +// be a list of results or a single ValueId if the instruction was simplified +// to an existing value. +pub(crate) enum InsertInstructionResult<'dfg> { + Results(&'dfg [ValueId]), + SimplifiedTo(ValueId), + InstructionRemoved, +} + +impl<'dfg> InsertInstructionResult<'dfg> { + /// Retrieve the first (and expected to be the only) result. + pub(crate) fn first(&self) -> ValueId { + match self { + InsertInstructionResult::SimplifiedTo(value) => *value, + InsertInstructionResult::Results(results) => results[0], + InsertInstructionResult::InstructionRemoved => { + panic!("Instruction was removed, no results") + } + } + } + + /// Return all the results contained in the internal results array. + /// This is used for instructions returning multiple results that were + /// not simplified - like function calls. + pub(crate) fn results(&self) -> &'dfg [ValueId] { + match self { + InsertInstructionResult::Results(results) => results, + InsertInstructionResult::SimplifiedTo(_) => { + panic!("InsertInstructionResult::results called on a simplified instruction") + } + InsertInstructionResult::InstructionRemoved => { + panic!("InsertInstructionResult::results called on a removed instruction") + } + } + } + + /// Returns the amount of ValueIds contained + pub(crate) fn len(&self) -> usize { + match self { + InsertInstructionResult::SimplifiedTo(_) => 1, + InsertInstructionResult::Results(results) => results.len(), + InsertInstructionResult::InstructionRemoved => 0, + } + } +} + #[cfg(test)] mod tests { use super::DataFlowGraph; diff --git a/crates/noirc_evaluator/src/ssa_refactor/ir/instruction.rs b/crates/noirc_evaluator/src/ssa_refactor/ir/instruction.rs index 812d12b23a..42968568de 100644 --- a/crates/noirc_evaluator/src/ssa_refactor/ir/instruction.rs +++ b/crates/noirc_evaluator/src/ssa_refactor/ir/instruction.rs @@ -1,7 +1,13 @@ -use acvm::acir::BlackBoxFunc; +use acvm::{acir::BlackBoxFunc, FieldElement}; use iter_extended::vecmap; -use super::{basic_block::BasicBlockId, map::Id, types::Type, value::ValueId}; +use super::{ + basic_block::BasicBlockId, + dfg::DataFlowGraph, + map::Id, + types::Type, + value::{Value, ValueId}, +}; /// Reference to an instruction /// @@ -151,6 +157,48 @@ 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 { + match self { + Instruction::Binary(binary) => binary.simplify(dfg), + Instruction::Cast(value, typ) => (*typ == dfg.type_of_value(*value)).then_some(*value), + Instruction::Not(value) => { + match &dfg[*value] { + // Limit optimizing ! on constants to only booleans. If we tried it on fields, + // there is no Not on FieldElement, so we'd need to convert between u128. This + // 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())) + } + Value::Instruction { instruction, .. } => { + // !!v => v + match &dfg[*instruction] { + Instruction::Not(value) => Some(*value), + _ => None, + } + } + _ => None, + } + } + 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)); + } + } + None + } + Instruction::Truncate { .. } => None, + Instruction::Call { .. } => None, + Instruction::Allocate { .. } => None, + Instruction::Load { .. } => None, + Instruction::Store { .. } => None, + } + } } /// The possible return values for Instruction::return_types @@ -219,6 +267,156 @@ impl Binary { _ => InstructionResultType::Operand(self.lhs), } } + + /// Try to simplify this binary instruction, returning the new value if possible. + fn simplify(&self, dfg: &mut DataFlowGraph) -> Option { + 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); + } + + let lhs_is_zero = lhs.map_or(false, |lhs| lhs.is_zero()); + let rhs_is_zero = rhs.map_or(false, |rhs| rhs.is_zero()); + + let lhs_is_one = lhs.map_or(false, |lhs| lhs.is_one()); + let rhs_is_one = rhs.map_or(false, |rhs| rhs.is_one()); + + match self.operator { + BinaryOp::Add => { + if lhs_is_zero { + return Some(self.rhs); + } + if rhs_is_zero { + return Some(self.lhs); + } + } + BinaryOp::Sub => { + if rhs_is_zero { + return Some(self.lhs); + } + } + BinaryOp::Mul => { + if lhs_is_one { + return Some(self.rhs); + } + if rhs_is_one { + return Some(self.lhs); + } + } + BinaryOp::Div => { + if rhs_is_one { + return Some(self.lhs); + } + } + BinaryOp::Mod => { + if rhs_is_one { + return Some(self.lhs); + } + } + BinaryOp::Eq => { + if self.lhs == self.rhs { + return Some(dfg.make_constant(FieldElement::one(), Type::bool())); + } + } + BinaryOp::Lt => { + if self.lhs == self.rhs { + return Some(dfg.make_constant(FieldElement::zero(), Type::bool())); + } + } + BinaryOp::And => { + if lhs_is_zero || rhs_is_zero { + return Some(dfg.make_constant(FieldElement::zero(), operand_type)); + } + } + BinaryOp::Or => { + if lhs_is_zero { + return Some(self.rhs); + } + if rhs_is_zero { + return Some(self.lhs); + } + } + BinaryOp::Xor => (), + BinaryOp::Shl => { + if rhs_is_zero { + return Some(self.lhs); + } + } + BinaryOp::Shr => { + if rhs_is_zero { + return Some(self.lhs); + } + } + } + None + } + + /// Evaluate the two constants with the operation specified by self.operator. + /// Pushes the resulting value to the given DataFlowGraph's constants and returns it. + fn eval_constants( + &self, + dfg: &mut DataFlowGraph, + lhs: FieldElement, + rhs: FieldElement, + operand_type: Type, + ) -> Option> { + let value = match self.operator { + BinaryOp::Add => lhs + rhs, + BinaryOp::Sub => lhs - rhs, + BinaryOp::Mul => lhs * rhs, + BinaryOp::Div => lhs / rhs, + BinaryOp::Eq => (lhs == rhs).into(), + BinaryOp::Lt => (lhs < rhs).into(), + + // The rest of the operators we must try to convert to u128 first + BinaryOp::Mod => self.eval_constant_u128_operations(lhs, rhs)?, + BinaryOp::And => self.eval_constant_u128_operations(lhs, rhs)?, + BinaryOp::Or => self.eval_constant_u128_operations(lhs, rhs)?, + BinaryOp::Xor => self.eval_constant_u128_operations(lhs, rhs)?, + BinaryOp::Shl => self.eval_constant_u128_operations(lhs, rhs)?, + BinaryOp::Shr => self.eval_constant_u128_operations(lhs, rhs)?, + }; + // TODO: Keep original type of constant + Some(dfg.make_constant(value, operand_type)) + } + + /// Try to evaluate the given operands as u128s for operators that are only valid on u128s, + /// like the bitwise operators and modulus. + fn eval_constant_u128_operations( + &self, + lhs: FieldElement, + rhs: FieldElement, + ) -> Option { + let lhs = lhs.try_into_u128()?; + let rhs = rhs.try_into_u128()?; + match self.operator { + BinaryOp::Mod => Some((lhs % rhs).into()), + BinaryOp::And => Some((lhs & rhs).into()), + BinaryOp::Or => Some((lhs | rhs).into()), + BinaryOp::Shr => Some((lhs >> rhs).into()), + // Check for overflow and return None if anything does overflow + BinaryOp::Shl => { + let rhs = rhs.try_into().ok()?; + lhs.checked_shl(rhs).map(Into::into) + } + + // Converting a field xor to a u128 xor would be incorrect since we wouldn't have the + // extra bits of the field. So we don't optimize it here. + BinaryOp::Xor => None, + + op @ (BinaryOp::Add + | BinaryOp::Sub + | BinaryOp::Mul + | BinaryOp::Div + | BinaryOp::Eq + | BinaryOp::Lt) => panic!( + "eval_constant_u128_operations invalid for {op:?} use eval_constants instead" + ), + } + } } /// Binary Operations allowed in the IR. diff --git a/crates/noirc_evaluator/src/ssa_refactor/opt/inlining.rs b/crates/noirc_evaluator/src/ssa_refactor/opt/inlining.rs index 50c97b765b..c63cac520b 100644 --- a/crates/noirc_evaluator/src/ssa_refactor/opt/inlining.rs +++ b/crates/noirc_evaluator/src/ssa_refactor/opt/inlining.rs @@ -9,6 +9,7 @@ use iter_extended::vecmap; use crate::ssa_refactor::{ ir::{ basic_block::BasicBlockId, + dfg::InsertInstructionResult, function::{Function, FunctionId}, instruction::{Instruction, InstructionId, TerminatorInstruction}, value::{Value, ValueId}, @@ -343,7 +344,8 @@ impl<'function> PerFunctionContext<'function> { let old_results = self.source_function.dfg.instruction_results(call_id); let arguments = vecmap(arguments, |arg| self.translate_value(*arg)); let new_results = self.context.inline_function(ssa, function, &arguments); - Self::insert_new_instruction_results(&mut self.values, old_results, &new_results); + let new_results = InsertInstructionResult::Results(&new_results); + Self::insert_new_instruction_results(&mut self.values, old_results, new_results); } /// Push the given instruction from the source_function into the current block of the @@ -365,11 +367,20 @@ impl<'function> PerFunctionContext<'function> { fn insert_new_instruction_results( values: &mut HashMap, old_results: &[ValueId], - new_results: &[ValueId], + new_results: InsertInstructionResult, ) { assert_eq!(old_results.len(), new_results.len()); - for (old_result, new_result) in old_results.iter().zip(new_results) { - values.insert(*old_result, *new_result); + + 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 => (), } } 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 f621503e59..6037909752 100644 --- a/crates/noirc_evaluator/src/ssa_refactor/ssa_builder/mod.rs +++ b/crates/noirc_evaluator/src/ssa_refactor/ssa_builder/mod.rs @@ -11,6 +11,7 @@ use crate::ssa_refactor::ir::{ use super::{ ir::{ basic_block::BasicBlock, + dfg::InsertInstructionResult, instruction::{InstructionId, Intrinsic}, }, ssa_gen::Ssa, @@ -108,10 +109,8 @@ impl FunctionBuilder { &mut self, instruction: Instruction, ctrl_typevars: Option>, - ) -> &[ValueId] { - let id = self.current_function.dfg.make_instruction(instruction, ctrl_typevars); - self.current_function.dfg.insert_instruction_in_block(self.current_block, id); - self.current_function.dfg.instruction_results(id) + ) -> InsertInstructionResult { + self.current_function.dfg.insert_instruction(instruction, self.current_block, ctrl_typevars) } /// Switch to inserting instructions in the given block. @@ -130,7 +129,7 @@ impl FunctionBuilder { /// given amount of field elements. Returns the result of the allocate instruction, /// which is always a Reference to the allocated data. pub(crate) fn insert_allocate(&mut self, size_to_allocate: u32) -> ValueId { - self.insert_instruction(Instruction::Allocate { size: size_to_allocate }, None)[0] + self.insert_instruction(Instruction::Allocate { size: size_to_allocate }, None).first() } /// Insert a Load instruction at the end of the current block, loading from the given offset @@ -147,7 +146,7 @@ impl FunctionBuilder { type_to_load: Type, ) -> ValueId { address = self.insert_binary(address, BinaryOp::Add, offset); - self.insert_instruction(Instruction::Load { address }, Some(vec![type_to_load]))[0] + self.insert_instruction(Instruction::Load { address }, Some(vec![type_to_load])).first() } /// Insert a Store instruction at the end of the current block, storing the given element @@ -166,19 +165,19 @@ impl FunctionBuilder { rhs: ValueId, ) -> ValueId { let instruction = Instruction::Binary(Binary { lhs, rhs, operator }); - self.insert_instruction(instruction, None)[0] + self.insert_instruction(instruction, None).first() } /// Insert a not instruction at the end of the current block. /// Returns the result of the instruction. pub(crate) fn insert_not(&mut self, rhs: ValueId) -> ValueId { - self.insert_instruction(Instruction::Not(rhs), None)[0] + self.insert_instruction(Instruction::Not(rhs), None).first() } /// Insert a cast instruction at the end of the current block. /// Returns the result of the cast instruction. pub(crate) fn insert_cast(&mut self, value: ValueId, typ: Type) -> ValueId { - self.insert_instruction(Instruction::Cast(value, typ), None)[0] + self.insert_instruction(Instruction::Cast(value, typ), None).first() } /// Insert a constrain instruction at the end of the current block. @@ -194,7 +193,7 @@ impl FunctionBuilder { arguments: Vec, result_types: Vec, ) -> &[ValueId] { - self.insert_instruction(Instruction::Call { func, arguments }, Some(result_types)) + self.insert_instruction(Instruction::Call { func, arguments }, Some(result_types)).results() } /// Terminates the current block with the given terminator instruction