[CVP] [LVI] Add LVI::getPredicateInBlock and use it in CVP.

[jlebar]

Summary

LVI::getPredicateInBlock is like LVI::getPredicateAt, but uses the more
powerful getValueInBlock API, instead of getValueAt. This lets it solve
things like x & y < z.

This came out of a discussion in D44102, in which sanjoy@ realized that
getConstantRange, which calls getValueInBlock, was more powerful than
getPredicateAt, which calls getValueAt.

[reames]

Sorry for just getting around to reviewing this now. Should be a bit more response now.

Looking at the approach, it appears you're effectively working around a implementation limitation in getPredicateAt. As an alternative, why not simply extend getPredicateAt to compute facts known at the begining of the block specified by the context pointer, and then merge that with the facts from within the block?

[sanjoy]

llvm/include/llvm/Analysis/LazyValueInfo.h

84

Can we pick the block from CxtI?

llvm/include/llvm/Analysis/LazyValueInfo.h

public:
   /// Determine whether the specified value comparison with a constant is known
   /// to be true or false on the specified CFG edge.
   /// Pred is a CmpInst predicate.
   Tristate getPredicateOnEdge(unsigned Pred, Value *V, Constant *C,
                               BasicBlock *FromBB, BasicBlock *ToBB,
                               Instruction *CxtI = nullptr);
 
   /// Determine whether the specified value comparison with a constant is known
-  /// to be true or false at the specified instruction
-  /// (from an assume intrinsic). Pred is a CmpInst predicate.
+  /// to be true or false at the specified instruction Pred is a CmpInst
+  /// predicate.
+  ///
+  /// Because this function does not take a basic block, it generally only looks
+  /// at assume intrinsics and range metadata.  This makes it significantly less
+  /// powerful than getConstantRange or getPredicateAtBlock.
   Tristate getPredicateAt(unsigned Pred, Value *V, Constant *C,
                           Instruction *CxtI);
 
+  /// Like getPredicateAt, tries to determine whether the V <Pred> C is true.
+  /// This version evaluates V at the end of BB.
+  ///
+  /// This is more powerful than getPredicateAt; it doesn't just look at assume
+  /// intrinsics and range metadata.
+  Tristate getPredicateInBlock(unsigned Pred, Value *V, Constant *C,
+                               BasicBlock *BB, Instruction *CxtI = nullptr);
+
   /// Determine whether the specified value is known to be a
   /// constant at the end of the specified block.  Return null if not.
   Constant *getConstant(Value *V, BasicBlock *BB, Instruction *CxtI = nullptr);
 
   /// Return the ConstantRange constraint that is known to hold for the
   /// specified value at the end of the specified block. This may only be called
   /// on integer-typed Values.
   ConstantRange getConstantRange(Value *V, BasicBlock *BB, Instruction *CxtI = nullptr);

llvm/lib/Analysis/LazyValueInfo.cpp

LazyValueInfo::getPredicateOnEdge(unsigned Pred, Value *V, Constant *C,
                                   Instruction *CxtI) {
   const DataLayout &DL = FromBB->getModule()->getDataLayout();
   ValueLatticeElement Result =
       getImpl(PImpl, AC, &DL, DT).getValueOnEdge(V, FromBB, ToBB, CxtI);
 
   return getPredicateResult(Pred, C, Result, DL, TLI);
 }
 
