[LVI][CVP] Make use of condition known at use

[nikic]

Summary

When an instruction is only used in a select or phi operand, we might be able to make use of additional information from the select/branch condition. For example in

%sub = call i16 @llvm.usub.sat.i16(i16 %x, i16 10)
%cmp = icmp uge i16 %x, 10
%sel = select i1 %cmp, i16 %sub, i16 42

the usub.sat is only used in a select where %x uge 10 is known to hold, so we can fold it based on that knowledge.

This addresses the regression reported at https://reviews.llvm.org/D140798#4039748, but also provides a solution to a recurring problem we've had, where we fail to make use of range information after a branch+phi has been converted into a select. Our current solution to this is to hope that IPSCCP can perform the fold before that happens, but handling this in LVI is a somewhat more general solution.

Currently we only make use of this for the willNotOverflow() fold, but I plan to adjust other folds to use the new API as well.

[TimNN]

Thanks a lot @nikic! I've verified that this unbreaks Rust's codegen test.

[spatel]

Seems like a nice improvement to me, but adding some more potential reviewers in case there are other opinions.

If the intent is to eventually call from other places, the most basic case would be something like this?

define i16 @sel_true_cond(i16 %x) {
  %cmp = icmp eq i16 %x, 10
  %sel = select i1 %cmp, i16 %x, i16 42   ; can replace %x with 10
  ret i16 %sel
}

If so, might want to adjust the test names because patterns with an overflow intrinsic are not the minimal case.

llvm/lib/Analysis/LazyValueInfo.cpp

1662

The meaning of this loop wasn't clear to me on first reading.

If the intent is to experiment with optimization power vs. compile-time, make this a cl::opt or give the limit a name like static const unsigned MaxUserRangesToIntersect?

[nikic]

In D141482#4044348, @spatel wrote:

Seems like a nice improvement to me, but adding some more potential reviewers in case there are other opinions.

If the intent is to eventually call from other places, the most basic case would be something like this?

define i16 @sel_true_cond(i16 %x) {
  %cmp = icmp eq i16 %x, 10
  %sel = select i1 %cmp, i16 %x, i16 42   ; can replace %x with 10
  ret i16 %sel
}

Basically yes, but CVP doesn't actually perform folds that simple (just replacing a value with a known constant), it's generally some kind of instruction simplification (we do this in InstCombine though, in foldSelectValueEquivalence).

[nikic]

Give magic constant 2 a name.

[StephenFan]

llvm/lib/Analysis/LazyValueInfo.cpp

1680

Is limiting the number of uses a consideration of compile-time? If we can give it a mutable limit and support handling multiple uses?

[nikic]

llvm/lib/Analysis/LazyValueInfo.cpp

1680

It's a correctness requirement for the current implementation. For multiple uses, we would have to intersect with the union of conditions at all uses. This is more complex, expensive and unlikely to be useful.

[spatel]

LGTM

llvm/lib/Analysis/LazyValueInfo.cpp

1680

Since there was a question about the hasOneUse(), it would be good to include that reasoning as a code comment.

llvm/include/llvm/Analysis/LazyValueInfo.h

public:
   Constant *getConstant(Value *V, Instruction *CxtI);
 
   /// Return the ConstantRange constraint that is known to hold for the
   /// specified value at the specified instruction. This may only be called
   /// on integer-typed Values.
   ConstantRange getConstantRange(Value *V, Instruction *CxtI,
                                  bool UndefAllowed = true);
 
+  /// Return the ConstantRange constraint that is known to hold for the value
+  /// at a specific use-site.
+  ConstantRange getConstantRangeAtUse(const Use &U, bool UndefAllowed = true);
+
   /// Determine whether the specified value is known to be a
   /// constant on the specified edge.  Return null if not.
   Constant *getConstantOnEdge(Value *V, BasicBlock *FromBB, BasicBlock *ToBB,
                               Instruction *CxtI = nullptr);
 
   /// Return the ConstantRage constraint that is known to hold for the
   /// specified value on the specified edge. This may be only be called
   /// on integer-typed Values.

llvm/lib/Analysis/LazyValueInfo.cpp

if (Result.isConstantRange(UndefAllowed))
     return Result.getConstantRange(UndefAllowed);
   // We represent ConstantInt constants as constant ranges but other kinds
   // of integer constants, i.e. ConstantExpr will be tagged as constants
   assert(!(Result.isConstant() && isa<ConstantInt>(Result.getConstant())) &&
          "ConstantInt value must be represented as constantrange");
   return ConstantRange::getFull(Width);
 }
 
