public io size finalize in keygen time

ceno_cli/src/commands/common_args/ceno.rs

@@ -56,6 +56,10 @@ pub struct CenoOptions {
     #[arg(long, value_parser, num_args = 1.., value_delimiter = ',')]
     public_io: Option<Vec<Word>>,
 
+    /// pub io size in byte
+    #[arg(long, default_value = "1k", value_parser = parse_size)]
+    public_io_size: u32,
+
     /// The preset configuration to use.
     #[arg(short, long, value_enum, default_value_t = SecurityLevel::default())]
     security_level: SecurityLevel,
@@ -295,7 +299,7 @@ impl CenoOptions {
                     self,
                     compilation_options,
                     elf_path,
-                    Checkpoint::PrepVerify, // FIXME: when whir and babybear is ready
+                    Checkpoint::Complete,
                 )
             }
             (PcsKind::Whir, FieldType::Goldilocks) => {
@@ -337,26 +341,29 @@ fn run_elf_inner<
     let public_io = options
         .read_public_io()
         .context("failed to read public io")?;
-    // estimate required pub io size, which is required in platform/key setup phase
-    let pub_io_size: u32 = ((public_io.len() * WORD_SIZE) as u32)
-        .next_power_of_two()
-        .max(16);
+    let public_io_size = options.public_io_size;
+    assert!(
+        public_io.len() <= public_io_size as usize / WORD_SIZE,
+        "require pub io length {} < max public_io_size {}",
+        public_io.len(),
+        public_io_size as usize / WORD_SIZE
+    );
 
     let platform = if compilation_options.release {
         setup_platform(
             options.platform,
             &program,
             options.stack_size(),
             options.heap_size(),
-            pub_io_size,
+            public_io_size,
         )
     } else {
         setup_platform_debug(
             options.platform,
             &program,
             options.stack_size(),
             options.heap_size(),
-            pub_io_size,
+            public_io_size,
         )
     };
     tracing::info!("Running on platform {:?} {}", options.platform, platform);

ceno_zkvm/src/bin/e2e.rs

@@ -101,6 +101,10 @@ struct Args {
     #[arg(long, value_parser, num_args = 1.., value_delimiter = ',')]
     public_io: Option<Vec<Word>>,
 
+    /// pub io size in byte
+    #[arg(long, default_value = "1k", value_parser = parse_size)]
+    public_io_size: u32,
+
     /// The security level to use.
     #[arg(short, long, value_enum, default_value_t = SecurityLevel::default())]
     security_level: SecurityLevel,
@@ -168,11 +172,12 @@ fn main() {
             }
         })
         .unwrap_or_default();
-
-    // estimate required pub io size, which is required in platform/key setup phase
-    let pub_io_size: u32 = ((public_io.len() * WORD_SIZE) as u32)
-        .next_power_of_two()
-        .max(16);
+    assert!(
+        public_io.len() <= args.public_io_size as usize / WORD_SIZE,
+        "require pub io length {} < max public_io_size {}",
+        public_io.len(),
+        args.public_io_size as usize / WORD_SIZE
+    );
 
     tracing::info!("Loading ELF file: {}", args.elf.display());
     let elf_bytes = fs::read(&args.elf).expect("read elf file");
@@ -183,15 +188,15 @@ fn main() {
             &program,
             args.stack_size,
             args.heap_size,
-            pub_io_size,
+            args.public_io_size,
         )
     } else {
         setup_platform(
             args.platform,
             &program,
             args.stack_size,
             args.heap_size,
-            pub_io_size,
+            args.public_io_size,
         )
     };
     tracing::info!("Running on platform {:?} {}", args.platform, platform);

ceno_zkvm/src/e2e.rs

@@ -349,13 +349,17 @@ fn setup_platform_inner(
         heap.start..heap_end as u32
     };
 
+    assert!(
+        pub_io_size.is_power_of_two(),
+        "pub io size {pub_io_size} must be a power of two"
+    );
     let platform = Platform {
         rom: program.base_address
             ..program.base_address + (program.instructions.len() * WORD_SIZE) as u32,
         prog_data,
         stack,
         heap,
-        public_io: preset.public_io.start..preset.public_io.start + pub_io_size.next_power_of_two(),
+        public_io: preset.public_io.start..preset.public_io.start + pub_io_size,
         ..preset
     };
     assert!(