Addressing comments

* Mask is optional for vplan recipe
* Left load+update to DCE
* Documentation for new recipe
* Stored histogram info by reference
* Allow GEPs with more than one index value, provided that all
  but the last are constant.
* Cleanup to remove some unnecessary local variables.
* Minor nits addressed.

Author: huntergr-arm

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

@@ -1050,7 +1050,7 @@ bool LoopVectorizationLegality::canVectorizeInstrs() {
   return true;
 }
 
-/// Find Histogram counts that match high-level code in loops:
+/// Find histogram operations that match high-level code in loops:
 /// \code
 /// buckets[indices[i]]+=step;
 /// \endcode
@@ -1063,10 +1063,9 @@ bool LoopVectorizationLegality::canVectorizeInstrs() {
 /// regardless of hardware support. When there is support, it additionally
 /// stores the BinOp/Load pairs in \p HistogramCounts, as well the pointers
 /// used to update histogram in \p HistogramPtrs.
-
-static bool findHistograms(LoadInst *LI, StoreInst *HSt, Loop *TheLoop,
-                           const PredicatedScalarEvolution &PSE,
-                           SmallVectorImpl<HistogramInfo> &Histograms) {
+static bool findHistogram(LoadInst *LI, StoreInst *HSt, Loop *TheLoop,
+                          const PredicatedScalarEvolution &PSE,
+                          SmallVectorImpl<HistogramInfo> &Histograms) {
 
   // Store value must come from a Binary Operation.
   Instruction *HPtrInstr = nullptr;
@@ -1088,23 +1087,27 @@ static bool findHistograms(LoadInst *LI, StoreInst *HSt, Loop *TheLoop,
     return false;
 
   // The address to store is calculated through a GEP Instruction.
-  // FIXME: Support GEPs with more operands.
-  GetElementPtrInst *HPtr = dyn_cast<GetElementPtrInst>(HPtrInstr);
-  if (!HPtr || HPtr->getNumOperands() > 2)
+  GetElementPtrInst *GEP = dyn_cast<GetElementPtrInst>(HPtrInstr);
+  if (!GEP)
+    return false;
+
+  // Restrict address calculation to constant indices except for the last term.
+  Value *HIdx = nullptr;
+  for (Value *Index : GEP->indices()) {
+    if (HIdx)
+      return false;
+    if (!isa<ConstantInt>(Index))
+      HIdx = Index;
+  }
+
+  if (!HIdx)
     return false;
 
   // Check that the index is calculated by loading from another array. Ignore
   // any extensions.
   // FIXME: Support indices from other sources that a linear load from memory?
-  Value *HIdx = HPtr->getOperand(1);
-  Instruction *IdxInst = nullptr;
-  if (!match(HIdx, m_ZExtOrSExtOrSelf(m_Instruction(IdxInst))))
-    return false;
-
-  // Currently restricting this to linear addressing when loading indices.
-  LoadInst *VLoad = dyn_cast<LoadInst>(IdxInst);
   Value *VPtrVal;
-  if (!VLoad || !match(VLoad, m_Load(m_Value(VPtrVal))))
+  if (!match(HIdx, m_ZExtOrSExtOrSelf(m_Load(m_Value(VPtrVal)))))
     return false;
 
   if (!isa<SCEVAddRecExpr>(PSE.getSE()->getSCEV(VPtrVal)))
@@ -1134,7 +1137,7 @@ bool LoopVectorizationLegality::canVectorizeUnknownDependences() {
   LAI = &LAIs.getInfo(*TheLoop);
   const MemoryDepChecker &DepChecker = LAI->getDepChecker();
   const auto *Deps = DepChecker.getDependences();
-  if (!Deps || Deps->size() > 1)
+  if (!Deps || Deps->size() != 1)
     return false;
 
   const MemoryDepChecker::Dependence &Dep = (*Deps).front();
@@ -1151,7 +1154,7 @@ bool LoopVectorizationLegality::canVectorizeUnknownDependences() {
     return false;
 
   LLVM_DEBUG(dbgs() << "LV: Checking for a histogram on: " << *SI << "\n");
-  return findHistograms(LI, SI, TheLoop, LAI->getPSE(), Histograms);
+  return findHistogram(LI, SI, TheLoop, LAI->getPSE(), Histograms);
 }
 
 bool LoopVectorizationLegality::canVectorizeMemory() {

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -4660,10 +4660,7 @@ bool LoopVectorizationPlanner::isCandidateForEpilogueVectorization(
     return false;
 
   // Loops containing histograms are not currently supported.
-  if (!Legal->getHistograms().empty())
-    return false;
-
-  return true;
+  return Legal->getHistograms().empty();
 }
 
 bool LoopVectorizationCostModel::isEpilogueVectorizationProfitable(
@@ -8296,17 +8293,11 @@ VPRecipeBuilder::tryToWidenHistogram(const HistogramInfo *HI,
   HGramOps.push_back(getVPValueOrAddLiveIn(HI->Update->getOperand(1), Plan));
 
   // In case of predicated execution (due to tail-folding, or conditional
-  // execution, or both), pass the relevant mask. When there is no such mask,
-  // generate an all-true mask.
-  VPValue *Mask = nullptr;
+  // execution, or both), pass the relevant mask.
   if (Legal->isMaskRequired(HI->Store))
-    Mask = getBlockInMask(HI->Store->getParent());
-  else
-    Mask = Plan.getOrAddLiveIn(
-        ConstantInt::getTrue(IntegerType::getInt1Ty(HI->Load->getContext())));
-  HGramOps.push_back(Mask);
+    HGramOps.push_back(getBlockInMask(HI->Store->getParent()));
 
-  return new VPHistogramRecipe(HI, Opcode,
+  return new VPHistogramRecipe(*HI, Opcode,
                                make_range(HGramOps.begin(), HGramOps.end()),
                                HI->Store->getDebugLoc());
 }
@@ -8711,15 +8702,6 @@ LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(VFRange &Range) {
         Operands = {OpRange.begin(), OpRange.end()};
       }
 
-      // If this is a load instruction or a binop associated with a histogram,
-      // leave it until the store instruction to emit a combined intrinsic.
-      // Note that if the initial VF is scalar, we need to generate the normal
-      // clone recipe for these instructions. A histogram recipe will only be
-      // generated when minVF > 1.
-      if (Legal->getHistogramInfo(Instr) && !isa<StoreInst>(Instr) &&
-          !Range.Start.isScalar())
-        continue;
-
       // Invariant stores inside loop will be deleted and a single store
       // with the final reduction value will be added to the exit block
       StoreInst *SI;

llvm/lib/Transforms/Vectorize/VPlan.h

@@ -1577,13 +1577,18 @@ class VPWidenCallRecipe : public VPSingleDefRecipe {
 #endif
 };
 
+/// A recipe representing a sequence of load -> update -> store as part of
+/// a histogram operation. This means there may be aliasing between vector
+/// lanes, which is handled by the llvm.experimental.vector.histogram family
+/// of intrinsics. The only update operations currently supported are
+/// 'add' and 'sub' where the other term is loop-invariant.
 class VPHistogramRecipe : public VPRecipeBase {
-  const HistogramInfo *Info;
+  const HistogramInfo &Info;
   unsigned Opcode;
 
 public:
   template <typename IterT>
-  VPHistogramRecipe(const HistogramInfo *HI, unsigned Opcode,
+  VPHistogramRecipe(const HistogramInfo &HI, unsigned Opcode,
                     iterator_range<IterT> Operands, DebugLoc DL = {})
       : VPRecipeBase(VPDef::VPHistogramSC, Operands, DL), Info(HI),
         Opcode(Opcode) {}
@@ -1596,12 +1601,20 @@ class VPHistogramRecipe : public VPRecipeBase {
 
   VP_CLASSOF_IMPL(VPDef::VPHistogramSC);
 
-  // Produce a histogram operation with widened ingredients
+  /// Produce a vectorized histogram operation.
   void execute(VPTransformState &State) override;
 
   unsigned getOpcode() const { return Opcode; }
 
-  const HistogramInfo *getHistogramInfo() const { return Info; }
+  const HistogramInfo &getHistogramInfo() const { return Info; }
+
+  /// Return the mask operand if one was provided, or a null pointer if all
+  /// lanes should be executed unconditionally.
+  VPValue *getMask() const {
+    if (getNumOperands() == 3)
+      return getOperand(2);
+    return nullptr;
+  }
 
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
   /// Print the recipe

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

@@ -284,7 +284,7 @@ static Instruction *getInstructionForCost(const VPRecipeBase *R) {
   // FIXME: Override the cost method properly to take gather/scatter cost
   //        into account, instead of just the intrinsic via the legacy model.
   if (auto *HG = dyn_cast<VPHistogramRecipe>(R))
-    return HG->getHistogramInfo()->Update;
+    return HG->getHistogramInfo().Update;
   return nullptr;
 }
 
@@ -973,7 +973,17 @@ void VPHistogramRecipe::execute(VPTransformState &State) {
   for (unsigned Part = 0; Part < State.UF; ++Part) {
     Value *Address = State.get(getOperand(0), Part);
     Value *IncVec = State.get(getOperand(1), Part);
-    Value *Mask = State.get(getOperand(2), Part);
+    VectorType *VTy = cast<VectorType>(Address->getType());
+
+    // The histogram intrinsic requires a mask even if the recipe doesn't;
+    // if the mask operand was omitted then all lanes should be executed and
+    // we just need to synthesize an all-true mask.
+    Value *Mask = nullptr;
+    if (VPValue *VPMask = getMask())
+      Mask = State.get(VPMask, Part);
+    else
+      Mask = Builder.CreateVectorSplat(
+          VTy->getElementCount(), ConstantInt::getTrue(Builder.getInt1Ty()));
 
     // Not sure how to make IncAmt stay scalar yet. For now just extract the
     // first element and tidy up later.
@@ -985,9 +995,11 @@ void VPHistogramRecipe::execute(VPTransformState &State) {
     // add a separate intrinsic in future, but for now we'll try this.
     if (Opcode == Instruction::Sub)
       IncAmt = Builder.CreateNeg(IncAmt);
+    else
+      assert(Opcode == Instruction::Add);
 
     State.Builder.CreateIntrinsic(Intrinsic::experimental_vector_histogram_add,
-                                  {Address->getType(), IncAmt->getType()},
+                                  {VTy, IncAmt->getType()},
                                   {Address, IncAmt, Mask});
   }
 }
@@ -997,10 +1009,19 @@ void VPHistogramRecipe::print(raw_ostream &O, const Twine &Indent,
                               VPSlotTracker &SlotTracker) const {
   O << Indent << "WIDEN-HISTOGRAM buckets: ";
   getOperand(0)->printAsOperand(O, SlotTracker);
-  O << ", inc: ";
+
+  if (Opcode == Instruction::Sub)
+    O << ", dec: ";
+  else {
+    assert(Opcode == Instruction::Add);
+    O << ", inc: ";
+  }
   getOperand(1)->printAsOperand(O, SlotTracker);
-  O << ", mask: ";
-  getOperand(2)->printAsOperand(O, SlotTracker);
+
+  if (VPValue *Mask = getMask()) {
+    O << ", mask: ";
+    Mask->printAsOperand(O, SlotTracker);
+  }
 }
 
 void VPWidenSelectRecipe::print(raw_ostream &O, const Twine &Indent,