[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 01:58:12 PST 2024
https://github.com/dtcxzyw created https://github.com/llvm/llvm-project/pull/77068
This patch canonicalizes `switch(X << C)` to `switch(X)`. If the shift may wrap, an and instruction will be created to mask out all of the shifted bits.
Alive2: https://alive2.llvm.org/ce/z/wSsL5y
NOTE: We can relax the one-use constraint. But I don't see any benefit in my benchmark.
>From eb1f93f044f8851f6d8d835592e15bb3cfc41086 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 634ef0786f8066b71b374459278aa5a00d59c472 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 f3181dc14792c8..e4a2c4b66c5b5a 100644
--- a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
@@ -3208,6 +3208,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