[InstCombine] Whitelist non-refining folds in SimplifyWithOpReplaced

[nikic]

Summary

This is an alternative to D98391/D98585, playing things more conservatively. If AllowRefinement == false, then we don't use InstSimplify methods at all, and instead explicitly implement some non-refining folds. Most cases are handled by constant folding, and I only had to add two folds to cover our test cases. While this may lose some optimization power, I think it is safer to approach from this direction, given how many issues this has already caused.

[aqjune]

It is impressive that leaving the few patterns was enough to cover all existing unit tests.
If more patterns are needed, we can factor out the non-refined foldings from SimplifyXXInst as a static function and call it here.

Since ConstantFoldInstOperands can return a refined value as well, I gave a shot and disabled the folding if AllowRefinement = false.
However, it resulted in many unit test failures (implying that it will prevent a lot of optimizations). I couldn't write an input that involves constant folding but results in miscompilation either; so I think it is fine.

LGTM, but it would be good if someone else can double-check the diffs as well.

llvm/lib/Analysis/InstructionSimplify.cpp

4042

This might make the result more refined as well, is it? ConstantFoldBinaryInstruction may fold undef/poison into a well-defined value.
I think it is better to leave a comment here saying the issue.

[spatel]

LGTM - but see inline for general question/comment.

llvm/lib/Analysis/InstructionSimplify.cpp

3935–3936

I lost track of exactly what we mean by "refinement". The existing example in the code comment below + the test diffs in this patch suggest it is something like:
"Substitution with another value can't let poison flow where it did not before." ?
Are we getting in trouble with selects only or are there other opcodes with problems?

It can be independent of this patch, but we should make "refinement" clearer in the code comments here or in the header file comment where SimplifyWithOpReplaced() is declared.

[nikic]

Add fold for binop + identity element. This is the only other relevant fold I found in test-suite. Shuffle comments a bit.

[nikic]

llvm/lib/Analysis/InstructionSimplify.cpp

3935–3936

Refinement means making a replacement like poison -> 0. This is legal. The select transform in question will take the simplification result (0) and replace it with the original instruction (poison). If the original simplification was a refinement, then this is a de-refinement, which is not legal. We can only make replacements that are valid in both directions, which requires the simplification result to be identical to the input. I believe the select optimization is the only place where we try to do something like this.

4042

There is a TODO For this in the canCreatePoison comment, but it ended up a bit far from the code it affects. I've moved it directly before the constant folding code now, which is hopefully clearer.

[aqjune]

llvm/lib/Analysis/InstructionSimplify.cpp

4019–4020

ok, looks good.

llvm/lib/Analysis/InstructionSimplify.cpp

