Merge pull request #233 from mooori/emit-inst-profiling

feat(bench): add `analyze-zkasm` CLI and `instrument-inst` command

Cargo.toml

@@ -113,6 +113,7 @@ members = [
   "cranelift/isle/fuzz",
   "cranelift/isle/islec",
   "cranelift/serde",
+  "cranelift/zkasm_data/analyze-zkasm",
   "crates/bench-api",
   "crates/c-api/artifact",
   "crates/environ/fuzz",
@@ -440,4 +441,4 @@ overflow-checks = false
 incremental = false
 debug-assertions = false
 overflow-checks = false
-opt-level = 's'
+opt-level = 's'

cranelift/codegen/src/isa/zkasm/inst/emit.rs

@@ -374,7 +374,10 @@ impl Inst {
             // Labels are handled separately since benchmarking will provide a separate command to
             // analyze labels.
             &MInst::Label { .. } => {}
-            _ => put_string(";; TODO(mooori) add instruction instrumentation\n", sink),
+            _ => put_string(
+                ";; TODO(mooori) call the helper to trace `MInst` execution\n",
+                sink,
+            ),
         }
     }
 }

cranelift/filetests/src/lib.rs

@@ -16,7 +16,8 @@ mod match_directive;
 mod runner;
 mod runone;
 mod subtest;
-mod zkasm_codegen;
+// TODO(#237) Move `zkasm_codegen` now that it's used outside of filetests too.
+pub mod zkasm_codegen;
 mod zkasm_runner;
 
 mod test_alias_analysis;

cranelift/filetests/src/zkasm_codegen.rs

@@ -1,3 +1,5 @@
+//! zkASM code generation
+
 use std::collections::HashMap;
 use std::sync::Arc;
 
@@ -20,7 +22,7 @@ pub struct ZkasmSettings {
     pub emit_profiling_info: bool,
 }
 
-#[allow(dead_code)]
+/// Generates zkASM for the provided `wasm_module`.
 pub fn generate_zkasm(settings: &ZkasmSettings, wasm_module: &[u8]) -> String {
     let flag_builder = settings::builder();
     let mut isa_builder = zkasm::isa_builder("zkasm-unknown-unknown".parse().unwrap());
@@ -83,7 +85,7 @@ fn handle_zkasm_settings(
     }
 }
 
