[llvm] [SelectionDAG] Fold (avg x, 0) -> x >> 1 (PR #85581)

[via llvm-commits]

https://github.com/AtariDreams updated https://github.com/llvm/llvm-project/pull/85581

>From 0952dc09c173b20c0571da90c78871bec2881f30 Mon Sep 17 00:00:00 2001
From: Rose <gfunni234 at gmail.com>
Date: Sun, 17 Mar 2024 15:31:39 -0400
Subject: [PATCH] [SelectionDAG] Fold (avg x, 0) -> x >> 1

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

diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index 5eb53d57c9c2bf..84554bdc68a7eb 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -5083,7 +5083,19 @@ SDValue DAGCombiner::visitAVG(SDNode *N) {
   if (N0 == N1 && Level >= AfterLegalizeTypes)
     return N0;
 
-  // TODO If we use avg for scalars anywhere, we can add (avgfl x, 0) -> x >> 1
+  // fold (avgfloor x, 0) -> x >> 1
+  if (isNullOrNullSplat(N1)) {
+    if (Opcode == ISD::AVGFLOORS ||
+        (Opcode == ISD::AVGCEILS &&
+         DAG.MaskedValueIsZero(
+             N0, APInt::getOneBitSet(VT.getScalarSizeInBits(), 1))))
+      return DAG.getNode(ISD::SRA, DL, VT, N0, DAG.getConstant(1, DL, VT));
+    if (Opcode == ISD::AVGFLOORU ||
+        (Opcode == ISD::AVGCEILU &&
+         DAG.MaskedValueIsZero(
+             N0, APInt::getOneBitSet(VT.getScalarSizeInBits(), 1)))))
+      return DAG.getNode(ISD::SRL, DL, VT, N0, DAG.getConstant(1, DL, VT));
+  }
 
   return SDValue();
 }