[DAG] fold avgu(zext(x), zext(y)) -> zext(avgu(x, y))

[c8ef]

close: #86301

[llvmbot]

@llvm/pr-subscribers-backend-aarch64

@llvm/pr-subscribers-llvm-selectiondag

Author: None (c8ef)

Changes

close: #86301

---

Full diff: https://github.com/llvm/llvm-project/pull/95134.diff

1 Files Affected:

• (modified) llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp (+15)

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();
 }
 

[RKSimon]

For tests you can start with these: https://simd.godbolt.org/z/xfrWo8Tc7

[dtcxzyw]

LGTM. Thank you!

[RKSimon]

LGTM - cheers

[jayfoad]

Any plans to do avgs(sext(x), sext(y)) -> sext(avgs(x, y)) ?

[RKSimon]

Any plans to do avgs(sext(x), sext(y)) -> sext(avgs(x, y)) ?

IIRC the constraints on that pattern aren't as easy as the unsigned cases.

[jayfoad]

Any plans to do avgs(sext(x), sext(y)) -> sext(avgs(x, y)) ?

IIRC the constraints on that pattern aren't as easy as the unsigned cases.

Really? My instinct is that it should Just Work, but I have not thought about it deeply.

[RKSimon]

Any plans to do avgs(sext(x), sext(y)) -> sext(avgs(x, y)) ?

IIRC the constraints on that pattern aren't as easy as the unsigned cases.

Really? My instinct is that it should Just Work, but I have not thought about it deeply.

I'm miss-remembering - it looks like it just needs an explicit freeze in the alive test where the unsigned cases didn't https://alive2.llvm.org/ce/z/qgp7bF

[c8ef]

I would like to continue working on the signed case. Do we need to explicitly use SelectionDAG::getFreeze for the signed extended source value to match the alive test?

[dtcxzyw]

Any plans to do avgs(sext(x), sext(y)) -> sext(avgs(x, y)) ?

IIRC the constraints on that pattern aren't as easy as the unsigned cases.

Really? My instinct is that it should Just Work, but I have not thought about it deeply.

I'm miss-remembering - it looks like it just needs an explicit freeze in the alive test where the unsigned cases didn't https://alive2.llvm.org/ce/z/qgp7bF

I don't think we need freeze here because this fold doesn't introduce multiple uses as xor + or does. BTW, we already gave up on correct undef semantics in #90097.