From c83fef855db08864494a1d894646c365b59ee010 Mon Sep 17 00:00:00 2001 From: Thomas Coratger Date: Sat, 16 Aug 2025 19:17:40 +0200 Subject: [PATCH] instructions: update the precompiles --- Cargo.toml | 14 ++-- crates/leanVm/Cargo.toml | 2 + .../src/bytecode/instruction/dot_product.rs | 76 ++++++++++++++++++ .../src/bytecode/instruction/extension_mul.rs | 74 ----------------- crates/leanVm/src/bytecode/instruction/mod.rs | 25 +++--- .../bytecode/instruction/multilinear_eval.rs | 79 +++++++++++++++++++ crates/leanVm/src/core.rs | 8 -- crates/leanVm/src/errors/vm.rs | 2 + 8 files changed, 181 insertions(+), 99 deletions(-) create mode 100644 crates/leanVm/src/bytecode/instruction/dot_product.rs delete mode 100644 crates/leanVm/src/bytecode/instruction/extension_mul.rs create mode 100644 crates/leanVm/src/bytecode/instruction/multilinear_eval.rs diff --git a/Cargo.toml b/Cargo.toml index b04347f3..14e93eb2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -36,12 +36,14 @@ lean-isa = { path = "crates/leanIsa" } lean-snark = { path = "crates/leanSnark" } lean-vm = { path = "crates/leanVm" } -p3-field = { git = "https://github.com/Plonky3/Plonky3.git", rev = "d0c4a36" } -p3-symmetric = { git = "https://github.com/Plonky3/Plonky3.git", rev = "d0c4a36" } -p3-baby-bear = { git = "https://github.com/Plonky3/Plonky3.git", rev = "d0c4a36" } -p3-koala-bear = { git = "https://github.com/Plonky3/Plonky3.git", rev = "d0c4a36" } -p3-air = { git = "https://github.com/Plonky3/Plonky3.git", rev = "d0c4a36" } -p3-matrix = { git = "https://github.com/Plonky3/Plonky3.git", rev = "d0c4a36" } +p3-field = { git = "https://github.com/Plonky3/Plonky3.git", rev = "5ebf8e4" } +p3-symmetric = { git = "https://github.com/Plonky3/Plonky3.git", rev = "5ebf8e4" } +p3-baby-bear = { git = "https://github.com/Plonky3/Plonky3.git", rev = "5ebf8e4" } +p3-koala-bear = { git = "https://github.com/Plonky3/Plonky3.git", rev = "5ebf8e4" } +p3-air = { git = "https://github.com/Plonky3/Plonky3.git", rev = "5ebf8e4" } +p3-matrix = { git = "https://github.com/Plonky3/Plonky3.git", rev = "5ebf8e4" } + +whir-p3 = { git = "https://github.com/tcoratger/whir-p3.git", rev = "df2241d" } thiserror = "2.0" proptest = "1.7" diff --git a/crates/leanVm/Cargo.toml b/crates/leanVm/Cargo.toml index 70c1b3e7..e4dc950f 100644 --- a/crates/leanVm/Cargo.toml +++ b/crates/leanVm/Cargo.toml @@ -16,6 +16,8 @@ p3-symmetric.workspace = true p3-air.workspace = true p3-matrix.workspace = true +whir-p3.workspace = true + thiserror.workspace = true num-traits.workspace = true diff --git a/crates/leanVm/src/bytecode/instruction/dot_product.rs b/crates/leanVm/src/bytecode/instruction/dot_product.rs new file mode 100644 index 00000000..327f28fe --- /dev/null +++ b/crates/leanVm/src/bytecode/instruction/dot_product.rs @@ -0,0 +1,76 @@ +use p3_field::{BasedVectorSpace, dot_product}; + +use crate::{ + bytecode::operand::{MemOrConstant, MemOrFp}, + constant::{DIMENSION, EF, F}, + context::run_context::RunContext, + errors::vm::VirtualMachineError, + memory::{address::MemoryAddress, manager::MemoryManager}, +}; + +/// An instruction to compute the dot product of two vectors of extension field elements. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct DotProductInstruction { + /// The first operand, a pointer to the start of the first vector. + pub arg0: MemOrConstant, + /// The second operand, a pointer to the start of the second vector. + pub arg1: MemOrConstant, + /// The destination pointer for the result. + pub res: MemOrFp, + /// The number of elements in each vector. + pub size: usize, +} + +impl DotProductInstruction { + /// Executes the `DotProductExtensionExtension` instruction. + /// + /// This function performs the following steps: + /// 1. Resolves the pointers to the two input vectors and the output location from memory. + /// 2. Reads the vector data for both operands from memory. + /// 3. Computes the dot product of the two vectors in the extension field. + /// 4. Writes the resulting extension field element back to the specified memory location. + pub fn execute( + &self, + run_context: &RunContext, + memory_manager: &mut MemoryManager, + ) -> Result<(), VirtualMachineError> { + // Resolve the memory addresses for the two input vectors and the result. + let ptr_arg_0: MemoryAddress = run_context + .value_from_mem_or_constant(&self.arg0, memory_manager)? + .try_into()?; + let ptr_arg_1: MemoryAddress = run_context + .value_from_mem_or_constant(&self.arg1, memory_manager)? + .try_into()?; + let ptr_res: MemoryAddress = run_context + .value_from_mem_or_fp(&self.res, memory_manager)? + .try_into()?; + + // Read the first vector slice from memory. + let slice_0: Vec = (0..self.size) + .map(|i| { + let addr = (ptr_arg_0 + i)?; + let vector_coeffs = memory_manager.memory.get_array_as::(addr)?; + EF::from_basis_coefficients_slice(&vector_coeffs) + .ok_or(VirtualMachineError::InvalidExtensionField) + }) + .collect::>()?; + + // Read the second vector slice from memory. + let slice_1: Vec = (0..self.size) + .map(|i| { + let addr = (ptr_arg_1 + i)?; + let vector_coeffs = memory_manager.memory.get_array_as::(addr)?; + EF::from_basis_coefficients_slice(&vector_coeffs) + .ok_or(VirtualMachineError::InvalidExtensionField) + }) + .collect::>()?; + + // Compute the dot product of the two slices by converting them into iterators. + let dot_product_res = dot_product::(slice_0.into_iter(), slice_1.into_iter()); + + // Write the resulting vector back to memory. + memory_manager.load_data::(ptr_res, dot_product_res.as_basis_coefficients_slice())?; + + Ok(()) + } +} diff --git a/crates/leanVm/src/bytecode/instruction/extension_mul.rs b/crates/leanVm/src/bytecode/instruction/extension_mul.rs deleted file mode 100644 index d723b4f1..00000000 --- a/crates/leanVm/src/bytecode/instruction/extension_mul.rs +++ /dev/null @@ -1,74 +0,0 @@ -use p3_field::BasedVectorSpace; - -use crate::{ - constant::{DIMENSION, EF, F}, - context::run_context::RunContext, - errors::vm::VirtualMachineError, - memory::{address::MemoryAddress, manager::MemoryManager}, -}; - -/// Degree-8 extension field multiplication of 2 elements -> 1 result. -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct ExtensionMulInstruction { - /// An array of three offsets from `fp`. These point to the start of the 8-cell memory blocks - /// for the two input extension field elements and the resulting output element. - args: [usize; 3], -} - -impl ExtensionMulInstruction { - /// Executes the `ExtensionMul` instruction. - /// - /// Multiplies two extension field elements stored in memory and writes the result back. - /// - /// ## Memory Layout - /// The instruction takes three argument offsets `[a, b, c]` from the frame pointer `fp`: - /// - /// - `fp + a`: address of a pointer to the first operand - /// - `fp + b`: address of a pointer to the second operand - /// - `fp + c`: address of a pointer to the output location - /// - /// The multiplication is: - /// - /// ```text - /// m_vec[ptr_out] = EF::from(m_vec[ptr_lhs]) * EF::from(m_vec[ptr_rhs]) - /// ``` - /// - /// where `ptr_lhs`, `ptr_rhs`, and `ptr_out` are memory addresses stored at `fp + args[0..=2]`. - pub fn execute( - &self, - run_context: &RunContext, - memory_manager: &mut MemoryManager, - ) -> Result<(), VirtualMachineError> { - // Get the frame pointer. - let fp = run_context.fp; - - // Read the memory addresses where the operands and result will reside. - let ptr_lhs: MemoryAddress = memory_manager.memory.get_as((fp + self.args[0])?)?; - let ptr_rhs: MemoryAddress = memory_manager.memory.get_as((fp + self.args[1])?)?; - let ptr_out: MemoryAddress = memory_manager.memory.get_as((fp + self.args[2])?)?; - - // Load the `[F; EF::DIMENSION]` input arrays from memory and convert them into EF elements. - let lhs = EF::from_basis_coefficients_slice( - &memory_manager - .memory - .get_array_as::(ptr_lhs)?, - ) - .unwrap(); - - let rhs = EF::from_basis_coefficients_slice( - &memory_manager - .memory - .get_array_as::(ptr_rhs)?, - ) - .unwrap(); - - // Perform the field multiplication in the extension field. - let product = lhs * rhs; - - // Write the result converted back into `[F; EF::DIMENSION]` to memory. - let result_coeffs: &[F] = product.as_basis_coefficients_slice(); - memory_manager.load_data(ptr_out, result_coeffs)?; - - Ok(()) - } -} diff --git a/crates/leanVm/src/bytecode/instruction/mod.rs b/crates/leanVm/src/bytecode/instruction/mod.rs index 34a73c84..5444a9ac 100644 --- a/crates/leanVm/src/bytecode/instruction/mod.rs +++ b/crates/leanVm/src/bytecode/instruction/mod.rs @@ -1,7 +1,8 @@ use computation::ComputationInstruction; use deref::DerefInstruction; -use extension_mul::ExtensionMulInstruction; +use dot_product::DotProductInstruction; use jump::JumpIfNotZeroInstruction; +use multilinear_eval::MultilinearEvalInstruction; use p3_symmetric::Permutation; use poseidon16::Poseidon2_16Instruction; use poseidon24::Poseidon2_24Instruction; @@ -15,15 +16,13 @@ use crate::{ pub mod computation; pub mod deref; -pub mod extension_mul; +pub mod dot_product; pub mod jump; +pub mod multilinear_eval; pub mod poseidon16; pub mod poseidon24; /// Defines the instruction set for this zkVM, specialized for the `AggregateMerge` logic. -/// -/// The ISA is minimal and includes basic arithmetic, memory operations, control flow, -/// and powerful precompiles for hashing and extension field arithmetic. #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] pub enum Instruction { /// Performs a basic arithmetic computation: `res = arg_a op arg_b`. @@ -53,10 +52,12 @@ pub enum Instruction { /// /// It reads 6 pointers from memory, starting at `m[fp+shift]`. Poseidon2_24(Poseidon2_24Instruction), - /// **Precompile** for multiplication in the degree-8 extension field. - /// - /// This is important for speeding up recursive proof verification (`snark_verify`). - ExtensionMul(ExtensionMulInstruction), + + /// Dot product of two vectors of extension field elements. + DotProduct(DotProductInstruction), + + /// Evaluation of a multilinear polynomial over the extension field. + MultilinearEval(MultilinearEvalInstruction), } impl Instruction { @@ -100,8 +101,10 @@ impl Instruction { Self::Poseidon2_24(instruction) => { instruction.execute(run_context, memory_manager, perm24) } - // Handle the extension field multiplication precompile. - Self::ExtensionMul(instruction) => instruction.execute(run_context, memory_manager), + // Handle the dot product precompile. + Self::DotProduct(instruction) => instruction.execute(run_context, memory_manager), + // Handle the multilinear evaluation precompile. + Self::MultilinearEval(instruction) => instruction.execute(run_context, memory_manager), } } } diff --git a/crates/leanVm/src/bytecode/instruction/multilinear_eval.rs b/crates/leanVm/src/bytecode/instruction/multilinear_eval.rs new file mode 100644 index 00000000..3e1c6474 --- /dev/null +++ b/crates/leanVm/src/bytecode/instruction/multilinear_eval.rs @@ -0,0 +1,79 @@ +use p3_field::BasedVectorSpace; +use whir_p3::poly::{coeffs::CoefficientList, multilinear::MultilinearPoint}; + +use crate::{ + bytecode::operand::{MemOrConstant, MemOrFp}, + constant::{DIMENSION, EF, F}, + context::run_context::RunContext, + errors::vm::VirtualMachineError, + memory::{address::MemoryAddress, manager::MemoryManager}, +}; + +/// An instruction to evaluate a multilinear polynomial at a point in the extension field. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct MultilinearEvalInstruction { + /// A pointer to the polynomial's coefficients. + pub coeffs: MemOrConstant, + /// A pointer to the evaluation point's coordinates. + pub point: MemOrConstant, + /// The destination pointer for the result. + pub res: MemOrFp, + /// The number of variables in the multilinear polynomial. + pub n_vars: usize, +} + +impl MultilinearEvalInstruction { + /// Executes the `MultilinearEval` instruction. + /// + /// This function performs the following steps: + /// 1. Resolves pointers to the polynomial coefficients, evaluation point, and output location. + /// 2. Reads the polynomial coefficients (base field elements) from memory. + /// 3. Reads the evaluation point (extension field elements) from memory. + /// 4. Evaluates the polynomial at the given point. + /// 5. Writes the resulting extension field element back to memory. + pub fn execute( + &self, + run_context: &RunContext, + memory_manager: &mut MemoryManager, + ) -> Result<(), VirtualMachineError> { + // Resolve the memory addresses for the coefficients, point, and result. + let ptr_coeffs: MemoryAddress = run_context + .value_from_mem_or_constant(&self.coeffs, memory_manager)? + .try_into()?; + let ptr_point: MemoryAddress = run_context + .value_from_mem_or_constant(&self.point, memory_manager)? + .try_into()?; + let ptr_res: MemoryAddress = run_context + .value_from_mem_or_fp(&self.res, memory_manager)? + .try_into()?; + + // Read the polynomial coefficients from memory. + // + // The total number of coefficients is 2^n_vars. + let num_coeffs = 1 << self.n_vars; + let slice_coeffs: Vec = (0..num_coeffs) + .map(|i| { + let addr = (ptr_coeffs + i)?; + memory_manager.memory.get_as(addr) + }) + .collect::>()?; + + // Read the evaluation point from memory. + let point: Vec = (0..self.n_vars) + .map(|i| { + let addr = (ptr_point + i)?; + let vector_coeffs = memory_manager.memory.get_array_as::(addr)?; + EF::from_basis_coefficients_slice(&vector_coeffs) + .ok_or(VirtualMachineError::InvalidExtensionField) + }) + .collect::>()?; + + // Evaluate the multilinear polynomial. + let eval = CoefficientList::new(slice_coeffs).evaluate(&MultilinearPoint(point)); + + // Write the resulting vector back to memory. + memory_manager.load_data::(ptr_res, eval.as_basis_coefficients_slice())?; + + Ok(()) + } +} diff --git a/crates/leanVm/src/core.rs b/crates/leanVm/src/core.rs index b1b2e0c0..1c834b4a 100644 --- a/crates/leanVm/src/core.rs +++ b/crates/leanVm/src/core.rs @@ -29,7 +29,6 @@ pub struct VirtualMachine { pub(crate) cpu_cycles: u64, pub(crate) poseidon16_calls: u64, pub(crate) poseidon24_calls: u64, - pub(crate) ext_mul_calls: u64, // A string buffer to hold output from the Print hint. pub(crate) std_out: String, } @@ -52,7 +51,6 @@ where cpu_cycles: 0, poseidon16_calls: 0, poseidon24_calls: 0, - ext_mul_calls: 0, std_out: String::new(), } } @@ -289,7 +287,6 @@ where self.cpu_cycles = 0; self.poseidon16_calls = 0; self.poseidon24_calls = 0; - self.ext_mul_calls = 0; self.std_out.clear(); } @@ -298,7 +295,6 @@ where match instruction { Instruction::Poseidon2_16(_) => self.poseidon16_calls += 1, Instruction::Poseidon2_24(_) => self.poseidon24_calls += 1, - Instruction::ExtensionMul(_) => self.ext_mul_calls += 1, _ => (), // Other instructions do not have special counters. } } @@ -412,10 +408,6 @@ where self.cpu_cycles / total_poseidon_calls ); } - - if self.ext_mul_calls > 0 { - println!("ExtensionMul calls: {}", self.ext_mul_calls,); - } } } diff --git a/crates/leanVm/src/errors/vm.rs b/crates/leanVm/src/errors/vm.rs index b3ced7e7..85029a0e 100644 --- a/crates/leanVm/src/errors/vm.rs +++ b/crates/leanVm/src/errors/vm.rs @@ -24,4 +24,6 @@ pub enum VirtualMachineError { TooManyUnknownOperands, #[error("Program counter (pc) is out of bounds.")] PCOutOfBounds, + #[error("Invalid extensionn field.")] + InvalidExtensionField, }