Specialize functions for operation id. (#2440)

chriseth · web-flow · commit c5202a9d834a · 2025-02-10T18:36:24.000Z
If one item in the RHS of the connection has `operation_id` in its name
and has a known concrete value, create a specialized function for this
value.
diff --git a/executor/src/witgen/jit/block_machine_processor.rs b/executor/src/witgen/jit/block_machine_processor.rs
@@ -54,6 +54,7 @@ impl<'a, T: FieldElement> BlockMachineProcessor<'a, T> {
         can_process: impl CanProcessCall<T>,
         identity_id: u64,
         known_args: &BitVec,
+        known_concrete: Option<(usize, T)>,
     ) -> Result<(ProcessorResult<T>, Vec<ProverFunction<'a, T>>), String> {
         let connection = self.machine_parts.connections[&identity_id];
         assert_eq!(connection.right.expressions.len(), known_args.len());
@@ -77,6 +78,15 @@ impl<'a, T: FieldElement> BlockMachineProcessor<'a, T> {
             T::one(),
         ));
 
+        if let Some((index, value)) = known_concrete {
+            // Set the known argument to the concrete value.
+            assignments.push(Assignment::assign_constant(
+                &connection.right.expressions[index],
+                self.latch_row as i32,
+                value,
+            ));
+        }
+
         // Set all other selectors to 0 in the latch row.
         for other_connection in self.machine_parts.connections.values() {
             let other_selector = &other_connection.right.selector;
@@ -263,7 +273,7 @@ mod test {
         );
 
         processor
-            .generate_code(&mutable_state, connection_id, &known_values)
+            .generate_code(&mutable_state, connection_id, &known_values, None)
             .map(|(result, _)| result)
     }
 
diff --git a/executor/src/witgen/jit/function_cache.rs b/executor/src/witgen/jit/function_cache.rs
@@ -26,8 +26,11 @@ use super::{
 };
 
 #[derive(Debug, Clone, Hash, PartialEq, Eq)]
-struct CacheKey {
+struct CacheKey<T: FieldElement> {
     identity_id: u64,
+    /// If `Some((index, value))`, then this function is used only if the
+    /// `index`th argument is set to `value`.
+    known_concrete: Option<(usize, T)>,
     known_args: BitVec,
 }
 
@@ -37,7 +40,7 @@ pub struct FunctionCache<'a, T: FieldElement> {
     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, Option<CacheEntry<T>>>,
+    witgen_functions: HashMap<CacheKey<T>, Option<CacheEntry<T>>>,
     column_layout: ColumnLayout,
     block_size: usize,
     machine_name: String,
@@ -72,23 +75,26 @@ impl<'a, T: FieldElement> FunctionCache<'a, T> {
         }
     }
 
-    /// Compiles the JIT function for the given identity and known arguments.
+    /// Compiles the JIT function for the given identity and known arguments, and a potentially
+    /// fully known argument.
     /// Returns the function and the output range constraints if the function was successfully compiled.
     pub fn compile_cached(
         &mut self,
         can_process: impl CanProcessCall<T>,
         identity_id: u64,
         known_args: &BitVec,
+        known_concrete: Option<(usize, T)>,
     ) -> Option<&CacheEntry<T>> {
         let cache_key = CacheKey {
             identity_id,
             known_args: known_args.clone(),
+            known_concrete,
         };
         self.ensure_cache(can_process, &cache_key);
         self.witgen_functions.get(&cache_key).unwrap().as_ref()
     }
 
-    fn ensure_cache(&mut self, can_process: impl CanProcessCall<T>, cache_key: &CacheKey) {
+    fn ensure_cache(&mut self, can_process: impl CanProcessCall<T>, cache_key: &CacheKey<T>) {
         if self.witgen_functions.contains_key(cache_key) {
             return;
         }
@@ -108,13 +114,17 @@ impl<'a, T: FieldElement> FunctionCache<'a, T> {
     fn compile_witgen_function(
         &self,
         can_process: impl CanProcessCall<T>,
-        cache_key: &CacheKey,
+        cache_key: &CacheKey<T>,
     ) -> Option<CacheEntry<T>> {
         log::debug!(
-            "Compiling JIT function for\n  Machine: {}\n  Connection: {}\n   Inputs: {:?}",
+            "Compiling JIT function for\n  Machine: {}\n  Connection: {}\n   Inputs: {:?}{}",
             self.machine_name,
             self.parts.connections[&cache_key.identity_id],
-            cache_key.known_args
+            cache_key.known_args,
+            cache_key
+                .known_concrete
+                .map(|(i, v)| format!("\n   Input {i} = {v}"))
+                .unwrap_or_default()
         );
 
         let (
@@ -125,7 +135,12 @@ impl<'a, T: FieldElement> FunctionCache<'a, T> {
             prover_functions,
         ) = self
             .processor
-            .generate_code(can_process, cache_key.identity_id, &cache_key.known_args)
+            .generate_code(
+                can_process,
+                cache_key.identity_id,
+                &cache_key.known_args,
+                cache_key.known_concrete,
+            )
             .map_err(|e| {
                 // These errors can be pretty verbose and are quite common currently.
                 log::debug!(
@@ -202,6 +217,7 @@ impl<'a, T: FieldElement> FunctionCache<'a, T> {
         connection_id: u64,
         values: &mut [LookupCell<'c, T>],
         data: CompactDataRef<'d, T>,
+        known_concrete: Option<(usize, T)>,
     ) -> Result<bool, EvalError<T>> {
         let known_args = values
             .iter()
@@ -211,8 +227,12 @@ impl<'a, T: FieldElement> FunctionCache<'a, T> {
         let cache_key = CacheKey {
             identity_id: connection_id,
             known_args,
+            known_concrete,
         };
 
+        // TODO If the function is not in the cache, we should also try with
+        // known_concrete set to None.
+
         self.witgen_functions
             .get(&cache_key)
             .expect("Need to call compile_cached() first!")
diff --git a/executor/src/witgen/jit/interpreter.rs b/executor/src/witgen/jit/interpreter.rs
@@ -535,7 +535,7 @@ mod test {
 
         // TODO we cannot compile the prover functions here, but we can evaluate them.
         let (result, _prover_functions) = processor
-            .generate_code(&mutable_state, connection_id, &known_values)
+            .generate_code(&mutable_state, connection_id, &known_values, None)
             .unwrap();
 
         let known_inputs = (0..12).map(Variable::Param).collect::<Vec<_>>();
diff --git a/executor/src/witgen/machines/block_machine.rs b/executor/src/witgen/machines/block_machine.rs
@@ -170,12 +170,17 @@ impl<'a, T: FieldElement> Machine<'a, T> for BlockMachine<'a, T> {
         can_process: impl CanProcessCall<T>,
         identity_id: u64,
         known_arguments: &BitVec,
-        _range_constraints: &[RangeConstraint<T>],
+        range_constraints: &[RangeConstraint<T>],
     ) -> Option<Vec<RangeConstraint<T>>> {
-        // We do not use the input range constraints because then we would need
-        // to generate new code depending on the range constraints as well.
+        // We use the input range constraints to see if there is a column
+        // containing the substring "operation_id" which is constrained to a
+        // single value and use that value as part of the cache key.
+        let operation_id = self.find_operation_id(identity_id).and_then(|index| {
+            let v = range_constraints[index].try_to_single_value()?;
+            Some((index, v))
+        });
         self.function_cache
-            .compile_cached(can_process, identity_id, known_arguments)
+            .compile_cached(can_process, identity_id, known_arguments, operation_id)
             .map(|r| r.range_constraints.clone())
     }
 
@@ -189,12 +194,23 @@ impl<'a, T: FieldElement> Machine<'a, T> for BlockMachine<'a, T> {
             return Err(EvalError::RowsExhausted(self.name.clone()));
         }
 
+        let operation_id =
+            self.find_operation_id(identity_id)
+                .and_then(|index| match &values[index] {
+                    LookupCell::Input(v) => Some((index, **v)),
+                    LookupCell::Output(_) => None,
+                });
+
         self.data.finalize_all();
         let data = self.data.append_new_finalized_rows(self.block_size);
 
-        let success =
-            self.function_cache
-                .process_lookup_direct(mutable_state, identity_id, values, data)?;
+        let success = self.function_cache.process_lookup_direct(
+            mutable_state,
+            identity_id,
+            values,
+            data,
+            operation_id,
+        )?;
         assert!(success);
         self.block_count_jit += 1;
         Ok(true)
@@ -443,9 +459,13 @@ impl<'a, T: FieldElement> BlockMachine<'a, T> {
             .iter()
             .map(|e| e.is_constant())
             .collect();
+        let operation_id = self.find_operation_id(identity_id).and_then(|index| {
+            let v = arguments[index].constant_value()?;
+            Some((index, v))
+        });
         if self
             .function_cache
-            .compile_cached(mutable_state, identity_id, &known_inputs)
+            .compile_cached(mutable_state, identity_id, &known_inputs, operation_id)
             .is_some()
         {
             let updates = self.process_lookup_via_jit(mutable_state, identity_id, outer_query)?;
@@ -510,26 +530,41 @@ impl<'a, T: FieldElement> BlockMachine<'a, T> {
         identity_id: u64,
         outer_query: OuterQuery<'a, 'b, T>,
     ) -> EvalResult<'a, T> {
-        let mut values = CallerData::from(&outer_query);
-
         assert!(
             (self.rows() + self.block_size as DegreeType) <= self.degree,
             "Block machine is full (this should have been checked before)"
         );
         self.data.finalize_all();
+
+        let mut values = CallerData::from(&outer_query);
+        let mut lookup_cells = values.as_lookup_cells();
+        let operation_id =
+            self.find_operation_id(identity_id)
+                .and_then(|index| match &lookup_cells[index] {
+                    LookupCell::Input(v) => Some((index, **v)),
+                    LookupCell::Output(_) => None,
+                });
         let data = self.data.append_new_finalized_rows(self.block_size);
 
         let success = self.function_cache.process_lookup_direct(
             mutable_state,
             identity_id,
-            &mut values.as_lookup_cells(),
+            &mut lookup_cells,
             data,
+            operation_id,
         )?;
         assert!(success);
 
         values.into()
     }
 
+    fn find_operation_id(&self, identity_id: u64) -> Option<usize> {
+        let right = &self.parts.connections[&identity_id].right.expressions;
+        right.iter().position(|r| {
+            try_to_simple_poly(r).is_some_and(|poly| poly.name.contains("operation_id"))
+        })
+    }
+
     fn process<'b, Q: QueryCallback<T>>(
         &self,
         mutable_state: &MutableState<'a, T, Q>,