-#[allow(dead_code)]
+/// Generates a preamble.
 pub fn generate_preamble(
     start_func_index: usize,
     globals: &[(cranelift_wasm::GlobalIndex, cranelift_wasm::GlobalInit)],
@@ -257,6 +259,7 @@ fn optimize_labels(code: &[&str], func_index: usize) -> Vec<String> {
 }
 
 // TODO: fix same label names in different functions
+/// Compiles a clif function.
 pub fn compile_clif_function(func: &Function) -> Vec<String> {
     let flag_builder = settings::builder();
     let isa_builder = zkasm::isa_builder("zkasm-unknown-unknown".parse().unwrap());
@@ -289,6 +292,7 @@ pub fn compile_clif_function(func: &Function) -> Vec<String> {
 // Simple progam which don't contain globals or some other speciefic preamble\postamble
 // Program don't need helper functions (for example 2-exp.zkasm)
 // How to fix it? Use generate_preamble and provide correct inputs for it.
+/// Builds zkASM used in filetests.
 pub fn build_test_zkasm(functions: Vec<Vec<String>>, invocations: Vec<Vec<String>>) -> String {
     // TODO: use generate_preamble to get preamble
     let preamble = "\
@@ -316,6 +320,7 @@ start:
     program.join("\n")
 }
 
+/// Compiles a invocation.
 pub fn compile_invocation(
     invoke: Invocation,
     compare: Comparison,

cranelift/zkasm_data/analyze-zkasm/Cargo.toml

@@ -0,0 +1,12 @@
+[package]
+name = "analyze-zkasm"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+anyhow = { workspace = true }
+clap = { workspace = true, default-features = true, features = ["std", "derive"] }
+cranelift-filetests = { workspace = true }
+wat = { workspace = true }

cranelift/zkasm_data/analyze-zkasm/src/main.rs

@@ -0,0 +1,64 @@
+use std::path::PathBuf;
+
+use clap::{Parser, Subcommand};
+
+use cranelift_filetests::zkasm_codegen::{generate_zkasm, ZkasmSettings};
+
+/// A CLI to analyze zkasm generated by cranelift.
+#[derive(Parser)]
+#[command(version, about, long_about = None)]
+#[command(propagate_version = true)]
+struct Cli {
+    #[command(subcommand)]
+    command: Command,
+}
+
+#[derive(Subcommand)]
+enum Command {
+    /// Generate zkasm that calls zkevm-proverjs helpers to trace executed instructions.
+    #[command[arg_required_else_help = true]]
+    InstrumentInst {
+        /// Path to the input WAT file.
+        wat_path: PathBuf,
+        /// Path to the file where generated zkasm is written.
+        out_path: PathBuf,
+    },
+}
+
+fn main() -> anyhow::Result<()> {
+    let cli = Cli::parse();
+
+    match &cli.command {
+        Command::InstrumentInst { wat_path, out_path } => {
+            let zkasm = instrument_inst(wat_path)?;
+            std::fs::write(out_path, zkasm.as_bytes())?;
+            println!("wrote instrumented zkASM to {}", out_path.display());
+        }
+    }
+
+    Ok(())
+}
+
+fn instrument_inst(wat_path: &PathBuf) -> anyhow::Result<String> {
+    let wasm_module = wat::parse_file(wat_path)?;
+    let zkasm_settings = ZkasmSettings {
+        emit_profiling_info: true,
+    };
+    Ok(generate_zkasm(&zkasm_settings, &wasm_module))
+}
+
+#[cfg(test)]
+mod tests {
+    use std::path::PathBuf;
+
+    use super::instrument_inst;
+
+    #[test]
+    fn test_instrument_inst() -> anyhow::Result<()> {
+        let zkasm = instrument_inst(&PathBuf::from("./testfiles/simple.wat"))?;
+        let expected_zkasm =
+            std::fs::read_to_string("./testfiles/simple_instrumented.zkasm").unwrap();
+        assert_eq!(zkasm, expected_zkasm);
+        Ok(())
+    }
+}

cranelift/zkasm_data/analyze-zkasm/testfiles/README.md

@@ -0,0 +1,12 @@
+# About
+
+Contains files used in tests.
+
+# Updating files containing expected zkASM
+
+Use the CLI, for instance:
+
+```
+# In the analyze-zkasm directory run
+cargo run -- instrument-inst ./testfiles/simple.wat ./testfiles/simple_instrumented.zkasm
+```

cranelift/zkasm_data/analyze-zkasm/testfiles/simple.wat

@@ -0,0 +1,4 @@
+(module
+    (func $main 
+        nop)
+    (start $main))

cranelift/zkasm_data/analyze-zkasm/testfiles/simple_instrumented.zkasm

@@ -0,0 +1,21 @@
+start:
+  0xffff => SP
+  zkPC + 2 => RR
+  :JMP(function_0)
+  :JMP(finalizeExecution)
+function_0:
+  ;; TODO(mooori) call the helper to trace `MInst` execution
+  SP - 1 => SP
+  ;; TODO(mooori) call the helper to trace `MInst` execution
+  RR :MSTORE(SP)
+  ;; TODO(mooori) call the helper to trace `MInst` execution
+  ;; TODO(mooori) call the helper to trace `MInst` execution
+  $ => RR :MLOAD(SP)
+  ;; TODO(mooori) call the helper to trace `MInst` execution
+  SP + 1 => SP
+  ;; TODO(mooori) call the helper to trace `MInst` execution
+  :JMP(RR)
+finalizeExecution:
+  ${beforeLast()}  :JMPN(finalizeExecution)
+                   :JMP(start)
+INCLUDE "helpers/2-exp.zkasm"