From 542e492b7d12bfc0d8e68b6fe7f1b42c7590722b Mon Sep 17 00:00:00 2001 From: Kevin Jue Date: Mon, 19 Feb 2024 07:07:36 -0800 Subject: [PATCH] added fri fold precompile to runtime --- core/src/runtime/syscall.rs | 9 ++ .../syscall/precompiles/fri_fold/execute.rs | 109 ++++++++++++++++++ core/src/syscall/precompiles/fri_fold/mod.rs | 9 ++ core/src/syscall/precompiles/mod.rs | 1 + zkvm/src/syscalls/fri_fold.rs | 22 ++++ zkvm/src/syscalls/mod.rs | 4 + 6 files changed, 154 insertions(+) create mode 100644 core/src/syscall/precompiles/fri_fold/execute.rs create mode 100644 core/src/syscall/precompiles/fri_fold/mod.rs create mode 100644 zkvm/src/syscalls/fri_fold.rs diff --git a/core/src/runtime/syscall.rs b/core/src/runtime/syscall.rs index b37ce1c537..1e1fd3ea02 100644 --- a/core/src/runtime/syscall.rs +++ b/core/src/runtime/syscall.rs @@ -5,6 +5,7 @@ use crate::runtime::{Register, Runtime}; use crate::syscall::precompiles::blake3::Blake3CompressInnerChip; use crate::syscall::precompiles::edwards::EdAddAssignChip; use crate::syscall::precompiles::edwards::EdDecompressChip; +use crate::syscall::precompiles::fri_fold::FriFoldChip; use crate::syscall::precompiles::k256::K256DecompressChip; use crate::syscall::precompiles::keccak256::KeccakPermuteChip; use crate::syscall::precompiles::sha256::{ShaCompressChip, ShaExtendChip}; @@ -60,6 +61,9 @@ pub enum SyscallCode { /// Executes the `BLAKE3_COMPRESS_INNER` precompile. BLAKE3_COMPRESS_INNER = 112, + /// Executes fri fold precompile. + FRI_FOLD = 113, + WRITE = 999, } @@ -80,6 +84,7 @@ impl SyscallCode { 110 => SyscallCode::ENTER_UNCONSTRAINED, 111 => SyscallCode::EXIT_UNCONSTRAINED, 112 => SyscallCode::BLAKE3_COMPRESS_INNER, + 113 => SyscallCode::FRI_FOLD, 999 => SyscallCode::WRITE, _ => panic!("invalid syscall number: {}", value), } @@ -217,6 +222,10 @@ pub fn default_syscall_map() -> HashMap> { SyscallCode::BLAKE3_COMPRESS_INNER, Rc::new(Blake3CompressInnerChip::new()), ); + syscall_map.insert( + SyscallCode::BLAKE3_COMPRESS_INNER, + Rc::new(FriFoldChip::new()), + ); syscall_map.insert( SyscallCode::ENTER_UNCONSTRAINED, Rc::new(SyscallEnterUnconstrained::new()), diff --git a/core/src/syscall/precompiles/fri_fold/execute.rs b/core/src/syscall/precompiles/fri_fold/execute.rs new file mode 100644 index 0000000000..46209a9857 --- /dev/null +++ b/core/src/syscall/precompiles/fri_fold/execute.rs @@ -0,0 +1,109 @@ +use p3_baby_bear::BabyBear; +use p3_field::{ + extension::BinomialExtensionField, AbstractExtensionField, AbstractField, PrimeField32, +}; + +use crate::{ + runtime::{Register, Syscall}, + syscall::precompiles::SyscallContext, +}; + +use super::FriFoldChip; + +impl Syscall for FriFoldChip { + fn num_extra_cycles(&self) -> u32 { + 8 + } + + fn execute(&self, rt: &mut SyscallContext) -> u32 { + // TODO: these will have to be be constrained, but can do it later. + // Read `input_mem_ptr` from register a0. + let input_mem_ptr = rt.register_unsafe(Register::X10); + if input_mem_ptr % 4 != 0 { + panic!(); + } + // Read `output_mem_ptr` from register a1. + let output_mem_ptr = rt.register_unsafe(Register::X11); + if output_mem_ptr % 4 != 0 { + panic!(); + } + let (input_read_records, input_values) = rt.mr_slice(input_mem_ptr, 14); + + let x = BabyBear::from_canonical_u32(input_values[0]); + let alpha = BinomialExtensionField::::from_base_slice( + input_values[1..5] + .iter() + .map(|x| BabyBear::from_canonical_u32(*x)) + .collect::>() + .as_slice(), + ); + let z = BinomialExtensionField::::from_base_slice( + input_values[5..9] + .iter() + .map(|x| BabyBear::from_canonical_u32(*x)) + .collect::>() + .as_slice(), + ); + let p_at_z = BinomialExtensionField::::from_base_slice( + input_values[9..13] + .iter() + .map(|x| BabyBear::from_canonical_u32(*x)) + .collect::>() + .as_slice(), + ); + let p_at_x = BabyBear::from_canonical_u32(input_values[13]); + + // Read ro[log_height] and alpha_pow[log_height] address + let (output_read_records, output_addresses) = rt.mr_slice(output_mem_ptr, 2); + let ro_address = output_addresses[0]; + let alpha_pow_address = output_addresses[1]; + + let (ro_read_records, ro_values) = rt.mr_slice(ro_address, 4); + let mut ro = BinomialExtensionField::::from_base_slice( + ro_values + .iter() + .map(|&x| BabyBear::from_canonical_u32(x)) + .collect::>() + .as_slice(), + ); + + let (alpha_read_records, alpha_values) = rt.mr_slice(alpha_pow_address, 4); + let mut alpha_pow = BinomialExtensionField::::from_base_slice( + alpha_values + .iter() + .map(|&x| BabyBear::from_canonical_u32(x)) + .collect::>() + .as_slice(), + ); + + rt.clk += 4; + + let quotient = (-p_at_z + p_at_x) / (-z + x); + ro += alpha_pow * quotient; + alpha_pow *= alpha; + + let ro_write_records = rt.mw_slice( + ro_address, + ro.as_base_slice() + .iter() + .map(|x: &BabyBear| x.as_canonical_u32()) + .collect::>() + .as_slice(), + ); + let alpha_pow_write_records = rt.mw_slice( + alpha_pow_address, + alpha_pow + .as_base_slice() + .iter() + .map(|x: &BabyBear| x.as_canonical_u32()) + .collect::>() + .as_slice(), + ); + + rt.clk += 4; + + // TODO. Push the fri fold event. + + input_mem_ptr + } +} diff --git a/core/src/syscall/precompiles/fri_fold/mod.rs b/core/src/syscall/precompiles/fri_fold/mod.rs new file mode 100644 index 0000000000..0a816fad1e --- /dev/null +++ b/core/src/syscall/precompiles/fri_fold/mod.rs @@ -0,0 +1,9 @@ +mod execute; + +pub struct FriFoldChip {} + +impl FriFoldChip { + pub fn new() -> Self { + Self {} + } +} diff --git a/core/src/syscall/precompiles/mod.rs b/core/src/syscall/precompiles/mod.rs index afac749ffc..1a2d6625ef 100644 --- a/core/src/syscall/precompiles/mod.rs +++ b/core/src/syscall/precompiles/mod.rs @@ -1,5 +1,6 @@ pub mod blake3; pub mod edwards; +pub mod fri_fold; pub mod k256; pub mod keccak256; pub mod sha256; diff --git a/zkvm/src/syscalls/fri_fold.rs b/zkvm/src/syscalls/fri_fold.rs new file mode 100644 index 0000000000..88c324bd00 --- /dev/null +++ b/zkvm/src/syscalls/fri_fold.rs @@ -0,0 +1,22 @@ +#[cfg(target_os = "zkvm")] +use core::arch::asm; + +/// Fri fold operation. +/// +/// The result is written to the addresses in the output mem entries. +#[allow(unused_variables)] +#[no_mangle] +pub extern "C" fn syscall_fri_fold(input_mem_ptr: *mut u32, output_mem_ptr: *mut u32) { + #[cfg(target_os = "zkvm")] + unsafe { + asm!( + "ecall", + in("t0") crate::syscalls::FRI_FOLD, + in("a0") input_mem_ptr, + in("a1") output_mem_ptr + ); + } + + #[cfg(not(target_os = "zkvm"))] + unreachable!() +} diff --git a/zkvm/src/syscalls/mod.rs b/zkvm/src/syscalls/mod.rs index cb72fc4c1c..14dc704756 100644 --- a/zkvm/src/syscalls/mod.rs +++ b/zkvm/src/syscalls/mod.rs @@ -49,6 +49,7 @@ extern "C" { pub fn syscall_secp256k1_decompress(point: &mut [u8; 64], is_odd: bool); pub fn syscall_keccak_permute(state: *mut u64); pub fn syscall_blake3_compress_inner(p: *mut u32, q: *const u32); + pub fn syscall_fri_fold(input_mem_ptr: *mut u32, output_mem_ptr: *mut u32); pub fn syscall_enter_unconstrained() -> bool; pub fn syscall_exit_unconstrained(); pub fn sys_alloc_aligned(bytes: usize, align: usize) -> *mut u8; @@ -96,5 +97,8 @@ pub const EXIT_UNCONSTRAINED: u32 = 111; /// Executes `BLAKE3_COMPRESS_INNER`. pub const BLAKE3_COMPRESS_INNER: u32 = 112; +/// Executes fri fold operation. +pub const FRI_FOLD: u32 = 113; + /// Writes to a file descriptor. Currently only used for `STDOUT/STDERR`. pub const WRITE: u32 = 999;