[NARY-REASSOCIATE][NFC] Simplify min/max handling

[ebrevnov]

Summary

In order to explore different variants of reassociation current implementation uses "swap in a loop" approach. Unfortunately, the implementation is more complicated than it could be. This is an attempt to streamline the code. New approach is to extract core functionality into a helper function and call it explicitly as many times as required.

[nikic]

LGTM, thanks!

llvm/lib/Transforms/Scalar/NaryReassociate.cpp

NaryReassociatePass::matchAndReassociateMinOrMax(Instruction *I,
   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));
+    if (auto *NewMinMax = dyn_cast_or_null<Instruction>(
+            tryReassociateMinOrMax(I, MinMaxMatcher, LHS, RHS)))
+      return NewMinMax;
+    if (auto *NewMinMax = dyn_cast_or_null<Instruction>(
+            tryReassociateMinOrMax(I, MinMaxMatcher, RHS, LHS)))
+      return NewMinMax;
   }
   return nullptr;
 }
 
 Instruction *NaryReassociatePass::tryReassociate(Instruction * I,
                                                  const SCEV *&OrigSCEV) {
 
   if (!SE->isSCEVable(I->getType()))

 //  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)) {
-      Value *C = RHS;
-      const SCEV *AExpr = SE->getSCEV(A), *BExpr = SE->getSCEV(B);
-      const SCEV *CExpr = SE->getSCEV(C);
-      for (unsigned int j = 0; j < 2; ++j) {
-        if (j == 0) {
-          if (BExpr == CExpr)
-            continue;
-          // Transform 'I = (A op B) op C' to 'I = (A op C) op B' on the
-          // first iteration.
-          std::swap(BExpr, CExpr);
-          std::swap(B, C);
-        } else {
-          if (AExpr == CExpr)
-            continue;
-          // Transform 'I = (A op C) op B' to 'I = (B op C) op A' on the second
-          // iteration.
-          std::swap(AExpr, CExpr);
-          std::swap(A, C);
-        }
 
-        // 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;
+  if (LHS->hasNUsesOrMore(3) ||
+      llvm::any_of(LHS->users(),
+                    [&](auto *U) {
+                      return U != I &&
+                             !(U->hasOneUser() && *U->users().begin() == I);
+                    }) ||
+      !match(LHS, m_MaxMin))
+    return nullptr;
 
-        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);
+  auto tryCombination = [&](Value *A, const SCEV *AExpr, Value *B,
+                            const SCEV *BExpr, Value *C,
+                            const SCEV *CExpr) -> Value * {
+    SmallVector<const SCEV *, 2> Ops1{BExpr, AExpr};
 
-        Instruction *R1MinMax = findClosestMatchingDominator(R1Expr, I);
 
-        if (!R1MinMax)
-          continue;
+      return nullptr;
 
-        LLVM_DEBUG(dbgs() << "NARY: Found common sub-expr: " << *R1MinMax
-                          << "\n");
+    LLVM_DEBUG(dbgs() << "NARY: Found common sub-expr: " << *R1MinMax << "\n");
 
-        SmallVector<const SCEV *, 2> Ops2{SE->getUnknown(C),
-                                          SE->getUnknown(R1MinMax)};
-        const SCEV *R2Expr = SE->getMinMaxExpr(SCEVType, Ops2);
 
-        Value *NewMinMax = Expander.expandCodeFor(R2Expr, I->getType(), I);
-        NewMinMax->setName(Twine(I->getName()).concat(".nary"));
+    SCEVExpander Expander(*SE, *DL, "nary-reassociate");
 
-        LLVM_DEBUG(dbgs() << "NARY: Deleting:  " << *I << "\n"
-                          << "NARY: Inserting: " << *NewMinMax << "\n");
-        return NewMinMax;
-      }
-    }
-    std::swap(LHS, RHS);
-  }
+  };
+
+  const SCEV *AExpr = SE->getSCEV(A);
+  const SCEV *BExpr = SE->getSCEV(B);
+  const SCEV *RHSExpr = SE->getSCEV(RHS);
+
+  if (BExpr != RHSExpr) {
+    // Try (A op RHS) op B
+    if (auto *NewMinMax = tryCombination(A, AExpr, RHS, RHSExpr, B, BExpr))
+      return NewMinMax;
+
+  if (AExpr != RHSExpr) {
+    // Try (RHS op B) op A
+    if (auto *NewMinMax = tryCombination(RHS, RHSExpr, B, BExpr, A, AExpr))
+      return NewMinMax;
+
   return nullptr;
 }