Skip to content

Commit

Permalink
simplify zkmips commands
Browse files Browse the repository at this point in the history
  • Loading branch information
@weilzkm
weilzkm committed Aug 17, 2024
1 parent 8c25355 commit d961255
Show file tree
Hide file tree
Showing 2 changed files with 46 additions and 132 deletions.
30 changes: 9 additions & 21 deletions prover/examples/README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,32 +15,20 @@ GOOS=linux GOARCH=mips GOMIPS=softfloat go build hello.go
```
BASEDIR=./emulator/test-vectors RUST_LOG=info ELF_PATH=./emulator/test-vectors/minigeth BLOCK_NO=13284491 SEG_OUTPUT=/tmp/output SEG_SIZE=1024 ARGS="" \
cargo run --release --example zkmips split
OR
RUST_LOG=info ELF_PATH=./emulator/test-vectors/hello SEG_OUTPUT=/tmp/output SEG_SIZE=1024 ARGS="" \
cargo run --release --example zkmips split_without_preimage
```

* Generate proof for each segment

```
BASEDIR=./emulator/test-vectors RUST_LOG=info BLOCK_NO=13284491 SEG_FILE="/tmp/output/0" SEG_SIZE=1024 \
cargo run --release --example zkmips prove
```

* Aggregate proof
* Generate proof for specific segment (Set SEG_START_ID to specific segment id and set SEG_NUM to 1)

```
BASEDIR=./emulator/test-vectors RUST_LOG=info BLOCK_NO=13284491 SEG_FILE="/tmp/output/0" SEG_FILE2="/tmp/output/1" SEG_SIZE=1024 \
cargo run --release --example zkmips aggregate_proof
BASEDIR=./emulator/test-vectors RUST_LOG=info BLOCK_NO=13284491 SEG_FILE_DIR="/tmp/output" SEG_START_ID=0 SEG_NUM=1 SEG_SIZE=1024 \
cargo run --release --example zkmips prove_segments
```

* Aggregate proof all
* Aggregate proof all segments (Set SEG_START_ID to 0, and set SEG_NUM to the total segments number)

```
BASEDIR=./emulator/test-vectors RUST_LOG=info BLOCK_NO=13284491 SEG_FILE_DIR="/tmp/output" SEG_FILE_NUM=299 SEG_SIZE=1024 \
cargo run --release --example zkmips aggregate_proof_all
BASEDIR=./emulator/test-vectors RUST_LOG=info BLOCK_NO=13284491 SEG_FILE_DIR="/tmp/output" SEG_START_ID=0 SEG_NUM=299 SEG_SIZE=1024 \
cargo run --release --example zkmips prove_segments
```

### Prove Go sdk code
Expand All @@ -57,7 +45,7 @@ cd ../../
* Run the host program

```
RUST_LOG=info ELF_PATH=examples/add-go/go-add HOST_PROGRAM=add_example SEG_OUTPUT=/tmp/output SEG_SIZE=262144 cargo run --release --example zkmips host
RUST_LOG=info ELF_PATH=examples/add-go/go-add HOST_PROGRAM=add_example SEG_OUTPUT=/tmp/output SEG_SIZE=262144 cargo run --release --example zkmips prove_host_program
```

