[llvm] [SROA] Extend SROA to support dynamic indexing (PR #217188)
via llvm-commits
llvm-commits at lists.llvm.org
Tue Aug 18 18:47:22 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-llvm-ir
Author: Stephen Verderame (stephenverderame)
<details>
<summary>Changes</summary>
Instead of bailing completely when non-constant GEP indices are encountered, this patch will allow SROA to handle non-constant, but bounded indices by considering an operation as potentially acessing memory between the minimum and maximum address. This would allow splitting aggregates where only part of it is accessed dynamically, which can reduce the size of the alloca.
---
Patch is 109.33 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/217188.diff
7 Files Affected:
- (modified) llvm/include/llvm/Analysis/PtrUseVisitor.h (+30-3)
- (modified) llvm/include/llvm/IR/Instructions.h (+12-7)
- (modified) llvm/lib/Analysis/PtrUseVisitor.cpp (+27-8)
- (modified) llvm/lib/IR/Instructions.cpp (+5-3)
- (modified) llvm/lib/Transforms/Scalar/SROA.cpp (+454-153)
- (added) llvm/test/Transforms/SROA/non-constant.ll (+900)
- (modified) llvm/test/Transforms/SROA/select-gep.ll (+8-6)
``````````diff
diff --git a/llvm/include/llvm/Analysis/PtrUseVisitor.h b/llvm/include/llvm/Analysis/PtrUseVisitor.h
index 304c147190728..50887c7443527 100644
--- a/llvm/include/llvm/Analysis/PtrUseVisitor.h
+++ b/llvm/include/llvm/Analysis/PtrUseVisitor.h
@@ -26,6 +26,7 @@
#include "llvm/ADT/PointerIntPair.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/ADT/SmallVector.h"
+#include "llvm/IR/ConstantRange.h"
#include "llvm/IR/DerivedTypes.h"
#include "llvm/IR/InstVisitor.h"
#include "llvm/IR/IntrinsicInst.h"
@@ -134,6 +135,7 @@ class PtrUseVisitorBase {
UseAndIsOffsetKnownPair UseAndIsOffsetKnown;
APInt Offset;
+ APInt HighOffset;
};
/// The worklist of to-visit uses.
@@ -158,6 +160,9 @@ class PtrUseVisitorBase {
/// The constant offset of the use if that is known.
APInt Offset;
+ /// The maximum constant offset of the use if that is known.
+ APInt HighOffset;
+
/// @}
/// Note that the constructor is protected because this class must be a base
@@ -174,7 +179,14 @@ class PtrUseVisitorBase {
///
/// This routine does the heavy lifting of the pointer walk by computing
/// offsets and looking through GEPs.
- LLVM_ABI bool adjustOffsetForGEP(GetElementPtrInst &GEPI);
+ ///
+ /// If `RangeAnalysis` is provided, it is queried for the signed range of
+ /// each non-constant index, and the low and high bounds of that range are
+ /// accumulated into `Offset` and `HighOffset` respectively.
+ LLVM_ABI bool adjustOffsetForGEP(
+ GetElementPtrInst &GEPI,
+ function_ref<bool(const Value &, ConstantRange &)> RangeAnalysis =
+ nullptr);
};
} // end namespace detail
@@ -230,6 +242,7 @@ class PtrUseVisitor : protected InstVisitor<DerivedT>,
IntegerType *IntIdxTy = cast<IntegerType>(DL.getIndexType(I.getType()));
IsOffsetKnown = true;
Offset = APInt(IntIdxTy->getBitWidth(), 0);
+ HighOffset = Offset;
PI.reset();
// Enqueue the uses of this pointer.
@@ -240,8 +253,10 @@ class PtrUseVisitor : protected InstVisitor<DerivedT>,
UseToVisit ToVisit = Worklist.pop_back_val();
U = ToVisit.UseAndIsOffsetKnown.getPointer();
IsOffsetKnown = ToVisit.UseAndIsOffsetKnown.getInt();
- if (IsOffsetKnown)
+ if (IsOffsetKnown) {
Offset = std::move(ToVisit.Offset);
+ HighOffset = std::move(ToVisit.HighOffset);
+ }
Instruction *I = cast<Instruction>(U->getUser());
static_cast<DerivedT*>(this)->visit(I);
@@ -274,9 +289,13 @@ class PtrUseVisitor : protected InstVisitor<DerivedT>,
return;
// If we can't walk the GEP, clear the offset.
- if (!adjustOffsetForGEP(GEPI)) {
+ if (!adjustOffsetForGEP(GEPI, [&](const Value &V, ConstantRange &CR) {
+ return static_cast<DerivedT *>(this)->getNonConstantGepIndexRange(
+ GEPI, V, CR);
+ })) {
IsOffsetKnown = false;
Offset = APInt();
+ HighOffset = APInt();
}
// Enqueue the users now that the offset has been adjusted.
@@ -309,6 +328,14 @@ class PtrUseVisitor : protected InstVisitor<DerivedT>,
PI.setEscaped(&CB);
Base::visitCallBase(CB);
}
+
+ /// Report the signed range a non-constant GEP index can take, if known.
+ ///
+ /// \returns false if the range cannot be determined.
+ bool getNonConstantGepIndexRange(const GetElementPtrInst &, const Value &,
+ ConstantRange &) const {
+ return false;
+ }
};
} // end namespace llvm
diff --git a/llvm/include/llvm/IR/Instructions.h b/llvm/include/llvm/IR/Instructions.h
index e0b26c62d7854..e6a281fee04d2 100644
--- a/llvm/include/llvm/IR/Instructions.h
+++ b/llvm/include/llvm/IR/Instructions.h
@@ -1198,13 +1198,18 @@ class GetElementPtrInst : public Instruction {
/// Accumulate the constant address offset of this GEP if possible.
///
/// This routine accepts an APInt into which it will accumulate the constant
- /// offset of this GEP if the GEP is in fact constant. If the GEP is not
- /// all-constant, it returns false and the value of the offset APInt is
- /// undefined (it is *not* preserved!). The APInt passed into this routine
- /// must be at least as wide as the IntPtr type for the address space of
- /// the base GEP pointer.
- LLVM_ABI bool accumulateConstantOffset(const DataLayout &DL,
- APInt &Offset) const;
+ /// offset of this GEP. If the GEP is not all-constant and `ExternalAnalysis`
+ /// is null or cannot provide a value for any offset, it returns false and the
+ /// value of the offset APInt is undefined (it is *not* preserved!). The APInt
+ /// passed into this routine must be at least as wide as the IntPtr type for
+ /// the address space of the base GEP pointer. If an ExternalAnalysis is
+ /// supplied, it will be used to determine the value of any non-constant
+ /// indicies. If a value can be provided, `ExternalAnalysis` should return
+ /// true and set the value of the APInt. Otherwise, it should return false,
+ /// which will cause this routine to return false.
+ LLVM_ABI bool accumulateConstantOffset(
+ const DataLayout &DL, APInt &Offset,
+ function_ref<bool(Value &, APInt &)> ExternalAnalysis = nullptr) const;
LLVM_ABI bool
collectOffset(const DataLayout &DL, unsigned BitWidth,
SmallMapVector<Value *, APInt, 4> &VariableOffsets,
diff --git a/llvm/lib/Analysis/PtrUseVisitor.cpp b/llvm/lib/Analysis/PtrUseVisitor.cpp
index 9c79546f491ef..e8e9b8683655b 100644
--- a/llvm/lib/Analysis/PtrUseVisitor.cpp
+++ b/llvm/lib/Analysis/PtrUseVisitor.cpp
@@ -20,24 +20,43 @@ using namespace llvm;
void detail::PtrUseVisitorBase::enqueueUsers(Value &I) {
for (Use &U : I.uses()) {
if (VisitedUses.insert(&U).second) {
- UseToVisit NewU = {
- UseToVisit::UseAndIsOffsetKnownPair(&U, IsOffsetKnown),
- Offset
- };
+ UseToVisit NewU = {UseToVisit::UseAndIsOffsetKnownPair(&U, IsOffsetKnown),
+ Offset, HighOffset};
Worklist.push_back(std::move(NewU));
}
}
}
-bool detail::PtrUseVisitorBase::adjustOffsetForGEP(GetElementPtrInst &GEPI) {
+bool detail::PtrUseVisitorBase::adjustOffsetForGEP(
+ GetElementPtrInst &GEPI,
+ function_ref<bool(const Value &, ConstantRange &)> RangeAnalysis) {
if (!IsOffsetKnown)
return false;
- APInt TmpOffset(DL.getIndexTypeSizeInBits(GEPI.getType()), 0);
- if (GEPI.accumulateConstantOffset(DL, TmpOffset)) {
+ if (!RangeAnalysis) {
+ APInt TmpOffset(DL.getIndexTypeSizeInBits(GEPI.getType()), 0);
+ if (!GEPI.accumulateConstantOffset(DL, TmpOffset))
+ return false;
Offset += TmpOffset.sextOrTrunc(Offset.getBitWidth());
+ HighOffset += TmpOffset.sextOrTrunc(HighOffset.getBitWidth());
return true;
}
- return false;
+ auto AccumulateBound = [&](APInt &Accum, bool IsUpperBound) {
+ auto ExternalAnalysis = [&](Value &V, APInt &Index) {
+ ConstantRange CR(Index.getBitWidth(), /*isFullSet=*/false);
+ if (!RangeAnalysis(V, CR))
+ return false;
+ Index = IsUpperBound ? CR.getSignedMax() : CR.getSignedMin();
+ return true;
+ };
+ APInt Tmp(DL.getIndexTypeSizeInBits(GEPI.getType()), 0);
+ if (!GEPI.accumulateConstantOffset(DL, Tmp, ExternalAnalysis))
+ return false;
+ Accum += Tmp.sextOrTrunc(Accum.getBitWidth());
+ return true;
+ };
+
+ return AccumulateBound(Offset, /*IsUpperBound=*/false) &&
+ AccumulateBound(HighOffset, /*IsUpperBound=*/true);
}
diff --git a/llvm/lib/IR/Instructions.cpp b/llvm/lib/IR/Instructions.cpp
index 6ca12da2454cc..409fa203abc01 100644
--- a/llvm/lib/IR/Instructions.cpp
+++ b/llvm/lib/IR/Instructions.cpp
@@ -1688,10 +1688,12 @@ bool GetElementPtrInst::hasNoUnsignedWrap() const {
return cast<GEPOperator>(this)->hasNoUnsignedWrap();
}
-bool GetElementPtrInst::accumulateConstantOffset(const DataLayout &DL,
- APInt &Offset) const {
+bool GetElementPtrInst::accumulateConstantOffset(
+ const DataLayout &DL, APInt &Offset,
+ function_ref<bool(Value &, APInt &)> ExternalAnalysis) const {
// Delegate to the generic GEPOperator implementation.
- return cast<GEPOperator>(this)->accumulateConstantOffset(DL, Offset);
+ return cast<GEPOperator>(this)->accumulateConstantOffset(DL, Offset,
+ ExternalAnalysis);
}
bool GetElementPtrInst::collectOffset(
diff --git a/llvm/lib/Transforms/Scalar/SROA.cpp b/llvm/lib/Transforms/Scalar/SROA.cpp
index e4770154e2998..5da912ea346ac 100644
--- a/llvm/lib/Transforms/Scalar/SROA.cpp
+++ b/llvm/lib/Transforms/Scalar/SROA.cpp
@@ -48,6 +48,7 @@
#include "llvm/IR/BasicBlock.h"
#include "llvm/IR/Constant.h"
#include "llvm/IR/ConstantFolder.h"
+#include "llvm/IR/ConstantRange.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/DIBuilder.h"
#include "llvm/IR/DataLayout.h"
@@ -333,8 +334,12 @@ static DebugVariable getAggregateVariable(DbgVariableRecord *DVR) {
/// \param OldAlloca Alloca for the variable before splitting.
/// \param IsSplit True if the store (not necessarily alloca)
/// is being split.
-/// \param OldAllocaOffsetInBits Offset of the slice taken from OldAlloca.
-/// \param SliceSizeInBits New number of bits being written to.
+/// \param OldAllocaOffsetInBits Offset of the slice taken from OldAlloca, if
+/// constant. Must not be nullopt if the store is
+/// being split.
+/// \param SliceSizeInBits New number of bits being written to, if
+/// constant. Must not be nullopt if the store is
+/// being split.
/// \param OldInst Instruction that is being split.
/// \param Inst New instruction performing this part of the
/// split store.
@@ -342,10 +347,10 @@ static DebugVariable getAggregateVariable(DbgVariableRecord *DVR) {
/// \param Value Stored value.
/// \param DL Datalayout.
static void migrateDebugInfo(AllocaInst *OldAlloca, bool IsSplit,
- uint64_t OldAllocaOffsetInBits,
- uint64_t SliceSizeInBits, Instruction *OldInst,
- Instruction *Inst, Value *Dest, Value *Value,
- const DataLayout &DL) {
+ std::optional<uint64_t> OldAllocaOffsetInBits,
+ std::optional<uint64_t> SliceSizeInBits,
+ Instruction *OldInst, Instruction *Inst,
+ Value *Dest, Value *Value, const DataLayout &DL) {
// If we want allocas to be migrated using this helper then we need to ensure
// that the BaseFragments map code still works. A simple solution would be
// to choose to always clone alloca dbg_assigns (rather than sometimes
@@ -391,6 +396,7 @@ static void migrateDebugInfo(AllocaInst *OldAlloca, bool IsSplit,
bool SetKillLocation = false;
if (IsSplit) {
+ assert(OldAllocaOffsetInBits && SliceSizeInBits);
std::optional<DIExpression::FragmentInfo> BaseFragment;
{
auto R = BaseFragments.find(getAggregateVariable(DbgAssign));
@@ -402,7 +408,7 @@ static void migrateDebugInfo(AllocaInst *OldAlloca, bool IsSplit,
Expr->getFragmentInfo();
DIExpression::FragmentInfo NewFragment;
FragCalcResult Result = calculateFragment(
- DbgAssign->getVariable(), OldAllocaOffsetInBits, SliceSizeInBits,
+ DbgAssign->getVariable(), *OldAllocaOffsetInBits, *SliceSizeInBits,
BaseFragment, CurrentFragment, NewFragment);
if (Result == Skip)
@@ -533,12 +539,16 @@ class Slice {
/// split.
PointerIntPair<Use *, 1, bool> UseAndIsSplittable;
+ /// Whether the slice is indexed by non-constant indices.
+ bool IsDynamic;
+
public:
Slice() = default;
- Slice(uint64_t BeginOffset, uint64_t EndOffset, Use *U, bool IsSplittable)
+ Slice(uint64_t BeginOffset, uint64_t EndOffset, Use *U, bool IsSplittable,
+ bool IsDynamic)
: BeginOffset(BeginOffset), EndOffset(EndOffset),
- UseAndIsSplittable(U, IsSplittable) {}
+ UseAndIsSplittable(U, IsSplittable), IsDynamic(IsDynamic) {}
uint64_t beginOffset() const { return BeginOffset; }
uint64_t endOffset() const { return EndOffset; }
@@ -551,6 +561,8 @@ class Slice {
bool isDead() const { return getUse() == nullptr; }
void kill() { UseAndIsSplittable.setPointer(nullptr); }
+ bool isDynamic() const { return IsDynamic; }
+
/// Support for ordering ranges.
///
/// This provides an ordering over ranges such that start offsets are
@@ -584,6 +596,36 @@ class Slice {
bool operator!=(const Slice &RHS) const { return !operator==(RHS); }
};
+/// The indices and strides used to offset a pointer from an alloca.
+class AccumulatedGEPIndices {
+ APInt ConstantOffset;
+ SmallVector<std::pair<Value *, APInt>, 4> VariableIndices;
+
+public:
+ explicit AccumulatedGEPIndices(unsigned BitWidth)
+ : ConstantOffset(BitWidth, 0, true) {}
+
+ /// Adds a GEP index of `V` with stride `Stride`.
+ void addIndex(Value *V, const APInt &Stride) {
+ VariableIndices.emplace_back(
+ V, Stride.sextOrTrunc(ConstantOffset.getBitWidth()));
+ }
+
+ /// Gets the constant offset of the pointer relative to `Base`.
+ APInt getAdjustedConstantOffset(uint64_t Base) const {
+ return ConstantOffset - APInt(ConstantOffset.getBitWidth(), Base);
+ }
+
+ /// Increments the constant offset by `V`.
+ void incConstantOffset(const APInt &V) {
+ ConstantOffset += V.sextOrTrunc(ConstantOffset.getBitWidth());
+ }
+
+ auto indices() const {
+ return make_range(VariableIndices.begin(), VariableIndices.end());
+ }
+};
+
/// Representation of the alloca slices.
///
/// This class represents the slices of an alloca which are formed by its
@@ -594,7 +636,8 @@ class Slice {
class AllocaSlices {
public:
/// Construct the slices of a particular alloca.
- AllocaSlices(const DataLayout &DL, AllocaInst &AI);
+ AllocaSlices(const DataLayout &DL, AllocaInst &AI, AssumptionCache &AC,
+ DominatorTree &DT);
/// Test whether a pointer to the allocation escapes our analysis.
///
@@ -655,6 +698,12 @@ class AllocaSlices {
/// need to replace with undef.
ArrayRef<Use *> getDeadOperands() const { return DeadOperands; }
+ /// Gets the accumulated offset of `Ptr` from its alloca.
+ const AccumulatedGEPIndices *getAccumulatedGEPIndices(Value *Ptr) const {
+ auto It = PtrOffsets.find(Ptr);
+ return It == PtrOffsets.end() ? nullptr : &It->second;
+ }
+
#if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
void print(raw_ostream &OS, const_iterator I, StringRef Indent = " ") const;
void printSlice(raw_ostream &OS, const_iterator I,
@@ -694,6 +743,9 @@ class AllocaSlices {
/// details.
SmallVector<Slice, 8> Slices;
+ /// Symbolic offsets of each pointer value derived from the alloca.
+ SmallDenseMap<Value *, AccumulatedGEPIndices> PtrOffsets;
+
/// Instructions which will become dead if we rewrite the alloca.
///
/// Note that these are not separated by slice. This is because we expect an
@@ -1023,17 +1075,26 @@ class AllocaSlices::SliceBuilder : public PtrUseVisitor<SliceBuilder> {
const uint64_t AllocSize;
AllocaSlices &AS;
+ AssumptionCache &AC;
+ DominatorTree &DT;
SmallDenseMap<Instruction *, unsigned> MemTransferSliceMap;
SmallDenseMap<Instruction *, uint64_t> PHIOrSelectSizes;
+ /// Map from GEP instruction and non-constant index to the range of values
+ /// that index can take, if such a range can be determined.
+ SmallDenseMap<std::pair<const Instruction *, const Value *>, ConstantRange>
+ DynGepRanges;
+
/// Set to de-duplicate dead instructions found in the use walk.
SmallPtrSet<Instruction *, 4> VisitedDeadInsts;
public:
- SliceBuilder(const DataLayout &DL, AllocaInst &AI, AllocaSlices &AS)
+ SliceBuilder(const DataLayout &DL, AllocaInst &AI, AllocaSlices &AS,
+ AssumptionCache &AC, DominatorTree &DT)
: PtrUseVisitor<SliceBuilder>(DL),
- AllocSize(AI.getAllocationSize(DL)->getFixedValue()), AS(AS) {}
+ AllocSize(AI.getAllocationSize(DL)->getFixedValue()), AS(AS), AC(AC),
+ DT(DT) {}
private:
void markAsDead(Instruction &I) {
@@ -1041,6 +1102,17 @@ class AllocaSlices::SliceBuilder : public PtrUseVisitor<SliceBuilder> {
AS.DeadUsers.push_back(&I);
}
+ bool isSliceDynamic() const { return Offset != HighOffset; }
+
+ bool isInvalidDynamicUse() {
+ return isSliceDynamic() && !AS.getAccumulatedGEPIndices(U->get());
+ }
+
+ bool isInvalidDynamicMemIntrinsic(ConstantInt *Length) {
+ return isSliceDynamic() &&
+ (!AS.getAccumulatedGEPIndices(U->get()) || !Length);
+ }
+
void insertUse(Instruction &I, const APInt &Offset, uint64_t Size,
bool IsSplittable = false) {
// Completely skip uses which have a zero size or start either before or
@@ -1056,7 +1128,7 @@ class AllocaSlices::SliceBuilder : public PtrUseVisitor<SliceBuilder> {
}
uint64_t BeginOffset = Offset.getZExtValue();
- uint64_t EndOffset = BeginOffset + Size;
+ uint64_t EndOffset = HighOffset.getZExtValue() + Size;
// Clamp the end offset to the end of the allocation. Note that this is
// formulated to handle even the case where "BeginOffset + Size" overflows.
@@ -1065,7 +1137,8 @@ class AllocaSlices::SliceBuilder : public PtrUseVisitor<SliceBuilder> {
// some instructions are dead but not others. We can't completely ignore
// them, and so have to record at least the information here.
assert(AllocSize >= BeginOffset); // Established above.
- if (Size > AllocSize - BeginOffset) {
+ if (HighOffset.uge(AllocSize) ||
+ Size > AllocSize - HighOffset.getZExtValue()) {
LLVM_DEBUG(dbgs() << "WARNING: Clamping a " << Size << " byte use @"
<< Offset << " to remain within the " << AllocSize
<< " byte alloca:\n"
@@ -1074,13 +1147,18 @@ class AllocaSlices::SliceBuilder : public PtrUseVisitor<SliceBuilder> {
EndOffset = AllocSize;
}
- AS.Slices.push_back(Slice(BeginOffset, EndOffset, U, IsSplittable));
+ bool IsDynamic = isSliceDynamic();
+ // We cannot split an alloca that is dynamically indexed.
+ IsSplittable &= !IsDynamic;
+ AS.Slices.push_back(
+ Slice(BeginOffset, EndOffset, U, IsSplittable, IsDynamic));
}
void visitBitCastInst(BitCastInst &BC) {
if (BC.use_empty())
return markAsDead(BC);
+ AS.PtrOffsets.try_emplace(&BC, getOrCreateOffsetsOf(BC.getOperand(0)));
return Base::visitBitCastInst(BC);
}
@@ -1088,6 +1166,8 @@ class AllocaSlices::SliceBuilder : public PtrUseVisitor<SliceBuilder> {
if (ASC.use_empty())
return markAsDead(ASC);
+ AS.PtrOffsets.try_emplace(&ASC,
+ getOrCreateOffsetsOf(ASC.getPointerOperand()));
return Base::visitAddrSpaceCastInst(ASC);
}
@@ -1095,6 +1175,7 @@ class AllocaSlices::SliceBuilder : public PtrUseVisitor<SliceBuilder> {
if (GEPI.use_empty())
return markAsDead(GEPI);
+ computeNonConstantGepRanges(GEPI);
return Base::visitGetElementPtrInst(GEPI);
}
@@ -1106,6 +1187,9 @@ class AllocaSlices::SliceBuilder : public PtrUseVisitor<SliceBuilder> {
bool IsSplittable =
Ty->isIntegerTy() && !IsVolatile && DL.typeSizeEqualsStoreSize(Ty);
+ if (isInvalidDynamicUse())
+ return PI.setAborted(&I);
+
insertUse(I, Offset, Size, IsSplittable);
}
@@ -1178,7 +1262,7 @@ class AllocaSlices::SliceBuilder : public PtrUseVisitor<SliceBuilder> {
// Zero-length mem transfer intrinsics can be ignored entirely.
return markAsDead(II);
- if (!IsOffsetKnown)
+ if (!IsOffsetKnown || isInvalidDynamicMemIntrinsic(Length))
return PI.setAborted(&II);
insertUse(II, Offset,
@@ -1198,7 +1282,7 @@ class AllocaSlices::SliceBuilder : public PtrUseVisitor<SliceBuilder> {
if (VisitedDeadInsts.count(&II))
return;
- if (!IsOffsetKnown)
+ if (!IsOffsetKnown || isInvalidDynamicMemIntrinsic(Length))
return PI.setAborted(&II);
// This side of the tr...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/217188
More information about the llvm-commits
mailing list