[llvm] Matched some basic ISD::AVGFLOORU patterns (PR #84903)

[Shourya Goel via llvm-commits]

https://github.com/Sh0g0-1758 created https://github.com/llvm/llvm-project/pull/84903

Fixes: #84749

NOTE: The Tests are not written properly as of now, please suggest changes in them. 

>From cd98974c7e086153fe97ad93902b6bb5cd7d3367 Mon Sep 17 00:00:00 2001
From: Sh0g0-1758 <shouryagoel10000 at gmail.com>
Date: Tue, 12 Mar 2024 15:05:36 +0530
Subject: [PATCH] Matched some basic ISD::AVGFLOORU patterns

---
 llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp | 34 +++++++++++++++++++
 llvm/test/CodeGen/AArch64/dag-fixedwidth.ll   | 11 ++++++
 2 files changed, 45 insertions(+)
 create mode 100644 llvm/test/CodeGen/AArch64/dag-fixedwidth.ll

diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index 5476ef87971436..f6132e9627bf8e 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -2826,6 +2826,36 @@ SDValue DAGCombiner::visitADDLike(SDNode *N) {
   return SDValue();
 }
 
+// Attempt to form ext(avgflooru(A, B)) from add(and(A, B), lshr(xor(A, B), 1))
+static SDValue combineFixedwidthToAVG(SDNode *N, SelectionDAG &DAG) {
+  assert(N->getOpcode() == ISD::ADD && "ADD node is required here");
+  SDValue And = N->getOperand(0);
+  SDValue Lshr = N->getOperand(1);
+  if (And.getOpcode() != ISD::AND || Lshr.getOpcode() != ISD::SRL)
+    return SDValue();
+  SDValue Xor = Lshr.getOperand(0);
+  if (Xor.getOpcode() != ISD::XOR)
+    return SDValue();
+  SDValue And1 = And.getOperand(0);
+  SDValue And2 = And.getOperand(1);
+  SDValue Xor1 = Xor.getOperand(0);
+  SDValue Xor2 = Xor.getOperand(1);
+  if(Xor1 != And1 && Xor2 != And2)
+    return SDValue();
+  // Is the right shift using an immediate value of 1?
+  ConstantSDNode *N1C = isConstOrConstSplat(Lshr.getOperand(1));
+  if (!N1C || N1C->getAPIntValue() != 1)
+    return SDValue();
+  EVT VT = And.getValueType();
+  SDLoc DL(N);
+  const TargetLowering &TLI = DAG.getTargetLoweringInfo();
+  if (!TLI.isOperationLegalOrCustom(ISD::AVGFLOORU, VT))
+    return SDValue();
+  return DAG.getNode(ISD::AVGFLOORU, DL, VT,
+                     DAG.getExtOrTrunc(false, And1, DL, VT),
+                     DAG.getExtOrTrunc(false, And2, DL, VT));
+}
+
 SDValue DAGCombiner::visitADD(SDNode *N) {
   SDValue N0 = N->getOperand(0);
   SDValue N1 = N->getOperand(1);
@@ -2841,6 +2871,10 @@ SDValue DAGCombiner::visitADD(SDNode *N) {
   if (SDValue V = foldAddSubOfSignBit(N, DAG))
     return V;
 
+  // Try to match AVG fixedwidth pattern
+  if (SDValue V = combineFixedwidthToAVG(N, DAG))
+    return V;
+
   // fold (a+b) -> (a|b) iff a and b share no bits.
   if ((!LegalOperations || TLI.isOperationLegal(ISD::OR, VT)) &&
       DAG.haveNoCommonBitsSet(N0, N1))
diff --git a/llvm/test/CodeGen/AArch64/dag-fixedwidth.ll b/llvm/test/CodeGen/AArch64/dag-fixedwidth.ll
new file mode 100644
index 00000000000000..7000da89900ec3
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/dag-fixedwidth.ll
@@ -0,0 +1,11 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc < %s -mtriple=aarch64 | FileCheck %s
+
+define i4 @fixedwidth(i4 %a0, i4 %a1)  {
+; CHECK-LABEL: fixedwidth:
+  %and = and i4 %a0, %a1
+  %xor = xor i4 %a0, %a1
+  %srl = lshr i4 %xor, 1
+  %res = add i4 %and, %srl
+  ret i4 %res
+}
\ No newline at end of file