With continuations, don't load initial memory state via code.

This commit is contained in:
Lucas Clemente Vella
2024-02-27 13:23:04 +00:00
parent 490b301382
commit b87928df8e
6 changed files with 165 additions and 71 deletions

View File

@@ -14,18 +14,16 @@ pub enum SingleDataValue<'a> {
}
struct WordWriter<'a, 'b> {
data_writer: &'a mut dyn FnMut(u32, SingleDataValue) -> Vec<String>,
data_writer: &'a mut dyn FnMut(Option<&str>, u32, SingleDataValue),
partial: u32,
current_pos: u32,
generated_code: Vec<String>,
latest_label: Option<&'b str>,
}
impl<'a, 'b> WordWriter<'a, 'b> {
fn new(
starting_pos: u32,
data_writer: &'a mut dyn FnMut(u32, SingleDataValue) -> Vec<String>,
data_writer: &'a mut dyn FnMut(Option<&str>, u32, SingleDataValue),
) -> Self {
// sanitary alignment to 8 bytes
let current_pos = next_aligned(starting_pos as usize, 8) as u32;
@@ -33,7 +31,6 @@ impl<'a, 'b> WordWriter<'a, 'b> {
partial: 0,
current_pos,
data_writer,
generated_code: Vec::new(),
latest_label: None,
}
}
@@ -52,14 +49,11 @@ impl<'a, 'b> WordWriter<'a, 'b> {
// if changed words, flush
let curr_word = self.current_pos & (!0b11);
if (next_pos & (!0b11) != curr_word) && (self.partial != 0) {
if let Some(label) = std::mem::take(&mut self.latest_label) {
self.generated_code.push(format!("// data {label}"));
}
self.generated_code.extend((*self.data_writer)(
(*self.data_writer)(
std::mem::take(&mut self.latest_label),
curr_word,
SingleDataValue::Value(self.partial),
));
);
self.partial = 0;
}
self.current_pos = next_pos;
@@ -90,28 +84,27 @@ impl<'a, 'b> WordWriter<'a, 'b> {
"reference to code labels in misaligned data section is not supported"
);
self.generated_code.extend((*self.data_writer)(
(*self.data_writer)(
std::mem::take(&mut self.latest_label),
self.current_pos,
SingleDataValue::LabelReference(label),
));
);
assert_eq!(self.partial, 0);
self.current_pos += 4;
}
fn finish(mut self) -> Vec<String> {
fn finish(mut self) {
// ensure the latest partial word is written
self.advance(4);
self.generated_code
}
}
pub fn store_data_objects(
sections: Vec<Vec<(Option<String>, Vec<DataValue>)>>,
memory_start: u32,
code_gen: &mut dyn FnMut(u32, SingleDataValue) -> Vec<String>,
) -> (Vec<String>, BTreeMap<String, u32>) {
code_gen: &mut dyn FnMut(Option<&str>, u32, SingleDataValue),
) -> BTreeMap<String, u32> {
let mut writer = WordWriter::new(memory_start, code_gen);
let positions = {
@@ -157,6 +150,7 @@ pub fn store_data_objects(
}
}
}
writer.finish();
(writer.finish(), positions)
positions
}

View File

