Add option to build and run pulley pre-compiled modules or components

Signed-off-by: Doru Blânzeanu <dblnz@pm.me>

Author: dblnz

Cargo.lock

@@ -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/Cargo.toml

@@ -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/build.rs

@@ -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/Cargo.toml

@@ -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,17 @@ 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 = match get_aot_target(&bytes) {
+                Ok(target) => target,
+                Err(e) => {
+                    eprintln!(
+                        "Error - {} is not a valid precompiled Wasmtime module: {}",
+                        file, e
+                    );
+                    std::process::exit(1)
+                }
+            };
+            let config = get_config(debug, false, &target);
             let engine = Engine::new(&config).unwrap();
             match Engine::detect_precompiled(&bytes) {
                 Some(pre_compiled) => {
@@ -132,8 +170,8 @@ 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();
@@ -145,8 +183,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 +206,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 +232,49 @@ fn get_config(debug: bool, minimal: bool) -> Config {
 
     config
 }
+
+/// Parses the AOT compiled file as an ELF file and extracts the target triple
+/// NOTE: These flag bits must match Wasmtime's EF_WASMTIME_PULLEY{64,32} values
+/// used when emitting RISC-V ELF object files. If Wasmtime changes these values,
+/// this code must be updated accordingly to correctly detect pulley targets.
+/// Source of definitions:
+/// https://github.com/bytecodealliance/wasmtime/blob/release-42.0.0/crates/environ/src/obj.rs#L26
+/// Source of logic for detecting pulley targets:
+/// https://github.com/bytecodealliance/wasmtime/blob/release-42.0.0/src/commands/objdump.rs#L408
+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".to_string())
+            }
+            Architecture::S390x => {
+                Err("Unsupported architecture S390x in AOT compiled file".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 {
+                    Ok(SupportedTarget::WasmtimePulley64)
+                } else if e_flags & EF_WASMTIME_PULLEY32 != 0 {
+                    Err("Unsupported Riscv64 AOT compiled file: pulley32 artifacts are not supported".to_string())
+                } else {
+                    Err(
+                        "Unsupported Riscv64 AOT compiled file, 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/hyperlight_wasm_aot/src/main.rs

@@ -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/Cargo.toml

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

src/wasm_runtime/build.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/component.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)?;