Skip to content

Commit d91318b

Browse files
Sterling-AugustineSterling-Augustine
authored
[SandboxVectorizer] New class to actually collect and manage seeds (#112979)
There are many more tests to add, but I would like to get this reviewed and the details sorted out before it grows too big.
1 parent 3acc58c commit d91318b

File tree

4 files changed

+292
-0
lines changed

4 files changed

+292
-0
lines changed

‎llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -284,6 +284,33 @@ class SeedContainer {
284284
#endif // NDEBUG
285285
};
286286

287+
class SeedCollector {
288+
SeedContainer StoreSeeds;
289+
SeedContainer LoadSeeds;
290+
Context &Ctx;
291+
292+
/// \Returns the number of SeedBundle groups for all seed types.
293+
/// This is to be used for limiting compilation time.
294+
unsigned totalNumSeedGroups() const {
295+
return StoreSeeds.size() + LoadSeeds.size();
296+
}
297+
298+
public:
299+
SeedCollector(BasicBlock *BB, ScalarEvolution &SE);
300+
~SeedCollector();
301+
302+
iterator_range<SeedContainer::iterator> getStoreSeeds() {
303+
return {StoreSeeds.begin(), StoreSeeds.end()};
304+
}
305+
iterator_range<SeedContainer::iterator> getLoadSeeds() {
306+
return {LoadSeeds.begin(), LoadSeeds.end()};
307+
}
308+
#ifndef NDEBUG
309+
void print(raw_ostream &OS) const;
310+
LLVM_DUMP_METHOD void dump() const;
311+
#endif
312+
};
313+
287314
} // namespace llvm::sandboxir
288315

289316
#endif // LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_SEEDCOLLECTOR_H

‎llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/VecUtils.h‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
//===- VecUtils.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+
// Collector for SandboxVectorizer related convenience functions that don't
10+
// belong in other classes.
11+
12+
#ifndef LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_VECUTILS_H
13+
#define LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_VECUTILS_H
14+
15+
class Utils {
16+
public:
17+
/// \Returns the number of elements in \p Ty. That is the number of lanes if a
18+
/// fixed vector or 1 if scalar. ScalableVectors have unknown size and
19+
/// therefore are unsupported.
20+
static int getNumElements(Type *Ty) {
21+
assert(!isa<ScalableVectorType>(Ty));
22+
return Ty->isVectorTy() ? cast<FixedVectorType>(Ty)->getNumElements() : 1;
23+
}
24+
/// Returns \p Ty if scalar or its element type if vector.
25+
static Type *getElementType(Type *Ty) {
26+
return Ty->isVectorTy() ? cast<FixedVectorType>(Ty)->getElementType() : Ty;
27+
}
28+
}
29+
30+
#endif LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_VECUTILS_H

