zkVM: change segment storage behavior

risc0/circuit/rv32im/src/prove/emu/exec/mod.rs

@@ -57,6 +57,11 @@ use crate::{
 
 pub const DEFAULT_SEGMENT_LIMIT_PO2: usize = 20;
 
+// Segments can either be stored in ram or on disk.
+// This value specifies the amount of segments that should be stored in ram
+// before switching to disk in order to avoid depleting memory.
+pub const DEFAULT_SEGMENT_RAM_STORAGE_LIMIT: u32 = 64;
+
 /// A host-side implementation of a system call.
 pub trait Syscall {
     /// Invokes the system call.

risc0/zkvm/src/host/api/client.rs

@@ -439,6 +439,7 @@ impl Client {
                 .as_ref()
                 .map(|x| x.path().to_string_lossy().into())
                 .unwrap_or_default(),
+            segment_limit_ram_storage: env.segment_limit_ram_storage,
         })
     }
 

risc0/zkvm/src/host/api/server.rs

@@ -16,16 +16,19 @@ use std::{
     error::Error as StdError,
     io::{BufReader, Error as IoError, ErrorKind as IoErrorKind, Read, Write},
     path::{Path, PathBuf},
+    sync::Arc,
 };
 
 use anyhow::{anyhow, bail, Result};
 use bytes::Bytes;
 use prost::Message;
