[ValueTracking] Add support to deduce a PHI node being a power of 2 if each incoming value is a power of 2.

[spatel]

llvm/lib/Analysis/ValueTracking.cpp

2196

As mentioned in the earlier patch (D123748), the Depth should be limited here to prevent a compile-time explosion. If the phi has a large number of operands, this could consume a lot of time - O(Operands^Depth).

The code that this patch is copying from was added with D88276, and there is more discussion about it there.
Note that the limit was derived from this code, but it dropped the explanatory code comment:
https://github.com/llvm/llvm-project/blob/d2166076b882e38becf3657ea830ffd2b6a5695e/llvm/lib/Analysis/ValueTracking.cpp#L1574-L1606

I'm not sure why it is using std::max instead of just passing MaxAnalysisRecursionDepth - 1. If I'm seeing the tests in this patch correctly, we will not see a functional difference on either of those tests by limiting the recursion.

The other non-obvious part of this patch is the setting of the context instruction; that is copied/derived from:
D71181

All of this copying and mutation suggests that we really need to create a helper function, so the logic is consistent for phi analysis across these functions. If that isn't feasible, at least copy the code comments, so we don't have to dig all over the place to figure out why this code does what it does. :)

[huangjd]

llvm/lib/Analysis/ValueTracking.cpp

2196

Passing MaxAnalysisRecursionDepth - 1 leads to incorrect behavior, as it defeats the purpose of limiting the depth of recursion. Every time the code execution calls the function on this line the same depth Passing MaxAnalysisRecursionDepth - 1 is passed in, so line 2065 will never be true, and the function never terminates due to depth exceeding limit.

[huangjd]

Added back recursion limit
Added comment on instruction context

[spatel]

Why was D124890 pulled back into this patch? I still think it would be better to have 2 patches to reduce risk.

[huangjd]

Added back D124890 (now D125332)

[davidxl]

This one looks fine to me. Wait to see if other reviewers have further comment.

[huangjd]

Updated parent D125332

[huangjd]

Restore D124889 for review

[davidxl]

looks like all comments have been addressed. Wait a few days to commit. (also abandon the duplicate https://reviews.llvm.org/D126020)

[nikic]

llvm/lib/Analysis/ValueTracking.cpp

2193

Rather than requiring Q.IIQ.UseInstrInfo here, it would be better to use Q.IIQ.hasNoUnsignedWrap() etc in the function. Presence of nowrap flags is not required for all cases.

[huangjd]

This comment was removed by huangjd.

[huangjd]

[ValueTracking] Relaxed requirement of UseInstrInfo in power of 2
recurrece check

[nikic]

llvm/lib/Analysis/ValueTracking.cpp

2057

Hm, I think that for recurrences, this context instruction adjustment is not necessary. I believe the reason why we need it in general for phis is that we might follow a backedge, in which case we will be working with values from a previous loop iterations. This can't happen when we're looking at Start, because it must dominate the whole loop.

2089

Similar to the nowrap cases above, can't this (and div) use OrZero || Q.IIQ.isExact(BO)?

[huangjd]

llvm/lib/Analysis/ValueTracking.cpp

2057

I saw similar usage in computeKnownBitsFromOperator. at line 1464. The comment said

... TODO: It may be correct to use the original context.  IF warranted, explore and
 add sufficient tests to cover.

I am worry about some unusual cases such as the "Start" value is actually coming from a BB dominated by PHI's BB, and matchSimpleRecurrence is not strict on checking value domination, so I would be cautious here

[huangjd]

Add isExact check for power of 2 in recursion

[huangjd]

Updated by merging main
Just to confirm it still passes all tests

llvm/lib/Analysis/ValueTracking.cpp

#endif
   // Therefore, we run them after computeKnownBitsFromOperator.
 
   // Check whether a nearby assume intrinsic can determine some known bits.
   computeKnownBitsFromAssume(V, Known, Depth, Q);
 
   assert((Known.Zero & Known.One) == 0 && "Bits known to be one AND zero?");
 }
 
