FEX-Emu--FEX/External/FEXCore/Source/Interface/IR/AOTIR.cpp

#include "Interface/Context/Context.h"
#include "Interface/IR/AOTIR.h"

#include <FEXCore/IR/IntrusiveIRList.h>
#include <FEXCore/IR/RegisterAllocationData.h>
#include <FEXCore/Utils/Allocator.h>
#include <FEXCore/HLE/SyscallHandler.h>
#include <Interface/Core/LookupCache.h>
#include <FEXCore/HLE/SourcecodeResolver.h>

#include <cstddef>
#include <cstdint>
#include <filesystem>
#include <fstream>
#include <mutex>
#include <sys/mman.h>
#include <sys/stat.h>
#include <unistd.h>
#include <xxhash.h>


#if defined(GDB_SYMBOLS_ENABLED)

#include <gdb/jit-reader.h>

struct blocks_t {
  char name[512];
  GDB_CORE_ADDR start;
  GDB_CORE_ADDR end;
};

struct info_t {
  char filename[512];
  long blocks_ofs;
  long lines_ofs;

  int nblocks;
  int nlines;
};


extern "C" {
enum jit_actions_t
{
  JIT_NOACTION = 0,
  JIT_REGISTER_FN,
  JIT_UNREGISTER_FN
};

struct jit_code_entry
{
  jit_code_entry *next_entry;
  jit_code_entry *prev_entry;
  const char *symfile_addr;
  uint64_t symfile_size;
};

struct jit_descriptor
{
  uint32_t version;
  /* This type should be jit_actions_t, but we use uint32_t
     to be explicit about the bitwidth.  */
  uint32_t action_flag;
  jit_code_entry *relevant_entry;
  jit_code_entry *first_entry;
};

/* Make sure to specify the version statically, because the
   debugger may check the version before we can set it.  */

constinit jit_descriptor __jit_debug_descriptor = { .version = 1 };

/* GDB puts a breakpoint in this function.  */
void __attribute__((noinline)) __jit_debug_register_code() { asm volatile (""::"r"(&__jit_debug_descriptor)); };
}
#endif

namespace FEXCore::IR {
  AOTIRInlineEntry *AOTIRInlineIndex::GetInlineEntry(uint64_t DataOffset) {
    uintptr_t This = (uintptr_t)this;

    return (AOTIRInlineEntry*)(This + DataBase + DataOffset);
  }

  AOTIRInlineEntry *AOTIRInlineIndex::Find(uint64_t GuestStart) {
    ssize_t l = 0;
    ssize_t r = Count - 1;

    while (l <= r) {
      size_t m = l + (r - l) / 2;

      if (Entries[m].GuestStart == GuestStart)
        return GetInlineEntry(Entries[m].DataOffset);
      else if (Entries[m].GuestStart < GuestStart)
        l = m + 1;
      else
        r = m - 1;
    }

    return nullptr;
  }

  IR::RegisterAllocationData *AOTIRInlineEntry::GetRAData() {
    return (IR::RegisterAllocationData *)InlineData;
  }

  IR::IRListView *AOTIRInlineEntry::GetIRData() {
    auto RAData = GetRAData();
    auto Offset = RAData->Size(RAData->MapCount);

    return (IR::IRListView *)&InlineData[Offset];
  }

  void AOTIRCaptureCacheEntry::AppendAOTIRCaptureCache(uint64_t GuestRIP, uint64_t Start, uint64_t Length, uint64_t Hash, FEXCore::IR::IRListView *IRList, FEXCore::IR::RegisterAllocationData *RAData) {
    auto Inserted = Index.emplace(GuestRIP, Stream->tellp());

    if (Inserted.second) {
      //GuestHash
      Stream->write((const char*)&Hash, sizeof(Hash));

      //GuestLength
      Stream->write((const char*)&Length, sizeof(Length));

      RAData->Serialize(*Stream);

      // IRData (inline)
      IRList->Serialize(*Stream);
    }
  }

  static bool readAll(int fd, void *data, size_t size) {
    int rv = read(fd, data, size);

    if (rv != size)
      return false;
    else
      return true;
  }

