[InstSimplify] Fix simplifyAndOrWithICmpEq with undef refinement

The final case in Simplify (where Res == Absorber and the predicate
is inverted) is not generally safe when the simplification is a
refinement. In particular, we may simplify assuming a specific value
for undef, but then chose a different one later.

However, it *is* safe to refine poison in this context, unlike in
the equivalent select folds. This is the reason why this fold did
not use AllowRefinement=false in the first place, and using that
option would introduce a lot of test regressions.

This patch takes the middle path of disabling undef refinements
in particular using the getWithoutUndef() SimplifyQuery option.
However, this option doesn't actually work in this case, because
the problematic fold is inside constant folding, and we currently
don't propagate this option all the way from InstSimplify over
ConstantFolding to ConstantFold. Work around this by explicitly
checking for undef operands in simplifyWithOpReplaced().

Finally, make sure that places where AllowRefinement=false also
use Q.getWithoutUndef(). I don't have a specific test case for
this (the original one does not work because we don't simplify
selects with constant condition in this mode in the first place)
but this seems like the correct thing to do to be conservative.

Fixes llvm#98753.

Author: nikic

llvm/lib/Analysis/InstructionSimplify.cpp

@@ -1975,13 +1975,16 @@ static Value *simplifyAndOrWithICmpEq(unsigned Opcode, Value *Op0, Value *Op1,
     return nullptr;
   };
 
-  if (Value *Res =
-          simplifyWithOpReplaced(Op1, A, B, Q, /* AllowRefinement */ true,
-                                 /* DropFlags */ nullptr, MaxRecurse))
+  // In the final case (Res == Absorber with inverted predicate), it is safe to
+  // refine poison during simplification, but not undef. For simplicity always
+  // disable undef-based folds here.
+  if (Value *Res = simplifyWithOpReplaced(Op1, A, B, Q.getWithoutUndef(),
+                                          /* AllowRefinement */ true,
+                                          /* DropFlags */ nullptr, MaxRecurse))
     return Simplify(Res);
-  if (Value *Res =
-          simplifyWithOpReplaced(Op1, B, A, Q, /* AllowRefinement */ true,
-                                 /* DropFlags */ nullptr, MaxRecurse))
+  if (Value *Res = simplifyWithOpReplaced(Op1, B, A, Q.getWithoutUndef(),
+                                          /* AllowRefinement */ true,
+                                          /* DropFlags */ nullptr, MaxRecurse))
     return Simplify(Res);
 
   return nullptr;
@@ -4300,6 +4303,9 @@ static Value *simplifyWithOpReplaced(Value *V, Value *Op, Value *RepOp,
                                      bool AllowRefinement,
                                      SmallVectorImpl<Instruction *> *DropFlags,
                                      unsigned MaxRecurse) {
+  assert((AllowRefinement || !Q.CanUseUndef) &&
+         "If AllowRefinement=false then CanUseUndef=false");
+
   // Trivial replacement.
   if (V == Op)
     return RepOp;
@@ -4347,6 +4353,11 @@ static Value *simplifyWithOpReplaced(Value *V, Value *Op, Value *RepOp,
     } else {
       NewOps.push_back(InstOp);
     }
+
+    // Bail out if any operand is undef and SimplifyQuery disables undef
+    // simplification. Constant folding currently doesn't respect this option.
+    if (isa<UndefValue>(NewOps.back()) && !Q.CanUseUndef)
+      return nullptr;
   }
 
   if (!AnyReplaced)
@@ -4467,6 +4478,11 @@ Value *llvm::simplifyWithOpReplaced(Value *V, Value *Op, Value *RepOp,
                                     const SimplifyQuery &Q,
                                     bool AllowRefinement,
                                     SmallVectorImpl<Instruction *> *DropFlags) {
+  // If refinement is disabled, also disable undef simplifications (which are
+  // always refinements) in SimplifyQuery.
+  if (!AllowRefinement)
+    return ::simplifyWithOpReplaced(V, Op, RepOp, Q.getWithoutUndef(),
+                                    AllowRefinement, DropFlags, RecursionLimit);
   return ::simplifyWithOpReplaced(V, Op, RepOp, Q, AllowRefinement, DropFlags,
                                   RecursionLimit);
 }
@@ -4606,7 +4622,7 @@ static Value *simplifySelectWithICmpEq(Value *CmpLHS, Value *CmpRHS,
                                        Value *TrueVal, Value *FalseVal,
                                        const SimplifyQuery &Q,
                                        unsigned MaxRecurse) {
-  if (simplifyWithOpReplaced(FalseVal, CmpLHS, CmpRHS, Q,
+  if (simplifyWithOpReplaced(FalseVal, CmpLHS, CmpRHS, Q.getWithoutUndef(),
                              /* AllowRefinement */ false,
                              /* DropFlags */ nullptr, MaxRecurse) == TrueVal)
     return FalseVal;

llvm/test/Transforms/InstSimplify/and-or-implied-cond.ll

@@ -331,13 +331,13 @@ define i1 @and_is_constant(ptr %arg, ptr %arg2) {
   ret i1 %and
 }
 
-; FIXME: This is a miscompile.
 define i1 @pr98753(i32 noundef %x, i32 %y) {
 ; CHECK-LABEL: @pr98753(
 ; CHECK-NEXT:    [[CMP1:%.*]] = icmp ne i32 [[X:%.*]], 0
 ; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP1]], i32 [[Y:%.*]], i32 undef
 ; CHECK-NEXT:    [[CMP2:%.*]] = icmp sgt i32 [[SEL]], 0
-; CHECK-NEXT:    ret i1 [[CMP2]]
+; CHECK-NEXT:    [[AND:%.*]] = and i1 [[CMP1]], [[CMP2]]
+; CHECK-NEXT:    ret i1 [[AND]]
 ;
   %cmp1 = icmp ne i32 %x, 0
   %sel = select i1 %cmp1, i32 %y, i32 undef