[feat] Add the functionality for custom RISC-V intrinsics, support hints in RISC-V

toolchain/build/src/lib.rs

@@ -257,11 +257,11 @@ pub fn build_guest_package<P>(
         "--target-dir",
         target_dir.as_ref().to_str().unwrap(),
     ]);
-    tty_println(&format!("cargo command: {:?}", cmd));
 
     if !is_debug() {
         cmd.args(["--release"]);
     }
+    tty_println(&format!("cargo command: {:?}", cmd));
 
     let mut child = cmd
         .stderr(Stdio::piped())

toolchain/instructions/src/phantom.rs

@@ -13,6 +13,8 @@ pub enum PhantomInstruction {
     PrintF,
     /// Prepare the next input vector for hinting.
     HintInput,
+    /// Prepare the next input vector for hinting, but prepend it with a 4-byte decomposition of its length instead of one field element.
+    HintInputRv32,
     /// Prepare the little-endian bit decomposition of a variable for hinting.
     HintBits,
     /// Start tracing

toolchain/riscv/axvm/src/intrinsics/io.rs

@@ -0,0 +1,57 @@
+// use crate::custom_insn_i;
+
+use alloc::{alloc::Layout, vec::Vec};
+
+use axvm_platform::{custom_insn_i, intrinsics::CUSTOM_0};
+
+/// Store the next 4 bytes from the hint stream to [[rd] + imm]_2.
+#[macro_export]
+macro_rules! hint_store_u32 {
+    ($x:ident, $imm:expr) => {
+        custom_insn_i!(CUSTOM_0, 0b001, $x, "x0", $imm)
+    };
+}
+
+/// Read the next 4 bytes from the hint stream.
+#[allow(asm_sub_register)]
+pub fn read_u32() -> u32 {
+    let ptr = unsafe { alloc::alloc::alloc(Layout::from_size_align(4, 4).unwrap()) };
+    let addr = ptr as u32;
+    hint_store_u32!(addr, 0);
+    let result: u32;
+    unsafe {
+        core::arch::asm!("lw {rd}, ({rs1})", rd = out(reg) result, rs1 = in(reg) addr);
+    }
+    result
+}
+
+/// Read the next `len` bytes from the hint stream into a vector.
+fn read_vec_by_len(len: usize) -> Vec<u8> {
+    // Note: this expect message doesn't matter until our panic handler actually cares about it
+    let layout = Layout::from_size_align((len + 3) / 4 * 4, 4).expect("vec is too large");
+    let ptr = unsafe { alloc::alloc::alloc(layout) };
+    let mut x: u32 = 0;
+    // Note: if len % 4 != 0, this will discard some last bytes
+    for i in 0..len {
+        if i % 4 == 0 {
+            // TODO: it probably makes sense not to juggle the data between registers and memory here.
+            // On the other hand, maybe it's not a big deal.
+            x = read_u32();
+        }
+        unsafe {
+            ptr.add(i).write_volatile((x & 255) as u8);
+        }
+        x >>= 8;
+    }
+    unsafe { Vec::from_raw_parts(ptr, len, len) }
+}
+
+/// Read `size: u32` and then `size` bytes from the hint stream into a vector.
+pub fn read_vec() -> Vec<u8> {
+    read_vec_by_len(read_u32() as usize)
+}
+
+/// Reset the hint stream with the next hint.
+pub fn hint_input() {
+    custom_insn_i!(CUSTOM_0, 0b011, "x0", "x0", 0);
+}

toolchain/riscv/axvm/src/intrinsics/mod.rs

@@ -1,5 +1,18 @@
 //! Functions that call custom instructions that use axVM intrinsic instructions.
 
 mod hash;
+/// Library functions for user input/output.
+pub mod io;
 
 pub use hash::*;
+pub use io::*;
+
+/// Exit the program with exit code 0.
+pub fn exit() {
+    axvm_platform::rust_rt::terminate::<0>();
+}
+
+/// Exit the program with exit code 1.
+pub fn panic() {
+    axvm_platform::rust_rt::terminate::<1>();
+}

toolchain/riscv/axvm/src/lib.rs

@@ -79,6 +79,7 @@
 #![deny(rustdoc::broken_intra_doc_links)]
 #![deny(missing_docs)]
 #![cfg_attr(docsrs, feature(doc_cfg, doc_auto_cfg))]
+#![feature(asm_const)]
 
 extern crate alloc;
 

toolchain/riscv/examples/hint/program/.cargo/config.toml

@@ -0,0 +1,5 @@
+[build]
+target = "riscv32im-risc0-zkvm-elf"
+
+[unstable]
+build-std = ["core", "alloc", "proc_macro", "panic_abort"]

toolchain/riscv/examples/hint/program/Cargo.toml

@@ -0,0 +1,8 @@
+[workspace]
+[package]
+version = "0.1.0"
+name = "axvm-hint-program"
+edition = "2021"
+
+[dependencies]
+axvm = { path = "../../../axvm" }

toolchain/riscv/examples/hint/program/src/main.rs

@@ -0,0 +1,12 @@
+#![no_main]
+#![no_std]
+
+axvm::entry!(main);
+
+pub fn main() {
+    axvm::intrinsics::io::hint_input();
+    let vec = axvm::intrinsics::io::read_size_and_vec();
+    if vec.iter().sum::<u8>() == 0 {
+        axvm::intrinsics::panic();
+    }
+}

toolchain/riscv/platform/src/intrinsics/mod.rs

@@ -0,0 +1,53 @@
+pub const CUSTOM_0: u8 = 0x0b;
+pub const CUSTOM_1: u8 = 0x2b;
+
+#[macro_export]
+macro_rules! custom_insn_i {
+    ($opcode:expr, $funct3:expr, $rd:literal, $rs1:literal, $imm:expr) => {
+        unsafe {
+            core::arch::asm!(concat!(
+                ".insn i {opcode}, {funct3}, ",
+                $rd,
+                ", ",
+                $rs1,
+                ", {imm}",
+            ), opcode = const $opcode, funct3 = const $funct3, imm = const $imm)
+        }
+    };
+    ($opcode:expr, $funct3:expr, $x:ident, $rs1:literal, $imm:expr) => {
+        unsafe {
+            core::arch::asm!(concat!(
+                ".insn i {opcode}, {funct3}, {rd}, ",
+                $rs1,
+                ", {imm}",
+            ), opcode = const $opcode, funct3 = const $funct3, rd = in(reg) $x, imm = const $imm)
+        }
+    };
+}
+
+#[macro_export]
+macro_rules! custom_insn_r {
+    ($opcode:expr, $funct3:expr, $rd:literal, $rs1:literal, $rs2:literal) => {
+        unsafe {
+            core::arch::asm!(concat!(
+                ".insn r {opcode}, {funct3}, ",
+                $rd,
+                ", ",
+                $rs1,
+                ", ",
+                $rs2,
+            ), opcode = const $opcode, funct3 = const $funct3)
+        }
+    };
+    ($opcode:expr, $funct3:expr, $x:ident, $rs1:literal, $rs2:literal) => {
+        unsafe {
+            core::arch::asm!(concat!(
+                ".insn r {opcode}, {funct3}, {rd}, ",
+                $rs1,
+                ", ",
+                $rs2,
+            ), opcode = const $opcode, funct3 = const $funct3, rd = out(reg) $x)
+        }
+    };
+    // TODO: implement more variants with like rs1 = in(reg) $y etc
+}

toolchain/riscv/platform/src/lib.rs

@@ -26,6 +26,7 @@ pub mod memory;
 // mod getrandom;
 #[cfg(all(feature = "rust-runtime", target_os = "zkvm"))]
 pub mod heap;
+pub mod intrinsics;
 #[cfg(all(feature = "export-libm", target_os = "zkvm"))]
 mod libm_extern;
 #[cfg(feature = "rust-runtime")]

toolchain/riscv/platform/src/rust_rt.rs

@@ -23,17 +23,17 @@
 #[cfg(target_os = "zkvm")]
 use core::arch::asm;
 
+#[cfg(target_os = "zkvm")]
+use crate::{custom_insn_i, custom_insn_r, intrinsics::CUSTOM_0};
+
 extern crate alloc;
 
 #[inline(always)]
 pub fn terminate<const EXIT_CODE: u8>() {
     #[cfg(target_os = "zkvm")]
-    unsafe {
-        asm!(".insn i 0x0b, 0, x0, x0, {ec}", ec = const EXIT_CODE)
-    };
+    custom_insn_i!(CUSTOM_0, 0, "x0", "x0", EXIT_CODE);
     #[cfg(not(target_os = "zkvm"))]
     {
-        core::hint::black_box(());
         unimplemented!()
     }
 }

toolchain/riscv/transpiler/src/rrs.rs

@@ -9,8 +9,10 @@ use axvm_circuit::{
     rv32im::adapters::RV32_REGISTER_NUM_LIMBS,
 };
 use axvm_instructions::{
-    instruction::Instruction, riscv::RvIntrinsic, EccOpcode, Rv32ModularArithmeticOpcode,
+    instruction::Instruction, riscv::RvIntrinsic, EccOpcode, PhantomInstruction,
+    Rv32HintStoreOpcode, Rv32ModularArithmeticOpcode,
 };
+use axvm_platform::intrinsics::{CUSTOM_0, CUSTOM_1};
 use p3_field::PrimeField32;
 use rrs_lib::{
     instruction_formats::{BType, IType, ITypeShamt, JType, RType, SType, UType},
@@ -245,24 +247,32 @@ fn process_custom_instruction<F: PrimeField32>(instruction_u32: u32) -> Instruct
     let funct3 = ((instruction_u32 >> 12) & 0b111) as u8; // All our instructions are R- or I-type
 
     match opcode {
-        0x0b => {
-            match funct3 {
-                0b000 => {
-                    let imm = (instruction_u32 >> 20) & 0xfff;
-                    Some(terminate(imm.try_into().expect("exit code must be byte")))
-                }
-                0b001 => {
-                    // keccak or poseidon
-                    None
-                }
-                0b010 => {
-                    // u256
-                    todo!("Implement u256 transpiler");
-                }
-                _ => None,
+        CUSTOM_0 => match funct3 {
+            0b000 => {
+                let imm = (instruction_u32 >> 20) & 0xfff;
+                Some(terminate(imm.try_into().expect("exit code must be byte")))
             }
-        }
-        0x2b => {
+            0b001 => {
+                let rd = (instruction_u32 >> 7) & 0x1f;
+                let imm = (instruction_u32 >> 20) & 0xfff;
+                Some(Instruction::from_isize(
+                    Rv32HintStoreOpcode::HINT_STOREW.with_default_offset(),
+                    0,
+                    (RV32_REGISTER_NUM_LIMBS * rd as usize) as isize,
+                    imm as isize,
+                    1,
+                    2,
+                ))
+            }
+            0b011 => Some(Instruction::phantom(
+                PhantomInstruction::HintInputRv32,
+                F::zero(),
+                F::zero(),
+                0,
+            )),
+            _ => unimplemented!(),
+        },
+        CUSTOM_1 => {
             match funct3 {
                 Rv32ModularArithmeticOpcode::FUNCT3 => {
                     // mod operations
@@ -315,7 +325,6 @@ 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 {
-        // TODO: we probably want to forbid such instructions, but for now we just skip them
         assert!(*instruction_u32 != 115, "ecall is not supported");
         let instruction = process_instruction(&mut transpiler, *instruction_u32)
             .unwrap_or_else(|| process_custom_instruction(*instruction_u32));

toolchain/riscv/transpiler/src/tests.rs

@@ -12,6 +12,7 @@ use axvm_circuit::{
 use axvm_platform::memory::MEM_SIZE;
 use eyre::Result;
 use p3_baby_bear::BabyBear;
+use p3_field::AbstractField;
 use tempfile::tempdir;
 use test_case::test_case;
 
@@ -91,6 +92,30 @@ fn test_rv32i_prove(examples_path: &str, min_segments: usize) -> Result<()> {
     Ok(())
 }
 
+#[test]
+fn test_rv32i_prove_with_hint() -> Result<()> {
+    let pkg = get_package(get_examples_dir().join("hint/program"));
+    let target_dir = tempdir()?;
+    let guest_opts = GuestOptions::default().into();
+    build_guest_package(&pkg, &target_dir, &guest_opts, None);
+    let elf_path = guest_methods(&pkg, &target_dir, &[]).pop().unwrap();
+    let config = VmConfig {
+        max_segment_len: (1 << 18) - 1,
+        ..VmConfig::rv32i()
+    };
+    let (_, exe) = setup_executor_from_elf(elf_path, config.clone())?;
+    air_test_with_min_segments(
+        config,
+        exe,
+        vec![[1, 0, 0, 0]
+            .iter()
+            .map(|x| F::from_canonical_u32(*x))
+            .collect()],
+        1,
+    );
+    Ok(())
+}
+
 #[test_case("data/rv32im-intrin-from-as")]
 fn test_intrinsic_runtime(elf_path: &str) -> Result<()> {
     setup_tracing();

vm/src/system/phantom/mod.rs

@@ -137,7 +137,7 @@ impl<F: PrimeField32> InstructionExecutor<F> for PhantomChip<F> {
                 let value = RefCell::borrow(&self.memory).unsafe_read_cell(addr_space, a);
                 println!("{}", value);
             }
-            PhantomInstruction::HintInput => {
+            PhantomInstruction::HintInput | PhantomInstruction::HintInputRv32 => {
                 let mut streams = self.streams.get().unwrap().lock();
                 let hint = match streams.input_stream.pop_front() {
                     Some(hint) => hint,
@@ -146,9 +146,18 @@ impl<F: PrimeField32> InstructionExecutor<F> for PhantomChip<F> {
                     }
                 };
                 streams.hint_stream.clear();
-                streams
-                    .hint_stream
-                    .push_back(F::from_canonical_usize(hint.len()));
+                if phantom == PhantomInstruction::HintInputRv32 {
+                    streams.hint_stream.extend(
+                        (hint.len() as u32)
+                            .to_le_bytes()
+                            .iter()
+                            .map(|b| F::from_canonical_u8(*b)),
+                    );
+                } else {
+                    streams
+                        .hint_stream
+                        .push_back(F::from_canonical_usize(hint.len()));
+                }
                 streams.hint_stream.extend(hint);
                 drop(streams);
             }