[JSC] Compile wasm stubs in compiler threads

https://bugs.webkit.org/show_bug.cgi?id=274721
rdar://128746101

Reviewed by Justin Michaud.

This patch moves wasm stub compilation from the main thread to one of compiler thread.
When EntryPlan is running, the first thread which starts compiling will compile stubs too.
If no thread compiled it, it gets compiled in the complete phase.

Since we can complete stub compilation in the compiler thread and share it in all wasm users,
we can make JSWebAssemblyModule creation simple. Now we change JSWebAssemblyModule::createStub
to JSWebAssemblyModule::create, and which just creates JSWebAssemblyModule and it does not fail
anymore.

We also change WasmEntryPlan's batch mechanism: instead of checking for each function compilation,
we now compute size apriori and grab range of function indice to compile.

And we also stop computing costly thing for wasm tail-calls which is not enabled. It should be done
more efficiently when enabling it.

* Source/JavaScriptCore/wasm/WasmBBQPlan.cpp:
(JSC::Wasm::BBQPlan::didCompleteCompilation):
* Source/JavaScriptCore/wasm/WasmEntryPlan.cpp:
(JSC::Wasm::EntryPlan::prepare):
(JSC::Wasm::EntryPlan::compileFunctions):
(JSC::Wasm::EntryPlan::didCompleteCompilation):
(JSC::Wasm::EntryPlan::generateWasmToWasmStubs):
(JSC::Wasm::EntryPlan::generateWasmToJSStubs):
* Source/JavaScriptCore/wasm/WasmEntryPlan.h:
* Source/JavaScriptCore/wasm/WasmIPIntPlan.cpp:
(JSC::Wasm::IPIntPlan::didCompleteCompilation):
* Source/JavaScriptCore/wasm/WasmLLIntPlan.cpp:
(JSC::Wasm::LLIntPlan::didCompleteCompilation):
* Source/JavaScriptCore/wasm/WasmModule.h:
* Source/JavaScriptCore/wasm/WasmStreamingCompiler.cpp:
(JSC::Wasm::StreamingCompiler::didComplete):
* Source/JavaScriptCore/wasm/js/JSWebAssembly.cpp:
(JSC::JSWebAssembly::webAssemblyModuleValidateAsync):
(JSC::compileAndInstantiate):
* Source/JavaScriptCore/wasm/js/JSWebAssemblyModule.cpp:
(JSC::JSWebAssemblyModule::create):
(JSC::JSWebAssemblyModule::createStub): Deleted.
* Source/JavaScriptCore/wasm/js/JSWebAssemblyModule.h:
* Source/JavaScriptCore/wasm/js/WebAssemblyModuleConstructor.cpp:
(JSC::WebAssemblyModuleConstructor::createModule):
* Source/WebCore/bindings/js/SerializedScriptValue.cpp:
(WebCore::CloneDeserializer::readTerminal):

Canonical link: https://commits.webkit.org/279352@main

Source/JavaScriptCore/wasm/WasmBBQPlan.cpp

@@ -58,6 +58,8 @@ BBQPlan::BBQPlan(VM& vm, Ref<ModuleInformation> moduleInformation, uint32_t func
     ASSERT(Options::useBBQJIT());
     setMode(m_calleeGroup->mode());
     dataLogLnIf(WasmBBQPlanInternal::verbose, "Starting BBQ plan for ", functionIndex);
+    m_areWasmToJSStubsCompiled = true;
+    m_areWasmToWasmStubsCompiled = true;
 }
 
 FunctionAllowlist& BBQPlan::ensureGlobalBBQAllowlist()
