Add new interface, switch for histogram support

Adds a separate interface to distinguish between the existing
semantics of canVectorizeMemory(), and the ability to do so if
the transform pass can handle histograms.

Also adds a flag to enable histogram support in LoopVecLegality,
defaulting to off for now.

Author: huntergr-arm

llvm/include/llvm/Analysis/LoopAccessAnalysis.h

@@ -639,6 +639,13 @@ class RuntimePointerChecking {
 /// Checks for both memory dependences and the SCEV predicates contained in the
 /// PSE must be emitted in order for the results of this analysis to be valid.
 class LoopAccessInfo {
+  /// Represents whether the memory access dependencies in the loop:
+  ///   * Prohibit vectorization
+  ///   * Allow for vectorization (possibly with runtime checks)
+  ///   * Allow for vectorization (possibly with runtime checks),
+  ///     as long as histogram operations are supported.
+  enum VecMemPossible { CantVec = 0, NormalVec = 1, HistogramVec = 2 };
+
 public:
   LoopAccessInfo(Loop *L, ScalarEvolution *SE, const TargetTransformInfo *TTI,
                  const TargetLibraryInfo *TLI, AAResults *AA, DominatorTree *DT,
@@ -650,7 +657,11 @@ class LoopAccessInfo {
   /// hasStoreStoreDependenceInvolvingLoopInvariantAddress and
   /// hasLoadStoreDependenceInvolvingLoopInvariantAddress also need to be
   /// checked.
-  bool canVectorizeMemory() const { return CanVecMem; }
+  bool canVectorizeMemory() const { return CanVecMem == NormalVec; }
+
+  bool canVectorizeMemoryWithHistogram() const {
+    return CanVecMem == NormalVec || CanVecMem == HistogramVec;
+  }
 
   /// Return true if there is a convergent operation in the loop. There may
   /// still be reported runtime pointer checks that would be required, but it is
@@ -733,8 +744,8 @@ class LoopAccessInfo {
 private:
   /// Analyze the loop. Returns true if all memory access in the loop can be
   /// vectorized.
-  bool analyzeLoop(AAResults *AA, LoopInfo *LI, const TargetLibraryInfo *TLI,
-                   DominatorTree *DT);
+  VecMemPossible analyzeLoop(AAResults *AA, LoopInfo *LI,
+                             const TargetLibraryInfo *TLI, DominatorTree *DT);
 
   /// Check if the structure of the loop allows it to be analyzed by this
   /// pass.
@@ -775,7 +786,7 @@ class LoopAccessInfo {
   unsigned NumStores = 0;
 
   /// Cache the result of analyzeLoop.
-  bool CanVecMem = false;
+  VecMemPossible CanVecMem = CantVec;
   bool HasConvergentOp = false;
 
   /// Indicator that there are two non vectorizable stores to the same uniform

llvm/lib/Analysis/LoopAccessAnalysis.cpp

@@ -2426,9 +2426,9 @@ bool LoopAccessInfo::canAnalyzeLoop() {
   return true;
 }
 
-bool LoopAccessInfo::analyzeLoop(AAResults *AA, LoopInfo *LI,
-                                 const TargetLibraryInfo *TLI,
-                                 DominatorTree *DT) {
+LoopAccessInfo::VecMemPossible
+LoopAccessInfo::analyzeLoop(AAResults *AA, LoopInfo *LI,
+                            const TargetLibraryInfo *TLI, DominatorTree *DT) {
   // Holds the Load and Store instructions.
   SmallVector<LoadInst *, 16> Loads;
   SmallVector<StoreInst *, 16> Stores;
@@ -2468,7 +2468,7 @@ bool LoopAccessInfo::analyzeLoop(AAResults *AA, LoopInfo *LI,
       // With both a non-vectorizable memory instruction and a convergent
       // operation, found in this loop, no reason to continue the search.
       if (HasComplexMemInst && HasConvergentOp)
-        return false;
+        return CantVec;
 
       // Avoid hitting recordAnalysis multiple times.
       if (HasComplexMemInst)
@@ -2544,7 +2544,7 @@ bool LoopAccessInfo::analyzeLoop(AAResults *AA, LoopInfo *LI,
   } // Next block.
 
   if (HasComplexMemInst)
-    return false;
+    return CantVec;
 
   // Now we have two lists that hold the loads and the stores.
   // Next, we find the pointers that they use.
@@ -2553,7 +2553,7 @@ bool LoopAccessInfo::analyzeLoop(AAResults *AA, LoopInfo *LI,
   // care if the pointers are *restrict*.
   if (!Stores.size()) {
     LLVM_DEBUG(dbgs() << "LAA: Found a read-only loop!\n");
-    return true;
+    return NormalVec;
   }
 
   MemoryDepChecker::DepCandidates DependentAccesses;
@@ -2606,7 +2606,7 @@ bool LoopAccessInfo::analyzeLoop(AAResults *AA, LoopInfo *LI,
     LLVM_DEBUG(
         dbgs() << "LAA: A loop annotated parallel, ignore memory dependency "
                << "checks.\n");
-    return true;
+    return NormalVec;
   }
 
   for (LoadInst *LD : Loads) {
@@ -2653,7 +2653,7 @@ bool LoopAccessInfo::analyzeLoop(AAResults *AA, LoopInfo *LI,
   // other reads in this loop then is it safe to vectorize.
   if (NumReadWrites == 1 && NumReads == 0) {
     LLVM_DEBUG(dbgs() << "LAA: Found a write-only loop!\n");
-    return true;
+    return NormalVec;
   }
 
   // Build dependence sets and check whether we need a runtime pointer bounds
@@ -2675,7 +2675,7 @@ bool LoopAccessInfo::analyzeLoop(AAResults *AA, LoopInfo *LI,
         << "cannot identify array bounds";
     LLVM_DEBUG(dbgs() << "LAA: We can't vectorize because we can't find "
                       << "the array bounds.\n");
-    return false;
+    return CantVec;
   }
 
   LLVM_DEBUG(
@@ -2708,7 +2708,7 @@ bool LoopAccessInfo::analyzeLoop(AAResults *AA, LoopInfo *LI,
         recordAnalysis("CantCheckMemDepsAtRunTime", I)
             << "cannot check memory dependencies at runtime";
         LLVM_DEBUG(dbgs() << "LAA: Can't vectorize with memory checks\n");
-        return false;
+        return CantVec;
       }
       DepsAreSafe = true;
     }
@@ -2719,19 +2719,19 @@ bool LoopAccessInfo::analyzeLoop(AAResults *AA, LoopInfo *LI,
         << "cannot add control dependency to convergent operation";
     LLVM_DEBUG(dbgs() << "LAA: We can't vectorize because a runtime check "
                          "would be needed with a convergent operation\n");
-    return false;
+    return CantVec;
   }
 
   if (DepsAreSafe) {
     LLVM_DEBUG(
         dbgs() << "LAA: No unsafe dependent memory operations in loop.  We"
                << (PtrRtChecking->Need ? "" : " don't")
                << " need runtime memory checks.\n");
-    return true;
+    return Histograms.empty() ? NormalVec : HistogramVec;
   }
 
   emitUnsafeDependenceRemark();
-  return false;
+  return CantVec;
 }
 
 void LoopAccessInfo::emitUnsafeDependenceRemark() {
@@ -3065,7 +3065,7 @@ LoopAccessInfo::LoopAccessInfo(Loop *L, ScalarEvolution *SE,
 }
 
 void LoopAccessInfo::print(raw_ostream &OS, unsigned Depth) const {
-  if (CanVecMem) {
+  if (CanVecMem != CantVec) {
     OS.indent(Depth) << "Memory dependences are safe";
     const MemoryDepChecker &DC = getDepChecker();
     if (!DC.isSafeForAnyVectorWidth())

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

@@ -78,6 +78,10 @@ static cl::opt<LoopVectorizeHints::ScalableForceKind>
                 "Scalable vectorization is available and favored when the "
                 "cost is inconclusive.")));
 
+static cl::opt<bool> EnableHistogramVectorization(
+    "enable-histogram-loop-vectorization", cl::init(false), cl::Hidden,
+    cl::desc("Enables autovectorization of some loops containing histograms"));
+
 /// Maximum vectorization interleave count.
 static const unsigned MaxInterleaveFactor = 16;
 
@@ -1065,7 +1069,9 @@ bool LoopVectorizationLegality::canVectorizeMemory() {
   }
 
   if (!LAI->canVectorizeMemory())
-    return false;
+    if (!EnableHistogramVectorization ||
+        !LAI->canVectorizeMemoryWithHistogram())
+      return false;
 
   if (LAI->hasLoadStoreDependenceInvolvingLoopInvariantAddress()) {
     reportVectorizationFailure("We don't allow storing to uniform addresses",

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

@@ -1,5 +1,5 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 3
-; RUN: opt < %s -passes=loop-vectorize -force-vector-interleave=1 -sve-gather-overhead=2 -sve-scatter-overhead=2 -S | FileCheck %s
+; RUN: opt < %s -passes=loop-vectorize -force-vector-interleave=1 -enable-histogram-loop-vectorization -sve-gather-overhead=2 -sve-scatter-overhead=2 -S | FileCheck %s
 
 target triple = "aarch64-unknown-linux-gnu"