[llvm] [InstCombine] Infer nsw flag for `(X <<nuw C1) >>u C --> X << (C1 - C)` (PR #72407)

[Yingwei Zheng via llvm-commits]

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

Alive2: https://alive2.llvm.org/ce/z/nnHAPy
This missed optimization is discovered with the help of https://github.com/AliveToolkit/alive2/pull/962.

>From 05abbaaf9415aa6fb64257cfdcc692024133ffc9 Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Wed, 15 Nov 2023 23:27:26 +0800
Subject: [PATCH] [InstCombine] Infer nsw flag for `(X <<nuw C1) >>u C --> X <<
 (C1 - C)`

---
 llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp     | 3 ++-
 .../InstCombine/canonicalize-shl-lshr-to-masking.ll       | 4 ++--
 llvm/test/Transforms/InstCombine/shift.ll                 | 8 ++++----
 3 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp b/llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp
index f3f4183efbc7895..a234b916f72c0e2 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp
@@ -1293,9 +1293,10 @@ Instruction *InstCombinerImpl::visitLShr(BinaryOperator &I) {
         unsigned ShlAmtC = C1->getZExtValue();
         Constant *ShiftDiff = ConstantInt::get(Ty, ShlAmtC - ShAmtC);
         if (cast<BinaryOperator>(Op0)->hasNoUnsignedWrap()) {
-          // (X <<nuw C1) >>u C --> X <<nuw (C1 - C)
+          // (X <<nuw C1) >>u C --> X <<nuw/nsw (C1 - C)
           auto *NewShl = BinaryOperator::CreateShl(X, ShiftDiff);
           NewShl->setHasNoUnsignedWrap(true);
+          NewShl->setHasNoSignedWrap(ShAmtC > 0);
           return NewShl;
         }
         if (Op0->hasOneUse()) {
diff --git a/llvm/test/Transforms/InstCombine/canonicalize-shl-lshr-to-masking.ll b/llvm/test/Transforms/InstCombine/canonicalize-shl-lshr-to-masking.ll
index a12cea49786b60b..7c1c18cc214a5c5 100644
--- a/llvm/test/Transforms/InstCombine/canonicalize-shl-lshr-to-masking.ll
+++ b/llvm/test/Transforms/InstCombine/canonicalize-shl-lshr-to-masking.ll
@@ -90,7 +90,7 @@ define i32 @positive_sameconst_shlnuw(i32 %x) {
 
 define i32 @positive_biggerShl_shlnuw(i32 %x) {
 ; CHECK-LABEL: @positive_biggerShl_shlnuw(
-; CHECK-NEXT:    [[RET:%.*]] = shl nuw i32 [[X:%.*]], 5
+; CHECK-NEXT:    [[RET:%.*]] = shl nuw nsw i32 [[X:%.*]], 5
 ; CHECK-NEXT:    ret i32 [[RET]]
 ;
   %t0 = shl nuw i32 %x, 10
@@ -288,7 +288,7 @@ define i32 @positive_biggerShl_shlnuw_multiuse(i32 %x) {
 ; CHECK-LABEL: @positive_biggerShl_shlnuw_multiuse(
 ; CHECK-NEXT:    [[T0:%.*]] = shl nuw i32 [[X:%.*]], 10
 ; CHECK-NEXT:    call void @use32(i32 [[T0]])
-; CHECK-NEXT:    [[RET:%.*]] = shl nuw i32 [[X]], 5
+; CHECK-NEXT:    [[RET:%.*]] = shl nuw nsw i32 [[X]], 5
 ; CHECK-NEXT:    ret i32 [[RET]]
 ;
   %t0 = shl nuw i32 %x, 10
diff --git a/llvm/test/Transforms/InstCombine/shift.ll b/llvm/test/Transforms/InstCombine/shift.ll
index f7135971cf9eebd..e928d99d2feb127 100644
--- a/llvm/test/Transforms/InstCombine/shift.ll
+++ b/llvm/test/Transforms/InstCombine/shift.ll
@@ -951,11 +951,11 @@ define <2 x i32> @test52_splat_vec(<2 x i32> %x) {
   ret <2 x i32> %B
 }
 
-; (X <<nuw C1) >>u C2 --> X <<nuw (C1 - C2)
+; (X <<nuw C1) >>u C2 --> X <<nuw/nsw (C1 - C2)
 
 define i32 @test53(i32 %x) {
 ; CHECK-LABEL: @test53(
-; CHECK-NEXT:    [[B:%.*]] = shl nuw i32 [[X:%.*]], 2
+; CHECK-NEXT:    [[B:%.*]] = shl nuw nsw i32 [[X:%.*]], 2
 ; CHECK-NEXT:    ret i32 [[B]]
 ;
   %A = shl nuw i32 %x, 3
@@ -963,11 +963,11 @@ define i32 @test53(i32 %x) {
   ret i32 %B
 }
 
-; (X <<nuw C1) >>u C2 --> X <<nuw (C1 - C2)
+; (X <<nuw C1) >>u C2 --> X <<nuw/nsw (C1 - C2)
 
 define <2 x i32> @test53_splat_vec(<2 x i32> %x) {
 ; CHECK-LABEL: @test53_splat_vec(
-; CHECK-NEXT:    [[B:%.*]] = shl nuw <2 x i32> [[X:%.*]], <i32 2, i32 2>
+; CHECK-NEXT:    [[B:%.*]] = shl nuw nsw <2 x i32> [[X:%.*]], <i32 2, i32 2>
 ; CHECK-NEXT:    ret <2 x i32> [[B]]
 ;
   %A = shl nuw <2 x i32> %x, <i32 3, i32 3>