Skip to content

[SLP]Prefer segmented/deinterleaved loads to strided and fix codegen #135058

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 6 commits into from
Apr 22, 2025
Merged

[SLP]Prefer segmented/deinterleaved loads to strided and fix codegen #135058

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
1fabde4
[𝘀𝗽𝗿] initial version
alexey-bataev Apr 9, 2025
780d8a2
Fix formatting
alexey-bataev Apr 9, 2025
c28eeac
Rebase
alexey-bataev Apr 10, 2025
f01c105
Rebase
alexey-bataev Apr 16, 2025
bc22a15
Rebase, drop masked deinterleaved loads as not ready
alexey-bataev Apr 21, 2025
a5f65e5
Rebase
alexey-bataev Apr 22, 2025
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
65 changes: 36 additions & 29 deletions llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5923,9 +5923,9 @@ static bool isMaskedLoadCompress(
// Check for very large distances between elements.
if (*Diff / Sz >= MaxRegSize / 8)
return false;
Align CommonAlignment = computeCommonAlignment<LoadInst>(VL);
LoadVecTy = getWidenedType(ScalarTy, *Diff + 1);
auto *LI = cast<LoadInst>(Order.empty() ? VL.front() : VL[Order.front()]);
Align CommonAlignment = LI->getAlign();
IsMasked = !isSafeToLoadUnconditionally(
Ptr0, LoadVecTy, CommonAlignment, DL,
cast<LoadInst>(Order.empty() ? VL.back() : VL[Order.back()]), &AC, &DT,
Expand Down Expand Up @@ -5964,26 +5964,28 @@ static bool isMaskedLoadCompress(
TTI.getMaskedMemoryOpCost(Instruction::Load, LoadVecTy, CommonAlignment,
LI->getPointerAddressSpace(), CostKind);
} else {
CommonAlignment = LI->getAlign();
LoadCost =
TTI.getMemoryOpCost(Instruction::Load, LoadVecTy, CommonAlignment,
LI->getPointerAddressSpace(), CostKind);
}
if (IsStrided) {
if (IsStrided && !IsMasked) {
// Check for potential segmented(interleaved) loads.
if (TTI.isLegalInterleavedAccessType(LoadVecTy, CompressMask[1],
auto *AlignedLoadVecTy = getWidenedType(
ScalarTy, getFullVectorNumberOfElements(TTI, ScalarTy, *Diff + 1));
if (TTI.isLegalInterleavedAccessType(AlignedLoadVecTy, CompressMask[1],
CommonAlignment,
LI->getPointerAddressSpace())) {
InstructionCost InterleavedCost =
VectorGEPCost + TTI.getInterleavedMemoryOpCost(
Instruction::Load, LoadVecTy, CompressMask[1],
std::nullopt, CommonAlignment,
Instruction::Load, AlignedLoadVecTy,
CompressMask[1], std::nullopt, CommonAlignment,
LI->getPointerAddressSpace(), CostKind, IsMasked);
if (!Mask.empty())
InterleavedCost += ::getShuffleCost(TTI, TTI::SK_PermuteSingleSrc,
VecTy, Mask, CostKind);
if (InterleavedCost < GatherCost) {
InterleaveFactor = CompressMask[1];
LoadVecTy = AlignedLoadVecTy;
return true;
}
}
Expand All @@ -6001,6 +6003,24 @@ static bool isMaskedLoadCompress(
return TotalVecCost < GatherCost;
}

/// Checks if the \p VL can be transformed to a (masked)load + compress or
/// (masked) interleaved load.
static bool
isMaskedLoadCompress(ArrayRef<Value *> VL, ArrayRef<Value *> PointerOps,
ArrayRef<unsigned> Order, const TargetTransformInfo &TTI,
const DataLayout &DL, ScalarEvolution &SE,
AssumptionCache &AC, const DominatorTree &DT,
const TargetLibraryInfo &TLI,
const function_ref<bool(Value *)> AreAllUsersVectorized) {
bool IsMasked;
unsigned InterleaveFactor;
SmallVector<int> CompressMask;
VectorType *LoadVecTy;
return isMaskedLoadCompress(VL, PointerOps, Order, TTI, DL, SE, AC, DT, TLI,
AreAllUsersVectorized, IsMasked, InterleaveFactor,
CompressMask, LoadVecTy);
}

/// Checks if strided loads can be generated out of \p VL loads with pointers \p
/// PointerOps:
/// 1. Target with strided load support is detected.
Expand Down Expand Up @@ -6137,6 +6157,12 @@ BoUpSLP::canVectorizeLoads(ArrayRef<Value *> VL, const Value *VL0,
// Check that the sorted loads are consecutive.
if (static_cast<unsigned>(*Diff) == Sz - 1)
return LoadsState::Vectorize;
if (isMaskedLoadCompress(VL, PointerOps, Order, *TTI, *DL, *SE, *AC, *DT,
*TLI, [&](Value *V) {
return areAllUsersVectorized(
cast<Instruction>(V), UserIgnoreList);
}))
return LoadsState::CompressVectorize;
// Simple check if not a strided access - clear order.
bool IsPossibleStrided = *Diff % (Sz - 1) == 0;
// Try to generate strided load node.
Expand All @@ -6150,18 +6176,6 @@ BoUpSLP::canVectorizeLoads(ArrayRef<Value *> VL, const Value *VL0,
isStridedLoad(VL, PointerOps, Order, *TTI, *DL, *SE,
IsAnyPointerUsedOutGraph, *Diff))
return LoadsState::StridedVectorize;
bool IsMasked;
unsigned InterleaveFactor;
SmallVector<int> CompressMask;
VectorType *LoadVecTy;
if (isMaskedLoadCompress(
VL, PointerOps, Order, *TTI, *DL, *SE, *AC, *DT, *TLI,
[&](Value *V) {
return areAllUsersVectorized(cast<Instruction>(V),
UserIgnoreList);
},
IsMasked, InterleaveFactor, CompressMask, LoadVecTy))
return LoadsState::CompressVectorize;
}
if (!TTI->isLegalMaskedGather(VecTy, CommonAlignment) ||
TTI->forceScalarizeMaskedGather(VecTy, CommonAlignment))
Expand Down Expand Up @@ -13439,11 +13453,7 @@ BoUpSLP::getEntryCost(const TreeEntry *E, ArrayRef<Value *> VectorizedVals,
assert(IsVectorized && "Expected to be vectorized");
CompressEntryToData.try_emplace(E, CompressMask, LoadVecTy,
InterleaveFactor, IsMasked);
Align CommonAlignment;
if (IsMasked)
CommonAlignment = computeCommonAlignment<LoadInst>(VL);
else
CommonAlignment = LI0->getAlign();
Align CommonAlignment = LI0->getAlign();
if (InterleaveFactor) {
VecLdCost = TTI->getInterleavedMemoryOpCost(
Instruction::Load, LoadVecTy, InterleaveFactor, std::nullopt,
Expand Down Expand Up @@ -18049,14 +18059,11 @@ Value *BoUpSLP::vectorizeTree(TreeEntry *E) {
PointerOps[I] = cast<LoadInst>(V)->getPointerOperand();
auto [CompressMask, LoadVecTy, InterleaveFactor, IsMasked] =
CompressEntryToData.at(E);
Align CommonAlignment;
if (IsMasked)
CommonAlignment = computeCommonAlignment<LoadInst>(E->Scalars);
else
CommonAlignment = LI->getAlign();
Align CommonAlignment = LI->getAlign();
Copy link
Contributor

Choose a reason for hiding this comment

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

Curious: Why was computeCommonAlignment() needed before?

Copy link
Member Author

Choose a reason for hiding this comment

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

That was a mistake. It is mostly required for masked gathers, for consecutive loads it is safe to use the alignment for the beginning element

if (IsMasked) {
unsigned VF = getNumElements(LoadVecTy);
SmallVector<Constant *> MaskValues(
getNumElements(LoadVecTy) / getNumElements(LI->getType()),
VF / getNumElements(LI->getType()),
ConstantInt::getFalse(VecTy->getContext()));
for (int I : CompressMask)
MaskValues[I] = ConstantInt::getTrue(VecTy->getContext());
Expand Down
22 changes: 12 additions & 10 deletions llvm/test/Transforms/SLPVectorizer/RISCV/segmented-loads-simple.ll
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,11 +5,12 @@ define i32 @sum_of_abs_stride_2(ptr noalias %a, ptr noalias %b) {
; CHECK-LABEL: define i32 @sum_of_abs_stride_2
; CHECK-SAME: (ptr noalias [[A:%.*]], ptr noalias [[B:%.*]]) #[[ATTR0:[0-9]+]] {
; CHECK-NEXT: entry:
; CHECK-NEXT: [[TMP0:%.*]] = call <8 x i8> @llvm.experimental.vp.strided.load.v8i8.p0.i64(ptr align 1 [[A]], i64 2, <8 x i1> splat (i1 true), i32 8)
; CHECK-NEXT: [[TMP1:%.*]] = call <8 x i8> @llvm.abs.v8i8(<8 x i8> [[TMP0]], i1 false)
; CHECK-NEXT: [[TMP2:%.*]] = sext <8 x i8> [[TMP1]] to <8 x i32>
; CHECK-NEXT: [[TMP3:%.*]] = call i32 @llvm.vector.reduce.add.v8i32(<8 x i32> [[TMP2]])
; CHECK-NEXT: ret i32 [[TMP3]]
; CHECK-NEXT: [[TMP0:%.*]] = call <15 x i8> @llvm.masked.load.v15i8.p0(ptr [[A]], i32 1, <15 x i1> <i1 true, i1 false, i1 true, i1 false, i1 true, i1 false, i1 true, i1 false, i1 true, i1 false, i1 true, i1 false, i1 true, i1 false, i1 true>, <15 x i8> poison)
; CHECK-NEXT: [[TMP1:%.*]] = shufflevector <15 x i8> [[TMP0]], <15 x i8> poison, <8 x i32> <i32 0, i32 2, i32 4, i32 6, i32 8, i32 10, i32 12, i32 14>
; CHECK-NEXT: [[TMP2:%.*]] = call <8 x i8> @llvm.abs.v8i8(<8 x i8> [[TMP1]], i1 false)
; CHECK-NEXT: [[TMP3:%.*]] = sext <8 x i8> [[TMP2]] to <8 x i32>
; CHECK-NEXT: [[TMP4:%.*]] = call i32 @llvm.vector.reduce.add.v8i32(<8 x i32> [[TMP3]])
; CHECK-NEXT: ret i32 [[TMP4]]
;
entry:
%0 = load i8, ptr %a, align 1
Expand Down Expand Up @@ -57,11 +58,12 @@ define i32 @sum_of_abs_stride_3(ptr noalias %a, ptr noalias %b) {
; CHECK-LABEL: define i32 @sum_of_abs_stride_3
; CHECK-SAME: (ptr noalias [[A:%.*]], ptr noalias [[B:%.*]]) #[[ATTR0]] {
; CHECK-NEXT: entry:
; CHECK-NEXT: [[TMP0:%.*]] = call <8 x i8> @llvm.experimental.vp.strided.load.v8i8.p0.i64(ptr align 1 [[A]], i64 3, <8 x i1> splat (i1 true), i32 8)
; CHECK-NEXT: [[TMP1:%.*]] = call <8 x i8> @llvm.abs.v8i8(<8 x i8> [[TMP0]], i1 false)
; CHECK-NEXT: [[TMP2:%.*]] = sext <8 x i8> [[TMP1]] to <8 x i32>
; CHECK-NEXT: [[TMP3:%.*]] = call i32 @llvm.vector.reduce.add.v8i32(<8 x i32> [[TMP2]])
; CHECK-NEXT: ret i32 [[TMP3]]
; CHECK-NEXT: [[TMP0:%.*]] = call <22 x i8> @llvm.masked.load.v22i8.p0(ptr [[A]], i32 1, <22 x i1> <i1 true, i1 false, i1 false, i1 true, i1 false, i1 false, i1 true, i1 false, i1 false, i1 true, i1 false, i1 false, i1 true, i1 false, i1 false, i1 true, i1 false, i1 false, i1 true, i1 false, i1 false, i1 true>, <22 x i8> poison)
; CHECK-NEXT: [[TMP1:%.*]] = shufflevector <22 x i8> [[TMP0]], <22 x i8> poison, <8 x i32> <i32 0, i32 3, i32 6, i32 9, i32 12, i32 15, i32 18, i32 21>
; CHECK-NEXT: [[TMP2:%.*]] = call <8 x i8> @llvm.abs.v8i8(<8 x i8> [[TMP1]], i1 false)
; CHECK-NEXT: [[TMP3:%.*]] = sext <8 x i8> [[TMP2]] to <8 x i32>
; CHECK-NEXT: [[TMP4:%.*]] = call i32 @llvm.vector.reduce.add.v8i32(<8 x i32> [[TMP3]])
; CHECK-NEXT: ret i32 [[TMP4]]
;
entry:
%0 = load i8, ptr %a, align 1
Expand Down
2 changes: 1 addition & 1 deletion llvm/test/Transforms/SLPVectorizer/X86/reorder-reused-masked-gather.ll
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
define void @test(ptr noalias %0, ptr %p) {
; CHECK-LABEL: @test(
; CHECK-NEXT: [[TMP5:%.*]] = getelementptr inbounds float, ptr [[TMP0:%.*]], i64 2
; CHECK-NEXT: [[TMP3:%.*]] = call <16 x float> @llvm.masked.load.v16f32.p0(ptr [[P:%.*]], i32 4, <16 x i1> <i1 true, i1 false, i1 true, i1 false, i1 true, i1 true, i1 true, i1 true, i1 true, i1 false, i1 false, i1 false, i1 false, i1 false, i1 false, i1 true>, <16 x float> poison)
; CHECK-NEXT: [[TMP3:%.*]] = call <16 x float> @llvm.masked.load.v16f32.p0(ptr [[P:%.*]], i32 16, <16 x i1> <i1 true, i1 false, i1 true, i1 false, i1 true, i1 true, i1 true, i1 true, i1 true, i1 false, i1 false, i1 false, i1 false, i1 false, i1 false, i1 true>, <16 x float> poison)
; CHECK-NEXT: [[TMP6:%.*]] = shufflevector <16 x float> [[TMP3]], <16 x float> poison, <8 x i32> <i32 15, i32 4, i32 5, i32 0, i32 2, i32 6, i32 7, i32 8>
; CHECK-NEXT: [[TMP7:%.*]] = shufflevector <16 x float> [[TMP3]], <16 x float> poison, <16 x i32> <i32 15, i32 4, i32 5, i32 15, i32 4, i32 5, i32 15, i32 0, i32 5, i32 2, i32 6, i32 7, i32 8, i32 6, i32 7, i32 8>
; CHECK-NEXT: [[TMP9:%.*]] = shufflevector <16 x float> [[TMP3]], <16 x float> <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, float poison, float 0.000000e+00, float poison, float 0.000000e+00, float 0.000000e+00, float 0.000000e+00, float 0.000000e+00, float 0.000000e+00, float 0.000000e+00>, <16 x i32> <i32 16, i32 17, i32 18, i32 19, i32 20, i32 21, i32 22, i32 4, i32 24, i32 15, i32 26, i32 27, i32 28, i32 29, i32 30, i32 31>
Expand Down
Loading