Merge pull request #2463 from OffchainLabs/stylus_target_wasm_store

stylus: add compilation target config, support multiple targets in wasm store

arbitrator/prover/src/programs/config.rs

@@ -17,7 +17,7 @@ use {
         meter::Meter, start::StartMover, MiddlewareWrapper,
     },
     std::sync::Arc,
-    wasmer::{Cranelift, CraneliftOptLevel, Engine, Store},
+    wasmer::{Cranelift, CraneliftOptLevel, Engine, Store, Target},
     wasmer_compiler_singlepass::Singlepass,
 };
 
@@ -180,17 +180,19 @@ impl CompileConfig {
     }
 
     #[cfg(feature = "native")]
-    pub fn store(&self) -> Store {
-        let mut compiler: Box<dyn wasmer::CompilerConfig> = match self.debug.cranelift {
+    pub fn engine(&self, target: Target) -> Engine {
+        use wasmer::sys::EngineBuilder;
+
+        let mut wasmer_config: Box<dyn wasmer::CompilerConfig> = match self.debug.cranelift {
             true => {
-                let mut compiler = Cranelift::new();
-                compiler.opt_level(CraneliftOptLevel::Speed);
-                Box::new(compiler)
+                let mut wasmer_config = Cranelift::new();
+                wasmer_config.opt_level(CraneliftOptLevel::Speed);
+                Box::new(wasmer_config)
             }
             false => Box::new(Singlepass::new()),
         };
-        compiler.canonicalize_nans(true);
-        compiler.enable_verifier();
+        wasmer_config.canonicalize_nans(true);
+        wasmer_config.enable_verifier();
 
         let start = MiddlewareWrapper::new(StartMover::new(self.debug.debug_info));
         let meter = MiddlewareWrapper::new(Meter::new(&self.pricing));
@@ -200,22 +202,24 @@ impl CompileConfig {
 
         // add the instrumentation in the order of application
         // note: this must be consistent with the prover
-        compiler.push_middleware(Arc::new(start));
-        compiler.push_middleware(Arc::new(meter));
-        compiler.push_middleware(Arc::new(dygas));
-        compiler.push_middleware(Arc::new(depth));
-        compiler.push_middleware(Arc::new(bound));
+        wasmer_config.push_middleware(Arc::new(start));
+        wasmer_config.push_middleware(Arc::new(meter));
+        wasmer_config.push_middleware(Arc::new(dygas));
+        wasmer_config.push_middleware(Arc::new(depth));
+        wasmer_config.push_middleware(Arc::new(bound));
 
         if self.debug.count_ops {
             let counter = Counter::new();
-            compiler.push_middleware(Arc::new(MiddlewareWrapper::new(counter)));
+            wasmer_config.push_middleware(Arc::new(MiddlewareWrapper::new(counter)));
         }
 
-        Store::new(compiler)
+        EngineBuilder::new(wasmer_config)
+            .set_target(Some(target))
+            .into()
     }
 
     #[cfg(feature = "native")]
-    pub fn engine(&self) -> Engine {
-        self.store().engine().clone()
+    pub fn store(&self, target: Target) -> Store {
+        Store::new(self.engine(target))
     }
 }

arbitrator/stylus/src/cache.rs

@@ -10,6 +10,8 @@ use prover::programs::config::CompileConfig;
 use std::{collections::HashMap, num::NonZeroUsize};
 use wasmer::{Engine, Module, Store};
 
+use crate::target_cache::target_native;
+
 lazy_static! {
     static ref INIT_CACHE: Mutex<InitCache> = Mutex::new(InitCache::new(256));
 }
@@ -120,7 +122,7 @@ impl InitCache {
         }
         drop(cache);
 
-        let engine = CompileConfig::version(version, debug).engine();
+        let engine = CompileConfig::version(version, debug).engine(target_native());
         let module = unsafe { Module::deserialize_unchecked(&engine, module)? };
 
         let item = CacheItem::new(module, engine);

arbitrator/stylus/src/lib.rs

@@ -18,6 +18,7 @@ use native::NativeInstance;
 use prover::programs::{prelude::*, StylusData};
 use run::RunProgram;
 use std::{marker::PhantomData, mem, ptr};
+use target_cache::{target_cache_get, target_cache_set};
 
 pub use brotli;
 pub use prover;
@@ -29,6 +30,7 @@ pub mod run;
 
 mod cache;
 mod evm_api;
+mod target_cache;
 mod util;
 
 #[cfg(test)]
@@ -122,9 +124,9 @@ impl RustBytes {
     }
 }
 
-/// Instruments and "activates" a user wasm.
+/// "activates" a user wasm.
 ///
-/// The `output` is either the serialized asm & module pair or an error string.
+/// The `output` is either the module or an error string.
 /// Returns consensus info such as the module hash and footprint on success.
 ///
 /// Note that this operation costs gas and is limited by the amount supplied via the `gas` pointer.
@@ -140,7 +142,6 @@ pub unsafe extern "C" fn stylus_activate(
     version: u16,
     debug: bool,
     output: *mut RustBytes,
-    asm_len: *mut usize,
     codehash: *const Bytes32,
     module_hash: *mut Bytes32,
     stylus_data: *mut StylusData,
@@ -152,18 +153,97 @@ pub unsafe extern "C" fn stylus_activate(
     let codehash = &*codehash;
     let gas = &mut *gas;
 
-    let (asm, module, info) =
-        match native::activate(wasm, codehash, version, page_limit, debug, gas) {
-            Ok(val) => val,
-            Err(err) => return output.write_err(err),
-        };
-    *asm_len = asm.len();
+    let (module, info) = match native::activate(wasm, codehash, version, page_limit, debug, gas) {
+        Ok(val) => val,
+        Err(err) => return output.write_err(err),
+    };
+
     *module_hash = module.hash();
     *stylus_data = info;
 
-    let mut data = asm;
-    data.extend(&*module.into_bytes());
-    output.write(data);
+    output.write(module.into_bytes());
+    UserOutcomeKind::Success
+}
+
+/// "compiles" a user wasm.
+///
+/// The `output` is either the asm or an error string.
+/// Returns consensus info such as the module hash and footprint on success.
+///
+/// # Safety
+///
+/// `output` must not be null.
+#[no_mangle]
+pub unsafe extern "C" fn stylus_compile(
+    wasm: GoSliceData,
+    version: u16,
+    debug: bool,
+    name: GoSliceData,
+    output: *mut RustBytes,
+) -> UserOutcomeKind {
+    let wasm = wasm.slice();
+    let output = &mut *output;
+    let name = match String::from_utf8(name.slice().to_vec()) {
+        Ok(val) => val,
+        Err(err) => return output.write_err(err.into()),
+    };
+    let target = match target_cache_get(&name) {
+        Ok(val) => val,
+        Err(err) => return output.write_err(err),
+    };
+
+    let asm = match native::compile(wasm, version, debug, target) {
+        Ok(val) => val,
+        Err(err) => return output.write_err(err),
+    };
+
+    output.write(asm);
+    UserOutcomeKind::Success
+}
+
+#[no_mangle]
+/// # Safety
+///
+/// `output` must not be null.
+pub unsafe extern "C" fn wat_to_wasm(wat: GoSliceData, output: *mut RustBytes) -> UserOutcomeKind {
+    let output = &mut *output;
+    let wasm = match wasmer::wat2wasm(wat.slice()) {
+        Ok(val) => val,
+        Err(err) => return output.write_err(err.into()),
+    };
+    output.write(wasm.into_owned());
+    UserOutcomeKind::Success
+}
+
+/// sets target index to a string
+///
+/// String format is: Triple+CpuFeature+CpuFeature..
+///
+/// # Safety
+///
+/// `output` must not be null.
+#[no_mangle]
+pub unsafe extern "C" fn stylus_target_set(
+    name: GoSliceData,
+    description: GoSliceData,
+    output: *mut RustBytes,
+    native: bool,
+) -> UserOutcomeKind {
+    let output = &mut *output;
+    let name = match String::from_utf8(name.slice().to_vec()) {
+        Ok(val) => val,
+        Err(err) => return output.write_err(err.into()),
+    };
+
+    let desc_str = match String::from_utf8(description.slice().to_vec()) {
+        Ok(val) => val,
+        Err(err) => return output.write_err(err.into()),
+    };
+
+    if let Err(err) = target_cache_set(name, desc_str, native) {
+        return output.write_err(err);
+    };
+
     UserOutcomeKind::Success
 }
 

arbitrator/stylus/src/native.rs

@@ -4,7 +4,7 @@
 use crate::{
     cache::InitCache,
     env::{MeterData, WasmEnv},
-    host, util,
+    host,
 };
 use arbutil::{
     evm::{
@@ -33,11 +33,13 @@ use std::{
     ops::{Deref, DerefMut},
 };
 use wasmer::{
-    imports, AsStoreMut, Function, FunctionEnv, Instance, Memory, Module, Pages, Store,
+    imports, AsStoreMut, Function, FunctionEnv, Instance, Memory, Module, Pages, Store, Target,
     TypedFunction, Value, WasmTypeList,
 };
 use wasmer_vm::VMExtern;
 
+use crate::target_cache::target_native;
+
 #[derive(Debug)]
 pub struct NativeInstance<D: DataReader, E: EvmApi<D>> {
     pub instance: Instance,
@@ -98,7 +100,7 @@ impl<D: DataReader, E: EvmApi<D>> NativeInstance<D, E> {
         evm_data: EvmData,
     ) -> Result<Self> {
         let env = WasmEnv::new(compile, None, evm, evm_data);
-        let store = env.compile.store();
+        let store = env.compile.store(target_native());
         let module = unsafe { Module::deserialize_unchecked(&store, module)? };
         Self::from_module(module, store, env)
     }
@@ -137,9 +139,10 @@ impl<D: DataReader, E: EvmApi<D>> NativeInstance<D, E> {
         evm_data: EvmData,
         compile: &CompileConfig,
         config: StylusConfig,
+        target: Target,
     ) -> Result<Self> {
         let env = WasmEnv::new(compile.clone(), Some(config), evm_api, evm_data);
-        let store = env.compile.store();
+        let store = env.compile.store(target);
         let wat_or_wasm = std::fs::read(path)?;
         let module = Module::new(&store, wat_or_wasm)?;
         Self::from_module(module, store, env)
@@ -347,8 +350,8 @@ impl<D: DataReader, E: EvmApi<D>> StartlessMachine for NativeInstance<D, E> {
     }
 }
 
-pub fn module(wasm: &[u8], compile: CompileConfig) -> Result<Vec<u8>> {
-    let mut store = compile.store();
+pub fn module(wasm: &[u8], compile: CompileConfig, target: Target) -> Result<Vec<u8>> {
+    let mut store = compile.store(target);
     let module = Module::new(&store, wasm)?;
     macro_rules! stub {
         (u8 <- $($types:tt)+) => {
@@ -428,7 +431,6 @@ pub fn module(wasm: &[u8], compile: CompileConfig) -> Result<Vec<u8>> {
         imports.define("console", "tee_f64", stub!(f64 <- |_: f64|));
         imports.define("debug", "null_host", stub!(||));
     }
-    Instance::new(&mut store, &module, &imports)?;
 
     let module = module.serialize()?;
     Ok(module.to_vec())
@@ -441,14 +443,14 @@ pub fn activate(
     page_limit: u16,
     debug: bool,
     gas: &mut u64,
-) -> Result<(Vec<u8>, ProverModule, StylusData)> {
-    let compile = CompileConfig::version(version, debug);
+) -> Result<(ProverModule, StylusData)> {
     let (module, stylus_data) =
         ProverModule::activate(wasm, codehash, version, page_limit, debug, gas)?;
 
-    let asm = match self::module(wasm, compile) {
-        Ok(asm) => asm,
-        Err(err) => util::panic_with_wasm(wasm, err),
-    };
-    Ok((asm, module, stylus_data))
+    Ok((module, stylus_data))
+}
+
+pub fn compile(wasm: &[u8], version: u16, debug: bool, target: Target) -> Result<Vec<u8>> {
+    let compile = CompileConfig::version(version, debug);
+    self::module(wasm, compile, target)
 }

arbitrator/stylus/src/target_cache.rs

@@ -0,0 +1,81 @@
+// Copyright 2022-2024, Offchain Labs, Inc.
+// For license information, see https://github.com/OffchainLabs/nitro/blob/master/LICENSE
+
+use eyre::{eyre, OptionExt, Result};
+use lazy_static::lazy_static;
+use parking_lot::RwLock;
+use std::{collections::HashMap, str::FromStr};
+use wasmer_types::{CpuFeature, Target, Triple};
+
+lazy_static! {
+    static ref TARGET_CACHE: RwLock<HashMap<String, Target>> = RwLock::new(HashMap::new());
+    static ref TARGET_NATIVE: RwLock<Target> = RwLock::new(Target::default());
+}
+
+fn target_from_string(input: String) -> Result<Target> {
+    if input.is_empty() {
+        return Ok(Target::default());
+    }
+    let mut parts = input.split('+');
+
+    let Some(triple_string) = parts.next() else {
+        return Err(eyre!("no architecture"));
+    };
+
+    let triple = match Triple::from_str(triple_string) {
+        Ok(val) => val,
+        Err(e) => return Err(eyre!(e)),
+    };
+
+    let mut features = CpuFeature::set();
+    for flag in parts {
+        features.insert(CpuFeature::from_str(flag)?);
+    }
+
+    Ok(Target::new(triple, features))
+}
+
+/// Populates `TARGET_CACHE` inserting target specified by `description` under `name` key.
+/// Additionally, if `native` is set it sets `TARGET_NATIVE` to the specified target.
+pub fn target_cache_set(name: String, description: String, native: bool) -> Result<()> {
+    let target = target_from_string(description)?;
+
+    if native {
+        if !target.is_native() {
+            return Err(eyre!("arch not native"));
+        }
+        let flags_not_supported = Target::default()
+            .cpu_features()
+            .complement()
+            .intersection(*target.cpu_features());
+        if !flags_not_supported.is_empty() {
+            let mut err_message = String::new();
+            err_message.push_str("cpu flags not supported on local cpu for: ");
+            for item in flags_not_supported.iter() {
+                err_message.push('+');
+                err_message.push_str(&item.to_string());
+            }
+            return Err(eyre!(err_message));
+        }
+        *TARGET_NATIVE.write() = target.clone();
+    }
+
+    TARGET_CACHE.write().insert(name, target);
+
+    Ok(())
+}
+
+pub fn target_native() -> Target {
+    TARGET_NATIVE.read().clone()
+}
+
+pub fn target_cache_get(name: &str) -> Result<Target> {
+    if name.is_empty() {
+        return Ok(TARGET_NATIVE.read().clone());
+    }
+    TARGET_CACHE
+        .read()
+        .get(name)
+        .cloned()
+        .ok_or_eyre("arch not set")
+}