[DAG] fold `avgs(sext(x), sext(y))` -> `sext(avgs(x, y))` #95365

c8ef · 2024-06-13T07:32:02Z

Follow up of #95134.

Context: #95134 (comment).

llvmbot · 2024-06-13T07:32:34Z

@llvm/pr-subscribers-backend-aarch64

@llvm/pr-subscribers-llvm-selectiondag

Author: None (c8ef)

Changes

Follow up of #95134.

Context: #95134 (comment).

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

2 Files Affected:

(modified) llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp (+15)
(modified) llvm/test/CodeGen/AArch64/avg.ll (+84)

diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index 78970bc4fe4ab..0d4df4a7ecda5 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -5237,6 +5237,7 @@ SDValue DAGCombiner::visitAVG(SDNode *N) {
                        DAG.getShiftAmountConstant(1, VT, DL));
 
   // fold avgu(zext(x), zext(y)) -> zext(avgu(x, y))
+  // fold avgu(sext(x), sext(y)) -> sext(avgu(x, y))
   if (sd_match(
           N, m_BinOp(ISD::AVGFLOORU, m_ZExt(m_Value(X)), m_ZExt(m_Value(Y)))) &&
       X.getValueType() == Y.getValueType() &&
@@ -5251,6 +5252,20 @@ SDValue DAGCombiner::visitAVG(SDNode *N) {
     SDValue AvgCeilU = DAG.getNode(ISD::AVGCEILU, DL, X.getValueType(), X, Y);
     return DAG.getNode(ISD::ZERO_EXTEND, DL, VT, AvgCeilU);
   }
+  if (sd_match(
+          N, m_BinOp(ISD::AVGFLOORU, m_SExt(m_Value(X)), m_SExt(m_Value(Y)))) &&
+      X.getValueType() == Y.getValueType() &&
+      hasOperation(ISD::AVGFLOORU, X.getValueType())) {
+    SDValue AvgFloorU = DAG.getNode(ISD::AVGFLOORU, DL, X.getValueType(), X, Y);
+    return DAG.getNode(ISD::SIGN_EXTEND, DL, VT, AvgFloorU);
+  }
+  if (sd_match(
+          N, m_BinOp(ISD::AVGCEILU, m_SExt(m_Value(X)), m_SExt(m_Value(Y)))) &&
+      X.getValueType() == Y.getValueType() &&
+      hasOperation(ISD::AVGCEILU, X.getValueType())) {
+    SDValue AvgCeilU = DAG.getNode(ISD::AVGCEILU, DL, X.getValueType(), X, Y);
+    return DAG.getNode(ISD::SIGN_EXTEND, DL, VT, AvgCeilU);
+  }
 
   // Fold avgflooru(x,y) -> avgceilu(x,y-1) iff y != 0
   // Fold avgflooru(x,y) -> avgceilu(x-1,y) iff x != 0
diff --git a/llvm/test/CodeGen/AArch64/avg.ll b/llvm/test/CodeGen/AArch64/avg.ll
index dc87708555987..e61b47772b7d7 100644
--- a/llvm/test/CodeGen/AArch64/avg.ll
+++ b/llvm/test/CodeGen/AArch64/avg.ll
@@ -68,3 +68,87 @@ define <16 x i16> @zext_avgceilu_mismatch(<16 x i4> %a0, <16 x i8> %a1) {
   %avg = sub <16 x i16> %or, %shift
   ret <16 x i16> %avg
 }
+
+define <16 x i16> @sext_avgflooru(<16 x i8> %a0, <16 x i8> %a1) {
+; CHECK-LABEL: sext_avgflooru:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    sshll v2.8h, v0.8b, #0
+; CHECK-NEXT:    sshll2 v0.8h, v0.16b, #0
+; CHECK-NEXT:    sshll v3.8h, v1.8b, #0
+; CHECK-NEXT:    sshll2 v1.8h, v1.16b, #0
+; CHECK-NEXT:    shadd v1.8h, v0.8h, v1.8h
+; CHECK-NEXT:    shadd v0.8h, v2.8h, v3.8h
+; CHECK-NEXT:    ret
+  %x0 = sext <16 x i8> %a0 to <16 x i16>
+  %x1 = sext <16 x i8> %a1 to <16 x i16>
+  %and = and <16 x i16> %x0, %x1
+  %xor = xor <16 x i16> %x0, %x1
+  %shift = ashr <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> @sext_avgflooru_mismatch(<16 x i8> %a0, <16 x i4> %a1) {
+; CHECK-LABEL: sext_avgflooru_mismatch:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ushll2 v2.8h, v1.16b, #0
+; CHECK-NEXT:    ushll v1.8h, v1.8b, #0
+; CHECK-NEXT:    sshll v3.8h, v0.8b, #0
+; CHECK-NEXT:    sshll2 v0.8h, v0.16b, #0
+; CHECK-NEXT:    shl v1.8h, v1.8h, #12
+; CHECK-NEXT:    shl v2.8h, v2.8h, #12
+; CHECK-NEXT:    sshr v4.8h, v1.8h, #12
+; CHECK-NEXT:    sshr v1.8h, v2.8h, #12
+; CHECK-NEXT:    shadd v1.8h, v0.8h, v1.8h
+; CHECK-NEXT:    shadd v0.8h, v3.8h, v4.8h
+; CHECK-NEXT:    ret
+  %x0 = sext <16 x i8> %a0 to <16 x i16>
+  %x1 = sext <16 x i4> %a1 to <16 x i16>
+  %and = and <16 x i16> %x0, %x1
+  %xor = xor <16 x i16> %x0, %x1
+  %shift = ashr <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> @sext_avgceilu(<16 x i8> %a0, <16 x i8> %a1) {
+; CHECK-LABEL: sext_avgceilu:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    sshll v2.8h, v0.8b, #0
+; CHECK-NEXT:    sshll2 v0.8h, v0.16b, #0
+; CHECK-NEXT:    sshll v3.8h, v1.8b, #0
+; CHECK-NEXT:    sshll2 v1.8h, v1.16b, #0
+; CHECK-NEXT:    srhadd v1.8h, v0.8h, v1.8h
+; CHECK-NEXT:    srhadd v0.8h, v2.8h, v3.8h
+; CHECK-NEXT:    ret
+  %x0 = sext <16 x i8> %a0 to <16 x i16>
+  %x1 = sext <16 x i8> %a1 to <16 x i16>
+  %or = or <16 x i16> %x0, %x1
+  %xor = xor <16 x i16> %x0, %x1
+  %shift = ashr <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
+}
+
+define <16 x i16> @sext_avgceilu_mismatch(<16 x i4> %a0, <16 x i8> %a1) {
+; CHECK-LABEL: sext_avgceilu_mismatch:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ushll v2.8h, v0.8b, #0
+; CHECK-NEXT:    ushll2 v0.8h, v0.16b, #0
+; CHECK-NEXT:    sshll v3.8h, v1.8b, #0
+; CHECK-NEXT:    sshll2 v1.8h, v1.16b, #0
+; CHECK-NEXT:    shl v2.8h, v2.8h, #12
+; CHECK-NEXT:    shl v0.8h, v0.8h, #12
+; CHECK-NEXT:    sshr v2.8h, v2.8h, #12
+; CHECK-NEXT:    sshr v0.8h, v0.8h, #12
+; CHECK-NEXT:    srhadd v1.8h, v0.8h, v1.8h
+; CHECK-NEXT:    srhadd v0.8h, v2.8h, v3.8h
+; CHECK-NEXT:    ret
+  %x0 = sext <16 x i4> %a0 to <16 x i16>
+  %x1 = sext <16 x i8> %a1 to <16 x i16>
+  %or = or <16 x i16> %x0, %x1
+  %xor = xor <16 x i16> %x0, %x1
+  %shift = ashr <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
+}

jayfoad · 2024-06-13T07:51:54Z

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

No, this needs to be: avgs(sext(x), sext(y)) -> sext(avgs(x, y))

c8ef · 2024-06-13T08:03:37Z

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

No, this needs to be: avgs(sext(x), sext(y)) -> sext(avgs(x, y))

Sorry, my bad.

dtcxzyw

LGTM. Thank you!

llvm/test/CodeGen/AArch64/arm64-vhadd.ll

Context: #95365 (comment) The current implementation of `m_SExt` matches both `ISD::SIGN_EXTEND` and `ISD::SIGN_EXTEND_INREG`. However, in cases where we specifically need to match _only_ `ISD::SIGN_EXTEND`, such as in the SelectionDAG graph below, this can lead to issues and unintended combinations. ``` SelectionDAG has 13 nodes: t0: ch,glue = EntryToken t2: v2i32,ch = CopyFromReg t0, Register:v2i32 %0 t21: v2i32 = sign_extend_inreg t2, ValueType:ch:v2i8 t4: v2i32,ch = CopyFromReg t0, Register:v2i32 %1 t22: v2i32 = sign_extend_inreg t4, ValueType:ch:v2i8 t23: v2i32 = avgfloors t21, t22 t24: v2i32 = sign_extend_inreg t23, ValueType:ch:v2i8 t15: ch,glue = CopyToReg t0, Register:v2i32 $d0, t24 t16: ch = AArch64ISD::RET_GLUE t15, Register:v2i32 $d0, t15:1 ```

RKSimon

Please can you rebase + regenerate the tests now that #95415 has landed

c8ef · 2024-06-14T13:25:52Z

Please can you rebase + regenerate the tests now that #95415 has landed

Done.

RKSimon

LGTM - cheers

fold avgu(sext(x), sext(y)) -> sext(avgu(x, y))

df7448f

llvmbot added backend:AArch64 llvm:SelectionDAG SelectionDAGISel as well labels Jun 13, 2024

c8ef changed the title ~~[DAG] fold avgu(sext(x), sext(y)) -> sext(avgu(x, y))~~ [DAG] fold avgs(sext(x), sext(y)) -> sext(avgs(x, y)) Jun 13, 2024

handle avgs, not avgu

fbe39f4

dtcxzyw approved these changes Jun 13, 2024

View reviewed changes

dtcxzyw requested review from RKSimon and jayfoad June 13, 2024 08:32

regen test

2cc9ea1

RKSimon reviewed Jun 13, 2024

View reviewed changes

llvm/test/CodeGen/AArch64/arm64-vhadd.ll Outdated Show resolved Hide resolved

c8ef mentioned this pull request Jun 13, 2024

[SDPatternMatch] Only match ISD::SIGN_EXTEND in m_SExt #95415

Merged

RKSimon requested changes Jun 14, 2024

View reviewed changes

c8ef added 2 commits June 14, 2024 17:55

Merge branch 'main' into fold-avg

5f549d7

regen test

a7db1cf

c8ef requested a review from RKSimon June 14, 2024 13:25

RKSimon approved these changes Jun 14, 2024

View reviewed changes

dtcxzyw merged commit 08fae46 into llvm:main Jun 14, 2024
5 of 6 checks passed

c8ef deleted the fold-avg branch June 14, 2024 14:32

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[DAG] fold `avgs(sext(x), sext(y))` -> `sext(avgs(x, y))` #95365

[DAG] fold `avgs(sext(x), sext(y))` -> `sext(avgs(x, y))` #95365

Uh oh!

c8ef commented Jun 13, 2024

Uh oh!

llvmbot commented Jun 13, 2024 •

edited

Loading

Uh oh!

jayfoad commented Jun 13, 2024

Uh oh!

c8ef commented Jun 13, 2024

Uh oh!

dtcxzyw left a comment

Uh oh!

Uh oh!

RKSimon left a comment

Uh oh!

c8ef commented Jun 14, 2024

Uh oh!

RKSimon left a comment

Uh oh!

Uh oh!

Uh oh!

[DAG] fold avgs(sext(x), sext(y)) -> sext(avgs(x, y)) #95365

[DAG] fold avgs(sext(x), sext(y)) -> sext(avgs(x, y)) #95365

Uh oh!

Conversation

c8ef commented Jun 13, 2024

Uh oh!

llvmbot commented Jun 13, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jayfoad commented Jun 13, 2024

Uh oh!

c8ef commented Jun 13, 2024

Uh oh!

dtcxzyw left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

RKSimon left a comment

Choose a reason for hiding this comment

Uh oh!

c8ef commented Jun 14, 2024

Uh oh!

RKSimon left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

[DAG] fold `avgs(sext(x), sext(y))` -> `sext(avgs(x, y))` #95365

[DAG] fold `avgs(sext(x), sext(y))` -> `sext(avgs(x, y))` #95365

llvmbot commented Jun 13, 2024 •

edited

Loading