Skip to content
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.

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

[AIE2P] Combine G_SHUFFLE_VECTOR into Unmerge+Concat #357

Open
wants to merge 1 commit into
base: aie-public
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
66 changes: 46 additions & 20 deletions llvm/lib/Target/AIE/AIECombinerHelper.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -2004,6 +2004,20 @@ buildExtractSubvector(MachineIRBuilder &B, MachineRegisterInfo &MRI,
return B.buildInstr(Opc, {DstVecReg}, {NewSrcReg, Cst});
}

static void buildUnmergeVector(MachineIRBuilder &B, MachineRegisterInfo &MRI,
Register DstReg, Register SrcReg,
unsigned NumSubVectors, unsigned SubIdx) {
const LLT DstTy = MRI.getType(DstReg);
SmallVector<Register, 4> SubVecs;
for (unsigned I = 0; I < NumSubVectors; I++) {
if (I == (unsigned)SubIdx)
SubVecs.push_back(DstReg);
else
SubVecs.push_back(MRI.createGenericVirtualRegister(DstTy));
}
B.buildUnmerge(SubVecs, SrcReg);
}

/// Match something like this:
/// %1:_(<16 x s32>) = COPY $x0
/// %2:_(<16 x s32>) = COPY $x1
Expand All @@ -2020,18 +2034,10 @@ static bool matchShuffleToUnmerge(MachineInstr &MI, MachineRegisterInfo &MRI,
unsigned NumSubVectors) {
const Register DstReg = MI.getOperand(0).getReg();
const Register Src1Reg = MI.getOperand(1).getReg();
const LLT DstTy = MRI.getType(DstReg);

// TODO: Select into G_EXTRACT_SUBVECTOR once it is more widely supported
MatchInfo = [=, &MRI](MachineIRBuilder &B) {
SmallVector<Register, 4> SubVecs;
for (unsigned I = 0; I < NumSubVectors; I++) {
if (I == (unsigned)SubIdx)
SubVecs.push_back(DstReg);
else
SubVecs.push_back(MRI.createGenericVirtualRegister(DstTy));
}
B.buildUnmerge(SubVecs, Src1Reg);
buildUnmergeVector(B, MRI, DstReg, Src1Reg, NumSubVectors, SubIdx);
};
return true;
}
Expand Down Expand Up @@ -2247,20 +2253,40 @@ static bool matchShuffleToSubvecBroadcast(MachineInstr &MI,
const LLT ElemTy = Src1Ty.getElementType();
const LLT DstSubvecType =
LLT::fixed_vector(SplatMaskLen, ElemTy.getSizeInBits());
const unsigned SubIdx = SplatMaskStart / SplatMaskLen;
Register ExtractSubvecDstReg =
MRI.createGenericVirtualRegister(DstSubvecType);

// Check whether we can extract the subvector
if (!checkExtractSubvectorPrerequisites(TII, DstSubvecType, Src1Ty))
return false;
const bool CanExtractSubvector =
checkExtractSubvectorPrerequisites(TII, DstSubvecType, Src1Ty);
if (CanExtractSubvector) {
MatchInfo = [=, &MRI, &TII](MachineIRBuilder &B) {
auto Extract = buildExtractSubvector(B, MRI, TII, ExtractSubvecDstReg,
Src1Reg, SubIdx);
buildBroadcastVector(B, MRI, Extract.getReg(0), DstReg);
};
return true;
}

MatchInfo = [=, &MRI, &TII](MachineIRBuilder &B) {
Register ExtractSubvecDstReg =
MRI.createGenericVirtualRegister(DstSubvecType);
auto Extract =
buildExtractSubvector(B, MRI, TII, ExtractSubvecDstReg, Src1Reg,
SplatMaskStart / SplatMaskLen);
buildBroadcastVector(B, MRI, Extract.getReg(0), DstReg);
};
return true;
// If we cannot extract the subvector, we try to apply UNMERGE + CONCAT
const unsigned NumSubVectors = NumSrcElems / SplatMaskLen;
// Don't try to unmerge when we have just one subvector.
// We can overcome with a copy, but other combiners can do a
// better job for this case.
if (NumSubVectors > 1 && NumDstElems == SplatMaskLen * 2) {
MatchInfo = [=, &MRI](MachineIRBuilder &B) {
buildUnmergeVector(B, MRI, ExtractSubvecDstReg, Src1Reg, NumSubVectors,
SubIdx);

const SmallVector<Register, 2> ConcatOps = {ExtractSubvecDstReg,
ExtractSubvecDstReg};
B.buildConcatVectors({DstReg}, ConcatOps);
};
return true;
}

return false;
}

/// Match something like this:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -990,3 +990,42 @@ body: |
%0:_(<8 x s32>) = G_SHUFFLE_VECTOR %1(<8 x s32>), %2(<8 x s32>), shufflemask(0, 1, 2, 3, 4, 5, 6, 2)
PseudoRET implicit $lr, implicit %0
...

