[llvm] [DAG] fold avgu(zext(x), zext(y)) -> zext(avgu(x, y)) (PR #95134)

[via llvm-commits]

https://github.com/c8ef updated https://github.com/llvm/llvm-project/pull/95134

>From 4ef1f7bb730381b57a7b9cd74e142ddfa9e2235e Mon Sep 17 00:00:00 2001
From: c8ef <c8ef at outlook.com>
Date: Tue, 11 Jun 2024 23:35:21 +0800
Subject: [PATCH 1/2] fold avgu(zext(x), zext(y)) -> zext(avgu(x, y))

---
 llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index 4fcbe08e4b2b9..0a78803357410 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -5236,6 +5236,21 @@ SDValue DAGCombiner::visitAVG(SDNode *N) {
     return DAG.getNode(ISD::SRL, DL, VT, X,
                        DAG.getShiftAmountConstant(1, VT, DL));
 
+  // fold avgu(zext(x), zext(y)) -> zext(avgu(x, y))
+  SDValue A;
+  SDValue B;
+  if (hasOperation(ISD::AVGFLOORU, VT) &&
+      sd_match(N, m_c_BinOp(ISD::AVGFLOORU, m_ZExt(m_Value(A)),
+                            m_ZExt(m_Value(B))))) {
+    SDValue AvgFloorU = DAG.getNode(ISD::AVGFLOORU, DL, A.getValueType(), A, B);
+    return DAG.getNode(ISD::ZERO_EXTEND, DL, VT, AvgFloorU);
+  }
+  if (hasOperation(ISD::AVGCEILU, VT) &&
+      sd_match(N, m_c_BinOp(ISD::AVGCEILU, m_ZExt(m_Value(A)),
+                            m_ZExt(m_Value(B))))) {
+    SDValue AvgCeilU = DAG.getNode(ISD::AVGCEILU, DL, A.getValueType(), A, B);
+    return DAG.getNode(ISD::ZERO_EXTEND, DL, VT, AvgCeilU);
+  }
   return SDValue();
 }
 

>From 8e04c8ea04bb165ee3fd7239dc0ef10ceaba44f8 Mon Sep 17 00:00:00 2001
From: c8ef <c8ef at outlook.com>
Date: Tue, 11 Jun 2024 23:54:02 +0800
Subject: [PATCH 2/2] check extend from same type; check the original type

---
 llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index 0a78803357410..65661ad005fad 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -5239,15 +5239,17 @@ SDValue DAGCombiner::visitAVG(SDNode *N) {
   // fold avgu(zext(x), zext(y)) -> zext(avgu(x, y))
   SDValue A;
   SDValue B;
-  if (hasOperation(ISD::AVGFLOORU, VT) &&
-      sd_match(N, m_c_BinOp(ISD::AVGFLOORU, m_ZExt(m_Value(A)),
-                            m_ZExt(m_Value(B))))) {
+  if (sd_match(N, m_c_BinOp(ISD::AVGFLOORU, m_ZExt(m_Value(A)),
+                            m_ZExt(m_Value(B)))) &&
+      A.getValueType() == B.getValueType() &&
+      hasOperation(ISD::AVGFLOORU, A.getValueType())) {
     SDValue AvgFloorU = DAG.getNode(ISD::AVGFLOORU, DL, A.getValueType(), A, B);
     return DAG.getNode(ISD::ZERO_EXTEND, DL, VT, AvgFloorU);
   }
-  if (hasOperation(ISD::AVGCEILU, VT) &&
-      sd_match(N, m_c_BinOp(ISD::AVGCEILU, m_ZExt(m_Value(A)),
-                            m_ZExt(m_Value(B))))) {
+  if (sd_match(N, m_c_BinOp(ISD::AVGCEILU, m_ZExt(m_Value(A)),
+                            m_ZExt(m_Value(B)))) &&
+      A.getValueType() == B.getValueType() &&
+      hasOperation(ISD::AVGCEILU, A.getValueType())) {
     SDValue AvgCeilU = DAG.getNode(ISD::AVGCEILU, DL, A.getValueType(), A, B);
     return DAG.getNode(ISD::ZERO_EXTEND, DL, VT, AvgCeilU);
   }