!fixup address comments, thanks!

Author: fhahn

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

@@ -1339,9 +1339,18 @@ bool LoopVectorizationLegality::canVectorizeWithIfConvert() {
 
   // Collect the blocks that need predication.
   for (BasicBlock *BB : TheLoop->blocks()) {
-    // We don't support switch statements inside loops.
-    if (!isa<BranchInst, SwitchInst>(BB->getTerminator())) {
-      reportVectorizationFailure("Loop contains an unsupported termaintor",
+    // We support only branches and switch statements as terminators inside the
+    // loop.
+    if (isa<SwitchInst>(BB->getTerminator())) {
+      if (TheLoop->isLoopExiting(BB)) {
+        reportVectorizationFailure("Loop contains an unsupported switch",
+                                   "loop contains an unsupported switch",
+                                   "LoopContainsUnsupportedSwitch", ORE,
+                                   TheLoop, BB->getTerminator());
+        return false;
+      }
+    } else if (!isa<BranchInst>(BB->getTerminator())) {
+      reportVectorizationFailure("Loop contains an unsupported terminator",
                                  "loop contains an unsupported terminator",
                                  "LoopContainsUnsupportedTerminator", ORE,
                                  TheLoop, BB->getTerminator());

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -6456,6 +6456,17 @@ LoopVectorizationCostModel::getInstructionCost(Instruction *I,
     // a predicated block since it will become a fall-through, although we
     // may decide in the future to call TTI for all branches.
   }
+  case Instruction::Switch: {
+    if (VF.isScalar())
+      return TTI.getCFInstrCost(Instruction::Switch, CostKind);
+    auto *Switch = cast<SwitchInst>(I);
+    return Switch->getNumCases() *
+           TTI.getCmpSelInstrCost(
+               Instruction::ICmp,
+               ToVectorTy(Switch->getCondition()->getType(), VF),
+               ToVectorTy(Type::getInt1Ty(I->getContext()), VF),
+               CmpInst::BAD_ICMP_PREDICATE, CostKind);
+  }
   case Instruction::PHI: {
     auto *Phi = cast<PHINode>(I);
 
@@ -7843,38 +7854,58 @@ VPValue *VPRecipeBuilder::createEdgeMask(BasicBlock *Src, BasicBlock *Dst) {
   VPValue *SrcMask = getBlockInMask(Src);
 
   if (auto *SI = dyn_cast<SwitchInst>(Src->getTerminator())) {
-    // Create mask where the terminator in Src is a switch. We need to handle 2
-    // separate cases:
-    // 1. Dst is not the default desintation. Dst is reached if any of the cases
+    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;
+    }
+
     // 2. Dst is the default destination. Dst is reached if none of the cases
     // with destination != Dst are taken. Join the conditions for each case
     // where the destination is != Dst using a logical OR and negate it.
-    VPValue *Mask = nullptr;
-    VPValue *Cond = getVPValueOrAddLiveIn(SI->getCondition(), Plan);
-    bool IsDefault = SI->getDefaultDest() == Dst;
-    for (auto &C : SI->cases()) {
-      if (IsDefault) {
-        if (C.getCaseSuccessor() == Dst)
-          continue;
-      } else if (C.getCaseSuccessor() != Dst)
+    VPValue *DefaultMask = nullptr;
+    for (const auto &[Dst, Conds] : Map) {
+      if (Dst == DefaultDst)
+        continue;
+      if (!DefaultMask) {
+        DefaultMask = EdgeMaskCache[{Src, Dst}];
         continue;
-
-      VPValue *Eq = EdgeMaskCache.lookup({Src, C.getCaseSuccessor()});
-      if (!Eq) {
-        VPValue *V = getVPValueOrAddLiveIn(C.getCaseValue(), Plan);
-        Eq = Builder.createICmp(CmpInst::ICMP_EQ, Cond, V);
       }
-      if (Mask)
-        Mask = Builder.createOr(Mask, Eq);
-      else
-        Mask = Eq;
+      DefaultMask = Builder.createOr(DefaultMask, EdgeMaskCache[{Src, Dst}]);
+    }
+    if (DefaultMask) {
+      DefaultMask = Builder.createNot(DefaultMask);
+      if (SrcMask)
+        DefaultMask = Builder.createLogicalAnd(SrcMask, DefaultMask);
     }
-    if (IsDefault)
-      Mask = Builder.createNot(Mask);
-    assert(Mask && "mask must be created");
-    return EdgeMaskCache[Edge] = Mask;
+    EdgeMaskCache[{Src, DefaultDst}] = DefaultMask;
+    assert(EdgeMaskCache.contains(Edge) && "Mask for Edge not created?");
+    return EdgeMaskCache[Edge];
   }
 
   // The terminator has to be a branch inst!