Stack check size can be zero if omg skips stack checks.

https://bugs.webkit.org/show_bug.cgi?id=271011
rdar://124390384

Reviewed by Yusuke Suzuki.

For leaf functions that have really small stacks, this stack check can
be skipped and the ASSERT(stackCheckSize()) is wrong.

We change the assert to ensure that the stack check size is set, but
if it is not needed, we can skip the stack check.

* Source/JavaScriptCore/wasm/WasmB3IRGenerator.cpp:
(JSC::Wasm::parseAndCompileB3):
* Source/JavaScriptCore/wasm/WasmCallee.h:

Originally-landed-as: 272448.753@safari-7618-branch (aef9332). rdar://124390384
Canonical link: https://commits.webkit.org/276682@main

JSTests/wasm/stress/omg-osr-stack-check-2.js

@@ -0,0 +1,95 @@
+//@ runDefault("--jitPolicyScale=0", "--useConcurrentJIT=0")
+let tools;
+if (globalThis.callerIsBBQOrOMGCompiled) {
+  function instantiateJsc(filename, importObject) {
+    let bytes = read(filename, 'binary');
+    return WebAssembly.instantiate(bytes, importObject, 'x');
+  }
+  const log = function() {};
+  const report = $.agent.report;
+  const isJIT = callerIsBBQOrOMGCompiled;
+  tools = {log, report, isJIT, instantiate: instantiateJsc};
+} else {
+  function instantiateBrowser(filename, importObject) {
+    let bytes = fetch(filename);
+    return WebAssembly.instantiateStreaming(bytes, importObject, 'x');
+  }
+  const log = function() {};
+  const report = console.log;
+  const isJIT = () => 1;
+  tools = {log, report, isJIT, instantiate: instantiateBrowser};
+}
+const {log, report, isJIT, instantiate} = tools;
+const extra = {isJIT};
+(async function () {
+let memory0 = new WebAssembly.Memory({initial: 3947, shared: false, maximum: 6209});
+let tag3 = new WebAssembly.Tag({parameters: []});
+let global0 = new WebAssembly.Global({value: 'f64', mutable: true}, 882640.3220068762);
+let global1 = new WebAssembly.Global({value: 'f32', mutable: true}, 162294.89036678328);
+let global2 = new WebAssembly.Global({value: 'f64', mutable: true}, 50173.96827009934);
+let table0 = new WebAssembly.Table({initial: 6, element: 'externref'});
+let m1 = {global0, memory0, tag3};
+let m0 = {global1, global2};
+let m2 = {table0};
+let importObject0 = /** @type {Imports2} */ ({m0, m1, m2});
+let i0 = await instantiate('omg-osr-stack-check-2.wasm', importObject0);
+let {fn0, global3, global4, memory1, table1, table2, table3, table4, table5, table6, table7, tag0, tag1, tag2} = /**
+  @type {{
+fn0: () => void,
+global3: WebAssembly.Global,
+global4: WebAssembly.Global,
+memory1: WebAssembly.Memory,
+table1: WebAssembly.Table,
+table2: WebAssembly.Table,
+table3: WebAssembly.Table,
+table4: WebAssembly.Table,
+table5: WebAssembly.Table,
+table6: WebAssembly.Table,
+table7: WebAssembly.Table,
+tag0: WebAssembly.Tag,
+tag1: WebAssembly.Tag,
+tag2: WebAssembly.Tag
+  }} */ (i0.instance.exports);
+table4.set(6, table7);
+table4.set(44, table1);
+global4.value = 0;
+log('calling fn0');
+report('progress');
+try {
+  for (let k=0; k<21; k++) {
+  let zzz = fn0();
+  if (zzz !== undefined) { throw new Error('expected undefined but return value is '+zzz); }
+  }
+} catch (e) {
+  if (e instanceof WebAssembly.Exception) {
+  log(e); if (e.stack) { log(e.stack); }
+  } else if (e instanceof TypeError) {
+  if (e.message === 'an exported wasm function cannot contain a v128 parameter or return value') { log(e); } else { throw e; }
+  } else if (e instanceof WebAssembly.RuntimeError || e instanceof RangeError) { log(e); } else { throw e; }
+}
+log('calling fn0');
+report('progress');
+try {
+  for (let k=0; k<19; k++) {
+  let zzz = fn0();
+  if (zzz !== undefined) { throw new Error('expected undefined but return value is '+zzz); }
+  }
+} catch (e) {
+  if (e instanceof WebAssembly.Exception) {
+  log(e); if (e.stack) { log(e.stack); }
+  } else if (e instanceof TypeError) {
+  if (e.message === 'an exported wasm function cannot contain a v128 parameter or return value') { log(e); } else { throw e; }
+  } else if (e instanceof WebAssembly.RuntimeError || e instanceof RangeError) { log(e); } else { throw e; }
+}
+let tables = [table0, table7, table5, table1, table4, table3, table6, table2];
+for (let table of tables) {
+for (let k=0; k < table.length; k++) { table.get(k)?.toString(); }
+}
+})().then(() => {
+  log('after')
+  report('after');
+}).catch(e => {
+  log(e)
+  log('error')
+  report('error');
+})

Source/JavaScriptCore/wasm/WasmCallee.h

@@ -243,6 +243,9 @@ class OptimizingJITCallee : public JITCallee {
     Vector<Ref<NameSection>, 0> nameSections;
 };
 
