diff --git a/evm/src/cpu/kernel/cost_estimator.rs b/evm/src/cpu/kernel/cost_estimator.rs new file mode 100644 index 00000000..3dfcf63a --- /dev/null +++ b/evm/src/cpu/kernel/cost_estimator.rs @@ -0,0 +1,37 @@ +use crate::cpu::kernel::assembler::BYTES_PER_OFFSET; +use crate::cpu::kernel::ast::Item; +use crate::cpu::kernel::ast::Item::*; +use crate::cpu::kernel::ast::PushTarget::*; +use crate::cpu::kernel::utils::u256_to_trimmed_be_bytes; + +pub(crate) fn is_code_improved(before: &[Item], after: &[Item]) -> bool { + cost_estimate(after) < cost_estimate(before) +} + +fn cost_estimate(code: &[Item]) -> u32 { + code.iter().map(cost_estimate_item).sum() +} + +fn cost_estimate_item(item: &Item) -> u32 { + match item { + MacroDef(_, _, _) => 0, + GlobalLabelDeclaration(_) => 0, + LocalLabelDeclaration(_) => 0, + Push(Literal(n)) => cost_estimate_push(u256_to_trimmed_be_bytes(n).len()), + Push(Label(_)) => cost_estimate_push(BYTES_PER_OFFSET as usize), + ProverInput(_) => 1, + StandardOp(op) => cost_estimate_standard_op(op.as_str()), + _ => panic!("Unexpected item: {:?}", item), + } +} + +fn cost_estimate_standard_op(_op: &str) -> u32 { + // For now we just treat any standard operation as having the same cost. This is pretty naive, + // but should work fine with our current set of simple optimization rules. + 1 +} + +fn cost_estimate_push(num_bytes: usize) -> u32 { + // TODO: Once PUSH is actually implemented, check if this needs to be revised. + num_bytes as u32 +} diff --git a/evm/src/cpu/kernel/mod.rs b/evm/src/cpu/kernel/mod.rs index d87c1e13..59caff76 100644 --- a/evm/src/cpu/kernel/mod.rs +++ b/evm/src/cpu/kernel/mod.rs @@ -1,6 +1,7 @@ pub mod aggregator; pub mod assembler; mod ast; +mod cost_estimator; pub(crate) mod keccak_util; mod opcodes; mod optimizer; diff --git a/evm/src/cpu/kernel/optimizer.rs b/evm/src/cpu/kernel/optimizer.rs index 6fe9d496..2a1db6d3 100644 --- a/evm/src/cpu/kernel/optimizer.rs +++ b/evm/src/cpu/kernel/optimizer.rs @@ -5,6 +5,7 @@ use PushTarget::Literal; use crate::cpu::kernel::ast::Item::{GlobalLabelDeclaration, LocalLabelDeclaration}; use crate::cpu::kernel::ast::PushTarget::Label; use crate::cpu::kernel::ast::{Item, PushTarget}; +use crate::cpu::kernel::cost_estimator::is_code_improved; use crate::cpu::kernel::utils::{replace_windows, u256_from_bool}; pub(crate) fn optimize_asm(code: &mut Vec) { @@ -30,7 +31,7 @@ fn optimize_asm_once(code: &mut Vec) { /// Constant propagation. fn constant_propagation(code: &mut Vec) { // Constant propagation for unary ops: `[PUSH x, UNARYOP] -> [PUSH UNARYOP(x)]` - replace_windows(code, |window| { + replace_windows_if_better(code, |window| { if let [Push(Literal(x)), StandardOp(op)] = window { match op.as_str() { "ISZERO" => Some(vec![Push(Literal(u256_from_bool(x.is_zero())))]), @@ -43,7 +44,7 @@ fn constant_propagation(code: &mut Vec) { }); // Constant propagation for binary ops: `[PUSH y, PUSH x, BINOP] -> [PUSH BINOP(x, y)]` - replace_windows(code, |window| { + replace_windows_if_better(code, |window| { if let [Push(Literal(y)), Push(Literal(x)), StandardOp(op)] = window { match op.as_str() { "ADD" => Some(x.overflowing_add(y).0), @@ -129,6 +130,17 @@ fn remove_ignored_values(code: &mut Vec) { }); } +/// Like `replace_windows`, but specifically for code, and only makes replacements if our cost +/// estimator thinks that the new code is more efficient. +fn replace_windows_if_better(code: &mut Vec, maybe_replace: F) +where + F: Fn([Item; W]) -> Option>, +{ + replace_windows(code, |window| { + maybe_replace(window.clone()).filter(|suggestion| is_code_improved(&window, suggestion)) + }) +} + #[cfg(test)] mod tests { use super::*; @@ -153,13 +165,18 @@ mod tests { #[test] fn test_constant_propagation_sub_underflowing() { - let mut code = vec![ + let original = vec![ Push(Literal(U256::one())), Push(Literal(U256::zero())), StandardOp("SUB".into()), ]; + let mut code = original.clone(); constant_propagation(&mut code); - assert_eq!(code, vec![Push(Literal(U256::max_value()))]); + // Constant propagation could replace the code with [PUSH U256::MAX], but that's actually + // more expensive, so the code shouldn't be changed. + // (The code could also be replaced with [PUSH 0; NOT], which would be an improvement, but + // our optimizer isn't smart enough yet.) + assert_eq!(code, original); } #[test]