[llvm] [InstCombine] Propagate NUW flags for `shl (lshr X, C1), C2 -> shl X, C2-C1` (PR #72525)

[Yingwei Zheng via llvm-commits]

https://github.com/dtcxzyw created https://github.com/llvm/llvm-project/pull/72525

Alive2: https://alive2.llvm.org/ce/z/KNXNQA

This missed optimization is discovered with the help of https://github.com/AliveToolkit/alive2/pull/962.

>From ee113d1a959ccb98b0c951b37cafe8d26848ac4b Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Thu, 16 Nov 2023 23:30:14 +0800
Subject: [PATCH] [InstCombine] Propagate NUW flags for `shl (lshr X, C1), C2
 -> shl X, C2-C1`

---
 .../lib/Transforms/InstCombine/InstCombineShifts.cpp | 12 ++++++++++--
 .../InstCombine/canonicalize-lshr-shl-to-masking.ll  |  6 +++---
 .../Transforms/InstCombine/known-signbit-shift.ll    |  3 +--
 3 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp b/llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp
index a234b916f72c0e2..496891d4857a7ca 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp
@@ -1045,7 +1045,11 @@ Instruction *InstCombinerImpl::visitShl(BinaryOperator &I) {
         // If C1 < C: (X >>?,exact C1) << C --> X << (C - C1)
         Constant *ShiftDiff = ConstantInt::get(Ty, ShAmtC - ShrAmt);
         auto *NewShl = BinaryOperator::CreateShl(X, ShiftDiff);
-        NewShl->setHasNoUnsignedWrap(I.hasNoUnsignedWrap());
+        NewShl->setHasNoUnsignedWrap(
+            I.hasNoUnsignedWrap() ||
+            (ShrAmt &&
+             cast<Instruction>(Op0)->getOpcode() == Instruction::LShr &&
+             I.hasNoSignedWrap()));
         NewShl->setHasNoSignedWrap(I.hasNoSignedWrap());
         return NewShl;
       }
@@ -1066,7 +1070,11 @@ Instruction *InstCombinerImpl::visitShl(BinaryOperator &I) {
         // If C1 < C: (X >>? C1) << C --> (X << (C - C1)) & (-1 << C)
         Constant *ShiftDiff = ConstantInt::get(Ty, ShAmtC - ShrAmt);
         auto *NewShl = BinaryOperator::CreateShl(X, ShiftDiff);
-        NewShl->setHasNoUnsignedWrap(I.hasNoUnsignedWrap());
+        NewShl->setHasNoUnsignedWrap(
+            I.hasNoUnsignedWrap() ||
+            (ShrAmt &&
+             cast<Instruction>(Op0)->getOpcode() == Instruction::LShr &&
+             I.hasNoSignedWrap()));
         NewShl->setHasNoSignedWrap(I.hasNoSignedWrap());
         Builder.Insert(NewShl);
         APInt Mask(APInt::getHighBitsSet(BitWidth, BitWidth - ShAmtC));
diff --git a/llvm/test/Transforms/InstCombine/canonicalize-lshr-shl-to-masking.ll b/llvm/test/Transforms/InstCombine/canonicalize-lshr-shl-to-masking.ll
index 0b8165e577133e9..ddd72c320f92bb7 100644
--- a/llvm/test/Transforms/InstCombine/canonicalize-lshr-shl-to-masking.ll
+++ b/llvm/test/Transforms/InstCombine/canonicalize-lshr-shl-to-masking.ll
@@ -141,8 +141,8 @@ define i8 @positive_biggerlshr_shlnsw(i8 %x) {
 
 define i8 @positive_biggershl_shlnsw(i8 %x) {
 ; CHECK-LABEL: @positive_biggershl_shlnsw(
-; CHECK-NEXT:    [[TMP1:%.*]] = shl nsw i8 [[X:%.*]], 3
-; CHECK-NEXT:    [[RET:%.*]] = and i8 [[TMP1]], -64
+; CHECK-NEXT:    [[TMP1:%.*]] = shl nuw nsw i8 [[X:%.*]], 3
+; CHECK-NEXT:    [[RET:%.*]] = and i8 [[TMP1]], 64
 ; CHECK-NEXT:    ret i8 [[RET]]
 ;
   %t0 = lshr i8 %x, 3
@@ -273,7 +273,7 @@ define i8 @positive_biggerlshr_shlnsw_lshrexact(i8 %x) {
 
 define i8 @positive_biggershl_shlnsw_lshrexact(i8 %x) {
 ; CHECK-LABEL: @positive_biggershl_shlnsw_lshrexact(
-; CHECK-NEXT:    [[RET:%.*]] = shl nsw i8 [[X:%.*]], 3
+; CHECK-NEXT:    [[RET:%.*]] = shl nuw nsw i8 [[X:%.*]], 3
 ; CHECK-NEXT:    ret i8 [[RET]]
 ;
   %t0 = lshr exact i8 %x, 3
diff --git a/llvm/test/Transforms/InstCombine/known-signbit-shift.ll b/llvm/test/Transforms/InstCombine/known-signbit-shift.ll
index 17c7fbab12dc470..7492bf96b4fbcb5 100644
--- a/llvm/test/Transforms/InstCombine/known-signbit-shift.ll
+++ b/llvm/test/Transforms/InstCombine/known-signbit-shift.ll
@@ -5,8 +5,7 @@
 ; with nsw flag should also be non-negative
 define i1 @test_shift_nonnegative(i32 %a) {
 ; CHECK-LABEL: @test_shift_nonnegative(
-; CHECK-NEXT:    [[CMP:%.*]] = icmp sgt i32 [[A:%.*]], -1
-; CHECK-NEXT:    ret i1 [[CMP]]
+; CHECK-NEXT:    ret i1 true
 ;
   %b = lshr i32 %a, 2
   %shift = shl nsw i32 %b, 3