  static bool LoadAOTIRCache(AOTIRCacheEntry *Entry, int streamfd) {
    uint64_t tag;

    if (!readAll(streamfd, (char*)&tag, sizeof(tag)) || tag != FEXCore::IR::AOTIR_COOKIE)
      return false;

    std::string Module;
    uint64_t ModSize;
    uint64_t IndexSize;

    lseek(streamfd, -sizeof(ModSize), SEEK_END);

    if (!readAll(streamfd,  (char*)&ModSize, sizeof(ModSize)))
      return false;

    Module.resize(ModSize);

    lseek(streamfd, -sizeof(ModSize) - ModSize, SEEK_END);

    if (!readAll(streamfd,  (char*)&Module[0], Module.size()))
      return false;

    if (Entry->FileId != Module) {
      return false;
    }

    lseek(streamfd, -sizeof(ModSize) - ModSize - sizeof(IndexSize), SEEK_END);

    if (!readAll(streamfd,  (char*)&IndexSize, sizeof(IndexSize)))
      return false;

    struct stat fileinfo;
    if (fstat(streamfd, &fileinfo) < 0)
      return false;
    size_t Size = (fileinfo.st_size + 4095) & ~4095;

    size_t IndexOffset = fileinfo.st_size - IndexSize -sizeof(ModSize) - ModSize - sizeof(IndexSize);

    void *FilePtr = FEXCore::Allocator::mmap(nullptr, Size, PROT_READ, MAP_SHARED, streamfd, 0);

    if (FilePtr == MAP_FAILED) {
      return false;
    }

    auto Array = (AOTIRInlineIndex *)((char*)FilePtr + IndexOffset);

    LOGMAN_THROW_A_FMT(Entry->Array == nullptr && Entry->FilePtr == nullptr, "Entry must not be initialized here");
    Entry->Array = Array;
    Entry->FilePtr = FilePtr;
    Entry->Size = Size;

    LogMan::Msg::DFmt("AOTIR: Module {} has {} functions", Module, Array->Count);

    return true;
  }

  void AOTIRCaptureCache::FinalizeAOTIRCache() {
    AOTIRCaptureCacheWriteoutQueue_Flush();

    std::unique_lock lk(AOTIRCacheLock);

    for (auto& [String, Entry] : AOTIRCaptureCacheMap) {
      if (!Entry.Stream) {
        continue;
      }

      const auto ModSize = String.size();
      auto &stream = Entry.Stream;

      // pad to 32 bytes
      constexpr char Zero = 0;
      while(stream->tellp() & 31)
        stream->write(&Zero, 1);

      // AOTIRInlineIndex
      const auto FnCount = Entry.Index.size();
      const size_t DataBase = -stream->tellp();

      stream->write((const char*)&FnCount, sizeof(FnCount));
      stream->write((const char*)&DataBase, sizeof(DataBase));

      for (const auto& [GuestStart, DataOffset] : Entry.Index) {
        //AOTIRInlineIndexEntry

        // GuestStart
        stream->write((const char*)&GuestStart, sizeof(GuestStart));

        // DataOffset
        stream->write((const char*)&DataOffset, sizeof(DataOffset));
      }

      // End of file header
      const auto IndexSize = FnCount * sizeof(FEXCore::IR::AOTIRInlineIndexEntry) + sizeof(DataBase) + sizeof(FnCount);
      stream->write((const char*)&IndexSize, sizeof(IndexSize));
      stream->write(String.c_str(), ModSize);
      stream->write((const char*)&ModSize, sizeof(ModSize));

      // Close the stream
      stream->close();

      // Rename the file to atomically update the cache with the temporary file
      AOTIRRenamer(String);
    }
  }

  void AOTIRCaptureCache::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_FMT("Must never get here");
  }

  void AOTIRCaptureCache::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();
    }
  }

  void AOTIRCaptureCache::WriteFilesWithCode(std::function<void(const std::string& fileid, const std::string& filename)> Writer) {
    std::shared_lock lk(AOTIRCacheLock);
    for( const auto &Entry: AOTIRCache) {
      if (Entry.second.ContainsCode) {
        Writer(Entry.second.FileId, Entry.second.Filename);
      }
    }
  }

  AOTIRCaptureCache::PreGenerateIRFetchResult AOTIRCaptureCache::PreGenerateIRFetch(uint64_t GuestRIP, FEXCore::IR::IRListView *IRList) {
    auto AOTIRCacheEntry = CTX->SyscallHandler->LookupAOTIRCacheEntry(GuestRIP);

    PreGenerateIRFetchResult Result{};

    if (AOTIRCacheEntry.Entry) {
      AOTIRCacheEntry.Entry->ContainsCode = true;

      if (IRList == nullptr && CTX->Config.AOTIRLoad()) {
        auto Mod = AOTIRCacheEntry.Entry->Array;

        if (Mod != nullptr)
        {
          auto AOTEntry = Mod->Find(GuestRIP - AOTIRCacheEntry.VAFileStart);

          if (AOTEntry) {
            // verify hash
            auto MappedStart = GuestRIP;
            auto hash = XXH3_64bits((void*)MappedStart, AOTEntry->GuestLength);
            if (hash == AOTEntry->GuestHash) {
              Result.IRList = AOTEntry->GetIRData();
              //LogMan::Msg::DFmt("using {} + {:x} -> {:x}\n", file->second.fileid, AOTEntry->first, GuestRIP);

              Result.RAData = AOTEntry->GetRAData()->CreateCopy();
              Result.DebugData = new FEXCore::Core::DebugData();
              Result.StartAddr = MappedStart;
              Result.Length = AOTEntry->GuestLength;
              Result.GeneratedIR = true;
            } else {
              LogMan::Msg::IFmt("AOTIR: hash check failed {:x}\n", MappedStart);
            }
          } else {
            //LogMan::Msg::IFmt("AOTIR: Failed to find {:x}, {:x}, {}\n", GuestRIP, GuestRIP - file->second.Start + file->second.Offset, file->second.fileid);
          }
        }
      }
    }

    return Result;
  }

