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

via llvm-commits llvm-commits at lists.llvm.org
Tue Jun 11 18:08:26 PDT 2024


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/4] 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/4] 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/4] 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);

>From 4bf4ee451f0004c9e88b7f4a4cc14efb1f8c189f Mon Sep 17 00:00:00 2001
From: c8ef <c8ef at outlook.com>
Date: Wed, 12 Jun 2024 01:08:14 +0000
Subject: [PATCH 4/4] fix & add test

---
 .../AArch64/aarch64-known-bits-hadd.ll        |  6 +--
 llvm/test/CodeGen/AArch64/avg.ll              | 40 +++++++++++++++++++
 2 files changed, 42 insertions(+), 4 deletions(-)
 create mode 100644 llvm/test/CodeGen/AArch64/avg.ll

diff --git a/llvm/test/CodeGen/AArch64/aarch64-known-bits-hadd.ll b/llvm/test/CodeGen/AArch64/aarch64-known-bits-hadd.ll
index f36b8440fe4bf..b2cf089d8145f 100644
--- a/llvm/test/CodeGen/AArch64/aarch64-known-bits-hadd.ll
+++ b/llvm/test/CodeGen/AArch64/aarch64-known-bits-hadd.ll
@@ -9,9 +9,8 @@ declare <8 x i16> @llvm.aarch64.neon.srhadd.v8i16(<8 x i16>, <8 x i16>)
 define <8 x i16> @haddu_zext(<8 x i8> %a0, <8 x i8> %a1) {
 ; CHECK-LABEL: haddu_zext:
 ; CHECK:       // %bb.0:
+; CHECK-NEXT:    uhadd v0.8b, v0.8b, v1.8b
 ; CHECK-NEXT:    ushll v0.8h, v0.8b, #0
-; CHECK-NEXT:    ushll v1.8h, v1.8b, #0
-; CHECK-NEXT:    uhadd v0.8h, v0.8h, v1.8h
 ; CHECK-NEXT:    ret
   %x0 = zext <8 x i8> %a0 to <8 x i16>
   %x1 = zext <8 x i8> %a1 to <8 x i16>
@@ -23,9 +22,8 @@ define <8 x i16> @haddu_zext(<8 x i8> %a0, <8 x i8> %a1) {
 define <8 x i16> @rhaddu_zext(<8 x i8> %a0, <8 x i8> %a1) {
 ; CHECK-LABEL: rhaddu_zext:
 ; CHECK:       // %bb.0:
+; CHECK-NEXT:    urhadd v0.8b, v0.8b, v1.8b
 ; CHECK-NEXT:    ushll v0.8h, v0.8b, #0
-; CHECK-NEXT:    ushll v1.8h, v1.8b, #0
-; CHECK-NEXT:    urhadd v0.8h, v0.8h, v1.8h
 ; CHECK-NEXT:    ret
   %x0 = zext <8 x i8> %a0 to <8 x i16>
   %x1 = zext <8 x i8> %a1 to <8 x i16>
diff --git a/llvm/test/CodeGen/AArch64/avg.ll b/llvm/test/CodeGen/AArch64/avg.ll
new file mode 100644
index 0000000000000..5971e7d2ec7e3
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/avg.ll
@@ -0,0 +1,40 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 4
+; RUN: llc -mtriple=aarch64 < %s | FileCheck %s
+
+define <16 x i16> @zext_avgflooru(<16 x i8> %a0, <16 x i8> %a1) {
+; CHECK-LABEL: zext_avgflooru:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ext v2.16b, v0.16b, v0.16b, #8
+; CHECK-NEXT:    ext v3.16b, v1.16b, v1.16b, #8
+; CHECK-NEXT:    uhadd v0.8b, v0.8b, v1.8b
+; CHECK-NEXT:    ushll v0.8h, v0.8b, #0
+; CHECK-NEXT:    uhadd v1.8b, v2.8b, v3.8b
+; CHECK-NEXT:    ushll v1.8h, v1.8b, #0
+; CHECK-NEXT:    ret
+  %x0 = zext <16 x i8> %a0 to <16 x i16>
+  %x1 = zext <16 x i8> %a1 to <16 x i16>
+  %and = and <16 x i16> %x0, %x1
+  %xor = xor <16 x i16> %x0, %x1
+  %shift = lshr <16 x i16> %xor, <i16 1, i16 1, i16 1, i16 1, i16 1, i16 1, i16 1, i16 1, i16 1, i16 1, i16 1, i16 1, i16 1, i16 1, i16 1, i16 1>
+  %avg = add <16 x i16> %and, %shift
+  ret <16 x i16> %avg
+}
+
+define <16 x i16> @zext_avgceilu(<16 x i8> %a0, <16 x i8> %a1) {
+; CHECK-LABEL: zext_avgceilu:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ext v2.16b, v0.16b, v0.16b, #8
+; CHECK-NEXT:    ext v3.16b, v1.16b, v1.16b, #8
+; CHECK-NEXT:    urhadd v0.8b, v0.8b, v1.8b
+; CHECK-NEXT:    ushll v0.8h, v0.8b, #0
+; CHECK-NEXT:    urhadd v1.8b, v2.8b, v3.8b
+; CHECK-NEXT:    ushll v1.8h, v1.8b, #0
+; CHECK-NEXT:    ret
+  %x0 = zext <16 x i8> %a0 to <16 x i16>
+  %x1 = zext <16 x i8> %a1 to <16 x i16>
+  %or = or <16 x i16> %x0, %x1
+  %xor = xor <16 x i16> %x0, %x1
+  %shift = lshr <16 x i16> %xor, <i16 1, i16 1, i16 1, i16 1, i16 1, i16 1, i16 1, i16 1, i16 1, i16 1, i16 1, i16 1, i16 1, i16 1, i16 1, i16 1>
+  %avg = sub <16 x i16> %or, %shift
+  ret <16 x i16> %avg
+}



More information about the llvm-commits mailing list