Wasm-opt integration

Cargo.toml

@@ -22,6 +22,8 @@ walrus = "0.19.0"
 candid = "0.9.0-beta.2"
 rustc-demangle = "0.1"
 thiserror = "1.0.35"
+wasm-opt = "0.112.0"
+tempfile = "3.5.0"
 
 anyhow = { version = "1.0.34", optional = true }
 clap = { version = "4.1", features = ["derive", "cargo"], optional = true }

README.md

@@ -17,7 +17,7 @@ Manage metadata in the Wasm module.
 Usage: `ic-wasm <input.wasm> [-o <output.wasm>] metadata [name] [-d <text content> | -f <file content>] [-v <public|private>]`
 
 * List current metadata sections
-``` 
+```
 $ ic-wasm input.wasm metadata
 ```
 
@@ -56,7 +56,7 @@ Usage: `ic-wasm <input.wasm> -o <output.wasm> resource --remove_cycles_transfer
 
 ### Instrument (experimental)
 
-Instrument canister method to emit execution trace to stable memory. 
+Instrument canister method to emit execution trace to stable memory.
 
 Usage: `ic-wasm <input.wasm> -o <output.wasm> instrument --trace-only func1 --trace-only func2`
 
@@ -81,6 +81,29 @@ Current limitations:
 * We cannot measure query calls.
 * No concurrent calls
 
+### Optimize
+
+Invoke wasm optimizations from [`wasm-opt`](https://github.com/WebAssembly/binaryen).
+
+The optimizer exposes different optimization levels to choose from.
+
+Performance levels (optimizes for runtime):
+- O4
+- O3 (default setting: best for minimizing cycle usage)
+- O2
+- O1
+- O0 (no optimizations)
+
+Code size levels (optimizes for binary size):
+- Oz (best for minimizing code size)
+- Os
+
+The recommended setting (O3) reduces cycle usage for Motoko programs by ~10% and  Rust programs by ~4%. The code size for both languages is reduced by ~16%.
+
+Note: The `icp` metadata sections are preserved through the optimizations.
+
+
+Usage: `ic-wasm <input.wasm> -o <output.wasm> optimize --level <optional level>`
 
 ## Library
 

src/bin/main.rs

@@ -55,6 +55,11 @@ enum SubCommand {
         #[clap(short, long)]
         trace_only: Option<Vec<String>>,
     },
+    /// Invoke wasm optimizations from wasm-opt
+    Optimize {
+        #[clap(short, long, value_parser = ["O0", "O1", "O2", "O3", "O4", "Os", "Oz"], default_value = "O3")]
+        level: String,
+    },
 }
 
 fn main() -> anyhow::Result<()> {
@@ -121,6 +126,10 @@ fn main() -> anyhow::Result<()> {
                 return Ok(());
             }
         }
