perf(v2.1): split rv64 load/store by operation and width by shuklaayush · Pull Request #2922 · openvm-org/openvm

shuklaayush · 2026-06-22T18:19:55Z

Summary

split RV64 memory accesses into separate load and store adapters, execution paths, AIRs, chips, and CUDA tracegen
split unsigned loads and stores by access width, with byte, halfword, word, and doubleword variants for each operation
split signed-load AIRs and chips into byte, halfword, and word sign-extension variants
keep width-specific columns, range checks, and memory interactions local to the chips that need them
register the split chips in the CPU and CUDA extension builders
update load, store, and signed-load tests to cover the split CPU and CUDA variants

Testing

cargo +nightly fmt --all
RUSTC_WRAPPER= cargo test -p openvm-riscv-circuit --no-run
OPENVM_SKIP_DEBUG=1 RUSTC_WRAPPER= cargo test -p openvm-riscv-circuit negative_split -- --nocapture
OPENVM_SKIP_DEBUG=1 RUSTC_WRAPPER= cargo test -p openvm-riscv-circuit rand_load -- --nocapture
OPENVM_SKIP_DEBUG=1 RUSTC_WRAPPER= cargo test -p openvm-riscv-circuit rand_store -- --nocapture
OPENVM_SKIP_DEBUG=1 RUSTC_WRAPPER= cargo test -p openvm-riscv-circuit rand_load_sign_extend -- --nocapture
CUDA load/store and signed-load tests

reth benchmark comparison

resolve int-7671, int-8104

github-actions · 2026-06-29T15:27:35Z

Code review

No issues found. Checked for bugs and CLAUDE.md compliance.

GunaDD · 2026-07-01T23:59:55Z

@@ -0,0 +1 @@
+pub mod core;


maybe we shouldn't put width_aligned in its own folder if its mod.rs is only one line

GunaDD · 2026-07-02T00:00:21Z

+    LOAD_SIGN_EXTEND_WORD_SELECTOR_WIDTH,
+>;
+pub type LoadSignExtendWordFiller = LoadSignExtendWidthAlignedFiller<
+    crate::adapters::Rv64LoadAdapterFiller,


maybe import can be put above

GunaDD · 2026-07-02T00:01:04Z

+        LOADW => KIND_WORD,
+        LOADH => KIND_HALFWORD,
+        LOADB => KIND_BYTE,


KIND doesn't sound right here

GunaDD · 2026-07-02T00:17:27Z

 /// register). Numerically equal to [`RV64_PTR_U16_LIMBS`], but named for arithmetic-word use.
 pub const RV64_WORD_U16_LIMBS: usize = RV64_WORD_NUM_LIMBS / 2;

+pub(crate) const RV64_ACCESS_SIZE_BYTE: usize = 0;


I think the naming of this constant doesn't actually represent what it means

GunaDD · 2026-07-02T00:18:49Z

+#[inline(always)]
+pub fn memory_read_u16<const N: usize>(
+    memory: &GuestMemory,
+    address_space: u32,
+    ptr: u32,
+) -> [u16; N] {
+    debug_assert!(
+        address_space == RV64_REGISTER_AS
+            || address_space == RV64_MEMORY_AS
+            || address_space == PUBLIC_VALUES_AS,
+    );
+
+    // SAFETY: these address spaces are u16-celled and `ptr` is an AS-native cell pointer.
+    unsafe { memory.read::<u16, N>(address_space, ptr) }
+}
+


this helper feels weird to me. I don't think we need a helper for this function

GunaDD · 2026-07-02T00:20:02Z

+#[inline(always)]
+pub(crate) fn u16_cell_byte(cell: u16, byte_idx: usize) -> u16 {
+    u16::from(cell.to_le_bytes()[byte_idx])
+}


better to not have this function for readability since this function is just one line

GunaDD · 2026-07-02T00:21:41Z

+pub(crate) const RV64_BYTE_MASK: u16 = (1 << RV64_BYTE_BITS) - 1;
+pub(crate) const RV64_BYTE_SIGN_BIT: u16 = 1 << (RV64_BYTE_BITS - 1);
+pub(crate) const RV64_U16_SIGN_BIT: u16 = 1 << (U16_BITS - 1);


maybe better to not have this constants ? either use another constant if its already there or just inline the definition of this constant. I feel like AI likes to define constants for everything

GunaDD · 2026-07-02T00:24:02Z