## Prove the Rust code
Expand Down Expand Up @@ -92,9 +80,9 @@ cargo build --target=mips-unknown-linux-musl
cd ../..
ARGS="711e9609339e92b03ddc0a211827dba421f38f9ed8b9d806e1ffdd8c15ffa03d world!" RUST_LOG=info ELF_PATH=examples/sha2/target/mips-unknown-linux-musl/debug/sha2-bench HOST_PROGRAM=sha2_bench SEG_OUTPUT=/tmp/output cargo run --release --example zkmips host
ARGS="711e9609339e92b03ddc0a211827dba421f38f9ed8b9d806e1ffdd8c15ffa03d world!" RUST_LOG=info ELF_PATH=examples/sha2/target/mips-unknown-linux-musl/debug/sha2-bench HOST_PROGRAM=sha2_bench SEG_OUTPUT=/tmp/output cargo run --release --example zkmips prove_host_program
Or
RUST_LOG=info ELF_PATH=examples/revme/target/mips-unknown-linux-musl/debug/evm HOST_PROGRAM=revm JSON_PATH=../emulator/test-vectors/test.json SEG_OUTPUT=/tmp/output SEG_SIZE=262144 cargo run --release --example zkmips host
RUST_LOG=info ELF_PATH=examples/revme/target/mips-unknown-linux-musl/debug/evm HOST_PROGRAM=revm JSON_PATH=../emulator/test-vectors/test.json SEG_OUTPUT=/tmp/output SEG_SIZE=262144 cargo run --release --example zkmips prove_host_program
```
148 changes: 37 additions & 111 deletions prover/examples/zkmips.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ use zkm_prover::verifier::verify_proof;

const DEGREE_BITS_RANGE: [Range<usize>; 6] = [10..21, 12..22, 12..21, 8..21, 6..21, 13..23];

fn split_segs(load_preimage: bool) {
fn split_segments() {
// 1. split ELF into segs
let basedir = env::var("BASEDIR").unwrap_or("/tmp/cannon".to_string());
let elf_path = env::var("ELF_PATH").expect("ELF file is missing");
Expand All @@ -40,7 +40,7 @@ fn split_segs(load_preimage: bool) {

let mut state = load_elf_with_patch(&elf_path, args);
let block_path = get_block_path(&basedir, &block_no, "");
if load_preimage {
if block_no != "" {
state.load_input(&block_path);
}
let _ = split_prog_into_segs(state, &seg_path, &block_path, seg_size);
Expand Down Expand Up @@ -84,6 +84,7 @@ fn prove_multi_seg_common(
file: &str,
seg_size: usize,
seg_file_number: usize,
seg_start_id: usize,
) -> anyhow::Result<()> {
type InnerParameters = DefaultParameters;
type OuterParameters = Groth16WrapperParameters;
Expand All @@ -103,8 +104,8 @@ fn prove_multi_seg_common(
let all_circuits =
AllRecursiveCircuits::<F, C, D>::new(&all_stark, &DEGREE_BITS_RANGE, &config);

let seg_file = format!("{}/{}", seg_dir, 0);
log::info!("Process segment 0");
let seg_file = format!("{}/{}", seg_dir, seg_start_id);
log::info!("Process segment {}", seg_file);
let seg_reader = BufReader::new(File::open(seg_file)?);
let input_first = segment_kernel(basedir, block, file, seg_reader, seg_size);
let mut timing = TimingTree::new("prove root first", log::Level::Info);
Expand All @@ -114,12 +115,12 @@ fn prove_multi_seg_common(
timing.filter(Duration::from_millis(100)).print();
all_circuits.verify_root(agg_proof.clone())?;

let mut base_seg = 1;
let mut base_seg = seg_start_id + 1;
let mut is_agg = false;

if seg_file_number % 2 == 0 {
let seg_file = format!("{}/{}", seg_dir, 1);
log::info!("Process segment 1");
let seg_file = format!("{}/{}", seg_dir, seg_start_id + 1);
log::info!("Process segment {}", seg_file);
let seg_reader = BufReader::new(File::open(seg_file)?);
let input = segment_kernel(basedir, block, file, seg_reader, seg_size);
timing = TimingTree::new("prove root second", log::Level::Info);
Expand Down Expand Up @@ -148,12 +149,12 @@ fn prove_multi_seg_common(
all_circuits.verify_aggregation(&agg_proof)?;

is_agg = true;
base_seg = 2;
base_seg = seg_start_id + 2;
}

for i in 0..(seg_file_number - base_seg) / 2 {
let seg_file = format!("{}/{}", seg_dir, base_seg + (i << 1));
log::info!("Process segment {}", base_seg + (i << 1));
log::info!("Process segment {}", seg_file);
let seg_reader = BufReader::new(File::open(&seg_file)?);
let input_first = segment_kernel(basedir, block, file, seg_reader, seg_size);
let mut timing = TimingTree::new("prove root first", log::Level::Info);
Expand All @@ -164,7 +165,7 @@ fn prove_multi_seg_common(
all_circuits.verify_root(root_proof_first.clone())?;

let seg_file = format!("{}/{}", seg_dir, base_seg + (i << 1) + 1);
log::info!("Process segment {}", base_seg + (i << 1) + 1);
log::info!("Process segment {}", seg_file);
let seg_reader = BufReader::new(File::open(&seg_file)?);
let input = segment_kernel(basedir, block, file, seg_reader, seg_size);
let mut timing = TimingTree::new("prove root second", log::Level::Info);
Expand Down Expand Up @@ -301,20 +302,10 @@ fn prove_revm() {
let seg_file = format!("{seg_path}/{}", 0);
prove_single_seg_common(&seg_file, "", "", "", total_steps)
} else {
prove_multi_seg_common(&seg_path, "", "", "", seg_size, seg_num).unwrap()
prove_multi_seg_common(&seg_path, "", "", "", seg_size, seg_num, 0).unwrap()
}
}

fn prove_single_seg() {
let basedir = env::var("BASEDIR").unwrap_or("/tmp/cannon".to_string());
let block = env::var("BLOCK_NO").unwrap_or("".to_string());
let file = env::var("BLOCK_FILE").unwrap_or(String::from(""));
let seg_file = env::var("SEG_FILE").expect("Segment file is missing");
let seg_size = env::var("SEG_SIZE").unwrap_or(format!("{SEGMENT_STEPS}"));
let seg_size = seg_size.parse::<_>().unwrap_or(SEGMENT_STEPS);
prove_single_seg_common(&seg_file, &basedir, &block, &file, seg_size)
}

#[derive(Debug, Clone, Deserialize, Serialize)]
pub enum DataId {
TYPE1,
Expand Down Expand Up @@ -397,14 +388,10 @@ fn prove_add_example() {
let seg_file = format!("{seg_path}/{}", 0);
prove_single_seg_common(&seg_file, "", "", "", total_steps)
} else {
prove_multi_seg_common(&seg_path, "", "", "", seg_size, seg_num).unwrap()
prove_multi_seg_common(&seg_path, "", "", "", seg_size, seg_num, 0).unwrap()
}
}

fn prove_groth16() {
todo!()
}

fn prove_host() {
let host_program = env::var("HOST_PROGRAM").expect("host_program name is missing");
match host_program.as_str() {
Expand All @@ -415,12 +402,32 @@ fn prove_host() {
};
}

fn prove_segments() {
let basedir = env::var("BASEDIR").unwrap_or("/tmp/cannon".to_string());
let block = env::var("BLOCK_NO").unwrap_or("".to_string());
let file = env::var("BLOCK_FILE").unwrap_or(String::from(""));
let seg_dir = env::var("SEG_FILE_DIR").expect("segment file dir is missing");
let seg_num = env::var("SEG_NUM").unwrap_or("1".to_string());
let seg_num = seg_num.parse::<_>().unwrap_or(1usize);
let seg_start_id = env::var("SEG_START_ID").unwrap_or("0".to_string());
let seg_start_id = seg_start_id.parse::<_>().unwrap_or(0usize);
let seg_size = env::var("SEG_SIZE").unwrap_or(format!("{SEGMENT_STEPS}"));
let seg_size = seg_size.parse::<_>().unwrap_or(SEGMENT_STEPS);

if seg_num == 1 {
let seg_file = format!("{seg_dir}/{}", seg_start_id);
prove_single_seg_common(&seg_file, &basedir, &block, &file, seg_size)
} else {
prove_multi_seg_common(&seg_dir, &basedir, &block, &file, seg_size, seg_num, seg_start_id).unwrap()
}
}

fn main() {
env_logger::try_init().unwrap_or_default();
let args: Vec<String> = env::args().collect();
let helper = || {
log::info!(
"Help: {} split | split_without_preimage | prove | aggregate_proof | aggregate_proof_all | prove_groth16 | host",
"Help: {} split | prove_segments | prove_host_program",
args[0]
);
std::process::exit(-1);
Expand All @@ -429,92 +436,11 @@ fn main() {
helper();
}
match args[1].as_str() {
"split" => split_segs(true),
"split_without_preimage" => split_segs(false),
"prove" => prove_single_seg(),
"aggregate_proof" => aggregate_proof().unwrap(),
"aggregate_proof_all" => aggregate_proof_all().unwrap(),
"prove_groth16" => prove_groth16(),
"host" => prove_host(),
"split" => split_segments(),
"prove_segments" => prove_segments(),
"prove_host_program" => prove_host(),
_ => helper(),
};
}

fn aggregate_proof() -> anyhow::Result<()> {
type F = GoldilocksField;
const D: usize = 2;
type C = PoseidonGoldilocksConfig;

let basedir = env::var("BASEDIR").unwrap_or("/tmp/cannon".to_string());
let block = env::var("BLOCK_NO").unwrap_or("".to_string());
let file = env::var("BLOCK_FILE").unwrap_or(String::from(""));
let seg_file = env::var("SEG_FILE").expect("first segment file is missing");
let seg_file2 = env::var("SEG_FILE2").expect("The next segment file is missing");
let seg_size = env::var("SEG_SIZE").unwrap_or(format!("{SEGMENT_STEPS}"));
let seg_size = seg_size.parse::<_>().unwrap_or(SEGMENT_STEPS);
let all_stark = AllStark::<F, D>::default();
let config = StarkConfig::standard_fast_config();
// Preprocess all circuits.
let all_circuits =
AllRecursiveCircuits::<F, C, D>::new(&all_stark, &DEGREE_BITS_RANGE, &config);

let seg_reader = BufReader::new(File::open(seg_file)?);
let input_first = segment_kernel(&basedir, &block, &file, seg_reader, seg_size);
let mut timing = TimingTree::new("prove root first", log::Level::Info);
let (root_proof_first, first_public_values) =
all_circuits.prove_root(&all_stark, &input_first, &config, &mut timing)?;

timing.filter(Duration::from_millis(100)).print();
all_circuits.verify_root(root_proof_first.clone())?;

let seg_reader = BufReader::new(File::open(seg_file2)?);
let input = segment_kernel(&basedir, &block, &file, seg_reader, seg_size);
let mut timing = TimingTree::new("prove root second", log::Level::Info);
let (root_proof, public_values) =
all_circuits.prove_root(&all_stark, &input, &config, &mut timing)?;
timing.filter(Duration::from_millis(100)).print();

all_circuits.verify_root(root_proof.clone())?;

// Update public values for the aggregation.
let agg_public_values = PublicValues {
roots_before: first_public_values.roots_before,
roots_after: public_values.roots_after,
userdata: public_values.userdata,
};

// We can duplicate the proofs here because the state hasn't mutated.
let timing = TimingTree::new("prove aggregation", log::Level::Info);
let (agg_proof, updated_agg_public_values) = all_circuits.prove_aggregation(
false,
&root_proof_first,
false,
&root_proof,
agg_public_values,
)?;
timing.filter(Duration::from_millis(100)).print();
all_circuits.verify_aggregation(&agg_proof)?;

let timing = TimingTree::new("prove block", log::Level::Info);
let (block_proof, _block_public_values) =
all_circuits.prove_block(None, &agg_proof, updated_agg_public_values)?;
timing.filter(Duration::from_millis(100)).print();

log::info!(
"proof size: {:?}",
serde_json::to_string(&block_proof.proof).unwrap().len()
);
all_circuits.verify_block(&block_proof)
}

fn aggregate_proof_all() -> anyhow::Result<()> {
let basedir = env::var("BASEDIR").unwrap_or("/tmp/cannon".to_string());
let block = env::var("BLOCK_NO").unwrap_or("".to_string());
let file = env::var("BLOCK_FILE").unwrap_or(String::from(""));
let seg_dir = env::var("SEG_FILE_DIR").expect("segment file dir is missing");
let seg_file_number = env::var("SEG_FILE_NUM").expect("The segment file number is missing");
let seg_file_number = seg_file_number.parse::<_>().unwrap_or(2usize);
let seg_size = env::var("SEG_SIZE").unwrap_or(format!("{SEGMENT_STEPS}"));
let seg_size = seg_size.parse::<_>().unwrap_or(SEGMENT_STEPS);
prove_multi_seg_common(&seg_dir, &basedir, &block, &file, seg_size, seg_file_number)
}

0 comments on commit d961255

Please sign in to comment.