Addressing review comments:

* Improve test ir names
* Explain current limitation on indices
* Check index address is not loop-invariant (with test)
* Check for null pointer on dependences (with test)
* Remove HistogramInfo from recipe, remove legality header from VPlanRecipes.cpp
* Implement overridden computeCost for VPHistogramRecipe
* Separated VPlan checks from main test file, used regex matching
* Added a 4-op GEP test featuring a struct
* Added 64b histogram test without index extension
* Added comment to target-generic test file explaining it doesn't currently vectorize

Author: huntergr-arm

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

@@ -1110,12 +1110,17 @@ static bool findHistogram(LoadInst *LI, StoreInst *HSt, Loop *TheLoop,
 
   // 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?
+  // FIXME: Support indices from other sources than a linear load from memory?
+  //        We're currently trying to match an operation looping over an array
+  //        of indices, but there could be additional levels of indirection
+  //        in place, or possibly some additional calculation to form the index
+  //        from the loaded data.
   Value *VPtrVal;
   if (!match(HIdx, m_ZExtOrSExtOrSelf(m_Load(m_Value(VPtrVal)))))
     return false;
 
-  if (!isa<SCEVAddRecExpr>(PSE.getSE()->getSCEV(VPtrVal)))
+  if (!isa<SCEVAddRecExpr>(PSE.getSE()->getSCEV(VPtrVal)) ||
+      TheLoop->isLoopInvariant(VPtrVal))
     return false;
 
   // Ensure we'll have the same mask by checking that all parts of the histogram
@@ -1143,6 +1148,11 @@ bool LoopVectorizationLegality::canVectorizeIndirectUnsafeDependences() {
   const MemoryDepChecker::Dependence *IUDep = nullptr;
   const MemoryDepChecker &DepChecker = LAI->getDepChecker();
   const auto *Deps = DepChecker.getDependences();
+  // If there were too many dependences, LAA abandons recording them. We can't
+  // proceed safely if we don't know what the dependences are.
+  if (!Deps)
+    return false;
+
   for (const MemoryDepChecker::Dependence &Dep : *Deps) {
     // Ignore dependencies that are either known to be safe or can be
     // checked at runtime.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -8477,7 +8477,7 @@ VPRecipeBuilder::tryToWidenHistogram(const HistogramInfo *HI,
   if (Legal->isMaskRequired(HI->Store))
     HGramOps.push_back(getBlockInMask(HI->Store->getParent()));
 
-  return new VPHistogramRecipe(*HI, Opcode,
+  return new VPHistogramRecipe(Opcode,
                                make_range(HGramOps.begin(), HGramOps.end()),
                                HI->Store->getDebugLoc());
 }

llvm/lib/Transforms/Vectorize/VPlan.h

@@ -1698,15 +1698,13 @@ class VPWidenCallRecipe : public VPSingleDefRecipe {
 /// 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;
   unsigned Opcode;
 
 public:
   template <typename IterT>
-  VPHistogramRecipe(const HistogramInfo &HI, unsigned Opcode,
-                    iterator_range<IterT> Operands, DebugLoc DL = {})
-      : VPRecipeBase(VPDef::VPHistogramSC, Operands, DL), Info(HI),
-        Opcode(Opcode) {}
+  VPHistogramRecipe(unsigned Opcode, iterator_range<IterT> Operands,
+                    DebugLoc DL = {})
+      : VPRecipeBase(VPDef::VPHistogramSC, Operands, DL), Opcode(Opcode) {}
 
   ~VPHistogramRecipe() override = default;
 
@@ -1719,9 +1717,11 @@ class VPHistogramRecipe : public VPRecipeBase {
   /// Produce a vectorized histogram operation.
   void execute(VPTransformState &State) override;
 
-  unsigned getOpcode() const { return Opcode; }
+  /// Return the cost of this VPHistogramRecipe.
+  InstructionCost computeCost(ElementCount VF,
+                              VPCostContext &Ctx) const override;
 
-  const HistogramInfo &getHistogramInfo() const { return Info; }
+  unsigned getOpcode() const { return Opcode; }
 
   /// Return the mask operand if one was provided, or a null pointer if all
   /// lanes should be executed unconditionally.

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

@@ -33,7 +33,6 @@
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/LoopUtils.h"
 #include "llvm/Transforms/Utils/ScalarEvolutionExpander.h"
-#include "llvm/Transforms/Vectorize/LoopVectorizationLegality.h"
 #include <cassert>
 
 using namespace llvm;
@@ -286,10 +285,6 @@ static Instruction *getInstructionForCost(const VPRecipeBase *R) {
   // the legacy model.
   if (auto *WidenMem = dyn_cast<VPWidenMemoryRecipe>(R))
     return &WidenMem->getIngredient();
-  // 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 nullptr;
 }
 
@@ -1087,6 +1082,39 @@ void VPHistogramRecipe::execute(VPTransformState &State) {
   }
 }
 
+InstructionCost VPHistogramRecipe::computeCost(ElementCount VF,
+                                               VPCostContext &Ctx) const {
+  // FIXME: Take the gather and scatter into account as well. For now we're
+  //        generating the same cost as the fallback path, but we'll likely
+  //        need to create a new TTI method for determining the cost, including
+  //        whether we can use base + vec-of-smaller-indices or just
+  //        vec-of-pointers.
+  assert(VF.isVector() && "Invalid VF for histogram cost");
+  Value *Address = getOperand(0)->getUnderlyingValue();
+  Value *IncAmt = getOperand(1)->getUnderlyingValue();
+  Type *IncTy = IncAmt->getType();
+  VectorType *VTy = VectorType::get(IncTy, VF);
+
+  // Assume that a non-constant update value (or a constant != 1) requires
+  // a multiply, and add that into the cost.
+  InstructionCost MulCost = TTI::TCC_Free;
+  ConstantInt *RHS = dyn_cast<ConstantInt>(IncAmt);
+  if (!RHS || RHS->getZExtValue() != 1)
+    MulCost = Ctx.TTI.getArithmeticInstrCost(Instruction::Mul, VTy);
+
+  // Find the cost of the histogram operation itself.
+  Type *PtrTy = VectorType::get(Address->getType(), VF);
+  Type *MaskTy = VectorType::get(Type::getInt1Ty(Ctx.LLVMCtx), VF);
+  IntrinsicCostAttributes ICA(Intrinsic::experimental_vector_histogram_add,
+                              Type::getVoidTy(Ctx.LLVMCtx),
+                              {PtrTy, IncTy, MaskTy});
+
+  // Add the costs together with the add/sub operation.
+  return Ctx.TTI.getIntrinsicInstrCost(
+             ICA, TargetTransformInfo::TCK_RecipThroughput) +
+         MulCost + Ctx.TTI.getArithmeticInstrCost(Opcode, VTy);
+}
+
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
 void VPHistogramRecipe::print(raw_ostream &O, const Twine &Indent,
                               VPSlotTracker &SlotTracker) const {

llvm/test/Transforms/LoopVectorize/AArch64/sve2-histcnt-epilogue.ll

@@ -59,13 +59,13 @@ entry:
 
 for.body:
   %iv = phi i64 [ 0, %entry ], [ %iv.next, %for.body ]
-  %arrayidx = getelementptr inbounds i32, ptr %indices, i64 %iv
-  %0 = load i32, ptr %arrayidx, align 4
-  %idxprom1 = zext i32 %0 to i64
-  %arrayidx2 = getelementptr inbounds i32, ptr %buckets, i64 %idxprom1
-  %1 = load i32, ptr %arrayidx2, align 4
-  %inc = add nsw i32 %1, 1
-  store i32 %inc, ptr %arrayidx2, align 4
+  %gep.indices = getelementptr inbounds i32, ptr %indices, i64 %iv
+  %l.idx = load i32, ptr %gep.indices, align 4
+  %idxprom1 = zext i32 %l.idx to i64
+  %gep.bucket = getelementptr inbounds i32, ptr %buckets, i64 %idxprom1
+  %l.bucket = load i32, ptr %gep.bucket, align 4
+  %inc = add nsw i32 %l.bucket, 1
+  store i32 %inc, ptr %gep.bucket, align 4
   %iv.next = add nuw nsw i64 %iv, 1
   %exitcond = icmp eq i64 %iv.next, %N
   br i1 %exitcond, label %for.exit, label %for.body, !llvm.loop !0

llvm/test/Transforms/LoopVectorize/AArch64/sve2-histcnt-no-scalar-interleave.ll

@@ -32,13 +32,13 @@ entry:
 
 for.body:
   %iv = phi i64 [ 0, %entry ], [ %iv.next, %for.body ]
-  %arrayidx = getelementptr inbounds i32, ptr %indices, i64 %iv
-  %0 = load i32, ptr %arrayidx, align 4
-  %idxprom1 = zext i32 %0 to i64
-  %arrayidx2 = getelementptr inbounds i32, ptr %buckets, i64 %idxprom1
-  %1 = load i32, ptr %arrayidx2, align 4
-  %inc = add nsw i32 %1, 1
-  store i32 %inc, ptr %arrayidx2, align 4
+  %gep.indices = getelementptr inbounds i32, ptr %indices, i64 %iv
+  %l.idx = load i32, ptr %gep.indices, align 4
+  %idxprom1 = zext i32 %l.idx to i64
+  %gep.bucket = getelementptr inbounds i32, ptr %buckets, i64 %idxprom1
+  %l.bucket = load i32, ptr %gep.bucket, align 4
+  %inc = add nsw i32 %l.bucket, 1
+  store i32 %inc, ptr %gep.bucket, align 4
   %iv.next = add nuw nsw i64 %iv, 1
   %exitcond = icmp eq i64 %iv.next, %N
   br i1 %exitcond, label %for.exit, label %for.body, !llvm.loop !0

llvm/test/Transforms/LoopVectorize/AArch64/sve2-histcnt-outerloop-scevaddrec.ll

@@ -0,0 +1,42 @@
+; RUN: opt < %s -passes=loop-vectorize,instcombine -enable-histogram-loop-vectorization -sve-gather-overhead=2 -sve-scatter-overhead=2 -debug-only=loop-vectorize -S 2>&1 | FileCheck %s
+; REQUIRES: asserts
+
+target triple = "aarch64-unknown-linux-gnu"
+
+;; Make sure we don't detect a histogram operation if the index address is
+;; loop invariant.
+; CHECK: LV: Checking for a histogram on:   store i32 %inc, ptr %gep.bucket, align 4
+; CHECK-NEXT: LV: Can't vectorize due to memory conflicts
+; CHECK-NEXT: LV: Not vectorizing: Cannot prove legality.
+
+define void @outer_loop_scevaddrec(ptr noalias %buckets, ptr readonly %indices, i64 %N, i64 %M) #0 {
+entry:
+  br label %outer.header
+
+outer.header:
+  %outer.iv = phi i64 [ 0, %entry ], [ %outer.iv.next, %outer.latch ]
+  %gep.indices = getelementptr inbounds i32, ptr %indices, i64 %outer.iv
+  br label %inner.body
+
+inner.body:
+  %iv = phi i64 [ 0, %outer.header ], [ %iv.next, %inner.body ]
+  %l.idx = load i32, ptr %gep.indices, align 4
+  %idxprom1 = zext i32 %l.idx to i64
+  %gep.bucket = getelementptr inbounds i32, ptr %buckets, i64 %idxprom1
+  %l.bucket = load i32, ptr %gep.bucket, align 4
+  %inc = add nsw i32 %l.bucket, 1
+  store i32 %inc, ptr %gep.bucket, align 4
+  %iv.next = add nuw nsw i64 %iv, 1
+  %exitcond = icmp eq i64 %iv.next, %N
+  br i1 %exitcond, label %outer.latch, label %inner.body
+
+outer.latch:
+  %outer.iv.next = add nuw nsw i64 %outer.iv, 1
+  %outer.exitcond = icmp eq i64 %outer.iv.next, %M
+  br i1 %outer.exitcond, label %outer.exit, label %outer.header
+
+outer.exit:
+  ret void
+}
+
+attributes #0 = { "target-features"="+sve2" vscale_range(1,16) }