[knownbits] Preserve known bits for small shift recurrences

[reames]

Summary

The motivation for this is that I'm looking at an example that uses shifts as induction variables. There's lots of other omissions, but one of the first I noticed is that we can't compute tight known bits. (This indirectly causes SCEV's range analysis to produce very poor results as well.)

One specific question to reviewers. This patch essentially treats all shifts as if they were not exact/nsw/nuw. This is a valid interpretation, but I'm concerned it might differ from that taken elsewhere in the optimizer in a problematic way. What is the best practice for this in known bits computation?

Compile time wise, this should be reasonable. We locally match the recurrence using simple IR pattern matching, and only recurse on the operand once we found a recurrence.

[nikic]

llvm/lib/Analysis/ValueTracking.cpp

1343

Why do we need to do this both starting from the shifts and starting from the phi? The existing code for add recurrences only starts from phis.

[reames]

llvm/lib/Analysis/ValueTracking.cpp

1343

Great question! Turns out I don't, I just never thought to check after adding the phi logic. Will remove these three call sites.

[reames]

Remove unneeded lines as pointed out by Nikita.

[nikic]

llvm/lib/Analysis/ValueTracking.cpp

1031

The context instruction here needs to be adjusted to the terminator of the block for the incoming value StartV. Context can't be preserved when looking through phis.

1475

The surrounding code here is already handling other types of recurrences of the same form -- combining it with RefineRecurrence is rather odd, because it does it's own recurrence matching that is partially redundant with the code here. I think these should be integrated, or at least factored differently. The scaffolding here should work for the shift case as well (restricted to LL == I, as for Sub below), just with different KnownBits logic being applied.

[aqjune]

One specific question to reviewers. This patch essentially treats all shifts as if they were not exact/nsw/nuw. This is a valid interpretation, but I'm concerned it might differ from that taken elsewhere in the optimizer in a problematic way. What is the best practice for this in known bits computation?

I think this is okay because IIUC analyses in ValueTracking (computeKnownXX/isKnownXX) can answer anything if the value was poison.
For example, isGEPKnownNonNull(I) is returning true if I = gep inbounds X, ofs where ofs > 0.
Users of the analyses should assume that the value can be poison; this is consistent with the recent change in the definition of nonnull/align attribute.

[reames]

llvm/lib/Analysis/ValueTracking.cpp

1031

In this particular case, I believe the current usage is actually correct. The general problem with phis is that you can end up traversing a value which is not control dependent on the conditions controlling the original context. In this particular case, we know that the phi dominates the binop which dominates the original context (or at least it had better), and thus we should be okay to preserve the initial context.

Having said all that, I hadn't intended to be fancy here and don't care about the potential quality gain from preserving the context. I'm going to switch to using the phi one as a defensive programming measure anyways. It's much more obviously correct.

1475

I acknowledge this ends up being slightly odd for this caller. I want to preserve the matching logic structure anyways. Specifically, immediately after this lands, I have several other use cases for recognizing shift recurrences.

If you really want, I can invert and generalize the existing code here and defer the addition of the new matcher to a following patch. I wouldn't bother myself, but don't care strongly enough to really argue against it either.

[reames]

Adjust context instruction.

Another diff updated expected shortly. Going to land tests, and rebase. Also want to try out the code structure suggested by Nikita to see how much cleaner that is. Depending on that, may update that as well.

[reames]

Rebase over landed tests and integrate into single remaining caller as suggested by Nikita.

[nikic]

LGTM

llvm/lib/Analysis/ValueTracking.cpp

1031

Yeah, I think you're right that keeping the context instruction would be correct in this case.

1499

nit: switch (

llvm/test/Analysis/ValueTracking/shift-recurrence-knownbits.ll

64

I'd suggest another negative case where the wrong operand of the shift is used (lshr i64 C, %phi style).

And another positive test where the order of the phi operands is swapped.

llvm/lib/Analysis/ValueTracking.cpp

if (Q.ORE)
                                           CxtI)
                << "Detected conflicting code assumptions. Program may "
                   "have undefined behavior, or compiler may have "
                   "internal error.";
       });
   }
 }
 
