[InstCombine] fold min/max tree with common operand (PR35717)

[spatel]

Summary

This might be over the edge of what we want from instcombine, but I'm not sure. There is precedence for factorization transforms in instcombine for FP ops with fast-math.

I think it would take more work to add this to reassociate because that's specialized for binops, and min/max are not binops (or even single instructions). Also, I don't have evidence that larger min/max trees than this exist in real code.

In the motivating example from https://bugs.llvm.org/show_bug.cgi?id=35717 , we have:

int test(int xc, int xm, int xy) {
  int xk;
  if (xc < xm)
    xk = xc < xy ? xc : xy;
  else
    xk = xm < xy ? xm : xy;
  return xk;
}

This patch solves that problem because we recognize more min/max patterns after rL321672

https://rise4fun.com/Alive/Qjne
https://rise4fun.com/Alive/3yg

[spatel]

Ping.

Patch updated:
I adjusted the use checks to be '<='. This accounts for non-canonical min/max patterns as seen in the motivating test (the compare operands are not the same as the select operands). So now this patch should be the last step towards solving PR35717.

[hfinkel]

If I'm right about the easy FP cases, I'd like to see those handled in follow-up. Regardless, this LGTM.

lib/Transforms/InstCombine/InstCombineSelect.cpp

1297

I think that you can easily handle most of the interesting FP cases by checking all of the select pattern results for NaNBehavior == SPNB_NA || NaNBehavior == SPNB_RETURNS_ANY. The other cases will require further thought (so make sure we had the correct relative ordering of all the potentially-NaN inputs).

[spatel]

In D41603#969436, @hfinkel wrote:

If I'm right about the easy FP cases, I'd like to see those handled in follow-up. Regardless, this LGTM.

Thanks - yes, I'm planning on that; still working on the similar FP min/max enhancement for early-cse ( D41136 ) too.

lib/Transforms/InstCombine/InstCombineSelect.cpp

if (auto *X = isExtractFromCmpXchg(SI.getFalseValue(), 0))
     if (X == CmpXchg && X->getCompareOperand() == SI.getTrueValue()) {
       SI.setTrueValue(SI.getFalseValue());
       return &SI;
     }
 
   return nullptr;
 }
 
