[ValueTracking] isKnownNonZero, computeKnownBits for freeze

[aqjune]

Summary

This implements support for isKnownNonZero, computeKnownBits when freeze is involved.

  br (x != 0), BB1, BB2
BB1:
  y = freeze x

In the above program, we can say that y is non-zero. The reason is as follows:

(1) If x was poison, br (x != 0) raised UB
(2) If x was fully undef, the branch again raised UB
(3) If x was non-zero partially undef, say undef | 1, freeze x will return a nondeterministic value which is also non-zero.
(4) If x was just a concrete value, it is trivial

[aqjune]

Consider alloca/undef too, minor update

[jdoerfert]

llvm/include/llvm/Analysis/ValueTracking.h

580 ↗

(On Diff #248936)

We should consider having 3 API functions now. isUndef, isPoison, isUndefOrPoison, with a common impl that takes two booleans. It makes it clearer at the call sites (IMHO).

[reames]

llvm/include/llvm/Analysis/ValueTracking.h

580 ↗

(On Diff #248936)

I agree, this is a very confusing API.

llvm/lib/Analysis/ValueTracking.cpp

1875

I don't think this is worth adding. Such a freeze is trivial, and I'd expect the optimizer to simply remove it instead.

4781–4805

Please test these additions through an existing user, separate and commit separately please.

4866

Same here.

4873

We've ended up with potentially very deep recursion here. In general, ValueTracking works with a small fixed depth window. I think we need to add the same here. (Should be done as a separate patch.)

Note: I don't think we can have a cycle in the graph since we only allow constant arguments to phis.

[aqjune]

llvm/lib/Analysis/ValueTracking.cpp

4873

Made a separate patch here: https://reviews.llvm.org/D76010

[aqjune]

Leave value analysis-related parts only

[aqjune]

Update

[fpichet]

What is the status of this patch?
I've seen cases in a out of tree target where an optimization dependent on ValueTracking was not performed because there was a freeze in the way.

[aqjune]

In D75808#2259962, @fpichet wrote:

What is the status of this patch?
I've seen cases in a out of tree target where an optimization dependent on ValueTracking was not performed because there was a freeze in the way.

Hi, this patch has a dependency on D84242; I'll rebase the patch.

llvm/lib/Analysis/ValueTracking.cpp

1875

KnownBits analysis on freeze can be done when the operand is partially undef. (isGuaranteedNotToBePoison vs. isGuaranteedNotToBeUndefOrPoison)

[nikic]

llvm/lib/Analysis/ValueTracking.cpp

1875

And we can do that because "computeKnownBits" provides a "KnownBits or poison" result and not a "KnownBits or undef/poison" result, right?

[aqjune]

llvm/lib/Analysis/ValueTracking.cpp

1875

Yes, that's true. Take this example:

y = or i8 x, 1
y.fr = freeze y

The current implementation of computeKnownBits(y) will answer that the LSB of y is one, and this means that y is either poison or any value. y can be poison because or poison, 1 is poison. However, y cannot be full undef because or x, 1 cannot return full undef.

[fpichet]

Another thing: I noticed that SelectionDAG::computeKnownBits doesn't deal with ISD::FREEZE.
Would it be ok to just call computeKnownBits (Op.getOperand(0)) to make it work.
(Since Select_FREEZE just lower to a copy)

[aqjune]

In D75808#2260866, @fpichet wrote:

Another thing: I noticed that SelectionDAG::computeKnownBits doesn't deal with ISD::FREEZE.
Would it be ok to just call computeKnownBits (Op.getOperand(0)) to make it work.
(Since Select_FREEZE just lower to a copy)

It depends on whether SelDag has a poison value - does anyone know whether it has one?

[spatel]

In D75808#2260906, @aqjune wrote:

In D75808#2260866, @fpichet wrote:

Another thing: I noticed that SelectionDAG::computeKnownBits doesn't deal with ISD::FREEZE.
Would it be ok to just call computeKnownBits (Op.getOperand(0)) to make it work.
(Since Select_FREEZE just lower to a copy)

It depends on whether SelDag has a poison value - does anyone know whether it has one?

We have nsw/nuw/exact flags on nodes, so that means poison exists?

[aqjune]

Rebase

[aqjune]

In D75808#2261305, @spatel wrote:

In D75808#2260906, @aqjune wrote:

It depends on whether SelDag has a poison value - does anyone know whether it has one?

We have nsw/nuw/exact flags on nodes, so that means poison exists?

I think so. Whether calling computeKnownBits (Op.getOperand(0)) is valid in SelDag seems non trivial then.

[nikic]

LGTM. This does seem useful as at least assumes and dominating conditions always imply non-poison.

llvm/lib/Analysis/ValueTracking.cpp

1876

Should be Depth + 1 here?

2586

Please fix this tidy warning.

[aqjune]

Address comments

llvm/lib/Analysis/ValueTracking.cpp

if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(I->getOperand(0))) {
         case Intrinsic::smul_with_overflow:
           computeKnownBitsMul(II->getArgOperand(0), II->getArgOperand(1), false,
                               DemandedElts, Known, Known2, Depth, Q);
           break;
         }
       }
     }
     break;