‎llvm/lib/Transforms/Vectorize/SandboxVectorizer/SeedCollector.cpp‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,16 @@ namespace llvm::sandboxir {
2222
cl::opt<unsigned> SeedBundleSizeLimit(
2323
"sbvec-seed-bundle-size-limit", cl::init(32), cl::Hidden,
2424
cl::desc("Limit the size of the seed bundle to cap compilation time."));
25+
#define LoadSeedsDef "loads"
26+
#define StoreSeedsDef "stores"
27+
cl::opt<std::string> CollectSeeds(
28+
"sbvec-collect-seeds", cl::init(LoadSeedsDef "," StoreSeedsDef), cl::Hidden,
29+
cl::desc("Collect these seeds. Use empty for none or a comma-separated "
30+
"list of '" LoadSeedsDef "' and '" StoreSeedsDef "'."));
31+
cl::opt<unsigned> SeedGroupsLimit(
32+
"sbvec-seed-groups-limit", cl::init(256), cl::Hidden,
33+
cl::desc("Limit the number of collected seeds groups in a BB to "
34+
"cap compilation time."));
2535

2636
MutableArrayRef<Instruction *> SeedBundle::getSlice(unsigned StartIdx,
2737
unsigned MaxVecRegBits,
@@ -131,4 +141,61 @@ void SeedContainer::print(raw_ostream &OS) const {
131141
LLVM_DUMP_METHOD void SeedContainer::dump() const { print(dbgs()); }
132142
#endif // NDEBUG
133143

144+
template <typename LoadOrStoreT> static bool isValidMemSeed(LoadOrStoreT *LSI) {
145+
if (LSI->isSimple())
146+
return true;
147+
auto *Ty = Utils::getExpectedType(LSI);
148+
// Omit types that are architecturally unvectorizable
149+
if (Ty->isX86_FP80Ty() || Ty->isPPC_FP128Ty())
150+
return false;
151+
// Omit vector types without compile-time-known lane counts
152+
if (isa<ScalableVectorType>(Ty))
153+
return false;
154+
if (auto *VTy = dyn_cast<FixedVectorType>(Ty))
155+
return VectorType::isValidElementType(VTy->getElementType());
156+
return VectorType::isValidElementType(Ty);
157+
}
158+
159+
template bool isValidMemSeed<LoadInst>(LoadInst *LSI);
160+
template bool isValidMemSeed<StoreInst>(StoreInst *LSI);
161+
162+
SeedCollector::SeedCollector(BasicBlock *BB, ScalarEvolution &SE)
163+
: StoreSeeds(SE), LoadSeeds(SE), Ctx(BB->getContext()) {
164+
// TODO: Register a callback for updating the Collector data structures upon
165+
// instr removal
166+
167+
bool CollectStores = CollectSeeds.find(StoreSeedsDef) != std::string::npos;
168+
bool CollectLoads = CollectSeeds.find(LoadSeedsDef) != std::string::npos;
169+
if (!CollectStores && !CollectLoads)
170+
return;
171+
// Actually collect the seeds.
172+
for (auto &I : *BB) {
173+
if (StoreInst *SI = dyn_cast<StoreInst>(&I))
174+
if (CollectStores && isValidMemSeed(SI))
175+
StoreSeeds.insert(SI);
176+
if (LoadInst *LI = dyn_cast<LoadInst>(&I))
177+
if (CollectLoads && isValidMemSeed(LI))
178+
LoadSeeds.insert(LI);
179+
// Cap compilation time.
180+
if (totalNumSeedGroups() > SeedGroupsLimit)
181+
break;
182+
}
183+
}
184+
185+
SeedCollector::~SeedCollector() {
186+
// TODO: Unregister the callback for updating the seed datastructures upon
187+
// instr removal
188+
}
189+
190+
#ifndef NDEBUG
191+
void SeedCollector::print(raw_ostream &OS) const {
192+
OS << "=== StoreSeeds ===\n";
193+
StoreSeeds.print(OS);
194+
OS << "=== LoadSeeds ===\n";
195+
LoadSeeds.print(OS);
196+
}
197+
198+
void SeedCollector::dump() const { print(dbgs()); }
199+
#endif
200+
134201
} // namespace llvm::sandboxir

‎llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SeedCollectorTest.cpp‎

Lines changed: 168 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -268,3 +268,171 @@ define void @foo(ptr %ptrA, float %val, ptr %ptrB) {
268268
}
269269
EXPECT_EQ(Cnt, 0u);
270270
}
271+
272+
TEST_F(SeedBundleTest, ConsecutiveStores) {
273+
// Where "Consecutive" means the stores address consecutive locations in
274+
// memory, but not in program order. Check to see that the collector puts them
275+
// in the proper order for vectorization.
276+
parseIR(C, R"IR(
277+
define void @foo(ptr noalias %ptr, float %val) {
278+
bb:
279+
%ptr0 = getelementptr float, ptr %ptr, i32 0
280+
%ptr1 = getelementptr float, ptr %ptr, i32 1
281+
%ptr2 = getelementptr float, ptr %ptr, i32 2
282+
%ptr3 = getelementptr float, ptr %ptr, i32 3
283+
store float %val, ptr %ptr0
284+
store float %val, ptr %ptr2
285+
store float %val, ptr %ptr1
286+
store float %val, ptr %ptr3
287+
ret void
288+
}
289+
)IR");
290+
Function &LLVMF = *M->getFunction("foo");
291+
DominatorTree DT(LLVMF);
292+
TargetLibraryInfoImpl TLII;
293+
TargetLibraryInfo TLI(TLII);
294+
DataLayout DL(M->getDataLayout());
295+
LoopInfo LI(DT);
296+
AssumptionCache AC(LLVMF);
297+
ScalarEvolution SE(LLVMF, TLI, AC, DT, LI);
298+
299+
sandboxir::Context Ctx(C);
300+
auto &F = *Ctx.createFunction(&LLVMF);
301+
auto BB = F.begin();
302+
sandboxir::SeedCollector SC(&*BB, SE);
303+
304+
// Find the stores
305+
auto It = std::next(BB->begin(), 4);
306+
// StX with X as the order by offset in memory
307+
auto *St0 = &*It++;
308+
auto *St2 = &*It++;
309+
auto *St1 = &*It++;
310+
auto *St3 = &*It++;
311+
312+
auto StoreSeedsRange = SC.getStoreSeeds();
313+
auto &SB = *StoreSeedsRange.begin();
314+
// Expect just one vector of store seeds
315+
EXPECT_EQ(range_size(StoreSeedsRange), 1u);
316+
EXPECT_THAT(SB, testing::ElementsAre(St0, St1, St2, St3));
317+
}
318+
319+
TEST_F(SeedBundleTest, StoresWithGaps) {
320+
parseIR(C, R"IR(
321+
define void @foo(ptr noalias %ptr, float %val) {
322+
bb:
323+
%ptr0 = getelementptr float, ptr %ptr, i32 0
324+
%ptr1 = getelementptr float, ptr %ptr, i32 3
325+
%ptr2 = getelementptr float, ptr %ptr, i32 5
326+
%ptr3 = getelementptr float, ptr %ptr, i32 7
327+
store float %val, ptr %ptr0
328+
store float %val, ptr %ptr2
329+
store float %val, ptr %ptr1
330+
store float %val, ptr %ptr3
331+
ret void
332+
}
333+
)IR");
334+
Function &LLVMF = *M->getFunction("foo");
335+
DominatorTree DT(LLVMF);
336+
TargetLibraryInfoImpl TLII;
337+
TargetLibraryInfo TLI(TLII);
338+
DataLayout DL(M->getDataLayout());
339+
LoopInfo LI(DT);
340+
AssumptionCache AC(LLVMF);
341+
ScalarEvolution SE(LLVMF, TLI, AC, DT, LI);
342+
343+
sandboxir::Context Ctx(C);
344+
auto &F = *Ctx.createFunction(&LLVMF);
345+
auto BB = F.begin();
346+
sandboxir::SeedCollector SC(&*BB, SE);
347+
348+
// Find the stores
349+
auto It = std::next(BB->begin(), 4);
350+
// StX with X as the order by offset in memory
351+
auto *St0 = &*It++;
352+
auto *St2 = &*It++;
353+
auto *St1 = &*It++;
354+
auto *St3 = &*It++;
355+
356+
auto StoreSeedsRange = SC.getStoreSeeds();
357+
auto &SB = *StoreSeedsRange.begin();
358+
// Expect just one vector of store seeds
359+
EXPECT_EQ(range_size(StoreSeedsRange), 1u);
360+
EXPECT_THAT(SB, testing::ElementsAre(St0, St1, St2, St3));
361+
}
362+
363+
TEST_F(SeedBundleTest, VectorStores) {
364+
parseIR(C, R"IR(
365+
define void @foo(ptr noalias %ptr, <2 x float> %val) {
366+
bb:
367+
%ptr0 = getelementptr float, ptr %ptr, i32 0
368+
%ptr1 = getelementptr float, ptr %ptr, i32 1
369+
store <2 x float> %val, ptr %ptr1
370+
store <2 x float> %val, ptr %ptr0
371+
ret void
372+
}
373+
)IR");
374+
Function &LLVMF = *M->getFunction("foo");
375+
DominatorTree DT(LLVMF);
376+
TargetLibraryInfoImpl TLII;
377+
TargetLibraryInfo TLI(TLII);
378+
DataLayout DL(M->getDataLayout());
379+
LoopInfo LI(DT);
380+
AssumptionCache AC(LLVMF);
381+
ScalarEvolution SE(LLVMF, TLI, AC, DT, LI);
382+
383+
sandboxir::Context Ctx(C);
384+
auto &F = *Ctx.createFunction(&LLVMF);
385+
auto BB = F.begin();
386+
sandboxir::SeedCollector SC(&*BB, SE);
387+
388+
// Find the stores
389+
auto It = std::next(BB->begin(), 2);
390+
// StX with X as the order by offset in memory
391+
auto *St1 = &*It++;
392+
auto *St0 = &*It++;
393+
394+
auto StoreSeedsRange = SC.getStoreSeeds();
395+
EXPECT_EQ(range_size(StoreSeedsRange), 1u);
396+
auto &SB = *StoreSeedsRange.begin();
397+
EXPECT_THAT(SB, testing::ElementsAre(St0, St1));
398+
}
399+
400+
TEST_F(SeedBundleTest, MixedScalarVectors) {
401+
parseIR(C, R"IR(
402+
define void @foo(ptr noalias %ptr, float %v, <2 x float> %val) {
403+
bb:
404+
%ptr0 = getelementptr float, ptr %ptr, i32 0
405+
%ptr1 = getelementptr float, ptr %ptr, i32 1
406+
%ptr3 = getelementptr float, ptr %ptr, i32 3
407+
store float %v, ptr %ptr0
408+
store float %v, ptr %ptr3
409+
store <2 x float> %val, ptr %ptr1
410+
ret void
411+
}
412+
)IR");
413+
Function &LLVMF = *M->getFunction("foo");
414+
DominatorTree DT(LLVMF);
415+
TargetLibraryInfoImpl TLII;
416+
TargetLibraryInfo TLI(TLII);
417+
DataLayout DL(M->getDataLayout());
418+
LoopInfo LI(DT);
419+
AssumptionCache AC(LLVMF);
420+
ScalarEvolution SE(LLVMF, TLI, AC, DT, LI);
421+
422+
sandboxir::Context Ctx(C);
423+
auto &F = *Ctx.createFunction(&LLVMF);
424+
auto BB = F.begin();
425+
sandboxir::SeedCollector SC(&*BB, SE);
426+
427+
// Find the stores
428+
auto It = std::next(BB->begin(), 3);
429+
// StX with X as the order by offset in memory
430+
auto *St0 = &*It++;
431+
auto *St3 = &*It++;
432+
auto *St1 = &*It++;
433+
434+
auto StoreSeedsRange = SC.getStoreSeeds();
435+
EXPECT_EQ(range_size(StoreSeedsRange), 1u);
436+
auto &SB = *StoreSeedsRange.begin();
437+
EXPECT_THAT(SB, testing::ElementsAre(St0, St1, St3));
438+
}

0 commit comments

Comments
 (0)