!fixup address latest comments, thanks!

Author: fhahn

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -7850,6 +7850,60 @@ VPRecipeBuilder::mapToVPValues(User::op_range Operands) {
   return map_range(Operands, Fn);
 }
 
+void VPRecipeBuilder::createSwitchEdgeMasks(SwitchInst *SI) {
+  BasicBlock *Src = SI->getParent();
+  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 *>> Dst2Compares;
+  for (auto &C : SI->cases()) {
+    // Cases whose destination is the same as default are redundant and can be
+    // ignored - they will get there anyhow.
+    if (C.getCaseSuccessor() == DefaultDst)
+      continue;
+    auto I = Dst2Compares.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 below for all entries in Dst2Compares,
+  // which excludes destinations matching the default destination.
+  VPValue *SrcMask = getBlockInMask(Src);
+  VPValue *DefaultMask = nullptr;
+  for (const auto &[Dst, Conds] : Dst2Compares) {
+    // 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
+    // whose destination == Dst using an OR.
+    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. Create the mask for the default destination, which is reached if none
+    // of the cases with destination != default destination are taken. Join the
+    // conditions for each case where the destination is != Dst using an OR and
+    // negate it.
+    DefaultMask = DefaultMask ? Builder.createOr(DefaultMask, Mask) : Mask;
+  }
+
+  if (DefaultMask) {
+    DefaultMask = Builder.createNot(DefaultMask);
+    if (SrcMask)
+      DefaultMask = Builder.createLogicalAnd(SrcMask, DefaultMask);
+  }
+  EdgeMaskCache[{Src, DefaultDst}] = DefaultMask;
+}
+
 VPValue *VPRecipeBuilder::createEdgeMask(BasicBlock *Src, BasicBlock *Dst) {
   assert(is_contained(predecessors(Dst), Src) && "Invalid edge");
 
@@ -7859,67 +7913,17 @@ VPValue *VPRecipeBuilder::createEdgeMask(BasicBlock *Src, BasicBlock *Dst) {
   if (ECEntryIt != EdgeMaskCache.end())
     return ECEntryIt->second;
 
-  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;
-    }
-
-    // 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 *DefaultMask = nullptr;
-    for (const auto &[Dst, Conds] : Map) {
-      if (Dst == DefaultDst)
-        continue;
-      if (!DefaultMask) {
-        DefaultMask = EdgeMaskCache[{Src, Dst}];
-        continue;
-      }
-      DefaultMask = Builder.createOr(DefaultMask, EdgeMaskCache[{Src, Dst}]);
-    }
-    if (DefaultMask) {
-      DefaultMask = Builder.createNot(DefaultMask);
-      if (SrcMask)
-        DefaultMask = Builder.createLogicalAnd(SrcMask, DefaultMask);
-    }
-    EdgeMaskCache[{Src, DefaultDst}] = DefaultMask;
+    createSwitchEdgeMasks(SI);
     assert(EdgeMaskCache.contains(Edge) && "Mask for Edge not created?");
     return EdgeMaskCache[Edge];
   }
 
+  VPValue *SrcMask = getBlockInMask(Src);
+
   // The terminator has to be a branch inst!
   BranchInst *BI = dyn_cast<BranchInst>(Src->getTerminator());
   assert(BI && "Unexpected terminator found");
-
   if (!BI->isConditional() || BI->getSuccessor(0) == BI->getSuccessor(1))
     return EdgeMaskCache[Edge] = SrcMask;
 

llvm/lib/Transforms/Vectorize/VPRecipeBuilder.h

@@ -134,6 +134,10 @@ class VPRecipeBuilder {
   /// Returns the *entry* mask for the block \p BB.
   VPValue *getBlockInMask(BasicBlock *BB) const;
 
+  /// Create masks for all cases with destination different than the default
+  /// destination, and a mask for the default destination.
+  void createSwitchEdgeMasks(SwitchInst *SI);
+
   /// A helper function that computes the predicate of the edge between SRC
   /// and DST.
   VPValue *createEdgeMask(BasicBlock *Src, BasicBlock *Dst);