[llvm] 32f5ee8 - [Attributor] Fixup block addresses after rewriting function signature

[Sergey Dmitriev via llvm-commits]

Author: Sergey Dmitriev
Date: 2020-05-12T13:53:04-07:00
New Revision: 32f5ee830b0f0e2426552d33174170554d8d336e

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

LOG: [Attributor] Fixup block addresses after rewriting function signature

Reviewers: jdoerfert, sstefan1, uenoku

Reviewed By: jdoerfert

Subscribers: hiraditya, uenoku, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D79801

Added: 
    

Modified: 
    llvm/lib/Transforms/IPO/Attributor.cpp
    llvm/test/Transforms/Attributor/misc_crash.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/IPO/Attributor.cpp b/llvm/lib/Transforms/IPO/Attributor.cpp
index c1a089b73488..a4f0965c27cf 100644
--- a/llvm/lib/Transforms/IPO/Attributor.cpp
+++ b/llvm/lib/Transforms/IPO/Attributor.cpp
@@ -1504,6 +1504,14 @@ ChangeStatus Attributor::rewriteFunctionSignatures(
     NewFn->getBasicBlockList().splice(NewFn->begin(),
                                       OldFn->getBasicBlockList());
 
+    // Fixup block addresses to reference new function.
+    SmallVector<BlockAddress *, 8u> BlockAddresses;
+    for (User *U : OldFn->users())
+      if (auto *BA = dyn_cast<BlockAddress>(U))
+        BlockAddresses.push_back(BA);
+    for (auto *BA : BlockAddresses)
+      BA->replaceAllUsesWith(BlockAddress::get(NewFn, BA->getBasicBlock()));
+
     // Set of all "call-like" instructions that invoke the old function mapped
     // to their new replacements.
     SmallVector<std::pair<CallBase *, CallBase *>, 8> CallSitePairs;

diff  --git a/llvm/test/Transforms/Attributor/misc_crash.ll b/llvm/test/Transforms/Attributor/misc_crash.ll
index 989b2bd8e204..9f15cdb24943 100644
--- a/llvm/test/Transforms/Attributor/misc_crash.ll
+++ b/llvm/test/Transforms/Attributor/misc_crash.ll
@@ -75,3 +75,36 @@ ret_bb:
   %2 = load i32, i32* @var2
   ret i32 %2
 }
+
+define void @func4() {
+; CHECK-LABEL: define {{[^@]+}}@func4
+; CHECK-NEXT:    call void @func5()
+; CHECK-NEXT:    ret void
+;
+  call void @func5(i32 0)
+  ret void
+}
+
+define internal void @func5(i32 %0) {
+; CHECK-LABEL: define {{[^@]+}}@func5
+; CHECK-NEXT:    [[TMP1:%.*]] = alloca i8*
+; CHECK-NEXT:    br label %block
+; CHECK:       block:
+; CHECK-NEXT:    store i8* blockaddress(@func5, %block), i8** [[TMP1]]
+; CHECK-NEXT:    [[TMP2:%.*]] = load i8*, i8** [[TMP1]]
+; CHECK-NEXT:    call void @func6(i8* [[TMP2]])
+; CHECK-NEXT:    ret void
+;
+  %tmp = alloca i8*
+  br label %block
+
+block:
+  store i8* blockaddress(@func5, %block), i8** %tmp
+  %addr = load i8*, i8** %tmp
+  call void @func6(i8* %addr)
+  ret void
+}
+
+; CHECK-LABEL: declare {{[^@]+}}@func6
+; CHECK-SAME: (i8*)
+declare void @func6(i8*)