[clang] [llvm] [LV] Support generating masks for switch terminators. (PR #99808)

[Florian Hahn via cfe-commits]

================
@@ -7842,6 +7853,61 @@ VPValue *VPRecipeBuilder::createEdgeMask(BasicBlock *Src, BasicBlock *Dst) {
 
   VPValue *SrcMask = getBlockInMask(Src);
 
+  if (auto *SI = dyn_cast<SwitchInst>(Src->getTerminator())) {
+    assert(!OrigLoop->isLoopExiting(Src) &&
+           all_of(successors(Src),
+                  [this](BasicBlock *Succ) {
+                    return OrigLoop->getHeader() != Succ;
+                  }) &&
+           "unsupported switch either exiting loop or continuing to header");
+    // Create masks where the terminator in Src is a switch. We create mask for
+    // all edges at the same time. This is more efficient, as we can create and
+    // collect compares for all cases once.
+    VPValue *Cond = getVPValueOrAddLiveIn(SI->getCondition(), Plan);
+    BasicBlock *DefaultDst = SI->getDefaultDest();
+    MapVector<BasicBlock *, SmallVector<VPValue *>> Map;
+    for (auto &C : SI->cases()) {
+      auto I = Map.insert({C.getCaseSuccessor(), {}});
+      VPValue *V = getVPValueOrAddLiveIn(C.getCaseValue(), Plan);
+      I.first->second.push_back(Builder.createICmp(CmpInst::ICMP_EQ, Cond, V));
+    }
+
+    // We need to handle 2 separate cases:
+    // 1. Dst is not the default destination. Dst is reached if any of the cases
+    // with destination == Dst are taken. Join the conditions for each case
+    // where destination == Dst using a logical OR.
+    for (const auto &[Dst, Conds] : Map) {
+      VPValue *Mask = Conds[0];
+      for (VPValue *V : ArrayRef<VPValue *>(Conds).drop_front())
+        Mask = Builder.createOr(Mask, V);
+      if (SrcMask)
+        Mask = Builder.createLogicalAnd(SrcMask, Mask);
+      EdgeMaskCache[{Src, Dst}] = Mask;
----------------
fhahn wrote:

Updated, thanks!

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