feat: GPU witness generation (RV32IM + Keccak + ShardRam)

Cargo.toml

@@ -127,20 +127,20 @@ lto = "thin"
 #ceno_crypto_primitives = { path = "../ceno-patch/crypto-primitives", package = "ceno_crypto_primitives" }
 #ceno_syscall = { path = "../ceno-patch/syscall", package = "ceno_syscall" }
 
-#[patch."https://github.com/scroll-tech/ceno-gpu-mock.git"]
-#ceno_gpu = { path = "../ceno-gpu/cuda_hal", package = "cuda_hal", default-features = false, features = ["bb31"] }
-
-#[patch."https://github.com/scroll-tech/gkr-backend"]
-#ff_ext = { path = "../gkr-backend/crates/ff_ext", package = "ff_ext" }
-#mpcs = { path = "../gkr-backend/crates/mpcs", package = "mpcs" }
-#multilinear_extensions = { path = "../gkr-backend/crates/multilinear_extensions", package = "multilinear_extensions" }
-#p3 = { path = "../gkr-backend/crates/p3", package = "p3" }
-#poseidon = { path = "../gkr-backend/crates/poseidon", package = "poseidon" }
-#sp1-curves = { path = "../gkr-backend/crates/curves", package = "sp1-curves" }
-#sumcheck = { path = "../gkr-backend/crates/sumcheck", package = "sumcheck" }
-#transcript = { path = "../gkr-backend/crates/transcript", package = "transcript" }
-#whir = { path = "../gkr-backend/crates/whir", package = "whir" }
-#witness = { path = "../gkr-backend/crates/witness", package = "witness" }
+# [patch."https://github.com/scroll-tech/ceno-gpu-mock.git"]
+# ceno_gpu = { path = "../ceno-gpu/cuda_hal", package = "cuda_hal", default-features = false, features = ["bb31"] }
+
+# [patch."https://github.com/scroll-tech/gkr-backend"]
+# ff_ext = { path = "../gkr-backend/crates/ff_ext", package = "ff_ext" }
+# mpcs = { path = "../gkr-backend/crates/mpcs", package = "mpcs" }
+# multilinear_extensions = { path = "../gkr-backend/crates/multilinear_extensions", package = "multilinear_extensions" }
+# p3 = { path = "../gkr-backend/crates/p3", package = "p3" }
+# poseidon = { path = "../gkr-backend/crates/poseidon", package = "poseidon" }
+# sp1-curves = { path = "../gkr-backend/crates/curves", package = "sp1-curves" }
+# sumcheck = { path = "../gkr-backend/crates/sumcheck", package = "sumcheck" }
+# transcript = { path = "../gkr-backend/crates/transcript", package = "transcript" }
+# whir = { path = "../gkr-backend/crates/whir", package = "whir" }
+# witness = { path = "../gkr-backend/crates/witness", package = "witness" }
 
 # [patch."https://github.com/scroll-tech/openvm.git"]
 # openvm = { path = "../openvm-scroll-tech/crates/toolchain/openvm", default-features = false }

ceno_emul/src/addr.rs

@@ -30,12 +30,14 @@ pub type Word = u32;
 pub type SWord = i32;
 pub type Addr = u32;
 pub type Cycle = u64;
-pub type RegIdx = usize;
+pub type RegIdx = u8;
 
 #[derive(Clone, Copy, Default, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
+#[repr(C)]
 pub struct ByteAddr(pub u32);
 
 #[derive(Clone, Copy, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
