Skip to content

[AArch64] Combine getActiveLaneMask with vector_extract #81139

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
May 10, 2024
Merged

[AArch64] Combine getActiveLaneMask with vector_extract #81139

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
90ef28b
[AArch64] Combine getActiveLaneMask with vector_extract
momchil-velikov Dec 22, 2023
34d0edc
[fixup] Don't enable the transformation for SME and simplify some cod…
momchil-velikov Apr 18, 2024
450b580
[fixup] Check return type of vector_extract
momchil-velikov May 2, 2024
ced70f6
[fixup] Guard against extracting a fixed-length vector
momchil-velikov May 3, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
62 changes: 62 additions & 0 deletions llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20535,6 +20535,66 @@ static SDValue convertMergedOpToPredOp(SDNode *N, unsigned Opc,
return SDValue();
}

static SDValue tryCombineWhileLo(SDNode *N,
TargetLowering::DAGCombinerInfo &DCI,
const AArch64Subtarget *Subtarget) {
if (DCI.isBeforeLegalize())
return SDValue();

if (!Subtarget->hasSVE2p1())
return SDValue();

if (!N->hasNUsesOfValue(2, 0))
return SDValue();

const uint64_t HalfSize = N->getValueType(0).getVectorMinNumElements() / 2;
if (HalfSize < 2)
return SDValue();

auto It = N->use_begin();
SDNode *Lo = *It++;
SDNode *Hi = *It;

if (Lo->getOpcode() != ISD::EXTRACT_SUBVECTOR ||
Hi->getOpcode() != ISD::EXTRACT_SUBVECTOR)
return SDValue();

uint64_t OffLo = Lo->getConstantOperandVal(1);
uint64_t OffHi = Hi->getConstantOperandVal(1);

if (OffLo > OffHi) {
std::swap(Lo, Hi);
std::swap(OffLo, OffHi);
}

if (OffLo != 0 || OffHi != HalfSize)
return SDValue();

EVT HalfVec = Lo->getValueType(0);
if (HalfVec != Hi->getValueType(0) ||
HalfVec.getVectorElementCount() != ElementCount::getScalable(HalfSize))
return SDValue();

SelectionDAG &DAG = DCI.DAG;
SDLoc DL(N);
SDValue ID =
DAG.getTargetConstant(Intrinsic::aarch64_sve_whilelo_x2, DL, MVT::i64);
SDValue Idx = N->getOperand(1);
SDValue TC = N->getOperand(2);
if (Idx.getValueType() != MVT::i64) {
Idx = DAG.getZExtOrTrunc(Idx, DL, MVT::i64);
TC = DAG.getZExtOrTrunc(TC, DL, MVT::i64);
}
auto R =
DAG.getNode(ISD::INTRINSIC_WO_CHAIN, DL,
{Lo->getValueType(0), Hi->getValueType(0)}, {ID, Idx, TC});

DCI.CombineTo(Lo, R.getValue(0));
DCI.CombineTo(Hi, R.getValue(1));

return SDValue(N, 0);
}

static SDValue performIntrinsicCombine(SDNode *N,
TargetLowering::DAGCombinerInfo &DCI,
const AArch64Subtarget *Subtarget) {
Expand Down Expand Up @@ -20832,6 +20892,8 @@ static SDValue performIntrinsicCombine(SDNode *N,
case Intrinsic::aarch64_sve_ptest_last:
return getPTest(DAG, N->getValueType(0), N->getOperand(1), N->getOperand(2),
AArch64CC::LAST_ACTIVE);
case Intrinsic::aarch64_sve_whilelo:
return tryCombineWhileLo(N, DCI, Subtarget);
}
return SDValue();
}
Expand Down
156 changes: 156 additions & 0 deletions llvm/test/CodeGen/AArch64/get-active-lane-mask-extract.ll
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,156 @@
; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 4
; RUN: llc -mattr=+sve < %s | FileCheck %s -check-prefix CHECK-SVE
; RUN: llc -mattr=+sve2p1 < %s | FileCheck %s -check-prefix CHECK-SVE2p1
target triple = "aarch64-linux"

; Test combining of getActiveLaneMask with a pair of extract_vector operations.

