1use std::borrow::{Borrow, BorrowMut};
2
3use openvm_circuit::{
4 arch::*,
5 system::memory::{
6 offline_checker::{
7 MemoryBridge, MemoryReadAuxCols, MemoryReadAuxRecord, MemoryWriteAuxCols,
8 MemoryWriteBytesAuxRecord,
9 },
10 online::TracingMemory,
11 MemoryAddress, MemoryAuxColsFactory,
12 },
13};
14use openvm_circuit_primitives::{
15 bitwise_op_lookup::{BitwiseOperationLookupBus, SharedBitwiseOperationLookupChip},
16 utils::not,
17 ColumnsAir, StructReflection, StructReflectionHelper,
18};
19use openvm_circuit_primitives_derive::{AlignedBorrow, AlignedBytesBorrow};
20use openvm_instructions::{
21 instruction::Instruction,
22 program::DEFAULT_PC_STEP,
23 riscv::{RV32_CELL_BITS, RV32_MEMORY_AS, RV32_REGISTER_AS, RV32_REGISTER_NUM_LIMBS},
24 LocalOpcode,
25};
26use openvm_rv32im_transpiler::{
27 Rv32HintStoreOpcode,
28 Rv32HintStoreOpcode::{HINT_BUFFER, HINT_STOREW},
29 MAX_HINT_BUFFER_WORDS, MAX_HINT_BUFFER_WORDS_BITS,
30};
31use openvm_stark_backend::{
32 interaction::InteractionBuilder,
33 p3_air::{Air, AirBuilder, BaseAir},
34 p3_field::{Field, PrimeCharacteristicRing, PrimeField32},
35 p3_matrix::{dense::RowMajorMatrix, Matrix},
36 p3_maybe_rayon::prelude::*,
37 BaseAirWithPublicValues, PartitionedBaseAir,
38};
39
40use crate::adapters::{read_rv32_register, tracing_read, tracing_write};
41
42mod execution;
43
44#[cfg(feature = "cuda")]
45mod cuda;
46#[cfg(feature = "cuda")]
47pub use cuda::*;
48
49#[cfg(test)]
50mod tests;
51
52const REM_WORD_NUM_ZERO_LIMBS: usize = 2;
53
54#[repr(C)]
55#[derive(AlignedBorrow, StructReflection, Debug)]
56pub struct Rv32HintStoreCols<T> {
57 pub is_single: T,
59 pub is_buffer: T,
60 pub rem_words_limbs: [T; RV32_REGISTER_NUM_LIMBS],
62
63 pub from_state: ExecutionState<T>,
64 pub mem_ptr_ptr: T,
65 pub mem_ptr_limbs: [T; RV32_REGISTER_NUM_LIMBS],
66 pub mem_ptr_aux_cols: MemoryReadAuxCols<T>,
67
68 pub write_aux: MemoryWriteAuxCols<T, RV32_REGISTER_NUM_LIMBS>,
69 pub data: [T; RV32_REGISTER_NUM_LIMBS],
70
71 pub is_buffer_start: T,
73 pub num_words_ptr: T,
74 pub num_words_aux_cols: MemoryReadAuxCols<T>,
75}
76
77#[derive(Copy, Clone, Debug, derive_new::new, ColumnsAir)]
78#[columns_via(Rv32HintStoreCols<u8>)]
79pub struct Rv32HintStoreAir {
80 pub execution_bridge: ExecutionBridge,
81 pub memory_bridge: MemoryBridge,
82 pub bitwise_operation_lookup_bus: BitwiseOperationLookupBus,
83 pub offset: usize,
84 pointer_max_bits: usize,
85}
86
87impl<F: Field> BaseAir<F> for Rv32HintStoreAir {
88 fn width(&self) -> usize {
89 Rv32HintStoreCols::<F>::width()
90 }
91}
92
93impl<F: Field> BaseAirWithPublicValues<F> for Rv32HintStoreAir {}
94impl<F: Field> PartitionedBaseAir<F> for Rv32HintStoreAir {}
95
96impl<AB: InteractionBuilder> Air<AB> for Rv32HintStoreAir {
97 fn eval(&self, builder: &mut AB) {
98 let main = builder.main();
99 let local = main.row_slice(0).expect("window should have two elements");
100 let local_cols: &Rv32HintStoreCols<AB::Var> = (*local).borrow();
101 let next = main.row_slice(1).expect("window should have two elements");
102 let next_cols: &Rv32HintStoreCols<AB::Var> = (*next).borrow();
103
104 let timestamp: AB::Var = local_cols.from_state.timestamp;
105 let mut timestamp_delta: usize = 0;
106 let mut timestamp_pp = || {
107 timestamp_delta += 1;
108 timestamp + AB::Expr::from_usize(timestamp_delta - 1)
109 };
110
111 builder.assert_bool(local_cols.is_single);
112 builder.assert_bool(local_cols.is_buffer);
113 builder.assert_bool(local_cols.is_buffer_start);
114 builder
115 .when(local_cols.is_buffer_start)
116 .assert_one(local_cols.is_buffer);
117 builder.assert_bool(local_cols.is_single + local_cols.is_buffer);
118
119 let is_valid = local_cols.is_single + local_cols.is_buffer;
120 let is_start = local_cols.is_single + local_cols.is_buffer_start;
121 let is_end = not::<AB::Expr>(next_cols.is_buffer) + next_cols.is_buffer_start;
125
126 let mut rem_words = AB::Expr::ZERO;
127 let mut next_rem_words = AB::Expr::ZERO;
128 let mut mem_ptr = AB::Expr::ZERO;
129 let mut next_mem_ptr = AB::Expr::ZERO;
130 for i in (0..RV32_REGISTER_NUM_LIMBS).rev() {
131 rem_words =
132 rem_words * AB::F::from_u32(1 << RV32_CELL_BITS) + local_cols.rem_words_limbs[i];
133 next_rem_words = next_rem_words * AB::F::from_u32(1 << RV32_CELL_BITS)
134 + next_cols.rem_words_limbs[i];
135 mem_ptr = mem_ptr * AB::F::from_u32(1 << RV32_CELL_BITS) + local_cols.mem_ptr_limbs[i];
136 next_mem_ptr =
137 next_mem_ptr * AB::F::from_u32(1 << RV32_CELL_BITS) + next_cols.mem_ptr_limbs[i];
138 }
139
140 builder
142 .when_transition()
143 .when(not::<AB::Expr>(is_valid.clone()))
144 .assert_zero(next_cols.is_single + next_cols.is_buffer);
145
146 builder
148 .when(local_cols.is_single)
149 .assert_one(is_end.clone());
150 builder
151 .when_first_row()
152 .assert_one(not::<AB::Expr>(local_cols.is_buffer) + local_cols.is_buffer_start);
153
154 self.memory_bridge
156 .read(
157 MemoryAddress::new(AB::F::from_u32(RV32_REGISTER_AS), local_cols.mem_ptr_ptr),
158 local_cols.mem_ptr_limbs,
159 timestamp_pp(),
160 &local_cols.mem_ptr_aux_cols,
161 )
162 .eval(builder, is_start.clone());
163
164 self.memory_bridge
166 .read(
167 MemoryAddress::new(AB::F::from_u32(RV32_REGISTER_AS), local_cols.num_words_ptr),
168 local_cols.rem_words_limbs,
169 timestamp_pp(),
170 &local_cols.num_words_aux_cols,
171 )
172 .eval(builder, local_cols.is_buffer_start);
173
174 self.memory_bridge
176 .write(
177 MemoryAddress::new(AB::F::from_u32(RV32_MEMORY_AS), mem_ptr.clone()),
178 local_cols.data,
179 timestamp_pp(),
180 &local_cols.write_aux,
181 )
182 .eval(builder, is_valid.clone());
183 let expected_opcode = (local_cols.is_single
184 * AB::F::from_usize(HINT_STOREW as usize + self.offset))
185 + (local_cols.is_buffer * AB::F::from_usize(HINT_BUFFER as usize + self.offset));
186
187 self.execution_bridge
188 .execute_and_increment_pc(
189 expected_opcode,
190 [
191 local_cols.is_buffer * (local_cols.num_words_ptr),
192 local_cols.mem_ptr_ptr.into(),
193 AB::Expr::ZERO,
194 AB::Expr::from_u32(RV32_REGISTER_AS),
195 AB::Expr::from_u32(RV32_MEMORY_AS),
196 ],
197 local_cols.from_state,
198 rem_words.clone() * AB::F::from_usize(timestamp_delta),
199 )
200 .eval(builder, is_start.clone());
201
202 debug_assert!(
205 (8..16).contains(&MAX_HINT_BUFFER_WORDS_BITS),
206 "MAX_HINT_BUFFER_WORDS_BITS must be in [16, 23] for these constraints to work"
207 );
208 for i in 1..=REM_WORD_NUM_ZERO_LIMBS {
213 builder.assert_zero(local_cols.rem_words_limbs[RV32_REGISTER_NUM_LIMBS - i]);
214 }
215
216 self.bitwise_operation_lookup_bus
219 .send_range(
220 local_cols.mem_ptr_limbs[RV32_REGISTER_NUM_LIMBS - 1]
221 * AB::F::from_usize(
222 1 << (RV32_REGISTER_NUM_LIMBS * RV32_CELL_BITS - self.pointer_max_bits),
223 ),
224 local_cols.rem_words_limbs[RV32_REGISTER_NUM_LIMBS - 1 - REM_WORD_NUM_ZERO_LIMBS]
225 * AB::F::from_usize(
226 1 << ((RV32_REGISTER_NUM_LIMBS - REM_WORD_NUM_ZERO_LIMBS) * RV32_CELL_BITS
227 - MAX_HINT_BUFFER_WORDS_BITS),
228 ),
229 )
230 .eval(builder, is_start.clone());
231
232 for i in 0..RV32_REGISTER_NUM_LIMBS / 2 {
234 self.bitwise_operation_lookup_bus
235 .send_range(local_cols.data[2 * i], local_cols.data[(2 * i) + 1])
236 .eval(builder, is_valid.clone());
237 }
238
239 builder
245 .when(is_valid)
246 .when(is_end.clone())
247 .assert_one(rem_words.clone());
248
249 let mut when_buffer_transition = builder.when(not::<AB::Expr>(is_end.clone()));
250 when_buffer_transition.assert_one(rem_words.clone() - next_rem_words.clone());
259 when_buffer_transition.assert_eq(
266 next_mem_ptr.clone() - mem_ptr.clone(),
267 AB::F::from_usize(RV32_REGISTER_NUM_LIMBS),
268 );
269 when_buffer_transition.assert_eq(
270 timestamp + AB::F::from_usize(timestamp_delta),
271 next_cols.from_state.timestamp,
272 );
273 }
274}
275
276#[derive(Copy, Clone, Debug)]
277pub struct Rv32HintStoreMetadata {
278 num_words: usize,
279}
280
281impl MultiRowMetadata for Rv32HintStoreMetadata {
282 #[inline(always)]
283 fn get_num_rows(&self) -> usize {
284 self.num_words
285 }
286}
287
288pub type Rv32HintStoreLayout = MultiRowLayout<Rv32HintStoreMetadata>;
289
290#[repr(C)]
292#[derive(AlignedBytesBorrow, Debug)]
293pub struct Rv32HintStoreRecordHeader {
294 pub num_words: u32,
295
296 pub from_pc: u32,
297 pub timestamp: u32,
298
299 pub mem_ptr_ptr: u32,
300 pub mem_ptr: u32,
301 pub mem_ptr_aux_record: MemoryReadAuxRecord,
302
303 pub num_words_ptr: u32,
305 pub num_words_read: MemoryReadAuxRecord,
306}
307
308#[repr(C)]
310#[derive(AlignedBytesBorrow, Debug)]
311pub struct Rv32HintStoreVar {
312 pub data_write_aux: MemoryWriteBytesAuxRecord<RV32_REGISTER_NUM_LIMBS>,
313 pub data: [u8; RV32_REGISTER_NUM_LIMBS],
314}
315
316#[derive(Debug)]
320pub struct Rv32HintStoreRecordMut<'a> {
321 pub inner: &'a mut Rv32HintStoreRecordHeader,
322 pub var: &'a mut [Rv32HintStoreVar],
323}
324
325impl<'a> CustomBorrow<'a, Rv32HintStoreRecordMut<'a>, Rv32HintStoreLayout> for [u8] {
330 fn custom_borrow(&'a mut self, layout: Rv32HintStoreLayout) -> Rv32HintStoreRecordMut<'a> {
331 let (header_buf, rest) =
336 unsafe { self.split_at_mut_unchecked(size_of::<Rv32HintStoreRecordHeader>()) };
337
338 let (_, vars, _) = unsafe { rest.align_to_mut::<Rv32HintStoreVar>() };
343 Rv32HintStoreRecordMut {
344 inner: header_buf.borrow_mut(),
345 var: &mut vars[..layout.metadata.num_words],
346 }
347 }
348
349 unsafe fn extract_layout(&self) -> Rv32HintStoreLayout {
350 let header: &Rv32HintStoreRecordHeader = self.borrow();
351 MultiRowLayout::new(Rv32HintStoreMetadata {
352 num_words: header.num_words as usize,
353 })
354 }
355}
356
357impl SizedRecord<Rv32HintStoreLayout> for Rv32HintStoreRecordMut<'_> {
358 fn size(layout: &Rv32HintStoreLayout) -> usize {
359 let mut total_len = size_of::<Rv32HintStoreRecordHeader>();
360 total_len = total_len.next_multiple_of(align_of::<Rv32HintStoreVar>());
362 total_len += size_of::<Rv32HintStoreVar>() * layout.metadata.num_words;
363 total_len
364 }
365
366 fn alignment(_layout: &Rv32HintStoreLayout) -> usize {
367 align_of::<Rv32HintStoreRecordHeader>()
368 }
369}
370
371#[derive(Clone, Copy, derive_new::new)]
372pub struct Rv32HintStoreExecutor {
373 pub pointer_max_bits: usize,
374 pub offset: usize,
375}
376
377#[derive(Clone, derive_new::new)]
378pub struct Rv32HintStoreFiller {
379 pointer_max_bits: usize,
380 bitwise_lookup_chip: SharedBitwiseOperationLookupChip<RV32_CELL_BITS>,
381}
382
383impl<F, RA> PreflightExecutor<F, RA> for Rv32HintStoreExecutor
384where
385 F: PrimeField32,
386 for<'buf> RA:
387 RecordArena<'buf, MultiRowLayout<Rv32HintStoreMetadata>, Rv32HintStoreRecordMut<'buf>>,
388{
389 fn get_opcode_name(&self, opcode: usize) -> String {
390 if opcode == HINT_STOREW.global_opcode().as_usize() {
391 String::from("HINT_STOREW")
392 } else if opcode == HINT_BUFFER.global_opcode().as_usize() {
393 String::from("HINT_BUFFER")
394 } else {
395 unreachable!("unsupported opcode: {opcode}")
396 }
397 }
398
399 fn execute(
400 &self,
401 state: VmStateMut<F, TracingMemory, RA>,
402 instruction: &Instruction<F>,
403 ) -> Result<(), ExecutionError> {
404 let &Instruction {
405 opcode, a, b, d, e, ..
406 } = instruction;
407
408 let a = a.as_canonical_u32();
409 let b = b.as_canonical_u32();
410 debug_assert_eq!(d.as_canonical_u32(), RV32_REGISTER_AS);
411 debug_assert_eq!(e.as_canonical_u32(), RV32_MEMORY_AS);
412
413 let local_opcode = Rv32HintStoreOpcode::from_usize(opcode.local_opcode_idx(self.offset));
414
415 let num_words = if local_opcode == HINT_STOREW {
417 1
418 } else {
419 read_rv32_register(state.memory.data(), a)
420 };
421
422 if num_words == 0 {
424 return Err(ExecutionError::HintBufferZeroWords { pc: *state.pc });
425 }
426 if num_words > MAX_HINT_BUFFER_WORDS as u32 {
427 return Err(ExecutionError::HintBufferTooLarge {
428 pc: *state.pc,
429 num_words,
430 max_hint_buffer_words: MAX_HINT_BUFFER_WORDS as u32,
431 });
432 }
433
434 let record = state.ctx.alloc(MultiRowLayout::new(Rv32HintStoreMetadata {
435 num_words: num_words as usize,
436 }));
437
438 record.inner.from_pc = *state.pc;
439 record.inner.timestamp = state.memory.timestamp;
440 record.inner.mem_ptr_ptr = b;
441
442 record.inner.mem_ptr = u32::from_le_bytes(tracing_read(
443 state.memory,
444 RV32_REGISTER_AS,
445 b,
446 &mut record.inner.mem_ptr_aux_record.prev_timestamp,
447 ));
448
449 debug_assert!(record.inner.mem_ptr <= (1 << self.pointer_max_bits));
450 debug_assert!(num_words <= (1 << self.pointer_max_bits));
451
452 record.inner.num_words = num_words;
453 if local_opcode == HINT_STOREW {
454 state.memory.increment_timestamp();
455 record.inner.num_words_ptr = u32::MAX;
456 } else {
457 record.inner.num_words_ptr = a;
458 tracing_read::<RV32_REGISTER_NUM_LIMBS>(
459 state.memory,
460 RV32_REGISTER_AS,
461 record.inner.num_words_ptr,
462 &mut record.inner.num_words_read.prev_timestamp,
463 );
464 };
465
466 if state.streams.hint_stream.len() < RV32_REGISTER_NUM_LIMBS * num_words as usize {
467 return Err(ExecutionError::HintOutOfBounds { pc: *state.pc });
468 }
469
470 for idx in 0..(num_words as usize) {
471 if idx != 0 {
472 state.memory.increment_timestamp();
473 state.memory.increment_timestamp();
474 }
475
476 let data_f: [F; RV32_REGISTER_NUM_LIMBS] =
477 std::array::from_fn(|_| state.streams.hint_stream.pop_front().unwrap());
478 let data: [u8; RV32_REGISTER_NUM_LIMBS] =
479 data_f.map(|byte| byte.as_canonical_u32() as u8);
480
481 record.var[idx].data = data;
482
483 tracing_write(
484 state.memory,
485 RV32_MEMORY_AS,
486 record.inner.mem_ptr + (RV32_REGISTER_NUM_LIMBS * idx) as u32,
487 data,
488 &mut record.var[idx].data_write_aux.prev_timestamp,
489 &mut record.var[idx].data_write_aux.prev_data,
490 );
491 }
492 *state.pc = state.pc.wrapping_add(DEFAULT_PC_STEP);
493
494 Ok(())
495 }
496}
497
498impl<F: PrimeField32> TraceFiller<F> for Rv32HintStoreFiller {
499 fn fill_trace(
500 &self,
501 mem_helper: &MemoryAuxColsFactory<F>,
502 trace: &mut RowMajorMatrix<F>,
503 rows_used: usize,
504 ) {
505 if rows_used == 0 {
506 return;
507 }
508
509 let width = trace.width;
510 debug_assert_eq!(width, size_of::<Rv32HintStoreCols<u8>>());
511 let mut trace = &mut trace.values[..width * rows_used];
512 let mut sizes = Vec::with_capacity(rows_used);
513 let mut chunks = Vec::with_capacity(rows_used);
514
515 while !trace.is_empty() {
516 let record: &Rv32HintStoreRecordHeader =
521 unsafe { get_record_from_slice(&mut trace, ()) };
522 let (chunk, rest) = trace.split_at_mut(width * record.num_words as usize);
523 sizes.push(record.num_words);
524 chunks.push(chunk);
525 trace = rest;
526 }
527
528 let msl_rshift: u32 = ((RV32_REGISTER_NUM_LIMBS - 1) * RV32_CELL_BITS) as u32;
529 let msl_lshift: u32 =
530 (RV32_REGISTER_NUM_LIMBS * RV32_CELL_BITS - self.pointer_max_bits) as u32;
531
532 let rem_words_msl_lshift: u32 = ((RV32_REGISTER_NUM_LIMBS - REM_WORD_NUM_ZERO_LIMBS)
534 * RV32_CELL_BITS
535 - MAX_HINT_BUFFER_WORDS_BITS) as u32;
536
537 chunks
538 .par_iter_mut()
539 .zip(sizes.par_iter())
540 .for_each(|(chunk, &num_words)| {
541 let record: Rv32HintStoreRecordMut = unsafe {
548 get_record_from_slice(
549 chunk,
550 MultiRowLayout::new(Rv32HintStoreMetadata {
551 num_words: num_words as usize,
552 }),
553 )
554 };
555 debug_assert!(
560 num_words <= MAX_HINT_BUFFER_WORDS as u32,
561 "num_words must be <= MAX_HINT_BUFFER_WORDS"
562 );
563 self.bitwise_lookup_chip.request_range(
564 (record.inner.mem_ptr >> msl_rshift) << msl_lshift,
565 ((num_words
566 >> (RV32_CELL_BITS
567 * (RV32_REGISTER_NUM_LIMBS - 1 - REM_WORD_NUM_ZERO_LIMBS)))
568 & 0xFF)
569 << rem_words_msl_lshift,
570 );
571
572 let mut timestamp = record.inner.timestamp + num_words * 3;
573 let mut mem_ptr = record.inner.mem_ptr + num_words * RV32_REGISTER_NUM_LIMBS as u32;
574
575 chunk
580 .rchunks_exact_mut(width)
581 .zip(record.var.iter().enumerate().rev())
582 .for_each(|(row, (idx, var))| {
583 for pair in var.data.chunks_exact(2) {
584 self.bitwise_lookup_chip
585 .request_range(pair[0] as u32, pair[1] as u32);
586 }
587
588 let cols: &mut Rv32HintStoreCols<F> = row.borrow_mut();
589 let is_single = record.inner.num_words_ptr == u32::MAX;
590 timestamp -= 3;
591 if idx == 0 && !is_single {
592 mem_helper.fill(
593 record.inner.num_words_read.prev_timestamp,
594 timestamp + 1,
595 cols.num_words_aux_cols.as_mut(),
596 );
597 cols.num_words_ptr = F::from_u32(record.inner.num_words_ptr);
598 } else {
599 mem_helper.fill_zero(cols.num_words_aux_cols.as_mut());
600 cols.num_words_ptr = F::ZERO;
601 }
602
603 cols.is_buffer_start = F::from_bool(idx == 0 && !is_single);
604
605 cols.data = var.data.map(|x| F::from_u8(x));
607
608 cols.write_aux
609 .set_prev_data(var.data_write_aux.prev_data.map(|x| F::from_u8(x)));
610 mem_helper.fill(
611 var.data_write_aux.prev_timestamp,
612 timestamp + 2,
613 cols.write_aux.as_mut(),
614 );
615
616 if idx == 0 {
617 mem_helper.fill(
618 record.inner.mem_ptr_aux_record.prev_timestamp,
619 timestamp,
620 cols.mem_ptr_aux_cols.as_mut(),
621 );
622 } else {
623 mem_helper.fill_zero(cols.mem_ptr_aux_cols.as_mut());
624 }
625
626 mem_ptr -= RV32_REGISTER_NUM_LIMBS as u32;
627 cols.mem_ptr_limbs = mem_ptr.to_le_bytes().map(|x| F::from_u8(x));
628 cols.mem_ptr_ptr = F::from_u32(record.inner.mem_ptr_ptr);
629
630 cols.from_state.timestamp = F::from_u32(timestamp);
631 cols.from_state.pc = F::from_u32(record.inner.from_pc);
632
633 cols.rem_words_limbs = (num_words - idx as u32)
634 .to_le_bytes()
635 .map(|x| F::from_u8(x));
636 cols.is_buffer = F::from_bool(!is_single);
637 cols.is_single = F::from_bool(is_single);
638 });
639 })
640 }
641}
642
643pub type Rv32HintStoreChip<F> = VmChipWrapper<F, Rv32HintStoreFiller>;