diff --git a/llvm/include/llvm/Transforms/Utils/SSAUpdaterBulk.h b/llvm/include/llvm/Transforms/Utils/SSAUpdaterBulk.h --- a/llvm/include/llvm/Transforms/Utils/SSAUpdaterBulk.h +++ b/llvm/include/llvm/Transforms/Utils/SSAUpdaterBulk.h @@ -37,8 +37,12 @@ /// updates of different uses (which is not the case when traditional SSAUpdater /// is used). class SSAUpdaterBulk { + struct BlockInfo { + Value *Define = nullptr; + Value *Phi = nullptr; + }; struct RewriteInfo { - DenseMap Defines; + DenseMap Blocks; SmallVector Uses; StringRef Name; Type *Ty; @@ -49,7 +53,8 @@ PredIteratorCache PredCache; - Value *computeValueAt(BasicBlock *BB, RewriteInfo &R, DominatorTree *DT); + Value *computeValueAt(BasicBlock *BB, bool AtEnd, RewriteInfo &R, + DominatorTree *DT); public: explicit SSAUpdaterBulk() = default; diff --git a/llvm/lib/Transforms/Utils/SSAUpdaterBulk.cpp b/llvm/lib/Transforms/Utils/SSAUpdaterBulk.cpp --- a/llvm/lib/Transforms/Utils/SSAUpdaterBulk.cpp +++ b/llvm/lib/Transforms/Utils/SSAUpdaterBulk.cpp @@ -53,7 +53,7 @@ LLVM_DEBUG(dbgs() << "SSAUpdater: Var=" << Var << ": added new available value" << *V << " in " << BB->getName() << "\n"); - Rewrites[Var].Defines[BB] = V; + Rewrites[Var].Blocks[BB].Define = V; } /// Record a use of the symbolic value. This use will be updated with a @@ -65,19 +65,44 @@ Rewrites[Var].Uses.push_back(U); } -// Compute value at the given block BB. We either should already know it, or we -// should be able to recursively reach it going up dominator tree. -Value *SSAUpdaterBulk::computeValueAt(BasicBlock *BB, RewriteInfo &R, - DominatorTree *DT) { - if (!R.Defines.count(BB)) { - if (DT->isReachableFromEntry(BB) && PredCache.get(BB).size()) { - BasicBlock *IDom = DT->getNode(BB)->getIDom()->getBlock(); - Value *V = computeValueAt(IDom, R, DT); - R.Defines[BB] = V; - } else - R.Defines[BB] = UndefValue::get(R.Ty); +// Compute value for the given block BB. Depending on AtEnd, return either the +// value that is valid at the end of the block (i.e., after the definition in +// the block, if any) or at the beginning of the block (i.e., before any uses +// added in the block). +// +// We either should already know it, or we should be able to recursively reach +// it going up dominator tree. +Value *SSAUpdaterBulk::computeValueAt(BasicBlock *BB, bool AtEnd, + RewriteInfo &R, DominatorTree *DT) { + auto InfoIt = R.Blocks.find(BB); + if (InfoIt != R.Blocks.end()) { + BlockInfo &Info = InfoIt->second; + if (AtEnd) { + if (Info.Define) + return Info.Define; + + assert(Info.Phi != nullptr && "avoid creating unnecessary BlockInfos!"); + Info.Define = Info.Phi; + return Info.Define; + } + + if (Info.Phi) + return Info.Phi; } - return R.Defines[BB]; + + Value *V; + if (DT->isReachableFromEntry(BB) && PredCache.get(BB).size()) { + BasicBlock *IDom = DT->getNode(BB)->getIDom()->getBlock(); + V = computeValueAt(IDom, true, R, DT); + } else { + V = UndefValue::get(R.Ty); + } + + BlockInfo &Info = R.Blocks[BB]; + Info.Phi = V; + if (AtEnd) + Info.Define = V; + return V; } /// Given sets of UsingBlocks and DefBlocks, compute the set of LiveInBlocks. @@ -132,7 +157,7 @@ << " use(s)\n"); SmallPtrSet DefBlocks; - for (auto &Def : R.Defines) + for (auto &Def : R.Blocks) DefBlocks.insert(Def.first); IDF.setDefiningBlocks(DefBlocks); @@ -152,7 +177,7 @@ for (auto *FrontierBB : IDFBlocks) { IRBuilder<> B(FrontierBB, FrontierBB->begin()); PHINode *PN = B.CreatePHI(R.Ty, 0, R.Name); - R.Defines[FrontierBB] = PN; + R.Blocks[FrontierBB].Phi = PN; InsertedPHIsForVar.push_back(PN); if (InsertedPHIs) InsertedPHIs->push_back(PN); @@ -162,7 +187,7 @@ for (auto *PN : InsertedPHIsForVar) { BasicBlock *PBB = PN->getParent(); for (BasicBlock *Pred : PredCache.get(PBB)) - PN->addIncoming(computeValueAt(Pred, R, DT), Pred); + PN->addIncoming(computeValueAt(Pred, true, R, DT), Pred); } // Rewrite actual uses with the inserted definitions. @@ -170,7 +195,7 @@ for (Use *U : R.Uses) { if (!ProcessedUses.insert(U).second) continue; - Value *V = computeValueAt(getUserBB(U), R, DT); + Value *V = computeValueAt(getUserBB(U), false, R, DT); Value *OldVal = U->get(); assert(OldVal && "Invalid use!"); // Notify that users of the existing value that it is being replaced. diff --git a/llvm/unittests/Transforms/Utils/SSAUpdaterBulkTest.cpp b/llvm/unittests/Transforms/Utils/SSAUpdaterBulkTest.cpp --- a/llvm/unittests/Transforms/Utils/SSAUpdaterBulkTest.cpp +++ b/llvm/unittests/Transforms/Utils/SSAUpdaterBulkTest.cpp @@ -192,3 +192,76 @@ EXPECT_EQ(UpdatePhi->getIncomingValueForBlock(LoopStartBB), AddOp2); EXPECT_EQ(UpdatePhi->getIncomingValueForBlock(IfBB), UndefValue::get(I32Ty)); } + +TEST(SSAUpdaterBulk, Loop) { + SSAUpdaterBulk Updater; + LLVMContext C; + Module M("SSAUpdaterTest", C); + IRBuilder<> B(C); + Type *I32Ty = B.getInt32Ty(); + auto *F = Function::Create(FunctionType::get(I32Ty, {I32Ty}, false), + GlobalValue::ExternalLinkage, "F", &M); + + // Generate a small program with a simple loop: + // entry: + // br label %loop + // + // loop: + // %1 = add i32 %0, 1 + // br i1 true, label %loop, label %end + // + // end: + // ret i32 %0 + Argument *FirstArg = &*F->arg_begin(); + BasicBlock *Entry = BasicBlock::Create(C, "entry", F); + BasicBlock *Loop = BasicBlock::Create(C, "loop", F); + BasicBlock *End = BasicBlock::Create(C, "end", F); + + B.SetInsertPoint(Entry); + B.CreateBr(Loop); + + B.SetInsertPoint(Loop); + auto *Add = cast(B.CreateAdd(FirstArg, B.getInt32(0))); + B.CreateCondBr(B.getTrue(), Loop, End); + + B.SetInsertPoint(End); + ReturnInst *Return = B.CreateRet(FirstArg); + + // Now rewrite uses in instructions %1 and 'ret i32 %0'. The expected result + // is: + // entry: + // br label %loop + // + // loop: + // %1 = phi i32 [ 0, %entry ], [ %2, %loop ] + // %2 = add i32 %1, 1 + // br i1 true, label %loop, label %end + // + // end: + // ret i32 %2 + + // Add definitions and uses. + unsigned VarNum = Updater.AddVariable("c", I32Ty); + Updater.AddAvailableValue(VarNum, Entry, B.getInt32(0)); + Updater.AddAvailableValue(VarNum, Loop, Add); + + Updater.AddUse(VarNum, &Add->getOperandUse(0)); + Updater.AddUse(VarNum, &Return->getOperandUse(0)); + + // Save all inserted phis into a vector. + SmallVector Inserted; + DominatorTree DT(*F); + Updater.RewriteAllUses(&DT, &Inserted); + + // Only one phi should have been inserted, in the loop block. + EXPECT_EQ(Inserted.size(), 1u); + + PHINode *Phi = Inserted[0]; + EXPECT_EQ(Phi->getParent(), Loop); + EXPECT_EQ(Phi->getIncomingValueForBlock(Entry), B.getInt32(0)); + EXPECT_EQ(Phi->getIncomingValueForBlock(Loop), Add); + + // Verify rewrites of uses. + EXPECT_EQ(Add->getOperand(0), Phi); + EXPECT_EQ(Return->getOperand(0), Add); +}