[llvm] [DAG] Fold mismatched widened avg idioms to narrow form (#147946) (PR #163366)

[Simon Pilgrim via llvm-commits]

================
@@ -16481,10 +16481,57 @@ SDValue DAGCombiner::visitTRUNCATE(SDNode *N) {
                                  DAG, DL);
     }
     break;
-  case ISD::AVGFLOORS:
-  case ISD::AVGFLOORU:
   case ISD::AVGCEILS:
   case ISD::AVGCEILU:
+    // trunc (avgceilu (sext (x), sext (y))) -> avgceils(x, y)
+    // trunc (avgceils (zext (x), zext (y))) -> avgceilu(x, y)
+    if (N0.hasOneUse()) {
+      SDValue Op0 = N0.getOperand(0);
+      SDValue Op1 = N0.getOperand(1);
+      if (N0.getOpcode() == ISD::AVGCEILU) {
+        if (TLI.isOperationLegalOrCustom(ISD::AVGCEILS, VT) &&
+            Op0.getOperand(0).getValueType() == VT &&
+            Op1.getOperand(0).getValueType() == VT) {
+          if (Op0.getOpcode() == ISD::SIGN_EXTEND &&
+              Op1.getOpcode() == ISD::SIGN_EXTEND)
+            return DAG.getNode(ISD::AVGCEILS, DL, VT, Op0.getOperand(0),
+                               Op1.getOperand(0));
+
+          if (Op0.getOpcode() == ISD::SIGN_EXTEND_INREG &&
+              Op1.getOpcode() == ISD::SIGN_EXTEND_INREG) {
+            EVT VT0 = cast<VTSDNode>(Op0.getOperand(1))->getVT();
+            EVT VT1 = cast<VTSDNode>(Op1.getOperand(1))->getVT();
+            if (VT0 == VT && VT1 == VT) {
+              SDValue Op0Input = Op0.getOperand(0);
+              SDValue Op1Input = Op1.getOperand(0);
+              if (Op0Input.getOpcode() == ISD::TRUNCATE &&
+                  Op1Input.getOpcode() == ISD::TRUNCATE) {
+                SDValue Op0Pre = Op0Input.getOperand(0);
+                SDValue Op1Pre = Op1Input.getOperand(0);
+                if (Op0Pre.getOpcode() == ISD::SIGN_EXTEND &&
+                    Op1Pre.getOpcode() == ISD::SIGN_EXTEND) {
+                  SDValue X = Op0Pre.getOperand(0);
+                  SDValue Y = Op1Pre.getOperand(0);
+                  if (X.getValueType() == VT && Y.getValueType() == VT)
+                    return DAG.getNode(ISD::AVGCEILS, DL, VT, X, Y);
+                }
+              }
+            }
+          }
+        }
+      } else {
+        if (TLI.isOperationLegalOrCustom(ISD::AVGCEILU, VT) &&
+            Op0.getOperand(0).getValueType() == VT &&
+            Op1.getOperand(0).getValueType() == VT &&
+            Op0.getOpcode() == ISD::ZERO_EXTEND &&
+            Op1.getOpcode() == ISD::ZERO_EXTEND)
----------------
RKSimon wrote:

Check getOpcode BEFORE getOperand() - its unlikely for operand(0) but we don't know if the SDValue has any operands

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