define void @test_2x8bit_mask_with_32bit_index_and_trip_count(i32 %i, i32 %n) #0 {
; CHECK-SVE-LABEL: test_2x8bit_mask_with_32bit_index_and_trip_count:
; CHECK-SVE: // %bb.0:
; CHECK-SVE-NEXT: whilelo p1.b, w0, w1
; CHECK-SVE-NEXT: punpklo p0.h, p1.b
; CHECK-SVE-NEXT: punpkhi p1.h, p1.b
; CHECK-SVE-NEXT: b use
;
; CHECK-SVE2p1-LABEL: test_2x8bit_mask_with_32bit_index_and_trip_count:
; CHECK-SVE2p1: // %bb.0:
; CHECK-SVE2p1-NEXT: mov w8, w1
; CHECK-SVE2p1-NEXT: mov w9, w0
; CHECK-SVE2p1-NEXT: whilelo { p0.h, p1.h }, x9, x8
; CHECK-SVE2p1-NEXT: b use
%r = call <vscale x 16 x i1> @llvm.get.active.lane.mask.nxv16i1.i32(i32 %i, i32 %n)
%v0 = call <vscale x 8 x i1> @llvm.vector.extract.nxv8i1.nxv16i1.i64(<vscale x 16 x i1> %r, i64 0)
%v1 = call <vscale x 8 x i1> @llvm.vector.extract.nxv8i1.nxv16i1.i64(<vscale x 16 x i1> %r, i64 8)
tail call void @use(<vscale x 8 x i1> %v0, <vscale x 8 x i1> %v1)
ret void
}

define void @test_2x8bit_mask_with_64bit_index_and_trip_count(i64 %i, i64 %n) #0 {
; CHECK-SVE-LABEL: test_2x8bit_mask_with_64bit_index_and_trip_count:
; CHECK-SVE: // %bb.0:
; CHECK-SVE-NEXT: whilelo p1.b, x0, x1
; CHECK-SVE-NEXT: punpklo p0.h, p1.b
; CHECK-SVE-NEXT: punpkhi p1.h, p1.b
; CHECK-SVE-NEXT: b use
;
; CHECK-SVE2p1-LABEL: test_2x8bit_mask_with_64bit_index_and_trip_count:
; CHECK-SVE2p1: // %bb.0:
; CHECK-SVE2p1-NEXT: whilelo { p0.h, p1.h }, x0, x1
; CHECK-SVE2p1-NEXT: b use
%r = call <vscale x 16 x i1> @llvm.get.active.lane.mask.nxv16i1.i64(i64 %i, i64 %n)
%v0 = call <vscale x 8 x i1> @llvm.vector.extract.nxv8i1.nxv16i1.i64(<vscale x 16 x i1> %r, i64 0)
%v1 = call <vscale x 8 x i1> @llvm.vector.extract.nxv8i1.nxv16i1.i64(<vscale x 16 x i1> %r, i64 8)
tail call void @use(<vscale x 8 x i1> %v0, <vscale x 8 x i1> %v1)
ret void
}

define void @test_edge_case_2x1bit_mask(i64 %i, i64 %n) #0 {
; CHECK-SVE-LABEL: test_edge_case_2x1bit_mask:
; CHECK-SVE: // %bb.0:
; CHECK-SVE-NEXT: whilelo p1.d, x0, x1
; CHECK-SVE-NEXT: punpklo p0.h, p1.b
; CHECK-SVE-NEXT: punpkhi p1.h, p1.b
; CHECK-SVE-NEXT: b use
;
; CHECK-SVE2p1-LABEL: test_edge_case_2x1bit_mask:
; CHECK-SVE2p1: // %bb.0:
; CHECK-SVE2p1-NEXT: whilelo p1.d, x0, x1
; CHECK-SVE2p1-NEXT: punpklo p0.h, p1.b
; CHECK-SVE2p1-NEXT: punpkhi p1.h, p1.b
; CHECK-SVE2p1-NEXT: b use
%r = call <vscale x 2 x i1> @llvm.get.active.lane.mask.nxv2i1.i64(i64 %i, i64 %n)
%v0 = call <vscale x 1 x i1> @llvm.vector.extract.nxv1i1.nxv2i1.i64(<vscale x 2 x i1> %r, i64 0)
%v1 = call <vscale x 1 x i1> @llvm.vector.extract.nxv1i1.nxv2i1.i64(<vscale x 2 x i1> %r, i64 1)
tail call void @use(<vscale x 1 x i1> %v0, <vscale x 1 x i1> %v1)
ret void
}

define void @test_boring_case_2x2bit_mask(i64 %i, i64 %n) #0 {
; CHECK-SVE-LABEL: test_boring_case_2x2bit_mask:
; CHECK-SVE: // %bb.0:
; CHECK-SVE-NEXT: whilelo p1.s, x0, x1
; CHECK-SVE-NEXT: punpklo p0.h, p1.b
; CHECK-SVE-NEXT: punpkhi p1.h, p1.b
; CHECK-SVE-NEXT: b use
;
; CHECK-SVE2p1-LABEL: test_boring_case_2x2bit_mask:
; CHECK-SVE2p1: // %bb.0:
; CHECK-SVE2p1-NEXT: whilelo { p0.d, p1.d }, x0, x1
; CHECK-SVE2p1-NEXT: b use
%r = call <vscale x 4 x i1> @llvm.get.active.lane.mask.nxv4i1.i64(i64 %i, i64 %n)
%v0 = call <vscale x 2 x i1> @llvm.vector.extract.nxv2i1.nxv4i1.i64(<vscale x 4 x i1> %r, i64 0)
%v1 = call <vscale x 2 x i1> @llvm.vector.extract.nxv2i1.nxv4i1.i64(<vscale x 4 x i1> %r, i64 2)
tail call void @use(<vscale x 2 x i1> %v0, <vscale x 2 x i1> %v1)
ret void
}

