Bug 1866491 - Try to batch Ion compilation tasks to free in a single IonFreeTask. r=jonco

When we cancel Ion compilations, there can be a lot of tasks on the lazy-link list.
Instead of creating a separate `IonFreeTask` for each of those, we can add the
compilation tasks to a vector and then create a single `IonFreeTask` to free them.

On Speedometer 3 there are a few cases where we now batch more than 70 compilation
tasks into a single `IonFreeTask`.

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

js/src/jit/Ion.cpp

@@ -376,10 +376,8 @@ void jit::LinkIonScript(JSContext* cx, HandleScript calleeScript) {
     }
   }
 
-  {
-    AutoLockHelperThreadState lock;
-    FinishOffThreadTask(cx->runtime(), task, lock);
-  }
+  AutoStartIonFreeTask freeTask;
+  FinishOffThreadTask(cx->runtime(), freeTask, task);
 }
 
 uint8_t* jit::LazyLinkTopActivation(JSContext* cx,

js/src/jit/IonCompileTask.cpp

@@ -151,7 +151,7 @@ void jit::AttachFinishedCompilations(JSContext* cx) {
   MOZ_ASSERT(!rt->jitRuntime()->numFinishedOffThreadTasks());
 }
 
-void jit::FreeIonCompileTask(IonCompileTask* task) {
+static void FreeIonCompileTask(IonCompileTask* task) {
   // The task is allocated into its LifoAlloc, so destroying that will
   // destroy the task and all other data accumulated during compilation,
   // except any final codegen (which includes an assembler and needs to be
@@ -160,18 +160,27 @@ void jit::FreeIonCompileTask(IonCompileTask* task) {
   js_delete(task->alloc().lifoAlloc());
 }
 
+void jit::FreeIonCompileTasks(const IonFreeCompileTasks& tasks) {
+  MOZ_ASSERT(!tasks.empty());
+  for (auto* task : tasks) {
+    FreeIonCompileTask(task);
+  }
+}
+
 void IonFreeTask::runHelperThreadTask(AutoLockHelperThreadState& locked) {
   {
     AutoUnlockHelperThreadState unlock(locked);
-    jit::FreeIonCompileTask(task_);
+    jit::FreeIonCompileTasks(compileTasks());
   }
 
   js_delete(this);
 }
 
-void jit::FinishOffThreadTask(JSRuntime* runtime, IonCompileTask* task,
-                              const AutoLockHelperThreadState& locked) {
+void jit::FinishOffThreadTask(JSRuntime* runtime,
+                              AutoStartIonFreeTask& freeTask,
+                              IonCompileTask* task) {
   MOZ_ASSERT(runtime);
+  MOZ_ASSERT(CurrentThreadCanAccessRuntime(runtime));
 
   JSScript* script = task->script();
 
@@ -196,8 +205,9 @@ void jit::FinishOffThreadTask(JSRuntime* runtime, IonCompileTask* task,
     }
   }
 
-  // Free Ion LifoAlloc off-thread. Free on the main thread if this OOMs.
-  if (!StartOffThreadIonFree(task, locked)) {
+  // Try to free the Ion LifoAlloc off-thread. Free on the main thread if this
+  // OOMs.
+  if (!freeTask.appendCompileTask(task)) {
     FreeIonCompileTask(task);
   }
 }

js/src/jit/IonCompileTask.h

@@ -67,21 +67,23 @@ class IonCompileTask final : public HelperThreadTask,
 };
 
 class IonFreeTask : public HelperThreadTask {
+  IonFreeCompileTasks tasks_;
+
  public:
-  explicit IonFreeTask(IonCompileTask* task) : task_(task) {}
-  IonCompileTask* compileTask() { return task_; }
+  explicit IonFreeTask(IonFreeCompileTasks&& tasks) : tasks_(std::move(tasks)) {
+    MOZ_ASSERT(!tasks_.empty());
+  }
+
+  const IonFreeCompileTasks& compileTasks() const { return tasks_; }
 
   ThreadType threadType() override { return THREAD_TYPE_ION_FREE; }
   void runHelperThreadTask(AutoLockHelperThreadState& locked) override;
-
- private:
-  IonCompileTask* task_;
 };
 
 void AttachFinishedCompilations(JSContext* cx);
-void FinishOffThreadTask(JSRuntime* runtime, IonCompileTask* task,
-                         const AutoLockHelperThreadState& lock);
-void FreeIonCompileTask(IonCompileTask* task);
+void FinishOffThreadTask(JSRuntime* runtime, AutoStartIonFreeTask& freeTask,
+                         IonCompileTask* task);
+void FreeIonCompileTasks(const IonFreeCompileTasks& tasks);
 
 }  // namespace jit
 }  // namespace js

js/src/vm/HelperThreadState.h

@@ -412,7 +412,7 @@ class GlobalHelperThreadState {
  public:
   bool submitTask(wasm::UniqueTier2GeneratorTask task);
   bool submitTask(wasm::CompileTask* task, wasm::CompileMode mode);
-  bool submitTask(UniquePtr<jit::IonFreeTask> task,
+  bool submitTask(UniquePtr<jit::IonFreeTask>&& task,
                   const AutoLockHelperThreadState& lock);
   bool submitTask(jit::IonCompileTask* task,
                   const AutoLockHelperThreadState& locked);

js/src/vm/HelperThreads.cpp

@@ -258,19 +258,30 @@ bool GlobalHelperThreadState::submitTask(
   return true;
 }
 
-bool js::StartOffThreadIonFree(jit::IonCompileTask* task,
-                               const AutoLockHelperThreadState& lock) {
-  js::UniquePtr<jit::IonFreeTask> freeTask =
-      js::MakeUnique<jit::IonFreeTask>(task);
-  if (!freeTask) {
-    return false;
+js::AutoStartIonFreeTask::~AutoStartIonFreeTask() {
+  if (tasks_.empty()) {
+    return;
   }
 
-  return HelperThreadState().submitTask(std::move(freeTask), lock);
+  auto freeTask = js::MakeUnique<jit::IonFreeTask>(std::move(tasks_));
+  if (!freeTask) {
+    // Free compilation data on the main thread instead.
+    MOZ_ASSERT(!tasks_.empty(), "shouldn't have moved tasks_ on OOM");
+    jit::FreeIonCompileTasks(tasks_);
+    return;
+  }
+
+  AutoLockHelperThreadState lock;
+  if (!HelperThreadState().submitTask(std::move(freeTask), lock)) {
+    // If submitTask OOMs, then freeTask hasn't been moved so we can still use
+    // its task list.
+    jit::FreeIonCompileTasks(freeTask->compileTasks());
+  }
 }
 
 bool GlobalHelperThreadState::submitTask(
-    UniquePtr<jit::IonFreeTask> task, const AutoLockHelperThreadState& locked) {
+    UniquePtr<jit::IonFreeTask>&& task,
+    const AutoLockHelperThreadState& locked) {
   MOZ_ASSERT(isInitialized(locked));
 
   if (!ionFreeList(locked).append(std::move(task))) {
@@ -343,64 +354,72 @@ static bool IonCompileTaskMatches(const CompilationSelector& selector,
   return selector.match(TaskMatches{task});
 }
 
-static void CancelOffThreadIonCompileLocked(const CompilationSelector& selector,
-                                            AutoLockHelperThreadState& lock) {
-  if (jit::IsPortableBaselineInterpreterEnabled()) {
+void js::CancelOffThreadIonCompile(const CompilationSelector& selector) {
+  if (!JitDataStructuresExist(selector)) {
     return;
   }
 
-  if (!HelperThreadState().isInitialized(lock)) {
+  if (jit::IsPortableBaselineInterpreterEnabled()) {
     return;
   }
 
   MOZ_ASSERT(GetSelectorRuntime(selector)->jitRuntime() != nullptr);
 
-  /* Cancel any pending entries for which processing hasn't started. */
-  GlobalHelperThreadState::IonCompileTaskVector& worklist =
-      HelperThreadState().ionWorklist(lock);
-  for (size_t i = 0; i < worklist.length(); i++) {
-    jit::IonCompileTask* task = worklist[i];
-    if (IonCompileTaskMatches(selector, task)) {
-      // Once finished, tasks are added to a Linked list which is
-      // allocated with the IonCompileTask class. The IonCompileTask is
-      // allocated in the LifoAlloc so we need the LifoAlloc to be mutable.
-      worklist[i]->alloc().lifoAlloc()->setReadWrite();
+  AutoStartIonFreeTask freeTask;
 
-      FinishOffThreadIonCompile(task, lock);
-      HelperThreadState().remove(worklist, &i);
+  {
+    AutoLockHelperThreadState lock;
+    if (!HelperThreadState().isInitialized(lock)) {
+      return;
     }
-  }
 
-  /* Wait for in progress entries to finish up. */
-  bool cancelled;
-  do {
-    cancelled = false;
-    for (auto* helper : HelperThreadState().helperTasks(lock)) {
-      if (!helper->is<jit::IonCompileTask>()) {
-        continue;
-      }
+    /* Cancel any pending entries for which processing hasn't started. */
+    GlobalHelperThreadState::IonCompileTaskVector& worklist =
+        HelperThreadState().ionWorklist(lock);
+    for (size_t i = 0; i < worklist.length(); i++) {
+      jit::IonCompileTask* task = worklist[i];
+      if (IonCompileTaskMatches(selector, task)) {
+        // Once finished, tasks are added to a Linked list which is
+        // allocated with the IonCompileTask class. The IonCompileTask is
+        // allocated in the LifoAlloc so we need the LifoAlloc to be mutable.
+        worklist[i]->alloc().lifoAlloc()->setReadWrite();
 
-      jit::IonCompileTask* ionCompileTask = helper->as<jit::IonCompileTask>();
-      if (IonCompileTaskMatches(selector, ionCompileTask)) {
-        ionCompileTask->mirGen().cancel();
-        cancelled = true;
+        FinishOffThreadIonCompile(task, lock);
+        HelperThreadState().remove(worklist, &i);
       }
     }
-    if (cancelled) {
-      HelperThreadState().wait(lock);
-    }
-  } while (cancelled);
 
-  /* Cancel code generation for any completed entries. */
-  GlobalHelperThreadState::IonCompileTaskVector& finished =
-      HelperThreadState().ionFinishedList(lock);
-  for (size_t i = 0; i < finished.length(); i++) {
-    jit::IonCompileTask* task = finished[i];
-    if (IonCompileTaskMatches(selector, task)) {
-      JSRuntime* rt = task->script()->runtimeFromAnyThread();
-      rt->jitRuntime()->numFinishedOffThreadTasksRef(lock)--;
-      jit::FinishOffThreadTask(rt, task, lock);
-      HelperThreadState().remove(finished, &i);
+    /* Wait for in progress entries to finish up. */
+    bool cancelled;
+    do {
+      cancelled = false;
+      for (auto* helper : HelperThreadState().helperTasks(lock)) {
+        if (!helper->is<jit::IonCompileTask>()) {
+          continue;
+        }
+
+        jit::IonCompileTask* ionCompileTask = helper->as<jit::IonCompileTask>();
+        if (IonCompileTaskMatches(selector, ionCompileTask)) {
+          ionCompileTask->mirGen().cancel();
+          cancelled = true;
+        }
+      }
+      if (cancelled) {
+        HelperThreadState().wait(lock);
+      }
+    } while (cancelled);
+
+    /* Cancel code generation for any completed entries. */
+    GlobalHelperThreadState::IonCompileTaskVector& finished =
+        HelperThreadState().ionFinishedList(lock);
+    for (size_t i = 0; i < finished.length(); i++) {
+      jit::IonCompileTask* task = finished[i];
+      if (IonCompileTaskMatches(selector, task)) {
+        JSRuntime* rt = task->script()->runtimeFromAnyThread();
+        rt->jitRuntime()->numFinishedOffThreadTasksRef(lock)--;
+        jit::FinishOffThreadTask(rt, freeTask, task);
+        HelperThreadState().remove(finished, &i);
+      }
     }
   }
 
@@ -411,21 +430,12 @@ static void CancelOffThreadIonCompileLocked(const CompilationSelector& selector,
   while (task) {
     jit::IonCompileTask* next = task->getNext();
     if (IonCompileTaskMatches(selector, task)) {
-      jit::FinishOffThreadTask(runtime, task, lock);
+      jit::FinishOffThreadTask(runtime, freeTask, task);
     }
     task = next;
   }
 }
 
-void js::CancelOffThreadIonCompile(const CompilationSelector& selector) {
-  if (!JitDataStructuresExist(selector)) {
-    return;
-  }
-
-  AutoLockHelperThreadState lock;
-  CancelOffThreadIonCompileLocked(selector, lock);
-}
-
 #ifdef DEBUG
 bool js::HasOffThreadIonCompile(Zone* zone) {
   if (jit::IsPortableBaselineInterpreterEnabled()) {
@@ -817,7 +827,7 @@ void GlobalHelperThreadState::finish(AutoLockHelperThreadState& lock) {
   while (!freeList.empty()) {
     UniquePtr<jit::IonFreeTask> task = std::move(freeList.back());
     freeList.popBack();
-    jit::FreeIonCompileTask(task->compileTask());
+    jit::FreeIonCompileTasks(task->compileTasks());
   }
 }
 
@@ -998,8 +1008,9 @@ void GlobalHelperThreadState::addSizeOfIncludingThis(
     htStats.ionCompileTask += task->sizeOfExcludingThis(mallocSizeOf);
   }
   for (const auto& task : ionFreeList_) {
-    htStats.ionCompileTask +=
-        task->compileTask()->sizeOfExcludingThis(mallocSizeOf);
+    for (auto* compileTask : task->compileTasks()) {
+      htStats.ionCompileTask += compileTask->sizeOfExcludingThis(mallocSizeOf);
+    }
   }
 
   // Report wasm::CompileTasks on wait lists

js/src/vm/HelperThreads.h

@@ -13,8 +13,10 @@
 
 #include "mozilla/Variant.h"
 
+#include "js/AllocPolicy.h"
 #include "js/shadow/Zone.h"
 #include "js/UniquePtr.h"
+#include "js/Vector.h"
 #include "threading/LockGuard.h"
 #include "threading/Mutex.h"
 #include "wasm/WasmConstants.h"
@@ -49,6 +51,7 @@ class GCRuntime;
 namespace jit {
 class IonCompileTask;
 class IonFreeTask;
+using IonFreeCompileTasks = Vector<IonCompileTask*, 8, SystemAllocPolicy>;
 }  // namespace jit
 
 namespace wasm {
@@ -140,15 +143,21 @@ bool StartOffThreadPromiseHelperTask(PromiseHelperTask* task);
 bool StartOffThreadIonCompile(jit::IonCompileTask* task,
                               const AutoLockHelperThreadState& lock);
 
-/*
- * Schedule deletion of Ion compilation data.
- */
-bool StartOffThreadIonFree(jit::IonCompileTask* task,
-                           const AutoLockHelperThreadState& lock);
-
 void FinishOffThreadIonCompile(jit::IonCompileTask* task,
                                const AutoLockHelperThreadState& lock);
 
+// RAII class to submit an IonFreeTask for a list of Ion compilation tasks.
+class MOZ_RAII AutoStartIonFreeTask {
+  jit::IonFreeCompileTasks tasks_;
+
+ public:
+  ~AutoStartIonFreeTask();
+
+  [[nodiscard]] bool appendCompileTask(jit::IonCompileTask* task) {
+    return tasks_.append(task);
+  }
+};
+
 struct ZonesInState {
   JSRuntime* runtime;
   JS::shadow::Zone::GCState state;