Vectorize/SandboxVectorizer/Legality.h

0001 //===- Legality.h -----------------------------------------------*- C++ -*-===//
0002 //
0003 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
0004 // See https://llvm.org/LICENSE.txt for license information.
0005 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
0006 //
0007 //===----------------------------------------------------------------------===//
0008 //
0009 // Legality checks for the Sandbox Vectorizer.
0010 //
0011
0012 #ifndef LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_LEGALITY_H
0013 #define LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_LEGALITY_H
0014
0015 #include "llvm/ADT/ArrayRef.h"
0016 #include "llvm/Analysis/ScalarEvolution.h"
0017 #include "llvm/IR/DataLayout.h"
0018 #include "llvm/Support/Casting.h"
0019 #include "llvm/Support/raw_ostream.h"
0020 #include "llvm/Transforms/Vectorize/SandboxVectorizer/Scheduler.h"
0021
0022 namespace llvm::sandboxir {
0023
0024 class LegalityAnalysis;
0025 class Value;
0026 class InstrMaps;
0027
0028 class ShuffleMask {
0029 public:
0030   using IndicesVecT = SmallVector<int, 8>;
0031
0032 private:
0033   IndicesVecT Indices;
0034
0035 public:
0036   ShuffleMask(SmallVectorImpl<int> &&Indices) : Indices(std::move(Indices)) {}
0037   ShuffleMask(std::initializer_list<int> Indices) : Indices(Indices) {}
0038   explicit ShuffleMask(ArrayRef<int> Indices) : Indices(Indices) {}
0039   operator ArrayRef<int>() const { return Indices; }
0040   /// Creates and returns an identity shuffle mask of size \p Sz.
0041   /// For example if Sz == 4 the returned mask is {0, 1, 2, 3}.
0042   static ShuffleMask getIdentity(unsigned Sz) {
0043     IndicesVecT Indices;
0044     Indices.reserve(Sz);
0045     for (auto Idx : seq<int>(0, (int)Sz))
0046       Indices.push_back(Idx);
0047     return ShuffleMask(std::move(Indices));
0048   }
0049   /// \Returns true if the mask is a perfect identity mask with consecutive
0050   /// indices, i.e., performs no lane shuffling, like 0,1,2,3...
0051   bool isIdentity() const {
0052     for (auto [Idx, Elm] : enumerate(Indices)) {
0053       if ((int)Idx != Elm)
0054         return false;
0055     }
0056     return true;
0057   }
0058   bool operator==(const ShuffleMask &Other) const {
0059     return Indices == Other.Indices;
0060   }
0061   bool operator!=(const ShuffleMask &Other) const { return !(*this == Other); }
0062   size_t size() const { return Indices.size(); }
0063   int operator[](int Idx) const { return Indices[Idx]; }
0064   using const_iterator = IndicesVecT::const_iterator;
0065   const_iterator begin() const { return Indices.begin(); }
0066   const_iterator end() const { return Indices.end(); }
0067 #ifndef NDEBUG
0068   friend raw_ostream &operator<<(raw_ostream &OS, const ShuffleMask &Mask) {
0069     Mask.print(OS);
0070     return OS;
0071   }
0072   void print(raw_ostream &OS) const {
0073     interleave(Indices, OS, [&OS](auto Elm) { OS << Elm; }, ",");
0074   }
0075   LLVM_DUMP_METHOD void dump() const;
0076 #endif
0077 };
0078
0079 enum class LegalityResultID {
0080   Pack,                    ///> Collect scalar values.
0081   Widen,                   ///> Vectorize by combining scalars to a vector.
0082   DiamondReuse,            ///> Don't generate new code, reuse existing vector.
0083   DiamondReuseWithShuffle, ///> Reuse the existing vector but add a shuffle.
0084   DiamondReuseMultiInput,  ///> Reuse more than one vector and/or scalars.
0085 };
0086
0087 /// The reason for vectorizing or not vectorizing.
0088 enum class ResultReason {
0089   NotInstructions,
0090   DiffOpcodes,
0091   DiffTypes,
0092   DiffMathFlags,
0093   DiffWrapFlags,
0094   DiffBBs,
0095   NotConsecutive,
0096   CantSchedule,
0097   Unimplemented,
0098   Infeasible,
0099 };
0100
0101 #ifndef NDEBUG
0102 struct ToStr {
0103   static const char *getLegalityResultID(LegalityResultID ID) {
0104     switch (ID) {
0105     case LegalityResultID::Pack:
0106       return "Pack";
0107     case LegalityResultID::Widen:
0108       return "Widen";
0109     case LegalityResultID::DiamondReuse:
0110       return "DiamondReuse";
0111     case LegalityResultID::DiamondReuseWithShuffle:
0112       return "DiamondReuseWithShuffle";
0113     case LegalityResultID::DiamondReuseMultiInput:
0114       return "DiamondReuseMultiInput";
0115     }
0116     llvm_unreachable("Unknown LegalityResultID enum");
0117   }
0118
0119   static const char *getVecReason(ResultReason Reason) {
0120     switch (Reason) {
0121     case ResultReason::NotInstructions:
0122       return "NotInstructions";
0123     case ResultReason::DiffOpcodes:
0124       return "DiffOpcodes";
0125     case ResultReason::DiffTypes:
0126       return "DiffTypes";
0127     case ResultReason::DiffMathFlags:
0128       return "DiffMathFlags";
0129     case ResultReason::DiffWrapFlags:
0130       return "DiffWrapFlags";
0131     case ResultReason::DiffBBs:
0132       return "DiffBBs";
0133     case ResultReason::NotConsecutive:
0134       return "NotConsecutive";
0135     case ResultReason::CantSchedule:
0136       return "CantSchedule";
0137     case ResultReason::Unimplemented:
0138       return "Unimplemented";
0139     case ResultReason::Infeasible:
0140       return "Infeasible";
0141     }
0142     llvm_unreachable("Unknown ResultReason enum");
0143   }
0144 };
0145 #endif // NDEBUG
0146
0147 /// The legality outcome is represented by a class rather than an enum class
0148 /// because in some cases the legality checks are expensive and look for a
0149 /// particular instruction that can be passed along to the vectorizer to avoid
0150 /// repeating the same expensive computation.
0151 class LegalityResult {
0152 protected:
0153   LegalityResultID ID;
0154   /// Only Legality can create LegalityResults.
0155   LegalityResult(LegalityResultID ID) : ID(ID) {}
0156   friend class LegalityAnalysis;
0157
0158   /// We shouldn't need copies.
0159   LegalityResult(const LegalityResult &) = delete;
0160   LegalityResult &operator=(const LegalityResult &) = delete;
0161
0162 public:
0163   virtual ~LegalityResult() {}
0164   LegalityResultID getSubclassID() const { return ID; }
0165 #ifndef NDEBUG
0166   virtual void print(raw_ostream &OS) const {
0167     OS << ToStr::getLegalityResultID(ID);
0168   }
0169   LLVM_DUMP_METHOD void dump() const;
0170   friend raw_ostream &operator<<(raw_ostream &OS, const LegalityResult &LR) {
0171     LR.print(OS);
0172     return OS;
0173   }
0174 #endif // NDEBUG
0175 };
0176
0177 /// Base class for results with reason.
0178 class LegalityResultWithReason : public LegalityResult {
0179   [[maybe_unused]] ResultReason Reason;
0180   LegalityResultWithReason(LegalityResultID ID, ResultReason Reason)
0181       : LegalityResult(ID), Reason(Reason) {}
0182   friend class Pack; // For constructor.
0183
0184 public:
0185   ResultReason getReason() const { return Reason; }
0186 #ifndef NDEBUG
0187   void print(raw_ostream &OS) const override {
0188     LegalityResult::print(OS);
0189     OS << " Reason: " << ToStr::getVecReason(Reason);
0190   }
0191 #endif
0192 };
0193
0194 class Widen final : public LegalityResult {
0195   friend class LegalityAnalysis;
0196   Widen() : LegalityResult(LegalityResultID::Widen) {}
0197
0198 public:
0199   static bool classof(const LegalityResult *From) {
0200     return From->getSubclassID() == LegalityResultID::Widen;
0201   }
0202 };
0203
0204 class DiamondReuse final : public LegalityResult {
0205   friend class LegalityAnalysis;
0206   Value *Vec;
0207   DiamondReuse(Value *Vec)
0208       : LegalityResult(LegalityResultID::DiamondReuse), Vec(Vec) {}
0209
0210 public:
0211   static bool classof(const LegalityResult *From) {
0212     return From->getSubclassID() == LegalityResultID::DiamondReuse;
0213   }
0214   Value *getVector() const { return Vec; }
0215 };
0216
0217 class DiamondReuseWithShuffle final : public LegalityResult {
0218   friend class LegalityAnalysis;
0219   Value *Vec;
0220   ShuffleMask Mask;
0221   DiamondReuseWithShuffle(Value *Vec, const ShuffleMask &Mask)
0222       : LegalityResult(LegalityResultID::DiamondReuseWithShuffle), Vec(Vec),
0223         Mask(Mask) {}
0224
0225 public:
0226   static bool classof(const LegalityResult *From) {
0227     return From->getSubclassID() == LegalityResultID::DiamondReuseWithShuffle;
0228   }
0229   Value *getVector() const { return Vec; }
0230   const ShuffleMask &getMask() const { return Mask; }
0231 };
0232
0233 class Pack final : public LegalityResultWithReason {
0234   Pack(ResultReason Reason)
0235       : LegalityResultWithReason(LegalityResultID::Pack, Reason) {}
0236   friend class LegalityAnalysis; // For constructor.
0237
0238 public:
0239   static bool classof(const LegalityResult *From) {
0240     return From->getSubclassID() == LegalityResultID::Pack;
0241   }
0242 };
0243
0244 /// Describes how to collect the values needed by each lane.
0245 class CollectDescr {
0246 public:
0247   /// Describes how to get a value element. If the value is a vector then it
0248   /// also provides the index to extract it from.
0249   class ExtractElementDescr {
0250     Value *V;
0251     /// The index in `V` that the value can be extracted from.
0252     /// This is nullopt if we need to use `V` as a whole.
0253     std::optional<int> ExtractIdx;
0254
0255   public:
0256     ExtractElementDescr(Value *V, int ExtractIdx)
0257         : V(V), ExtractIdx(ExtractIdx) {}
0258     ExtractElementDescr(Value *V) : V(V), ExtractIdx(std::nullopt) {}
0259     Value *getValue() const { return V; }
0260     bool needsExtract() const { return ExtractIdx.has_value(); }
0261     int getExtractIdx() const { return *ExtractIdx; }
0262   };
0263
0264   using DescrVecT = SmallVector<ExtractElementDescr, 4>;
0265   DescrVecT Descrs;
0266
0267 public:
0268   CollectDescr(SmallVectorImpl<ExtractElementDescr> &&Descrs)
0269       : Descrs(std::move(Descrs)) {}
0270   /// If all elements come from a single vector input, then return that vector
0271   /// and also the shuffle mask required to get them in order.
0272   std::optional<std::pair<Value *, ShuffleMask>> getSingleInput() const {
0273     const auto &Descr0 = *Descrs.begin();
0274     Value *V0 = Descr0.getValue();
0275     if (!Descr0.needsExtract())
0276       return std::nullopt;
0277     ShuffleMask::IndicesVecT MaskIndices;
0278     MaskIndices.push_back(Descr0.getExtractIdx());
0279     for (const auto &Descr : drop_begin(Descrs)) {
0280       if (!Descr.needsExtract())
0281         return std::nullopt;
0282       if (Descr.getValue() != V0)
0283         return std::nullopt;
0284       MaskIndices.push_back(Descr.getExtractIdx());
0285     }
0286     return std::make_pair(V0, ShuffleMask(std::move(MaskIndices)));
0287   }
0288   bool hasVectorInputs() const {
0289     return any_of(Descrs, [](const auto &D) { return D.needsExtract(); });
0290   }
0291   const SmallVector<ExtractElementDescr, 4> &getDescrs() const {
0292     return Descrs;
0293   }
0294 };
0295
0296 class DiamondReuseMultiInput final : public LegalityResult {
0297   friend class LegalityAnalysis;
0298   CollectDescr Descr;
0299   DiamondReuseMultiInput(CollectDescr &&Descr)
0300       : LegalityResult(LegalityResultID::DiamondReuseMultiInput),
0301         Descr(std::move(Descr)) {}
0302
0303 public:
0304   static bool classof(const LegalityResult *From) {
0305     return From->getSubclassID() == LegalityResultID::DiamondReuseMultiInput;
0306   }
0307   const CollectDescr &getCollectDescr() const { return Descr; }
0308 };
0309
0310 /// Performs the legality analysis and returns a LegalityResult object.
0311 class LegalityAnalysis {
0312   Scheduler Sched;
0313   /// Owns the legality result objects created by createLegalityResult().
0314   SmallVector<std::unique_ptr<LegalityResult>> ResultPool;
0315   /// Checks opcodes, types and other IR-specifics and returns a ResultReason
0316   /// object if not vectorizable, or nullptr otherwise.
0317   std::optional<ResultReason>
0318   notVectorizableBasedOnOpcodesAndTypes(ArrayRef<Value *> Bndl);
0319
0320   ScalarEvolution &SE;
0321   const DataLayout &DL;
0322   InstrMaps &IMaps;
0323
0324   /// Finds how we can collect the values in \p Bndl from the vectorized or
0325   /// non-vectorized code. It returns a map of the value we should extract from
0326   /// and the corresponding shuffle mask we need to use.
0327   CollectDescr getHowToCollectValues(ArrayRef<Value *> Bndl) const;
0328
0329 public:
0330   LegalityAnalysis(AAResults &AA, ScalarEvolution &SE, const DataLayout &DL,
0331                    Context &Ctx, InstrMaps &IMaps)
0332       : Sched(AA, Ctx), SE(SE), DL(DL), IMaps(IMaps) {}
0333   /// A LegalityResult factory.
0334   template <typename ResultT, typename... ArgsT>
0335   ResultT &createLegalityResult(ArgsT &&...Args) {
0336     ResultPool.push_back(
0337         std::unique_ptr<ResultT>(new ResultT(std::move(Args)...)));
0338     return cast<ResultT>(*ResultPool.back());
0339   }
0340   /// Checks if it's legal to vectorize the instructions in \p Bndl.
0341   /// \Returns a LegalityResult object owned by LegalityAnalysis.
0342   /// \p SkipScheduling skips the scheduler check and is only meant for testing.
0343   // TODO: Try to remove the SkipScheduling argument by refactoring the tests.
0344   const LegalityResult &canVectorize(ArrayRef<Value *> Bndl,
0345                                      bool SkipScheduling = false);
0346   void clear();
0347 };
0348
0349 } // namespace llvm::sandboxir
0350
0351 #endif // LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_LEGALITY_H