[BUG] Metal proof generation seems to get stuck for quite a while in `eval_check`

[weikengchen]

Bug Report

I used Apple M2 Ultra and ran the proof generation for examples/ecdsa.

I can do the proof generation in CPU without issues. But, when it comes to Metal, although I do get significant performance improvement for the Merkle tree commitments (even though it is Poseidon), the evaluation check took me 307s. This step would take my CPU about 20s to finish.

The evaluation check is this step in zkp/src/prove/prover.rs

        let groups: Vec<&_> = self
            .groups
            .iter()
            .map(|pg| &pg.as_ref().unwrap().evaluated)
            .collect();
        circuit_hal.eval_check(
            &check_poly,
            groups.as_slice(),
            globals,
            poly_mix,
            self.po2,
            self.cycles,
        );

where the Metal would be computing the quotient polynomials.

Steps to Reproduce

cargo run --features=metal on example/ecdsa

Expected behavior

Since CPU took 20s to do eval check, Metal shouldn't take this long.

Your Environment

• risc0-zkvm version: v0.21.0-alpha.1, aka main
• Rust version: 1.73
• Platform/OS: MacOS Sonoma Public Beta 14.3

[weikengchen]

Here is the GPU history from Activity Monitor.

The eval check was the part in the middle that is sort of not moving.

Here is my breakdown in Metal:

segment 1 (of size 1048576)

• execute the VM takes 5.11s
• commit CODE takes 0.40s
• commit DATA takes 3.44s
• commit ACCUM takes 0.75s
• eval check takes 307s
• commit CHECK takes 0.41s

segment 2 (smaller, of size likely 262144)

• execute the VM takes 1.05s
• commit CODE takes 0.20s
• commit DATA takes 0.82s
• commit ACCUM takes 0.19s
• eval check takes 0.40s
• commit CHECK takes 0.09s

This can be compared with CPU:

segment 1

• execute the VM takes 11.2s
• commit CODE takes 3.52s
• commit DATA takes 28.6s
• commit ACCUM takes 6.79s
• eval check takes 19.59s
• commit CHECK takes 3.43s

segment 2

• execute the VM takes 1.22s
• commit CODE takes 0.87s
• commit DATA takes 6.51s
• commit ACCUM takes 1.71s
• eval check takes 5.17s
• commit CHECK takes 0.86s

[weikengchen]

I am actually going to close this issue since I can no longer reproduce it after I close my Chrome browser with 30+ tabs. Now, the GPU history caused by a proof generation is just this:

Highly suspect that this is just memory being used up and eval check has to invoke virtual memory.

[flaub]

I recently landed a change to some of the circuit generated files which will cause Metal and other compilation caches to become invalid. After a change of this nature, it's expected that the first run on Metal/CUDA will take sometimes a long time to JIT. This is likely why you can't reproduce the slow down. If you catch it again, you can confirm on a mac by looking for a MTLCompilerService process running in the background.

[weikengchen]

Can I ask a question about the estimated performance breakdown from the RISC Zero team?

I am getting 5 : 5 : 2 on my M2 Ultra.

• The first 5---aka witness generation---is the step that populates the code column with the code, executes each instruction one by one in order to fill in the data column, and then fills in more numbers that we will omit there.
• The second 5---aka polynomial commitment---includes the overhead to commit the multiple trees. My local experiment shows that one can make it lower to by switching the entire hash function to SHA-256, so that it becomes 5 : 3 : 2, but this is not really helpful.
• The last 2---aka quotient polynomial---is to run eval_check for each point (or roughly "each cycle").

Is it an open problem to make the witness generation faster? Can a two-pass solution---a sequential pass and a parallel pass shine that leverages multi cores?

[flaub]

Yes, we are planning to make the witness generation parallelized in the near future by having a preflight step to record all the 'back' information to make parallel execution possible.

[weikengchen]

Got it. I can see that the recursion circuit already has preflight.