; Negative test for when not extracting exactly two halves of the source vector
define void @test_partial_extract(i64 %i, i64 %n) #0 {
; CHECK-SVE-LABEL: test_partial_extract:
; CHECK-SVE: // %bb.0:
; CHECK-SVE-NEXT: whilelo p0.h, x0, x1
; CHECK-SVE-NEXT: punpklo p1.h, p0.b
; CHECK-SVE-NEXT: punpkhi p2.h, p0.b
; CHECK-SVE-NEXT: punpklo p0.h, p1.b
; CHECK-SVE-NEXT: punpklo p1.h, p2.b
; CHECK-SVE-NEXT: b use
;
; CHECK-SVE2p1-LABEL: test_partial_extract:
; CHECK-SVE2p1: // %bb.0:
; CHECK-SVE2p1-NEXT: whilelo p0.h, x0, x1
; CHECK-SVE2p1-NEXT: punpklo p1.h, p0.b
; CHECK-SVE2p1-NEXT: punpkhi p2.h, p0.b
; CHECK-SVE2p1-NEXT: punpklo p0.h, p1.b
; CHECK-SVE2p1-NEXT: punpklo p1.h, p2.b
; CHECK-SVE2p1-NEXT: b use
%r = call <vscale x 8 x i1> @llvm.get.active.lane.mask.nxv8i1.i64(i64 %i, i64 %n)
%v0 = call <vscale x 2 x i1> @llvm.vector.extract.nxv2i1.nxv8i1.i64(<vscale x 8 x i1> %r, i64 0)
%v1 = call <vscale x 2 x i1> @llvm.vector.extract.nxv2i1.nxv8i1.i64(<vscale x 8 x i1> %r, i64 4)
tail call void @use(<vscale x 2 x i1> %v0, <vscale x 2 x i1> %v1)
ret void
}

;; Negative test for when extracting a fixed-length vector.
define void @test_fixed_extract(i64 %i, i64 %n) #0 {
; CHECK-SVE-LABEL: test_fixed_extract:
; CHECK-SVE: // %bb.0:
; CHECK-SVE-NEXT: whilelo p0.h, x0, x1
; CHECK-SVE-NEXT: cset w8, mi
; CHECK-SVE-NEXT: mov z0.h, p0/z, #1 // =0x1
; CHECK-SVE-NEXT: umov w9, v0.h[4]
; CHECK-SVE-NEXT: umov w10, v0.h[1]
; CHECK-SVE-NEXT: umov w11, v0.h[5]
; CHECK-SVE-NEXT: fmov s0, w8
; CHECK-SVE-NEXT: fmov s1, w9
; CHECK-SVE-NEXT: mov v0.s[1], w10
; CHECK-SVE-NEXT: // kill: def $d0 killed $d0 killed $q0
; CHECK-SVE-NEXT: mov v1.s[1], w11
; CHECK-SVE-NEXT: // kill: def $d1 killed $d1 killed $q1
; CHECK-SVE-NEXT: b use
;
; CHECK-SVE2p1-LABEL: test_fixed_extract:
; CHECK-SVE2p1: // %bb.0:
; CHECK-SVE2p1-NEXT: whilelo p0.h, x0, x1
; CHECK-SVE2p1-NEXT: cset w8, mi
; CHECK-SVE2p1-NEXT: mov z0.h, p0/z, #1 // =0x1
; CHECK-SVE2p1-NEXT: umov w9, v0.h[4]
; CHECK-SVE2p1-NEXT: umov w10, v0.h[1]
; CHECK-SVE2p1-NEXT: umov w11, v0.h[5]
; CHECK-SVE2p1-NEXT: fmov s0, w8
; CHECK-SVE2p1-NEXT: fmov s1, w9
; CHECK-SVE2p1-NEXT: mov v0.s[1], w10
; CHECK-SVE2p1-NEXT: // kill: def $d0 killed $d0 killed $q0
; CHECK-SVE2p1-NEXT: mov v1.s[1], w11
; CHECK-SVE2p1-NEXT: // kill: def $d1 killed $d1 killed $q1
; CHECK-SVE2p1-NEXT: b use
%r = call <vscale x 8 x i1> @llvm.get.active.lane.mask.nxv8i1.i64(i64 %i, i64 %n)
%v0 = call <2 x i1> @llvm.vector.extract.v2i1.nxv8i1.i64(<vscale x 8 x i1> %r, i64 0)
%v1 = call <2 x i1> @llvm.vector.extract.v2i1.nxv8i1.i64(<vscale x 8 x i1> %r, i64 4)
tail call void @use(<2 x i1> %v0, <2 x i1> %v1)
ret void
}

declare void @use(...)

attributes #0 = { nounwind }