[InstSimplify] Propagate AllowRefinement info from SimplifyWithOpReplaced

[aqjune]

Summary

SimplifyWithOpReplaced has AllowRefinement argument, but when it further calls SimplifyXXX,
the flag isn't properly passed.
This causes the simplifying calls to fold possibly-poison values into non-poison constant
even if AllowRefinement is set to false.

This causes miscompilation https://llvm.org/pr49495.
simplifySelectWithICmpCond calls SimplifyWithOpReplaced with AllowRefinement set to false.
However, the information is lost, causing computePointerICmp to do folding based on inbounds flag.
computePointerICmp transforms (gep inbounds x, -1) >u x into -1 >s 0 which is correct only if
inbounds is set (https://alive2.llvm.org/ce/z/mGwYxH vs. https://alive2.llvm.org/ce/z/e-2KQh ).

To fix this, this patch updates simplifySelectWithICmpCond to call SimplifyWithOpReplaced with Q.getWithoutRefinement().
SimplifyQuery::getWithoutRefinement() is a new method that simply returns a new object with relevant flags set to false.
If this flag is set, simplifications cannot fold the value into a more defined one.

This wasn't enough to fix the pointer folding, and I updated computePointerICmp and its relevant calls to respect relevant flags.

[aqjune]

llvm/test/Transforms/InstCombine/select.ll

2766 ↗

(On Diff #329838)

This is a regression; after this patch is landed, I'll make a fix for this.

[aqjune]

Fix deoptimized case

[aqjune]

Confirmed that this resolves the regression (https://reviews.llvm.org/D95026#2610334) from two-stage build!

[nikic]

I'm not sure this is moving in the right direction. It should be noted that the primary refinement problem in SimplifyWithOpReplaced, which we have encountered in various forms in the past, is along the lines of folding add nuw UINT_MAX, 1 to 0 and thus refining poision to zero. The key here is that this would be valid without the nuw flag. This is the converse problem of what UseInstrInfo is solving. In fact, if add nuw UINT_MAX, 1 is folded to poison, then that result is okay to use, as no refinement occurs. This can also be seen in the changes to stripPointerCasts you are making -- effectively you are preventing optimization across inbounds GEP, even though the point of UseInstrInfo is to optimize without making use of the inbounds information.

I don't think there's any easy principled solution here, so what I'd do instead is to just extend the existing canCreatePoison() hack, which has served us well enough in practice, to also go one level up for compare instructions. (The more definitive way to prevent issues is to check that it isGuaranteedNotUndefOrPoison, but I expect that this would break many transforms.)

[aqjune]

TBH, I don't feel comfortable with extending the canCreatePoison() check. :( Even if it requires inspection of foldings in InstSimplify, I think a safer way is to add a flag that disallows foldings returning refined values.
Let me update this patch and write following patches for updating them.

[aqjune]

Add AllowRefinement field

[aqjune]

minor updates

[aqjune]

Make computePointerICmp slightly more conservative if Q.AllowRefinement is false

It is folding possibly poison/undef value into a well-defined constant.

[aqjune]

Does this direction look okay? @nikic
One benefit of this approach is that it detects unsafe replacements that have non-poison-creating instructions as well.
For example, this transformation is illegal:

  %v = select (x == 0), 0, (x & y)
=>
  %v = x & y

If x = 0 and y was poison this makes the result more undefined. However, the current check relying on canCreatePoison cannot detect this.
D98585 fixes this by making the relevant transformation respect Q.AllowRefinement.

This change is about making the existing check safer than before: if the change causes a problem with generating optimal code I'll look into it.

[aqjune]

In terms of maintenance: if reviewers think the complexity this patch is introducing is big to maintain, I'll turn this patch to using canCreatePoison check since it will fix the bug (pr49495) anyway.

[nikic]

I've put up D99027 as a possible alternative, though I'm not really sure yet which approach is best here. My premise is that for the AllowRefinement=false case the main thing we care about are the constant folding cases, plus a small set of additional folds. This avoids infecting the rest of InstSimplify with another parameter, and we can't run into issues when we forgot the check the flag somewhere. On the other hand, this approach may lose optimization power.

[aqjune]

Closed in favor of D99027

llvm/include/llvm/Analysis/InstructionSimplify.h

 //===-- InstructionSimplify.h - Fold instrs into simpler forms --*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
struct InstrInfoQuery {
   }
 
   template <class InstT> bool hasNoSignedWrap(const InstT *Op) const {
     if (UseInstrInfo)
       return Op->hasNoSignedWrap();
     return false;
   }
 
+  bool isInBounds(const GEPOperator *Op) const {
+    return UseInstrInfo && Op->isInBounds();
+  }
+
   bool isExact(const BinaryOperator *Op) const {
     if (UseInstrInfo && isa<PossiblyExactOperator>(Op))
       return cast<PossiblyExactOperator>(Op)->isExact();
     return false;
   }
 };
 
 struct SimplifyQuery {
   const DataLayout &DL;
   const TargetLibraryInfo *TLI = nullptr;
   const DominatorTree *DT = nullptr;
   AssumptionCache *AC = nullptr;
   const Instruction *CxtI = nullptr;
 
   // Wrapper to query additional information for instructions like metadata or
   // keywords like nsw, which provides conservative results if those cannot
   // be safely used.
-  const InstrInfoQuery IIQ;
+  InstrInfoQuery IIQ;
 
   /// Controls whether simplifications are allowed to constrain the range of
   /// possible values for uses of undef. If it is false, simplifications are not
   /// allowed to assume a particular value for a use of undef for example.
   bool CanUseUndef = true;
 
+  /// Controls whether returning more defined value is allowed. If it is false,
+  /// simplifications must not make the value more defined. These foldings
+  /// become illegal, for example:
+  ///   undef            => 0
+  ///   (x +nsw 1) >s x  => true
+  ///   (x +nsw 1) - 1   => x
+  bool AllowRefinement = true;
+
   SimplifyQuery(const DataLayout &DL, const Instruction *CXTI = nullptr)
       : DL(DL), CxtI(CXTI) {}
 
   SimplifyQuery(const DataLayout &DL, const TargetLibraryInfo *TLI,
                 const DominatorTree *DT = nullptr,
                 AssumptionCache *AC = nullptr,
                 const Instruction *CXTI = nullptr, bool UseInstrInfo = true,
                 bool CanUseUndef = true)
       : DL(DL), TLI(TLI), DT(DT), AC(AC), CxtI(CXTI), IIQ(UseInstrInfo),
         CanUseUndef(CanUseUndef) {}
   SimplifyQuery getWithInstruction(Instruction *I) const {
     SimplifyQuery Copy(*this);
     Copy.CxtI = I;
     return Copy;
   }
   SimplifyQuery getWithoutUndef() const {
     SimplifyQuery Copy(*this);
     Copy.CanUseUndef = false;
     return Copy;
   }
+  SimplifyQuery getWithoutRefinement() const {
+    SimplifyQuery Copy = getWithoutUndef();
+    Copy.AllowRefinement = false;
+    /// Actually, setting IIQ.UseInstrInfo to false disables more
+    /// transformations than necessary. For example,
+    ///   (INT_MAX +nsw 1) => poison
+    /// is okay because it does not make the result more defined, but setting
+    /// the flag to false turns this folding off.
+    /// However, there are quite a few refining transformations that use
+    /// nsw/inbounds/... (e.g. "x +nsw 1 >s x" => true), and IIQ.UseInstrInfo
+    /// is set to false to easily opt out such transformations.
+    Copy.IIQ.UseInstrInfo = false;
+    return Copy;
+  }
 
   /// If CanUseUndef is true, returns whether \p V is undef.
   /// Otherwise always return false.
   bool isUndefValue(Value *V) const {
     if (!CanUseUndef)
       return false;
     return isa<UndefValue>(V);
   }

 Value *SimplifyInstruction(Instruction *I, const SimplifyQuery &Q,
                            OptimizationRemarkEmitter *ORE = nullptr);
 
 /// See if V simplifies when its operand Op is replaced with RepOp. If not,
 /// return null.
 /// AllowRefinement specifies whether the simplification can be a refinement,
 /// or whether it needs to be strictly identical.
 Value *SimplifyWithOpReplaced(Value *V, Value *Op, Value *RepOp,
-                              const SimplifyQuery &Q, bool AllowRefinement);
+                              const SimplifyQuery &Q);
 
 /// Replace all uses of 'I' with 'SimpleV' and simplify the uses recursively.
 ///
 /// This first performs a normal RAUW of I with SimpleV. It then recursively
 /// attempts to simplify those users updated by the operation. The 'I'
 /// instruction must not be equal to the simplified value 'SimpleV'.
 /// If UnsimplifiedUsers is provided, instructions that could not be simplified
 /// are added to it.

llvm/lib/Analysis/InstructionSimplify.cpp

 //===- InstructionSimplify.cpp - Fold instruction operands ----------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
static Constant *stripAndComputeConstantOffsets(const DataLayout &DL, Value *&V,
   Constant *OffsetIntPtr = ConstantInt::get(IntIdxTy, Offset);
   if (VectorType *VecTy = dyn_cast<VectorType>(V->getType()))
     return ConstantVector::getSplat(VecTy->getElementCount(), OffsetIntPtr);
   return OffsetIntPtr;
 }
 
 /// Compute the constant difference between two pointer values.
 /// If the difference is not a constant, returns zero.
-static Constant *computePointerDifference(const DataLayout &DL, Value *LHS,
+static Constant *computePointerDifference(const SimplifyQuery &Q, Value *LHS,
                                           Value *RHS) {
-  Constant *LHSOffset = stripAndComputeConstantOffsets(DL, LHS);
-  Constant *RHSOffset = stripAndComputeConstantOffsets(DL, RHS);
+  if (!Q.AllowRefinement)
+    // Folding e.g., '(gep p 1) - p' to '1' may result in returning more
+    // defined value. Conservatively bail out if AllowRefinement is false.
+    return nullptr;
+
+  Constant *LHSOffset = stripAndComputeConstantOffsets(Q.DL, LHS);
+  Constant *RHSOffset = stripAndComputeConstantOffsets(Q.DL, RHS);
 
   // If LHS and RHS are not related via constant offsets to the same base
   // value, there is nothing we can do here.
   if (LHS != RHS)
     return nullptr;
 
   // Otherwise, the difference of LHS - RHS can be computed as:
   //    LHS - RHS

 // equal to the beginning of a stack object in the callee.
 //
 // If the C and C++ standards are ever made sufficiently restrictive in this
 // area, it may be possible to update LLVM's semantics accordingly and reinstate
 // this optimization.
 static Constant *
 computePointerICmp(CmpInst::Predicate Pred, Value *LHS, Value *RHS,
                    const SimplifyQuery &Q) {
+  if (!Q.AllowRefinement)
+    // Conservatively return nullptr since it may result in returning more
+    // defined value.
+    return nullptr;
+
   const DataLayout &DL = Q.DL;
   const TargetLibraryInfo *TLI = Q.TLI;
   const DominatorTree *DT = Q.DT;
   const Instruction *CxtI = Q.CxtI;
   const InstrInfoQuery &IIQ = Q.IIQ;
 
   // First, skip past any trivial no-ops.
   LHS = LHS->stripPointerCasts();
if (auto *CRHS = dyn_cast<PtrToIntOperator>(RHS))
                                          CRHS->getPointerOperand(), Q))
           return C;
 
   if (GetElementPtrInst *GLHS = dyn_cast<GetElementPtrInst>(LHS)) {
     if (GEPOperator *GRHS = dyn_cast<GEPOperator>(RHS)) {
       if (GLHS->getPointerOperand() == GRHS->getPointerOperand() &&
           GLHS->hasAllConstantIndices() && GRHS->hasAllConstantIndices() &&
           (ICmpInst::isEquality(Pred) ||
-           (GLHS->isInBounds() && GRHS->isInBounds() &&
+           (Q.IIQ.isInBounds(cast<GEPOperator>(GLHS)) &&
+            Q.IIQ.isInBounds(GRHS) &&
             Pred == ICmpInst::getSignedPredicate(Pred)))) {
         // The bases are equal and the indices are constant.  Build a constant
         // expression GEP with the same indices and a null base pointer to see
         // what constant folding can make out of it.
         Constant *Null = Constant::getNullValue(GLHS->getPointerOperandType());
         SmallVector<Value *, 4> IndicesLHS(GLHS->indices());
         Constant *NewLHS = ConstantExpr::getGetElementPtr(
             GLHS->getSourceElementType(), Null, IndicesLHS);

 
 Value *llvm::SimplifyFCmpInst(unsigned Predicate, Value *LHS, Value *RHS,
                               FastMathFlags FMF, const SimplifyQuery &Q) {
   return ::SimplifyFCmpInst(Predicate, LHS, RHS, FMF, Q, RecursionLimit);
 }
 
 static Value *SimplifyWithOpReplaced(Value *V, Value *Op, Value *RepOp,
                                      const SimplifyQuery &Q,
-                                     bool AllowRefinement,
                                      unsigned MaxRecurse) {
   // Trivial replacement.
   if (V == Op)
     return RepOp;
 
   // We cannot replace a constant, and shouldn't even try.
   if (isa<Constant>(Op))
     return nullptr;
 
   auto *I = dyn_cast<Instruction>(V);
   if (!I)
     return nullptr;
 
   // Consider:
   //   %cmp = icmp eq i32 %x, 2147483647
   //   %add = add nsw i32 %x, 1
   //   %sel = select i1 %cmp, i32 -2147483648, i32 %add
   //
   // We can't replace %sel with %add unless we strip away the flags (which will
   // be done in InstCombine).
   // TODO: This is unsound, because it only catches some forms of refinement.
-  if (!AllowRefinement && canCreatePoison(cast<Operator>(I)))
+  if (!Q.AllowRefinement && canCreatePoison(cast<Operator>(I)))
     return nullptr;
 
   // The simplification queries below may return the original value. Consider:
   //   %div = udiv i32 %arg, %arg2
   //   %mul = mul nsw i32 %div, %arg2
   //   %cmp = icmp eq i32 %mul, %arg
   //   %sel = select i1 %cmp, i32 %div, i32 undef
   // Replacing %arg by %mul, %div becomes "udiv i32 %mul, %arg2", which
if (ConstOps.size() == I->getNumOperands()) {
       return ConstantFoldInstOperands(I, ConstOps, Q.DL, Q.TLI);
     }
   }
 
   return nullptr;
 }
 
 Value *llvm::SimplifyWithOpReplaced(Value *V, Value *Op, Value *RepOp,
-                                    const SimplifyQuery &Q,
-                                    bool AllowRefinement) {
-  return ::SimplifyWithOpReplaced(V, Op, RepOp, Q, AllowRefinement,
-                                  RecursionLimit);
+                                    const SimplifyQuery &Q) {
+  return ::SimplifyWithOpReplaced(V, Op, RepOp, Q, RecursionLimit);
 }
 
 /// Try to simplify a select instruction when its condition operand is an
 /// integer comparison where one operand of the compare is a constant.
 static Value *simplifySelectBitTest(Value *TrueVal, Value *FalseVal, Value *X,
                                     const APInt *Y, bool TrueWhenUnset) {
   const APInt *C;
 
static Value *simplifySelectWithICmpCond(Value *CondVal, Value *TrueVal,
   if (Value *V = simplifySelectWithFakeICmpEq(CmpLHS, CmpRHS, Pred,
                                               TrueVal, FalseVal))
     return V;
 
   // If we have an equality comparison, then we know the value in one of the
   // arms of the select. See if substituting this value into the arm and
   // simplifying the result yields the same value as the other arm.
   if (Pred == ICmpInst::ICMP_EQ) {
-    if (SimplifyWithOpReplaced(FalseVal, CmpLHS, CmpRHS, Q,
-                               /* AllowRefinement */ false, MaxRecurse) ==
-            TrueVal ||
-        SimplifyWithOpReplaced(FalseVal, CmpRHS, CmpLHS, Q,
-                               /* AllowRefinement */ false, MaxRecurse) ==
-            TrueVal)
+    if (SimplifyWithOpReplaced(FalseVal, CmpLHS, CmpRHS,
+                               Q.getWithoutRefinement(),
+                               MaxRecurse) == TrueVal ||
+        SimplifyWithOpReplaced(FalseVal, CmpRHS, CmpLHS,
+                               Q.getWithoutRefinement(), MaxRecurse) == TrueVal)
       return FalseVal;
-    if (SimplifyWithOpReplaced(TrueVal, CmpLHS, CmpRHS, Q,
-                               /* AllowRefinement */ true, MaxRecurse) ==
+    if (SimplifyWithOpReplaced(TrueVal, CmpLHS, CmpRHS, Q, MaxRecurse) ==
             FalseVal ||
-        SimplifyWithOpReplaced(TrueVal, CmpRHS, CmpLHS, Q,
-                               /* AllowRefinement */ true, MaxRecurse) ==
+        SimplifyWithOpReplaced(TrueVal, CmpRHS, CmpLHS, Q, MaxRecurse) ==
             FalseVal)
       return FalseVal;
   }
 
   return nullptr;
 }
 
 /// Try to simplify a select instruction when its condition operand is a

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

 //===- InstCombineSelect.cpp ----------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
Instruction *InstCombinerImpl::foldSelectValueEquivalence(SelectInst &Sel,
   // In X == Y ? f(X) : Z, try to evaluate f(Y) and replace the operand.
   // Make sure Y cannot be undef though, as we might pick different values for
   // undef in the icmp and in f(Y). Additionally, take care to avoid replacing
   // X == Y ? X : Z with X == Y ? Y : Z, as that would lead to an infinite
   // replacement cycle.
   Value *CmpLHS = Cmp.getOperand(0), *CmpRHS = Cmp.getOperand(1);
   if (TrueVal != CmpLHS &&
       isGuaranteedNotToBeUndefOrPoison(CmpRHS, SQ.AC, &Sel, &DT)) {
-    if (Value *V = SimplifyWithOpReplaced(TrueVal, CmpLHS, CmpRHS, SQ,
-                                          /* AllowRefinement */ true))
+    if (Value *V = SimplifyWithOpReplaced(TrueVal, CmpLHS, CmpRHS, SQ))
       return replaceOperand(Sel, Swapped ? 2 : 1, V);
 
     // Even if TrueVal does not simplify, we can directly replace a use of
     // CmpLHS with CmpRHS, as long as the instruction is not used anywhere
     // else and is safe to speculatively execute (we may end up executing it
     // with different operands, which should not cause side-effects or trigger
     // undefined behavior). Only do this if CmpRHS is a constant, as
     // profitability is not clear for other cases.
     // FIXME: The replacement could be performed recursively.
     if (match(CmpRHS, m_ImmConstant()) && !match(CmpLHS, m_ImmConstant()))
       if (auto *I = dyn_cast<Instruction>(TrueVal))
         if (I->hasOneUse() && isSafeToSpeculativelyExecute(I))
           for (Use &U : I->operands())
             if (U == CmpLHS) {
               replaceUse(U, CmpRHS);
               return &Sel;
             }
   }
   if (TrueVal != CmpRHS &&
       isGuaranteedNotToBeUndefOrPoison(CmpLHS, SQ.AC, &Sel, &DT))
-    if (Value *V = SimplifyWithOpReplaced(TrueVal, CmpRHS, CmpLHS, SQ,
-                                          /* AllowRefinement */ true))
+    if (Value *V = SimplifyWithOpReplaced(TrueVal, CmpRHS, CmpLHS, SQ))
       return replaceOperand(Sel, Swapped ? 2 : 1, V);
 
   auto *FalseInst = dyn_cast<Instruction>(FalseVal);
   if (!FalseInst)
     return nullptr;
 
   // InstSimplify already performed this fold if it was possible subject to
   // current poison-generating flags. Try the transform again with
if (auto *GEP = dyn_cast<GetElementPtrInst>(FalseVal)) {
     WasInBounds = GEP->isInBounds();
     GEP->setIsInBounds(false);
   }
 
   // Try each equivalence substitution possibility.
   // We have an 'EQ' comparison, so the select's false value will propagate.
   // Example:
   // (X == 42) ? 43 : (X + 1) --> (X == 42) ? (X + 1) : (X + 1) --> X + 1
-  if (SimplifyWithOpReplaced(FalseVal, CmpLHS, CmpRHS, SQ,
-                             /* AllowRefinement */ false) == TrueVal ||
-      SimplifyWithOpReplaced(FalseVal, CmpRHS, CmpLHS, SQ,
-                             /* AllowRefinement */ false) == TrueVal) {
+  if (SimplifyWithOpReplaced(FalseVal, CmpLHS, CmpRHS,
+                             SQ.getWithoutRefinement()) == TrueVal ||
+      SimplifyWithOpReplaced(FalseVal, CmpRHS, CmpLHS,
+                             SQ.getWithoutRefinement()) == TrueVal) {
     return replaceInstUsesWith(Sel, FalseVal);
   }
 
   // Restore poison-generating flags if the transform did not apply.
   if (WasNUW)
     FalseInst->setHasNoUnsignedWrap();
   if (WasNSW)
     FalseInst->setHasNoSignedWrap();

llvm/test/Transforms/InstSimplify/pr49495.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -instsimplify -S | FileCheck %s
 
 ; The first comparison (a != b) should not be dropped
 define i1 @test1(i8* %a, i8* %b) {
 ; CHECK-LABEL: @test1(
-; CHECK-NEXT:    [[A2:%.*]] = getelementptr inbounds i8, i8* [[A:%.*]], i64 -1
-; CHECK-NEXT:    [[COND2:%.*]] = icmp ugt i8* [[A2]], [[B:%.*]]
-; CHECK-NEXT:    ret i1 [[COND2]]
+; CHECK-NEXT:    [[COND1:%.*]] = icmp ne i8* [[A:%.*]], [[B:%.*]]
+; CHECK-NEXT:    [[A2:%.*]] = getelementptr inbounds i8, i8* [[A]], i64 -1
+; CHECK-NEXT:    [[COND2:%.*]] = icmp ugt i8* [[A2]], [[B]]
+; CHECK-NEXT:    [[RES:%.*]] = select i1 [[COND1]], i1 [[COND2]], i1 false
+; CHECK-NEXT:    ret i1 [[RES]]
 ;
   %cond1 = icmp ne i8* %a, %b
   %a2 = getelementptr inbounds i8, i8* %a, i64 -1
   %cond2 = icmp ugt i8* %a2, %b
   %res = select i1 %cond1, i1 %cond2, i1 false
   ret i1 %res
 }
 
 ; The first comparison (a != b) should not be dropped
 define i1 @test2(i32 %a, i32 %b) {
 ; CHECK-LABEL: @test2(
-; CHECK-NEXT:    [[A2:%.*]] = add nuw i32 [[A:%.*]], 1
-; CHECK-NEXT:    [[COND2:%.*]] = icmp ult i32 [[A2]], [[B:%.*]]
-; CHECK-NEXT:    ret i1 [[COND2]]
+; CHECK-NEXT:    [[COND1:%.*]] = icmp ne i32 [[A:%.*]], [[B:%.*]]
+; CHECK-NEXT:    [[A2:%.*]] = add nuw i32 [[A]], 1
+; CHECK-NEXT:    [[COND2:%.*]] = icmp ult i32 [[A2]], [[B]]
+; CHECK-NEXT:    [[RES:%.*]] = select i1 [[COND1]], i1 [[COND2]], i1 false
+; CHECK-NEXT:    ret i1 [[RES]]
 ;
   %cond1 = icmp ne i32 %a, %b
   %a2 = add nuw i32 %a, 1
   %cond2 = icmp ult i32 %a2, %b
   %res = select i1 %cond1, i1 %cond2, i1 false
   ret i1 %res
 }