[LV][VPlan] Add initial support for CSA vectorization

llvm/include/llvm/Analysis/IVDescriptors.h

@@ -6,7 +6,8 @@
 //
 //===----------------------------------------------------------------------===//
 //
-// This file "describes" induction and recurrence variables.
+// This file "describes" induction, recurrence, and conditional scalar
+// assignment variables.
 //
 //===----------------------------------------------------------------------===//
 
@@ -423,6 +424,69 @@ class InductionDescriptor {
   SmallVector<Instruction *, 2> RedundantCasts;
 };
 
+/// A Conditional Scalar Assignment is an assignment from an initial
+/// scalar that may or may not occur.
+class ConditionalScalarAssignmentDescriptor {
+  /// If the conditional assignment occurs inside a loop, then Phi chooses
+  /// the value of the assignment from the entry block or the loop body block.
+  PHINode *Phi = nullptr;
+
+  /// The initial value of the ConditionalScalarAssignment. If the condition
+  /// guarding the assignment is not met, then the assignment retains this
+  /// value.
+  Value *InitScalar = nullptr;
+
+  /// The Instruction that conditionally assigned to inside the loop.
+  SelectInst *Assignment = nullptr;
+
+  /// Create a ConditionalScalarAssignmentDescriptor that models a valid
+  /// conditional scalar assignment with its members initialized correctly.
+  ConditionalScalarAssignmentDescriptor(PHINode *Phi, SelectInst *Assignment,
+                                        Value *InitScalar)
+      : Phi(Phi), InitScalar(InitScalar), Assignment(Assignment) {}
+
+public:
+  /// Create a ConditionalScalarAssignmentDescriptor that models an invalid
+  /// ConditionalScalarAssignment.
+  ConditionalScalarAssignmentDescriptor() = default;
+
+  /// If Phi is the root of a ConditionalScalarAssignment, set
+  /// ConditionalScalarAssignmentDesc as the ConditionalScalarAssignment rooted
+  /// by Phi. Otherwise, return a false, leaving ConditionalScalarAssignmentDesc
+  /// unmodified.
+  static bool
+  isConditionalScalarAssignmentPhi(PHINode *Phi, Loop *TheLoop,
+                                   ConditionalScalarAssignmentDescriptor &Desc);
+
+  operator bool() const { return isValid(); }
+
+  /// Returns whether SI is the Assignment in ConditionalScalarAssignment
+  static bool isConditionalScalarAssignmentSelect(
+      ConditionalScalarAssignmentDescriptor Desc, SelectInst *SI) {
+    return Desc.getAssignment() == SI;
+  }
+
+  /// Return whether this ConditionalScalarAssignmentDescriptor models a valid
+  /// ConditionalScalarAssignment.
+  bool isValid() const { return Phi && InitScalar && Assignment; }
+
+  /// Return the PHI that roots this ConditionalScalarAssignment.
+  PHINode *getPhi() const { return Phi; }
+
+  /// Return the initial value of the ConditionalScalarAssignment. This is the
+  /// value if the conditional assignment does not occur.
+  Value *getInitScalar() const { return InitScalar; }
+
+  /// The Instruction that is used after the loop
+  SelectInst *getAssignment() const { return Assignment; }
+
+  /// Return the condition that this ConditionalScalarAssignment is conditional
+  /// upon.
+  Value *getCond() const {
+    return Assignment ? Assignment->getCondition() : nullptr;
+  }
+};
+
 } // end namespace llvm
 
 #endif // LLVM_ANALYSIS_IVDESCRIPTORS_H

llvm/include/llvm/Analysis/TargetTransformInfo.h

@@ -1859,6 +1859,10 @@ class TargetTransformInfo {
         : EVLParamStrategy(EVLParamStrategy), OpStrategy(OpStrategy) {}
   };
 
