diff --git a/llvm/lib/Target/AIE/AIECombinerHelper.cpp b/llvm/lib/Target/AIE/AIECombinerHelper.cpp index e834612c9eda..37a8bacb1726 100644 --- a/llvm/lib/Target/AIE/AIECombinerHelper.cpp +++ b/llvm/lib/Target/AIE/AIECombinerHelper.cpp @@ -52,6 +52,69 @@ cl::opt 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 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 Mask) { + for (unsigned I = 0; I < Mask.size(); ++I) { + if (Mask[I] != -1) + return false; + } + return true; +} + +std::optional MaskMatch::getHeight(ArrayRef 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 MaskMatch::getUniqueIndex(ArrayRef Mask) { + std::optional 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, @@ -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) { @@ -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. @@ -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 getUniqueIndex(ArrayRef Mask) { - std::optional 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.: @@ -1856,7 +1904,7 @@ static bool matchShuffleToVecEltBroadcast(MachineInstr &MI, BuildFnTy &MatchInfo) { ArrayRef Mask = MI.getOperand(3).getShuffleMask(); - std::optional UniqOpIdx = getUniqueIndex(Mask); + std::optional UniqOpIdx = MaskMatch::getUniqueIndex(Mask); if (!UniqOpIdx) return false; @@ -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 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, @@ -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 { 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; @@ -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 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; }; @@ -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); @@ -2288,6 +2309,12 @@ bool llvm::matchShuffleToBroadcast(MachineInstr &MI, MachineRegisterInfo &MRI, const AIEBaseInstrInfo &TII, BuildFnTy &MatchInfo) { assert(MI.getOpcode() == TargetOpcode::G_SHUFFLE_VECTOR); + + ArrayRef 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)) @@ -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 &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); }; diff --git a/llvm/lib/Target/AIE/AIECombinerHelper.h b/llvm/lib/Target/AIE/AIECombinerHelper.h index caeb2c5d5452..75ada55c0692 100644 --- a/llvm/lib/Target/AIE/AIECombinerHelper.h +++ b/llvm/lib/Target/AIE/AIECombinerHelper.h @@ -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 Mask) const; + unsigned getHeight() const { return Height; } + + static bool isMaskWithAllUndefs(ArrayRef Mask); + static std::optional getHeight(ArrayRef Mask, unsigned Period); + static std::optional getUniqueIndex(ArrayRef 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