[llvm] [X86][PartialReduction] Lower zext-byte add reductions to vpsadbw (PR #201076)

[via llvm-commits]

================
@@ -353,6 +355,141 @@ bool X86PartialReduction::trySADReplacement(Instruction *Op) {
   return true;
 }
 
+bool X86PartialReduction::tryByteSumReplacement(Instruction *Op) {
+  // Gate on SSE2.
+  if (!ST->hasSSE2())
+    return false;
+
+  // Element type must be i32 or i64.
+  // - i32 is the leaf of the canonical i32 add-reduction tree LV emits
+  //   when the source code uses a uint32_t accumulator.
+  // - i64 is what LV emits directly (skipping the i32 intermediate) when
+  //   the source code uses a uint64_t accumulator: see for example the
+  //   `uint64_t sum += a[i:i8]` shape, where LV produces a wide
+  //   `zext <N x i8> to <N x i64>` feeding an i64 add-reduce.  psadbw's
+  //   natural output is already i64-per-8-byte-lane, so handling this
+  //   case requires only skipping the i32 bitcast below.
+  auto *OpTy = dyn_cast<FixedVectorType>(Op->getType());
+  if (!OpTy)
+    return false;
+  unsigned ElemBits = OpTy->getElementType()->getScalarSizeInBits();
+  if (ElemBits != 32 && ElemBits != 64)
+    return false;
+
+  // Must be a zext from i8.
+  auto *ZExt = dyn_cast<ZExtInst>(Op);
+  if (!ZExt)
+    return false;
+
+  auto *SrcTy = dyn_cast<FixedVectorType>(ZExt->getOperand(0)->getType());
+  if (!SrcTy || !SrcTy->getElementType()->isIntegerTy(8))
+    return false;
+
+  unsigned NumElts = OpTy->getNumElements();
+
+  // Only match shapes that are at least one native psadbw width wide.
+  // Below that (e.g. <8 x i8>) SelectionDAG's own SAD matcher already
+  // produces optimal code and padding up to 16 elements here can actually
+  // regress small cases by forcing wider horizontal reduction sequences.
+  if (NumElts < 16)
+    return false;
+
+  // Select the widest psadbw intrinsic the subtarget supports that can cover
+  // the input in whole 16/32/64-byte slices.  Same dispatch shape as the
+  // sibling trySADReplacement matcher.
+  unsigned IntrinsicNumElts;
+  Intrinsic::ID IID;
+  if (ST->hasBWI() && NumElts >= 64) {
----------------
mbhade-amd wrote:

Thanks @topperc, i will make this change.

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