From 497b26dee6a219b7b28df037fd35c59708d07400 Mon Sep 17 00:00:00 2001 From: Daniel Lubarov Date: Mon, 25 Jul 2022 09:36:26 -0700 Subject: [PATCH] Some simple optimization rules Depends on #647. --- evm/src/cpu/kernel/aggregator.rs | 2 +- evm/src/cpu/kernel/asm/util/basic_macros.asm | 8 ++- evm/src/cpu/kernel/assembler.rs | 22 ++++-- evm/src/cpu/kernel/mod.rs | 3 +- evm/src/cpu/kernel/optimizer.rs | 71 ++++++++++++++++++++ evm/src/cpu/kernel/parser.rs | 2 +- evm/src/cpu/kernel/utils.rs | 39 +++++++++++ 7 files changed, 135 insertions(+), 12 deletions(-) create mode 100644 evm/src/cpu/kernel/optimizer.rs diff --git a/evm/src/cpu/kernel/aggregator.rs b/evm/src/cpu/kernel/aggregator.rs index e5e1f29f..73d1797b 100644 --- a/evm/src/cpu/kernel/aggregator.rs +++ b/evm/src/cpu/kernel/aggregator.rs @@ -74,7 +74,7 @@ pub(crate) fn combined_kernel() -> Kernel { ]; let parsed_files = files.iter().map(|f| parse(f)).collect_vec(); - assemble(parsed_files, evm_constants()) + assemble(parsed_files, evm_constants(), true) } #[cfg(test)] diff --git a/evm/src/cpu/kernel/asm/util/basic_macros.asm b/evm/src/cpu/kernel/asm/util/basic_macros.asm index e266b2cb..58ab5d2c 100644 --- a/evm/src/cpu/kernel/asm/util/basic_macros.asm +++ b/evm/src/cpu/kernel/asm/util/basic_macros.asm @@ -120,7 +120,9 @@ // stack: input, ... PUSH $c // stack: c, input, ... - GE // Check it backwards: (input <= c) == (c >= input) + %add_const(1) // This will be optimized out. + // stack: c + 1, input, ... + GE // Check it backwards: (input <= c) == (c + 1 > input) // stack: input <= c, ... %endmacro @@ -136,7 +138,9 @@ // stack: input, ... PUSH $c // stack: c, input, ... - LE // Check it backwards: (input >= c) == (c <= input) + %sub_const(1) // This will be optimized out. + // stack: c - 1, input, ... + LT // Check it backwards: (input >= c) == (c - 1 < input) // stack: input >= c, ... %endmacro diff --git a/evm/src/cpu/kernel/assembler.rs b/evm/src/cpu/kernel/assembler.rs index 6e98b22c..0a1232e1 100644 --- a/evm/src/cpu/kernel/assembler.rs +++ b/evm/src/cpu/kernel/assembler.rs @@ -7,6 +7,7 @@ use log::debug; use super::ast::PushTarget; use crate::cpu::kernel::ast::StackReplacement; use crate::cpu::kernel::keccak_util::hash_kernel; +use crate::cpu::kernel::optimizer::optimize_asm; use crate::cpu::kernel::prover_input::ProverInputFn; use crate::cpu::kernel::stack_manipulation::expand_stack_manipulation; use crate::cpu::kernel::utils::u256_to_trimmed_be_bytes; @@ -64,7 +65,11 @@ impl Macro { } } -pub(crate) fn assemble(files: Vec, constants: HashMap) -> Kernel { +pub(crate) fn assemble( + files: Vec, + constants: HashMap, + optimize: bool, +) -> Kernel { let macros = find_macros(&files); let mut global_labels = HashMap::new(); let mut prover_inputs = HashMap::new(); @@ -75,7 +80,10 @@ pub(crate) fn assemble(files: Vec, constants: HashMap) -> Ke let expanded_file = expand_macros(file.body, ¯os); let expanded_file = expand_repeats(expanded_file); let expanded_file = inline_constants(expanded_file, &constants); - let expanded_file = expand_stack_manipulation(expanded_file); + let mut expanded_file = expand_stack_manipulation(expanded_file); + if optimize { + optimize_asm(&mut expanded_file); + } local_labels.push(find_labels( &expanded_file, &mut offset, @@ -381,7 +389,7 @@ mod tests { let expected_kernel = Kernel::new(expected_code, expected_global_labels, HashMap::new()); let program = vec![file_1, file_2]; - assert_eq!(assemble(program, HashMap::new()), expected_kernel); + assert_eq!(assemble(program, HashMap::new(), false), expected_kernel); } #[test] @@ -399,7 +407,7 @@ mod tests { Item::StandardOp("JUMPDEST".to_string()), ], }; - assemble(vec![file_1, file_2], HashMap::new()); + assemble(vec![file_1, file_2], HashMap::new(), false); } #[test] @@ -413,7 +421,7 @@ mod tests { Item::StandardOp("ADD".to_string()), ], }; - assemble(vec![file], HashMap::new()); + assemble(vec![file], HashMap::new(), false); } #[test] @@ -421,7 +429,7 @@ mod tests { let file = File { body: vec![Item::Bytes(vec![0x12, 42]), Item::Bytes(vec![0xFE, 255])], }; - let code = assemble(vec![file], HashMap::new()).code; + let code = assemble(vec![file], HashMap::new(), false).code; assert_eq!(code, vec![0x12, 42, 0xfe, 255]); } @@ -530,6 +538,6 @@ mod tests { constants: HashMap, ) -> Kernel { let parsed_files = files.iter().map(|f| parse(f)).collect_vec(); - assemble(parsed_files, constants) + assemble(parsed_files, constants, false) } } diff --git a/evm/src/cpu/kernel/mod.rs b/evm/src/cpu/kernel/mod.rs index f0247f93..d87c1e13 100644 --- a/evm/src/cpu/kernel/mod.rs +++ b/evm/src/cpu/kernel/mod.rs @@ -3,6 +3,7 @@ pub mod assembler; mod ast; pub(crate) mod keccak_util; mod opcodes; +mod optimizer; mod parser; pub mod prover_input; mod stack_manipulation; @@ -23,6 +24,6 @@ use crate::cpu::kernel::aggregator::evm_constants; /// This is for debugging the kernel only. pub fn assemble_to_bytes(files: &[String]) -> Vec { let parsed_files: Vec<_> = files.iter().map(|f| parse(f)).collect(); - let kernel = assemble(parsed_files, evm_constants()); + let kernel = assemble(parsed_files, evm_constants(), true); kernel.code } diff --git a/evm/src/cpu/kernel/optimizer.rs b/evm/src/cpu/kernel/optimizer.rs new file mode 100644 index 00000000..2b162792 --- /dev/null +++ b/evm/src/cpu/kernel/optimizer.rs @@ -0,0 +1,71 @@ +use ethereum_types::U256; +use Item::{Push, StandardOp}; +use PushTarget::Literal; + +use crate::cpu::kernel::ast::Item::LocalLabelDeclaration; +use crate::cpu::kernel::ast::PushTarget::Label; +use crate::cpu::kernel::ast::{Item, PushTarget}; +use crate::cpu::kernel::utils::replace_windows; + +pub(crate) fn optimize_asm(code: &mut Vec) { + constant_propagation(code); + + // Remove no-op jumps: [PUSH label, JUMP, label:] -> [label:] + replace_windows(code, |window| { + if let [Push(Label(l1)), StandardOp(jump), LocalLabelDeclaration(l2)] = window + && l1 == l2 + && &jump == "JUMP" + { + Some(vec![LocalLabelDeclaration(l2)]) + } else { + None + } + }); + + // Remove swaps: [PUSH x, PUSH y, SWAP1] -> [PUSH y, PUSH x] + replace_windows(code, |window| { + if let [Push(Literal(x)), Push(Literal(y)), StandardOp(swap1)] = window + && &swap1 == "SWAP1" { + Some(vec![Push(Literal(y)), Push(Literal(x))]) + } else { + None + } + }); +} + +fn constant_propagation(code: &mut Vec) { + // Constant propagation for unary ops: [PUSH x, UNARYOP] -> [PUSH UNARYOP(x)] + replace_windows(code, |window| { + if let [Push(Literal(x)), StandardOp(op)] = window { + match op.as_str() { + "ISZERO" => Some(vec![Push(Literal(if x.is_zero() { + U256::one() + } else { + U256::zero() + }))]), + "NOT" => Some(vec![Push(Literal(!x))]), + _ => None, + } + } else { + None + } + }); + + // Constant propagation for binary ops: [PUSH x, PUSH y, BINOP] -> [PUSH BINOP(x, y)] + replace_windows(code, |window| { + if let [Push(Literal(x)), Push(Literal(y)), StandardOp(op)] = window { + match op.as_str() { + "ADD" => Some(vec![Push(Literal(x + y))]), + "SUB" => Some(vec![Push(Literal(x - y))]), + "MUL" => Some(vec![Push(Literal(x * y))]), + "DIV" => Some(vec![Push(Literal(x / y))]), + _ => None, + } + } else { + None + } + }); +} + +#[cfg(test)] +mod tests {} diff --git a/evm/src/cpu/kernel/parser.rs b/evm/src/cpu/kernel/parser.rs index 860dc19d..66bf0757 100644 --- a/evm/src/cpu/kernel/parser.rs +++ b/evm/src/cpu/kernel/parser.rs @@ -45,7 +45,7 @@ fn parse_item(item: Pair) -> Item { .collect::>() .into(), ), - Rule::nullary_instruction => Item::StandardOp(item.as_str().into()), + Rule::nullary_instruction => Item::StandardOp(item.as_str().to_uppercase()), _ => panic!("Unexpected {:?}", item.as_rule()), } } diff --git a/evm/src/cpu/kernel/utils.rs b/evm/src/cpu/kernel/utils.rs index d9682679..cff10430 100644 --- a/evm/src/cpu/kernel/utils.rs +++ b/evm/src/cpu/kernel/utils.rs @@ -1,6 +1,30 @@ +use std::fmt::Debug; + use ethereum_types::U256; use plonky2_util::ceil_div_usize; +/// Enumerate the length `W` windows of `vec`, and run `maybe_replace` on each one. +/// +/// Whenever `maybe_replace` returns `Some(replacement)`, the given replacement will be applied. +pub(crate) fn replace_windows(vec: &mut Vec, maybe_replace: F) +where + T: Clone + Debug, + F: Fn([T; W]) -> Option>, +{ + let mut start = 0; + while start + W <= vec.len() { + let range = start..start + W; + let window = vec[range.clone()].to_vec().try_into().unwrap(); + if let Some(replacement) = maybe_replace(window) { + vec.splice(range, replacement); + // Go back to the earliest window that changed. + start = start.saturating_sub(W - 1); + } else { + start += 1; + } + } +} + pub(crate) fn u256_to_trimmed_be_bytes(u256: &U256) -> Vec { let num_bytes = ceil_div_usize(u256.bits(), 8).max(1); // `byte` is little-endian, so we manually reverse it. @@ -11,6 +35,21 @@ pub(crate) fn u256_to_trimmed_be_bytes(u256: &U256) -> Vec { mod tests { use super::*; + #[test] + fn test_replace_windows() { + // This replacement function adds pairs of integers together. + let mut vec = vec![1, 2, 3, 4, 5]; + replace_windows(&mut vec, |[x, y]| Some(vec![x + y])); + assert_eq!(vec, vec![15u32]); + + // This replacement function splits each composite integer into two factors. + let mut vec = vec![9, 1, 6, 8, 15, 7, 9]; + replace_windows(&mut vec, |[n]| { + (2..n).filter(|d| n % d == 0).next().map(|d| vec![d, n / d]) + }); + assert_eq!(vec, vec![3, 3, 1, 2, 3, 2, 2, 2, 3, 5, 7, 3, 3]); + } + #[test] fn literal_to_be_bytes() { assert_eq!(u256_to_trimmed_be_bytes(&0.into()), vec![0x00]);