Merge pull request powdr-labs#678 from powdr-labs/poseidon-gl-rust

poseidon gl rust

compiler/benches/executor_benchmark.rs

@@ -8,13 +8,15 @@ use mktemp::Temp;
 use number::{FieldElement, GoldilocksField};
 use riscv::{compile_rust_crate_to_riscv_asm, compiler};
 
+use riscv::CoProcessors;
+
 type T = GoldilocksField;
 
 fn get_pil() -> Analyzed<T> {
     let tmp_dir = Temp::new_dir().unwrap();
     let riscv_asm_files =
         compile_rust_crate_to_riscv_asm("../riscv/tests/riscv_data/keccak/Cargo.toml", &tmp_dir);
-    let contents = compiler::compile(riscv_asm_files);
+    let contents = compiler::compile(riscv_asm_files, &CoProcessors::base());
     let parsed = parser::parse_asm::<T>(None, &contents).unwrap();
     let resolved = importer::resolve(None, parsed).unwrap();
     let analyzed = analyze(resolved).unwrap();

powdr_cli/src/main.rs

@@ -122,6 +122,10 @@ enum Commands {
         #[arg(short, long)]
         #[arg(value_parser = clap_enum_variants!(BackendType))]
         prove_with: Option<BackendType>,
+
+        /// Comma-separated list of coprocessors.
+        #[arg(long)]
+        coprocessors: Option<String>,
     },
 
     /// Compiles riscv assembly to powdr assembly and then to PIL
@@ -156,6 +160,10 @@ enum Commands {
         #[arg(short, long)]
         #[arg(value_parser = clap_enum_variants!(BackendType))]
         prove_with: Option<BackendType>,
+
+        /// Comma-separated list of coprocessors.
+        #[arg(long)]
+        coprocessors: Option<String>,
     },
 
     Prove {
@@ -285,13 +293,21 @@ fn run_command(command: Commands) {
             output_directory,
             force,
             prove_with,
+            coprocessors,
         } => {
+            let coprocessors = match coprocessors {
+                Some(list) => {
+                    riscv::CoProcessors::try_from(list.split(',').collect::<Vec<_>>()).unwrap()
+                }
+                None => riscv::CoProcessors::base(),
+            };
             if let Err(errors) = call_with_field!(run_rust::<field>(
                 &file,
                 split_inputs(&inputs),
                 Path::new(&output_directory),
                 force,
-                prove_with
+                prove_with,
+                coprocessors
             )) {
                 eprintln!("Errors:");
                 for e in errors {
@@ -306,6 +322,7 @@ fn run_command(command: Commands) {
             output_directory,
             force,
             prove_with,
+            coprocessors,
         } => {
             assert!(!files.is_empty());
             let name = if files.len() == 1 {
@@ -314,13 +331,20 @@ fn run_command(command: Commands) {
                 Cow::Borrowed("output")
             };
 
+            let coprocessors = match coprocessors {
+                Some(list) => {
+                    riscv::CoProcessors::try_from(list.split(',').collect::<Vec<_>>()).unwrap()
+                }
+                None => riscv::CoProcessors::base(),
+            };
             if let Err(errors) = call_with_field!(run_riscv_asm::<field>(
                 &name,
                 files.into_iter(),
                 split_inputs(&inputs),
                 Path::new(&output_directory),
                 force,
-                prove_with
+                prove_with,
+                coprocessors
             )) {
                 eprintln!("Errors:");
                 for e in errors {
@@ -386,7 +410,7 @@ fn run_command(command: Commands) {
         } => {
             call_with_field!(setup::<field>(size, dir, backend));
         }
-    }
+    };
 }
 
 fn setup<F: FieldElement>(size: u64, dir: String, backend_type: BackendType) {
@@ -403,15 +427,18 @@ fn write_backend_to_fs<F: FieldElement>(be: &dyn Backend<F>, output_dir: &Path)
     params_writer.flush().unwrap();
     log::info!("Wrote params.bin.");
 }
+
 fn run_rust<F: FieldElement>(
     file_name: &str,
     inputs: Vec<F>,
     output_dir: &Path,
     force_overwrite: bool,
     prove_with: Option<BackendType>,
+    coprocessors: riscv::CoProcessors,
 ) -> Result<(), Vec<String>> {
-    let (asm_file_path, asm_contents) = compile_rust(file_name, output_dir, force_overwrite)
-        .ok_or_else(|| vec!["could not compile rust".to_string()])?;
+    let (asm_file_path, asm_contents) =
+        compile_rust(file_name, output_dir, force_overwrite, &coprocessors)
+            .ok_or_else(|| vec!["could not compile rust".to_string()])?;
 
     compile_asm_string(
         asm_file_path.to_str().unwrap(),
@@ -431,10 +458,16 @@ fn run_riscv_asm<F: FieldElement>(
     output_dir: &Path,
     force_overwrite: bool,
     prove_with: Option<BackendType>,
+    coprocessors: riscv::CoProcessors,
 ) -> Result<(), Vec<String>> {
-    let (asm_file_path, asm_contents) =
-        compile_riscv_asm(original_file_name, file_names, output_dir, force_overwrite)
-            .ok_or_else(|| vec!["could not compile RISC-V assembly".to_string()])?;
+    let (asm_file_path, asm_contents) = compile_riscv_asm(
+        original_file_name,
+        file_names,
+        output_dir,
+        force_overwrite,
+        &coprocessors,
+    )
+    .ok_or_else(|| vec!["could not compile RISC-V assembly".to_string()])?;
 
     compile_asm_string(
         asm_file_path.to_str().unwrap(),

riscv/runtime/src/coprocessors.rs

@@ -1,12 +1,36 @@
 // This is a dummy implementation of Poseidon hash,
-// which will be replaced with a call to the poseidon
-// coporocessor during compilation.
+// which will be replaced with a call to the Poseidon
+// coprocessor during compilation.
 // The function itself will be removed by the compiler
 // during the reachability analysis.
 extern "C" {
-    fn poseidon_coprocessor(a: u32, b: u32) -> u32;
+    fn poseidon_gl_coprocessor(data: *mut [u64; 12]);
 }
 
-pub fn poseidon_hash(a: u32, b: u32) -> u32 {
-    unsafe { poseidon_coprocessor(a, b) }
+const GOLDILOCKS: u64 = 0xffffffff00000001;
+
+/// Calls the low level Poseidon coprocessor in PIL, where
+/// the last 4 elements are the "cap"
+/// and the return value is placed in data[0:4].
+/// The safe version below also checks that each u64 element
+/// is less than the Goldilocks field.
+/// The unsafe version does not perform such checks.
+pub fn poseidon_gl(mut data: [u64; 12]) -> [u64; 4] {
+    for &n in data.iter() {
+        assert!(n < GOLDILOCKS);
+    }
+
+    unsafe {
+        poseidon_gl_coprocessor(&mut data as *mut [u64; 12]);
+    }
+
+    [data[0], data[1], data[2], data[3]]
+}
+
+pub fn poseidon_gl_unsafe(mut data: [u64; 12]) -> [u64; 4] {
+    unsafe {
+        poseidon_gl_coprocessor(&mut data as *mut [u64; 12]);
+    }
+
+    [data[0], data[1], data[2], data[3]]
 }