Refactor platform impls a bit (#261)

Cargo.lock

@@ -10461,6 +10461,7 @@ dependencies = [
  "getrandom 0.2.16",
  "hex",
  "lazy_static",
+ "libm",
  "log",
  "p3-baby-bear 0.1.0 (git+https://github.com/brevis-network/Plonky3.git?rev=411a80d)",
  "p3-challenger 0.1.0 (git+https://github.com/brevis-network/Plonky3.git?rev=411a80d)",
@@ -16815,9 +16816,13 @@ source = "git+https://github.com/ProjectZKM/Ziren.git?tag=v1.2.3#be43a854e4be21c
 dependencies = [
  "bincode 1.3.3",
  "cfg-if",
+ "critical-section",
+ "embedded-alloc",
  "getrandom 0.2.16",
  "lazy_static",
  "libm",
+ "p3-field 0.1.0 (git+https://github.com/ProjectZKM/Plonky3)",
+ "p3-koala-bear 0.1.0 (git+https://github.com/ProjectZKM/Plonky3)",
  "rand 0.8.5",
  "serde",
  "sha2",

Cargo.toml

@@ -84,7 +84,7 @@ wait-timeout = "0.2.1"
 
 # Airbender dependencies
 airbender_execution_utils = { git = "https://github.com/matter-labs/zksync-airbender", package = "execution_utils", tag = "v0.5.2" }
-airbender_riscv_common = { git = "https://github.com/matter-labs/zksync-airbender", package = "riscv_common", tag = "v0.5.2" }
+airbender_riscv_common = { git = "https://github.com/matter-labs/zksync-airbender", package = "riscv_common", tag = "v0.5.2", default-features = false }
 
 # Jolt dependencies
 jolt-ark-serialize = { git = "https://github.com/a16z/arkworks-algebra", package = "ark-serialize", branch = "feat/fewer-reductions" }
@@ -104,7 +104,7 @@ miden-verifier = { git = "https://github.com/0xPolygonMiden/miden-vm.git", tag =
 nexus-core = { git = "https://github.com/nexus-xyz/nexus-zkvm.git", tag = "v0.3.5" }
 nexus-sdk = { git = "https://github.com/nexus-xyz/nexus-zkvm.git", tag = "v0.3.5" }
 nexus-vm = { git = "https://github.com/nexus-xyz/nexus-zkvm.git", tag = "v0.3.5" }
-nexus-rt = { git = "https://github.com/nexus-xyz/nexus-zkvm.git", tag = "v0.3.5" }
+nexus-rt = { git = "https://github.com/nexus-xyz/nexus-zkvm.git", tag = "v0.3.5", default-features = false }
 
 # OpenVM dependencies
 openvm-build = { git = "https://github.com/openvm-org/openvm.git", tag = "v1.4.2" }
@@ -118,7 +118,7 @@ openvm = { git = "https://github.com/openvm-org/openvm.git", tag = "v1.4.2", def
 # Pico dependencies
 pico-p3-field = { git = "https://github.com/brevis-network/Plonky3.git", package = "p3-field", rev = "411a80d" }
 pico-vm = { git = "https://github.com/brevis-network/pico.git", tag = "v1.2.2" }
-pico-sdk = { git = "https://github.com/brevis-network/pico", tag = "v1.2.2" }
+pico-sdk = { git = "https://github.com/brevis-network/pico", tag = "v1.2.2", default-features = false }
 
 # Risc0 dependencies
 risc0-binfmt = { version = "3.0.3", default-features = false }
@@ -129,11 +129,11 @@ risc0-zkvm-platform = { version = "2.2.1", default-features = false }
 
 # SP1 dependencies
 sp1-sdk = "5.2.4"
-sp1-zkvm = "5.2.4"
+sp1-zkvm = { version = "5.2.4", default-features = false }
 
 # Ziren dependencies
 zkm-sdk = { git = "https://github.com/ProjectZKM/Ziren.git", tag = "v1.2.3" }
-zkm-zkvm = { git = "https://github.com/ProjectZKM/Ziren.git", tag = "v1.2.3" }
+zkm-zkvm = { git = "https://github.com/ProjectZKM/Ziren.git", tag = "v1.2.3", default-features = false }
 
 # ZisK dependencies
 ziskos = { git = "https://github.com/0xPolygonHermez/zisk.git", tag = "v0.15.0" }

README.md

@@ -383,9 +383,11 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
 
 ## Environment Variables
 
-| Variable          | Description                                                                                                                             | Default |
-| ----------------- | --------------------------------------------------------------------------------------------------------------------------------------- | ------- |
-| `ERE_GPU_DEVICES` | Specifies which GPU devices to use when running Docker containers for GPU-enabled zkVMs. The value is passed to Docker's `--gpus` flag. | `all`   |
+| Variable                         | Description                                                                                                                             | Default |
+| -------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------- | ------- |
+| `ERE_IMAGE_REGISTRY`             | Specifies docker image registry of the images. When specified, it will try to pull image from the registry and possibly skip building.  | ``      |
+| `ERE_FORCE_REBUILD_DOCKER_IMAGE` | Force to rebuild docker images locally even they exist, it also prevents pulling image from registry.                                   | `false` |
+| `ERE_GPU_DEVICES`                | Specifies which GPU devices to use when running Docker containers for GPU-enabled zkVMs. The value is passed to Docker's `--gpus` flag. | `all`   |
 
 Example usage:
 

crates/dockerized/src/compiler.rs

@@ -1,8 +1,10 @@
 use crate::{
     CompilerKind,
-    image::{base_image, base_zkvm_image, compiler_zkvm_image},
+    image::{base_image, base_zkvm_image, compiler_zkvm_image, image_registry},
     util::{
-        docker::{DockerBuildCmd, DockerRunCmd, docker_image_exists, force_rebuild},
+        docker::{
+            DockerBuildCmd, DockerRunCmd, docker_image_exists, docker_pull_image, force_rebuild,
+        },
         workspace_dir,
     },
     zkVMKind,
@@ -34,9 +36,19 @@ fn build_compiler_image(zkvm_kind: zkVMKind) -> Result<(), Error> {
     let base_zkvm_image = base_zkvm_image(zkvm_kind, false);
     let compiler_zkvm_image = compiler_zkvm_image(zkvm_kind);
 
-    if !force_rebuild && docker_image_exists(&compiler_zkvm_image)? {
-        info!("Image {compiler_zkvm_image} exists, skip building");
-        return Ok(());
+    if !force_rebuild {
+        if docker_image_exists(&compiler_zkvm_image)? {
+            info!("Image {compiler_zkvm_image} exists, skip building");
+            return Ok(());
+        }
+
+        if image_registry().is_some()
+            && docker_pull_image(&compiler_zkvm_image).is_ok()
+            && docker_image_exists(&compiler_zkvm_image)?
+        {
+            info!("Image {compiler_zkvm_image} pulled, skip building");
+            return Ok(());
+        }
     }
 
     let workspace_dir = workspace_dir()?;

crates/dockerized/src/util/docker.rs

@@ -281,6 +281,25 @@ pub fn docker_container_exists(container_name: impl AsRef<str>) -> Result<bool,
     Ok(stdout.trim() == container_name.as_ref())
 }
 
+pub fn docker_pull_image(image: impl AsRef<str>) -> Result<(), CommonError> {
+    let mut cmd = Command::new("docker");
+    let output = cmd
+        .args(["image", "pull", image.as_ref()])
+        .stdout(Stdio::inherit())
+        .output()
+        .map_err(|err| CommonError::command(&cmd, err))?;
+
+    if !output.status.success() {
+        Err(CommonError::command_exit_non_zero(
+            &cmd,
+            output.status,
+            Some(&output),
+        ))?
+    }
+
+    Ok(())
+}
+
 pub fn docker_image_exists(image: impl AsRef<str>) -> Result<bool, CommonError> {
     let mut cmd = Command::new("docker");
     let output = cmd

crates/dockerized/src/zkvm.rs

@@ -1,11 +1,11 @@
 use crate::{
     compiler::SerializedProgram,
-    image::{base_image, base_zkvm_image, server_zkvm_image},
+    image::{base_image, base_zkvm_image, image_registry, server_zkvm_image},
     util::{
         cuda::cuda_arch,
         docker::{
             DockerBuildCmd, DockerRunCmd, docker_container_exists, docker_image_exists,
-            force_rebuild, stop_docker_container,
+            docker_pull_image, force_rebuild, stop_docker_container,
         },
         home_dir, workspace_dir,
     },
@@ -45,9 +45,19 @@ fn build_server_image(zkvm_kind: zkVMKind, gpu: bool) -> Result<(), Error> {
     let base_zkvm_image = base_zkvm_image(zkvm_kind, gpu);
     let server_zkvm_image = server_zkvm_image(zkvm_kind, gpu);
 
-    if !force_rebuild && docker_image_exists(&server_zkvm_image)? {
-        info!("Image {server_zkvm_image} exists, skip building");
-        return Ok(());
+    if !force_rebuild {
+        if docker_image_exists(&server_zkvm_image)? {
+            info!("Image {server_zkvm_image} exists, skip building");
+            return Ok(());
+        }
+
+        if image_registry().is_some()
+            && docker_pull_image(&server_zkvm_image).is_ok()
+            && docker_image_exists(&server_zkvm_image)?
+        {
+            info!("Image {server_zkvm_image} pulled, skip building");
+            return Ok(());
+        }
     }
 
     let workspace_dir = workspace_dir()?;

crates/zkvm/airbender/platform/src/lib.rs

@@ -30,7 +30,9 @@ impl Platform for AirbenderPlatform {
             "Maximum output size is 32 bytes, got {} bytes",
             output.len()
         );
-        let words = array::from_fn(|i| u32::from_le_bytes(array::from_fn(|j| output[4 * i + j])));
+        let words = array::from_fn(|i| {
+            u32::from_le_bytes(array::from_fn(|j| *output.get(4 * i + j).unwrap_or(&0)))
+        });
         riscv_common::zksync_os_finish_success(&words);
     }
 

crates/zkvm/pico/platform/Cargo.toml

@@ -7,13 +7,14 @@ license.workspace = true
 
 [dependencies]
 # Pico dependencies
-pico-sdk.workspace = true
+pico-sdk = { workspace = true, features = ["kb"] }
 
 # Local dependencies
 ere-platform-trait.workspace = true
 
 [features]
 default = []
+libm = ["pico-sdk/libm"]
 
 [lints]
 workspace = true

crates/zkvm/sp1/platform/Cargo.toml

@@ -7,7 +7,7 @@ license.workspace = true
 
 [dependencies]
 # SP1 dependencies
-sp1-zkvm.workspace = true
+sp1-zkvm = { workspace = true, features = ["lib"] }
 
 # Local dependencies
 ere-platform-trait.workspace = true

crates/zkvm/ziren/platform/Cargo.toml

@@ -13,7 +13,11 @@ zkm-zkvm.workspace = true
 ere-platform-trait.workspace = true
 
 [features]
-default = []
+default = ["libm", "bump"]
+bump = ["zkm-zkvm/bump"]
+embedded = ["zkm-zkvm/embedded"]
+libm = ["zkm-zkvm/libm"]
+verify = ["zkm-zkvm/verify"]
 
 [lints]
 workspace = true