[llvm] [InstSimplify] Fix simplifyAndOrWithICmpEq with undef refinement (PR #98898)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Nikita Popov (nikic)

<details>
<summary>Changes</summary>

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 https://github.com/llvm/llvm-project/issues/98753.

---
Full diff: https://github.com/llvm/llvm-project/pull/98898.diff


2 Files Affected:

- (modified) llvm/lib/Analysis/InstructionSimplify.cpp (+23-7) 
- (modified) llvm/test/Transforms/InstSimplify/and-or-implied-cond.ll (+2-2) 


``````````diff
diff --git a/llvm/lib/Analysis/InstructionSimplify.cpp b/llvm/lib/Analysis/InstructionSimplify.cpp
index 242c200f7ef15..3a7ae577bb068 100644
--- a/llvm/lib/Analysis/InstructionSimplify.cpp
+++ b/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;
diff --git a/llvm/test/Transforms/InstSimplify/and-or-implied-cond.ll b/llvm/test/Transforms/InstSimplify/and-or-implied-cond.ll
index 8620dce059ae7..99e1dd4528697 100644
--- a/llvm/test/Transforms/InstSimplify/and-or-implied-cond.ll
+++ b/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

``````````

</details>


https://github.com/llvm/llvm-project/pull/98898