[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

1239

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

1239

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

1012

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.

1368

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

1012

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.

1368

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

1012

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

1392

nit: switch (

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

86

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

static void computeKnownBitsFromShiftOperator(
 
   // Early exit if we can't constrain any well-defined shift amount.
   if (!(ShiftAmtKZ & (PowerOf2Ceil(BitWidth) - 1)) &&
       !(ShiftAmtKO & (PowerOf2Ceil(BitWidth) - 1))) {
     ShifterOperandIsNonZero =
         isKnownNonZero(I->getOperand(1), DemandedElts, Depth + 1, Q);
     if (!*ShifterOperandIsNonZero)
       return;
   }
 
   Known.Zero.setAllBits();
   Known.One.setAllBits();
   for (unsigned ShiftAmt = 0; ShiftAmt < BitWidth; ++ShiftAmt) {
     // Combine the shifted known input bits only for those shift amounts
     // compatible with its known constraints.
     if ((ShiftAmt & ~ShiftAmtKZ) != ShiftAmt)
       continue;
static void computeKnownBitsFromOperator(const Operator *I,
   case Instruction::AShr: {
     auto KF = [](const KnownBits &KnownVal, const KnownBits &KnownAmt) {
       return KnownBits::ashr(KnownVal, KnownAmt);
     };
     computeKnownBitsFromShiftOperator(I, DemandedElts, Known, Known2, Depth, Q,
                                       KF);
     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;
   }
   case Instruction::Add: {
     bool NSW = Q.IIQ.hasNoSignedWrap(cast<OverflowingBinaryOperator>(I));
     computeKnownBitsAddSub(true, I->getOperand(0), I->getOperand(1), NSW,
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 a shift recurrence, we know the bits being shifted in.
+        // We can combine that with information about the start value of the
+        // recurrence to conclude facts about the result.
+        if (Opcode == Instruction::LShr ||
+            Opcode == Instruction::AShr ||
+            Opcode == Instruction::Shl) {
+          Value *LL = LU->getOperand(0);
+          Value *LR = LU->getOperand(1);
+          // Find a recurrence.
+          if (LL == I)
+            L = LR;
+          else
+            continue; // Check for recurrence with L and R flipped.
+
+          // We have matched a recurrence of the form:
+          // %iv = [R, %entry], [%iv.next, %backedge]
+          // %iv.next = shift_op %iv, L
+
+          // Recurse with the phi context to avoid concern about whether facts
+          // inferred hold at original context instruction.  TODO: It may be
+          // correct to use the original context.  IF warranted, explore and
+          // add sufficient tests to cover.
+          Query RecQ = Q;
+          RecQ.CxtI = P;
+          computeKnownBits(R, DemandedElts, Known2, Depth + 1, RecQ);
+          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());
+            break;
+          };
+        }
+
         // 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) {

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:    [[IV_LSHR:%.*]] = phi i64 [ 1023, [[ENTRY:%.*]] ], [ [[IV_LSHR_NEXT:%.*]], [[LOOP]] ]
-; CHECK-NEXT:    [[IV_LSHR_NEXT]] = lshr i64 [[IV_LSHR]], 1
 ; CHECK-NEXT:    br i1 undef, label [[EXIT:%.*]], label [[LOOP]]
 ; CHECK:       exit:
-; CHECK-NEXT:    [[RES:%.*]] = or i64 [[IV_LSHR]], 1023
-; CHECK-NEXT:    ret i64 [[RES]]
+; 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:    [[IV_ASHR:%.*]] = phi i64 [ 1023, [[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]]
+; 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:    [[IV_ASHR:%.*]] = phi i64 [ -1023, [[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]]
+; 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
 }
 
+; Same as previous, but swapped operands to phi
+define i64 @test_ashr_ones2() {
+; CHECK-LABEL: @test_ashr_ones2(
+; 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 [%iv.ashr.next, %loop], [-1023, %entry]
+  %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
 }
 
+; Negative case where we don't have a (shift) recurrence because the operands
+; of the ashr are swapped.  (This does end up being a divide recurrence.)
+define i64 @test_ashr_wrong_op(i64 %start) {
+; CHECK-LABEL: @test_ashr_wrong_op(
+; 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]] = lshr i64 1, [[IV_ASHR]]
+; 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 1, %iv.ashr
+  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:    [[IV_SHL:%.*]] = phi i64 [ 8, [[ENTRY:%.*]] ], [ [[IV_SHL_NEXT:%.*]], [[LOOP]] ]
-; CHECK-NEXT:    [[IV_SHL_NEXT]] = shl i64 [[IV_SHL]], 1
 ; CHECK-NEXT:    br i1 undef, label [[EXIT:%.*]], label [[LOOP]]
 ; CHECK:       exit:
-; CHECK-NEXT:    [[RES:%.*]] = and i64 [[IV_SHL]], 6
-; CHECK-NEXT:    ret i64 [[RES]]
+; 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
 }