LowerTypeTests: Fix quadratic complexity (try 2).

[pcc]

Currently we have quadratic complexity in LowerTypeTests because
ScopedSaveAliaseesAndUsed loops over all aliases for each disjoint
set, and the number of aliases and number of disjoint sets is
roughly proportional to the program size. Fix that by moving
ScopedSaveAliaseesAndUsed to LowerTypeTestsModule::lower() so that
we do this only once.

Reland of #135875 with fix for bug that caused check-lld test failures.
The fix is to only remove functions from llvm.used/llvm.compiler.used
because buildBitSetsFromGlobalVariables, which now runs while
ScopedSaveAliaseesAndUsed is in scope, will delete global variables,
which would otherwise lead to a use-after-free when they are added
back to llvm.used or llvm.compiler.used.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Peter Collingbourne (pcc)

Changes

Currently we have quadratic complexity in LowerTypeTests because
ScopedSaveAliaseesAndUsed loops over all aliases for each disjoint
set, and the number of aliases and number of disjoint sets is
roughly proportional to the program size. Fix that by moving
ScopedSaveAliaseesAndUsed to LowerTypeTestsModule::lower() so that
we do this only once.

Reland of #135875 with fix for bug that caused check-lld test failures.
The fix is to only remove functions from llvm.used/llvm.compiler.used
because buildBitSetsFromGlobalVariables, which now runs while
ScopedSaveAliaseesAndUsed is in scope, will delete global variables,
which would otherwise lead to a use-after-free when they are added
back to llvm.used or llvm.compiler.used.

---

Full diff: https://github.com/llvm/llvm-project/pull/136053.diff

1 Files Affected:

• (modified) llvm/lib/Transforms/IPO/LowerTypeTests.cpp (+104-87)

