Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: batch sized recursion #785

Merged
merged 4 commits into from
May 23, 2024
Merged

feat: batch sized recursion #785

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
6ee15b1
fix: update release workflow (#777)
ratankaliani May 21, 2024
5a2cf9f
feat: release on main (#779)
ratankaliani May 21, 2024
9425a9a
test
jtguibas May 21, 2024
03db43d
Merge branch 'dev' into john/batch-size-compress
jtguibas May 23, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 11 additions & 0 deletions .github/workflows/release.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@ name: release

on:
push:
branches:
- main
tags:
- "v*.*.*"
workflow_dispatch:
Expand Down Expand Up @@ -119,6 +121,15 @@ jobs:
with:
key: ${{ matrix.target }}
cache-on-failure: true

- name: Install go
uses: actions/setup-go@v5
with:
go-version: '^1.22.1'

- name: Check go installation
run: |
go version

- name: Set up git private repo access
run: |
Expand Down
159 changes: 88 additions & 71 deletions prover/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -390,7 +390,6 @@ impl SP1Prover {
// Run the recursion and reduce programs.

// Run the recursion programs.
let mut records = Vec::new();

let (core_inputs, deferred_inputs) = self.get_first_layer_inputs(
vk,
Expand All @@ -400,62 +399,71 @@ impl SP1Prover {
batch_size,
);

for input in core_inputs {
let mut runtime = RecursionRuntime::<Val<InnerSC>, Challenge<InnerSC>, _>::new(
&self.recursion_program,
self.compress_machine.config().perm.clone(),
);

let mut witness_stream = Vec::new();
witness_stream.extend(input.write());

runtime.witness_stream = witness_stream.into();
runtime.run();
runtime.print_stats();
let mut first_layer_proofs = Vec::new();
let shard_batch_size = sp1_core::utils::env::shard_batch_size() as usize;
for inputs in core_inputs.chunks(shard_batch_size) {
let proofs = inputs
.into_par_iter()
.map(|input| {
let mut runtime = RecursionRuntime::<Val<InnerSC>, Challenge<InnerSC>, _>::new(
&self.recursion_program,
self.compress_machine.config().perm.clone(),
);

records.push((runtime.record, ReduceProgramType::Core));
let mut witness_stream = Vec::new();
witness_stream.extend(input.write());

runtime.witness_stream = witness_stream.into();
runtime.run();
runtime.print_stats();

let pk = &self.rec_pk;
let mut recursive_challenger = self.compress_machine.config().challenger();
(
self.compress_machine.prove::<LocalProver<_, _>>(
pk,
runtime.record,
&mut recursive_challenger,
),
ReduceProgramType::Core,
)
})
.collect::<Vec<_>>();
first_layer_proofs.extend(proofs);
}

// Run the deferred proofs programs.
for input in deferred_inputs {
let mut runtime = RecursionRuntime::<Val<InnerSC>, Challenge<InnerSC>, _>::new(
&self.deferred_program,
self.compress_machine.config().perm.clone(),
);

let mut witness_stream = Vec::new();
witness_stream.extend(input.write());

runtime.witness_stream = witness_stream.into();
runtime.run();
runtime.print_stats();
for inputs in deferred_inputs.chunks(shard_batch_size) {
let proofs = inputs
.into_par_iter()
.map(|input| {
let mut runtime = RecursionRuntime::<Val<InnerSC>, Challenge<InnerSC>, _>::new(
&self.deferred_program,
self.compress_machine.config().perm.clone(),
);

records.push((runtime.record, ReduceProgramType::Deferred));
let mut witness_stream = Vec::new();
witness_stream.extend(input.write());

runtime.witness_stream = witness_stream.into();
runtime.run();
runtime.print_stats();

let pk = &self.deferred_pk;
let mut recursive_challenger = self.compress_machine.config().challenger();
(
self.compress_machine.prove::<LocalProver<_, _>>(
pk,
runtime.record,
&mut recursive_challenger,
),
ReduceProgramType::Deferred,
)
})
.collect::<Vec<_>>();
first_layer_proofs.extend(proofs);
}

// Prove all recursion programs and recursion deferred programs and verify the proofs.

// Make the recursive proofs for core and deferred proofs.
let first_layer_proofs = records
.into_par_iter()
.map(|(record, kind)| {
let pk = match kind {
ReduceProgramType::Core => &self.rec_pk,
ReduceProgramType::Deferred => &self.deferred_pk,
ReduceProgramType::Reduce => unreachable!(),
};
let mut recursive_challenger = self.compress_machine.config().challenger();
(
self.compress_machine.prove::<LocalProver<_, _>>(
pk,
record,
&mut recursive_challenger,
),
kind,
)
})
.collect::<Vec<_>>();

// Chain all the individual shard proofs.
let mut reduce_proofs = first_layer_proofs
.into_iter()
Expand All @@ -467,28 +475,37 @@ impl SP1Prover {
loop {
tracing::debug!("Recursive proof layer size: {}", reduce_proofs.len());
is_complete = reduce_proofs.len() <= batch_size;
reduce_proofs = reduce_proofs
.par_chunks(batch_size)
.map(|batch| {
let (shard_proofs, kinds) =
batch.iter().cloned().unzip::<_, _, Vec<_>, Vec<_>>();

let input = SP1ReduceMemoryLayout {
compress_vk: &self.compress_vk,
recursive_machine: &self.compress_machine,
shard_proofs,
kinds,
is_complete,
};

let proof = self.compress_machine_proof(
input,
&self.compress_program,
&self.compress_pk,
);
(proof, ReduceProgramType::Reduce)

let compress_inputs = reduce_proofs.chunks(batch_size).collect::<Vec<_>>();
let batched_compress_inputs =
compress_inputs.chunks(shard_batch_size).collect::<Vec<_>>();
reduce_proofs = batched_compress_inputs
.into_iter()
.flat_map(|batches| {
batches
.par_iter()
.map(|batch| {
let (shard_proofs, kinds) =
batch.iter().cloned().unzip::<_, _, Vec<_>, Vec<_>>();

let input = SP1ReduceMemoryLayout {
compress_vk: &self.compress_vk,
recursive_machine: &self.compress_machine,
shard_proofs,
kinds,
is_complete,
};

let proof = self.compress_machine_proof(
input,
&self.compress_program,
&self.compress_pk,
);
(proof, ReduceProgramType::Reduce)
})
.collect::<Vec<_>>()
})
.collect();
.collect::<Vec<_>>();

if reduce_proofs.len() == 1 {
break;
Expand Down
7 changes: 2 additions & 5 deletions sp1up/sp1up
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -80,13 +80,10 @@ main() {

# Install by downloading binaries
if [[ "$SP1UP_REPO" == "succinctlabs/sp1" && -z "$SP1UP_BRANCH" && -z "$SP1UP_COMMIT" ]]; then
SP1UP_VERSION=${SP1UP_VERSION:-nightly}
SP1UP_VERSION=${SP1UP_VERSION:-main}
SP1UP_TAG=$SP1UP_VERSION

# Normalize versions (handle channels, versions without v prefix
if [[ "$SP1UP_VERSION" =~ ^nightly ]]; then
SP1UP_VERSION="nightly"
elif [[ "$SP1UP_VERSION" == [[:digit:]]* ]]; then
if [[ "$SP1UP_VERSION" == [[:digit:]]* ]]; then
# Add v prefix
SP1UP_VERSION="v${SP1UP_VERSION}"
SP1UP_TAG="${SP1UP_VERSION}"
Expand Down
Loading