fix: Reduce memory usage in mem2reg

[aakoshh]

Description

Problem*

Resolves #7001

Summary*

Changes mem2reg::block::Block to only use the mutable APIs of OrdMap if unifying the alias sets would cause something to change. That should save some memory, because doing self.aliases.entry(...).and_modify(...).or_insert_with(...) causes OrdMap::get_mut to be called, which calls Arc::make_mut to prepare the entry for modification, which involves making a clone and a new Arc if the Arc is shared, which it will be because we cloned a Block to get the folding started.

This resulted in the memory required for mem2reg to go down to ~1GB, with later peaks corresponding to convert_ssa_instruction during ACIR generation:

With #7057 merged the results are even better:

TODO in followup PR:

• See if blocks and aliases can be freed up as we go
• Pre-process functions before inlining them (Tom suggests: inline, unroll and mem2reg)

Additional Context

Profiling on Mac

I couldn't compile the heaptrack_gui. Instead tried cargo-instruments which requires Xcode to be installed, but it stopped working after a while, and I ended up using Instruments directly (a profiler tool on Mac, part of Xcode).

Executing the following command:

./target/release/nargo --program-dir ../aztec-packages/noir-projects/noir-protocol-circuits/crates/rollup-base-public compile  --silence-warnings --skip-underconstrained-check --skip-brillig-constraints-check --force

Compiled with:

cargo build -p nargo_cli --release

The following is needed for Instruments to be able to start the executable:

codesign -s - -v -f --entitlements ../debug.plist target/release/nargo

where debug.plist has the following content:

<?xml version="1.0" encoding="UTF-8"?><!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "https://www.apple.com/DTDs/PropertyList-1.0.dtd">
<plist version="1.0"><dict><key>com.apple.security.get-task-allow</key><true/></dict></plist>

Memory profile

Original memory profile over time:

EXPERIMENT: Changed AliasSet to use OrdSet instead of BTreeSet to utilise structural sharing.

When I tried this on aztec-packages it was promising but perhaps I didn't run it long enough I accidentally got the 2nd mem2reg commented out. Profiling again locally on noir revealed this version uses 21GB of memory 😖 :

The increase seems to come from the allocation of a large number of Arcs:

After having a second look at the size of the AliasSets, they seem to actually only contain a single item, so I was barking up the wrong tree.

Other tweaks I tried:

• Replace Expression in aliases: OrdMap<Expression, AliasSet> with u64 calculated by fxhash::hash64, just to see the impact of not storing clones of Expressions; the result is the exact same memory usage, because as it turns out the memory allocated by OrdMap::get_mut call is still for Arcs, it has always been.
• Replace aliases: OrdMap<Expression, AliasSet> with std::HashMap or std::BTreeMap; the memory usage in this case went from 5GB to 11GB; this is because of the increased cost of cloning these maps
• Replace im with im_rc to see if Arc vs Rc makes any difference, although it's probably just faster than smaller. Can't because of rayon.

Documentation*

Check one:

• No documentation needed.
• Documentation included in this PR.
• [For Experimental Features] Documentation to be submitted in a separate PR.

PR Checklist*

• I have tested the changes locally.
• I have formatted the changes with Prettier and/or cargo fmt on default settings.

[github-actions]

Compilation Report

Program	Compilation Time	%
sha256_regression	1.050s	-17%
regression_4709	0.833s	5%
ram_blowup_regression	16.200s	-5%
rollup-root	3.504s	5%
rollup-merge	1.844s	0%
rollup-block-merge	3.324s	0%
rollup-base-public	31.240s	-14%
rollup-base-private	11.740s	-7%
private-kernel-tail	0.940s	-4%
private-kernel-reset	6.086s	-3%
private-kernel-inner	2.148s	7%

[github-actions]

Execution Report

Program	Execution Time	%
sha256_regression	0.052s	0%
regression_4709	0.001s	0%
ram_blowup_regression	0.598s	-2%
rollup-root	0.104s	0%
rollup-merge	0.006s	0%
rollup-block-merge	0.104s	0%
rollup-base-public	1.222s	0%
rollup-base-private	0.462s	2%
private-kernel-tail	0.019s	0%
private-kernel-reset	0.310s	-2%
private-kernel-inner	0.069s	1%

[github-actions]

Execution Memory Report

Program	Peak Memory
keccak256	74.67M
workspace	123.84M
regression_4709	316.00M
ram_blowup_regression	512.61M
rollup-base-public	478.87M
rollup-base-private	325.48M
private-kernel-tail	180.47M
private-kernel-reset	245.16M
private-kernel-inner	208.55M

[github-actions]

Compilation Memory Report

Program	Peak Memory	%
keccak256	77.72M	0%
workspace	123.71M	0%
regression_4709	424.09M	0%
ram_blowup_regression	1.58G	0%
rollup-base-public	2.51G	-34%
rollup-base-private	901.10M	-9%
private-kernel-tail	207.18M	0%
private-kernel-reset	669.26M	0%
private-kernel-inner	294.40M	0%

[TomAFrench]

Nice!

[aakoshh]

Now merged #7058 the results are very similar 👍

[TomAFrench]

Sorry, didn't see the usage of AbstractVecSet when fixing those conflicts.