1//===- SeedCollector.cpp -------------------------------------------------===//
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#include "llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h"
10#include "../VectorizeOptions.h"
11#include "llvm/Analysis/LoopAccessAnalysis.h"
12#include "llvm/Analysis/ValueTracking.h"
13#include "llvm/IR/Type.h"
14#include "llvm/SandboxIR/Instruction.h"
15#include "llvm/SandboxIR/Utils.h"
16#include "llvm/Support/Compiler.h"
17#include "llvm/Support/Debug.h"
18
19using namespace llvm;
20namespace llvm::sandboxir {
21
22ArrayRef<Instruction *> SeedBundle::getSlice(unsigned StartIdx,
23 unsigned MaxVecRegBits,
24 bool ForcePowerOf2) {
25 // Use uint32_t here for compatibility with IsPowerOf2_32
26
27 // BitCount tracks the size of the working slice. From that we can tell
28 // when the working slice's size is a power-of-two and when it exceeds
29 // the legal size in MaxVecBits.
30 uint32_t BitCount = 0;
31 uint32_t NumElements = 0;
32 // Tracks the most recent slice where NumElements gave a power-of-2 BitCount
33 uint32_t NumElementsPowerOfTwo = 0;
34 uint32_t BitCountPowerOfTwo = 0;
35 // Can't start a slice with a used instruction.
36 assert(!isUsed(StartIdx) && "Expected unused at StartIdx");
37 for (Instruction *S : drop_begin(RangeOrContainer&: Seeds, N: StartIdx)) {
38 // Stop if this instruction is used. This needs to be done before
39 // getNumBits() because a "used" instruction may have been erased.
40 if (isUsed(Element: StartIdx + NumElements))
41 break;
42 uint32_t InstBits = Utils::getNumBits(I: S);
43 // Stop if adding it puts the slice over the limit.
44 if (BitCount + InstBits > MaxVecRegBits)
45 break;
46 NumElements++;
47 BitCount += InstBits;
48 if (ForcePowerOf2 && isPowerOf2_32(Value: BitCount)) {
49 NumElementsPowerOfTwo = NumElements;
50 BitCountPowerOfTwo = BitCount;
51 }
52 }
53 if (ForcePowerOf2) {
54 NumElements = NumElementsPowerOfTwo;
55 BitCount = BitCountPowerOfTwo;
56 }
57
58 // Return any non-empty slice
59 if (NumElements > 1) {
60 assert((!ForcePowerOf2 || isPowerOf2_32(BitCount)) &&
61 "Must be a power of two");
62 return ArrayRef<Instruction *>(&Seeds[StartIdx], NumElements);
63 }
64 return {};
65}
66
67template <typename LoadOrStoreT>
68SeedContainer::KeyT SeedContainer::getKey(LoadOrStoreT *LSI,
69 bool AllowDiffTypes) const {
70 assert((isa<LoadInst>(LSI) || isa<StoreInst>(LSI)) &&
71 "Expected Load or Store!");
72 Value *Ptr = Utils::getMemInstructionBase(LSI);
73 Instruction::Opcode Op = LSI->getOpcode();
74 Type *Ty;
75 if (AllowDiffTypes) {
76 Ty = nullptr;
77 } else {
78 Ty = Utils::getExpectedType(V: LSI);
79 if (auto *VTy = dyn_cast<VectorType>(Val: Ty))
80 Ty = VTy->getElementType();
81 }
82 return {Ptr, Ty, Op};
83}
84
85// Explicit instantiations
86template SeedContainer::KeyT
87SeedContainer::getKey<LoadInst>(LoadInst *LSI, bool AllowDiffTypes) const;
88template SeedContainer::KeyT
89SeedContainer::getKey<StoreInst>(StoreInst *LSI, bool AllowDiffTypes) const;
90
91bool SeedContainer::erase(Instruction *I) {
92 assert((isa<LoadInst>(I) || isa<StoreInst>(I)) && "Expected Load or Store!");
93 auto It = SeedLookupMap.find(Val: I);
94 if (It == SeedLookupMap.end())
95 return false;
96 SeedBundle *Bndl = It->second;
97 Bndl->setUsed(I);
98 return true;
99}
100
101SeedContainer::SeedContainer(ScalarEvolution &SE) : SE(SE) {
102 const VectorizeOptions &Opts =
103 SE.getContext().getLibraryOptions().get<VectorizeOptions>();
104 BundleSizeLimit = Opts.sbvec_seed_bundle_size_limit;
105}
106
107template <typename LoadOrStoreT>
108void SeedContainer::insert(LoadOrStoreT *LSI, bool AllowDiffTypes) {
109 // Find the bundle containing seeds for this symbol and type-of-access.
110 auto &BundleVec = Bundles[getKey(LSI, AllowDiffTypes)];
111 // Fill this vector of bundles front to back so that only the last bundle in
112 // the vector may have available space. This avoids iteration to find one with
113 // space.
114 if (BundleVec.empty() || BundleVec.back()->size() == BundleSizeLimit)
115 BundleVec.emplace_back(std::make_unique<MemSeedBundle<LoadOrStoreT>>(LSI));
116 else
117 BundleVec.back()->tryInsert(LSI, SE);
118
119 SeedLookupMap[LSI] = BundleVec.back().get();
120}
121
122// Explicit instantiations
123template LLVM_EXPORT_TEMPLATE void SeedContainer::insert<LoadInst>(LoadInst *,
124 bool);
125template LLVM_EXPORT_TEMPLATE void SeedContainer::insert<StoreInst>(StoreInst *,
126 bool);
127
128#ifndef NDEBUG
129void SeedContainer::print(raw_ostream &OS) const {
130 for (const auto &Pair : Bundles) {
131 auto [I, Ty, Opc] = Pair.first;
132 const auto &SeedsVec = Pair.second;
133 std::string RefType = dyn_cast<LoadInst>(I) ? "Load"
134 : dyn_cast<StoreInst>(I) ? "Store"
135 : "Other";
136 OS << "[Inst=" << *I << " Ty=" << Ty << " " << RefType << "]\n";
137 for (const auto &SeedPtr : SeedsVec) {
138 SeedPtr->dump(OS);
139 OS << "\n";
140 }
141 }
142 OS << "\n";
143}
144
145LLVM_DUMP_METHOD void SeedContainer::dump() const { print(dbgs()); }
146#endif // NDEBUG
147
148template <typename LoadOrStoreT> static bool isValidMemSeed(LoadOrStoreT *LSI) {
149 if (!LSI->isSimple())
150 return false;
151 auto *Ty = Utils::getExpectedType(V: LSI);
152 // Omit types that are architecturally unvectorizable
153 if (Ty->isX86_FP80Ty() || Ty->isPPC_FP128Ty())
154 return false;
155 // Omit vector types without compile-time-known lane counts
156 if (isa<ScalableVectorType>(Ty))
157 return false;
158 if (auto *VTy = dyn_cast<FixedVectorType>(Ty))
159 return VectorType::isValidElementType(ElemTy: VTy->getElementType());
160 return VectorType::isValidElementType(ElemTy: Ty);
161}
162
163template bool isValidMemSeed<LoadInst>(LoadInst *LSI);
164template bool isValidMemSeed<StoreInst>(StoreInst *LSI);
165
166SeedCollector::SeedCollector(BasicBlock *BB, ScalarEvolution &SE,
167 bool CollectStores, bool CollectLoads,
168 bool AllowDiffTypes)
169 : StoreSeeds(SE), LoadSeeds(SE), Ctx(BB->getContext()) {
170
171 if (!CollectStores && !CollectLoads)
172 return;
173
174 EraseCallbackID = Ctx.registerEraseInstrCallback(CB: [this](Instruction *I) {
175 if (auto SI = dyn_cast<StoreInst>(Val: I))
176 StoreSeeds.erase(I: SI);
177 else if (auto LI = dyn_cast<LoadInst>(Val: I))
178 LoadSeeds.erase(I: LI);
179 });
180
181 // Actually collect the seeds.
182 const VectorizeOptions &Opts =
183 Ctx.getLibraryOptions().get<VectorizeOptions>();
184 for (auto &I : *BB) {
185 if (StoreInst *SI = dyn_cast<StoreInst>(Val: &I))
186 if (CollectStores && isValidMemSeed(LSI: SI))
187 StoreSeeds.insert(LSI: SI, AllowDiffTypes);
188 if (LoadInst *LI = dyn_cast<LoadInst>(Val: &I))
189 if (CollectLoads && isValidMemSeed(LSI: LI))
190 LoadSeeds.insert(LSI: LI, AllowDiffTypes);
191 // Cap compilation time.
192 if (totalNumSeedGroups() > Opts.sbvec_seed_groups_limit)
193 break;
194 }
195}
196
197SeedCollector::~SeedCollector() {
198 Ctx.unregisterEraseInstrCallback(ID: EraseCallbackID);
199}
200
201#ifndef NDEBUG
202void SeedCollector::print(raw_ostream &OS) const {
203 OS << "=== StoreSeeds ===\n";
204 StoreSeeds.print(OS);
205 OS << "=== LoadSeeds ===\n";
206 LoadSeeds.print(OS);
207}
208
209void SeedCollector::dump() const { print(dbgs()); }
210#endif
211
212} // namespace llvm::sandboxir
213