[SimplifyCFG] Consider preds to switch in `simplifyDuplicateSwitchArms` #118955

antoniofrighetto · 2024-12-06T11:39:00Z

Allow a duplicate basic block with multiple predecessors to the jump table to be simplified, by considering that the same basic block may appear in more switch cases.

llvmbot · 2024-12-06T11:39:31Z

@llvm/pr-subscribers-llvm-transforms

Author: Antonio Frighetto (antoniofrighetto)

Changes

Allow a duplicate basic block with multiple predecessors to the jump table to be simplified, by considering that the same basic block may appear in more switch cases.

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

3 Files Affected:

(modified) llvm/lib/Transforms/Utils/SimplifyCFG.cpp (+12-7)
(modified) llvm/test/Transforms/SimplifyCFG/X86/switch_to_lookup_table.ll (+21-13)
(modified) llvm/test/Transforms/SimplifyCFG/switch-dup-bbs.ll (+32)

diff --git a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
index c7e814bced57db..162e60c0c58f7f 100644
--- a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
+++ b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
@@ -7474,9 +7474,6 @@ static bool simplifySwitchOfCmpIntrinsic(SwitchInst *SI, IRBuilderBase &Builder,
 /// IncomingValue and add it in the Wrapper so isEqual can do O(1) checking
 /// of the incoming values.
 struct SwitchSuccWrapper {
-  // Keep so we can use SwitchInst::setSuccessor to do the replacement. It won't
-  // be important to equality though.
-  unsigned SuccNum;
   BasicBlock *Dest;
   DenseMap<PHINode *, SmallDenseMap<BasicBlock *, Value *, 8>> *PhiPredIVs;
 };
@@ -7563,6 +7560,7 @@ bool SimplifyCFGOpt::simplifyDuplicateSwitchArms(SwitchInst *SI,
   SmallPtrSet<PHINode *, 8> Phis;
   SmallPtrSet<BasicBlock *, 8> Seen;
   DenseMap<PHINode *, SmallDenseMap<BasicBlock *, Value *, 8>> PhiPredIVs;
+  DenseMap<BasicBlock *, SmallVector<unsigned, 4>> BBToSuccessorIndexes;
   SmallVector<SwitchSuccWrapper> Cases;
   Cases.reserve(SI->getNumSuccessors());
 
@@ -7575,8 +7573,10 @@ bool SimplifyCFGOpt::simplifyDuplicateSwitchArms(SwitchInst *SI,
       continue;
 
     // FIXME: This case needs some extra care because the terminators other than
-    // SI need to be updated.
-    if (BB->hasNPredecessorsOrMore(2))
+    // SI need to be updated. For now, consider only backedges to the SI.
+    if (pred_size(BB) > 4 || any_of(predecessors(BB), [&](BasicBlock *Pred) {
+          return Pred != SI->getParent();
+        }))
       continue;
 
     // FIXME: Relax that the terminator is a BranchInst by checking for equality
@@ -7591,8 +7591,11 @@ bool SimplifyCFGOpt::simplifyDuplicateSwitchArms(SwitchInst *SI,
       for (BasicBlock *Succ : BI->successors())
         for (PHINode &Phi : Succ->phis())
           Phis.insert(&Phi);
+      // Add the successor only if not previously visited.
+      Cases.emplace_back(SwitchSuccWrapper{BB, &PhiPredIVs});
     }
-    Cases.emplace_back(SwitchSuccWrapper{I, BB, &PhiPredIVs});
+
+    BBToSuccessorIndexes[BB].emplace_back(I);
   }
 
   // Precompute a data structure to improve performance of isEqual for
@@ -7627,7 +7630,9 @@ bool SimplifyCFGOpt::simplifyDuplicateSwitchArms(SwitchInst *SI,
       // We know that SI's parent BB no longer dominates the old case successor
       // since we are making it dead.
       Updates.push_back({DominatorTree::Delete, SI->getParent(), SSW.Dest});
-      SI->setSuccessor(SSW.SuccNum, (*It)->Dest);
+      const auto &Successors = BBToSuccessorIndexes.at(SSW.Dest);
+      for (unsigned Idx : Successors)
+        SI->setSuccessor(Idx, (*It)->Dest);
       MadeChange = true;
     }
   }
diff --git a/llvm/test/Transforms/SimplifyCFG/X86/switch_to_lookup_table.ll b/llvm/test/Transforms/SimplifyCFG/X86/switch_to_lookup_table.ll
index 9549ccdbfe9ec4..7f484e2ec29d7d 100644
--- a/llvm/test/Transforms/SimplifyCFG/X86/switch_to_lookup_table.ll
+++ b/llvm/test/Transforms/SimplifyCFG/X86/switch_to_lookup_table.ll
@@ -314,20 +314,20 @@ lor.end:
 define i32 @overflow(i32 %type) {
 ; CHECK-LABEL: @overflow(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    switch i32 [[TYPE:%.*]], label [[IF_END:%.*]] [
+; CHECK-NEXT:    switch i32 [[TYPE:%.*]], label [[SW_DEFAULT:%.*]] [
 ; CHECK-NEXT:      i32 3, label [[SW_BB3:%.*]]
 ; CHECK-NEXT:      i32 -2147483645, label [[SW_BB3]]
-; CHECK-NEXT:      i32 1, label [[SW_BB1:%.*]]
+; CHECK-NEXT:      i32 1, label [[IF_END:%.*]]
 ; CHECK-NEXT:      i32 2, label [[SW_BB2:%.*]]
 ; CHECK-NEXT:    ]
-; CHECK:       sw.bb1:
-; CHECK-NEXT:    br label [[IF_END]]
 ; CHECK:       sw.bb2:
 ; CHECK-NEXT:    br label [[IF_END]]
 ; CHECK:       sw.bb3:
 ; CHECK-NEXT:    br label [[IF_END]]
+; CHECK:       sw.default:
+; CHECK-NEXT:    br label [[IF_END]]
 ; CHECK:       if.end:
-; CHECK-NEXT:    [[DIRENT_TYPE_0:%.*]] = phi i32 [ 6, [[SW_BB3]] ], [ 5, [[SW_BB2]] ], [ 0, [[SW_BB1]] ], [ 3, [[ENTRY:%.*]] ]
+; CHECK-NEXT:    [[DIRENT_TYPE_0:%.*]] = phi i32 [ 3, [[SW_DEFAULT]] ], [ 6, [[SW_BB3]] ], [ 5, [[SW_BB2]] ], [ 0, [[ENTRY:%.*]] ]
 ; CHECK-NEXT:    ret i32 [[DIRENT_TYPE_0]]
 ;
 entry:
@@ -340,15 +340,23 @@ entry:
   i32 3, label %sw.bb3
   ]
 
-sw.bb: br label %if.end
-sw.bb1: br label %if.end
-sw.bb2: br label %if.end
-sw.bb3: br label %if.end
-sw.default: br label %if.end
-if.else: br label %if.end
+sw.bb:                                            ; preds = %entry, %entry
+  br label %if.end
 
-if.end:
-  %dirent_type.0 = phi i32 [ 3, %sw.default ], [ 6, %sw.bb3 ], [ 5, %sw.bb2 ], [ 0, %sw.bb1 ], [ 3, %sw.bb ], [ 0, %if.else ]
+sw.bb1:                                           ; preds = %entry
+  br label %if.end
+
+sw.bb2:                                           ; preds = %entry
+  br label %if.end
+
+sw.bb3:                                           ; preds = %entry, %entry
+  br label %if.end
+
+sw.default:                                       ; preds = %entry
+  br label %if.end
+
+if.end:                                           ; preds = %sw.default, %sw.bb3, %sw.bb2, %sw.bb1, %sw.bb
+  %dirent_type.0 = phi i32 [ 3, %sw.default ], [ 6, %sw.bb3 ], [ 5, %sw.bb2 ], [ 0, %sw.bb1 ], [ 3, %sw.bb ]
   ret i32 %dirent_type.0
 }
 
diff --git a/llvm/test/Transforms/SimplifyCFG/switch-dup-bbs.ll b/llvm/test/Transforms/SimplifyCFG/switch-dup-bbs.ll
index 0df4deba1a156f..d9590252308e60 100644
--- a/llvm/test/Transforms/SimplifyCFG/switch-dup-bbs.ll
+++ b/llvm/test/Transforms/SimplifyCFG/switch-dup-bbs.ll
@@ -127,3 +127,35 @@ define i32 @switch_dup_default(i32 %0, i32 %1, i32 %2, i32 %3) {
   %9 = phi i32 [ %3, %5 ], [ %2, %6 ], [ %2, %7 ]
   ret i32 %9
 }
+
+define i64 @switch_dup_exit(i32 %val) {
+; SIMPLIFY-CFG-LABEL: define i64 @switch_dup_exit(
+; SIMPLIFY-CFG-SAME: i32 [[VAL:%.*]]) {
+; SIMPLIFY-CFG-NEXT:  [[ENTRY:.*:]]
+; SIMPLIFY-CFG-NEXT:    [[SWITCH_SELECTCMP_CASE1:%.*]] = icmp eq i32 [[VAL]], 1
+; SIMPLIFY-CFG-NEXT:    [[SWITCH_SELECTCMP_CASE2:%.*]] = icmp eq i32 [[VAL]], 11
+; SIMPLIFY-CFG-NEXT:    [[SWITCH_SELECTCMP:%.*]] = or i1 [[SWITCH_SELECTCMP_CASE1]], [[SWITCH_SELECTCMP_CASE2]]
+; SIMPLIFY-CFG-NEXT:    [[TMP0:%.*]] = select i1 [[SWITCH_SELECTCMP]], i64 1, i64 0
+; SIMPLIFY-CFG-NEXT:    ret i64 [[TMP0]]
+;
+entry:
+  switch i32 %val, label %default [
+  i32 1, label %bb2
+  i32 11, label %exit
+  i32 13, label %bb1
+  i32 0, label %bb1
+  ]
+
+bb1:                                              ; preds = %entry, %entry
+  br label %exit
+
+bb2:                                              ; preds = %entry
+  br label %exit
+
+default:                                          ; preds = %entry
+  br label %exit
+
+exit:                                             ; preds = %entry, %default, %bb2, %bb1
+  %ret = phi i64 [ 0, %default ], [ 0, %bb1 ], [ 1, %entry ], [ 1, %bb2 ]
+  ret i64 %ret
+}

dtcxzyw

LGTM. Please wait for additional approval from other reviewers.

dtcxzyw · 2024-12-06T14:30:53Z

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

-    // SI need to be updated.
-    if (BB->hasNPredecessorsOrMore(2))
+    // SI need to be updated. For now, consider only backedges to the SI.
+    if (pred_size(BB) > 4 || any_of(predecessors(BB), [&](BasicBlock *Pred) {


Use hasNPredecessorsOrMore(4) and BB->getUniquePredecessor() == SI->getParent().

What is special about having >= 4 predecessors?

It is used to avoid O(N^2) lookup. We may choose a larger threshold if the compile-time impact is negligible.

Correct, this way we keep the time complexity constant. Happy to tune this differently if compile-time remains fine.

michaelmaitland · 2024-12-06T15:05:03Z

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

-    // SI need to be updated.
-    if (BB->hasNPredecessorsOrMore(2))
+    // SI need to be updated. For now, consider only backedges to the SI.
+    if (pred_size(BB) > 4 || any_of(predecessors(BB), [&](BasicBlock *Pred) {


What is special about having >= 4 predecessors?

michaelmaitland · 2024-12-06T15:06:15Z

llvm/test/Transforms/SimplifyCFG/X86/switch_to_lookup_table.ll

-sw.bb3: br label %if.end
-sw.default: br label %if.end
-if.else: br label %if.end
+sw.bb:                                            ; preds = %entry, %entry


Did you mean to change these lines?

I took the chance to regenerate the test too, as the if.else branch was dead.

nikic

Can you please pre-commit the tests and add a test where the transform triggers but does not convert into a select afterwards (only reduces blocks)?

michaelmaitland

LGTM

Allow a duplicate basic block with multiple predecessors to the jump table to be simplified, by considering that the same basic block may appear in more switch cases.

antoniofrighetto requested review from XChy, dtcxzyw, michaelmaitland, mshockwave and nikic December 6, 2024 11:39

llvmbot added the llvm:transforms label Dec 6, 2024

This was referenced Dec 6, 2024

Task submission dtcxzyw/llvm-opt-benchmark#1312

Open

pre-commit: PR118955 dtcxzyw/llvm-opt-benchmark#1831

Closed

dtcxzyw approved these changes Dec 6, 2024

View reviewed changes

michaelmaitland reviewed Dec 6, 2024

View reviewed changes

nikic reviewed Dec 11, 2024

View reviewed changes

antoniofrighetto force-pushed the feature/scfg-extend-dup-arms branch from c04a12e to 727ae16 Compare December 11, 2024 18:44

michaelmaitland approved these changes Dec 11, 2024

View reviewed changes

antoniofrighetto added 2 commits December 13, 2024 09:07

[SimplifyCFG] Precommit tests for PR118955 (NFC)

e32c428

[SimplifyCFG] Consider preds to switch in simplifyDuplicateSwitchArms

d26df32

Allow a duplicate basic block with multiple predecessors to the jump table to be simplified, by considering that the same basic block may appear in more switch cases.

antoniofrighetto force-pushed the feature/scfg-extend-dup-arms branch from 727ae16 to d26df32 Compare December 13, 2024 08:07

antoniofrighetto merged commit d26df32 into llvm:main Dec 13, 2024
5 of 7 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[SimplifyCFG] Consider preds to switch in `simplifyDuplicateSwitchArms` #118955

[SimplifyCFG] Consider preds to switch in `simplifyDuplicateSwitchArms` #118955

Uh oh!

antoniofrighetto commented Dec 6, 2024

Uh oh!

llvmbot commented Dec 6, 2024

Uh oh!

dtcxzyw left a comment

Uh oh!

dtcxzyw Dec 6, 2024

Uh oh!

michaelmaitland Dec 6, 2024

Uh oh!

dtcxzyw Dec 6, 2024

Uh oh!

antoniofrighetto Dec 10, 2024

Uh oh!

michaelmaitland Dec 6, 2024

Uh oh!

michaelmaitland Dec 6, 2024

Uh oh!

antoniofrighetto Dec 10, 2024

Uh oh!

nikic left a comment

Uh oh!

michaelmaitland left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

[SimplifyCFG] Consider preds to switch in simplifyDuplicateSwitchArms #118955

[SimplifyCFG] Consider preds to switch in simplifyDuplicateSwitchArms #118955

Uh oh!

Conversation

antoniofrighetto commented Dec 6, 2024

Uh oh!

llvmbot commented Dec 6, 2024

Uh oh!

dtcxzyw left a comment

Choose a reason for hiding this comment

Uh oh!

dtcxzyw Dec 6, 2024

Choose a reason for hiding this comment

Uh oh!

michaelmaitland Dec 6, 2024

Choose a reason for hiding this comment

Uh oh!

dtcxzyw Dec 6, 2024

Choose a reason for hiding this comment

Uh oh!

antoniofrighetto Dec 10, 2024

Choose a reason for hiding this comment

Uh oh!

michaelmaitland Dec 6, 2024

Choose a reason for hiding this comment

Uh oh!

michaelmaitland Dec 6, 2024

Choose a reason for hiding this comment

Uh oh!

antoniofrighetto Dec 10, 2024

Choose a reason for hiding this comment

Uh oh!

nikic left a comment

Choose a reason for hiding this comment

Uh oh!

michaelmaitland left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

[SimplifyCFG] Consider preds to switch in `simplifyDuplicateSwitchArms` #118955

[SimplifyCFG] Consider preds to switch in `simplifyDuplicateSwitchArms` #118955