+  /// \returns true if the loop vectorizer should vectorize conditional
+  /// scalar assignments for the target.
+  bool enableConditionalScalarAssignmentVectorization() const;
+
   /// \returns How the target needs this vector-predicated operation to be
   /// transformed.
   VPLegalization getVPLegalizationStrategy(const VPIntrinsic &PI) const;
@@ -2326,6 +2330,7 @@ class TargetTransformInfo::Concept {
                              SmallVectorImpl<Use *> &OpsToSink) const = 0;
 
   virtual bool isVectorShiftByScalarCheap(Type *Ty) const = 0;
+  virtual bool enableConditionalScalarAssignmentVectorization() const = 0;
   virtual VPLegalization
   getVPLegalizationStrategy(const VPIntrinsic &PI) const = 0;
   virtual bool hasArmWideBranch(bool Thumb) const = 0;
@@ -3157,6 +3162,10 @@ class TargetTransformInfo::Model final : public TargetTransformInfo::Concept {
     return Impl.isVectorShiftByScalarCheap(Ty);
   }
 
+  bool enableConditionalScalarAssignmentVectorization() const override {
+    return Impl.enableConditionalScalarAssignmentVectorization();
+  }
+
   VPLegalization
   getVPLegalizationStrategy(const VPIntrinsic &PI) const override {
     return Impl.getVPLegalizationStrategy(PI);

llvm/include/llvm/Analysis/TargetTransformInfoImpl.h

@@ -1030,6 +1030,8 @@ class TargetTransformInfoImplBase {
 
   bool isVectorShiftByScalarCheap(Type *Ty) const { return false; }
 
+  bool enableConditionalScalarAssignmentVectorization() const { return false; }
+
   TargetTransformInfo::VPLegalization
   getVPLegalizationStrategy(const VPIntrinsic &PI) const {
     return TargetTransformInfo::VPLegalization(

llvm/include/llvm/Transforms/Vectorize/LoopVectorizationLegality.h

@@ -269,6 +269,12 @@ class LoopVectorizationLegality {
   /// induction descriptor.
   using InductionList = MapVector<PHINode *, InductionDescriptor>;
 
+  /// ConditionalScalarAssignmentList contains the
+  /// ConditionalScalarAssignmentDescriptors for all the conditional scalar
+  /// assignments  that were found in the loop, rooted by their phis.
+  using ConditionalScalarAssignmentList =
+      MapVector<PHINode *, ConditionalScalarAssignmentDescriptor>;
+
   /// RecurrenceSet contains the phi nodes that are recurrences other than
   /// inductions and reductions.
   using RecurrenceSet = SmallPtrSet<const PHINode *, 8>;
@@ -321,6 +327,18 @@ class LoopVectorizationLegality {
   /// Returns True if V is a Phi node of an induction variable in this loop.
   bool isInductionPhi(const Value *V) const;
 
+  /// Returns the conditional scalar assignments found in the loop.
+  const ConditionalScalarAssignmentList &
+  getConditionalScalarAssignments() const {
+    return ConditionalScalarAssignments;
+  }
+
+  /// Returns true if Phi is the root of a conditional scalar assignments in the
+  /// loop.
+  bool isConditionalScalarAssignmentPhi(PHINode *Phi) const {
+    return ConditionalScalarAssignments.count(Phi) != 0;
+  }
+
   /// Returns a pointer to the induction descriptor, if \p Phi is an integer or
   /// floating point induction.
   const InductionDescriptor *getIntOrFpInductionDescriptor(PHINode *Phi) const;
@@ -545,6 +563,12 @@ class LoopVectorizationLegality {
   void addInductionPhi(PHINode *Phi, const InductionDescriptor &ID,
                        SmallPtrSetImpl<Value *> &AllowedExit);
 
+  /// Updates the vetorization state by adding \p Phi to the
+  /// ConditionalScalarAssignment list.
+  void addConditionalScalarAssignmentPhi(
+      PHINode *Phi, const ConditionalScalarAssignmentDescriptor &Desc,
+      SmallPtrSetImpl<Value *> &AllowedExit);
+
   /// The loop that we evaluate.
   Loop *TheLoop;
 
@@ -589,6 +613,9 @@ class LoopVectorizationLegality {
   /// variables can be pointers.
   InductionList Inductions;
 
+  /// Holds the conditional scalar assignments
+  ConditionalScalarAssignmentList ConditionalScalarAssignments;
+
   /// Holds all the casts that participate in the update chain of the induction
   /// variables, and that have been proven to be redundant (possibly under a
   /// runtime guard). These casts can be ignored when creating the vectorized

llvm/lib/Analysis/IVDescriptors.cpp

@@ -6,7 +6,8 @@
 //
 //===----------------------------------------------------------------------===//
 //
-// This file "describes" induction and recurrence variables.
+// This file "describes" induction, recurrence, and conditional scalar
+// assignment variables.
 //
 //===----------------------------------------------------------------------===//
 
@@ -1570,3 +1571,60 @@ bool InductionDescriptor::isInductionPHI(
   D = InductionDescriptor(StartValue, IK_PtrInduction, Step);
   return true;
 }
+
+/// Return ConditionalScalarAssignmentDescriptor that describes a
+/// ConditionalScalarAssignment that matches one of these patterns:
+///   phi loop_inv, (select cmp, value, phi)
+///   phi loop_inv, (select cmp, phi, value)
+///   phi (select cmp, value, phi), loop_inv
+///   phi (select cmp, phi, value), loop_inv
+/// If the ConditionalScalarAssignment does not match any of these paterns,
+/// return a ConditionalScalarAssignmentDescriptor that describes an
+/// InvalidConditionalScalarAssignment.
+bool ConditionalScalarAssignmentDescriptor::isConditionalScalarAssignmentPhi(
+    PHINode *Phi, Loop *TheLoop, ConditionalScalarAssignmentDescriptor &Desc) {
+
+  // Must be a scalar.
+  Type *Type = Phi->getType();
+  if (!Type->isIntegerTy() && !Type->isFloatingPointTy() &&
+      !Type->isPointerTy())
+    return false;
+
+  // Match phi loop_inv, (select cmp, value, phi)
+  //    or phi loop_inv, (select cmp, phi, value)
+  //    or phi (select cmp, value, phi), loop_inv
+  //    or phi (select cmp, phi, value), loop_inv
+  if (Phi->getNumIncomingValues() != 2)
+    return false;
+  auto SelectInstIt = find_if(Phi->incoming_values(), [&Phi](const Use &U) {
+    return match(U.get(), m_Select(m_Value(), m_Specific(Phi), m_Value())) ||
+           match(U.get(), m_Select(m_Value(), m_Value(), m_Specific(Phi)));
+  });
+  if (SelectInstIt == Phi->incoming_values().end())
+    return false;
+  auto LoopInvIt = find_if(Phi->incoming_values(), [&](Use &U) {
+    return U.get() != *SelectInstIt && TheLoop->isLoopInvariant(U.get());
+  });
+  if (LoopInvIt == Phi->incoming_values().end())
+    return false;
+
+  // Phi or Sel must be used only outside the loop,
+  // excluding if Phi use Sel or Sel use Phi
+  auto IsOnlyUsedOutsideLoop = [&](Value *V, Value *Ignore) {
+    return all_of(V->users(), [Ignore, TheLoop](User *U) {
+      if (U == Ignore)
+        return true;
+      if (auto *I = dyn_cast<Instruction>(U))
+        return !TheLoop->contains(I);
+      return true;
+    });
+  };
+  SelectInst *Select = cast<SelectInst>(SelectInstIt->get());
+  Value *LoopInv = LoopInvIt->get();
+  if (!IsOnlyUsedOutsideLoop(Phi, Select) ||
+      !IsOnlyUsedOutsideLoop(Select, Phi))
+    return false;
+
+  Desc = ConditionalScalarAssignmentDescriptor(Phi, Select, LoopInv);
+  return true;
+}

llvm/lib/Analysis/TargetTransformInfo.cpp

@@ -1374,6 +1374,11 @@ bool TargetTransformInfo::preferEpilogueVectorization() const {
   return TTIImpl->preferEpilogueVectorization();
 }
 
+bool TargetTransformInfo::enableConditionalScalarAssignmentVectorization()
+    const {
+  return TTIImpl->enableConditionalScalarAssignmentVectorization();
+}
+
 TargetTransformInfo::VPLegalization
 TargetTransformInfo::getVPLegalizationStrategy(const VPIntrinsic &VPI) const {
   return TTIImpl->getVPLegalizationStrategy(VPI);

llvm/lib/Target/RISCV/RISCVTargetTransformInfo.cpp

@@ -2535,6 +2535,11 @@ bool RISCVTTIImpl::isLegalMaskedExpandLoad(Type *DataTy, Align Alignment) {
   return true;
 }
 
+bool RISCVTTIImpl::enableConditionalScalarAssignmentVectorization() const {
+  return ST->hasVInstructions() &&
+         ST->getProcFamily() == RISCVSubtarget::SiFive7;
+}
+
 bool RISCVTTIImpl::isLegalMaskedCompressStore(Type *DataTy, Align Alignment) {
   auto *VTy = dyn_cast<VectorType>(DataTy);
   if (!VTy || VTy->isScalableTy())

llvm/lib/Target/RISCV/RISCVTargetTransformInfo.h

@@ -314,6 +314,10 @@ class RISCVTTIImpl : public BasicTTIImplBase<RISCVTTIImpl> {
     return TLI->isVScaleKnownToBeAPowerOfTwo();
   }
 
+  /// \returns true if the loop vectorizer should vectorize conditional
+  /// scalar assignments for the target.
+  bool enableConditionalScalarAssignmentVectorization() const;
+
   /// \returns How the target needs this vector-predicated operation to be
   /// transformed.
   TargetTransformInfo::VPLegalization

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

@@ -83,6 +83,10 @@ static cl::opt<bool> EnableHistogramVectorization(
     "enable-histogram-loop-vectorization", cl::init(false), cl::Hidden,
     cl::desc("Enables autovectorization of some loops containing histograms"));
 
+static cl::opt<bool> EnableConditionalScalarAssignment(
+    "enable-csa-vectorization", cl::init(false), cl::Hidden,
+    cl::desc("Control whether loop vectorization is enabled"));
+
 /// Maximum vectorization interleave count.
 static const unsigned MaxInterleaveFactor = 16;
 
@@ -749,6 +753,18 @@ bool LoopVectorizationLegality::setupOuterLoopInductions() {
   return llvm::all_of(Header->phis(), IsSupportedPhi);
 }
 
+void LoopVectorizationLegality::addConditionalScalarAssignmentPhi(
+    PHINode *Phi, const ConditionalScalarAssignmentDescriptor &Desc,
+    SmallPtrSetImpl<Value *> &AllowedExit) {
+  assert(Desc.isValid() &&
+         "Expected Valid ConditionalScalarAssignmentDescriptor");
+  LLVM_DEBUG(
+      dbgs() << "LV: found legal conditional scalar assignment opportunity"
+             << *Phi << "\n");
+  AllowedExit.insert(Phi);
+  ConditionalScalarAssignments.insert({Phi, Desc});
+}
+
 /// Checks if a function is scalarizable according to the TLI, in
 /// the sense that it should be vectorized and then expanded in
 /// multiple scalar calls. This is represented in the
@@ -878,14 +894,27 @@ bool LoopVectorizationLegality::canVectorizeInstrs() {
           continue;
         }
 
-        // As a last resort, coerce the PHI to a AddRec expression
-        // and re-try classifying it a an induction PHI.
+        // Try to coerce the PHI to a AddRec expression and re-try classifying
+        // it a an induction PHI.
         if (InductionDescriptor::isInductionPHI(Phi, TheLoop, PSE, ID, true) &&
             !IsDisallowedStridedPointerInduction(ID)) {
           addInductionPhi(Phi, ID, AllowedExit);
           continue;
         }
 
+        // Check if the PHI can be classified as a conditional scalar assignment
+        // PHI.
+        if (EnableConditionalScalarAssignment ||
+            (TTI->enableConditionalScalarAssignmentVectorization() &&
+             EnableConditionalScalarAssignment.getNumOccurrences() == 0)) {
+          ConditionalScalarAssignmentDescriptor Desc;
+          if (ConditionalScalarAssignmentDescriptor::
+                  isConditionalScalarAssignmentPhi(Phi, TheLoop, Desc)) {
+            addConditionalScalarAssignmentPhi(Phi, Desc, AllowedExit);
+            continue;
+          }
+        }
+
         reportVectorizationFailure("Found an unidentified PHI",
             "value that could not be identified as "
             "reduction is used outside the loop",
@@ -1885,11 +1914,15 @@ bool LoopVectorizationLegality::canFoldTailByMasking() const {
   for (const auto &Reduction : getReductionVars())
     ReductionLiveOuts.insert(Reduction.second.getLoopExitInstr());
 
+  SmallPtrSet<const Value *, 8> CSALiveOuts;
+  for (const auto &CSA : getConditionalScalarAssignments())
+    CSALiveOuts.insert(CSA.second.getAssignment());
+
   // TODO: handle non-reduction outside users when tail is folded by masking.
   for (auto *AE : AllowedExit) {
     // Check that all users of allowed exit values are inside the loop or
-    // are the live-out of a reduction.
-    if (ReductionLiveOuts.count(AE))
+    // are the live-out of a reduction or conditional scalar assignment.
+    if (ReductionLiveOuts.count(AE) || CSALiveOuts.count(AE))
       continue;
     for (User *U : AE->users()) {
       Instruction *UI = cast<Instruction>(U);

llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h

@@ -175,8 +175,8 @@ class VPBuilder {
         new VPInstruction(Opcode, Operands, WrapFlags, DL, Name));
   }
 
-  VPValue *createNot(VPValue *Operand, DebugLoc DL = {},
-                     const Twine &Name = "") {
+  VPInstruction *createNot(VPValue *Operand, DebugLoc DL = {},
+                           const Twine &Name = "") {
     return createInstruction(VPInstruction::Not, {Operand}, DL, Name);
   }
 
@@ -262,6 +262,26 @@ class VPBuilder {
         FPBinOp ? FPBinOp->getFastMathFlags() : FastMathFlags()));
   }
 
+  VPInstruction *createConditionalScalarAssignmentMaskPhi(VPValue *InitMask,
+                                                          DebugLoc DL,
+                                                          const Twine &Name) {
+    return createInstruction(VPInstruction::ConditionalScalarAssignmentMaskPhi,
+                             {InitMask}, DL, Name);
+  }
+
+  VPInstruction *createAnyOf(VPValue *Cond, DebugLoc DL, const Twine &Name) {
+    return createInstruction(VPInstruction::AnyOf, {Cond}, DL, Name);
+  }
+
+  VPInstruction *createConditionalScalarAssignmentMaskSel(VPValue *Cond,
+                                                          VPValue *MaskPhi,
+                                                          VPValue *AnyOf,
+                                                          DebugLoc DL,
+                                                          const Twine &Name) {
+    return createInstruction(VPInstruction::ConditionalScalarAssignmentMaskSel,
+                             {Cond, MaskPhi, AnyOf}, DL, Name);
+  }
+
   //===--------------------------------------------------------------------===//
   // RAII helpers.
   //===--------------------------------------------------------------------===//