Skip to content

[SandboxVec][SeedCollector] Implement collection of seeds with different types #146171

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Aug 23, 2025
Merged

[SandboxVec][SeedCollector] Implement collection of seeds with different types #146171

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 8 additions & 5 deletions llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -191,7 +191,8 @@ class SeedContainer {

ScalarEvolution &SE;

template <typename LoadOrStoreT> KeyT getKey(LoadOrStoreT *LSI) const;
template <typename LoadOrStoreT>
KeyT getKey(LoadOrStoreT *LSI, bool AllowDiffTypes) const;
Copy link
Collaborator

@slackito slackito Aug 11, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What's the motivation for having AllowDiffTypes as a configurable thing? Wouldn't we want it to be always on to enable more vectorization opportunities?

Copy link
Contributor Author

@vporpo vporpo Aug 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Allowing different types is only beneficial for very specific (and limited) vectorization patterns, like load-store pairs of different types. In the general case we can't really vectorize instructions of different types unless we legally convert them into the same type. So allowing different types is not always desired, which is why we have this as an option.

Copy link
Collaborator

@slackito slackito Aug 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the explanation!


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

// Explicit instantiations
extern template LLVM_TEMPLATE_ABI void
SeedContainer::insert<LoadInst>(LoadInst *);
SeedContainer::insert<LoadInst>(LoadInst *, bool);
extern template LLVM_TEMPLATE_ABI void
SeedContainer::insert<StoreInst>(StoreInst *);
SeedContainer::insert<StoreInst>(StoreInst *, bool);

class SeedCollector {
SeedContainer StoreSeeds;
Expand All @@ -308,7 +310,8 @@ class SeedCollector {

public:
LLVM_ABI SeedCollector(BasicBlock *BB, ScalarEvolution &SE,
bool CollectStores, bool CollectLoads);
bool CollectStores, bool CollectLoads,
bool AllowDiffTypes = false);
LLVM_ABI ~SeedCollector();

iterator_range<SeedContainer::iterator> getStoreSeeds() {
Expand Down
37 changes: 23 additions & 14 deletions llvm/lib/Transforms/Vectorize/SandboxVectorizer/SeedCollector.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -73,22 +73,28 @@ ArrayRef<Instruction *> SeedBundle::getSlice(unsigned StartIdx,
}

template <typename LoadOrStoreT>
SeedContainer::KeyT SeedContainer::getKey(LoadOrStoreT *LSI) const {
SeedContainer::KeyT SeedContainer::getKey(LoadOrStoreT *LSI,
bool AllowDiffTypes) const {
assert((isa<LoadInst>(LSI) || isa<StoreInst>(LSI)) &&
"Expected Load or Store!");
Value *Ptr = Utils::getMemInstructionBase(LSI);
Instruction::Opcode Op = LSI->getOpcode();
Type *Ty = Utils::getExpectedType(LSI);
if (auto *VTy = dyn_cast<VectorType>(Ty))
Ty = VTy->getElementType();
Type *Ty;
if (AllowDiffTypes) {
Ty = nullptr;
} else {
Ty = Utils::getExpectedType(LSI);
if (auto *VTy = dyn_cast<VectorType>(Ty))
Ty = VTy->getElementType();
}
return {Ptr, Ty, Op};
}

// Explicit instantiations
template SeedContainer::KeyT
SeedContainer::getKey<LoadInst>(LoadInst *LSI) const;
SeedContainer::getKey<LoadInst>(LoadInst *LSI, bool AllowDiffTypes) const;
template SeedContainer::KeyT
SeedContainer::getKey<StoreInst>(StoreInst *LSI) const;
SeedContainer::getKey<StoreInst>(StoreInst *LSI, bool AllowDiffTypes) const;

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

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

// Explicit instantiations
template LLVM_EXPORT_TEMPLATE void SeedContainer::insert<LoadInst>(LoadInst *);
template LLVM_EXPORT_TEMPLATE void
SeedContainer::insert<StoreInst>(StoreInst *);
template LLVM_EXPORT_TEMPLATE void SeedContainer::insert<LoadInst>(LoadInst *,
bool);
template LLVM_EXPORT_TEMPLATE void SeedContainer::insert<StoreInst>(StoreInst *,
bool);

#ifndef NDEBUG
void SeedContainer::print(raw_ostream &OS) const {
Expand Down Expand Up @@ -158,7 +166,8 @@ template bool isValidMemSeed<LoadInst>(LoadInst *LSI);
template bool isValidMemSeed<StoreInst>(StoreInst *LSI);

SeedCollector::SeedCollector(BasicBlock *BB, ScalarEvolution &SE,
bool CollectStores, bool CollectLoads)
bool CollectStores, bool CollectLoads,
bool AllowDiffTypes)
: StoreSeeds(SE), LoadSeeds(SE), Ctx(BB->getContext()) {

if (!CollectStores && !CollectLoads)
Expand All @@ -175,10 +184,10 @@ SeedCollector::SeedCollector(BasicBlock *BB, ScalarEvolution &SE,
for (auto &I : *BB) {
if (StoreInst *SI = dyn_cast<StoreInst>(&I))
if (CollectStores && isValidMemSeed(SI))
StoreSeeds.insert(SI);
StoreSeeds.insert(SI, AllowDiffTypes);
if (LoadInst *LI = dyn_cast<LoadInst>(&I))
if (CollectLoads && isValidMemSeed(LI))
LoadSeeds.insert(LI);
LoadSeeds.insert(LI, AllowDiffTypes);
// Cap compilation time.
if (totalNumSeedGroups() > SeedGroupsLimit)
break;
Expand Down
47 changes: 43 additions & 4 deletions llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SeedCollectorTest.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -259,10 +259,10 @@ define void @foo(ptr %ptrA, float %val, ptr %ptrB) {
// Check begin() end() when empty.
EXPECT_EQ(SC.begin(), SC.end());

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

TEST_F(SeedBundleTest, DiffTypes) {
parseIR(C, R"IR(
define void @foo(ptr noalias %ptr, i8 %v, i16 %v16) {
bb:
%ptr0 = getelementptr i8, ptr %ptr, i32 0
%ptr1 = getelementptr i8, ptr %ptr, i32 1
%ptr3 = getelementptr i8, ptr %ptr, i32 3
store i8 %v, ptr %ptr0
store i8 %v, ptr %ptr3
store i16 %v16, ptr %ptr1
ret void
}
)IR");
Function &LLVMF = *M->getFunction("foo");
DominatorTree DT(LLVMF);
TargetLibraryInfoImpl TLII(M->getTargetTriple());
TargetLibraryInfo TLI(TLII);
DataLayout DL(M->getDataLayout());
LoopInfo LI(DT);
AssumptionCache AC(LLVMF);
ScalarEvolution SE(LLVMF, TLI, AC, DT, LI);

sandboxir::Context Ctx(C);
auto &F = *Ctx.createFunction(&LLVMF);
auto BB = F.begin();
auto It = std::next(BB->begin(), 3);
auto *St0 = &*It++;
auto *St3 = &*It++;
auto *St1 = &*It++;

sandboxir::SeedCollector SC(&*BB, SE, /*CollectStores=*/true,
/*CollectLoads=*/false, /*AllowDiffTypes=*/true);

auto StoreSeedsRange = SC.getStoreSeeds();
EXPECT_EQ(range_size(StoreSeedsRange), 1u);
auto &SB = *StoreSeedsRange.begin();
ExpectThatElementsAre(SB, {St0, St1, St3});
}

TEST_F(SeedBundleTest, VectorLoads) {
parseIR(C, R"IR(
define void @foo(ptr noalias %ptr, <2 x float> %val0) {
Expand Down
Loading