Skip to content

InferAddressSpaces: Factor replacement loop into function [NFC] #104430

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 5 commits into from
Aug 20, 2024
Merged

InferAddressSpaces: Factor replacement loop into function [NFC] #104430

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
feef401
InferAddressSpaces: Factor replacement loop into function
arsenm Aug 15, 2024
2288448
Early return
arsenm Aug 20, 2024
c257409
Comments
arsenm Aug 20, 2024
b75da36
More early return
arsenm Aug 20, 2024
ded31f9
Formatting
arsenm Aug 20, 2024
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
211 changes: 111 additions & 100 deletions llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -184,6 +184,7 @@ class InferAddressSpaces : public FunctionPass {

class InferAddressSpacesImpl {
AssumptionCache ∾
Function *F = nullptr;
const DominatorTree *DT = nullptr;
const TargetTransformInfo *TTI = nullptr;
const DataLayout *DL = nullptr;
Expand Down Expand Up @@ -212,14 +213,17 @@ class InferAddressSpacesImpl {
const PredicatedAddrSpaceMapTy &PredicatedAS,
SmallVectorImpl<const Use *> *PoisonUsesToFix) const;

void performPointerReplacement(
Value *V, Value *NewV, Use &U, ValueToValueMapTy &ValueWithNewAddrSpace,
SmallVectorImpl<Instruction *> &DeadInstructions) const;

// Changes the flat address expressions in function F to point to specific
// address spaces if InferredAddrSpace says so. Postorder is the postorder of
// all flat expressions in the use-def graph of function F.
bool
rewriteWithNewAddressSpaces(ArrayRef<WeakTrackingVH> Postorder,
const ValueToAddrSpaceMapTy &InferredAddrSpace,
const PredicatedAddrSpaceMapTy &PredicatedAS,
Function *F) const;
bool rewriteWithNewAddressSpaces(
ArrayRef<WeakTrackingVH> Postorder,
const ValueToAddrSpaceMapTy &InferredAddrSpace,
const PredicatedAddrSpaceMapTy &PredicatedAS) const;

void appendsFlatAddressExpressionToPostorderStack(
Value *V, PostorderStackTy &PostorderStack,
Expand Down Expand Up @@ -842,8 +846,9 @@ unsigned InferAddressSpacesImpl::joinAddressSpaces(unsigned AS1,
return (AS1 == AS2) ? AS1 : FlatAddrSpace;
}

bool InferAddressSpacesImpl::run(Function &F) {
DL = &F.getDataLayout();
bool InferAddressSpacesImpl::run(Function &CurFn) {
F = &CurFn;
DL = &F->getDataLayout();

if (AssumeDefaultIsFlatAddressSpace)
FlatAddrSpace = 0;
Expand All @@ -855,7 +860,7 @@ bool InferAddressSpacesImpl::run(Function &F) {
}

// Collects all flat address expressions in postorder.
std::vector<WeakTrackingVH> Postorder = collectFlatAddressExpressions(F);
std::vector<WeakTrackingVH> Postorder = collectFlatAddressExpressions(*F);

// Runs a data-flow analysis to refine the address spaces of every expression
// in Postorder.
Expand All @@ -865,8 +870,8 @@ bool InferAddressSpacesImpl::run(Function &F) {

// Changes the address spaces of the flat address expressions who are inferred
// to point to a specific address space.
return rewriteWithNewAddressSpaces(Postorder, InferredAddrSpace, PredicatedAS,
&F);
return rewriteWithNewAddressSpaces(Postorder, InferredAddrSpace,
PredicatedAS);
}

// Constants need to be tracked through RAUW to handle cases with nested
Expand Down Expand Up @@ -1168,10 +1173,103 @@ static Value::use_iterator skipToNextUser(Value::use_iterator I,
return I;
}

void InferAddressSpacesImpl::performPointerReplacement(
Value *V, Value *NewV, Use &U, ValueToValueMapTy &ValueWithNewAddrSpace,
SmallVectorImpl<Instruction *> &DeadInstructions) const {

User *CurUser = U.getUser();

unsigned AddrSpace = V->getType()->getPointerAddressSpace();
if (replaceIfSimplePointerUse(*TTI, CurUser, AddrSpace, V, NewV))
return;

// Skip if the current user is the new value itself.
if (CurUser == NewV)
return;

auto *CurUserI = dyn_cast<Instruction>(CurUser);
if (!CurUserI || CurUserI->getFunction() != F)
return;

// Handle more complex cases like intrinsic that need to be remangled.
if (auto *MI = dyn_cast<MemIntrinsic>(CurUser)) {
if (!MI->isVolatile() && handleMemIntrinsicPtrUse(MI, V, NewV))
return;
}

if (auto *II = dyn_cast<IntrinsicInst>(CurUser)) {
if (rewriteIntrinsicOperands(II, V, NewV))
return;
}

if (ICmpInst *Cmp = dyn_cast<ICmpInst>(CurUserI)) {
Copy link
Contributor

Choose a reason for hiding this comment

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

nit: could use auto here

// If we can infer that both pointers are in the same addrspace,
// transform e.g.
// %cmp = icmp eq float* %p, %q
// into
// %cmp = icmp eq float addrspace(3)* %new_p, %new_q

unsigned NewAS = NewV->getType()->getPointerAddressSpace();
int SrcIdx = U.getOperandNo();
int OtherIdx = (SrcIdx == 0) ? 1 : 0;
Value *OtherSrc = Cmp->getOperand(OtherIdx);

if (Value *OtherNewV = ValueWithNewAddrSpace.lookup(OtherSrc)) {
if (OtherNewV->getType()->getPointerAddressSpace() == NewAS) {
Cmp->setOperand(OtherIdx, OtherNewV);
Cmp->setOperand(SrcIdx, NewV);
return;
}
}

// Even if the type mismatches, we can cast the constant.
if (auto *KOtherSrc = dyn_cast<Constant>(OtherSrc)) {
if (isSafeToCastConstAddrSpace(KOtherSrc, NewAS)) {
Cmp->setOperand(SrcIdx, NewV);
Cmp->setOperand(OtherIdx, ConstantExpr::getAddrSpaceCast(
KOtherSrc, NewV->getType()));
return;
}
}
}

if (AddrSpaceCastInst *ASC = dyn_cast<AddrSpaceCastInst>(CurUserI)) {
Copy link
Contributor

Choose a reason for hiding this comment

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

ditto

unsigned NewAS = NewV->getType()->getPointerAddressSpace();
if (ASC->getDestAddressSpace() == NewAS) {
ASC->replaceAllUsesWith(NewV);
DeadInstructions.push_back(ASC);
return;
}
}

// Otherwise, replaces the use with flat(NewV).
if (Instruction *VInst = dyn_cast<Instruction>(V)) {
Copy link
Contributor

Choose a reason for hiding this comment

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

ditto

// Don't create a copy of the original addrspacecast.
if (U == V && isa<AddrSpaceCastInst>(V))
return;

// Insert the addrspacecast after NewV.
BasicBlock::iterator InsertPos;
if (Instruction *NewVInst = dyn_cast<Instruction>(NewV))
Copy link
Contributor

Choose a reason for hiding this comment

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

ditto

InsertPos = std::next(NewVInst->getIterator());
else
InsertPos = std::next(VInst->getIterator());

while (isa<PHINode>(InsertPos))
++InsertPos;
// This instruction may contain multiple uses of V, update them all.
CurUser->replaceUsesOfWith(
V, new AddrSpaceCastInst(NewV, V->getType(), "", InsertPos));
} else {
CurUserI->replaceUsesOfWith(
V, ConstantExpr::getAddrSpaceCast(cast<Constant>(NewV), V->getType()));
}
}

bool InferAddressSpacesImpl::rewriteWithNewAddressSpaces(
ArrayRef<WeakTrackingVH> Postorder,
const ValueToAddrSpaceMapTy &InferredAddrSpace,
const PredicatedAddrSpaceMapTy &PredicatedAS, Function *F) const {
const PredicatedAddrSpaceMapTy &PredicatedAS) const {
// For each address expression to be modified, creates a clone of it with its
// pointer operands converted to the new address space. Since the pointer
// operands are converted, the clone is naturally in the new address space by
Expand Down Expand Up @@ -1262,100 +1360,13 @@ bool InferAddressSpacesImpl::rewriteWithNewAddressSpaces(
Value::use_iterator I, E, Next;
for (I = V->use_begin(), E = V->use_end(); I != E;) {
Use &U = *I;
User *CurUser = U.getUser();

// Some users may see the same pointer operand in multiple operands. Skip
// to the next instruction.
I = skipToNextUser(I, E);

unsigned AddrSpace = V->getType()->getPointerAddressSpace();
if (replaceIfSimplePointerUse(*TTI, CurUser, AddrSpace, V, NewV))
continue;

// Skip if the current user is the new value itself.
if (CurUser == NewV)
continue;

if (auto *CurUserI = dyn_cast<Instruction>(CurUser);
CurUserI && CurUserI->getFunction() != F)
continue;

// Handle more complex cases like intrinsic that need to be remangled.
if (auto *MI = dyn_cast<MemIntrinsic>(CurUser)) {
if (!MI->isVolatile() && handleMemIntrinsicPtrUse(MI, V, NewV))
continue;
}

if (auto *II = dyn_cast<IntrinsicInst>(CurUser)) {
if (rewriteIntrinsicOperands(II, V, NewV))
continue;
}

if (isa<Instruction>(CurUser)) {
if (ICmpInst *Cmp = dyn_cast<ICmpInst>(CurUser)) {
// If we can infer that both pointers are in the same addrspace,
// transform e.g.
// %cmp = icmp eq float* %p, %q
// into
// %cmp = icmp eq float addrspace(3)* %new_p, %new_q

unsigned NewAS = NewV->getType()->getPointerAddressSpace();
int SrcIdx = U.getOperandNo();
int OtherIdx = (SrcIdx == 0) ? 1 : 0;
Value *OtherSrc = Cmp->getOperand(OtherIdx);

if (Value *OtherNewV = ValueWithNewAddrSpace.lookup(OtherSrc)) {
if (OtherNewV->getType()->getPointerAddressSpace() == NewAS) {
Cmp->setOperand(OtherIdx, OtherNewV);
Cmp->setOperand(SrcIdx, NewV);
continue;
}
}

// Even if the type mismatches, we can cast the constant.
if (auto *KOtherSrc = dyn_cast<Constant>(OtherSrc)) {
if (isSafeToCastConstAddrSpace(KOtherSrc, NewAS)) {
Cmp->setOperand(SrcIdx, NewV);
Cmp->setOperand(OtherIdx, ConstantExpr::getAddrSpaceCast(
KOtherSrc, NewV->getType()));
continue;
}
}
}

if (AddrSpaceCastInst *ASC = dyn_cast<AddrSpaceCastInst>(CurUser)) {
unsigned NewAS = NewV->getType()->getPointerAddressSpace();
if (ASC->getDestAddressSpace() == NewAS) {
ASC->replaceAllUsesWith(NewV);
DeadInstructions.push_back(ASC);
continue;
}
}

// Otherwise, replaces the use with flat(NewV).
if (Instruction *VInst = dyn_cast<Instruction>(V)) {
// Don't create a copy of the original addrspacecast.
if (U == V && isa<AddrSpaceCastInst>(V))
continue;

// Insert the addrspacecast after NewV.
BasicBlock::iterator InsertPos;
if (Instruction *NewVInst = dyn_cast<Instruction>(NewV))
InsertPos = std::next(NewVInst->getIterator());
else
InsertPos = std::next(VInst->getIterator());

while (isa<PHINode>(InsertPos))
++InsertPos;
// This instruction may contain multiple uses of V, update them all.
CurUser->replaceUsesOfWith(
V, new AddrSpaceCastInst(NewV, V->getType(), "", InsertPos));
} else {
CurUser->replaceUsesOfWith(
V, ConstantExpr::getAddrSpaceCast(cast<Constant>(NewV),
V->getType()));
}
}
performPointerReplacement(V, NewV, U, ValueWithNewAddrSpace,
DeadInstructions);
}

if (V->use_empty()) {
Expand Down
Loading