[LV] Autovectorization for the all-in-one histogram intrinsic (#91458)

This patch implements limited loop vectorization support for the 'all-in-one' histogram intrinsic. The feature is disabled by default, and when enabled will only vectorize if there are no other users of values in the gather-modify-scatter sequence.

Author: huntergr-arm

llvm/include/llvm/Analysis/LoopAccessAnalysis.h

@@ -144,7 +144,9 @@ class MemoryDepChecker {
       // on MinDepDistBytes.
       BackwardVectorizable,
       // Same, but may prevent store-to-load forwarding.
-      BackwardVectorizableButPreventsForwarding
+      BackwardVectorizableButPreventsForwarding,
+      // Access is to a loop loaded value, but is part of a histogram operation.
+      Histogram
     };
 
     /// String version of the types.
@@ -201,7 +203,8 @@ class MemoryDepChecker {
   /// Only checks sets with elements in \p CheckDeps.
   bool areDepsSafe(DepCandidates &AccessSets, MemAccessInfoList &CheckDeps,
                    const DenseMap<Value *, SmallVector<const Value *, 16>>
-                       &UnderlyingObjects);
+                       &UnderlyingObjects,
+                   const SmallPtrSetImpl<const Value *> &HistogramPtrs);
 
   /// No memory dependence was encountered that would inhibit
   /// vectorization.
@@ -352,7 +355,8 @@ class MemoryDepChecker {
   isDependent(const MemAccessInfo &A, unsigned AIdx, const MemAccessInfo &B,
               unsigned BIdx,
               const DenseMap<Value *, SmallVector<const Value *, 16>>
-                  &UnderlyingObjects);
+                  &UnderlyingObjects,
+              const SmallPtrSetImpl<const Value *> &HistogramPtrs);
 
   /// Check whether the data dependence could prevent store-load
   /// forwarding.
@@ -393,7 +397,8 @@ class MemoryDepChecker {
       const MemAccessInfo &A, Instruction *AInst, const MemAccessInfo &B,
       Instruction *BInst,
       const DenseMap<Value *, SmallVector<const Value *, 16>>
-          &UnderlyingObjects);
+          &UnderlyingObjects,
+      const SmallPtrSetImpl<const Value *> &HistogramPtrs);
 };
 
 class RuntimePointerChecking;
@@ -445,6 +450,15 @@ struct PointerDiffInfo {
         NeedsFreeze(NeedsFreeze) {}
 };
 
+struct HistogramInfo {
+  LoadInst *Load;
+  Instruction *Update;
+  StoreInst *Store;
+
+  HistogramInfo(LoadInst *Load, Instruction *Update, StoreInst *Store)
+      : Load(Load), Update(Update), Store(Store) {}
+};
+
 /// Holds information about the memory runtime legality checks to verify
 /// that a group of pointers do not overlap.
 class RuntimePointerChecking {
@@ -625,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,
@@ -636,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
@@ -664,6 +689,10 @@ class LoopAccessInfo {
   unsigned getNumStores() const { return NumStores; }
   unsigned getNumLoads() const { return NumLoads;}
 
+  const SmallVectorImpl<HistogramInfo> &getHistograms() const {
+    return Histograms;
+  }
+
   /// The diagnostics report generated for the analysis.  E.g. why we
   /// couldn't analyze the loop.
   const OptimizationRemarkAnalysis *getReport() const { return Report.get(); }
@@ -715,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.
@@ -757,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
@@ -777,6 +806,13 @@ class LoopAccessInfo {
   /// If an access has a symbolic strides, this maps the pointer value to
   /// the stride symbol.
   DenseMap<Value *, const SCEV *> SymbolicStrides;
+
+  /// Holds the load, update, and store instructions for all histogram-style
+  /// operations found in the loop.
+  SmallVector<HistogramInfo, 2> Histograms;
+
+  /// Storing Histogram Pointers
+  SmallPtrSet<const Value *, 2> HistogramPtrs;
 };
 
 /// Return the SCEV corresponding to a pointer with the symbolic stride

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

@@ -390,6 +390,23 @@ class LoopVectorizationLegality {
   unsigned getNumStores() const { return LAI->getNumStores(); }
   unsigned getNumLoads() const { return LAI->getNumLoads(); }
 
+  std::optional<const HistogramInfo *> getHistogramInfo(Instruction *I) const {
+    for (const HistogramInfo &HGram : LAI->getHistograms())
+      if (HGram.Load == I || HGram.Update == I || HGram.Store == I)
+        return &HGram;
+
+    return std::nullopt;
+  }
+
+  std::optional<const HistogramInfo *>
+  getHistogramForStore(StoreInst *SI) const {
+    for (const HistogramInfo &HGram : LAI->getHistograms())
+      if (HGram.Store == SI)
+        return &HGram;
+
+    return std::nullopt;
+  }
+
   PredicatedScalarEvolution *getPredicatedScalarEvolution() const {
     return &PSE;
   }