+ConstantRange LazyValueInfo::getConstantRangeAtUse(const Use &U,
+                                                   bool UndefAllowed) {
+  Value *V = U.get();
+  ConstantRange CR =
+      getConstantRange(V, cast<Instruction>(U.getUser()), UndefAllowed);
+
+  // Check whether the only (possibly transitive) use of the value is in a
+  // position where V can be constrained by a select or branch condition.
+  const Use *CurrU = &U;
+  // TODO: Increase limit?
+  const unsigned MaxUsesToInspect = 2;
+  for (unsigned I = 0; I < MaxUsesToInspect; ++I) {
+    std::optional<ValueLatticeElement> CondVal;
+    auto *CurrI = cast<Instruction>(CurrU->getUser());
+    if (auto *SI = dyn_cast<SelectInst>(CurrI)) {
+      if (CurrU->getOperandNo() == 1)
+        CondVal = getValueFromCondition(V, SI->getCondition(), true);
+      else if (CurrU->getOperandNo() == 2)
+        CondVal = getValueFromCondition(V, SI->getCondition(), false);
+    } else if (auto *PHI = dyn_cast<PHINode>(CurrI)) {
+      // TODO: Use non-local query?
+      CondVal =
+          getEdgeValueLocal(V, PHI->getIncomingBlock(*CurrU), PHI->getParent());
+    }
+    if (CondVal && CondVal->isConstantRange())
+      CR = CR.intersectWith(CondVal->getConstantRange());
+
+    // Only follow one-use chain, to allow direct intersection of conditions.
+    // If there are multiple uses, we would have to intersect with the union of
+    // all conditions at different uses.
+    if (!CurrI->hasOneUse())
+      break;
+    CurrU = &*CurrI->use_begin();
+  }
+  return CR;
+}
+
 /// Determine whether the specified value is known to be a
 /// constant on the specified edge. Return null if not.
 Constant *LazyValueInfo::getConstantOnEdge(Value *V, BasicBlock *FromBB,
                                            BasicBlock *ToBB,
                                            Instruction *CxtI) {
   Module *M = FromBB->getModule();
   ValueLatticeElement Result =
       getImpl(PImpl, AC, M).getValueOnEdge(V, FromBB, ToBB, CxtI);

llvm/lib/Transforms/Scalar/CorrelatedValuePropagation.cpp

if (Changed)
     // by a branch then do so now.
     ConstantFoldTerminator(BB, /*DeleteDeadConditions = */ false,
                            /*TLI = */ nullptr, &DTU);
   return Changed;
 }
 
 // See if we can prove that the given binary op intrinsic will not overflow.
 static bool willNotOverflow(BinaryOpIntrinsic *BO, LazyValueInfo *LVI) {
-  ConstantRange LRange = LVI->getConstantRange(BO->getLHS(), BO);
-  ConstantRange RRange = LVI->getConstantRange(BO->getRHS(), BO);
+  ConstantRange LRange = LVI->getConstantRangeAtUse(BO->getOperandUse(0));
+  ConstantRange RRange = LVI->getConstantRangeAtUse(BO->getOperandUse(1));
   ConstantRange NWRegion = ConstantRange::makeGuaranteedNoWrapRegion(
       BO->getBinaryOp(), RRange, BO->getNoWrapKind());
   return NWRegion.contains(LRange);
 }
 
 static void setDeducedOverflowingFlags(Value *V, Instruction::BinaryOps Opcode,
                                        bool NewNSW, bool NewNUW) {
   Statistic *OpcNW, *OpcNSW, *OpcNUW;

llvm/test/Transforms/CorrelatedValuePropagation/cond-at-use.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -S -passes=correlated-propagation < %s | FileCheck %s
 
 declare void @use.i16(i16)
 declare i16 @llvm.usub.sat.i16(i16, i16)
 declare {i16, i1} @llvm.usub.with.overflow.i16(i16, i16)
 
 define i16 @sel_true_cond(i16 %x) {
 ; CHECK-LABEL: @sel_true_cond(
-; CHECK-NEXT:    [[SUB:%.*]] = call i16 @llvm.usub.sat.i16(i16 [[X:%.*]], i16 10)
+; CHECK-NEXT:    [[SUB1:%.*]] = sub nuw i16 [[X:%.*]], 10
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp uge i16 [[X]], 10
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i16 [[SUB]], i16 42
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i16 [[SUB1]], i16 42
 ; CHECK-NEXT:    ret i16 [[SEL]]
 ;
   %sub = call i16 @llvm.usub.sat.i16(i16 %x, i16 10)
   %cmp = icmp uge i16 %x, 10
   %sel = select i1 %cmp, i16 %sub, i16 42
   ret i16 %sel
 }
 

 ;
   %sub = call i16 @llvm.usub.sat.i16(i16 %x, i16 10)
   call void @use.i16(i16 %sub)
   %cmp = icmp uge i16 %x, 10
   %sel = select i1 %cmp, i16 %sub, i16 42
   ret i16 %sel
 }
 
+; TODO: We could handle this case by raising the limit on the number of
+; instructions we look through.
 define i16 @sel_true_cond_longer_chain(i16 %x) {
 ; CHECK-LABEL: @sel_true_cond_longer_chain(
 ; CHECK-NEXT:    [[SUB:%.*]] = call i16 @llvm.usub.sat.i16(i16 [[X:%.*]], i16 10)
 ; CHECK-NEXT:    [[EXTRA:%.*]] = mul i16 [[SUB]], 3
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp uge i16 [[X]], 10
 ; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i16 [[EXTRA]], i16 42
 ; CHECK-NEXT:    ret i16 [[SEL]]
 ;
   %sub = call i16 @llvm.usub.sat.i16(i16 %x, i16 10)
   %extra = mul i16 %sub, 3
   %cmp = icmp uge i16 %x, 10
   %sel = select i1 %cmp, i16 %extra, i16 42
   ret i16 %sel
 }
 
 define i16 @sel_false_cond(i16 %x) {
 ; CHECK-LABEL: @sel_false_cond(
-; CHECK-NEXT:    [[SUB:%.*]] = call i16 @llvm.usub.sat.i16(i16 [[X:%.*]], i16 10)
+; CHECK-NEXT:    [[SUB1:%.*]] = sub nuw i16 [[X:%.*]], 10
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i16 [[X]], 10
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i16 42, i16 [[SUB]]
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i16 42, i16 [[SUB1]]
 ; CHECK-NEXT:    ret i16 [[SEL]]
 ;
   %sub = call i16 @llvm.usub.sat.i16(i16 %x, i16 10)
   %cmp = icmp ult i16 %x, 10
   %sel = select i1 %cmp, i16 42, i16 %sub
   ret i16 %sel
 }
 
 define i16 @sel_false_cond_insufficient(i16 %x) {
 ; CHECK-LABEL: @sel_false_cond_insufficient(
 ; CHECK-NEXT:    [[SUB:%.*]] = call i16 @llvm.usub.sat.i16(i16 [[X:%.*]], i16 10)
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i16 [[X]], 9
 ; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i16 42, i16 [[SUB]]
 ; CHECK-NEXT:    ret i16 [[SEL]]
 ;
   %sub = call i16 @llvm.usub.sat.i16(i16 %x, i16 10)
   %cmp = icmp ult i16 %x, 9
   %sel = select i1 %cmp, i16 42, i16 %sub
   ret i16 %sel
 }
 
 define i16 @phi_true_cond(i16 %x) {
 ; CHECK-LABEL: @phi_true_cond(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[SUB:%.*]] = call i16 @llvm.usub.sat.i16(i16 [[X:%.*]], i16 10)
+; CHECK-NEXT:    [[SUB1:%.*]] = sub nuw i16 [[X:%.*]], 10
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp uge i16 [[X]], 10
 ; CHECK-NEXT:    br i1 [[CMP]], label [[JOIN:%.*]], label [[SPLIT:%.*]]
 ; CHECK:       split:
 ; CHECK-NEXT:    br label [[JOIN]]
 ; CHECK:       join:
-; CHECK-NEXT:    [[PHI:%.*]] = phi i16 [ [[SUB]], [[ENTRY:%.*]] ], [ 42, [[SPLIT]] ]
+; CHECK-NEXT:    [[PHI:%.*]] = phi i16 [ [[SUB1]], [[ENTRY:%.*]] ], [ 42, [[SPLIT]] ]
 ; CHECK-NEXT:    ret i16 [[PHI]]
 ;
 entry:
   %sub = call i16 @llvm.usub.sat.i16(i16 %x, i16 10)
   %cmp = icmp uge i16 %x, 10
   br i1 %cmp, label %join, label %split
 
 split:

 split:
   br label %join
 
 join:
   %phi = phi i16 [ %sub, %entry ], [ 42, %split ]
   ret i16 %phi
 }
 
+; TODO: We could handle this by using conditions that are not directly on the
+; phi edge.
 define i16 @phi_true_cond_non_local(i16 %x) {
 ; CHECK-LABEL: @phi_true_cond_non_local(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[SUB:%.*]] = call i16 @llvm.usub.sat.i16(i16 [[X:%.*]], i16 10)
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp uge i16 [[X]], 10
 ; CHECK-NEXT:    br i1 [[CMP]], label [[SPLIT:%.*]], label [[JOIN:%.*]]
 ; CHECK:       split:
 ; CHECK-NEXT:    br label [[JOIN]]

 join:
   %phi = phi i16 [ %sub, %split ], [ 42, %entry ]
   ret i16 %phi
 }
 
 define i16 @phi_false_cond(i16 %x) {
 ; CHECK-LABEL: @phi_false_cond(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[SUB:%.*]] = call i16 @llvm.usub.sat.i16(i16 [[X:%.*]], i16 10)
+; CHECK-NEXT:    [[SUB1:%.*]] = sub nuw i16 [[X:%.*]], 10
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i16 [[X]], 10
 ; CHECK-NEXT:    br i1 [[CMP]], label [[SPLIT:%.*]], label [[JOIN:%.*]]
 ; CHECK:       split:
 ; CHECK-NEXT:    br label [[JOIN]]
 ; CHECK:       join:
-; CHECK-NEXT:    [[PHI:%.*]] = phi i16 [ [[SUB]], [[ENTRY:%.*]] ], [ 42, [[SPLIT]] ]
+; CHECK-NEXT:    [[PHI:%.*]] = phi i16 [ [[SUB1]], [[ENTRY:%.*]] ], [ 42, [[SPLIT]] ]
 ; CHECK-NEXT:    ret i16 [[PHI]]
 ;
 entry:
   %sub = call i16 @llvm.usub.sat.i16(i16 %x, i16 10)
   %cmp = icmp ult i16 %x, 10
   br i1 %cmp, label %split, label %join
 
 split:

 split:
   br label %join
 
 join:
   %phi = phi i16 [ %sub, %entry ], [ 42, %split ]
   ret i16 %phi
 }
 
+; TODO: We could handle this by using conditions that are not directly on the
+; phi edge.
 define i16 @phi_false_cond_non_local(i16 %x) {
 ; CHECK-LABEL: @phi_false_cond_non_local(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[SUB:%.*]] = call i16 @llvm.usub.sat.i16(i16 [[X:%.*]], i16 10)
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i16 [[X]], 10
 ; CHECK-NEXT:    br i1 [[CMP]], label [[JOIN:%.*]], label [[SPLIT:%.*]]
 ; CHECK:       split:
 ; CHECK-NEXT:    br label [[JOIN]]
join:
   ret i16 %phi
 }
 
 define i16 @loop_cond() {
 ; CHECK-LABEL: @loop_cond(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
-; CHECK-NEXT:    [[IV:%.*]] = phi i16 [ 1000, [[ENTRY:%.*]] ], [ [[IV_NEXT:%.*]], [[LOOP]] ]
+; CHECK-NEXT:    [[IV:%.*]] = phi i16 [ 1000, [[ENTRY:%.*]] ], [ [[IV_NEXT1:%.*]], [[LOOP]] ]
 ; CHECK-NEXT:    [[COUNT:%.*]] = phi i16 [ 0, [[ENTRY]] ], [ [[COUNT_NEXT:%.*]], [[LOOP]] ]
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i16 [[IV]], 0
-; CHECK-NEXT:    [[IV_NEXT]] = call i16 @llvm.usub.sat.i16(i16 [[IV]], i16 1)
+; CHECK-NEXT:    [[IV_NEXT1]] = sub nuw i16 [[IV]], 1
 ; CHECK-NEXT:    [[COUNT_NEXT]] = add i16 [[COUNT]], 1
 ; CHECK-NEXT:    br i1 [[CMP]], label [[EXIT:%.*]], label [[LOOP]]
 ; CHECK:       exit:
 ; CHECK-NEXT:    ret i16 [[COUNT]]
 ;
 entry:
   br label %loop