Skip to content

Commit

Permalink
[Assignment Tracking] Fix migrateDebuginfo in SROA
Browse files Browse the repository at this point in the history 
Without this patch, migrateDebugInfo doesn't understand how to handle existing
fragments that are smaller than the to-be-split store. This can occur
if. e.g. a vector store (1 dbg.assign) is split (many dbg.assigns - 1 fragment
for each scalar) and later those stores are re-vectorized (many dbg.assigns),
and then SROA runs on that.

The approach taken in this patch is to drop intrinsics with fragments outside
of the slice.

For example, starting with:

  store <2 x float> %v, ptr %dest !DIAssignID !1
  call void @llvm.dbg.assign(..., DIExpression(DW_OP_LLVM_fragment, 0, 32), !1, ...)
  call void @llvm.dbg.assign(..., DIExpression(DW_OP_LLVM_fragment, 32, 32), !1, ...)

When visiting the slice of bits 0 to 31 we get:

  store float %v.extract.0, ptr %dest !DIAssignID !2
  call void @llvm.dbg.assign(..., DIExpression(DW_OP_LLVM_fragment, 0, 32), !2, ...)

The other dbg.assign associated with the currently-split store is dropped for
this split part. And visiting bits 32 to 63 we get the following:

  store float %v.extract.1, ptr %adjusted.dest !DIAssignID !3
  call void @llvm.dbg.assign(..., DIExpression(DW_OP_LLVM_fragment, 32, 32), !3, ...)

I've added two tests that cover this case.

Implementing this meant re-writing the fragment-calculation part of
migrateDebugInfo to work with the absolute offset of the new slice in terms of
the base alloca (instead of the offset of the slice into the new alloca), the
fragment (if any) of the variable associated with the base alloca, and the
fragment associated with the split store. Because we need the offset into the
base alloca for the variables being split, some careful wiring is required for
memory intrinsics due to the fact that memory intrinsics can be split when
either the source or dest allocas are split. In the case where the source
alloca drives the splitting, we need to be careful to pass migrateDebugInfo the
information in relation to the dest alloca.

Reviewed By: StephenTozer

Differential Revision: https://reviews.llvm.org/D143146
  • Loading branch information
