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}