Simplify using saturated add/sub

jayfoad · jayfoad · commit c854484260db · 2024-03-04T13:44:31.000Z
diff --git a/llvm/lib/Support/KnownBits.cpp b/llvm/lib/Support/KnownBits.cpp
@@ -57,181 +57,69 @@ KnownBits KnownBits::computeForAddCarry(
 KnownBits KnownBits::computeForAddSub(bool Add, bool NSW, bool NUW,
                                       const KnownBits &LHS,
                                       const KnownBits &RHS) {
-  // This can be a relatively expensive helper, so optimistically save some
-  // work.
-  if (LHS.isUnknown() && RHS.isUnknown())
-    return LHS;
-  KnownBits KnownOut;
-  if (Add) {
-    // Sum = LHS + RHS + 0
-    KnownOut =
-        ::computeForAddCarry(LHS, RHS, /*CarryZero*/ true, /*CarryOne*/ false);
-  } else {
-    // Sum = LHS + ~RHS + 1
-    KnownBits NotRHS = RHS;
-    std::swap(NotRHS.Zero, NotRHS.One);
-    KnownOut = ::computeForAddCarry(LHS, NotRHS, /*CarryZero*/ false,
-                                    /*CarryOne*/ true);
-  }
-  if (!NSW && !NUW)
-    return KnownOut;
-
-  auto GetMinMaxVal = [Add](bool ForNSW, bool ForMax, const KnownBits &L,
-                            const KnownBits &R, bool &OV) {
-    APInt LVal = ForMax ? L.getMaxValue() : L.getMinValue();
-    APInt RVal = Add == ForMax ? R.getMaxValue() : R.getMinValue();
-
-    if (ForNSW) {
-      LVal.clearSignBit();
-      RVal.clearSignBit();
-    }
-    APInt Res = Add ? LVal.uadd_ov(RVal, OV) : LVal.usub_ov(RVal, OV);
-    if (ForNSW) {
-      OV = Res.isSignBitSet();
-      Res.clearSignBit();
-      if (Res.getBitWidth() > 1 && Res[Res.getBitWidth() - 2])
-        Res.setSignBit();
+  unsigned BitWidth = LHS.getBitWidth();
+  KnownBits KnownOut(LHS.getBitWidth());
+  if (!LHS.isUnknown() && !RHS.isUnknown()) {
+    if (Add) {
+      // Sum = LHS + RHS + 0
+      KnownOut = ::computeForAddCarry(LHS, RHS, /*CarryZero=*/true,
+                                      /*CarryOne=*/false);
+    } else {
+      // Sum = LHS + ~RHS + 1
+      KnownBits NotRHS = RHS;
+      std::swap(NotRHS.Zero, NotRHS.One);
+      KnownOut = ::computeForAddCarry(LHS, NotRHS, /*CarryZero=*/false,
+                                      /*CarryOne=*/true);
     }
-    return Res;
-  };
-
-  auto GetMaxVal = [&GetMinMaxVal](bool ForNSW, const KnownBits &L,
-                                   const KnownBits &R, bool &OV) {
-    return GetMinMaxVal(ForNSW, /*ForMax=*/true, L, R, OV);
-  };
-
-  auto GetMinVal = [&GetMinMaxVal](bool ForNSW, const KnownBits &L,
-                                   const KnownBits &R, bool &OV) {
-    return GetMinMaxVal(ForNSW, /*ForMax=*/false, L, R, OV);
-  };
-
-  auto ForceNegative = [](KnownBits &Known) {
-    Known.Zero.clearSignBit();
-    Known.One.setSignBit();
-  };
-
-  auto ForcePositive = [](KnownBits &Known) {
-    Known.One.clearSignBit();
-    Known.Zero.setSignBit();
-  };
+  }
 
-  // Handle add/sub given nsw and/or nuw.
-  //
-  // Possible TODO: Add/Sub implementations mirror one another in many ways.
-  // They could probably be compressed into a single implementation of roughly
-  // half the total LOC. Leaving seperate for now to increase clarity.
-  // NB: We handle NSW by essentially treating as nuw of bitwidth - 1 then
-  // deducing bits based on the known sign result.
-  if (Add) {
-    if (NUW || (LHS.isNonNegative() && RHS.isNonNegative())) {
-      bool OverflowMin;
-      APInt MinVal;
+  if (NUW) {
+    if (Add) {
+      // (add nuw X, Y)
+      APInt MinVal = LHS.getMinValue().uadd_sat(RHS.getMinValue());
       if (NSW) {
-        MinVal = GetMinVal(/*ForNSW=*/true, LHS, RHS, OverflowMin);
-        // (add nsw nuw) or (add nsw PosX, PosY)
-
-        // None of the adds can end up overflowing, so min consecutive
-        // highbits in minimum possible of X + Y must all remain set.
-        KnownOut.One.setHighBits(MinVal.countLeadingOnes());
-
-        // NSW and Positive arguments leads to positive result.
-        if (LHS.isNonNegative() && RHS.isNonNegative())
-          ForcePositive(KnownOut);
+        unsigned NumBits = MinVal.trunc(BitWidth - 1).countl_one();
+        KnownOut.One.setBits(BitWidth - 1 - NumBits, BitWidth - 1);
       }
-      if (NUW) {
-        KnownOut.One.clearSignBit();
-        // (add nuw X, Y)
-        MinVal = GetMinVal(/*ForNSW=*/false, LHS, RHS, OverflowMin);
-        // Same as (add nsw PosX, PosY), basically since we can't overflow,
-        // the high bits of minimum possible X + Y must remain set.
-        KnownOut.One.setHighBits(MinVal.countLeadingOnes());
+      KnownOut.One.setHighBits(MinVal.countLeadingOnes());
+    } else {
+      // (sub nuw X, Y)
+      APInt MaxVal = LHS.getMaxValue().usub_sat(RHS.getMinValue());
+      if (NSW) {
+        unsigned NumBits = MaxVal.trunc(BitWidth - 1).countl_zero();
+        KnownOut.Zero.setBits(BitWidth - 1 - NumBits, BitWidth - 1);
       }
-    } else if (LHS.isNegative() && RHS.isNegative()) {
-      bool OverflowMax;
-      APInt MaxVal = GetMaxVal(/*ForNSW=*/true, LHS, RHS, OverflowMax);
-      // (add nsw NegX, NegY)
-
-      // We need to re-overflow the signbit, so we are looking for sequence
-      // of 0s from consecutive overflows.
       KnownOut.Zero.setHighBits(MaxVal.countLeadingZeros());
-      ForceNegative(KnownOut);
-    } else if (!KnownOut.isSignUnknown()) {
-      // Pass, avoid extra work if we already know the sign bit.
-    } else if (LHS.isNonNegative() || RHS.isNonNegative()) {
-      bool OverflowMin;
-      (void)GetMinVal(/*ForNSW=*/true, LHS, RHS, OverflowMin);
-      // (add nsw PosX, ?Y)
-
-      // If the minimal possible of X + Y overflows the signbit, then Y must
-      // have been signed (which will cause unsigned overflow otherwise nsw
-      // will be violated) leading to unsigned result.
-      if (OverflowMin)
-        KnownOut.makeNonNegative();
-    } else if (LHS.isNegative() || RHS.isNegative()) {
-      bool OverflowMax;
-      (void)GetMaxVal(/*ForNSW=*/true, LHS, RHS, OverflowMax);
-      // (add nsw NegX, ?Y)
-
-      // If the maximum possible of X + Y doesn't overflows the signbit,
-      // then Y must have been unsigned (otherwise nsw violated) so NegX +
-      // PosY w.o overflowing the signbit results in Negative.
-      if (!OverflowMax)
-        KnownOut.makeNegative();
     }
-  } else {
-    if (NUW || (LHS.isNegative() && RHS.isNonNegative())) {
-      bool OverflowMax;
-      APInt MaxVal;
-      if (NSW) {
-        MaxVal = GetMaxVal(/*ForNSW=*/true, LHS, RHS, OverflowMax);
-        // (sub nsw nuw) or (sub nsw NegX, PosY)
-
-        // None of the subs can overflow at any point, so any common high bits
-        // will subtract away and result in zeros.
-        KnownOut.Zero.setHighBits(MaxVal.countLeadingZeros());
-        if (LHS.isNegative() && RHS.isNonNegative())
-          ForceNegative(KnownOut);
-      }
-      if (NUW) {
-        KnownOut.Zero.clearSignBit();
-        // (sub nuw X, Y)
-        MaxVal = GetMaxVal(/*ForNSW=*/false, LHS, RHS, OverflowMax);
-
-        // Basically all common high bits between X/Y will cancel out as
-        // leading zeros.
-        KnownOut.Zero.setHighBits(MaxVal.countLeadingZeros());
-      }
-    } else if (LHS.isNonNegative() && RHS.isNegative()) {
-      bool OverflowMin;
-      APInt MinVal = GetMinVal(/*ForNSW=*/true, LHS, RHS, OverflowMin);
-      // (sub nsw PosX, NegY)
+  }
 
-      // Opposite case of above, we must "re-overflow" the signbit, so
-      // minimal set of high bits will be fixed.
-      KnownOut.One.setHighBits(MinVal.countLeadingOnes());
-      ForcePositive(KnownOut);
-    } else if (!KnownOut.isSignUnknown()) {
-      // Pass, avoid extra work if we already know the sign bit.
-    } else if (LHS.isNegative() || RHS.isNonNegative()) {
-      bool OverflowMax;
-      (void)GetMaxVal(/*ForNSW=*/true, LHS, RHS, OverflowMax);
-      // (sub nsw NegX/?X, ?Y/PosY)
-      if (OverflowMax)
-        KnownOut.makeNegative();
-    } else if (LHS.isNonNegative() || RHS.isNegative()) {
-      bool OverflowMin;
-      (void)GetMinVal(/*ForNSW=*/true, LHS, RHS, OverflowMin);
-      // (sub nsw PosX/?X, ?Y/NegY)
-      if (!OverflowMin)
-        KnownOut.makeNonNegative();
+  if (NSW) {
+    APInt MinVal;
+    APInt MaxVal;
+    if (Add) {
+      // (add nsw X, Y)
+      MinVal = LHS.getSignedMinValue().sadd_sat(RHS.getSignedMinValue());
+      MaxVal = LHS.getSignedMaxValue().sadd_sat(RHS.getSignedMaxValue());
+    } else {
+      // (sub nsw X, Y)
+      MinVal = LHS.getSignedMinValue().ssub_sat(RHS.getSignedMaxValue());
+      MaxVal = LHS.getSignedMaxValue().ssub_sat(RHS.getSignedMinValue());
+    }
+    if (MinVal.isNonNegative()) {
+      unsigned NumBits = MinVal.trunc(BitWidth - 1).countl_one();
+      KnownOut.One.setBits(BitWidth - 1 - NumBits, BitWidth - 1);
+      KnownOut.Zero.setSignBit();
+    }
+    if (MaxVal.isNegative()) {
+      unsigned NumBits = MaxVal.trunc(BitWidth - 1).countl_zero();
+      KnownOut.Zero.setBits(BitWidth - 1 - NumBits, BitWidth - 1);
+      KnownOut.One.setSignBit();
     }
   }
 
   // Just return 0 if the nsw/nuw is violated and we have poison.
-  if (KnownOut.hasConflict()) {
+  if (KnownOut.hasConflict())
     KnownOut.setAllZero();
-    return KnownOut;
-  }
 
   return KnownOut;
 }
