[CodeExtractor] Fix multiple bugs under certain shape of extracted region

[myhsu]

Summary

If the extracted region has multiple exported data flows toward the same BB which is not included in the region, correct resotre instructions and PHI nodes won't be generated inside the exitStub. The solution is simply put the restore instructions right after the definition of output values instead of putting in exitStub.
Unittest for this bug is included

[davide]

I'm a little nervous about updating the dominator manually, still. I wonder if we can use the new API that Jakub implemented (cc: @kuhar / @dberlin )

[• dberlin]

Please don't update it manually.
Pretty much every manual update LLVM does has been proven incorrect over time.
Please just express the added/removed edges if you have to.
If you are cutting out a region, this should not be difficult.

[kuhar]

I have to agree with Danny, I'm not sure this version is correct. It should be fairly easy to use the batch updater here.
If we discover that doing incremental update makes the code too slow, that it's also great, as I still haven't come up a good benchmark to tell me what needs more optimization.

Nit: please run clang-format on your changes.

lib/Transforms/Utils/CodeExtractor.cpp

237 ↗

(On Diff #115402)

Does this handle the self-loop case?

244 ↗

(On Diff #115402)

In general, we prefer to update DFS numbers lazily and it's better to just invalidate them.

[kuhar]

lib/Transforms/Utils/CodeExtractor.cpp

538 ↗

(On Diff #115402)

You can use auto here, as the type is obvious looking at the right hand side.

[myhsu]

It turn out that the DomTree updating issue has actually been fixed in D32308
The unittest also pass after I applying that change
Sorry I didn't work incrementally on master branch at first place

[myhsu]

@kuhar @dberlin
Could you help reviewing my update?
Thanks : )

[chandlerc]

Can you revert the unrelated formatting changes? Otherwise reviewing this is a bit hard.

[myhsu]

@chandlerc
Temporarily revert to version before applying clang-format

[kuhar]

In D37902#878726, @myhsu wrote:

@chandlerc
Temporarily revert to version before applying clang-format

I think the point was to only format your changes, not the surrounding code. You can use the clang-format-diff.py script to to that.

lib/Transforms/Utils/CodeExtractor.cpp

785 ↗

(On Diff #116334)

Why is it the case that incrementing OAI works here and gets to OutputArgEnd? Same with the GEP::Create, why is passing the OAI pointer there valid if it doesn't get incremented and you use outputs[i] the same time?

This is not clear to me, could you consider adding a comment?
I think that the comment from the previous version of this snippet got lost here.

unittests/Transforms/Utils/CodeExtractor.cpp

56 ↗

(On Diff #116334)

nit: s/require/requires
Maybe it'd be better to phrase is like CodeExtractor requires the first element to dominate all the other ones?. Isn't an element really a BasicBlock here?

67 ↗

(On Diff #116334)

Why is FALSE expected here?

[myhsu]

@kuhar
Fix things mentioned in the inline comments

[myhsu]

lib/Transforms/Utils/CodeExtractor.cpp

785 ↗

(On Diff #116334)

The original code didn't check the OutputArgEnd either, however I add an assertion anyway.

OAI iterator shouldn't be increased in each iteration if AggregateArgs is true. Since under the configuration of using aggregate argument, there would be only one struct argument that accommodate all of the output values in struct fields, and OAI should always point to that struct argument.

The code comment in the original code seems to be not related to the aggregate argument condition, but I add few comment regarding this concern anyway

unittests/Transforms/Utils/CodeExtractor.cpp

67 ↗

(On Diff #116334)

verifyFunction would return false if no error

[myhsu]

Fix typo

[kuhar]

lib/Transforms/Utils/CodeExtractor.cpp

739 ↗

(On Diff #117387)

nit: s/reference/reference.
Comments should be full sentences and end with ..

764 ↗

(On Diff #117387)

nit: s/value/value.

776 ↗

(On Diff #117387)

nit: maybe s/amount/number?

786 ↗

(On Diff #117387)

Maybe s/increase/increment or even s/increase/advance?

787 ↗

(On Diff #117387)

nit: s/point/points
s/case/case.

unittests/Transforms/Utils/CodeExtractor.cpp

58 ↗

(On Diff #117387)

nit: s/ones/ones.

67 ↗

(On Diff #116334)

Maybe that's obvious for people more familiar with this part of the codebase, but I'd appreciate a comment explaining that here :)

[myhsu]

@kuhar
Fix typos

[kuhar]

Other than the two minor nits, the patch looks good to me.

lib/Transforms/Utils/CodeExtractor.cpp

768 ↗

(On Diff #117581)

nit: s/point/point.

771 ↗

(On Diff #117581)

nit: s/PHIs/PHIs.
You can also move this piece of comment to the line above, as you won't hit the 80-column limit (I think :P).

[myhsu]

Fix comment

[myhsu]

In D37902#887576, @kuhar wrote:

Other than the two minor nits, the patch looks good to me.

Thanks @kuhar :)

[myhsu]

@chandlerc @kuhar
I'm not pretty familiar the procedure yet, could someone help me commit the patch or merge it?
Thanks

[kuhar]

In D37902#889231, @myhsu wrote:

@chandlerc @kuhar
I'm not pretty familiar the procedure yet, could someone help me commit the patch or merge it?
Thanks

If you have commit access, you can go ahead and submit the patch. Otherwise, I can submit it for you.
Please take a look here: Obtaining Commit Access.

[myhsu]

In D37902#889297, @kuhar wrote:

In D37902#889231, @myhsu wrote:

@chandlerc @kuhar
I'm not pretty familiar the procedure yet, could someone help me commit the patch or merge it?
Thanks

If you have commit access, you can go ahead and submit the patch. Otherwise, I can submit it for you.
Please take a look here: Obtaining Commit Access.

Since this is my first time committing a patch, I guess I don't have enough record for being granted the commit access
Please submit this changed for me, thanks!

[kuhar]

Can you please rebase the patch so that is applies cleanly to the current ToT?

[myhsu]

Rebase to master

[kuhar]

With your patch, FindPhiPredForUseInBlock becomes unused and it causes the following warning to be emitted:

llvm/lib/Transforms/Utils/CodeExtractor.cpp:657:20: warning: unused function 'FindPhiPredForUseInBlock' [-Wunused-function]
static BasicBlock* FindPhiPredForUseInBlock(Value* Used, BasicBlock* BB) {

And that would make many buildbots very grumpy :)
I think that the function should be now deleted, as it has no users.

[myhsu]

Remove 'FindPhiPredForUsedInBlock' function

llvm/trunk/lib/Transforms/Utils/CodeExtractor.cpp

for (unsigned i = 0, e = Users.size(); i != e; ++i)
     if (TerminatorInst *TI = dyn_cast<TerminatorInst>(Users[i]))
       if (!Blocks.count(TI->getParent()) &&
           TI->getParent()->getParent() == oldFunction)
         TI->replaceUsesOfWith(header, newHeader);
 
   return newFunction;
 }
 
-/// FindPhiPredForUseInBlock - Given a value and a basic block, find a PHI
-/// that uses the value within the basic block, and return the predecessor
-/// block associated with that use, or return 0 if none is found.
-static BasicBlock* FindPhiPredForUseInBlock(Value* Used, BasicBlock* BB) {
-  for (Use &U : Used->uses()) {
-     PHINode *P = dyn_cast<PHINode>(U.getUser());
-     if (P && P->getParent() == BB)
-       return P->getIncomingBlock(U);
-  }
-
-  return nullptr;
-}
-
 /// emitCallAndSwitchStatement - This method sets up the caller side by adding
 /// the call instruction, splitting any PHI nodes in the header block as
 /// necessary.
 void CodeExtractor::
 emitCallAndSwitchStatement(Function *newFunction, BasicBlock *codeReplacer,
                            ValueSet &inputs, ValueSet &outputs) {
   // Emit a call to the new function, passing in: *pointer to struct (if
   // aggregating parameters), or plan inputs and allocated memory for outputs
CallInst *call = CallInst::Create(newFunction, params,
                                     NumExitBlocks > 1 ? "targetBlock" : "");
   codeReplacer->getInstList().push_back(call);
 
   Function::arg_iterator OutputArgBegin = newFunction->arg_begin();
   unsigned FirstOut = inputs.size();
   if (!AggregateArgs)
     std::advance(OutputArgBegin, inputs.size());
 
-  // Reload the outputs passed in by reference
+  // Reload the outputs passed in by reference.
+  Function::arg_iterator OAI = OutputArgBegin;
   for (unsigned i = 0, e = outputs.size(); i != e; ++i) {
     Value *Output = nullptr;
     if (AggregateArgs) {
       Value *Idx[2];
       Idx[0] = Constant::getNullValue(Type::getInt32Ty(Context));
       Idx[1] = ConstantInt::get(Type::getInt32Ty(Context), FirstOut + i);
       GetElementPtrInst *GEP = GetElementPtrInst::Create(
           StructArgTy, Struct, Idx, "gep_reload_" + outputs[i]->getName());
       codeReplacer->getInstList().push_back(GEP);
       Output = GEP;
     } else {
       Output = ReloadOutputs[i];
     }
     LoadInst *load = new LoadInst(Output, outputs[i]->getName()+".reload");
     Reloads.push_back(load);
     codeReplacer->getInstList().push_back(load);
     std::vector<User*> Users(outputs[i]->user_begin(), outputs[i]->user_end());
     for (unsigned u = 0, e = Users.size(); u != e; ++u) {
       Instruction *inst = cast<Instruction>(Users[u]);
       if (!Blocks.count(inst->getParent()))
         inst->replaceUsesOfWith(outputs[i], load);
     }
+
+    // Store to argument right after the definition of output value.
+    auto *OutI = dyn_cast<Instruction>(outputs[i]);
+    if (!OutI)
+      continue;
+    // Find proper insertion point.
+    Instruction *InsertPt = OutI->getNextNode();
+    // Let's assume that there is no other guy interleave non-PHI in PHIs.
+    if (isa<PHINode>(InsertPt))
+      InsertPt = InsertPt->getParent()->getFirstNonPHI();
+
+    assert(OAI != newFunction->arg_end() &&
+           "Number of output arguments should match "
+           "the amount of defined values");
+    if (AggregateArgs) {
+      Value *Idx[2];
+      Idx[0] = Constant::getNullValue(Type::getInt32Ty(Context));
+      Idx[1] = ConstantInt::get(Type::getInt32Ty(Context), FirstOut + i);
+      GetElementPtrInst *GEP = GetElementPtrInst::Create(
+          StructArgTy, &*OAI, Idx, "gep_" + outputs[i]->getName(), InsertPt);
+      new StoreInst(outputs[i], GEP, InsertPt);
+      // Since there should be only one struct argument aggregating
+      // all the output values, we shouldn't increment OAI, which always
+      // points to the struct argument, in this case.
+    } else {
+      new StoreInst(outputs[i], &*OAI, InsertPt);
+      ++OAI;
+    }
   }
 
   // Now we can emit a switch statement using the call as a value.
   SwitchInst *TheSwitch =
       SwitchInst::Create(Constant::getNullValue(Type::getInt16Ty(Context)),
                          codeReplacer, 0, codeReplacer);
 
   // Since there may be multiple exits from the original region, make the new
for (unsigned i = 0, e = TI->getNumSuccessors(); i != e; ++i)
           case 2:         // Conditional branch, return a bool
             brVal = ConstantInt::get(Type::getInt1Ty(Context), !SuccNum);
             break;
           default:
             brVal = ConstantInt::get(Type::getInt16Ty(Context), SuccNum);
             break;
           }
 
-          ReturnInst *NTRet = ReturnInst::Create(Context, brVal, NewTarget);
+          ReturnInst::Create(Context, brVal, NewTarget);
 
           // Update the switch instruction.
           TheSwitch->addCase(ConstantInt::get(Type::getInt16Ty(Context),
                                               SuccNum),
                              OldTarget);
 
-          // Restore values just before we exit
-          Function::arg_iterator OAI = OutputArgBegin;
-          for (unsigned out = 0, e = outputs.size(); out != e; ++out) {
-            // For an invoke, the normal destination is the only one that is
-            // dominated by the result of the invocation
-            BasicBlock *DefBlock = cast<Instruction>(outputs[out])->getParent();
-
-            bool DominatesDef = true;
-
-            BasicBlock *NormalDest = nullptr;
-            if (auto *Invoke = dyn_cast<InvokeInst>(outputs[out]))
-              NormalDest = Invoke->getNormalDest();
-
-            if (NormalDest) {
-              DefBlock = NormalDest;
-
-              // Make sure we are looking at the original successor block, not
-              // at a newly inserted exit block, which won't be in the dominator
-              // info.
-              for (const auto &I : ExitBlockMap)
-                if (DefBlock == I.second) {
-                  DefBlock = I.first;
-                  break;
-                }
-
-              // In the extract block case, if the block we are extracting ends
-              // with an invoke instruction, make sure that we don't emit a
-              // store of the invoke value for the unwind block.
-              if (!DT && DefBlock != OldTarget)
-                DominatesDef = false;
-            }
-
-            if (DT) {
-              DominatesDef = DT->dominates(DefBlock, OldTarget);
-              
-              // If the output value is used by a phi in the target block,
-              // then we need to test for dominance of the phi's predecessor
-              // instead.  Unfortunately, this a little complicated since we
-              // have already rewritten uses of the value to uses of the reload.
-              BasicBlock* pred = FindPhiPredForUseInBlock(Reloads[out], 
-                                                          OldTarget);
-              if (pred && DT && DT->dominates(DefBlock, pred))
-                DominatesDef = true;
-            }
-
-            if (DominatesDef) {
-              if (AggregateArgs) {
-                Value *Idx[2];
-                Idx[0] = Constant::getNullValue(Type::getInt32Ty(Context));
-                Idx[1] = ConstantInt::get(Type::getInt32Ty(Context),
-                                          FirstOut+out);
-                GetElementPtrInst *GEP = GetElementPtrInst::Create(
-                    StructArgTy, &*OAI, Idx, "gep_" + outputs[out]->getName(),
-                    NTRet);
-                new StoreInst(outputs[out], GEP, NTRet);
-              } else {
-                new StoreInst(outputs[out], &*OAI, NTRet);
-              }
-            }
-            // Advance output iterator even if we don't emit a store
-            if (!AggregateArgs) ++OAI;
-          }
         }
 
         // rewrite the original branch instruction with this new target
         TI->setSuccessor(i, NewTarget);
       }
   }
 
   // Now that we've done the deed, simplify the switch instruction.

llvm/trunk/unittests/Transforms/Utils/CMakeLists.txt

 set(LLVM_LINK_COMPONENTS
   Analysis
   AsmParser
   Core
   Support
   TransformUtils
   )
 
 add_llvm_unittest(UtilsTests
   ASanStackFrameLayoutTest.cpp
   Cloning.cpp
+  CodeExtractor.cpp
   FunctionComparator.cpp
   IntegerDivision.cpp
   Local.cpp
   OrderedInstructions.cpp
   ValueMapperTest.cpp
   )

llvm/trunk/unittests/Transforms/Utils/CodeExtractor.cpp

+//===- CodeExtractor.cpp - Unit tests for CodeExtractor -------------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Transforms/Utils/CodeExtractor.h"
+#include "llvm/AsmParser/Parser.h"
+#include "llvm/IR/BasicBlock.h"
+#include "llvm/IR/Dominators.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IR/Verifier.h"
+#include "llvm/IRReader/IRReader.h"
+#include "llvm/Support/SourceMgr.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+
+namespace {
+TEST(CodeExtractor, ExitStub) {
+  LLVMContext Ctx;
+  SMDiagnostic Err;
+  std::unique_ptr<Module> M(parseAssemblyString(R"invalid(
+    define i32 @foo(i32 %x, i32 %y, i32 %z) {
+    header:
+      %0 = icmp ugt i32 %x, %y
+      br i1 %0, label %body1, label %body2
+
+    body1:
+      %1 = add i32 %z, 2
+      br label %notExtracted
+
+    body2:
+      %2 = mul i32 %z, 7
+      br label %notExtracted
+
+    notExtracted:
+      %3 = phi i32 [ %1, %body1 ], [ %2, %body2 ]
+      %4 = add i32 %3, %x
+      ret i32 %4
+    }
+  )invalid",
+                                                Err, Ctx));
+
+  Function *Func = M->getFunction("foo");
+  SmallVector<BasicBlock *, 3> Candidates;
+  for (auto &BB : *Func) {
+    if (BB.getName() == "body1")
+      Candidates.push_back(&BB);
+    if (BB.getName() == "body2")
+      Candidates.push_back(&BB);
+  }
+  // CodeExtractor requires the first basic block
+  // to dominate all the other ones.
+  Candidates.insert(Candidates.begin(), &Func->getEntryBlock());
+
+  DominatorTree DT(*Func);
+  CodeExtractor CE(Candidates, &DT);
+  EXPECT_TRUE(CE.isEligible());
+
+  Function *Outlined = CE.extractCodeRegion();
+  EXPECT_TRUE(Outlined);
+  EXPECT_FALSE(verifyFunction(*Outlined));
+}
+} // end anonymous namespace