Skip to content

Commit

Permalink
[AIEX] Combine G_SHUFFLE_VECTOR to UNMERGE
Browse files Browse the repository at this point in the history
  • Loading branch information
@katerynamuts
katerynamuts committed Feb 6, 2025
1 parent 9af5b07 commit 7425bfd
Show file tree
Hide file tree
Showing 3 changed files with 197 additions and 263 deletions.
158 changes: 109 additions & 49 deletions llvm/lib/Target/AIE/AIECombinerHelper.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1907,6 +1907,38 @@ static bool checkSequentialMask(const ArrayRef<int> Mask, unsigned StartValue,
return true;
}

/// Match something like this:
/// %1:_(<16 x s32>) = COPY $x0
/// %2:_(<16 x s32>) = COPY $x1
/// %0:_(<8 x s32>) = G_SHUFFLE_VECTOR %1(<16 x s32>), %2(<16 x s32>),
/// shufflemask(8, 9, 10, 11, 12, 13, 14, 15)
/// PseudoRET implicit $lr, implicit %0

/// To convert to:
/// %1:_(<16 x s32>) = COPY $x0
/// %2:_(<8 x s32>), %3:_(<8 x s32>) = G_UNMERGE_VALUES %1(<16 x s32>)
/// PseudoRET implicit $lr, implicit %3(<8 x s32>)
static bool matchShuffleToUnmerge(MachineInstr &MI, MachineRegisterInfo &MRI,
BuildFnTy &MatchInfo, unsigned SubIdx,
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);
};
return true;
}

