Revised instruction-count bounding with fewer checks and loads/stores.

This PR revises the simple instruction-count binding work, which
initially loaded and stored the instruction-count field from memory on
every basic block and emitted an extra conditional branch and yielding
block after every basic block, with two optimizations:

1. The instruction count is kept in a local variable (which the register
   allocator will ideally keep in a register unless spills are
   necessary), and this count is saved back to the field in memory only
   when necessary, e.g., before calls and returns.

2. The checks and conditional branches are inserted only in places where
   necessary to ensure bounded time between checks, at the cost of
   slightly more overrun/approximation in the bound: specifically,
   before calls and returns (to avoid unbounded runtime due to
   recursion) and at loop backedges.

lucet-module/src/runtime.rs

@@ -11,16 +11,10 @@ pub struct InstanceRuntimeData {
     /// `instruction_count_adj` set to some negative value and
     /// `instruction_count_bound` adjusted upward in compensation.
     /// `instruction_count_adj` is incremented as execution proceeds; on each
-    /// increment, the Wasm code checks the carry flag. If the value crosses
-    /// zero (becomes positive), then we have exceeded the bound and we must
-    /// yield. At any point, the `adj` value can be adjusted downward by
-    /// transferring the count to the `bound`.
-    ///
-    /// Note that the bound-yield is only triggered if the `adj` value
-    /// transitions from negative to non-negative; in other words, it is
-    /// edge-triggered, not level-triggered. So entering code that has been
-    /// instrumented for instruction counting with `adj` >= 0 will result in no
-    /// bound ever triggered (until 2^64 instructions execute).
+    /// increment, the Wasm code checks the sign. If the value is greater than
+    /// zero, then we have exceeded the bound and we must yield. At any point,
+    /// the `adj` value can be adjusted downward by transferring the count to
+    /// the `bound`.
     pub instruction_count_adj: i64,
     pub instruction_count_bound: i64,
     pub stack_limit: u64,

lucet-runtime/lucet-runtime-internals/src/future.rs

@@ -1,5 +1,6 @@
 use crate::error::Error;
 use crate::instance::{InstanceHandle, InternalRunResult, RunResult, State, TerminationDetails};
+use crate::module::FunctionHandle;
 use crate::val::{UntypedRetVal, Val};
 use crate::vmctx::{Vmctx, VmctxInternal};
 use std::any::Any;
@@ -130,6 +131,46 @@ impl InstanceHandle {
         entrypoint: &'a str,
         args: &'a [Val],
         runtime_bound: Option<u64>,
+    ) -> Result<UntypedRetVal, Error> {
+        let func = self.module.get_export_func(entrypoint)?;
+        self.run_async_internal(func, args, runtime_bound).await
+    }
+
+    /// Run the module's [start function][start], if one exists.
+    ///
+    /// If there is no start function in the module, this does nothing.
+    ///
+    /// All of the other restrictions on the start function, what it may do, and
+    /// the requirement that it must be invoked first, are described in the
+    /// documentation for `Instance::run_start()`. This async version of that
+    /// function satisfies the requirement to run the start function first, as
+    /// long as the async function fully returns (not just yields).
+    ///
+    /// This method is similar to `Instance::run_start()`, except that it bounds
+    /// runtime between async future yields (invocations of `.poll()` on the
+    /// underlying generated future) if `runtime_bound` is provided. This
+    /// behaves the same way as `Instance::run_async()`.
+    pub async fn run_async_start<'a>(
+        &'a mut self,
+        runtime_bound: Option<u64>,
+    ) -> Result<(), Error> {
+        println!("run_async_start: get_start_func = {:?}", self.module.get_start_func());
+        if let Some(start) = self.module.get_start_func()? {
+            if !self.is_not_started() {
+                return Err(Error::StartAlreadyRun);
+            }
+            self.run_async_internal(start, &[], runtime_bound).await?;
+        }
+        Ok(())
+    }
+
+    /// Shared async run-loop implementation for both `run_async()` and
+    /// `run_start_async()`.
+    async fn run_async_internal<'a>(
+        &'a mut self,
+        func: FunctionHandle,
+        args: &'a [Val],
+        runtime_bound: Option<u64>,
     ) -> Result<UntypedRetVal, Error> {
         if self.is_yielded() {
             return Err(Error::Unsupported(
@@ -158,7 +199,6 @@ impl InstanceHandle {
                 )
             } else {
                 // This is the first iteration, call the entrypoint:
-                let func = self.module.get_export_func(entrypoint)?;
                 self.run_func(func, args, true, runtime_bound)
             };
             match run_result? {

lucet-runtime/lucet-runtime-internals/src/instance.rs

@@ -427,12 +427,10 @@ pub(crate) enum InternalRunResult {
 }
 
 impl InternalRunResult {
-    pub(crate) fn unwrap(self) -> Result<RunResult, Error> {
+    pub(crate) fn unwrap(self) -> RunResult {
         match self {
-            InternalRunResult::Normal(result) => Ok(result),
-            InternalRunResult::BoundExpired => Err(Error::InvalidArgument(
-                "should not have had a runtime bound",
-            )),
+            InternalRunResult::Normal(result) => result,
+            InternalRunResult::BoundExpired => panic!("should not have had a runtime bound"),
         }
     }
 }
@@ -517,7 +515,7 @@ impl Instance {
     /// in the future.
     pub fn run(&mut self, entrypoint: &str, args: &[Val]) -> Result<RunResult, Error> {
         let func = self.module.get_export_func(entrypoint)?;
-        self.run_func(func, &args, false, None)?.unwrap()
+        Ok(self.run_func(func, &args, false, None)?.unwrap())
     }
 
     /// Run a function with arguments in the guest context from the [WebAssembly function
@@ -533,7 +531,7 @@ impl Instance {
         args: &[Val],
     ) -> Result<RunResult, Error> {
         let func = self.module.get_func_from_idx(table_idx, func_idx)?;
-        self.run_func(func, &args, false, None)?.unwrap()
+        Ok(self.run_func(func, &args, false, None)?.unwrap())
     }
 
     /// Resume execution of an instance that has yielded without providing a value to the guest.
@@ -564,7 +562,7 @@ impl Instance {
     /// The foreign code safety caveat of [`Instance::run()`](struct.Instance.html#method.run)
     /// applies.
     pub fn resume_with_val<A: Any + 'static>(&mut self, val: A) -> Result<RunResult, Error> {
-        self.resume_with_val_impl(val, false, None)?.unwrap()
+        Ok(self.resume_with_val_impl(val, false, None)?.unwrap())
     }
 
     pub(crate) fn resume_with_val_impl<A: Any + 'static>(
@@ -587,7 +585,7 @@ impl Instance {
 
         self.resumed_val = Some(Box::new(val) as Box<dyn Any + 'static>);
 
-        self.set_instruction_bound_delta(max_insn_count.unwrap_or(0));
+        self.set_instruction_bound_delta(max_insn_count);
         self.swap_and_return(async_context)
     }
 
@@ -611,7 +609,7 @@ impl Instance {
                 "can only call resume_bounded() on an instance that hit an instruction bound",
             ));
         }
-        self.set_instruction_bound_delta(max_insn_count);
+        self.set_instruction_bound_delta(Some(max_insn_count));
         self.swap_and_return(true)
     }
 
@@ -650,10 +648,7 @@ impl Instance {
             if !self.is_not_started() {
                 return Err(Error::StartAlreadyRun);
             }
-            let res = self.run_func(start, &[], false, None)?.unwrap()?;
-            if res.is_yielded() {
-                return Err(Error::StartYielded);
-            }
+            self.run_func(start, &[], false, None)?;
         }
         Ok(())
     }
@@ -958,11 +953,12 @@ impl Instance {
     /// value is *crossed*, but not if execution *begins* with the value exceeded. Hence `delta`
     /// must be greater than zero for this to set up the instance state to trigger a yield.
     #[inline]
-    pub fn set_instruction_bound_delta(&mut self, delta: u64) {
+    pub fn set_instruction_bound_delta(&mut self, delta: Option<u64>) {
         let implicits = self.get_instance_implicits_mut();
         let sum = implicits.instruction_count_adj + implicits.instruction_count_bound;
+        let delta = delta.unwrap_or(i64::MAX as u64);
         let delta = i64::try_from(delta).expect("delta too large");
-        implicits.instruction_count_bound = sum + delta;
+        implicits.instruction_count_bound = sum.wrapping_add(delta);
         implicits.instruction_count_adj = -delta;
     }
 
@@ -1140,7 +1136,7 @@ impl Instance {
         let mut args_with_vmctx = vec![Val::from(self.alloc.slot().heap)];
         args_with_vmctx.extend_from_slice(args);
 
-        self.set_instruction_bound_delta(inst_count_bound.unwrap_or(0));
+        self.set_instruction_bound_delta(inst_count_bound);
 
         let self_ptr = self as *mut _;
         Context::init_with_callback(

lucet-runtime/tests/instruction_counting.rs

@@ -26,7 +26,7 @@ pub fn wasm_test<P: AsRef<Path>>(
     Ok(dlmodule)
 }
 
-pub fn get_instruction_count_test_files() -> Vec<DirEntry> {
+pub fn get_instruction_count_test_files(want_start_function: bool) -> Vec<DirEntry> {
     std::fs::read_dir("./tests/instruction_counting")
         .expect("can iterate test files")
         .map(|ent| {
@@ -37,12 +37,14 @@ pub fn get_instruction_count_test_files() -> Vec<DirEntry> {
             );
             ent
         })
+        .filter(|ent| want_start_function == ent.path().to_str().unwrap().contains("_start.wat"))
         .collect()
 }
 
 #[test]
 pub fn check_instruction_count_off() {
-    let files: Vec<DirEntry> = get_instruction_count_test_files();
+    let files: Vec<DirEntry> =
+        get_instruction_count_test_files(/* want_start_function = */ false);
 
     assert!(
         !files.is_empty(),
@@ -71,7 +73,8 @@ pub fn check_instruction_count_off() {
 
 #[test]
 pub fn check_instruction_count() {
-    let files: Vec<DirEntry> = get_instruction_count_test_files();
+    let files: Vec<DirEntry> =
+        get_instruction_count_test_files(/* want_start_function = */ false);
 
     assert!(
         !files.is_empty(),
@@ -136,7 +139,16 @@ fn dummy_waker() -> Waker {
 
 #[test]
 pub fn check_instruction_count_with_periodic_yields() {
-    let files: Vec<DirEntry> = get_instruction_count_test_files();
+    check_instruction_count_with_periodic_yields_internal(/* want_start_function = */ false);
+}
+
+#[test]
+pub fn check_instruction_count_with_periodic_yields_start_func() {
+    check_instruction_count_with_periodic_yields_internal(/* want_start_function = */ true);
+}
+
+fn check_instruction_count_with_periodic_yields_internal(want_start_function: bool) {
+    let files: Vec<DirEntry> = get_instruction_count_test_files(want_start_function);
 
     assert!(
         !files.is_empty(),
@@ -154,15 +166,13 @@ pub fn check_instruction_count_with_periodic_yields() {
             .new_instance(module)
             .expect("instance can be created");
 
-        let yields = {
+        fn future_loop(mut future: std::pin::Pin<Box<impl Future>>) -> u64 {
             let mut yields = 0;
-            let mut future = Box::pin(inst.run_async("test_function", &[], Some(1000)));
             let waker = dummy_waker();
             let mut context = Context::from_waker(&waker);
             loop {
                 match future.as_mut().poll(&mut context) {
-                    Poll::Ready(val) => {
-                        val.expect("instance runs");
+                    Poll::Ready(_) => {
                         break;
                     }
                     Poll::Pending => {
@@ -174,6 +184,14 @@ pub fn check_instruction_count_with_periodic_yields() {
                 }
             }
             yields
+        }
+
+        let yields = if want_start_function {
+            let future = Box::pin(inst.run_async_start(Some(1000)));
+            future_loop(future)
+        } else {
+            let future = Box::pin(inst.run_async("test_function", &[], Some(1000)));
+            future_loop(future)
         };
 
         let instruction_count = inst

lucet-runtime/tests/instruction_counting/long_loop_start.wat

@@ -0,0 +1,17 @@
+(module
+  (start $start)
+  (func $start (export "_start") (local i32)
+      loop
+        local.get 0
+        i32.const 1
+	i32.add
+	local.tee 0
+	i32.const 10000
+	i32.ne
+	br_if 0
+      end
+  )
+  (func $instruction_count (export "instruction_count") (result i64)
+    i64.const 70000
+  )
+)

lucetc/src/compiler.rs

@@ -341,6 +341,22 @@ impl<'a> Compiler<'a> {
         self.decls.get_module_data(self.module_features())
     }
 
+    fn get_local_count(body: &FunctionBody, name: &str) -> Result<u32, Error> {
+        let error_mapper = |e| Error::FunctionTranslation {
+            symbol: name.to_string(),
+            source: Box::new(Error::from(e)),
+        };
+        Ok(body
+            .get_locals_reader()
+            .map_err(error_mapper)?
+            .into_iter()
+            .map(|result| result.map_err(error_mapper))
+            .collect::<Result<Vec<(u32, _)>, Error>>()?
+            .into_iter()
+            .map(|(count, _ty)| count)
+            .sum())
+    }
+
     pub fn object_file(self) -> Result<ObjectFile, Error> {
         let mut function_manifest_ctx = ClifDataContext::new();
         let mut function_manifest_bytes = Cursor::new(Vec::new());
@@ -359,7 +375,15 @@ impl<'a> Compiler<'a> {
                 let func = decls
                     .get_func(unique_func_ix)
                     .expect("decl exists for func body");
-                let mut func_info = FuncInfo::new(&decls, &codegen_context, count_instructions);
+                let arg_count = func.signature.params.len() as u32;
+                let local_count = Self::get_local_count(&func_body, func.name.symbol())?;
+                let mut func_info = FuncInfo::new(
+                    &decls,
+                    &codegen_context,
+                    count_instructions,
+                    arg_count,
+                    local_count,
+                );
                 let mut clif_context = ClifContext::new();
                 clif_context.func.name = func.name.as_externalname();
                 clif_context.func.signature = func.signature.clone();
@@ -536,8 +560,15 @@ impl<'a> Compiler<'a> {
                 .decls
                 .get_func(unique_func_ix)
                 .expect("decl exists for func body");
-            let mut func_info =
-                FuncInfo::new(&self.decls, &self.codegen_context, self.count_instructions);
+            let arg_count = func.signature.params.len() as u32;
+            let local_count = Self::get_local_count(&body, func.name.symbol())?;
+            let mut func_info = FuncInfo::new(
+                &self.decls,
+                &self.codegen_context,
+                self.count_instructions,
+                arg_count,
+                local_count,
+            );
             let mut clif_context = ClifContext::new();
             clif_context.func.name = func.name.as_externalname();
             clif_context.func.signature = func.signature.clone();

lucetc/src/error.rs

@@ -1,6 +1,7 @@
 use crate::types::SignatureError;
 use crate::validate::Error as ValidationError;
 use cranelift_module::ModuleError as ClifModuleError;
+use cranelift_wasm::wasmparser::BinaryReaderError as ClifWasmReaderError;
 use cranelift_wasm::WasmError as ClifWasmError;
 use lucet_module::error::Error as LucetModuleError;
 use object;
@@ -27,6 +28,8 @@ pub enum Error {
     MissingWasmPreamble,
     #[error("Wasm validation: {0}")]
     WasmValidation(#[from] wasmparser::BinaryReaderError),
+    #[error("Wasm validation: {0}")]
+    ClifWasmValidation(#[from] ClifWasmReaderError),
     #[error("Wat input: {0}")]
     WatInput(#[from] wabt::Error),
     #[error("Object artifact: {1}. {0:?}")]