@@ -284,6 +284,7 @@ mod builder {
/// in Err.
pub fn new<T: FieldElement>(
main: &'a Machine<T>,
mem: MemoryState,
batch_to_line_map: &'b [u32],
max_rows_len: usize,
mode: ExecMode,
@@ -326,7 +327,7 @@ mod builder {
batch_to_line_map,
max_rows: max_rows_len,
regs,
mem: HashMap::new(),
mem,
mode,
};
@@ -472,7 +473,7 @@ mod builder {
}
}
fn get_main_machine<T: FieldElement>(program: &AnalysisASMFile<T>) -> &Machine<T> {
pub fn get_main_machine<T: FieldElement>(program: &AnalysisASMFile<T>) -> &Machine<T> {
for (name, m) in program.items.iter() {
if name.len() == 1 && name.parts().next() == Some("Main") {
let Item::Machine(m) = m else {
@@ -847,6 +848,7 @@ impl<'a, 'b, F: FieldElement> Executor<'a, 'b, F> {
pub fn execute_ast<T: FieldElement>(
program: &AnalysisASMFile<T>,
initial_memory: MemoryState,
inputs: &Callback<T>,
bootloader_inputs: &[Elem<T>],
max_steps_to_execute: usize,
@@ -860,8 +862,13 @@ pub fn execute_ast<T: FieldElement>(
debug_files,
} = preprocess_main_function(main_machine);
let proc = match TraceBuilder::new(main_machine, &batch_to_line_map, max_steps_to_execute, mode)
{
let proc = match TraceBuilder::new(
main_machine,
initial_memory,
&batch_to_line_map,
max_steps_to_execute,
mode,
) {
Ok(proc) => proc,
Err(ret) => return *ret,
};
@@ -941,7 +948,14 @@ pub fn execute<F: FieldElement>(
let analyzed = powdr_analysis::analyze(resolved).unwrap();
log::info!("Executing...");
execute_ast(&analyzed, inputs, bootloader_inputs, usize::MAX, mode)
execute_ast(
&analyzed,
MemoryState::new(),
inputs,
bootloader_inputs,
usize::MAX,
mode,
)
}
fn to_u32<F: FieldElement>(val: &F) -> Option<u32> {

View File

@@ -144,34 +144,59 @@ pub fn compile(
// for compilation, and will not be called.
statements = replace_coprocessor_stubs(statements, coprocessors).collect::<Vec<_>>();
let (data_code, data_positions) =
store_data_objects(data_sections, data_start, &mut |addr, value| match value {
SingleDataValue::Value(v) => {
vec![format!("mstore 0x{addr:x}, 0x{v:x};")]
let mut initial_mem = Vec::new();
let mut data_code = Vec::new();
let data_positions =
store_data_objects(data_sections, data_start, &mut |label, addr, value| {
if let Some(label) = label {
let comment = format!(" // data {label}");
if with_bootloader && !matches!(value, SingleDataValue::LabelReference(_)) {
&mut initial_mem
} else {
&mut data_code
}
.push(comment);
}
SingleDataValue::LabelReference(sym) => {
// TODO should be possible without temporary
vec![
format!("tmp1 <== load_label({});", escape_label(sym)),
format!("mstore 0x{addr:x}, tmp1;"),
]
}
SingleDataValue::Offset(_, _) => {
unimplemented!();
/*
object_code.push(format!("addr <=X= 0x{pos:x};"));
match value {
SingleDataValue::Value(v) => {
if with_bootloader {
// Instead of generating the data loading code, we store it
// in the variable that will be used as the initial memory
// snapshot, committed by the bootloader.
initial_mem.push(format!("(0x{addr:x}, 0x{v:x})"));
} else {
// There is no bootloader to commit to memory, so we have to
// load it explicitly.
data_code.push(format!("mstore 0x{addr:x}, 0x{v:x};"));
}
}
SingleDataValue::LabelReference(sym) => {
// The label value is not known at this point, so we have to
// load it via code, irrespectively of bootloader availability.
//
// TODO should be possible without temporary
data_code.extend([
format!("tmp1 <== load_label({});", escape_label(sym)),
format!("mstore 0x{addr:x}, tmp1;"),
]);
}
SingleDataValue::Offset(_, _) => {
unimplemented!();
/*
object_code.push(format!("addr <=X= 0x{pos:x};"));
I think this solution should be fine but hard to say without
an actual code snippet that uses it.
I think this solution should be fine but hard to say without
an actual code snippet that uses it.
// TODO should be possible without temporary
object_code.extend([
format!("tmp1 <== load_label({});", escape_label(a)),
format!("tmp2 <== load_label({});", escape_label(b)),
// TODO check if registers match
"mstore wrap(tmp1 - tmp2);".to_string(),
]);
*/
// TODO should be possible without temporary
object_code.extend([
format!("tmp1 <== load_label({});", escape_label(a)),
format!("tmp2 <== load_label({});", escape_label(b)),
// TODO check if registers match
"mstore wrap(tmp1 - tmp2);".to_string(),
]);
*/
}
}
});
@@ -187,28 +212,33 @@ pub fn compile(
submachine_init
};
let program: Vec<String> = file_ids
let mut program: Vec<String> = file_ids
.into_iter()
.map(|(id, dir, file)| format!(".debug file {id} {} {};", quote(&dir), quote(&file)))
.chain(bootloader_and_shutdown_routine_lines)
.chain(["x1 <== jump(__data_init);".to_string()])
.chain([
format!("// Set stack pointer\nx2 <=X= {stack_start};"),
"x1 <== jump(__runtime_start);".to_string(),
"return;".to_string(), // This is not "riscv ret", but "return from powdr asm function".
])
.chain(
substitute_symbols_with_values(statements, &data_positions)
.into_iter()
.flat_map(|v| process_statement(v, coprocessors)),
)
.chain(["// This is the data initialization routine.\n__data_init:".to_string()])
.collect();
if !data_code.is_empty() {
program.push("x1 <== jump(__data_init);".to_string());
}
program.extend([
format!("// Set stack pointer\nx2 <=X= {stack_start};"),
"x1 <== jump(__runtime_start);".to_string(),
"return;".to_string(), // This is not "riscv ret", but "return from powdr asm function".
]);
program.extend(
substitute_symbols_with_values(statements, &data_positions)
.into_iter()
.flat_map(|v| process_statement(v, coprocessors)),
);
if !data_code.is_empty() {
program.extend(
["// This is the data initialization routine.\n__data_init:".to_string()].into_iter()
.chain(data_code)
.chain([
"// This is the end of the data initialization routine.\ntmp1 <== jump_dyn(x1);"
.to_string(),
])
.collect();
]));
}
// The program ROM needs to fit the degree, so we use the next power of 2.
let degree = program.len().ilog2() + 1;
@@ -229,6 +259,7 @@ pub fn compile(
riscv_machine(
&coprocessors.machine_imports(),
&preamble(degree, coprocessors, with_bootloader),
initial_mem,
&coprocessors.declarations(),
program,
)
@@ -407,6 +438,7 @@ fn substitute_symbols_with_values(
fn riscv_machine(
machines: &[&str],
preamble: &str,
initial_memory: Vec<String>,
submachines: &[(&str, &str)],
program: Vec<String>,
) -> String {
@@ -418,6 +450,10 @@ machine Main {{
{}
let initial_memory = [
{}
];
function main {{
{}
}}
@@ -430,6 +466,7 @@ machine Main {{
.collect::<Vec<_>>()
.join("\n"),
preamble,
initial_memory.join(",\n"),
program
.into_iter()
.map(|line| format!("\t\t{line}"))

View File

@@ -2,11 +2,11 @@ use std::collections::{BTreeSet, HashMap};
use powdr_ast::{
asm_analysis::{AnalysisASMFile, RegisterTy},
parsed::asm::parse_absolute_path,
parsed::{asm::parse_absolute_path, Expression, PilStatement},
};
use powdr_number::FieldElement;
use powdr_number::{BigInt, FieldElement};
use powdr_pipeline::Pipeline;
use powdr_riscv_executor::{Elem, ExecutionTrace};
use powdr_riscv_executor::{get_main_machine, Elem, ExecutionTrace, MemoryState};
pub mod bootloader;
mod memory_merkle_tree;
@@ -134,15 +134,49 @@ fn sanity_check<T>(program: &AnalysisASMFile<T>) {
assert_eq!(machine_registers, expected_registers);
}
fn load_initial_memory<F: FieldElement>(program: &AnalysisASMFile<F>) -> MemoryState {
let machine = get_main_machine(program);
let Some(expr) = machine.pil.iter().find_map(|v| match v {
PilStatement::LetStatement(_, n, None, expr) if n == "initial_memory" => expr.as_ref(),
_ => None,
}) else {
log::warn!("No initial_memory variable found in the machine. Assuming zeroed memory.");
return MemoryState::default();
};
let Expression::ArrayLiteral(array) = expr else {
panic!("initial_memory is not an array literal");
};
array
.items
.iter()
.map(|entry| {
let Expression::Tuple(tuple) = entry else {
panic!("initial_memory entry is not a tuple");
};
assert_eq!(tuple.len(), 2);
let Expression::Number(key, None) = &tuple[0] else {
panic!("initial_memory entry key is not a number");
};
let Expression::Number(value, None) = &tuple[1] else {
panic!("initial_memory entry value is not a number");
};
(
key.to_integer().try_into_u32().unwrap(),
value.to_integer().try_into_u32().unwrap(),
)
})
.collect()
}
/// Runs the entire execution using the RISC-V executor. For each chunk, it collects:
/// - The inputs to the bootloader, needed to restore the correct state.
/// - The number of rows after which the prover should jump to the shutdown routine.
pub fn rust_continuations_dry_run<F: FieldElement>(
pipeline: &mut Pipeline<F>,
) -> Vec<(Vec<F>, u64)> {
log::info!("Initializing memory merkle tree...");
let mut merkle_tree = MerkleTree::<F>::new();
// All inputs for all chunks.
let mut bootloader_inputs_and_num_rows = vec![];
@@ -152,10 +186,24 @@ pub fn rust_continuations_dry_run<F: FieldElement>(
let program = pipeline.compute_analyzed_asm().unwrap().clone();
sanity_check(&program);
log::info!("Initializing memory merkle tree...");
// Get initial memory contents from the special variable "initial_memory".
// In the first full run, we use it as the memory contents of the executor;
// on the independent chunk runs, the executor uses zeroed initial memory,
// and the pages are loaded via the bootloader.
let initial_memory = load_initial_memory(&program);
let mut merkle_tree = MerkleTree::<F>::new();
merkle_tree.update(initial_memory.iter().map(|(k, v)| (*k, *v)));
// TODO: commit to the merkle_tree root in the verifier.
log::info!("Executing powdr-asm...");
let (full_trace, memory_accesses) = {
let trace = powdr_riscv_executor::execute_ast::<F>(
&program,
initial_memory,
pipeline.data_callback().unwrap(),
// Run full trace without any accessed pages. This would actually violate the
// constraints, but the executor does the right thing (read zero if the memory
@@ -248,6 +296,7 @@ pub fn rust_continuations_dry_run<F: FieldElement>(
let (chunk_trace, memory_snapshot_update) = {
let (trace, memory_snapshot_update) = powdr_riscv_executor::execute_ast::<F>(
&program,
MemoryState::new(),
pipeline.data_callback().unwrap(),
&bootloader_inputs,
num_rows,

View File

@@ -93,7 +93,6 @@ impl<T: FieldElement, const N_LEAVES_LOG: usize, const WORDS_PER_PAGE: usize>
/// Applies updates, given an iterator of (memory address, value) pairs.
/// Memory addresses are assumed to be word-aligned.
#[allow(dead_code)]
pub fn update(&mut self, updates: impl Iterator<Item = (u32, u32)>) {
for (page_index, updates) in self.organize_updates_by_page(updates) {
self.update_page(page_index, updates.into_iter());

View File

@@ -13,6 +13,7 @@ pub fn verify_riscv_asm_string(file_name: &str, contents: &str, inputs: Vec<Gold
let analyzed = pipeline.compute_analyzed_asm().unwrap().clone();
powdr_riscv_executor::execute_ast(
&analyzed,
Default::default(),
pipeline.data_callback().unwrap(),
// Assume the RISC-V program was compiled without a bootloader, otherwise this will fail.
&[],