[knownbits] Preserve known bits for small shift recurrences

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.)

Differential Revision: https://reviews.llvm.org/D96440

Author: preames

llvm/lib/Analysis/ValueTracking.cpp

@@ -1362,6 +1362,51 @@ static void computeKnownBitsFromOperator(const Operator *I,
         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.

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

@@ -6,12 +6,9 @@ define i64 @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
@@ -29,12 +26,9 @@ define i64 @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
@@ -52,12 +46,9 @@ define i64 @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
@@ -70,6 +61,28 @@ exit:
   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(
@@ -94,17 +107,40 @@ exit:
   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