Add support for metadata.code.branch_hint section

CMakeLists.txt

@@ -119,6 +119,11 @@ if (NOT DEFINED WAMR_BUILD_REF_TYPES)
   set (WAMR_BUILD_REF_TYPES 1)
 endif ()
 
+if (NOT DEFINED WAMR_BUILD_BRANCH_HINTS)
+  # Enable branch hints by default
+  set (WAMR_BUILD_BRANCH_HINTS 1)
+endif ()
+
 set (WAMR_ROOT_DIR ${CMAKE_CURRENT_SOURCE_DIR})
 
 include (${WAMR_ROOT_DIR}/build-scripts/runtime_lib.cmake)

build-scripts/SConscript_config

@@ -128,6 +128,10 @@ if GetDepend(['WAMR_BUILD_REF_TYPES']):
     CPPDEFINES += ['WASM_ENABLE_REF_TYPES=1']
     print('[WAMR] enable ref types')
 
+if GetDepend(['WAMR_BUILD_BRANCH_HINTS']):
+    CPPDEFINES += ['WASM_ENABLE_BRANCH_HINTS=1']
+    print('[WAMR] enable branch hints')
+
 CPPDEFINES += ['BH_MALLOC=wasm_runtime_malloc']
 CPPDEFINES += ['BH_FREE=wasm_runtime_free']
 

build-scripts/config_common.cmake

@@ -215,6 +215,10 @@ if (NOT DEFINED WAMR_BUILD_EXTENDED_CONST_EXPR)
   set (WAMR_BUILD_EXTENDED_CONST_EXPR 0)
 endif ()
 
+if (NOT DEFINED WAMR_BUILD_BRANCH_HINTS)
+  set (WAMR_BUILD_BRANCH_HINTS 1)
+endif ()
+
 ########################################
 # Compilation options to marco
 ########################################
@@ -419,6 +423,12 @@ endif ()
 if (WAMR_BUILD_REF_TYPES EQUAL 1)
   add_definitions (-DWASM_ENABLE_REF_TYPES=1)
 endif ()
+if (WAMR_BUILD_BRANCH_HINTS EQUAL 1)
+  add_definitions (-DWASM_ENABLE_BRANCH_HINTS=1)
+  message ("     branch hints enabled")
+else ()
+  message ("     branch hints disabled")
+endif ()
 if (WAMR_BUILD_GC EQUAL 1)
   if (WAMR_TEST_GC EQUAL 1)
     message("      GC testing enabled")
