[llvm] 6dc519b - [AMDGPU] Fix value swap for potentially aliased out arguments in RewriteOutArguments (#202922)
via llvm-commits
llvm-commits at lists.llvm.org
Wed Jun 10 05:54:21 PDT 2026
Author: Arseniy Obolenskiy
Date: 2026-06-10T14:54:16+02:00
New Revision: 6dc519bbfc956da873ab527d6d5e3a214b902496
URL: https://github.com/llvm/llvm-project/commit/6dc519bbfc956da873ab527d6d5e3a214b902496
DIFF: https://github.com/llvm/llvm-project/commit/6dc519bbfc956da873ab527d6d5e3a214b902496.diff
LOG: [AMDGPU] Fix value swap for potentially aliased out arguments in RewriteOutArguments (#202922)
When two out argument pointers may alias, MemoryDependence returns the
last aliasing store for both, so each argument was paired with the other
stored value
Match each store pointer to its argument and store them in pairs instead
Added:
Modified:
llvm/lib/Target/AMDGPU/AMDGPURewriteOutArguments.cpp
llvm/test/CodeGen/AMDGPU/rewrite-out-arguments.ll
Removed:
################################################################################
diff --git a/llvm/lib/Target/AMDGPU/AMDGPURewriteOutArguments.cpp b/llvm/lib/Target/AMDGPU/AMDGPURewriteOutArguments.cpp
index a2e16c7f873f7..c0802de2763dc 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPURewriteOutArguments.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPURewriteOutArguments.cpp
@@ -182,7 +182,10 @@ bool AMDGPURewriteOutArguments::runOnFunction(Function &F) {
MDA = &getAnalysis<MemoryDependenceWrapperPass>().getMemDep();
unsigned ReturnNumRegs = 0;
- SmallDenseMap<int, Type *, 4> OutArgIndexes;
+ // Maps an out-argument number to its field index in the return struct.
+ // Fields are in processing order, which the retry loop below can reorder
+ // relative to argument order, so the index must be tracked, not recomputed.
+ SmallDenseMap<unsigned, unsigned, 4> OutArgIndexes;
SmallVector<Type *, 4> ReturnTypes;
Type *RetTy = F.getReturnType();
if (!RetTy->isVoidTy()) {
@@ -226,11 +229,10 @@ bool AMDGPURewriteOutArguments::runOnFunction(Function &F) {
// Keep retrying if we are able to successfully eliminate an argument. This
// helps with cases with multiple arguments which may alias, such as in a
- // sincos implementation. If we have 2 stores to arguments, on the first
- // attempt the MDA query will succeed for the second store but not the
- // first. On the second iteration we've removed that out clobbering argument
- // (by effectively moving it into another function) and will find the second
- // argument is OK to move.
+ // sincos implementation. With 2 stores to may-aliasing arguments, MDA
+ // returns the second store for the first argument too; the identity guard
+ // below rejects it, and a later iteration folds the first argument once the
+ // second store has been removed.
for (const auto &Pair : OutArgs) {
bool ThisReplaceable = true;
SmallVector<std::pair<ReturnInst *, StoreInst *>, 4> ReplaceableStores;
@@ -258,6 +260,11 @@ bool AMDGPURewriteOutArguments::runOnFunction(Function &F) {
if (Q.isDef())
SI = dyn_cast<StoreInst>(Q.getInst());
+ // MDA stops at the first may-aliasing store, which need not be to this
+ // argument; only fold a store whose pointer is exactly OutArg.
+ if (SI && SI->getPointerOperand() != OutArg)
+ SI = nullptr;
+
if (SI) {
LLVM_DEBUG(dbgs() << "Found out argument store: " << *SI << '\n');
ReplaceableStores.emplace_back(RI, SI);
@@ -288,8 +295,8 @@ bool AMDGPURewriteOutArguments::runOnFunction(Function &F) {
}
if (ThisReplaceable) {
+ OutArgIndexes.insert({OutArg->getArgNo(), ReturnTypes.size()});
ReturnTypes.push_back(ArgTy);
- OutArgIndexes.insert({OutArg->getArgNo(), ArgTy});
++NumOutArgumentsReplaced;
Changing = true;
}
@@ -334,15 +341,17 @@ bool AMDGPURewriteOutArguments::runOnFunction(Function &F) {
IRBuilder<> B(RI);
B.SetCurrentDebugLocation(RI->getDebugLoc());
- int RetIdx = 0;
Value *NewRetVal = PoisonValue::get(NewRetTy);
Value *RetVal = RI->getReturnValue();
if (RetVal)
- NewRetVal = B.CreateInsertValue(NewRetVal, RetVal, RetIdx++);
+ NewRetVal = B.CreateInsertValue(NewRetVal, RetVal, 0);
- for (std::pair<Argument *, Value *> ReturnPoint : Replacement.second)
- NewRetVal = B.CreateInsertValue(NewRetVal, ReturnPoint.second, RetIdx++);
+ // Use OutArgIndexes so body and stub agree on the field for each argument.
+ for (std::pair<Argument *, Value *> ReturnPoint : Replacement.second) {
+ unsigned FieldIdx = OutArgIndexes.lookup(ReturnPoint.first->getArgNo());
+ NewRetVal = B.CreateInsertValue(NewRetVal, ReturnPoint.second, FieldIdx);
+ }
if (RetVal)
RI->setOperand(0, NewRetVal);
@@ -367,17 +376,17 @@ bool AMDGPURewriteOutArguments::runOnFunction(Function &F) {
IRBuilder<> B(StubBB);
CallInst *StubCall = B.CreateCall(NewFunc, StubCallArgs);
- int RetIdx = RetTy->isVoidTy() ? 0 : 1;
for (Argument &Arg : F.args()) {
auto It = OutArgIndexes.find(Arg.getArgNo());
if (It == OutArgIndexes.end())
continue;
- Type *EltTy = It->second;
+ unsigned FieldIdx = It->second;
+ Type *EltTy = NewRetTy->getElementType(FieldIdx);
const auto Align =
DL->getValueOrABITypeAlignment(Arg.getParamAlign(), EltTy);
- Value *Val = B.CreateExtractValue(StubCall, RetIdx++);
+ Value *Val = B.CreateExtractValue(StubCall, FieldIdx);
B.CreateAlignedStore(Val, &Arg, Align);
}
diff --git a/llvm/test/CodeGen/AMDGPU/rewrite-out-arguments.ll b/llvm/test/CodeGen/AMDGPU/rewrite-out-arguments.ll
index 5b48005aaaeb6..fc0ce680b8cf5 100644
--- a/llvm/test/CodeGen/AMDGPU/rewrite-out-arguments.ll
+++ b/llvm/test/CodeGen/AMDGPU/rewrite-out-arguments.ll
@@ -677,9 +677,9 @@ attributes #2 = { alwaysinline nounwind }
; CHECK-LABEL: define {{[^@]+}}@multiple_same_return_mayalias
; CHECK-SAME: (ptr addrspace(5) [[TMP0:%.*]], ptr addrspace(5) [[TMP1:%.*]]) #[[ATTR2]] {
; CHECK-NEXT: [[TMP3:%.*]] = call [[MULTIPLE_SAME_RETURN_MAYALIAS:%.*]] @[[MULTIPLE_SAME_RETURN_MAYALIAS_BODY:[a-zA-Z0-9_$\"\\.-]*[a-zA-Z_$\"\\.-][a-zA-Z0-9_$\"\\.-]*]](ptr addrspace(5) poison, ptr addrspace(5) poison)
-; CHECK-NEXT: [[TMP4:%.*]] = extractvalue [[MULTIPLE_SAME_RETURN_MAYALIAS]] [[TMP3]], 0
+; CHECK-NEXT: [[TMP4:%.*]] = extractvalue [[MULTIPLE_SAME_RETURN_MAYALIAS]] [[TMP3]], 1
; CHECK-NEXT: store i32 [[TMP4]], ptr addrspace(5) [[TMP0]], align 4
-; CHECK-NEXT: [[TMP5:%.*]] = extractvalue [[MULTIPLE_SAME_RETURN_MAYALIAS]] [[TMP3]], 1
+; CHECK-NEXT: [[TMP5:%.*]] = extractvalue [[MULTIPLE_SAME_RETURN_MAYALIAS]] [[TMP3]], 0
; CHECK-NEXT: store i32 [[TMP5]], ptr addrspace(5) [[TMP1]], align 4
; CHECK-NEXT: ret void
;
@@ -861,9 +861,9 @@ attributes #2 = { alwaysinline nounwind }
; CHECK-LABEL: define {{[^@]+}}@num_regs_reach_limit_leftover
; CHECK-SAME: (ptr addrspace(5) [[TMP0:%.*]], ptr addrspace(5) [[TMP1:%.*]], i32 [[TMP2:%.*]]) #[[ATTR2]] {
; CHECK-NEXT: [[TMP4:%.*]] = call [[NUM_REGS_REACH_LIMIT_LEFTOVER:%.*]] @[[NUM_REGS_REACH_LIMIT_LEFTOVER_BODY:[a-zA-Z0-9_$\"\\.-]*[a-zA-Z_$\"\\.-][a-zA-Z0-9_$\"\\.-]*]](ptr addrspace(5) poison, ptr addrspace(5) poison, i32 [[TMP2]])
-; CHECK-NEXT: [[TMP5:%.*]] = extractvalue [[NUM_REGS_REACH_LIMIT_LEFTOVER]] [[TMP4]], 1
+; CHECK-NEXT: [[TMP5:%.*]] = extractvalue [[NUM_REGS_REACH_LIMIT_LEFTOVER]] [[TMP4]], 2
; CHECK-NEXT: store i32 [[TMP5]], ptr addrspace(5) [[TMP0]], align 4
-; CHECK-NEXT: [[TMP6:%.*]] = extractvalue [[NUM_REGS_REACH_LIMIT_LEFTOVER]] [[TMP4]], 2
+; CHECK-NEXT: [[TMP6:%.*]] = extractvalue [[NUM_REGS_REACH_LIMIT_LEFTOVER]] [[TMP4]], 1
; CHECK-NEXT: store i32 [[TMP6]], ptr addrspace(5) [[TMP1]], align 4
; CHECK-NEXT: [[TMP7:%.*]] = extractvalue [[NUM_REGS_REACH_LIMIT_LEFTOVER]] [[TMP4]], 0
; CHECK-NEXT: ret [15 x i32] [[TMP7]]
More information about the llvm-commits
mailing list