[DAGCombiner] Fold avg(x, x) --> x

[junaire]

Summary

Signed-off-by: Jun Zhang <[e-mail]>

[RKSimon]

Can you add coverage to aarch64 so we catch all ISD::AVG node types?

[junaire]

Rebase

[junaire]

arc diff git merge-base HEAD origin --update D147404 to fix patch application failure.

[junaire]

In D147404#4239123, @RKSimon wrote:

Can you add coverage to aarch64 so we catch all ISD::AVG node types?

I've created a parent revision. Let me know if that's what you want, please :)

[junaire]

I screw it up with arc, sorry.

[dmgreen]

Thanks for the patch. I'm not sure about the X86 changes but I think if you add extra tests then all the AArch64 variants looks good.

[junaire]

Rebase

[RKSimon]

llvm/test/CodeGen/AArch64/avg.ll

2 ↗

(On Diff #510331)

Maybe just add coverage to https://github.com/llvm/llvm-project/blob/main/llvm/test/CodeGen/AArch64/hadd-combine.ll where the rest of the combines are already for ? AArch64 uses the term 'hadd' for halving-add averages AVGFLOORS/AVGFLOORU/AVGCEILS/AVGCEILU.

[pengfei]

llvm/test/CodeGen/X86/avg.ll

758–759

This and below are regressions.

[junaire]

llvm/test/CodeGen/X86/avg.ll

758–759

What makes AVX512F/BW so special? I have no clue about what happened.

[junaire]

Perform combine after the types are legalized.

[pengfei]

LGTM, thanks!

[junaire]

Update aarch64 tests

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

SDValue DAGCombiner::visitAVG(SDNode *N) {
   }
 
   // fold (avg x, undef) -> x
   if (N0.isUndef())
     return N1;
   if (N1.isUndef())
     return N0;
 
+  // Fold (avg x, x) --> x
+  if (N0 == N1 && Level >= AfterLegalizeTypes)
+    return N0;
+
   // TODO If we use avg for scalars anywhere, we can add (avgfl x, 0) -> x >> 1
 
   return SDValue();
 }
 
 SDValue DAGCombiner::visitABD(SDNode *N) {
   unsigned Opcode = N->getOpcode();
   SDValue N0 = N->getOperand(0);

llvm/test/CodeGen/AArch64/hadd-combine.ll

; CHECK-NEXT:    ret
   %result = call <8 x i16> @llvm.aarch64.neon.srhadd.v8i16(<8 x i16> undef, <8 x i16> %src1)
   ret <8 x i16> %result
 }
 
 
 define <8 x i8> @shadd_v8i8(<8 x i8> %x) {
 ; CHECK-LABEL: shadd_v8i8:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    shadd v0.8b, v0.8b, v0.8b
 ; CHECK-NEXT:    ret
   %r = tail call <8 x i8> @llvm.aarch64.neon.shadd.v8i8(<8 x i8> %x, <8 x i8> %x)
   ret <8 x i8> %r
 }
 
 define <4 x i16> @shadd_v4i16(<4 x i16> %x) {
 ; CHECK-LABEL: shadd_v4i16:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    shadd v0.4h, v0.4h, v0.4h
 ; CHECK-NEXT:    ret
   %r = tail call <4 x i16> @llvm.aarch64.neon.shadd.v4i16(<4 x i16> %x, <4 x i16> %x)
   ret <4 x i16> %r
 }
 
 define <2 x i32> @shadd_v2i32(<2 x i32> %x) {
 ; CHECK-LABEL: shadd_v2i32:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    shadd v0.2s, v0.2s, v0.2s
 ; CHECK-NEXT:    ret
   %r = tail call <2 x i32> @llvm.aarch64.neon.shadd.v2i32(<2 x i32> %x, <2 x i32> %x)
   ret <2 x i32> %r
 }
 
 define <16 x i8> @shadd_v16i8(<16 x i8> %x) {
 ; CHECK-LABEL: shadd_v16i8:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    shadd v0.16b, v0.16b, v0.16b
 ; CHECK-NEXT:    ret
   %r = tail call <16 x i8> @llvm.aarch64.neon.shadd.v16i8(<16 x i8> %x, <16 x i8> %x)
   ret <16 x i8> %r
 }
 
 define <8 x i16> @shadd_v8i16(<8 x i16> %x) {
 ; CHECK-LABEL: shadd_v8i16:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    shadd v0.8h, v0.8h, v0.8h
 ; CHECK-NEXT:    ret
   %r = tail call <8 x i16> @llvm.aarch64.neon.shadd.v8i16(<8 x i16> %x, <8 x i16> %x)
   ret <8 x i16> %r
 }
 
 define <4 x i32> @shadd_v4i32(<4 x i32> %x) {
 ; CHECK-LABEL: shadd_v4i32:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    shadd v0.4s, v0.4s, v0.4s
 ; CHECK-NEXT:    ret
   %r = tail call <4 x i32> @llvm.aarch64.neon.shadd.v4i32(<4 x i32> %x, <4 x i32> %x)
   ret <4 x i32> %r
 }
 
 define <8 x i8> @uhadd_v8i8(<8 x i8> %x) {
 ; CHECK-LABEL: uhadd_v8i8:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    uhadd v0.8b, v0.8b, v0.8b
 ; CHECK-NEXT:    ret
   %r = tail call <8 x i8> @llvm.aarch64.neon.uhadd.v8i8(<8 x i8> %x, <8 x i8> %x)
   ret <8 x i8> %r
 }
 
 define <4 x i16> @uhadd_v4i16(<4 x i16> %x) {
 ; CHECK-LABEL: uhadd_v4i16:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    uhadd v0.4h, v0.4h, v0.4h
 ; CHECK-NEXT:    ret
   %r = tail call <4 x i16> @llvm.aarch64.neon.uhadd.v4i16(<4 x i16> %x, <4 x i16> %x)
   ret <4 x i16> %r
 }
 
 define <2 x i32> @uhadd_v2i32(<2 x i32> %x) {
 ; CHECK-LABEL: uhadd_v2i32:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    uhadd v0.2s, v0.2s, v0.2s
 ; CHECK-NEXT:    ret
   %r = tail call <2 x i32> @llvm.aarch64.neon.uhadd.v2i32(<2 x i32> %x, <2 x i32> %x)
   ret <2 x i32> %r
 }
 
 define <16 x i8> @uhadd_v16i8(<16 x i8> %x) {
 ; CHECK-LABEL: uhadd_v16i8:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    uhadd v0.16b, v0.16b, v0.16b
 ; CHECK-NEXT:    ret
   %r = tail call <16 x i8> @llvm.aarch64.neon.uhadd.v16i8(<16 x i8> %x, <16 x i8> %x)
   ret <16 x i8> %r
 }
 
 define <8 x i16> @uhadd_v8i16(<8 x i16> %x) {
 ; CHECK-LABEL: uhadd_v8i16:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    uhadd v0.8h, v0.8h, v0.8h
 ; CHECK-NEXT:    ret
   %r = tail call <8 x i16> @llvm.aarch64.neon.uhadd.v8i16(<8 x i16> %x, <8 x i16> %x)
   ret <8 x i16> %r
 }
 
 define <4 x i32> @uhadd_v4i32(<4 x i32> %x) {
 ; CHECK-LABEL: uhadd_v4i32:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    uhadd v0.4s, v0.4s, v0.4s
 ; CHECK-NEXT:    ret
   %r = tail call <4 x i32> @llvm.aarch64.neon.uhadd.v4i32(<4 x i32> %x, <4 x i32> %x)
   ret <4 x i32> %r
 }
 define <8 x i8> @srhadd_v8i8(<8 x i8> %x) {
 ; CHECK-LABEL: srhadd_v8i8:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    srhadd v0.8b, v0.8b, v0.8b
 ; CHECK-NEXT:    ret
   %r = tail call <8 x i8> @llvm.aarch64.neon.srhadd.v8i8(<8 x i8> %x, <8 x i8> %x)
   ret <8 x i8> %r
 }
 
 define <4 x i16> @srhadd_v4i16(<4 x i16> %x) {
 ; CHECK-LABEL: srhadd_v4i16:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    srhadd v0.4h, v0.4h, v0.4h
 ; CHECK-NEXT:    ret
   %r = tail call <4 x i16> @llvm.aarch64.neon.srhadd.v4i16(<4 x i16> %x, <4 x i16> %x)
   ret <4 x i16> %r
 }
 
 define <2 x i32> @srhadd_v2i32(<2 x i32> %x) {
 ; CHECK-LABEL: srhadd_v2i32:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    srhadd v0.2s, v0.2s, v0.2s
 ; CHECK-NEXT:    ret
   %r = tail call <2 x i32> @llvm.aarch64.neon.srhadd.v2i32(<2 x i32> %x, <2 x i32> %x)
   ret <2 x i32> %r
 }
 
 define <16 x i8> @srhadd_v16i8(<16 x i8> %x) {
 ; CHECK-LABEL: srhadd_v16i8:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    srhadd v0.16b, v0.16b, v0.16b
 ; CHECK-NEXT:    ret
   %r = tail call <16 x i8> @llvm.aarch64.neon.srhadd.v16i8(<16 x i8> %x, <16 x i8> %x)
   ret <16 x i8> %r
 }
 
 define <8 x i16> @srhadd_v8i16(<8 x i16> %x) {
 ; CHECK-LABEL: srhadd_v8i16:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    srhadd v0.8h, v0.8h, v0.8h
 ; CHECK-NEXT:    ret
   %r = tail call <8 x i16> @llvm.aarch64.neon.srhadd.v8i16(<8 x i16> %x, <8 x i16> %x)
   ret <8 x i16> %r
 }
 
 define <4 x i32> @srhadd_v4i32(<4 x i32> %x) {
 ; CHECK-LABEL: srhadd_v4i32:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    srhadd v0.4s, v0.4s, v0.4s
 ; CHECK-NEXT:    ret
   %r = tail call <4 x i32> @llvm.aarch64.neon.srhadd.v4i32(<4 x i32> %x, <4 x i32> %x)
   ret <4 x i32> %r
 }
 
 define <8 x i8> @urhadd_v8i8(<8 x i8> %x) {
 ; CHECK-LABEL: urhadd_v8i8:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    urhadd v0.8b, v0.8b, v0.8b
 ; CHECK-NEXT:    ret
   %r = tail call <8 x i8> @llvm.aarch64.neon.urhadd.v8i8(<8 x i8> %x, <8 x i8> %x)
   ret <8 x i8> %r
 }
 
 define <4 x i16> @urhadd_v4i16(<4 x i16> %x) {
 ; CHECK-LABEL: urhadd_v4i16:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    urhadd v0.4h, v0.4h, v0.4h
 ; CHECK-NEXT:    ret
   %r = tail call <4 x i16> @llvm.aarch64.neon.urhadd.v4i16(<4 x i16> %x, <4 x i16> %x)
   ret <4 x i16> %r
 }
 
 define <2 x i32> @urhadd_v2i32(<2 x i32> %x) {
 ; CHECK-LABEL: urhadd_v2i32:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    urhadd v0.2s, v0.2s, v0.2s
 ; CHECK-NEXT:    ret
   %r = tail call <2 x i32> @llvm.aarch64.neon.urhadd.v2i32(<2 x i32> %x, <2 x i32> %x)
   ret <2 x i32> %r
 }
 
 define <16 x i8> @urhadd_v16i8(<16 x i8> %x) {
 ; CHECK-LABEL: urhadd_v16i8:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    urhadd v0.16b, v0.16b, v0.16b
 ; CHECK-NEXT:    ret
   %r = tail call <16 x i8> @llvm.aarch64.neon.urhadd.v16i8(<16 x i8> %x, <16 x i8> %x)
   ret <16 x i8> %r
 }
 
 define <8 x i16> @urhadd_v8i16(<8 x i16> %x) {
 ; CHECK-LABEL: urhadd_v8i16:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    urhadd v0.8h, v0.8h, v0.8h
 ; CHECK-NEXT:    ret
   %r = tail call <8 x i16> @llvm.aarch64.neon.urhadd.v8i16(<8 x i16> %x, <8 x i16> %x)
   ret <8 x i16> %r
 }
 
 define <4 x i32> @urhadd_v4i32(<4 x i32> %x) {
 ; CHECK-LABEL: urhadd_v4i32:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    urhadd v0.4s, v0.4s, v0.4s
 ; CHECK-NEXT:    ret
   %r = tail call <4 x i32> @llvm.aarch64.neon.urhadd.v4i32(<4 x i32> %x, <4 x i32> %x)
   ret <4 x i32> %r
 }
 
 declare <8 x i8> @llvm.aarch64.neon.shadd.v8i8(<8 x i8>, <8 x i8>)
 declare <4 x i16> @llvm.aarch64.neon.shadd.v4i16(<4 x i16>, <4 x i16>)
 declare <2 x i32> @llvm.aarch64.neon.shadd.v2i32(<2 x i32>, <2 x i32>)

