fix(gpu_prover): improve memory allocation/deallocation ordering (#6)

## What ❔

Enforce allocation order for stage1&2 allocations.

## Why ❔

Explicit ordering helps to lower memory fragmentation and removes
allocation peak in stage 2.

## Is this a breaking change?
- [ ] Yes
- [x] No

## Checklist

- [x] PR title corresponds to the body of PR (we generate changelog
entries from PRs).
- [x] Code has been formatted.

Author: robik75

gpu_prover/src/prover/proof.rs

@@ -202,18 +202,51 @@ where
     setup.ensure_commitment_produced(context)?;
     setup_range.end(stream)?;
 
-    // stage 1
-    let stage_1_range = device_tracing::Range::new("stage_1")?;
-    stage_1_range.start(stream)?;
-    let mut stage_1_output = StageOneOutput::new(
+    let mut stage_1_output = StageOneOutput::allocate_trace_holders(
         circuit,
-        setup,
-        tracing_data_transfer,
         log_lde_factor,
         log_tree_cap_size,
+        context,
+    )?;
+    #[cfg(feature = "print_gpu_mem_usage")]
+    {
+        print!("after stage_1.allocate_trace_holders ");
+        context.print_mem_pool_stats()?;
+    }
+
+    let mut stage_2_output = StageTwoOutput::allocate_trace_evaluations(
+        circuit,
+        log_lde_factor,
+        log_tree_cap_size,
+        context,
+    )?;
+    #[cfg(feature = "print_gpu_mem_usage")]
+    {
+        print!("after stage_2.allocate_trace_evaluations ");
+        context.print_mem_pool_stats()?;
+    }
+
+    // witness_generation
+    let witness_generation_range = device_tracing::Range::new("witness_generation")?;
+    witness_generation_range.start(stream)?;
+    stage_1_output.generate_witness(
+        circuit,
+        setup,
+        tracing_data_transfer,
         circuit_sequence,
         context,
     )?;
+    witness_generation_range.end(stream)?;
+    #[cfg(feature = "print_gpu_mem_usage")]
+    {
+        print!("after generate_witness ");
+        context.print_mem_pool_stats()?;
+    }
+
+    // stage 1
+    let stage_1_range = device_tracing::Range::new("stage_1")?;
+    stage_1_range.start(stream)?;
+    stage_1_output.commit_witness(circuit, context)?;
     stage_1_range.end(stream)?;
     #[cfg(feature = "print_gpu_mem_usage")]
     {
@@ -222,9 +255,6 @@ where
     }
 
     setup.trace_holder.produce_tree_caps(context)?;
-    stage_1_output.witness_holder.produce_tree_caps(context)?;
-    stage_1_output.memory_holder.produce_tree_caps(context)?;
-    stage_1_output.produce_public_inputs(circuit, context)?;
 
     // seed
     let seed = initialize_seed::<C>(
@@ -241,18 +271,15 @@ where
     // stage 2
     let stage_2_range = device_tracing::Range::new("stage_2")?;
     stage_2_range.start(stream)?;
-    let stage_2_output = StageTwoOutput::new(
+    stage_2_output.generate(
         seed.clone(),
         circuit,
         &cached_data_values,
         setup,
         &mut stage_1_output,
-        log_lde_factor,
-        log_tree_cap_size,
         context,
     )?;
     stage_2_range.end(stream)?;
-
     #[cfg(feature = "print_gpu_mem_usage")]
     {
         print!("after stage_2 ");
@@ -395,8 +422,8 @@ where
     let is_finished_event = CudaEvent::create_with_flags(CudaEventCreateFlags::DISABLE_TIMING)?;
     is_finished_event.record(stream)?;
 
-    callbacks.extend(stage_1_output.callbacks);
-    callbacks.extend(stage_2_output.callbacks);
+    callbacks.extend(stage_1_output.callbacks.unwrap());
+    callbacks.extend(stage_2_output.callbacks.unwrap());
     callbacks.extend(stage_3_output.callbacks);
     callbacks.extend(stage_4_output.callbacks);
     callbacks.extend(stage_5_output.callbacks);
@@ -411,20 +438,20 @@ where
         memory_tree_caps: stage_1_output.memory_holder.get_tree_caps(),
         setup_tree_caps: setup.trace_holder.get_tree_caps(),
         stage_2_tree_caps: stage_2_output.trace_holder.get_tree_caps(),
-        stage_2_last_row: stage_2_output.last_row.clone(),
+        stage_2_last_row: stage_2_output.last_row.unwrap(),
         stage_2_offset_for_memory_grand_product_poly: stage_2_output.offset_for_grand_product_poly,
         stage_2_offset_for_delegation_argument_poly: stage_2_output
             .offset_for_sum_over_delegation_poly,
         quotient_tree_caps: stage_3_output.trace_holder.get_tree_caps(),
-        evaluations_at_random_points: stage_4_output.values_at_z.clone(),
+        evaluations_at_random_points: stage_4_output.values_at_z,
         deep_poly_caps: stage_4_output.trace_holder.get_tree_caps(),
         intermediate_fri_oracle_caps: stage_5_output
             .fri_oracles
-            .iter()
-            .map(|o| o.tree_caps.clone())
+            .into_iter()
+            .map(|o| o.tree_caps)
             .collect_vec(),
-        last_fri_step_plain_leaf_values: stage_5_output.last_fri_step_plain_leaf_values.clone(),
-        final_monomial_form: stage_5_output.final_monomials.clone(),
+        last_fri_step_plain_leaf_values: stage_5_output.last_fri_step_plain_leaf_values,
+        final_monomial_form: stage_5_output.final_monomials,
         pow_output,
         queries_output,
         circuit_sequence: circuit_sequence as u16,

gpu_prover/src/prover/stage_1.rs

@@ -27,26 +27,22 @@ pub(crate) struct StageOneOutput<'a, C: ProverContext> {
     pub witness_holder: TraceHolder<BF, C>,
     pub memory_holder: TraceHolder<BF, C>,
     pub generic_lookup_mapping: Option<C::Allocation<u32>>,
-    pub callbacks: Callbacks<'a>,
+    pub callbacks: Option<Callbacks<'a>>,
     pub public_inputs: Option<Arc<Mutex<Vec<BF>>>>,
 }
 
 impl<'a, C: ProverContext> StageOneOutput<'a, C> {
-    pub fn new(
+    pub fn allocate_trace_holders(
         circuit: &CompiledCircuitArtifact<BF>,
-        setup: &SetupPrecomputations<C>,
-        tracing_data_transfer: TracingDataTransfer<'a, C>,
         log_lde_factor: u32,
         log_tree_cap_size: u32,
-        circuit_sequence: usize,
         context: &C,
     ) -> CudaResult<Self> {
         let trace_len = circuit.trace_len;
         assert!(trace_len.is_power_of_two());
         let log_domain_size = trace_len.trailing_zeros();
-        let witness_subtree = &circuit.witness_layout;
-        let witness_columns_count = witness_subtree.total_width;
-        let mut witness_holder = TraceHolder::new(
+        let witness_columns_count = circuit.witness_layout.total_width;
+        let witness_holder = TraceHolder::new(
             log_domain_size,
             log_lde_factor,
             0,
@@ -55,9 +51,8 @@ impl<'a, C: ProverContext> StageOneOutput<'a, C> {
             true,
             context,
         )?;
-        let memory_subtree = &circuit.memory_layout;
-        let memory_columns_count = memory_subtree.total_width;
-        let mut memory_holder = TraceHolder::new(
+        let memory_columns_count = circuit.memory_layout.total_width;
+        let memory_holder = TraceHolder::new(
             log_domain_size,
             log_lde_factor,
             0,
@@ -66,6 +61,28 @@ impl<'a, C: ProverContext> StageOneOutput<'a, C> {
             true,
             context,
         )?;
+        Ok(Self {
+            witness_holder,
+            memory_holder,
+            generic_lookup_mapping: None,
+            callbacks: None,
+            public_inputs: None,
+        })
+    }
+
+    pub fn generate_witness(
+        &mut self,
+        circuit: &CompiledCircuitArtifact<BF>,
+        setup: &SetupPrecomputations<C>,
+        tracing_data_transfer: TracingDataTransfer<'a, C>,
+        circuit_sequence: usize,
+        context: &C,
+    ) -> CudaResult<()> {
+        let trace_len = circuit.trace_len;
+        assert!(trace_len.is_power_of_two());
+        let log_domain_size = trace_len.trailing_zeros();
+        let witness_subtree = &circuit.witness_layout;
+        let memory_subtree = &circuit.memory_layout;
         let generic_lookup_mapping_size = witness_subtree.width_3_lookups.len() << log_domain_size;
         let mut generic_lookup_mapping = context.alloc(generic_lookup_mapping_size)?;
         let TracingDataTransfer {
@@ -75,6 +92,7 @@ impl<'a, C: ProverContext> StageOneOutput<'a, C> {
             transfer,
         } = tracing_data_transfer;
         transfer.ensure_transferred(context)?;
+        self.callbacks = Some(transfer.callbacks);
         let stream = context.get_exec_stream();
         assert_eq!(COMMON_TABLE_WIDTH, 3);
         assert_eq!(NUM_COLUMNS_FOR_COMMON_TABLE_WIDTH_SETUP, 4);
@@ -100,6 +118,8 @@ impl<'a, C: ProverContext> StageOneOutput<'a, C> {
                 + timestamp_range_check_multiplicities_columns.num_elements,
             generic_multiplicities_columns.start
         );
+        let witness_holder = &mut self.witness_holder;
+        let memory_holder = &mut self.memory_holder;
         match data_device {
             TracingDataDevice::Main {
                 setup_and_teardown,
@@ -172,15 +192,26 @@ impl<'a, C: ProverContext> StageOneOutput<'a, C> {
             trace_len,
             context,
         )?;
-        memory_holder.make_evaluations_sum_to_zero_extend_and_commit(context)?;
-        witness_holder.make_evaluations_sum_to_zero_extend_and_commit(context)?;
-        Ok(Self {
-            witness_holder,
-            memory_holder,
-            generic_lookup_mapping: Some(generic_lookup_mapping),
-            callbacks: transfer.callbacks,
-            public_inputs: None,
-        })
+        self.generic_lookup_mapping = Some(generic_lookup_mapping);
+        Ok(())
+    }
+
+    pub fn commit_witness(
+        &mut self,
+        circuit: &'a CompiledCircuitArtifact<BF>,
+        context: &C,
+    ) -> CudaResult<()>
+    where
+        C::HostAllocator: 'a,
+    {
+        self.memory_holder
+            .make_evaluations_sum_to_zero_extend_and_commit(context)?;
+        self.memory_holder.produce_tree_caps(context)?;
+        self.witness_holder
+            .make_evaluations_sum_to_zero_extend_and_commit(context)?;
+        self.witness_holder.produce_tree_caps(context)?;
+        self.produce_public_inputs(circuit, context)?;
+        Ok(())
     }
 
     pub fn produce_public_inputs(
@@ -257,7 +288,10 @@ impl<'a, C: ProverContext> StageOneOutput<'a, C> {
             guard.extend(first_row_public_inputs);
             guard.extend(one_before_last_row_public_inputs);
         };
-        self.callbacks.schedule(function, stream)?;
+        self.callbacks
+            .as_mut()
+            .unwrap()
+            .schedule(function, stream)?;
         self.public_inputs = Some(public_inputs);
         Ok(())
     }