Block machine: only use JIT on Goldilocks

Author: georgwiese

executor/src/witgen/eval_result.rs

@@ -47,6 +47,8 @@ pub enum IncompleteCause<K = usize> {
     SymbolicEvaluationOfChallenge,
     /// Some knowledge was learnt, but not a concrete value. Example: `Y = X` if we know that `Y` is boolean. We learn that `X` is boolean, but not its exact value.
     NotConcrete,
+    /// The JIT compiler was not able to generate a function that computes a unique witness.
+    JitCompilationFailed,
     Multiple(Vec<IncompleteCause<K>>),
 }
 

executor/src/witgen/jit/function_cache.rs

@@ -36,8 +36,8 @@ pub struct FunctionCache<'a, T: FieldElement> {
     /// The processor that generates the JIT code
     processor: BlockMachineProcessor<'a, T>,
     /// The cache of JIT functions and the returned range constraints.
-    /// If the entry is None, we attempted to generate the function but failed.
-    witgen_functions: HashMap<CacheKey<T>, Option<CacheEntry<T>>>,
+    /// If the entry is Err, we attempted to generate the function but failed.
+    witgen_functions: HashMap<CacheKey<T>, Result<CacheEntry<T>, CompilationError>>,
     column_layout: ColumnLayout,
     block_size: usize,
     machine_name: String,
@@ -49,6 +49,12 @@ pub struct CacheEntry<T: FieldElement> {
     pub range_constraints: Vec<RangeConstraint<T>>,
 }
 
+#[derive(Debug)]
+pub enum CompilationError {
+    UnsupportedField,
+    Other(String),
+}
+
 impl<'a, T: FieldElement> FunctionCache<'a, T> {
     pub fn new(
         fixed_data: &'a FixedData<'a, T>,
@@ -81,15 +87,15 @@ impl<'a, T: FieldElement> FunctionCache<'a, T> {
         bus_id: T,
         known_args: &BitVec,
         known_concrete: Option<(usize, T)>,
-    ) -> &Option<CacheEntry<T>> {
+    ) -> &Result<CacheEntry<T>, CompilationError> {
         // First try the generic version, then the specific.
         let mut key = CacheKey {
             bus_id,
             known_args: known_args.clone(),
             known_concrete: None,
         };
 
-        if self.ensure_cache(can_process.clone(), &key).is_none() && known_concrete.is_some() {
+        if self.ensure_cache(can_process.clone(), &key).is_err() && known_concrete.is_some() {
             key = CacheKey {
                 bus_id,
                 known_args: known_args.clone(),
@@ -104,15 +110,15 @@ impl<'a, T: FieldElement> FunctionCache<'a, T> {
         &mut self,
         can_process: impl CanProcessCall<T>,
         cache_key: &CacheKey<T>,
-    ) -> &Option<CacheEntry<T>> {
+    ) -> &Result<CacheEntry<T>, CompilationError> {
         if !self.witgen_functions.contains_key(cache_key) {
             record_start("Auto-witgen code derivation");
             let f = match T::known_field() {
                 // Currently, we only support the Goldilocks fields
                 Some(KnownField::GoldilocksField) => {
                     self.compile_witgen_function(can_process, cache_key)
                 }
-                _ => None,
+                _ => Err(CompilationError::UnsupportedField),
             };
             assert!(self.witgen_functions.insert(cache_key.clone(), f).is_none());
             record_end("Auto-witgen code derivation");
@@ -124,7 +130,7 @@ impl<'a, T: FieldElement> FunctionCache<'a, T> {
         &self,
         can_process: impl CanProcessCall<T>,
         cache_key: &CacheKey<T>,
-    ) -> Option<CacheEntry<T>> {
+    ) -> Result<CacheEntry<T>, CompilationError> {
         log::debug!(
             "Compiling JIT function for\n  Machine: {}\n  Connection: {}\n   Inputs: {:?}{}",
             self.machine_name,
@@ -154,10 +160,10 @@ impl<'a, T: FieldElement> FunctionCache<'a, T> {
                 // These errors can be pretty verbose and are quite common currently.
                 log::debug!(
                     "=> Error generating JIT code: {}\n...",
-                    e.to_string().lines().take(5).join("\n")
+                    e.lines().take(5).join("\n")
                 );
-            })
-            .ok()?;
+                CompilationError::Other(e)
+            })?;
 
         log::debug!("=> Success!");
         let out_of_bounds_vars = code
@@ -198,7 +204,7 @@ impl<'a, T: FieldElement> FunctionCache<'a, T> {
         .unwrap();
         log::trace!("Compilation done.");
 
-        Some(CacheEntry {
+        Ok(CacheEntry {
             function,
             range_constraints,
         })

executor/src/witgen/machines/block_machine.rs

@@ -13,7 +13,7 @@ use crate::witgen::data_structures::caller_data::CallerData;
 use crate::witgen::data_structures::finalizable_data::FinalizableData;
 use crate::witgen::data_structures::mutable_state::MutableState;
 use crate::witgen::global_constraints::RangeConstraintSet;
-use crate::witgen::jit::function_cache::FunctionCache;
+use crate::witgen::jit::function_cache::{CompilationError, FunctionCache};
 use crate::witgen::jit::witgen_inference::CanProcessCall;
 use crate::witgen::processor::{OuterQuery, Processor, SolverState};
 use crate::witgen::range_constraints::RangeConstraint;
@@ -183,8 +183,8 @@ impl<'a, T: FieldElement> Machine<'a, T> for BlockMachine<'a, T> {
             known_arguments,
             fixed_first_input,
         ) {
-            Some(entry) => (true, entry.range_constraints.clone()),
-            None => (false, range_constraints),
+            Ok(entry) => (true, entry.range_constraints.clone()),
+            Err(_) => (false, range_constraints),
         }
     }
 
@@ -454,16 +454,27 @@ impl<'a, T: FieldElement> BlockMachine<'a, T> {
         let fixed_first_input = arguments
             .first()
             .and_then(|a| a.constant_value().map(|v| (0, v)));
-        if self
-            .function_cache
-            .compile_cached(mutable_state, bus_id, &known_inputs, fixed_first_input)
-            .is_some()
-        {
-            let caller_data = CallerData::new(arguments, range_constraints);
-            let updates = self.process_lookup_via_jit(mutable_state, bus_id, caller_data)?;
-            assert!(updates.is_complete());
-            self.block_count_jit += 1;
-            return Ok(updates);
+        match self.function_cache.compile_cached(
+            mutable_state,
+            bus_id,
+            &known_inputs,
+            fixed_first_input,
+        ) {
+            Ok(_) => {
+                let caller_data = CallerData::new(arguments, range_constraints);
+                let updates = self.process_lookup_via_jit(mutable_state, bus_id, caller_data)?;
+                assert!(updates.is_complete());
+                self.block_count_jit += 1;
+                return Ok(updates);
+            }
+            Err(CompilationError::Other(_e)) => {
+                // Assuming the JIT compiler is feature-complete, this means that the witness is not
+                // unique, which could happen e.g. if not all required arguments are provided.
+                return Ok(EvalValue::incomplete(IncompleteCause::JitCompilationFailed));
+            }
+            // If we're on an unsupported field, this won't be fixed in future invocations.
+            // Fall back to run-time witgen.
+            Err(CompilationError::UnsupportedField) => {}
         }
 
         let outer_query = OuterQuery::new(