static Value *SimplifyWithOpReplaced(Value *V, Value *Op, Value *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;
 
+  if (!AllowRefinement) {
+    // General InstSimplify functions may refine the result. To avoid this,
+    // implement only a few non-refining but profitable transforms here.
+
+    // x & x -> x, x | x -> x
+    if (match(I, m_c_And(m_Specific(Op), m_Specific(RepOp))))
+      return RepOp;
+    if (match(I, m_c_Or(m_Specific(Op), m_Specific(RepOp))))
+      return RepOp;
+
+    if (auto *GEP = dyn_cast<GetElementPtrInst>(I)) {
+      // getelementptr x, 0 -> x
+      if (GEP->getNumIndices() == 1 && *GEP->idx_begin() == Op &&
+          match(RepOp, m_Zero()) && !GEP->isInBounds())
+        return GEP->getPointerOperand();
+    }
+
     // 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)))
+    // We can't replace %sel with %add unless we strip away the flags (which
+    // will be done in InstCombine).
+    // TODO: Now that this only guards the constant folding case, it might
+    // be possible to drop this check.
+    if (canCreatePoison(cast<Operator>(I)))
       return nullptr;
-
+  } else {
     // 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
     // simplifies back to %arg. This can only happen because %mul does not
     // dominate %div. To ensure a consistent return value contract, we make sure
     // that this case returns nullptr as well.
     auto PreventSelfSimplify = [V](Value *Simplified) {
       return Simplified != V ? Simplified : nullptr;
     };
 
     // If this is a binary operator, try to simplify it with the replaced op.
     if (auto *B = dyn_cast<BinaryOperator>(I)) {
       if (MaxRecurse) {
         if (B->getOperand(0) == Op)
           return PreventSelfSimplify(SimplifyBinOp(B->getOpcode(), RepOp,
                                                    B->getOperand(1), Q,
                                                    MaxRecurse - 1));
         if (B->getOperand(1) == Op)
           return PreventSelfSimplify(SimplifyBinOp(B->getOpcode(),
                                                    B->getOperand(0), RepOp, Q,
                                                    MaxRecurse - 1));
       }
     }
 
     // Same for CmpInsts.
     if (CmpInst *C = dyn_cast<CmpInst>(I)) {
       if (MaxRecurse) {
         if (C->getOperand(0) == Op)
           return PreventSelfSimplify(SimplifyCmpInst(C->getPredicate(), RepOp,
                                                      C->getOperand(1), Q,
                                                      MaxRecurse - 1));
         if (C->getOperand(1) == Op)
           return PreventSelfSimplify(SimplifyCmpInst(C->getPredicate(),
                                                      C->getOperand(0), RepOp, Q,
                                                      MaxRecurse - 1));
       }
     }
 
     // Same for GEPs.
     if (auto *GEP = dyn_cast<GetElementPtrInst>(I)) {
       if (MaxRecurse) {
         SmallVector<Value *, 8> NewOps(GEP->getNumOperands());
         transform(GEP->operands(), NewOps.begin(),
                   [&](Value *V) { return V == Op ? RepOp : V; });
         return PreventSelfSimplify(SimplifyGEPInst(GEP->getSourceElementType(),
                                                    NewOps, Q, MaxRecurse - 1));
       }
     }
 
     // TODO: We could hand off more cases to instsimplify here.
+  }
 
   // If all operands are constant after substituting Op for RepOp then we can
   // constant fold the instruction.
   if (Constant *CRepOp = dyn_cast<Constant>(RepOp)) {
     // Build a list of all constant operands.
     SmallVector<Constant *, 8> ConstOps;
     for (unsigned i = 0, e = I->getNumOperands(); i != e; ++i) {
       if (I->getOperand(i) == Op)
         ConstOps.push_back(CRepOp);
       else if (Constant *COp = dyn_cast<Constant>(I->getOperand(i)))
         ConstOps.push_back(COp);
       else
         break;
     }
 
     // All operands were constants, fold it.
     if (ConstOps.size() == I->getNumOperands()) {
       if (CmpInst *C = dyn_cast<CmpInst>(I))
         return ConstantFoldCompareInstOperands(C->getPredicate(), ConstOps[0],
                                                ConstOps[1], Q.DL, Q.TLI);
 
       if (LoadInst *LI = dyn_cast<LoadInst>(I))
         if (!LI->isVolatile())
           return ConstantFoldLoadFromConstPtr(ConstOps[0], LI->getType(), Q.DL);
 
       return ConstantFoldInstOperands(I, ConstOps, Q.DL, Q.TLI);
     }
   }
 
   return nullptr;
 }
 
 Value *llvm::SimplifyWithOpReplaced(Value *V, Value *Op, Value *RepOp,
                                     const SimplifyQuery &Q,

llvm/test/Transforms/InstCombine/select.ll

;
   %conv = zext i16 %x to i32
   %cond = select i1 %tobool, i32 0, i32 %conv
   ret i32 %cond
 }
 
 define i32 @test57(i32 %x, i32 %y) {
 ; CHECK-LABEL: @test57(
 ; CHECK-NEXT:    [[AND:%.*]] = and i32 [[X:%.*]], [[Y:%.*]]
-; CHECK-NEXT:    ret i32 [[AND]]
+; CHECK-NEXT:    [[TOBOOL:%.*]] = icmp eq i32 [[X]], 0
+; CHECK-NEXT:    [[DOTAND:%.*]] = select i1 [[TOBOOL]], i32 0, i32 [[AND]]
+; CHECK-NEXT:    ret i32 [[DOTAND]]
 ;
   %and = and i32 %x, %y
   %tobool = icmp eq i32 %x, 0
   %.and = select i1 %tobool, i32 0, i32 %and
   ret i32 %.and
 }
 
 define i32 @test58(i16 %x) {
;
   %v = select i1 %c.fr, i32 %x, i32 %y
   call void @use_i1_i32(i1 %c.fr, i32 %v)
   ret void
 }
 
 define i32 @pr47322_more_poisonous_replacement(i32 %arg) {
 ; CHECK-LABEL: @pr47322_more_poisonous_replacement(
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[ARG:%.*]], 0
-; CHECK-NEXT:    [[TRAILING:%.*]] = call i32 @llvm.cttz.i32(i32 [[ARG]], i1 immarg true), [[RNG0:!range !.*]]
+; CHECK-NEXT:    [[TRAILING:%.*]] = call i32 @llvm.cttz.i32(i32 [[ARG]], i1 immarg true), !range [[RNG0:![0-9]+]]
 ; CHECK-NEXT:    [[SHIFTED:%.*]] = lshr i32 [[ARG]], [[TRAILING]]
 ; CHECK-NEXT:    [[R1_SROA_0_1:%.*]] = select i1 [[CMP]], i32 0, i32 [[SHIFTED]]
 ; CHECK-NEXT:    ret i32 [[R1_SROA_0_1]]
 ;
   %cmp = icmp eq i32 %arg, 0
   %trailing = call i32 @llvm.cttz.i32(i32 %arg, i1 immarg true)
   %shifted = lshr i32 %arg, %trailing
   %r1.sroa.0.1 = select i1 %cmp, i32 0, i32 %shifted

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
 }