[KnownBits] Make shl/lshr/ashr implementations optimal

[nikic]

Summary

The implementations for shifts were suboptimal in the case where the max shift amount was >= bitwidth. In that case we should still use the usual code clamped to BitWidth-1 rather than just giving up entirely.

Additionally, there was an implementation bug where the known zero bits for the individual shift amounts were not set in the shl/lshr implementations. I think after these changes, we'll be able to drop some of the code in ValueTracking which *also* evaluates all possible shift amounts and has been papering over this issue.

For the "all poison" case I've opted to return an unknown value for now. It would be better to return zero, but this has fairly substantial test fallout, so I figured it's best to not mix it into this change. (The "correct" return value would be a conflict, but given that a lot of our APIs assert conflict-freedom, that's probably not the best idea to actually return.)

[nikic]

llvm/test/CodeGen/AMDGPU/amdgpu.private-memory.ll

224

I believe this is correct, because SCALED_IDX is IDX << 4 and as such at least 16. As such, 0x10000 >> SCALED_IDX is either zero or one and the and mask can be narrowed to 1.

llvm/test/Transforms/InstCombine/not-add.ll

175

A nominal regression because I did not try to preserve the exact behavior for "always poison" and always return unknown. If we switch to returning zero this whole code folds away (as well as code in many other tests).

[foad]

The "correct" return value would be a conflict,

Yes!

but given that a lot of our APIs assert conflict-freedom

That is a shame.

[goldstein.w.n]

llvm/test/CodeGen/AMDGPU/amdgpu.private-memory.ll

224

Was the old value buggy then?

[nikic]

llvm/test/CodeGen/AMDGPU/amdgpu.private-memory.ll

224

No, it's also correct, the constant is just unnecessarily wide. I believe this is done as part of demanded bits simplification.

[goldstein.w.n]

llvm/lib/Support/KnownBits.cpp

243–248

Not for this patch, but in future installments, maybe a third argument to specify behavior here? In the X86 backend, for example, we remove masks on shiftamt b.c its known processor will just modulo shiftamt by bitwidth. Likewise some targets just return zero. If we could specify the behavior, it may be usable in TargetLowering::computeKnownBitsForTargetNode.

llvm/lib/Support/KnownBits.cpp

if (RHS.isConstant() && RHS.getConstant().ult(BitWidth)) {
     // Low bits are known zero.
     Known.Zero.setLowBits(Shift);
     return Known;
   }
 
   // No matter the shift amount, the trailing zeros will stay zero.
   unsigned MinTrailingZeros = LHS.countMinTrailingZeros();
 
-  // Minimum shift amount low bits are known zero.
   APInt MinShiftAmount = RHS.getMinValue();
-  if (MinShiftAmount.ult(BitWidth)) {
-    MinTrailingZeros += MinShiftAmount.getZExtValue();
-    MinTrailingZeros = std::min(MinTrailingZeros, BitWidth);
-  }
+  if (MinShiftAmount.uge(BitWidth))
+    // Always poison. Return unknown because we don't like returning conflict.
+    return Known;
+
+  // Minimum shift amount low bits are known zero.
 
   // If the maximum shift is in range, then find the common bits from all
   // possible shifts.
   APInt MaxShiftAmount = RHS.getMaxValue();
-  if (MaxShiftAmount.ult(BitWidth) && !LHS.isUnknown()) {
+  if (!LHS.isUnknown()) {
     uint64_t ShiftAmtZeroMask = (~RHS.Zero).getZExtValue();
     uint64_t ShiftAmtOneMask = RHS.One.getZExtValue();
     assert(MinShiftAmount.ult(MaxShiftAmount) && "Illegal shift range");
     Known.Zero.setAllBits();
     Known.One.setAllBits();
     for (uint64_t ShiftAmt = MinShiftAmount.getZExtValue(),
-                  MaxShiftAmt = MaxShiftAmount.getZExtValue();
+                  MaxShiftAmt = MaxShiftAmount.getLimitedValue(BitWidth - 1);
          ShiftAmt <= MaxShiftAmt; ++ShiftAmt) {
       // Skip if the shift amount is impossible.
       if ((ShiftAmtZeroMask & ShiftAmt) != ShiftAmt ||
           (ShiftAmtOneMask | ShiftAmt) != ShiftAmt)
         continue;
       KnownBits SpecificShift;
       SpecificShift.Zero = LHS.Zero << ShiftAmt;
+      SpecificShift.Zero.setLowBits(ShiftAmt);
       SpecificShift.One = LHS.One << ShiftAmt;
       Known = KnownBits::commonBits(Known, SpecificShift);
       if (Known.isUnknown())
         break;
     }
   }
 
   Known.Zero.setLowBits(MinTrailingZeros);
if (RHS.isConstant() && RHS.getConstant().ult(BitWidth)) {
     return Known;
   }
 
   // No matter the shift amount, the leading zeros will stay zero.
   unsigned MinLeadingZeros = LHS.countMinLeadingZeros();
 
   // Minimum shift amount high bits are known zero.
   APInt MinShiftAmount = RHS.getMinValue();
-  if (MinShiftAmount.ult(BitWidth)) {
-    MinLeadingZeros += MinShiftAmount.getZExtValue();
-    MinLeadingZeros = std::min(MinLeadingZeros, BitWidth);
-  }
+  if (MinShiftAmount.uge(BitWidth))
+    // Always poison. Return unknown because we don't like returning conflict.
+    return Known;
+
 
   // If the maximum shift is in range, then find the common bits from all
   // possible shifts.
   APInt MaxShiftAmount = RHS.getMaxValue();
-  if (MaxShiftAmount.ult(BitWidth) && !LHS.isUnknown()) {
+  if (!LHS.isUnknown()) {
     uint64_t ShiftAmtZeroMask = (~RHS.Zero).getZExtValue();
     uint64_t ShiftAmtOneMask = RHS.One.getZExtValue();
     assert(MinShiftAmount.ult(MaxShiftAmount) && "Illegal shift range");
     Known.Zero.setAllBits();
     Known.One.setAllBits();
     for (uint64_t ShiftAmt = MinShiftAmount.getZExtValue(),
-                  MaxShiftAmt = MaxShiftAmount.getZExtValue();
+                  MaxShiftAmt = MaxShiftAmount.getLimitedValue(BitWidth - 1);
          ShiftAmt <= MaxShiftAmt; ++ShiftAmt) {
       // Skip if the shift amount is impossible.
       if ((ShiftAmtZeroMask & ShiftAmt) != ShiftAmt ||
           (ShiftAmtOneMask | ShiftAmt) != ShiftAmt)
         continue;
       KnownBits SpecificShift = LHS;
       SpecificShift.Zero.lshrInPlace(ShiftAmt);
+      SpecificShift.Zero.setHighBits(ShiftAmt);
       SpecificShift.One.lshrInPlace(ShiftAmt);
       Known = KnownBits::commonBits(Known, SpecificShift);
       if (Known.isUnknown())
         break;
     }
   }
 
   Known.Zero.setHighBits(MinLeadingZeros);
KnownBits KnownBits::ashr(const KnownBits &LHS, const KnownBits &RHS) {
   }
 
   // No matter the shift amount, the leading sign bits will stay.
   unsigned MinLeadingZeros = LHS.countMinLeadingZeros();
   unsigned MinLeadingOnes = LHS.countMinLeadingOnes();
 
   // Minimum shift amount high bits are known sign bits.
   APInt MinShiftAmount = RHS.getMinValue();
-  if (MinShiftAmount.ult(BitWidth)) {
-    if (MinLeadingZeros) {
-      MinLeadingZeros += MinShiftAmount.getZExtValue();
-      MinLeadingZeros = std::min(MinLeadingZeros, BitWidth);
-    }
-    if (MinLeadingOnes) {
-      MinLeadingOnes += MinShiftAmount.getZExtValue();
-      MinLeadingOnes = std::min(MinLeadingOnes, BitWidth);
-    }
-  }
+  if (MinShiftAmount.uge(BitWidth))
+    // Always poison. Return unknown because we don't like returning conflict.
+    return Known;
+
 
   // If the maximum shift is in range, then find the common bits from all
   // possible shifts.
   APInt MaxShiftAmount = RHS.getMaxValue();
-  if (MaxShiftAmount.ult(BitWidth) && !LHS.isUnknown()) {
+  if (!LHS.isUnknown()) {
     uint64_t ShiftAmtZeroMask = (~RHS.Zero).getZExtValue();
     uint64_t ShiftAmtOneMask = RHS.One.getZExtValue();
     assert(MinShiftAmount.ult(MaxShiftAmount) && "Illegal shift range");
     Known.Zero.setAllBits();
     Known.One.setAllBits();
     for (uint64_t ShiftAmt = MinShiftAmount.getZExtValue(),
-                  MaxShiftAmt = MaxShiftAmount.getZExtValue();
+                  MaxShiftAmt = MaxShiftAmount.getLimitedValue(BitWidth - 1);
          ShiftAmt <= MaxShiftAmt; ++ShiftAmt) {
       // Skip if the shift amount is impossible.
       if ((ShiftAmtZeroMask & ShiftAmt) != ShiftAmt ||
           (ShiftAmtOneMask | ShiftAmt) != ShiftAmt)
         continue;
       KnownBits SpecificShift = LHS;
       SpecificShift.Zero.ashrInPlace(ShiftAmt);
       SpecificShift.One.ashrInPlace(ShiftAmt);

llvm/test/CodeGen/AMDGPU/amdgpu.private-memory.ll

 ; SI-ALLOCA-DAG: buffer_store_short v{{[0-9]+}}, off, s[{{[0-9]+:[0-9]+}}], 0 offset:6 ; encoding: [0x06,0x00,0x68,0xe0
 ; SI-ALLOCA-DAG: buffer_store_short v{{[0-9]+}}, off, s[{{[0-9]+:[0-9]+}}], 0 offset:4 ; encoding: [0x04,0x00,0x68,0xe0
 ; Loaded value is 0 or 1, so sext will become zext, so we get buffer_load_ushort instead of buffer_load_sshort.
 ; SI-ALLOCA: buffer_load_sshort v{{[0-9]+}}, v{{[0-9]+}}, s[{{[0-9]+:[0-9]+}}], 0
 
 ; SI-PROMOTE-VECT: s_load_dword [[IDX:s[0-9]+]]
 ; SI-PROMOTE-VECT: s_lshl_b32 [[SCALED_IDX:s[0-9]+]], [[IDX]], 4
 ; SI-PROMOTE-VECT: s_lshr_b32 [[SREG:s[0-9]+]], 0x10000, [[SCALED_IDX]]
-; SI-PROMOTE-VECT: s_and_b32 s{{[0-9]+}}, [[SREG]], 0xffff
+; SI-PROMOTE-VECT: s_and_b32 s{{[0-9]+}}, [[SREG]], 1
 define amdgpu_kernel void @short_array(ptr addrspace(1) %out, i32 %index) #0 {
 entry:
   %0 = alloca [2 x i16], addrspace(5)
   %1 = getelementptr inbounds [2 x i16], ptr addrspace(5) %0, i32 0, i32 1
   store i16 0, ptr addrspace(5) %0
   store i16 1, ptr addrspace(5) %1
   %2 = getelementptr inbounds [2 x i16], ptr addrspace(5) %0, i32 0, i32 %index
   %3 = load i16, ptr addrspace(5) %2

llvm/test/Transforms/InstCombine/not-add.ll

 }
 
 @g = extern_weak global i32
 define void @pr50370(i32 %x) {
 ; CHECK-LABEL: @pr50370(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[XOR:%.*]] = xor i32 [[X:%.*]], 1
 ; CHECK-NEXT:    [[B15:%.*]] = srem i32 ashr (i32 65536, i32 or (i32 zext (i1 icmp eq (ptr @g, ptr null) to i32), i32 65537)), [[XOR]]
-; CHECK-NEXT:    [[B12:%.*]] = add nuw nsw i32 [[B15]], ashr (i32 65536, i32 or (i32 zext (i1 icmp eq (ptr @g, ptr null) to i32), i32 65537))
+; CHECK-NEXT:    [[B12:%.*]] = add nsw i32 [[B15]], ashr (i32 65536, i32 or (i32 zext (i1 icmp eq (ptr @g, ptr null) to i32), i32 65537))
 ; CHECK-NEXT:    [[B:%.*]] = xor i32 [[B12]], -1
 ; CHECK-NEXT:    store i32 [[B]], ptr undef, align 4
 ; CHECK-NEXT:    ret void
 ;
 entry:
   %xor = xor i32 %x, 1
   %or4 = or i32 or (i32 zext (i1 icmp eq (ptr @g, ptr null) to i32), i32 1), 65536
   %B6 = ashr i32 65536, %or4

llvm/unittests/Support/KnownBitsTest.cpp

testBinaryOpExhaustive(
       },
       [](const APInt &N1, const APInt &N2) -> std::optional<APInt> {
         if (N2.isZero())
           return std::nullopt;
         return N1.srem(N2);
       },
       checkCorrectnessOnlyBinary);
 
-  // TODO: Make optimal for non-constant cases.
   testBinaryOpExhaustive(
       [](const KnownBits &Known1, const KnownBits &Known2) {
         return KnownBits::shl(Known1, Known2);
       },
       [](const APInt &N1, const APInt &N2) -> std::optional<APInt> {
         if (N2.uge(N2.getBitWidth()))
           return std::nullopt;
         return N1.shl(N2);
-      },
-      [](const KnownBits &, const KnownBits &Known) {
-        return Known.isConstant();
       });
   testBinaryOpExhaustive(
       [](const KnownBits &Known1, const KnownBits &Known2) {
         return KnownBits::lshr(Known1, Known2);
       },
       [](const APInt &N1, const APInt &N2) -> std::optional<APInt> {
         if (N2.uge(N2.getBitWidth()))
           return std::nullopt;
         return N1.lshr(N2);
-      },
-      [](const KnownBits &, const KnownBits &Known) {
-        return Known.isConstant();
       });
   testBinaryOpExhaustive(
       [](const KnownBits &Known1, const KnownBits &Known2) {
         return KnownBits::ashr(Known1, Known2);
       },
       [](const APInt &N1, const APInt &N2) -> std::optional<APInt> {
         if (N2.uge(N2.getBitWidth()))
           return std::nullopt;
         return N1.ashr(N2);
-      },
-      [](const KnownBits &, const KnownBits &Known) {
-        return Known.isConstant();
       });
 
   testBinaryOpExhaustive(
       [](const KnownBits &Known1, const KnownBits &Known2) {
         return KnownBits::mul(Known1, Known2);
       },
       [](const APInt &N1, const APInt &N2) { return N1 * N2; },
       checkCorrectnessOnlyBinary);