[CloneFunction] Support BB == PredBB in DuplicateInstructionsInSplit.

[fhahn]

Summary

In case PredBB == BB and StopAt == BB's terminator, StopAt != &*BI will
fail, because BB's terminator instruction gets replaced.

By using BB.getTerminator() we get the current terminator which we can use to compare.

[junbuml]

Just want to make sure if this is the only case where StopAt is changed in the middle ?

[junbuml]

lib/Transforms/Utils/CloneFunction.cpp

816 ↗

(On Diff #136082)

Since StopAt is Instruction*, this function cannot clone the terminator itself. We can clone up to one instruction before the terminator. Curious if this is what this function was initially intended.

[fhahn]

In D43822#1020933, @junbuml wrote:

Just want to make sure if this is the only case where StopAt is changed in the middle ?

AFAIK this is the only case due to how SplitEdge behaves when BB == PredBB. The other instructions should not be moved around while splitting.

lib/Transforms/Utils/CloneFunction.cpp

816 ↗

(On Diff #136082)

Not sure in which cases it would make sense to clone the terminator, as the New BB needs it's own terminator to branch to BB.

[anna]

lib/Transforms/Utils/CloneFunction.cpp

814 ↗

(On Diff #136082)

Basically, StopAt was pointing to the old terminator of BB before it was replaced. SplitEdge does not update BB's terminator. So, where exactly, is the update happening for BB's terminator now?

If PredBB == BB, then BB's terminator should be updated from "br i1 <cond> BB, otherBB" to "br i1 <cond> NewBB, otherBB".

[fhahn]

lib/Transforms/Utils/CloneFunction.cpp

814 ↗

(On Diff #136082)

The case that caused the crash was PredBB == BB and BB has an unconditional branch as a terminator. Then we split BB at the terminator, which replaces the terminator of BB (relevant line in SplitEdge https://reviews.llvm.org/diffusion/L/browse/llvm/trunk/lib/Transforms/Utils/BasicBlockUtils.cpp;326450$273 , BasicBlock::splitBasicBlock, which is used by SplitBlock https://reviews.llvm.org/diffusion/L/browse/llvm/trunk/lib/IR/BasicBlock.cpp;326450$372)

[anna]

LGTM

lib/Transforms/Utils/CloneFunction.cpp

814 ↗

(On Diff #136082)

yes, that makes sense.

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

for (; PHINode *PN = dyn_cast<PHINode>(BI); ++BI)
     ValueMapping[PN] = PN->getIncomingValueForBlock(PredBB);
 
   BasicBlock *NewBB = SplitEdge(PredBB, BB);
   NewBB->setName(PredBB->getName() + ".split");
   Instruction *NewTerm = NewBB->getTerminator();
 
   // Clone the non-phi instructions of BB into NewBB, keeping track of the
   // mapping and using it to remap operands in the cloned instructions.
-  for (; StopAt != &*BI; ++BI) {
+  // Stop once we see the terminator too. This covers the case where BB's
+  // terminator gets replaced and StopAt == BB's terminator.
+  for (; StopAt != &*BI && BB->getTerminator() != &*BI; ++BI) {
     Instruction *New = BI->clone();
     New->setName(BI->getName());
     New->insertBefore(NewTerm);
     ValueMapping[&*BI] = New;
 
     // Remap operands to patch up intra-block references.
     for (unsigned i = 0, e = New->getNumOperands(); i != e; ++i)
       if (Instruction *Inst = dyn_cast<Instruction>(New->getOperand(i))) {
         auto I = ValueMapping.find(Inst);
         if (I != ValueMapping.end())
           New->setOperand(i, I->second);
       }
   }
 
   return NewBB;
 }

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

TEST_F(CloneInstruction, DuplicateInstructionsToSplit) {
   EXPECT_EQ(MulSplit->getParent(), Split);
 
   EXPECT_EQ(AddSplit->getNextNode(), MulSplit);
   EXPECT_EQ(MulSplit->getNextNode(), Split->getTerminator());
 
   delete F;
 }
 
+TEST_F(CloneInstruction, DuplicateInstructionsToSplitBlocksEq1) {
+  Type *ArgTy1[] = {Type::getInt32PtrTy(context)};
+  FunctionType *FT = FunctionType::get(Type::getVoidTy(context), ArgTy1, false);
+  V = new Argument(Type::getInt32Ty(context));
+
+  Function *F = Function::Create(FT, Function::ExternalLinkage);
+
+  BasicBlock *BB1 = BasicBlock::Create(context, "", F);
+  IRBuilder<> Builder1(BB1);
+
+  BasicBlock *BB2 = BasicBlock::Create(context, "", F);
+  IRBuilder<> Builder2(BB2);
+
+  Builder1.CreateBr(BB2);
+
+  Instruction *AddInst = cast<Instruction>(Builder2.CreateAdd(V, V));
+  Instruction *MulInst = cast<Instruction>(Builder2.CreateMul(AddInst, V));
+  Instruction *SubInst = cast<Instruction>(Builder2.CreateSub(MulInst, V));
+  Builder2.CreateBr(BB2);
+
+  ValueToValueMapTy Mapping;
+
+  auto Split = DuplicateInstructionsInSplitBetween(BB2, BB2, BB2->getTerminator(), Mapping);
+
+  EXPECT_TRUE(Split);
+  EXPECT_EQ(Mapping.size(), 3u);
+  EXPECT_TRUE(Mapping.find(AddInst) != Mapping.end());
+  EXPECT_TRUE(Mapping.find(MulInst) != Mapping.end());
+  EXPECT_TRUE(Mapping.find(SubInst) != Mapping.end());
+
+  auto AddSplit = dyn_cast<Instruction>(Mapping[AddInst]);
+  EXPECT_TRUE(AddSplit);
+  EXPECT_EQ(AddSplit->getOperand(0), V);
+  EXPECT_EQ(AddSplit->getOperand(1), V);
+  EXPECT_EQ(AddSplit->getParent(), Split);
+
+  auto MulSplit = dyn_cast<Instruction>(Mapping[MulInst]);
+  EXPECT_TRUE(MulSplit);
+  EXPECT_EQ(MulSplit->getOperand(0), AddSplit);
+  EXPECT_EQ(MulSplit->getOperand(1), V);
+  EXPECT_EQ(MulSplit->getParent(), Split);
+
+  auto SubSplit = dyn_cast<Instruction>(Mapping[SubInst]);
+  EXPECT_EQ(MulSplit->getNextNode(), SubSplit);
+  EXPECT_EQ(SubSplit->getNextNode(), Split->getTerminator());
+  EXPECT_EQ(Split->getSingleSuccessor(), BB2);
+  EXPECT_EQ(BB2->getSingleSuccessor(), Split);
+
+  delete F;
+}
+
+TEST_F(CloneInstruction, DuplicateInstructionsToSplitBlocksEq2) {
+  Type *ArgTy1[] = {Type::getInt32PtrTy(context)};
+  FunctionType *FT = FunctionType::get(Type::getVoidTy(context), ArgTy1, false);
+  V = new Argument(Type::getInt32Ty(context));
+
+  Function *F = Function::Create(FT, Function::ExternalLinkage);
+
+  BasicBlock *BB1 = BasicBlock::Create(context, "", F);
+  IRBuilder<> Builder1(BB1);
+
+  BasicBlock *BB2 = BasicBlock::Create(context, "", F);
+  IRBuilder<> Builder2(BB2);
+
+  Builder1.CreateBr(BB2);
+
+  Instruction *AddInst = cast<Instruction>(Builder2.CreateAdd(V, V));
+  Instruction *MulInst = cast<Instruction>(Builder2.CreateMul(AddInst, V));
+  Instruction *SubInst = cast<Instruction>(Builder2.CreateSub(MulInst, V));
+  Builder2.CreateBr(BB2);
+
+  ValueToValueMapTy Mapping;
+
+  auto Split = DuplicateInstructionsInSplitBetween(BB2, BB2, SubInst, Mapping);
+
+  EXPECT_TRUE(Split);
+  EXPECT_EQ(Mapping.size(), 2u);
+  EXPECT_TRUE(Mapping.find(AddInst) != Mapping.end());
+  EXPECT_TRUE(Mapping.find(MulInst) != Mapping.end());
+
+  auto AddSplit = dyn_cast<Instruction>(Mapping[AddInst]);
+  EXPECT_TRUE(AddSplit);
+  EXPECT_EQ(AddSplit->getOperand(0), V);
+  EXPECT_EQ(AddSplit->getOperand(1), V);
+  EXPECT_EQ(AddSplit->getParent(), Split);
+
+  auto MulSplit = dyn_cast<Instruction>(Mapping[MulInst]);
+  EXPECT_TRUE(MulSplit);
+  EXPECT_EQ(MulSplit->getOperand(0), AddSplit);
+  EXPECT_EQ(MulSplit->getOperand(1), V);
+  EXPECT_EQ(MulSplit->getParent(), Split);
+  EXPECT_EQ(MulSplit->getNextNode(), Split->getTerminator());
+  EXPECT_EQ(Split->getSingleSuccessor(), BB2);
+  EXPECT_EQ(BB2->getSingleSuccessor(), Split);
+
+  delete F;
+}
+
 class CloneFunc : public ::testing::Test {
 protected:
   void SetUp() override {
     SetupModule();
     CreateOldFunc();
     CreateNewFunc();
     SetupFinder();
   }