[NewPassManager] Resolve assertion in CGSCCPassManager when CallCounts change.

[asbirlea]

Summary

If the CallCounts change after an iteration of the DevirtSCCRepeatedPass, this is not reflected in the local CallCounts structure triggering the assertion checking the before/after sizes.
Since it is valid for the size to change and this only uses the CallCounts for the devirtualizing heuristic, keep a <Function*, CallCount> map instead, and make the devirtualizing decision using the counts for the functions that exist both before and after the pass.

Resolves PR42726.

[asbirlea]

Update test.

[chandlerc]

LGTM! Thanks for fixing this nasty bug (and sorry I wrote it).

include/llvm/Analysis/CGSCCPassManager.h

 ///    entity's shape changes, it and everything after it in post-order needs
 ///    to be visited to observe that shape.
 ///
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_ANALYSIS_CGSCCPASSMANAGER_H
 #define LLVM_ANALYSIS_CGSCCPASSMANAGER_H
 
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/PriorityWorklist.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Analysis/LazyCallGraph.h"
 #include "llvm/IR/CallSite.h"
 #include "llvm/IR/Function.h"
PreservedAnalyses run(LazyCallGraph::SCC &InitialC, CGSCCAnalysisManager &AM,
     };
 
     // Put value handles on all of the indirect calls and return the number of
     // direct calls for each function in the SCC.
     auto ScanSCC = [](LazyCallGraph::SCC &C,
                       SmallVectorImpl<WeakTrackingVH> &CallHandles) {
       assert(CallHandles.empty() && "Must start with a clear set of handles.");
 
-      SmallVector<CallCount, 4> CallCounts;
+      SmallDenseMap<Function *, CallCount> CallCounts;
+      CallCount CountLocal = {0, 0};
       for (LazyCallGraph::Node &N : C) {
-        CallCounts.push_back({0, 0});
-        CallCount &Count = CallCounts.back();
+        CallCount &Count =
+            CallCounts.insert(std::make_pair(&N.getFunction(), CountLocal))
+                .first->second;
         for (Instruction &I : instructions(N.getFunction()))
           if (auto CS = CallSite(&I)) {
             if (CS.getCalledFunction()) {
               ++Count.Direct;
             } else {
               ++Count.Indirect;
               CallHandles.push_back(WeakTrackingVH(&I));
             }
for (int Iteration = 0;; ++Iteration) {
       if (UR.UpdatedC && UR.UpdatedC != C) {
         PA.intersect(std::move(PassPA));
         break;
       }
 
       // Check that we didn't miss any update scenario.
       assert(!UR.InvalidatedSCCs.count(C) && "Processing an invalid SCC!");
       assert(C->begin() != C->end() && "Cannot have an empty SCC!");
-      assert((int)CallCounts.size() == C->size() &&
-             "Cannot have changed the size of the SCC!");
 
       // Check whether any of the handles were devirtualized.
       auto IsDevirtualizedHandle = [&](WeakTrackingVH &CallH) {
         if (!CallH)
           return false;
         auto CS = CallSite(CallH);
         if (!CS)
           return false;
 
         // If the call is still indirect, leave it alone.
         Function *F = CS.getCalledFunction();
         if (!F)
           return false;
 
-        LLVM_DEBUG(dbgs() << "Found devirutalized call from "
+        LLVM_DEBUG(dbgs() << "Found devirtualized call from "
                           << CS.getParent()->getParent()->getName() << " to "
                           << F->getName() << "\n");
 
         // We now have a direct call where previously we had an indirect call,
         // so iterate to process this devirtualization site.
         return true;
       };
       bool Devirt = llvm::any_of(CallHandles, IsDevirtualizedHandle);
 
       // Rescan to build up a new set of handles and count how many direct
       // calls remain. If we decide to iterate, this also sets up the input to
       // the next iteration.
       CallHandles.clear();
       auto NewCallCounts = ScanSCC(*C, CallHandles);
 
       // If we haven't found an explicit devirtualization already see if we
       // have decreased the number of indirect calls and increased the number
       // of direct calls for any function in the SCC. This can be fooled by all
       // manner of transformations such as DCE and other things, but seems to
       // work well in practice.
       if (!Devirt)
-        for (int i = 0, Size = C->size(); i < Size; ++i)
-          if (CallCounts[i].Indirect > NewCallCounts[i].Indirect &&
-              CallCounts[i].Direct < NewCallCounts[i].Direct) {
+        // Iterate over the keys in NewCallCounts, if Function also exists in
+        // CallCounts, make the check below.
+        for (auto &Pair : NewCallCounts) {
+          auto &CallCountNew = Pair.second;
+          auto CountIt = CallCounts.find(Pair.first);
+          if (CountIt != CallCounts.end()) {
+            const auto &CallCountOld = CountIt->second;
+            if (CallCountOld.Indirect > CallCountNew.Indirect &&
+                CallCountOld.Direct < CallCountNew.Direct) {
               Devirt = true;
               break;
             }
+          }
+        }
 
       if (!Devirt) {
         PA.intersect(std::move(PassPA));
         break;
       }
 
       // Otherwise, if we've already hit our max, we're done.
       if (Iteration >= MaxIterations) {

test/Other/new-pm-pr42726-cgscc.ll

+; RUN: opt -aa-pipeline=default -passes="default<O1>" %s -S | FileCheck %s
+
+declare void @bar()
+declare void @baz(i32*)
+
+; CHECK-LABEL: @foo1()
+define void @foo1() {
+entry:
+  %tag = alloca i32, align 4
+  call void @baz(i32* %tag)
+  %tmp = load i32, i32* %tag, align 4
+  switch i32 %tmp, label %sw.bb799 [
+    i32 10, label %sw.bb239
+  ]
+
+sw.bb239:
+  call void @foo2()
+  br label %cleanup871
+
+sw.bb799:
+  call void @foo3(i32 undef)
+  br label %cleanup871
+
+cleanup871:
+  call void @bar()
+  unreachable
+}
+
+define void @foo2() {
+  call void @foo4()
+  unreachable
+}
+
+define void @foo3(i32 %ptr) {
+  call void @foo1()
+  unreachable
+}
+
+define void @foo4() {
+entry:
+  %tag = alloca i32, align 4
+  call void @baz(i32* %tag)
+  %tmp = load i32, i32* %tag, align 4
+  switch i32 %tmp, label %sw.bb442 [
+    i32 16, label %sw.bb352
+  ]
+
+sw.bb352:
+  call void @foo3(i32 undef)
+  unreachable
+
+sw.bb442:
+  call void @foo2()
+  unreachable
+}
+