feat: optimize segment proof aggregation

[atanmarko]

Resolves #399
Resolves #403

[atanmarko]

PR limits the number of the segment proofs that are proven in parallel to limit the memory consumption with the small penalty of async engine utilization. Chunk size could be tweaked if needed.

[Nashtare]

I'm afraid this doesn't quite solve the problem.
The current "issue" causing large memory overhead with continuations, is that we want to process a large batch of txns at once, to reduce redundant work in between (re-calculations of intermediate trie roots).
Note that I admit the names here not being changed are confusing. When we iterate over txs, we don't have individual txn that are passed as inputs of the SegmentDataIterator but a Vec of txns, of length possibly greater than 1 given the above.

The most CPU cycle optimized approach to reduce the mentioned redundant computations is to have a single payload representing all txns of a block (by payload here, I mean the GenerationInputs fed to the prover).

As such, the below loop can be simplified at a high-level as:

            let tx_proof_futs: FuturesUnordered<_> = {
                let data_iterator = SegmentDataIterator {
                    partial_next_data: None,
                    inputs: single_payload, // <-- all txns are contained here!
                    max_cpu_len_log: Some(max_cpu_len_log),
                };

                Directive::map(IndexedStream::from(data_iterator), &seg_ops)
                    .fold(&agg_ops)
                    .run(runtime)
                    .map(move |e| {
                        e.map(|p| (idx, proof_gen::proof_types::TxnAggregatableProof::from(p)))
                    })
            };

As such, the problematic iterator here is data_iterator which is running unboundedly, generating possibly hundreds of segments that stay in memory until being consumed (because generating them is orders of magnitude faster than proving them).

[atanmarko]

At the moment, generate_next_segment is executed on GENERATION_INPUTS_AGG_CHUNK_SIZE_INNER GenerationInputs. 4 such jobs (so 16 GenerationInputs segments with SegmentProof proving) would run in parallel, and then be aggregated to chunk transaction proof.
Eventually all the chunk resulting transaction proofs would be aggregated to final transaction proof. Chunk sizes are there to limit memory consumption (we will keep in memory input GenerationSegmentData for GENERATION_INPUTS_AGG_CHUNK_SIZE GenerationInputs proving jobs). If that is not what we want, what is the alternative?

Original implementation was keeping all the GenerationSegmentData and executed first all the SegmentProof jobs before it switch to the proof aggregation (e.g.if block has 200 GenerationInpus then one data_iterator and SegmentProof for all segments in the block running in parallel), because paladin does not switch to other operation in chained operation setup before it finishes the previous one.

[Nashtare]

if block has 200 GenerationInpus then one data_iterator and SegmentProof for all segments in the block running in parallel

Put it differently, think that we now have a single GenerationInputs representing 200 txns. The issue is the internal SegmentDataIterator. There is only of them, but it is responsible for generating all segments, which cannot be processed as fast as they are being generated.

We still have the outer loop iterating over txns because we want the flexibility here. This is configured by the -b batch size argument when calling the proving mode. But theoretically, the best is when the batch covers a whole block at once, effectively removing the outer loop.

[atanmarko]

Put it differently, think that we now have a single GenerationInputs representing 200 txns. The issue is the internal SegmentDataIterator. There is only of them, but it is responsible for generating all segments, which cannot be processed as fast as they are being generated.

We still have the outer loop iterating over txns because we want the flexibility here. This is configured by the -b batch size argument when calling the proving mode. But theoretically, the best is when the batch covers a whole block at once, effectively removing the outer loop.

Ok, I had wrong understanding of the problem. I'll look how to configurable limit data_iterator

[atanmarko]

Returned it to draft, CI tests seem pretty slow (maybe there is some other reason), but I need to test this PR more with the real prover tomorrow.

[atanmarko]

PR is tested. Currently we have fixed parameters for cpu len and batch size of 10, where premature segment chunk aggregation does not make a big improvement in memory used, but takes significant CPU performance hit. We give up this approach.
Alternative that would be better but would probably take much more effort (if possible) would be to change Paladin to be able to run chained operations (map/fold) without waiting for the first operation to finish from all the workers.