[InstCombine] Simplify switch with selects #84143

dtcxzyw · 2024-03-06T08:59:37Z

An example from https://github.com/image-rs/image:

define void @test_ult_rhsc(i8 %x) {
  %val = add nsw i8 %x, -2
  %cmp = icmp ult i8 %val, 11
  %cond = select i1 %cmp, i8 %val, i8 6
  switch i8 %cond, label %bb1 [
  i8 0, label %bb2
  i8 10, label %bb3
  ]

bb1:
  call void @func1()
  unreachable
bb2:
  call void @func2()
  unreachable
bb3:
  call void @func3()
  unreachable
}

When %cmp evaluates to false, we can prove that the range of %val is [11, umax]. Thus we can safely replace %cond with %val since both switch 6 and switch %val go to the default dest %bb1.

Alive2: https://alive2.llvm.org/ce/z/uSTj6w
Godbolt: https://godbolt.org/z/MGrG84bzr

This patch will benefit many rust applications and some C/C++ applications (e.g., cvc5).

llvmbot · 2024-03-06T09:00:04Z

@llvm/pr-subscribers-llvm-transforms

Author: Yingwei Zheng (dtcxzyw)

Changes

An example from https://github.com/image-rs/image:

define void @<!-- -->test_ult_rhsc(i8 %x) {
  %val = add nsw i8 %x, -2
  %cmp = icmp ult i8 %val, 11
  %cond = select i1 %cmp, i8 %val, i8 6
  switch i8 %cond, label %bb1 [
  i8 0, label %bb2
  i8 10, label %bb3
  ]

bb1:
  call void @<!-- -->func1()
  unreachable
bb2:
  call void @<!-- -->func2()
  unreachable
bb3:
  call void @<!-- -->func3()
  unreachable
}

When %cmp evaluates to false, we can prove that the range of %val is [11, umax]. Thus we can safely replace %cond with %val since both switch 6 and switch %val go to the default dest %bb1.

Alive2: https://alive2.llvm.org/ce/z/uSTj6w

This patch will benefit many rust applications and some C/C++ applications (e.g., cvc5).

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

2 Files Affected:

(modified) llvm/lib/Transforms/InstCombine/InstructionCombining.cpp (+41)
(added) llvm/test/Transforms/InstCombine/switch-select.ll (+159)

diff --git a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
index 80ce0c9275b2cb..3d231f9fc7ff4b 100644
--- a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
@@ -3334,6 +3334,37 @@ Instruction *InstCombinerImpl::visitBranchInst(BranchInst &BI) {
   return nullptr;
 }
 
+// Replaces (switch (select cond, X, C)/(select cond, C, X)) with (switch X) if
+// we can prove that both (switch C) and (switch X) go to the default when cond
+// is false/true.
+static Value *simplifySwitchOnSelectUsingRanges(SwitchInst &SI,
+                                                SelectInst *Select,
+                                                unsigned CstOpIdx) {
+  auto *C = dyn_cast<ConstantInt>(Select->getOperand(CstOpIdx));
+  if (!C)
+    return nullptr;
+
+  BasicBlock *CstBB = SI.findCaseValue(C)->getCaseSuccessor();
+  if (CstBB != SI.getDefaultDest())
+    return nullptr;
+  Value *X = Select->getOperand(3 - CstOpIdx);
+  ICmpInst::Predicate Pred;
+  const APInt *RHSC;
+  if (!match(Select->getCondition(),
+             m_ICmp(Pred, m_Specific(X), m_APInt(RHSC))))
+    return nullptr;
+  if (CstOpIdx == 1)
+    Pred = ICmpInst::getInversePredicate(Pred);
+
+  // See whether we can replace the select with X
+  ConstantRange CR = ConstantRange::makeExactICmpRegion(Pred, *RHSC);
+  for (auto Case : SI.cases())
+    if (!CR.contains(Case.getCaseValue()->getValue()))
+      return nullptr;
+
+  return X;
+}
+
 Instruction *InstCombinerImpl::visitSwitchInst(SwitchInst &SI) {
   Value *Cond = SI.getCondition();
   Value *Op0;
@@ -3407,6 +3438,16 @@ Instruction *InstCombinerImpl::visitSwitchInst(SwitchInst &SI) {
     }
   }
 