@OCHyams
OCHyams committed Feb 10, 2023
1 parent c52255d commit 295f5fa
Show file tree
Hide file tree
Showing 4 changed files with 353 additions and 57 deletions.
5 changes: 5 additions & 0 deletions llvm/include/llvm/IR/DebugInfoMetadata.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2776,6 +2776,11 @@ class DIExpression : public MDNode {
struct FragmentInfo {
uint64_t SizeInBits;
uint64_t OffsetInBits;
/// Return the index of the first bit of the fragment.
uint64_t startInBits() const { return OffsetInBits; }
/// Return the index of the bit after the end of the fragment, e.g. for
/// fragment offset=16 and size=32 return their sum, 48.
uint64_t endInBits() const { return OffsetInBits + SizeInBits; }
};

/// Retrieve the details of this fragment expression.
Expand Down
204 changes: 147 additions & 57 deletions llvm/lib/Transforms/Scalar/SROA.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -118,22 +118,74 @@ STATISTIC(NumVectorized, "Number of vectorized aggregates");
/// GEPs.
static cl::opt<bool> SROAStrictInbounds("sroa-strict-inbounds", cl::init(false),
cl::Hidden);
/// Disable running mem2reg during SROA in order to test or debug SROA.
static cl::opt<bool> SROASkipMem2Reg("sroa-skip-mem2reg", cl::init(false),
cl::Hidden);
namespace {

/// Calculate the fragment of a variable to use when slicing a store
/// based on the slice dimensions, existing fragment, and base storage
/// fragment.
/// Note that a returned value of std::nullopt indicates that there is
/// no appropriate fragment available (rather than meaning use the whole
/// variable, which is a common usage). Because the store is being sliced
/// we always expect a fragment - there's never a case where the whole
/// variable should be used.
static std::optional<DIExpression::FragmentInfo>
calculateFragment(uint64_t NewStorageSliceOffsetInBits,
uint64_t NewStorageSliceSizeInBits,
std::optional<DIExpression::FragmentInfo> StorageFragment,
std::optional<DIExpression::FragmentInfo> CurrentFragment) {
DIExpression::FragmentInfo Target;
// If the base storage describes part of the variable apply the offset and
// the size constraint.
if (StorageFragment) {
Target.SizeInBits =
std::min(NewStorageSliceSizeInBits, StorageFragment->SizeInBits);
Target.OffsetInBits =
NewStorageSliceOffsetInBits + StorageFragment->OffsetInBits;
} else {
Target.SizeInBits = NewStorageSliceSizeInBits;
Target.OffsetInBits = NewStorageSliceOffsetInBits;
}

// No additional work to do if there isn't a fragment already, or there is
// but it already exactly describes the new assignment.
if (!CurrentFragment || *CurrentFragment == Target)
return Target;

// Reject the target fragment if it doesn't fit wholly within the current
// fragment. TODO: We could instead chop up the target to fit in the case of
// a partial overlap.
if (Target.startInBits() < CurrentFragment->startInBits() ||
Target.endInBits() > CurrentFragment->endInBits())
return std::nullopt;

// Target fits within the current fragment, return it.
return Target;
}

static DebugVariable getAggregateVariable(DbgVariableIntrinsic *DVI) {
return DebugVariable(DVI->getVariable(), std::nullopt,
DVI->getDebugLoc().getInlinedAt());
}

/// Find linked dbg.assign and generate a new one with the correct
/// FragmentInfo. Link Inst to the new dbg.assign. If Value is nullptr the
/// value component is copied from the old dbg.assign to the new.
/// \param OldAlloca Alloca for the variable before splitting.
/// \param RelativeOffsetInBits Offset into \p OldAlloca relative to the
/// offset prior to splitting (change in offset).
/// \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 OldInst Instruction that is being split.
/// \param Inst New instruction performing this part of the
/// split store.
/// \param Dest Store destination.
/// \param Value Stored value.
/// \param DL Datalayout.
static void migrateDebugInfo(AllocaInst *OldAlloca,
uint64_t RelativeOffsetInBits,
static void migrateDebugInfo(AllocaInst *OldAlloca, bool IsSplit,
uint64_t OldAllocaOffsetInBits,
uint64_t SliceSizeInBits, Instruction *OldInst,
Instruction *Inst, Value *Dest, Value *Value,
const DataLayout &DL) {
Expand All @@ -144,52 +196,68 @@ static void migrateDebugInfo(AllocaInst *OldAlloca,

LLVM_DEBUG(dbgs() << " migrateDebugInfo\n");
LLVM_DEBUG(dbgs() << " OldAlloca: " << *OldAlloca << "\n");
LLVM_DEBUG(dbgs() << " RelativeOffset: " << RelativeOffsetInBits << "\n");
LLVM_DEBUG(dbgs() << " IsSplit: " << IsSplit << "\n");
LLVM_DEBUG(dbgs() << " OldAllocaOffsetInBits: " << OldAllocaOffsetInBits
<< "\n");
LLVM_DEBUG(dbgs() << " SliceSizeInBits: " << SliceSizeInBits << "\n");
LLVM_DEBUG(dbgs() << " OldInst: " << *OldInst << "\n");
LLVM_DEBUG(dbgs() << " Inst: " << *Inst << "\n");
LLVM_DEBUG(dbgs() << " Dest: " << *Dest << "\n");
if (Value)
LLVM_DEBUG(dbgs() << " Value: " << *Value << "\n");

/// Map of aggregate variables to their fragment associated with OldAlloca.
DenseMap<DebugVariable, std::optional<DIExpression::FragmentInfo>>
BaseFragments;
for (auto *DAI : at::getAssignmentMarkers(OldAlloca))
BaseFragments[getAggregateVariable(DAI)] =
DAI->getExpression()->getFragmentInfo();

// The new inst needs a DIAssignID unique metadata tag (if OldInst has
// one). It shouldn't already have one: assert this assumption.
assert(!Inst->getMetadata(LLVMContext::MD_DIAssignID));
DIAssignID *NewID = nullptr;
auto &Ctx = Inst->getContext();
DIBuilder DIB(*OldInst->getModule(), /*AllowUnresolved*/ false);
uint64_t AllocaSizeInBits = *OldAlloca->getAllocationSizeInBits(DL);
assert(OldAlloca->isStaticAlloca());

for (DbgAssignIntrinsic *DbgAssign : MarkerRange) {
LLVM_DEBUG(dbgs() << " existing dbg.assign is: " << *DbgAssign
<< "\n");
auto *Expr = DbgAssign->getExpression();

// Check if the dbg.assign already describes a fragment.
auto GetCurrentFragSize = [AllocaSizeInBits, DbgAssign,
Expr]() -> uint64_t {
if (auto FI = Expr->getFragmentInfo())
return FI->SizeInBits;
if (auto VarSize = DbgAssign->getVariable()->getSizeInBits())
return *VarSize;
// The variable type has an unspecified size. This can happen in the
// case of DW_TAG_unspecified_type types, e.g. std::nullptr_t. Because
// there is no fragment and we do not know the size of the variable type,
// we'll guess by looking at the alloca.
return AllocaSizeInBits;
};
uint64_t CurrentFragSize = GetCurrentFragSize();
bool MakeNewFragment = CurrentFragSize != SliceSizeInBits;
assert(MakeNewFragment || RelativeOffsetInBits == 0);

assert(SliceSizeInBits <= AllocaSizeInBits);
if (MakeNewFragment) {
assert(RelativeOffsetInBits + SliceSizeInBits <= CurrentFragSize);
auto E = DIExpression::createFragmentExpression(
Expr, RelativeOffsetInBits, SliceSizeInBits);
assert(E && "Failed to create fragment expr!");
Expr = *E;
if (IsSplit) {
std::optional<DIExpression::FragmentInfo> BaseFragment = std::nullopt;
{
auto R = BaseFragments.find(getAggregateVariable(DbgAssign));
if (R == BaseFragments.end())
continue;
BaseFragment = R->second;
}
std::optional<DIExpression::FragmentInfo> CurrentFragment =
Expr->getFragmentInfo();
std::optional<DIExpression::FragmentInfo> NewFragment =
calculateFragment(OldAllocaOffsetInBits, SliceSizeInBits,
BaseFragment, CurrentFragment);
// Note that std::nullopt here means "skip this fragment" rather than
// "there is no fragment / use the whole variable".
if (!NewFragment)
continue;

if (!(NewFragment == CurrentFragment)) {
if (CurrentFragment) {
// Rewrite NewFragment to be relative to the existing one (this is
// what createFragmentExpression wants). CalculateFragment has
// already resolved the size for us. FIXME: Should it return the
// relative fragment too?
NewFragment->OffsetInBits -= CurrentFragment->OffsetInBits;
}

auto E = DIExpression::createFragmentExpression(
Expr, NewFragment->OffsetInBits, NewFragment->SizeInBits);
assert(E && "Failed to create fragment expr!");
Expr = *E;
}
}

// If we haven't created a DIAssignID ID do that now and attach it to Inst.
Expand Down Expand Up @@ -2334,7 +2402,6 @@ class llvm::sroa::AllocaSliceRewriter
// original alloca.
uint64_t NewBeginOffset = 0, NewEndOffset = 0;

uint64_t RelativeOffset = 0;
uint64_t SliceSize = 0;
bool IsSplittable = false;
bool IsSplit = false;
Expand Down Expand Up @@ -2408,14 +2475,13 @@ class llvm::sroa::AllocaSliceRewriter
NewBeginOffset = std::max(BeginOffset, NewAllocaBeginOffset);
NewEndOffset = std::min(EndOffset, NewAllocaEndOffset);

RelativeOffset = NewBeginOffset - BeginOffset;
SliceSize = NewEndOffset - NewBeginOffset;
LLVM_DEBUG(dbgs() << " Begin:(" << BeginOffset << ", " << EndOffset
<< ") NewBegin:(" << NewBeginOffset << ", "
<< NewEndOffset << ") NewAllocaBegin:("
<< NewAllocaBeginOffset << ", " << NewAllocaEndOffset
<< ")\n");
assert(IsSplit || RelativeOffset == 0);
assert(IsSplit || NewBeginOffset == BeginOffset);
OldUse = I->getUse();
OldPtr = cast<Instruction>(OldUse->get());

Expand Down Expand Up @@ -2678,8 +2744,8 @@ class llvm::sroa::AllocaSliceRewriter
Pass.DeadInsts.push_back(&SI);

// NOTE: Careful to use OrigV rather than V.
migrateDebugInfo(&OldAI, RelativeOffset * 8, SliceSize * 8, &SI, Store,
Store->getPointerOperand(), OrigV, DL);
migrateDebugInfo(&OldAI, IsSplit, NewBeginOffset * 8, SliceSize * 8, &SI,
Store, Store->getPointerOperand(), OrigV, DL);
LLVM_DEBUG(dbgs() << " to: " << *Store << "\n");
return true;
}
Expand All @@ -2703,8 +2769,9 @@ class llvm::sroa::AllocaSliceRewriter
if (AATags)
Store->setAAMetadata(AATags.shift(NewBeginOffset - BeginOffset));

migrateDebugInfo(&OldAI, RelativeOffset * 8, SliceSize * 8, &SI, Store,
Store->getPointerOperand(), Store->getValueOperand(), DL);
migrateDebugInfo(&OldAI, IsSplit, NewBeginOffset * 8, SliceSize * 8, &SI,
Store, Store->getPointerOperand(),
Store->getValueOperand(), DL);

Pass.DeadInsts.push_back(&SI);
LLVM_DEBUG(dbgs() << " to: " << *Store << "\n");
Expand Down Expand Up @@ -2782,8 +2849,9 @@ class llvm::sroa::AllocaSliceRewriter
if (NewSI->isAtomic())
NewSI->setAlignment(SI.getAlign());

migrateDebugInfo(&OldAI, RelativeOffset * 8, SliceSize * 8, &SI, NewSI,
NewSI->getPointerOperand(), NewSI->getValueOperand(), DL);
migrateDebugInfo(&OldAI, IsSplit, NewBeginOffset * 8, SliceSize * 8, &SI,
NewSI, NewSI->getPointerOperand(),
NewSI->getValueOperand(), DL);

Pass.DeadInsts.push_back(&SI);
deleteIfTriviallyDead(OldOp);
Expand Down Expand Up @@ -2883,8 +2951,8 @@ class llvm::sroa::AllocaSliceRewriter
if (AATags)
New->setAAMetadata(AATags.shift(NewBeginOffset - BeginOffset));

migrateDebugInfo(&OldAI, RelativeOffset * 8, SliceSize * 8, &II, New,
New->getRawDest(), nullptr, DL);
migrateDebugInfo(&OldAI, IsSplit, NewBeginOffset * 8, SliceSize * 8, &II,
New, New->getRawDest(), nullptr, DL);

LLVM_DEBUG(dbgs() << " to: " << *New << "\n");
return false;
Expand Down Expand Up @@ -2959,8 +3027,8 @@ class llvm::sroa::AllocaSliceRewriter
if (AATags)
New->setAAMetadata(AATags.shift(NewBeginOffset - BeginOffset));

migrateDebugInfo(&OldAI, RelativeOffset * 8, SliceSize * 8, &II, New,
New->getPointerOperand(), V, DL);
migrateDebugInfo(&OldAI, IsSplit, NewBeginOffset * 8, SliceSize * 8, &II,
New, New->getPointerOperand(), V, DL);

LLVM_DEBUG(dbgs() << " to: " << *New << "\n");
return !II.isVolatile();
Expand Down Expand Up @@ -3088,8 +3156,16 @@ class llvm::sroa::AllocaSliceRewriter
if (AATags)
New->setAAMetadata(AATags.shift(NewBeginOffset - BeginOffset));

migrateDebugInfo(&OldAI, RelativeOffset * 8, SliceSize * 8, &II, New,
DestPtr, nullptr, DL);
APInt Offset(DL.getIndexTypeSizeInBits(DestPtr->getType()), 0);
if (IsDest) {
migrateDebugInfo(&OldAI, IsSplit, NewBeginOffset * 8, SliceSize * 8,
&II, New, DestPtr, nullptr, DL);
} else if (AllocaInst *Base = dyn_cast<AllocaInst>(
DestPtr->stripAndAccumulateConstantOffsets(
DL, Offset, /*AllowNonInbounds*/ true))) {
migrateDebugInfo(Base, IsSplit, Offset.getZExtValue() * 8,
SliceSize * 8, &II, New, DestPtr, nullptr, DL);
}
LLVM_DEBUG(dbgs() << " to: " << *New << "\n");
return false;
}
Expand Down Expand Up @@ -3177,8 +3253,18 @@ class llvm::sroa::AllocaSliceRewriter
if (AATags)
Store->setAAMetadata(AATags.shift(NewBeginOffset - BeginOffset));

migrateDebugInfo(&OldAI, RelativeOffset * 8, SliceSize * 8, &II, Store,
DstPtr, Src, DL);
APInt Offset(DL.getIndexTypeSizeInBits(DstPtr->getType()), 0);
if (IsDest) {

migrateDebugInfo(&OldAI, IsSplit, NewBeginOffset * 8, SliceSize * 8, &II,
Store, DstPtr, Src, DL);
} else if (AllocaInst *Base = dyn_cast<AllocaInst>(
DstPtr->stripAndAccumulateConstantOffsets(
DL, Offset, /*AllowNonInbounds*/ true))) {
migrateDebugInfo(Base, IsSplit, Offset.getZExtValue() * 8, SliceSize * 8,
&II, Store, DstPtr, Src, DL);
}

LLVM_DEBUG(dbgs() << " to: " << *Store << "\n");
return !II.isVolatile();
}
Expand Down Expand Up @@ -3540,23 +3626,22 @@ class AggLoadStoreRewriter : public InstVisitor<AggLoadStoreRewriter, bool> {

APInt Offset(
DL.getIndexSizeInBits(Ptr->getType()->getPointerAddressSpace()), 0);
if (AATags &&
GEPOperator::accumulateConstantOffset(BaseTy, GEPIndices, DL, Offset))
GEPOperator::accumulateConstantOffset(BaseTy, GEPIndices, DL, Offset);
if (AATags)
Store->setAAMetadata(AATags.shift(Offset.getZExtValue()));

// migrateDebugInfo requires the base Alloca. Walk to it from this gep.
// If we cannot (because there's an intervening non-const or unbounded
// gep) then we wouldn't expect to see dbg.assign intrinsics linked to
// this instruction.
APInt OffsetInBytes(DL.getTypeSizeInBits(Ptr->getType()), false);
Value *Base = InBoundsGEP->stripAndAccumulateInBoundsConstantOffsets(
DL, OffsetInBytes);
Value *Base = AggStore->getPointerOperand()->stripInBoundsOffsets();
if (auto *OldAI = dyn_cast<AllocaInst>(Base)) {
uint64_t SizeInBits =
DL.getTypeSizeInBits(Store->getValueOperand()->getType());
migrateDebugInfo(OldAI, OffsetInBytes.getZExtValue() * 8, SizeInBits,
AggStore, Store, Store->getPointerOperand(),
Store->getValueOperand(), DL);
migrateDebugInfo(OldAI, /*IsSplit*/ true, Offset.getZExtValue() * 8,
SizeInBits, AggStore, Store,
Store->getPointerOperand(), Store->getValueOperand(),
DL);
} else {
assert(at::getAssignmentMarkers(Store).empty() &&
"AT: unexpected debug.assign linked to store through "
Expand Down Expand Up @@ -4890,8 +4975,13 @@ bool SROAPass::promoteAllocas(Function &F) {

NumPromoted += PromotableAllocas.size();

LLVM_DEBUG(dbgs() << "Promoting allocas with mem2reg...\n");
PromoteMemToReg(PromotableAllocas, DTU->getDomTree(), AC);
if (SROASkipMem2Reg) {
LLVM_DEBUG(dbgs() << "Not promoting allocas with mem2reg!\n");
} else {
LLVM_DEBUG(dbgs() << "Promoting allocas with mem2reg...\n");
PromoteMemToReg(PromotableAllocas, DTU->getDomTree(), AC);
}

PromotableAllocas.clear();
return true;
}
Expand Down

0 comments on commit 295f5fa

Please sign in to comment.