Revert "[CGData] Lazy loading support for stable function map (#151660)"

This reverts commit 76dd742.

Author: kyulee-com

llvm/include/llvm/CGData/CodeGenData.h

@@ -285,9 +285,6 @@ enum CGDataVersion {
   // Version 3 adds the total size of the Names in the stable function map so
   // we can skip reading them into the memory for non-assertion builds.
   Version3 = 3,
-  // Version 4 adjusts the structure of stable function merging map for
-  // efficient lazy loading support.
-  Version4 = 4,
   CurrentVersion = CG_DATA_INDEX_VERSION
 };
 const uint64_t Version = CGDataVersion::CurrentVersion;

llvm/include/llvm/CGData/CodeGenData.inc

@@ -49,4 +49,4 @@ CG_DATA_SECT_ENTRY(CG_merge, CG_DATA_QUOTE(CG_DATA_MERGE_COMMON),
 #endif
 
 /* Indexed codegen data format version (start from 1). */
-#define CG_DATA_INDEX_VERSION 4
+#define CG_DATA_INDEX_VERSION 3

llvm/include/llvm/CGData/StableFunctionMap.h

@@ -20,8 +20,6 @@
 #include "llvm/ADT/StringMap.h"
 #include "llvm/IR/StructuralHash.h"
 #include "llvm/Support/Compiler.h"
-#include "llvm/Support/MemoryBuffer.h"
-#include <mutex>
 
 namespace llvm {
 
@@ -74,37 +72,11 @@ struct StableFunctionMap {
           IndexOperandHashMap(std::move(IndexOperandHashMap)) {}
   };
 
-  using StableFunctionEntries =
-      SmallVector<std::unique_ptr<StableFunctionEntry>>;
-
-  /// In addition to the deserialized StableFunctionEntry, the struct stores
-  /// the offsets of corresponding serialized stable function entries, and a
-  /// once flag for safe lazy loading in a multithreaded environment.
-  struct EntryStorage {
-    /// The actual storage of deserialized stable function entries. If the map
-    /// is lazily loaded, this will be empty until the first access by the
-    /// corresponding function hash.
-    StableFunctionEntries Entries;
-
-  private:
-    /// This is used to deserialize the entry lazily. Each element is the
-    /// corresponding serialized stable function entry's offset in the memory
-    /// buffer (StableFunctionMap::Buffer).
-    /// The offsets are only populated when loading the map lazily, otherwise
-    /// it is empty.
-    SmallVector<uint64_t> Offsets;
-    std::once_flag LazyLoadFlag;
-    friend struct StableFunctionMap;
-    friend struct StableFunctionMapRecord;
-  };
-
-  // Note: DenseMap requires value type to be copyable even if only using
-  // in-place insertion. Use STL instead. This also affects the
-  // deletion-while-iteration in finalize().
-  using HashFuncsMapType = std::unordered_map<stable_hash, EntryStorage>;
+  using HashFuncsMapType =
+      DenseMap<stable_hash, SmallVector<std::unique_ptr<StableFunctionEntry>>>;
 
   /// Get the HashToFuncs map for serialization.
-  const HashFuncsMapType &getFunctionMap() const;
+  const HashFuncsMapType &getFunctionMap() const { return HashToFuncs; }
 
   /// Get the NameToId vector for serialization.
   ArrayRef<std::string> getNames() const { return IdToName; }
@@ -127,19 +99,6 @@ struct StableFunctionMap {
   /// \returns true if there is no stable function entry.
   bool empty() const { return size() == 0; }
 
-  /// \returns true if there is an entry for the given function hash.
-  /// This does not trigger lazy loading.
-  bool contains(HashFuncsMapType::key_type FunctionHash) const {
-    return HashToFuncs.count(FunctionHash) > 0;
-  }
-
-  /// \returns the stable function entries for the given function hash. If the
-  /// map is lazily loaded, it will deserialize the entries if it is not already
-  /// done, other requests to the same hash at the same time will be blocked
-  /// until the entries are deserialized.
-  const StableFunctionEntries &
-  at(HashFuncsMapType::key_type FunctionHash) const;
-
   enum SizeType {
     UniqueHashCount,        // The number of unique hashes in HashToFuncs.
     TotalFunctionCount,     // The number of total functions in HashToFuncs.
@@ -160,31 +119,17 @@ struct StableFunctionMap {
   /// `StableFunctionEntry` is ready for insertion.
   void insert(std::unique_ptr<StableFunctionEntry> FuncEntry) {
     assert(!Finalized && "Cannot insert after finalization");
-    HashToFuncs[FuncEntry->Hash].Entries.emplace_back(std::move(FuncEntry));
+    HashToFuncs[FuncEntry->Hash].emplace_back(std::move(FuncEntry));
   }
 
-  void deserializeLazyLoadingEntry(HashFuncsMapType::iterator It) const;
-
-  /// Eagerly deserialize all the unloaded entries in the lazy loading map.
-  void deserializeLazyLoadingEntries() const;
-
-  bool isLazilyLoaded() const { return (bool)Buffer; }
-
   /// A map from a stable_hash to a vector of functions with that hash.
-  mutable HashFuncsMapType HashToFuncs;
+  HashFuncsMapType HashToFuncs;
   /// A vector of strings to hold names.
   SmallVector<std::string> IdToName;
   /// A map from StringRef (name) to an ID.
   StringMap<unsigned> NameToId;
   /// True if the function map is finalized with minimal content.
   bool Finalized = false;
-  /// The memory buffer that contains the serialized stable function map for
-  /// lazy loading.
-  /// Non-empty only if this StableFunctionMap is created from a MemoryBuffer
-  /// (i.e. by IndexedCodeGenDataReader::read()) and lazily deserialized.
-  std::shared_ptr<MemoryBuffer> Buffer;
-  /// Whether to read stable function names from the buffer.
-  bool ReadStableFunctionMapNames = true;
 
   friend struct StableFunctionMapRecord;
 };

llvm/include/llvm/CGData/StableFunctionMapRecord.h

@@ -24,26 +24,6 @@
 
 namespace llvm {
 
-/// The structure of the serialized stable function map is as follows:
-/// - Number of unique function/module names
-/// - Total size of unique function/module names for opt-in skipping
-/// - Unique function/module names
-/// - Padding to align to 4 bytes
-/// - Number of StableFunctionEntries
-/// - Hashes of each StableFunctionEntry
-/// - Fixed-size fields for each StableFunctionEntry (the order is consistent
-///   with the hashes above):
-///   - FunctionNameId
-///   - ModuleNameId
-///   - InstCount
-///   - Relative offset to the beginning of IndexOperandHashes for this entry
-/// - Total size of variable-sized IndexOperandHashes for lazy-loading support
-/// - Variable-sized IndexOperandHashes for each StableFunctionEntry:
-///   - Number of IndexOperandHashes
-///   - Contents of each IndexOperandHashes
-///     - InstIndex
-///     - OpndIndex
-///     - OpndHash
 struct StableFunctionMapRecord {
   std::unique_ptr<StableFunctionMap> FunctionMap;
 
@@ -60,25 +40,13 @@ struct StableFunctionMapRecord {
                                  const StableFunctionMap *FunctionMap,
                                  std::vector<CGDataPatchItem> &PatchItems);
 
-  /// A static helper function to deserialize the stable function map entry.
-  /// Ptr should be pointing to the start of the fixed-sized fields of the
-  /// entry when passed in.
-  LLVM_ABI static void deserializeEntry(const unsigned char *Ptr,
-                                        stable_hash Hash,
-                                        StableFunctionMap *FunctionMap);
-
   /// Serialize the stable function map to a raw_ostream.
   LLVM_ABI void serialize(raw_ostream &OS,
                           std::vector<CGDataPatchItem> &PatchItems) const;
 
   /// Deserialize the stable function map from a raw_ostream.
-  LLVM_ABI void deserialize(const unsigned char *&Ptr);
-
-  /// Lazily deserialize the stable function map from `Buffer` starting at
-  /// `Offset`. The individual stable function entry would be read lazily from
-  /// `Buffer` when the function map is accessed.
-  LLVM_ABI void lazyDeserialize(std::shared_ptr<MemoryBuffer> Buffer,
-                                uint64_t Offset);
+  LLVM_ABI void deserialize(const unsigned char *&Ptr,
+                            bool ReadStableFunctionMapNames = true);
 
   /// Serialize the stable function map to a YAML stream.
   LLVM_ABI void serializeYAML(yaml::Output &YOS) const;
@@ -102,18 +70,6 @@ struct StableFunctionMapRecord {
     yaml::Output YOS(OS);
     serializeYAML(YOS);
   }
-
-  /// Set whether to read stable function names from the buffer.
-  /// Has no effect if the function map is read from a YAML stream.
-  void setReadStableFunctionMapNames(bool Read) {
-    assert(
-        FunctionMap->empty() &&
-        "Cannot change ReadStableFunctionMapNames after the map is populated");
-    FunctionMap->ReadStableFunctionMapNames = Read;
-  }
-
-private:
-  void deserialize(const unsigned char *&Ptr, bool Lazy);
 };
 
 } // namespace llvm

llvm/lib/CGData/CodeGenData.cpp

@@ -186,7 +186,7 @@ Expected<Header> Header::readFromBuffer(const unsigned char *Curr) {
     return make_error<CGDataError>(cgdata_error::unsupported_version);
   H.DataKind = endian::readNext<uint32_t, endianness::little, unaligned>(Curr);
 
-  static_assert(IndexedCGData::CGDataVersion::CurrentVersion == Version4,
+  static_assert(IndexedCGData::CGDataVersion::CurrentVersion == Version3,
                 "Please update the offset computation below if a new field has "
                 "been added to the header.");
   H.OutlinedHashTreeOffset =

llvm/lib/CGData/CodeGenDataReader.cpp

@@ -26,12 +26,6 @@ static cl::opt<bool> IndexedCodeGenDataReadFunctionMapNames(
              "disabled to save memory and time for final consumption of the "
              "indexed CodeGenData in production."));
 
-cl::opt<bool> IndexedCodeGenDataLazyLoading(
-    "indexed-codegen-data-lazy-loading", cl::init(false), cl::Hidden,
-    cl::desc(
-        "Lazily load indexed CodeGenData. Enable to save memory and time "
-        "for final consumption of the indexed CodeGenData in production."));
-
 namespace llvm {
 
 static Expected<std::unique_ptr<MemoryBuffer>>
@@ -115,20 +109,11 @@ Error IndexedCodeGenDataReader::read() {
       return error(cgdata_error::eof);
     HashTreeRecord.deserialize(Ptr);
   }
-
-  // TODO: lazy loading support for outlined hash tree.
-  std::shared_ptr<MemoryBuffer> SharedDataBuffer = std::move(DataBuffer);
   if (hasStableFunctionMap()) {
     const unsigned char *Ptr = Start + Header.StableFunctionMapOffset;
     if (Ptr >= End)
       return error(cgdata_error::eof);
-    FunctionMapRecord.setReadStableFunctionMapNames(
-        IndexedCodeGenDataReadFunctionMapNames);
-    if (IndexedCodeGenDataLazyLoading)
-      FunctionMapRecord.lazyDeserialize(SharedDataBuffer,
-                                        Header.StableFunctionMapOffset);
-    else
-      FunctionMapRecord.deserialize(Ptr);
+    FunctionMapRecord.deserialize(Ptr, IndexedCodeGenDataReadFunctionMapNames);
   }
 
   return success();

llvm/lib/CGData/StableFunctionMap.cpp

@@ -15,10 +15,8 @@
 
 #include "llvm/CGData/StableFunctionMap.h"
 #include "llvm/ADT/SmallSet.h"
-#include "llvm/CGData/StableFunctionMapRecord.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
-#include <mutex>
 
 #define DEBUG_TYPE "stable-function-map"
 
@@ -95,10 +93,9 @@ void StableFunctionMap::insert(const StableFunction &Func) {
 
 void StableFunctionMap::merge(const StableFunctionMap &OtherMap) {
   assert(!Finalized && "Cannot merge after finalization");
-  deserializeLazyLoadingEntries();
   for (auto &[Hash, Funcs] : OtherMap.HashToFuncs) {
-    auto &ThisFuncs = HashToFuncs[Hash].Entries;
-    for (auto &Func : Funcs.Entries) {
+    auto &ThisFuncs = HashToFuncs[Hash];
+    for (auto &Func : Funcs) {
       auto FuncNameId =
           getIdOrCreateForName(*OtherMap.getNameForId(Func->FunctionNameId));
       auto ModuleNameId =
@@ -117,63 +114,25 @@ size_t StableFunctionMap::size(SizeType Type) const {
   case UniqueHashCount:
     return HashToFuncs.size();
   case TotalFunctionCount: {
-    deserializeLazyLoadingEntries();
     size_t Count = 0;
     for (auto &Funcs : HashToFuncs)
-      Count += Funcs.second.Entries.size();
+      Count += Funcs.second.size();
     return Count;
   }
   case MergeableFunctionCount: {
-    deserializeLazyLoadingEntries();
     size_t Count = 0;
     for (auto &[Hash, Funcs] : HashToFuncs)
-      if (Funcs.Entries.size() >= 2)
-        Count += Funcs.Entries.size();
+      if (Funcs.size() >= 2)
+        Count += Funcs.size();
     return Count;
   }
   }
   llvm_unreachable("Unhandled size type");
 }
 
-const StableFunctionMap::StableFunctionEntries &
-StableFunctionMap::at(HashFuncsMapType::key_type FunctionHash) const {
-  auto It = HashToFuncs.find(FunctionHash);
-  if (isLazilyLoaded())
-    deserializeLazyLoadingEntry(It);
-  return It->second.Entries;
-}
-
-void StableFunctionMap::deserializeLazyLoadingEntry(
-    HashFuncsMapType::iterator It) const {
-  assert(isLazilyLoaded() && "Cannot deserialize non-lazily-loaded map");
-  auto &[Hash, Storage] = *It;
-  std::call_once(Storage.LazyLoadFlag,
-                 [this, HashArg = Hash, &StorageArg = Storage]() {
-                   for (auto Offset : StorageArg.Offsets)
-                     StableFunctionMapRecord::deserializeEntry(
-                         reinterpret_cast<const unsigned char *>(Offset),
-                         HashArg, const_cast<StableFunctionMap *>(this));
-                 });
-}
-
-void StableFunctionMap::deserializeLazyLoadingEntries() const {
-  if (!isLazilyLoaded())
-    return;
-  for (auto It = HashToFuncs.begin(); It != HashToFuncs.end(); ++It)
-    deserializeLazyLoadingEntry(It);
-}
-
-const StableFunctionMap::HashFuncsMapType &
-StableFunctionMap::getFunctionMap() const {
-  // Ensure all entries are deserialized before returning the raw map.
-  if (isLazilyLoaded())
-    deserializeLazyLoadingEntries();
-  return HashToFuncs;
-}
-
 using ParamLocs = SmallVector<IndexPair>;
-static void
-removeIdenticalIndexPair(StableFunctionMap::StableFunctionEntries &SFS) {
+static void removeIdenticalIndexPair(
+    SmallVector<std::unique_ptr<StableFunctionMap::StableFunctionEntry>> &SFS) {
   auto &RSF = SFS[0];
   unsigned StableFunctionCount = SFS.size();
 
@@ -200,7 +159,9 @@ removeIdenticalIndexPair(StableFunctionMap::StableFunctionEntries &SFS) {
       SF->IndexOperandHashMap->erase(Pair);
 }
 
-static bool isProfitable(const StableFunctionMap::StableFunctionEntries &SFS) {
+static bool isProfitable(
+    const SmallVector<std::unique_ptr<StableFunctionMap::StableFunctionEntry>>
+        &SFS) {
   unsigned StableFunctionCount = SFS.size();
   if (StableFunctionCount < GlobalMergingMinMerges)
     return false;
@@ -241,11 +202,8 @@ static bool isProfitable(const StableFunctionMap::StableFunctionEntries &SFS) {
 }
 
 void StableFunctionMap::finalize(bool SkipTrim) {
-  deserializeLazyLoadingEntries();
-  SmallVector<HashFuncsMapType::iterator> ToDelete;
   for (auto It = HashToFuncs.begin(); It != HashToFuncs.end(); ++It) {
-    auto &[StableHash, Storage] = *It;
-    auto &SFS = Storage.Entries;
+    auto &[StableHash, SFS] = *It;
 
     // Group stable functions by ModuleIdentifier.
     llvm::stable_sort(SFS, [&](const std::unique_ptr<StableFunctionEntry> &L,
@@ -278,7 +236,7 @@ void StableFunctionMap::finalize(bool SkipTrim) {
       }
     }
     if (Invalid) {
-      ToDelete.push_back(It);
+      HashToFuncs.erase(It);
       continue;
     }
 
@@ -290,10 +248,8 @@ void StableFunctionMap::finalize(bool SkipTrim) {
     removeIdenticalIndexPair(SFS);
 
     if (!isProfitable(SFS))
-      ToDelete.push_back(It);
+      HashToFuncs.erase(It);
   }
-  for (auto It : ToDelete)
-    HashToFuncs.erase(It);
 
   Finalized = true;
 }