[IR] Track users of comdat

[nikic]

Summary

Track all GlobalObjects that reference a given comdat, which allows determining whether a function in a comdat is dead without scanning the whole module.

In particular, this makes filterDeadComdatFunctions() have complexity O(#DeadFunctions) rather than O(#SymbolsInModule), which addressed half of the compile-time issue exposed by D115545.

[nikic]

Add unit test for comdat user tracking.

[aeubanks]

presubmit is showing some lld test failures

[nikic]

In D115864#3200207, @aeubanks wrote:

presubmit is showing some lld test failures

Looks like the bitcode reader is doing some magic with comdats: https://github.com/llvm/llvm-project/blob/f78c1b07cb263e77731b6f43eedd345bdcc3be08/llvm/lib/Bitcode/Reader/BitcodeReader.cpp#L3296

[nikic]

Rebase over https://github.com/llvm/llvm-project/commit/18ab892ff7e9032914ff7fdb07685d5945c84fef.

[aeubanks]

very nice!

llvm/lib/Transforms/Utils/ModuleUtils.cpp

201

or doesn't have a comdat

llvm/include/llvm/IR/Comdat.h

 /// single COMDAT in LLVM.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_IR_COMDAT_H
 #define LLVM_IR_COMDAT_H
 
 #include "llvm-c/Types.h"
+#include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/Support/CBindingWrapping.h"
 
 namespace llvm {
 
+class GlobalObject;
 class raw_ostream;
 class StringRef;
 template <typename ValueTy> class StringMapEntry;
 
 // This is a Name X SelectionKind pair. The reason for having this be an
 // independent object instead of just adding the name and the SelectionKind
 // to a GlobalObject is that it is invalid to have two Comdats with the same
 // name but different SelectionKind. This structure makes that unrepresentable.
public:
   Comdat(const Comdat &) = delete;
   Comdat(Comdat &&C);
 
   SelectionKind getSelectionKind() const { return SK; }
   void setSelectionKind(SelectionKind Val) { SK = Val; }
   StringRef getName() const;
   void print(raw_ostream &OS, bool IsForDebug = false) const;
   void dump() const;
+  const SmallPtrSetImpl<GlobalObject *> &getUsers() const { return Users; }
 
 private:
   friend class Module;
+  friend class GlobalObject;
 
   Comdat();
+  void addUser(GlobalObject *GO);
+  void removeUser(GlobalObject *GO);
 
   // Points to the map in Module.
   StringMapEntry<Comdat> *Name = nullptr;
   SelectionKind SK = Any;
+  // Globals using this comdat.
+  SmallPtrSet<GlobalObject *, 2> Users;
 };
 
 // Create wrappers for C Binding types (see CBindingWrapping.h).
 DEFINE_SIMPLE_CONVERSION_FUNCTIONS(Comdat, LLVMComdatRef)
 
 inline raw_ostream &operator<<(raw_ostream &OS, const Comdat &C) {
   C.print(OS);
   return OS;
 }
 
 } // end namespace llvm
 
 #endif // LLVM_IR_COMDAT_H

llvm/include/llvm/IR/GlobalObject.h

 protected:
   GlobalObject(Type *Ty, ValueTy VTy, Use *Ops, unsigned NumOps,
                LinkageTypes Linkage, const Twine &Name,
                unsigned AddressSpace = 0)
       : GlobalValue(Ty, VTy, Ops, NumOps, Linkage, Name, AddressSpace),
         ObjComdat(nullptr) {
     setGlobalValueSubClassData(0);
   }
+  ~GlobalObject();
 
   Comdat *ObjComdat;
   enum {
     LastAlignmentBit = 5,
     HasSectionHashEntryBit,
 
     GlobalObjectBits,
   };
public:
   ///
   /// Setting the section to the empty string tells LLVM to choose an
   /// appropriate default object file section.
   void setSection(StringRef S);
 
   bool hasComdat() const { return getComdat() != nullptr; }
   const Comdat *getComdat() const { return ObjComdat; }
   Comdat *getComdat() { return ObjComdat; }
