[NFC][NARY-REASSOCIATE] Restructure code to aviod isPotentiallyReassociatable

[ebrevnov]

Summary

Currently we have to duplicate the same checks in isPotentiallyReassociatable and tryReassociate. With simple pattern like add/mul this may be not a big deal. But the situation gets much worse when I try to add support for min/max. Min/Max may be represented by several instructions and can take different forms. In order reduce complexity for upcoming min/max support we need to restructure the code a bit to avoid mentioned code duplication.

[mkazantsev]

LGTM

llvm/lib/Transforms/Scalar/NaryReassociate.cpp

290

Return after unreachable not needed.

[ebrevnov]

Rebase

[ebrevnov]

Rebase

llvm/include/llvm/Transforms/Scalar/NaryReassociate.h

bool runImpl(Function &F, AssumptionCache *AC_, DominatorTree *DT_,
                TargetTransformInfo *TTI_);
 
 private:
   // Runs only one iteration of the dominator-based algorithm. See the header
   // comments for why we need multiple iterations.
   bool doOneIteration(Function &F);
 
   // Reassociates I for better CSE.
-  Instruction *tryReassociate(Instruction *I);
+  Instruction *tryReassociate(Instruction *I, const SCEV *&OrigSCEV);
 
   // Reassociate GEP for better CSE.
   Instruction *tryReassociateGEP(GetElementPtrInst *GEP);
 
   // Try splitting GEP at the I-th index and see whether either part can be
   // CSE'ed. This is a helper function for tryReassociateGEP.
   //
   // \p IndexedType The element type indexed by GEP's I-th index. This is

llvm/lib/Transforms/Scalar/NaryReassociate.cpp

bool NaryReassociatePass::runImpl(Function &F, AssumptionCache *AC_,
   bool Changed = false, ChangedInThisIteration;
   do {
     ChangedInThisIteration = doOneIteration(F);
     Changed |= ChangedInThisIteration;
   } while (ChangedInThisIteration);
   return Changed;
 }
 
