openvm_rv32im_circuit/mul/
execution.rs

1use std::{
2    borrow::{Borrow, BorrowMut},
3    mem::size_of,
4};
5
6use openvm_circuit::{arch::*, system::memory::online::GuestMemory};
7use openvm_circuit_primitives_derive::AlignedBytesBorrow;
8use openvm_instructions::{
9    instruction::Instruction,
10    program::DEFAULT_PC_STEP,
11    riscv::{RV32_REGISTER_AS, RV32_REGISTER_NUM_LIMBS},
12    LocalOpcode,
13};
14use openvm_rv32im_transpiler::MulOpcode;
15use openvm_stark_backend::p3_field::PrimeField32;
16
17#[cfg(feature = "aot")]
18use crate::common::*;
19use crate::MultiplicationExecutor;
20
21#[derive(AlignedBytesBorrow, Clone)]
22#[repr(C)]
23struct MultiPreCompute {
24    a: u8,
25    b: u8,
26    c: u8,
27}
28
29impl<A, const LIMB_BITS: usize> MultiplicationExecutor<A, { RV32_REGISTER_NUM_LIMBS }, LIMB_BITS> {
30    fn pre_compute_impl<F: PrimeField32>(
31        &self,
32        pc: u32,
33        inst: &Instruction<F>,
34        data: &mut MultiPreCompute,
35    ) -> Result<(), StaticProgramError> {
36        if MulOpcode::from_usize(inst.opcode.local_opcode_idx(self.offset)) != MulOpcode::MUL {
37            return Err(StaticProgramError::InvalidInstruction(pc));
38        }
39        if inst.d.as_canonical_u32() != RV32_REGISTER_AS {
40            return Err(StaticProgramError::InvalidInstruction(pc));
41        }
42
43        *data = MultiPreCompute {
44            a: inst.a.as_canonical_u32() as u8,
45            b: inst.b.as_canonical_u32() as u8,
46            c: inst.c.as_canonical_u32() as u8,
47        };
48        Ok(())
49    }
50}
51
52impl<F, A, const LIMB_BITS: usize> InterpreterExecutor<F>
53    for MultiplicationExecutor<A, { RV32_REGISTER_NUM_LIMBS }, LIMB_BITS>
54where
55    F: PrimeField32,
56{
57    fn pre_compute_size(&self) -> usize {
58        size_of::<MultiPreCompute>()
59    }
60    #[cfg(not(feature = "tco"))]
61    fn pre_compute<Ctx>(
62        &self,
63        pc: u32,
64        inst: &Instruction<F>,
65        data: &mut [u8],
66    ) -> Result<ExecuteFunc<F, Ctx>, StaticProgramError>
67    where
68        Ctx: ExecutionCtxTrait,
69    {
70        let pre_compute: &mut MultiPreCompute = data.borrow_mut();
71        self.pre_compute_impl(pc, inst, pre_compute)?;
72        Ok(execute_e1_impl)
73    }
74
75    #[cfg(feature = "tco")]
76    fn handler<Ctx>(
77        &self,
78        pc: u32,
79        inst: &Instruction<F>,
80        data: &mut [u8],
81    ) -> Result<Handler<F, Ctx>, StaticProgramError>
82    where
83        Ctx: ExecutionCtxTrait,
84    {
85        let pre_compute: &mut MultiPreCompute = data.borrow_mut();
86        self.pre_compute_impl(pc, inst, pre_compute)?;
87        Ok(execute_e1_handler)
88    }
89}
90
91#[cfg(feature = "aot")]
92impl<F, A, const LIMB_BITS: usize> AotExecutor<F>
93    for MultiplicationExecutor<A, { RV32_REGISTER_NUM_LIMBS }, LIMB_BITS>
94where
95    F: PrimeField32,
96{
97    fn is_aot_supported(&self, inst: &Instruction<F>) -> bool {
98        inst.opcode == MulOpcode::MUL.global_opcode()
99    }
100
101    fn generate_x86_asm(&self, inst: &Instruction<F>, _pc: u32) -> Result<String, AotError> {
102        let to_i16 = |c: F| -> i16 {
103            let c_u24 = (c.as_canonical_u64() & 0xFFFFFF) as u32;
104            let c_i24 = ((c_u24 << 8) as i32) >> 8;
105            c_i24 as i16
106        };
107        let a = to_i16(inst.a);
108        let b = to_i16(inst.b);
109        let c = to_i16(inst.c);
110
111        if a % 4 != 0 || b % 4 != 0 || c % 4 != 0 {
112            return Err(AotError::InvalidInstruction);
113        }
114
115        let mut asm_str = String::new();
116
117        let str_reg_a = if RISCV_TO_X86_OVERRIDE_MAP[(a / 4) as usize].is_some() {
118            RISCV_TO_X86_OVERRIDE_MAP[(a / 4) as usize].unwrap()
119        } else {
120            REG_A_W
121        };
122
123        if a == c {
124            // a = b * c; commutative, so don't need to write to tmp, but should copy c to a first
125            let (gpr_reg_c, delta_str_c) = xmm_to_gpr((c / 4) as u8, str_reg_a, true);
126            asm_str += &delta_str_c;
127            let (gpr_reg_b, delta_str_b) = xmm_to_gpr((b / 4) as u8, REG_C_W, false);
128            asm_str += &delta_str_b;
129            asm_str += &format!("   imul {gpr_reg_c}, {gpr_reg_b}\n");
130            asm_str += &gpr_to_xmm(&gpr_reg_c, (a / 4) as u8);
131        } else {
132            let (gpr_reg_b, delta_str_b) = xmm_to_gpr((b / 4) as u8, str_reg_a, true);
133            asm_str += &delta_str_b; // data is now in gpr_reg_b
134            let (gpr_reg_c, delta_str_c) = xmm_to_gpr((c / 4) as u8, REG_C_W, false); // data is in gpr_reg_c now
135            asm_str += &delta_str_c; // have to get a return value here, since it modifies further registers too
136            asm_str += &format!("   imul {gpr_reg_b}, {gpr_reg_c}\n");
137            asm_str += &gpr_to_xmm(&gpr_reg_b, (a / 4) as u8);
138        }
139
140        Ok(asm_str)
141    }
142}
143
144impl<F, A, const LIMB_BITS: usize> InterpreterMeteredExecutor<F>
145    for MultiplicationExecutor<A, { RV32_REGISTER_NUM_LIMBS }, LIMB_BITS>
146where
147    F: PrimeField32,
148{
149    fn metered_pre_compute_size(&self) -> usize {
150        size_of::<E2PreCompute<MultiPreCompute>>()
151    }
152
153    #[cfg(not(feature = "tco"))]
154    fn metered_pre_compute<Ctx>(
155        &self,
156        chip_idx: usize,
157        pc: u32,
158        inst: &Instruction<F>,
159        data: &mut [u8],
160    ) -> Result<ExecuteFunc<F, Ctx>, StaticProgramError>
161    where
162        Ctx: MeteredExecutionCtxTrait,
163    {
164        let pre_compute: &mut E2PreCompute<MultiPreCompute> = data.borrow_mut();
165        pre_compute.chip_idx = chip_idx as u32;
166        self.pre_compute_impl(pc, inst, &mut pre_compute.data)?;
167        Ok(execute_e2_impl)
168    }
169
170    #[cfg(feature = "tco")]
171    fn metered_handler<Ctx>(
172        &self,
173        chip_idx: usize,
174        pc: u32,
175        inst: &Instruction<F>,
176        data: &mut [u8],
177    ) -> Result<Handler<F, Ctx>, StaticProgramError>
178    where
179        Ctx: MeteredExecutionCtxTrait,
180    {
181        let pre_compute: &mut E2PreCompute<MultiPreCompute> = data.borrow_mut();
182        pre_compute.chip_idx = chip_idx as u32;
183        self.pre_compute_impl(pc, inst, &mut pre_compute.data)?;
184        Ok(execute_e2_handler)
185    }
186}
187
188#[cfg(feature = "aot")]
189impl<F, A, const LIMB_BITS: usize> AotMeteredExecutor<F>
190    for MultiplicationExecutor<A, { RV32_REGISTER_NUM_LIMBS }, LIMB_BITS>
191where
192    F: PrimeField32,
193{
194    fn is_aot_metered_supported(&self, _inst: &Instruction<F>) -> bool {
195        true
196    }
197    fn generate_x86_metered_asm(
198        &self,
199        inst: &Instruction<F>,
200        pc: u32,
201        chip_idx: usize,
202        _config: &SystemConfig,
203    ) -> Result<String, AotError> {
204        let mut asm_str = self.generate_x86_asm(inst, pc)?;
205
206        asm_str += &update_height_change_asm(chip_idx, 1)?;
207
208        Ok(asm_str)
209    }
210}
211#[inline(always)]
212unsafe fn execute_e12_impl<F: PrimeField32, CTX: ExecutionCtxTrait>(
213    pre_compute: &MultiPreCompute,
214    exec_state: &mut VmExecState<F, GuestMemory, CTX>,
215) {
216    let rs1: [u8; RV32_REGISTER_NUM_LIMBS] =
217        exec_state.vm_read(RV32_REGISTER_AS, pre_compute.b as u32);
218    let rs2: [u8; RV32_REGISTER_NUM_LIMBS] =
219        exec_state.vm_read(RV32_REGISTER_AS, pre_compute.c as u32);
220    let rs1 = u32::from_le_bytes(rs1);
221    let rs2 = u32::from_le_bytes(rs2);
222    let rd = rs1.wrapping_mul(rs2);
223    exec_state.vm_write(RV32_REGISTER_AS, pre_compute.a as u32, &rd.to_le_bytes());
224
225    let pc = exec_state.pc();
226    exec_state.set_pc(pc.wrapping_add(DEFAULT_PC_STEP));
227}
228
229#[create_handler]
230#[inline(always)]
231unsafe fn execute_e1_impl<F: PrimeField32, CTX: ExecutionCtxTrait>(
232    pre_compute: *const u8,
233    exec_state: &mut VmExecState<F, GuestMemory, CTX>,
234) {
235    let pre_compute: &MultiPreCompute =
236        std::slice::from_raw_parts(pre_compute, size_of::<MultiPreCompute>()).borrow();
237    execute_e12_impl(pre_compute, exec_state);
238}
239
240#[create_handler]
241#[inline(always)]
242unsafe fn execute_e2_impl<F: PrimeField32, CTX: MeteredExecutionCtxTrait>(
243    pre_compute: *const u8,
244    exec_state: &mut VmExecState<F, GuestMemory, CTX>,
245) {
246    let pre_compute: &E2PreCompute<MultiPreCompute> =
247        std::slice::from_raw_parts(pre_compute, size_of::<E2PreCompute<MultiPreCompute>>())
248            .borrow();
249    exec_state
250        .ctx
251        .on_height_change(pre_compute.chip_idx as usize, 1);
252    execute_e12_impl(&pre_compute.data, exec_state);
253}