[llvm] [InstCombine][ValueTracking] Improve logic for shift flags. (PR #67758)

[via llvm-commits]

================
@@ -941,6 +941,60 @@ Instruction *InstCombinerImpl::foldLShrOverflowBit(BinaryOperator &I) {
   return new ZExtInst(Overflow, Ty);
 }
 
+// Try to set nuw/nsw flags on shl or exact flag on lshr/ashr using knownbits.
+static bool setShiftFlags(BinaryOperator &I, const SimplifyQuery &Q) {
+  assert(I.isShift() && "Expected a shift as input");
+  // We already have all the flags.
+  if (I.getOpcode() == Instruction::Shl) {
+    if (I.hasNoUnsignedWrap() && I.hasNoSignedWrap())
+      return false;
+  } else {
+    if (I.isExact())
+      return false;
+  }
+
+  // Compute what we know about shift count.
+  KnownBits KnownCnt =
+      computeKnownBits(I.getOperand(1), Q.DL, /*Depth*/ 0, Q.AC, Q.CxtI, Q.DT);
+  // If we know nothing about shift count or its a poison shift, we won't be
+  // able to prove anything so return before computing shift amount.
+  if (KnownCnt.isUnknown())
+    return false;
+  unsigned BitWidth = KnownCnt.getBitWidth();
+  APInt MaxCnt = KnownCnt.getMaxValue();
+  if (MaxCnt.uge(BitWidth))
+    return false;
+
+  KnownBits KnownAmt =
+      computeKnownBits(I.getOperand(0), Q.DL, /*Depth*/ 0, Q.AC, Q.CxtI, Q.DT);
+  bool Changed = false;
+
+  if (I.getOpcode() == Instruction::Shl) {
+    // If we have as many leading zeros than maximum shift cnt we have nuw.
+    if (!I.hasNoUnsignedWrap() && MaxCnt.ule(KnownAmt.countMinLeadingZeros())) {
+      I.setHasNoUnsignedWrap();
+      Changed = true;
+    }
+    // If we have more sign bits than maximum shift cnt we have nsw.
+    if (!I.hasNoSignedWrap()) {
+      if (MaxCnt.ult(KnownAmt.countMinSignBits()) ||
+          MaxCnt.ult(ComputeNumSignBits(I.getOperand(0), Q.DL, /*Depth*/ 0,
+                                        Q.AC, Q.CxtI, Q.DT))) {
+        I.setHasNoSignedWrap();
+        Changed = true;
+      }
+    }
+    return Changed;
+  }
+  if (!I.isExact()) {
+    // If we have at least as many trailing zeros as maximum count then we have
+    // exact.
+    Changed = MaxCnt.ule(KnownAmt.countMinTrailingZeros());
+    I.setIsExact(Changed);
+  }
+  return Changed;
----------------
goldsteinn wrote:

Thats not the same, this will cause inf loop if the instruction was already exact (we will always return changed).

https://github.com/llvm/llvm-project/pull/67758