[llvm] [AMDGPU] Fix LDS access via flat pointer argument in amdgpu-sw-lower-lds (PR #209842)
via llvm-commits
llvm-commits at lists.llvm.org
Fri Jul 31 04:39:32 PDT 2026
https://github.com/skc7 updated https://github.com/llvm/llvm-project/pull/209842
>From c256325bb70344fd067978a14ce77aa8f7d04af7 Mon Sep 17 00:00:00 2001
From: skc7 <Krishna.Sankisa at amd.com>
Date: Wed, 15 Jul 2026 22:57:15 +0530
Subject: [PATCH 1/4] [AMDGPU] Fix LDS access via flat pointer argument in
amdgpu-sw-lower-lds
---
llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp | 82 +++++++++++++++++++
.../amdgpu-sw-lower-lds-flat-ptr-arg-asan.ll | 73 +++++++++++++++++
2 files changed, 155 insertions(+)
create mode 100644 llvm/test/CodeGen/AMDGPU/amdgpu-sw-lower-lds-flat-ptr-arg-asan.ll
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp b/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp
index 2b78094837c68..e2668272ea29c 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp
@@ -94,6 +94,7 @@
#include "llvm/ADT/StringRef.h"
#include "llvm/Analysis/CallGraph.h"
#include "llvm/Analysis/DomTreeUpdater.h"
+#include "llvm/Analysis/ValueTracking.h"
#include "llvm/CodeGen/TargetPassConfig.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/DIBuilder.h"
@@ -204,6 +205,7 @@ class AMDGPUSwLowerLDS {
void lowerNonKernelLDSAccesses(Function *Func,
SetVector<GlobalVariable *> &LDSGlobals,
NonKernelLDSParameters &NKLDSParams);
+ void lowerFlatToLocalRoundTrips(Function *Func);
void
updateMallocSizeForDynamicLDS(Function *Func, Value **CurrMallocSize,
Value *HiddenDynLDSSize,
@@ -790,6 +792,80 @@ void AMDGPUSwLowerLDS::translateLDSMemoryOperationsToGlobalMemory(
}
}
+void AMDGPUSwLowerLDS::lowerFlatToLocalRoundTrips(Function *Func) {
+ // After SW LDS lowering, storage that used to live in LDS now lives in global
+ // memory, but it may still be handed to a callee as a flat (generic) pointer.
+ // Such a callee can re-derive an LDS pointer with a round-trip pattern:
+ // %p = getelementptr ..., ptr %flat_arg ; flat, backed by global mem
+ // %c = addrspacecast ptr %p to ptr addrspace(3)
+ // load/store ... ptr addrspace(3) %c
+ // The addrspace(3) access would then hit dead LDS. Redo the access through
+ // the flat source pointer, which still points at the global backing buffer.
+ SmallVector<Instruction *> ToErase;
+ SmallPtrSet<AddrSpaceCastInst *, 8> Casts;
+ for (BasicBlock &BB : *Func) {
+ for (Instruction &Inst : BB) {
+ Value *Ptr = nullptr;
+ if (auto *LI = dyn_cast<LoadInst>(&Inst))
+ Ptr = LI->getPointerOperand();
+ else if (auto *SI = dyn_cast<StoreInst>(&Inst))
+ Ptr = SI->getPointerOperand();
+ else if (auto *RMW = dyn_cast<AtomicRMWInst>(&Inst))
+ Ptr = RMW->getPointerOperand();
+ else if (auto *XCHG = dyn_cast<AtomicCmpXchgInst>(&Inst))
+ Ptr = XCHG->getPointerOperand();
+ else
+ continue;
+
+ auto *ASC = dyn_cast<AddrSpaceCastInst>(Ptr);
+ if (!ASC || ASC->getSrcAddressSpace() != AMDGPUAS::FLAT_ADDRESS ||
+ ASC->getDestAddressSpace() != AMDGPUAS::LOCAL_ADDRESS)
+ continue;
+
+ // Only collapse pointers that originate from a flat function argument,
+ // i.e. storage passed in from a caller. Genuine in-function LDS pointers
+ // are lowered through the base/offset table machinery instead.
+ Value *FlatPtr = ASC->getPointerOperand();
+ if (!isa<Argument>(getUnderlyingObject(FlatPtr)))
+ continue;
+
+ IRB.SetInsertPoint(&Inst);
+ Instruction *NewInst = nullptr;
+ if (auto *LI = dyn_cast<LoadInst>(&Inst)) {
+ NewInst = IRB.CreateLoad(LI->getType(), FlatPtr, LI->getProperties());
+ } else if (auto *SI = dyn_cast<StoreInst>(&Inst)) {
+ NewInst = IRB.CreateStore(SI->getValueOperand(), FlatPtr,
+ SI->getProperties());
+ } else if (auto *RMW = dyn_cast<AtomicRMWInst>(&Inst)) {
+ auto *NewRMW = IRB.CreateAtomicRMW(
+ RMW->getOperation(), FlatPtr, RMW->getValOperand(), RMW->getAlign(),
+ RMW->getOrdering(), RMW->getSyncScopeID());
+ NewRMW->setVolatile(RMW->isVolatile());
+ NewInst = NewRMW;
+ } else {
+ auto *XCHG = cast<AtomicCmpXchgInst>(&Inst);
+ auto *NewXCHG = IRB.CreateAtomicCmpXchg(
+ FlatPtr, XCHG->getCompareOperand(), XCHG->getNewValOperand(),
+ XCHG->getAlign(), XCHG->getSuccessOrdering(),
+ XCHG->getFailureOrdering(), XCHG->getSyncScopeID());
+ NewXCHG->setVolatile(XCHG->isVolatile());
+ NewInst = NewXCHG;
+ }
+ // Flat access into the global backing buffer must be sanitized.
+ AsanInfo.Instructions.insert(NewInst);
+ Inst.replaceAllUsesWith(NewInst);
+ ToErase.push_back(&Inst);
+ Casts.insert(ASC);
+ }
+ }
+ for (Instruction *I : ToErase)
+ I->eraseFromParent();
+ // Drop the now-dead round-trip casts.
+ for (AddrSpaceCastInst *ASC : Casts)
+ if (ASC->use_empty())
+ ASC->eraseFromParent();
+}
+
void AMDGPUSwLowerLDS::poisonRedzones(Function *Func, Value *MallocPtr) {
auto &LDSParams = FuncLDSAccessInfo.KernelToLDSParametersMap[Func];
Type *Int64Ty = IRB.getInt64Ty();
@@ -1323,6 +1399,12 @@ bool AMDGPUSwLowerLDS::run() {
if (!Changed)
return Changed;
+ // Fix up callees that still access lowered storage through the
+ // flat-arg -> addrspace(3) round-trip pattern.
+ for (Function &F : M)
+ if (!F.isDeclaration())
+ lowerFlatToLocalRoundTrips(&F);
+
for (auto &GV : make_early_inc_range(M.globals())) {
if (AMDGPU::isLDSVariableToLower(GV)) {
// probably want to remove from used lists
diff --git a/llvm/test/CodeGen/AMDGPU/amdgpu-sw-lower-lds-flat-ptr-arg-asan.ll b/llvm/test/CodeGen/AMDGPU/amdgpu-sw-lower-lds-flat-ptr-arg-asan.ll
new file mode 100644
index 0000000000000..646371a5bde24
--- /dev/null
+++ b/llvm/test/CodeGen/AMDGPU/amdgpu-sw-lower-lds-flat-ptr-arg-asan.ll
@@ -0,0 +1,73 @@
+; RUN: opt < %s -passes=amdgpu-sw-lower-lds -amdgpu-asan-instrument-lds=false -S -mtriple=amdgcn-amd-amdhsa | FileCheck %s
+
+; A non-kernel receives lowered LDS storage as a flat (generic) pointer and
+; re-derives an LDS pointer via a flat -> addrspace(3) round-trip. After SW LDS
+; lowering the storage lives in global memory, so the addrspace(3) access would
+; hit dead LDS. The round-trip must be collapsed into a flat access on the
+; incoming argument, which still points at the global backing buffer. Check the
+; store, load, atomicrmw and cmpxchg cases.
+
+ at lds_var = internal addrspace(3) global [64 x i32] poison, align 4
+
+define void @store_case(ptr %storage) sanitize_address {
+; CHECK-LABEL: define void @store_case(
+; CHECK-SAME: ptr [[STORAGE:%.*]]){{.*}} {
+; CHECK-NOT: addrspace(3)
+; CHECK: store i32 42, ptr [[STORAGE]], align 4
+; CHECK-NOT: addrspace(3)
+; CHECK: ret void
+ %cast = addrspacecast ptr %storage to ptr addrspace(3)
+ store i32 42, ptr addrspace(3) %cast, align 4
+ ret void
+}
+
+define i32 @load_case(ptr %storage) sanitize_address {
+; CHECK-LABEL: define i32 @load_case(
+; CHECK-SAME: ptr [[STORAGE:%.*]]){{.*}} {
+; CHECK-NOT: addrspace(3)
+; CHECK: [[VAL:%.*]] = load i32, ptr [[STORAGE]], align 4
+; CHECK-NOT: addrspace(3)
+; CHECK: ret i32 [[VAL]]
+ %cast = addrspacecast ptr %storage to ptr addrspace(3)
+ %val = load i32, ptr addrspace(3) %cast, align 4
+ ret i32 %val
+}
+
+define i32 @atomicrmw_case(ptr %storage) sanitize_address {
+; CHECK-LABEL: define i32 @atomicrmw_case(
+; CHECK-SAME: ptr [[STORAGE:%.*]]){{.*}} {
+; CHECK-NOT: addrspace(3)
+; CHECK: [[OLD:%.*]] = atomicrmw add ptr [[STORAGE]], i32 1 seq_cst, align 4
+; CHECK-NOT: addrspace(3)
+; CHECK: ret i32 [[OLD]]
+ %cast = addrspacecast ptr %storage to ptr addrspace(3)
+ %old = atomicrmw add ptr addrspace(3) %cast, i32 1 seq_cst, align 4
+ ret i32 %old
+}
+
+define i32 @cmpxchg_case(ptr %storage) sanitize_address {
+; CHECK-LABEL: define i32 @cmpxchg_case(
+; CHECK-SAME: ptr [[STORAGE:%.*]]){{.*}} {
+; CHECK-NOT: addrspace(3)
+; CHECK: [[RES:%.*]] = cmpxchg ptr [[STORAGE]], i32 0, i32 1 seq_cst seq_cst, align 4
+; CHECK-NOT: addrspace(3)
+; CHECK: [[VAL:%.*]] = extractvalue { i32, i1 } [[RES]], 0
+; CHECK: ret i32 [[VAL]]
+ %cast = addrspacecast ptr %storage to ptr addrspace(3)
+ %res = cmpxchg ptr addrspace(3) %cast, i32 0, i32 1 seq_cst seq_cst, align 4
+ %val = extractvalue { i32, i1 } %res, 0
+ ret i32 %val
+}
+
+define amdgpu_kernel void @kernel() sanitize_address {
+; CHECK-LABEL: define amdgpu_kernel void @kernel(
+ %flat = addrspacecast ptr addrspace(3) @lds_var to ptr
+ call void @store_case(ptr %flat)
+ %l = call i32 @load_case(ptr %flat)
+ %a = call i32 @atomicrmw_case(ptr %flat)
+ %c = call i32 @cmpxchg_case(ptr %flat)
+ ret void
+}
+
+!llvm.module.flags = !{!0}
+!0 = !{i32 4, !"nosanitize_address", i32 1}
>From 25d57b23b33ff21483d4ec5e3c322f02bf1a22c5 Mon Sep 17 00:00:00 2001
From: skc7 <Krishna.Sankisa at amd.com>
Date: Fri, 17 Jul 2026 21:47:54 +0530
Subject: [PATCH 2/4] update
---
llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp | 325 +++++++++++++-----
.../amdgpu-sw-lower-lds-flat-ptr-arg-asan.ll | 243 +++++++++++--
2 files changed, 444 insertions(+), 124 deletions(-)
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp b/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp
index e2668272ea29c..25c82883bd419 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp
@@ -81,6 +81,15 @@
// The Offset into the base "SW LDS" is obtained from
// corresponding element in offset table. With this information, replacement
// value is obtained.
+//
+// Replacement of LDS accessed via flat pointer arguments:
+// An LDS pointer can reach a non-kernel as a flat (generic) pointer
+// argument, e.g. after an addrspacecast of an LDS global to flat. A
+// fixed-point call-graph walk identifies the functions and argument
+// indices that carry such LDS-derived flat pointers. Memory operations
+// that reach LDS through these arguments are redirected to the relocated
+// global memory by reconstructing an equivalent flat pointer, instead of
+// accessing the now-dead hardware LDS.
//===----------------------------------------------------------------------===//
#include "AMDGPU.h"
@@ -94,7 +103,6 @@
#include "llvm/ADT/StringRef.h"
#include "llvm/Analysis/CallGraph.h"
#include "llvm/Analysis/DomTreeUpdater.h"
-#include "llvm/Analysis/ValueTracking.h"
#include "llvm/CodeGen/TargetPassConfig.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/DIBuilder.h"
@@ -104,10 +112,13 @@
#include "llvm/IR/Instructions.h"
#include "llvm/IR/IntrinsicsAMDGPU.h"
#include "llvm/IR/MDBuilder.h"
+#include "llvm/IR/Operator.h"
#include "llvm/IR/ReplaceConstant.h"
+#include "llvm/IR/ValueHandle.h"
#include "llvm/Pass.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Transforms/Instrumentation/AddressSanitizerCommon.h"
+#include "llvm/Transforms/Utils/Local.h"
#include "llvm/Transforms/Utils/ModuleUtils.h"
#include <algorithm>
@@ -170,6 +181,7 @@ struct FunctionsAndLDSAccess {
SetVector<Function *> NonKernelsWithLDSArgument;
SetVector<GlobalVariable *> AllNonKernelLDSAccess;
FunctionVariableMap NonKernelToLDSAccessMap;
+ DenseMap<Function *, SmallPtrSet<Argument *, 4>> NonKernelsWithLDSFlatArg;
};
class AMDGPUSwLowerLDS {
@@ -179,6 +191,7 @@ class AMDGPUSwLowerLDS {
bool run();
void getUsesOfLDSByNonKernels();
void getNonKernelsWithLDSArguments(const CallGraph &CG);
+ void getNonKernelsWithLDSFlatArguments();
SetVector<Function *>
getOrderedIndirectLDSAccessingKernels(SetVector<Function *> &Kernels);
SetVector<GlobalVariable *>
@@ -205,7 +218,8 @@ class AMDGPUSwLowerLDS {
void lowerNonKernelLDSAccesses(Function *Func,
SetVector<GlobalVariable *> &LDSGlobals,
NonKernelLDSParameters &NKLDSParams);
- void lowerFlatToLocalRoundTrips(Function *Func);
+ void lowerNonKernelLDSFlatArgAccesses(Function *Func);
+ Value *getFlatPtrForRoundTripLDSAccess(Function *Func, Value *LDSPtr);
void
updateMallocSizeForDynamicLDS(Function *Func, Value **CurrMallocSize,
Value *HiddenDynLDSSize,
@@ -289,6 +303,99 @@ void AMDGPUSwLowerLDS::getNonKernelsWithLDSArguments(const CallGraph &CG) {
}
}
+// True if flat pointer V is derived from an LDS-carrying argument in LDSArgs,
+// looking through GEP/bitcast/phi/select; Visited breaks phi cycles. When
+// AcceptLocalCast is set, a flat<-local addrspacecast also counts as an origin.
+// Detection (getNonKernelsWithLDSFlatArguments) sets it, since that cast is how
+// LDS first enters a flat pointer; the rewrite
+// (getFlatPtrForRoundTripLDSAccess) leaves it clear so only argument origins
+// qualify.
+static bool flatPtrDerivesFromLDS(Value *V,
+ const SmallPtrSetImpl<Argument *> &LDSArgs,
+ SmallPtrSetImpl<Value *> &Visited,
+ bool AcceptLocalCast) {
+ if (!V->getType()->isPointerTy() || !Visited.insert(V).second)
+ return false;
+ if (auto *A = dyn_cast<Argument>(V))
+ return LDSArgs.contains(A);
+ auto *Op = dyn_cast<Operator>(V);
+ if (!Op)
+ return false;
+ switch (Op->getOpcode()) {
+ case Instruction::AddrSpaceCast: {
+ if (!AcceptLocalCast)
+ return false;
+ Value *Src = Op->getOperand(0);
+ if (Src->getType()->getPointerAddressSpace() == AMDGPUAS::LOCAL_ADDRESS)
+ return true;
+ return flatPtrDerivesFromLDS(Src, LDSArgs, Visited, AcceptLocalCast);
+ }
+ case Instruction::GetElementPtr:
+ case Instruction::BitCast:
+ return flatPtrDerivesFromLDS(Op->getOperand(0), LDSArgs, Visited,
+ AcceptLocalCast);
+ case Instruction::PHI:
+ for (Value *In : cast<PHINode>(Op)->incoming_values())
+ if (flatPtrDerivesFromLDS(In, LDSArgs, Visited, AcceptLocalCast))
+ return true;
+ return false;
+ case Instruction::Select: {
+ auto *SI = cast<SelectInst>(Op);
+ return flatPtrDerivesFromLDS(SI->getTrueValue(), LDSArgs, Visited,
+ AcceptLocalCast) ||
+ flatPtrDerivesFromLDS(SI->getFalseValue(), LDSArgs, Visited,
+ AcceptLocalCast);
+ }
+ default:
+ return false;
+ }
+}
+
+void AMDGPUSwLowerLDS::getNonKernelsWithLDSFlatArguments() {
+ // A kernel or non-kernel may pass lowered LDS storage to a non-kernel as a
+ // flat pointer instead of addrspace(3). Record, per callee, which flat
+ // parameters carry LDS. Must run before lowering rewrites the call-site
+ // casts.
+ auto &FlatArgMap = FuncLDSAccessInfo.NonKernelsWithLDSFlatArg;
+ bool Changed = true;
+ while (Changed) {
+ Changed = false;
+ for (Function &F : M) {
+ if (F.isDeclaration())
+ continue;
+ SmallPtrSet<Argument *, 4> CallerLDSArgs;
+ if (auto It = FlatArgMap.find(&F); It != FlatArgMap.end())
+ CallerLDSArgs = It->second;
+ for (BasicBlock &BB : F) {
+ for (Instruction &I : BB) {
+ auto *CB = dyn_cast<CallBase>(&I);
+ if (!CB)
+ continue;
+ Function *Callee = CB->getCalledFunction();
+ if (!Callee || Callee->isDeclaration() || AMDGPU::isKernel(*Callee))
+ continue;
+ unsigned NumArgs =
+ std::min(CB->arg_size(), (unsigned)Callee->arg_size());
+ for (unsigned ArgNo = 0; ArgNo < NumArgs; ++ArgNo) {
+ Argument *CalleeArg = Callee->getArg(ArgNo);
+ Type *ArgTy = CalleeArg->getType();
+ if (!ArgTy->isPointerTy() ||
+ ArgTy->getPointerAddressSpace() != AMDGPUAS::FLAT_ADDRESS)
+ continue;
+ Value *Actual = CB->getArgOperand(ArgNo);
+ SmallPtrSet<Value *, 8> Visited;
+ if (!flatPtrDerivesFromLDS(Actual, CallerLDSArgs, Visited,
+ /*AcceptLocalCast=*/true))
+ continue;
+ if (FlatArgMap[Callee].insert(CalleeArg).second)
+ Changed = true;
+ }
+ }
+ }
+ }
+ }
+}
+
void AMDGPUSwLowerLDS::getUsesOfLDSByNonKernels() {
for (GlobalVariable *GV : FuncLDSAccessInfo.AllNonKernelLDSAccess) {
if (!AMDGPU::isLDSVariableToLower(*GV))
@@ -690,31 +797,51 @@ void AMDGPUSwLowerLDS::translateLDSMemoryOperationsToGlobalMemory(
SetVector<Instruction *> &LDSInstructions) {
LLVM_DEBUG(dbgs() << "Translating LDS memory operations to global memory : "
<< Func->getName());
+ // Map an LDS pointer to its global-memory equivalent: a flat-argument access
+ // uses the flat source pointer, otherwise the base/offset table. nullptr
+ // means neither applies, so the operation is left unchanged.
+ auto TranslatePtr = [&](Value *LDSPtr) -> Value * {
+ if (Value *Flat = getFlatPtrForRoundTripLDSAccess(Func, LDSPtr))
+ return Flat;
+ if (!LoadMallocPtr)
+ return nullptr;
+ return getTranslatedGlobalMemoryPtrOfLDS(LoadMallocPtr, LDSPtr);
+ };
+ // Old LDS pointer operands to clean up. Flat-arg round-trips leave the cast
+ // and local GEPs dead; the base/offset path keeps them live via ptrtoint, so
+ // cleanup is a no-op there. WeakTrackingVH tolerates chains freed
+ // recursively.
+ SmallVector<WeakTrackingVH, 8> MaybeDeadPtrs;
for (Instruction *Inst : LDSInstructions) {
IRB.SetInsertPoint(Inst);
if (LoadInst *LI = dyn_cast<LoadInst>(Inst)) {
Value *LIOperand = LI->getPointerOperand();
- Value *Replacement =
- getTranslatedGlobalMemoryPtrOfLDS(LoadMallocPtr, LIOperand);
+ Value *Replacement = TranslatePtr(LIOperand);
+ if (!Replacement)
+ continue;
LoadInst *NewLI =
IRB.CreateLoad(LI->getType(), Replacement, LI->getProperties());
AsanInfo.Instructions.insert(NewLI);
LI->replaceAllUsesWith(NewLI);
LI->eraseFromParent();
+ MaybeDeadPtrs.push_back(LIOperand);
} else if (StoreInst *SI = dyn_cast<StoreInst>(Inst)) {
Value *SIOperand = SI->getPointerOperand();
- Value *Replacement =
- getTranslatedGlobalMemoryPtrOfLDS(LoadMallocPtr, SIOperand);
+ Value *Replacement = TranslatePtr(SIOperand);
+ if (!Replacement)
+ continue;
StoreInst *NewSI = IRB.CreateStore(SI->getValueOperand(), Replacement,
SI->getProperties());
AsanInfo.Instructions.insert(NewSI);
SI->replaceAllUsesWith(NewSI);
SI->eraseFromParent();
+ MaybeDeadPtrs.push_back(SIOperand);
} else if (AtomicRMWInst *RMW = dyn_cast<AtomicRMWInst>(Inst)) {
Value *RMWPtrOperand = RMW->getPointerOperand();
Value *RMWValOperand = RMW->getValOperand();
- Value *Replacement =
- getTranslatedGlobalMemoryPtrOfLDS(LoadMallocPtr, RMWPtrOperand);
+ Value *Replacement = TranslatePtr(RMWPtrOperand);
+ if (!Replacement)
+ continue;
AtomicRMWInst *NewRMW = IRB.CreateAtomicRMW(
RMW->getOperation(), Replacement, RMWValOperand, RMW->getAlign(),
RMW->getOrdering(), RMW->getSyncScopeID());
@@ -722,10 +849,12 @@ void AMDGPUSwLowerLDS::translateLDSMemoryOperationsToGlobalMemory(
AsanInfo.Instructions.insert(NewRMW);
RMW->replaceAllUsesWith(NewRMW);
RMW->eraseFromParent();
+ MaybeDeadPtrs.push_back(RMWPtrOperand);
} else if (AtomicCmpXchgInst *XCHG = dyn_cast<AtomicCmpXchgInst>(Inst)) {
Value *XCHGPtrOperand = XCHG->getPointerOperand();
- Value *Replacement =
- getTranslatedGlobalMemoryPtrOfLDS(LoadMallocPtr, XCHGPtrOperand);
+ Value *Replacement = TranslatePtr(XCHGPtrOperand);
+ if (!Replacement)
+ continue;
AtomicCmpXchgInst *NewXCHG = IRB.CreateAtomicCmpXchg(
Replacement, XCHG->getCompareOperand(), XCHG->getNewValOperand(),
XCHG->getAlign(), XCHG->getSuccessOrdering(),
@@ -734,10 +863,16 @@ void AMDGPUSwLowerLDS::translateLDSMemoryOperationsToGlobalMemory(
AsanInfo.Instructions.insert(NewXCHG);
XCHG->replaceAllUsesWith(NewXCHG);
XCHG->eraseFromParent();
+ MaybeDeadPtrs.push_back(XCHGPtrOperand);
} else if (AnyMemIntrinsic *MI = dyn_cast<AnyMemIntrinsic>(Inst)) {
- Value *NewDest = MI->getRawDest();
- if (MI->getDestAddressSpace() == AMDGPUAS::LOCAL_ADDRESS)
- NewDest = getTranslatedGlobalMemoryPtrOfLDS(LoadMallocPtr, NewDest);
+ Value *OldRawDest = MI->getRawDest();
+ Value *NewDest = OldRawDest;
+ if (MI->getDestAddressSpace() == AMDGPUAS::LOCAL_ADDRESS) {
+ NewDest = TranslatePtr(NewDest);
+ if (!NewDest)
+ continue;
+ MaybeDeadPtrs.push_back(OldRawDest);
+ }
CallInst *NewMI = nullptr;
if (AnyMemSetInst *MSI = dyn_cast<AnyMemSetInst>(MI)) {
if (MI->isAtomic()) {
@@ -750,9 +885,14 @@ void AMDGPUSwLowerLDS::translateLDSMemoryOperationsToGlobalMemory(
cast<MemSetInst>(MI)->isVolatile());
}
} else if (AnyMemTransferInst *MTI = dyn_cast<AnyMemTransferInst>(MI)) {
- Value *NewSrc = MTI->getRawSource();
- if (MTI->getSourceAddressSpace() == AMDGPUAS::LOCAL_ADDRESS)
- NewSrc = getTranslatedGlobalMemoryPtrOfLDS(LoadMallocPtr, NewSrc);
+ Value *OldRawSrc = MTI->getRawSource();
+ Value *NewSrc = OldRawSrc;
+ if (MTI->getSourceAddressSpace() == AMDGPUAS::LOCAL_ADDRESS) {
+ NewSrc = TranslatePtr(NewSrc);
+ if (!NewSrc)
+ continue;
+ MaybeDeadPtrs.push_back(OldRawSrc);
+ }
if (MI->isAtomic()) {
if (MI->getIntrinsicID() ==
Intrinsic::memmove_element_unordered_atomic) {
@@ -779,91 +919,78 @@ void AMDGPUSwLowerLDS::translateLDSMemoryOperationsToGlobalMemory(
MI->eraseFromParent();
} else if (AddrSpaceCastInst *ASC = dyn_cast<AddrSpaceCastInst>(Inst)) {
Value *AIOperand = ASC->getPointerOperand();
- Value *Replacement =
- getTranslatedGlobalMemoryPtrOfLDS(LoadMallocPtr, AIOperand);
+ Value *Replacement = TranslatePtr(AIOperand);
+ if (!Replacement)
+ continue;
Value *NewAI = IRB.CreateAddrSpaceCast(Replacement, ASC->getType());
// Note: No need to add the instruction to AsanInfo instructions to be
// instrumented list. FLAT_ADDRESS ptr would have been already
// instrumented by asan pass prior to this pass.
ASC->replaceAllUsesWith(NewAI);
ASC->eraseFromParent();
+ MaybeDeadPtrs.push_back(AIOperand);
} else
report_fatal_error("Unimplemented LDS lowering instruction");
}
-}
-
-void AMDGPUSwLowerLDS::lowerFlatToLocalRoundTrips(Function *Func) {
- // After SW LDS lowering, storage that used to live in LDS now lives in global
- // memory, but it may still be handed to a callee as a flat (generic) pointer.
- // Such a callee can re-derive an LDS pointer with a round-trip pattern:
- // %p = getelementptr ..., ptr %flat_arg ; flat, backed by global mem
- // %c = addrspacecast ptr %p to ptr addrspace(3)
- // load/store ... ptr addrspace(3) %c
- // The addrspace(3) access would then hit dead LDS. Redo the access through
- // the flat source pointer, which still points at the global backing buffer.
- SmallVector<Instruction *> ToErase;
- SmallPtrSet<AddrSpaceCastInst *, 8> Casts;
- for (BasicBlock &BB : *Func) {
- for (Instruction &Inst : BB) {
- Value *Ptr = nullptr;
- if (auto *LI = dyn_cast<LoadInst>(&Inst))
- Ptr = LI->getPointerOperand();
- else if (auto *SI = dyn_cast<StoreInst>(&Inst))
- Ptr = SI->getPointerOperand();
- else if (auto *RMW = dyn_cast<AtomicRMWInst>(&Inst))
- Ptr = RMW->getPointerOperand();
- else if (auto *XCHG = dyn_cast<AtomicCmpXchgInst>(&Inst))
- Ptr = XCHG->getPointerOperand();
- else
- continue;
- auto *ASC = dyn_cast<AddrSpaceCastInst>(Ptr);
- if (!ASC || ASC->getSrcAddressSpace() != AMDGPUAS::FLAT_ADDRESS ||
- ASC->getDestAddressSpace() != AMDGPUAS::LOCAL_ADDRESS)
- continue;
+ // Drop the now-dead flat->local casts and local GEPs.
+ RecursivelyDeleteTriviallyDeadInstructionsPermissive(MaybeDeadPtrs);
+}
- // Only collapse pointers that originate from a flat function argument,
- // i.e. storage passed in from a caller. Genuine in-function LDS pointers
- // are lowered through the base/offset table machinery instead.
- Value *FlatPtr = ASC->getPointerOperand();
- if (!isa<Argument>(getUnderlyingObject(FlatPtr)))
- continue;
+Value *AMDGPUSwLowerLDS::getFlatPtrForRoundTripLDSAccess(Function *Func,
+ Value *LDSPtr) {
+ // If LDSPtr is lowered storage reached via a flat argument, i.e.
+ // addrspacecast(flat->local) optionally followed by local GEPs, return the
+ // equivalent flat pointer into the global backing buffer. Otherwise nullptr,
+ // so the caller falls back to the base/offset-table translation.
+ auto It = FuncLDSAccessInfo.NonKernelsWithLDSFlatArg.find(Func);
+ if (It == FuncLDSAccessInfo.NonKernelsWithLDSFlatArg.end())
+ return nullptr;
+ const SmallPtrSet<Argument *, 4> &LDSArgs = It->second;
+
+ // Peel off local GEPs sitting between the cast and the access.
+ SmallVector<GEPOperator *, 4> GEPs;
+ Value *Cur = LDSPtr;
+ while (auto *GEP = dyn_cast<GEPOperator>(Cur)) {
+ if (GEP->getPointerAddressSpace() != AMDGPUAS::LOCAL_ADDRESS)
+ return nullptr;
+ GEPs.push_back(GEP);
+ Cur = GEP->getPointerOperand();
+ }
- IRB.SetInsertPoint(&Inst);
- Instruction *NewInst = nullptr;
- if (auto *LI = dyn_cast<LoadInst>(&Inst)) {
- NewInst = IRB.CreateLoad(LI->getType(), FlatPtr, LI->getProperties());
- } else if (auto *SI = dyn_cast<StoreInst>(&Inst)) {
- NewInst = IRB.CreateStore(SI->getValueOperand(), FlatPtr,
- SI->getProperties());
- } else if (auto *RMW = dyn_cast<AtomicRMWInst>(&Inst)) {
- auto *NewRMW = IRB.CreateAtomicRMW(
- RMW->getOperation(), FlatPtr, RMW->getValOperand(), RMW->getAlign(),
- RMW->getOrdering(), RMW->getSyncScopeID());
- NewRMW->setVolatile(RMW->isVolatile());
- NewInst = NewRMW;
- } else {
- auto *XCHG = cast<AtomicCmpXchgInst>(&Inst);
- auto *NewXCHG = IRB.CreateAtomicCmpXchg(
- FlatPtr, XCHG->getCompareOperand(), XCHG->getNewValOperand(),
- XCHG->getAlign(), XCHG->getSuccessOrdering(),
- XCHG->getFailureOrdering(), XCHG->getSyncScopeID());
- NewXCHG->setVolatile(XCHG->isVolatile());
- NewInst = NewXCHG;
- }
- // Flat access into the global backing buffer must be sanitized.
- AsanInfo.Instructions.insert(NewInst);
- Inst.replaceAllUsesWith(NewInst);
- ToErase.push_back(&Inst);
- Casts.insert(ASC);
- }
+ auto *ASC = dyn_cast<Operator>(Cur);
+ if (!ASC || ASC->getOpcode() != Instruction::AddrSpaceCast)
+ return nullptr;
+ Value *FlatSrc = ASC->getOperand(0);
+ if (FlatSrc->getType()->getPointerAddressSpace() != AMDGPUAS::FLAT_ADDRESS)
+ return nullptr;
+
+ // The flat source must trace back to an LDS-carrying arg.
+ SmallPtrSet<Value *, 8> Visited;
+ if (!flatPtrDerivesFromLDS(FlatSrc, LDSArgs, Visited,
+ /*AcceptLocalCast=*/false))
+ return nullptr;
+
+ // Re-apply the peeled GEPs in the flat address space.
+ Value *Flat = FlatSrc;
+ for (GEPOperator *GEP : reverse(GEPs)) {
+ SmallVector<Value *, 4> Indices(GEP->idx_begin(), GEP->idx_end());
+ Flat = IRB.CreateGEP(GEP->getSourceElementType(), Flat, Indices, "",
+ GEP->getNoWrapFlags());
}
- for (Instruction *I : ToErase)
- I->eraseFromParent();
- // Drop the now-dead round-trip casts.
- for (AddrSpaceCastInst *ASC : Casts)
- if (ASC->use_empty())
- ASC->eraseFromParent();
+ return Flat;
+}
+
+void AMDGPUSwLowerLDS::lowerNonKernelLDSFlatArgAccesses(Function *Func) {
+ // Lower a non-kernel that only reaches lowered LDS through a flat argument.
+ // No base/offset table is needed: the flat pointer already points at the
+ // global backing buffer, so the shared translation reuses the flat source.
+ SetVector<Instruction *> LDSInstructions;
+ getLDSMemoryInstructions(Func, LDSInstructions);
+ if (LDSInstructions.empty())
+ return;
+ translateLDSMemoryOperationsToGlobalMemory(Func, /*LoadMallocPtr=*/nullptr,
+ LDSInstructions);
}
void AMDGPUSwLowerLDS::poisonRedzones(Function *Func, Value *MallocPtr) {
@@ -1334,6 +1461,11 @@ bool AMDGPUSwLowerLDS::run() {
// Get address sanitizer scale.
initAsanInfo();
+ // Discover non-kernels that receive lowered LDS via a flat pointer argument.
+ // This must happen before lowering rewrites the call-site address-space
+ // casts.
+ getNonKernelsWithLDSFlatArguments();
+
for (auto &K : FuncLDSAccessInfo.KernelToLDSParametersMap) {
Function *Func = K.first;
auto &LDSParams = FuncLDSAccessInfo.KernelToLDSParametersMap[Func];
@@ -1396,15 +1528,20 @@ bool AMDGPUSwLowerLDS::run() {
Changed = true;
}
+ // Lower non-kernels that reach lowered LDS only through a flat argument.
+ // Functions also handled above already had these accesses translated there.
+ for (auto &K : FuncLDSAccessInfo.NonKernelsWithLDSFlatArg) {
+ Function *Func = K.first;
+ if (FuncLDSAccessInfo.NonKernelToLDSAccessMap.contains(Func) ||
+ FuncLDSAccessInfo.NonKernelsWithLDSArgument.contains(Func))
+ continue;
+ lowerNonKernelLDSFlatArgAccesses(Func);
+ Changed = true;
+ }
+
if (!Changed)
return Changed;
- // Fix up callees that still access lowered storage through the
- // flat-arg -> addrspace(3) round-trip pattern.
- for (Function &F : M)
- if (!F.isDeclaration())
- lowerFlatToLocalRoundTrips(&F);
-
for (auto &GV : make_early_inc_range(M.globals())) {
if (AMDGPU::isLDSVariableToLower(GV)) {
// probably want to remove from used lists
diff --git a/llvm/test/CodeGen/AMDGPU/amdgpu-sw-lower-lds-flat-ptr-arg-asan.ll b/llvm/test/CodeGen/AMDGPU/amdgpu-sw-lower-lds-flat-ptr-arg-asan.ll
index 646371a5bde24..f69ebe54cb694 100644
--- a/llvm/test/CodeGen/AMDGPU/amdgpu-sw-lower-lds-flat-ptr-arg-asan.ll
+++ b/llvm/test/CodeGen/AMDGPU/amdgpu-sw-lower-lds-flat-ptr-arg-asan.ll
@@ -1,21 +1,18 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
; RUN: opt < %s -passes=amdgpu-sw-lower-lds -amdgpu-asan-instrument-lds=false -S -mtriple=amdgcn-amd-amdhsa | FileCheck %s
-; A non-kernel receives lowered LDS storage as a flat (generic) pointer and
-; re-derives an LDS pointer via a flat -> addrspace(3) round-trip. After SW LDS
-; lowering the storage lives in global memory, so the addrspace(3) access would
-; hit dead LDS. The round-trip must be collapsed into a flat access on the
-; incoming argument, which still points at the global backing buffer. Check the
-; store, load, atomicrmw and cmpxchg cases.
-
@lds_var = internal addrspace(3) global [64 x i32] poison, align 4
+declare void @llvm.memset.p3.i64(ptr addrspace(3), i8, i64, i1)
+declare void @llvm.memcpy.p3.p1.i64(ptr addrspace(3), ptr addrspace(1), i64, i1)
+declare void @llvm.memcpy.p1.p3.i64(ptr addrspace(1), ptr addrspace(3), i64, i1)
+
define void @store_case(ptr %storage) sanitize_address {
; CHECK-LABEL: define void @store_case(
-; CHECK-SAME: ptr [[STORAGE:%.*]]){{.*}} {
-; CHECK-NOT: addrspace(3)
-; CHECK: store i32 42, ptr [[STORAGE]], align 4
-; CHECK-NOT: addrspace(3)
-; CHECK: ret void
+; CHECK-SAME: ptr [[STORAGE:%.*]]) #[[ATTR2:[0-9]+]] {
+; CHECK-NEXT: store i32 42, ptr [[STORAGE]], align 4
+; CHECK-NEXT: ret void
+;
%cast = addrspacecast ptr %storage to ptr addrspace(3)
store i32 42, ptr addrspace(3) %cast, align 4
ret void
@@ -23,11 +20,10 @@ define void @store_case(ptr %storage) sanitize_address {
define i32 @load_case(ptr %storage) sanitize_address {
; CHECK-LABEL: define i32 @load_case(
-; CHECK-SAME: ptr [[STORAGE:%.*]]){{.*}} {
-; CHECK-NOT: addrspace(3)
-; CHECK: [[VAL:%.*]] = load i32, ptr [[STORAGE]], align 4
-; CHECK-NOT: addrspace(3)
-; CHECK: ret i32 [[VAL]]
+; CHECK-SAME: ptr [[STORAGE:%.*]]) #[[ATTR2]] {
+; CHECK-NEXT: [[TMP1:%.*]] = load i32, ptr [[STORAGE]], align 4
+; CHECK-NEXT: ret i32 [[TMP1]]
+;
%cast = addrspacecast ptr %storage to ptr addrspace(3)
%val = load i32, ptr addrspace(3) %cast, align 4
ret i32 %val
@@ -35,11 +31,10 @@ define i32 @load_case(ptr %storage) sanitize_address {
define i32 @atomicrmw_case(ptr %storage) sanitize_address {
; CHECK-LABEL: define i32 @atomicrmw_case(
-; CHECK-SAME: ptr [[STORAGE:%.*]]){{.*}} {
-; CHECK-NOT: addrspace(3)
-; CHECK: [[OLD:%.*]] = atomicrmw add ptr [[STORAGE]], i32 1 seq_cst, align 4
-; CHECK-NOT: addrspace(3)
-; CHECK: ret i32 [[OLD]]
+; CHECK-SAME: ptr [[STORAGE:%.*]]) #[[ATTR2]] {
+; CHECK-NEXT: [[TMP1:%.*]] = atomicrmw add ptr [[STORAGE]], i32 1 seq_cst, align 4
+; CHECK-NEXT: ret i32 [[TMP1]]
+;
%cast = addrspacecast ptr %storage to ptr addrspace(3)
%old = atomicrmw add ptr addrspace(3) %cast, i32 1 seq_cst, align 4
ret i32 %old
@@ -47,25 +42,213 @@ define i32 @atomicrmw_case(ptr %storage) sanitize_address {
define i32 @cmpxchg_case(ptr %storage) sanitize_address {
; CHECK-LABEL: define i32 @cmpxchg_case(
-; CHECK-SAME: ptr [[STORAGE:%.*]]){{.*}} {
-; CHECK-NOT: addrspace(3)
-; CHECK: [[RES:%.*]] = cmpxchg ptr [[STORAGE]], i32 0, i32 1 seq_cst seq_cst, align 4
-; CHECK-NOT: addrspace(3)
-; CHECK: [[VAL:%.*]] = extractvalue { i32, i1 } [[RES]], 0
-; CHECK: ret i32 [[VAL]]
+; CHECK-SAME: ptr [[STORAGE:%.*]]) #[[ATTR2]] {
+; CHECK-NEXT: [[TMP1:%.*]] = cmpxchg ptr [[STORAGE]], i32 0, i32 1 seq_cst seq_cst, align 4
+; CHECK-NEXT: [[VAL:%.*]] = extractvalue { i32, i1 } [[TMP1]], 0
+; CHECK-NEXT: ret i32 [[VAL]]
+;
%cast = addrspacecast ptr %storage to ptr addrspace(3)
%res = cmpxchg ptr addrspace(3) %cast, i32 0, i32 1 seq_cst seq_cst, align 4
%val = extractvalue { i32, i1 } %res, 0
ret i32 %val
}
-define amdgpu_kernel void @kernel() sanitize_address {
+; GEP applied in the local address space after the cast must be rebuilt in flat.
+define void @gep_chain_case(ptr %storage) sanitize_address {
+; CHECK-LABEL: define void @gep_chain_case(
+; CHECK-SAME: ptr [[STORAGE:%.*]]) #[[ATTR2]] {
+; CHECK-NEXT: [[TMP1:%.*]] = getelementptr i32, ptr [[STORAGE]], i64 5
+; CHECK-NEXT: store i32 7, ptr [[TMP1]], align 4
+; CHECK-NEXT: ret void
+;
+ %cast = addrspacecast ptr %storage to ptr addrspace(3)
+ %elem = getelementptr i32, ptr addrspace(3) %cast, i64 5
+ store i32 7, ptr addrspace(3) %elem, align 4
+ ret void
+}
+
+; A flat LDS argument forwarded through a non-kernel must be handled in the
+; callee too.
+define void @forwarded_callee(ptr %p) sanitize_address {
+; CHECK-LABEL: define void @forwarded_callee(
+; CHECK-SAME: ptr [[P:%.*]]) #[[ATTR2]] {
+; CHECK-NEXT: store i32 9, ptr [[P]], align 4
+; CHECK-NEXT: ret void
+;
+ %cast = addrspacecast ptr %p to ptr addrspace(3)
+ store i32 9, ptr addrspace(3) %cast, align 4
+ ret void
+}
+
+define void @forwarder(ptr %q) sanitize_address {
+; CHECK-LABEL: define void @forwarder(
+; CHECK-SAME: ptr [[Q:%.*]]) #[[ATTR2]] {
+; CHECK-NEXT: call void @forwarded_callee(ptr [[Q]])
+; CHECK-NEXT: ret void
+;
+ call void @forwarded_callee(ptr %q)
+ ret void
+}
+
+; Memory intrinsics with the lowered LDS as the destination must be redirected
+; through the flat argument.
+define void @memset_case(ptr %storage) sanitize_address {
+; CHECK-LABEL: define void @memset_case(
+; CHECK-SAME: ptr [[STORAGE:%.*]]) #[[ATTR2]] {
+; CHECK-NEXT: call void @llvm.memset.p0.i64(ptr [[STORAGE]], i8 0, i64 32, i1 false)
+; CHECK-NEXT: ret void
+;
+ %cast = addrspacecast ptr %storage to ptr addrspace(3)
+ call void @llvm.memset.p3.i64(ptr addrspace(3) %cast, i8 0, i64 32, i1 false)
+ ret void
+}
+
+define void @memcpy_dst_case(ptr %storage, ptr addrspace(1) %src) sanitize_address {
+; CHECK-LABEL: define void @memcpy_dst_case(
+; CHECK-SAME: ptr [[STORAGE:%.*]], ptr addrspace(1) [[SRC:%.*]]) #[[ATTR2]] {
+; CHECK-NEXT: call void @llvm.memcpy.p0.p1.i64(ptr [[STORAGE]], ptr addrspace(1) [[SRC]], i64 32, i1 false)
+; CHECK-NEXT: ret void
+;
+ %cast = addrspacecast ptr %storage to ptr addrspace(3)
+ call void @llvm.memcpy.p3.p1.i64(ptr addrspace(3) %cast, ptr addrspace(1) %src, i64 32, i1 false)
+ ret void
+}
+
+; ... and with the lowered LDS as the source.
+define void @memcpy_src_case(ptr %storage, ptr addrspace(1) %dst) sanitize_address {
+; CHECK-LABEL: define void @memcpy_src_case(
+; CHECK-SAME: ptr [[STORAGE:%.*]], ptr addrspace(1) [[DST:%.*]]) #[[ATTR2]] {
+; CHECK-NEXT: call void @llvm.memcpy.p1.p0.i64(ptr addrspace(1) [[DST]], ptr [[STORAGE]], i64 32, i1 false)
+; CHECK-NEXT: ret void
+;
+ %cast = addrspacecast ptr %storage to ptr addrspace(3)
+ call void @llvm.memcpy.p1.p3.i64(ptr addrspace(1) %dst, ptr addrspace(3) %cast, i64 32, i1 false)
+ ret void
+}
+
+; A flat pointer merged with a select before the cast must still be recognized
+; as an LDS-carrying flat argument (both operands trace to %storage).
+define void @select_case(ptr %storage, i1 %c) sanitize_address {
+; CHECK-LABEL: define void @select_case(
+; CHECK-SAME: ptr [[STORAGE:%.*]], i1 [[C:%.*]]) #[[ATTR2]] {
+; CHECK-NEXT: [[G:%.*]] = getelementptr i32, ptr [[STORAGE]], i64 1
+; CHECK-NEXT: [[SEL:%.*]] = select i1 [[C]], ptr [[STORAGE]], ptr [[G]]
+; CHECK-NEXT: store i32 5, ptr [[SEL]], align 4
+; CHECK-NEXT: ret void
+;
+ %g = getelementptr i32, ptr %storage, i64 1
+ %sel = select i1 %c, ptr %storage, ptr %g
+ %cast = addrspacecast ptr %sel to ptr addrspace(3)
+ store i32 5, ptr addrspace(3) %cast, align 4
+ ret void
+}
+
+; Loop-carried flat pointer: the phi lives in the flat address space and the
+; addrspace(3) pointer is re-derived each iteration.
+define void @loop_case(ptr %storage, i32 %n) sanitize_address {
+; CHECK-LABEL: define void @loop_case(
+; CHECK-SAME: ptr [[STORAGE:%.*]], i32 [[N:%.*]]) #[[ATTR2]] {
+; CHECK-NEXT: [[ENTRY:.*]]:
+; CHECK-NEXT: br label %[[LOOP:.*]]
+; CHECK: [[LOOP]]:
+; CHECK-NEXT: [[I:%.*]] = phi i32 [ 0, %[[ENTRY]] ], [ [[I_NEXT:%.*]], %[[LOOP]] ]
+; CHECK-NEXT: [[P:%.*]] = phi ptr [ [[STORAGE]], %[[ENTRY]] ], [ [[P_NEXT:%.*]], %[[LOOP]] ]
+; CHECK-NEXT: store i32 [[I]], ptr [[P]], align 4
+; CHECK-NEXT: [[P_NEXT]] = getelementptr i32, ptr [[P]], i64 1
+; CHECK-NEXT: [[I_NEXT]] = add i32 [[I]], 1
+; CHECK-NEXT: [[COND:%.*]] = icmp slt i32 [[I_NEXT]], [[N]]
+; CHECK-NEXT: br i1 [[COND]], label %[[LOOP]], label %[[EXIT:.*]]
+; CHECK: [[EXIT]]:
+; CHECK-NEXT: ret void
+;
+entry:
+ br label %loop
+loop:
+ %i = phi i32 [ 0, %entry ], [ %i.next, %loop ]
+ %p = phi ptr [ %storage, %entry ], [ %p.next, %loop ]
+ %cast = addrspacecast ptr %p to ptr addrspace(3)
+ store i32 %i, ptr addrspace(3) %cast, align 4
+ %p.next = getelementptr i32, ptr %p, i64 1
+ %i.next = add i32 %i, 1
+ %cond = icmp slt i32 %i.next, %n
+ br i1 %cond, label %loop, label %exit
+exit:
+ ret void
+}
+
+define amdgpu_kernel void @kernel(i1 %c, i32 %n, ptr addrspace(1) %buf) sanitize_address {
; CHECK-LABEL: define amdgpu_kernel void @kernel(
+; CHECK-SAME: i1 [[C:%.*]], i32 [[N:%.*]], ptr addrspace(1) [[BUF:%.*]]) #[[ATTR3:[0-9]+]] {
+; CHECK-NEXT: [[WID:.*]]:
+; CHECK-NEXT: [[TMP0:%.*]] = call i32 @llvm.amdgcn.workitem.id.x()
+; CHECK-NEXT: [[TMP1:%.*]] = call i32 @llvm.amdgcn.workitem.id.y()
+; CHECK-NEXT: [[TMP2:%.*]] = call i32 @llvm.amdgcn.workitem.id.z()
+; CHECK-NEXT: [[TMP3:%.*]] = or i32 [[TMP0]], [[TMP1]]
+; CHECK-NEXT: [[TMP4:%.*]] = or i32 [[TMP3]], [[TMP2]]
+; CHECK-NEXT: [[TMP5:%.*]] = icmp eq i32 [[TMP4]], 0
+; CHECK-NEXT: br i1 [[TMP5]], label %[[MALLOC:.*]], label %[[BB18:.*]]
+; CHECK: [[MALLOC]]:
+; CHECK-NEXT: [[TMP6:%.*]] = load i32, ptr addrspace(1) getelementptr inbounds ([[LLVM_AMDGCN_SW_LDS_KERNEL_MD_TYPE:%.*]], ptr addrspace(1) @llvm.amdgcn.sw.lds.kernel.md, i32 0, i32 1, i32 0), align 4
+; CHECK-NEXT: [[TMP7:%.*]] = load i32, ptr addrspace(1) getelementptr inbounds ([[LLVM_AMDGCN_SW_LDS_KERNEL_MD_TYPE]], ptr addrspace(1) @llvm.amdgcn.sw.lds.kernel.md, i32 0, i32 1, i32 2), align 4
+; CHECK-NEXT: [[TMP8:%.*]] = add i32 [[TMP6]], [[TMP7]]
+; CHECK-NEXT: [[TMP9:%.*]] = zext i32 [[TMP8]] to i64
+; CHECK-NEXT: [[TMP10:%.*]] = call ptr @llvm.returnaddress.p0(i32 0)
+; CHECK-NEXT: [[TMP11:%.*]] = ptrtoint ptr [[TMP10]] to i64
+; CHECK-NEXT: [[TMP12:%.*]] = call i64 @__asan_malloc_impl(i64 [[TMP9]], i64 [[TMP11]])
+; CHECK-NEXT: [[TMP13:%.*]] = inttoptr i64 [[TMP12]] to ptr addrspace(1)
+; CHECK-NEXT: store ptr addrspace(1) [[TMP13]], ptr addrspace(3) @llvm.amdgcn.sw.lds.kernel, align 8
+; CHECK-NEXT: [[TMP14:%.*]] = getelementptr inbounds i8, ptr addrspace(1) [[TMP13]], i64 8
+; CHECK-NEXT: [[TMP15:%.*]] = ptrtoint ptr addrspace(1) [[TMP14]] to i64
+; CHECK-NEXT: call void @__asan_poison_region(i64 [[TMP15]], i64 24)
+; CHECK-NEXT: [[TMP16:%.*]] = getelementptr inbounds i8, ptr addrspace(1) [[TMP13]], i64 288
+; CHECK-NEXT: [[TMP17:%.*]] = ptrtoint ptr addrspace(1) [[TMP16]] to i64
+; CHECK-NEXT: call void @__asan_poison_region(i64 [[TMP17]], i64 64)
+; CHECK-NEXT: br label %[[BB18]]
+; CHECK: [[BB18]]:
+; CHECK-NEXT: [[XYZCOND:%.*]] = phi i1 [ false, %[[WID]] ], [ true, %[[MALLOC]] ]
+; CHECK-NEXT: call void @llvm.amdgcn.s.barrier()
+; CHECK-NEXT: [[TMP19:%.*]] = load ptr addrspace(1), ptr addrspace(3) @llvm.amdgcn.sw.lds.kernel, align 8
+; CHECK-NEXT: [[TMP20:%.*]] = load i32, ptr addrspace(1) getelementptr inbounds ([[LLVM_AMDGCN_SW_LDS_KERNEL_MD_TYPE]], ptr addrspace(1) @llvm.amdgcn.sw.lds.kernel.md, i32 0, i32 1, i32 0), align 4
+; CHECK-NEXT: [[TMP21:%.*]] = getelementptr inbounds i8, ptr addrspace(3) @llvm.amdgcn.sw.lds.kernel, i32 [[TMP20]]
+; CHECK-NEXT: [[TMP22:%.*]] = ptrtoint ptr addrspace(3) [[TMP21]] to i32
+; CHECK-NEXT: [[TMP23:%.*]] = getelementptr inbounds i8, ptr addrspace(1) [[TMP19]], i32 [[TMP22]]
+; CHECK-NEXT: [[TMP24:%.*]] = addrspacecast ptr addrspace(1) [[TMP23]] to ptr
+; CHECK-NEXT: call void @store_case(ptr [[TMP24]])
+; CHECK-NEXT: [[L:%.*]] = call i32 @load_case(ptr [[TMP24]])
+; CHECK-NEXT: [[A:%.*]] = call i32 @atomicrmw_case(ptr [[TMP24]])
+; CHECK-NEXT: [[X:%.*]] = call i32 @cmpxchg_case(ptr [[TMP24]])
+; CHECK-NEXT: call void @gep_chain_case(ptr [[TMP24]])
+; CHECK-NEXT: call void @forwarder(ptr [[TMP24]])
+; CHECK-NEXT: call void @memset_case(ptr [[TMP24]])
+; CHECK-NEXT: call void @memcpy_dst_case(ptr [[TMP24]], ptr addrspace(1) [[BUF]])
+; CHECK-NEXT: call void @memcpy_src_case(ptr [[TMP24]], ptr addrspace(1) [[BUF]])
+; CHECK-NEXT: call void @select_case(ptr [[TMP24]], i1 [[C]])
+; CHECK-NEXT: call void @loop_case(ptr [[TMP24]], i32 [[N]])
+; CHECK-NEXT: br label %[[CONDFREE:.*]]
+; CHECK: [[CONDFREE]]:
+; CHECK-NEXT: call void @llvm.amdgcn.s.barrier()
+; CHECK-NEXT: br i1 [[XYZCOND]], label %[[FREE:.*]], label %[[END:.*]]
+; CHECK: [[FREE]]:
+; CHECK-NEXT: [[TMP25:%.*]] = call ptr @llvm.returnaddress.p0(i32 0)
+; CHECK-NEXT: [[TMP26:%.*]] = ptrtoint ptr [[TMP25]] to i64
+; CHECK-NEXT: [[TMP27:%.*]] = ptrtoint ptr addrspace(1) [[TMP19]] to i64
+; CHECK-NEXT: call void @__asan_free_impl(i64 [[TMP27]], i64 [[TMP26]])
+; CHECK-NEXT: br label %[[END]]
+; CHECK: [[END]]:
+; CHECK-NEXT: ret void
+;
%flat = addrspacecast ptr addrspace(3) @lds_var to ptr
call void @store_case(ptr %flat)
%l = call i32 @load_case(ptr %flat)
%a = call i32 @atomicrmw_case(ptr %flat)
- %c = call i32 @cmpxchg_case(ptr %flat)
+ %x = call i32 @cmpxchg_case(ptr %flat)
+ call void @gep_chain_case(ptr %flat)
+ call void @forwarder(ptr %flat)
+ call void @memset_case(ptr %flat)
+ call void @memcpy_dst_case(ptr %flat, ptr addrspace(1) %buf)
+ call void @memcpy_src_case(ptr %flat, ptr addrspace(1) %buf)
+ call void @select_case(ptr %flat, i1 %c)
+ call void @loop_case(ptr %flat, i32 %n)
ret void
}
>From 6e56fa3a874da74147771e97c092a1793627fc68 Mon Sep 17 00:00:00 2001
From: skc7 <Krishna.Sankisa at amd.com>
Date: Thu, 30 Jul 2026 20:41:22 +0530
Subject: [PATCH 3/4] changes as per review feedback
---
llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp | 63 +++++++++++----------
1 file changed, 33 insertions(+), 30 deletions(-)
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp b/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp
index 25c82883bd419..69d89b8862230 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp
@@ -303,17 +303,21 @@ void AMDGPUSwLowerLDS::getNonKernelsWithLDSArguments(const CallGraph &CG) {
}
}
-// True if flat pointer V is derived from an LDS-carrying argument in LDSArgs,
-// looking through GEP/bitcast/phi/select; Visited breaks phi cycles. When
-// AcceptLocalCast is set, a flat<-local addrspacecast also counts as an origin.
-// Detection (getNonKernelsWithLDSFlatArguments) sets it, since that cast is how
-// LDS first enters a flat pointer; the rewrite
-// (getFlatPtrForRoundTripLDSAccess) leaves it clear so only argument origins
-// qualify.
-static bool flatPtrDerivesFromLDS(Value *V,
- const SmallPtrSetImpl<Argument *> &LDSArgs,
- SmallPtrSetImpl<Value *> &Visited,
- bool AcceptLocalCast) {
+/// Returns true if the flat pointer \p V is derived from an LDS-carrying
+/// argument in \p LDSArgs, looking through GEP/bitcast/phi/select.
+///
+/// \param V The flat pointer value to test.
+/// \param LDSArgs Arguments known to carry LDS-backed storage.
+/// \param Visited Set of already-visited values, used to break phi cycles.
+/// \param AcceptLocalCast If set, a flat<-local addrspacecast also counts as
+/// an origin. Detection (getNonKernelsWithLDSFlatArguments) sets this,
+/// since that cast is how LDS first enters a flat pointer; the
+/// rewrite (getFlatPtrForRoundTripLDSAccess) leaves it clear so only
+/// argument origins qualify.
+static bool isFlatPtrDerivedFromLDS(Value *V,
+ const SmallPtrSetImpl<Argument *> &LDSArgs,
+ SmallPtrSetImpl<Value *> &Visited,
+ bool AcceptLocalCast) {
if (!V->getType()->isPointerTy() || !Visited.insert(V).second)
return false;
if (auto *A = dyn_cast<Argument>(V))
@@ -328,23 +332,23 @@ static bool flatPtrDerivesFromLDS(Value *V,
Value *Src = Op->getOperand(0);
if (Src->getType()->getPointerAddressSpace() == AMDGPUAS::LOCAL_ADDRESS)
return true;
- return flatPtrDerivesFromLDS(Src, LDSArgs, Visited, AcceptLocalCast);
+ return isFlatPtrDerivedFromLDS(Src, LDSArgs, Visited, AcceptLocalCast);
}
case Instruction::GetElementPtr:
case Instruction::BitCast:
- return flatPtrDerivesFromLDS(Op->getOperand(0), LDSArgs, Visited,
- AcceptLocalCast);
+ return isFlatPtrDerivedFromLDS(Op->getOperand(0), LDSArgs, Visited,
+ AcceptLocalCast);
case Instruction::PHI:
for (Value *In : cast<PHINode>(Op)->incoming_values())
- if (flatPtrDerivesFromLDS(In, LDSArgs, Visited, AcceptLocalCast))
+ if (isFlatPtrDerivedFromLDS(In, LDSArgs, Visited, AcceptLocalCast))
return true;
return false;
case Instruction::Select: {
auto *SI = cast<SelectInst>(Op);
- return flatPtrDerivesFromLDS(SI->getTrueValue(), LDSArgs, Visited,
- AcceptLocalCast) ||
- flatPtrDerivesFromLDS(SI->getFalseValue(), LDSArgs, Visited,
- AcceptLocalCast);
+ return isFlatPtrDerivedFromLDS(SI->getTrueValue(), LDSArgs, Visited,
+ AcceptLocalCast) ||
+ isFlatPtrDerivedFromLDS(SI->getFalseValue(), LDSArgs, Visited,
+ AcceptLocalCast);
}
default:
return false;
@@ -374,8 +378,8 @@ void AMDGPUSwLowerLDS::getNonKernelsWithLDSFlatArguments() {
Function *Callee = CB->getCalledFunction();
if (!Callee || Callee->isDeclaration() || AMDGPU::isKernel(*Callee))
continue;
- unsigned NumArgs =
- std::min(CB->arg_size(), (unsigned)Callee->arg_size());
+ unsigned NumArgs = std::min(
+ CB->arg_size(), static_cast<unsigned>(Callee->arg_size()));
for (unsigned ArgNo = 0; ArgNo < NumArgs; ++ArgNo) {
Argument *CalleeArg = Callee->getArg(ArgNo);
Type *ArgTy = CalleeArg->getType();
@@ -384,8 +388,8 @@ void AMDGPUSwLowerLDS::getNonKernelsWithLDSFlatArguments() {
continue;
Value *Actual = CB->getArgOperand(ArgNo);
SmallPtrSet<Value *, 8> Visited;
- if (!flatPtrDerivesFromLDS(Actual, CallerLDSArgs, Visited,
- /*AcceptLocalCast=*/true))
+ if (!isFlatPtrDerivedFromLDS(Actual, CallerLDSArgs, Visited,
+ /*AcceptLocalCast=*/true))
continue;
if (FlatArgMap[Callee].insert(CalleeArg).second)
Changed = true;
@@ -958,17 +962,16 @@ Value *AMDGPUSwLowerLDS::getFlatPtrForRoundTripLDSAccess(Function *Func,
Cur = GEP->getPointerOperand();
}
- auto *ASC = dyn_cast<Operator>(Cur);
- if (!ASC || ASC->getOpcode() != Instruction::AddrSpaceCast)
- return nullptr;
- Value *FlatSrc = ASC->getOperand(0);
- if (FlatSrc->getType()->getPointerAddressSpace() != AMDGPUAS::FLAT_ADDRESS)
+ auto *ASC = dyn_cast<AddrSpaceCastOperator>(Cur);
+ if (!ASC || ASC->getSrcAddressSpace() != AMDGPUAS::FLAT_ADDRESS ||
+ ASC->getDestAddressSpace() != AMDGPUAS::LOCAL_ADDRESS)
return nullptr;
+ Value *FlatSrc = ASC->getPointerOperand();
// The flat source must trace back to an LDS-carrying arg.
SmallPtrSet<Value *, 8> Visited;
- if (!flatPtrDerivesFromLDS(FlatSrc, LDSArgs, Visited,
- /*AcceptLocalCast=*/false))
+ if (!isFlatPtrDerivedFromLDS(FlatSrc, LDSArgs, Visited,
+ /*AcceptLocalCast=*/false))
return nullptr;
// Re-apply the peeled GEPs in the flat address space.
>From d27d924bd05ad42ddfcc83634ffee991afd0d90d Mon Sep 17 00:00:00 2001
From: skc7 <Krishna.Sankisa at amd.com>
Date: Fri, 31 Jul 2026 17:01:26 +0530
Subject: [PATCH 4/4] [AMDGPU] Lower LDS flat round trips without provenance
analysis
---
llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp | 193 +++++---------------
1 file changed, 42 insertions(+), 151 deletions(-)
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp b/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp
index 69d89b8862230..ec2d4a24586a6 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp
@@ -82,14 +82,13 @@
// corresponding element in offset table. With this information, replacement
// value is obtained.
//
-// Replacement of LDS accessed via flat pointer arguments:
-// An LDS pointer can reach a non-kernel as a flat (generic) pointer
-// argument, e.g. after an addrspacecast of an LDS global to flat. A
-// fixed-point call-graph walk identifies the functions and argument
-// indices that carry such LDS-derived flat pointers. Memory operations
-// that reach LDS through these arguments are redirected to the relocated
-// global memory by reconstructing an equivalent flat pointer, instead of
-// accessing the now-dead hardware LDS.
+// Replacement of LDS accessed via a flat round trip:
+// An LDS pointer can reach a non-kernel as a flat (generic) pointer, e.g.
+// after an addrspacecast of an LDS global to flat that the callee casts
+// back to addrspace(3). Such a flat<->local round trip is collapsed into a
+// direct flat access into the relocated global memory, instead of accessing
+// the now-dead hardware LDS. This is valid regardless of the pointer's
+// origin, so no provenance analysis is needed.
//===----------------------------------------------------------------------===//
#include "AMDGPU.h"
@@ -181,7 +180,6 @@ struct FunctionsAndLDSAccess {
SetVector<Function *> NonKernelsWithLDSArgument;
SetVector<GlobalVariable *> AllNonKernelLDSAccess;
FunctionVariableMap NonKernelToLDSAccessMap;
- DenseMap<Function *, SmallPtrSet<Argument *, 4>> NonKernelsWithLDSFlatArg;
};
class AMDGPUSwLowerLDS {
@@ -191,7 +189,6 @@ class AMDGPUSwLowerLDS {
bool run();
void getUsesOfLDSByNonKernels();
void getNonKernelsWithLDSArguments(const CallGraph &CG);
- void getNonKernelsWithLDSFlatArguments();
SetVector<Function *>
getOrderedIndirectLDSAccessingKernels(SetVector<Function *> &Kernels);
SetVector<GlobalVariable *>
@@ -218,8 +215,8 @@ class AMDGPUSwLowerLDS {
void lowerNonKernelLDSAccesses(Function *Func,
SetVector<GlobalVariable *> &LDSGlobals,
NonKernelLDSParameters &NKLDSParams);
- void lowerNonKernelLDSFlatArgAccesses(Function *Func);
- Value *getFlatPtrForRoundTripLDSAccess(Function *Func, Value *LDSPtr);
+ bool lowerNonKernelLDSFlatArgAccesses(Function *Func);
+ Value *getFlatPtrForRoundTripLDSAccess(Value *LDSPtr);
void
updateMallocSizeForDynamicLDS(Function *Func, Value **CurrMallocSize,
Value *HiddenDynLDSSize,
@@ -303,103 +300,6 @@ void AMDGPUSwLowerLDS::getNonKernelsWithLDSArguments(const CallGraph &CG) {
}
}
-/// Returns true if the flat pointer \p V is derived from an LDS-carrying
-/// argument in \p LDSArgs, looking through GEP/bitcast/phi/select.
-///
-/// \param V The flat pointer value to test.
-/// \param LDSArgs Arguments known to carry LDS-backed storage.
-/// \param Visited Set of already-visited values, used to break phi cycles.
-/// \param AcceptLocalCast If set, a flat<-local addrspacecast also counts as
-/// an origin. Detection (getNonKernelsWithLDSFlatArguments) sets this,
-/// since that cast is how LDS first enters a flat pointer; the
-/// rewrite (getFlatPtrForRoundTripLDSAccess) leaves it clear so only
-/// argument origins qualify.
-static bool isFlatPtrDerivedFromLDS(Value *V,
- const SmallPtrSetImpl<Argument *> &LDSArgs,
- SmallPtrSetImpl<Value *> &Visited,
- bool AcceptLocalCast) {
- if (!V->getType()->isPointerTy() || !Visited.insert(V).second)
- return false;
- if (auto *A = dyn_cast<Argument>(V))
- return LDSArgs.contains(A);
- auto *Op = dyn_cast<Operator>(V);
- if (!Op)
- return false;
- switch (Op->getOpcode()) {
- case Instruction::AddrSpaceCast: {
- if (!AcceptLocalCast)
- return false;
- Value *Src = Op->getOperand(0);
- if (Src->getType()->getPointerAddressSpace() == AMDGPUAS::LOCAL_ADDRESS)
- return true;
- return isFlatPtrDerivedFromLDS(Src, LDSArgs, Visited, AcceptLocalCast);
- }
- case Instruction::GetElementPtr:
- case Instruction::BitCast:
- return isFlatPtrDerivedFromLDS(Op->getOperand(0), LDSArgs, Visited,
- AcceptLocalCast);
- case Instruction::PHI:
- for (Value *In : cast<PHINode>(Op)->incoming_values())
- if (isFlatPtrDerivedFromLDS(In, LDSArgs, Visited, AcceptLocalCast))
- return true;
- return false;
- case Instruction::Select: {
- auto *SI = cast<SelectInst>(Op);
- return isFlatPtrDerivedFromLDS(SI->getTrueValue(), LDSArgs, Visited,
- AcceptLocalCast) ||
- isFlatPtrDerivedFromLDS(SI->getFalseValue(), LDSArgs, Visited,
- AcceptLocalCast);
- }
- default:
- return false;
- }
-}
-
-void AMDGPUSwLowerLDS::getNonKernelsWithLDSFlatArguments() {
- // A kernel or non-kernel may pass lowered LDS storage to a non-kernel as a
- // flat pointer instead of addrspace(3). Record, per callee, which flat
- // parameters carry LDS. Must run before lowering rewrites the call-site
- // casts.
- auto &FlatArgMap = FuncLDSAccessInfo.NonKernelsWithLDSFlatArg;
- bool Changed = true;
- while (Changed) {
- Changed = false;
- for (Function &F : M) {
- if (F.isDeclaration())
- continue;
- SmallPtrSet<Argument *, 4> CallerLDSArgs;
- if (auto It = FlatArgMap.find(&F); It != FlatArgMap.end())
- CallerLDSArgs = It->second;
- for (BasicBlock &BB : F) {
- for (Instruction &I : BB) {
- auto *CB = dyn_cast<CallBase>(&I);
- if (!CB)
- continue;
- Function *Callee = CB->getCalledFunction();
- if (!Callee || Callee->isDeclaration() || AMDGPU::isKernel(*Callee))
- continue;
- unsigned NumArgs = std::min(
- CB->arg_size(), static_cast<unsigned>(Callee->arg_size()));
- for (unsigned ArgNo = 0; ArgNo < NumArgs; ++ArgNo) {
- Argument *CalleeArg = Callee->getArg(ArgNo);
- Type *ArgTy = CalleeArg->getType();
- if (!ArgTy->isPointerTy() ||
- ArgTy->getPointerAddressSpace() != AMDGPUAS::FLAT_ADDRESS)
- continue;
- Value *Actual = CB->getArgOperand(ArgNo);
- SmallPtrSet<Value *, 8> Visited;
- if (!isFlatPtrDerivedFromLDS(Actual, CallerLDSArgs, Visited,
- /*AcceptLocalCast=*/true))
- continue;
- if (FlatArgMap[Callee].insert(CalleeArg).second)
- Changed = true;
- }
- }
- }
- }
- }
-}
-
void AMDGPUSwLowerLDS::getUsesOfLDSByNonKernels() {
for (GlobalVariable *GV : FuncLDSAccessInfo.AllNonKernelLDSAccess) {
if (!AMDGPU::isLDSVariableToLower(*GV))
@@ -801,11 +701,11 @@ void AMDGPUSwLowerLDS::translateLDSMemoryOperationsToGlobalMemory(
SetVector<Instruction *> &LDSInstructions) {
LLVM_DEBUG(dbgs() << "Translating LDS memory operations to global memory : "
<< Func->getName());
- // Map an LDS pointer to its global-memory equivalent: a flat-argument access
- // uses the flat source pointer, otherwise the base/offset table. nullptr
- // means neither applies, so the operation is left unchanged.
+ // Map an LDS pointer to its global-memory equivalent: a flat<->local round
+ // trip collapses to its flat source pointer, otherwise use the base/offset
+ // table. nullptr means neither applies, so the operation is left unchanged.
auto TranslatePtr = [&](Value *LDSPtr) -> Value * {
- if (Value *Flat = getFlatPtrForRoundTripLDSAccess(Func, LDSPtr))
+ if (Value *Flat = getFlatPtrForRoundTripLDSAccess(LDSPtr))
return Flat;
if (!LoadMallocPtr)
return nullptr;
@@ -941,18 +841,17 @@ void AMDGPUSwLowerLDS::translateLDSMemoryOperationsToGlobalMemory(
RecursivelyDeleteTriviallyDeadInstructionsPermissive(MaybeDeadPtrs);
}
-Value *AMDGPUSwLowerLDS::getFlatPtrForRoundTripLDSAccess(Function *Func,
- Value *LDSPtr) {
- // If LDSPtr is lowered storage reached via a flat argument, i.e.
- // addrspacecast(flat->local) optionally followed by local GEPs, return the
- // equivalent flat pointer into the global backing buffer. Otherwise nullptr,
- // so the caller falls back to the base/offset-table translation.
- auto It = FuncLDSAccessInfo.NonKernelsWithLDSFlatArg.find(Func);
- if (It == FuncLDSAccessInfo.NonKernelsWithLDSFlatArg.end())
- return nullptr;
- const SmallPtrSet<Argument *, 4> &LDSArgs = It->second;
-
- // Peel off local GEPs sitting between the cast and the access.
+Value *AMDGPUSwLowerLDS::getFlatPtrForRoundTripLDSAccess(Value *LDSPtr) {
+ // If LDSPtr is an addrspacecast(flat->local) optionally followed by local
+ // GEPs, return the equivalent pointer with the GEPs re-applied in the flat
+ // address space; otherwise nullptr so the caller falls back to the
+ // base/offset table.
+ //
+ // No provenance check is needed: collapsing a flat<->local round trip into a
+ // direct flat access is always valid. When the flat pointer is in the LDS
+ // aperture the flat access reaches the same memory as the local one; when it
+ // is not (e.g. it points into the relocated global buffer) the local cast was
+ // already undefined, so the flat access is strictly better.
SmallVector<GEPOperator *, 4> GEPs;
Value *Cur = LDSPtr;
while (auto *GEP = dyn_cast<GEPOperator>(Cur)) {
@@ -966,16 +865,9 @@ Value *AMDGPUSwLowerLDS::getFlatPtrForRoundTripLDSAccess(Function *Func,
if (!ASC || ASC->getSrcAddressSpace() != AMDGPUAS::FLAT_ADDRESS ||
ASC->getDestAddressSpace() != AMDGPUAS::LOCAL_ADDRESS)
return nullptr;
- Value *FlatSrc = ASC->getPointerOperand();
-
- // The flat source must trace back to an LDS-carrying arg.
- SmallPtrSet<Value *, 8> Visited;
- if (!isFlatPtrDerivedFromLDS(FlatSrc, LDSArgs, Visited,
- /*AcceptLocalCast=*/false))
- return nullptr;
// Re-apply the peeled GEPs in the flat address space.
- Value *Flat = FlatSrc;
+ Value *Flat = ASC->getPointerOperand();
for (GEPOperator *GEP : reverse(GEPs)) {
SmallVector<Value *, 4> Indices(GEP->idx_begin(), GEP->idx_end());
Flat = IRB.CreateGEP(GEP->getSourceElementType(), Flat, Indices, "",
@@ -984,16 +876,17 @@ Value *AMDGPUSwLowerLDS::getFlatPtrForRoundTripLDSAccess(Function *Func,
return Flat;
}
-void AMDGPUSwLowerLDS::lowerNonKernelLDSFlatArgAccesses(Function *Func) {
- // Lower a non-kernel that only reaches lowered LDS through a flat argument.
- // No base/offset table is needed: the flat pointer already points at the
- // global backing buffer, so the shared translation reuses the flat source.
+bool AMDGPUSwLowerLDS::lowerNonKernelLDSFlatArgAccesses(Function *Func) {
+ // Lower a non-kernel that reaches lowered LDS only through a flat round trip
+ // (addrspacecast flat->local). No base/offset table is needed: collapsing the
+ // round trip yields a direct flat access into the relocated buffer.
SetVector<Instruction *> LDSInstructions;
getLDSMemoryInstructions(Func, LDSInstructions);
if (LDSInstructions.empty())
- return;
+ return false;
translateLDSMemoryOperationsToGlobalMemory(Func, /*LoadMallocPtr=*/nullptr,
LDSInstructions);
+ return true;
}
void AMDGPUSwLowerLDS::poisonRedzones(Function *Func, Value *MallocPtr) {
@@ -1464,11 +1357,6 @@ bool AMDGPUSwLowerLDS::run() {
// Get address sanitizer scale.
initAsanInfo();
- // Discover non-kernels that receive lowered LDS via a flat pointer argument.
- // This must happen before lowering rewrites the call-site address-space
- // casts.
- getNonKernelsWithLDSFlatArguments();
-
for (auto &K : FuncLDSAccessInfo.KernelToLDSParametersMap) {
Function *Func = K.first;
auto &LDSParams = FuncLDSAccessInfo.KernelToLDSParametersMap[Func];
@@ -1531,15 +1419,18 @@ bool AMDGPUSwLowerLDS::run() {
Changed = true;
}
- // Lower non-kernels that reach lowered LDS only through a flat argument.
- // Functions also handled above already had these accesses translated there.
- for (auto &K : FuncLDSAccessInfo.NonKernelsWithLDSFlatArg) {
- Function *Func = K.first;
- if (FuncLDSAccessInfo.NonKernelToLDSAccessMap.contains(Func) ||
- FuncLDSAccessInfo.NonKernelsWithLDSArgument.contains(Func))
+ // Lower any remaining non-kernel that reaches lowered LDS through a flat
+ // round trip (addrspacecast flat->local). Functions handled above already had
+ // these accesses translated there, so skip them. No detection is needed:
+ // collapsing the round trip is valid regardless of the pointer's origin.
+ for (Function &F : M) {
+ if (F.isDeclaration() || AMDGPU::isKernel(F))
continue;
- lowerNonKernelLDSFlatArgAccesses(Func);
- Changed = true;
+ if (FuncLDSAccessInfo.NonKernelToLDSAccessMap.contains(&F) ||
+ FuncLDSAccessInfo.NonKernelsWithLDSArgument.contains(&F))
+ continue;
+ if (lowerNonKernelLDSFlatArgAccesses(&F))
+ Changed = true;
}
if (!Changed)
More information about the llvm-commits
mailing list