[LV] Build VPlan for CSA

Author: michaelmaitland

llvm/include/llvm/Analysis/CSADescriptors.h

@@ -6,7 +6,8 @@
 //
 //===----------------------------------------------------------------------===//
 //
-// This file "describes" induction and recurrence variables.
+// This file "describes" induction, recurrence, and conditional scalar
+// assignment (CSA) variables.
 //
 //===----------------------------------------------------------------------===//
 
@@ -423,6 +424,61 @@ class InductionDescriptor {
   SmallVector<Instruction *, 2> RedundantCasts;
 };
 
+/// A Conditional Scalar Assignment (CSA) is an assignment from an initial
+/// scalar that may or may not occur.
+class CSADescriptor {
+  /// 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 CSA. 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.
+  Instruction *Assignment = nullptr;
+
+  /// Create a CSA Descriptor that models a valid CSA with its members
+  /// initialized correctly.
+  CSADescriptor(PHINode *Phi, Instruction *Assignment, Value *InitScalar)
+      : Phi(Phi), InitScalar(InitScalar), Assignment(Assignment) {}
+
+public:
+  /// Create a CSA Descriptor that models an invalid CSA.
+  CSADescriptor() = default;
+
+  /// If Phi is the root of a CSA, set CSADesc as the CSA rooted by
+  /// Phi. Otherwise, return a false, leaving CSADesc unmodified.
+  static bool isCSAPhi(PHINode *Phi, Loop *TheLoop, CSADescriptor &CSADesc);
+
+  operator bool() const { return isValid(); }
+
+  /// Returns whether SI is the Assignment in CSA
+  static bool isCSASelect(CSADescriptor Desc, SelectInst *SI) {
+    return Desc.getAssignment() == SI;
+  }
+
+  /// Return whether this CSADescriptor models a valid CSA.
+  bool isValid() const { return Phi && InitScalar && Assignment; }
+
+  /// Return the PHI that roots this CSA.
+  PHINode *getPhi() const { return Phi; }
+
+  /// Return the initial value of the CSA. This is the value if the conditional
+  /// assignment does not occur.
+  Value *getInitScalar() const { return InitScalar; }
+
+  /// The Instruction that is used after the loop
+  Instruction *getAssignment() const { return Assignment; }
+
+  /// Return the condition that this CSA is conditional upon.
+  Value *getCond() const {
+    if (auto *SI = dyn_cast_or_null<SelectInst>(Assignment))
+      return SI->getCondition();
+    return nullptr;
+  }
+};
+
 } // end namespace llvm
 
 #endif // LLVM_ANALYSIS_IVDESCRIPTORS_H

llvm/include/llvm/Analysis/IVDescriptors.h

@@ -27,7 +27,6 @@
 #define LLVM_TRANSFORMS_VECTORIZE_LOOPVECTORIZATIONLEGALITY_H
 
 #include "llvm/ADT/MapVector.h"
-#include "llvm/Analysis/CSADescriptors.h"
 #include "llvm/Analysis/LoopAccessAnalysis.h"
 #include "llvm/Support/TypeSize.h"
 #include "llvm/Transforms/Utils/LoopUtils.h"

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

@@ -46,7 +46,6 @@ add_llvm_component_library(LLVMAnalysis
   CostModel.cpp
   CodeMetrics.cpp
   ConstantFolding.cpp
-  CSADescriptors.cpp
   CtxProfAnalysis.cpp
   CycleAnalysis.cpp
   DDG.cpp

llvm/lib/Analysis/CMakeLists.txt

@@ -6,7 +6,8 @@
 //
 //===----------------------------------------------------------------------===//
 //
-// This file "describes" induction and recurrence variables.
+// This file "describes" induction, recurrence, and conditional scalar
+// assignment (CSA) variables.
 //
 //===----------------------------------------------------------------------===//
 
@@ -1570,3 +1571,58 @@ bool InductionDescriptor::isInductionPHI(
   D = InductionDescriptor(StartValue, IK_PtrInduction, Step);
   return true;
 }
+
+/// Return CSADescriptor that describes a CSA 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 CSA does not match any of these paterns, return a CSADescriptor
+/// that describes an InvalidCSA.
+bool CSADescriptor::isCSAPhi(PHINode *Phi, Loop *TheLoop, CSADescriptor &CSA) {
+
+  // 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;
+    });
+  };
+  Instruction *Select = cast<SelectInst>(SelectInstIt->get());
+  Value *LoopInv = LoopInvIt->get();
+  if (!IsOnlyUsedOutsideLoop(Phi, Select) ||
+      !IsOnlyUsedOutsideLoop(Select, Phi))
+    return false;
+
+  CSA = CSADescriptor(Phi, Select, LoopInv);
+  return true;
+}