@@ -322,6 +324,8 @@ std::unique_ptr<InternalFunction> BBQPlan::compileFunction(uint32_t functionInde
 
 void BBQPlan::didCompleteCompilation()
 {
+    generateStubsIfNecessary();
+
     for (uint32_t functionIndex = 0; functionIndex < m_moduleInformation->functions.size(); functionIndex++) {
         CompilationContext& context = m_compilationContexts[functionIndex];
         TypeIndex typeIndex = m_moduleInformation->internalFunctionTypeIndices[functionIndex];

Source/JavaScriptCore/wasm/WasmEntryPlan.cpp

@@ -125,13 +125,6 @@ void EntryPlan::prepare()
 
     m_unlinkedWasmToWasmCalls.resize(functions.size());
 
-#if ENABLE(JIT)
-    if (UNLIKELY(!generateWasmToWasmStubs()))
-        return;
-    if (UNLIKELY(!generateWasmToJSStubs()))
-        return;
-#endif
-
     const uint32_t importFunctionCount = m_moduleInformation->importFunctionCount();
     for (const auto& exp : m_moduleInformation->exports) {
         if (exp.kindIndex >= importFunctionCount)
@@ -196,6 +189,9 @@ void EntryPlan::compileFunctions(CompilationEffort effort)
             return;
 
         uint32_t functionIndex;
+        uint32_t functionIndexEnd;
+        bool areWasmToWasmStubsCompiled = false;
+        bool areWasmToJSStubsCompiled = false;
         {
             Locker locker { m_lock };
             if (m_currentIndex >= m_numberOfFunctions) {
@@ -204,11 +200,41 @@ void EntryPlan::compileFunctions(CompilationEffort effort)
                 return;
             }
             functionIndex = m_currentIndex;
-            ++m_currentIndex;
+            functionIndexEnd = m_numberOfFunctions;
+            for (uint32_t index = functionIndex; index < m_numberOfFunctions; ++index) {
+                bytesCompiled += m_moduleInformation->functions[index].data.size();
+                if (bytesCompiled >= Options::webAssemblyPartialCompileLimit()) {
+                    functionIndexEnd = index + 1;
+                    break;
+                }
+            }
+            m_currentIndex = functionIndexEnd;
+            areWasmToWasmStubsCompiled = std::exchange(m_areWasmToWasmStubsCompiled, true);
+            areWasmToJSStubsCompiled = std::exchange(m_areWasmToJSStubsCompiled, true);
+        }
+
+        for (uint32_t index = functionIndex; index < functionIndexEnd; ++index)
+            compileFunction(index);
+
+        if (!areWasmToWasmStubsCompiled) {
+#if ENABLE(JIT)
+            if (UNLIKELY(!generateWasmToWasmStubs())) {
+                Locker locker { m_lock };
+                fail(makeString("Out of executable memory at stub generation"_s));
+                return;
+            }
+#endif
         }
 
-        compileFunction(functionIndex);
-        bytesCompiled += m_moduleInformation->functions[functionIndex].data.size();
+        if (!areWasmToJSStubsCompiled) {
+#if ENABLE(JIT)
+            if (UNLIKELY(!generateWasmToJSStubs())) {
+                Locker locker { m_lock };
+                fail(makeString("Out of executable memory at stub generation"_s));
+                return;
+            }
+#endif
+        }
     }
 }
 
@@ -221,11 +247,33 @@ void EntryPlan::complete()
         didCompleteCompilation();
 
     if (!isComplete()) {
+        generateStubsIfNecessary();
         moveToState(State::Completed);
         runCompletionTasks();
     }
 }
 
+void EntryPlan::generateStubsIfNecessary()
+{
+    if (!std::exchange(m_areWasmToWasmStubsCompiled, true)) {
+#if ENABLE(JIT)
+        if (UNLIKELY(!generateWasmToWasmStubs())) {
+            fail(makeString("Out of executable memory at stub generation"_s));
+            return;
+        }
+#endif
+    }
+
+    if (!std::exchange(m_areWasmToJSStubsCompiled, true)) {
+#if ENABLE(JIT)
+        if (UNLIKELY(!generateWasmToJSStubs())) {
+            fail(makeString("Out of executable memory at stub generation"_s));
+            return;
+        }
+#endif
+    }
+}
+
 #if ENABLE(JIT)
 
 bool EntryPlan::generateWasmToWasmStubs()
@@ -238,17 +286,8 @@ bool EntryPlan::generateWasmToWasmStubs()
             continue;
         dataLogLnIf(WasmEntryPlanInternal::verbose, "Processing import function number "_s, importFunctionIndex, ": "_s, makeString(import->module), ": "_s, makeString(import->field));
         auto binding = wasmToWasm(importFunctionIndex);
-        if (UNLIKELY(!binding)) {
-            switch (binding.error()) {
-            case BindingFailure::OutOfMemory: {
-                Locker locker { m_lock };
-                m_wasmToWasmExitStubs.resize(importFunctionIndex);
-                fail(makeString("Out of executable memory at import "_s, importIndex));
-                return false;
-            }
-            }
-            RELEASE_ASSERT_NOT_REACHED();
-        }
+        if (UNLIKELY(!binding))
+            return false;
         m_wasmToWasmExitStubs[importFunctionIndex++] = binding.value();
     }
     ASSERT(importFunctionIndex == m_wasmToWasmExitStubs.size());
