[CallGraph] Update callback call sites in RefreshCallGraph

[jdoerfert]

Summary

Since D82572, we keep "reference" edges for callback call sites. While
not strictly necessary they can improve the traversal order. However, we
did not update them properly in case a pass removed the callback call
site which caused a verification error (PR46687). With this patch we
update these reference edges properly during the invocation of
CallGraphSCCPass::RefreshCallGraph in non-checking mode.

[jdoerfert]

The test is actually not testing what it should. I failed to create one so far.

[ye-luo]

Assertion failure in PR46687 goes away with this patch.

[sdmitriev]

LGTM

llvm/lib/Analysis/CallGraphSCCPass.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Analysis/CallGraphSCCPass.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/SCCIterator.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/CallGraph.h"
+#include "llvm/IR/AbstractCallSite.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRPrintingPasses.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/LegacyPassManagers.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/OptBisect.h"
 #include "llvm/IR/PassTimingInfo.h"
for (CallGraphSCC::iterator SCCIdx = CurSCC.begin(), E = CurSCC.end();
 
     // Walk the function body looking for call sites.  Sync up the call sites in
     // CGN with those actually in the function.
 
     // Keep track of the number of direct and indirect calls that were
     // invalidated and removed.
     unsigned NumDirectRemoved = 0, NumIndirectRemoved = 0;
 
+    CallGraphNode::iterator CGNEnd = CGN->end();
+
+    auto RemoveAndCheckForDone = [&](CallGraphNode::iterator I) {
+      // Just remove the edge from the set of callees, keep track of whether
+      // I points to the last element of the vector.
+      bool WasLast = I + 1 == CGNEnd;
+      CGN->removeCallEdge(I);
+
+      // If I pointed to the last element of the vector, we have to bail out:
+      // iterator checking rejects comparisons of the resultant pointer with
+      // end.
+      if (WasLast)
+        return true;
+
+      CGNEnd = CGN->end();
+      return false;
+    };
+
     // Get the set of call sites currently in the function.
-    for (CallGraphNode::iterator I = CGN->begin(), E = CGN->end(); I != E; ) {
-      // Skip "reference" call records that do not have call instruction.
+    for (CallGraphNode::iterator I = CGN->begin(); I != CGNEnd;) {
+      // Delete "reference" call records that do not have call instruction. We
+      // reinsert them as needed later. However, keep them in checking mode.
       if (!I->first) {
+        if (CheckingMode) {
           ++I;
           continue;
         }
+        if (RemoveAndCheckForDone(I))
+          break;
+        continue;
+      }
 
       // If this call site is null, then the function pass deleted the call
       // entirely and the WeakTrackingVH nulled it out.
       auto *Call = dyn_cast_or_null<CallBase>(*I->first);
       if (!Call ||
           // If we've already seen this call site, then the FunctionPass RAUW'd
           // one call with another, which resulted in two "uses" in the edge
           // list of the same call.
for (CallGraphNode::iterator I = CGN->begin(); I != CGNEnd;) {
                "CallGraphSCCPass did not update the CallGraph correctly!");
 
         // If this was an indirect call site, count it.
         if (!I->second->getFunction())
           ++NumIndirectRemoved;
         else
           ++NumDirectRemoved;
 
-        // Just remove the edge from the set of callees, keep track of whether
-        // I points to the last element of the vector.
-        bool WasLast = I + 1 == E;
-        CGN->removeCallEdge(I);
-
-        // If I pointed to the last element of the vector, we have to bail out:
-        // iterator checking rejects comparisons of the resultant pointer with
-        // end.
-        if (WasLast)
+        if (RemoveAndCheckForDone(I))
           break;
-        E = CGN->end();
         continue;
       }
 
       assert(!Calls.count(Call) && "Call site occurs in node multiple times");
 
       if (Call) {
         Function *Callee = Call->getCalledFunction();
         // Ignore intrinsics because they're not really function calls.
for (BasicBlock &BB : *F)
       for (Instruction &I : BB) {
         auto *Call = dyn_cast<CallBase>(&I);
         if (!Call)
           continue;
         Function *Callee = Call->getCalledFunction();
         if (Callee && Callee->isIntrinsic())
           continue;
 
+        // If we are not in checking mode, insert potential callback calls as
+        // references. This is not a requirement but helps to iterate over the
+        // functions in the right order.
+        if (!CheckingMode) {
+          forEachCallbackFunction(*Call, [&](Function *CB) {
+            CGN->addCalledFunction(nullptr, CG.getOrInsertFunction(CB));
+          });
+        }
+
         // If this call site already existed in the callgraph, just verify it
         // matches up to expectations and remove it from Calls.
         DenseMap<Value *, CallGraphNode *>::iterator ExistingIt =
             Calls.find(Call);
         if (ExistingIt != Calls.end()) {
           CallGraphNode *ExistingNode = ExistingIt->second;
 
           // Remove from Calls since we have now seen it.

llvm/test/Transforms/OpenMP/parallel_deletion_cg_update.ll

+; RUN: opt < %s -instcombine -attributor-cgscc -print-callgraph -disable-output -verify 2>&1 | FileCheck %s
+
+; CHECK: Call graph node <<null function>><<{{.*}}>>  #uses=0
+; CHECK:   CS<None> calls function 'dead_fork_call'
+; CHECK:   CS<None> calls function 'd'
+; CHECK:   CS<None> calls function '__kmpc_fork_call'
+; CHECK:   CS<None> calls function 'live_fork_call'
+; CHECK:   CS<None> calls function '.omp_outlined..1'
+;
+; CHECK: Call graph node for function: '.omp_outlined..1'<<{{.*}}>>  #uses=3
+; CHECK:   CS<{{.*}}> calls function 'd'
+;
+; CHECK: Call graph node for function: '__kmpc_fork_call'<<{{.*}}>>  #uses=3
+; CHECK:   CS<None> calls external node
+;
+; CHECK: Call graph node for function: 'd'<<{{.*}}>>  #uses=2
+; CHECK:   CS<None> calls external node
+;
+; CHECK: Call graph node for function: 'dead_fork_call'<<{{.*}}>>  #uses=1
+;
+; CHECK: Call graph node for function: 'dead_fork_call2'<<{{.*}}>>  #uses=0
+; CHECK:   CS<{{.*}}> calls function '__kmpc_fork_call'
+; CHECK:   CS<None> calls function '.omp_outlined..1'
+;
+; CHECK: Call graph node for function: 'live_fork_call'<<{{.*}}>>  #uses=1
+; CHECK:   CS<{{.*}}> calls function '__kmpc_fork_call'
+; CHECK:   CS<None> calls function '.omp_outlined..1'
+
+
+%struct.ident_t = type { i32, i32, i32, i32, i8* }
+
+@.str = private unnamed_addr constant [23 x i8] c";unknown;unknown;0;0;;\00", align 1
+@0 = private unnamed_addr global %struct.ident_t { i32 0, i32 2, i32 0, i32 0, i8* getelementptr inbounds ([23 x i8], [23 x i8]* @.str, i32 0, i32 0) }, align 8
+
+define dso_local void @dead_fork_call() {
+entry:
+  br i1 true, label %if.then, label %if.else
+
+if.then:                                          ; preds = %entry
+  br label %if.end
+
+if.else:                                          ; preds = %entry
+  call void @dead_fork_call2()
+  call void (%struct.ident_t*, i32, void (i32*, i32*, ...)*, ...) @__kmpc_fork_call(%struct.ident_t* @0, i32 0, void (i32*, i32*, ...)* bitcast (void (i32*, i32*)* @.omp_outlined..0 to void (i32*, i32*, ...)*))
+  br label %if.end
+
+if.end:                                           ; preds = %if.else, %if.then
+  ret void
+}
+
+define internal void @dead_fork_call2() {
+entry:
+  call void (%struct.ident_t*, i32, void (i32*, i32*, ...)*, ...) @__kmpc_fork_call(%struct.ident_t* @0, i32 0, void (i32*, i32*, ...)* bitcast (void (i32*, i32*)* @.omp_outlined..1 to void (i32*, i32*, ...)*))
+  ret void
+}
+
+define internal void @.omp_outlined..0(i32* noalias %.global_tid., i32* noalias %.bound_tid.) {
+entry:
+  %.global_tid..addr = alloca i32*, align 8
+  %.bound_tid..addr = alloca i32*, align 8
+  store i32* %.global_tid., i32** %.global_tid..addr, align 8
+  store i32* %.bound_tid., i32** %.bound_tid..addr, align 8
+  ret void
+}
+
+declare !callback !2 void @__kmpc_fork_call(%struct.ident_t*, i32, void (i32*, i32*, ...)*, ...)
+
+define dso_local void @live_fork_call() {
+entry:
+  call void (%struct.ident_t*, i32, void (i32*, i32*, ...)*, ...) @__kmpc_fork_call(%struct.ident_t* @0, i32 0, void (i32*, i32*, ...)* bitcast (void (i32*, i32*)* @.omp_outlined..1 to void (i32*, i32*, ...)*))
+  ret void
+}
+
+define internal void @.omp_outlined..1(i32* noalias %.global_tid., i32* noalias %.bound_tid.) {
+entry:
+  %.global_tid..addr = alloca i32*, align 8
+  %.bound_tid..addr = alloca i32*, align 8
+  store i32* %.global_tid., i32** %.global_tid..addr, align 8
+  store i32* %.bound_tid., i32** %.bound_tid..addr, align 8
+  call void (...) @d()
+  ret void
+}
+
+declare dso_local void @d(...)
+
+!0 = !{i32 1, !"wchar_size", i32 4}
+!1 = !{!"clang version 11.0.0"}
+!2 = !{!3}
+!3 = !{i64 2, i64 -1, i64 -1, i1 true}