Skip to content

Commit 88b3683

Browse files
TomWambsgansTomWambsgans
committed
precompute some poseidons in parallel before sequential witness generation
1 parent baf67a7 commit 88b3683

File tree

18 files changed

+387
-148
lines changed

18 files changed

+387
-148
lines changed

‎crates/lean_compiler/src/c_compile_final.rs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,8 @@ struct Compiler {
3636

3737
pub fn compile_to_low_level_bytecode(
3838
mut intermediate_bytecode: IntermediateBytecode,
39+
program: String,
40+
function_locations: BTreeMap<usize, String>,
3941
) -> Result<Bytecode, String> {
4042
intermediate_bytecode.bytecode.insert(
4143
Label::EndProgram,
@@ -121,6 +123,8 @@ pub fn compile_to_low_level_bytecode(
121123
hints,
122124
starting_frame_memory,
123125
ending_pc,
126+
program,
127+
function_locations,
124128
})
125129
}
126130

‎crates/lean_compiler/src/lib.rs‎

Lines changed: 9 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,3 @@
1-
use std::collections::BTreeMap;
2-
31
use lean_vm::*;
42

53
use crate::{
@@ -16,38 +14,35 @@ mod parser;
1614
mod precompiles;
1715
pub use precompiles::PRECOMPILES;
1816

19-
pub fn compile_program(program: &str) -> (Bytecode, BTreeMap<usize, String>) {
20-
let (parsed_program, function_locations) = parse_program(program).unwrap();
17+
pub fn compile_program(program: String) -> Bytecode {
18+
let (parsed_program, function_locations) = parse_program(&program).unwrap();
2119
// println!("Parsed program: {}", parsed_program.to_string());
2220
let simple_program = simplify_program(parsed_program);
2321
// println!("Simplified program: {}", simple_program.to_string());
2422
let intermediate_bytecode = compile_to_intermediate_bytecode(simple_program).unwrap();
2523
// println!("Intermediate Bytecode:\n\n{}", intermediate_bytecode.to_string());
26-
let compiled = compile_to_low_level_bytecode(intermediate_bytecode).unwrap();
24+
2725
// println!("Function Locations: \n");
2826
// for (loc, name) in function_locations.iter() {
2927
// println!("{name}: {loc}");
3028
// }
3129
// println!("\n\nCompiled Program:\n\n{compiled}");
32-
(compiled, function_locations)
30+
compile_to_low_level_bytecode(intermediate_bytecode, program, function_locations).unwrap()
3331
}
3432

3533
pub fn compile_and_run(
36-
program: &str,
37-
public_input: &[F],
38-
private_input: &[F],
34+
program: String,
35+
(public_input, private_input): (&[F], &[F]),
3936
no_vec_runtime_memory: usize, // size of the "non-vectorized" runtime memory
4037
profiler: bool,
4138
) {
42-
let (bytecode, function_locations) = compile_program(program);
39+
let bytecode = compile_program(program);
4340
execute_bytecode(
4441
&bytecode,
45-
public_input,
46-
private_input,
47-
program,
48-
&function_locations,
42+
(public_input, private_input),
4943
no_vec_runtime_memory,
5044
(profiler, true),
45+
(&vec![], &vec![]),
5146
);
5247
}
5348

‎crates/lean_compiler/tests/test_compiler.rs‎

Lines changed: 107 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,12 @@ fn test_duplicate_function_name() {
2222
return;
2323
}
2424
"#;
25-
compile_and_run(program, &[], &[], DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
25+
compile_and_run(
26+
program.to_string(),
27+
(&[], &[]),
28+
DEFAULT_NO_VEC_RUNTIME_MEMORY,
29+
false,
30+
);
2631
}
2732

2833
#[test]
@@ -36,7 +41,12 @@ fn test_duplicate_constant_name() {
3641
return;
3742
}
3843
"#;
39-
compile_and_run(program, &[], &[], DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
44+
compile_and_run(
45+
program.to_string(),
46+
(&[], &[]),
47+
DEFAULT_NO_VEC_RUNTIME_MEMORY,
48+
false,
49+
);
4050
}
4151

4252
#[test]
@@ -57,7 +67,12 @@ fn test_fibonacci_program() {
5767
return;
5868
}
5969
"#;
60-
compile_and_run(program, &[], &[], DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
70+
compile_and_run(
71+
program.to_string(),
72+
(&[], &[]),
73+
DEFAULT_NO_VEC_RUNTIME_MEMORY,
74+
false,
75+
);
6176
}
6277

6378
#[test]
@@ -75,7 +90,12 @@ fn test_edge_case_0() {
7590
return;
7691
}
7792
"#;
78-
compile_and_run(program, &[], &[], DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
93+
compile_and_run(
94+
program.to_string(),
95+
(&[], &[]),
96+
DEFAULT_NO_VEC_RUNTIME_MEMORY,
97+
false,
98+
);
7999
}
80100

81101
#[test]
@@ -88,7 +108,12 @@ fn test_edge_case_1() {
88108
return;
89109
}
90110
"#;
91-
compile_and_run(program, &[], &[], DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
111+
compile_and_run(
112+
program.to_string(),
113+
(&[], &[]),
114+
DEFAULT_NO_VEC_RUNTIME_MEMORY,
115+
false,
116+
);
92117
}
93118

94119
#[test]
@@ -106,7 +131,12 @@ fn test_edge_case_2() {
106131
return;
107132
}
108133
"#;
109-
compile_and_run(program, &[], &[], DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
134+
compile_and_run(
135+
program.to_string(),
136+
(&[], &[]),
137+
DEFAULT_NO_VEC_RUNTIME_MEMORY,
138+
false,
139+
);
110140
}
111141

112142
#[test]
@@ -122,7 +152,12 @@ fn test_decompose_bits() {
122152
return;
123153
}
124154
"#;
125-
compile_and_run(program, &[], &[], DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
155+
compile_and_run(
156+
program.to_string(),
157+
(&[], &[]),
158+
DEFAULT_NO_VEC_RUNTIME_MEMORY,
159+
false,
160+
);
126161
}
127162

128163
#[test]
@@ -138,7 +173,12 @@ fn test_unroll() {
138173
return;
139174
}
140175
"#;
141-
compile_and_run(program, &[], &[], DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
176+
compile_and_run(
177+
program.to_string(),
178+
(&[], &[]),
179+
DEFAULT_NO_VEC_RUNTIME_MEMORY,
180+
false,
181+
);
142182
}
143183

144184
#[test]
@@ -150,7 +190,12 @@ fn test_rev_unroll() {
150190
return;
151191
}
152192
"#;
153-
compile_and_run(program, &[], &[], DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
193+
compile_and_run(
194+
program.to_string(),
195+
(&[], &[]),
196+
DEFAULT_NO_VEC_RUNTIME_MEMORY,
197+
false,
198+
);
154199
}
155200

156201
#[test]
@@ -170,7 +215,12 @@ fn test_mini_program_0() {
170215
return;
171216
}
172217
"#;
173-
compile_and_run(program, &[], &[], DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
218+
compile_and_run(
219+
program.to_string(),
220+
(&[], &[]),
221+
DEFAULT_NO_VEC_RUNTIME_MEMORY,
222+
false,
223+
);
174224
}
175225

176226
#[test]
@@ -213,7 +263,12 @@ fn test_mini_program_1() {
213263
return;
214264
}
215265
"#;
216-
compile_and_run(program, &[], &[], DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
266+
compile_and_run(
267+
program.to_string(),
268+
(&[], &[]),
269+
DEFAULT_NO_VEC_RUNTIME_MEMORY,
270+
false,
271+
);
217272
}
218273

219274
#[test]
@@ -241,7 +296,12 @@ fn test_mini_program_2() {
241296
return sum, product;
242297
}
243298
"#;
244-
compile_and_run(program, &[], &[], DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
299+
compile_and_run(
300+
program.to_string(),
301+
(&[], &[]),
302+
DEFAULT_NO_VEC_RUNTIME_MEMORY,
303+
false,
304+
);
245305
}
246306

247307
#[test]
@@ -269,9 +329,8 @@ fn test_mini_program_3() {
269329
"#;
270330
let mut public_input: [F; 16] = (0..16).map(F::new).collect::<Vec<F>>().try_into().unwrap();
271331
compile_and_run(
272-
program,
273-
&public_input,
274-
&[],
332+
program.to_string(),
333+
(&public_input, &[]),
275334
DEFAULT_NO_VEC_RUNTIME_MEMORY,
276335
false,
277336
);
@@ -298,9 +357,8 @@ fn test_mini_program_4() {
298357
"#;
299358
let mut public_input: [F; 24] = (0..24).map(F::new).collect::<Vec<F>>().try_into().unwrap();
300359
compile_and_run(
301-
program,
302-
&public_input,
303-
&[],
360+
program.to_string(),
361+
(&public_input, &[]),
304362
DEFAULT_NO_VEC_RUNTIME_MEMORY,
305363
false,
306364
);
@@ -383,7 +441,12 @@ fn test_inlined() {
383441
return;
384442
}
385443
"#;
386-
compile_and_run(program, &[], &[], DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
444+
compile_and_run(
445+
program.to_string(),
446+
(&[], &[]),
447+
DEFAULT_NO_VEC_RUNTIME_MEMORY,
448+
false,
449+
);
387450
}
388451

389452
#[test]
@@ -435,7 +498,12 @@ fn test_match() {
435498
return x * x * x * x * x * x;
436499
}
437500
"#;
438-
compile_and_run(program, &[], &[], DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
501+
compile_and_run(
502+
program.to_string(),
503+
(&[], &[]),
504+
DEFAULT_NO_VEC_RUNTIME_MEMORY,
505+
false,
506+
);
439507
}
440508

441509
// #[test]
@@ -453,7 +521,7 @@ fn test_match() {
453521
// return 1;
454522
// }
455523
// "#;
456-
// compile_and_run(program, &[], &[]);
524+
// compile_and_run(program.to_string(), (&[], &[]));
457525
// }
458526

459527
#[test]
@@ -476,7 +544,12 @@ fn test_const_functions_calling_const_functions() {
476544
}
477545
"#;
478546

479-
compile_and_run(program, &[], &[], DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
547+
compile_and_run(
548+
program.to_string(),
549+
(&[], &[]),
550+
DEFAULT_NO_VEC_RUNTIME_MEMORY,
551+
false,
552+
);
480553
}
481554

482555
#[test]
@@ -499,7 +572,12 @@ fn test_inline_functions_calling_inline_functions() {
499572
}
500573
"#;
501574

502-
compile_and_run(program, &[], &[], DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
575+
compile_and_run(
576+
program.to_string(),
577+
(&[], &[]),
578+
DEFAULT_NO_VEC_RUNTIME_MEMORY,
579+
false,
580+
);
503581
}
504582

505583
#[test]
@@ -526,5 +604,10 @@ fn test_nested_inline_functions() {
526604
}
527605
"#;
528606

529-
compile_and_run(program, &[], &[], DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
607+
compile_and_run(
608+
program.to_string(),
609+
(&[], &[]),
610+
DEFAULT_NO_VEC_RUNTIME_MEMORY,
611+
false,
612+
);
530613
}

‎crates/lean_prover/src/prove_execution.rs‎

Lines changed: 6 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,6 @@ use p3_field::Field;
1111
use p3_field::PrimeCharacteristicRing;
1212
use p3_util::{log2_ceil_usize, log2_strict_usize};
1313
use packed_pcs::*;
14-
use std::collections::BTreeMap;
1514
use tracing::info_span;
1615
use utils::ToUsize;
1716
use utils::dot_product_with_base;
@@ -24,15 +23,15 @@ use vm_air::*;
2423
use whir_p3::{
2524
WhirConfig, WhirConfigBuilder, precompute_dft_twiddles, second_batched_whir_config_builder,
2625
};
26+
use xmss::{Poseidon16History, Poseidon24History};
2727

2828
pub fn prove_execution(
2929
bytecode: &Bytecode,
30-
source_code: &str, // debug purpose
31-
function_locations: &BTreeMap<usize, String>, // debug purpose
3230
(public_input, private_input): (&[F], &[F]),
3331
whir_config_builder: WhirConfigBuilder,
3432
no_vec_runtime_memory: usize, // size of the "non-vectorized" runtime memory
3533
vm_profiler: bool,
34+
(poseidons_16_precomputed, poseidons_24_precomputed): (&Poseidon16History, &Poseidon24History),
3635
) -> (Vec<PF<EF>>, usize) {
3736
let ExecutionTrace {
3837
full_trace,
@@ -49,17 +48,14 @@ pub fn prove_execution(
4948
let execution_result = info_span!("Executing bytecode").in_scope(|| {
5049
execute_bytecode(
5150
bytecode,
52-
public_input,
53-
private_input,
54-
source_code,
55-
function_locations,
51+
(public_input, private_input),
5652
no_vec_runtime_memory,
5753
(vm_profiler, false),
54+
(poseidons_16_precomputed, poseidons_24_precomputed),
5855
)
5956
});
60-
info_span!("Building execution trace").in_scope(|| {
61-
get_execution_trace(bytecode, execution_result)
62-
})
57+
info_span!("Building execution trace")
58+
.in_scope(|| get_execution_trace(bytecode, execution_result))
6359
});
6460

6561
let public_memory = &memory[..public_memory_size];

0 commit comments

Comments
 (0)