[llvm] d2d2fb1 - [Attributor][FIX] Allow dead users of rewritten function

[Johannes Doerfert via llvm-commits]

Author: Johannes Doerfert
Date: 2020-01-03T10:43:40-06:00
New Revision: d2d2fb19f7ea3b192247ca2d6995e66229366cec

URL: https://github.com/llvm/llvm-project/commit/d2d2fb19f7ea3b192247ca2d6995e66229366cec
DIFF: https://github.com/llvm/llvm-project/commit/d2d2fb19f7ea3b192247ca2d6995e66229366cec.diff

LOG: [Attributor][FIX] Allow dead users of rewritten function

If we replace a function with a new one because we rewrite the
signature, dead users may still refer to the old version. With this
patch we reuse the code that deals with dead functions, which the old
versions are, to avoid problems.

Added: 
    llvm/test/Transforms/Attributor/ArgumentPromotion/live_called_from_dead.ll

Modified: 
    llvm/lib/Transforms/IPO/Attributor.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/IPO/Attributor.cpp b/llvm/lib/Transforms/IPO/Attributor.cpp
index a8d7eb35b98e..ef52775c759d 100644
--- a/llvm/lib/Transforms/IPO/Attributor.cpp
+++ b/llvm/lib/Transforms/IPO/Attributor.cpp
@@ -5727,11 +5727,12 @@ ChangeStatus Attributor::run(Module &M) {
     }
   }
 
+  STATS_DECL(AAIsDead, Function, "Number of dead functions deleted.");
+  BUILD_STAT_NAME(AAIsDead, Function) += ToBeDeletedFunctions.size();
+
   // Rewrite the functions as requested during manifest.
   ManifestChange = ManifestChange | rewriteFunctionSignatures();
 
-  STATS_DECL(AAIsDead, Function, "Number of dead functions deleted.");
-  BUILD_STAT_NAME(AAIsDead, Function) += ToBeDeletedFunctions.size();
   for (Function *Fn : ToBeDeletedFunctions) {
     Fn->deleteBody();
     Fn->replaceAllUsesWith(UndefValue::get(Fn->getType()));
@@ -5890,8 +5891,9 @@ ChangeStatus Attributor::rewriteFunctionSignatures() {
     NewFn->getBasicBlockList().splice(NewFn->begin(),
                                       OldFn->getBasicBlockList());
 
-    // Set of all "call-like" instructions that invoke the old function.
-    SmallPtrSet<Instruction *, 8> OldCallSites;
+    // Set of all "call-like" instructions that invoke the old function mapped
+    // to their new replacements.
+    SmallVector<std::pair<CallBase *, CallBase *>, 8> CallSitePairs;
 
     // Callback to create a new "call-like" instruction for a given one.
     auto CallSiteReplacementCreator = [&](AbstractCallSite ACS) {
@@ -5943,7 +5945,6 @@ ChangeStatus Attributor::rewriteFunctionSignatures() {
       }
 
       // Copy over various properties and the new attributes.
-      OldCB->replaceAllUsesWith(NewCB);
       uint64_t W;
       if (OldCB->extractProfTotalWeight(W))
         NewCB->setProfWeight(W);
@@ -5954,10 +5955,7 @@ ChangeStatus Attributor::rewriteFunctionSignatures() {
           Ctx, OldCallAttributeList.getFnAttributes(),
           OldCallAttributeList.getRetAttributes(), NewArgOperandAttributes));
 
-      bool Inserted = OldCallSites.insert(OldCB).second;
-      assert(Inserted && "Call site was old twice!");
-      (void)Inserted;
-
+      CallSitePairs.push_back({OldCB, NewCB});
       return true;
     };
 
@@ -5984,11 +5982,15 @@ ChangeStatus Attributor::rewriteFunctionSignatures() {
     }
 
     // Eliminate the instructions *after* we visited all of them.
-    for (Instruction *OldCallSite : OldCallSites)
-      OldCallSite->eraseFromParent();
+    for (auto &CallSitePair : CallSitePairs) {
+      CallBase &OldCB = *CallSitePair.first;
+      CallBase &NewCB = *CallSitePair.second;
+      OldCB.replaceAllUsesWith(&NewCB);
+      OldCB.eraseFromParent();
+    }
+
+    ToBeDeletedFunctions.insert(OldFn);
 
-    assert(OldFn->getNumUses() == 0 && "Unexpected leftover uses!");
-    OldFn->eraseFromParent();
     Changed = ChangeStatus::CHANGED;
   }
 

diff  --git a/llvm/test/Transforms/Attributor/ArgumentPromotion/live_called_from_dead.ll b/llvm/test/Transforms/Attributor/ArgumentPromotion/live_called_from_dead.ll
new file mode 100644
index 000000000000..b81d35491f94
--- /dev/null
+++ b/llvm/test/Transforms/Attributor/ArgumentPromotion/live_called_from_dead.ll
@@ -0,0 +1,55 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -S -basicaa -attributor -attributor-disable=false -attributor-max-iterations-verify -attributor-max-iterations=3 < %s | FileCheck %s --check-prefixes=CHECK,OLDPM_MODULE
+; RUN: opt -S -passes='attributor' -aa-pipeline='basic-aa' -attributor-disable=false -attributor-max-iterations-verify -attributor-max-iterations=3 < %s | FileCheck %s --check-prefixes=CHECK,NEWPM_MODULE
+
+; OLDPM_MODULE-NOT: @dead
+; NEWPM_MODULE-NOT: @dead
+
+define internal void @dead() {
+  call i32 @test(i32* null, i32* null)
+  ret void
+}
+
+define internal i32 @test(i32* %X, i32* %Y) {
+; CHECK-LABEL: define {{[^@]+}}@test()
+; CHECK-NEXT:    br i1 true, label [[LIVE:%.*]], label [[DEAD:%.*]]
+; CHECK:       live:
+; CHECK-NEXT:    ret i32 0
+; CHECK:       dead:
+; CHECK-NEXT:    unreachable
+;
+  br i1 true, label %live, label %dead
+live:
+  ret i32 0
+dead:
+  call i32 @caller(i32* null)
+  call void @dead()
+  ret i32 1
+}
+
+define internal i32 @caller(i32* %B) {
+; CHECK-LABEL: define {{[^@]+}}@caller()
+; CHECK-NEXT:    [[A:%.*]] = alloca i32
+; CHECK-NEXT:    store i32 1, i32* [[A]], align 4
+; CHECK-NEXT:    [[C:%.*]] = call i32 @test()
+; CHECK-NEXT:    ret i32 0
+;
+  %A = alloca i32
+  store i32 1, i32* %A
+  %C = call i32 @test(i32* %A, i32* %B)
+  ret i32 %C
+}
+
+define i32 @callercaller() {
+; CHECK-LABEL: define {{[^@]+}}@callercaller()
+; CHECK-NEXT:    [[B:%.*]] = alloca i32
+; CHECK-NEXT:    store i32 2, i32* [[B]], align 4
+; CHECK-NEXT:    [[X:%.*]] = call i32 @caller()
+; CHECK-NEXT:    ret i32 0
+;
+  %B = alloca i32
+  store i32 2, i32* %B
+  %X = call i32 @caller(i32* %B)
+  ret i32 %X
+}
+