feat(zkvms_host_io): More standard CSV metrics output format and milliseconds flag

2 files changed, 23 insertions, 21 deletions

diff --git a/zkvms_host_io/input_macros/src/lib.rs b/zkvms_host_io/input_macros/src/lib.rs
index 80a6a54..10ff43b 100644
--- a/zkvms_host_io/input_macros/src/lib.rs
+++ b/zkvms_host_io/input_macros/src/lib.rs
@@ -132,31 +132,28 @@ pub fn benchmarkable(item: TokenStream) -> TokenStream {
                  }}
              }}
 
-             let mut output = String::new();
-
              if run_info.benchmarking {{
-                 output += &format!("Total Duration: {{}}",
-                     (*ends.last().unwrap() - *starts.first().unwrap())
-                     .as_secs());
-             }}
-             if run_info.repeats > 1 {{
+                 let mut output = String::new();
+
+                 let duration = *ends.last().unwrap() - *starts.first().unwrap();
+                 let duration = if run_info.millis {{ duration.as_millis() }} else {{ duration.as_secs().into() }};
+                 output += &format!("duration,{{duration}}\n");
+
                  let durations = starts
                      .into_iter()
                      .zip(ends.into_iter())
-                     .map(|(s,e)| (e - s).as_secs())
-                     .collect::<Vec<u64>>();
-                 output += &format!(";Average: {{}}",
-                     durations.iter().sum::<u64>() / durations.len() as u64);
-             }}
-
-             if run_info.benchmarking {{
-                if let Some(file) = run_info.output_file {{
-                    let mut outfile = File::create(file).unwrap();
-                    writeln!(outfile, "{{}}", output);
-                }}
-                else {{
-                    println!("{{}}", output);
-                }}
+                     .map(|(s,e)| if run_info.millis {{ (e - s).as_millis() }} else {{ (e - s).as_secs().into() }})
+                     .collect::<Vec<u128>>();
+                 let average = durations.iter().sum::<u128>() / durations.len() as u128;
+                 output += &format!("repeats,{{}}\naverage,{{average}}\n", run_info.repeats);
+
+                 if let Some(file) = run_info.output_file {{
+                     let mut outfile = File::create(file).unwrap();
+                     write!(outfile, "{{}}", output);
+                 }}
+                 else {{
+                     print!("{{}}", output);
+                 }}
              }}
         }}
     "#).parse().unwrap()
diff --git a/zkvms_host_io/src/lib.rs b/zkvms_host_io/src/lib.rs
index 46c4740..fbdd8f9 100644
--- a/zkvms_host_io/src/lib.rs
+++ b/zkvms_host_io/src/lib.rs
@@ -26,6 +26,9 @@ struct Cli {
     repeat: Option<usize>,
 
     #[arg(short, long, requires = "benchmark")]
+    millis: bool,
+
+    #[arg(short, long, requires = "benchmark")]
     metrics_output: Option<String>,
 }
 
@@ -41,6 +44,7 @@ pub struct RunWith {
     pub run_type: RunType,
     pub benchmarking: bool,
     pub repeats: usize,
+    pub millis: bool,
     pub output_file: Option<String>,
 
     pub input: Input,
@@ -92,6 +96,7 @@ pub fn read_args() -> RunWith {
         run_type: cli.run_type,
         benchmarking: cli.benchmark,
         repeats: cli.repeat.unwrap_or(1),
+        millis: cli.millis,
         output_file: cli.metrics_output,
 
         input,