+/// Try to detect a recurrence that the value of the induction variable is
+/// always a power of two (or zero).
+static bool isPowerOfTwoRecurrence(const PHINode *PN, bool OrZero,
+                                   unsigned Depth, Query &Q) {
+  BinaryOperator *BO = nullptr;
+  Value *Start = nullptr, *Step = nullptr;
+  if (!matchSimpleRecurrence(PN, BO, Start, Step))
+    return false;
+
+  // Initial value must be a power of two.
+  for (const Use &U : PN->operands()) {
+    if (U.get() == Start) {
+      // Initial value comes from a different BB, need to adjust context
+      // instruction for analysis.
+      Q.CxtI = PN->getIncomingBlock(U)->getTerminator();
+      if (!isKnownToBeAPowerOfTwo(Start, OrZero, Depth, Q))
+        return false;
+    }
+  }
+
+  // Except for Mul, the induction variable must be on the left side of the
+  // increment expression, otherwise its value can be arbitrary.
+  if (BO->getOpcode() != Instruction::Mul && BO->getOperand(1) != Step)
+    return false;
+
+  Q.CxtI = BO->getParent()->getTerminator();
+  switch (BO->getOpcode()) {
+  case Instruction::Mul:
+    // Power of two is closed under multiplication.
+    return (OrZero || Q.IIQ.hasNoUnsignedWrap(BO) ||
+            Q.IIQ.hasNoSignedWrap(BO)) &&
+           isKnownToBeAPowerOfTwo(Step, OrZero, Depth, Q);
+  case Instruction::SDiv:
+    // Start value must not be signmask for signed division, so simply being a
+    // power of two is not sufficient, and it has to be a constant.
+    if (!match(Start, m_Power2()) || match(Start, m_SignMask()))
+      return false;
+    LLVM_FALLTHROUGH;
+  case Instruction::UDiv:
+    // Divisor must be a power of two.
+    // If OrZero is false, cannot guarantee induction variable is non-zero after
+    // division, same for Shr, unless it is exact division.
+    return (OrZero || Q.IIQ.isExact(BO)) &&
+           isKnownToBeAPowerOfTwo(Step, false, Depth, Q);
+  case Instruction::Shl:
+    return OrZero || Q.IIQ.hasNoUnsignedWrap(BO) || Q.IIQ.hasNoSignedWrap(BO);
+  case Instruction::AShr:
+    if (!match(Start, m_Power2()) || match(Start, m_SignMask()))
+      return false;
+    LLVM_FALLTHROUGH;
+  case Instruction::LShr:
+    return OrZero || Q.IIQ.isExact(BO);
+  default:
+    return false;
+  }
+}
+
 /// Return true if the given value is known to have exactly one
 /// bit set when defined. For vectors return true if every element is known to
 /// be a power of two when defined. Supports values with integer or pointer
 /// types and vectors of integers.
 bool isKnownToBeAPowerOfTwo(const Value *V, bool OrZero, unsigned Depth,
                             const Query &Q) {
   assert(Depth <= MaxAnalysisRecursionDepth && "Limit Search Depth");
 
if (OrZero || Q.IIQ.hasNoUnsignedWrap(VOBO) ||
       if ((~(LHSBits.Zero & RHSBits.Zero)).isPowerOf2())
         // If OrZero isn't set, we cannot give back a zero result.
         // Make sure either the LHS or RHS has a bit set.
         if (OrZero || RHSBits.One.getBoolValue() || LHSBits.One.getBoolValue())
           return true;
     }
   }
 
