feat: rv64 mul mulh divrem

crates/toolchain/instructions/src/riscv.rs

@@ -12,6 +12,8 @@ pub const RV32_NUM_REGISTERS: usize = 32;
 /// 64-bit register stored as 8 bytes (8 limbs of 8-bits) in OpenVM memory.
 pub const RV64_REGISTER_NUM_LIMBS: usize = 8;
 pub const RV64_CELL_BITS: usize = 8;
+/// 32-bit word stored as 4 bytes (4 limbs of 8-bits), i.e. half a 64-bit register.
+pub const RV64_WORD_NUM_LIMBS: usize = RV64_REGISTER_NUM_LIMBS / 2;
 
 pub const RV64_IMM_AS: u32 = 0;
 pub const RV64_REGISTER_AS: u32 = 1;

extensions/riscv/circuit/src/adapters/jalr.rs

@@ -1,4 +1,7 @@
-use std::borrow::{Borrow, BorrowMut};
+use std::{
+    borrow::{Borrow, BorrowMut},
+    mem::size_of,
+};
 
 use openvm_circuit::{
     arch::{
@@ -17,50 +20,50 @@ use openvm_circuit::{
 use openvm_circuit_primitives::{utils::not, AlignedBytesBorrow};
 use openvm_circuit_primitives_derive::AlignedBorrow;
 use openvm_instructions::{
-    instruction::Instruction, program::DEFAULT_PC_STEP, riscv::RV32_REGISTER_AS,
+    instruction::Instruction, program::DEFAULT_PC_STEP, riscv::RV64_REGISTER_AS,
 };
 use openvm_stark_backend::{
     interaction::InteractionBuilder,
     p3_air::{AirBuilder, BaseAir},
     p3_field::{Field, FieldAlgebra, PrimeField32},
 };
 
-use super::RV32_REGISTER_NUM_LIMBS;
+use super::RV64_REGISTER_NUM_LIMBS;
 use crate::adapters::{tracing_read, tracing_write};
 
 #[repr(C)]
 #[derive(Debug, Clone, AlignedBorrow)]
-pub struct Rv32JalrAdapterCols<T> {
+pub struct Rv64JalrAdapterCols<T> {
     pub from_state: ExecutionState<T>,
     pub rs1_ptr: T,
     pub rs1_aux_cols: MemoryReadAuxCols<T>,
     pub rd_ptr: T,
-    pub rd_aux_cols: MemoryWriteAuxCols<T, RV32_REGISTER_NUM_LIMBS>,
+    pub rd_aux_cols: MemoryWriteAuxCols<T, RV64_REGISTER_NUM_LIMBS>,
     /// Only writes if `needs_write`.
     /// Sets `needs_write` to 0 iff `rd == x0`
     pub needs_write: T,
 }
 
 #[derive(Clone, Copy, Debug, derive_new::new)]
-pub struct Rv32JalrAdapterAir {
+pub struct Rv64JalrAdapterAir {
     pub(super) memory_bridge: MemoryBridge,
     pub(super) execution_bridge: ExecutionBridge,
 }
 
-impl<F: Field> BaseAir<F> for Rv32JalrAdapterAir {
+impl<F: Field> BaseAir<F> for Rv64JalrAdapterAir {
     fn width(&self) -> usize {
-        Rv32JalrAdapterCols::<F>::width()
+        Rv64JalrAdapterCols::<F>::width()
     }
 }
 
-impl<AB: InteractionBuilder> VmAdapterAir<AB> for Rv32JalrAdapterAir {
+impl<AB: InteractionBuilder> VmAdapterAir<AB> for Rv64JalrAdapterAir {
     type Interface = BasicAdapterInterface<
         AB::Expr,
         SignedImmInstruction<AB::Expr>,
         1,
         1,
-        RV32_REGISTER_NUM_LIMBS,
-        RV32_REGISTER_NUM_LIMBS,
+        RV64_REGISTER_NUM_LIMBS,
+        RV64_REGISTER_NUM_LIMBS,
     >;
 
     fn eval(
@@ -69,7 +72,7 @@ impl<AB: InteractionBuilder> VmAdapterAir<AB> for Rv32JalrAdapterAir {
         local: &[AB::Var],
         ctx: AdapterAirContext<AB::Expr, Self::Interface>,
     ) {
-        let local_cols: &Rv32JalrAdapterCols<AB::Var> = local.borrow();
+        let local_cols: &Rv64JalrAdapterCols<AB::Var> = local.borrow();
 
         let timestamp: AB::Var = local_cols.from_state.timestamp;
         let mut timestamp_delta: usize = 0;
@@ -88,7 +91,7 @@ impl<AB: InteractionBuilder> VmAdapterAir<AB> for Rv32JalrAdapterAir {
         self.memory_bridge
             .read(
                 MemoryAddress::new(
-                    AB::F::from_canonical_u32(RV32_REGISTER_AS),
+                    AB::F::from_canonical_u32(RV64_REGISTER_AS),
                     local_cols.rs1_ptr,
                 ),
                 ctx.reads[0].clone(),
@@ -100,7 +103,7 @@ impl<AB: InteractionBuilder> VmAdapterAir<AB> for Rv32JalrAdapterAir {
         self.memory_bridge
             .write(
                 MemoryAddress::new(
-                    AB::F::from_canonical_u32(RV32_REGISTER_AS),
+                    AB::F::from_canonical_u32(RV64_REGISTER_AS),
                     local_cols.rd_ptr,
                 ),
                 ctx.writes[0].clone(),
@@ -121,7 +124,7 @@ impl<AB: InteractionBuilder> VmAdapterAir<AB> for Rv32JalrAdapterAir {
                     local_cols.rd_ptr.into(),
                     local_cols.rs1_ptr.into(),
                     ctx.instruction.immediate,
-                    AB::Expr::from_canonical_u32(RV32_REGISTER_AS),
+                    AB::Expr::from_canonical_u32(RV64_REGISTER_AS),
                     AB::Expr::ZERO,
                     write_count.into(),
                     ctx.instruction.imm_sign,
@@ -136,14 +139,14 @@ impl<AB: InteractionBuilder> VmAdapterAir<AB> for Rv32JalrAdapterAir {
     }
 
     fn get_from_pc(&self, local: &[AB::Var]) -> AB::Var {
-        let cols: &Rv32JalrAdapterCols<_> = local.borrow();
+        let cols: &Rv64JalrAdapterCols<_> = local.borrow();
         cols.from_state.pc
     }
 }
 
 #[repr(C)]
 #[derive(AlignedBytesBorrow, Debug)]
-pub struct Rv32JalrAdapterRecord {
+pub struct Rv64JalrAdapterRecord {
     pub from_pc: u32,
     pub from_timestamp: u32,
 
@@ -152,24 +155,24 @@ pub struct Rv32JalrAdapterRecord {
     pub rd_ptr: u32,
 
     pub reads_aux: MemoryReadAuxRecord,
-    pub writes_aux: MemoryWriteBytesAuxRecord<RV32_REGISTER_NUM_LIMBS>,
+    pub writes_aux: MemoryWriteBytesAuxRecord<RV64_REGISTER_NUM_LIMBS>,
 }
 
-// This adapter reads from [b:4]_d (rs1) and writes to [a:4]_d (rd)
+// This adapter reads from [b:8]_d (rs1) and writes to [a:8]_d (rd)
 #[derive(Clone, Copy, derive_new::new)]
-pub struct Rv32JalrAdapterExecutor;
+pub struct Rv64JalrAdapterExecutor;
 
 #[derive(Clone, Copy, derive_new::new)]
-pub struct Rv32JalrAdapterFiller;
+pub struct Rv64JalrAdapterFiller;
 
-impl<F> AdapterTraceExecutor<F> for Rv32JalrAdapterExecutor
+impl<F> AdapterTraceExecutor<F> for Rv64JalrAdapterExecutor
 where
     F: PrimeField32,
 {
-    const WIDTH: usize = size_of::<Rv32JalrAdapterCols<u8>>();
-    type ReadData = [u8; RV32_REGISTER_NUM_LIMBS];
-    type WriteData = [u8; RV32_REGISTER_NUM_LIMBS];
-    type RecordMut<'a> = &'a mut Rv32JalrAdapterRecord;
+    const WIDTH: usize = size_of::<Rv64JalrAdapterCols<u8>>();
+    type ReadData = [u8; RV64_REGISTER_NUM_LIMBS];
+    type WriteData = [u8; RV64_REGISTER_NUM_LIMBS];
+    type RecordMut<'a> = &'a mut Rv64JalrAdapterRecord;
 
     #[inline(always)]
     fn start(pc: u32, memory: &TracingMemory, record: &mut Self::RecordMut<'_>) {
@@ -186,12 +189,12 @@ where
     ) -> Self::ReadData {
         let &Instruction { b, d, .. } = instruction;
 
-        debug_assert_eq!(d.as_canonical_u32(), RV32_REGISTER_AS);
+        debug_assert_eq!(d.as_canonical_u32(), RV64_REGISTER_AS);
 
         record.rs1_ptr = b.as_canonical_u32();
         tracing_read(
             memory,
-            RV32_REGISTER_AS,
+            RV64_REGISTER_AS,
             b.as_canonical_u32(),
             &mut record.reads_aux.prev_timestamp,
         )
@@ -209,14 +212,14 @@ where
             a, d, f: enabled, ..
         } = instruction;
 
-        debug_assert_eq!(d.as_canonical_u32(), RV32_REGISTER_AS);
+        debug_assert_eq!(d.as_canonical_u32(), RV64_REGISTER_AS);
 
         if enabled.is_one() {
             record.rd_ptr = a.as_canonical_u32();
 
             tracing_write(
                 memory,
-                RV32_REGISTER_AS,
+                RV64_REGISTER_AS,
                 a.as_canonical_u32(),
                 data,
                 &mut record.writes_aux.prev_timestamp,
@@ -229,17 +232,17 @@ where
     }
 }
 
-impl<F: PrimeField32> AdapterTraceFiller<F> for Rv32JalrAdapterFiller {
-    const WIDTH: usize = size_of::<Rv32JalrAdapterCols<u8>>();
+impl<F: PrimeField32> AdapterTraceFiller<F> for Rv64JalrAdapterFiller {
+    const WIDTH: usize = size_of::<Rv64JalrAdapterCols<u8>>();
 
     #[inline(always)]
     fn fill_trace_row(&self, mem_helper: &MemoryAuxColsFactory<F>, mut adapter_row: &mut [F]) {
         // SAFETY:
         // - caller ensures `adapter_row` contains a valid record representation that was previously
         //   written by the executor
-        // - get_record_from_slice correctly interprets the bytes as Rv32JalrAdapterRecord
-        let record: &Rv32JalrAdapterRecord = unsafe { get_record_from_slice(&mut adapter_row, ()) };
-        let adapter_row: &mut Rv32JalrAdapterCols<F> = adapter_row.borrow_mut();
+        // - get_record_from_slice correctly interprets the bytes as Rv64JalrAdapterRecord
+        let record: &Rv64JalrAdapterRecord = unsafe { get_record_from_slice(&mut adapter_row, ()) };
+        let adapter_row: &mut Rv64JalrAdapterCols<F> = adapter_row.borrow_mut();
 
         // We must assign in reverse
         adapter_row.needs_write = F::from_bool(record.rd_ptr != u32::MAX);

extensions/riscv/circuit/src/adapters/mod.rs

@@ -24,6 +24,7 @@ pub use loadstore::*;
 pub use mul::*;
 pub use openvm_instructions::riscv::{
     RV32_CELL_BITS, RV32_REGISTER_NUM_LIMBS, RV64_CELL_BITS, RV64_REGISTER_NUM_LIMBS,
+    RV64_WORD_NUM_LIMBS,
 };
 pub use rdwrite::*;