[AArch64] Expand llvm.histogram intrinsic to support umax, umin, and uadd.sat operations (llvm#138447)

This patch extends the llvm.histogram intrinsic to support additional
update operations beyond the existing add. Specifically, the new
supported operations are:

* umax: unsigned maximum

* umin: unsigned minimum

* uadd.sat: unsigned saturated addition

Based on the discussion from:


https://discourse.llvm.org/t/rfc-expanding-the-experimental-histogram-intrinsic/84673

Author: RonDahan101

llvm/docs/LangRef.rst

@@ -20545,6 +20545,9 @@ More update operation types may be added in the future.
 
     declare void @llvm.experimental.vector.histogram.add.v8p0.i32(<8 x ptr> %ptrs, i32 %inc, <8 x i1> %mask)
     declare void @llvm.experimental.vector.histogram.add.nxv2p0.i64(<vscale x 2 x ptr> %ptrs, i64 %inc, <vscale x 2 x i1> %mask)
+    declare void @llvm.experimental.vector.histogram.uadd.sat.v8p0.i32(<8 x ptr> %ptrs, i32 %inc, <8 x i1> %mask)
+    declare void @llvm.experimental.vector.histogram.umax.v8p0.i32(<8 x ptr> %ptrs, i32 %val, <8 x i1> %mask)
+    declare void @llvm.experimental.vector.histogram.umin.v8p0.i32(<8 x ptr> %ptrs, i32 %val, <8 x i1> %mask)
 
 Arguments:
 """"""""""

llvm/include/llvm/IR/Intrinsics.td

@@ -1968,6 +1968,24 @@ def int_experimental_vector_histogram_add : DefaultAttrsIntrinsic<[],
                                LLVMScalarOrSameVectorWidth<0, llvm_i1_ty>], // Mask
                              [ IntrArgMemOnly ]>;
 
+def int_experimental_vector_histogram_uadd_sat : DefaultAttrsIntrinsic<[],
+                             [ llvm_anyvector_ty, // Vector of pointers
+                               llvm_anyint_ty,    // Increment
+                               LLVMScalarOrSameVectorWidth<0, llvm_i1_ty>], // Mask
+                             [ IntrArgMemOnly ]>;
+
+def int_experimental_vector_histogram_umin : DefaultAttrsIntrinsic<[],
+                             [ llvm_anyvector_ty, // Vector of pointers
+                               llvm_anyint_ty,    // Update value
+                               LLVMScalarOrSameVectorWidth<0, llvm_i1_ty>], // Mask
+                             [ IntrArgMemOnly ]>;
+
+def int_experimental_vector_histogram_umax : DefaultAttrsIntrinsic<[],
+                             [ llvm_anyvector_ty, // Vector of pointers
+                               llvm_anyint_ty,    // Update value
+                               LLVMScalarOrSameVectorWidth<0, llvm_i1_ty>], // Mask
+                             [ IntrArgMemOnly ]>;
+
 // Experimental match
 def int_experimental_vector_match : DefaultAttrsIntrinsic<
                              [ LLVMScalarOrSameVectorWidth<0, llvm_i1_ty> ],

llvm/lib/Transforms/Scalar/ScalarizeMaskedMemIntrin.cpp

@@ -968,6 +968,29 @@ static void scalarizeMaskedVectorHistogram(const DataLayout &DL, CallInst *CI,
 
   // FIXME: Do we need to add an alignment parameter to the intrinsic?
   unsigned VectorWidth = AddrType->getNumElements();
+  auto CreateHistogramUpdateValue = [&](IntrinsicInst *CI, Value *Load,
+                                        Value *Inc) -> Value * {
+    Value *UpdateOp;
+    switch (CI->getIntrinsicID()) {
+    case Intrinsic::experimental_vector_histogram_add:
+      UpdateOp = Builder.CreateAdd(Load, Inc);
+      break;
+    case Intrinsic::experimental_vector_histogram_uadd_sat:
+      UpdateOp =
+          Builder.CreateIntrinsic(Intrinsic::uadd_sat, {EltTy}, {Load, Inc});
+      break;
+    case Intrinsic::experimental_vector_histogram_umin:
+      UpdateOp = Builder.CreateIntrinsic(Intrinsic::umin, {EltTy}, {Load, Inc});
+      break;
+    case Intrinsic::experimental_vector_histogram_umax:
+      UpdateOp = Builder.CreateIntrinsic(Intrinsic::umax, {EltTy}, {Load, Inc});
+      break;
+
+    default:
+      llvm_unreachable("Unexpected histogram intrinsic");
+    }
+    return UpdateOp;
+  };
 
   // Shorten the way if the mask is a vector of constants.
   if (isConstantIntVector(Mask)) {
@@ -976,8 +999,9 @@ static void scalarizeMaskedVectorHistogram(const DataLayout &DL, CallInst *CI,
         continue;
       Value *Ptr = Builder.CreateExtractElement(Ptrs, Idx, "Ptr" + Twine(Idx));
       LoadInst *Load = Builder.CreateLoad(EltTy, Ptr, "Load" + Twine(Idx));
-      Value *Add = Builder.CreateAdd(Load, Inc);
-      Builder.CreateStore(Add, Ptr);
+      Value *Update =
+          CreateHistogramUpdateValue(cast<IntrinsicInst>(CI), Load, Inc);
+      Builder.CreateStore(Update, Ptr);
     }
     CI->eraseFromParent();
     return;
@@ -997,8 +1021,9 @@ static void scalarizeMaskedVectorHistogram(const DataLayout &DL, CallInst *CI,
     Builder.SetInsertPoint(CondBlock->getTerminator());
     Value *Ptr = Builder.CreateExtractElement(Ptrs, Idx, "Ptr" + Twine(Idx));
     LoadInst *Load = Builder.CreateLoad(EltTy, Ptr, "Load" + Twine(Idx));
-    Value *Add = Builder.CreateAdd(Load, Inc);
-    Builder.CreateStore(Add, Ptr);
+    Value *UpdateOp =
+        CreateHistogramUpdateValue(cast<IntrinsicInst>(CI), Load, Inc);
+    Builder.CreateStore(UpdateOp, Ptr);
 
     // Create "else" block, fill it in the next iteration
     BasicBlock *NewIfBlock = ThenTerm->getSuccessor(0);
@@ -1089,6 +1114,9 @@ static bool optimizeCallInst(CallInst *CI, bool &ModifiedDT,
     default:
       break;
     case Intrinsic::experimental_vector_histogram_add:
+    case Intrinsic::experimental_vector_histogram_uadd_sat:
+    case Intrinsic::experimental_vector_histogram_umin:
+    case Intrinsic::experimental_vector_histogram_umax:
       if (TTI.isLegalMaskedVectorHistogram(CI->getArgOperand(0)->getType(),
                                            CI->getArgOperand(1)->getType()))
         return false;