[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 7d142d48444e3a1a176096cb800c3822de1b2c7f 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 | 13 ++++++++++++-
 llvm/test/CodeGen/AArch64/hadd-combine.ll     |  6 ++----
 2 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index 5eb53d57c9c2bf..6a5509dd902449 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -5083,7 +5083,18 @@ 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 (avg x, 0) -> x >> 1
+  if (isNullOrNullSplat(N0))
+    return DAG.getNode((Opcode == ISD::AVGFLOORS || Opcode == ISD::AVGCEILS)
+                           ? ISD::SRA
+                           : ISD::SRL,
+                       DL, VT, N1, DAG.getConstant(1, DL, VT));
+
+  if (isNullOrNullSplat(N1))
+    return DAG.getNode((Opcode == ISD::AVGFLOORS || Opcode == ISD::AVGCEILS)
+                           ? ISD::SRA
+                           : ISD::SRL,
+                       DL, VT, N0, DAG.getConstant(1, DL, VT));
 
   return SDValue();
 }
diff --git a/llvm/test/CodeGen/AArch64/hadd-combine.ll b/llvm/test/CodeGen/AArch64/hadd-combine.ll
index 2269d75cdbb9ed..65b7e7e8b5188c 100644
--- a/llvm/test/CodeGen/AArch64/hadd-combine.ll
+++ b/llvm/test/CodeGen/AArch64/hadd-combine.ll
@@ -464,8 +464,7 @@ define <8 x i16> @rhaddu_i_const_lhs(<8 x i16> %src1) {
 define <8 x i16> @rhaddu_i_const_zero(<8 x i16> %src1) {
 ; CHECK-LABEL: rhaddu_i_const_zero:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    movi v1.2d, #0000000000000000
-; CHECK-NEXT:    urhadd v0.8h, v0.8h, v1.8h
+; CHECK-NEXT:    ushr v0.8h, v0.8h, #1
 ; CHECK-NEXT:    ret
   %result = call <8 x i16> @llvm.aarch64.neon.urhadd.v8i16(<8 x i16> <i16 0, i16 0, i16 0, i16 0, i16 0, i16 0, i16 0, i16 0>, <8 x i16> %src1)
   ret <8 x i16> %result
@@ -633,8 +632,7 @@ define <8 x i16> @rhadds_i_const_lhs(<8 x i16> %src1) {
 define <8 x i16> @rhadds_i_const_zero(<8 x i16> %src1) {
 ; CHECK-LABEL: rhadds_i_const_zero:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    movi v1.2d, #0000000000000000
-; CHECK-NEXT:    srhadd v0.8h, v0.8h, v1.8h
+; CHECK-NEXT:    sshr v0.8h, v0.8h, #1
 ; CHECK-NEXT:    ret
   %result = call <8 x i16> @llvm.aarch64.neon.srhadd.v8i16(<8 x i16> <i16 0, i16 0, i16 0, i16 0, i16 0, i16 0, i16 0, i16 0>, <8 x i16> %src1)
   ret <8 x i16> %result