Merge pull request #1032 from FEX-Emu/skmp/aotir-gen-mt

Multi threaded AOTGen

External/FEXCore/Source/Interface/Context/Context.cpp

@@ -50,8 +50,8 @@ namespace FEXCore::Context {
     CTX->Step();
   }
 
-  void CompileRIP(FEXCore::Context::Context *CTX, uint64_t GuestRIP) {
-    CTX->CompileBlock(CTX->ParentThread->CurrentFrame, GuestRIP);
+  void CompileRIP(FEXCore::Core::InternalThreadState *Thread, uint64_t GuestRIP) {
+    Thread->CTX->CompileBlock(Thread->CurrentFrame, GuestRIP);
   }
 
   FEXCore::Context::ExitReason RunUntilExit(FEXCore::Context::Context *CTX) {

External/FEXCore/Source/Interface/Core/Core.cpp

@@ -43,6 +43,7 @@ $end_info$
 #include <sys/stat.h>
 
 #include "Interface/Core/GdbServer.h"
+#include <shared_mutex>
 
 namespace FEXCore::CPU {
   bool CreateCPUCore(FEXCore::Context::Context *CTX) {
@@ -52,7 +53,60 @@ namespace FEXCore::CPU {
   }
 }
 
-static std::mutex AOTIRCacheLock;
+static std::shared_mutex AOTIRCacheLock;
+static std::shared_mutex AOTIRCaptureCacheWriteoutLock;
+static std::queue<std::function<void()>> AOTIRCaptureCacheWriteoutQueue;
+static std::atomic<bool> AOTIRCaptureCacheWriteoutFlusing;
+
+void AOTIRCaptureCacheWriteoutQueue_Flush() {
+
+  {
+    std::shared_lock lk{AOTIRCaptureCacheWriteoutLock};
+    if (AOTIRCaptureCacheWriteoutQueue.size() == 0) {
+      AOTIRCaptureCacheWriteoutFlusing.store(false);
+      return;    
+    }
+  }
+  
+  for (;;) {
+    AOTIRCaptureCacheWriteoutLock.lock();
+    std::function<void()> fn = std::move(AOTIRCaptureCacheWriteoutQueue.front());
+    bool MaybeEmpty = false;
+    AOTIRCaptureCacheWriteoutQueue.pop();
+    MaybeEmpty = AOTIRCaptureCacheWriteoutQueue.size() == 0;
+    AOTIRCaptureCacheWriteoutLock.unlock();
+
+    fn();
+    if (MaybeEmpty) {
+      std::shared_lock lk{AOTIRCaptureCacheWriteoutLock};
+      if (AOTIRCaptureCacheWriteoutQueue.size() == 0) {
+        AOTIRCaptureCacheWriteoutFlusing.store(false);
+        return;
+      }
+    }
+  }
+
+  LOGMAN_MSG_A("Must never get here");
+}
+
+void AOTIRCaptureCacheWriteoutQueue_Append(const std::function<void()> &fn) {
+  bool Flush = false;
+
+  {
+    std::unique_lock lk{AOTIRCaptureCacheWriteoutLock};
+    AOTIRCaptureCacheWriteoutQueue.push(fn);
+    if (AOTIRCaptureCacheWriteoutQueue.size() > 10000) {
+      Flush = true;
+    }
+  }
+
+  bool test_val = false;
+  if (Flush && AOTIRCaptureCacheWriteoutFlusing.compare_exchange_strong(test_val, true)) {
+    AOTIRCaptureCacheWriteoutQueue_Flush();
+  }
+}
+
+
 
 namespace FEXCore::Core {
 struct ThreadLocalData {
@@ -434,8 +488,6 @@ namespace FEXCore::Context {
   }
 
   void Context::InitializeThread(FEXCore::Core::InternalThreadState *Thread) {
-    InitializeThreadData(Thread);
-
     // This will create the execution thread but it won't actually start executing
     ExecutionThreadHandler *Arg = reinterpret_cast<ExecutionThreadHandler*>(FEXCore::Allocator::malloc(sizeof(ExecutionThreadHandler)));
     Arg->This = this;
@@ -514,6 +566,7 @@ namespace FEXCore::Context {
     Thread->ThreadManager.parent_tid = ParentTID;
 
     InitializeCompiler(Thread, false);
+    InitializeThreadData(Thread);
 
     return Thread;
   }
@@ -817,23 +870,25 @@ namespace FEXCore::Context {
   }
 
   void AOTIRCaptureCacheEntry::AppendAOTIRCaptureCache(uint64_t GuestRIP, uint64_t Start, uint64_t Length, uint64_t Hash, FEXCore::IR::IRListView *IRList, FEXCore::IR::RegisterAllocationData *RAData) {
-    Index.emplace(GuestRIP, Stream->tellp());
+    auto Inserted = Index.emplace(GuestRIP, Stream->tellp());
 
-    //GuestHash
-    Stream->write((char*)&Hash, sizeof(Hash));
+    if (Inserted.second) {
+      //GuestHash
+      Stream->write((char*)&Hash, sizeof(Hash));
 
-    //GuestLength
-    Stream->write((char*)&Length, sizeof(Length));
-    
-    // RAData (inline)
-    // In file, IsShared is always set
-    auto Shared = RAData->IsShared;
-    RAData->IsShared = true;
-    Stream->write((char*)RAData, RAData->Size(RAData->MapCount));
-    RAData->IsShared = Shared;
-    
-    // IRData (inline)
-    IRList->Serialize(*Stream);
+      //GuestLength
+      Stream->write((char*)&Length, sizeof(Length));
+      
+      // RAData (inline)
+      // In file, IsShared is always set
+      auto Shared = RAData->IsShared;
+      RAData->IsShared = true;
+      Stream->write((char*)RAData, RAData->Size(RAData->MapCount));
+      RAData->IsShared = Shared;
+      
+      // IRData (inline)
+      IRList->Serialize(*Stream);
+    }
   }
 
   std::tuple<void *, FEXCore::IR::IRListView *, FEXCore::Core::DebugData *, FEXCore::IR::RegisterAllocationData *, bool, uint64_t, uint64_t> Context::CompileCode(FEXCore::Core::InternalThreadState *Thread, uint64_t GuestRIP) {
@@ -860,7 +915,7 @@ namespace FEXCore::Context {
     }
 
     {
-      std::lock_guard<std::mutex> lk(AOTIRCacheLock);
+      std::shared_lock lk(AOTIRCacheLock);
       auto file = AddrToFile.lower_bound(GuestRIP);
       if (file != AddrToFile.begin()) {
         --file;
@@ -872,7 +927,7 @@ namespace FEXCore::Context {
     }
 
     if (IRList == nullptr && Config.AOTIRLoad) {
-      std::lock_guard<std::mutex> lk(AOTIRCacheLock);
+      std::shared_lock lk(AOTIRCacheLock);
       auto file = AddrToFile.lower_bound(GuestRIP);
       if (file != AddrToFile.begin()) {
         --file;
@@ -999,14 +1054,16 @@ namespace FEXCore::Context {
   }
 
   void Context::WriteFilesWithCode(std::function<void(const std::string& fileid, const std::string& filename)> Writer) {
-    std::lock_guard<std::mutex> lk(AOTIRCacheLock);
+    std::shared_lock lk(AOTIRCacheLock);
     for( const auto &File: FilesWithCode) {
       Writer(File.first, File.second);
     }
   }
 
   void Context::FinalizeAOTIRCache() {
-    std::lock_guard<std::mutex> lk(AOTIRCacheLock);
+    AOTIRCaptureCacheWriteoutQueue_Flush();
+
+    std::unique_lock lk(AOTIRCacheLock);
 
     for (auto &AOTModule: AOTIRCaptureCache) {
       if (!AOTModule.second.Stream) {
@@ -1054,7 +1111,7 @@ namespace FEXCore::Context {
       abort();
     }
   }
-  
+
   uintptr_t Context::CompileBlock(FEXCore::Core::CpuStateFrame *Frame, uint64_t GuestRIP) {
     auto Thread = Frame->Thread;
 
@@ -1130,28 +1187,34 @@ namespace FEXCore::Context {
     if (GeneratedIR) {
       // Add to AOT cache if aot generation is enabled
       if ((Config.AOTIRCapture() || Config.AOTIRGenerate()) && RAData) {
-        std::lock_guard<std::mutex> lk(AOTIRCacheLock);
-
         auto hash = XXH3_64bits((void*)StartAddr, Length);
 
+        std::shared_lock lk(AOTIRCacheLock);
+
         auto file = AddrToFile.lower_bound(StartAddr);
         if (file != AddrToFile.begin()) {
           --file;
           if (file->second.Start <= StartAddr && (file->second.Start + file->second.Len) >= (StartAddr + Length)) {
-            auto *AotFile = &AOTIRCaptureCache[file->second.fileid];
-            if (!AotFile->Stream) {
-              AotFile->Stream = AOTIRWriter(file->second.fileid);
-              uint64_t tag = 0xDEADBEEFC0D30004;
-              AotFile->Stream->write((char*)&tag, sizeof(tag));
-            }
-            AotFile->AppendAOTIRCaptureCache(GuestRIP - file->second.Start + file->second.Offset, StartAddr - file->second.Start + file->second.Offset, Length, hash, IRList, RAData);
+            auto LocalRIP = GuestRIP - file->second.Start + file->second.Offset;
+            auto LocalStartAddr = StartAddr - file->second.Start + file->second.Offset;
+            auto fileid = file->second.fileid;
+            AOTIRCaptureCacheWriteoutQueue_Append([this, LocalRIP, LocalStartAddr, Length, hash, IRList, RAData, fileid]() {
+              auto *AotFile = &AOTIRCaptureCache[fileid];
+
+              if (!AotFile->Stream) {
+                AotFile->Stream = AOTIRWriter(fileid);
+                uint64_t tag = 0xDEADBEEFC0D30004;
+                AotFile->Stream->write((char*)&tag, sizeof(tag));
+              }
+              AotFile->AppendAOTIRCaptureCache(LocalRIP, LocalStartAddr, Length, hash, IRList, RAData);
+              delete IRList;
+              FEXCore::Allocator::free(RAData);
+            });
           }
         }
 
         if (Config.AOTIRGenerate()) {
           // cleanup memory and early exit here -- we're not running the application
-          delete IRList;
-          FEXCore::Allocator::free(RAData);
 
           if (DecrementRefCount)
             --Thread->CompileBlockReentrantRefCount;
@@ -1313,7 +1376,7 @@ namespace FEXCore::Context {
       fileid += Config.ABILocalFlags ? "L" : "l";
       fileid += Config.ABINoPF ? "p" : "P";
 
-      std::lock_guard<std::mutex> lk(AOTIRCacheLock);
+      std::unique_lock lk(AOTIRCacheLock);
 
       AddrToFile.insert({ Base, { Base, Size, Offset, fileid, filename, nullptr, false} });
 
@@ -1328,13 +1391,13 @@ namespace FEXCore::Context {
   }
 
   void Context::RemoveNamedRegion(uintptr_t Base, uintptr_t Size) {
-    std::lock_guard<std::mutex> lk(AOTIRCacheLock);
+    std::unique_lock lk(AOTIRCacheLock);
     // TODO: Support partial removing
     AddrToFile.erase(Base);
   }
 
-  void ConfigureAOTGen(FEXCore::Context::Context *CTX, std::set<uint64_t> *ExternalBranches, uint64_t SectionMaxAddress) {
-    CTX->ParentThread->FrontendDecoder->SetExternalBranches(ExternalBranches);
-    CTX->ParentThread->FrontendDecoder->SetSectionMaxAddress(SectionMaxAddress);
+  void ConfigureAOTGen(FEXCore::Core::InternalThreadState *Thread, std::set<uint64_t> *ExternalBranches, uint64_t SectionMaxAddress) {
+    Thread->FrontendDecoder->SetExternalBranches(ExternalBranches);
+    Thread->FrontendDecoder->SetSectionMaxAddress(SectionMaxAddress);
   }
 }

External/FEXCore/include/FEXCore/Core/Context.h

@@ -119,7 +119,7 @@ namespace FEXCore::Context {
    */
   __attribute__((visibility("default"))) ExitReason RunUntilExit(FEXCore::Context::Context *CTX);
 
-  __attribute__((visibility("default"))) void CompileRIP(FEXCore::Context::Context *CTX, uint64_t GuestRIP);
+  __attribute__((visibility("default"))) void CompileRIP(FEXCore::Core::InternalThreadState *Thread, uint64_t GuestRIP);
 
   /**
    * @brief Gets the program exit status
@@ -235,5 +235,5 @@ namespace FEXCore::Context {
   __attribute__((visibility("default"))) void WriteFilesWithCode(FEXCore::Context::Context *CTX, std::function<void(const std::string& fileid, const std::string& filename)> Writer);
   __attribute__((visibility("default"))) void FlushCodeRange(FEXCore::Core::InternalThreadState *Thread, uint64_t Start, uint64_t Length);
 
-  __attribute__((visibility("default"))) void ConfigureAOTGen(FEXCore::Context::Context *CTX, std::set<uint64_t> *ExternalBranches, uint64_t SectionMaxAddress);
+  __attribute__((visibility("default"))) void ConfigureAOTGen(FEXCore::Core::InternalThreadState *Thread, std::set<uint64_t> *ExternalBranches, uint64_t SectionMaxAddress);
 }

Source/Tests/FEXLoader.cpp

@@ -30,6 +30,10 @@ $end_info$
 #include <filesystem>
 #include <algorithm>
 #include <set>
+#include <thread>
+#include <queue>
+
+#include <sys/sysinfo.h>
 
 namespace {
 static bool SilentLog;
@@ -160,6 +164,152 @@ bool IsInterpreterInstalled() {
          std::filesystem::exists("/proc/sys/fs/binfmt_misc/FEX-x86_64");
 }
 
+void AOTGenSection(FEXCore::Context::Context *CTX, ELFCodeLoader2::LoadedSection &Section) {
+
+  // Make sure this section is executable and big enough
+  if (!Section.Executable || Section.Size < 16)
+    return;
+  
+  std::set<uintptr_t> InitialBranchTargets;
+
+  // Load the ELF again with symbol parsing this time
+  ELFLoader::ELFContainer container{Section.Filename, "", false};
+
+  // Add symbols to the branch targets list
+  container.AddSymbols([&](ELFLoader::ELFSymbol* sym) {
+    auto Destination = sym->Address + Section.ElfBase;
+
+    if (! (Destination >= Section.Base && Destination <= (Section.Base + Section.Size)) ) {
+      return; // outside of current section, unlikely to be real code
+    }
+
+    InitialBranchTargets.insert(Destination);
+  });
+
+  LogMan::Msg::I("Symbol seed: %ld", InitialBranchTargets.size());
+
+  // Add unwind entries to the branch target list
+  container.AddUnwindEntries([&](uintptr_t Entry) {
+    auto Destination = Entry + Section.ElfBase;
+
+    if (! (Destination >= Section.Base && Destination <= (Section.Base + Section.Size)) ) {
+      return; // outside of current section, unlikely to be real code
+    }
+
+    InitialBranchTargets.insert(Destination);
+  });
+
+  LogMan::Msg::I("Symbol + Unwind seed: %ld", InitialBranchTargets.size());
+
+  // Scan the executable section and try to find function entries
+  for (size_t Offset = 0; Offset < (Section.Size - 16); Offset++) {
+    uint8_t *pCode = (uint8_t *)(Section.Base + Offset);
+
+    // Possible CALL <disp32>
+    if (*pCode == 0xE8) {
+      uintptr_t Destination = (int)(pCode[1] | (pCode[2] << 8) | (pCode[3] << 16) | (pCode[4] << 24));
+      Destination += (uintptr_t)pCode + 5;
+
+      auto DestinationPtr = (uint8_t*)Destination;
+      
+      if (! (Destination >= Section.Base && Destination <= (Section.Base + Section.Size)) )
+        continue; // outside of current section, unlikely to be real code
+
+      if (DestinationPtr[0] == 0 && DestinationPtr[1] == 0)
+        continue; // add al, [rax], unlikely to be real code
+
+      InitialBranchTargets.insert(Destination);
+    }
+
+    // endbr64 marker marks an indirect branch destination
+    if (pCode[0] == 0xf3 && pCode[1] == 0x0f && pCode[2] == 0x1e && pCode[3] == 0xfa) {
+      InitialBranchTargets.insert((uintptr_t)pCode);
+    }
+  }
+
+  uint64_t SectionMaxAddress = Section.Base + Section.Size;
+
+  std::set<uint64_t> Compiled;
+  std::atomic<int> counter = 0;
+
+  std::queue<uint64_t> BranchTargets;
+  
+  // Setup BranchTargets, Compiled sets from InitiaBranchTargets
+
+  Compiled.insert(InitialBranchTargets.begin(), InitialBranchTargets.end());
+  for (auto BranchTarget: InitialBranchTargets) {
+    BranchTargets.push(BranchTarget);
+  }
+
+  InitialBranchTargets.clear();
+
+
+  std::mutex QueueMutex;
+  std::vector<std::thread> ThreadPool;
+
+  for (int i = 0; i < get_nprocs_conf(); i++) {
+    std::thread thd([&BranchTargets, CTX, &counter, &Compiled, &Section, &QueueMutex, SectionMaxAddress]() {
+
+      // Setup thread - Each compilation thread uses its own backing FEX thread
+      FEXCore::Core::CPUState state;
+      auto Thread = FEXCore::Context::CreateThread(CTX, &state, gettid());
+      std::set<uint64_t> ExternalBranchesLocal;
+      FEXCore::Context::ConfigureAOTGen(Thread, &ExternalBranchesLocal, SectionMaxAddress);
+
+
+      for (;;) {
+        uint64_t BranchTarget;
+
+        // Get a entrypoint to process from the queue
+        QueueMutex.lock();
+        if (BranchTargets.empty()) {
+          QueueMutex.unlock();
+          break; // no entrypoint to process - exit
+        }
+
+        BranchTarget = BranchTargets.front();
+        BranchTargets.pop();
+        QueueMutex.unlock();
+
+        // Compile entrypoint
+        counter++;
+        FEXCore::Context::CompileRIP(Thread, BranchTarget);
+
+        // Are there more branches?
+        if (ExternalBranchesLocal.size() > 0) {
+          // Add them to the "to process" list
+          QueueMutex.lock();
+          for(auto Destination: ExternalBranchesLocal) {
+              if (! (Destination >= Section.Base && Destination <= (Section.Base + Section.Size)) )
+                continue;
+              if (Compiled.contains(Destination))
+                continue;
+              Compiled.insert(Destination);
+            BranchTargets.push(Destination);
+          }
+          QueueMutex.unlock();
+          ExternalBranchesLocal.clear();
+        }
+      }
+
+      // All entryproints processed, cleanup this thread
+      FEXCore::Context::DestroyThread(CTX, Thread);
+    });
+
+    // Add to the thread pool
+    ThreadPool.push_back(std::move(thd));
+  }
+
+  // Make sure all threads are finished
+  for (auto & Thread: ThreadPool) {
+    Thread.join();
+  }
+
+  ThreadPool.clear();
+
+  LogMan::Msg::I("\nAll Done: %d", counter.load());
+}
+
 int main(int argc, char **argv, char **const envp) {
   bool IsInterpreter = RanAsInterpreter(argv[0]);
   LogMan::Throw::InstallHandler(AssertHandler);
@@ -356,93 +506,8 @@ int main(int argc, char **argv, char **const envp) {
   }
 
   if (AOTIRGenerate()) {
-    for(auto Section: Loader.Sections) {
-      if (Section.Executable && Section.Size > 16) {
-        ELFLoader::ELFContainer container{Section.Filename, "", false};
-
-        std::set<uintptr_t> BranchTargets;
-
-        container.AddSymbols([&](ELFLoader::ELFSymbol* sym) {
-          auto Destination = sym->Address + Section.ElfBase;
-
-          if (! (Destination >= Section.Base && Destination <= (Section.Base + Section.Size)) ) {
-            //printf("Sym : %lx %lx out of range\n", sym->Address, Destination);
-            return; // outside of current section, unlikely to be real code
-          }
-
-          BranchTargets.insert(Destination);
-        });
-
-        LogMan::Msg::I("Symbol seed: %ld", BranchTargets.size());
-
-        container.AddUnwindEntries([&](uintptr_t Entry) {
-          auto Destination = Entry + Section.ElfBase;
-
-          if (! (Destination >= Section.Base && Destination <= (Section.Base + Section.Size)) ) {
-            //printf("Sym : %lx %lx out of range\n", sym->Address, Destination);
-            return; // outside of current section, unlikely to be real code
-          }
-
-          BranchTargets.insert(Destination);
-        });
-
-
-        LogMan::Msg::I("Symbol + Unwind seed: %ld", BranchTargets.size());
-
-        for (size_t Offset = 0; Offset < (Section.Size - 16); Offset++) {
-          uint8_t *pCode = (uint8_t *)(Section.Base + Offset);
-
-          if (*pCode == 0xE8) {
-            uintptr_t Destination = (int)(pCode[1] | (pCode[2] << 8) | (pCode[3] << 16) | (pCode[4] << 24));
-            Destination += (uintptr_t)pCode + 5;
-
-            auto DestinationPtr = (uint8_t*)Destination;
-            
-            if (! (Destination >= Section.Base && Destination <= (Section.Base + Section.Size)) )
-              continue; // outside of current section, unlikely to be real code
-
-            if (DestinationPtr[0] == 0 && DestinationPtr[1] == 0)
-              continue; // add al, [rax], unlikely to be real code
-/*
-            if (DestinationPtr[0] == 0x44 && DestinationPtr[1] == 0x0f && DestinationPtr[2] == 0x6f)
-              continue; // REX.W + movq leads to frontend bugs
-*/
-            BranchTargets.insert(Destination);
-          }
-
-          if (pCode[0] == 0xf3 && pCode[1] == 0x0f && pCode[2] == 0x1e && pCode[3] == 0xfa) {
-            BranchTargets.insert((uintptr_t)pCode);
-          }
-        }
-
-        uint64_t SectionMaxAddress = Section.Base + Section.Size;
-        std::set<uint64_t> ExternalBranches;
-
-        FEXCore::Context::ConfigureAOTGen(CTX, &ExternalBranches, SectionMaxAddress);
-
-        std::set<uint64_t> Compiled;
-        int counter = 0;
-        do {        
-          LogMan::Msg::I("Discovered %ld Branch Targets in this pass", BranchTargets.size());
-          for (auto RIP: BranchTargets) {
-            if ((counter++) % 1000 == 0)
-              LogMan::Msg::I("Compiling %d %lX", counter, RIP - Section.ElfBase);
-            FEXCore::Context::CompileRIP(CTX, RIP);
-            Compiled.insert(RIP);
-          }
-          LogMan::Msg::I("\nPass Done");
-          BranchTargets.clear();
-          for (auto Destination: ExternalBranches) {
-            if (! (Destination >= Section.Base && Destination <= (Section.Base + Section.Size)) )
-              continue;
-            if (Compiled.contains(Destination))
-              continue;
-            BranchTargets.insert(Destination);
-          }
-          ExternalBranches.clear();
-        } while (BranchTargets.size() > 0);
-        LogMan::Msg::I("\nAll Done: %d", counter);
-      }
+    for(auto &Section: Loader.Sections) {
+      AOTGenSection(CTX, Section);
     }
   } else {
     FEXCore::Context::RunUntilExit(CTX);