Skip to content

[AArch64] Lower alias mask to a whilewr #100769

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 19 commits into from
Aug 2, 2024
Merged

[AArch64] Lower alias mask to a whilewr #100769

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
19 commits
Select commit Hold shift + click to select a range
c943b04
[AArch64] Lower alias mask to a whilewr
SamTebbs33 Jun 26, 2024
54129dc
Add codegen test
SamTebbs33 Jul 29, 2024
a667dcd
format
SamTebbs33 Jul 29, 2024
9899a41
Use Log2_64
SamTebbs33 Jul 30, 2024
5f739bd
Fix comment formatting
SamTebbs33 Jul 30, 2024
e69f3be
Use SDValue instead of auto
SamTebbs33 Jul 30, 2024
653d6d2
Check for OR operands being the other way around
SamTebbs33 Jul 30, 2024
487ff3c
Replace dyn_cast and assert by cast
SamTebbs33 Jul 30, 2024
ab22dd1
Fix eltsize comparison and add test for it
SamTebbs33 Jul 30, 2024
0c378a3
Remove O3 from test
SamTebbs33 Jul 30, 2024
f114dc7
Remove dso_local
SamTebbs33 Jul 30, 2024
bbcb935
Move loop tests to whilewr.ll
SamTebbs33 Jul 30, 2024
6c700e7
Check the divide's first operand
SamTebbs33 Jul 31, 2024
207c3bc
Format
SamTebbs33 Jul 31, 2024
82da889
Pass subtarget from caller
SamTebbs33 Aug 1, 2024
4959eba
Remove dso_local from loop tests
SamTebbs33 Aug 1, 2024
2072294
Use isNullConstant
SamTebbs33 Aug 1, 2024
46602d8
Fix new lines
SamTebbs33 Aug 1, 2024
c30aa10
format
SamTebbs33 Aug 1, 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
83 changes: 83 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 @@ -94,6 +94,7 @@
#include <bitset>
#include <cassert>
#include <cctype>
#include <cmath>
#include <cstdint>
#include <cstdlib>
#include <iterator>
Expand Down Expand Up @@ -1523,6 +1524,7 @@ AArch64TargetLowering::AArch64TargetLowering(const TargetMachine &TM,
setOperationAction(ISD::VECREDUCE_AND, VT, Custom);
setOperationAction(ISD::VECREDUCE_OR, VT, Custom);
setOperationAction(ISD::VECREDUCE_XOR, VT, Custom);
setOperationAction(ISD::OR, VT, Custom);

setOperationAction(ISD::SELECT_CC, VT, Expand);
setOperationAction(ISD::INSERT_VECTOR_ELT, VT, Custom);
Expand Down Expand Up @@ -13782,8 +13784,89 @@ static SDValue tryLowerToSLI(SDNode *N, SelectionDAG &DAG) {
return ResultSLI;
}

/// Try to lower the construction of a pointer alias mask to a WHILEWR.
/// The mask's enabled lanes represent the elements that will not overlap across
/// one loop iteration. This tries to match: or (splat (setcc_lt (sub ptrA,
/// ptrB), -(element_size - 1))),
/// (get_active_lane_mask 0, (div (sub ptrA, ptrB), element_size))
SDValue tryWhileWRFromOR(SDValue Op, SelectionDAG &DAG) {
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pass Subtarget from the caller.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done.

if (!DAG.getSubtarget<AArch64Subtarget>().hasSVE2())
return SDValue();
auto LaneMask = Op.getOperand(0);
auto Splat = Op.getOperand(1);

if (LaneMask.getOpcode() != ISD::INTRINSIC_WO_CHAIN ||
LaneMask.getConstantOperandVal(0) != Intrinsic::get_active_lane_mask ||
Splat.getOpcode() != ISD::SPLAT_VECTOR)
return SDValue();

auto Cmp = Splat.getOperand(0);
if (Cmp.getOpcode() != ISD::SETCC)
return SDValue();

CondCodeSDNode *Cond = dyn_cast<CondCodeSDNode>(Cmp.getOperand(2));
assert(Cond && "SETCC doesn't have a condition code");

auto ComparatorConst = dyn_cast<ConstantSDNode>(Cmp.getOperand(1));
if (!ComparatorConst || ComparatorConst->getSExtValue() > 0 ||
Cond->get() != ISD::CondCode::SETLT)
return SDValue();
unsigned CompValue = std::abs(ComparatorConst->getSExtValue());
unsigned EltSize = CompValue + 1;
if (!isPowerOf2_64(EltSize) || EltSize > 64)
return SDValue();

auto Diff = Cmp.getOperand(0);
if (Diff.getOpcode() != ISD::SUB || Diff.getValueType() != MVT::i64)
return SDValue();

auto LaneMaskConst = dyn_cast<ConstantSDNode>(LaneMask.getOperand(1));
if (!LaneMaskConst || LaneMaskConst->getZExtValue() != 0 ||
(EltSize != 1 && LaneMask.getOperand(2).getOpcode() != ISD::SRA))
return SDValue();

// An alias mask for i8 elements omits the division because it would just
// divide by 1
if (EltSize > 1) {
auto DiffDiv = LaneMask.getOperand(2);
auto DiffDivConst = dyn_cast<ConstantSDNode>(DiffDiv.getOperand(1));
if (!DiffDivConst || DiffDivConst->getZExtValue() != std::log2(EltSize))
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

std::log -> Log2_64 if it is integer.
Should something check DiffDiv.getOperand(0) == Diff?

Copy link
Collaborator Author

@SamTebbs33 SamTebbs33 Jul 31, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There is indeed a check of the divide operand missing. It's a bit more involved since Diff can be negative, so some extra selects and such are inserted to turn it positive before dividing. The 16 bit element case is a little different to the others and so needs some different matching. I hope I've explained it well enough in the comments in the commit I just pushed to address your suggestion.

return SDValue();
} else if (LaneMask.getOperand(2) != Diff)
return SDValue();

auto StorePtr = Diff.getOperand(0);
auto ReadPtr = Diff.getOperand(1);

unsigned IntrinsicID = 0;
switch (EltSize) {
case 1:
IntrinsicID = Intrinsic::aarch64_sve_whilewr_b;
break;
case 2:
IntrinsicID = Intrinsic::aarch64_sve_whilewr_h;
break;
case 4:
IntrinsicID = Intrinsic::aarch64_sve_whilewr_s;
break;
case 8:
IntrinsicID = Intrinsic::aarch64_sve_whilewr_d;
break;
default:
return SDValue();
}
SDLoc DL(Op);
SDValue ID = DAG.getConstant(IntrinsicID, DL, MVT::i32);
auto N = DAG.getNode(ISD::INTRINSIC_WO_CHAIN, DL, Op.getValueType(), ID,
StorePtr, ReadPtr);
return N;
}

SDValue AArch64TargetLowering::LowerVectorOR(SDValue Op,
SelectionDAG &DAG) const {

if (SDValue SV = tryWhileWRFromOR(Op, DAG))
return SV;
if (useSVEForFixedLengthVectorVT(Op.getValueType(),
!Subtarget->isNeonAvailable()))
return LowerToScalableOp(Op, DAG);
Expand Down
127 changes: 127 additions & 0 deletions llvm/test/CodeGen/AArch64/whilewr.ll
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,127 @@
; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 4
; RUN: llc %s -mtriple=aarch64-linux-gnu -O3 -mattr=+sve2 -o - | FileCheck %s
; RUN: llc %s -mtriple=aarch64-linux-gnu -O3 -mattr=+sve -o - | FileCheck %s --check-prefix=CHECK-NOSVE2
define dso_local <vscale x 16 x i1> @whilewr_8(ptr noalias %a, ptr %b, ptr %c, i32 %n) {
; CHECK-LABEL: whilewr_8:
; CHECK: // %bb.0: // %entry
; CHECK-NEXT: whilewr p0.b, x1, x2
; CHECK-NEXT: ret
;
; CHECK-NOSVE2-LABEL: whilewr_8:
; CHECK-NOSVE2: // %bb.0: // %entry
; CHECK-NOSVE2-NEXT: sub x8, x1, x2
; CHECK-NOSVE2-NEXT: cmp x8, #0
; CHECK-NOSVE2-NEXT: cset w9, lt
; CHECK-NOSVE2-NEXT: whilelo p0.b, xzr, x8
; CHECK-NOSVE2-NEXT: sbfx x8, x9, #0, #1
; CHECK-NOSVE2-NEXT: whilelo p1.b, xzr, x8
; CHECK-NOSVE2-NEXT: sel p0.b, p0, p0.b, p1.b
; CHECK-NOSVE2-NEXT: ret
entry:
%c14 = ptrtoint ptr %c to i64
%b15 = ptrtoint ptr %b to i64
%sub.diff = sub i64 %b15, %c14
%neg.compare = icmp slt i64 %sub.diff, 0
%.splatinsert = insertelement <vscale x 16 x i1> poison, i1 %neg.compare, i64 0
%.splat = shufflevector <vscale x 16 x i1> %.splatinsert, <vscale x 16 x i1> poison, <vscale x 16 x i32> zeroinitializer
%ptr.diff.lane.mask = tail call <vscale x 16 x i1> @llvm.get.active.lane.mask.nxv16i1.i64(i64 0, i64 %sub.diff)
%active.lane.mask.alias = or <vscale x 16 x i1> %ptr.diff.lane.mask, %.splat
ret <vscale x 16 x i1> %active.lane.mask.alias
}

define dso_local <vscale x 8 x i1> @whilewr_16(ptr noalias %a, ptr %b, ptr %c, i32 %n) {
; CHECK-LABEL: whilewr_16:
; CHECK: // %bb.0: // %entry
; CHECK-NEXT: whilewr p0.h, x1, x2
; CHECK-NEXT: ret
;
; CHECK-NOSVE2-LABEL: whilewr_16:
; CHECK-NOSVE2: // %bb.0: // %entry
; CHECK-NOSVE2-NEXT: sub x8, x1, x2
; CHECK-NOSVE2-NEXT: cmn x8, #1
; CHECK-NOSVE2-NEXT: add x8, x8, x8, lsr #63
; CHECK-NOSVE2-NEXT: cset w9, lt
; CHECK-NOSVE2-NEXT: sbfx x9, x9, #0, #1
; CHECK-NOSVE2-NEXT: asr x8, x8, #1
; CHECK-NOSVE2-NEXT: whilelo p0.h, xzr, x9
; CHECK-NOSVE2-NEXT: whilelo p1.h, xzr, x8
; CHECK-NOSVE2-NEXT: mov p0.b, p1/m, p1.b
; CHECK-NOSVE2-NEXT: ret
entry:
%b14 = ptrtoint ptr %b to i64
%c15 = ptrtoint ptr %c to i64
%sub.diff = sub i64 %b14, %c15
%diff = sdiv i64 %sub.diff, 2
%neg.compare = icmp slt i64 %sub.diff, -1
%.splatinsert = insertelement <vscale x 8 x i1> poison, i1 %neg.compare, i64 0
%.splat = shufflevector <vscale x 8 x i1> %.splatinsert, <vscale x 8 x i1> poison, <vscale x 8 x i32> zeroinitializer
%ptr.diff.lane.mask = tail call <vscale x 8 x i1> @llvm.get.active.lane.mask.nxv8i1.i64(i64 0, i64 %diff)
%active.lane.mask.alias = or <vscale x 8 x i1> %ptr.diff.lane.mask, %.splat
ret <vscale x 8 x i1> %active.lane.mask.alias
}

define dso_local <vscale x 4 x i1> @whilewr_32(ptr noalias %a, ptr %b, ptr %c, i32 %n) {
; CHECK-LABEL: whilewr_32:
; CHECK: // %bb.0: // %entry
; CHECK-NEXT: whilewr p0.s, x1, x2
; CHECK-NEXT: ret
;
; CHECK-NOSVE2-LABEL: whilewr_32:
; CHECK-NOSVE2: // %bb.0: // %entry
; CHECK-NOSVE2-NEXT: sub x8, x1, x2
; CHECK-NOSVE2-NEXT: add x9, x8, #3
; CHECK-NOSVE2-NEXT: cmp x8, #0
; CHECK-NOSVE2-NEXT: csel x9, x9, x8, lt
; CHECK-NOSVE2-NEXT: cmn x8, #3
; CHECK-NOSVE2-NEXT: cset w8, lt
; CHECK-NOSVE2-NEXT: asr x9, x9, #2
; CHECK-NOSVE2-NEXT: sbfx x8, x8, #0, #1
; CHECK-NOSVE2-NEXT: whilelo p1.s, xzr, x9
; CHECK-NOSVE2-NEXT: whilelo p0.s, xzr, x8
; CHECK-NOSVE2-NEXT: mov p0.b, p1/m, p1.b
; CHECK-NOSVE2-NEXT: ret
entry:
%b12 = ptrtoint ptr %b to i64
%c13 = ptrtoint ptr %c to i64
%sub.diff = sub i64 %b12, %c13
%diff = sdiv i64 %sub.diff, 4
%neg.compare = icmp slt i64 %sub.diff, -3
%.splatinsert = insertelement <vscale x 4 x i1> poison, i1 %neg.compare, i64 0
%.splat = shufflevector <vscale x 4 x i1> %.splatinsert, <vscale x 4 x i1> poison, <vscale x 4 x i32> zeroinitializer
%ptr.diff.lane.mask = tail call <vscale x 4 x i1> @llvm.get.active.lane.mask.nxv4i1.i64(i64 0, i64 %diff)
%active.lane.mask.alias = or <vscale x 4 x i1> %ptr.diff.lane.mask, %.splat
ret <vscale x 4 x i1> %active.lane.mask.alias
}

define dso_local <vscale x 2 x i1> @whilewr_64(ptr noalias %a, ptr %b, ptr %c, i32 %n) {
; CHECK-LABEL: whilewr_64:
; CHECK: // %bb.0: // %entry
; CHECK-NEXT: whilewr p0.d, x1, x2
; CHECK-NEXT: ret
;
; CHECK-NOSVE2-LABEL: whilewr_64:
; CHECK-NOSVE2: // %bb.0: // %entry
; CHECK-NOSVE2-NEXT: sub x8, x1, x2
; CHECK-NOSVE2-NEXT: add x9, x8, #7
; CHECK-NOSVE2-NEXT: cmp x8, #0
; CHECK-NOSVE2-NEXT: csel x9, x9, x8, lt
; CHECK-NOSVE2-NEXT: cmn x8, #7
; CHECK-NOSVE2-NEXT: cset w8, lt
; CHECK-NOSVE2-NEXT: asr x9, x9, #3
; CHECK-NOSVE2-NEXT: sbfx x8, x8, #0, #1
; CHECK-NOSVE2-NEXT: whilelo p1.d, xzr, x9
; CHECK-NOSVE2-NEXT: whilelo p0.d, xzr, x8
; CHECK-NOSVE2-NEXT: mov p0.b, p1/m, p1.b
; CHECK-NOSVE2-NEXT: ret
entry:
%b12 = ptrtoint ptr %b to i64
%c13 = ptrtoint ptr %c to i64
%sub.diff = sub i64 %b12, %c13
%diff = sdiv i64 %sub.diff, 8
%neg.compare = icmp slt i64 %sub.diff, -7
%.splatinsert = insertelement <vscale x 2 x i1> poison, i1 %neg.compare, i64 0
%.splat = shufflevector <vscale x 2 x i1> %.splatinsert, <vscale x 2 x i1> poison, <vscale x 2 x i32> zeroinitializer
%ptr.diff.lane.mask = tail call <vscale x 2 x i1> @llvm.get.active.lane.mask.nxv2i1.i64(i64 0, i64 %diff)
%active.lane.mask.alias = or <vscale x 2 x i1> %ptr.diff.lane.mask, %.splat
ret <vscale x 2 x i1> %active.lane.mask.alias
}
Loading
Loading