2 files changed, 3 insertions, 7 deletions
diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index be867b24ba6..d242d3c159a 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -7128,11 +7128,7 @@ SDValue DAGCombiner::visitFunnelShift(SDNode *N) {
       return IsFSHL ? N0 : N1;
 
   auto IsUndefOrZero = [](SDValue V) {
-    if (V.isUndef())
-      return true;
-    if (ConstantSDNode *Cst = isConstOrConstSplat(V, /*AllowUndefs*/true))
-      return Cst->getAPIntValue() == 0;
-    return false;
+    return V.isUndef() || isNullOrNullSplat(V, /*AllowUndefs*/ true);
   };
 
   if (ConstantSDNode *Cst = isConstOrConstSplat(N2)) {
diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
index 2e92f9582be..9b00a550e33 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
@@ -8622,9 +8622,9 @@ ConstantFPSDNode *llvm::isConstOrConstSplatFP(SDValue N, bool AllowUndefs) {
   return nullptr;
 }
 
-bool llvm::isNullOrNullSplat(SDValue N) {
+bool llvm::isNullOrNullSplat(SDValue N, bool AllowUndefs) {
   // TODO: may want to use peekThroughBitcast() here.
-  ConstantSDNode *C = isConstOrConstSplat(N);
+  ConstantSDNode *C = isConstOrConstSplat(N, AllowUndefs);
   return C && C->isNullValue();
 }