-  void setComdat(Comdat *C) { ObjComdat = C; }
+  void setComdat(Comdat *C);
 
   using Value::addMetadata;
   using Value::clearMetadata;
   using Value::eraseMetadata;
   using Value::getAllMetadata;
   using Value::getMetadata;
   using Value::hasMetadata;
   using Value::setMetadata;

llvm/lib/IR/Comdat.cpp

 using namespace llvm;
 
 Comdat::Comdat(Comdat &&C) : Name(C.Name), SK(C.SK) {}
 
 Comdat::Comdat() = default;
 
 StringRef Comdat::getName() const { return Name->first(); }
 
+void Comdat::addUser(GlobalObject *GO) { Users.insert(GO); }
+
+void Comdat::removeUser(GlobalObject *GO) { Users.erase(GO); }
+
 LLVMComdatRef LLVMGetOrInsertComdat(LLVMModuleRef M, const char *Name) {
   return wrap(unwrap(M)->getOrInsertComdat(Name));
 }
 
 LLVMComdatRef LLVMGetComdat(LLVMValueRef V) {
   GlobalObject *G = unwrap<GlobalObject>(V);
   return wrap(G->getComdat());
 }

llvm/lib/IR/Globals.cpp

case Value::NAME##Val:                                                       \
     return static_cast<NAME *>(this)->eraseFromParent();
 #include "llvm/IR/Value.def"
   default:
     break;
   }
   llvm_unreachable("not a global");
 }
 
+GlobalObject::~GlobalObject() { setComdat(nullptr); }
+
 bool GlobalValue::isInterposable() const {
   if (isInterposableLinkage(getLinkage()))
     return true;
   return getParent() && getParent()->getSemanticInterposition() &&
          !isDSOLocal();
 }
 
 bool GlobalValue::canBenefitFromLocalAlias() const {
if (auto *GA = dyn_cast<GlobalAlias>(this)) {
     return nullptr;
   }
   // ifunc and its resolver are separate things so don't use resolver comdat.
   if (isa<GlobalIFunc>(this))
     return nullptr;
   return cast<GlobalObject>(this)->getComdat();
 }
 
