Add support for metadata.code.branch_hint section

Author: Lukasdoe

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;
     }