Vectorize/SandboxVectorizer/InstrMaps.h

0001 //===- InstrMaps.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 #ifndef LLVM_TRANSFORMS_VECTORIZE_SANDBOXVEC_PASSES_INSTRMAPS_H
0010 #define LLVM_TRANSFORMS_VECTORIZE_SANDBOXVEC_PASSES_INSTRMAPS_H
0011
0012 #include "llvm/ADT/ArrayRef.h"
0013 #include "llvm/ADT/DenseMap.h"
0014 #include "llvm/ADT/SmallSet.h"
0015 #include "llvm/ADT/SmallVector.h"
0016 #include "llvm/SandboxIR/Context.h"
0017 #include "llvm/SandboxIR/Instruction.h"
0018 #include "llvm/SandboxIR/Value.h"
0019 #include "llvm/Support/Casting.h"
0020 #include "llvm/Support/raw_ostream.h"
0021 #include <algorithm>
0022
0023 namespace llvm::sandboxir {
0024
0025 /// Maps the original instructions to the vectorized instrs and the reverse.
0026 /// For now an original instr can only map to a single vector.
0027 class InstrMaps {
0028   /// A map from the original values that got combined into vectors, to the
0029   /// vector value(s).
0030   DenseMap<Value *, Value *> OrigToVectorMap;
0031   /// A map from the vector value to a map of the original value to its lane.
0032   /// Please note that for constant vectors, there may multiple original values
0033   /// with the same lane, as they may be coming from vectorizing different
0034   /// original values.
0035   DenseMap<Value *, DenseMap<Value *, unsigned>> VectorToOrigLaneMap;
0036   Context &Ctx;
0037   std::optional<Context::CallbackID> EraseInstrCB;
0038
0039 private:
0040   void notifyEraseInstr(Value *V) {
0041     // We don't know if V is an original or a vector value.
0042     auto It = OrigToVectorMap.find(V);
0043     if (It != OrigToVectorMap.end()) {
0044       // V is an original value.
0045       // Remove it from VectorToOrigLaneMap.
0046       Value *Vec = It->second;
0047       VectorToOrigLaneMap[Vec].erase(V);
0048       // Now erase V from OrigToVectorMap.
0049       OrigToVectorMap.erase(It);
0050     } else {
0051       // V is a vector value.
0052       // Go over the original values it came from and remove them from
0053       // OrigToVectorMap.
0054       for (auto [Orig, Lane] : VectorToOrigLaneMap[V])
0055         OrigToVectorMap.erase(Orig);
0056       // Now erase V from VectorToOrigLaneMap.
0057       VectorToOrigLaneMap.erase(V);
0058     }
0059   }
0060
0061 public:
0062   InstrMaps(Context &Ctx) : Ctx(Ctx) {
0063     EraseInstrCB = Ctx.registerEraseInstrCallback(
0064         [this](Instruction *I) { notifyEraseInstr(I); });
0065   }
0066   ~InstrMaps() { Ctx.unregisterEraseInstrCallback(*EraseInstrCB); }
0067   /// \Returns the vector value that we got from vectorizing \p Orig, or
0068   /// nullptr if not found.
0069   Value *getVectorForOrig(Value *Orig) const {
0070     auto It = OrigToVectorMap.find(Orig);
0071     return It != OrigToVectorMap.end() ? It->second : nullptr;
0072   }
0073   /// \Returns the lane of \p Orig before it got vectorized into \p Vec, or
0074   /// nullopt if not found.
0075   std::optional<unsigned> getOrigLane(Value *Vec, Value *Orig) const {
0076     auto It1 = VectorToOrigLaneMap.find(Vec);
0077     if (It1 == VectorToOrigLaneMap.end())
0078       return std::nullopt;
0079     const auto &OrigToLaneMap = It1->second;
0080     auto It2 = OrigToLaneMap.find(Orig);
0081     if (It2 == OrigToLaneMap.end())
0082       return std::nullopt;
0083     return It2->second;
0084   }
0085   /// Update the map to reflect that \p Origs got vectorized into \p Vec.
0086   void registerVector(ArrayRef<Value *> Origs, Value *Vec) {
0087     auto &OrigToLaneMap = VectorToOrigLaneMap[Vec];
0088     for (auto [Lane, Orig] : enumerate(Origs)) {
0089       auto Pair = OrigToVectorMap.try_emplace(Orig, Vec);
0090       assert(Pair.second && "Orig already exists in the map!");
0091       (void)Pair;
0092       OrigToLaneMap[Orig] = Lane;
0093     }
0094   }
0095   void clear() {
0096     OrigToVectorMap.clear();
0097     VectorToOrigLaneMap.clear();
0098   }
0099 #ifndef NDEBUG
0100   void print(raw_ostream &OS) const {
0101     OS << "OrigToVectorMap:\n";
0102     for (auto [Orig, Vec] : OrigToVectorMap)
0103       OS << *Orig << " : " << *Vec << "\n";
0104   }
0105   LLVM_DUMP_METHOD void dump() const;
0106 #endif
0107 };
0108 } // namespace llvm::sandboxir
0109
0110 #endif // LLVM_TRANSFORMS_VECTORIZE_SANDBOXVEC_PASSES_INSTRMAPS_H