+/// Match a small linear recurrence cycle involving shift opcodes.
+/// Specifically, recognize the following:
+///   %iv = phi iN [%start, %?], [%iv.next, %backedge]
+///   %iv.next = shift_op %iv, %step
+///
+/// TODO: This can fairly trivially be generalized for other
+/// binary instructions if that were useful.
+static bool MatchShiftRecurrence(const Operator *O,
+                                 Value *&StartV, Value *&StepV) {
+  PHINode* PN = nullptr;
+  switch (O->getOpcode()) {
+  case Instruction::Shl:
+  case Instruction::AShr:
+  case Instruction::LShr:
+    PN = dyn_cast<PHINode>(O->getOperand(0));
+    StepV = O->getOperand(1);
+    break;
+  default:
+    return false;
+  }
+  if (!PN)
+    return false;
+
+  if (PN->getNumIncomingValues() != 2)
+    return false;
+
+  StartV = nullptr;
+  for (Value *IncV : PN->incoming_values()) {
+    if (IncV == O)
+      continue;
+    if (!StartV) {
+      StartV = IncV;
+      continue;
+    }
+    if (StartV != IncV)
+      return false;
+  }
+
+  if (!StartV)
+    return false;
+
+  return true;
+}
+
+/// If operator is part of a small recurrence cycle, use facts about the
+/// recurrence itself to refine the known bits of the result.
+/// As is idiomatic, Known2 is used as scratch.  Known is refined.
+static void RefineRecurrence(const Operator *O, const APInt &DemandedElts,
+                             KnownBits &Known, KnownBits &Known2,
+                             unsigned Depth, const Query &Q) {
+  unsigned Opcode = O->getOpcode();
+
+  // For the moment, we only handle the shift cases here.  This can
+  // probably be extended to most binary instructions.
+  if (Opcode != Instruction::LShr &&
+      Opcode != Instruction::AShr &&
+      Opcode != Instruction::Shl)
+    return;
+
+  Value *StartV = nullptr, *StepV = nullptr;
+  if (MatchShiftRecurrence(O, StartV, StepV)) {
+    computeKnownBits(StartV, DemandedElts, Known2, Depth + 1, Q);
+    switch(Opcode) {
+    case Instruction::Shl:
+      // A shl recurrence will only increase the tailing zeros
+      Known.Zero.setLowBits(Known2.countMinTrailingZeros());
+      break;
+    case Instruction::LShr:
+      // A lshr recurrence will preserve the leading zeros of the start value
+      Known.Zero.setHighBits(Known2.countMinLeadingZeros());
+      break;
+    case Instruction::AShr:
+      // An ashr recurrence will extend the initial sign bit
+      Known.Zero.setHighBits(Known2.countMinLeadingZeros());
+      Known.One.setHighBits(Known2.countMinLeadingOnes());
+    };
+  }
+}
+
 /// Compute known bits from a shift operator, including those with a
 /// non-constant shift amount. Known is the output of this function. Known2 is a
 /// pre-allocated temporary with the same bit width as Known and on return
 /// contains the known bit of the shift value source. KF is an
 /// operator-specific function that, given the known-bits and a shift amount,
 /// compute the implied known-bits of the shift operator's result respectively
 /// for that shift amount. The results from calling KF are conservatively
 /// combined for all permitted shift amounts.
auto KF = [NSW](const KnownBits &KnownVal, const KnownBits &KnownAmt) {
       return Result;
     };
     computeKnownBitsFromShiftOperator(I, DemandedElts, Known, Known2, Depth, Q,
                                       KF);
     // Trailing zeros of a right-shifted constant never decrease.
     const APInt *C;
     if (match(I->getOperand(0), m_APInt(C)))
       Known.Zero.setLowBits(C->countTrailingZeros());
+
+    // A shl recurrence will only increase the tailing zeros
+    RefineRecurrence(I, DemandedElts, Known, Known2, Depth, Q);
     break;
   }
   case Instruction::LShr: {
     auto KF = [](const KnownBits &KnownVal, const KnownBits &KnownAmt) {
       return KnownBits::lshr(KnownVal, KnownAmt);
     };
     computeKnownBitsFromShiftOperator(I, DemandedElts, Known, Known2, Depth, Q,
                                       KF);
     // Leading zeros of a left-shifted constant never decrease.
     const APInt *C;
     if (match(I->getOperand(0), m_APInt(C)))
       Known.Zero.setHighBits(C->countLeadingZeros());
+
+    // A lshr recurrence will preserve the leading zeros of the start value
+    RefineRecurrence(I, DemandedElts, Known, Known2, Depth, Q);
     break;
   }
   case Instruction::AShr: {
     auto KF = [](const KnownBits &KnownVal, const KnownBits &KnownAmt) {
       return KnownBits::ashr(KnownVal, KnownAmt);
     };
     computeKnownBitsFromShiftOperator(I, DemandedElts, Known, Known2, Depth, Q,
                                       KF);
+
+    // An ashr recurrence will extend the initial sign bit
+    RefineRecurrence(I, DemandedElts, Known, Known2, Depth, Q);
     break;
   }
   case Instruction::Sub: {
     bool NSW = Q.IIQ.hasNoSignedWrap(cast<OverflowingBinaryOperator>(I));
     computeKnownBitsAddSub(false, I->getOperand(0), I->getOperand(1), NSW,
                            DemandedElts, Known, Known2, Depth, Q);
     break;
   }
