[lld][common][lld-macho] Support per-thread allocators and StringSavers

[oontvoo]

Summary

Details:

(draft for discussion only | benchmarking both envoy and chromium showed no difference in perf)
branched from D122922, with the difference being:

• keep TLS Context ptr
• need a lock only shortly when pushing the context into the vector

[oontvoo]

s/__thread/thread_local

[oontvoo]

folded into orig patch - closing this

lld/Common/CommonLinkerContext.cpp

 }
 
 CommonLinkerContext::~CommonLinkerContext() {
   assert(lctx);
   // Explicitly call the destructors since we created the objects with placement
   // new in SpecificAlloc::create().
   for (auto &it : instances)
     it.second->~SpecificAllocBase();
+
+  for (PerThreadContext **context : perThreadContexts) {
+    for (auto &instance : (*context)->perThreadInstances)
+      instance.second->~SpecificAllocBase();
+    delete *context;
+    // Reset the tls variable 's value to sentinel in case we reuse the threads.
+    *context = nullptr;
+  }
   lctx = nullptr;
 }
 
 CommonLinkerContext &lld::commonContext() {
   assert(lctx);
   return *lctx;
 }
 
 bool lld::hasContext() { return lctx != nullptr; }
 
 void CommonLinkerContext::destroy() {
   if (lctx == nullptr)
     return;
   delete lctx;
+  lctx = nullptr;
+}
+
+llvm::StringSaver &lld::perThreadSaver() {
+  return commonContext().perThreadContext()->perThreadSaver;
 }

lld/Common/Memory.cpp

 //===- Memory.cpp ---------------------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
-#include "lld/Common/Memory.h"
+#include <atomic>
+
 #include "lld/Common/CommonLinkerContext.h"
+#include "lld/Common/Memory.h"
 
 using namespace llvm;
 using namespace lld;
 
+namespace {
+// TODO: may need to move these to CommonLinkerContext class when we support
+// multiple processes.
+// Pointer to this thread's context.
+thread_local PerThreadContext *threadContext = nullptr;
+} // namespace
+
 SpecificAllocBase *
 lld::SpecificAllocBase::getOrCreate(void *tag, size_t size, size_t align,
                                     SpecificAllocBase *(&creator)(void *)) {
   auto &instances = context().instances;
   auto &instance = instances[tag];
   if (instance == nullptr) {
     void *storage = context().bAlloc.Allocate(size, align);
     instance = creator(storage);
   }
   return instance;
 }
+
+PerThreadContext *CommonLinkerContext::perThreadContext() {
+  if (!threadContext) {
+    // Context didn't exist yet for this thread, so create a new one.
+    threadContext = new PerThreadContext;
+    llvm::sys::ScopedWriter lock(contextMutex);
+    perThreadContexts.push_back(&threadContext);
+  }
+  return threadContext;
+}
+
+SpecificAllocBase *lld::SpecificAllocBase::getOrCreatePerThread(
+    void *tag, size_t size, size_t align,
+    SpecificAllocBase *(&creator)(void *)) {
+  PerThreadContext *threadContext = context().perThreadContext();
+  auto &instances = threadContext->perThreadInstances;
+  auto &instance = instances[tag];
+  if (instance == nullptr) {
+    void *storage = threadContext->bAlloc.Allocate(size, align);
+    instance = creator(storage);
+  }
+  return instance;
+}

lld/include/lld/Common/CommonLinkerContext.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLD_COMMON_COMMONLINKINGCONTEXT_H
 #define LLD_COMMON_COMMONLINKINGCONTEXT_H
 
 #include "lld/Common/ErrorHandler.h"
 #include "lld/Common/Memory.h"
+#include "llvm/Support/RWMutex.h"
 #include "llvm/Support/StringSaver.h"
