1 //===- InstrMaps.h ----------------------------------------------*- C++ -*-===// 2 // 3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. 4 // See https://llvm.org/LICENSE.txt for license information. 5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception 6 // 7 //===----------------------------------------------------------------------===// 8 9 #ifndef LLVM_TRANSFORMS_VECTORIZE_SANDBOXVEC_PASSES_INSTRMAPS_H 10 #define LLVM_TRANSFORMS_VECTORIZE_SANDBOXVEC_PASSES_INSTRMAPS_H 11 12 #include "llvm/ADT/ArrayRef.h" 13 #include "llvm/ADT/DenseMap.h" 14 #include "llvm/ADT/SmallSet.h" 15 #include "llvm/ADT/SmallVector.h" 16 #include "llvm/SandboxIR/Context.h" 17 #include "llvm/SandboxIR/Instruction.h" 18 #include "llvm/SandboxIR/Value.h" 19 #include "llvm/Support/Casting.h" 20 #include "llvm/Support/raw_ostream.h" 21 #include "llvm/Transforms/Vectorize/SandboxVectorizer/VecUtils.h" 22 #include <algorithm> 23 24 namespace llvm::sandboxir { 25 26 class LegalityResult; 27 28 struct Action { 29 unsigned Idx = 0; 30 const LegalityResult *LegalityRes = nullptr; 31 SmallVector<Value *, 4> Bndl; 32 SmallVector<Value *> UserBndl; 33 unsigned Depth; 34 SmallVector<Action *> Operands; 35 Value *Vec = nullptr; ActionAction36 Action(const LegalityResult *LR, ArrayRef<Value *> B, ArrayRef<Value *> UB, 37 unsigned Depth) 38 : LegalityRes(LR), Bndl(B), UserBndl(UB), Depth(Depth) {} 39 #ifndef NDEBUG 40 void print(raw_ostream &OS) const; 41 void dump() const; 42 friend raw_ostream &operator<<(raw_ostream &OS, const Action &A) { 43 A.print(OS); 44 return OS; 45 } 46 #endif // NDEBUG 47 }; 48 49 /// Maps the original instructions to the vectorized instrs and the reverse. 50 /// For now an original instr can only map to a single vector. 51 class InstrMaps { 52 /// A map from the original values that got combined into vectors, to the 53 /// vectorization Action. 54 DenseMap<Value *, Action *> OrigToVectorMap; 55 /// A map from the vec Action to a map of the original value to its lane. 56 /// Please note that for constant vectors, there may multiple original values 57 /// with the same lane, as they may be coming from vectorizing different 58 /// original values. 59 DenseMap<Action *, DenseMap<Value *, unsigned>> VectorToOrigLaneMap; 60 std::optional<Context::CallbackID> EraseInstrCB; 61 62 public: 63 InstrMaps() = default; 64 ~InstrMaps() = default; 65 /// \Returns the vector value that we got from vectorizing \p Orig, or 66 /// nullptr if not found. getVectorForOrig(Value * Orig)67 Action *getVectorForOrig(Value *Orig) const { 68 auto It = OrigToVectorMap.find(Orig); 69 return It != OrigToVectorMap.end() ? It->second : nullptr; 70 } 71 /// \Returns the lane of \p Orig before it got vectorized into \p Vec, or 72 /// nullopt if not found. getOrigLane(Action * Vec,Value * Orig)73 std::optional<unsigned> getOrigLane(Action *Vec, Value *Orig) const { 74 auto It1 = VectorToOrigLaneMap.find(Vec); 75 if (It1 == VectorToOrigLaneMap.end()) 76 return std::nullopt; 77 const auto &OrigToLaneMap = It1->second; 78 auto It2 = OrigToLaneMap.find(Orig); 79 if (It2 == OrigToLaneMap.end()) 80 return std::nullopt; 81 return It2->second; 82 } 83 /// Update the map to reflect that \p Origs got vectorized into \p Vec. registerVector(ArrayRef<Value * > Origs,Action * Vec)84 void registerVector(ArrayRef<Value *> Origs, Action *Vec) { 85 auto &OrigToLaneMap = VectorToOrigLaneMap[Vec]; 86 unsigned Lane = 0; 87 for (Value *Orig : Origs) { 88 auto Pair = OrigToVectorMap.try_emplace(Orig, Vec); 89 assert(Pair.second && "Orig already exists in the map!"); 90 (void)Pair; 91 OrigToLaneMap[Orig] = Lane; 92 Lane += VecUtils::getNumLanes(Orig); 93 } 94 } clear()95 void clear() { 96 OrigToVectorMap.clear(); 97 VectorToOrigLaneMap.clear(); 98 } 99 #ifndef NDEBUG print(raw_ostream & OS)100 void print(raw_ostream &OS) const { 101 OS << "OrigToVectorMap:\n"; 102 for (auto [Orig, Vec] : OrigToVectorMap) 103 OS << *Orig << " : " << *Vec << "\n"; 104 } 105 LLVM_DUMP_METHOD void dump() const; 106 #endif 107 }; 108 } // namespace llvm::sandboxir 109 110 #endif // LLVM_TRANSFORMS_VECTORIZE_SANDBOXVEC_PASSES_INSTRMAPS_H 111