Add support for fine-grained operator costs

crates/cranelift/src/func_environ.rs

@@ -420,23 +420,7 @@ impl<'module_environment> FuncEnvironment<'module_environment> {
             return;
         }
 
-        self.fuel_consumed += match op {
-            // Nop and drop generate no code, so don't consume fuel for them.
-            Operator::Nop | Operator::Drop => 0,
-
-            // Control flow may create branches, but is generally cheap and
-            // free, so don't consume fuel. Note the lack of `if` since some
-            // cost is incurred with the conditional check.
-            Operator::Block { .. }
-            | Operator::Loop { .. }
-            | Operator::Unreachable
-            | Operator::Return
-            | Operator::Else
-            | Operator::End => 0,
-
-            // everything else, just call it one operation.
-            _ => 1,
-        };
+        self.fuel_consumed += self.tunables.operator_cost.cost(op);
 
         match op {
             // Exiting a function (via a return or unreachable) or otherwise

crates/environ/src/tunables.rs

@@ -3,6 +3,7 @@ use crate::{IndexType, Limits, Memory, TripleExt};
 use core::{fmt, str::FromStr};
 use serde_derive::{Deserialize, Serialize};
 use target_lexicon::{PointerWidth, Triple};
+use wasmparser::Operator;
 
 macro_rules! define_tunables {
     (
@@ -46,8 +47,8 @@ macro_rules! define_tunables {
             /// Configure the `Tunables` provided.
             pub fn configure(&self, tunables: &mut Tunables) {
                 $(
-                    if let Some(val) = self.$field {
-                        tunables.$field = val;
+                    if let Some(val) = &self.$field {
+                        tunables.$field = val.clone();
                     }
                 )*
             }
@@ -87,6 +88,9 @@ define_tunables! {
         /// will be consumed every time a wasm instruction is executed.
         pub consume_fuel: bool,
 
+        /// The cost of each operator. If fuel is not enabled, this is ignored.
+        pub operator_cost: OperatorCostStrategy,
+
         /// Whether or not we use epoch-based interruption.
         pub epoch_interruption: bool,
 
@@ -193,7 +197,7 @@ impl Tunables {
     }
 
     /// Returns the default set of tunables for running under MIRI.
-    pub const fn default_miri() -> Tunables {
+    pub fn default_miri() -> Tunables {
         Tunables {
             collector: None,
 
@@ -208,6 +212,7 @@ impl Tunables {
             debug_native: false,
             parse_wasm_debuginfo: true,
             consume_fuel: false,
+            operator_cost: OperatorCostStrategy::Default,
             epoch_interruption: false,
             memory_may_move: true,
             guard_before_linear_memory: true,
@@ -229,7 +234,7 @@ impl Tunables {
     }
 
     /// Returns the default set of tunables for running under a 32-bit host.
-    pub const fn default_u32() -> Tunables {
+    pub fn default_u32() -> Tunables {
         Tunables {
             // For 32-bit we scale way down to 10MB of reserved memory. This
             // impacts performance severely but allows us to have more than a
@@ -244,7 +249,7 @@ impl Tunables {
     }
 
     /// Returns the default set of tunables for running under a 64-bit host.
-    pub const fn default_u64() -> Tunables {
+    pub fn default_u64() -> Tunables {
         Tunables {
             // 64-bit has tons of address space to static memories can have 4gb
             // address space reservations liberally by default, allowing us to
@@ -326,3 +331,134 @@ impl FromStr for IntraModuleInlining {
         }
     }
 }
+
+/// The cost of each operator.
+///
+/// Note: a more dynamic approach (e.g. a user-supplied callback) can be
+/// added as a variant in the future if needed.
+#[derive(Clone, Hash, Serialize, Deserialize, Debug, PartialEq, Eq, Default)]
+pub enum OperatorCostStrategy {
+    /// A table of operator costs.
+    Table(Box<OperatorCost>),
+
+    /// Each cost defaults to 1 fuel unit, except `Nop`, `Drop` and
+    /// a few control flow operators.
+    #[default]
+    Default,
+}
+
+impl OperatorCostStrategy {
+    /// Create a new operator cost strategy with a table of costs.
+    pub fn table(cost: OperatorCost) -> Self {
+        OperatorCostStrategy::Table(Box::new(cost))
+    }
+
+    /// Get the cost of an operator.
+    pub fn cost(&self, op: &Operator) -> i64 {
+        match self {
+            OperatorCostStrategy::Table(cost) => cost.cost(op),
+            OperatorCostStrategy::Default => default_operator_cost(op),
+        }
+    }
+}
+
+const fn default_operator_cost(op: &Operator) -> i64 {
+    match op {
+        // Nop and drop generate no code, so don't consume fuel for them.
+        Operator::Nop | Operator::Drop => 0,
+
+        // Control flow may create branches, but is generally cheap and
+        // free, so don't consume fuel. Note the lack of `if` since some
+        // cost is incurred with the conditional check.
+        Operator::Block { .. }
+        | Operator::Loop { .. }
+        | Operator::Unreachable
+        | Operator::Return
+        | Operator::Else
+        | Operator::End => 0,
+
+        // Everything else, just call it one operation.
+        _ => 1,
+    }
+}
+
+macro_rules! default_cost {
+    // Nop and drop generate no code, so don't consume fuel for them.
+    (Nop) => {
+        0
+    };
+    (Drop) => {
+        0
+    };
+
+    // Control flow may create branches, but is generally cheap and
+    // free, so don't consume fuel. Note the lack of `if` since some
+    // cost is incurred with the conditional check.
+    (Block) => {
+        0
+    };
+    (Loop) => {
+        0
+    };
+    (Unreachable) => {
+        0
+    };
+    (Return) => {
+        0
+    };
+    (Else) => {
+        0
+    };
+    (End) => {
+        0
+    };
+
+    // Everything else, just call it one operation.
+    ($op:ident) => {
+        1
+    };
+}
+
+macro_rules! define_operator_cost {
+    ($(@$proposal:ident $op:ident $({ $($arg:ident: $argty:ty),* })? => $visit:ident ($($ann:tt)*) )*) => {
+        /// The fuel cost of each operator in a table.
+        #[derive(Clone, Hash, Serialize, Deserialize, Debug, PartialEq, Eq)]
+        #[allow(missing_docs, non_snake_case, reason = "to avoid triggering clippy lints")]
+        pub struct OperatorCost {
+            $(
+                pub $op: u8,
+            )*
+        }
+
+        impl OperatorCost {
+            /// Returns the cost of the given operator.
+            pub fn cost(&self, op: &Operator) -> i64 {
+                match op {
+                    $(
+                        Operator::$op $({ $($arg: _),* })? => self.$op as i64,
+                    )*
+                    unknown => panic!("unknown op: {unknown:?}"),
+                }
+            }
+        }
+
+        impl OperatorCost {
+            /// Creates a new `OperatorCost` table with default costs for each operator.
+            pub const fn new() -> Self {
+                Self {
+                    $(
+                        $op: default_cost!($op),
+                    )*
+                }
+            }
+        }
+
+        impl Default for OperatorCost {
+            fn default() -> Self {
+                Self::new()
+            }
+        }
+    }
+}
+
+wasmparser::for_each_operator!(define_operator_cost);

crates/wasmtime/src/config.rs

@@ -7,7 +7,7 @@ use core::str::FromStr;
 #[cfg(any(feature = "cranelift", feature = "winch"))]
 use std::path::Path;
 pub use wasmparser::WasmFeatures;
-use wasmtime_environ::{ConfigTunables, TripleExt, Tunables};
+use wasmtime_environ::{ConfigTunables, OperatorCost, OperatorCostStrategy, TripleExt, Tunables};
 
 #[cfg(feature = "runtime")]
 use crate::memory::MemoryCreator;
@@ -607,6 +607,14 @@ impl Config {
         self
     }
 
+    /// Configures the fuel cost of each WebAssembly operator.
+    ///
+    /// This is only relevant when [`Config::consume_fuel`] is enabled.
+    pub fn operator_cost(&mut self, cost: OperatorCost) -> &mut Self {
+        self.tunables.operator_cost = Some(OperatorCostStrategy::table(cost));
+        self
+    }
+
     /// Enables epoch-based interruption.
     ///
     /// When executing code in async mode, we sometimes want to

crates/wasmtime/src/engine/serialization.rs

@@ -34,8 +34,7 @@ use object::{
     read::elf::{ElfFile64, FileHeader, SectionHeader},
 };
 use serde_derive::{Deserialize, Serialize};
-use wasmtime_environ::obj;
-use wasmtime_environ::{FlagValue, ObjectKind, Tunables, collections};
+use wasmtime_environ::{FlagValue, ObjectKind, OperatorCostStrategy, Tunables, collections, obj};
 
 const VERSION: u8 = 0;
 
@@ -276,6 +275,22 @@ impl Metadata<'_> {
         );
     }
 
+    fn check_cost(
+        consume_fuel: bool,
+        found: &OperatorCostStrategy,
+        expected: &OperatorCostStrategy,
+    ) -> Result<()> {
+        if !consume_fuel {
+            return Ok(());
+        }
+
+        if found != expected {
+            bail!("Module costs are incompatible");
+        }
+
+        Ok(())
+    }
+
     fn check_tunables(&mut self, other: &Tunables) -> Result<()> {
         let Tunables {
             collector,
@@ -285,6 +300,7 @@ impl Metadata<'_> {
             debug_guest,
             parse_wasm_debuginfo,
             consume_fuel,
+            ref operator_cost,
             epoch_interruption,
             memory_may_move,
             guard_before_linear_memory,
@@ -336,6 +352,7 @@ impl Metadata<'_> {
             "WebAssembly backtrace support",
         )?;
         Self::check_bool(consume_fuel, other.consume_fuel, "fuel support")?;
+        Self::check_cost(consume_fuel, operator_cost, &other.operator_cost)?;
         Self::check_bool(
             epoch_interruption,
             other.epoch_interruption,

crates/wasmtime/src/lib.rs

@@ -515,6 +515,7 @@ mod sync_nostd;
 #[cfg(not(feature = "std"))]
 use sync_nostd as sync;
 
+pub use wasmtime_environ::OperatorCost;
 #[doc(inline)]
 pub use wasmtime_environ::error;
 

crates/wasmtime/src/runtime/vm/memory/malloc.rs

@@ -157,16 +157,18 @@ mod tests {
     };
 
     // Valid tunables that can be used to create a `MallocMemory`.
-    const TUNABLES: Tunables = Tunables {
-        memory_reservation: 0,
-        memory_guard_size: 0,
-        memory_init_cow: false,
-        ..Tunables::default_miri()
-    };
+    fn tunables() -> Tunables {
+        Tunables {
+            memory_reservation: 0,
+            memory_guard_size: 0,
+            memory_init_cow: false,
+            ..Tunables::default_miri()
+        }
+    }
 
     #[test]
     fn simple() {
-        let mut memory = MallocMemory::new(&TY, &TUNABLES, 10).unwrap();
+        let mut memory = MallocMemory::new(&TY, &tunables(), 10).unwrap();
         assert_eq!(memory.storage.len(), 1);
         assert_valid(&memory);
 
@@ -191,7 +193,7 @@ mod tests {
     fn reservation_not_initialized() {
         let tunables = Tunables {
             memory_reservation_for_growth: 1 << 20,
-            ..TUNABLES
+            ..tunables()
         };
         let mut memory = MallocMemory::new(&TY, &tunables, 10).unwrap();
         assert_eq!(memory.storage.len(), 1);

tests/all/fuel.rs

@@ -382,3 +382,50 @@ fn ensure_stack_alignment(config: &mut Config) -> Result<()> {
     );
     Ok(())
 }
+
+#[wasmtime_test]
+#[cfg_attr(miri, ignore)]
+fn custom_operator_cost(config: &mut Config) -> Result<()> {
+    config.consume_fuel(true);
+    let op_cost = OperatorCost {
+        I32Const: 12,
+        I32Add: 23,
+        I64Const: 64,
+        I64Add: 128,
+        Drop: 5,
+        ..Default::default()
+    };
+    config.operator_cost(op_cost.clone());
+    let engine = Engine::new(config)?;
+    let module = Module::new(
+        &engine,
+        r#"
+            (module
+              (func (export "main")
+                ;; i32: 1 + 2
+                (drop (i32.add (i32.const 1) (i32.const 2)))
+
+                ;; i64: 3 + 4
+                (drop (i64.add (i64.const 3) (i64.const 4)))
+              )
+            )
+        "#,
+    )?;
+    let mut store = Store::new(&engine, ());
+    store.set_fuel(10_000)?;
+
+    let instance = Instance::new(&mut store, &module, &[])?;
+    let main = instance.get_typed_func::<(), ()>(&mut store, "main")?;
+
+    let initial_fuel = store.get_fuel()?;
+    main.call(&mut store, ())?;
+    let cost_of_execution = u64::from(op_cost.I32Add)
+        + u64::from(op_cost.I64Add)
+        + u64::from(op_cost.I32Const) * 2
+        + u64::from(op_cost.I64Const) * 2
+        + u64::from(op_cost.Drop) * 2
+        + 1;
+    assert_eq!(store.get_fuel()?, initial_fuel - cost_of_execution);
+
+    Ok(())
+}