1329018 Wasm: Move validation to helper threads, also 1317678

1329018 - Part 1: Wasm: Add plumbing for compile threads to report validation errors.

1329018 - Part 2: Wasm: Move function validation to helper threads.

Also: 1317678 - don't unconditionally unroll stack probe loop

js/src/jit/x64/MacroAssembler-x64.cpp

@@ -427,7 +427,7 @@ MacroAssembler::subFromStackPtr(Imm32 imm32)
         // it as we go.
         //
         // When the amount is quite large, which it can be, we emit an actual loop, in order
-        // to keep the function prologue compact. Compactness is a requirement for e.g.
+        // to keep the function prologue compact.  Compactness is a requirement for eg
         // Wasm's CodeRange data structure, which can encode only 8-bit offsets.
         uint32_t amountLeft = imm32.value;
         uint32_t fullPages = amountLeft / 4096;

js/src/jit/x86/MacroAssembler-x86.cpp

@@ -320,7 +320,7 @@ MacroAssembler::subFromStackPtr(Imm32 imm32)
         // it as we go.
         //
         // When the amount is quite large, which it can be, we emit an actual loop, in order
-        // to keep the function prologue compact. Compactness is a requirement for e.g.
+        // to keep the function prologue compact.  Compactness is a requirement for eg
         // Wasm's CodeRange data structure, which can encode only 8-bit offsets.
         uint32_t amountLeft = imm32.value;
         uint32_t fullPages = amountLeft / 4096;

js/src/vm/HelperThreads.cpp

@@ -1470,11 +1470,12 @@ HelperThread::handleWasmWorkload(AutoLockHelperThreadState& locked)
 
     currentTask.emplace(HelperThreadState().wasmWorklist(locked).popCopy());
     bool success = false;
+    UniqueChars error;
 
     wasm::IonCompileTask* task = wasmTask();
     {
         AutoUnlockHelperThreadState unlock(locked);
-        success = wasm::CompileFunction(task);
+        success = wasm::CompileFunction(task, &error);
     }
 
     // On success, try to move work to the finished list.
@@ -1482,8 +1483,10 @@ HelperThread::handleWasmWorkload(AutoLockHelperThreadState& locked)
         success = HelperThreadState().wasmFinishedList(locked).append(task);
 
     // On failure, note the failure for harvesting by the parent.
-    if (!success)
+    if (!success) {
         HelperThreadState().noteWasmFailure(locked);
+        HelperThreadState().setWasmError(locked, Move(error));
+	}
 
     // Notify the main thread in case it's waiting.
     HelperThreadState().notifyAll(GlobalHelperThreadState::CONSUMER, locked);

js/src/vm/HelperThreads.h

@@ -220,10 +220,17 @@ class GlobalHelperThreadState
         uint32_t n = numWasmFailedJobs;
         numWasmFailedJobs = 0;
         return n;
