[ConstraintElimination] Add eq/ne facts to signed constraint system (#121423)

Facts of eq/ne were added to unsigned system only, causing some missing
optimizations. This patch adds eq/ne facts to both signed & unsigned
constraint system.

Fixes #117961.

Author: zsrkmyn

llvm/lib/Transforms/Scalar/ConstraintElimination.cpp

@@ -313,7 +313,8 @@ class ConstraintInfo {
   /// New variables that need to be added to the system are collected in
   /// \p NewVariables.
   ConstraintTy getConstraint(CmpInst::Predicate Pred, Value *Op0, Value *Op1,
-                             SmallVectorImpl<Value *> &NewVariables) const;
+                             SmallVectorImpl<Value *> &NewVariables,
+                             bool ForceSignedSystem = false) const;
 
   /// Turns a comparison of the form \p Op0 \p Pred \p Op1 into a vector of
   /// constraints using getConstraint. Returns an empty constraint if the result
@@ -330,6 +331,14 @@ class ConstraintInfo {
   void transferToOtherSystem(CmpInst::Predicate Pred, Value *A, Value *B,
                              unsigned NumIn, unsigned NumOut,
                              SmallVectorImpl<StackEntry> &DFSInStack);
+
+private:
+  /// Adds facts into constraint system. \p ForceSignedSystem can be set when
+  /// the \p Pred is eq/ne, and signed constraint system is used when it's
+  /// specified.
+  void addFactImpl(CmpInst::Predicate Pred, Value *A, Value *B, unsigned NumIn,
+                   unsigned NumOut, SmallVectorImpl<StackEntry> &DFSInStack,
+                   bool ForceSignedSystem);
 };
 
 /// Represents a (Coefficient * Variable) entry after IR decomposition.
@@ -636,8 +645,12 @@ static Decomposition decompose(Value *V,
 
 ConstraintTy
 ConstraintInfo::getConstraint(CmpInst::Predicate Pred, Value *Op0, Value *Op1,
-                              SmallVectorImpl<Value *> &NewVariables) const {
+                              SmallVectorImpl<Value *> &NewVariables,
+                              bool ForceSignedSystem) const {
   assert(NewVariables.empty() && "NewVariables must be empty when passed in");
+  assert((!ForceSignedSystem || CmpInst::isEquality(Pred)) &&
+         "signed system can only be forced on eq/ne");
+
   bool IsEq = false;
   bool IsNe = false;
 
@@ -652,15 +665,15 @@ ConstraintInfo::getConstraint(CmpInst::Predicate Pred, Value *Op0, Value *Op1,
     break;
   }
   case CmpInst::ICMP_EQ:
-    if (match(Op1, m_Zero())) {
+    if (!ForceSignedSystem && match(Op1, m_Zero())) {
       Pred = CmpInst::ICMP_ULE;
     } else {
       IsEq = true;
       Pred = CmpInst::ICMP_ULE;
     }
     break;
   case CmpInst::ICMP_NE:
-    if (match(Op1, m_Zero())) {
+    if (!ForceSignedSystem && match(Op1, m_Zero())) {
       Pred = CmpInst::getSwappedPredicate(CmpInst::ICMP_UGT);
       std::swap(Op0, Op1);
     } else {
@@ -677,7 +690,7 @@ ConstraintInfo::getConstraint(CmpInst::Predicate Pred, Value *Op0, Value *Op1,
     return {};
 
   SmallVector<ConditionTy, 4> Preconditions;
-  bool IsSigned = CmpInst::isSigned(Pred);
+  bool IsSigned = ForceSignedSystem || CmpInst::isSigned(Pred);
   auto &Value2Index = getValue2Index(IsSigned);
   auto ADec = decompose(Op0->stripPointerCastsSameRepresentation(),
                         Preconditions, IsSigned, DL);
@@ -737,7 +750,7 @@ ConstraintInfo::getConstraint(CmpInst::Predicate Pred, Value *Op0, Value *Op1,
   int64_t OffsetSum;
   if (AddOverflow(Offset1, Offset2, OffsetSum))
     return {};
-  if (Pred == (IsSigned ? CmpInst::ICMP_SLT : CmpInst::ICMP_ULT))
+  if (Pred == CmpInst::ICMP_SLT || Pred == CmpInst::ICMP_ULT)
     if (AddOverflow(OffsetSum, int64_t(-1), OffsetSum))
       return {};
   R[0] = OffsetSum;
@@ -1580,10 +1593,20 @@ static bool checkOrAndOpImpliedByOther(
 void ConstraintInfo::addFact(CmpInst::Predicate Pred, Value *A, Value *B,
                              unsigned NumIn, unsigned NumOut,
                              SmallVectorImpl<StackEntry> &DFSInStack) {
+  addFactImpl(Pred, A, B, NumIn, NumOut, DFSInStack, false);
+  // If the Pred is eq/ne, also add the fact to signed system.
+  if (CmpInst::isEquality(Pred))
+    addFactImpl(Pred, A, B, NumIn, NumOut, DFSInStack, true);
+}
+
+void ConstraintInfo::addFactImpl(CmpInst::Predicate Pred, Value *A, Value *B,
+                                 unsigned NumIn, unsigned NumOut,
+                                 SmallVectorImpl<StackEntry> &DFSInStack,
+                                 bool ForceSignedSystem) {
   // If the constraint has a pre-condition, skip the constraint if it does not
   // hold.
   SmallVector<Value *> NewVariables;
-  auto R = getConstraint(Pred, A, B, NewVariables);
+  auto R = getConstraint(Pred, A, B, NewVariables, ForceSignedSystem);
 
   // TODO: Support non-equality for facts as well.
   if (!R.isValid(*this) || R.isNe())

llvm/test/Transforms/ConstraintElimination/eq.ll

@@ -424,3 +424,53 @@ bc_equal:
 not_eq:
   ret i1 false
 }
+
+define i1 @test_eq_for_signed_cmp(i32 noundef %v0, i32 noundef %v1, i32 noundef %v2)  {
+; CHECK-LABEL: @test_eq_for_signed_cmp(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[V2:%.*]], [[V0:%.*]]
+; CHECK-NEXT:    [[CMP1:%.*]] = icmp sge i32 [[V0]], [[V1:%.*]]
+; CHECK-NEXT:    [[AND0:%.*]] = and i1 [[CMP1]], [[CMP]]
+; CHECK-NEXT:    [[CMP4:%.*]] = icmp sgt i32 [[V1]], [[V2]]
+; CHECK-NEXT:    [[AND1:%.*]] = and i1 false, [[AND0]]
+; CHECK-NEXT:    ret i1 [[AND1]]
+;
+entry:
+  %cmp = icmp eq i32 %v2, %v0
+  %cmp1 = icmp sge i32 %v0, %v1
+  %and0 = and i1 %cmp1, %cmp
+  %cmp4 = icmp sgt i32 %v1, %v2
+  %and1 = and i1 %cmp4, %and0
+  ret i1 %and1
+}
+
+define i1 @test_eq_for_signed_cmp_with_decompsition(i32 noundef %v0, i32 noundef %v1, i32 noundef %v2, i32 noundef %addend0, i32 noundef %addend1)  {
+; CHECK-LABEL: @test_eq_for_signed_cmp_with_decompsition(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[V0ADD:%.*]] = add nsw i32 [[V0:%.*]], [[ADDEND0:%.*]]
+; CHECK-NEXT:    [[V1ADD:%.*]] = add nsw i32 [[V1:%.*]], [[ADDEND1:%.*]]
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[V2:%.*]], [[V0ADD]]
+; CHECK-NEXT:    [[CMP1:%.*]] = icmp sge i32 [[V0ADD]], [[V1ADD]]
+; CHECK-NEXT:    [[CMP2:%.*]] = icmp sge i32 [[ADDEND0]], 0
+; CHECK-NEXT:    [[CMP3:%.*]] = icmp slt i32 [[ADDEND0]], [[ADDEND1]]
+; CHECK-NEXT:    [[AND0:%.*]] = and i1 [[CMP1]], [[CMP]]
+; CHECK-NEXT:    [[AND1:%.*]] = and i1 [[AND0]], [[CMP2]]
+; CHECK-NEXT:    [[AND2:%.*]] = and i1 [[AND1]], [[CMP3]]
+; CHECK-NEXT:    [[CMP4:%.*]] = icmp sgt i32 [[V1]], [[V2]]
+; CHECK-NEXT:    [[AND3:%.*]] = and i1 false, [[AND2]]
+; CHECK-NEXT:    ret i1 [[AND3]]
+;
+entry:
+  %v0add = add nsw i32 %v0, %addend0
+  %v1add = add nsw i32 %v1, %addend1
+  %cmp = icmp eq i32 %v2, %v0add
+  %cmp1 = icmp sge i32 %v0add, %v1add
+  %cmp2 = icmp sge i32 %addend0, 0
+  %cmp3 = icmp slt i32 %addend0, %addend1
+  %and0 = and i1 %cmp1, %cmp
+  %and1 = and i1 %and0, %cmp2
+  %and2 = and i1 %and1, %cmp3
+  %cmp4 = icmp sgt i32 %v1, %v2
+  %and3 = and i1 %cmp4, %and2
+  ret i1 %and3
+}

llvm/test/Transforms/ConstraintElimination/ne.ll

@@ -71,8 +71,7 @@ define i1 @test_ne_eq_0(i8 %a, i8 %b) {
 ; CHECK-NEXT:    [[RES_13:%.*]] = xor i1 [[RES_12]], false
 ; CHECK-NEXT:    [[RES_14:%.*]] = xor i1 [[RES_13]], false
 ; CHECK-NEXT:    [[RES_15:%.*]] = xor i1 [[RES_14]], false
-; CHECK-NEXT:    [[C_12:%.*]] = icmp sgt i8 [[A]], 0
-; CHECK-NEXT:    [[RES_16:%.*]] = xor i1 [[RES_15]], [[C_12]]
+; CHECK-NEXT:    [[RES_16:%.*]] = xor i1 [[RES_15]], false
 ; CHECK-NEXT:    ret i1 [[RES_16]]
 ;
 entry:
@@ -209,8 +208,7 @@ define i1 @test_ne_eq_1(i8 %a, i8 %b) {
 ; CHECK-NEXT:    [[RES_13:%.*]] = xor i1 [[RES_12]], true
 ; CHECK-NEXT:    [[RES_14:%.*]] = xor i1 [[RES_13]], true
 ; CHECK-NEXT:    [[RES_15:%.*]] = xor i1 [[RES_14]], false
-; CHECK-NEXT:    [[C_12:%.*]] = icmp sgt i8 [[A]], 0
-; CHECK-NEXT:    [[RES_16:%.*]] = xor i1 [[RES_15]], [[C_12]]
+; CHECK-NEXT:    [[RES_16:%.*]] = xor i1 [[RES_15]], true
 ; CHECK-NEXT:    ret i1 [[RES_16]]
 ;
 entry:

llvm/test/Transforms/ConstraintElimination/pr105785.ll

@@ -15,8 +15,7 @@ define void @pr105785(ptr %p) {
 ; CHECK-NEXT:    [[CMP2:%.*]] = icmp ult i32 [[FOR_IND2]], 3
 ; CHECK-NEXT:    br i1 [[CMP2]], label %[[FOR_BODY3]], label %[[FOR_COND]]
 ; CHECK:       [[FOR_BODY3]]:
-; CHECK-NEXT:    [[SCMP:%.*]] = call i32 @llvm.scmp.i32.i32(i32 [[FOR_IND]], i32 1)
-; CHECK-NEXT:    store i32 [[SCMP]], ptr [[P]], align 4
+; CHECK-NEXT:    store i32 -1, ptr [[P]], align 4
 ; CHECK-NEXT:    [[INC]] = add nuw nsw i32 [[FOR_IND2]], 1
 ; CHECK-NEXT:    br label %[[FOR_COND1]]
 ; CHECK:       [[FOR_END6]]: