(draft) feat: GPU witness generation

Cargo.toml

@@ -127,8 +127,8 @@ 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/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" }

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,7 +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,
+    PackedNextAccessEntry, PreflightTracer, PreflightTracerConfig, ReadOp, ShardPlanBuilder,
+    StepCellExtractor, StepIndex,
     StepRecord, Tracer, WriteOp,
 };
 
@@ -34,7 +35,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));