+  case Instruction::Freeze:
+    if (isGuaranteedNotToBePoison(I->getOperand(0), Q.CxtI, Q.DT, Depth + 1))
+      computeKnownBits(I->getOperand(0), Known, Depth + 1, Q);
+    break;
   }
 }
 
 /// Determine which bits of V are known to be either zero or one and return
 /// them.
 KnownBits computeKnownBits(const Value *V, const APInt &DemandedElts,
                            unsigned Depth, const Query &Q) {
   KnownBits Known(getBitWidth(V->getType(), Q.DL));
else if (const auto *EEI = dyn_cast<ExtractElementInst>(V)) {
     if (auto *VecTy = dyn_cast<FixedVectorType>(Vec->getType())) {
       unsigned NumElts = VecTy->getNumElements();
       APInt DemandedVecElts = APInt::getAllOnesValue(NumElts);
       if (CIdx && CIdx->getValue().ult(NumElts))
         DemandedVecElts = APInt::getOneBitSet(NumElts, CIdx->getZExtValue());
       return isKnownNonZero(Vec, DemandedVecElts, Depth, Q);
     }
   }
+  // Freeze
+  else if (const FreezeInst *FI = dyn_cast<FreezeInst>(V)) {
+    auto *Op = FI->getOperand(0);
+    if (isKnownNonZero(Op, Depth, Q) &&
+        isGuaranteedNotToBePoison(Op, Q.CxtI, Q.DT, Depth))
+      return true;
+  }
 
   KnownBits Known(BitWidth);
   computeKnownBits(V, DemandedElts, Known, Depth, Q);
   return Known.One != 0;
 }
 
 bool isKnownNonZero(const Value* V, unsigned Depth, const Query& Q) {
   // FIXME: We currently have no way to represent the DemandedElts of a scalable
if (const auto *ExactOp = dyn_cast<PossiblyExactOperator>(Op))
     if (ExactOp->isExact())
       return true;
   if (const auto *FP = dyn_cast<FPMathOperator>(Op)) {
     auto FMF = FP->getFastMathFlags();
     if (FMF.noNaNs() || FMF.noInfs())
       return true;
   }
 
   unsigned Opcode = Op->getOpcode();
 
   // Check whether opcode is a poison/undef-generating operation
   switch (Opcode) {
   case Instruction::Shl:
   case Instruction::AShr:
   case Instruction::LShr: {
     // Shifts return poison if shiftwidth is larger than the bitwidth.
     if (auto *C = dyn_cast<Constant>(Op->getOperand(1))) {
       SmallVector<Constant *, 4> ShiftAmounts;
       if (auto *FVTy = dyn_cast<FixedVectorType>(C->getType())) {
         unsigned NumElts = FVTy->getNumElements();
         for (unsigned i = 0; i < NumElts; ++i)
           ShiftAmounts.push_back(C->getAggregateElement(i));
       } else if (isa<ScalableVectorType>(C->getType()))
         return true; // Can't tell, just return true to be safe
       else
         ShiftAmounts.push_back(C);
 
       bool Safe = llvm::all_of(ShiftAmounts, [](Constant *C) {
         auto *CI = dyn_cast<ConstantInt>(C);
         return CI && CI->getZExtValue() < C->getType()->getIntegerBitWidth();
       });
       return !Safe;
     }
     return true;
   }
   case Instruction::FPToSI:
   case Instruction::FPToUI:
     // fptosi/ui yields poison if the resulting value does not fit in the
     // destination type.
     return true;
   case Instruction::Call:
default: {
     else if (Instruction::isBinaryOp(Opcode))
       return false;
     // Be conservative and return true.
     return true;
   }
   }
 }
 
 bool llvm::canCreateUndefOrPoison(const Operator *Op) {
   return ::canCreateUndefOrPoison(Op, /*PoisonOnly=*/false);
 }
 
 bool llvm::canCreatePoison(const Operator *Op) {
   return ::canCreateUndefOrPoison(Op, /*PoisonOnly=*/true);
 }
 
 static bool programUndefinedIfUndefOrPoison(const Instruction *Inst,
                                             bool PoisonOnly);
 
 static bool isGuaranteedNotToBeUndefOrPoison(const Value *V,
                                              const Instruction *CtxI,
                                              const DominatorTree *DT,
                                              unsigned Depth, bool PoisonOnly) {
   if (Depth >= MaxAnalysisRecursionDepth)

llvm/test/Transforms/InstSimplify/known-non-zero.ll

for.body:                                         ; preds = %for.cond
   %conv = zext i8 %0 to i32
   %mul = shl nuw nsw i32 %shift.0, 3
   %shl = shl nuw nsw i32 %conv, %mul
   %1 = trunc i32 %shl to i8
   %conv2 = or i8 %result.0, %1
   %inc = add nuw nsw i32 %shift.0, 1
   br label %for.cond
 }
