[WIP] Restructure the memory pipeline #118

winston-h-zhang · 2023-11-14T04:33:40Z

This PR restructures the way arecibo approaches memory allocation and restructures the entire memory pipeline. See the notion design doc for more info.

Notable improvements

The critical sections in prove_step and R1CSShape::commit_T no longer clone the large witness.
Downstream in lurk-rs, this PR fixes the strange regression we observed when using loaded public parameters. This is a strong indicator that inefficient memory allocations in arecibo was creating (and could've created) very unpredictable performance regressions.
Generally, the memory consumption of Nova is now very predictable, and this new memory pipeline is much safer to scale.

To-dos and other outstanding issues

We only refactor the Nova side of arecibo, and left the SuperNova untouched. This contains the scope of the PR. In the future, we should be interested in converting SuperNova to the same memory strategy as well.
The ResourceSink structure was temporarily created to manage the extra buffers prove_step needs. This should be integrated into the RecursiveSNARK API. Maybe with some sort of RecursiveSNARKEngine / RecursiveSNARKEngineTrait to manage folding.
There is one last inefficiency, which is that we recompute R1CS multiplication against the running witness in commit_T. The Z vectors should be moved into the ResourceSink to de-duplicate this.
Audit the security of these changes. @gabriel-barrett brought up possible security concerns when I re-introduced l_w_primary and l_u_primary into RecursiveSNARK, pointing out the revisiting Nova paper. I think ResourceSink fixes this, but we should make sure.
Refactor WitnessCS in upstream bellpepper to unify with WitnessViewCS, which is redundant.

adr1anh · 2023-11-14T08:55:11Z

Nice work getting rid of these clones!

Why does RecursiveSNARK::new need to return a sink, instead of keeping it as a member of the struct?

winston-h-zhang · 2023-11-14T19:07:24Z

@adr1anh I originally did that, but then I didn't want to pollute the RecursiveSNARK type, so I moved it out. There should be a nice way to manage everything together, like how CommitmentEngineTrait manages generating CommitmentKeys -- some sort of RecursiveSNARKEngineTrait.

winston-h-zhang · 2023-11-14T22:32:22Z

src/bellpepper/r1cs.rs

-    let W = R1CSWitness::<G>::new(shape, self.aux_assignment())?;
-    let X = &self.input_assignment()[1..];
+    let W = R1CSWitness::<G>::new(shape, self.aux_assignment().to_vec())?;

    let comm_W = W.commit(ck);

-    let instance = R1CSInstance::<G>::new(shape, &comm_W, X)?;
+    let instance = R1CSInstance::<G>::new(shape, comm_W, self.input_assignment().to_vec())?;


Since two people have asked me about why there are still clones here, I want to clarify. This function, r1cs_instance_and_witness, is no longer being called in prove_step. This change is purely atheistic -- it moves the .to_vec() call inside R1CSWitness::new/R1CSInstance::new to be explicit on construction -- so the compiler is happy.

We do not call this function, there are no extra clones.

It's a good change, and follows C-CALLER-CONTROL

winston-h-zhang · 2023-11-14T22:35:12Z

src/lib.rs

+    let (u_primary, w_primary) = r1cs::instance_and_witness(
+      r1cs_primary,
+      &pp.ck_primary,
+      input_assignment,
+      aux_assignment,
+    )?;


This new r1cs::instance_and_witness function eats the inputs instead of cloning.

huitseeker · 2023-11-14T19:43:31Z

src/bellpepper/solver.rs

+/// `setup = true`. After the initial step, every next Nova step has a fixed shape, so the buffers in
+/// `R1CSWitness` and `R1CSInstance` have the exact capacity they need. To be memory efficient,
+/// [`WitnessViewCS`] is flagged as `setup = false` and we no longer allow the buffers to resize.
+pub struct WitnessViewCS<'a, Scalar>


You should be able to use https://github.com/lurk-lab/bellpepper/blob/dev/crates/bellpepper/src/util_cs/witness_cs.rs now, if you want.

huitseeker · 2023-11-14T23:09:19Z

src/bellpepper/r1cs.rs

-    let W = R1CSWitness::<G>::new(shape, self.aux_assignment())?;
-    let X = &self.input_assignment()[1..];
+    let W = R1CSWitness::<G>::new(shape, self.aux_assignment().to_vec())?;

    let comm_W = W.commit(ck);

-    let instance = R1CSInstance::<G>::new(shape, &comm_W, X)?;
+    let instance = R1CSInstance::<G>::new(shape, comm_W, self.input_assignment().to_vec())?;


It's a good change, and follows C-CALLER-CONTROL

huitseeker · 2023-11-15T16:10:53Z

src/r1cs/sparse.rs

+    self.multiply_witness_unchecked(W, u_and_X)
+  }
+
+  /// Multiply by a witness representing a dense vector; uses rayon/gpu.


Not sure this uses the GPU yet.

This is a mis-comment; it indeed doesn't use GPU

huitseeker · 2023-11-15T16:15:55Z