+use risc0_circuit_rv32im::prove::emu::exec::DEFAULT_SEGMENT_RAM_STORAGE_LIMIT;
+use tempfile::tempdir;
 
 use super::{malformed_err, path_to_string, pb, ConnectionWrapper, Connector, TcpConnector};
 use crate::{
     get_prover_server, get_version,
-    host::{client::slice_io::SliceIo, server::session::NullSegmentRef},
+    host::{client::env::SegmentPath, client::slice_io::SliceIo, server::session::NullSegmentRef},
     receipt_claim::{MaybePruned, ReceiptClaim},
     ExecutorEnv, ExecutorImpl, ProverOpts, Receipt, Segment, SegmentReceipt, SuccinctReceipt,
     TraceCallback, TraceEvent, VerifierContext,
@@ -259,18 +262,39 @@ impl Server {
             request: pb::api::ExecuteRequest,
         ) -> Result<pb::api::ServerReply> {
             let env_request = request.env.ok_or(malformed_err())?;
-            let env = build_env(conn, &env_request)?;
+            let mut env = build_env(conn, &env_request)?;
+
+            if env.segment_path.is_none() {
+                env.segment_path = Some(SegmentPath::TempDir(Arc::new(tempdir()?)));
+            }
 
             let binary = env_request.binary.ok_or(malformed_err())?;
 
-            let segments_out = request.segments_out.ok_or(malformed_err())?;
+            let segment_path = env.segment_path.clone().unwrap();
+            let segment_path = segment_path.path().to_str().unwrap();
+
+            let segment_ram_storage_limit = env
+                .segment_limit_ram_storage
+                .unwrap_or(DEFAULT_SEGMENT_RAM_STORAGE_LIMIT);
             let bytes = binary.as_bytes()?;
             let mut exec = ExecutorImpl::from_elf(env, &bytes)?;
 
             let session = exec.run_with_callback(|segment| {
                 let segment_bytes = bincode::serialize(&segment)?;
+                let segment_storage: pb::api::AssetRequest = if segment.index
+                    < segment_ram_storage_limit
+                {
+                    pb::api::AssetRequest {
+                        kind: Some(pb::api::asset_request::Kind::Inline(())),
+                    }
+                } else {
+                    pb::api::AssetRequest {
+                        kind: Some(pb::api::asset_request::Kind::Path(segment_path.to_string())),
+                    }
+                };
+
                 let asset = pb::api::Asset::from_bytes(
-                    &segments_out,
+                    &segment_storage,
                     segment_bytes.into(),
                     format!("segment-{}", segment.index),
                 )?;

risc0/zkvm/src/host/client/env.rs

@@ -91,6 +91,7 @@ pub struct ExecutorEnv<'a> {
     pub(crate) segment_path: Option<SegmentPath>,
     pub(crate) pprof_out: Option<PathBuf>,
     pub(crate) input_digest: Option<Digest>,
+    pub(crate) segment_limit_ram_storage: Option<u32>,
 }
 
 impl<'a> ExecutorEnv<'a> {
@@ -361,6 +362,12 @@ impl<'a> ExecutorEnvBuilder<'a> {
         self
     }
 
+    /// Set the path where segments will be stored.
+    pub fn segment_limit_ram_storage(&mut self, limit: u32) -> &mut Self {
+        self.inner.segment_limit_ram_storage = Some(limit);
+        self
+    }
+
     /// Enable the profiler and output results to the specified path.
     pub fn enable_profiler<P: AsRef<Path>>(&mut self, path: P) -> &mut Self {
         self.inner.pprof_out = Some(path.as_ref().to_path_buf());

risc0/zkvm/src/host/protos/api.proto

@@ -140,6 +140,7 @@ message ExecutorEnv {
   string pprof_out = 10;
   repeated Assumption assumptions = 11;
   string segment_path = 12;
+  optional uint32 segment_limit_ram_storage = 13;
 }
 
 message Assumption {

risc0/zkvm/src/host/server/exec/executor.rs

@@ -20,7 +20,7 @@ use risc0_circuit_rv32im::prove::emu::{
     addr::ByteAddr,
     exec::{
         Executor, Syscall as NewSyscall, SyscallContext as NewSyscallContext,
-        DEFAULT_SEGMENT_LIMIT_PO2,
+        DEFAULT_SEGMENT_LIMIT_PO2, DEFAULT_SEGMENT_RAM_STORAGE_LIMIT,
     },
 };
 use risc0_zkp::core::digest::Digest;
@@ -29,7 +29,7 @@ use tempfile::tempdir;
 
 use crate::{
     host::client::env::SegmentPath, Assumption, Assumptions, ExecutorEnv, FileSegmentRef, Output,
-    Segment, SegmentRef, Session,
+    Segment, SegmentRef, Session, SimpleSegmentRef,
 };
 
 use super::{
@@ -107,12 +107,24 @@ impl<'a> ExecutorImpl<'a> {
     /// This will run the executor to get a [Session] which contain the results
     /// of the execution.
     pub fn run(&mut self) -> Result<Session> {
+        let ram_segment_limit: u32 = self
+            .env
+            .segment_limit_ram_storage
+            .unwrap_or(DEFAULT_SEGMENT_RAM_STORAGE_LIMIT);
+
         if self.env.segment_path.is_none() {
             self.env.segment_path = Some(SegmentPath::TempDir(Arc::new(tempdir()?)));
         }
 
         let path = self.env.segment_path.clone().unwrap();
-        self.run_with_callback(|segment| Ok(Box::new(FileSegmentRef::new(&segment, &path)?)))
+        self.run_with_callback(|segment| {
+            if segment.index < ram_segment_limit {
+                Ok(Box::new(SimpleSegmentRef::new(segment)))
+            } else {
+                // After creating more than the limit, store it in files to save RAM.
+                Ok(Box::new(FileSegmentRef::new(&segment, &path)?))
+            }
+        })
     }
 
     /// Run the executor until [crate::ExitCode::Halted] or

risc0/zkvm/src/host/server/prove/tests.rs

@@ -28,6 +28,7 @@ use test_log::test;
 
 use super::{get_prover_server, HalPair, ProverImpl};
 use crate::{
+    default_prover,
     host::server::testutils,
     serde::{from_slice, to_vec},
     ExecutorEnv, ExecutorImpl, ExitCode, ProveInfo, ProverOpts, ProverServer, Receipt, ReceiptKind,
@@ -64,6 +65,17 @@ fn prove_nothing(hashfn: &str) -> Result<ProveInfo> {
     get_prover_server(&opts).unwrap().prove(env, MULTI_TEST_ELF)
 }
 
+#[test]
+fn prove_no_ram() {
+    let env = ExecutorEnv::builder()
+        .write(&MultiTestSpec::DoNothing)
+        .unwrap()
+        .segment_limit_ram_storage(0)
+        .build()
+        .unwrap();
+    default_prover().prove(env, MULTI_TEST_ELF).unwrap();
+}
+
 #[test]
 fn prove_nothing_succinct() {
     let env = ExecutorEnv::builder()