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

@@ -695,6 +695,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"
@@ -709,7 +710,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"
 "       \"Import/Export of Mutable Globals\"\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 @@ aot_compile_op_br(AOTCompContext *comp_ctx, AOTFuncContext *func_ctx,
 
 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 @@ aot_compile_conditional_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)
 {
     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