llvm/test/CodeGen/X86/avg.ll

; AVX512-NEXT:    retq
   %8 = trunc <32 x i32> %7 to <32 x i8>
   store <32 x i8> %8, ptr undef, align 4
   ret void
 }
 
 define void @avg_v64i8_2(ptr %a, ptr %b) nounwind {
 ; SSE2-LABEL: avg_v64i8_2:
 ; SSE2:       # %bb.0:
-; SSE2-NEXT:    movdqa (%rsi), %xmm0
-; SSE2-NEXT:    movdqa 16(%rsi), %xmm1
-; SSE2-NEXT:    movdqa 32(%rsi), %xmm2
-; SSE2-NEXT:    movdqa 48(%rsi), %xmm3
-; SSE2-NEXT:    pavgb %xmm0, %xmm0
-; SSE2-NEXT:    pavgb %xmm1, %xmm1
-; SSE2-NEXT:    pavgb %xmm2, %xmm2
-; SSE2-NEXT:    pavgb %xmm3, %xmm3
-; SSE2-NEXT:    movdqu %xmm3, (%rax)
-; SSE2-NEXT:    movdqu %xmm2, (%rax)
-; SSE2-NEXT:    movdqu %xmm1, (%rax)
-; SSE2-NEXT:    movdqu %xmm0, (%rax)
+; SSE2-NEXT:    movaps (%rsi), %xmm0
+; SSE2-NEXT:    movaps 16(%rsi), %xmm1
+; SSE2-NEXT:    movaps 32(%rsi), %xmm2
+; SSE2-NEXT:    movaps 48(%rsi), %xmm3
+; SSE2-NEXT:    movups %xmm3, (%rax)
+; SSE2-NEXT:    movups %xmm2, (%rax)
+; SSE2-NEXT:    movups %xmm1, (%rax)
+; SSE2-NEXT:    movups %xmm0, (%rax)
 ; SSE2-NEXT:    retq
 ;
 ; AVX1-LABEL: avg_v64i8_2:
 ; AVX1:       # %bb.0:
-; AVX1-NEXT:    vmovdqa (%rsi), %xmm0
-; AVX1-NEXT:    vmovdqa 16(%rsi), %xmm1
-; AVX1-NEXT:    vmovdqa 32(%rsi), %xmm2
-; AVX1-NEXT:    vmovdqa 48(%rsi), %xmm3
-; AVX1-NEXT:    vpavgb %xmm0, %xmm0, %xmm0
-; AVX1-NEXT:    vpavgb %xmm1, %xmm1, %xmm1
-; AVX1-NEXT:    vpavgb %xmm2, %xmm2, %xmm2
-; AVX1-NEXT:    vpavgb %xmm3, %xmm3, %xmm3
-; AVX1-NEXT:    vmovdqu %xmm3, (%rax)
-; AVX1-NEXT:    vmovdqu %xmm2, (%rax)
-; AVX1-NEXT:    vmovdqu %xmm1, (%rax)
-; AVX1-NEXT:    vmovdqu %xmm0, (%rax)
+; AVX1-NEXT:    vmovaps (%rsi), %ymm0
+; AVX1-NEXT:    vmovaps 32(%rsi), %ymm1
+; AVX1-NEXT:    vmovups %ymm1, (%rax)
+; AVX1-NEXT:    vmovups %ymm0, (%rax)
+; AVX1-NEXT:    vzeroupper
 ; AVX1-NEXT:    retq
 ;
 ; AVX2-LABEL: avg_v64i8_2:
 ; AVX2:       # %bb.0:
-; AVX2-NEXT:    vmovdqa (%rsi), %ymm0
-; AVX2-NEXT:    vmovdqa 32(%rsi), %ymm1
-; AVX2-NEXT:    vpavgb %ymm0, %ymm0, %ymm0
-; AVX2-NEXT:    vpavgb %ymm1, %ymm1, %ymm1
-; AVX2-NEXT:    vmovdqu %ymm1, (%rax)
-; AVX2-NEXT:    vmovdqu %ymm0, (%rax)
+; AVX2-NEXT:    vmovaps (%rsi), %ymm0
+; AVX2-NEXT:    vmovaps 32(%rsi), %ymm1
+; AVX2-NEXT:    vmovups %ymm1, (%rax)
+; AVX2-NEXT:    vmovups %ymm0, (%rax)
 ; AVX2-NEXT:    vzeroupper
 ; AVX2-NEXT:    retq
 ;
-; AVX512F-LABEL: avg_v64i8_2:
-; AVX512F:       # %bb.0:
-; AVX512F-NEXT:    vmovdqa (%rsi), %ymm0
-; AVX512F-NEXT:    vmovdqa 32(%rsi), %ymm1
-; AVX512F-NEXT:    vpavgb %ymm0, %ymm0, %ymm0
-; AVX512F-NEXT:    vpavgb %ymm1, %ymm1, %ymm1
-; AVX512F-NEXT:    vmovdqu %ymm1, (%rax)
-; AVX512F-NEXT:    vmovdqu %ymm0, (%rax)
-; AVX512F-NEXT:    vzeroupper
-; AVX512F-NEXT:    retq
-;
-; AVX512BW-LABEL: avg_v64i8_2:
-; AVX512BW:       # %bb.0:
-; AVX512BW-NEXT:    vmovdqa64 (%rsi), %zmm0
-; AVX512BW-NEXT:    vpavgb %zmm0, %zmm0, %zmm0
-; AVX512BW-NEXT:    vmovdqu64 %zmm0, (%rax)
-; AVX512BW-NEXT:    vzeroupper
-; AVX512BW-NEXT:    retq
+; AVX512-LABEL: avg_v64i8_2:
+; AVX512:       # %bb.0:
+; AVX512-NEXT:    vmovaps (%rsi), %zmm0
+; AVX512-NEXT:    vmovups %zmm0, (%rax)
+; AVX512-NEXT:    vzeroupper
+; AVX512-NEXT:    retq
   %1 = load <64 x i8>, ptr %a
   %2 = load <64 x i8>, ptr %b
   %3 = zext <64 x i8> %1 to <64 x i32>
   %4 = zext <64 x i8> %2 to <64 x i32>
   %5 = add nuw nsw <64 x i32> %4, %4
   %6 = add nuw nsw <64 x i32> %5, <i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1>
   %7 = lshr <64 x i32> %6, <i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1>
   %8 = trunc <64 x i32> %7 to <64 x i8>

llvm/test/CodeGen/X86/combine-pavg.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc < %s -mtriple=x86_64-unknown-unknown -mattr=+sse4.1 | FileCheck %s --check-prefix=SSE
 ; RUN: llc < %s -mtriple=x86_64-unknown-unknown -mattr=+avx | FileCheck %s --check-prefixes=AVX
 ; RUN: llc < %s -mtriple=x86_64-unknown-unknown -mattr=+avx2 | FileCheck %s --check-prefixes=AVX
 
 declare <16 x i8> @llvm.x86.sse2.pavg.b(<16 x i8>, <16 x i8>) nounwind readnone
 declare <8 x i16> @llvm.x86.sse2.pavg.w(<8 x i16>, <8 x i16>) nounwind readnone
 
-; TODO: AVG(X,X) -> X
 define <16 x i8> @combine_pavgb_self(<16 x i8> %a0) {
 ; SSE-LABEL: combine_pavgb_self:
 ; SSE:       # %bb.0:
-; SSE-NEXT:    pavgb %xmm0, %xmm0
 ; SSE-NEXT:    retq
 ;
 ; AVX-LABEL: combine_pavgb_self:
 ; AVX:       # %bb.0:
-; AVX-NEXT:    vpavgb %xmm0, %xmm0, %xmm0
 ; AVX-NEXT:    retq
   %1 = call <16 x i8> @llvm.x86.sse2.pavg.b(<16 x i8> %a0, <16 x i8> %a0)
   ret <16 x i8> %1
 }
 
 define <16 x i8> @combine_pavgw_knownbits(<8 x i16> %a0, <8 x i16> %a1, <8 x i16> %a2, <8 x i16> %a3) {
 ; SSE-LABEL: combine_pavgw_knownbits:
 ; SSE:       # %bb.0: