1 files changed, 121 insertions, 10 deletions

diff --git a/llvm/lib/Transforms/Scalar/NaryReassociate.cpp b/llvm/lib/Transforms/Scalar/NaryReassociate.cpp
index 32bb62129e8f..ded5caf53b5a 100644
--- a/llvm/lib/Transforms/Scalar/NaryReassociate.cpp
+++ b/llvm/lib/Transforms/Scalar/NaryReassociate.cpp
@@ -80,6 +80,7 @@
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/ScalarEvolution.h"
+#include "llvm/Analysis/ScalarEvolutionExpressions.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/Analysis/ValueTracking.h"
@@ -106,6 +107,7 @@
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Transforms/Scalar.h"
 #include "llvm/Transforms/Utils/Local.h"
+#include "llvm/Transforms/Utils/ScalarEvolutionExpander.h"
 #include <cassert>
 #include <cstdint>
 
@@ -222,15 +224,14 @@ bool NaryReassociatePass::doOneIteration(Function &F) {
   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;
+    for (Instruction &OrigI : *BB) {
       const SCEV *OrigSCEV = nullptr;
-      if (Instruction *NewI = tryReassociate(OrigI, OrigSCEV)) {
+      if (Instruction *NewI = tryReassociate(&OrigI, OrigSCEV)) {
         Changed = true;
-        OrigI->replaceAllUsesWith(NewI);
+        OrigI.replaceAllUsesWith(NewI);
 
         // Add 'OrigI' to the list of dead instructions.
-        DeadInsts.push_back(WeakTrackingVH(OrigI));
+        DeadInsts.push_back(WeakTrackingVH(&OrigI));
         // Add the rewritten instruction to SeenExprs; the original
         // instruction is deleted.
         const SCEV *NewSCEV = SE->getSCEV(NewI);
@@ -258,7 +259,7 @@ bool NaryReassociatePass::doOneIteration(Function &F) {
         if (NewSCEV != OrigSCEV)
           SeenExprs[OrigSCEV].push_back(WeakTrackingVH(NewI));
       } else if (OrigSCEV)
-        SeenExprs[OrigSCEV].push_back(WeakTrackingVH(OrigI));
+        SeenExprs[OrigSCEV].push_back(WeakTrackingVH(&OrigI));
     }
   }
   // Delete all dead instructions from 'DeadInsts'.
@@ -269,6 +270,24 @@ bool NaryReassociatePass::doOneIteration(Function &F) {
   return Changed;
 }
 
+template <typename PredT>
+Instruction *
+NaryReassociatePass::matchAndReassociateMinOrMax(Instruction *I,
+                                                 const SCEV *&OrigSCEV) {
+  Value *LHS = nullptr;
+  Value *RHS = nullptr;
+
+  auto MinMaxMatcher =
+      MaxMin_match<ICmpInst, bind_ty<Value>, bind_ty<Value>, PredT>(
+          m_Value(LHS), m_Value(RHS));
+  if (match(I, MinMaxMatcher)) {
+    OrigSCEV = SE->getSCEV(I);
+    return dyn_cast_or_null<Instruction>(
+        tryReassociateMinOrMax(I, MinMaxMatcher, LHS, RHS));
+  }
+  return nullptr;
+}
+
 Instruction *NaryReassociatePass::tryReassociate(Instruction * I,
                                                  const SCEV *&OrigSCEV) {
 
@@ -284,10 +303,21 @@ Instruction *NaryReassociatePass::tryReassociate(Instruction * I,
     OrigSCEV = SE->getSCEV(I);
     return tryReassociateGEP(cast<GetElementPtrInst>(I));
   default:
-    return nullptr;
+    break;
   }
 
-  llvm_unreachable("should not be reached");
+  // Try to match signed/unsigned Min/Max.
+  Instruction *ResI = nullptr;
+  // TODO: Currently min/max reassociation is restricted to integer types only
+  // due to use of SCEVExpander which my introduce incompatible forms of min/max
+  // for pointer types.
+  if (I->getType()->isIntegerTy())
+    if ((ResI = matchAndReassociateMinOrMax<umin_pred_ty>(I, OrigSCEV)) ||
+        (ResI = matchAndReassociateMinOrMax<smin_pred_ty>(I, OrigSCEV)) ||
+        (ResI = matchAndReassociateMinOrMax<umax_pred_ty>(I, OrigSCEV)) ||
+        (ResI = matchAndReassociateMinOrMax<smax_pred_ty>(I, OrigSCEV)))
+      return ResI;
+
   return nullptr;
 }
 
@@ -364,8 +394,8 @@ NaryReassociatePass::tryReassociateGEPAtIndex(GetElementPtrInst *GEP,
   // 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));
+  for (Use &Index : GEP->indices())
+    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) &&
@@ -540,3 +570,84 @@ NaryReassociatePass::findClosestMatchingDominator(const SCEV *CandidateExpr,
   }
   return nullptr;
 }
+
+template <typename MaxMinT> static SCEVTypes convertToSCEVype(MaxMinT &MM) {
+  if (std::is_same<smax_pred_ty, typename MaxMinT::PredType>::value)
+    return scSMaxExpr;
+  else if (std::is_same<umax_pred_ty, typename MaxMinT::PredType>::value)
+    return scUMaxExpr;
+  else if (std::is_same<smin_pred_ty, typename MaxMinT::PredType>::value)
+    return scSMinExpr;
+  else if (std::is_same<umin_pred_ty, typename MaxMinT::PredType>::value)
+    return scUMinExpr;
+
+  llvm_unreachable("Can't convert MinMax pattern to SCEV type");
+  return scUnknown;
+}
+
+// Parameters:
+//  I - instruction matched by MaxMinMatch matcher
+//  MaxMinMatch - min/max idiom matcher
+//  LHS - first operand of I
+//  RHS - second operand of I
+template <typename MaxMinT>
+Value *NaryReassociatePass::tryReassociateMinOrMax(Instruction *I,
+                                                   MaxMinT MaxMinMatch,
+                                                   Value *LHS, Value *RHS) {
+  Value *A = nullptr, *B = nullptr;
+  MaxMinT m_MaxMin(m_Value(A), m_Value(B));
+  for (unsigned int i = 0; i < 2; ++i) {
+    if (!LHS->hasNUsesOrMore(3) && match(LHS, m_MaxMin)) {
+      const SCEV *AExpr = SE->getSCEV(A), *BExpr = SE->getSCEV(B);
+      const SCEV *RHSExpr = SE->getSCEV(RHS);
+      for (unsigned int j = 0; j < 2; ++j) {
+        if (j == 0) {
+          if (BExpr == RHSExpr)
+            continue;
+          // Transform 'I = (A op B) op RHS' to 'I = (A op RHS) op B' on the
+          // first iteration.
+          std::swap(BExpr, RHSExpr);
+        } else {
+          if (AExpr == RHSExpr)
+            continue;
+          // Transform 'I = (A op RHS) op B' 'I = (B op RHS) op A' on the second
+          // iteration.
+          std::swap(AExpr, RHSExpr);
+        }
+
+        // The optimization is profitable only if LHS can be removed in the end.
+        // In other words LHS should be used (directly or indirectly) by I only.
+        if (llvm::any_of(LHS->users(), [&](auto *U) {
+              return U != I && !(U->hasOneUser() && *U->users().begin() == I);
+            }))
+          continue;
+
+        SCEVExpander Expander(*SE, *DL, "nary-reassociate");
+        SmallVector<const SCEV *, 2> Ops1{ BExpr, AExpr };
+        const SCEVTypes SCEVType = convertToSCEVype(m_MaxMin);
+        const SCEV *R1Expr = SE->getMinMaxExpr(SCEVType, Ops1);
+
+        Instruction *R1MinMax = findClosestMatchingDominator(R1Expr, I);
+
+        if (!R1MinMax)
+          continue;
+
+        LLVM_DEBUG(dbgs() << "NARY: Found common sub-expr: " << *R1MinMax
+                          << "\n");
+
+        R1Expr = SE->getUnknown(R1MinMax);
+        SmallVector<const SCEV *, 2> Ops2{ RHSExpr, R1Expr };
+        const SCEV *R2Expr = SE->getMinMaxExpr(SCEVType, Ops2);
+
+        Value *NewMinMax = Expander.expandCodeFor(R2Expr, I->getType(), I);
+        NewMinMax->setName(Twine(I->getName()).concat(".nary"));
+
+        LLVM_DEBUG(dbgs() << "NARY: Deleting:  " << *I << "\n"
+                          << "NARY: Inserting: " << *NewMinMax << "\n");
+        return NewMinMax;
+      }
+    }
+    std::swap(LHS, RHS);
+  }
+  return nullptr;
+}