diff --git a/llvm/unittests/Support/KnownBitsTest.cpp b/llvm/unittests/Support/KnownBitsTest.cpp
--- a/llvm/unittests/Support/KnownBitsTest.cpp
+++ b/llvm/unittests/Support/KnownBitsTest.cpp
@@ -28,7 +28,11 @@
 using BinaryCheckFn =
     llvm::function_ref<bool(const KnownBits &, const KnownBits &)>;
 
+static bool checkOptimalityUnary(const KnownBits &) { return true; }
 static bool checkCorrectnessOnlyUnary(const KnownBits &) { return false; }
+static bool checkOptimalityBinary(const KnownBits &, const KnownBits &) {
+  return true;
+}
 static bool checkCorrectnessOnlyBinary(const KnownBits &, const KnownBits &) {
   return false;
 }
@@ -62,9 +66,9 @@
   return Result;
 }
 
-static void testUnaryOpExhaustive(
-    UnaryBitsFn BitsFn, UnaryIntFn IntFn,
-    UnaryCheckFn CheckOptimalityFn = [](const KnownBits &) { return true; }) {
+static void
+testUnaryOpExhaustive(UnaryBitsFn BitsFn, UnaryIntFn IntFn,
+                      UnaryCheckFn CheckOptimalityFn = checkOptimalityUnary) {
   unsigned Bits = 4;
   ForeachKnownBits(Bits, [&](const KnownBits &Known) {
     KnownBits Computed = BitsFn(Known);
@@ -89,11 +93,10 @@
   });
 }
 
-static void testBinaryOpExhaustive(
-    BinaryBitsFn BitsFn, BinaryIntFn IntFn,
-    BinaryCheckFn CheckOptimalityFn = [](const KnownBits &, const KnownBits &) {
-      return true;
-    }) {
+static void
+testBinaryOpExhaustive(BinaryBitsFn BitsFn, BinaryIntFn IntFn,
+                       BinaryCheckFn CheckOptimalityFn = checkOptimalityBinary,
+                       bool RefinePoisonToZero = false) {
   unsigned Bits = 4;
   ForeachKnownBits(Bits, [&](const KnownBits &Known1) {
     ForeachKnownBits(Bits, [&](const KnownBits &Known2) {
@@ -118,6 +121,10 @@
       if (CheckOptimalityFn(Known1, Known2) && !Exact.hasConflict()) {
         EXPECT_TRUE(isOptimal(Exact, Computed, {Known1, Known2}));
       }
+      // In some cases we choose to return zero if the result is always poison.
+      if (RefinePoisonToZero && Exact.hasConflict()) {
+        EXPECT_TRUE(Computed.isZero());
+      }
     });
   });
 }
@@ -342,7 +349,8 @@
         if (N2.uge(N2.getBitWidth()))
           return std::nullopt;
         return N1.shl(N2);
-      });
+      },
+      checkOptimalityBinary, /* RefinePoisonToZero */ true);
   testBinaryOpExhaustive(
       [](const KnownBits &Known1, const KnownBits &Known2) {
         return KnownBits::shl(Known1, Known2, /* NUW */ true);
@@ -353,7 +361,8 @@
         if (Overflow)
           return std::nullopt;
         return Res;
-      });
+      },
+      checkOptimalityBinary, /* RefinePoisonToZero */ true);
   testBinaryOpExhaustive(
       [](const KnownBits &Known1, const KnownBits &Known2) {
         return KnownBits::shl(Known1, Known2, /* NUW */ false, /* NSW */ true);
@@ -364,7 +373,8 @@
         if (Overflow)
           return std::nullopt;
         return Res;
-      });
+      },
+      checkOptimalityBinary, /* RefinePoisonToZero */ true);
   testBinaryOpExhaustive(
       [](const KnownBits &Known1, const KnownBits &Known2) {
         return KnownBits::shl(Known1, Known2, /* NUW */ true, /* NSW */ true);
@@ -376,7 +386,8 @@
         if (OverflowUnsigned || OverflowSigned)
           return std::nullopt;
         return Res;
-      });
+      },
+      checkOptimalityBinary, /* RefinePoisonToZero */ true);
 
   testBinaryOpExhaustive(
       [](const KnownBits &Known1, const KnownBits &Known2) {
@@ -386,7 +397,8 @@
         if (N2.uge(N2.getBitWidth()))
           return std::nullopt;
         return N1.lshr(N2);
-      });
+      },
+      checkOptimalityBinary, /* RefinePoisonToZero */ true);
   testBinaryOpExhaustive(
       [](const KnownBits &Known1, const KnownBits &Known2) {
         return KnownBits::ashr(Known1, Known2);
@@ -395,7 +407,8 @@
         if (N2.uge(N2.getBitWidth()))
           return std::nullopt;
         return N1.ashr(N2);
-      });
+      },
+      checkOptimalityBinary, /* RefinePoisonToZero */ true);
 
   testBinaryOpExhaustive(
       [](const KnownBits &Known1, const KnownBits &Known2) {