[llvm] [InstCombine] Turn AShr into LShr more often in SimplifyDemandedUseBits (PR #99155)

[Björn Pettersson via llvm-commits]

https://github.com/bjope updated https://github.com/llvm/llvm-project/pull/99155

>From b8c4c58ecf186dd91f40bdff4d1bdad403435789 Mon Sep 17 00:00:00 2001
From: Bjorn Pettersson <bjorn.a.pettersson at ericsson.com>
Date: Mon, 8 Jul 2024 01:35:53 +0200
Subject: [PATCH] [InstCombine] Turn AShr into LShr more often in
 SimplifyDemandedUseBits (#99155)

The functional change here is to undo "llvm-svn: 311773", aka D36936,
aka commit 22178dd33b3460207b8. That patch avoided to convert AShr
into LShr in SimplifyDemandedUseBits based on known sign bits
analysis. Even if it would be legal to turn the shift into a
logical shift (given by the fact that the shifted in bits wasn't
demanded), that patch prevented converting the shift into LShr when
any of the original sign bits were demanded.
One side effect of the reverted functionalty was that the better we
were at computing number of sign bits, the less likely it was that
we would replace AShr by LShr during SimplifyDemandedUseBits. This
was seen in https://github.com/llvm/llvm-project/pull/97693/ when
an improvement of ComputeNumSignBits resulted in regressions due
to no longer rewriting AShr to LShr.
The test case from D36936 still passes after this commit. So it seems
like at least the compiler has been taught how to optimize that
scenario even if we do the AShr->LShr transform more aggressively.
---
 .../InstCombineSimplifyDemanded.cpp           | 26 ++++++++-----------
 .../Transforms/InstCombine/ashr-demand.ll     |  9 +++----
 2 files changed, 14 insertions(+), 21 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp b/llvm/lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp
index 081e783c964fd..8a6ec3076ac62 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp
@@ -806,34 +806,30 @@ Value *InstCombinerImpl::SimplifyDemandedUseBits(Instruction *I,
 
       // Signed shift right.
       APInt DemandedMaskIn(DemandedMask.shl(ShiftAmt));
-      // If any of the high bits are demanded, we should set the sign bit as
-      // demanded.
-      if (DemandedMask.countl_zero() <= ShiftAmt)
+      // If any of the bits being shifted in are demanded, then we should set
+      // the sign bit as demanded.
+      bool ShiftedInBitsDemanded = DemandedMask.countl_zero() < ShiftAmt;
+      if (ShiftedInBitsDemanded)
         DemandedMaskIn.setSignBit();
-
       if (SimplifyDemandedBits(I, 0, DemandedMaskIn, Known, Depth + 1, Q)) {
         // exact flag may not longer hold.
         I->dropPoisonGeneratingFlags();
         return I;
       }
 
-      Known = KnownBits::ashr(
-          Known, KnownBits::makeConstant(APInt(BitWidth, ShiftAmt)),
-          ShiftAmt != 0, I->isExact());
-
-      // If the input sign bit is known to be zero, or if none of the top bits
-      // are demanded, turn this into an unsigned shift right.
-      assert(BitWidth > ShiftAmt && "Shift amount not saturated?");
-      APInt HighBits(APInt::getHighBitsSet(
-          BitWidth, std::min(SignBits + ShiftAmt - 1, BitWidth)));
-      if (Known.Zero[BitWidth-ShiftAmt-1] ||
-          !DemandedMask.intersects(HighBits)) {
+      // If the input sign bit is known to be zero, or if none of the shifted in
+      // bits are demanded, turn this into an unsigned shift right.
+      if (Known.Zero[BitWidth - 1] || !ShiftedInBitsDemanded) {
         BinaryOperator *LShr = BinaryOperator::CreateLShr(I->getOperand(0),
                                                           I->getOperand(1));
         LShr->setIsExact(cast<BinaryOperator>(I)->isExact());
         LShr->takeName(I);
         return InsertNewInstWith(LShr, I->getIterator());
       }
+
+      Known = KnownBits::ashr(
+          Known, KnownBits::makeConstant(APInt(BitWidth, ShiftAmt)),
+          ShiftAmt != 0, I->isExact());
     } else {
       llvm::computeKnownBits(I, Known, Depth, Q);
     }
diff --git a/llvm/test/Transforms/InstCombine/ashr-demand.ll b/llvm/test/Transforms/InstCombine/ashr-demand.ll
index be7496d38f272..a0e2af93b809b 100644
--- a/llvm/test/Transforms/InstCombine/ashr-demand.ll
+++ b/llvm/test/Transforms/InstCombine/ashr-demand.ll
@@ -74,12 +74,9 @@ define i16 @ashr_can_be_lshr(i32 %a) {
 ; optimizations.
 define i32 @ashr_can_be_lshr_2(i32 %a) {
 ; CHECK-LABEL: @ashr_can_be_lshr_2(
-; CHECK-NEXT:    [[TMP1:%.*]] = or i32 [[A:%.*]], 1056964608
-; CHECK-NEXT:    [[OR:%.*]] = zext i32 [[TMP1]] to i64
-; CHECK-NEXT:    [[SHL:%.*]] = shl i64 [[OR]], 34
-; CHECK-NEXT:    [[ASHR:%.*]] = ashr exact i64 [[SHL]], 32
-; CHECK-NEXT:    [[TRUNC:%.*]] = trunc nsw i64 [[ASHR]] to i32
-; CHECK-NEXT:    ret i32 [[TRUNC]]
+; CHECK-NEXT:    [[TMP1:%.*]] = shl i32 [[A:%.*]], 2
+; CHECK-NEXT:    [[TMP2:%.*]] = or i32 [[TMP1]], -67108864
+; CHECK-NEXT:    ret i32 [[TMP2]]
 ;
   %ext = zext i32 %a to i64
   %or = or i64 %ext, 4278190080