llvm/lib/Analysis/CSADescriptors.cpp

@@ -891,7 +891,8 @@ bool LoopVectorizationLegality::canVectorizeInstrs() {
         // Check if the PHI can be classified as a CSA PHI.
         if (EnableCSA || (TTI->enableCSAVectorization() &&
                           EnableCSA.getNumOccurrences() == 0)) {
-          if (auto CSADesc = CSADescriptor::isCSAPhi(Phi, TheLoop)) {
+          CSADescriptor CSADesc;
+          if (CSADescriptor::isCSAPhi(Phi, TheLoop, CSADesc)) {
             addCSAPhi(Phi, CSADesc, AllowedExit);
             continue;
           }
@@ -1887,7 +1888,7 @@ bool LoopVectorizationLegality::canFoldTailByMasking() const {
   // 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 or a CSA
+    // are the live-out of a reduction or a CSA.
     if (ReductionLiveOuts.count(AE) || CSALiveOuts.count(AE))
       continue;
     for (User *U : AE->users()) {

llvm/lib/Analysis/IVDescriptors.cpp

@@ -174,8 +174,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);
   }
 
@@ -231,6 +231,37 @@ class VPBuilder {
         Ptr, Offset, VPRecipeWithIRFlags::GEPFlagsTy(true), DL, Name));
   }
 
+  VPInstruction *createCSAMaskPhi(VPValue *InitMask, DebugLoc DL,
+                                  const Twine &Name) {
+    return createInstruction(VPInstruction::CSAMaskPhi, {InitMask}, DL, Name);
+  }
+
+  VPInstruction *createAnyOf(VPValue *Cond, DebugLoc DL, const Twine &Name) {
+    return createInstruction(VPInstruction::AnyOf, {Cond}, DL, Name);
+  }
+
+  VPInstruction *createCSAMaskSel(VPValue *Cond, VPValue *MaskPhi,
+                                  VPValue *AnyOf, DebugLoc DL,
+                                  const Twine &Name) {
+    return createInstruction(VPInstruction::CSAMaskSel, {Cond, MaskPhi, AnyOf},
+                             DL, Name);
+  }
+
+  VPInstruction *createAnyOfEVL(VPValue *Cond, VPValue *EVL, DebugLoc DL,
+                                const Twine &Name) {
+    return createInstruction(VPInstruction::AnyOfEVL, {Cond, EVL}, DL, Name);
+  }
+
+  VPInstruction *createCSAVLPhi(DebugLoc DL, const Twine &Name) {
+    return createInstruction(VPInstruction::CSAVLPhi, {}, DL, Name);
+  }
+
+  VPInstruction *createCSAVLSel(VPValue *AnyOfEVL, VPValue *VLPhi, VPValue *EVL,
+                                DebugLoc DL, const Twine &Name) {
+    return createInstruction(VPInstruction::CSAVLSel, {AnyOfEVL, VLPhi, EVL},
+                             DL, Name);
+  }
+
   VPDerivedIVRecipe *createDerivedIV(InductionDescriptor::InductionKind Kind,
                                      FPMathOperator *FPBinOp, VPValue *Start,
                                      VPCanonicalIVPHIRecipe *CanonicalIV,