[DAG] visitTRUNCATE - merge the trunc(abd) and trunc(avg) handling which are almost identical (#154301)

CC @houngkoungting

Author: RKSimon

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

@@ -16279,40 +16279,6 @@ SDValue DAGCombiner::visitTRUNCATE(SDNode *N) {
   // because targets may prefer a wider type during later combines and invert
   // this transform.
   switch (N0.getOpcode()) {
-  case ISD::AVGCEILU:
-  case ISD::AVGFLOORU:
-    if (!LegalOperations && N0.hasOneUse() &&
-        TLI.isOperationLegal(N0.getOpcode(), VT)) {
-      SDValue X = N0.getOperand(0);
-      SDValue Y = N0.getOperand(1);
-      unsigned SrcBits = X.getScalarValueSizeInBits();
-      unsigned DstBits = VT.getScalarSizeInBits();
-      APInt UpperBits = APInt::getBitsSetFrom(SrcBits, DstBits);
-      if (DAG.MaskedValueIsZero(X, UpperBits) &&
-          DAG.MaskedValueIsZero(Y, UpperBits)) {
-        SDValue Tx = DAG.getNode(ISD::TRUNCATE, DL, VT, X);
-        SDValue Ty = DAG.getNode(ISD::TRUNCATE, DL, VT, Y);
-        return DAG.getNode(N0.getOpcode(), DL, VT, Tx, Ty);
-      }
-    }
-    break;
-  case ISD::AVGCEILS:
-  case ISD::AVGFLOORS:
-    if (!LegalOperations && N0.hasOneUse() &&
-        TLI.isOperationLegal(N0.getOpcode(), VT)) {
-      SDValue X = N0.getOperand(0);
-      SDValue Y = N0.getOperand(1);
-      unsigned SrcBits = X.getScalarValueSizeInBits();
-      unsigned DstBits = VT.getScalarSizeInBits();
-      unsigned NeededSignBits = SrcBits - DstBits + 1;
-      if (DAG.ComputeNumSignBits(X) >= NeededSignBits &&
-          DAG.ComputeNumSignBits(Y) >= NeededSignBits) {
-        SDValue Tx = DAG.getNode(ISD::TRUNCATE, DL, VT, X);
-        SDValue Ty = DAG.getNode(ISD::TRUNCATE, DL, VT, Y);
-        return DAG.getNode(N0.getOpcode(), DL, VT, Tx, Ty);
-      }
-    }
-    break;
   case ISD::ADD:
   case ISD::SUB:
   case ISD::MUL:
@@ -16361,10 +16327,15 @@ SDValue DAGCombiner::visitTRUNCATE(SDNode *N) {
                                  DAG, DL);
     }
     break;
-  case ISD::ABDU:
+  case ISD::AVGFLOORS:
+  case ISD::AVGFLOORU:
+  case ISD::AVGCEILS:
+  case ISD::AVGCEILU:
   case ISD::ABDS:
+  case ISD::ABDU:
+    // (trunc (avg a, b)) -> (avg (trunc a), (trunc b))
     // (trunc (abdu/abds a, b)) -> (abdu/abds (trunc a), (trunc b))
-    if ((!LegalOperations || N0.hasOneUse()) &&
+    if (!LegalOperations && N0.hasOneUse() &&
         TLI.isOperationLegal(N0.getOpcode(), VT)) {
       EVT TruncVT = VT;
       unsigned SrcBits = SrcVT.getScalarSizeInBits();
@@ -16374,7 +16345,8 @@ SDValue DAGCombiner::visitTRUNCATE(SDNode *N) {
       SDValue B = N0.getOperand(1);
       bool CanFold = false;
 
-      if (N0.getOpcode() == ISD::ABDU) {
+      if (N0.getOpcode() == ISD::AVGFLOORU || N0.getOpcode() == ISD::AVGCEILU ||
+          N0.getOpcode() == ISD::ABDU) {
         APInt UpperBits = APInt::getBitsSetFrom(SrcBits, TruncBits);
         CanFold = DAG.MaskedValueIsZero(B, UpperBits) &&
                   DAG.MaskedValueIsZero(A, UpperBits);