[Support] Add KnownBits::computeForSubBorrow

llvm/include/llvm/Support/KnownBits.h

@@ -332,6 +332,11 @@ struct KnownBits {
   static KnownBits computeForAddSub(bool Add, bool NSW, const KnownBits &LHS,
                                     KnownBits RHS);
 
+  /// Compute known bits results from subtracting RHS from LHS with 1-bit
+  /// Borrow.
+  static KnownBits computeForSubBorrow(const KnownBits &LHS, KnownBits RHS,
+                                       const KnownBits &Borrow);
+
   /// Compute knownbits resulting from llvm.sadd.sat(LHS, RHS)
   static KnownBits sadd_sat(const KnownBits &LHS, const KnownBits &RHS);
 

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

@@ -3741,14 +3741,19 @@ KnownBits SelectionDAG::computeKnownBits(SDValue Op, const APInt &DemandedElts,
     assert(Op.getResNo() == 0 &&
            "We only compute knownbits for the difference here.");
 
-    // TODO: Compute influence of the carry operand.
-    if (Opcode == ISD::USUBO_CARRY || Opcode == ISD::SSUBO_CARRY)
-      break;
+    // With USUBO_CARRY and SSUBO_CARRY a borrow bit may be added in.
+    KnownBits Borrow(1);
+    if (Opcode == ISD::USUBO_CARRY || Opcode == ISD::SSUBO_CARRY) {
+      Borrow = computeKnownBits(Op.getOperand(2), DemandedElts, Depth + 1);
+      // Borrow has bit width 1
+      Borrow = Borrow.trunc(1);
+    } else {
+      Borrow.setAllZero();
+    }
 
     Known = computeKnownBits(Op.getOperand(0), DemandedElts, Depth + 1);
     Known2 = computeKnownBits(Op.getOperand(1), DemandedElts, Depth + 1);
-    Known = KnownBits::computeForAddSub(/* Add */ false, /* NSW */ false,
-                                        Known, Known2);
+    Known = KnownBits::computeForSubBorrow(Known, Known2, Borrow);
     break;
   }
   case ISD::UADDO:
@@ -3773,15 +3778,13 @@ KnownBits SelectionDAG::computeKnownBits(SDValue Op, const APInt &DemandedElts,
     if (Opcode == ISD::ADDE)
       // Can't track carry from glue, set carry to unknown.
       Carry.resetAll();
-    else if (Opcode == ISD::UADDO_CARRY || Opcode == ISD::SADDO_CARRY)
-      // TODO: Compute known bits for the carry operand. Not sure if it is worth
-      // the trouble (how often will we find a known carry bit). And I haven't
-      // tested this very much yet, but something like this might work:
-      //   Carry = computeKnownBits(Op.getOperand(2), DemandedElts, Depth + 1);
-      //   Carry = Carry.zextOrTrunc(1, false);
-      Carry.resetAll();
-    else
+    else if (Opcode == ISD::UADDO_CARRY || Opcode == ISD::SADDO_CARRY) {
+      Carry = computeKnownBits(Op.getOperand(2), DemandedElts, Depth + 1);
+      // Carry has bit width 1
+      Carry = Carry.trunc(1);
+    } else {
       Carry.setAllZero();
+    }
 
     Known = computeKnownBits(Op.getOperand(0), DemandedElts, Depth + 1);
     Known2 = computeKnownBits(Op.getOperand(1), DemandedElts, Depth + 1);

llvm/lib/Support/KnownBits.cpp

@@ -85,6 +85,18 @@ KnownBits KnownBits::computeForAddSub(bool Add, bool NSW,
   return KnownOut;
 }
 
+KnownBits KnownBits::computeForSubBorrow(const KnownBits &LHS, KnownBits RHS,
+                                         const KnownBits &Borrow) {
+  assert(Borrow.getBitWidth() == 1 && "Borrow must be 1-bit");
+
+  // LHS - RHS = LHS + ~RHS + 1
+  // Carry 1 - Borrow in ::computeForAddCarry
+  std::swap(RHS.Zero, RHS.One);
+  return ::computeForAddCarry(LHS, RHS,
+                              /*CarryZero=*/Borrow.One.getBoolValue(),
+                              /*CarryOne=*/Borrow.Zero.getBoolValue());
+}
+
 KnownBits KnownBits::sextInReg(unsigned SrcBitWidth) const {
   unsigned BitWidth = getBitWidth();
   assert(0 < SrcBitWidth && SrcBitWidth <= BitWidth &&

llvm/unittests/CodeGen/AArch64SelectionDAGTest.cpp

@@ -254,6 +254,59 @@ TEST_F(AArch64SelectionDAGTest, ComputeKnownBits_ADD) {
   EXPECT_EQ(Known.One, APInt(8, 0x55));
 }
 
+// Piggy-backing on the AArch64 tests to verify SelectionDAG::computeKnownBits.
+TEST_F(AArch64SelectionDAGTest, ComputeKnownBits_UADDO_CARRY) {
+  SDLoc Loc;
+  auto IntVT = EVT::getIntegerVT(Context, 8);
+  auto UnknownOp = DAG->getRegister(0, IntVT);
+  auto Mask_Zero = DAG->getConstant(0x28, Loc, IntVT);
+  auto Mask_One = DAG->getConstant(0x20, Loc, IntVT);
+  auto N0 = DAG->getNode(ISD::AND, Loc, IntVT, Mask_Zero, UnknownOp);
+  N0 = DAG->getNode(ISD::OR, Loc, IntVT, Mask_One, N0);
+  auto N1 = DAG->getConstant(0x65, Loc, IntVT);
+
+  KnownBits Known;
+
+  auto UnknownBorrow = DAG->getRegister(1, IntVT);
+  auto OpUnknownBorrow =
+      DAG->getNode(ISD::UADDO_CARRY, Loc, IntVT, N0, N1, UnknownBorrow);
+  // N0 = 0010?000
+  // N1 = 01100101
+  // B  =        ?
+  //  =>
+  // Known.Zero = 01110000 (0x70)
+  // Known.One  = 10000100 (0x84)
+  Known = DAG->computeKnownBits(OpUnknownBorrow);
+  EXPECT_EQ(Known.Zero, APInt(8, 0x70));
+  EXPECT_EQ(Known.One, APInt(8, 0x84));
+
+  auto ZeroBorrow = DAG->getConstant(0x0, Loc, IntVT);
+  auto OpZeroBorrow =
+      DAG->getNode(ISD::UADDO_CARRY, Loc, IntVT, N0, N1, ZeroBorrow);
+  // N0 = 0010?000
+  // N1 = 01100101
+  // B  =        0
+  //  =>
+  // Known.Zero = 01110010 (0x72)
+  // Known.One  = 10000101 (0x85)
+  Known = DAG->computeKnownBits(OpZeroBorrow);
+  EXPECT_EQ(Known.Zero, APInt(8, 0x72));
+  EXPECT_EQ(Known.One, APInt(8, 0x85));
+
+  auto OneBorrow = DAG->getConstant(0x1, Loc, IntVT);
+  auto OpOneBorrow =
+      DAG->getNode(ISD::UADDO_CARRY, Loc, IntVT, N0, N1, OneBorrow);
+  // N0 = 0010?000
+  // N1 = 01100101
+  // B  =        1
+  //  =>
+  // Known.Zero = 01110001 (0x71)
+  // Known.One  = 10000110 (0x86)
+  Known = DAG->computeKnownBits(OpOneBorrow);
+  EXPECT_EQ(Known.Zero, APInt(8, 0x71));
+  EXPECT_EQ(Known.One, APInt(8, 0x86));
+}
+
 // Piggy-backing on the AArch64 tests to verify SelectionDAG::computeKnownBits.
 TEST_F(AArch64SelectionDAGTest, ComputeKnownBits_SUB) {
   SDLoc Loc;
@@ -273,6 +326,60 @@ TEST_F(AArch64SelectionDAGTest, ComputeKnownBits_SUB) {
   EXPECT_EQ(Known.One, APInt(8, 0x1));
 }
 
+// Piggy-backing on the AArch64 tests to verify SelectionDAG::computeKnownBits.
+TEST_F(AArch64SelectionDAGTest, ComputeKnownBits_USUBO_CARRY) {
+  SDLoc Loc;
+  auto IntVT = EVT::getIntegerVT(Context, 8);
+  auto N0 = DAG->getConstant(0x5a, Loc, IntVT);
+  auto UnknownOp = DAG->getRegister(0, IntVT);         // ????????
+  auto Mask1_Zero = DAG->getConstant(0x8, Loc, IntVT); // 00001000
+  auto Mask1_One = DAG->getConstant(0x20, Loc, IntVT); // 00100000
+  // N1 = (???????? & 00001000) | 00100000 = 0010?000
+  auto N1 = DAG->getNode(ISD::AND, Loc, IntVT, Mask1_Zero, UnknownOp);
+  N1 = DAG->getNode(ISD::OR, Loc, IntVT, Mask1_One, N1);
+
+  KnownBits Known;
+
+  auto UnknownBorrow = DAG->getRegister(1, IntVT);
+  auto OpUnknownBorrow =
+      DAG->getNode(ISD::USUBO_CARRY, Loc, IntVT, N0, N1, UnknownBorrow);
+  // N0 = 01011010
+  // N1 = 0010?000
+  // B  =        ?
+  //  =>
+  // Known.Zero = 11000100 (0xc4)
+  // Known.One  = 00110000 (0x30)
+  Known = DAG->computeKnownBits(OpUnknownBorrow);
+  EXPECT_EQ(Known.Zero, APInt(8, 0xc4));
+  EXPECT_EQ(Known.One, APInt(8, 0x30));
+
+  auto ZeroBorrow = DAG->getConstant(0x0, Loc, IntVT);
+  auto OpZeroBorrow =
+      DAG->getNode(ISD::USUBO_CARRY, Loc, IntVT, N0, N1, ZeroBorrow);
+  // N0 = 01011010
+  // N1 = 0010?000
+  // B  =        0
+  //  =>
+  // Known.Zero = 11000101 (0xc5)
+  // Known.One  = 00110010 (0x32)
+  Known = DAG->computeKnownBits(OpZeroBorrow);
+  EXPECT_EQ(Known.Zero, APInt(8, 0xc5));
+  EXPECT_EQ(Known.One, APInt(8, 0x32));
+
+  auto OneBorrow = DAG->getConstant(0x1, Loc, IntVT);
+  auto OpOneBorrow =
+      DAG->getNode(ISD::USUBO_CARRY, Loc, IntVT, N0, N1, OneBorrow);
+  // N0 = 01011010
+  // N1 = 0010?000
+  // B  =        1
+  //  =>
+  // Known.Zero = 11000110 (0xc6)
+  // Known.One  = 00110001 (0x31)
+  Known = DAG->computeKnownBits(OpOneBorrow);
+  EXPECT_EQ(Known.Zero, APInt(8, 0xc6));
+  EXPECT_EQ(Known.One, APInt(8, 0x31));
+}
+
 TEST_F(AArch64SelectionDAGTest, isSplatValue_Fixed_BUILD_VECTOR) {
   TargetLowering TL(*TM);
 

llvm/unittests/Support/KnownBitsTest.cpp

@@ -213,6 +213,37 @@ TEST(KnownBitsTest, AddSubExhaustive) {
   TestAddSubExhaustive(false);
 }
 
+TEST(KnownBitsTest, SubBorrowExhaustive) {
+  unsigned Bits = 4;
+  ForeachKnownBits(Bits, [&](const KnownBits &Known1) {
+    ForeachKnownBits(Bits, [&](const KnownBits &Known2) {
+      ForeachKnownBits(1, [&](const KnownBits &KnownBorrow) {
+        // Explicitly compute known bits of the subtraction by trying all
+        // possibilities.
+        KnownBits Known(Bits);
+        Known.Zero.setAllBits();
+        Known.One.setAllBits();
+        ForeachNumInKnownBits(Known1, [&](const APInt &N1) {
+          ForeachNumInKnownBits(Known2, [&](const APInt &N2) {
+            ForeachNumInKnownBits(KnownBorrow, [&](const APInt &Borrow) {
+              APInt Sub = N1 - N2;
+              if (Borrow.getBoolValue())
+                --Sub;
+
+              Known.One &= Sub;
+              Known.Zero &= ~Sub;
+            });
+          });
+        });
+
+        KnownBits KnownComputed =
+            KnownBits::computeForSubBorrow(Known1, Known2, KnownBorrow);
+        EXPECT_EQ(Known, KnownComputed);
+      });
+    });
+  });
+}
+
 TEST(KnownBitsTest, BinaryExhaustive) {
   testBinaryOpExhaustive(
       [](const KnownBits &Known1, const KnownBits &Known2) {