Precompile Backend #562

Roee-87 · 2025-01-16T07:50:32Z

The purpose of this PR is to add code to support precompiles in the zkVM.

…re method

Roee-87 · 2025-01-16T18:35:03Z

common/src/rv_trace.rs

+        if self.precompile_outputs.len() <= internal_address {
+            self.precompile_outputs.resize(internal_address + 1, 0);
+        }
+
        self.outputs[internal_address] = value;


The store method only takes one value as an input. Should there be additional Option inputs for precompile input and output?

Roee-87 · 2025-01-16T18:48:31Z

jolt-core/src/jolt/instruction/ecall.rs

+
+        let mut virtual_trace = vec![];
+
+        let precompile_output: &[u32; 16] = &[0u32; 16];  // compute precompile output based on t0 register


Is the precompile output going to be stored in rs1?

moodlezoup

Nice work! I think we're almost there

moodlezoup · 2025-02-13T18:56:32Z

jolt-core/src/jolt/precompiles/mod.rs

+    pub input:[u32; 16] // 512 bits
+    pub output:[u32; 16] // 512 bits
+}
+struct PrecompileWitness {


let's instead make PrecompileWitness an associated type on the Precompile trait

I added PrecompileWitness and PrecompileProof as associated types with an example implementation (for PrecompileWitness) in the Bn254.rs file: 37da785

Should PrecompileIO also be an associated type? Or will it be identical for all precompiles?

moodlezoup · 2025-02-14T16:28:57Z

tracer/src/emulator/mmu.rs

+    /// Gets the precompile input from the correct location in memory designated for precompile inputs.
+    pub fn get_precompile_input_word(&mut self, v_address: u64) -> Result<u32, Trap> {
+        // Load the world from virtual memory.
+        let input_word = self.load_word(v_address);
+        input_word
+    }
+
+    /// Sets the precompile output to the correct location in memory designated for precompile outputs.
+    pub fn set_precompile_output_word(&mut self, output_word: u32, p_address: u64) {
+        // Store the word in the correct precompile output memory location
+        self.store_word(p_address, output_word);
+    }


I think what we should do here is use self.jolt_device.memory_layout (which contains the start/end addresses of the precompile input/output regions) to determine where to read/write the precompile inputs/outputs.

That way, we don't need to pass in v_address or p_address to these methods; we'll always read/write 512 bits (16 words) from the respective memory regions.

fixed using fields in memory_layout: 0782e6e

moodlezoup · 2025-02-14T16:29:29Z

tracer/src/emulator/cpu.rs

+            let precompile_input_address = 0_u64; // @TODO: which address should be supplied to fetch the input?            
+            let inputs = cpu.mmu.get_precompile_input(precompile_input_address); // which address should be supplied to fetch the input?
+            let output = precompile.execute(input);
+            let precompile_output_address = 0_u64; // @TODO: which address should be supplied to store the output?
+            cpu.get_mut_mmu.set_precompile_output(output, precompile_output_address);


re: which addresses to supply, see comment below

Set using memory_layout struct: 0782e6e

moodlezoup · 2025-02-14T16:44:57Z

jolt-core/src/jolt/instruction/ecall.rs

+
+        let mut virtual_trace = vec![];
+
+        let precompile_output: &[u32; 16] = &[0u32; 16];  // compute precompile output based on t0 register


No, we'll just be writing the precompile output to memory

Separately, we'll need to invoke the precompile (e.g. bn254_add) here on the input.

Actually I guess we'll need to add some fields on RVTraceRow to capture the inputs during tracing. Namely, we should add an Option<[u32; 16]> for precompile inputs
And we'll probably need an Option<u64> for the precompile output address

moodlezoup · 2025-02-14T16:50:36Z

jolt-core/src/jolt/precompiles/mod.rs

+    pub fn prove(&self, witness: PrecompileWitness) -> JoltProof;  // Need to include Preprocessing struct as an input?
+    pub fn verify(&self, proof: JoltProof) -> Result<(), ProofVerifyError>;


instead of JoltProof, I think we'll need another associated type PrecompileProof

moodlezoup · 2025-02-14T16:53:32Z

jolt-sdk/macros/src/lib.rs

I don't think we need to make any changes to this file for the time being; the main prove function can just invoke the individual precompile provers

reverted changes: 5407655

moodlezoup · 2025-02-14T17:37:27Z

jolt-core/src/jolt/instruction/ecall.rs

+        // Ecall source registers
+        let r_x = trace_row.instruction.rs1;
+        let r_y = trace_row.instruction.rs2;


There should only be on source register (t0), specifying which precompile to execute

moodlezoup · 2025-02-14T17:49:38Z

jolt-core/src/jolt/instruction/ecall.rs

+        // Virtual registers used in sequence
+        let v_0 = Some(virtual_register_index(0));
+        let v_ao0: Option<u64> = Some(virtual_register_index(1));
+        let v_ao1: Option<u64> = Some(virtual_register_index(2));
+        let v_ao2: Option<u64> = Some(virtual_register_index(3));
+        let v_ao3: Option<u64> = Some(virtual_register_index(4));
+        let v_ao4: Option<u64> = Some(virtual_register_index(5));
+        let v_ao5: Option<u64> = Some(virtual_register_index(6));
+        let v_ao6: Option<u64> = Some(virtual_register_index(7));
+        let v_ao7: Option<u64> = Some(virtual_register_index(8));
+        let v_ao8: Option<u64> = Some(virtual_register_index(9));
+        let v_ao9: Option<u64> = Some(virtual_register_index(10));
+        let v_ao10: Option<u64> = Some(virtual_register_index(11));
+        let v_ao11: Option<u64> = Some(virtual_register_index(12));
+        let v_ao12: Option<u64> = Some(virtual_register_index(13));
+        let v_ao13: Option<u64> = Some(virtual_register_index(14));
+        let v_ao14: Option<u64> = Some(virtual_register_index(15));
+        let v_ao15: Option<u64> = Some(virtual_register_index(16));


we shouldn't need any virtual registers, as we'll be writing the outputs to memory, not registers

Removed and updated the virtual instruction to write the relevant precompile output word to memory: 8864710

moodlezoup · 2025-02-14T17:51:02Z

jolt-core/src/jolt/instruction/ecall.rs

+        virtual_trace.push(RVTraceRow {
+            instruction: ELFInstruction {
+                address: trace_row.instruction.address,
+                opcode: RV32IM::VIRTUAL_ADVICE,


I think I was wrong before; we will in fact need a new RV32IM::VIRTUAL_ADVICE_SW instruction that writes a word to memory, as opposed to VIRTUAL_ADVICE which writes to a register. This also entails some changes in rv_trace.rs for the new instruction which we can talk about

…ry_laout struct

…e-backend pulled in latest main

Roee-87 added 19 commits January 13, 2025 23:25

added precompile code for common directory

9d21aa7

added comments for jolt-core/jolt/instruction/ecall.rs

e7e56aa

added prcompile directory to jotl-core/jolt/

c2d1abb

added precompile template code to jolt-sdk

337af2a

added skeleton code for ecall.rs

379541c

removed extra two virtual advice registers

8b15c33

added prove and verify placeholder functions in jolt-sdk/macros/lib.rs

ab4d42d

added skeleton code to tracer/src/emulator/cpu.rs

9d0231a

added set_precompile_output_word() to mmu.rs

053d7f3

updated mmu precompile methods

da40ee3

updated cpu.rs

3f38bd5

added ok(()) to ecall method

f65d5f3

updated precompile mod.rs for jolt-core

e85ddbe

added 16 advice instruction to ecall.rs

691efe6

added 16 advice instructions to ecall.rs

b70ce3c

added conditional for precompile output to load method

8f06579

fixed typo

e813eac

unsure of how to handle store() method in rv_trace.rs

b3b57ae

refactored rv_trace.rs to include conditionals for precompiles in sto…

9499b18

…re method

Roee-87 commented Jan 16, 2025

View reviewed changes

added correct address for panic mem region

c972066

Roee-87 commented Jan 16, 2025

View reviewed changes

Roee-87 and others added 6 commits February 5, 2025 23:56

updated jolt-sdk

fd62abf

added comment

952d96d

updated instructions for ecall and precompile

3600f9c

added precompile to list of RISCV32 enum

0c5466a

updated rv_trace.rs to include code for setting precompile flag

4748f73

Merge branch 'a16z:main' into rr-jolt-precompile-backend

9af1042

Roee-87 marked this pull request as ready for review February 6, 2025 08:53

fixed typo for precompile_input(s)/output(s)

d562e2c

Roee-87 added 2 commits February 6, 2025 11:54

fixed ordering for RV32IM instructions in rv_tracer.rs

5334965

added WORD_SIZE type to advice instruction for ecall.rs

e264e85

moodlezoup requested changes Feb 14, 2025

View reviewed changes

Roee-87 added 19 commits February 15, 2025 15:10

added notes

d00f4da

updated to latest main

1439ad1

reverted changes to jolt-sdk

5407655

modified associated types for precompile trait

37da785

fixed formatting

7c4ceaf

set precompile input and output memory address using jolt_device.memo…

0782e6e

…ry_laout struct

made precompile directory a moddule in jolt/common/

a8ef501

modified precompile input/output methods for cpu.rs and mmu.rs

6b6f96e

code cleanup

e26b960

fixed typos

86d90c7

removed try_into() from bn254_add inputs

93503d4

debug

060e239

hopefully last clippy complaint...

344a9ca

fixed clippy errors for mmu.rs and cpu.rs

159b7e1

updated ecall.rs

450708c

added sequence_output placeholder function to ecall.rs

2e7c0f2

Merge branch 'main' of github.com:Roee-87/jolt into rr-jolt-precompil…

5ad76d8

…e-backend pulled in latest main

updated ecall.rs RVTraceRow fields

c528616

modified ecall.rs to write precompile output to memory

8864710

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Precompile Backend #562

Precompile Backend #562

Roee-87 commented Jan 16, 2025

Roee-87 Jan 16, 2025

Roee-87 Jan 16, 2025

moodlezoup left a comment

moodlezoup Feb 13, 2025

Roee-87 Feb 26, 2025

moodlezoup Feb 14, 2025

Roee-87 Feb 26, 2025

moodlezoup Feb 14, 2025

Roee-87 Feb 26, 2025

moodlezoup Feb 14, 2025

moodlezoup Feb 14, 2025

moodlezoup Feb 14, 2025

moodlezoup Feb 14, 2025

Roee-87 Feb 26, 2025

moodlezoup Feb 14, 2025

moodlezoup Feb 14, 2025

Roee-87 Mar 6, 2025

moodlezoup Feb 14, 2025


		let mut virtual_trace = vec![];

		let precompile_output: &[u32; 16] = &[0u32; 16]; // compute precompile output based on t0 register

		pub fn prove(&self, witness: PrecompileWitness) -> JoltProof; // Need to include Preprocessing struct as an input?
		pub fn verify(&self, proof: JoltProof) -> Result<(), ProofVerifyError>;

Precompile Backend #562

Are you sure you want to change the base?

Precompile Backend #562

Conversation

Roee-87 commented Jan 16, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

moodlezoup left a 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

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