+        SubCommand::Optimize { level } => {
+            use ic_wasm::optimize::optimize;
+            optimize(&mut m, keep_name_section, level)
+        }
     };
     if let Some(output) = opts.output {
         m.emit_wasm_file(output)?;

src/lib.rs

@@ -2,6 +2,7 @@ pub mod info;
 pub mod instrumentation;
 pub mod limit_resource;
 pub mod metadata;
+pub mod optimize;
 pub mod shrink;
 pub mod utils;
 

src/optimize.rs

@@ -0,0 +1,79 @@
+use crate::metadata::*;
+use crate::utils::*;
+use tempfile::NamedTempFile;
+use walrus::*;
+use wasm_opt::OptimizationOptions;
+
+pub fn optimize(m: &mut Module, keep_name_section: bool, level: &str) {
+    // recursively optimize embedded modules in Motoko actor classes
+    if is_motoko_canister(m) {
+        let data = get_motoko_wasm_data_sections(m);
+        for (id, mut module) in data.into_iter() {
+            optimize(&mut module, keep_name_section, level);
+            let blob = encode_module_as_data_section(module);
+            m.data.get_mut(id).value = blob;
+        }
+    }
+
+    // pull out a copy of the custom sections to preserve
+    let m_copy = parse_wasm(&m.emit_wasm(), keep_name_section).unwrap();
+    let mut metadata_sections = Vec::new();
+    list_metadata(&m_copy).iter().for_each(|full_name| {
+        match full_name.strip_prefix("icp:public ") {
+            Some(name) => {
+                metadata_sections.push(("public", name, get_metadata(&m_copy, name).unwrap()))
+            }
+            None => match full_name.strip_prefix("icp:private ") {
+                Some(name) => {
+                    metadata_sections.push(("private", name, get_metadata(&m_copy, name).unwrap()))
+                }
+                None => unreachable!(),
+            },
+        };
+    });
+
+    // write to temp file
+    let temp_file = NamedTempFile::new().unwrap_or_else(|e| {
+        eprintln!("unable to create temp file: {}", e);
+        std::process::exit(1);
+    });
+    m.emit_wasm_file(temp_file.path()).unwrap_or_else(|e| {
+        eprintln!("unable to write to temp file: {}", e);
+        std::process::exit(1);
+    });
+
+    // read in from temp file and optimize
+    match level {
+        "O0" => OptimizationOptions::new_opt_level_0(),
+        "O1" => OptimizationOptions::new_opt_level_1(),
+        "O2" => OptimizationOptions::new_opt_level_2(),
+        "O3" => OptimizationOptions::new_opt_level_3(),
+        "O4" => OptimizationOptions::new_opt_level_4(),
+        "Os" => OptimizationOptions::new_optimize_for_size(),
+        "Oz" => OptimizationOptions::new_optimize_for_size_aggressively(),
+        _ => unreachable!(),
+    }
+    .run(temp_file.path(), temp_file.path())
+    .unwrap_or_else(|e| {
+        eprintln!("unable to optimize wasm: {}", e);
+        std::process::exit(1);
+    });
+
+    // read optimized wasm back in from temp file
+    *m = parse_wasm_file(temp_file.path().to_path_buf(), keep_name_section).unwrap_or_else(|e| {
+        eprintln!("unable to read optimized wasm: {}", e);
+        std::process::exit(1);
+    });
+
+    // re-insert the custom sections
+    metadata_sections
+        .iter()
+        .for_each(|(visibility, name, data)| {
+            let visibility = match *visibility {
+                "public" => Kind::Public,
+                "private" => Kind::Private,
+                _ => unreachable!(),
+            };
+            add_metadata(m, visibility, name, data.to_vec());
+        });
+}

tests/deployable.ic-repl.sh

@@ -21,6 +21,12 @@ function motoko(wasm) {
   call S.inc();
   call S.get();
   assert _ == (43 : nat);
+
+  call S.inc();
+  call S.inc();
+  call S.get();
+  call S.get();
+  assert _ == (45 : nat);
   S
 };
 function rust(wasm) {
@@ -29,6 +35,12 @@ function rust(wasm) {
   call S.inc();
   call S.read();
   assert _ == (43 : nat);
+
+  call S.inc();
+  call S.inc();
+  call S.read();
+  call S.read();
+  assert _ == (45 : nat);
   S
 };
 function wat(wasm) {
@@ -37,6 +49,12 @@ function wat(wasm) {
   call S.inc();
   call S.get();
   assert _ == (43 : int64);
+
+  call S.inc();
+  call S.inc();
+  call S.get();
+  call S.get();
+  assert _ == (45 : int64);
   S
 };
 function classes(wasm) {
@@ -46,6 +64,12 @@ function classes(wasm) {
   call S.put(42, "text");
   call S.get(42);
   assert _ == opt "text";
+
+  call S.put(40, "text0");
+  call S.put(41, "text1");
+  call S.put(42, "text2");
+  call S.get(42);
+  assert _ == opt "text2";
   S
 };
 function classes_limit(wasm) {
@@ -67,26 +91,39 @@ function classes_redirect(wasm) {
 
 let S = motoko(file("ok/motoko-instrument.wasm"));
 call S.__get_cycles();
-assert _ == (7199 : int64);
+assert _ == (9003 : int64);
 let S = motoko(file("ok/motoko-gc-instrument.wasm"));
 call S.__get_cycles();
-assert _ == (177 : int64);
+assert _ == (295 : int64);
 motoko(file("ok/motoko-shrink.wasm"));
 motoko(file("ok/motoko-limit.wasm"));
 
 let S = rust(file("ok/rust-instrument.wasm"));
 call S.__get_cycles();
-assert _ == (66016 : int64);
+assert _ == (136378 : int64);
 rust(file("ok/rust-shrink.wasm"));
 rust(file("ok/rust-limit.wasm"));
 
 let S = wat(file("ok/wat-instrument.wasm"));
 call S.__get_cycles();
-assert _ == (121 : int64);
+assert _ == (189 : int64);
 wat(file("ok/wat-shrink.wasm"));
 wat(file("ok/wat-limit.wasm"));
 
 classes(file("ok/classes-shrink.wasm"));
 classes_limit(file("ok/classes-limit.wasm"));
 classes_redirect(file("ok/classes-redirect.wasm"));
 classes(file("ok/classes-nop-redirect.wasm"));
+
+motoko(file("ok/motoko-optimize.wasm"));
+rust(file("ok/rust-optimize.wasm"));
+wat(file("ok/wat-optimize.wasm"));
+classes(file("ok/classes-optimize.wasm"));
+motoko(file("ok/motoko-optimize-level-4.wasm"));
+rust(file("ok/rust-optimize-level-4.wasm"));
+wat(file("ok/wat-optimize-level-4.wasm"));
+classes(file("ok/classes-optimize-level-4.wasm"));
+motoko(file("ok/motoko-optimize-level-z.wasm"));
+rust(file("ok/rust-optimize-level-z.wasm"));
+wat(file("ok/wat-optimize-level-z.wasm"));
+classes(file("ok/classes-optimize-level-z.wasm"));