[llvm] [SimplifyCFG] Simplify switch instruction that has duplicate arms (PR #114262)

[Michael Maitland via llvm-commits]

================
@@ -7436,6 +7437,110 @@ static bool simplifySwitchOfCmpIntrinsic(SwitchInst *SI, IRBuilderBase &Builder,
   return true;
 }
 
+bool SimplifyCFGOpt::simplifyDuplicateSwitchArms(SwitchInst *SI) {
+  // Simplify the case where multiple arms contain only a terminator, the
+  // terminators are the same, and their sucessor PHIS incoming values are the
+  // same.
+
+  // Find BBs that are candidates for simplification.
+  SmallPtrSet<BasicBlock *, 8> BBs;
+  for (auto &Case : SI->cases()) {
+    BasicBlock *BB = Case.getCaseSuccessor();
+
+    // FIXME: This case needs some extra care because the terminators other than
+    // SI need to be updated.
+    if (!BB->hasNPredecessors(1))
+      continue;
+
+    // FIXME: Relax that the terminator is a BranchInst by checking for equality
+    // on other kinds of terminators.
+    Instruction *T = BB->getTerminator();
+    if (T && isa<BranchInst>(T))
+      BBs.insert(BB);
+  }
+
+  auto IsBranchEq = [](BranchInst *A, BranchInst *B) {
+    if (A->isConditional() != B->isConditional())
+      return false;
+
+    if (A->isConditional()) {
+      // If the conditions are instructions, check equality up to commutativity.
+      // Otherwise, check that the two Values are the same.
+      Value *AC = A->getCondition();
+      Value *BC = B->getCondition();
+      auto *ACI = dyn_cast<Instruction>(AC);
+      auto *BCI = dyn_cast<Instruction>(BC);
+      if ((ACI && BCI && !areIdenticalUpToCommutativity(ACI, BCI)) && AC != BC)
+        return false;
+    }
+
+    if (A->getNumSuccessors() != B->getNumSuccessors())
+      return false;
+
+    for (unsigned I = 0; I < A->getNumSuccessors(); ++I) {
+      BasicBlock *ASucc = A->getSuccessor(I);
+      if (ASucc != B->getSuccessor(I))
+        return false;
+      // Need to check that PHIs in sucessors have matching values
+      for (PHINode &Phi : ASucc->phis())
+        if (Phi.getIncomingValueForBlock(A->getParent()) !=
+            Phi.getIncomingValueForBlock(B->getParent()))
+          return false;
+    }
+
+    return true;
+  };
+
+  auto IsBBEqualTo = [&IsBranchEq](BasicBlock *A, BasicBlock *B) {
+    // FIXME: Support more than just a single BranchInst. One way we could do
+    // this is by taking a hashing approach.
+    if (A->size() != 1 || B->size() != 1)
+      return false;
+
+    return IsBranchEq(cast<BranchInst>(A->getTerminator()),
+                      cast<BranchInst>(B->getTerminator()));
+  };
+
+  // Construct a map from candidate basic block to an equivalent basic block
+  // to replace it with. All equivalent basic blocks should be replaced with
+  // the same basic block. To do this, if there is no equivalent BB in the map,
+  // then insert into the map BB -> BB. Otherwise, we should check only elements
+  // in the map for equivalence to ensure that all equivalent BB get replaced
+  // by the BB in the map. Replacing BB with BB has no impact, so we skip
+  // a call to setSuccessor when we do the actual replacement.
+  DenseMap<BasicBlock *, BasicBlock *> ReplaceWith;
+  for (BasicBlock *BB : BBs) {
+    bool Inserted = false;
+    for (auto KV : ReplaceWith) {
+      if (IsBBEqualTo(BB, KV.first)) {
+        ReplaceWith[BB] = KV.first;
+        Inserted = true;
+        break;
+      }
+    }
+    if (!Inserted)
+      ReplaceWith[BB] = BB;
+  }
----------------
michaelmaitland wrote:

Could you please help me understand why you propose this as a suggestion? Do you think the current approach is incorrect, slow, uses too much memory, or is overcomplicated?

I tried to implement your suggestion, but I'm definitely misunderstanding what you mean.

This is what I *think* you are saying:

```
DenseMap<BasicBlock *, SmallVector<BasicBlock *>> ReplaceWith;
// Snip: create a map where values are the basic blocks that should 
// be replaced with the key basic block. Must implement
// DenseMapInfo to make it work.

// Iterate over all key BBs in ReplaceWith and replace all items in the
// SmallVector with the key BB.
for (auto KV : ReplaceWith) 
  for (BasicBlock *BB : KV.second)
    for (&Case : SI->cases())
      if (Case.getCaseSuccessor() = BB)
        Case.setSuccessor(KV.first);
```

Again, there is no way to get BasicBlock -> Cases mapping without looping over all the cases and checking that the CaseSuccessor is equal to BB. This approach is less performant because we have to iterate over cases for every (KV, BB) pair. In the current approach, we have O(1) lookup  -> ReplaceWith[Case.getCaseSuccessor].

We could change the loop to:

```
for (&Case : SI->cases()) {
  for (auto KV : ReplaceWith)
    for (BasicBlock *BB : KV.second)
      if (Case.getCaseSuccessor() == BB)
        Case.setSuccessor(KV.first);
```

In this case, we still have worse performance since we don't have O(1) lookup.

https://github.com/llvm/llvm-project/pull/114262