[SCCP] Handle assume predicates

[nikic]

Summary

Take assume predicates into account when visiting ssa.copy. The handling is the same as for branch predicates, with the difference that we're always on the true edge.

[nikic]

llvm/lib/Transforms/Scalar/SCCP.cpp

1326

We could set MayIncludeUndef=false for assumes (as undef/poison has always been UB there), but I didn't think it worthwhile to make the distinction, as we plan to flip this for branches in the future anyway.

[fhahn]

LGTM, thanks!

llvm/lib/Transforms/Scalar/SCCP.cpp

1326

It's probably not worth the trouble and we should just prioritize flipping the flag in general.

llvm/lib/Transforms/Scalar/SCCP.cpp

void SCCPSolver::handleCallResult(CallBase &CB) {
   Function *F = CB.getCalledFunction();
 
   if (auto *II = dyn_cast<IntrinsicInst>(&CB)) {
     if (II->getIntrinsicID() == Intrinsic::ssa_copy) {
       if (ValueState[&CB].isOverdefined())
         return;
 
       Value *CopyOf = CB.getOperand(0);
-      auto *PI = getPredicateInfoFor(&CB);
-      auto *PBranch = dyn_cast_or_null<PredicateBranch>(PI);
       ValueLatticeElement OriginalVal = getValueState(CopyOf);
-      if (!PI || !PBranch) {
+      auto *PI = getPredicateInfoFor(&CB);
+      assert(PI && "Missing predicate info for ssa.copy");
+
+      CmpInst *Cmp;
+      bool TrueEdge;
+      if (auto *PBranch = dyn_cast<PredicateBranch>(PI)) {
+        Cmp = dyn_cast<CmpInst>(PBranch->Condition);
+        TrueEdge = PBranch->TrueEdge;
+      } else if (auto *PAssume = dyn_cast<PredicateAssume>(PI)) {
+        Cmp = dyn_cast<CmpInst>(PAssume->Condition);
+        TrueEdge = true;
+      } else {
         mergeInValue(ValueState[&CB], &CB, OriginalVal);
         return;
       }
 
       // Everything below relies on the condition being a comparison.
-      auto *Cmp = dyn_cast<CmpInst>(PBranch->Condition);
       if (!Cmp) {
         mergeInValue(ValueState[&CB], &CB, OriginalVal);
         return;
       }
 
       Value *CmpOp0 = Cmp->getOperand(0);
       Value *CmpOp1 = Cmp->getOperand(1);
       if (CopyOf != CmpOp0 && CopyOf != CmpOp1) {
         mergeInValue(ValueState[&CB], &CB, OriginalVal);
         return;
       }
 
       auto Pred = Cmp->getPredicate();
       if (CmpOp0 != CopyOf) {
         std::swap(CmpOp0, CmpOp1);
         Pred = Cmp->getSwappedPredicate();
       }
 
       // Wait until CmpOp1 is resolved.
       if (getValueState(CmpOp1).isUnknown()) {
         addAdditionalUser(CmpOp1, &CB);
         return;
       }
 
-      if (!PBranch->TrueEdge)
+      if (!TrueEdge)
         Pred = CmpInst::getInversePredicate(Pred);
 
       ValueLatticeElement CondVal = getValueState(CmpOp1);
       ValueLatticeElement &IV = ValueState[&CB];
       if (CondVal.isConstantRange() || OriginalVal.isConstantRange()) {
         auto NewCR =
             ConstantRange::getFull(DL.getTypeSizeInBits(CopyOf->getType()));
 
if (II->getIntrinsicID() == Intrinsic::ssa_copy) {
         // undef is as follows below:
         // The new range is based on a branch condition. That guarantees that
         // neither of the compare operands can be undef in the branch targets,
         // unless we have conditions that are always true/false (e.g. icmp ule
         // i32, %a, i32_max). For the latter overdefined/empty range will be
         // inferred, but the branch will get folded accordingly anyways.
         mergeInValue(
             IV, &CB,
             ValueLatticeElement::getRange(NewCR, /*MayIncludeUndef=*/true));
         return;
       } else if (Pred == CmpInst::ICMP_EQ && CondVal.isConstant()) {
         // For non-integer values or integer constant expressions, only
         // propagate equal constants.
         addAdditionalUser(CmpOp1, &CB);
         mergeInValue(IV, &CB, CondVal);
         return;
       }

llvm/test/Transforms/SCCP/assume.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -ipsccp -S | FileCheck %s
+
+declare void @use(i1)
+declare void @llvm.assume(i1)
+
+define void @basic(i32 %v) {
+; CHECK-LABEL: @basic(
+; CHECK-NEXT:    [[A1:%.*]] = icmp ult i32 [[V:%.*]], 10
+; CHECK-NEXT:    call void @llvm.assume(i1 [[A1]])
+; CHECK-NEXT:    [[A2:%.*]] = icmp ugt i32 [[V]], 5
+; CHECK-NEXT:    call void @llvm.assume(i1 [[A2]])
+; CHECK-NEXT:    call void @use(i1 true)
+; CHECK-NEXT:    [[C2:%.*]] = icmp ult i32 [[V]], 9
+; CHECK-NEXT:    call void @use(i1 [[C2]])
+; CHECK-NEXT:    call void @use(i1 false)
+; CHECK-NEXT:    [[C4:%.*]] = icmp ugt i32 [[V]], 8
+; CHECK-NEXT:    call void @use(i1 [[C4]])
+; CHECK-NEXT:    call void @use(i1 true)
+; CHECK-NEXT:    [[C6:%.*]] = icmp ugt i32 [[V]], 6
+; CHECK-NEXT:    call void @use(i1 [[C6]])
+; CHECK-NEXT:    call void @use(i1 false)
+; CHECK-NEXT:    [[C8:%.*]] = icmp ult i32 [[V]], 7
+; CHECK-NEXT:    call void @use(i1 [[C8]])
+; CHECK-NEXT:    ret void
+;
+  %a1 = icmp ult i32 %v, 10
+  call void @llvm.assume(i1 %a1)
+  %a2 = icmp ugt i32 %v, 5
+  call void @llvm.assume(i1 %a2)
+  %c1 = icmp ult i32 %v, 10
+  call void @use(i1 %c1)
+  %c2 = icmp ult i32 %v, 9
+  call void @use(i1 %c2)
+  %c3 = icmp ugt i32 %v, 9
+  call void @use(i1 %c3)
+  %c4 = icmp ugt i32 %v, 8
+  call void @use(i1 %c4)
+  %c5 = icmp ugt i32 %v, 5
+  call void @use(i1 %c5)
+  %c6 = icmp ugt i32 %v, 6
+  call void @use(i1 %c6)
+  %c7 = icmp ult i32 %v, 6
+  call void @use(i1 %c7)
+  %c8 = icmp ult i32 %v, 7
+  call void @use(i1 %c8)
+  ret void
+}