[llvm] [InstCombine] Canonicalize `switch(X << C)` into `switch(X)` (PR #77068)

Yingwei Zheng via llvm-commits llvm-commits at lists.llvm.org
Fri Jan 5 07:59:35 PST 2024


https://github.com/dtcxzyw updated https://github.com/llvm/llvm-project/pull/77068

>From c2f6e5d124c64cfaf2d72d2a84f14c7d953fbbaf Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Fri, 5 Jan 2024 17:25:37 +0800
Subject: [PATCH 1/2] [InstCombine] Add pre-commit tests. NFC.

---
 .../test/Transforms/InstCombine/switch-shl.ll | 185 ++++++++++++++++++
 1 file changed, 185 insertions(+)
 create mode 100644 llvm/test/Transforms/InstCombine/switch-shl.ll

diff --git a/llvm/test/Transforms/InstCombine/switch-shl.ll b/llvm/test/Transforms/InstCombine/switch-shl.ll
new file mode 100644
index 00000000000000..474bd56c596625
--- /dev/null
+++ b/llvm/test/Transforms/InstCombine/switch-shl.ll
@@ -0,0 +1,185 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 4
+; RUN: opt < %s -passes=instcombine -S | FileCheck %s
+
+define i1 @test_switch_with_shl_mask(i32 %a) {
+; CHECK-LABEL: define i1 @test_switch_with_shl_mask(
+; CHECK-SAME: i32 [[A:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[B:%.*]] = shl i32 [[A]], 24
+; CHECK-NEXT:    switch i32 [[B]], label [[SW_DEFAULT:%.*]] [
+; CHECK-NEXT:      i32 0, label [[SW_BB:%.*]]
+; CHECK-NEXT:      i32 16777216, label [[SW_BB]]
+; CHECK-NEXT:      i32 -2147483648, label [[SW_BB]]
+; CHECK-NEXT:    ]
+; CHECK:       sw.bb:
+; CHECK-NEXT:    ret i1 true
+; CHECK:       sw.default:
+; CHECK-NEXT:    ret i1 false
+;
+entry:
+  %b = shl i32 %a, 24
+  switch i32 %b, label %sw.default [
+  i32 0, label %sw.bb
+  i32 16777216, label %sw.bb
+  i32 2147483648, label %sw.bb
+  ]
+
+sw.bb:
+  ret i1 true
+sw.default:
+  ret i1 false
+}
+
+define i1 @test_switch_with_shl_nuw_multiuse(i32 %a) {
+; CHECK-LABEL: define i1 @test_switch_with_shl_nuw_multiuse(
+; CHECK-SAME: i32 [[A:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[B:%.*]] = shl nuw i32 [[A]], 24
+; CHECK-NEXT:    call void @use(i32 [[B]])
+; CHECK-NEXT:    switch i32 [[B]], label [[SW_DEFAULT:%.*]] [
+; CHECK-NEXT:      i32 0, label [[SW_BB:%.*]]
+; CHECK-NEXT:      i32 16777216, label [[SW_BB]]
+; CHECK-NEXT:      i32 -2147483648, label [[SW_BB]]
+; CHECK-NEXT:    ]
+; CHECK:       sw.bb:
+; CHECK-NEXT:    ret i1 true
+; CHECK:       sw.default:
+; CHECK-NEXT:    ret i1 false
+;
+entry:
+  %b = shl nuw i32 %a, 24
+  call void @use(i32 %b)
+  switch i32 %b, label %sw.default [
+  i32 0, label %sw.bb
+  i32 16777216, label %sw.bb
+  i32 2147483648, label %sw.bb
+  ]
+
+sw.bb:
+  ret i1 true
+sw.default:
+  ret i1 false
+}
+
+define i1 @test_switch_with_shl_nsw_multiuse(i32 %a) {
+; CHECK-LABEL: define i1 @test_switch_with_shl_nsw_multiuse(
+; CHECK-SAME: i32 [[A:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[B:%.*]] = shl nsw i32 [[A]], 24
+; CHECK-NEXT:    call void @use(i32 [[B]])
+; CHECK-NEXT:    switch i32 [[B]], label [[SW_DEFAULT:%.*]] [
+; CHECK-NEXT:      i32 0, label [[SW_BB:%.*]]
+; CHECK-NEXT:      i32 16777216, label [[SW_BB]]
+; CHECK-NEXT:      i32 -2147483648, label [[SW_BB]]
+; CHECK-NEXT:    ]
+; CHECK:       sw.bb:
+; CHECK-NEXT:    ret i1 true
+; CHECK:       sw.default:
+; CHECK-NEXT:    ret i1 false
+;
+entry:
+  %b = shl nsw i32 %a, 24
+  call void @use(i32 %b)
+  switch i32 %b, label %sw.default [
+  i32 0, label %sw.bb
+  i32 16777216, label %sw.bb
+  i32 2147483648, label %sw.bb
+  ]
+
+sw.bb:
+  ret i1 true
+sw.default:
+  ret i1 false
+}
+
+; Negative tests
+
+define i1 @test_switch_with_shl_mask_multiuse(i32 %a) {
+; CHECK-LABEL: define i1 @test_switch_with_shl_mask_multiuse(
+; CHECK-SAME: i32 [[A:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[B:%.*]] = shl i32 [[A]], 24
+; CHECK-NEXT:    call void @use(i32 [[B]])
+; CHECK-NEXT:    switch i32 [[B]], label [[SW_DEFAULT:%.*]] [
+; CHECK-NEXT:      i32 0, label [[SW_BB:%.*]]
+; CHECK-NEXT:      i32 16777216, label [[SW_BB]]
+; CHECK-NEXT:      i32 -2147483648, label [[SW_BB]]
+; CHECK-NEXT:    ]
+; CHECK:       sw.bb:
+; CHECK-NEXT:    ret i1 true
+; CHECK:       sw.default:
+; CHECK-NEXT:    ret i1 false
+;
+entry:
+  %b = shl i32 %a, 24
+  call void @use(i32 %b)
+  switch i32 %b, label %sw.default [
+  i32 0, label %sw.bb
+  i32 16777216, label %sw.bb
+  i32 2147483648, label %sw.bb
+  ]
+
+sw.bb:
+  ret i1 true
+sw.default:
+  ret i1 false
+}
+
+define i1 @test_switch_with_shl_mask_unknown_shamt(i32 %a, i32 %shamt) {
+; CHECK-LABEL: define i1 @test_switch_with_shl_mask_unknown_shamt(
+; CHECK-SAME: i32 [[A:%.*]], i32 [[SHAMT:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[B:%.*]] = shl i32 [[A]], [[SHAMT]]
+; CHECK-NEXT:    switch i32 [[B]], label [[SW_DEFAULT:%.*]] [
+; CHECK-NEXT:      i32 0, label [[SW_BB:%.*]]
+; CHECK-NEXT:      i32 16777216, label [[SW_BB]]
+; CHECK-NEXT:      i32 -2147483648, label [[SW_BB]]
+; CHECK-NEXT:    ]
+; CHECK:       sw.bb:
+; CHECK-NEXT:    ret i1 true
+; CHECK:       sw.default:
+; CHECK-NEXT:    ret i1 false
+;
+entry:
+  %b = shl i32 %a, %shamt
+  switch i32 %b, label %sw.default [
+  i32 0, label %sw.bb
+  i32 16777216, label %sw.bb
+  i32 2147483648, label %sw.bb
+  ]
+
+sw.bb:
+  ret i1 true
+sw.default:
+  ret i1 false
+}
+
+define i1 @test_switch_with_shl_mask_poison(i32 %a) {
+; CHECK-LABEL: define i1 @test_switch_with_shl_mask_poison(
+; CHECK-SAME: i32 [[A:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    switch i32 poison, label [[SW_DEFAULT:%.*]] [
+; CHECK-NEXT:      i32 0, label [[SW_BB:%.*]]
+; CHECK-NEXT:      i32 16777216, label [[SW_BB]]
+; CHECK-NEXT:      i32 -2147483648, label [[SW_BB]]
+; CHECK-NEXT:    ]
+; CHECK:       sw.bb:
+; CHECK-NEXT:    ret i1 true
+; CHECK:       sw.default:
+; CHECK-NEXT:    ret i1 false
+;
+entry:
+  %b = shl i32 %a, 32
+  switch i32 %b, label %sw.default [
+  i32 0, label %sw.bb
+  i32 16777216, label %sw.bb
+  i32 2147483648, label %sw.bb
+  ]
+
+sw.bb:
+  ret i1 true
+sw.default:
+  ret i1 false
+}
+
+declare void @use(i32)

>From 6dd783b9f90ae5f258102d732953567d7e317c02 Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Fri, 5 Jan 2024 17:47:53 +0800
Subject: [PATCH 2/2] [InstCombine] Canonicalize `switch(X << C)` into
 `switch(X)`

---
 .../InstCombine/InstructionCombining.cpp      | 27 +++++++++++++++++++
 .../test/Transforms/InstCombine/switch-shl.ll | 22 +++++++--------
 2 files changed, 38 insertions(+), 11 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
index b45c6784b85b6f..e845bf38df2dfd 100644
--- a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
@@ -3220,6 +3220,33 @@ Instruction *InstCombinerImpl::visitSwitchInst(SwitchInst &SI) {
     return replaceOperand(SI, 0, Op0);
   }
 
+  uint64_t ShiftAmt;
+  if (match(Cond, m_Shl(m_Value(Op0), m_ConstantInt(ShiftAmt))) &&
+      ShiftAmt < Op0->getType()->getScalarSizeInBits() &&
+      all_of(SI.cases(), [&](const auto &Case) {
+        return Case.getCaseValue()->getValue().countr_zero() >= ShiftAmt;
+      })) {
+    // Change 'switch (X << 2) case 4:' into 'switch (X) case 1:'.
+    OverflowingBinaryOperator *Shl = cast<OverflowingBinaryOperator>(Cond);
+    if (Shl->hasNoUnsignedWrap() || Shl->hasNoSignedWrap() ||
+        Shl->hasOneUse()) {
+      Value *NewCond = Op0;
+      if (!Shl->hasNoUnsignedWrap() && !Shl->hasNoSignedWrap()) {
+        // If the shift may wrap, we need to mask off the shifted bits.
+        unsigned BitWidth = Op0->getType()->getScalarSizeInBits();
+        NewCond = Builder.CreateAnd(
+            Op0, APInt::getLowBitsSet(BitWidth, BitWidth - ShiftAmt));
+      }
+      for (auto Case : SI.cases()) {
+        const APInt &CaseVal = Case.getCaseValue()->getValue();
+        APInt ShiftedCase = Shl->hasNoSignedWrap() ? CaseVal.ashr(ShiftAmt)
+                                                   : CaseVal.lshr(ShiftAmt);
+        Case.setValue(ConstantInt::get(SI.getContext(), ShiftedCase));
+      }
+      return replaceOperand(SI, 0, NewCond);
+    }
+  }
+
   KnownBits Known = computeKnownBits(Cond, 0, &SI);
   unsigned LeadingKnownZeros = Known.countMinLeadingZeros();
   unsigned LeadingKnownOnes = Known.countMinLeadingOnes();
diff --git a/llvm/test/Transforms/InstCombine/switch-shl.ll b/llvm/test/Transforms/InstCombine/switch-shl.ll
index 474bd56c596625..037ef37f97a31a 100644
--- a/llvm/test/Transforms/InstCombine/switch-shl.ll
+++ b/llvm/test/Transforms/InstCombine/switch-shl.ll
@@ -5,11 +5,11 @@ define i1 @test_switch_with_shl_mask(i32 %a) {
 ; CHECK-LABEL: define i1 @test_switch_with_shl_mask(
 ; CHECK-SAME: i32 [[A:%.*]]) {
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[B:%.*]] = shl i32 [[A]], 24
-; CHECK-NEXT:    switch i32 [[B]], label [[SW_DEFAULT:%.*]] [
-; CHECK-NEXT:      i32 0, label [[SW_BB:%.*]]
-; CHECK-NEXT:      i32 16777216, label [[SW_BB]]
-; CHECK-NEXT:      i32 -2147483648, label [[SW_BB]]
+; CHECK-NEXT:    [[TRUNC:%.*]] = trunc i32 [[A]] to i8
+; CHECK-NEXT:    switch i8 [[TRUNC]], label [[SW_DEFAULT:%.*]] [
+; CHECK-NEXT:      i8 0, label [[SW_BB:%.*]]
+; CHECK-NEXT:      i8 1, label [[SW_BB]]
+; CHECK-NEXT:      i8 -128, label [[SW_BB]]
 ; CHECK-NEXT:    ]
 ; CHECK:       sw.bb:
 ; CHECK-NEXT:    ret i1 true
@@ -36,10 +36,10 @@ define i1 @test_switch_with_shl_nuw_multiuse(i32 %a) {
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[B:%.*]] = shl nuw i32 [[A]], 24
 ; CHECK-NEXT:    call void @use(i32 [[B]])
-; CHECK-NEXT:    switch i32 [[B]], label [[SW_DEFAULT:%.*]] [
+; CHECK-NEXT:    switch i32 [[A]], label [[SW_DEFAULT:%.*]] [
 ; CHECK-NEXT:      i32 0, label [[SW_BB:%.*]]
-; CHECK-NEXT:      i32 16777216, label [[SW_BB]]
-; CHECK-NEXT:      i32 -2147483648, label [[SW_BB]]
+; CHECK-NEXT:      i32 1, label [[SW_BB]]
+; CHECK-NEXT:      i32 128, label [[SW_BB]]
 ; CHECK-NEXT:    ]
 ; CHECK:       sw.bb:
 ; CHECK-NEXT:    ret i1 true
@@ -67,10 +67,10 @@ define i1 @test_switch_with_shl_nsw_multiuse(i32 %a) {
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[B:%.*]] = shl nsw i32 [[A]], 24
 ; CHECK-NEXT:    call void @use(i32 [[B]])
-; CHECK-NEXT:    switch i32 [[B]], label [[SW_DEFAULT:%.*]] [
+; CHECK-NEXT:    switch i32 [[A]], label [[SW_DEFAULT:%.*]] [
 ; CHECK-NEXT:      i32 0, label [[SW_BB:%.*]]
-; CHECK-NEXT:      i32 16777216, label [[SW_BB]]
-; CHECK-NEXT:      i32 -2147483648, label [[SW_BB]]
+; CHECK-NEXT:      i32 1, label [[SW_BB]]
+; CHECK-NEXT:      i32 -128, label [[SW_BB]]
 ; CHECK-NEXT:    ]
 ; CHECK:       sw.bb:
 ; CHECK-NEXT:    ret i1 true



More information about the llvm-commits mailing list