+void GlobalObject::setComdat(Comdat *C) {
+  if (ObjComdat)
+    ObjComdat->removeUser(this);
+  ObjComdat = C;
+  if (C)
+    C->addUser(this);
+}
+
 StringRef GlobalValue::getPartition() const {
   if (!hasPartition())
     return "";
   return getContext().pImpl->GlobalValuePartitions[this];
 }
 
 void GlobalValue::setPartition(StringRef S) {
   // Do nothing if we're clearing the partition and it is already empty.

llvm/lib/Transforms/Utils/ModuleUtils.cpp

llvm::getOrCreateSanitizerCtorAndInitFunctions(
   std::tie(Ctor, InitFunction) = llvm::createSanitizerCtorAndInitFunctions(
       M, CtorName, InitName, InitArgTypes, InitArgs, VersionCheckName);
   FunctionsCreatedCallback(Ctor, InitFunction);
   return std::make_pair(Ctor, InitFunction);
 }
 
 void llvm::filterDeadComdatFunctions(
     Module &M, SmallVectorImpl<Function *> &DeadComdatFunctions) {
-  // Build a map from the comdat to the number of entries in that comdat we
-  // think are dead. If this fully covers the comdat group, then the entire
-  // group is dead. If we find another entry in the comdat group though, we'll
-  // have to preserve the whole group.
-  SmallDenseMap<Comdat *, int, 16> ComdatEntriesCovered;
+  SmallPtrSet<Function *, 32> MaybeDeadFunctions;
+  SmallPtrSet<Comdat *, 32> MaybeDeadComdats;
   for (Function *F : DeadComdatFunctions) {
-    Comdat *C = F->getComdat();
-    assert(C && "Expected all input GVs to be in a comdat!");
-    ComdatEntriesCovered[C] += 1;
+    MaybeDeadFunctions.insert(F);
+    if (Comdat *C = F->getComdat())
+      MaybeDeadComdats.insert(C);
   }
 
-  auto CheckComdat = [&](Comdat &C) {
-    auto CI = ComdatEntriesCovered.find(&C);
-    if (CI == ComdatEntriesCovered.end())
-      return;
-
-    // If this could have been covered by a dead entry, just subtract one to
-    // account for it.
-    if (CI->second > 0) {
-      CI->second -= 1;
-      return;
-    }
-
-    // If we've already accounted for all the entries that were dead, the
-    // entire comdat is alive so remove it from the map.
-    ComdatEntriesCovered.erase(CI);
+  // Find comdats for which all users are dead now.
+  SmallPtrSet<Comdat *, 32> DeadComdats;
+  for (Comdat *C : MaybeDeadComdats) {
+    auto IsUserDead = [&](GlobalObject *GO) {
+      auto *F = dyn_cast<Function>(GO);
+      return F && MaybeDeadFunctions.contains(F);
     };
-
-  auto CheckAllComdats = [&] {
-    for (Function &F : M.functions())
-      if (Comdat *C = F.getComdat()) {
-        CheckComdat(*C);
-        if (ComdatEntriesCovered.empty())
-          return;
-      }
-    for (GlobalVariable &GV : M.globals())
-      if (Comdat *C = GV.getComdat()) {
-        CheckComdat(*C);
-        if (ComdatEntriesCovered.empty())
-          return;
-      }
-    for (GlobalAlias &GA : M.aliases())
-      if (Comdat *C = GA.getComdat()) {
-        CheckComdat(*C);
-        if (ComdatEntriesCovered.empty())
-          return;
-      }
-  };
-  CheckAllComdats();
+    if (all_of(C->getUsers(), IsUserDead))
+      DeadComdats.insert(C);
 
-  if (ComdatEntriesCovered.empty()) {
-    DeadComdatFunctions.clear();
-    return;
-  }
-
-  // Remove the entries that were not covering.
-  erase_if(DeadComdatFunctions, [&](GlobalValue *GV) {
-    return ComdatEntriesCovered.find(GV->getComdat()) ==
-           ComdatEntriesCovered.end();
+  // Only keep functions for which the comdat is dead.
+  erase_if(DeadComdatFunctions, [&](Function *F) {
+    Comdat *C = F->getComdat();
+    return C && !DeadComdats.contains(C);
   });
 }
 
 std::string llvm::getUniqueModuleId(Module *M) {
   MD5 Md5;
   bool ExportsSymbols = false;
   auto AddGlobal = [&](GlobalValue &GV) {
     if (GV.isDeclaration() || GV.getName().startswith("llvm.") ||

llvm/unittests/IR/ConstantsTest.cpp

for (auto EC : {ScalableEC, FixedEC}) {
           EXPECT_NE(nullptr, SplatVal);
           EXPECT_EQ(SplatVal, C);
         }
       }
     }
   }
 }
 
+TEST(ConstantsTest, ComdatUserTracking) {
+  LLVMContext Context;
+  Module M("MyModule", Context);
+
+  Comdat *C = M.getOrInsertComdat("comdat");
+  const SmallPtrSetImpl<GlobalObject *> &Users = C->getUsers();
+  EXPECT_TRUE(Users.size() == 0);
+
+  Type *Ty = Type::getInt8Ty(Context);
+  GlobalVariable *GV1 = cast<GlobalVariable>(M.getOrInsertGlobal("gv1", Ty));
+  GV1->setComdat(C);
+  EXPECT_TRUE(Users.size() == 1);
+  EXPECT_TRUE(Users.contains(GV1));
+
+  GlobalVariable *GV2 = cast<GlobalVariable>(M.getOrInsertGlobal("gv2", Ty));
+  GV2->setComdat(C);
+  EXPECT_TRUE(Users.size() == 2);
+  EXPECT_TRUE(Users.contains(GV2));
+
+  GV1->eraseFromParent();
+  EXPECT_TRUE(Users.size() == 1);
+  EXPECT_TRUE(Users.contains(GV2));
+
+  GV2->eraseFromParent();
+  EXPECT_TRUE(Users.size() == 0);
+}
+
 } // end anonymous namespace
 } // end namespace llvm