-  // A PHI node is power of two if all incoming values are power of two.
+  // A PHI node is power of two if all incoming values are power of two, or if
+  // it is an induction variable where in each step its value is a power of two.
   if (const PHINode *PN = dyn_cast<PHINode>(V)) {
     Query RecQ = Q;
 
+    // Check if it is an induction variable and always power of two.
+    if (isPowerOfTwoRecurrence(PN, OrZero, Depth, RecQ))
+      return true;
+
     // Recursively check all incoming values. Limit recursion to 2 levels, so
     // that search complexity is limited to number of operands^2.
     unsigned NewDepth = std::max(Depth, MaxAnalysisRecursionDepth - 1);
     return llvm::all_of(PN->operands(), [&](const Use &U) {
       // Value is power of 2 if it is coming from PHI node itself by induction.
       if (U.get() == PN)
         return true;
 
       // Change the context instruction to the incoming block where it is

llvm/test/Analysis/ValueTracking/known-power-of-two-urem.ll

 define i64 @known_power_of_two_urem_loop_mul(i64 %size, i64 %a) {
 ; CHECK-LABEL: @known_power_of_two_urem_loop_mul(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[START:%.*]] = shl nuw i64 1, [[A:%.*]]
 ; CHECK-NEXT:    br label [[FOR_BODY:%.*]]
 ; CHECK:       for.body:
 ; CHECK-NEXT:    [[PHI:%.*]] = phi i64 [ [[START]], [[ENTRY:%.*]] ], [ [[I:%.*]], [[FOR_BODY]] ]
 ; CHECK-NEXT:    [[SUM:%.*]] = phi i64 [ 0, [[ENTRY]] ], [ [[ADD:%.*]], [[FOR_BODY]] ]
-; CHECK-NEXT:    [[UREM:%.*]] = urem i64 [[SIZE:%.*]], [[PHI]]
+; CHECK-NEXT:    [[TMP0:%.*]] = add i64 [[PHI]], -1
+; CHECK-NEXT:    [[UREM:%.*]] = and i64 [[TMP0]], [[SIZE:%.*]]
 ; CHECK-NEXT:    [[ADD]] = add nuw i64 [[SUM]], [[UREM]]
 ; CHECK-NEXT:    [[I]] = shl nuw i64 [[PHI]], 2
 ; CHECK-NEXT:    [[ICMP:%.*]] = icmp ult i64 [[PHI]], 25000000
 ; CHECK-NEXT:    br i1 [[ICMP]], label [[FOR_BODY]], label [[FOR_END:%.*]]
 ; CHECK:       for.end:
 ; CHECK-NEXT:    ret i64 [[SUM]]
 ;
 entry:

 define i64 @known_power_of_two_urem_loop_shl(i64 %size, i64 %a) {
 ; CHECK-LABEL: @known_power_of_two_urem_loop_shl(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[START:%.*]] = shl nuw i64 1, [[A:%.*]]
 ; CHECK-NEXT:    br label [[FOR_BODY:%.*]]
 ; CHECK:       for.body:
 ; CHECK-NEXT:    [[PHI:%.*]] = phi i64 [ [[START]], [[ENTRY:%.*]] ], [ [[I:%.*]], [[FOR_BODY]] ]
 ; CHECK-NEXT:    [[SUM:%.*]] = phi i64 [ 0, [[ENTRY]] ], [ [[ADD:%.*]], [[FOR_BODY]] ]
-; CHECK-NEXT:    [[UREM:%.*]] = urem i64 [[SIZE:%.*]], [[PHI]]
+; CHECK-NEXT:    [[TMP0:%.*]] = add i64 [[PHI]], -1
+; CHECK-NEXT:    [[UREM:%.*]] = and i64 [[TMP0]], [[SIZE:%.*]]
 ; CHECK-NEXT:    [[ADD]] = add nuw i64 [[SUM]], [[UREM]]
 ; CHECK-NEXT:    [[I]] = shl nuw i64 [[PHI]], 1
 ; CHECK-NEXT:    [[ICMP:%.*]] = icmp ult i64 [[PHI]], 50000000
 ; CHECK-NEXT:    br i1 [[ICMP]], label [[FOR_BODY]], label [[FOR_END:%.*]]
 ; CHECK:       for.end:
 ; CHECK-NEXT:    ret i64 [[SUM]]
 ;
 entry:

 define i64 @known_power_of_two_urem_loop_lshr(i64 %size, i64 %a) {
 ; CHECK-LABEL: @known_power_of_two_urem_loop_lshr(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[START:%.*]] = shl nuw i64 1, [[A:%.*]]
 ; CHECK-NEXT:    br label [[FOR_BODY:%.*]]
 ; CHECK:       for.body:
 ; CHECK-NEXT:    [[PHI:%.*]] = phi i64 [ [[START]], [[ENTRY:%.*]] ], [ [[I:%.*]], [[FOR_BODY]] ]
 ; CHECK-NEXT:    [[SUM:%.*]] = phi i64 [ 0, [[ENTRY]] ], [ [[ADD:%.*]], [[FOR_BODY]] ]
-; CHECK-NEXT:    [[UREM:%.*]] = urem i64 [[SIZE:%.*]], [[PHI]]
+; CHECK-NEXT:    [[TMP0:%.*]] = add i64 [[PHI]], -1
+; CHECK-NEXT:    [[UREM:%.*]] = and i64 [[TMP0]], [[SIZE:%.*]]
 ; CHECK-NEXT:    [[ADD]] = add nuw i64 [[SUM]], [[UREM]]
 ; CHECK-NEXT:    [[I]] = lshr i64 [[PHI]], 1
 ; CHECK-NEXT:    [[ICMP_NOT:%.*]] = icmp ult i64 [[PHI]], 2
 ; CHECK-NEXT:    br i1 [[ICMP_NOT]], label [[FOR_END:%.*]], label [[FOR_BODY]]
 ; CHECK:       for.end:
 ; CHECK-NEXT:    ret i64 [[SUM]]
 ;
 entry:

 
 define i64 @known_power_of_two_urem_loop_ashr(i64 %size, i64 %a) {
 ; CHECK-LABEL: @known_power_of_two_urem_loop_ashr(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    br label [[FOR_BODY:%.*]]
 ; CHECK:       for.body:
 ; CHECK-NEXT:    [[PHI:%.*]] = phi i64 [ 4096, [[ENTRY:%.*]] ], [ [[I:%.*]], [[FOR_BODY]] ]
 ; CHECK-NEXT:    [[SUM:%.*]] = phi i64 [ 0, [[ENTRY]] ], [ [[ADD:%.*]], [[FOR_BODY]] ]
-; CHECK-NEXT:    [[UREM:%.*]] = urem i64 [[SIZE:%.*]], [[PHI]]
-; CHECK-NEXT:    [[ADD]] = add nuw nsw i64 [[SUM]], [[UREM]]
+; CHECK-NEXT:    [[TMP0:%.*]] = add nsw i64 [[PHI]], -1
+; CHECK-NEXT:    [[UREM:%.*]] = and i64 [[TMP0]], [[SIZE:%.*]]
+; CHECK-NEXT:    [[ADD]] = add nsw i64 [[SUM]], [[UREM]]
 ; CHECK-NEXT:    [[I]] = lshr i64 [[PHI]], [[A:%.*]]
 ; CHECK-NEXT:    [[ICMP_NOT:%.*]] = icmp eq i64 [[I]], 0
 ; CHECK-NEXT:    br i1 [[ICMP_NOT]], label [[FOR_END:%.*]], label [[FOR_BODY]]
 ; CHECK:       for.end:
 ; CHECK-NEXT:    ret i64 [[SUM]]
 ;
 entry:
   br label %for.body