# Test G_SHUFFLE_VECTOR to UNMERGE+CONCAT
---
name: shuffle_vector_to_unmerge_concat_lo
tracksRegLiveness: true
body: |
bb.1:
liveins: $x0, $x1
; CHECK-LABEL: name: shuffle_vector_to_unmerge_concat_lo
; CHECK: liveins: $x0, $x1
; CHECK-NEXT: {{ $}}
; CHECK-NEXT: [[COPY:%[0-9]+]]:_(<16 x s32>) = COPY $x0
; CHECK-NEXT: [[AIE_UNPAD_VECTOR:%[0-9]+]]:_(<8 x s32>) = G_AIE_UNPAD_VECTOR [[COPY]](<16 x s32>)
; CHECK-NEXT: [[CONCAT_VECTORS:%[0-9]+]]:_(<16 x s32>) = G_CONCAT_VECTORS [[AIE_UNPAD_VECTOR]](<8 x s32>), [[AIE_UNPAD_VECTOR]](<8 x s32>)
; CHECK-NEXT: PseudoRET implicit $lr, implicit [[CONCAT_VECTORS]](<16 x s32>)
%1:_(<16 x s32>) = COPY $x0
%2:_(<16 x s32>) = COPY $x1
%0:_(<16 x s32>) = G_SHUFFLE_VECTOR %1(<16 x s32>), %2(<16 x s32>), shufflemask(0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7)
PseudoRET implicit $lr, implicit %0
...

---
name: shuffle_vector_to_unmerge_concat_hi
tracksRegLiveness: true
body: |
bb.1:
liveins: $x0, $x1
; CHECK-LABEL: name: shuffle_vector_to_unmerge_concat_hi
; CHECK: liveins: $x0, $x1
; CHECK-NEXT: {{ $}}
; CHECK-NEXT: [[COPY:%[0-9]+]]:_(<16 x s32>) = COPY $x0
; CHECK-NEXT: [[UV:%[0-9]+]]:_(<8 x s32>), [[UV1:%[0-9]+]]:_(<8 x s32>) = G_UNMERGE_VALUES [[COPY]](<16 x s32>)
; CHECK-NEXT: [[CONCAT_VECTORS:%[0-9]+]]:_(<16 x s32>) = G_CONCAT_VECTORS [[UV1]](<8 x s32>), [[UV1]](<8 x s32>)
; CHECK-NEXT: PseudoRET implicit $lr, implicit [[CONCAT_VECTORS]](<16 x s32>)
%1:_(<16 x s32>) = COPY $x0
%2:_(<16 x s32>) = COPY $x1
%0:_(<16 x s32>) = G_SHUFFLE_VECTOR %1(<16 x s32>), %2(<16 x s32>), shufflemask(8, 9, 10, 11, 12, 13, 14, 15, 8, 9, 10, 11, 12, 13, 14, 15)
PseudoRET implicit $lr, implicit %0
...
30 changes: 30 additions & 0 deletions llvm/test/CodeGen/AIE/aie2p/shufflevec.ll
Original file line number Diff line number Diff line change
Expand Up @@ -580,3 +580,33 @@ entry:
%shuffle = shufflevector <8 x i32> %a, <8 x i32> %b, <8 x i32> <i32 4, i32 5, i32 6, i32 7, i32 8, i32 9, i32 10, i32 11>
ret <8 x i32> %shuffle
}

define <16 x i32> @shuffle_vector_to_unmerge_concat_lo(<16 x i32> noundef %a, <16 x i32> noundef %b) {
; CHECK-LABEL: shuffle_vector_to_unmerge_concat_lo:
; CHECK: .p2align 4
; CHECK-NEXT: // %bb.0: // %entry
; CHECK-NEXT: nopa ; nopb ; nops ; ret lr; nopm ; nopv
; CHECK-NEXT: nopx // Delay Slot 5
; CHECK-NEXT: nop // Delay Slot 4
; CHECK-NEXT: vmov x0, x2 // Delay Slot 3
; CHECK-NEXT: vmov wh0, wl0 // Delay Slot 2
; CHECK-NEXT: nop // Delay Slot 1
entry:
%shuffle = shufflevector <16 x i32> %a, <16 x i32> %b, <16 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7, i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7>
ret <16 x i32> %shuffle
}

define <16 x i32> @shuffle_vector_to_unmerge_concat_hi(<16 x i32> noundef %a, <16 x i32> noundef %b) {
; CHECK-LABEL: shuffle_vector_to_unmerge_concat_hi:
; CHECK: .p2align 4
; CHECK-NEXT: // %bb.0: // %entry
; CHECK-NEXT: nopa ; nopb ; nops ; ret lr; nopm ; nopv
; CHECK-NEXT: nopx // Delay Slot 5
; CHECK-NEXT: nop // Delay Slot 4
; CHECK-NEXT: vmov x0, x2 // Delay Slot 3
; CHECK-NEXT: vmov wl0, wh0 // Delay Slot 2
; CHECK-NEXT: nop // Delay Slot 1
entry:
%shuffle = shufflevector <16 x i32> %a, <16 x i32> %b, <16 x i32> <i32 8, i32 9, i32 10, i32 11, i32 12, i32 13, i32 14, i32 15, i32 8, i32 9, i32 10, i32 11, i32 12, i32 13, i32 14, i32 15>
ret <16 x i32> %shuffle
}