Bug 1837683 - wasm: test case for branch hinting. r=rhunt

Differential Revision: https://phabricator.services.mozilla.com/D201263

js/src/builtin/TestingFunctions.cpp

@@ -2125,7 +2125,7 @@ static bool WasmDumpIon(JSContext* cx, unsigned argc, Value* vp) {
   return true;
 }
 
-enum class Flag { Tier2Complete, Deserialized };
+enum class Flag { Tier2Complete, Deserialized, ParsedBranchHints };
 
 static bool WasmReturnFlag(JSContext* cx, unsigned argc, Value* vp, Flag flag) {
   CallArgs args = CallArgsFromVp(argc, vp);
@@ -2150,6 +2150,9 @@ static bool WasmReturnFlag(JSContext* cx, unsigned argc, Value* vp, Flag flag) {
     case Flag::Deserialized:
       b = module->module().loggingDeserialized();
       break;
+    case Flag::ParsedBranchHints:
+      b = module->module().metadata().parsedBranchHints;
+      break;
   }
 
   args.rval().set(BooleanValue(b));
@@ -2221,6 +2224,12 @@ static bool WasmLoadedFromCache(JSContext* cx, unsigned argc, Value* vp) {
   return WasmReturnFlag(cx, argc, vp, Flag::Deserialized);
 }
 
+#ifdef ENABLE_WASM_BRANCH_HINTING
+static bool WasmParsedBranchHints(JSContext* cx, unsigned argc, Value* vp) {
+  return WasmReturnFlag(cx, argc, vp, Flag::ParsedBranchHints);
+}
+#endif  // ENABLE_WASM_BRANCH_HINTING
+
 static bool WasmBuiltinI8VecMul(JSContext* cx, unsigned argc, Value* vp) {
   if (!wasm::HasSupport(cx)) {
     JS_ReportErrorASCII(cx, "wasm support unavailable");
@@ -9956,6 +9965,14 @@ JS_FOR_WASM_FEATURES(WASM_FEATURE)
 "  Gets the length of a WebAssembly GC array."),
 #endif // ENABLE_WASM_GC
 
+#ifdef ENABLE_WASM_BRANCH_HINTING
+    JS_FN_HELP("wasmParsedBranchHints", WasmParsedBranchHints, 1, 0,
+"wasmParsedBranchHints(module)",
+"  Returns a boolean indicating whether a given module has successfully parsed a\n"
+"  custom branch hinting section."),
+
+#endif // ENABLE_WASM_BRANCH_HINTING
+
     JS_FN_HELP("largeArrayBufferSupported", LargeArrayBufferSupported, 0, 0,
 "largeArrayBufferSupported()",
 "  Returns true if array buffers larger than 2GB can be allocated."),

js/src/jit-test/tests/wasm/branch-hinting/directives.txt

@@ -0,0 +1 @@
+|jit-test| --setpref=wasm_branch_hinting=true; --wasm-compiler=ion; test-also=--wasm-compiler=baseline;skip-if: !wasmBranchHintingEnabled(); include:wasm.js

js/src/jit-test/tests/wasm/branch-hinting/parsing.js

@@ -0,0 +1,154 @@
+// Make sure we are correctly parsing this custom section.
+var code =`
+(module
+    (func $$dummy)
+    (func $main (param i32) (result i32)
+        i32.const 0
+        local.get 0
+        i32.eq
+        ;; Only allowed on br_if and if
+        (@metadata.code.branch_hint "\\00") if
+            call $$dummy
+            i32.const 1
+            return
+        else
+            call $$dummy
+            i32.const 0
+            return
+        end
+        i32.const 3
+        return
+    )
+
+    (export "_main" (func $main))
+)`;
+
+let branchHintsModule = new WebAssembly.Module(wasmTextToBinary(code));
+assertEq(WebAssembly.Module.customSections(branchHintsModule, "metadata.code.branch_hint").length, 1);
+assertEq(wasmParsedBranchHints(branchHintsModule), true);
+
+let instance = new WebAssembly.Instance(branchHintsModule);
+assertEq(instance.exports._main(0), 1);
+
+// Testing branch hints parsing on `if` and `br_if`
+branchHintsModule = new WebAssembly.Module(wasmTextToBinary(`
+(module
+    (func $main
+      i32.const 0
+      (@metadata.code.branch_hint "\\00")
+      if
+        i32.const 0
+        (@metadata.code.branch_hint "\\01")
+        br_if 0
+      end
+    )
+    (export "_main" (func $main))
+)`));
+assertEq(wasmParsedBranchHints(branchHintsModule), true);
+instance = new WebAssembly.Instance(branchHintsModule);
+instance.exports._main();
+
+let m = new WebAssembly.Module(wasmTextToBinary(`
+(module
+  (type (;0;) (func))
+  (type (;1;) (func (param i32) (result i32)))
+  (type (;2;) (func (result i32)))
+  (func $__wasm_nullptr (type 0)
+    unreachable)
+  (func $main (type 2) (result i32)
+    (local i32 i32 i32 i32)
+    i32.const 0
+    local.tee 2
+    local.set 3
+    loop
+      local.get 2
+      i32.const 50000
+      i32.eq
+      (@metadata.code.branch_hint "\\00") if
+        i32.const 1
+        local.set 3
+      end
+      local.get 2
+      i32.const 1
+      i32.add
+      local.tee 2
+      i32.const 100000
+      i32.ne
+      (@metadata.code.branch_hint "\\01") br_if 0 (;@1;)
+    end
+    local.get 3)
+  (table (;0;) 1 1 funcref)
+  (memory (;0;) 17 128)
+  (global (;0;) (mut i32) (i32.const 42))
+  (export "memory" (memory 0))
+  (export "_main" (func $main))
+  (elem (;0;) (i32.const 0) func $__wasm_nullptr)
+  (type (;0;) (func (param i32)))
+)`));
+
+assertEq(wasmParsedBranchHints(m), true);
+instance = new WebAssembly.Instance(m);
+assertEq(instance.exports._main(0), 1);
+
+// Testing invalid values for branch hints
+assertErrorMessage(() => new WebAssembly.Module(wasmTextToBinary(`
+(module
+    (func $main (param i32) (result i32)
+      i32.const 0
+      (@metadata.code.branch_hint "\\0000000") if
+        i32.const 1
+        return
+      end
+      i32.const 42
+      return
+    )
+  )
+`)), SyntaxError, /invalid value for branch hint/);
+
+assertErrorMessage(() => new WebAssembly.Module(wasmTextToBinary(`
+(module
+    (func $main (param i32) (result i32)
+      i32.const 0
+      (@metadata.code.branch_hint "\\02") if
+        i32.const 1
+        return
+      end
+      i32.const 42
+      return
+    )
+  )
+`)), SyntaxError, /invalid value for branch hint/);
+
+assertErrorMessage(() => new WebAssembly.Module(wasmTextToBinary(`
+(module
+    (func $main (param i32) (result i32)
+      i32.const 0
+      (@metadata.code.branch_hint "\\aaaa") if
+        i32.const 1
+        return
+      end
+      i32.const 42
+      return
+    )
+  )
+`)), SyntaxError, /wasm text error/);
+
+assertErrorMessage(() => new WebAssembly.Module(wasmTextToBinary(`
+(module
+    (func $main (param i32) (result i32)
+      i32.const 0
+      (@metadata.code.branch_hint) if
+        i32.const 1
+        return
+      end
+      i32.const 42
+      return
+    )
+  )
+`)), SyntaxError, /wasm text error/);
+
+assertErrorMessage(() => new WebAssembly.Module(wasmTextToBinary(`
+(module
+    (@metadata.code.branch_hint)
+)
+`)), SyntaxError, /wasm text error/);

js/src/jit-test/tests/wasm/branch-hinting/simple_example.js

@@ -0,0 +1,35 @@
+// Branch Hinting proposal
+
+function runModule(hint) {
+    let code =`
+    (module
+        (func $$dummy)
+        (func $main (param i32) (result i32)
+            i32.const 0
+            local.get 0
+            i32.eq
+            ;; Only allowed on br_if and if
+            (@metadata.code.branch_hint "${hint}") if
+                call $$dummy
+                i32.const 1
+                return
+            else
+                call $$dummy
+                i32.const 0
+                return
+            end
+            i32.const 3
+            return
+        )
+        (export "_main" (func $main))
+    )`;
+    let branchHintsModule = new WebAssembly.Module(wasmTextToBinary(code));
+    assertEq(wasmParsedBranchHints(branchHintsModule), true);
+
+    let instance = new WebAssembly.Instance(branchHintsModule);
+    assertEq(instance.exports._main(0), 1);
+}
+
+// Ensure that we have the same result with different branch hints.
+runModule("\\00");
+runModule("\\01");

js/src/wasm/WasmCode.h

@@ -364,6 +364,7 @@ struct MetadataCacheablePod {
   BuiltinModuleIds builtinModules;
   FeatureUsage featureUsage;
   bool filenameIsURL;
+  bool parsedBranchHints;
   uint32_t typeDefsOffsetStart;
   uint32_t memoriesOffsetStart;
   uint32_t tablesOffsetStart;
@@ -372,15 +373,16 @@ struct MetadataCacheablePod {
 
   WASM_CHECK_CACHEABLE_POD(kind, instanceDataLength, startFuncIndex,
                            nameCustomSectionIndex, builtinModules, featureUsage,
-                           filenameIsURL, typeDefsOffsetStart,
-                           memoriesOffsetStart, tablesOffsetStart,
-                           tagsOffsetStart)
+                           filenameIsURL, parsedBranchHints,
+                           typeDefsOffsetStart, memoriesOffsetStart,
+                           tablesOffsetStart, tagsOffsetStart)
 
   explicit MetadataCacheablePod(ModuleKind kind)
       : kind(kind),
         instanceDataLength(0),
         featureUsage(FeatureUsage::None),
         filenameIsURL(false),
+        parsedBranchHints(false),
         typeDefsOffsetStart(UINT32_MAX),
         memoriesOffsetStart(UINT32_MAX),
         tablesOffsetStart(UINT32_MAX),

js/src/wasm/WasmGenerator.cpp

@@ -1135,6 +1135,7 @@ SharedMetadata ModuleGenerator::finishMetadata(const Bytes& bytecode) {
   metadata_->nameCustomSectionIndex = moduleEnv_->nameCustomSectionIndex;
   metadata_->moduleName = moduleEnv_->moduleName;
   metadata_->funcNames = std::move(moduleEnv_->funcNames);
+  metadata_->parsedBranchHints = moduleEnv_->parsedBranchHints;
 
   // Copy over additional debug information.
 

js/src/wasm/WasmValidate.cpp

@@ -3008,7 +3008,7 @@ static bool DecodeBranchHintingSection(Decoder& d, ModuleEnvironment* env) {
   }
 
   // Skip this custom section if errors are encountered during parsing.
-  (void)ParseBranchHintingSection(d, env);
+  env->parsedBranchHints = ParseBranchHintingSection(d, env);
 
   d.finishCustomSection(BranchHintingSectionName, *range);
   return true;

js/src/wasm/WasmValidate.h

@@ -94,6 +94,9 @@ struct ModuleEnvironment {
   Maybe<Name> moduleName;
   NameVector funcNames;
 
+  // Indicates whether the branch hint section was successfully parsed.
+  bool parsedBranchHints;
+
   explicit ModuleEnvironment(FeatureArgs features,
                              ModuleKind kind = ModuleKind::Wasm)
       : kind(kind),
@@ -104,7 +107,8 @@ struct ModuleEnvironment {
         typeDefsOffsetStart(UINT32_MAX),
         memoriesOffsetStart(UINT32_MAX),
         tablesOffsetStart(UINT32_MAX),
-        tagsOffsetStart(UINT32_MAX) {}
+        tagsOffsetStart(UINT32_MAX),
+        parsedBranchHints(false) {}
 
   [[nodiscard]] bool init() {
     types = js_new<TypeContext>(features);