-// Explicitly list the instruction types NaryReassociate handles for now.
-static bool isPotentiallyNaryReassociable(Instruction *I) {
-  switch (I->getOpcode()) {
-  case Instruction::Add:
-  case Instruction::GetElementPtr:
-  case Instruction::Mul:
-    return true;
-  default:
-    return false;
-  }
-}
-
 bool NaryReassociatePass::doOneIteration(Function &F) {
   bool Changed = false;
   SeenExprs.clear();
   // Process the basic blocks in a depth first traversal of the dominator
   // tree. This order ensures that all bases of a candidate are in Candidates
   // when we process it.
   SmallVector<WeakTrackingVH, 16> DeadInsts;
   for (const auto Node : depth_first(DT)) {
     BasicBlock *BB = Node->getBlock();
     for (auto I = BB->begin(); I != BB->end(); ++I) {
       Instruction *OrigI = &*I;
-
-      if (!SE->isSCEVable(OrigI->getType()) ||
-          !isPotentiallyNaryReassociable(OrigI))
-        continue;
-
-      const SCEV *OrigSCEV = SE->getSCEV(OrigI);
-      if (Instruction *NewI = tryReassociate(OrigI)) {
+      const SCEV *OrigSCEV = nullptr;
+      if (Instruction *NewI = tryReassociate(OrigI, OrigSCEV)) {
         Changed = true;
         OrigI->replaceAllUsesWith(NewI);
 
         // Add 'OrigI' to the list of dead instructions.
         DeadInsts.push_back(WeakTrackingVH(OrigI));
         // Add the rewritten instruction to SeenExprs; the original
         // instruction is deleted.
         const SCEV *NewSCEV = SE->getSCEV(NewI);
for (auto I = BB->begin(); I != BB->end(); ++I) {
         // To alleviate this issue of ScalarEvolution not always capturing
         // equivalence, we add I to SeenExprs[OldSCEV] as well so that we can
         // map both SCEV before and after tryReassociate(I) to I.
         //
         // This improvement is exercised in @reassociate_gep_nsw in
         // nary-gep.ll.
         if (NewSCEV != OrigSCEV)
           SeenExprs[OrigSCEV].push_back(WeakTrackingVH(NewI));
-      } else
+      } else if (OrigSCEV)
         SeenExprs[OrigSCEV].push_back(WeakTrackingVH(OrigI));
-    }
-  }
   // Delete all dead instructions from 'DeadInsts'.
   // Please note ScalarEvolution is updated along the way.
   RecursivelyDeleteTriviallyDeadInstructionsPermissive(
       DeadInsts, TLI, nullptr, [this](Value *V) { SE->forgetValue(V); });
 
   return Changed;
 }
 
-Instruction *NaryReassociatePass::tryReassociate(Instruction *I) {
+Instruction *NaryReassociatePass::tryReassociate(Instruction * I,
+                                                 const SCEV * &OrigSCEV) {
+
+  if (!SE->isSCEVable(I->getType()))
+    return nullptr;
+
   switch (I->getOpcode()) {
   case Instruction::Add:
   case Instruction::Mul:
+    OrigSCEV = SE->getSCEV(I);
     return tryReassociateBinaryOp(cast<BinaryOperator>(I));
   case Instruction::GetElementPtr:
+    OrigSCEV = SE->getSCEV(I);
     return tryReassociateGEP(cast<GetElementPtrInst>(I));
   default:
-    llvm_unreachable("should be filtered out by isPotentiallyNaryReassociable");
+    return nullptr;
   }
+
+  llvm_unreachable("should not be reached");
+  return nullptr;
 }
 
 static bool isGEPFoldable(GetElementPtrInst *GEP,
                           const TargetTransformInfo *TTI) {
   SmallVector<const Value*, 4> Indices;
   for (auto I = GEP->idx_begin(); I != GEP->idx_end(); ++I)
     Indices.push_back(*I);
   return TTI->getGEPCost(GEP->getSourceElementType(), GEP->getPointerOperand(),
                          Indices) == TargetTransformInfo::TCC_Free;
 }
 
 Instruction *NaryReassociatePass::tryReassociateGEP(GetElementPtrInst *GEP) {
   // Not worth reassociating GEP if it is foldable.
   if (isGEPFoldable(GEP, TTI))
     return nullptr;
 
   gep_type_iterator GTI = gep_type_begin(*GEP);
   for (unsigned I = 1, E = GEP->getNumOperands(); I != E; ++I, ++GTI) {
     if (GTI.isSequential()) {
       if (auto *NewGEP = tryReassociateGEPAtIndex(GEP, I - 1,
                                                   GTI.getIndexedType())) {
         return NewGEP;
       }
     }
   }
   return nullptr;
 }
 
 bool NaryReassociatePass::requiresSignExtension(Value *Index,
                                                 GetElementPtrInst *GEP) {
   unsigned PointerSizeInBits =
       DL->getPointerSizeInBits(GEP->getType()->getPointerAddressSpace());
   return cast<IntegerType>(Index->getType())->getBitWidth() < PointerSizeInBits;
 }
 
 GetElementPtrInst *
 NaryReassociatePass::tryReassociateGEPAtIndex(GetElementPtrInst *GEP,
                                               unsigned I, Type *IndexedType) {
   Value *IndexToSplit = GEP->getOperand(I + 1);
   if (SExtInst *SExt = dyn_cast<SExtInst>(IndexToSplit)) {
     IndexToSplit = SExt->getOperand(0);
   } else if (ZExtInst *ZExt = dyn_cast<ZExtInst>(IndexToSplit)) {
     // zext can be treated as sext if the source is non-negative.
     if (isKnownNonNegative(ZExt->getOperand(0), *DL, 0, AC, GEP, DT))
       IndexToSplit = ZExt->getOperand(0);
   }
if (AddOperator *AO = dyn_cast<AddOperator>(IndexToSplit)) {
     }
   }
   return nullptr;
 }
 
 GetElementPtrInst *
 NaryReassociatePass::tryReassociateGEPAtIndex(GetElementPtrInst *GEP,
                                               unsigned I, Value *LHS,
                                               Value *RHS, Type *IndexedType) {
   // Look for GEP's closest dominator that has the same SCEV as GEP except that
   // the I-th index is replaced with LHS.
   SmallVector<const SCEV *, 4> IndexExprs;
   for (auto Index = GEP->idx_begin(); Index != GEP->idx_end(); ++Index)
     IndexExprs.push_back(SE->getSCEV(*Index));
   // Replace the I-th index with LHS.
   IndexExprs[I] = SE->getSCEV(LHS);
   if (isKnownNonNegative(LHS, *DL, 0, AC, GEP, DT) &&
NaryReassociatePass::tryReassociateGEPAtIndex(GetElementPtrInst *GEP,
   }
   GetElementPtrInst *NewGEP = cast<GetElementPtrInst>(
       Builder.CreateGEP(GEP->getResultElementType(), Candidate, RHS));
   NewGEP->setIsInBounds(GEP->isInBounds());
   NewGEP->takeName(GEP);
   return NewGEP;
 }
 
 Instruction *NaryReassociatePass::tryReassociateBinaryOp(BinaryOperator *I) {
   Value *LHS = I->getOperand(0), *RHS = I->getOperand(1);
   // There is no need to reassociate 0.
   if (SE->getSCEV(I)->isZero())
     return nullptr;
   if (auto *NewI = tryReassociateBinaryOp(LHS, RHS, I))
     return NewI;
   if (auto *NewI = tryReassociateBinaryOp(RHS, LHS, I))
     return NewI;
   return nullptr;
 }
 
 Instruction *NaryReassociatePass::tryReassociateBinaryOp(Value *LHS, Value *RHS,
                                                          BinaryOperator *I) {
   Value *A = nullptr, *B = nullptr;
   // To be conservative, we reassociate I only when it is the only user of (A op
   // B).
   if (LHS->hasOneUse() && matchTernaryOp(I, LHS, A, B)) {
     // I = (A op B) op RHS
     //   = (A op RHS) op B or (B op RHS) op A
     const SCEV *AExpr = SE->getSCEV(A), *BExpr = SE->getSCEV(B);
     const SCEV *RHSExpr = SE->getSCEV(RHS);
     if (BExpr != RHSExpr) {
       if (auto *NewI =
               tryReassociatedBinaryOp(getBinarySCEV(I, AExpr, RHSExpr), B, I))
         return NewI;
     }
     if (AExpr != RHSExpr) {
       if (auto *NewI =
               tryReassociatedBinaryOp(getBinarySCEV(I, BExpr, RHSExpr), A, I))
         return NewI;
     }
   }
   return nullptr;
 }
 
 Instruction *NaryReassociatePass::tryReassociatedBinaryOp(const SCEV *LHSExpr,
                                                           Value *RHS,
                                                           BinaryOperator *I) {
   // Look for the closest dominator LHS of I that computes LHSExpr, and replace
   // I with LHS op RHS.
   auto *LHS = findClosestMatchingDominator(LHSExpr, I);
   if (LHS == nullptr)
     return nullptr;
 
   Instruction *NewI = nullptr;
   switch (I->getOpcode()) {
   case Instruction::Add:
     NewI = BinaryOperator::CreateAdd(LHS, RHS, "", I);
     break;
   case Instruction::Mul:
     NewI = BinaryOperator::CreateMul(LHS, RHS, "", I);
     break;
   default:
     llvm_unreachable("Unexpected instruction.");
   }
   NewI->takeName(I);
   return NewI;
 }
 
 bool NaryReassociatePass::matchTernaryOp(BinaryOperator *I, Value *V,
                                          Value *&Op1, Value *&Op2) {
   switch (I->getOpcode()) {
   case Instruction::Add:
     return match(V, m_Add(m_Value(Op1), m_Value(Op2)));
   case Instruction::Mul:
     return match(V, m_Mul(m_Value(Op1), m_Value(Op2)));
   default:
     llvm_unreachable("Unexpected instruction.");
   }
   return false;
 }
 
 const SCEV *NaryReassociatePass::getBinarySCEV(BinaryOperator *I,
                                                const SCEV *LHS,
                                                const SCEV *RHS) {
   switch (I->getOpcode()) {
   case Instruction::Add:
     return SE->getAddExpr(LHS, RHS);
   case Instruction::Mul:
     return SE->getMulExpr(LHS, RHS);
   default:
     llvm_unreachable("Unexpected instruction.");
   }
   return nullptr;
 }
 
 Instruction *
 NaryReassociatePass::findClosestMatchingDominator(const SCEV *CandidateExpr,
                                                   Instruction *Dominatee) {
   auto Pos = SeenExprs.find(CandidateExpr);
   if (Pos == SeenExprs.end())
     return nullptr;
 
   auto &Candidates = Pos->second;
   // Because we process the basic blocks in pre-order of the dominator tree, a
   // candidate that doesn't dominate the current instruction won't dominate any
   // future instruction either. Therefore, we pop it out of the stack. This
   // optimization makes the algorithm O(n).
   while (!Candidates.empty()) {
     // Candidates stores WeakTrackingVHs, so a candidate can be nullptr if it's
     // removed
     // during rewriting.
     if (Value *Candidate = Candidates.back()) {
       Instruction *CandidateInstruction = cast<Instruction>(Candidate);
       if (DT->dominates(CandidateInstruction, Dominatee))
         return CandidateInstruction;
     }
     Candidates.pop_back();
   }
   return nullptr;
 }