Index: llvm/lib/Transforms/Utils/CloneFunction.cpp =================================================================== --- llvm/lib/Transforms/Utils/CloneFunction.cpp +++ llvm/lib/Transforms/Utils/CloneFunction.cpp @@ -765,37 +765,38 @@ DT->addNewBlock(NewPH, LoopDomBB); for (Loop *CurLoop : OrigLoop->getLoopsInPreorder()) { - for (BasicBlock *BB : CurLoop->getBlocks()) { - if (CurLoop != LI->getLoopFor(BB)) - continue; + Loop *&NewLoop = LMap[CurLoop]; + if (!NewLoop) { + NewLoop = LI->AllocateLoop(); - Loop *&NewLoop = LMap[CurLoop]; - if (!NewLoop) { - NewLoop = LI->AllocateLoop(); + // Establish the parent/child relationship. + Loop *OrigParent = CurLoop->getParentLoop(); + assert(OrigParent && "Could not find the original parent loop"); + Loop *NewParentLoop = LMap[OrigParent]; + assert(NewParentLoop && "Could not find the new parent loop"); - // Establish the parent/child relationship. - Loop *OrigParent = CurLoop->getParentLoop(); - assert(OrigParent && "Could not find the original parent loop"); - Loop *NewParentLoop = LMap[OrigParent]; - assert(NewParentLoop && "Could not find the new parent loop"); + NewParentLoop->addChildLoop(NewLoop); + } + } - NewParentLoop->addChildLoop(NewLoop); - } + for (BasicBlock *BB : OrigLoop->getBlocks()) { + Loop *CurLoop = LI->getLoopFor(BB); + Loop *&NewLoop = LMap[CurLoop]; + assert(NewLoop && "Expecting new loop to be allocated"); - BasicBlock *NewBB = CloneBasicBlock(BB, VMap, NameSuffix, F); - VMap[BB] = NewBB; + BasicBlock *NewBB = CloneBasicBlock(BB, VMap, NameSuffix, F); + VMap[BB] = NewBB; - // Update LoopInfo. - NewLoop->addBasicBlockToLoop(NewBB, *LI); - if (BB == CurLoop->getHeader()) - NewLoop->moveToHeader(NewBB); + // Update LoopInfo. + NewLoop->addBasicBlockToLoop(NewBB, *LI); + if (BB == CurLoop->getHeader()) + NewLoop->moveToHeader(NewBB); - // Add DominatorTree node. After seeing all blocks, update to correct - // IDom. - DT->addNewBlock(NewBB, NewPH); + // Add DominatorTree node. After seeing all blocks, update to correct + // IDom. + DT->addNewBlock(NewBB, NewPH); - Blocks.push_back(NewBB); - } + Blocks.push_back(NewBB); } for (BasicBlock *BB : OrigLoop->getBlocks()) { Index: llvm/unittests/Transforms/Utils/CloningTest.cpp =================================================================== --- llvm/unittests/Transforms/Utils/CloningTest.cpp +++ llvm/unittests/Transforms/Utils/CloningTest.cpp @@ -10,6 +10,8 @@ #include "llvm/ADT/STLExtras.h" #include "llvm/ADT/SmallPtrSet.h" #include "llvm/Analysis/DomTreeUpdater.h" +#include "llvm/Analysis/LoopInfo.h" +#include "llvm/AsmParser/Parser.h" #include "llvm/IR/Argument.h" #include "llvm/IR/Constant.h" #include "llvm/IR/DIBuilder.h" @@ -355,6 +357,91 @@ delete F; } +static void runWithLoopInfoAndDominatorTree( + Module &M, StringRef FuncName, + function_ref Test) { + auto *F = M.getFunction(FuncName); + ASSERT_NE(F, nullptr) << "Could not find " << FuncName; + + DominatorTree DT(*F); + LoopInfo LI(DT); + + Test(*F, LI, DT); +} + +static std::unique_ptr parseIR(LLVMContext &C, const char *IR) { + SMDiagnostic Err; + std::unique_ptr Mod = parseAssemblyString(IR, Err, C); + if (!Mod) + Err.print("CloneLoop", errs()); + return Mod; +} + +TEST(CloneLoop, CloneLoopNest) { + // Parse the module. + LLVMContext Context; + + std::unique_ptr M = parseIR( + Context, + R"(define void @foo(i32* %A, i32 %ub) { +entry: + %guardcmp = icmp slt i32 0, %ub + br i1 %guardcmp, label %for.outer.preheader, label %for.end +for.outer.preheader: + br label %for.outer +for.outer: + %j = phi i32 [ 0, %for.outer.preheader ], [ %inc.outer, %for.outer.latch ] + br i1 %guardcmp, label %for.inner.preheader, label %for.outer.latch +for.inner.preheader: + br label %for.inner +for.inner: + %i = phi i32 [ 0, %for.inner.preheader ], [ %inc, %for.inner ] + %idxprom = sext i32 %i to i64 + %arrayidx = getelementptr inbounds i32, i32* %A, i64 %idxprom + store i32 %i, i32* %arrayidx, align 4 + %inc = add nsw i32 %i, 1 + %cmp = icmp slt i32 %inc, %ub + br i1 %cmp, label %for.inner, label %for.inner.exit +for.inner.exit: + br label %for.outer.latch +for.outer.latch: + %inc.outer = add nsw i32 %j, 1 + %cmp.outer = icmp slt i32 %inc.outer, %ub + br i1 %cmp.outer, label %for.outer, label %for.outer.exit +for.outer.exit: + br label %for.end +for.end: + ret void +})" + ); + + runWithLoopInfoAndDominatorTree( + *M, "foo", [&](Function &F, LoopInfo &LI, DominatorTree &DT) { + Function::iterator FI = F.begin(); + // First basic block is entry - skip it. + BasicBlock *Preheader = &*(++FI); + BasicBlock *Header = &*(++FI); + assert(Header->getName() == "for.outer"); + Loop *L = LI.getLoopFor(Header); + EXPECT_NE(L, nullptr); + EXPECT_EQ(Header, L->getHeader()); + EXPECT_EQ(Preheader, L->getLoopPreheader()); + + ValueToValueMapTy VMap; + SmallVector ClonedLoopBlocks; + Loop *NewLoop = cloneLoopWithPreheader(Preheader, Preheader, L, VMap, + "", &LI, &DT, ClonedLoopBlocks); + EXPECT_NE(NewLoop, nullptr); + EXPECT_EQ(NewLoop->getSubLoops().size(), 1u); + Loop::block_iterator BI = NewLoop->block_begin(); + EXPECT_TRUE((*BI)->getName().startswith("for.outer")); + EXPECT_TRUE((*(++BI))->getName().startswith("for.inner.preheader")); + EXPECT_TRUE((*(++BI))->getName().startswith("for.inner")); + EXPECT_TRUE((*(++BI))->getName().startswith("for.inner.exit")); + EXPECT_TRUE((*(++BI))->getName().startswith("for.outer.latch")); + }); +} + class CloneFunc : public ::testing::Test { protected: void SetUp() override {