Index: lib/Transforms/Vectorize/LoopVectorize.cpp =================================================================== --- lib/Transforms/Vectorize/LoopVectorize.cpp +++ lib/Transforms/Vectorize/LoopVectorize.cpp @@ -6841,7 +6841,7 @@ // Note: first set Entry as region entry and then connect successors starting // from it in order, to propagate the "parent" of each VPBasicBlock. - VPBlockUtils::insertTwoBlocksAfter(Pred, Exit, Entry); + VPBlockUtils::insertTwoBlocksAfter(Pred, Exit, BlockInMask, Entry); VPBlockUtils::connectBlocks(Pred, Exit); return Region; Index: lib/Transforms/Vectorize/VPlan.h =================================================================== --- lib/Transforms/Vectorize/VPlan.h +++ lib/Transforms/Vectorize/VPlan.h @@ -339,6 +339,9 @@ /// List of successor blocks. SmallVector Successors; + /// Successor selector, null for zero or single successor blocks. + VPValue *CondBit = nullptr; + /// Add \p Successor as the last successor to this block. void appendSuccessor(VPBlockBase *Successor) { assert(Successor && "Cannot add nullptr successor!"); @@ -470,6 +473,13 @@ return getEnclosingBlockWithPredecessors()->getSinglePredecessor(); } + /// \return the condition bit selecting the successor. + VPValue *getCondBit() { return CondBit; } + + const VPValue *getCondBit() const { return CondBit; } + + void setCondBit(VPValue *CV) { CondBit = CV; } + /// Set a given VPBlockBase \p Successor as the single successor of this /// VPBlockBase. This VPBlockBase is not added as predecessor of \p Successor. /// This VPBlockBase must have no successors. @@ -479,11 +489,14 @@ } /// Set two given VPBlockBases \p IfTrue and \p IfFalse to be the two - /// successors of this VPBlockBase. This VPBlockBase is not added as - /// predecessor of \p IfTrue or \p IfFalse. This VPBlockBase must have no - /// successors. - void setTwoSuccessors(VPBlockBase *IfTrue, VPBlockBase *IfFalse) { + /// successors of this VPBlockBase. \p Condition is set as the successor + /// selector. This VPBlockBase is not added as predecessor of \p IfTrue or \p + /// IfFalse. This VPBlockBase must have no successors. + void setTwoSuccessors(VPBlockBase *IfTrue, VPBlockBase *IfFalse, + VPValue *Condition) { assert(Successors.empty() && "Setting two successors when others exist."); + assert(Condition && "Setting two successors without condition!"); + CondBit = Condition; appendSuccessor(IfTrue); appendSuccessor(IfFalse); } @@ -1265,9 +1278,10 @@ VPBlockUtils() = delete; /// Insert disconnected VPBlockBase \p NewBlock after \p BlockPtr. Add \p - /// NewBlock as successor of \p BlockPtr and \p Block as predecessor of \p - /// NewBlock, and propagate \p BlockPtr parent to \p NewBlock. \p NewBlock - /// must have neither successors nor predecessors. + /// NewBlock as successor of \p BlockPtr and \p BlockPtr as predecessor of \p + /// NewBlock, and propagate \p BlockPtr parent to \p NewBlock. If \p BlockPtr + /// has more than one successor, its conditional bit is propagated to \p + /// NewBlock. \p NewBlock must have neither successors nor predecessors. static void insertBlockAfter(VPBlockBase *NewBlock, VPBlockBase *BlockPtr) { assert(NewBlock->getSuccessors().empty() && "Can't insert new block with successors."); @@ -1282,16 +1296,16 @@ /// Insert disconnected VPBlockBases \p IfTrue and \p IfFalse after \p /// BlockPtr. Add \p IfTrue and \p IfFalse as succesors of \p BlockPtr and \p /// BlockPtr as predecessor of \p IfTrue and \p IfFalse. Propagate \p BlockPtr - /// parent to \p IfTrue and \p IfFalse. \p BlockPtr must have no successors - /// and \p IfTrue and \p IfFalse must have neither successors nor - /// predecessors. + /// parent to \p IfTrue and \p IfFalse. \p Condition is set as the successor + /// selector. \p BlockPtr must have no successors and \p IfTrue and \p IfFalse + /// must have neither successors nor predecessors. static void insertTwoBlocksAfter(VPBlockBase *IfTrue, VPBlockBase *IfFalse, - VPBlockBase *BlockPtr) { + VPValue *Condition, VPBlockBase *BlockPtr) { assert(IfTrue->getSuccessors().empty() && "Can't insert IfTrue with successors."); assert(IfFalse->getSuccessors().empty() && "Can't insert IfFalse with successors."); - BlockPtr->setTwoSuccessors(IfTrue, IfFalse); + BlockPtr->setTwoSuccessors(IfTrue, IfFalse, Condition); IfTrue->setPredecessors({BlockPtr}); IfFalse->setPredecessors({BlockPtr}); IfTrue->setParent(BlockPtr->getParent()); Index: lib/Transforms/Vectorize/VPlan.cpp =================================================================== --- lib/Transforms/Vectorize/VPlan.cpp +++ lib/Transforms/Vectorize/VPlan.cpp @@ -457,6 +457,18 @@ bumpIndent(1); for (const VPRecipeBase &Recipe : *BasicBlock) Recipe.print(OS, Indent); + + // Dump the condition bit. + const VPValue *CBV = BasicBlock->getCondBit(); + if (CBV) { + OS << " +\n" << Indent << " \"CondBit: "; + if (const VPInstruction *CBI = dyn_cast(CBV)) { + CBI->printAsOperand(OS); + OS << " (" << DOT::EscapeString(CBI->getParent()->getName()) << ")\\l\""; + } else + CBV->printAsOperand(OS); + } + bumpIndent(-2); OS << "\n" << Indent << "]\n"; dumpEdges(BasicBlock); Index: lib/Transforms/Vectorize/VPlanHCFGBuilder.cpp =================================================================== --- lib/Transforms/Vectorize/VPlanHCFGBuilder.cpp +++ lib/Transforms/Vectorize/VPlanHCFGBuilder.cpp @@ -77,12 +77,6 @@ }; } // anonymous namespace -// Return true if \p Inst is an incoming Instruction to be ignored in the VPlan -// representation. -static bool isInstructionToIgnore(Instruction *Inst) { - return isa(Inst); -} - // Set predecessors of \p VPBB in the same order as they are in \p BB. \p VPBB // must have no predecessors. void PlainCFGBuilder::setVPBBPredsFromBB(VPBasicBlock *VPBB, BasicBlock *BB) { @@ -197,16 +191,24 @@ VPIRBuilder.setInsertPoint(VPBB); for (Instruction &InstRef : *BB) { Instruction *Inst = &InstRef; - if (isInstructionToIgnore(Inst)) - continue; - // There should't be any VPValue for Inst at this point. Otherwise, we + // There shouldn't be any VPValue for Inst at this point. Otherwise, we // visited Inst when we shouldn't, breaking the RPO traversal order. assert(!IRDef2VPValue.count(Inst) && "Instruction shouldn't have been visited."); + if (auto *Br = dyn_cast(Inst)) { + // Branch instruction is not explicitly represented in VPlan but we need + // to represent its condition bit when it's conditional. + if (Br->isConditional()) + getOrCreateVPOperand(Br->getCondition()); + + // Skip the rest of the Instruction processing for Branch instructions. + continue; + } + VPInstruction *NewVPInst; - if (PHINode *Phi = dyn_cast(Inst)) { + if (auto *Phi = dyn_cast(Inst)) { // Phi node's operands may have not been visited at this point. We create // an empty VPInstruction that we will fix once the whole plain CFG has // been built. @@ -279,7 +281,19 @@ assert(SuccVPBB0 && "Successor 0 not found."); VPBasicBlock *SuccVPBB1 = getOrCreateVPBB(TI->getSuccessor(1)); assert(SuccVPBB1 && "Successor 1 not found."); - VPBB->setTwoSuccessors(SuccVPBB0, SuccVPBB1); + + // Get VPBB's condition bit. + assert(isa(TI) && "Unsupported terminator!"); + auto *Br = cast(TI); + Value *BrCond = Br->getCondition(); + // Look up the branch condition to get the corresponding VPValue + // representing the condition bit in VPlan (which may be in another VPBB). + assert(IRDef2VPValue.count(BrCond) && + "Missing condition bit in IRDef2VPValue!"); + VPValue *VPCondBit = IRDef2VPValue[BrCond]; + + // Link successors using condition bit. + VPBB->setTwoSuccessors(SuccVPBB0, SuccVPBB1, VPCondBit); } else llvm_unreachable("Number of successors not supported."); Index: lib/Transforms/Vectorize/VPlanVerifier.cpp =================================================================== --- lib/Transforms/Vectorize/VPlanVerifier.cpp +++ lib/Transforms/Vectorize/VPlanVerifier.cpp @@ -48,6 +48,12 @@ // Check block's parent. assert(VPB->getParent() == Region && "VPBlockBase has wrong parent"); + // Check block's condition bit. + if (VPB->getNumSuccessors() > 1) + assert(VPB->getCondBit() && "Missing condition bit!"); + else + assert(!VPB->getCondBit() && "Unexpected condition bit!"); + // Check block's successors. const auto &Successors = VPB->getSuccessors(); // There must be only one instance of a successor in block's successor list. Index: unittests/Transforms/Vectorize/VPlanHCFGTest.cpp =================================================================== --- unittests/Transforms/Vectorize/VPlanHCFGTest.cpp +++ unittests/Transforms/Vectorize/VPlanHCFGTest.cpp @@ -61,10 +61,11 @@ BasicBlock *LoopHeader = F->getEntryBlock().getSingleSuccessor(); auto Plan = doBuildPlan(LoopHeader); - VPBlockBase *Entry = Plan->getEntry()->getEntryBasicBlock(); + VPBasicBlock *Entry = Plan->getEntry()->getEntryBasicBlock(); EXPECT_NE(nullptr, Entry->getSingleSuccessor()); EXPECT_EQ(0u, Entry->getNumPredecessors()); EXPECT_EQ(1u, Entry->getNumSuccessors()); + EXPECT_EQ(nullptr, Entry->getCondBit()); VPBasicBlock *VecBB = Entry->getSingleSuccessor()->getEntryBasicBlock(); EXPECT_EQ(7u, VecBB->size()); @@ -105,6 +106,7 @@ EXPECT_EQ(Instruction::ICmp, ICmp->getOpcode()); EXPECT_EQ(2u, ICmp->getNumOperands()); EXPECT_EQ(IndvarAdd, ICmp->getOperand(0)); + EXPECT_EQ(VecBB->getCondBit(), ICmp); LoopVectorizationLegality::InductionList Inductions; SmallPtrSet DeadInstructions;