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

* Add custom insn macros

* Wip

* Wip

* Add "read-vec" function

* Add the new HintInputRv32 phantom instruction, add the corresponding intrinsic, support transpiling it

* Add a test

* Optimize reading a vector

* Address comments, add a test, fix transpiler

* Fix rust-v hint program

* chore: mv exit,panic to process

* feat: make `read_vec` more optimal

* fix: () vs !

* fix import

* fix: ptr was not ptr_start

---------

Co-authored-by: Jonathan Wang <[email protected]>

Author: Golovanov399

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/env/mod.rs

@@ -0,0 +1,17 @@
+// use crate::custom_insn_i;
+
+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) => {
+        axvm_platform::custom_insn_i!(axvm_platform::intrinsics::CUSTOM_0, 0b001, $x, "x0", $imm)
+    };
+}
+
+/// Reset the hint stream with the next hint.
+#[inline(always)]
+pub fn hint_input() {
+    custom_insn_i!(CUSTOM_0, 0b011, "x0", "x0", 0);
+}

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

@@ -1,5 +1,8 @@
 //! 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::*;

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

@@ -0,0 +1,47 @@
+//! User IO functions
+
+use alloc::vec::Vec;
+use core::alloc::Layout;
+
+use crate::{hint_store_u32, intrinsics::hint_input};
+
+/// Read `size: u32` and then `size` bytes from the hint stream into a vector.
+pub fn read_vec() -> Vec<u8> {
+    hint_input();
+    read_vec_by_len(read_u32() as usize)
+}
+
+/// Read the next 4 bytes from the hint stream into a register.
+/// Because [hint_store_u32] stores a word to memory, this function first reads to memory and then
+/// loads from memory to register.
+#[inline(always)]
+#[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> {
+    let num_words = (len + 3) / 4;
+    let capacity = num_words * 4;
+    // Allocate a buffer of the required length that is 4 byte aligned
+    // Note: this expect message doesn't matter until our panic handler actually cares about it
+    let layout = Layout::from_size_align(capacity, 4).expect("vec is too large");
+    // SAFETY: We populate a `Vec<u8>` by hintstore-ing `num_words` 4 byte words. We set the length to `len` and don't care about the extra `capacity - len` bytes stored.
+    let ptr_start = unsafe { alloc::alloc::alloc(layout) };
+    let mut ptr = ptr_start;
+
+    // Note: if len % 4 != 0, this will discard some last bytes
+    for _ in 0..num_words {
+        hint_store_u32!(ptr, 0);
+        ptr = unsafe { ptr.add(4) };
+    }
+    unsafe { Vec::from_raw_parts(ptr_start, len, capacity) }
+}

toolchain/riscv/axvm/src/io.rs

@@ -1,89 +1,16 @@
-// Copyright 2024 RISC Zero, Inc.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-//! The RISC Zero zkVM's guest-side RISC-V API.
-//!
-//! Code that is validated by the [RISC Zero zkVM](crate) is run inside the guest. In almost all
-//! practical cases, the guest will want to read private input data from the host and write public
-//! data to the journal. This can be done with [env::read] and [env::commit], respectively;
-//! additional I/O functionality is also available in [mod@env].
-//!
-//! ## Installation
-//!
-//! To build and run RISC Zero zkVM code, you will need to install the RISC Zero
-//! toolchain, which can be done using the rzup utility:
-//!
-//! ```sh
-//! curl -L https://risczero.com/install | bash
-//! rzup install
-//! ```
-//!
-//! ## Example
-//!
-//! The following guest code[^starter-ex] proves a number is
-//! composite by multiplying two unsigned integers, and panicking if either is
-//! `1` or if the multiplication overflows:
-//!
-//! ```ignore
-//! #![no_main]
-//! #![no_std]
-//!
-//! use risc0_zkvm::guest::env;
-//!
-//! risc0_zkvm::guest::entry!(main);
-//!
-//! fn main() {
-//!     // Load the first number from the host
-//!     let a: u64 = env::read();
-//!     // Load the second number from the host
-//!     let b: u64 = env::read();
-//!     // Verify that neither of them are 1 (i.e. nontrivial factors)
-//!     if a == 1 || b == 1 {
-//!         panic!("Trivial factors")
-//!     }
-//!     // Compute the product while being careful with integer overflow
-//!     let product = a.checked_mul(b).expect("Integer overflow");
-//!     env::commit(&product);
-//! }
-//! ```
-//!
-//! Notice how [env::read] is used to load the two factors, and [env::commit] is used to make their
-//! composite product publicly available. All input an output of your guest is private except for
-//! what is written to the journal with [env::commit].
-//!
-//! By default, the guest only has the Rust `core` libraries and not `std`. A partial
-//! implementation of the Rust standard libraries can be enabled with the `std` feature on this [crate].
-//! When this feature is not enabled, the lines including `#![no_std]` and `#![no_main]` are
-//! required, as well as the use of the [crate::guest::entry] macro. When `std` is enabled, these
-//! three lines can be omitted and many features of `std` can be used.
-//!
-//! If you encounter problems building zkVM guest code, you can see if we have a
-//! known workaround for your issue by looking in our
-//! [rust guest workarounds](https://github.com/risc0/risc0/issues?q=is%3Aissue+is%3Aopen+label%3A%22rust+guest+workarounds%22)
-//! tag on GitHub.
-//!
-//! [^starter-ex]: The example is based on the [Factors example](https://github.com/risc0/risc0/tree/main/examples/factors).
+//! # axVM
 
 #![cfg_attr(not(feature = "std"), no_std)]
 #![deny(rustdoc::broken_intra_doc_links)]
 #![deny(missing_docs)]
 #![cfg_attr(docsrs, feature(doc_cfg, doc_auto_cfg))]
+#![feature(asm_const)]
 
 extern crate alloc;
 
-pub mod env;
 pub mod intrinsics;
+pub mod io;
+pub mod process;
 pub mod serde;
 
 #[cfg(target_os = "zkvm")]
@@ -165,7 +92,8 @@ unsafe extern "C" fn __start() -> ! {
         main()
     }
 
-    env::exit::<0>();
+    process::exit();
+    unreachable!()
 }
 
 #[cfg(target_os = "zkvm")]

toolchain/riscv/axvm/src/lib.rs

@@ -0,0 +1,11 @@
+//! System exit and panic functions.
+
+/// 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/process.rs

@@ -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/config.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/Cargo.toml

@@ -0,0 +1,13 @@
+#![no_main]
+#![no_std]
+use axvm::io::read_vec;
+
+axvm::entry!(main);
+
+pub fn main() {
+    let vec = read_vec();
+    assert_eq!(vec.len(), 4);
+    for i in 0..4 {
+        assert_eq!(vec[i], i as u8);
+    }
+}

toolchain/riscv/examples/hint/program/src/main.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/intrinsics/mod.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")]