@@ -708,8 +718,8 @@ message (
 "       \"Tail call\" via WAMR_BUILD_TAIL_CALL: ${WAMR_BUILD_TAIL_CALL}\n"
 "       \"Threads\" via WAMR_BUILD_SHARED_MEMORY: ${WAMR_BUILD_SHARED_MEMORY}\n"
 "       \"Typed Function References\" via WAMR_BUILD_GC: ${WAMR_BUILD_GC}\n"
+"       \"Branch Hinting\" via WAMR_BUILD_BRANCH_HINTS: ${WAMR_BUILD_BRANCH_HINTS}\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/aot/aot_loader.c

@@ -505,6 +505,13 @@ check_feature_flags(char *error_buf, uint32 error_buf_size,
     }
 #endif
 
+#if WASM_ENABLE_BRANCH_HINTS == 0
+    if (feature_flags & WASM_ENABLE_BRANCH_HINTS) {
+        LOG_WARNING(
+            "branch hints not enabled, but wasm file contains branch hints");
+    }
+#endif
+
 #if WASM_ENABLE_GC == 0
     if (feature_flags & WASM_FEATURE_GARBAGE_COLLECTION) {
         set_error_buf(error_buf, error_buf_size,

core/iwasm/aot/aot_runtime.h

@@ -35,6 +35,7 @@ extern "C" {
  * and not at the beginning of each function call */
 #define WASM_FEATURE_FRAME_PER_FUNCTION (1 << 12)
 #define WASM_FEATURE_FRAME_NO_FUNC_IDX (1 << 13)
+#define WASM_FEATURE_BRANCH_HINTS (1 << 14)
 
 typedef enum AOTSectionType {
     AOT_SECTION_TYPE_TARGET_INFO = 0,
@@ -57,6 +58,7 @@ typedef enum AOTCustomSectionType {
     AOT_CUSTOM_SECTION_ACCESS_CONTROL = 2,
     AOT_CUSTOM_SECTION_NAME = 3,
     AOT_CUSTOM_SECTION_STRING_LITERAL = 4,
+    AOT_CUSTOM_SECTION_CODE_METADATA = 5,
 } AOTCustomSectionType;
 
 typedef struct AOTObjectDataSection {

core/iwasm/compilation/aot.c

@@ -392,6 +392,7 @@ aot_create_funcs(const WASMModule *module, uint32 pointer_size)
         memset(aot_func, 0, sizeof(AOTFunc));
 
         func_type = aot_func->func_type = func->func_type;
+        aot_func->func_index = i + module->import_function_count;
 
         /* Resolve function type index */
         for (j = 0; j < module->type_count; j++) {
@@ -416,6 +417,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 +876,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

@@ -207,6 +207,7 @@ typedef struct AOTImportFunc {
 typedef struct AOTFunc {
     AOTFuncType *func_type;
     uint32 func_type_index;
+    uint32 func_index;
     uint32 local_count;
     uint8 *local_types_wp;
     uint16 param_cell_num;
@@ -217,6 +218,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 +300,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,16 @@ aot_compile_func(AOTCompContext *comp_ctx, uint32 func_index)
 
             case WASM_OP_BR_IF:
             {
+                // ip is advanced by one byte for the opcode
+#if WASM_ENABLE_BRANCH_HINTS != 0
+                uint32 instr_offset =
+                    (frame_ip - 0x1) - (func_ctx->aot_func->code_body_begin);
+#else
+                uint32 instr_offset = 0;
+#endif
                 read_leb_uint32(frame_ip, frame_ip_end, br_depth);
                 if (!aot_compile_op_br_if(comp_ctx, func_ctx, br_depth,
-                                          &frame_ip))
+                                          &frame_ip, instr_offset))
                     return false;
                 break;
             }

core/iwasm/compilation/aot_emit_aot_file.c

@@ -4501,6 +4501,9 @@ aot_obj_data_create(AOTCompContext *comp_ctx)
     if (comp_ctx->enable_ref_types) {
         obj_data->target_info.feature_flags |= WASM_FEATURE_REF_TYPES;
     }
+    if (comp_ctx->enable_branch_hints) {
+        obj_data->target_info.feature_flags |= WASM_FEATURE_BRANCH_HINTS;
+    }
     if (comp_ctx->enable_gc) {
         obj_data->target_info.feature_flags |= WASM_FEATURE_GARBAGE_COLLECTION;
     }

core/iwasm/compilation/aot_emit_control.c

@@ -87,6 +87,16 @@ 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),                             \
+            !instr) {                                                        \
+            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,40 @@ 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(LLVMContextRef ctxt, AOTFuncContext *func_ctx,
+                     uint32 offset, LLVMValueRef br_if_instr)
+{
+    struct WASMCompilationHintBranchHint *hint =
+        (struct WASMCompilationHintBranchHint *)func_ctx->function_hints;
+    while (hint != NULL) {
+        if (hint->type == WASM_COMPILATION_BRANCH_HINT
+            && hint->offset == offset) {
+            break;
+        }
+        hint = hint->next;
+    }
+    if (hint != NULL) {
+        LLVMMetadataRef header = LLVMMDStringInContext2(
+            ctxt, "branch_weights", 14 /* strlen("branch_hint") */);
+        // same weight llvm MDBuilder::createLikelyBranchWeights assigns
+        const uint32_t likely_weight = (1U << 20) - 1;
+        const uint32_t unlikely_weight = 1;
+        LLVMMetadataRef true_w = LLVMValueAsMetadata(LLVMConstInt(
+            LLVMInt32TypeInContext(ctxt),
+            hint->is_likely ? likely_weight : unlikely_weight, false));
+        LLVMMetadataRef false_w = LLVMValueAsMetadata(LLVMConstInt(
+            LLVMInt32TypeInContext(ctxt),
+            hint->is_likely ? unlikely_weight : likely_weight, false));
+        LLVMMetadataRef mds[] = { header, true_w, false_w };
+        LLVMMetadataRef md = LLVMMDNodeInContext2(ctxt, mds, 3);
+        LLVMValueRef md_val = LLVMMetadataAsValue(ctxt, md);
+        LLVMSetMetadata(br_if_instr, LLVMGetMDKindID("prof", 4), md_val);
+    }
+}
+#endif
+
 static bool
 handle_next_reachable_block(AOTCompContext *comp_ctx, AOTFuncContext *func_ctx,
                             uint8 **p_frame_ip)
@@ -673,13 +717,33 @@ 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->context, 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->context, 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,
@@ -1027,7 +1091,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)
+                           uint8 **p_frame_ip, uint32 off)
 {
     AOTBlock *block_dst;
     LLVMValueRef value, *values = NULL;
@@ -1108,8 +1172,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->context, func_ctx, off, 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 +1223,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->context, func_ctx, off, 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,14 +1255,14 @@ 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)
+                     uint32 br_depth, uint8 **p_frame_ip, uint32 off)
 {
     LLVMValueRef value_cmp;
 
     POP_COND(value_cmp);
 
     return aot_compile_conditional_br(comp_ctx, func_ctx, br_depth, value_cmp,
-                                      p_frame_ip);
+                                      p_frame_ip, off);
 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);
+                     uint32 br_depth, uint8 **p_frame_ip, uint32 offset);
 
 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:
@@ -2733,6 +2740,9 @@ aot_create_comp_context(const AOTCompData *comp_data, aot_comp_option_t option)
     if (option->enable_ref_types)
         comp_ctx->enable_ref_types = true;
 
+    if (option->enable_branch_hints)
+        comp_ctx->enable_branch_hints = true;
+
     comp_ctx->aux_stack_frame_type = option->aux_stack_frame_type;
     comp_ctx->call_stack_features = option->call_stack_features;
 

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;
 
@@ -449,6 +452,9 @@ typedef struct AOTCompContext {
     /* Reference Types */
     bool enable_ref_types;
 
+    /* Branch Hinting */
+    bool enable_branch_hints;
+
     /* Disable LLVM built-in intrinsics */
     bool disable_llvm_intrinsics;
 

core/iwasm/include/aot_comp_option.h

@@ -66,6 +66,7 @@ typedef struct AOTCompOption {
     bool enable_tail_call;
     bool enable_simd;
     bool enable_ref_types;
+    bool enable_branch_hints;
     bool enable_gc;
     bool enable_aux_stack_check;
     bool enable_extended_const;