+  // Fold switch(select cond, X, Y) into switch(X/Y) if possible
+  if (auto *Select = dyn_cast<SelectInst>(Cond)) {
+    if (Value *V =
+            simplifySwitchOnSelectUsingRanges(SI, Select, /*CstOpIdx=*/1))
+      return replaceOperand(SI, 0, V);
+    if (Value *V =
+            simplifySwitchOnSelectUsingRanges(SI, Select, /*CstOpIdx=*/2))
+      return replaceOperand(SI, 0, V);
+  }
+
   KnownBits Known = computeKnownBits(Cond, 0, &SI);
   unsigned LeadingKnownZeros = Known.countMinLeadingZeros();
   unsigned LeadingKnownOnes = Known.countMinLeadingOnes();
diff --git a/llvm/test/Transforms/InstCombine/switch-select.ll b/llvm/test/Transforms/InstCombine/switch-select.ll
new file mode 100644
index 00000000000000..60757c5d22527f
--- /dev/null
+++ b/llvm/test/Transforms/InstCombine/switch-select.ll
@@ -0,0 +1,159 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 4
+; RUN: opt < %s -passes=instcombine -S | FileCheck %s
+
+define void @test_ult_rhsc(i8 %x) {
+; CHECK-LABEL: define void @test_ult_rhsc(
+; CHECK-SAME: i8 [[X:%.*]]) {
+; CHECK-NEXT:    switch i8 [[X]], label [[BB1:%.*]] [
+; CHECK-NEXT:      i8 2, label [[BB2:%.*]]
+; CHECK-NEXT:      i8 12, label [[BB3:%.*]]
+; CHECK-NEXT:    ]
+; CHECK:       bb1:
+; CHECK-NEXT:    call void @func1()
+; CHECK-NEXT:    unreachable
+; CHECK:       bb2:
+; CHECK-NEXT:    call void @func2()
+; CHECK-NEXT:    unreachable
+; CHECK:       bb3:
+; CHECK-NEXT:    call void @func3()
+; CHECK-NEXT:    unreachable
+;
+  %val = add nsw i8 %x, -2
+  %cmp = icmp ult i8 %val, 11
+  %cond = select i1 %cmp, i8 %val, i8 6
+  switch i8 %cond, label %bb1 [
+  i8 0, label %bb2
+  i8 10, label %bb3
+  ]
+
+bb1:
+  call void @func1()
+  unreachable
+bb2:
+  call void @func2()
+  unreachable
+bb3:
+  call void @func3()
+  unreachable
+}
+
+define void @test_eq_lhsc(i8 %x) {
+; CHECK-LABEL: define void @test_eq_lhsc(
+; CHECK-SAME: i8 [[X:%.*]]) {
+; CHECK-NEXT:    switch i8 [[X]], label [[BB1:%.*]] [
+; CHECK-NEXT:      i8 0, label [[BB2:%.*]]
+; CHECK-NEXT:      i8 10, label [[BB3:%.*]]
+; CHECK-NEXT:    ]
+; CHECK:       bb1:
+; CHECK-NEXT:    call void @func1()
+; CHECK-NEXT:    unreachable
+; CHECK:       bb2:
+; CHECK-NEXT:    call void @func2()
+; CHECK-NEXT:    unreachable
+; CHECK:       bb3:
+; CHECK-NEXT:    call void @func3()
+; CHECK-NEXT:    unreachable
+;
+  %cmp = icmp eq i8 %x, 4
+  %cond = select i1 %cmp, i8 6, i8 %x
+  switch i8 %cond, label %bb1 [
+  i8 0, label %bb2
+  i8 10, label %bb3
+  ]
+
+bb1:
+  call void @func1()
+  unreachable
+bb2:
+  call void @func2()
+  unreachable
+bb3:
+  call void @func3()
+  unreachable
+}
+
+define void @test_ult_rhsc_invalid_cond(i8 %x, i8 %y) {
+; CHECK-LABEL: define void @test_ult_rhsc_invalid_cond(
+; CHECK-SAME: i8 [[X:%.*]], i8 [[Y:%.*]]) {
+; CHECK-NEXT:    [[VAL:%.*]] = add nsw i8 [[X]], -2
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i8 [[Y]], 11
+; CHECK-NEXT:    [[COND:%.*]] = select i1 [[CMP]], i8 [[VAL]], i8 6
+; CHECK-NEXT:    switch i8 [[COND]], label [[BB1:%.*]] [
+; CHECK-NEXT:      i8 0, label [[BB2:%.*]]
+; CHECK-NEXT:      i8 10, label [[BB3:%.*]]
+; CHECK-NEXT:      i8 13, label [[BB3]]
+; CHECK-NEXT:    ]
+; CHECK:       bb1:
+; CHECK-NEXT:    call void @func1()
+; CHECK-NEXT:    unreachable
+; CHECK:       bb2:
+; CHECK-NEXT:    call void @func2()
+; CHECK-NEXT:    unreachable
+; CHECK:       bb3:
+; CHECK-NEXT:    call void @func3()
+; CHECK-NEXT:    unreachable
+;
+  %val = add nsw i8 %x, -2
+  %cmp = icmp ult i8 %y, 11
+  %cond = select i1 %cmp, i8 %val, i8 6
+  switch i8 %cond, label %bb1 [
+  i8 0, label %bb2
+  i8 10, label %bb3
+  i8 13, label %bb3
+  ]
+
+bb1:
+  call void @func1()
+  unreachable
+bb2:
+  call void @func2()
+  unreachable
+bb3:
+  call void @func3()
+  unreachable
+}
+
+define void @test_ult_rhsc_fail(i8 %x) {
+; CHECK-LABEL: define void @test_ult_rhsc_fail(
+; CHECK-SAME: i8 [[X:%.*]]) {
+; CHECK-NEXT:    [[VAL:%.*]] = add nsw i8 [[X]], -2
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i8 [[VAL]], 11
+; CHECK-NEXT:    [[COND:%.*]] = select i1 [[CMP]], i8 [[VAL]], i8 6
+; CHECK-NEXT:    switch i8 [[COND]], label [[BB1:%.*]] [
+; CHECK-NEXT:      i8 0, label [[BB2:%.*]]
+; CHECK-NEXT:      i8 10, label [[BB3:%.*]]
+; CHECK-NEXT:      i8 13, label [[BB3]]
+; CHECK-NEXT:    ]
+; CHECK:       bb1:
+; CHECK-NEXT:    call void @func1()
+; CHECK-NEXT:    unreachable
+; CHECK:       bb2:
+; CHECK-NEXT:    call void @func2()
+; CHECK-NEXT:    unreachable
+; CHECK:       bb3:
+; CHECK-NEXT:    call void @func3()
+; CHECK-NEXT:    unreachable
+;
+  %val = add nsw i8 %x, -2
+  %cmp = icmp ult i8 %val, 11
+  %cond = select i1 %cmp, i8 %val, i8 6
+  switch i8 %cond, label %bb1 [
+  i8 0, label %bb2
+  i8 10, label %bb3
+  i8 13, label %bb3
+  ]
+
+bb1:
+  call void @func1()
+  unreachable
+bb2:
+  call void @func2()
+  unreachable
+bb3:
+  call void @func3()
+  unreachable
+}
+
+declare void @func1()
+declare void @func2()
+declare void @func3()

PR Link: llvm/llvm-project#84143

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

goldsteinn · 2024-03-06T17:15:23Z

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

+  ICmpInst::Predicate Pred;
+  const APInt *RHSC;
+  if (!match(Select->getCondition(),
+             m_ICmp(Pred, m_Specific(X), m_APInt(RHSC))))


Is there a reason to only handle RHSC?

We can improve this by using KnownBits/ConstantRange in the future. But I guess this patch covers most of the cases in rust applications.

Think you could also technically generalize to RHSC (or Other arm really) implying any single case and not condition implying the same condition for the arm, although not sure if that would practically ever come up. Think we would just simplify the select.

dianqk

LGTM

dtcxzyw · 2024-04-11T14:16:42Z

Ping.

goldsteinn · 2024-04-11T15:37:38Z

LGTM to me as well.

nikic

LGTM

@func3

An example from https://github.com/image-rs/image: ``` define void @test_ult_rhsc(i8 %x) { %val = add nsw i8 %x, -2 %cmp = icmp ult i8 %val, 11 %cond = select i1 %cmp, i8 %val, i8 6 switch i8 %cond, label %bb1 [ i8 0, label %bb2 i8 10, label %bb3 ] bb1: call void @func1() unreachable bb2: call void @func2() unreachable bb3: call void @func3() unreachable } ``` When `%cmp` evaluates to false, we can prove that the range of `%val` is [11, umax]. Thus we can safely replace `%cond` with `%val` since both `switch 6` and `switch %val` go to the default dest `%bb1`. Alive2: https://alive2.llvm.org/ce/z/uSTj6w Godbolt: https://godbolt.org/z/MGrG84bzr This patch will benefit many rust applications and some C/C++ applications (e.g., cvc5).

@func3

An example from https://github.com/image-rs/image: ``` define void @test_ult_rhsc(i8 %x) { %val = add nsw i8 %x, -2 %cmp = icmp ult i8 %val, 11 %cond = select i1 %cmp, i8 %val, i8 6 switch i8 %cond, label %bb1 [ i8 0, label %bb2 i8 10, label %bb3 ] bb1: call void @func1() unreachable bb2: call void @func2() unreachable bb3: call void @func3() unreachable } ``` When `%cmp` evaluates to false, we can prove that the range of `%val` is [11, umax]. Thus we can safely replace `%cond` with `%val` since both `switch 6` and `switch %val` go to the default dest `%bb1`. Alive2: https://alive2.llvm.org/ce/z/uSTj6w Godbolt: https://godbolt.org/z/MGrG84bzr This patch will benefit many rust applications and some C/C++ applications (e.g., cvc5).

dtcxzyw requested review from dianqk and goldsteinn March 6, 2024 08:59

dtcxzyw requested a review from nikic as a code owner March 6, 2024 08:59

llvmbot added the llvm:transforms label Mar 6, 2024

dtcxzyw added a commit to dtcxzyw/llvm-opt-benchmark that referenced this pull request Mar 6, 2024

pre-commit: test PR84143

3347467

PR Link: llvm/llvm-project#84143

dtcxzyw mentioned this pull request Mar 6, 2024

pre-commit: test PR84143 dtcxzyw/llvm-opt-benchmark#294

Closed

goldsteinn reviewed Mar 6, 2024

View reviewed changes

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp Outdated Show resolved Hide resolved

goldsteinn reviewed Mar 6, 2024

View reviewed changes

dtcxzyw added 3 commits March 22, 2024 14:30

[InstCombine] Add pre-commit tests. NFC.

0c4c262

[InstCombine] Simplify switch with selects

28ed09c

[InstCombine] Address comments.

1828c72

dtcxzyw force-pushed the perf/switch-on-select branch from f782c5e to 1828c72 Compare March 22, 2024 07:14

dianqk reviewed Mar 22, 2024

View reviewed changes

nikic approved these changes Apr 15, 2024

View reviewed changes

dtcxzyw merged commit 5fe1466 into llvm:main Apr 15, 2024

dtcxzyw deleted the perf/switch-on-select branch April 15, 2024 08:40

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[InstCombine] Simplify switch with selects #84143

[InstCombine] Simplify switch with selects #84143

Uh oh!

dtcxzyw commented Mar 6, 2024 •

edited

Loading

Uh oh!

llvmbot commented Mar 6, 2024

Uh oh!

Uh oh!

goldsteinn Mar 6, 2024

Uh oh!

dtcxzyw Mar 6, 2024

Uh oh!

goldsteinn Mar 6, 2024

Uh oh!

dianqk left a comment

Uh oh!

dtcxzyw commented Apr 11, 2024

Uh oh!

goldsteinn commented Apr 11, 2024

Uh oh!

nikic left a comment

Uh oh!

Uh oh!

[InstCombine] Simplify switch with selects #84143

[InstCombine] Simplify switch with selects #84143

Uh oh!

Conversation

dtcxzyw commented Mar 6, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

llvmbot commented Mar 6, 2024

Uh oh!

Uh oh!

goldsteinn Mar 6, 2024

Choose a reason for hiding this comment

Uh oh!

dtcxzyw Mar 6, 2024

Choose a reason for hiding this comment

Uh oh!

goldsteinn Mar 6, 2024

Choose a reason for hiding this comment

Uh oh!

dianqk left a comment

Choose a reason for hiding this comment

Uh oh!

dtcxzyw commented Apr 11, 2024

Uh oh!

goldsteinn commented Apr 11, 2024

Uh oh!

nikic left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

dtcxzyw commented Mar 6, 2024 •

edited

Loading