Add support for metadata.code.branch_hint section (#4460)

Add support for metadata.code.branch_hint section

README.md

@@ -47,7 +47,7 @@ WebAssembly Micro Runtime (WAMR) is a lightweight standalone WebAssembly (Wasm)
 - [128-bit SIMD](https://github.com/WebAssembly/simd), ref to [samples/workload](samples/workload)
 - [Reference Types](https://github.com/WebAssembly/reference-types), ref to [document](doc/ref_types.md) and [sample](samples/ref-types)
 - [Bulk memory operations](https://github.com/WebAssembly/bulk-memory-operations), [Shared memory](https://github.com/WebAssembly/threads/blob/main/proposals/threads/Overview.md#shared-linear-memory), [Memory64](https://github.com/WebAssembly/memory64)
-- [Tail-call](https://github.com/WebAssembly/tail-call), [Garbage Collection](https://github.com/WebAssembly/gc), [Exception Handling](https://github.com/WebAssembly/exception-handling)
+- [Tail-call](https://github.com/WebAssembly/tail-call), [Garbage Collection](https://github.com/WebAssembly/gc), [Exception Handling](https://github.com/WebAssembly/exception-handling), [Branch Hinting](https://github.com/WebAssembly/branch-hinting)
 - [Extended Constant Expressions](https://github.com/WebAssembly/extended-const)
 
 ### Supported architectures and platforms
@@ -117,4 +117,3 @@ Any contributions you make will be under the same license.
 - [WAMR Blogs](https://bytecodealliance.github.io/wamr.dev/blog/)
 - [Community news and events](https://bytecodealliance.github.io/wamr.dev/events/)
 - [WAMR TSC meetings](https://github.com/bytecodealliance/wasm-micro-runtime/wiki/TSC-meeting-notes)
-

build-scripts/config_common.cmake

@@ -739,6 +739,7 @@ message (
 "       \"Non-trapping float-to-int Conversions\"\n"
 "       \"Sign-extension Operators\"\n"
 "       \"WebAssembly C and C++ API\"\n"
+"       \"Branch Hinting\"\n"
 "     Configurable. 0 is OFF. 1 is ON:\n"
 "       \"Bulk Memory Operation\" via WAMR_BUILD_BULK_MEMORY: ${WAMR_BUILD_BULK_MEMORY}\n"
 "       \"Extended Constant Expressions\" via WAMR_BUILD_EXTENDED_CONST_EXPR: ${WAMR_BUILD_EXTENDED_CONST_EXPR}\n"
@@ -753,7 +754,6 @@ message (
 "       \"Threads\" via WAMR_BUILD_SHARED_MEMORY: ${WAMR_BUILD_SHARED_MEMORY}\n"
 "       \"Typed Function References\" via WAMR_BUILD_GC: ${WAMR_BUILD_GC}\n"
 "     Unsupported (>= Phase4):\n"
-"       \"Branch Hinting\"\n"
 "       \"Custom Annotation Syntax in the Text Format\"\n"
 "       \"Exception Handling\"\n"
 "       \"JS String Builtins\"\n"

core/config.h

@@ -579,6 +579,10 @@ unless used elsewhere */
 #define WASM_ENABLE_REF_TYPES 0
 #endif
 
+#ifndef WASM_ENABLE_BRANCH_HINTS
+#define WASM_ENABLE_BRANCH_HINTS 0
+#endif
+
 #ifndef WASM_ENABLE_GC
 #define WASM_ENABLE_GC 0
 #endif

core/iwasm/compilation/aot.c

@@ -416,6 +416,9 @@ aot_create_funcs(const WASMModule *module, uint32 pointer_size)
         aot_func->local_types_wp = func->local_types;
         aot_func->code = func->code;
         aot_func->code_size = func->code_size;
+#if WASM_ENABLE_BRANCH_HINTS != 0
+        aot_func->code_body_begin = func->code_body_begin;
+#endif
 
         /* Resolve local offsets */
         for (j = 0; j < func_type->param_count; j++) {
@@ -872,6 +875,10 @@ aot_create_comp_data(WASMModule *module, const char *target_arch,
     comp_data->name_section_buf_end = module->name_section_buf_end;
 #endif
 
+#if WASM_ENABLE_BRANCH_HINTS != 0
+    comp_data->function_hints = module->function_hints;
+#endif
+
     aot_init_aux_data(comp_data, module);
 
     comp_data->wasm_module = module;

core/iwasm/compilation/aot.h

@@ -217,6 +217,9 @@ typedef struct AOTFunc {
     /* offset of each local, including function parameters
        and local variables */
     uint16 *local_offsets;
+#if WASM_ENABLE_BRANCH_HINTS != 0
+    uint8 *code_body_begin;
+#endif
 } AOTFunc;
 
 typedef struct AOTCompData {
@@ -296,6 +299,10 @@ typedef struct AOTCompData {
 #if WASM_ENABLE_DEBUG_AOT != 0
     dwarf_extractor_handle_t extractor;
 #endif
+
+#if WASM_ENABLE_BRANCH_HINTS != 0
+    struct WASMCompilationHint **function_hints;
+#endif
 } AOTCompData;
 
 typedef struct AOTNativeSymbol {

core/iwasm/compilation/aot_compiler.c

@@ -1158,9 +1158,7 @@ aot_compile_func(AOTCompContext *comp_ctx, uint32 func_index)
 
             case WASM_OP_BR_IF:
             {
-                read_leb_uint32(frame_ip, frame_ip_end, br_depth);
-                if (!aot_compile_op_br_if(comp_ctx, func_ctx, br_depth,
-                                          &frame_ip))
+                if (!aot_compile_op_br_if(comp_ctx, func_ctx, &frame_ip))
                     return false;
                 break;
             }

core/iwasm/compilation/aot_emit_control.c

@@ -12,6 +12,7 @@
 #endif
 #include "../aot/aot_runtime.h"
 #include "../interpreter/wasm_loader.h"
+#include "../common/wasm_loader_common.h"
 
 #if WASM_ENABLE_DEBUG_AOT != 0
 #include "debug/dwarf_extractor.h"
@@ -87,6 +88,15 @@ format_block_name(char *name, uint32 name_size, uint32 block_index,
         }                                                             \
     } while (0)
 
+#define BUILD_COND_BR_V(value_if, block_then, block_else, instr)               \
+    do {                                                                       \
+        if (!(instr = LLVMBuildCondBr(comp_ctx->builder, value_if, block_then, \
+                                      block_else))) {                          \
+            aot_set_last_error("llvm build cond br failed.");                  \
+            goto fail;                                                         \
+        }                                                                      \
+    } while (0)
+
 #define SET_BUILDER_POS(llvm_block) \
     LLVMPositionBuilderAtEnd(comp_ctx->builder, llvm_block)
 
@@ -255,6 +265,36 @@ restore_frame_sp_for_op_end(AOTBlock *block, AOTCompFrame *aot_frame)
     aot_frame->sp = block->frame_sp_begin;
 }
 
+#if WASM_ENABLE_BRANCH_HINTS != 0
+static void
+aot_emit_branch_hint(AOTCompContext *comp_ctx, AOTFuncContext *func_ctx,
+                     uint32 offset, LLVMValueRef br_if_instr)
+{
+    struct WASMCompilationHint *hint = func_ctx->function_hints;
+    while (hint != NULL) {
+        if (hint->type == WASM_COMPILATION_BRANCH_HINT
+            && ((struct WASMCompilationHintBranchHint *)hint)->offset
+                   == offset) {
+            break;
+        }
+        hint = hint->next;
+    }
+    if (hint != NULL) {
+        // same weight llvm MDBuilder::createLikelyBranchWeights assigns
+        const uint32_t likely_weight = (1U << 20) - 1;
+        const uint32_t unlikely_weight = 1;
+        aot_set_cond_br_weights(
+            comp_ctx, br_if_instr,
+            ((struct WASMCompilationHintBranchHint *)hint)->is_likely
+                ? likely_weight
+                : unlikely_weight,
+            ((struct WASMCompilationHintBranchHint *)hint)->is_likely
+                ? unlikely_weight
+                : likely_weight);
+    }
+}
+#endif
+
 static bool
 handle_next_reachable_block(AOTCompContext *comp_ctx, AOTFuncContext *func_ctx,
                             uint8 **p_frame_ip)
@@ -673,13 +713,31 @@ aot_compile_op_block(AOTCompContext *comp_ctx, AOTFuncContext *func_ctx,
                 MOVE_BLOCK_AFTER(block->llvm_else_block,
                                  block->llvm_entry_block);
                 /* Create condition br IR */
+#if WASM_ENABLE_BRANCH_HINTS != 0
+                LLVMValueRef br_if_val = NULL;
+                BUILD_COND_BR_V(value, block->llvm_entry_block,
+                                block->llvm_else_block, br_if_val);
+                const uint32 off =
+                    *p_frame_ip - func_ctx->aot_func->code_body_begin;
+                aot_emit_branch_hint(comp_ctx, func_ctx, off, br_if_val);
+#else
                 BUILD_COND_BR(value, block->llvm_entry_block,
                               block->llvm_else_block);
+#endif
             }
             else {
                 /* Create condition br IR */
+#if WASM_ENABLE_BRANCH_HINTS != 0
+                LLVMValueRef br_if_val = NULL;
+                BUILD_COND_BR_V(value, block->llvm_entry_block,
+                                block->llvm_end_block, br_if_val);
+                const uint32 off =
+                    *p_frame_ip - func_ctx->aot_func->code_body_begin;
+                aot_emit_branch_hint(comp_ctx, func_ctx, off, br_if_val);
+#else
                 BUILD_COND_BR(value, block->llvm_entry_block,
                               block->llvm_end_block);
+#endif
                 block->is_reachable = true;
             }
             if (!push_aot_block_to_stack_and_pass_params(comp_ctx, func_ctx,
@@ -1026,8 +1084,7 @@ fail:
 
 static bool
 aot_compile_conditional_br(AOTCompContext *comp_ctx, AOTFuncContext *func_ctx,
-                           uint32 br_depth, LLVMValueRef value_cmp,
-                           uint8 **p_frame_ip)
+                           LLVMValueRef value_cmp, uint8 **p_frame_ip)
 {
     AOTBlock *block_dst;
     LLVMValueRef value, *values = NULL;
@@ -1036,6 +1093,17 @@ aot_compile_conditional_br(AOTCompContext *comp_ctx, AOTFuncContext *func_ctx,
     uint32 i, param_index, result_index;
     uint64 size;
 
+    // ip is advanced by one byte for the opcode
+#if WASM_ENABLE_BRANCH_HINTS != 0
+    uint32 instr_offset =
+        (*p_frame_ip - 0x1) - (func_ctx->aot_func->code_body_begin);
+#else
+    uint32 instr_offset = 0;
+#endif
+    uint64 br_depth;
+    if (!read_leb(p_frame_ip, *p_frame_ip + 5, 32, false, &br_depth, NULL, 0))
+        return false;
+
     if (!(block_dst = get_target_block(func_ctx, br_depth))) {
         return false;
     }
@@ -1108,8 +1176,15 @@ aot_compile_conditional_br(AOTCompContext *comp_ctx, AOTFuncContext *func_ctx,
                 values = NULL;
             }
 
+#if WASM_ENABLE_BRANCH_HINTS != 0
+            LLVMValueRef br_if_val = NULL;
+            BUILD_COND_BR_V(value_cmp, block_dst->llvm_entry_block,
+                            llvm_else_block, br_if_val);
+            aot_emit_branch_hint(comp_ctx, func_ctx, instr_offset, br_if_val);
+#else
             BUILD_COND_BR(value_cmp, block_dst->llvm_entry_block,
                           llvm_else_block);
+#endif
 
             /* Move builder to else block */
             SET_BUILDER_POS(llvm_else_block);
@@ -1152,9 +1227,15 @@ aot_compile_conditional_br(AOTCompContext *comp_ctx, AOTFuncContext *func_ctx,
             }
 
             /* Condition jump to end block */
+#if WASM_ENABLE_BRANCH_HINTS != 0
+            LLVMValueRef br_if_val = NULL;
+            BUILD_COND_BR_V(value_cmp, block_dst->llvm_end_block,
+                            llvm_else_block, br_if_val);
+            aot_emit_branch_hint(comp_ctx, func_ctx, instr_offset, br_if_val);
+#else
             BUILD_COND_BR(value_cmp, block_dst->llvm_end_block,
                           llvm_else_block);
-
+#endif
             /* Move builder to else block */
             SET_BUILDER_POS(llvm_else_block);
         }
@@ -1178,13 +1259,13 @@ fail:
 
 bool
 aot_compile_op_br_if(AOTCompContext *comp_ctx, AOTFuncContext *func_ctx,
-                     uint32 br_depth, uint8 **p_frame_ip)
+                     uint8 **p_frame_ip)
 {
     LLVMValueRef value_cmp;
 
     POP_COND(value_cmp);
 
-    return aot_compile_conditional_br(comp_ctx, func_ctx, br_depth, value_cmp,
+    return aot_compile_conditional_br(comp_ctx, func_ctx, value_cmp,
                                       p_frame_ip);
 fail:
     return false;

core/iwasm/compilation/aot_emit_control.h

@@ -32,7 +32,7 @@ aot_compile_op_br(AOTCompContext *comp_ctx, AOTFuncContext *func_ctx,
 
 bool
 aot_compile_op_br_if(AOTCompContext *comp_ctx, AOTFuncContext *func_ctx,
-                     uint32 br_depth, uint8 **p_frame_ip);
+                     uint8 **p_frame_ip);
 
 bool
 aot_compile_op_br_table(AOTCompContext *comp_ctx, AOTFuncContext *func_ctx,

core/iwasm/compilation/aot_llvm.c

@@ -1963,6 +1963,13 @@ aot_create_func_context(const AOTCompData *comp_data, AOTCompContext *comp_ctx,
         goto fail;
     }
 
+#if WASM_ENABLE_BRANCH_HINTS != 0
+    func_ctx->function_hints =
+        comp_ctx->comp_data->function_hints
+            ? comp_ctx->comp_data->function_hints[func_index]
+            : NULL;
+#endif
+
     return func_ctx;
 
 fail:

core/iwasm/compilation/aot_llvm.h

@@ -270,6 +270,9 @@ typedef struct AOTFuncContext {
 #if WASM_ENABLE_DEBUG_AOT != 0
     LLVMMetadataRef debug_func;
 #endif
+#if WASM_ENABLE_BRANCH_HINTS != 0
+    struct WASMCompilationHint *function_hints;
+#endif
 
     unsigned int stack_consumption_for_func_call;
 

core/iwasm/interpreter/wasm.h

@@ -751,6 +751,10 @@ struct WASMFunction {
     void *call_to_fast_jit_from_llvm_jit;
 #endif
 #endif
+
+#if WASM_ENABLE_BRANCH_HINTS != 0
+    uint8 *code_body_begin;
+#endif
 };
 
 #if WASM_ENABLE_TAGS != 0
@@ -761,6 +765,23 @@ struct WASMTag {
 };
 #endif
 
+#if WASM_ENABLE_BRANCH_HINTS != 0
+enum WASMCompilationHintType {
+    DUMMY = 0,
+    WASM_COMPILATION_BRANCH_HINT = 0,
+};
+struct WASMCompilationHint {
+    struct WASMCompilationHint *next;
+    enum WASMCompilationHintType type;
+};
+struct WASMCompilationHintBranchHint {
+    struct WASMCompilationHint *next;
+    enum WASMCompilationHintType type;
+    uint32 offset;
+    bool is_likely;
+};
+#endif
+
 struct WASMGlobal {
     WASMGlobalType type;
 #if WASM_ENABLE_GC != 0
@@ -1049,6 +1070,10 @@ struct WASMModule {
     const uint8 *name_section_buf_end;
 #endif
 
+#if WASM_ENABLE_BRANCH_HINTS != 0
+    struct WASMCompilationHint **function_hints;
+#endif
+
 #if WASM_ENABLE_LOAD_CUSTOM_SECTION != 0
     WASMCustomSection *custom_section_list;
 #endif

core/iwasm/interpreter/wasm_loader.c

@@ -3908,6 +3908,9 @@ load_function_section(const uint8 *buf, const uint8 *buf_end,
 
             /* Resolve local set count */
             p_code_end = p_code + code_size;
+#if WASM_ENABLE_BRANCH_HINTS != 0
+            uint8 *p_body_start = (uint8 *)p_code;
+#endif
             local_count = 0;
             read_leb_uint32(p_code, buf_code_end, local_set_count);
             p_code_save = p_code;
@@ -3988,6 +3991,9 @@ load_function_section(const uint8 *buf, const uint8 *buf_end,
             if (local_count > 0)
                 func->local_types = (uint8 *)func + sizeof(WASMFunction);
             func->code_size = code_size;
+#if WASM_ENABLE_BRANCH_HINTS != 0
+            func->code_body_begin = p_body_start;
+#endif
             /*
              * we shall make a copy of code body [p_code, p_code + code_size]
              * when we are worrying about inappropriate releasing behaviour.
@@ -5560,6 +5566,88 @@ fail:
 }
 #endif
 
+#if WASM_ENABLE_BRANCH_HINTS != 0
+static bool
+handle_branch_hint_section(const uint8 *buf, const uint8 *buf_end,
+                           WASMModule *module, char *error_buf,
+                           uint32 error_buf_size)
+{
+    if (module->function_hints == NULL) {
+        module->function_hints = loader_malloc(
+            sizeof(struct WASMCompilationHint) * module->function_count,
+            error_buf, error_buf_size);
+    }
+    uint32 numFunctionHints = 0;
+    read_leb_uint32(buf, buf_end, numFunctionHints);
+    for (uint32 i = 0; i < numFunctionHints; ++i) {
+        uint32 func_idx;
+        read_leb_uint32(buf, buf_end, func_idx);
+        if (!check_function_index(module, func_idx, error_buf,
+                                  error_buf_size)) {
+            goto fail;
+        }
+        if (func_idx < module->import_function_count) {
+            set_error_buf(error_buf, error_buf_size,
+                          "branch hint for imported function is not allowed");
+            goto fail;
+        }
+
+        struct WASMCompilationHint *current_hint =
+            (struct WASMCompilationHint *)&module
+                ->function_hints[func_idx - module->import_function_count];
+        while (current_hint->next != NULL) {
+            current_hint = current_hint->next;
+        }
+
+        uint32 num_hints;
+        read_leb_uint32(buf, buf_end, num_hints);
+        struct WASMCompilationHintBranchHint *new_hints = loader_malloc(
+            sizeof(struct WASMCompilationHintBranchHint) * num_hints, error_buf,
+            error_buf_size);
+        for (uint32 j = 0; j < num_hints; ++j) {
+            struct WASMCompilationHintBranchHint *new_hint = &new_hints[j];
+            new_hint->next = NULL;
+            new_hint->type = WASM_COMPILATION_BRANCH_HINT;
+            read_leb_uint32(buf, buf_end, new_hint->offset);
+
+            uint32 size;
+            read_leb_uint32(buf, buf_end, size);
+            if (size != 1) {
+                set_error_buf_v(error_buf, error_buf_size,
+                                "invalid branch hint size, expected 1, got %d.",
+                                size);
+                wasm_runtime_free(new_hint);
+                goto fail;
+            }
+
+            uint8 data = *buf++;
+            if (data == 0x00)
+                new_hint->is_likely = false;
+            else if (data == 0x01)
+                new_hint->is_likely = true;
+            else {
+                set_error_buf_v(error_buf, error_buf_size,
+                                "invalid branch hint, expected 0 or 1, got %d",
+                                data);
+                wasm_runtime_free(new_hint);
+                goto fail;
+            }
+
+            current_hint->next = (struct WASMCompilationHint *)new_hint;
+            current_hint = (struct WASMCompilationHint *)new_hint;
+        }
+    }
+    if (buf != buf_end) {
+        set_error_buf(error_buf, error_buf_size,
+                      "invalid branch hint section, not filled until end");
+        goto fail;
+    }
+    return true;
+fail:
+    return false;
+}
+#endif
+
 static bool
 load_user_section(const uint8 *buf, const uint8 *buf_end, WASMModule *module,
                   bool is_load_from_file_buf, char *error_buf,
@@ -5609,6 +5697,24 @@ load_user_section(const uint8 *buf, const uint8 *buf_end, WASMModule *module,
     }
 #endif
 
+#if WASM_ENABLE_BRANCH_HINTS != 0
+    if (name_len == 25
+        && memcmp((const char *)p, "metadata.code.branch_hint", 25) == 0) {
+        p += name_len;
+        if (!handle_branch_hint_section(p, p_end, module, error_buf,
+                                        error_buf_size)) {
+            return false;
+        }
+        LOG_VERBOSE("Load branch hint section success.");
+    }
+#else
+    if (name_len == 25
+        && memcmp((const char *)p, "metadata.code.branch_hint", 25) == 0) {
+        LOG_VERBOSE("Found branch hint section, but branch hints are disabled "
+                    "in this build, skipping.");
+    }
+#endif
+
 #if WASM_ENABLE_LOAD_CUSTOM_SECTION != 0
     {
         WASMCustomSection *section =
@@ -7388,7 +7494,17 @@ wasm_loader_unload(WASMModule *module)
     }
 #endif
 #endif
-
+#if WASM_ENABLE_BRANCH_HINTS != 0
+    for (size_t i = 0; i < module->function_count; i++) {
+        // be carefull when adding more hints. This only works as long as
+        // the hint structs have been allocated all at once as an array.
+        // With only branch-hints at the moment, this is the case.
+        if (module->function_hints != NULL && module->function_hints[i] != NULL)
+            wasm_runtime_free(module->function_hints[i]);
+    }
+    if (module->function_hints != NULL)
+        wasm_runtime_free(module->function_hints);
+#endif
     wasm_runtime_free(module);
 }
 

doc/stability_wasm_proposals.md

@@ -12,16 +12,17 @@ Users can turn those features on or off by using compilation options. If a relev
 
 ## On-by-default Wasm Proposals
 
-| Proposal                              | >= Phase 4 | Compilation Option       |
-| ------------------------------------- | ---------- | ------------------------ |
-| Bulk Memory Operations                | Yes        | `WAMR_BUILD_BULK_MEMORY` |
-| Fixed-width SIMD[^1]                  | Yes        | `WAMR_BUILD_SIMD`        |
-| Import/Export of Mutable Globals[^2]  | Yes        | N/A                      |
-| Multi-value                           | Yes        | N/A                      |
-| Non-trapping float-to-int Conversions | Yes        | N/A                      |
-| Reference Types                       | Yes        | `WAMR_BUILD_REF_TYPES`   |
-| Sign-extension Operators              | Yes        | N/A                      |
-| WebAssembly C and C++ API             | No         | N/A                      |
+| Proposal                              | >= Phase 4 | Compilation Option         |
+| ------------------------------------- | ---------- |----------------------------|
+| Bulk Memory Operations                | Yes        | `WAMR_BUILD_BULK_MEMORY`   |
+| Fixed-width SIMD[^1]                  | Yes        | `WAMR_BUILD_SIMD`          |
+| Import/Export of Mutable Globals[^2]  | Yes        | N/A                        |
+| Multi-value                           | Yes        | N/A                        |
+| Non-trapping float-to-int Conversions | Yes        | N/A                        |
+| Reference Types                       | Yes        | `WAMR_BUILD_REF_TYPES`     |
+| Sign-extension Operators              | Yes        | N/A                        |
+| WebAssembly C and C++ API             | No         | N/A                        |
+| Branch Hinting                        | Yes        | `WASM_ENABLE_BRANCH_HINTS` |
 
 [^1]: llvm-jit and aot only.
 
@@ -54,7 +55,6 @@ Users can turn those features on or off by using compilation options. If a relev
 
 | Proposal                                    | >= Phase 4 |
 | ------------------------------------------- | ---------- |
-| Branch Hinting                              | Yes        |
 | Custom Annotation Syntax in the Text Format | Yes        |
 | Exception Handling[^6]                      | Yes        |
 | JS String Builtins                          | Yes        |

wamr-compiler/CMakeLists.txt

@@ -46,6 +46,7 @@ add_definitions(-DWASM_ENABLE_SHARED_MEMORY=1)
 add_definitions(-DWASM_ENABLE_THREAD_MGR=1)
 add_definitions(-DWASM_ENABLE_TAIL_CALL=1)
 add_definitions(-DWASM_ENABLE_REF_TYPES=1)
+add_definitions(-DWASM_ENABLE_BRANCH_HINTS=1)
 add_definitions(-DWASM_ENABLE_CUSTOM_NAME_SECTION=1)
 add_definitions(-DWASM_ENABLE_AOT_STACK_FRAME=1)
 add_definitions(-DWASM_ENABLE_DUMP_CALL_STACK=1)