diff --git a/executor/src/witgen/jit/compiler.rs b/executor/src/witgen/jit/compiler.rs index b15e4ba84c..4fbdc98eba 100644 --- a/executor/src/witgen/jit/compiler.rs +++ b/executor/src/witgen/jit/compiler.rs @@ -223,7 +223,6 @@ fn witgen_code( let load_known_inputs = known_inputs .iter() .map(|v| { - let var_name = variable_to_string(v); let value = match v { Variable::WitnessCell(c) => { format!("get(data, row_offset, {}, {})", c.row_offset, c.id) @@ -237,7 +236,8 @@ fn witgen_code( unreachable!("Machine call variables should not be pre-known.") } }; - format!(" let {var_name} = {value};") + + indent(set(v, &value, true, false), 1) }) .format("\n"); @@ -252,12 +252,11 @@ fn witgen_code( }) .unique() .map(|(var, cell)| { - format!( - " let {} = get_fixed_value(fixed_data, {}, (row_offset + {}));", - variable_to_string(var), - cell.id, - cell.row_offset, - ) + let value = format!( + "get_fixed_value(fixed_data, {}, (row_offset + {}))", + cell.id, cell.row_offset + ); + indent(set(var, &value, true, false), 1) }) .format("\n"); @@ -270,12 +269,9 @@ fn witgen_code( let store_values = vars_known .iter() .filter_map(|var| { - let value = variable_to_string(var); + let value = get(var); match var { - Variable::WitnessCell(cell) => Some(format!( - " set(data, row_offset, {}, {}, {value});", - cell.row_offset, cell.id, - )), + Variable::WitnessCell(_) => None, Variable::Param(i) => Some(format!(" set_param(params, {i}, {value});")), Variable::FixedCell(_) => panic!("Fixed columns should not be written to."), Variable::IntermediateCell(_) => { @@ -363,14 +359,7 @@ fn format_effects_inner( fn format_effect(effect: &Effect, is_top_level: bool) -> String { match effect { - Effect::Assignment(var, e) => { - format!( - "{}{} = {};", - if is_top_level { "let " } else { "" }, - variable_to_string(var), - format_expression(e) - ) - } + Effect::Assignment(var, e) => set(var, &format_expression(e), is_top_level, false), Effect::RangeConstraint(..) => { unreachable!("Final code should not contain pure range constraints.") } @@ -390,22 +379,30 @@ fn format_effect(effect: &Effect, is_top_level: bo .iter() .zip_eq(known) .map(|(v, known)| { - let var_name = variable_to_string(v); if known { - format!("LookupCell::Input(&{var_name})") + let input = get(v); + format!("LookupCell::Input(&{input})") } else { if is_top_level { - result_vars.push(var_name.clone()); + result_vars.push(v); } + // Assumes that get returns a simple variable name. + let var_name = get(v); format!("LookupCell::Output(&mut {var_name})") } }) .format(", ") .to_string(); - let var_decls = result_vars - .iter() - .map(|var_name| format!("let mut {var_name} = FieldElement::default();\n")) - .format(""); + let var_decls = if result_vars.is_empty() { + "".to_string() + } else { + result_vars + .iter() + .map(|var_name| set(var_name, "FieldElement::default()", is_top_level, true)) + .format("\n") + .to_string() + + "\n" + }; format!( "{var_decls}assert!(call_machine(mutable_state, {id}.into(), MutSlice::from((&mut [{args}]).as_mut_slice())));" ) @@ -416,12 +413,17 @@ fn format_effect(effect: &Effect, is_top_level: bo row_offset, inputs, }) => { - format!( - "{}[{}] = prover_function_{function_index}(mutable_state, input_from_channel, output_to_channel, row_offset + {row_offset}, &[{}]);", - if is_top_level { "let " } else { "" }, - targets.iter().map(variable_to_string).format(", "), - inputs.iter().map(variable_to_string).format(", ") - ) + let function_call = format!( + "let result = prover_function_{function_index}(mutable_state, input_from_channel, output_to_channel, row_offset + {row_offset}, &[{}]);", + inputs.iter().map(get).format(", ") + ); + let store_results = targets + .iter() + .enumerate() + .map(|(i, v)| set(v, &format!("result[{i}]"), is_top_level, false)) + .format("\n"); + let block = format!("{function_call}\n{store_results}"); + format!("{{\n{}\n}}", indent(block, 1)) } Effect::Branch(condition, first, second) => { let var_decls = if is_top_level { @@ -434,14 +436,9 @@ fn format_effect(effect: &Effect, is_top_level: bo .sorted() .dedup() .map(|(v, needs_mut)| { - let v = variable_to_string(v); - if needs_mut { - format!("let mut {v} = FieldElement::default();\n") - } else { - format!("let {v};\n") - } + set(v, "FieldElement::default()", is_top_level, needs_mut) }) - .format("") + .format("\n") .to_string() } else { "".to_string() @@ -465,7 +462,7 @@ fn format_effect(effect: &Effect, is_top_level: bo fn format_expression(e: &SymbolicExpression) -> String { match e { SymbolicExpression::Concrete(v) => format!("FieldElement::from({v})"), - SymbolicExpression::Symbol(symbol, _) => variable_to_string(symbol), + SymbolicExpression::Symbol(symbol, _) => get(symbol), SymbolicExpression::BinaryOperation(left, op, right, _) => { let left = format_expression(left); let right = format_expression(right); @@ -498,7 +495,7 @@ fn format_condition( condition, }: &BranchCondition, ) -> String { - let var = format!("IntType::from({})", variable_to_string(variable)); + let var = format!("IntType::from({})", get(variable)); let (min, max) = condition.range(); match min.cmp(&max) { Ordering::Equal => format!("{var} == {min}",), @@ -507,6 +504,33 @@ fn format_condition( } } +fn set(v: &Variable, value: &str, is_top_level: bool, needs_mut: bool) -> String { + match v { + Variable::WitnessCell(cell) => { + format!( + "set(data, row_offset, {}, {}, {value});", + cell.row_offset, cell.id + ) + } + _ => format!( + "{}{}{} = {};", + if is_top_level { "let " } else { "" }, + if needs_mut { "mut " } else { "" }, + variable_to_string(v), + value + ), + } +} + +fn get(v: &Variable) -> String { + match v { + Variable::WitnessCell(cell) => { + format!("get(data, row_offset, {}, {})", cell.row_offset, cell.id) + } + _ => variable_to_string(v), + } +} + /// Returns the name of a local (stack) variable for the given expression variable. fn variable_to_string(v: &Variable) -> String { match v { diff --git a/executor/src/witgen/jit/function_cache.rs b/executor/src/witgen/jit/function_cache.rs index 74ab4878e9..633850d854 100644 --- a/executor/src/witgen/jit/function_cache.rs +++ b/executor/src/witgen/jit/function_cache.rs @@ -125,7 +125,7 @@ impl<'a, T: FieldElement> FunctionCache<'a, T> { can_process: impl CanProcessCall, cache_key: &CacheKey, ) -> Option> { - log::debug!( + log::info!( "Compiling JIT function for\n Machine: {}\n Connection: {}\n Inputs: {:?}{}", self.machine_name, self.parts.bus_receives[&cache_key.bus_id], @@ -152,14 +152,14 @@ impl<'a, T: FieldElement> FunctionCache<'a, T> { ) .map_err(|e| { // These errors can be pretty verbose and are quite common currently. - log::debug!( + log::info!( "=> Error generating JIT code: {}\n...", e.to_string().lines().take(5).join("\n") ); }) .ok()?; - log::debug!("=> Success!"); + log::info!("=> Success!"); let out_of_bounds_vars = code .iter() .flat_map(|effect| effect.referenced_variables()) @@ -179,7 +179,7 @@ impl<'a, T: FieldElement> FunctionCache<'a, T> { ); } - log::trace!("Generated code ({} steps)", code.len()); + log::info!("Generated code ({} steps)", code.len()); let known_inputs = cache_key .known_args .iter() @@ -187,7 +187,7 @@ impl<'a, T: FieldElement> FunctionCache<'a, T> { .filter_map(|(i, b)| if b { Some(Variable::Param(i)) } else { None }) .collect::>(); - log::trace!("Compiling effects..."); + log::info!("Compiling effects..."); let function = compile_effects( self.fixed_data.analyzed, self.column_layout.clone(), @@ -196,7 +196,7 @@ impl<'a, T: FieldElement> FunctionCache<'a, T> { prover_functions, ) .unwrap(); - log::trace!("Compilation done."); + log::info!("Compilation done."); Some(CacheEntry { function, diff --git a/jit-compiler/src/compiler.rs b/jit-compiler/src/compiler.rs index 55bfd2c48c..517b774637 100644 --- a/jit-compiler/src/compiler.rs +++ b/jit-compiler/src/compiler.rs @@ -1,6 +1,7 @@ use mktemp::Temp; use std::{ fs::{self}, + path::PathBuf, process::Command, str::from_utf8, sync::Arc, @@ -83,11 +84,24 @@ fn cargo_toml(opt_level: Option) -> String { } } +const DEBUG: bool = false; + /// Compiles the given code and returns the path to the /// temporary directory containing the compiled library /// and the path to the compiled library. pub fn call_cargo(code: &str, opt_level: Option) -> Result { - let dir = mktemp::Temp::new_dir().unwrap(); + let dir_tmp = mktemp::Temp::new_dir().unwrap(); + + let dir = if DEBUG { + let dir = PathBuf::from("../cargo_dir"); + // rm -r cargo_dir/* + fs::remove_dir_all(&dir).ok(); + fs::create_dir(&dir).unwrap(); + dir + } else { + dir_tmp.as_path().to_path_buf() + }; + fs::write(dir.join("Cargo.toml"), cargo_toml(opt_level)).unwrap(); fs::create_dir(dir.join("src")).unwrap(); fs::write(dir.join("src").join("lib.rs"), code).unwrap(); @@ -132,7 +146,7 @@ pub fn call_cargo(code: &str, opt_level: Option) -> Result