[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/3] 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/3] 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);
   }

>From f32b0c660b7d4a5fbdd412d57abaf27b91f1c2c5 Mon Sep 17 00:00:00 2001
From: c8ef <c8ef at outlook.com>
Date: Wed, 12 Jun 2024 08:01:33 +0800
Subject: [PATCH 3/3] use m_BinOp instead of m_c_BinOp

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

diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index 65661ad005fad..36fd8c136a3c5 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -5239,15 +5239,15 @@ SDValue DAGCombiner::visitAVG(SDNode *N) {
   // fold avgu(zext(x), zext(y)) -> zext(avgu(x, y))
   SDValue A;
   SDValue B;
-  if (sd_match(N, m_c_BinOp(ISD::AVGFLOORU, m_ZExt(m_Value(A)),
-                            m_ZExt(m_Value(B)))) &&
+  if (sd_match(
+          N, m_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 (sd_match(N, m_c_BinOp(ISD::AVGCEILU, m_ZExt(m_Value(A)),
-                            m_ZExt(m_Value(B)))) &&
+  if (sd_match(
+          N, m_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);