GlobalISel: Start adding computeNumSignBits to GISelKnownBits

3 files changed, 152 insertions, 0 deletions

diff --git a/llvm/include/llvm/CodeGen/GlobalISel/GISelKnownBits.h b/llvm/include/llvm/CodeGen/GlobalISel/GISelKnownBits.h
index dfe5a7f3177..d44612f54ae 100644
--- a/llvm/include/llvm/CodeGen/GlobalISel/GISelKnownBits.h
+++ b/llvm/include/llvm/CodeGen/GlobalISel/GISelKnownBits.h
@@ -40,6 +40,10 @@ public:
                                     const APInt &DemandedElts,
                                     unsigned Depth = 0);
 
+  unsigned computeNumSignBits(Register R, const APInt &DemandedElts,
+                              unsigned Depth = 0);
+  unsigned computeNumSignBits(Register R, unsigned Depth = 0);
+
   // KnownBitsAPI
   KnownBits getKnownBits(Register R);
   // Calls getKnownBits for first operand def of MI.
diff --git a/llvm/lib/CodeGen/GlobalISel/GISelKnownBits.cpp b/llvm/lib/CodeGen/GlobalISel/GISelKnownBits.cpp
index 41ef631906e..64023ecfad8 100644
--- a/llvm/lib/CodeGen/GlobalISel/GISelKnownBits.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/GISelKnownBits.cpp
@@ -373,6 +373,76 @@ void GISelKnownBits::computeKnownBitsImpl(Register R, KnownBits &Known,
                     << Known.One.toString(16, false) << "\n");
 }
 