/// Match something like this:
/// %1:_(<16 x s16>) = COPY $wl0
/// %2:_(<16 x s16>) = COPY $wl1
Expand All @@ -1917,73 +1949,36 @@ static bool checkSequentialMask(const ArrayRef<int> Mask, unsigned StartValue,
/// %1:_(<16 x s16>) = COPY $wl0
/// %2:_(s32) = G_CONSTANT i32 1
/// %3:_(<4 x s16>) = G_AIE_EXTRACT_SUBVECTOR %1(<16 x s16>), %2(s32)
bool llvm::matchShuffleToExtractSubvec(MachineInstr &MI,
MachineRegisterInfo &MRI,
const AIEBaseInstrInfo &TII,
BuildFnTy &MatchInfo) {
assert(MI.getOpcode() == TargetOpcode::G_SHUFFLE_VECTOR);

/// NOTE: This combine works ONLY for 32- and 64-bit outputs!
static bool matchShuffleToAIEExtractSubvec(
MachineInstr &MI, MachineRegisterInfo &MRI, const AIEBaseInstrInfo &TII,
BuildFnTy &MatchInfo, unsigned SubIdx, unsigned NumSubVectors) {
const unsigned GPRSize = TII.getScalarRegSize();
const unsigned VecRegSize = TII.getBasicVecRegSize();
const unsigned ExtractSubvecNativeSrcSize = TII.getBasicVectorBitSize();

const Register DstReg = MI.getOperand(0).getReg();
const Register Src1Reg = MI.getOperand(1).getReg();
const ArrayRef<int> Mask = MI.getOperand(3).getShuffleMask();

const LLT DstTy = MRI.getType(DstReg);
const LLT Src1Ty = MRI.getType(Src1Reg);
const unsigned Src1TySize = Src1Ty.getSizeInBits();

if (!DstTy.isVector() || !Src1Ty.isVector() || Src1TySize < GPRSize ||
(DstTy.getSizeInBits() != GPRSize &&
DstTy.getSizeInBits() != 2 * GPRSize))
if (Src1TySize < GPRSize || (DstTy.getSizeInBits() != GPRSize &&
DstTy.getSizeInBits() != 2 * GPRSize))
return false;

// Currently, we cannot concat vectors of the size less than vector register
// size.
// Currently, we cannot concat vectors for the case when the size of the
// source vector is less than vector register size.
if (Src1TySize < VecRegSize)
return false;

// This should be handled by a separate combine that copies Src1Reg to
// DstReg.
if (Src1TySize == DstTy.getSizeInBits())
return false;

const unsigned NumDstElems = DstTy.getNumElements();
const unsigned NumSrc1Elems = Src1Ty.getNumElements();
const unsigned NumSubVectors = NumSrc1Elems / NumDstElems;

// Not an extract pattern
if (NumSrc1Elems < NumDstElems)
return false;

// Unlikely to select into a subregister copy
if (NumSrc1Elems % NumDstElems != 0)
return false;

auto GetSubIdx = [=, &Mask]() -> std::optional<unsigned> {
for (unsigned SubVecIdx = 0; SubVecIdx < NumSubVectors; ++SubVecIdx) {
if (checkSequentialMask(Mask, SubVecIdx * NumDstElems, NumDstElems)) {
return SubVecIdx;
}
}

return std::nullopt;
};

std::optional<unsigned> SubIdx = GetSubIdx();

// Not an extract pattern
if (!SubIdx)
return false;

const unsigned Opc = TII.getGenericExtractSubvectorOpcode();

// Natively supported source vector type
if (Src1TySize == ExtractSubvecNativeSrcSize) {
MatchInfo = [=](MachineIRBuilder &B) {
auto Cst = B.buildConstant(LLT::scalar(GPRSize), SubIdx.value());
auto Cst = B.buildConstant(LLT::scalar(GPRSize), SubIdx);
B.buildInstr(Opc, {DstReg}, {Src1Reg, Cst});
};

Expand All @@ -2006,7 +2001,7 @@ bool llvm::matchShuffleToExtractSubvec(MachineInstr &MI,
ConcatOps.push_back(ImplicitDef);
}
B.buildConcatVectors({NewSrcReg}, ConcatOps);
auto Cst = B.buildConstant(LLT::scalar(GPRSize), SubIdx.value());
auto Cst = B.buildConstant(LLT::scalar(GPRSize), SubIdx);
B.buildInstr(Opc, {DstReg}, {NewSrcReg, Cst});
};
return true;
Expand All @@ -2016,10 +2011,10 @@ bool llvm::matchShuffleToExtractSubvec(MachineInstr &MI,
MatchInfo = [=, &MRI](MachineIRBuilder &B) {
const unsigned SizeCoefficient = Src1TySize / ExtractSubvecNativeSrcSize;
const unsigned NumSubVectorsNativeSize = NumSubVectors / SizeCoefficient;
unsigned NewSubIdx = SubIdx.value() % NumSubVectorsNativeSize;
unsigned NewSubIdx = SubIdx % NumSubVectorsNativeSize;

SmallVector<Register, 4> SubRegs;
unsigned NewSrcRegPosition = SubIdx.value() / NumSubVectorsNativeSize;
unsigned NewSrcRegPosition = SubIdx / NumSubVectorsNativeSize;
for (unsigned I = 0; I < SizeCoefficient; ++I) {
if (I == NewSrcRegPosition)
SubRegs.push_back(NewSrcReg);
Expand All @@ -2033,3 +2028,68 @@ bool llvm::matchShuffleToExtractSubvec(MachineInstr &MI,
};
return true;
}

/// The method does some checks and calls matchShuffleToAIEExtractSubvec and
/// matchShuffleToUnmerge which extract subvectors is possible.
bool llvm::matchShuffleToExtractSubvec(MachineInstr &MI,
MachineRegisterInfo &MRI,
const AIEBaseInstrInfo &TII,
BuildFnTy &MatchInfo) {
assert(MI.getOpcode() == TargetOpcode::G_SHUFFLE_VECTOR);

const Register DstReg = MI.getOperand(0).getReg();
const Register Src1Reg = MI.getOperand(1).getReg();
ArrayRef<int> Mask = MI.getOperand(3).getShuffleMask();

const LLT DstTy = MRI.getType(DstReg);
const LLT Src1Ty = MRI.getType(Src1Reg);
const unsigned Src1TySize = Src1Ty.getSizeInBits();

if (!DstTy.isVector() || !Src1Ty.isVector())
return false;

// Boolean vectors are unlikely to select into subregister copy
if (DstTy.getElementType() == LLT::scalar(1))
return false;

// This should be handled by a separate combine that copies Src1Reg to
// DstReg.
if (Src1TySize == DstTy.getSizeInBits())
return false;

const unsigned NumDstElems = DstTy.getNumElements();
const unsigned NumSrc1Elems = Src1Ty.getNumElements();
const unsigned NumSubVectors = NumSrc1Elems / NumDstElems;

// Not an extract pattern
if (NumSrc1Elems <= NumDstElems)
return false;

// Unlikely to select into a subregister copy
if (NumSrc1Elems % NumDstElems != 0)
return false;

auto GetSubIdx = [=, &Mask]() -> std::optional<unsigned> {
for (unsigned SubVecIdx = 0; SubVecIdx < NumSubVectors; ++SubVecIdx) {
if (checkSequentialMask(Mask, SubVecIdx * NumDstElems, NumDstElems)) {
return SubVecIdx;
}
}

return std::nullopt;
};

std::optional<unsigned> SubIdx = GetSubIdx();

// Not an extract pattern
if (!SubIdx)
return false;

if (matchShuffleToAIEExtractSubvec(MI, MRI, TII, MatchInfo, SubIdx.value(),
NumSubVectors))
return true;
if (matchShuffleToUnmerge(MI, MRI, MatchInfo, SubIdx.value(), NumSubVectors))
return true;

return false;
}
54 changes: 54 additions & 0 deletions llvm/test/CodeGen/AIE/GlobalISel/prelegalizercombiner-shuffle-vector.mir
View file
Original file line number Diff line number Diff line change
Expand Up @@ -615,3 +615,57 @@ body: |
%0:_(<8 x s64>) = G_SHUFFLE_VECTOR %1(<8 x s64>), %2(<8 x s64>), shufflemask(-1, 5, -1, 3, 4, 5, 6, 7)
PseudoRET implicit $lr, implicit %0
...

# Test G_SHUFFLE_VECTOR to UNMERGE
---
name: shuffle_vector_unmerge_lo
tracksRegLiveness: true
body: |
bb.1:
liveins: $x0, $x1
; CHECK-LABEL: name: shuffle_vector_unmerge_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: PseudoRET implicit $lr, implicit [[AIE_UNPAD_VECTOR]](<8 x s32>)
%1:_(<16 x s32>) = COPY $x0
%2:_(<16 x s32>) = COPY $x1
%0:_(<8 x s32>) = G_SHUFFLE_VECTOR %1(<16 x s32>), %2(<16 x s32>), shufflemask(0, 1, 2, 3, 4, 5, 6, 7)
PseudoRET implicit $lr, implicit %0
...
---
name: shuffle_vector_unmerge_hi
tracksRegLiveness: true
body: |
bb.1:
liveins: $x0, $x1
; CHECK-LABEL: name: shuffle_vector_unmerge_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: PseudoRET implicit $lr, implicit [[UV1]](<8 x s32>)
%1:_(<16 x s32>) = COPY $x0
%2:_(<16 x s32>) = COPY $x1
%0:_(<8 x s32>) = G_SHUFFLE_VECTOR %1(<16 x s32>), %2(<16 x s32>), shufflemask(8, 9, 10, 11, 12, 13, 14, 15)
PseudoRET implicit $lr, implicit %0
...
---
name: shuffle_vector_subreg_boundary_cross_invalid
tracksRegLiveness: true
body: |
bb.1:
liveins: $x0, $x1
; CHECK-LABEL: name: shuffle_vector_subreg_boundary_cross_invalid
; CHECK: liveins: $x0, $x1
; CHECK-NEXT: {{ $}}
; CHECK-NEXT: [[COPY:%[0-9]+]]:_(<16 x s32>) = COPY $x0
; CHECK-NEXT: [[COPY1:%[0-9]+]]:_(<16 x s32>) = COPY $x1
; CHECK-NEXT: [[SHUF:%[0-9]+]]:_(<8 x s32>) = G_SHUFFLE_VECTOR [[COPY]](<16 x s32>), [[COPY1]], shufflemask(4, 5, 6, 7, 8, 9, 10, 11)
; CHECK-NEXT: PseudoRET implicit $lr, implicit [[SHUF]](<8 x s32>)
%1:_(<16 x s32>) = COPY $x0
%2:_(<16 x s32>) = COPY $x1
%0:_(<8 x s32>) = G_SHUFFLE_VECTOR %1(<16 x s32>), %2(<16 x s32>), shufflemask(4, 5, 6, 7, 8, 9, 10, 11)
PseudoRET implicit $lr, implicit %0
...
Loading

0 comments on commit 7425bfd

Please sign in to comment.