#if defined(GDB_SYMBOLS_ENABLED)
  static void GDBJITRegister(AOTIRCacheEntry *Entry, uintptr_t VAFileStart, uint64_t GuestRIP, uintptr_t HostEntry, FEXCore::Core::DebugData *DebugData) {
    auto map = Entry->SourcecodeMap.get();

    if (map) {
      auto FileOffset = GuestRIP - VAFileStart;

      auto Sym = map->FindSymbolMapping(FileOffset);

      std::string SymName = HLE::SourcecodeSymbolMapping::SymName(Sym, Entry->Filename, HostEntry, FileOffset);

      std::vector<gdb_line_mapping> Lines;
      for (const auto &GuestOpcode : DebugData->GuestOpcodes) {
        auto Line = map->FindLineMapping(GuestRIP + GuestOpcode.GuestEntryOffset - VAFileStart);
        if (Line) {
          Lines.push_back({Line->LineNumber, HostEntry + GuestOpcode.HostEntryOffset});
        }
      }

      size_t size = sizeof(info_t) + 1 * sizeof(blocks_t) + Lines.size() * sizeof(gdb_line_mapping);

      auto mem = (uint8_t *)malloc(size);
      auto base = mem;
      info_t *info = (info_t *)mem;
      mem += sizeof(info_t);

      strncpy(info->filename, map->SourceFile.c_str(), 511);

      info->nblocks = 1;

      auto blocks = (blocks_t *)mem;
      info->blocks_ofs = mem - base;

      mem += info->nblocks * sizeof(blocks_t);

      for (int i = 0; i < info->nblocks; i++) {
        strncpy(blocks[i].name, SymName.c_str(), 511);
        blocks[i].start = HostEntry;
        blocks[i].end = HostEntry + DebugData->HostCodeSize;
      }

      info->nlines = Lines.size();

      auto lines = (gdb_line_mapping *)mem;
      info->lines_ofs = mem - base;
      mem += info->nlines * sizeof(gdb_line_mapping);

      if (info->nlines) {
        memcpy(lines, &Lines.at(0), info->nlines * sizeof(gdb_line_mapping));
      }

      auto entry = new jit_code_entry{0, 0, 0, 0};

      entry->symfile_addr = (const char *)info;
      entry->symfile_size = size;

      if (__jit_debug_descriptor.first_entry) {
        __jit_debug_descriptor.relevant_entry->next_entry = entry;
        entry->prev_entry = __jit_debug_descriptor.relevant_entry;
      } else {
        __jit_debug_descriptor.first_entry = entry;
      }

      __jit_debug_descriptor.relevant_entry = entry;
      __jit_debug_descriptor.action_flag = JIT_REGISTER_FN;
      __jit_debug_register_code();
    }
  }
  #endif

  bool AOTIRCaptureCache::PostCompileCode(
    FEXCore::Core::InternalThreadState *Thread,
    void* CodePtr,
    uint64_t GuestRIP,
    uint64_t StartAddr,
    uint64_t Length,
    FEXCore::IR::RegisterAllocationData::UniquePtr RAData,
    FEXCore::IR::IRListView *IRList,
    FEXCore::Core::DebugData *DebugData,
    bool GeneratedIR) {

    // Both generated ir and LibraryJITName need a named region lookup
    if (GeneratedIR || CTX->Config.LibraryJITNaming() || CTX->Config.GDBSymbols()) {

      auto AOTIRCacheEntry = CTX->SyscallHandler->LookupAOTIRCacheEntry(GuestRIP);

      if (AOTIRCacheEntry.Entry) {
        if (DebugData && CTX->Config.LibraryJITNaming()) {
          CTX->Symbols.RegisterNamedRegion(CodePtr, DebugData->HostCodeSize, AOTIRCacheEntry.Entry->Filename);
        }

        if (CTX->Config.GDBSymbols()) {
          #if defined(GDB_SYMBOLS_ENABLED)
            GDBJITRegister(AOTIRCacheEntry.Entry, AOTIRCacheEntry.VAFileStart, GuestRIP, (uintptr_t)CodePtr, DebugData);
          #else
            ERROR_AND_DIE_FMT("GDBSymbols support not compiled in");
          #endif
        }

        // Add to AOT cache if aot generation is enabled
        if (GeneratedIR && RAData &&
            (CTX->Config.AOTIRCapture() || CTX->Config.AOTIRGenerate())) {

          auto hash = XXH3_64bits((void*)StartAddr, Length);

          auto LocalRIP = GuestRIP - AOTIRCacheEntry.VAFileStart;
          auto LocalStartAddr = StartAddr - AOTIRCacheEntry.VAFileStart;
          auto FileId = AOTIRCacheEntry.Entry->FileId;
          // The underlying pointer and the unique_ptr deleter for RAData must
          // be marshalled separately to the lambda below. Otherwise, the
          // lambda can't be used as an std::function due to being non-copyable
          auto RADataCopy = RAData->CreateCopy();
          auto RADataCopyDeleter = RADataCopy.get_deleter();
          auto IRListCopy = IRList->CreateCopy();
          AOTIRCaptureCacheWriteoutQueue_Append([this, LocalRIP, LocalStartAddr, Length, hash, IRListCopy, RADataCopy=RADataCopy.release(), RADataCopyDeleter, FileId]() {

            // It is guaranteed via AOTIRCaptureCacheWriteoutLock and AOTIRCaptureCacheWriteoutFlusing that this will not run concurrently
            // Memory coherency is guaranteed via AOTIRCaptureCacheWriteoutLock
            
            auto *AotFile = &AOTIRCaptureCacheMap[FileId];

            if (!AotFile->Stream) {
              AotFile->Stream = AOTIRWriter(FileId);
              uint64_t tag = FEXCore::IR::AOTIR_COOKIE;
              AotFile->Stream->write((char*)&tag, sizeof(tag));
            }
            AotFile->AppendAOTIRCaptureCache(LocalRIP, LocalStartAddr, Length, hash, IRListCopy, RADataCopy);
            RADataCopyDeleter(RADataCopy);
            delete IRListCopy;
          });

          if (CTX->Config.AOTIRGenerate()) {
            // cleanup memory and early exit here -- we're not running the application
            Thread->CPUBackend->ClearCache();
            return true;
          }
        }
      }

      // Insert to caches if we generated IR
      if (GeneratedIR) {
        if (CTX->GetGdbServerStatus()) {
          // Add to thread local ir cache
          Core::LocalIREntry Entry = {StartAddr, Length, decltype(Entry.IR)(IRList), std::move(RAData), decltype(Entry.DebugData)(DebugData)};
          
          std::lock_guard<std::recursive_mutex> lk(Thread->LookupCache->WriteLock);
          Thread->DebugStore.insert({GuestRIP, std::move(Entry)});
        }
        else {
          // If the IR doesn't need to be retained then we can just delete it now
          delete DebugData;
          delete IRList;
        }
      }
    }

    return false;
  }

  AOTIRCacheEntry *AOTIRCaptureCache::LoadAOTIRCacheEntry(const std::string &filename) {
    auto base_filename = std::filesystem::path(filename).filename().string();

    if (!base_filename.empty()) {
      auto filename_hash = XXH3_64bits(filename.c_str(), filename.size());

      auto fileid = base_filename + "-" + std::to_string(filename_hash) + "-";

      // append optimization flags to the fileid
      fileid += (CTX->Config.SMCChecks == FEXCore::Config::CONFIG_SMC_FULL) ? "S" : "s";
      fileid += CTX->Config.TSOEnabled ? "T" : "t";
      fileid += CTX->Config.ABILocalFlags ? "L" : "l";
      fileid += CTX->Config.ABINoPF ? "p" : "P";

      std::unique_lock lk(AOTIRCacheLock);

      auto Inserted = AOTIRCache.insert({fileid, AOTIRCacheEntry { .FileId = fileid, .Filename = filename }});
      auto Entry = &(Inserted.first->second);

      LOGMAN_THROW_A_FMT(Entry->Array == nullptr, "Duplicate LoadAOTIRCacheEntry");

      if (CTX->Config.AOTIRLoad && AOTIRLoader) {
        auto streamfd = AOTIRLoader(fileid);
        if (streamfd != -1) {
          FEXCore::IR::LoadAOTIRCache(Entry, streamfd);
          close(streamfd);
        }
      }
      return Entry;
    }

    return nullptr;
  }

  void AOTIRCaptureCache::UnloadAOTIRCacheEntry(AOTIRCacheEntry *Entry) {
    LOGMAN_THROW_A_FMT(Entry != nullptr, "Removing not existing entry");

    if (Entry->Array) {
      FEXCore::Allocator::munmap(Entry->FilePtr, Entry->Size);
      Entry->Array = nullptr;
      Entry->FilePtr = nullptr;
      Entry->Size = 0;
    }
  }
}