Skip to content

Commit a312799

Browse files
woruyuworuyu
committed
[DAG] Fold trunc(abdu(x,y)) and trunc(abds(x,y)) if they have sufficient leading zero/sign bits
1 parent 7250b66 commit a312799

File tree

3 files changed

+123
-29
lines changed

3 files changed

+123
-29
lines changed

‎llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -607,6 +607,7 @@ namespace {
607607
SDValue foldLogicOfSetCCs(bool IsAnd, SDValue N0, SDValue N1,
608608
const SDLoc &DL);
609609
SDValue foldSubToUSubSat(EVT DstVT, SDNode *N, const SDLoc &DL);
610+
SDValue foldAbdToNarrowType(EVT VT, SDNode *N, const SDLoc &DL);
610611
SDValue foldABSToABD(SDNode *N, const SDLoc &DL);
611612
SDValue foldSelectToABD(SDValue LHS, SDValue RHS, SDValue True,
612613
SDValue False, ISD::CondCode CC, const SDLoc &DL);
@@ -3925,6 +3926,38 @@ SDValue DAGCombiner::foldSubToUSubSat(EVT DstVT, SDNode *N, const SDLoc &DL) {
39253926
return SDValue();
39263927
}
39273928

3929+
// trunc (ABDU/S A, B)) → ABDU/S (trunc A), (trunc B)
3930+
SDValue DAGCombiner::foldAbdToNarrowType(EVT VT, SDNode *N, const SDLoc &DL) {
3931+
SDValue Op = N->getOperand(0);
3932+
unsigned Opcode = Op.getOpcode();
3933+
if (Opcode != ISD::ABDU && Opcode != ISD::ABDS)
3934+
return SDValue();
3935+
3936+
EVT SrcVT = Op.getValueType();
3937+
EVT TruncVT = N->getValueType(0);
3938+
unsigned NumSrcBits = SrcVT.getScalarSizeInBits();
3939+
unsigned NumTruncBits = TruncVT.getScalarSizeInBits();
3940+
unsigned NeededBits = NumSrcBits - NumTruncBits;
3941+
3942+
bool CanFold = false;
3943+
3944+
if (Opcode == ISD::ABDU) {
3945+
KnownBits Known = DAG.computeKnownBits(Op);
3946+
CanFold = Known.countMinLeadingZeros() >= NeededBits;
3947+
} else {
3948+
unsigned SignBits = DAG.ComputeNumSignBits(Op);
3949+
CanFold = SignBits >= NeededBits;
3950+
}
3951+
3952+
if (CanFold) {
3953+
SDValue NewOp0 = DAG.getNode(ISD::TRUNCATE, DL, TruncVT, Op.getOperand(0));
3954+
SDValue NewOp1 = DAG.getNode(ISD::TRUNCATE, DL, TruncVT, Op.getOperand(1));
3955+
return DAG.getNode(Opcode, DL, TruncVT, NewOp0, NewOp1);
3956+
}
3957+
3958+
return SDValue();
3959+
}
3960+
39283961
// Refinement of DAG/Type Legalisation (promotion) when CTLZ is used for
39293962
// counting leading ones. Broadly, it replaces the substraction with a left
39303963
// shift.
@@ -16275,6 +16308,10 @@ SDValue DAGCombiner::visitTRUNCATE(SDNode *N) {
1627516308
if (SDValue NewVSel = matchVSelectOpSizesWithSetCC(N))
1627616309
return NewVSel;
1627716310

16311+
// fold trunc (ABDU/S A, B)) → ABDU/S (trunc A), (trunc B)
16312+
if (SDValue V = foldAbdToNarrowType(VT, N, SDLoc(N)))
16313+
return V;
16314+
1627816315
// Narrow a suitable binary operation with a non-opaque constant operand by
1627916316
// moving it ahead of the truncate. This is limited to pre-legalization
1628016317
// because targets may prefer a wider type during later combines and invert

‎llvm/test/CodeGen/AArch64/abd-combine.ll

Lines changed: 20 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -17,12 +17,9 @@ define <8 x i16> @abdu_base(<8 x i16> %src1, <8 x i16> %src2) {
1717
define <8 x i16> @abdu_const(<8 x i16> %src1) {
1818
; CHECK-LABEL: abdu_const:
1919
; CHECK: // %bb.0:
20-
; CHECK-NEXT: movi v1.4s, #1
21-
; CHECK-NEXT: ushll v2.4s, v0.4h, #0
22-
; CHECK-NEXT: ushll2 v0.4s, v0.8h, #0
23-
; CHECK-NEXT: uabd v0.4s, v0.4s, v1.4s
24-
; CHECK-NEXT: uabd v1.4s, v2.4s, v1.4s
25-
; CHECK-NEXT: uzp1 v0.8h, v1.8h, v0.8h
20+
; CHECK-NEXT: movi v1.4h, #1
21+
; CHECK-NEXT: mov v1.d[1], v1.d[0]
22+
; CHECK-NEXT: sabd v0.8h, v0.8h, v1.8h
2623
; CHECK-NEXT: ret
2724
%zextsrc1 = zext <8 x i16> %src1 to <8 x i32>
2825
%sub = sub <8 x i32> %zextsrc1, <i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1>
@@ -34,12 +31,9 @@ define <8 x i16> @abdu_const(<8 x i16> %src1) {
3431
define <8 x i16> @abdu_const_lhs(<8 x i16> %src1) {
3532
; CHECK-LABEL: abdu_const_lhs:
3633
; CHECK: // %bb.0:
37-
; CHECK-NEXT: movi v1.4s, #1
38-
; CHECK-NEXT: ushll v2.4s, v0.4h, #0
39-
; CHECK-NEXT: ushll2 v0.4s, v0.8h, #0
40-
; CHECK-NEXT: uabd v0.4s, v0.4s, v1.4s
41-
; CHECK-NEXT: uabd v1.4s, v2.4s, v1.4s
42-
; CHECK-NEXT: uzp1 v0.8h, v1.8h, v0.8h
34+
; CHECK-NEXT: movi v1.4h, #1
35+
; CHECK-NEXT: mov v1.d[1], v1.d[0]
36+
; CHECK-NEXT: sabd v0.8h, v0.8h, v1.8h
4337
; CHECK-NEXT: ret
4438
%zextsrc1 = zext <8 x i16> %src1 to <8 x i32>
4539
%sub = sub <8 x i32> <i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1>, %zextsrc1
@@ -51,6 +45,10 @@ define <8 x i16> @abdu_const_lhs(<8 x i16> %src1) {
5145
define <8 x i16> @abdu_const_zero(<8 x i16> %src1) {
5246
; CHECK-LABEL: abdu_const_zero:
5347
; CHECK: // %bb.0:
48+
; CHECK-NEXT: ext v1.16b, v0.16b, v0.16b, #8
49+
; CHECK-NEXT: abs v0.4h, v0.4h
50+
; CHECK-NEXT: abs v1.4h, v1.4h
51+
; CHECK-NEXT: mov v0.d[1], v1.d[0]
5452
; CHECK-NEXT: ret
5553
%zextsrc1 = zext <8 x i16> %src1 to <8 x i32>
5654
%sub = sub <8 x i32> <i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0>, %zextsrc1
@@ -318,12 +316,9 @@ define <8 x i16> @abds_base(<8 x i16> %src1, <8 x i16> %src2) {
318316
define <8 x i16> @abds_const(<8 x i16> %src1) {
319317
; CHECK-LABEL: abds_const:
320318
; CHECK: // %bb.0:
321-
; CHECK-NEXT: movi v1.4s, #1
322-
; CHECK-NEXT: sshll v2.4s, v0.4h, #0
323-
; CHECK-NEXT: sshll2 v0.4s, v0.8h, #0
324-
; CHECK-NEXT: sabd v0.4s, v0.4s, v1.4s
325-
; CHECK-NEXT: sabd v1.4s, v2.4s, v1.4s
326-
; CHECK-NEXT: uzp1 v0.8h, v1.8h, v0.8h
319+
; CHECK-NEXT: movi v1.4h, #1
320+
; CHECK-NEXT: mov v1.d[1], v1.d[0]
321+
; CHECK-NEXT: sabd v0.8h, v0.8h, v1.8h
327322
; CHECK-NEXT: ret
328323
%zextsrc1 = sext <8 x i16> %src1 to <8 x i32>
329324
%sub = sub <8 x i32> %zextsrc1, <i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1>
@@ -335,12 +330,9 @@ define <8 x i16> @abds_const(<8 x i16> %src1) {
335330
define <8 x i16> @abds_const_lhs(<8 x i16> %src1) {
336331
; CHECK-LABEL: abds_const_lhs:
337332
; CHECK: // %bb.0:
338-
; CHECK-NEXT: movi v1.4s, #1
339-
; CHECK-NEXT: sshll v2.4s, v0.4h, #0
340-
; CHECK-NEXT: sshll2 v0.4s, v0.8h, #0
341-
; CHECK-NEXT: sabd v0.4s, v0.4s, v1.4s
342-
; CHECK-NEXT: sabd v1.4s, v2.4s, v1.4s
343-
; CHECK-NEXT: uzp1 v0.8h, v1.8h, v0.8h
333+
; CHECK-NEXT: movi v1.4h, #1
334+
; CHECK-NEXT: mov v1.d[1], v1.d[0]
335+
; CHECK-NEXT: sabd v0.8h, v0.8h, v1.8h
344336
; CHECK-NEXT: ret
345337
%zextsrc1 = sext <8 x i16> %src1 to <8 x i32>
346338
%sub = sub <8 x i32> <i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1>, %zextsrc1
@@ -352,11 +344,10 @@ define <8 x i16> @abds_const_lhs(<8 x i16> %src1) {
352344
define <8 x i16> @abds_const_zero(<8 x i16> %src1) {
353345
; CHECK-LABEL: abds_const_zero:
354346
; CHECK: // %bb.0:
355-
; CHECK-NEXT: sshll v1.4s, v0.4h, #0
356-
; CHECK-NEXT: sshll2 v0.4s, v0.8h, #0
357-
; CHECK-NEXT: abs v0.4s, v0.4s
358-
; CHECK-NEXT: abs v1.4s, v1.4s
359-
; CHECK-NEXT: uzp1 v0.8h, v1.8h, v0.8h
347+
; CHECK-NEXT: ext v1.16b, v0.16b, v0.16b, #8
348+
; CHECK-NEXT: abs v0.4h, v0.4h
349+
; CHECK-NEXT: abs v1.4h, v1.4h
350+
; CHECK-NEXT: mov v0.d[1], v1.d[0]
360351
; CHECK-NEXT: ret
361352
%zextsrc1 = sext <8 x i16> %src1 to <8 x i32>
362353
%sub = sub <8 x i32> <i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0>, %zextsrc1

‎llvm/test/CodeGen/AArch64/arm64-neon-aba-abd.ll

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -575,3 +575,69 @@ define <4 x i32> @knownbits_sabd_and_mul_mask(<4 x i32> %a0, <4 x i32> %a1) {
575575
%6 = shufflevector <4 x i32> %5, <4 x i32> undef, <4 x i32> <i32 0, i32 0, i32 3, i32 3>
576576
ret <4 x i32> %6
577577
}
578+
579+
define <4 x i16> @trunc_abdu_foldable(<4 x i16> %a, <4 x i16> %b) {
580+
; CHECK-SD-LABEL: trunc_abdu_foldable:
581+
; CHECK-SD: // %bb.0:
582+
; CHECK-SD-NEXT: uabd v0.4h, v0.4h, v1.4h
583+
; CHECK-SD-NEXT: ret
584+
;
585+
; CHECK-GI-LABEL: trunc_abdu_foldable:
586+
; CHECK-GI: // %bb.0:
587+
; CHECK-GI-NEXT: ushll v0.4s, v0.4h, #0
588+
; CHECK-GI-NEXT: ushll v1.4s, v1.4h, #0
589+
; CHECK-GI-NEXT: uabd v0.4s, v0.4s, v1.4s
590+
; CHECK-GI-NEXT: xtn v0.4h, v0.4s
591+
; CHECK-GI-NEXT: ret
592+
%ext_a = zext <4 x i16> %a to <4 x i32>
593+
%ext_b = zext <4 x i16> %b to <4 x i32>
594+
%abd = call <4 x i32> @llvm.aarch64.neon.uabd.v4i32(<4 x i32> %ext_a, <4 x i32> %ext_b)
595+
%trunc = trunc <4 x i32> %abd to <4 x i16>
596+
ret <4 x i16> %trunc
597+
}
598+
599+
define <4 x i16> @trunc_abds_foldable(<4 x i16> %a, <4 x i16> %b) {
600+
; CHECK-SD-LABEL: trunc_abds_foldable:
601+
; CHECK-SD: // %bb.0:
602+
; CHECK-SD-NEXT: sabd v0.4h, v0.4h, v1.4h
603+
; CHECK-SD-NEXT: ret
604+
;
605+
; CHECK-GI-LABEL: trunc_abds_foldable:
606+
; CHECK-GI: // %bb.0:
607+
; CHECK-GI-NEXT: sshll v0.4s, v0.4h, #0
608+
; CHECK-GI-NEXT: sshll v1.4s, v1.4h, #0
609+
; CHECK-GI-NEXT: sabd v0.4s, v0.4s, v1.4s
610+
; CHECK-GI-NEXT: xtn v0.4h, v0.4s
611+
; CHECK-GI-NEXT: ret
612+
%a32 = sext <4 x i16> %a to <4 x i32>
613+
%b32 = sext <4 x i16> %b to <4 x i32>
614+
%abd32 = call <4 x i32> @llvm.aarch64.neon.sabd.v4i32(<4 x i32> %a32, <4 x i32> %b32)
615+
%res16 = trunc <4 x i32> %abd32 to <4 x i16>
616+
ret <4 x i16> %res16
617+
}
618+
619+
define <4 x i16> @trunc_abdu_not_foldable(<4 x i16> %a, <4 x i32> %b) {
620+
; CHECK-LABEL: trunc_abdu_not_foldable:
621+
; CHECK: // %bb.0:
622+
; CHECK-NEXT: ushll v0.4s, v0.4h, #0
623+
; CHECK-NEXT: uabd v0.4s, v0.4s, v1.4s
624+
; CHECK-NEXT: xtn v0.4h, v0.4s
625+
; CHECK-NEXT: ret
626+
%ext_a = zext <4 x i16> %a to <4 x i32>
627+
%abd = call <4 x i32> @llvm.aarch64.neon.uabd.v4i32(<4 x i32> %ext_a, <4 x i32> %b)
628+
%trunc = trunc <4 x i32> %abd to <4 x i16>
629+
ret <4 x i16> %trunc
630+
}
631+
632+
define <4 x i16> @truncate_abds_testcase1(<4 x i16> %a, <4 x i32> %b) {
633+
; CHECK-LABEL: truncate_abds_testcase1:
634+
; CHECK: // %bb.0:
635+
; CHECK-NEXT: sshll v0.4s, v0.4h, #0
636+
; CHECK-NEXT: sabd v0.4s, v0.4s, v1.4s
637+
; CHECK-NEXT: xtn v0.4h, v0.4s
638+
; CHECK-NEXT: ret
639+
%a32 = sext <4 x i16> %a to <4 x i32>
640+
%abd32 = call <4 x i32> @llvm.aarch64.neon.sabd.v4i32(<4 x i32> %a32, <4 x i32> %b)
641+
%res16 = trunc <4 x i32> %abd32 to <4 x i16>
642+
ret <4 x i16> %res16
643+
}

0 commit comments

Comments
 (0)