-LazyValueInfo::Tristate
-LazyValueInfo::getPredicateAt(unsigned Pred, Value *V, Constant *C,
-                              Instruction *CxtI) {
-  // Is or is not NonNull are common predicates being queried. If
-  // isKnownNonZero can tell us the result of the predicate, we can
-  // return it quickly. But this is only a fastpath, and falling
-  // through would still be correct.
-  const DataLayout &DL = CxtI->getModule()->getDataLayout();
+// If Pred and C are a null or not-null check, quickly checks whether V is known
+// null/not-null.
+//
+// This is a fast approximation, used because these checks are common.
+static LazyValueInfo::Tristate fastCheckIsOrIsNotNull(unsigned Pred, Value *V,
+                                                      Constant *C,
+                                                      const DataLayout &DL) {
   if (V->getType()->isPointerTy() && C->isNullValue() &&
       isKnownNonZero(V->stripPointerCasts(), DL)) {
     if (Pred == ICmpInst::ICMP_EQ)
       return LazyValueInfo::False;
     else if (Pred == ICmpInst::ICMP_NE)
       return LazyValueInfo::True;
   }
-  ValueLatticeElement Result = getImpl(PImpl, AC, &DL, DT).getValueAt(V, CxtI);
-  Tristate Ret = getPredicateResult(Pred, C, Result, DL, TLI);
-  if (Ret != Unknown)
-    return Ret;
+  return LazyValueInfo::Unknown;
+}
 
-  // Note: The following bit of code is somewhat distinct from the rest of LVI;
-  // LVI as a whole tries to compute a lattice value which is conservatively
-  // correct at a given location.  In this case, we have a predicate which we
-  // weren't able to prove about the merged result, and we're pushing that
-  // predicate back along each incoming edge to see if we can prove it
-  // separately for each input.  As a motivating example, consider:
-  // bb1:
-  //   %v1 = ... ; constantrange<1, 5>
-  //   br label %merge
-  // bb2:
-  //   %v2 = ... ; constantrange<10, 20>
-  //   br label %merge
-  // merge:
-  //   %phi = phi [%v1, %v2] ; constantrange<1,20>
-  //   %pred = icmp eq i32 %phi, 8
-  // We can't tell from the lattice value for '%phi' that '%pred' is false
-  // along each path, but by checking the predicate over each input separately,
-  // we can.
-  // We limit the search to one step backwards from the current BB and value.
-  // We could consider extending this to search further backwards through the
-  // CFG and/or value graph, but there are non-obvious compile time vs quality
-  // tradeoffs.
+// Tries to determine the value of the given predicate by walking the CFG
+// backwards.
+//
+// The approach here is somewhat distinct from the rest of LVI.  LVI as a whole
+// tries to compute a lattice value which is conservatively correct at a given
+// location, but  in this case, we have a predicate which we weren't able to
+// prove about the merged result, and we're pushing that predicate back along
+// each incoming edge to see if we can prove it separately for each input.
+//
+// As a motivating example, consider:
+//
+//   bb1:
+//     %v1 = ... ; constantrange<1, 5>
+//     br label %merge
+//   bb2:
+//     %v2 = ... ; constantrange<10, 20>
+//     br label %merge
+//   merge:
+//     %phi = phi [%v1, %v2] ; constantrange<1,20>
+//     %pred = icmp eq i32 %phi, 8
+//
+// We can't tell from the lattice value for '%phi' that '%pred' is false along
+// each path, but by checking the predicate over each input separately, we can.
+//
+// We limit the search to one step backwards from the current BB and value.  We
+// could consider extending this to search further backwards through the CFG
+// and/or value graph, but there are non-obvious compile time vs quality
+static LazyValueInfo::Tristate
+getPredicateFromEdges(unsigned Pred, Value *V, Constant *C, BasicBlock *BB,
+                      Instruction *CxtI, LazyValueInfo *LVI) {
+  using Tristate = LazyValueInfo::Tristate;
+  constexpr auto Unknown = LazyValueInfo::Unknown;
+
   if (CxtI) {
     BasicBlock *BB = CxtI->getParent();
 
     // Function entry or an unreachable block.  Bail to avoid confusing
     // analysis below.
     pred_iterator PI = pred_begin(BB), PE = pred_end(BB);
     if (PI == PE)
       return Unknown;
 
     // If V is a PHI node in the same block as the context, we need to ask
     // questions about the predicate as applied to the incoming value along
     // each edge. This is useful for eliminating cases where the predicate is
     // known along all incoming edges.
     if (auto *PHI = dyn_cast<PHINode>(V))
       if (PHI->getParent() == BB) {
         Tristate Baseline = Unknown;
         for (unsigned i = 0, e = PHI->getNumIncomingValues(); i < e; i++) {
           Value *Incoming = PHI->getIncomingValue(i);
           BasicBlock *PredBB = PHI->getIncomingBlock(i);
           // Note that PredBB may be BB itself.
-          Tristate Result = getPredicateOnEdge(Pred, Incoming, C, PredBB, BB,
-                                               CxtI);
+          Tristate Result =
+              LVI->getPredicateOnEdge(Pred, Incoming, C, PredBB, BB, CxtI);
 
           // Keep going as long as we've seen a consistent known result for
           // all inputs.
           Baseline = (i == 0) ? Result /* First iteration */
             : (Baseline == Result ? Baseline : Unknown); /* All others */
           if (Baseline == Unknown)
             break;
         }
         if (Baseline != Unknown)
           return Baseline;
       }
 
     // For a comparison where the V is outside this block, it's possible
     // that we've branched on it before. Look to see if the value is known
     // on all incoming edges.
     if (!isa<Instruction>(V) ||
         cast<Instruction>(V)->getParent() != BB) {
       // For predecessor edge, determine if the comparison is true or false
       // on that edge. If they're all true or all false, we can conclude
       // the value of the comparison in this block.
-      Tristate Baseline = getPredicateOnEdge(Pred, V, C, *PI, BB, CxtI);
+      Tristate Baseline = LVI->getPredicateOnEdge(Pred, V, C, *PI, BB, CxtI);
       if (Baseline != Unknown) {
         // Check that all remaining incoming values match the first one.
         while (++PI != PE) {
-          Tristate Ret = getPredicateOnEdge(Pred, V, C, *PI, BB, CxtI);
+          Tristate Ret = LVI->getPredicateOnEdge(Pred, V, C, *PI, BB, CxtI);
           if (Ret != Baseline) break;
         }
         // If we terminated early, then one of the values didn't match.
         if (PI == PE) {
           return Baseline;
         }
       }
     }
   }
   return Unknown;
 }
 
+LazyValueInfo::Tristate LazyValueInfo::getPredicateAt(unsigned Pred, Value *V,
+                                                      Constant *C,
+                                                      Instruction *CxtI) {
+  const DataLayout &DL = CxtI->getModule()->getDataLayout();
+  Tristate Ret = fastCheckIsOrIsNotNull(Pred, V, C, DL);
+  if (Ret != Unknown)
+    return Ret;
+
+  Ret = getPredicateResult(
+      Pred, C, getImpl(PImpl, AC, &DL, DT).getValueAt(V, CxtI), DL, TLI);
+  if (Ret != Unknown)
+    return Ret;
+
+  return getPredicateFromEdges(Pred, V, C, /*BB=*/nullptr, CxtI, this);
+}
+
+LazyValueInfo::Tristate
+LazyValueInfo::getPredicateInBlock(unsigned Pred, Value *V, Constant *C,
+                                   BasicBlock *BB, Instruction *CxtI) {
+  const DataLayout &DL = CxtI->getModule()->getDataLayout();
+  Tristate Ret = fastCheckIsOrIsNotNull(Pred, V, C, DL);
+  if (Ret != Unknown)
+    return Ret;
+
+  Ret = getPredicateResult(
+      Pred, C, getImpl(PImpl, AC, &DL, DT).getValueInBlock(V, BB, CxtI), DL,
+      TLI);
+  if (Ret != Unknown)
+    return Ret;
+
+  return getPredicateFromEdges(Pred, V, C, /*BB=*/nullptr, CxtI, this);
+}
+
 void LazyValueInfo::threadEdge(BasicBlock *PredBB, BasicBlock *OldSucc,
                                BasicBlock *NewSucc) {
   if (PImpl) {
     const DataLayout &DL = PredBB->getModule()->getDataLayout();
     getImpl(PImpl, AC, &DL, DT).threadEdge(PredBB, OldSucc, NewSucc);
   }
 }
 

llvm/lib/Transforms/Scalar/CorrelatedValuePropagation.cpp

static bool processCmp(CmpInst *C, LazyValueInfo *LVI) {
   // about such cases, it's not its primary purpose.  We do make sure to do
   // the block local query for uses from terminator instructions, but that's
   // handled in the code for each terminator.
   auto *I = dyn_cast<Instruction>(Op0);
   if (I && I->getParent() == C->getParent())
     return false;
 
   LazyValueInfo::Tristate Result =
-    LVI->getPredicateAt(C->getPredicate(), Op0, Op1, C);
+      LVI->getPredicateInBlock(C->getPredicate(), Op0, Op1, C->getParent(), C);
   if (Result == LazyValueInfo::Unknown) return false;
 
   ++NumCmps;
   if (Result == LazyValueInfo::True)
     C->replaceAllUsesWith(ConstantInt::getTrue(C->getContext()));
   else
     C->replaceAllUsesWith(ConstantInt::getFalse(C->getContext()));
   C->eraseFromParent();
static bool processCallSite(CallSite CS, LazyValueInfo *LVI) {
   }
 
   for (Value *V : CS.args()) {
     PointerType *Type = dyn_cast<PointerType>(V->getType());
     // Try to mark pointer typed parameters as non-null.  We skip the
     // relatively expensive analysis for constants which are obviously either
     // null or non-null to start with.
     if (Type && !CS.paramHasAttr(ArgNo, Attribute::NonNull) &&
-        !isa<Constant>(V) && 
-        LVI->getPredicateAt(ICmpInst::ICMP_EQ, V,
-                            ConstantPointerNull::get(Type),
+        !isa<Constant>(V) &&
+        LVI->getPredicateInBlock(ICmpInst::ICMP_EQ, V,
+                                 ConstantPointerNull::get(Type), CS.getParent(),
                                  CS.getInstruction()) == LazyValueInfo::False)
       ArgNos.push_back(ArgNo);
     ArgNo++;
   }
 
   assert(ArgNo == CS.arg_size() && "sanity check");
 
   if (ArgNos.empty())
     return false;
 
   AttributeList AS = CS.getAttributes();
   LLVMContext &Ctx = CS.getInstruction()->getContext();
   AS = AS.addParamAttribute(Ctx, ArgNos,
                             Attribute::get(Ctx, Attribute::NonNull));
   CS.setAttributes(AS);
 
   return true;
 }
 
 static bool hasPositiveOperands(BinaryOperator *SDI, LazyValueInfo *LVI) {
   Constant *Zero = ConstantInt::get(SDI->getType(), 0);
   for (Value *O : SDI->operands()) {
-    auto Result = LVI->getPredicateAt(ICmpInst::ICMP_SGE, O, Zero, SDI);
+    auto Result = LVI->getPredicateInBlock(ICmpInst::ICMP_SGE, O, Zero,
+                                           SDI->getParent(), SDI);
     if (Result != LazyValueInfo::True)
       return false;
   }
   return true;
 }
 
 /// Try to shrink a udiv/urem's width down to the smallest power of two that's
 /// sufficient to contain its operands.
static bool processSDiv(BinaryOperator *SDI, LazyValueInfo *LVI) {
   return true;
 }
 
 static bool processAShr(BinaryOperator *SDI, LazyValueInfo *LVI) {
   if (SDI->getType()->isVectorTy())
     return false;
 
   Constant *Zero = ConstantInt::get(SDI->getType(), 0);
-  if (LVI->getPredicateAt(ICmpInst::ICMP_SGE, SDI->getOperand(0), Zero, SDI) !=
-      LazyValueInfo::True)
+  if (LVI->getPredicateInBlock(ICmpInst::ICMP_SGE, SDI->getOperand(0), Zero,
+                               SDI->getParent(), SDI) != LazyValueInfo::True)
     return false;
 
   ++NumAShrs;
   auto *BO = BinaryOperator::CreateLShr(SDI->getOperand(0), SDI->getOperand(1),
                                         SDI->getName(), SDI);
   BO->setIsExact(SDI->isExact());
   SDI->replaceAllUsesWith(BO);
   SDI->eraseFromParent();
static Constant *getConstantAt(Value *V, Instruction *At, LazyValueInfo *LVI) {
   // TODO: The following really should be sunk inside LVI's core algorithm, or
   // at least the outer shims around such.
   auto *C = dyn_cast<CmpInst>(V);
   if (!C) return nullptr;
 
   Value *Op0 = C->getOperand(0);
   Constant *Op1 = dyn_cast<Constant>(C->getOperand(1));
   if (!Op1) return nullptr;
-  
-  LazyValueInfo::Tristate Result =
-    LVI->getPredicateAt(C->getPredicate(), Op0, Op1, At);
+  LazyValueInfo::Tristate Result = LVI->getPredicateInBlock(
+      C->getPredicate(), Op0, Op1, At->getParent(), At);
   if (Result == LazyValueInfo::Unknown)
     return nullptr;
   
   return (Result == LazyValueInfo::True) ?
     ConstantInt::getTrue(C->getContext()) :
     ConstantInt::getFalse(C->getContext());
 }
 

llvm/test/Transforms/CorrelatedValuePropagation/ashr.ll

; CHECK: lshr i32 %a, 1
   call void @llvm.assume(i1 %cond)
   %shr = ashr i32 %a, 1
   %loopcond = icmp sgt i32 %shr, 8
   br i1 %loopcond, label %loop, label %exit
 
 exit:
   ret void
 }
+
+; CHECK-LABEL: @mask
+define void @mask(i8 %n) {
+  %m = and i8 %n, 31
+; CHECK: lshr i8 %m, 1
+  %shr = ashr i8 %m, 1
+  ret void
+}

llvm/test/Transforms/CorrelatedValuePropagation/icmp.ll

merge:
   unreachable
 
 bb_false:
   %tmp48 = icmp sle i64 %tmp35, 0
   tail call void @check2(i1 %tmp48) #4
   unreachable
 }
 
+; CHECK-LABEL: @mask
+define i1 @mask(i8 %val) {
+  %v = and i8 %val, 31
+  %pred = icmp sle i8 %v, 31
+; CHECK: ret i1 true
+  ret i1 %pred
+}
+
 attributes #4 = { noreturn }

llvm/test/Transforms/CorrelatedValuePropagation/non-null.ll

null:
   br label %merge
 
 merge:
   %merged_arg = phi i8* [%another_arg, %null], [%arg1, %non_null]
   call void @test12_helper(i8* %merged_arg)
   ; CHECK: call void @test12_helper(i8* nonnull %merged_arg)
   ret void
 }