+}
+
+impl Rv64StoreAdapterRecord {
+    pub(crate) fn effective_ptr(&self) -> u32 {


I feel like there has to be a better name for this function.

GunaDD · 2026-07-02T00:28:33Z

+/// Wrapper type for u8 so typed VM memory reads and writes can use opcode-specific helpers.
+#[derive(Copy, Clone, Debug, Default)]
+struct U8(u8);


hmm what's with the U8 struct here

GunaDD · 2026-07-02T00:33:52Z

+pub(crate) fn u16_block_to_f_bytes(block: [u16; BLOCK_FE_WIDTH]) -> [F; 8] {
+    rv64_u16_block_to_bytes(block).map(F::from_u8)
+}


do we not have a function for this already

GunaDD · 2026-07-02T00:35:30Z

+    let a = rng.random_range(0..(max_addr - 8)) / 8 * 8;
+    let b = rng.random_range(0..(max_addr - 8)) / 8 * 8;


use the gen_pointer helper here

GunaDD · 2026-07-02T00:36:43Z

+pub(crate) fn b(bytes: [u8; 8]) -> [u16; BLOCK_FE_WIDTH] {
+    rv64_bytes_to_u16_block(bytes)
+}


I don't think we need a new function for this

GunaDD · 2026-07-02T00:55:22Z

+pub mod mul_cuda {
+
+    use super::*;
+


nit: remove the space here ?

GunaDD · 2026-07-02T00:56:13Z

-  - The instruction is correctly fetched from the program ROM at address `from_pc` and the program counter is set to `from_pc + 4`
+- A memory read from register `rs1` is performed
+- A memory read from register `rs2` is performed
+- A memory write to `mem_as` is performed at address `val(rs1) + imm`


maybe specify mem_as is only as 2 or 3 here

GunaDD · 2026-07-02T13:22:34Z

+static_assert(sizeof(Rv64LoadAdapterRecord) == 44);
+static_assert(sizeof(LoadRecord) == 8);
+static_assert(sizeof(Rv64LoadRecord) == 52);
+static_assert(offsetof(LoadRecord, read_data) == 0);
+static_assert(offsetof(Rv64LoadRecord, core) == 44);


I don't think we have this static assert in the other cuda tracegens, should we remove this for consistency?

GunaDD · 2026-07-02T13:27:08Z

+enum Rv64LoadSignExtendOpcode {
+    LOADB = 8,
+    LOADH = 9,
+    LOADW = 10,
+};


seems like a dead code ? same in the load.cuh file about the opcode enum

GunaDD · 2026-07-02T13:28:56Z

+enum Rv64StoreOpcode {
+    STORED = 4,
+    STOREW = 5,
+    STOREH = 6,
+    STOREB = 7,
+};


seems to be dead code

GunaDD · 2026-07-02T13:32:15Z

+    }
+};
+
+__global__ void rv64_load_byte_tracegen_kernel(


I don't think we usually call it _kernel. We usually name it rv64_load_byte_tracegen and _rv64_load_byte_tracegen

GunaDD · 2026-07-02T13:32:55Z

+    }
+};
+
+__global__ void rv64_load_doubleword_tracegen_kernel(


same here and all other cuda tracegen files in this PR

GunaDD · 2026-07-02T13:34:20Z

+template <typename T> struct LoadByteCoreCols {
+    T selector[LOAD_BYTE_SELECTOR_WIDTH];
+    T is_valid;
+    T read_cell_bytes[2];
+    T read_data[BLOCK_FE_WIDTH];
+};
+
+template <typename T> struct Rv64LoadByteCols {
+    Rv64LoadAdapterCols<T> adapter;
+    LoadByteCoreCols<T> core;
+};


should we move the definition of the columns in the load_byte.cuh files ? all the other .cu file doesn't have the definition of the columns. one idea would be to split the header .cuh file for consistency with other chips ?

GunaDD · 2026-07-02T13:44:44Z

+    mem_config.addr_spaces[RV64_REGISTER_AS as usize].num_cells = 1 << 29;
+    mem_config.addr_spaces[PUBLIC_VALUES_AS as usize].num_cells = 1 << 29;


why did it change the num_cells here ? and doesn't this exceed the currently allowed num_cells ? it should be at most 1 << 28. I think it should just use VmChipTestBuilder::from_config(MemoryConfig::default()) ?

GunaDD · 2026-07-02T13:46:56Z

+        LOADW => 2,
+        LOADH => 1,
+        LOADB => 0,
+        _ => unreachable!("signed load test only supports LOADB/LOADH/LOADW"),


alignment_bit instead of alignment

GunaDD · 2026-07-02T13:49:50Z

+        store_adapter_context::<AB, I>(
+            cols.is_valid.into(),
+            expected_opcode,
+            shift_amount,
+            cols.read_data,
+            cols.prev_data,
+            write_data,
+        )


wanted to note that this is inconsistent with the behavior of other chips which just returns the AirAdapterContext without using a helper function

GunaDD · 2026-07-02T13:51:44Z

+fn access_cells<const KIND: usize>() -> usize {
+    match KIND {
+        KIND_DOUBLEWORD => 4,
+        KIND_WORD => 2,
+        KIND_HALFWORD => 1,
+        _ => unreachable!("unsupported width-aligned store kind"),
+    }
+}


the naming of this function is confusing

GunaDD · 2026-07-02T13:53:19Z

+    fn fill_dummy_trace_row(&self, row_slice: &mut [F]) {
+        let (adapter_row, _) = unsafe {
+            row_slice
+                .split_at_mut_unchecked(<Rv64StoreAdapterFiller as AdapterTraceFiller<F>>::WIDTH)
+        };
+        let adapter_row: &mut Rv64StoreAdapterCols<F> = adapter_row.borrow_mut();
+        adapter_row.mem_as = F::from_u32(2);
+    }


hmmm why did it need to write a custom fill_dummy_trace_row function for store ?

GunaDD · 2026-07-02T13:57:32Z

+        let d_records = tracing::info_span!("trace_gen.h2d_records")
+            .in_scope(|| records.to_device_on(device_ctx))
+            .unwrap();
+        let d_trace = DeviceMatrix::<F>::with_capacity_on(trace_height, trace_width, device_ctx);


its unusual to have this tracing in the generate_proving_ctx. should we remove this ?

GunaDD · 2026-07-02T14:00:44Z

+pub const WORD_STORE_CASES: usize = 2;
+pub const WORD_STORE_SELECTOR_WIDTH: usize = 1;
+


the constant name isn't good. maybe STOREW_CASES ?

but I also think the CASES constant name isn't good as well in the comment before

GunaDD · 2026-07-02T14:03:07Z

+#[test]
+fn run_storew_sanity_test() {
+    let read_data = b([138, 45, 202, 76, 131, 74, 186, 29]);
+    let prev_data = b([159, 213, 89, 34, 142, 67, 210, 88]);
+    assert_eq!(
+        store_write_data(STOREW, read_data, prev_data, 0),
+        b([138, 45, 202, 76, 142, 67, 210, 88])
+    );
+    assert_eq!(
+        store_write_data(STOREW, read_data, prev_data, 4),
+        b([159, 213, 89, 34, 138, 45, 202, 76])
+    );
+}


more tests maybe ?

GunaDD · 2026-07-02T14:04:36Z

+    }
+}
+
+pub(crate) fn store_kind_for_opcode(opcode: Rv64LoadStoreOpcode) -> usize {


this function name is confusing

GunaDD

left comments mostly on the code structure and naming of variables which is confusing. the logic of the chip itself seems correct to me!

shuklaayush force-pushed the codex/int-7671-loadstore-split-rv64 branch from f32022a to 0c58aa7 Compare June 24, 2026 07:47

This comment has been minimized.

Sign in to view

shuklaayush force-pushed the codex/int-7671-loadstore-split-rv64 branch 2 times, most recently from 100fb7f to a209580 Compare June 26, 2026 18:41

shuklaayush changed the title ~~perf(v2.1): split rv64 loadstore by cell width~~ perf(v2.1): split rv64 load/store by operation and width Jun 26, 2026

This comment has been minimized.

Sign in to view

shuklaayush added 11 commits June 28, 2026 16:08

perf(riscv): split rv64 loadstore by cell width

edc16eb

feat(riscv): add cuda tracegen for split loadstore

a80efaf

docs(riscv): split loadstore README sections

b549b8d

refactor(riscv): colocate loadstore cuda wrappers

c1fccdc

refactor(riscv): split loadstore tests by width

e308147

refactor(riscv): colocate loadstore width tests

2811e85

refactor(riscv): split loadstore cuda kernels

7e3cbd8

refactor(riscv): split loadstore chips by operation

db9b301

fix(riscv): import cuda test traits

c3a3662

fix(riscv): order memory test helpers for clippy

2917a1b

fix(deps): update quinn-proto audit advisory

9c4933c

shuklaayush force-pushed the codex/int-7671-loadstore-split-rv64 branch from 26202ea to 9c4933c Compare June 28, 2026 14:08

This comment has been minimized.

Sign in to view

refactor(riscv): clarify width-aligned memory cores

d3318be

shuklaayush marked this pull request as ready for review June 29, 2026 15:01

This comment has been minimized.

Sign in to view

refactor(riscv): split load store cuda kernels

037e349

GunaDD reviewed Jul 1, 2026

View reviewed changes

GunaDD reviewed Jul 2, 2026

View reviewed changes

Comment thread extensions/riscv/circuit/src/README.md

GunaDD reviewed Jul 2, 2026

View reviewed changes

		let a = rng.random_range(0..(max_addr - 8)) / 8 * 8;
		let b = rng.random_range(0..(max_addr - 8)) / 8 * 8;

		mem_config.addr_spaces[RV64_REGISTER_AS as usize].num_cells = 1 << 29;
		mem_config.addr_spaces[PUBLIC_VALUES_AS as usize].num_cells = 1 << 29;

		pub const WORD_STORE_CASES: usize = 2;
		pub const WORD_STORE_SELECTOR_WIDTH: usize = 1;

Uh oh!

Conversation

shuklaayush commented Jun 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Testing

Uh oh!

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

github-actions Bot commented Jun 29, 2026

Code review

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

shuklaayush commented Jun 22, 2026 •

edited

Loading