[llvm] LV: Expand llvm.histogram intrinsic to support umax, umin, and uadd.sat operations (PR #127399)

[via llvm-commits]

================
@@ -1032,34 +1032,44 @@ bool LoopVectorizationLegality::canVectorizeInstr(Instruction &I) {
 
 /// Find histogram operations that match high-level code in loops:
 /// \code
-/// buckets[indices[i]]+=step;
+/// buckets[indices[i]] = UpdateOp(buckets[indices[i]], Val);
 /// \endcode
+/// Where UpdateOp can be add, sub, uadd.sat, umax, or umin.
 ///
 /// It matches a pattern starting from \p HSt, which Stores to the 'buckets'
-/// array the computed histogram. It uses a BinOp to sum all counts, storing
-/// them using a loop-variant index Load from the 'indices' input array.
+/// array the computed histogram. It uses an update instruction to update all
+/// counts, storing them using a loop-variant index Load from the 'indices'
+/// input array.
 ///
 /// On successful matches it updates the STATISTIC 'HistogramsDetected',
 /// regardless of hardware support. When there is support, it additionally
-/// stores the BinOp/Load pairs in \p HistogramCounts, as well the pointers
+/// stores the UpdateOp/Load pairs in \p HistogramCounts, as well the pointers
 /// used to update histogram in \p HistogramPtrs.
 static bool findHistogram(LoadInst *LI, StoreInst *HSt, Loop *TheLoop,
                           const PredicatedScalarEvolution &PSE,
                           SmallVectorImpl<HistogramInfo> &Histograms) {
 
-  // Store value must come from a Binary Operation.
+  // Store value must come from an update operation (binary op or intrinsic).
   Instruction *HPtrInstr = nullptr;
-  BinaryOperator *HBinOp = nullptr;
-  if (!match(HSt, m_Store(m_BinOp(HBinOp), m_Instruction(HPtrInstr))))
+  Instruction *HUpdateOp = nullptr;
+  if (!match(HSt, m_Store(m_Instruction(HUpdateOp), m_Instruction(HPtrInstr))))
     return false;
 
-  // BinOp must be an Add or a Sub modifying the bucket value by a
-  // loop invariant amount.
+  // The update operation must modify the bucket value by a loop invariant
+  // amount. Supported operations: add, sub, uadd.sat, umax, umin.
   // FIXME: We assume the loop invariant term is on the RHS.
   //        Fine for an immediate/constant, but maybe not a generic value?
   Value *HIncVal = nullptr;
-  if (!match(HBinOp, m_Add(m_Load(m_Specific(HPtrInstr)), m_Value(HIncVal))) &&
-      !match(HBinOp, m_Sub(m_Load(m_Specific(HPtrInstr)), m_Value(HIncVal))))
+  if (!match(HUpdateOp,
+             m_Add(m_Load(m_Specific(HPtrInstr)), m_Value(HIncVal))) &&
+      !match(HUpdateOp,
+             m_Sub(m_Load(m_Specific(HPtrInstr)), m_Value(HIncVal))) &&
+      !match(HUpdateOp, m_Intrinsic<Intrinsic::uadd_sat>(
+                            m_Load(m_Specific(HPtrInstr)), m_Value(HIncVal))) &&
+      !match(HUpdateOp, m_Intrinsic<Intrinsic::umax>(
----------------
RonDahan101 wrote:

Switched to  m_UMax / m_UMin . I tried variadic  m_CombineOr  too, but it only returns  bool , so recovering which arm matched needs a bind variable per arm plus a chain testing which is non-null — 30 lines vs 17. Since the kind is now derived at the match site, I kept the  else if  chain. Happy to switch if you prefer.

https://github.com/llvm/llvm-project/pull/127399