Skip to content

Commit bc9d06c

Browse files
vporpovporpo
committed
[SandboxVec][SeedCollector] Implement collection of seeds with different types
Up until now the seed collector could only collect seeds with the same element type. For example, `i32` and <2 x i32>`. This patch implements the collection of seeds with different types, like `i32` and `i8`.
1 parent 171aa34 commit bc9d06c

File tree

3 files changed

+73
-23
lines changed

3 files changed

+73
-23
lines changed

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

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -191,7 +191,8 @@ class SeedContainer {
191191

192192
ScalarEvolution &SE;
193193

194-
template <typename LoadOrStoreT> KeyT getKey(LoadOrStoreT *LSI) const;
194+
template <typename LoadOrStoreT>
195+
KeyT getKey(LoadOrStoreT *LSI, bool AllowDiffTypes) const;
195196

196197
public:
197198
SeedContainer(ScalarEvolution &SE) : SE(SE) {}
@@ -267,7 +268,8 @@ class SeedContainer {
267268
bool operator!=(const iterator &Other) const { return !(*this == Other); }
268269
};
269270
using const_iterator = BundleMapT::const_iterator;
270-
template <typename LoadOrStoreT> void insert(LoadOrStoreT *LSI);
271+
template <typename LoadOrStoreT>
272+
void insert(LoadOrStoreT *LSI, bool AllowDiffTypes);
271273
// To support constant-time erase, these just mark the element used, rather
272274
// than actually removing them from the bundle.
273275
LLVM_ABI bool erase(Instruction *I);
@@ -291,9 +293,9 @@ class SeedContainer {
291293

292294
// Explicit instantiations
293295
extern template LLVM_TEMPLATE_ABI void
294-
SeedContainer::insert<LoadInst>(LoadInst *);
296+
SeedContainer::insert<LoadInst>(LoadInst *, bool);
295297
extern template LLVM_TEMPLATE_ABI void
296-
SeedContainer::insert<StoreInst>(StoreInst *);
298+
SeedContainer::insert<StoreInst>(StoreInst *, bool);
297299

298300
class SeedCollector {
299301
SeedContainer StoreSeeds;
@@ -308,7 +310,8 @@ class SeedCollector {
308310

309311
public:
310312
LLVM_ABI SeedCollector(BasicBlock *BB, ScalarEvolution &SE,
311-
bool CollectStores, bool CollectLoads);
313+
bool CollectStores, bool CollectLoads,
314+
bool AllowDiffTypes = false);
312315
LLVM_ABI ~SeedCollector();
313316

314317
iterator_range<SeedContainer::iterator> getStoreSeeds() {

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

Lines changed: 22 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -73,22 +73,28 @@ ArrayRef<Instruction *> SeedBundle::getSlice(unsigned StartIdx,
7373
}
7474

7575
template <typename LoadOrStoreT>
76-
SeedContainer::KeyT SeedContainer::getKey(LoadOrStoreT *LSI) const {
76+
SeedContainer::KeyT SeedContainer::getKey(LoadOrStoreT *LSI,
77+
bool AllowDiffTypes) const {
7778
assert((isa<LoadInst>(LSI) || isa<StoreInst>(LSI)) &&
7879
"Expected Load or Store!");
7980
Value *Ptr = Utils::getMemInstructionBase(LSI);
8081
Instruction::Opcode Op = LSI->getOpcode();
81-
Type *Ty = Utils::getExpectedType(LSI);
82-
if (auto *VTy = dyn_cast<VectorType>(Ty))
83-
Ty = VTy->getElementType();
82+
Type *Ty;
83+
if (AllowDiffTypes) {
84+
Ty = nullptr;
85+
} else {
86+
Ty = Utils::getExpectedType(LSI);
87+
if (auto *VTy = dyn_cast<VectorType>(Ty))
88+
Ty = VTy->getElementType();
89+
}
8490
return {Ptr, Ty, Op};
8591
}
8692

8793
// Explicit instantiations
8894
template SeedContainer::KeyT
89-
SeedContainer::getKey<LoadInst>(LoadInst *LSI) const;
95+
SeedContainer::getKey<LoadInst>(LoadInst *LSI, bool AllowDiffTypes) const;
9096
template SeedContainer::KeyT
91-
SeedContainer::getKey<StoreInst>(StoreInst *LSI) const;
97+
SeedContainer::getKey<StoreInst>(StoreInst *LSI, bool AllowDiffTypes) const;
9298

9399
bool SeedContainer::erase(Instruction *I) {
94100
assert((isa<LoadInst>(I) || isa<StoreInst>(I)) && "Expected Load or Store!");
@@ -100,9 +106,10 @@ bool SeedContainer::erase(Instruction *I) {
100106
return true;
101107
}
102108

103-
template <typename LoadOrStoreT> void SeedContainer::insert(LoadOrStoreT *LSI) {
109+
template <typename LoadOrStoreT>
110+
void SeedContainer::insert(LoadOrStoreT *LSI, bool AllowDiffTypes) {
104111
// Find the bundle containing seeds for this symbol and type-of-access.
105-
auto &BundleVec = Bundles[getKey(LSI)];
112+
auto &BundleVec = Bundles[getKey(LSI, AllowDiffTypes)];
106113
// Fill this vector of bundles front to back so that only the last bundle in
107114
// the vector may have available space. This avoids iteration to find one with
108115
// space.
@@ -115,9 +122,9 @@ template <typename LoadOrStoreT> void SeedContainer::insert(LoadOrStoreT *LSI) {
115122
}
116123

117124
// Explicit instantiations
118-
template LLVM_EXPORT_TEMPLATE void SeedContainer::insert<LoadInst>(LoadInst *);
119-
template LLVM_EXPORT_TEMPLATE void
120-
SeedContainer::insert<StoreInst>(StoreInst *);
125+
template LLVM_EXPORT_TEMPLATE void SeedContainer::insert<LoadInst>(LoadInst *, bool);
126+
template LLVM_EXPORT_TEMPLATE void SeedContainer::insert<StoreInst>(StoreInst *,
127+
bool);
121128

122129
#ifndef NDEBUG
123130
void SeedContainer::print(raw_ostream &OS) const {
@@ -158,7 +165,8 @@ template bool isValidMemSeed<LoadInst>(LoadInst *LSI);
158165
template bool isValidMemSeed<StoreInst>(StoreInst *LSI);
159166

160167
SeedCollector::SeedCollector(BasicBlock *BB, ScalarEvolution &SE,
161-
bool CollectStores, bool CollectLoads)
168+
bool CollectStores, bool CollectLoads,
169+
bool AllowDiffTypes)
162170
: StoreSeeds(SE), LoadSeeds(SE), Ctx(BB->getContext()) {
163171

164172
if (!CollectStores && !CollectLoads)
@@ -175,10 +183,10 @@ SeedCollector::SeedCollector(BasicBlock *BB, ScalarEvolution &SE,
175183
for (auto &I : *BB) {
176184
if (StoreInst *SI = dyn_cast<StoreInst>(&I))
177185
if (CollectStores && isValidMemSeed(SI))
178-
StoreSeeds.insert(SI);
186+
StoreSeeds.insert(SI, AllowDiffTypes);
179187
if (LoadInst *LI = dyn_cast<LoadInst>(&I))
180188
if (CollectLoads && isValidMemSeed(LI))
181-
LoadSeeds.insert(LI);
189+
LoadSeeds.insert(LI, AllowDiffTypes);
182190
// Cap compilation time.
183191
if (totalNumSeedGroups() > SeedGroupsLimit)
184192
break;

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

Lines changed: 43 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -259,10 +259,10 @@ define void @foo(ptr %ptrA, float %val, ptr %ptrB) {
259259
// Check begin() end() when empty.
260260
EXPECT_EQ(SC.begin(), SC.end());
261261

262-
SC.insert(S0);
263-
SC.insert(S1);
264-
SC.insert(S2);
265-
SC.insert(S3);
262+
SC.insert(S0, /*AllowDiffTypes=*/false);
263+
SC.insert(S1, /*AllowDiffTypes=*/false);
264+
SC.insert(S2, /*AllowDiffTypes=*/false);
265+
SC.insert(S3, /*AllowDiffTypes=*/false);
266266
unsigned Cnt = 0;
267267
SmallVector<sandboxir::SeedBundle *> Bndls;
268268
for (auto &SeedBndl : SC) {
@@ -480,6 +480,45 @@ define void @foo(ptr noalias %ptr, float %v, <2 x float> %val) {
480480
ExpectThatElementsAre(SB, {St0, St1, St3});
481481
}
482482

483+
TEST_F(SeedBundleTest, DiffTypes) {
484+
parseIR(C, R"IR(
485+
define void @foo(ptr noalias %ptr, i8 %v, i16 %v16) {
486+
bb:
487+
%ptr0 = getelementptr i8, ptr %ptr, i32 0
488+
%ptr1 = getelementptr i8, ptr %ptr, i32 1
489+
%ptr3 = getelementptr i8, ptr %ptr, i32 3
490+
store i8 %v, ptr %ptr0
491+
store i8 %v, ptr %ptr3
492+
store i16 %v16, ptr %ptr1
493+
ret void
494+
}
495+
)IR");
496+
Function &LLVMF = *M->getFunction("foo");
497+
DominatorTree DT(LLVMF);
498+
TargetLibraryInfoImpl TLII(M->getTargetTriple());
499+
TargetLibraryInfo TLI(TLII);
500+
DataLayout DL(M->getDataLayout());
501+
LoopInfo LI(DT);
502+
AssumptionCache AC(LLVMF);
503+
ScalarEvolution SE(LLVMF, TLI, AC, DT, LI);
504+
505+
sandboxir::Context Ctx(C);
506+
auto &F = *Ctx.createFunction(&LLVMF);
507+
auto BB = F.begin();
508+
auto It = std::next(BB->begin(), 3);
509+
auto *St0 = &*It++;
510+
auto *St3 = &*It++;
511+
auto *St1 = &*It++;
512+
513+
sandboxir::SeedCollector SC(&*BB, SE, /*CollectStores=*/true,
514+
/*CollectLoads=*/false, /*AllowDiffTypes=*/true);
515+
516+
auto StoreSeedsRange = SC.getStoreSeeds();
517+
EXPECT_EQ(range_size(StoreSeedsRange), 1u);
518+
auto &SB = *StoreSeedsRange.begin();
519+
ExpectThatElementsAre(SB, {St0, St1, St3});
520+
}
521+
483522
TEST_F(SeedBundleTest, VectorLoads) {
484523
parseIR(C, R"IR(
485524
define void @foo(ptr noalias %ptr, <2 x float> %val0) {

0 commit comments

Comments
 (0)