[SLP] Allow targets to add cost for nonstandard conditions

[jrbyrnes]

There are conditions in which vectorization is profitable, but are not expressible by the current cost model. As an example, the vectorization profit may entirely be based on conditions of the users of the tree entry.

This gives targets a chance to express things of this nature.

[llvmbot]

@llvm/pr-subscribers-backend-amdgpu
@llvm/pr-subscribers-llvm-transforms

@llvm/pr-subscribers-llvm-analysis

Author: Jeffrey Byrnes (jrbyrnes)

Changes

There are conditions in which vectorization is profitable, but are not expressible by the current cost model. As an example, the vectorization profit may entirely be based on conditions of the users of the tree entry.

This gives targets a chance to express things of this nature.

---

Full diff: https://github.com/llvm/llvm-project/pull/95328.diff

5 Files Affected:

• (modified) llvm/include/llvm/Analysis/TargetTransformInfo.h (+16)
• (modified) llvm/include/llvm/Analysis/TargetTransformInfoImpl.h (+5)
• (modified) llvm/include/llvm/CodeGen/BasicTTIImpl.h (+5)
• (modified) llvm/lib/Analysis/TargetTransformInfo.cpp (+5)
• (modified) llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp (+8)

diff --git a/llvm/include/llvm/Analysis/TargetTransformInfo.h b/llvm/include/llvm/Analysis/TargetTransformInfo.h
index f55f21c94a85a..49117ca8c74c5 100644
--- a/llvm/include/llvm/Analysis/TargetTransformInfo.h
+++ b/llvm/include/llvm/Analysis/TargetTransformInfo.h
@@ -891,6 +891,11 @@ class TargetTransformInfo {
                                            bool Insert, bool Extract,
                                            TTI::TargetCostKind CostKind) const;
 
+  /// Whether or not there is any target-specific condition that imposes an
+  /// overhead for scalarization
+  bool hasScalarizationOverhead(ArrayRef<Value *> VL,
+                                std::pair<bool, bool> &ScalarizationKind) const;
+
   /// Estimate the overhead of scalarizing an instructions unique
   /// non-constant operands. The (potentially vector) types to use for each of
   /// argument are passes via Tys.
