[SandboxVec][LoadStoreVec] Add support for constants (#189769)

Up until now the pass would only vectorize load-store pairs. This patch
implements vectorization of constant-store pairs.
This commit is contained in:
vporpo 2026-04-06 11:25:20 -07:00 committed by GitHub
parent e11a31f4c7
commit 8d442bc5b5
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 125 additions and 39 deletions

View File

@ -36,7 +36,7 @@ class LoadStoreVec final : public RegionPass {
Scheduler &Sched);
void tryEraseDeadInstrs(ArrayRef<Instruction *> Stores,
ArrayRef<Instruction *> Loads);
ArrayRef<Value *> Operands);
public:
LoadStoreVec() : RegionPass("load-store-vec") {}

View File

@ -48,7 +48,7 @@ std::optional<Type *> LoadStoreVec::canVectorize(ArrayRef<Instruction *> Bndl,
}
void LoadStoreVec::tryEraseDeadInstrs(ArrayRef<Instruction *> Stores,
ArrayRef<Instruction *> Loads) {
ArrayRef<Value *> Operands) {
SmallPtrSet<Instruction *, 8> DeadCandidates;
for (auto *SI : Stores) {
if (auto *PtrI =
@ -56,11 +56,13 @@ void LoadStoreVec::tryEraseDeadInstrs(ArrayRef<Instruction *> Stores,
DeadCandidates.insert(PtrI);
SI->eraseFromParent();
}
for (auto *LI : Loads) {
if (auto *PtrI =
dyn_cast<Instruction>(cast<LoadInst>(LI)->getPointerOperand()))
DeadCandidates.insert(PtrI);
cast<LoadInst>(LI)->eraseFromParent();
for (auto *Op : Operands) {
if (auto *LI = dyn_cast<LoadInst>(Op)) {
if (auto *PtrI =
dyn_cast<Instruction>(cast<LoadInst>(LI)->getPointerOperand()))
DeadCandidates.insert(PtrI);
cast<LoadInst>(LI)->eraseFromParent();
}
}
for (auto *PtrI : DeadCandidates)
if (!PtrI->hasNUsesOrMore(1))
@ -91,46 +93,76 @@ bool LoadStoreVec::runOnRegion(Region &Rgn, const Analyses &A) {
// TODO: For now we only support load operands.
// TODO: For now we don't cross BBs.
// TODO: For now don't vectorize if the loads have external uses.
if (!all_of(Operands, [BB](Value *V) {
auto *LI = dyn_cast<LoadInst>(V);
if (LI == nullptr)
return false;
if (LI->getParent() != BB)
return false;
if (LI->hasNUsesOrMore(2))
return false;
return true;
}))
return false;
// TODO: Try to avoid the extra copy to an instruction vector.
SmallVector<Instruction *, 8> Loads;
Loads.reserve(Operands.size());
for (Value *Op : Operands)
Loads.push_back(cast<Instruction>(Op));
bool Consecutive = VecUtils::areConsecutive<LoadInst, Instruction>(
Loads, A.getScalarEvolution(), *DL);
if (!Consecutive)
return false;
if (!canVectorize(Loads, Sched))
bool AllLoads = all_of(Operands, [BB](Value *V) {
auto *LI = dyn_cast<LoadInst>(V);
if (LI == nullptr)
return false;
// TODO: For now we don't cross BBs.
if (LI->getParent() != BB)
return false;
if (LI->hasNUsesOrMore(2))
return false;
return true;
});
bool AllConstants =
all_of(Operands, [](Value *V) { return isa<Constant>(V); });
if (!AllLoads && !AllConstants)
return false;
// Generate vector store and vector load
Type *Ty = VecUtils::getCombinedVectorTypeFor(Bndl, *DL);
Value *LdPtr = cast<LoadInst>(Loads[0])->getPointerOperand();
// TODO: Compute alignment.
Align LdAlign(1);
auto LdWhereIt = std::next(VecUtils::getLowest(Loads)->getIterator());
auto *VecLd =
LoadInst::create(Ty, LdPtr, LdAlign, LdWhereIt, Ctx, "VecIinitL");
Value *VecOp = nullptr;
if (AllLoads) {
// TODO: Try to avoid the extra copy to an instruction vector.
SmallVector<Instruction *, 8> Loads;
Loads.reserve(Operands.size());
for (Value *Op : Operands)
Loads.push_back(cast<Instruction>(Op));
bool Consecutive = VecUtils::areConsecutive<LoadInst, Instruction>(
Loads, A.getScalarEvolution(), *DL);
if (!Consecutive)
return false;
if (!canVectorize(Loads, Sched))
return false;
// Generate vector load.
Type *Ty = VecUtils::getCombinedVectorTypeFor(Bndl, *DL);
Value *LdPtr = cast<LoadInst>(Loads[0])->getPointerOperand();
// TODO: Compute alignment.
Align LdAlign(1);
auto LdWhereIt = std::next(VecUtils::getLowest(Loads)->getIterator());
VecOp = LoadInst::create(Ty, LdPtr, LdAlign, LdWhereIt, Ctx, "VecIinitL");
} else if (AllConstants) {
SmallVector<Constant *, 8> Constants;
Constants.reserve(Operands.size());
for (Value *Op : Operands) {
auto *COp = cast<Constant>(Op);
if (auto *AggrCOp = dyn_cast<ConstantAggregate>(COp)) {
// If the operand is a constant aggregate, then append all its elements.
for (Value *Elm : AggrCOp->operands())
Constants.push_back(cast<Constant>(Elm));
} else if (auto *SeqCOp = dyn_cast<ConstantDataSequential>(COp)) {
for (auto ElmIdx : seq<unsigned>(SeqCOp->getNumElements()))
Constants.push_back(SeqCOp->getElementAsConstant(ElmIdx));
} else if (auto *Zero = dyn_cast<ConstantAggregateZero>(COp)) {
auto *ZeroElm = Zero->getSequentialElement();
for ([[maybe_unused]] auto Cnt :
seq<unsigned>(Zero->getElementCount().getFixedValue()))
Constants.push_back(ZeroElm);
} else {
Constants.push_back(COp);
}
}
VecOp = ConstantVector::get(Constants);
}
// Generate vector store.
Value *StPtr = cast<StoreInst>(Bndl[0])->getPointerOperand();
// TODO: Compute alignment.
Align StAlign(1);
auto StWhereIt = std::next(VecUtils::getLowest(Bndl)->getIterator());
StoreInst::create(VecLd, StPtr, StAlign, StWhereIt, Ctx);
StoreInst::create(VecOp, StPtr, StAlign, StWhereIt, Ctx);
tryEraseDeadInstrs(Bndl, Loads);
tryEraseDeadInstrs(Bndl, Operands);
return true;
}

View File

@ -301,6 +301,57 @@ define void @load_store_vec_external_uses(ptr %ptr) {
ret void
}
; Store-constant pattern.
define void @load_store_vec_constants(ptr %ptr) {
; CHECK-LABEL: define void @load_store_vec_constants(
; CHECK-SAME: ptr [[PTR:%.*]]) {
; CHECK-NEXT: [[PTR0:%.*]] = getelementptr i8, ptr [[PTR]], i32 0
; CHECK-NEXT: store <3 x i8> <i8 42, i8 43, i8 44>, ptr [[PTR0]], align 1, !sandboxvec [[META15:![0-9]+]]
; CHECK-NEXT: ret void
;
%ptr0 = getelementptr i8, ptr %ptr, i32 0
%ptr1 = getelementptr i8, ptr %ptr, i32 1
%ptr2 = getelementptr i8, ptr %ptr, i32 3
store i8 42, ptr %ptr0
store i16 43, ptr %ptr1
store i8 44, ptr %ptr2
ret void
}
; Same but with ConstantDataSequential.
define void @load_store_vec_constants_CDS(ptr %ptr) {
; CHECK-LABEL: define void @load_store_vec_constants_CDS(
; CHECK-SAME: ptr [[PTR:%.*]]) {
; CHECK-NEXT: [[PTR0:%.*]] = getelementptr i8, ptr [[PTR]], i32 0
; CHECK-NEXT: store <4 x i8> <i8 0, i8 1, i8 2, i8 3>, ptr [[PTR0]], align 1, !sandboxvec [[META16:![0-9]+]]
; CHECK-NEXT: ret void
;
%ptr0 = getelementptr i8, ptr %ptr, i32 0
%ptr1 = getelementptr i8, ptr %ptr, i32 1
%ptr2 = getelementptr i8, ptr %ptr, i32 3
store i8 0, ptr %ptr0
store <2 x i8> <i8 1, i8 2>, ptr %ptr1
store i8 3, ptr %ptr2
ret void
}
; Same but with floats
define void @load_store_vec_constants_CDS_float(ptr %ptr) {
; CHECK-LABEL: define void @load_store_vec_constants_CDS_float(
; CHECK-SAME: ptr [[PTR:%.*]]) {
; CHECK-NEXT: [[PTR0:%.*]] = getelementptr float, ptr [[PTR]], i32 0
; CHECK-NEXT: store <8 x float> <float 1.000000e+00, float 0.000000e+00, float 0.000000e+00, float 0.000000e+00, float 0.000000e+00, float 0.000000e+00, float 0.000000e+00, float 0.000000e+00>, ptr [[PTR0]], align 1, !sandboxvec [[META17:![0-9]+]]
; CHECK-NEXT: ret void
;
%ptr0 = getelementptr float, ptr %ptr, i32 0
%ptr1 = getelementptr float, ptr %ptr, i32 1
%ptr2 = getelementptr float, ptr %ptr, i32 3
store float 1.0, ptr %ptr0
store <2 x float> zeroinitializer, ptr %ptr1
store <5 x float> zeroinitializer, ptr %ptr2
ret void
}
;.
; CHECK: [[META0]] = distinct !{!"sandboxregion"}
; CHECK: [[META1]] = distinct !{!"sandboxregion"}
@ -317,4 +368,7 @@ define void @load_store_vec_external_uses(ptr %ptr) {
; CHECK: [[META12]] = distinct !{!"sandboxregion"}
; CHECK: [[META13]] = distinct !{!"sandboxregion"}
; CHECK: [[META14]] = distinct !{!"sandboxregion"}
; CHECK: [[META15]] = distinct !{!"sandboxregion"}
; CHECK: [[META16]] = distinct !{!"sandboxregion"}
; CHECK: [[META17]] = distinct !{!"sandboxregion"}
;.