+#include <array>
+#include <limits>
 
 namespace llvm {
 class raw_ostream;
 } // namespace llvm
 
 namespace lld {
 struct SpecificAllocBase;
+
+// Max number of threads ever created in a process.
+constexpr uint32_t MAX_THREADS = std::numeric_limits<uint32_t>::max() - 1;
+
+// TODO Maybe just use the CommonLinkerContext class - but I dont want
+// unnecessary nesting of the mutexes
+struct PerThreadContext {
+  llvm::BumpPtrAllocator bAlloc;
+  llvm::DenseMap<void *, SpecificAllocBase *> perThreadInstances;
+  llvm::StringSaver perThreadSaver{bAlloc};
+};
+
 class CommonLinkerContext {
 public:
   CommonLinkerContext();
   virtual ~CommonLinkerContext();
 
   static void destroy();
 
   llvm::BumpPtrAllocator bAlloc;
   llvm::StringSaver saver{bAlloc};
   llvm::DenseMap<void *, SpecificAllocBase *> instances;
 
   ErrorHandler e;
+
+  // Per thread allocs and savers.
+  PerThreadContext *perThreadContext();
+
+  // Vector of the TLS context's addresses.
+  // We store the TLS variable's addresses so that we can reset the variable's
+  // value from another thread.
+  std::vector<PerThreadContext **> perThreadContexts;
+  llvm::sys::RWMutex contextMutex;
 };
 
 // Retrieve the global state. Currently only one state can exist per process,
 // but in the future we plan on supporting an arbitrary number of LLD instances
 // in a single process.
+// TODO: How would this affect the TLS variable's addresses?
+// (hopefully not much)
 CommonLinkerContext &commonContext();
 
 template <typename T = CommonLinkerContext> T &context() {
   return static_cast<T &>(commonContext());
 }
 
 bool hasContext();
 
+llvm::StringSaver &perThreadSaver();
+llvm::BumpPtrAllocator &perThreadbAlloc();
+
 inline llvm::StringSaver &saver() { return context().saver; }
 inline llvm::BumpPtrAllocator &bAlloc() { return context().bAlloc; }
 } // namespace lld
 
 #endif

lld/include/lld/Common/Memory.h

 // Most objects are allocated using the arena allocators defined by this file.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLD_COMMON_MEMORY_H
 #define LLD_COMMON_MEMORY_H
 
 #include "llvm/Support/Allocator.h"
+#include "llvm/Support/RWMutex.h"
 
 namespace lld {
+
 // A base class only used by the CommonLinkerContext to keep track of the
 // SpecificAlloc<> instances.
 struct SpecificAllocBase {
   virtual ~SpecificAllocBase() = default;
   static SpecificAllocBase *getOrCreate(void *tag, size_t size, size_t align,
                                         SpecificAllocBase *(&creator)(void *));
+
+  static SpecificAllocBase *
+  getOrCreatePerThread(void *tag, size_t size, size_t align,
+                       SpecificAllocBase *(&creator)(void *));
 };
 
 // An arena of specific types T, created on-demand.
 template <class T> struct SpecificAlloc : public SpecificAllocBase {
   static SpecificAllocBase *create(void *storage) {
     return new (storage) SpecificAlloc<T>();
   }
   llvm::SpecificBumpPtrAllocator<T> alloc;

 template <typename T>
 inline llvm::SpecificBumpPtrAllocator<T> &getSpecificAllocSingleton() {
   SpecificAllocBase *instance = SpecificAllocBase::getOrCreate(
       &SpecificAlloc<T>::tag, sizeof(SpecificAlloc<T>),
       alignof(SpecificAlloc<T>), SpecificAlloc<T>::create);
   return ((SpecificAlloc<T> *)instance)->alloc;
 }
 
+template <typename T>
+inline llvm::SpecificBumpPtrAllocator<T> &getSpecificAllocSingletonPerThread() {
+  SpecificAllocBase *instance = SpecificAllocBase::getOrCreatePerThread(
+      &SpecificAlloc<T>::tag, sizeof(SpecificAlloc<T>),
+      alignof(SpecificAlloc<T>), SpecificAlloc<T>::create);
+  return ((SpecificAlloc<T> *)instance)->alloc;
+}
 // Creates new instances of T off a (almost) contiguous arena/object pool. The
 // instances are destroyed whenever lldMain() goes out of scope.
-template <typename T, typename... U> T *make(U &&... args) {
+template <typename T, typename... U> T *makeUnsafe(U &&...args) {
+  return new (getSpecificAllocSingleton<T>().Allocate())
+      T(std::forward<U>(args)...);
+}
+
+// Creates new instances of T off a (almost) contiguous arena/object pool
+// specific for this thread. The instances are destroyed whenever lldMain() goes
+// out of scope.
+template <typename T, typename... U> T *makePerThread(U &&...args) {
+  return new (getSpecificAllocSingletonPerThread<T>().Allocate())
+      T(std::forward<U>(args)...);
+}
+
+template <typename T, typename... U> T *make(U &&...args) {
+#ifndef THREAD_SAFE_MEMORY
   return new (getSpecificAllocSingleton<T>().Allocate())
       T(std::forward<U>(args)...);
+#else
+  return new (getSpecificAllocSingletonPerThread<T>().Allocate())
+      T(std::forward<U>(args)...);
+#endif
 }
 
 } // namespace lld
 
 #endif