@@ -1921,6 +1926,10 @@ class TargetTransformInfo::Concept {
   getOperandsScalarizationOverhead(ArrayRef<const Value *> Args,
                                    ArrayRef<Type *> Tys,
                                    TargetCostKind CostKind) = 0;
+
+  virtual bool
+  hasScalarizationOverhead(ArrayRef<Value *> VL,
+                           std::pair<bool, bool> &ScalarizationKind) = 0;
   virtual bool supportsEfficientVectorElementLoadStore() = 0;
   virtual bool supportsTailCalls() = 0;
   virtual bool supportsTailCallFor(const CallBase *CB) = 0;
@@ -2456,6 +2465,13 @@ class TargetTransformInfo::Model final : public TargetTransformInfo::Concept {
     return Impl.getScalarizationOverhead(Ty, DemandedElts, Insert, Extract,
                                          CostKind);
   }
+
+  bool
+  hasScalarizationOverhead(ArrayRef<Value *> VL,
+                           std::pair<bool, bool> &ScalarizationKind) override {
+    return Impl.hasScalarizationOverhead(VL, ScalarizationKind);
+  }
+
   InstructionCost
   getOperandsScalarizationOverhead(ArrayRef<const Value *> Args,
                                    ArrayRef<Type *> Tys,
diff --git a/llvm/include/llvm/Analysis/TargetTransformInfoImpl.h b/llvm/include/llvm/Analysis/TargetTransformInfoImpl.h
index 7828bdc1f1f43..1d3e6752006d9 100644
--- a/llvm/include/llvm/Analysis/TargetTransformInfoImpl.h
+++ b/llvm/include/llvm/Analysis/TargetTransformInfoImpl.h
@@ -371,6 +371,11 @@ class TargetTransformInfoImplBase {
     return 0;
   }
 
+  bool hasScalarizationOverhead(ArrayRef<Value *> VL,
+                                std::pair<bool, bool> &ScalarizationKind) {
+    return false;
+  }
+
   InstructionCost
   getOperandsScalarizationOverhead(ArrayRef<const Value *> Args,
                                    ArrayRef<Type *> Tys,
diff --git a/llvm/include/llvm/CodeGen/BasicTTIImpl.h b/llvm/include/llvm/CodeGen/BasicTTIImpl.h
index 9f8d3ded9b3c1..2aa36c724bc03 100644
--- a/llvm/include/llvm/CodeGen/BasicTTIImpl.h
+++ b/llvm/include/llvm/CodeGen/BasicTTIImpl.h
@@ -807,6 +807,11 @@ class BasicTTIImplBase : public TargetTransformInfoImplCRTPBase<T> {
                                              CostKind);
   }
 
+  bool hasScalarizationOverhead(ArrayRef<Value *> VL,
+                                std::pair<bool, bool> &ScalarizationKind) {
+    return false;
+  }
+
   /// Estimate the overhead of scalarizing an instructions unique
   /// non-constant operands. The (potentially vector) types to use for each of
   /// argument are passes via Tys.
diff --git a/llvm/lib/Analysis/TargetTransformInfo.cpp b/llvm/lib/Analysis/TargetTransformInfo.cpp
index 7e721cbc87f3f..81b1e6b181bb0 100644
--- a/llvm/lib/Analysis/TargetTransformInfo.cpp
+++ b/llvm/lib/Analysis/TargetTransformInfo.cpp
@@ -594,6 +594,11 @@ InstructionCost TargetTransformInfo::getScalarizationOverhead(
                                            CostKind);
 }
 
+bool TargetTransformInfo::hasScalarizationOverhead(
+    ArrayRef<Value *> VL, std::pair<bool, bool> &ScalarizeKind) const {
+  return TTIImpl->hasScalarizationOverhead(VLm ScalarizeKind);
+}
+
 InstructionCost TargetTransformInfo::getOperandsScalarizationOverhead(
     ArrayRef<const Value *> Args, ArrayRef<Type *> Tys,
     TTI::TargetCostKind CostKind) const {
diff --git a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
index ae0819c964bef..f189e9b6ba14b 100644
--- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -9084,6 +9084,14 @@ BoUpSLP::getEntryCost(const TreeEntry *E, ArrayRef<Value *> VectorizedVals,
         E, ScalarTy, *TTI, VectorizedVals, *this, CheckedExtracts);
   }
   InstructionCost CommonCost = 0;
+  std::pair<bool, bool> ScalarizationKind(false, false);
+  if (TTI->hasScalarizationOverhead(VL, ScalarizationKind)) {
+    APInt DemandedElts = APInt::getAllOnes(VL.size());
+    CommonCost -= TTI->getScalarizationOverhead(
+        VecTy, DemandedElts,
+        /*Insert*/ ScalarizationKind.first,
+        /*Extract*/ ScalarizationKind.second, CostKind);
+  }
   SmallVector<int> Mask;
   bool IsReverseOrder = isReverseOrder(E->ReorderIndices);
   if (!E->ReorderIndices.empty() &&

[arsenm]

Include target implementation and test?

[arsenm]

What can this express that getOperandsScalarizationOverhead and getScalarizationOverhead do not?

I'm already confused by every function in TTI having multiple versions of ~everything

[jrbyrnes]

What can this express that getOperandsScalarizationOverhead and getScalarizationOverhead do not?

Those functions are used to calculate the cost of the scalarized sequence based on the inserts/extracts needed and the legalization costs based on TLI. For the purpose here, we can use either to calculate the cost of scalarization, but we need a mechanism to control whether or not to include this scalarization overhead for a particular tree entry (which is what this new hook is used for).

Even if we wanted to implement this control in those functions, I think we would still need to query targets as to whether or not there needs to be a cost accounting for the SelectionDAG issue workaround (assuming we could generalize the conditions in which the SelectionDAG issue occurred). Perhaps renaming the hook would help (e.g. hasSelectionDAGScalarizationOverhead)?

[jrbyrnes]

Add implementation of hasScalarizationOverhead and test.

Rebased on top of #91016 in order to facilitate the test