[feat] Add exit code in the terminate transpilation (#662)

* Add exit code in the terminate transpiler * chore: make terminate I-type * compile terminate.S into ELF * fix:remove terminate.dump * fix: terminate * feat: AlignedBorrow on MerkleMemoryCols (#661) * Add CommittedProgram (#658) * chore: runs-on for benchmark (#651) * chore: runs-on for benchmark * chore: fix path * chore: clean up * feat: Clean up core chip (#663) * Add `is_valid` to CoreCols * Remove DUMMY opcode * Remove `f` and `g` operands * [feat] generic field expression vm chip (#653) * generic field expression vm chip * add ne simplified * double * fix * update * use struct * chore: fix runs_on (#665) * chore: fix git config * chore: add back merge to main control * Expose is_terminate as a public value (#666) * feat: Always increment timestamp (#667) * fix (#668) * Chore: transpile shamts in a more natural way (same functionality, but not a goldberg machine) (#669) * fix: Nop needs to increment timestamp now * fix: test * fix: build ELF and add command --------- Co-authored-by: Jonathan Wang <[email protected]> Co-authored-by: Arayi <[email protected]> Co-authored-by: Zach Langley <[email protected]> Co-authored-by: Xinding Wei <[email protected]> Co-authored-by: PangZhi <[email protected]> Co-authored-by: luffykai <[email protected]>
openvm-org · Oct 25, 2024 · deeb3d6 · deeb3d6
1 parent f3f2753
commit deeb3d6
Show file tree

Hide file tree

Showing 8 changed files with 80 additions and 38 deletions.
diff --git a/toolchain/riscv/transpiler/data/rv32im-terminate-from-as b/toolchain/riscv/transpiler/data/rv32im-terminate-from-as
diff --git a/toolchain/riscv/transpiler/data/terminate.S b/toolchain/riscv/transpiler/data/terminate.S
@@ -0,0 +1,23 @@
+#define CUSTOM_0 0x0b
+#define CUSTOM_1 0x2b
+
+.macro addmod_1 rd, rs1, rs2
+	.insn r CUSTOM_1, 0, 0, \rd, \rs1, \rs2
+.endm
+
+.macro addmod_2 rd, rs1, rs2
+	.insn r CUSTOM_1, 0, 4, \rd, \rs1, \rs2
+.endm
+
+.macro terminate ec
+	.insn i CUSTOM_0, 0, x0, x0, \ec
+.endm
+
+.global _start
+
+_start:
+	li zero, 1
+	add a0, a0, zero
+	bne a0, a1, 8
+	terminate 0
+	terminate 1
diff --git a/toolchain/riscv/transpiler/src/rrs.rs b/toolchain/riscv/transpiler/src/rrs.rs
@@ -245,12 +245,15 @@ impl<F: PrimeField32> InstructionProcessor for InstructionTranspiler<F> {
 
 fn process_custom_instruction<F: PrimeField32>(instruction_u32: u32) -> Instruction<F> {
     let opcode = (instruction_u32 & 0x7f) as u8;
-    let funct3 = ((instruction_u32 >> 12) & 3) as u8; // All our instructions are R- or I-type
+    let funct3 = ((instruction_u32 >> 12) & 0b111) as u8; // All our instructions are R- or I-type
 
     match opcode {
         0x0b => {
             match funct3 {
-                0b000 => Some(terminate()),
+                0b000 => {
+                    let imm = (instruction_u32 >> 20) & 0xfff;
+                    Some(terminate(imm.try_into().expect("exit code must be byte")))
+                }
                 0b001 => {
                     // keccak or poseidon
                     None
@@ -315,10 +318,11 @@ pub(crate) fn transpile<F: PrimeField32>(instructions_u32: &[u32]) -> Vec<Instru
     let mut instructions = Vec::new();
     let mut transpiler = InstructionTranspiler::<F>(PhantomData);
     for instruction_u32 in instructions_u32 {
+        dbg!(instruction_u32);
         // TODO: we probably want to forbid such instructions, but for now we just skip them
         if *instruction_u32 == 115 {
             eprintln!("trying to transpile ecall ({:x})", instruction_u32);
-            instructions.push(terminate());
+            instructions.push(terminate(1));
             continue;
         }
         let instruction = process_instruction(&mut transpiler, *instruction_u32)

diff --git a/toolchain/riscv/transpiler/src/tests.rs b/toolchain/riscv/transpiler/src/tests.rs
@@ -10,10 +10,20 @@ use stark_vm::{
 };
 use test_case::test_case;
 
-use crate::{elf::Elf, rrs::transpile, AxVmExe};
+use crate::{elf::Elf, rrs::transpile, AxVmExe, Program};
 
 type F = BabyBear;
 
+fn setup_vm_from_elf(elf_path: &str, config: VmConfig) -> Result<(VirtualMachine<F>, Program<F>)> {
+    let dir = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
+    let data = read(dir.join(elf_path))?;
+    let elf = Elf::decode(&data, MEM_SIZE as u32)?;
+    dbg!(&elf.instructions);
+    let exe = AxVmExe::<F>::from_elf(elf);
+    let vm = VirtualMachine::new(config).with_initial_memory(exe.memory_image);
+    Ok((vm, exe.program))
+}
+
 #[test]
 fn test_decode_elf() -> Result<()> {
     let dir = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
@@ -23,6 +33,9 @@ fn test_decode_elf() -> Result<()> {
     Ok(())
 }
 
+// To create ELF directly from .S file, `brew install riscv-gnu-toolchain` and run
+// `riscv64-unknown-elf-gcc -march=rv32im -mabi=ilp32 -nostartfiles -e _start -Ttext 0 fib.S -o rv32im-fib-from-as`
+// riscv64-unknown-elf-gcc supports rv32im if you set -march target
 #[test_case("data/rv32im-fib-from-as")]
 #[test_case("data/rv32im-intrin-from-as")]
 fn test_generate_program(elf_path: &str) -> Result<()> {
@@ -40,43 +53,35 @@ fn test_generate_program(elf_path: &str) -> Result<()> {
 #[test_case("data/rv32im-exp-from-as")]
 #[test_case("data/rv32im-fib-from-as")]
 fn test_rv32im_runtime(elf_path: &str) -> Result<()> {
-    let dir = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
-    let data = read(dir.join(elf_path))?;
-    let elf = Elf::decode(&data, MEM_SIZE as u32)?;
-    let exe = AxVmExe::<F>::from_elf(elf);
     setup_tracing();
     let config = VmConfig::rv32im();
-    let vm = VirtualMachine::new(config).with_initial_memory(exe.memory_image);
-
-    // TODO: use "execute_and_generate" when it's implemented
-
-    vm.execute(exe.program)?;
+    let (vm, program) = setup_vm_from_elf(elf_path, config)?;
+    vm.execute(program)?;
     Ok(())
 }
 
 #[test_case("data/rv32im-fibonacci-program-elf-release")]
 fn test_rv32i_prove(elf_path: &str) -> Result<()> {
-    let dir = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
-    let data = read(dir.join(elf_path))?;
-    let elf = Elf::decode(&data, MEM_SIZE as u32)?;
-    let exe = AxVmExe::from_elf(elf);
     let config = VmConfig::rv32i();
-    let vm = VirtualMachine::new(config).with_initial_memory(exe.memory_image);
-
-    air_test(vm, exe.program);
+    let (vm, program) = setup_vm_from_elf(elf_path, config)?;
+    air_test(vm, program);
     Ok(())
 }
 
 #[test_case("data/rv32im-intrin-from-as")]
 fn test_intrinsic_runtime(elf_path: &str) -> Result<()> {
-    let dir = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
-    let data = read(dir.join(elf_path))?;
-    let elf = Elf::decode(&data, MEM_SIZE as u32)?;
-    let exe = AxVmExe::<F>::from_elf(elf);
     setup_tracing();
     let config = VmConfig::rv32im().add_canonical_modulus();
-    let vm = VirtualMachine::new(config).with_initial_memory(exe.memory_image);
+    let (vm, program) = setup_vm_from_elf(elf_path, config)?;
+    vm.execute(program)?;
+    Ok(())
+}
 
-    vm.execute(exe.program)?;
+#[test]
+fn test_terminate_runtime() -> Result<()> {
+    setup_tracing();
+    let config = VmConfig::rv32i();
+    let (vm, program) = setup_vm_from_elf("data/rv32im-terminate-from-as", config)?;
+    air_test(vm, program);
     Ok(())
 }
diff --git a/toolchain/riscv/transpiler/src/util.rs b/toolchain/riscv/transpiler/src/util.rs
@@ -173,9 +173,10 @@ pub fn nop<F: PrimeField32>() -> Instruction<F> {
     }
 }
 
-pub fn terminate<F: PrimeField32>() -> Instruction<F> {
+pub fn terminate<F: PrimeField32>(code: u8) -> Instruction<F> {
     Instruction {
         opcode: TerminateOpcode::TERMINATE.with_default_offset(),
+        c: F::from_canonical_u8(code),
         ..Default::default()
     }
 }

diff --git a/vm/src/common/nop/mod.rs b/vm/src/common/nop/mod.rs
@@ -14,8 +14,9 @@ use p3_field::{AbstractField, Field, PrimeField32};
 use p3_matrix::{dense::RowMajorMatrix, Matrix};
 
 use crate::{
-    arch::{ExecutionBridge, ExecutionBus, ExecutionState, InstructionExecutor},
+    arch::{ExecutionBridge, ExecutionBus, ExecutionState, InstructionExecutor, PcIncOrSet},
     system::{
+        memory::MemoryControllerRef,
         program::{ExecutionError, Instruction, ProgramBus},
         DEFAULT_PC_STEP,
     },
@@ -56,33 +57,38 @@ impl<AB: AirBuilder + InteractionBuilder> Air<AB> for NopAir {
         } = (*local).borrow();
 
         self.execution_bridge
-            .execute(
+            .execute_and_increment_or_set_pc(
                 AB::Expr::from_canonical_usize(self.nop_opcode),
                 iter::empty::<AB::Expr>(),
                 ExecutionState::<AB::Expr>::new(pc, timestamp),
-                ExecutionState::<AB::Expr>::new(
-                    pc + AB::Expr::from_canonical_u32(DEFAULT_PC_STEP),
-                    timestamp,
-                ),
+                AB::Expr::one(),
+                PcIncOrSet::Inc(AB::Expr::from_canonical_u32(DEFAULT_PC_STEP)),
             )
             .eval(builder, is_valid);
     }
 }
 
-pub struct NopChip<F> {
+pub struct NopChip<F: Field> {
     pub air: NopAir,
     pub rows: Vec<NopCols<F>>,
     pub nop_opcode: usize,
+    memory: MemoryControllerRef<F>,
 }
 
-impl<F> NopChip<F> {
-    pub fn new(execution_bus: ExecutionBus, program_bus: ProgramBus, offset: usize) -> Self {
+impl<F: Field> NopChip<F> {
+    pub fn new(
+        execution_bus: ExecutionBus,
+        program_bus: ProgramBus,
+        memory_controller: MemoryControllerRef<F>,
+        offset: usize,
+    ) -> Self {
         Self {
             air: NopAir {
                 execution_bridge: ExecutionBridge::new(execution_bus, program_bus),
                 nop_opcode: offset + NopOpcode::NOP.as_usize(),
             },
             rows: vec![],
+            memory: memory_controller,
             nop_opcode: offset,
         }
     }
@@ -101,9 +107,10 @@ impl<F: PrimeField32> InstructionExecutor<F> for NopChip<F> {
             timestamp: F::from_canonical_u32(from_state.timestamp),
             is_valid: F::one(),
         });
+        self.memory.borrow_mut().increment_timestamp();
         Ok(ExecutionState::new(
             from_state.pc + DEFAULT_PC_STEP,
-            from_state.timestamp,
+            from_state.timestamp + 1,
         ))
     }
 

diff --git a/vm/src/common/nop/tests.rs b/vm/src/common/nop/tests.rs
@@ -16,6 +16,7 @@ fn test_nops_and_terminate() {
     let mut chip = NopChip::<F>::new(
         tester.execution_bus(),
         tester.program_bus(),
+        tester.memory_controller(),
         NopOpcode::default_offset(),
     );
 
@@ -26,7 +27,7 @@ fn test_nops_and_terminate() {
         tester.execute_with_pc(&mut chip, nop.clone(), state.pc.as_canonical_u32());
         let new_state = tester.execution.records.last().unwrap().final_state;
         assert_eq!(state.pc + F::from_canonical_usize(4), new_state.pc);
-        assert_eq!(state.timestamp, new_state.timestamp);
+        assert_eq!(state.timestamp + F::one(), new_state.timestamp);
         state = new_state;
     }
 

diff --git a/vm/src/system/vm/chip_set.rs b/vm/src/system/vm/chip_set.rs
@@ -291,6 +291,7 @@ impl VmConfig {
                     let nop_chip = Rc::new(RefCell::new(NopChip::new(
                         execution_bus,
                         program_bus,
+                        memory_controller.clone(),
                         offset,
                     )));
                     for opcode in range {