openvm_bigint_circuit/
mult.rs

1use std::borrow::{Borrow, BorrowMut};
2
3use openvm_bigint_transpiler::Rv32Mul256Opcode;
4use openvm_circuit::{arch::*, system::memory::online::GuestMemory};
5use openvm_circuit_primitives_derive::AlignedBytesBorrow;
6use openvm_instructions::{
7    instruction::Instruction,
8    program::DEFAULT_PC_STEP,
9    riscv::{RV32_MEMORY_AS, RV32_REGISTER_AS},
10    LocalOpcode,
11};
12use openvm_rv32im_circuit::MultiplicationExecutor;
13use openvm_rv32im_transpiler::MulOpcode;
14use openvm_stark_backend::p3_field::PrimeField32;
15
16use crate::{
17    common::{bytes_to_u32_array, read_int256, u32_array_to_bytes, write_int256},
18    AluAdapterExecutor, Rv32Multiplication256Executor, INT256_NUM_LIMBS,
19};
20
21impl Rv32Multiplication256Executor {
22    pub fn new(adapter: AluAdapterExecutor, offset: usize) -> Self {
23        Self(MultiplicationExecutor::new(adapter, offset))
24    }
25}
26
27#[derive(AlignedBytesBorrow, Clone)]
28#[repr(C)]
29struct MultPreCompute {
30    a: u8,
31    b: u8,
32    c: u8,
33}
34
35impl<F: PrimeField32> InterpreterExecutor<F> for Rv32Multiplication256Executor {
36    fn pre_compute_size(&self) -> usize {
37        size_of::<MultPreCompute>()
38    }
39
40    #[cfg(not(feature = "tco"))]
41    fn pre_compute<Ctx>(
42        &self,
43        pc: u32,
44        inst: &Instruction<F>,
45        data: &mut [u8],
46    ) -> Result<ExecuteFunc<F, Ctx>, StaticProgramError>
47    where
48        Ctx: ExecutionCtxTrait,
49    {
50        let data: &mut MultPreCompute = data.borrow_mut();
51        self.pre_compute_impl(pc, inst, data)?;
52        Ok(execute_e1_impl)
53    }
54
55    #[cfg(feature = "tco")]
56    fn handler<Ctx>(
57        &self,
58        pc: u32,
59        inst: &Instruction<F>,
60        data: &mut [u8],
61    ) -> Result<Handler<F, Ctx>, StaticProgramError>
62    where
63        Ctx: ExecutionCtxTrait,
64    {
65        let data: &mut MultPreCompute = data.borrow_mut();
66        self.pre_compute_impl(pc, inst, data)?;
67        Ok(execute_e1_handler)
68    }
69}
70
71#[cfg(feature = "aot")]
72impl<F: PrimeField32> AotExecutor<F> for Rv32Multiplication256Executor {}
73
74impl<F: PrimeField32> InterpreterMeteredExecutor<F> for Rv32Multiplication256Executor {
75    fn metered_pre_compute_size(&self) -> usize {
76        size_of::<E2PreCompute<MultPreCompute>>()
77    }
78
79    #[cfg(not(feature = "tco"))]
80    fn metered_pre_compute<Ctx>(
81        &self,
82        chip_idx: usize,
83        pc: u32,
84        inst: &Instruction<F>,
85        data: &mut [u8],
86    ) -> Result<ExecuteFunc<F, Ctx>, StaticProgramError>
87    where
88        Ctx: MeteredExecutionCtxTrait,
89    {
90        let data: &mut E2PreCompute<MultPreCompute> = data.borrow_mut();
91        data.chip_idx = chip_idx as u32;
92        self.pre_compute_impl(pc, inst, &mut data.data)?;
93        Ok(execute_e2_impl)
94    }
95
96    #[cfg(feature = "tco")]
97    fn metered_handler<Ctx>(
98        &self,
99        chip_idx: usize,
100        pc: u32,
101        inst: &Instruction<F>,
102        data: &mut [u8],
103    ) -> Result<Handler<F, Ctx>, StaticProgramError>
104    where
105        Ctx: MeteredExecutionCtxTrait,
106    {
107        let data: &mut E2PreCompute<MultPreCompute> = data.borrow_mut();
108        data.chip_idx = chip_idx as u32;
109        self.pre_compute_impl(pc, inst, &mut data.data)?;
110        Ok(execute_e2_handler)
111    }
112}
113
114#[cfg(feature = "aot")]
115impl<F: PrimeField32> AotMeteredExecutor<F> for Rv32Multiplication256Executor {}
116
117#[inline(always)]
118unsafe fn execute_e12_impl<F: PrimeField32, CTX: ExecutionCtxTrait>(
119    pre_compute: &MultPreCompute,
120    exec_state: &mut VmExecState<F, GuestMemory, CTX>,
121) {
122    let rs1_ptr = exec_state.vm_read::<u8, 4>(RV32_REGISTER_AS, pre_compute.b as u32);
123    let rs2_ptr = exec_state.vm_read::<u8, 4>(RV32_REGISTER_AS, pre_compute.c as u32);
124    let rd_ptr = exec_state.vm_read::<u8, 4>(RV32_REGISTER_AS, pre_compute.a as u32);
125    let rs1 = read_int256(exec_state, RV32_MEMORY_AS, u32::from_le_bytes(rs1_ptr));
126    let rs2 = read_int256(exec_state, RV32_MEMORY_AS, u32::from_le_bytes(rs2_ptr));
127    let rd = u256_mul(rs1, rs2);
128    write_int256(exec_state, RV32_MEMORY_AS, u32::from_le_bytes(rd_ptr), &rd);
129
130    let pc = exec_state.pc();
131    exec_state.set_pc(pc.wrapping_add(DEFAULT_PC_STEP));
132}
133
134#[create_handler]
135#[inline(always)]
136unsafe fn execute_e1_impl<F: PrimeField32, CTX: ExecutionCtxTrait>(
137    pre_compute: *const u8,
138    exec_state: &mut VmExecState<F, GuestMemory, CTX>,
139) {
140    let pre_compute: &MultPreCompute =
141        std::slice::from_raw_parts(pre_compute, size_of::<MultPreCompute>()).borrow();
142    execute_e12_impl(pre_compute, exec_state);
143}
144
145#[create_handler]
146#[inline(always)]
147unsafe fn execute_e2_impl<F: PrimeField32, CTX: MeteredExecutionCtxTrait>(
148    pre_compute: *const u8,
149    exec_state: &mut VmExecState<F, GuestMemory, CTX>,
150) {
151    let pre_compute: &E2PreCompute<MultPreCompute> =
152        std::slice::from_raw_parts(pre_compute, size_of::<E2PreCompute<MultPreCompute>>()).borrow();
153    exec_state
154        .ctx
155        .on_height_change(pre_compute.chip_idx as usize, 1);
156    execute_e12_impl(&pre_compute.data, exec_state);
157}
158
159impl Rv32Multiplication256Executor {
160    fn pre_compute_impl<F: PrimeField32>(
161        &self,
162        pc: u32,
163        inst: &Instruction<F>,
164        data: &mut MultPreCompute,
165    ) -> Result<(), StaticProgramError> {
166        let Instruction {
167            opcode,
168            a,
169            b,
170            c,
171            d,
172            e,
173            ..
174        } = inst;
175        let e_u32 = e.as_canonical_u32();
176        if d.as_canonical_u32() != RV32_REGISTER_AS || e_u32 != RV32_MEMORY_AS {
177            return Err(StaticProgramError::InvalidInstruction(pc));
178        }
179        let local_opcode =
180            MulOpcode::from_usize(opcode.local_opcode_idx(Rv32Mul256Opcode::CLASS_OFFSET));
181        if local_opcode != MulOpcode::MUL {
182            return Err(StaticProgramError::InvalidInstruction(pc));
183        }
184        *data = MultPreCompute {
185            a: a.as_canonical_u32() as u8,
186            b: b.as_canonical_u32() as u8,
187            c: c.as_canonical_u32() as u8,
188        };
189        Ok(())
190    }
191}
192
193#[inline(always)]
194pub(crate) fn u256_mul(
195    rs1: [u8; INT256_NUM_LIMBS],
196    rs2: [u8; INT256_NUM_LIMBS],
197) -> [u8; INT256_NUM_LIMBS] {
198    let rs1_u64: [u32; 8] = bytes_to_u32_array(rs1);
199    let rs2_u64: [u32; 8] = bytes_to_u32_array(rs2);
200    let mut rd = [0u32; 8];
201    for i in 0..8 {
202        let mut carry = 0u64;
203        for j in 0..(8 - i) {
204            let res = rs1_u64[i] as u64 * rs2_u64[j] as u64 + rd[i + j] as u64 + carry;
205            rd[i + j] = res as u32;
206            carry = res >> 32;
207        }
208    }
209    u32_array_to_bytes(rd)
210}
211
212#[cfg(test)]
213mod tests {
214    use alloy_primitives::U256;
215    use rand::{prelude::StdRng, Rng, SeedableRng};
216
217    use crate::{common::u64_array_to_bytes, mult::u256_mul, INT256_NUM_LIMBS};
218
219    #[test]
220    fn test_u256_mul() {
221        let mut rng = StdRng::from_seed([42; 32]);
222        for _ in 0..10000 {
223            let limbs_a: [u64; 4] = rng.random();
224            let limbs_b: [u64; 4] = rng.random();
225            let a = U256::from_limbs(limbs_a);
226            let b = U256::from_limbs(limbs_b);
227            let a_u8: [u8; INT256_NUM_LIMBS] = u64_array_to_bytes(limbs_a);
228            let b_u8: [u8; INT256_NUM_LIMBS] = u64_array_to_bytes(limbs_b);
229            assert_eq!(U256::from_le_bytes(u256_mul(a_u8, b_u8)), a.wrapping_mul(b));
230        }
231    }
232}