+constexpr int32_t stackCheckUnset = 0;
+constexpr int32_t stackCheckNotNeeded = -1;
+
 class OSREntryCallee final : public OptimizingJITCallee {
     WTF_MAKE_TZONE_ALLOCATED(OSREntryCallee);
 public:
@@ -261,8 +264,18 @@ class OSREntryCallee final : public OptimizingJITCallee {
         OptimizingJITCallee::setEntrypoint(WTFMove(entrypoint), WTFMove(unlinkedCalls), WTFMove(stackmaps), WTFMove(exceptionHandlers), WTFMove(exceptionHandlerLocations));
     }
 
-    void setStackCheckSize(unsigned stackCheckSize) { m_stackCheckSize = stackCheckSize; }
-    unsigned stackCheckSize() const { return m_stackCheckSize; }
+    void setStackCheckSize(int32_t stackCheckSize)
+    {
+        ASSERT(m_stackCheckSize == stackCheckUnset);
+        ASSERT(stackCheckSize > 0 || stackCheckSize == stackCheckNotNeeded);
+        m_stackCheckSize = stackCheckSize;
+    }
+
+    int32_t stackCheckSize() const
+    {
+        ASSERT(m_stackCheckSize > 0 || m_stackCheckSize == stackCheckNotNeeded);
+        return m_stackCheckSize;
+    }
 
 private:
     OSREntryCallee(CompilationMode compilationMode, size_t index, std::pair<const Name*, RefPtr<NameSection>>&& name, uint32_t loopIndex)
@@ -273,7 +286,7 @@ class OSREntryCallee final : public OptimizingJITCallee {
 
     unsigned m_osrEntryScratchBufferSize { 0 };
     uint32_t m_loopIndex;
-    unsigned m_stackCheckSize { 0 };
+    int32_t m_stackCheckSize { stackCheckUnset };
 };
 
 #endif // ENABLE(WEBASSEMBLY_BBQJIT) || ENABLE(WEBASSEMBLY_OMGJIT)
@@ -355,8 +368,17 @@ class BBQCallee final : public OptimizingJITCallee {
         return m_switchJumpTables.last().ptr();
     }
 
-    void setStackCheckSize(unsigned stackCheckSize) { m_stackCheckSize = stackCheckSize; }
-    unsigned stackCheckSize() const { return m_stackCheckSize; }
+    void setStackCheckSize(unsigned stackCheckSize)
+    {
+        ASSERT(m_stackCheckSize == stackCheckUnset);
+        ASSERT(stackCheckSize > 0 || int32_t(stackCheckSize) == stackCheckNotNeeded);
+        m_stackCheckSize = stackCheckSize;
+    }
+    int32_t stackCheckSize() const
+    {
+        ASSERT(m_stackCheckSize > 0 || int32_t(m_stackCheckSize) == stackCheckNotNeeded);
+        return m_stackCheckSize;
+    }
 
 private:
     BBQCallee(size_t index, std::pair<const Name*, RefPtr<NameSection>>&& name, std::unique_ptr<TierUpCount>&& tierUpCount, SavedFPWidth savedFPWidth)

Source/JavaScriptCore/wasm/WasmOMGIRGenerator.cpp

@@ -5314,6 +5314,8 @@ Expected<std::unique_ptr<InternalFunction>, String> parseAndCompileOMG(Compilati
         bool needsOverflowCheck = false;
         irGenerator.computeStackCheckSize(needsOverflowCheck, checkSize);
         ASSERT(checkSize || !needsOverflowCheck);
+        if (!needsOverflowCheck)
+            checkSize = stackCheckNotNeeded;
         static_cast<OSREntryCallee*>(&callee)->setStackCheckSize(checkSize);
     }
 

Source/JavaScriptCore/wasm/WasmOperations.cpp

@@ -328,6 +328,8 @@ JSC_DEFINE_JIT_OPERATION(operationWasmTriggerOSREntryNow, void, (Probe::Context&
     auto doStackCheck = [instance](OSREntryCallee* callee) -> bool {
         uintptr_t stackPointer = reinterpret_cast<uintptr_t>(currentStackPointer());
         ASSERT(callee->stackCheckSize());
+        if (callee->stackCheckSize() == stackCheckNotNeeded)
+            return true;
         uintptr_t stackExtent = stackPointer - callee->stackCheckSize();
         uintptr_t stackLimit = reinterpret_cast<uintptr_t>(instance->softStackLimit());
         if (UNLIKELY(stackExtent >= stackPointer || stackExtent <= stackLimit)) {

Source/JavaScriptCore/wasm/WasmSlowPaths.cpp

@@ -305,6 +305,17 @@ WASM_SLOW_PATH_DECL(loop_osr)
 
             size_t osrEntryScratchBufferSize = osrEntryCallee->osrEntryScratchBufferSize();
             RELEASE_ASSERT(osrEntryScratchBufferSize == osrEntryData.values.size());
+
+            if (osrEntryCallee->stackCheckSize() != Wasm::stackCheckNotNeeded) {
+                uintptr_t stackPointer = reinterpret_cast<uintptr_t>(currentStackPointer());
+                uintptr_t stackExtent = stackPointer - osrEntryCallee->stackCheckSize();
+                uintptr_t stackLimit = reinterpret_cast<uintptr_t>(instance->softStackLimit());
+                if (UNLIKELY(stackExtent >= stackPointer || stackExtent <= stackLimit)) {
+                    dataLogLnIf(Options::verboseOSR(), "Skipping OMG loop tier up due to stack check; ", RawHex(stackPointer), " -> ", RawHex(stackExtent), " is past soft limit ", RawHex(stackLimit));
+                    WASM_RETURN_TWO(nullptr, nullptr);
+                }
+            }
+
             uint64_t* buffer = instance->vm().wasmContext.scratchBufferForSize(osrEntryScratchBufferSize);
             if (!buffer)
                 WASM_RETURN_TWO(nullptr, nullptr);