[PATCH] D33839: Correctly remap BlockAddress when using CodeExtractor

[serge via Phabricator via llvm-commits]

serge-sans-paille updated this revision to Diff 101356.

https://reviews.llvm.org/D33839

Files:
  lib/Transforms/Utils/CodeExtractor.cpp
  test/Transforms/CodeExtractor/BlockAddressreference.ll


Index: test/Transforms/CodeExtractor/BlockAddressreference.ll
===================================================================
--- test/Transforms/CodeExtractor/BlockAddressreference.ll
+++ test/Transforms/CodeExtractor/BlockAddressreference.ll
@@ -0,0 +1,33 @@
+; RUN: opt < %s -loop-extract -disable-output
+
+ at label = common local_unnamed_addr global i8* null
+
+define i32 @sterix(i32 %n) {
+entry:
+  %tobool = icmp ne i32 %n, 0
+  ; this blockaddress references a basic block that goes in the extracted loop
+  %cond = select i1 %tobool, i8* blockaddress(@sterix, %for.cond), i8* blockaddress(@sterix, %exit)
+  store i8* %cond, i8** @label
+  %cmp5 = icmp sgt i32 %n, 0
+  br i1 %cmp5, label %for.body, label %exit
+
+for.cond:
+  %mul = shl nsw i32 %s.06, 1
+  %exitcond = icmp eq i32 %inc, %n
+  br i1 %exitcond, label %exit.loopexit, label %for.body
+
+for.body:
+  %i.07 = phi i32 [ %inc, %for.cond ], [ 0, %entry ]
+  %s.06 = phi i32 [ %mul, %for.cond ], [ 1, %entry ]
+  %inc = add nuw nsw i32 %i.07, 1
+  br label %for.cond
+
+exit.loopexit:
+  %phitmp = icmp ne i32 %s.06, 2
+  %phitmp8 = zext i1 %phitmp to i32
+  br label %exit
+
+exit:
+  %s.1 = phi i32 [ 1, %entry ], [ %phitmp8, %exit.loopexit ]
+  ret i32 %s.1
+}
Index: lib/Transforms/Utils/CodeExtractor.cpp
===================================================================
--- lib/Transforms/Utils/CodeExtractor.cpp
+++ lib/Transforms/Utils/CodeExtractor.cpp
@@ -871,6 +871,15 @@
   if (BFI && NumExitBlocks > 1)
     calculateNewCallTerminatorWeights(codeReplacer, ExitWeights, BPI);
 
+  // If some blocks are captured by a block address, remap to new Function
+  for (auto &BB : *newFunction) {
+    SmallVector<User *, 8> Users(BB.user_begin(), BB.user_end());
+    for (auto *U : Users) {
+      if (auto *BA = dyn_cast<BlockAddress>(U))
+        BA->replaceAllUsesWith(BlockAddress::get(newFunction, &BB));
+    }
+  }
+
   // Loop over all of the PHI nodes in the header block, and change any
   // references to the old incoming edge to be the new incoming edge.
   for (BasicBlock::iterator I = header->begin(); isa<PHINode>(I); ++I) {


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D33839.101356.patch
Type: text/x-patch
Size: 2124 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20170604/fc56906f/attachment.bin>