Skip to content

Commit

Permalink
[AIEX] Add MaskMatch classes in AIECombineHelper
Browse files Browse the repository at this point in the history
  • Loading branch information
katerynamuts committed Feb 13, 2025
1 parent d269033 commit 18f174f
Show file tree
Hide file tree
Showing 2 changed files with 133 additions and 88 deletions.
193 changes: 105 additions & 88 deletions llvm/lib/Target/AIE/AIECombinerHelper.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,69 @@ cl::opt<bool> CombineVecShiftByZero(
"aie-combine-vec-shift-by-zero", cl::init(true), cl::Hidden,
cl::desc("Combine vectors shift by zero into copies."));

bool MaskMatch::isValidMask(ArrayRef<int> Mask) const {
bool FirstNotUndef = true;
for (unsigned Idx = 0; Idx < Mask.size(); ++Idx) {
if (Mask[Idx] == -1)
continue;

// Find the start value of the mask
if (FirstNotUndef) {
// Get the start value
const unsigned MaskStart = Mask[Idx] - Idx % (Period == 0 ? 1 : Period);

if (MaskStart != Height)
return false;

FirstNotUndef = false;
}

// Check not undef values (not -1) of the mask
if ((unsigned)Mask[Idx] != getMaskValue(Idx))
return false;
}

return true;
}

bool MaskMatch::isMaskWithAllUndefs(ArrayRef<int> Mask) {
for (unsigned I = 0; I < Mask.size(); ++I) {
if (Mask[I] != -1)
return false;
}
return true;
}

std::optional<unsigned> MaskMatch::getHeight(ArrayRef<int> Mask,
unsigned Period) {
for (unsigned I = 0; I < Mask.size(); ++I) {
if (Mask[I] != -1)
return Mask[I] - I % (Period == 0 ? 1 : Period);
}
return std::nullopt;
}

/// This function returns the unique index in the shuffle mask \p Mask if the
/// unique index exists.
std::optional<int> MaskMatch::getUniqueIndex(ArrayRef<int> Mask) {
std::optional<int> UniqOpIdx;
for (unsigned I = 0; I < Mask.size(); I++) {
int Idx = Mask[I];
if (Idx == -1)
continue;

if (!UniqOpIdx) {
UniqOpIdx = Idx;
continue;
}

if (UniqOpIdx != Idx) {
return std::nullopt;
}
}
return UniqOpIdx;
}

