[llvm] Add target hook for automatic histogram vectorization (PR #128414)

[via llvm-commits]

https://github.com/RonDahan101 created https://github.com/llvm/llvm-project/pull/128414

Now, each target, depending on the load/store instruction, can indicate whether it is safe to vectorize histogram operations.

For example, this allows distinguishing between different pointer address spaces.

>From 3928d941b2444abd803839bef6b688cab2fe1c88 Mon Sep 17 00:00:00 2001
From: rond <ron.dahan at mobileye.com>
Date: Sun, 23 Feb 2025 12:28:47 +0200
Subject: [PATCH] Add target hook for automatic histogram vectorization

Now, each target, depending on the load/store instruction, can indicate
whether it is safe to vectorize histogram operations.

For example, this allows distinguishing between different pointer address spaces.
---
 llvm/include/llvm/Analysis/TargetTransformInfo.h |  9 +++++++++
 .../llvm/Analysis/TargetTransformInfoImpl.h      | 16 ++++++++++++++++
 llvm/lib/Analysis/TargetTransformInfo.cpp        |  5 +++++
 .../Vectorize/LoopVectorizationLegality.cpp      |  3 +++
 4 files changed, 33 insertions(+)

diff --git a/llvm/include/llvm/Analysis/TargetTransformInfo.h b/llvm/include/llvm/Analysis/TargetTransformInfo.h
index 9048481b49189..94a565a99ab41 100644
--- a/llvm/include/llvm/Analysis/TargetTransformInfo.h
+++ b/llvm/include/llvm/Analysis/TargetTransformInfo.h
@@ -805,6 +805,9 @@ class TargetTransformInfo {
   bool isLegalMaskedScatter(Type *DataType, Align Alignment) const;
   /// Return true if the target supports masked gather.
   bool isLegalMaskedGather(Type *DataType, Align Alignment) const;
+  /// Return true if the target supports vectorization of histograms.
+  bool isLegalForHistogramVectorization(const LoadInst *LI,
+                                        const StoreInst *SI) const;
   /// Return true if the target forces scalarizing of llvm.masked.gather
   /// intrinsics.
   bool forceScalarizeMaskedGather(VectorType *Type, Align Alignment) const;
@@ -2028,6 +2031,8 @@ class TargetTransformInfo::Concept {
                                     ElementCount NumElements) const = 0;
   virtual bool isLegalMaskedScatter(Type *DataType, Align Alignment) = 0;
   virtual bool isLegalMaskedGather(Type *DataType, Align Alignment) = 0;
+  virtual bool isLegalForHistogramVectorization(const LoadInst *LI,
+                                                const StoreInst *SI) = 0;
   virtual bool forceScalarizeMaskedGather(VectorType *DataType,
                                           Align Alignment) = 0;
   virtual bool forceScalarizeMaskedScatter(VectorType *DataType,
@@ -2589,6 +2594,10 @@ class TargetTransformInfo::Model final : public TargetTransformInfo::Concept {
   bool isLegalMaskedGather(Type *DataType, Align Alignment) override {
     return Impl.isLegalMaskedGather(DataType, Alignment);
   }
+  bool isLegalForHistogramVectorization(const LoadInst *LI,
+                                        const StoreInst *SI) override {
+    return Impl.isLegalForHistogramVectorization(LI, SI);
+  }
   bool forceScalarizeMaskedGather(VectorType *DataType,
                                   Align Alignment) override {
     return Impl.forceScalarizeMaskedGather(DataType, Alignment);
diff --git a/llvm/include/llvm/Analysis/TargetTransformInfoImpl.h b/llvm/include/llvm/Analysis/TargetTransformInfoImpl.h
index a8d6dd18266bb..de93612ffffda 100644
--- a/llvm/include/llvm/Analysis/TargetTransformInfoImpl.h
+++ b/llvm/include/llvm/Analysis/TargetTransformInfoImpl.h
@@ -310,6 +310,22 @@ class TargetTransformInfoImplBase {
     return false;
   }
 
+  bool isLegalForHistogramVectorization(const LoadInst *LI,
+                                        const StoreInst *SI) const {
+    // TODO: Currently, when a target wants to use histogram intrinsics, it adds
+    // the
+    // `-enable-histogram-loop-vectorization` flag to Clang.
+    //
+    // Once the histogram option is enabled by default, we will need to update
+    // the default hook to return `false`, and each target that wants automatic
+    // histogram vectorization will need to override it to return `true`.
+    //
+    // Additionally, we will need to deprecate the
+    // `-enable-histogram-loop-vectorization` flag, as it will no longer be
+    // necessary.
+    return true;
+  }
+
   bool forceScalarizeMaskedGather(VectorType *DataType, Align Alignment) const {
     return false;
   }
diff --git a/llvm/lib/Analysis/TargetTransformInfo.cpp b/llvm/lib/Analysis/TargetTransformInfo.cpp
index 7df7038f6dd47..7a020d55abad2 100644
--- a/llvm/lib/Analysis/TargetTransformInfo.cpp
+++ b/llvm/lib/Analysis/TargetTransformInfo.cpp
@@ -492,6 +492,11 @@ bool TargetTransformInfo::isLegalMaskedGather(Type *DataType,
   return TTIImpl->isLegalMaskedGather(DataType, Alignment);
 }
 
+bool TargetTransformInfo::isLegalForHistogramVectorization(
+    const LoadInst *LI, const StoreInst *SI) const {
+  return TTIImpl->isLegalForHistogramVectorization(LI, SI);
+}
+
 bool TargetTransformInfo::isLegalAltInstr(
     VectorType *VecTy, unsigned Opcode0, unsigned Opcode1,
     const SmallBitVector &OpcodeMask) const {
diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp
index 420cbc5384ce4..4045fd315d5bd 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp
@@ -1193,6 +1193,9 @@ bool LoopVectorizationLegality::canVectorizeIndirectUnsafeDependences() {
   if (!LI || !SI)
     return false;
 
+  if (!TTI->isLegalForHistogramVectorization(LI, SI))
+    return false;
+
   LLVM_DEBUG(dbgs() << "LV: Checking for a histogram on: " << *SI << "\n");
   return findHistogram(LI, SI, TheLoop, LAI->getPSE(), Histograms);
 }