+
+declare void @test_bittwiddle_helper(i8* %arg)
+define void @test_bittwiddle(i8* %arg) {
+; CHECK-LABEL: @test_bittwiddle
+  %raw_ptr = ptrtoint i8* %arg to i64
+  %twiddled = or i64 %raw_ptr, 1
+  %twiddled_ptr = inttoptr i64 %twiddled to i8*
+; CHECK: call void @test_bittwiddle_helper(i8* nonnull %twiddled_ptr)
+  call void @test_bittwiddle_helper(i8* %twiddled_ptr)
+  ret void
+}

llvm/test/Transforms/CorrelatedValuePropagation/sdiv.ll

; CHECK: udiv i32 %a, 6
   call void @llvm.assume(i1 %cond)
   %div = sdiv i32 %a, 6
   %loopcond = icmp sgt i32 %div, 8
   br i1 %loopcond, label %loop, label %exit
 
 exit:
   ret void
 }
+
+; CHECK-LABEL: @mask
+define void @mask(i8 %n) {
+  %m = and i8 %n, 31
+; CHECK: udiv i8 %m, 42
+  %div = sdiv i8 %m, 42
+  ret void
+}

llvm/test/Transforms/CorrelatedValuePropagation/srem.ll

; CHECK: urem i32 %a, 6
   call void @llvm.assume(i1 %cond)
   %rem = srem i32 %a, 6
   %loopcond = icmp sgt i32 %rem, 8
   br i1 %loopcond, label %loop, label %exit
 
 exit:
   ret void
 }
+
+; CHECK-LABEL: @mask
+define void @mask(i8 %n) {
+  %m = and i8 %n, 31
+; CHECK: urem i8 %m, 42
+  %rem = srem i8 %m, 42
+  ret void
+}