Improve logging for JIT

powdr-labs · Jan 9, 2025 · 4ffc6e6 · 4ffc6e6
1 parent e4d9fcb
commit 4ffc6e6
Show file tree

Hide file tree

Showing 5 changed files with 26 additions and 13 deletions.
diff --git a/executor/src/witgen/jit/block_machine_processor.rs b/executor/src/witgen/jit/block_machine_processor.rs
@@ -87,7 +87,7 @@ impl<'a, T: FieldElement> BlockMachineProcessor<'a, T> {
                     .join("\n  ");
                 log::debug!("Known arguments:\n  {known_args_str}");
                 log::debug!("Error:\n  {e}");
-                log::debug!(
+                log::trace!(
                     "The following code was generated so far:\n{}",
                     format_code(witgen.code())
                 );

diff --git a/executor/src/witgen/jit/function_cache.rs b/executor/src/witgen/jit/function_cache.rs
@@ -29,6 +29,7 @@ pub struct FunctionCache<'a, T: FieldElement> {
     witgen_functions: HashMap<CacheKey, Option<WitgenFunction<T>>>,
     column_layout: ColumnLayout,
     block_size: usize,
+    machine_name: String,
 }
 
 impl<'a, T: FieldElement> FunctionCache<'a, T> {
@@ -38,6 +39,7 @@ impl<'a, T: FieldElement> FunctionCache<'a, T> {
         block_size: usize,
         latch_row: usize,
         metadata: ColumnLayout,
+        machine_name: String,
     ) -> Self {
         let processor =
             BlockMachineProcessor::new(fixed_data, parts.clone(), block_size, latch_row);
@@ -47,6 +49,7 @@ impl<'a, T: FieldElement> FunctionCache<'a, T> {
             column_layout: metadata,
             witgen_functions: HashMap::new(),
             block_size,
+            machine_name,
         }
     }
 
@@ -90,10 +93,18 @@ impl<'a, T: FieldElement> FunctionCache<'a, T> {
         mutable_state: &MutableState<'a, T, Q>,
         cache_key: &CacheKey,
     ) -> Option<WitgenFunction<T>> {
-        log::trace!("Compiling JIT function for {:?}", cache_key);
+        log::debug!(
+            "Compiling JIT function for {}\n   {:?}",
+            self.machine_name,
+            cache_key
+        );
 
         self.processor
             .generate_code(mutable_state, cache_key.identity_id, &cache_key.known_args)
+            .map_err(|e| {
+                // log::error!("Error generating JIT code: {:?}", e);
+                e
+            })
             .ok()
             .map(|code| {
                 let is_in_bounds = code

diff --git a/executor/src/witgen/machines/block_machine.rs b/executor/src/witgen/machines/block_machine.rs
@@ -118,6 +118,14 @@ impl<'a, T: FieldElement> BlockMachine<'a, T> {
             fixed_data,
         );
         let layout = data.layout();
+        let function_cache = FunctionCache::new(
+            fixed_data,
+            parts.clone(),
+            block_size,
+            latch_row,
+            layout,
+            name.clone(),
+        );
         Some(BlockMachine {
             name,
             degree_range,
@@ -134,13 +142,7 @@ impl<'a, T: FieldElement> BlockMachine<'a, T> {
                 latch_row,
                 parts.identities.len(),
             ),
-            function_cache: FunctionCache::new(
-                fixed_data,
-                parts.clone(),
-                block_size,
-                latch_row,
-                layout,
-            ),
+            function_cache,
             block_count_jit: 0,
             block_count_runtime: 0,
         })

diff --git a/executor/src/witgen/mod.rs b/executor/src/witgen/mod.rs
@@ -227,7 +227,7 @@ impl<'a, 'b, T: FieldElement> WitnessGenerator<'a, 'b, T> {
                 let discard = references_later_stage_challenge || references_later_stage_witness;
 
                 if discard {
-                    log::debug!("Skipping identity that references later-stage items: {identity}",);
+                    log::trace!("Skipping identity that references later-stage items: {identity}",);
                 }
                 !discard
             })

diff --git a/executor/src/witgen/processor.rs b/executor/src/witgen/processor.rs
@@ -662,15 +662,15 @@ Known values in current row (local: {row_index}, global {global_row_index}):
             .process_identity(identity, &row_pair)
             .is_err()
         {
-            log::debug!(
+            log::trace!(
                 "Previous {}",
                 self.data[row_index - 1].render_values(true, self.parts)
             );
-            log::debug!(
+            log::trace!(
                 "Proposed {:?}",
                 proposed_row.render_values(true, self.parts)
             );
-            log::debug!("Failed on identity: {}", identity);
+            log::trace!("Failed on identity: {}", identity);
 
             return false;
         }