Merge branch 'mk/wasm_metering_a' into 'master'

RUN-328, RUN-633: New Wasm Instrumentation Part 1

The new instrumentation significantly improves the performance of several applications as well as improves the accuracy of instruction counting in some (corner) cases. 

See merge request dfinity-lab/public/ic!12391

rs/config/src/embedders.rs

@@ -88,6 +88,14 @@ impl Default for FeatureFlags {
     }
 }
 
+#[derive(Copy, Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
+pub enum MeteringType {
+    Old,
+    New,
+    /// for testing and benchmarking
+    None,
+}
+
 #[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
 pub struct Config {
     pub max_wasm_stack_size: usize,
@@ -125,6 +133,9 @@ pub struct Config {
     /// Flags to enable or disable features that are still experimental.
     pub feature_flags: FeatureFlags,
 
+    /// Instruction counting strategy
+    pub metering_type: MeteringType,
+
     // Maximum number of stable memory dirty pages that a single message execution
     // is allowed to produce.
     pub stable_memory_dirty_page_limit: NumPages,
@@ -174,6 +185,7 @@ impl Config {
             cost_to_compile_wasm_instruction: DEFAULT_COST_TO_COMPILE_WASM_INSTRUCTION,
             num_rayon_compilation_threads: DEFAULT_WASMTIME_RAYON_COMPILATION_THREADS,
             feature_flags: FeatureFlags::const_default(),
+            metering_type: MeteringType::Old,
             stable_memory_dirty_page_limit: NumPages::new(STABLE_MEMORY_DIRTY_PAGE_LIMIT),
             stable_memory_accessed_page_limit: NumPages::new(STABLE_MEMORY_ACCESSED_PAGE_LIMIT),
             min_sandbox_count: DEFAULT_MIN_SANDBOX_COUNT,

rs/embedders/src/wasm_utils.rs

@@ -207,6 +207,7 @@ fn validate_and_instrument(
         config.cost_to_compile_wasm_instruction,
         config.feature_flags.write_barrier,
         config.feature_flags.wasm_native_stable_memory,
+        config.metering_type,
         config.subnet_type,
         config.dirty_page_overhead,
     )?;

rs/embedders/src/wasm_utils/instrumentation.rs

@@ -113,6 +113,7 @@
 use super::system_api_replacements::replacement_functions;
 use super::validation::API_VERSION_IC0;
 use super::{InstrumentationOutput, Segments, SystemApiFunc};
+use ic_config::embedders::MeteringType;
 use ic_config::flag_status::FlagStatus;
 use ic_registry_subnet_type::SubnetType;
 use ic_replicated_state::NumWasmPages;
@@ -155,7 +156,7 @@ impl InjectedImports {
 }
 
 // Gets the cost of an instruction.
-fn instruction_to_cost(i: &Operator) -> u64 {
+pub fn instruction_to_cost(i: &Operator) -> u64 {
     match i {
         // The following instructions are mostly signaling the start/end of code blocks,
         // so we assign 0 cost to them.
@@ -387,6 +388,7 @@ pub(super) fn instrument(
     cost_to_compile_wasm_instruction: NumInstructions,
     write_barrier: FlagStatus,
     wasm_native_stable_memory: FlagStatus,
+    metering_type: MeteringType,
     subnet_type: SubnetType,
     dirty_page_overhead: NumInstructions,
 ) -> Result<InstrumentationOutput, WasmInstrumentationError> {
@@ -448,7 +450,7 @@ pub(super) fn instrument(
 
     // inject instructions counter decrementation
     for func_body in &mut module.code_sections {
-        inject_metering(&mut func_body.instructions, &special_indices);
+        inject_metering(&mut func_body.instructions, &special_indices, metering_type);
     }
 
     // Collect all the function types of the locally defined functions inside the
@@ -880,8 +882,16 @@ impl InjectionPoint {
 // - we insert a function call before each dynamic cost instruction which
 //   performs an overflow check and then decrements the counter by the value at
 //   the top of the stack.
-fn inject_metering(code: &mut Vec<Operator>, export_data_module: &SpecialIndices) {
-    let points = injections(code);
+fn inject_metering(
+    code: &mut Vec<Operator>,
+    export_data_module: &SpecialIndices,
+    metering_type: MeteringType,
+) {
+    let points = match metering_type {
+        MeteringType::Old => injections_old(code),
+        MeteringType::None => Vec::new(),
+        MeteringType::New => injections_new(code),
+    };
     let points = points.iter().filter(|point| match point.cost_detail {
         InjectionPointCostDetail::StaticCost {
             scope: Scope::ReentrantBlockStart,
@@ -1218,8 +1228,9 @@ fn inject_update_available_memory(func_body: &mut wasm_transform::Body, func_typ
 // at the beginning of every basic block (straight-line sequence of instructions
 // with no branches) and before each bulk memory instruction. An injection point
 // contains a "hint" about the context of every basic block, specifically if
-// it's re-entrant or not.
-fn injections(code: &[Operator]) -> Vec<InjectionPoint> {
+// it's re-entrant or not. This version over-estimates the cost of code with
+// returns and jumps.
+fn injections_old(code: &[Operator]) -> Vec<InjectionPoint> {
     let mut res = Vec::new();
     let mut stack = Vec::new();
     use Operator::*;
@@ -1271,6 +1282,63 @@ fn injections(code: &[Operator]) -> Vec<InjectionPoint> {
     res
 }
 
+// This function scans through the Wasm code and creates an injection point
+// at the beginning of every basic block (straight-line sequence of instructions
+// with no branches) and before each bulk memory instruction. An injection point
+// contains a "hint" about the context of every basic block, specifically if
+// it's re-entrant or not.
+fn injections_new(code: &[Operator]) -> Vec<InjectionPoint> {
+    let mut res = Vec::new();
+    use Operator::*;
+    // The function itself is a re-entrant code block.
+    let mut curr = InjectionPoint::new_static_cost(0, Scope::ReentrantBlockStart);
+    for (position, i) in code.iter().enumerate() {
+        curr.cost_detail.increment_cost(instruction_to_cost(i));
+        match i {
+            // Start of a re-entrant code block.
+            Loop { .. } => {
+                res.push(curr);
+                curr = InjectionPoint::new_static_cost(position + 1, Scope::ReentrantBlockStart);
+            }
+            // Start of a non re-entrant code block.
+            If { .. } => {
+                res.push(curr);
+                curr = InjectionPoint::new_static_cost(position + 1, Scope::NonReentrantBlockStart);
+            }
+            // End of a code block but still more code left.
+            Else | Br { .. } | BrIf { .. } | BrTable { .. } => {
+                res.push(curr);
+                curr = InjectionPoint::new_static_cost(position + 1, Scope::BlockEnd);
+            }
+            End => {
+                res.push(curr);
+                curr = InjectionPoint::new_static_cost(position + 1, Scope::BlockEnd);
+            }
+            Return | Unreachable | ReturnCall { .. } | ReturnCallIndirect { .. } => {
+                res.push(curr);
+                // This injection point will be unreachable itself (most likely empty)
+                // but we create it to keep the algorithm uniform
+                curr = InjectionPoint::new_static_cost(position + 1, Scope::BlockEnd);
+            }
+            // Bulk memory instructions require injected metering __before__ the instruction
+            // executes so that size arguments can be read from the stack at runtime.
+            MemoryFill { .. }
+            | MemoryCopy { .. }
+            | MemoryInit { .. }
+            | TableCopy { .. }
+            | TableInit { .. }
+            | TableFill { .. } => {
+                res.push(InjectionPoint::new_dynamic_cost(position));
+            }
+            // Nothing special to be done for other instructions.
+            _ => (),
+        }
+    }
+
+    res.sort_by_key(|k| k.position);
+    res
+}
+
 // Looks for the data section and if it is present, converts it to a vector of
 // tuples (heap offset, bytes) and then deletes the section.
 fn get_data(