+unsigned GISelKnownBits::computeNumSignBits(Register R,
+                                            const APInt &DemandedElts,
+                                            unsigned Depth) {
+  MachineInstr &MI = *MRI.getVRegDef(R);
+  unsigned Opcode = MI.getOpcode();
+
+  if (Opcode == TargetOpcode::G_CONSTANT)
+    return MI.getOperand(1).getCImm()->getValue().getNumSignBits();
+
+  if (Depth == getMaxDepth())
+    return 1;
+
+  if (!DemandedElts)
+    return 1; // No demanded elts, better to assume we don't know anything.
+
+  LLT DstTy = MRI.getType(R);
+
+  // Handle the case where this is called on a register that does not have a
+  // type constraint. This is unlikely to occur except by looking through copies
+  // but it is possible for the initial register being queried to be in this
+  // state.
+  if (!DstTy.isValid())
+    return 1;
+
+  switch (Opcode) {
+  case TargetOpcode::COPY: {
+    MachineOperand &Src = MI.getOperand(1);
+    if (Src.getReg().isVirtual() && Src.getSubReg() == 0 &&
+        MRI.getType(Src.getReg()).isValid()) {
+      // Don't increment Depth for this one since we didn't do any work.
+      return computeNumSignBits(Src.getReg(), DemandedElts, Depth);
+    }
+
+    return 1;
+  }
+  case TargetOpcode::G_SEXT: {
+    Register Src = MI.getOperand(1).getReg();
+    LLT SrcTy = MRI.getType(Src);
+    unsigned Tmp = DstTy.getScalarSizeInBits() - SrcTy.getScalarSizeInBits();
+    return computeNumSignBits(Src, DemandedElts, Depth + 1) + Tmp;
+  }
+  case TargetOpcode::G_TRUNC: {
+    Register Src = MI.getOperand(1).getReg();
+    LLT SrcTy = MRI.getType(Src);
+
+    // Check if the sign bits of source go down as far as the truncated value.
+    unsigned DstTyBits = DstTy.getScalarSizeInBits();
+    unsigned NumSrcBits = SrcTy.getScalarSizeInBits();
+    unsigned NumSrcSignBits = computeNumSignBits(Src, DemandedElts, Depth + 1);
+    if (NumSrcSignBits > (NumSrcBits - DstTyBits))
+      return NumSrcSignBits - (NumSrcBits - DstTyBits);
+    break;
+  }
+  default:
+    break;
+  }
+
+  // TODO: Handle target instructions
+  // TODO: Fall back to known bits
+  return 1;
+}
+
+unsigned GISelKnownBits::computeNumSignBits(Register R, unsigned Depth) {
+  LLT Ty = MRI.getType(R);
+  APInt DemandedElts = Ty.isVector()
+                           ? APInt::getAllOnesValue(Ty.getNumElements())
+                           : APInt(1, 1);
+  return computeNumSignBits(R, DemandedElts, Depth);
+}
+
 void GISelKnownBitsAnalysis::getAnalysisUsage(AnalysisUsage &AU) const {
   AU.setPreservesAll();
   MachineFunctionPass::getAnalysisUsage(AU);
diff --git a/llvm/unittests/CodeGen/GlobalISel/KnownBitsTest.cpp b/llvm/unittests/CodeGen/GlobalISel/KnownBitsTest.cpp
index f7b037866e0..e6c92f39e44 100644
--- a/llvm/unittests/CodeGen/GlobalISel/KnownBitsTest.cpp
+++ b/llvm/unittests/CodeGen/GlobalISel/KnownBitsTest.cpp
@@ -133,3 +133,81 @@ TEST_F(GISelMITest, TestSignBitIsZero) {
   EXPECT_TRUE(KnownBits.signBitIsZero(Zero.getReg(0)));
   EXPECT_FALSE(KnownBits.signBitIsZero(SignBit.getReg(0)));
 }
+
+TEST_F(GISelMITest, TestNumSignBitsConstant) {
+  StringRef MIRString = "  %3:_(s8) = G_CONSTANT i8 1\n"
+                        "  %4:_(s8) = COPY %3\n"
+
+                        "  %5:_(s8) = G_CONSTANT i8 -1\n"
+                        "  %6:_(s8) = COPY %5\n"
+
+                        "  %7:_(s8) = G_CONSTANT i8 127\n"
+                        "  %8:_(s8) = COPY %7\n"
+
+                        "  %9:_(s8) = G_CONSTANT i8 32\n"
+                        "  %10:_(s8) = COPY %9\n"
+
+                        "  %11:_(s8) = G_CONSTANT i8 -32\n"
+                        "  %12:_(s8) = COPY %11\n";
+  setUp(MIRString);
+  if (!TM)
+    return;
+  Register CopyReg1 = Copies[Copies.size() - 5];
+  Register CopyRegNeg1 = Copies[Copies.size() - 4];
+  Register CopyReg127 = Copies[Copies.size() - 3];
+  Register CopyReg32 = Copies[Copies.size() - 2];
+  Register CopyRegNeg32 = Copies[Copies.size() - 1];
+
+  GISelKnownBits Info(*MF);
+  EXPECT_EQ(7u, Info.computeNumSignBits(CopyReg1));
+  EXPECT_EQ(8u, Info.computeNumSignBits(CopyRegNeg1));
+  EXPECT_EQ(1u, Info.computeNumSignBits(CopyReg127));
+  EXPECT_EQ(2u, Info.computeNumSignBits(CopyReg32));
+  EXPECT_EQ(3u, Info.computeNumSignBits(CopyRegNeg32));
+}
+
+TEST_F(GISelMITest, TestNumSignBitsSext) {
+  StringRef MIRString = "  %3:_(p0) = G_IMPLICIT_DEF\n"
+                        "  %4:_(s8) = G_LOAD %3 :: (load 1)\n"
+                        "  %5:_(s32) = G_SEXT %4\n"
+                        "  %6:_(s32) = COPY %5\n"
+
+                        "  %7:_(s8) = G_CONSTANT i8 -1\n"
+                        "  %8:_(s32) = G_SEXT %7\n"
+                        "  %9:_(s32) = COPY %8\n";
+  setUp(MIRString);
+  if (!TM)
+    return;
+  Register CopySextLoad = Copies[Copies.size() - 2];
+  Register CopySextNeg1 = Copies[Copies.size() - 1];
+
+  GISelKnownBits Info(*MF);
+  EXPECT_EQ(25u, Info.computeNumSignBits(CopySextLoad));
+  EXPECT_EQ(32u, Info.computeNumSignBits(CopySextNeg1));
+}
+
+TEST_F(GISelMITest, TestNumSignBitsTrunc) {
+  StringRef MIRString = "  %3:_(p0) = G_IMPLICIT_DEF\n"
+                        "  %4:_(s32) = G_LOAD %3 :: (load 4)\n"
+                        "  %5:_(s8) = G_TRUNC %4\n"
+                        "  %6:_(s8) = COPY %5\n"
+
+                        "  %7:_(s32) = G_CONSTANT i32 -1\n"
+                        "  %8:_(s8) = G_TRUNC %7\n"
+                        "  %9:_(s8) = COPY %8\n"
+
+                        "  %10:_(s32) = G_CONSTANT i32 7\n"
+                        "  %11:_(s8) = G_TRUNC %10\n"
+                        "  %12:_(s8) = COPY %11\n";
+  setUp(MIRString);
+  if (!TM)
+    return;
+  Register CopyTruncLoad = Copies[Copies.size() - 3];
+  Register CopyTruncNeg1 = Copies[Copies.size() - 2];
+  Register CopyTrunc7 = Copies[Copies.size() - 1];
+
+  GISelKnownBits Info(*MF);
+  EXPECT_EQ(1u, Info.computeNumSignBits(CopyTruncLoad));
+  EXPECT_EQ(8u, Info.computeNumSignBits(CopyTruncNeg1));
+  EXPECT_EQ(5u, Info.computeNumSignBits(CopyTrunc7));
+}