@@ -261,17 +300,8 @@ bool EntryPlan::generateWasmToJSStubs()
     for (unsigned importIndex = 0; importIndex < m_moduleInformation->importFunctionCount(); ++importIndex) {
         Wasm::TypeIndex typeIndex = m_moduleInformation->importFunctionTypeIndices.at(importIndex);
         auto binding = wasmToJS(typeIndex, importIndex);
-        if (UNLIKELY(!binding)) {
-            switch (binding.error()) {
-            case BindingFailure::OutOfMemory: {
-                Locker locker { m_lock };
-                m_wasmToJSExitStubs.resize(importIndex);
-                fail(makeString("Out of executable memory at import "_s, importIndex));
-                return false;
-            }
-            }
-            RELEASE_ASSERT_NOT_REACHED();
-        }
+        if (UNLIKELY(!binding))
+            return false;
         m_wasmToJSExitStubs[importIndex] = binding.value();
     }
     return true;

Source/JavaScriptCore/wasm/WasmEntryPlan.h

@@ -125,6 +125,7 @@ class EntryPlan : public Plan, public StreamingParserClient {
     bool generateWasmToJSStubs();
     bool generateWasmToWasmStubs();
 #endif
+    void generateStubsIfNecessary() WTF_REQUIRES_LOCK(m_lock);
 
     Vector<uint8_t> m_source;
     Vector<MacroAssemblerCodeRef<WasmEntryPtrTag>> m_wasmToWasmExitStubs;
@@ -135,6 +136,8 @@ class EntryPlan : public Plan, public StreamingParserClient {
     StreamingParser m_streamingParser;
     State m_state;
 
+    bool m_areWasmToWasmStubsCompiled { false };
+    bool m_areWasmToJSStubsCompiled { false };
     const CompilerMode m_compilerMode;
     uint8_t m_numberOfActiveThreads { 0 };
     uint32_t m_currentIndex { 0 };

Source/JavaScriptCore/wasm/WasmIPIntPlan.cpp

@@ -55,6 +55,7 @@ IPIntPlan::IPIntPlan(VM& vm, Ref<ModuleInformation> info, const Ref<IPIntCallee>
     , m_callees(callees)
 {
     ASSERT(m_callees || !m_moduleInformation->functions.size());
+    m_areWasmToJSStubsCompiled = true;
     prepare();
     m_currentIndex = m_moduleInformation->functions.size();
 }
@@ -81,7 +82,7 @@ void IPIntPlan::compileFunction(uint32_t functionIndex)
     const auto& function = m_moduleInformation->functions[functionIndex];
     TypeIndex typeIndex = m_moduleInformation->internalFunctionTypeIndices[functionIndex];
     const TypeDefinition& signature = TypeInformation::get(typeIndex);
-    unsigned functionIndexSpace = m_wasmToWasmExitStubs.size() + functionIndex;
+    unsigned functionIndexSpace = m_moduleInformation->importFunctionTypeIndices.size() + functionIndex;
     ASSERT_UNUSED(functionIndexSpace, m_moduleInformation->typeIndexFromFunctionIndexSpace(functionIndexSpace) == typeIndex);
 
     m_unlinkedWasmToWasmCalls[functionIndex] = Vector<UnlinkedWasmToWasmCall>();
@@ -97,19 +98,23 @@ void IPIntPlan::compileFunction(uint32_t functionIndex)
         return;
     }
 
-    Locker locker { m_lock };
+    if (Options::useWebAssemblyTailCalls()) {
+        Locker locker { m_lock };
 
-    for (auto successor : parseAndCompileResult->get()->tailCallSuccessors())
-        addTailCallEdge(m_moduleInformation->importFunctionCount() + parseAndCompileResult->get()->functionIndex(), successor);
+        for (auto successor : parseAndCompileResult->get()->tailCallSuccessors())
+            addTailCallEdge(m_moduleInformation->importFunctionCount() + parseAndCompileResult->get()->functionIndex(), successor);
 
-    if (parseAndCompileResult->get()->tailCallClobbersInstance())
-        m_moduleInformation->addClobberingTailCall(m_moduleInformation->importFunctionCount() + parseAndCompileResult->get()->functionIndex());
+        if (parseAndCompileResult->get()->tailCallClobbersInstance())
+            m_moduleInformation->addClobberingTailCall(m_moduleInformation->importFunctionCount() + parseAndCompileResult->get()->functionIndex());
+    }
 
     m_wasmInternalFunctions[functionIndex] = WTFMove(*parseAndCompileResult);
 }
 
 void IPIntPlan::didCompleteCompilation()
 {
+    generateStubsIfNecessary();
+
 #if ENABLE(JIT)
     unsigned functionCount = m_wasmInternalFunctions.size();
     if (!m_callees && functionCount) {

Source/JavaScriptCore/wasm/WasmLLIntPlan.cpp

@@ -57,6 +57,7 @@ LLIntPlan::LLIntPlan(VM& vm, Ref<ModuleInformation> info, const Ref<LLIntCallee>
     , m_callees(callees)
 {
     ASSERT(m_callees || !m_moduleInformation->functions.size());
+    m_areWasmToJSStubsCompiled = true;
     prepare();
     m_currentIndex = m_moduleInformation->functions.size();
 }
@@ -84,7 +85,7 @@ void LLIntPlan::compileFunction(uint32_t functionIndex)
     const auto& function = m_moduleInformation->functions[functionIndex];
     TypeIndex typeIndex = m_moduleInformation->internalFunctionTypeIndices[functionIndex];
     const TypeDefinition& signature = TypeInformation::get(typeIndex);
-    unsigned functionIndexSpace = m_wasmToWasmExitStubs.size() + functionIndex;
+    unsigned functionIndexSpace = m_moduleInformation->importFunctionTypeIndices.size() + functionIndex;
     ASSERT_UNUSED(functionIndexSpace, m_moduleInformation->typeIndexFromFunctionIndexSpace(functionIndexSpace) == typeIndex);
 
     m_unlinkedWasmToWasmCalls[functionIndex] = Vector<UnlinkedWasmToWasmCall>();
@@ -100,13 +101,15 @@ void LLIntPlan::compileFunction(uint32_t functionIndex)
         return;
     }
 
-    Locker locker { m_lock };
+    if (Options::useWebAssemblyTailCalls()) {
+        Locker locker { m_lock };
 
-    for (auto successor : parseAndCompileResult->get()->tailCallSuccessors())
-        addTailCallEdge(m_moduleInformation->importFunctionCount() + parseAndCompileResult->get()->functionIndex(), successor);
+        for (auto successor : parseAndCompileResult->get()->tailCallSuccessors())
+            addTailCallEdge(m_moduleInformation->importFunctionCount() + parseAndCompileResult->get()->functionIndex(), successor);
 
-    if (parseAndCompileResult->get()->tailCallClobbersInstance())
-        m_moduleInformation->addClobberingTailCall(m_moduleInformation->importFunctionCount() + parseAndCompileResult->get()->functionIndex());
+        if (parseAndCompileResult->get()->tailCallClobbersInstance())
+            m_moduleInformation->addClobberingTailCall(m_moduleInformation->importFunctionCount() + parseAndCompileResult->get()->functionIndex());
+    }
 
     m_wasmInternalFunctions[functionIndex] = WTFMove(*parseAndCompileResult);
 }
@@ -246,6 +249,8 @@ bool LLIntPlan::makeInterpretedJSToWasmCallee(unsigned) { return false; }
 
 void LLIntPlan::didCompleteCompilation()
 {
+    generateStubsIfNecessary();
+
     unsigned functionCount = m_wasmInternalFunctions.size();
     if (!m_callees && functionCount) {
         m_calleesVector.resize(functionCount);

Source/JavaScriptCore/wasm/WasmModule.h

@@ -49,7 +49,7 @@ enum class BindingFailure;
 
 class Module : public ThreadSafeRefCounted<Module> {
 public:
-    using ValidationResult = Expected<RefPtr<Module>, String>;
+    using ValidationResult = Expected<Ref<Module>, String>;
     typedef void CallbackType(ValidationResult&&);
     using AsyncValidationCallback = RefPtr<SharedTask<CallbackType>>;
 

Source/JavaScriptCore/wasm/WasmStreamingCompiler.cpp

@@ -145,12 +145,14 @@ void StreamingCompiler::didComplete()
             VM& vm = globalObject->vm();
             auto scope = DECLARE_THROW_SCOPE(vm);
 
-            JSWebAssemblyModule* module = JSWebAssemblyModule::createStub(vm, globalObject, globalObject->webAssemblyModuleStructure(), WTFMove(result));
-            if (UNLIKELY(scope.exception())) {
+            if (UNLIKELY(!result.has_value())) {
+                throwException(globalObject, scope, createJSWebAssemblyCompileError(globalObject, vm, result.error()));
                 promise->rejectWithCaughtException(globalObject, scope);
                 return;
             }
 
+            JSWebAssemblyModule* module = JSWebAssemblyModule::create(vm, globalObject->webAssemblyModuleStructure(), WTFMove(result.value()));
+
             scope.release();
             promise->resolve(globalObject, module);
         });
@@ -165,12 +167,13 @@ void StreamingCompiler::didComplete()
             VM& vm = globalObject->vm();
             auto scope = DECLARE_THROW_SCOPE(vm);
 
-            JSWebAssemblyModule* module = JSWebAssemblyModule::createStub(vm, globalObject, globalObject->webAssemblyModuleStructure(), WTFMove(result));
-            if (UNLIKELY(scope.exception())) {
+            if (UNLIKELY(!result.has_value())) {
+                throwException(globalObject, scope, createJSWebAssemblyCompileError(globalObject, vm, result.error()));
                 promise->rejectWithCaughtException(globalObject, scope);
                 return;
             }
 
+            JSWebAssemblyModule* module = JSWebAssemblyModule::create(vm, globalObject->webAssemblyModuleStructure(), WTFMove(result.value()));
             JSWebAssembly::instantiateForStreaming(vm, globalObject, promise, module, importObject);
             if (UNLIKELY(scope.exception())) {
                 promise->rejectWithCaughtException(globalObject, scope);

Source/JavaScriptCore/wasm/js/JSWebAssembly.cpp

@@ -39,6 +39,7 @@
 #include "JSModuleNamespaceObject.h"
 #include "JSObjectInlines.h"
 #include "JSPromise.h"
+#include "JSWebAssemblyCompileError.h"
 #include "JSWebAssemblyHelpers.h"
 #include "JSWebAssemblyInstance.h"
 #include "JSWebAssemblyModule.h"
@@ -151,12 +152,15 @@ void JSWebAssembly::webAssemblyModuleValidateAsync(JSGlobalObject* globalObject,
     Wasm::Module::validateAsync(vm, WTFMove(source), createSharedTask<Wasm::Module::CallbackType>([ticket, promise, globalObject, &vm] (Wasm::Module::ValidationResult&& result) mutable {
         vm.deferredWorkTimer->scheduleWorkSoon(ticket, [promise, globalObject, result = WTFMove(result), &vm](DeferredWorkTimer::Ticket) mutable {
             auto scope = DECLARE_THROW_SCOPE(vm);
-            JSValue module = JSWebAssemblyModule::createStub(vm, globalObject, globalObject->webAssemblyModuleStructure(), WTFMove(result));
-            if (UNLIKELY(scope.exception())) {
+
+            if (UNLIKELY(!result.has_value())) {
+                throwException(globalObject, scope, createJSWebAssemblyCompileError(globalObject, vm, result.error()));
                 promise->rejectWithCaughtException(globalObject, scope);
                 return;
             }
 
+            JSValue module = JSWebAssemblyModule::create(vm, globalObject->webAssemblyModuleStructure(), WTFMove(result.value()));
+
             scope.release();
             promise->resolve(globalObject, module);
         });
@@ -242,12 +246,15 @@ static void compileAndInstantiate(VM& vm, JSGlobalObject* globalObject, JSPromis
     Wasm::Module::validateAsync(vm, WTFMove(source), createSharedTask<Wasm::Module::CallbackType>([ticket, promise, importObject, moduleKeyCell, globalObject, resolveKind, creationMode, &vm] (Wasm::Module::ValidationResult&& result) mutable {
         vm.deferredWorkTimer->scheduleWorkSoon(ticket, [promise, importObject, moduleKeyCell, globalObject, result = WTFMove(result), resolveKind, creationMode, &vm](DeferredWorkTimer::Ticket) mutable {
             auto scope = DECLARE_THROW_SCOPE(vm);
-            JSWebAssemblyModule* module = JSWebAssemblyModule::createStub(vm, globalObject, globalObject->webAssemblyModuleStructure(), WTFMove(result));
-            if (UNLIKELY(scope.exception())) {
+
+            if (UNLIKELY(!result.has_value())) {
+                throwException(globalObject, scope, createJSWebAssemblyCompileError(globalObject, vm, result.error()));
                 promise->rejectWithCaughtException(globalObject, scope);
                 return;
             }
 
+            JSWebAssemblyModule* module = JSWebAssemblyModule::create(vm, globalObject->webAssemblyModuleStructure(), WTFMove(result.value()));
+
             const Identifier moduleKey = JSValue(moduleKeyCell).toPropertyKey(globalObject);
             if (UNLIKELY(scope.exception())) {
                 promise->rejectWithCaughtException(globalObject, scope);

Source/JavaScriptCore/wasm/js/JSWebAssemblyModule.cpp

@@ -41,14 +41,9 @@ namespace JSC {
 
 const ClassInfo JSWebAssemblyModule::s_info = { "WebAssembly.Module"_s, &Base::s_info, nullptr, nullptr, CREATE_METHOD_TABLE(JSWebAssemblyModule) };
 
-JSWebAssemblyModule* JSWebAssemblyModule::createStub(VM& vm, JSGlobalObject* globalObject, Structure* structure, Wasm::Module::ValidationResult&& result)
+JSWebAssemblyModule* JSWebAssemblyModule::create(VM& vm, Structure* structure, Ref<Wasm::Module>&& result)
 {
-    auto scope = DECLARE_THROW_SCOPE(vm);
-    if (!result.has_value()) {
-        throwException(globalObject, scope, createJSWebAssemblyCompileError(globalObject, vm, result.error()));
-        return nullptr;
-    }
-    auto* module = new (NotNull, allocateCell<JSWebAssemblyModule>(vm)) JSWebAssemblyModule(vm, structure, result.value().releaseNonNull());
+    auto* module = new (NotNull, allocateCell<JSWebAssemblyModule>(vm)) JSWebAssemblyModule(vm, structure, WTFMove(result));
     module->finishCreation(vm);
     return module;
 }

Source/JavaScriptCore/wasm/js/JSWebAssemblyModule.h

@@ -63,7 +63,7 @@ class JSWebAssemblyModule final : public JSNonFinalObject {
 
     DECLARE_EXPORT_INFO;
 
-    JS_EXPORT_PRIVATE static JSWebAssemblyModule* createStub(VM&, JSGlobalObject*, Structure*, Expected<RefPtr<Wasm::Module>, String>&&);
+    JS_EXPORT_PRIVATE static JSWebAssemblyModule* create(VM&, Structure*, Ref<Wasm::Module>&&);
     static Structure* createStructure(VM&, JSGlobalObject*, JSValue);
 
     const Wasm::ModuleInformation& moduleInformation() const;