[llvm] [DAG] Fold trunc(abdu(x, y)) and trunc(abds(x, y)) if they have sufficient leading zero/sign bits (PR #151471)

[Simon Pilgrim via llvm-commits]

================
@@ -3925,6 +3926,47 @@ SDValue DAGCombiner::foldSubToUSubSat(EVT DstVT, SDNode *N, const SDLoc &DL) {
   return SDValue();
 }
 
+// trunc (ABDU/S A, B)) → ABDU/S (trunc A), (trunc B)
+SDValue DAGCombiner::foldAbdToNarrowType(EVT VT, SDNode *N, const SDLoc &DL) {
+  SDValue Op = N->getOperand(0);
+
+  unsigned Opcode = Op.getOpcode();
+  if (Opcode != ISD::ABDU && Opcode != ISD::ABDS)
+    return SDValue();
+
+  SDValue Operand0 = Op.getOperand(0);
+  SDValue Operand1 = Op.getOperand(1);
+
+  // Early exit if either operand is zero.
+  if (ISD::isBuildVectorAllZeros(Operand0.getNode()) ||
+      ISD::isBuildVectorAllZeros(Operand1.getNode()))
+    return SDValue();
+
+  EVT SrcVT = Op.getValueType();
+  EVT TruncVT = N->getValueType(0);
+  unsigned NumSrcBits = SrcVT.getScalarSizeInBits();
+  unsigned NumTruncBits = TruncVT.getScalarSizeInBits();
+  unsigned NeededBits = NumSrcBits - NumTruncBits;
+
+  bool CanFold = false;
+
+  if (Opcode == ISD::ABDU) {
+    KnownBits Known = DAG.computeKnownBits(Op);
+    CanFold = Known.countMinLeadingZeros() >= NeededBits;
+  } else {
+    unsigned SignBits = DAG.ComputeNumSignBits(Op);
+    CanFold = SignBits >= NeededBits;
+  }
+
+  if (CanFold) {
----------------
RKSimon wrote:

Just because we can fold doesn't mean we should - look at the switch statement at the bottom of visitTRUNCATE - it has a series of legality/profitability checks for different opcodes so that we don't always fold trunc(abd(x,y)) -> abd(trunc(x),trunc(y)) if it'd be more costly.

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