From 78bf93f07afb2c13b9650fe6a74d1dcbb55dba41 Mon Sep 17 00:00:00 2001 From: Ayush Shukla Date: Sun, 21 Jun 2026 22:34:22 +0000 Subject: [PATCH 01/26] perf(riscv): split rv64 loadstore by cell width --- .../circuit/cuda/src/load_sign_extend.cu | 165 --- .../riscv/circuit/cuda/src/loadstore.cu | 221 ---- extensions/riscv/circuit/src/README.md | 8 +- .../riscv/circuit/src/adapters/loadstore.rs | 112 +- extensions/riscv/circuit/src/adapters/mod.rs | 16 + extensions/riscv/circuit/src/cuda_abi.rs | 88 -- .../riscv/circuit/src/extension/cuda.rs | 62 +- extensions/riscv/circuit/src/extension/mod.rs | 220 +++- .../src/load_sign_extend/aligned/core.rs | 287 +++++ .../src/load_sign_extend/aligned/mod.rs | 1 + .../circuit/src/load_sign_extend/byte/core.rs | 233 ++++ .../circuit/src/load_sign_extend/byte/mod.rs | 15 + .../circuit/src/load_sign_extend/core.rs | 450 ------- .../circuit/src/load_sign_extend/cuda.rs | 81 +- .../circuit/src/load_sign_extend/execution.rs | 344 ----- .../src/load_sign_extend/halfword/mod.rs | 28 + .../riscv/circuit/src/load_sign_extend/mod.rs | 31 +- .../circuit/src/load_sign_extend/tests.rs | 801 +++++------- .../circuit/src/load_sign_extend/word/mod.rs | 28 + .../circuit/src/loadstore/aligned/core.rs | 338 +++++ .../circuit/src/loadstore/aligned/mod.rs | 2 + .../riscv/circuit/src/loadstore/byte/core.rs | 399 ++++++ .../riscv/circuit/src/loadstore/byte/mod.rs | 13 + .../riscv/circuit/src/loadstore/common.rs | 200 +++ .../riscv/circuit/src/loadstore/core.rs | 578 --------- .../riscv/circuit/src/loadstore/cuda.rs | 94 +- .../circuit/src/loadstore/doubleword/core.rs | 134 ++ .../circuit/src/loadstore/doubleword/mod.rs | 15 + .../riscv/circuit/src/loadstore/execution.rs | 284 ++-- .../circuit/src/loadstore/halfword/mod.rs | 26 + extensions/riscv/circuit/src/loadstore/mod.rs | 24 +- .../riscv/circuit/src/loadstore/tests.rs | 1138 ++++++++--------- .../riscv/circuit/src/loadstore/word/mod.rs | 24 + 33 files changed, 3193 insertions(+), 3267 deletions(-) delete mode 100644 extensions/riscv/circuit/cuda/src/load_sign_extend.cu delete mode 100644 extensions/riscv/circuit/cuda/src/loadstore.cu create mode 100644 extensions/riscv/circuit/src/load_sign_extend/aligned/core.rs create mode 100644 extensions/riscv/circuit/src/load_sign_extend/aligned/mod.rs create mode 100644 extensions/riscv/circuit/src/load_sign_extend/byte/core.rs create mode 100644 extensions/riscv/circuit/src/load_sign_extend/byte/mod.rs delete mode 100644 extensions/riscv/circuit/src/load_sign_extend/core.rs delete mode 100644 extensions/riscv/circuit/src/load_sign_extend/execution.rs create mode 100644 extensions/riscv/circuit/src/load_sign_extend/halfword/mod.rs create mode 100644 extensions/riscv/circuit/src/load_sign_extend/word/mod.rs create mode 100644 extensions/riscv/circuit/src/loadstore/aligned/core.rs create mode 100644 extensions/riscv/circuit/src/loadstore/aligned/mod.rs create mode 100644 extensions/riscv/circuit/src/loadstore/byte/core.rs create mode 100644 extensions/riscv/circuit/src/loadstore/byte/mod.rs create mode 100644 extensions/riscv/circuit/src/loadstore/common.rs delete mode 100644 extensions/riscv/circuit/src/loadstore/core.rs create mode 100644 extensions/riscv/circuit/src/loadstore/doubleword/core.rs create mode 100644 extensions/riscv/circuit/src/loadstore/doubleword/mod.rs create mode 100644 extensions/riscv/circuit/src/loadstore/halfword/mod.rs create mode 100644 extensions/riscv/circuit/src/loadstore/word/mod.rs diff --git a/extensions/riscv/circuit/cuda/src/load_sign_extend.cu b/extensions/riscv/circuit/cuda/src/load_sign_extend.cu deleted file mode 100644 index 9ba1ffa3ba..0000000000 --- a/extensions/riscv/circuit/cuda/src/load_sign_extend.cu +++ /dev/null @@ -1,165 +0,0 @@ -#include "launcher.cuh" -#include "primitives/buffer_view.cuh" -#include "primitives/constants.h" -#include "primitives/histogram.cuh" -#include "primitives/trace_access.h" -#include "riscv/adapters/loadstore.cuh" - -using namespace riscv; -using namespace program; - -template struct LoadSignExtendCoreCols { - /// This chip treats each (opcode, inner_shift) pair as a different instruction - T opcode_loadb_flag0; - T opcode_loadb_flag1; - T opcode_loadb_flag2; - T opcode_loadb_flag3; - T opcode_loadh_flag0; - T opcode_loadh_flag2; - T opcode_loadw_flag; - - T shift_most_sig_bit; - // The bit that is extended to the remaining bits - T data_most_sig_bit; - - T shifted_read_data[NUM_CELLS]; - T prev_data[NUM_CELLS]; -}; - -template struct LoadSignExtendCoreRecord { - bool is_byte; - bool is_word; - uint8_t shift_amount; - uint8_t read_data[NUM_CELLS]; - uint8_t prev_data[NUM_CELLS]; -}; - -template struct LoadSignExtendCore { - VariableRangeChecker range_checker; - BitwiseOperationLookup bitwise_lookup; - - template using Cols = LoadSignExtendCoreCols; - - __device__ LoadSignExtendCore( - VariableRangeChecker range_checker, - BitwiseOperationLookup bitwise_lookup - ) - : range_checker(range_checker), bitwise_lookup(bitwise_lookup) {} - - __device__ void fill_trace_row(RowSlice row, LoadSignExtendCoreRecord record) { - uint8_t shift = record.shift_amount; - uint8_t shift_most_sig_bit = (shift >> 2) & 1; - uint8_t inner_shift = shift & 3; - uint8_t rotate = shift_most_sig_bit * (NUM_CELLS / 2); - - uint8_t shifted_read_data[NUM_CELLS]; -#pragma unroll - for (size_t i = 0; i < NUM_CELLS; i++) { - shifted_read_data[i] = record.read_data[(i + rotate) % NUM_CELLS]; - } - - uint8_t most_sig_limb; - if (record.is_byte) { - most_sig_limb = shifted_read_data[inner_shift]; - } else if (record.is_word) { - most_sig_limb = shifted_read_data[NUM_CELLS / 2 - 1]; - } else { - most_sig_limb = shifted_read_data[inner_shift + 1]; - } - - uint8_t most_sig_bit = most_sig_limb & (1u << (RV64_BYTE_BITS - 1)); - bool is_word = record.is_word; - bool is_half = !record.is_byte && !is_word; - - range_checker.add_count(most_sig_limb - most_sig_bit, RV64_BYTE_BITS - 1); -#pragma unroll - for (size_t i = 0; i < NUM_CELLS; i += 2) { - bitwise_lookup.add_range(shifted_read_data[i], shifted_read_data[i + 1]); - } - COL_WRITE_VALUE(row, Cols, opcode_loadb_flag0, record.is_byte && inner_shift == 0); - COL_WRITE_VALUE(row, Cols, opcode_loadb_flag1, record.is_byte && inner_shift == 1); - COL_WRITE_VALUE(row, Cols, opcode_loadb_flag2, record.is_byte && inner_shift == 2); - COL_WRITE_VALUE(row, Cols, opcode_loadb_flag3, record.is_byte && inner_shift == 3); - COL_WRITE_VALUE(row, Cols, opcode_loadh_flag0, is_half && inner_shift == 0); - COL_WRITE_VALUE(row, Cols, opcode_loadh_flag2, is_half && inner_shift == 2); - COL_WRITE_VALUE(row, Cols, opcode_loadw_flag, is_word); - - COL_WRITE_VALUE(row, Cols, data_most_sig_bit, most_sig_bit != 0); - COL_WRITE_VALUE(row, Cols, shift_most_sig_bit, shift_most_sig_bit == 1); - COL_WRITE_ARRAY(row, Cols, shifted_read_data, shifted_read_data); - COL_WRITE_ARRAY(row, Cols, prev_data, record.prev_data); - } -}; - -// [Adapter + Core] columns and record -template struct Rv64LoadSignExtendCols { - Rv64LoadStoreAdapterCols adapter; - LoadSignExtendCoreCols core; -}; - -struct Rv64LoadSignExtendRecord { - Rv64LoadStoreAdapterRecord adapter; - LoadSignExtendCoreRecord core; -}; - -__global__ void rv64_load_sign_extend_tracegen( - Fp *trace, - size_t height, - size_t width, - DeviceBufferConstView records, - size_t pointer_max_bits, - uint32_t *range_checker_ptr, - uint32_t range_checker_num_bins, - uint32_t *bitwise_lookup_ptr, - uint32_t timestamp_max_bits -) { - uint32_t idx = blockIdx.x * blockDim.x + threadIdx.x; - RowSlice row(trace + idx, height); - if (idx < records.len()) { - auto const &record = records[idx]; - - auto adapter = Rv64LoadStoreAdapter( - pointer_max_bits, - VariableRangeChecker(range_checker_ptr, range_checker_num_bins), - timestamp_max_bits - ); - adapter.fill_trace_row(row, record.adapter); - - auto core = LoadSignExtendCore( - VariableRangeChecker(range_checker_ptr, range_checker_num_bins), - BitwiseOperationLookup(bitwise_lookup_ptr) - ); - core.fill_trace_row(row.slice_from(COL_INDEX(Rv64LoadSignExtendCols, core)), record.core); - } else { - row.fill_zero(0, sizeof(Rv64LoadSignExtendCols)); - } -} - -extern "C" int _rv64_load_sign_extend_tracegen( - Fp *__restrict__ d_trace, - size_t height, - size_t width, - DeviceBufferConstView d_records, - size_t pointer_max_bits, - uint32_t *__restrict__ d_range_checker, - uint32_t range_checker_num_bins, - uint32_t *__restrict__ d_bitwise_lookup, - uint32_t timestamp_max_bits, - cudaStream_t stream -) { - assert(width == sizeof(Rv64LoadSignExtendCols)); - auto [grid, block] = kernel_launch_params(height, 512); - - rv64_load_sign_extend_tracegen<<>>( - d_trace, - height, - width, - d_records, - pointer_max_bits, - d_range_checker, - range_checker_num_bins, - d_bitwise_lookup, - timestamp_max_bits - ); - return CHECK_KERNEL(); -} diff --git a/extensions/riscv/circuit/cuda/src/loadstore.cu b/extensions/riscv/circuit/cuda/src/loadstore.cu deleted file mode 100644 index bf3ec53a57..0000000000 --- a/extensions/riscv/circuit/cuda/src/loadstore.cu +++ /dev/null @@ -1,221 +0,0 @@ -#include "launcher.cuh" -#include "primitives/buffer_view.cuh" -#include "primitives/constants.h" -#include "primitives/encoder.cuh" -#include "primitives/histogram.cuh" -#include "primitives/trace_access.h" -#include "riscv/adapters/loadstore.cuh" - -using namespace riscv; -using namespace program; - -constexpr uint32_t LOADSTORE_SELECTOR_CASES = 30; -constexpr uint32_t LOADSTORE_SELECTOR_MAX_DEGREE = 2; -constexpr size_t LOADSTORE_SELECTOR_WIDTH = 7; - -template struct LoadStoreCoreCols { - T selector[LOADSTORE_SELECTOR_WIDTH]; - /// we need to keep the degree of is_valid and is_load to 1 - T is_valid; - T is_load; - - T read_data[NUM_CELLS]; - T prev_data[NUM_CELLS]; - /// write_data will be constrained against read_data and prev_data - /// depending on the opcode and the shift amount - T write_data[NUM_CELLS]; -}; - -template struct LoadStoreCoreRecord { - uint8_t local_opcode; - uint8_t shift_amount; - uint8_t read_data[NUM_CELLS]; - uint8_t prev_data[NUM_CELLS]; -}; - -enum Rv64LoadStoreOpcode { - LOADD, - LOADBU, - LOADHU, - LOADWU, - STORED, - STOREW, - STOREH, - STOREB, -}; - -// Lookup table mapping (opcode, shift) -> InstructionCase index for the Encoder. -// Indexed as INSTRUCTION_CASE[opcode][shift]. Invalid entries are 0xFF. -// clang-format off -__device__ constexpr uint8_t INSTRUCTION_CASE[8][8] = { - // LOADD: shift=0 - { 0, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF }, - // LOADBU: shift=0..7 - { 7, 8, 9, 10, 11, 12, 13, 14 }, - // LOADHU: shift=0,2,4,6 - { 3, 0xFF, 4, 0xFF, 5, 0xFF, 6, 0xFF }, - // LOADWU: shift=0,4 - { 1, 0xFF, 0xFF, 0xFF, 2, 0xFF, 0xFF, 0xFF }, - // STORED: shift=0 - { 15, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF }, - // STOREW: shift=0,4 - { 16, 0xFF, 0xFF, 0xFF, 17, 0xFF, 0xFF, 0xFF }, - // STOREH: shift=0,2,4,6 - { 18, 0xFF, 19, 0xFF, 20, 0xFF, 21, 0xFF }, - // STOREB: shift=0..7 - { 22, 23, 24, 25, 26, 27, 28, 29 }, -}; -// clang-format on - -__device__ __forceinline__ uint32_t instruction_case_from_opcode_shift( - Rv64LoadStoreOpcode opcode, - uint8_t shift -) { - uint8_t idx = INSTRUCTION_CASE[opcode][shift]; - assert(idx != 0xFF); - return idx; -} - -__device__ constexpr uint32_t LOADSTORE_WIDTH[] = { - // LOADD, LOADBU, LOADHU, LOADWU, STORED, STOREW, STOREH, STOREB - 8, 1, 2, 4, 8, 4, 2, 1, -}; - -template -__device__ __forceinline__ void run_write_data( - uint8_t (&write_data)[NUM_CELLS], - Rv64LoadStoreOpcode opcode, - const uint8_t (&read_data)[NUM_CELLS], - const uint8_t (&prev_data)[NUM_CELLS], - uint8_t shift -) { - bool is_store = opcode >= STORED; - uint32_t width = LOADSTORE_WIDTH[opcode]; - - if (is_store) { -#pragma unroll - for (size_t i = 0; i < NUM_CELLS; i++) { - bool in_range = (i >= shift) && (i < shift + width); - write_data[i] = in_range ? read_data[i - shift] : prev_data[i]; - } - } else { -#pragma unroll - for (size_t i = 0; i < NUM_CELLS; i++) { - write_data[i] = (i < width) ? read_data[i + shift] : 0u; - } - } -} - -template struct LoadStoreCore { - BitwiseOperationLookup bitwise_lookup; - - __device__ LoadStoreCore(BitwiseOperationLookup bitwise_lookup) - : bitwise_lookup(bitwise_lookup) {} - - template using Cols = LoadStoreCoreCols; - - __device__ void fill_trace_row(RowSlice row, LoadStoreCoreRecord record) { - Rv64LoadStoreOpcode opcode = static_cast(record.local_opcode); - Encoder encoder( - LOADSTORE_SELECTOR_CASES, - LOADSTORE_SELECTOR_MAX_DEGREE, - true - ); - uint8_t shift = record.shift_amount; - uint8_t write_data[NUM_CELLS] = {0}; - - COL_WRITE_VALUE(row, Cols, is_valid, 1); - COL_WRITE_VALUE( - row, - Cols, - is_load, - (opcode == LOADD || opcode == LOADWU || opcode == LOADHU || opcode == LOADBU) - ); - encoder.write_flag_pt( - row.slice_from(COL_INDEX(Cols, selector)), - instruction_case_from_opcode_shift(opcode, shift) - ); - COL_WRITE_ARRAY(row, Cols, read_data, record.read_data); - COL_WRITE_ARRAY(row, Cols, prev_data, record.prev_data); -#pragma unroll - for (size_t i = 0; i < NUM_CELLS; i += 2) { - bitwise_lookup.add_range(record.read_data[i], record.read_data[i + 1]); - bitwise_lookup.add_range(record.prev_data[i], record.prev_data[i + 1]); - } - - run_write_data(write_data, opcode, record.read_data, record.prev_data, shift); - COL_WRITE_ARRAY(row, Cols, write_data, write_data); - } -}; - -// [Adapter + Core] columns and record -template struct Rv64LoadStoreCols { - Rv64LoadStoreAdapterCols adapter; - LoadStoreCoreCols core; -}; - -struct Rv64LoadStoreRecord { - Rv64LoadStoreAdapterRecord adapter; - LoadStoreCoreRecord core; -}; - -__global__ void rv64_load_store_tracegen( - Fp *trace, - size_t height, - size_t width, - DeviceBufferConstView records, - size_t pointer_max_bits, - uint32_t *range_checker_ptr, - uint32_t range_checker_num_bins, - uint32_t *bitwise_lookup_ptr, - uint32_t timestamp_max_bits -) { - uint32_t idx = blockIdx.x * blockDim.x + threadIdx.x; - RowSlice row(trace + idx, height); - if (idx < records.len()) { - auto const &record = records[idx]; - - auto adapter = Rv64LoadStoreAdapter( - pointer_max_bits, - VariableRangeChecker(range_checker_ptr, range_checker_num_bins), - timestamp_max_bits - ); - adapter.fill_trace_row(row, record.adapter); - - auto core = LoadStoreCore( - BitwiseOperationLookup(bitwise_lookup_ptr) - ); - core.fill_trace_row(row.slice_from(COL_INDEX(Rv64LoadStoreCols, core)), record.core); - } else { - row.fill_zero(0, sizeof(Rv64LoadStoreCols)); - } -} - -extern "C" int _rv64_load_store_tracegen( - Fp *d_trace, - size_t height, - size_t width, - DeviceBufferConstView d_records, - size_t pointer_max_bits, - uint32_t *d_range_checker, - uint32_t range_checker_num_bins, - uint32_t *d_bitwise_lookup, - uint32_t timestamp_max_bits, - cudaStream_t stream -) { - assert(width == sizeof(Rv64LoadStoreCols)); - auto [grid, block] = kernel_launch_params(height); - - rv64_load_store_tracegen<<>>( - d_trace, - height, - width, - d_records, - pointer_max_bits, - d_range_checker, - range_checker_num_bins, - d_bitwise_lookup, - timestamp_max_bits - ); - return CHECK_KERNEL(); -} diff --git a/extensions/riscv/circuit/src/README.md b/extensions/riscv/circuit/src/README.md index 33775c6ea8..ec15b5fde0 100644 --- a/extensions/riscv/circuit/src/README.md +++ b/extensions/riscv/circuit/src/README.md @@ -294,15 +294,17 @@ This circuit proves that: - If `opcode` is `sltu` and `compose(b) < compose(c)` (unsigned comparison), then `a` is 1. - Otherwise, `a` is 0. -#### 10. [Load sign extend](./load_sign_extend/core.rs) and [Loadstore](./loadstore/core.rs) +#### 10. [Load/store](./loadstore/mod.rs) and [Load sign extend](./load_sign_extend/mod.rs) + +The RV64 load/store circuit is split by access width: [byte](./loadstore/byte/core.rs), [halfword](./loadstore/halfword/mod.rs), [word](./loadstore/word/mod.rs), and [doubleword](./loadstore/doubleword/core.rs). Signed loads stay in separate sign-extension chips split by width: [byte](./load_sign_extend/byte/core.rs), [halfword](./load_sign_extend/halfword/mod.rs), and [word](./load_sign_extend/word/mod.rs). Given: - `read_data` is the data read from `mem_as[aligned(val(rs1) + imm)]` if the instruction is load, otherwise it is the data read from register `rd` -- `write_data` is the data to be written to register `rd` if the instruction is load, otherwise it is the data to be written to `mem_as[aligned(val(rs1) + imm)]` +- `prev_data` is the previous write target block, either register `rd` for loads or `mem_as[aligned(val(rs1) + imm)]` for stores - `opcode` indicates the operation to be performed -This circuit proves that `write_data` equals `shift(read_data)`, where the shift amount is adjusted according to the instruction. +These circuits prove that the block written by the adapter matches the width-specific load extension or store merge for the selected opcode and alignment. #### 11. [Multiplication](./mul/core.rs) diff --git a/extensions/riscv/circuit/src/adapters/loadstore.rs b/extensions/riscv/circuit/src/adapters/loadstore.rs index 7deea67f61..d07ea5853f 100644 --- a/extensions/riscv/circuit/src/adapters/loadstore.rs +++ b/extensions/riscv/circuit/src/adapters/loadstore.rs @@ -8,11 +8,10 @@ use openvm_circuit::{ arch::{ get_record_from_slice, AdapterAirContext, AdapterTraceExecutor, AdapterTraceFiller, ExecutionBridge, ExecutionState, VmAdapterAir, VmAdapterInterface, BLOCK_FE_WIDTH, - MEMORY_BLOCK_BYTES, }, system::memory::{ offline_checker::{ - pack_u8_block, MemoryBaseAuxCols, MemoryBridge, MemoryReadAuxCols, MemoryReadAuxRecord, + MemoryBaseAuxCols, MemoryBridge, MemoryReadAuxCols, MemoryReadAuxRecord, MemoryWriteAuxInput, }, online::TracingMemory, @@ -39,11 +38,12 @@ use openvm_stark_backend::{ }; use super::{ - byte_ptr_to_u16_ptr, expand_to_rv64_block, ptr_to_field_u16_limbs, ptr_to_u16_limbs, - rv64_address_add_imm, sign_extend_imm16, try_rv64_bytes_to_u32, RV64_PTR_BITS, - RV64_PTR_U16_LIMBS, RV64_REGISTER_NUM_LIMBS, U16_BITS, + byte_ptr_to_u16_ptr, byte_ptr_to_u16_ptr_value, expand_to_rv64_block, ptr_to_field_u16_limbs, + ptr_to_u16_limbs, sign_extend_imm16, RV64_PTR_U16_LIMBS, RV64_REGISTER_NUM_LIMBS, U16_BITS, +}; +use crate::adapters::{ + memory_read_u16, rv64_bytes_to_u32, timed_write_u16, tracing_read, tracing_read_u16, }; -use crate::adapters::{memory_read, timed_write, tracing_read}; /// LoadStore Adapter handles all memory and register operations, so it must be aware /// of the instruction type, specifically whether it is a load or store. @@ -73,11 +73,8 @@ pub struct Rv64LoadStoreAdapterAirInterface(PhantomData< /// Using AB::Var for prev_data and AB::Expr for read_data impl VmAdapterInterface for Rv64LoadStoreAdapterAirInterface { - type Reads = ( - [AB::Var; RV64_REGISTER_NUM_LIMBS], - [AB::Expr; RV64_REGISTER_NUM_LIMBS], - ); - type Writes = [[AB::Expr; RV64_REGISTER_NUM_LIMBS]; 1]; + type Reads = ([AB::Var; BLOCK_FE_WIDTH], [AB::Expr; BLOCK_FE_WIDTH]); + type Writes = [[AB::Expr; BLOCK_FE_WIDTH]; 1]; type ProcessedInstruction = LoadStoreInstruction; } @@ -192,9 +189,6 @@ impl VmAdapterAir for Rv64LoadStoreAdapterAir { let imm_extend_limb = local_cols.imm_sign * AB::F::from_u32(u16::MAX as u32); let carry = (limbs_23 + imm_extend_limb + carry - local_cols.mem_ptr_limbs[1]) * inv; builder.when(is_valid.clone()).assert_bool(carry.clone()); - builder - .when(is_valid.clone()) - .assert_eq(carry, local_cols.imm_sign); // preventing mem_ptr overflow self.range_bus @@ -215,12 +209,10 @@ impl VmAdapterAir for Rv64LoadStoreAdapterAir { let mem_ptr = local_cols.mem_ptr_limbs[0] + local_cols.mem_ptr_limbs[1] * AB::F::from_u32(1u32 << U16_BITS); - // Constrain loads to address space 2 and stores to address spaces 2 or 3. - let mem_as_minus_two = local_cols.mem_as - AB::Expr::TWO; let is_store = is_valid.clone() - is_load.clone(); - builder - .when(is_valid.clone()) - .assert_zero(mem_as_minus_two.clone() * (mem_as_minus_two.clone() - is_store)); + // constrain mem_as to be in {0, 1, 2} if the instruction is a load, + // and in {2, 3, 4} if the instruction is a store + builder.assert_tern(local_cols.mem_as - is_store * AB::Expr::TWO); builder .when(not::(is_valid.clone())) .assert_zero(local_cols.mem_as); @@ -243,7 +235,7 @@ impl VmAdapterAir for Rv64LoadStoreAdapterAir { self.memory_bridge .read( MemoryAddress::new(read_as, byte_ptr_to_u16_ptr::(read_ptr)), - pack_u8_block::(&ctx.reads.1), + ctx.reads.1, timestamp_pp(), &local_cols.read_data_aux, ) @@ -260,17 +252,14 @@ impl VmAdapterAir for Rv64LoadStoreAdapterAir { let write_ptr = select::(is_load.clone(), local_cols.rd_rs2_ptr, mem_ptr.clone()) - store_shift_amount; - // The core supplies the previous write bytes; this adapter stores only the base aux - // columns. - let prev_data_expr: [AB::Expr; MEMORY_BLOCK_BYTES] = ctx.reads.0.map(Into::into); self.memory_bridge .write( MemoryAddress::new(write_as, byte_ptr_to_u16_ptr::(write_ptr)), - pack_u8_block::(&ctx.writes[0].clone()), + ctx.writes[0].clone(), timestamp_pp(), MemoryWriteAuxInput::from_prev_data_exprs( &local_cols.write_base_aux, - pack_u8_block::(&prev_data_expr), + ctx.reads.0.map(Into::into), ), ) .eval(builder, write_count); @@ -344,11 +333,8 @@ where F: PrimeField32, { const WIDTH: usize = size_of::>(); - type ReadData = ( - ([u8; RV64_REGISTER_NUM_LIMBS], [u8; RV64_REGISTER_NUM_LIMBS]), - u8, - ); - type WriteData = [u8; RV64_REGISTER_NUM_LIMBS]; + type ReadData = (([u16; BLOCK_FE_WIDTH], [u16; BLOCK_FE_WIDTH]), u8); + type WriteData = [u16; BLOCK_FE_WIDTH]; type RecordMut<'a> = &'a mut Rv64LoadStoreAdapterRecord; #[inline(always)] @@ -381,57 +367,55 @@ where ); record.rs1_ptr = b.as_canonical_u32(); - let rs1_bytes = tracing_read( + record.rs1_val = rv64_bytes_to_u32(tracing_read( memory, RV64_REGISTER_AS, record.rs1_ptr, &mut record.rs1_aux_record.prev_timestamp, - ); - let rs1_val = try_rv64_bytes_to_u32(rs1_bytes).expect("upper 4 bytes must be zero"); - record.rs1_val = rs1_val; + )); record.imm = c.as_canonical_u32() as u16; record.imm_sign = g.is_one(); let imm_extended = sign_extend_imm16(record.imm as u32, record.imm_sign as u32); - let addr = rv64_address_add_imm(rs1_val, imm_extended); - let ptr_val = u32::try_from(addr) - .ok() - .filter(|&ptr| { - self.pointer_max_bits >= RV64_PTR_BITS - || u64::from(ptr) < (1u64 << self.pointer_max_bits) - }) - .expect("effective address exceeds implemented memory address space"); + let ptr_val = record.rs1_val.wrapping_add(imm_extended); let shift_amount = ptr_val & (RV64_REGISTER_NUM_LIMBS as u32 - 1); - let aligned_ptr = ptr_val - shift_amount; + let ptr_val = ptr_val - shift_amount; + debug_assert!((ptr_val as u64) < (1u64 << self.pointer_max_bits)); + let ptr_cell = byte_ptr_to_u16_ptr_value(ptr_val); - // prev_data: We need to keep values of some cells to keep them unchanged when writing to - // those cells let (read_data, prev_data) = match local_opcode { - LOADD | LOADW | LOADB | LOADH | LOADBU | LOADHU | LOADWU => { + LOADD | LOADW | LOADH | LOADB | LOADWU | LOADHU | LOADBU => { debug_assert_eq!(e, F::from_u32(RV64_MEMORY_AS)); record.mem_as = RV64_MEMORY_AS as u8; - let read_data = tracing_read( + let read_data = tracing_read_u16( memory, RV64_MEMORY_AS, - aligned_ptr, + ptr_cell, &mut record.read_data_aux.prev_timestamp, ); - let prev_data = memory_read(memory.data(), RV64_REGISTER_AS, a.as_canonical_u32()); + let prev_data = memory_read_u16( + memory.data(), + RV64_REGISTER_AS, + byte_ptr_to_u16_ptr_value(a.as_canonical_u32()), + ); (read_data, prev_data) } STORED | STOREW | STOREH | STOREB => { let e = e.as_canonical_u32(); debug_assert_ne!(e, RV64_IMM_AS); debug_assert_ne!(e, RV64_REGISTER_AS); + if e == DEFERRAL_AS { + unreachable!("STORE to DEFERRAL_AS is unsupported"); + } record.mem_as = e as u8; - let read_data = tracing_read( + let read_data = tracing_read_u16( memory, RV64_REGISTER_AS, - a.as_canonical_u32(), + byte_ptr_to_u16_ptr_value(a.as_canonical_u32()), &mut record.read_data_aux.prev_timestamp, ); - let prev_data = memory_read(memory.data(), e, aligned_ptr); + let prev_data = memory_read_u16(memory.data(), e, ptr_cell); (read_data, prev_data) } }; @@ -466,20 +450,30 @@ where if enabled != F::ZERO { record.rd_rs2_ptr = a.as_canonical_u32(); - record.write_prev_timestamp = match local_opcode { STORED | STOREW | STOREH | STOREB => { let imm_extended = sign_extend_imm16(record.imm as u32, record.imm_sign as u32); let ptr = record.rs1_val.wrapping_add(imm_extended) & !(RV64_REGISTER_NUM_LIMBS as u32 - 1); if record.mem_as == DEFERRAL_AS as u8 { - // TODO: Remove loadstore read/write support for DEFERRAL_AS. unreachable!("STORE to DEFERRAL_AS is unsupported"); } - timed_write(memory, record.mem_as as u32, ptr, data).0 + timed_write_u16( + memory, + record.mem_as as u32, + byte_ptr_to_u16_ptr_value(ptr), + data, + ) + .0 } - LOADD | LOADW | LOADB | LOADH | LOADWU | LOADBU | LOADHU => { - timed_write(memory, RV64_REGISTER_AS, record.rd_rs2_ptr, data).0 + LOADD | LOADW | LOADH | LOADB | LOADWU | LOADHU | LOADBU => { + timed_write_u16( + memory, + RV64_REGISTER_AS, + byte_ptr_to_u16_ptr_value(record.rd_rs2_ptr), + data, + ) + .0 } }; } else { @@ -522,9 +516,11 @@ impl AdapterTraceFiller for Rv64LoadStoreAdapterFiller { let ptr = record .rs1_val .wrapping_add(sign_extend_imm16(record.imm as u32, record.imm_sign as u32)); + let ptr_limbs = ptr_to_u16_limbs(ptr).map(u32::from); + let shift_amount = ptr & (RV64_REGISTER_NUM_LIMBS as u32 - 1); self.range_checker_chip - .add_count(ptr_limbs[0] >> 3, U16_BITS - 3); + .add_count((ptr_limbs[0] - shift_amount) >> 3, U16_BITS - 3); self.range_checker_chip .add_count(ptr_limbs[1], self.pointer_max_bits - U16_BITS); adapter_row.mem_ptr_limbs = ptr_limbs.map(F::from_u32); diff --git a/extensions/riscv/circuit/src/adapters/mod.rs b/extensions/riscv/circuit/src/adapters/mod.rs index 035e294d80..cee099a78e 100644 --- a/extensions/riscv/circuit/src/adapters/mod.rs +++ b/extensions/riscv/circuit/src/adapters/mod.rs @@ -371,6 +371,22 @@ pub fn memory_read(memory: &GuestMemory, address_space: u32, ptr unsafe { memory.read_bytes::(address_space, ptr) } } +#[inline(always)] +pub fn memory_read_u16( + 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::(address_space, ptr) } +} + #[inline(always)] pub fn memory_write( memory: &mut GuestMemory, diff --git a/extensions/riscv/circuit/src/cuda_abi.rs b/extensions/riscv/circuit/src/cuda_abi.rs index 4636117d3f..3045c55915 100644 --- a/extensions/riscv/circuit/src/cuda_abi.rs +++ b/extensions/riscv/circuit/src/cuda_abi.rs @@ -430,94 +430,6 @@ pub mod bitwise_logic_cuda { } } -pub mod loadstore_cuda { - use super::*; - - extern "C" { - pub fn _rv64_load_store_tracegen( - d_trace: *mut F, - height: usize, - width: usize, - d_records: DeviceBufferView, - pointer_max_bits: usize, - d_range_checker: *mut u32, - range_checker_num_bins: u32, - d_bitwise_lookup: *mut u32, - timestamp_max_bits: u32, - stream: cudaStream_t, - ) -> i32; - } - - pub unsafe fn tracegen( - d_trace: &DeviceBuffer, - height: usize, - width: usize, - d_records: &DeviceBuffer, - pointer_max_bits: usize, - d_range_checker: &DeviceBuffer, - d_bitwise_lookup: &DeviceBuffer, - timestamp_max_bits: u32, - stream: cudaStream_t, - ) -> Result<(), CudaError> { - CudaError::from_result(_rv64_load_store_tracegen( - d_trace.as_mut_ptr(), - height, - width, - d_records.view(), - pointer_max_bits, - d_range_checker.as_mut_ptr() as *mut u32, - d_range_checker.len() as u32, - d_bitwise_lookup.as_mut_ptr() as *mut u32, - timestamp_max_bits, - stream, - )) - } -} - -pub mod load_sign_extend_cuda { - use super::*; - - extern "C" { - pub fn _rv64_load_sign_extend_tracegen( - d_trace: *mut F, - height: usize, - width: usize, - d_records: DeviceBufferView, - pointer_max_bits: usize, - d_range_checker: *mut u32, - range_checker_num_bins: u32, - d_bitwise_lookup: *mut u32, - timestamp_max_bits: u32, - stream: cudaStream_t, - ) -> i32; - } - - pub unsafe fn tracegen( - d_trace: &DeviceBuffer, - height: usize, - width: usize, - d_records: &DeviceBuffer, - pointer_max_bits: usize, - d_range_checker: &DeviceBuffer, - d_bitwise_lookup: &DeviceBuffer, - timestamp_max_bits: u32, - stream: cudaStream_t, - ) -> Result<(), CudaError> { - CudaError::from_result(_rv64_load_sign_extend_tracegen( - d_trace.as_mut_ptr(), - height, - width, - d_records.view(), - pointer_max_bits, - d_range_checker.as_mut_ptr() as *mut u32, - d_range_checker.len() as u32, - d_bitwise_lookup.as_mut_ptr() as *mut u32, - timestamp_max_bits, - stream, - )) - } -} - pub mod jal_lui_cuda { use super::*; diff --git a/extensions/riscv/circuit/src/extension/cuda.rs b/extensions/riscv/circuit/src/extension/cuda.rs index f0155254c8..ae478120c3 100644 --- a/extensions/riscv/circuit/src/extension/cuda.rs +++ b/extensions/riscv/circuit/src/extension/cuda.rs @@ -16,9 +16,13 @@ use crate::{ Rv64BranchEqualChipGpu, Rv64BranchLessThanAir, Rv64BranchLessThanChipGpu, Rv64DivRemAir, Rv64DivRemChipGpu, Rv64DivRemWAir, Rv64DivRemWChipGpu, Rv64HintStoreAir, Rv64HintStoreChipGpu, Rv64I, Rv64Io, Rv64JalLuiAir, Rv64JalLuiChipGpu, Rv64JalrAir, Rv64JalrChipGpu, Rv64LessThanAir, - Rv64LessThanChipGpu, Rv64LoadSignExtendAir, Rv64LoadSignExtendChipGpu, Rv64LoadStoreAir, - Rv64LoadStoreChipGpu, Rv64M, Rv64MulHAir, Rv64MulHChipGpu, Rv64MulWAir, Rv64MulWChipGpu, - Rv64MultiplicationAir, Rv64MultiplicationChipGpu, Rv64ShiftLogicalAir, Rv64ShiftLogicalChipGpu, + Rv64LessThanChipGpu, Rv64LoadSignExtendByteAir, Rv64LoadSignExtendByteChipGpu, + Rv64LoadSignExtendHalfwordAir, Rv64LoadSignExtendHalfwordChipGpu, Rv64LoadSignExtendWordAir, + Rv64LoadSignExtendWordChipGpu, Rv64LoadStoreByteAir, Rv64LoadStoreByteChipGpu, + Rv64LoadStoreDoublewordAir, Rv64LoadStoreDoublewordChipGpu, Rv64LoadStoreHalfwordAir, + Rv64LoadStoreHalfwordChipGpu, Rv64LoadStoreWordAir, Rv64LoadStoreWordChipGpu, Rv64M, + Rv64MulHAir, Rv64MulHChipGpu, Rv64MulWAir, Rv64MulWChipGpu, Rv64MultiplicationAir, + Rv64MultiplicationChipGpu, Rv64ShiftLogicalAir, Rv64ShiftLogicalChipGpu, Rv64ShiftRightArithmeticAir, Rv64ShiftRightArithmeticChipGpu, Rv64ShiftWLogicalAir, Rv64ShiftWLogicalChipGpu, Rv64ShiftWRightArithmeticAir, Rv64ShiftWRightArithmeticChipGpu, }; @@ -80,23 +84,63 @@ impl VmProverExtension for Rv64ShiftWRightArithmeticChipGpu::new(range_checker.clone(), timestamp_max_bits); inventory.add_executor_chip(shift_w_right_arithmetic); - inventory.next_air::()?; - let load_store_chip = Rv64LoadStoreChipGpu::new( + inventory.next_air::()?; + let load_sign_extend_byte = Rv64LoadSignExtendByteChipGpu::new( range_checker.clone(), bitwise_lu.clone(), byte_ptr_max_bits, timestamp_max_bits, ); - inventory.add_executor_chip(load_store_chip); + inventory.add_executor_chip(load_sign_extend_byte); - inventory.next_air::()?; - let load_sign_extend = Rv64LoadSignExtendChipGpu::new( + inventory.next_air::()?; + let load_store_byte = Rv64LoadStoreByteChipGpu::new( range_checker.clone(), bitwise_lu.clone(), byte_ptr_max_bits, timestamp_max_bits, ); - inventory.add_executor_chip(load_sign_extend); + inventory.add_executor_chip(load_store_byte); + + inventory.next_air::()?; + let load_sign_extend_halfword = Rv64LoadSignExtendHalfwordChipGpu::new( + range_checker.clone(), + byte_ptr_max_bits, + timestamp_max_bits, + ); + inventory.add_executor_chip(load_sign_extend_halfword); + + inventory.next_air::()?; + let load_store_halfword = Rv64LoadStoreHalfwordChipGpu::new( + range_checker.clone(), + byte_ptr_max_bits, + timestamp_max_bits, + ); + inventory.add_executor_chip(load_store_halfword); + + inventory.next_air::()?; + let load_sign_extend_word = Rv64LoadSignExtendWordChipGpu::new( + range_checker.clone(), + byte_ptr_max_bits, + timestamp_max_bits, + ); + inventory.add_executor_chip(load_sign_extend_word); + + inventory.next_air::()?; + let load_store_word = Rv64LoadStoreWordChipGpu::new( + range_checker.clone(), + byte_ptr_max_bits, + timestamp_max_bits, + ); + inventory.add_executor_chip(load_store_word); + + inventory.next_air::()?; + let load_store_doubleword = Rv64LoadStoreDoublewordChipGpu::new( + range_checker.clone(), + byte_ptr_max_bits, + timestamp_max_bits, + ); + inventory.add_executor_chip(load_store_doubleword); inventory.next_air::()?; let beq = Rv64BranchEqualChipGpu::new(range_checker.clone(), timestamp_max_bits); diff --git a/extensions/riscv/circuit/src/extension/mod.rs b/extensions/riscv/circuit/src/extension/mod.rs index b3275e0b62..8c179e3088 100644 --- a/extensions/riscv/circuit/src/extension/mod.rs +++ b/extensions/riscv/circuit/src/extension/mod.rs @@ -128,8 +128,10 @@ pub enum Rv64IExecutor { JalLui(Rv64JalLuiExecutor), Jalr(Rv64JalrExecutor), Auipc(Rv64AuipcExecutor), - LoadStore(Rv64LoadStoreExecutor), - LoadSignExtend(Rv64LoadSignExtendExecutor), + LoadStoreByte(Rv64LoadStoreByteExecutor), + LoadStoreHalfword(Rv64LoadStoreHalfwordExecutor), + LoadStoreWord(Rv64LoadStoreWordExecutor), + LoadStoreDoubleword(Rv64LoadStoreDoublewordExecutor), } /// RISC-V 64-bit Multiplication Extension (RV64M) Instruction Executors @@ -222,27 +224,67 @@ impl VmExecutionExtension for Rv64I { [ShiftWOpcode::SRAW].map(|x| x.global_opcode()), )?; - let load_store = Rv64LoadStoreExecutor::new( + let load_sign_extend_byte = Rv64LoadSignExtendByteExecutor::new( Rv64LoadStoreAdapterExecutor::new(byte_ptr_max_bits), Rv64LoadStoreOpcode::CLASS_OFFSET, ); inventory.add_executor( - load_store, - Rv64LoadStoreOpcode::iter() - .take(Rv64LoadStoreOpcode::STOREB as usize + 1) - .map(|x| x.global_opcode()), + load_sign_extend_byte, + [Rv64LoadStoreOpcode::LOADB].map(|x| x.global_opcode()), )?; - let load_sign_extend = - Rv64LoadSignExtendExecutor::new(Rv64LoadStoreAdapterExecutor::new(byte_ptr_max_bits)); + let load_store_byte = Rv64LoadStoreByteExecutor::new( + Rv64LoadStoreAdapterExecutor::new(byte_ptr_max_bits), + Rv64LoadStoreOpcode::CLASS_OFFSET, + ); + inventory.add_executor( + load_store_byte, + [Rv64LoadStoreOpcode::LOADBU, Rv64LoadStoreOpcode::STOREB].map(|x| x.global_opcode()), + )?; + + let load_sign_extend_halfword = Rv64LoadSignExtendHalfwordExecutor::new( + Rv64LoadStoreAdapterExecutor::new(byte_ptr_max_bits), + Rv64LoadStoreOpcode::CLASS_OFFSET, + ); + inventory.add_executor( + load_sign_extend_halfword, + [Rv64LoadStoreOpcode::LOADH].map(|x| x.global_opcode()), + )?; + + let load_store_halfword = Rv64LoadStoreHalfwordExecutor::new( + Rv64LoadStoreAdapterExecutor::new(byte_ptr_max_bits), + Rv64LoadStoreOpcode::CLASS_OFFSET, + ); + inventory.add_executor( + load_store_halfword, + [Rv64LoadStoreOpcode::LOADHU, Rv64LoadStoreOpcode::STOREH].map(|x| x.global_opcode()), + )?; + + let load_sign_extend_word = Rv64LoadSignExtendWordExecutor::new( + Rv64LoadStoreAdapterExecutor::new(byte_ptr_max_bits), + Rv64LoadStoreOpcode::CLASS_OFFSET, + ); + inventory.add_executor( + load_sign_extend_word, + [Rv64LoadStoreOpcode::LOADW].map(|x| x.global_opcode()), + )?; + + let load_store_word = Rv64LoadStoreWordExecutor::new( + Rv64LoadStoreAdapterExecutor::new(byte_ptr_max_bits), + Rv64LoadStoreOpcode::CLASS_OFFSET, + ); + inventory.add_executor( + load_store_word, + [Rv64LoadStoreOpcode::LOADWU, Rv64LoadStoreOpcode::STOREW].map(|x| x.global_opcode()), + )?; + + let load_store_doubleword = Rv64LoadStoreDoublewordExecutor::new( + Rv64LoadStoreAdapterExecutor::new(byte_ptr_max_bits), + Rv64LoadStoreOpcode::CLASS_OFFSET, + ); inventory.add_executor( - load_sign_extend, - [ - Rv64LoadStoreOpcode::LOADB, - Rv64LoadStoreOpcode::LOADH, - Rv64LoadStoreOpcode::LOADW, - ] - .map(|x| x.global_opcode()), + load_store_doubleword, + [Rv64LoadStoreOpcode::LOADD, Rv64LoadStoreOpcode::STORED].map(|x| x.global_opcode()), )?; let beq = BranchEqualExecutor::new( @@ -363,27 +405,86 @@ impl VmCircuitExtension for Rv64I { ); inventory.add_air(shift_w_right_arithmetic); - let load_store = Rv64LoadStoreAir::new( + let load_sign_extend_byte = Rv64LoadSignExtendByteAir::new( + Rv64LoadStoreAdapterAir::new( + memory_bridge, + exec_bridge, + range_checker, + byte_ptr_max_bits, + ), + LoadSignExtendByteCoreAir::new( + Rv64LoadStoreOpcode::CLASS_OFFSET, + bitwise_lu, + range_checker, + ), + ); + inventory.add_air(load_sign_extend_byte); + + let load_store_byte = Rv64LoadStoreByteAir::new( + Rv64LoadStoreAdapterAir::new( + memory_bridge, + exec_bridge, + range_checker, + byte_ptr_max_bits, + ), + LoadStoreByteCoreAir::new(Rv64LoadStoreOpcode::CLASS_OFFSET, bitwise_lu, range_checker), + ); + inventory.add_air(load_store_byte); + + let load_sign_extend_halfword = Rv64LoadSignExtendHalfwordAir::new( + Rv64LoadStoreAdapterAir::new( + memory_bridge, + exec_bridge, + range_checker, + byte_ptr_max_bits, + ), + LoadSignExtendHalfwordCoreAir::new(Rv64LoadStoreOpcode::CLASS_OFFSET, range_checker), + ); + inventory.add_air(load_sign_extend_halfword); + + let load_store_halfword = Rv64LoadStoreHalfwordAir::new( Rv64LoadStoreAdapterAir::new( memory_bridge, exec_bridge, range_checker, byte_ptr_max_bits, ), - LoadStoreCoreAir::new(Rv64LoadStoreOpcode::CLASS_OFFSET, bitwise_lu), + LoadStoreHalfwordCoreAir::new(Rv64LoadStoreOpcode::CLASS_OFFSET, range_checker), ); - inventory.add_air(load_store); + inventory.add_air(load_store_halfword); - let load_sign_extend = Rv64LoadSignExtendAir::new( + let load_sign_extend_word = Rv64LoadSignExtendWordAir::new( Rv64LoadStoreAdapterAir::new( memory_bridge, exec_bridge, range_checker, byte_ptr_max_bits, ), - LoadSignExtendCoreAir::new(range_checker, bitwise_lu), + LoadSignExtendWordCoreAir::new(Rv64LoadStoreOpcode::CLASS_OFFSET, range_checker), ); - inventory.add_air(load_sign_extend); + inventory.add_air(load_sign_extend_word); + + let load_store_word = Rv64LoadStoreWordAir::new( + Rv64LoadStoreAdapterAir::new( + memory_bridge, + exec_bridge, + range_checker, + byte_ptr_max_bits, + ), + LoadStoreWordCoreAir::new(Rv64LoadStoreOpcode::CLASS_OFFSET, range_checker), + ); + inventory.add_air(load_store_word); + + let load_store_doubleword = Rv64LoadStoreDoublewordAir::new( + Rv64LoadStoreAdapterAir::new( + memory_bridge, + exec_bridge, + range_checker, + byte_ptr_max_bits, + ), + LoadStoreDoublewordCoreAir::new(Rv64LoadStoreOpcode::CLASS_OFFSET, range_checker), + ); + inventory.add_air(load_store_doubleword); let beq = Rv64BranchEqualAir::new( Rv64BranchAdapterAir::new(exec_bridge, memory_bridge), @@ -544,27 +645,84 @@ where ); inventory.add_executor_chip(shift_w_right_arithmetic); - inventory.next_air::()?; - let load_store_chip = Rv64LoadStoreChip::new( - LoadStoreFiller::new( + inventory.next_air::()?; + let load_sign_extend_byte_chip = Rv64LoadSignExtendByteChip::new( + LoadSignExtendByteFiller::new( Rv64LoadStoreAdapterFiller::new(byte_ptr_max_bits, range_checker.clone()), Rv64LoadStoreOpcode::CLASS_OFFSET, bitwise_lu.clone(), + range_checker.clone(), ), mem_helper.clone(), ); - inventory.add_executor_chip(load_store_chip); + inventory.add_executor_chip(load_sign_extend_byte_chip); - inventory.next_air::()?; - let load_sign_extend = Rv64LoadSignExtendChip::new( - LoadSignExtendFiller::new( + inventory.next_air::()?; + let load_store_byte_chip = Rv64LoadStoreByteChip::new( + LoadStoreByteFiller::new( Rv64LoadStoreAdapterFiller::new(byte_ptr_max_bits, range_checker.clone()), - range_checker.clone(), + Rv64LoadStoreOpcode::CLASS_OFFSET, bitwise_lu.clone(), + range_checker.clone(), + ), + mem_helper.clone(), + ); + inventory.add_executor_chip(load_store_byte_chip); + + inventory.next_air::()?; + let load_sign_extend_halfword_chip = Rv64LoadSignExtendHalfwordChip::new( + LoadSignExtendHalfwordFiller::new( + Rv64LoadStoreAdapterFiller::new(byte_ptr_max_bits, range_checker.clone()), + Rv64LoadStoreOpcode::CLASS_OFFSET, + range_checker.clone(), + ), + mem_helper.clone(), + ); + inventory.add_executor_chip(load_sign_extend_halfword_chip); + + inventory.next_air::()?; + let load_store_halfword_chip = Rv64LoadStoreHalfwordChip::new( + LoadStoreHalfwordFiller::new( + Rv64LoadStoreAdapterFiller::new(byte_ptr_max_bits, range_checker.clone()), + Rv64LoadStoreOpcode::CLASS_OFFSET, + range_checker.clone(), + ), + mem_helper.clone(), + ); + inventory.add_executor_chip(load_store_halfword_chip); + + inventory.next_air::()?; + let load_sign_extend_word_chip = Rv64LoadSignExtendWordChip::new( + LoadSignExtendWordFiller::new( + Rv64LoadStoreAdapterFiller::new(byte_ptr_max_bits, range_checker.clone()), + Rv64LoadStoreOpcode::CLASS_OFFSET, + range_checker.clone(), + ), + mem_helper.clone(), + ); + inventory.add_executor_chip(load_sign_extend_word_chip); + + inventory.next_air::()?; + let load_store_word_chip = Rv64LoadStoreWordChip::new( + LoadStoreWordFiller::new( + Rv64LoadStoreAdapterFiller::new(byte_ptr_max_bits, range_checker.clone()), + Rv64LoadStoreOpcode::CLASS_OFFSET, + range_checker.clone(), + ), + mem_helper.clone(), + ); + inventory.add_executor_chip(load_store_word_chip); + + inventory.next_air::()?; + let load_store_doubleword_chip = Rv64LoadStoreDoublewordChip::new( + LoadStoreDoublewordFiller::new( + Rv64LoadStoreAdapterFiller::new(byte_ptr_max_bits, range_checker.clone()), + Rv64LoadStoreOpcode::CLASS_OFFSET, + range_checker.clone(), ), mem_helper.clone(), ); - inventory.add_executor_chip(load_sign_extend); + inventory.add_executor_chip(load_store_doubleword_chip); inventory.next_air::()?; let beq = Rv64BranchEqualChip::new( diff --git a/extensions/riscv/circuit/src/load_sign_extend/aligned/core.rs b/extensions/riscv/circuit/src/load_sign_extend/aligned/core.rs new file mode 100644 index 0000000000..05954fdc6c --- /dev/null +++ b/extensions/riscv/circuit/src/load_sign_extend/aligned/core.rs @@ -0,0 +1,287 @@ +use std::borrow::{Borrow, BorrowMut}; + +use openvm_circuit::{arch::*, system::memory::MemoryAuxColsFactory}; +use openvm_circuit_primitives::{ + encoder::Encoder, + var_range::{SharedVariableRangeCheckerChip, VariableRangeCheckerBus}, + AlignedBorrow, ColumnsAir, StructReflection, StructReflectionHelper, SubAir, +}; +use openvm_instructions::LocalOpcode; +use openvm_riscv_transpiler::Rv64LoadStoreOpcode::{self, *}; +use openvm_stark_backend::{ + interaction::InteractionBuilder, + p3_air::BaseAir, + p3_field::{Field, PrimeCharacteristicRing, PrimeField32}, + BaseAirWithPublicValues, +}; + +use crate::{ + adapters::{LoadStoreInstruction, Rv64LoadStoreAdapterFiller, U16_BITS}, + loadstore::common::{adapter_context, LoadStoreRecord, KIND_HALFWORD, KIND_WORD, SIGN_U16}, +}; + +const SELECTOR_MAX_DEGREE: u32 = 2; + +#[derive(Clone, Copy)] +pub(crate) struct SignedAlignedCase { + opcode: Rv64LoadStoreOpcode, + byte_shift: usize, +} + +impl SignedAlignedCase { + fn cell_shift(self) -> usize { + self.byte_shift / 2 + } +} + +const WORD_CASES: [SignedAlignedCase; 2] = [ + SignedAlignedCase { + opcode: LOADW, + byte_shift: 0, + }, + SignedAlignedCase { + opcode: LOADW, + byte_shift: 4, + }, +]; + +const HALFWORD_CASES: [SignedAlignedCase; 4] = [ + SignedAlignedCase { + opcode: LOADH, + byte_shift: 0, + }, + SignedAlignedCase { + opcode: LOADH, + byte_shift: 2, + }, + SignedAlignedCase { + opcode: LOADH, + byte_shift: 4, + }, + SignedAlignedCase { + opcode: LOADH, + byte_shift: 6, + }, +]; + +pub(crate) fn signed_aligned_cases() -> &'static [SignedAlignedCase] { + match KIND { + KIND_WORD => &WORD_CASES, + KIND_HALFWORD => &HALFWORD_CASES, + _ => unreachable!("unsupported signed aligned load kind"), + } +} + +fn access_cells() -> usize { + match KIND { + KIND_WORD => 2, + KIND_HALFWORD => 1, + _ => unreachable!("unsupported signed aligned load kind"), + } +} + +fn encoder() -> Encoder { + let encoder = Encoder::new(CASES, SELECTOR_MAX_DEGREE, true); + debug_assert_eq!(encoder.width(), SELECTOR_WIDTH); + encoder +} + +#[repr(C)] +#[derive(Debug, Clone, AlignedBorrow, StructReflection)] +pub struct LoadSignExtendAlignedCoreCols { + pub selector: [T; SELECTOR_WIDTH], + pub is_valid: T, + pub data_most_sig_bit: T, + pub read_data: [T; BLOCK_FE_WIDTH], + pub prev_data: [T; BLOCK_FE_WIDTH], +} + +#[derive(Debug, Clone, ColumnsAir)] +#[columns_via(LoadSignExtendAlignedCoreCols)] +pub struct LoadSignExtendAlignedCoreAir< + const KIND: usize, + const CASES: usize, + const SELECTOR_WIDTH: usize, +> { + pub offset: usize, + encoder: Encoder, + range_bus: VariableRangeCheckerBus, +} + +impl + LoadSignExtendAlignedCoreAir +{ + pub fn new(offset: usize, range_bus: VariableRangeCheckerBus) -> Self { + debug_assert_eq!(signed_aligned_cases::().len(), CASES); + Self { + offset, + encoder: encoder::(), + range_bus, + } + } +} + +impl BaseAir + for LoadSignExtendAlignedCoreAir +{ + fn width(&self) -> usize { + LoadSignExtendAlignedCoreCols::::width() + } +} + +impl + BaseAirWithPublicValues for LoadSignExtendAlignedCoreAir +{ +} + +impl VmCoreAir + for LoadSignExtendAlignedCoreAir +where + AB: InteractionBuilder, + I: VmAdapterInterface, + I::Reads: From<([AB::Var; BLOCK_FE_WIDTH], [AB::Expr; BLOCK_FE_WIDTH])>, + I::Writes: From<[[AB::Expr; BLOCK_FE_WIDTH]; 1]>, + I::ProcessedInstruction: From>, +{ + fn eval( + &self, + builder: &mut AB, + local_core: &[AB::Var], + _from_pc: AB::Var, + ) -> AdapterAirContext { + let cols: &LoadSignExtendAlignedCoreCols = (*local_core).borrow(); + let cases = signed_aligned_cases::(); + let width = access_cells::(); + + self.encoder.eval(builder, &cols.selector); + let flags = self.encoder.flags::(&cols.selector); + let is_valid = self.encoder.is_valid::(&cols.selector); + + builder.assert_eq(cols.is_valid, is_valid.clone()); + builder.assert_bool(cols.data_most_sig_bit); + + let sign_cell = cases + .iter() + .enumerate() + .fold(AB::Expr::ZERO, |acc, (i, case)| { + acc + flags[i].clone() * cols.read_data[case.cell_shift() + width - 1] + }); + self.range_bus + .range_check( + sign_cell - cols.data_most_sig_bit * AB::Expr::from_u32(SIGN_U16 as u32), + U16_BITS - 1, + ) + .eval(builder, is_valid.clone()); + + let expected_opcode = cases + .iter() + .enumerate() + .fold(AB::Expr::ZERO, |acc, (i, case)| { + acc + flags[i].clone() * AB::Expr::from_u8(case.opcode as u8) + }); + let expected_opcode = VmCoreAir::::expr_to_global_expr(self, expected_opcode); + let load_shift_amount = cases + .iter() + .enumerate() + .fold(AB::Expr::ZERO, |acc, (i, case)| { + acc + flags[i].clone() * AB::Expr::from_usize(case.byte_shift) + }); + + let sign_extend = cols.data_most_sig_bit * AB::Expr::from_u32(u16::MAX as u32); + let write_data = std::array::from_fn(|i| { + cases + .iter() + .enumerate() + .fold(AB::Expr::ZERO, |acc, (case_idx, case)| { + let shift = case.cell_shift(); + let term = if i < width { + cols.read_data[i + shift].into() + } else { + sign_extend.clone() + }; + acc + flags[case_idx].clone() * term + }) + }); + adapter_context::( + cols.is_valid.into(), + cols.is_valid.into(), + expected_opcode, + load_shift_amount, + AB::Expr::ZERO, + cols.read_data, + cols.prev_data, + write_data, + ) + } + + fn start_offset(&self) -> usize { + self.offset + } +} + +#[derive(Clone)] +pub struct LoadSignExtendAlignedFiller< + A = Rv64LoadStoreAdapterFiller, + const KIND: usize = KIND_WORD, + const CASES: usize = 2, + const SELECTOR_WIDTH: usize = 2, +> { + adapter: A, + pub offset: usize, + encoder: Encoder, + range_checker_chip: SharedVariableRangeCheckerChip, +} + +impl + LoadSignExtendAlignedFiller +{ + pub fn new( + adapter: A, + offset: usize, + range_checker_chip: SharedVariableRangeCheckerChip, + ) -> Self { + debug_assert_eq!(signed_aligned_cases::().len(), CASES); + Self { + adapter, + offset, + encoder: encoder::(), + range_checker_chip, + } + } +} + +impl TraceFiller + for LoadSignExtendAlignedFiller +where + F: PrimeField32, + A: 'static + AdapterTraceFiller, +{ + fn fill_trace_row(&self, mem_helper: &MemoryAuxColsFactory, row_slice: &mut [F]) { + let (adapter_row, mut core_row) = unsafe { row_slice.split_at_mut_unchecked(A::WIDTH) }; + self.adapter.fill_trace_row(mem_helper, adapter_row); + + let record: &LoadStoreRecord = unsafe { get_record_from_slice(&mut core_row, ()) }; + let opcode = Rv64LoadStoreOpcode::from_usize(record.local_opcode as usize); + let shift = record.shift_amount as usize; + let read_data = record.read_data; + let prev_data = record.prev_data; + let core_row: &mut LoadSignExtendAlignedCoreCols = core_row.borrow_mut(); + let cases = signed_aligned_cases::(); + let case_idx = cases + .iter() + .position(|case| case.opcode == opcode && case.byte_shift == shift) + .expect("invalid signed aligned load opcode/shift"); + + let width = access_cells::(); + let sign_cell = read_data[shift / 2 + width - 1]; + let sign_bit = sign_cell & SIGN_U16; + self.range_checker_chip + .add_count((sign_cell - sign_bit) as u32, U16_BITS - 1); + core_row.data_most_sig_bit = F::from_bool(sign_bit != 0); + core_row.read_data = read_data.map(F::from_u16); + core_row.prev_data = prev_data.map(F::from_u16); + core_row.is_valid = F::ONE; + let pt: [u32; SELECTOR_WIDTH] = self.encoder.get_flag_pt(case_idx).try_into().unwrap(); + core_row.selector = pt.map(F::from_u32); + } +} diff --git a/extensions/riscv/circuit/src/load_sign_extend/aligned/mod.rs b/extensions/riscv/circuit/src/load_sign_extend/aligned/mod.rs new file mode 100644 index 0000000000..5a7ca06a4f --- /dev/null +++ b/extensions/riscv/circuit/src/load_sign_extend/aligned/mod.rs @@ -0,0 +1 @@ +pub mod core; diff --git a/extensions/riscv/circuit/src/load_sign_extend/byte/core.rs b/extensions/riscv/circuit/src/load_sign_extend/byte/core.rs new file mode 100644 index 0000000000..f45ef3d6b1 --- /dev/null +++ b/extensions/riscv/circuit/src/load_sign_extend/byte/core.rs @@ -0,0 +1,233 @@ +use std::borrow::{Borrow, BorrowMut}; + +use openvm_circuit::{arch::*, system::memory::MemoryAuxColsFactory}; +use openvm_circuit_primitives::{ + bitwise_op_lookup::{BitwiseOperationLookupBus, SharedBitwiseOperationLookupChip}, + encoder::Encoder, + var_range::{SharedVariableRangeCheckerChip, VariableRangeCheckerBus}, + AlignedBorrow, ColumnsAir, StructReflection, StructReflectionHelper, SubAir, +}; +use openvm_instructions::LocalOpcode; +use openvm_riscv_transpiler::Rv64LoadStoreOpcode::{self, *}; +use openvm_stark_backend::{ + interaction::InteractionBuilder, + p3_air::BaseAir, + p3_field::{Field, PrimeCharacteristicRing, PrimeField32}, + BaseAirWithPublicValues, +}; + +use crate::{ + adapters::{LoadStoreInstruction, Rv64LoadStoreAdapterFiller, RV64_BYTE_BITS}, + loadstore::common::{adapter_context, byte_from_cell, LoadStoreRecord, BYTE_BITS, SIGN_BYTE}, +}; + +const LOAD_SIGN_EXTEND_BYTE_CASES: usize = 8; +const LOAD_SIGN_EXTEND_BYTE_SELECTOR_MAX_DEGREE: u32 = 2; +pub(crate) const LOAD_SIGN_EXTEND_BYTE_SELECTOR_WIDTH: usize = 3; + +fn encoder() -> Encoder { + let encoder = Encoder::new( + LOAD_SIGN_EXTEND_BYTE_CASES, + LOAD_SIGN_EXTEND_BYTE_SELECTOR_MAX_DEGREE, + true, + ); + debug_assert_eq!(encoder.width(), LOAD_SIGN_EXTEND_BYTE_SELECTOR_WIDTH); + encoder +} + +#[repr(C)] +#[derive(Debug, Clone, AlignedBorrow, StructReflection)] +pub struct LoadSignExtendByteCoreCols { + pub selector: [T; LOAD_SIGN_EXTEND_BYTE_SELECTOR_WIDTH], + pub is_valid: T, + pub data_most_sig_bit: T, + pub read_cell_bytes: [T; 2], + pub read_data: [T; BLOCK_FE_WIDTH], + pub prev_data: [T; BLOCK_FE_WIDTH], +} + +#[derive(Debug, Clone, ColumnsAir)] +#[columns_via(LoadSignExtendByteCoreCols)] +pub struct LoadSignExtendByteCoreAir { + pub offset: usize, + encoder: Encoder, + bitwise_lookup_bus: BitwiseOperationLookupBus, + range_bus: VariableRangeCheckerBus, +} + +impl LoadSignExtendByteCoreAir { + pub fn new( + offset: usize, + bitwise_lookup_bus: BitwiseOperationLookupBus, + range_bus: VariableRangeCheckerBus, + ) -> Self { + Self { + offset, + encoder: encoder(), + bitwise_lookup_bus, + range_bus, + } + } +} + +impl BaseAir for LoadSignExtendByteCoreAir { + fn width(&self) -> usize { + LoadSignExtendByteCoreCols::::width() + } +} + +impl BaseAirWithPublicValues for LoadSignExtendByteCoreAir {} + +impl VmCoreAir for LoadSignExtendByteCoreAir +where + AB: InteractionBuilder, + I: VmAdapterInterface, + I::Reads: From<([AB::Var; BLOCK_FE_WIDTH], [AB::Expr; BLOCK_FE_WIDTH])>, + I::Writes: From<[[AB::Expr; BLOCK_FE_WIDTH]; 1]>, + I::ProcessedInstruction: From>, +{ + fn eval( + &self, + builder: &mut AB, + local_core: &[AB::Var], + _from_pc: AB::Var, + ) -> AdapterAirContext { + let cols: &LoadSignExtendByteCoreCols = (*local_core).borrow(); + self.encoder.eval(builder, &cols.selector); + let flags = self.encoder.flags::(&cols.selector); + let is_valid = self.encoder.is_valid::(&cols.selector); + + builder.assert_eq(cols.is_valid, is_valid.clone()); + builder.assert_bool(cols.data_most_sig_bit); + + self.bitwise_lookup_bus + .send_range(cols.read_cell_bytes[0], cols.read_cell_bytes[1]) + .eval(builder, is_valid.clone()); + + let read_cell = + cols.read_cell_bytes[0] + cols.read_cell_bytes[1] * AB::Expr::from_u32(1 << BYTE_BITS); + let expected_read_cell = flags + .iter() + .enumerate() + .fold(AB::Expr::ZERO, |acc, (shift, flag)| { + acc + flag.clone() * cols.read_data[shift / 2] + }); + builder.assert_eq(read_cell, expected_read_cell); + + let selected_byte = flags + .iter() + .enumerate() + .fold(AB::Expr::ZERO, |acc, (shift, flag)| { + let byte = if shift % 2 == 0 { + cols.read_cell_bytes[0].into() + } else { + cols.read_cell_bytes[1].into() + }; + acc + flag.clone() * byte + }); + self.range_bus + .range_check( + selected_byte.clone() + - cols.data_most_sig_bit * AB::Expr::from_u32(SIGN_BYTE as u32), + RV64_BYTE_BITS - 1, + ) + .eval(builder, is_valid.clone()); + + let sign_cell = cols.data_most_sig_bit * AB::Expr::from_u32(u16::MAX as u32); + let write_data = std::array::from_fn(|i| { + if i == 0 { + selected_byte.clone() + cols.data_most_sig_bit * AB::Expr::from_u32(0xff00) + } else { + sign_cell.clone() + } + }); + let load_shift_amount = flags + .iter() + .enumerate() + .fold(AB::Expr::ZERO, |acc, (shift, flag)| { + acc + flag.clone() * AB::Expr::from_usize(shift) + }); + let expected_opcode = VmCoreAir::::expr_to_global_expr( + self, + is_valid.clone() * AB::Expr::from_u8(LOADB as u8), + ); + + adapter_context::( + cols.is_valid.into(), + cols.is_valid.into(), + expected_opcode, + load_shift_amount, + AB::Expr::ZERO, + cols.read_data, + cols.prev_data, + write_data, + ) + } + + fn start_offset(&self) -> usize { + self.offset + } +} + +#[derive(Clone)] +pub struct LoadSignExtendByteFiller { + adapter: A, + pub offset: usize, + encoder: Encoder, + bitwise_lookup_chip: SharedBitwiseOperationLookupChip, + range_checker_chip: SharedVariableRangeCheckerChip, +} + +impl LoadSignExtendByteFiller { + pub fn new( + adapter: A, + offset: usize, + bitwise_lookup_chip: SharedBitwiseOperationLookupChip, + range_checker_chip: SharedVariableRangeCheckerChip, + ) -> Self { + Self { + adapter, + offset, + encoder: encoder(), + bitwise_lookup_chip, + range_checker_chip, + } + } +} + +impl TraceFiller for LoadSignExtendByteFiller +where + F: PrimeField32, + A: 'static + AdapterTraceFiller, +{ + fn fill_trace_row(&self, mem_helper: &MemoryAuxColsFactory, row_slice: &mut [F]) { + let (adapter_row, mut core_row) = unsafe { row_slice.split_at_mut_unchecked(A::WIDTH) }; + self.adapter.fill_trace_row(mem_helper, adapter_row); + + let record: &LoadStoreRecord = unsafe { get_record_from_slice(&mut core_row, ()) }; + let opcode = Rv64LoadStoreOpcode::from_usize(record.local_opcode as usize); + let shift = record.shift_amount as usize; + let read_data = record.read_data; + let prev_data = record.prev_data; + debug_assert_eq!(opcode, LOADB); + let core_row: &mut LoadSignExtendByteCoreCols = core_row.borrow_mut(); + + let read_cell = read_data[shift / 2]; + let read_cell_bytes = [byte_from_cell(read_cell, 0), byte_from_cell(read_cell, 1)]; + self.bitwise_lookup_chip + .request_range(read_cell_bytes[0] as u32, read_cell_bytes[1] as u32); + core_row.read_cell_bytes = read_cell_bytes.map(F::from_u16); + + let byte = read_cell_bytes[shift % 2]; + let sign_bit = byte & SIGN_BYTE; + self.range_checker_chip + .add_count((byte - sign_bit) as u32, RV64_BYTE_BITS - 1); + core_row.data_most_sig_bit = F::from_bool(sign_bit != 0); + core_row.read_data = read_data.map(F::from_u16); + core_row.prev_data = prev_data.map(F::from_u16); + core_row.is_valid = F::ONE; + let pt: [u32; LOAD_SIGN_EXTEND_BYTE_SELECTOR_WIDTH] = + self.encoder.get_flag_pt(shift).try_into().unwrap(); + core_row.selector = pt.map(F::from_u32); + } +} diff --git a/extensions/riscv/circuit/src/load_sign_extend/byte/mod.rs b/extensions/riscv/circuit/src/load_sign_extend/byte/mod.rs new file mode 100644 index 0000000000..cdeab45e2e --- /dev/null +++ b/extensions/riscv/circuit/src/load_sign_extend/byte/mod.rs @@ -0,0 +1,15 @@ +use openvm_circuit::arch::{VmAirWrapper, VmChipWrapper}; + +use crate::{ + adapters::{Rv64LoadStoreAdapterAir, Rv64LoadStoreAdapterExecutor}, + loadstore::common::{LoadStoreExecutor, KIND_BYTE}, +}; + +mod core; +pub use core::*; + +pub type Rv64LoadSignExtendByteAir = + VmAirWrapper; +pub type Rv64LoadSignExtendByteExecutor = + LoadStoreExecutor; +pub type Rv64LoadSignExtendByteChip = VmChipWrapper; diff --git a/extensions/riscv/circuit/src/load_sign_extend/core.rs b/extensions/riscv/circuit/src/load_sign_extend/core.rs deleted file mode 100644 index ef9c78164c..0000000000 --- a/extensions/riscv/circuit/src/load_sign_extend/core.rs +++ /dev/null @@ -1,450 +0,0 @@ -use std::{ - array, - borrow::{Borrow, BorrowMut}, -}; - -use openvm_circuit::{ - arch::*, - system::memory::{online::TracingMemory, MemoryAuxColsFactory}, -}; -use openvm_circuit_primitives::{ - bitwise_op_lookup::{BitwiseOperationLookupBus, SharedBitwiseOperationLookupChip}, - utils::select, - var_range::{SharedVariableRangeCheckerChip, VariableRangeCheckerBus}, - AlignedBytesBorrow, ColumnsAir, StructReflection, StructReflectionHelper, -}; -use openvm_circuit_primitives_derive::AlignedBorrow; -use openvm_instructions::{ - instruction::Instruction, - program::DEFAULT_PC_STEP, - riscv::{RV64_BYTE_BITS, RV64_MEMORY_AS, RV64_REGISTER_AS, RV64_REGISTER_NUM_LIMBS}, - LocalOpcode, -}; -use openvm_riscv_transpiler::Rv64LoadStoreOpcode::{self, *}; -use openvm_stark_backend::{ - interaction::InteractionBuilder, - p3_air::BaseAir, - p3_field::{Field, PrimeCharacteristicRing, PrimeField32}, - BaseAirWithPublicValues, -}; - -use crate::adapters::{LoadStoreInstruction, Rv64LoadStoreAdapterFiller}; - -/// LoadSignExtend Core Chip handles byte/halfword/word into doubleword conversions through sign -/// extend. This chip uses read_data to construct write_data. -/// prev_data columns are not used in constraints defined in the CoreAir, but are used in -/// constraints by the Adapter. shifted_read_data is the read_data shifted by (shift_amount & 4), -/// this reduces the number of opcode flags needed. Using this shifted data we can generate the -/// write_data as if the shift_amount was 0..3 for loadb, 0 or 2 for loadh, and 0 for loadw. -#[repr(C)] -#[derive(Debug, Clone, AlignedBorrow, StructReflection)] -pub struct LoadSignExtendCoreCols { - /// This chip treats each (opcode, inner_shift) pair as a different instruction - pub opcode_loadb_flag0: T, - pub opcode_loadb_flag1: T, - pub opcode_loadb_flag2: T, - pub opcode_loadb_flag3: T, - pub opcode_loadh_flag0: T, - pub opcode_loadh_flag2: T, - pub opcode_loadw_flag: T, - - pub shift_most_sig_bit: T, - // The bit that is extended to the remaining bits - pub data_most_sig_bit: T, - - pub shifted_read_data: [T; NUM_CELLS], - pub prev_data: [T; NUM_CELLS], -} - -#[derive(Debug, Clone, ColumnsAir)] -#[columns_via(LoadSignExtendCoreCols)] -pub struct LoadSignExtendCoreAir { - pub range_bus: VariableRangeCheckerBus, - pub bitwise_lookup_bus: BitwiseOperationLookupBus, -} - -impl LoadSignExtendCoreAir { - pub fn new( - range_bus: VariableRangeCheckerBus, - bitwise_lookup_bus: BitwiseOperationLookupBus, - ) -> Self { - assert!(NUM_CELLS.is_multiple_of(2)); - Self { - range_bus, - bitwise_lookup_bus, - } - } -} - -impl BaseAir - for LoadSignExtendCoreAir -{ - fn width(&self) -> usize { - LoadSignExtendCoreCols::::width() - } -} - -impl BaseAirWithPublicValues - for LoadSignExtendCoreAir -{ -} - -impl VmCoreAir - for LoadSignExtendCoreAir -where - AB: InteractionBuilder, - I: VmAdapterInterface, - I::Reads: From<([AB::Var; NUM_CELLS], [AB::Expr; NUM_CELLS])>, - I::Writes: From<[[AB::Expr; NUM_CELLS]; 1]>, - I::ProcessedInstruction: From>, -{ - fn eval( - &self, - builder: &mut AB, - local_core: &[AB::Var], - _from_pc: AB::Var, - ) -> AdapterAirContext { - let cols: &LoadSignExtendCoreCols = (*local_core).borrow(); - let LoadSignExtendCoreCols:: { - shifted_read_data, - prev_data, - opcode_loadb_flag0: is_loadb0, - opcode_loadb_flag1: is_loadb1, - opcode_loadb_flag2: is_loadb2, - opcode_loadb_flag3: is_loadb3, - opcode_loadh_flag0: is_loadh0, - opcode_loadh_flag2: is_loadh2, - opcode_loadw_flag: is_loadw, - data_most_sig_bit, - shift_most_sig_bit, - } = *cols; - - let flags = [ - is_loadb0, is_loadb1, is_loadb2, is_loadb3, is_loadh0, is_loadh2, is_loadw, - ]; - - let is_valid = flags.iter().fold(AB::Expr::ZERO, |acc, &flag| { - builder.assert_bool(flag); - acc + flag - }); - - builder.assert_bool(is_valid.clone()); - builder.assert_bool(data_most_sig_bit); - builder.assert_bool(shift_most_sig_bit); - - let expected_opcode = (is_loadb0 + is_loadb1 + is_loadb2 + is_loadb3) - * AB::F::from_u8(LOADB as u8) - + (is_loadh0 + is_loadh2) * AB::F::from_u8(LOADH as u8) - + is_loadw * AB::F::from_u8(LOADW as u8) - + AB::Expr::from_usize(Rv64LoadStoreOpcode::CLASS_OFFSET); - - let limb_mask = data_most_sig_bit * AB::Expr::from_u32((1 << LIMB_BITS) - 1); - - let sd = shifted_read_data; - - // there are four parts to write_data: - // - 1st limb is the sign-extended byte (selected by opcode and inner_shift) - // - 2nd limb is read_data if loadh/loadw and sign extended if loadb - // - 3rd to 4th limbs are read_data if loadw and sign extended otherwise - // - 5th to last limbs are always sign extended limbs - let write_data: [AB::Expr; NUM_CELLS] = array::from_fn(|i| { - if i == 0 { - (is_loadb0 + is_loadh0 + is_loadw) * sd[0] - + is_loadb1 * sd[1] - + (is_loadb2 + is_loadh2) * sd[2] - + is_loadb3 * sd[3] - } else if i == 1 { - (is_loadh0 + is_loadw) * sd[1] - + is_loadh2 * sd[3] - + (is_loadb0 + is_loadb1 + is_loadb2 + is_loadb3) * limb_mask.clone() - } else if i < 4 { - is_loadw * sd[i] + (is_valid.clone() - is_loadw) * limb_mask.clone() - } else { - limb_mask.clone() - } - }); - - // Constrain that most_sig_bit is correct - let most_sig_limb = is_loadb0 * sd[0] - + (is_loadb1 + is_loadh0) * sd[1] - + is_loadb2 * sd[2] - + (is_loadb3 + is_loadh2 + is_loadw) * sd[3]; - - self.range_bus - .range_check( - most_sig_limb - data_most_sig_bit * AB::Expr::from_u32(1 << (LIMB_BITS - 1)), - LIMB_BITS - 1, - ) - .eval(builder, is_valid.clone()); - for pair in shifted_read_data.chunks_exact(2) { - self.bitwise_lookup_bus - .send_range(pair[0], pair[1]) - .eval(builder, is_valid.clone()); - } - - // Unshift the shifted_read_data to get the original read_data - let read_data: [AB::Expr; NUM_CELLS] = array::from_fn(|i| { - select( - shift_most_sig_bit, - sd[(i + NUM_CELLS - 4) % NUM_CELLS], - sd[i], - ) - }); - - let load_shift_amount = shift_most_sig_bit * AB::Expr::from_u32(4) - + is_loadb1 - + (is_loadb2 + is_loadh2) * AB::Expr::TWO - + is_loadb3 * AB::Expr::from_u32(3); - - AdapterAirContext { - to_pc: None, - reads: (prev_data, read_data).into(), - writes: [write_data].into(), - instruction: LoadStoreInstruction { - is_valid: is_valid.clone(), - opcode: expected_opcode, - is_load: is_valid, - load_shift_amount, - store_shift_amount: AB::Expr::ZERO, - } - .into(), - } - } - - fn start_offset(&self) -> usize { - Rv64LoadStoreOpcode::CLASS_OFFSET - } -} - -#[repr(C)] -#[derive(AlignedBytesBorrow, Debug)] -pub struct LoadSignExtendCoreRecord { - pub is_byte: bool, - pub is_word: bool, - pub shift_amount: u8, - pub read_data: [u8; NUM_CELLS], - pub prev_data: [u8; NUM_CELLS], -} - -#[derive(Clone, Copy, derive_new::new)] -pub struct LoadSignExtendExecutor { - adapter: A, -} - -#[derive(Clone)] -pub struct LoadSignExtendFiller< - A = Rv64LoadStoreAdapterFiller, - const NUM_CELLS: usize = RV64_REGISTER_NUM_LIMBS, - const LIMB_BITS: usize = RV64_BYTE_BITS, -> { - adapter: A, - pub range_checker_chip: SharedVariableRangeCheckerChip, - pub bitwise_lookup_chip: SharedBitwiseOperationLookupChip, -} - -impl - LoadSignExtendFiller -{ - pub fn new( - adapter: A, - range_checker_chip: SharedVariableRangeCheckerChip, - bitwise_lookup_chip: SharedBitwiseOperationLookupChip, - ) -> Self { - assert!(NUM_CELLS.is_multiple_of(2)); - Self { - adapter, - range_checker_chip, - bitwise_lookup_chip, - } - } -} - -impl PreflightExecutor - for LoadSignExtendExecutor -where - F: PrimeField32, - A: 'static - + AdapterTraceExecutor< - F, - ReadData = (([u8; NUM_CELLS], [u8; NUM_CELLS]), u8), - WriteData = [u8; NUM_CELLS], - >, - for<'buf> RA: RecordArena< - 'buf, - EmptyAdapterCoreLayout, - ( - A::RecordMut<'buf>, - &'buf mut LoadSignExtendCoreRecord, - ), - >, -{ - fn get_opcode_name(&self, opcode: usize) -> String { - format!( - "{:?}", - Rv64LoadStoreOpcode::from_usize(opcode - Rv64LoadStoreOpcode::CLASS_OFFSET) - ) - } - - fn execute( - &self, - state: VmStateMut, - instruction: &Instruction, - ) -> Result<(), ExecutionError> { - let Instruction { opcode, d, e, .. } = instruction; - - let local_opcode = Rv64LoadStoreOpcode::from_usize( - opcode.local_opcode_idx(Rv64LoadStoreOpcode::CLASS_OFFSET), - ); - debug_assert_eq!(d.as_canonical_u32(), RV64_REGISTER_AS); - debug_assert_eq!(e.as_canonical_u32(), RV64_MEMORY_AS); - - let (mut adapter_record, core_record) = state.ctx.alloc(EmptyAdapterCoreLayout::new()); - - A::start(*state.pc, state.memory, &mut adapter_record); - - let tmp = self - .adapter - .read(state.memory, instruction, &mut adapter_record); - - core_record.is_byte = local_opcode == LOADB; - core_record.is_word = local_opcode == LOADW; - core_record.prev_data = tmp.0 .0; - core_record.read_data = tmp.0 .1; - core_record.shift_amount = tmp.1; - - let write_data = run_write_data_sign_extend::( - local_opcode, - core_record.read_data, - core_record.shift_amount as usize, - ); - - self.adapter - .write(state.memory, instruction, write_data, &mut adapter_record); - - *state.pc = state.pc.wrapping_add(DEFAULT_PC_STEP); - - Ok(()) - } -} - -impl TraceFiller - for LoadSignExtendFiller -where - F: PrimeField32, - A: 'static + AdapterTraceFiller, -{ - fn fill_trace_row(&self, mem_helper: &MemoryAuxColsFactory, row_slice: &mut [F]) { - // SAFETY: row_slice is guaranteed by the caller to have at least A::WIDTH + - // LoadSignExtendCoreCols::width() elements - let (adapter_row, mut core_row) = unsafe { row_slice.split_at_mut_unchecked(A::WIDTH) }; - self.adapter.fill_trace_row(mem_helper, adapter_row); - // SAFETY: core_row contains a valid LoadSignExtendCoreRecord written by the executor - // during trace generation - let record: &LoadSignExtendCoreRecord = - unsafe { get_record_from_slice(&mut core_row, ()) }; - - let core_row: &mut LoadSignExtendCoreCols = core_row.borrow_mut(); - - let shift = record.shift_amount as usize; - let shift_most_sig_bit = (shift >> 2) & 1; - let inner_shift = shift & 3; - - let mut shifted = record.read_data; - shifted.rotate_left(shift_most_sig_bit * (NUM_CELLS / 2)); - - let most_sig_limb = if record.is_byte { - shifted[inner_shift] - } else if record.is_word { - shifted[NUM_CELLS / 2 - 1] - } else { - shifted[inner_shift + 1] - }; - - let most_sig_bit = most_sig_limb & (1 << (LIMB_BITS - 1)); - self.range_checker_chip - .add_count((most_sig_limb - most_sig_bit) as u32, LIMB_BITS - 1); - for pair in shifted.chunks_exact(2) { - self.bitwise_lookup_chip - .request_range(pair[0] as u32, pair[1] as u32); - } - - core_row.prev_data = record.prev_data.map(F::from_u8); - core_row.shifted_read_data = shifted.map(F::from_u8); - - core_row.data_most_sig_bit = F::from_bool(most_sig_bit != 0); - core_row.shift_most_sig_bit = F::from_bool(shift_most_sig_bit == 1); - - let is_byte = record.is_byte; - let is_word = record.is_word; - let is_half = !is_byte && !is_word; - - core_row.opcode_loadb_flag0 = F::from_bool(is_byte && inner_shift == 0); - core_row.opcode_loadb_flag1 = F::from_bool(is_byte && inner_shift == 1); - core_row.opcode_loadb_flag2 = F::from_bool(is_byte && inner_shift == 2); - core_row.opcode_loadb_flag3 = F::from_bool(is_byte && inner_shift == 3); - core_row.opcode_loadh_flag0 = F::from_bool(is_half && inner_shift == 0); - core_row.opcode_loadh_flag2 = F::from_bool(is_half && inner_shift == 2); - core_row.opcode_loadw_flag = F::from_bool(is_word); - } -} - -// Returns write_data -#[inline(always)] -pub(super) fn run_write_data_sign_extend( - opcode: Rv64LoadStoreOpcode, - read_data: [u8; NUM_CELLS], - shift: usize, -) -> [u8; NUM_CELLS] { - const { assert!(NUM_CELLS == RV64_REGISTER_NUM_LIMBS) }; - let word_width = NUM_CELLS / 2; - let half_width = NUM_CELLS / 4; - match opcode { - LOADW => { - assert!( - shift == 0 || shift == word_width, - "LOADW requires {word_width}-byte aligned shift, got {shift}" - ); - assert!(shift + word_width <= NUM_CELLS); - let ext = (read_data[shift + word_width - 1] >> (LIMB_BITS - 1)) * u8::MAX; - array::from_fn(|i| { - if i < word_width { - read_data[i + shift] - } else { - ext - } - }) - } - LOADH => { - assert!( - shift.is_multiple_of(half_width), - "LOADH requires {half_width}-byte aligned shift, got {shift}" - ); - debug_assert!(shift + half_width <= NUM_CELLS); - let ext = (read_data[shift + half_width - 1] >> (LIMB_BITS - 1)) * u8::MAX; - array::from_fn(|i| { - if i < half_width { - read_data[i + shift] - } else { - ext - } - }) - } - LOADB => { - debug_assert!(shift < NUM_CELLS); - let ext = (read_data[shift] >> (LIMB_BITS - 1)) * u8::MAX; - array::from_fn(|i| { - if i == 0 { - read_data[i + shift] - } else { - ext - } - }) - } - // Currently the adapter AIR requires `ptr_val` to be aligned to the data size in bytes. - // The circuit requires that `shift = ptr_val % 8` so that `ptr_val - shift` is a multiple of 8. - // This requirement is non-trivial to remove, because we use it to ensure that `ptr_val - shift + 8 <= 2^pointer_max_bits`. - _ => unreachable!( - "unaligned memory access not supported by this execution environment: {opcode:?}, shift: {shift}" - ), - } -} diff --git a/extensions/riscv/circuit/src/load_sign_extend/cuda.rs b/extensions/riscv/circuit/src/load_sign_extend/cuda.rs index a0f710b42b..6095dc82b3 100644 --- a/extensions/riscv/circuit/src/load_sign_extend/cuda.rs +++ b/extensions/riscv/circuit/src/load_sign_extend/cuda.rs @@ -1,67 +1,58 @@ -use std::{mem::size_of, sync::Arc}; +use std::sync::Arc; use derive_new::new; -use openvm_circuit::{arch::DenseRecordArena, utils::next_power_of_two_or_zero}; +use openvm_circuit::arch::DenseRecordArena; use openvm_circuit_primitives::{ bitwise_op_lookup::BitwiseOperationLookupChipGPU, var_range::VariableRangeCheckerChipGPU, Chip, }; -use openvm_cuda_backend::{base::DeviceMatrix, prelude::F, GpuBackend}; -use openvm_cuda_common::copy::MemCopyH2D; -use openvm_instructions::riscv::RV64_REGISTER_NUM_LIMBS; +use openvm_cuda_backend::{base::DeviceMatrix, GpuBackend}; use openvm_stark_backend::prover::AirProvingContext; -use crate::{ - adapters::{Rv64LoadStoreAdapterCols, Rv64LoadStoreAdapterRecord, RV64_BYTE_BITS}, - cuda_abi::load_sign_extend_cuda::tracegen, - LoadSignExtendCoreCols, LoadSignExtendCoreRecord, -}; +use crate::adapters::RV64_BYTE_BITS; + +fn unsupported_split_signed_load_ctx(arena: DenseRecordArena) -> AirProvingContext { + if arena.allocated().is_empty() { + return AirProvingContext::simple_no_pis(DeviceMatrix::dummy()); + } + unimplemented!("CUDA trace generation for split RV64 signed load chips is not implemented") +} #[derive(new)] -pub struct Rv64LoadSignExtendChipGpu { +pub struct Rv64LoadSignExtendByteChipGpu { pub range_checker: Arc, pub bitwise_lookup: Arc>, pub pointer_max_bits: usize, pub timestamp_max_bits: usize, } -impl Chip for Rv64LoadSignExtendChipGpu { +impl Chip for Rv64LoadSignExtendByteChipGpu { fn generate_proving_ctx(&self, arena: DenseRecordArena) -> AirProvingContext { - const RECORD_SIZE: usize = size_of::<( - Rv64LoadStoreAdapterRecord, - LoadSignExtendCoreRecord, - )>(); - let records = arena.allocated(); - if records.is_empty() { - return AirProvingContext::simple_no_pis(DeviceMatrix::dummy()); - } - debug_assert_eq!(records.len() % RECORD_SIZE, 0); + unsupported_split_signed_load_ctx(arena) + } +} - let trace_width = Rv64LoadStoreAdapterCols::::width() - + LoadSignExtendCoreCols::::width(); - let height = records.len() / RECORD_SIZE; - let padded_height = next_power_of_two_or_zero(height); - let device_ctx = &self.range_checker.device_ctx; +#[derive(new)] +pub struct Rv64LoadSignExtendHalfwordChipGpu { + pub range_checker: Arc, + pub pointer_max_bits: usize, + pub timestamp_max_bits: usize, +} - let d_records = tracing::info_span!("trace_gen.h2d_records") - .in_scope(|| records.to_device_on(device_ctx)) - .unwrap(); - let d_trace = DeviceMatrix::::with_capacity_on(padded_height, trace_width, device_ctx); +impl Chip for Rv64LoadSignExtendHalfwordChipGpu { + fn generate_proving_ctx(&self, arena: DenseRecordArena) -> AirProvingContext { + unsupported_split_signed_load_ctx(arena) + } +} - unsafe { - tracegen( - d_trace.buffer(), - padded_height, - trace_width, - &d_records, - self.pointer_max_bits, - &self.range_checker.count, - &self.bitwise_lookup.count, - self.timestamp_max_bits as u32, - device_ctx.stream.as_raw(), - ) - .unwrap(); - } +#[derive(new)] +pub struct Rv64LoadSignExtendWordChipGpu { + pub range_checker: Arc, + pub pointer_max_bits: usize, + pub timestamp_max_bits: usize, +} - AirProvingContext::simple_no_pis(d_trace) +impl Chip for Rv64LoadSignExtendWordChipGpu { + fn generate_proving_ctx(&self, arena: DenseRecordArena) -> AirProvingContext { + unsupported_split_signed_load_ctx(arena) } } diff --git a/extensions/riscv/circuit/src/load_sign_extend/execution.rs b/extensions/riscv/circuit/src/load_sign_extend/execution.rs deleted file mode 100644 index e75a81e9eb..0000000000 --- a/extensions/riscv/circuit/src/load_sign_extend/execution.rs +++ /dev/null @@ -1,344 +0,0 @@ -use std::{ - array, - borrow::{Borrow, BorrowMut}, - mem::size_of, -}; - -use openvm_circuit::{arch::*, system::memory::online::GuestMemory}; -use openvm_circuit_primitives_derive::AlignedBytesBorrow; -use openvm_instructions::{ - instruction::Instruction, - program::DEFAULT_PC_STEP, - riscv::{RV64_MEMORY_AS, RV64_REGISTER_AS, RV64_REGISTER_NUM_LIMBS}, - LocalOpcode, -}; -use openvm_riscv_transpiler::Rv64LoadStoreOpcode::{self, *}; -use openvm_stark_backend::p3_field::PrimeField32; - -use super::core::LoadSignExtendExecutor; -use crate::adapters::{rv64_address_add_imm, rv64_bytes_to_u32}; - -#[derive(AlignedBytesBorrow, Clone)] -#[repr(C)] -struct LoadSignExtendPreCompute { - imm_extended: u32, - a: u8, - b: u8, - e: u8, -} - -impl LoadSignExtendExecutor { - /// Return (local_opcode, enabled) - fn pre_compute_impl_rv64( - &self, - pc: u32, - inst: &Instruction, - data: &mut LoadSignExtendPreCompute, - ) -> Result<(Rv64LoadStoreOpcode, bool), StaticProgramError> { - let Instruction { - opcode, - a, - b, - c, - d, - e, - f, - g, - .. - } = inst; - - let local_opcode = Rv64LoadStoreOpcode::from_usize( - opcode.local_opcode_idx(Rv64LoadStoreOpcode::CLASS_OFFSET), - ); - let e_u32 = e.as_canonical_u32(); - if d.as_canonical_u32() != RV64_REGISTER_AS || e_u32 != RV64_MEMORY_AS { - return Err(StaticProgramError::InvalidInstruction(pc)); - } - debug_assert!( - matches!(local_opcode, LOADB | LOADH | LOADW), - "LoadSignExtendExecutor should only handle LOADB/LOADH/LOADW opcodes" - ); - - let imm = c.as_canonical_u32(); - let imm_sign = g.as_canonical_u32(); - let imm_extended = imm + imm_sign * 0xffff0000; - - *data = LoadSignExtendPreCompute { - imm_extended, - a: a.as_canonical_u32() as u8, - b: b.as_canonical_u32() as u8, - e: e_u32 as u8, - }; - let enabled = !f.is_zero(); - Ok((local_opcode, enabled)) - } -} - -macro_rules! dispatch { - ($execute_impl:ident, $local_opcode:ident, $enabled:ident) => { - match ($local_opcode, $enabled) { - (LOADB, true) => Ok($execute_impl::<_, _, LoadBOp, true>), - (LOADB, false) => Ok($execute_impl::<_, _, LoadBOp, false>), - (LOADH, true) => Ok($execute_impl::<_, _, LoadHOp, true>), - (LOADH, false) => Ok($execute_impl::<_, _, LoadHOp, false>), - (LOADW, true) => Ok($execute_impl::<_, _, LoadWOp, true>), - (LOADW, false) => Ok($execute_impl::<_, _, LoadWOp, false>), - _ => unreachable!(), - } - }; -} - -impl InterpreterExecutor - for LoadSignExtendExecutor -where - F: PrimeField32, -{ - fn pre_compute_size(&self) -> usize { - size_of::() - } - - #[cfg(not(feature = "tco"))] - #[inline(always)] - fn pre_compute( - &self, - pc: u32, - inst: &Instruction, - data: &mut [u8], - ) -> Result, StaticProgramError> { - let pre_compute: &mut LoadSignExtendPreCompute = data.borrow_mut(); - let (local_opcode, enabled) = self.pre_compute_impl_rv64(pc, inst, pre_compute)?; - dispatch!(execute_e1_handler, local_opcode, enabled) - } - - #[cfg(feature = "tco")] - fn handler( - &self, - pc: u32, - inst: &Instruction, - data: &mut [u8], - ) -> Result, StaticProgramError> - where - Ctx: ExecutionCtxTrait, - { - let pre_compute: &mut LoadSignExtendPreCompute = data.borrow_mut(); - let (local_opcode, enabled) = self.pre_compute_impl_rv64(pc, inst, pre_compute)?; - dispatch!(execute_e1_handler, local_opcode, enabled) - } -} - -impl InterpreterMeteredExecutor - for LoadSignExtendExecutor -where - F: PrimeField32, -{ - fn metered_pre_compute_size(&self) -> usize { - size_of::>() - } - - #[cfg(not(feature = "tco"))] - fn metered_pre_compute( - &self, - chip_idx: usize, - pc: u32, - inst: &Instruction, - data: &mut [u8], - ) -> Result, StaticProgramError> - where - Ctx: MeteredExecutionCtxTrait, - { - let pre_compute: &mut E2PreCompute = data.borrow_mut(); - pre_compute.chip_idx = chip_idx as u32; - let (local_opcode, enabled) = - self.pre_compute_impl_rv64(pc, inst, &mut pre_compute.data)?; - dispatch!(execute_e2_handler, local_opcode, enabled) - } - - #[cfg(feature = "tco")] - fn metered_handler( - &self, - chip_idx: usize, - pc: u32, - inst: &Instruction, - data: &mut [u8], - ) -> Result, StaticProgramError> - where - Ctx: MeteredExecutionCtxTrait, - { - let pre_compute: &mut E2PreCompute = data.borrow_mut(); - pre_compute.chip_idx = chip_idx as u32; - let (local_opcode, enabled) = - self.pre_compute_impl_rv64(pc, inst, &mut pre_compute.data)?; - dispatch!(execute_e2_handler, local_opcode, enabled) - } -} - -trait LoadSignExtOp { - fn compute_write_data( - read_data: &[u8; RV64_REGISTER_NUM_LIMBS], - shift_amount: u32, - ) -> Option<[u8; RV64_REGISTER_NUM_LIMBS]>; -} - -struct LoadBOp; -struct LoadHOp; -struct LoadWOp; - -impl LoadSignExtOp for LoadBOp { - #[inline(always)] - fn compute_write_data( - read_data: &[u8; RV64_REGISTER_NUM_LIMBS], - shift_amount: u32, - ) -> Option<[u8; RV64_REGISTER_NUM_LIMBS]> { - let byte = read_data[shift_amount as usize]; - Some(((byte as i8) as i64).to_le_bytes()) - } -} - -impl LoadSignExtOp for LoadHOp { - #[inline(always)] - fn compute_write_data( - read_data: &[u8; RV64_REGISTER_NUM_LIMBS], - shift_amount: u32, - ) -> Option<[u8; RV64_REGISTER_NUM_LIMBS]> { - if !shift_amount.is_multiple_of(2) { - return None; - } - let half: [u8; 2] = array::from_fn(|i| read_data[shift_amount as usize + i]); - Some((i16::from_le_bytes(half) as i64).to_le_bytes()) - } -} - -impl LoadSignExtOp for LoadWOp { - #[inline(always)] - fn compute_write_data( - read_data: &[u8; RV64_REGISTER_NUM_LIMBS], - shift_amount: u32, - ) -> Option<[u8; RV64_REGISTER_NUM_LIMBS]> { - if shift_amount != 0 && shift_amount != 4 { - return None; - } - let word: [u8; 4] = array::from_fn(|i| read_data[shift_amount as usize + i]); - Some((i32::from_le_bytes(word) as i64).to_le_bytes()) - } -} - -#[inline(always)] -unsafe fn execute_e12_impl< - F: PrimeField32, - CTX: ExecutionCtxTrait, - OP: LoadSignExtOp, - const ENABLED: bool, ->( - pre_compute: &LoadSignExtendPreCompute, - exec_state: &mut VmExecState, -) -> Result<(), ExecutionError> { - let pc = exec_state.pc(); - let rs1_bytes: [u8; RV64_REGISTER_NUM_LIMBS] = - exec_state.vm_read_bytes(RV64_REGISTER_AS, pre_compute.b as u32); - let rs1_val = rv64_bytes_to_u32(rs1_bytes); - let addr = rv64_address_add_imm(rs1_val, pre_compute.imm_extended); - debug_assert!((addr as usize) < RV64_MEMORY_BYTES); - let ptr_val = addr as u32; - - let shift_amount = ptr_val % RV64_REGISTER_NUM_LIMBS as u32; - let ptr_val = ptr_val - shift_amount; // aligned ptr - - let read_data: [u8; RV64_REGISTER_NUM_LIMBS] = - exec_state.vm_read_bytes(pre_compute.e as u32, ptr_val); - - let write_data = - OP::compute_write_data(&read_data, shift_amount).ok_or(ExecutionError::Fail { - pc, - msg: "Invalid LoadSignExtendOp", - })?; - - if ENABLED { - exec_state.vm_write_bytes(RV64_REGISTER_AS, pre_compute.a as u32, &write_data); - } - - exec_state.set_pc(pc.wrapping_add(DEFAULT_PC_STEP)); - - Ok(()) -} - -#[create_handler] -#[inline(always)] -unsafe fn execute_e1_impl< - F: PrimeField32, - CTX: ExecutionCtxTrait, - OP: LoadSignExtOp, - const ENABLED: bool, ->( - pre_compute: *const u8, - exec_state: &mut VmExecState, -) -> Result<(), ExecutionError> { - let pre_compute: &LoadSignExtendPreCompute = - std::slice::from_raw_parts(pre_compute, size_of::()).borrow(); - execute_e12_impl::(pre_compute, exec_state) -} - -#[create_handler] -#[inline(always)] -unsafe fn execute_e2_impl< - F: PrimeField32, - CTX: MeteredExecutionCtxTrait, - OP: LoadSignExtOp, - const ENABLED: bool, ->( - pre_compute: *const u8, - exec_state: &mut VmExecState, -) -> Result<(), ExecutionError> { - let pre_compute: &E2PreCompute = std::slice::from_raw_parts( - pre_compute, - size_of::>(), - ) - .borrow(); - exec_state - .ctx - .on_height_change(pre_compute.chip_idx as usize, 1); - execute_e12_impl::(&pre_compute.data, exec_state) -} - -#[cfg(test)] -mod tests { - use openvm_circuit::arch::StaticProgramError; - use openvm_instructions::{ - instruction::Instruction, riscv::RV64_REGISTER_AS, LocalOpcode, DEFERRAL_AS, - }; - use openvm_riscv_transpiler::Rv64LoadStoreOpcode::LOADW; - use openvm_stark_sdk::p3_baby_bear::BabyBear; - - use super::LoadSignExtendPreCompute; - use crate::{adapters::Rv64LoadStoreAdapterExecutor, Rv64LoadSignExtendExecutor}; - - #[test] - fn precompute_enforces_address_space_domain() { - const PC: u32 = 0x100; - - let executor = Rv64LoadSignExtendExecutor::new(Rv64LoadStoreAdapterExecutor::new(29)); - let inst = Instruction::::from_usize( - LOADW.global_opcode(), - [ - 8, - 16, - 0, - RV64_REGISTER_AS as usize, - DEFERRAL_AS as usize, - 1, - 0, - ], - ); - let mut data = LoadSignExtendPreCompute { - imm_extended: 0, - a: 0, - b: 0, - e: 0, - }; - - let err = executor - .pre_compute_impl_rv64(PC, &inst, &mut data) - .expect_err("load address-space domain should be enforced"); - - assert!(matches!(err, StaticProgramError::InvalidInstruction(PC))); - } -} diff --git a/extensions/riscv/circuit/src/load_sign_extend/halfword/mod.rs b/extensions/riscv/circuit/src/load_sign_extend/halfword/mod.rs new file mode 100644 index 0000000000..c921d160ec --- /dev/null +++ b/extensions/riscv/circuit/src/load_sign_extend/halfword/mod.rs @@ -0,0 +1,28 @@ +use openvm_circuit::arch::{VmAirWrapper, VmChipWrapper}; + +use crate::{ + adapters::{Rv64LoadStoreAdapterAir, Rv64LoadStoreAdapterExecutor}, + load_sign_extend::aligned::core::{LoadSignExtendAlignedCoreAir, LoadSignExtendAlignedFiller}, + loadstore::common::{LoadStoreExecutor, KIND_HALFWORD}, +}; + +pub const LOAD_SIGN_EXTEND_HALFWORD_CASES: usize = 4; +pub const LOAD_SIGN_EXTEND_HALFWORD_SELECTOR_WIDTH: usize = 2; + +pub type LoadSignExtendHalfwordCoreAir = LoadSignExtendAlignedCoreAir< + KIND_HALFWORD, + LOAD_SIGN_EXTEND_HALFWORD_CASES, + LOAD_SIGN_EXTEND_HALFWORD_SELECTOR_WIDTH, +>; +pub type LoadSignExtendHalfwordFiller = LoadSignExtendAlignedFiller< + crate::adapters::Rv64LoadStoreAdapterFiller, + KIND_HALFWORD, + LOAD_SIGN_EXTEND_HALFWORD_CASES, + LOAD_SIGN_EXTEND_HALFWORD_SELECTOR_WIDTH, +>; + +pub type Rv64LoadSignExtendHalfwordAir = + VmAirWrapper; +pub type Rv64LoadSignExtendHalfwordExecutor = + LoadStoreExecutor; +pub type Rv64LoadSignExtendHalfwordChip = VmChipWrapper; diff --git a/extensions/riscv/circuit/src/load_sign_extend/mod.rs b/extensions/riscv/circuit/src/load_sign_extend/mod.rs index e0eff8846e..4cf81cfcbc 100644 --- a/extensions/riscv/circuit/src/load_sign_extend/mod.rs +++ b/extensions/riscv/circuit/src/load_sign_extend/mod.rs @@ -1,33 +1,16 @@ -use openvm_circuit::arch::{VmAirWrapper, VmChipWrapper}; -use openvm_instructions::riscv::{RV64_BYTE_BITS, RV64_REGISTER_NUM_LIMBS}; +pub(crate) mod aligned; +mod byte; +mod halfword; +mod word; -use crate::adapters::{Rv64LoadStoreAdapterAir, Rv64LoadStoreAdapterExecutor}; - -mod core; -mod execution; -pub use core::*; +pub use byte::*; +pub use halfword::*; +pub use word::*; #[cfg(feature = "cuda")] mod cuda; #[cfg(feature = "cuda")] pub use cuda::*; -#[cfg(feature = "aot")] -mod aot; #[cfg(test)] mod tests; - -pub type Rv64LoadSignExtendAir = VmAirWrapper< - Rv64LoadStoreAdapterAir, - LoadSignExtendCoreAir, ->; -pub type Rv64LoadSignExtendExecutor = - LoadSignExtendExecutor; -pub type Rv64LoadSignExtendChip = VmChipWrapper< - F, - LoadSignExtendFiller< - crate::adapters::Rv64LoadStoreAdapterFiller, - RV64_REGISTER_NUM_LIMBS, - RV64_BYTE_BITS, - >, ->; diff --git a/extensions/riscv/circuit/src/load_sign_extend/tests.rs b/extensions/riscv/circuit/src/load_sign_extend/tests.rs index 1c19f4b6d3..22258518fd 100644 --- a/extensions/riscv/circuit/src/load_sign_extend/tests.rs +++ b/extensions/riscv/circuit/src/load_sign_extend/tests.rs @@ -6,22 +6,15 @@ use openvm_circuit::{ memory::gen_pointer, TestBuilder, TestChipHarness, VmChipTestBuilder, BITWISE_OP_LOOKUP_BUS, }, - Arena, ExecutionBridge, PreflightExecutor, + Arena, MemoryConfig, PreflightExecutor, BLOCK_FE_WIDTH, }, - system::memory::{offline_checker::MemoryBridge, SharedMemoryHelper}, + system::memory::merkle::public_values::PUBLIC_VALUES_AS, }; -use openvm_circuit_primitives::{ - bitwise_op_lookup::{ - BitwiseOperationLookupAir, BitwiseOperationLookupBus, BitwiseOperationLookupChip, - SharedBitwiseOperationLookupChip, - }, - var_range::VariableRangeCheckerChip, -}; -use openvm_instructions::{ - instruction::Instruction, - riscv::{RV64_BYTE_BITS, RV64_REGISTER_NUM_LIMBS}, - LocalOpcode, +use openvm_circuit_primitives::bitwise_op_lookup::{ + BitwiseOperationLookupAir, BitwiseOperationLookupBus, BitwiseOperationLookupChip, + SharedBitwiseOperationLookupChip, }; +use openvm_instructions::{instruction::Instruction, riscv::RV64_REGISTER_AS, LocalOpcode}; use openvm_riscv_transpiler::Rv64LoadStoreOpcode::{self, *}; use openvm_stark_backend::{ p3_air::BaseAir, @@ -35,146 +28,193 @@ use openvm_stark_backend::{ use openvm_stark_sdk::{p3_baby_bear::BabyBear, utils::create_seeded_rng}; use rand::{rngs::StdRng, Rng}; use test_case::test_case; -#[cfg(feature = "cuda")] -use { - crate::{ - adapters::Rv64LoadStoreAdapterRecord, LoadSignExtendCoreRecord, Rv64LoadSignExtendChipGpu, + +use super::{ + byte::{ + LoadSignExtendByteCoreAir, LoadSignExtendByteCoreCols, LoadSignExtendByteFiller, + Rv64LoadSignExtendByteAir, Rv64LoadSignExtendByteChip, Rv64LoadSignExtendByteExecutor, }, - openvm_circuit::arch::{ - testing::{ - default_bitwise_lookup_bus, default_var_range_checker_bus, dummy_range_checker, - GpuChipTestBuilder, GpuTestChipHarness, - }, - EmptyAdapterCoreLayout, + halfword::{ + LoadSignExtendHalfwordCoreAir, LoadSignExtendHalfwordFiller, Rv64LoadSignExtendHalfwordAir, + Rv64LoadSignExtendHalfwordChip, Rv64LoadSignExtendHalfwordExecutor, + LOAD_SIGN_EXTEND_HALFWORD_SELECTOR_WIDTH, + }, + word::{ + LoadSignExtendWordCoreAir, LoadSignExtendWordFiller, Rv64LoadSignExtendWordAir, + Rv64LoadSignExtendWordChip, Rv64LoadSignExtendWordExecutor, + LOAD_SIGN_EXTEND_WORD_SELECTOR_WIDTH, }, }; - -use super::{run_write_data_sign_extend, LoadSignExtendCoreAir}; use crate::{ adapters::{ - rv64_bytes_to_u32, Rv64LoadStoreAdapterAir, Rv64LoadStoreAdapterExecutor, - Rv64LoadStoreAdapterFiller, + rv64_bytes_to_u16_block, rv64_bytes_to_u32, rv64_u16_block_to_bytes, sign_extend_imm16, + Rv64LoadStoreAdapterAir, Rv64LoadStoreAdapterExecutor, Rv64LoadStoreAdapterFiller, + RV64_BYTE_BITS, }, - load_sign_extend::LoadSignExtendCoreCols, - LoadSignExtendFiller, Rv64LoadSignExtendAir, Rv64LoadSignExtendChip, - Rv64LoadSignExtendExecutor, + load_sign_extend::aligned::core::LoadSignExtendAlignedCoreCols, + loadstore::common::run_write_data, }; const IMM_BITS: usize = 16; const MAX_INS_CAPACITY: usize = 128; -type Harness = TestChipHarness< - F, - Rv64LoadSignExtendExecutor, - Rv64LoadSignExtendAir, - Rv64LoadSignExtendChip, ->; type F = BabyBear; -fn create_harness_fields( - memory_bridge: MemoryBridge, - execution_bridge: ExecutionBridge, - range_checker_chip: Arc, - bitwise_chip: SharedBitwiseOperationLookupChip, - memory_helper: SharedMemoryHelper, - address_bits: usize, -) -> ( - Rv64LoadSignExtendAir, - Rv64LoadSignExtendExecutor, - Rv64LoadSignExtendChip, -) { - let air = Rv64LoadSignExtendAir::new( - Rv64LoadStoreAdapterAir::new( - memory_bridge, - execution_bridge, - range_checker_chip.bus(), - address_bits, - ), - LoadSignExtendCoreAir::new(range_checker_chip.bus(), bitwise_chip.bus()), - ); - let executor = Rv64LoadSignExtendExecutor::new(Rv64LoadStoreAdapterExecutor::new(address_bits)); - let chip = Rv64LoadSignExtendChip::::new( - LoadSignExtendFiller::new( - Rv64LoadStoreAdapterFiller::new(address_bits, range_checker_chip.clone()), - range_checker_chip, - bitwise_chip, - ), - memory_helper, - ); - (air, executor, chip) -} +type ByteHarness = TestChipHarness< + F, + Rv64LoadSignExtendByteExecutor, + Rv64LoadSignExtendByteAir, + Rv64LoadSignExtendByteChip, +>; +type HalfwordHarness = TestChipHarness< + F, + Rv64LoadSignExtendHalfwordExecutor, + Rv64LoadSignExtendHalfwordAir, + Rv64LoadSignExtendHalfwordChip, +>; +type WordHarness = TestChipHarness< + F, + Rv64LoadSignExtendWordExecutor, + Rv64LoadSignExtendWordAir, + Rv64LoadSignExtendWordChip, +>; -fn create_test_chip( +fn create_byte_harness( tester: &mut VmChipTestBuilder, ) -> ( - Harness, + ByteHarness, ( BitwiseOperationLookupAir, SharedBitwiseOperationLookupChip, ), ) { + let range_checker = tester.range_checker(); let bitwise_bus = BitwiseOperationLookupBus::new(BITWISE_OP_LOOKUP_BUS); let bitwise_chip = Arc::new(BitwiseOperationLookupChip::::new( bitwise_bus, )); - let (air, executor, chip) = create_harness_fields( - tester.memory_bridge(), - tester.execution_bridge(), - tester.range_checker(), - bitwise_chip.clone(), + let air = Rv64LoadSignExtendByteAir::new( + Rv64LoadStoreAdapterAir::new( + tester.memory_bridge(), + tester.execution_bridge(), + range_checker.bus(), + tester.address_bits(), + ), + LoadSignExtendByteCoreAir::new( + Rv64LoadStoreOpcode::CLASS_OFFSET, + bitwise_chip.bus(), + range_checker.bus(), + ), + ); + let executor = Rv64LoadSignExtendByteExecutor::new( + Rv64LoadStoreAdapterExecutor::new(tester.address_bits()), + Rv64LoadStoreOpcode::CLASS_OFFSET, + ); + let chip = Rv64LoadSignExtendByteChip::::new( + LoadSignExtendByteFiller::new( + Rv64LoadStoreAdapterFiller::new(tester.address_bits(), range_checker.clone()), + Rv64LoadStoreOpcode::CLASS_OFFSET, + bitwise_chip.clone(), + range_checker, + ), tester.memory_helper(), - tester.address_bits(), ); ( - Harness::with_capacity(executor, air, chip, MAX_INS_CAPACITY), + ByteHarness::with_capacity(executor, air, chip, MAX_INS_CAPACITY), (bitwise_chip.air, bitwise_chip), ) } +fn create_halfword_harness(tester: &mut VmChipTestBuilder) -> HalfwordHarness { + let range_checker = tester.range_checker(); + let air = Rv64LoadSignExtendHalfwordAir::new( + Rv64LoadStoreAdapterAir::new( + tester.memory_bridge(), + tester.execution_bridge(), + range_checker.bus(), + tester.address_bits(), + ), + LoadSignExtendHalfwordCoreAir::new(Rv64LoadStoreOpcode::CLASS_OFFSET, range_checker.bus()), + ); + let executor = Rv64LoadSignExtendHalfwordExecutor::new( + Rv64LoadStoreAdapterExecutor::new(tester.address_bits()), + Rv64LoadStoreOpcode::CLASS_OFFSET, + ); + let chip = Rv64LoadSignExtendHalfwordChip::::new( + LoadSignExtendHalfwordFiller::new( + Rv64LoadStoreAdapterFiller::new(tester.address_bits(), range_checker.clone()), + Rv64LoadStoreOpcode::CLASS_OFFSET, + range_checker, + ), + tester.memory_helper(), + ); + HalfwordHarness::with_capacity(executor, air, chip, MAX_INS_CAPACITY) +} + +fn create_word_harness(tester: &mut VmChipTestBuilder) -> WordHarness { + let range_checker = tester.range_checker(); + let air = Rv64LoadSignExtendWordAir::new( + Rv64LoadStoreAdapterAir::new( + tester.memory_bridge(), + tester.execution_bridge(), + range_checker.bus(), + tester.address_bits(), + ), + LoadSignExtendWordCoreAir::new(Rv64LoadStoreOpcode::CLASS_OFFSET, range_checker.bus()), + ); + let executor = Rv64LoadSignExtendWordExecutor::new( + Rv64LoadStoreAdapterExecutor::new(tester.address_bits()), + Rv64LoadStoreOpcode::CLASS_OFFSET, + ); + let chip = Rv64LoadSignExtendWordChip::::new( + LoadSignExtendWordFiller::new( + Rv64LoadStoreAdapterFiller::new(tester.address_bits(), range_checker.clone()), + Rv64LoadStoreOpcode::CLASS_OFFSET, + range_checker, + ), + tester.memory_helper(), + ); + WordHarness::with_capacity(executor, air, chip, MAX_INS_CAPACITY) +} + #[allow(clippy::too_many_arguments)] fn set_and_execute>( - tester: &mut impl TestBuilder, + tester: &mut VmChipTestBuilder, executor: &mut E, arena: &mut RA, rng: &mut StdRng, opcode: Rv64LoadStoreOpcode, - read_data: Option<[u8; RV64_REGISTER_NUM_LIMBS]>, - rs1: Option<[u8; RV64_REGISTER_NUM_LIMBS]>, + rs1: Option<[u8; 8]>, imm: Option, imm_sign: Option, ) { - let imm = imm.unwrap_or(rng.random_range(0..(1 << IMM_BITS))); - let imm_sign = imm_sign.unwrap_or(rng.random_range(0..2)); - let imm_ext = imm + imm_sign * (0xffff0000); - + let imm = imm.unwrap_or_else(|| rng.random_range(0..(1 << IMM_BITS))); + let imm_sign = imm_sign.unwrap_or_else(|| rng.random_range(0..2)); + let imm_ext = sign_extend_imm16(imm, imm_sign); let alignment = match opcode { - LOADB => 0, - LOADH => 1, LOADW => 2, - _ => unreachable!(), + LOADH => 1, + LOADB => 0, + _ => unreachable!("signed load test only supports LOADB/LOADH/LOADW"), }; - - let ptr_val: u32 = rng.random_range(0..(1 << (tester.address_bits() - alignment))) << alignment; - // rs1 is 8 bytes, but only low 4 bytes used for address + let ptr_val: u32 = + rng.random_range(0..(1u32 << (tester.address_bits() - alignment))) << alignment; let rs1 = rs1.unwrap_or_else(|| { let low4 = ptr_val.wrapping_sub(imm_ext).to_le_bytes(); [low4[0], low4[1], low4[2], low4[3], 0, 0, 0, 0] }); let ptr_val = imm_ext.wrapping_add(rv64_bytes_to_u32(rs1)); + let shift_amount = ptr_val % 8; let a = gen_pointer(rng, 8); let b = gen_pointer(rng, 8); - - let shift_amount = ptr_val % 8; - tester.write_bytes(1, b, rs1.map(F::from_u8)); - - let some_prev_data: [F; RV64_REGISTER_NUM_LIMBS] = if a != 0 { + let read_data: [u8; 8] = array::from_fn(|_| rng.random()); + let prev_data: [F; 8] = if a != 0 { array::from_fn(|_| F::from_u8(rng.random())) } else { - [F::ZERO; RV64_REGISTER_NUM_LIMBS] + [F::ZERO; 8] }; - let read_data: [u8; RV64_REGISTER_NUM_LIMBS] = - read_data.unwrap_or(array::from_fn(|_| rng.random())); - tester.write_bytes(1, a, some_prev_data); + tester.write_bytes(RV64_REGISTER_AS as usize, b, rs1.map(F::from_u8)); + tester.write_bytes(RV64_REGISTER_AS as usize, a, prev_data); tester.write_bytes( 2, (ptr_val - shift_amount) as usize, @@ -190,7 +230,7 @@ fn set_and_execute>( a, b, imm as usize, - 1, + RV64_REGISTER_AS as usize, 2, (a != 0) as usize, imm_sign as usize, @@ -198,32 +238,37 @@ fn set_and_execute>( ), ); - let write_data = run_write_data_sign_extend::( + let expected = run_write_data( opcode, - read_data, + rv64_bytes_to_u16_block(read_data), + [0; BLOCK_FE_WIDTH], shift_amount as usize, ); if a != 0 { - assert_eq!(write_data.map(F::from_u8), tester.read_bytes::<8>(1, a)); + assert_eq!( + rv64_u16_block_to_bytes(expected).map(F::from_u8), + tester.read_bytes::<8>(RV64_REGISTER_AS as usize, a) + ); } else { - assert_eq!([F::ZERO; 8], tester.read_bytes::<8>(1, a)); + assert_eq!( + [F::ZERO; 8], + tester.read_bytes::<8>(RV64_REGISTER_AS as usize, a) + ); } } -/////////////////////////////////////////////////////////////////////////////////////// -/// POSITIVE TESTS -/// -/// Randomly generate computations and execute, ensuring that the generated trace -/// passes all constraints. -/////////////////////////////////////////////////////////////////////////////////////// +fn memory_config_for() -> MemoryConfig { + let mut mem_config = MemoryConfig::default(); + 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; + mem_config +} + #[test_case(LOADB, 100)] -#[test_case(LOADH, 100)] -#[test_case(LOADW, 100)] -fn rand_load_sign_extend_test(opcode: Rv64LoadStoreOpcode, num_ops: usize) { +fn rand_load_sign_extend_byte_test(opcode: Rv64LoadStoreOpcode, num_ops: usize) { let mut rng = create_seeded_rng(); - let mut tester = VmChipTestBuilder::default(); - - let (mut harness, bitwise) = create_test_chip(&mut tester); + let mut tester = VmChipTestBuilder::from_config(memory_config_for()); + let (mut harness, bitwise) = create_byte_harness(&mut tester); for _ in 0..num_ops { set_and_execute( &mut tester, @@ -234,481 +279,235 @@ fn rand_load_sign_extend_test(opcode: Rv64LoadStoreOpcode, num_ops: usize) { None, None, None, + ); + } + tester + .build() + .load(harness) + .load_periphery(bitwise) + .finalize() + .simple_test() + .unwrap(); +} + +#[test_case(LOADH, 100)] +fn rand_load_sign_extend_halfword_test(opcode: Rv64LoadStoreOpcode, num_ops: usize) { + let mut rng = create_seeded_rng(); + let mut tester = VmChipTestBuilder::from_config(memory_config_for()); + let mut harness = create_halfword_harness(&mut tester); + for _ in 0..num_ops { + set_and_execute( + &mut tester, + &mut harness.executor, + &mut harness.arena, + &mut rng, + opcode, + None, + None, None, ); } + tester + .build() + .load(harness) + .finalize() + .simple_test() + .unwrap(); +} - let tester = tester +#[test_case(LOADW, 100)] +fn rand_load_sign_extend_word_test(opcode: Rv64LoadStoreOpcode, num_ops: usize) { + let mut rng = create_seeded_rng(); + let mut tester = VmChipTestBuilder::from_config(memory_config_for()); + let mut harness = create_word_harness(&mut tester); + for _ in 0..num_ops { + set_and_execute( + &mut tester, + &mut harness.executor, + &mut harness.arena, + &mut rng, + opcode, + None, + None, + None, + ); + } + tester .build() .load(harness) - .load_periphery(bitwise) - .finalize(); - tester.simple_test().expect("Verification failed"); + .finalize() + .simple_test() + .unwrap(); } #[test] fn positive_loadb_shift7_test() { let mut rng = create_seeded_rng(); - let mut tester = VmChipTestBuilder::default(); - let (mut harness, bitwise) = create_test_chip(&mut tester); - + let mut tester = VmChipTestBuilder::from_config(memory_config_for()); + let (mut harness, bitwise) = create_byte_harness(&mut tester); set_and_execute( &mut tester, &mut harness.executor, &mut harness.arena, &mut rng, LOADB, - None, Some([7, 0, 0, 0, 0, 0, 0, 0]), Some(0), Some(0), ); - - let tester = tester + tester .build() .load(harness) .load_periphery(bitwise) - .finalize(); - tester.simple_test().expect("Verification failed"); + .finalize() + .simple_test() + .unwrap(); } #[test] fn positive_loadh_shift6_test() { let mut rng = create_seeded_rng(); - let mut tester = VmChipTestBuilder::default(); - let (mut harness, bitwise) = create_test_chip(&mut tester); - + let mut tester = VmChipTestBuilder::from_config(memory_config_for()); + let mut harness = create_halfword_harness(&mut tester); set_and_execute( &mut tester, &mut harness.executor, &mut harness.arena, &mut rng, LOADH, - None, Some([6, 0, 0, 0, 0, 0, 0, 0]), Some(0), Some(0), ); - - let tester = tester + tester .build() .load(harness) - .load_periphery(bitwise) - .finalize(); - tester.simple_test().expect("Verification failed"); + .finalize() + .simple_test() + .unwrap(); } #[test] fn positive_loadw_shift4_test() { let mut rng = create_seeded_rng(); - let mut tester = VmChipTestBuilder::default(); - let (mut harness, bitwise) = create_test_chip(&mut tester); - + let mut tester = VmChipTestBuilder::from_config(memory_config_for()); + let mut harness = create_word_harness(&mut tester); set_and_execute( &mut tester, &mut harness.executor, &mut harness.arena, &mut rng, LOADW, - None, Some([4, 0, 0, 0, 0, 0, 0, 0]), Some(0), Some(0), ); - - let tester = tester + tester .build() .load(harness) - .load_periphery(bitwise) - .finalize(); - tester.simple_test().expect("Verification failed"); -} - -////////////////////////////////////////////////////////////////////////////////////// -// NEGATIVE TESTS -// -// Given a fake trace of a single operation, setup a chip and run the test. We replace -// part of the trace and check that the chip throws the expected error. -////////////////////////////////////////////////////////////////////////////////////// - -#[derive(Clone, Copy, Default, PartialEq)] -struct LoadSignExtPrankValues { - data_most_sig_bit: Option, - shift_most_sig_bit: Option, - opcode_flags: Option<[bool; 7]>, + .finalize() + .simple_test() + .unwrap(); } -#[allow(clippy::too_many_arguments)] -fn run_negative_load_sign_extend_test( - opcode: Rv64LoadStoreOpcode, - read_data: Option<[u8; RV64_REGISTER_NUM_LIMBS]>, - rs1: Option<[u8; RV64_REGISTER_NUM_LIMBS]>, - imm: Option, - imm_sign: Option, - prank_vals: LoadSignExtPrankValues, - _interaction_error: bool, -) { +fn assert_pranked_byte_fails(prank: impl Fn(&mut LoadSignExtendByteCoreCols)) { let mut rng = create_seeded_rng(); - let mut tester = VmChipTestBuilder::default(); - let (mut harness, bitwise) = create_test_chip(&mut tester); - + let mut tester = VmChipTestBuilder::from_config(memory_config_for()); + let (mut harness, bitwise) = create_byte_harness(&mut tester); set_and_execute( &mut tester, &mut harness.executor, &mut harness.arena, &mut rng, - opcode, - read_data, - rs1, - imm, - imm_sign, + LOADB, + None, + None, + None, ); - let adapter_width = BaseAir::::width(&harness.air.adapter); - let modify_trace = |trace: &mut DenseMatrix| { + let modify_trace = |trace: &mut DenseMatrix| { let mut trace_row = trace.row_slice(0).unwrap().to_vec(); let (_, core_row) = trace_row.split_at_mut(adapter_width); - - let core_cols: &mut LoadSignExtendCoreCols = - core_row.borrow_mut(); - if let Some(shifted_read_data) = read_data { - core_cols.shifted_read_data = shifted_read_data.map(F::from_u8); - } - if let Some(data_most_sig_bit) = prank_vals.data_most_sig_bit { - core_cols.data_most_sig_bit = F::from_u32(data_most_sig_bit); - } - if let Some(shift_most_sig_bit) = prank_vals.shift_most_sig_bit { - core_cols.shift_most_sig_bit = F::from_u32(shift_most_sig_bit); - } - if let Some(opcode_flags) = prank_vals.opcode_flags { - core_cols.opcode_loadb_flag0 = F::from_bool(opcode_flags[0]); - core_cols.opcode_loadb_flag1 = F::from_bool(opcode_flags[1]); - core_cols.opcode_loadb_flag2 = F::from_bool(opcode_flags[2]); - core_cols.opcode_loadb_flag3 = F::from_bool(opcode_flags[3]); - core_cols.opcode_loadh_flag0 = F::from_bool(opcode_flags[4]); - core_cols.opcode_loadh_flag2 = F::from_bool(opcode_flags[5]); - core_cols.opcode_loadw_flag = F::from_bool(opcode_flags[6]); - } - + prank(core_row.borrow_mut()); *trace = RowMajorMatrix::new(trace_row, trace.width()); }; - disable_debug_builder(); - let tester = tester + tester .build() .load_and_prank_trace(harness, modify_trace) .load_periphery(bitwise) - .finalize(); - tester + .finalize() .simple_test() - .expect_err("Expected verification to fail, but it passed"); + .expect_err("pranked signed byte load trace should fail"); } -#[test] -fn loadstore_negative_tests() { - // Pranking shifted_read_data + data_most_sig_bit breaks memory read interaction - run_negative_load_sign_extend_test( - LOADB, - Some([233, 187, 145, 238, 12, 55, 200, 99]), - None, - None, - None, - LoadSignExtPrankValues { - data_most_sig_bit: Some(0), - ..Default::default() - }, - true, - ); - - // Pranking shift_most_sig_bit breaks the read_data unrotation → memory interaction error - run_negative_load_sign_extend_test( +fn assert_pranked_halfword_fails( + prank: impl Fn(&mut LoadSignExtendAlignedCoreCols), +) { + let mut rng = create_seeded_rng(); + let mut tester = VmChipTestBuilder::from_config(memory_config_for()); + let mut harness = create_halfword_harness(&mut tester); + set_and_execute( + &mut tester, + &mut harness.executor, + &mut harness.arena, + &mut rng, LOADH, None, - Some([202, 109, 183, 26, 0, 0, 0, 0]), - Some(31212), - None, - LoadSignExtPrankValues { - shift_most_sig_bit: Some(0), - ..Default::default() - }, - true, - ); - - // Pranking opcode_flags to wrong value → execution bridge interaction error - run_negative_load_sign_extend_test( - LOADB, None, - Some([250, 132, 77, 5, 0, 0, 0, 0]), - Some(47741), None, - LoadSignExtPrankValues { - opcode_flags: Some([true, false, false, false, false, false, false]), - ..Default::default() - }, - true, - ); -} - -/////////////////////////////////////////////////////////////////////////////////////// -/// SANITY TESTS -/// -/// Ensure that solve functions produce the correct results. -/////////////////////////////////////////////////////////////////////////////////////// - -#[test] -fn solve_loadh_extend_sign_sanity_test() { - let read_data = [34, 159, 237, 151, 100, 200, 50, 25]; - let write_data0 = - run_write_data_sign_extend::(LOADH, read_data, 0); - let write_data2 = - run_write_data_sign_extend::(LOADH, read_data, 2); - let write_data4 = - run_write_data_sign_extend::(LOADH, read_data, 4); - let write_data6 = - run_write_data_sign_extend::(LOADH, read_data, 6); - - assert_eq!(write_data0, [34, 159, 255, 255, 255, 255, 255, 255]); - assert_eq!(write_data2, [237, 151, 255, 255, 255, 255, 255, 255]); - assert_eq!(write_data4, [100, 200, 255, 255, 255, 255, 255, 255]); - assert_eq!(write_data6, [50, 25, 0, 0, 0, 0, 0, 0]); -} - -#[test] -fn solve_loadh_extend_zero_sanity_test() { - let read_data = [34, 121, 237, 97, 10, 20, 30, 40]; - let write_data0 = - run_write_data_sign_extend::(LOADH, read_data, 0); - let write_data2 = - run_write_data_sign_extend::(LOADH, read_data, 2); - let write_data4 = - run_write_data_sign_extend::(LOADH, read_data, 4); - let write_data6 = - run_write_data_sign_extend::(LOADH, read_data, 6); - - assert_eq!(write_data0, [34, 121, 0, 0, 0, 0, 0, 0]); - assert_eq!(write_data2, [237, 97, 0, 0, 0, 0, 0, 0]); - assert_eq!(write_data4, [10, 20, 0, 0, 0, 0, 0, 0]); - assert_eq!(write_data6, [30, 40, 0, 0, 0, 0, 0, 0]); -} - -#[test] -fn solve_loadb_extend_sign_sanity_test() { - let read_data = [45, 82, 99, 127, 200, 150, 180, 210]; - for shift in 0..8 { - let write_data = run_write_data_sign_extend::( - LOADB, read_data, shift, - ); - let byte = read_data[shift]; - let expected = (byte as i8 as i64).to_le_bytes(); - assert_eq!(write_data, expected, "LOADB shift={shift}"); - } -} - -#[test] -fn solve_loadb_extend_zero_sanity_test() { - let read_data = [173, 210, 227, 255, 128, 250, 200, 190]; - for shift in 0..8 { - let write_data = run_write_data_sign_extend::( - LOADB, read_data, shift, - ); - let byte = read_data[shift]; - let expected = (byte as i8 as i64).to_le_bytes(); - assert_eq!(write_data, expected, "LOADB shift={shift}"); - } -} - -#[test] -fn solve_loadw_extend_sign_sanity_test() { - // shift=0: word = [0x01, 0x02, 0x03, 0x84] => 0x84030201 (negative) - let read_data = [0x01, 0x02, 0x03, 0x84, 0xAA, 0xBB, 0xCC, 0xDD]; - let write_data0 = - run_write_data_sign_extend::(LOADW, read_data, 0); - assert_eq!( - write_data0, - [0x01, 0x02, 0x03, 0x84, 0xFF, 0xFF, 0xFF, 0xFF] - ); - - // shift=4: word = [0xAA, 0xBB, 0xCC, 0xDD] => 0xDDCCBBAA (negative) - let write_data4 = - run_write_data_sign_extend::(LOADW, read_data, 4); - assert_eq!( - write_data4, - [0xAA, 0xBB, 0xCC, 0xDD, 0xFF, 0xFF, 0xFF, 0xFF] - ); -} - -#[test] -fn solve_loadw_extend_zero_sanity_test() { - // shift=0: word = [0x01, 0x02, 0x03, 0x04] => 0x04030201 (positive) - let read_data = [0x01, 0x02, 0x03, 0x04, 0xAA, 0xBB, 0xCC, 0xDD]; - let write_data0 = - run_write_data_sign_extend::(LOADW, read_data, 0); - assert_eq!( - write_data0, - [0x01, 0x02, 0x03, 0x04, 0x00, 0x00, 0x00, 0x00] - ); - - // shift=4: word = [0xAA, 0xBB, 0xCC, 0x7D] => 0x7DCCBBAA (positive) - let read_data2 = [0x01, 0x02, 0x03, 0x04, 0xAA, 0xBB, 0xCC, 0x7D]; - let write_data4 = - run_write_data_sign_extend::(LOADW, read_data2, 4); - assert_eq!( - write_data4, - [0xAA, 0xBB, 0xCC, 0x7D, 0x00, 0x00, 0x00, 0x00] ); + let adapter_width = BaseAir::::width(&harness.air.adapter); + let modify_trace = |trace: &mut DenseMatrix| { + let mut trace_row = trace.row_slice(0).unwrap().to_vec(); + let (_, core_row) = trace_row.split_at_mut(adapter_width); + prank(core_row.borrow_mut()); + *trace = RowMajorMatrix::new(trace_row, trace.width()); + }; + disable_debug_builder(); + tester + .build() + .load_and_prank_trace(harness, modify_trace) + .finalize() + .simple_test() + .expect_err("pranked signed halfword load trace should fail"); } -#[test] -#[should_panic(expected = "LOADW requires 4-byte aligned shift")] -fn solve_loadw_rejects_shift_2() { - let read_data = [0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08]; - run_write_data_sign_extend::(LOADW, read_data, 2); -} - -#[test] -#[should_panic(expected = "LOADW requires 4-byte aligned shift")] -fn solve_loadw_rejects_shift_6() { - let read_data = [0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08]; - run_write_data_sign_extend::(LOADW, read_data, 6); -} - -#[test] -#[should_panic(expected = "LOADH requires 2-byte aligned shift")] -fn solve_loadh_rejects_shift_1() { - let read_data = [0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08]; - run_write_data_sign_extend::(LOADH, read_data, 1); -} - -#[test] -#[should_panic(expected = "LOADH requires 2-byte aligned shift")] -fn solve_loadh_rejects_shift_3() { - let read_data = [0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08]; - run_write_data_sign_extend::(LOADH, read_data, 3); -} - -#[test] -#[should_panic(expected = "LOADH requires 2-byte aligned shift")] -fn solve_loadh_rejects_shift_5() { - let read_data = [0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08]; - run_write_data_sign_extend::(LOADH, read_data, 5); -} - -#[test] -#[should_panic(expected = "LOADH requires 2-byte aligned shift")] -fn solve_loadh_rejects_shift_7() { - let read_data = [0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08]; - run_write_data_sign_extend::(LOADH, read_data, 7); -} - -/// Assert the full set of accepted shifts per opcode: -/// LOADB: 0..7, LOADH: {0,2,4,6}, LOADW: {0,4} -#[test] -fn accepted_shift_sets() { - let read_data = [0x10, 0x20, 0x30, 0x40, 0x50, 0x60, 0x70, 0x80]; - - // LOADB accepts all shifts 0..7 - for shift in 0..8 { - let _ = run_write_data_sign_extend::( - LOADB, read_data, shift, - ); - } - - // LOADH accepts even shifts {0, 2, 4, 6} - for shift in [0, 2, 4, 6] { - let _ = run_write_data_sign_extend::( - LOADH, read_data, shift, - ); - } - - // LOADW accepts only {0, 4} - for shift in [0, 4] { - let _ = run_write_data_sign_extend::( - LOADW, read_data, shift, - ); - } -} - -// //////////////////////////////////////////////////////////////////////////////////// -// CUDA TESTS -// -// Ensure GPU tracegen is equivalent to CPU tracegen -// //////////////////////////////////////////////////////////////////////////////////// - -#[cfg(feature = "cuda")] -type GpuHarness = GpuTestChipHarness< - F, - Rv64LoadSignExtendExecutor, - Rv64LoadSignExtendAir, - Rv64LoadSignExtendChipGpu, - Rv64LoadSignExtendChip, ->; - -#[cfg(feature = "cuda")] -fn create_cuda_harness(tester: &GpuChipTestBuilder) -> GpuHarness { - let dummy_range_checker_chip = dummy_range_checker(default_var_range_checker_bus()); - let dummy_bitwise_chip = Arc::new(BitwiseOperationLookupChip::::new( - default_bitwise_lookup_bus(), - )); - - let (air, executor, cpu_chip) = create_harness_fields( - tester.memory_bridge(), - tester.execution_bridge(), - dummy_range_checker_chip, - dummy_bitwise_chip, - tester.dummy_memory_helper(), - tester.address_bits(), - ); - let gpu_chip = Rv64LoadSignExtendChipGpu::new( - tester.range_checker(), - tester.bitwise_op_lookup(), - tester.address_bits(), - tester.timestamp_max_bits(), - ); - - GpuTestChipHarness::with_capacity(executor, air, gpu_chip, cpu_chip, MAX_INS_CAPACITY) -} - -#[cfg(feature = "cuda")] -#[test_case(LOADB, 100)] -#[test_case(LOADH, 100)] -#[test_case(LOADW, 100)] -fn test_cuda_rand_load_sign_extend_tracegen(opcode: Rv64LoadStoreOpcode, num_ops: usize) { +fn assert_pranked_word_fails( + prank: impl Fn(&mut LoadSignExtendAlignedCoreCols), +) { let mut rng = create_seeded_rng(); - let mut tester = - GpuChipTestBuilder::default().with_bitwise_op_lookup(default_bitwise_lookup_bus()); - - let mut harness = create_cuda_harness(&tester); - for _ in 0..num_ops { - set_and_execute( - &mut tester, - &mut harness.executor, - &mut harness.dense_arena, - &mut rng, - opcode, - None, - None, - None, - None, - ); - } - - type Record<'a> = ( - &'a mut Rv64LoadStoreAdapterRecord, - &'a mut LoadSignExtendCoreRecord, + let mut tester = VmChipTestBuilder::from_config(memory_config_for()); + let mut harness = create_word_harness(&mut tester); + set_and_execute( + &mut tester, + &mut harness.executor, + &mut harness.arena, + &mut rng, + LOADW, + None, + None, + None, ); - - harness - .dense_arena - .get_record_seeker::() - .transfer_to_matrix_arena( - &mut harness.matrix_arena, - EmptyAdapterCoreLayout::::new(), - ); - + let adapter_width = BaseAir::::width(&harness.air.adapter); + let modify_trace = |trace: &mut DenseMatrix| { + let mut trace_row = trace.row_slice(0).unwrap().to_vec(); + let (_, core_row) = trace_row.split_at_mut(adapter_width); + prank(core_row.borrow_mut()); + *trace = RowMajorMatrix::new(trace_row, trace.width()); + }; + disable_debug_builder(); tester .build() - .load_gpu_harness(harness) + .load_and_prank_trace(harness, modify_trace) .finalize() .simple_test() - .unwrap(); + .expect_err("pranked signed word load trace should fail"); +} + +#[test] +fn negative_split_signed_load_tests() { + assert_pranked_byte_fails(|core| core.data_most_sig_bit += F::ONE); + assert_pranked_halfword_fails(|core| core.data_most_sig_bit += F::ONE); + assert_pranked_word_fails(|core| core.data_most_sig_bit += F::ONE); } diff --git a/extensions/riscv/circuit/src/load_sign_extend/word/mod.rs b/extensions/riscv/circuit/src/load_sign_extend/word/mod.rs new file mode 100644 index 0000000000..485fc4ac7d --- /dev/null +++ b/extensions/riscv/circuit/src/load_sign_extend/word/mod.rs @@ -0,0 +1,28 @@ +use openvm_circuit::arch::{VmAirWrapper, VmChipWrapper}; + +use crate::{ + adapters::{Rv64LoadStoreAdapterAir, Rv64LoadStoreAdapterExecutor}, + load_sign_extend::aligned::core::{LoadSignExtendAlignedCoreAir, LoadSignExtendAlignedFiller}, + loadstore::common::{LoadStoreExecutor, KIND_WORD}, +}; + +pub const LOAD_SIGN_EXTEND_WORD_CASES: usize = 2; +pub const LOAD_SIGN_EXTEND_WORD_SELECTOR_WIDTH: usize = 1; + +pub type LoadSignExtendWordCoreAir = LoadSignExtendAlignedCoreAir< + KIND_WORD, + LOAD_SIGN_EXTEND_WORD_CASES, + LOAD_SIGN_EXTEND_WORD_SELECTOR_WIDTH, +>; +pub type LoadSignExtendWordFiller = LoadSignExtendAlignedFiller< + crate::adapters::Rv64LoadStoreAdapterFiller, + KIND_WORD, + LOAD_SIGN_EXTEND_WORD_CASES, + LOAD_SIGN_EXTEND_WORD_SELECTOR_WIDTH, +>; + +pub type Rv64LoadSignExtendWordAir = + VmAirWrapper; +pub type Rv64LoadSignExtendWordExecutor = + LoadStoreExecutor; +pub type Rv64LoadSignExtendWordChip = VmChipWrapper; diff --git a/extensions/riscv/circuit/src/loadstore/aligned/core.rs b/extensions/riscv/circuit/src/loadstore/aligned/core.rs new file mode 100644 index 0000000000..3eb87235ec --- /dev/null +++ b/extensions/riscv/circuit/src/loadstore/aligned/core.rs @@ -0,0 +1,338 @@ +use std::borrow::{Borrow, BorrowMut}; + +use openvm_circuit::{arch::*, system::memory::MemoryAuxColsFactory}; +use openvm_circuit_primitives::{ + encoder::Encoder, + var_range::{SharedVariableRangeCheckerChip, VariableRangeCheckerBus}, + AlignedBorrow, ColumnsAir, StructReflection, StructReflectionHelper, SubAir, +}; +use openvm_instructions::LocalOpcode; +use openvm_riscv_transpiler::Rv64LoadStoreOpcode::{self, *}; +use openvm_stark_backend::{ + interaction::InteractionBuilder, + p3_air::BaseAir, + p3_field::{Field, PrimeCharacteristicRing, PrimeField32}, + BaseAirWithPublicValues, +}; + +use crate::{ + adapters::{LoadStoreInstruction, Rv64LoadStoreAdapterFiller}, + loadstore::common::{ + adapter_context, LoadStoreRecord, KIND_DOUBLEWORD, KIND_HALFWORD, KIND_WORD, + }, +}; + +const SELECTOR_MAX_DEGREE: u32 = 2; + +#[derive(Clone, Copy)] +pub(crate) struct AlignedCase { + opcode: Rv64LoadStoreOpcode, + byte_shift: usize, +} + +impl AlignedCase { + fn is_load(self) -> bool { + matches!(self.opcode, LOADD | LOADWU | LOADHU) + } + + fn cell_shift(self) -> usize { + self.byte_shift / 2 + } +} + +const DOUBLEWORD_CASES: [AlignedCase; 2] = [ + AlignedCase { + opcode: LOADD, + byte_shift: 0, + }, + AlignedCase { + opcode: STORED, + byte_shift: 0, + }, +]; + +const WORD_CASES: [AlignedCase; 4] = [ + AlignedCase { + opcode: LOADWU, + byte_shift: 0, + }, + AlignedCase { + opcode: LOADWU, + byte_shift: 4, + }, + AlignedCase { + opcode: STOREW, + byte_shift: 0, + }, + AlignedCase { + opcode: STOREW, + byte_shift: 4, + }, +]; + +const HALFWORD_CASES: [AlignedCase; 8] = [ + AlignedCase { + opcode: LOADHU, + byte_shift: 0, + }, + AlignedCase { + opcode: LOADHU, + byte_shift: 2, + }, + AlignedCase { + opcode: LOADHU, + byte_shift: 4, + }, + AlignedCase { + opcode: LOADHU, + byte_shift: 6, + }, + AlignedCase { + opcode: STOREH, + byte_shift: 0, + }, + AlignedCase { + opcode: STOREH, + byte_shift: 2, + }, + AlignedCase { + opcode: STOREH, + byte_shift: 4, + }, + AlignedCase { + opcode: STOREH, + byte_shift: 6, + }, +]; + +pub(crate) fn aligned_cases() -> &'static [AlignedCase] { + match KIND { + KIND_DOUBLEWORD => &DOUBLEWORD_CASES, + KIND_WORD => &WORD_CASES, + KIND_HALFWORD => &HALFWORD_CASES, + _ => unreachable!("unsupported aligned loadstore kind"), + } +} + +fn access_cells() -> usize { + match KIND { + KIND_DOUBLEWORD => 4, + KIND_WORD => 2, + KIND_HALFWORD => 1, + _ => unreachable!("unsupported aligned loadstore kind"), + } +} + +fn encoder() -> Encoder { + let encoder = Encoder::new(CASES, SELECTOR_MAX_DEGREE, true); + debug_assert_eq!(encoder.width(), SELECTOR_WIDTH); + encoder +} + +#[repr(C)] +#[derive(Debug, Clone, AlignedBorrow, StructReflection)] +pub struct LoadStoreAlignedCoreCols { + pub selector: [T; SELECTOR_WIDTH], + pub is_valid: T, + pub is_load: T, + pub read_data: [T; BLOCK_FE_WIDTH], + pub prev_data: [T; BLOCK_FE_WIDTH], +} + +#[derive(Debug, Clone, ColumnsAir)] +#[columns_via(LoadStoreAlignedCoreCols)] +pub struct LoadStoreAlignedCoreAir< + const KIND: usize, + const CASES: usize, + const SELECTOR_WIDTH: usize, +> { + pub offset: usize, + encoder: Encoder, +} + +impl + LoadStoreAlignedCoreAir +{ + pub fn new(offset: usize, _range_bus: VariableRangeCheckerBus) -> Self { + debug_assert_eq!(aligned_cases::().len(), CASES); + Self { + offset, + encoder: encoder::(), + } + } +} + +impl BaseAir + for LoadStoreAlignedCoreAir +{ + fn width(&self) -> usize { + LoadStoreAlignedCoreCols::::width() + } +} + +impl + BaseAirWithPublicValues for LoadStoreAlignedCoreAir +{ +} + +impl VmCoreAir + for LoadStoreAlignedCoreAir +where + AB: InteractionBuilder, + I: VmAdapterInterface, + I::Reads: From<([AB::Var; BLOCK_FE_WIDTH], [AB::Expr; BLOCK_FE_WIDTH])>, + I::Writes: From<[[AB::Expr; BLOCK_FE_WIDTH]; 1]>, + I::ProcessedInstruction: From>, +{ + fn eval( + &self, + builder: &mut AB, + local_core: &[AB::Var], + _from_pc: AB::Var, + ) -> AdapterAirContext { + let cols: &LoadStoreAlignedCoreCols = (*local_core).borrow(); + let cases = aligned_cases::(); + let width = access_cells::(); + + self.encoder.eval(builder, &cols.selector); + let selector_flags = self.encoder.flags::(&cols.selector); + let is_valid = self.encoder.is_valid::(&cols.selector); + let is_load = cases + .iter() + .enumerate() + .fold(AB::Expr::ZERO, |acc, (i, case)| { + if case.is_load() { + acc + selector_flags[i].clone() + } else { + acc + } + }); + + builder.assert_eq(cols.is_valid, is_valid.clone()); + builder.assert_eq(cols.is_load, is_load.clone()); + + let expected_opcode = cases + .iter() + .enumerate() + .fold(AB::Expr::ZERO, |acc, (i, case)| { + acc + selector_flags[i].clone() * AB::Expr::from_u8(case.opcode as u8) + }); + let expected_opcode = VmCoreAir::::expr_to_global_expr(self, expected_opcode); + + let load_shift_amount = cases + .iter() + .enumerate() + .fold(AB::Expr::ZERO, |acc, (i, case)| { + if case.is_load() { + acc + selector_flags[i].clone() * AB::Expr::from_usize(case.byte_shift) + } else { + acc + } + }); + let store_shift_amount = cases + .iter() + .enumerate() + .fold(AB::Expr::ZERO, |acc, (i, case)| { + if case.is_load() { + acc + } else { + acc + selector_flags[i].clone() * AB::Expr::from_usize(case.byte_shift) + } + }); + + let write_data = std::array::from_fn(|i| { + cases + .iter() + .enumerate() + .fold(AB::Expr::ZERO, |acc, (case_idx, case)| { + let shift = case.cell_shift(); + let term = if case.is_load() { + if i < width { + cols.read_data[i + shift].into() + } else { + AB::Expr::ZERO + } + } else if i >= shift && i < shift + width { + cols.read_data[i - shift].into() + } else { + cols.prev_data[i].into() + }; + acc + selector_flags[case_idx].clone() * term + }) + }); + adapter_context::( + cols.is_valid.into(), + cols.is_load.into(), + expected_opcode, + load_shift_amount, + store_shift_amount, + cols.read_data, + cols.prev_data, + write_data, + ) + } + + fn start_offset(&self) -> usize { + self.offset + } +} + +#[derive(Clone)] +pub struct LoadStoreAlignedFiller< + A = Rv64LoadStoreAdapterFiller, + const KIND: usize = KIND_WORD, + const CASES: usize = 4, + const SELECTOR_WIDTH: usize = 2, +> { + adapter: A, + pub offset: usize, + encoder: Encoder, +} + +impl + LoadStoreAlignedFiller +{ + pub fn new( + adapter: A, + offset: usize, + _range_checker_chip: SharedVariableRangeCheckerChip, + ) -> Self { + debug_assert_eq!(aligned_cases::().len(), CASES); + Self { + adapter, + offset, + encoder: encoder::(), + } + } +} + +impl TraceFiller + for LoadStoreAlignedFiller +where + F: PrimeField32, + A: 'static + AdapterTraceFiller, +{ + fn fill_trace_row(&self, mem_helper: &MemoryAuxColsFactory, row_slice: &mut [F]) { + let (adapter_row, mut core_row) = unsafe { row_slice.split_at_mut_unchecked(A::WIDTH) }; + self.adapter.fill_trace_row(mem_helper, adapter_row); + + let record: &LoadStoreRecord = unsafe { get_record_from_slice(&mut core_row, ()) }; + let opcode = Rv64LoadStoreOpcode::from_usize(record.local_opcode as usize); + let shift = record.shift_amount as usize; + let read_data = record.read_data; + let prev_data = record.prev_data; + let core_row: &mut LoadStoreAlignedCoreCols = core_row.borrow_mut(); + let cases = aligned_cases::(); + let case_idx = cases + .iter() + .position(|case| case.opcode == opcode && case.byte_shift == shift) + .expect("invalid aligned loadstore opcode/shift"); + + core_row.read_data = read_data.map(F::from_u16); + core_row.prev_data = prev_data.map(F::from_u16); + core_row.is_load = F::from_bool(matches!(opcode, LOADD | LOADWU | LOADHU)); + core_row.is_valid = F::ONE; + let pt: [u32; SELECTOR_WIDTH] = self.encoder.get_flag_pt(case_idx).try_into().unwrap(); + core_row.selector = pt.map(F::from_u32); + } +} diff --git a/extensions/riscv/circuit/src/loadstore/aligned/mod.rs b/extensions/riscv/circuit/src/loadstore/aligned/mod.rs new file mode 100644 index 0000000000..8244e2f2b3 --- /dev/null +++ b/extensions/riscv/circuit/src/loadstore/aligned/mod.rs @@ -0,0 +1,2 @@ +pub mod core; +pub use core::*; diff --git a/extensions/riscv/circuit/src/loadstore/byte/core.rs b/extensions/riscv/circuit/src/loadstore/byte/core.rs new file mode 100644 index 0000000000..8aeb8a9f55 --- /dev/null +++ b/extensions/riscv/circuit/src/loadstore/byte/core.rs @@ -0,0 +1,399 @@ +use std::borrow::{Borrow, BorrowMut}; + +use openvm_circuit::{arch::*, system::memory::MemoryAuxColsFactory}; +use openvm_circuit_primitives::{ + bitwise_op_lookup::{BitwiseOperationLookupBus, SharedBitwiseOperationLookupChip}, + encoder::Encoder, + var_range::{SharedVariableRangeCheckerChip, VariableRangeCheckerBus}, + AlignedBorrow, ColumnsAir, StructReflection, StructReflectionHelper, SubAir, +}; +use openvm_instructions::LocalOpcode; +use openvm_riscv_transpiler::Rv64LoadStoreOpcode::{self, *}; +use openvm_stark_backend::{ + interaction::InteractionBuilder, + p3_air::BaseAir, + p3_field::{Field, PrimeCharacteristicRing, PrimeField32}, + BaseAirWithPublicValues, +}; + +use crate::{ + adapters::{LoadStoreInstruction, Rv64LoadStoreAdapterFiller, RV64_BYTE_BITS}, + loadstore::common::{ + adapter_context, byte_from_cell, replace_byte, run_write_data, LoadStoreRecord, BYTE_BITS, + }, +}; + +const BYTE_CASES: usize = 16; +const BYTE_SELECTOR_MAX_DEGREE: u32 = 2; +pub(crate) const BYTE_SELECTOR_WIDTH: usize = 5; + +#[derive(Clone, Copy)] +struct ByteCase { + opcode: Rv64LoadStoreOpcode, + byte_shift: usize, +} + +impl ByteCase { + fn is_load(self) -> bool { + self.opcode == LOADBU + } + + fn cell_shift(self) -> usize { + self.byte_shift / 2 + } + + fn byte_idx(self) -> usize { + self.byte_shift % 2 + } + + fn read_cell_idx(self) -> usize { + if self.opcode == STOREB { + 0 + } else { + self.cell_shift() + } + } +} + +const CASES: [ByteCase; BYTE_CASES] = [ + ByteCase { + opcode: LOADBU, + byte_shift: 0, + }, + ByteCase { + opcode: LOADBU, + byte_shift: 1, + }, + ByteCase { + opcode: LOADBU, + byte_shift: 2, + }, + ByteCase { + opcode: LOADBU, + byte_shift: 3, + }, + ByteCase { + opcode: LOADBU, + byte_shift: 4, + }, + ByteCase { + opcode: LOADBU, + byte_shift: 5, + }, + ByteCase { + opcode: LOADBU, + byte_shift: 6, + }, + ByteCase { + opcode: LOADBU, + byte_shift: 7, + }, + ByteCase { + opcode: STOREB, + byte_shift: 0, + }, + ByteCase { + opcode: STOREB, + byte_shift: 1, + }, + ByteCase { + opcode: STOREB, + byte_shift: 2, + }, + ByteCase { + opcode: STOREB, + byte_shift: 3, + }, + ByteCase { + opcode: STOREB, + byte_shift: 4, + }, + ByteCase { + opcode: STOREB, + byte_shift: 5, + }, + ByteCase { + opcode: STOREB, + byte_shift: 6, + }, + ByteCase { + opcode: STOREB, + byte_shift: 7, + }, +]; + +fn encoder() -> Encoder { + let encoder = Encoder::new(BYTE_CASES, BYTE_SELECTOR_MAX_DEGREE, true); + debug_assert_eq!(encoder.width(), BYTE_SELECTOR_WIDTH); + encoder +} + +#[repr(C)] +#[derive(Debug, Clone, AlignedBorrow, StructReflection)] +pub struct LoadStoreByteCoreCols { + pub selector: [T; BYTE_SELECTOR_WIDTH], + pub is_valid: T, + pub is_load: T, + pub read_cell_bytes: [T; 2], + pub prev_cell_bytes: [T; 2], + pub read_data: [T; BLOCK_FE_WIDTH], + pub prev_data: [T; BLOCK_FE_WIDTH], +} + +#[derive(Debug, Clone, ColumnsAir)] +#[columns_via(LoadStoreByteCoreCols)] +pub struct LoadStoreByteCoreAir { + pub offset: usize, + encoder: Encoder, + bitwise_lookup_bus: BitwiseOperationLookupBus, +} + +impl LoadStoreByteCoreAir { + pub fn new( + offset: usize, + bitwise_lookup_bus: BitwiseOperationLookupBus, + _range_bus: VariableRangeCheckerBus, + ) -> Self { + Self { + offset, + encoder: encoder(), + bitwise_lookup_bus, + } + } +} + +impl BaseAir for LoadStoreByteCoreAir { + fn width(&self) -> usize { + LoadStoreByteCoreCols::::width() + } +} + +impl BaseAirWithPublicValues for LoadStoreByteCoreAir {} + +impl VmCoreAir for LoadStoreByteCoreAir +where + AB: InteractionBuilder, + I: VmAdapterInterface, + I::Reads: From<([AB::Var; BLOCK_FE_WIDTH], [AB::Expr; BLOCK_FE_WIDTH])>, + I::Writes: From<[[AB::Expr; BLOCK_FE_WIDTH]; 1]>, + I::ProcessedInstruction: From>, +{ + fn eval( + &self, + builder: &mut AB, + local_core: &[AB::Var], + _from_pc: AB::Var, + ) -> AdapterAirContext { + let cols: &LoadStoreByteCoreCols = (*local_core).borrow(); + self.encoder.eval(builder, &cols.selector); + let flags = self.encoder.flags::(&cols.selector); + let is_valid = self.encoder.is_valid::(&cols.selector); + let is_load = CASES + .iter() + .enumerate() + .fold(AB::Expr::ZERO, |acc, (i, case)| { + if case.is_load() { + acc + flags[i].clone() + } else { + acc + } + }); + let is_store = is_valid.clone() - is_load.clone(); + + builder.assert_eq(cols.is_valid, is_valid.clone()); + builder.assert_eq(cols.is_load, is_load.clone()); + + self.bitwise_lookup_bus + .send_range(cols.read_cell_bytes[0], cols.read_cell_bytes[1]) + .eval(builder, is_valid.clone()); + self.bitwise_lookup_bus + .send_range(cols.prev_cell_bytes[0], cols.prev_cell_bytes[1]) + .eval(builder, is_store.clone()); + + let read_cell = + cols.read_cell_bytes[0] + cols.read_cell_bytes[1] * AB::Expr::from_u32(1 << BYTE_BITS); + let expected_read_cell = CASES + .iter() + .enumerate() + .fold(AB::Expr::ZERO, |acc, (i, case)| { + acc + flags[i].clone() * cols.read_data[case.read_cell_idx()] + }); + builder.assert_eq(read_cell.clone(), expected_read_cell); + + let prev_cell = + cols.prev_cell_bytes[0] + cols.prev_cell_bytes[1] * AB::Expr::from_u32(1 << BYTE_BITS); + let expected_prev_cell = CASES + .iter() + .enumerate() + .fold(AB::Expr::ZERO, |acc, (i, case)| { + if case.opcode == STOREB { + acc + flags[i].clone() * cols.prev_data[case.cell_shift()] + } else { + acc + } + }); + builder.assert_eq(is_store.clone() * prev_cell, expected_prev_cell); + + let expected_opcode = CASES + .iter() + .enumerate() + .fold(AB::Expr::ZERO, |acc, (i, case)| { + acc + flags[i].clone() * AB::Expr::from_u8(case.opcode as u8) + }); + let expected_opcode = VmCoreAir::::expr_to_global_expr(self, expected_opcode); + let load_shift_amount = CASES + .iter() + .enumerate() + .fold(AB::Expr::ZERO, |acc, (i, case)| { + if case.is_load() { + acc + flags[i].clone() * AB::Expr::from_usize(case.byte_shift) + } else { + acc + } + }); + let store_shift_amount = CASES + .iter() + .enumerate() + .fold(AB::Expr::ZERO, |acc, (i, case)| { + if case.is_load() { + acc + } else { + acc + flags[i].clone() * AB::Expr::from_usize(case.byte_shift) + } + }); + + let write_data = std::array::from_fn(|i| { + CASES + .iter() + .enumerate() + .fold(AB::Expr::ZERO, |acc, (case_idx, case)| { + let term = match case.opcode { + LOADBU => { + let selected_case_byte = if case.byte_idx() == 0 { + cols.read_cell_bytes[0].into() + } else { + cols.read_cell_bytes[1].into() + }; + if i == 0 { + selected_case_byte + } else { + AB::Expr::ZERO + } + } + STOREB => { + if i == case.cell_shift() { + if case.byte_idx() == 0 { + cols.read_cell_bytes[0] + + cols.prev_cell_bytes[1] + * AB::Expr::from_u32(1 << BYTE_BITS) + } else { + cols.prev_cell_bytes[0] + + cols.read_cell_bytes[0] + * AB::Expr::from_u32(1 << BYTE_BITS) + } + } else { + cols.prev_data[i].into() + } + } + _ => unreachable!(), + }; + acc + flags[case_idx].clone() * term + }) + }); + adapter_context::( + cols.is_valid.into(), + cols.is_load.into(), + expected_opcode, + load_shift_amount, + store_shift_amount, + cols.read_data, + cols.prev_data, + write_data, + ) + } + + fn start_offset(&self) -> usize { + self.offset + } +} + +#[derive(Clone)] +pub struct LoadStoreByteFiller { + adapter: A, + pub offset: usize, + encoder: Encoder, + bitwise_lookup_chip: SharedBitwiseOperationLookupChip, +} + +impl LoadStoreByteFiller { + pub fn new( + adapter: A, + offset: usize, + bitwise_lookup_chip: SharedBitwiseOperationLookupChip, + _range_checker_chip: SharedVariableRangeCheckerChip, + ) -> Self { + Self { + adapter, + offset, + encoder: encoder(), + bitwise_lookup_chip, + } + } +} + +impl TraceFiller for LoadStoreByteFiller +where + F: PrimeField32, + A: 'static + AdapterTraceFiller, +{ + fn fill_trace_row(&self, mem_helper: &MemoryAuxColsFactory, row_slice: &mut [F]) { + let (adapter_row, mut core_row) = unsafe { row_slice.split_at_mut_unchecked(A::WIDTH) }; + self.adapter.fill_trace_row(mem_helper, adapter_row); + + let record: &LoadStoreRecord = unsafe { get_record_from_slice(&mut core_row, ()) }; + let opcode = Rv64LoadStoreOpcode::from_usize(record.local_opcode as usize); + let shift = record.shift_amount as usize; + let read_data = record.read_data; + let prev_data = record.prev_data; + let core_row: &mut LoadStoreByteCoreCols = core_row.borrow_mut(); + let cell_shift = shift / 2; + let byte_idx = shift % 2; + let case_idx = CASES + .iter() + .position(|case| case.opcode == opcode && case.byte_shift == shift) + .expect("invalid byte loadstore opcode/shift"); + + let read_cell = if opcode == STOREB { + read_data[0] + } else { + read_data[cell_shift] + }; + let read_cell_bytes = [byte_from_cell(read_cell, 0), byte_from_cell(read_cell, 1)]; + self.bitwise_lookup_chip + .request_range(read_cell_bytes[0] as u32, read_cell_bytes[1] as u32); + core_row.read_cell_bytes = read_cell_bytes.map(F::from_u16); + + if opcode == STOREB { + let prev_cell_bytes = [ + byte_from_cell(prev_data[cell_shift], 0), + byte_from_cell(prev_data[cell_shift], 1), + ]; + self.bitwise_lookup_chip + .request_range(prev_cell_bytes[0] as u32, prev_cell_bytes[1] as u32); + core_row.prev_cell_bytes = prev_cell_bytes.map(F::from_u16); + debug_assert_eq!( + run_write_data(opcode, read_data, prev_data, shift)[cell_shift], + replace_byte(prev_data[cell_shift], byte_idx, read_cell_bytes[0]) + ); + } else { + core_row.prev_cell_bytes = [F::ZERO; 2]; + } + + core_row.read_data = read_data.map(F::from_u16); + core_row.prev_data = prev_data.map(F::from_u16); + core_row.is_load = F::from_bool(opcode == LOADBU); + core_row.is_valid = F::ONE; + let pt: [u32; BYTE_SELECTOR_WIDTH] = self.encoder.get_flag_pt(case_idx).try_into().unwrap(); + core_row.selector = pt.map(F::from_u32); + } +} diff --git a/extensions/riscv/circuit/src/loadstore/byte/mod.rs b/extensions/riscv/circuit/src/loadstore/byte/mod.rs new file mode 100644 index 0000000000..98c6f2d6e8 --- /dev/null +++ b/extensions/riscv/circuit/src/loadstore/byte/mod.rs @@ -0,0 +1,13 @@ +use openvm_circuit::arch::{VmAirWrapper, VmChipWrapper}; + +use crate::{ + adapters::{Rv64LoadStoreAdapterAir, Rv64LoadStoreAdapterExecutor}, + loadstore::common::{LoadStoreExecutor, KIND_BYTE}, +}; + +mod core; +pub use core::*; + +pub type Rv64LoadStoreByteAir = VmAirWrapper; +pub type Rv64LoadStoreByteExecutor = LoadStoreExecutor; +pub type Rv64LoadStoreByteChip = VmChipWrapper; diff --git a/extensions/riscv/circuit/src/loadstore/common.rs b/extensions/riscv/circuit/src/loadstore/common.rs new file mode 100644 index 0000000000..c778540dea --- /dev/null +++ b/extensions/riscv/circuit/src/loadstore/common.rs @@ -0,0 +1,200 @@ +use openvm_circuit::{arch::*, system::memory::online::TracingMemory}; +use openvm_circuit_primitives::AlignedBytesBorrow; +use openvm_instructions::{instruction::Instruction, program::DEFAULT_PC_STEP, LocalOpcode}; +use openvm_riscv_transpiler::Rv64LoadStoreOpcode::{self, *}; +use openvm_stark_backend::p3_field::PrimeField32; + +use crate::adapters::LoadStoreInstruction; + +pub const KIND_BYTE: usize = 0; +pub const KIND_HALFWORD: usize = 1; +pub const KIND_WORD: usize = 2; +pub const KIND_DOUBLEWORD: usize = 3; + +pub(crate) const BYTE_BITS: usize = 8; +pub(crate) const BYTE_MASK: u16 = 0xff; +pub(crate) const SIGN_BYTE: u16 = 1 << 7; +pub(crate) const SIGN_U16: u16 = 1 << 15; + +#[repr(C)] +#[derive(AlignedBytesBorrow, Debug)] +pub struct LoadStoreRecord { + pub local_opcode: u8, + pub shift_amount: u8, + pub read_data: [u16; BLOCK_FE_WIDTH], + pub prev_data: [u16; BLOCK_FE_WIDTH], +} + +#[derive(Clone, Copy, derive_new::new)] +pub struct LoadStoreExecutor { + adapter: A, + pub offset: usize, +} + +impl PreflightExecutor for LoadStoreExecutor +where + F: PrimeField32, + A: 'static + + AdapterTraceExecutor< + F, + ReadData = (([u16; BLOCK_FE_WIDTH], [u16; BLOCK_FE_WIDTH]), u8), + WriteData = [u16; BLOCK_FE_WIDTH], + >, + for<'buf> RA: RecordArena< + 'buf, + EmptyAdapterCoreLayout, + (A::RecordMut<'buf>, &'buf mut LoadStoreRecord), + >, +{ + fn get_opcode_name(&self, opcode: usize) -> String { + format!( + "{:?}", + Rv64LoadStoreOpcode::from_usize(opcode - self.offset) + ) + } + + fn execute( + &self, + state: VmStateMut, + instruction: &Instruction, + ) -> Result<(), ExecutionError> { + let Instruction { opcode, .. } = instruction; + let (mut adapter_record, core_record) = state.ctx.alloc(EmptyAdapterCoreLayout::new()); + + A::start(*state.pc, state.memory, &mut adapter_record); + ( + (core_record.prev_data, core_record.read_data), + core_record.shift_amount, + ) = self + .adapter + .read(state.memory, instruction, &mut adapter_record); + + let local_opcode = Rv64LoadStoreOpcode::from_usize(opcode.local_opcode_idx(self.offset)); + core_record.local_opcode = local_opcode as u8; + + let write_data = run_write_data( + local_opcode, + core_record.read_data, + core_record.prev_data, + core_record.shift_amount as usize, + ); + self.adapter + .write(state.memory, instruction, write_data, &mut adapter_record); + + *state.pc = state.pc.wrapping_add(DEFAULT_PC_STEP); + Ok(()) + } +} + +pub(crate) fn byte_from_cell(cell: u16, byte_idx: usize) -> u16 { + (cell >> (BYTE_BITS * byte_idx)) & BYTE_MASK +} + +pub(crate) fn replace_byte(cell: u16, byte_idx: usize, byte: u16) -> u16 { + debug_assert!(byte <= BYTE_MASK); + let shift = BYTE_BITS * byte_idx; + (cell & !(BYTE_MASK << shift)) | (byte << shift) +} + +pub(crate) fn run_write_data( + opcode: Rv64LoadStoreOpcode, + read_data: [u16; BLOCK_FE_WIDTH], + prev_data: [u16; BLOCK_FE_WIDTH], + byte_shift: usize, +) -> [u16; BLOCK_FE_WIDTH] { + let cell_shift = byte_shift / 2; + match opcode { + LOADD if byte_shift == 0 => read_data, + LOADWU if byte_shift == 0 || byte_shift == 4 => [ + read_data[cell_shift], + read_data[cell_shift + 1], + 0, + 0, + ], + LOADW if byte_shift == 0 || byte_shift == 4 => { + let sign = if read_data[cell_shift + 1] & SIGN_U16 != 0 { + u16::MAX + } else { + 0 + }; + [read_data[cell_shift], read_data[cell_shift + 1], sign, sign] + } + LOADHU if byte_shift == 0 || byte_shift == 2 || byte_shift == 4 || byte_shift == 6 => { + [read_data[cell_shift], 0, 0, 0] + } + LOADH if byte_shift == 0 || byte_shift == 2 || byte_shift == 4 || byte_shift == 6 => { + let sign = if read_data[cell_shift] & SIGN_U16 != 0 { + u16::MAX + } else { + 0 + }; + [read_data[cell_shift], sign, sign, sign] + } + LOADBU if byte_shift < 8 => { + let byte = byte_from_cell(read_data[cell_shift], byte_shift % 2); + [byte, 0, 0, 0] + } + LOADB if byte_shift < 8 => { + let byte = byte_from_cell(read_data[cell_shift], byte_shift % 2); + if byte & SIGN_BYTE != 0 { + [byte | 0xff00, u16::MAX, u16::MAX, u16::MAX] + } else { + [byte, 0, 0, 0] + } + } + STORED if byte_shift == 0 => read_data, + STOREW if byte_shift == 0 || byte_shift == 4 => { + let mut write_data = prev_data; + write_data[cell_shift] = read_data[0]; + write_data[cell_shift + 1] = read_data[1]; + write_data + } + STOREH if byte_shift == 0 || byte_shift == 2 || byte_shift == 4 || byte_shift == 6 => { + let mut write_data = prev_data; + write_data[cell_shift] = read_data[0]; + write_data + } + STOREB if byte_shift < 8 => { + let mut write_data = prev_data; + let byte = byte_from_cell(read_data[0], 0); + write_data[cell_shift] = replace_byte(prev_data[cell_shift], byte_shift % 2, byte); + write_data + } + _ => unreachable!( + "unaligned memory access not supported by this execution environment: {opcode:?}, byte_shift: {byte_shift}" + ), + } +} + +#[allow(clippy::too_many_arguments)] +pub(crate) fn adapter_context( + is_valid: AB::Expr, + is_load: AB::Expr, + expected_opcode: AB::Expr, + load_shift_amount: AB::Expr, + store_shift_amount: AB::Expr, + read_data: [AB::Var; BLOCK_FE_WIDTH], + prev_data: [AB::Var; BLOCK_FE_WIDTH], + write_data: [AB::Expr; BLOCK_FE_WIDTH], +) -> AdapterAirContext +where + AB: openvm_stark_backend::interaction::InteractionBuilder, + I: VmAdapterInterface, + I::Reads: From<([AB::Var; BLOCK_FE_WIDTH], [AB::Expr; BLOCK_FE_WIDTH])>, + I::Writes: From<[[AB::Expr; BLOCK_FE_WIDTH]; 1]>, + I::ProcessedInstruction: From>, +{ + AdapterAirContext { + to_pc: None, + reads: (prev_data, read_data.map(|x| x.into())).into(), + writes: [write_data].into(), + instruction: LoadStoreInstruction { + is_valid, + opcode: expected_opcode, + is_load, + load_shift_amount, + store_shift_amount, + } + .into(), + } +} diff --git a/extensions/riscv/circuit/src/loadstore/core.rs b/extensions/riscv/circuit/src/loadstore/core.rs deleted file mode 100644 index ec15d304b3..0000000000 --- a/extensions/riscv/circuit/src/loadstore/core.rs +++ /dev/null @@ -1,578 +0,0 @@ -use std::{ - array, - borrow::{Borrow, BorrowMut}, - fmt::Debug, -}; - -use openvm_circuit::{ - arch::*, - system::memory::{online::TracingMemory, MemoryAuxColsFactory}, -}; -use openvm_circuit_primitives::{ - bitwise_op_lookup::{BitwiseOperationLookupBus, SharedBitwiseOperationLookupChip}, - encoder::Encoder, - AlignedBorrow, AlignedBytesBorrow, ColumnsAir, StructReflection, StructReflectionHelper, - SubAir, -}; -use openvm_instructions::{ - instruction::Instruction, - program::DEFAULT_PC_STEP, - riscv::{RV64_MEMORY_AS, RV64_REGISTER_AS, RV64_REGISTER_NUM_LIMBS}, - LocalOpcode, PUBLIC_VALUES_AS, -}; -use openvm_riscv_transpiler::Rv64LoadStoreOpcode::{self, *}; -use openvm_stark_backend::{ - interaction::InteractionBuilder, - p3_air::BaseAir, - p3_field::{Field, PrimeCharacteristicRing, PrimeField32}, - BaseAirWithPublicValues, -}; - -use crate::adapters::{LoadStoreInstruction, Rv64LoadStoreAdapterFiller}; - -const LOADSTORE_SELECTOR_CASES: usize = 30; -const LOADSTORE_SELECTOR_MAX_DEGREE: u32 = 2; -pub(crate) const LOADSTORE_SELECTOR_WIDTH: usize = 7; - -#[derive(Debug, Clone, Copy)] -enum InstructionCase { - LoadD0, - LoadWu0, - LoadWu4, - LoadHu0, - LoadHu2, - LoadHu4, - LoadHu6, - LoadBu0, - LoadBu1, - LoadBu2, - LoadBu3, - LoadBu4, - LoadBu5, - LoadBu6, - LoadBu7, - StoreD0, - StoreW0, - StoreW4, - StoreH0, - StoreH2, - StoreH4, - StoreH6, - StoreB0, - StoreB1, - StoreB2, - StoreB3, - StoreB4, - StoreB5, - StoreB6, - StoreB7, -} - -use InstructionCase::*; - -impl InstructionCase { - const ALL: [Self; LOADSTORE_SELECTOR_CASES] = [ - LoadD0, LoadWu0, LoadWu4, LoadHu0, LoadHu2, LoadHu4, LoadHu6, LoadBu0, LoadBu1, LoadBu2, - LoadBu3, LoadBu4, LoadBu5, LoadBu6, LoadBu7, StoreD0, StoreW0, StoreW4, StoreH0, StoreH2, - StoreH4, StoreH6, StoreB0, StoreB1, StoreB2, StoreB3, StoreB4, StoreB5, StoreB6, StoreB7, - ]; - - fn opcode(self) -> Rv64LoadStoreOpcode { - match self { - LoadD0 => LOADD, - LoadWu0 | LoadWu4 => LOADWU, - LoadHu0 | LoadHu2 | LoadHu4 | LoadHu6 => LOADHU, - LoadBu0 | LoadBu1 | LoadBu2 | LoadBu3 | LoadBu4 | LoadBu5 | LoadBu6 | LoadBu7 => LOADBU, - StoreD0 => STORED, - StoreW0 | StoreW4 => STOREW, - StoreH0 | StoreH2 | StoreH4 | StoreH6 => STOREH, - StoreB0 | StoreB1 | StoreB2 | StoreB3 | StoreB4 | StoreB5 | StoreB6 | StoreB7 => STOREB, - } - } - - fn shift(self) -> usize { - match self { - LoadD0 | StoreD0 => 0, - LoadWu0 | StoreW0 | LoadHu0 | StoreH0 | LoadBu0 | StoreB0 => 0, - LoadBu1 | StoreB1 => 1, - LoadHu2 | StoreH2 | LoadBu2 | StoreB2 => 2, - LoadBu3 | StoreB3 => 3, - LoadWu4 | StoreW4 | LoadHu4 | StoreH4 | LoadBu4 | StoreB4 => 4, - LoadBu5 | StoreB5 => 5, - LoadHu6 | StoreH6 | LoadBu6 | StoreB6 => 6, - LoadBu7 | StoreB7 => 7, - } - } - - fn is_load(self) -> bool { - matches!( - self, - LoadD0 - | LoadWu0 - | LoadWu4 - | LoadHu0 - | LoadHu2 - | LoadHu4 - | LoadHu6 - | LoadBu0 - | LoadBu1 - | LoadBu2 - | LoadBu3 - | LoadBu4 - | LoadBu5 - | LoadBu6 - | LoadBu7 - ) - } - - fn width(self) -> usize { - match self.opcode() { - LOADD | STORED => 8, - LOADWU | STOREW => 4, - LOADHU | STOREH => 2, - LOADBU | STOREB => 1, - _ => unreachable!("loadstore core should not handle sign-extension loads"), - } - } - - fn from_opcode_shift(opcode: Rv64LoadStoreOpcode, shift: usize) -> Self { - match (opcode, shift) { - (LOADD, 0) => LoadD0, - (LOADWU, 0) => LoadWu0, - (LOADWU, 4) => LoadWu4, - (LOADHU, 0) => LoadHu0, - (LOADHU, 2) => LoadHu2, - (LOADHU, 4) => LoadHu4, - (LOADHU, 6) => LoadHu6, - (LOADBU, 0) => LoadBu0, - (LOADBU, 1) => LoadBu1, - (LOADBU, 2) => LoadBu2, - (LOADBU, 3) => LoadBu3, - (LOADBU, 4) => LoadBu4, - (LOADBU, 5) => LoadBu5, - (LOADBU, 6) => LoadBu6, - (LOADBU, 7) => LoadBu7, - (STORED, 0) => StoreD0, - (STOREW, 0) => StoreW0, - (STOREW, 4) => StoreW4, - (STOREH, 0) => StoreH0, - (STOREH, 2) => StoreH2, - (STOREH, 4) => StoreH4, - (STOREH, 6) => StoreH6, - (STOREB, 0) => StoreB0, - (STOREB, 1) => StoreB1, - (STOREB, 2) => StoreB2, - (STOREB, 3) => StoreB3, - (STOREB, 4) => StoreB4, - (STOREB, 5) => StoreB5, - (STOREB, 6) => StoreB6, - (STOREB, 7) => StoreB7, - _ => unreachable!( - "unaligned memory access not supported by this execution environment: {opcode:?}, shift: {shift}" - ), - } - } -} - -fn loadstore_encoder() -> Encoder { - let encoder = Encoder::new( - LOADSTORE_SELECTOR_CASES, - LOADSTORE_SELECTOR_MAX_DEGREE, - true, - ); - debug_assert_eq!(encoder.width(), LOADSTORE_SELECTOR_WIDTH); - encoder -} - -#[cfg(test)] -pub(crate) fn selector_point_for_opcode_shift( - opcode: Rv64LoadStoreOpcode, - shift: usize, -) -> [u32; LOADSTORE_SELECTOR_WIDTH] { - loadstore_encoder() - .get_flag_pt(InstructionCase::from_opcode_shift(opcode, shift) as usize) - .try_into() - .unwrap() -} - -/// LoadStore Core Chip handles byte/halfword/word into doubleword conversions and unsigned -/// extends. This chip uses read_data and prev_data to constrain the write_data. It also handles -/// the shifting in case of not 8 byte aligned instructions. This chip treats each `(opcode, -/// shift)` pair as a separate instruction. -#[repr(C)] -#[derive(Debug, Clone, AlignedBorrow, StructReflection)] -pub struct LoadStoreCoreCols { - pub selector: [T; LOADSTORE_SELECTOR_WIDTH], - /// we need to keep the degree of is_valid and is_load to 1 - pub is_valid: T, - pub is_load: T, - - pub read_data: [T; NUM_CELLS], - pub prev_data: [T; NUM_CELLS], - /// write_data will be constrained against read_data and prev_data - /// depending on the opcode and the shift amount - pub write_data: [T; NUM_CELLS], -} - -#[derive(Debug, Clone, ColumnsAir)] -#[columns_via(LoadStoreCoreCols)] -pub struct LoadStoreCoreAir { - pub offset: usize, - encoder: Encoder, - bitwise_lookup_bus: BitwiseOperationLookupBus, -} - -impl LoadStoreCoreAir { - pub fn new(offset: usize, bitwise_lookup_bus: BitwiseOperationLookupBus) -> Self { - assert!(NUM_CELLS.is_multiple_of(2)); - Self { - offset, - encoder: loadstore_encoder(), - bitwise_lookup_bus, - } - } -} - -impl BaseAir for LoadStoreCoreAir { - fn width(&self) -> usize { - LoadStoreCoreCols::::width() - } -} - -impl BaseAirWithPublicValues for LoadStoreCoreAir {} - -impl VmCoreAir for LoadStoreCoreAir -where - AB: InteractionBuilder, - I: VmAdapterInterface, - I::Reads: From<([AB::Var; NUM_CELLS], [AB::Expr; NUM_CELLS])>, - I::Writes: From<[[AB::Expr; NUM_CELLS]; 1]>, - I::ProcessedInstruction: From>, -{ - fn eval( - &self, - builder: &mut AB, - local_core: &[AB::Var], - _from_pc: AB::Var, - ) -> AdapterAirContext { - let cols: &LoadStoreCoreCols = (*local_core).borrow(); - let LoadStoreCoreCols:: { - selector, - is_valid, - is_load, - read_data, - prev_data, - write_data, - } = *cols; - - self.encoder.eval(builder, &selector); - - let selector_flags = self.encoder.flags::(&selector); - let expected_is_valid = self.encoder.is_valid::(&selector); - let expected_is_load = InstructionCase::ALL - .iter() - .fold(AB::Expr::ZERO, |acc, &case| { - if case.is_load() { - acc + selector_flags[case as usize].clone() - } else { - acc - } - }); - - builder.assert_eq(is_valid, expected_is_valid.clone()); - builder.assert_eq(is_load, expected_is_load.clone()); - - for pair in read_data.chunks_exact(2) { - self.bitwise_lookup_bus - .send_range(pair[0], pair[1]) - .eval(builder, is_valid.into()); - } - for pair in prev_data.chunks_exact(2) { - self.bitwise_lookup_bus - .send_range(pair[0], pair[1]) - .eval(builder, is_valid.into()); - } - - let expected_opcode = InstructionCase::ALL - .iter() - .fold(AB::Expr::ZERO, |acc, &case| { - acc + selector_flags[case as usize].clone() * AB::Expr::from_u8(case.opcode() as u8) - }); - let expected_opcode = VmCoreAir::::expr_to_global_expr(self, expected_opcode); - - let load_shift_amount = InstructionCase::ALL - .iter() - .fold(AB::Expr::ZERO, |acc, &case| { - if case.is_load() { - acc + selector_flags[case as usize].clone() * AB::Expr::from_usize(case.shift()) - } else { - acc - } - }); - let store_shift_amount = InstructionCase::ALL - .iter() - .fold(AB::Expr::ZERO, |acc, &case| { - if case.is_load() { - acc - } else { - acc + selector_flags[case as usize].clone() * AB::Expr::from_usize(case.shift()) - } - }); - - for (i, cell) in write_data.iter().enumerate() { - let expected = InstructionCase::ALL - .iter() - .fold(AB::Expr::ZERO, |acc, &case| { - let width = case.width(); - let shift = case.shift(); - debug_assert!(shift + width <= NUM_CELLS); - let term = if case.is_load() { - if i < width { - read_data[i + shift].into() - } else { - AB::Expr::ZERO - } - } else if i >= shift && i < shift + width { - read_data[i - shift].into() - } else { - prev_data[i].into() - }; - acc + selector_flags[case as usize].clone() * term - }); - builder.assert_eq(*cell, expected); - } - - AdapterAirContext { - to_pc: None, - reads: (prev_data, read_data.map(|x| x.into())).into(), - writes: [write_data.map(|x| x.into())].into(), - instruction: LoadStoreInstruction { - is_valid: is_valid.into(), - opcode: expected_opcode, - is_load: is_load.into(), - load_shift_amount, - store_shift_amount, - } - .into(), - } - } - - fn start_offset(&self) -> usize { - self.offset - } -} - -#[repr(C)] -#[derive(AlignedBytesBorrow, Debug)] -pub struct LoadStoreCoreRecord { - pub local_opcode: u8, - pub shift_amount: u8, - pub read_data: [u8; NUM_CELLS], - pub prev_data: [u8; NUM_CELLS], -} - -#[derive(Clone, Copy, derive_new::new)] -pub struct LoadStoreExecutor { - adapter: A, - pub offset: usize, -} - -#[derive(Clone)] -pub struct LoadStoreFiller< - A = Rv64LoadStoreAdapterFiller, - const NUM_CELLS: usize = RV64_REGISTER_NUM_LIMBS, -> { - adapter: A, - pub offset: usize, - encoder: Encoder, - bitwise_lookup_chip: SharedBitwiseOperationLookupChip<8>, -} - -impl LoadStoreFiller { - pub fn new( - adapter: A, - offset: usize, - bitwise_lookup_chip: SharedBitwiseOperationLookupChip<8>, - ) -> Self { - assert!(NUM_CELLS.is_multiple_of(2)); - Self { - adapter, - offset, - encoder: loadstore_encoder(), - bitwise_lookup_chip, - } - } -} - -impl PreflightExecutor for LoadStoreExecutor -where - F: PrimeField32, - A: 'static - + AdapterTraceExecutor< - F, - ReadData = (([u8; NUM_CELLS], [u8; NUM_CELLS]), u8), - WriteData = [u8; NUM_CELLS], - >, - for<'buf> RA: RecordArena< - 'buf, - EmptyAdapterCoreLayout, - (A::RecordMut<'buf>, &'buf mut LoadStoreCoreRecord), - >, -{ - fn get_opcode_name(&self, opcode: usize) -> String { - format!( - "{:?}", - Rv64LoadStoreOpcode::from_usize(opcode - self.offset) - ) - } - - fn execute( - &self, - state: VmStateMut, - instruction: &Instruction, - ) -> Result<(), ExecutionError> { - let Instruction { opcode, d, e, .. } = instruction; - let local_opcode = Rv64LoadStoreOpcode::from_usize(opcode.local_opcode_idx(self.offset)); - debug_assert_eq!(d.as_canonical_u32(), RV64_REGISTER_AS); - debug_assert!(match local_opcode { - LOADD | LOADWU | LOADHU | LOADBU => e.as_canonical_u32() == RV64_MEMORY_AS, - STORED | STOREW | STOREH | STOREB => - e.as_canonical_u32() == RV64_MEMORY_AS || e.as_canonical_u32() == PUBLIC_VALUES_AS, - _ => false, - }); - - let (mut adapter_record, core_record) = state.ctx.alloc(EmptyAdapterCoreLayout::new()); - - A::start(*state.pc, state.memory, &mut adapter_record); - - let ((prev_data, read_data), shift_amount) = - self.adapter - .read(state.memory, instruction, &mut adapter_record); - core_record.prev_data = prev_data; - core_record.read_data = read_data; - core_record.shift_amount = shift_amount; - - core_record.local_opcode = local_opcode as u8; - - let write_data = run_write_data( - local_opcode, - core_record.read_data, - core_record.prev_data, - core_record.shift_amount as usize, - ); - self.adapter - .write(state.memory, instruction, write_data, &mut adapter_record); - - *state.pc = state.pc.wrapping_add(DEFAULT_PC_STEP); - - Ok(()) - } -} - -impl TraceFiller for LoadStoreFiller -where - F: PrimeField32, - A: 'static + AdapterTraceFiller, -{ - fn fill_trace_row(&self, mem_helper: &MemoryAuxColsFactory, row_slice: &mut [F]) { - // SAFETY: row_slice is guaranteed by the caller to have at least A::WIDTH + - // LoadStoreCoreCols::width() elements - let (adapter_row, mut core_row) = unsafe { row_slice.split_at_mut_unchecked(A::WIDTH) }; - self.adapter.fill_trace_row(mem_helper, adapter_row); - // SAFETY: core_row contains a valid LoadStoreCoreRecord written by the executor - // during trace generation - let record: &LoadStoreCoreRecord = - unsafe { get_record_from_slice(&mut core_row, ()) }; - let core_row: &mut LoadStoreCoreCols = core_row.borrow_mut(); - - let opcode = Rv64LoadStoreOpcode::from_usize(record.local_opcode as usize); - let shift = record.shift_amount as usize; - let write_data = run_write_data(opcode, record.read_data, record.prev_data, shift); - - for pair in record.read_data.chunks_exact(2) { - self.bitwise_lookup_chip - .request_range(pair[0] as u32, pair[1] as u32); - } - for pair in record.prev_data.chunks_exact(2) { - self.bitwise_lookup_chip - .request_range(pair[0] as u32, pair[1] as u32); - } - - core_row.write_data = write_data.map(F::from_u8); - core_row.prev_data = record.prev_data.map(F::from_u8); - core_row.read_data = record.read_data.map(F::from_u8); - core_row.is_load = F::from_bool(matches!(opcode, LOADD | LOADWU | LOADHU | LOADBU)); - core_row.is_valid = F::ONE; - let pt: [u32; LOADSTORE_SELECTOR_WIDTH] = self - .encoder - .get_flag_pt(InstructionCase::from_opcode_shift(opcode, shift) as usize) - .try_into() - .unwrap(); - core_row.selector = pt.map(F::from_u32); - } -} - -// Returns the write data -#[inline(always)] -pub(super) fn run_write_data( - opcode: Rv64LoadStoreOpcode, - read_data: [u8; NUM_CELLS], - prev_data: [u8; NUM_CELLS], - shift: usize, -) -> [u8; NUM_CELLS] { - const { assert!(NUM_CELLS == RV64_REGISTER_NUM_LIMBS) }; - let word_width = NUM_CELLS / 2; - let half_width = NUM_CELLS / 4; - - match opcode { - LOADD if shift == 0 => read_data, - LOADWU if shift == 0 || shift == word_width => array::from_fn(|i| { - if i < word_width { - read_data[i + shift] - } else { - 0 - } - }), - LOADHU if [0, half_width, word_width, word_width + half_width].contains(&shift) => { - array::from_fn(|i| { - if i < half_width { - read_data[i + shift] - } else { - 0 - } - }) - } - LOADBU if shift < NUM_CELLS => array::from_fn(|i| { - if i == 0 { - read_data[shift] - } else { - 0 - } - }), - STORED if shift == 0 => read_data, - STOREW if shift == 0 || shift == word_width => array::from_fn(|i| { - if i >= shift && i < shift + word_width { - read_data[i - shift] - } else { - prev_data[i] - } - }), - STOREH if [0, half_width, word_width, word_width + half_width].contains(&shift) => { - array::from_fn(|i| { - if i >= shift && i < shift + half_width { - read_data[i - shift] - } else { - prev_data[i] - } - }) - } - STOREB if shift < NUM_CELLS => { - let mut write_data = prev_data; - write_data[shift] = read_data[0]; - write_data - } - _ => unreachable!( - "unaligned memory access not supported by this execution environment: {opcode:?}, shift: {shift}" - ), - } -} diff --git a/extensions/riscv/circuit/src/loadstore/cuda.rs b/extensions/riscv/circuit/src/loadstore/cuda.rs index a4ea1c3668..979def4160 100644 --- a/extensions/riscv/circuit/src/loadstore/cuda.rs +++ b/extensions/riscv/circuit/src/loadstore/cuda.rs @@ -1,67 +1,71 @@ -use std::{mem::size_of, sync::Arc}; +use std::sync::Arc; use derive_new::new; -use openvm_circuit::{arch::DenseRecordArena, utils::next_power_of_two_or_zero}; +use openvm_circuit::arch::DenseRecordArena; use openvm_circuit_primitives::{ bitwise_op_lookup::BitwiseOperationLookupChipGPU, var_range::VariableRangeCheckerChipGPU, Chip, }; -use openvm_cuda_backend::{base::DeviceMatrix, prelude::F, GpuBackend}; -use openvm_cuda_common::copy::MemCopyH2D; -use openvm_instructions::riscv::RV64_REGISTER_NUM_LIMBS; +use openvm_cuda_backend::{base::DeviceMatrix, GpuBackend}; use openvm_stark_backend::prover::AirProvingContext; -use crate::{ - adapters::{Rv64LoadStoreAdapterCols, Rv64LoadStoreAdapterRecord, RV64_BYTE_BITS}, - cuda_abi::loadstore_cuda::tracegen, - LoadStoreCoreCols, LoadStoreCoreRecord, -}; +use crate::adapters::RV64_BYTE_BITS; + +fn unsupported_split_loadstore_ctx(arena: DenseRecordArena) -> AirProvingContext { + if arena.allocated().is_empty() { + return AirProvingContext::simple_no_pis(DeviceMatrix::dummy()); + } + unimplemented!("CUDA trace generation for split RV64 loadstore chips is not implemented") +} #[derive(new)] -pub struct Rv64LoadStoreChipGpu { +pub struct Rv64LoadStoreByteChipGpu { pub range_checker: Arc, pub bitwise_lookup: Arc>, pub pointer_max_bits: usize, pub timestamp_max_bits: usize, } -impl Chip for Rv64LoadStoreChipGpu { +impl Chip for Rv64LoadStoreByteChipGpu { fn generate_proving_ctx(&self, arena: DenseRecordArena) -> AirProvingContext { - const RECORD_SIZE: usize = size_of::<( - Rv64LoadStoreAdapterRecord, - LoadStoreCoreRecord, - )>(); - let records = arena.allocated(); - if records.is_empty() { - return AirProvingContext::simple_no_pis(DeviceMatrix::dummy()); - } - debug_assert_eq!(records.len() % RECORD_SIZE, 0); + unsupported_split_loadstore_ctx(arena) + } +} + +#[derive(new)] +pub struct Rv64LoadStoreHalfwordChipGpu { + pub range_checker: Arc, + pub pointer_max_bits: usize, + pub timestamp_max_bits: usize, +} + +impl Chip for Rv64LoadStoreHalfwordChipGpu { + fn generate_proving_ctx(&self, arena: DenseRecordArena) -> AirProvingContext { + unsupported_split_loadstore_ctx(arena) + } +} - let trace_width = Rv64LoadStoreAdapterCols::::width() - + LoadStoreCoreCols::::width(); - let height = records.len() / RECORD_SIZE; - let padded_height = next_power_of_two_or_zero(height); - let device_ctx = &self.range_checker.device_ctx; +#[derive(new)] +pub struct Rv64LoadStoreWordChipGpu { + pub range_checker: Arc, + pub pointer_max_bits: usize, + pub timestamp_max_bits: usize, +} - let d_records = tracing::info_span!("trace_gen.h2d_records") - .in_scope(|| records.to_device_on(device_ctx)) - .unwrap(); - let d_trace = DeviceMatrix::::with_capacity_on(padded_height, trace_width, device_ctx); +impl Chip for Rv64LoadStoreWordChipGpu { + fn generate_proving_ctx(&self, arena: DenseRecordArena) -> AirProvingContext { + unsupported_split_loadstore_ctx(arena) + } +} - unsafe { - tracegen( - d_trace.buffer(), - padded_height, - trace_width, - &d_records, - self.pointer_max_bits, - &self.range_checker.count, - &self.bitwise_lookup.count, - self.timestamp_max_bits as u32, - device_ctx.stream.as_raw(), - ) - .unwrap(); - } +#[derive(new)] +pub struct Rv64LoadStoreDoublewordChipGpu { + pub range_checker: Arc, + pub pointer_max_bits: usize, + pub timestamp_max_bits: usize, +} - AirProvingContext::simple_no_pis(d_trace) +impl Chip for Rv64LoadStoreDoublewordChipGpu { + fn generate_proving_ctx(&self, arena: DenseRecordArena) -> AirProvingContext { + unsupported_split_loadstore_ctx(arena) } } diff --git a/extensions/riscv/circuit/src/loadstore/doubleword/core.rs b/extensions/riscv/circuit/src/loadstore/doubleword/core.rs new file mode 100644 index 0000000000..df5b4a092a --- /dev/null +++ b/extensions/riscv/circuit/src/loadstore/doubleword/core.rs @@ -0,0 +1,134 @@ +use std::borrow::{Borrow, BorrowMut}; + +use openvm_circuit::{arch::*, system::memory::MemoryAuxColsFactory}; +use openvm_circuit_primitives::{ + var_range::{SharedVariableRangeCheckerChip, VariableRangeCheckerBus}, + AlignedBorrow, ColumnsAir, StructReflection, StructReflectionHelper, +}; +use openvm_instructions::LocalOpcode; +use openvm_riscv_transpiler::Rv64LoadStoreOpcode::{self, *}; +use openvm_stark_backend::{ + interaction::InteractionBuilder, + p3_air::BaseAir, + p3_field::{Field, PrimeCharacteristicRing, PrimeField32}, + BaseAirWithPublicValues, +}; + +use crate::{ + adapters::{LoadStoreInstruction, Rv64LoadStoreAdapterFiller}, + loadstore::common::{adapter_context, LoadStoreRecord}, +}; + +#[repr(C)] +#[derive(Debug, Clone, AlignedBorrow, StructReflection)] +pub struct LoadStoreDoublewordCoreCols { + pub is_valid: T, + pub is_load: T, + pub read_data: [T; BLOCK_FE_WIDTH], + pub prev_data: [T; BLOCK_FE_WIDTH], +} + +#[derive(Debug, Clone, ColumnsAir)] +#[columns_via(LoadStoreDoublewordCoreCols)] +pub struct LoadStoreDoublewordCoreAir { + pub offset: usize, +} + +impl LoadStoreDoublewordCoreAir { + pub fn new(offset: usize, _range_bus: VariableRangeCheckerBus) -> Self { + Self { offset } + } +} + +impl BaseAir for LoadStoreDoublewordCoreAir { + fn width(&self) -> usize { + LoadStoreDoublewordCoreCols::::width() + } +} + +impl BaseAirWithPublicValues for LoadStoreDoublewordCoreAir {} + +impl VmCoreAir for LoadStoreDoublewordCoreAir +where + AB: InteractionBuilder, + I: VmAdapterInterface, + I::Reads: From<([AB::Var; BLOCK_FE_WIDTH], [AB::Expr; BLOCK_FE_WIDTH])>, + I::Writes: From<[[AB::Expr; BLOCK_FE_WIDTH]; 1]>, + I::ProcessedInstruction: From>, +{ + fn eval( + &self, + builder: &mut AB, + local_core: &[AB::Var], + _from_pc: AB::Var, + ) -> AdapterAirContext { + let cols: &LoadStoreDoublewordCoreCols = (*local_core).borrow(); + + builder.assert_bool(cols.is_valid); + builder.assert_bool(cols.is_load); + builder.assert_zero(cols.is_load * (AB::Expr::ONE - cols.is_valid)); + + let is_store = cols.is_valid - cols.is_load; + let expected_opcode = cols.is_load * AB::Expr::from_u8(LOADD as u8) + + is_store * AB::Expr::from_u8(STORED as u8); + let expected_opcode = VmCoreAir::::expr_to_global_expr(self, expected_opcode); + let write_data = cols.read_data.map(Into::into); + + adapter_context::( + cols.is_valid.into(), + cols.is_load.into(), + expected_opcode, + AB::Expr::ZERO, + AB::Expr::ZERO, + cols.read_data, + cols.prev_data, + write_data, + ) + } + + fn start_offset(&self) -> usize { + self.offset + } +} + +#[derive(Clone)] +pub struct LoadStoreDoublewordFiller { + adapter: A, + pub offset: usize, +} + +impl LoadStoreDoublewordFiller { + pub fn new( + adapter: A, + offset: usize, + _range_checker_chip: SharedVariableRangeCheckerChip, + ) -> Self { + Self { adapter, offset } + } +} + +impl TraceFiller for LoadStoreDoublewordFiller +where + F: PrimeField32, + A: 'static + AdapterTraceFiller, +{ + fn fill_trace_row(&self, mem_helper: &MemoryAuxColsFactory, row_slice: &mut [F]) { + let (adapter_row, mut core_row) = unsafe { row_slice.split_at_mut_unchecked(A::WIDTH) }; + self.adapter.fill_trace_row(mem_helper, adapter_row); + + let record: &LoadStoreRecord = unsafe { get_record_from_slice(&mut core_row, ()) }; + let opcode = Rv64LoadStoreOpcode::from_usize(record.local_opcode as usize); + let read_data = record.read_data; + let prev_data = record.prev_data; + debug_assert_eq!(record.shift_amount, 0); + let core_row: &mut LoadStoreDoublewordCoreCols = core_row.borrow_mut(); + + if !matches!(opcode, LOADD | STORED) { + unreachable!("doubleword loadstore core only handles LOADD/STORED"); + } + core_row.is_valid = F::ONE; + core_row.is_load = F::from_bool(opcode == LOADD); + core_row.read_data = read_data.map(F::from_u16); + core_row.prev_data = prev_data.map(F::from_u16); + } +} diff --git a/extensions/riscv/circuit/src/loadstore/doubleword/mod.rs b/extensions/riscv/circuit/src/loadstore/doubleword/mod.rs new file mode 100644 index 0000000000..720722b24a --- /dev/null +++ b/extensions/riscv/circuit/src/loadstore/doubleword/mod.rs @@ -0,0 +1,15 @@ +use openvm_circuit::arch::{VmAirWrapper, VmChipWrapper}; + +use crate::{ + adapters::{Rv64LoadStoreAdapterAir, Rv64LoadStoreAdapterExecutor}, + loadstore::common::{LoadStoreExecutor, KIND_DOUBLEWORD}, +}; + +mod core; +pub use core::*; + +pub type Rv64LoadStoreDoublewordAir = + VmAirWrapper; +pub type Rv64LoadStoreDoublewordExecutor = + LoadStoreExecutor; +pub type Rv64LoadStoreDoublewordChip = VmChipWrapper; diff --git a/extensions/riscv/circuit/src/loadstore/execution.rs b/extensions/riscv/circuit/src/loadstore/execution.rs index 3c74702e9a..1be6381105 100644 --- a/extensions/riscv/circuit/src/loadstore/execution.rs +++ b/extensions/riscv/circuit/src/loadstore/execution.rs @@ -9,14 +9,14 @@ use openvm_circuit_primitives::AlignedBytesBorrow; use openvm_instructions::{ instruction::Instruction, program::DEFAULT_PC_STEP, - riscv::{RV64_MEMORY_AS, RV64_REGISTER_AS, RV64_REGISTER_NUM_LIMBS}, - LocalOpcode, PUBLIC_VALUES_AS, + riscv::{RV64_IMM_AS, RV64_REGISTER_AS, RV64_REGISTER_NUM_LIMBS}, + LocalOpcode, DEFERRAL_AS, }; use openvm_riscv_transpiler::Rv64LoadStoreOpcode::{self, *}; use openvm_stark_backend::p3_field::PrimeField32; -use super::core::LoadStoreExecutor; -use crate::adapters::{rv64_address_add_imm, rv64_bytes_to_u32, sign_extend_imm16}; +use super::common::LoadStoreExecutor; +use crate::adapters::{rv64_bytes_to_u32, sign_extend_imm16}; #[derive(AlignedBytesBorrow, Clone)] #[repr(C)] @@ -27,14 +27,14 @@ struct LoadStorePreCompute { e: u8, } -impl LoadStoreExecutor { - /// Return (local_opcode, enabled) +impl LoadStoreExecutor { + /// Return (local_opcode, enabled, is_native_store) fn pre_compute_impl( &self, pc: u32, inst: &Instruction, data: &mut LoadStorePreCompute, - ) -> Result<(Rv64LoadStoreOpcode, bool), StaticProgramError> { + ) -> Result<(Rv64LoadStoreOpcode, bool, bool), StaticProgramError> { let Instruction { opcode, a, @@ -48,33 +48,27 @@ impl LoadStoreExecutor { } = inst; let enabled = !f.is_zero(); - let local_opcode = Rv64LoadStoreOpcode::from_usize( - opcode.local_opcode_idx(Rv64LoadStoreOpcode::CLASS_OFFSET), - ); let e_u32 = e.as_canonical_u32(); - let valid_address_space = match local_opcode { - LOADD | LOADWU | LOADHU | LOADBU => e_u32 == RV64_MEMORY_AS, - STORED | STOREW | STOREH | STOREB => { - e_u32 == RV64_MEMORY_AS || e_u32 == PUBLIC_VALUES_AS - } - _ => false, - }; - if d.as_canonical_u32() != RV64_REGISTER_AS || !valid_address_space { + if d.as_canonical_u32() != RV64_REGISTER_AS || e_u32 == RV64_IMM_AS { return Err(StaticProgramError::InvalidInstruction(pc)); } + + let local_opcode = Rv64LoadStoreOpcode::from_usize( + opcode.local_opcode_idx(Rv64LoadStoreOpcode::CLASS_OFFSET), + ); match local_opcode { - LOADD | LOADWU | LOADHU | LOADBU => {} + LOADD | LOADW | LOADH | LOADWU | LOADHU | LOADBU | LOADB => {} STORED | STOREW | STOREH | STOREB => { if !enabled { return Err(StaticProgramError::InvalidInstruction(pc)); } } - _ => unreachable!("LoadStoreExecutor should not handle sign-extension load opcodes"), } let imm = c.as_canonical_u32(); let imm_sign = g.as_canonical_u32(); let imm_extended = sign_extend_imm16(imm, imm_sign); + let is_native_store = e_u32 == DEFERRAL_AS; *data = LoadStorePreCompute { imm_extended, @@ -82,35 +76,48 @@ impl LoadStoreExecutor { b: b.as_canonical_u32() as u8, e: e_u32 as u8, }; - Ok((local_opcode, enabled)) + Ok((local_opcode, enabled, is_native_store)) } } macro_rules! dispatch { - ($execute_impl:ident, $local_opcode:ident, $enabled:ident) => { - match ($local_opcode, $enabled) { - (LOADD, true) => Ok($execute_impl::<_, _, U8, LoadDOp, true>), - (LOADD, false) => Ok($execute_impl::<_, _, U8, LoadDOp, false>), - (LOADWU, true) => Ok($execute_impl::<_, _, U8, LoadWUOp, true>), - (LOADWU, false) => Ok($execute_impl::<_, _, U8, LoadWUOp, false>), - (LOADHU, true) => Ok($execute_impl::<_, _, U8, LoadHUOp, true>), - (LOADHU, false) => Ok($execute_impl::<_, _, U8, LoadHUOp, false>), - (LOADBU, true) => Ok($execute_impl::<_, _, U8, LoadBUOp, true>), - (LOADBU, false) => Ok($execute_impl::<_, _, U8, LoadBUOp, false>), - (STORED, true) => Ok($execute_impl::<_, _, U8, StoreDOp, true>), - (STORED, false) => Ok($execute_impl::<_, _, U8, StoreDOp, false>), - (STOREW, true) => Ok($execute_impl::<_, _, U8, StoreWOp, true>), - (STOREW, false) => Ok($execute_impl::<_, _, U8, StoreWOp, false>), - (STOREH, true) => Ok($execute_impl::<_, _, U8, StoreHOp, true>), - (STOREH, false) => Ok($execute_impl::<_, _, U8, StoreHOp, false>), - (STOREB, true) => Ok($execute_impl::<_, _, U8, StoreBOp, true>), - (STOREB, false) => Ok($execute_impl::<_, _, U8, StoreBOp, false>), - (_, _) => unreachable!(), + ($execute_impl:ident, $local_opcode:ident, $enabled:ident, $is_native_store:ident) => { + match ($local_opcode, $enabled, $is_native_store) { + (LOADD, true, _) => Ok($execute_impl::<_, _, U8, LoadDOp, true>), + (LOADD, false, _) => Ok($execute_impl::<_, _, U8, LoadDOp, false>), + (LOADWU, true, _) => Ok($execute_impl::<_, _, U8, LoadWUOp, true>), + (LOADWU, false, _) => Ok($execute_impl::<_, _, U8, LoadWUOp, false>), + (LOADW, true, _) => Ok($execute_impl::<_, _, U8, LoadWOp, true>), + (LOADW, false, _) => Ok($execute_impl::<_, _, U8, LoadWOp, false>), + (LOADHU, true, _) => Ok($execute_impl::<_, _, U8, LoadHUOp, true>), + (LOADHU, false, _) => Ok($execute_impl::<_, _, U8, LoadHUOp, false>), + (LOADH, true, _) => Ok($execute_impl::<_, _, U8, LoadHOp, true>), + (LOADH, false, _) => Ok($execute_impl::<_, _, U8, LoadHOp, false>), + (LOADBU, true, _) => Ok($execute_impl::<_, _, U8, LoadBUOp, true>), + (LOADBU, false, _) => Ok($execute_impl::<_, _, U8, LoadBUOp, false>), + (LOADB, true, _) => Ok($execute_impl::<_, _, U8, LoadBOp, true>), + (LOADB, false, _) => Ok($execute_impl::<_, _, U8, LoadBOp, false>), + (STORED, true, false) => Ok($execute_impl::<_, _, U8, StoreDOp, true>), + (STORED, false, false) => Ok($execute_impl::<_, _, U8, StoreDOp, false>), + (STORED, true, true) => Ok($execute_impl::<_, _, F, StoreDOp, true>), + (STORED, false, true) => Ok($execute_impl::<_, _, F, StoreDOp, false>), + (STOREW, true, false) => Ok($execute_impl::<_, _, U8, StoreWOp, true>), + (STOREW, false, false) => Ok($execute_impl::<_, _, U8, StoreWOp, false>), + (STOREW, true, true) => Ok($execute_impl::<_, _, F, StoreWOp, true>), + (STOREW, false, true) => Ok($execute_impl::<_, _, F, StoreWOp, false>), + (STOREH, true, false) => Ok($execute_impl::<_, _, U8, StoreHOp, true>), + (STOREH, false, false) => Ok($execute_impl::<_, _, U8, StoreHOp, false>), + (STOREH, true, true) => Ok($execute_impl::<_, _, F, StoreHOp, true>), + (STOREH, false, true) => Ok($execute_impl::<_, _, F, StoreHOp, false>), + (STOREB, true, false) => Ok($execute_impl::<_, _, U8, StoreBOp, true>), + (STOREB, false, false) => Ok($execute_impl::<_, _, U8, StoreBOp, false>), + (STOREB, true, true) => Ok($execute_impl::<_, _, F, StoreBOp, true>), + (STOREB, false, true) => Ok($execute_impl::<_, _, F, StoreBOp, false>), } }; } -impl InterpreterExecutor for LoadStoreExecutor +impl InterpreterExecutor for LoadStoreExecutor where F: PrimeField32, { @@ -128,8 +135,9 @@ where data: &mut [u8], ) -> Result, StaticProgramError> { let pre_compute: &mut LoadStorePreCompute = data.borrow_mut(); - let (local_opcode, enabled) = self.pre_compute_impl(pc, inst, pre_compute)?; - dispatch!(execute_e1_handler, local_opcode, enabled) + let (local_opcode, enabled, is_native_store) = + self.pre_compute_impl(pc, inst, pre_compute)?; + dispatch!(execute_e1_handler, local_opcode, enabled, is_native_store) } #[cfg(feature = "tco")] @@ -143,12 +151,13 @@ where Ctx: ExecutionCtxTrait, { let pre_compute: &mut LoadStorePreCompute = data.borrow_mut(); - let (local_opcode, enabled) = self.pre_compute_impl(pc, inst, pre_compute)?; - dispatch!(execute_e1_handler, local_opcode, enabled) + let (local_opcode, enabled, is_native_store) = + self.pre_compute_impl(pc, inst, pre_compute)?; + dispatch!(execute_e1_handler, local_opcode, enabled, is_native_store) } } -impl InterpreterMeteredExecutor for LoadStoreExecutor +impl InterpreterMeteredExecutor for LoadStoreExecutor where F: PrimeField32, { @@ -169,8 +178,9 @@ where { let pre_compute: &mut E2PreCompute = data.borrow_mut(); pre_compute.chip_idx = chip_idx as u32; - let (local_opcode, enabled) = self.pre_compute_impl(pc, inst, &mut pre_compute.data)?; - dispatch!(execute_e2_handler, local_opcode, enabled) + let (local_opcode, enabled, is_native_store) = + self.pre_compute_impl(pc, inst, &mut pre_compute.data)?; + dispatch!(execute_e2_handler, local_opcode, enabled, is_native_store) } #[cfg(feature = "tco")] @@ -186,8 +196,9 @@ where { let pre_compute: &mut E2PreCompute = data.borrow_mut(); pre_compute.chip_idx = chip_idx as u32; - let (local_opcode, enabled) = self.pre_compute_impl(pc, inst, &mut pre_compute.data)?; - dispatch!(execute_e2_handler, local_opcode, enabled) + let (local_opcode, enabled, is_native_store) = + self.pre_compute_impl(pc, inst, &mut pre_compute.data)?; + dispatch!(execute_e2_handler, local_opcode, enabled, is_native_store) } } @@ -206,9 +217,9 @@ unsafe fn execute_e12_impl< let rs1_bytes: [u8; RV64_REGISTER_NUM_LIMBS] = exec_state.vm_read_bytes(RV64_REGISTER_AS, pre_compute.b as u32); let rs1_val = rv64_bytes_to_u32(rs1_bytes); - let addr = rv64_address_add_imm(rs1_val, pre_compute.imm_extended); - debug_assert!((addr as usize) < RV64_MEMORY_BYTES); - let ptr_val = addr as u32; + let ptr_val = rs1_val.wrapping_add(pre_compute.imm_extended); + // sign_extend([r64{c,g}(b):2]_e) + debug_assert!((ptr_val as usize) < RV64_MEMORY_BYTES); let shift_amount = ptr_val % RV64_REGISTER_NUM_LIMBS as u32; let ptr_val = ptr_val - shift_amount; // aligned ptr @@ -304,7 +315,10 @@ struct U8(u8); struct LoadDOp; struct LoadWUOp; struct LoadHUOp; +struct LoadHOp; struct LoadBUOp; +struct LoadBOp; +struct LoadWOp; struct StoreDOp; struct StoreWOp; struct StoreHOp; @@ -365,6 +379,49 @@ impl LoadStoreOp for LoadHUOp { } } +impl LoadStoreOp for LoadHOp { + const IS_LOAD: bool = true; + const HOST_READ: bool = false; + + #[inline(always)] + fn compute_write_data( + write_data: &mut [U8; RV64_REGISTER_NUM_LIMBS], + read_data: [u8; RV64_REGISTER_NUM_LIMBS], + shift_amount: usize, + ) -> bool { + if shift_amount != 0 && shift_amount != 2 && shift_amount != 4 && shift_amount != 6 { + return false; + } + let half = i16::from_le_bytes([read_data[shift_amount], read_data[shift_amount + 1]]); + *write_data = (half as i64).to_le_bytes().map(U8); + true + } +} + +impl LoadStoreOp for LoadWOp { + const IS_LOAD: bool = true; + const HOST_READ: bool = false; + + #[inline(always)] + fn compute_write_data( + write_data: &mut [U8; RV64_REGISTER_NUM_LIMBS], + read_data: [u8; RV64_REGISTER_NUM_LIMBS], + shift_amount: usize, + ) -> bool { + if shift_amount != 0 && shift_amount != 4 { + return false; + } + let word = i32::from_le_bytes([ + read_data[shift_amount], + read_data[shift_amount + 1], + read_data[shift_amount + 2], + read_data[shift_amount + 3], + ]); + *write_data = (word as i64).to_le_bytes().map(U8); + true + } +} + impl LoadStoreOp for LoadBUOp { const IS_LOAD: bool = true; const HOST_READ: bool = false; @@ -380,6 +437,22 @@ impl LoadStoreOp for LoadBUOp { } } +impl LoadStoreOp for LoadBOp { + const IS_LOAD: bool = true; + const HOST_READ: bool = false; + + #[inline(always)] + fn compute_write_data( + write_data: &mut [U8; RV64_REGISTER_NUM_LIMBS], + read_data: [u8; RV64_REGISTER_NUM_LIMBS], + shift_amount: usize, + ) -> bool { + let byte = read_data[shift_amount] as i8; + *write_data = (byte as i64).to_le_bytes().map(U8); + true + } +} + impl LoadStoreOp for StoreDOp { const IS_LOAD: bool = false; const HOST_READ: bool = false; @@ -450,51 +523,72 @@ impl LoadStoreOp for StoreBOp { } } -#[cfg(test)] -mod tests { - use openvm_circuit::arch::StaticProgramError; - use openvm_instructions::{ - instruction::Instruction, riscv::RV64_REGISTER_AS, LocalOpcode, DEFERRAL_AS, - }; - use openvm_riscv_transpiler::Rv64LoadStoreOpcode::{LOADWU, STORED}; - use openvm_stark_sdk::p3_baby_bear::BabyBear; +impl LoadStoreOp for StoreDOp { + const IS_LOAD: bool = false; + const HOST_READ: bool = false; - use super::{LoadStorePreCompute, Rv64LoadStoreOpcode}; - use crate::{adapters::Rv64LoadStoreAdapterExecutor, Rv64LoadStoreExecutor}; + #[inline(always)] + fn compute_write_data( + write_data: &mut [F; RV64_REGISTER_NUM_LIMBS], + read_data: [u8; RV64_REGISTER_NUM_LIMBS], + _shift_amount: usize, + ) -> bool { + *write_data = read_data.map(F::from_u8); + true + } +} - #[test] - fn precompute_enforces_address_space_domain() { - const PC: u32 = 0x100; +impl LoadStoreOp for StoreWOp { + const IS_LOAD: bool = false; + const HOST_READ: bool = true; - let executor = Rv64LoadStoreExecutor::new( - Rv64LoadStoreAdapterExecutor::new(29), - Rv64LoadStoreOpcode::CLASS_OFFSET, - ); - for opcode in [LOADWU, STORED] { - let inst = Instruction::::from_usize( - opcode.global_opcode(), - [ - 8, - 16, - 0, - RV64_REGISTER_AS as usize, - DEFERRAL_AS as usize, - 1, - 0, - ], - ); - let mut data = LoadStorePreCompute { - imm_extended: 0, - a: 0, - b: 0, - e: 0, - }; - - let err = executor - .pre_compute_impl(PC, &inst, &mut data) - .expect_err("load/store address-space domain should be enforced"); - - assert!(matches!(err, StaticProgramError::InvalidInstruction(PC))); + #[inline(always)] + fn compute_write_data( + write_data: &mut [F; RV64_REGISTER_NUM_LIMBS], + read_data: [u8; RV64_REGISTER_NUM_LIMBS], + shift_amount: usize, + ) -> bool { + if shift_amount != 0 && shift_amount != 4 { + return false; + } + write_data[shift_amount] = F::from_u8(read_data[0]); + write_data[shift_amount + 1] = F::from_u8(read_data[1]); + write_data[shift_amount + 2] = F::from_u8(read_data[2]); + write_data[shift_amount + 3] = F::from_u8(read_data[3]); + true + } +} + +impl LoadStoreOp for StoreHOp { + const IS_LOAD: bool = false; + const HOST_READ: bool = true; + + #[inline(always)] + fn compute_write_data( + write_data: &mut [F; RV64_REGISTER_NUM_LIMBS], + read_data: [u8; RV64_REGISTER_NUM_LIMBS], + shift_amount: usize, + ) -> bool { + if shift_amount != 0 && shift_amount != 2 && shift_amount != 4 && shift_amount != 6 { + return false; } + write_data[shift_amount] = F::from_u8(read_data[0]); + write_data[shift_amount + 1] = F::from_u8(read_data[1]); + true + } +} + +impl LoadStoreOp for StoreBOp { + const IS_LOAD: bool = false; + const HOST_READ: bool = true; + + #[inline(always)] + fn compute_write_data( + write_data: &mut [F; RV64_REGISTER_NUM_LIMBS], + read_data: [u8; RV64_REGISTER_NUM_LIMBS], + shift_amount: usize, + ) -> bool { + write_data[shift_amount] = F::from_u8(read_data[0]); + true } } diff --git a/extensions/riscv/circuit/src/loadstore/halfword/mod.rs b/extensions/riscv/circuit/src/loadstore/halfword/mod.rs new file mode 100644 index 0000000000..77e4736fcd --- /dev/null +++ b/extensions/riscv/circuit/src/loadstore/halfword/mod.rs @@ -0,0 +1,26 @@ +use openvm_circuit::arch::{VmAirWrapper, VmChipWrapper}; + +use crate::{ + adapters::{Rv64LoadStoreAdapterAir, Rv64LoadStoreAdapterExecutor}, + loadstore::{ + aligned::{LoadStoreAlignedCoreAir, LoadStoreAlignedFiller}, + common::{LoadStoreExecutor, KIND_HALFWORD}, + }, +}; + +pub const HALFWORD_CASES: usize = 8; +pub const HALFWORD_SELECTOR_WIDTH: usize = 3; + +pub type LoadStoreHalfwordCoreAir = + LoadStoreAlignedCoreAir; +pub type LoadStoreHalfwordFiller = LoadStoreAlignedFiller< + crate::adapters::Rv64LoadStoreAdapterFiller, + KIND_HALFWORD, + HALFWORD_CASES, + HALFWORD_SELECTOR_WIDTH, +>; + +pub type Rv64LoadStoreHalfwordAir = VmAirWrapper; +pub type Rv64LoadStoreHalfwordExecutor = + LoadStoreExecutor; +pub type Rv64LoadStoreHalfwordChip = VmChipWrapper; diff --git a/extensions/riscv/circuit/src/loadstore/mod.rs b/extensions/riscv/circuit/src/loadstore/mod.rs index d01e5cfe5d..b12a609aa5 100644 --- a/extensions/riscv/circuit/src/loadstore/mod.rs +++ b/extensions/riscv/circuit/src/loadstore/mod.rs @@ -1,25 +1,25 @@ -use openvm_circuit::arch::{VmAirWrapper, VmChipWrapper}; +pub mod aligned; +pub mod byte; +pub mod common; +pub mod doubleword; +pub mod halfword; +pub mod word; -use super::adapters::RV64_REGISTER_NUM_LIMBS; -use crate::adapters::{Rv64LoadStoreAdapterAir, Rv64LoadStoreAdapterExecutor}; +pub use byte::*; +pub use common::{LoadStoreExecutor, LoadStoreRecord}; +pub use doubleword::*; +pub use halfword::*; +pub use word::*; -mod core; mod execution; -pub use core::*; - #[cfg(feature = "cuda")] mod cuda; #[cfg(feature = "cuda")] pub use cuda::*; + #[cfg(feature = "aot")] mod aot; #[cfg(test)] mod tests; - -pub type Rv64LoadStoreAir = - VmAirWrapper>; -pub type Rv64LoadStoreExecutor = - LoadStoreExecutor; -pub type Rv64LoadStoreChip = VmChipWrapper; diff --git a/extensions/riscv/circuit/src/loadstore/tests.rs b/extensions/riscv/circuit/src/loadstore/tests.rs index ab93679f2c..85ea55381c 100644 --- a/extensions/riscv/circuit/src/loadstore/tests.rs +++ b/extensions/riscv/circuit/src/loadstore/tests.rs @@ -3,24 +3,19 @@ use std::{array, borrow::BorrowMut, sync::Arc}; use openvm_circuit::{ arch::{ testing::{TestBuilder, TestChipHarness, VmChipTestBuilder, BITWISE_OP_LOOKUP_BUS}, - Arena, ExecutionBridge, MemoryConfig, PreflightExecutor, - }, - system::memory::{ - merkle::public_values::PUBLIC_VALUES_AS, offline_checker::MemoryBridge, SharedMemoryHelper, + Arena, MemoryConfig, PreflightExecutor, BLOCK_FE_WIDTH, }, + system::memory::merkle::public_values::PUBLIC_VALUES_AS, }; -use openvm_circuit_primitives::{ - bitwise_op_lookup::{ - BitwiseOperationLookupAir, BitwiseOperationLookupBus, BitwiseOperationLookupChip, - SharedBitwiseOperationLookupChip, - }, - var_range::VariableRangeCheckerChip, +use openvm_circuit_primitives::bitwise_op_lookup::{ + BitwiseOperationLookupAir, BitwiseOperationLookupBus, BitwiseOperationLookupChip, + SharedBitwiseOperationLookupChip, }; use openvm_instructions::{instruction::Instruction, riscv::RV64_REGISTER_AS, LocalOpcode}; use openvm_riscv_transpiler::Rv64LoadStoreOpcode::{self, *}; use openvm_stark_backend::{ p3_air::BaseAir, - p3_field::{PrimeCharacteristicRing, PrimeField32}, + p3_field::PrimeCharacteristicRing, p3_matrix::{ dense::{DenseMatrix, RowMajorMatrix}, Matrix, @@ -30,123 +25,207 @@ use openvm_stark_backend::{ use openvm_stark_sdk::{p3_baby_bear::BabyBear, utils::create_seeded_rng}; use rand::{rngs::StdRng, seq::IndexedRandom, Rng}; use test_case::test_case; -#[cfg(feature = "cuda")] -use { - crate::{adapters::Rv64LoadStoreAdapterRecord, LoadStoreCoreRecord, Rv64LoadStoreChipGpu}, - openvm_circuit::arch::{ - testing::{ - default_bitwise_lookup_bus, default_var_range_checker_bus, dummy_range_checker, - GpuChipTestBuilder, GpuTestChipHarness, - }, - EmptyAdapterCoreLayout, - }, - openvm_instructions::riscv::RV64_MEMORY_AS, -}; use super::{ - run_write_data, selector_point_for_opcode_shift, LoadStoreCoreAir, LoadStoreCoreCols, - Rv64LoadStoreChip, LOADSTORE_SELECTOR_WIDTH, -}; -use crate::{ - adapters::{ - pack_u8_pair_u32, rv64_bytes_to_u32, sign_extend_imm16, Rv64LoadStoreAdapterAir, - Rv64LoadStoreAdapterCols, Rv64LoadStoreAdapterExecutor, Rv64LoadStoreAdapterFiller, - RV64_BYTE_BITS, RV64_REGISTER_NUM_LIMBS, RV64_WORD_NUM_LIMBS, + aligned::LoadStoreAlignedCoreCols, + byte::{LoadStoreByteCoreAir, LoadStoreByteCoreCols, LoadStoreByteFiller}, + common::run_write_data, + doubleword::{ + LoadStoreDoublewordCoreAir, LoadStoreDoublewordCoreCols, LoadStoreDoublewordFiller, }, - LoadStoreFiller, Rv64LoadStoreAir, Rv64LoadStoreExecutor, + halfword::{LoadStoreHalfwordCoreAir, LoadStoreHalfwordFiller, HALFWORD_SELECTOR_WIDTH}, + word::{LoadStoreWordCoreAir, LoadStoreWordFiller, WORD_SELECTOR_WIDTH}, + Rv64LoadStoreByteAir, Rv64LoadStoreByteChip, Rv64LoadStoreByteExecutor, + Rv64LoadStoreDoublewordAir, Rv64LoadStoreDoublewordChip, Rv64LoadStoreDoublewordExecutor, + Rv64LoadStoreHalfwordAir, Rv64LoadStoreHalfwordChip, Rv64LoadStoreHalfwordExecutor, + Rv64LoadStoreWordAir, Rv64LoadStoreWordChip, Rv64LoadStoreWordExecutor, +}; +use crate::adapters::{ + rv64_bytes_to_u16_block, rv64_bytes_to_u32, rv64_u16_block_to_bytes, sign_extend_imm16, + Rv64LoadStoreAdapterAir, Rv64LoadStoreAdapterExecutor, Rv64LoadStoreAdapterFiller, + RV64_BYTE_BITS, }; const IMM_BITS: usize = 16; const MAX_INS_CAPACITY: usize = 128; - type F = BabyBear; -type Harness = TestChipHarness>; - -fn create_harness_fields( - memory_bridge: MemoryBridge, - execution_bridge: ExecutionBridge, - range_checker_chip: Arc, - bitwise_chip: SharedBitwiseOperationLookupChip, - memory_helper: SharedMemoryHelper, - address_bits: usize, -) -> ( - Rv64LoadStoreAir, - Rv64LoadStoreExecutor, - Rv64LoadStoreChip, -) { - let air = Rv64LoadStoreAir::new( - Rv64LoadStoreAdapterAir::new( - memory_bridge, - execution_bridge, - range_checker_chip.bus(), - address_bits, - ), - LoadStoreCoreAir::new(Rv64LoadStoreOpcode::CLASS_OFFSET, bitwise_chip.bus()), - ); - let executor = Rv64LoadStoreExecutor::new( - Rv64LoadStoreAdapterExecutor::new(address_bits), - Rv64LoadStoreOpcode::CLASS_OFFSET, - ); - let chip = Rv64LoadStoreChip::::new( - LoadStoreFiller::new( - Rv64LoadStoreAdapterFiller::new(address_bits, range_checker_chip.clone()), - Rv64LoadStoreOpcode::CLASS_OFFSET, - bitwise_chip, - ), - memory_helper, - ); - (air, executor, chip) + +type ByteHarness = + TestChipHarness>; +type HalfwordHarness = TestChipHarness< + F, + Rv64LoadStoreHalfwordExecutor, + Rv64LoadStoreHalfwordAir, + Rv64LoadStoreHalfwordChip, +>; +type WordHarness = + TestChipHarness>; +type DoublewordHarness = TestChipHarness< + F, + Rv64LoadStoreDoublewordExecutor, + Rv64LoadStoreDoublewordAir, + Rv64LoadStoreDoublewordChip, +>; + +fn u16_block_to_f_bytes(block: [u16; BLOCK_FE_WIDTH]) -> [F; 8] { + rv64_u16_block_to_bytes(block).map(F::from_u8) } -fn create_harness( +fn create_byte_harness( tester: &mut VmChipTestBuilder, ) -> ( - Harness, + ByteHarness, ( BitwiseOperationLookupAir, SharedBitwiseOperationLookupChip, ), ) { + let range_checker = tester.range_checker(); let bitwise_bus = BitwiseOperationLookupBus::new(BITWISE_OP_LOOKUP_BUS); let bitwise_chip = Arc::new(BitwiseOperationLookupChip::::new( bitwise_bus, )); - let (air, executor, chip) = create_harness_fields( - tester.memory_bridge(), - tester.execution_bridge(), - tester.range_checker(), - bitwise_chip.clone(), + let air = Rv64LoadStoreByteAir::new( + Rv64LoadStoreAdapterAir::new( + tester.memory_bridge(), + tester.execution_bridge(), + range_checker.bus(), + tester.address_bits(), + ), + LoadStoreByteCoreAir::new( + Rv64LoadStoreOpcode::CLASS_OFFSET, + bitwise_chip.bus(), + range_checker.bus(), + ), + ); + let executor = Rv64LoadStoreByteExecutor::new( + Rv64LoadStoreAdapterExecutor::new(tester.address_bits()), + Rv64LoadStoreOpcode::CLASS_OFFSET, + ); + let chip = Rv64LoadStoreByteChip::::new( + LoadStoreByteFiller::new( + Rv64LoadStoreAdapterFiller::new(tester.address_bits(), range_checker.clone()), + Rv64LoadStoreOpcode::CLASS_OFFSET, + bitwise_chip.clone(), + range_checker, + ), tester.memory_helper(), - tester.address_bits(), ); ( - Harness::with_capacity(executor, air, chip, MAX_INS_CAPACITY), + ByteHarness::with_capacity(executor, air, chip, MAX_INS_CAPACITY), (bitwise_chip.air, bitwise_chip), ) } +fn create_halfword_harness(tester: &mut VmChipTestBuilder) -> HalfwordHarness { + let range_checker = tester.range_checker(); + let air = Rv64LoadStoreHalfwordAir::new( + Rv64LoadStoreAdapterAir::new( + tester.memory_bridge(), + tester.execution_bridge(), + range_checker.bus(), + tester.address_bits(), + ), + LoadStoreHalfwordCoreAir::new(Rv64LoadStoreOpcode::CLASS_OFFSET, range_checker.bus()), + ); + let executor = Rv64LoadStoreHalfwordExecutor::new( + Rv64LoadStoreAdapterExecutor::new(tester.address_bits()), + Rv64LoadStoreOpcode::CLASS_OFFSET, + ); + let chip = Rv64LoadStoreHalfwordChip::::new( + LoadStoreHalfwordFiller::new( + Rv64LoadStoreAdapterFiller::new(tester.address_bits(), range_checker.clone()), + Rv64LoadStoreOpcode::CLASS_OFFSET, + range_checker, + ), + tester.memory_helper(), + ); + HalfwordHarness::with_capacity(executor, air, chip, MAX_INS_CAPACITY) +} + +fn create_word_harness(tester: &mut VmChipTestBuilder) -> WordHarness { + let range_checker = tester.range_checker(); + let air = Rv64LoadStoreWordAir::new( + Rv64LoadStoreAdapterAir::new( + tester.memory_bridge(), + tester.execution_bridge(), + range_checker.bus(), + tester.address_bits(), + ), + LoadStoreWordCoreAir::new(Rv64LoadStoreOpcode::CLASS_OFFSET, range_checker.bus()), + ); + let executor = Rv64LoadStoreWordExecutor::new( + Rv64LoadStoreAdapterExecutor::new(tester.address_bits()), + Rv64LoadStoreOpcode::CLASS_OFFSET, + ); + let chip = Rv64LoadStoreWordChip::::new( + LoadStoreWordFiller::new( + Rv64LoadStoreAdapterFiller::new(tester.address_bits(), range_checker.clone()), + Rv64LoadStoreOpcode::CLASS_OFFSET, + range_checker, + ), + tester.memory_helper(), + ); + WordHarness::with_capacity(executor, air, chip, MAX_INS_CAPACITY) +} + +fn create_doubleword_harness(tester: &mut VmChipTestBuilder) -> DoublewordHarness { + let range_checker = tester.range_checker(); + let air = Rv64LoadStoreDoublewordAir::new( + Rv64LoadStoreAdapterAir::new( + tester.memory_bridge(), + tester.execution_bridge(), + range_checker.bus(), + tester.address_bits(), + ), + LoadStoreDoublewordCoreAir::new(Rv64LoadStoreOpcode::CLASS_OFFSET, range_checker.bus()), + ); + let executor = Rv64LoadStoreDoublewordExecutor::new( + Rv64LoadStoreAdapterExecutor::new(tester.address_bits()), + Rv64LoadStoreOpcode::CLASS_OFFSET, + ); + let chip = Rv64LoadStoreDoublewordChip::::new( + LoadStoreDoublewordFiller::new( + Rv64LoadStoreAdapterFiller::new(tester.address_bits(), range_checker.clone()), + Rv64LoadStoreOpcode::CLASS_OFFSET, + range_checker, + ), + tester.memory_helper(), + ); + DoublewordHarness::with_capacity(executor, air, chip, MAX_INS_CAPACITY) +} + #[allow(clippy::too_many_arguments)] fn set_and_execute>( - tester: &mut impl TestBuilder, + tester: &mut VmChipTestBuilder, + executor: &mut E, + arena: &mut RA, + rng: &mut StdRng, + opcode: Rv64LoadStoreOpcode, +) { + set_and_execute_with(tester, executor, arena, rng, opcode, None, None, None, None); +} + +#[allow(clippy::too_many_arguments)] +fn set_and_execute_with>( + tester: &mut VmChipTestBuilder, executor: &mut E, arena: &mut RA, rng: &mut StdRng, opcode: Rv64LoadStoreOpcode, - rs1: Option<[u8; RV64_REGISTER_NUM_LIMBS]>, + rs1: Option<[u8; 8]>, imm: Option, imm_sign: Option, mem_as: Option, ) { - let imm = imm.unwrap_or(rng.random_range(0..(1 << IMM_BITS))); - let imm_sign = imm_sign.unwrap_or(rng.random_range(0..2)); + let imm = imm.unwrap_or_else(|| rng.random_range(0..(1 << IMM_BITS))); + let imm_sign = imm_sign.unwrap_or_else(|| rng.random_range(0..2)); let imm_ext = sign_extend_imm16(imm, imm_sign); - let alignment = match opcode { LOADD | STORED => 3, - LOADWU | STOREW => 2, - LOADHU | STOREH => 1, - LOADBU | STOREB => 0, - _ => unreachable!("loadstore tests should not handle sign-extension load opcodes"), + LOADW | LOADWU | STOREW => 2, + LOADH | LOADHU | STOREH => 1, + LOADB | LOADBU | STOREB => 0, }; let ptr_val: u32 = rng.random_range(0..(1 << (tester.address_bits() - alignment))) << alignment; @@ -154,45 +233,51 @@ fn set_and_execute>( let rs1 = rs1.unwrap_or([ptr[0], ptr[1], ptr[2], ptr[3], 0, 0, 0, 0]); let rs1_low = rv64_bytes_to_u32(rs1); let ptr_val = imm_ext.wrapping_add(rs1_low); - let shift_amount = (ptr_val as usize) & (RV64_REGISTER_NUM_LIMBS - 1); + let shift_amount = (ptr_val as usize) & 7; + let base_ptr = (ptr_val as usize) - shift_amount; let max_addr = 1usize << tester.address_bits(); - let a = rng.random_range(0..(max_addr - RV64_REGISTER_NUM_LIMBS)) / RV64_REGISTER_NUM_LIMBS - * RV64_REGISTER_NUM_LIMBS; - let b = rng.random_range(0..(max_addr - RV64_REGISTER_NUM_LIMBS)) / RV64_REGISTER_NUM_LIMBS - * RV64_REGISTER_NUM_LIMBS; - - let is_load = [LOADD, LOADWU, LOADHU, LOADBU].contains(&opcode); - // Store tests choose writable u16-celled address spaces. - let mem_as = mem_as.unwrap_or(if is_load { - 2 - } else { - *[2, 3].choose(rng).unwrap() + let a = rng.random_range(0..(max_addr - 8)) / 8 * 8; + let b = rng.random_range(0..(max_addr - 8)) / 8 * 8; + let is_load = matches!( + opcode, + LOADD | LOADW | LOADH | LOADB | LOADWU | LOADHU | LOADBU + ); + let mem_as: usize = mem_as.unwrap_or_else(|| { + if is_load { + 2 + } else { + *[2usize, PUBLIC_VALUES_AS as usize].choose(rng).unwrap() + } }); - tester.write_bytes(1, b, rs1.map(F::from_u8)); - - let mut prev_data: [F; RV64_REGISTER_NUM_LIMBS] = - array::from_fn(|_| F::from_u32(rng.random_range(0..(1 << RV64_BYTE_BITS)))); - let mut read_data: [F; RV64_REGISTER_NUM_LIMBS] = - array::from_fn(|_| F::from_u32(rng.random_range(0..(1 << RV64_BYTE_BITS)))); + tester.write_bytes(RV64_REGISTER_AS as usize, b, rs1.map(F::from_u8)); + let mut prev_data: [u16; BLOCK_FE_WIDTH] = array::from_fn(|_| rng.random()); + let mut read_data: [u16; BLOCK_FE_WIDTH] = array::from_fn(|_| rng.random()); if is_load { if a == 0 { - prev_data = [F::ZERO; RV64_REGISTER_NUM_LIMBS]; + prev_data = [0; BLOCK_FE_WIDTH]; } - tester.write_bytes(1, a, prev_data); - tester.write_bytes(mem_as, (ptr_val as usize) - shift_amount, read_data); + tester.write_bytes( + RV64_REGISTER_AS as usize, + a, + u16_block_to_f_bytes(prev_data), + ); + tester.write_bytes(mem_as, base_ptr, u16_block_to_f_bytes(read_data)); } else { if a == 0 { - read_data = [F::ZERO; RV64_REGISTER_NUM_LIMBS]; + read_data = [0; BLOCK_FE_WIDTH]; } - tester.write_bytes(mem_as, (ptr_val as usize) - shift_amount, prev_data); - tester.write_bytes(1, a, read_data); + tester.write_bytes(mem_as, base_ptr, u16_block_to_f_bytes(prev_data)); + tester.write_bytes( + RV64_REGISTER_AS as usize, + a, + u16_block_to_f_bytes(read_data), + ); } - let enabled_write = !(is_load & (a == 0)); - + let enabled_write = !(is_load && a == 0); tester.execute( executor, arena, @@ -202,7 +287,7 @@ fn set_and_execute>( a, b, imm as usize, - 1, + RV64_REGISTER_AS as usize, mem_as, enabled_write as usize, imm_sign as usize, @@ -210,57 +295,43 @@ fn set_and_execute>( ), ); - let write_data = run_write_data( - opcode, - read_data.map(|x| x.as_canonical_u32() as u8), - prev_data.map(|x| x.as_canonical_u32() as u8), - shift_amount, - ) - .map(F::from_u8); + let write_data = run_write_data(opcode, read_data, prev_data, shift_amount); if is_load { - if enabled_write { - assert_eq!( - write_data, - tester.read_bytes::(1, a) - ); + let expected = if enabled_write { + u16_block_to_f_bytes(write_data) } else { - assert_eq!( - [F::ZERO; RV64_REGISTER_NUM_LIMBS], - tester.read_bytes::(1, a) - ); - } + [F::ZERO; 8] + }; + assert_eq!( + expected, + tester.read_bytes::<8>(RV64_REGISTER_AS as usize, a) + ); } else { assert_eq!( - write_data, - tester.read_bytes::(mem_as, (ptr_val as usize) - shift_amount) + u16_block_to_f_bytes(write_data), + tester.read_bytes::<8>(mem_as, base_ptr) ); } } -/////////////////////////////////////////////////////////////////////////////////////// -/// POSITIVE TESTS -/// -/// Randomly generate computations and execute, ensuring that the generated trace -/// passes all constraints. -/////////////////////////////////////////////////////////////////////////////////////// -#[test_case(LOADBU, 100)] -#[test_case(LOADHU, 100)] -#[test_case(LOADD, 100)] -#[test_case(LOADWU, 100)] -#[test_case(STOREB, 100)] -#[test_case(STOREH, 100)] -#[test_case(STORED, 100)] -#[test_case(STOREW, 100)] -fn rand_loadstore_test(opcode: Rv64LoadStoreOpcode, num_ops: usize) { - let mut rng = create_seeded_rng(); +fn memory_config_for(opcodes: &[Rv64LoadStoreOpcode]) -> MemoryConfig { let mut mem_config = MemoryConfig::default(); mem_config.addr_spaces[RV64_REGISTER_AS as usize].num_cells = 1 << 29; - if [STORED, STOREW, STOREB, STOREH].contains(&opcode) { + if opcodes + .iter() + .any(|opcode| matches!(opcode, STORED | STOREW | STOREH | STOREB)) + { mem_config.addr_spaces[PUBLIC_VALUES_AS as usize].num_cells = 1 << 29; } - let mut tester = VmChipTestBuilder::from_config(mem_config); - let (mut harness, bitwise) = create_harness(&mut tester); + mem_config +} +#[test_case(LOADBU, 100)] +#[test_case(STOREB, 100)] +fn rand_loadstore_byte_test(opcode: Rv64LoadStoreOpcode, num_ops: usize) { + let mut rng = create_seeded_rng(); + let mut tester = VmChipTestBuilder::from_config(memory_config_for(&[opcode])); + let (mut harness, bitwise) = create_byte_harness(&mut tester); for _ in 0..num_ops { set_and_execute( &mut tester, @@ -268,30 +339,23 @@ fn rand_loadstore_test(opcode: Rv64LoadStoreOpcode, num_ops: usize) { &mut harness.arena, &mut rng, opcode, - None, - None, - None, - None, ); } - - let tester = tester + tester .build() .load(harness) .load_periphery(bitwise) - .finalize(); - tester.simple_test().expect("Verification failed"); + .finalize() + .simple_test() + .unwrap(); } #[test] fn positive_loadwu_shift4_test() { let mut rng = create_seeded_rng(); - let mut mem_config = MemoryConfig::default(); - mem_config.addr_spaces[RV64_REGISTER_AS as usize].num_cells = 1 << 29; - let mut tester = VmChipTestBuilder::from_config(mem_config); - let (mut harness, bitwise) = create_harness(&mut tester); - - set_and_execute( + let mut tester = VmChipTestBuilder::from_config(memory_config_for(&[LOADWU])); + let mut harness = create_word_harness(&mut tester); + set_and_execute_with( &mut tester, &mut harness.executor, &mut harness.arena, @@ -302,24 +366,20 @@ fn positive_loadwu_shift4_test() { Some(0), Some(2), ); - - let tester = tester + tester .build() .load(harness) - .load_periphery(bitwise) - .finalize(); - tester.simple_test().expect("Verification failed"); + .finalize() + .simple_test() + .unwrap(); } #[test] fn positive_loadhu_shift6_test() { let mut rng = create_seeded_rng(); - let mut mem_config = MemoryConfig::default(); - mem_config.addr_spaces[RV64_REGISTER_AS as usize].num_cells = 1 << 29; - let mut tester = VmChipTestBuilder::from_config(mem_config); - let (mut harness, bitwise) = create_harness(&mut tester); - - set_and_execute( + let mut tester = VmChipTestBuilder::from_config(memory_config_for(&[LOADHU])); + let mut harness = create_halfword_harness(&mut tester); + set_and_execute_with( &mut tester, &mut harness.executor, &mut harness.arena, @@ -330,25 +390,20 @@ fn positive_loadhu_shift6_test() { Some(0), Some(2), ); - - let tester = tester + tester .build() .load(harness) - .load_periphery(bitwise) - .finalize(); - tester.simple_test().expect("Verification failed"); + .finalize() + .simple_test() + .unwrap(); } #[test] fn positive_storew_public_values_test() { let mut rng = create_seeded_rng(); - let mut mem_config = MemoryConfig::default(); - 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; - let mut tester = VmChipTestBuilder::from_config(mem_config); - let (mut harness, bitwise) = create_harness(&mut tester); - - set_and_execute( + let mut tester = VmChipTestBuilder::from_config(memory_config_for(&[STOREW])); + let mut harness = create_word_harness(&mut tester); + set_and_execute_with( &mut tester, &mut harness.executor, &mut harness.arena, @@ -357,547 +412,434 @@ fn positive_storew_public_values_test() { Some([4, 0, 0, 0, 0, 0, 0, 0]), Some(0), Some(0), - Some(3), + Some(PUBLIC_VALUES_AS as usize), ); - - let tester = tester + tester .build() .load(harness) - .load_periphery(bitwise) - .finalize(); - tester.simple_test().expect("Verification failed"); -} - -////////////////////////////////////////////////////////////////////////////////////// -// NEGATIVE TESTS -// -// Given a fake trace of a single operation, setup a chip and run the test. We replace -// part of the trace and check that the chip throws the expected error. -////////////////////////////////////////////////////////////////////////////////////// - -#[derive(Clone, Copy, Default, PartialEq)] -struct LoadStorePrankValues { - rs1_data: Option<[u32; RV64_WORD_NUM_LIMBS]>, - read_data: Option<[u32; RV64_REGISTER_NUM_LIMBS]>, - prev_data: Option<[u32; RV64_REGISTER_NUM_LIMBS]>, - write_data: Option<[u32; RV64_REGISTER_NUM_LIMBS]>, - flags: Option<[u32; LOADSTORE_SELECTOR_WIDTH]>, - is_load: Option, - mem_as: Option, + .finalize() + .simple_test() + .unwrap(); } -#[allow(clippy::too_many_arguments)] -fn run_negative_loadstore_test( - opcode: Rv64LoadStoreOpcode, - rs1: Option<[u8; RV64_REGISTER_NUM_LIMBS]>, - imm: Option, - imm_sign: Option, - prank_vals: LoadStorePrankValues, - _interaction_error: bool, -) { +#[test_case(LOADHU, 100)] +#[test_case(STOREH, 100)] +fn rand_loadstore_halfword_test(opcode: Rv64LoadStoreOpcode, num_ops: usize) { let mut rng = create_seeded_rng(); - let mut mem_config = MemoryConfig::default(); - 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; - let mut tester = VmChipTestBuilder::from_config(mem_config); - let (mut harness, bitwise) = create_harness(&mut tester); - - set_and_execute( - &mut tester, - &mut harness.executor, - &mut harness.arena, - &mut rng, - opcode, - rs1, - imm, - imm_sign, - None, - ); - - let adapter_width = BaseAir::::width(&harness.air.adapter); - - let modify_trace = |trace: &mut DenseMatrix| { - let mut trace_row = trace.row_slice(0).unwrap().to_vec(); - let (adapter_row, core_row) = trace_row.split_at_mut(adapter_width); - let adapter_cols: &mut Rv64LoadStoreAdapterCols = adapter_row.borrow_mut(); - let core_cols: &mut LoadStoreCoreCols = core_row.borrow_mut(); - - if let Some(rs1_data) = prank_vals.rs1_data { - adapter_cols.rs1_data = - array::from_fn(|i| pack_u8_pair_u32(rs1_data[2 * i], rs1_data[2 * i + 1])); - } - if let Some(read_data) = prank_vals.read_data { - core_cols.read_data = read_data.map(F::from_u32); - } - if let Some(prev_data) = prank_vals.prev_data { - core_cols.prev_data = prev_data.map(F::from_u32); - } - if let Some(write_data) = prank_vals.write_data { - core_cols.write_data = write_data.map(F::from_u32); - } - if let Some(flags) = prank_vals.flags { - core_cols.selector = flags.map(F::from_u32); - } - if let Some(is_load) = prank_vals.is_load { - core_cols.is_load = F::from_bool(is_load); - } - if let Some(mem_as) = prank_vals.mem_as { - adapter_cols.mem_as = F::from_u32(mem_as); - } - - *trace = RowMajorMatrix::new(trace_row, trace.width()); - }; - - disable_debug_builder(); - let tester = tester - .build() - .load_and_prank_trace(harness, modify_trace) - .load_periphery(bitwise) - .finalize(); + let mut tester = VmChipTestBuilder::from_config(memory_config_for(&[opcode])); + let mut harness = create_halfword_harness(&mut tester); + for _ in 0..num_ops { + set_and_execute( + &mut tester, + &mut harness.executor, + &mut harness.arena, + &mut rng, + opcode, + ); + } tester + .build() + .load(harness) + .finalize() .simple_test() - .expect_err("Expected verification to fail, but it passed"); + .unwrap(); } -#[test] -fn negative_wrong_opcode_tests() { - run_negative_loadstore_test( - LOADD, - None, - None, - None, - LoadStorePrankValues { - is_load: Some(false), - ..Default::default() - }, - false, - ); - - run_negative_loadstore_test( - LOADBU, - Some([4, 0, 0, 0, 0, 0, 0, 0]), - Some(1), - None, - LoadStorePrankValues { - flags: Some(selector_point_for_opcode_shift(LOADBU, 0)), - ..Default::default() - }, - false, - ); - - run_negative_loadstore_test( - STOREH, - Some([11, 169, 76, 28, 0, 0, 0, 0]), - Some(37121), - None, - LoadStorePrankValues { - flags: Some(selector_point_for_opcode_shift(STOREH, 0)), - is_load: Some(true), - ..Default::default() - }, - false, - ); - - run_negative_loadstore_test( - LOADWU, - Some([4, 0, 0, 0, 0, 0, 0, 0]), - Some(0), - Some(0), - LoadStorePrankValues { - flags: Some(selector_point_for_opcode_shift(LOADWU, 0)), - ..Default::default() - }, - false, - ); +#[test_case(LOADWU, 100)] +#[test_case(STOREW, 100)] +fn rand_loadstore_word_test(opcode: Rv64LoadStoreOpcode, num_ops: usize) { + let mut rng = create_seeded_rng(); + let mut tester = VmChipTestBuilder::from_config(memory_config_for(&[opcode])); + let mut harness = create_word_harness(&mut tester); + for _ in 0..num_ops { + set_and_execute( + &mut tester, + &mut harness.executor, + &mut harness.arena, + &mut rng, + opcode, + ); + } + tester + .build() + .load(harness) + .finalize() + .simple_test() + .unwrap(); } -#[test] -fn negative_write_data_tests() { - run_negative_loadstore_test( - LOADHU, - Some([13, 11, 156, 23, 0, 0, 0, 0]), - Some(43641), - None, - LoadStorePrankValues { - rs1_data: None, - read_data: Some([175, 33, 198, 250, 131, 74, 186, 29]), - prev_data: Some([90, 121, 64, 205, 159, 213, 89, 34]), - write_data: Some([175, 33, 0, 0, 0, 0, 0, 0]), - flags: Some(selector_point_for_opcode_shift(LOADHU, 0)), - is_load: Some(true), - mem_as: None, - }, - true, - ); - - run_negative_loadstore_test( - STOREB, - Some([45, 123, 87, 24, 0, 0, 0, 0]), - Some(28122), - Some(0), - LoadStorePrankValues { - rs1_data: None, - read_data: Some([175, 33, 198, 250, 131, 74, 186, 29]), - prev_data: Some([90, 121, 64, 205, 159, 213, 89, 34]), - write_data: Some([175, 121, 64, 205, 159, 213, 89, 34]), - flags: Some(selector_point_for_opcode_shift(STOREB, 3)), - is_load: None, - mem_as: None, - }, - false, - ); - - run_negative_loadstore_test( - LOADWU, - Some([4, 0, 0, 0, 0, 0, 0, 0]), - Some(0), - Some(0), - LoadStorePrankValues { - rs1_data: None, - read_data: Some([138, 45, 202, 76, 131, 74, 186, 29]), - prev_data: Some([159, 213, 89, 34, 142, 67, 210, 88]), - write_data: Some([138, 45, 202, 76, 0, 0, 0, 0]), - flags: Some(selector_point_for_opcode_shift(LOADWU, 4)), - is_load: Some(true), - mem_as: None, - }, - false, - ); +#[test_case(LOADD, 100)] +#[test_case(STORED, 100)] +fn rand_loadstore_doubleword_test(opcode: Rv64LoadStoreOpcode, num_ops: usize) { + let mut rng = create_seeded_rng(); + let mut tester = VmChipTestBuilder::from_config(memory_config_for(&[opcode])); + let mut harness = create_doubleword_harness(&mut tester); + for _ in 0..num_ops { + set_and_execute( + &mut tester, + &mut harness.executor, + &mut harness.arena, + &mut rng, + opcode, + ); + } + tester + .build() + .load(harness) + .finalize() + .simple_test() + .unwrap(); } -#[test] -#[should_panic(expected = "effective address exceeds implemented memory address space")] -fn negative_32_bit_address_wraparound_test() { - run_negative_loadstore_test( - LOADBU, - Some([0xf8, 0xff, 0xff, 0xff, 0, 0, 0, 0]), - Some(16), - Some(0), - LoadStorePrankValues::default(), - false, - ); +fn b(bytes: [u8; 8]) -> [u16; BLOCK_FE_WIDTH] { + rv64_bytes_to_u16_block(bytes) } -#[test] -fn negative_wrong_address_space_tests() { - run_negative_loadstore_test( - LOADD, - None, - None, - None, - LoadStorePrankValues { - mem_as: Some(3), - ..Default::default() - }, - false, - ); - - run_negative_loadstore_test( - LOADWU, - None, - None, - None, - LoadStorePrankValues { - mem_as: Some(4), - ..Default::default() - }, - false, - ); - - run_negative_loadstore_test( - STOREW, - None, - None, - None, - LoadStorePrankValues { - mem_as: Some(1), - ..Default::default() - }, - false, - ); - - run_negative_loadstore_test( - STORED, - None, - None, - None, - LoadStorePrankValues { - mem_as: Some(1), - ..Default::default() - }, - false, - ); - - run_negative_loadstore_test( - STOREW, - None, - None, - None, - LoadStorePrankValues { - mem_as: Some(4), - ..Default::default() - }, - false, - ); -} - -/////////////////////////////////////////////////////////////////////////////////////// -/// SANITY TESTS -/// -/// Ensure that solve functions produce the correct results. -/////////////////////////////////////////////////////////////////////////////////////// #[test] fn run_loadd_stored_sanity_test() { - let read_data = [138, 45, 202, 76, 131, 74, 186, 29]; - let prev_data = [159, 213, 89, 34, 142, 67, 210, 88]; + 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!(run_write_data(LOADD, read_data, prev_data, 0), read_data); assert_eq!(run_write_data(STORED, read_data, prev_data, 0), read_data); } #[test] fn run_loadwu_storew_sanity_test() { - let read_data = [138, 45, 202, 76, 131, 74, 186, 29]; - let prev_data = [159, 213, 89, 34, 142, 67, 210, 88]; - + 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!( run_write_data(LOADWU, read_data, prev_data, 0), - [138, 45, 202, 76, 0, 0, 0, 0] + b([138, 45, 202, 76, 0, 0, 0, 0]) ); assert_eq!( run_write_data(LOADWU, read_data, prev_data, 4), - [131, 74, 186, 29, 0, 0, 0, 0] + b([131, 74, 186, 29, 0, 0, 0, 0]) ); assert_eq!( run_write_data(STOREW, read_data, prev_data, 0), - [138, 45, 202, 76, 142, 67, 210, 88] + b([138, 45, 202, 76, 142, 67, 210, 88]) ); assert_eq!( run_write_data(STOREW, read_data, prev_data, 4), - [159, 213, 89, 34, 138, 45, 202, 76] + b([159, 213, 89, 34, 138, 45, 202, 76]) ); } #[test] fn run_storeh_sanity_test() { - let read_data = [250, 123, 67, 198, 175, 33, 198, 250]; - let prev_data = [144, 56, 175, 92, 90, 121, 64, 205]; - + let read_data = b([250, 123, 67, 198, 175, 33, 198, 250]); + let prev_data = b([144, 56, 175, 92, 90, 121, 64, 205]); assert_eq!( run_write_data(STOREH, read_data, prev_data, 0), - [250, 123, 175, 92, 90, 121, 64, 205] + b([250, 123, 175, 92, 90, 121, 64, 205]) ); assert_eq!( run_write_data(STOREH, read_data, prev_data, 2), - [144, 56, 250, 123, 90, 121, 64, 205] + b([144, 56, 250, 123, 90, 121, 64, 205]) ); assert_eq!( run_write_data(STOREH, read_data, prev_data, 4), - [144, 56, 175, 92, 250, 123, 64, 205] + b([144, 56, 175, 92, 250, 123, 64, 205]) ); assert_eq!( run_write_data(STOREH, read_data, prev_data, 6), - [144, 56, 175, 92, 90, 121, 250, 123] + b([144, 56, 175, 92, 90, 121, 250, 123]) ); } #[test] fn run_storeb_sanity_test() { - let read_data = [221, 104, 58, 147, 175, 33, 198, 250]; - let prev_data = [199, 83, 243, 12, 90, 121, 64, 205]; - + let read_data = b([221, 104, 58, 147, 175, 33, 198, 250]); + let prev_data = b([199, 83, 243, 12, 90, 121, 64, 205]); assert_eq!( run_write_data(STOREB, read_data, prev_data, 0), - [221, 83, 243, 12, 90, 121, 64, 205] + b([221, 83, 243, 12, 90, 121, 64, 205]) ); assert_eq!( run_write_data(STOREB, read_data, prev_data, 1), - [199, 221, 243, 12, 90, 121, 64, 205] + b([199, 221, 243, 12, 90, 121, 64, 205]) ); assert_eq!( run_write_data(STOREB, read_data, prev_data, 2), - [199, 83, 221, 12, 90, 121, 64, 205] + b([199, 83, 221, 12, 90, 121, 64, 205]) ); assert_eq!( run_write_data(STOREB, read_data, prev_data, 3), - [199, 83, 243, 221, 90, 121, 64, 205] + b([199, 83, 243, 221, 90, 121, 64, 205]) ); assert_eq!( run_write_data(STOREB, read_data, prev_data, 4), - [199, 83, 243, 12, 221, 121, 64, 205] + b([199, 83, 243, 12, 221, 121, 64, 205]) ); assert_eq!( run_write_data(STOREB, read_data, prev_data, 5), - [199, 83, 243, 12, 90, 221, 64, 205] + b([199, 83, 243, 12, 90, 221, 64, 205]) ); assert_eq!( run_write_data(STOREB, read_data, prev_data, 6), - [199, 83, 243, 12, 90, 121, 221, 205] + b([199, 83, 243, 12, 90, 121, 221, 205]) ); assert_eq!( run_write_data(STOREB, read_data, prev_data, 7), - [199, 83, 243, 12, 90, 121, 64, 221] + b([199, 83, 243, 12, 90, 121, 64, 221]) ); } #[test] fn run_loadhu_sanity_test() { - let read_data = [175, 33, 198, 250, 131, 74, 186, 29]; - let prev_data = [90, 121, 64, 205, 142, 67, 210, 88]; - + let read_data = b([175, 33, 198, 250, 131, 74, 186, 29]); + let prev_data = b([90, 121, 64, 205, 142, 67, 210, 88]); assert_eq!( run_write_data(LOADHU, read_data, prev_data, 0), - [175, 33, 0, 0, 0, 0, 0, 0] + b([175, 33, 0, 0, 0, 0, 0, 0]) ); assert_eq!( run_write_data(LOADHU, read_data, prev_data, 2), - [198, 250, 0, 0, 0, 0, 0, 0] + b([198, 250, 0, 0, 0, 0, 0, 0]) ); assert_eq!( run_write_data(LOADHU, read_data, prev_data, 4), - [131, 74, 0, 0, 0, 0, 0, 0] + b([131, 74, 0, 0, 0, 0, 0, 0]) ); assert_eq!( run_write_data(LOADHU, read_data, prev_data, 6), - [186, 29, 0, 0, 0, 0, 0, 0] + b([186, 29, 0, 0, 0, 0, 0, 0]) ); } #[test] fn run_loadbu_sanity_test() { - let read_data = [131, 74, 186, 29, 138, 45, 202, 76]; - let prev_data = [142, 67, 210, 88, 159, 213, 89, 34]; + let read_data = b([131, 74, 186, 29, 138, 45, 202, 76]); + let prev_data = b([142, 67, 210, 88, 159, 213, 89, 34]); + for (shift, expected) in [ + (0, [131, 0, 0, 0, 0, 0, 0, 0]), + (1, [74, 0, 0, 0, 0, 0, 0, 0]), + (2, [186, 0, 0, 0, 0, 0, 0, 0]), + (3, [29, 0, 0, 0, 0, 0, 0, 0]), + (4, [138, 0, 0, 0, 0, 0, 0, 0]), + (5, [45, 0, 0, 0, 0, 0, 0, 0]), + (6, [202, 0, 0, 0, 0, 0, 0, 0]), + (7, [76, 0, 0, 0, 0, 0, 0, 0]), + ] { + assert_eq!( + run_write_data(LOADBU, read_data, prev_data, shift), + b(expected) + ); + } +} +#[test] +fn load_sign_extend_sanity_tests() { + let read_data = b([34, 159, 237, 151, 100, 200, 50, 25]); assert_eq!( - run_write_data(LOADBU, read_data, prev_data, 0), - [131, 0, 0, 0, 0, 0, 0, 0] + run_write_data(LOADH, read_data, [0; BLOCK_FE_WIDTH], 0), + b([34, 159, 255, 255, 255, 255, 255, 255]) ); assert_eq!( - run_write_data(LOADBU, read_data, prev_data, 1), - [74, 0, 0, 0, 0, 0, 0, 0] + run_write_data(LOADH, read_data, [0; BLOCK_FE_WIDTH], 2), + b([237, 151, 255, 255, 255, 255, 255, 255]) ); assert_eq!( - run_write_data(LOADBU, read_data, prev_data, 2), - [186, 0, 0, 0, 0, 0, 0, 0] + run_write_data(LOADH, read_data, [0; BLOCK_FE_WIDTH], 4), + b([100, 200, 255, 255, 255, 255, 255, 255]) ); assert_eq!( - run_write_data(LOADBU, read_data, prev_data, 3), - [29, 0, 0, 0, 0, 0, 0, 0] + run_write_data(LOADH, read_data, [0; BLOCK_FE_WIDTH], 6), + b([50, 25, 0, 0, 0, 0, 0, 0]) ); + + let read_data = b([45, 82, 99, 127, 200, 150, 180, 210]); + for shift in 0..8 { + let byte = rv64_u16_block_to_bytes(read_data)[shift]; + assert_eq!( + rv64_u16_block_to_bytes(run_write_data(LOADB, read_data, [0; BLOCK_FE_WIDTH], shift)), + (byte as i8 as i64).to_le_bytes(), + "LOADB shift={shift}" + ); + } + + let read_data = b([0x01, 0x02, 0x03, 0x84, 0xAA, 0xBB, 0xCC, 0xDD]); assert_eq!( - run_write_data(LOADBU, read_data, prev_data, 4), - [138, 0, 0, 0, 0, 0, 0, 0] + run_write_data(LOADW, read_data, [0; BLOCK_FE_WIDTH], 0), + b([0x01, 0x02, 0x03, 0x84, 0xFF, 0xFF, 0xFF, 0xFF]) ); assert_eq!( - run_write_data(LOADBU, read_data, prev_data, 5), - [45, 0, 0, 0, 0, 0, 0, 0] + run_write_data(LOADW, read_data, [0; BLOCK_FE_WIDTH], 4), + b([0xAA, 0xBB, 0xCC, 0xDD, 0xFF, 0xFF, 0xFF, 0xFF]) ); + + let read_data = b([0x01, 0x02, 0x03, 0x04, 0xAA, 0xBB, 0xCC, 0x7D]); assert_eq!( - run_write_data(LOADBU, read_data, prev_data, 6), - [202, 0, 0, 0, 0, 0, 0, 0] + run_write_data(LOADW, read_data, [0; BLOCK_FE_WIDTH], 0), + b([0x01, 0x02, 0x03, 0x04, 0, 0, 0, 0]) ); assert_eq!( - run_write_data(LOADBU, read_data, prev_data, 7), - [76, 0, 0, 0, 0, 0, 0, 0] + run_write_data(LOADW, read_data, [0; BLOCK_FE_WIDTH], 4), + b([0xAA, 0xBB, 0xCC, 0x7D, 0, 0, 0, 0]) ); } -// //////////////////////////////////////////////////////////////////////////////////// -// CUDA TESTS -// -// Ensure GPU tracegen is equivalent to CPU tracegen -// //////////////////////////////////////////////////////////////////////////////////// +#[test] +#[should_panic] +fn solve_loadw_rejects_shift_2() { + run_write_data(LOADW, b([1, 2, 3, 4, 5, 6, 7, 8]), [0; BLOCK_FE_WIDTH], 2); +} -#[cfg(feature = "cuda")] -type GpuHarness = GpuTestChipHarness< - F, - Rv64LoadStoreExecutor, - Rv64LoadStoreAir, - Rv64LoadStoreChipGpu, - Rv64LoadStoreChip, ->; +#[test] +#[should_panic] +fn solve_loadw_rejects_shift_6() { + run_write_data(LOADW, b([1, 2, 3, 4, 5, 6, 7, 8]), [0; BLOCK_FE_WIDTH], 6); +} -#[cfg(feature = "cuda")] -fn create_cuda_harness(tester: &GpuChipTestBuilder) -> GpuHarness { - let range_bus = default_var_range_checker_bus(); - let dummy_range_checker_chip = dummy_range_checker(range_bus); - let bitwise_bus = default_bitwise_lookup_bus(); - let dummy_bitwise_chip = Arc::new(BitwiseOperationLookupChip::::new( - bitwise_bus, - )); +#[test] +fn accepted_shift_sets() { + let read_data = b([0x10, 0x20, 0x30, 0x40, 0x50, 0x60, 0x70, 0x80]); + for shift in 0..8 { + let _ = run_write_data(LOADB, read_data, [0; BLOCK_FE_WIDTH], shift); + } + for shift in [0, 2, 4, 6] { + let _ = run_write_data(LOADH, read_data, [0; BLOCK_FE_WIDTH], shift); + let _ = run_write_data(LOADHU, read_data, [0; BLOCK_FE_WIDTH], shift); + } + for shift in [0, 4] { + let _ = run_write_data(LOADW, read_data, [0; BLOCK_FE_WIDTH], shift); + let _ = run_write_data(LOADWU, read_data, [0; BLOCK_FE_WIDTH], shift); + } +} - let (air, executor, cpu_chip) = create_harness_fields( - tester.memory_bridge(), - tester.execution_bridge(), - dummy_range_checker_chip, - dummy_bitwise_chip, - tester.dummy_memory_helper(), - tester.address_bits(), - ); - let gpu_chip = Rv64LoadStoreChipGpu::new( - tester.range_checker(), - tester.bitwise_op_lookup(), - tester.address_bits(), - tester.timestamp_max_bits(), +fn assert_pranked_byte_fails( + opcode: Rv64LoadStoreOpcode, + prank: impl Fn(&mut LoadStoreByteCoreCols), +) { + let mut rng = create_seeded_rng(); + let mut tester = VmChipTestBuilder::from_config(memory_config_for(&[opcode])); + let (mut harness, bitwise) = create_byte_harness(&mut tester); + set_and_execute( + &mut tester, + &mut harness.executor, + &mut harness.arena, + &mut rng, + opcode, ); - - GpuTestChipHarness::with_capacity(executor, air, gpu_chip, cpu_chip, MAX_INS_CAPACITY) + let adapter_width = BaseAir::::width(&harness.air.adapter); + let modify_trace = |trace: &mut DenseMatrix| { + let mut trace_row = trace.row_slice(0).unwrap().to_vec(); + let (_, core_row) = trace_row.split_at_mut(adapter_width); + prank(core_row.borrow_mut()); + *trace = RowMajorMatrix::new(trace_row, trace.width()); + }; + disable_debug_builder(); + tester + .build() + .load_and_prank_trace(harness, modify_trace) + .load_periphery(bitwise) + .finalize() + .simple_test() + .expect_err("pranked byte loadstore trace should fail"); } -#[cfg(feature = "cuda")] -#[test_case(LOADD, 100)] -#[test_case(LOADBU, 100)] -#[test_case(LOADHU, 100)] -#[test_case(LOADWU, 100)] -#[test_case(STORED, 100)] -#[test_case(STOREW, 100)] -#[test_case(STOREB, 100)] -#[test_case(STOREH, 100)] -fn test_cuda_rand_load_store_tracegen(opcode: Rv64LoadStoreOpcode, num_ops: usize) { +fn assert_pranked_halfword_fails( + opcode: Rv64LoadStoreOpcode, + prank: impl Fn(&mut LoadStoreAlignedCoreCols), +) { let mut rng = create_seeded_rng(); - let mut mem_config = MemoryConfig { - pointer_max_bits: 20, - ..Default::default() + let mut tester = VmChipTestBuilder::from_config(memory_config_for(&[opcode])); + let mut harness = create_halfword_harness(&mut tester); + set_and_execute( + &mut tester, + &mut harness.executor, + &mut harness.arena, + &mut rng, + opcode, + ); + let adapter_width = BaseAir::::width(&harness.air.adapter); + let modify_trace = |trace: &mut DenseMatrix| { + let mut trace_row = trace.row_slice(0).unwrap().to_vec(); + let (_, core_row) = trace_row.split_at_mut(adapter_width); + prank(core_row.borrow_mut()); + *trace = RowMajorMatrix::new(trace_row, trace.width()); }; - mem_config.addr_spaces[RV64_REGISTER_AS as usize].num_cells = 1 << 20; - mem_config.addr_spaces[RV64_MEMORY_AS as usize].num_cells = 1 << 20; - if [STORED, STOREW, STOREB, STOREH].contains(&opcode) { - mem_config.addr_spaces[PUBLIC_VALUES_AS as usize].num_cells = 1 << 20; - } - let mut tester = GpuChipTestBuilder::new(mem_config, default_var_range_checker_bus()) - .with_bitwise_op_lookup(default_bitwise_lookup_bus()); - - let mut harness = create_cuda_harness(&tester); - for _ in 0..num_ops { - set_and_execute( - &mut tester, - &mut harness.executor, - &mut harness.dense_arena, - &mut rng, - opcode, - None, - None, - None, - None, - ); - } + disable_debug_builder(); + tester + .build() + .load_and_prank_trace(harness, modify_trace) + .finalize() + .simple_test() + .expect_err("pranked halfword loadstore trace should fail"); +} - type Record<'a> = ( - &'a mut Rv64LoadStoreAdapterRecord, - &'a mut LoadStoreCoreRecord, +fn assert_pranked_word_fails( + opcode: Rv64LoadStoreOpcode, + prank: impl Fn(&mut LoadStoreAlignedCoreCols), +) { + let mut rng = create_seeded_rng(); + let mut tester = VmChipTestBuilder::from_config(memory_config_for(&[opcode])); + let mut harness = create_word_harness(&mut tester); + set_and_execute( + &mut tester, + &mut harness.executor, + &mut harness.arena, + &mut rng, + opcode, ); + let adapter_width = BaseAir::::width(&harness.air.adapter); + let modify_trace = |trace: &mut DenseMatrix| { + let mut trace_row = trace.row_slice(0).unwrap().to_vec(); + let (_, core_row) = trace_row.split_at_mut(adapter_width); + prank(core_row.borrow_mut()); + *trace = RowMajorMatrix::new(trace_row, trace.width()); + }; + disable_debug_builder(); + tester + .build() + .load_and_prank_trace(harness, modify_trace) + .finalize() + .simple_test() + .expect_err("pranked word loadstore trace should fail"); +} - harness - .dense_arena - .get_record_seeker::() - .transfer_to_matrix_arena( - &mut harness.matrix_arena, - EmptyAdapterCoreLayout::::new(), - ); - +fn assert_pranked_doubleword_fails( + opcode: Rv64LoadStoreOpcode, + prank: impl Fn(&mut LoadStoreDoublewordCoreCols), +) { + let mut rng = create_seeded_rng(); + let mut tester = VmChipTestBuilder::from_config(memory_config_for(&[opcode])); + let mut harness = create_doubleword_harness(&mut tester); + set_and_execute( + &mut tester, + &mut harness.executor, + &mut harness.arena, + &mut rng, + opcode, + ); + let adapter_width = BaseAir::::width(&harness.air.adapter); + let modify_trace = |trace: &mut DenseMatrix| { + let mut trace_row = trace.row_slice(0).unwrap().to_vec(); + let (_, core_row) = trace_row.split_at_mut(adapter_width); + prank(core_row.borrow_mut()); + *trace = RowMajorMatrix::new(trace_row, trace.width()); + }; + disable_debug_builder(); tester .build() - .load_gpu_harness(harness) + .load_and_prank_trace(harness, modify_trace) .finalize() .simple_test() - .unwrap(); + .expect_err("pranked doubleword loadstore trace should fail"); +} + +#[test] +fn negative_split_write_data_tests() { + assert_pranked_byte_fails(STOREB, |core| core.read_data[0] += F::ONE); + assert_pranked_halfword_fails(LOADHU, |core| core.read_data[0] += F::ONE); + assert_pranked_word_fails(LOADWU, |core| core.read_data[0] += F::ONE); + assert_pranked_doubleword_fails(LOADD, |core| core.read_data[0] += F::ONE); +} + +#[test] +fn negative_split_opcode_role_tests() { + assert_pranked_byte_fails(LOADBU, |core| core.is_load = F::ZERO); + assert_pranked_halfword_fails(STOREH, |core| core.is_load = F::ONE); + assert_pranked_word_fails(LOADWU, |core| core.is_load = F::ZERO); + assert_pranked_doubleword_fails(LOADD, |core| core.is_load = F::ZERO); } diff --git a/extensions/riscv/circuit/src/loadstore/word/mod.rs b/extensions/riscv/circuit/src/loadstore/word/mod.rs new file mode 100644 index 0000000000..c93d34f82d --- /dev/null +++ b/extensions/riscv/circuit/src/loadstore/word/mod.rs @@ -0,0 +1,24 @@ +use openvm_circuit::arch::{VmAirWrapper, VmChipWrapper}; + +use crate::{ + adapters::{Rv64LoadStoreAdapterAir, Rv64LoadStoreAdapterExecutor}, + loadstore::{ + aligned::{LoadStoreAlignedCoreAir, LoadStoreAlignedFiller}, + common::{LoadStoreExecutor, KIND_WORD}, + }, +}; + +pub const WORD_CASES: usize = 4; +pub const WORD_SELECTOR_WIDTH: usize = 2; + +pub type LoadStoreWordCoreAir = LoadStoreAlignedCoreAir; +pub type LoadStoreWordFiller = LoadStoreAlignedFiller< + crate::adapters::Rv64LoadStoreAdapterFiller, + KIND_WORD, + WORD_CASES, + WORD_SELECTOR_WIDTH, +>; + +pub type Rv64LoadStoreWordAir = VmAirWrapper; +pub type Rv64LoadStoreWordExecutor = LoadStoreExecutor; +pub type Rv64LoadStoreWordChip = VmChipWrapper; From 3431f19c1f4824802b26f9b559ed16e2720ecdc1 Mon Sep 17 00:00:00 2001 From: Ayush Shukla Date: Wed, 24 Jun 2026 09:46:29 +0200 Subject: [PATCH 02/26] feat(riscv): add cuda tracegen for split loadstore --- .../circuit/cuda/src/load_sign_extend.cu | 302 ++++++++++++++ .../riscv/circuit/cuda/src/loadstore.cu | 371 ++++++++++++++++++ extensions/riscv/circuit/src/cuda_abi.rs | 286 ++++++++++++++ .../circuit/src/load_sign_extend/cuda.rs | 120 +++++- .../circuit/src/load_sign_extend/tests.rs | 260 +++++++++++- .../riscv/circuit/src/loadstore/cuda.rs | 151 ++++++- .../riscv/circuit/src/loadstore/tests.rs | 339 +++++++++++++++- 7 files changed, 1797 insertions(+), 32 deletions(-) create mode 100644 extensions/riscv/circuit/cuda/src/load_sign_extend.cu create mode 100644 extensions/riscv/circuit/cuda/src/loadstore.cu diff --git a/extensions/riscv/circuit/cuda/src/load_sign_extend.cu b/extensions/riscv/circuit/cuda/src/load_sign_extend.cu new file mode 100644 index 0000000000..7db0355d20 --- /dev/null +++ b/extensions/riscv/circuit/cuda/src/load_sign_extend.cu @@ -0,0 +1,302 @@ +#include "launcher.cuh" +#include "primitives/buffer_view.cuh" +#include "primitives/constants.h" +#include "primitives/encoder.cuh" +#include "primitives/histogram.cuh" +#include "primitives/trace_access.h" +#include "riscv/adapters/loadstore.cuh" +#include + +using namespace riscv; +using namespace program; + +enum Rv64LoadStoreOpcode { + LOADD, + LOADBU, + LOADHU, + LOADWU, + STORED, + STOREW, + STOREH, + STOREB, + LOADB, + LOADH, + LOADW, +}; + +constexpr size_t LOAD_SIGN_EXTEND_BYTE_SELECTOR_WIDTH = 3; +constexpr uint32_t LOAD_SIGN_EXTEND_BYTE_CASES = 8; +constexpr size_t LOAD_SIGN_EXTEND_HALFWORD_SELECTOR_WIDTH = 2; +constexpr uint32_t LOAD_SIGN_EXTEND_HALFWORD_CASES = 4; +constexpr size_t LOAD_SIGN_EXTEND_WORD_SELECTOR_WIDTH = 1; +constexpr uint32_t LOAD_SIGN_EXTEND_WORD_CASES = 2; +constexpr uint32_t LOAD_SIGN_EXTEND_SELECTOR_MAX_DEGREE = 2; +constexpr uint16_t SIGN_BYTE = 1 << (RV64_BYTE_BITS - 1); +constexpr uint16_t SIGN_U16 = 1 << (U16_BITS - 1); +struct LoadStoreRecord { + uint8_t local_opcode; + uint8_t shift_amount; + uint16_t read_data[BLOCK_FE_WIDTH]; + uint16_t prev_data[BLOCK_FE_WIDTH]; +}; + +template struct LoadSignExtendByteCoreCols { + T selector[LOAD_SIGN_EXTEND_BYTE_SELECTOR_WIDTH]; + T is_valid; + T data_most_sig_bit; + T read_cell_bytes[2]; + T read_data[BLOCK_FE_WIDTH]; + T prev_data[BLOCK_FE_WIDTH]; +}; + +template struct LoadSignExtendAlignedCoreCols { + T selector[SELECTOR_WIDTH]; + T is_valid; + T data_most_sig_bit; + T read_data[BLOCK_FE_WIDTH]; + T prev_data[BLOCK_FE_WIDTH]; +}; + +template struct Rv64LoadSignExtendByteCols { + Rv64LoadStoreAdapterCols adapter; + LoadSignExtendByteCoreCols core; +}; + +template struct Rv64LoadSignExtendHalfwordCols { + Rv64LoadStoreAdapterCols adapter; + LoadSignExtendAlignedCoreCols core; +}; + +template struct Rv64LoadSignExtendWordCols { + Rv64LoadStoreAdapterCols adapter; + LoadSignExtendAlignedCoreCols core; +}; + +constexpr size_t RV64_LOAD_SIGN_EXTEND_BYTE_WIDTH = + sizeof(Rv64LoadSignExtendByteCols); +constexpr size_t RV64_LOAD_SIGN_EXTEND_HALFWORD_WIDTH = + sizeof(Rv64LoadSignExtendHalfwordCols); +constexpr size_t RV64_LOAD_SIGN_EXTEND_WORD_WIDTH = + sizeof(Rv64LoadSignExtendWordCols); + +struct Rv64LoadSignExtendRecord { + Rv64LoadStoreAdapterRecord adapter; + LoadStoreRecord core; +}; + +static_assert(sizeof(Rv64LoadStoreAdapterRecord) == 36); +static_assert(sizeof(LoadStoreRecord) == 18); +static_assert(sizeof(Rv64LoadSignExtendRecord) == 56); +static_assert(offsetof(Rv64LoadSignExtendRecord, core) == 36); + +__device__ __forceinline__ uint16_t byte_from_cell(uint16_t cell, uint8_t byte_idx) { + return (cell >> (RV64_BYTE_BITS * byte_idx)) & 0xff; +} + +struct LoadSignExtendByteCore { + VariableRangeChecker range_checker; + BitwiseOperationLookup bitwise_lookup; + + __device__ LoadSignExtendByteCore( + VariableRangeChecker range_checker, + BitwiseOperationLookup bitwise_lookup + ) + : range_checker(range_checker), bitwise_lookup(bitwise_lookup) {} + + __device__ void fill_trace_row(RowSlice row, LoadStoreRecord record) { + assert(record.local_opcode == LOADB); + uint8_t shift = record.shift_amount; + uint16_t read_cell = record.read_data[shift >> 1]; + uint16_t read_cell_bytes[2] = { + byte_from_cell(read_cell, 0), + byte_from_cell(read_cell, 1), + }; + uint16_t selected_byte = read_cell_bytes[shift & 1]; + uint16_t sign_bit = selected_byte & SIGN_BYTE; + + bitwise_lookup.add_range(read_cell_bytes[0], read_cell_bytes[1]); + range_checker.add_count(selected_byte - sign_bit, RV64_BYTE_BITS - 1); + + Encoder encoder( + LOAD_SIGN_EXTEND_BYTE_CASES, + LOAD_SIGN_EXTEND_SELECTOR_MAX_DEGREE, + true, + LOAD_SIGN_EXTEND_BYTE_SELECTOR_WIDTH + ); + encoder.write_flag_pt( + row.slice_from(COL_INDEX(LoadSignExtendByteCoreCols, selector)), + shift + ); + COL_WRITE_VALUE(row, LoadSignExtendByteCoreCols, is_valid, 1); + COL_WRITE_VALUE(row, LoadSignExtendByteCoreCols, data_most_sig_bit, sign_bit != 0); + COL_WRITE_ARRAY(row, LoadSignExtendByteCoreCols, read_cell_bytes, read_cell_bytes); + COL_WRITE_ARRAY(row, LoadSignExtendByteCoreCols, read_data, record.read_data); + COL_WRITE_ARRAY(row, LoadSignExtendByteCoreCols, prev_data, record.prev_data); + } +}; + +template struct LoadSignExtendAlignedCore { + VariableRangeChecker range_checker; + + __device__ LoadSignExtendAlignedCore(VariableRangeChecker range_checker) + : range_checker(range_checker) {} + + __device__ void fill_trace_row(RowSlice row, LoadStoreRecord record) { + Rv64LoadStoreOpcode opcode = static_cast(record.local_opcode); + uint8_t shift = record.shift_amount; + uint32_t access_cells; + uint32_t case_idx; + if constexpr (CASES == LOAD_SIGN_EXTEND_HALFWORD_CASES) { + assert(opcode == LOADH); + access_cells = 1; + case_idx = shift >> 1; + } else { + assert(opcode == LOADW); + access_cells = 2; + case_idx = shift >> 2; + } + uint16_t sign_cell = record.read_data[(shift >> 1) + access_cells - 1]; + uint16_t sign_bit = sign_cell & SIGN_U16; + range_checker.add_count(sign_cell - sign_bit, U16_BITS - 1); + + Encoder encoder(CASES, LOAD_SIGN_EXTEND_SELECTOR_MAX_DEGREE, true, SELECTOR_WIDTH); + encoder.write_flag_pt(row, case_idx); + row[SELECTOR_WIDTH] = 1; + row[SELECTOR_WIDTH + 1] = sign_bit != 0; + row.write_array(SELECTOR_WIDTH + 2, BLOCK_FE_WIDTH, record.read_data); + row.write_array( + SELECTOR_WIDTH + 2 + BLOCK_FE_WIDTH, + BLOCK_FE_WIDTH, + record.prev_data + ); + } +}; + +template