+
+define i1 @freeze_nonzero(i8 %x, i8 %mask) {
+; CHECK-LABEL: @freeze_nonzero(
+; CHECK-NEXT:    [[Y:%.*]] = or i8 [[X:%.*]], [[MASK:%.*]]
+; CHECK-NEXT:    [[C:%.*]] = icmp ne i8 [[Y]], 0
+; CHECK-NEXT:    br i1 [[C]], label [[A:%.*]], label [[B:%.*]]
+; CHECK:       A:
+; CHECK-NEXT:    ret i1 false
+; CHECK:       B:
+; CHECK-NEXT:    ret i1 false
+;
+  %y = or i8 %x, %mask
+  %c = icmp ne i8 %y, 0
+  br i1 %c, label %A, label %B
+A:
+  %fr = freeze i8 %y
+  %c2 = icmp eq i8 %fr, 0
+  ret i1 %c2
+B:
+  ret i1 0
+}

llvm/unittests/Analysis/ValueTrackingTest.cpp

parseAssembly(
       "declare void @llvm.assume(i1)\n");
   AssumptionCache AC(*F);
   KnownBits Known = computeKnownBits(
       A, M->getDataLayout(), /* Depth */ 0, &AC, F->front().getTerminator());
   EXPECT_EQ(Known.Zero.getZExtValue(), 31u);
   EXPECT_EQ(Known.One.getZExtValue(), 0u);
 }
 
+TEST_F(ComputeKnownBitsTest, ComputeKnownBitsFreeze) {
+  parseAssembly("define void @test() {\n"
+                "  %m = call i32 @any_num()\n"
+                "  %A = freeze i32 %m\n"
+                "  %n = and i32 %m, 31\n"
+                "  %c = icmp eq i32 %n, 0\n"
+                "  call void @llvm.assume(i1 %c)\n"
+                "  ret void\n"
+                "}\n"
+                "declare void @llvm.assume(i1)\n"
+                "declare i32 @any_num()\n");
+  AssumptionCache AC(*F);
+  KnownBits Known = computeKnownBits(A, M->getDataLayout(), /* Depth */ 0, &AC,
+                                     F->front().getTerminator());
+  EXPECT_EQ(Known.Zero.getZExtValue(), 31u);
+  EXPECT_EQ(Known.One.getZExtValue(), 0u);
+}
+
 class IsBytewiseValueTest : public ValueTrackingTest,
                             public ::testing::WithParamInterface<
                                 std::pair<const char *, const char *>> {
 protected:
 };
 
 const std::pair<const char *, const char *> IsBytewiseValueTests[] = {
     {