[InstCombine] Propagate flags when folding consecutative shifts

[goldsteinn]

• [InstCombine] Add tests for propagating flags when folding consecutative shifts; NFC
• [InstCombine] Propagate flags when folding consecutative shifts

When we fold (shift (shift C0, x), C1) we can propagate flags that
are common to both shifts.

Proofs: https://alive2.llvm.org/ce/z/LkEzXD

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: None (goldsteinn)

Changes

• [InstCombine] Add tests for propagating flags when folding consecutative shifts; NFC
• [InstCombine] Propagate flags when folding consecutative shifts

---

Full diff: https://github.com/llvm/llvm-project/pull/94872.diff

2 Files Affected:

• (modified) llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp (+12-3)
• (modified) llvm/test/Transforms/InstCombine/shift.ll (+80)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp b/llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp
index 9ff817da79368..4a014ab6e044e 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp
@@ -767,11 +767,20 @@ Instruction *InstCombinerImpl::FoldShiftByConstant(Value *Op0, Constant *C1,
   // (C2 >> X) >> C1 --> (C2 >> C1) >> X
   Constant *C2;
   Value *X;
-  if (match(Op0, m_BinOp(I.getOpcode(), m_ImmConstant(C2), m_Value(X))))
-    return BinaryOperator::Create(
+  bool IsLeftShift = I.getOpcode() == Instruction::Shl;
+  if (match(Op0, m_BinOp(I.getOpcode(), m_ImmConstant(C2), m_Value(X)))) {
+    Instruction *R = BinaryOperator::Create(
         I.getOpcode(), Builder.CreateBinOp(I.getOpcode(), C2, C1), X);
+    BinaryOperator *BO0 = cast<BinaryOperator>(Op0);
+    if (IsLeftShift) {
+      R->setHasNoUnsignedWrap(I.hasNoUnsignedWrap() &&
+                              BO0->hasNoUnsignedWrap());
+      R->setHasNoSignedWrap(I.hasNoSignedWrap() && BO0->hasNoSignedWrap());
+    } else
+      R->setIsExact(I.isExact() && BO0->isExact());
+    return R;
+  }
 
-  bool IsLeftShift = I.getOpcode() == Instruction::Shl;
   Type *Ty = I.getType();
   unsigned TypeBits = Ty->getScalarSizeInBits();
 
diff --git a/llvm/test/Transforms/InstCombine/shift.ll b/llvm/test/Transforms/InstCombine/shift.ll
index 8da52e0746373..03536f37fe762 100644
--- a/llvm/test/Transforms/InstCombine/shift.ll
+++ b/llvm/test/Transforms/InstCombine/shift.ll
@@ -2240,4 +2240,84 @@ define i129 @shift_zext_not_nneg(i8 %arg) {
   ret i129 %shl
 }
 
+define i8 @src_shl_nsw(i8 %x) {
+; CHECK-LABEL: @src_shl_nsw(
+; CHECK-NEXT:    [[R:%.*]] = shl nsw i8 32, [[X:%.*]]
+; CHECK-NEXT:    ret i8 [[R]]
+;
+  %sh = shl nsw i8 1, %x
+  %r = shl nsw i8 %sh, 5
+  ret i8 %r
+}
+
+define i8 @src_shl_nsw_fail(i8 %x) {
+; CHECK-LABEL: @src_shl_nsw_fail(
+; CHECK-NEXT:    [[R:%.*]] = shl i8 32, [[X:%.*]]
+; CHECK-NEXT:    ret i8 [[R]]
+;
+  %sh = shl nsw i8 1, %x
+  %r = shl i8 %sh, 5
+  ret i8 %r
+}
+
+define i8 @src_shl_nuw(i8 %x) {
+; CHECK-LABEL: @src_shl_nuw(
+; CHECK-NEXT:    [[R:%.*]] = shl nuw i8 12, [[X:%.*]]
+; CHECK-NEXT:    ret i8 [[R]]
+;
+  %sh = shl nuw i8 3, %x
+  %r = shl nuw i8 %sh, 2
+  ret i8 %r
+}
+
+define i8 @src_shl_nuw_fail(i8 %x) {
+; CHECK-LABEL: @src_shl_nuw_fail(
+; CHECK-NEXT:    [[R:%.*]] = shl i8 12, [[X:%.*]]
+; CHECK-NEXT:    ret i8 [[R]]
+;
+  %sh = shl i8 3, %x
+  %r = shl nuw i8 %sh, 2
+  ret i8 %r
+}
+
+define i8 @src_lshr_exact(i8 %x) {
+; CHECK-LABEL: @src_lshr_exact(
+; CHECK-NEXT:    [[R:%.*]] = lshr exact i8 48, [[X:%.*]]
+; CHECK-NEXT:    ret i8 [[R]]
+;
+  %sh = lshr exact i8 96, %x
+  %r = lshr exact i8 %sh, 1
+  ret i8 %r
+}
+
+define i8 @src_lshr_exact_fail(i8 %x) {
+; CHECK-LABEL: @src_lshr_exact_fail(
+; CHECK-NEXT:    [[R:%.*]] = lshr i8 48, [[X:%.*]]
+; CHECK-NEXT:    ret i8 [[R]]
+;
+  %sh = lshr exact i8 96, %x
+  %r = lshr i8 %sh, 1
+  ret i8 %r
+}
+
+define i8 @src_ashr_exact(i8 %x) {
+; CHECK-LABEL: @src_ashr_exact(
+; CHECK-NEXT:    [[R:%.*]] = lshr exact i8 8, [[X:%.*]]
+; CHECK-NEXT:    ret i8 [[R]]
+;
+  %sh = ashr exact i8 32, %x
+  %r = ashr exact i8 %sh, 2
+  ret i8 %r
+}
+
+define i8 @src_ashr_exact_fail(i8 %x) {
+; CHECK-LABEL: @src_ashr_exact_fail(
+; CHECK-NEXT:    [[R:%.*]] = lshr i8 8, [[X:%.*]]
+; CHECK-NEXT:    ret i8 [[R]]
+;
+  %sh = ashr i8 32, %x
+  %r = ashr exact i8 %sh, 2
+  ret i8 %r
+}
+
 declare i16 @llvm.umax.i16(i16, i16)

[nikic]

LGTM

(It looks like your alive2 link also has the proofs for flag-preservation for a related case that's not part of this PR.)

[goldsteinn]

LGTM

(It looks like your alive2 link also has the proofs for flag-preservation for a related case that's not part of this PR.)

Bah sorry for the confusion. Realized that case is already handled after writing the proofs.

[dtcxzyw]

LGTM.

[dtcxzyw]

Can you replace LHS with a negative number?