[InstCombine] Optimize overflow check base on uadd.with.overflow result (PR40846)

[nikic]

Summary

Fix for https://bugs.llvm.org/show_bug.cgi?id=40846.

This adds a combine for cases where a (a + b) < a style overflow check is performed, but with a + b being the result of uadd.with.overflow, so the overflow result is also already available and we can just use it. Subsequently GVN/CSE will deduplicate the extracts.

We can run into this situation if you have both a uadd.with.overflow and a manual add + overflow check in the same function (on the same operands), in which case GVN will rewrite the add to the with.overflow result and leave you with this pattern.

The implementation is a bit ugly because I'm handling the various canonicalization edge cases. I'm wondering if we couldn't introduce something like m_ICmpULT() to match both ULT and the UGT commutation as well as the canonicalization edge-cases for constant operands.

[spatel]

Probably should not be included in the same patch, but if we're seeing this pattern, then is it likely that we'd also see code that checks the inverse condition (that the overflow did not occur)?

define i1 @uadd_res_uge_x(i32 %x, i32 %y, i1* %p) {
  %a = tail call { i32, i1 } @llvm.uadd.with.overflow.i32(i32 %x, i32 %y)
  %b = extractvalue { i32, i1 } %a, 1
  store i1 %b, i1* %p, align 1
  %c = extractvalue { i32, i1 } %a, 0
  %d = icmp uge i32 %c, %x
  ret i1 %d
}

[lebedev.ri]

@nikic reverse-ping?

[nikic]

Rebase and use new matchers.

[spatel]

LGTM

[lebedev.ri]

This LG, thanks.

[lebedev.ri]

Forgot to add, what about other predicates +xor (extract %agg, 1), -1 ?

[nikic]

@lebedev.ri Right, the negated version should also be handled. I've committed this, but left https://bugs.llvm.org/show_bug.cgi?id=40846 open to track that.

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

if (match(LHS, m_ShuffleVector(m_Value(V1), m_Undef(), m_Constant(M))) &&
     CmpInst::Predicate P = Cmp.getPredicate();
     Value *NewCmp = isa<ICmpInst>(Cmp) ? Builder.CreateICmp(P, V1, V2)
                                        : Builder.CreateFCmp(P, V1, V2);
     return new ShuffleVectorInst(NewCmp, UndefValue::get(NewCmp->getType()), M);
   }
   return nullptr;
 }
 
+// extract(uadd.with.overflow(A, B), 0) ult A
+//  -> extract(uadd.with.overflow(A, B), 1)
+static Instruction *foldICmpOfUAddOv(ICmpInst &I) {
+  CmpInst::Predicate Pred = I.getPredicate();
+  Value *Op0 = I.getOperand(0), *Op1 = I.getOperand(1);
+
+  Value *UAddOv;
+  Value *A, *B;
+  auto UAddOvResultPat = m_ExtractValue<0>(
+      m_Intrinsic<Intrinsic::uadd_with_overflow>(m_Value(A), m_Value(B)));
+  if (match(Op0, UAddOvResultPat) &&
+      ((Pred == ICmpInst::ICMP_ULT && (Op1 == A || Op1 == B)) ||
+       (Pred == ICmpInst::ICMP_EQ && match(Op1, m_ZeroInt()) &&
+        (match(A, m_One()) || match(B, m_One()))) ||
+       (Pred == ICmpInst::ICMP_NE && match(Op1, m_AllOnes()) &&
+        (match(A, m_AllOnes()) || match(B, m_AllOnes())))))
+    // extract(uadd.with.overflow(A, B), 0) < A
+    // extract(uadd.with.overflow(A, 1), 0) == 0
+    // extract(uadd.with.overflow(A, -1), 0) != -1
+    UAddOv = cast<ExtractValueInst>(Op0)->getAggregateOperand();
+  else if (match(Op1, UAddOvResultPat) &&
+           Pred == ICmpInst::ICMP_UGT && (Op0 == A || Op0 == B))
+    // A > extract(uadd.with.overflow(A, B), 0)
+    UAddOv = cast<ExtractValueInst>(Op1)->getAggregateOperand();
+  else
+    return nullptr;
+
+  return ExtractValueInst::Create(UAddOv, 1);
+}
+
 Instruction *InstCombiner::visitICmpInst(ICmpInst &I) {
   bool Changed = false;
   const SimplifyQuery Q = SQ.getWithInstruction(&I);
   Value *Op0 = I.getOperand(0), *Op1 = I.getOperand(1);
   unsigned Op0Cplxity = getComplexity(Op0);
   unsigned Op1Cplxity = getComplexity(Op1);
 
   /// Orders the operands of the compare so that they are listed from most
if (match(Op1, m_Mul(m_ZExt(m_Value(A)), m_ZExt(m_Value(B))))) {
       if (Instruction *R = processUMulZExtIdiom(I, Op1, Op0, *this))
         return R;
     }
   }
 
   if (Instruction *Res = foldICmpEquality(I))
     return Res;
 
