[DAG] Fold trunc(avg(x,y)) for avgceil/floor u/s nodes if they have sufficient leading zero/sign bits

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

@@ -16279,6 +16279,40 @@ SDValue DAGCombiner::visitTRUNCATE(SDNode *N) {
   // because targets may prefer a wider type during later combines and invert
   // this transform.
   switch (N0.getOpcode()) {
+  case ISD::AVGCEILU:
+  case ISD::AVGFLOORU:
+    if (!LegalOperations && N0.hasOneUse() &&
+        TLI.isOperationLegal(N0.getOpcode(), VT)) {
+      SDValue X = N0.getOperand(0);
+      SDValue Y = N0.getOperand(1);
+      unsigned SrcBits = X.getScalarValueSizeInBits();
+      unsigned DstBits = VT.getScalarSizeInBits();
+      APInt UpperBits = APInt::getBitsSetFrom(SrcBits, DstBits);
+      if (DAG.MaskedValueIsZero(X, UpperBits) &&
+          DAG.MaskedValueIsZero(Y, UpperBits)) {
+        SDValue Tx = DAG.getNode(ISD::TRUNCATE, DL, VT, X);
+        SDValue Ty = DAG.getNode(ISD::TRUNCATE, DL, VT, Y);
+        return DAG.getNode(N0.getOpcode(), DL, VT, Tx, Ty);
+      }
+    }
+    break;
+  case ISD::AVGCEILS:
+  case ISD::AVGFLOORS:
+    if (!LegalOperations && N0.hasOneUse() &&
+        TLI.isOperationLegal(N0.getOpcode(), VT)) {
+      SDValue X = N0.getOperand(0);
+      SDValue Y = N0.getOperand(1);
+      unsigned SrcBits = X.getScalarValueSizeInBits();
+      unsigned DstBits = VT.getScalarSizeInBits();
+      unsigned NeededSignBits = SrcBits - DstBits + 1;
+      if (DAG.ComputeNumSignBits(X) >= NeededSignBits &&
+          DAG.ComputeNumSignBits(Y) >= NeededSignBits) {
+        SDValue Tx = DAG.getNode(ISD::TRUNCATE, DL, VT, X);
+        SDValue Ty = DAG.getNode(ISD::TRUNCATE, DL, VT, Y);
+        return DAG.getNode(N0.getOpcode(), DL, VT, Tx, Ty);
+      }
+    }
+    break;
   case ISD::ADD:
   case ISD::SUB:
   case ISD::MUL:

llvm/test/CodeGen/AArch64/trunc-avg-fold.ll

@@ -0,0 +1,53 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
+; RUN: llc -mtriple=aarch64-- -O2 -mattr=+neon < %s | FileCheck %s
+
+define <8 x i8> @avgceil_u_i8_to_i16(<8 x i8> %a, <8 x i8> %b) {
+; CHECK-LABEL: avgceil_u_i8_to_i16:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    urhadd v0.8b, v0.8b, v1.8b
+; CHECK-NEXT:    ret
+  %a16 = zext <8 x i8> %a to <8 x i16>
+  %b16 = zext <8 x i8> %b to <8 x i16>
+  %avg16 = call <8 x i16> @llvm.aarch64.neon.urhadd.v8i16(<8 x i16> %a16, <8 x i16> %b16)
+  %r = trunc <8 x i16> %avg16 to <8 x i8>
+  ret <8 x i8> %r
+}
+
+
+define <8 x i8> @test_avgceil_s(<8 x i8> %a, <8 x i8> %b) {
+; CHECK-LABEL: test_avgceil_s:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    srhadd v0.8b, v0.8b, v1.8b
+; CHECK-NEXT:    ret
+  %a16 = sext <8 x i8> %a to <8 x i16>
+  %b16 = sext <8 x i8> %b to <8 x i16>
+  %avg16 = call <8 x i16> @llvm.aarch64.neon.srhadd.v8i16(<8 x i16> %a16, <8 x i16> %b16)
+  %res  = trunc <8 x i16> %avg16 to <8 x i8>
+  ret <8 x i8> %res
+}
+
+define <8 x i8> @avgfloor_u_i8_to_i16(<8 x i8> %a, <8 x i8> %b) {
+; CHECK-LABEL: avgfloor_u_i8_to_i16:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    uhadd v0.8b, v0.8b, v1.8b
+; CHECK-NEXT:    ret
+  %a16 = zext  <8 x i8>  %a to <8 x i16>
+  %b16 = zext  <8 x i8>  %b to <8 x i16>
+  %avg16 = call <8 x i16> @llvm.aarch64.neon.uhadd.v8i16(<8 x i16> %a16, <8 x i16> %b16)
+  %res = trunc <8 x i16> %avg16 to <8 x i8>
+  ret <8 x i8> %res
+}
+
+define <8 x i8> @test_avgfloor_s(<8 x i8> %a, <8 x i8> %b) {
+; CHECK-LABEL: test_avgfloor_s:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    shadd v0.8b, v0.8b, v1.8b
+; CHECK-NEXT:    ret
+  %a16 = sext  <8 x i8>  %a to <8 x i16>
+  %b16 = sext  <8 x i8>  %b to <8 x i16>
+  %avg16 = call <8 x i16> @llvm.aarch64.neon.shadd.v8i16(<8 x i16> %a16, <8 x i16> %b16)
+  %res  = trunc <8 x i16> %avg16 to <8 x i8>
+  ret <8 x i8> %res
+}
+
+