+#[repr(C)]
 pub struct WordAddr(pub u32);
 
 impl From<ByteAddr> for WordAddr {

ceno_emul/src/disassemble/mod.rs

@@ -1,4 +1,7 @@
-use crate::rv32im::{InsnKind, Instruction};
+use crate::{
+    addr::RegIdx,
+    rv32im::{InsnKind, Instruction},
+};
 use itertools::izip;
 use rrs_lib::{
     InstructionProcessor,
@@ -19,9 +22,9 @@ impl Instruction {
     pub const fn from_r_type(kind: InsnKind, dec_insn: &RType, raw: u32) -> Self {
         Self {
             kind,
-            rd: dec_insn.rd,
-            rs1: dec_insn.rs1,
-            rs2: dec_insn.rs2,
+            rd: dec_insn.rd as RegIdx,
+            rs1: dec_insn.rs1 as RegIdx,
+            rs2: dec_insn.rs2 as RegIdx,
             imm: 0,
             raw,
         }
@@ -32,8 +35,8 @@ impl Instruction {
     pub const fn from_i_type(kind: InsnKind, dec_insn: &IType, raw: u32) -> Self {
         Self {
             kind,
-            rd: dec_insn.rd,
-            rs1: dec_insn.rs1,
+            rd: dec_insn.rd as RegIdx,
+            rs1: dec_insn.rs1 as RegIdx,
             imm: dec_insn.imm,
             rs2: 0,
             raw,
@@ -45,8 +48,8 @@ impl Instruction {
     pub const fn from_i_type_shamt(kind: InsnKind, dec_insn: &ITypeShamt, raw: u32) -> Self {
         Self {
             kind,
-            rd: dec_insn.rd,
-            rs1: dec_insn.rs1,
+            rd: dec_insn.rd as RegIdx,
+            rs1: dec_insn.rs1 as RegIdx,
             imm: dec_insn.shamt as i32,
             rs2: 0,
             raw,
@@ -59,8 +62,8 @@ impl Instruction {
         Self {
             kind,
             rd: 0,
-            rs1: dec_insn.rs1,
-            rs2: dec_insn.rs2,
+            rs1: dec_insn.rs1 as RegIdx,
+            rs2: dec_insn.rs2 as RegIdx,
             imm: dec_insn.imm,
             raw,
         }
@@ -72,8 +75,8 @@ impl Instruction {
         Self {
             kind,
             rd: 0,
-            rs1: dec_insn.rs1,
-            rs2: dec_insn.rs2,
+            rs1: dec_insn.rs1 as RegIdx,
+            rs2: dec_insn.rs2 as RegIdx,
             imm: dec_insn.imm,
             raw,
         }
@@ -231,7 +234,7 @@ impl InstructionProcessor for InstructionTranspiler {
     fn process_jal(&mut self, dec_insn: JType) -> Self::InstructionResult {
         Instruction {
             kind: InsnKind::JAL,
-            rd: dec_insn.rd,
+            rd: dec_insn.rd as RegIdx,
             rs1: 0,
             rs2: 0,
             imm: dec_insn.imm,
@@ -242,8 +245,8 @@ impl InstructionProcessor for InstructionTranspiler {
     fn process_jalr(&mut self, dec_insn: IType) -> Self::InstructionResult {
         Instruction {
             kind: InsnKind::JALR,
-            rd: dec_insn.rd,
-            rs1: dec_insn.rs1,
+            rd: dec_insn.rd as RegIdx,
+            rs1: dec_insn.rs1 as RegIdx,
             rs2: 0,
             imm: dec_insn.imm,
             raw: self.word,
@@ -265,7 +268,7 @@ impl InstructionProcessor for InstructionTranspiler {
             // See [`InstructionTranspiler::process_auipc`] for more background on the conversion.
             Instruction {
                 kind: InsnKind::ADDI,
-                rd: dec_insn.rd,
+                rd: dec_insn.rd as RegIdx,
                 rs1: 0,
                 rs2: 0,
                 imm: dec_insn.imm,
@@ -276,7 +279,7 @@ impl InstructionProcessor for InstructionTranspiler {
         {
             Instruction {
                 kind: InsnKind::LUI,
-                rd: dec_insn.rd,
+                rd: dec_insn.rd as RegIdx,
                 rs1: 0,
                 rs2: 0,
                 imm: dec_insn.imm,
@@ -311,7 +314,7 @@ impl InstructionProcessor for InstructionTranspiler {
             // real world scenarios like a `reth` run.
             Instruction {
                 kind: InsnKind::ADDI,
-                rd: dec_insn.rd,
+                rd: dec_insn.rd as RegIdx,
                 rs1: 0,
                 rs2: 0,
                 imm: dec_insn.imm.wrapping_add(pc as i32),
@@ -322,7 +325,7 @@ impl InstructionProcessor for InstructionTranspiler {
         {
             Instruction {
                 kind: InsnKind::AUIPC,
-                rd: dec_insn.rd,
+                rd: dec_insn.rd as RegIdx,
                 rs1: 0,
                 rs2: 0,
                 imm: dec_insn.imm,

ceno_emul/src/lib.rs

@@ -11,8 +11,8 @@ pub use platform::{CENO_PLATFORM, Platform};
 mod tracer;
 pub use tracer::{
     Change, FullTracer, FullTracerConfig, LatestAccesses, MemOp, NextAccessPair, NextCycleAccess,
-    PreflightTracer, PreflightTracerConfig, ReadOp, ShardPlanBuilder, StepCellExtractor, StepIndex,
-    StepRecord, Tracer, WriteOp,
+    PackedNextAccessEntry, PreflightTracer, PreflightTracerConfig, ReadOp, ShardPlanBuilder,
+    StepCellExtractor, StepIndex, StepRecord, Tracer, WriteOp,
 };
 
 mod vm_state;
@@ -34,7 +34,7 @@ pub use syscalls::{
     BN254_FP_MUL, BN254_FP2_ADD, BN254_FP2_MUL, KECCAK_PERMUTE, SECP256K1_ADD,
     SECP256K1_DECOMPRESS, SECP256K1_DOUBLE, SECP256K1_SCALAR_INVERT, SECP256R1_ADD,
     SECP256R1_DECOMPRESS, SECP256R1_DOUBLE, SECP256R1_SCALAR_INVERT, SHA_EXTEND, SyscallSpec,
-    UINT256_MUL,
+    SyscallWitness, UINT256_MUL,
     bn254::{
         BN254_FP_WORDS, BN254_FP2_WORDS, BN254_POINT_WORDS, Bn254AddSpec, Bn254DoubleSpec,
         Bn254Fp2AddSpec, Bn254Fp2MulSpec, Bn254FpAddSpec, Bn254FpMulSpec,

ceno_emul/src/platform.rs

@@ -134,7 +134,7 @@ impl Platform {
     /// Virtual address of a register.
     pub const fn register_vma(index: RegIdx) -> Addr {
         // Register VMAs are aligned, cannot be confused with indices, and readable in hex.
-        (index << 8) as Addr
+        (index as Addr) << 8
     }
 
     /// Register index from a virtual address (unchecked).
@@ -220,7 +220,7 @@ mod tests {
         // Registers do not overlap with ROM or RAM.
         for reg in [
             Platform::register_vma(0),
-            Platform::register_vma(VMState::<PreflightTracer>::REG_COUNT - 1),
+            Platform::register_vma((VMState::<PreflightTracer>::REG_COUNT - 1) as RegIdx),
         ] {
             assert!(!p.is_rom(reg));
             assert!(!p.is_ram(reg));

ceno_emul/src/rv32im.rs

@@ -29,9 +29,9 @@ use super::addr::{ByteAddr, RegIdx, WORD_SIZE, Word, WordAddr};
 pub const fn encode_rv32(kind: InsnKind, rs1: u32, rs2: u32, rd: u32, imm: i32) -> Instruction {
     Instruction {
         kind,
-        rs1: rs1 as usize,
-        rs2: rs2 as usize,
-        rd: rd as usize,
+        rs1: rs1 as RegIdx,
+        rs2: rs2 as RegIdx,
+        rd: rd as RegIdx,
         imm,
         raw: 0,
     }
@@ -43,9 +43,9 @@ pub const fn encode_rv32(kind: InsnKind, rs1: u32, rs2: u32, rd: u32, imm: i32)
 pub const fn encode_rv32u(kind: InsnKind, rs1: u32, rs2: u32, rd: u32, imm: u32) -> Instruction {
     Instruction {
         kind,
-        rs1: rs1 as usize,
-        rs2: rs2 as usize,
-        rd: rd as usize,
+        rs1: rs1 as RegIdx,
+        rs2: rs2 as RegIdx,
+        rd: rd as RegIdx,
         imm: imm as i32,
         raw: 0,
     }
@@ -113,6 +113,7 @@ pub enum TrapCause {
 }
 
 #[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord)]
+#[repr(C)]
 pub struct Instruction {
     pub kind: InsnKind,
     pub rs1: RegIdx,
@@ -162,6 +163,7 @@ use InsnFormat::*;
     ToPrimitive,
     Default,
 )]
+#[repr(u8)]
 #[allow(clippy::upper_case_acronyms)]
 pub enum InsnKind {
     #[default]
@@ -425,7 +427,7 @@ fn step_compute<M: EmuContext>(ctx: &mut M, kind: InsnKind, insn: &Instruction)
     if !new_pc.is_aligned() {
         return ctx.trap(TrapCause::InstructionAddressMisaligned);
     }
-    ctx.store_register(insn.rd_internal() as usize, out)?;
+    ctx.store_register(insn.rd_internal() as RegIdx, out)?;
     ctx.set_pc(new_pc);
     Ok(true)
 }
@@ -502,7 +504,7 @@ fn step_load<M: EmuContext>(ctx: &mut M, kind: InsnKind, decoded: &Instruction)
         }
         _ => unreachable!(),
     };
-    ctx.store_register(decoded.rd_internal() as usize, out)?;
+    ctx.store_register(decoded.rd_internal() as RegIdx, out)?;
     ctx.set_pc(ctx.get_pc() + WORD_SIZE);
     Ok(true)
 }

ceno_emul/src/syscalls.rs

@@ -60,19 +60,15 @@ pub fn handle_syscall<T: Tracer>(vm: &VMState<T>, function_code: u32) -> Result<
 /// A syscall event, available to the circuit witness generators.
 /// TODO: separate mem_ops into two stages: reads-and-writes
 #[derive(Clone, Debug, Default, PartialEq, Eq)]
+#[non_exhaustive]
 pub struct SyscallWitness {
     pub mem_ops: Vec<WriteOp>,
     pub reg_ops: Vec<WriteOp>,
-    _marker: (),
 }
 
 impl SyscallWitness {
     fn new(mem_ops: Vec<WriteOp>, reg_ops: Vec<WriteOp>) -> SyscallWitness {
-        SyscallWitness {
-            mem_ops,
-            reg_ops,
-            _marker: (),
-        }
+        SyscallWitness { mem_ops, reg_ops }
     }
 }
 

ceno_emul/src/test_utils.rs

@@ -1,10 +1,12 @@
 use crate::{
     CENO_PLATFORM, InsnKind, Instruction, Platform, Program, StepRecord, VMState, encode_rv32,
-    encode_rv32u, syscalls::KECCAK_PERMUTE,
+    encode_rv32u,
+    syscalls::{KECCAK_PERMUTE, SyscallWitness},
+    tracer::FullTracerConfig,
 };
 use anyhow::Result;
 
-pub fn keccak_step() -> (StepRecord, Vec<Instruction>) {
+pub fn keccak_step() -> (StepRecord, Vec<Instruction>, Vec<SyscallWitness>) {
     let instructions = vec![
         // Call Keccak-f.
         load_immediate(Platform::reg_arg0() as u32, CENO_PLATFORM.heap.start),
@@ -23,11 +25,16 @@ pub fn keccak_step() -> (StepRecord, Vec<Instruction>) {
         instructions.clone(),
         Default::default(),
     );
-    let mut vm = VMState::new(CENO_PLATFORM.clone(), program.into());
+    let mut vm: VMState = VMState::new_with_tracer_config(
+        CENO_PLATFORM.clone(),
+        program.into(),
+        FullTracerConfig { max_step_shard: 10 },
+    );
     vm.iter_until_halt().collect::<Result<Vec<_>>>().unwrap();
     let steps = vm.tracer().recorded_steps();
+    let syscall_witnesses = vm.tracer().syscall_witnesses().to_vec();
 
-    (steps[2].clone(), instructions)
+    (steps[2], instructions, syscall_witnesses)
 }
 
 const fn load_immediate(rd: u32, imm: u32) -> Instruction {