[llvm] [InstCombine] Canonicalize `switch(X << C)` into `switch(X)` (PR #77068)
via llvm-commits
llvm-commits at lists.llvm.org
Fri Jan 5 01:58:38 PST 2024
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-llvm-transforms
Author: Yingwei Zheng (dtcxzyw)
<details>
<summary>Changes</summary>
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.
---
Full diff: https://github.com/llvm/llvm-project/pull/77068.diff
2 Files Affected:
- (modified) llvm/lib/Transforms/InstCombine/InstructionCombining.cpp (+27)
- (added) llvm/test/Transforms/InstCombine/switch-shl.ll (+185)
``````````diff
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
new file mode 100644
index 00000000000000..037ef37f97a31a
--- /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: [[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
+; 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 [[A]], label [[SW_DEFAULT:%.*]] [
+; CHECK-NEXT: i32 0, 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
+; 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 [[A]], label [[SW_DEFAULT:%.*]] [
+; CHECK-NEXT: i32 0, 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
+; 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)
``````````
</details>
https://github.com/llvm/llvm-project/pull/77068
More information about the llvm-commits
mailing list