Merge pull request #1097 from powdr-labs/poseidon-regs

Use state registers directly in Poseidon

Author: Leo

riscv-executor/src/lib.rs

@@ -714,7 +714,7 @@ impl<'a, 'b, F: FieldElement> Executor<'a, 'b, F> {
             "shl" => vec![(args[0].u() << args[1].u()).into()],
             "shr" => vec![(args[0].u() >> args[1].u()).into()],
             "split_gl" => {
-                let value = args[0].fe().to_integer();
+                let value = args[0].into_fe().to_integer();
                 // This instruction is only for Goldilocks, so the value must
                 // fit into a u64.
                 let value = value.try_into_u64().unwrap();
@@ -724,13 +724,16 @@ impl<'a, 'b, F: FieldElement> Executor<'a, 'b, F> {
                 vec![lo.into(), hi.into()]
             }
             "poseidon_gl" => {
-                let inputs = args
-                    .iter()
-                    .take(12)
-                    .map(|arg| arg.into_fe())
+                assert!(args.is_empty());
+                let inputs = (0..12)
+                    .map(|i| self.proc.get_reg(format!("P{}", i).as_str()).into_fe())
                     .collect::<Vec<_>>();
                 let result = poseidon_gl::poseidon_gl(&inputs);
-                result.into_iter().map(Elem::Field).collect()
+                (0..4).for_each(|i| {
+                    self.proc
+                        .set_reg(format!("P{}", i).as_str(), Elem::Field(result[i]))
+                });
+                vec![]
             }
             instr => {
                 panic!("unknown instruction: {instr}");

riscv/src/continuations/bootloader.rs

@@ -228,11 +228,7 @@ mstore_bootloader x3 * {PAGE_SIZE_BYTES} + {i} * {BYTES_PER_WORD}, P{reg_index};
 
         // Hash if buffer is full
         if i % 4 == 3 {
-            bootloader.push_str(
-                r#"
-P0, P1, P2, P3 <== poseidon_gl(P0, P1, P2, P3, P4, P5, P6, P7, P8, P9, P10, P11);
-"#,
-            );
+            bootloader.push_str("poseidon_gl;");
         }
     }
 
@@ -294,7 +290,7 @@ P1 <== load_bootloader_input(x2 * {BOOTLOADER_INPUTS_PER_PAGE} + {PAGE_INPUTS_OF
 P2 <== load_bootloader_input(x2 * {BOOTLOADER_INPUTS_PER_PAGE} + {PAGE_INPUTS_OFFSET} + 1 + {WORDS_PER_PAGE} + 4 + {i} * 4 + 2);
 P3 <== load_bootloader_input(x2 * {BOOTLOADER_INPUTS_PER_PAGE} + {PAGE_INPUTS_OFFSET} + 1 + {WORDS_PER_PAGE} + 4 + {i} * 4 + 3);
 bootloader_level_{i}_end:
-P0, P1, P2, P3 <== poseidon_gl(P0, P1, P2, P3, P4, P5, P6, P7, P8, P9, P10, P11);
+    poseidon_gl;
 "#
         ));
     }
@@ -455,9 +451,7 @@ P11 <=X= 0;
 
         // Hash if buffer is full
         if i % 4 == 3 {
-            bootloader.push_str(
-                "P0, P1, P2, P3 <== poseidon_gl(P0, P1, P2, P3, P4, P5, P6, P7, P8, P9, P10, P11);\n",
-            );
+            bootloader.push_str("poseidon_gl;\n");
         }
     }
 

riscv/src/coprocessors.rs

@@ -59,7 +59,7 @@ static POSEIDON_GL_COPROCESSOR: CoProcessor = CoProcessor {
     import: "use std::hash::poseidon_gl::PoseidonGL;",
     instructions: r#"
 // ================== hashing instructions ==============
-instr poseidon_gl A0, A1, A2, A3, A4, A5, A6, A7, A8, A9, A10, A11 -> X, Y, Z, W = poseidon_gl.poseidon_permutation;
+instr poseidon_gl = poseidon_gl.poseidon_permutation P0, P1, P2, P3, P4, P5, P6, P7, P8, P9, P10, P11 -> P0, P1, P2, P3;
 
 "#,
     runtime_function_impl: Some(("poseidon_gl_coprocessor", poseidon_gl_call)),
@@ -199,18 +199,16 @@ impl CoProcessors {
     pub fn registers(&self) -> String {
         // Poseidon has 12 inputs and 4 outputs.
         // The base RISCV machine has 4 assignment registers.
-        // Therefore we need to add 12 assignment registers when using Poseidon.
-        // Moreover, we also need 12 extra general purpose registers to store the
-        // input values.
+        // We need 12 extra general purpose registers to store the
+        // input and values directly.
 
         if !self.coprocessors.contains_key(POSEIDON_GL_COPROCESSOR.name) {
             return String::new();
         }
 
-        let a_regs: Vec<String> = (0..12).map(|i| format!("reg A{}[<=];", i)).collect();
         let p_regs: Vec<String> = (0..12).map(|i| format!("reg P{};", i)).collect();
 
-        [a_regs, p_regs].concat().join("\n")
+        p_regs.join("\n")
     }
 }
 
@@ -244,7 +242,7 @@ fn poseidon_gl_call() -> String {
         )
     };
 
-    let call = "P0, P1, P2, P3 <== poseidon_gl(P0, P1, P2, P3, P4, P5, P6, P7, P8, P9, P10, P11);";
+    let call = "poseidon_gl;";
 
     (0..12)
         .map(decoding)
@@ -273,18 +271,14 @@ pub fn call_every_submachine(coprocessors: &CoProcessors) -> Vec<String> {
     }
     if coprocessors.has(POSEIDON_GL_COPROCESSOR.name) {
         calls.extend(vec![
-            "P0, P1, P2, x10 <== poseidon_gl(P0, P1, P2, P3, P4, P5, P6, P7, P8, P9, P10, P11);"
-                .to_string(),
+            "poseidon_gl;".to_string(),
             "P0 <=X= 0;".to_string(),
             "P1 <=X= 0;".to_string(),
             "P2 <=X= 0;".to_string(),
+            "P3 <=X= 0;".to_string(),
         ]);
     }
     if coprocessors.has(SPLIT_GL_COPROCESSOR.name) {
-        // SplitGL only makes sense when used with PoseidonGL, because there is
-        // no other way to get a value in the full range of Goldilocks into a
-        // register of the RISC-V zkVM.
-        assert!(coprocessors.has(POSEIDON_GL_COPROCESSOR.name));
         calls.push("x10, x11 <== split_gl(x10);".to_string());
     }