Add option to enable pulley support for pre-compilation and runtime of modules and components

Justfile

@@ -43,13 +43,25 @@ build-rust-wasm-examples target=default-target features="": (mkdir-redist target
     cd ./src/rust_wasm_samples && cargo build --target wasm32-unknown-unknown --profile={{ if target == "debug" {"dev"} else { target } }}
     cargo run {{ if features =="" {''} else if features=="no-default-features" {"--no-default-features" } else {"--features " + features } }} -p hyperlight-wasm-aot compile {{ if features =~ "gdb" {"--debug"} else {""} }} ./src/rust_wasm_samples/target/wasm32-unknown-unknown/{{ target }}/rust_wasm_samples.wasm ./x64/{{ target }}/rust_wasm_samples.aot
 
+build-pulley-rust-wasm-examples target=default-target features="": (mkdir-redist target)
+    rustup target add wasm32-unknown-unknown
+    cd ./src/rust_wasm_samples && cargo build --target wasm32-unknown-unknown --profile={{ if target == "debug" {"dev"} else { target } }}
+    cargo run {{ if features =="" {''} else if features=="no-default-features" {"--no-default-features" } else {"--features " + features } }} -p hyperlight-wasm-aot compile --pulley {{ if features =~ "gdb" {"--debug"} else {""} }} ./src/rust_wasm_samples/target/wasm32-unknown-unknown/{{ target }}/rust_wasm_samples.wasm ./x64/{{ target }}/rust_wasm_samples.aot
+
 build-rust-component-examples target=default-target features="": (compile-wit)
     # use cargo component so we don't get all the wasi imports https://github.com/bytecodealliance/cargo-component?tab=readme-ov-file#relationship-with-wasm32-wasip2
     # we also explicitly target wasm32-unknown-unknown since cargo component might try to pull in wasi imports https://github.com/bytecodealliance/cargo-component/issues/290
     rustup target add wasm32-unknown-unknown
     cd ./src/component_sample && cargo component build --target wasm32-unknown-unknown --profile={{ if target == "debug" {"dev"} else { target } }}
     cargo run {{ if features =="" {''} else if features=="no-default-features" {"--no-default-features" } else {"--features " + features } }} -p hyperlight-wasm-aot compile {{ if features =~ "gdb" {"--debug"} else {""} }} --component ./src/component_sample/target/wasm32-unknown-unknown/{{ target }}/component_sample.wasm ./x64/{{ target }}/component_sample.aot
 
+build-pulley-rust-component-examples target=default-target features="": (compile-wit)
+    # use cargo component so we don't get all the wasi imports https://github.com/bytecodealliance/cargo-component?tab=readme-ov-file#relationship-with-wasm32-wasip2
+    # we also explicitly target wasm32-unknown-unknown since cargo component might try to pull in wasi imports https://github.com/bytecodealliance/cargo-component/issues/290
+    rustup target add wasm32-unknown-unknown
+    cd ./src/component_sample && cargo component build --target wasm32-unknown-unknown --profile={{ if target == "debug" {"dev"} else { target } }}
+    cargo run {{ if features =="" {''} else if features=="no-default-features" {"--no-default-features" } else {"--features " + features } }} -p hyperlight-wasm-aot compile --pulley {{ if features =~ "gdb" {"--debug"} else {""} }} --component ./src/component_sample/target/wasm32-unknown-unknown/{{ target }}/component_sample.wasm ./x64/{{ target }}/component_sample.aot
+
 check target=default-target:
     cargo check --profile={{ if target == "debug" {"dev"} else { target } }}
     cd src/rust_wasm_samples  && cargo check --profile={{ if target == "debug" {"dev"} else { target } }}
@@ -97,12 +109,18 @@ examples-ci target=default-target features="": (build-rust-wasm-examples target)
     cargo run {{ if features =="" {''} else {"--no-default-features -F " + features } }} --profile={{ if target == "debug" {"dev"} else { target } }} --example rust_wasm_examples
     cargo run {{ if features =="" {''} else {"--no-default-features -F " + features } }} --profile={{ if target == "debug" {"dev"} else { target } }} --example interruption
     cargo run {{ if features =="" {''} else {"--no-default-features -F function_call_metrics," + features } }} --profile={{ if target == "debug" {"dev"} else { target } }} --example metrics
-    cargo run {{ if features =="" {"--no-default-features --features kvm,mshv3"} else {"--no-default-features -F function_call_metrics," + features } }} --profile={{ if target == "debug" {"dev"} else { target } }} --example metrics 
+    cargo run {{ if features =="" {"--no-default-features --features kvm,mshv3"} else {"--no-default-features -F function_call_metrics," + features } }} --profile={{ if target == "debug" {"dev"} else { target } }} --example metrics
+    just examples-pulley {{ target }} {{ features }}
 
 examples-components target=default-target features="": (build-rust-component-examples target) 
     {{ wit-world }} cargo run {{ if features =="" {''} else {"--no-default-features -F kvm -F " + features } }} --profile={{ if target == "debug" {"dev"} else { target } }} --example component_example
     {{ wit-world-c }} cargo run {{ if features =="" {''} else {"--no-default-features -F kvm -F " + features } }} --profile={{ if target == "debug" {"dev"} else { target } }} --example c-component
 
+# Test a component and a module compiled with pulley
+examples-pulley target=default-target features="": (build-pulley-rust-component-examples target) (build-pulley-rust-wasm-examples target)
+    {{ wit-world }} cargo run {{ if features =="" {'-F pulley'} else {"--no-default-features -F kvm,pulley -F " + features } }} --profile={{ if target == "debug" {"dev"} else { target } }} --example component_example
+    cargo run {{ if features =="" {'-F pulley'} else {"--no-default-features -F pulley -F " + features } }} --profile={{ if target == "debug" {"dev"} else { target } }} --example rust_wasm_examples
+
 # warning, compares to and then OVERWRITES the given baseline
 bench-ci baseline target="release" features="":
     cd src/hyperlight_wasm && cargo bench --profile={{ if target == "debug" {"dev"} else { target } }} {{ if features =="" {''} else { "--features " + features } }} --bench benchmarks -- --verbose --save-baseline {{baseline}}

src/hyperlight_wasm/Cargo.toml

@@ -114,6 +114,7 @@ crashdump = ["hyperlight-host/crashdump"]
 gdb = ["hyperlight-host/gdb"]
 kvm = ["hyperlight-host/kvm"]
 mshv3 = ["hyperlight-host/mshv3"]
+pulley = []
 trace_guest = ["hyperlight-host/trace_guest"]
 
 [[bench]]

src/hyperlight_wasm/build.rs

@@ -128,6 +128,10 @@ fn build_wasm_runtime() -> PathBuf {
     if std::env::var("CARGO_FEATURE_GDB").is_ok() {
         cmd = cmd.arg("--features").arg("gdb");
     }
+    // Add --features pulley if the pulley feature is enabled
+    if std::env::var("CARGO_FEATURE_PULLEY").is_ok() {
+        cmd = cmd.arg("--features").arg("pulley");
+    }
     // Enable the "trace_guest" feature if the corresponding Cargo feature is enabled
     if std::env::var("CARGO_FEATURE_TRACE_GUEST").is_ok() {
         cmd = cmd.arg("--features").arg("trace_guest");

src/hyperlight_wasm/examples/rust_wasm_examples/main.rs

@@ -40,7 +40,7 @@ fn main() -> Result<()> {
 
         let mut proto_wasm_sandbox = SandboxBuilder::new()
             .with_guest_input_buffer_size(256 * 1024)
-            .with_guest_heap_size(768 * 1024)
+            .with_guest_heap_size(1280 * 1024)
             .build()?;
 
         proto_wasm_sandbox.register("TestHostFunc", host_func)?;

src/hyperlight_wasm_aot/Cargo.toml

@@ -12,10 +12,11 @@ Application to precompile WebAssembly binaries to for hyperlight-wasm.
 """
 
 [dependencies]
-wasmtime = { version = "36.0.5", default-features = false, features = ["cranelift", "runtime", "component-model" ] }
+wasmtime = { version = "36.0.5", default-features = false, features = ["cranelift", "pulley", "runtime", "component-model" ] }
 clap = { version = "4.5", features = ["derive"] }
 cargo_metadata = "0.23"
 cargo-util-schemas = "0.10.1"
+object = { version = "0.38.1", default-features = false, features = ["read_core", "elf"] }
 
 [features]
 gdb = ["wasmtime/debug-builtins"]

src/hyperlight_wasm_aot/src/main.rs

@@ -14,13 +14,31 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
+use std::fmt::Display;
 use std::path::Path;
 
 use cargo_metadata::{MetadataCommand, Package};
 use cargo_util_schemas::manifest::PackageName;
 use clap::{Parser, Subcommand};
+use object::read::elf::ElfFile64;
+use object::{Architecture, Endianness, FileFlags, Object};
 use wasmtime::{Config, Engine, Module, OptLevel, Precompiled};
 
+#[derive(Debug)]
+enum SupportedTarget {
+    X86_64UnknownNone,
+    WasmtimePulley64,
+}
+
+impl Display for SupportedTarget {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            SupportedTarget::X86_64UnknownNone => write!(f, "x86_64-unknown-none"),
+            SupportedTarget::WasmtimePulley64 => write!(f, "pulley64"),
+        }
+    }
+}
+
 #[derive(Parser)]
 #[command(name = "hyperlight-wasm-aot")]
 #[command(version = env!("CARGO_PKG_VERSION"))]
@@ -51,6 +69,10 @@ enum Commands {
         /// Disable address map and native unwind info for smaller binaries
         #[arg(long)]
         minimal: bool,
+
+        /// Pre-compile for the pulley64 target
+        #[arg(long)]
+        pulley: bool,
     },
 
     /// Check which Wasmtime version was used to precompile a module
@@ -74,6 +96,7 @@ fn main() {
             component,
             debug,
             minimal,
+            pulley,
         } => {
             let outfile = match output {
                 Some(s) => s,
@@ -83,15 +106,20 @@ fn main() {
                     path.to_str().unwrap().to_string()
                 }
             };
+            let target = if pulley {
+                SupportedTarget::WasmtimePulley64
+            } else {
+                SupportedTarget::X86_64UnknownNone
+            };
             if debug {
                 println!(
-                    "Aot Compiling {} to {} with debug info and optimizations off",
-                    input, outfile
+                    "Aot Compiling {} to [{}]: {} with debug info and optimizations off",
+                    input, target, outfile
                 );
             } else {
-                println!("Aot Compiling {} to {}", input, outfile);
+                println!("Aot Compiling {} to [{}]: {}", input, target, outfile);
             }
-            let config = get_config(debug, minimal);
+            let config = get_config(debug, minimal, &target);
             let engine = Engine::new(&config).unwrap();
             let bytes = std::fs::read(&input).unwrap();
             let serialized = if component {
@@ -121,7 +149,8 @@ fn main() {
             }
             // load the file into wasmtime, check that it is aot compiled and extract the version of wasmtime used to compile it from its metadata
             let bytes = std::fs::read(&file).unwrap();
-            let config = get_config(debug, false);
+            let target = get_aot_target(&bytes).unwrap();
+            let config = get_config(debug, false, &target);
             let engine = Engine::new(&config).unwrap();
             match Engine::detect_precompiled(&bytes) {
                 Some(pre_compiled) => {
@@ -132,11 +161,12 @@ fn main() {
                             // so we will try and load it and then catch the error and parse the version from the error message :-(
                             match unsafe { Module::deserialize(&engine, bytes) } {
                                 Ok(_) => println!(
-                                    "File {} was AOT compiled with wasmtime version: {}",
-                                    file, version_number
+                                    "File {} was AOT compiled to '{}' with wasmtime version: {}",
+                                    file, target, version_number
                                 ),
                                 Err(e) => {
                                     let error_message = e.to_string();
+                                    println!("{}", error_message);
                                     if !error_message.starts_with(
                                         "Module was compiled with incompatible Wasmtime version",
                                     ) {
@@ -145,8 +175,8 @@ fn main() {
                                     }
                                     let version = error_message.trim_start_matches("Module was compiled with incompatible Wasmtime version ").trim();
                                     println!(
-                                        "File {} was AOT compiled with wasmtime version: {}",
-                                        file, version
+                                        "File {} was AOT compiled to '{}' with wasmtime version: {}",
+                                        file, target, version
                                     );
                                 }
                             };
@@ -168,9 +198,18 @@ fn main() {
 }
 
 /// Returns a new `Config` for the Wasmtime engine with additional settings for AOT compilation.
-fn get_config(debug: bool, minimal: bool) -> Config {
+fn get_config(debug: bool, minimal: bool, target: &SupportedTarget) -> Config {
     let mut config = Config::new();
-    config.target("x86_64-unknown-none").unwrap();
+
+    // Compile for the pulley64 target if specified
+    match target {
+        SupportedTarget::X86_64UnknownNone => {
+            config.target("x86_64-unknown-none").unwrap();
+        }
+        SupportedTarget::WasmtimePulley64 => {
+            config.target("pulley64").unwrap();
+        }
+    }
 
     // Enable the default features for the Wasmtime engine.
     if debug {
@@ -185,3 +224,42 @@ fn get_config(debug: bool, minimal: bool) -> Config {
 
     config
 }
+
+/// Parses the AOT compiled file as an ELF file and extracts the target triple
+fn get_aot_target(bytes: &[u8]) -> Result<SupportedTarget, String> {
+    const EF_WASMTIME_PULLEY64: u32 = 1 << 3;
+    const EF_WASMTIME_PULLEY32: u32 = 1 << 2;
+
+    if let Ok(elf) = ElfFile64::<Endianness>::parse(bytes) {
+        match elf.architecture() {
+            Architecture::X86_64 => Ok(SupportedTarget::X86_64UnknownNone),
+            Architecture::Aarch64 => Err(
+                "Unsupported architecture Aarch64 in AOT compiled file, was expecting pulley64"
+                    .to_string(),
+            ),
+            Architecture::S390x => Err(
+                "Unsupported architecture S390x in AOT compiled file, was expecting pulley64"
+                    .to_string(),
+            ),
+            Architecture::Riscv64 => {
+                let e_flags = match elf.flags() {
+                    FileFlags::Elf { e_flags, .. } => e_flags,
+                    _ => return Err("Unsupported file format in AOT compiled file".to_string()),
+                };
+
+                if (e_flags & EF_WASMTIME_PULLEY64 != 0) || (e_flags & EF_WASMTIME_PULLEY32) != 0 {
+                    Ok(SupportedTarget::WasmtimePulley64)
+                } else {
+                    Err(
+                        "Unsupported Riscv64 AOT compiled file, was expecting pulley64, missing expected e_flags in elf header".to_string()
+                    )
+                }
+            }
+            other => Err(format!(
+                "Unsupported architecture {other:?} in AOT compiled file"
+            )),
+        }
+    } else {
+        Err("Failed to parse AOT compiled file as ELF".to_string())
+    }
+}

src/wasm_runtime/Cargo.toml

@@ -31,4 +31,5 @@ reqwest =  {version = "0.12", default-features = false, features = ["blocking","
 [features]
 default = []
 gdb = ["wasmtime/debug-builtins"]
+pulley = ["wasmtime/pulley"]
 trace_guest = ["hyperlight-common/trace_guest", "hyperlight-guest/trace_guest", "hyperlight-guest-bin/trace_guest"]

src/wasm_runtime/build.rs

@@ -67,5 +67,6 @@ fn main() {
 
     cfg_aliases::cfg_aliases! {
         gdb: { all(feature = "gdb", debug_assertions) },
+        pulley: { feature = "pulley" },
     }
 }

src/wasm_runtime/src/component.rs

@@ -111,6 +111,16 @@ pub extern "C" fn hyperlight_main() {
     config.with_custom_code_memory(Some(alloc::sync::Arc::new(platform::WasmtimeCodeMemory {})));
     #[cfg(gdb)]
     config.debug_info(true);
+    #[cfg(pulley)]
+    config
+        .target("pulley64")
+        .map_err(|_| {
+            HyperlightGuestError::new(
+                ErrorCode::GuestError,
+                "Failed to set wasmtime target: pulley64".to_string(),
+            )
+        })
+        .unwrap();
     let engine = Engine::new(&config).unwrap();
     let linker = Linker::new(&engine);
     *CUR_ENGINE.lock() = Some(engine);

src/wasm_runtime/src/module.rs

@@ -101,6 +101,13 @@ fn init_wasm_runtime() -> Result<Vec<u8>> {
     config.with_custom_code_memory(Some(alloc::sync::Arc::new(platform::WasmtimeCodeMemory {})));
     #[cfg(gdb)]
     config.debug_info(true);
+    #[cfg(pulley)]
+    config.target("pulley64").map_err(|_| {
+        HyperlightGuestError::new(
+            ErrorCode::GuestError,
+            "Failed to set wasmtime target: pulley64".to_string(),
+        )
+    })?;
     let engine = Engine::new(&config)?;
     let mut linker = Linker::new(&engine);
     wasip1::register_handlers(&mut linker)?;