if (P->getNumIncomingValues() == 2) {
         Value *L = P->getIncomingValue(i);
         Value *R = P->getIncomingValue(!i);
         Instruction *RInst = P->getIncomingBlock(!i)->getTerminator();
         Instruction *LInst = P->getIncomingBlock(i)->getTerminator();
         Operator *LU = dyn_cast<Operator>(L);
         if (!LU)
           continue;
         unsigned Opcode = LU->getOpcode();
+
+        // If this is an small recurrence cycle, use facts about the
+        // recurrence itself to refine the result.
+        RefineRecurrence(LU, DemandedElts, Known, Known2, Depth, Q);
+
         // Check for operations that have the property that if
         // both their operands have low zero bits, the result
         // will have low zero bits.
         if (Opcode == Instruction::Add ||
             Opcode == Instruction::Sub ||
             Opcode == Instruction::And ||
             Opcode == Instruction::Or ||
             Opcode == Instruction::Mul) {
           Value *LL = LU->getOperand(0);
           Value *LR = LU->getOperand(1);
           // Find a recurrence.
           if (LL == I)
             L = LR;
           else if (LR == I)
             L = LL;
           else
             continue; // Check for recurrence with L and R flipped.
 
           // Change the context instruction to the "edge" that flows into the
           // phi. This is important because that is where the value is actually
           // "evaluated" even though it is used later somewhere else. (see also
           // D69571).
           Query RecQ = Q;
 
           // Ok, we have a PHI of the form L op= R. Check for low
           // zero bits.
           RecQ.CxtI = RInst;
           computeKnownBits(R, Known2, Depth + 1, RecQ);
 
           // We need to take the minimum number of known bits
           KnownBits Known3(BitWidth);

llvm/test/Analysis/ValueTracking/shift-recurrence-knownbits.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -instcombine -S | FileCheck %s
+
+define i64 @test_lshr() {
+; CHECK-LABEL: @test_lshr(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    br i1 undef, label [[EXIT:%.*]], label [[LOOP]]
+; CHECK:       exit:
+; CHECK-NEXT:    ret i64 1023
+;
+entry:
+  br label %loop
+loop:
+  %iv.lshr = phi i64 [1023, %entry], [%iv.lshr.next, %loop]
+  %iv.lshr.next = lshr i64 %iv.lshr, 1
+  br i1 undef, label %exit, label %loop
+exit:
+  %res = or i64 %iv.lshr, 1023
+  ret i64 %res
+}
+
+define i64 @test_ashr_zeros() {
+; CHECK-LABEL: @test_ashr_zeros(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    br i1 undef, label [[EXIT:%.*]], label [[LOOP]]
+; CHECK:       exit:
+; CHECK-NEXT:    ret i64 1023
+;
+entry:
+  br label %loop
+loop:
+  %iv.ashr = phi i64 [1023, %entry], [%iv.ashr.next, %loop]
+  %iv.ashr.next = ashr i64 %iv.ashr, 1
+  br i1 undef, label %exit, label %loop
+exit:
+  %res = or i64 %iv.ashr, 1023
+  ret i64 %res
+}
+
+define i64 @test_ashr_ones() {
+; CHECK-LABEL: @test_ashr_ones(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    br i1 undef, label [[EXIT:%.*]], label [[LOOP]]
+; CHECK:       exit:
+; CHECK-NEXT:    ret i64 -1
+;
+entry:
+  br label %loop
+loop:
+  %iv.ashr = phi i64 [-1023, %entry], [%iv.ashr.next, %loop]
+  %iv.ashr.next = ashr i64 %iv.ashr, 1
+  br i1 undef, label %exit, label %loop
+exit:
+  %res = or i64 %iv.ashr, 1023
+  ret i64 %res
+}
+
+; negative case for when start is unknown
+define i64 @test_ashr_unknown(i64 %start) {
+; CHECK-LABEL: @test_ashr_unknown(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    [[IV_ASHR:%.*]] = phi i64 [ [[START:%.*]], [[ENTRY:%.*]] ], [ [[IV_ASHR_NEXT:%.*]], [[LOOP]] ]
+; CHECK-NEXT:    [[IV_ASHR_NEXT]] = ashr i64 [[IV_ASHR]], 1
+; CHECK-NEXT:    br i1 undef, label [[EXIT:%.*]], label [[LOOP]]
+; CHECK:       exit:
+; CHECK-NEXT:    [[RES:%.*]] = or i64 [[IV_ASHR]], 1023
+; CHECK-NEXT:    ret i64 [[RES]]
+;
+entry:
+  br label %loop
+loop:
+  %iv.ashr = phi i64 [%start, %entry], [%iv.ashr.next, %loop]
+  %iv.ashr.next = ashr i64 %iv.ashr, 1
+  br i1 undef, label %exit, label %loop
+exit:
+  %res = or i64 %iv.ashr, 1023
+  ret i64 %res
+}
+
+define i64 @test_shl() {
+; CHECK-LABEL: @test_shl(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    br i1 undef, label [[EXIT:%.*]], label [[LOOP]]
+; CHECK:       exit:
+; CHECK-NEXT:    ret i64 0
+;
+entry:
+  br label %loop
+loop:
+  %iv.shl = phi i64 [8, %entry], [%iv.shl.next, %loop]
+  %iv.shl.next = shl i64 %iv.shl, 1
+  br i1 undef, label %exit, label %loop
+exit:
+  %res = and i64 %iv.shl, 7
+  ret i64 %res
+}