diff --git a/llvm/lib/Transforms/IPO/LowerTypeTests.cpp b/llvm/lib/Transforms/IPO/LowerTypeTests.cpp
index 7cf7d74acfcfa..604ca7f8973b4 100644
--- a/llvm/lib/Transforms/IPO/LowerTypeTests.cpp
+++ b/llvm/lib/Transforms/IPO/LowerTypeTests.cpp
@@ -352,6 +352,27 @@ struct ScopedSaveAliaseesAndUsed {
   std::vector<std::pair<GlobalAlias *, Function *>> FunctionAliases;
   std::vector<std::pair<GlobalIFunc *, Function *>> ResolverIFuncs;
 
+  // This function only removes functions from llvm.used and llvm.compiler.used.
+  // We cannot remove global variables because they need to follow RAUW, as
+  // they may be deleted by buildBitSetsFromGlobalVariables.
+  void collectAndEraseUsedFunctions(Module &M,
+                                    SmallVectorImpl<GlobalValue *> &Vec,
+                                    bool CompilerUsed) {
+    auto *GV = collectUsedGlobalVariables(M, Vec, CompilerUsed);
+    if (!GV)
+      return;
+    GV->eraseFromParent();
+    auto NonFuncBegin =
+        std::stable_partition(Vec.begin(), Vec.end(), [](GlobalValue *GV) {
+          return isa<Function>(GV);
+        });
+    if (CompilerUsed)
+      appendToCompilerUsed(M, {NonFuncBegin, Vec.end()});
+    else
+      appendToUsed(M, {NonFuncBegin, Vec.end()});
+    Vec.resize(NonFuncBegin - Vec.begin());
+  }
+
   ScopedSaveAliaseesAndUsed(Module &M) : M(M) {
     // The users of this class want to replace all function references except
     // for aliases and llvm.used/llvm.compiler.used with references to a jump
@@ -365,10 +386,8 @@ struct ScopedSaveAliaseesAndUsed {
     // llvm.used/llvm.compiler.used and aliases, erase the used lists, let RAUW
     // replace the aliasees and then set them back to their original values at
     // the end.
-    if (GlobalVariable *GV = collectUsedGlobalVariables(M, Used, false))
-      GV->eraseFromParent();
-    if (GlobalVariable *GV = collectUsedGlobalVariables(M, CompilerUsed, true))
-      GV->eraseFromParent();
+    collectAndEraseUsedFunctions(M, Used, false);
+    collectAndEraseUsedFunctions(M, CompilerUsed, true);
 
     for (auto &GA : M.aliases()) {
       // FIXME: This should look past all aliases not just interposable ones,
@@ -1669,61 +1688,55 @@ void LowerTypeTestsModule::buildBitSetsFromFunctionsNative(
 
   lowerTypeTestCalls(TypeIds, JumpTable, GlobalLayout);
 
-  {
-    ScopedSaveAliaseesAndUsed S(M);
+  // Build aliases pointing to offsets into the jump table, and replace
+  // references to the original functions with references to the aliases.
+  for (unsigned I = 0; I != Functions.size(); ++I) {
+    Function *F = cast<Function>(Functions[I]->getGlobal());
+    bool IsJumpTableCanonical = Functions[I]->isJumpTableCanonical();
 
-    // Build aliases pointing to offsets into the jump table, and replace
-    // references to the original functions with references to the aliases.
-    for (unsigned I = 0; I != Functions.size(); ++I) {
-      Function *F = cast<Function>(Functions[I]->getGlobal());
-      bool IsJumpTableCanonical = Functions[I]->isJumpTableCanonical();
-
-      Constant *CombinedGlobalElemPtr = ConstantExpr::getInBoundsGetElementPtr(
-          JumpTableType, JumpTable,
-          ArrayRef<Constant *>{ConstantInt::get(IntPtrTy, 0),
-                               ConstantInt::get(IntPtrTy, I)});
-
-      const bool IsExported = Functions[I]->isExported();
-      if (!IsJumpTableCanonical) {
-        GlobalValue::LinkageTypes LT = IsExported
-                                           ? GlobalValue::ExternalLinkage
-                                           : GlobalValue::InternalLinkage;
-        GlobalAlias *JtAlias = GlobalAlias::create(F->getValueType(), 0, LT,
-                                                   F->getName() + ".cfi_jt",
-                                                   CombinedGlobalElemPtr, &M);
-        if (IsExported)
-          JtAlias->setVisibility(GlobalValue::HiddenVisibility);
-        else
-          appendToUsed(M, {JtAlias});
-      }
+    Constant *CombinedGlobalElemPtr = ConstantExpr::getInBoundsGetElementPtr(
+        JumpTableType, JumpTable,
+        ArrayRef<Constant *>{ConstantInt::get(IntPtrTy, 0),
+                             ConstantInt::get(IntPtrTy, I)});
+
+    const bool IsExported = Functions[I]->isExported();
+    if (!IsJumpTableCanonical) {
+      GlobalValue::LinkageTypes LT = IsExported ? GlobalValue::ExternalLinkage
+                                                : GlobalValue::InternalLinkage;
+      GlobalAlias *JtAlias = GlobalAlias::create(F->getValueType(), 0, LT,
+                                                 F->getName() + ".cfi_jt",
+                                                 CombinedGlobalElemPtr, &M);
+      if (IsExported)
+        JtAlias->setVisibility(GlobalValue::HiddenVisibility);
+      else
+        appendToUsed(M, {JtAlias});
+    }
 
-      if (IsExported) {
-        if (IsJumpTableCanonical)
-          ExportSummary->cfiFunctionDefs().emplace(F->getName());
-        else
-          ExportSummary->cfiFunctionDecls().emplace(F->getName());
-      }
+    if (IsExported) {
+      if (IsJumpTableCanonical)
+        ExportSummary->cfiFunctionDefs().emplace(F->getName());
+      else
+        ExportSummary->cfiFunctionDecls().emplace(F->getName());
+    }
 
-      if (!IsJumpTableCanonical) {
-        if (F->hasExternalWeakLinkage())
-          replaceWeakDeclarationWithJumpTablePtr(F, CombinedGlobalElemPtr,
-                                                 IsJumpTableCanonical);
-        else
-          replaceCfiUses(F, CombinedGlobalElemPtr, IsJumpTableCanonical);
-      } else {
-        assert(F->getType()->getAddressSpace() == 0);
-
-        GlobalAlias *FAlias =
-            GlobalAlias::create(F->getValueType(), 0, F->getLinkage(), "",
-                                CombinedGlobalElemPtr, &M);
-        FAlias->setVisibility(F->getVisibility());
-        FAlias->takeName(F);
-        if (FAlias->hasName())
-          F->setName(FAlias->getName() + ".cfi");
-        replaceCfiUses(F, FAlias, IsJumpTableCanonical);
-        if (!F->hasLocalLinkage())
-          F->setVisibility(GlobalVariable::HiddenVisibility);
-      }
+    if (!IsJumpTableCanonical) {
+      if (F->hasExternalWeakLinkage())
+        replaceWeakDeclarationWithJumpTablePtr(F, CombinedGlobalElemPtr,
+                                               IsJumpTableCanonical);
+      else
+        replaceCfiUses(F, CombinedGlobalElemPtr, IsJumpTableCanonical);
+    } else {
+      assert(F->getType()->getAddressSpace() == 0);
+
+      GlobalAlias *FAlias = GlobalAlias::create(
+          F->getValueType(), 0, F->getLinkage(), "", CombinedGlobalElemPtr, &M);
+      FAlias->setVisibility(F->getVisibility());
+      FAlias->takeName(F);
+      if (FAlias->hasName())
+        F->setName(FAlias->getName() + ".cfi");
+      replaceCfiUses(F, FAlias, IsJumpTableCanonical);
+      if (!F->hasLocalLinkage())
+        F->setVisibility(GlobalVariable::HiddenVisibility);
     }
   }
 
@@ -2339,39 +2352,43 @@ bool LowerTypeTestsModule::lower() {
   if (GlobalClasses.empty())
     return false;
 
-  // For each disjoint set we found...
-  for (const auto &C : GlobalClasses) {
-    if (!C->isLeader())
-      continue;
-
-    ++NumTypeIdDisjointSets;
-    // Build the list of type identifiers in this disjoint set.
-    std::vector<Metadata *> TypeIds;
-    std::vector<GlobalTypeMember *> Globals;
-    std::vector<ICallBranchFunnel *> ICallBranchFunnels;
-    for (auto M : GlobalClasses.members(*C)) {
-      if (isa<Metadata *>(M))
-        TypeIds.push_back(cast<Metadata *>(M));
-      else if (isa<GlobalTypeMember *>(M))
-        Globals.push_back(cast<GlobalTypeMember *>(M));
-      else
-        ICallBranchFunnels.push_back(cast<ICallBranchFunnel *>(M));
-    }
-
-    // Order type identifiers by unique ID for determinism. This ordering is
-    // stable as there is a one-to-one mapping between metadata and unique IDs.
-    llvm::sort(TypeIds, [&](Metadata *M1, Metadata *M2) {
-      return TypeIdInfo[M1].UniqueId < TypeIdInfo[M2].UniqueId;
-    });
+  {
+    ScopedSaveAliaseesAndUsed S(M);
+    // For each disjoint set we found...
+    for (const auto &C : GlobalClasses) {
+      if (!C->isLeader())
+        continue;
 
-    // Same for the branch funnels.
-    llvm::sort(ICallBranchFunnels,
-               [&](ICallBranchFunnel *F1, ICallBranchFunnel *F2) {
-                 return F1->UniqueId < F2->UniqueId;
-               });
+      ++NumTypeIdDisjointSets;
+      // Build the list of type identifiers in this disjoint set.
+      std::vector<Metadata *> TypeIds;
+      std::vector<GlobalTypeMember *> Globals;
+      std::vector<ICallBranchFunnel *> ICallBranchFunnels;
+      for (auto M : GlobalClasses.members(*C)) {
+        if (isa<Metadata *>(M))
+          TypeIds.push_back(cast<Metadata *>(M));
+        else if (isa<GlobalTypeMember *>(M))
+          Globals.push_back(cast<GlobalTypeMember *>(M));
+        else
+          ICallBranchFunnels.push_back(cast<ICallBranchFunnel *>(M));
+      }
 
-    // Build bitsets for this disjoint set.
-    buildBitSetsFromDisjointSet(TypeIds, Globals, ICallBranchFunnels);
+      // Order type identifiers by unique ID for determinism. This ordering is
+      // stable as there is a one-to-one mapping between metadata and unique
+      // IDs.
+      llvm::sort(TypeIds, [&](Metadata *M1, Metadata *M2) {
+        return TypeIdInfo[M1].UniqueId < TypeIdInfo[M2].UniqueId;
+      });
+
+      // Same for the branch funnels.
+      llvm::sort(ICallBranchFunnels,
+                 [&](ICallBranchFunnel *F1, ICallBranchFunnel *F2) {
+                   return F1->UniqueId < F2->UniqueId;
+                 });
+
+      // Build bitsets for this disjoint set.
+      buildBitSetsFromDisjointSet(TypeIds, Globals, ICallBranchFunnels);
+    }
   }
 
   allocateByteArrays();

[fmayer]

please confirm the ScopedSaveAliaseesAndUsed use on line 2063 (of RHS, 2050 of LHS) is OK with the semantics change

[pcc]

please confirm the ScopedSaveAliaseesAndUsed use on line 2063 (of RHS, 2050 of LHS) is OK with the semantics change

This change actually fixes a crash in the affected case. Previously this would crash:

$ cat au.c
__attribute__((alias("g"), used)) void f(void);
void g(void) {}
$ cat b.c
__attribute__((weak)) void f(void) {}
void f(void);

int main(int argc, char **argv) {
  argv[0] = (char *)f;
  void (*fp)(void) = (void (*)(void))argv;
  fp();
}
$ ~/l2/ra/bin/clang au.c b.c -flto=thin -fuse-ld=lld -fsanitize=cfi-icall -Wl,--save-temps
While deleting: ptr %
Use still stuck around after Def is destroyed:[1 x ptr] [ptr <badref>]

Now we at least don't crash but fail to link instead (this is a separate known issue):

~/l3/ra/bin/clang au.c b.c -flto=thin -fuse-ld=lld -fsanitize=cfi-icall -Wl,--save-temps
ld.lld: error: undefined hidden symbol: f.cfi
>>> referenced by ld-temp.o
>>>               a.out.lto.o:(g)
>>> did you mean: g.cfi
>>> defined in: /tmp/a.out.lto.au-fd4132.o
clang: error: linker command failed with exit code 1 (use -v to see invocation)

[fmayer]

Thanks. SGTM