MachineInstr *findPreIncMatch(MachineInstr &MemI, MachineRegisterInfo &MRI,
CombinerHelper &Helper,
AIELoadStoreCombineMatchData &MatchData,
Expand Down Expand Up @@ -1765,6 +1828,15 @@ bool llvm::matchBroadcastElement(MachineInstr &MI, MachineRegisterInfo &MRI,
return true;
}

/// \returns true if it is possible to combine the shuffle vector to VSEL.
/// E.g.:
/// From : %0:_(<16 x s32>) = COPY $x0
/// %1:_(<16 x s32>) = COPY $x1
/// %2:_(<16 x s32>) = G_SHUFFLE_VECTOR %X(<16 x s32>), %1(<16 x s32>),
/// shufflemask(0, 1, 2, 3, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30,
/// 31)
/// To : %3:_(s32) = G_CONSTANT i32 65520
/// %4:_(<16 x s32>) = G_AIE_VSEL %0, %1, %3(s32)
bool llvm::matchShuffleToVSel(MachineInstr &MI, MachineRegisterInfo &MRI,
const AIEBaseInstrInfo &TII,
BuildFnTy &MatchInfo) {
Expand All @@ -1785,13 +1857,10 @@ bool llvm::matchShuffleToVSel(MachineInstr &MI, MachineRegisterInfo &MRI,
if (NumDstElems != NumSrcElems)
return false;

// Check that the shuffle mask can be converted into VSel mask:
// The mask contains only -1
if (std::all_of(Mask.begin(), Mask.end(),
[&](int Value) { return Value == -1; })) {
if (MaskMatch::isMaskWithAllUndefs(Mask))
return false;
}

// Check that the shuffle mask can be converted into VSel mask:
// 1. The shuffle mask doesn't contain indices that correspond to the same
// index in Src1 and Src2, i.e., for each i only the i-th element from Src1 or
// the i-th element from Src2 is used.
Expand Down Expand Up @@ -1820,27 +1889,6 @@ bool llvm::matchShuffleToVSel(MachineInstr &MI, MachineRegisterInfo &MRI,
return true;
}

/// This function returns the unique index in the shuffle mask \p Mask if the
/// unique index exists.
static std::optional<int> getUniqueIndex(ArrayRef<int> Mask) {
std::optional<int> UniqOpIdx;
for (unsigned I = 0; I < Mask.size(); I++) {
int Idx = Mask[I];
if (Idx < 0)
continue;

if (!UniqOpIdx) {
UniqOpIdx = Idx;
continue;
}

if (UniqOpIdx != Idx) {
return std::nullopt;
}
}
return UniqOpIdx;
}

/// \returns true if it is possible to combine the shuffle vector with a mask
/// that extracts an element from the first source vector and broadcasts
/// it. E.g.:
Expand All @@ -1856,7 +1904,7 @@ static bool matchShuffleToVecEltBroadcast(MachineInstr &MI,
BuildFnTy &MatchInfo) {
ArrayRef<int> Mask = MI.getOperand(3).getShuffleMask();

std::optional<int> UniqOpIdx = getUniqueIndex(Mask);
std::optional<int> UniqOpIdx = MaskMatch::getUniqueIndex(Mask);
if (!UniqOpIdx)
return false;

Expand All @@ -1873,24 +1921,6 @@ static bool matchShuffleToVecEltBroadcast(MachineInstr &MI,
return true;
}

/// A sequential mask with \p StartValue and \p NumElems is generated. If \p
/// Mask is equivalent to the generated sequential mask, the method returns
/// true. Otherwise, false.
static bool checkSequentialMask(const ArrayRef<int> Mask, unsigned StartValue,
unsigned NumElems) {
if (Mask.size() != NumElems)
return false;

auto SeqMask = createSequentialMask(StartValue, NumElems, 0);

for (unsigned I = 0; I < NumElems; ++I) {
if (Mask[I] != -1 && Mask[I] != SeqMask[I])
return false;
}

return true;
}

/// Check prerequisites to extract a subvector
static bool checkExtractSubvectorPrerequisites(const AIEBaseInstrInfo &TII,
const LLT DstTy,
Expand Down Expand Up @@ -2124,12 +2154,15 @@ bool llvm::matchShuffleToExtractSubvec(MachineInstr &MI,
if (NumSrc1Elems % NumDstElems != 0)
return false;

if (MaskMatch::isMaskWithAllUndefs(Mask))
return false;

const unsigned NumSubVectors = NumSrc1Elems / NumDstElems;
auto GetSubvecExtractIdx = [=, &Mask]() -> std::optional<unsigned> {
for (unsigned SubVecIdx = 0; SubVecIdx < NumSubVectors; ++SubVecIdx) {
if (checkSequentialMask(Mask, SubVecIdx * NumDstElems, NumDstElems)) {
MaskMatch SequentialMask{/*Height*/ SubVecIdx * NumDstElems};
if (SequentialMask.isValidMask(Mask))
return SubVecIdx;
}
}

return std::nullopt;
Expand Down Expand Up @@ -2189,30 +2222,17 @@ static bool matchShuffleToSubvecBroadcast(MachineInstr &MI,
if (Mask[0] != -1 && Mask[0] % SplatMaskLen != 0)
return std::nullopt;

// Find the start value of the splat mask and check that the mask is valid
bool ValidMask = true;
int SplatMaskStart = -1;
for (unsigned I = 0; I < MaskSize; ++I) {
if (Mask[I] == -1)
continue;

if (SplatMaskStart == -1) {
// First Mask[I]!=-1
// Get the start value
SplatMaskStart = Mask[I] - I % SplatMaskLen;

if (SplatMaskStart % SplatMaskLen != 0)
return std::nullopt;

} else if ((unsigned)Mask[I] != SplatMaskStart + I % SplatMaskLen) {
// Check the rest not undef values (not -1) of the mask
ValidMask = false;
break;
}
}
// Get Height (start value)
std::optional<unsigned> Height =
MaskMatch::getHeight(Mask, /*Period*/ SplatMaskLen);
if (!Height)
return std::nullopt;

if (ValidMask)
return std::make_pair(SplatMaskStart, SplatMaskLen);
// Check the mask
MaskMatch SequentialPeriodicMask{/*Height*/ Height.value(),
/*Period*/ SplatMaskLen};
if (SequentialPeriodicMask.isValidMask(Mask))
return std::make_pair(Height.value(), SplatMaskLen);
}
return std::nullopt;
};
Expand Down Expand Up @@ -2272,10 +2292,11 @@ static bool matchShuffleToVecBroadcast(MachineInstr &MI,
return false;
}

for (unsigned I = 0; I < Mask.size(); ++I) {
if (Mask[I] != -1 && (unsigned)Mask[I] != I % NumSrcElems)
return false;
}
// Check the mask
MaskMatch SequentialPeriodicMask{/*Height*/ 0,
/*Period*/ NumSrcElems};
if (!SequentialPeriodicMask.isValidMask(Mask))
return false;

MatchInfo = [=, &MRI](MachineIRBuilder &B) {
buildBroadcastVector(B, MRI, Src1Reg, DstReg);
Expand All @@ -2288,6 +2309,12 @@ bool llvm::matchShuffleToBroadcast(MachineInstr &MI, MachineRegisterInfo &MRI,
const AIEBaseInstrInfo &TII,
BuildFnTy &MatchInfo) {
assert(MI.getOpcode() == TargetOpcode::G_SHUFFLE_VECTOR);

ArrayRef<int> Mask = MI.getOperand(3).getShuffleMask();

if (MaskMatch::isMaskWithAllUndefs(Mask))
return false;

if (matchShuffleToVecBroadcast(MI, MRI, TII, MatchInfo))
return true;
if (matchShuffleToVecEltBroadcast(MI, MRI, TII, MatchInfo))
Expand Down Expand Up @@ -2322,23 +2349,13 @@ bool llvm::matchShuffleToCopy(MachineInstr &MI, MachineRegisterInfo &MRI,
if (Mask.size() != NumSrcElems)
return false;

// If the mask has only -1 (undef), do nothing
auto AllUndefs = [NumSrcElems](const ArrayRef<int> &Mask) -> bool {
for (unsigned I = 0; I < NumSrcElems; ++I) {
if (Mask[I] != -1)
return false;
}
return true;
};

if (AllUndefs(Mask))
if (MaskMatch::isMaskWithAllUndefs(Mask))
return false;

// Check that the mask is sequential
for (unsigned I = 0; I < NumSrcElems; ++I) {
if (Mask[I] != -1 && Mask[I] != (int)I)
return false;
}
MaskMatch SequentialMask{/*Height*/ 0};
if (!SequentialMask.isValidMask(Mask))
return false;

MatchInfo = [=](MachineIRBuilder &B) { B.buildCopy(DstReg, Src1Reg); };

Expand Down
28 changes: 28 additions & 0 deletions llvm/lib/Target/AIE/AIECombinerHelper.h
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,34 @@ struct AIELoadStoreCombineMatchData {
bool RemoveInstr;
};

/// The mask is represented by a sawtooth function F with Period, Height and
/// Amplitude, i.e., F(idx + Period) = F(idx) = Height + idx * Amplitude, where
/// idx >= 0.
/// Example: mask = (4, 6, 8, 4, 6, 8) <=> Height=4, Amplitude=2, Period=3
class MaskMatch {
public:
MaskMatch(unsigned MaskHeight, unsigned MaskPeriod = 0, int MaskAmplitude = 1)
: Period{MaskPeriod}, Height{MaskHeight}, Amplitude{MaskAmplitude} {}

bool isValidMask(ArrayRef<int> Mask) const;
unsigned getHeight() const { return Height; }

static bool isMaskWithAllUndefs(ArrayRef<int> Mask);
static std::optional<unsigned> getHeight(ArrayRef<int> Mask, unsigned Period);
static std::optional<int> getUniqueIndex(ArrayRef<int> Mask);

protected:
unsigned getMaskValue(unsigned Idx) const {
unsigned BaseIdx = Period == 0 ? Idx : Idx % Period;
return Height + BaseIdx * Amplitude;
}

unsigned Period = 0;
unsigned Height = 0;
/// Negative amplitude can be used for reverse mask patterns.
int Amplitude = 1;
};

/// Look for any PtrAdd instruction that use the same base as \a MI that can be
/// combined with it and stores it in \a MatchData
/// \return true if an instruction is found
Expand Down

0 comments on commit 18f174f

Please sign in to comment.