+    }
+	UniqueChars harvestWasmError(const AutoLockHelperThreadState&) {
+        return Move(firstWasmError);
     }
     void noteWasmFailure(const AutoLockHelperThreadState&) {
         // Be mindful to signal the main thread after calling this function.
         numWasmFailedJobs++;
+    }
+	void setWasmError(const AutoLockHelperThreadState&, UniqueChars error) {
+        if (!firstWasmError)
+          firstWasmError = Move(error);
     }
     bool wasmFailed(const AutoLockHelperThreadState&) {
         return bool(numWasmFailedJobs);
@@ -245,6 +252,12 @@ class GlobalHelperThreadState
      */
     uint32_t numWasmFailedJobs;
 
+    /*
+     * Error string from wasm validation. Arbitrarily choose to keep the first one that gets
+     * reported. Nondeterministic if multiple threads have errors.
+     */
+    UniqueChars firstWasmError;
+
   public:
     JSScript* finishScriptParseTask(JSContext* cx, void* token);
     JSScript* finishScriptDecodeTask(JSContext* cx, void* token);
@@ -400,7 +413,15 @@ PauseCurrentHelperThread();
 
 /* Perform MIR optimization and LIR generation on a single function. */
 bool
-StartOffThreadWasmCompile(wasm::IonCompileTask* task);
+StartOffThreadWasmCompile(wasm::CompileTask* task);
+
+namespace wasm {
+
+// Performs MIR optimization and LIR generation on one or several functions.
+[[nodiscard]]  bool
+CompileFunction(CompileTask* task, UniqueChars* error);
+
+}
 
 /*
  * If helper threads are available, start executing the given PromiseTask on a

js/src/wasm/AsmJS.cpp

@@ -1465,6 +1465,7 @@ class MOZ_STACK_CLASS ModuleValidator
         importMap_(cx),
         arrayViews_(cx),
         atomicsPresent_(false),
+        mg_(nullptr),
         errorString_(nullptr),
         errorOffset_(UINT32_MAX),
         errorOverRecursed_(false)

js/src/wasm/WasmBaselineCompile.cpp

@@ -8251,12 +8251,11 @@ js::wasm::BaselineCanCompile(const FunctionGenerator* fg)
         return false;
 #endif
 
-#if defined(JS_CODEGEN_X64) || defined(JS_CODEGEN_X86) || defined(JS_CODEGEN_ARM) || \
-    defined(JS_CODEGEN_LOONGARCH64)
-    if (fg->usesAtomics())
-        return false;
-
-    if (fg->usesSimd())
+#if defined(JS_CODEGEN_X64) || defined(JS_CODEGEN_X86) || defined(JS_CODEGEN_ARM)
+    // AsmJS code may use SIMD or atomics, which Baseline doesn't currently
+    // handle. Since we haven't yet validated the function, we don't know
+    // whether it actually uses those features. Assume the worst.
+    if (fg->isAsmJS())
         return false;
 
     return true;
@@ -8266,14 +8265,19 @@ js::wasm::BaselineCanCompile(const FunctionGenerator* fg)
 }
 
 bool
-js::wasm::BaselineCompileFunction(IonCompileTask* task)
+js::wasm::BaselineCompileFunction(CompileTask* task, FuncCompileUnit* unit, UniqueChars *error)
 {
     MOZ_ASSERT(task->mode() == IonCompileTask::CompileMode::Baseline);
 
-    const FuncBytes& func = task->func();
-    FuncCompileResults& results = task->results();
+    const FuncBytes& func = unit->func();
+    uint32_t bodySize = func.bytes().length();
 
-    Decoder d(func.bytes());
+    Decoder d(func.bytes(), error);
+
+    if (!ValidateFunctionBody(task->env(), func.index(), bodySize, d))
+        return false;
+
+    d.rollbackPosition(d.begin());
 
     // Build the local types vector.
 

js/src/wasm/WasmBaselineCompile.h

@@ -18,7 +18,7 @@
 #ifndef asmjs_wasm_baseline_compile_h
 #define asmjs_wasm_baseline_compile_h
 
-#include "wasm/WasmIonCompile.h"
+#include "wasm/WasmTypes.h"
 
 namespace js {
 namespace wasm {
@@ -39,7 +39,7 @@ BaselineCanCompile(const FunctionGenerator* fg);
 
 // Generate adequate code quickly.
 bool
-BaselineCompileFunction(IonCompileTask* task);
+BaselineCompileFunction(CompileTask* task, FuncCompileUnit* unit, UniqueChars* error);
 
 } // namespace wasm
 } // namespace js

js/src/wasm/WasmCompile.cpp

@@ -35,22 +35,17 @@ DecodeFunctionBody(Decoder& d, ModuleGenerator& mg, uint32_t funcIndex)
     if (!d.readVarU32(&bodySize))
         return d.fail("expected number of function body bytes");
 
-    if (d.bytesRemain() < bodySize)
-        return d.fail("function body length too big");
-
-    const uint8_t* bodyBegin = d.currentPosition();
     const size_t offsetInModule = d.currentOffset();
 
+    // Skip over the function body; we'll validate it later.
+    const uint8_t* bodyBegin;
+    if (!d.readBytes(bodySize, &bodyBegin))
+        return d.fail("function body length too big");
+
     FunctionGenerator fg;
     if (!mg.startFuncDef(offsetInModule, &fg))
         return false;
 
-    if (!ValidateFunctionBody(mg.env(), funcIndex, d))
-        return false;
-
-    if (d.currentPosition() != bodyBegin + bodySize)
-        return d.fail("function body length mismatch");
-
     if (!fg.bytes().resize(bodySize))
         return false;
 
@@ -118,7 +113,7 @@ wasm::Compile(const ShareableBytes& bytecode, const CompileArgs& args, UniqueCha
     if (!DecodeModuleEnvironment(d, env.get()))
         return nullptr;
 
-    ModuleGenerator mg;
+    ModuleGenerator mg(error);
     if (!mg.init(Move(env), args))
         return nullptr;
 

js/src/wasm/WasmGenerator.cpp

@@ -44,8 +44,9 @@ static const unsigned GENERATOR_LIFO_DEFAULT_CHUNK_SIZE = 4 * 1024;
 static const unsigned COMPILATION_LIFO_DEFAULT_CHUNK_SIZE = 64 * 1024;
 static const uint32_t BAD_CODE_RANGE = UINT32_MAX;
 
-ModuleGenerator::ModuleGenerator()
+ModuleGenerator::ModuleGenerator(UniqueChars* error)
   : alwaysBaseline_(false),
+    error_(error),
     numSigs_(0),
     numTables_(0),
     lifo_(GENERATOR_LIFO_DEFAULT_CHUNK_SIZE),
@@ -238,9 +239,13 @@ ModuleGenerator::finishOutstandingTask()
         while (true) {
             MOZ_ASSERT(outstanding_ > 0);
 
-            if (HelperThreadState().wasmFailed(lock))
+            if (HelperThreadState().wasmFailed(lock)) {
+                if (error_) {
+                  MOZ_ASSERT(!*error_, "Should have stopped earlier");
+                  *error_ = Move(HelperThreadState().harvestWasmError(lock));
+                }
                 return false;
-
+            }
             if (!HelperThreadState().wasmFinishedList(lock).empty()) {
                 outstanding_--;
                 task = HelperThreadState().wasmFinishedList(lock).popCopy();
@@ -880,6 +885,32 @@ ModuleGenerator::startFuncDef(uint32_t lineOrBytecode, FunctionGenerator* fg)
     return true;
 }
 
+bool
+ModuleGenerator::launchBatchCompile()
+{
+    MOZ_ASSERT(currentTask_);
+
+    size_t numBatchedFuncs = currentTask_->units().length();
+    MOZ_ASSERT(numBatchedFuncs);
+
+   if (parallel_) {
+        if (!StartOffThreadWasmCompile(currentTask_))
+            return false;
+        outstanding_++;
+    } else {
+        if (!CompileFunction(currentTask_, error_))
+            return false;
+        if (!finishTask(currentTask_))
+            return false;
+    }
+
+    currentTask_ = nullptr;
+    batchedBytecode_ = 0;
+
+    numFinishedFuncDefs_ += numBatchedFuncs;
+    return true;
+}
+
 bool
 ModuleGenerator::finishFuncDef(uint32_t funcIndex, FunctionGenerator* fg)
 {
@@ -1116,3 +1147,25 @@ ModuleGenerator::finish(const ShareableBytes& bytecode)
                                        *metadata_,
                                        bytecode));
 }
+
+bool
+wasm::CompileFunction(CompileTask* task, UniqueChars* error)
+{
+    TraceLoggerThread* logger = TraceLoggerForCurrentThread();
+    AutoTraceLog logCompile(logger, TraceLogger_WasmCompilation);
+
+    for (FuncCompileUnit& unit : task->units()) {
+        switch (unit.mode()) {
+          case CompileMode::Ion:
+            if (!IonCompileFunction(task, &unit, error))
+                return false;
+            break;
+          case CompileMode::Baseline:
+            if (!BaselineCompileFunction(task, &unit, error))
+                return false;
+            break;
+        }
+    }
+
+    return true;
+}

js/src/wasm/WasmGenerator.h

@@ -205,6 +205,7 @@ class MOZ_STACK_CLASS ModuleGenerator
 
     // Constant parameters
     bool                            alwaysBaseline_;
+    UniqueChars*                    error_;
 
     // Data that is moved into the result of finish()
     Assumptions                     assumptions_;
@@ -258,7 +259,7 @@ private:
     [[nodiscard]] bool launchBatchCompile();
 
   public:
-    explicit ModuleGenerator();
+    explicit ModuleGenerator(UniqueChars* error);
     ~ModuleGenerator();
 
     [[nodiscard]] bool init(UniqueModuleEnvironment env, const CompileArgs& args,
@@ -349,6 +350,11 @@ class MOZ_STACK_CLASS FunctionGenerator
         usesAtomics_ = true;
     }
 
+    bool isAsmJS() const {
+      return m_->isAsmJS();
+    }
+
+
     Bytes& bytes() {
         return bytes_;
     }

js/src/wasm/WasmIonCompile.cpp

@@ -2857,14 +2857,23 @@ EmitExpr(FunctionCompiler& f)
 }
 
 bool
-wasm::IonCompileFunction(IonCompileTask* task)
+wasm::IonCompileFunction(CompileTask* task, FuncCompileUnit* unit, UniqueChars* error)
 {
     MOZ_ASSERT(task->mode() == IonCompileTask::CompileMode::Ion);
 
     const FuncBytes& func = unit->func();
     const ModuleEnvironment& env = task->env();
+    uint32_t bodySize = func.bytes().length();
+
+    Decoder d(func.bytes(), error);
+
+    if (!env.isAsmJS()) {
+        if (!ValidateFunctionBody(task->env(), func.index(), bodySize, d))
+          return false;
+
+        d.rollbackPosition(d.begin());
+    }
 
-    Decoder d(func.bytes());
 
     // Build the local types vector.
 

js/src/wasm/WasmIonCompile.h

@@ -21,6 +21,8 @@
 #include "jit/MacroAssembler.h"
 #include "wasm/WasmTypes.h"
 
+#include "wasm/WasmTypes.h"
+
 namespace js {
 namespace wasm {
 
@@ -30,127 +32,8 @@ typedef Vector<jit::MIRType, 8, SystemAllocPolicy> MIRTypeVector;
 typedef jit::ABIArgIter<MIRTypeVector> ABIArgMIRTypeIter;
 typedef jit::ABIArgIter<ValTypeVector> ABIArgValTypeIter;
 
-// The FuncBytes class represents a single, concurrently-compilable function.
-// A FuncBytes object is composed of the wasm function body bytes along with the
-// ambient metadata describing the function necessary to compile it.
-
-class FuncBytes
-{
-    Bytes            bytes_;
-    uint32_t         index_;
-    const SigWithId& sig_;
-    uint32_t         lineOrBytecode_;
-    Uint32Vector     callSiteLineNums_;
-
-  public:
-    FuncBytes(Bytes&& bytes,
-              uint32_t index,
-              const SigWithId& sig,
-              uint32_t lineOrBytecode,
-              Uint32Vector&& callSiteLineNums)
-      : bytes_(Move(bytes)),
-        index_(index),
-        sig_(sig),
-        lineOrBytecode_(lineOrBytecode),
-        callSiteLineNums_(Move(callSiteLineNums))
-    {}
-
-    Bytes& bytes() { return bytes_; }
-    const Bytes& bytes() const { return bytes_; }
-    uint32_t index() const { return index_; }
-    const SigWithId& sig() const { return sig_; }
-    uint32_t lineOrBytecode() const { return lineOrBytecode_; }
-    const Uint32Vector& callSiteLineNums() const { return callSiteLineNums_; }
-};
-
-typedef UniquePtr<FuncBytes> UniqueFuncBytes;
-
-// The FuncCompileResults class contains the results of compiling a single
-// function body, ready to be merged into the whole-module MacroAssembler.
-
-class FuncCompileResults
-{
-    jit::TempAllocator alloc_;
-    jit::MacroAssembler masm_;
-    FuncOffsets offsets_;
-
-    FuncCompileResults(const FuncCompileResults&) = delete;
-    FuncCompileResults& operator=(const FuncCompileResults&) = delete;
-
-  public:
-    explicit FuncCompileResults(LifoAlloc& lifo)
-      : alloc_(&lifo),
-        masm_(jit::MacroAssembler::WasmToken(), alloc_)
-    {}
-
-    jit::TempAllocator& alloc() { return alloc_; }
-    jit::MacroAssembler& masm() { return masm_; }
-    FuncOffsets& offsets() { return offsets_; }
-};
-
-// An IonCompileTask represents the task of compiling a single function body. An
-// IonCompileTask is filled with the wasm code to be compiled on the main
-// validation thread, sent off to an Ion compilation helper thread which creates
-// the FuncCompileResults, and finally sent back to the validation thread. To
-// save time allocating and freeing memory, IonCompileTasks are reset() and
-// reused.
-
-class IonCompileTask
-{
-  public:
-    enum class CompileMode { None, Baseline, Ion };
-
-  private:
-    const ModuleGeneratorData& mg_;
-    LifoAlloc                  lifo_;
-    UniqueFuncBytes            func_;
-    CompileMode                mode_;
-    Maybe<FuncCompileResults>  results_;
-
-    IonCompileTask(const IonCompileTask&) = delete;
-    IonCompileTask& operator=(const IonCompileTask&) = delete;
-
-  public:
-    IonCompileTask(const ModuleGeneratorData& mg, size_t defaultChunkSize)
-      : mg_(mg), lifo_(defaultChunkSize), func_(nullptr), mode_(CompileMode::None)
-    {}
-    LifoAlloc& lifo() {
-        return lifo_;
-    }
-    const ModuleGeneratorData& mg() const {
-        return mg_;
-    }
-    void init(UniqueFuncBytes func, CompileMode mode) {
-        MOZ_ASSERT(!func_);
-        func_ = Move(func);
-        results_.emplace(lifo_);
-        mode_ = mode;
-    }
-    CompileMode mode() const {
-        return mode_;
-    }
-    const FuncBytes& func() const {
-        MOZ_ASSERT(func_);
-        return *func_;
-    }
-    FuncCompileResults& results() {
-        return *results_;
-    }
-    void reset(Bytes* recycled) {
-        if (func_)
-            *recycled = Move(func_->bytes());
-        func_.reset(nullptr);
-        results_.reset();
-        lifo_.releaseAll();
-        mode_ = CompileMode::None;
-    }
-};
-
-MOZ_MUST_USE bool
-IonCompileFunction(IonCompileTask* task);
-
-bool
-CompileFunction(IonCompileTask* task);
+[[nodiscard]] bool
+IonCompileFunction(CompileTask* task, FuncCompileUnit* unit, UniqueChars* error);
 
 } // namespace wasm
 } // namespace js

js/src/wasm/WasmValidate.cpp

@@ -664,13 +664,16 @@ DecodeFunctionBodyExprs(FunctionDecoder& f)
 }
 
 bool
-wasm::ValidateFunctionBody(const ModuleEnvironment& env, uint32_t funcIndex, Decoder& d)
+wasm::ValidateFunctionBody(const ModuleEnvironment& env, uint32_t funcIndex, uint32_t bodySize,
+                           Decoder& d)
 {
     ValTypeVector locals;
     const Sig& sig = *env.funcSigs[funcIndex];
     if (!locals.appendAll(sig.args()))
         return false;
 
+    const uint8_t* bodyBegin = d.currentPosition();
+
     if (!DecodeLocalEntries(d, ModuleKind::Wasm, &locals))
         return false;
 
@@ -682,7 +685,13 @@ wasm::ValidateFunctionBody(const ModuleEnvironment& env, uint32_t funcIndex, Dec
     if (!DecodeFunctionBodyExprs(f))
         return false;
 
-    return f.iter().readFunctionEnd();
+    if (!f.iter().readFunctionEnd())
+        return false;
+
+    if (d.currentPosition() != bodyBegin + bodySize)
+        return d.fail("function body length mismatch");
+
+    return true;
 }
 // Section macros.
 
@@ -1449,14 +1458,9 @@ DecodeFunctionBody(Decoder& d, const ModuleEnvironment& env, uint32_t funcIndex)
     if (d.bytesRemain() < bodySize)
         return d.fail("function body length too big");
 
-    const uint8_t* bodyBegin = d.currentPosition();
-
-    if (!ValidateFunctionBody(env, funcIndex, d))
+    if (!ValidateFunctionBody(env, funcIndex, bodySize, d))
         return false;
 
-    if (d.currentPosition() != bodyBegin + bodySize)
-        return d.fail("function body length mismatch");
-
     return true;
 }
 

js/src/wasm/WasmValidate.h

@@ -612,7 +612,8 @@ DecodeLocalEntries(Decoder& d, ModuleKind kind, ValTypeVector* locals);
 DecodeModuleEnvironment(Decoder& d, ModuleEnvironment* env);
 
 [[nodiscard]] bool
-ValidateFunctionBody(const ModuleEnvironment& env, uint32_t funcIndex, Decoder& d);
+ValidateFunctionBody(const ModuleEnvironment& env, uint32_t funcIndex, uint32_t bodySize,
+                     Decoder& d);
 
 [[nodiscard]] bool
 DecodeModuleTail(Decoder& d, ModuleEnvironment* env);