src/r1cs/mod.rs

+    let mut W = vec![G::Scalar::ZERO; S.num_vars];
+    W.shrink_to_fit();
+    let mut E = vec![G::Scalar::ZERO; S.num_cons];
+    E.shrink_to_fit();


what's the capacity allocated by the vec! macro? how does that compare to the length of the vector created by the vec! macro?

Given the answers to these questions, what is the effect of the call to shrink_to_fit()?

I'm not sure, I couldn't comfirm from the vec! documentation that vec![x; n] initializes with_capacity(n). So I just redundantly called shrink_to_fit

I'm not sure, I couldn't comfirm from the vec! documentation that vec![x; n] initializes with_capacity(n)

Why? This is the part of the documentation where this is confirmed.

Alternatively, a quick test might have helped answer your question using the capcity method.

Ok, I just tested and indeed there's no need for shrink_to_fit

huitseeker · 2023-11-15T17:08:06Z

src/r1cs/mod.rs

+pub fn instance_and_witness<G: Group>(
+  shape: &R1CSShape<G>,
+  ck: &CommitmentKey<G>,
+  input_assignment: Vec<G::Scalar>,
+  aux_assignment: Vec<G::Scalar>,
+) -> Result<(R1CSInstance<G>, R1CSWitness<G>), NovaError> {
+  let W = R1CSWitness::<G>::new(shape, aux_assignment)?;
+  let comm_W = W.commit(ck);
+  let instance = R1CSInstance::<G>::new(shape, comm_W, input_assignment)?;
+
+  Ok((instance, W))
+}


The refactor of r1cs_instance_and_witness in #121 should supersede this.

#121)" This reverts commit fdf8296. This should allow us to investigate the effects of components of #118 without noise. We can re-apply this later if we see it as an improvement.

#121)" (#125) This reverts commit fdf8296. This should allow us to investigate the effects of components of #118 without noise. We can re-apply this later if we see it as an improvement.

winston-h-zhang · 2023-11-29T22:38:11Z

Closed after #137

winston-h-zhang changed the title ~~Restructure the memory pipeline~~ [WIP] Restructure the memory pipeline Nov 14, 2023

winston-h-zhang added 3 commits November 14, 2023 10:51

-- wip --

5cc36cf

rfc to u_and_X and one_and_X

d45ba90

bug fixes

e797c96

winston-h-zhang force-pushed the remove-r1cs-cloning branch from e7e7767 to ca12283 Compare November 14, 2023 18:57

rfc out ResourceSink

f701fab

winston-h-zhang force-pushed the remove-r1cs-cloning branch from ca12283 to f701fab Compare November 14, 2023 19:04

winston-h-zhang commented Nov 14, 2023

View reviewed changes

huitseeker mentioned this pull request Nov 15, 2023

refactor: Refactor R1CS for in-place operations and Vec inputs #121

Merged

huitseeker reviewed Nov 15, 2023

View reviewed changes

huitseeker mentioned this pull request Nov 15, 2023

Pre-allocate in SatisfyingAssignment<G> #124

Closed

huitseeker mentioned this pull request Nov 15, 2023

Revert "refactor: Refactor R1CS for in-place operations and Vec inputs (#121)" #125

Merged

This was referenced Nov 27, 2023

Optimizing the memory pipeline #137

Closed

Preallocate witness before synthesis #143

Merged

winston-h-zhang closed this Nov 29, 2023

huitseeker mentioned this pull request Dec 6, 2023

Streamline Nova allocations (Arecibo backport) microsoft/Nova#277

Closed

winston-h-zhang deleted the remove-r1cs-cloning branch December 6, 2023 21:41

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[WIP] Restructure the memory pipeline #118

[WIP] Restructure the memory pipeline #118

winston-h-zhang commented Nov 14, 2023 •

edited

Loading

adr1anh commented Nov 14, 2023

winston-h-zhang commented Nov 14, 2023

winston-h-zhang Nov 14, 2023

huitseeker Nov 14, 2023

winston-h-zhang Nov 14, 2023

huitseeker Nov 14, 2023

huitseeker Nov 14, 2023

huitseeker Nov 15, 2023

winston-h-zhang Nov 15, 2023

huitseeker Nov 15, 2023

winston-h-zhang Nov 15, 2023

huitseeker Nov 15, 2023 •

edited

Loading

winston-h-zhang Nov 15, 2023

huitseeker Nov 15, 2023

winston-h-zhang commented Nov 29, 2023

[WIP] Restructure the memory pipeline #118

[WIP] Restructure the memory pipeline #118

Conversation

winston-h-zhang commented Nov 14, 2023 • edited Loading

Notable improvements

To-dos and other outstanding issues

adr1anh commented Nov 14, 2023

winston-h-zhang commented Nov 14, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

huitseeker Nov 15, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

winston-h-zhang commented Nov 29, 2023

winston-h-zhang commented Nov 14, 2023 •

edited

Loading

huitseeker Nov 15, 2023 •

edited

Loading