+/// Reduce a sequence of min/max with a common operand.
+static Instruction *factorizeMinMaxTree(SelectPatternFlavor SPF, Value *LHS,
+                                        Value *RHS,
+                                        InstCombiner::BuilderTy &Builder) {
+  assert(SelectPatternResult::isMinOrMax(SPF) && "Expected a min/max");
+  // TODO: Allow FP min/max with nnan/nsz.
+  if (!LHS->getType()->isIntOrIntVectorTy())
+    return nullptr;
+
+  // Match 3 of the same min/max ops. Example: umin(umin(), umin()).
+  Value *A, *B, *C, *D;
+  SelectPatternResult L = matchSelectPattern(LHS, A, B);
+  SelectPatternResult R = matchSelectPattern(RHS, C, D);
+  if (SPF != L.Flavor || L.Flavor != R.Flavor)
+    return nullptr;
+
+  // Look for a common operand. The use checks are different than usual because
+  // a min/max pattern typically has 2 uses of each op: 1 by the cmp and 1 by
+  // the select.
+  Value *MinMaxOp = nullptr;
+  Value *ThirdOp = nullptr;
+  if (LHS->getNumUses() <= 2 && RHS->getNumUses() > 2) {
+    // If the LHS is only used in this chain and the RHS is used outside of it,
+    // reuse the RHS min/max because that will eliminate the LHS.
+    if (D == A || C == A) {
+      // min(min(a, b), min(c, a)) --> min(min(c, a), b)
+      // min(min(a, b), min(a, d)) --> min(min(a, d), b)
+      MinMaxOp = RHS;
+      ThirdOp = B;
+    } else if (D == B || C == B) {
+      // min(min(a, b), min(c, b)) --> min(min(c, b), a)
+      // min(min(a, b), min(b, d)) --> min(min(b, d), a)
+      MinMaxOp = RHS;
+      ThirdOp = A;
+    }
+  } else if (RHS->getNumUses() <= 2) {
+    // Reuse the LHS. This will eliminate the RHS.
+    if (D == A || D == B) {
+      // min(min(a, b), min(c, a)) --> min(min(a, b), c)
+      // min(min(a, b), min(c, b)) --> min(min(a, b), c)
+      MinMaxOp = LHS;
+      ThirdOp = C;
+    } else if (C == A || C == B) {
+      // min(min(a, b), min(b, d)) --> min(min(a, b), d)
+      // min(min(a, b), min(c, b)) --> min(min(a, b), d)
+      MinMaxOp = LHS;
+      ThirdOp = D;
+    }
+  }
+  if (!MinMaxOp || !ThirdOp)
+    return nullptr;
+
+  CmpInst::Predicate P = getCmpPredicateForMinMax(SPF);
+  Value *CmpABC = Builder.CreateICmp(P, MinMaxOp, ThirdOp);
+  return SelectInst::Create(CmpABC, MinMaxOp, ThirdOp);
+}
+
 Instruction *InstCombiner::visitSelectInst(SelectInst &SI) {
   Value *CondVal = SI.getCondition();
   Value *TrueVal = SI.getTrueValue();
   Value *FalseVal = SI.getFalseValue();
   Type *SelType = SI.getType();
 
   // FIXME: Remove this workaround when freeze related patches are done.
   // For select with undef operand which feeds into an equality comparison,
if (SelectPatternResult::isMinOrMax(SPF)) {
 
         Value *NewSI = Builder.CreateSelect(Cmp, LHS, RHS, SI.getName(), &SI);
         if (!IsCastNeeded)
           return replaceInstUsesWith(SI, NewSI);
 
         Value *NewCast = Builder.CreateCast(CastOp, NewSI, SelType);
         return replaceInstUsesWith(SI, NewCast);
       }
+
+      if (Instruction *I = factorizeMinMaxTree(SPF, LHS, RHS, Builder))
+        return I;
     }
 
     if (SPF) {
       // MAX(MAX(a, b), a) -> MAX(a, b)
       // MIN(MIN(a, b), a) -> MIN(a, b)
       // MAX(MIN(a, b), a) -> a
       // MIN(MAX(a, b), a) -> a
       // ABS(ABS(a)) -> ABS(a)

test/Transforms/InstCombine/minmax-fold.ll

 ; PR35717: https://bugs.llvm.org/show_bug.cgi?id=35717
 
 ; min(min(a, b), min(b, c)) --> min(min(a, b), c)
 
 define i32 @common_factor_smin(i32 %a, i32 %b, i32 %c) {
 ; CHECK-LABEL: @common_factor_smin(
 ; CHECK-NEXT:    [[CMP_AB:%.*]] = icmp slt i32 %a, %b
 ; CHECK-NEXT:    [[MIN_AB:%.*]] = select i1 [[CMP_AB]], i32 %a, i32 %b
-; CHECK-NEXT:    [[CMP_BC:%.*]] = icmp slt i32 %b, %c
-; CHECK-NEXT:    [[MIN_BC:%.*]] = select i1 [[CMP_BC]], i32 %b, i32 %c
-; CHECK-NEXT:    [[CMP_AB_BC:%.*]] = icmp slt i32 [[MIN_AB]], [[MIN_BC]]
-; CHECK-NEXT:    [[MIN_ABC:%.*]] = select i1 [[CMP_AB_BC]], i32 [[MIN_AB]], i32 [[MIN_BC]]
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp slt i32 [[MIN_AB]], %c
+; CHECK-NEXT:    [[MIN_ABC:%.*]] = select i1 [[TMP1]], i32 [[MIN_AB]], i32 %c
 ; CHECK-NEXT:    ret i32 [[MIN_ABC]]
 ;
   %cmp_ab = icmp slt i32 %a, %b
   %min_ab = select i1 %cmp_ab, i32 %a, i32 %b
   %cmp_bc = icmp slt i32 %b, %c
   %min_bc = select i1 %cmp_bc, i32 %b, i32 %c
   %cmp_ab_bc = icmp slt i32 %min_ab, %min_bc
   %min_abc = select i1 %cmp_ab_bc, i32 %min_ab, i32 %min_bc
   ret i32 %min_abc
 }
 
 ; max(max(a, b), max(c, b)) --> max(max(a, b), c)
 
 define <2 x i32> @common_factor_smax(<2 x i32> %a, <2 x i32> %b, <2 x i32> %c) {
 ; CHECK-LABEL: @common_factor_smax(
 ; CHECK-NEXT:    [[CMP_AB:%.*]] = icmp sgt <2 x i32> %a, %b
 ; CHECK-NEXT:    [[MAX_AB:%.*]] = select <2 x i1> [[CMP_AB]], <2 x i32> %a, <2 x i32> %b
-; CHECK-NEXT:    [[CMP_CB:%.*]] = icmp sgt <2 x i32> %c, %b
-; CHECK-NEXT:    [[MAX_CB:%.*]] = select <2 x i1> [[CMP_CB]], <2 x i32> %c, <2 x i32> %b
-; CHECK-NEXT:    [[CMP_AB_CB:%.*]] = icmp sgt <2 x i32> [[MAX_AB]], [[MAX_CB]]
-; CHECK-NEXT:    [[MAX_ABC:%.*]] = select <2 x i1> [[CMP_AB_CB]], <2 x i32> [[MAX_AB]], <2 x i32> [[MAX_CB]]
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp sgt <2 x i32> [[MAX_AB]], %c
+; CHECK-NEXT:    [[MAX_ABC:%.*]] = select <2 x i1> [[TMP1]], <2 x i32> [[MAX_AB]], <2 x i32> %c
 ; CHECK-NEXT:    ret <2 x i32> [[MAX_ABC]]
 ;
   %cmp_ab = icmp sgt <2 x i32> %a, %b
   %max_ab = select <2 x i1> %cmp_ab, <2 x i32> %a, <2 x i32> %b
   %cmp_cb = icmp sgt <2 x i32> %c, %b
   %max_cb = select <2 x i1> %cmp_cb, <2 x i32> %c, <2 x i32> %b
   %cmp_ab_cb = icmp sgt <2 x i32> %max_ab, %max_cb
   %max_abc = select <2 x i1> %cmp_ab_cb, <2 x i32> %max_ab, <2 x i32> %max_cb
   ret <2 x i32> %max_abc
 }
 
 ; min(min(b, c), min(a, b)) --> min(min(b, c), a)
 
 define <2 x i32> @common_factor_umin(<2 x i32> %a, <2 x i32> %b, <2 x i32> %c) {
 ; CHECK-LABEL: @common_factor_umin(
 ; CHECK-NEXT:    [[CMP_BC:%.*]] = icmp ult <2 x i32> %b, %c
 ; CHECK-NEXT:    [[MIN_BC:%.*]] = select <2 x i1> [[CMP_BC]], <2 x i32> %b, <2 x i32> %c
-; CHECK-NEXT:    [[CMP_AB:%.*]] = icmp ult <2 x i32> %a, %b
-; CHECK-NEXT:    [[MIN_AB:%.*]] = select <2 x i1> [[CMP_AB]], <2 x i32> %a, <2 x i32> %b
-; CHECK-NEXT:    [[CMP_BC_AB:%.*]] = icmp ult <2 x i32> [[MIN_BC]], [[MIN_AB]]
-; CHECK-NEXT:    [[MIN_ABC:%.*]] = select <2 x i1> [[CMP_BC_AB]], <2 x i32> [[MIN_BC]], <2 x i32> [[MIN_AB]]
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp ult <2 x i32> [[MIN_BC]], %a
+; CHECK-NEXT:    [[MIN_ABC:%.*]] = select <2 x i1> [[TMP1]], <2 x i32> [[MIN_BC]], <2 x i32> %a
 ; CHECK-NEXT:    ret <2 x i32> [[MIN_ABC]]
 ;
   %cmp_bc = icmp ult <2 x i32> %b, %c
   %min_bc = select <2 x i1> %cmp_bc, <2 x i32> %b, <2 x i32> %c
   %cmp_ab = icmp ult <2 x i32> %a, %b
   %min_ab = select <2 x i1> %cmp_ab, <2 x i32> %a, <2 x i32> %b
   %cmp_bc_ab = icmp ult <2 x i32> %min_bc, %min_ab
   %min_abc = select <2 x i1> %cmp_bc_ab, <2 x i32> %min_bc, <2 x i32> %min_ab
   ret <2 x i32> %min_abc
 }
 
 ; max(max(b, c), max(b, a)) --> max(max(b, c), a)
 
 define i32 @common_factor_umax(i32 %a, i32 %b, i32 %c) {
 ; CHECK-LABEL: @common_factor_umax(
 ; CHECK-NEXT:    [[CMP_BC:%.*]] = icmp ugt i32 %b, %c
 ; CHECK-NEXT:    [[MAX_BC:%.*]] = select i1 [[CMP_BC]], i32 %b, i32 %c
-; CHECK-NEXT:    [[CMP_BA:%.*]] = icmp ugt i32 %b, %a
-; CHECK-NEXT:    [[MAX_BA:%.*]] = select i1 [[CMP_BA]], i32 %b, i32 %a
-; CHECK-NEXT:    [[CMP_BC_BA:%.*]] = icmp ugt i32 [[MAX_BC]], [[MAX_BA]]
-; CHECK-NEXT:    [[MAX_ABC:%.*]] = select i1 [[CMP_BC_BA]], i32 [[MAX_BC]], i32 [[MAX_BA]]
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp ugt i32 [[MAX_BC]], %a
+; CHECK-NEXT:    [[MAX_ABC:%.*]] = select i1 [[TMP1]], i32 [[MAX_BC]], i32 %a
 ; CHECK-NEXT:    ret i32 [[MAX_ABC]]
 ;
   %cmp_bc = icmp ugt i32 %b, %c
   %max_bc = select i1 %cmp_bc, i32 %b, i32 %c
   %cmp_ba = icmp ugt i32 %b, %a
   %max_ba = select i1 %cmp_ba, i32 %b, i32 %a
   %cmp_bc_ba = icmp ugt i32 %max_bc, %max_ba
   %max_abc = select i1 %cmp_bc_ba, i32 %max_bc, i32 %max_ba
   ret i32 %max_abc
 }
 
 declare void @extra_use(i32)
 
 define i32 @common_factor_umax_extra_use_lhs(i32 %a, i32 %b, i32 %c) {
 ; CHECK-LABEL: @common_factor_umax_extra_use_lhs(
 ; CHECK-NEXT:    [[CMP_BC:%.*]] = icmp ugt i32 %b, %c
 ; CHECK-NEXT:    [[MAX_BC:%.*]] = select i1 [[CMP_BC]], i32 %b, i32 %c
-; CHECK-NEXT:    [[CMP_BA:%.*]] = icmp ugt i32 %b, %a
-; CHECK-NEXT:    [[MAX_BA:%.*]] = select i1 [[CMP_BA]], i32 %b, i32 %a
-; CHECK-NEXT:    [[CMP_BC_BA:%.*]] = icmp ugt i32 [[MAX_BC]], [[MAX_BA]]
-; CHECK-NEXT:    [[MAX_ABC:%.*]] = select i1 [[CMP_BC_BA]], i32 [[MAX_BC]], i32 [[MAX_BA]]
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp ugt i32 [[MAX_BC]], %a
+; CHECK-NEXT:    [[MAX_ABC:%.*]] = select i1 [[TMP1]], i32 [[MAX_BC]], i32 %a
 ; CHECK-NEXT:    call void @extra_use(i32 [[MAX_BC]])
 ; CHECK-NEXT:    ret i32 [[MAX_ABC]]
 ;
   %cmp_bc = icmp ugt i32 %b, %c
   %max_bc = select i1 %cmp_bc, i32 %b, i32 %c
   %cmp_ba = icmp ugt i32 %b, %a
   %max_ba = select i1 %cmp_ba, i32 %b, i32 %a
   %cmp_bc_ba = icmp ugt i32 %max_bc, %max_ba
   %max_abc = select i1 %cmp_bc_ba, i32 %max_bc, i32 %max_ba
   call void @extra_use(i32 %max_bc)
   ret i32 %max_abc
 }
 
 define i32 @common_factor_umax_extra_use_rhs(i32 %a, i32 %b, i32 %c) {
 ; CHECK-LABEL: @common_factor_umax_extra_use_rhs(
-; CHECK-NEXT:    [[CMP_BC:%.*]] = icmp ugt i32 %b, %c
-; CHECK-NEXT:    [[MAX_BC:%.*]] = select i1 [[CMP_BC]], i32 %b, i32 %c
 ; CHECK-NEXT:    [[CMP_BA:%.*]] = icmp ugt i32 %b, %a
 ; CHECK-NEXT:    [[MAX_BA:%.*]] = select i1 [[CMP_BA]], i32 %b, i32 %a
-; CHECK-NEXT:    [[CMP_BC_BA:%.*]] = icmp ugt i32 [[MAX_BC]], [[MAX_BA]]
-; CHECK-NEXT:    [[MAX_ABC:%.*]] = select i1 [[CMP_BC_BA]], i32 [[MAX_BC]], i32 [[MAX_BA]]
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp ugt i32 [[MAX_BA]], %c
+; CHECK-NEXT:    [[MAX_ABC:%.*]] = select i1 [[TMP1]], i32 [[MAX_BA]], i32 %c
 ; CHECK-NEXT:    call void @extra_use(i32 [[MAX_BA]])
 ; CHECK-NEXT:    ret i32 [[MAX_ABC]]
 ;
   %cmp_bc = icmp ugt i32 %b, %c
   %max_bc = select i1 %cmp_bc, i32 %b, i32 %c
   %cmp_ba = icmp ugt i32 %b, %a
   %max_ba = select i1 %cmp_ba, i32 %b, i32 %a
   %cmp_bc_ba = icmp ugt i32 %max_bc, %max_ba