+  if (Instruction *Res = foldICmpOfUAddOv(I))
+    return Res;
+
   // The 'cmpxchg' instruction returns an aggregate containing the old value and
   // an i1 which indicates whether or not we successfully did the swap.
   //
   // Replace comparisons between the old value and the expected value with the
   // indicator that 'cmpxchg' returns.
   //
   // N.B.  This transform is only valid when the 'cmpxchg' is not permitted to
   // spuriously fail.  In those cases, the old value may equal the expected

llvm/test/Transforms/InstCombine/with_overflow.ll

;
   ret { i32, i1 } %x
 }
 
 define i1 @uadd_res_ult_x(i32 %x, i32 %y, i1* %p) nounwind {
 ; CHECK-LABEL: @uadd_res_ult_x(
 ; CHECK-NEXT:    [[A:%.*]] = call { i32, i1 } @llvm.uadd.with.overflow.i32(i32 [[X:%.*]], i32 [[Y:%.*]])
 ; CHECK-NEXT:    [[B:%.*]] = extractvalue { i32, i1 } [[A]], 1
 ; CHECK-NEXT:    store i1 [[B]], i1* [[P:%.*]], align 1
-; CHECK-NEXT:    [[C:%.*]] = extractvalue { i32, i1 } [[A]], 0
-; CHECK-NEXT:    [[D:%.*]] = icmp ult i32 [[C]], [[X]]
+; CHECK-NEXT:    [[D:%.*]] = extractvalue { i32, i1 } [[A]], 1
 ; CHECK-NEXT:    ret i1 [[D]]
 ;
   %a = call { i32, i1 } @llvm.uadd.with.overflow.i32(i32 %x, i32 %y)
   %b = extractvalue { i32, i1 } %a, 1
   store i1 %b, i1* %p
   %c = extractvalue { i32, i1 } %a, 0
   %d = icmp ult i32 %c, %x
   ret i1 %d
 }
 
 define i1 @uadd_res_ult_y(i32 %x, i32 %y, i1* %p) nounwind {
 ; CHECK-LABEL: @uadd_res_ult_y(
 ; CHECK-NEXT:    [[A:%.*]] = call { i32, i1 } @llvm.uadd.with.overflow.i32(i32 [[X:%.*]], i32 [[Y:%.*]])
 ; CHECK-NEXT:    [[B:%.*]] = extractvalue { i32, i1 } [[A]], 1
 ; CHECK-NEXT:    store i1 [[B]], i1* [[P:%.*]], align 1
-; CHECK-NEXT:    [[C:%.*]] = extractvalue { i32, i1 } [[A]], 0
-; CHECK-NEXT:    [[D:%.*]] = icmp ult i32 [[C]], [[Y]]
+; CHECK-NEXT:    [[D:%.*]] = extractvalue { i32, i1 } [[A]], 1
 ; CHECK-NEXT:    ret i1 [[D]]
 ;
   %a = call { i32, i1 } @llvm.uadd.with.overflow.i32(i32 %x, i32 %y)
   %b = extractvalue { i32, i1 } %a, 1
   store i1 %b, i1* %p
   %c = extractvalue { i32, i1 } %a, 0
   %d = icmp ult i32 %c, %y
   ret i1 %d
 }
 
 define i1 @uadd_res_ugt_x(i32 %xx, i32 %y, i1* %p) nounwind {
 ; CHECK-LABEL: @uadd_res_ugt_x(
 ; CHECK-NEXT:    [[X:%.*]] = urem i32 42, [[XX:%.*]]
 ; CHECK-NEXT:    [[A:%.*]] = call { i32, i1 } @llvm.uadd.with.overflow.i32(i32 [[X]], i32 [[Y:%.*]])
 ; CHECK-NEXT:    [[B:%.*]] = extractvalue { i32, i1 } [[A]], 1
 ; CHECK-NEXT:    store i1 [[B]], i1* [[P:%.*]], align 1
-; CHECK-NEXT:    [[C:%.*]] = extractvalue { i32, i1 } [[A]], 0
-; CHECK-NEXT:    [[D:%.*]] = icmp ugt i32 [[X]], [[C]]
+; CHECK-NEXT:    [[D:%.*]] = extractvalue { i32, i1 } [[A]], 1
 ; CHECK-NEXT:    ret i1 [[D]]
 ;
   %x = urem i32 42, %xx ; Thwart complexity-based canonicalization
   %a = call { i32, i1 } @llvm.uadd.with.overflow.i32(i32 %x, i32 %y)
   %b = extractvalue { i32, i1 } %a, 1
   store i1 %b, i1* %p
   %c = extractvalue { i32, i1 } %a, 0
   %d = icmp ugt i32 %x, %c
   ret i1 %d
 }
 
 define i1 @uadd_res_ugt_y(i32 %x, i32 %yy, i1* %p) nounwind {
 ; CHECK-LABEL: @uadd_res_ugt_y(
 ; CHECK-NEXT:    [[Y:%.*]] = urem i32 42, [[YY:%.*]]
 ; CHECK-NEXT:    [[A:%.*]] = call { i32, i1 } @llvm.uadd.with.overflow.i32(i32 [[X:%.*]], i32 [[Y]])
 ; CHECK-NEXT:    [[B:%.*]] = extractvalue { i32, i1 } [[A]], 1
 ; CHECK-NEXT:    store i1 [[B]], i1* [[P:%.*]], align 1
-; CHECK-NEXT:    [[C:%.*]] = extractvalue { i32, i1 } [[A]], 0
-; CHECK-NEXT:    [[D:%.*]] = icmp ugt i32 [[Y]], [[C]]
+; CHECK-NEXT:    [[D:%.*]] = extractvalue { i32, i1 } [[A]], 1
 ; CHECK-NEXT:    ret i1 [[D]]
 ;
   %y = urem i32 42, %yy ; Thwart complexity-based canonicalization
   %a = call { i32, i1 } @llvm.uadd.with.overflow.i32(i32 %x, i32 %y)
   %b = extractvalue { i32, i1 } %a, 1
   store i1 %b, i1* %p
   %c = extractvalue { i32, i1 } %a, 0
   %d = icmp ugt i32 %y, %c
   ret i1 %d
 }
 
 define i1 @uadd_res_ult_const(i32 %x, i1* %p) nounwind {
 ; CHECK-LABEL: @uadd_res_ult_const(
 ; CHECK-NEXT:    [[A:%.*]] = call { i32, i1 } @llvm.uadd.with.overflow.i32(i32 [[X:%.*]], i32 42)
 ; CHECK-NEXT:    [[B:%.*]] = extractvalue { i32, i1 } [[A]], 1
 ; CHECK-NEXT:    store i1 [[B]], i1* [[P:%.*]], align 1
-; CHECK-NEXT:    [[C:%.*]] = extractvalue { i32, i1 } [[A]], 0
-; CHECK-NEXT:    [[D:%.*]] = icmp ult i32 [[C]], 42
+; CHECK-NEXT:    [[D:%.*]] = extractvalue { i32, i1 } [[A]], 1
 ; CHECK-NEXT:    ret i1 [[D]]
 ;
   %a = call { i32, i1 } @llvm.uadd.with.overflow.i32(i32 %x, i32 42)
   %b = extractvalue { i32, i1 } %a, 1
   store i1 %b, i1* %p
   %c = extractvalue { i32, i1 } %a, 0
   %d = icmp ult i32 %c, 42
   ret i1 %d
 }
 
 define i1 @uadd_res_ult_const_one(i32 %x, i1* %p) nounwind {
 ; CHECK-LABEL: @uadd_res_ult_const_one(
 ; CHECK-NEXT:    [[A:%.*]] = call { i32, i1 } @llvm.uadd.with.overflow.i32(i32 [[X:%.*]], i32 1)
 ; CHECK-NEXT:    [[B:%.*]] = extractvalue { i32, i1 } [[A]], 1
 ; CHECK-NEXT:    store i1 [[B]], i1* [[P:%.*]], align 1
-; CHECK-NEXT:    [[C:%.*]] = extractvalue { i32, i1 } [[A]], 0
-; CHECK-NEXT:    [[D:%.*]] = icmp eq i32 [[C]], 0
+; CHECK-NEXT:    [[D:%.*]] = extractvalue { i32, i1 } [[A]], 1
 ; CHECK-NEXT:    ret i1 [[D]]
 ;
   %a = call { i32, i1 } @llvm.uadd.with.overflow.i32(i32 %x, i32 1)
   %b = extractvalue { i32, i1 } %a, 1
   store i1 %b, i1* %p
   %c = extractvalue { i32, i1 } %a, 0
   %d = icmp ult i32 %c, 1
   ret i1 %d
 }
 
 define i1 @uadd_res_ult_const_minus_one(i32 %x, i1* %p) nounwind {
 ; CHECK-LABEL: @uadd_res_ult_const_minus_one(
 ; CHECK-NEXT:    [[A:%.*]] = call { i32, i1 } @llvm.uadd.with.overflow.i32(i32 [[X:%.*]], i32 -1)
 ; CHECK-NEXT:    [[B:%.*]] = extractvalue { i32, i1 } [[A]], 1
 ; CHECK-NEXT:    store i1 [[B]], i1* [[P:%.*]], align 1
-; CHECK-NEXT:    [[C:%.*]] = extractvalue { i32, i1 } [[A]], 0
-; CHECK-NEXT:    [[D:%.*]] = icmp ne i32 [[C]], -1
+; CHECK-NEXT:    [[D:%.*]] = extractvalue { i32, i1 } [[A]], 1
 ; CHECK-NEXT:    ret i1 [[D]]
 ;
   %a = call { i32, i1 } @llvm.uadd.with.overflow.i32(i32 %x, i32 -1)
   %b = extractvalue { i32, i1 } %a, 1
   store i1 %b, i1* %p
   %c = extractvalue { i32, i1 } %a, 0
   %d = icmp ult i32 %c, -1
   ret i1 %d