Implement the branch-hinting proposal (parse + Cranelift cold blocks)

RELEASES.md

@@ -4,6 +4,12 @@ Unreleased.
 
 ### Added
 
+* Added opt-in support for the WebAssembly branch-hinting proposal: the
+  `metadata.code.branch_hint` custom section is parsed and used to mark cold
+  blocks during Cranelift compilation, behind `Config::wasm_branch_hinting`
+  (off by default).
+  [#13459](https://github.com/bytecodealliance/wasmtime/pull/13459)
+
 ### Changed
 
 --------------------------------------------------------------------------------

crates/c-api/include/wasmtime/config.h

@@ -305,6 +305,14 @@ WASMTIME_CONFIG_PROP(void, wasm_memory64, bool)
  */
 WASMTIME_CONFIG_PROP(void, wasm_wide_arithmetic, bool)
 
+/**
+ * \brief Configures whether the WebAssembly branch-hinting proposal is
+ * enabled.
+ *
+ * This setting is `false` by default.
+ */
+WASMTIME_CONFIG_PROP(void, wasm_branch_hinting, bool)
+
 #ifdef WASMTIME_FEATURE_GC
 
 /**

crates/c-api/src/config.rs

@@ -500,6 +500,11 @@ pub extern "C" fn wasmtime_config_wasm_wide_arithmetic_set(c: &mut wasm_config_t
     c.config.wasm_wide_arithmetic(enable);
 }
 
+#[unsafe(no_mangle)]
+pub extern "C" fn wasmtime_config_wasm_branch_hinting_set(c: &mut wasm_config_t, enable: bool) {
+    c.config.wasm_branch_hinting(enable);
+}
+
 #[unsafe(no_mangle)]
 #[cfg(feature = "gc")]
 pub extern "C" fn wasmtime_config_wasm_exceptions_set(c: &mut wasm_config_t, enable: bool) {

crates/cli-flags/src/lib.rs

@@ -451,6 +451,8 @@ wasmtime_option_group! {
         pub custom_page_sizes: Option<bool>,
         /// Configure support for the wide-arithmetic proposal.
         pub wide_arithmetic: Option<bool>,
+        /// Configure support for the branch-hinting proposal.
+        pub branch_hinting: Option<bool>,
         /// Configure support for the extended-const proposal.
         pub extended_const: Option<bool>,
         /// Configure support for the exceptions proposal.
@@ -1202,6 +1204,10 @@ impl CommonOptions {
         if let Some(enable) = self.wasm.wide_arithmetic.or(all) {
             config.wasm_wide_arithmetic(enable);
         }
+        // Not included in `all_proposals`: off by default until fuzzed.
+        if let Some(enable) = self.wasm.branch_hinting {
+            config.wasm_branch_hinting(enable);
+        }
         if let Some(enable) = self.wasm.extended_const.or(all) {
             config.wasm_extended_const(enable);
         }

crates/cranelift/src/compiler.rs

@@ -413,7 +413,7 @@ impl Compiler {
         context.func.signature = wasm_call_signature(&*self.isa, ty, &self.tunables);
         let (namespace, index) = key.into_raw_parts();
         context.func.name = UserFuncName::User(UserExternalName { namespace, index });
-        let mut func_env = FuncEnvironment::new(self, translation, types, ty, key);
+        let mut func_env = FuncEnvironment::new(self, translation, types, ty, key, None, 0);
         compiler
             .cx
             .func_translator
@@ -486,7 +486,20 @@ impl wasmtime_environ::Compiler for Compiler {
             context.func.collect_debug_info();
         }
 
-        let mut func_env = FuncEnvironment::new(self, translation, types, wasm_func_ty, key);
+        // Branch hints are keyed by function-body-relative offset, so the body's
+        // module-relative start is needed to convert source locations later.
+        let FunctionBodyData { validator, body } = input;
+        let func_body_offset = body.get_binary_reader().original_position();
+
+        let mut func_env = FuncEnvironment::new(
+            self,
+            translation,
+            types,
+            wasm_func_ty,
+            key,
+            Some(func_index),
+            func_body_offset,
+        );
 
         // The `stack_limit` global value below is the implementation of stack
         // overflow checks in Wasmtime.
@@ -554,7 +567,6 @@ impl wasmtime_environ::Compiler for Compiler {
                 func_env.stack_limit_at_function_entry = Some(stack_limit);
             }
         }
-        let FunctionBodyData { validator, body } = input;
         let mut validator =
             validator.into_validator(mem::take(&mut compiler.cx.validator_allocations));
         compiler.cx.func_translator.translate_body(

crates/cranelift/src/func_environ.rs

@@ -25,8 +25,11 @@ use cranelift_entity::{EntityRef, PrimaryMap, SecondaryMap};
 use cranelift_frontend::Variable;
 use cranelift_frontend::{FuncInstBuilder, FunctionBuilder};
 use smallvec::{SmallVec, smallvec};
+use std::iter::Peekable;
 use std::mem;
-use wasmparser::{FuncValidator, Operator, WasmFeatures, WasmModuleResources};
+use wasmparser::{
+    BranchHint, FuncValidator, Operator, SectionLimitedIntoIter, WasmFeatures, WasmModuleResources,
+};
 use wasmtime_core::math::f64_cvt_to_int_bounds;
 use wasmtime_environ::{
     BuiltinFunctionIndex, ComponentPC, ConstExpr, ConstOp, DataIndex, DefinedFuncIndex,
@@ -235,6 +238,13 @@ pub struct FuncEnvironment<'module_environment> {
     /// to e.g. record the return-address of a callsite for debuginfo.
     pub(crate) next_srcloc: ir::SourceLoc,
 
+    /// Lazily-decoded branch hints for the current function, in ascending
+    /// `func_offset` order (as the proposal requires). `None` when the function
+    /// carries no hints.
+    branch_hints: Option<Peekable<SectionLimitedIntoIter<'module_environment, BranchHint>>>,
+    /// Module-relative byte offset of the current function body's start.
+    func_body_offset: usize,
+
     /// Cached alias regions for alias analysis.
     heap_alias_region: Option<ir::AliasRegion>,
     table_alias_region: Option<ir::AliasRegion>,
@@ -252,10 +262,18 @@ impl<'module_environment> FuncEnvironment<'module_environment> {
         types: &'module_environment ModuleTypesBuilder,
         wasm_func_ty: &'module_environment WasmFuncType,
         key: FuncKey,
+        func_index: Option<FuncIndex>,
+        func_body_offset: usize,
     ) -> Self {
         let tunables = compiler.tunables();
         let builtin_functions = BuiltinFunctions::new(compiler);
 
+        // Synthesized functions without a wasm body (e.g. module startup) pass
+        // `None`, yielding no hints.
+        let branch_hints = func_index
+            .and_then(|func_index| translation.branch_hints(func_index))
+            .map(|reader| reader.into_iter().peekable());
+
         // This isn't used during translation, so squash the warning about this
         // being unused from the compiler.
         let _ = BuiltinFunctions::raise;
@@ -305,12 +323,40 @@ impl<'module_environment> FuncEnvironment<'module_environment> {
             state_slot: None,
             next_srcloc: ir::SourceLoc::default(),
             wasm_module_offset: translation.wasm_module_offset,
+
+            branch_hints,
+            func_body_offset,
             heap_alias_region: None,
             table_alias_region: None,
             vmctx_alias_region: None,
         }
     }
 
+    /// Consume the branch hint for the instruction at module-relative `offset`
+    /// (i.e. `builder.srcloc().bits()`), if any. The lazy decoder only moves
+    /// forward, making this O(n) over a function body.
+    pub(crate) fn take_branch_hint(&mut self, offset: usize) -> Option<BranchHint> {
+        // Fast path: no hints (always so when the proposal is off), and this
+        // runs for every `if`/`br_if`.
+        let hints = self.branch_hints.as_mut()?;
+        let rel = u32::try_from(offset.checked_sub(self.func_body_offset)?).ok()?;
+        loop {
+            // Hint bytes were validated when the section was decoded, so an error
+            // here is unexpected; treat it like exhaustion (end of hints).
+            let hint = *hints.peek()?.as_ref().ok()?;
+            if hint.func_offset < rel {
+                hints.next();
+                continue;
+            }
+            if hint.func_offset == rel {
+                hints.next();
+                return Some(hint);
+            }
+            // The next hint is for a later offset; nothing for this branch.
+            return None;
+        }
+    }
+
     pub(crate) fn pointer_type(&self) -> ir::Type {
         self.isa.pointer_type()
     }

crates/cranelift/src/translate/code_translator.rs

@@ -286,6 +286,9 @@ pub fn translate_operator(
             environ.translate_loop_header(builder)?;
         }
         Operator::If { blockty } => {
+            // Read the hint before `environ` is borrowed mutably below.
+            let branch_hint = environ.take_branch_hint(builder.srcloc().bits() as usize);
+
             let val = environ.stacks.pop1();
 
             let next_block = builder.create_block();
@@ -330,6 +333,24 @@ pub fn translate_operator(
                 (destination, ElseData::WithElse { else_block })
             };
 
+            // Mark the unlikely successor cold per the branch hint. A likely
+            // condition makes the else block cold; when it is allocated lazily
+            // (`NoElse`) defer that to `Operator::Else` via `else_is_cold`.
+            let else_is_cold = match branch_hint {
+                Some(hint) if !hint.taken => {
+                    builder.set_cold_block(next_block);
+                    false
+                }
+                Some(_) => match &else_data {
+                    ElseData::WithElse { else_block } => {
+                        builder.set_cold_block(*else_block);
+                        false
+                    }
+                    ElseData::NoElse { .. } => true,
+                },
+                None => false,
+            };
+
             builder.seal_block(next_block); // Only predecessor is the current block.
             builder.switch_to_block(next_block);
 
@@ -345,6 +366,7 @@ pub fn translate_operator(
                 params.len(),
                 results.len(),
                 *blockty,
+                else_is_cold,
             );
         }
         Operator::Else => {
@@ -358,6 +380,7 @@ pub fn translate_operator(
                     num_return_values,
                     blocktype,
                     destination,
+                    else_is_cold,
                     ..
                 } => {
                     // We finished the consequent, so record its final
@@ -407,6 +430,12 @@ pub fn translate_operator(
                             }
                         };
 
+                        // Apply a likely-taken hint deferred from a lazily
+                        // allocated else block (see `Operator::If`).
+                        if else_is_cold {
+                            builder.set_cold_block(else_block);
+                        }
+
                         // You might be expecting that we push the parameters for this
                         // `else` block here, something like this:
                         //
@@ -3339,6 +3368,7 @@ fn translate_unreachable_operator(
                 0,
                 0,
                 blockty,
+                false,
             );
         }
         Operator::Loop { blockty: _ }
@@ -3984,9 +4014,22 @@ fn translate_br_if(
     builder: &mut FunctionBuilder,
     env: &mut FuncEnvironment<'_>,
 ) {
+    // Read the hint before `env` is borrowed mutably below.
+    let branch_hint = env.take_branch_hint(builder.srcloc().bits() as usize);
+
     let val = env.stacks.pop1();
     let (br_destination, inputs) = translate_br_if_args(relative_depth, env);
     let next_block = builder.create_block();
+
+    if let Some(hint) = branch_hint {
+        // Likely taken => the fallthrough is cold, else the branch target is.
+        builder.set_cold_block(if hint.taken {
+            next_block
+        } else {
+            br_destination
+        });
+    }
+
     canonicalise_brif(builder, val, br_destination, inputs, next_block, &[]);
 
     builder.seal_block(next_block); // The only predecessor is the current block.

crates/cranelift/src/translate/stack.rs

@@ -58,6 +58,9 @@ pub enum ControlStackFrame {
         original_stack_size: usize,
         exit_is_branched_to: bool,
         blocktype: wasmparser::BlockType,
+        /// Whether the `else` block, when allocated lazily, should be marked
+        /// cold (from a branch hint that the condition is likely true).
+        else_is_cold: bool,
         /// Was the head of the `if` reachable?
         head_is_reachable: bool,
         /// What was the reachability at the end of the consequent?
@@ -507,6 +510,7 @@ impl FuncTranslationStacks {
         num_param_types: usize,
         num_result_types: usize,
         blocktype: wasmparser::BlockType,
+        else_is_cold: bool,
     ) {
         debug_assert!(num_param_types <= self.stack.len());
         self.assert_debug_stack_is_synced();
@@ -541,6 +545,7 @@ impl FuncTranslationStacks {
             head_is_reachable: self.reachable,
             consequent_ends_reachable: None,
             blocktype,
+            else_is_cold,
         });
     }
 

crates/environ/src/compile/module_environ.rs

@@ -116,6 +116,11 @@ pub struct ModuleTranslation<'data> {
     /// validation process.
     types: Option<Types>,
 
+    /// Per-function [`BranchHintReader`]s from the `metadata.code.branch_hint`
+    /// section, keyed by function index. Populated only when
+    /// [`Tunables::branch_hinting`] is enabled.
+    branch_hints: HashMap<FuncIndex, BranchHintReader<'data>>,
+
     /// The WebAssembly `start` function, if defined.
     pub start_func: Option<FuncIndex>,
 
@@ -173,6 +178,11 @@ pub enum MemorySegmentOffset {
     Static(u64),
 }
 
+/// Lazy decoder over the branch hints attached to a single function in the
+/// `metadata.code.branch_hint` custom section
+/// ([branch-hinting proposal](https://github.com/WebAssembly/branch-hinting)).
+pub type BranchHintReader<'a> = wasmparser::SectionLimited<'a, wasmparser::BranchHint>;
+
 impl<'data> ModuleTranslation<'data> {
     /// Create a new translation for the module with the given index.
     pub fn new(module_index: StaticModuleIndex) -> Self {
@@ -191,6 +201,7 @@ impl<'data> ModuleTranslation<'data> {
             types: None,
             runtime_data_map: Default::default(),
             passive_elem_map: Default::default(),
+            branch_hints: HashMap::default(),
             start_func: None,
             global_initializers: Vec::new(),
             passive_elements: Default::default(),
@@ -200,6 +211,11 @@ impl<'data> ModuleTranslation<'data> {
         }
     }
 
+    /// Returns the [`BranchHintReader`] for `func`, if the section attached any.
+    pub fn branch_hints(&self, func: FuncIndex) -> Option<BranchHintReader<'data>> {
+        self.branch_hints.get(&func).cloned()
+    }
+
     /// Returns a reference to the type information of the current module.
     pub fn get_types(&self) -> &Types {
         self.types
@@ -802,6 +818,26 @@ and for re-adding support for interface types you can see this issue:
                     log::warn!("failed to parse name section {e:?}");
                 }
             }
+            KnownCustom::BranchHints(reader) if self.tunables.branch_hinting => {
+                // Branch hints are advisory and this section is never validated;
+                // it is decoded lazily during compilation, so record only the
+                // per-function sub-readers here. Discard the whole section if any
+                // entry is malformed rather than applying it partially.
+                let mut hints = HashMap::new();
+                let result: wasmparser::Result<()> = reader.into_iter().try_for_each(|func| {
+                    let func = func?;
+                    // A well-formed section lists each function at most once; keep
+                    // the first entry deterministically if it repeats.
+                    hints
+                        .entry(FuncIndex::from_u32(func.func))
+                        .or_insert(func.hints);
+                    Ok(())
+                });
+                match result {
+                    Ok(()) => self.result.branch_hints = hints,
+                    Err(e) => log::warn!("failed to parse branch-hint section {e:?}"),
+                }
+            }
             _ => {
                 let name = section.name().trim_end_matches(".dwo");
                 if name.starts_with(".debug_") {

crates/environ/src/tunables.rs

@@ -188,6 +188,10 @@ define_tunables! {
         /// Boolean to track whether compiled code retains metadata necessary to
         /// report extra information on gc heap corruption being detected.
         pub metadata_for_gc_heap_corruption: bool,
+
+        /// Whether `metadata.code.branch_hint` sections are parsed and used to
+        /// mark cold blocks during compilation.
+        pub branch_hinting: bool,
     }
 
     pub struct ConfigTunables {
@@ -273,6 +277,7 @@ impl Tunables {
             gc_heap_may_move: true,
             metadata_for_internal_asserts: false,
             metadata_for_gc_heap_corruption: true,
+            branch_hinting: false,
         }
     }